diff --git a/config.example.yaml b/config.example.yaml index 7ee54309..498d37cd 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -142,15 +142,21 @@ force-model-prefix: false passthrough-headers: false # 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. +# its eligible credentials. Round 0 is the initial round; round r only admits +# credentials whose effective request-retry is at least r. Explicit non-negative +# credential/provider overrides take precedence; omitted or negative overrides +# inherit this global value, and explicit 0 only admits round 0. New CPA nodes +# send retry_round=0 for the initial round and increment it for additional rounds; +# legacy dispatch methods omit the field and keep old semantics. # 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 in each credential retry round. -# Set to 0 to keep legacy behavior (try all available credentials). +# Maximum number of different credentials to try in each credential retry round +# after per-credential round filtering. Set to 0 to try all available +# credentials. Credentials skipped by this cap still age with the global round, +# so the cap does not guarantee a fixed number of actual retries per credential. max-retry-credentials: 0 # Maximum cooldown wait in seconds between retry rounds. diff --git a/internal/home/client.go b/internal/home/client.go index c0424e97..e4bdfa4c 100644 --- a/internal/home/client.go +++ b/internal/home/client.go @@ -1255,9 +1255,9 @@ func newAuthDispatchRequest(requestedModel string, sessionID string, headers htt } 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. + // 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. New Home servers apply retry_round eligibility remotely. count = 1 values := append([]string{}, (*excludedAuthIDs)...) excludedAuthIDsCopy = &values @@ -1275,28 +1275,49 @@ func newAuthDispatchRequest(requestedModel string, sessionID string, headers htt } } +func newAuthDispatchRequestWithRetryRound(requestedModel string, sessionID string, headers http.Header, count int, credentialPolicy string, retryRound int, excludedAuthIDs *[]string, pinnedAuthID string) authDispatchRequest { + req := newAuthDispatchRequest(requestedModel, sessionID, headers, count, credentialPolicy, excludedAuthIDs, pinnedAuthID) + if retryRound < 0 { + retryRound = 0 + } + req.RetryRound = &retryRound + return req +} + 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, "", nil, "") + return c.rPopAuth(ctx, requestedModel, sessionID, headers, count, "", nil, 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, nil, "") + return c.rPopAuth(ctx, requestedModel, sessionID, headers, count, credentialPolicy, nil, 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) + return c.rPopAuth(ctx, requestedModel, sessionID, headers, count, "", nil, &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) + return c.rPopAuth(ctx, requestedModel, sessionID, headers, count, credentialPolicy, nil, &excludedAuthIDs, pinnedAuthID) +} + +// RPopAuthWithRetryRoundConstraints requests a credential with the retry round, +// current-round exclusions, and optional pinned credential constraint. +func (c *Client) RPopAuthWithRetryRoundConstraints(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int, retryRound int, excludedAuthIDs []string, pinnedAuthID string) ([]byte, error) { + return c.rPopAuth(ctx, requestedModel, sessionID, headers, count, "", &retryRound, &excludedAuthIDs, pinnedAuthID) +} + +// RPopAuthWithPolicyAndRetryRoundConstraints combines a credential policy with +// the retry round, current-round exclusions, and optional pin. +func (c *Client) RPopAuthWithPolicyAndRetryRoundConstraints(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int, credentialPolicy string, retryRound int, excludedAuthIDs []string, pinnedAuthID string) ([]byte, error) { + return c.rPopAuth(ctx, requestedModel, sessionID, headers, count, credentialPolicy, &retryRound, &excludedAuthIDs, pinnedAuthID) } -func (c *Client) rPopAuth(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int, credentialPolicy string, excludedAuthIDs *[]string, pinnedAuthID string) ([]byte, error) { +func (c *Client) rPopAuth(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int, credentialPolicy string, retryRound *int, excludedAuthIDs *[]string, pinnedAuthID string) ([]byte, error) { if c == nil || c.dispatchFenced.Load() { return nil, ErrDispatchFenced } @@ -1310,7 +1331,12 @@ 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, excludedAuthIDs, pinnedAuthID) + var req authDispatchRequest + if retryRound == nil { + req = newAuthDispatchRequest(requestedModel, sessionID, headers, count, credentialPolicy, excludedAuthIDs, pinnedAuthID) + } else { + req = newAuthDispatchRequestWithRetryRound(requestedModel, sessionID, headers, count, credentialPolicy, *retryRound, 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 6677acd5..07187764 100644 --- a/internal/home/client_test.go +++ b/internal/home/client_test.go @@ -131,6 +131,54 @@ func TestAuthDispatchRequestIncludesPinnedAuthID(t *testing.T) { } } +func TestAuthDispatchRequestDistinguishesLegacyAndRetryRoundProtocol(t *testing.T) { + excludedAuthIDs := []string{"auth-a"} + legacy := newAuthDispatchRequest("gpt-5.4", "", nil, 3, "", &excludedAuthIDs, "") + legacyRaw, errMarshal := json.Marshal(&legacy) + if errMarshal != nil { + t.Fatalf("marshal legacy auth dispatch request: %v", errMarshal) + } + var legacyPayload map[string]any + if errUnmarshal := json.Unmarshal(legacyRaw, &legacyPayload); errUnmarshal != nil { + t.Fatalf("unmarshal legacy auth dispatch request: %v", errUnmarshal) + } + if _, present := legacyPayload["retry_round"]; present { + t.Fatalf("legacy request unexpectedly included retry_round: %#v", legacyPayload["retry_round"]) + } + + initial := newAuthDispatchRequestWithRetryRound("gpt-5.4", "", nil, 3, "", 0, &excludedAuthIDs, "") + initialRaw, errMarshal := json.Marshal(&initial) + if errMarshal != nil { + t.Fatalf("marshal initial auth dispatch request: %v", errMarshal) + } + var initialPayload map[string]any + if errUnmarshal := json.Unmarshal(initialRaw, &initialPayload); errUnmarshal != nil { + t.Fatalf("unmarshal initial auth dispatch request: %v", errUnmarshal) + } + if got := int(initialPayload["retry_round"].(float64)); got != 0 { + t.Fatalf("initial retry_round = %d, want explicit 0", got) + } + if got := int(initialPayload["count"].(float64)); got != 1 { + t.Fatalf("initial retry-contract count = %d, want 1", got) + } + + additional := newAuthDispatchRequestWithRetryRound("gpt-5.4", "", nil, 3, "", 2, &excludedAuthIDs, "") + additionalRaw, errMarshal := json.Marshal(&additional) + if errMarshal != nil { + t.Fatalf("marshal additional auth dispatch request: %v", errMarshal) + } + var additionalPayload map[string]any + if errUnmarshal := json.Unmarshal(additionalRaw, &additionalPayload); errUnmarshal != nil { + t.Fatalf("unmarshal additional auth dispatch request: %v", errUnmarshal) + } + if got := int(additionalPayload["retry_round"].(float64)); got != 2 { + t.Fatalf("retry_round = %d, want 2", got) + } + if got := additionalPayload["excluded_auth_ids"].([]any); len(got) != 1 || got[0] != "auth-a" { + t.Fatalf("excluded_auth_ids = %#v, want [auth-a]", additionalPayload["excluded_auth_ids"]) + } +} + func TestRedisOptionsHomeTLSDisabled(t *testing.T) { client := New(config.HomeConfig{ Enabled: true, diff --git a/internal/home/requests.go b/internal/home/requests.go index 30c7c5ff..a19a81d0 100644 --- a/internal/home/requests.go +++ b/internal/home/requests.go @@ -10,6 +10,7 @@ type authDispatchRequest struct { SessionID string `json:"session_id,omitempty"` Headers map[string]string `json:"headers,omitempty"` CredentialPolicy string `json:"credential_policy,omitempty"` + RetryRound *int `json:"retry_round,omitempty"` ExcludedAuthIDs *[]string `json:"excluded_auth_ids,omitempty"` PinnedAuthID string `json:"pinned_auth_id,omitempty"` } diff --git a/sdk/cliproxy/auth/conductor_execution.go b/sdk/cliproxy/auth/conductor_execution.go index 45b405d4..c98508ab 100644 --- a/sdk/cliproxy/auth/conductor_execution.go +++ b/sdk/cliproxy/auth/conductor_execution.go @@ -49,12 +49,12 @@ func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxye return resp, unwrapRequestStopError(errHome) } - _, maxRetryCredentials, maxWait := m.retrySettings() + defaultRequestRetry, maxRetryCredentials, maxWait := m.retrySettings() var lastErr error retryModel := authSelectionModelFromOptions(opts, req.Model) for attempt := 0; ; attempt++ { - resp, errExec := m.executeMixedOnce(ctx, normalized, req, opts, maxRetryCredentials) + resp, errExec := m.executeMixedOnce(ctx, normalized, req, opts, maxRetryCredentials, attempt, defaultRequestRetry) if errExec == nil { return resp, nil } @@ -62,7 +62,7 @@ func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxye return cliproxyexecutor.Response{}, unwrapRequestStopError(errExec) } lastErr = errExec - wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errExec, attempt, normalized, retryModel, maxWait, -1) + wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errExec, attempt, normalized, retryModel, maxWait, -1, defaultRequestRetry) if !shouldRetry { break } @@ -96,12 +96,12 @@ func (m *Manager) ExecuteCount(ctx context.Context, providers []string, req clip return resp, unwrapRequestStopError(errHome) } - _, maxRetryCredentials, maxWait := m.retrySettings() + defaultRequestRetry, maxRetryCredentials, maxWait := m.retrySettings() var lastErr error retryModel := authSelectionModelFromOptions(opts, req.Model) for attempt := 0; ; attempt++ { - resp, errExec := m.executeCountMixedOnce(ctx, normalized, req, opts, maxRetryCredentials) + resp, errExec := m.executeCountMixedOnce(ctx, normalized, req, opts, maxRetryCredentials, attempt, defaultRequestRetry) if errExec == nil { return resp, nil } @@ -109,7 +109,7 @@ func (m *Manager) ExecuteCount(ctx context.Context, providers []string, req clip return cliproxyexecutor.Response{}, unwrapRequestStopError(errExec) } lastErr = errExec - wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errExec, attempt, normalized, retryModel, maxWait, -1) + wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errExec, attempt, normalized, retryModel, maxWait, -1, defaultRequestRetry) if !shouldRetry { break } @@ -137,7 +137,7 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli return nil, &Error{Code: "provider_not_found", Message: "no provider supplied"} } - _, maxRetryCredentials, maxWait := m.retrySettings() + defaultRequestRetry, maxRetryCredentials, maxWait := m.retrySettings() var lastErr error homeRetryLimit := -1 @@ -146,7 +146,7 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli retryRoundPending := false retryRoundWaited := false for { - result, errStream := m.executeStreamMixedOnce(ctx, normalized, req, opts, maxRetryCredentials, &homeRetryLimit) + result, errStream := m.executeStreamMixedOnce(ctx, normalized, req, opts, maxRetryCredentials, &homeRetryLimit, attempt, defaultRequestRetry) if errStream == nil { return result, nil } @@ -168,7 +168,7 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli return nil, unwrapRequestStopError(errStream) } lastErr = errStream - wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errStream, attempt, normalized, retryModel, maxWait, homeRetryLimit) + wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errStream, attempt, normalized, retryModel, maxWait, homeRetryLimit, defaultRequestRetry) if !shouldRetry { break } @@ -300,7 +300,7 @@ func mergeRequestHeaders(current, updates http.Header, clear []string) http.Head return out } -func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int) (cliproxyexecutor.Response, error) { +func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int, retryRound int, defaultRequestRetry int) (cliproxyexecutor.Response, error) { if len(providers) == 0 { return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"} } @@ -310,6 +310,11 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req homeMode := m.HomeEnabled() homeAuthCount := 1 tried := make(map[string]struct{}) + if !homeMode { + for authID := range m.requestRetryRoundExclusions(retryRound, defaultRequestRetry) { + tried[authID] = struct{}{} + } + } attempted := make(map[string]struct{}) var lastErr error for { @@ -321,7 +326,8 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req } pickOpts := opts if homeMode { - pickOpts = withHomeAuthCount(opts, homeAuthCount) + pickOpts = withHomeRetryRound(pickOpts, retryRound) + pickOpts = withHomeAuthCount(pickOpts, homeAuthCount) pickOpts = withHomeExcludedAuthIDs(pickOpts, tried) } auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried) @@ -468,7 +474,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req } } -func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int) (cliproxyexecutor.Response, error) { +func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int, retryRound int, defaultRequestRetry int) (cliproxyexecutor.Response, error) { if len(providers) == 0 { return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"} } @@ -478,6 +484,11 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, homeMode := m.HomeEnabled() homeAuthCount := 1 tried := make(map[string]struct{}) + if !homeMode { + for authID := range m.requestRetryRoundExclusions(retryRound, defaultRequestRetry) { + tried[authID] = struct{}{} + } + } attempted := make(map[string]struct{}) var lastErr error for { @@ -489,7 +500,8 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, } pickOpts := opts if homeMode { - pickOpts = withHomeAuthCount(opts, homeAuthCount) + pickOpts = withHomeRetryRound(pickOpts, retryRound) + pickOpts = withHomeAuthCount(pickOpts, homeAuthCount) pickOpts = withHomeExcludedAuthIDs(pickOpts, tried) } auth, executor, provider, errPick := m.pickNextMixed(ctx, providers, routeModel, pickOpts, tried) @@ -640,7 +652,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, homeRetryLimit *int) (*cliproxyexecutor.StreamResult, error) { +func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, maxRetryCredentials int, homeRetryLimit *int, retryRound int, defaultRequestRetry int) (*cliproxyexecutor.StreamResult, error) { if len(providers) == 0 { return nil, &Error{Code: "provider_not_found", Message: "no provider supplied"} } @@ -651,6 +663,11 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string homeMode := m.HomeEnabled() homeAuthCount := 1 tried := make(map[string]struct{}) + if !homeMode { + for authID := range m.requestRetryRoundExclusions(retryRound, defaultRequestRetry) { + tried[authID] = struct{}{} + } + } homeExcludedAuthIDs := make(map[string]struct{}) homeSameAuthRetries := make(map[string]int) lastHomeAuthID := "" @@ -672,7 +689,8 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string } pickOpts := opts if homeMode { - pickOpts = withHomeAuthCount(opts, homeAuthCount) + pickOpts = withHomeRetryRound(pickOpts, retryRound) + pickOpts = withHomeAuthCount(pickOpts, homeAuthCount) pickOpts = withHomeExcludedAuthIDs(pickOpts, homeExcludedAuthIDs) } @@ -1004,6 +1022,20 @@ func withHomeAuthCount(opts cliproxyexecutor.Options, count int) cliproxyexecuto return opts } +func withHomeRetryRound(opts cliproxyexecutor.Options, retryRound int) cliproxyexecutor.Options { + meta := make(map[string]any, len(opts.Metadata)+1) + for key, value := range opts.Metadata { + meta[key] = value + } + if retryRound > 0 { + meta[homeRetryRoundMetadataKey] = retryRound + } else { + delete(meta, homeRetryRoundMetadataKey) + } + opts.Metadata = meta + 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 { diff --git a/sdk/cliproxy/auth/conductor_home.go b/sdk/cliproxy/auth/conductor_home.go index 92721e16..ce956a17 100644 --- a/sdk/cliproxy/auth/conductor_home.go +++ b/sdk/cliproxy/auth/conductor_home.go @@ -19,7 +19,8 @@ import ( ) const ( - homeAuthCountMetadataKey = "__cliproxy_home_auth_count" + homeAuthCountMetadataKey = "__cliproxy_home_auth_count" + homeRetryRoundMetadataKey = "request_retry_round" // ExcludedAuthIDsMetadataKey stores credential IDs already attempted in the // current request retry round. ExcludedAuthIDsMetadataKey = "excluded_auth_ids" @@ -357,6 +358,10 @@ type homeDispatchConstraintsDispatcher interface { RPopAuthWithConstraints(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int, excludedAuthIDs []string, pinnedAuthID string) ([]byte, error) } +type homeDispatchRetryRoundConstraintsDispatcher interface { + RPopAuthWithRetryRoundConstraints(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int, retryRound 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) } @@ -365,6 +370,10 @@ type homeCredentialPolicyConstraintsDispatcher interface { RPopAuthWithPolicyAndConstraints(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int, credentialPolicy string, excludedAuthIDs []string, pinnedAuthID string) ([]byte, error) } +type homeCredentialPolicyRetryRoundConstraintsDispatcher interface { + RPopAuthWithPolicyAndRetryRoundConstraints(ctx context.Context, requestedModel string, sessionID string, headers http.Header, count int, credentialPolicy string, retryRound int, excludedAuthIDs []string, pinnedAuthID string) ([]byte, error) +} + var currentHomeDispatcher = func() homeAuthDispatcher { return home.Current() } @@ -507,7 +516,7 @@ func (m *Manager) retainedHomeSessionSelection(ctx context.Context, opts cliprox routeModel, validRouteModel := validCanonicalHomeConcurrencyModelKey(model) var retained *HomeDispatchSelection var ended []*HomeDispatchSelection - fallbackAttempt := homeAuthCountFromMetadata(opts.Metadata) > 1 + fallbackAttempt := homeAuthCountFromMetadata(opts.Metadata) > 1 || homeRetryRoundFromMetadata(opts.Metadata) > 0 m.mu.Lock() selections := m.homeSessionSelections[sessionID] for key, selection := range selections { @@ -910,6 +919,7 @@ func (m *Manager) pickHomeDispatchSelection(ctx context.Context, model string, o requestedModel = requestedModelFromMetadata(opts.Metadata, model) } pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata) + retryRound := homeRetryRoundFromMetadata(opts.Metadata) excludedAuthIDList := homeExcludedAuthIDsFromMetadata(opts.Metadata) excludedAuthIDs := make(map[string]struct{}, len(excludedAuthIDList)) for _, authID := range excludedAuthIDList { @@ -955,11 +965,15 @@ func (m *Manager) pickHomeDispatchSelection(ctx context.Context, model string, o var raw []byte var errRPop error if credentialPolicy == "" { - if constrainedClient, okConstraints := client.(homeDispatchConstraintsDispatcher); okConstraints { + if retryRoundClient, okRetryRound := client.(homeDispatchRetryRoundConstraintsDispatcher); okRetryRound { + raw, errRPop = retryRoundClient.RPopAuthWithRetryRoundConstraints(ctx, requestedModel, sessionID, dispatchHeaders, homeAuthCountFromMetadata(opts.Metadata), retryRound, excludedAuthIDList, pinnedAuthID) + } else 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 retryRoundPolicyClient, okRetryRound := client.(homeCredentialPolicyRetryRoundConstraintsDispatcher); okRetryRound { + raw, errRPop = retryRoundPolicyClient.RPopAuthWithPolicyAndRetryRoundConstraints(ctx, requestedModel, sessionID, dispatchHeaders, homeAuthCountFromMetadata(opts.Metadata), credentialPolicy, retryRound, excludedAuthIDList, pinnedAuthID) } else if policyClient, okPolicy := client.(homeCredentialPolicyDispatcher); okPolicy { if constrainedClient, okConstraints := client.(homeCredentialPolicyConstraintsDispatcher); okConstraints { raw, errRPop = constrainedClient.RPopAuthWithPolicyAndConstraints(ctx, requestedModel, sessionID, dispatchHeaders, homeAuthCountFromMetadata(opts.Metadata), credentialPolicy, excludedAuthIDList, pinnedAuthID) @@ -1144,6 +1158,27 @@ func (m *Manager) pickHomeDispatchSelection(ctx context.Context, model string, o return selection, nil } +func homeRetryRoundFromMetadata(metadata map[string]any) int { + if metadata == nil { + return 0 + } + switch value := metadata[homeRetryRoundMetadataKey].(type) { + case int: + if value > 0 { + return value + } + case int64: + if value > 0 { + return int(value) + } + case float64: + if value > 0 && value == float64(int(value)) { + return int(value) + } + } + return 0 +} + func requestedModelFromMetadata(metadata map[string]any, fallback string) string { if metadata != nil { if v, ok := metadata[cliproxyexecutor.RequestedModelMetadataKey]; ok { diff --git a/sdk/cliproxy/auth/conductor_home_execution.go b/sdk/cliproxy/auth/conductor_home_execution.go index 204e7a32..d965712a 100644 --- a/sdk/cliproxy/auth/conductor_home_execution.go +++ b/sdk/cliproxy/auth/conductor_home_execution.go @@ -15,14 +15,14 @@ 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() + defaultRequestRetry, 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) + response, errExecute := m.executeHomeOnce(ctx, providers, req, opts, countTokens, maxRetryCredentials, &homeRetryLimit, attempt) if errExecute == nil { return response, nil } @@ -43,7 +43,7 @@ func (m *Manager) executeHome(ctx context.Context, providers []string, req clipr if isRequestTerminatedError(errExecute) || isRequestStopError(errExecute) { return cliproxyexecutor.Response{}, unwrapRequestStopError(errExecute) } - wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errExecute, attempt, providers, retryModel, maxWait, homeRetryLimit) + wait, shouldRetry := m.shouldRetryAfterErrorWithHomeRetryLimit(ctx, opts, errExecute, attempt, providers, retryModel, maxWait, homeRetryLimit, defaultRequestRetry) if !shouldRetry { return cliproxyexecutor.Response{}, unwrapRequestStopError(errExecute) } @@ -56,7 +56,11 @@ func (m *Manager) executeHome(ctx context.Context, providers []string, req clipr } } -func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, countTokens bool, maxRetryCredentials int, homeRetryLimit *int) (cliproxyexecutor.Response, error) { +func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, countTokens bool, maxRetryCredentials int, homeRetryLimit *int, retryRounds ...int) (cliproxyexecutor.Response, error) { + retryRound := 0 + if len(retryRounds) > 0 { + retryRound = retryRounds[0] + } routeModel := authSelectionModelFromOptions(opts, req.Model) responseAlias := requestedModelAliasFromOptions(opts, routeModel) executionModel, restoreExecutionModel := executionModelForAuthSelection(opts, req.Model) @@ -72,7 +76,8 @@ func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req c } return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"} } - pickOpts := withHomeAuthCount(opts, homeAuthCount) + pickOpts := withHomeRetryRound(opts, retryRound) + pickOpts = withHomeAuthCount(pickOpts, homeAuthCount) pickOpts = withHomeExcludedAuthIDs(pickOpts, tried) selection, errSelection := m.pickHomeDispatchSelection(ctx, routeModel, pickOpts) if errSelection != nil { diff --git a/sdk/cliproxy/auth/conductor_overrides_test.go b/sdk/cliproxy/auth/conductor_overrides_test.go index 69b51406..e8e635af 100644 --- a/sdk/cliproxy/auth/conductor_overrides_test.go +++ b/sdk/cliproxy/auth/conductor_overrides_test.go @@ -271,7 +271,7 @@ func TestManager_ShouldRetryAfterError_IgnoresRequestIneligibleOverrides(t *test } } - wait, shouldRetry := manager.shouldRetryAfterErrorWithHomeRetryLimit(test.ctx, test.opts, &Error{HTTPStatus: http.StatusBadGateway, Message: "retryable failure"}, 0, []string{"codex"}, model, 0, -1) + wait, shouldRetry := manager.shouldRetryAfterErrorWithHomeRetryLimit(test.ctx, test.opts, &Error{HTTPStatus: http.StatusBadGateway, Message: "retryable failure"}, 0, []string{"codex"}, model, 0, -1, 0) if shouldRetry || wait != 0 { t.Fatalf("request-ineligible override retry = (%v, %t), want (0, false)", wait, shouldRetry) } diff --git a/sdk/cliproxy/auth/conductor_retry_round_test.go b/sdk/cliproxy/auth/conductor_retry_round_test.go new file mode 100644 index 00000000..1d65b837 --- /dev/null +++ b/sdk/cliproxy/auth/conductor_retry_round_test.go @@ -0,0 +1,455 @@ +package auth + +import ( + "context" + "encoding/json" + "net/http" + "sort" + "sync" + "testing" + + 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 retryRoundCallExecutor struct { + identifier string + mu sync.Mutex + executeIDs []string + streamIDs []string + countIDs []string +} + +func (e *retryRoundCallExecutor) Identifier() string { return e.identifier } + +func (e *retryRoundCallExecutor) Execute(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + e.mu.Lock() + e.executeIDs = append(e.executeIDs, auth.ID) + e.mu.Unlock() + return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusInternalServerError, Message: "retry-round test failure"} +} + +func (e *retryRoundCallExecutor) ExecuteStream(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + e.mu.Lock() + e.streamIDs = append(e.streamIDs, auth.ID) + e.mu.Unlock() + return nil, &Error{HTTPStatus: http.StatusInternalServerError, Message: "retry-round test failure"} +} + +func (*retryRoundCallExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) { + return auth, nil +} + +func (e *retryRoundCallExecutor) CountTokens(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + e.mu.Lock() + e.countIDs = append(e.countIDs, auth.ID) + e.mu.Unlock() + return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusInternalServerError, Message: "retry-round test failure"} +} + +func (*retryRoundCallExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func (e *retryRoundCallExecutor) ids(kind string) []string { + e.mu.Lock() + defer e.mu.Unlock() + var source []string + switch kind { + case "execute": + source = e.executeIDs + case "stream": + source = e.streamIDs + case "count": + source = e.countIDs + } + return append([]string(nil), source...) +} + +func registerRetryRoundLocalAuths(t *testing.T, manager *Manager, provider, model string, limits map[string]int) []string { + t.Helper() + ids := make([]string, 0, len(limits)) + for id := range limits { + ids = append(ids, id) + } + sort.Strings(ids) + reg := registry.GetGlobalRegistry() + for _, id := range ids { + reg.RegisterClient(id, provider, []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { reg.UnregisterClient(id) }) + if _, errRegister := manager.Register(context.Background(), &Auth{ + ID: id, + Provider: provider, + Metadata: map[string]any{"request_retry": limits[id], "disable_cooling": true}, + }); errRegister != nil { + t.Fatalf("register %s: %v", id, errRegister) + } + } + return ids +} + +func countRetryRoundIDs(ids []string) map[string]int { + counts := make(map[string]int, len(ids)) + for _, id := range ids { + counts[id]++ + } + return counts +} + +func TestExecuteRetryRoundCredentialWindows(t *testing.T) { + tests := []struct { + name string + invoke func(*Manager, cliproxyexecutor.Request) error + kind string + }{ + { + name: "non-stream", + invoke: func(manager *Manager, req cliproxyexecutor.Request) error { + _, errExecute := manager.Execute(context.Background(), []string{"retry-round-test"}, req, cliproxyexecutor.Options{}) + return errExecute + }, + kind: "execute", + }, + { + name: "count-tokens", + invoke: func(manager *Manager, req cliproxyexecutor.Request) error { + _, errExecute := manager.ExecuteCount(context.Background(), []string{"retry-round-test"}, req, cliproxyexecutor.Options{}) + return errExecute + }, + kind: "count", + }, + { + name: "stream", + invoke: func(manager *Manager, req cliproxyexecutor.Request) error { + _, errExecute := manager.ExecuteStream(context.Background(), []string{"retry-round-test"}, req, cliproxyexecutor.Options{Stream: true}) + return errExecute + }, + kind: "stream", + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetRetryConfig(3, 0, 0) + executor := &retryRoundCallExecutor{identifier: "retry-round-test"} + manager.RegisterExecutor(executor) + registerRetryRoundLocalAuths(t, manager, "retry-round-test", "retry-round-model", map[string]int{ + "retry-round-a": 3, + "retry-round-b": 2, + "retry-round-c": 2, + }) + + if errExecute := test.invoke(manager, cliproxyexecutor.Request{Model: "retry-round-model"}); errExecute == nil { + t.Fatal("execution error = nil, want terminal retry error") + } + counts := countRetryRoundIDs(executor.ids(test.kind)) + if counts["retry-round-a"] != 4 || counts["retry-round-b"] != 3 || counts["retry-round-c"] != 3 { + t.Fatalf("credential call counts = %#v, want A=4 B=3 C=3; calls=%v", counts, executor.ids(test.kind)) + } + }) + } +} + +func TestExecuteRetryRoundMaxCredentialsAgesSkippedAuths(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetRetryConfig(2, 0, 3) + executor := &retryRoundCallExecutor{identifier: "retry-round-test"} + manager.RegisterExecutor(executor) + registerRetryRoundLocalAuths(t, manager, "retry-round-test", "retry-round-model-cap", map[string]int{ + "retry-cap-a": 1, + "retry-cap-b": 1, + "retry-cap-c": 1, + "retry-cap-d": 2, + }) + + if _, errExecute := manager.Execute(context.Background(), []string{"retry-round-test"}, cliproxyexecutor.Request{Model: "retry-round-model-cap"}, cliproxyexecutor.Options{}); errExecute == nil { + t.Fatal("execution error = nil, want terminal retry error") + } + calls := executor.ids("execute") + if len(calls) != 7 { + t.Fatalf("credential calls = %v, want three initial, three round-1, and one round-2 call", calls) + } + counts := countRetryRoundIDs(calls) + if counts["retry-cap-a"] > 2 || counts["retry-cap-b"] > 2 || counts["retry-cap-c"] > 2 || counts["retry-cap-d"] > 3 { + t.Fatalf("credential call counts exceed their round windows: %#v", counts) + } + if calls[len(calls)-1] != "retry-cap-d" { + t.Fatalf("last retry call = %q, want retry-cap-d; calls=%v", calls[len(calls)-1], calls) + } +} + +type retryConfigMutationExecutor struct { + identifier string + manager *Manager + nextDefault int + + mu sync.Mutex + calls int +} + +func (e *retryConfigMutationExecutor) Identifier() string { return e.identifier } + +func (e *retryConfigMutationExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + if e.recordCall() == 1 { + return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusInternalServerError, Message: "retry config changed"} + } + return cliproxyexecutor.Response{Payload: []byte("ok")}, nil +} + +func (e *retryConfigMutationExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { + if e.recordCall() == 1 { + return nil, &Error{HTTPStatus: http.StatusInternalServerError, Message: "retry config changed"} + } + chunks := make(chan cliproxyexecutor.StreamChunk) + close(chunks) + return &cliproxyexecutor.StreamResult{Chunks: chunks}, nil +} + +func (*retryConfigMutationExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) { + return auth, nil +} + +func (e *retryConfigMutationExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { + if e.recordCall() == 1 { + return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusInternalServerError, Message: "retry config changed"} + } + return cliproxyexecutor.Response{Payload: []byte("ok")}, nil +} + +func (*retryConfigMutationExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { + return nil, nil +} + +func (e *retryConfigMutationExecutor) recordCall() int { + e.mu.Lock() + defer e.mu.Unlock() + e.calls++ + if e.calls == 1 { + e.manager.SetRetryConfig(e.nextDefault, 0, 0) + } + return e.calls +} + +func (e *retryConfigMutationExecutor) callCount() int { + e.mu.Lock() + defer e.mu.Unlock() + return e.calls +} + +func TestExecuteSnapshotsDefaultRequestRetry(t *testing.T) { + paths := []struct { + name string + invoke func(*Manager) error + }{ + { + name: "non-stream", + invoke: func(manager *Manager) error { + _, errExecute := manager.Execute(context.Background(), []string{"retry-config-snapshot"}, cliproxyexecutor.Request{Model: "retry-config-snapshot-model"}, cliproxyexecutor.Options{}) + return errExecute + }, + }, + { + name: "count-tokens", + invoke: func(manager *Manager) error { + _, errExecute := manager.ExecuteCount(context.Background(), []string{"retry-config-snapshot"}, cliproxyexecutor.Request{Model: "retry-config-snapshot-model"}, cliproxyexecutor.Options{}) + return errExecute + }, + }, + { + name: "stream", + invoke: func(manager *Manager) error { + _, errExecute := manager.ExecuteStream(context.Background(), []string{"retry-config-snapshot"}, cliproxyexecutor.Request{Model: "retry-config-snapshot-model"}, cliproxyexecutor.Options{Stream: true}) + return errExecute + }, + }, + } + scenarios := []struct { + name string + initial int + next int + wantCalls int + wantSuccess bool + }{ + {name: "decrease after request starts", initial: 1, next: 0, wantCalls: 2, wantSuccess: true}, + {name: "increase after request starts", initial: 0, next: 1, wantCalls: 1, wantSuccess: false}, + } + + for _, path := range paths { + for _, scenario := range scenarios { + t.Run(path.name+"/"+scenario.name, func(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetRetryConfig(scenario.initial, 0, 0) + executor := &retryConfigMutationExecutor{ + identifier: "retry-config-snapshot", + manager: manager, + nextDefault: scenario.next, + } + manager.RegisterExecutor(executor) + + const authID = "retry-config-snapshot-auth" + reg := registry.GetGlobalRegistry() + reg.RegisterClient(authID, "retry-config-snapshot", []*registry.ModelInfo{{ID: "retry-config-snapshot-model"}}) + t.Cleanup(func() { reg.UnregisterClient(authID) }) + if _, errRegister := manager.Register(context.Background(), &Auth{ + ID: authID, + Provider: "retry-config-snapshot", + Metadata: map[string]any{"disable_cooling": true}, + }); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + errExecute := path.invoke(manager) + if scenario.wantSuccess && errExecute != nil { + t.Fatalf("execution error = %v, want success", errExecute) + } + if !scenario.wantSuccess && statusCodeFromError(errExecute) != http.StatusInternalServerError { + t.Fatalf("execution error = %v, want HTTP 500", errExecute) + } + if calls := executor.callCount(); calls != scenario.wantCalls { + t.Fatalf("executor calls = %d, want %d", calls, scenario.wantCalls) + } + }) + } + } +} + +type retryRoundHomeDispatcher struct { + mu sync.Mutex + limits map[string]int + rounds []int + newCalls int + oldCalls int +} + +func (*retryRoundHomeDispatcher) HeartbeatOK() bool { return true } + +func (d *retryRoundHomeDispatcher) RPopAuth(ctx context.Context, model, sessionID string, headers http.Header, count int) ([]byte, error) { + return d.RPopAuthWithRetryRoundConstraints(ctx, model, sessionID, headers, count, 0, nil, "") +} + +func (d *retryRoundHomeDispatcher) RPopAuthWithConstraints(ctx context.Context, model, sessionID string, headers http.Header, count int, excluded []string, pinned string) ([]byte, error) { + d.mu.Lock() + d.oldCalls++ + d.mu.Unlock() + return d.RPopAuthWithRetryRoundConstraints(ctx, model, sessionID, headers, count, 0, excluded, pinned) +} + +func (d *retryRoundHomeDispatcher) RPopAuthWithRetryRoundConstraints(_ context.Context, _ string, _ string, _ http.Header, _ int, retryRound int, excluded []string, pinned string) ([]byte, error) { + d.mu.Lock() + defer d.mu.Unlock() + d.newCalls++ + d.rounds = append(d.rounds, retryRound) + excludedSet := make(map[string]struct{}, len(excluded)) + for _, id := range excluded { + excludedSet[id] = struct{}{} + } + ids := make([]string, 0, len(d.limits)) + maxRetry := 0 + for id, limit := range d.limits { + if limit >= retryRound { + ids = append(ids, id) + if limit > maxRetry { + maxRetry = limit + } + } + } + sort.Strings(ids) + for _, id := range ids { + if _, okExcluded := excludedSet[id]; okExcluded || (pinned != "" && pinned != id) { + continue + } + return json.Marshal(homeAuthDispatchResponse{ + RequestRetry: func() *int { value := maxRetry; return &value }(), + Auth: Auth{ID: id, Provider: "retry-round-home", Status: StatusActive, Metadata: map[string]any{"request_retry": d.limits[id]}}, + }) + } + return nil, home.ErrAuthNotFound +} + +func (*retryRoundHomeDispatcher) AbortAmbiguousDispatch() {} + +func (d *retryRoundHomeDispatcher) roundsSeen() []int { + d.mu.Lock() + defer d.mu.Unlock() + return append([]int(nil), d.rounds...) +} + +func (d *retryRoundHomeDispatcher) dispatchMethodCalls() (int, int) { + d.mu.Lock() + defer d.mu.Unlock() + return d.newCalls, d.oldCalls +} + +func TestExecuteHomeRetryRoundCredentialWindows(t *testing.T) { + tests := []struct { + name string + invoke func(*Manager, cliproxyexecutor.Request) error + }{ + { + name: "non-stream", + invoke: func(manager *Manager, req cliproxyexecutor.Request) error { + _, errExecute := manager.Execute(context.Background(), []string{"retry-round-home"}, req, cliproxyexecutor.Options{}) + return errExecute + }, + }, + { + name: "count-tokens", + invoke: func(manager *Manager, req cliproxyexecutor.Request) error { + _, errExecute := manager.ExecuteCount(context.Background(), []string{"retry-round-home"}, req, cliproxyexecutor.Options{}) + return errExecute + }, + }, + { + name: "stream", + invoke: func(manager *Manager, req cliproxyexecutor.Request) error { + _, errExecute := manager.ExecuteStream(context.Background(), []string{"retry-round-home"}, req, cliproxyexecutor.Options{Stream: true}) + return errExecute + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + manager := NewManager(nil, nil, nil) + manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) + manager.SetRetryConfig(3, 0, 3) + dispatcher := &retryRoundHomeDispatcher{limits: map[string]int{ + "retry-round-a": 3, + "retry-round-b": 2, + "retry-round-c": 2, + }} + manager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1) + executor := &retryRoundCallExecutor{identifier: "retry-round-home"} + manager.RegisterExecutor(executor) + + if errExecute := test.invoke(manager, cliproxyexecutor.Request{Model: "retry-round-model"}); errExecute == nil { + t.Fatal("execution error = nil, want terminal retry error") + } + counts := countRetryRoundIDs(executor.ids(map[string]string{"non-stream": "execute", "count-tokens": "count", "stream": "stream"}[test.name])) + if counts["retry-round-a"] != 4 || counts["retry-round-b"] != 3 || counts["retry-round-c"] != 3 { + t.Fatalf("credential call counts = %#v, want A=4 B=3 C=3", counts) + } + rounds := dispatcher.roundsSeen() + if len(rounds) == 0 || rounds[0] != 0 { + t.Fatalf("Home retry rounds = %v, want initial round 0", rounds) + } + foundRoundOne := false + foundRoundTwo := false + foundRoundThree := false + for _, round := range rounds { + foundRoundOne = foundRoundOne || round == 1 + foundRoundTwo = foundRoundTwo || round == 2 + foundRoundThree = foundRoundThree || round == 3 + } + if !foundRoundOne || !foundRoundTwo || !foundRoundThree { + t.Fatalf("Home retry rounds = %v, want rounds 0,1,2,3", rounds) + } + newCalls, oldCalls := dispatcher.dispatchMethodCalls() + if newCalls == 0 || oldCalls != 0 { + t.Fatalf("Home dispatcher method calls = new %d, old %d; want new interface only", newCalls, oldCalls) + } + }) + } +} diff --git a/sdk/cliproxy/auth/conductor_selection.go b/sdk/cliproxy/auth/conductor_selection.go index 90682466..cc9dc751 100644 --- a/sdk/cliproxy/auth/conductor_selection.go +++ b/sdk/cliproxy/auth/conductor_selection.go @@ -709,6 +709,37 @@ func (m *Manager) retrySettings() (int, int, time.Duration) { return int(m.requestRetry.Load()), int(m.maxRetryCredentials.Load()), time.Duration(m.maxRetryInterval.Load()) } +func effectiveRequestRetryLimit(auth *Auth, defaultRetry int) int { + if defaultRetry < 0 { + defaultRetry = 0 + } + if override, ok := auth.RequestRetryOverride(); ok { + return override + } + return defaultRetry +} + +func (m *Manager) requestRetryRoundExclusions(retryRound int, defaultRequestRetry int) map[string]struct{} { + excluded := make(map[string]struct{}) + if m == nil || retryRound <= 0 { + return excluded + } + if defaultRequestRetry < 0 { + defaultRequestRetry = 0 + } + m.mu.RLock() + defer m.mu.RUnlock() + for _, auth := range m.auths { + if auth == nil || strings.TrimSpace(auth.ID) == "" { + continue + } + if effectiveRequestRetryLimit(auth, defaultRequestRetry) < retryRound { + excluded[auth.ID] = struct{}{} + } + } + return excluded +} + func retryRoundAvailabilityForAuth(auth *Auth, model string, now time.Time) (bool, time.Time) { blocked, reason, next := isAuthBlockedForModel(auth, model, now) if !blocked { @@ -757,14 +788,13 @@ func credentialRetryRoundStateEligible(lastErr *Error, quotaExceeded bool) bool return isCredentialRetryRoundStatus(statusCodeFromResult(lastErr)) } -func (m *Manager) closestCooldownWait(providers []string, model string, attempt int, eligibility authSelectionEligibility, pinnedAuthID string) (time.Duration, bool) { +func (m *Manager) closestCooldownWait(providers []string, model string, attempt int, eligibility authSelectionEligibility, pinnedAuthID string, defaultRequestRetry int) (time.Duration, bool) { if m == nil || len(providers) == 0 { return 0, false } now := time.Now() - defaultRetry := int(m.requestRetry.Load()) - if defaultRetry < 0 { - defaultRetry = 0 + if defaultRequestRetry < 0 { + defaultRequestRetry = 0 } providerSet := make(map[string]struct{}, len(providers)) for i := range providers { @@ -798,13 +828,7 @@ func (m *Manager) closestCooldownWait(providers []string, model string, attempt if model != "" && !m.authSupportsRouteModel(registryRef, auth, model) { continue } - effectiveRetry := defaultRetry - if override, ok := auth.RequestRetryOverride(); ok { - effectiveRetry = override - } - if effectiveRetry < 0 { - effectiveRetry = 0 - } + effectiveRetry := effectiveRequestRetryLimit(auth, defaultRequestRetry) if attempt >= effectiveRetry { continue } @@ -831,14 +855,13 @@ func (m *Manager) closestCooldownWait(providers []string, model string, attempt return minWait, found } -func (m *Manager) retryAllowed(attempt int, providers []string, model string, eligibility authSelectionEligibility, pinnedAuthID string) bool { +func (m *Manager) retryAllowed(attempt int, providers []string, model string, eligibility authSelectionEligibility, pinnedAuthID string, defaultRequestRetry int) bool { if m == nil || attempt < 0 || len(providers) == 0 { return false } now := time.Now() - defaultRetry := int(m.requestRetry.Load()) - if defaultRetry < 0 { - defaultRetry = 0 + if defaultRequestRetry < 0 { + defaultRequestRetry = 0 } providerSet := make(map[string]struct{}, len(providers)) for i := range providers { @@ -872,13 +895,7 @@ func (m *Manager) retryAllowed(attempt int, providers []string, model string, el if model != "" && !m.authSupportsRouteModel(registryRef, auth, model) { continue } - effectiveRetry := defaultRetry - if override, ok := auth.RequestRetryOverride(); ok { - effectiveRetry = override - } - if effectiveRetry < 0 { - effectiveRetry = 0 - } + effectiveRetry := effectiveRequestRetryLimit(auth, defaultRequestRetry) if attempt >= effectiveRetry { continue } @@ -894,7 +911,8 @@ func (m *Manager) retryAllowed(attempt int, providers []string, model string, el } 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) + defaultRequestRetry, _, _ := m.retrySettings() + return m.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), cliproxyexecutor.Options{}, err, attempt, providers, model, maxWait, -1, defaultRequestRetry) } // maxWait limits only positive cooldown waits between credential retry rounds. @@ -902,7 +920,7 @@ func (m *Manager) shouldRetryAfterError(err error, attempt int, providers []stri // 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) { +func (m *Manager) shouldRetryAfterErrorWithHomeRetryLimit(ctx context.Context, opts cliproxyexecutor.Options, err error, attempt int, providers []string, model string, maxWait time.Duration, homeRetryLimit int, defaultRequestRetry int) (time.Duration, bool) { if err == nil { return 0, false } @@ -953,10 +971,10 @@ func (m *Manager) shouldRetryAfterErrorWithHomeRetryLimit(ctx context.Context, o } eligibility := authSelectionEligibilityForRequest(ctx, opts) pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata) - if !isCredentialRetryRoundStatus(status) || !m.retryAllowed(attempt, providers, model, eligibility, pinnedAuthID) { + if !isCredentialRetryRoundStatus(status) || !m.retryAllowed(attempt, providers, model, eligibility, pinnedAuthID, defaultRequestRetry) { return 0, false } - wait, found := m.closestCooldownWait(providers, model, attempt, eligibility, pinnedAuthID) + wait, found := m.closestCooldownWait(providers, model, attempt, eligibility, pinnedAuthID, defaultRequestRetry) if found { if wait > 0 && (maxWait <= 0 || wait > maxWait) { return 0, false diff --git a/sdk/cliproxy/auth/home_retry_contract_test.go b/sdk/cliproxy/auth/home_retry_contract_test.go index c9b2d4df..5f62c652 100644 --- a/sdk/cliproxy/auth/home_retry_contract_test.go +++ b/sdk/cliproxy/auth/home_retry_contract_test.go @@ -563,23 +563,23 @@ func TestHomeRetryPolicyUsesRemoteCredentialOverrideBeforeSelection(t *testing.T hasRequestRetry: true, } - wait, shouldRetry := manager.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), cliproxyexecutor.Options{}, errRemoteCooldown, 0, []string{"home-retry-contract"}, "gpt", time.Second, -1) + wait, shouldRetry := manager.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), cliproxyexecutor.Options{}, errRemoteCooldown, 0, []string{"home-retry-contract"}, "gpt", time.Second, -1, 0) 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 { + if _, shouldRetry = manager.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), cliproxyexecutor.Options{}, errRemoteCooldown, 1, []string{"home-retry-contract"}, "gpt", time.Second, -1, 0); 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 { + if _, shouldRetry = manager.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), pinnedOpts, errRemoteCooldown, 0, []string{"home-retry-contract"}, "gpt", time.Second, -1, 0); 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 { + if _, shouldRetry = manager.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), cliproxyexecutor.Options{}, errRemoteCooldown, 0, []string{"home-retry-contract"}, "gpt", time.Second, -1, 0); shouldRetry { t.Fatal("explicit remote credential override 0 did not suppress the global retry setting") } retryLimit := 3 @@ -779,7 +779,7 @@ func TestHomeRetryRoundUsesAuthoritativeRemoteCooldown(t *testing.T) { { 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) + _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}, 2, retryLimit, 0, 0) return errExecute }, }, @@ -831,7 +831,7 @@ func TestHomeCooldownClassificationPreservesNonRetryableRoundStatus(t *testing.T { 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) + _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}, 2, retryLimit, 0, 0) return errExecute }, }, @@ -860,7 +860,7 @@ func TestHomeCooldownClassificationPreservesNonRetryableRoundStatus(t *testing.T 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 { + if wait, shouldRetry := manager.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), cliproxyexecutor.Options{}, errExecute, 0, []string{"home-retry-contract"}, "gpt", time.Second, retryLimit, 0); shouldRetry || wait != 0 { t.Fatalf("401 round retry = (%v, %t), want (0, false)", wait, shouldRetry) } }) @@ -882,7 +882,7 @@ func TestHomeRetryRoundStartsImmediatelyWhenHomeReportsAvailableNextRound(t *tes { 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) + _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}, 0, retryLimit, 0, 0) return errExecute }, }, @@ -911,7 +911,7 @@ func TestHomeRetryRoundStartsImmediatelyWhenHomeReportsAvailableNextRound(t *tes 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) + wait, shouldRetry := manager.shouldRetryAfterErrorWithHomeRetryLimit(context.Background(), cliproxyexecutor.Options{}, errExecute, 0, []string{"home-retry-contract"}, "gpt", 10*time.Second, retryLimit, 0) if !shouldRetry || wait != 0 { t.Fatalf("next-round retry = (%v, %t), want immediate", wait, shouldRetry) } @@ -934,7 +934,7 @@ func TestHomeRetryRoundUsesRemoteCooldownWhenAttemptedErrorHasNoTiming(t *testin { 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) + _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}, 2, retryLimit, 0, 0) return errExecute }, }, @@ -1274,7 +1274,7 @@ func TestHomeLocalSelectionRejectionWaitsForReleaseAcknowledgement(t *testing.T) 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) + _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}, maxRetryCredentials, retryLimit, 0, 0) return errExecute }, }, @@ -1289,7 +1289,7 @@ func TestHomeLocalSelectionRejectionWaitsForReleaseAcknowledgement(t *testing.T) 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) + _, errExecute := manager.executeStreamMixedOnce(context.Background(), []string{"home-retry-contract"}, cliproxyexecutor.Request{Model: "gpt"}, cliproxyexecutor.Options{Stream: true}, maxRetryCredentials, retryLimit, 0, 0) return errExecute }, },