From 73294372970c8f5d285cd17170918f465bda36ca Mon Sep 17 00:00:00 2001 From: sususu98 <33882693+sususu98@users.noreply.github.com> Date: Mon, 20 Jul 2026 16:18:13 +0800 Subject: [PATCH] fix(websocket): propagate upstream 1009 without credential fallback (#4444) Preserve Codex and XAI upstream message-too-big closes across the Responses WebSocket boundary while treating them as request-scoped errors. Prevent CPA reconnect and credential cooldown/fallback, scope close state to the active connection, retain terminal errors under backpressure, and ensure downstream disconnect cleanup cannot wait behind a blocked writer. --- .../executor/codex_websockets_executor.go | 141 ++++++- .../codex_websockets_executor_test.go | 172 +++++++++ .../executor/xai_websockets_executor.go | 64 +++- .../executor/xai_websockets_executor_test.go | 130 +++++++ .../openai/openai_responses_websocket.go | 194 ++++++++-- .../openai/openai_responses_websocket_test.go | 354 ++++++++++++++++-- sdk/cliproxy/auth/conductor_overrides_test.go | 21 +- 7 files changed, 984 insertions(+), 92 deletions(-) diff --git a/internal/runtime/executor/codex_websockets_executor.go b/internal/runtime/executor/codex_websockets_executor.go index 823726cb..ae503fa7 100644 --- a/internal/runtime/executor/codex_websockets_executor.go +++ b/internal/runtime/executor/codex_websockets_executor.go @@ -79,8 +79,11 @@ type codexWebsocketSession struct { readerConn *websocket.Conn - upstreamDisconnectOnce sync.Once - upstreamDisconnectCh chan error + upstreamDisconnectOnce sync.Once + upstreamDisconnectCh chan error + upstreamDisconnectErrMu sync.RWMutex + upstreamDisconnectErrConn *websocket.Conn + upstreamDisconnectErr error } func NewCodexWebsocketsExecutor(cfg *config.Config) *CodexWebsocketsExecutor { @@ -169,16 +172,43 @@ func (s *codexWebsocketSession) writeMessage(conn *websocket.Conn, msgType int, return conn.WriteMessage(msgType, payload) } +// sendTerminalWebsocketRead reports whether it invalidated a full channel's connection before waiting. +func sendTerminalWebsocketRead(ch chan<- codexWebsocketRead, done <-chan struct{}, event codexWebsocketRead, invalidate func()) bool { + select { + case ch <- event: + return false + case <-done: + return false + default: + } + + invalidated := invalidate != nil + if invalidated { + invalidate() + } + select { + case ch <- event: + case <-done: + } + return invalidated +} + func (s *codexWebsocketSession) configureConn(conn *websocket.Conn) { if s == nil || conn == nil { return } + s.resetUpstreamDisconnectError(conn) conn.SetPingHandler(func(appData string) error { s.writeMu.Lock() defer s.writeMu.Unlock() // Reply pongs from the same write lock to avoid concurrent writes. return conn.WriteControl(websocket.PongMessage, []byte(appData), time.Now().Add(10*time.Second)) }) + defaultCloseHandler := conn.CloseHandler() + conn.SetCloseHandler(func(code int, text string) error { + s.setUpstreamDisconnectError(conn, &websocket.CloseError{Code: code, Text: text}) + return defaultCloseHandler(code, text) + }) } func websocketSessionTargetChanged(sess *codexWebsocketSession, authID string, wsURL string) bool { @@ -215,6 +245,39 @@ func detachMismatchedWebsocketSessionConn(sess *codexWebsocketSession, authID st return conn, previousAuthID, previousWSURL } +func (s *codexWebsocketSession) resetUpstreamDisconnectError(conn *websocket.Conn) { + if s == nil || conn == nil { + return + } + s.upstreamDisconnectErrMu.Lock() + s.upstreamDisconnectErrConn = conn + s.upstreamDisconnectErr = nil + s.upstreamDisconnectErrMu.Unlock() +} + +func (s *codexWebsocketSession) setUpstreamDisconnectError(conn *websocket.Conn, err error) { + if s == nil || conn == nil || err == nil { + return + } + s.upstreamDisconnectErrMu.Lock() + if s.upstreamDisconnectErrConn == conn && s.upstreamDisconnectErr == nil { + s.upstreamDisconnectErr = err + } + s.upstreamDisconnectErrMu.Unlock() +} + +func (s *codexWebsocketSession) upstreamDisconnectError(conn *websocket.Conn) error { + if s == nil || conn == nil { + return nil + } + s.upstreamDisconnectErrMu.RLock() + defer s.upstreamDisconnectErrMu.RUnlock() + if s.upstreamDisconnectErrConn != conn { + return nil + } + return s.upstreamDisconnectErr +} + func (s *codexWebsocketSession) notifyUpstreamDisconnect(err error) { if s == nil { return @@ -368,8 +431,13 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut } if errSend := writeCodexWebsocketMessage(sess, conn, wsReqBody); errSend != nil { + errSend = mapCodexWebsocketWriteError(sess, conn, errSend) if sess != nil { e.invalidateUpstreamConn(sess, conn, "send_error", errSend) + if !shouldRetryCodexWebsocketSend(errSend) { + helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend) + return resp, errSend + } // Retry once with a fresh websocket connection. This is mainly to handle // upstream closing the socket between sequential requests within the same @@ -395,6 +463,7 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut if errSendRetry := writeCodexWebsocketMessage(sess, conn, wsReqBodyRetry); errSendRetry == nil { wsReqBody = wsReqBodyRetry } else { + errSendRetry = mapCodexWebsocketWriteError(sess, connRetry, errSendRetry) e.invalidateUpstreamConn(sess, connRetry, "send_error", errSendRetry) helps.RecordAPIWebsocketError(ctx, e.cfg, "send_retry", errSendRetry) return resp, errSendRetry @@ -608,9 +677,15 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr } if errSend := writeCodexWebsocketMessage(sess, conn, wsReqBody); errSend != nil { + errSend = mapCodexWebsocketWriteError(sess, conn, errSend) helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend) if sess != nil { e.invalidateUpstreamConn(sess, conn, "send_error", errSend) + if !shouldRetryCodexWebsocketSend(errSend) { + sess.clearActive(conn, readCh) + sess.reqMu.Unlock() + return nil, errSend + } // Retry once with a new websocket connection for the same execution session. connRetry, respHSRetry, errDialRetry := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) @@ -638,6 +713,7 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr recordAPIWebsocketHandshake(ctx, e.cfg, respHSRetry) reporter.StartResponseTTFT() if errSendRetry := writeCodexWebsocketMessage(sess, conn, wsReqBodyRetry); errSendRetry != nil { + errSendRetry = mapCodexWebsocketWriteError(sess, conn, errSendRetry) helps.RecordAPIWebsocketError(ctx, e.cfg, "send_retry", errSendRetry) e.invalidateUpstreamConn(sess, conn, "send_error", errSendRetry) sess.clearActive(conn, readCh) @@ -847,13 +923,44 @@ func writeCodexWebsocketMessage(sess *codexWebsocketSession, conn *websocket.Con return conn.WriteMessage(websocket.TextMessage, payload) } +func mapCodexWebsocketWriteError(sess *codexWebsocketSession, conn *websocket.Conn, err error) error { + if err == nil || sess == nil || conn == nil { + return err + } + upstreamErr := sess.upstreamDisconnectError(conn) + var closeErr *websocket.CloseError + if !errors.As(upstreamErr, &closeErr) || closeErr.Code != websocket.CloseMessageTooBig { + return err + } + return mapCodexWebsocketReadError(upstreamErr) +} + +func shouldRetryCodexWebsocketSend(err error) bool { + if err == nil { + return false + } + var requestErr cliproxyexecutor.RequestScopedError + return !errors.As(err, &requestErr) || !requestErr.IsRequestScoped() +} + +type codexWebsocketMessageTooBigError struct { + statusErr +} + +func (codexWebsocketMessageTooBigError) IsRequestScoped() bool { + return true +} + func mapCodexWebsocketReadError(err error) error { if err == nil { return nil } var closeErr *websocket.CloseError if errors.As(err, &closeErr) && closeErr.Code == websocket.CloseMessageTooBig { - return statusErr{code: http.StatusRequestEntityTooLarge, msg: `{"error":{"message":"upstream websocket message too big","type":"invalid_request_error","code":"message_too_big"}}`} + return codexWebsocketMessageTooBigError{statusErr: statusErr{ + code: http.StatusRequestEntityTooLarge, + msg: `{"error":{"message":"upstream websocket message too big","type":"invalid_request_error","code":"message_too_big"}}`, + }} } return err } @@ -1600,36 +1707,40 @@ func (e *CodexWebsocketsExecutor) readUpstreamLoop(sess *codexWebsocketSession, _ = conn.SetReadDeadline(time.Now().Add(codexResponsesWebsocketIdleTimeout)) msgType, payload, errRead := conn.ReadMessage() if errRead != nil { + invalidate := func() { + e.invalidateUpstreamConn(sess, conn, "upstream_disconnected", errRead) + } + invalidated := false ch, done := sess.activeForConn(conn) if ch != nil { - select { - case ch <- codexWebsocketRead{conn: conn, err: errRead}: - case <-done: - default: - } + invalidated = sendTerminalWebsocketRead(ch, done, codexWebsocketRead{conn: conn, err: errRead}, invalidate) if sess.clearActive(conn, ch) { close(ch) } } - e.invalidateUpstreamConn(sess, conn, "upstream_disconnected", errRead) + if !invalidated { + invalidate() + } return } if msgType != websocket.TextMessage { if msgType == websocket.BinaryMessage { errBinary := fmt.Errorf("codex websockets executor: unexpected binary message") + invalidate := func() { + e.invalidateUpstreamConn(sess, conn, "unexpected_binary", errBinary) + } + invalidated := false ch, done := sess.activeForConn(conn) if ch != nil { - select { - case ch <- codexWebsocketRead{conn: conn, err: errBinary}: - case <-done: - default: - } + invalidated = sendTerminalWebsocketRead(ch, done, codexWebsocketRead{conn: conn, err: errBinary}, invalidate) if sess.clearActive(conn, ch) { close(ch) } } - e.invalidateUpstreamConn(sess, conn, "unexpected_binary", errBinary) + if !invalidated { + invalidate() + } return } continue diff --git a/internal/runtime/executor/codex_websockets_executor_test.go b/internal/runtime/executor/codex_websockets_executor_test.go index 555f8f62..2f3103ba 100644 --- a/internal/runtime/executor/codex_websockets_executor_test.go +++ b/internal/runtime/executor/codex_websockets_executor_test.go @@ -498,6 +498,173 @@ func TestCodexWebsocketsExecuteStreamPropagatesUpstreamErrorForDownstreamWebsock } } +func TestSendTerminalWebsocketReadInvalidatesBeforeWaitingForCapacity(t *testing.T) { + terminalErr := &websocket.CloseError{Code: websocket.CloseMessageTooBig} + + t.Run("available channel keeps fast path ordering", func(t *testing.T) { + ch := make(chan codexWebsocketRead, 1) + done := make(chan struct{}) + invalidateCalls := 0 + invalidated := sendTerminalWebsocketRead(ch, done, codexWebsocketRead{err: terminalErr}, func() { + invalidateCalls++ + }) + if invalidated { + t.Fatal("available channel should not invalidate before delivery") + } + if invalidateCalls != 0 { + t.Fatalf("invalidate calls = %d, want 0", invalidateCalls) + } + event := <-ch + if !errors.Is(event.err, terminalErr) { + t.Fatalf("terminal error = %v, want %v", event.err, terminalErr) + } + }) + + t.Run("full channel invalidates before waiting", func(t *testing.T) { + ch := make(chan codexWebsocketRead, 1) + ch <- codexWebsocketRead{payload: []byte("queued")} + done := make(chan struct{}) + invalidateCalled := make(chan struct{}) + result := make(chan bool, 1) + + go func() { + result <- sendTerminalWebsocketRead(ch, done, codexWebsocketRead{err: terminalErr}, func() { + close(invalidateCalled) + }) + }() + + select { + case <-invalidateCalled: + case <-time.After(time.Second): + t.Fatal("invalidation did not happen before waiting for channel capacity") + } + select { + case <-result: + t.Fatal("terminal sender returned before capacity was released") + default: + } + + <-ch + select { + case event := <-ch: + if !errors.Is(event.err, terminalErr) { + t.Fatalf("terminal error = %v, want %v", event.err, terminalErr) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for terminal read") + } + select { + case invalidated := <-result: + if !invalidated { + t.Fatal("full channel should report early invalidation") + } + case <-time.After(time.Second): + t.Fatal("terminal sender did not finish") + } + }) + + t.Run("full channel stops when invalidation cancels active read", func(t *testing.T) { + ch := make(chan codexWebsocketRead, 1) + ch <- codexWebsocketRead{payload: []byte("queued")} + done := make(chan struct{}) + invalidated := sendTerminalWebsocketRead(ch, done, codexWebsocketRead{err: terminalErr}, func() { + close(done) + }) + if !invalidated { + t.Fatal("full channel should report early invalidation") + } + if len(ch) != 1 { + t.Fatalf("channel length = %d, want queued payload only", len(ch)) + } + }) +} + +func TestMapCodexWebsocketWriteErrorStopsRetryForMessageTooBig(t *testing.T) { + networkWriteErr := errors.New("write: broken pipe") + tests := []struct { + name string + closeCode int + writeErr error + wantStatus int + wantRetry bool + }{ + { + name: "close sent after message too big is request scoped", + closeCode: websocket.CloseMessageTooBig, + writeErr: websocket.ErrCloseSent, + wantStatus: http.StatusRequestEntityTooLarge, + wantRetry: false, + }, + { + name: "network write error after message too big is request scoped", + closeCode: websocket.CloseMessageTooBig, + writeErr: networkWriteErr, + wantStatus: http.StatusRequestEntityTooLarge, + wantRetry: false, + }, + { + name: "other close keeps stale connection retry", + closeCode: websocket.CloseNormalClosure, + writeErr: websocket.ErrCloseSent, + wantRetry: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + sess := &codexWebsocketSession{} + conn := &websocket.Conn{} + sess.resetUpstreamDisconnectError(conn) + sess.setUpstreamDisconnectError(conn, &websocket.CloseError{Code: tt.closeCode}) + + mappedErr := mapCodexWebsocketWriteError(sess, conn, tt.writeErr) + if got := shouldRetryCodexWebsocketSend(mappedErr); got != tt.wantRetry { + t.Fatalf("shouldRetryCodexWebsocketSend() = %v, want %v; err=%v", got, tt.wantRetry, mappedErr) + } + if tt.wantStatus == 0 { + if !errors.Is(mappedErr, tt.writeErr) { + t.Fatalf("mapped error = %v, want %v", mappedErr, tt.writeErr) + } + return + } + statusErr, ok := mappedErr.(interface{ StatusCode() int }) + if !ok || statusErr.StatusCode() != tt.wantStatus { + t.Fatalf("mapped status = %v, want %d; err=%v", statusErr, tt.wantStatus, mappedErr) + } + requestErr, ok := mappedErr.(interface{ IsRequestScoped() bool }) + if !ok || !requestErr.IsRequestScoped() { + t.Fatalf("mapped error should be request scoped, got %T", mappedErr) + } + }) + } +} + +func TestMapCodexWebsocketWriteErrorDoesNotReusePriorConnectionClose(t *testing.T) { + sess := &codexWebsocketSession{} + priorConn := &websocket.Conn{} + replacementConn := &websocket.Conn{} + + sess.resetUpstreamDisconnectError(priorConn) + sess.setUpstreamDisconnectError(priorConn, &websocket.CloseError{Code: websocket.CloseMessageTooBig}) + priorErr := mapCodexWebsocketWriteError(sess, priorConn, websocket.ErrCloseSent) + if shouldRetryCodexWebsocketSend(priorErr) { + t.Fatalf("prior connection 1009 should not retry, got %v", priorErr) + } + + sess.resetUpstreamDisconnectError(replacementConn) + // A late close callback from the prior connection must not overwrite the + // replacement connection's close state. + sess.setUpstreamDisconnectError(priorConn, &websocket.CloseError{Code: websocket.CloseMessageTooBig}) + sess.setUpstreamDisconnectError(replacementConn, &websocket.CloseError{Code: websocket.CloseNormalClosure}) + replacementErr := mapCodexWebsocketWriteError(sess, replacementConn, websocket.ErrCloseSent) + if !errors.Is(replacementErr, websocket.ErrCloseSent) { + t.Fatalf("replacement connection error = %v, want %v", replacementErr, websocket.ErrCloseSent) + } + if !shouldRetryCodexWebsocketSend(replacementErr) { + t.Fatalf("replacement connection should keep stale-connection retry, got %v", replacementErr) + } +} + func TestCodexWebsocketsExecuteStreamMapsMessageTooBigClose(t *testing.T) { upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -555,6 +722,10 @@ func TestCodexWebsocketsExecuteStreamMapsMessageTooBigClose(t *testing.T) { if got := gjson.Get(chunk.Err.Error(), "error.code").String(); got != "message_too_big" { t.Fatalf("error code = %q, want message_too_big; err=%v", got, chunk.Err) } + requestErr, ok := chunk.Err.(interface{ IsRequestScoped() bool }) + if !ok || !requestErr.IsRequestScoped() { + t.Fatalf("message-too-big error should be request scoped, got %T", chunk.Err) + } case <-time.After(5 * time.Second): t.Fatal("timed out waiting for error stream chunk") } @@ -585,6 +756,7 @@ func TestCodexWebsocketsUpstreamDisconnectChanSignalsOnInvalidate(t *testing.T) defer func() { _ = conn.Close() }() exec := NewCodexWebsocketsExecutor(&config.Config{}) + exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} sessionID := "sess-1" disconnectCh := exec.UpstreamDisconnectChan(sessionID) if disconnectCh == nil { diff --git a/internal/runtime/executor/xai_websockets_executor.go b/internal/runtime/executor/xai_websockets_executor.go index 7fd8cdb5..1fe240a1 100644 --- a/internal/runtime/executor/xai_websockets_executor.go +++ b/internal/runtime/executor/xai_websockets_executor.go @@ -499,9 +499,15 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox } if errSend := writeCodexWebsocketMessage(sess, conn, wsReqBody); errSend != nil { + errSend = mapXAIWebsocketWriteError(sess, conn, errSend) helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend) if sess != nil { e.invalidateUpstreamConn(sess, conn, "send_error", errSend) + if !shouldRetryXAIWebsocketSend(errSend) { + sess.clearActive(conn, readCh) + sess.reqMu.Unlock() + return nil, errSend + } connRetry, respHSRetry, errDialRetry := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) if errDialRetry != nil || connRetry == nil { bodyErrRetry := websocketHandshakeBody(respHSRetry) @@ -532,8 +538,9 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox recordAPIWebsocketHandshake(ctx, e.cfg, respHSRetry) reporter.StartResponseTTFT() if errSendRetry := writeCodexWebsocketMessage(sess, conn, wsReqBodyRetry); errSendRetry != nil { + errSendRetry = mapXAIWebsocketWriteError(sess, connRetry, errSendRetry) helps.RecordAPIWebsocketError(ctx, e.cfg, "send_retry", errSendRetry) - e.invalidateUpstreamConn(sess, conn, "send_error", errSendRetry) + e.invalidateUpstreamConn(sess, connRetry, "send_error", errSendRetry) sess.clearActive(conn, readCh) sess.reqMu.Unlock() return nil, errSendRetry @@ -599,11 +606,12 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox _ = send(cliproxyexecutor.StreamChunk{Err: ctx.Err()}) return } + mappedErr := mapXAIWebsocketReadError(errRead) terminateReason = "read_error" - terminateErr = errRead - helps.RecordAPIWebsocketError(ctx, e.cfg, "read", errRead) - reporter.PublishFailure(ctx, errRead) - _ = send(cliproxyexecutor.StreamChunk{Err: errRead}) + terminateErr = mappedErr + helps.RecordAPIWebsocketError(ctx, e.cfg, "read", mappedErr) + reporter.PublishFailure(ctx, mappedErr) + _ = send(cliproxyexecutor.StreamChunk{Err: mappedErr}) return } if msgType != websocket.TextMessage { @@ -994,11 +1002,29 @@ func configureXAIWebsocketConn(sess *codexWebsocketSession, conn *websocket.Conn if sess == nil || conn == nil { return } + sess.resetUpstreamDisconnectError(conn) conn.SetPingHandler(func(appData string) error { sess.writeMu.Lock() defer sess.writeMu.Unlock() return conn.WriteControl(websocket.PongMessage, []byte(appData), time.Time{}) }) + defaultCloseHandler := conn.CloseHandler() + conn.SetCloseHandler(func(code int, text string) error { + sess.setUpstreamDisconnectError(conn, &websocket.CloseError{Code: code, Text: text}) + return defaultCloseHandler(code, text) + }) +} + +func mapXAIWebsocketReadError(err error) error { + return mapCodexWebsocketReadError(err) +} + +func mapXAIWebsocketWriteError(sess *codexWebsocketSession, conn *websocket.Conn, err error) error { + return mapCodexWebsocketWriteError(sess, conn, err) +} + +func shouldRetryXAIWebsocketSend(err error) bool { + return shouldRetryCodexWebsocketSend(err) } func readXAIWebsocketMessage(ctx context.Context, sess *codexWebsocketSession, conn *websocket.Conn, readCh chan codexWebsocketRead) (int, []byte, error) { @@ -1044,36 +1070,40 @@ func (e *XAIWebsocketsExecutor) readUpstreamLoop(sess *codexWebsocketSession, co for { msgType, payload, errRead := conn.ReadMessage() if errRead != nil { + invalidate := func() { + e.invalidateUpstreamConn(sess, conn, "upstream_disconnected", errRead) + } + invalidated := false ch, done := sess.activeForConn(conn) if ch != nil { - select { - case ch <- codexWebsocketRead{conn: conn, err: errRead}: - case <-done: - default: - } + invalidated = sendTerminalWebsocketRead(ch, done, codexWebsocketRead{conn: conn, err: errRead}, invalidate) if sess.clearActive(conn, ch) { close(ch) } } - e.invalidateUpstreamConn(sess, conn, "upstream_disconnected", errRead) + if !invalidated { + invalidate() + } return } if msgType != websocket.TextMessage { if msgType == websocket.BinaryMessage { errBinary := fmt.Errorf("xai websockets executor: unexpected binary message") + invalidate := func() { + e.invalidateUpstreamConn(sess, conn, "unexpected_binary", errBinary) + } + invalidated := false ch, done := sess.activeForConn(conn) if ch != nil { - select { - case ch <- codexWebsocketRead{conn: conn, err: errBinary}: - case <-done: - default: - } + invalidated = sendTerminalWebsocketRead(ch, done, codexWebsocketRead{conn: conn, err: errBinary}, invalidate) if sess.clearActive(conn, ch) { close(ch) } } - e.invalidateUpstreamConn(sess, conn, "unexpected_binary", errBinary) + if !invalidated { + invalidate() + } return } continue diff --git a/internal/runtime/executor/xai_websockets_executor_test.go b/internal/runtime/executor/xai_websockets_executor_test.go index 60fe4b14..6df8ade6 100644 --- a/internal/runtime/executor/xai_websockets_executor_test.go +++ b/internal/runtime/executor/xai_websockets_executor_test.go @@ -3,6 +3,7 @@ package executor import ( "bytes" "context" + "errors" "fmt" "io" "net/http" @@ -33,6 +34,135 @@ func TestXAIWebsocketsEnabledForConfigAPIKey(t *testing.T) { } } +func TestMapXAIWebsocketWriteErrorStopsRetryForMessageTooBig(t *testing.T) { + networkWriteErr := errors.New("write: broken pipe") + tests := []struct { + name string + closeCode int + writeErr error + wantStatus int + wantRetry bool + }{ + { + name: "close sent after message too big is request scoped", + closeCode: websocket.CloseMessageTooBig, + writeErr: websocket.ErrCloseSent, + wantStatus: http.StatusRequestEntityTooLarge, + wantRetry: false, + }, + { + name: "network write error after message too big is request scoped", + closeCode: websocket.CloseMessageTooBig, + writeErr: networkWriteErr, + wantStatus: http.StatusRequestEntityTooLarge, + wantRetry: false, + }, + { + name: "other close keeps stale connection retry", + closeCode: websocket.CloseNormalClosure, + writeErr: websocket.ErrCloseSent, + wantRetry: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + sess := &codexWebsocketSession{} + conn := &websocket.Conn{} + sess.resetUpstreamDisconnectError(conn) + sess.setUpstreamDisconnectError(conn, &websocket.CloseError{Code: tt.closeCode}) + + mappedErr := mapXAIWebsocketWriteError(sess, conn, tt.writeErr) + if got := shouldRetryXAIWebsocketSend(mappedErr); got != tt.wantRetry { + t.Fatalf("shouldRetryXAIWebsocketSend() = %v, want %v; err=%v", got, tt.wantRetry, mappedErr) + } + if tt.wantStatus == 0 { + if !errors.Is(mappedErr, tt.writeErr) { + t.Fatalf("mapped error = %v, want %v", mappedErr, tt.writeErr) + } + return + } + statusErr, ok := mappedErr.(interface{ StatusCode() int }) + if !ok || statusErr.StatusCode() != tt.wantStatus { + t.Fatalf("mapped status = %v, want %d; err=%v", statusErr, tt.wantStatus, mappedErr) + } + requestErr, ok := mappedErr.(interface{ IsRequestScoped() bool }) + if !ok || !requestErr.IsRequestScoped() { + t.Fatalf("mapped error should be request scoped, got %T", mappedErr) + } + }) + } +} + +func TestXAIWebsocketsExecuteStreamMapsMessageTooBigClose(t *testing.T) { + upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade websocket: %v", err) + return + } + defer func() { _ = conn.Close() }() + + if _, _, errRead := conn.ReadMessage(); errRead != nil { + t.Errorf("read upstream websocket message: %v", errRead) + return + } + deadline := time.Now().Add(time.Second) + closeMessage := websocket.FormatCloseMessage(websocket.CloseMessageTooBig, "message too big") + if errWrite := conn.WriteControl(websocket.CloseMessage, closeMessage, deadline); errWrite != nil { + t.Errorf("write close websocket message: %v", errWrite) + } + })) + defer server.Close() + + exec := NewXAIWebsocketsExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + Provider: "xai", + Attributes: map[string]string{ + "base_url": server.URL, + "websockets": "true", + }, + Metadata: map[string]any{"access_token": "xai-token"}, + } + req := cliproxyexecutor.Request{ + Model: "grok-4.5", + Payload: []byte(`{"model":"grok-4.5","input":[{"type":"message","role":"user","content":"hello"}]}`), + } + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FormatOpenAIResponse, + ResponseFormat: sdktranslator.FormatOpenAIResponse, + } + + result, err := exec.ExecuteStream(context.Background(), auth, req, opts) + if err != nil { + t.Fatalf("ExecuteStream() error = %v", err) + } + + select { + case chunk, ok := <-result.Chunks: + if !ok { + t.Fatal("stream closed before error chunk") + } + if chunk.Err == nil { + t.Fatal("error chunk Err = nil, want message-too-big error") + } + statusErr, ok := chunk.Err.(interface{ StatusCode() int }) + if !ok || statusErr.StatusCode() != http.StatusRequestEntityTooLarge { + t.Fatalf("status error = %v, want %d; err=%v", statusErr, http.StatusRequestEntityTooLarge, chunk.Err) + } + if got := gjson.Get(chunk.Err.Error(), "error.code").String(); got != "message_too_big" { + t.Fatalf("error code = %q, want message_too_big; err=%v", got, chunk.Err) + } + requestErr, ok := chunk.Err.(interface{ IsRequestScoped() bool }) + if !ok || !requestErr.IsRequestScoped() { + t.Fatalf("message-too-big error should be request scoped, got %T", chunk.Err) + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for error stream chunk") + } +} + func TestXAIWebsocketsExecuteStreamSendsResponseCreateWithPreviousResponseID(t *testing.T) { upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} capturedPayload := make(chan []byte, 1) diff --git a/sdk/api/handlers/openai/openai_responses_websocket.go b/sdk/api/handlers/openai/openai_responses_websocket.go index 3c63a543..b1f962df 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket.go +++ b/sdk/api/handlers/openai/openai_responses_websocket.go @@ -4,13 +4,17 @@ import ( "bytes" "context" "encoding/json" + "errors" "fmt" "io" "net/http" "sort" "strconv" "strings" + "sync" + "sync/atomic" "time" + "unicode/utf8" "github.com/gin-gonic/gin" "github.com/google/uuid" @@ -29,14 +33,15 @@ import ( ) const ( - wsRequestTypeCreate = "response.create" - wsRequestTypeAppend = "response.append" - wsEventTypeError = "error" - wsEventTypeCompleted = "response.completed" - wsEventTypeDone = "response.done" - wsDoneMarker = "[DONE]" - wsTurnStateHeader = "x-codex-turn-state" - wsTimelineBodyKey = "WEBSOCKET_TIMELINE_OVERRIDE" + wsRequestTypeCreate = "response.create" + wsRequestTypeAppend = "response.append" + wsEventTypeError = "error" + wsEventTypeCompleted = "response.completed" + wsEventTypeDone = "response.done" + wsDoneMarker = "[DONE]" + wsTurnStateHeader = "x-codex-turn-state" + wsTimelineBodyKey = "WEBSOCKET_TIMELINE_OVERRIDE" + wsCloseReasonMaxBytes = 123 codexLocalCompactionSummaryPrefix = "Another language model started to solve this problem and produced a summary of its thinking process. You also have access to the state of the tools that were used by that language model. Use this to build on the work that has already been done and avoid duplicating work. Here is the summary produced by the other language model, use the information in this summary to assist with your own analysis:" ) @@ -49,6 +54,133 @@ var responsesWebsocketUpgrader = websocket.Upgrader{ }, } +// writeWebsocketCloseForUpstreamError mirrors transport-level upstream close +// codes to the downstream WebSocket client before the connection is torn down. +// Without this the client only observes an abnormal closure (1006) and cannot +// apply its own close-code based handling (e.g. falling back to SSE on 1009). +func writeWebsocketCloseForUpstreamError(conn *websocket.Conn, err error) (bool, error) { + if conn == nil { + return false, nil + } + matched, payload := websocketClosePayloadForUpstreamError(err) + if !matched { + return false, nil + } + return true, conn.WriteControl(websocket.CloseMessage, payload, time.Time{}) +} + +func websocketClosePayloadForUpstreamError(err error) (bool, []byte) { + if err == nil { + return false, nil + } + + code := 0 + reason := "" + var closeErr *websocket.CloseError + if errors.As(err, &closeErr) && closeErr.Code == websocket.CloseMessageTooBig { + code = closeErr.Code + reason = closeErr.Text + } else { + type statusCoder interface { + StatusCode() int + } + var statusErr statusCoder + errText := err.Error() + if !errors.As(err, &statusErr) || statusErr.StatusCode() != http.StatusRequestEntityTooLarge || + gjson.Get(errText, "error.code").String() != "message_too_big" { + return false, nil + } + code = websocket.CloseMessageTooBig + reason = strings.TrimSpace(gjson.Get(errText, "error.message").String()) + } + if reason == "" { + reason = "message too big" + } + reason = truncateWebsocketCloseReason(reason, wsCloseReasonMaxBytes) + return true, websocket.FormatCloseMessage(code, reason) +} + +type responsesWebsocketWriter struct { + conn *websocket.Conn + writeMu sync.Mutex + closing atomic.Bool +} + +func newResponsesWebsocketWriter(conn *websocket.Conn) *responsesWebsocketWriter { + return &responsesWebsocketWriter{conn: conn} +} + +// closeForUpstreamError sends a best-effort close frame without waiting behind +// an active downstream data writer. If a data write already owns writeMu, the +// connection is closed immediately so the blocked writer and session can exit. +func (w *responsesWebsocketWriter) closeForUpstreamError(err error) (bool, error) { + if w == nil || w.conn == nil { + return false, nil + } + matched, payload := websocketClosePayloadForUpstreamError(err) + if !matched { + return false, nil + } + if !w.closing.CompareAndSwap(false, true) { + return true, nil + } + if !w.writeMu.TryLock() { + return true, w.conn.Close() + } + defer w.writeMu.Unlock() + + errWrite := w.conn.WriteControl(websocket.CloseMessage, payload, time.Time{}) + errClose := w.conn.Close() + if errWrite != nil { + return true, errWrite + } + return true, errClose +} + +func (w *responsesWebsocketWriter) closeForUpstreamDisconnect(err error) { + if w == nil || w.conn == nil { + return + } + if matched, _ := w.closeForUpstreamError(err); matched { + return + } + _ = w.conn.Close() +} + +func truncateWebsocketCloseReason(reason string, maxBytes int) string { + if maxBytes <= 0 { + return "" + } + if len(reason) <= maxBytes && utf8.ValidString(reason) { + return reason + } + + // Decode from the front so work and output stay bounded by maxBytes. + var truncated strings.Builder + truncated.Grow(min(len(reason), maxBytes)) + remaining := maxBytes + runeErrorSize := utf8.RuneLen(utf8.RuneError) + for len(reason) > 0 && remaining > 0 { + r, size := utf8.DecodeRuneInString(reason) + if r == utf8.RuneError && size == 1 { + if runeErrorSize > remaining { + break + } + truncated.WriteRune(utf8.RuneError) + reason = reason[1:] + remaining -= runeErrorSize + continue + } + if size > remaining { + break + } + truncated.WriteString(reason[:size]) + reason = reason[size:] + remaining -= size + } + return truncated.String() +} + type websocketTimelineAppender interface { Append(eventType string, payload []byte, timestamp time.Time) } @@ -222,6 +354,7 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { if err != nil { return } + writer := newResponsesWebsocketWriter(conn) passthroughSessionID := uuid.NewString() downstreamSessionKey := websocketDownstreamSessionKey(c.Request) retainResponsesWebsocketToolCaches(downstreamSessionKey) @@ -250,8 +383,8 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { select { case <-wsDone: return - case <-disconnectCh: - _ = conn.Close() + case disconnectErr := <-disconnectCh: + writer.closeForUpstreamDisconnect(disconnectErr) } }() } @@ -418,7 +551,7 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { if errMsg != nil { h.LoggingAPIResponseError(context.WithValue(context.Background(), "gin", c), errMsg) markAPIResponseTimestamp(c) - errorPayload, errWrite := writeResponsesWebsocketError(conn, wsTimelineLog, errMsg) + errorPayload, errWrite := writeResponsesWebsocketError(writer, wsTimelineLog, errMsg) log.Infof( "responses websocket: downstream_out id=%s type=%d event=%s payload=%s", passthroughSessionID, @@ -448,7 +581,7 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { lastResponseOutput = []byte("[]") lastResponseID = "" lastResponsePendingToolCallIDs = nil - if errWrite := writeResponsesWebsocketSyntheticPrewarm(c, conn, requestJSON, wsTimelineLog, passthroughSessionID); errWrite != nil { + if errWrite := writeResponsesWebsocketSyntheticPrewarm(c, writer, requestJSON, wsTimelineLog, passthroughSessionID); errWrite != nil { wsTerminateErr = errWrite return } @@ -502,7 +635,7 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { } dataChan, _, errChan := h.ExecuteStreamWithAuthManager(cliCtx, h.HandlerType(), modelName, requestJSON, "") - completedOutput, completedResponseID, completedPendingToolCallIDs, forwardErrMsg, errForward := h.forwardResponsesWebsocket(c, conn, cliCancel, dataChan, errChan, wsTimelineLog, passthroughSessionID) + completedOutput, completedResponseID, completedPendingToolCallIDs, forwardErrMsg, errForward := h.forwardResponsesWebsocket(c, writer, cliCancel, dataChan, errChan, wsTimelineLog, passthroughSessionID) if errForward != nil { wsTerminateErr = errForward log.Warnf("responses websocket: forward failed id=%s error=%v", passthroughSessionID, errForward) @@ -1290,7 +1423,7 @@ func shouldHandleResponsesWebsocketPrewarmLocally(rawJSON []byte, lastRequest [] func writeResponsesWebsocketSyntheticPrewarm( c *gin.Context, - conn *websocket.Conn, + writer *responsesWebsocketWriter, requestJSON []byte, wsTimelineLog websocketTimelineAppender, sessionID string, @@ -1308,7 +1441,7 @@ func writeResponsesWebsocketSyntheticPrewarm( // websocketPayloadEventType(payloads[i]), // websocketPayloadPreview(payloads[i]), // ) - if errWrite := writeResponsesWebsocketPayload(conn, wsTimelineLog, payloads[i], time.Now()); errWrite != nil { + if errWrite := writeResponsesWebsocketPayload(writer, wsTimelineLog, payloads[i], time.Now()); errWrite != nil { log.Warnf( "responses websocket: downstream_out write failed id=%s event=%s error=%v", sessionID, @@ -1439,7 +1572,7 @@ func normalizeJSONArrayRaw(raw []byte) string { func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( c *gin.Context, - conn *websocket.Conn, + writer *responsesWebsocketWriter, cancel handlers.APIHandlerCancelFunc, data <-chan []byte, errs <-chan *interfaces.ErrorMessage, @@ -1470,7 +1603,14 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( if errMsg != nil { h.LoggingAPIResponseError(context.WithValue(context.Background(), "gin", c), errMsg) markAPIResponseTimestamp(c) - errorPayload, errWrite := writeResponsesWebsocketError(conn, wsTimelineLog, errMsg) + if matched, errClose := writer.closeForUpstreamError(errMsg.Error); matched { + cancel(errMsg.Error) + if errClose != nil { + return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), errMsg, errClose + } + return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), errMsg, websocket.ErrCloseSent + } + errorPayload, errWrite := writeResponsesWebsocketError(writer, wsTimelineLog, errMsg) log.Infof( "responses websocket: downstream_out id=%s type=%d event=%s payload=%s", sessionID, @@ -1504,7 +1644,7 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( } h.LoggingAPIResponseError(context.WithValue(context.Background(), "gin", c), errMsg) markAPIResponseTimestamp(c) - errorPayload, errWrite := writeResponsesWebsocketError(conn, wsTimelineLog, errMsg) + errorPayload, errWrite := writeResponsesWebsocketError(writer, wsTimelineLog, errMsg) log.Infof( "responses websocket: downstream_out id=%s type=%d event=%s payload=%s", sessionID, @@ -1557,7 +1697,7 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( // websocketPayloadEventType(payloads[i]), // websocketPayloadPreview(payloads[i]), // ) - if errWrite := writeResponsesWebsocketPayload(conn, wsTimelineLog, payloads[i], time.Now()); errWrite != nil { + if errWrite := writeResponsesWebsocketPayload(writer, wsTimelineLog, payloads[i], time.Now()); errWrite != nil { log.Warnf( "responses websocket: downstream_out write failed id=%s event=%s error=%v", sessionID, @@ -1759,7 +1899,7 @@ func websocketJSONPayloadsFromChunk(chunk []byte) [][]byte { return payloads } -func writeResponsesWebsocketError(conn *websocket.Conn, wsTimelineLog websocketTimelineAppender, errMsg *interfaces.ErrorMessage) ([]byte, error) { +func writeResponsesWebsocketError(writer *responsesWebsocketWriter, wsTimelineLog websocketTimelineAppender, errMsg *interfaces.ErrorMessage) ([]byte, error) { status := http.StatusInternalServerError errText := http.StatusText(status) if errMsg != nil { @@ -1829,7 +1969,7 @@ func writeResponsesWebsocketError(conn *websocket.Conn, wsTimelineLog websocketT } } - return payload, writeResponsesWebsocketPayload(conn, wsTimelineLog, payload, time.Now()) + return payload, writeResponsesWebsocketPayload(writer, wsTimelineLog, payload, time.Now()) } func appendWebsocketEvent(builder *strings.Builder, eventType string, payload []byte) { @@ -1909,11 +2049,19 @@ func setWebsocketBody(c *gin.Context, key string, body string) { c.Set(key, []byte(trimmedBody)) } -func writeResponsesWebsocketPayload(conn *websocket.Conn, wsTimelineLog websocketTimelineAppender, payload []byte, timestamp time.Time) error { +func writeResponsesWebsocketPayload(writer *responsesWebsocketWriter, wsTimelineLog websocketTimelineAppender, payload []byte, timestamp time.Time) error { if wsTimelineLog != nil { wsTimelineLog.Append("response", payload, timestamp) } - return conn.WriteMessage(websocket.TextMessage, payload) + if writer == nil || writer.conn == nil { + return fmt.Errorf("responses websocket: writer is nil") + } + writer.writeMu.Lock() + defer writer.writeMu.Unlock() + if writer.closing.Load() { + return websocket.ErrCloseSent + } + return writer.conn.WriteMessage(websocket.TextMessage, payload) } func appendWebsocketTimelineDisconnect(timeline websocketTimelineAppender, err error, timestamp time.Time) { diff --git a/sdk/api/handlers/openai/openai_responses_websocket_test.go b/sdk/api/handlers/openai/openai_responses_websocket_test.go index c5f54b06..20203800 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket_test.go +++ b/sdk/api/handlers/openai/openai_responses_websocket_test.go @@ -11,6 +11,7 @@ import ( "sync" "testing" "time" + "unicode/utf8" "github.com/gin-gonic/gin" "github.com/gorilla/websocket" @@ -24,6 +25,275 @@ import ( "github.com/tidwall/gjson" ) +func TestWriteWebsocketCloseForUpstreamErrorMirrorsMessageTooBig(t *testing.T) { + tests := []struct { + name string + err error + reason string + }{ + { + name: "raw close error", + err: &websocket.CloseError{ + Code: websocket.CloseMessageTooBig, + Text: "message too big", + }, + reason: "message too big", + }, + { + name: "mapped stream error", + err: websocketPinnedFailoverStatusError{ + status: http.StatusRequestEntityTooLarge, + msg: `{"error":{"message":"upstream websocket message too big","code":"message_too_big"}}`, + }, + reason: "upstream websocket message too big", + }, + { + name: "multibyte reason stays valid", + err: &websocket.CloseError{ + Code: websocket.CloseMessageTooBig, + Text: strings.Repeat("🙂", 31), + }, + reason: strings.Repeat("🙂", 30), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + serverErr := make(chan error, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := responsesWebsocketUpgrader.Upgrade(w, r, nil) + if err != nil { + serverErr <- err + return + } + matched, errWrite := writeWebsocketCloseForUpstreamError(conn, tt.err) + if !matched && errWrite == nil { + errWrite = errors.New("message-too-big error did not match") + } + if errClose := conn.Close(); errWrite == nil { + errWrite = errClose + } + serverErr <- errWrite + })) + defer server.Close() + + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("dial websocket: %v", err) + } + defer func() { _ = conn.Close() }() + + if err = conn.SetReadDeadline(time.Now().Add(time.Second)); err != nil { + t.Fatalf("set read deadline: %v", err) + } + _, _, err = conn.ReadMessage() + var closeErr *websocket.CloseError + if !errors.As(err, &closeErr) { + t.Fatalf("expected websocket close error, got %v", err) + } + if closeErr.Code != websocket.CloseMessageTooBig { + t.Fatalf("expected close code 1009, got %d", closeErr.Code) + } + if closeErr.Text != tt.reason { + t.Fatalf("expected close reason %q, got %q", tt.reason, closeErr.Text) + } + if err = <-serverErr; err != nil { + t.Fatalf("close server websocket: %v", err) + } + }) + } +} + +func TestResponsesWebsocketWriterCloseDoesNotWaitForActiveDataWriter(t *testing.T) { + serverErrCh := make(chan error, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := responsesWebsocketUpgrader.Upgrade(w, r, nil) + if err != nil { + serverErrCh <- err + return + } + writer := newResponsesWebsocketWriter(conn) + + // Holding writeMu models a data writer blocked inside WriteMessage. The + // upstream-close path must hard-close the socket instead of waiting for it. + writer.writeMu.Lock() + closeDone := make(chan error, 1) + go func() { + matched, errClose := writer.closeForUpstreamError(&websocket.CloseError{ + Code: websocket.CloseMessageTooBig, + Text: "message too big", + }) + if !matched && errClose == nil { + errClose = errors.New("message-too-big error did not match") + } + closeDone <- errClose + }() + + select { + case errClose := <-closeDone: + writer.writeMu.Unlock() + serverErrCh <- errClose + case <-time.After(time.Second): + writer.writeMu.Unlock() + serverErrCh <- errors.New("close waited behind active data writer") + } + })) + defer server.Close() + + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("dial websocket: %v", err) + } + defer func() { _ = conn.Close() }() + if err = conn.SetReadDeadline(time.Now().Add(time.Second)); err != nil { + t.Fatalf("set read deadline: %v", err) + } + if _, _, err = conn.ReadMessage(); err == nil { + t.Fatal("client read succeeded, want connection closure") + } + if errServer := <-serverErrCh; errServer != nil { + t.Fatalf("server error: %v", errServer) + } +} + +func TestTruncateWebsocketCloseReason(t *testing.T) { + tests := []struct { + name string + reason string + maxBytes int + want string + }{ + { + name: "non-positive limit", + reason: "message too big", + maxBytes: 0, + want: "", + }, + { + name: "short valid reason unchanged", + reason: "message too big", + maxBytes: wsCloseReasonMaxBytes, + want: "message too big", + }, + { + name: "long ascii reason", + reason: strings.Repeat("x", 1<<20), + maxBytes: wsCloseReasonMaxBytes, + want: strings.Repeat("x", wsCloseReasonMaxBytes), + }, + { + name: "long invalid reason", + reason: strings.Repeat("\xff", 1<<20), + maxBytes: wsCloseReasonMaxBytes, + want: strings.Repeat("�", wsCloseReasonMaxBytes/utf8.RuneLen(utf8.RuneError)), + }, + { + name: "multibyte rune does not fit", + reason: "ab🙂cd", + maxBytes: 5, + want: "ab", + }, + { + name: "invalid bytes become replacement runes", + reason: string([]byte{'a', 0xff, 0xfe, 'b'}), + maxBytes: 8, + want: "a��b", + }, + { + name: "invalid replacement does not cross limit", + reason: string([]byte{'a', 'b', 0xff, 'c'}), + maxBytes: 4, + want: "ab", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := truncateWebsocketCloseReason(tt.reason, tt.maxBytes) + if got != tt.want { + t.Fatalf("truncateWebsocketCloseReason() = %q, want %q", got, tt.want) + } + if !utf8.ValidString(got) { + t.Fatalf("truncateWebsocketCloseReason() returned invalid UTF-8: %q", got) + } + if tt.maxBytes > 0 && len(got) > tt.maxBytes { + t.Fatalf("truncateWebsocketCloseReason() returned %d bytes, limit %d", len(got), tt.maxBytes) + } + }) + } +} + +func TestForwardResponsesWebsocketMirrorsMappedMessageTooBig(t *testing.T) { + gin.SetMode(gin.TestMode) + + serverErrCh := make(chan error, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := responsesWebsocketUpgrader.Upgrade(w, r, nil) + if err != nil { + serverErrCh <- err + return + } + defer func() { _ = conn.Close() }() + + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Request = r + data := make(chan []byte) + errCh := make(chan *interfaces.ErrorMessage, 1) + errCh <- &interfaces.ErrorMessage{ + StatusCode: http.StatusRequestEntityTooLarge, + Error: websocketPinnedFailoverStatusError{ + status: http.StatusRequestEntityTooLarge, + msg: `{"error":{"message":"upstream websocket message too big","code":"message_too_big"}}`, + }, + } + + h := NewOpenAIResponsesAPIHandler(handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil)) + _, _, _, errMsg, errForward := h.forwardResponsesWebsocket( + ctx, + newResponsesWebsocketWriter(conn), + func(...interface{}) {}, + data, + errCh, + newInMemoryWebsocketTimelineLog(), + "session-1", + ) + if errMsg == nil || errMsg.StatusCode != http.StatusRequestEntityTooLarge { + serverErrCh <- fmt.Errorf("forward error message = %#v, want status %d", errMsg, http.StatusRequestEntityTooLarge) + return + } + if !errors.Is(errForward, websocket.ErrCloseSent) { + serverErrCh <- fmt.Errorf("forward error = %v, want %v", errForward, websocket.ErrCloseSent) + return + } + serverErrCh <- nil + })) + defer server.Close() + + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("dial websocket: %v", err) + } + defer func() { _ = conn.Close() }() + + if err = conn.SetReadDeadline(time.Now().Add(time.Second)); err != nil { + t.Fatalf("set read deadline: %v", err) + } + _, _, err = conn.ReadMessage() + var closeErr *websocket.CloseError + if !errors.As(err, &closeErr) { + t.Fatalf("expected websocket close error, got %v", err) + } + if closeErr.Code != websocket.CloseMessageTooBig { + t.Fatalf("close code = %d, want %d", closeErr.Code, websocket.CloseMessageTooBig) + } + if err = <-serverErrCh; err != nil { + t.Fatalf("server error: %v", err) + } +} + type websocketCaptureExecutor struct { streamCalls int payloads [][]byte @@ -234,11 +504,17 @@ func (e *websocketDirectCaptureExecutor) AuthIDs() []string { type websocketUpstreamDisconnectExecutor struct { mu sync.Mutex + provider string subscribed chan string sessions map[string]chan error } -func (e *websocketUpstreamDisconnectExecutor) Identifier() string { return "codex" } +func (e *websocketUpstreamDisconnectExecutor) Identifier() string { + if provider := strings.TrimSpace(e.provider); provider != "" { + return provider + } + return "codex" +} func (e *websocketUpstreamDisconnectExecutor) UpstreamDisconnectChan(sessionID string) <-chan error { sessionID = strings.TrimSpace(sessionID) @@ -1369,7 +1645,7 @@ func TestForwardResponsesWebsocketRestoresAndForwardsCompletedOutput(t *testing. timelineLog := newInMemoryWebsocketTimelineLog() completedOutput, completedResponseID, pendingToolCallIDs, errMsg, err := (*OpenAIResponsesAPIHandler)(nil).forwardResponsesWebsocket( ctx, - conn, + newResponsesWebsocketWriter(conn), func(...interface{}) {}, data, errCh, @@ -1472,7 +1748,7 @@ func TestForwardResponsesWebsocketTreatsResponseDoneAsTerminalWithoutRewriting(t timelineLog := newInMemoryWebsocketTimelineLog() completedOutput, completedResponseID, pendingToolCallIDs, errMsg, err := (*OpenAIResponsesAPIHandler)(nil).forwardResponsesWebsocket( ctx, - conn, + newResponsesWebsocketWriter(conn), func(...interface{}) {}, data, errCh, @@ -1556,7 +1832,7 @@ func TestForwardResponsesWebsocketTreatsErrorPayloadAsTerminal(t *testing.T) { _, _, _, errMsg, err := (*OpenAIResponsesAPIHandler)(nil).forwardResponsesWebsocket( ctx, - conn, + newResponsesWebsocketWriter(conn), func(...interface{}) {}, data, errCh, @@ -1647,7 +1923,7 @@ func TestForwardResponsesWebsocketLogsAttemptedResponseOnWriteFailure(t *testing _, _, _, _, err = (*OpenAIResponsesAPIHandler)(nil).forwardResponsesWebsocket( ctx, - conn, + newResponsesWebsocketWriter(conn), func(...interface{}) {}, data, errCh, @@ -1747,40 +2023,54 @@ func TestResponsesWebsocketTimelineRecordsDisconnectEvent(t *testing.T) { } } -func TestResponsesWebsocketClosesOnCodexUpstreamDisconnect(t *testing.T) { +func TestResponsesWebsocketMirrorsUpstreamMessageTooBigDisconnect(t *testing.T) { gin.SetMode(gin.TestMode) - executor := &websocketUpstreamDisconnectExecutor{subscribed: make(chan string, 1)} - manager := coreauth.NewManager(nil, nil, nil) - manager.RegisterExecutor(executor) - base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager) - h := NewOpenAIResponsesAPIHandler(base) + for _, provider := range []string{"codex", "xai"} { + t.Run(provider, func(t *testing.T) { + executor := &websocketUpstreamDisconnectExecutor{provider: provider, subscribed: make(chan string, 1)} + manager := coreauth.NewManager(nil, nil, nil) + manager.RegisterExecutor(executor) + base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager) + h := NewOpenAIResponsesAPIHandler(base) - router := gin.New() - router.GET("/v1/responses/ws", h.ResponsesWebsocket) - server := httptest.NewServer(router) - defer server.Close() + router := gin.New() + router.GET("/v1/responses/ws", h.ResponsesWebsocket) + server := httptest.NewServer(router) + defer server.Close() - wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/v1/responses/ws" - conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) - if err != nil { - t.Fatalf("dial websocket: %v", err) - } - defer func() { _ = conn.Close() }() + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/v1/responses/ws" + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("dial websocket: %v", err) + } + defer func() { _ = conn.Close() }() - var sessionID string - select { - case sessionID = <-executor.subscribed: - case <-time.After(5 * time.Second): - t.Fatal("timed out waiting for upstream disconnect subscription") - } + var sessionID string + select { + case sessionID = <-executor.subscribed: + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for upstream disconnect subscription") + } - executor.TriggerDisconnect(sessionID, errors.New("upstream disconnected")) + executor.TriggerDisconnect(sessionID, &websocket.CloseError{ + Code: websocket.CloseMessageTooBig, + Text: "message too big", + }) - _ = conn.SetReadDeadline(time.Now().Add(2 * time.Second)) - _, _, err = conn.ReadMessage() - if err == nil { - t.Fatalf("expected downstream websocket to close after upstream disconnect") + _ = conn.SetReadDeadline(time.Now().Add(2 * time.Second)) + _, _, err = conn.ReadMessage() + var closeErr *websocket.CloseError + if !errors.As(err, &closeErr) { + t.Fatalf("expected downstream websocket close error, got %v", err) + } + if closeErr.Code != websocket.CloseMessageTooBig { + t.Fatalf("downstream close code = %d, want %d", closeErr.Code, websocket.CloseMessageTooBig) + } + if closeErr.Text != "message too big" { + t.Fatalf("downstream close reason = %q, want message too big", closeErr.Text) + } + }) } } diff --git a/sdk/cliproxy/auth/conductor_overrides_test.go b/sdk/cliproxy/auth/conductor_overrides_test.go index c4b606f8..3ee3ccfe 100644 --- a/sdk/cliproxy/auth/conductor_overrides_test.go +++ b/sdk/cliproxy/auth/conductor_overrides_test.go @@ -1162,6 +1162,10 @@ func TestManager_RequestScopedErrorStopsCredentialFallbackWithoutSuspendingAuth( status: http.StatusRequestTimeout, message: "stream error: stream disconnected before completion: stream closed before response.completed", } + messageTooBigErr := &requestScopedStatusError{ + status: http.StatusRequestEntityTooLarge, + message: `{"error":{"message":"upstream websocket message too big","type":"invalid_request_error","code":"message_too_big"}}`, + } invalidRequestErr := &Error{ HTTPStatus: http.StatusBadRequest, Message: `{"error":{"type":"invalid_request_error","code":"invalid_value","message":"Invalid input."}}`, @@ -1172,12 +1176,15 @@ func TestManager_RequestScopedErrorStopsCredentialFallbackWithoutSuspendingAuth( } tests := []struct { name string + provider string stream bool err error wantStatus int }{ {name: "non-streaming incomplete", err: incompleteErr, wantStatus: http.StatusRequestTimeout}, {name: "streaming incomplete", stream: true, err: incompleteErr, wantStatus: http.StatusRequestTimeout}, + {name: "streaming codex websocket message too big", provider: "codex", stream: true, err: messageTooBigErr, wantStatus: http.StatusRequestEntityTooLarge}, + {name: "streaming xai websocket message too big", provider: "xai", stream: true, err: messageTooBigErr, wantStatus: http.StatusRequestEntityTooLarge}, {name: "non-streaming invalid request", err: invalidRequestErr, wantStatus: http.StatusBadRequest}, {name: "streaming invalid request", stream: true, err: invalidRequestErr, wantStatus: http.StatusBadRequest}, {name: "non-streaming bad request", err: badRequestErr, wantStatus: http.StatusBadRequest}, @@ -1186,10 +1193,14 @@ func TestManager_RequestScopedErrorStopsCredentialFallbackWithoutSuspendingAuth( for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { + provider := tc.provider + if provider == "" { + provider = "codex" + } m := NewManager(nil, nil, nil) m.SetRetryConfig(2, 30*time.Second, 0) - executor := &authFallbackExecutor{id: "codex"} + executor := &authFallbackExecutor{id: provider} if tc.stream { executor.streamFirstErrors = map[string]error{"aa-bad-auth": tc.err} } else { @@ -1198,8 +1209,8 @@ func TestManager_RequestScopedErrorStopsCredentialFallbackWithoutSuspendingAuth( m.RegisterExecutor(executor) model := "gpt-5.5" - badAuth := &Auth{ID: "aa-bad-auth", Provider: "codex"} - goodAuth := &Auth{ID: "bb-good-auth", Provider: "codex"} + badAuth := &Auth{ID: "aa-bad-auth", Provider: provider} + goodAuth := &Auth{ID: "bb-good-auth", Provider: provider} reg := registry.GetGlobalRegistry() reg.RegisterClient(badAuth.ID, badAuth.Provider, []*registry.ModelInfo{{ID: model}}) @@ -1218,14 +1229,14 @@ func TestManager_RequestScopedErrorStopsCredentialFallbackWithoutSuspendingAuth( var errExecute error if tc.stream { - result, errStream := m.ExecuteStream(context.Background(), []string{"codex"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{Stream: true}) + result, errStream := m.ExecuteStream(context.Background(), []string{provider}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{Stream: true}) if result != nil { for range result.Chunks { } } errExecute = errStream } else { - _, errExecute = m.Execute(context.Background(), []string{"codex"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + _, errExecute = m.Execute(context.Background(), []string{provider}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) } if errExecute == nil { t.Fatal("expected request-scoped stream error")