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") +}