From a1b27c65c99405b5e3a4a54e89a44d9714018819 Mon Sep 17 00:00:00 2001 From: nanashiwang Date: Sun, 13 Sep 2026 17:32:03 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=E7=8B=AC=E4=BA=AB=E8=B4=A6=E5=8F=B7?= =?UTF-8?q?=E9=A2=9D=E5=BA=A6=E8=80=97=E5=B0=BD=E6=97=B6=E5=AE=89=E5=85=A8?= =?UTF-8?q?=E5=88=87=E6=8D=A2=E5=90=8C=E7=BB=84=E7=A9=BA=E9=97=B2=E8=B4=A6?= =?UTF-8?q?=E5=8F=B7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs/ACCOUNT_POOL_LEASES_CN.md | 13 +- internal/poollease/store.go | 47 ++++++ internal/poollease/store_test.go | 55 +++++++ sdk/cliproxy/auth/account_pools_test.go | 4 + sdk/cliproxy/auth/conductor_execution.go | 40 ++++- sdk/cliproxy/auth/conductor_stream.go | 20 +-- sdk/cliproxy/auth/pool_leases.go | 183 +++++++++++++++++++---- sdk/cliproxy/auth/pool_leases_test.go | 180 ++++++++++++++++++++++ 8 files changed, 498 insertions(+), 44 deletions(-) diff --git a/docs/ACCOUNT_POOL_LEASES_CN.md b/docs/ACCOUNT_POOL_LEASES_CN.md index be845ced149..e518be0b2da 100644 --- a/docs/ACCOUNT_POOL_LEASES_CN.md +++ b/docs/ACCOUNT_POOL_LEASES_CN.md @@ -42,8 +42,8 @@ account-pools: - Key 绑定实例后只使用被授权的租约池；未绑定实例的普通 Key 只能使用非租约池，包括 scope=all 的 Key。 - `X-CPA-Instance-ID` 必须与 Key 绑定实例匹配，`X-CPA-User-ID` 必须是正整数。身份仅在客户端鉴权后读取，专用头随即从请求中移除，不转发给模型上游。 - 归属按“实例 + 用户 ID”的摘要确定，与 CPA Key 和下游令牌无关。同一用户换渠道或 Key 不会再占一个池；若新 Key 无权使用原池，返回权限错误。 -- 固定 60 分钟，从首次分配开始；普通调用不续期。首次从授权组内选择一个可用账号，此后只使用该账号。账号暂时不可用或不支持后续请求模型时返回错误，不再切换到组内其他账号。 -- 无空闲且可用的授权账号返回 503 `pool_busy`，不排无界队列、不抢占。已绑定账号失败时不自动换号。 +- 固定 60 分钟，从首次分配开始；普通调用不续期。首次从授权组内选择一个可用账号，正常情况下继续使用该账号。明确额度耗尽、刷新后仍未授权或账号失效时，可安全切换同组空闲账号，保留原到期时间。临时限流、服务过载、请求错误或模型不支持不会触发换号。 +- 无空闲且可用的授权账号返回 503 `pool_busy`，不排无界队列、不抢占。没有同组空闲账号或仍有其他在途请求时，返回 503 `pool_lease_failover_pending`，后续请求再尝试分配。 - 到期但仍有请求执行时，租约进入待释放状态；旧请求可以结束，新请求不能继续使用该租约，其他用户也不能接管。 - 流式取消后仍等待上游生产通道结束才释放在途计数，不以断开下游连接为理由提前把池交给别人。若上游未能正常结束，需要先排查/结束旧请求，不能直接清零计数。 - 带 `previous_response_id` 的请求在已无有效租约时返回 409，要求开始新会话，避免续接到另一账号。 @@ -67,3 +67,12 @@ account-pools: 包括：并发首请求单次分配、同组多账号分配给不同用户、跨 Key 复用、请求不续期、到期在途保护、流式取消、上游失败不换池、状态恢复/缺失/损坏、文件锁、落盘失败、分组变更与失效会话，以及 New API 身份上下文传递和专用头移除。 本阶段为本地代码接入，不自动发布或修改线上配置。New API 与 CPA 两边部署并完成实例/Key 配置后才生效。 + +## 额度耗尽时的账号替换 + +- 普通请求、计数请求及流式启动失败可在安全边界自动换号并重试一次；入口也检查此前已记录的账号故障。每次执行循环最多换号一次，避免无限重试。 +- 替换必须原子落盘成功才生效，更新租约 ID 以隔离旧的会话路由缓存；用户和原到期时间保持不变。其他用户占用的账号和其他分组均不可选。 +- 旧账号仍有并发请求或被丢弃的上游流尚未结束时禁止替换，防止旧账号过早释放。已经向客户端输出的流不会自动重放；下一次完整请求可换号。 +- 携带 previous_response_id 的会话无法安全跨账号，返回 409 `pool_lease_session_expired`，需使用完整历史新建会话。当前采用保守保护：发生过换号的租约在剩余有效期内不接受 previous_response_id，避免误把旧账号响应 ID 交给新账号；完整历史请求不受影响。 +- 短暂的通用 429 仍按现有退避处理；只有明确 usage_limit_reached、insufficient_quota、usage limit 等账号额度错误才会换号。裸 403 或内容策略错误不会触发换号。 +- 此版本不改变前端 API 契约，可继续使用 v1.22.2-cpa.21。线上升级须由管理员明确执行，不随发布自动部署。 diff --git a/internal/poollease/store.go b/internal/poollease/store.go index 0dfdf7bbcae..523beed66e4 100644 --- a/internal/poollease/store.go +++ b/internal/poollease/store.go @@ -24,6 +24,7 @@ type Candidate struct { } type Lease struct { + Reassigned bool `json:"reassigned,omitempty"` Credential string `json:"credential,omitempty"` LegacyGroup bool `json:"legacy_group,omitempty"` ID string `json:"id"` @@ -276,3 +277,49 @@ func (s *Store) Snapshot(now time.Time) []Lease { }) return rows } + +// Replace moves a lease only when its sole in-flight request has stopped using +// the old credential. Rotating the ID invalidates old routing namespaces. +func (s *Store) Replace(expected Lease, candidates []Candidate, now time.Time) (Lease, func(), error) { + s.mu.Lock() + defer s.mu.Unlock() + if err := s.check(expected.Policy, now); err != nil { + return Lease{}, nil, err + } + old := s.leases[expected.ID] + if old == nil || old.Owner != expected.Owner || old.Credential != expected.Credential || old.LegacyGroup || old.Active != 1 || !now.Before(old.Expires) { + return Lease{}, nil, ErrBusy + } + for _, c := range candidates { + if c.Group != old.Group || c.Credential == "" || c.Credential == old.Credential || s.occupied(c, "") { + continue + } + id := make([]byte, 16) + if _, err := rand.Read(id); err != nil { + return Lease{}, nil, err + } + next := *old + next.ID, next.Credential, next.Reassigned = hex.EncodeToString(id), c.Credential, true + delete(s.leases, old.ID) + s.leases[next.ID] = &next + if err := s.save(); err != nil { + delete(s.leases, next.ID) + s.leases[old.ID] = old + return Lease{}, nil, err + } + return next, s.releaser(next.ID), nil + } + return Lease{}, nil, ErrBusy +} + +// Hold retains a live lease while an abandoned upstream producer is drained. +func (s *Store) Hold(id string) (func(), bool) { + s.mu.Lock() + defer s.mu.Unlock() + l := s.leases[id] + if l == nil || l.Active == 0 { + return nil, false + } + l.Active++ + return s.releaser(id), true +} diff --git a/internal/poollease/store_test.go b/internal/poollease/store_test.go index 979e0a7ad81..513866adfee 100644 --- a/internal/poollease/store_test.go +++ b/internal/poollease/store_test.go @@ -306,3 +306,58 @@ func TestStateRejectsDuplicateAccountReservations(t *testing.T) { t.Fatal("duplicate account accepted") } } + +func TestReplacementPreservesExpiryAndExclusiveOwnership(t *testing.T) { + s := testStore(t) + now := time.Now() + allowed := map[string]bool{"a": true, "b": true} + candidates := []Candidate{{"a", "one"}, {"a", "two"}, {"a", "three"}, {"b", "other"}} + old, done, err := s.Acquire("alice", "v1", allowed, candidates, now) + if err != nil { + t.Fatal(err) + } + defer done() + _, done2, _ := s.Acquire("bob", "v1", allowed, candidates, now) + defer done2() + _, parallel, _ := s.Acquire("alice", "v1", allowed, candidates, now) + if _, _, err = s.Replace(old, candidates, now); !errors.Is(err, ErrBusy) { + t.Fatal("replaced during concurrent execution", err) + } + parallel() + next, release, err := s.Replace(old, candidates, now.Add(time.Minute)) + if err != nil { + t.Fatal(err) + } + if next.Credential != "three" || next.Group != old.Group || !next.Expires.Equal(old.Expires) || next.ID == old.ID { + t.Fatalf("bad replacement %+v", next) + } + done() + if rows := s.Snapshot(now); len(rows) != 2 { + t.Fatal(rows) + } + if _, _, err = s.Replace(old, candidates, now); !errors.Is(err, ErrBusy) { + t.Fatal("stale replacement accepted") + } + if _, _, err = s.Replace(next, []Candidate{{"a", "two"}, {"b", "other"}}, now); !errors.Is(err, ErrBusy) { + t.Fatal("cross group or occupied replacement") + } + release() +} + +func TestReplacementPersistenceFailureRestoresOldLease(t *testing.T) { + s := testStore(t) + now := time.Now() + c := []Candidate{{"a", "one"}, {"a", "two"}} + old, done, _ := s.Acquire("alice", "v1", map[string]bool{"a": true}, c, now) + defer done() + original := s.path + s.path = filepath.Join(t.TempDir(), "missing", "state") + if _, _, err := s.Replace(old, c, now); err == nil { + t.Fatal("expected save failure") + } + s.path = original + rows := s.Snapshot(now) + if len(rows) != 1 || rows[0].ID != old.ID || rows[0].Credential != old.Credential { + t.Fatal(rows) + } +} diff --git a/sdk/cliproxy/auth/account_pools_test.go b/sdk/cliproxy/auth/account_pools_test.go index 8cd471953ec..84247ce1060 100644 --- a/sdk/cliproxy/auth/account_pools_test.go +++ b/sdk/cliproxy/auth/account_pools_test.go @@ -19,6 +19,7 @@ type poolCaptureExecutor struct { mu sync.Mutex ids []string failures map[string]bool + errors map[string]error } func (*poolCaptureExecutor) Identifier() string { return "pool-test" } @@ -26,6 +27,9 @@ func (e *poolCaptureExecutor) Execute(_ context.Context, a *Auth, _ coreexecutor e.mu.Lock() defer e.mu.Unlock() e.ids = append(e.ids, a.ID) + if err := e.errors[a.ID]; err != nil { + return coreexecutor.Response{}, err + } if e.failures[a.ID] { return coreexecutor.Response{}, &Error{HTTPStatus: 503, Message: "upstream unavailable"} } diff --git a/sdk/cliproxy/auth/conductor_execution.go b/sdk/cliproxy/auth/conductor_execution.go index 3d4ce90d301..c458a2ab36c 100644 --- a/sdk/cliproxy/auth/conductor_execution.go +++ b/sdk/cliproxy/auth/conductor_execution.go @@ -44,7 +44,7 @@ func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxye if errLease != nil { return cliproxyexecutor.Response{}, errLease } - defer releaseLease() + defer func() { releaseLease() }() req, opts = cliproxysession.Enrich(req, opts) normalized := m.normalizeProviders(providers) @@ -59,6 +59,7 @@ func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxye _, maxRetryCredentials, maxWait := m.retrySettings() var lastErr error + leaseReplaced := false retryModel := authSelectionModelFromOptions(opts, req.Model) for attempt := 0; ; attempt++ { resp, errExec := m.executeMixedOnce(ctx, normalized, req, opts, maxRetryCredentials) @@ -68,6 +69,17 @@ func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxye if isRequestTerminatedError(errExec) || isRequestStopError(errExec) { return cliproxyexecutor.Response{}, unwrapRequestStopError(errExec) } + if !leaseReplaced { + next, done, changed, err := m.replaceFailedPoolLease(ctx, normalized, req, opts, errExec) + if err != nil { + return cliproxyexecutor.Response{}, err + } + if changed { + releaseLease() + ctx, releaseLease, leaseReplaced = next, done, true + continue + } + } lastErr = errExec wait, shouldRetry := m.shouldRetryAfterError(errExec, attempt, normalized, retryModel, maxWait) if !shouldRetry { @@ -101,7 +113,7 @@ func (m *Manager) ExecuteCount(ctx context.Context, providers []string, req clip if errLease != nil { return cliproxyexecutor.Response{}, errLease } - defer releaseLease() + defer func() { releaseLease() }() req, opts = cliproxysession.Enrich(req, opts) normalized := m.normalizeProviders(providers) @@ -116,6 +128,7 @@ func (m *Manager) ExecuteCount(ctx context.Context, providers []string, req clip _, maxRetryCredentials, maxWait := m.retrySettings() var lastErr error + leaseReplaced := false retryModel := authSelectionModelFromOptions(opts, req.Model) for attempt := 0; ; attempt++ { resp, errExec := m.executeCountMixedOnce(ctx, normalized, req, opts, maxRetryCredentials) @@ -125,6 +138,17 @@ func (m *Manager) ExecuteCount(ctx context.Context, providers []string, req clip if isRequestTerminatedError(errExec) || isRequestStopError(errExec) { return cliproxyexecutor.Response{}, unwrapRequestStopError(errExec) } + if !leaseReplaced { + next, done, changed, err := m.replaceFailedPoolLease(ctx, normalized, req, opts, errExec) + if err != nil { + return cliproxyexecutor.Response{}, err + } + if changed { + releaseLease() + ctx, releaseLease, leaseReplaced = next, done, true + continue + } + } lastErr = errExec wait, shouldRetry := m.shouldRetryAfterError(errExec, attempt, normalized, retryModel, maxWait) if !shouldRetry { @@ -172,6 +196,7 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli _, maxRetryCredentials, maxWait := m.retrySettings() var lastErr error + leaseReplaced := false retryModel := authSelectionModelFromOptions(opts, req.Model) for attempt := 0; ; attempt++ { result, errStream := m.executeStreamMixedOnce(ctx, normalized, req, opts, maxRetryCredentials) @@ -182,6 +207,17 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli if isRequestTerminatedError(errStream) || isRequestStopError(errStream) { return nil, unwrapRequestStopError(errStream) } + if !leaseReplaced { + next, done, changed, err := m.replaceFailedPoolLease(ctx, normalized, req, opts, errStream) + if err != nil { + return nil, err + } + if changed { + releaseLease() + ctx, releaseLease, leaseReplaced = next, done, true + continue + } + } lastErr = errStream wait, shouldRetry := m.shouldRetryAfterError(errStream, attempt, normalized, retryModel, maxWait) if !shouldRetry { diff --git a/sdk/cliproxy/auth/conductor_stream.go b/sdk/cliproxy/auth/conductor_stream.go index 1fd48be0b75..980bcb90620 100644 --- a/sdk/cliproxy/auth/conductor_stream.go +++ b/sdk/cliproxy/auth/conductor_stream.go @@ -172,13 +172,13 @@ func (m *Manager) wrapStreamResult(ctx context.Context, auth *Auth, provider, re } for _, chunk := range buffered { if ok := emit(chunk); !ok { - discardStreamChunks(remaining) + m.discardLeasedStream(ctx, remaining) return } } for chunk := range remaining { if ok := emit(chunk); !ok { - discardStreamChunks(remaining) + m.discardLeasedStream(ctx, remaining) return } } @@ -302,7 +302,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi buffered, closed, bootstrapErr := readStreamBootstrap(ctx, streamResult.Chunks) if bootstrapErr != nil { if errCtx := ctx.Err(); errCtx != nil { - discardStreamChunks(streamResult.Chunks) + m.discardLeasedStream(ctx, streamResult.Chunks) return nil, errCtx } if allowRetry { @@ -316,11 +316,11 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi } } if errRefresh != nil { - discardStreamChunks(streamResult.Chunks) + m.discardLeasedStream(ctx, streamResult.Chunks) bootstrapErr = errRefresh streamResult = &cliproxyexecutor.StreamResult{} } else if okRefresh { - discardStreamChunks(streamResult.Chunks) + m.discardLeasedStream(ctx, streamResult.Chunks) auth = refreshed m.replaceHomeExecutionLifecycleAuth(execOpts.ExecutionLifecycle, auth) publishSelectedAuthMetadata(execOpts.Metadata, auth) @@ -342,7 +342,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi } if !ephemeralResult { if errCancel := claudeOAuthRequestCancellation(ctx, auth, bootstrapErr); errCancel != nil { - discardStreamChunks(streamResult.Chunks) + m.discardLeasedStream(ctx, streamResult.Chunks) return nil, errCancel } } @@ -357,7 +357,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi } applyRequestScopedActionToResult(action, okAction, &result) m.recordExecutionResult(ctx, result, auth, ephemeralResult) - discardStreamChunks(streamResult.Chunks) + m.discardLeasedStream(ctx, streamResult.Chunks) if isRequestScopedStop(action, okAction) { return nil, wrapRequestStopError(bootstrapErr) } @@ -375,7 +375,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi result.CredentialScope = true } m.recordExecutionResult(ctx, result, auth, ephemeralResult) - discardStreamChunks(streamResult.Chunks) + m.discardLeasedStream(ctx, streamResult.Chunks) return nil, bootstrapErr } if idx < len(execModels)-1 { @@ -386,7 +386,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi result.CredentialScope = true } m.recordExecutionResult(ctx, result, auth, ephemeralResult) - discardStreamChunks(streamResult.Chunks) + m.discardLeasedStream(ctx, streamResult.Chunks) lastErr = bootstrapErr if result.CredentialScope { return nil, newStreamBootstrapError(bootstrapErr, streamResult.Headers) @@ -400,7 +400,7 @@ func (m *Manager) executeStreamWithModelPool(ctx context.Context, executor Provi result.CredentialScope = true } m.recordExecutionResult(ctx, result, auth, ephemeralResult) - discardStreamChunks(streamResult.Chunks) + m.discardLeasedStream(ctx, streamResult.Chunks) return nil, newStreamBootstrapError(bootstrapErr, streamResult.Headers) } diff --git a/sdk/cliproxy/auth/pool_leases.go b/sdk/cliproxy/auth/pool_leases.go index 8220c609469..e6335b58e78 100644 --- a/sdk/cliproxy/auth/pool_leases.go +++ b/sdk/cliproxy/auth/pool_leases.go @@ -5,11 +5,13 @@ import ( "crypto/sha256" "encoding/hex" "errors" + log "github.com/sirupsen/logrus" "github.com/tidwall/gjson" "net/http" "os" "path/filepath" "sort" + "strings" "sync" "time" @@ -153,39 +155,14 @@ func (m *Manager) beginPoolLease(ctx context.Context, providers []string, req co } if previous, exists, e := store.Lookup(owner, runtimeLeaseRevision(cfg), time.Now()); e != nil { return ctx, noop, e - } else if (!exists || previous.LegacyGroup) && (gjson.GetBytes(req.Payload, "previous_response_id").String() != "" || gjson.GetBytes(opts.OriginalRequest, "previous_response_id").String() != "") { + } else if (!exists || previous.LegacyGroup || previous.Reassigned) && (gjson.GetBytes(req.Payload, "previous_response_id").String() != "" || gjson.GetBytes(opts.OriginalRequest, "previous_response_id").String() != "") { return ctx, noop, leaseError("pool_lease_session_expired", "start a new conversation after the pool lease expires", http.StatusConflict) } - groups := scope.LeaseGroups() + candidates := m.poolLeaseCandidates(ctx, scope, providers, req, opts) allowed := map[string]bool{} - for _, g := range groups { - allowed[g] = true + for _, group := range scope.LeaseGroups() { + allowed[group] = true } - candidates := []poollease.Candidate{} - providerSet := map[string]bool{} - for _, p := range m.normalizeProviders(providers) { - providerSet[p] = true - } - model := authSelectionModelFromOptions(opts, req.Model) - eligibility := authSelectionEligibilityForRequest(ctx, opts) - pinned := pinnedAuthIDFromMetadata(opts.Metadata) - m.mu.RLock() - for _, a := range m.auths { - if a == nil || a.Disabled || !providerSet[executorKeyFromAuth(a)] || !eligibility.allows(a) || (pinned != "" && a.ID != pinned) { - continue - } - group := scope.GroupForCredential(a.ID) - if !allowed[group] || !m.supportsPoolModel(a, model) { - continue - } - if available, e := getAvailableAuths([]*Auth{a}, a.Provider, model, time.Now()); e == nil && len(available) > 0 { - candidates = append(candidates, poollease.Candidate{Group: group, Credential: a.ID}) - } - } - m.mu.RUnlock() - sort.Slice(candidates, func(i, j int) bool { - return candidates[i].Group < candidates[j].Group || (candidates[i].Group == candidates[j].Group && candidates[i].Credential < candidates[j].Credential) - }) lease, release, err := store.Acquire(owner, runtimeLeaseRevision(cfg), allowed, candidates, time.Now()) if err != nil { if errors.Is(err, poollease.ErrDenied) { @@ -196,7 +173,19 @@ func (m *Manager) beginPoolLease(ctx context.Context, providers []string, req co } return ctx, noop, leaseError("pool_lease_unavailable", "lease allocation failed", 503) } - return context.WithValue(ctx, requestPoolLeaseKey{}, &requestPoolLease{owner, lease}), release, nil + ctx = context.WithValue(ctx, requestPoolLeaseKey{}, &requestPoolLease{owner, lease}) + // A previous request may have recorded a terminal failure while another + // in-flight request prevented replacement. Retry allocation at admission. + next, done, changed, err := m.replaceFailedPoolLease(ctx, providers, req, opts, nil) + if err != nil { + release() + return ctx, noop, err + } + if changed { + release() + return next, done, nil + } + return ctx, release, nil } // Keep the in-flight lease until the upstream producer closes. Cancellation stops @@ -256,3 +245,137 @@ func runtimeLeaseRevision(cfg *config.Config) string { } return "disabled" } + +func (m *Manager) poolLeaseCandidates(ctx context.Context, scope *config.AccountPoolScope, providers []string, req coreexecutor.Request, opts coreexecutor.Options) []poollease.Candidate { + groups := scope.LeaseGroups() + allowed := map[string]bool{} + for _, g := range groups { + allowed[g] = true + } + candidates := []poollease.Candidate{} + providerSet := map[string]bool{} + for _, p := range m.normalizeProviders(providers) { + providerSet[p] = true + } + model := authSelectionModelFromOptions(opts, req.Model) + eligibility := authSelectionEligibilityForRequest(ctx, opts) + pinned := pinnedAuthIDFromMetadata(opts.Metadata) + m.mu.RLock() + for _, a := range m.auths { + if a == nil || a.Disabled || !providerSet[executorKeyFromAuth(a)] || !eligibility.allows(a) || (pinned != "" && a.ID != pinned) { + continue + } + group := scope.GroupForCredential(a.ID) + if !allowed[group] || !m.supportsPoolModel(a, model) { + continue + } + if available, e := getAvailableAuths([]*Auth{a}, a.Provider, model, time.Now()); e == nil && len(available) > 0 { + candidates = append(candidates, poollease.Candidate{Group: group, Credential: a.ID}) + } + } + m.mu.RUnlock() + sort.Slice(candidates, func(i, j int) bool { + return candidates[i].Group < candidates[j].Group || (candidates[i].Group == candidates[j].Group && candidates[i].Credential < candidates[j].Credential) + }) + return candidates +} + +// Only explicit account quota/authentication failures permit replacement. +// Generic 429s, transport errors, overload and request policy errors do not. +func poolLeaseTerminalFailure(err error) bool { + if err == nil { + return false + } + var e *Error + if errors.As(err, &e) && e.Code == ErrorCodeRequestScoped { + return false + } + if IsTerminalAuthError(err) { + return true + } + status := 0 + var sc interface{ StatusCode() int } + if errors.As(err, &sc) { + status = sc.StatusCode() + } + msg := strings.ToLower(err.Error()) + if status == 429 { + return strings.Contains(msg, "usage_limit_reached") || strings.Contains(msg, "insufficient_quota") || strings.Contains(msg, "you've hit your usage limit") + } + if status == 401 { + return true + } + return status == 403 && (strings.Contains(msg, "account_deactivated") || strings.Contains(msg, "account_disabled") || strings.Contains(msg, "token_revoked")) +} + +func (m *Manager) replaceFailedPoolLease(ctx context.Context, providers []string, req coreexecutor.Request, opts coreexecutor.Options, failure error) (context.Context, func(), bool, error) { + binding, ok := ctx.Value(requestPoolLeaseKey{}).(*requestPoolLease) + if !ok || ctx.Err() != nil { + return ctx, nil, false, nil + } + terminal := poolLeaseTerminalFailure(failure) + m.mu.RLock() + if a := m.auths[binding.Lease.Credential]; a != nil { + if _, unavailable := getAvailableAuths([]*Auth{a}, a.Provider, authSelectionModelFromOptions(opts, req.Model), time.Now()); unavailable != nil || a.Disabled { + terminal = terminal || poolLeaseTerminalFailure(a.LastError) + } + } + m.mu.RUnlock() + if !terminal { + return ctx, nil, false, nil + } + if gjson.GetBytes(req.Payload, "previous_response_id").String() != "" || gjson.GetBytes(opts.OriginalRequest, "previous_response_id").String() != "" { + return ctx, nil, false, leaseError("pool_lease_session_expired", "leased account is unavailable; start a new conversation with full history", http.StatusConflict) + } + scope, err := m.AccountPoolScope(ctx) + if err != nil { + return ctx, nil, false, err + } + store, err := m.checkLeaseConfiguration(m.runtimeConfigSnapshot()) + if err != nil { + return ctx, nil, false, err + } + if store == nil || scope == nil { + return ctx, nil, false, nil + } + next, release, err := store.Replace(binding.Lease, m.poolLeaseCandidates(ctx, scope, providers, req, opts), time.Now()) + if errors.Is(err, poollease.ErrBusy) { + return ctx, nil, false, leaseError("pool_lease_failover_pending", "no free account in the leased group or previous requests are still running; retry later", 503) + } + if err != nil { + return ctx, nil, false, leaseError("pool_lease_unavailable", "cannot persist account replacement", 503) + } + log.WithFields(log.Fields{"owner": next.Owner[:10], "group": next.Group, "lease": next.ID, "expires_at": next.Expires}).Info("exclusive account lease replaced after quota or authentication failure") + return context.WithValue(ctx, requestPoolLeaseKey{}, &requestPoolLease{binding.Owner, next}), release, true, nil +} + +func (m *Manager) discardLeasedStream(ctx context.Context, chunks <-chan coreexecutor.StreamChunk) { + if chunks == nil { + return + } + if ctx == nil { + discardStreamChunks(chunks) + return + } + binding, leased := ctx.Value(requestPoolLeaseKey{}).(*requestPoolLease) + if !leased { + discardStreamChunks(chunks) + return + } + m.poolLeaseRuntime.mu.Lock() + store := m.poolLeaseRuntime.store + m.poolLeaseRuntime.mu.Unlock() + if store != nil { + if done, ok := store.Hold(binding.Lease.ID); ok { + go func() { + defer done() + for range chunks { + } + }() + return + } + } + // If retention fails, do not release the caller's lease before draining. + for range chunks { + } +} diff --git a/sdk/cliproxy/auth/pool_leases_test.go b/sdk/cliproxy/auth/pool_leases_test.go index 77a9f103a8e..2e0a7d0259d 100644 --- a/sdk/cliproxy/auth/pool_leases_test.go +++ b/sdk/cliproxy/auth/pool_leases_test.go @@ -152,3 +152,183 @@ func TestLeaseStreamFailoverAfterTargetedCooldownCannotEscape(t *testing.T) { } } } + +func TestLeaseTerminalFailureReplacesAccountInEveryMode(t *testing.T) { + for _, mode := range []string{"normal", "count", "stream"} { + t.Run(mode, func(t *testing.T) { + m, _, e := leaseManager(t) + ctx := leaseCaller("key-all", "1") + req := coreexecutor.Request{Model: "pool-model"} + if _, err := m.Execute(ctx, []string{"pool-test"}, req, coreexecutor.Options{}); err != nil { + t.Fatal(err) + } + before, _ := m.AccountPoolLeases() + e.errors = map[string]error{e.ids[0]: &Error{HTTPStatus: 429, Message: `{"error":{"type":"usage_limit_reached"}}`}} + var err error + switch mode { + case "normal": + _, err = m.Execute(ctx, []string{"pool-test"}, req, coreexecutor.Options{}) + case "count": + _, err = m.ExecuteCount(ctx, []string{"pool-test"}, req, coreexecutor.Options{}) + case "stream": + var r *coreexecutor.StreamResult + r, err = m.ExecuteStream(ctx, []string{"pool-test"}, req, coreexecutor.Options{}) + if err == nil { + for c := range r.Chunks { + if c.Err != nil { + t.Fatal(c.Err) + } + } + } + } + if err != nil { + t.Fatal(err) + } + after, _ := m.AccountPoolLeases() + if len(after) != 1 || after[0].Credential == before[0].Credential || after[0].Group != before[0].Group || after[0].ID == before[0].ID || !after[0].Expires.Equal(before[0].Expires) { + t.Fatalf("invalid replacement: %+v -> %+v", before, after) + } + _, err = m.Execute(ctx, []string{"pool-test"}, coreexecutor.Request{Model: "pool-model", Payload: []byte(`{"previous_response_id":"old"}`)}, coreexecutor.Options{}) + var pe *Error + if !errors.As(err, &pe) || pe.Code != "pool_lease_session_expired" { + t.Fatalf("unsafe continuation: %v", err) + } + }) + } +} + +func TestLeaseGenericRateLimitKeepsAccount(t *testing.T) { + m, _, e := leaseManager(t) + ctx := leaseCaller("key-all", "1") + req := coreexecutor.Request{Model: "pool-model"} + m.Execute(ctx, []string{"pool-test"}, req, coreexecutor.Options{}) + before, _ := m.AccountPoolLeases() + e.errors = map[string]error{e.ids[0]: &Error{HTTPStatus: 429, Message: "rate_limit_exceeded"}} + if _, err := m.Execute(ctx, []string{"pool-test"}, req, coreexecutor.Options{}); err == nil { + t.Fatal("expected rate limit") + } + after, _ := m.AccountPoolLeases() + if before[0].ID != after[0].ID { + t.Fatal("transient rate limit replaced account") + } +} + +func TestLeaseReplacementWaitsForOtherRequestAndResumesOnAdmission(t *testing.T) { + m, _, e := leaseManager(t) + ctx := leaseCaller("key-all", "1") + req := coreexecutor.Request{Model: "pool-model"} + held, done, err := m.beginPoolLease(ctx, []string{"pool-test"}, req, coreexecutor.Options{}) + if err != nil { + t.Fatal(err) + } + defer done() + old := held.Value(requestPoolLeaseKey{}).(*requestPoolLease).Lease + e.errors = map[string]error{old.Credential: &Error{HTTPStatus: 429, Message: "usage_limit_reached"}} + _, err = m.Execute(ctx, []string{"pool-test"}, req, coreexecutor.Options{}) + var pe *Error + if !errors.As(err, &pe) || pe.Code != "pool_lease_failover_pending" { + t.Fatal(err) + } + done() + if _, err = m.Execute(ctx, []string{"pool-test"}, req, coreexecutor.Options{}); err != nil { + t.Fatal(err) + } + after, _ := m.AccountPoolLeases() + if after[0].Credential == old.Credential { + t.Fatal("failed account was retained") + } +} + +func TestLeaseTerminalFailureClassification(t *testing.T) { + for _, tc := range []struct { + status int + msg string + want bool + }{ + {429, "usage_limit_reached", true}, {429, "You've hit your usage limit", true}, {429, "rate_limit_exceeded", false}, {401, "unauthorized", true}, {403, "account_deactivated", true}, {403, "content policy", false}, {503, "usage_limit_reached", false}, + } { + if got := poolLeaseTerminalFailure(&Error{HTTPStatus: tc.status, Message: tc.msg}); got != tc.want { + t.Fatalf("%+v = %v", tc, got) + } + } +} + +func TestLeaseDiscardedProducerBlocksReplacement(t *testing.T) { + m, _, _ := leaseManager(t) + ctx, done, err := m.beginPoolLease(leaseCaller("key-all", "1"), []string{"pool-test"}, coreexecutor.Request{Model: "pool-model"}, coreexecutor.Options{}) + if err != nil { + t.Fatal(err) + } + defer done() + old := ctx.Value(requestPoolLeaseKey{}).(*requestPoolLease).Lease + ch := make(chan coreexecutor.StreamChunk) + m.discardLeasedStream(ctx, ch) + store := m.poolLeaseRuntime.store + l, _, _ := store.Lookup(old.Owner, old.Policy, time.Now()) + if l.Active != 2 { + t.Fatal(l) + } + if _, _, err = store.Replace(old, nil, time.Now()); err == nil { + t.Fatal("replaced with active producer") + } + close(ch) + deadline := time.Now().Add(time.Second) + for time.Now().Before(deadline) { + l, _, _ = store.Lookup(old.Owner, old.Policy, time.Now()) + if l.Active == 1 { + return + } + time.Sleep(time.Millisecond) + } + t.Fatal("producer lease retention leaked") +} + +type leaseCommittedErrorExecutor struct{ *poolCaptureExecutor } + +func (e *leaseCommittedErrorExecutor) ExecuteStream(ctx context.Context, a *Auth, req coreexecutor.Request, opts coreexecutor.Options) (*coreexecutor.StreamResult, error) { + if _, err := e.Execute(ctx, a, req, opts); err != nil { + return nil, err + } + ch := make(chan coreexecutor.StreamChunk, 2) + ch <- coreexecutor.StreamChunk{Payload: []byte("already emitted")} + ch <- coreexecutor.StreamChunk{Err: &Error{HTTPStatus: 429, Message: "usage_limit_reached"}} + close(ch) + return &coreexecutor.StreamResult{Chunks: ch}, nil +} +func TestLeaseCommittedStreamIsNotReplayed(t *testing.T) { + m, _, e := leaseManager(t) + m.RegisterExecutor(&leaseCommittedErrorExecutor{e}) + req := coreexecutor.Request{Model: "pool-model"} + ctx := leaseCaller("key-all", "1") + r, err := m.ExecuteStream(ctx, []string{"pool-test"}, req, coreexecutor.Options{}) + if err != nil { + t.Fatal(err) + } + chunks := 0 + for c := range r.Chunks { + chunks++ + if chunks == 1 && c.Err != nil { + t.Fatal(c.Err) + } + } + if chunks != 2 || len(e.ids) != 1 { + t.Fatalf("stream replayed: chunks=%d ids=%v", chunks, e.ids) + } + before, _ := m.AccountPoolLeases() + deadline := time.Now().Add(time.Second) + for time.Now().Before(deadline) { + ls, _ := m.AccountPoolLeases() + if ls[0].Active == 0 { + break + } + time.Sleep(time.Millisecond) + } + m.RegisterExecutor(e) + if _, err = m.Execute(ctx, []string{"pool-test"}, req, coreexecutor.Options{}); err != nil { + t.Fatal(err) + } + after, _ := m.AccountPoolLeases() + if before[0].Credential == after[0].Credential { + t.Fatal("next request did not replace failed account") + } +}