diff --git a/config.example.yaml b/config.example.yaml index dd93a0125..c2ff1edab 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -126,7 +126,11 @@ access: # Routing strategy for selecting credentials when multiple match. routing: - strategy: "round-robin" # round-robin (default), weighted-round-robin, fill-first + strategy: "round-robin" # round-robin (default), weighted-round-robin, fill-first, soonest-reset + # soonest-reset prefers the credential whose weekly (7-day) quota window resets soonest, so + # quota is used before it expires. Reset times are learned from Claude/Codex response headers. + # A credential with an exhausted window is skipped until that window resets. An unknown Claude + # reset is tried first; an unknown Codex reset last. Other providers behave like fill-first. # weighted-round-robin uses each credential's integer weight (default 1, maximum 1,000,000). # Non-positive weights exclude the credential while this strategy is active. # For OAuth/file credentials, add a top-level numeric "weight" field to the auth JSON. @@ -182,7 +186,10 @@ routing: disable-cooling: false # When true, persist per-auth cooldown status as .cds files next to auth files. - # Default is false; when false, cooldown status is kept in memory only. + # The latest observed quota (rate-limit headers from Claude/Codex responses) is saved + # with it on every observation, so quota-aware routing such as soonest-reset and the + # management quota view keep their data across restarts. + # Default is false; when false, cooldown status and observed quota are kept in memory only. save-cooldown-status: false # Cooldown duration in seconds for transient upstream errors (408/500/502/503/504/520-526). diff --git a/e2e/hermetic/mock-upstream.mjs b/e2e/hermetic/mock-upstream.mjs index 600db5e92..695d737aa 100644 --- a/e2e/hermetic/mock-upstream.mjs +++ b/e2e/hermetic/mock-upstream.mjs @@ -1,9 +1,12 @@ // Mock Codex and Claude upstream. A credential whose key ends in `-limited` answers Codex // 429 usage_limit_reached; every other request gets a complete SSE stream (or a Claude JSON // message when the request does not stream). GET /_mock/requests returns what it received. +// POST /_mock/script takes `{ [key]: { headers?, limited? } }`: `headers` replaces that +// credential's quota headers and `limited: true` makes its Codex requests answer 429. import { createServer } from 'node:http'; const requests = []; +const scripts = new Map(); const resetAt = () => Math.floor(Date.now() / 1000) + 3600; const sse = (res, headers, events) => { @@ -13,7 +16,7 @@ const sse = (res, headers, events) => { }; const codex = (res, key, body) => { - if (key.endsWith('-limited')) { + if (key.endsWith('-limited') || scripts.get(key)?.limited) { res.writeHead(429, { 'content-type': 'application/json' }); const error = { type: 'usage_limit_reached', message: 'You have hit your usage limit.' }; res.end(JSON.stringify({ error: { ...error, resets_at: resetAt(), resets_in_seconds: 3600 } })); @@ -30,7 +33,7 @@ const codex = (res, key, body) => { const usage = { input_tokens: 5, output_tokens: 3, total_tokens: 8 }; sse( res, - { + scripts.get(key)?.headers ?? { 'x-codex-primary-used-percent': '10', 'x-codex-primary-window-minutes': '300', 'x-codex-primary-reset-at': String(resetAt()), @@ -53,8 +56,8 @@ const codex = (res, key, body) => { ); }; -const claude = (res, body) => { - const headers = { +const claude = (res, key, body) => { + const headers = scripts.get(key)?.headers ?? { 'anthropic-ratelimit-unified-status': 'allowed', 'anthropic-ratelimit-unified-5h-status': 'allowed', 'anthropic-ratelimit-unified-5h-utilization': '0.1', @@ -104,9 +107,14 @@ export const startMock = () => 'Bearer ', '', ); + if (req.url === '/_mock/script') { + for (const [scripted, script] of Object.entries(body)) scripts.set(scripted, script); + res.writeHead(204).end(); + return; + } requests.push({ path: req.url, headers: req.headers, body, key }); if (req.url === '/v1/responses') codex(res, key, body); - else claude(res, body); + else claude(res, key, body); }); server.listen(0, '127.0.0.1', () => resolve(`http://127.0.0.1:${server.address().port}`)); }); diff --git a/e2e/hermetic/start.mjs b/e2e/hermetic/start.mjs index dd5fd4d5b..e636ad6f8 100644 --- a/e2e/hermetic/start.mjs +++ b/e2e/hermetic/start.mjs @@ -12,28 +12,42 @@ writeFileSync('.e2e/mock.json', JSON.stringify({ url })); const dir = mkdtempSync(join(tmpdir(), 'cliproxy-e2e-')); mkdirSync(join(dir, 'auths')); -const codex = (key, model, priority = 0) => ({ +const models = (names) => [names].flat().map((name) => ({ name })); +const codex = (key, names, priority = 0) => ({ 'api-key': key, 'base-url': `${url}/v1`, priority, - models: [{ name: model }], + models: models(names), }); -// YAML accepts JSON. The failover and cooldown tests own their models and credentials; -// failover-a has the higher priority, so the proxy tries it first. +const claude = (key, names) => ({ 'api-key': key, 'base-url': url, models: models(names) }); +// A selector scenario's credentials share its model; `-` reaches one of them. +const scenario = (make, model, credentials) => + credentials.map((name) => make(`${model}-${name}`, [model, `${model}-${name}`])); +// YAML accepts JSON. The failover, cooldown and selector tests own their models and +// credentials; failover-a has the higher priority, so the proxy tries it first. const config = { host: '127.0.0.1', port: Number(process.env.PORT), 'auth-dir': join(dir, 'auths'), 'api-keys': ['e2e-client-key'], + routing: { strategy: 'soonest-reset', 'session-affinity': true }, + 'save-cooldown-status': true, 'codex-api-key': [ codex('codex-main', 'gpt-5.6-sol'), codex('failover-a-limited', 'e2e-failover', 1), codex('failover-b', 'e2e-failover'), codex('cooled-a-limited', 'e2e-cooled'), codex('cooled-b-limited', 'e2e-cooled'), + ...scenario(codex, 'e2e-s1', ['a', 'b', 'c']), + ...scenario(codex, 'e2e-s2', ['a', 'b']), + ...scenario(codex, 'e2e-s5', ['known', 'unknown']), + ...scenario(codex, 'e2e-s7', ['a', 'b']), + ...scenario(codex, 'e2e-s8', ['a', 'b']), ], 'claude-api-key': [ - { 'api-key': 'claude-main', 'base-url': url, models: [{ name: 'claude-sonnet-4-6' }] }, + claude('claude-main', 'claude-sonnet-4-6'), + ...scenario(claude, 'e2e-s3', ['a', 'b']), + ...scenario(claude, 'e2e-s5c', ['x', 'y', 'z']), ], }; writeFileSync(join(dir, 'config.yaml'), JSON.stringify(config, null, 2)); diff --git a/e2e/lib/proxy.ts b/e2e/lib/proxy.ts index 0018e3d93..c0919830f 100644 --- a/e2e/lib/proxy.ts +++ b/e2e/lib/proxy.ts @@ -30,3 +30,14 @@ export async function upstreamRequests(marker: string): Promise JSON.stringify(request.body).includes(marker)); } + +export type Script = { headers?: Record; limited?: boolean }; + +/** Sets what the mock answers for each credential key (hermetic/mock-upstream.mjs). */ +export async function script(scripts: Record): Promise { + const response = await fetch(`${mockUrl()}/_mock/script`, { + method: 'POST', + body: JSON.stringify(scripts), + }); + if (!response.ok) throw new Error(`mock script failed: ${response.status}`); +} diff --git a/e2e/tests/failover.e2e.ts b/e2e/tests/failover.e2e.ts index 6e85b3c93..5f34a5e88 100644 --- a/e2e/tests/failover.e2e.ts +++ b/e2e/tests/failover.e2e.ts @@ -41,7 +41,6 @@ test( expect((await response.json()).error).toMatchObject({ code: 'model_cooldown', model: 'e2e-cooled', - provider: 'codex', }); expect(await upstreamRequests(marker)).toEqual([]); }, diff --git a/e2e/tests/models.e2e.ts b/e2e/tests/models.e2e.ts index c2b733aca..f308185ab 100644 --- a/e2e/tests/models.e2e.ts +++ b/e2e/tests/models.e2e.ts @@ -17,7 +17,9 @@ test( }); expect(response.status).toBe(200); const { data } = await response.json(); - expect(data.map((model: { id: string }) => model.id).sort()).toEqual([ + const ids: string[] = data.map((model: { id: string }) => model.id); + // Selector scenarios (tests/selector.e2e.ts) add their own e2e-s* models. + expect(ids.filter((id) => !id.startsWith('e2e-s')).sort()).toEqual([ 'claude-sonnet-4-6', 'e2e-cooled', 'e2e-failover', diff --git a/e2e/tests/selector.e2e.ts b/e2e/tests/selector.e2e.ts new file mode 100644 index 000000000..e7a1f4681 --- /dev/null +++ b/e2e/tests/selector.e2e.ts @@ -0,0 +1,169 @@ +// Soonest-reset selector scenarios (MASTRA-765, MASTRA-768). Each scenario owns its model and +// credentials (hermetic/start.mjs): credential `-` serves `` and its own +// `-`, which seeds that credential's quota observation. Exhaustion is seeded with +// 200 responses carrying exhausted headers, so the selector's header rule decides, not cooldown. +// Several new sessions must land on one credential, which round-robin cannot do by chance. +import { randomUUID } from 'node:crypto'; +import { expect, test } from 'e2e'; +import { clientKey, input, post, script, upstreamRequests } from '../lib/proxy.ts'; + +const hermetic = { platforms: ['hermetic'] }; +const hour = 3600; +const day = 24 * hour; +const at = (seconds: number) => String(Math.floor(Date.now() / 1000) + seconds); + +// Codex Pro reports its weekly window as the primary one. +const codexWeekly = (usedPercent: number, resetIn: number) => ({ + 'x-codex-primary-used-percent': String(usedPercent), + 'x-codex-primary-window-minutes': '10080', + 'x-codex-primary-reset-at': at(resetIn), +}); +const claudeWindows = ( + fiveHourUtilization: number, + fiveHourResetIn: number, + weekResetIn: number, +) => ({ + 'anthropic-ratelimit-unified-status': 'allowed', + 'anthropic-ratelimit-unified-5h-utilization': String(fiveHourUtilization), + 'anthropic-ratelimit-unified-5h-reset': at(fiveHourResetIn), + 'anthropic-ratelimit-unified-7d-utilization': '0.5', + 'anthropic-ratelimit-unified-7d-reset': at(weekResetIn), +}); + +/** Sends one Responses request in `session`; returns the credential keys the mock saw. */ +async function responses(baseUrl: string | undefined, model: string, session = randomUUID()) { + const marker = randomUUID(); + const response = await post( + baseUrl, + '/v1/responses', + { model, input: input(marker) }, + { Session_id: session }, + ); + expect(response.status).toBe(200); + await response.text(); + return (await upstreamRequests(marker)).map((request) => request.key); +} + +/** Sends one Messages request in `session`; returns the credential keys the mock saw. */ +async function messages(baseUrl: string | undefined, model: string, session = randomUUID()) { + const marker = randomUUID(); + const response = await post( + baseUrl, + '/v1/messages', + { model, max_tokens: 16, messages: [{ role: 'user', content: marker }] }, + { + 'x-api-key': clientKey, + 'anthropic-version': '2023-06-01', + 'X-Claude-Code-Session-Id': session, + }, + ); + expect(response.status).toBe(200); + await response.text(); + return (await upstreamRequests(marker)).map((request) => request.key); +} + +test('S1: new Codex sessions go to the soonest weekly reset', hermetic, async ({ app }) => { + await script({ + 'e2e-s1-a': { headers: codexWeekly(50, 6 * day) }, + 'e2e-s1-b': { headers: codexWeekly(50, 1 * day) }, + 'e2e-s1-c': { headers: codexWeekly(50, 3 * day) }, + }); + for (const name of ['a', 'b', 'c']) await responses(app.baseUrl, `e2e-s1-${name}`); + for (let i = 0; i < 3; i++) expect(await responses(app.baseUrl, 'e2e-s1')).toEqual(['e2e-s1-b']); +}); + +test( + 'S2: a small remainder resetting soon goes before a fresh late window, then fails over', + hermetic, + async ({ app }) => { + await script({ + 'e2e-s2-a': { headers: codexWeekly(98, 6 * hour) }, + 'e2e-s2-b': { headers: codexWeekly(0, 5 * day) }, + }); + for (const name of ['a', 'b']) await responses(app.baseUrl, `e2e-s2-${name}`); + for (let i = 0; i < 2; i++) { + expect(await responses(app.baseUrl, 'e2e-s2')).toEqual(['e2e-s2-a']); + } + await script({ 'e2e-s2-a': { limited: true } }); + expect(await responses(app.baseUrl, 'e2e-s2')).toEqual(['e2e-s2-a', 'e2e-s2-b']); + }, +); + +test( + 'S3: an exhausted Claude 5-hour window is skipped until it resets', + hermetic, + async ({ app }) => { + const fiveHourReset = 4; + await script({ + 'e2e-s3-a': { headers: claudeWindows(1, fiveHourReset, 1 * day) }, + 'e2e-s3-b': { headers: claudeWindows(0.1, 4 * hour, 3 * day) }, + }); + for (const name of ['a', 'b']) await messages(app.baseUrl, `e2e-s3-${name}`); + for (let i = 0; i < 2; i++) expect(await messages(app.baseUrl, 'e2e-s3')).toEqual(['e2e-s3-b']); + await new Promise((resolve) => setTimeout(resolve, (fiveHourReset + 1) * 1000)); + expect(await messages(app.baseUrl, 'e2e-s3')).toEqual(['e2e-s3-a']); + }, +); + +test( + 'S5: an unknown Codex reset sorts last and an unknown Claude reset sorts first', + hermetic, + async ({ app }) => { + await script({ + 'e2e-s5-known': { headers: codexWeekly(50, 6 * day) }, + 'e2e-s5-unknown': { headers: {} }, + 'e2e-s5c-x': { headers: claudeWindows(0.1, 4 * hour, 6 * day) }, + 'e2e-s5c-y': { headers: claudeWindows(0.1, 4 * hour, 1 * day) }, + 'e2e-s5c-z': { headers: claudeWindows(0.1, 4 * hour, 3 * day) }, + }); + await responses(app.baseUrl, 'e2e-s5-known'); + for (let i = 0; i < 2; i++) { + expect(await responses(app.baseUrl, 'e2e-s5')).toEqual(['e2e-s5-known']); + } + + for (const name of ['y', 'z']) await messages(app.baseUrl, `e2e-s5c-${name}`); + expect(await messages(app.baseUrl, 'e2e-s5c')).toEqual(['e2e-s5c-x']); + // x's first response taught its late reset, so y (soonest) now leads. + for (let i = 0; i < 2; i++) + expect(await messages(app.baseUrl, 'e2e-s5c')).toEqual(['e2e-s5c-y']); + }, +); + +test( + 'S7 (regression): a bound session stays on its credential when another becomes better', + hermetic, + async ({ app }) => { + await script({ + 'e2e-s7-a': { headers: codexWeekly(50, 3 * day) }, + 'e2e-s7-b': { headers: codexWeekly(50, 3 * day) }, + }); + const session = randomUUID(); + const [bound] = await responses(app.baseUrl, 'e2e-s7', session); + const other = bound === 'e2e-s7-a' ? 'e2e-s7-b' : 'e2e-s7-a'; + await script({ + [bound]: { headers: codexWeekly(50, 6 * day) }, + [other]: { headers: codexWeekly(50, 1 * hour) }, + }); + await responses(app.baseUrl, other); + for (let i = 0; i < 3; i++) { + expect(await responses(app.baseUrl, 'e2e-s7', session)).toEqual([bound]); + } + }, +); + +test( + 'S8 (regression): a burst of new sessions at exhaustion completes without client errors', + hermetic, + async ({ app }) => { + await script({ + 'e2e-s8-a': { headers: codexWeekly(99, 1 * hour) }, + 'e2e-s8-b': { headers: codexWeekly(0, 5 * day) }, + }); + for (const name of ['a', 'b']) await responses(app.baseUrl, `e2e-s8-${name}`); + await script({ 'e2e-s8-a': { limited: true } }); + const served = await Promise.all( + Array.from({ length: 6 }, () => responses(app.baseUrl, 'e2e-s8')), + ); + for (const keys of served) expect(keys.at(-1)).toBe('e2e-s8-b'); + }, +); diff --git a/internal/clienterror/client_error.go b/internal/clienterror/client_error.go index f5d520ccb..e5af82b76 100644 --- a/internal/clienterror/client_error.go +++ b/internal/clienterror/client_error.go @@ -25,6 +25,10 @@ var requestFaultCodes = map[string]struct{}{ "unsupported_value": {}, "invalid_request_error": {}, "previous_response_not_found": {}, + // Anthropic: server-side thread state for previous_message_id is missing on + // this account. Another credential cannot have it either; the client recovers + // by replaying the conversation as a new thread. + "thread_not_found": {}, } var requestFaultTypes = map[string]struct{}{ @@ -170,7 +174,7 @@ func hasRequestFaultBody(err error) bool { if body == "" || !json.Valid([]byte(body)) { return false } - for _, path := range []string{"error.code", "code", "response.error.code", "body.error.code"} { + for _, path := range []string{"error.code", "code", "response.error.code", "body.error.code", "error.details.error_code"} { code := strings.ToLower(strings.TrimSpace(gjson.Get(body, path).String())) if _, ok := requestFaultCodes[code]; ok { return true diff --git a/internal/clienterror/client_error_test.go b/internal/clienterror/client_error_test.go index 12ee78c11..9d2e581cc 100644 --- a/internal/clienterror/client_error_test.go +++ b/internal/clienterror/client_error_test.go @@ -140,6 +140,18 @@ func TestIsRequestFaultStructuredIdentifiers(t *testing.T) { } } +func TestIsRequestFaultClaudeThreadNotFound(t *testing.T) { + err := errors.New(`{"type":"error","error":{"type":"not_found_error","message":"No thread state was found for the requested ` + "`previous_message_id`" + `.","details":{"error_code":"thread_not_found"}}}`) + if !IsRequestFault(http.StatusNotFound, err) { + t.Fatal("thread_not_found was not classified as a request fault") + } + + other := errors.New(`{"type":"error","error":{"type":"not_found_error","message":"missing","details":{"error_code":"something_else"}}}`) + if IsRequestFault(http.StatusNotFound, other) { + t.Fatal("unrelated not_found_error was classified as a request fault") + } +} + func TestIsRequestFault(t *testing.T) { tests := []struct { name string diff --git a/internal/config/config.go b/internal/config/config.go index ca3767a5d..f7e6c866b 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -72,7 +72,8 @@ type Config struct { // DisableCooling disables auth/model cooldown scheduling when true unless a credential or provider overrides it. DisableCooling bool `yaml:"disable-cooling" json:"disable-cooling"` - // SaveCooldownStatus persists runtime cooldown status next to auth files when true. + // SaveCooldownStatus persists runtime cooldown status, and the latest observed quota + // snapshot, next to auth files when true. SaveCooldownStatus bool `yaml:"save-cooldown-status" json:"save-cooldown-status"` // TransientErrorCooldownSeconds controls cooldowns for transient upstream errors (408/500/502/503/504/520-526). diff --git a/internal/config/config_types.go b/internal/config/config_types.go index c3584aa11..e5f2ecb69 100644 --- a/internal/config/config_types.go +++ b/internal/config/config_types.go @@ -351,7 +351,7 @@ type QuotaExceeded struct { // RoutingConfig configures how credentials are selected for requests. type RoutingConfig struct { // Strategy selects the credential selection strategy. - // Supported values: "round-robin" (default), "weighted-round-robin", "fill-first". + // Supported values: "round-robin" (default), "weighted-round-robin", "fill-first", "soonest-reset". Strategy string `yaml:"strategy,omitempty" json:"strategy,omitempty"` // SessionAffinity enables universal session-sticky routing for all clients. diff --git a/sdk/cliproxy/auth/conductor_cooldown.go b/sdk/cliproxy/auth/conductor_cooldown.go index 33f8b5d8d..567c1e3c9 100644 --- a/sdk/cliproxy/auth/conductor_cooldown.go +++ b/sdk/cliproxy/auth/conductor_cooldown.go @@ -323,6 +323,14 @@ func (m *Manager) RestoreCooldownStates(ctx context.Context) error { m.mu.Lock() for _, record := range records { + if m.restoreQuotaObservationLocked(record) { + if auth := m.auths[strings.TrimSpace(record.AuthID)]; auth != nil { + snapshotsByID[auth.ID] = auth.Clone() + } + } + if record.Status == cooldownStateStatusObserved { + continue + } if strings.TrimSpace(record.Model) == "" { authLevelRecords = append(authLevelRecords, record) continue @@ -611,7 +619,7 @@ func (m *Manager) cooldownStateRecordsSnapshot() []CooldownStateRecord { m.mu.RLock() for _, auth := range m.auths { - records = append(records, m.cooldownStateRecordsForAuthLocked(auth, now)...) + records = append(records, withQuotaObservations(auth, m.cooldownStateRecordsForAuthLocked(auth, now))...) } m.mu.RUnlock() @@ -646,6 +654,76 @@ func (m *Manager) cooldownStateRecordsForAuthLocked(auth *Auth, now time.Time) [ return records } +// withQuotaObservations attaches each passive quota observation to the +// persisted record for the same auth/model, adding an observation-only record +// where no cooldown exists. Stores key records by auth and model, so one record +// carries both. Change detection does not use this; MarkResult saves whenever it +// records an observation. +func withQuotaObservations(auth *Auth, records []CooldownStateRecord) []CooldownStateRecord { + if auth == nil || auth.ID == "" || auth.Disabled || auth.Status == StatusDisabled { + return records + } + attach := func(model string, quota QuotaState, updatedAt time.Time) { + if quota.ObservedAt.IsZero() || len(quota.Signals) == 0 { + return + } + observation := quota.Clone() + for i := range records { + if records[i].Model == model { + records[i].Quota.ObservedAt = observation.ObservedAt + records[i].Quota.Signals = observation.Signals + return + } + } + records = append(records, CooldownStateRecord{ + Provider: strings.TrimSpace(auth.Provider), + AuthID: auth.ID, + AuthFile: cooldownAuthFile(auth), + Model: model, + Status: cooldownStateStatusObserved, + Quota: QuotaState{ObservedAt: observation.ObservedAt, Signals: observation.Signals}, + UpdatedAt: updatedAt, + }) + } + attach("", auth.Quota, auth.UpdatedAt) + for model, state := range auth.ModelStates { + if state != nil && strings.TrimSpace(model) != "" { + attach(strings.TrimSpace(model), state.Quota, state.UpdatedAt) + } + } + sort.Slice(records, func(i, j int) bool { + return records[i].Model < records[j].Model + }) + return records +} + +// restoreQuotaObservationLocked merges a saved observation into the auth or model +// state. The newest observation wins, so a live response after startup is kept. +func (m *Manager) restoreQuotaObservationLocked(record CooldownStateRecord) bool { + quota := record.Quota + if quota.ObservedAt.IsZero() || len(quota.Signals) == 0 { + return false + } + auth := m.auths[strings.TrimSpace(record.AuthID)] + if auth == nil || auth.Disabled || auth.Status == StatusDisabled || !ProviderSupportsQuotaObservation(auth.Provider) { + return false + } + observation := QuotaState{ObservedAt: quota.ObservedAt, Signals: quota.Clone().Signals} + model := strings.TrimSpace(record.Model) + if model == "" { + before := auth.Quota.ObservedAt + auth.Quota = mergeQuotaObservation(auth.Quota, observation) + return !auth.Quota.ObservedAt.Equal(before) + } + state := ensureModelState(auth, model) + if state == nil { + return false + } + before := state.Quota.ObservedAt + state.Quota = mergeQuotaObservation(state.Quota, observation) + return !state.Quota.ObservedAt.Equal(before) +} + func cooldownStateRecordsEqual(a, b []CooldownStateRecord) bool { if len(a) != len(b) { return false @@ -1009,10 +1087,11 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { auth.Generation++ auth.UpdatedAt = now + observed := false if !result.SkipQuotaObservation { - auth.Quota.ObserveResponseHeadersForProvider(result.Provider, responseHeaders, now) - if modelState != nil { - modelState.Quota.ObserveResponseHeadersForProvider(result.Provider, responseHeaders, now) + observed = auth.Quota.ObserveResponseHeadersForProvider(result.Provider, responseHeaders, now) + if modelState != nil && modelState.Quota.ObserveResponseHeadersForProvider(result.Provider, responseHeaders, now) { + observed = true } } @@ -1021,6 +1100,7 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { if trackCooldownState { cooldownRecordsAfter := m.cooldownStateRecordsForAuthLocked(auth, now) cooldownStateChanged = !cooldownStateRecordsEqual(cooldownRecordsBefore, cooldownRecordsAfter) + cooldownStateChanged = cooldownStateChanged || observed } } m.mu.Unlock() diff --git a/sdk/cliproxy/auth/conductor_overrides_test.go b/sdk/cliproxy/auth/conductor_overrides_test.go index 2149cb7e8..075b2c674 100644 --- a/sdk/cliproxy/auth/conductor_overrides_test.go +++ b/sdk/cliproxy/auth/conductor_overrides_test.go @@ -2335,6 +2335,62 @@ func TestManager_RequestScopedNotFoundStopsRetryWithoutSuspendingAuth(t *testing } } +// A Claude thread lives on the account that created it, so thread_not_found must +// reach the client (which replays the conversation) instead of rotating +// credentials or putting the model into the generic 12h not-found cooldown. +func TestManager_ClaudeThreadNotFoundStopsRetryWithoutCooldown(t *testing.T) { + const threadNotFound = `{"type":"error","error":{"type":"not_found_error","message":"No thread state was found for the requested previous_message_id.","details":{"error_code":"thread_not_found"}}}` + m := NewManager(nil, nil, nil) + executor := &authFallbackExecutor{ + id: "claude", + executeErrors: map[string]error{ + "aa-thread-auth": &Error{HTTPStatus: http.StatusNotFound, Message: threadNotFound}, + }, + } + m.RegisterExecutor(executor) + + model := "claude-thread-test-model" + threadAuth := &Auth{ID: "aa-thread-auth", Provider: "claude"} + otherAuth := &Auth{ID: "bb-other-auth", Provider: "claude"} + + reg := registry.GetGlobalRegistry() + reg.RegisterClient(threadAuth.ID, "claude", []*registry.ModelInfo{{ID: model}}) + reg.RegisterClient(otherAuth.ID, "claude", []*registry.ModelInfo{{ID: model}}) + t.Cleanup(func() { + reg.UnregisterClient(threadAuth.ID) + reg.UnregisterClient(otherAuth.ID) + }) + + for _, auth := range []*Auth{threadAuth, otherAuth} { + if _, errRegister := m.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("register %s: %v", auth.ID, errRegister) + } + } + + _, errExecute := m.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) + if errExecute == nil { + t.Fatal("expected thread_not_found error") + } + if errResult, ok := errExecute.(*Error); !ok || errResult.HTTPStatus != http.StatusNotFound { + t.Fatalf("error = %#v, want 404 *Error", errExecute) + } + + if got := executor.ExecuteCalls(); len(got) != 1 || got[0] != threadAuth.ID { + t.Fatalf("execute calls = %v, want only %s", got, threadAuth.ID) + } + + updated, ok := m.GetByID(threadAuth.ID) + if !ok || updated == nil { + t.Fatal("expected thread auth to remain registered") + } + if updated.Unavailable || !updated.NextRetryAfter.IsZero() { + t.Fatalf("thread auth was suspended: unavailable=%v nextRetryAfter=%v", updated.Unavailable, updated.NextRetryAfter) + } + if state := updated.ModelStates[model]; state != nil && !state.NextRetryAfter.IsZero() { + t.Fatalf("thread auth model cooled down until %v", state.NextRetryAfter) + } +} + func TestManager_MarkResult_RequestFaultBodyDoesNotCooldownModelOrAuth(t *testing.T) { previous := quotaCooldownDisabled.Load() quotaCooldownDisabled.Store(false) diff --git a/sdk/cliproxy/auth/cooldown_state.go b/sdk/cliproxy/auth/cooldown_state.go index 830a2e6e9..c6685115e 100644 --- a/sdk/cliproxy/auth/cooldown_state.go +++ b/sdk/cliproxy/auth/cooldown_state.go @@ -29,6 +29,12 @@ type CooldownStateRecord struct { UpdatedAt time.Time `json:"updated_at"` } +// cooldownStateStatusObserved marks a record that carries only a passive quota +// observation (Quota.ObservedAt and Quota.Signals) and no active cooldown. Saving +// observations lets quota-aware routing and the management quota view keep their +// data across restarts instead of relearning it from live traffic. +const cooldownStateStatusObserved = "observed" + // CooldownStateStore persists runtime cooldown state independently from auth tokens. type CooldownStateStore interface { Load(context.Context) ([]CooldownStateRecord, error) diff --git a/sdk/cliproxy/auth/quota_observation_persist_test.go b/sdk/cliproxy/auth/quota_observation_persist_test.go new file mode 100644 index 000000000..fb3cefd71 --- /dev/null +++ b/sdk/cliproxy/auth/quota_observation_persist_test.go @@ -0,0 +1,133 @@ +package auth + +import ( + "context" + "strconv" + "testing" + "time" +) + +func newObservationPersistManager(t *testing.T, store CooldownStateStore, authID string) *Manager { + t.Helper() + m := NewManager(nil, nil, nil) + m.SetCooldownStateStore(store) + if _, err := m.Register(context.Background(), &Auth{ID: authID, Provider: "claude", Status: StatusActive}); err != nil { + t.Fatalf("register: %v", err) + } + return m +} + +func TestQuotaObservationsSurviveRestart(t *testing.T) { + ctx := context.Background() + store := NewFileCooldownStateStore(t.TempDir()) + now := time.Now() + weeklyReset := now.Add(48 * time.Hour).Truncate(time.Second) + fableReset := now.Add(20 * time.Hour).Truncate(time.Second) + + first := newObservationPersistManager(t, store, "claude-a.json") + first.mu.Lock() + auth := first.auths["claude-a.json"] + auth.Quota.ObservedAt = now.Add(-time.Minute) + auth.Quota.Signals = map[string]string{ + "Anthropic-Ratelimit-Unified-7d-Reset": strconv.FormatInt(weeklyReset.Unix(), 10), + } + ensureModelState(auth, "claude-fable-5").Quota = QuotaState{ + ObservedAt: now.Add(-time.Minute), + Signals: map[string]string{ + "Anthropic-Ratelimit-Unified-7d_oi-Reset": strconv.FormatInt(fableReset.Unix(), 10), + }, + } + first.mu.Unlock() + first.persistCooldownStates(ctx) + + second := newObservationPersistManager(t, store, "claude-a.json") + if err := second.RestoreCooldownStates(ctx); err != nil { + t.Fatalf("restore: %v", err) + } + restored, ok := second.GetByID("claude-a.json") + if !ok { + t.Fatal("auth missing after restore") + } + if got, want := restored.Quota.Signals["Anthropic-Ratelimit-Unified-7d-Reset"], strconv.FormatInt(weeklyReset.Unix(), 10); got != want { + t.Fatalf("credential weekly reset = %q, want %q", got, want) + } + if got, want := restored.ModelStates["claude-fable-5"].Quota.Signals["Anthropic-Ratelimit-Unified-7d_oi-Reset"], strconv.FormatInt(fableReset.Unix(), 10); got != want { + t.Fatalf("fable weekly reset = %q, want %q", got, want) + } + if restored.Unavailable || restored.ModelStates["claude-fable-5"].Unavailable { + t.Fatal("restoring an observation must not make the auth or model unavailable") + } +} + +func TestQuotaObservationSharesRecordWithCooldown(t *testing.T) { + ctx := context.Background() + store := NewFileCooldownStateStore(t.TempDir()) + now := time.Now() + retryAt := now.Add(time.Hour) + + first := newObservationPersistManager(t, store, "claude-b.json") + first.mu.Lock() + state := ensureModelState(first.auths["claude-b.json"], "claude-opus") + state.Unavailable = true + state.Status = StatusError + state.NextRetryAfter = retryAt + state.Quota = QuotaState{ + Exceeded: true, + Reason: "quota", + ObservedAt: now, + Signals: map[string]string{"Anthropic-Ratelimit-Unified-7d-Utilization": "1.0"}, + } + first.mu.Unlock() + + records := first.cooldownStateRecordsSnapshot() + if len(records) != 1 { + t.Fatalf("records = %+v, want one combined record", records) + } + if records[0].NextRetryAfter.IsZero() || records[0].Quota.Signals["Anthropic-Ratelimit-Unified-7d-Utilization"] != "1.0" { + t.Fatalf("record lost cooldown or observation: %+v", records[0]) + } + first.persistCooldownStates(ctx) + + second := newObservationPersistManager(t, store, "claude-b.json") + if err := second.RestoreCooldownStates(ctx); err != nil { + t.Fatalf("restore: %v", err) + } + restored, _ := second.GetByID("claude-b.json") + got := restored.ModelStates["claude-opus"] + if got == nil || !got.Unavailable || !got.NextRetryAfter.Equal(retryAt) { + t.Fatalf("cooldown not restored: %+v", got) + } + if got.Quota.Signals["Anthropic-Ratelimit-Unified-7d-Utilization"] != "1.0" { + t.Fatalf("observation not restored: %+v", got.Quota) + } +} + +func TestRestoreKeepsNewerLiveObservation(t *testing.T) { + ctx := context.Background() + store := NewFileCooldownStateStore(t.TempDir()) + now := time.Now() + + first := newObservationPersistManager(t, store, "claude-d.json") + first.mu.Lock() + first.auths["claude-d.json"].Quota = QuotaState{ + ObservedAt: now.Add(-time.Hour), + Signals: map[string]string{"Anthropic-Ratelimit-Unified-7d-Utilization": "0.9"}, + } + first.mu.Unlock() + first.persistCooldownStates(ctx) + + second := newObservationPersistManager(t, store, "claude-d.json") + second.mu.Lock() + second.auths["claude-d.json"].Quota = QuotaState{ + ObservedAt: now, + Signals: map[string]string{"Anthropic-Ratelimit-Unified-7d-Utilization": "0.1"}, + } + second.mu.Unlock() + if err := second.RestoreCooldownStates(ctx); err != nil { + t.Fatalf("restore: %v", err) + } + restored, _ := second.GetByID("claude-d.json") + if got := restored.Quota.Signals["Anthropic-Ratelimit-Unified-7d-Utilization"]; got != "0.1" { + t.Fatalf("utilization = %q, want the newer live value 0.1", got) + } +} diff --git a/sdk/cliproxy/auth/selector_soonest_reset.go b/sdk/cliproxy/auth/selector_soonest_reset.go new file mode 100644 index 000000000..b29a1ba00 --- /dev/null +++ b/sdk/cliproxy/auth/selector_soonest_reset.go @@ -0,0 +1,152 @@ +package auth + +import ( + "context" + "net/http" + "strconv" + "strings" + "time" + + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" +) + +const ( + // weeklyWindowMinutes is the length of the weekly quota window advertised by Codex. + weeklyWindowMinutes = 7 * 24 * 60 + week = 7 * 24 * time.Hour +) + +// unknownResetLast orders a credential after every credential with a known reset. +var unknownResetLast = time.Unix(1<<62, 0) + +// SoonestResetSelector prefers the credential whose weekly quota window resets soonest, so +// quota that would otherwise expire unused at the reset is consumed first. For one shared +// weekly limit no other order wastes less quota. +// +// Windows come from the passive quota snapshot (Quota.Signals) captured from upstream +// response headers. A credential is skipped while one of its observed windows is exhausted +// and that window's reset is still ahead; when every candidate is skipped this way the +// ordering alone decides and upstream 429 handling takes over. +// +// An unknown Claude weekly reset sorts first: the window runs on a fixed per-account +// schedule, so one request learns it. A passed Claude reset rolls forward by whole weeks. +// An unknown or passed Codex weekly reset sorts last: that window is usually unstarted and +// starting it early can waste it. Ties keep input order (ID order), so providers without +// quota signals behave like fill-first. +type SoonestResetSelector struct{} + +// Pick selects the available auth whose weekly quota window resets soonest. +func (s *SoonestResetSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) { + _ = opts + now := time.Now() + available, err := getSelectorAvailableAuths(ctx, auths, provider, model, now) + if err != nil { + return nil, err + } + available = preferCodexWebsocketAuths(ctx, provider, available) + var picked *Auth + var pickedGated bool + var pickedReset time.Time + for _, auth := range available { + gated, resetAt := soonestResetRank(auth, now) + if picked == nil || pickedGated && !gated || gated == pickedGated && resetAt.Before(pickedReset) { + picked, pickedGated, pickedReset = auth, gated, resetAt + } + } + return picked, nil +} + +// soonestResetRank reports whether an observed window of auth is exhausted until a future +// reset, and the weekly reset that orders auth among the candidates. +func soonestResetRank(auth *Auth, now time.Time) (gated bool, weeklyReset time.Time) { + claude := strings.EqualFold(strings.TrimSpace(auth.Provider), "claude") + weeklyReset = unknownResetLast + if claude { + weeklyReset = time.Time{} + } + for _, window := range observedQuotaWindows(auth) { + if window.exhausted && window.resetAt.After(now) { + gated = true + } + if !window.weekly { + continue + } + switch { + case window.resetAt.After(now): + weeklyReset = window.resetAt + case claude: + weeklyReset = window.resetAt.Add(week * (now.Sub(window.resetAt)/week + 1)) + } + } + return gated, weeklyReset +} + +// quotaWindow is one observed quota window of a credential. +type quotaWindow struct { + weekly bool + exhausted bool + usedPercent float64 + resetAt time.Time +} + +// observedQuotaWindows reads the Claude 5h and 7d windows, or the Codex primary and +// secondary windows classified by length, from the credential-wide quota snapshot. +func observedQuotaWindows(auth *Auth) []quotaWindow { + signal := func(name string) string { + return strings.TrimSpace(auth.Quota.Signals[http.CanonicalHeaderKey(name)]) + } + var windows []quotaWindow + switch strings.ToLower(strings.TrimSpace(auth.Provider)) { + case "claude": + for _, name := range []string{"5h", "7d"} { + prefix := "Anthropic-Ratelimit-Unified-" + name + "-" + resetAt := parseQuotaResetAt(signal(prefix+"Reset"), "", time.Time{}) + if resetAt.IsZero() { + continue + } + utilization, _ := strconv.ParseFloat(signal(prefix+"Utilization"), 64) + rejected := strings.EqualFold(signal(prefix+"Status"), "rejected") + windows = append(windows, quotaWindow{weekly: name == "7d", exhausted: utilization >= 1 || rejected, resetAt: resetAt}) + } + case "codex": + // Codex Pro reports its weekly window as primary, so windows are keyed by length. + mostUsed := 0.0 + for _, name := range []string{"Primary", "Secondary"} { + prefix := "X-Codex-" + name + "-" + resetAt := parseQuotaResetAt(signal(prefix+"Reset-At"), signal(prefix+"Reset-After-Seconds"), auth.Quota.ObservedAt) + if resetAt.IsZero() { + continue + } + // A missing window length keeps the window; it only means the length is unknown. + minutes, _ := strconv.ParseInt(signal(prefix+"Window-Minutes"), 10, 64) + used, _ := strconv.ParseFloat(signal(prefix+"Used-Percent"), 64) + mostUsed = max(mostUsed, used) + windows = append(windows, quotaWindow{weekly: minutes >= weeklyWindowMinutes, usedPercent: used, resetAt: resetAt}) + } + // The credential-wide limit-reached flag names no window; it is the most used one. + limitReached, _ := strconv.ParseBool(signal("X-Codex-Limit-Reached")) + for i := range windows { + used := windows[i].usedPercent + windows[i].exhausted = used >= 100 || limitReached && used == mostUsed + } + } + return windows +} + +// parseQuotaResetAt resolves a reset time from an absolute unix timestamp or, failing +// that, from a relative seconds value anchored at the observation time. +func parseQuotaResetAt(resetAt, resetAfterSeconds string, observedAt time.Time) time.Time { + if unix, errParse := strconv.ParseFloat(resetAt, 64); errParse == nil && unix > 0 { + return time.Unix(int64(unix), 0) + } + if parsed, errParse := time.Parse(time.RFC3339, resetAt); errParse == nil { + return parsed + } + if observedAt.IsZero() { + return time.Time{} + } + if seconds, errParse := strconv.ParseFloat(resetAfterSeconds, 64); errParse == nil && seconds >= 0 { + return observedAt.Add(time.Duration(seconds * float64(time.Second))) + } + return time.Time{} +} diff --git a/sdk/cliproxy/auth/selector_soonest_reset_test.go b/sdk/cliproxy/auth/selector_soonest_reset_test.go new file mode 100644 index 000000000..4f1c8abe3 --- /dev/null +++ b/sdk/cliproxy/auth/selector_soonest_reset_test.go @@ -0,0 +1,377 @@ +package auth + +import ( + "context" + "net/http" + "strconv" + "sync" + "testing" + "time" + + internallogging "github.com/router-for-me/CLIProxyAPI/v8/internal/logging" + "github.com/router-for-me/CLIProxyAPI/v8/internal/registry" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" +) + +const day = 24 * time.Hour + +func unixString(at time.Time) string { + return strconv.FormatInt(at.Unix(), 10) +} + +func claudeAuthResettingAt(id string, resetAt time.Time) *Auth { + return &Auth{ + ID: id, + Provider: "claude", + Quota: QuotaState{ + ObservedAt: resetAt.Add(-time.Hour), + Signals: map[string]string{ + "Anthropic-Ratelimit-Unified-7d-Reset": unixString(resetAt), + }, + }, + } +} + +// claudeAuthWithWindows observes a 5-hour window and a weekly window at 50% use. +func claudeAuthWithWindows(id string, fiveHourUtilization float64, fiveHourReset, weekReset time.Time) *Auth { + auth := claudeAuthResettingAt(id, weekReset) + auth.Quota.Signals["Anthropic-Ratelimit-Unified-7d-Utilization"] = "0.5" + auth.Quota.Signals["Anthropic-Ratelimit-Unified-5h-Utilization"] = strconv.FormatFloat(fiveHourUtilization, 'f', -1, 64) + auth.Quota.Signals["Anthropic-Ratelimit-Unified-5h-Reset"] = unixString(fiveHourReset) + return auth +} + +func codexWindowSignals(position string, minutes int, usedPercent float64, resetAt time.Time) map[string]string { + prefix := "X-Codex-" + position + "-" + return map[string]string{ + prefix + "Window-Minutes": strconv.Itoa(minutes), + prefix + "Used-Percent": strconv.FormatFloat(usedPercent, 'f', -1, 64), + prefix + "Reset-At": unixString(resetAt), + } +} + +func codexAuthWithSignals(id string, signalSets ...map[string]string) *Auth { + signals := map[string]string{} + for _, set := range signalSets { + for key, value := range set { + signals[key] = value + } + } + return &Auth{ID: id, Provider: "codex", Quota: QuotaState{ObservedAt: time.Now(), Signals: signals}} +} + +// codexWeeklyAuth observes a Codex Pro account, whose weekly window is the primary one. +func codexWeeklyAuth(id string, usedPercent float64, resetAt time.Time) *Auth { + return codexAuthWithSignals(id, codexWindowSignals("Primary", weeklyWindowMinutes, usedPercent, resetAt)) +} + +func pickSoonestReset(t *testing.T, provider string, auths ...*Auth) string { + t.Helper() + got, err := (&SoonestResetSelector{}).Pick(context.Background(), provider, "", cliproxyexecutor.Options{}, auths) + if err != nil { + t.Fatalf("Pick() error = %v", err) + } + return got.ID +} + +func TestSoonestResetSelectorPick_PrefersSoonestWeeklyReset(t *testing.T) { + t.Parallel() + + now := time.Now() + if got := pickSoonestReset(t, "claude", + claudeAuthResettingAt("a", now.Add(5*day)), + claudeAuthResettingAt("b", now.Add(2*time.Hour)), + claudeAuthResettingAt("c", now.Add(3*day)), + ); got != "b" { + t.Fatalf("Pick() = %q, want b", got) + } +} + +// S1: Codex order follows the weekly window, which is chosen by length, not position. +func TestSoonestResetSelectorPick_CodexWeeklyWindowByLength(t *testing.T) { + t.Parallel() + + now := time.Now() + if got := pickSoonestReset(t, "codex", + codexWeeklyAuth("a", 50, now.Add(6*day)), + codexWeeklyAuth("b", 50, now.Add(1*day)), + codexWeeklyAuth("c", 50, now.Add(3*day)), + ); got != "b" { + t.Fatalf("Pick() = %q, want b", got) + } + + // x's 5-hour primary window resets first, but y's weekly window resets before x's. + x := codexAuthWithSignals("x", + codexWindowSignals("Primary", 300, 10, now.Add(time.Hour)), + codexWindowSignals("Secondary", weeklyWindowMinutes, 50, now.Add(5*day)), + ) + y := codexAuthWithSignals("y", codexWindowSignals("Secondary", weeklyWindowMinutes, 50, now.Add(2*day))) + if got := pickSoonestReset(t, "codex", x, y); got != "y" { + t.Fatalf("Pick() = %q, want y", got) + } +} + +// S2: a small remainder resetting soon is used before a fresh window resetting late. +func TestSoonestResetSelectorPick_SmallRemainderResettingSoonFirst(t *testing.T) { + t.Parallel() + + now := time.Now() + if got := pickSoonestReset(t, "codex", + codexWeeklyAuth("a", 0, now.Add(5*day)), + codexWeeklyAuth("b", 98, now.Add(6*time.Hour)), + ); got != "b" { + t.Fatalf("Pick() = %q, want b", got) + } +} + +func TestSoonestResetSelectorPick_SkipsCoolingDownAuth(t *testing.T) { + t.Parallel() + + now := time.Now() + exhausted := claudeAuthResettingAt("b", now.Add(2*time.Hour)) + exhausted.Quota.Exceeded = true + exhausted.Quota.Reason = "credential_quota" + exhausted.Quota.NextRecoverAt = now.Add(2 * time.Hour) + if got := pickSoonestReset(t, "claude", + claudeAuthResettingAt("a", now.Add(5*day)), + exhausted, + claudeAuthResettingAt("c", now.Add(3*day)), + ); got != "c" { + t.Fatalf("Pick() = %q, want c", got) + } +} + +// S3: an exhausted window skips the credential until that window resets. +func TestSoonestResetSelectorPick_ExhaustionGate(t *testing.T) { + t.Parallel() + + now := time.Now() + later := claudeAuthWithWindows("b", 0.1, now.Add(4*time.Hour), now.Add(3*day)) + rejected := claudeAuthWithWindows("a", 0.5, now.Add(time.Hour), now.Add(day)) + rejected.Quota.Signals["Anthropic-Ratelimit-Unified-5h-Status"] = "rejected" + limitReached := codexAuthWithSignals("a", + codexWindowSignals("Primary", 300, 90, now.Add(2*time.Hour)), + codexWindowSignals("Secondary", weeklyWindowMinutes, 40, now.Add(day)), + map[string]string{"X-Codex-Limit-Reached": "true"}, + ) + + tests := []struct { + name string + provider string + auths []*Auth + want string + }{ + { + name: "claude 5h utilization", + provider: "claude", + auths: []*Auth{claudeAuthWithWindows("a", 1, now.Add(time.Hour), now.Add(day)), later}, + want: "b", + }, + {name: "claude 5h rejected", provider: "claude", auths: []*Auth{rejected, later}, want: "b"}, + { + name: "claude 5h reset passed", + provider: "claude", + auths: []*Auth{claudeAuthWithWindows("a", 1, now.Add(-time.Minute), now.Add(day)), later}, + want: "a", + }, + { + name: "every candidate gated keeps the ordering", + provider: "claude", + auths: []*Auth{ + claudeAuthWithWindows("a", 1, now.Add(time.Hour), now.Add(day)), + claudeAuthWithWindows("b", 1, now.Add(time.Hour), now.Add(3*day)), + }, + want: "a", + }, + { + name: "codex used 100", + provider: "codex", + auths: []*Auth{codexWeeklyAuth("a", 100, now.Add(day)), codexWeeklyAuth("b", 10, now.Add(3*day))}, + want: "b", + }, + { + name: "codex limit reached", + provider: "codex", + auths: []*Auth{limitReached, codexWeeklyAuth("b", 10, now.Add(3*day))}, + want: "b", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := pickSoonestReset(t, tt.provider, tt.auths...); got != tt.want { + t.Fatalf("Pick() = %q, want %q", got, tt.want) + } + }) + } +} + +// A Codex window without a Window-Minutes header still gates the credential when exhausted. +func TestSoonestResetSelectorPick_CodexWindowWithoutLengthStillGates(t *testing.T) { + t.Parallel() + + now := time.Now() + primary := codexWindowSignals("Primary", 0, 100, now.Add(time.Hour)) + delete(primary, "X-Codex-Primary-Window-Minutes") + a := codexAuthWithSignals("a", primary, codexWindowSignals("Secondary", weeklyWindowMinutes, 40, now.Add(day))) + if got := pickSoonestReset(t, "codex", a, codexWeeklyAuth("b", 10, now.Add(3*day))); got != "b" { + t.Fatalf("Pick() = %q, want b", got) + } +} + +// S5: an unknown Claude reset sorts first, an unknown or passed Codex reset sorts last, and a +// passed Claude reset rolls forward by whole weeks. +func TestSoonestResetSelectorPick_UnknownResets(t *testing.T) { + t.Parallel() + + now := time.Now() + if got := pickSoonestReset(t, "codex", + &Auth{ID: "a", Provider: "codex"}, + codexWeeklyAuth("b", 50, now.Add(-time.Hour)), + codexWeeklyAuth("c", 50, now.Add(6*day)), + ); got != "c" { + t.Fatalf("codex Pick() = %q, want c", got) + } + if got := pickSoonestReset(t, "claude", + claudeAuthResettingAt("a", now.Add(time.Hour)), + &Auth{ID: "b", Provider: "claude"}, + ); got != "b" { + t.Fatalf("claude Pick() = %q, want b", got) + } + // a's reset passed six days ago, so its next one is in one day. + if got := pickSoonestReset(t, "claude", + claudeAuthResettingAt("a", now.Add(2*day)), + claudeAuthResettingAt("b", now.Add(-6*day)), + ); got != "b" { + t.Fatalf("claude rolled-forward Pick() = %q, want b", got) + } + _, rolled := soonestResetRank(claudeAuthResettingAt("b", now.Add(-6*day)), now) + if want := time.Unix(now.Add(-6*day).Unix(), 0).Add(week); !rolled.Equal(want) { + t.Fatalf("rolled reset = %v, want %v", rolled, want) + } +} + +func newSoonestResetManager(t *testing.T, store CooldownStateStore, auths ...*Auth) *Manager { + t.Helper() + ctx := context.Background() + manager := NewManager(nil, nil, nil) + affinity := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{Fallback: &SoonestResetSelector{}, TTL: time.Hour}) + t.Cleanup(affinity.Stop) + manager.SetSelector(affinity) + manager.SetCooldownStateStore(store) + manager.RegisterExecutor(schedulerTestExecutor{provider: "codex"}) + for _, auth := range auths { + if _, err := manager.Register(WithSkipPersist(ctx), auth); err != nil { + t.Fatalf("Register(%s): %v", auth.ID, err) + } + registry.GetGlobalRegistry().RegisterClient(auth.ID, "codex", []*registry.ModelInfo{{ID: "soonest-reset-model"}}) + t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient(auth.ID) }) + } + return manager +} + +func pickSession(t *testing.T, manager *Manager, session string) string { + t.Helper() + opts := cliproxyexecutor.Options{Metadata: map[string]any{cliproxyexecutor.DerivedSessionIDMetadataKey: session}} + auth, _, err := manager.pickNext(context.Background(), "codex", "soonest-reset-model", opts, nil) + if err != nil { + t.Fatalf("pickNext() error = %v", err) + } + return auth.ID +} + +func markWithHeaders(manager *Manager, authID string, headers http.Header) { + ctx := internallogging.WithResponseHeadersHolder(context.Background()) + internallogging.SetResponseHeaders(ctx, headers) + manager.MarkResult(ctx, Result{AuthID: authID, Provider: "codex", Model: "soonest-reset-model", Success: true}) +} + +func codexWeeklyHeaders(resetAt time.Time) http.Header { + headers := http.Header{} + for key, value := range codexWindowSignals("Primary", weeklyWindowMinutes, 50, resetAt) { + headers.Set(key, value) + } + return headers +} + +// Routing reaches the selector through session affinity, and a bound session stays put. +func TestSoonestResetSelector_RoutedThroughSessionAffinity(t *testing.T) { + now := time.Now() + manager := newSoonestResetManager(t, nil, + codexWeeklyAuth("soonest-routing-a", 50, now.Add(3*day)), + codexWeeklyAuth("soonest-routing-b", 50, now.Add(day)), + ) + if manager.useSchedulerFastPath() { + t.Fatal("session affinity must use the legacy pick path") + } + if got := pickSession(t, manager, "first"); got != "soonest-routing-b" { + t.Fatalf("new session = %q, want soonest-routing-b", got) + } + + markWithHeaders(manager, "soonest-routing-a", codexWeeklyHeaders(now.Add(time.Hour))) + if got := pickSession(t, manager, "first"); got != "soonest-routing-b" { + t.Fatalf("bound session = %q, want soonest-routing-b", got) + } + if got := pickSession(t, manager, "second"); got != "soonest-routing-a" { + t.Fatalf("new session after observation = %q, want soonest-routing-a", got) + } +} + +// S10: observations saved before a restart order new sessions after it, with no new headers. +func TestSoonestResetSelector_PersistedObservationsSurviveRestart(t *testing.T) { + store := NewFileCooldownStateStore(t.TempDir()) + now := time.Now() + ids := []string{"soonest-restart-a", "soonest-restart-b"} + resets := []time.Time{now.Add(6 * day), now.Add(day)} + fresh := func() []*Auth { + auths := make([]*Auth, len(ids)) + for i, id := range ids { + auths[i] = &Auth{ID: id, Provider: "codex", Status: StatusActive} + } + return auths + } + + first := newSoonestResetManager(t, store, fresh()...) + for i, id := range ids { + markWithHeaders(first, id, codexWeeklyHeaders(resets[i])) + } + + second := newSoonestResetManager(t, store, fresh()...) + if got := pickSession(t, second, "before-restore"); got != ids[0] { + t.Fatalf("pick with no observations = %q, want %q (ID order)", got, ids[0]) + } + if err := second.RestoreCooldownStates(context.Background()); err != nil { + t.Fatalf("RestoreCooldownStates() error = %v", err) + } + if got := pickSession(t, second, "after-restore"); got != ids[1] { + t.Fatalf("pick after restore = %q, want %q", got, ids[1]) + } +} + +// Quota signals written by MarkResult and read by the selector do not race. +func TestSoonestResetSelector_ConcurrentObservationAndPick(t *testing.T) { + now := time.Now() + manager := newSoonestResetManager(t, nil, + &Auth{ID: "soonest-race-a", Provider: "codex", Status: StatusActive}, + &Auth{ID: "soonest-race-b", Provider: "codex", Status: StatusActive}, + ) + var wg sync.WaitGroup + for i := 0; i < 2; i++ { + wg.Add(2) + go func() { + defer wg.Done() + for j := 0; j < 50; j++ { + markWithHeaders(manager, "soonest-race-a", codexWeeklyHeaders(now.Add(time.Duration(j)*time.Hour))) + } + }() + go func() { + defer wg.Done() + for j := 0; j < 50; j++ { + opts := cliproxyexecutor.Options{Metadata: map[string]any{cliproxyexecutor.DerivedSessionIDMetadataKey: strconv.Itoa(j)}} + if _, _, err := manager.pickNext(context.Background(), "codex", "soonest-reset-model", opts, nil); err != nil { + t.Error(err) + } + } + }() + } + wg.Wait() +} diff --git a/sdk/cliproxy/service_config.go b/sdk/cliproxy/service_config.go index 9cac9e47a..91c6be59f 100644 --- a/sdk/cliproxy/service_config.go +++ b/sdk/cliproxy/service_config.go @@ -46,6 +46,8 @@ func normalizedRoutingRuntimeState(cfg *config.Config) routingRuntimeState { state.strategy = "weighted-round-robin" case "fill-first", "fillfirst", "ff": state.strategy = "fill-first" + case "soonest-reset", "soonestreset", "reset-first": + state.strategy = "soonest-reset" } state.sessionAffinity = cfg.Routing.SessionAffinity if ttl := strings.TrimSpace(cfg.Routing.SessionAffinityTTL); ttl != "" { @@ -69,6 +71,8 @@ func newRoutingSelector(state routingRuntimeState) coreauth.Selector { selector = &coreauth.WeightedRoundRobinSelector{} case "fill-first": selector = &coreauth.FillFirstSelector{} + case "soonest-reset": + selector = &coreauth.SoonestResetSelector{} default: selector = &coreauth.RoundRobinSelector{} } diff --git a/sdk/cliproxy/service_config_soonest_reset_test.go b/sdk/cliproxy/service_config_soonest_reset_test.go new file mode 100644 index 000000000..8415e087d --- /dev/null +++ b/sdk/cliproxy/service_config_soonest_reset_test.go @@ -0,0 +1,42 @@ +package cliproxy + +import ( + "context" + "strconv" + "testing" + "time" + + internalconfig "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + coreauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor" +) + +func TestSoonestResetRoutingSelectorWrappedBySessionAffinity(t *testing.T) { + selector := newRoutingSelector(normalizedRoutingRuntimeState(&internalconfig.Config{ + Routing: internalconfig.RoutingConfig{Strategy: "soonest-reset", SessionAffinity: true}, + })) + affinity, ok := selector.(*coreauth.SessionAffinitySelector) + if !ok { + t.Fatalf("selector type = %T, want *auth.SessionAffinitySelector", selector) + } + defer affinity.Stop() + + weekly := func(id string, resetAt time.Time) *coreauth.Auth { + return &coreauth.Auth{ID: id, Provider: "codex", Quota: coreauth.QuotaState{ + ObservedAt: time.Now(), + Signals: map[string]string{ + "X-Codex-Primary-Window-Minutes": "10080", + "X-Codex-Primary-Reset-At": strconv.FormatInt(resetAt.Unix(), 10), + }, + }} + } + auths := []*coreauth.Auth{weekly("a", time.Now().Add(72*time.Hour)), weekly("b", time.Now().Add(time.Hour))} + opts := cliproxyexecutor.Options{Metadata: map[string]any{cliproxyexecutor.DerivedSessionIDMetadataKey: "s"}} + got, err := affinity.Pick(context.Background(), "codex", "m", opts, auths) + if err != nil { + t.Fatalf("Pick() error = %v", err) + } + if got.ID != "b" { + t.Fatalf("Pick() = %q, want b", got.ID) + } +}