From 601ca43090f7aa41b2d16733ce92ee3d9fe0ab35 Mon Sep 17 00:00:00 2001 From: hkfires <10558748+hkfires@users.noreply.github.com> Date: Fri, 21 Aug 2026 18:45:02 +0800 Subject: [PATCH] feat(auth): add credential retry round contract Redefine request-retry as additional credential retry rounds and enforce max-retry-credentials per round. Home dispatch now carries excluded and pinned auth constraints, supports remote retry limits, and propagates cooldown retry-after metadata across exhausted rounds. Move Antigravity upstream retries under conductor ownership to avoid double-consuming retry attempts. Update configuration comments and add coverage for Home retry rounds, cooldown handling, pinned credentials, and legacy dispatcher compatibility. --- config.example.yaml | 29 +- internal/config/config.go | 11 +- internal/config/config_types.go | 8 +- internal/config/vertex_compat.go | 2 +- internal/home/client.go | 33 +- internal/home/client_test.go | 66 +- internal/home/requests.go | 2 + .../executor/antigravity_executor_credits.go | 20 - .../antigravity_executor_credits_test.go | 28 +- .../executor/antigravity_executor_execute.go | 8 +- .../executor/antigravity_executor_stream.go | 4 +- sdk/cliproxy/auth/conductor_execution.go | 238 ++- sdk/cliproxy/auth/conductor_home.go | 213 ++- sdk/cliproxy/auth/conductor_home_execution.go | 75 +- sdk/cliproxy/auth/conductor_lifecycle.go | 2 +- sdk/cliproxy/auth/conductor_overrides_test.go | 254 +++ sdk/cliproxy/auth/conductor_selection.go | 224 ++- ...conductor_stream_overload_failover_test.go | 10 +- sdk/cliproxy/auth/home_concurrency.go | 34 +- .../auth/home_execution_paths_test.go | 66 +- sdk/cliproxy/auth/home_retry_contract_test.go | 1477 +++++++++++++++++ sdk/cliproxy/auth/home_selection.go | 2 + 22 files changed, 2676 insertions(+), 130 deletions(-) create mode 100644 sdk/cliproxy/auth/home_retry_contract_test.go diff --git a/config.example.yaml b/config.example.yaml index 18bc0ca3..7ee54309 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -141,16 +141,21 @@ force-model-prefix: false # Default is false (disabled). passthrough-headers: false -# Number of times to retry a request. Retries will occur if the HTTP response code is 403, 408, 500, 502, 503, or 504. -# Individual API-key entries can override this via request-retry. OAuth/token JSON files can override via request_retry. -# 0 disables retries for that credential; omit the field or set a negative value to keep this global setting. +# Number of additional credential retry rounds after the first round exhausts +# its eligible credentials. A value of 0 still allows the first round to try +# multiple credentials. +# Additional rounds apply to HTTP 403, 408, 429, 500, 502, 503, and 504 failures. +# Individual credential/provider overrides take precedence; 0 disables additional +# rounds, while an omitted or negative override inherits this global setting. request-retry: 3 -# Maximum number of different credentials to try for one failed request. +# Maximum number of different credentials to try in each credential retry round. # Set to 0 to keep legacy behavior (try all available credentials). max-retry-credentials: 0 -# Maximum wait time in seconds for a cooled-down credential before triggering a retry. +# Maximum cooldown wait in seconds between retry rounds. +# Set to 0 or below to never wait for credential cooldown. +# Retry rounds that need no wait remain controlled by request-retry. max-retry-interval: 30 # When true, disable auth/model cooldown scheduling globally (prevents blackout windows after failure states). @@ -309,7 +314,7 @@ nonstream-keepalive-interval: 0 # weight: 5 # optional: weighted-round-robin share; omitted defaults to 1; maximum 1,000,000 # prefix: "test" # optional: require calls like "test/gemini-3-pro-preview" to target this credential # disable-cooling: false # optional override: true disables cooling, false enables it; omit to inherit global -# request-retry: 3 # optional: per-auth override of the global request-retry; 0 disables retries; omit or set < 0 to use the global value +# request-retry: 3 # optional per-auth override; 0 disables additional rounds; omit or set < 0 to inherit global # request-scoped-errors: # optional: custom rules to classify upstream errors by status and body patterns # - status: 400 # HTTP status code to match # match: # optional: string contains matching @@ -351,7 +356,7 @@ nonstream-keepalive-interval: 0 # weight: 5 # optional: weighted-round-robin share; omitted defaults to 1; maximum 1,000,000 # prefix: "native" # optional: require calls like "native/gemini-3-pro-preview" to target this credential # disable-cooling: false # optional override: true disables cooling, false enables it; omit to inherit global -# request-retry: 3 # optional: per-auth override of the global request-retry; 0 disables retries; omit or set < 0 to use the global value +# request-retry: 3 # optional per-auth override; 0 disables additional rounds; omit or set < 0 to inherit global # request-scoped-errors: # optional: custom rules to classify upstream errors by status and body patterns # - status: 400 # match: @@ -381,7 +386,7 @@ nonstream-keepalive-interval: 0 # weight: 5 # optional: weighted-round-robin share; omitted defaults to 1; maximum 1,000,000 # prefix: "test" # optional: require calls like "test/gpt-5-codex" to target this credential # disable-cooling: false # optional override: true disables cooling, false enables it; omit to inherit global -# request-retry: 3 # optional: per-auth override of the global request-retry; 0 disables retries; omit or set < 0 to use the global value +# request-retry: 3 # optional per-auth override; 0 disables additional rounds; omit or set < 0 to inherit global # request-scoped-errors: # optional: custom rules to classify upstream errors by status and body patterns # - status: 400 # match: @@ -423,7 +428,7 @@ nonstream-keepalive-interval: 0 # weight: 5 # optional: weighted-round-robin share; omitted defaults to 1; maximum 1,000,000 # prefix: "xai" # optional: require calls like "xai/grok-4.5" to target this credential # disable-cooling: false # optional override: true disables cooling, false enables it; omit to inherit global -# request-retry: 3 # optional: per-auth override of the global request-retry; 0 disables retries; omit or set < 0 to use the global value +# request-retry: 3 # optional per-auth override; 0 disables additional rounds; omit or set < 0 to inherit global # request-scoped-errors: # optional: custom rules to classify upstream errors by status and body patterns # - status: 400 # match: @@ -458,7 +463,7 @@ nonstream-keepalive-interval: 0 # weight: 5 # optional: weighted-round-robin share; omitted defaults to 1; maximum 1,000,000 # prefix: "test" # optional: require calls like "test/claude-sonnet-latest" to target this credential # disable-cooling: false # optional override: true disables cooling, false enables it; omit to inherit global -# request-retry: 3 # optional: per-auth override of the global request-retry; 0 disables retries; omit or set < 0 to use the global value +# request-retry: 3 # optional per-auth override; 0 disables additional rounds; omit or set < 0 to inherit global # request-scoped-errors: # optional: custom rules to classify upstream errors by status and body patterns # - status: 400 # match: @@ -603,7 +608,7 @@ nonstream-keepalive-interval: 0 # base-url: "https://openrouter.ai/api/v1" # The base URL of the provider. # support-prompt-cache-key: false # optional: derive prompt_cache_key for requests from all input protocols # disable-cooling: false # optional provider override: true disables cooling, false enables it; omit to inherit global -# request-retry: 3 # optional: per-provider override of the global request-retry; 0 disables retries; omit or set < 0 to use the global value +# request-retry: 3 # optional per-provider override; 0 disables additional rounds; omit or set < 0 to inherit global # request-scoped-errors: # optional: custom rules to classify upstream errors by status and body patterns # - status: 400 # match: @@ -653,7 +658,7 @@ nonstream-keepalive-interval: 0 # weight: 5 # optional: weighted-round-robin share; omitted defaults to 1; maximum 1,000,000 # prefix: "test" # optional: require calls like "test/vertex-pro" to target this credential # disable-cooling: false # optional override: true disables cooling, false enables it; omit to inherit global -# request-retry: 3 # optional: per-auth override of the global request-retry; 0 disables retries; omit or set < 0 to use the global value +# request-retry: 3 # optional per-auth override; 0 disables additional rounds; omit or set < 0 to inherit global # base-url: "https://example.com/api" # optional, e.g. https://zenmux.ai/api; falls back to Google Vertex when omitted # proxy-url: "socks5://proxy.example.com:1080" # optional per-key proxy override # # proxy-url: "direct" # optional: explicit direct connect for this credential diff --git a/internal/config/config.go b/internal/config/config.go index 2b3ef28a..d8f7c247 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -76,12 +76,17 @@ type Config struct { // When <= 0, the default worker count is used. AuthAutoRefreshWorkers int `yaml:"auth-auto-refresh-workers" json:"auth-auto-refresh-workers"` - // RequestRetry defines the retry times when the request failed. + // RequestRetry defines the number of additional credential retry rounds after + // the first round has exhausted its eligible credentials. RequestRetry int `yaml:"request-retry" json:"request-retry"` - // MaxRetryCredentials defines the maximum number of credentials to try for a failed request. + // MaxRetryCredentials defines the maximum number of different credentials to + // try in each credential retry round. // Set to 0 or a negative value to keep trying all available credentials (legacy behavior). MaxRetryCredentials int `yaml:"max-retry-credentials" json:"max-retry-credentials"` - // MaxRetryInterval defines the maximum wait time in seconds before retrying a cooled-down credential. + // MaxRetryInterval defines the maximum positive cooldown wait, in seconds, + // allowed before starting another credential retry round. A non-positive value + // forbids positive cooldown waits; it does not disable same-round credential + // failover or immediate additional rounds allowed by RequestRetry. MaxRetryInterval int `yaml:"max-retry-interval" json:"max-retry-interval"` // QuotaExceeded defines the behavior when a quota is exceeded. diff --git a/internal/config/config_types.go b/internal/config/config_types.go index 50c468f6..6228284c 100644 --- a/internal/config/config_types.go +++ b/internal/config/config_types.go @@ -378,7 +378,7 @@ type ClaudeKey struct { DisableCooling *bool `yaml:"disable-cooling,omitempty" json:"disable-cooling,omitempty"` // RequestRetry optionally overrides the global request-retry for this credential. - // Nil or a negative value means "use the global request-retry". 0 disables retries. + // Nil or a negative value means "use the global request-retry". 0 disables additional retry rounds. RequestRetry *int `yaml:"request-retry,omitempty" json:"request-retry,omitempty"` // RequestScopedErrors configures custom classification rules for upstream errors. @@ -496,7 +496,7 @@ type CodexKey struct { DisableCooling *bool `yaml:"disable-cooling,omitempty" json:"disable-cooling,omitempty"` // RequestRetry optionally overrides the global request-retry for this credential. - // Nil or a negative value means "use the global request-retry". 0 disables retries. + // Nil or a negative value means "use the global request-retry". 0 disables additional retry rounds. RequestRetry *int `yaml:"request-retry,omitempty" json:"request-retry,omitempty"` // RequestScopedErrors configures custom classification rules for upstream errors. @@ -593,7 +593,7 @@ type GeminiKey struct { DisableCooling *bool `yaml:"disable-cooling,omitempty" json:"disable-cooling,omitempty"` // RequestRetry optionally overrides the global request-retry for this credential. - // Nil or a negative value means "use the global request-retry". 0 disables retries. + // Nil or a negative value means "use the global request-retry". 0 disables additional retry rounds. RequestRetry *int `yaml:"request-retry,omitempty" json:"request-retry,omitempty"` // RequestScopedErrors configures custom classification rules for upstream errors. @@ -680,7 +680,7 @@ type OpenAICompatibility struct { DisableCooling *bool `yaml:"disable-cooling,omitempty" json:"disable-cooling,omitempty"` // RequestRetry optionally overrides the global request-retry for this provider. - // Nil or a negative value means "use the global request-retry". 0 disables retries. + // Nil or a negative value means "use the global request-retry". 0 disables additional retry rounds. RequestRetry *int `yaml:"request-retry,omitempty" json:"request-retry,omitempty"` // RequestScopedErrors configures custom classification rules for upstream errors. diff --git a/internal/config/vertex_compat.go b/internal/config/vertex_compat.go index 0adefb49..8a7a76a0 100644 --- a/internal/config/vertex_compat.go +++ b/internal/config/vertex_compat.go @@ -51,7 +51,7 @@ type VertexCompatKey struct { DisableCooling *bool `yaml:"disable-cooling,omitempty" json:"disable-cooling,omitempty"` // RequestRetry optionally overrides the global request-retry for this credential. - // Nil or a negative value means "use the global request-retry". 0 disables retries. + // Nil or a negative value means "use the global request-retry". 0 disables additional retry rounds. RequestRetry *int `yaml:"request-retry,omitempty" json:"request-retry,omitempty"` } diff --git a/internal/home/client.go b/internal/home/client.go index 62964f18..c0424e97 100644 --- a/internal/home/client.go +++ b/internal/home/client.go @@ -1249,10 +1249,19 @@ func queryToLowerMap(query url.Values) map[string]string { return out } -func newAuthDispatchRequest(requestedModel string, sessionID string, headers http.Header, count int, credentialPolicy string) authDispatchRequest { +func newAuthDispatchRequest(requestedModel string, sessionID string, headers http.Header, count int, credentialPolicy string, excludedAuthIDs *[]string, pinnedAuthID string) authDispatchRequest { if count <= 0 { count = 1 } + var excludedAuthIDsCopy *[]string + if excludedAuthIDs != nil { + // CPA enforces the new retry-round contract locally. Keep count at one so + // older Home servers that ignore excluded_auth_ids do not apply their + // legacy count-based retry cap before CPA can rotate credentials. + count = 1 + values := append([]string{}, (*excludedAuthIDs)...) + excludedAuthIDsCopy = &values + } return authDispatchRequest{ Type: "auth", Model: requestedModel, @@ -1261,19 +1270,33 @@ func newAuthDispatchRequest(requestedModel string, sessionID string, headers htt SessionID: strings.TrimSpace(sessionID), Headers: headersToLowerMap(headers), CredentialPolicy: strings.TrimSpace(credentialPolicy), + ExcludedAuthIDs: excludedAuthIDsCopy, + PinnedAuthID: strings.TrimSpace(pinnedAuthID), } } func (c *Client) RPopAuth(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int) ([]byte, error) { - return c.rPopAuth(ctx, requestedModel, sessionID, headers, count, "") + return c.rPopAuth(ctx, requestedModel, sessionID, headers, count, "", nil, "") } // RPopAuthWithPolicy requests a Home credential constrained by the supplied fixed policy. func (c *Client) RPopAuthWithPolicy(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int, credentialPolicy string) ([]byte, error) { - return c.rPopAuth(ctx, requestedModel, sessionID, headers, count, credentialPolicy) + return c.rPopAuth(ctx, requestedModel, sessionID, headers, count, credentialPolicy, nil, "") +} + +// RPopAuthWithConstraints requests a credential using the current retry-round +// exclusions and optional pinned credential constraint. +func (c *Client) RPopAuthWithConstraints(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int, excludedAuthIDs []string, pinnedAuthID string) ([]byte, error) { + return c.rPopAuth(ctx, requestedModel, sessionID, headers, count, "", &excludedAuthIDs, pinnedAuthID) +} + +// RPopAuthWithPolicyAndConstraints combines a fixed credential policy with the +// current retry-round exclusions and optional pinned credential constraint. +func (c *Client) RPopAuthWithPolicyAndConstraints(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int, credentialPolicy string, excludedAuthIDs []string, pinnedAuthID string) ([]byte, error) { + return c.rPopAuth(ctx, requestedModel, sessionID, headers, count, credentialPolicy, &excludedAuthIDs, pinnedAuthID) } -func (c *Client) rPopAuth(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int, credentialPolicy string) ([]byte, error) { +func (c *Client) rPopAuth(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int, credentialPolicy string, excludedAuthIDs *[]string, pinnedAuthID string) ([]byte, error) { if c == nil || c.dispatchFenced.Load() { return nil, ErrDispatchFenced } @@ -1287,7 +1310,7 @@ func (c *Client) rPopAuth(ctx context.Context, requestedModel string, sessionID if requestedModel == "" { return nil, fmt.Errorf("home: requested model is empty") } - req := newAuthDispatchRequest(requestedModel, sessionID, headers, count, credentialPolicy) + req := newAuthDispatchRequest(requestedModel, sessionID, headers, count, credentialPolicy, excludedAuthIDs, pinnedAuthID) keyBytes, errMarshal := json.Marshal(&req) if errMarshal != nil { return nil, errMarshal diff --git a/internal/home/client_test.go b/internal/home/client_test.go index ccf2db3f..6677acd5 100644 --- a/internal/home/client_test.go +++ b/internal/home/client_test.go @@ -26,7 +26,7 @@ import ( ) func TestAuthDispatchRequestIncludesCount(t *testing.T) { - req := newAuthDispatchRequest("gpt-5.4", "session-1", http.Header{"Authorization": {"Bearer test"}}, 2, "") + req := newAuthDispatchRequest("gpt-5.4", "session-1", http.Header{"Authorization": {"Bearer test"}}, 2, "", nil, "") raw, err := json.Marshal(&req) if err != nil { @@ -43,10 +43,13 @@ func TestAuthDispatchRequestIncludesCount(t *testing.T) { if got := int(payload["concurrency_protocol"].(float64)); got != 1 { t.Fatalf("concurrency_protocol = %d, want 1", got) } + if _, present := payload["excluded_auth_ids"]; present { + t.Fatalf("legacy request unexpectedly included excluded_auth_ids: %#v", payload["excluded_auth_ids"]) + } } func TestAuthDispatchRequestDefaultsCountToOne(t *testing.T) { - req := newAuthDispatchRequest("gpt-5.4", "", nil, 0, "") + req := newAuthDispatchRequest("gpt-5.4", "", nil, 0, "", nil, "") if req.Count != 1 { t.Fatalf("count = %d, want 1", req.Count) @@ -57,7 +60,7 @@ func TestAuthDispatchRequestDefaultsCountToOne(t *testing.T) { } func TestAuthDispatchRequestIncludesCredentialPolicy(t *testing.T) { - req := newAuthDispatchRequest("gpt-5.4", "", nil, 1, "codex_alpha_search_v1") + req := newAuthDispatchRequest("gpt-5.4", "", nil, 1, "codex_alpha_search_v1", nil, "") raw, errMarshal := json.Marshal(&req) if errMarshal != nil { t.Fatalf("marshal auth dispatch request: %v", errMarshal) @@ -71,6 +74,63 @@ func TestAuthDispatchRequestIncludesCredentialPolicy(t *testing.T) { } } +func TestAuthDispatchRequestIncludesExcludedAuthIDs(t *testing.T) { + excludedAuthIDs := []string{"auth-a", "auth-b"} + req := newAuthDispatchRequest("gpt-5.4", "", nil, 2, "", &excludedAuthIDs, "") + if req.Count != 1 { + t.Fatalf("new retry-contract count = %d, want 1 for legacy Home compatibility", req.Count) + } + raw, errMarshal := json.Marshal(&req) + if errMarshal != nil { + t.Fatalf("marshal auth dispatch request: %v", errMarshal) + } + var payload map[string]any + if errUnmarshal := json.Unmarshal(raw, &payload); errUnmarshal != nil { + t.Fatalf("unmarshal auth dispatch request: %v", errUnmarshal) + } + got, ok := payload["excluded_auth_ids"].([]any) + if !ok || len(got) != 2 || got[0] != "auth-a" || got[1] != "auth-b" { + t.Fatalf("excluded_auth_ids = %#v, want [auth-a auth-b]", payload["excluded_auth_ids"]) + } +} + +func TestAuthDispatchRequestIncludesEmptyExcludedAuthIDs(t *testing.T) { + excludedAuthIDs := []string{} + req := newAuthDispatchRequest("gpt-5.4", "", nil, 2, "", &excludedAuthIDs, "") + if req.Count != 1 { + t.Fatalf("new retry-contract count = %d, want 1 for legacy Home compatibility", req.Count) + } + raw, errMarshal := json.Marshal(&req) + if errMarshal != nil { + t.Fatalf("marshal auth dispatch request: %v", errMarshal) + } + var payload map[string]any + if errUnmarshal := json.Unmarshal(raw, &payload); errUnmarshal != nil { + t.Fatalf("unmarshal auth dispatch request: %v", errUnmarshal) + } + got, ok := payload["excluded_auth_ids"].([]any) + if !ok || len(got) != 0 { + t.Fatalf("excluded_auth_ids = %#v, want []", payload["excluded_auth_ids"]) + } +} + +func TestAuthDispatchRequestIncludesPinnedAuthID(t *testing.T) { + excludedAuthIDs := []string{} + req := newAuthDispatchRequest("gpt-5.4", "", nil, 2, "", &excludedAuthIDs, " auth-pinned ") + + raw, errMarshal := json.Marshal(&req) + if errMarshal != nil { + t.Fatalf("marshal auth dispatch request: %v", errMarshal) + } + var payload map[string]any + if errUnmarshal := json.Unmarshal(raw, &payload); errUnmarshal != nil { + t.Fatalf("unmarshal auth dispatch request: %v", errUnmarshal) + } + if got := payload["pinned_auth_id"]; got != "auth-pinned" { + t.Fatalf("pinned_auth_id = %#v, want auth-pinned", got) + } +} + func TestRedisOptionsHomeTLSDisabled(t *testing.T) { client := New(config.HomeConfig{ Enabled: true, diff --git a/internal/home/requests.go b/internal/home/requests.go index e013eb43..30c7c5ff 100644 --- a/internal/home/requests.go +++ b/internal/home/requests.go @@ -10,6 +10,8 @@ type authDispatchRequest struct { SessionID string `json:"session_id,omitempty"` Headers map[string]string `json:"headers,omitempty"` CredentialPolicy string `json:"credential_policy,omitempty"` + ExcludedAuthIDs *[]string `json:"excluded_auth_ids,omitempty"` + PinnedAuthID string `json:"pinned_auth_id,omitempty"` } type modelsRequest struct { diff --git a/internal/runtime/executor/antigravity_executor_credits.go b/internal/runtime/executor/antigravity_executor_credits.go index 0eff72c7..bb049f04 100644 --- a/internal/runtime/executor/antigravity_executor_credits.go +++ b/internal/runtime/executor/antigravity_executor_credits.go @@ -525,26 +525,6 @@ func (e *AntigravityExecutor) updateAntigravityCreditsBalance(ctx context.Contex return } } -func antigravityRetryAttempts(auth *cliproxyauth.Auth, cfg *config.Config) int { - retry := 0 - if cfg != nil { - retry = cfg.RequestRetry - } - if auth != nil { - if override, ok := auth.RequestRetryOverride(); ok { - retry = override - } - } - if retry < 0 { - retry = 0 - } - attempts := retry + 1 - if attempts < 1 { - return 1 - } - return attempts -} - func antigravityShouldRetryNoCapacity(statusCode int, body []byte) bool { if statusCode != http.StatusServiceUnavailable { return false diff --git a/internal/runtime/executor/antigravity_executor_credits_test.go b/internal/runtime/executor/antigravity_executor_credits_test.go index a763d19f..3223c1bb 100644 --- a/internal/runtime/executor/antigravity_executor_credits_test.go +++ b/internal/runtime/executor/antigravity_executor_credits_test.go @@ -278,27 +278,19 @@ func TestParseRetryDelay_HumanReadableDuration(t *testing.T) { } } -func TestAntigravityExecute_RetriesTransient429ResourceExhausted(t *testing.T) { +func TestAntigravityExecute_DoesNotUseRequestRetryForInternalRetries(t *testing.T) { resetAntigravityCreditsRetryState() t.Cleanup(resetAntigravityCreditsRetryState) var requestCount int server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { requestCount++ - switch requestCount { - case 1: - w.WriteHeader(http.StatusTooManyRequests) - _, _ = w.Write([]byte(`{"error":{"code":429,"message":"Resource has been exhausted (e.g. check quota).","status":"RESOURCE_EXHAUSTED"}}`)) - case 2: - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{"response":{"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]}}],"usageMetadata":{"promptTokenCount":1,"candidatesTokenCount":1,"totalTokenCount":2}}}`)) - default: - t.Fatalf("unexpected request count %d", requestCount) - } + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(`{"error":{"code":429,"message":"Resource has been exhausted (e.g. check quota).","status":"RESOURCE_EXHAUSTED"}}`)) })) defer server.Close() - exec := NewAntigravityExecutor(&config.Config{RequestRetry: 1}) + exec := NewAntigravityExecutor(&config.Config{RequestRetry: 3}) auth := &cliproxyauth.Auth{ ID: "auth-transient-429", Attributes: map[string]string{ @@ -317,14 +309,14 @@ func TestAntigravityExecute_RetriesTransient429ResourceExhausted(t *testing.T) { }, cliproxyexecutor.Options{ SourceFormat: sdktranslator.FormatAntigravity, }) - if err != nil { - t.Fatalf("Execute() error = %v", err) + if err == nil { + t.Fatalf("Execute() error = nil, want upstream 429") } - if len(resp.Payload) == 0 { - t.Fatal("Execute() returned empty payload") + if len(resp.Payload) != 0 { + t.Fatalf("Execute() returned payload %q, want empty payload", resp.Payload) } - if requestCount != 2 { - t.Fatalf("request count = %d, want 2", requestCount) + if requestCount != 1 { + t.Fatalf("request count = %d, want 1", requestCount) } } diff --git a/internal/runtime/executor/antigravity_executor_execute.go b/internal/runtime/executor/antigravity_executor_execute.go index 58419750..bc64a854 100644 --- a/internal/runtime/executor/antigravity_executor_execute.go +++ b/internal/runtime/executor/antigravity_executor_execute.go @@ -85,7 +85,9 @@ func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Au baseURLs := antigravityBaseURLFallbackOrder(auth) httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) httpClient = reporter.TrackHTTPClient(httpClient) - attempts := antigravityRetryAttempts(auth, e.cfg) + // Credential retry rounds are owned by the conductor. Keep one upstream + // attempt per credential so request-retry is not consumed twice. + attempts := 1 attemptLoop: for attempt := 0; attempt < attempts; attempt++ { @@ -313,7 +315,9 @@ func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth * httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) httpClient = reporter.TrackHTTPClient(httpClient) - attempts := antigravityRetryAttempts(auth, e.cfg) + // Credential retry rounds are owned by the conductor. Keep one upstream + // attempt per credential so request-retry is not consumed twice. + attempts := 1 attemptLoop: for attempt := 0; attempt < attempts; attempt++ { diff --git a/internal/runtime/executor/antigravity_executor_stream.go b/internal/runtime/executor/antigravity_executor_stream.go index 9a3f9bb2..96e44d2f 100644 --- a/internal/runtime/executor/antigravity_executor_stream.go +++ b/internal/runtime/executor/antigravity_executor_stream.go @@ -82,7 +82,9 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) httpClient = reporter.TrackHTTPClient(httpClient) - attempts := antigravityRetryAttempts(auth, e.cfg) + // Credential retry rounds are owned by the conductor. Keep one upstream + // attempt per credential so request-retry is not consumed twice. + attempts := 1 attemptLoop: for attempt := 0; attempt < attempts; attempt++ { diff --git a/sdk/cliproxy/auth/conductor_execution.go b/sdk/cliproxy/auth/conductor_execution.go index 53754046..45b405d4 100644 --- a/sdk/cliproxy/auth/conductor_execution.go +++ b/sdk/cliproxy/auth/conductor_execution.go @@ -8,6 +8,7 @@ import ( "io" "net/http" "path/filepath" + "sort" "strconv" "strings" "sync" @@ -61,7 +62,7 @@ func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxye return cliproxyexecutor.Response{}, unwrapRequestStopError(errExec) } lastErr = errExec - wait, shouldRetry := m.shouldRetryAfterError(errExec, attempt, normalized, retryModel, maxWait) + wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errExec, attempt, normalized, retryModel, maxWait, -1) if !shouldRetry { break } @@ -108,7 +109,7 @@ func (m *Manager) ExecuteCount(ctx context.Context, providers []string, req clip return cliproxyexecutor.Response{}, unwrapRequestStopError(errExec) } lastErr = errExec - wait, shouldRetry := m.shouldRetryAfterError(errExec, attempt, normalized, retryModel, maxWait) + wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errExec, attempt, normalized, retryModel, maxWait, -1) if !shouldRetry { break } @@ -139,23 +140,44 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli _, maxRetryCredentials, maxWait := m.retrySettings() var lastErr error + homeRetryLimit := -1 retryModel := authSelectionModelFromOptions(opts, req.Model) - for attempt := 0; ; attempt++ { - result, errStream := m.executeStreamMixedOnce(ctx, normalized, req, opts, maxRetryCredentials) + attempt := 0 + retryRoundPending := false + retryRoundWaited := false + for { + result, errStream := m.executeStreamMixedOnce(ctx, normalized, req, opts, maxRetryCredentials, &homeRetryLimit) if errStream == nil { return result, nil } + if m.HomeEnabled() && retryRoundPending { + if wait, okWait := pendingHomeRetryRoundDelay(errStream, maxWait, &homeRetryLimit, pinnedAuthIDFromMetadata(opts.Metadata) == ""); okWait && m.homeRetryAllowed(attempt-1, homeRetryLimit) { + if retryRoundWaited { + return nil, errStream + } + if errWait := waitForCooldown(ctx, wait, maxWait); errWait != nil { + return nil, errWait + } + retryRoundWaited = true + continue + } + } + retryRoundPending = false + retryRoundWaited = false if isRequestTerminatedError(errStream) || isRequestStopError(errStream) { return nil, unwrapRequestStopError(errStream) } lastErr = errStream - wait, shouldRetry := m.shouldRetryAfterError(errStream, attempt, normalized, retryModel, maxWait) + wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errStream, attempt, normalized, retryModel, maxWait, homeRetryLimit) if !shouldRetry { break } if errWait := waitForCooldown(ctx, wait, maxWait); errWait != nil { return nil, errWait } + attempt++ + retryRoundPending = m.HomeEnabled() + retryRoundWaited = false } if lastErr != nil { lastErr = unwrapRequestStopError(lastErr) @@ -168,7 +190,7 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli } var bootstrapErr *streamBootstrapError if errors.As(lastErr, &bootstrapErr) && bootstrapErr != nil { - return streamErrorResult(bootstrapErr.Headers(), bootstrapErr.cause), nil + return streamErrorResult(bootstrapErr.Headers(), lastErr), nil } return nil, lastErr } @@ -291,7 +313,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req attempted := make(map[string]struct{}) var lastErr error for { - if !homeMode && maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials { + if maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials { if lastErr != nil { return cliproxyexecutor.Response{}, lastErr } @@ -300,6 +322,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req pickOpts := opts if homeMode { pickOpts = withHomeAuthCount(opts, homeAuthCount) + pickOpts = withHomeExcludedAuthIDs(pickOpts, tried) } auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried) if errPick != nil { @@ -458,7 +481,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, attempted := make(map[string]struct{}) var lastErr error for { - if !homeMode && maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials { + if maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials { if lastErr != nil { return cliproxyexecutor.Response{}, lastErr } @@ -467,6 +490,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, pickOpts := opts if homeMode { pickOpts = withHomeAuthCount(opts, homeAuthCount) + pickOpts = withHomeExcludedAuthIDs(pickOpts, tried) } auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried) if errPick != nil { @@ -616,7 +640,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, } } -func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int) (*cliproxyexecutor.StreamResult, error) { +func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int, homeRetryLimit *int) (*cliproxyexecutor.StreamResult, error) { if len(providers) == 0 { return nil, &Error{Code: "provider_not_found", Message: "no provider supplied"} } @@ -627,12 +651,21 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string homeMode := m.HomeEnabled() homeAuthCount := 1 tried := make(map[string]struct{}) + homeExcludedAuthIDs := make(map[string]struct{}) + homeSameAuthRetries := make(map[string]int) + lastHomeAuthID := "" + homeSameAuthRetryPending := false attempted := make(map[string]struct{}) unauthorizedRefreshTried := make(map[string]struct{}) var lastErr error + var roundTiming homeRetryRoundTiming for { - if !homeMode && maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials { + allowSameAuthRetry := homeMode && homeSameAuthRetryPending && lastHomeAuthID != "" && homeSameAuthRetries[lastHomeAuthID] == 0 + if maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials && !allowSameAuthRetry { if lastErr != nil { + if homeMode { + return nil, markHomeRetryRoundExhausted(lastErr, roundTiming.RetryAfter(), true) + } return nil, lastErr } return nil, &Error{Code: "auth_not_found", Message: "no auth available"} @@ -640,6 +673,7 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string pickOpts := opts if homeMode { pickOpts = withHomeAuthCount(opts, homeAuthCount) + pickOpts = withHomeExcludedAuthIDs(pickOpts, homeExcludedAuthIDs) } var selection *HomeDispatchSelection @@ -658,7 +692,15 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string auth, executor, provider, errPick = m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried) } if errPick != nil { + var homeCooldown *homeDispatchRetryAfterError + if homeMode && lastErr != nil && errors.As(errPick, &homeCooldown) && homeCooldown != nil { + observeHomeCooldownRetryLimit(homeCooldown, homeRetryLimit, pinnedAuthIDFromMetadata(opts.Metadata) == "") + return nil, markHomeRetryRoundExhausted(lastErr, homeCooldown.RetryAfter(), false) + } if shouldReturnLastErrorOnPickFailure(homeMode, lastErr, errPick) { + if homeMode { + return nil, markHomeRetryRoundExhausted(lastErr, roundTiming.RetryAfter(), isHomeNextRoundImmediatelyAvailable(errPick)) + } return nil, lastErr } return nil, errPick @@ -669,13 +711,56 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string } return nil, &Error{Code: "executor_not_found", Message: "executor not registered"} } + if homeMode { + m.observeHomeRetryLimit(auth, selection, homeRetryLimit) + } + if selection != nil && allowSameAuthRetry && maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials && auth.ID != lastHomeAuthID { + if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "max_retry_credentials"); errEnd != nil { + return nil, errEnd + } + if lastErr != nil { + return nil, markHomeRetryRoundExhausted(lastErr, roundTiming.RetryAfter(), true) + } + return nil, &Error{Code: "auth_not_found", Message: "no auth available"} + } + if homeMode && lastHomeAuthID != "" && auth.ID != lastHomeAuthID { + homeSameAuthRetryPending = false + } if selection != nil { + // A legacy Home may ignore excluded_auth_ids and return the same + // credential again. Reject credentials explicitly excluded from this + // round while retaining the explicit same-auth retry path, which + // intentionally leaves the credential out of homeExcludedAuthIDs. + if _, alreadyTried := tried[auth.ID]; alreadyTried { + if _, excluded := homeExcludedAuthIDs[auth.ID]; excluded { + if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "repeated_excluded_auth"); errEnd != nil { + return nil, errEnd + } + if lastErr != nil { + return nil, markHomeRetryRoundExhausted(lastErr, roundTiming.RetryAfter(), false) + } + return nil, repeatedHomeAuthError() + } else { + homeSameAuthRetries[auth.ID]++ + if homeSameAuthRetries[auth.ID] > 1 { + // A fresh Home selection may retry the same auth once for + // connection lifecycle or authorization recovery. Repeated + // failures must still rotate away from this credential. + homeExcludedAuthIDs[auth.ID] = struct{}{} + if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "repeated_same_auth"); errEnd != nil { + return nil, errEnd + } + continue + } + } + } if _, refreshedAlready := unauthorizedRefreshTried[auth.ID]; refreshedAlready { - selection.End("repeated_refresh_auth") - if lastErr != nil { - return nil, lastErr + homeExcludedAuthIDs[auth.ID] = struct{}{} + homeSameAuthRetryPending = false + if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "repeated_refresh_auth"); errEnd != nil { + return nil, errEnd } - return nil, repeatedHomeAuthError() + continue } } @@ -712,6 +797,9 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string } if len(models) == 0 { if selection != nil { + homeExcludedAuthIDs[auth.ID] = struct{}{} + lastHomeAuthID = auth.ID + homeSameAuthRetryPending = false releaseAttempt() if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "no_execution_models"); errEnd != nil { return nil, errEnd @@ -727,6 +815,17 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string auth, errPrepare = m.prepareRequestAuth(execCtx, executor, auth) } if errPrepare != nil { + if selection != nil { + excludeAuth := shouldExcludeHomeAuthAfterStreamError(execCtx, auth, errPrepare) + if _, refreshedAlready := unauthorizedRefreshTried[auth.ID]; refreshedAlready || homeSameAuthRetries[auth.ID] > 0 { + excludeAuth = true + } + if excludeAuth { + homeExcludedAuthIDs[auth.ID] = struct{}{} + } + lastHomeAuthID = auth.ID + homeSameAuthRetryPending = !excludeAuth + } if selection == nil { if errCancel := claudeOAuthRequestCancellation(execCtx, auth, errPrepare); errCancel != nil { return nil, errCancel @@ -740,6 +839,9 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string m.MarkResult(execCtx, result) } lastErr = errPrepare + if homeMode { + roundTiming.Observe(lastErr) + } if selection != nil { if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "prepare_failed"); errEnd != nil { return nil, errEnd @@ -762,6 +864,17 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string } streamResult, errStream := m.executeStreamWithModelPool(execCtx, executor, auth, provider, execReq, execOpts, routeModel, streamExecutionModel, models, pooled, aliasResult, routing, !homeMode || selection != nil, selection != nil, unauthorizedRefreshTried) if errStream != nil { + if selection != nil { + excludeAuth := shouldExcludeHomeAuthAfterStreamError(execCtx, auth, errStream) + if _, refreshedAlready := unauthorizedRefreshTried[auth.ID]; refreshedAlready || homeSameAuthRetries[auth.ID] > 0 { + excludeAuth = true + } + if excludeAuth { + homeExcludedAuthIDs[auth.ID] = struct{}{} + } + lastHomeAuthID = auth.ID + homeSameAuthRetryPending = !excludeAuth + } if selection != nil { releaseAttempt() if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "stream_start_failed"); errEnd != nil { @@ -777,6 +890,9 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string return nil, wrapRequestStopError(errStream) } lastErr = errStream + if homeMode { + roundTiming.Observe(lastErr) + } if homeMode { homeAuthCount++ } @@ -786,6 +902,9 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string return nil, errStream } lastErr = errStream + if homeMode { + roundTiming.Observe(lastErr) + } if homeMode { homeAuthCount++ } @@ -801,6 +920,19 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string } } +func shouldExcludeHomeAuthAfterStreamError(ctx context.Context, auth *Auth, err error) bool { + if err == nil || isConnectionLifecycleError(err) { + return false + } + // A 426 during a downstream websocket attempt is a transport fallback + // signal. OAuth authorization failures may also recover after a refresh. + // Both paths may retry the same credential once. + if cliproxyexecutor.DownstreamWebsocket(ctx) && statusCodeFromError(err) == http.StatusUpgradeRequired { + return false + } + return !isUnauthorizedError(err) || auth == nil || auth.AuthKind() != AuthKindOAuth +} + func cloneRequestMetadata(src map[string]any) map[string]any { if len(src) == 0 { return make(map[string]any, 4) @@ -872,6 +1004,34 @@ func withHomeAuthCount(opts cliproxyexecutor.Options, count int) cliproxyexecuto return opts } +func withHomeExcludedAuthIDs(opts cliproxyexecutor.Options, tried map[string]struct{}) cliproxyexecutor.Options { + meta := make(map[string]any, len(opts.Metadata)+1) + for key, value := range opts.Metadata { + meta[key] = value + } + excluded := make(map[string]struct{}) + for _, authID := range homeExcludedAuthIDsFromMetadata(meta) { + excluded[authID] = struct{}{} + } + for authID := range tried { + if authID = strings.TrimSpace(authID); authID != "" { + excluded[authID] = struct{}{} + } + } + if len(excluded) == 0 { + delete(meta, ExcludedAuthIDsMetadataKey) + } else { + ids := make([]string, 0, len(excluded)) + for authID := range excluded { + ids = append(ids, authID) + } + sort.Strings(ids) + meta[ExcludedAuthIDsMetadataKey] = ids + } + opts.Metadata = meta + return opts +} + func homeAuthCountFromMetadata(meta map[string]any) int { if len(meta) == 0 { return 1 @@ -893,6 +1053,56 @@ func homeAuthCountFromMetadata(meta map[string]any) int { return 1 } +func homeExcludedAuthIDsFromMetadata(meta map[string]any) []string { + if len(meta) == 0 { + return nil + } + raw, ok := meta[ExcludedAuthIDsMetadataKey] + if !ok { + return nil + } + seen := make(map[string]struct{}) + ids := make([]string, 0) + appendID := func(value string) { + value = strings.TrimSpace(value) + if value == "" { + return + } + if _, exists := seen[value]; exists { + return + } + seen[value] = struct{}{} + ids = append(ids, value) + } + switch values := raw.(type) { + case []string: + for _, value := range values { + appendID(value) + } + case []any: + for _, value := range values { + if text, okText := value.(string); okText { + appendID(text) + } + } + case map[string]struct{}: + for value := range values { + appendID(value) + } + case map[string]bool: + for value, enabled := range values { + if enabled { + appendID(value) + } + } + } + if len(ids) == 0 { + return nil + } + sort.Strings(ids) + return ids +} + func hasRequestedModelMetadata(meta map[string]any) bool { if len(meta) == 0 { return false diff --git a/sdk/cliproxy/auth/conductor_home.go b/sdk/cliproxy/auth/conductor_home.go index 7825f91a..92721e16 100644 --- a/sdk/cliproxy/auth/conductor_home.go +++ b/sdk/cliproxy/auth/conductor_home.go @@ -20,6 +20,9 @@ import ( const ( homeAuthCountMetadataKey = "__cliproxy_home_auth_count" + // ExcludedAuthIDsMetadataKey stores credential IDs already attempted in the + // current request retry round. + ExcludedAuthIDsMetadataKey = "excluded_auth_ids" // CloseAllExecutionSessionsID asks an executor to release all active execution sessions. // Executors that do not support this marker may ignore it. CloseAllExecutionSessionsID = "__all_execution_sessions__" @@ -120,6 +123,139 @@ type homeErrorDetail struct { Code string `json:"code,omitempty"` Retryable bool `json:"retryable,omitempty"` RetryAfterMS int64 `json:"retry_after_ms,omitempty"` + RequestRetry *int `json:"request_retry,omitempty"` +} + +type homeDispatchRetryAfterError struct { + cause *Error + retryAfter time.Duration + requestRetry int + hasRequestRetry bool +} + +// homeRetryRoundExhaustedError marks a terminal error produced after the +// current Home credential round has been exhausted. The wrapped error retains +// its status and retry-after metadata for the outer request retry policy. +type homeRetryRoundExhaustedError struct { + cause error + retryAfter time.Duration + hasRetryAfter bool + retryNow bool +} + +func (e *homeRetryRoundExhaustedError) Error() string { + if e == nil || e.cause == nil { + return "" + } + return e.cause.Error() +} + +func (e *homeRetryRoundExhaustedError) Unwrap() error { + if e == nil { + return nil + } + return e.cause +} + +func (e *homeRetryRoundExhaustedError) RetryAfter() *time.Duration { + if e == nil || !e.hasRetryAfter { + return nil + } + value := e.retryAfter + return &value +} + +func markHomeRetryRoundExhausted(err error, retryAfter *time.Duration, retryNow bool) error { + if err == nil { + return nil + } + marked := &homeRetryRoundExhaustedError{cause: err, retryNow: retryNow} + if retryAfter != nil { + marked.retryAfter = *retryAfter + marked.hasRetryAfter = true + } + return marked +} + +func isHomeRetryRoundExhausted(err error) bool { + if err == nil { + return false + } + var marker *homeRetryRoundExhaustedError + return errors.As(err, &marker) && marker != nil +} + +type homeRetryRoundTiming struct { + retryAfter time.Duration + immediate bool + invalid bool +} + +func (t *homeRetryRoundTiming) Observe(err error) { + if t == nil || err == nil || t.immediate || t.invalid { + return + } + retryAfter := retryAfterFromError(err) + if retryAfter == nil { + return + } + if *retryAfter == 0 { + t.retryAfter = 0 + t.immediate = true + return + } + if *retryAfter < 0 { + t.retryAfter = *retryAfter + t.invalid = true + return + } + if t.retryAfter <= 0 || *retryAfter < t.retryAfter { + t.retryAfter = *retryAfter + } +} + +func (t *homeRetryRoundTiming) RetryAfter() *time.Duration { + if t == nil || t.immediate || (!t.invalid && t.retryAfter <= 0) { + return nil + } + value := t.retryAfter + return &value +} + +func (e *homeDispatchRetryAfterError) Error() string { + if e == nil || e.cause == nil { + return "" + } + return e.cause.Error() +} + +func (e *homeDispatchRetryAfterError) Unwrap() error { + if e == nil { + return nil + } + return e.cause +} + +func (e *homeDispatchRetryAfterError) StatusCode() int { + if e == nil || e.cause == nil { + return 0 + } + return e.cause.HTTPStatus +} + +func (e *homeDispatchRetryAfterError) RetryAfter() *time.Duration { + if e == nil || e.retryAfter <= 0 { + return nil + } + value := e.retryAfter + return &value +} + +func (e *homeDispatchRetryAfterError) RequestRetryLimit() (int, bool) { + if e == nil || !e.hasRequestRetry { + return 0, false + } + return e.requestRetry, true } const ( @@ -159,6 +295,30 @@ func shouldReturnLastErrorOnPickFailure(homeMode bool, lastErr error, errPick er } } +func isHomeNextRoundImmediatelyAvailable(err error) bool { + var authErr *Error + if !errors.As(err, &authErr) || authErr == nil { + return false + } + return strings.EqualFold(strings.TrimSpace(authErr.Code), "auth_unavailable") +} + +func pendingHomeRetryRoundDelay(err error, maxWait time.Duration, retryLimit *int, acceptRemoteRetryLimit bool) (time.Duration, bool) { + if err == nil || isHomeRetryRoundExhausted(err) { + return 0, false + } + var homeCooldown *homeDispatchRetryAfterError + if !errors.As(err, &homeCooldown) || homeCooldown == nil { + return 0, false + } + observeHomeCooldownRetryLimit(homeCooldown, retryLimit, acceptRemoteRetryLimit) + retryAfter := homeCooldown.RetryAfter() + if retryAfter == nil || *retryAfter <= 0 || maxWait <= 0 || *retryAfter > maxWait { + return 0, false + } + return *retryAfter, true +} + func homeAuthAlreadyTried(tried map[string]struct{}, authID string) bool { authID = strings.TrimSpace(authID) if authID == "" || len(tried) == 0 { @@ -181,6 +341,7 @@ type homeAuthDispatchResponse struct { Provider string `json:"provider"` AuthIndex string `json:"auth_index"` UserAPIKey string `json:"user_api_key"` + RequestRetry *int `json:"request_retry,omitempty"` ForceMapping bool `json:"force_mapping"` OriginalAlias string `json:"original_alias"` Auth Auth `json:"auth"` @@ -192,10 +353,18 @@ type homeAuthDispatcher interface { AbortAmbiguousDispatch() } +type homeDispatchConstraintsDispatcher interface { + RPopAuthWithConstraints(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int, excludedAuthIDs []string, pinnedAuthID string) ([]byte, error) +} + type homeCredentialPolicyDispatcher interface { RPopAuthWithPolicy(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int, credentialPolicy string) ([]byte, error) } +type homeCredentialPolicyConstraintsDispatcher interface { + RPopAuthWithPolicyAndConstraints(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int, credentialPolicy string, excludedAuthIDs []string, pinnedAuthID string) ([]byte, error) +} + var currentHomeDispatcher = func() homeAuthDispatcher { return home.Current() } @@ -325,7 +494,7 @@ func (m *Manager) lockHomeWebsocketSession(ctx context.Context, opts cliproxyexe return mutex.Unlock } -func (m *Manager) retainedHomeSessionSelection(ctx context.Context, opts cliproxyexecutor.Options, model string) (*HomeDispatchSelection, bool, error) { +func (m *Manager) retainedHomeSessionSelection(ctx context.Context, opts cliproxyexecutor.Options, model string, excludedAuthIDs map[string]struct{}) (*HomeDispatchSelection, bool, error) { if m == nil || !cliproxyexecutor.DownstreamWebsocket(ctx) { return nil, false, nil } @@ -348,7 +517,8 @@ func (m *Manager) retainedHomeSessionSelection(ctx context.Context, opts cliprox } matchesCredential := credentialID == "" || key.credentialID == credentialID matchesRoute := validRouteModel && key.routeModel == routeModel - if !fallbackAttempt && matchesCredential && selection.Active() && matchesRoute && retained == nil { + _, excluded := excludedAuthIDs[strings.TrimSpace(key.credentialID)] + if !fallbackAttempt && !excluded && matchesCredential && selection.Active() && matchesRoute && retained == nil { retained = selection continue } @@ -711,7 +881,7 @@ func (m *Manager) pickNextViaHome(ctx context.Context, model string, opts clipro if ctx == nil { ctx = context.Background() } - selection, errSelection := m.pickHomeDispatchSelection(ctx, model, opts) + selection, errSelection := m.pickHomeDispatchSelection(ctx, model, withHomeExcludedAuthIDs(opts, tried)) if errSelection != nil { return nil, nil, "", errSelection } @@ -739,7 +909,13 @@ func (m *Manager) pickHomeDispatchSelection(ctx context.Context, model string, o if requestedModel == "" { requestedModel = requestedModelFromMetadata(opts.Metadata, model) } - retained, retainedOK, errRetained := m.retainedHomeSessionSelection(ctx, opts, requestedModel) + pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata) + excludedAuthIDList := homeExcludedAuthIDsFromMetadata(opts.Metadata) + excludedAuthIDs := make(map[string]struct{}, len(excludedAuthIDList)) + for _, authID := range excludedAuthIDList { + excludedAuthIDs[authID] = struct{}{} + } + retained, retainedOK, errRetained := m.retainedHomeSessionSelection(ctx, opts, requestedModel, excludedAuthIDs) if errRetained != nil { return nil, errRetained } @@ -747,8 +923,8 @@ func (m *Manager) pickHomeDispatchSelection(ctx context.Context, model string, o return retained, nil } if sessionID := homeExecutionSessionIDFromMetadata(opts.Metadata); sessionID != "" { - if credentialID := pinnedAuthIDFromMetadata(opts.Metadata); credentialID != "" { - if errEnd := m.endMismatchedHomeSessionSelections(ctx, sessionID, credentialID, requestedModel, true); errEnd != nil { + if pinnedAuthID != "" { + if errEnd := m.endMismatchedHomeSessionSelections(ctx, sessionID, pinnedAuthID, requestedModel, true); errEnd != nil { return nil, errEnd } } @@ -763,6 +939,11 @@ func (m *Manager) pickHomeDispatchSelection(ctx context.Context, model string, o if !client.HeartbeatOK() { return nil, &Error{Code: "home_unavailable", Message: "home control center unavailable", HTTPStatus: http.StatusServiceUnavailable} } + if pinnedAuthID != "" { + if _, excluded := excludedAuthIDs[pinnedAuthID]; excluded { + return nil, &Error{Code: "auth_not_found", Message: "pinned auth is unavailable in the current retry round", HTTPStatus: http.StatusServiceUnavailable} + } + } pending, errBegin := registry.BeginDispatch() if errBegin != nil { return nil, &Error{Code: "home_unavailable", Message: "home execution registry unavailable", Retryable: true, HTTPStatus: http.StatusServiceUnavailable} @@ -774,9 +955,17 @@ func (m *Manager) pickHomeDispatchSelection(ctx context.Context, model string, o var raw []byte var errRPop error if credentialPolicy == "" { - raw, errRPop = client.RPopAuth(ctx, requestedModel, sessionID, dispatchHeaders, homeAuthCountFromMetadata(opts.Metadata)) + if constrainedClient, okConstraints := client.(homeDispatchConstraintsDispatcher); okConstraints { + raw, errRPop = constrainedClient.RPopAuthWithConstraints(ctx, requestedModel, sessionID, dispatchHeaders, homeAuthCountFromMetadata(opts.Metadata), excludedAuthIDList, pinnedAuthID) + } else { + raw, errRPop = client.RPopAuth(ctx, requestedModel, sessionID, dispatchHeaders, homeAuthCountFromMetadata(opts.Metadata)) + } } else if policyClient, okPolicy := client.(homeCredentialPolicyDispatcher); okPolicy { - raw, errRPop = policyClient.RPopAuthWithPolicy(ctx, requestedModel, sessionID, dispatchHeaders, homeAuthCountFromMetadata(opts.Metadata), credentialPolicy) + if constrainedClient, okConstraints := client.(homeCredentialPolicyConstraintsDispatcher); okConstraints { + raw, errRPop = constrainedClient.RPopAuthWithPolicyAndConstraints(ctx, requestedModel, sessionID, dispatchHeaders, homeAuthCountFromMetadata(opts.Metadata), credentialPolicy, excludedAuthIDList, pinnedAuthID) + } else { + raw, errRPop = policyClient.RPopAuthWithPolicy(ctx, requestedModel, sessionID, dispatchHeaders, homeAuthCountFromMetadata(opts.Metadata), credentialPolicy) + } } else { pending.End() return nil, &Error{Code: "home_unavailable", Message: "home dispatcher does not support credential policies", HTTPStatus: http.StatusServiceUnavailable} @@ -887,6 +1076,10 @@ func (m *Manager) pickHomeDispatchSelection(ctx context.Context, model string, o endScope() return nil, &Error{Code: "invalid_auth", Message: "home returned auth without id", HTTPStatus: http.StatusBadGateway} } + if pinnedAuthID != "" && strings.TrimSpace(auth.ID) != pinnedAuthID { + endScope() + return nil, &Error{Code: "auth_not_found", Message: "home returned an auth that does not match the pinned credential", HTTPStatus: http.StatusServiceUnavailable} + } if errIdentity := verifyAccountedHomeConcurrencyIdentity(envelope.Tuple, &auth, dispatch.AuthIndex); errIdentity != nil { endScope() return nil, errIdentity @@ -935,6 +1128,10 @@ func (m *Manager) pickHomeDispatchSelection(ctx context.Context, model string, o endScope() return nil, &Error{Code: "home_unavailable", Message: "home execution registry unavailable", Retryable: true, HTTPStatus: http.StatusServiceUnavailable} } + if pinnedAuthID == "" && dispatch.RequestRetry != nil && *dispatch.RequestRetry >= 0 { + selection.requestRetry = *dispatch.RequestRetry + selection.hasRequestRetry = true + } if envelope.Present { selection.accountedModel = envelope.Tuple.Model } diff --git a/sdk/cliproxy/auth/conductor_home_execution.go b/sdk/cliproxy/auth/conductor_home_execution.go index 2e3eec65..204e7a32 100644 --- a/sdk/cliproxy/auth/conductor_home_execution.go +++ b/sdk/cliproxy/auth/conductor_home_execution.go @@ -2,6 +2,7 @@ package auth import ( "context" + "errors" "fmt" "sync" "time" @@ -14,17 +15,74 @@ func (m *Manager) executeHome(ctx context.Context, providers []string, req clipr if unlockSession := m.lockHomeWebsocketSession(ctx, opts); unlockSession != nil { defer unlockSession() } + _, maxRetryCredentials, maxWait := m.retrySettings() + retryModel := authSelectionModelFromOptions(opts, req.Model) + homeRetryLimit := -1 + attempt := 0 + retryRoundPending := false + retryRoundWaited := false + for { + response, errExecute := m.executeHomeOnce(ctx, providers, req, opts, countTokens, maxRetryCredentials, &homeRetryLimit) + if errExecute == nil { + return response, nil + } + if retryRoundPending { + if wait, okWait := pendingHomeRetryRoundDelay(errExecute, maxWait, &homeRetryLimit, pinnedAuthIDFromMetadata(opts.Metadata) == ""); okWait && m.homeRetryAllowed(attempt-1, homeRetryLimit) { + if retryRoundWaited { + return cliproxyexecutor.Response{}, errExecute + } + if errWait := waitForCooldown(ctx, wait, maxWait); errWait != nil { + return cliproxyexecutor.Response{}, errWait + } + retryRoundWaited = true + continue + } + } + retryRoundPending = false + retryRoundWaited = false + if isRequestTerminatedError(errExecute) || isRequestStopError(errExecute) { + return cliproxyexecutor.Response{}, unwrapRequestStopError(errExecute) + } + wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errExecute, attempt, providers, retryModel, maxWait, homeRetryLimit) + if !shouldRetry { + return cliproxyexecutor.Response{}, unwrapRequestStopError(errExecute) + } + if errWait := waitForCooldown(ctx, wait, maxWait); errWait != nil { + return cliproxyexecutor.Response{}, errWait + } + attempt++ + retryRoundPending = true + retryRoundWaited = false + } +} + +func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, countTokens bool, maxRetryCredentials int, homeRetryLimit *int) (cliproxyexecutor.Response, error) { routeModel := authSelectionModelFromOptions(opts, req.Model) responseAlias := requestedModelAliasFromOptions(opts, routeModel) executionModel, restoreExecutionModel := executionModelForAuthSelection(opts, req.Model) opts = ensureRequestedModelMetadata(opts, routeModel) tried := make(map[string]struct{}) + attempted := make(map[string]struct{}) var lastErr error + var roundTiming homeRetryRoundTiming for homeAuthCount := 1; ; homeAuthCount++ { - selection, errSelection := m.pickHomeDispatchSelection(ctx, routeModel, withHomeAuthCount(opts, homeAuthCount)) + if maxRetryCredentials > 0 && len(attempted) >= maxRetryCredentials { + if lastErr != nil { + return cliproxyexecutor.Response{}, markHomeRetryRoundExhausted(lastErr, roundTiming.RetryAfter(), true) + } + return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"} + } + pickOpts := withHomeAuthCount(opts, homeAuthCount) + pickOpts = withHomeExcludedAuthIDs(pickOpts, tried) + selection, errSelection := m.pickHomeDispatchSelection(ctx, routeModel, pickOpts) if errSelection != nil { + var homeCooldown *homeDispatchRetryAfterError + if lastErr != nil && errors.As(errSelection, &homeCooldown) && homeCooldown != nil { + observeHomeCooldownRetryLimit(homeCooldown, homeRetryLimit, pinnedAuthIDFromMetadata(opts.Metadata) == "") + return cliproxyexecutor.Response{}, markHomeRetryRoundExhausted(lastErr, homeCooldown.RetryAfter(), false) + } if shouldReturnLastErrorOnPickFailure(true, lastErr, errSelection) { - return cliproxyexecutor.Response{}, lastErr + return cliproxyexecutor.Response{}, markHomeRetryRoundExhausted(lastErr, roundTiming.RetryAfter(), isHomeNextRoundImmediatelyAvailable(errSelection)) } return cliproxyexecutor.Response{}, errSelection } @@ -33,13 +91,18 @@ func (m *Manager) executeHome(ctx context.Context, providers []string, req clipr selection.End("missing_execution_target") return cliproxyexecutor.Response{}, &Error{Code: "executor_not_found", Message: "executor not registered"} } + m.observeHomeRetryLimit(auth, selection, homeRetryLimit) if _, seen := tried[auth.ID]; seen { - selection.End("repeated_auth") + if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "repeated_auth"); errEnd != nil { + return cliproxyexecutor.Response{}, errEnd + } if lastErr != nil { - return cliproxyexecutor.Response{}, lastErr + return cliproxyexecutor.Response{}, markHomeRetryRoundExhausted(lastErr, roundTiming.RetryAfter(), false) } return cliproxyexecutor.Response{}, repeatedHomeAuthError() } + tried[auth.ID] = struct{}{} + attempted[auth.ID] = struct{}{} entry := logEntryWithRequestID(ctx) debugLogAuthSelection(entry, auth, selection.Provider, routeModel) if errRuntimeAuth := m.bindHomeSelectionRuntimeAuth(ctx, opts, selection); errRuntimeAuth != nil { @@ -47,7 +110,6 @@ func (m *Manager) executeHome(ctx context.Context, providers []string, req clipr return cliproxyexecutor.Response{}, errRuntimeAuth } publishSelectedAuthMetadata(opts.Metadata, auth) - tried[auth.ID] = struct{}{} execCtx, releaseAttempt, errBind := homeExecutionAttemptContext(ctx, selection) if errBind != nil { selection.End("attempt_bind_failed") @@ -73,6 +135,7 @@ func (m *Manager) executeHome(ctx context.Context, providers []string, req clipr return cliproxyexecutor.Response{}, errEnd } lastErr = &Error{Code: "auth_not_found", Message: "no execution models available"} + roundTiming.Observe(lastErr) continue } preparedAuth, errPrepare := m.prepareHomeRequestAuth(execCtx, selection.Executor, selection) @@ -83,6 +146,7 @@ func (m *Manager) executeHome(ctx context.Context, providers []string, req clipr return cliproxyexecutor.Response{}, errEnd } lastErr = errPrepare + roundTiming.Observe(lastErr) continue } didRefreshOnUnauthorized := false @@ -216,6 +280,7 @@ func (m *Manager) executeHome(ctx context.Context, providers []string, req clipr break } } + roundTiming.Observe(lastErr) releaseAttempt() if errEnd := m.endHomeSelectionBeforeRedispatch(ctx, selection, "execution_failed"); errEnd != nil { return cliproxyexecutor.Response{}, errEnd diff --git a/sdk/cliproxy/auth/conductor_lifecycle.go b/sdk/cliproxy/auth/conductor_lifecycle.go index bc9c81d6..bcf7846c 100644 --- a/sdk/cliproxy/auth/conductor_lifecycle.go +++ b/sdk/cliproxy/auth/conductor_lifecycle.go @@ -10,7 +10,7 @@ import ( internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" ) -// SetRetryConfig updates retry attempts, credential retry limit and cooldown wait interval. +// SetRetryConfig updates additional credential retry rounds, the per-round credential limit, and the cooldown wait interval. func (m *Manager) SetRetryConfig(retry int, maxRetryInterval time.Duration, maxRetryCredentials int) { if m == nil { return diff --git a/sdk/cliproxy/auth/conductor_overrides_test.go b/sdk/cliproxy/auth/conductor_overrides_test.go index 52f2c05b..69b51406 100644 --- a/sdk/cliproxy/auth/conductor_overrides_test.go +++ b/sdk/cliproxy/auth/conductor_overrides_test.go @@ -36,9 +36,12 @@ func TestManager_ShouldRetryAfterError_RespectsAuthRequestRetryOverride(t *testi Unavailable: true, Status: StatusError, NextRetryAfter: next, + LastError: &Error{HTTPStatus: http.StatusInternalServerError, Message: "upstream unavailable"}, }, }, } + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) if _, errRegister := m.Register(context.Background(), auth); errRegister != nil { t.Fatalf("register auth: %v", errRegister) } @@ -83,6 +86,255 @@ func TestManager_ShouldRetryAfterError_SkipsWrappedHomeConcurrencyBusy(t *testin } } +func TestManager_ShouldRetryAfterError_RetriesLocalRoundWithoutCooldown(t *testing.T) { + m := NewManager(nil, nil, nil) + m.SetRetryConfig(1, 0, 0) + model := "gpt-retry-without-cooldown-" + uuid.NewString() + registry.GetGlobalRegistry().RegisterClient("retry-auth", "codex", []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient("retry-auth") }) + if _, errRegister := m.Register(context.Background(), &Auth{ID: "retry-auth", Provider: "codex"}); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + for _, status := range []int{http.StatusTooManyRequests, http.StatusBadGateway} { + wait, shouldRetry := m.shouldRetryAfterError(&Error{HTTPStatus: status, Message: "retryable failure"}, 0, []string{"codex"}, model, 0) + if !shouldRetry || wait != 0 { + t.Fatalf("status %d retry = (%v, %t), want (0, true)", status, wait, shouldRetry) + } + if _, shouldRetry = m.shouldRetryAfterError(&Error{HTTPStatus: status, Message: "retryable failure"}, 1, []string{"codex"}, model, 0); shouldRetry { + t.Fatalf("status %d retried after the configured additional round", status) + } + } +} + +func TestManager_ShouldRetryAfterError_DoesNotWaitWhenAnotherCredentialIsAvailable(t *testing.T) { + m := NewManager(nil, nil, nil) + m.SetRetryConfig(1, time.Minute, 1) + model := "retry-available-credential-" + uuid.NewString() + next := time.Now().Add(30 * time.Second) + auths := []*Auth{ + { + ID: "cooling-" + uuid.NewString(), + Provider: "codex", + ModelStates: map[string]*ModelState{ + model: {Unavailable: true, Status: StatusError, NextRetryAfter: next}, + }, + }, + {ID: "available-" + uuid.NewString(), Provider: "codex"}, + } + for _, auth := range auths { + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + if _, errRegister := m.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("register auth %s: %v", auth.ID, errRegister) + } + } + + wait, shouldRetry := m.shouldRetryAfterError(&Error{HTTPStatus: http.StatusTooManyRequests, Message: "rate limited"}, 0, []string{"codex"}, model, time.Minute) + if !shouldRetry || wait != 0 { + t.Fatalf("retry with available credential = (%v, %t), want immediate retry", wait, shouldRetry) + } +} + +func TestManager_ShouldRetryAfterError_IgnoresUnrelatedModelOverride(t *testing.T) { + m := NewManager(nil, nil, nil) + m.SetRetryConfig(0, 0, 0) + targetModel := "retry-target-" + uuid.NewString() + unrelatedModel := "retry-unrelated-" + uuid.NewString() + registryRef := registry.GetGlobalRegistry() + registryRef.RegisterClient("target-auth", "codex", []*registry.ModelInfo{{ID: targetModel}}) + registryRef.RegisterClient("unrelated-auth", "codex", []*registry.ModelInfo{{ID: unrelatedModel}}) + t.Cleanup(func() { + registryRef.UnregisterClient("target-auth") + registryRef.UnregisterClient("unrelated-auth") + }) + if _, errRegister := m.Register(context.Background(), &Auth{ID: "target-auth", Provider: "codex", Metadata: map[string]any{"request_retry": 0}}); errRegister != nil { + t.Fatalf("register target auth: %v", errRegister) + } + if _, errRegister := m.Register(context.Background(), &Auth{ID: "unrelated-auth", Provider: "codex", Metadata: map[string]any{"request_retry": 2}}); errRegister != nil { + t.Fatalf("register unrelated auth: %v", errRegister) + } + + if wait, shouldRetry := m.shouldRetryAfterError(&Error{HTTPStatus: http.StatusBadGateway, Message: "retryable failure"}, 0, []string{"codex"}, targetModel, 0); shouldRetry || wait != 0 { + t.Fatalf("unrelated model override retry = (%v, %t), want (0, false)", wait, shouldRetry) + } +} + +func TestManager_ShouldRetryAfterError_IgnoresDisabledRetryOverride(t *testing.T) { + m := NewManager(nil, nil, nil) + m.SetRetryConfig(0, 0, 0) + if _, errRegister := m.Register(context.Background(), &Auth{ID: "active-auth", Provider: "codex", Metadata: map[string]any{"request_retry": 0}}); errRegister != nil { + t.Fatalf("register active auth: %v", errRegister) + } + if _, errRegister := m.Register(context.Background(), &Auth{ID: "disabled-auth", Provider: "codex", Disabled: true, Metadata: map[string]any{"request_retry": 2}}); errRegister != nil { + t.Fatalf("register disabled auth: %v", errRegister) + } + + if wait, shouldRetry := m.shouldRetryAfterError(&Error{HTTPStatus: http.StatusBadGateway, Message: "retryable failure"}, 0, []string{"codex"}, "", 0); shouldRetry || wait != 0 { + t.Fatalf("disabled override retry = (%v, %t), want (0, false)", wait, shouldRetry) + } +} + +func TestManager_ShouldRetryAfterError_IgnoresNonRoundCooldownOverrides(t *testing.T) { + tests := []struct { + name string + state *ModelState + }{ + {name: "model disabled", state: &ModelState{Status: StatusDisabled}}, + {name: "unauthorized", state: &ModelState{Status: StatusError, Unavailable: true, LastError: &Error{HTTPStatus: http.StatusUnauthorized, Message: "unauthorized"}}}, + {name: "payment required", state: &ModelState{Status: StatusError, Unavailable: true, LastError: &Error{HTTPStatus: http.StatusPaymentRequired, Message: "payment required"}}}, + {name: "not found", state: &ModelState{Status: StatusError, Unavailable: true, LastError: &Error{HTTPStatus: http.StatusNotFound, Message: "not found"}}}, + {name: "model unsupported", state: &ModelState{Status: StatusError, Unavailable: true, LastError: &Error{HTTPStatus: http.StatusBadRequest, Message: "model not supported"}}}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetRetryConfig(0, time.Minute, 0) + model := "retry-non-round-" + uuid.NewString() + if test.state.Status != StatusDisabled { + test.state.NextRetryAfter = time.Now().Add(time.Minute) + } + auths := []*Auth{ + {ID: "retry-round-eligible-" + uuid.NewString(), Provider: "codex", Metadata: map[string]any{"request_retry": 0}}, + { + ID: "retry-round-ineligible-" + uuid.NewString(), + Provider: "codex", + Metadata: map[string]any{"request_retry": 2}, + ModelStates: map[string]*ModelState{model: test.state}, + }, + } + for _, auth := range auths { + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("register %s: %v", auth.ID, errRegister) + } + } + + if wait, shouldRetry := manager.shouldRetryAfterError(&Error{HTTPStatus: http.StatusBadGateway, Message: "upstream unavailable"}, 0, []string{"codex"}, model, time.Minute); shouldRetry || wait != 0 { + t.Fatalf("non-round cooldown override retry = (%v, %t), want (0, false)", wait, shouldRetry) + } + }) + } +} + +func TestManager_ShouldRetryAfterError_IgnoresRequestIneligibleOverrides(t *testing.T) { + tests := []struct { + name string + ctx context.Context + opts cliproxyexecutor.Options + eligible *Auth + ineligible *Auth + }{ + { + name: "credential policy", + ctx: withCredentialPolicy(context.Background(), CredentialPolicyCodexAlphaSearchV1), + eligible: &Auth{ + ID: "retry-policy-eligible", + Provider: "codex", + Attributes: map[string]string{"auth_kind": "oauth"}, + Metadata: map[string]any{"request_retry": 0}, + }, + ineligible: &Auth{ + ID: "retry-policy-ineligible", + Provider: "codex", + Attributes: map[string]string{"api_key": "ordinary"}, + Metadata: map[string]any{"request_retry": 2}, + }, + }, + { + name: "pinned credential", + ctx: context.Background(), + opts: cliproxyexecutor.Options{Metadata: map[string]any{cliproxyexecutor.PinnedAuthMetadataKey: "retry-pinned-eligible"}}, + eligible: &Auth{ + ID: "retry-pinned-eligible", + Provider: "codex", + Metadata: map[string]any{"request_retry": 0}, + }, + ineligible: &Auth{ + ID: "retry-pinned-ineligible", + Provider: "codex", + Metadata: map[string]any{"request_retry": 2}, + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetRetryConfig(0, 0, 0) + model := "retry-eligibility-" + uuid.NewString() + for _, auth := range []*Auth{test.eligible, test.ineligible} { + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("register %s: %v", auth.ID, errRegister) + } + } + + wait, shouldRetry := manager.shouldRetryAfterErrorWithHomeRetryLimit(test.ctx, test.opts, &Error{HTTPStatus: http.StatusBadGateway, Message: "retryable failure"}, 0, []string{"codex"}, model, 0, -1) + if shouldRetry || wait != 0 { + t.Fatalf("request-ineligible override retry = (%v, %t), want (0, false)", wait, shouldRetry) + } + }) + } +} + +func TestManager_RequestRetryRunsAdditionalLocalRoundWithoutCooldown(t *testing.T) { + previousDisableCooling := quotaCooldownDisabled.Load() + quotaCooldownDisabled.Store(false) + t.Cleanup(func() { quotaCooldownDisabled.Store(previousDisableCooling) }) + + tests := []struct { + name string + execute func(*Manager, cliproxyexecutor.Request) error + }{ + { + name: "nonstream", + execute: func(m *Manager, req cliproxyexecutor.Request) error { + _, errExecute := m.Execute(context.Background(), []string{"claude"}, req, cliproxyexecutor.Options{}) + return errExecute + }, + }, + { + name: "count tokens", + execute: func(m *Manager, req cliproxyexecutor.Request) error { + _, errExecute := m.ExecuteCount(context.Background(), []string{"claude"}, req, cliproxyexecutor.Options{}) + return errExecute + }, + }, + { + name: "stream", + execute: func(m *Manager, req cliproxyexecutor.Request) error { + _, errExecute := m.ExecuteStream(context.Background(), []string{"claude"}, req, cliproxyexecutor.Options{Stream: true}) + return errExecute + }, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + m := NewManager(nil, nil, nil) + m.SetRetryConfig(1, 0, 0) + executor := &credentialRetryLimitExecutor{id: "claude"} + m.RegisterExecutor(executor) + authID := uuid.NewString() + model := "retry-model-" + authID + auth := &Auth{ID: authID, Provider: "claude", Metadata: map[string]any{"disable_cooling": true}} + registry.GetGlobalRegistry().RegisterClient(authID, "claude", []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(authID) }) + if _, errRegister := m.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + if errExecute := tc.execute(m, cliproxyexecutor.Request{Model: model}); errExecute == nil || statusCodeFromError(errExecute) != http.StatusInternalServerError { + t.Fatalf("execute error = %v, want status 500", errExecute) + } + if got := executor.Calls(); got != 2 { + t.Fatalf("executor calls = %d, want initial round plus one additional round", got) + } + }) + } +} + func TestManager_ShouldRetryAfterError_UsesOAuthModelAliasForCooldown(t *testing.T) { m := NewManager(nil, nil, nil) m.SetRetryConfig(3, 30*time.Second, 0) @@ -112,6 +364,8 @@ func TestManager_ShouldRetryAfterError_UsesOAuthModelAliasForCooldown(t *testing }, }, } + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: upstreamModel}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) if _, errRegister := m.Register(context.Background(), auth); errRegister != nil { t.Fatalf("register auth: %v", errRegister) } diff --git a/sdk/cliproxy/auth/conductor_selection.go b/sdk/cliproxy/auth/conductor_selection.go index 248ad08b..90682466 100644 --- a/sdk/cliproxy/auth/conductor_selection.go +++ b/sdk/cliproxy/auth/conductor_selection.go @@ -709,7 +709,55 @@ func (m *Manager) retrySettings() (int, int, time.Duration) { return int(m.requestRetry.Load()), int(m.maxRetryCredentials.Load()), time.Duration(m.maxRetryInterval.Load()) } -func (m *Manager) closestCooldownWait(providers []string, model string, attempt int) (time.Duration, bool) { +func retryRoundAvailabilityForAuth(auth *Auth, model string, now time.Time) (bool, time.Time) { + blocked, reason, next := isAuthBlockedForModel(auth, model, now) + if !blocked { + return true, time.Time{} + } + if auth == nil || next.IsZero() || reason == blockReasonDisabled { + return false, time.Time{} + } + if auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" && auth.Quota.NextRecoverAt.After(now) { + return credentialRetryRoundStateEligible(auth.LastError, true), next + } + + modelKey := canonicalModelKey(model) + if modelKey != "" && len(auth.ModelStates) > 0 { + matchedBlocked := false + for stateModel, state := range auth.ModelStates { + if state == nil || canonicalModelKey(stateModel) != modelKey { + continue + } + if state.Status == StatusDisabled { + return false, time.Time{} + } + stateBlocked, _, stateNext := availabilityBlock(state.Unavailable, state.Quota.Exceeded, state.NextRetryAfter, state.Quota.NextRecoverAt, now) + if !stateBlocked { + continue + } + matchedBlocked = true + if stateNext.IsZero() || !credentialRetryRoundStateEligible(state.LastError, state.Quota.Exceeded) { + return false, time.Time{} + } + } + if matchedBlocked { + return true, next + } + } + if !credentialRetryRoundStateEligible(auth.LastError, auth.Quota.Exceeded) { + return false, time.Time{} + } + return true, next +} + +func credentialRetryRoundStateEligible(lastErr *Error, quotaExceeded bool) bool { + if lastErr == nil { + return quotaExceeded + } + return isCredentialRetryRoundStatus(statusCodeFromResult(lastErr)) +} + +func (m *Manager) closestCooldownWait(providers []string, model string, attempt int, eligibility authSelectionEligibility, pinnedAuthID string) (time.Duration, bool) { if m == nil || len(providers) == 0 { return 0, false } @@ -726,6 +774,7 @@ func (m *Manager) closestCooldownWait(providers []string, model string, attempt } providerSet[key] = struct{}{} } + registryRef := registry.GetGlobalRegistry() m.mu.RLock() defer m.mu.RUnlock() var ( @@ -733,13 +782,22 @@ func (m *Manager) closestCooldownWait(providers []string, model string, attempt minWait time.Duration ) for _, auth := range m.auths { - if auth == nil { + if auth == nil || auth.Disabled || auth.Status == StatusDisabled { + continue + } + if pinnedAuthID != "" && auth.ID != pinnedAuthID { + continue + } + if !eligibility.allows(auth) { continue } providerKey := executorKeyFromAuth(auth) if _, ok := providerSet[providerKey]; !ok { continue } + if model != "" && !m.authSupportsRouteModel(registryRef, auth, model) { + continue + } effectiveRetry := defaultRetry if override, ok := auth.RequestRetryOverride(); ok { effectiveRetry = override @@ -754,10 +812,13 @@ func (m *Manager) closestCooldownWait(providers []string, model string, attempt if strings.TrimSpace(model) != "" { checkModel = m.selectionModelForAuth(auth, model) } - blocked, reason, next := isAuthBlockedForModel(auth, checkModel, now) - if !blocked || next.IsZero() || reason == blockReasonDisabled { + retryEligible, next := retryRoundAvailabilityForAuth(auth, checkModel, now) + if !retryEligible { continue } + if next.IsZero() { + return 0, true + } wait := next.Sub(now) if wait < 0 { continue @@ -770,10 +831,11 @@ func (m *Manager) closestCooldownWait(providers []string, model string, attempt return minWait, found } -func (m *Manager) retryAllowed(attempt int, providers []string) bool { +func (m *Manager) retryAllowed(attempt int, providers []string, model string, eligibility authSelectionEligibility, pinnedAuthID string) bool { if m == nil || attempt < 0 || len(providers) == 0 { return false } + now := time.Now() defaultRetry := int(m.requestRetry.Load()) if defaultRetry < 0 { defaultRetry = 0 @@ -790,16 +852,26 @@ func (m *Manager) retryAllowed(attempt int, providers []string) bool { return false } + registryRef := registry.GetGlobalRegistry() m.mu.RLock() defer m.mu.RUnlock() for _, auth := range m.auths { - if auth == nil { + if auth == nil || auth.Disabled || auth.Status == StatusDisabled { + continue + } + if pinnedAuthID != "" && auth.ID != pinnedAuthID { + continue + } + if !eligibility.allows(auth) { continue } providerKey := executorKeyFromAuth(auth) if _, ok := providerSet[providerKey]; !ok { continue } + if model != "" && !m.authSupportsRouteModel(registryRef, auth, model) { + continue + } effectiveRetry := defaultRetry if override, ok := auth.RequestRetryOverride(); ok { effectiveRetry = override @@ -807,7 +879,14 @@ func (m *Manager) retryAllowed(attempt int, providers []string) bool { if effectiveRetry < 0 { effectiveRetry = 0 } - if attempt < effectiveRetry { + if attempt >= effectiveRetry { + continue + } + checkModel := model + if strings.TrimSpace(model) != "" { + checkModel = m.selectionModelForAuth(auth, model) + } + if retryEligible, _ := retryRoundAvailabilityForAuth(auth, checkModel, now); retryEligible { return true } } @@ -815,6 +894,15 @@ func (m *Manager) retryAllowed(attempt int, providers []string) bool { } func (m *Manager) shouldRetryAfterError(err error, attempt int, providers []string, model string, maxWait time.Duration) (time.Duration, bool) { + return m.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), cliproxyexecutor.Options{}, err, attempt, providers, model, maxWait, -1) +} + +// maxWait limits only positive cooldown waits between credential retry rounds. +// A non-positive value means no waiting: it does not disable same-round +// credential failover or an additional round that request-retry permits to start +// immediately. If every eligible credential still needs a positive cooldown, +// retry stops without waiting. +func (m *Manager) shouldRetryAfterErrorWithHomeRetryLimit(ctx context.Context, opts cliproxyexecutor.Options, err error, attempt int, providers []string, model string, maxWait time.Duration, homeRetryLimit int) (time.Duration, bool) { if err == nil { return 0, false } @@ -822,9 +910,6 @@ func (m *Manager) shouldRetryAfterError(err error, attempt int, providers []stri if errors.As(err, &homeBusy) && homeBusy != nil { return 0, false } - if maxWait <= 0 { - return 0, false - } status := statusCodeFromError(err) if status == http.StatusOK { return 0, false @@ -832,24 +917,119 @@ func (m *Manager) shouldRetryAfterError(err error, attempt int, providers []stri if isRequestInvalidError(err) || isRequestStopError(err) { return 0, false } - wait, found := m.closestCooldownWait(providers, model, attempt) + if m.HomeEnabled() { + var cooldownErr *homeDispatchRetryAfterError + if errors.As(err, &cooldownErr) && cooldownErr != nil { + observeHomeCooldownRetryLimit(cooldownErr, &homeRetryLimit, pinnedAuthIDFromMetadata(opts.Metadata) == "") + } + } + var exhausted *homeRetryRoundExhaustedError + if m.HomeEnabled() && errors.As(err, &exhausted) && exhausted != nil { + if !isCredentialRetryRoundStatus(status) || !m.homeRetryAllowed(attempt, homeRetryLimit) { + return 0, false + } + if exhausted.retryNow { + return 0, true + } + if retryAfter := retryAfterFromError(err); retryAfter != nil { + if *retryAfter < 0 || (*retryAfter > 0 && (maxWait <= 0 || *retryAfter > maxWait)) { + return 0, false + } + return *retryAfter, true + } + // Home will provide a cooldown error on the next round if all + // credentials are still cooling down; otherwise retry immediately. + return 0, true + } + if m.HomeEnabled() { + if status != http.StatusTooManyRequests || !m.homeRetryAllowed(attempt, homeRetryLimit) { + return 0, false + } + retryAfter := retryAfterFromError(err) + if retryAfter == nil || *retryAfter <= 0 || (maxWait <= 0 || *retryAfter > maxWait) { + return 0, false + } + return *retryAfter, true + } + eligibility := authSelectionEligibilityForRequest(ctx, opts) + pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata) + if !isCredentialRetryRoundStatus(status) || !m.retryAllowed(attempt, providers, model, eligibility, pinnedAuthID) { + return 0, false + } + wait, found := m.closestCooldownWait(providers, model, attempt, eligibility, pinnedAuthID) if found { - if wait > maxWait { + if wait > 0 && (maxWait <= 0 || wait > maxWait) { return 0, false } return wait, true } - if status != http.StatusTooManyRequests { - return 0, false + if retryAfter := retryAfterFromError(err); retryAfter != nil { + if *retryAfter < 0 || (*retryAfter > 0 && (maxWait <= 0 || *retryAfter > maxWait)) { + return 0, false + } + return *retryAfter, true } - if !m.retryAllowed(attempt, providers) { - return 0, false + return 0, true +} + +func (m *Manager) homeRetryAllowed(attempt int, retryLimit int) bool { + if m == nil || !m.HomeEnabled() || attempt < 0 { + return false } - retryAfter := retryAfterFromError(err) - if retryAfter == nil || *retryAfter <= 0 || *retryAfter > maxWait { - return 0, false + if retryLimit < 0 { + retryLimit = int(m.requestRetry.Load()) + if retryLimit < 0 { + retryLimit = 0 + } + } + return attempt < retryLimit +} + +func (m *Manager) observeHomeRetryLimit(auth *Auth, selection *HomeDispatchSelection, retryLimit *int) { + if m == nil || retryLimit == nil { + return + } + if selection != nil && selection.hasRequestRetry { + *retryLimit = selection.requestRetry + return + } + if auth == nil { + return + } + limit := int(m.requestRetry.Load()) + if override, ok := auth.RequestRetryOverride(); ok { + limit = override + } + if limit < 0 { + limit = 0 + } + if *retryLimit < 0 || limit > *retryLimit { + *retryLimit = limit + } +} + +func observeHomeCooldownRetryLimit(cooldown *homeDispatchRetryAfterError, retryLimit *int, acceptRemoteRetryLimit bool) { + if cooldown == nil || retryLimit == nil || !acceptRemoteRetryLimit { + return + } + if remoteLimit, ok := cooldown.RequestRetryLimit(); ok { + *retryLimit = remoteLimit + } +} + +func isCredentialRetryRoundStatus(status int) bool { + switch status { + case http.StatusForbidden, + http.StatusRequestTimeout, + http.StatusTooManyRequests, + http.StatusInternalServerError, + http.StatusBadGateway, + http.StatusServiceUnavailable, + http.StatusGatewayTimeout: + return true + default: + return false } - return *retryAfter, true } // cooldownWaitJitterCap bounds the random jitter added to cooldown waits so a @@ -860,7 +1040,7 @@ const cooldownWaitJitterCap = 2 * time.Second // concurrent requests waiting on the same recovery deadline do not wake in // lockstep and stampede the first credential that recovers. The jitter never // pushes the total wait past maxWait, which callers have already enforced as -// the retry ceiling; maxWait <= 0 means no ceiling. +// the retry ceiling; maxWait <= 0 is reserved for immediate retries. func jitteredCooldownWait(wait, maxWait time.Duration) time.Duration { if wait <= 0 { return wait @@ -1199,6 +1379,7 @@ func (m *Manager) SelectHomeAuthWithCredentialPolicy(ctx context.Context, provid tried := make(map[string]struct{}) for { selectionOpts := withHomeAuthCount(opts, homeAuthCount) + selectionOpts = withHomeExcludedAuthIDs(selectionOpts, tried) selection, errSelection := m.pickHomeDispatchSelection(selectionCtx, model, selectionOpts) if errSelection != nil { return nil, errSelection @@ -1245,6 +1426,7 @@ func (m *Manager) SelectHomeAuthByKind(ctx context.Context, provider string, mod tried := make(map[string]struct{}) for { selectionOpts := withHomeAuthCount(opts, homeAuthCount) + selectionOpts = withHomeExcludedAuthIDs(selectionOpts, tried) selection, errSelection := m.pickHomeDispatchSelection(ctx, model, selectionOpts) if errSelection != nil { return nil, errSelection diff --git a/sdk/cliproxy/auth/conductor_stream_overload_failover_test.go b/sdk/cliproxy/auth/conductor_stream_overload_failover_test.go index d0b1c2c3..8736ecf8 100644 --- a/sdk/cliproxy/auth/conductor_stream_overload_failover_test.go +++ b/sdk/cliproxy/auth/conductor_stream_overload_failover_test.go @@ -126,7 +126,7 @@ func TestExecuteStream_BootstrapOverload_StopsAtCredentialBudget(t *testing.T) { t.Cleanup(func() { quotaCooldownDisabled.Store(previous) }) m := NewManager(nil, nil, nil) - // Six credentials exist but only four may be attempted. + // Six credentials exist and only four may be attempted in one round. m.SetRetryConfig(5, 0, 4) registerOverloadAuths(t, m, 6) @@ -159,10 +159,10 @@ func TestExecuteStream_BootstrapOverload_StopsAtCredentialBudget(t *testing.T) { 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) + // A no-wait retry round may consume the remaining two credentials after the + // first four-credential sweep, but it must not exceed the available set. + if attempts < 4 || attempts > 6 { + t.Fatalf("attempts = %d, want between 4 and 6", attempts) } t.Logf("total upstream attempts across retry sweeps: %d", attempts) } diff --git a/sdk/cliproxy/auth/home_concurrency.go b/sdk/cliproxy/auth/home_concurrency.go index d3f91774..e5e06821 100644 --- a/sdk/cliproxy/auth/home_concurrency.go +++ b/sdk/cliproxy/auth/home_concurrency.go @@ -237,6 +237,17 @@ func decodeHomeDispatchError(raw []byte) error { switch strings.ToLower(code) { case "model_not_found": result.HTTPStatus = http.StatusNotFound + case "model_cooldown": + result.HTTPStatus = http.StatusTooManyRequests + cooldownErr := &homeDispatchRetryAfterError{cause: result} + if detail.RetryAfterMS > 0 { + cooldownErr.retryAfter = time.Duration(detail.RetryAfterMS) * time.Millisecond + } + if detail.RequestRetry != nil && *detail.RequestRetry >= 0 { + cooldownErr.requestRetry = *detail.RequestRetry + cooldownErr.hasRequestRetry = true + } + return cooldownErr case "authentication_error", "unauthorized", "no_credentials", "invalid_credential": result.HTTPStatus = http.StatusUnauthorized case "credential_concurrency_exceeded", "credential_model_concurrency_exceeded": @@ -263,13 +274,30 @@ func verifyAccountedHomeConcurrencyIdentity(tuple homeConcurrencyTuple, auth *Au return nil } -// SafeResponseHeaders returns trusted response headers only for CPA's concrete Home busy error. +// SafeResponseHeaders returns trusted response headers only for concrete +// Home-generated retry errors. func SafeResponseHeaders(err error) http.Header { var busy *HomeConcurrencyBusyError - if !errors.As(err, &busy) || busy == nil { + if errors.As(err, &busy) && busy != nil { + return busy.SafeResponseHeaders() + } + var exhausted *homeRetryRoundExhaustedError + if errors.As(err, &exhausted) && exhausted != nil { + retryAfter := exhausted.RetryAfter() + if retryAfter == nil { + return nil + } + return safeRetryAfterHeader(*retryAfter) + } + var cooldown *homeDispatchRetryAfterError + if !errors.As(err, &cooldown) || cooldown == nil { + return nil + } + retryAfter := cooldown.RetryAfter() + if retryAfter == nil { return nil } - return busy.SafeResponseHeaders() + return safeRetryAfterHeader(*retryAfter) } func safeRetryAfterHeader(retryAfter time.Duration) http.Header { diff --git a/sdk/cliproxy/auth/home_execution_paths_test.go b/sdk/cliproxy/auth/home_execution_paths_test.go index bba4055f..adaa5560 100644 --- a/sdk/cliproxy/auth/home_execution_paths_test.go +++ b/sdk/cliproxy/auth/home_execution_paths_test.go @@ -492,7 +492,15 @@ type lifecycleRetryDispatcher struct { } func (d *lifecycleRetryDispatcher) HeartbeatOK() bool { return true } -func (d *lifecycleRetryDispatcher) RPopAuth(_ context.Context, _ string, _ string, _ http.Header, _ int) ([]byte, error) { +func (d *lifecycleRetryDispatcher) RPopAuth(ctx context.Context, model string, sessionID string, headers http.Header, count int) ([]byte, error) { + return d.RPopAuthWithConstraints(ctx, model, sessionID, headers, count, nil, "") +} +func (d *lifecycleRetryDispatcher) RPopAuthWithConstraints(_ context.Context, _ string, _ string, _ http.Header, _ int, excludedAuthIDs []string, _ string) ([]byte, error) { + for _, authID := range excludedAuthIDs { + if authID == "home-auth" { + return nil, home.ErrAuthNotFound + } + } if d.calls.Add(1) == 2 && d.executor.first != nil { d.firstEndedBeforeRedispatch.Store(!d.executor.first.Active() && d.executor.firstCtx.Err() != nil) } @@ -535,6 +543,7 @@ func TestHomeStreamLifecycleFailureEndsBeforeFreshDispatch(t *testing.T) { registry := executionregistry.New() manager := NewManager(nil, nil, nil) manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(0, time.Second, 1) manager.PublishHomeDispatch(dispatcher, registry, 1) manager.RegisterExecutor(executor) ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) @@ -580,15 +589,54 @@ func TestHomeSelectionCancellationPreventsExecute(t *testing.T) { } } -type retryingHomeStreamExecutor struct { +type freshHomeStreamSelectionDispatcher struct { calls atomic.Int32 } +func (*freshHomeStreamSelectionDispatcher) HeartbeatOK() bool { return true } + +func (d *freshHomeStreamSelectionDispatcher) RPopAuth(ctx context.Context, model string, sessionID string, headers http.Header, count int) ([]byte, error) { + return d.RPopAuthWithConstraints(ctx, model, sessionID, headers, count, nil, "") +} + +func (d *freshHomeStreamSelectionDispatcher) RPopAuthWithConstraints(_ context.Context, _ string, _ string, _ http.Header, _ int, excludedAuthIDs []string, _ string) ([]byte, error) { + d.calls.Add(1) + excluded := make(map[string]struct{}, len(excludedAuthIDs)) + for _, authID := range excludedAuthIDs { + excluded[authID] = struct{}{} + } + for _, authID := range []string{"home-auth-a", "home-auth-b"} { + if _, okExcluded := excluded[authID]; okExcluded { + continue + } + return json.Marshal(homeAuthDispatchResponse{Auth: Auth{ + ID: authID, + Provider: "home-execution", + Status: StatusActive, + Attributes: map[string]string{ + AttributeAuthKind: AuthKindAPIKey, + }, + }}) + } + return nil, home.ErrAuthNotFound +} + +func (*freshHomeStreamSelectionDispatcher) AbortAmbiguousDispatch() {} + +type retryingHomeStreamExecutor struct { + mu sync.Mutex + calls atomic.Int32 + authIDs []string +} + func (*retryingHomeStreamExecutor) Identifier() string { return "home-execution" } func (*retryingHomeStreamExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { return cliproxyexecutor.Response{}, nil } -func (e *retryingHomeStreamExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { +func (e *retryingHomeStreamExecutor) ExecuteStream(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + e.mu.Lock() + e.authIDs = append(e.authIDs, auth.ID) + e.mu.Unlock() if e.calls.Add(1) == 1 { return nil, &Error{HTTPStatus: http.StatusUnauthorized, Message: "expired"} } @@ -607,10 +655,17 @@ func (*retryingHomeStreamExecutor) HttpRequest(context.Context, *Auth, *http.Req return nil, nil } +func (e *retryingHomeStreamExecutor) AuthIDs() []string { + e.mu.Lock() + defer e.mu.Unlock() + return append([]string(nil), e.authIDs...) +} + func TestHomeStreamRetryUsesFreshSelection(t *testing.T) { - dispatcher := &retainingHomeExecutionDispatcher{} + dispatcher := &freshHomeStreamSelectionDispatcher{} manager := NewManager(nil, nil, nil) manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(0, time.Second, 2) manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) executor := &retryingHomeStreamExecutor{} manager.RegisterExecutor(executor) @@ -624,6 +679,9 @@ func TestHomeStreamRetryUsesFreshSelection(t *testing.T) { if got := dispatcher.calls.Load(); got != 2 { t.Fatalf("Home RPOP calls = %d, want 2 for retrying stream invocations", got) } + if got := executor.AuthIDs(); len(got) != 2 || got[0] != "home-auth-a" || got[1] != "home-auth-b" { + t.Fatalf("executor auth IDs = %v, want [home-auth-a home-auth-b]", got) + } } type cancellationBarrierExecutor struct { diff --git a/sdk/cliproxy/auth/home_retry_contract_test.go b/sdk/cliproxy/auth/home_retry_contract_test.go new file mode 100644 index 00000000..c9b2d4df --- /dev/null +++ b/sdk/cliproxy/auth/home_retry_contract_test.go @@ -0,0 +1,1477 @@ +package auth + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/home" + "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +type retryContractHomeDispatcher struct { + mu sync.Mutex + authIDs []string + excluded [][]string + metadata map[string]any + requestRetry *int + websocket bool + exhaustedPayload []byte +} + +type legacyRepeatedStreamDispatcher struct { + calls atomic.Int32 +} + +type retryRoundStartCooldownDispatcher struct { + calls atomic.Int32 +} + +type retryRoundRepeatedCooldownDispatcher struct { + calls atomic.Int32 +} + +type retryRoundLimitDownshiftDispatcher struct { + calls atomic.Int32 +} + +type aggregateRetryHomeDispatcher struct { + calls atomic.Int32 +} + +func (*legacyRepeatedStreamDispatcher) HeartbeatOK() bool { return true } + +func (d *legacyRepeatedStreamDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) { + d.calls.Add(1) + return json.Marshal(homeAuthDispatchResponse{Auth: Auth{ + ID: "home-retry-a", + Provider: "home-retry-contract", + Status: StatusActive, + }}) +} + +func (*legacyRepeatedStreamDispatcher) AbortAmbiguousDispatch() {} + +func (*retryRoundStartCooldownDispatcher) HeartbeatOK() bool { return true } + +func (d *retryRoundStartCooldownDispatcher) RPopAuth(ctx context.Context, model string, sessionID string, headers http.Header, count int) ([]byte, error) { + return d.RPopAuthWithConstraints(ctx, model, sessionID, headers, count, nil, "") +} + +func (d *retryRoundStartCooldownDispatcher) RPopAuthWithConstraints(_ context.Context, _ string, _ string, _ http.Header, _ int, _ []string, _ string) ([]byte, error) { + if d.calls.Add(1) == 2 { + return []byte(`{"error":{"type":"model_cooldown","message":"credential is cooling down","retryable":true,"retry_after_ms":1,"request_retry":1}}`), nil + } + return json.Marshal(homeAuthDispatchResponse{Auth: Auth{ + ID: "home-retry-a", + Provider: "home-retry-contract", + Status: StatusActive, + }}) +} + +func (*retryRoundStartCooldownDispatcher) AbortAmbiguousDispatch() {} + +func (*retryRoundRepeatedCooldownDispatcher) HeartbeatOK() bool { return true } + +func (d *retryRoundRepeatedCooldownDispatcher) RPopAuth(ctx context.Context, model string, sessionID string, headers http.Header, count int) ([]byte, error) { + return d.RPopAuthWithConstraints(ctx, model, sessionID, headers, count, nil, "") +} + +func (d *retryRoundRepeatedCooldownDispatcher) RPopAuthWithConstraints(_ context.Context, _ string, _ string, _ http.Header, _ int, _ []string, _ string) ([]byte, error) { + if d.calls.Add(1) > 1 { + return []byte(`{"error":{"type":"model_cooldown","message":"credential is cooling down","retryable":true,"retry_after_ms":1,"request_retry":1}}`), nil + } + return json.Marshal(homeAuthDispatchResponse{Auth: Auth{ + ID: "home-retry-a", + Provider: "home-retry-contract", + Status: StatusActive, + }}) +} + +func (*retryRoundRepeatedCooldownDispatcher) AbortAmbiguousDispatch() {} + +func (*retryRoundLimitDownshiftDispatcher) HeartbeatOK() bool { return true } + +func (d *retryRoundLimitDownshiftDispatcher) RPopAuth(ctx context.Context, model string, sessionID string, headers http.Header, count int) ([]byte, error) { + return d.RPopAuthWithConstraints(ctx, model, sessionID, headers, count, nil, "") +} + +func (d *retryRoundLimitDownshiftDispatcher) RPopAuthWithConstraints(_ context.Context, _ string, _ string, _ http.Header, _ int, _ []string, _ string) ([]byte, error) { + switch d.calls.Add(1) { + case 1: + retryLimit := 1 + return json.Marshal(homeAuthDispatchResponse{RequestRetry: &retryLimit, Auth: Auth{ + ID: "home-retry-a", + Provider: "home-retry-contract", + Status: StatusActive, + }}) + case 2: + return []byte(`{"error":{"type":"model_cooldown","message":"remaining credentials are cooling down","retryable":true,"retry_after_ms":1,"request_retry":0}}`), nil + default: + return json.Marshal(homeAuthDispatchResponse{Auth: Auth{ + ID: "home-retry-b", + Provider: "home-retry-contract", + Status: StatusActive, + }}) + } +} + +func (*retryRoundLimitDownshiftDispatcher) AbortAmbiguousDispatch() {} + +func (*aggregateRetryHomeDispatcher) HeartbeatOK() bool { return true } + +func (d *aggregateRetryHomeDispatcher) RPopAuth(ctx context.Context, model string, sessionID string, headers http.Header, count int) ([]byte, error) { + return d.RPopAuthWithConstraints(ctx, model, sessionID, headers, count, nil, "") +} + +func (d *aggregateRetryHomeDispatcher) RPopAuthWithConstraints(_ context.Context, _ string, _ string, _ http.Header, _ int, _ []string, _ string) ([]byte, error) { + authID := "home-retry-a" + override := 0 + if d.calls.Add(1) > 1 { + authID = "home-retry-b" + override = 2 + } + requestRetry := 2 + return json.Marshal(homeAuthDispatchResponse{ + RequestRetry: &requestRetry, + Auth: Auth{ + ID: authID, + Provider: "home-retry-contract", + Status: StatusActive, + Metadata: map[string]any{"request_retry": override}, + }, + }) +} + +func (*aggregateRetryHomeDispatcher) AbortAmbiguousDispatch() {} + +func (*retryContractHomeDispatcher) HeartbeatOK() bool { return true } + +func (d *retryContractHomeDispatcher) RPopAuth(ctx context.Context, model string, sessionID string, headers http.Header, count int) ([]byte, error) { + return d.RPopAuthWithConstraints(ctx, model, sessionID, headers, count, nil, "") +} + +func (d *retryContractHomeDispatcher) RPopAuthWithConstraints(_ context.Context, _ string, _ string, _ http.Header, _ int, excludedAuthIDs []string, pinnedAuthID string) ([]byte, error) { + d.mu.Lock() + defer d.mu.Unlock() + d.excluded = append(d.excluded, append([]string(nil), excludedAuthIDs...)) + excluded := make(map[string]struct{}, len(excludedAuthIDs)) + for _, authID := range excludedAuthIDs { + excluded[authID] = struct{}{} + } + for _, authID := range d.authIDs { + if pinnedAuthID != "" && authID != pinnedAuthID { + continue + } + if _, okExcluded := excluded[authID]; okExcluded { + continue + } + attributes := map[string]string{} + if d.websocket { + attributes["websockets"] = "true" + } + return json.Marshal(homeAuthDispatchResponse{ + RequestRetry: d.requestRetry, + Auth: Auth{ + ID: authID, + Provider: "home-retry-contract", + Status: StatusActive, + Metadata: d.metadata, + Attributes: attributes, + }, + }) + } + if len(d.exhaustedPayload) > 0 { + return append([]byte(nil), d.exhaustedPayload...), nil + } + return nil, home.ErrAuthNotFound +} + +func (*retryContractHomeDispatcher) AbortAmbiguousDispatch() {} + +func (d *retryContractHomeDispatcher) Excluded() [][]string { + d.mu.Lock() + defer d.mu.Unlock() + result := make([][]string, len(d.excluded)) + for index := range d.excluded { + result[index] = append([]string(nil), d.excluded[index]...) + } + return result +} + +type retryContractHomeExecutor struct { + mu sync.Mutex + calls []string + failAll bool + failure error + failures map[string]error + streamBootstrap bool + streamHeaders http.Header +} + +func (*retryContractHomeExecutor) Identifier() string { return "home-retry-contract" } + +func (e *retryContractHomeExecutor) Execute(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + e.mu.Lock() + e.calls = append(e.calls, auth.ID) + e.mu.Unlock() + if auth.ID == "home-retry-a" || e.failAll { + return cliproxyexecutor.Response{}, e.failureError(auth.ID) + } + return cliproxyexecutor.Response{Payload: []byte(auth.ID)}, nil +} + +func (e *retryContractHomeExecutor) ExecuteStream(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + e.mu.Lock() + e.calls = append(e.calls, auth.ID) + e.mu.Unlock() + if auth.ID == "home-retry-a" || e.failAll { + errFailure := e.failureError(auth.ID) + if e.streamBootstrap { + chunks := make(chan cliproxyexecutor.StreamChunk, 1) + chunks <- cliproxyexecutor.StreamChunk{Err: errFailure} + close(chunks) + return &cliproxyexecutor.StreamResult{Headers: e.streamHeaders.Clone(), Chunks: chunks}, nil + } + return nil, errFailure + } + chunks := make(chan cliproxyexecutor.StreamChunk, 1) + chunks <- cliproxyexecutor.StreamChunk{Payload: []byte(auth.ID)} + close(chunks) + return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil +} + +func (e *retryContractHomeExecutor) failureError(authID string) error { + if failure := e.failures[authID]; failure != nil { + return failure + } + if e.failure != nil { + return e.failure + } + return retryContractRateLimitError{} +} + +func (*retryContractHomeExecutor) Refresh(context.Context, *Auth) (*Auth, error) { return nil, nil } + +func (e *retryContractHomeExecutor) CountTokens(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + e.mu.Lock() + e.calls = append(e.calls, auth.ID) + e.mu.Unlock() + if auth.ID == "home-retry-a" || e.failAll { + return cliproxyexecutor.Response{}, e.failureError(auth.ID) + } + return cliproxyexecutor.Response{Payload: []byte(auth.ID)}, nil +} + +func (*retryContractHomeExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func (e *retryContractHomeExecutor) Calls() []string { + e.mu.Lock() + defer e.mu.Unlock() + return append([]string(nil), e.calls...) +} + +type retryContractRateLimitError struct { + retryAfter time.Duration +} + +func (retryContractRateLimitError) Error() string { return "credential rate limited" } + +func (retryContractRateLimitError) StatusCode() int { return http.StatusTooManyRequests } + +func (e retryContractRateLimitError) RetryAfter() *time.Duration { + value := e.retryAfter + if value == 0 { + value = time.Millisecond + } + return &value +} + +type retainingRetryContractHomeExecutor struct { + *retryContractHomeExecutor +} + +func (e *retainingRetryContractHomeExecutor) Execute(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + if lifecycle, ok := opts.ExecutionLifecycle.(interface{ Retain() }); ok { + lifecycle.Retain() + } + return cliproxyexecutor.Response{Payload: []byte(auth.ID)}, nil +} + +func TestHomePinnedAuthRejectsMismatchedDispatch(t *testing.T) { + dispatcher := &retryContractHomeDispatcher{authIDs: []string{"home-retry-b"}} + executor := &retryContractHomeExecutor{} + registry := executionregistry.New() + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.PublishHomeDispatch(dispatcher, registry, 1) + manager.RegisterExecutor(executor) + + _, errExecute := manager.Execute(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Metadata: map[string]any{ + cliproxyexecutor.PinnedAuthMetadataKey: "home-retry-a", + }}) + var authErr *Error + if !errors.As(errExecute, &authErr) || authErr == nil || authErr.Code != "auth_not_found" { + t.Fatalf("Execute() error = %T %v, want pinned auth_not_found", errExecute, errExecute) + } + if got := executor.Calls(); len(got) != 0 { + t.Fatalf("executor calls = %v, want no mismatched credential execution", got) + } + if errDrain := registry.Drain(context.Background()); errDrain != nil { + t.Fatalf("Drain() error = %v", errDrain) + } +} + +func TestHomePinnedAuthRetriesOnlyPinnedCredential(t *testing.T) { + aggregateRetry := 3 + dispatcher := &retryContractHomeDispatcher{ + authIDs: []string{"home-retry-b", "home-retry-a"}, + metadata: map[string]any{"request_retry": 1}, + requestRetry: &aggregateRetry, + } + executor := &retryContractHomeExecutor{ + failAll: true, + failure: &Error{HTTPStatus: http.StatusBadGateway, Message: "upstream unavailable"}, + } + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(0, time.Second, 0) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + _, errExecute := manager.Execute(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Metadata: map[string]any{ + cliproxyexecutor.PinnedAuthMetadataKey: "home-retry-a", + }}) + if errExecute == nil { + t.Fatal("Execute() error = nil, want terminal upstream error") + } + if got := executor.Calls(); len(got) != 2 || got[0] != "home-retry-a" || got[1] != "home-retry-a" { + t.Fatalf("executor calls = %v, want pinned auth once in each of two rounds", got) + } + excluded := dispatcher.Excluded() + if len(excluded) != 2 || len(excluded[0]) != 0 || len(excluded[1]) != 0 { + t.Fatalf("Home excluded auth IDs = %v, want a fresh pinned selection in each round", excluded) + } +} + +func TestHomeExcludedCredentialEndsRetainedWebsocketSelection(t *testing.T) { + dispatcher := &retryContractHomeDispatcher{ + authIDs: []string{"home-retry-a", "home-retry-b"}, + websocket: true, + } + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(0, time.Second, 2) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(&retainingRetryContractHomeExecutor{retryContractHomeExecutor: &retryContractHomeExecutor{}}) + + ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) + opts := cliproxyexecutor.Options{Metadata: map[string]any{ + cliproxyexecutor.ExecutionSessionMetadataKey: "home-retry-session", + }} + if _, errExecute := manager.Execute(ctx, []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, opts); errExecute != nil { + t.Fatalf("first Execute() error = %v", errExecute) + } + + pickOpts := withHomeExcludedAuthIDs(opts, map[string]struct{}{"home-retry-a": {}}) + selection, errPick := manager.pickHomeDispatchSelection(ctx, "gpt", pickOpts) + if errPick != nil { + t.Fatalf("pickHomeDispatchSelection() error = %v", errPick) + } + defer selection.End("test_complete") + if auth := selection.CloneAuth(); auth == nil || auth.ID != "home-retry-b" { + t.Fatalf("selected auth = %#v, want home-retry-b", auth) + } + + excluded := dispatcher.Excluded() + if len(excluded) != 2 || len(excluded[0]) != 0 || len(excluded[1]) != 1 || excluded[1][0] != "home-retry-a" { + t.Fatalf("Home excluded auth IDs = %v, want [[], [home-retry-a]]", excluded) + } +} + +func TestHomeRetryRoundTriesFreshCredentialWhenRequestRetryIsZero(t *testing.T) { + for _, stream := range []bool{false, true} { + t.Run(map[bool]string{false: "nonstream", true: "stream"}[stream], func(t *testing.T) { + dispatcher := &retryContractHomeDispatcher{authIDs: []string{"home-retry-a", "home-retry-b"}} + executor := &retryContractHomeExecutor{} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(0, time.Second, 2) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + if stream { + result, errExecute := manager.ExecuteStream(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}) + if errExecute != nil { + t.Fatalf("ExecuteStream() error = %v", errExecute) + } + for range result.Chunks { + } + } else { + response, errExecute := manager.Execute(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}) + if errExecute != nil { + t.Fatalf("Execute() error = %v", errExecute) + } + if string(response.Payload) != "home-retry-b" { + t.Fatalf("response payload = %q, want home-retry-b", string(response.Payload)) + } + } + + if got := executor.Calls(); len(got) != 2 || got[0] != "home-retry-a" || got[1] != "home-retry-b" { + t.Fatalf("executor calls = %v, want [home-retry-a home-retry-b]", got) + } + excluded := dispatcher.Excluded() + if len(excluded) != 2 || len(excluded[0]) != 0 || len(excluded[1]) != 1 || excluded[1][0] != "home-retry-a" { + t.Fatalf("Home excluded auth IDs = %v, want [[], [home-retry-a]]", excluded) + } + }) + } +} + +func TestHomeCountTokensTriesFreshCredentialWhenRequestRetryIsZero(t *testing.T) { + dispatcher := &retryContractHomeDispatcher{authIDs: []string{"home-retry-a", "home-retry-b"}} + executor := &retryContractHomeExecutor{} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(0, time.Second, 2) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + response, errExecute := manager.ExecuteCount(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}) + if errExecute != nil { + t.Fatalf("ExecuteCount() error = %v", errExecute) + } + if string(response.Payload) != "home-retry-b" { + t.Fatalf("response payload = %q, want home-retry-b", string(response.Payload)) + } + if got := executor.Calls(); len(got) != 2 || got[0] != "home-retry-a" || got[1] != "home-retry-b" { + t.Fatalf("executor calls = %v, want [home-retry-a home-retry-b]", got) + } + excluded := dispatcher.Excluded() + if len(excluded) != 2 || len(excluded[0]) != 0 || len(excluded[1]) != 1 || excluded[1][0] != "home-retry-a" { + t.Fatalf("Home excluded auth IDs = %v, want [[], [home-retry-a]]", excluded) + } +} + +func TestHomeRetryPolicyAllowsRemoteCooldownWithoutLocalCredentials(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(1, time.Second, 0) + errRemoteCooldown := &homeDispatchRetryAfterError{ + cause: &Error{HTTPStatus: http.StatusTooManyRequests, Message: "all Home credentials are cooling down"}, + retryAfter: 10 * time.Millisecond, + } + + wait, shouldRetry := manager.shouldRetryAfterError(errRemoteCooldown, 0, []string{"home-retry-contract"}, "gpt", time.Second) + if !shouldRetry || wait != 10*time.Millisecond { + t.Fatalf("shouldRetryAfterError() = (%v, %t), want (10ms, true)", wait, shouldRetry) + } + if _, shouldRetry = manager.shouldRetryAfterError(errRemoteCooldown, 1, []string{"home-retry-contract"}, "gpt", time.Second); shouldRetry { + t.Fatal("shouldRetryAfterError() retried after the configured Home retry round") + } + wait, shouldRetry = manager.shouldRetryAfterError(errRemoteCooldown, 0, []string{"home-retry-contract"}, "gpt", 0) + if shouldRetry || wait != 0 { + t.Fatalf("shouldRetryAfterError() with zero wait interval = (%v, %t), want (0, false)", wait, shouldRetry) + } + errRoundExhausted := markHomeRetryRoundExhausted(&Error{HTTPStatus: http.StatusBadGateway, Message: "upstream unavailable"}, nil, false) + wait, shouldRetry = manager.shouldRetryAfterError(errRoundExhausted, 0, []string{"home-retry-contract"}, "gpt", 0) + if !shouldRetry || wait != 0 { + t.Fatalf("shouldRetryAfterError() immediate round = (%v, %t), want (0, true)", wait, shouldRetry) + } + var invalidTiming homeRetryRoundTiming + invalidTiming.Observe(retryContractRateLimitError{retryAfter: -time.Millisecond}) + errInvalidWait := markHomeRetryRoundExhausted(retryContractRateLimitError{retryAfter: -time.Millisecond}, invalidTiming.RetryAfter(), false) + if wait, shouldRetry = manager.shouldRetryAfterError(errInvalidWait, 0, []string{"home-retry-contract"}, "gpt", 0); shouldRetry || wait != 0 { + t.Fatalf("shouldRetryAfterError() negative wait = (%v, %t), want (0, false)", wait, shouldRetry) + } +} + +func TestRetryIntervalFiltersCooldownCredentials(t *testing.T) { + tests := []struct { + name string + cooldowns []time.Duration + wantRetry bool + maxWantWait time.Duration + }{ + {name: "short and long cooldowns", cooldowns: []time.Duration{10 * time.Second, time.Minute}, wantRetry: true, maxWantWait: 10 * time.Second}, + {name: "only long cooldown", cooldowns: []time.Duration{time.Minute}}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + const ( + provider = "retry-interval-contract" + model = "gpt" + ) + manager := NewManager(nil, nil, nil) + manager.SetRetryConfig(1, 30*time.Second, 0) + now := time.Now() + for index, cooldown := range test.cooldowns { + authID := fmt.Sprintf("retry-interval-%s-%d", strings.ReplaceAll(test.name, " ", "-"), index) + deadline := now.Add(cooldown) + auth := &Auth{ + ID: authID, + Provider: provider, + Status: StatusActive, + ModelStates: map[string]*ModelState{ + model: { + Status: StatusError, + Unavailable: true, + NextRetryAfter: deadline, + LastError: &Error{HTTPStatus: http.StatusTooManyRequests}, + Quota: QuotaState{Exceeded: true, NextRecoverAt: deadline}, + }, + }, + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(authID) }) + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("Register() error = %v", errRegister) + } + } + + wait, shouldRetry := manager.shouldRetryAfterError(&Error{HTTPStatus: http.StatusTooManyRequests}, 0, []string{provider}, model, 30*time.Second) + if shouldRetry != test.wantRetry { + t.Fatalf("shouldRetryAfterError() = (%v, %t), want retry %t", wait, shouldRetry, test.wantRetry) + } + if test.wantRetry && (wait <= 0 || wait > test.maxWantWait) { + t.Fatalf("shouldRetryAfterError() wait = %v, want the earliest cooldown within %v", wait, test.maxWantWait) + } + }) + } +} + +func TestHomeRetryPolicyUsesRemoteCredentialOverrideBeforeSelection(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(0, time.Second, 0) + errRemoteCooldown := &homeDispatchRetryAfterError{ + cause: &Error{HTTPStatus: http.StatusTooManyRequests, Message: "all Home credentials are cooling down"}, + retryAfter: 10 * time.Millisecond, + requestRetry: 1, + hasRequestRetry: true, + } + + wait, shouldRetry := manager.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), cliproxyexecutor.Options{}, errRemoteCooldown, 0, []string{"home-retry-contract"}, "gpt", time.Second, -1) + if !shouldRetry || wait != 10*time.Millisecond { + t.Fatalf("remote credential override retry = (%v, %t), want (10ms, true)", wait, shouldRetry) + } + if _, shouldRetry = manager.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), cliproxyexecutor.Options{}, errRemoteCooldown, 1, []string{"home-retry-contract"}, "gpt", time.Second, -1); shouldRetry { + t.Fatal("remote credential override allowed more than one additional round") + } + pinnedOpts := cliproxyexecutor.Options{Metadata: map[string]any{ + cliproxyexecutor.PinnedAuthMetadataKey: "home-retry-a", + }} + if _, shouldRetry = manager.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), pinnedOpts, errRemoteCooldown, 0, []string{"home-retry-contract"}, "gpt", time.Second, -1); shouldRetry { + t.Fatal("aggregate retry limit from unpinned Home credentials affected a pinned request") + } + + errRemoteCooldown.requestRetry = 0 + manager.SetRetryConfig(3, time.Second, 0) + if _, shouldRetry = manager.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), cliproxyexecutor.Options{}, errRemoteCooldown, 0, []string{"home-retry-contract"}, "gpt", time.Second, -1); shouldRetry { + t.Fatal("explicit remote credential override 0 did not suppress the global retry setting") + } + retryLimit := 3 + observeHomeCooldownRetryLimit(errRemoteCooldown, &retryLimit, true) + if retryLimit != 0 { + t.Fatalf("observed remote cooldown retry limit = %d, want authoritative 0", retryLimit) + } +} + +func TestHomeRetryRoundCredentialLimitStartsNextRoundImmediately(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(1, time.Second, 1) + retryAfter := 5 * time.Second + errRoundExhausted := markHomeRetryRoundExhausted( + retryContractRateLimitError{retryAfter: retryAfter}, + &retryAfter, + true, + ) + + wait, shouldRetry := manager.shouldRetryAfterError(errRoundExhausted, 0, []string{"home-retry-contract"}, "gpt", time.Second) + if !shouldRetry || wait != 0 { + t.Fatalf("credential-limit retry = (%v, %t), want immediate next round", wait, shouldRetry) + } + if got := SafeResponseHeaders(errRoundExhausted).Get("Retry-After"); got != "5" { + t.Fatalf("safe Retry-After header = %q, want 5", got) + } +} + +func TestHomeCredentialLimitWaitsBeforeConsumingAdditionalRound(t *testing.T) { + for _, stream := range []bool{false, true} { + t.Run(map[bool]string{false: "nonstream", true: "stream"}[stream], func(t *testing.T) { + dispatcher := &retryRoundStartCooldownDispatcher{} + executor := &retryContractHomeExecutor{failAll: true} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(1, time.Second, 1) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + if stream { + result, errExecute := manager.ExecuteStream(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}) + if errExecute == nil || result != nil { + t.Fatalf("ExecuteStream() = result %#v, error %v; want terminal retry error", result, errExecute) + } + } else { + if _, errExecute := manager.Execute(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}); errExecute == nil { + t.Fatal("Execute() error = nil, want terminal retry error") + } + } + if got := executor.Calls(); len(got) != 2 { + t.Fatalf("executor calls = %v, want one execution in each of two rounds", got) + } + if got := dispatcher.calls.Load(); got != 3 { + t.Fatalf("Home dispatch calls = %d, want selection, cooldown wait, selection", got) + } + }) + } +} + +func TestHomePendingRetryRoundStopsWhenRemoteLimitDrops(t *testing.T) { + for _, stream := range []bool{false, true} { + t.Run(map[bool]string{false: "nonstream", true: "stream"}[stream], func(t *testing.T) { + dispatcher := &retryRoundLimitDownshiftDispatcher{} + executor := &retryContractHomeExecutor{} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(1, time.Second, 1) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + if stream { + result, errExecute := manager.ExecuteStream(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}) + if errExecute == nil || result != nil { + t.Fatalf("ExecuteStream() = result %#v, error %v; want terminal cooldown error", result, errExecute) + } + } else { + if _, errExecute := manager.Execute(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}); errExecute == nil { + t.Fatal("Execute() error = nil, want terminal cooldown error") + } + } + if got := executor.Calls(); len(got) != 1 || got[0] != "home-retry-a" { + t.Fatalf("executor calls = %v, want only the initial credential", got) + } + if got := dispatcher.calls.Load(); got != 2 { + t.Fatalf("Home dispatch calls = %d, want initial selection and one cooldown response", got) + } + }) + } +} + +func TestHomePendingRetryRoundStopsAfterRepeatedCooldown(t *testing.T) { + for _, stream := range []bool{false, true} { + t.Run(map[bool]string{false: "nonstream", true: "stream"}[stream], func(t *testing.T) { + dispatcher := &retryRoundRepeatedCooldownDispatcher{} + executor := &retryContractHomeExecutor{failAll: true} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(1, time.Second, 1) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + if stream { + result, errExecute := manager.ExecuteStream(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}) + if errExecute == nil || result != nil { + t.Fatalf("ExecuteStream() = result %#v, error %v; want terminal cooldown error", result, errExecute) + } + } else { + if _, errExecute := manager.Execute(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}); errExecute == nil { + t.Fatal("Execute() error = nil, want terminal cooldown error") + } + } + if got := executor.Calls(); len(got) != 1 { + t.Fatalf("executor calls = %v, want only the initial round execution", got) + } + if got := dispatcher.calls.Load(); got != 3 { + t.Fatalf("Home dispatch calls = %d, want selection and two cooldown responses", got) + } + }) + } +} + +func TestHomeRetryRoundUsesEarliestCredentialRetryAfter(t *testing.T) { + dispatcher := &retryContractHomeDispatcher{authIDs: []string{"home-retry-a", "home-retry-b"}} + executor := &retryContractHomeExecutor{ + failAll: true, + failures: map[string]error{ + "home-retry-a": retryContractRateLimitError{retryAfter: 5 * time.Millisecond}, + "home-retry-b": retryContractRateLimitError{retryAfter: 50 * time.Millisecond}, + }, + } + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(1, time.Second, 2) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + retryLimit := -1 + _, errExecute := manager.executeHomeOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}, false, 2, &retryLimit) + if !isHomeRetryRoundExhausted(errExecute) { + t.Fatalf("executeHomeOnce() error = %v, want exhausted retry round", errExecute) + } + retryAfter := retryAfterFromError(errExecute) + if retryAfter == nil || *retryAfter != 5*time.Millisecond { + t.Fatalf("retry after = %v, want earliest credential delay 5ms", retryAfter) + } +} + +func TestHomeStreamBootstrapErrorPreservesAggregatedRetryAfter(t *testing.T) { + dispatcher := &retryContractHomeDispatcher{authIDs: []string{"home-retry-a", "home-retry-b"}} + executor := &retryContractHomeExecutor{ + failAll: true, + streamBootstrap: true, + streamHeaders: http.Header{"Retry-After": {"30"}}, + failures: map[string]error{ + "home-retry-a": retryContractRateLimitError{retryAfter: 1500 * time.Millisecond}, + "home-retry-b": retryContractRateLimitError{retryAfter: 5 * time.Second}, + }, + } + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(0, time.Second, 2) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + result, errExecute := manager.ExecuteStream(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}) + if errExecute != nil { + t.Fatalf("ExecuteStream() error = %v", errExecute) + } + if result == nil { + t.Fatal("ExecuteStream() result = nil") + } + chunk, ok := <-result.Chunks + if !ok || chunk.Err == nil { + t.Fatalf("stream bootstrap chunk = %#v, %t; want terminal error", chunk, ok) + } + if !isHomeRetryRoundExhausted(chunk.Err) { + t.Fatalf("stream bootstrap error = %v, want exhausted retry round", chunk.Err) + } + if got := SafeResponseHeaders(chunk.Err).Get("Retry-After"); got != "2" { + t.Fatalf("safe Retry-After header = %q, want aggregated delay rounded to 2 seconds", got) + } +} + +func TestHomeRetryRoundUsesAuthoritativeRemoteCooldown(t *testing.T) { + tests := []struct { + name string + execute func(*Manager, *int) error + }{ + { + name: "nonstream", + execute: func(manager *Manager, retryLimit *int) error { + _, errExecute := manager.executeHomeOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}, false, 2, retryLimit) + return errExecute + }, + }, + { + name: "stream", + execute: func(manager *Manager, retryLimit *int) error { + _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}, 2, retryLimit) + return errExecute + }, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + dispatcher := &retryContractHomeDispatcher{ + authIDs: []string{"home-retry-a"}, + exhaustedPayload: []byte(`{"error":{"type":"model_cooldown","message":"remaining Home credentials are cooling down","retryable":true,"retry_after_ms":5000}}`), + } + executor := &retryContractHomeExecutor{ + failAll: true, + failure: retryContractRateLimitError{retryAfter: 1500 * time.Millisecond}, + } + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(1, time.Second, 2) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + retryLimit := -1 + errExecute := tc.execute(manager, &retryLimit) + if !isHomeRetryRoundExhausted(errExecute) { + t.Fatalf("execution error = %v, want exhausted retry round", errExecute) + } + retryAfter := retryAfterFromError(errExecute) + if retryAfter == nil || *retryAfter != 5*time.Second { + t.Fatalf("retry after = %v, want Home next-round cooldown 5s", retryAfter) + } + if got := SafeResponseHeaders(errExecute).Get("Retry-After"); got != "5" { + t.Fatalf("safe Retry-After header = %q, want Home next-round delay 5 seconds", got) + } + }) + } +} + +func TestHomeCooldownClassificationPreservesNonRetryableRoundStatus(t *testing.T) { + tests := []struct { + name string + execute func(*Manager, *int) error + }{ + { + name: "nonstream", + execute: func(manager *Manager, retryLimit *int) error { + _, errExecute := manager.executeHomeOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}, false, 2, retryLimit) + return errExecute + }, + }, + { + name: "stream", + execute: func(manager *Manager, retryLimit *int) error { + _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}, 2, retryLimit) + return errExecute + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + dispatcher := &retryContractHomeDispatcher{ + authIDs: []string{"home-retry-a"}, + exhaustedPayload: []byte(`{"error":{"type":"model_cooldown","message":"another credential is cooling down","retryable":true,"retry_after_ms":5,"request_retry":2}}`), + } + executor := &retryContractHomeExecutor{ + failAll: true, + failure: &Error{HTTPStatus: http.StatusUnauthorized, Message: "invalid credential"}, + } + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(3, time.Second, 2) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + retryLimit := -1 + errExecute := test.execute(manager, &retryLimit) + if !isHomeRetryRoundExhausted(errExecute) || statusCodeFromError(errExecute) != http.StatusUnauthorized { + t.Fatalf("execution error = %T %v, want exhausted 401 round", errExecute, errExecute) + } + if retryLimit != 2 { + t.Fatalf("observed retry limit = %d, want authoritative Home limit 2", retryLimit) + } + if wait, shouldRetry := manager.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), cliproxyexecutor.Options{}, errExecute, 0, []string{"home-retry-contract"}, "gpt", time.Second, retryLimit); shouldRetry || wait != 0 { + t.Fatalf("401 round retry = (%v, %t), want (0, false)", wait, shouldRetry) + } + }) + } +} + +func TestHomeRetryRoundStartsImmediatelyWhenHomeReportsAvailableNextRound(t *testing.T) { + tests := []struct { + name string + execute func(*Manager, *int) error + }{ + { + name: "nonstream", + execute: func(manager *Manager, retryLimit *int) error { + _, errExecute := manager.executeHomeOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}, false, 0, retryLimit) + return errExecute + }, + }, + { + name: "stream", + execute: func(manager *Manager, retryLimit *int) error { + _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}, 0, retryLimit) + return errExecute + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + dispatcher := &retryContractHomeDispatcher{ + authIDs: []string{"home-retry-a", "home-retry-b"}, + exhaustedPayload: []byte(`{"error":{"type":"auth_unavailable","message":"a credential is immediately available next round"}}`), + } + executor := &retryContractHomeExecutor{ + failAll: true, + failures: map[string]error{ + "home-retry-a": retryContractRateLimitError{retryAfter: 5 * time.Second}, + "home-retry-b": &Error{HTTPStatus: http.StatusBadGateway, Message: "upstream unavailable"}, + }, + } + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(1, 10*time.Second, 0) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + retryLimit := -1 + errExecute := test.execute(manager, &retryLimit) + if !isHomeRetryRoundExhausted(errExecute) { + t.Fatalf("execution error = %v, want exhausted retry round", errExecute) + } + wait, shouldRetry := manager.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), cliproxyexecutor.Options{}, errExecute, 0, []string{"home-retry-contract"}, "gpt", 10*time.Second, retryLimit) + if !shouldRetry || wait != 0 { + t.Fatalf("next-round retry = (%v, %t), want immediate", wait, shouldRetry) + } + }) + } +} + +func TestHomeRetryRoundUsesRemoteCooldownWhenAttemptedErrorHasNoTiming(t *testing.T) { + tests := []struct { + name string + execute func(*Manager, *int) error + }{ + { + name: "nonstream", + execute: func(manager *Manager, retryLimit *int) error { + _, errExecute := manager.executeHomeOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}, false, 2, retryLimit) + return errExecute + }, + }, + { + name: "stream", + execute: func(manager *Manager, retryLimit *int) error { + _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}, 2, retryLimit) + return errExecute + }, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + dispatcher := &retryContractHomeDispatcher{ + authIDs: []string{"home-retry-a"}, + exhaustedPayload: []byte(`{"error":{"type":"model_cooldown","message":"remaining Home credentials are cooling down","retryable":true,"retry_after_ms":1500}}`), + } + executor := &retryContractHomeExecutor{ + failAll: true, + failure: &Error{HTTPStatus: http.StatusBadGateway, Message: "upstream unavailable"}, + } + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(1, 2*time.Second, 2) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + retryLimit := -1 + errExecute := tc.execute(manager, &retryLimit) + if !isHomeRetryRoundExhausted(errExecute) { + t.Fatalf("execution error = %v, want exhausted retry round", errExecute) + } + retryAfter := retryAfterFromError(errExecute) + if retryAfter == nil || *retryAfter != 1500*time.Millisecond { + t.Fatalf("retry after = %v, want remote cooldown delay 1500ms", retryAfter) + } + if got := SafeResponseHeaders(errExecute).Get("Retry-After"); got != "2" { + t.Fatalf("safe Retry-After header = %q, want 2", got) + } + }) + } +} + +func TestHomeStreamOAuthUnauthorizedRotatesAfterRefreshRetry(t *testing.T) { + dispatcher := &retryContractHomeDispatcher{ + authIDs: []string{"home-retry-a", "home-retry-b"}, + metadata: map[string]any{ + "auth_kind": "oauth", + }, + } + executor := &retryContractHomeExecutor{ + failures: map[string]error{ + "home-retry-a": &Error{HTTPStatus: http.StatusUnauthorized, Message: "expired"}, + }, + } + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(0, time.Second, 2) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + result, errExecute := manager.ExecuteStream(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}) + if errExecute != nil { + t.Fatalf("ExecuteStream() error = %v", errExecute) + } + for range result.Chunks { + } + if got := executor.Calls(); len(got) != 3 || got[0] != "home-retry-a" || got[1] != "home-retry-a" || got[2] != "home-retry-b" { + t.Fatalf("executor calls = %v, want [home-retry-a home-retry-a home-retry-b]", got) + } + excluded := dispatcher.Excluded() + if len(excluded) != 2 || len(excluded[0]) != 0 || len(excluded[1]) != 1 || excluded[1][0] != "home-retry-a" { + t.Fatalf("Home excluded auth IDs = %v, want [[], [home-retry-a]]", excluded) + } +} + +func TestHomeStreamLifecycleRecoveryFailureRotatesWithoutExtraDispatch(t *testing.T) { + dispatcher := &retryContractHomeDispatcher{authIDs: []string{"home-retry-a", "home-retry-b"}} + executor := &retryContractHomeExecutor{ + failures: map[string]error{ + "home-retry-a": errors.New("unexpected EOF"), + }, + } + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(0, time.Second, 2) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + result, errExecute := manager.ExecuteStream(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}) + if errExecute != nil { + t.Fatalf("ExecuteStream() error = %v", errExecute) + } + for range result.Chunks { + } + if got := executor.Calls(); len(got) != 3 || got[0] != "home-retry-a" || got[1] != "home-retry-a" || got[2] != "home-retry-b" { + t.Fatalf("executor calls = %v, want [home-retry-a home-retry-a home-retry-b]", got) + } + excluded := dispatcher.Excluded() + if len(excluded) != 3 || len(excluded[0]) != 0 || len(excluded[1]) != 0 || len(excluded[2]) != 1 || excluded[2][0] != "home-retry-a" { + t.Fatalf("Home excluded auth IDs = %v, want [[], [], [home-retry-a]]", excluded) + } +} + +func TestRetryRoundAvailabilityRejectsStaleQuotaForNonRetryableStatus(t *testing.T) { + now := time.Now() + for _, test := range []struct { + name string + lastError *Error + want bool + }{ + {name: "implicit quota", want: true}, + {name: "rate limit", lastError: &Error{HTTPStatus: http.StatusTooManyRequests}, want: true}, + {name: "payment required", lastError: &Error{HTTPStatus: http.StatusPaymentRequired}, want: false}, + {name: "not found", lastError: &Error{HTTPStatus: http.StatusNotFound}, want: false}, + } { + t.Run(test.name, func(t *testing.T) { + nextRetry := now.Add(time.Minute) + auth := &Auth{ + ID: "retry-round-stale-quota", + Provider: "codex", + Status: StatusActive, + ModelStates: map[string]*ModelState{ + "gpt": { + Status: StatusError, + Unavailable: true, + NextRetryAfter: nextRetry, + LastError: test.lastError, + Quota: QuotaState{Exceeded: true, NextRecoverAt: nextRetry}, + }, + }, + } + got, next := retryRoundAvailabilityForAuth(auth, "gpt", now) + if got != test.want { + t.Fatalf("retryRoundAvailabilityForAuth() eligible = %t, want %t", got, test.want) + } + if got && !next.Equal(nextRetry) { + t.Fatalf("retryRoundAvailabilityForAuth() next = %v, want %v", next, nextRetry) + } + }) + } +} + +func TestHomeStreamAPIKeyUnauthorizedRotatesImmediately(t *testing.T) { + dispatcher := &retryContractHomeDispatcher{ + authIDs: []string{"home-retry-a", "home-retry-b"}, + metadata: map[string]any{ + "auth_kind": "apikey", + }, + } + executor := &retryContractHomeExecutor{ + failures: map[string]error{ + "home-retry-a": &Error{HTTPStatus: http.StatusUnauthorized, Message: "invalid api key"}, + }, + } + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(0, time.Second, 2) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + result, errExecute := manager.ExecuteStream(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}) + if errExecute != nil { + t.Fatalf("ExecuteStream() error = %v", errExecute) + } + for range result.Chunks { + } + if got := executor.Calls(); len(got) != 2 || got[0] != "home-retry-a" || got[1] != "home-retry-b" { + t.Fatalf("executor calls = %v, want [home-retry-a home-retry-b]", got) + } + excluded := dispatcher.Excluded() + if len(excluded) != 2 || len(excluded[0]) != 0 || len(excluded[1]) != 1 || excluded[1][0] != "home-retry-a" { + t.Fatalf("Home excluded auth IDs = %v, want [[], [home-retry-a]]", excluded) + } +} + +func TestHomeModelCooldownErrorPreservesRetryContract(t *testing.T) { + errDecoded := decodeHomeDispatchError([]byte(`{"error":{"type":"model_cooldown","message":"all credentials are cooling down","retryable":true,"retry_after_ms":1500,"request_retry":2}}`)) + var retryErr *homeDispatchRetryAfterError + if !errors.As(errDecoded, &retryErr) || retryErr == nil { + t.Fatalf("decodeHomeDispatchError() = %#v, want retry-after error", errDecoded) + } + if retryErr.StatusCode() != http.StatusTooManyRequests || retryErr.RetryAfter() == nil || *retryErr.RetryAfter() != 1500*time.Millisecond { + t.Fatalf("decoded Home cooldown = status %d retry-after %v, want 429/1500ms", retryErr.StatusCode(), retryErr.RetryAfter()) + } + if retryLimit, ok := retryErr.RequestRetryLimit(); !ok || retryLimit != 2 { + t.Fatalf("decoded Home request retry limit = (%d, %t), want (2, true)", retryLimit, ok) + } + var cause *Error + if !errors.As(errDecoded, &cause) || cause == nil || cause.Code != "model_cooldown" || !cause.Retryable { + t.Fatalf("decoded Home cooldown cause = %#v, want retryable model_cooldown", cause) + } + if got := SafeResponseHeaders(errDecoded).Get("Retry-After"); got != "2" { + t.Fatalf("safe Retry-After header = %q, want 2", got) + } +} + +func TestHomeRequestRetryCountsAdditionalCredentialRounds(t *testing.T) { + for _, stream := range []bool{false, true} { + t.Run(map[bool]string{false: "nonstream", true: "stream"}[stream], func(t *testing.T) { + dispatcher := &retryContractHomeDispatcher{authIDs: []string{"home-retry-a", "home-retry-b"}} + executor := &retryContractHomeExecutor{failAll: true} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(1, time.Second, 2) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + if stream { + result, errExecute := manager.ExecuteStream(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}) + if errExecute == nil || result != nil { + t.Fatalf("ExecuteStream() = result %#v, error %v; want terminal rate-limit error", result, errExecute) + } + } else { + _, errExecute := manager.Execute(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}) + if errExecute == nil { + t.Fatal("Execute() error = nil, want rate-limit error") + } + } + if got := executor.Calls(); len(got) != 4 { + t.Fatalf("executor calls = %v, want four calls across two rounds", got) + } + excluded := dispatcher.Excluded() + if len(excluded) != 4 || len(excluded[0]) != 0 || len(excluded[1]) != 1 || excluded[1][0] != "home-retry-a" || len(excluded[2]) != 0 || len(excluded[3]) != 1 || excluded[3][0] != "home-retry-a" { + t.Fatalf("Home excluded auth IDs = %v, want [[], [home-retry-a], [], [home-retry-a]]", excluded) + } + }) + } +} + +func TestHomeRequestRetryRoundDoesNotRequireRetryAfter(t *testing.T) { + for _, stream := range []bool{false, true} { + t.Run(map[bool]string{false: "nonstream", true: "stream"}[stream], func(t *testing.T) { + dispatcher := &retryContractHomeDispatcher{authIDs: []string{"home-retry-a", "home-retry-b"}} + executor := &retryContractHomeExecutor{ + failAll: true, + failure: &Error{HTTPStatus: http.StatusBadGateway, Message: "upstream unavailable"}, + } + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(1, time.Second, 2) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + if stream { + result, errExecute := manager.ExecuteStream(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}) + if errExecute == nil || result != nil { + t.Fatalf("ExecuteStream() = result %#v, error %v; want terminal upstream error", result, errExecute) + } + } else { + _, errExecute := manager.Execute(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}) + if errExecute == nil { + t.Fatal("Execute() error = nil, want upstream error") + } + } + if got := executor.Calls(); len(got) != 4 { + t.Fatalf("executor calls = %v, want four calls across two rounds", got) + } + }) + } +} + +func TestHomeStreamLegacyDispatcherDoesNotSpinOnIgnoredExclusions(t *testing.T) { + dispatcher := &legacyRepeatedStreamDispatcher{} + executor := &retryContractHomeExecutor{ + failAll: true, + failure: &Error{HTTPStatus: http.StatusBadGateway, Message: "upstream unavailable"}, + } + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(1, time.Second, 2) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + result, errExecute := manager.ExecuteStream(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}) + if result != nil || errExecute == nil { + t.Fatalf("ExecuteStream() = result %#v, error %v; want terminal upstream error", result, errExecute) + } + if got := len(executor.Calls()); got != 2 { + t.Fatalf("executor calls = %d, want one attempt in each of two rounds", got) + } + if got := dispatcher.calls.Load(); got != 4 { + t.Fatalf("legacy Home dispatch calls = %d, want two dispatches in each of two rounds", got) + } +} + +func TestHomeNonStreamLegacyDispatcherCompletesAdditionalRetryRound(t *testing.T) { + dispatcher := &legacyRepeatedStreamDispatcher{} + executor := &retryContractHomeExecutor{ + failAll: true, + failure: &Error{HTTPStatus: http.StatusBadGateway, Message: "upstream unavailable"}, + } + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(1, time.Second, 0) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + _, errExecute := manager.Execute(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}) + if errExecute == nil { + t.Fatal("Execute() error = nil, want terminal upstream error") + } + if got := len(executor.Calls()); got != 2 { + t.Fatalf("executor calls = %d, want one attempt in each of two rounds", got) + } + if got := dispatcher.calls.Load(); got != 4 { + t.Fatalf("legacy Home dispatch calls = %d, want two dispatches in each of two rounds", got) + } +} + +func TestHomeLocalSelectionRejectionWaitsForReleaseAcknowledgement(t *testing.T) { + tests := []struct { + name string + dispatcher homeAuthDispatcher + executor *retryContractHomeExecutor + maxRetryCredentials int + blockedGroup executionregistry.ReleaseGroup + blockedSequence int64 + execute func(*Manager, int, *int) error + }{ + { + name: "nonstream repeated auth", + dispatcher: &accountedHomeExecutionDispatcher{auths: []Auth{ + {ID: "home-retry-a", Provider: "home-retry-contract", Status: StatusActive}, + {ID: "home-retry-a", Provider: "home-retry-contract", Status: StatusActive}, + }}, + executor: &retryContractHomeExecutor{failure: &Error{HTTPStatus: http.StatusBadGateway, Message: "upstream unavailable"}}, + maxRetryCredentials: 0, + blockedGroup: executionregistry.ReleaseGroup{CredentialID: "home-retry-a", Model: "gpt"}, + blockedSequence: 2, + execute: func(manager *Manager, maxRetryCredentials int, retryLimit *int) error { + _, errExecute := manager.executeHomeOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}, false, maxRetryCredentials, retryLimit) + return errExecute + }, + }, + { + name: "stream repeated excluded auth", + dispatcher: &accountedHomeExecutionDispatcher{auths: []Auth{ + {ID: "home-retry-a", Provider: "home-retry-contract", Status: StatusActive}, + {ID: "home-retry-a", Provider: "home-retry-contract", Status: StatusActive}, + }}, + executor: &retryContractHomeExecutor{failure: &Error{HTTPStatus: http.StatusBadGateway, Message: "upstream unavailable"}}, + maxRetryCredentials: 0, + blockedGroup: executionregistry.ReleaseGroup{CredentialID: "home-retry-a", Model: "gpt"}, + blockedSequence: 2, + execute: func(manager *Manager, maxRetryCredentials int, retryLimit *int) error { + _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}, maxRetryCredentials, retryLimit) + return errExecute + }, + }, + { + name: "stream max retry credentials", + dispatcher: &accountedHomeExecutionDispatcher{auths: []Auth{ + {ID: "home-retry-a", Provider: "home-retry-contract", Status: StatusActive}, + {ID: "home-retry-b", Provider: "home-retry-contract", Status: StatusActive}, + }}, + executor: &retryContractHomeExecutor{failure: errors.New("unexpected EOF")}, + maxRetryCredentials: 1, + blockedGroup: executionregistry.ReleaseGroup{CredentialID: "home-retry-b", Model: "gpt"}, + blockedSequence: 1, + execute: func(manager *Manager, maxRetryCredentials int, retryLimit *int) error { + _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}, maxRetryCredentials, retryLimit) + return errExecute + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + registry := executionregistry.New() + acknowledged := make(chan struct{}) + close(acknowledged) + unacknowledged := make(chan struct{}) + var blockedReleaseSeen atomic.Bool + registry.SetReleaseSink(func(group executionregistry.ReleaseGroup, sequence int64) *executionregistry.ReleaseTicket { + done := (<-chan struct{})(acknowledged) + if group == test.blockedGroup && sequence == test.blockedSequence { + blockedReleaseSeen.Store(true) + done = unacknowledged + } + return executionregistry.NewReleaseTicket(group, sequence, done) + }) + + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{ + Home: internalconfig.HomeConfig{Enabled: true}, + CredentialConcurrency: internalconfig.CredentialConcurrencyConfig{CPACancelBound: 10 * time.Millisecond}, + }) + manager.PublishHomeDispatch(test.dispatcher, registry, 1) + manager.RegisterExecutor(test.executor) + + retryLimit := -1 + errExecute := test.execute(manager, test.maxRetryCredentials, &retryLimit) + if !blockedReleaseSeen.Load() { + t.Fatal("target release was not attempted") + } + var homeErr *Error + if !errors.As(errExecute, &homeErr) || homeErr == nil || homeErr.Code != "home_unavailable" { + t.Fatalf("execution error = %T %v, want Home release acknowledgement timeout", errExecute, errExecute) + } + }) + } +} + +func TestHomeRetryRoundHonorsCredentialRequestRetryOverride(t *testing.T) { + tests := []struct { + name string + globalRetry int + override int + wantCallCount int + }{ + {name: "override disables global rounds", globalRetry: 3, override: 0, wantCallCount: 2}, + {name: "override enables rounds over global", globalRetry: 0, override: 1, wantCallCount: 4}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + dispatcher := &retryContractHomeDispatcher{ + authIDs: []string{"home-retry-a", "home-retry-b"}, + metadata: map[string]any{"request_retry": tc.override}, + } + executor := &retryContractHomeExecutor{failAll: true} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(tc.globalRetry, time.Second, 2) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + _, errExecute := manager.Execute(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}) + if errExecute == nil { + t.Fatal("Execute() error = nil, want terminal rate-limit error") + } + if got := len(executor.Calls()); got != tc.wantCallCount { + t.Fatalf("executor call count = %d, want %d", got, tc.wantCallCount) + } + }) + } +} + +func TestHomeRetryRoundUsesSuccessfulDispatchAggregate(t *testing.T) { + tests := []struct { + name string + execute func(*Manager) error + }{ + { + name: "nonstream", + execute: func(manager *Manager) error { + _, errExecute := manager.Execute(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}) + return errExecute + }, + }, + { + name: "count tokens", + execute: func(manager *Manager) error { + _, errExecute := manager.ExecuteCount(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}) + return errExecute + }, + }, + { + name: "stream", + execute: func(manager *Manager) error { + result, errExecute := manager.ExecuteStream(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}) + if errExecute != nil { + return errExecute + } + for range result.Chunks { + } + return nil + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + dispatcher := &aggregateRetryHomeDispatcher{} + executor := &retryContractHomeExecutor{} + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(0, time.Second, 1) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + if errExecute := test.execute(manager); errExecute != nil { + t.Fatalf("execution error = %v", errExecute) + } + if got := executor.Calls(); len(got) != 2 || got[0] != "home-retry-a" || got[1] != "home-retry-b" { + t.Fatalf("executor calls = %v, want [home-retry-a home-retry-b]", got) + } + if got := dispatcher.calls.Load(); got != 2 { + t.Fatalf("Home dispatch calls = %d, want 2", got) + } + }) + } +} + +func TestHomeRetryRoundUsesAuthoritativeZeroAggregate(t *testing.T) { + tests := []struct { + name string + execute func(*Manager) error + }{ + { + name: "nonstream", + execute: func(manager *Manager) error { + _, errExecute := manager.Execute(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}) + return errExecute + }, + }, + { + name: "count tokens", + execute: func(manager *Manager) error { + _, errExecute := manager.ExecuteCount(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{}) + return errExecute + }, + }, + { + name: "stream", + execute: func(manager *Manager) error { + _, errExecute := manager.ExecuteStream(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}) + return errExecute + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + remoteRetry := 0 + dispatcher := &retryContractHomeDispatcher{ + authIDs: []string{"home-retry-a", "home-retry-b"}, + metadata: map[string]any{"request_retry": 3}, + requestRetry: &remoteRetry, + } + executor := &retryContractHomeExecutor{ + failAll: true, + failure: &Error{HTTPStatus: http.StatusBadGateway, Message: "upstream unavailable"}, + } + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(3, time.Second, 2) + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + manager.RegisterExecutor(executor) + + if errExecute := test.execute(manager); errExecute == nil { + t.Fatal("execution error = nil, want terminal first-round error") + } + if got := executor.Calls(); len(got) != 2 { + t.Fatalf("executor calls = %v, want only the two first-round credentials", got) + } + }) + } +} diff --git a/sdk/cliproxy/auth/home_selection.go b/sdk/cliproxy/auth/home_selection.go index 33584be6..815d6d36 100644 --- a/sdk/cliproxy/auth/home_selection.go +++ b/sdk/cliproxy/auth/home_selection.go @@ -144,6 +144,8 @@ type HomeDispatchSelection struct { authMu sync.RWMutex scope *executionregistry.Scope accountedModel string + requestRetry int + hasRequestRetry bool resources *executionResources attemptCancels *attemptCancels once sync.Once -- knotmirror 0.0.1