diff --git a/packages/opencode/src/index.ts b/packages/opencode/src/index.ts index a206a57..1691e63 100644 --- a/packages/opencode/src/index.ts +++ b/packages/opencode/src/index.ts @@ -1374,6 +1374,7 @@ export async function CodexAuthPlugin( async function writeRequestSidebarRouting( sessionId: string | undefined, + parentSessionId: string | undefined, activeId: string, route: RoutingMode, accounts: readonly { id: string; enabled?: boolean }[], @@ -1385,6 +1386,13 @@ export async function CodexAuthPlugin( accounts, boundSidebarFile, ) + if (parentSessionId && parentSessionId !== sessionId) { + await upsertSidebarActiveRouting( + { sessionId: parentSessionId, ...input }, + accounts, + boundSidebarFile, + ) + } return } await setSidebarLegacyRouting(input, boundSidebarFile) @@ -1987,9 +1995,10 @@ export async function CodexAuthPlugin( return { apiKey: OAUTH_DUMMY_KEY, async fetch(requestInput: RequestInfo | URL, init?: RequestInit) { - const sidebarSessionId = resolveSidebarSessionId( - effectiveRequestHeaders(requestInput, init), - ) + const requestHeaders = effectiveRequestHeaders(requestInput, init) + const sidebarSessionId = resolveSidebarSessionId(requestHeaders) + const sidebarParentSessionId = + requestHeaders.get('x-parent-session-id')?.trim() || undefined // Routing is purely mode-driven. The primary is ALWAYS the main // account; fallback-first is handled by a proactive gate below that // tries usable fallbacks before main. There is no per-account pin. @@ -2197,10 +2206,11 @@ export async function CodexAuthPlugin( await writeRequestSidebarRouting( sidebarSessionId, + sidebarParentSessionId, servedActiveId, mode, reqStorage?.accounts ?? [], - ) + ).catch(() => {}) return finalResponse }, async dispose() { diff --git a/packages/opencode/src/sidebar-state.ts b/packages/opencode/src/sidebar-state.ts index 0996bf1..83a25a1 100644 --- a/packages/opencode/src/sidebar-state.ts +++ b/packages/opencode/src/sidebar-state.ts @@ -295,6 +295,7 @@ export const ACTIVE_ROUTING_MAX_ENTRIES = 128 export type SidebarRoutingAccount = { id: string enabled?: boolean + killed?: boolean } export function isUsableRoutingEntry( @@ -308,11 +309,30 @@ export function isUsableRoutingEntry( return ( entry.activeId === 'main' || accounts.some( - (account) => account.enabled !== false && account.id === entry.activeId, + (account) => + account.enabled !== false && + account.killed !== true && + account.id === entry.activeId, ) ) } +export function isQuotaExhausted( + quota: AccountQuota | null | undefined, + now = Date.now(), +): boolean { + const primary = quota?.primary + if ( + typeof primary?.resetsAt !== 'string' || + !Number.isFinite(primary.usedPercent) || + primary.usedPercent < 100 + ) { + return false + } + const resetsAt = Date.parse(primary.resetsAt) + return Number.isFinite(resetsAt) && resetsAt > now +} + export function resolveSessionSidebarRouting( state: SidebarState, sessionId?: string, @@ -322,13 +342,24 @@ export function resolveSessionSidebarRouting( return { activeId: state.activeId ?? 'main', route: state.route } } const own = sessionId ? state.activeRouting?.[sessionId] : undefined - if (own && isUsableRoutingEntry(own, state.fallbacks, now)) { + const ownQuota = + own?.activeId === 'main' + ? state.main.quota + : state.fallbacks.find((account) => account.id === own?.activeId)?.quota + if ( + own && + isUsableRoutingEntry(own, state.fallbacks, now) && + !isQuotaExhausted(ownQuota, now) + ) { return { activeId: own.activeId, route: own.route } } - const fallback = state.fallbacks.find( + const enabledFallbacks = state.fallbacks.filter( (account) => account.enabled && !account.killed, ) + const fallback = + enabledFallbacks.find((account) => !isQuotaExhausted(account.quota, now)) ?? + enabledFallbacks[0] return { activeId: state.route === 'fallback-first' && fallback ? fallback.id : 'main', diff --git a/packages/opencode/src/tests/integration.test.ts b/packages/opencode/src/tests/integration.test.ts index 8d15364..ab193c8 100644 --- a/packages/opencode/src/tests/integration.test.ts +++ b/packages/opencode/src/tests/integration.test.ts @@ -1,5 +1,11 @@ import { afterEach, beforeEach, describe, expect, it, test } from 'bun:test' -import { mkdtempSync, readFileSync, writeFileSync } from 'node:fs' +import { + mkdirSync, + mkdtempSync, + readFileSync, + renameSync, + writeFileSync, +} from 'node:fs' import { readFile } from 'node:fs/promises' import { tmpdir } from 'node:os' import { join } from 'node:path' @@ -21,6 +27,7 @@ import { drainSidebarWrites, getSidebarStateFile, normalizeSidebarState, + resolveSessionSidebarRouting, type SidebarState, } from '../sidebar-state.ts' import { @@ -420,6 +427,106 @@ describe('integration: HTTP quota push', () => { } }) + it('records served routing for child and parent sessions', async () => { + writeFileSync( + configFile, + JSON.stringify({ + version: 1, + main: { type: 'opencode', provider: 'openai' }, + accounts: [], + }), + ) + + const originalFetch = globalThis.fetch + globalThis.fetch = (async () => + new Response('{}', { + status: 200, + headers: { 'content-type': 'application/json' }, + })) as unknown as typeof globalThis.fetch + + let hooks: Hooks | undefined + try { + hooks = await CodexAuthPlugin(createMockPluginInput(), { + experimentalWebSockets: false, + }) + const authHook = hooks.auth + if (!authHook?.loader) throw new Error('No auth loader') + const loaderResult = await authHook.loader( + async () => ({ + type: 'oauth' as const, + provider: 'openai', + access: accessToken, + refresh: refreshToken, + expires: Date.now() + 3600_000, + }), + { + id: 'openai', + label: 'OpenAI', + models: [], + } as unknown as Parameters>[1], + ) + const fetchOverride = (loaderResult as Record).fetch as + | ((url: RequestInfo | URL, init?: RequestInit) => Promise) + | undefined + if (!fetchOverride) throw new Error('No fetch in loader result') + + const serve = async (sessionId: string, parentId?: string) => { + const headers = new Headers({ + 'content-type': 'application/json', + 'session-id': sessionId, + }) + if (parentId) headers.set('x-parent-session-id', parentId) + const response = await fetchOverride( + 'https://api.openai.com/v1/responses', + { + method: 'POST', + headers, + body: JSON.stringify({ model: 'gpt-5.5', input: [] }), + }, + ) + expect(response.status).toBe(200) + await response.body?.cancel() + } + + await serve('child-session', 'parent-session') + await serve('child-only-session') + await serve('same-session', 'same-session') + await drainSidebarWrites() + + const sidebar = normalizeSidebarState( + JSON.parse(readFileSync(sidebarFile, 'utf8')), + ) + expect(Object.keys(sidebar.activeRouting ?? {}).sort()).toEqual([ + 'child-only-session', + 'child-session', + 'parent-session', + 'same-session', + ]) + expect(sidebar.activeRouting?.['child-session']).toMatchObject({ + activeId: 'main', + route: 'main-first', + }) + expect(sidebar.activeRouting?.['parent-session']).toEqual( + sidebar.activeRouting?.['child-session'], + ) + expect(sidebar.activeRouting?.['child-only-session']).toMatchObject({ + activeId: 'main', + route: 'main-first', + }) + expect(sidebar.activeRouting?.['same-session']).toMatchObject({ + activeId: 'main', + route: 'main-first', + }) + + renameSync(sidebarFile, `${sidebarFile}.saved`) + mkdirSync(sidebarFile) + await serve('unwritable-child', 'unwritable-parent') + } finally { + globalThis.fetch = originalFetch + await hooks?.dispose?.() + } + }) + it('a complete header frame reporting every window retired clears cached windows', async () => { const store = { version: 1, @@ -2627,6 +2734,88 @@ describe('integration: active fallback routing', () => { } }) + it('records fallback-served routing on the parent session', async () => { + seedStorage({ access: 'fallback-access-token' }) + const originalFetch = globalThis.fetch + globalThis.fetch = (async () => + new Response('{}', { status: 200 })) as unknown as typeof globalThis.fetch + + let hooks: Hooks | undefined + try { + const loaded = await loadFetchOverride( + createMockPluginInput(), + Date.now() + 3600_000, + ) + hooks = loaded.hooks + + await loaded.fetchOverride( + 'https://api.openai.com/v1/responses', + responseRequestInit({ + 'x-opencode-session': 'child-session', + 'x-parent-session-id': 'parent-session', + }), + ) + await drainSidebarWrites() + + const sidebar = normalizeSidebarState( + JSON.parse(readFileSync(sidebarFile, 'utf8')), + ) + expect(sidebar.activeRouting?.['child-session']).toMatchObject({ + activeId: 'fallback-1', + route: 'fallback-first', + }) + // A fallback (not main) served the child; the parent entry must mirror + // that same fallback so the parent's sidebar highlights the live account. + expect(sidebar.activeRouting?.['parent-session']).toMatchObject({ + activeId: 'fallback-1', + route: 'fallback-first', + }) + } finally { + globalThis.fetch = originalFetch + await hooks?.dispose?.() + } + }) + + it('uses legacy display routing when the request carries no session headers', async () => { + seedStorage({ access: 'fallback-access-token' }) + const originalFetch = globalThis.fetch + globalThis.fetch = (async () => + new Response('{}', { status: 200 })) as unknown as typeof globalThis.fetch + + let hooks: Hooks | undefined + try { + const loaded = await loadFetchOverride( + createMockPluginInput(), + Date.now() + 3600_000, + ) + hooks = loaded.hooks + + await loaded.fetchOverride( + 'https://api.openai.com/v1/responses', + responseRequestInit(), + ) + await drainSidebarWrites() + + const sidebar = normalizeSidebarState( + JSON.parse(readFileSync(sidebarFile, 'utf8')), + ) + // Sessionless requests write only the legacy display fields, never a + // per-session entry. + expect(sidebar.activeRouting).toBeUndefined() + expect(sidebar.activeId).toBe('fallback-1') + expect(sidebar.route).toBe('fallback-first') + // Resolving without a session reads those legacy fields and must yield + // defined routing rather than crash or return undefined. + expect(resolveSessionSidebarRouting(sidebar, undefined)).toEqual({ + activeId: 'fallback-1', + route: 'fallback-first', + }) + } finally { + globalThis.fetch = originalFetch + await hooks?.dispose?.() + } + }) + it('reads the sidebar session from a Request and strips it before the wire', async () => { seedEmptyAccountStorage() const originalFetch = globalThis.fetch diff --git a/packages/opencode/src/tests/sidebar-routing-tui.test.ts b/packages/opencode/src/tests/sidebar-routing-tui.test.ts index 6f73fd9..5ff29c7 100644 --- a/packages/opencode/src/tests/sidebar-routing-tui.test.ts +++ b/packages/opencode/src/tests/sidebar-routing-tui.test.ts @@ -5,6 +5,8 @@ import { join } from 'node:path' import { DEFAULT_SIDEBAR_STATE, drainSidebarWrites, + isQuotaExhausted, + isUsableRoutingEntry, normalizeSidebarState, removeSidebarActiveRouting, resolveSessionSidebarRouting, @@ -19,6 +21,67 @@ function state(overrides: Partial): SidebarState { const now = 2 * 60 * 60 * 1000 +describe('quota exhaustion guard', () => { + test('malformed persisted quota values stay fail-open', () => { + const future = new Date(now + 60_000).toISOString() + const window = { remainingPercent: 0, checkedAt: now, windowMinutes: 10080 } + expect( + isQuotaExhausted( + { + primary: { + ...window, + usedPercent: Number.NaN, + resetsAt: future, + }, + }, + now, + ), + ).toBe(false) + expect( + isQuotaExhausted( + { + primary: { + ...window, + usedPercent: '100' as unknown as number, + resetsAt: future, + }, + }, + now, + ), + ).toBe(false) + expect( + isQuotaExhausted( + { + primary: { + ...window, + usedPercent: 100, + resetsAt: 'not-a-date', + }, + }, + now, + ), + ).toBe(false) + expect( + isQuotaExhausted( + { + primary: { + ...window, + usedPercent: 100, + resetsAt: (now + 60_000) as unknown as string, + }, + }, + now, + ), + ).toBe(false) + expect( + isQuotaExhausted( + { primary: { ...window, usedPercent: 100, resetsAt: future } }, + now, + ), + ).toBe(true) + }) +}) + describe('session sidebar routing', () => { test('selects each session own usable routing entry', () => { const shared = state({ @@ -116,6 +179,193 @@ describe('session sidebar routing', () => { }) }) + test('derives a non-exhausted fallback when the session entry targets an exhausted fallback', () => { + const shared = state({ + route: 'fallback-first', + fallbacks: [ + { + id: 'work-alt', + label: 'Work', + quota: { + primary: { + usedPercent: 100, + remainingPercent: 0, + resetsAt: new Date(now + 60_000).toISOString(), + }, + }, + killed: false, + enabled: true, + }, + { + id: 'client-alt', + label: 'Client', + quota: { + primary: { + usedPercent: 6, + remainingPercent: 94, + resetsAt: new Date(now + 60_000).toISOString(), + }, + }, + killed: false, + enabled: true, + }, + ], + activeRouting: { + 'parent-session': { + activeId: 'work-alt', + route: 'fallback-first', + updatedAt: now, + }, + }, + }) + + expect(resolveSessionSidebarRouting(shared, 'parent-session', now)).toEqual( + { activeId: 'client-alt', route: 'fallback-first' }, + ) + }) + + test('derives a fallback when the session entry targets exhausted main', () => { + const shared = state({ + route: 'fallback-first', + main: { + quota: { + primary: { + usedPercent: 100, + remainingPercent: 0, + resetsAt: new Date(now + 60_000).toISOString(), + }, + }, + killed: false, + }, + fallbacks: [ + { + id: 'fallback-1', + label: 'Fallback', + quota: null, + killed: false, + enabled: true, + }, + ], + activeRouting: { + 'parent-session': { + activeId: 'main', + route: 'main-first', + updatedAt: now, + }, + }, + }) + + expect(resolveSessionSidebarRouting(shared, 'parent-session', now)).toEqual( + { activeId: 'fallback-1', route: 'fallback-first' }, + ) + }) + + test('honors an exhausted entry after its quota window reset has elapsed', () => { + const shared = state({ + route: 'fallback-first', + fallbacks: [ + { + id: 'fallback-1', + label: 'Fallback', + quota: { + primary: { + usedPercent: 100, + remainingPercent: 0, + resetsAt: new Date(now - 1).toISOString(), + }, + }, + killed: false, + enabled: true, + }, + ], + activeRouting: { + 'parent-session': { + activeId: 'fallback-1', + route: 'main-first', + updatedAt: now, + }, + }, + }) + + expect(resolveSessionSidebarRouting(shared, 'parent-session', now)).toEqual( + { activeId: 'fallback-1', route: 'main-first' }, + ) + }) + + test('honors entries when quota is missing or unknown', () => { + const shared = state({ + fallbacks: [ + { + id: 'missing-quota', + label: 'Missing quota', + quota: {}, + killed: false, + enabled: true, + }, + { + id: 'null-quota', + label: 'Null quota', + quota: null, + killed: false, + enabled: true, + }, + ], + activeRouting: { + missing: { + activeId: 'missing-quota', + route: 'fallback-first', + updatedAt: now, + }, + null: { + activeId: 'null-quota', + route: 'fallback-first', + updatedAt: now, + }, + }, + }) + + expect(resolveSessionSidebarRouting(shared, 'missing', now).activeId).toBe( + 'missing-quota', + ) + expect(resolveSessionSidebarRouting(shared, 'null', now).activeId).toBe( + 'null-quota', + ) + }) + + test('keeps the first enabled fallback when every fallback is exhausted', () => { + const exhaustedQuota = { + primary: { + usedPercent: 100, + remainingPercent: 0, + resetsAt: new Date(now + 60_000).toISOString(), + }, + } + const shared = state({ + route: 'fallback-first', + fallbacks: [ + { + id: 'fallback-1', + label: 'First', + quota: exhaustedQuota, + killed: false, + enabled: true, + }, + { + id: 'fallback-2', + label: 'Second', + quota: exhaustedQuota, + killed: false, + enabled: true, + }, + ], + }) + + expect(resolveSessionSidebarRouting(shared, 'missing', now)).toEqual({ + activeId: 'fallback-1', + route: 'fallback-first', + }) + }) + test('does not derive a killed fallback for a missing session', () => { const shared = state({ route: 'fallback-first', @@ -136,6 +386,16 @@ describe('session sidebar routing', () => { }) }) + test('treats a routing entry targeting a killed fallback as unusable', () => { + const accounts = [{ id: 'killed-fb', enabled: true, killed: true }] + const entry = { + activeId: 'killed-fb', + route: 'fallback-first', + updatedAt: now, + } + expect(isUsableRoutingEntry(entry, accounts, now)).toBe(false) + }) + test('present entry for a removed fallback derives instead of highlighting it', () => { const shared = state({ route: 'fallback-first',