diff --git a/internal/runtime/executor/codex_executor.go b/internal/runtime/executor/codex_executor.go index 4847fbf0..457a210b 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 266612cf..a63eea63 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 2f3103ba..6c9a611b 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 967172e3..1f04538e 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 6df8ade6..75ceae3c 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 a05b1e41..1eea484f 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 f631f1d4..4bd93a2c 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 11ffcb0c..a0aa52c9 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 4f093f47..f50f4469 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 367b507e..c18d3f68 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 00000000..1fa0d79e --- /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 00000000..f4327fb6 --- /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) + } +}