From 4b9d404fb04f3445a1435e21027257687b89f274 Mon Sep 17 00:00:00 2001 From: sususu98 <33882693+sususu98@users.noreply.github.com> Date: Thu, 20 Aug 2026 21:31:53 +0800 Subject: [PATCH] feat(codex): add opt-in stream bootstrap buffering and overload failover (#5115) * feat(codex): add opt-in stream bootstrap buffering The upstream smuggles capacity rejections into an HTTP 200 stream. The handshake events arrive normally and only a later event carries {"error":{"type":"service_unavailable_error","code": "server_is_overloaded"}}. By then the executor has already handed the first chunk downstream, the response is committed, and the conductor can no longer retry on another credential, so the request fails even though other credentials were available. When codex.stream-bootstrap-buffering is enabled the executor holds back the handshake events until it can tell whether the stream carries real output or a rejection. An overload rejection then fails the attempt before any chunk is delivered, letting the conductor retry on another credential; every other terminal failure is flushed in order and delivered in-stream exactly as before. Detection uses an event-type allow-list rather than a fixed count. On the websocket transport codex.rate_limits and codex.response.metadata arrive before response.created, making the first generated event the fifth frame, so a small counter would release the stream before the rejection is visible. Buffering is bounded and hitting the bound degrades to the original unbuffered behaviour. Two details are load-bearing. The error must be returned synchronously: delivering it as the first stream chunk makes ExecuteStream downgrade it into a committed 200 and the status is lost. And the websocket path must not signal an upstream disconnect for a rejection it intends to retry, because the downstream handler closes the client connection on that signal and the retry would have nowhere to deliver. The 503 status is produced only on this path rather than in the shared codexTerminalFailureStatus mapping, so disabling the feature restores the previous behaviour exactly, including cooldown classification and retry-after parsing. Defaults to false: response headers are withheld until generation starts, which can trip client or reverse-proxy read timeouts. * test(codex): pin bootstrap overload failover through the conductor Executor-level tests cannot show what the client finally receives. These exercise ExecuteStream end to end to pin three properties that are easy to regress: - consecutive overloaded credentials are skipped until one serves the request, and retries are capped by max-retry-credentials rather than multiplying with request-retry - exhausting the pool surfaces the upstream status instead of a committed 200 stream - with buffering disabled the rejection stays an in-stream error on a committed stream, which is the behaviour the feature must preserve The third case also documents why the executor returns its error synchronously: an error arriving as the first stream chunk is wrapped and downgraded into a committed 200, silently losing the status. --- config.example.yaml | 11 + internal/config/config_types.go | 8 + .../runtime/executor/codex_executor_stream.go | 151 +++++- .../codex_executor_stream_output_test.go | 7 + .../executor/codex_executor_terminal.go | 48 ++ .../codex_stream_bootstrap_buffering_test.go | 478 ++++++++++++++++++ .../executor/codex_websockets_stream.go | 206 +++++++- internal/watcher/diff/config_diff.go | 3 + ...conductor_stream_overload_failover_test.go | 168 ++++++ .../conductor_stream_overload_status_test.go | 138 +++++ 10 files changed, 1206 insertions(+), 12 deletions(-) create mode 100644 internal/runtime/executor/codex_stream_bootstrap_buffering_test.go create mode 100644 sdk/cliproxy/auth/conductor_stream_overload_failover_test.go create mode 100644 sdk/cliproxy/auth/conductor_stream_overload_status_test.go diff --git a/config.example.yaml b/config.example.yaml index f23526482..18bc0ca3b 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -228,6 +228,17 @@ codex: identity-confuse: false # Disable forcing the official Codex User-Agent and Originator headers on HTTP/SSE and WebSocket requests. disable-codex-cloaking: false + # Hold back the initial handshake events (response.created, response.in_progress and the + # websocket metadata frames) until the upstream emits its first generated event. + # Why: the upstream smuggles `server_is_overloaded` rejections *inside* an HTTP 200 stream, + # right after those handshake events, instead of returning 503 on the wire. Buffering them + # keeps the downstream response headers uncommitted long enough to transparently retry on + # another credential. Only overload/rate-limit rejections trigger failover; every other + # terminal failure is still delivered in-stream exactly as before. + # Trade-off: response headers are delayed until generation starts, which can trip client or + # reverse-proxy read timeouts (e.g. nginx proxy_read_timeout) on long reasoning requests. + # Default: false + stream-bootstrap-buffering: false # When true, optimize Codex Desktop, codex-tui, and codex_cli_rs requests for multi-agent v2. # This refreshes Codex spawn_agent model details, removes message parameter encryption, # normalizes encrypted agent_message content for Codex, and converts agent_message input diff --git a/internal/config/config_types.go b/internal/config/config_types.go index 783f666f7..50c468f6e 100644 --- a/internal/config/config_types.go +++ b/internal/config/config_types.go @@ -148,6 +148,14 @@ type CodexConfig struct { IdentityConfuse bool `yaml:"identity-confuse" json:"identity-confuse"` // DisableCodexCloaking disables forcing the official Codex identity headers on HTTP/SSE and WebSocket requests. DisableCodexCloaking bool `yaml:"disable-codex-cloaking" json:"disable-codex-cloaking"` + // StreamBootstrapBuffering holds back initial handshake events (response.created, + // response.in_progress and the websocket metadata frames) until the first generated event + // arrives. The upstream delivers server_is_overloaded rejections inside an HTTP 200 stream + // right after those handshake events instead of returning 503 on the wire, so buffering them + // keeps the downstream response headers uncommitted long enough to retry on another credential. + // Trade-off: the response headers are delayed until the upstream starts generating, which can + // trip client or reverse-proxy read timeouts. Default is false. + StreamBootstrapBuffering bool `yaml:"stream-bootstrap-buffering" json:"stream-bootstrap-buffering"` // OptimizeMultiAgentV2 optimizes official Codex multi-agent requests. OptimizeMultiAgentV2 bool `yaml:"optimize-multi-agent-v2" json:"optimize-multi-agent-v2"` // LiveMediaRelay terminates and relays Codex Live WebRTC media in this process. diff --git a/internal/runtime/executor/codex_executor_stream.go b/internal/runtime/executor/codex_executor_stream.go index 00dc46ad4..d6e4918f3 100644 --- a/internal/runtime/executor/codex_executor_stream.go +++ b/internal/runtime/executor/codex_executor_stream.go @@ -130,7 +130,150 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au err = newCodexStatusErr(httpResp.StatusCode, data) return nil, err } - out := make(chan cliproxyexecutor.StreamChunk) + + buffering := e.cfg != nil && e.cfg.Codex.StreamBootstrapBuffering + + scanner := bufio.NewScanner(httpResp.Body) + scanner.Buffer(nil, 52_428_800) // 50MB + claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload) + var param any + outputItemsByIndex := make(map[int64][]byte) + var outputItemsFallback [][]byte + + var bufferedChunks [][]byte + var initialChunks [][]byte + streamStarted := false + immediateTerminal := false + // bootstrapTerminalErr holds a non-overload terminal failure seen while buffering. It is + // delivered as an in-stream chunk after the buffered handshake so downstream behaviour stays + // identical to the unbuffered path instead of silently turning into a credential failover. + var bootstrapTerminalErr error + + closeBootstrapBody := func() { + if errClose := httpResp.Body.Close(); errClose != nil { + log.Errorf("codex executor: close response body error: %v", errClose) + } + } + + if buffering { + for scanner.Scan() { + line := applyCodexIdentityConfuseResponsePayload(scanner.Bytes(), identityState) + helps.AppendAPIResponseChunk(ctx, e.cfg, line) + translatedLine := bytes.Clone(line) + isHandshake := false + terminalSuccess := false + + if bytes.HasPrefix(line, dataTag) { + data := bytes.TrimSpace(line[5:]) + data = helps.RestoreCodexMultiAgentV2Response(data, optimizeMultiAgentV2) + translatedLine = append([]byte("data: "), data...) + eventType := gjson.GetBytes(data, "type").String() + if streamErr, terminalBody, ok := codexTerminalFailureErr(data); ok { + closeBootstrapBody() + if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, streamErr.StatusCode(), terminalBody); errClearReplay != nil { + helps.RecordAPIResponseError(ctx, e.cfg, errClearReplay) + reporter.PublishFailure(ctx, errClearReplay) + return nil, errClearReplay + } + helps.RecordAPIResponseError(ctx, e.cfg, streamErr) + reporter.PublishFailure(ctx, streamErr) + if isCodexOverloadBootstrapFailure(terminalBody) { + // Transient capacity rejection smuggled into an HTTP 200 stream. Fail the + // attempt before the downstream headers are committed so the conductor can + // transparently retry on another credential, and report the status the + // upstream refused to put on the wire. + helps.LogWithRequestID(ctx).Debugf("codex executor: bootstrap overload rejection after %d buffered handshake events, failing over", len(bufferedChunks)) + return nil, newCodexBootstrapOverloadErr(terminalBody) + } + bootstrapTerminalErr = streamErr + break + } + if isCodexHandshakeMetadataEvent(eventType) { + isHandshake = true + } + switch eventType { + case "response.output_item.done": + collectCodexOutputItemDone(data, outputItemsByIndex, &outputItemsFallback) + case "response.completed", "response.incomplete": + terminalSuccess = true + if detail, ok := helps.ParseCodexUsage(data); ok { + reporter.Publish(ctx, detail) + } + publishCodexImageToolUsage(ctx, reporter, body, data) + data = patchCodexCompletedOutput(data, outputItemsByIndex, outputItemsFallback) + if eventType == "response.completed" { + cacheCodexReasoningReplayFromCompleted(replayScope, data) + } + translatedLine = append([]byte("data: "), data...) + } + } else { + isHandshake = true + } + + translatedLine = applyCodexIdentityExposeResponsePayload(translatedLine, identityState) + chunks := helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, originalPayload, body, translatedLine, ¶m, claudeInputTokens) + if isHandshake && !terminalSuccess { + if len(bufferedChunks) < codexBootstrapMaxBufferedEvents { + bufferedChunks = append(bufferedChunks, chunks...) + continue + } + helps.LogWithRequestID(ctx).Debugf("codex executor: bootstrap buffer limit %d reached, releasing stream without overload probing", codexBootstrapMaxBufferedEvents) + } + + initialChunks = chunks + streamStarted = true + if terminalSuccess { + immediateTerminal = true + } + break + } + + if !streamStarted && bootstrapTerminalErr == nil { + closeBootstrapBody() + if errScan := scanner.Err(); errScan != nil { + // A cancelled downstream request must not be recorded as an upstream failure or + // penalise the credential; mirror the unbuffered goroutine's guard. + if ctx.Err() != nil { + return nil, ctx.Err() + } + helps.RecordAPIResponseError(ctx, e.cfg, errScan) + reporter.PublishFailure(ctx, errScan) + return nil, errScan + } + if ctx.Err() != nil { + return nil, ctx.Err() + } + streamErr := newCodexIncompleteStreamError() + helps.RecordAPIResponseError(ctx, e.cfg, streamErr) + reporter.PublishFailure(ctx, streamErr) + return nil, streamErr + } + } + + chanCapacity := len(bufferedChunks) + len(initialChunks) + if bootstrapTerminalErr != nil { + chanCapacity++ + } + out := make(chan cliproxyexecutor.StreamChunk, chanCapacity) + for _, chunk := range bufferedChunks { + out <- cliproxyexecutor.StreamChunk{Payload: chunk} + } + for _, chunk := range initialChunks { + out <- cliproxyexecutor.StreamChunk{Payload: chunk} + } + if bootstrapTerminalErr != nil { + // Buffered handshake payloads are flushed first so the conductor observes a committed + // stream and delivers this failure in-stream, exactly as the unbuffered path would. + out <- cliproxyexecutor.StreamChunk{Err: bootstrapTerminalErr} + close(out) + return &cliproxyexecutor.StreamResult{Headers: httpResp.Header.Clone(), Chunks: out}, nil + } + if immediateTerminal { + closeBootstrapBody() + close(out) + return &cliproxyexecutor.StreamResult{Headers: httpResp.Header.Clone(), Chunks: out}, nil + } + go func() { defer close(out) defer func() { @@ -138,12 +281,6 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au log.Errorf("codex executor: close response body error: %v", errClose) } }() - scanner := bufio.NewScanner(httpResp.Body) - scanner.Buffer(nil, 52_428_800) // 50MB - claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload) - var param any - outputItemsByIndex := make(map[int64][]byte) - var outputItemsFallback [][]byte for scanner.Scan() { line := applyCodexIdentityConfuseResponsePayload(scanner.Bytes(), identityState) helps.AppendAPIResponseChunk(ctx, e.cfg, line) diff --git a/internal/runtime/executor/codex_executor_stream_output_test.go b/internal/runtime/executor/codex_executor_stream_output_test.go index 3d1fbcf51..f40ef038d 100644 --- a/internal/runtime/executor/codex_executor_stream_output_test.go +++ b/internal/runtime/executor/codex_executor_stream_output_test.go @@ -576,6 +576,13 @@ func TestCodexTerminalFailureErrClassifiesStatus(t *testing.T) { event: `{"type":"response.failed","response":{"error":{"type":"upstream_error","code":"unknown","message":"Upstream failed."}}}`, wantStatus: http.StatusBadGateway, }, + // Overload rejections keep falling through to 502 here. The 503 restoration is scoped to + // the opt-in bootstrap buffering path so this shared mapping stays unchanged. + { + name: "overload stays a bad gateway without buffering", + event: `{"type":"error","error":{"type":"service_unavailable_error","code":"server_is_overloaded","message":"Our servers are currently overloaded. Please try again later."}}`, + wantStatus: http.StatusBadGateway, + }, } for _, tc := range tests { diff --git a/internal/runtime/executor/codex_executor_terminal.go b/internal/runtime/executor/codex_executor_terminal.go index f05727dc0..be69833cf 100644 --- a/internal/runtime/executor/codex_executor_terminal.go +++ b/internal/runtime/executor/codex_executor_terminal.go @@ -405,3 +405,51 @@ func parseCodexRetryAfter(statusCode int, errorBody []byte, now time.Time) *time } return nil } + +// codexBootstrapMaxBufferedEvents bounds how many handshake metadata events may be held +// back while probing for an upstream rejection embedded in an HTTP 200 stream. The websocket +// transport prefixes response events with codex.response.metadata and codex.rate_limits frames, +// so the limit must comfortably exceed the four handshake frames observed in practice. Once the +// limit is reached the stream is released and the original unbuffered semantics apply. +const codexBootstrapMaxBufferedEvents = 16 + +// isCodexHandshakeMetadataEvent reports whether an event carries no generated output and is +// therefore safe to hold back before the downstream response headers are committed. Keeping a type +// allow-list rather than a fixed event count matters for the websocket transport, where the +// handshake frames arrive before response.created and would otherwise exhaust a small counter +// before the rejection event is seen. +func isCodexHandshakeMetadataEvent(eventType string) bool { + switch eventType { + case "response.created", "response.in_progress", "codex.rate_limits", "codex.response.metadata": + return true + default: + return false + } +} + +// newCodexBootstrapOverloadErr reports a buffered overload rejection with its real status. +// +// The status is deliberately produced here instead of in codexTerminalFailureStatus: that mapping +// is shared with the unbuffered path, where the rejection is delivered in-stream and a status +// change would alter cooldown classification and retry-after parsing for everyone. Keeping 503 +// scoped to this path means disabling the feature restores the previous behaviour exactly. +func newCodexBootstrapOverloadErr(body []byte) statusErr { + return newCodexStatusErr(http.StatusServiceUnavailable, body) +} + +// isCodexOverloadBootstrapFailure reports whether a terminal failure delivered inside an HTTP 200 +// stream is a transient capacity rejection that a different credential may be able to serve. +// Only these failures justify replacing the whole attempt during bootstrap; every other terminal +// failure keeps the original in-stream delivery semantics so downstream behaviour is unchanged. +func isCodexOverloadBootstrapFailure(body []byte) bool { + errorType := strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "error.type").String())) + errorCode := strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "error.code").String())) + switch { + case errorType == "service_unavailable_error", errorCode == "server_is_overloaded": + return true + case errorType == "rate_limit_error", errorCode == "rate_limit_exceeded": + return true + default: + return false + } +} diff --git a/internal/runtime/executor/codex_stream_bootstrap_buffering_test.go b/internal/runtime/executor/codex_stream_bootstrap_buffering_test.go new file mode 100644 index 000000000..2a4e6e226 --- /dev/null +++ b/internal/runtime/executor/codex_stream_bootstrap_buffering_test.go @@ -0,0 +1,478 @@ +package executor + +import ( + "bytes" + "context" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/gorilla/websocket" + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + _ "github.com/router-for-me/CLIProxyAPI/v7/internal/translator" + cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" + sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" +) + +const ( + codexOverloadEvent = `{"type":"error","error":{"type":"service_unavailable_error","code":"server_is_overloaded","message":"Our servers are currently overloaded. Please try again later.","param":null},"sequence_number":2}` + codexInvalidEvent = `{"type":"error","error":{"type":"invalid_request_error","code":"invalid_value","message":"Invalid input."},"sequence_number":2}` + codexCreatedEvent = `{"type":"response.created","response":{"id":"resp_1","model":"gpt-5.6-terra"}}` + codexInProgressEvent = `{"type":"response.in_progress","response":{"id":"resp_1"}}` + codexOutputAddedEvent = `{"type":"response.output_item.added","item":{"id":"msg_1","type":"message","role":"assistant","content":[]},"output_index":0}` + codexCompletedEventBody = `{"type":"response.completed","response":{"id":"resp_1","status":"completed","output":[{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"output_text","text":"hello"}]}],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}` +) + +func codexBufferingConfig(enabled bool) *config.Config { + return &config.Config{Codex: config.CodexConfig{StreamBootstrapBuffering: enabled}} +} + +func codexTestAuth(baseURL string) *cliproxyauth.Auth { + return &cliproxyauth.Auth{Attributes: map[string]string{"base_url": baseURL, "api_key": "test"}} +} + +func codexTestRequest() (cliproxyexecutor.Request, cliproxyexecutor.Options) { + return cliproxyexecutor.Request{ + Model: "gpt-5.6-terra", + Payload: []byte(`{"model":"gpt-5.6-terra","input":"hello"}`), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai-response"), + Stream: true, + } +} + +// codexSSEServer streams the supplied event payloads as an HTTP 200 SSE response. +func codexSSEServer(events ...string) *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + for _, event := range events { + eventType := "message" + if parsed := strings.SplitN(event, `"type":"`, 2); len(parsed) == 2 { + eventType = strings.SplitN(parsed[1], `"`, 2)[0] + } + _, _ = w.Write([]byte("event: " + eventType + "\n")) + _, _ = w.Write([]byte("data: " + event + "\n\n")) + } + })) +} + +// codexWebsocketServer echoes the supplied frames after receiving the client request frame. +func codexWebsocketServer(t *testing.T, frames ...string) *httptest.Server { + t.Helper() + upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade websocket: %v", err) + return + } + defer func() { _ = conn.Close() }() + if _, _, errRead := conn.ReadMessage(); errRead != nil { + t.Errorf("read websocket message: %v", errRead) + return + } + for _, frame := range frames { + _ = conn.WriteMessage(websocket.TextMessage, []byte(frame)) + } + })) +} + +func codexWebsocketRequest() (cliproxyexecutor.Request, cliproxyexecutor.Options) { + return cliproxyexecutor.Request{ + Model: "gpt-5.6-terra", + Payload: []byte(`{"model":"gpt-5.6-terra","input":[{"type":"message","role":"user","content":"hello"}]}`), + }, cliproxyexecutor.Options{ + SourceFormat: sdktranslator.FromString("openai-response"), + } +} + +// drainChunks collects every payload and the first error from a stream result. +func drainChunks(result *cliproxyexecutor.StreamResult) (string, error) { + var payloads [][]byte + var streamErr error + for chunk := range result.Chunks { + if chunk.Err != nil { + if streamErr == nil { + streamErr = chunk.Err + } + continue + } + payloads = append(payloads, chunk.Payload) + } + return string(bytes.Join(payloads, []byte("\n"))), streamErr +} + +// An overload rejection smuggled into an HTTP 200 stream must fail the whole attempt before any +// downstream chunk escapes, so the conductor can retry on another credential. A nil StreamResult +// is the invariant: with no channel there is no way for the buffered handshake to reach the client. +func TestCodexExecutor_BootstrapBuffering_OverloadFailsAttemptWithoutLeakingHandshake(t *testing.T) { + server := codexSSEServer(codexCreatedEvent, codexInProgressEvent, codexOverloadEvent) + defer server.Close() + + req, opts := codexTestRequest() + result, err := NewCodexExecutor(codexBufferingConfig(true)).ExecuteStream(context.Background(), codexTestAuth(server.URL), req, opts) + + if err == nil { + t.Fatal("expected ExecuteStream to fail the attempt on an overload rejection") + } + if result != nil { + t.Fatal("expected nil result so no buffered handshake chunk can reach the client") + } + if got := statusCodeFromTestError(t, err); got != http.StatusServiceUnavailable { + t.Fatalf("status code = %d, want %d (upstream hides 503 behind HTTP 200)", got, http.StatusServiceUnavailable) + } +} + +// A non-overload terminal failure must keep the original in-stream delivery semantics: the +// buffered handshake is flushed first and the error arrives as a stream chunk, so the conductor +// sees a committed stream and does not burn another credential on a request-level fault. +func TestCodexExecutor_BootstrapBuffering_NonOverloadStaysInStream(t *testing.T) { + server := codexSSEServer(codexCreatedEvent, codexInProgressEvent, codexInvalidEvent) + defer server.Close() + + req, opts := codexTestRequest() + result, err := NewCodexExecutor(codexBufferingConfig(true)).ExecuteStream(context.Background(), codexTestAuth(server.URL), req, opts) + + if err != nil { + t.Fatalf("non-overload failure must not fail the attempt synchronously: %v", err) + } + if result == nil { + t.Fatal("expected a stream result for in-stream error delivery") + } + combined, streamErr := drainChunks(result) + if streamErr == nil { + t.Fatal("expected the invalid-request failure to arrive as an in-stream chunk error") + } + if !strings.Contains(combined, "response.created") { + t.Fatalf("buffered handshake must be flushed before the in-stream error: %s", combined) + } + if got := statusCodeFromTestError(t, streamErr); got != http.StatusBadRequest { + t.Fatalf("status code = %d, want %d", got, http.StatusBadRequest) + } +} + +// Once the buffer limit is exceeded the stream is released and overload probing stops, which +// bounds how long the downstream response headers can stay uncommitted. +func TestCodexExecutor_BootstrapBuffering_BufferLimitReleasesStream(t *testing.T) { + events := make([]string, 0, codexBootstrapMaxBufferedEvents+2) + for i := 0; i < codexBootstrapMaxBufferedEvents+1; i++ { + events = append(events, fmt.Sprintf(`{"type":"response.in_progress","response":{"id":"resp_%d"}}`, i)) + } + events = append(events, codexOverloadEvent) + server := codexSSEServer(events...) + defer server.Close() + + req, opts := codexTestRequest() + result, err := NewCodexExecutor(codexBufferingConfig(true)).ExecuteStream(context.Background(), codexTestAuth(server.URL), req, opts) + + if err != nil { + t.Fatalf("expected the stream to be released once the buffer limit is hit: %v", err) + } + if result == nil { + t.Fatal("expected a stream result after the buffer limit released the stream") + } + _, streamErr := drainChunks(result) + if streamErr == nil { + t.Fatal("expected the overload error to be delivered in-stream after the limit was hit") + } +} + +// Buffered handshake events must be replayed in upstream order ahead of the first generated event. +func TestCodexExecutor_BootstrapBuffering_FlushesInOrderOnFirstOutput(t *testing.T) { + server := codexSSEServer(codexCreatedEvent, codexInProgressEvent, codexOutputAddedEvent, codexCompletedEventBody) + defer server.Close() + + req, opts := codexTestRequest() + result, err := NewCodexExecutor(codexBufferingConfig(true)).ExecuteStream(context.Background(), codexTestAuth(server.URL), req, opts) + if err != nil { + t.Fatalf("unexpected ExecuteStream error: %v", err) + } + + combined, streamErr := drainChunks(result) + if streamErr != nil { + t.Fatalf("unexpected chunk error: %v", streamErr) + } + createdAt := strings.Index(combined, "response.created") + addedAt := strings.Index(combined, "response.output_item.added") + if createdAt < 0 || addedAt < 0 { + t.Fatalf("missing handshake or first generated event: %s", combined) + } + if createdAt > addedAt { + t.Fatalf("buffered handshake must be replayed before the first generated event: %s", combined) + } +} + +// With the feature disabled the overload rejection keeps its legacy in-stream delivery. +func TestCodexExecutor_BootstrapBuffering_DefaultDisabledPassthrough(t *testing.T) { + server := codexSSEServer(codexCreatedEvent, codexOverloadEvent) + defer server.Close() + + req, opts := codexTestRequest() + result, err := NewCodexExecutor(&config.Config{}).ExecuteStream(context.Background(), codexTestAuth(server.URL), req, opts) + if err != nil { + t.Fatalf("default unbuffered ExecuteStream returned error at call time: %v", err) + } + if result == nil { + t.Fatal("expected non-nil result in default unbuffered mode") + } + _, streamErr := drainChunks(result) + if streamErr == nil { + t.Fatal("expected stream error in chunks for default unbuffered mode") + } + // Disabling the feature must restore the previous behaviour exactly, status classification + // included: the 503 restoration is scoped to the buffered failover path, so an unbuffered + // overload still classifies as a bad gateway and keeps its old cooldown treatment. + if got := statusCodeFromTestError(t, streamErr); got != http.StatusBadGateway { + t.Fatalf("status code = %d, want %d while buffering is disabled", got, http.StatusBadGateway) + } +} + +// A cancelled downstream request must surface the context error rather than being recorded as an +// upstream failure that penalises the credential. +func TestCodexExecutor_BootstrapBuffering_ContextCancelDuringBootstrap(t *testing.T) { + server := codexSSEServer(codexCreatedEvent) + defer server.Close() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + req, opts := codexTestRequest() + _, err := NewCodexExecutor(codexBufferingConfig(true)).ExecuteStream(ctx, codexTestAuth(server.URL), req, opts) + if err == nil { + t.Fatal("expected an error for a cancelled bootstrap") + } + if !strings.Contains(err.Error(), context.Canceled.Error()) { + t.Fatalf("expected the context cancellation to surface, got: %v", err) + } +} + +func TestCodexWebsocketsExecutor_BootstrapBuffering_OverloadFailsAttempt(t *testing.T) { + server := codexWebsocketServer(t, codexCreatedEvent, codexInProgressEvent, codexOverloadEvent) + defer server.Close() + + req, opts := codexWebsocketRequest() + result, err := NewCodexWebsocketsExecutor(codexBufferingConfig(true)).ExecuteStream(context.Background(), codexTestAuth(server.URL), req, opts) + + if err == nil { + t.Fatal("expected ExecuteStream to fail the attempt on a websocket overload rejection") + } + if result != nil { + t.Fatal("expected nil result so no buffered handshake frame can reach the client") + } + if got := statusCodeFromTestError(t, err); got != http.StatusServiceUnavailable { + t.Fatalf("status code = %d, want %d", got, http.StatusServiceUnavailable) + } +} + +// The websocket transport prefixes response events with private metadata frames. Frame order +// below matches live wire capture: codex.rate_limits and codex.response.metadata both arrive +// *before* response.created, making the first generated event the fifth frame. They must be +// treated as handshake events, otherwise a fixed 3-event window would release the stream at +// response.created and never observe the rejection. +func TestCodexWebsocketsExecutor_BootstrapBuffering_PrivateHandshakeFramesDoNotExhaustWindow(t *testing.T) { + server := codexWebsocketServer(t, + `{"type":"codex.rate_limits","rate_limits":{"primary":{"used_percent":1}}}`, + `{"type":"codex.response.metadata","metadata":{"conversation_id":"conv_1"}}`, + codexCreatedEvent, + codexInProgressEvent, + codexOverloadEvent, + ) + defer server.Close() + + req, opts := codexWebsocketRequest() + result, err := NewCodexWebsocketsExecutor(codexBufferingConfig(true)).ExecuteStream(context.Background(), codexTestAuth(server.URL), req, opts) + + if err == nil { + t.Fatal("expected the overload rejection to be caught past the private handshake frames") + } + if result != nil { + t.Fatal("expected nil result so no buffered frame can reach the client") + } + if got := statusCodeFromTestError(t, err); got != http.StatusServiceUnavailable { + t.Fatalf("status code = %d, want %d", got, http.StatusServiceUnavailable) + } +} + +func TestCodexWebsocketsExecutor_BootstrapBuffering_NonOverloadStaysInStream(t *testing.T) { + server := codexWebsocketServer(t, codexCreatedEvent, codexInProgressEvent, codexInvalidEvent) + defer server.Close() + + req, opts := codexWebsocketRequest() + result, err := NewCodexWebsocketsExecutor(codexBufferingConfig(true)).ExecuteStream(context.Background(), codexTestAuth(server.URL), req, opts) + + if err != nil { + t.Fatalf("non-overload failure must not fail the attempt synchronously: %v", err) + } + if result == nil { + t.Fatal("expected a stream result for in-stream error delivery") + } + combined, streamErr := drainChunks(result) + if streamErr == nil { + t.Fatal("expected the invalid-request failure to arrive as an in-stream chunk error") + } + if !strings.Contains(combined, "response.created") { + t.Fatalf("buffered handshake must be flushed before the in-stream error: %s", combined) + } +} + +func TestCodexWebsocketsExecutor_BootstrapBuffering_FlushesInOrderOnFirstOutput(t *testing.T) { + server := codexWebsocketServer(t, + codexCreatedEvent, + codexInProgressEvent, + codexOutputAddedEvent, + `{"type":"response.completed","response":{"id":"resp_1","output":[],"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`, + ) + defer server.Close() + + req, opts := codexWebsocketRequest() + result, err := NewCodexWebsocketsExecutor(codexBufferingConfig(true)).ExecuteStream(context.Background(), codexTestAuth(server.URL), req, opts) + if err != nil { + t.Fatalf("unexpected ExecuteStream error: %v", err) + } + + combined, streamErr := drainChunks(result) + if streamErr != nil { + t.Fatalf("unexpected chunk error: %v", streamErr) + } + createdAt := strings.Index(combined, "response.created") + addedAt := strings.Index(combined, "response.output_item.added") + if createdAt < 0 || addedAt < 0 { + t.Fatalf("missing handshake or first generated event: %s", combined) + } + if createdAt > addedAt { + t.Fatalf("buffered handshake must be replayed before the first generated event: %s", combined) + } +} + +func TestCodexWebsocketsExecutor_BootstrapBuffering_DefaultDisabledPassthrough(t *testing.T) { + server := codexWebsocketServer(t, codexCreatedEvent, codexOverloadEvent) + defer server.Close() + + req, opts := codexWebsocketRequest() + result, err := NewCodexWebsocketsExecutor(&config.Config{}).ExecuteStream(context.Background(), codexTestAuth(server.URL), req, opts) + if err != nil { + t.Fatalf("default unbuffered ExecuteStream returned error at call time: %v", err) + } + if result == nil { + t.Fatal("expected non-nil result in default unbuffered mode") + } + _, streamErr := drainChunks(result) + if streamErr == nil { + t.Fatal("expected stream error in chunks for default unbuffered mode") + } + if got := statusCodeFromTestError(t, streamErr); got != http.StatusBadGateway { + t.Fatalf("status code = %d, want %d while buffering is disabled", got, http.StatusBadGateway) + } +} + +// The 503 restoration is scoped to the buffered failover path, so this only covers which +// rejections are eligible to replace the whole attempt. +func TestIsCodexOverloadBootstrapFailureRejectsRequestFaults(t *testing.T) { + notOverload := []string{ + `{"error":{"type":"invalid_request_error","code":"invalid_value"}}`, + `{"error":{"type":"authentication_error","code":"invalid_api_key"}}`, + `{"error":{"type":"upstream_error","code":"unknown"}}`, + } + for _, body := range notOverload { + if isCodexOverloadBootstrapFailure([]byte(body)) { + t.Fatalf("request-level fault must not trigger bootstrap failover: %s", body) + } + } + if !isCodexOverloadBootstrapFailure([]byte(`{"error":{"type":"rate_limit_error","code":"rate_limit_exceeded"}}`)) { + t.Fatal("rate limit rejections should be eligible for bootstrap failover") + } +} + +// codexWebsocketServerHoldingConnection behaves like codexWebsocketServer but keeps the upstream +// connection open after writing the frames, so the executor's own teardown path is the only +// source of session invalidation. With the plain helper the connection closes immediately, the +// reader goroutine observes EOF first and reports upstream_disconnected, which both masks the +// path under test and can make a disconnect assertion pass for the wrong reason. +func codexWebsocketServerHoldingConnection(t *testing.T, frames ...string) *httptest.Server { + t.Helper() + upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade websocket: %v", err) + return + } + defer func() { _ = conn.Close() }() + if _, _, errRead := conn.ReadMessage(); errRead != nil { + t.Errorf("read websocket message: %v", errRead) + return + } + for _, frame := range frames { + _ = conn.WriteMessage(websocket.TextMessage, []byte(frame)) + } + for { + if _, _, errRead := conn.ReadMessage(); errRead != nil { + return + } + } + })) +} + +// executeWebsocketStreamInSession runs ExecuteStream bound to a named execution session and +// reports whether the upstream teardown was signalled to the downstream handler. +// +// The downstream Responses WebSocket handler subscribes to UpstreamDisconnectChan and closes +// the client connection as soon as a disconnect is published. A bootstrap overload is retried +// on another credential, so publishing there would tear down the client connection before the +// retry can deliver anything, and the client would observe an abnormal close with zero frames. +func executeWebsocketStreamInSession(t *testing.T, frames ...string) (notified bool, err error) { + t.Helper() + + server := codexWebsocketServerHoldingConnection(t, frames...) + defer server.Close() + + exec := NewCodexWebsocketsExecutor(codexBufferingConfig(true)) + exec.store = &codexWebsocketSessionStore{sessions: make(map[string]*codexWebsocketSession)} + + const sessionID = "bootstrap-session" + disconnectCh := exec.UpstreamDisconnectChan(sessionID) + if disconnectCh == nil { + t.Fatal("expected a disconnect channel") + } + + req, opts := codexWebsocketRequest() + opts.Metadata = map[string]any{cliproxyexecutor.ExecutionSessionMetadataKey: sessionID} + _, err = exec.ExecuteStream(context.Background(), codexTestAuth(server.URL), req, opts) + + select { + case <-disconnectCh: + notified = true + default: + } + return notified, err +} + +func TestCodexWebsocketsExecutor_BootstrapOverload_DoesNotNotifyDownstreamDisconnect(t *testing.T) { + notified, err := executeWebsocketStreamInSession(t, codexCreatedEvent, codexInProgressEvent, codexOverloadEvent) + + if err == nil { + t.Fatal("expected the overload rejection to fail the attempt") + } + if got := statusCodeFromTestError(t, err); got != http.StatusServiceUnavailable { + t.Fatalf("status code = %d, want %d", got, http.StatusServiceUnavailable) + } + if notified { + t.Fatal("bootstrap overload must not signal a downstream disconnect: the conductor still has to retry on another credential, and signalling closes the client connection with zero frames delivered") + } +} + +// A non-overload terminal failure is delivered in-stream and genuinely ends the session, so it +// must keep signalling the disconnect exactly as it did before buffering existed. +func TestCodexWebsocketsExecutor_BootstrapNonOverload_StillNotifiesDownstreamDisconnect(t *testing.T) { + notified, err := executeWebsocketStreamInSession(t, codexCreatedEvent, codexInProgressEvent, codexInvalidEvent) + + if err != nil { + t.Fatalf("non-overload failures stay in-stream, got err = %v", err) + } + if !notified { + t.Fatal("a terminal failure that is delivered in-stream must still signal the downstream disconnect") + } +} diff --git a/internal/runtime/executor/codex_websockets_stream.go b/internal/runtime/executor/codex_websockets_stream.go index d094a894b..52279bc47 100644 --- a/internal/runtime/executor/codex_websockets_stream.go +++ b/internal/runtime/executor/codex_websockets_stream.go @@ -262,7 +262,207 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr sess.setMultiAgentV2Optimized(conn, optimizeMultiAgentV2 && !multiAgentV2Conflict) } - out := make(chan cliproxyexecutor.StreamChunk) + buffering := e.cfg != nil && e.cfg.Codex.StreamBootstrapBuffering + + claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload) + var param any + outputItemsByIndex := make(map[int64][]byte) + var outputItemsFallback [][]byte + + var bufferedChunks [][]byte + var initialChunks [][]byte + immediateTerminal := false + // bootstrapTerminalErr holds a non-overload terminal failure seen while buffering. It is + // delivered as an in-stream chunk after the buffered handshake so downstream behaviour stays + // identical to the unbuffered path instead of silently turning into a credential failover. + var bootstrapTerminalErr error + + if buffering { + for { + if ctx != nil && ctx.Err() != nil { + if sess != nil { + sess.clearActive(conn, readCh) + unlockStreamSession() + } else { + _ = closer.Close() + } + return nil, ctx.Err() + } + msgType, payload, errRead := readCodexWebsocketMessage(ctx, sess, conn, readCh) + if errRead != nil { + mappedErr := mapCodexWebsocketReadError(errRead) + if sess != nil { + e.invalidateUpstreamConn(sess, conn, "read_error", mappedErr) + sess.clearActive(conn, readCh) + unlockStreamSession() + } else { + logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, "read_error", mappedErr) + _ = closer.Close() + } + helps.RecordAPIWebsocketError(ctx, e.cfg, "read", mappedErr) + reporter.PublishFailure(ctx, mappedErr) + return nil, mappedErr + } + if msgType != websocket.TextMessage { + if msgType == websocket.BinaryMessage { + errBinary := fmt.Errorf("codex websockets executor: unexpected binary message") + if sess != nil { + e.invalidateUpstreamConn(sess, conn, "unexpected_binary", errBinary) + sess.clearActive(conn, readCh) + unlockStreamSession() + } else { + logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, "unexpected_binary", errBinary) + _ = closer.Close() + } + helps.RecordAPIWebsocketError(ctx, e.cfg, "unexpected_binary", errBinary) + reporter.PublishFailure(ctx, errBinary) + return nil, errBinary + } + continue + } + + payload = bytes.TrimSpace(payload) + if len(payload) == 0 { + continue + } + reporter.MarkFirstResponseByte() + payload = applyCodexIdentityConfuseResponsePayload(payload, identityState) + helps.AppendAPIWebsocketResponse(ctx, e.cfg, payload) + payload = helps.RestoreCodexMultiAgentV2Response(payload, restoreMultiAgentV2) + + if wsErr, ok := parseCodexWebsocketError(payload); ok { + if sess != nil { + e.invalidateUpstreamConn(sess, conn, "upstream_error", wsErr) + sess.clearActive(conn, readCh) + unlockStreamSession() + } else { + logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, "upstream_error", wsErr) + _ = closer.Close() + } + if errClearReplay := clearCodexReasoningReplayOnWebsocketError(ctx, replayScope, payload); errClearReplay != nil { + helps.RecordAPIWebsocketError(ctx, e.cfg, "replay_clear_error", errClearReplay) + reporter.PublishFailure(ctx, errClearReplay) + return nil, errClearReplay + } + helps.RecordAPIWebsocketError(ctx, e.cfg, "upstream_error", wsErr) + reporter.PublishFailure(ctx, wsErr) + return nil, wsErr + } + if streamErr, terminalBody, ok := codexTerminalFailureErr(payload); ok { + // A transient capacity rejection is retried on another credential, so the + // downstream websocket session must survive this upstream teardown. Notifying + // the disconnect here would close the client connection before the retry can + // deliver anything. Every other terminal failure is forwarded in-stream and + // legitimately terminates the session, so it keeps the notifying variant. + failoverPending := isCodexOverloadBootstrapFailure(terminalBody) + if sess != nil { + unlockStreamSession() + if failoverPending { + e.invalidateUpstreamConnWithoutDisconnectNotify(sess, conn, "terminal_failure", streamErr) + } else { + e.invalidateUpstreamConn(sess, conn, "terminal_failure", streamErr) + } + sess.clearActive(conn, readCh) + } else { + logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, "terminal_failure", streamErr) + _ = closer.Close() + } + if errClearReplay := clearCodexReasoningReplayOnInvalidSignature(ctx, replayScope, streamErr.StatusCode(), terminalBody); errClearReplay != nil { + helps.RecordAPIWebsocketError(ctx, e.cfg, "replay_clear_error", errClearReplay) + reporter.PublishFailure(ctx, errClearReplay) + return nil, errClearReplay + } + helps.RecordAPIWebsocketError(ctx, e.cfg, "upstream_error", streamErr) + reporter.PublishFailure(ctx, streamErr) + if failoverPending { + // Fail the attempt before the downstream headers are committed so the + // conductor can transparently retry on another credential, and report the + // status the upstream refused to put on the wire. + helps.LogWithRequestID(ctx).Debugf("codex websockets executor: bootstrap overload rejection after %d buffered handshake events, failing over", len(bufferedChunks)) + return nil, newCodexBootstrapOverloadErr(terminalBody) + } + bootstrapTerminalErr = streamErr + break + } + + eventType := gjson.GetBytes(payload, "type").String() + isTerminalEvent := eventType == "response.completed" || eventType == "response.done" || eventType == "error" + if eventType == "response.output_item.done" { + collectCodexOutputItemDone(payload, outputItemsByIndex, &outputItemsFallback) + } + completedPayload := payload + if eventType == "response.completed" || eventType == "response.done" { + completedPayload = normalizeCodexWebsocketCompletion(completedPayload) + completedPayload = patchCodexCompletedOutput(completedPayload, outputItemsByIndex, outputItemsFallback) + cacheCodexReasoningReplayFromCompleted(replayScope, completedPayload) + if detail, ok := helps.ParseCodexUsage(completedPayload); ok { + reporter.Publish(ctx, detail) + } + } + + var currentChunks [][]byte + if cliproxyexecutor.DownstreamWebsocket(ctx) { + clientPayload := applyCodexIdentityExposeResponsePayload(payload, identityState) + downstreamPayload := helps.EnsureResponsesUsageDetails(clientPayload) + currentChunks = [][]byte{downstreamPayload} + } else { + payload = normalizeCodexWebsocketCompletion(payload) + if eventType == "response.completed" || eventType == "response.done" { + payload = completedPayload + } + clientPayload := applyCodexIdentityExposeResponsePayload(payload, identityState) + line := encodeCodexWebsocketAsSSE(clientPayload) + currentChunks = helps.TranslateStreamWithClaudeInputTokens(ctx, to, responseFormat, req.Model, originalPayload, clientBody, line, ¶m, claudeInputTokens) + } + + if isCodexHandshakeMetadataEvent(eventType) && !isTerminalEvent { + if len(bufferedChunks) < codexBootstrapMaxBufferedEvents { + bufferedChunks = append(bufferedChunks, currentChunks...) + continue + } + helps.LogWithRequestID(ctx).Debugf("codex websockets executor: bootstrap buffer limit %d reached, releasing stream without overload probing", codexBootstrapMaxBufferedEvents) + } + + initialChunks = currentChunks + if isTerminalEvent { + immediateTerminal = true + } + break + } + } + + chanCapacity := len(bufferedChunks) + len(initialChunks) + if bootstrapTerminalErr != nil { + chanCapacity++ + } + out := make(chan cliproxyexecutor.StreamChunk, chanCapacity) + for _, chunk := range bufferedChunks { + out <- cliproxyexecutor.StreamChunk{Payload: chunk} + } + for _, chunk := range initialChunks { + out <- cliproxyexecutor.StreamChunk{Payload: chunk} + } + if bootstrapTerminalErr != nil { + // The upstream connection was already invalidated and released in the terminal-failure + // branch above, so only the buffered payloads plus the in-stream error remain to emit. + out <- cliproxyexecutor.StreamChunk{Err: bootstrapTerminalErr} + close(out) + return &cliproxyexecutor.StreamResult{Headers: upstreamHeaders, Chunks: out}, nil + } + if immediateTerminal { + if sess != nil { + sess.clearActive(conn, readCh) + unlockStreamSession() + } else { + logCodexWebsocketDisconnected(executionSessionID, authID, wsURL, "completed", nil) + if errClose := closer.Close(); errClose != nil { + log.Errorf("codex websockets executor: close websocket error: %v", errClose) + } + } + close(out) + return &cliproxyexecutor.StreamResult{Headers: upstreamHeaders, Chunks: out}, nil + } + go func() { terminateReason := "completed" var terminateErr error @@ -293,10 +493,6 @@ func (e *CodexWebsocketsExecutor) ExecuteStream(ctx context.Context, auth *clipr } } - claudeInputTokens := helps.NewClaudeInputTokenState(from, to, responseFormat, originalPayload) - var param any - outputItemsByIndex := make(map[int64][]byte) - var outputItemsFallback [][]byte for { if ctx != nil && ctx.Err() != nil { terminateReason = "context_done" diff --git a/internal/watcher/diff/config_diff.go b/internal/watcher/diff/config_diff.go index b19385bd7..b8a79720d 100644 --- a/internal/watcher/diff/config_diff.go +++ b/internal/watcher/diff/config_diff.go @@ -115,6 +115,9 @@ func BuildConfigChangeDetails(oldCfg, newCfg *config.Config) []string { if oldCfg.Codex.DisableCodexCloaking != newCfg.Codex.DisableCodexCloaking { changes = append(changes, fmt.Sprintf("codex.disable-codex-cloaking: %t -> %t", oldCfg.Codex.DisableCodexCloaking, newCfg.Codex.DisableCodexCloaking)) } + if oldCfg.Codex.StreamBootstrapBuffering != newCfg.Codex.StreamBootstrapBuffering { + changes = append(changes, fmt.Sprintf("codex.stream-bootstrap-buffering: %t -> %t", oldCfg.Codex.StreamBootstrapBuffering, newCfg.Codex.StreamBootstrapBuffering)) + } if oldCfg.Codex.OptimizeMultiAgentV2 != newCfg.Codex.OptimizeMultiAgentV2 { changes = append(changes, fmt.Sprintf("codex.optimize-multi-agent-v2: %t -> %t", oldCfg.Codex.OptimizeMultiAgentV2, newCfg.Codex.OptimizeMultiAgentV2)) } diff --git a/sdk/cliproxy/auth/conductor_stream_overload_failover_test.go b/sdk/cliproxy/auth/conductor_stream_overload_failover_test.go new file mode 100644 index 000000000..d0b1c2c3f --- /dev/null +++ b/sdk/cliproxy/auth/conductor_stream_overload_failover_test.go @@ -0,0 +1,168 @@ +package auth + +import ( + "context" + "fmt" + "net/http" + "sync" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +// registerOverloadAuths registers n active codex credentials with descending priority so the +// selection order is deterministic, and returns their IDs in expected pick order. +func registerOverloadAuths(t *testing.T, m *Manager, n int) []string { + t.Helper() + reg := registry.GetGlobalRegistry() + ids := make([]string, 0, n) + for i := 0; i < n; i++ { + id := fmt.Sprintf("auth-overload-%d", i+1) + auth := &Auth{ + ID: id, + Provider: "codex", + Status: StatusActive, + // Higher priority is picked first, so descending values keep the order stable. + Attributes: map[string]string{"priority": fmt.Sprintf("%d", 100-i)}, + } + reg.RegisterClient(id, "codex", []*registry.ModelInfo{{ID: "gpt-5.6-terra"}}) + if _, err := m.Register(context.Background(), auth); err != nil { + t.Fatalf("register %s: %v", id, err) + } + ids = append(ids, id) + } + t.Cleanup(func() { + for _, id := range ids { + reg.UnregisterClient(id) + } + }) + return ids +} + +func overloadStatusError() customStatusError { + return customStatusError{ + code: http.StatusServiceUnavailable, + msg: `{"error":{"type":"service_unavailable_error","code":"server_is_overloaded","message":"Our servers are currently overloaded. Please try again later.","param":null}}`, + } +} + +func successStreamResult() *cliproxyexecutor.StreamResult { + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: []byte(`data: {"type":"response.output_item.added"}`)} + ch <- cliproxyexecutor.StreamChunk{Payload: []byte(`data: {"type":"response.completed"}`)} + close(ch) + return &cliproxyexecutor.StreamResult{ + Headers: http.Header{"Content-Type": []string{"text/event-stream"}}, + Chunks: ch, + } +} + +// With stream-bootstrap-buffering enabled the codex executor returns the overload rejection +// synchronously instead of relaying it in-stream. This test pins the operational question: with +// request-retry=5 and max-retry-credentials=6, do three consecutive overloaded accounts get +// skipped so the fourth credential serves the request? +func TestExecuteStream_BootstrapOverload_SkipsConsecutiveOverloadedCredentials(t *testing.T) { + previous := quotaCooldownDisabled.Load() + quotaCooldownDisabled.Store(false) + t.Cleanup(func() { quotaCooldownDisabled.Store(previous) }) + + m := NewManager(nil, nil, nil) + m.SetRetryConfig(5, 0, 6) + ids := registerOverloadAuths(t, m, 6) + + var mu sync.Mutex + var order []string + overloaded := map[string]bool{ids[0]: true, ids[1]: true, ids[2]: true} + + m.RegisterExecutor(&customStreamMockExecutor{ + identifier: "codex", + streamFn: func(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + mu.Lock() + order = append(order, auth.ID) + mu.Unlock() + if overloaded[auth.ID] { + return nil, overloadStatusError() + } + return successStreamResult(), nil + }, + }) + + result, err := m.ExecuteStream(context.Background(), []string{"codex"}, + cliproxyexecutor.Request{Model: "gpt-5.6-terra"}, cliproxyexecutor.Options{}) + if err != nil { + t.Fatalf("expected the request to survive three overloaded credentials: %v", err) + } + if result == nil { + t.Fatal("expected a stream result from the fourth credential") + } + for chunk := range result.Chunks { + if chunk.Err != nil { + t.Fatalf("unexpected chunk error: %v", chunk.Err) + } + } + + mu.Lock() + defer mu.Unlock() + if len(order) != 4 { + t.Fatalf("attempted %d credentials (%v), want exactly 4", len(order), order) + } + for i := 0; i < 3; i++ { + if overloaded[order[i]] != true { + t.Fatalf("attempt %d used %s, expected one of the overloaded credentials", i+1, order[i]) + } + } + if overloaded[order[3]] { + t.Fatalf("final attempt used overloaded credential %s", order[3]) + } +} + +// The credential budget must be honoured: when every credential is overloaded the request fails +// after max-retry-credentials attempts rather than looping forever. +func TestExecuteStream_BootstrapOverload_StopsAtCredentialBudget(t *testing.T) { + previous := quotaCooldownDisabled.Load() + quotaCooldownDisabled.Store(false) + t.Cleanup(func() { quotaCooldownDisabled.Store(previous) }) + + m := NewManager(nil, nil, nil) + // Six credentials exist but only four may be attempted. + m.SetRetryConfig(5, 0, 4) + registerOverloadAuths(t, m, 6) + + var mu sync.Mutex + attempts := 0 + m.RegisterExecutor(&customStreamMockExecutor{ + identifier: "codex", + streamFn: func(_ context.Context, _ *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + mu.Lock() + attempts++ + mu.Unlock() + return nil, overloadStatusError() + }, + }) + + done := make(chan struct{}) + go func() { + defer close(done) + _, _ = m.ExecuteStream(context.Background(), []string{"codex"}, + cliproxyexecutor.Request{Model: "gpt-5.6-terra"}, cliproxyexecutor.Options{}) + }() + select { + case <-done: + case <-time.After(30 * time.Second): + t.Fatal("ExecuteStream did not terminate within the credential budget") + } + + mu.Lock() + defer mu.Unlock() + if attempts == 0 { + t.Fatal("expected at least one attempt") + } + // The outer request-retry loop may restart the credential sweep, so assert the per-sweep + // budget is respected rather than a single exact total. + if attempts%4 != 0 { + t.Fatalf("attempts = %d, expected a multiple of the 4-credential budget", attempts) + } + t.Logf("total upstream attempts across retry sweeps: %d", attempts) +} diff --git a/sdk/cliproxy/auth/conductor_stream_overload_status_test.go b/sdk/cliproxy/auth/conductor_stream_overload_status_test.go new file mode 100644 index 000000000..c8901b83b --- /dev/null +++ b/sdk/cliproxy/auth/conductor_stream_overload_status_test.go @@ -0,0 +1,138 @@ +package auth + +import ( + "context" + "net/http" + "testing" + + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +// When every credential is exhausted by overload rejections the caller must receive a real +// error carrying 503, not a committed 200 stream. This is what lets the downstream client and +// any upstream proxy see the true capacity signal. +func TestExecuteStream_AllCredentialsOverloaded_ReturnsStatusError(t *testing.T) { + previous := quotaCooldownDisabled.Load() + quotaCooldownDisabled.Store(false) + t.Cleanup(func() { quotaCooldownDisabled.Store(previous) }) + + m := NewManager(nil, nil, nil) + m.SetRetryConfig(5, 0, 3) + registerOverloadAuths(t, m, 3) + + m.RegisterExecutor(&customStreamMockExecutor{ + identifier: "codex", + streamFn: func(_ context.Context, _ *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + // Mirrors the buffering-enabled codex executor: the rejection is returned + // synchronously, before any downstream chunk is committed. + return nil, overloadStatusError() + }, + }) + + result, err := m.ExecuteStream(context.Background(), []string{"codex"}, + cliproxyexecutor.Request{Model: "gpt-5.6-terra"}, cliproxyexecutor.Options{}) + + if err == nil { + t.Fatalf("expected a hard error once every credential is overloaded, got result=%v", result) + } + statusErr, ok := err.(interface{ StatusCode() int }) + if !ok { + t.Fatalf("error %T does not expose StatusCode(): %v", err, err) + } + if got := statusErr.StatusCode(); got != http.StatusServiceUnavailable { + t.Fatalf("status code = %d, want %d", got, http.StatusServiceUnavailable) + } +} + +// Contrast: the unbuffered path commits response.created first, so the rejection can only be +// relayed inside an already-successful stream. The caller gets no error at all. +func TestExecuteStream_UnbufferedOverload_StaysCommittedStream(t *testing.T) { + previous := quotaCooldownDisabled.Load() + quotaCooldownDisabled.Store(false) + t.Cleanup(func() { quotaCooldownDisabled.Store(previous) }) + + m := NewManager(nil, nil, nil) + m.SetRetryConfig(5, 0, 3) + registerOverloadAuths(t, m, 3) + + m.RegisterExecutor(&customStreamMockExecutor{ + identifier: "codex", + streamFn: func(_ context.Context, _ *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + ch := make(chan cliproxyexecutor.StreamChunk, 2) + ch <- cliproxyexecutor.StreamChunk{Payload: []byte(`data: {"type":"response.created"}`)} + ch <- cliproxyexecutor.StreamChunk{Err: overloadStatusError()} + close(ch) + return &cliproxyexecutor.StreamResult{ + Headers: http.Header{"Content-Type": []string{"text/event-stream"}}, + Chunks: ch, + }, nil + }, + }) + + result, err := m.ExecuteStream(context.Background(), []string{"codex"}, + cliproxyexecutor.Request{Model: "gpt-5.6-terra"}, cliproxyexecutor.Options{}) + if err != nil { + t.Fatalf("unbuffered path should hand back a committed stream, got error: %v", err) + } + if result == nil { + t.Fatal("expected a committed stream result") + } + var sawErr bool + for chunk := range result.Chunks { + if chunk.Err != nil { + sawErr = true + } + } + if !sawErr { + t.Fatal("expected the overload rejection to arrive in-stream") + } +} + +// Critical distinction: if an executor surfaces the rejection as the *first* stream chunk instead +// of returning it synchronously, the conductor downgrades it to a committed stream carrying the +// error (streamErrorResult), and the caller again observes no error. Returning synchronously is +// therefore required to preserve the 503 status semantics. +func TestExecuteStream_ErrorAsFirstChunk_IsDowngradedToCommittedStream(t *testing.T) { + previous := quotaCooldownDisabled.Load() + quotaCooldownDisabled.Store(false) + t.Cleanup(func() { quotaCooldownDisabled.Store(previous) }) + + m := NewManager(nil, nil, nil) + m.SetRetryConfig(5, 0, 3) + registerOverloadAuths(t, m, 3) + + m.RegisterExecutor(&customStreamMockExecutor{ + identifier: "codex", + streamFn: func(_ context.Context, _ *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + ch := make(chan cliproxyexecutor.StreamChunk, 1) + ch <- cliproxyexecutor.StreamChunk{Err: overloadStatusError()} + close(ch) + return &cliproxyexecutor.StreamResult{ + Headers: http.Header{"Content-Type": []string{"text/event-stream"}}, + Chunks: ch, + }, nil + }, + }) + + result, err := m.ExecuteStream(context.Background(), []string{"codex"}, + cliproxyexecutor.Request{Model: "gpt-5.6-terra"}, cliproxyexecutor.Options{}) + + if err != nil { + t.Logf("first-chunk error surfaced as a hard error: %v", err) + t.Log("NOTE: this contradicts the streamErrorResult downgrade path; review if it changes") + return + } + if result == nil { + t.Fatal("expected either an error or a committed stream") + } + var sawErr bool + for chunk := range result.Chunks { + if chunk.Err != nil { + sawErr = true + } + } + if !sawErr { + t.Fatal("expected the rejection to be delivered in-stream after the downgrade") + } + t.Log("confirmed: an error delivered as the first chunk is downgraded to a committed stream") +}