From e05ae0942507e088589ba63de318afd2cb83167e Mon Sep 17 00:00:00 2001 From: sususu Date: Thu, 23 Jul 2026 13:06:33 +0800 Subject: [PATCH 1/3] fix(responses): preserve context across websocket transport changes --- internal/runtime/executor/codex_executor.go | 3 + .../executor/codex_websockets_executor.go | 88 +- .../codex_websockets_executor_test.go | 168 +++ .../executor/xai_websockets_executor.go | 30 +- .../executor/xai_websockets_executor_test.go | 44 + sdk/api/handlers/handlers.go | 26 +- .../handlers/handlers_model_router_test.go | 32 + .../openai/openai_responses_websocket.go | 291 ++++-- .../openai/openai_responses_websocket_test.go | 955 +++++++++++++++--- sdk/cliproxy/executor/context.go | 19 + sdk/cliproxy/executor/websocket.go | 29 + sdk/cliproxy/executor/websocket_test.go | 25 + 12 files changed, 1499 insertions(+), 211 deletions(-) create mode 100644 sdk/cliproxy/executor/websocket.go create mode 100644 sdk/cliproxy/executor/websocket_test.go diff --git a/internal/runtime/executor/codex_executor.go b/internal/runtime/executor/codex_executor.go index 4847fbf00..457a210bb 100644 --- a/internal/runtime/executor/codex_executor.go +++ b/internal/runtime/executor/codex_executor.go @@ -1139,6 +1139,7 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re body = helps.SetStringIfDifferent(body, "model", baseModel) body = helps.SetBoolIfDifferent(body, "stream", true) body, _ = sjson.DeleteBytes(body, "previous_response_id") + body, _ = sjson.DeleteBytes(body, "generate") body, _ = sjson.DeleteBytes(body, "prompt_cache_retention") body, _ = sjson.DeleteBytes(body, "safety_identifier") body, _ = sjson.DeleteBytes(body, "stream_options") @@ -1408,6 +1409,7 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au requestPath := helps.PayloadRequestPath(opts) body = helps.ApplyPayloadConfigWithRequest(e.cfg, baseModel, to.String(), from.String(), "", body, originalTranslated, requestedModel, requestPath, opts.Headers) body, _ = sjson.DeleteBytes(body, "previous_response_id") + body, _ = sjson.DeleteBytes(body, "generate") body, _ = sjson.DeleteBytes(body, "prompt_cache_retention") body, _ = sjson.DeleteBytes(body, "safety_identifier") body, _ = sjson.DeleteBytes(body, "stream_options") @@ -1580,6 +1582,7 @@ func (e *CodexExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth body = helps.SetStringIfDifferent(body, "model", baseModel) body, _ = sjson.DeleteBytes(body, "previous_response_id") + body, _ = sjson.DeleteBytes(body, "generate") body, _ = sjson.DeleteBytes(body, "prompt_cache_retention") body, _ = sjson.DeleteBytes(body, "safety_identifier") body, _ = sjson.DeleteBytes(body, "stream_options") diff --git a/internal/runtime/executor/codex_websockets_executor.go b/internal/runtime/executor/codex_websockets_executor.go index 266612cf1..a63eea635 100644 --- a/internal/runtime/executor/codex_websockets_executor.go +++ b/internal/runtime/executor/codex_websockets_executor.go @@ -224,6 +224,22 @@ func websocketSessionTargetChanged(sess *codexWebsocketSession, authID string, w return strings.TrimSpace(sess.authID) != strings.TrimSpace(authID) || strings.TrimSpace(sess.wsURL) != strings.TrimSpace(wsURL) } +func existingWebsocketSessionConn(sess *codexWebsocketSession, authID string, wsURL string) *websocket.Conn { + if sess == nil { + return nil + } + sess.connMu.Lock() + conn := sess.conn + matches := conn != nil && + strings.TrimSpace(sess.authID) == strings.TrimSpace(authID) && + strings.TrimSpace(sess.wsURL) == strings.TrimSpace(wsURL) + sess.connMu.Unlock() + if !matches || sess.upstreamDisconnectError(conn) != nil { + return nil + } + return conn +} + func detachMismatchedWebsocketSessionConn(sess *codexWebsocketSession, authID string, wsURL string) (*websocket.Conn, string, string) { if sess == nil { return nil, "", "" @@ -391,14 +407,24 @@ func (e *CodexWebsocketsExecutor) Execute(ctx context.Context, auth *cliproxyaut } helps.RecordAPIWebsocketRequest(ctx, e.cfg, wsReqLog) - conn, respHS, errDial := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) + var conn *websocket.Conn + var respHS *http.Response + var errDial error + if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) { + conn = existingWebsocketSessionConn(sess, authID, wsURL) + if conn == nil { + return resp, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() + } + } else { + conn, respHS, errDial = e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) + } if errDial != nil { bodyErr := websocketHandshakeBody(respHS) if respHS != nil { helps.RecordAPIWebsocketUpgradeRejection(ctx, e.cfg, websocketUpgradeRequestLog(wsReqLog), respHS.StatusCode, respHS.Header.Clone(), bodyErr) } if respHS != nil && respHS.StatusCode == http.StatusUpgradeRequired { - return e.CodexExecutor.Execute(ctx, auth, req, opts) + return resp, statusErr{code: http.StatusUpgradeRequired, msg: string(bodyErr)} } if respHS != nil && respHS.StatusCode > 0 { return resp, statusErr{code: respHS.StatusCode, msg: string(bodyErr)} @@ -433,6 +459,14 @@ 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 { + if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) { + e.invalidateUpstreamConnWithoutDisconnectNotify(sess, conn, "send_error", errSend) + if !shouldRetryCodexWebsocketSend(errSend) { + helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend) + return resp, errSend + } + return resp, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() + } e.invalidateUpstreamConn(sess, conn, "send_error", errSend) if !shouldRetryCodexWebsocketSend(errSend) { helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend) @@ -642,7 +676,20 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr } helps.RecordAPIWebsocketRequest(ctx, e.cfg, wsReqLog) - conn, respHS, errDial := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) + var conn *websocket.Conn + var respHS *http.Response + var errDial error + if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) { + conn = existingWebsocketSessionConn(sess, authID, wsURL) + if conn == nil { + if sess != nil { + sess.reqMu.Unlock() + } + return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() + } + } else { + conn, respHS, errDial = e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) + } var upstreamHeaders http.Header if respHS != nil { upstreamHeaders = respHS.Header.Clone() @@ -653,9 +700,15 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr helps.RecordAPIWebsocketUpgradeRejection(ctx, e.cfg, websocketUpgradeRequestLog(wsReqLog), respHS.StatusCode, respHS.Header.Clone(), bodyErr) } if respHS != nil && respHS.StatusCode == http.StatusUpgradeRequired { - return e.CodexExecutor.ExecuteStream(ctx, auth, req, opts) + if sess != nil { + sess.reqMu.Unlock() + } + return nil, statusErr{code: http.StatusUpgradeRequired, msg: string(bodyErr)} } if respHS != nil && respHS.StatusCode > 0 { + if sess != nil { + sess.reqMu.Unlock() + } return nil, statusErr{code: respHS.StatusCode, msg: string(bodyErr)} } helps.RecordAPIWebsocketError(ctx, e.cfg, "dial", errDial) @@ -680,6 +733,15 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr errSend = mapCodexWebsocketWriteError(sess, conn, errSend) helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend) if sess != nil { + if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) { + e.invalidateUpstreamConnWithoutDisconnectNotify(sess, conn, "send_error", errSend) + sess.clearActive(conn, readCh) + sess.reqMu.Unlock() + if !shouldRetryCodexWebsocketSend(errSend) { + return nil, errSend + } + return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() + } e.invalidateUpstreamConn(sess, conn, "send_error", errSend) if !shouldRetryCodexWebsocketSend(errSend) { sess.clearActive(conn, readCh) @@ -1759,6 +1821,14 @@ func (e *CodexWebsocketsExecutor) readUpstreamLoop(sess *codexWebsocketSession, } func (e *CodexWebsocketsExecutor) invalidateUpstreamConn(sess *codexWebsocketSession, conn *websocket.Conn, reason string, err error) { + e.invalidateUpstreamConnWithNotify(sess, conn, reason, err, true) +} + +func (e *CodexWebsocketsExecutor) invalidateUpstreamConnWithoutDisconnectNotify(sess *codexWebsocketSession, conn *websocket.Conn, reason string, err error) { + e.invalidateUpstreamConnWithNotify(sess, conn, reason, err, false) +} + +func (e *CodexWebsocketsExecutor) invalidateUpstreamConnWithNotify(sess *codexWebsocketSession, conn *websocket.Conn, reason string, err error, notify bool) { if sess == nil || conn == nil { return } @@ -1779,7 +1849,9 @@ func (e *CodexWebsocketsExecutor) invalidateUpstreamConn(sess *codexWebsocketSes sess.connMu.Unlock() logCodexWebsocketDisconnected(sessionID, authID, wsURL, reason, err) - sess.notifyUpstreamDisconnect(err) + if notify { + sess.notifyUpstreamDisconnect(err) + } if errClose := conn.Close(); errClose != nil { log.Errorf("codex websockets executor: close websocket error: %v", errClose) } @@ -1984,6 +2056,9 @@ func (e *CodexAutoExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth if cliproxyexecutor.DownstreamWebsocket(ctx) && codexWebsocketsEnabled(auth) { return e.wsExec.Execute(ctx, auth, req, opts) } + if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) { + return cliproxyexecutor.Response{}, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() + } return e.httpExec.Execute(ctx, auth, req, opts) } @@ -1994,6 +2069,9 @@ func (e *CodexAutoExecutor) ExecuteStream(ctx context.Context, auth *cliproxyaut if cliproxyexecutor.DownstreamWebsocket(ctx) && codexWebsocketsEnabled(auth) { return e.wsExec.ExecuteStream(ctx, auth, req, opts) } + if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) { + return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() + } return e.httpExec.ExecuteStream(ctx, auth, req, opts) } diff --git a/internal/runtime/executor/codex_websockets_executor_test.go b/internal/runtime/executor/codex_websockets_executor_test.go index 2f3103ba9..6c9a611ba 100644 --- a/internal/runtime/executor/codex_websockets_executor_test.go +++ b/internal/runtime/executor/codex_websockets_executor_test.go @@ -360,6 +360,174 @@ func TestCodexWebsocketsExecutePreservesPreviousResponseIDUpstream(t *testing.T) } } +func TestCodexWebsocketsExecuteStreamUpgradeRequiredReturnsWithoutLockingSession(t *testing.T) { + upgradeAttempts := make(chan struct{}, 2) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !strings.EqualFold(r.Header.Get("Upgrade"), "websocket") { + t.Errorf("unexpected HTTP fallback request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusInternalServerError) + return + } + upgradeAttempts <- struct{}{} + w.WriteHeader(http.StatusUpgradeRequired) + _, _ = w.Write([]byte(`{"error":{"message":"websocket unavailable"}}`)) + })) + defer server.Close() + + exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) + const executionSessionID = "ws-upgrade-required-session" + t.Cleanup(func() { exec.CloseExecutionSession(executionSessionID) }) + auth := &cliproxyauth.Auth{ + ID: "codex-test", + Provider: "codex", + Attributes: map[string]string{ + "api_key": "sk-test", + "base_url": server.URL, + }, + } + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai-response"), + ResponseFormat: sdktranslator.FromString("openai-response"), + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: executionSessionID, + }, + } + ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + + execute := func(payload string) { + t.Helper() + done := make(chan error, 1) + go func() { + _, errExecute := exec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{ + Model: "gpt-5.4", + Payload: []byte(payload), + }, opts) + done <- errExecute + }() + + select { + case errExecute := <-done: + if errExecute == nil { + t.Fatal("upgrade-required error = nil") + } + statusErr, ok := errExecute.(interface{ StatusCode() int }) + if !ok || statusErr.StatusCode() != http.StatusUpgradeRequired { + t.Fatalf("upgrade-required error = %T %v, want status 426", errExecute, errExecute) + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for upgrade-required error; execution session may still be locked") + } + } + + execute(`{"model":"gpt-5.4","generate":false,"input":[]}`) + execute(`{"model":"gpt-5.4","previous_response_id":"resp-1","input":[{"type":"message","id":"msg-2"}]}`) + + if got := len(upgradeAttempts); got != 2 { + t.Fatalf("websocket upgrade attempts = %d, want 2", got) + } +} + +func TestCodexWebsocketsExecuteStreamHandshakeErrorReturnsWithoutLockingSession(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"error":{"message":"unauthorized"}}`)) + })) + defer server.Close() + + exec := NewCodexWebsocketsExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) + const executionSessionID = "ws-handshake-error-session" + t.Cleanup(func() { exec.CloseExecutionSession(executionSessionID) }) + auth := &cliproxyauth.Auth{ + ID: "codex-test", + Provider: "codex", + Attributes: map[string]string{ + "api_key": "sk-test", + "base_url": server.URL, + }, + } + opts := cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai-response"), + Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: executionSessionID, + }, + } + + for i := 0; i < 2; i++ { + done := make(chan error, 1) + go func() { + _, errExecute := exec.ExecuteStream(context.Background(), auth, cliproxyexecutor.Request{ + Model: "gpt-5.4", + Payload: []byte(`{"model":"gpt-5.4","input":[{"type":"message","id":"msg-1"}]}`), + }, opts) + done <- errExecute + }() + select { + case errExecute := <-done: + statusErr, ok := errExecute.(interface{ StatusCode() int }) + if !ok || statusErr.StatusCode() != http.StatusUnauthorized { + t.Fatalf("attempt %d error = %T %v, want status 401", i+1, errExecute, errExecute) + } + case <-time.After(5 * time.Second): + t.Fatalf("attempt %d timed out; execution session remained locked", i+1) + } + } +} + +func TestExistingWebsocketSessionConnRequiresMatchingHealthyConnection(t *testing.T) { + conn := &websocket.Conn{} + sess := &codexWebsocketSession{ + conn: conn, + authID: "auth-a", + wsURL: "ws://example.test/responses", + } + sess.resetUpstreamDisconnectError(conn) + if got := existingWebsocketSessionConn(sess, "auth-a", "ws://example.test/responses"); got != conn { + t.Fatal("matching healthy websocket session was not reusable") + } + if got := existingWebsocketSessionConn(sess, "auth-b", "ws://example.test/responses"); got != nil { + t.Fatal("websocket session matched a different auth") + } + if got := existingWebsocketSessionConn(sess, "auth-a", "ws://other.test/responses"); got != nil { + t.Fatal("websocket session matched a different URL") + } + sess.setUpstreamDisconnectError(conn, errors.New("upstream disconnected")) + if got := existingWebsocketSessionConn(sess, "auth-a", "ws://example.test/responses"); got != nil { + t.Fatal("disconnected websocket session remained reusable") + } +} + +func TestCodexAutoExecutorRequiredUpstreamWebsocketRejectsHTTPFallback(t *testing.T) { + exec := NewCodexAutoExecutor(&config.Config{SDKConfig: config.SDKConfig{DisableImageGeneration: config.DisableImageGenerationAll}}) + auth := &cliproxyauth.Auth{ + ID: "codex-http-only", + Provider: "codex", + Attributes: map[string]string{ + "api_key": "sk-test", + }, + } + ctx := cliproxyexecutor.WithRequiredUpstreamWebsocket( + cliproxyexecutor.WithDownstreamWebsocket(context.Background()), + ) + _, errExecute := exec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{ + Model: "gpt-5.4", + Payload: []byte(`{"model":"gpt-5.4","previous_response_id":"resp-1","input":[{"type":"message","id":"msg-2"}]}`), + }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("openai-response")}) + if errExecute == nil { + t.Fatal("ExecuteStream() error = nil, want replay-required error") + } + statusErr, ok := errExecute.(interface{ StatusCode() int }) + if !ok || statusErr.StatusCode() != http.StatusUpgradeRequired { + t.Fatalf("ExecuteStream() error = %T %v, want status 426", errExecute, errExecute) + } + if got := gjson.Get(errExecute.Error(), "error.code").String(); got != "upstream_http_replay_required" { + t.Fatalf("ExecuteStream() error code = %q, want upstream_http_replay_required", got) + } + requestScoped, ok := errExecute.(cliproxyexecutor.RequestScopedError) + if !ok || !requestScoped.IsRequestScoped() { + t.Fatalf("ExecuteStream() error = %T, want request-scoped replay signal", errExecute) + } +} + func TestCodexWebsocketsExecuteStreamPassesThroughUpstreamWebsocketPayloadForDownstreamWebsocket(t *testing.T) { upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} capturedPayload := make(chan []byte, 1) diff --git a/internal/runtime/executor/xai_websockets_executor.go b/internal/runtime/executor/xai_websockets_executor.go index 967172e33..1f04538e6 100644 --- a/internal/runtime/executor/xai_websockets_executor.go +++ b/internal/runtime/executor/xai_websockets_executor.go @@ -398,6 +398,9 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox } idMapper := newXAIWebsocketRequestIDMapper(e.idStore, stateSessionID, req.Payload) if xaiInputHasItemType(req.Payload, "compaction_trigger") { + if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) { + return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() + } return e.executeCompactionTriggerFromWebsocketContext(ctx, auth, req, opts, idMapper) } @@ -463,7 +466,20 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox helps.RecordAPIWebsocketRequest(ctx, e.cfg, wsReqLog) logXAIWebsocketRequest(executionSessionID, authID, wsURL, wsReqBody) - conn, respHS, errDial := e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) + var conn *websocket.Conn + var respHS *http.Response + var errDial error + if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) { + conn = existingWebsocketSessionConn(sess, authID, wsURL) + if conn == nil { + if sess != nil { + sess.reqMu.Unlock() + } + return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() + } + } else { + conn, respHS, errDial = e.ensureUpstreamConn(ctx, auth, sess, authID, wsURL, wsHeaders) + } var upstreamHeaders http.Header if respHS != nil { upstreamHeaders = respHS.Header.Clone() @@ -501,6 +517,15 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox errSend = mapXAIWebsocketWriteError(sess, conn, errSend) helps.RecordAPIWebsocketError(ctx, e.cfg, "send", errSend) if sess != nil { + if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) { + e.invalidateUpstreamConnWithoutDisconnectNotify(sess, conn, "send_error", errSend) + sess.clearActive(conn, readCh) + sess.reqMu.Unlock() + if !shouldRetryXAIWebsocketSend(errSend) { + return nil, errSend + } + return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() + } e.invalidateUpstreamConn(sess, conn, "send_error", errSend) if !shouldRetryXAIWebsocketSend(errSend) { sess.clearActive(conn, readCh) @@ -1431,6 +1456,9 @@ func (e *XAIAutoExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth. if cliproxyexecutor.DownstreamWebsocket(ctx) && xaiWebsocketsEnabled(auth) { return e.wsExec.ExecuteStream(ctx, auth, req, opts) } + if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) { + return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() + } return e.httpExec.ExecuteStream(ctx, auth, req, opts) } diff --git a/internal/runtime/executor/xai_websockets_executor_test.go b/internal/runtime/executor/xai_websockets_executor_test.go index 6df8ade69..75ceae3c7 100644 --- a/internal/runtime/executor/xai_websockets_executor_test.go +++ b/internal/runtime/executor/xai_websockets_executor_test.go @@ -34,6 +34,50 @@ func TestXAIWebsocketsEnabledForConfigAPIKey(t *testing.T) { } } +func TestXAIAutoExecutorRequiredUpstreamWebsocketRejectsHTTPFallback(t *testing.T) { + exec := NewXAIAutoExecutor(&config.Config{}) + auth := &cliproxyauth.Auth{ + ID: "xai-http-only", + Provider: "xai", + Attributes: map[string]string{ + "api_key": "xai-key", + }, + } + ctx := cliproxyexecutor.WithRequiredUpstreamWebsocket( + cliproxyexecutor.WithDownstreamWebsocket(context.Background()), + ) + _, errExecute := exec.ExecuteStream(ctx, auth, cliproxyexecutor.Request{ + Model: "grok-4", + Payload: []byte(`{"model":"grok-4","previous_response_id":"resp-1","input":[{"type":"message","id":"msg-2"}]}`), + }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("openai-response")}) + if errExecute == nil { + t.Fatal("ExecuteStream() error = nil, want replay-required error") + } + statusErr, ok := errExecute.(interface{ StatusCode() int }) + if !ok || statusErr.StatusCode() != http.StatusUpgradeRequired { + t.Fatalf("ExecuteStream() error = %T %v, want status 426", errExecute, errExecute) + } + if got := gjson.Get(errExecute.Error(), "error.code").String(); got != "upstream_http_replay_required" { + t.Fatalf("ExecuteStream() error code = %q, want upstream_http_replay_required", got) + } + requestScoped, ok := errExecute.(cliproxyexecutor.RequestScopedError) + if !ok || !requestScoped.IsRequestScoped() { + t.Fatalf("ExecuteStream() error = %T, want request-scoped replay signal", errExecute) + } +} + +func TestXAIWebsocketsRequiredUpstreamRejectsCompactionHTTPFallback(t *testing.T) { + exec := NewXAIWebsocketsExecutor(&config.Config{}) + ctx := cliproxyexecutor.WithRequiredUpstreamWebsocket(context.Background()) + _, errExecute := exec.ExecuteStream(ctx, &cliproxyauth.Auth{}, cliproxyexecutor.Request{ + Model: "grok-4", + Payload: []byte(`{"model":"grok-4","input":[{"type":"compaction_trigger"}]}`), + }, cliproxyexecutor.Options{}) + if !cliproxyexecutor.IsUpstreamWebsocketReplayRequired(errExecute) { + t.Fatalf("ExecuteStream() error = %T %v, want replay-required", errExecute, errExecute) + } +} + func TestMapXAIWebsocketWriteErrorStopsRetryForMessageTooBig(t *testing.T) { networkWriteErr := errors.New("write: broken pipe") tests := []struct { diff --git a/sdk/api/handlers/handlers.go b/sdk/api/handlers/handlers.go index a05b1e416..1eea484f4 100644 --- a/sdk/api/handlers/handlers.go +++ b/sdk/api/handlers/handlers.go @@ -64,6 +64,7 @@ const ( type pinnedAuthContextKey struct{} type selectedAuthCallbackContextKey struct{} +type preparedModelRouteContextKey struct{} type executionSessionContextKey struct{} type disallowFreeAuthContextKey struct{} @@ -142,6 +143,26 @@ func WithSelectedAuthIDCallback(ctx context.Context, callback func(string)) cont return context.WithValue(ctx, selectedAuthCallbackContextKey{}, callback) } +// PrepareStreamModelRoute resolves a stream route once and stores it on the returned context for execution. +// The boolean reports whether the route overrides normal model-to-provider resolution. +func (h *BaseAPIHandler) PrepareStreamModelRoute(ctx context.Context, handlerType string, modelName string, rawJSON []byte) (context.Context, bool) { + if ctx == nil { + ctx = context.Background() + } + decision := h.applyModelRouter(ctx, handlerType, modelName, rawJSON, true, modelExecutionOptions{}) + ctx = context.WithValue(ctx, preparedModelRouteContextKey{}, decision) + hasOverride := strings.TrimSpace(decision.ExecutorPluginID) != "" || strings.TrimSpace(decision.Provider) != "" + return ctx, hasOverride +} + +func preparedModelRouteFromContext(ctx context.Context) (modelRouteDecision, bool) { + if ctx == nil { + return modelRouteDecision{}, false + } + decision, ok := ctx.Value(preparedModelRouteContextKey{}).(modelRouteDecision) + return decision, ok +} + // WithExecutionSessionID returns a child context tagged with a long-lived execution session ID. func WithExecutionSessionID(ctx context.Context, sessionID string) context.Context { sessionID = strings.TrimSpace(sessionID) @@ -1139,7 +1160,10 @@ func (h *BaseAPIHandler) executeStreamWithAuthManager(ctx context.Context, handl func (h *BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context, entryProtocol, exitProtocol, modelName string, rawJSON []byte, alt string, allowImageModel bool, execOptions modelExecutionOptions) (<-chan []byte, http.Header, <-chan *interfaces.ErrorMessage) { originalRequestedModel := modelName - routeDecision := h.applyModelRouter(ctx, entryProtocol, modelName, rawJSON, true, execOptions) + routeDecision, preparedRoute := preparedModelRouteFromContext(ctx) + if !preparedRoute { + routeDecision = h.applyModelRouter(ctx, entryProtocol, modelName, rawJSON, true, execOptions) + } responseProtocol := modelExecutionResponseProtocol(entryProtocol, exitProtocol) if errMsg := validateNativeInteractionsExecution(entryProtocol, execOptions, routeDecision); errMsg != nil { errChan := make(chan *interfaces.ErrorMessage, 1) diff --git a/sdk/api/handlers/handlers_model_router_test.go b/sdk/api/handlers/handlers_model_router_test.go index f631f1d46..4bd93a2c8 100644 --- a/sdk/api/handlers/handlers_model_router_test.go +++ b/sdk/api/handlers/handlers_model_router_test.go @@ -426,6 +426,38 @@ func TestHandlerModelRouterRoutesStreamBeforeRequestDetails(t *testing.T) { } } +func TestPrepareStreamModelRouteReusesDecisionDuringExecution(t *testing.T) { + const model = "prepared-router-model" + const targetPluginID = "prepared-stream-plugin" + routeCalls := 0 + host := &handlerDirectExecutorRouteHost{} + host.hasRouters = true + host.route = func(context.Context, pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) { + routeCalls++ + return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: targetPluginID}, true + } + handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil) + handler.SetModelRouterHost(host) + body := []byte(`{"model":"prepared-router-model","stream":true}`) + ctx, routedToPlugin := handler.PrepareStreamModelRoute(context.Background(), "openai", model, body) + if !routedToPlugin { + t.Fatal("PrepareStreamModelRoute() did not detect plugin executor route") + } + + dataChan, _, errChan := handler.ExecuteStreamWithAuthManager(ctx, "openai", model, body, "") + for range dataChan { + } + if errMsg := <-errChan; errMsg != nil { + t.Fatalf("ExecuteStreamWithAuthManager() error = %+v", errMsg) + } + if routeCalls != 1 { + t.Fatalf("model router calls = %d, want 1", routeCalls) + } + if host.lastPluginID != targetPluginID { + t.Fatalf("plugin id = %q, want %q", host.lastPluginID, targetPluginID) + } +} + func TestExecuteModelPropagatesRouterSkipPluginID(t *testing.T) { model := "model-execution-router-skip-model" requestBody := []byte(fmt.Sprintf(`{"model":%q}`, model)) diff --git a/sdk/api/handlers/openai/openai_responses_websocket.go b/sdk/api/handlers/openai/openai_responses_websocket.go index 11ffcb0c9..a0aa52c9b 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket.go +++ b/sdk/api/handlers/openai/openai_responses_websocket.go @@ -33,15 +33,19 @@ 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" - wsCloseReasonMaxBytes = 123 + 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 + wsHTTPReplayRequiredCloseReason = "upstream requires HTTP replay" + responsesWebsocketUpstreamModeUnknown = "" + responsesWebsocketUpstreamModeWS = "websocket" + responsesWebsocketUpstreamModeHTTP = "http" 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:" ) @@ -74,6 +78,14 @@ func websocketClosePayloadForUpstreamError(err error) (bool, []byte) { return false, nil } + errText := err.Error() + if cliproxyexecutor.IsUpstreamWebsocketReplayRequired(err) { + return true, websocket.FormatCloseMessage( + websocket.CloseServiceRestart, + truncateWebsocketCloseReason(wsHTTPReplayRequiredCloseReason, wsCloseReasonMaxBytes), + ) + } + code := 0 reason := "" var closeErr *websocket.CloseError @@ -85,7 +97,6 @@ func websocketClosePayloadForUpstreamError(err error) (bool, []byte) { 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 @@ -419,20 +430,35 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { // Preserve independent upstream auth affinity when a downstream session switches providers. pinnedAuthByProvider := make(map[string]responsesWebsocketPinnedAuthState) passthroughModelName := "" + upstreamMode := responsesWebsocketUpstreamModeUnknown + upstreamWebsocketAuthID := "" sessionAuthByIDWithSource := func(authID string) (*coreauth.Auth, bool, bool) { if h == nil || h.AuthManager == nil { return nil, false, false } + // Prefer the current manager view so hot-reloaded transport eligibility is + // observed even when the execution session still holds an older auth snapshot. + if auth, ok := h.AuthManager.GetByID(authID); ok { + return auth, false, true + } if auth, ok := h.AuthManager.GetExecutionSessionAuthByID(passthroughSessionID, authID); ok { return auth, true, true } - auth, ok := h.AuthManager.GetByID(authID) - return auth, false, ok + return nil, false, false } sessionAuthByID := func(authID string) (*coreauth.Auth, bool) { auth, _, ok := sessionAuthByIDWithSource(authID) return auth, ok } + upstreamModeForAuth := func(auth *coreauth.Auth) string { + if auth != nil && websocketUpstreamSupportsIncrementalInput(auth.Attributes, auth.Metadata) { + provider := strings.ToLower(strings.TrimSpace(auth.Provider)) + if provider == "codex" || provider == "xai" { + return responsesWebsocketUpstreamModeWS + } + } + return responsesWebsocketUpstreamModeHTTP + } rememberPinnedAuth := func(authID string, modelName string) { authID = strings.TrimSpace(authID) auth, ok := sessionAuthByID(authID) @@ -454,7 +480,6 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { } pinnedAuthID = "" } - forceTranscriptReplayNextRequest := false for { msgType, payload, errReadMessage := conn.ReadMessage() @@ -488,6 +513,13 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { if requestModelName == "" { requestModelName = strings.TrimSpace(gjson.GetBytes(lastRequest, "model").String()) } + executionParent := context.WithValue(c.Request.Context(), "gin", c) + executionParent, routeOverridesModelResolution := h.PrepareStreamModelRoute( + executionParent, + h.HandlerType(), + requestModelName, + payload, + ) if pinnedAuthID != "" { pinnedAuth, homeRuntime, ok := sessionAuthByIDWithSource(pinnedAuthID) providerKey := "" @@ -514,15 +546,40 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { } } useUpstreamWebsocketPassthrough := h.responsesWebsocketUsesUpstreamWebsocketPassthrough(requestModelName) + if pinnedAuthID != "" { + if pinnedAuth, ok := sessionAuthByID(pinnedAuthID); ok && responsesWebsocketAuthSupportsIncrementalInput(pinnedAuth) { + provider := strings.ToLower(strings.TrimSpace(pinnedAuth.Provider)) + useUpstreamWebsocketPassthrough = provider == "codex" || provider == "xai" + } + } + nativeWebsocketPassthrough := !routeOverridesModelResolution && responsesWebsocketNativePassthroughAllowed( + upstreamMode, + useUpstreamWebsocketPassthrough, + pinnedAuthID, + upstreamWebsocketAuthID, + ) + requestRequiresCurrentUpstreamWebsocket := responsesWebsocketRequestRequiresCurrentUpstream(payload) + if upstreamMode == responsesWebsocketUpstreamModeWS && !nativeWebsocketPassthrough { + if requestRequiresCurrentUpstreamWebsocket { + replayErr := responsesWebsocketHTTPReplayRequiredError() + wsTerminateErr = replayErr + matched, errClose := writer.closeForUpstreamError(replayErr) + if !matched { + _ = conn.Close() + } else if errClose != nil && !errors.Is(errClose, websocket.ErrCloseSent) { + log.Debugf("responses websocket: replay close failed id=%s error=%v", passthroughSessionID, errClose) + } + return + } + // A full response.create is already a self-contained reset and can safely + // establish a new upstream transport without another replay. + } if explicitRequestModelName != "" && !useUpstreamWebsocketPassthrough { passthroughModelName = "" } - allowIncrementalInputWithPreviousResponseID := false + allowCompactionReplayBypass := false - if !useUpstreamWebsocketPassthrough { - // Downstream websocket with CPA-mediated upstream (HTTP/SSE) always uses merged - // transcript replay. Incremental previous_response_id is reserved for end-to-end - // upstream websocket passthrough only. + if !nativeWebsocketPassthrough { if pinnedAuthID != "" { if pinnedAuth, ok := sessionAuthByID(pinnedAuthID); ok && pinnedAuth != nil { allowCompactionReplayBypass = responsesWebsocketAuthSupportsCompactionReplay(pinnedAuth) @@ -535,8 +592,10 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { var requestJSON []byte var updatedLastRequest []byte var errMsg *interfaces.ErrorMessage - if useUpstreamWebsocketPassthrough { + if nativeWebsocketPassthrough { requestJSON, errMsg = normalizeResponsesWebsocketPassthroughRequest(payload, requestModelName) + } else if len(lastRequest) == 0 && strings.TrimSpace(gjson.GetBytes(payload, "previous_response_id").String()) != "" { + errMsg = responsesWebsocketPreviousResponseNotFoundError() } else { requestJSON, updatedLastRequest, errMsg = normalizeResponsesWebsocketRequestWithIncrementalState( payload, @@ -544,7 +603,7 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { lastResponseOutput, lastResponseID, lastResponsePendingToolCallIDs, - allowIncrementalInputWithPreviousResponseID, + false, allowCompactionReplayBypass, ) } @@ -570,7 +629,7 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { } continue } - if !useUpstreamWebsocketPassthrough && shouldHandleResponsesWebsocketPrewarmLocally(payload, lastRequest, allowIncrementalInputWithPreviousResponseID) { + if !useUpstreamWebsocketPassthrough && shouldHandleResponsesWebsocketPrewarmLocally(payload, lastRequest, false) { if updated, errDelete := sjson.DeleteBytes(requestJSON, "generate"); errDelete == nil { requestJSON = updated } @@ -592,80 +651,112 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { previousLastResponseOutput := bytes.Clone(lastResponseOutput) previousLastResponseID := lastResponseID previousLastResponsePendingToolCallIDs := append([]string(nil), lastResponsePendingToolCallIDs...) - forcedTranscriptReplay := forceTranscriptReplayNextRequest - if useUpstreamWebsocketPassthrough { + if nativeWebsocketPassthrough { if modelName := strings.TrimSpace(gjson.GetBytes(requestJSON, "model").String()); modelName != "" { passthroughModelName = modelName } - if forcedTranscriptReplay { - forceTranscriptReplayNextRequest = false - } } else { requestJSON = repairResponsesWebsocketToolCalls(downstreamSessionKey, requestJSON) requestJSON = dedupeResponsesWebsocketInputItemsByID(requestJSON) updatedLastRequest = bytes.Clone(requestJSON) lastRequest = updatedLastRequest - if forcedTranscriptReplay { - forceTranscriptReplayNextRequest = false - } } modelName := gjson.GetBytes(requestJSON, "model").String() lastAttemptedAuthID := pinnedAuthID - cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background()) + attemptedUpstreamMode := responsesWebsocketUpstreamModeUnknown + selectedAuthObserved := false + pinnedAuthAttempted := false + cliCtx, cliCancel := h.GetContextWithCancel(h, c, executionParent) cliCtx = cliproxyexecutor.WithDownstreamWebsocket(cliCtx) + if nativeWebsocketPassthrough && requestRequiresCurrentUpstreamWebsocket { + cliCtx = cliproxyexecutor.WithRequiredUpstreamWebsocket(cliCtx) + } cliCtx = handlers.WithExecutionSessionID(cliCtx, passthroughSessionID) - if pinnedAuthID != "" { + cliCtx = handlers.WithSelectedAuthIDCallback(cliCtx, func(authID string) { + authID = strings.TrimSpace(authID) + if authID == "" || h == nil || h.AuthManager == nil { + return + } + lastAttemptedAuthID = authID + selectedAuthObserved = true + pinnedAuthAttempted = pinnedAuthAttempted || (pinnedAuthID != "" && authID == pinnedAuthID) + selectedAuth, ok := sessionAuthByID(authID) + if !ok || selectedAuth == nil { + return + } + attemptedUpstreamMode = upstreamModeForAuth(selectedAuth) + }) + if pinnedAuthID != "" && !routeOverridesModelResolution { cliCtx = handlers.WithPinnedAuthID(cliCtx, pinnedAuthID) - } else { - cliCtx = handlers.WithSelectedAuthIDCallback(cliCtx, func(authID string) { - authID = strings.TrimSpace(authID) - if authID == "" || h == nil || h.AuthManager == nil { - return - } - lastAttemptedAuthID = authID - selectedAuth, ok := sessionAuthByID(authID) - if !ok || selectedAuth == nil { - return - } - if websocketUpstreamSupportsIncrementalInput(selectedAuth.Attributes, selectedAuth.Metadata) { - rememberPinnedAuth(authID, modelName) - } - }) } dataChan, _, errChan := h.ExecuteStreamWithAuthManager(cliCtx, h.HandlerType(), modelName, requestJSON, "") + if !selectedAuthObserved { + // Plugin/alternate routes bypass auth selection. Keep canonical HTTP-mode + // state instead of inheriting the previous pinned websocket mode. + attemptedUpstreamMode = responsesWebsocketUpstreamModeHTTP + } + // A connection-scoped continuation cannot rotate credentials in place. Suppress + // credential errors and make the client replay the full turn on a new socket. + replayPinnedAuthFailure := func(errMsg *interfaces.ErrorMessage) bool { + return nativeWebsocketPassthrough && requestRequiresCurrentUpstreamWebsocket && pinnedAuthAttempted && + shouldReplayResponsesWebsocketPinnedAuthFailure(errMsg) + } - completedOutput, completedResponseID, completedPendingToolCallIDs, forwardErrMsg, errForward := h.forwardResponsesWebsocket(c, writer, cliCancel, dataChan, errChan, wsTimelineLog, passthroughSessionID) + completedOutput, completedResponseID, completedPendingToolCallIDs, forwardErrMsg, errForward := h.forwardResponsesWebsocket( + c, + writer, + cliCancel, + dataChan, + errChan, + wsTimelineLog, + passthroughSessionID, + responsesWebsocketForwardOptions{ + suppressError: replayPinnedAuthFailure, + }, + ) if errForward != nil { wsTerminateErr = errForward - log.Warnf("responses websocket: forward failed id=%s error=%v", passthroughSessionID, errForward) + if !errors.Is(errForward, websocket.ErrCloseSent) { + log.Warnf("responses websocket: forward failed id=%s error=%v", passthroughSessionID, errForward) + } return } - if forwardErrMsg == nil && !useUpstreamWebsocketPassthrough && lastAttemptedAuthID != "" { - if selectedAuth, ok := sessionAuthByID(lastAttemptedAuthID); ok && selectedAuth != nil { - if websocketUpstreamSupportsIncrementalInput(selectedAuth.Attributes, selectedAuth.Metadata) { - rememberPinnedAuth(lastAttemptedAuthID, modelName) - } else if pinnedAuthID != "" { - if pinnedAuth, ok := sessionAuthByID(pinnedAuthID); ok && pinnedAuth != nil && websocketUpstreamSupportsIncrementalInput(pinnedAuth.Attributes, pinnedAuth.Metadata) { - rememberPinnedAuth(lastAttemptedAuthID, modelName) - } - } + if forwardErrMsg != nil { + lastRequest = previousLastRequest + lastResponseOutput = previousLastResponseOutput + lastResponseID = previousLastResponseID + lastResponsePendingToolCallIDs = previousLastResponsePendingToolCallIDs + if pinnedAuthAttempted && shouldReleaseResponsesWebsocketPinnedAuth(forwardErrMsg) { + forgetPinnedAuth() } - } - if shouldReleaseResponsesWebsocketPinnedAuth(forwardErrMsg) { - forgetPinnedAuth() - forceTranscriptReplayNextRequest = true - if useUpstreamWebsocketPassthrough { - passthroughModelName = "" - } else { - lastRequest = previousLastRequest - lastResponseOutput = previousLastResponseOutput - lastResponseID = previousLastResponseID - lastResponsePendingToolCallIDs = previousLastResponsePendingToolCallIDs + if replayPinnedAuthFailure(forwardErrMsg) { + replayErr := responsesWebsocketHTTPReplayRequiredError() + wsTerminateErr = replayErr + matched, errClose := writer.closeForUpstreamError(replayErr) + if !matched { + _ = conn.Close() + } else if errClose != nil && !errors.Is(errClose, websocket.ErrCloseSent) { + log.Debugf("responses websocket: credential replay close failed id=%s error=%v", passthroughSessionID, errClose) + } + return } continue } - if !useUpstreamWebsocketPassthrough { + + upstreamMode = attemptedUpstreamMode + if upstreamMode == responsesWebsocketUpstreamModeWS { + upstreamWebsocketAuthID = lastAttemptedAuthID + if lastAttemptedAuthID != "" { + rememberPinnedAuth(lastAttemptedAuthID, modelName) + } + passthroughModelName = modelName + lastRequest = nil + lastResponseOutput = []byte("[]") + lastResponseID = "" + lastResponsePendingToolCallIDs = nil + } else { + upstreamWebsocketAuthID = "" lastResponseOutput = completedOutput lastResponseID = strings.TrimSpace(completedResponseID) lastResponsePendingToolCallIDs = append([]string(nil), completedPendingToolCallIDs...) @@ -673,6 +764,20 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { } } +func responsesWebsocketHTTPReplayRequiredError() error { + return cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() +} + +func responsesWebsocketRequestRequiresCurrentUpstream(payload []byte) bool { + return strings.TrimSpace(gjson.GetBytes(payload, "previous_response_id").String()) != "" || + strings.TrimSpace(gjson.GetBytes(payload, "type").String()) == wsRequestTypeAppend +} + +func responsesWebsocketNativePassthroughAllowed(upstreamMode string, useUpstreamWebsocket bool, pinnedAuthID string, upstreamAuthID string) bool { + return upstreamMode == responsesWebsocketUpstreamModeWS && useUpstreamWebsocket && + strings.TrimSpace(pinnedAuthID) != "" && strings.TrimSpace(pinnedAuthID) == strings.TrimSpace(upstreamAuthID) +} + func websocketClientAddress(c *gin.Context) string { if c == nil || c.Request == nil { return "" @@ -694,6 +799,15 @@ func websocketUpgradeHeaders(req *http.Request) http.Header { return headers } +func responsesWebsocketPreviousResponseNotFoundError() *interfaces.ErrorMessage { + return &interfaces.ErrorMessage{ + StatusCode: http.StatusConflict, + Error: errors.New( + `{"error":{"message":"Previous response is not available on this websocket; resend the full conversation input without previous_response_id","type":"invalid_request_error","code":"previous_response_not_found","param":"previous_response_id"}}`, + ), + } +} + func normalizeResponsesWebsocketRequest(rawJSON []byte, lastRequest []byte, lastResponseOutput []byte) ([]byte, []byte, *interfaces.ErrorMessage) { return normalizeResponsesWebsocketRequestWithMode(rawJSON, lastRequest, lastResponseOutput, true, true) } @@ -1570,6 +1684,10 @@ func normalizeJSONArrayRaw(raw []byte) string { return "[]" } +type responsesWebsocketForwardOptions struct { + suppressError func(*interfaces.ErrorMessage) bool +} + func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( c *gin.Context, writer *responsesWebsocketWriter, @@ -1578,7 +1696,12 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( errs <-chan *interfaces.ErrorMessage, wsTimelineLog websocketTimelineAppender, sessionID string, + options ...responsesWebsocketForwardOptions, ) ([]byte, string, []string, *interfaces.ErrorMessage, error) { + var opts responsesWebsocketForwardOptions + if len(options) > 0 { + opts = options[0] + } completed := false completedOutput := []byte("[]") completedResponseID := "" @@ -1602,6 +1725,10 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( } if errMsg != nil { h.LoggingAPIResponseError(context.WithValue(context.Background(), "gin", c), errMsg) + if opts.suppressError != nil && opts.suppressError(errMsg) { + cancel(errMsg.Error) + return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), errMsg, nil + } markAPIResponseTimestamp(c) if matched, errClose := writer.closeForUpstreamError(errMsg.Error); matched { cancel(errMsg.Error) @@ -1684,6 +1811,10 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( if h != nil { h.LoggingAPIResponseError(context.WithValue(context.Background(), "gin", c), payloadErrMsg) } + if opts.suppressError != nil && opts.suppressError(payloadErrMsg) { + cancel(payloadErrMsg.Error) + return completedOutput, completedResponseID, sortedStringSet(pendingToolCallIDs), payloadErrMsg, nil + } } else if isResponsesWebsocketCompletionEvent(eventType) { completed = true completedOutput = responseCompletedOutputFromPayload(payloads[i], outputItemsByIndex, outputItemsFallback) @@ -1716,9 +1847,9 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( } } -func shouldReleaseResponsesWebsocketPinnedAuth(errMsg *interfaces.ErrorMessage) bool { +func responsesWebsocketErrorStatus(errMsg *interfaces.ErrorMessage) int { if errMsg == nil { - return false + return 0 } status := errMsg.StatusCode if status <= 0 && errMsg.Error != nil { @@ -1726,7 +1857,23 @@ func shouldReleaseResponsesWebsocketPinnedAuth(errMsg *interfaces.ErrorMessage) status = se.StatusCode() } } - switch status { + return status +} + +func shouldReplayResponsesWebsocketPinnedAuthFailure(errMsg *interfaces.ErrorMessage) bool { + switch responsesWebsocketErrorStatus(errMsg) { + case http.StatusUnauthorized, http.StatusTooManyRequests: + return true + default: + return false + } +} + +func shouldReleaseResponsesWebsocketPinnedAuth(errMsg *interfaces.ErrorMessage) bool { + if errMsg == nil { + return false + } + switch responsesWebsocketErrorStatus(errMsg) { case http.StatusUnauthorized, http.StatusPaymentRequired, http.StatusForbidden, diff --git a/sdk/api/handlers/openai/openai_responses_websocket_test.go b/sdk/api/handlers/openai/openai_responses_websocket_test.go index 4f093f477..f50f44695 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket_test.go +++ b/sdk/api/handlers/openai/openai_responses_websocket_test.go @@ -7,6 +7,7 @@ import ( "fmt" "net/http" "net/http/httptest" + "strconv" "strings" "sync" "testing" @@ -22,9 +23,55 @@ import ( coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" "github.com/tidwall/gjson" ) +func TestWebsocketReplayCloseRequiresTypedSignal(t *testing.T) { + matched, payload := websocketClosePayloadForUpstreamError(responsesWebsocketHTTPReplayRequiredError()) + if !matched || len(payload) == 0 { + t.Fatalf("typed replay signal matched=%t payload_len=%d, want close payload", matched, len(payload)) + } + spoofed := websocketPinnedFailoverStatusError{ + status: http.StatusUpgradeRequired, + msg: `{"error":{"code":"upstream_http_replay_required"}}`, + } + if matched, _ := websocketClosePayloadForUpstreamError(spoofed); matched { + t.Fatal("untyped upstream error spoofed replay close") + } +} + +func TestResponsesWebsocketRequestRequiresCurrentUpstream(t *testing.T) { + cases := []struct { + name string + payload string + want bool + }{ + {name: "incremental create", payload: `{"type":"response.create","previous_response_id":"resp-1","input":[]}`, want: true}, + {name: "append", payload: `{"type":"response.append","input":[]}`, want: true}, + {name: "full create", payload: `{"type":"response.create","input":[]}`, want: false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := responsesWebsocketRequestRequiresCurrentUpstream([]byte(tc.payload)); got != tc.want { + t.Fatalf("responsesWebsocketRequestRequiresCurrentUpstream() = %t, want %t", got, tc.want) + } + }) + } +} + +func TestResponsesWebsocketNativePassthroughRequiresImmediatelyPreviousAuth(t *testing.T) { + if !responsesWebsocketNativePassthroughAllowed(responsesWebsocketUpstreamModeWS, true, "auth-a", "auth-a") { + t.Fatal("matching immediate websocket auth did not allow native passthrough") + } + if responsesWebsocketNativePassthroughAllowed(responsesWebsocketUpstreamModeWS, true, "auth-a", "auth-b") { + t.Fatal("restored auth from an older provider session allowed native passthrough") + } + if responsesWebsocketNativePassthroughAllowed(responsesWebsocketUpstreamModeHTTP, true, "auth-a", "auth-a") { + t.Fatal("HTTP mode allowed native websocket passthrough") + } +} + func TestWriteWebsocketCloseForUpstreamErrorMirrorsMessageTooBig(t *testing.T) { tests := []struct { name string @@ -304,6 +351,22 @@ type websocketProviderCaptureExecutor struct { websocketCaptureExecutor } +type websocketProviderRouteHost struct{} + +func (*websocketProviderRouteHost) HasModelRouters() bool { return true } + +func (*websocketProviderRouteHost) RouteModel(_ context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) { + if !gjson.GetBytes(req.Body, "route_to_claude").Bool() { + return pluginapi.ModelRouteResponse{}, false + } + return pluginapi.ModelRouteResponse{ + Handled: true, + TargetKind: pluginapi.ModelRouteTargetProvider, + Target: "claude", + TargetModel: "claude-provider-route-target", + }, true +} + type websocketCompactionCaptureExecutor struct { mu sync.Mutex streamPayloads [][]byte @@ -346,10 +409,11 @@ type websocketAuthCaptureExecutor struct { } type websocketPinnedFailoverExecutor struct { - mu sync.Mutex - authIDs []string - calls map[string]int - payloads map[string][][]byte + mu sync.Mutex + failStatus int + authIDs []string + calls map[string]int + payloads map[string][][]byte } type websocketBootstrapFallbackExecutor struct { @@ -359,12 +423,21 @@ type websocketBootstrapFallbackExecutor struct { } type websocketDirectCaptureExecutor struct { + mu sync.Mutex + provider string + failStatus int + authIDs []string + models []string + payloads [][]byte + requiredUpstreamWebsocket []bool + done chan struct{} + doneOnce sync.Once +} + +type websocketCanonicalRollbackExecutor struct { mu sync.Mutex - provider string - authIDs []string payloads [][]byte - done chan struct{} - doneOnce sync.Once + calls int } type websocketPinnedFailoverStatusError struct { @@ -399,7 +472,7 @@ func (e *websocketBootstrapFallbackExecutor) ExecuteStream(_ context.Context, au chunks := make(chan coreexecutor.StreamChunk, 1) if authID == "auth-ws" { chunks <- coreexecutor.StreamChunk{Err: websocketPinnedFailoverStatusError{ - status: http.StatusServiceUnavailable, + status: http.StatusUpgradeRequired, msg: `{"error":{"message":"websocket bootstrap failed","type":"server_error","code":"ws_failed"}}`, }} close(chunks) @@ -451,18 +524,29 @@ func (e *websocketDirectCaptureExecutor) Execute(context.Context, *coreauth.Auth return coreexecutor.Response{}, errors.New("not implemented") } -func (e *websocketDirectCaptureExecutor) ExecuteStream(_ context.Context, auth *coreauth.Auth, req coreexecutor.Request, _ coreexecutor.Options) (*coreexecutor.StreamResult, error) { +func (e *websocketDirectCaptureExecutor) ExecuteStream(ctx context.Context, auth *coreauth.Auth, req coreexecutor.Request, _ coreexecutor.Options) (*coreexecutor.StreamResult, error) { authID := "" if auth != nil { authID = auth.ID } e.mu.Lock() e.authIDs = append(e.authIDs, authID) + e.models = append(e.models, req.Model) e.payloads = append(e.payloads, bytes.Clone(req.Payload)) + e.requiredUpstreamWebsocket = append(e.requiredUpstreamWebsocket, coreexecutor.RequiredUpstreamWebsocket(ctx)) count := len(e.payloads) + failStatus := e.failStatus e.mu.Unlock() chunks := make(chan coreexecutor.StreamChunk, 1) + if failStatus > 0 { + chunks <- coreexecutor.StreamChunk{Err: websocketPinnedFailoverStatusError{ + status: failStatus, + msg: `{"error":{"message":"routed provider failed","type":"authentication_error","code":"invalid_api_key"}}`, + }} + close(chunks) + return &coreexecutor.StreamResult{Chunks: chunks}, nil + } responseID := fmt.Sprintf("resp-%d", count) chunks <- coreexecutor.StreamChunk{Payload: []byte(fmt.Sprintf(`{"type":"response.completed","response":{"id":%q,"output":[{"type":"message","id":"out-%d"}]}}`, responseID, count))} close(chunks) @@ -502,6 +586,67 @@ func (e *websocketDirectCaptureExecutor) AuthIDs() []string { return append([]string(nil), e.authIDs...) } +func (e *websocketDirectCaptureExecutor) Models() []string { + e.mu.Lock() + defer e.mu.Unlock() + return append([]string(nil), e.models...) +} + +func (e *websocketDirectCaptureExecutor) RequiredUpstreamWebsocketFlags() []bool { + e.mu.Lock() + defer e.mu.Unlock() + return append([]bool(nil), e.requiredUpstreamWebsocket...) +} + +func (e *websocketCanonicalRollbackExecutor) Identifier() string { return "xai" } + +func (e *websocketCanonicalRollbackExecutor) Execute(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { + return coreexecutor.Response{}, errors.New("not implemented") +} + +func (e *websocketCanonicalRollbackExecutor) ExecuteStream(_ context.Context, _ *coreauth.Auth, req coreexecutor.Request, _ coreexecutor.Options) (*coreexecutor.StreamResult, error) { + e.mu.Lock() + e.calls++ + call := e.calls + e.payloads = append(e.payloads, bytes.Clone(req.Payload)) + e.mu.Unlock() + + chunks := make(chan coreexecutor.StreamChunk, 1) + if call == 2 { + chunks <- coreexecutor.StreamChunk{Err: websocketPinnedFailoverStatusError{ + status: http.StatusBadRequest, + msg: `{"error":{"message":"bad turn","type":"invalid_request_error","code":"invalid_request"}}`, + }} + close(chunks) + return &coreexecutor.StreamResult{Chunks: chunks}, nil + } + chunks <- coreexecutor.StreamChunk{Payload: []byte(fmt.Sprintf(`{"type":"response.completed","response":{"id":"resp-%d","output":[{"type":"message","id":"out-%d"}]}}`, call, call))} + close(chunks) + return &coreexecutor.StreamResult{Chunks: chunks}, nil +} + +func (e *websocketCanonicalRollbackExecutor) Refresh(_ context.Context, auth *coreauth.Auth) (*coreauth.Auth, error) { + return auth, nil +} + +func (e *websocketCanonicalRollbackExecutor) CountTokens(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { + return coreexecutor.Response{}, errors.New("not implemented") +} + +func (e *websocketCanonicalRollbackExecutor) HttpRequest(context.Context, *coreauth.Auth, *http.Request) (*http.Response, error) { + return nil, errors.New("not implemented") +} + +func (e *websocketCanonicalRollbackExecutor) Payloads() [][]byte { + e.mu.Lock() + defer e.mu.Unlock() + out := make([][]byte, len(e.payloads)) + for i := range e.payloads { + out[i] = bytes.Clone(e.payloads[i]) + } + return out +} + type websocketUpstreamDisconnectExecutor struct { mu sync.Mutex provider string @@ -618,7 +763,7 @@ func (e *websocketAuthCaptureExecutor) AuthIDs() []string { return append([]string(nil), e.authIDs...) } -func (e *websocketPinnedFailoverExecutor) Identifier() string { return "test-provider" } +func (e *websocketPinnedFailoverExecutor) Identifier() string { return "xai" } func (e *websocketPinnedFailoverExecutor) Execute(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { return coreexecutor.Response{}, errors.New("not implemented") @@ -646,8 +791,8 @@ func (e *websocketPinnedFailoverExecutor) ExecuteStream(_ context.Context, auth if authID == "auth-a" && call == 2 { chunks := make(chan coreexecutor.StreamChunk, 1) chunks <- coreexecutor.StreamChunk{Err: websocketPinnedFailoverStatusError{ - status: http.StatusTooManyRequests, - msg: `{"error":{"message":"quota exhausted","type":"rate_limit_error","code":"rate_limit_exceeded"}}`, + status: e.failStatus, + msg: fmt.Sprintf(`{"error":{"message":"credential failed","status":%d}}`, e.failStatus), }} close(chunks) return &coreexecutor.StreamResult{Chunks: chunks}, nil @@ -2284,7 +2429,7 @@ func TestResponsesWebsocketCodexWebsocketPassthroughPassesCompactedRequestWithou } } -func TestResponsesWebsocketXAIWebsocketPassthroughCarriesPreviousResponseID(t *testing.T) { +func TestResponsesWebsocketXAIWebsocketPassthroughKeepsNativeIncrementalRequest(t *testing.T) { gin.SetMode(gin.TestMode) modelName := "xai-websocket-passthrough-model" @@ -2348,28 +2493,444 @@ func TestResponsesWebsocketXAIWebsocketPassthroughCarriesPreviousResponseID(t *t } secondPayload := payloads[1] if got := gjson.GetBytes(secondPayload, "type").String(); got != wsRequestTypeCreate { - t.Fatalf("second xai passthrough type = %s, want %s: %s", got, wsRequestTypeCreate, secondPayload) + t.Fatalf("incremental xai payload type = %q, want %q: %s", got, wsRequestTypeCreate, secondPayload) } if got := gjson.GetBytes(secondPayload, "model").String(); got != modelName { t.Fatalf("second xai payload model = %s, want %s", got, modelName) } if got := gjson.GetBytes(secondPayload, "previous_response_id").String(); got != "resp-1" { - t.Fatalf("second xai previous_response_id = %s, want resp-1: %s", got, secondPayload) + t.Fatalf("second xai previous_response_id = %q, want resp-1: %s", got, secondPayload) } input := gjson.GetBytes(secondPayload, "input").Array() - if len(input) != 1 { - t.Fatalf("second xai passthrough input len = %d, want 1: %s", len(input), secondPayload) - } - if input[0].Get("id").String() != "msg-2" { - t.Fatalf("second xai passthrough input must contain only the new turn: %s", secondPayload) - } - if bytes.Contains(secondPayload, []byte(`"id":"msg-1"`)) || bytes.Contains(secondPayload, []byte(`"id":"out-1"`)) { - t.Fatalf("second xai passthrough payload contains stale transcript state: %s", secondPayload) + if len(input) != 1 || input[0].Get("id").String() != "msg-2" { + t.Fatalf("second xai incremental input is not the client delta: %s", secondPayload) } authIDs := executor.AuthIDs() if len(authIDs) != 2 || authIDs[0] != "auth-xai-ws" || authIDs[1] != "auth-xai-ws" { t.Fatalf("xai websocket auth IDs = %v, want [auth-xai-ws auth-xai-ws]", authIDs) } + if got := executor.RequiredUpstreamWebsocketFlags(); len(got) != 2 || got[0] || !got[1] { + t.Fatalf("required upstream websocket flags = %v, want [false true]", got) + } +} + +func TestResponsesWebsocketFullRequestCanRouteFromNativeWebsocketToBuiltInProvider(t *testing.T) { + gin.SetMode(gin.TestMode) + + const sourceModel = "codex-provider-route-source" + const targetModel = "claude-provider-route-target" + codexExecutor := &websocketDirectCaptureExecutor{provider: "codex"} + claudeExecutor := &websocketDirectCaptureExecutor{provider: "claude"} + manager := coreauth.NewManager(nil, nil, nil) + manager.RegisterExecutor(codexExecutor) + manager.RegisterExecutor(claudeExecutor) + codexAuth := &coreauth.Auth{ + ID: "auth-codex-provider-route", + Provider: "codex", + Status: coreauth.StatusActive, + Attributes: map[string]string{"websockets": "true"}, + } + claudeAuth := &coreauth.Auth{ + ID: "auth-claude-provider-route", + Provider: "claude", + Status: coreauth.StatusActive, + } + for _, auth := range []*coreauth.Auth{codexAuth, claudeAuth} { + if _, err := manager.Register(context.Background(), auth); err != nil { + t.Fatalf("Register auth %s: %v", auth.ID, err) + } + } + registry.GetGlobalRegistry().RegisterClient(codexAuth.ID, codexAuth.Provider, []*registry.ModelInfo{{ID: sourceModel}}) + registry.GetGlobalRegistry().RegisterClient(claudeAuth.ID, claudeAuth.Provider, []*registry.ModelInfo{{ID: targetModel}}) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(codexAuth.ID) + registry.GetGlobalRegistry().UnregisterClient(claudeAuth.ID) + }) + + base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager) + base.SetModelRouterHost(&websocketProviderRouteHost{}) + h := NewOpenAIResponsesAPIHandler(base) + 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() }() + + firstRequest := []byte(fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-1"}]}`, sourceModel)) + if errWrite := conn.WriteMessage(websocket.TextMessage, firstRequest); errWrite != nil { + t.Fatalf("write first websocket message: %v", errWrite) + } + if _, _, errRead := conn.ReadMessage(); errRead != nil { + t.Fatalf("read first websocket response: %v", errRead) + } + + routedRequest := []byte(fmt.Sprintf(`{"type":"response.create","model":%q,"route_to_claude":true,"input":[{"type":"message","id":"msg-routed"}]}`, sourceModel)) + if errWrite := conn.WriteMessage(websocket.TextMessage, routedRequest); errWrite != nil { + t.Fatalf("write routed websocket message: %v", errWrite) + } + _, response, errRead := conn.ReadMessage() + if errRead != nil { + t.Fatalf("read routed websocket response: %v", errRead) + } + if got := gjson.GetBytes(response, "type").String(); got != wsEventTypeCompleted { + t.Fatalf("routed response type = %q, want %q: %s", got, wsEventTypeCompleted, response) + } + if got := len(codexExecutor.Payloads()); got != 1 { + t.Fatalf("codex payload count = %d, want 1", got) + } + claudePayloads := claudeExecutor.Payloads() + if len(claudePayloads) != 1 { + t.Fatalf("claude payload count = %d, want 1", len(claudePayloads)) + } + if got := claudeExecutor.Models(); len(got) != 1 || got[0] != targetModel { + t.Fatalf("routed models = %v, want [%s]", got, targetModel) + } +} + +func TestResponsesWebsocketFailedProviderRoutePreservesNativeWebsocketPin(t *testing.T) { + gin.SetMode(gin.TestMode) + + const sourceModel = "codex-provider-route-failure-source" + const targetModel = "claude-provider-route-target" + codexExecutor := &websocketDirectCaptureExecutor{provider: "codex"} + claudeExecutor := &websocketDirectCaptureExecutor{provider: "claude", failStatus: http.StatusUnauthorized} + manager := coreauth.NewManager(nil, nil, nil) + manager.RegisterExecutor(codexExecutor) + manager.RegisterExecutor(claudeExecutor) + codexAuth := &coreauth.Auth{ID: "auth-codex-provider-route-failure", Provider: "codex", Status: coreauth.StatusActive, Attributes: map[string]string{"websockets": "true"}} + claudeAuth := &coreauth.Auth{ID: "auth-claude-provider-route-failure", Provider: "claude", Status: coreauth.StatusActive} + for _, auth := range []*coreauth.Auth{codexAuth, claudeAuth} { + if _, err := manager.Register(context.Background(), auth); err != nil { + t.Fatalf("Register auth %s: %v", auth.ID, err) + } + } + registry.GetGlobalRegistry().RegisterClient(codexAuth.ID, codexAuth.Provider, []*registry.ModelInfo{{ID: sourceModel}}) + registry.GetGlobalRegistry().RegisterClient(claudeAuth.ID, claudeAuth.Provider, []*registry.ModelInfo{{ID: targetModel}}) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(codexAuth.ID) + registry.GetGlobalRegistry().UnregisterClient(claudeAuth.ID) + }) + + base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager) + base.SetModelRouterHost(&websocketProviderRouteHost{}) + h := NewOpenAIResponsesAPIHandler(base) + 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() }() + + requests := []string{ + fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-1"}]}`, sourceModel), + fmt.Sprintf(`{"type":"response.create","model":%q,"route_to_claude":true,"input":[{"type":"message","id":"msg-routed"}]}`, sourceModel), + `{"type":"response.create","previous_response_id":"resp-1","input":[{"type":"message","id":"msg-2"}]}`, + } + wantTypes := []string{wsEventTypeCompleted, wsEventTypeError, wsEventTypeCompleted} + for i, request := range requests { + if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(request)); errWrite != nil { + t.Fatalf("write request %d: %v", i+1, errWrite) + } + _, response, errRead := conn.ReadMessage() + if errRead != nil { + t.Fatalf("read response %d: %v", i+1, errRead) + } + if got := gjson.GetBytes(response, "type").String(); got != wantTypes[i] { + t.Fatalf("response %d type = %q, want %q: %s", i+1, got, wantTypes[i], response) + } + } + + codexPayloads := codexExecutor.Payloads() + if len(codexPayloads) != 2 { + t.Fatalf("codex payload count = %d, want 2", len(codexPayloads)) + } + if got := gjson.GetBytes(codexPayloads[1], "previous_response_id").String(); got != "resp-1" { + t.Fatalf("resumed codex previous_response_id = %q, want resp-1: %s", got, codexPayloads[1]) + } + if got := len(claudeExecutor.Payloads()); got != 1 { + t.Fatalf("claude payload count = %d, want 1", got) + } +} + +func TestResponsesWebsocketDeltaRouteToBuiltInProviderRequiresFullReplay(t *testing.T) { + gin.SetMode(gin.TestMode) + + const sourceModel = "codex-provider-route-delta-source" + const targetModel = "claude-provider-route-target" + codexExecutor := &websocketDirectCaptureExecutor{provider: "codex"} + claudeExecutor := &websocketDirectCaptureExecutor{provider: "claude"} + manager := coreauth.NewManager(nil, nil, nil) + manager.RegisterExecutor(codexExecutor) + manager.RegisterExecutor(claudeExecutor) + codexAuth := &coreauth.Auth{ID: "auth-codex-provider-route-delta", Provider: "codex", Status: coreauth.StatusActive, Attributes: map[string]string{"websockets": "true"}} + claudeAuth := &coreauth.Auth{ID: "auth-claude-provider-route-delta", Provider: "claude", Status: coreauth.StatusActive} + for _, auth := range []*coreauth.Auth{codexAuth, claudeAuth} { + if _, err := manager.Register(context.Background(), auth); err != nil { + t.Fatalf("Register auth %s: %v", auth.ID, err) + } + } + registry.GetGlobalRegistry().RegisterClient(codexAuth.ID, codexAuth.Provider, []*registry.ModelInfo{{ID: sourceModel}}) + registry.GetGlobalRegistry().RegisterClient(claudeAuth.ID, claudeAuth.Provider, []*registry.ModelInfo{{ID: targetModel}}) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(codexAuth.ID) + registry.GetGlobalRegistry().UnregisterClient(claudeAuth.ID) + }) + + base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager) + base.SetModelRouterHost(&websocketProviderRouteHost{}) + h := NewOpenAIResponsesAPIHandler(base) + 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() }() + + firstRequest := []byte(fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-1"}]}`, sourceModel)) + if errWrite := conn.WriteMessage(websocket.TextMessage, firstRequest); errWrite != nil { + t.Fatalf("write first websocket message: %v", errWrite) + } + if _, _, errRead := conn.ReadMessage(); errRead != nil { + t.Fatalf("read first websocket response: %v", errRead) + } + + routedDelta := []byte(`{"type":"response.create","route_to_claude":true,"previous_response_id":"resp-1","input":[{"type":"message","id":"msg-routed"}]}`) + if errWrite := conn.WriteMessage(websocket.TextMessage, routedDelta); errWrite != nil { + t.Fatalf("write routed delta: %v", errWrite) + } + _, _, errRead := conn.ReadMessage() + var closeErr *websocket.CloseError + if !errors.As(errRead, &closeErr) { + t.Fatalf("routed delta error = %v, want websocket close", errRead) + } + if closeErr.Code != websocket.CloseServiceRestart || closeErr.Text != wsHTTPReplayRequiredCloseReason { + t.Fatalf("routed delta close = %d %q, want %d %q", closeErr.Code, closeErr.Text, websocket.CloseServiceRestart, wsHTTPReplayRequiredCloseReason) + } + if got := len(codexExecutor.Payloads()); got != 1 { + t.Fatalf("codex payload count = %d, want 1", got) + } + if got := len(claudeExecutor.Payloads()); got != 0 { + t.Fatalf("claude payload count = %d, want 0 before full replay", got) + } +} + +func TestResponsesWebsocketClosesForHTTPReplayWhenWebsocketEligibilityChanges(t *testing.T) { + gin.SetMode(gin.TestMode) + + modelName := "xai-websocket-mode-change-model" + executor := &websocketDirectCaptureExecutor{provider: "xai", done: make(chan struct{})} + manager := coreauth.NewManager(nil, nil, nil) + manager.RegisterExecutor(executor) + auth := &coreauth.Auth{ + ID: "auth-xai-mode-change", + Provider: "xai", + Status: coreauth.StatusActive, + Attributes: map[string]string{"websockets": "true"}, + } + if _, err := manager.Register(context.Background(), auth); err != nil { + t.Fatalf("Register auth: %v", err) + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: modelName}}) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(auth.ID) + }) + + 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() + + 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() }() + + firstRequest := []byte(fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-1"}]}`, modelName)) + if errWrite := conn.WriteMessage(websocket.TextMessage, firstRequest); errWrite != nil { + t.Fatalf("write first websocket message: %v", errWrite) + } + if _, _, errRead := conn.ReadMessage(); errRead != nil { + t.Fatalf("read first websocket response: %v", errRead) + } + + secondRequest := []byte(`{"type":"response.create","previous_response_id":"resp-1","input":[{"type":"message","id":"msg-2"}]}`) + if errWrite := conn.WriteMessage(websocket.TextMessage, secondRequest); errWrite != nil { + t.Fatalf("write second websocket message: %v", errWrite) + } + if _, _, errRead := conn.ReadMessage(); errRead != nil { + t.Fatalf("read second websocket response: %v", errRead) + } + + updatedAuth := &coreauth.Auth{ + ID: auth.ID, + Provider: auth.Provider, + Status: coreauth.StatusActive, + } + if _, errUpdate := manager.Update(context.Background(), updatedAuth); errUpdate != nil { + t.Fatalf("Update auth: %v", errUpdate) + } + + thirdRequest := []byte(`{"type":"response.create","previous_response_id":"resp-2","input":[{"type":"message","id":"msg-3"}]}`) + if errWrite := conn.WriteMessage(websocket.TextMessage, thirdRequest); errWrite != nil { + t.Fatalf("write third websocket message: %v", errWrite) + } + _, _, errRead := conn.ReadMessage() + var closeErr *websocket.CloseError + if !errors.As(errRead, &closeErr) { + t.Fatalf("third response error = %v, want websocket close", errRead) + } + if closeErr.Code != websocket.CloseServiceRestart || closeErr.Text != wsHTTPReplayRequiredCloseReason { + t.Fatalf("third response close = %d %q, want %d %q", closeErr.Code, closeErr.Text, websocket.CloseServiceRestart, wsHTTPReplayRequiredCloseReason) + } + + payloads := executor.Payloads() + if len(payloads) != 2 { + t.Fatalf("executor payload count = %d, want 2; transport switch must not call HTTP upstream", len(payloads)) + } + second := payloads[1] + if got := gjson.GetBytes(second, "previous_response_id").String(); got != "resp-1" { + t.Fatalf("stable websocket previous_response_id = %q, want resp-1: %s", got, second) + } + if input := gjson.GetBytes(second, "input").Array(); len(input) != 1 || input[0].Get("id").String() != "msg-2" { + t.Fatalf("stable websocket payload is not incremental: %s", second) + } + + replayConn, _, errDialReplay := websocket.DefaultDialer.Dial(wsURL, nil) + if errDialReplay != nil { + t.Fatalf("dial replay websocket: %v", errDialReplay) + } + defer func() { _ = replayConn.Close() }() + fullReplay := []byte(fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-1"},{"type":"message","id":"out-1"},{"type":"message","id":"msg-2"},{"type":"message","id":"out-2"},{"type":"message","id":"msg-3"}]}`, modelName)) + if errWrite := replayConn.WriteMessage(websocket.TextMessage, fullReplay); errWrite != nil { + t.Fatalf("write full replay: %v", errWrite) + } + if _, _, errReadReplay := replayConn.ReadMessage(); errReadReplay != nil { + t.Fatalf("read full replay response: %v", errReadReplay) + } + deltaAfterReplay := []byte(`{"type":"response.create","previous_response_id":"resp-3","input":[{"type":"message","id":"msg-4"}]}`) + if errWrite := replayConn.WriteMessage(websocket.TextMessage, deltaAfterReplay); errWrite != nil { + t.Fatalf("write delta after replay: %v", errWrite) + } + if _, _, errReadReplay := replayConn.ReadMessage(); errReadReplay != nil { + t.Fatalf("read delta after replay response: %v", errReadReplay) + } + + payloads = executor.Payloads() + if len(payloads) != 4 { + t.Fatalf("executor payload count after replay = %d, want 4", len(payloads)) + } + httpDelta := payloads[3] + if gjson.GetBytes(httpDelta, "previous_response_id").Exists() { + t.Fatalf("HTTP-mode delta retained previous_response_id: %s", httpDelta) + } + input := gjson.GetBytes(httpDelta, "input").Array() + wantIDs := []string{"msg-1", "out-1", "msg-2", "out-2", "msg-3", "out-3", "msg-4"} + if len(input) != len(wantIDs) { + t.Fatalf("HTTP-mode canonical input len = %d, want %d: %s", len(input), len(wantIDs), httpDelta) + } + for i, wantID := range wantIDs { + if got := input[i].Get("id").String(); got != wantID { + t.Fatalf("HTTP-mode canonical input[%d].id = %q, want %q: %s", i, got, wantID, httpDelta) + } + } +} + +func TestResponsesWebsocketRejectsUnknownPreviousResponseOnNewSocket(t *testing.T) { + gin.SetMode(gin.TestMode) + + modelName := "xai-websocket-reconnect-model" + executor := &websocketDirectCaptureExecutor{provider: "xai"} + manager := coreauth.NewManager(nil, nil, nil) + manager.RegisterExecutor(executor) + auth := &coreauth.Auth{ + ID: "auth-xai-reconnect", + Provider: "xai", + Status: coreauth.StatusActive, + Attributes: map[string]string{"websockets": "true"}, + } + if _, err := manager.Register(context.Background(), auth); err != nil { + t.Fatalf("Register auth: %v", err) + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: modelName}}) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(auth.ID) + }) + + 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() + + 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() }() + + request := []byte(fmt.Sprintf(`{"type":"response.create","model":%q,"previous_response_id":"resp-old","input":[{"type":"message","id":"msg-2","role":"user","content":"second"}]}`, modelName)) + if errWrite := conn.WriteMessage(websocket.TextMessage, request); errWrite != nil { + t.Fatalf("write websocket message: %v", errWrite) + } + _, payload, errRead := conn.ReadMessage() + if errRead != nil { + t.Fatalf("read websocket response: %v", errRead) + } + if got := gjson.GetBytes(payload, "type").String(); got != wsEventTypeError { + t.Fatalf("response type = %q, want %q: %s", got, wsEventTypeError, payload) + } + if got := int(gjson.GetBytes(payload, "status").Int()); got != http.StatusConflict { + t.Fatalf("response status = %d, want %d: %s", got, http.StatusConflict, payload) + } + if got := gjson.GetBytes(payload, "error.code").String(); got != "previous_response_not_found" { + t.Fatalf("response error code = %q, want previous_response_not_found: %s", got, payload) + } + if got := len(executor.Payloads()); got != 0 { + t.Fatalf("executor payload count = %d, want 0", got) + } + + recoveryRequest := []byte(fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-1"},{"type":"message","id":"out-1","role":"assistant"},{"type":"message","id":"msg-2"}]}`, modelName)) + if errWrite := conn.WriteMessage(websocket.TextMessage, recoveryRequest); errWrite != nil { + t.Fatalf("write full recovery message: %v", errWrite) + } + _, recoveryPayload, errRead := conn.ReadMessage() + if errRead != nil { + t.Fatalf("read full recovery response: %v", errRead) + } + if got := gjson.GetBytes(recoveryPayload, "type").String(); got != wsEventTypeCompleted { + t.Fatalf("recovery response type = %q, want %q: %s", got, wsEventTypeCompleted, recoveryPayload) + } + payloads := executor.Payloads() + if len(payloads) != 1 { + t.Fatalf("executor payload count after recovery = %d, want 1", len(payloads)) + } + if got := len(gjson.GetBytes(payloads[0], "input").Array()); got != 3 { + t.Fatalf("full recovery input len = %d, want 3: %s", got, payloads[0]) + } } func TestResponsesWebsocketSwitchesPinnedAuthAcrossProviders(t *testing.T) { @@ -2422,7 +2983,6 @@ func TestResponsesWebsocketSwitchesPinnedAuthAcrossProviders(t *testing.T) { registry.GetGlobalRegistry().RegisterClient(xaiAuth.ID, xaiAuth.Provider, []*registry.ModelInfo{{ID: xaiModel}}) registry.GetGlobalRegistry().RegisterClient(codexAuth.ID, codexAuth.Provider, []*registry.ModelInfo{{ID: codexModel}}) registeredAuthIDs := []string{xaiAuth.ID, codexAuth.ID} - xaiAlternateAuthID := "" if testCase.xaiWebsockets { xaiAlternateAuth := &coreauth.Auth{ ID: "auth-alternate-" + xaiModel, @@ -2430,7 +2990,6 @@ func TestResponsesWebsocketSwitchesPinnedAuthAcrossProviders(t *testing.T) { Status: coreauth.StatusActive, Attributes: map[string]string{"websockets": "true"}, } - xaiAlternateAuthID = xaiAlternateAuth.ID selector.order = append(selector.order, xaiAlternateAuth.ID) if _, errRegister := manager.Register(context.Background(), xaiAlternateAuth); errRegister != nil { t.Fatalf("Register alternate xAI auth: %v", errRegister) @@ -2474,21 +3033,22 @@ func TestResponsesWebsocketSwitchesPinnedAuthAcrossProviders(t *testing.T) { `{"type":"response.create","input":[{"type":"message","id":"msg-xai-3"}]}`, } for index, request := range requests { + turn := index + 1 if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(request)); errWrite != nil { - t.Fatalf("write websocket message %d: %v", index+1, errWrite) + t.Fatalf("write websocket message %d: %v", turn, errWrite) } _, payload, errRead := conn.ReadMessage() if errRead != nil { - t.Fatalf("read websocket response %d: %v", index+1, errRead) + t.Fatalf("read websocket response %d: %v", turn, errRead) } if got := gjson.GetBytes(payload, "type").String(); got != wsEventTypeCompleted { - t.Fatalf("response %d type = %s, want %s: %s", index+1, got, wsEventTypeCompleted, payload) + t.Fatalf("response %d type = %s, want %s: %s", turn, got, wsEventTypeCompleted, payload) } } wantReturnAuthID := xaiAuth.ID if testCase.returnToDifferentXAIModel { - wantReturnAuthID = xaiAlternateAuthID + wantReturnAuthID = "auth-alternate-" + xaiModel } if got := xaiExecutor.AuthIDs(); len(got) != 3 || got[0] != xaiAuth.ID || got[1] != wantReturnAuthID || got[2] != wantReturnAuthID { t.Fatalf("xAI auth IDs = %v, want [%s %s %s]", got, xaiAuth.ID, wantReturnAuthID, wantReturnAuthID) @@ -2920,7 +3480,7 @@ func TestResponsesWebsocketDoesNotInjectPreviousResponseIDWhenPendingToolOutputM func TestResponsesWebsocketStripsGenerateWhenWebsocketAttemptFallsBackToHTTP(t *testing.T) { gin.SetMode(gin.TestMode) - selector := &orderedWebsocketSelector{order: []string{"auth-ws", "auth-http"}} + selector := &orderedWebsocketSelector{order: []string{"auth-ws", "auth-http", "auth-http"}} executor := &websocketBootstrapFallbackExecutor{} manager := coreauth.NewManager(nil, selector, nil) manager.RegisterExecutor(executor) @@ -2996,6 +3556,21 @@ func TestResponsesWebsocketStripsGenerateWhenWebsocketAttemptFallsBackToHTTP(t * if gjson.GetBytes(httpPayloads[0], "generate").Exists() { t.Fatalf("generate leaked after HTTP fallback: %s", httpPayloads[0]) } + + secondRequest := `{"type":"response.create","previous_response_id":"resp-http","input":[{"type":"message","id":"msg-2"}]}` + if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(secondRequest)); errWrite != nil { + t.Fatalf("write second websocket message: %v", errWrite) + } + _, secondPayload, errReadSecond := conn.ReadMessage() + if errReadSecond != nil { + t.Fatalf("read second websocket message: %v", errReadSecond) + } + if got := gjson.GetBytes(secondPayload, "type").String(); got != wsEventTypeCompleted { + t.Fatalf("second payload type = %s, want %s: %s", got, wsEventTypeCompleted, secondPayload) + } + if got := executor.AuthIDs(); len(got) != 3 || got[2] != "auth-http" { + t.Fatalf("selected auth IDs after HTTP retry = %v, want [auth-ws auth-http auth-http]", got) + } } func TestWebsocketClientAddressUsesGinClientIP(t *testing.T) { @@ -3093,45 +3668,38 @@ func TestResponsesWebsocketPinsOnlyWebsocketCapableAuth(t *testing.T) { } } -func TestResponsesWebsocketReleasesPinnedAuthAfterQuotaError(t *testing.T) { +func TestResponsesWebsocketUsesNativeIncrementalAfterPinningWebsocketAuthFromMixedPool(t *testing.T) { gin.SetMode(gin.TestMode) - selector := &orderedWebsocketSelector{order: []string{"auth-a", "auth-b"}} - executor := &websocketPinnedFailoverExecutor{} + modelName := "xai-mixed-pool-model" + selector := &orderedWebsocketSelector{order: []string{"auth-http", "auth-ws"}} + executor := &websocketDirectCaptureExecutor{provider: "xai"} manager := coreauth.NewManager(nil, selector, nil) manager.RegisterExecutor(executor) - - authA := &coreauth.Auth{ - ID: "auth-a", - Provider: executor.Identifier(), + authHTTP := &coreauth.Auth{ID: "auth-http", Provider: "xai", Status: coreauth.StatusActive} + if _, err := manager.Register(context.Background(), authHTTP); err != nil { + t.Fatalf("Register HTTP auth: %v", err) + } + authWS := &coreauth.Auth{ + ID: "auth-ws", + Provider: "xai", Status: coreauth.StatusActive, Attributes: map[string]string{"websockets": "true"}, } - if _, err := manager.Register(context.Background(), authA); err != nil { - t.Fatalf("Register auth A: %v", err) + if _, err := manager.Register(context.Background(), authWS); err != nil { + t.Fatalf("Register websocket auth: %v", err) } - authB := &coreauth.Auth{ - ID: "auth-b", - Provider: executor.Identifier(), - Status: coreauth.StatusActive, - Attributes: map[string]string{"websockets": "true"}, - } - if _, err := manager.Register(context.Background(), authB); err != nil { - t.Fatalf("Register auth B: %v", err) - } - - registry.GetGlobalRegistry().RegisterClient(authA.ID, authA.Provider, []*registry.ModelInfo{{ID: "quota-model"}}) - registry.GetGlobalRegistry().RegisterClient(authB.ID, authB.Provider, []*registry.ModelInfo{{ID: "quota-model"}}) + registry.GetGlobalRegistry().RegisterClient(authHTTP.ID, authHTTP.Provider, []*registry.ModelInfo{{ID: modelName}}) + registry.GetGlobalRegistry().RegisterClient(authWS.ID, authWS.Provider, []*registry.ModelInfo{{ID: modelName}}) t.Cleanup(func() { - registry.GetGlobalRegistry().UnregisterClient(authA.ID) - registry.GetGlobalRegistry().UnregisterClient(authB.ID) + registry.GetGlobalRegistry().UnregisterClient(authHTTP.ID) + registry.GetGlobalRegistry().UnregisterClient(authWS.ID) }) 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() @@ -3140,49 +3708,169 @@ func TestResponsesWebsocketReleasesPinnedAuthAfterQuotaError(t *testing.T) { if err != nil { t.Fatalf("dial websocket: %v", err) } - defer func() { - if errClose := conn.Close(); errClose != nil { - t.Fatalf("close websocket: %v", errClose) - } - }() + defer func() { _ = conn.Close() }() requests := []string{ - `{"type":"response.create","model":"quota-model","input":[{"type":"message","id":"msg-1"}]}`, - `{"type":"response.create","previous_response_id":"resp-auth-a-1","input":[{"type":"message","id":"msg-2"}]}`, - `{"type":"response.create","previous_response_id":"resp-auth-a-1","input":[{"type":"message","id":"msg-3"}]}`, + fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-1"}]}`, modelName), + `{"type":"response.create","previous_response_id":"resp-1","input":[{"type":"message","id":"msg-2"}]}`, + `{"type":"response.create","previous_response_id":"resp-2","input":[{"type":"message","id":"msg-3"}]}`, } - wantTypes := []string{wsEventTypeCompleted, wsEventTypeError, wsEventTypeCompleted} for i := range requests { if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(requests[i])); errWrite != nil { t.Fatalf("write websocket message %d: %v", i+1, errWrite) } - _, payload, errReadMessage := conn.ReadMessage() - if errReadMessage != nil { - t.Fatalf("read websocket message %d: %v", i+1, errReadMessage) - } - if got := gjson.GetBytes(payload, "type").String(); got != wantTypes[i] { - t.Fatalf("message %d payload type = %s, want %s: %s", i+1, got, wantTypes[i], payload) - } - if i == 1 && int(gjson.GetBytes(payload, "status").Int()) != http.StatusTooManyRequests { - t.Fatalf("quota payload status = %d, want %d: %s", gjson.GetBytes(payload, "status").Int(), http.StatusTooManyRequests, payload) + if _, _, errRead := conn.ReadMessage(); errRead != nil { + t.Fatalf("read websocket response %d: %v", i+1, errRead) } } - if got := executor.AuthIDs(); len(got) != 3 || got[0] != "auth-a" || got[1] != "auth-a" || got[2] != "auth-b" { - t.Fatalf("selected auth IDs = %v, want [auth-a auth-a auth-b]", got) + if got := executor.AuthIDs(); len(got) != 3 || got[0] != "auth-http" || got[1] != "auth-ws" || got[2] != "auth-ws" { + t.Fatalf("selected auth IDs = %v, want [auth-http auth-ws auth-ws]", got) } + payloads := executor.Payloads() + if len(payloads) != 3 { + t.Fatalf("payload count = %d, want 3", len(payloads)) + } + if gjson.GetBytes(payloads[1], "previous_response_id").Exists() || len(gjson.GetBytes(payloads[1], "input").Array()) != 3 { + t.Fatalf("first request on newly selected websocket auth must be canonical: %s", payloads[1]) + } + if got := gjson.GetBytes(payloads[2], "previous_response_id").String(); got != "resp-2" { + t.Fatalf("stable pinned websocket previous_response_id = %q, want resp-2: %s", got, payloads[2]) + } + input := gjson.GetBytes(payloads[2], "input").Array() + if len(input) != 1 || input[0].Get("id").String() != "msg-3" { + t.Fatalf("stable pinned websocket request is not incremental: %s", payloads[2]) + } +} - authBPayloads := executor.Payloads("auth-b") - if len(authBPayloads) != 1 { - t.Fatalf("auth-b payload count = %d, want 1", len(authBPayloads)) +func TestResponsesWebsocketReplaysImmediatelyAfterPinnedAuthFailure(t *testing.T) { + gin.SetMode(gin.TestMode) + + tests := []struct { + name string + status int + backupWebsocket bool + }{ + {name: "unauthorized to websocket", status: http.StatusUnauthorized, backupWebsocket: true}, + {name: "unauthorized to http", status: http.StatusUnauthorized, backupWebsocket: false}, + {name: "rate limit to websocket", status: http.StatusTooManyRequests, backupWebsocket: true}, + {name: "rate limit to http", status: http.StatusTooManyRequests, backupWebsocket: false}, } - authBPayload := authBPayloads[0] - if gjson.GetBytes(authBPayload, "previous_response_id").Exists() { - t.Fatalf("previous_response_id leaked after auth failover: %s", authBPayload) + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + modelName := fmt.Sprintf("credential-failure-%d-%t-model", tc.status, tc.backupWebsocket) + selector := &orderedWebsocketSelector{order: []string{"auth-a", "auth-b"}} + executor := &websocketPinnedFailoverExecutor{failStatus: tc.status} + manager := coreauth.NewManager(nil, selector, nil) + manager.RegisterExecutor(executor) + + authA := &coreauth.Auth{ + ID: "auth-a", + Provider: executor.Identifier(), + Status: coreauth.StatusActive, + Attributes: map[string]string{"websockets": "true"}, + } + if _, err := manager.Register(context.Background(), authA); err != nil { + t.Fatalf("Register auth A: %v", err) + } + authB := &coreauth.Auth{ + ID: "auth-b", + Provider: executor.Identifier(), + Status: coreauth.StatusActive, + Attributes: map[string]string{"websockets": strconv.FormatBool(tc.backupWebsocket)}, + } + if _, err := manager.Register(context.Background(), authB); err != nil { + t.Fatalf("Register auth B: %v", err) + } + + registry.GetGlobalRegistry().RegisterClient(authA.ID, authA.Provider, []*registry.ModelInfo{{ID: modelName}}) + registry.GetGlobalRegistry().RegisterClient(authB.ID, authB.Provider, []*registry.ModelInfo{{ID: modelName}}) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(authA.ID) + registry.GetGlobalRegistry().UnregisterClient(authB.ID) + }) + + 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() + + 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() }() + + firstRequest := fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-1"}]}`, modelName) + if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(firstRequest)); errWrite != nil { + t.Fatalf("write first websocket message: %v", errWrite) + } + if _, payload, errRead := conn.ReadMessage(); errRead != nil || gjson.GetBytes(payload, "type").String() != wsEventTypeCompleted { + t.Fatalf("first websocket response = %s, err=%v", payload, errRead) + } + + secondRequest := `{"type":"response.create","previous_response_id":"resp-auth-a-1","input":[{"type":"message","id":"msg-2"}]}` + if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(secondRequest)); errWrite != nil { + t.Fatalf("write second websocket message: %v", errWrite) + } + _, _, errReadClose := conn.ReadMessage() + var replayClose *websocket.CloseError + if !errors.As(errReadClose, &replayClose) || replayClose.Code != websocket.CloseServiceRestart || replayClose.Text != wsHTTPReplayRequiredCloseReason { + t.Fatalf("credential failure response = %v, want replay close %d %q", errReadClose, websocket.CloseServiceRestart, wsHTTPReplayRequiredCloseReason) + } + if got := executor.AuthIDs(); len(got) != 2 || got[0] != "auth-a" || got[1] != "auth-a" { + t.Fatalf("selected auth IDs before replay = %v, want [auth-a auth-a]", got) + } + + replayConn, _, errDialReplay := websocket.DefaultDialer.Dial(wsURL, nil) + if errDialReplay != nil { + t.Fatalf("dial replay websocket: %v", errDialReplay) + } + defer func() { _ = replayConn.Close() }() + fullReplay := fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-1"},{"type":"message","id":"out-auth-a-1"},{"type":"message","id":"msg-2"}]}`, modelName) + if errWrite := replayConn.WriteMessage(websocket.TextMessage, []byte(fullReplay)); errWrite != nil { + t.Fatalf("write full replay: %v", errWrite) + } + if _, replayPayload, errReadReplay := replayConn.ReadMessage(); errReadReplay != nil || gjson.GetBytes(replayPayload, "type").String() != wsEventTypeCompleted { + t.Fatalf("full replay response = %s, err=%v", replayPayload, errReadReplay) + } + if got := executor.AuthIDs(); len(got) != 3 || got[2] != "auth-b" { + t.Fatalf("selected auth IDs after replay = %v, want [auth-a auth-a auth-b]", got) + } + authBPayloads := executor.Payloads("auth-b") + if len(authBPayloads) != 1 { + t.Fatalf("auth-b payloads = %d, want 1", len(authBPayloads)) + } + authBPayload := authBPayloads[0] + if gjson.GetBytes(authBPayload, "previous_response_id").Exists() || len(gjson.GetBytes(authBPayload, "input").Array()) != 3 { + t.Fatalf("auth-b did not receive full replay: %s", authBPayload) + } + }) } - authBInput := gjson.GetBytes(authBPayload, "input").Raw - if !strings.Contains(authBInput, `"id":"msg-1"`) || !strings.Contains(authBInput, `"id":"msg-3"`) { - t.Fatalf("auth-b replay input missing expected transcript items: %s", authBInput) +} + +func TestShouldReplayResponsesWebsocketPinnedAuthFailure(t *testing.T) { + cases := []struct { + name string + err *interfaces.ErrorMessage + want bool + }{ + {name: "nil", err: nil, want: false}, + {name: "unauthorized", err: &interfaces.ErrorMessage{StatusCode: http.StatusUnauthorized}, want: true}, + {name: "rate limit", err: &interfaces.ErrorMessage{StatusCode: http.StatusTooManyRequests}, want: true}, + {name: "forbidden", err: &interfaces.ErrorMessage{StatusCode: http.StatusForbidden}, want: false}, + {name: "service unavailable", err: &interfaces.ErrorMessage{StatusCode: http.StatusServiceUnavailable}, want: false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := shouldReplayResponsesWebsocketPinnedAuthFailure(tc.err); got != tc.want { + t.Fatalf("shouldReplayResponsesWebsocketPinnedAuthFailure() = %v, want %v", got, tc.want) + } + }) } } @@ -3215,7 +3903,7 @@ type websocketPinnedPrematureCloseExecutor struct { payloads map[string][][]byte } -func (e *websocketPinnedPrematureCloseExecutor) Identifier() string { return "test-provider" } +func (e *websocketPinnedPrematureCloseExecutor) Identifier() string { return "xai" } func (e *websocketPinnedPrematureCloseExecutor) Execute(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) { return coreexecutor.Response{}, errors.New("not implemented") @@ -3335,66 +4023,69 @@ func TestResponsesWebsocketReleasesPinnedAuthAfterStreamClosed408(t *testing.T) } }() - requests := []string{ - `{"type":"response.create","model":"stream-model","input":[{"type":"message","id":"msg-1"}]}`, - `{"type":"response.create","previous_response_id":"resp-auth-a-1","input":[{"type":"message","id":"msg-2"}]}`, - `{"type":"response.create","previous_response_id":"resp-auth-a-1","input":[{"type":"message","id":"msg-3"}]}`, + firstRequest := `{"type":"response.create","model":"stream-model","input":[{"type":"message","id":"msg-1"}]}` + if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(firstRequest)); errWrite != nil { + t.Fatalf("write first websocket message: %v", errWrite) } - wantTypes := []string{wsEventTypeCompleted, wsEventTypeError, wsEventTypeCompleted} - for i := range requests { - if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(requests[i])); errWrite != nil { - t.Fatalf("write websocket message %d: %v", i+1, errWrite) + if _, payload, errRead := conn.ReadMessage(); errRead != nil || gjson.GetBytes(payload, "type").String() != wsEventTypeCompleted { + t.Fatalf("first websocket response = %s, err=%v", payload, errRead) + } + + secondRequest := `{"type":"response.create","previous_response_id":"resp-auth-a-1","input":[{"type":"message","id":"msg-2"}]}` + if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(secondRequest)); errWrite != nil { + t.Fatalf("write second websocket message: %v", errWrite) + } + for { + _, payload, errRead := conn.ReadMessage() + if errRead != nil { + t.Fatalf("read stream-closed response: %v", errRead) } - if i == 1 { - gotError := false - for { - _, payload, errReadMessage := conn.ReadMessage() - if errReadMessage != nil { - t.Fatalf("read websocket message %d: %v", i+1, errReadMessage) - } - got := gjson.GetBytes(payload, "type").String() - if got == wsEventTypeError { - if int(gjson.GetBytes(payload, "status").Int()) != http.StatusRequestTimeout { - t.Fatalf("stream-closed payload status = %d, want %d: %s", gjson.GetBytes(payload, "status").Int(), http.StatusRequestTimeout, payload) - } - gotError = true - break - } - if got == wsEventTypeCompleted { - t.Fatalf("message %d unexpectedly completed: %s", i+1, payload) - } + eventType := gjson.GetBytes(payload, "type").String() + if eventType == wsEventTypeError { + if got := int(gjson.GetBytes(payload, "status").Int()); got != http.StatusRequestTimeout { + t.Fatalf("stream-closed status = %d, want %d: %s", got, http.StatusRequestTimeout, payload) } - if !gotError { - t.Fatalf("message %d did not return stream-closed error", i+1) - } - continue + break } - _, payload, errReadMessage := conn.ReadMessage() - if errReadMessage != nil { - t.Fatalf("read websocket message %d: %v", i+1, errReadMessage) - } - if got := gjson.GetBytes(payload, "type").String(); got != wantTypes[i] { - t.Fatalf("message %d payload type = %s, want %s: %s", i+1, got, wantTypes[i], payload) + if eventType == wsEventTypeCompleted { + t.Fatalf("stream-closed turn unexpectedly completed: %s", payload) } } + thirdDelta := `{"type":"response.create","previous_response_id":"resp-auth-a-1","input":[{"type":"message","id":"msg-3"}]}` + if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(thirdDelta)); errWrite != nil { + t.Fatalf("write third websocket message: %v", errWrite) + } + _, _, errReadClose := conn.ReadMessage() + var replayClose *websocket.CloseError + if !errors.As(errReadClose, &replayClose) || replayClose.Code != websocket.CloseServiceRestart { + t.Fatalf("third websocket response error = %v, want replay close", errReadClose) + } + if got := executor.AuthIDs(); len(got) != 2 || got[0] != "auth-a" || got[1] != "auth-a" { + t.Fatalf("selected auth IDs before replay = %v, want [auth-a auth-a]", got) + } + + replayConn, _, errDialReplay := websocket.DefaultDialer.Dial(wsURL, nil) + if errDialReplay != nil { + t.Fatalf("dial replay websocket: %v", errDialReplay) + } + defer func() { _ = replayConn.Close() }() + fullReplay := `{"type":"response.create","model":"stream-model","input":[{"type":"message","id":"msg-1"},{"type":"message","id":"out-auth-a-1"},{"type":"message","id":"msg-3"}]}` + if errWrite := replayConn.WriteMessage(websocket.TextMessage, []byte(fullReplay)); errWrite != nil { + t.Fatalf("write full replay: %v", errWrite) + } + if _, replayResponse, errReadReplay := replayConn.ReadMessage(); errReadReplay != nil || gjson.GetBytes(replayResponse, "type").String() != wsEventTypeCompleted { + t.Fatalf("full replay response = %s, err=%v", replayResponse, errReadReplay) + } authIDs := executor.AuthIDs() if len(authIDs) != 3 || authIDs[0] != "auth-a" || authIDs[1] != "auth-a" { - t.Fatalf("selected auth IDs = %v, want auth-a for first two turns", authIDs) + t.Fatalf("selected auth IDs after replay = %v, want auth-a for the first two turns", authIDs) } - replayAuthID := authIDs[2] replayPayloads := executor.Payloads(replayAuthID) - if len(replayPayloads) == 0 { - t.Fatalf("replay auth %s has no payloads", replayAuthID) - } replayPayload := replayPayloads[len(replayPayloads)-1] - if gjson.GetBytes(replayPayload, "previous_response_id").Exists() { - t.Fatalf("previous_response_id leaked after stream-closed replay: %s", replayPayload) - } - replayInput := gjson.GetBytes(replayPayload, "input").Raw - if !strings.Contains(replayInput, `"id":"msg-1"`) || !strings.Contains(replayInput, `"id":"msg-3"`) { - t.Fatalf("replay input missing expected transcript items: %s", replayInput) + if gjson.GetBytes(replayPayload, "previous_response_id").Exists() || len(gjson.GetBytes(replayPayload, "input").Array()) != 3 { + t.Fatalf("replay auth %s did not receive full replay: %s", replayAuthID, replayPayload) } } diff --git a/sdk/cliproxy/executor/context.go b/sdk/cliproxy/executor/context.go index 367b507eb..c18d3f684 100644 --- a/sdk/cliproxy/executor/context.go +++ b/sdk/cliproxy/executor/context.go @@ -3,6 +3,7 @@ package executor import "context" type downstreamWebsocketContextKey struct{} +type requireUpstreamWebsocketContextKey struct{} // WithDownstreamWebsocket marks the current request as coming from a downstream websocket connection. func WithDownstreamWebsocket(ctx context.Context) context.Context { @@ -21,3 +22,21 @@ func DownstreamWebsocket(ctx context.Context) bool { enabled, ok := raw.(bool) return ok && enabled } + +// WithRequiredUpstreamWebsocket marks a request whose incremental context is valid only on the current upstream websocket. +func WithRequiredUpstreamWebsocket(ctx context.Context) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, requireUpstreamWebsocketContextKey{}, true) +} + +// RequiredUpstreamWebsocket reports whether falling back to an HTTP upstream would lose request context. +func RequiredUpstreamWebsocket(ctx context.Context) bool { + if ctx == nil { + return false + } + raw := ctx.Value(requireUpstreamWebsocketContextKey{}) + enabled, ok := raw.(bool) + return ok && enabled +} diff --git a/sdk/cliproxy/executor/websocket.go b/sdk/cliproxy/executor/websocket.go new file mode 100644 index 000000000..1fa0d79e8 --- /dev/null +++ b/sdk/cliproxy/executor/websocket.go @@ -0,0 +1,29 @@ +package executor + +import ( + "errors" + "net/http" +) + +// UpstreamWebsocketReplayRequiredError indicates that an incremental request +// cannot safely continue because its upstream websocket is no longer reusable. +type UpstreamWebsocketReplayRequiredError struct{} + +func (*UpstreamWebsocketReplayRequiredError) Error() string { + return `{"error":{"message":"upstream transport requires full HTTP replay","type":"server_error","code":"upstream_http_replay_required","status":426}}` +} + +func (*UpstreamWebsocketReplayRequiredError) StatusCode() int { return http.StatusUpgradeRequired } + +func (*UpstreamWebsocketReplayRequiredError) IsRequestScoped() bool { return true } + +// NewUpstreamWebsocketReplayRequiredError creates a request-scoped replay signal. +func NewUpstreamWebsocketReplayRequiredError() error { + return &UpstreamWebsocketReplayRequiredError{} +} + +// IsUpstreamWebsocketReplayRequired reports whether err is the internal replay signal. +func IsUpstreamWebsocketReplayRequired(err error) bool { + var replayErr *UpstreamWebsocketReplayRequiredError + return errors.As(err, &replayErr) +} diff --git a/sdk/cliproxy/executor/websocket_test.go b/sdk/cliproxy/executor/websocket_test.go new file mode 100644 index 000000000..f4327fb62 --- /dev/null +++ b/sdk/cliproxy/executor/websocket_test.go @@ -0,0 +1,25 @@ +package executor + +import ( + "fmt" + "net/http" + "testing" +) + +func TestUpstreamWebsocketReplayRequiredError(t *testing.T) { + err := NewUpstreamWebsocketReplayRequiredError() + if !IsUpstreamWebsocketReplayRequired(err) { + t.Fatal("replay error was not recognized") + } + if !IsUpstreamWebsocketReplayRequired(fmt.Errorf("wrapped: %w", err)) { + t.Fatal("wrapped replay error was not recognized") + } + statusErr, ok := err.(interface{ StatusCode() int }) + if !ok || statusErr.StatusCode() != http.StatusUpgradeRequired { + t.Fatalf("replay error = %T %v, want status 426", err, err) + } + requestErr, ok := err.(RequestScopedError) + if !ok || !requestErr.IsRequestScoped() { + t.Fatalf("replay error = %T, want request scoped", err) + } +} From 840ba5dcc1b8c0ab521650e68136a866588cb049 Mon Sep 17 00:00:00 2001 From: sususu Date: Thu, 23 Jul 2026 13:07:37 +0800 Subject: [PATCH 2/3] fix(xai): preserve compacted websocket transcript state --- .../executor/xai_websockets_executor.go | 139 +++++++++++++++--- .../executor/xai_websockets_executor_test.go | 121 +++++++++++++++ 2 files changed, 240 insertions(+), 20 deletions(-) diff --git a/internal/runtime/executor/xai_websockets_executor.go b/internal/runtime/executor/xai_websockets_executor.go index 1f04538e6..1e2e7c826 100644 --- a/internal/runtime/executor/xai_websockets_executor.go +++ b/internal/runtime/executor/xai_websockets_executor.go @@ -48,18 +48,21 @@ type xaiWebsocketIDStateStore struct { } type xaiWebsocketIDState struct { - mu sync.Mutex - downstreamToUpstream map[string]string - sequence int - transcriptInput []json.RawMessage + requestMu sync.Mutex + mu sync.Mutex + downstreamToUpstream map[string]string + sequence int + transcriptInput []json.RawMessage + replayCompactedTranscriptOnReset bool } type xaiWebsocketRequestIDMapper struct { - state *xaiWebsocketIDState - downstreamPreviousID string - upstreamPreviousID string - upstreamResponseID string - downstreamResponseID string + state *xaiWebsocketIDState + downstreamPreviousID string + upstreamPreviousID string + upstreamResponseID string + downstreamResponseID string + replayedCompactedTranscript bool } func NewXAIWebsocketsExecutor(cfg *config.Config) *XAIWebsocketsExecutor { @@ -177,20 +180,21 @@ func (s *xaiWebsocketIDState) prependTranscriptInput(payload []byte) []byte { return out } -func (s *xaiWebsocketIDState) recordTranscriptTurn(requestPayload []byte, completedPayload []byte) { +func (s *xaiWebsocketIDState) recordTranscriptTurn(requestPayload []byte, completedPayload []byte, reset bool) { if s == nil || len(requestPayload) == 0 || len(completedPayload) == 0 { return } inputItems := xaiJSONRawMessages(gjson.GetBytes(requestPayload, "input")) outputItems := xaiJSONRawMessages(gjson.GetBytes(completedPayload, "response.output")) - if len(inputItems) == 0 && len(outputItems) == 0 { - return - } s.mu.Lock() defer s.mu.Unlock() - if strings.TrimSpace(gjson.GetBytes(requestPayload, "previous_response_id").String()) == "" { + if reset { s.transcriptInput = nil + s.replayCompactedTranscriptOnReset = false + } + if len(inputItems) == 0 && len(outputItems) == 0 { + return } s.transcriptInput = append(s.transcriptInput, inputItems...) s.transcriptInput = append(s.transcriptInput, outputItems...) @@ -210,9 +214,34 @@ func (s *xaiWebsocketIDState) replaceTranscriptWithItems(items ...[]byte) { } s.mu.Lock() s.transcriptInput = next + s.replayCompactedTranscriptOnReset = len(next) > 0 s.mu.Unlock() } +func (s *xaiWebsocketIDState) prependCompactedTranscriptOnReset(payload []byte) ([]byte, bool) { + if s == nil || len(payload) == 0 { + return payload, false + } + s.mu.Lock() + if !s.replayCompactedTranscriptOnReset || len(s.transcriptInput) == 0 { + s.mu.Unlock() + return payload, false + } + prefix := make([]json.RawMessage, 0, len(s.transcriptInput)) + for _, item := range s.transcriptInput { + prefix = append(prefix, bytes.Clone(item)) + } + s.mu.Unlock() + + current := xaiJSONRawMessages(gjson.GetBytes(payload, "input")) + merged := append(prefix, current...) + out, errSet := sjson.SetRawBytes(payload, "input", xaiMarshalRawMessages(merged)) + if errSet != nil { + return payload, false + } + return out, true +} + func xaiJSONRawMessages(result gjson.Result) []json.RawMessage { if !result.Exists() || !result.IsArray() { return nil @@ -243,7 +272,16 @@ func xaiMarshalRawMessages(items []json.RawMessage) []byte { } func (m *xaiWebsocketRequestIDMapper) upstreamRequestPayload(payload []byte) []byte { - if m == nil || len(payload) == 0 || m.downstreamPreviousID == m.upstreamPreviousID { + if m == nil || len(payload) == 0 { + return payload + } + if m.downstreamPreviousID == m.upstreamPreviousID { + requestType := strings.TrimSpace(gjson.GetBytes(payload, "type").String()) + if m.downstreamPreviousID == "" && requestType == "response.append" && m.state != nil { + out, replayed := m.state.prependCompactedTranscriptOnReset(payload) + m.replayedCompactedTranscript = replayed + return out + } return payload } if m.upstreamPreviousID == "" { @@ -251,6 +289,7 @@ func (m *xaiWebsocketRequestIDMapper) upstreamRequestPayload(payload []byte) []b if errDelete == nil { if m.downstreamPreviousID != "" && m.state != nil { out = m.state.prependTranscriptInput(out) + m.replayedCompactedTranscript = true } return out } @@ -396,11 +435,30 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox if stateSessionID == "" { stateSessionID = executionSessionID } - idMapper := newXAIWebsocketRequestIDMapper(e.idStore, stateSessionID, req.Payload) + state := getXAIWebsocketIDState(e.idStore, stateSessionID) + stateRequestLocked := false + stateRequestLockTransferred := false + if executionSessionID == "" && state != nil { + state.requestMu.Lock() + stateRequestLocked = true + } + defer func() { + if stateRequestLocked && !stateRequestLockTransferred { + state.requestMu.Unlock() + } + }() if xaiInputHasItemType(req.Payload, "compaction_trigger") { if cliproxyexecutor.RequiredUpstreamWebsocket(ctx) { return nil, cliproxyexecutor.NewUpstreamWebsocketReplayRequiredError() } + if executionSessionID != "" { + sess := e.getOrCreateSession(executionSessionID) + if sess != nil { + sess.reqMu.Lock() + defer sess.reqMu.Unlock() + } + } + idMapper := newXAIWebsocketRequestIDMapper(e.idStore, stateSessionID, req.Payload) return e.executeCompactionTriggerFromWebsocketContext(ctx, auth, req, opts, idMapper) } @@ -440,6 +498,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox sess.reqMu.Lock() } } + idMapper := newXAIWebsocketRequestIDMapper(e.idStore, stateSessionID, req.Payload) if idMapper != nil { if websocketSessionTargetChanged(sess, authID, wsURL) { idMapper.upstreamPreviousID = "" @@ -450,6 +509,9 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox wsHeaders := applyXAIWebsocketHeaders(http.Header{}, auth, token, prepared.sessionID) wsReqBody := buildXAIWebsocketRequestBody(prepared.body) + requestType := strings.TrimSpace(gjson.GetBytes(req.Payload, "type").String()) + transcriptReset := strings.TrimSpace(gjson.GetBytes(wsReqBody, "previous_response_id").String()) == "" && + (requestType != "response.append" || (idMapper != nil && idMapper.replayedCompactedTranscript)) warmupRequest := xaiWebsocketGenerateFalse(wsReqBody) wsReqLog := helps.UpstreamRequestLog{ @@ -580,7 +642,13 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox } out := make(chan cliproxyexecutor.StreamChunk) + if stateRequestLocked { + stateRequestLockTransferred = true + } go func() { + if stateRequestLocked { + defer state.requestMu.Unlock() + } terminateReason := "completed" var terminateErr error @@ -687,6 +755,10 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox case "response.created": if warmupRequest { warmupCompletedPayload = buildXAIWebsocketWarmupCompletedPayload(payload) + if idMapper != nil && idMapper.state != nil && !recordedTranscript { + idMapper.state.recordTranscriptTurn(wsReqBody, warmupCompletedPayload, transcriptReset) + recordedTranscript = true + } logXAIWebsocketWarmupCompleted(executionSessionID, authID, wsURL, payload) } case "response.output_item.done": @@ -700,7 +772,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox payload = xaiNormalizeReasoningSummaryData(payload) cacheXAIReasoningReplayFromCompleted(ctx, prepared.replayScope, payload) if !warmupRequest && idMapper != nil && idMapper.state != nil && !recordedTranscript { - idMapper.state.recordTranscriptTurn(wsReqBody, payload) + idMapper.state.recordTranscriptTurn(wsReqBody, payload, transcriptReset) recordedTranscript = true } case "response.done": @@ -709,7 +781,7 @@ func (e *XAIWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *cliprox reporter.Publish(ctx, detail) } if !warmupRequest && idMapper != nil && idMapper.state != nil && !recordedTranscript { - idMapper.state.recordTranscriptTurn(wsReqBody, payload) + idMapper.state.recordTranscriptTurn(wsReqBody, payload, transcriptReset) recordedTranscript = true } } @@ -803,8 +875,11 @@ func (e *XAIWebsocketsExecutor) executeCompactionTriggerFromWebsocketContext(ctx return nil, err } - responseID := xaiCompactionResponseID(data) - idMapper.state.replaceTranscriptWithItems(xaiCompactionOutputItem(data, responseID)) + responseID, compactionItem, errValidate := validateXAIWebsocketCompactionResponse(data) + if errValidate != nil { + return nil, errValidate + } + idMapper.state.replaceTranscriptWithItems(compactionItem) idMapper.state.mapDownstreamToUpstream(responseID, "") headers = headers.Clone() @@ -822,6 +897,30 @@ func (e *XAIWebsocketsExecutor) executeCompactionTriggerFromWebsocketContext(ctx return &cliproxyexecutor.StreamResult{Headers: headers, Chunks: out}, nil } +func validateXAIWebsocketCompactionResponse(data []byte) (string, []byte, error) { + if len(data) == 0 || !json.Valid(data) { + return "", nil, statusErr{code: http.StatusBadGateway, msg: "xai websocket compaction returned invalid JSON"} + } + responseIDResult := gjson.GetBytes(data, "id") + output := gjson.GetBytes(data, "output") + if responseIDResult.Type != gjson.String || strings.TrimSpace(responseIDResult.String()) == "" || !output.Exists() || !output.IsArray() { + return "", nil, statusErr{code: http.StatusBadGateway, msg: "xai websocket compaction response is missing compacted state"} + } + items := output.Array() + if len(items) == 0 { + return "", nil, statusErr{code: http.StatusBadGateway, msg: "xai websocket compaction response is missing compacted state"} + } + item := items[0] + itemType := item.Get("type") + encryptedContent := item.Get("encrypted_content") + if item.Type != gjson.JSON || itemType.Type != gjson.String || strings.TrimSpace(itemType.String()) != "compaction" || + encryptedContent.Type != gjson.String || strings.TrimSpace(encryptedContent.String()) == "" { + return "", nil, statusErr{code: http.StatusBadGateway, msg: "xai websocket compaction response is missing compacted state"} + } + normalizedResponseID := xaiCompactionResponseID(data) + return normalizedResponseID, xaiCompactionOutputItem(data, normalizedResponseID), nil +} + func buildXAIWebsocketCompactionPayload(payload []byte, transcriptInput []byte) ([]byte, error) { if len(payload) == 0 { payload = []byte(`{}`) diff --git a/internal/runtime/executor/xai_websockets_executor_test.go b/internal/runtime/executor/xai_websockets_executor_test.go index 75ceae3c7..065e73f5a 100644 --- a/internal/runtime/executor/xai_websockets_executor_test.go +++ b/internal/runtime/executor/xai_websockets_executor_test.go @@ -1290,6 +1290,127 @@ func TestXAIWebsocketsExecuteStreamCompactionTriggerUsesHTTPCompactWithRecordedC } } +func TestXAIWebsocketPostCompactionAppendWithoutPreviousReplaysCompactedTranscript(t *testing.T) { + store := &xaiWebsocketIDStateStore{sessions: make(map[string]*xaiWebsocketIDState)} + state := getXAIWebsocketIDState(store, "post-compaction-append-session") + state.replaceTranscriptWithItems([]byte(`{"type":"compaction","encrypted_content":"compact-state"}`)) + state.mapDownstreamToUpstream("resp-compact", "") + + fullReset := []byte(`{"type":"response.create","model":"grok-4.3","input":[{"type":"message","id":"msg-full"}]}`) + fullMapper := newXAIWebsocketRequestIDMapper(store, "post-compaction-append-session", fullReset) + if full := fullMapper.upstreamRequestPayload(fullReset); len(gjson.GetBytes(full, "input").Array()) != 1 { + t.Fatalf("self-contained response.create unexpectedly replayed compacted transcript: %s", full) + } + + payload := []byte(`{"type":"response.append","model":"grok-4.3","input":[{"type":"message","id":"msg-2","role":"user","content":"second"}]}`) + mapper := newXAIWebsocketRequestIDMapper(store, "post-compaction-append-session", payload) + got := mapper.upstreamRequestPayload(payload) + input := gjson.GetBytes(got, "input").Array() + if len(input) != 2 { + t.Fatalf("post-compaction append input len = %d, want 2: %s", len(input), got) + } + if gotType := input[0].Get("type").String(); gotType != "compaction" { + t.Fatalf("post-compaction append input[0].type = %q, want compaction: %s", gotType, got) + } + if gotID := input[1].Get("id").String(); gotID != "msg-2" { + t.Fatalf("post-compaction append input[1].id = %q, want msg-2: %s", gotID, got) + } + + state.recordTranscriptTurn(got, []byte(`{"type":"response.completed","response":{"id":"resp-after-compact","output":[{"type":"message","id":"out-2"}]}}`), true) + nextPayload := []byte(`{"type":"response.create","model":"grok-4.3","input":[{"type":"message","id":"msg-3"}]}`) + nextMapper := newXAIWebsocketRequestIDMapper(store, "post-compaction-append-session", nextPayload) + next := nextMapper.upstreamRequestPayload(nextPayload) + if nextInput := gjson.GetBytes(next, "input").Array(); len(nextInput) != 1 || nextInput[0].Get("id").String() != "msg-3" { + t.Fatalf("compacted transcript replay was not cleared after success: %s", next) + } +} + +func TestXAIWebsocketPostCompactionWarmupPreservesTranscriptForLaterCompaction(t *testing.T) { + store := &xaiWebsocketIDStateStore{sessions: make(map[string]*xaiWebsocketIDState)} + state := getXAIWebsocketIDState(store, "warmup-reset-session") + state.replaceTranscriptWithItems([]byte(`{"type":"compaction","encrypted_content":"compact-state"}`)) + + warmupPayload := []byte(`{"type":"response.append","model":"grok-4.3","generate":false,"input":[{"type":"message","id":"warmup-context"}]}`) + warmupMapper := newXAIWebsocketRequestIDMapper(store, "warmup-reset-session", warmupPayload) + warmupUpstream := warmupMapper.upstreamRequestPayload(warmupPayload) + if !warmupMapper.replayedCompactedTranscript { + t.Fatal("post-compaction warmup did not mark full transcript replay") + } + state.recordTranscriptTurn( + warmupUpstream, + []byte(`{"type":"response.completed","response":{"id":"resp-warmup","output":[]}}`), + true, + ) + + appendPayload := []byte(`{"type":"response.append","model":"grok-4.3","input":[{"type":"message","id":"msg-after-warmup"}]}`) + appendMapper := newXAIWebsocketRequestIDMapper(store, "warmup-reset-session", appendPayload) + appendUpstream := appendMapper.upstreamRequestPayload(appendPayload) + input := gjson.GetBytes(appendUpstream, "input").Array() + if len(input) != 1 || input[0].Get("id").String() != "msg-after-warmup" { + t.Fatalf("warmup retained pending replay instead of native append: %s", appendUpstream) + } + state.recordTranscriptTurn( + appendUpstream, + []byte(`{"type":"response.completed","response":{"id":"resp-after-warmup","output":[{"type":"message","id":"out-after-warmup"}]}}`), + false, + ) + + transcript := gjson.ParseBytes(state.snapshotTranscriptInput()).Array() + wantTypes := []string{"compaction", "message", "message", "message"} + if len(transcript) != len(wantTypes) { + t.Fatalf("post-warmup transcript len = %d, want %d: %s", len(transcript), len(wantTypes), state.snapshotTranscriptInput()) + } + for i, wantType := range wantTypes { + if gotType := transcript[i].Get("type").String(); gotType != wantType { + t.Fatalf("post-warmup transcript[%d].type = %q, want %q: %s", i, gotType, wantType, state.snapshotTranscriptInput()) + } + } +} + +func TestXAIWebsocketEmptyFullResetClearsPendingCompactionReplay(t *testing.T) { + store := &xaiWebsocketIDStateStore{sessions: make(map[string]*xaiWebsocketIDState)} + state := getXAIWebsocketIDState(store, "empty-reset-session") + state.replaceTranscriptWithItems([]byte(`{"type":"compaction","encrypted_content":"stale-compact-state"}`)) + state.recordTranscriptTurn( + []byte(`{"type":"response.create","model":"grok-4.3","input":[]}`), + []byte(`{"type":"response.completed","response":{"id":"resp-empty","output":[]}}`), + true, + ) + + appendPayload := []byte(`{"type":"response.append","model":"grok-4.3","input":[{"type":"message","id":"msg-new"}]}`) + mapper := newXAIWebsocketRequestIDMapper(store, "empty-reset-session", appendPayload) + got := mapper.upstreamRequestPayload(appendPayload) + input := gjson.GetBytes(got, "input").Array() + if len(input) != 1 || input[0].Get("id").String() != "msg-new" { + t.Fatalf("empty full reset retained stale compaction replay: %s", got) + } +} + +func TestValidateXAIWebsocketCompactionResponse(t *testing.T) { + valid := []byte(`{"id":"resp_compact","output":[{"type":"compaction","encrypted_content":"opaque-state"}]}`) + responseID, item, err := validateXAIWebsocketCompactionResponse(valid) + if err != nil { + t.Fatalf("valid compaction response error: %v", err) + } + if responseID != "resp_compact" || gjson.GetBytes(item, "encrypted_content").String() != "opaque-state" { + t.Fatalf("validated compaction response = id:%q item:%s", responseID, item) + } + + for _, payload := range [][]byte{ + nil, + []byte(`{}`), + []byte(`{"id":"resp_empty","output":[]}`), + []byte(`{"id":123,"output":[{"type":"compaction","encrypted_content":"opaque"}]}`), + []byte(`{"id":"resp_object","output":{"0":{"type":"compaction","encrypted_content":"opaque"}}}`), + []byte(`{"id":"resp_numeric_state","output":[{"type":"compaction","encrypted_content":123}]}`), + []byte(`{"id":"resp_missing_state","output":[{"type":"compaction"}]}`), + } { + if _, _, errInvalid := validateXAIWebsocketCompactionResponse(payload); errInvalid == nil { + t.Fatalf("invalid compaction response accepted: %s", payload) + } + } +} + func TestBuildXAIWebsocketRequestBodySetsStoreAndKeepsPromptCacheKey(t *testing.T) { body := []byte(`{"model":"grok-4.3","stream":true,"stream_options":{"include_usage":true},"background":true,"prompt_cache_key":"cache-1","previous_response_id":"resp-prev","instructions":"system prompt","input":[{"type":"message","role":"user","content":"hello"}]}`) From a661172b1fe6594f7a0ec379d14d8dc731511e43 Mon Sep 17 00:00:00 2001 From: sususu Date: Thu, 23 Jul 2026 13:07:57 +0800 Subject: [PATCH 3/3] fix(responses): commit websocket tool cache atomically --- .../openai/openai_responses_websocket.go | 14 +- .../openai/openai_responses_websocket_test.go | 118 +++++++++++++++++ ...nai_responses_websocket_toolcall_repair.go | 123 +++++++++++++++++- 3 files changed, 248 insertions(+), 7 deletions(-) diff --git a/sdk/api/handlers/openai/openai_responses_websocket.go b/sdk/api/handlers/openai/openai_responses_websocket.go index a0aa52c9b..2d017d36a 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket.go +++ b/sdk/api/handlers/openai/openai_responses_websocket.go @@ -647,6 +647,7 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { continue } + toolCacheTurn := newResponsesWebsocketToolCacheTurn(downstreamSessionKey) previousLastRequest := bytes.Clone(lastRequest) previousLastResponseOutput := bytes.Clone(lastResponseOutput) previousLastResponseID := lastResponseID @@ -656,7 +657,8 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { passthroughModelName = modelName } } else { - requestJSON = repairResponsesWebsocketToolCalls(downstreamSessionKey, requestJSON) + toolCacheTurn.recordRequest(requestJSON) + requestJSON = repairResponsesWebsocketToolCallsWithoutRecording(downstreamSessionKey, requestJSON) requestJSON = dedupeResponsesWebsocketInputItemsByID(requestJSON) updatedLastRequest = bytes.Clone(requestJSON) lastRequest = updatedLastRequest @@ -712,6 +714,7 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { wsTimelineLog, passthroughSessionID, responsesWebsocketForwardOptions{ + toolCacheTurn: toolCacheTurn, suppressError: replayPinnedAuthFailure, }, ) @@ -744,6 +747,7 @@ func (h *OpenAIResponsesAPIHandler) ResponsesWebsocket(c *gin.Context) { continue } + toolCacheTurn.commit() upstreamMode = attemptedUpstreamMode if upstreamMode == responsesWebsocketUpstreamModeWS { upstreamWebsocketAuthID = lastAttemptedAuthID @@ -1685,6 +1689,7 @@ func normalizeJSONArrayRaw(raw []byte) string { } type responsesWebsocketForwardOptions struct { + toolCacheTurn *responsesWebsocketToolCacheTurn suppressError func(*interfaces.ErrorMessage) bool } @@ -1702,6 +1707,7 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( if len(options) > 0 { opts = options[0] } + toolCacheTurn := opts.toolCacheTurn completed := false completedOutput := []byte("[]") completedResponseID := "" @@ -1803,7 +1809,11 @@ func (h *OpenAIResponsesAPIHandler) forwardResponsesWebsocket( if isResponsesWebsocketCompletionEvent(eventType) { payloads[i] = restoreResponsesWebsocketCompletionOutput(payloads[i], outputItemsByIndex, outputItemsFallback) } - recordResponsesWebsocketToolCallsFromPayload(downstreamSessionKey, payloads[i]) + if toolCacheTurn != nil { + toolCacheTurn.recordResponse(payloads[i]) + } else { + recordResponsesWebsocketToolCallsFromPayload(downstreamSessionKey, payloads[i]) + } recordPendingToolCallIDsFromPayload(pendingToolCallIDs, payloads[i]) var payloadErrMsg *interfaces.ErrorMessage if eventType == wsEventTypeError { diff --git a/sdk/api/handlers/openai/openai_responses_websocket_test.go b/sdk/api/handlers/openai/openai_responses_websocket_test.go index f50f44695..99f426239 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket_test.go +++ b/sdk/api/handlers/openai/openai_responses_websocket_test.go @@ -1608,6 +1608,44 @@ func TestRepairResponsesWebsocketToolCallsInsertsCachedOutput(t *testing.T) { } } +func TestResponsesWebsocketToolCacheTurnCommitsOnlyOnSuccess(t *testing.T) { + const sessionKey = "tool-cache-turn-commit-session" + defer defaultWebsocketToolOutputCache.deleteSession(sessionKey) + defer defaultWebsocketToolCallCache.deleteSession(sessionKey) + + turn := newResponsesWebsocketToolCacheTurn(sessionKey) + turn.recordRequest([]byte(`{"input":[{"type":"function_call_output","id":"fco-1","call_id":"call-1","output":"cached result"}]}`)) + beforeCommit := repairResponsesWebsocketToolCallsWithoutRecording(sessionKey, []byte(`{"input":[{"type":"function_call","id":"fc-next","call_id":"call-1","name":"lookup","arguments":"{}"}]}`)) + if gjson.GetBytes(beforeCommit, "input.#").Int() != 0 { + t.Fatalf("uncommitted turn populated global cache: %s", beforeCommit) + } + + turn.commit() + afterCommit := repairResponsesWebsocketToolCallsWithoutRecording(sessionKey, []byte(`{"input":[{"type":"function_call","id":"fc-next","call_id":"call-1","name":"lookup","arguments":"{}"}]}`)) + input := gjson.GetBytes(afterCommit, "input").Array() + if len(input) != 2 || input[1].Get("output").String() != "cached result" { + t.Fatalf("committed turn was not available to tool repair: %s", afterCommit) + } +} + +func TestResponsesWebsocketToolCacheRetainPreventsOverlappingReleaseDeletion(t *testing.T) { + const sessionKey = "tool-cache-overlapping-retain-session" + retainResponsesWebsocketToolCaches(sessionKey) + retainResponsesWebsocketToolCaches(sessionKey) + turn := newResponsesWebsocketToolCacheTurn(sessionKey) + turn.recordRequest([]byte(`{"input":[{"type":"function_call_output","id":"fco-1","call_id":"call-1","output":"kept"}]}`)) + turn.commit() + + releaseResponsesWebsocketToolCaches(sessionKey) + if _, ok := defaultWebsocketToolOutputCache.get(sessionKey, "call-1"); !ok { + t.Fatal("first overlapping release deleted active session cache") + } + releaseResponsesWebsocketToolCaches(sessionKey) + if _, ok := defaultWebsocketToolOutputCache.get(sessionKey, "call-1"); ok { + t.Fatal("final release did not delete session cache") + } +} + func TestRepairResponsesWebsocketToolCallsDropsOrphanFunctionCall(t *testing.T) { cache := newWebsocketToolOutputCache(time.Minute, 10) sessionKey := "session-1" @@ -2933,6 +2971,86 @@ func TestResponsesWebsocketRejectsUnknownPreviousResponseOnNewSocket(t *testing. } } +func TestResponsesWebsocketRollsBackCanonicalTranscriptAfterNonRetryableError(t *testing.T) { + gin.SetMode(gin.TestMode) + + modelName := "xai-websocket-rollback-model" + executor := &websocketCanonicalRollbackExecutor{} + manager := coreauth.NewManager(nil, nil, nil) + manager.RegisterExecutor(executor) + auth := &coreauth.Auth{ + ID: "auth-xai-rollback", + Provider: "xai", + Status: coreauth.StatusActive, + } + if _, err := manager.Register(context.Background(), auth); err != nil { + t.Fatalf("Register auth: %v", err) + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: modelName}}) + t.Cleanup(func() { + registry.GetGlobalRegistry().UnregisterClient(auth.ID) + }) + + 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() + + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/v1/responses/ws" + conn, _, err := websocket.DefaultDialer.Dial(wsURL, http.Header{"Session-Id": []string{"rollback-tool-cache-session"}}) + if err != nil { + t.Fatalf("dial websocket: %v", err) + } + defer func() { _ = conn.Close() }() + + requests := []string{ + fmt.Sprintf(`{"type":"response.create","model":%q,"input":[{"type":"message","id":"msg-1"}]}`, modelName), + `{"type":"response.create","previous_response_id":"resp-1","input":[{"type":"function_call","id":"fc-failed","call_id":"failed-call","name":"failed_tool","arguments":"{}"},{"type":"function_call_output","id":"fco-failed","call_id":"failed-call","output":"must-not-survive"}]}`, + `{"type":"response.create","previous_response_id":"resp-1","input":[{"type":"message","id":"msg-3"},{"type":"function_call","id":"fc-retry","call_id":"failed-call","name":"failed_tool","arguments":"{}"}]}`, + } + wantTypes := []string{wsEventTypeCompleted, wsEventTypeError, wsEventTypeCompleted} + for i := range requests { + if errWrite := conn.WriteMessage(websocket.TextMessage, []byte(requests[i])); errWrite != nil { + t.Fatalf("write websocket message %d: %v", i+1, errWrite) + } + _, payload, errRead := conn.ReadMessage() + if errRead != nil { + t.Fatalf("read websocket response %d: %v", i+1, errRead) + } + if got := gjson.GetBytes(payload, "type").String(); got != wantTypes[i] { + t.Fatalf("response %d type = %q, want %q: %s", i+1, got, wantTypes[i], payload) + } + } + + payloads := executor.Payloads() + if len(payloads) != 3 { + t.Fatalf("executor payload count = %d, want 3", len(payloads)) + } + third := payloads[2] + if gjson.GetBytes(third, "previous_response_id").Exists() { + t.Fatalf("retry payload must not depend on previous_response_id: %s", third) + } + input := gjson.GetBytes(third, "input").Array() + if len(input) != 3 { + t.Fatalf("retry canonical input len = %d, want 3: %s", len(input), third) + } + wantIDs := []string{"msg-1", "out-1", "msg-3"} + for i, wantID := range wantIDs { + if got := input[i].Get("id").String(); got != wantID { + t.Fatalf("retry canonical input[%d].id = %q, want %q: %s", i, got, wantID, third) + } + } + if bytes.Contains(third, []byte(`"id":"fc-failed"`)) || bytes.Contains(third, []byte(`"id":"fco-failed"`)) { + t.Fatalf("failed turn leaked into retry transcript: %s", third) + } + if bytes.Contains(third, []byte(`"call_id":"failed-call"`)) || bytes.Contains(third, []byte("must-not-survive")) { + t.Fatalf("failed turn contaminated tool repair cache: %s", third) + } +} + func TestResponsesWebsocketSwitchesPinnedAuthAcrossProviders(t *testing.T) { for _, testCase := range []struct { name string diff --git a/sdk/api/handlers/openai/openai_responses_websocket_toolcall_repair.go b/sdk/api/handlers/openai/openai_responses_websocket_toolcall_repair.go index 0c3ab58d8..85260ad47 100644 --- a/sdk/api/handlers/openai/openai_responses_websocket_toolcall_repair.go +++ b/sdk/api/handlers/openai/openai_responses_websocket_toolcall_repair.go @@ -19,6 +19,7 @@ const ( var defaultWebsocketToolOutputCache = newWebsocketToolOutputCache(0, websocketToolOutputCacheMaxPerSession) var defaultWebsocketToolCallCache = newWebsocketToolOutputCache(0, websocketToolOutputCacheMaxPerSession) var defaultWebsocketToolSessionRefs = newWebsocketToolSessionRefCounter() +var defaultWebsocketToolCacheTransactionMu sync.RWMutex type websocketToolOutputCache struct { mu sync.Mutex @@ -33,6 +34,14 @@ type websocketToolOutputSession struct { order []string } +type responsesWebsocketToolCacheTurn struct { + sessionKey string + outputs map[string]json.RawMessage + outputOrder []string + calls map[string]json.RawMessage + callOrder []string +} + func newWebsocketToolOutputCache(ttl time.Duration, maxPerSession int) *websocketToolOutputCache { if ttl < 0 { ttl = websocketToolOutputCacheTTL @@ -196,6 +205,8 @@ func (c *websocketToolSessionRefCounter) release(sessionKey string) bool { } func retainResponsesWebsocketToolCaches(sessionKey string) { + defaultWebsocketToolCacheTransactionMu.Lock() + defer defaultWebsocketToolCacheTransactionMu.Unlock() if defaultWebsocketToolSessionRefs == nil { return } @@ -203,13 +214,14 @@ func retainResponsesWebsocketToolCaches(sessionKey string) { } func releaseResponsesWebsocketToolCaches(sessionKey string) { + defaultWebsocketToolCacheTransactionMu.Lock() + defer defaultWebsocketToolCacheTransactionMu.Unlock() if defaultWebsocketToolSessionRefs == nil { return } if !defaultWebsocketToolSessionRefs.release(sessionKey) { return } - if defaultWebsocketToolOutputCache != nil { defaultWebsocketToolOutputCache.deleteSession(sessionKey) } @@ -218,15 +230,114 @@ func releaseResponsesWebsocketToolCaches(sessionKey string) { } } +func newResponsesWebsocketToolCacheTurn(sessionKey string) *responsesWebsocketToolCacheTurn { + sessionKey = strings.TrimSpace(sessionKey) + if sessionKey == "" { + return nil + } + return &responsesWebsocketToolCacheTurn{ + sessionKey: sessionKey, + outputs: make(map[string]json.RawMessage), + calls: make(map[string]json.RawMessage), + } +} + +func (t *responsesWebsocketToolCacheTurn) recordRequest(payload []byte) { + if t == nil || len(payload) == 0 { + return + } + input := gjson.GetBytes(payload, "input") + if !input.Exists() || !input.IsArray() { + return + } + for _, item := range input.Array() { + t.recordItem(item) + } +} + +func (t *responsesWebsocketToolCacheTurn) recordResponse(payload []byte) { + if t == nil || len(payload) == 0 { + return + } + switch strings.TrimSpace(gjson.GetBytes(payload, "type").String()) { + case "response.completed": + output := gjson.GetBytes(payload, "response.output") + if !output.Exists() || !output.IsArray() { + return + } + for _, item := range output.Array() { + if isCompleteResponsesWebsocketToolCall(item) { + t.recordItem(item) + } + } + case "response.output_item.added", "response.output_item.done": + item := gjson.GetBytes(payload, "item") + if isCompleteResponsesWebsocketToolCall(item) { + t.recordItem(item) + } + } +} + +func (t *responsesWebsocketToolCacheTurn) recordItem(item gjson.Result) { + if t == nil || !item.Exists() { + return + } + callID := strings.TrimSpace(item.Get("call_id").String()) + if callID == "" || strings.TrimSpace(item.Raw) == "" { + return + } + raw := append(json.RawMessage(nil), item.Raw...) + switch { + case isResponsesToolCallOutputType(item.Get("type").String()): + if _, exists := t.outputs[callID]; !exists { + t.outputOrder = append(t.outputOrder, callID) + } + t.outputs[callID] = raw + case isResponsesToolCallType(item.Get("type").String()): + if _, exists := t.calls[callID]; !exists { + t.callOrder = append(t.callOrder, callID) + } + t.calls[callID] = raw + } +} + +func (t *responsesWebsocketToolCacheTurn) commit() { + if t == nil || t.sessionKey == "" { + return + } + defaultWebsocketToolCacheTransactionMu.Lock() + defer defaultWebsocketToolCacheTransactionMu.Unlock() + if defaultWebsocketToolOutputCache != nil { + for _, callID := range t.outputOrder { + defaultWebsocketToolOutputCache.record(t.sessionKey, callID, t.outputs[callID]) + } + } + if defaultWebsocketToolCallCache != nil { + for _, callID := range t.callOrder { + defaultWebsocketToolCallCache.record(t.sessionKey, callID, t.calls[callID]) + } + } +} + func repairResponsesWebsocketToolCalls(sessionKey string, payload []byte) []byte { return repairResponsesWebsocketToolCallsWithCaches(defaultWebsocketToolOutputCache, defaultWebsocketToolCallCache, sessionKey, payload) } +func repairResponsesWebsocketToolCallsWithoutRecording(sessionKey string, payload []byte) []byte { + defaultWebsocketToolCacheTransactionMu.RLock() + defer defaultWebsocketToolCacheTransactionMu.RUnlock() + return repairResponsesWebsocketToolCallsWithCachesMode(defaultWebsocketToolOutputCache, defaultWebsocketToolCallCache, sessionKey, payload, false) +} + func repairResponsesWebsocketToolCallsWithCache(cache *websocketToolOutputCache, sessionKey string, payload []byte) []byte { return repairResponsesWebsocketToolCallsWithCaches(cache, nil, sessionKey, payload) } func repairResponsesWebsocketToolCallsWithCaches(outputCache, callCache *websocketToolOutputCache, sessionKey string, payload []byte) []byte { + return repairResponsesWebsocketToolCallsWithCachesMode(outputCache, callCache, sessionKey, payload, true) +} + +func repairResponsesWebsocketToolCallsWithCachesMode(outputCache, callCache *websocketToolOutputCache, sessionKey string, payload []byte, record bool) []byte { sessionKey = strings.TrimSpace(sessionKey) if sessionKey == "" || outputCache == nil || len(payload) == 0 { return payload @@ -238,7 +349,7 @@ func repairResponsesWebsocketToolCallsWithCaches(outputCache, callCache *websock } allowOrphanOutputs := strings.TrimSpace(gjson.GetBytes(payload, "previous_response_id").String()) != "" - updatedRaw, errRepair := repairResponsesToolCallsArray(outputCache, callCache, sessionKey, input.Raw, allowOrphanOutputs) + updatedRaw, errRepair := repairResponsesToolCallsArray(outputCache, callCache, sessionKey, input.Raw, allowOrphanOutputs, record) if errRepair != nil || updatedRaw == "" || updatedRaw == input.Raw { return payload } @@ -250,7 +361,7 @@ func repairResponsesWebsocketToolCallsWithCaches(outputCache, callCache *websock return updated } -func repairResponsesToolCallsArray(outputCache, callCache *websocketToolOutputCache, sessionKey string, rawArray string, allowOrphanOutputs bool) (string, error) { +func repairResponsesToolCallsArray(outputCache, callCache *websocketToolOutputCache, sessionKey string, rawArray string, allowOrphanOutputs bool, record bool) (string, error) { rawArray = strings.TrimSpace(rawArray) if rawArray == "" { return "[]", nil @@ -276,14 +387,16 @@ func repairResponsesToolCallsArray(outputCache, callCache *websocketToolOutputCa continue } outputPresent[callID] = struct{}{} - outputCache.record(sessionKey, callID, item) + if record { + outputCache.record(sessionKey, callID, item) + } case isResponsesToolCallType(itemType): callID := strings.TrimSpace(gjson.GetBytes(item, "call_id").String()) if callID == "" { continue } callPresent[callID] = struct{}{} - if callCache != nil { + if record && callCache != nil { callCache.record(sessionKey, callID, item) } }