From 9f29ee17e9d8451e7529a797279284cbbc0d9f2c Mon Sep 17 00:00:00 2001 From: nanashiwang Date: Sun, 13 Sep 2026 18:19:06 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20New=20API=20=E7=94=A8=E6=88=B71?= =?UTF-8?q?=E6=94=B9=E4=B8=BA=E8=AF=B7=E6=B1=82=E6=9C=9F=E9=97=B4=E4=B8=B4?= =?UTF-8?q?=E6=97=B6=E5=8D=A0=E7=94=A8=E8=B4=A6=E5=8F=B7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs/ACCOUNT_POOL_LEASES_CN.md | 9 + .../api/handlers/management/account_pools.go | 2 +- internal/poollease/store.go | 85 +++++++++- internal/poollease/store_test.go | 114 +++++++++++++ sdk/cliproxy/auth/pool_leases.go | 28 ++- sdk/cliproxy/auth/pool_leases_test.go | 159 ++++++++++++++++-- 6 files changed, 374 insertions(+), 23 deletions(-) diff --git a/docs/ACCOUNT_POOL_LEASES_CN.md b/docs/ACCOUNT_POOL_LEASES_CN.md index e518be0b2da..c8c8ba3b641 100644 --- a/docs/ACCOUNT_POOL_LEASES_CN.md +++ b/docs/ACCOUNT_POOL_LEASES_CN.md @@ -76,3 +76,12 @@ account-pools: - 携带 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。线上升级须由管理员明确执行，不随发布自动部署。 + +## New API 用户 ID 为 1 的临时调用 + +- 仅对绑定了 New API 实例的租约 Key 生效。先校验可信身份的实例与 Key 配置一致，再对字符串精确等于 `1` 的用户 ID 使用临时占用；普通 Key、其他用户或错误实例不能借此绕过分组权限。 +- ID 为 1 的所有调用均适用，包括管理员正常聊天；不是只识别渠道测试。其他管理员 ID 的测试仍采用一小时租约。 +- 每个请求原子选择授权组内一个空闲账号，临时占用参与同一份账号排他判断，不能抢占普通用户或其他管理员请求使用的账号。同一管理员并发请求也不共用账号；无空闲账号返回 pool_busy。 +- 临时占用不写入持久化租约文件，不设一小时有效期。请求结束即释放，取消的流式请求仍等上游生产者排空再释放，避免提前复用。面板在请求期间显示“临时占用”。 +- 升级前 ID 为 1 的旧一小时租约在其下一次有可用账号的新请求时转换。先等待旧在途归零，再原子持久化移除旧长期租约；写入失败保留原状态，不能清空后继续分配。其他用户的租约不变。 +- 因为请求间不保证账号相同，ID 为 1 的调用不支持 previous_response_id，返回 409 pool_temporary_session_unsupported，需携带完整对话历史。临时调用的额度耗尽换号沿用同组空闲账号规则，换号后也只占用到请求结束。 diff --git a/internal/api/handlers/management/account_pools.go b/internal/api/handlers/management/account_pools.go index 458b9441fce..e315f8747d5 100644 --- a/internal/api/handlers/management/account_pools.go +++ b/internal/api/handlers/management/account_pools.go @@ -117,7 +117,7 @@ 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, "credential-id": l.Credential, "legacy-group": l.LegacyGroup, "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, "temporary": l.Temporary, "owner": l.Owner, "expires-at": l.Expires, "active": l.Active}) } } 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}) diff --git a/internal/poollease/store.go b/internal/poollease/store.go index 523beed66e4..083aea549d8 100644 --- a/internal/poollease/store.go +++ b/internal/poollease/store.go @@ -24,6 +24,7 @@ type Candidate struct { } type Lease struct { + Temporary bool `json:"-"` Reassigned bool `json:"reassigned,omitempty"` Credential string `json:"credential,omitempty"` LegacyGroup bool `json:"legacy_group,omitempty"` @@ -138,7 +139,7 @@ func (s *Store) Lookup(owner, policy string, now time.Time) (Lease, bool, error) return Lease{}, false, err } for _, l := range s.leases { - if l.Owner == owner { + if l.Owner == owner && !l.Temporary { return *l, true, nil } } @@ -162,7 +163,7 @@ func (s *Store) Acquire(owner, policy string, allowed map[string]bool, candidate return Lease{}, nil, err } for _, l := range s.leases { - if l.Owner != owner { + if l.Owner != owner || l.Temporary { continue } if !allowed[l.Group] { @@ -222,6 +223,9 @@ func (s *Store) releaser(id string) func() { defer s.mu.Unlock() if l := s.leases[id]; l != nil { l.Active-- + if l.Temporary && l.Active == 0 { + delete(s.leases, id) + } } }) } @@ -229,6 +233,9 @@ func (s *Store) releaser(id string) func() { func (s *Store) save() error { rows := make([]Lease, 0, len(s.leases)) for _, l := range s.leases { + if l.Temporary { + continue + } rows = append(rows, *l) } sort.Slice(rows, func(i, j int) bool { @@ -287,7 +294,7 @@ func (s *Store) Replace(expected Lease, candidates []Candidate, now time.Time) ( 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) { + if old == nil || old.Owner != expected.Owner || old.Credential != expected.Credential || old.LegacyGroup || old.Active != 1 || (!old.Temporary && !now.Before(old.Expires)) { return Lease{}, nil, ErrBusy } for _, c := range candidates { @@ -302,10 +309,12 @@ func (s *Store) Replace(expected Lease, candidates []Candidate, now time.Time) ( 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 + if !old.Temporary { + 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 } @@ -323,3 +332,65 @@ func (s *Store) Hold(id string) (func(), bool) { l.Active++ return s.releaser(id), true } + +// LookupID binds temporary requests independently, even for the same owner. +func (s *Store) LookupID(id, policy string, now time.Time) (Lease, bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + if err := s.check(policy, now); err != nil { + return Lease{}, false, err + } + l := s.leases[id] + if l == nil { + return Lease{}, false, nil + } + return *l, true, nil +} + +// AcquireTemporary reserves one free credential for a single request. Temporary +// reservations share the exclusion map but never enter persistent lease state. +func (s *Store) AcquireTemporary(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 + } + var old *Lease + for _, l := range s.leases { + if l.Owner == owner && !l.Temporary { + if !allowed[l.Group] { + return Lease{}, nil, ErrDenied + } + if l.Active != 0 { + return Lease{}, nil, ErrBusy + } + old = l + break + } + } + ignore := "" + if old != nil { + ignore = old.ID + } + for _, c := range candidates { + if c.Credential == "" || !allowed[c.Group] || s.occupied(c, ignore) { + continue + } + id := make([]byte, 16) + if _, err := rand.Read(id); err != nil { + return Lease{}, nil, err + } + next := &Lease{ID: hex.EncodeToString(id), Owner: owner, Group: c.Group, Credential: c.Credential, Policy: policy, Started: now, Expires: now, Active: 1, Temporary: true} + s.leases[next.ID] = next + if old != nil { + delete(s.leases, old.ID) + 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 +} diff --git a/internal/poollease/store_test.go b/internal/poollease/store_test.go index 513866adfee..95304abbad7 100644 --- a/internal/poollease/store_test.go +++ b/internal/poollease/store_test.go @@ -361,3 +361,117 @@ func TestReplacementPersistenceFailureRestoresOldLease(t *testing.T) { t.Fatal(rows) } } + +func TestTemporaryReservationsAreExclusiveAndNeverPersisted(t *testing.T) { + s := testStore(t) + now := time.Now() + allowed := map[string]bool{"a": true} + c := []Candidate{{"a", "one"}, {"a", "two"}, {"a", "three"}} + fixed, releaseFixed, err := s.Acquire("ordinary", "policy", allowed, c, now) + if err != nil { + t.Fatal(err) + } + defer releaseFixed() + before, _ := os.ReadFile(s.path) + first, done, err := s.AcquireTemporary("admin", "policy", allowed, c, now) + if err != nil { + t.Fatal(err) + } + second, done2, err := s.AcquireTemporary("admin", "policy", allowed, c, now) + if err != nil { + t.Fatal(err) + } + if !first.Temporary || first.Credential == fixed.Credential || second.Credential == first.Credential { + t.Fatal(first, second, fixed) + } + if _, _, err = s.Acquire("other", "policy", allowed, c, now); !errors.Is(err, ErrBusy) { + t.Fatal("stole temporary account", err) + } + if _, _, err = s.AcquireTemporary("admin", "policy", allowed, c, now); !errors.Is(err, ErrBusy) { + t.Fatal("stole occupied account", err) + } + if _, ok, _ := s.Lookup("admin", "policy", now); ok { + t.Fatal("temporary reservation became user lease") + } + hold, ok := s.Hold(first.ID) + if !ok { + t.Fatal("cannot hold temporary") + } + done() + if _, ok, _ := s.LookupID(first.ID, "policy", now.Add(2*time.Hour)); !ok { + t.Fatal("released before producer completed") + } + hold() + done2() + done() + if _, ok, _ := s.LookupID(first.ID, "policy", now); ok { + t.Fatal("temporary leaked") + } + after, _ := os.ReadFile(s.path) + if string(before) != string(after) { + t.Fatal("temporary state written to disk") + } + if rows := s.Snapshot(now); len(rows) != 1 || rows[0].ID != fixed.ID { + t.Fatal(rows) + } +} + +func TestTemporaryMigrationAndSaveFailure(t *testing.T) { + s := testStore(t) + now := time.Now() + a := map[string]bool{"a": true} + c := []Candidate{{"a", "one"}} + old, done, _ := s.Acquire("admin", "policy", a, c, now) + if _, _, err := s.AcquireTemporary("admin", "policy", a, c, now); !errors.Is(err, ErrBusy) { + t.Fatal("migrated active lease", err) + } + done() + original := s.path + s.path = filepath.Join(t.TempDir(), "missing", "state") + if _, _, err := s.AcquireTemporary("admin", "policy", a, c, now); err == nil { + t.Fatal("migration ignored persistence failure") + } + s.path = original + if l, ok, _ := s.Lookup("admin", "policy", now); !ok || l.ID != old.ID { + t.Fatal("lost original lease") + } + next, release, err := s.AcquireTemporary("admin", "policy", a, c, now) + if err != nil { + t.Fatal(err) + } + if !next.Temporary || next.Credential != old.Credential { + t.Fatal(next) + } + var state struct{ Leases []Lease } + raw, _ := os.ReadFile(s.path) + json.Unmarshal(raw, &state) + if len(state.Leases) != 0 { + t.Fatal("old administrator lease still persisted") + } + release() + if len(s.Snapshot(now)) != 0 { + t.Fatal("temporary lease remains") + } +} + +func TestTemporaryReplacementStaysTemporary(t *testing.T) { + s := testStore(t) + now := time.Now() + c := []Candidate{{"a", "one"}, {"a", "two"}} + old, done, err := s.AcquireTemporary("admin", "p", map[string]bool{"a": true}, c, now) + if err != nil { + t.Fatal(err) + } + next, release, err := s.Replace(old, c, now.Add(2*time.Hour)) + if err != nil { + t.Fatal(err) + } + if !next.Temporary || next.Credential == old.Credential { + t.Fatal(next) + } + done() + release() + if len(s.Snapshot(now)) != 0 { + t.Fatal("replaced temporary reservation leaked") + } +} diff --git a/sdk/cliproxy/auth/pool_leases.go b/sdk/cliproxy/auth/pool_leases.go index e6335b58e78..de54545760a 100644 --- a/sdk/cliproxy/auth/pool_leases.go +++ b/sdk/cliproxy/auth/pool_leases.go @@ -112,12 +112,20 @@ func (m *Manager) scopeWithLease(ctx context.Context, cfg *config.Config, scope if err != nil { return scope, nil } - lease, ok, err := store.Lookup(owner, runtimeLeaseRevision(cfg), time.Now()) + binding, _ := ctx.Value(requestPoolLeaseKey{}).(*requestPoolLease) + var lease poollease.Lease + var ok bool + if binding != nil { + lease, ok, err = store.LookupID(binding.Lease.ID, runtimeLeaseRevision(cfg), time.Now()) + } else if sdkaccess.GetGatewayIdentity(ctx).User == "1" { + return scope, nil + } else { + lease, ok, err = store.Lookup(owner, runtimeLeaseRevision(cfg), time.Now()) + } if err != nil { return nil, err } - binding, _ := ctx.Value(requestPoolLeaseKey{}).(*requestPoolLease) - if binding != nil && (!ok || binding.Owner != owner || binding.Lease.ID != lease.ID) { + if binding != nil && (!ok || binding.Owner != owner || lease.Owner != owner || binding.Lease.ID != lease.ID) { return nil, leaseError("pool_lease_expired", "pool lease is no longer valid", 403) } if !ok { @@ -126,7 +134,7 @@ func (m *Manager) scopeWithLease(ctx context.Context, cfg *config.Config, scope if !scope.AuthorizesGroup(lease.Group) { return nil, leaseError("pool_access_denied", "current lease is outside key permissions", 403) } - if !time.Now().Before(lease.Expires) { + if !lease.Temporary && !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, lease.Credential), nil @@ -153,6 +161,10 @@ 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) } + temporary := sdkaccess.GetGatewayIdentity(ctx).User == "1" + if temporary && (gjson.GetBytes(req.Payload, "previous_response_id").String() != "" || gjson.GetBytes(opts.OriginalRequest, "previous_response_id").String() != "") { + return ctx, noop, leaseError("pool_temporary_session_unsupported", "temporary account access requires full conversation history; previous_response_id is not supported", http.StatusConflict) + } if previous, exists, e := store.Lookup(owner, runtimeLeaseRevision(cfg), time.Now()); e != nil { return ctx, noop, e } else if (!exists || previous.LegacyGroup || previous.Reassigned) && (gjson.GetBytes(req.Payload, "previous_response_id").String() != "" || gjson.GetBytes(opts.OriginalRequest, "previous_response_id").String() != "") { @@ -163,7 +175,13 @@ func (m *Manager) beginPoolLease(ctx context.Context, providers []string, req co for _, group := range scope.LeaseGroups() { allowed[group] = true } - lease, release, err := store.Acquire(owner, runtimeLeaseRevision(cfg), allowed, candidates, time.Now()) + var lease poollease.Lease + var release func() + if temporary { + lease, release, err = store.AcquireTemporary(owner, runtimeLeaseRevision(cfg), allowed, candidates, time.Now()) + } else { + 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) diff --git a/sdk/cliproxy/auth/pool_leases_test.go b/sdk/cliproxy/auth/pool_leases_test.go index 2e0a7d0259d..2de575e03e3 100644 --- a/sdk/cliproxy/auth/pool_leases_test.go +++ b/sdk/cliproxy/auth/pool_leases_test.go @@ -38,7 +38,7 @@ func TestLeaseIdentityAndPoolIsolationAcrossKeys(t *testing.T) { if _, err := m.Execute(poolCaller("key-all"), []string{"pool-test"}, req, opts); err == nil { t.Fatal("accepted missing identity") } - for _, c := range []context.Context{leaseCaller("key-all", "1"), leaseCaller("key-multi", "1"), leaseCaller("key-all", "2")} { + for _, c := range []context.Context{leaseCaller("key-all", "101"), leaseCaller("key-multi", "101"), leaseCaller("key-all", "2")} { if _, err := m.Execute(c, []string{"pool-test"}, req, opts); err != nil { t.Fatal(err) } @@ -60,7 +60,7 @@ func TestLeaseIdentityAndPoolIsolationAcrossKeys(t *testing.T) { func TestLeaseRetryCannotEscapeAndConfigChangeBlocks(t *testing.T) { m, c, e := leaseManager(t) req := coreexecutor.Request{Model: "pool-model"} - ctx := leaseCaller("key-all", "1") + ctx := leaseCaller("key-all", "101") if _, err := m.Execute(ctx, []string{"pool-test"}, req, coreexecutor.Options{}); err != nil { t.Fatal(err) } @@ -82,7 +82,7 @@ func TestLeaseRetryCannotEscapeAndConfigChangeBlocks(t *testing.T) { } func TestLeaseStreamRetainsInFlightUntilProducerCloses(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{}) + ctx, done, err := m.beginPoolLease(leaseCaller("key-all", "101"), []string{"pool-test"}, coreexecutor.Request{Model: "pool-model"}, coreexecutor.Options{}) if err != nil { t.Fatal(err) } @@ -112,7 +112,7 @@ func TestLeaseStreamRetainsInFlightUntilProducerCloses(t *testing.T) { } func TestExpiredContinuationDoesNotAllocateAnotherPool(t *testing.T) { m, _, _ := leaseManager(t) - _, err := m.Execute(leaseCaller("key-all", "1"), []string{"pool-test"}, coreexecutor.Request{Model: "pool-model", Payload: []byte(`{"previous_response_id":"old"}`)}, coreexecutor.Options{}) + _, err := m.Execute(leaseCaller("key-all", "101"), []string{"pool-test"}, coreexecutor.Request{Model: "pool-model", Payload: []byte(`{"previous_response_id":"old"}`)}, coreexecutor.Options{}) var poolErr *Error if !errors.As(err, &poolErr) || poolErr.Code != "pool_lease_session_expired" { t.Fatal(err) @@ -121,7 +121,7 @@ func TestExpiredContinuationDoesNotAllocateAnotherPool(t *testing.T) { func TestLeaseStreamFailoverAfterTargetedCooldownCannotEscape(t *testing.T) { m, c, e := leaseManager(t) - ctx := leaseCaller("key-all", "1") + ctx := leaseCaller("key-all", "101") req := coreexecutor.Request{Model: "pool-model"} result, err := m.ExecuteStream(ctx, []string{"pool-test"}, req, coreexecutor.Options{}) if err != nil { @@ -157,7 +157,7 @@ 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") + ctx := leaseCaller("key-all", "101") req := coreexecutor.Request{Model: "pool-model"} if _, err := m.Execute(ctx, []string{"pool-test"}, req, coreexecutor.Options{}); err != nil { t.Fatal(err) @@ -199,7 +199,7 @@ func TestLeaseTerminalFailureReplacesAccountInEveryMode(t *testing.T) { func TestLeaseGenericRateLimitKeepsAccount(t *testing.T) { m, _, e := leaseManager(t) - ctx := leaseCaller("key-all", "1") + ctx := leaseCaller("key-all", "101") req := coreexecutor.Request{Model: "pool-model"} m.Execute(ctx, []string{"pool-test"}, req, coreexecutor.Options{}) before, _ := m.AccountPoolLeases() @@ -215,7 +215,7 @@ func TestLeaseGenericRateLimitKeepsAccount(t *testing.T) { func TestLeaseReplacementWaitsForOtherRequestAndResumesOnAdmission(t *testing.T) { m, _, e := leaseManager(t) - ctx := leaseCaller("key-all", "1") + ctx := leaseCaller("key-all", "101") req := coreexecutor.Request{Model: "pool-model"} held, done, err := m.beginPoolLease(ctx, []string{"pool-test"}, req, coreexecutor.Options{}) if err != nil { @@ -255,7 +255,7 @@ func TestLeaseTerminalFailureClassification(t *testing.T) { 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{}) + ctx, done, err := m.beginPoolLease(leaseCaller("key-all", "101"), []string{"pool-test"}, coreexecutor.Request{Model: "pool-model"}, coreexecutor.Options{}) if err != nil { t.Fatal(err) } @@ -299,7 +299,7 @@ func TestLeaseCommittedStreamIsNotReplayed(t *testing.T) { m, _, e := leaseManager(t) m.RegisterExecutor(&leaseCommittedErrorExecutor{e}) req := coreexecutor.Request{Model: "pool-model"} - ctx := leaseCaller("key-all", "1") + ctx := leaseCaller("key-all", "101") r, err := m.ExecuteStream(ctx, []string{"pool-test"}, req, coreexecutor.Options{}) if err != nil { t.Fatal(err) @@ -332,3 +332,142 @@ func TestLeaseCommittedStreamIsNotReplayed(t *testing.T) { t.Fatal("next request did not replace failed account") } } + +func TestUserOneUsesTemporaryAccountsInEveryMode(t *testing.T) { + for _, mode := range []string{"normal", "count", "stream"} { + t.Run(mode, func(t *testing.T) { + m, _, _ := leaseManager(t) + ctx := leaseCaller("key-all", "1") + req := coreexecutor.Request{Model: "pool-model"} + 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) + } + deadline := time.Now().Add(time.Second) + for time.Now().Before(deadline) { + ls, _ := m.AccountPoolLeases() + if len(ls) == 0 { + return + } + time.Sleep(time.Millisecond) + } + t.Fatal("user 1 left hour lease") + }) + } +} + +func TestTemporaryIdentityAndConcurrentIsolation(t *testing.T) { + m, _, _ := leaseManager(t) + req := coreexecutor.Request{Model: "pool-model"} + providers := []string{"pool-test"} + normal, done, err := m.beginPoolLease(leaseCaller("key-all", "2"), providers, req, coreexecutor.Options{}) + if err != nil { + t.Fatal(err) + } + defer done() + first, done1, err := m.beginPoolLease(leaseCaller("key-all", "1"), providers, req, coreexecutor.Options{}) + if err != nil { + t.Fatal(err) + } + defer done1() + second, done2, err := m.beginPoolLease(leaseCaller("key-all", "1"), providers, req, coreexecutor.Options{}) + if err != nil { + t.Fatal(err) + } + defer done2() + seen := map[string]bool{} + for _, ctx := range []context.Context{normal, first, second} { + l := ctx.Value(requestPoolLeaseKey{}).(*requestPoolLease).Lease + if seen[l.Credential] { + t.Fatal("concurrent accounts shared") + } + seen[l.Credential] = true + scope, err := m.AccountPoolScope(ctx) + if err != nil || !scope.Allows(l.Credential) { + t.Fatal("wrong request binding", err) + } + } + wrong := sdkaccess.WithGatewayIdentity(poolCaller("key-all"), "wrong-instance", "1") + if _, _, err = m.beginPoolLease(wrong, providers, req, coreexecutor.Options{}); err == nil { + t.Fatal("wrong instance bypassed identity validation") + } + done1() + done2() + rows, _ := m.AccountPoolLeases() + if len(rows) != 1 || rows[0].Temporary { + t.Fatal(rows) + } +} + +func TestTemporaryStreamCancellationKeepsReservationUntilDrain(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) + } + raw := make(chan coreexecutor.StreamChunk) + ctx, cancel := context.WithCancel(ctx) + result := holdPoolLeaseStream(ctx, &coreexecutor.StreamResult{Chunks: raw}, done) + cancel() + ls, _ := m.AccountPoolLeases() + if len(ls) != 1 || !ls[0].Temporary { + t.Fatal("released before upstream ended") + } + close(raw) + for range result.Chunks { + } + deadline := time.Now().Add(time.Second) + for time.Now().Before(deadline) { + ls, _ = m.AccountPoolLeases() + if len(ls) == 0 { + return + } + time.Sleep(time.Millisecond) + } + t.Fatal("temporary stream leaked") +} + +func TestTemporaryContinuationDoesNotAllocate(t *testing.T) { + m, _, _ := leaseManager(t) + _, err := m.Execute(leaseCaller("key-all", "1"), []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_temporary_session_unsupported" { + t.Fatal(err) + } + ls, _ := m.AccountPoolLeases() + if len(ls) != 0 { + t.Fatal("continuation allocated account") + } +} + +func TestTemporaryQuotaFailoverDoesNotLeaveHourLease(t *testing.T) { + m, c, e := leaseManager(t) + first := c.AccountPools.Groups[1].CredentialIDs[0] + e.errors = map[string]error{first: &Error{HTTPStatus: 429, Message: "usage_limit_reached"}} + if _, err := m.Execute(leaseCaller("key-all", "1"), []string{"pool-test"}, coreexecutor.Request{Model: "pool-model"}, coreexecutor.Options{}); err != nil { + t.Fatal(err) + } + if len(e.ids) != 2 || e.ids[0] == e.ids[1] { + t.Fatal("temporary failover did not occur", e.ids) + } + rows, _ := m.AccountPoolLeases() + if len(rows) != 0 { + t.Fatal("temporary failover retained lease", rows) + } +}