From 26ba47c3256ad1454cbe4619a57d13a75dd5d762 Mon Sep 17 00:00:00 2001 From: Eric Litman Date: Tue, 27 Oct 2026 11:31:31 -0400 Subject: [PATCH] HOUSE-79 part 01: upstream v8.0.44 changes for review --- sdk/cliproxy/auth/antigravity_credits.go | 10 + .../auth/api_key_model_capabilities.go | 30 +- sdk/cliproxy/auth/auth_change_watch.go | 46 ++ sdk/cliproxy/auth/auto_refresh_loop.go | 32 +- sdk/cliproxy/auth/classification.go | 9 + sdk/cliproxy/auth/codex_persist.go | 135 +++++ sdk/cliproxy/auth/codex_siwc_bind.go | 246 ++++++++ sdk/cliproxy/auth/conductor.go | 48 +- sdk/cliproxy/auth/conductor_cooldown.go | 501 ++++++++++++++- sdk/cliproxy/auth/conductor_execution.go | 51 +- sdk/cliproxy/auth/conductor_home.go | 18 +- sdk/cliproxy/auth/conductor_home_execution.go | 7 +- sdk/cliproxy/auth/conductor_lifecycle.go | 382 ++++++++++-- sdk/cliproxy/auth/conductor_persistence.go | 373 ++++++++++++ sdk/cliproxy/auth/conductor_refresh.go | 424 ++++++++++++- sdk/cliproxy/auth/conductor_selection.go | 205 ++++++- sdk/cliproxy/auth/conductor_stream.go | 57 +- sdk/cliproxy/auth/cooldown_state.go | 7 + sdk/cliproxy/auth/errors.go | 8 + sdk/cliproxy/auth/home_concurrency.go | 6 + sdk/cliproxy/auth/home_result.go | 10 + sdk/cliproxy/auth/metadata_keys.go | 4 + sdk/cliproxy/auth/metadata_merge.go | 65 +- .../auth/native_tool_search_capability.go | 64 ++ sdk/cliproxy/auth/persist_policy.go | 12 + sdk/cliproxy/auth/scheduler.go | 6 + sdk/cliproxy/auth/selector.go | 264 +++++++- sdk/cliproxy/auth/session_cache.go | 81 ++- sdk/cliproxy/auth/store.go | 6 + sdk/cliproxy/auth/types.go | 569 +++++++++++++++++- 30 files changed, 3465 insertions(+), 211 deletions(-) create mode 100644 sdk/cliproxy/auth/auth_change_watch.go create mode 100644 sdk/cliproxy/auth/codex_persist.go create mode 100644 sdk/cliproxy/auth/codex_siwc_bind.go create mode 100644 sdk/cliproxy/auth/conductor_persistence.go create mode 100644 sdk/cliproxy/auth/native_tool_search_capability.go diff --git a/sdk/cliproxy/auth/antigravity_credits.go b/sdk/cliproxy/auth/antigravity_credits.go index 06e5d285e..29524e2a0 100644 --- a/sdk/cliproxy/auth/antigravity_credits.go +++ b/sdk/cliproxy/auth/antigravity_credits.go @@ -26,6 +26,10 @@ func AntigravityCreditsRequested(ctx context.Context) bool { return v } +// AntigravityCreditsHintRefreshInterval defines how long a cached AI credits hint +// remains fresh before it should be re-validated. +const AntigravityCreditsHintRefreshInterval = 10 * time.Minute + // AntigravityCreditsHint stores the latest known AI credits state for one auth. type AntigravityCreditsHint struct { Known bool @@ -36,6 +40,12 @@ type AntigravityCreditsHint struct { UpdatedAt time.Time } +// IsAntigravityCreditsHintFresh reports whether a known credits hint is still +// within AntigravityCreditsHintRefreshInterval of now. +func IsAntigravityCreditsHintFresh(hint AntigravityCreditsHint, now time.Time) bool { + return hint.Known && !hint.UpdatedAt.IsZero() && now.Sub(hint.UpdatedAt) < AntigravityCreditsHintRefreshInterval +} + var antigravityCreditsHintByAuth sync.Map // SetAntigravityCreditsHint updates the latest known AI credits state for an auth. diff --git a/sdk/cliproxy/auth/api_key_model_capabilities.go b/sdk/cliproxy/auth/api_key_model_capabilities.go index 8d56296d0..bb0ae4248 100644 --- a/sdk/cliproxy/auth/api_key_model_capabilities.go +++ b/sdk/cliproxy/auth/api_key_model_capabilities.go @@ -398,7 +398,19 @@ func compileAPIKeyModelCapabilitiesForAuth(cfg *internalconfig.Config, auth *Aut } case "vertex": if entry := resolveVertexAPIKeyConfig(cfg, auth); entry != nil { - compileConfiguredModelCapabilities(out, entry.Models, "gemini") + for _, configured := range entry.Models { + support := configured.Thinking + openAPI := isVertexOpenAPIConfigModel(configured.Name) + if configured.OpenAPI != nil { + openAPI = *configured.OpenAPI + } else if entry.OpenAPI != nil { + openAPI = *entry.OpenAPI + } + if support == nil && openAPI && modelconfig.ResolveModelInfo(configured.Name, "gemini", nil).Thinking == nil { + support = ®istry.ThinkingSupport{Levels: []string{"low", "medium", "high"}} + } + addConfiguredModelCapability(out, configured.Name, configured.Alias, "gemini", support, false) + } } case "meta": if entry := resolveMetaAPIKeyConfig(cfg, auth); entry != nil { @@ -485,3 +497,19 @@ func addConfiguredModelCapability(out map[string][]apiKeyModelCapabilityRoute, n } return modelInfo } + +// isVertexOpenAPIConfigModel recognizes partner models without depending on executors. +func isVertexOpenAPIConfigModel(model string) bool { + base := strings.TrimSpace(thinking.ParseSuffix(model).ModelName) + if strings.HasPrefix(strings.ToLower(base), "publishers/") { + publisher, tail, ok := strings.Cut(base[len("publishers/"):], "/") + if !ok || !strings.HasPrefix(strings.ToLower(tail), "models/") { + return false + } + base = strings.TrimSpace(publisher) + "/" + strings.TrimSpace(tail[len("models/"):]) + } + publisher, name, ok := strings.Cut(base, "/") + publisher, name = strings.TrimSpace(publisher), strings.TrimSpace(name) + return ok && publisher != "" && name != "" && !strings.Contains(name, "/") && + !strings.EqualFold(publisher, "google") && !strings.EqualFold(publisher, "models") && !strings.EqualFold(publisher, "publishers") +} diff --git a/sdk/cliproxy/auth/auth_change_watch.go b/sdk/cliproxy/auth/auth_change_watch.go new file mode 100644 index 000000000..5659fa0c9 --- /dev/null +++ b/sdk/cliproxy/auth/auth_change_watch.go @@ -0,0 +1,46 @@ +package auth + +// WatchAuthChanges observes structural changes to one runtime credential, +// including token preparation/refresh, replacement, removal and store reload. +// Notifications are coalesced: consumers must inspect the latest manager state. +// The caller must unsubscribe when its work ends. Request result accounting does +// not notify watchers, so ordinary Generation increments do not invalidate work. +func (m *Manager) WatchAuthChanges(authID string) (<-chan struct{}, func()) { + changes := make(chan struct{}, 1) + if m == nil || authID == "" { + close(changes) + return changes, func() {} + } + m.mu.Lock() + if m.authChangeWatchers == nil { + m.authChangeWatchers = make(map[string]map[chan struct{}]struct{}) + } + if m.authChangeWatchers[authID] == nil { + m.authChangeWatchers[authID] = make(map[chan struct{}]struct{}) + } + m.authChangeWatchers[authID][changes] = struct{}{} + m.mu.Unlock() + return changes, func() { + m.mu.Lock() + defer m.mu.Unlock() + watchers := m.authChangeWatchers[authID] + if _, exists := watchers[changes]; !exists { + return + } + delete(watchers, changes) + if len(watchers) == 0 { + delete(m.authChangeWatchers, authID) + } + close(changes) + } +} + +// Caller holds m.mu. No user callback or blocking work runs under the manager lock. +func (m *Manager) notifyAuthChangeLocked(authID string) { + for changes := range m.authChangeWatchers[authID] { + select { + case changes <- struct{}{}: + default: + } + } +} diff --git a/sdk/cliproxy/auth/auto_refresh_loop.go b/sdk/cliproxy/auth/auto_refresh_loop.go index db33c5458..ab54ba483 100644 --- a/sdk/cliproxy/auth/auto_refresh_loop.go +++ b/sdk/cliproxy/auth/auto_refresh_loop.go @@ -400,14 +400,42 @@ func nextRefreshCheckAt(now time.Time, auth *Auth, interval time.Duration) (time return time.Time{}, false } - if auth.AuthKind() == AuthKindAPIKey { - return time.Time{}, false + // Quota-exhausted Devin credentials revalidate upstream status so early + // restoration can be detected (issue #6690); their tokens never expire, + // so the generic expiry/lead schedule below would drop them entirely. + // A custom RefreshEvaluator keeps precedence (matching shouldRefresh), + // otherwise an evaluator veto would reschedule an immediately-due entry. + if _, hasEvaluator := auth.Runtime.(RefreshEvaluator); hasEvaluator { + // Fall through to the evaluator branch below. + } else if next, due := devinQuotaRevalidateAt(auth, now); due { + if !auth.NextRefreshAfter.IsZero() && auth.NextRefreshAfter.After(next) { + next = auth.NextRefreshAfter + } + if !next.After(now) { + return now, true + } + return next, true } if !auth.NextRefreshAfter.IsZero() && now.Before(auth.NextRefreshAfter) { return auth.NextRefreshAfter, true } + if siwcDueAt, hasSIWC, _ := CodexSIWCNeedsRefresh(auth, now); hasSIWC { + minDue := siwcDueAt + if outerDueAt, hasOuter, _ := CodexOuterNeedsRefresh(auth, now); hasOuter && outerDueAt.Before(minDue) { + minDue = outerDueAt + } + if !minDue.After(now) { + return now, true + } + return minDue, true + } + + if auth.AuthKind() == AuthKindAPIKey { + return time.Time{}, false + } + if evaluator, ok := auth.Runtime.(RefreshEvaluator); ok && evaluator != nil { if interval <= 0 { interval = refreshCheckInterval diff --git a/sdk/cliproxy/auth/classification.go b/sdk/cliproxy/auth/classification.go index b3eb4d85b..ced4d0bf8 100644 --- a/sdk/cliproxy/auth/classification.go +++ b/sdk/cliproxy/auth/classification.go @@ -31,6 +31,9 @@ func (a *Auth) AuthKind() string { if a == nil { return "" } + if IsCodexUseSIWC(a) { + return AuthKindOAuth + } if kind := normalizeAuthKind(authAttribute(a, AttributeAuthKind)); kind != "" { return kind } @@ -119,6 +122,12 @@ func authHasOAuthMetadata(auth *Auth) bool { if token, ok := auth.Metadata["token"].(map[string]any); ok && len(token) > 0 { return true } + if siwc, ok := auth.Metadata["siwc"].(map[string]any); ok && len(siwc) > 0 { + return true + } + if sa, ok := auth.Metadata["service_account"].(map[string]any); ok && len(sa) > 0 { + return true + } return false } diff --git a/sdk/cliproxy/auth/codex_persist.go b/sdk/cliproxy/auth/codex_persist.go new file mode 100644 index 000000000..0215f7f6d --- /dev/null +++ b/sdk/cliproxy/auth/codex_persist.go @@ -0,0 +1,135 @@ +package auth + +import ( + "context" + "encoding/json" + "fmt" + "os" + "strings" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" +) + +// IsCodexAuthRecord reports whether auth represents a Codex credential record. +func IsCodexAuthRecord(auth *Auth) bool { + if auth == nil { + return false + } + if strings.EqualFold(strings.TrimSpace(auth.Provider), "codex") { + return true + } + if auth.Metadata != nil { + if t, ok := auth.Metadata["type"].(string); ok && strings.EqualFold(strings.TrimSpace(t), "codex") { + return true + } + } + if _, ok := auth.Storage.(*codex.CodexTokenStorage); ok { + return true + } + return false +} + +// PersistCodexAuthFile persists a Codex auth record to path using inter-process file locking, +// read-before-write merging, and scoped refresh patches when present. +// It returns handled=false when auth is not a Codex credential record. +func PersistCodexAuthFile(ctx context.Context, path string, auth *Auth) (bool, error) { + if !IsCodexAuthRecord(auth) { + return false, nil + } + + type metadataSetter interface { + SetMetadata(map[string]any) + } + + switch { + case auth.CodexRefreshPatch != nil: + errRefresh := codex.WithFileLock(path, func() error { + existing, errRead := os.ReadFile(path) + if errRead != nil { + if os.IsNotExist(errRead) { + if auth.Metadata == nil { + auth.Metadata = make(map[string]any) + } + auth.Metadata["disabled"] = auth.Disabled + NormalizeCredentialMetadata(auth.Metadata) + return codex.WriteJSONFileAtomic(path, auth.Metadata) + } + return fmt.Errorf("auth filestore: read refresh source failed: %w", errRead) + } + var diskMap map[string]any + if errUnmarshal := json.Unmarshal(existing, &diskMap); errUnmarshal != nil || diskMap == nil { + return fmt.Errorf("auth filestore: invalid refresh source JSON") + } + merged := auth.CodexRefreshPatch.Apply(diskMap) + NormalizeCredentialMetadata(merged) + if errWrite := codex.WriteJSONFileAtomic(path, merged); errWrite != nil { + return errWrite + } + auth.Metadata = merged + return nil + }) + if errRefresh != nil { + return true, errRefresh + } + auth.CodexRefreshPatch = nil + auth.Storage = nil + if disabled, ok := auth.Metadata["disabled"].(bool); ok { + auth.Disabled = disabled + } + return true, nil + + case auth.Storage != nil: + if auth.Metadata == nil { + auth.Metadata = make(map[string]any) + } + auth.Metadata["disabled"] = auth.Disabled + if storage, ok := auth.Storage.(*codex.CodexTokenStorage); ok && storage != nil { + clonedStorage := *storage + clonedStorage.ReplaceCredentials = HasAuthCreationIntent(ctx) + auth.Storage = &clonedStorage + } + if setter, ok := auth.Storage.(metadataSetter); ok { + setter.SetMetadata(auth.Metadata) + } + if errSave := auth.Storage.SaveTokenToFile(path); errSave != nil { + return true, errSave + } + if persistedBytes, errReadPersisted := os.ReadFile(path); errReadPersisted == nil && len(persistedBytes) > 0 { + var diskMap map[string]any + if errUnmarshal := json.Unmarshal(persistedBytes, &diskMap); errUnmarshal == nil && len(diskMap) > 0 { + auth.Metadata = codex.MergeCodexCredentialsFromDisk(diskMap, auth.Metadata) + NormalizeCredentialMetadata(auth.Metadata) + } + } + return true, nil + + case auth.Metadata != nil: + auth.Metadata["disabled"] = auth.Disabled + errLock := codex.WithFileLock(path, func() error { + existing, errRead := os.ReadFile(path) + if errRead == nil && len(existing) > 0 { + var diskMap map[string]any + if errUnmarshal := json.Unmarshal(existing, &diskMap); errUnmarshal == nil && len(diskMap) > 0 { + if !HasAuthCreationIntent(ctx) { + auth.Metadata = codex.MergeCodexCredentialsFromDisk(diskMap, auth.Metadata) + } + auth.Metadata["disabled"] = auth.Disabled + NormalizeCredentialMetadata(auth.Metadata) + if raw, errMarshal := json.Marshal(auth.Metadata); errMarshal == nil && string(existing) == string(raw) { + return nil + } + } + } else if errRead != nil && !os.IsNotExist(errRead) { + return fmt.Errorf("auth filestore: read existing failed: %w", errRead) + } + return codex.WriteJSONFileAtomic(path, auth.Metadata) + }) + if errLock != nil { + return true, errLock + } + return true, nil + + default: + return false, nil + } +} diff --git a/sdk/cliproxy/auth/codex_siwc_bind.go b/sdk/cliproxy/auth/codex_siwc_bind.go new file mode 100644 index 000000000..a2240c8c9 --- /dev/null +++ b/sdk/cliproxy/auth/codex_siwc_bind.go @@ -0,0 +1,246 @@ +package auth + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "strings" + "time" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/internal/util" +) + +// BindCodexSIWC explicitly replaces the SIWC partition of an existing file-backed +// Codex registration. Unlike ordinary persistence, a failed write never publishes +// credentials. expectedEpoch and expectedAccount must match the registration captured before OAuth. +func (m *Manager) BindCodexSIWC(ctx context.Context, id string, expectedEpoch uint64, expectedAccount string, tokenData *codex.CodexSIWCTokenData) (*Auth, error) { + if m == nil || strings.TrimSpace(id) == "" || expectedEpoch == 0 || tokenData == nil || strings.TrimSpace(tokenData.AccessToken) == "" || strings.TrimSpace(tokenData.RefreshToken) == "" || strings.TrimSpace(tokenData.ClientID) == "" { + return nil, fmt.Errorf("bind SIWC: invalid target or token data") + } + if ctx == nil { + ctx = context.Background() + } + release, err := m.lockAuthMutationContext(ctx, id) + if err != nil { + return nil, err + } + defer release() + m.mu.RLock() + base := m.auths[id].Clone() + m.mu.RUnlock() + if base == nil || base.RegistrationEpoch != expectedEpoch || !strings.EqualFold(strings.TrimSpace(base.Provider), "codex") || !codexBindingHasOuter(base.Metadata) { + return nil, fmt.Errorf("bind SIWC: target is missing, changed, or not a Codex OAuth credential") + } + path, err := m.codexBindingPath(base) + if err != nil { + return nil, err + } + var committed *Auth + cooldownChanged := false + err = codex.WithFileLock(path, func() error { + if errCtx := ctx.Err(); errCtx != nil { + return errCtx + } + checkedPath, errPath := m.codexBindingPath(base) + if errPath != nil || checkedPath != path { + return fmt.Errorf("bind SIWC: backing file changed") + } + info, errStat := os.Lstat(path) + if errStat != nil || !info.Mode().IsRegular() { + return fmt.Errorf("bind SIWC: backing file is not a regular file") + } + data, errRead := os.ReadFile(path) + if errRead != nil { + return fmt.Errorf("bind SIWC: read backing file: %w", errRead) + } + var disk map[string]any + if json.Unmarshal(data, &disk) != nil || !codexBindingHasOuter(disk) { + return fmt.Errorf("bind SIWC: invalid Codex OAuth file") + } + kind, _ := disk["type"].(string) + if !strings.EqualFold(strings.TrimSpace(kind), "codex") { + return fmt.Errorf("bind SIWC: backing file type changed") + } + if account, _ := base.Metadata["account_id"].(string); account != expectedAccount || (account != "" && disk["account_id"] != account) || (account == "" && disk["account_id"] != nil && disk["account_id"] != "") { + return fmt.Errorf("bind SIWC: backing account changed") + } + // Keep arbitrary disk values (including large JSON numbers) losslessly. Only + // the two explicitly bound fields are replaced; no metadata normalization runs. + var raw map[string]json.RawMessage + if errJSON := json.Unmarshal(data, &raw); errJSON != nil { + return errJSON + } + replacement := tokenData.ToMap() + siwc, errJSON := json.Marshal(replacement) + if errJSON != nil { + return errJSON + } + out := make(map[string]any, len(raw)+2) + for key, value := range raw { + out[key] = value + } + out["siwc"] = json.RawMessage(siwc) + out["use_siwc"] = true + if errCtx := ctx.Err(); errCtx != nil { + return errCtx + } + latestInfo, errStat := os.Lstat(path) + if errStat != nil || !os.SameFile(info, latestInfo) { + return fmt.Errorf("bind SIWC: backing file changed during binding") + } + latestData, errReadLatest := os.ReadFile(path) + if errReadLatest != nil || !bytes.Equal(data, latestData) { + return fmt.Errorf("bind SIWC: backing file changed during binding") + } + if errWrite := codex.WriteJSONFileAtomic(path, out); errWrite != nil { + return fmt.Errorf("bind SIWC: persist binding: %w", errWrite) + } + disk["siwc"] = replacement + disk["use_siwc"] = true + // The mutation gate fences Delete, reload, and refresh publication. Clone + // only now so counters and request-time runtime changes during I/O survive. + m.mu.Lock() + current := m.auths[id] + if current == nil || current.RegistrationEpoch != expectedEpoch { + m.mu.Unlock() + return fmt.Errorf("bind SIWC: registration changed before publication") + } + updated := current.Clone() + updated.Metadata = disk + updated.Storage = nil + updated.CodexRefreshPatch = nil + // The file owns this toggle; stale attribute overrides must not mask it + // or prevent a later file edit from disabling SIWC. + delete(updated.Attributes, "use_siwc") + if disabled, ok := disk["disabled"].(bool); ok { + updated.Disabled = disabled + } + if updated.Disabled { + updated.Status = StatusDisabled + } else if updated.Status == StatusDisabled { + updated.Status = StatusActive + } + now := time.Now() + if !updated.Disabled { + applyPreservedRejectedTokens(current, updated, false) + if updated.RejectedAccessToken == "" || authAccessToken(updated) != updated.RejectedAccessToken { + updated.RefreshFailures = 0 + updated.NextRefreshAfter = time.Time{} + if hasUnauthorizedAuthFailure(current) || + (updated.LastError != nil && (isUnauthorizedError(updated.LastError) || isInvalidGrantResultError(updated.LastError))) || + (updated.Unavailable && !updated.NextRetryAfter.After(now) && !updated.Quota.NextRecoverAt.After(now)) { + updated.Unavailable = false + updated.LastError = nil + updated.StatusMessage = "" + updated.Status = StatusActive + updated.NextRetryAfter = time.Time{} + cooldownChanged = true + } + clearUnauthorizedModelStates(updated, now) + } + // The helper can clear errors while preserving a model quota block. + if len(updated.ModelStates) > 0 { + cooldownChanged = true + } + if current.Quota.Exceeded && (current.Quota.NextRecoverAt.IsZero() || current.Quota.NextRecoverAt.After(now)) { + updated.Quota = current.Quota.Clone() + updated.Unavailable = current.Unavailable + updated.NextRetryAfter = current.NextRetryAfter + updated.Status = current.Status + } + } + updated.Generation++ + updated.CredentialVersion = max(current.CredentialVersion, 1) + 1 + updated.UpdatedAt = now + updated.ClearFileSnapshot() + m.auths[id] = updated + m.notifyAuthChangeLocked(id) + committed = updated.Clone() + m.mu.Unlock() + return nil + }) + if err != nil { + return nil, err + } + release() + m.finishAuthUpdate(ctx, committed, cooldownChanged) + return committed, nil +} + +func codexBindingHasOuter(metadata map[string]any) bool { + access, _ := metadata["access_token"].(string) + refresh, _ := metadata["refresh_token"].(string) + return strings.TrimSpace(access) != "" && strings.TrimSpace(refresh) != "" +} + +// Resolve confinement again at commit, not just when the management flow starts. +func (m *Manager) codexBindingPath(auth *Auth) (string, error) { + source := strings.ToLower(authAttribute(auth, AttributeSource)) + if auth == nil || auth.AuthSourceKind() != AuthSourceFile || source == AuthSourceConfig || strings.HasPrefix(source, AuthSourceConfig+":") || auth.AuthKind() != AuthKindOAuth { + return "", fmt.Errorf("bind SIWC: target is not a file-backed OAuth credential") + } + cfg, _ := m.runtimeConfig.Load().(*config.Config) + if cfg == nil || strings.TrimSpace(cfg.AuthDir) == "" { + return "", fmt.Errorf("bind SIWC: auth directory is not configured") + } + authDir, err := util.ResolveAuthDir(cfg.AuthDir) + if err != nil { + return "", fmt.Errorf("bind SIWC: resolve auth directory: %w", err) + } + authDir, err = filepath.Abs(authDir) + if err != nil { + return "", err + } + root, err := filepath.EvalSymlinks(authDir) + if err != nil { + return "", fmt.Errorf("bind SIWC: invalid auth directory") + } + root, err = filepath.Abs(root) + if err != nil { + return "", err + } + path := strings.TrimSpace(auth.Attributes["path"]) + if path == "" { + path = strings.TrimSpace(auth.Attributes["source"]) + } + if path != "" && !filepath.IsAbs(path) { + path, err = filepath.Abs(path) + if err != nil { + return "", err + } + } + if path == "" && strings.TrimSpace(auth.FileName) != "" { + fileName := strings.TrimSpace(auth.FileName) + if filepath.IsAbs(fileName) { + path = fileName + } else { + path = filepath.Join(root, fileName) + } + } + if path == "" { + return "", fmt.Errorf("bind SIWC: target is not file-backed") + } + resolved, err := filepath.EvalSymlinks(path) + if err != nil { + return "", fmt.Errorf("bind SIWC: backing file is missing or inaccessible") + } + resolved, err = filepath.Abs(resolved) + if err != nil { + return "", err + } + relative, err := filepath.Rel(root, resolved) + if err != nil || relative == "." || relative == ".." || strings.HasPrefix(relative, ".."+string(filepath.Separator)) { + return "", fmt.Errorf("bind SIWC: backing file is outside auth directory") + } + // Reject final-component symlinks rather than replacing their targets. + info, err := os.Lstat(path) + if err != nil || !info.Mode().IsRegular() { + return "", fmt.Errorf("bind SIWC: backing file is not regular") + } + return resolved, nil +} diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index 9c12c7101..3baead92f 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -2,6 +2,7 @@ package auth import ( "context" + "math" "net/http" "sync" "sync/atomic" @@ -10,6 +11,7 @@ import ( internalconfig "github.com/router-for-me/CLIProxyAPI/v8/internal/config" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" "github.com/router-for-me/CLIProxyAPI/v8/sdk/pluginapi" + "golang.org/x/sync/semaphore" ) // ProviderExecutor defines the contract required by Manager to execute provider calls. @@ -48,6 +50,18 @@ type ExecutionSessionCloser interface { CloseExecutionSession(sessionID string) } +// ExecutionSessionSharer reports whether a replaced executor shares its execution +// session store with the replacement executor so active sessions stay open. +type ExecutionSessionSharer interface { + SharesExecutionSessionsWith(other ProviderExecutor) bool +} + +// AuthExecutionSessionCloser allows executors to release execution sessions +// associated with a removed credential without tearing down other credentials' sessions. +type AuthExecutionSessionCloser interface { + CloseAuthExecutionSessions(authID string) +} + // Result captures execution outcome used to adjust auth state. type Result struct { // AuthID references the auth that produced this result. @@ -72,6 +86,18 @@ type Result struct { // observed watermark. Count-tokens requests reuse the credential but are not // generation traffic; their response headers are not a generation snapshot. SkipQuotaObservation bool + // AccessToken captures the active access token under which the request was executed. + AccessToken string + // CredentialVersion captures the credential version under which the request was executed. + CredentialVersion uint64 + // RegistrationEpoch captures the registration epoch under which the request was executed. + RegistrationEpoch uint64 + // ExecutionSeq is the manager-local monotonic sequence assigned when the + // execution was dispatched upstream. It orders in-flight executions so a + // success dispatched before a cooldown was recorded cannot clear that + // cooldown. Zero means unknown (direct MarkResult callers) and is exempt + // from the ordering guard. + ExecutionSeq uint64 } // Selector chooses an auth candidate for execution. @@ -154,11 +180,15 @@ type Manager struct { selectorMu sync.Mutex configCooldownMu sync.Mutex syncSchedulerMu sync.Mutex - structuralEpoch atomic.Uint64 - syncedVersion atomic.Uint64 - auths map[string]*Auth - authEpochs map[string]uint64 - scheduler *authScheduler + // executionSeq stamps upstream dispatches with a monotonic sequence used + // to reject cooldown clears from executions dispatched before the cooldown. + executionSeq atomic.Uint64 + structuralEpoch atomic.Uint64 + syncedVersion atomic.Uint64 + auths map[string]*Auth + authEpochs map[string]uint64 + authChangeWatchers map[string]map[chan struct{}]struct{} + scheduler *authScheduler // pluginScheduler runs outside m.mu before falling back to native selection. pluginScheduler PluginScheduler // homeRuntimeAuths retains legacy session auth lookups for non-execution callers. @@ -207,6 +237,13 @@ type Manager struct { refreshLocks sync.Map // persistLocks serializes disk persistence per auth ID and guards against out-of-order writes. persistLocks sync.Map + // authMutationLocks coordinate credential mutations without blocking unrelated readers. + authMutationLocks sync.Map + metadataPersistCancels sync.Map + metadataSaveGate sync.Map + // authLoadGate allows concurrent credential transactions, but excludes whole-store reloads. + // Mutations acquire one permit; Load acquires all permits before reading the store. + authLoadGate *semaphore.Weighted } // NewManager constructs a manager with optional custom selector and hook. @@ -219,6 +256,7 @@ func NewManager(store Store, selector Selector, hook Hook) *Manager { } manager := &Manager{ store: store, + authLoadGate: semaphore.NewWeighted(math.MaxInt64), executors: make(map[string]ProviderExecutor), selector: selector, hook: hook, diff --git a/sdk/cliproxy/auth/conductor_cooldown.go b/sdk/cliproxy/auth/conductor_cooldown.go index 4aa8a1aa0..b4d778f59 100644 --- a/sdk/cliproxy/auth/conductor_cooldown.go +++ b/sdk/cliproxy/auth/conductor_cooldown.go @@ -395,6 +395,20 @@ func (m *Manager) restoreCooldownRecordLocked(record CooldownStateRecord, now ti } state := ensureModelState(auth, model) + // Candidate independent layers for a wrapped record: the live shadow, or + // — when the live state is still unwrapped — the live state itself (a + // delayed store load must never erase a newer independent cooldown that + // was recorded after the snapshot was taken). A legacy pre-fix reporting + // state (reason "quota" whose own error is the explicit daily/weekly + // exhaustion message) is credential-derived and must not be captured; + // genuine model-scoped 429s carry different messages and stay eligible. + liveShadow := state.credentialQuotaShadow.cloneShallow() + if liveShadow == nil && !state.credentialQuotaPropagated && state.Quota.Reason != "credential_quota" && + isModelStateActiveCooldown(state, now) && !isDevinQuotaExhaustionMessage(devinStateErrorText(state)) { + liveShadow = state.cloneShallow() + } + wasLiveWrapped := state.credentialQuotaPropagated || state.Quota.Reason == "credential_quota" || + (auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" && auth.Quota.NextRecoverAt.After(now)) mergeModelState(state, &ModelState{ Unavailable: true, Status: StatusError, @@ -404,6 +418,55 @@ func (m *Manager) restoreCooldownRecordLocked(record CooldownStateRecord, now ti LastError: cloneError(record.LastError), UpdatedAt: updatedAt, }) + // The persisted independent layer joins the merge newest-wins; a newer + // live independent layer (shadow or freshly captured unwrapped state, + // e.g. an own failure recorded during a delayed store load) is never + // clobbered by the older record. A wrapped record is any record whose + // outer reason is the propagated credential quota. + isWrappedRecord := record.IndependentState != nil || record.CredentialQuotaPropagated || strings.TrimSpace(quota.Reason) == "credential_quota" + if !isWrappedRecord && wasLiveWrapped { + // An unwrapped independent record landing beneath a live wrap becomes + // the model's independent layer (unless it is itself a legacy + // credential-derived exhaustion state); the merged outer fields keep + // presenting the propagated credential quota. + if !isDevinQuotaExhaustionMessage(devinRecordErrorText(record)) { + recordShadow := &ModelState{ + Unavailable: true, + Status: StatusError, + StatusMessage: reason, + NextRetryAfter: record.NextRetryAfter, + Quota: quota, + LastError: cloneError(record.LastError), + UpdatedAt: updatedAt, + } + preferred := state.credentialQuotaShadow + if preferred == nil || (liveShadow != nil && !liveShadow.UpdatedAt.After(updatedAt)) { + preferred = recordShadow + } else if liveShadow != nil && liveShadow.UpdatedAt.After(updatedAt) { + preferred = liveShadow + } + state.credentialQuotaShadow = preferred + } + state.credentialQuotaPropagated = true + applyCooldownFields(&state.Quota, QuotaState{ + Exceeded: true, + Reason: "credential_quota", + NextRecoverAt: auth.Quota.NextRecoverAt, + BackoffLevel: auth.Quota.BackoffLevel, + }) + if state.NextRetryAfter.Before(auth.Quota.NextRecoverAt) { + state.NextRetryAfter = auth.Quota.NextRecoverAt + } + state.Unavailable = true + state.Status = StatusError + } else if isWrappedRecord { + state.credentialQuotaPropagated = true + if record.IndependentState != nil && (liveShadow == nil || !liveShadow.UpdatedAt.After(record.IndependentState.UpdatedAt)) { + state.credentialQuotaShadow = record.IndependentState.cloneShallow() + } else if liveShadow != nil { + state.credentialQuotaShadow = liveShadow + } + } auth.Generation++ auth.UpdatedAt = updatedAt updateAggregatedAvailability(auth, now) @@ -429,10 +492,12 @@ func clearCooldownStateForAuth(auth *Auth, now time.Time) bool { if state == nil { continue } - if state.Unavailable || !state.NextRetryAfter.IsZero() || state.Quota.Exceeded || !state.Quota.NextRecoverAt.IsZero() { + if state.Unavailable || !state.NextRetryAfter.IsZero() || state.Quota.Exceeded || !state.Quota.NextRecoverAt.IsZero() || state.credentialQuotaShadow != nil || state.credentialQuotaPropagated { state.Unavailable = false state.NextRetryAfter = time.Time{} applyCooldownFields(&state.Quota, QuotaState{}) + state.credentialQuotaShadow = nil + state.credentialQuotaPropagated = false state.UpdatedAt = now changed = true } @@ -483,6 +548,8 @@ func (m *Manager) ResetQuota(ctx context.Context, authID string) (*Auth, []strin registeredModels := modelsForRegisteredAuth(authID) cooldownStateChanged := false + releaseMutation := m.lockAuthMutation(authID) + defer releaseMutation() m.mu.Lock() auth, ok := m.auths[authID] if !ok || auth == nil { @@ -532,8 +599,9 @@ func (m *Manager) ResetQuota(ctx context.Context, authID string) (*Auth, []strin cooldownRecordsAfter := m.cooldownStateRecordsForAuthLocked(auth, now) cooldownStateChanged = !cooldownStateRecordsEqual(cooldownRecordsBefore, cooldownRecordsAfter) } - errPersist := m.persist(ctx, auth) + errPersist := m.persistLocked(ctx, auth) m.mu.Unlock() + releaseMutation() defer func() { if cooldownStateChanged { @@ -738,6 +806,9 @@ func cooldownStateRecordsEqual(a, b []CooldownStateRecord) bool { } func cooldownStateRecordEqual(a, b CooldownStateRecord) bool { + if a.CredentialQuotaPropagated != b.CredentialQuotaPropagated { + return false + } if a.Provider != b.Provider || a.AuthID != b.AuthID || a.AuthFile != b.AuthFile || @@ -749,7 +820,23 @@ func cooldownStateRecordEqual(a, b CooldownStateRecord) bool { !cooldownQuotaEqual(a.Quota, b.Quota) { return false } - return cooldownErrorEqual(a.LastError, b.LastError) + if !cooldownErrorEqual(a.LastError, b.LastError) { + return false + } + return cooldownIndependentStateEqual(a.IndependentState, b.IndependentState) +} + +func cooldownIndependentStateEqual(a, b *ModelState) bool { + if a == nil || b == nil { + return a == b + } + return a.Status == b.Status && + a.StatusMessage == b.StatusMessage && + a.Unavailable == b.Unavailable && + a.NextRetryAfter.Equal(b.NextRetryAfter) && + a.UpdatedAt.Equal(b.UpdatedAt) && + cooldownQuotaEqual(a.Quota, b.Quota) && + cooldownErrorEqual(a.LastError, b.LastError) } func cooldownQuotaEqual(a, b QuotaState) bool { @@ -791,7 +878,7 @@ func modelCooldownStateRecord(auth *Auth, model string, state *ModelState, now t if auth == nil || state == nil || model == "" || !state.Unavailable || state.NextRetryAfter.IsZero() || !state.NextRetryAfter.After(now) { return CooldownStateRecord{}, false } - return CooldownStateRecord{ + record := CooldownStateRecord{ Provider: strings.TrimSpace(auth.Provider), AuthID: auth.ID, AuthFile: cooldownAuthFile(auth), @@ -802,7 +889,12 @@ func modelCooldownStateRecord(auth *Auth, model string, state *ModelState, now t Quota: cooldownFieldsOf(state.Quota), LastError: cloneError(state.LastError), UpdatedAt: state.UpdatedAt, - }, true + } + if state.credentialQuotaShadow != nil { + record.IndependentState = state.credentialQuotaShadow.cloneShallow() + } + record.CredentialQuotaPropagated = state.credentialQuotaPropagated + return record, true } func cooldownReason(statusMessage string, quota QuotaState, lastErr *Error) string { @@ -823,6 +915,43 @@ func cooldownReason(statusMessage string, quota QuotaState, lastErr *Error) stri return "" } +// nextExecutionSeq stamps a new upstream dispatch with the next monotonic +// execution sequence. The stamp travels on the execution's Result so that +// MarkResult can order the execution against cooldowns recorded meanwhile. +func (m *Manager) nextExecutionSeq() uint64 { + if m == nil { + return 0 + } + return m.executionSeq.Add(1) +} + +// currentExecutionSeq returns the latest execution sequence issued so far. +// Capturing it as the cooldown watermark when a cooldown is recorded marks +// every execution dispatched up to that moment as predating the cooldown. +func (m *Manager) currentExecutionSeq() uint64 { + if m == nil { + return 0 + } + return m.executionSeq.Load() +} + +// isStaleExecutionResult reports whether an execution result originated from a superseded +// credential version or an earlier registration epoch than the currently active auth. +func isStaleExecutionResult(result Result, current *Auth) bool { + if current == nil { + return false + } + staleEpoch := result.RegistrationEpoch > 0 && result.RegistrationEpoch < current.RegistrationEpoch + if staleEpoch { + return true + } + if result.AccessToken != "" && authMatchesTokenPartition(current, result.AccessToken) { + return false + } + staleVersion := result.CredentialVersion < current.CredentialVersion && (result.CredentialVersion > 0 || current.CredentialVersion > 1) + return staleVersion +} + // MarkResult records an execution result and notifies hooks. func (m *Manager) MarkResult(ctx context.Context, result Result) { if result.AuthID == "" { @@ -843,10 +972,29 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { var authSnapshot *Auth cooldownStateChanged := false + scheduleUnauthorizedRefresh := false now := time.Now() + // A 401 condemns the access token, except one a -and-cooldown rule matched: that only cools the credential. + rejectsToken := !result.Success && (result.Error == nil || result.Error.Code != ErrorCodeForceCooldown) && !shouldSkipCredentialCooldown(result.Error) && (statusCodeFromResult(result.Error) == http.StatusUnauthorized || isUnauthorizedError(result.Error)) + releaseMutation := m.lockAuthMutation(result.AuthID) + defer releaseMutation() m.mu.Lock() if auth, ok := m.auths[result.AuthID]; ok && auth != nil { + if isStaleExecutionResult(result, auth) { + staleEpoch := result.RegistrationEpoch > 0 && result.RegistrationEpoch < auth.RegistrationEpoch + if !staleEpoch && rejectsToken && result.AccessToken != "" { + if recordRejectedTokenOnAuth(auth, result.AccessToken) { + auth.Generation++ + auth.UpdatedAt = now + } + } + m.mu.Unlock() + releaseMutation() + m.hook.OnResult(ctx, result) + m.publishErrorEvent(result, nil) + return + } if modelKey == "" && strings.TrimSpace(result.RouteModel) != "" { if m != nil { modelKey = m.selectionModelKeyForAuth(auth, result.RouteModel) @@ -871,6 +1019,22 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { } wasTerminalUnauthorized := hasUnauthorizedAuthFailure(auth) + // retainModelCooldown pins a still-live model cooldown against a success + // from an execution dispatched before that cooldown was recorded. Such a + // success predates the failure and proves nothing about recovery, so the + // state transition it normally causes is skipped entirely. + retainModelCooldown := modelKey != "" && !wasTerminalUnauthorized && result.ExecutionSeq > 0 + if retainModelCooldown { + if existing := existingModelState(auth, modelKey); existing != nil && isModelStateActiveCooldown(existing, now) && + existing.cooldownRecordedSeq > 0 && result.ExecutionSeq <= existing.cooldownRecordedSeq { + modelState = existing + } else { + retainModelCooldown = false + } + } + retainAuthCooldown := modelKey == "" && !wasTerminalUnauthorized && result.ExecutionSeq > 0 && + isAuthActiveCooldown(auth, now) && auth.cooldownRecordedSeq > 0 && result.ExecutionSeq <= auth.cooldownRecordedSeq + if result.Success { if wasTerminalUnauthorized { if modelKey != "" { @@ -880,6 +1044,8 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { } } else if auth.Quota.Reason == "credential_quota" && auth.Quota.NextRecoverAt.After(now) { // Retain active credential-scoped cooldown + } else if retainModelCooldown || retainAuthCooldown { + // Retain a cooldown recorded after this execution was dispatched. } else if modelKey != "" { state := ensureModelState(auth, modelKey) modelState = state @@ -902,6 +1068,41 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { } state := ensureModelState(auth, modelKey) modelState = state + // Normalize a naturally expired credential-quota wrap before + // applying a new result (issue #6690): reinstate the + // independent layer (or a clean slate) and drop stale wrap + // provenance so a future window snapshots fresh state. + if (state.credentialQuotaPropagated || state.Quota.Reason == "credential_quota") && + !(auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" && auth.Quota.NextRecoverAt.After(now)) { + if shadow := state.credentialQuotaShadow; shadow != nil && isModelStateActiveCooldown(shadow, now) { + preservedSeq := shadow.cooldownRecordedSeq + *state = *shadow.cloneShallow() + state.cooldownRecordedSeq = preservedSeq + } else { + resetModelState(state, now) + } + } + // Two-layer handling (issue #6690): while an auth-level + // credential quota wraps the credential, a non-credential + // own failure is applied to the independent layer (shadow) + // so its deadline/backoff never inherit the propagated + // quota; the outer fields are then composed from both + // layers. Credential-scoped failures refresh the outer + // layer and are never independent. + outerState := state + savedOuterState := state.cloneShallow() + ownLayerPending := !result.CredentialScope && + auth.Quota.Reason == "credential_quota" && auth.Quota.NextRecoverAt.After(now) + if ownLayerPending { + if state.credentialQuotaShadow != nil { + preservedSeq := state.credentialQuotaShadow.cooldownRecordedSeq + *state = *state.credentialQuotaShadow.cloneShallow() + state.cooldownRecordedSeq = preservedSeq + } else { + fresh := &ModelState{} + *state = *fresh + } + } state.Unavailable = true state.Status = StatusError state.UpdatedAt = now @@ -1003,9 +1204,59 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { NextRecoverAt: next, BackoffLevel: backoffLevel, }) + // Non-credential own failures re-derive their independent layer + // below, so a stale shadow is dropped. A credential-scoped + // report must not erase an existing independent layer: it only + // refreshes the propagated outer layer. + if !result.CredentialScope { + state.credentialQuotaShadow = nil + } else if !state.credentialQuotaPropagated && savedOuterState.Quota.Reason != "credential_quota" { + // First wrap of the reporting model: preserve any live prior + // own block (including an ordinary model-scoped 429, which + // quota restoration does not supersede). An outer + // credential_quota label instead proves a prior wrap whose + // state must never be snapshotted as independence. + priorIndependent := prevModelRetryAfter.After(now) && !isDevinQuotaExhaustionMessage(devinStateErrorText(savedOuterState)) + if priorIndependent { + preserved := savedOuterState.cloneShallow() + preserved.NextRetryAfter = prevModelRetryAfter + preserved.Unavailable = true + preserved.Status = StatusError + state.credentialQuotaShadow = preserved + } else { + state.credentialQuotaShadow = nil + } + } if result.CredentialScope && !disableCooling { for _, otherState := range auth.ModelStates { if otherState != nil && otherState != state { + // Normalize a naturally expired wrap first (issue #6690) so a + // stale marker cannot suppress snapshotting a fresh block. + if (otherState.credentialQuotaPropagated || otherState.Quota.Reason == "credential_quota") && + !(auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" && auth.Quota.NextRecoverAt.After(now)) { + if shadow := otherState.credentialQuotaShadow; shadow != nil && isModelStateActiveCooldown(shadow, now) { + preservedSeq := shadow.cooldownRecordedSeq + *otherState = *shadow.cloneShallow() + otherState.cooldownRecordedSeq = preservedSeq + } else { + resetModelState(otherState, now) + } + } + // On first wrap, preserve the sibling's independent layer + // beneath the propagated credential quota (issue #6690): + // a live own block is snapshotted; healthy/propagated-only + // states get none. The marker keeps repeated propagation + // from mistaking its own earlier layer for independence. + alreadyPropagated := otherState.credentialQuotaPropagated || + otherState.Quota.Reason == "credential_quota" + if !alreadyPropagated { + if isModelStateActiveCooldown(otherState, now) && !isDevinQuotaExhaustionMessage(devinStateErrorText(otherState)) { + otherState.credentialQuotaShadow = otherState.cloneShallow() + } else { + otherState.credentialQuotaShadow = nil + } + } + otherState.credentialQuotaPropagated = true otherState.Unavailable = true otherState.Status = StatusError otherQuotaNext := credentialNext @@ -1039,6 +1290,10 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { auth.Quota.NextRecoverAt = authNext auth.Quota.BackoffLevel = backoffLevel auth.NextRetryAfter = authNext + // The reporting model is now part of the wrap too: mark it + // so a later credential report from another model cannot + // snapshot this credential-derived state as independent. + state.credentialQuotaPropagated = true } } case 408, 500, 502, 503, 504, 520, 521, 522, 523, 524, 525, 526: @@ -1064,6 +1319,56 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { if !state.NextRetryAfter.IsZero() && prevModelRetryAfter.After(state.NextRetryAfter) && prevModelRetryAfter.After(now) { state.NextRetryAfter = prevModelRetryAfter } + // Capture the dispatch watermark before composing layers so the + // independent shadow carries this failure's stale-success + // protection (issue #6690). + cooldownSeq := m.currentExecutionSeq() + state.cooldownRecordedSeq = cooldownSeq + if ownLayerPending { + // Compose the outer presentation from the saved propagated + // layer and the independent layer just updated in place. + independent := state.cloneShallow() + outer := savedOuterState.cloneShallow() + if isModelStateActiveCooldown(independent, now) { + outer.credentialQuotaShadow = independent + } else { + outer.credentialQuotaShadow = nil + } + outer.cooldownRecordedSeq = cooldownSeq + // A composed outer state is wrapped by definition; stamp it so + // a later credential report cannot mistake it for unwrapped. + outer.credentialQuotaPropagated = true + outer.Unavailable = true + outer.Status = StatusError + outer.UpdatedAt = now + outer.LastError = cloneError(state.LastError) + outer.StatusMessage = state.StatusMessage + if state.NextRetryAfter.After(outer.NextRetryAfter) { + outer.NextRetryAfter = state.NextRetryAfter + } + if auth.Quota.NextRecoverAt.After(outer.NextRetryAfter) { + outer.NextRetryAfter = auth.Quota.NextRecoverAt + } + if !(outer.Quota.Exceeded && outer.Quota.Reason == "credential_quota") { + applyCooldownFields(&outer.Quota, QuotaState{ + Exceeded: true, + Reason: "credential_quota", + NextRecoverAt: auth.Quota.NextRecoverAt, + BackoffLevel: auth.Quota.BackoffLevel, + }) + } + *outerState = *outer + state = outerState + modelState = state + } + if result.CredentialScope && !disableCooling && statusCodeFromResult(result.Error) == 429 { + auth.cooldownRecordedSeq = cooldownSeq + for _, otherState := range auth.ModelStates { + if otherState != nil && otherState != state { + otherState.cooldownRecordedSeq = cooldownSeq + } + } + } auth.Status = StatusError updateAggregatedAvailability(auth, now) } @@ -1074,15 +1379,32 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { } if !wasTerminalUnauthorized { applyAuthFailureState(auth, result.Error, result.RetryAfter, now, disableCooling) + auth.cooldownRecordedSeq = m.currentExecutionSeq() } } } + if rejectsToken { + tokenToReject := result.AccessToken + if tokenToReject == "" { + tokenToReject = authAccessTokenForRequest(ctx, auth, result.Options.Headers) + } + _ = recordRejectedTokenOnAuth(auth, tokenToReject) + } if wasTerminalUnauthorized { auth.Unavailable = true auth.Status = StatusError auth.NextRefreshAfter = time.Time{} auth.NextRetryAfter = time.Time{} + } else if rejectsToken && authHasRefreshCredential(auth) { + if auth.NextRefreshAfter.IsZero() { + if !auth.LastRefreshedAt.IsZero() && now.Sub(auth.LastRefreshedAt) < refreshFailureBackoff { + auth.NextRefreshAfter = now.Add(refreshFailureBackoff) + } else { + auth.NextRefreshAfter = now + } + } + scheduleUnauthorizedRefresh = true } auth.Generation++ @@ -1096,7 +1418,7 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { } } - _ = m.persist(ctx, auth) + _ = m.persistLocked(ctx, auth) authSnapshot = auth.Clone() if trackCooldownState { cooldownRecordsAfter := m.cooldownStateRecordsForAuthLocked(auth, now) @@ -1105,19 +1427,30 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { } } m.mu.Unlock() + releaseMutation() if m.scheduler != nil && authSnapshot != nil { + credentialScopedResult := result.CredentialScope || scheduleUnauthorizedRefresh var targetModels []string - if !result.CredentialScope && modelKey != "" { + if !credentialScopedResult && modelKey != "" { targetModels = append(targetModels, modelKey) if routeKey := canonicalModelKey(result.RouteModel); routeKey != "" && routeKey != modelKey { targetModels = append(targetModels, routeKey) } } - m.scheduler.upsertAuthResult(authSnapshot, targetModels, result.CredentialScope) + m.scheduler.upsertAuthResult(authSnapshot, targetModels, credentialScopedResult) } if authSnapshot != nil && cooldownStateChanged { m.persistCooldownStates(context.Background()) } + // A newly exhausted Devin credential must join the refresh schedule so + // quota revalidation can detect early restoration (issue #6690). + if authSnapshot != nil { + if _, due := devinQuotaRevalidateAt(authSnapshot, now); due { + m.queueRefreshReschedule(result.AuthID) + } else if scheduleUnauthorizedRefresh { + m.queueRefreshReschedule(result.AuthID) + } + } supportedModels, regEpoch := registry.GetGlobalRegistry().GetModelsAndEpochForClient(result.AuthID) reg := registry.GetGlobalRegistry() @@ -1150,6 +1483,17 @@ func (m *Manager) updateSessionAffinity(result Result) { } func (m *Manager) recordExecutionResult(ctx context.Context, result Result, auth *Auth, ephemeral bool) { + if auth != nil { + if result.CredentialVersion == 0 { + result.CredentialVersion = auth.CredentialVersion + } + if result.RegistrationEpoch == 0 { + result.RegistrationEpoch = auth.RegistrationEpoch + } + if result.AccessToken == "" { + result.AccessToken = authAccessTokenForRequest(ctx, auth, result.Options.Headers) + } + } if !ephemeral { m.MarkResult(ctx, result) return @@ -1177,8 +1521,17 @@ func (m *Manager) recordAvailabilityNeutralResult(ctx context.Context, result Re } var authSnapshot *Auth + releaseMutation := m.lockAuthMutation(result.AuthID) + defer releaseMutation() m.mu.Lock() if auth, ok := m.auths[result.AuthID]; ok && auth != nil { + if isStaleExecutionResult(result, auth) { + m.mu.Unlock() + releaseMutation() + m.hook.OnResult(ctx, result) + m.publishErrorEvent(result, nil) + return + } now := time.Now() auth.recordRecentRequest(now, result.Success) if result.Success { @@ -1188,10 +1541,11 @@ func (m *Manager) recordAvailabilityNeutralResult(ctx context.Context, result Re } auth.Generation++ auth.UpdatedAt = now - _ = m.persist(ctx, auth) + _ = m.persistLocked(ctx, auth) authSnapshot = auth.Clone() } m.mu.Unlock() + releaseMutation() m.hook.OnResult(ctx, result) m.publishErrorEvent(result, authSnapshot) @@ -1291,6 +1645,14 @@ func mergeModelState(target, source *ModelState) *ModelState { if source.UpdatedAt.After(merged.UpdatedAt) { merged.UpdatedAt = source.UpdatedAt } + // Keep the strongest cooldown watermark from either side so merging a + // restored or alias-migrated state into a live auth cannot reopen the + // stale-success race. Both sides stay zero across a process restart, + // where restored cooldowns must remain clearable by any success. + merged.cooldownRecordedSeq = target.cooldownRecordedSeq + if source.cooldownRecordedSeq > merged.cooldownRecordedSeq { + merged.cooldownRecordedSeq = source.cooldownRecordedSeq + } if merged.StatusMessage == "" { merged.StatusMessage = fallback.StatusMessage } @@ -1305,6 +1667,22 @@ func mergeModelState(target, source *ModelState) *ModelState { } else if merged.Unavailable || merged.Quota.Exceeded { merged.Status = StatusError } + // Keep the independent layer and wrap provenance from either side + // (newer wins) while the merged state is wrapped. The reporting model + // carries outer reason "quota", so wrap detection must include the + // marker and shadow presence, not only the propagated label. + if merged.Quota.Reason == "credential_quota" || merged.Quota.Reason == "quota" || + target.credentialQuotaPropagated || source.credentialQuotaPropagated || + target.credentialQuotaShadow != nil || source.credentialQuotaShadow != nil { + shadow := fallback.credentialQuotaShadow + if preferred.credentialQuotaShadow != nil { + shadow = preferred.credentialQuotaShadow + } + if shadow != nil { + merged.credentialQuotaShadow = shadow.cloneShallow() + } + merged.credentialQuotaPropagated = target.credentialQuotaPropagated || source.credentialQuotaPropagated + } *target = merged return target } @@ -1319,6 +1697,9 @@ func resetModelState(state *ModelState, now time.Time) { state.NextRetryAfter = time.Time{} state.LastError = nil applyCooldownFields(&state.Quota, QuotaState{}) + state.cooldownRecordedSeq = 0 + state.credentialQuotaShadow = nil + state.credentialQuotaPropagated = false state.UpdatedAt = now } @@ -1539,9 +1920,28 @@ func clearAuthStateOnSuccess(auth *Auth, now time.Time) { auth.Quota.BackoffLevel = 0 auth.LastError = nil auth.NextRetryAfter = time.Time{} + auth.cooldownRecordedSeq = 0 auth.UpdatedAt = now } +// isAuthActiveCooldown reports whether the auth carries a still-live +// auth-level cooldown recorded without a model key. +func isAuthActiveCooldown(auth *Auth, now time.Time) bool { + if auth == nil || auth.Disabled || auth.Status == StatusDisabled { + return false + } + if !auth.NextRetryAfter.IsZero() && auth.NextRetryAfter.After(now) { + return true + } + if !auth.Quota.NextRecoverAt.IsZero() && auth.Quota.NextRecoverAt.After(now) { + return true + } + if auth.Quota.Exceeded && auth.Quota.NextRecoverAt.IsZero() { + return true + } + return false +} + func cloneError(err *Error) *Error { if err == nil { return nil @@ -1583,6 +1983,28 @@ func isRequestScopedError(err error) bool { return ok && requestErr != nil && requestErr.IsRequestScoped() } +// isStreamTruncatedError reports locally synthesized failures for accepted +// upstream streams that ended without a terminal frame. Unlike a received +// HTTP refusal, the generation was already submitted upstream. +func isStreamTruncatedError(err error) bool { + if err == nil { + return false + } + truncatedErr, ok := errors.AsType[cliproxyexecutor.StreamTruncatedError](err) + if ok && truncatedErr != nil && truncatedErr.IsStreamTruncated() { + return true + } + var authErr *Error + if errors.As(err, &authErr) && authErr != nil && authErr.Code == streamTruncatedErrorCode { + return true + } + return false +} + +func isStreamTruncatedResultError(err *Error) bool { + return err != nil && err.Code == streamTruncatedErrorCode +} + func resultErrorFromError(err error) *Error { if err == nil { return nil @@ -1598,6 +2020,10 @@ func resultErrorFromError(err error) *Error { resultErr.HTTPStatus = statusCodeFromError(err) } switch { + case isStreamTruncatedError(err): + // Preserve the local truncation classification alongside the public + // 502 status so MarkResult can skip credential cooldown. + resultErr.Code = streamTruncatedErrorCode case isExplicitModelNotFoundError(err, ""): if resultErr.Code == "" || resultErr.Code == requestScopedErrorCode { resultErr.Code = "model_not_found" @@ -1606,6 +2032,12 @@ func resultErrorFromError(err error) *Error { // Prefer true request-scoped faults (including Claude OAuth cancellation) // over the broader connection-lifecycle classification. resultErr.Code = requestScopedErrorCode + case isConfiguredTLSHandshakeTimeoutError(err): + // An internal TLS establishment deadline is a transient transport + // failure, not a caller-owned context lifecycle event. + if resultErr.Code == "" || resultErr.Code == transientTransportErrorCode { + resultErr.Code = transientTransportErrorCode + } case isConnectionLifecycleError(err): // Preserve lifecycle classification for MarkResult without making the error // request-scoped (which would also stop credential fallback). @@ -1627,11 +2059,11 @@ func shouldSkipCredentialCooldown(err *Error) bool { if err != nil && err.Code == ErrorCodeForceCooldown { return false } - return isRequestScopedResultError(err) || isConnectionLifecycleResultError(err) || isTransientTransportResultError(err) + return isRequestScopedResultError(err) || isConnectionLifecycleResultError(err) || isTransientTransportResultError(err) || isStreamTruncatedResultError(err) } // isConnectionLifecycleError reports transport/session lifecycle failures that must -// not cool credentials: client cancellation and WebSocket close/EOF disconnects. +// not cool credentials: client cancellation and WebSocket/HTTP/2 close/EOF disconnects. func isConnectionLifecycleError(err error) bool { if err == nil { return false @@ -1648,6 +2080,9 @@ func isConnectionLifecycleError(err error) bool { if statusCodeFromError(err) != 0 { return false } + if isConfiguredTLSHandshakeTimeoutError(err) { + return false + } // Client abort and request-scoped timeouts are not credential faults. if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) || errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) { return true @@ -1662,6 +2097,9 @@ func isConnectionLifecycleResultError(err *Error) bool { if err.Code == connectionLifecycleErrorCode { return true } + if err.Code == transientTransportErrorCode { + return false + } // Message fallback only when no HTTP status is attached, so 401/429/5xx // response bodies cannot suppress credential cooldown. if statusCodeFromResult(err) != 0 { @@ -1685,13 +2123,31 @@ func isConnectionLifecycleMessage(message string) bool { strings.Contains(lower, "websocket: close 1006") { return true } - // Wrapped transport EOF phrasing (e.g. "read tcp ...: unexpected EOF"). - if strings.Contains(lower, "unexpected eof") { + // Wrapped transport disconnects (e.g. "read tcp ...: unexpected EOF"). + if strings.Contains(lower, "unexpected eof") || isHTTP2ClosedConnectionMessage(lower) { return true } return false } +func isHTTP2ClosedConnectionMessage(lower string) bool { + // x/net/http2 keeps these errors unexported; allow common wrappers such as url.Error. + return strings.Contains(lower, "http2: client connection force closed via clientconn.close") || + strings.Contains(lower, "http2: client conn is closed") +} + +// isConfiguredTLSHandshakeTimeoutError identifies the typed internal +// establishment deadline produced by the configured TLS handshake bound. +// Delegating to the shared executor predicate ensures classification is +// based on error identity rather than message text and that any +// caller-owned context error outside a configured-timeout subtree wins. +func isConfiguredTLSHandshakeTimeoutError(err error) bool { + if err == nil || statusCodeFromError(err) != 0 { + return false + } + return cliproxyexecutor.IsConfiguredTLSHandshakeTimeoutError(err) +} + // isTransientTransportError reports pre-HTTP dial/TLS/DNS/reset failures that // should retry under request-retry without cooling the selected credential. func isTransientTransportError(err error) bool { @@ -1702,6 +2158,9 @@ func isTransientTransportError(err error) bool { if statusCodeFromError(err) != 0 { return false } + if isConfiguredTLSHandshakeTimeoutError(err) { + return true + } if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { return false } @@ -1736,6 +2195,14 @@ func isTransientTransportResultError(err *Error) bool { if statusCodeFromResult(err) != 0 { return false } + // Once an error has been reduced to a Result message, typed identity is + // no longer available. An explicit transient code assigned by + // resultErrorFromError is authoritative for configured handshake + // timeouts; matching text alone must not promote an external or + // caller-owned deadline that happens to use the same phrase. + if strings.Contains(strings.ToLower(err.Message), "tls handshake timeout") { + return false + } return isTransientTransportMessage(err.Message) } @@ -1774,7 +2241,8 @@ func isTransientTransportMessage(message string) bool { strings.Contains(lower, "broken pipe"), strings.Contains(lower, "connection aborted"), strings.Contains(lower, "use of closed network connection"), - strings.Contains(lower, "unexpected eof"): + strings.Contains(lower, "unexpected eof"), + isHTTP2ClosedConnectionMessage(lower): return true default: return false @@ -2289,6 +2757,11 @@ func isRequestInvalidError(err error) bool { if err == nil { return false } + if isStreamTruncatedError(err) { + // A truncated accepted stream must stop here: the generation was + // already submitted and must not be replayed on another credential. + return true + } if isRequestScopedError(err) { return true } diff --git a/sdk/cliproxy/auth/conductor_execution.go b/sdk/cliproxy/auth/conductor_execution.go index 9616a8217..ddd4066dc 100644 --- a/sdk/cliproxy/auth/conductor_execution.go +++ b/sdk/cliproxy/auth/conductor_execution.go @@ -407,7 +407,7 @@ func requestToFormat(provider string, executor ProviderExecutor, req cliproxyexe } } source := opts.SourceFormat.String() - if source == "openai-image" || source == "openai-video" { + if source == "openai-image" || source == "openai-video" || source == "openai-speech" { return opts.SourceFormat } if opts.Alt == "responses/compact" && !opts.Stream { @@ -505,6 +505,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req } return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"} } + opts.Metadata = cloneRequestMetadata(opts.Metadata) pickOpts := opts if homeMode { pickOpts = withHomeRetryRound(pickOpts, retryRound) @@ -529,7 +530,8 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt) execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt) } - execCtx = contextWithRequestedModelAlias(execCtx, opts, routeModel) + execCtx = contextWithRequestedModelAlias(execCtx, pickOpts, routeModel) + execCtx = coreusage.WithSelectionID(execCtx, selectionIDFromOptions(pickOpts)) execCtx = newUpstreamAttemptContext(execCtx) models, pooled, aliasResult, routing := m.preparedExecutionModelsWithAlias(auth, routeModel) @@ -547,7 +549,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req if stateModel == "" { stateModel = canonicalModelKey(routeModel) } - result := Result{AuthID: auth.ID, Provider: provider, Model: stateModel, RouteModel: routeModel, Success: false, Error: resultErrorFromError(errPrepare), Options: pickOpts} + result := Result{AuthID: auth.ID, Provider: provider, Model: stateModel, RouteModel: routeModel, Success: false, Error: resultErrorFromError(errPrepare), Options: pickOpts, AccessToken: authAccessTokenForRequest(execCtx, auth, pickOpts.Headers), CredentialVersion: auth.CredentialVersion, RegistrationEpoch: auth.RegistrationEpoch} m.MarkResult(execCtx, result) lastErr = errPrepare continue @@ -591,6 +593,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req } execReq = attachResolvedExecutionModelInfo(routing, execReq, auth, routeModel, upstreamModel, restoreExecutionModel) execCtx = syncMetadataSessionToContext(execCtx, execOpts.Metadata) + executionSeq := m.nextExecutionSeq() startExec := time.Now() resp, errExec := executor.Execute(execCtx, auth, execReq, execOpts) errExec = markUpstreamExecutionAttemptFromContext(execCtx, errExec) @@ -603,11 +606,12 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req return cliproxyexecutor.Response{}, errCtx } refreshCtx := newUpstreamAttemptContext(execCtx) - if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(refreshCtx, auth, errExec, didRefreshOnUnauthorized); okRefresh { + if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(refreshCtx, auth, errExec, didRefreshOnUnauthorized, execOpts.Headers); okRefresh { auth = refreshed didRefreshOnUnauthorized = true execCtx = newUpstreamAttemptContext(execCtx) execCtx = syncMetadataSessionToContext(execCtx, execOpts.Metadata) + executionSeq = m.nextExecutionSeq() startRetry := time.Now() resp, errExec = executor.Execute(execCtx, auth, execReq, execOpts) errExec = markUpstreamExecutionAttemptFromContext(execCtx, errExec) @@ -628,7 +632,7 @@ func (m *Manager) executeMixedOnce(ctx context.Context, providers []string, req if errCancel := claudeOAuthRequestCancellation(execCtx, auth, errExec); errCancel != nil { return cliproxyexecutor.Response{}, errCancel } - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: errExec == nil, Options: execOpts} + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: errExec == nil, Options: execOpts, AccessToken: authAccessTokenForRequest(execCtx, auth, execOpts.Headers), CredentialVersion: auth.CredentialVersion, RegistrationEpoch: auth.RegistrationEpoch, ExecutionSeq: executionSeq} if errExec != nil { result.Error = resultErrorFromError(errExec) if ra := retryAfterFromError(errExec); ra != nil { @@ -717,6 +721,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, } return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"} } + opts.Metadata = cloneRequestMetadata(opts.Metadata) pickOpts := opts if homeMode { pickOpts = withHomeRetryRound(pickOpts, retryRound) @@ -741,7 +746,8 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt) execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt) } - execCtx = contextWithRequestedModelAlias(execCtx, opts, routeModel) + execCtx = contextWithRequestedModelAlias(execCtx, pickOpts, routeModel) + execCtx = coreusage.WithSelectionID(execCtx, selectionIDFromOptions(pickOpts)) execCtx = newUpstreamAttemptContext(execCtx) models, pooled, aliasResult, routing := m.preparedExecutionModelsWithAlias(auth, routeModel) @@ -759,7 +765,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, if stateModel == "" { stateModel = canonicalModelKey(routeModel) } - result := Result{AuthID: auth.ID, Provider: provider, Model: stateModel, RouteModel: routeModel, Success: false, Error: resultErrorFromError(errPrepare), Options: pickOpts, SkipQuotaObservation: true} + result := Result{AuthID: auth.ID, Provider: provider, Model: stateModel, RouteModel: routeModel, Success: false, Error: resultErrorFromError(errPrepare), Options: pickOpts, SkipQuotaObservation: true, AccessToken: authAccessTokenForRequest(execCtx, auth, pickOpts.Headers), CredentialVersion: auth.CredentialVersion, RegistrationEpoch: auth.RegistrationEpoch} m.MarkResult(execCtx, result) lastErr = errPrepare continue @@ -803,6 +809,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, } execReq = attachResolvedExecutionModelInfo(routing, execReq, auth, routeModel, upstreamModel, restoreExecutionModel) execCtx = syncMetadataSessionToContext(execCtx, execOpts.Metadata) + executionSeq := m.nextExecutionSeq() startExec := time.Now() resp, errExec := executor.CountTokens(execCtx, auth, execReq, execOpts) errExec = markUpstreamExecutionAttemptFromContext(execCtx, errExec) @@ -815,11 +822,12 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, return cliproxyexecutor.Response{}, errCtx } refreshCtx := newUpstreamAttemptContext(execCtx) - if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(refreshCtx, auth, errExec, didRefreshOnUnauthorized); okRefresh { + if refreshed, okRefresh := m.tryRefreshAfterUnauthorized(refreshCtx, auth, errExec, didRefreshOnUnauthorized, execOpts.Headers); okRefresh { auth = refreshed didRefreshOnUnauthorized = true execCtx = newUpstreamAttemptContext(execCtx) execCtx = syncMetadataSessionToContext(execCtx, execOpts.Metadata) + executionSeq = m.nextExecutionSeq() startRetry := time.Now() resp, errExec = executor.CountTokens(execCtx, auth, execReq, execOpts) errExec = markUpstreamExecutionAttemptFromContext(execCtx, errExec) @@ -840,7 +848,7 @@ func (m *Manager) executeCountMixedOnce(ctx context.Context, providers []string, if errCancel := claudeOAuthRequestCancellation(execCtx, auth, errExec); errCancel != nil { return cliproxyexecutor.Response{}, errCancel } - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: errExec == nil, Options: execOpts, SkipQuotaObservation: true} + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: errExec == nil, Options: execOpts, SkipQuotaObservation: true, AccessToken: authAccessTokenForRequest(execCtx, auth, execOpts.Headers), CredentialVersion: auth.CredentialVersion, RegistrationEpoch: auth.RegistrationEpoch, ExecutionSeq: executionSeq} if errExec != nil { result.Error = resultErrorFromError(errExec) if ra := retryAfterFromError(errExec); ra != nil { @@ -944,6 +952,7 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string } return nil, &Error{Code: "auth_not_found", Message: "no auth available"} } + opts.Metadata = cloneRequestMetadata(opts.Metadata) pickOpts := opts if homeMode { pickOpts = withHomeRetryRound(pickOpts, retryRound) @@ -1058,9 +1067,18 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt) } // Enrich before auth preparation so prepare-stage usage records observe the client request. - execCtx = contextWithRequestedModelAlias(execCtx, opts, routeModel) + if homeMode { + clearSchedulerSelectionID(pickOpts.Metadata) + clearSchedulerSelectionID(opts.Metadata) + } + execCtx = contextWithRequestedModelAlias(execCtx, pickOpts, routeModel) + execCtx = coreusage.WithSelectionID(execCtx, selectionIDFromOptions(pickOpts)) execCtx = newUpstreamAttemptContext(execCtx) - models, pooled, aliasResult, routing := m.preparedExecutionModelsWithAlias(auth, routeModel) + models, pooled, aliasResult, routing := m.executionModelCandidatesWithAlias(auth, routeModel) + // Home owns dispatched credential availability; only local selections need state filtering. + if selection == nil { + models = m.filterExecutionModels(auth, routeModel, models, pooled) + } if selection != nil && aliasResult.ForceMapping && responseAlias != "" { aliasResult.OriginalAlias = responseAlias } @@ -1104,7 +1122,7 @@ func (m *Manager) executeStreamMixedOnce(ctx context.Context, providers []string if stateModel == "" { stateModel = canonicalModelKey(routeModel) } - result := Result{AuthID: auth.ID, Provider: provider, Model: stateModel, RouteModel: routeModel, Success: false, Error: resultErrorFromError(errPrepare), Options: pickOpts} + result := Result{AuthID: auth.ID, Provider: provider, Model: stateModel, RouteModel: routeModel, Success: false, Error: resultErrorFromError(errPrepare), Options: pickOpts, AccessToken: authAccessTokenForRequest(execCtx, auth, pickOpts.Headers), CredentialVersion: auth.CredentialVersion, RegistrationEpoch: auth.RegistrationEpoch} if selection != nil { m.reportHomeResult(execCtx, result, auth) releaseAttempt() @@ -1607,9 +1625,18 @@ func contextWithRequestedModelAlias(ctx context.Context, opts cliproxyexecutor.O ctx = coreusage.WithGenerate(ctx, generate) } ctx = coreusage.WithStream(ctx, opts.Stream) + if selectionID := selectionIDFromOptions(opts); selectionID != "" { + ctx = coreusage.WithSelectionID(ctx, selectionID) + } else if _, exists := opts.Metadata[cliproxyexecutor.SelectionIDMetadataKey]; exists { + ctx = coreusage.WithSelectionID(ctx, "") + } return ctx } +func selectionIDFromOptions(opts cliproxyexecutor.Options) string { + return stringMetadataValue(opts.Metadata, cliproxyexecutor.SelectionIDMetadataKey) +} + func requestedModelAliasFromOptions(opts cliproxyexecutor.Options, fallback string) string { fallback = strings.TrimSpace(fallback) if len(opts.Metadata) == 0 { diff --git a/sdk/cliproxy/auth/conductor_home.go b/sdk/cliproxy/auth/conductor_home.go index 403aa4ced..b3a0be969 100644 --- a/sdk/cliproxy/auth/conductor_home.go +++ b/sdk/cliproxy/auth/conductor_home.go @@ -16,6 +16,7 @@ import ( "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executionregistry" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" + coreusage "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/usage" log "github.com/sirupsen/logrus" ) @@ -281,7 +282,7 @@ func isHomeRequestRetryExceededError(err error) bool { } func shouldReturnLastErrorOnPickFailure(homeMode bool, lastErr error, errPick error) bool { - if lastErr == nil { + if lastErr == nil || isRequestStopError(errPick) { return false } if !homeMode { @@ -1314,6 +1315,7 @@ func (m *Manager) findAllAntigravityCreditsCandidateAuths(ctx context.Context, r } m.mu.RUnlock() + now := time.Now() var known []creditsCandidateEntry var unknown []creditsCandidateEntry for _, candidate := range candidates { @@ -1323,6 +1325,10 @@ func (m *Manager) findAllAntigravityCreditsCandidateAuths(ctx context.Context, r } if okHint && hint.Known { if !hint.Available { + if IsAntigravityCreditsHintFresh(hint, now) { + continue + } + unknown = append(unknown, candidate) continue } known = append(known, candidate) @@ -1355,7 +1361,7 @@ func hasAntigravityProvider(providers []string) bool { } func shouldAttemptAntigravityCreditsFallback(m *Manager, lastErr error, providers []string) bool { - if isRequestTerminatedError(lastErr) { + if isRequestTerminatedError(lastErr) || isRequestStopError(lastErr) { return false } status := statusCodeFromError(lastErr) @@ -1411,7 +1417,9 @@ func (m *Manager) tryAntigravityCreditsExecute(ctx context.Context, req cliproxy creditsCtx = context.WithValue(creditsCtx, "cliproxy.roundtripper", rt) } creditsOpts := ensureRequestedModelMetadata(opts, routeModel) + clearSchedulerSelectionID(creditsOpts.Metadata) creditsCtx = contextWithRequestedModelAlias(creditsCtx, creditsOpts, routeModel) + creditsCtx = coreusage.WithSelectionID(creditsCtx, "") preparedAuth, errPrepare := m.prepareRequestAuth(creditsCtx, c.executor, c.auth) if errPrepare != nil { continue @@ -1427,8 +1435,9 @@ func (m *Manager) tryAntigravityCreditsExecute(ctx context.Context, req cliproxy execReq := req execReq.Model = upstreamModel creditsCtx = syncMetadataSessionToContext(creditsCtx, creditsOpts.Metadata) + executionSeq := m.nextExecutionSeq() resp, errExec := c.executor.Execute(creditsCtx, c.auth, execReq, creditsOpts) - result := Result{AuthID: c.auth.ID, Provider: c.provider, Model: resultModel, RouteModel: routeModel, Success: errExec == nil, Options: creditsOpts} + result := Result{AuthID: c.auth.ID, Provider: c.provider, Model: resultModel, RouteModel: routeModel, Success: errExec == nil, Options: creditsOpts, CredentialVersion: c.auth.CredentialVersion, RegistrationEpoch: c.auth.RegistrationEpoch, ExecutionSeq: executionSeq} if errExec != nil { result.Error = resultErrorFromError(errExec) if ra := retryAfterFromError(errExec); ra != nil { @@ -1474,6 +1483,9 @@ func (m *Manager) tryAntigravityCreditsExecuteStream(ctx context.Context, req cl creditsCtx = context.WithValue(creditsCtx, "cliproxy.roundtripper", rt) } creditsOpts := ensureRequestedModelMetadata(opts, routeModel) + clearSchedulerSelectionID(creditsOpts.Metadata) + creditsCtx = contextWithRequestedModelAlias(creditsCtx, creditsOpts, routeModel) + creditsCtx = coreusage.WithSelectionID(creditsCtx, "") preparedAuth, errPrepare := m.prepareRequestAuth(creditsCtx, c.executor, c.auth) if errPrepare != nil { continue diff --git a/sdk/cliproxy/auth/conductor_home_execution.go b/sdk/cliproxy/auth/conductor_home_execution.go index 1b8f23387..32037eac2 100644 --- a/sdk/cliproxy/auth/conductor_home_execution.go +++ b/sdk/cliproxy/auth/conductor_home_execution.go @@ -8,6 +8,7 @@ import ( "time" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" + coreusage "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/usage" "github.com/tidwall/sjson" ) @@ -72,6 +73,7 @@ func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req c responseAlias := requestedModelAliasFromOptions(opts, routeModel) executionModel, restoreExecutionModel := executionModelForAuthSelection(opts, req.Model) opts = ensureRequestedModelMetadata(opts, routeModel) + clearSchedulerSelectionID(opts.Metadata) tried := make(map[string]struct{}) attempted := make(map[string]struct{}) var lastErr error @@ -130,13 +132,16 @@ func (m *Manager) executeHomeOnce(ctx context.Context, providers []string, req c return cliproxyexecutor.Response{}, errBind } // Enrich before auth preparation so prepare-stage usage records observe the client request. + clearSchedulerSelectionID(opts.Metadata) execCtx = contextWithRequestedModelAlias(execCtx, opts, routeModel) + execCtx = coreusage.WithSelectionID(execCtx, "") execCtx = newUpstreamAttemptContext(execCtx) if rt := m.roundTripperFor(auth); rt != nil { execCtx = context.WithValue(execCtx, roundTripperContextKey{}, rt) execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt) } - models, pooled, aliasResult, routing := m.preparedExecutionModelsWithAlias(auth, routeModel) + // Home owns credential availability; execute its dispatch without local state filtering. + models, pooled, aliasResult, routing := m.executionModelCandidatesWithAlias(auth, routeModel) if aliasResult.ForceMapping && responseAlias != "" { aliasResult.OriginalAlias = responseAlias } diff --git a/sdk/cliproxy/auth/conductor_lifecycle.go b/sdk/cliproxy/auth/conductor_lifecycle.go index c8f2f3173..55f990d89 100644 --- a/sdk/cliproxy/auth/conductor_lifecycle.go +++ b/sdk/cliproxy/auth/conductor_lifecycle.go @@ -3,6 +3,7 @@ package auth import ( "context" "fmt" + "math" "strings" "sync" "time" @@ -60,6 +61,9 @@ func (m *Manager) RegisterExecutor(executor ProviderExecutor) { if replaced == nil || replaced == executor { return } + if sharer, ok := replaced.(ExecutionSessionSharer); ok && sharer != nil && sharer.SharesExecutionSessionsWith(executor) { + return + } if closer, ok := replaced.(ExecutionSessionCloser); ok && closer != nil { closer.CloseExecutionSession(CloseAllExecutionSessionsID) } @@ -88,6 +92,11 @@ func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) { if auth.ID == "" { auth.ID = uuid.NewString() } + releaseMutation, errMutation := m.lockAuthMutationContext(ctx, auth.ID) + if errMutation != nil { + return nil, errMutation + } + defer releaseMutation() now := time.Now() if auth.Generation == 0 { auth.Generation = 1 @@ -113,15 +122,33 @@ func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) { } m.authEpochs[auth.ID]++ auth.RegistrationEpoch = m.authEpochs[auth.ID] + if existing, exists := m.auths[auth.ID]; exists && existing != nil { + auth.CredentialVersion = max(auth.CredentialVersion, existing.CredentialVersion) + 1 + } else if auth.CredentialVersion == 0 { + auth.CredentialVersion = 1 + } auth.Generation = 1 + // Serialize this credential, but release the manager lock during store I/O. + // Persist failures stay non-fatal, but must not be silent: a restart would + // lose the credential that only exists in memory. + if errPersist := m.persistLocked(ctx, auth); errPersist != nil { + log.WithFields(log.Fields{"auth_id": auth.ID, "credential": auth.ID, "provider": auth.Provider}).Warnf("failed to persist registered auth %s (%s): %v", auth.Provider, auth.ID, errPersist) + } + if m.authEpochs[auth.ID] != auth.RegistrationEpoch { + m.mu.Unlock() + return nil, fmt.Errorf("auth %s changed registration during persistence", auth.ID) + } + auth.ClearFileSnapshot() authClone := auth.Clone() m.auths[auth.ID] = authClone + m.notifyAuthChangeLocked(auth.ID) // Snapshot before unlocking: MarkResult mutates the published auth in place. var schedulerSnapshot *Auth if m.scheduler != nil { schedulerSnapshot = authClone.Clone() } m.mu.Unlock() + releaseMutation() if !shouldDeferAPIKeyModelAliasRebuild(ctx) { m.rebuildAPIKeyModelAliasFromRuntimeConfig() } @@ -130,11 +157,6 @@ func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) { } m.structuralEpoch.Add(1) m.queueRefreshReschedule(auth.ID) - // Persist failures stay non-fatal, but must not be silent: a restart would - // lose the credential that only exists in memory. - if errPersist := m.persist(ctx, auth); errPersist != nil { - log.WithFields(log.Fields{"auth_id": auth.ID, "credential": auth.ID, "provider": auth.Provider}).Warnf("failed to persist registered auth %s (%s): %v", auth.Provider, auth.ID, errPersist) - } m.hook.OnAuthRegistered(ctx, auth.Clone()) if cooldownStateChanged { m.persistCooldownStates(context.Background()) @@ -167,19 +189,128 @@ func (m *Manager) Update(ctx context.Context, auth *Auth) (*Auth, error) { return m.updateInternal(ctx, nil, auth, updateModeReplace) } +// Modify loads the latest auth for id while holding the per-auth mutation gate, +// applies edit to a clone of that auth, and commits the result through the +// standard persistence and publication path. It returns (nil, false, nil) if +// the auth is not present when the gate is acquired. +func (m *Manager) Modify(ctx context.Context, id string, expectedEpoch uint64, edit func(*Auth) error) (*Auth, bool, error) { + id = strings.TrimSpace(id) + if m == nil || id == "" || edit == nil { + return nil, false, nil + } + releaseMutation, errMutation := m.lockAuthMutationContext(ctx, id) + if errMutation != nil { + return nil, false, errMutation + } + defer releaseMutation() + + m.mu.RLock() + existing, ok := m.auths[id] + var current *Auth + var currentEpoch uint64 + if ok && existing != nil { + current = existing.Clone() + currentEpoch = existing.RegistrationEpoch + if m.authEpochs != nil && m.authEpochs[id] > currentEpoch { + currentEpoch = m.authEpochs[id] + } + } + m.mu.RUnlock() + if current == nil { + return nil, false, nil + } + if expectedEpoch != 0 && expectedEpoch < currentEpoch { + return nil, false, fmt.Errorf("update auth %s: stale registration epoch %d < %d", id, expectedEpoch, currentEpoch) + } + if errEdit := edit(current); errEdit != nil { + return nil, false, errEdit + } + if expectedEpoch != 0 { + current.RegistrationEpoch = expectedEpoch + } + committedAuth, didCommit, persistCooldown, errUpdate := m.updateInternalLocked(ctx, nil, current, updateModeReplace) + releaseMutation() + if errUpdate == nil && didCommit && committedAuth != nil { + m.finishAuthUpdate(ctx, committedAuth, persistCooldown) + } + return committedAuth, errUpdate == nil && didCommit && committedAuth != nil, errUpdate +} + func (m *Manager) updateInternal(ctx context.Context, base, auth *Auth, mode updateAuthMode) (*Auth, error) { if auth == nil || auth.ID == "" { return nil, nil } + releaseMutation, errMutation := m.lockAuthMutationContext(ctx, auth.ID) + if errMutation != nil { + return nil, errMutation + } + committedAuth, didCommit, persistCooldown, err := m.updateInternalLocked(ctx, base, auth, mode) + releaseMutation() + if err == nil && didCommit && committedAuth != nil { + m.finishAuthUpdate(ctx, committedAuth, persistCooldown) + } + return committedAuth, err +} + +// finishAuthUpdate runs post-commit side effects after the per-auth mutation +// gate has been released, preserving the original update ordering in which +// hooks and scheduler updates could re-enter manager operations. +func (m *Manager) finishAuthUpdate(ctx context.Context, committed *Auth, persistCooldown bool) { + if committed == nil { + return + } + if !shouldDeferAPIKeyModelAliasRebuild(ctx) { + m.rebuildAPIKeyModelAliasFromRuntimeConfig() + } + if m.scheduler != nil { + m.scheduler.upsertAuth(committed.Clone()) + } + m.structuralEpoch.Add(1) + m.queueRefreshReschedule(committed.ID) + m.hook.OnAuthUpdated(ctx, committed.Clone()) + if persistCooldown { + m.persistCooldownStates(context.Background()) + } +} + +// updateInternalLocked assumes the caller already holds the per-auth +// mutation lock. It only commits manager state; callers run +// finishAuthUpdate after releasing the gate. persistCooldownStates is also +// deferred by callers because it may re-enter manager operations. +// updateInternalLocked commits an update while the caller holds the gate. +// The second return value reports whether this call actually committed new +// state (versus returning the current snapshot for a discarded stale +// update); only real commits may run post-commit side effects. +func (m *Manager) updateInternalLocked(ctx context.Context, base, auth *Auth, mode updateAuthMode) (committedAuth *Auth, committed bool, persistCooldown bool, err error) { + if auth == nil || auth.ID == "" { + return nil, false, false, nil + } NormalizeCredentialMetadata(auth.Metadata) if errWeight := ValidateAuthWeight(auth); errWeight != nil { - return nil, fmt.Errorf("update auth: %w", errWeight) + return nil, false, false, fmt.Errorf("update auth: %w", errWeight) + } + isFileSnapshot := mode == updateModeReplace && auth.HasFileSnapshot() + if isFileSnapshot { + m.mu.RLock() + existing, ok := m.auths[auth.ID] + m.mu.RUnlock() + if !ok || existing == nil { + return nil, false, false, nil + } + // Validate file snapshot freshness under the per-auth mutation gate + // without holding m.mu across filesystem I/O. + if !auth.FileSnapshotMatchesDisk() { + m.mu.RLock() + current := m.auths[auth.ID].Clone() + m.mu.RUnlock() + return current, false, false, nil + } } m.mu.Lock() existing, ok := m.auths[auth.ID] if !ok || existing == nil { m.mu.Unlock() - return nil, nil + return nil, false, false, nil } if m.authEpochs == nil { m.authEpochs = make(map[string]uint64) @@ -189,7 +320,17 @@ func (m *Manager) updateInternal(ctx context.Context, base, auth *Auth, mode upd } if (mode == updateModeRefresh || mode == updateModePrepare) && base != nil && existing.RegistrationEpoch != base.RegistrationEpoch { m.mu.Unlock() - return nil, fmt.Errorf("update auth %s: stale registration epoch %d != %d", auth.ID, base.RegistrationEpoch, existing.RegistrationEpoch) + return nil, false, false, fmt.Errorf("update auth %s: stale registration epoch %d != %d", auth.ID, base.RegistrationEpoch, existing.RegistrationEpoch) + } + // Do not let an in-flight refresh overwrite credentials committed after its snapshot. + if mode == updateModeRefresh && base != nil && (existing.CredentialVersion != base.CredentialVersion || CredentialsChanged(base, existing)) { + if auth.CodexRefreshPatch == nil || !auth.CodexRefreshPatch.CanApplyTo(existing.Metadata) { + current := existing.Clone() + m.mu.Unlock() + // Discarded stale refresh: committed=false so callers do not + // run finishAuthUpdate for a state this call did not change. + return current, false, false, nil + } } if mode == updateModeRefresh { merged := MergeRefreshedAuth(base, existing, auth) @@ -206,7 +347,7 @@ func (m *Manager) updateInternal(ctx context.Context, base, auth *Auth, mode upd } if auth.RegistrationEpoch != 0 && auth.RegistrationEpoch < m.authEpochs[auth.ID] { m.mu.Unlock() - return nil, fmt.Errorf("update auth %s: stale registration epoch %d < %d", auth.ID, auth.RegistrationEpoch, m.authEpochs[auth.ID]) + return nil, false, false, fmt.Errorf("update auth %s: stale registration epoch %d < %d", auth.ID, auth.RegistrationEpoch, m.authEpochs[auth.ID]) } if auth.RegistrationEpoch >= m.authEpochs[auth.ID] { m.authEpochs[auth.ID] = auth.RegistrationEpoch @@ -217,6 +358,9 @@ func (m *Manager) updateInternal(ctx context.Context, base, auth *Auth, mode upd auth.Index = existing.Index auth.indexAssigned = existing.indexAssigned } + if isFileSnapshot && auth.CreatedAt.IsZero() { + auth.CreatedAt = existing.CreatedAt + } auth.Success = existing.Success auth.Failed = existing.Failed auth.recentRequests = existing.recentRequests @@ -225,83 +369,157 @@ func (m *Manager) updateInternal(ctx context.Context, base, auth *Auth, mode upd } else { auth.Generation++ } + existingVersion := existing.CredentialVersion + if existingVersion == 0 { + existingVersion = 1 + } + credChanged := CredentialsChanged(existing, auth) + if credChanged { + auth.CredentialVersion = max(auth.CredentialVersion, existingVersion) + 1 + } else { + auth.CredentialVersion = existingVersion + } cooldownStateChanged := false + now := time.Now() if !existing.Disabled && existing.Status != StatusDisabled && !auth.Disabled && auth.Status != StatusDisabled { - if len(auth.ModelStates) == 0 && len(existing.ModelStates) > 0 { + if isFileSnapshot { + auth.ModelStates = existing.ModelStates + auth.LastRefreshedAt = existing.LastRefreshedAt + if !credChanged { + auth.NextRefreshAfter = existing.NextRefreshAfter + auth.RefreshFailures = existing.RefreshFailures + } + } else if len(auth.ModelStates) == 0 && len(existing.ModelStates) > 0 { auth.ModelStates = existing.ModelStates } - // A reloaded credential file carries no quota observation. Keep the last - // one so soonest-reset ranking survives a token refresh. - if auth.Quota.ObservedAt.IsZero() { + if isFileSnapshot { + if (auth.Status == "" || auth.Status == StatusActive) && existing.Status != "" { + auth.Status = existing.Status + } + if auth.StatusMessage == "" { + auth.StatusMessage = existing.StatusMessage + } + auth.Unavailable = existing.Unavailable + auth.NextRetryAfter = existing.NextRetryAfter + if auth.LastError == nil && existing.LastError != nil { + auth.LastError = cloneError(existing.LastError) + } + observation := mergeQuotaObservation(auth.Quota, existing.Quota) + auth.Quota = existing.Quota.Clone() + auth.Quota.ObservedAt = observation.ObservedAt + auth.Quota.Signals = observation.Signals + auth.Quota.ReserveWindows = observation.ReserveWindows + auth.cooldownRecordedSeq = existing.cooldownRecordedSeq + } else if auth.Quota.ObservedAt.IsZero() { auth.Quota = mergeQuotaObservation(auth.Quota, existing.Quota) } - credChanged := CredentialsChanged(existing, auth) - if credChanged { - if hasUnauthorizedAuthFailure(existing) || (auth.LastError != nil && isUnauthorizedError(auth.LastError)) { - auth.Unavailable = false - auth.LastError = nil - auth.StatusMessage = "" - auth.Status = StatusActive - } - resumed := clearUnauthorizedModelStates(auth, time.Now()) - if len(resumed) > 0 { - cooldownStateChanged = true + if credChanged || mode == updateModeRefresh { + applyPreservedRejectedTokens(existing, auth, mode == updateModeRefresh) + if auth.RejectedAccessToken == "" || authAccessToken(auth) != auth.RejectedAccessToken { + auth.RefreshFailures = 0 + if credChanged && mode != updateModeRefresh && auth.NextRefreshAfter.Equal(existing.NextRefreshAfter) { + auth.NextRefreshAfter = time.Time{} + } + isAuthOrGrantFailure := hasUnauthorizedAuthFailure(existing) || + (auth.LastError != nil && (isUnauthorizedError(auth.LastError) || isInvalidGrantResultError(auth.LastError))) + if !auth.Quota.NextRecoverAt.After(now) && + (isAuthOrGrantFailure || (auth.Unavailable && !auth.NextRetryAfter.After(now))) { + auth.Unavailable = false + auth.NextRetryAfter = time.Time{} + auth.cooldownRecordedSeq = 0 + auth.LastError = nil + auth.StatusMessage = "" + auth.Status = StatusActive + cooldownStateChanged = true + } else if isAuthOrGrantFailure && auth.Quota.Exceeded && auth.Quota.NextRecoverAt.After(now) { + auth.LastError = nil + auth.StatusMessage = auth.Quota.Reason + cooldownStateChanged = true + } + resumed := clearUnauthorizedModelStates(auth, now) + if len(resumed) > 0 { + cooldownStateChanged = true + } } + } else { + applyPreservedRejectedTokens(existing, auth, false) } - if existing.Quota.Exceeded && existing.Quota.Reason == "credential_quota" && existing.Quota.NextRecoverAt.After(time.Now()) { + // A refresh that confirms quota restoration (issue #6690) releases the + // credential-wide cooldown early; otherwise the active credential + // quota below is preserved exactly as before. + devinQuotaRestored := mode == updateModeRefresh && applyDevinQuotaRestoration(auth, existing, now) + if devinQuotaRestored { + cooldownStateChanged = true + } else if existing.Quota.Exceeded && existing.Quota.Reason == "credential_quota" && existing.Quota.NextRecoverAt.After(now) { auth.Unavailable = existing.Unavailable auth.NextRetryAfter = existing.NextRetryAfter + // Preserve cooldown fields; keep the newest observation snapshot. + observation := mergeQuotaObservation(auth.Quota, existing.Quota) auth.Quota = existing.Quota + auth.Quota.ObservedAt = observation.ObservedAt + auth.Quota.Signals = observation.Signals + auth.Quota.ReserveWindows = observation.ReserveWindows + auth.cooldownRecordedSeq = existing.cooldownRecordedSeq if auth.Status == StatusActive { auth.Status = existing.Status } } } - now := time.Now() auth.UpdatedAt = now cooldownStateChanged = normalizeModelStates(auth) || cooldownStateChanged if m.cooldownDisabledForAuth(auth) || auth.Disabled || auth.Status == StatusDisabled { cooldownStateChanged = clearCooldownStateForAuth(auth, now) || cooldownStateChanged } + if existing.Metadata != nil { + if persistedHints, hasHints := existing.Metadata["antigravity_capability_hints"]; hasHints { + if auth.Metadata == nil { + auth.Metadata = make(map[string]any) + } + auth.Metadata["antigravity_capability_hints"] = persistedHints + } + } auth.EnsureIndex() - // A minted Meta key must reach the configured store before requests can use it. - // Keep the epoch check, save and installation together so a concurrent reload - // or removal cannot let an obsolete mint overwrite the credential on disk. + // Save before publication, including the transactional Meta mint path. + // Runtime-only changes made during I/O are merged below, never overwritten. + existingBeforeSave := existing.Clone() + refreshPatchBeforeSave := auth.CodexRefreshPatch persistMetaMint := (mode == updateModePrepare || mode == updateModeRefresh) && strings.EqualFold(strings.TrimSpace(auth.Provider), "meta") - if persistMetaMint { - if errPersist := m.persist(ctx, auth); errPersist != nil { + if errPersist := m.persistLocked(ctx, auth); errPersist != nil { + // A minted Meta key must reach the store before requests can use it. + if persistMetaMint { m.mu.Unlock() - return nil, fmt.Errorf("persist meta auth: %w", errPersist) + return nil, false, false, fmt.Errorf("persist meta auth: %w", errPersist) } + // Ordinary persistence failures remain non-fatal, but never silent. + log.WithFields(log.Fields{"auth_id": auth.ID, "credential": auth.ID, "provider": auth.Provider}).Warnf("failed to persist updated auth %s (%s): %v", auth.Provider, auth.ID, errPersist) + } + if m.authEpochs[auth.ID] != auth.RegistrationEpoch { + m.mu.Unlock() + return nil, false, false, fmt.Errorf("auth %s changed registration during persistence", auth.ID) + } + mergeAuthSaveDelta(auth, existingBeforeSave, existing, false) + if (auth.RejectedAccessToken == "" && existing.RejectedAccessToken != "") || (auth.RejectedSIWCAccessToken == "" && existing.RejectedSIWCAccessToken != "") { + probe := auth.Clone() + probe.CodexRefreshPatch = refreshPatchBeforeSave + applyPreservedRejectedTokens(existing, probe, mode == updateModeRefresh) + auth.RejectedAccessToken = probe.RejectedAccessToken + auth.RejectedSIWCAccessToken = probe.RejectedSIWCAccessToken + } + if existing.Generation >= auth.Generation { + auth.Generation = existing.Generation + 1 } + auth.ClearFileSnapshot() authClone := auth.Clone() m.auths[auth.ID] = authClone + m.notifyAuthChangeLocked(auth.ID) // Snapshot before unlocking: MarkResult mutates the published auth in place. var schedulerSnapshot *Auth if m.scheduler != nil { schedulerSnapshot = authClone.Clone() } m.mu.Unlock() - if !shouldDeferAPIKeyModelAliasRebuild(ctx) { - m.rebuildAPIKeyModelAliasFromRuntimeConfig() - } - if m.scheduler != nil { - m.scheduler.upsertAuth(schedulerSnapshot) - } - m.structuralEpoch.Add(1) - m.queueRefreshReschedule(auth.ID) - if !persistMetaMint { - // Persist failures stay non-fatal, but must not be silent: after a token - // refresh the rotated credentials only exist in memory until persisted. - if errPersist := m.persist(ctx, auth); errPersist != nil { - log.WithFields(log.Fields{"auth_id": auth.ID, "credential": auth.ID, "provider": auth.Provider}).Warnf("failed to persist updated auth %s (%s): %v", auth.Provider, auth.ID, errPersist) - } - } - m.hook.OnAuthUpdated(ctx, auth.Clone()) - if cooldownStateChanged { - m.persistCooldownStates(context.Background()) - } - return auth.Clone(), nil + _ = schedulerSnapshot + return auth.Clone(), true, cooldownStateChanged, nil } // Remove deletes an auth from runtime state without persisting. @@ -315,6 +533,8 @@ func (m *Manager) Remove(ctx context.Context, id string) { return } _ = ctx + releaseMutation := m.lockAuthMutation(id) + defer releaseMutation() m.mu.Lock() existing := m.auths[id] @@ -324,6 +544,7 @@ func (m *Manager) Remove(ctx context.Context, id string) { } provider := strings.TrimSpace(existing.Provider) delete(m.auths, id) + m.notifyAuthChangeLocked(id) if m.modelPoolOffsets != nil { delete(m.modelPoolOffsets, id) } @@ -345,6 +566,7 @@ func (m *Manager) Remove(ctx context.Context, id string) { m.authEpochs[id]++ tombstoneEpoch := m.authEpochs[id] m.mu.Unlock() + releaseMutation() if !shouldDeferAPIKeyModelAliasRebuild(ctx) { m.rebuildAPIKeyModelAliasFromRuntimeConfig() @@ -358,7 +580,9 @@ func (m *Manager) Remove(ctx context.Context, id string) { if provider != "" { if exec, ok := m.Executor(provider); ok && exec != nil { - if closer, okCloser := exec.(ExecutionSessionCloser); okCloser { + if authCloser, okAuthCloser := exec.(AuthExecutionSessionCloser); okAuthCloser && authCloser != nil { + authCloser.CloseAuthExecutionSessions(id) + } else if closer, okCloser := exec.(ExecutionSessionCloser); okCloser { closer.CloseExecutionSession(CloseAllExecutionSessionsID) } } @@ -378,16 +602,27 @@ func (m *Manager) invalidateSessionAffinity(authID string) { // Load resets manager state from the backing store. func (m *Manager) Load(ctx context.Context) error { - m.mu.Lock() - if m.store == nil { - m.mu.Unlock() + // Exclude complete Save/publication transactions, including new registrations + // not yet present in auths. Never wait for this barrier while holding m.mu. + lockCtx := ctx + if lockCtx == nil { + lockCtx = context.Background() + } + if errAcquire := m.authLoadGate.Acquire(lockCtx, math.MaxInt64); errAcquire != nil { + return errAcquire + } + defer m.authLoadGate.Release(math.MaxInt64) + m.mu.RLock() + store := m.store + m.mu.RUnlock() + if store == nil { return nil } - items, err := m.store.List(ctx) + items, err := store.List(ctx) if err != nil { - m.mu.Unlock() return err } + m.mu.Lock() previousAuths := m.auths m.auths = make(map[string]*Auth, len(items)) if m.authEpochs == nil { @@ -405,6 +640,15 @@ func (m *Manager) Load(ctx context.Context) error { m.authEpochs[auth.ID] = max(m.authEpochs[auth.ID], auth.RegistrationEpoch) + 1 auth.RegistrationEpoch = m.authEpochs[auth.ID] auth.Generation = 1 + if prev, exists := previousAuths[auth.ID]; exists && prev != nil { + auth.CredentialVersion = max(auth.CredentialVersion, prev.CredentialVersion) + if CredentialsChanged(prev, auth) { + auth.CredentialVersion++ + } + } + if auth.CredentialVersion == 0 { + auth.CredentialVersion = 1 + } m.auths[auth.ID] = auth.Clone() } @@ -428,6 +672,9 @@ func (m *Manager) Load(ctx context.Context) error { cfg = &internalconfig.Config{} } m.rebuildAPIKeyModelAliasLocked(cfg) + for id := range m.authChangeWatchers { + m.notifyAuthChangeLocked(id) + } m.mu.Unlock() if m.scheduler != nil { @@ -447,7 +694,10 @@ type authPersistLock struct { } func (m *Manager) persist(ctx context.Context, auth *Auth) error { - if m.store == nil || auth == nil { + m.mu.RLock() + store := m.store + m.mu.RUnlock() + if store == nil || auth == nil { return nil } if errWeight := ValidateAuthWeight(auth); errWeight != nil { @@ -482,13 +732,23 @@ func (m *Manager) persist(ctx context.Context, auth *Auth) error { if shouldSkipPersist(ctx) { return nil } - _, err := m.store.Save(ctx, auth) + releaseFg, errWait := m.beginForegroundStoreSave(ctx) + if errWait != nil { + return errWait + } + defer releaseFg() + _, err := store.Save(ctx, auth) return err } if shouldSkipPersist(ctx) { return nil } - _, err := m.store.Save(ctx, auth) + releaseFg, errWait := m.beginForegroundStoreSave(ctx) + if errWait != nil { + return errWait + } + defer releaseFg() + _, err := store.Save(ctx, auth) return err } diff --git a/sdk/cliproxy/auth/conductor_persistence.go b/sdk/cliproxy/auth/conductor_persistence.go new file mode 100644 index 000000000..a845f21f1 --- /dev/null +++ b/sdk/cliproxy/auth/conductor_persistence.go @@ -0,0 +1,373 @@ +package auth + +import ( + "context" + "reflect" + "sync" + "sync/atomic" + + log "github.com/sirupsen/logrus" +) + +// lockAuthMutation acquires the reload barrier, then the credential gate, before +// m.mu. Keep both until publication completes. Load never takes credential or +// refresh locks. Store callbacks may read the manager, but must not mutate it. +func (m *Manager) lockAuthMutation(id string) func() { + release, _ := m.lockAuthMutationContext(context.Background(), id) + return release +} + +type metadataPersistPreemptState struct { + mu sync.Mutex + waiters int + fgActive int + fgIdle chan struct{} + cancel context.CancelFunc + done chan struct{} +} + +type metadataPersistContextKey struct{} + +func withMetadataPersistContext(ctx context.Context) context.Context { + return context.WithValue(ctx, metadataPersistContextKey{}, true) +} + +func isMetadataPersistContext(ctx context.Context) bool { + if ctx == nil { + return false + } + v, _ := ctx.Value(metadataPersistContextKey{}).(bool) + return v +} + +func (m *Manager) metadataPersistPreempt(id string) *metadataPersistPreemptState { + val, _ := m.metadataPersistCancels.LoadOrStore(id, &metadataPersistPreemptState{}) + return val.(*metadataPersistPreemptState) +} + +func (m *Manager) globalMetadataSaveState() *metadataPersistPreemptState { + val, _ := m.metadataSaveGate.LoadOrStore("state", &metadataPersistPreemptState{}) + return val.(*metadataPersistPreemptState) +} + +func (m *Manager) acquireMetadataSaveGate(ctx context.Context, cancel context.CancelFunc) (func(), bool) { + if ctx == nil { + ctx = context.Background() + } + val, _ := m.metadataSaveGate.LoadOrStore("gate", make(chan struct{}, 1)) + gate := val.(chan struct{}) + select { + case gate <- struct{}{}: + if ctx.Err() != nil { + <-gate + return nil, false + } + case <-ctx.Done(): + return nil, false + } + + gState := m.globalMetadataSaveState() + for { + gState.mu.Lock() + if ctx.Err() != nil { + gState.mu.Unlock() + <-gate + return nil, false + } + if gState.fgActive == 0 { + done := make(chan struct{}) + gState.done = done + gState.cancel = cancel + gState.mu.Unlock() + var once sync.Once + release := func() { + once.Do(func() { + gState.mu.Lock() + gState.cancel = nil + gState.done = nil + gState.mu.Unlock() + close(done) + <-gate + }) + } + return release, true + } + fgIdle := gState.fgIdle + gState.mu.Unlock() + select { + case <-fgIdle: + case <-ctx.Done(): + <-gate + return nil, false + } + } +} + +func (m *Manager) beginForegroundStoreSave(ctx context.Context) (func(), error) { + if m == nil || isMetadataPersistContext(ctx) { + return func() {}, nil + } + if ctx == nil { + ctx = context.Background() + } + gState := m.globalMetadataSaveState() + gState.mu.Lock() + if gState.fgActive == 0 { + gState.fgIdle = make(chan struct{}) + } + gState.fgActive++ + done := gState.done + cancel := gState.cancel + gState.mu.Unlock() + + var once sync.Once + release := func() { + once.Do(func() { + gState.mu.Lock() + gState.fgActive-- + if gState.fgActive == 0 && gState.fgIdle != nil { + close(gState.fgIdle) + gState.fgIdle = nil + } + gState.mu.Unlock() + }) + } + if cancel != nil { + cancel() + } + if done != nil { + select { + case <-done: + if errCtx := ctx.Err(); errCtx != nil { + release() + return nil, errCtx + } + case <-ctx.Done(): + release() + return nil, ctx.Err() + } + } else if errCtx := ctx.Err(); errCtx != nil { + release() + return nil, errCtx + } + return release, nil +} + +// lockAuthMutationContext lets fallible mutations cancel either lock wait. +func (m *Manager) lockAuthMutationContext(ctx context.Context, id string) (func(), error) { + return m.lockAuthMutationWithPersistCancel(ctx, id, nil) +} + +func (m *Manager) lockAuthMutationWithPersistCancel(ctx context.Context, id string, persistCancel context.CancelFunc) (func(), error) { + if ctx == nil { + ctx = context.Background() + } + if errAcquire := m.authLoadGate.Acquire(ctx, 1); errAcquire != nil { + return nil, errAcquire + } + value, _ := m.authMutationLocks.LoadOrStore(id, make(chan struct{}, 1)) + gate := value.(chan struct{}) + state := m.metadataPersistPreempt(id) + if persistCancel == nil { + state.mu.Lock() + state.waiters++ + cancel := state.cancel + state.mu.Unlock() + if cancel != nil { + cancel() + } + defer func() { + state.mu.Lock() + state.waiters-- + state.mu.Unlock() + }() + } + select { + case gate <- struct{}{}: + // Cancellation may race with the gate becoming available. + if errContext := ctx.Err(); errContext != nil { + <-gate + m.authLoadGate.Release(1) + return nil, errContext + } + case <-ctx.Done(): + m.authLoadGate.Release(1) + return nil, ctx.Err() + } + if persistCancel != nil { + state.mu.Lock() + if state.waiters > 0 { + state.mu.Unlock() + persistCancel() + } else { + state.cancel = persistCancel + state.mu.Unlock() + } + } + var once sync.Once + return func() { + once.Do(func() { + if persistCancel != nil { + state.mu.Lock() + state.cancel = nil + state.mu.Unlock() + } + <-gate + m.authLoadGate.Release(1) + }) + }, nil +} + +// persistLocked requires m.mu and the credential mutation gate. Never pass a +// published object to Store.Save: readers and runtime-only writers remain live +// during I/O. Only the store's delta is merged back, preserving intervening edits. +func (m *Manager) persistLocked(ctx context.Context, auth *Auth) error { + before := auth.Clone() + saved := auth.Clone() + m.mu.Unlock() + err := m.persist(ctx, saved) + m.mu.Lock() + mergeAuthSaveDelta(auth, before, saved, true) + return err +} + +// mergeAuthSaveDelta handles all exported value fields, not a store-specific +// allowlist. Map entries are independent fields; nested values are atomic. +// Identity/version fields belong exclusively to the manager. Store-owned shared +// objects must be replaced rather than mutated in place (see Store.Save). +func mergeAuthSaveDelta(target, before, after *Auth, preserveConcurrent bool) { + dst := reflect.ValueOf(target).Elem() + base := reflect.ValueOf(before).Elem() + next := reflect.ValueOf(after).Elem() + for i := 0; i < dst.NumField(); i++ { + field := dst.Type().Field(i) + if field.PkgPath != "" || field.Name == "ID" || field.Name == "RegistrationEpoch" || field.Name == "CredentialVersion" || field.Name == "Generation" || field.Name == "RejectedAccessToken" || field.Name == "RejectedSIWCAccessToken" { + continue + } + d, b, n := dst.Field(i), base.Field(i), next.Field(i) + if reflect.DeepEqual(b.Interface(), n.Interface()) { + continue + } + if reflect.DeepEqual(d.Interface(), b.Interface()) { + d.Set(n) + continue + } + if d.Kind() == reflect.Map { + keys := make(map[any]reflect.Value) + for _, key := range b.MapKeys() { + keys[key.Interface()] = key + } + for _, key := range n.MapKeys() { + keys[key.Interface()] = key + } + for _, key := range keys { + oldValue, newValue := b.MapIndex(key), n.MapIndex(key) + if equalSaveValue(oldValue, newValue) || (preserveConcurrent && !equalSaveValue(d.MapIndex(key), oldValue)) { + continue + } + if d.IsNil() { + d.Set(reflect.MakeMap(d.Type())) + } + d.SetMapIndex(key, newValue) + } + continue + } + if !preserveConcurrent { + d.Set(n) + } + } +} + +func equalSaveValue(a, b reflect.Value) bool { + if !a.IsValid() || !b.IsValid() { + return a.IsValid() == b.IsValid() + } + return reflect.DeepEqual(a.Interface(), b.Interface()) +} + +// PersistMetadataEntry updates a single metadata entry on the runtime auth for +// the given registration epoch and persists it to the backing store without +// advancing Generation or firing credential-lifecycle hooks. If a foreground +// credential mutation contends for the same auth while store I/O is in flight, +// the metadata save is preempted so foreground operations do not stall. +func (m *Manager) PersistMetadataEntry(ctx context.Context, id string, expectedEpoch uint64, key string, value any, validate func(existing *Auth) bool) (bool, bool) { + if m == nil || id == "" || key == "" { + return false, false + } + ctx = WithoutSkipPersist(ctx) + if errContext := ctx.Err(); errContext != nil { + return false, false + } + m.mu.RLock() + existing := m.auths[id] + if existing == nil || (expectedEpoch != 0 && existing.RegistrationEpoch != expectedEpoch) || (validate != nil && !validate(existing)) { + m.mu.RUnlock() + return false, false + } + m.mu.RUnlock() + + persistCtx, cancelPersist := context.WithCancel(withMetadataPersistContext(ctx)) + defer cancelPersist() + var preempted atomic.Bool + preemptFn := func() { + preempted.Store(true) + cancelPersist() + } + releaseMutation, errMutation := m.lockAuthMutationWithPersistCancel(persistCtx, id, preemptFn) + if errMutation != nil { + return false, preempted.Load() && ctx.Err() == nil + } + defer releaseMutation() + + releaseSaveGate, okGate := m.acquireMetadataSaveGate(persistCtx, preemptFn) + if !okGate { + return false, preempted.Load() && ctx.Err() == nil + } + defer releaseSaveGate() + + m.mu.Lock() + existing = m.auths[id] + if existing == nil || (expectedEpoch != 0 && existing.RegistrationEpoch != expectedEpoch) || (validate != nil && !validate(existing)) { + m.mu.Unlock() + return false, false + } + if preempted.Load() { + m.mu.Unlock() + return false, ctx.Err() == nil + } + existingBeforeSave := existing.Clone() + updated := existing.Clone() + if updated.Metadata == nil { + updated.Metadata = make(map[string]any) + } + updated.Metadata[key] = value + errPersist := m.persistLocked(persistCtx, updated) + if errPersist != nil { + if !preempted.Load() { + log.WithFields(log.Fields{"auth_id": updated.ID, "credential": updated.ID, "provider": updated.Provider}).Warnf("failed to persist metadata for auth %s (%s): %v", updated.Provider, updated.ID, errPersist) + } + m.mu.Unlock() + return false, preempted.Load() && ctx.Err() == nil + } + current := m.auths[id] + if current == nil || current.RegistrationEpoch != updated.RegistrationEpoch { + m.mu.Unlock() + return false, false + } + if validate != nil && !validate(current) { + if current.Metadata != nil { + invalidated := current.Clone() + delete(invalidated.Metadata, key) + m.auths[id] = invalidated + } + m.mu.Unlock() + return false, false + } + mergeAuthSaveDelta(updated, existingBeforeSave, current, false) + if current.Generation > updated.Generation { + updated.Generation = current.Generation + } + m.auths[id] = updated.Clone() + m.mu.Unlock() + return true, false +} diff --git a/sdk/cliproxy/auth/conductor_refresh.go b/sdk/cliproxy/auth/conductor_refresh.go index a92e6b5c2..e21998ce5 100644 --- a/sdk/cliproxy/auth/conductor_refresh.go +++ b/sdk/cliproxy/auth/conductor_refresh.go @@ -127,6 +127,19 @@ func (m *Manager) shouldRefresh(a *Auth, now time.Time) bool { if evaluator, ok := a.Runtime.(RefreshEvaluator); ok && evaluator != nil { return evaluator.ShouldRefresh(now, a) } + if next, due := devinQuotaRevalidateAt(a, now); due { + return !next.After(now) + } + + if _, hasSIWC, needSIWC := CodexSIWCNeedsRefresh(a, now); hasSIWC { + if needSIWC { + return true + } + if _, hasOuter, needOuter := CodexOuterNeedsRefresh(a, now); hasOuter { + return needOuter + } + return false + } lastRefresh := a.LastRefreshedAt if lastRefresh.IsZero() { @@ -164,7 +177,7 @@ func (m *Manager) shouldRefresh(a *Auth, now time.Time) bool { return false } if hasExpiry && !expiry.IsZero() { - return time.Until(expiry) <= *lead + return expiry.Sub(now) <= *lead } if !lastRefresh.IsZero() { return now.Sub(lastRefresh) >= *lead @@ -385,13 +398,35 @@ type authRefreshLock struct { mu sync.Mutex } -func authAccessToken(auth *Auth) string { +func authOuterAccessToken(auth *Auth) string { if token := authMetadataString(auth, "access_token"); token != "" { return token } return authMetadataString(auth, "accessToken") } +func authAccessToken(auth *Auth) string { + if IsCodexUseSIWC(auth) { + return authSIWCTokenField(auth, "access_token") + } + return authOuterAccessToken(auth) +} + +func authAccessTokenForRequest(ctx context.Context, auth *Auth, clientHeaders ...http.Header) string { + if IsCodexUseSIWCForRequest(ctx, auth, clientHeaders...) { + return authSIWCTokenField(auth, "access_token") + } + return authOuterAccessToken(auth) +} + +func authMatchesTokenPartition(auth *Auth, token string) bool { + token = strings.TrimSpace(token) + if auth == nil || token == "" { + return false + } + return authAccessToken(auth) == token || authOuterAccessToken(auth) == token +} + func authRefreshToken(auth *Auth) string { if token := authMetadataString(auth, "refresh_token"); token != "" { return token @@ -406,6 +441,9 @@ func authHasRefreshCredential(auth *Auth) bool { if authMetadataString(auth, "refreshToken") != "" { return true } + if _, hasSIWC, _ := CodexSIWCNeedsRefresh(auth, time.Now()); hasSIWC { + return true + } // Meta exchanges its device token for a replacement API key after a 401. return auth != nil && strings.EqualFold(strings.TrimSpace(auth.Provider), "meta") && (authMetadataString(auth, "dca_token") != "" || strings.TrimSpace(auth.Attributes["dca_token"]) != "") @@ -417,7 +455,7 @@ func CredentialsChanged(existing, incoming *Auth) bool { if existing == nil || incoming == nil { return false } - if authAccessToken(existing) != authAccessToken(incoming) { + if authAccessToken(existing) != authAccessToken(incoming) || authOuterAccessToken(existing) != authOuterAccessToken(incoming) { return true } if authRefreshToken(existing) != authRefreshToken(incoming) { @@ -451,6 +489,125 @@ func CredentialsChanged(existing, incoming *Auth) bool { if existingKey != incomingKey { return true } + if authSIWCTokenField(existing, "access_token") != authSIWCTokenField(incoming, "access_token") || + authSIWCTokenField(existing, "refresh_token") != authSIWCTokenField(incoming, "refresh_token") || + authSIWCTokenField(existing, "id_token") != authSIWCTokenField(incoming, "id_token") || + authSIWCTokenField(existing, "client_id") != authSIWCTokenField(incoming, "client_id") { + return true + } + existingUseSIWC, existingHasUseSIWC := parseBoolAny(existing.Metadata["use_siwc"]) + incomingUseSIWC, incomingHasUseSIWC := parseBoolAny(incoming.Metadata["use_siwc"]) + if existingHasUseSIWC != incomingHasUseSIWC || existingUseSIWC != incomingUseSIWC { + return true + } + return false +} + +func authSIWCTokenField(auth *Auth, key string) string { + if auth == nil || auth.Metadata == nil { + return "" + } + siwcMap, ok := auth.Metadata["siwc"].(map[string]any) + if !ok || len(siwcMap) == 0 { + return "" + } + val, _ := siwcMap[key].(string) + return strings.TrimSpace(val) +} + +func effectiveRejectedSIWCToken(a *Auth) string { + if a == nil { + return "" + } + siwcTok := authSIWCTokenField(a, "access_token") + if siwcTok == "" { + return "" + } + if a.RejectedSIWCAccessToken != "" && a.RejectedSIWCAccessToken == siwcTok { + return siwcTok + } + if a.RejectedAccessToken != "" && a.RejectedAccessToken == siwcTok { + return siwcTok + } + return "" +} + +func shouldPreserveRejectedSIWCAccessToken(existing, updated *Auth) (string, bool) { + rejectedSIWC := effectiveRejectedSIWCToken(existing) + if rejectedSIWC == "" || updated == nil { + return "", false + } + if authSIWCTokenField(updated, "access_token") == rejectedSIWC { + return rejectedSIWC, true + } + return "", false +} + +func applyPreservedRejectedTokens(existing, updated *Auth, isRefresh bool) { + if updated == nil { + return + } + if rejectedSIWC, ok := shouldPreserveRejectedSIWCAccessToken(existing, updated); ok { + updated.RejectedSIWCAccessToken = rejectedSIWC + } else { + updated.RejectedSIWCAccessToken = "" + } + if shouldPreserveRejectedAccessToken(existing, updated, isRefresh) { + updated.RejectedAccessToken = existing.RejectedAccessToken + } else if IsCodexUseSIWC(updated) && updated.RejectedSIWCAccessToken != "" { + updated.RejectedAccessToken = updated.RejectedSIWCAccessToken + } else { + updated.RejectedAccessToken = "" + } +} + +func recordRejectedTokenOnAuth(auth *Auth, failedToken string) bool { + if auth == nil { + return false + } + failedToken = strings.TrimSpace(failedToken) + if failedToken == "" { + return false + } + changed := false + if siwcTok := authSIWCTokenField(auth, "access_token"); siwcTok != "" && siwcTok == failedToken { + if auth.RejectedSIWCAccessToken != failedToken { + auth.RejectedSIWCAccessToken = failedToken + changed = true + } + } + if (authAccessToken(auth) == failedToken || authOuterAccessToken(auth) == failedToken) && auth.RejectedAccessToken != failedToken { + auth.RejectedAccessToken = failedToken + changed = true + } + return changed +} + +func shouldPreserveRejectedAccessToken(existing, updated *Auth, isRefresh bool) bool { + if existing == nil || updated == nil { + return false + } + rejected := strings.TrimSpace(existing.RejectedAccessToken) + if rejected == "" { + return false + } + existingSIWC := authSIWCTokenField(existing, "access_token") + if existingSIWC != "" && existingSIWC == rejected { + return authSIWCTokenField(updated, "access_token") == rejected + } + existingOuter := authOuterAccessToken(existing) + if existingOuter != "" && existingOuter == rejected { + if authOuterAccessToken(updated) != rejected { + return false + } + if !isRefresh { + return true + } + if updated.CodexRefreshPatch != nil { + return !updated.CodexRefreshPatch.RefreshedOuter() + } + return false + } return false } @@ -459,31 +616,184 @@ func ClearUnauthorizedModelStates(auth *Auth, now time.Time) []string { return clearUnauthorizedModelStates(auth, now) } +// devinQuotaRevalidateInterval bounds how often a quota-exhausted Devin +// credential re-queries upstream status for early restoration (issue #6690). +const devinQuotaRevalidateInterval = time.Hour + +// devinQuotaRevalidateAt returns the next quota revalidation instant for a +// Devin auth under an active credential quota. Devin session tokens never +// expire, so without this the auto-refresh loop would never re-observe quota +// and early restoration could never be detected. It is scoped strictly to +// the exhausted state; healthy credentials are unaffected. +func devinQuotaRevalidateAt(a *Auth, now time.Time) (time.Time, bool) { + if a == nil || !strings.EqualFold(strings.TrimSpace(a.Provider), "devin") { + return time.Time{}, false + } + if !a.Quota.Exceeded || a.Quota.Reason != "credential_quota" || !a.Quota.NextRecoverAt.After(now) { + return time.Time{}, false + } + since := a.LastRefreshedAt + if since.IsZero() { + since = a.Quota.ObservedAt + } + if since.IsZero() { + return now, true + } + next := since.Add(devinQuotaRevalidateInterval) + if next.After(a.Quota.NextRecoverAt) { + next = a.Quota.NextRecoverAt + } + return next, true +} + +// devinQuotaObservationRestored reports whether the refreshed auth carries a +// newer Devin quota observation confirming daily and weekly quota are +// available again while existing is under an active credential quota. The +// observation must be strictly newer than both the stored snapshot and the +// moment the exhausted state was last written, so a stale or replayed +// refresh can never resurrect an exhausted credential. +func devinQuotaObservationRestored(auth, existing *Auth, now time.Time) bool { + if auth == nil || existing == nil { + return false + } + if !strings.EqualFold(strings.TrimSpace(existing.Provider), "devin") { + return false + } + if !existing.Quota.Exceeded || existing.Quota.Reason != "credential_quota" || !existing.Quota.NextRecoverAt.After(now) { + return false + } + observedAt := auth.Quota.ObservedAt + if observedAt.IsZero() || observedAt.After(now) { + return false + } + if !observedAt.After(existing.Quota.ObservedAt) || !existing.UpdatedAt.IsZero() && observedAt.Before(existing.UpdatedAt) { + return false + } + daily, okDaily := devinRemainingPercent(auth.Quota.Signals["daily_quota_remaining_percent"]) + weekly, okWeekly := devinRemainingPercent(auth.Quota.Signals["weekly_quota_remaining_percent"]) + return okDaily && okWeekly && daily > 0 && weekly > 0 +} + +// applyDevinQuotaRestoration releases a Devin credential-wide quota cooldown +// early when the refreshed auth carries an authoritative, newer observation +// showing both usage quotas available again (issue #6690). It never touches +// unauthorized failures, non-quota failure states, or model states whose +// block is not the auth-level credential quota propagation; unrelated +// cooldowns survive via aggregated availability recomputation. +func applyDevinQuotaRestoration(auth, existing *Auth, now time.Time) bool { + if !devinQuotaObservationRestored(auth, existing, now) { + return false + } + if hasUnauthorizedAuthFailure(existing) { + return false + } + + applyCooldownFields(&auth.Quota, QuotaState{}) + for _, state := range auth.ModelStates { + if state == nil { + continue + } + reason := strings.TrimSpace(state.Quota.Reason) + wrapped := state.credentialQuotaPropagated || reason == "credential_quota" + if !state.Quota.Exceeded || (reason != "quota" && !wrapped) { + continue + } + // A wrapped sibling carries its independent layer in the shadow. + // Reinstate that exact block (if still live) instead of erasing it; + // the reporting model ("quota" unwrapped) and pure propagation + // targets (no shadow) fully release. Deadline or status heuristics + // cannot recover this provenance reliably. + if shadow := state.credentialQuotaShadow; shadow != nil && isModelStateActiveCooldown(shadow, now) { + *state = *shadow.cloneShallow() + state.UpdatedAt = now + continue + } + resetModelState(state, now) + } + auth.Status = StatusActive + auth.StatusMessage = "" + auth.LastError = nil + auth.Unavailable = false + auth.NextRetryAfter = time.Time{} + if len(auth.ModelStates) > 0 { + updateAggregatedAvailability(auth, now) + } + return true +} + +// isDevinQuotaExhaustionMessage reports whether text is Devin's explicit +// daily/weekly usage-quota exhaustion message. It mirrors the executor's +// narrow classifier and is used only to tell credential-derived legacy +// "quota" states apart from genuine model-scoped rate limits. +func isDevinQuotaExhaustionMessage(text string) bool { + lower := strings.ToLower(text) + if !strings.Contains(lower, "quota") || !strings.Contains(lower, "exhausted") { + return false + } + return strings.Contains(lower, "weekly") || strings.Contains(lower, "daily") +} + +func devinRecordErrorText(record CooldownStateRecord) string { + if record.LastError != nil && strings.TrimSpace(record.LastError.Message) != "" { + return record.LastError.Message + } + return "" +} + +func devinStateErrorText(state *ModelState) string { + if state == nil { + return "" + } + if state.LastError != nil && strings.TrimSpace(state.LastError.Message) != "" { + return state.LastError.Message + } + return state.StatusMessage +} + +// devinRemainingPercent parses a Devin quota signal such as "55%". +func devinRemainingPercent(raw string) (int64, bool) { + value, errParse := strconv.ParseInt(strings.TrimSpace(strings.TrimSuffix(strings.TrimSpace(raw), "%")), 10, 64) + if errParse != nil { + return 0, false + } + return value, true +} + func clearUnauthorizedModelStates(auth *Auth, now time.Time) []string { if auth == nil || len(auth.ModelStates) == 0 { return nil } var resumed []string + anyChanged := false for model, state := range auth.ModelStates { if state == nil { continue } isUnauth := false if state.LastError != nil { - if state.LastError.StatusCode() == http.StatusUnauthorized || strings.EqualFold(state.LastError.Code, "unauthorized") || isUnauthorizedError(state.LastError) { + if state.LastError.StatusCode() == http.StatusUnauthorized || strings.EqualFold(state.LastError.Code, "unauthorized") || isUnauthorizedError(state.LastError) || isInvalidGrantResultError(state.LastError) { isUnauth = true } } - if !isUnauth && strings.Contains(strings.ToLower(state.StatusMessage), "unauthorized") { + if !isUnauth && (strings.Contains(strings.ToLower(state.StatusMessage), "unauthorized") || isInvalidGrantErrorMessage(state.StatusMessage)) { isUnauth = true } if !isUnauth { continue } + anyChanged = true + if state.Quota.Exceeded && (state.Quota.NextRecoverAt.IsZero() || state.Quota.NextRecoverAt.After(now)) { + state.LastError = nil + if strings.Contains(strings.ToLower(state.StatusMessage), "unauthorized") || isInvalidGrantErrorMessage(state.StatusMessage) { + state.StatusMessage = state.Quota.Reason + } + state.UpdatedAt = now + continue + } resetModelState(state, now) resumed = append(resumed, model) } - if len(resumed) > 0 { + if anyChanged { updateAggregatedAvailability(auth, now) } return resumed @@ -491,7 +801,7 @@ func clearUnauthorizedModelStates(auth *Auth, now time.Time) []string { // RefreshHomeSelectionAfterUnauthorized only reuses a newer snapshot already // installed by Home. It never refreshes or mutates Home-owned credentials. -func (m *Manager) RefreshHomeSelectionAfterUnauthorized(_ context.Context, selection *HomeDispatchSelection, failedAuth *Auth) (*Auth, bool, error) { +func (m *Manager) RefreshHomeSelectionAfterUnauthorized(ctx context.Context, selection *HomeDispatchSelection, failedAuth *Auth) (*Auth, bool, error) { if m == nil || selection == nil { return nil, false, nil } @@ -500,8 +810,8 @@ func (m *Manager) RefreshHomeSelectionAfterUnauthorized(_ context.Context, selec failedAuth = current } if current != nil && failedAuth != nil && current.ID == failedAuth.ID { - currentToken := authAccessToken(current) - failedToken := authAccessToken(failedAuth) + currentToken := authAccessTokenForRequest(ctx, current) + failedToken := authAccessTokenForRequest(ctx, failedAuth) if currentToken != "" && failedToken != "" && currentToken != failedToken { return current, true, nil } @@ -511,27 +821,35 @@ func (m *Manager) RefreshHomeSelectionAfterUnauthorized(_ context.Context, selec // tryRefreshAfterUnauthorized refreshes local OAuth credentials once after a // 401 so the current auth can be retried before fallback/suspend. -func (m *Manager) tryRefreshAfterUnauthorized(ctx context.Context, auth *Auth, execErr error, alreadyTried bool) (*Auth, bool) { - if m == nil || auth == nil || alreadyTried || execErr == nil { +func (m *Manager) tryRefreshAfterUnauthorized(ctx context.Context, auth *Auth, execErr error, alreadyTried bool, clientHeaders ...http.Header) (*Auth, bool) { + if m == nil || auth == nil || execErr == nil { return auth, false } // Request-scoped failures describe this request, not stale credentials. // Refreshing would turn a direct error response into an implicit retry. - if isRequestScopedError(execErr) { + if isRequestScopedError(execErr) || !isUnauthorizedError(execErr) { return auth, false } - if !isUnauthorizedError(execErr) || !authHasRefreshCredential(auth) { + // A 401 matched by a configured request-scoped rule is the rule's to handle: no rejection, no refresh. + if _, ruleMatched := matchRequestScopedErrorAction(auth, execErr, m.runtimeConfigSnapshot()); ruleMatched { return auth, false } - if hasUnauthorizedAuthFailure(auth) { + failedToken := authAccessTokenForRequest(ctx, auth, clientHeaders...) + if failedToken != "" { + m.markRejectedAccessToken(auth.ID, failedToken) + } + if alreadyTried || !authHasRefreshCredential(auth) || hasUnauthorizedAuthFailure(auth) { return auth, false } log.Debugf("unauthorized response for %s (%s), refreshing credentials before fallback", auth.Provider, auth.ID) - refreshed, errRefresh := m.refreshAuthForRequest(ctx, auth.ID, authAccessToken(auth)) + refreshed, errRefresh := m.refreshAuthForRequest(ctx, auth.ID, failedToken) if errRefresh != nil || refreshed == nil { log.Debugf("credential refresh before fallback failed for %s (%s): %v", auth.Provider, auth.ID, errRefresh) return auth, false } + if failedToken != "" && authAccessTokenForRequest(ctx, refreshed, clientHeaders...) == failedToken && (refreshed.RejectedAccessToken == failedToken || refreshed.RejectedSIWCAccessToken == failedToken) { + return auth, false + } return refreshed, true } @@ -563,14 +881,16 @@ func (m *Manager) refreshAuthForRequest(ctx context.Context, id, failedAccessTok type forceRefreshContextKey struct{} -func withForceRefresh(ctx context.Context) context.Context { +// WithForceRefresh marks ctx as an explicit forced credential refresh. +func WithForceRefresh(ctx context.Context) context.Context { if ctx == nil { ctx = context.Background() } return context.WithValue(ctx, forceRefreshContextKey{}, true) } -func isForceRefreshContext(ctx context.Context) bool { +// IsForceRefreshContext reports whether ctx carries an explicit forced credential refresh marker. +func IsForceRefreshContext(ctx context.Context) bool { if ctx == nil { return false } @@ -578,6 +898,32 @@ func isForceRefreshContext(ctx context.Context) bool { return ok && v } +func withForceRefresh(ctx context.Context) context.Context { + return WithForceRefresh(ctx) +} + +func isForceRefreshContext(ctx context.Context) bool { + return IsForceRefreshContext(ctx) +} + +func (m *Manager) markRejectedAccessToken(id, failedAccessToken string) { + if m == nil || id == "" || failedAccessToken == "" { + return + } + releaseMutation := m.lockAuthMutation(id) + defer releaseMutation() + m.mu.Lock() + if current := m.auths[id]; current != nil && recordRejectedTokenOnAuth(current, failedAccessToken) { + current.Generation++ + current.UpdatedAt = time.Now() + m.auths[id] = current + if m.scheduler != nil { + m.scheduler.upsertAuth(current.Clone()) + } + } + m.mu.Unlock() +} + func (m *Manager) refreshAuthForRequestAtEpoch(ctx context.Context, id, failedAccessToken string, registrationEpoch uint64) (*Auth, error) { if m == nil { return nil, errors.New("auth manager is nil") @@ -594,6 +940,9 @@ func (m *Manager) refreshAuthForRequestAtEpoch(ctx context.Context, id, failedAc if id == "" { return nil, errors.New("auth id is empty") } + if failedAccessToken != "" { + m.markRejectedAccessToken(id, failedAccessToken) + } lockValue, _ := m.refreshLocks.LoadOrStore(id, &authRefreshLock{}) lock, _ := lockValue.(*authRefreshLock) @@ -628,19 +977,31 @@ func (m *Manager) refreshAuthForRequestAtEpoch(ctx context.Context, id, failedAc // Another request may already have refreshed this credential. if failedAccessToken != "" { - if currentToken := authAccessToken(auth); currentToken != "" && currentToken != failedAccessToken { + if !authMatchesTokenPartition(auth, failedAccessToken) && (authAccessToken(auth) != "" || authOuterAccessToken(auth) != "") { return auth.Clone(), nil } } base := auth.Clone() updated, err := exec.Refresh(ctx, base.Clone()) + now := time.Now() if err != nil && errors.Is(err, context.Canceled) { log.Debugf("refresh canceled for %s, %s", auth.Provider, auth.ID) + m.mu.Lock() + if current := m.auths[id]; current != nil { + if current.NextRefreshAfter.IsZero() || current.NextRefreshAfter.Before(now) { + current.NextRefreshAfter = now.Add(time.Second) + } + m.auths[id] = current + if m.scheduler != nil { + m.scheduler.upsertAuth(current.Clone()) + } + } + m.mu.Unlock() + m.queueRefreshReschedule(id) return nil, err } log.Debugf("refreshed %s, %s, %v", auth.Provider, auth.ID, err) - now := time.Now() if err != nil { unauthorized := isUnauthorizedError(err) invalidGrant := isInvalidGrantError(err) @@ -685,7 +1046,7 @@ func (m *Manager) refreshAuthForRequestAtEpoch(ctx context.Context, id, failedAc hasValidAccessToken := current.HasValidAccessToken(now) // failedAccessToken is set only when upstream rejected this exact access // token. Its expiry time no longer proves it is usable. - accessTokenRejected := failedAccessToken != "" && authAccessToken(current) == failedAccessToken + accessTokenRejected := failedAccessToken != "" && authMatchesTokenPartition(current, failedAccessToken) if isDisabled && invalidGrant { current.Unavailable = true current.Status = StatusDisabled @@ -713,7 +1074,7 @@ func (m *Manager) refreshAuthForRequestAtEpoch(ctx context.Context, id, failedAc current.LastError = &Error{Code: "unauthorized", Message: err.Error(), HTTPStatus: http.StatusUnauthorized} current.StatusMessage = "unauthorized (refresh token invalid)" shouldUnschedule = true - } else if !hasValidAccessToken { + } else if !hasValidAccessToken || accessTokenRejected { current.Unavailable = true current.Status = StatusError if unauthorized { @@ -774,19 +1135,24 @@ func (m *Manager) refreshAuthForRequestAtEpoch(ctx context.Context, id, failedAc } updated.LastRefreshedAt = now updated.NextRefreshAfter = time.Time{} - updated.LastError = nil - updated.StatusMessage = "" - updated.Unavailable = false - updated.RefreshFailures = 0 - if updated.Status == StatusError || updated.Status == "" { - updated.Status = StatusActive + applyPreservedRejectedTokens(auth, updated, true) + if updated.RejectedAccessToken == "" || authAccessToken(updated) != updated.RejectedAccessToken { + updated.LastError = nil + updated.StatusMessage = "" + updated.Unavailable = false + updated.RefreshFailures = 0 + if updated.Status == StatusError || updated.Status == "" { + updated.Status = StatusActive + } + _ = clearUnauthorizedModelStates(updated, now) } updated.UpdatedAt = now - _ = clearUnauthorizedModelStates(updated, now) if m.shouldRefresh(updated, now) { updated.NextRefreshAfter = now.Add(refreshIneffectiveBackoff) } - saved, errUpdate := m.UpdateRefreshedAuth(ctx, base, updated) + // A successful refresh may have rotated upstream tokens after cancellation. + // Commit its result with context values intact; keep epoch checks and save ordering. + saved, errUpdate := m.UpdateRefreshedAuth(context.WithoutCancel(ctx), base, updated) if errUpdate != nil { // Warn, not debug: a restart after a lost persist fails with invalid_grant. log.WithFields(log.Fields{"auth_id": auth.ID, "credential": auth.ID, "provider": auth.Provider}).Warnf("persist refreshed auth %s (%s) failed: %v", auth.Provider, auth.ID, errUpdate) diff --git a/sdk/cliproxy/auth/conductor_selection.go b/sdk/cliproxy/auth/conductor_selection.go index 073142e42..4420c4ca4 100644 --- a/sdk/cliproxy/auth/conductor_selection.go +++ b/sdk/cliproxy/auth/conductor_selection.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "maps" "math/rand/v2" "net/http" "reflect" @@ -11,6 +12,7 @@ import ( "strings" "time" + "github.com/google/uuid" "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" "github.com/router-for-me/CLIProxyAPI/v8/internal/thinking" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" @@ -42,25 +44,31 @@ func (m *Manager) hasPluginScheduler() bool { return true } -func (m *Manager) pluginSchedulerWantsAcrossPrioritiesLocked() bool { - if m == nil || m.pluginScheduler == nil { +func pluginSchedulerWantsAcrossPriorities(scheduler PluginScheduler) bool { + if scheduler == nil { return false } - if opt, ok := m.pluginScheduler.(PluginSchedulerAcrossPriorities); ok && opt != nil { + if opt, ok := scheduler.(PluginSchedulerAcrossPriorities); ok && opt != nil { return opt.SchedulerWantsAcrossPriorities() } return false } -// PluginSchedulerWantsAcrossPriorities reports whether the configured plugin scheduler -// opted into receiving candidates across all priority tiers. -func (m *Manager) PluginSchedulerWantsAcrossPriorities() bool { +func (m *Manager) snapshotPluginScheduler() (PluginScheduler, bool) { if m == nil { - return false + return nil, false } m.mu.RLock() - defer m.mu.RUnlock() - return m.pluginSchedulerWantsAcrossPrioritiesLocked() + scheduler := m.pluginScheduler + m.mu.RUnlock() + return scheduler, pluginSchedulerWantsAcrossPriorities(scheduler) +} + +// PluginSchedulerWantsAcrossPriorities reports whether the configured plugin scheduler +// opted into receiving candidates across all priority tiers. +func (m *Manager) PluginSchedulerWantsAcrossPriorities() bool { + _, wants := m.snapshotPluginScheduler() + return wants } func isBuiltInSelector(selector Selector) bool { @@ -213,6 +221,8 @@ func (m *Manager) ReconcileRegistryModelStates(ctx context.Context, authID strin cooldownStateChanged bool ) + releaseMutation := m.lockAuthMutation(authID) + defer releaseMutation() m.mu.Lock() auth, ok := m.auths[authID] if ok && auth != nil { @@ -347,7 +357,7 @@ func (m *Manager) ReconcileRegistryModelStates(ctx context.Context, authID strin } auth.Generation++ auth.UpdatedAt = now - if errPersist := m.persist(context.Background(), auth); errPersist != nil { + if errPersist := m.persistLocked(context.Background(), auth); errPersist != nil { logEntryWithRequestID(ctx).WithField("auth_id", auth.ID).Warnf("failed to persist auth changes during model state reconciliation: %v", errPersist) } } @@ -361,6 +371,7 @@ func (m *Manager) ReconcileRegistryModelStates(ctx context.Context, authID strin } } m.mu.Unlock() + releaseMutation() if snapshot == nil { return @@ -617,9 +628,8 @@ func (m *Manager) availableAuthsForRouteModelWithPriorityMode(auths []*Auth, pro // the plugin scheduler, plus the candidates handed to the configured selector. Both are equal // unless session affinity or an across-priorities scheduler is active, in which case the selector // or scheduler additionally receives lower priority tiers. -func (m *Manager) availableAuthsForSelector(selector Selector, auths []*Auth, provider, routeModel string, now time.Time) (priorityAuths, selectorAuths []*Auth, err error) { +func (m *Manager) availableAuthsForSelector(selector Selector, auths []*Auth, provider, routeModel string, now time.Time, schedulerAcross bool) (priorityAuths, selectorAuths []*Auth, err error) { _, sessionAffinity := selector.(*SessionAffinitySelector) - schedulerAcross := m.pluginSchedulerWantsAcrossPrioritiesLocked() if !sessionAffinity && !schedulerAcross { priorityAuths, err = m.availableAuthsForRouteModel(auths, provider, routeModel, now) @@ -672,6 +682,30 @@ func selectorContextForAvailableAuths(ctx context.Context, selector Selector, ro return context.WithValue(ctx, prevalidatedAuthCandidatesKey{}, true) } +func (m *Manager) transientlyUnavailableSessionAffinityAuthIDs(selector Selector, affinityCandidates []*Auth, routeModel string, now time.Time) map[string]struct{} { + affinity, ok := selector.(*SessionAffinitySelector) + if !ok || affinity == nil || !affinity.rebindOnQuotaOnly || len(affinityCandidates) == 0 { + return nil + } + if _, weighted := affinity.fallback.(*WeightedRoundRobinSelector); weighted { + affinityCandidates = positiveWeightAuths(affinityCandidates) + } + var authIDs map[string]struct{} + for _, candidate := range affinityCandidates { + if candidate == nil || candidate.ID == "" { + continue + } + checkModel := m.selectionModelForAuth(candidate, routeModel) + if !isAuthBlockedByQuotaOrAuthForModel(candidate, checkModel, now) { + if authIDs == nil { + authIDs = make(map[string]struct{}, len(affinityCandidates)) + } + authIDs[candidate.ID] = struct{}{} + } + } + return authIDs +} + func restoreModelCooldownErrorModel(err error, requestedModel string) error { if err == nil || requestedModel == "" { return err @@ -961,10 +995,31 @@ func (m *Manager) pickViaBuiltinScheduler(ctx context.Context, strategy schedule return selected, true, nil } +func clearSchedulerSelectionID(meta map[string]any) { + if len(meta) == 0 { + return + } + delete(meta, cliproxyexecutor.SelectionIDMetadataKey) +} + func (m *Manager) pickViaPluginScheduler(ctx context.Context, scheduler PluginScheduler, provider string, providers []string, model string, opts cliproxyexecutor.Options, tried map[string]struct{}, candidates []*Auth) (*Auth, bool, error) { if scheduler == nil || len(candidates) == 0 { + clearSchedulerSelectionID(opts.Metadata) return nil, false, nil } + if state, ok := scheduler.(pluginSchedulerState); ok && !state.HasScheduler() { + clearSchedulerSelectionID(opts.Metadata) + return nil, false, nil + } + opts.EnsureMetadata() + selectionID := uuid.NewString() + opts.Metadata[cliproxyexecutor.SelectionIDMetadataKey] = selectionID + keepSelectionID := false + defer func() { + if !keepSelectionID { + clearSchedulerSelectionID(opts.Metadata) + } + }() providerKey := strings.ToLower(strings.TrimSpace(provider)) requestProvider := providerKey if providerKey == "mixed" { @@ -994,9 +1049,10 @@ func (m *Manager) pickViaPluginScheduler(ctx context.Context, scheduler PluginSc if rejectMessage == "" { rejectMessage = "scheduler rejected candidate selection" } - return nil, true, &Error{Code: rejectCode, Message: rejectMessage} + return nil, true, wrapRequestStopError(&Error{Code: rejectCode, Message: rejectMessage}) } if selected := pickSchedulerAuthByID(candidates, resp.AuthID); selected != nil { + keepSelectionID = true return selected, true, nil } @@ -1004,7 +1060,11 @@ func (m *Manager) pickViaPluginScheduler(ctx context.Context, scheduler PluginSc if !okStrategy { return nil, false, nil } - return m.pickViaBuiltinScheduler(ctx, strategy, providerKey, providers, model, opts, tried) + selected, handledBuiltin, errBuiltin := m.pickViaBuiltinScheduler(ctx, strategy, providerKey, providers, model, opts, tried) + if errBuiltin == nil && handledBuiltin && selected != nil { + keepSelectionID = true + } + return selected, handledBuiltin, errBuiltin } func (m *Manager) authSupportsRouteModel(registryRef *registry.ModelRegistry, auth *Auth, routeModel string) bool { @@ -1550,6 +1610,69 @@ func (m *Manager) GetByID(id string) (*Auth, bool) { return auth.Clone(), true } +// PublishIfRegistrationCurrent serializes the final model-registration +// publication against identity mutations for this auth, then calls publish +// only if the manager still holds auth with the same model-registration +// inputs. It does not persist or clone auth state. publish must be fast and +// must not re-enter the manager. +func (m *Manager) PublishIfRegistrationCurrent(id string, registrationEpoch uint64, prefix, provider string, disabled bool, authKind string, attributes map[string]string, publish func() bool) bool { + if m == nil || id == "" || publish == nil { + return false + } + releaseMutation := m.lockAuthMutation(id) + defer releaseMutation() + m.mu.RLock() + current, ok := m.auths[id] + matches := ok && current != nil && modelRegistrationInputsMatch(current, registrationEpoch, prefix, provider, disabled, authKind, attributes) + m.mu.RUnlock() + if !matches { + return false + } + return publish() +} + +// UpdateIfRegistrationCurrent applies a plugin-returned auth update only +// while the manager still holds the exact registration inputs the plugin +// computation started from. It prevents an old plugin result from adopting +// a removed or re-registered identity merely because its zero registration +// epoch would otherwise inherit the current epoch during Update. +func (m *Manager) UpdateIfRegistrationCurrent(ctx context.Context, base *Auth, updated *Auth) (*Auth, bool, error) { + if m == nil || base == nil || updated == nil || base.ID == "" || updated.ID != base.ID { + return nil, false, nil + } + releaseMutation, errMutation := m.lockAuthMutationContext(ctx, base.ID) + if errMutation != nil { + return nil, false, errMutation + } + defer releaseMutation() + m.mu.RLock() + current, ok := m.auths[base.ID] + matches := ok && current != nil && modelRegistrationInputsMatch(current, base.RegistrationEpoch, base.Prefix, base.Provider, base.Disabled || base.Status == StatusDisabled, base.AuthKind(), base.Attributes) + m.mu.RUnlock() + if !matches { + return nil, false, nil + } + committedAuth, didCommit, persistCooldown, errUpdate := m.updateInternalLocked(ctx, nil, updated, updateModeReplace) + releaseMutation() + if errUpdate == nil && didCommit && committedAuth != nil { + m.finishAuthUpdate(ctx, committedAuth, persistCooldown) + } + return committedAuth, errUpdate == nil && didCommit && committedAuth != nil, errUpdate +} + +func modelRegistrationInputsMatch(current *Auth, registrationEpoch uint64, prefix, provider string, disabled bool, authKind string, attributes map[string]string) bool { + if current == nil { + return false + } + currentDisabled := current.Disabled || current.Status == StatusDisabled + return current.RegistrationEpoch == registrationEpoch && + current.Prefix == prefix && + current.Provider == provider && + currentDisabled == disabled && + current.AuthKind() == authKind && + maps.Equal(current.Attributes, attributes) +} + // GetExecutionSessionAuthByID retrieves a Home runtime auth scoped to an execution session. func (m *Manager) GetExecutionSessionAuthByID(sessionID string, authID string) (*Auth, bool) { sessionID = strings.TrimSpace(sessionID) @@ -1738,20 +1861,22 @@ func (m *Manager) routeAwareSelectionRequired(auth *Auth, routeModel string) boo func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, error) { if m.HomeEnabled() { + clearSchedulerSelectionID(opts.Metadata) auth, exec, _, err := m.pickNextViaHome(ctx, model, opts, tried) return auth, exec, err } opts.EnsureMetadata() + clearSchedulerSelectionID(opts.Metadata) opts.Metadata[cliproxyexecutor.SessionAffinityProviderMetadataKey] = provider pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata) eligibility := authSelectionEligibilityForRequest(ctx, opts) + pluginScheduler, schedulerAcross := m.snapshotPluginScheduler() m.mu.RLock() selector := m.selector opts.Metadata[cliproxyexecutor.SessionAffinityModelMetadataKey] = selectionArgForSelector(selector, model) - pluginScheduler := m.pluginScheduler executor, okExecutor := m.executorLocked(provider) if !okExecutor { m.mu.RUnlock() @@ -1768,6 +1893,12 @@ func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, op } registryRef := registry.GetGlobalRegistry() targetKey := canonicalSchedulingProvider(provider) + affinitySelector, hasAffinity := selector.(*SessionAffinitySelector) + trackTransientAffinity := hasAffinity && affinitySelector != nil && affinitySelector.rebindOnQuotaOnly + var affinityCandidates []*Auth + if trackTransientAffinity { + affinityCandidates = make([]*Auth, 0, len(m.auths)) + } for _, candidate := range m.auths { if candidate == nil || canonicalSchedulingProvider(executorKeyFromAuth(candidate)) != targetKey || candidate.Disabled { continue @@ -1781,10 +1912,13 @@ func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, op if m.quotaReserveReached(candidate, time.Now()) { continue } - if _, used := tried[candidate.ID]; used { + if modelKey != "" && !m.authSupportsRouteModel(registryRef, candidate, model) { continue } - if modelKey != "" && !m.authSupportsRouteModel(registryRef, candidate, model) { + if trackTransientAffinity { + affinityCandidates = append(affinityCandidates, candidate) + } + if _, used := tried[candidate.ID]; used { continue } candidates = append(candidates, candidate) @@ -1793,12 +1927,14 @@ func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, op m.mu.RUnlock() return nil, nil, &Error{Code: "auth_not_found", Message: "no auth available"} } - available, selectorAuths, errAvailable := m.availableAuthsForSelector(selector, candidates, provider, model, time.Now()) + now := time.Now() + available, selectorAuths, errAvailable := m.availableAuthsForSelector(selector, candidates, provider, model, now, schedulerAcross) if errAvailable != nil { m.mu.RUnlock() m.warnLogAuthUnavailable(ctx, []string{provider}, model, opts, tried, errAvailable) return nil, nil, errAvailable } + transientAuthIDs := m.transientlyUnavailableSessionAffinityAuthIDs(selector, affinityCandidates, model, now) m.mu.RUnlock() selected, handled, errPick := m.pickViaPluginScheduler(ctx, pluginScheduler, provider, []string{provider}, model, opts, tried, available) @@ -1808,6 +1944,7 @@ func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, op } if !handled { selectorCtx := selectorContextForAvailableAuths(ctx, selector, model) + selectorCtx = withSessionAffinityTransientAuthIDs(selectorCtx, transientAuthIDs) selected, errPick = selector.Pick(selectorCtx, provider, selectionArgForSelector(selector, model), opts, selectorAuths) if errPick != nil { if isBuiltInSelector(selector) { @@ -1840,7 +1977,7 @@ func (m *Manager) SelectAuth(ctx context.Context, provider, model string, opts c } selected, _, errPick := m.pickNextLegacy(ctx, provider, model, opts, nil) if errPick != nil { - return nil, errPick + return nil, unwrapExecutionBoundaryError(errPick) } if m.HomeEnabled() { return nil, &Error{Code: "home_unavailable", Message: "legacy auth selection is unavailable while Home is enabled", HTTPStatus: http.StatusServiceUnavailable} @@ -1862,7 +1999,7 @@ func (m *Manager) SelectAuthByKind(ctx context.Context, provider, model, require selectionCtx := withRequiredAuthKind(ctx, requiredKind) selected, _, errPick := m.pickNextLegacy(selectionCtx, provider, model, opts, nil) if errPick != nil { - return nil, errPick + return nil, unwrapExecutionBoundaryError(errPick) } if selected == nil { return nil, &Error{Code: "auth_not_found", Message: "selector returned no auth"} @@ -1888,7 +2025,7 @@ func (m *Manager) SelectAuthWithCredentialPolicy(ctx context.Context, provider, selectionCtx := withCredentialPolicy(ctx, policy) selected, _, errPick := m.pickNextLegacy(selectionCtx, provider, model, opts, nil) if errPick != nil { - return nil, errPick + return nil, unwrapExecutionBoundaryError(errPick) } if selected == nil || !credentialPolicyAllows(policy, selected) { return nil, &Error{Code: "auth_not_found", Message: "selector returned no eligible auth"} @@ -1999,6 +2136,7 @@ func (m *Manager) SelectHomeAuthByKind(ctx context.Context, provider string, mod func (m *Manager) pickNext(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, error) { opts.EnsureMetadata() + clearSchedulerSelectionID(opts.Metadata) if m.HomeEnabled() { auth, exec, _, err := m.pickNextViaHome(ctx, model, opts, tried) return auth, exec, err @@ -2060,10 +2198,12 @@ func (m *Manager) pickNext(ctx context.Context, provider, model string, opts cli func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, string, error) { if m.HomeEnabled() { + clearSchedulerSelectionID(opts.Metadata) return m.pickNextViaHome(ctx, model, opts, tried) } opts.EnsureMetadata() + clearSchedulerSelectionID(opts.Metadata) opts.Metadata[cliproxyexecutor.SessionAffinityProviderMetadataKey] = "mixed" pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata) @@ -2080,11 +2220,11 @@ func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, m if len(providerSet) == 0 { return nil, nil, "", &Error{Code: "provider_not_found", Message: "no provider supplied"} } + pluginScheduler, schedulerAcross := m.snapshotPluginScheduler() m.mu.RLock() selector := m.selector opts.Metadata[cliproxyexecutor.SessionAffinityModelMetadataKey] = selectionArgForSelector(selector, model) - pluginScheduler := m.pluginScheduler candidates := make([]*Auth, 0, len(m.auths)) modelKey := strings.TrimSpace(model) // Always use base model name (without thinking suffix) for auth matching. @@ -2095,6 +2235,12 @@ func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, m } } registryRef := registry.GetGlobalRegistry() + affinitySelector, hasAffinity := selector.(*SessionAffinitySelector) + trackTransientAffinity := hasAffinity && affinitySelector != nil && affinitySelector.rebindOnQuotaOnly + var affinityCandidates []*Auth + if trackTransientAffinity { + affinityCandidates = make([]*Auth, 0, len(m.auths)) + } for _, candidate := range m.auths { if candidate == nil || candidate.Disabled { continue @@ -2115,27 +2261,32 @@ func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, m if _, ok := providerSet[providerKey]; !ok { continue } - if _, used := tried[candidate.ID]; used { - continue - } if _, ok := m.executorLocked(providerKey); !ok { continue } if modelKey != "" && !m.authSupportsRouteModel(registryRef, candidate, model) { continue } + if trackTransientAffinity { + affinityCandidates = append(affinityCandidates, candidate) + } + if _, used := tried[candidate.ID]; used { + continue + } candidates = append(candidates, candidate) } if len(candidates) == 0 { m.mu.RUnlock() return nil, nil, "", &Error{Code: "auth_not_found", Message: "no auth available"} } - available, selectorAuths, errAvailable := m.availableAuthsForSelector(selector, candidates, "mixed", model, time.Now()) + now := time.Now() + available, selectorAuths, errAvailable := m.availableAuthsForSelector(selector, candidates, "mixed", model, now, schedulerAcross) if errAvailable != nil { m.mu.RUnlock() m.warnLogAuthUnavailable(ctx, providers, model, opts, tried, errAvailable) return nil, nil, "", errAvailable } + transientAuthIDs := m.transientlyUnavailableSessionAffinityAuthIDs(selector, affinityCandidates, model, now) m.mu.RUnlock() selected, handled, errPick := m.pickViaPluginScheduler(ctx, pluginScheduler, "mixed", providers, model, opts, tried, available) @@ -2145,6 +2296,7 @@ func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, m } if !handled { selectorCtx := selectorContextForAvailableAuths(ctx, selector, model) + selectorCtx = withSessionAffinityTransientAuthIDs(selectorCtx, transientAuthIDs) selected, errPick = selector.Pick(selectorCtx, "mixed", selectionArgForSelector(selector, model), opts, selectorAuths) if errPick != nil { if isBuiltInSelector(selector) { @@ -2176,6 +2328,7 @@ func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, m func (m *Manager) pickNextMixed(ctx context.Context, providers []string, model string, opts cliproxyexecutor.Options, tried map[string]struct{}) (*Auth, ProviderExecutor, string, error) { opts.EnsureMetadata() + clearSchedulerSelectionID(opts.Metadata) if m.HomeEnabled() { return m.pickNextViaHome(ctx, model, opts, tried) } diff --git a/sdk/cliproxy/auth/conductor_stream.go b/sdk/cliproxy/auth/conductor_stream.go index 1e0137b9a..744950f60 100644 --- a/sdk/cliproxy/auth/conductor_stream.go +++ b/sdk/cliproxy/auth/conductor_stream.go @@ -99,10 +99,16 @@ func readStreamBootstrap(ctx context.Context, ch <-chan cliproxyexecutor.StreamC ok bool ) if ctx != nil { + // Prefer a chunk that is already available so an upstream error delivered together + // with a context cancellation is not discarded. select { - case <-ctx.Done(): - return nil, false, ctx.Err() case chunk, ok = <-ch: + default: + select { + case <-ctx.Done(): + return nil, false, ctx.Err() + case chunk, ok = <-ch: + } } } else { chunk, ok = <-ch @@ -120,7 +126,23 @@ func readStreamBootstrap(ctx context.Context, ch <-chan cliproxyexecutor.StreamC } } -func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, resultModel, routeModel string, headers http.Header, buffered []cliproxyexecutor.StreamChunk, remaining <-chan cliproxyexecutor.StreamChunk, aliasResult OAuthModelAliasResult, ephemeralResult bool, opts cliproxyexecutor.Options) *cliproxyexecutor.StreamResult { +// recordCanceledUpstreamCooldown keeps an upstream-confirmed cooldown (an error carrying +// RetryAfter) when the request context was canceled while the executor was returning, for +// example because the downstream connection was closed in response to the same upstream error. +// Errors without RetryAfter, such as plain client cancellations, are not recorded. +func (m *Manager) recordCanceledUpstreamCooldown(ctx context.Context, auth *Auth, provider, resultModel, routeModel string, opts cliproxyexecutor.Options, err error, ephemeral bool, executionSeq uint64) { + retryAfter := retryAfterFromError(err) + if m == nil || auth == nil || retryAfter == nil { + return + } + action, okAction := matchRequestScopedErrorAction(auth, err, m.runtimeConfigSnapshot()) + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: false, Error: resultErrorFromError(err), Options: opts, RetryAfter: retryAfter, CredentialVersion: auth.CredentialVersion, RegistrationEpoch: auth.RegistrationEpoch, ExecutionSeq: executionSeq} + result.CredentialScope = isCredentialScopedError(err) + applyRequestScopedActionToResult(action, okAction, &result) + m.recordExecutionResult(context.WithoutCancel(ctx), result, auth, ephemeral) +} + +func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, resultModel, routeModel string, headers http.Header, buffered []cliproxyexecutor.StreamChunk, remaining <-chan cliproxyexecutor.StreamChunk, aliasResult OAuthModelAliasResult, ephemeralResult bool, opts cliproxyexecutor.Options, executionSeq uint64) *cliproxyexecutor.StreamResult { out := make(chan cliproxyexecutor.StreamChunk) streamStart := time.Now() go func() { @@ -138,7 +160,7 @@ func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, re warnLogUpstreamFailure(ctx, entry, provider, resultModel, auth, time.Since(streamStart), chunk.Err) rerr := resultErrorFromError(chunk.Err) action, okAction := matchRequestScopedErrorAction(auth, chunk.Err, m.runtimeConfigSnapshot()) - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: false, Error: rerr, Options: opts} + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: false, Error: rerr, Options: opts, CredentialVersion: auth.CredentialVersion, RegistrationEpoch: auth.RegistrationEpoch, ExecutionSeq: executionSeq} result.RetryAfter = retryAfterFromError(chunk.Err) result.CredentialScope = isCredentialScopedError(chunk.Err) applyRequestScopedActionToResult(action, okAction, &result) @@ -199,7 +221,7 @@ func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, re } } if !failed && (ephemeralResult || claudeOAuthRequestCancellation(ctx, auth, nil) == nil) { - m.recordExecutionResult(ctx, Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: true, Options: opts}, auth, ephemeralResult) + m.recordExecutionResult(ctx, Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: true, Options: opts, CredentialVersion: auth.CredentialVersion, RegistrationEpoch: auth.RegistrationEpoch, ExecutionSeq: executionSeq}, auth, ephemeralResult) } }() return &cliproxyexecutor.StreamResult{Headers: headers, Chunks: out} @@ -239,6 +261,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi } execOpts.Metadata = ensureCanonicalSessionMetadata(execOpts.Metadata, execOpts.Headers, payload) ctx = syncMetadataSessionToContext(ctx, execOpts.Metadata) + executionSeq := m.nextExecutionSeq() startStream := time.Now() streamResult, errStream := executor.ExecuteStream(ctx, auth, execReq, execOpts) errStream = markUpstreamExecutionAttemptFromContext(ctx, errStream) @@ -248,17 +271,19 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi durationStream := time.Since(startStream) if errStream != nil { if errCtx := ctx.Err(); errCtx != nil { + m.recordCanceledUpstreamCooldown(ctx, auth, provider, resultModel, routeModel, execOpts, errStream, ephemeralResult, executionSeq) return nil, errCtx } if allowRetry && !ephemeralResult { alreadyTried := didRefreshOnUnauthorized - refreshed, okRefresh := m.tryRefreshAfterUnauthorized(newUpstreamAttemptContext(ctx), auth, errStream, alreadyTried) + refreshed, okRefresh := m.tryRefreshAfterUnauthorized(newUpstreamAttemptContext(ctx), auth, errStream, alreadyTried, execOpts.Headers) if okRefresh { auth = refreshed publishSelectedAuthMetadata(execOpts.Metadata, auth) didRefreshOnUnauthorized = true ctx = newUpstreamAttemptContext(ctx) ctx = syncMetadataSessionToContext(ctx, execOpts.Metadata) + executionSeq = m.nextExecutionSeq() startRetry := time.Now() streamResult, errStream = executor.ExecuteStream(ctx, auth, execReq, execOpts) errStream = markUpstreamExecutionAttemptFromContext(ctx, errStream) @@ -269,6 +294,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi if errStream != nil { warnLogUpstreamFailure(ctx, entry, provider, execModel, auth, durationRetry, errStream) if errCtx := ctx.Err(); errCtx != nil { + m.recordCanceledUpstreamCooldown(ctx, auth, provider, resultModel, routeModel, execOpts, errStream, ephemeralResult, executionSeq) return nil, errCtx } } @@ -289,7 +315,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi if errStream != nil { rerr := resultErrorFromError(errStream) action, okAction := matchRequestScopedErrorAction(auth, errStream, m.runtimeConfigSnapshot()) - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: false, Error: rerr, Options: execOpts} + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: false, Error: rerr, Options: execOpts, CredentialVersion: auth.CredentialVersion, RegistrationEpoch: auth.RegistrationEpoch, ExecutionSeq: executionSeq} result.RetryAfter = retryAfterFromError(errStream) if isCredentialScopedError(errStream) { result.CredentialScope = true @@ -323,18 +349,20 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi } if bootstrapErr != nil { if errCtx := ctx.Err(); errCtx != nil { + m.recordCanceledUpstreamCooldown(ctx, auth, provider, resultModel, routeModel, execOpts, bootstrapErr, ephemeralResult, executionSeq) discardStreamChunks(streamResult.Chunks) return nil, errCtx } if allowRetry && !ephemeralResult { alreadyTried := didRefreshOnUnauthorized - refreshed, okRefresh := m.tryRefreshAfterUnauthorized(newUpstreamAttemptContext(ctx), auth, bootstrapErr, alreadyTried) + refreshed, okRefresh := m.tryRefreshAfterUnauthorized(newUpstreamAttemptContext(ctx), auth, bootstrapErr, alreadyTried, execOpts.Headers) if okRefresh { discardStreamChunks(streamResult.Chunks) auth = refreshed publishSelectedAuthMetadata(execOpts.Metadata, auth) didRefreshOnUnauthorized = true ctx = newUpstreamAttemptContext(ctx) + executionSeq = m.nextExecutionSeq() startRetry := time.Now() retryStream, retryErr := executor.ExecuteStream(ctx, auth, execReq, execOpts) retryErr = markUpstreamExecutionAttemptFromContext(ctx, retryErr) @@ -342,6 +370,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi retryErr = markUpstreamExecutionAttemptFromContext(ctx, retryErr) if retryErr != nil { if errCtx := ctx.Err(); errCtx != nil { + m.recordCanceledUpstreamCooldown(ctx, auth, provider, resultModel, routeModel, execOpts, retryErr, ephemeralResult, executionSeq) return nil, errCtx } bootstrapErr = retryErr @@ -375,7 +404,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi action, okAction := matchRequestScopedErrorAction(auth, bootstrapErr, m.runtimeConfigSnapshot()) if okAction { rerr := resultErrorFromError(bootstrapErr) - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: false, Error: rerr, Options: execOpts} + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: false, Error: rerr, Options: execOpts, CredentialVersion: auth.CredentialVersion, RegistrationEpoch: auth.RegistrationEpoch, ExecutionSeq: executionSeq} result.RetryAfter = retryAfterFromError(bootstrapErr) if isCredentialScopedError(bootstrapErr) { result.CredentialScope = true @@ -395,7 +424,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi } if isRequestInvalidError(bootstrapErr) { rerr := resultErrorFromError(bootstrapErr) - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: false, Error: rerr, Options: execOpts} + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: false, Error: rerr, Options: execOpts, CredentialVersion: auth.CredentialVersion, RegistrationEpoch: auth.RegistrationEpoch, ExecutionSeq: executionSeq} result.RetryAfter = retryAfterFromError(bootstrapErr) if isCredentialScopedError(bootstrapErr) { result.CredentialScope = true @@ -406,7 +435,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi } if idx < len(execModels)-1 { rerr := resultErrorFromError(bootstrapErr) - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: false, Error: rerr, Options: execOpts} + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: false, Error: rerr, Options: execOpts, CredentialVersion: auth.CredentialVersion, RegistrationEpoch: auth.RegistrationEpoch, ExecutionSeq: executionSeq} result.RetryAfter = retryAfterFromError(bootstrapErr) if isCredentialScopedError(bootstrapErr) { result.CredentialScope = true @@ -421,7 +450,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi continue } rerr := resultErrorFromError(bootstrapErr) - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: false, Error: rerr, Options: execOpts} + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: false, Error: rerr, Options: execOpts, CredentialVersion: auth.CredentialVersion, RegistrationEpoch: auth.RegistrationEpoch, ExecutionSeq: executionSeq} result.RetryAfter = retryAfterFromError(bootstrapErr) if isCredentialScopedError(bootstrapErr) { result.CredentialScope = true @@ -439,7 +468,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi upstreamErr = currentErr } warnLogUpstreamFailure(ctx, entry, provider, execModel, auth, time.Since(startStream), emptyErr) - result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: false, Error: resultErrorFromError(emptyErr), Options: execOpts} + result := Result{AuthID: auth.ID, Provider: provider, Model: resultModel, RouteModel: routeModel, Success: false, Error: resultErrorFromError(emptyErr), Options: execOpts, CredentialVersion: auth.CredentialVersion, RegistrationEpoch: auth.RegistrationEpoch, ExecutionSeq: executionSeq} m.recordExecutionResult(ctx, result, auth, ephemeralResult) if idx < len(execModels)-1 { lastErr = emptyErr @@ -455,7 +484,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi remaining = closedCh } attemptAliasResult := resolveAttemptAliasResult(routing, auth, routeModel, execModel, aliasResult) - return m.wrapStreamResult(ctx, auth.Clone(), provider, resultModel, routeModel, streamResult.Headers, buffered, remaining, attemptAliasResult, ephemeralResult, execOpts), nil + return m.wrapStreamResult(ctx, auth.Clone(), provider, resultModel, routeModel, streamResult.Headers, buffered, remaining, attemptAliasResult, ephemeralResult, execOpts, executionSeq), nil } if lastErr == nil { lastErr = &Error{Code: "auth_not_found", Message: "no upstream model available"} diff --git a/sdk/cliproxy/auth/cooldown_state.go b/sdk/cliproxy/auth/cooldown_state.go index c7cc5ae3f..b469beb94 100644 --- a/sdk/cliproxy/auth/cooldown_state.go +++ b/sdk/cliproxy/auth/cooldown_state.go @@ -27,6 +27,13 @@ type CooldownStateRecord struct { Quota QuotaState `json:"quota,omitempty"` LastError *Error `json:"last_error,omitempty"` UpdatedAt time.Time `json:"updated_at"` + // IndependentState preserves a model's own cooldown beneath an + // auth-level credential-quota layer so a restart cannot lose the + // distinction and quota restoration keeps reinstate-exactly semantics. + IndependentState *ModelState `json:"independent_state,omitempty"` + // CredentialQuotaPropagated preserves wrap provenance for models that + // carry no independent layer (e.g. the reporting model itself). + CredentialQuotaPropagated bool `json:"credential_quota_propagated,omitempty"` } const cooldownStateStatusObserved = "observed" diff --git a/sdk/cliproxy/auth/errors.go b/sdk/cliproxy/auth/errors.go index d11d7061d..268953d04 100644 --- a/sdk/cliproxy/auth/errors.go +++ b/sdk/cliproxy/auth/errors.go @@ -30,6 +30,14 @@ const transientTransportErrorCode = ErrorCodeTransientTransport // ErrorCodeForceCooldown marks failures that must enforce credential cooldown. const ErrorCodeForceCooldown = "force_cooldown" +// ErrorCodeStreamTruncated marks locally synthesized failures for accepted +// upstream streams that ended without a terminal frame. The public status +// stays 502, but the credential must not cool and the generation must not be +// replayed on another credential. +const ErrorCodeStreamTruncated = "stream_truncated" + +const streamTruncatedErrorCode = ErrorCodeStreamTruncated + // Error describes an authentication related failure in a provider agnostic format. type Error struct { // Code is a short machine readable identifier. diff --git a/sdk/cliproxy/auth/home_concurrency.go b/sdk/cliproxy/auth/home_concurrency.go index 17fc93227..8a26fb997 100644 --- a/sdk/cliproxy/auth/home_concurrency.go +++ b/sdk/cliproxy/auth/home_concurrency.go @@ -309,6 +309,12 @@ func SafeResponseHeaders(err error) http.Header { if errors.As(err, &unavailable) && unavailable != nil { return unavailable.Headers() } + var safeHeaderProvider interface{ SafeResponseHeaders() http.Header } + if errors.As(err, &safeHeaderProvider) && safeHeaderProvider != nil { + if value := safeHeaderProvider.SafeResponseHeaders().Get("Retry-After"); value != "" { + return http.Header{"Retry-After": []string{value}} + } + } return nil } diff --git a/sdk/cliproxy/auth/home_result.go b/sdk/cliproxy/auth/home_result.go index 793f26d44..d2567f5b5 100644 --- a/sdk/cliproxy/auth/home_result.go +++ b/sdk/cliproxy/auth/home_result.go @@ -46,12 +46,22 @@ func (m *Manager) reportHomeUnauthorized(ctx context.Context, auth *Auth, provid if alias == "" { alias = model } + accessProvider := coreusage.AccessProviderFromContext(ctx) + isNativeKey := false + if nativeExplicit, ok := coreusage.IsNativeKeyFromContext(ctx); ok { + isNativeKey = nativeExplicit + } else if accessProvider != "" { + isNativeKey = coreusage.IsNativeAccessProvider(accessProvider) + } clientMeta := logging.GetClientRequestMetadata(ctx) coreusage.PublishRecord(ctx, coreusage.Record{ Provider: provider, ExecutorType: homeResultExecutorType, Model: model, Alias: alias, + APIKey: coreusage.APIKeyFromContext(ctx), + IsNativeKey: isNativeKey, + AccessProvider: accessProvider, SessionID: clientMeta.SessionID, ParentSessionID: clientMeta.ParentSessionID, AuthID: auth.ID, diff --git a/sdk/cliproxy/auth/metadata_keys.go b/sdk/cliproxy/auth/metadata_keys.go index e861cfedb..ffb662ac2 100644 --- a/sdk/cliproxy/auth/metadata_keys.go +++ b/sdk/cliproxy/auth/metadata_keys.go @@ -16,6 +16,10 @@ func CanonicalCredentialMetadataKey(key string) string { return "fingerprint_profile" case "model-aliases": return "model_aliases" + case "openapi-location": + return "openapi_location" + case "openapi-models": + return "openapi_models" case "proxy-url": return "proxy_url" case "request-retry": diff --git a/sdk/cliproxy/auth/metadata_merge.go b/sdk/cliproxy/auth/metadata_merge.go index 590567f42..f6592944b 100644 --- a/sdk/cliproxy/auth/metadata_merge.go +++ b/sdk/cliproxy/auth/metadata_merge.go @@ -4,6 +4,8 @@ import ( "reflect" "strings" "time" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" ) // IsAuthTokenPayloadKey returns true if key is a credential or token lifecycle field @@ -70,7 +72,10 @@ func MergeRefreshedAuth(base, current, updated *Auth) *Auth { return merged } + merged.CodexRefreshPatch = updated.CodexRefreshPatch + // 1. Refresh Lifecycle Timestamps + applyPreservedRejectedTokens(current, merged, true) if !updated.LastRefreshedAt.IsZero() { merged.LastRefreshedAt = updated.LastRefreshedAt } @@ -119,7 +124,7 @@ func MergeRefreshedAuth(base, current, updated *Auth) *Auth { merged.Status = current.Status merged.Unavailable = current.Unavailable merged.StatusMessage = current.StatusMessage - } else if current.Quota.Exceeded && current.Quota.Reason == "credential_quota" && current.Quota.NextRecoverAt.After(time.Now()) { + } else if current.Quota.Exceeded && current.Quota.Reason == "credential_quota" && current.Quota.NextRecoverAt.After(time.Now()) && !devinQuotaObservationRestored(merged, current, time.Now()) { // Preserve active credential quota merged.Unavailable = current.Unavailable merged.Status = current.Status @@ -129,7 +134,7 @@ func MergeRefreshedAuth(base, current, updated *Auth) *Auth { merged.Unavailable = current.Unavailable merged.Status = current.Status merged.StatusMessage = current.StatusMessage - } else if updated.Status == StatusActive || updated.Status == "" { + } else if (updated.Status == StatusActive || updated.Status == "") && (merged.RejectedAccessToken == "" || authAccessToken(merged) != merged.RejectedAccessToken) { // Successful refresh clears previous auth error and restores active merged.Status = StatusActive merged.Unavailable = false @@ -138,7 +143,28 @@ func MergeRefreshedAuth(base, current, updated *Auth) *Auth { } } + // 2b. Quota observations: adopt the refreshed snapshot when it is newer. + // Without this, executor Refresh results (e.g. Devin quota reset signals) + // never reach the stored auth and cannot drive scheduling. Cooldown + // fields stay with current so concurrent failures are never lost. + // The merge happens after the status branches above (they only copy + // scalar fields), so ordering relative to them is safe. + merged.Quota = mergeQuotaObservation(merged.Quota, updated.Quota) + // 3. ModelStates: three-way merge to preserve concurrent cooldown/quota + // Independent layers follow the same current-wins rule: executor refresh + // results are derived from the base snapshot and never move them. + for model, mergedState := range merged.ModelStates { + if mergedState == nil { + continue + } + if currentState := current.ModelStates[model]; currentState != nil { + if currentState.credentialQuotaShadow != nil && mergedState.credentialQuotaShadow == nil { + mergedState.credentialQuotaShadow = currentState.credentialQuotaShadow.cloneShallow() + } + mergedState.credentialQuotaPropagated = mergedState.credentialQuotaPropagated || currentState.credentialQuotaPropagated + } + } var baseModels map[string]*ModelState if base != nil { baseModels = base.ModelStates @@ -208,6 +234,34 @@ func mergeAuthContent(base, current, updated *Auth) *Auth { if strings.EqualFold(strings.TrimSpace(k), "proxy_url") { continue } + if updated.CodexRefreshPatch != nil && codex.IsCodexCredentialKey(strings.ToLower(strings.TrimSpace(k))) { + continue + } + if strings.EqualFold(strings.TrimSpace(k), "siwc") { + siwcUpd, okUpd := v.(map[string]any) + if !okUpd || len(siwcUpd) == 0 { + continue + } + baseVal, hadInBase := baseMeta[k] + changedByExecutor := !hadInBase || !reflect.DeepEqual(baseVal, v) + if changedByExecutor { + mergedSIWC := make(map[string]any) + if curSIWC, okCur := current.Metadata[k].(map[string]any); okCur { + for sk, sv := range curSIWC { + mergedSIWC[sk] = sv + } + } else if bSIWC, okBase := baseVal.(map[string]any); okBase { + for sk, sv := range bSIWC { + mergedSIWC[sk] = sv + } + } + for sk, sv := range siwcUpd { + mergedSIWC[sk] = sv + } + merged.Metadata[k] = mergedSIWC + } + continue + } baseVal, hadInBase := baseMeta[k] currentVal, hadInCurrent := current.Metadata[k] @@ -224,7 +278,8 @@ func mergeAuthContent(base, current, updated *Auth) *Auth { // Deletions by executor: only delete if user didn't modify the field concurrently if baseMeta != nil { for k, baseVal := range baseMeta { - if strings.EqualFold(strings.TrimSpace(k), "proxy_url") { + trimmedKey := strings.ToLower(strings.TrimSpace(k)) + if trimmedKey == "proxy_url" || trimmedKey == "siwc" || trimmedKey == "use_siwc" { continue } if _, inUpdated := updated.Metadata[k]; !inUpdated { @@ -238,6 +293,10 @@ func mergeAuthContent(base, current, updated *Auth) *Auth { } } + if updated.CodexRefreshPatch != nil { + merged.Metadata = updated.CodexRefreshPatch.Apply(merged.Metadata) + } + // 2. Storage and Runtime if updated.Storage != nil { merged.Storage = updated.Storage diff --git a/sdk/cliproxy/auth/native_tool_search_capability.go b/sdk/cliproxy/auth/native_tool_search_capability.go new file mode 100644 index 000000000..cf585ceb7 --- /dev/null +++ b/sdk/cliproxy/auth/native_tool_search_capability.go @@ -0,0 +1,64 @@ +package auth + +import ( + "strings" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" +) + +// NativeToolSearchSupport optionally reports whether an executor can send Codex +// tool_search to a credential as Anthropic native tool search. It does not +// change ProviderExecutor or the plugin ABI. +type NativeToolSearchSupport interface { + SupportsNativeToolSearch(auth *Auth, upstreamModel string) bool +} + +// SupportsNativeToolSearchForModel requires native tool search support from +// every enabled credential that serves the public model. A prefixed ID that no +// credential registers exactly falls back to its unprefixed form. A model with +// no serving credential is not supported. +func (m *Manager) SupportsNativeToolSearchForModel(model string) bool { + model = strings.TrimSpace(model) + if m == nil || model == "" { + return false + } + auths := m.List() + serving := nativeToolSearchServingAuths(auths, model) + if len(serving) == 0 { + if slash := strings.Index(model, "/"); slash >= 0 { + model = strings.TrimSpace(model[slash+1:]) + serving = nativeToolSearchServingAuths(auths, model) + } + } + if len(serving) == 0 { + return false + } + for _, auth := range serving { + executor, okExecutor := m.Executor(auth.Provider) + if !okExecutor { + return false + } + support, okSupport := executor.(NativeToolSearchSupport) + if !okSupport || !support.SupportsNativeToolSearch(auth, m.ResolveExecutionModel(auth, model)) { + return false + } + } + return true +} + +func nativeToolSearchServingAuths(auths []*Auth, model string) []*Auth { + if model == "" { + return nil + } + modelRegistry := registry.GetGlobalRegistry() + var serving []*Auth + for _, auth := range auths { + if auth == nil || auth.Disabled || auth.Status == StatusDisabled { + continue + } + if modelRegistry.ClientSupportsModel(auth.ID, model) { + serving = append(serving, auth) + } + } + return serving +} diff --git a/sdk/cliproxy/auth/persist_policy.go b/sdk/cliproxy/auth/persist_policy.go index 26bcc92bd..6a4615d3a 100644 --- a/sdk/cliproxy/auth/persist_policy.go +++ b/sdk/cliproxy/auth/persist_policy.go @@ -37,6 +37,18 @@ func WithSkipPersist(ctx context.Context) context.Context { return context.WithValue(ctx, skipPersistContextKey{}, true) } +// WithoutSkipPersist returns a derived context that re-enables persistence even +// when the parent context was marked with WithSkipPersist. +func WithoutSkipPersist(ctx context.Context) context.Context { + if ctx == nil { + return context.Background() + } + if !shouldSkipPersist(ctx) { + return ctx + } + return context.WithValue(ctx, skipPersistContextKey{}, false) +} + func shouldSkipPersist(ctx context.Context) bool { if ctx == nil { return false diff --git a/sdk/cliproxy/auth/scheduler.go b/sdk/cliproxy/auth/scheduler.go index 39c138e74..752a52d7d 100644 --- a/sdk/cliproxy/auth/scheduler.go +++ b/sdk/cliproxy/auth/scheduler.go @@ -891,6 +891,12 @@ func isCredentialBlocked(auth *Auth, supportedModelCount int, now time.Time) boo if auth.Disabled || auth.Status == StatusDisabled { return true } + if hasUnauthorizedAuthFailure(auth) { + return true + } + if exp, ok := auth.AccessTokenExpirationTime(); ok && !exp.IsZero() && !exp.After(now) { + return true + } if auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" && auth.Quota.NextRecoverAt.After(now) { return true } diff --git a/sdk/cliproxy/auth/selector.go b/sdk/cliproxy/auth/selector.go index bece8ec26..fc7805413 100644 --- a/sdk/cliproxy/auth/selector.go +++ b/sdk/cliproxy/auth/selector.go @@ -910,10 +910,134 @@ func availabilityBlock(unavailable, quotaExceeded bool, nextRetryAfter, nextReco // It extracts session ID from multiple sources and maintains session-to-auth // mappings with automatic failover when the bound auth becomes unavailable. type SessionAffinitySelector struct { - fallback Selector - cache *SessionCache - matcher *cliproxysession.MerklePrefixMatcher - subagentAffinity bool + fallback Selector + cache *SessionCache + matcher *cliproxysession.MerklePrefixMatcher + subagentAffinity bool + rebindOnQuotaOnly bool +} + +const ( + // SessionAffinityRebindOnAny rebinds the session on any credential-attributed failure (default). + SessionAffinityRebindOnAny = "any" + // SessionAffinityRebindOnQuota preserves the existing session binding across transient failures + // and only rebinds on quota, auth, model-support, or permanent unavailability. + SessionAffinityRebindOnQuota = "quota" +) + +func normalizeSessionAffinityRebindOn(value string) string { + switch strings.ToLower(strings.TrimSpace(value)) { + case SessionAffinityRebindOnQuota: + return SessionAffinityRebindOnQuota + default: + return SessionAffinityRebindOnAny + } +} + +type sessionAffinityTransientAuthIDsKey struct{} + +func withSessionAffinityTransientAuthIDs(ctx context.Context, authIDs map[string]struct{}) context.Context { + if len(authIDs) == 0 { + return ctx + } + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, sessionAffinityTransientAuthIDsKey{}, authIDs) +} + +func sessionAffinityTransientAuthIDsFromContext(ctx context.Context) map[string]struct{} { + if ctx == nil { + return nil + } + authIDs, _ := ctx.Value(sessionAffinityTransientAuthIDsKey{}).(map[string]struct{}) + return authIDs +} + +func shouldRebindSessionAffinityOnError(err *Error) bool { + if err == nil || shouldSkipCredentialCooldown(err) { + return false + } + if isModelSupportResultError(err) || isCloudflareChallengeResultError(err) || isInvalidGrantResultError(err) { + return true + } + switch statusCodeFromResult(err) { + case http.StatusUnauthorized, http.StatusPaymentRequired, http.StatusForbidden, http.StatusNotFound, http.StatusTooManyRequests: + return true + default: + return false + } +} + +func isAuthBlockedByQuotaOrAuthForModel(auth *Auth, model string, now time.Time) bool { + if auth == nil { + return true + } + if auth.Disabled || auth.Status == StatusDisabled { + return true + } + if hasUnauthorizedAuthFailure(auth) { + return true + } + if exp, ok := auth.AccessTokenExpirationTime(); ok && !exp.IsZero() && !exp.After(now) { + return true + } + if auth.Quota.Exceeded && auth.Quota.Reason == "credential_quota" && (auth.Quota.NextRecoverAt.IsZero() || auth.Quota.NextRecoverAt.After(now)) { + return true + } + if model != "" && len(auth.ModelStates) > 0 { + modelKey := canonicalModelKey(model) + for stateModel, state := range auth.ModelStates { + if state == nil || canonicalModelKey(stateModel) != modelKey { + continue + } + if state.Status == StatusDisabled { + return true + } + if isStateBlockedByQuotaOrAuth(state.Unavailable, state.Quota, state.NextRetryAfter, state.LastError, now) { + return true + } + } + return false + } + quota := auth.Quota + if len(auth.ModelStates) > 0 && auth.Quota.Reason != "credential_quota" && !auth.Unavailable { + quota.Exceeded = false + } + return isStateBlockedByQuotaOrAuth(auth.Unavailable, quota, auth.NextRetryAfter, auth.LastError, now) +} + +func isStateBlockedByQuotaOrAuth(unavailable bool, quota QuotaState, nextRetryAfter time.Time, lastErr *Error, now time.Time) bool { + if quota.Exceeded && (quota.NextRecoverAt.IsZero() || quota.NextRecoverAt.After(now)) { + return true + } + blocked, _, _ := availabilityBlock(unavailable, quota.Exceeded, nextRetryAfter, quota.NextRecoverAt, now) + if !blocked { + return lastErr != nil && shouldRebindSessionAffinityOnError(lastErr) + } + if lastErr != nil { + return shouldRebindSessionAffinityOnError(lastErr) + } + return nextRetryAfter.IsZero() && quota.NextRecoverAt.IsZero() +} + +func (s *SessionAffinitySelector) shouldPreserveBindingWhenUnavailable(ctx context.Context, candidates []*Auth, model, authID string, now time.Time) bool { + if s == nil || !s.rebindOnQuotaOnly || authID == "" { + return false + } + if ctx != nil { + if validated, _ := ctx.Value(prevalidatedAuthCandidatesKey{}).(bool); validated { + transientAuthIDs := sessionAffinityTransientAuthIDsFromContext(ctx) + _, ok := transientAuthIDs[authID] + return ok + } + } + for _, candidate := range candidates { + if candidate != nil && candidate.ID == authID { + return !isAuthBlockedByQuotaOrAuthForModel(candidate, model, now) + } + } + return false } // SessionAffinityConfig configures the session affinity selector. @@ -921,6 +1045,7 @@ type SessionAffinityConfig struct { Fallback Selector TTL time.Duration SubagentAffinity *bool + RebindOn string } // NewSessionAffinitySelector creates a new session-aware selector. @@ -943,11 +1068,13 @@ func NewSessionAffinitySelectorWithConfig(cfg SessionAffinityConfig) *SessionAff if cfg.SubagentAffinity != nil { subagentAffinity = *cfg.SubagentAffinity } + rebindOnQuotaOnly := normalizeSessionAffinityRebindOn(cfg.RebindOn) == SessionAffinityRebindOnQuota return &SessionAffinitySelector{ - fallback: cfg.Fallback, - cache: NewSessionCache(cfg.TTL), - matcher: cliproxysession.NewMerklePrefixMatcher(cfg.TTL), - subagentAffinity: subagentAffinity, + fallback: cfg.Fallback, + cache: NewSessionCache(cfg.TTL), + matcher: cliproxysession.NewMerklePrefixMatcher(cfg.TTL), + subagentAffinity: subagentAffinity, + rebindOnQuotaOnly: rebindOnQuotaOnly, } } @@ -1048,18 +1175,31 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri if fallbackID != "" && fallbackID != primaryID { fallbackKey = provider + "::" + fallbackID + "::" + modelKey } - bind := func(authID string) { + bind := func(authID string, recordSelection bool) { + keys := []string{cacheKey} if fallbackKey != "" && !isSubagent && !isFork { - s.cache.SetAliases(authID, cacheKey, fallbackKey) - } else { - s.cache.Set(cacheKey, authID) + keys = append(keys, fallbackKey) + } + if !recordSelection { + s.cache.SetAliases(authID, keys...) + return } + bindingProvider := canonicalSchedulingProvider(provider) + for _, candidate := range auths { + if candidate != nil && candidate.ID == authID { + if candidateProvider := canonicalSchedulingProvider(executorKeyFromAuth(candidate)); candidateProvider != "" { + bindingProvider = candidateProvider + } + break + } + } + s.cache.setSessionAliases(authID, bindingProvider, primaryID, keys...) } if cachedAuthID, ok := s.cache.GetAndRefresh(cacheKey); ok { for _, auth := range available { if auth.ID == cachedAuthID { - bind(auth.ID) + bind(auth.ID, true) entry.Infof("session-affinity: cache hit | session=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), auth.ID, provider, model) return auth, nil } @@ -1072,17 +1212,45 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri if auth == nil { return nil, nil } - bind(auth.ID) + if s.shouldPreserveBindingWhenUnavailable(ctx, availabilityCandidates, model, cachedAuthID, now) { + bind(cachedAuthID, false) + entry.Infof("session-affinity: cache hit but auth transiently unavailable, failover without rebinding | session=%s bound_auth=%s failover_auth=%s provider=%s model=%s", truncateSessionID(primaryID), cachedAuthID, auth.ID, provider, model) + return auth, nil + } + bind(auth.ID, true) entry.Infof("session-affinity: cache hit but auth unavailable, reselected | session=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), auth.ID, provider, model) return auth, nil } + if isSubagent && s.subagentAffinity && isGuardianRequest(opts.Headers) { + providers := make(map[string]struct{}) + schedulingProvider := canonicalSchedulingProvider(provider) + for _, candidate := range available { + candidateProvider := canonicalSchedulingProvider(executorKeyFromAuth(candidate)) + if candidateProvider == "" { + candidateProvider = schedulingProvider + } + if schedulingProvider == "mixed" || candidateProvider == schedulingProvider { + providers[candidateProvider] = struct{}{} + } + } + if parentAuthID := s.cache.latestSessionAuth(fallbackID, providers); parentAuthID != "" { + for _, auth := range available { + if auth.ID == parentAuthID { + bind(auth.ID, true) + entry.Infof("session-affinity: guardian inherited parent binding | session=%s parent=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), truncateSessionID(fallbackID), auth.ID, provider, model) + return auth, nil + } + } + } + } + if fallbackKey != "" { if cachedAuthID, ok := s.cache.Get(fallbackKey); ok { for _, auth := range available { if auth.ID == cachedAuthID { if !isSubagent || s.subagentAffinity { - bind(auth.ID) + bind(auth.ID, true) if isFork { entry.Infof("session-affinity: fork cache hit | session=%s parent=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), truncateSessionID(fallbackID), auth.ID, provider, model) } else { @@ -1092,6 +1260,18 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri } } } + if !isSubagent && !isFork && s.shouldPreserveBindingWhenUnavailable(ctx, availabilityCandidates, model, cachedAuthID, now) { + auth, err := s.fallback.Pick(ctx, provider, model, opts, fallbackAuths) + if err != nil { + return nil, err + } + if auth == nil { + return nil, nil + } + bind(cachedAuthID, false) + entry.Infof("session-affinity: fallback cache hit but auth transiently unavailable, failover without rebinding | session=%s fallback=%s bound_auth=%s failover_auth=%s provider=%s model=%s", truncateSessionID(primaryID), truncateSessionID(fallbackID), cachedAuthID, auth.ID, provider, model) + return auth, nil + } } } @@ -1102,7 +1282,7 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri if auth == nil { return nil, nil } - bind(auth.ID) + bind(auth.ID, true) if isFork && fallbackID != "" { entry.Infof("session-affinity: fork bound to new auth | session=%s parent=%s auth=%s provider=%s model=%s", truncateSessionID(primaryID), truncateSessionID(fallbackID), auth.ID, provider, model) } else { @@ -1138,11 +1318,13 @@ func (s *SessionAffinitySelector) pickLCP(ctx context.Context, provider, model s if _, weighted := s.fallback.(*WeightedRoundRobinSelector); weighted { availabilityCandidates = positiveWeightAuths(auths) } - available, errAvailable := getSelectorAvailableAuthsAcrossPriorities(ctx, availabilityCandidates, provider, model, time.Now()) + now := time.Now() + available, errAvailable := getSelectorAvailableAuthsAcrossPriorities(ctx, availabilityCandidates, provider, model, now) if errAvailable != nil { return nil, true, errAvailable } + var preservedAuthID string if match, ok := s.matcher.MatchFingerprintsWithContext(namespace, fingerprints, tailFingerprints, envDigest, minPrefixLength); ok { for _, auth := range available { if auth == nil || auth.ID != match.AuthID { @@ -1184,6 +1366,9 @@ func (s *SessionAffinitySelector) pickLCP(ctx context.Context, provider, model s } return auth, true, nil } + if !match.IsFork && s.shouldPreserveBindingWhenUnavailable(ctx, availabilityCandidates, model, match.AuthID, now) { + preservedAuthID = match.AuthID + } } fallbackAuths := highestPriorityAuths(available) @@ -1194,7 +1379,11 @@ func (s *SessionAffinitySelector) pickLCP(ctx context.Context, provider, model s if auth == nil { return nil, true, &Error{Code: "auth_not_found", Message: "selector returned no auth"} } - if bindRes := s.matcher.BindFingerprintsWithContext(namespace, fingerprints, tailFingerprints, envDigest, minPrefixLength, auth.ID); bindRes.SessionID != "" { + bindAuthID := auth.ID + if preservedAuthID != "" { + bindAuthID = preservedAuthID + } + if bindRes := s.matcher.BindFingerprintsWithContext(namespace, fingerprints, tailFingerprints, envDigest, minPrefixLength, bindAuthID); bindRes.SessionID != "" { opts.Metadata[cliproxyexecutor.LCPAffinitySessionIDMetadataKey] = bindRes.SessionID opts.Metadata[cliproxyexecutor.CanonicalSessionIDMetadataKey] = bindRes.SessionID if bindRes.ParentSessionID != "" { @@ -1205,7 +1394,20 @@ func (s *SessionAffinitySelector) pickLCP(ctx context.Context, provider, model s if bindRes.AccessNumber > 0 && opts.Metadata != nil { opts.Metadata[cliproxyexecutor.LCPAccessGenerationMetadataKey] = bindRes.AccessNumber } - if bindRes.IsFork { + if preservedAuthID != "" { + if bindRes.IsCompaction { + if opts.Metadata != nil { + opts.Metadata[cliproxyexecutor.IsCompactionMetadataKey] = true + delete(opts.Metadata, cliproxyexecutor.IsForkMetadataKey) + opts.Metadata[cliproxyexecutor.NodeKindMetadataKey] = "compaction" + } + } else if opts.Metadata != nil { + delete(opts.Metadata, cliproxyexecutor.IsForkMetadataKey) + delete(opts.Metadata, cliproxyexecutor.IsCompactionMetadataKey) + delete(opts.Metadata, cliproxyexecutor.NodeKindMetadataKey) + } + entry.Infof("session-affinity: LCP cache hit but auth transiently unavailable, failover without rebinding | session=%s bound_auth=%s failover_auth=%s provider=%s model=%s", truncateSessionID(bindRes.SessionID), preservedAuthID, auth.ID, provider, model) + } else if bindRes.IsFork { if opts.Metadata != nil { opts.Metadata[cliproxyexecutor.IsForkMetadataKey] = true delete(opts.Metadata, cliproxyexecutor.IsCompactionMetadataKey) @@ -1481,6 +1683,11 @@ func (s *SessionAffinitySelector) OnResult(res Result) { // selected credential is unhealthy, so preserve both explicit and LCP bindings. return } + if !res.Success && s.rebindOnQuotaOnly && !shouldRebindSessionAffinityOnError(res.Error) { + // Under session-affinity-rebind-on: quota, transient failures fail over for + // the current request without unbinding the session from its bound credential. + return + } // LCP bindings are independent from explicit harness bindings. A successful // extension is recorded as a new sequence while credential-attributed failures @@ -1504,7 +1711,14 @@ func (s *SessionAffinitySelector) OnResult(res Result) { generation = gen } } - s.matcher.RemoveFingerprintsBefore(namespace, fingerprints, res.AuthID, generation) + if !s.matcher.RemoveFingerprintsBefore(namespace, fingerprints, res.AuthID, generation) && res.Options.Metadata != nil { + isFork, _ := res.Options.Metadata[cliproxyexecutor.IsForkMetadataKey].(bool) + if !isFork { + if lcpSessionID := sessionMetadataString(res.Options.Metadata, cliproxyexecutor.LCPAffinitySessionIDMetadataKey); lcpSessionID != "" { + s.matcher.RemoveSessionBefore(namespace, lcpSessionID, res.AuthID, generation) + } + } + } } } } @@ -1551,6 +1765,16 @@ func (s *SessionAffinitySelector) OnResult(res Result) { } } +// isGuardianRequest deliberately matches the header, not the configurable review model. +func isGuardianRequest(headers http.Header) bool { + for key, values := range headers { + if strings.EqualFold(key, "X-Openai-Subagent") && len(values) > 0 { + return strings.EqualFold(strings.TrimSpace(values[0]), "guardian") + } + } + return false +} + // normalizedSessionCandidate validates an explicit client-provided session signal. // It keeps opaque printable IDs intact while rejecting values that are unsafe or // implausibly large for routing keys and logs. diff --git a/sdk/cliproxy/auth/session_cache.go b/sdk/cliproxy/auth/session_cache.go index b136bced5..eb3482550 100644 --- a/sdk/cliproxy/auth/session_cache.go +++ b/sdk/cliproxy/auth/session_cache.go @@ -12,11 +12,19 @@ const ( defaultMaxSessionEntries = 65536 ) +// sessionBinding records selection order independently of TTL refreshes and eviction order. +type sessionBinding struct { + sessionID string + provider string + sequence uint64 +} + // sessionEntry stores an auth binding, its identifier aliases, and expiration. type sessionEntry struct { authID string expiresAt time.Time aliases []string + binding sessionBinding } // SessionCache provides TTL-based session to auth mapping with automatic cleanup. @@ -30,6 +38,8 @@ type SessionCache struct { ttl time.Duration stopCh chan struct{} stopOnce sync.Once + sessionGroups map[string]map[string]struct{} + bindingSequence uint64 } // NewSessionCache creates a cache with the specified TTL. @@ -60,6 +70,9 @@ func NewSessionCacheWithCapacity(ttl time.Duration, maxEntries int) *SessionCach } func (c *SessionCache) ensureInitializedLocked() { + if c.sessionGroups == nil { + c.sessionGroups = make(map[string]map[string]struct{}) + } if c.entries == nil { c.entries = make(map[string]sessionEntry) } @@ -126,7 +139,7 @@ func (c *SessionCache) GetAndRefresh(sessionID string) (string, bool) { } aliases := compactSessionAliases(mergeSessionAliases([]string{sessionID}, entry.aliases...)) - c.replaceAliasGroupsLocked(entry.authID, now.Add(c.ttl), aliases, entry) + c.replaceAliasGroupsLocked(entry.authID, now.Add(c.ttl), aliases, entry.binding, entry) return entry.authID, true } @@ -141,6 +154,16 @@ func (c *SessionCache) Set(sessionID, authID string) { // SetAliases binds multiple identifiers for one logical session to an auth ID. func (c *SessionCache) SetAliases(authID string, sessionIDs ...string) { + c.setAliases(authID, sessionBinding{}, sessionIDs...) +} + +// setSessionAliases records a real selection, including reuse of an existing binding. +// Result-driven Touch calls preserve this order rather than becoming new selections. +func (c *SessionCache) setSessionAliases(authID, provider, sessionID string, keys ...string) { + c.setAliases(authID, sessionBinding{sessionID: sessionID, provider: provider}, keys...) +} + +func (c *SessionCache) setAliases(authID string, binding sessionBinding, sessionIDs ...string) { if c == nil || authID == "" { return } @@ -167,10 +190,20 @@ func (c *SessionCache) SetAliases(authID string, sessionIDs ...string) { if len(aliases) == 0 { return } - c.replaceAliasGroupsLocked(authID, now.Add(c.ttl), aliases, previousGroups...) + if binding.sessionID != "" && binding.provider != "" { + c.bindingSequence++ + binding.sequence = c.bindingSequence + } else { + for _, previous := range previousGroups { + if previous.binding.sequence > binding.sequence { + binding = previous.binding + } + } + } + c.replaceAliasGroupsLocked(authID, now.Add(c.ttl), aliases, binding, previousGroups...) } -func (c *SessionCache) replaceAliasGroupsLocked(authID string, expiresAt time.Time, aliases []string, previousGroups ...sessionEntry) { +func (c *SessionCache) replaceAliasGroupsLocked(authID string, expiresAt time.Time, aliases []string, binding sessionBinding, previousGroups ...sessionEntry) { for _, previous := range previousGroups { c.removeAliasGroupLocked(previous) } @@ -181,8 +214,14 @@ func (c *SessionCache) replaceAliasGroupsLocked(authID string, expiresAt time.Ti if existing, ok := c.groups[primaryKey]; ok { c.removeAliasGroupLocked(existing) } - entry := sessionEntry{authID: authID, expiresAt: expiresAt, aliases: append([]string(nil), aliases...)} + entry := sessionEntry{authID: authID, expiresAt: expiresAt, aliases: append([]string(nil), aliases...), binding: binding} c.groups[primaryKey] = entry + if binding.sequence != 0 { + if c.sessionGroups[binding.sessionID] == nil { + c.sessionGroups[binding.sessionID] = make(map[string]struct{}) + } + c.sessionGroups[binding.sessionID][primaryKey] = struct{}{} + } for _, alias := range aliases { c.entries[alias] = entry } @@ -216,6 +255,12 @@ func (c *SessionCache) removeAliasGroupLocked(entry sessionEntry) { primaryKey := entry.aliases[0] if currentGroup, ok := c.groups[primaryKey]; ok && sameSessionEntryGroup(currentGroup, entry) { delete(c.groups, primaryKey) + if groups := c.sessionGroups[entry.binding.sessionID]; groups != nil { + delete(groups, primaryKey) + if len(groups) == 0 { + delete(c.sessionGroups, entry.binding.sessionID) + } + } if elem, ok := c.evictionElements[primaryKey]; ok { c.evictionOrder.Remove(elem) delete(c.evictionElements, primaryKey) @@ -309,6 +354,28 @@ func mergeSessionAliases(existing []string, candidates ...string) []string { return aliases } +// latestSessionAuth observes the newest live selection on the requested real providers. +// It does not refresh TTL or eviction order and visits only this session's groups. +func (c *SessionCache) latestSessionAuth(sessionID string, providers map[string]struct{}) string { + if c == nil || sessionID == "" || len(providers) == 0 { + return "" + } + c.mu.RLock() + defer c.mu.RUnlock() + now := time.Now() + var latest sessionEntry + for key := range c.sessionGroups[sessionID] { + entry := c.groups[key] + if _, ok := providers[entry.binding.provider]; !ok || !now.Before(entry.expiresAt) { + continue + } + if entry.binding.sequence > latest.binding.sequence { + latest = entry + } + } + return latest.authID +} + // Touch refreshes the expiration for a session binding if it currently matches expectedAuthID. func (c *SessionCache) Touch(sessionID, expectedAuthID string) bool { if c == nil || sessionID == "" || expectedAuthID == "" { @@ -323,7 +390,7 @@ func (c *SessionCache) Touch(sessionID, expectedAuthID string) bool { return false } aliases := compactSessionAliases(mergeSessionAliases([]string{sessionID}, entry.aliases...)) - c.replaceAliasGroupsLocked(expectedAuthID, now.Add(c.ttl), aliases, entry) + c.replaceAliasGroupsLocked(expectedAuthID, now.Add(c.ttl), aliases, entry.binding, entry) return true } @@ -348,7 +415,7 @@ func (c *SessionCache) CompareAndDelete(sessionID, expectedAuthID string) bool { } } if len(surviving) > 0 { - c.replaceAliasGroupsLocked(entry.authID, entry.expiresAt, surviving) + c.replaceAliasGroupsLocked(entry.authID, entry.expiresAt, surviving, entry.binding) } return true } @@ -375,7 +442,7 @@ func (c *SessionCache) Invalidate(sessionID string) { } } if len(surviving) > 0 { - c.replaceAliasGroupsLocked(entry.authID, entry.expiresAt, surviving) + c.replaceAliasGroupsLocked(entry.authID, entry.expiresAt, surviving, entry.binding) } } diff --git a/sdk/cliproxy/auth/store.go b/sdk/cliproxy/auth/store.go index 0594a77dd..a903a46bf 100644 --- a/sdk/cliproxy/auth/store.go +++ b/sdk/cliproxy/auth/store.go @@ -7,6 +7,12 @@ type Store interface { // List returns all auth records stored in the backend. List(ctx context.Context) ([]*Auth, error) // Save persists the provided auth record, replacing any existing one with same ID. + // It may enrich exported fields, but must not change ID, RegistrationEpoch, + // CredentialVersion or Generation. The manager merges these changes even on non-fatal save errors. + // Do not retain auth or mutate shared nested values in place: replace metadata + // values, pointers, Storage and Runtime instead. Auth.Clone shares these objects. + // Save may call manager read APIs, but must not synchronously mutate the + // manager or call Load (the reload barrier remains held until publication). Save(ctx context.Context, auth *Auth) (string, error) // Delete removes the auth record identified by id. Delete(ctx context.Context, id string) error diff --git a/sdk/cliproxy/auth/types.go b/sdk/cliproxy/auth/types.go index 702ea010c..dc224010f 100644 --- a/sdk/cliproxy/auth/types.go +++ b/sdk/cliproxy/auth/types.go @@ -1,20 +1,25 @@ package auth import ( + "bytes" "context" "crypto/sha256" "encoding/base64" "encoding/hex" "encoding/json" + "io" "net/http" "net/url" + "os" "path/filepath" "strconv" "strings" "sync" "time" + "github.com/gin-gonic/gin" baseauth "github.com/router-for-me/CLIProxyAPI/v8/internal/auth" + "github.com/router-for-me/CLIProxyAPI/v8/internal/auth/codex" ) // PostAuthHook defines a function that is called after an Auth record is created @@ -50,6 +55,8 @@ type Auth struct { ID string `json:"id"` // RegistrationEpoch tracks monotonic registration cycles across unregister/re-register. RegistrationEpoch uint64 `json:"registration_epoch,omitempty"` + // CredentialVersion tracks monotonic credential replacements for secret material. + CredentialVersion uint64 `json:"credential_version,omitempty"` // Generation tracks monotonic mutations to resolve scheduler/reconcile snapshot races. Generation uint64 `json:"generation,omitempty"` // Index is a stable runtime identifier derived from auth metadata (not persisted). @@ -62,6 +69,8 @@ type Auth struct { FileName string `json:"-"` // Storage holds the token persistence implementation used during login flows. Storage baseauth.TokenStorage `json:"-"` + // CodexRefreshPatch carries an immutable pending credential update, never configuration. + CodexRefreshPatch *codex.CredentialRefreshPatch `json:"-"` // Label is an optional human readable label for logging. Label string `json:"label,omitempty"` // Status is the lifecycle status managed by the AuthManager. @@ -100,11 +109,138 @@ type Auth struct { // Runtime carries non-serialisable data used during execution (in-memory only). Runtime any `json:"-"` + // RejectedAccessToken tracks the access token rejected by upstream 401 until a refresh succeeds. + RejectedAccessToken string `json:"-"` + // RejectedSIWCAccessToken tracks the nested SIWC access token rejected by upstream 401 until replaced. + RejectedSIWCAccessToken string `json:"-"` + Success int64 `json:"-"` Failed int64 `json:"-"` + // cooldownRecordedSeq is the execution sequence watermark captured when + // the current auth-level cooldown was recorded. Successes from executions + // stamped at or below it were dispatched before the cooldown existed and + // must not clear it. Runtime only: zero after restart so restored + // cooldowns can still be cleared by any success. + cooldownRecordedSeq uint64 `json:"-"` + recentRequests recentRequestRing `json:"-"` indexAssigned bool `json:"-"` + // fileSnapshotPath and fileSnapshotDigest bind an unversioned file-synthesized + // Auth snapshot to the source file path and SHA-256 byte digest it observed. + // Cleared when committed to a Manager so runtime clones are not treated as + // unversioned file snapshots. + fileSnapshotPath string `json:"-"` + fileSnapshotDigest string `json:"-"` +} + +var readAuthSnapshotFileFunc = os.ReadFile + +func canonicalAuthFileDigest(data []byte) string { + trimmed := bytes.TrimSpace(data) + if len(trimmed) == 0 { + return "" + } + decoder := json.NewDecoder(bytes.NewReader(trimmed)) + decoder.UseNumber() + var parsed any + if errDecode := decoder.Decode(&parsed); errDecode == nil { + var trailing any + if errTrailing := decoder.Decode(&trailing); errTrailing == io.EOF { + if canonical, errMarshal := json.Marshal(parsed); errMarshal == nil { + sum := sha256.Sum256(canonical) + return hex.EncodeToString(sum[:]) + } + } + } + sum := sha256.Sum256(trimmed) + return hex.EncodeToString(sum[:]) +} + +// BindFileSnapshot binds an unversioned file-synthesized Auth snapshot to the +// source file path and canonical SHA-256 digest of the bytes from which it was parsed. +func (a *Auth) BindFileSnapshot(path string, data []byte) { + if a == nil || len(data) == 0 { + return + } + digest := canonicalAuthFileDigest(data) + if digest == "" { + return + } + a.fileSnapshotPath = strings.TrimSpace(path) + a.fileSnapshotDigest = digest +} + +const invalidatedFileSnapshotDigest = "invalidated" + +// ClearFileSnapshot removes any unversioned file-observation binding once the +// Auth is committed to runtime state or normalized for comparison. +func (a *Auth) ClearFileSnapshot() { + if a == nil { + return + } + a.fileSnapshotPath = "" + a.fileSnapshotDigest = "" +} + +// InvalidateFileSnapshot marks an unversioned file snapshot as discarded so a +// subsequent re-read of identical or reformatted file content is not suppressed. +func (a *Auth) InvalidateFileSnapshot() { + if !a.HasFileSnapshot() { + return + } + a.fileSnapshotDigest = invalidatedFileSnapshotDigest +} + +// HasFileSnapshot reports whether this Auth is an unversioned file-synthesized +// snapshot bound to a specific file observation digest. +func (a *Auth) HasFileSnapshot() bool { + return a != nil && a.Generation == 0 && a.RegistrationEpoch == 0 && a.fileSnapshotDigest != "" +} + +// FileSnapshotPath returns the backing file path bound to this file snapshot. +func (a *Auth) FileSnapshotPath() string { + if a == nil { + return "" + } + path := strings.TrimSpace(a.fileSnapshotPath) + if path == "" && a.Attributes != nil { + path = strings.TrimSpace(a.Attributes[AttributePath]) + if path == "" { + path = strings.TrimSpace(a.Attributes[AttributeSource]) + } + } + return path +} + +// SameFileSnapshot reports whether two unversioned file snapshots observed the +// same canonical file content. +func (a *Auth) SameFileSnapshot(other *Auth) bool { + if !a.HasFileSnapshot() || !other.HasFileSnapshot() { + return false + } + if a.fileSnapshotDigest == invalidatedFileSnapshotDigest || other.fileSnapshotDigest == invalidatedFileSnapshotDigest { + return false + } + return a.fileSnapshotDigest == other.fileSnapshotDigest && a.fileSnapshotPath == other.fileSnapshotPath +} + +// FileSnapshotMatchesDisk reports whether the backing file on disk still +// exists, is readable, and matches the canonical digest observed when this +// snapshot was synthesized. +func (a *Auth) FileSnapshotMatchesDisk() bool { + if !a.HasFileSnapshot() { + return true + } + path := a.FileSnapshotPath() + if path == "" { + return false + } + data, errRead := readAuthSnapshotFileFunc(path) + if errRead != nil || len(data) == 0 { + return false + } + return canonicalAuthFileDigest(data) == a.fileSnapshotDigest } const ( @@ -227,6 +363,28 @@ type ModelState struct { Quota QuotaState `json:"quota"` // UpdatedAt tracks the last update timestamp for this model state. UpdatedAt time.Time `json:"updated_at"` + // cooldownRecordedSeq is the execution sequence watermark captured when + // the current model cooldown was recorded. Successes from executions + // stamped at or below it were dispatched before the cooldown existed and + // must not clear it. Runtime only: zero after restart so restored + // cooldowns can still be cleared by any success. + cooldownRecordedSeq uint64 `json:"-"` + // credentialQuotaShadow holds this model's independent (own-failure) + // cooldown layer while the outer fields present the auth-level + // credential quota propagated over it, so quota restoration can drop + // only the propagated layer and reinstate precisely what the model's + // own failures justify (issue #6690). It is nil when no independent + // live block exists beneath an outer credential_quota label. Persisted + // via CooldownStateRecord.IndependentState (unexported fields are + // invisible to JSON, so the record owns the serialized form). + credentialQuotaShadow *ModelState `toml:"credential_quota_shadow,omitempty"` + // credentialQuotaPropagated marks that credential-quota propagation has + // already wrapped this state, distinguishing "no independent block" + // from "not yet propagated" so repeated propagation never mistakes its + // own earlier layer for an independent block. Persisted via + // CooldownStateRecord.CredentialQuotaPropagated across restarts, + // including for shadow-less reporting models. + credentialQuotaPropagated bool `toml:"credential_quota_propagated,omitempty"` } func recentRequestBucketID(now time.Time) int64 { @@ -307,10 +465,7 @@ func (a *Auth) Clone() *Auth { } } if a.Metadata != nil { - copyAuth.Metadata = make(map[string]any, len(a.Metadata)) - for key, value := range a.Metadata { - copyAuth.Metadata[key] = value - } + copyAuth.Metadata = cloneMetadataMap(a.Metadata) } if a.ModelStates != nil { copyAuth.ModelStates = make(map[string]*ModelState, len(a.ModelStates)) @@ -322,6 +477,30 @@ func (a *Auth) Clone() *Auth { return ©Auth } +func cloneMetadataMap(src map[string]any) map[string]any { + if src == nil { + return nil + } + dst := make(map[string]any, len(src)) + for key, value := range src { + switch typed := value.(type) { + case map[string]any: + dst[key] = cloneMetadataMap(typed) + case []string: + s := make([]string, len(typed)) + copy(s, typed) + dst[key] = s + case []any: + s := make([]any, len(typed)) + copy(s, typed) + dst[key] = s + default: + dst[key] = value + } + } + return dst +} + func stableAuthIndex(seed string) string { seed = strings.TrimSpace(seed) if seed == "" { @@ -453,6 +632,28 @@ func (m *ModelState) Clone() *ModelState { HTTPStatus: m.LastError.HTTPStatus, } } + if m.credentialQuotaShadow != nil { + copyState.credentialQuotaShadow = m.credentialQuotaShadow.cloneShallow() + } + return ©State +} + +// cloneShallow copies a shadow state one level deep; shadows never nest. +func (m *ModelState) cloneShallow() *ModelState { + if m == nil { + return nil + } + copyState := *m + copyState.Quota = m.Quota.Clone() + copyState.credentialQuotaShadow = nil + if m.LastError != nil { + copyState.LastError = &Error{ + Code: m.LastError.Code, + Message: m.LastError.Message, + Retryable: m.LastError.Retryable, + HTTPStatus: m.LastError.HTTPStatus, + } + } return ©State } @@ -604,6 +805,14 @@ func (a *Auth) AccountInfo() (string, string) { return "oauth", email } } + if sa, ok := a.Metadata["service_account"].(map[string]any); ok { + if v, okEmail := sa["client_email"].(string); okEmail { + email := strings.TrimSpace(v) + if email != "" { + return "oauth", email + } + } + } } return "oauth", "" case AuthKindAPIKey: @@ -616,15 +825,97 @@ func (a *Auth) AccountInfo() (string, string) { } } -// ExpirationTime attempts to extract the credential expiration timestamp from metadata. -// It inspects common absolute expiry keys, expires_in plus timestamp, and nested -// token objects to remain compatible with legacy auth file formats. -// If the access_token contains a valid JWT exp claim, it is given priority. -func (a *Auth) ExpirationTime() (time.Time, bool) { +// IsCodexUseSIWC reports whether a Codex auth record is configured to use its nested "siwc" section. +func IsCodexUseSIWC(a *Auth) bool { + if a == nil { + return false + } + enabled := false + if len(a.Attributes) > 0 { + if raw := strings.TrimSpace(a.Attributes["use_siwc"]); raw != "" { + if parsed, ok := parseBoolAny(raw); ok { + enabled = parsed + } + } + } + if !enabled && a.Metadata != nil { + if parsed, ok := parseBoolAny(a.Metadata["use_siwc"]); ok { + enabled = parsed + } + } + return enabled +} + +// IsCodexClientUserAgent reports whether userAgent identifies a Codex client. +func IsCodexClientUserAgent(userAgent string) bool { + ua := strings.ToLower(strings.TrimSpace(userAgent)) + if ua == "" { + return false + } + return strings.Contains(ua, "codex-desktop") || + strings.Contains(ua, "codex desktop") || + strings.Contains(ua, "codex-rs") || + strings.Contains(ua, "codex_cli_rs") || + strings.Contains(ua, "codex-tui") || + strings.Contains(ua, "codex_exec") +} + +func headerValueFromHTTPHeader(headers http.Header, name string) string { + if headers == nil { + return "" + } + if value := strings.TrimSpace(headers.Get(name)); value != "" { + return value + } + for key, values := range headers { + if !strings.EqualFold(key, name) { + continue + } + for _, value := range values { + if value = strings.TrimSpace(value); value != "" { + return value + } + } + } + return "" +} + +// IsCodexClientRequest reports whether the request originates from a Codex client. +func IsCodexClientRequest(ctx context.Context, clientHeaders ...http.Header) bool { + if ctx != nil { + if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil { + if IsCodexClientUserAgent(headerValueFromHTTPHeader(ginCtx.Request.Header, "User-Agent")) { + return true + } + } + } + for _, headers := range clientHeaders { + if IsCodexClientUserAgent(headerValueFromHTTPHeader(headers, "User-Agent")) { + return true + } + } + return false +} + +// IsCodexUseSIWCForRequest reports whether a request should use the nested SIWC credential section. +// Requests from Codex clients ignore use_siwc and always use non-SIWC credentials. +func IsCodexUseSIWCForRequest(ctx context.Context, a *Auth, clientHeaders ...http.Header) bool { + if IsCodexClientRequest(ctx, clientHeaders...) { + return false + } + return IsCodexUseSIWC(a) +} + +// OuterExpirationTime extracts the outer credential expiration timestamp, ignoring the nested "siwc" section. +func (a *Auth) OuterExpirationTime() (time.Time, bool) { if a == nil { return time.Time{}, false } - if tokenStr := authAccessToken(a); tokenStr != "" { + tokenStr := authOuterAccessToken(a) + if tokenStr != "" && a.RejectedAccessToken != "" && a.RejectedAccessToken == tokenStr { + return time.Unix(0, 0), true + } + if tokenStr != "" { if jwtExp, ok := parseJWTExp(tokenStr); ok { return jwtExp, true } @@ -635,6 +926,246 @@ func (a *Auth) ExpirationTime() (time.Time, bool) { return time.Time{}, false } +// ExpirationTime attempts to extract the active credential expiration timestamp from metadata. +// When use_siwc is enabled on a Codex credential, the active SIWC section expiration is returned. +func (a *Auth) ExpirationTime() (time.Time, bool) { + if a == nil { + return time.Time{}, false + } + if IsCodexUseSIWC(a) { + siwcToken := authSIWCTokenField(a, "access_token") + if a.isSIWCAccessTokenRejected(siwcToken) { + return time.Unix(0, 0), true + } + return a.SIWCExpirationTime() + } + return a.OuterExpirationTime() +} + +// CodexOuterRefreshLead is the conservative proactive refresh lead time before +// the nominal 10-day Codex OAuth JWT expiration, refreshing tokens well ahead +// of early upstream access-token invalidation. +const CodexOuterRefreshLead = 96 * time.Hour + +// CodexSIWCRefreshLead is the minimum remaining validity required before refreshing +// a Sign in with ChatGPT (SIWC) token, matching pi's DEFAULT_OAUTH_MINIMUM_VALIDITY_MS (5 minutes). +const CodexSIWCRefreshLead = 5 * time.Minute + +// SIWCStatus represents the read-only validity state of a Codex SIWC credential. +type SIWCStatus string + +const ( + SIWCStatusOff SIWCStatus = "off" + SIWCStatusValid SIWCStatus = "valid" + SIWCStatusInvalid SIWCStatus = "invalid" +) + +// SIWCRealExpirationTime extracts the real (unbuffered) expiration timestamp for the nested +// "siwc" access token. It reads "expires_at" (saved from OAuth expires_in) or the JWT "exp" +// claim of "access_token", without using the margin-deducted "expired" field. +func (a *Auth) SIWCRealExpirationTime() (time.Time, bool) { + if a == nil || a.Metadata == nil { + return time.Time{}, false + } + siwcMap, ok := a.Metadata["siwc"].(map[string]any) + if !ok || len(siwcMap) == 0 { + return time.Time{}, false + } + if ts, okExp := lookupMetadataTime(siwcMap, "expires_at", "expiresAt"); okExp && !ts.IsZero() { + return ts, true + } + if expiresIn, okExpiresIn := parseRelativeExpirySeconds(siwcMap); okExpiresIn { + if ts, okTs := lookupMetadataTime(siwcMap, "last_refresh", "lastRefresh", "timestamp", "issued_at", "issuedAt"); okTs && !ts.IsZero() { + return ts.Add(time.Duration(expiresIn) * time.Second), true + } + } + if rawToken, okTok := siwcMap["access_token"].(string); okTok && strings.TrimSpace(rawToken) != "" { + if jwtExp, okJWT := parseJWTExp(rawToken); okJWT && !jwtExp.IsZero() { + return jwtExp, true + } + } + return time.Time{}, false +} + +func (a *Auth) isSIWCAccessTokenRejected(siwcToken string) bool { + if a == nil { + return false + } + siwcToken = strings.TrimSpace(siwcToken) + if siwcToken == "" { + return false + } + return (a.RejectedSIWCAccessToken != "" && a.RejectedSIWCAccessToken == siwcToken) || + (a.RejectedAccessToken != "" && a.RejectedAccessToken == siwcToken) +} + +// CodexSIWCStatus evaluates the local SIWC credential status ("off", "valid", or "invalid") +// at the given time without triggering any refresh or upstream request. +func (a *Auth) CodexSIWCStatus(now time.Time) SIWCStatus { + if a == nil || !IsCodexUseSIWC(a) { + return SIWCStatusOff + } + siwcToken := authSIWCTokenField(a, "access_token") + if siwcToken == "" { + return SIWCStatusInvalid + } + if now.IsZero() { + now = time.Now() + } + realExp, okExp := a.SIWCRealExpirationTime() + if !okExp || realExp.IsZero() || !realExp.After(now) { + return SIWCStatusInvalid + } + if a.isSIWCAccessTokenRejected(siwcToken) { + return SIWCStatusInvalid + } + return SIWCStatusValid +} + +// SIWCExpirationTime extracts the expiration timestamp from the nested "siwc" metadata section. +func (a *Auth) SIWCExpirationTime() (time.Time, bool) { + if a == nil || a.Metadata == nil { + return time.Time{}, false + } + siwcMap, ok := a.Metadata["siwc"].(map[string]any) + if !ok || len(siwcMap) == 0 { + return time.Time{}, false + } + if ts, okExp := expirationFromMap(siwcMap); okExp && !ts.IsZero() { + return ts, true + } + if rawToken, okTok := siwcMap["access_token"].(string); okTok && strings.TrimSpace(rawToken) != "" { + if jwtExp, okJWT := parseJWTExp(rawToken); okJWT && !jwtExp.IsZero() { + return jwtExp, true + } + } + return time.Time{}, false +} + +// CodexOuterNeedsRefresh evaluates the outer credential refresh due time and whether it needs refresh at now. +func CodexOuterNeedsRefresh(a *Auth, now time.Time) (dueAt time.Time, hasRefresh bool, needsRefresh bool) { + if a == nil { + return time.Time{}, false, false + } + if authRefreshToken(a) == "" { + return time.Time{}, false, false + } + hasRefresh = true + + outerToken := authOuterAccessToken(a) + if outerToken == "" || (a.RejectedAccessToken != "" && a.RejectedAccessToken == outerToken) { + return now, true, true + } + + var lastRefresh time.Time + if ts, ok := authLastRefreshTimestamp(a); ok && !ts.IsZero() { + lastRefresh = ts + } else { + lastRefresh = a.LastRefreshedAt + } + expiry, hasExpiry := a.OuterExpirationTime() + + if pref := authPreferredInterval(a); pref > 0 { + candidates := make([]time.Time, 0, 2) + if hasExpiry && !expiry.IsZero() { + candidates = append(candidates, expiry.Add(-pref)) + } + if lastRefresh.IsZero() { + return now, true, true + } + candidates = append(candidates, lastRefresh.Add(pref)) + next := candidates[0] + for _, candidate := range candidates[1:] { + if candidate.Before(next) { + next = candidate + } + } + if !next.After(now) { + return now, true, true + } + return next, true, false + } + + provider := strings.ToLower(strings.TrimSpace(a.Provider)) + if provider == "" { + provider = strings.ToLower(authMetadataString(a, "type")) + } + lead := ProviderRefreshLead(provider, a.Runtime) + if lead == nil && provider == "codex" { + d := CodexOuterRefreshLead + lead = &d + } + if lead == nil { + return time.Time{}, false, false + } + if *lead <= 0 { + if hasExpiry && !expiry.IsZero() { + if !expiry.After(now) { + return now, true, true + } + return expiry, true, false + } + return time.Time{}, false, false + } + if hasExpiry && !expiry.IsZero() { + due := expiry.Add(-*lead) + if !due.After(now) { + return now, true, true + } + return due, true, false + } + if !lastRefresh.IsZero() { + due := lastRefresh.Add(*lead) + if !due.After(now) { + return now, true, true + } + return due, true, false + } + return now, true, true +} + +// CodexSIWCNeedsRefresh evaluates the nested "siwc" credential section refresh due time and whether it needs refresh at now. +func CodexSIWCNeedsRefresh(a *Auth, now time.Time) (dueAt time.Time, hasRefresh bool, needsRefresh bool) { + if a == nil || a.Metadata == nil { + return time.Time{}, false, false + } + siwcMap, ok := a.Metadata["siwc"].(map[string]any) + if !ok || len(siwcMap) == 0 { + return time.Time{}, false, false + } + refreshToken, _ := siwcMap["refresh_token"].(string) + clientID, _ := siwcMap["client_id"].(string) + if strings.TrimSpace(refreshToken) == "" || strings.TrimSpace(clientID) == "" { + return time.Time{}, false, false + } + hasRefresh = true + + siwcToken, _ := siwcMap["access_token"].(string) + siwcToken = strings.TrimSpace(siwcToken) + if siwcToken == "" { + return now, true, true + } + if a.isSIWCAccessTokenRejected(siwcToken) { + return now, true, true + } + expiry, hasExpiry := a.SIWCExpirationTime() + if hasExpiry && !expiry.IsZero() { + due := expiry.Add(-CodexSIWCRefreshLead) + if !due.After(now) { + return now, true, true + } + return due, true, false + } + if lastRefresh, okLast := lookupMetadataTime(siwcMap, "last_refresh", "lastRefresh"); okLast && !lastRefresh.IsZero() { + due := lastRefresh.Add(CodexSIWCRefreshLead) + if !due.After(now) { + return now, true, true + } + return due, true, false + } + return now, true, true +} + // AccessTokenExpirationTime returns the expiration time of the specific access_token. // If the access_token is a JWT, its exp claim takes strict precedence. func (a *Auth) AccessTokenExpirationTime() (time.Time, bool) { @@ -645,9 +1176,17 @@ func (a *Auth) AccessTokenExpirationTime() (time.Time, bool) { if tokenStr == "" { return time.Time{}, false } + if (a.RejectedAccessToken != "" && a.RejectedAccessToken == tokenStr) || (IsCodexUseSIWC(a) && a.isSIWCAccessTokenRejected(tokenStr)) { + return time.Unix(0, 0), true + } if jwtExp, ok := parseJWTExp(tokenStr); ok { return jwtExp, true } + if IsCodexUseSIWC(a) { + if realExp, ok := a.SIWCRealExpirationTime(); ok && !realExp.IsZero() { + return realExp, true + } + } return a.ExpirationTime() } @@ -813,11 +1352,15 @@ func ProviderRefreshLead(provider string, runtime any) *time.Duration { refreshLeadMu.RLock() factory := refreshLeadFactories[provider] refreshLeadMu.RUnlock() - if factory == nil { + if factory != nil { + if lead := factory(); lead != nil && *lead > 0 { + return lead + } return nil } - if lead := factory(); lead != nil && *lead > 0 { - return lead + if provider == "codex" { + d := CodexOuterRefreshLead + return &d } return nil }