diff --git a/docs/ACCOUNT_POOL_LEASES_CN.md b/docs/ACCOUNT_POOL_LEASES_CN.md index 2ad94a3123f..be845ced149 100644 --- a/docs/ACCOUNT_POOL_LEASES_CN.md +++ b/docs/ACCOUNT_POOL_LEASES_CN.md @@ -1,4 +1,4 @@ -# New API 与 CPA 一小时号池租约接入 +# New API 与 CPA 一小时账号租约接入 ## 配置步骤 @@ -27,18 +27,28 @@ account-pools: lease-instance: newapi-main ``` +## 从旧版整组租约升级 + +新租约锁定一个凭据 ID，而不是整个分组。例如一个分组有 8 个健康且支持请求模型的账号，最多可同时分配给 8 位用户。 + +状态文件升级为 version 2，记录 credential。version 1 旧租约不会直接丢弃：原用户下一次不带旧 previous_response_id 的请求到来，且旧租约没有在途请求时，按可用候选账号原子收窄为一个账号；保留原租约 ID、开始及到期时间，不续期。收窄失败则保留旧整组保护。原用户不再请求时，旧整组保护持续至到期。 + +迁移期间，依赖旧 previous_response_id 的续接会明确要求新会话，避免猜测旧账号。不要删除状态文件解除保护。旧二进制不支持 version 2；回滚前必须停止新流量并等待所有租约及在途结束，不能恢复过时备份覆盖活跃的账号租约。 + +管理接口返回 `lease-unit: account`；新面板发现旧后端时会提示升级，避免把整组租约误认成账号租约。 + ## 实际规则 - 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`，不排无界队列、不抢占。已绑定账号失败时不自动换号。 - 到期但仍有请求执行时，租约进入待释放状态；旧请求可以结束，新请求不能继续使用该租约，其他用户也不能接管。 - 流式取消后仍等待上游生产通道结束才释放在途计数，不以断开下游连接为理由提前把池交给别人。若上游未能正常结束，需要先排查/结束旧请求，不能直接清零计数。 - 带 `previous_response_id` 的请求在已无有效租约时返回 409，要求开始新会话，避免续接到另一账号。 - 模型列表查询不分配租约；Token 计数等需要凭据的执行路径受租约约束。 -- 界面显示的是用户摘要、到期时间与在途数；状态随页面刷新更新。配额与使用统计保持原有功能，本阶段不新增按租户复制的历史统计。 +- 界面显示的是每个已租账号的用户摘要、到期时间与在途数；状态随页面刷新更新。配额与使用统计保持原有功能，本阶段不新增按租户复制的历史统计。 ## 持久化与变更限制 @@ -54,6 +64,6 @@ account-pools: ## 验证范围 -包括：并发首请求单次分配、不同用户争用、跨 Key 复用、请求不续期、到期在途保护、流式取消、上游失败不换池、状态恢复/缺失/损坏、文件锁、落盘失败、分组变更与失效会话，以及 New API 身份上下文传递和专用头移除。 +包括：并发首请求单次分配、同组多账号分配给不同用户、跨 Key 复用、请求不续期、到期在途保护、流式取消、上游失败不换池、状态恢复/缺失/损坏、文件锁、落盘失败、分组变更与失效会话，以及 New API 身份上下文传递和专用头移除。 本阶段为本地代码接入，不自动发布或修改线上配置。New API 与 CPA 两边部署并完成实例/Key 配置后才生效。 diff --git a/internal/api/handlers/management/account_pools.go b/internal/api/handlers/management/account_pools.go index c9f0f709f05..458b9441fce 100644 --- a/internal/api/handlers/management/account_pools.go +++ b/internal/api/handlers/management/account_pools.go @@ -117,10 +117,10 @@ func (h *Handler) GetAccountPools(c *gin.Context) { leaseError = "lease state unavailable" } for _, l := range rows { - leases = append(leases, gin.H{"id": l.ID, "group-id": l.Group, "owner": l.Owner, "expires-at": l.Expires, "active": l.Active}) + leases = append(leases, gin.H{"id": l.ID, "group-id": l.Group, "credential-id": l.Credential, "legacy-group": l.LegacyGroup, "owner": l.Owner, "expires-at": l.Expires, "active": l.Active}) } } - c.JSON(http.StatusOK, gin.H{"leases": leases, "lease-error": leaseError, "config": pools, "keys": keys, "credentials": credentials, "revision": poolConfigRevision(cfg), "home-enabled": cfg.Home.Enabled}) + c.JSON(http.StatusOK, gin.H{"lease-unit": "account", "leases": leases, "lease-error": leaseError, "config": pools, "keys": keys, "credentials": credentials, "revision": poolConfigRevision(cfg), "home-enabled": cfg.Home.Enabled}) } // PutAccountPools publishes membership and key authorization as one versioned change. diff --git a/internal/config/account_pools.go b/internal/config/account_pools.go index 546b066d058..272df426fe1 100644 --- a/internal/config/account_pools.go +++ b/internal/config/account_pools.go @@ -118,13 +118,14 @@ type AccountPoolPolicy struct { invalid bool } type AccountPoolScope struct { - policy *AccountPoolPolicy - leaseInstance string - leaseGroup string - leaseID string - keyID string - all bool - groups map[string]bool + leaseCredential string + policy *AccountPoolPolicy + leaseInstance string + leaseGroup string + leaseID string + keyID string + all bool + groups map[string]bool } func (cfg *Config) CompileAccountPoolPolicy() *AccountPoolPolicy { @@ -185,7 +186,7 @@ func (s *AccountPoolScope) Allows(id string) bool { return true } group := s.policy.GroupForCredential(id) - return s.AuthorizesGroup(group) && ((!s.policy.leased[group] && s.leaseInstance == "") || s.leaseGroup == group) + return s.AuthorizesGroup(group) && ((!s.policy.leased[group] && s.leaseInstance == "") || (s.leaseGroup == group && (s.leaseCredential == "" || s.leaseCredential == id))) } func (s *AccountPoolScope) Namespace() string { return s.keyID + ":" + s.policy.revision + ":" + s.leaseID @@ -229,10 +230,11 @@ func (s *AccountPoolScope) LeaseGroups() []string { sort.Strings(result) return result } -func (s *AccountPoolScope) WithLease(group, id string) *AccountPoolScope { +func (s *AccountPoolScope) WithLease(group, id, credential string) *AccountPoolScope { copy := *s copy.leaseGroup = group copy.leaseID = id + copy.leaseCredential = credential return © } func (s *AccountPoolScope) AuthorizesCredential(id string) bool { @@ -241,3 +243,5 @@ func (s *AccountPoolScope) AuthorizesCredential(id string) bool { func (p *AccountPoolPolicy) LeaseRevision() string { return p.leaseRevision } func (p *AccountPoolPolicy) HasLeasePools() bool { return p != nil && p.hasLeases } + +func (s *AccountPoolScope) LeasedCredential() string { return s.leaseCredential } diff --git a/internal/poollease/store.go b/internal/poollease/store.go index 3c01b57119c..0dfdf7bbcae 100644 --- a/internal/poollease/store.go +++ b/internal/poollease/store.go @@ -1,4 +1,4 @@ -// Package poollease implements persistent, single-process exclusive pool leases. +// Package poollease implements persistent, single-process exclusive account leases. package poollease import ( @@ -18,14 +18,21 @@ var ErrBusy = errors.New("pool_busy") var ErrDenied = errors.New("pool_lease_access_denied") var ErrPolicy = errors.New("pool_lease_policy_changed") +type Candidate struct { + Group string + Credential string +} + type Lease struct { - ID string `json:"id"` - Owner string `json:"owner"` - Group string `json:"group"` - Policy string `json:"policy"` - Started time.Time `json:"started_at"` - Expires time.Time `json:"expires_at"` - Active int `json:"-"` + Credential string `json:"credential,omitempty"` + LegacyGroup bool `json:"legacy_group,omitempty"` + ID string `json:"id"` + Owner string `json:"owner"` + Group string `json:"group"` + Policy string `json:"policy"` + Started time.Time `json:"started_at"` + Expires time.Time `json:"expires_at"` + Active int `json:"-"` } type Store struct { mu sync.Mutex @@ -66,19 +73,37 @@ func Open(path string) (*Store, error) { Version int `json:"version"` Leases []Lease `json:"leases"` } - if err = json.Unmarshal(raw, &state); err != nil || state.Version != 1 { + if err = json.Unmarshal(raw, &state); err != nil || (state.Version != 1 && state.Version != 2) { s.Close() return nil, errors.New("invalid pool lease state") } owners := map[string]bool{} + credentials := map[string]bool{} + groups := map[string]bool{} + legacyGroups := map[string]bool{} for _, l := range state.Leases { - if l.ID == "" || l.Owner == "" || l.Group == "" || l.Policy == "" || !l.Expires.After(l.Started) || s.leases[l.Group] != nil || owners[l.Owner] { + if state.Version == 1 { + if l.Credential != "" { + s.Close() + return nil, errors.New("invalid legacy lease") + } + l.LegacyGroup = true + } + invalidResource := (l.Credential == "") != l.LegacyGroup + overlap := legacyGroups[l.Group] || (l.LegacyGroup && groups[l.Group]) || (l.Credential != "" && credentials[l.Credential]) + if l.ID == "" || l.Owner == "" || l.Group == "" || l.Policy == "" || !l.Expires.After(l.Started) || s.leases[l.ID] != nil || owners[l.Owner] || invalidResource || overlap { s.Close() return nil, errors.New("invalid pool lease record") } copy := l - s.leases[l.Group] = © + s.leases[l.ID] = © owners[l.Owner] = true + groups[l.Group] = true + if l.LegacyGroup { + legacyGroups[l.Group] = true + } else { + credentials[l.Credential] = true + } } } return s, nil @@ -118,49 +143,83 @@ func (s *Store) Lookup(owner, policy string, now time.Time) (Lease, bool, error) } return Lease{}, false, nil } -func (s *Store) Acquire(owner, policy string, allowed map[string]bool, candidates []string, now time.Time) (Lease, func(), error) { + +// occupied protects account IDs globally; legacy v1 records temporarily protect +// their entire group until the owner safely narrows the lease or it expires. +func (s *Store) occupied(candidate Candidate, ignore string) bool { + for _, l := range s.leases { + if l.ID != ignore && (l.Credential == candidate.Credential || (l.LegacyGroup && l.Group == candidate.Group)) { + return true + } + } + return false +} +func (s *Store) Acquire(owner, policy string, allowed map[string]bool, candidates []Candidate, now time.Time) (Lease, func(), error) { s.mu.Lock() defer s.mu.Unlock() if err := s.check(policy, now); err != nil { return Lease{}, nil, err } for _, l := range s.leases { - if l.Owner == owner { - if !allowed[l.Group] { - return Lease{}, nil, ErrDenied + if l.Owner != owner { + continue + } + if !allowed[l.Group] { + return Lease{}, nil, ErrDenied + } + if !now.Before(l.Expires) { + return Lease{}, nil, ErrBusy + } + if l.LegacyGroup { + if l.Active != 0 { + return Lease{}, nil, ErrBusy } - if !now.Before(l.Expires) { + chosen := "" + for _, c := range candidates { + if c.Group == l.Group && c.Credential != "" && !s.occupied(c, l.ID) { + chosen = c.Credential + break + } + } + if chosen == "" { return Lease{}, nil, ErrBusy } - l.Active++ - return *l, s.releaser(l.Group, l.ID), nil + l.Credential = chosen + l.LegacyGroup = false + if err := s.save(); err != nil { + l.Credential = "" + l.LegacyGroup = true + return Lease{}, nil, err + } } + l.Active++ + return *l, s.releaser(l.ID), nil } - for _, group := range candidates { - if !allowed[group] || s.leases[group] != nil { + for _, c := range candidates { + if c.Credential == "" || !allowed[c.Group] || s.occupied(c, "") { continue } id := make([]byte, 16) if _, err := rand.Read(id); err != nil { return Lease{}, nil, err } - l := &Lease{ID: hex.EncodeToString(id), Owner: owner, Group: group, Policy: policy, Started: now, Expires: now.Add(time.Hour), Active: 1} - s.leases[group] = l + l := &Lease{ID: hex.EncodeToString(id), Owner: owner, Group: c.Group, Credential: c.Credential, Policy: policy, Started: now, Expires: now.Add(time.Hour), Active: 1} + s.leases[l.ID] = l if err := s.save(); err != nil { - delete(s.leases, group) + delete(s.leases, l.ID) return Lease{}, nil, err } - return *l, s.releaser(group, l.ID), nil + return *l, s.releaser(l.ID), nil } return Lease{}, nil, ErrBusy } -func (s *Store) releaser(group, id string) func() { +func (s *Store) releaser(id string) func() { var once sync.Once return func() { once.Do(func() { s.mu.Lock() defer s.mu.Unlock() - if l := s.leases[group]; l != nil && l.ID == id { + if l := s.leases[id]; l != nil { l.Active-- } }) @@ -171,11 +230,13 @@ func (s *Store) save() error { for _, l := range s.leases { rows = append(rows, *l) } - sort.Slice(rows, func(i, j int) bool { return rows[i].Group < rows[j].Group }) + sort.Slice(rows, func(i, j int) bool { + return rows[i].Group < rows[j].Group || (rows[i].Group == rows[j].Group && rows[i].ID < rows[j].ID) + }) raw, err := json.Marshal(struct { Version int `json:"version"` Leases []Lease `json:"leases"` - }{1, rows}) + }{2, rows}) if err != nil { return err } @@ -210,6 +271,8 @@ func (s *Store) Snapshot(now time.Time) []Lease { for _, l := range s.leases { rows = append(rows, *l) } - sort.Slice(rows, func(i, j int) bool { return rows[i].Group < rows[j].Group }) + sort.Slice(rows, func(i, j int) bool { + return rows[i].Group < rows[j].Group || (rows[i].Group == rows[j].Group && rows[i].ID < rows[j].ID) + }) return rows } diff --git a/internal/poollease/store_test.go b/internal/poollease/store_test.go index aa2d88a0016..979e0a7ad81 100644 --- a/internal/poollease/store_test.go +++ b/internal/poollease/store_test.go @@ -1,7 +1,9 @@ package poollease import ( + "encoding/json" "errors" + "fmt" "os" "path/filepath" "sync" @@ -28,7 +30,7 @@ func TestConcurrentUsersAndFixedExpiry(t *testing.T) { wg.Add(1) go func() { defer wg.Done() - l, done, e := s.Acquire("alice", "v1", allowed, []string{"a", "b"}, now) + l, done, e := s.Acquire("alice", "v1", allowed, testCandidates("a", "b"), now) if e != nil { t.Error(e) return @@ -46,7 +48,7 @@ func TestConcurrentUsersAndFixedExpiry(t *testing.T) { } id = got } - l, done, e := s.Acquire("alice", "v1", allowed, []string{"b"}, now.Add(30*time.Minute)) + l, done, e := s.Acquire("alice", "v1", allowed, testCandidates("b"), now.Add(30*time.Minute)) if e != nil { t.Fatal(e) } @@ -54,7 +56,7 @@ func TestConcurrentUsersAndFixedExpiry(t *testing.T) { t.Fatal("request renewed the fixed lease") } done() - b, done, e := s.Acquire("bob", "v1", allowed, []string{"a", "b"}, now) + b, done, e := s.Acquire("bob", "v1", allowed, testCandidates("a", "b"), now) if e != nil { t.Fatal(e) } @@ -62,7 +64,7 @@ func TestConcurrentUsersAndFixedExpiry(t *testing.T) { t.Fatal("users shared pool") } done() - _, _, e = s.Acquire("charlie", "v1", allowed, []string{"a", "b"}, now) + _, _, e = s.Acquire("charlie", "v1", allowed, testCandidates("a", "b"), now) if !errors.Is(e, ErrBusy) { t.Fatal(e) } @@ -71,17 +73,17 @@ func TestExpiredInFlightLeaseCannotBeReassigned(t *testing.T) { s := testStore(t) now := time.Now() allowed := map[string]bool{"a": true} - _, done, e := s.Acquire("alice", "v1", allowed, []string{"a"}, now) + _, done, e := s.Acquire("alice", "v1", allowed, testCandidates("a"), now) if e != nil { t.Fatal(e) } - _, _, e = s.Acquire("bob", "v1", allowed, []string{"a"}, now.Add(61*time.Minute)) + _, _, e = s.Acquire("bob", "v1", allowed, testCandidates("a"), now.Add(61*time.Minute)) if !errors.Is(e, ErrBusy) { t.Fatal("reassigned in-flight pool", e) } done() done() - l, release, e := s.Acquire("bob", "v1", allowed, []string{"a"}, now.Add(61*time.Minute)) + l, release, e := s.Acquire("bob", "v1", allowed, testCandidates("a"), now.Add(61*time.Minute)) if e != nil { t.Fatal(e) } @@ -97,7 +99,7 @@ func TestRecoveryLockAndPolicyProtection(t *testing.T) { t.Fatal(e) } now := time.Now() - l, done, e := s.Acquire("alice", "v1", map[string]bool{"a": true}, []string{"a"}, now) + l, done, e := s.Acquire("alice", "v1", map[string]bool{"a": true}, testCandidates("a"), now) if e != nil { t.Fatal(e) } @@ -139,7 +141,7 @@ func TestCorruptStateAndFailedPersistenceFailClosed(t *testing.T) { defer s.Close() os.Remove(path) os.Mkdir(path, 0700) - _, _, e = s.Acquire("alice", "v1", map[string]bool{"a": true}, []string{"a"}, time.Now()) + _, _, e = s.Acquire("alice", "v1", map[string]bool{"a": true}, testCandidates("a"), time.Now()) if e == nil { t.Fatal("persistence failure accepted") } @@ -163,3 +165,144 @@ func TestMissingStateIsNotTreatedAsAnEmptyPool(t *testing.T) { t.Fatal("missing state reset all leases") } } + +func testCandidates(groups ...string) []Candidate { + out := []Candidate{} + for _, g := range groups { + out = append(out, Candidate{Group: g, Credential: "account-" + g}) + } + return out +} + +func TestEightAccountsInOneGroupServeEightUsers(t *testing.T) { + s := testStore(t) + now := time.Now() + allowed := map[string]bool{"pool": true} + candidates := make([]Candidate, 8) + for i := range candidates { + candidates[i] = Candidate{Group: "pool", Credential: fmt.Sprint("account-", i)} + } + var wg sync.WaitGroup + results := make(chan Lease, 8) + for i := 0; i < 8; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + l, done, err := s.Acquire(fmt.Sprint("user-", i), "v1", allowed, candidates, now) + if err != nil { + t.Error(err) + return + } + defer done() + results <- l + }(i) + } + wg.Wait() + close(results) + seen := map[string]bool{} + for l := range results { + if l.Credential == "" || seen[l.Credential] || l.Group != "pool" { + t.Fatal("account shared", l) + } + seen[l.Credential] = true + } + if len(seen) != 8 { + t.Fatal(len(seen)) + } + if _, _, err := s.Acquire("ninth", "v1", allowed, candidates, now); !errors.Is(err, ErrBusy) { + t.Fatal(err) + } + before, err := os.Stat(s.path) + if err != nil { + t.Fatal(err) + } + l, done, err := s.Acquire("user-0", "v1", allowed, nil, now.Add(30*time.Minute)) + if err != nil { + t.Fatal(err) + } + done() + after, err := os.Stat(s.path) + if err != nil { + t.Fatal(err) + } + if !os.SameFile(before, after) || !l.Expires.Equal(now.Add(time.Hour)) { + t.Fatal("reuse rewrote or renewed lease") + } +} + +func TestLegacyLeaseNarrowsWithoutRenewalAndPersists(t *testing.T) { + path := filepath.Join(t.TempDir(), "leases") + now := time.Now().UTC() + legacy := Lease{ID: "old", Owner: "alice", Group: "pool", Policy: "v1", Started: now, Expires: now.Add(time.Hour)} + raw, _ := json.Marshal(struct { + Version int `json:"version"` + Leases []Lease `json:"leases"` + }{1, []Lease{legacy}}) + if err := os.WriteFile(path, raw, 0600); err != nil { + t.Fatal(err) + } + s, err := Open(path) + if err != nil { + t.Fatal(err) + } + allowed := map[string]bool{"pool": true} + candidates := []Candidate{{"pool", "one"}, {"pool", "two"}} + if _, _, err = s.Acquire("bob", "v1", allowed, candidates, now); !errors.Is(err, ErrBusy) { + t.Fatal("legacy pool exposed", err) + } + // Failed narrowing must not free the legacy reservation. + backup := s.path + s.path = filepath.Join(t.TempDir(), "missing", "state") + if _, _, err = s.Acquire("alice", "v1", allowed, candidates, now); err == nil { + t.Fatal("ignored persistence failure") + } + s.path = backup + l, ok, err := s.Lookup("alice", "v1", now) + if err != nil || !ok || !l.LegacyGroup || l.Credential != "" { + t.Fatal(l, err) + } + l, done, err := s.Acquire("alice", "v1", allowed, candidates, now.Add(10*time.Minute)) + if err != nil { + t.Fatal(err) + } + done() + if l.Credential != "one" || l.LegacyGroup || l.ID != legacy.ID || !l.Expires.Equal(legacy.Expires) { + t.Fatal(l) + } + b, done, err := s.Acquire("bob", "v1", allowed, candidates, now) + if err != nil { + t.Fatal(err) + } + done() + if b.Credential != "two" { + t.Fatal(b) + } + s.Close() + s, err = Open(path) + if err != nil { + t.Fatal(err) + } + defer s.Close() + l, ok, err = s.Lookup("alice", "v1", now) + if err != nil || !ok || l.Credential != "one" { + t.Fatal(l, err) + } +} + +func TestStateRejectsDuplicateAccountReservations(t *testing.T) { + path := filepath.Join(t.TempDir(), "leases") + now := time.Now().UTC() + a := Lease{ID: "a", Owner: "alice", Group: "pool", Credential: "same", Policy: "v1", Started: now, Expires: now.Add(time.Hour)} + b := a + b.ID = "b" + b.Owner = "bob" + raw, _ := json.Marshal(struct { + Version int `json:"version"` + Leases []Lease `json:"leases"` + }{2, []Lease{a, b}}) + os.WriteFile(path, raw, 0600) + if s, err := Open(path); err == nil { + s.Close() + t.Fatal("duplicate account accepted") + } +} diff --git a/sdk/cliproxy/auth/account_pools.go b/sdk/cliproxy/auth/account_pools.go index f5fb96ad7ef..28e7139caf8 100644 --- a/sdk/cliproxy/auth/account_pools.go +++ b/sdk/cliproxy/auth/account_pools.go @@ -81,6 +81,9 @@ func (m *Manager) CheckAccountPoolAccess(ctx context.Context, authID string) err func poolUnavailable(ctx context.Context) *Error { if scope, _ := ctx.Value(accountPoolScopeContextKey{}).(*config.AccountPoolScope); scope != nil { + if scope.LeasedCredential() != "" { + return &Error{Code: "leased_account_unavailable", Message: "the leased account cannot currently serve this request", HTTPStatus: http.StatusServiceUnavailable} + } return &Error{Code: "pool_unavailable", Message: "no available credential in the authorized account groups", HTTPStatus: http.StatusServiceUnavailable} } return &Error{Code: "auth_not_found", Message: "no auth available"} diff --git a/sdk/cliproxy/auth/pool_leases.go b/sdk/cliproxy/auth/pool_leases.go index 0ff41288ced..8220c609469 100644 --- a/sdk/cliproxy/auth/pool_leases.go +++ b/sdk/cliproxy/auth/pool_leases.go @@ -127,7 +127,7 @@ func (m *Manager) scopeWithLease(ctx context.Context, cfg *config.Config, scope if !time.Now().Before(lease.Expires) { return nil, leaseError("pool_lease_expired", "pool lease expired; finish the current request before allocating again", 503) } - return scope.WithLease(lease.Group, lease.ID), nil + return scope.WithLease(lease.Group, lease.ID, lease.Credential), nil } func (m *Manager) beginPoolLease(ctx context.Context, providers []string, req coreexecutor.Request, opts coreexecutor.Options) (context.Context, func(), error) { @@ -151,9 +151,9 @@ func (m *Manager) beginPoolLease(ctx context.Context, providers []string, req co if store == nil { return ctx, noop, leaseError("pool_unavailable", "no lease pools configured", 503) } - if _, exists, e := store.Lookup(owner, runtimeLeaseRevision(cfg), time.Now()); e != nil { + if previous, exists, e := store.Lookup(owner, runtimeLeaseRevision(cfg), time.Now()); e != nil { return ctx, noop, e - } else if !exists && (gjson.GetBytes(req.Payload, "previous_response_id").String() != "" || gjson.GetBytes(opts.OriginalRequest, "previous_response_id").String() != "") { + } else if (!exists || previous.LegacyGroup) && (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() @@ -161,7 +161,7 @@ func (m *Manager) beginPoolLease(ctx context.Context, providers []string, req co for _, g := range groups { allowed[g] = true } - ready := map[string]bool{} + candidates := []poollease.Candidate{} providerSet := map[string]bool{} for _, p := range m.normalizeProviders(providers) { providerSet[p] = true @@ -179,22 +179,20 @@ func (m *Manager) beginPoolLease(ctx context.Context, providers []string, req co continue } if available, e := getAvailableAuths([]*Auth{a}, a.Provider, model, time.Now()); e == nil && len(available) > 0 { - ready[group] = true + candidates = append(candidates, poollease.Candidate{Group: group, Credential: a.ID}) } } m.mu.RUnlock() - candidates := []string{} - for g := range ready { - candidates = append(candidates, g) - } - sort.Strings(candidates) + 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) { return ctx, noop, leaseError("pool_access_denied", "current lease is outside key permissions", 403) } if errors.Is(err, poollease.ErrBusy) { - return ctx, noop, leaseError("pool_busy", "no free authorized pool; retry later", 503) + return ctx, noop, leaseError("pool_busy", "no free authorized account for the requested model; retry later", 503) } return ctx, noop, leaseError("pool_lease_unavailable", "lease allocation failed", 503) } diff --git a/sdk/cliproxy/auth/pool_leases_test.go b/sdk/cliproxy/auth/pool_leases_test.go index 876c1f68bd9..77a9f103a8e 100644 --- a/sdk/cliproxy/auth/pool_leases_test.go +++ b/sdk/cliproxy/auth/pool_leases_test.go @@ -44,11 +44,14 @@ func TestLeaseIdentityAndPoolIsolationAcrossKeys(t *testing.T) { } } policy := m.runtimeConfigSnapshot().AccountPoolPolicy - if policy.GroupForCredential(e.ids[0]) != policy.GroupForCredential(e.ids[1]) || policy.GroupForCredential(e.ids[0]) == policy.GroupForCredential(e.ids[2]) { + if e.ids[0] != e.ids[1] || e.ids[0] == e.ids[2] || policy.GroupForCredential(e.ids[0]) != policy.GroupForCredential(e.ids[2]) { t.Fatal("lease isolation failed", e.ids) } - if _, err := m.Execute(leaseCaller("key-all", "3"), []string{"pool-test"}, req, opts); err == nil { - t.Fatal("third user stole occupied pool") + if _, err := m.Execute(leaseCaller("key-all", "3"), []string{"pool-test"}, req, opts); err != nil { + t.Fatal(err) + } + if _, err := m.Execute(leaseCaller("key-all", "4"), []string{"pool-test"}, req, opts); err == nil { + t.Fatal("fourth user stole occupied account") } if _, err := m.Execute(poolCaller("key-a"), []string{"pool-test"}, req, opts); err == nil { t.Fatal("static key bypassed lease") @@ -61,14 +64,13 @@ func TestLeaseRetryCannotEscapeAndConfigChangeBlocks(t *testing.T) { if _, err := m.Execute(ctx, []string{"pool-test"}, req, coreexecutor.Options{}); err != nil { t.Fatal(err) } - for _, id := range c.AccountPools.Groups[1].CredentialIDs { - e.failures[id] = true - } + leasedID := e.ids[0] + e.failures[leasedID] = true if _, err := m.Execute(ctx, []string{"pool-test"}, req, coreexecutor.Options{}); err == nil { t.Fatal("failure escaped pool") } for _, id := range e.ids { - if m.runtimeConfigSnapshot().AccountPoolPolicy.GroupForCredential(id) == "b" { + if id != leasedID { t.Fatal("used other pool") } }