reliability improvements

Author: Thomas Stromberg

pkg/client/client.go

@@ -57,6 +57,7 @@ type Config struct {
 	OnConnect      func()
 	ServerURL      string
 	Token          string
+	TokenProvider  func() (string, error) // Optional: dynamically provide fresh tokens for reconnection
 	Organization   string
 	EventTypes     []string
 	PullRequests   []string
@@ -98,8 +99,8 @@ func New(config Config) (*Client, error) {
 	if config.Organization == "" && len(config.PullRequests) == 0 {
 		return nil, errors.New("organization or pull requests required")
 	}
-	if config.Token == "" {
-		return nil, errors.New("token is required")
+	if config.Token == "" && config.TokenProvider == nil {
+		return nil, errors.New("token or tokenProvider is required")
 	}
 
 	// Set defaults
@@ -237,6 +238,17 @@ func (c *Client) Stop() {
 func (c *Client) connect(ctx context.Context) error {
 	c.logger.Info("Establishing WebSocket connection")
 
+	// Get fresh token if TokenProvider is configured
+	token := c.config.Token
+	if c.config.TokenProvider != nil {
+		t, err := c.config.TokenProvider()
+		if err != nil {
+			return fmt.Errorf("token provider: %w", err)
+		}
+		token = t
+		c.logger.Debug("Using fresh token from TokenProvider")
+	}
+
 	// Create WebSocket config with appropriate origin
 	origin := "http://localhost/"
 	if strings.HasPrefix(c.config.ServerURL, "wss://") {
@@ -249,7 +261,7 @@ func (c *Client) connect(ctx context.Context) error {
 
 	// Add Authorization header
 	wsConfig.Header = make(map[string][]string)
-	wsConfig.Header["Authorization"] = []string{fmt.Sprintf("Bearer %s", c.config.Token)}
+	wsConfig.Header["Authorization"] = []string{fmt.Sprintf("Bearer %s", token)}
 
 	// Dial the server
 	ws, err := websocket.DialConfig(wsConfig)
@@ -361,10 +373,10 @@ func (c *Client) connect(ctx context.Context) error {
 		c.logger.Error("SUBSCRIPTION REJECTED BY SERVER!", "error_code", errorCode, "message", message)
 		c.logger.Error(separatorLine)
 
-		// Return AuthenticationError for access denied errors to prevent retries
-		if errorCode == "access_denied" {
+		// Return AuthenticationError for authentication/authorization errors to prevent retries
+		if errorCode == "access_denied" || errorCode == "authentication_failed" {
 			return &AuthenticationError{
-				message: fmt.Sprintf("Access denied: %s", message),
+				message: fmt.Sprintf("Authentication/authorization failed: %s", message),
 			}
 		}
 
@@ -537,11 +549,6 @@ func (c *Client) readEvents(ctx context.Context, ws *websocket.Conn) error {
 			return fmt.Errorf("read: %w", err)
 		}
 
-		// Clear read timeout after successful read
-		if err := ws.SetReadDeadline(time.Time{}); err != nil {
-			c.logger.Warn("Failed to clear read timeout", "error", err)
-		}
-
 		// Check message type
 		responseType, ok := response[msgTypeField].(string)
 		if !ok {

pkg/srv/client.go

@@ -20,6 +20,7 @@ import (
 type Client struct {
 	conn         *websocket.Conn
 	send         chan Event
+	control      chan map[string]any // Control messages (pongs, shutdown notices)
 	hub          *Hub
 	done         chan struct{}
 	userOrgs     map[string]bool
@@ -49,7 +50,8 @@ func NewClient(id string, sub Subscription, conn *websocket.Conn, hub *Hub, user
 		ID:           id,
 		subscription: sub,
 		conn:         conn,
-		send:         make(chan Event, 100), // Increased buffer to reduce dropped messages
+		send:         make(chan Event, 100),        // Increased buffer to reduce dropped messages
+		control:      make(chan map[string]any, 5), // Buffer for control messages (pongs, shutdown)
 		hub:          hub,
 		done:         make(chan struct{}),
 		userOrgs:     orgsMap,
@@ -98,6 +100,18 @@ func (c *Client) Run(ctx context.Context, pingInterval, writeTimeout time.Durati
 				return
 			}
 
+		case ctrl, ok := <-c.control:
+			if !ok {
+				log.Printf("client %s: control channel closed", c.ID)
+				return
+			}
+
+			// Send control message (pong, shutdown notice, etc.)
+			if err := c.write(ctrl, writeTimeout); err != nil {
+				log.Printf("client %s: control message send failed: %v", c.ID, err)
+				return
+			}
+
 		case event, ok := <-c.send:
 			if !ok {
 				log.Printf("client %s: send channel closed", c.ID)
@@ -133,5 +147,6 @@ func (c *Client) Close() {
 	c.closeOnce.Do(func() {
 		close(c.done)
 		close(c.send)
+		close(c.control)
 	})
 }

pkg/srv/websocket.go

@@ -233,6 +233,10 @@ func sendErrorResponse(ws *websocket.Conn, errInfo errorInfo, ip string) error {
 		return err
 	}
 
+	// Allow time for client to receive error before connection closes.
+	// Without this delay, TCP close can race with message delivery, causing clients to see EOF.
+	time.Sleep(100 * time.Millisecond)
+
 	return nil
 }
 
@@ -418,20 +422,20 @@ func (h *WebSocketHandler) validateAuth(ctx context.Context, ws *websocket.Conn,
 }
 
 // closeWebSocket gracefully closes a WebSocket connection with cleanup.
-func closeWebSocket(ws *websocket.Conn) {
+// If client is provided, shutdown message is sent via control channel to avoid race.
+func closeWebSocket(ws *websocket.Conn, client *Client) {
 	clientIP := security.ClientIP(ws.Request())
 	log.Printf("WebSocket Handle() cleanup - closing connection for IP %s", clientIP)
 
-	// Send a final shutdown message to allow graceful client disconnect
-	shutdownMsg := map[string]string{"type": "server_closing", "code": "1001"}
-	if err := ws.SetWriteDeadline(time.Now().Add(200 * time.Millisecond)); err != nil {
-		log.Printf("failed to set write deadline for shutdown message: %v", err)
-	}
-	if err := websocket.JSON.Send(ws, shutdownMsg); err != nil {
-		// Expected during abrupt disconnection - don't log common cases
-		if !strings.Contains(err.Error(), "use of closed network connection") &&
-			!strings.Contains(err.Error(), "broken pipe") {
-			log.Printf("failed to send shutdown message: %v", err)
+	// Send shutdown message via control channel if client exists
+	if client != nil {
+		shutdownMsg := map[string]any{"type": "server_closing", "code": "1001"}
+		select {
+		case client.control <- shutdownMsg:
+			// Give brief time for shutdown message to be sent
+			time.Sleep(100 * time.Millisecond)
+		case <-time.After(200 * time.Millisecond):
+			log.Printf("Timeout sending shutdown message to client %s", client.ID)
 		}
 	}
 
@@ -460,8 +464,11 @@ func (h *WebSocketHandler) Handle(ws *websocket.Conn) {
 	ctx, cancel := context.WithCancel(ws.Request().Context())
 	defer cancel()
 
-	// Ensure WebSocket is properly closed
-	defer closeWebSocket(ws)
+	// Ensure WebSocket is properly closed (client will be set later if connection succeeds)
+	var client *Client
+	defer func() {
+		closeWebSocket(ws, client)
+	}()
 
 	// Get client IP
 	ip := security.ClientIP(ws.Request())
@@ -645,7 +652,7 @@ func (h *WebSocketHandler) Handle(ws *websocket.Conn) {
 		}
 		id[i] = charset[n.Int64()]
 	}
-	client := NewClient(
+	client = NewClient(
 		string(id),
 		sub,
 		ws,
@@ -765,17 +772,17 @@ func (h *WebSocketHandler) Handle(ws *websocket.Conn) {
 					// (which happens for ANY message including pong) keeps the connection alive
 					continue
 				case "ping":
-					// Client sent us a ping, send pong back
+					// Client sent us a ping, send pong back via control channel to avoid race
 					pong := map[string]any{"type": "pong"}
 					if seq, ok := msgMap["seq"]; ok {
 						pong["seq"] = seq
 					}
-					if err := ws.SetWriteDeadline(time.Now().Add(writeTimeout)); err != nil {
-						log.Printf("failed to set write deadline for pong to client %s: %v", client.ID, err)
-						continue
-					}
-					if err := websocket.JSON.Send(ws, pong); err != nil {
-						log.Printf("failed to send pong to client %s: %v", client.ID, err)
+					// Non-blocking send to avoid deadlock if control channel is full
+					select {
+					case client.control <- pong:
+						// Pong queued successfully
+					default:
+						log.Printf("WARNING: client %s control channel full, dropping pong", client.ID)
 					}
 					continue
 				case "keepalive", "heartbeat":