Simplify WebSocket pool code structure

- Promote lastResponseID from wsConnection to wsPool so it naturally
  survives all connection transitions without manual threading
- Extract closeLocked(), dialLocked(), invalidateConn() helpers to
  eliminate duplicated connection lifecycle logic in Stream()
- Replace loop-of-one reconnect with a single dialLocked() call
- Extract sendResponseCreate() to deduplicate marshal+send between
  dialWebSocket() and sendOnExisting()
- Remove wsMaxReconnectAttempts constant (was always 1)
- Simplify wsConnection struct to just conn + createdAt

Net result: -18 lines, fewer code paths, same behavior.

Assisted-By: docker-agent

Author: dgageot

pkg/model/provider/openai/ws_pool.go

@@ -2,7 +2,6 @@ package openai
 
 import (
 	"context"
-	"encoding/json"
 	"fmt"
 	"log/slog"
 	"net/http"
@@ -17,23 +16,13 @@ const (
 	// wsMaxConnectionAge is the maximum lifetime of a WebSocket connection.
 	// OpenAI enforces a 60-minute limit; we reconnect slightly earlier.
 	wsMaxConnectionAge = 55 * time.Minute
-
-	// wsMaxReconnectAttempts is the maximum number of times a broken
-	// connection will be replaced with a fresh one within a single
-	// Stream call before the error is propagated to the caller.
-	wsMaxReconnectAttempts = 1
 )
 
 // wsConnection holds a WebSocket connection together with bookkeeping
 // metadata for the connection pool.
 type wsConnection struct {
 	conn      *websocket.Conn
 	createdAt time.Time
-
-	// lastResponseID is the ID of the most recent response completed on
-	// this connection. It can be passed as previous_response_id in subsequent
-	// requests to enable server-side context caching.
-	lastResponseID string
 }
 
 // isExpired returns true when the connection has been open longer than
@@ -50,6 +39,12 @@ type wsPool struct {
 	mu   sync.Mutex
 	conn *wsConnection
 
+	// lastResponseID is the ID of the most recent response completed on
+	// this pool. It can be passed as previous_response_id in subsequent
+	// requests to enable server-side context caching.
+	// It lives on the pool (not wsConnection) so it survives reconnections.
+	lastResponseID string
+
 	// wsURL is the WebSocket endpoint (e.g. wss://api.openai.com/v1/responses).
 	wsURL string
 
@@ -76,99 +71,89 @@ func (p *wsPool) Stream(
 	p.mu.Lock()
 	defer p.mu.Unlock()
 
-	// Close stale connections, preserving the last response ID.
-	var prevResponseID string
+	// Close stale connections.
 	if p.conn != nil && p.conn.isExpired() {
 		slog.Debug("Closing expired WebSocket connection",
 			"age", time.Since(p.conn.createdAt))
-		prevResponseID = p.conn.lastResponseID
-		_ = p.conn.conn.Close()
-		p.conn = nil
+		p.closeLocked()
 	}
 
 	// Establish a new connection if needed.
 	if p.conn == nil {
-		headers, err := p.headerFn(ctx)
-		if err != nil {
-			return nil, fmt.Errorf("websocket pool: headers: %w", err)
-		}
-
-		stream, err := dialWebSocket(ctx, p.wsURL, headers, params)
-		if err != nil {
-			return nil, err
-		}
-
-		p.conn = &wsConnection{
-			conn:           stream.conn,
-			createdAt:      time.Now(),
-			lastResponseID: prevResponseID,
-		}
-
-		return &pooledStream{pool: p, inner: stream}, nil
+		return p.dialLocked(ctx, params)
 	}
 
 	// Reuse existing connection: send a new response.create.
 	stream, err := sendOnExisting(p.conn.conn, params)
 	if err != nil {
-		// Connection is broken; tear down and retry with a fresh one.
-		// We only attempt wsMaxReconnectAttempts reconnections to avoid
-		// unbounded loops if the server keeps rejecting connections.
+		// Connection is broken; tear down and reconnect once.
 		slog.Warn("Existing WebSocket connection failed, reconnecting", "error", err)
-		prevResponseID := p.conn.lastResponseID
-		_ = p.conn.conn.Close()
-		p.conn = nil
-
-		var lastErr error
-		for attempt := range wsMaxReconnectAttempts {
-			headers, err2 := p.headerFn(ctx)
-			if err2 != nil {
-				lastErr = fmt.Errorf("websocket pool: headers on reconnect (attempt %d/%d): %w", attempt+1, wsMaxReconnectAttempts, err2)
-				continue
-			}
-			stream, err2 = dialWebSocket(ctx, p.wsURL, headers, params)
-			if err2 != nil {
-				lastErr = fmt.Errorf("websocket pool: reconnect (attempt %d/%d): %w", attempt+1, wsMaxReconnectAttempts, err2)
-				continue
-			}
-			p.conn = &wsConnection{
-				conn:           stream.conn,
-				createdAt:      time.Now(),
-				lastResponseID: prevResponseID,
-			}
-			return &pooledStream{pool: p, inner: stream}, nil
-		}
-		return nil, lastErr
+		p.closeLocked()
+		return p.dialLocked(ctx, params)
+	}
+
+	return &pooledStream{pool: p, inner: stream}, nil
+}
+
+// dialLocked opens a fresh WebSocket connection and stores it in the pool.
+// Caller must hold p.mu.
+func (p *wsPool) dialLocked(
+	ctx context.Context,
+	params responses.ResponseNewParams,
+) (*pooledStream, error) {
+	headers, err := p.headerFn(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("websocket pool: headers: %w", err)
+	}
+
+	stream, err := dialWebSocket(ctx, p.wsURL, headers, params)
+	if err != nil {
+		return nil, err
+	}
+
+	p.conn = &wsConnection{
+		conn:      stream.conn,
+		createdAt: time.Now(),
 	}
 
 	return &pooledStream{pool: p, inner: stream}, nil
 }
 
+// closeLocked closes the current connection. lastResponseID is preserved
+// on the pool so it survives reconnections. Caller must hold p.mu.
+func (p *wsPool) closeLocked() {
+	if p.conn == nil {
+		return
+	}
+	_ = p.conn.conn.Close()
+	p.conn = nil
+}
+
+// invalidateConn tears down the pooled connection if it matches conn.
+// Called by pooledStream.Close when the stream encountered an error,
+// so the pool does not hand out a broken connection.
+func (p *wsPool) invalidateConn(conn *websocket.Conn) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	if p.conn != nil && p.conn.conn == conn {
+		p.closeLocked()
+	}
+}
+
 // Close shuts down the pooled connection.
 func (p *wsPool) Close() {
 	p.mu.Lock()
 	defer p.mu.Unlock()
 
-	if p.conn != nil {
-		_ = p.conn.conn.Close()
-		p.conn = nil
-	}
+	p.closeLocked()
 }
 
 // sendOnExisting sends a response.create on an already-open connection and
 // returns a wsStream that reads events from it.
 func sendOnExisting(conn *websocket.Conn, params responses.ResponseNewParams) (*wsStream, error) {
-	paramsJSON, err := json.Marshal(params)
-	if err != nil {
-		return nil, fmt.Errorf("websocket: marshal params: %w", err)
-	}
-
-	msg := wsCreateMessage{
-		Type:   "response.create",
-		Params: paramsJSON,
-	}
-
-	if err := conn.WriteJSON(msg); err != nil {
-		return nil, fmt.Errorf("websocket: write response.create: %w", err)
+	if err := sendResponseCreate(conn, params); err != nil {
+		return nil, err
 	}
 
 	slog.Debug("WebSocket response.create sent (reused connection)")
@@ -196,9 +181,7 @@ func (s *pooledStream) Next() bool {
 	event := s.inner.Current()
 	if isTerminalEvent(event.Type) && event.Response.ID != "" {
 		s.pool.mu.Lock()
-		if s.pool.conn != nil {
-			s.pool.conn.lastResponseID = event.Response.ID
-		}
+		s.pool.lastResponseID = event.Response.ID
 		s.pool.mu.Unlock()
 	}
 
@@ -220,16 +203,8 @@ func (s *pooledStream) Close() error {
 	s.inner.done = true
 
 	if s.inner.Err() != nil {
-		// Connection is likely broken; tear it down so the pool
-		// doesn't hand out a dead socket.
-		s.pool.mu.Lock()
-		if s.pool.conn != nil && s.pool.conn.conn == s.inner.conn {
-			_ = s.pool.conn.conn.Close()
-			s.pool.conn = nil
-		}
-		s.pool.mu.Unlock()
+		s.pool.invalidateConn(s.inner.conn)
 	}
 
-	// Do NOT close the WebSocket connection when healthy—it stays in the pool.
 	return nil
 }

pkg/model/provider/openai/ws_stream.go

@@ -60,6 +60,26 @@ type wsStream struct {
 // Compile-time check: wsStream satisfies responseEventStream.
 var _ responseEventStream = (*wsStream)(nil)
 
+// sendResponseCreate marshals params and writes a response.create message
+// on the given WebSocket connection.
+func sendResponseCreate(conn *websocket.Conn, params responses.ResponseNewParams) error {
+	paramsJSON, err := json.Marshal(params)
+	if err != nil {
+		return fmt.Errorf("websocket: marshal params: %w", err)
+	}
+
+	msg := wsCreateMessage{
+		Type:   "response.create",
+		Params: paramsJSON,
+	}
+
+	if err := conn.WriteJSON(msg); err != nil {
+		return fmt.Errorf("websocket: write response.create: %w", err)
+	}
+
+	return nil
+}
+
 // dialWebSocket opens a WebSocket connection, sends the response.create
 // message, and returns a stream that yields server events.
 func dialWebSocket(
@@ -84,22 +104,9 @@ func dialWebSocket(
 		return nil, fmt.Errorf("websocket dial %s: %w", wsURL, err)
 	}
 
-	// Marshal the params using the SDK's MarshalJSON so all field
-	// encodings (omitzero, unions, etc.) are handled correctly.
-	paramsJSON, err := json.Marshal(params)
-	if err != nil {
-		conn.Close()
-		return nil, fmt.Errorf("websocket: marshal params: %w", err)
-	}
-
-	msg := wsCreateMessage{
-		Type:   "response.create",
-		Params: paramsJSON,
-	}
-
-	if err := conn.WriteJSON(msg); err != nil {
+	if err := sendResponseCreate(conn, params); err != nil {
 		conn.Close()
-		return nil, fmt.Errorf("websocket: write response.create: %w", err)
+		return nil, err
 	}
 
 	slog.Debug("WebSocket response.create sent", "url", wsURL)