Fix Codex websocket retry metadata

This commit is contained in:
Kenny
2026-05-03 19:01:44 -07:00
parent c19ae1d5be
commit 08b0fe6816
2 changed files with 30 additions and 3 deletions

View File

@@ -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
}

View File

@@ -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 {