diff --git a/internal/runtime/executor/codex_websockets_executor.go b/internal/runtime/executor/codex_websockets_executor.go index 87ae0efe4..d6f1de86b 100644 --- a/internal/runtime/executor/codex_websockets_executor.go +++ b/internal/runtime/executor/codex_websockets_executor.go @@ -868,7 +868,7 @@ func applyCodexWebsocketHeaders(ctx context.Context, headers http.Header, auth * if auth != nil && auth.Metadata != nil { if accountID, ok := auth.Metadata["account_id"].(string); ok { if trimmed := strings.TrimSpace(accountID); trimmed != "" { - headers.Set("ChatGPT-Account-ID", trimmed) + setHeaderCasePreserved(headers, "ChatGPT-Account-ID", trimmed) } } } @@ -1040,7 +1040,9 @@ func parseCodexWebsocketError(payload []byte) (error, bool) { out := buildCodexWebsocketErrorPayload(payload, status) headers := parseCodexWebsocketErrorHeaders(payload) statusError := statusErr{code: status, msg: string(out)} - if isCodexWebsocketConnectionLimitError(payload) { + if retryAfter := parseCodexRetryAfter(status, out, time.Now()); retryAfter != nil { + statusError.retryAfter = retryAfter + } else if isCodexWebsocketConnectionLimitError(payload) { retryAfter := time.Duration(0) statusError.retryAfter = &retryAfter } diff --git a/internal/runtime/executor/codex_websockets_executor_test.go b/internal/runtime/executor/codex_websockets_executor_test.go index 0b7a546e9..bf12ef786 100644 --- a/internal/runtime/executor/codex_websockets_executor_test.go +++ b/internal/runtime/executor/codex_websockets_executor_test.go @@ -296,9 +296,16 @@ func TestApplyCodexWebsocketHeadersUsesCanonicalAccountHeader(t *testing.T) { headers := applyCodexWebsocketHeaders(context.Background(), http.Header{}, auth, "", nil) - if got := headers.Get("ChatGPT-Account-ID"); got != "acct-1" { + if got := headerValueCaseInsensitive(headers, "ChatGPT-Account-ID"); got != "acct-1" { t.Fatalf("ChatGPT-Account-ID = %s, want acct-1", got) } + values, ok := headers["ChatGPT-Account-ID"] + if !ok { + t.Fatalf("expected exact ChatGPT-Account-ID key, got %#v", headers) + } + if len(values) != 1 || values[0] != "acct-1" { + t.Fatalf("ChatGPT-Account-ID values = %#v, want [acct-1]", values) + } } func TestBuildCodexResponsesWebsocketURLRequiresHTTPURL(t *testing.T) { @@ -326,12 +333,30 @@ func TestParseCodexWebsocketErrorMarksConnectionLimitRetryable(t *testing.T) { if !ok || retryable.RetryAfter() == nil { t.Fatalf("expected retryable websocket connection limit error") } + if got := *retryable.RetryAfter(); got != 0 { + t.Fatalf("retryAfter = %v, want connection-limit fallback 0", got) + } withHeaders, ok := err.(interface{ Headers() http.Header }) if !ok || withHeaders.Headers().Get("retry-after") != "1" { t.Fatalf("headers = %#v, want retry-after", err) } } +func TestParseCodexWebsocketErrorUsesUsageLimitRetryMetadata(t *testing.T) { + err, ok := parseCodexWebsocketError([]byte(`{"type":"error","status":429,"body":{"error":{"type":"usage_limit_reached","message":"usage limit reached","resets_in_seconds":7}}}`)) + if !ok { + t.Fatalf("expected websocket error") + } + + retryable, ok := err.(interface{ RetryAfter() *time.Duration }) + if !ok || retryable.RetryAfter() == nil { + t.Fatalf("expected retryable usage limit websocket error") + } + if got := *retryable.RetryAfter(); got != 7*time.Second { + t.Fatalf("retryAfter = %v, want 7s", got) + } +} + func TestParseCodexWebsocketErrorPreservesWrappedBodyAndHeaders(t *testing.T) { err, ok := parseCodexWebsocketError([]byte(`{"type":"error","status":429,"body":{"error":{"code":"websocket_connection_limit_reached","type":"server_error","message":"too many websocket connections"}},"headers":{"x-request-id":"req-1"}}`)) if !ok {