diff --git a/packages/opencode/src/auth/methods.ts b/packages/opencode/src/auth/methods.ts index 51848fa..008cd5c 100644 --- a/packages/opencode/src/auth/methods.ts +++ b/packages/opencode/src/auth/methods.ts @@ -60,6 +60,12 @@ export interface AuthMethodDependencies { readStoreIds: typeof readStoreIds openBrowser(url: string): boolean | undefined | Promise now(): number + custodyQuotaDeps: Pick< + RefreshAllQuotaDeps, + | 'isFallbackRefreshInert' + | 'resolveFallbackAccess' + | 'reportCustodyAuthFailure' + > } export interface CreateAuthMethodsOptions { @@ -187,6 +193,7 @@ export function createAuthMethods({ readStoreIds: dependencies?.readStoreIds ?? readStoreIds, openBrowser: dependencies?.openBrowser ?? openBrowserForMenu, now: dependencies?.now ?? Date.now, + custodyQuotaDeps: dependencies?.custodyQuotaDeps ?? {}, } const readAuth = async (): Promise => @@ -345,6 +352,7 @@ export function createAuthMethods({ whamFn: whamUsageFn, respectBackoff: false, readSidebarState: async () => ({ main: {}, fallbacks: [] }), + ...deps.custodyQuotaDeps, }) printQuotaResults(results) } diff --git a/packages/opencode/src/index.ts b/packages/opencode/src/index.ts index 0ee745d..93aedd6 100644 --- a/packages/opencode/src/index.ts +++ b/packages/opencode/src/index.ts @@ -1302,21 +1302,30 @@ export async function CodexAuthPlugin( let loaderGetAuth: | Parameters['loader']>>[0] | undefined + const custodyQuotaDepsForAuthMenu: Pick< + Parameters[0], + | 'isFallbackRefreshInert' + | 'resolveFallbackAccess' + | 'reportCustodyAuthFailure' + > = {} const authMethods = createAuthMethods({ client: input.client, getAuth: async () => loaderGetAuth?.(), fetchImpl: fetch, - dependencies: custodyOptions?.authorize - ? { - authorizeBrowser: custodyAuthorize( - custodyOptions.authorize.browser, - 'Complete authorization in your browser. This window will close automatically.', - ), - authorizeHeadless: custodyAuthorize( - custodyOptions.authorize.headless, - ), - } - : undefined, + dependencies: { + ...(custodyOptions?.authorize + ? { + authorizeBrowser: custodyAuthorize( + custodyOptions.authorize.browser, + 'Complete authorization in your browser. This window will close automatically.', + ), + authorizeHeadless: custodyAuthorize( + custodyOptions.authorize.headless, + ), + } + : {}), + custodyQuotaDeps: custodyQuotaDepsForAuthMenu, + }, }) const wrapCustodyAuthorize = ( @@ -1793,6 +1802,32 @@ export async function CodexAuthPlugin( }, }) } + async function resolveMainAccessForCustody( + currentStorage: Awaited>, + ): ReturnType { + if (claustrumMode(currentStorage) !== 'claustrum') { + return CUSTODY_EXCLUDED + } + const manifest = await readCustodyManifest() + const handle = lookupManifestHandle(manifest, 'main') + const cache = custodyRuntimeForDeps.getCache() + const refuse = ( + reason: 'no-handle' | 'blocked' | 'reauth' | 'cache-miss', + ): typeof CUSTODY_REFUSE => { + custodyLogger.warn('custody main request refused', { reason }) + return CUSTODY_REFUSE + } + if (!handle || !cache) return refuse('no-handle') + const now = (custodyOptions?.now ?? Date.now)() + if (cache.isBlocked(handle)) return refuse('blocked') + if (cache.isReauth(handle, now)) return refuse('reauth') + const served = await cache.peek(handle) + if (!served || served.expiresAtMs <= now) return refuse('cache-miss') + return { + token: served.payload.access, + provenance: { handle, recordVersion: served.recordVersion }, + } + } async function reportAuthFailureForCustody(params: { handle: string providerStatus: number @@ -1806,6 +1841,11 @@ export async function CodexAuthPlugin( recordVersion: params.recordVersion, }) } + Object.assign(custodyQuotaDepsForAuthMenu, { + isFallbackRefreshInert: isFallbackAccountRefreshInert, + resolveFallbackAccess: resolveAccountAccessForCustody, + reportCustodyAuthFailure: reportAuthFailureForCustody, + }) function buildRefreshAllQuotaDeps( overrides: Partial< Pick< @@ -3139,6 +3179,8 @@ export async function CodexAuthPlugin( sidebarState: SidebarState primaryAccess: string mainAccountIdentity?: string + mainCustodyRefused: boolean + primaryProvenance?: VaultProvenance }): Promise { const killswitchEnabled = isKillswitchEnabled(input.storage) const killswitchNow = Date.now() @@ -3164,22 +3206,27 @@ export async function CodexAuthPlugin( killswitchNow, ) : undefined - const roster: StickyRouteCandidate[] = [ - { - accountId: 'main', - wireAccountId: input.mainAccountIdentity, - access: input.primaryAccess, - keepwarmAccountKey: 'main', - quota: mainFreshest.quota, - quotaCheckedAt: mainFreshest.quotaCheckedAt, - reservePercent: getKillswitchThresholdsForAccount(input.storage), - configuredOrder: 0, - resetCreditsApplicable: resetCreditsApplicable( - mainFreshest.quota, - ), - killswitchPasses: mainKillswitchPasses, - }, - ] + const roster: StickyRouteCandidate[] = input.mainCustodyRefused + ? [] + : [ + { + accountId: 'main', + wireAccountId: input.mainAccountIdentity, + access: input.primaryAccess, + provenance: input.primaryProvenance, + keepwarmAccountKey: 'main', + quota: mainFreshest.quota, + quotaCheckedAt: mainFreshest.quotaCheckedAt, + reservePercent: getKillswitchThresholdsForAccount( + input.storage, + ), + configuredOrder: 0, + resetCreditsApplicable: resetCreditsApplicable( + mainFreshest.quota, + ), + killswitchPasses: mainKillswitchPasses, + }, + ] const usableFallbacks = await fallbackManager.getUsableFallbackAccounts(input.storage) if (!input.storage) return roster @@ -3738,31 +3785,47 @@ export async function CodexAuthPlugin( const myGeneration = ++mainIdentityGeneration if (currentAuth.type !== 'oauth') return fetch(requestInput, init) init = await materializeRequestInit(requestInput, init) - let primaryAccess: string = currentAuth.access ?? '' - - // Refresh expired main tokens and mirror them into opencode's slot. - if ( - !currentAuth.access || - (currentAuth.expires ?? 0) < Date.now() - ) { - logR.debug('token refresh triggered', { - pid: process.pid, - hasAccess: Boolean(currentAuth.access), - expiresInMs: currentAuth.expires - ? currentAuth.expires - Date.now() - : undefined, - }) - try { - const refreshed = await refreshMainWithLease() - currentAuth.access = refreshed.access - currentAuth.refresh = refreshed.refresh - currentAuth.expires = refreshed.expires - } catch (error) { - if (isAuthPersistError(error)) throw error - // Use stale token on refresh failure + const mainCustodyOwned = + recognizedMainTombstone && + claustrumMode(reqStorage) === 'claustrum' + let primaryAccess = '' + let primaryProvenance: VaultProvenance | undefined + let mainCustodyRefused = false + + if (mainCustodyOwned) { + const access = await resolveMainAccessForCustody(reqStorage) + if (access === CUSTODY_REFUSE || access === CUSTODY_EXCLUDED) { + mainCustodyRefused = true + } else { + primaryAccess = access.token + primaryProvenance = + access.provenance === 'local' ? undefined : access.provenance + } + } else { + // Refresh expired main tokens and mirror them into opencode's slot. + if ( + !currentAuth.access || + (currentAuth.expires ?? 0) < Date.now() + ) { + logR.debug('token refresh triggered', { + pid: process.pid, + hasAccess: Boolean(currentAuth.access), + expiresInMs: currentAuth.expires + ? currentAuth.expires - Date.now() + : undefined, + }) + try { + const refreshed = await refreshMainWithLease() + currentAuth.access = refreshed.access + currentAuth.refresh = refreshed.refresh + currentAuth.expires = refreshed.expires + } catch (error) { + if (isAuthPersistError(error)) throw error + // Use stale token on refresh failure + } } + primaryAccess = currentAuth.access ?? '' } - primaryAccess = currentAuth.access ?? '' const authWithAccount = currentAuth as typeof currentAuth & { accountId?: string @@ -3804,6 +3867,8 @@ export async function CodexAuthPlugin( sidebarState, primaryAccess, mainAccountIdentity, + mainCustodyRefused, + primaryProvenance, }) let stickyCandidate = await resolveStickyRouteCandidate({ sessionId: sidebarSessionId, @@ -4061,14 +4126,16 @@ export async function CodexAuthPlugin( ) } else { // Send through the main account. - response = await sendWithAccessToken( - requestInput, - init, - primaryAccess, - mainAccountIdentity, - 'main', - undefined, - ) + response = mainCustodyRefused + ? new Response(null, { status: 401 }) + : await sendWithAccessToken( + requestInput, + init, + primaryAccess, + mainAccountIdentity, + 'main', + primaryProvenance, + ) } } diff --git a/packages/opencode/src/tests/auth-menu.test.ts b/packages/opencode/src/tests/auth-menu.test.ts index 9d41c9c..81a525d 100644 --- a/packages/opencode/src/tests/auth-menu.test.ts +++ b/packages/opencode/src/tests/auth-menu.test.ts @@ -496,6 +496,72 @@ describe('OpenCode auth menu', () => { await expectMenuCompletionFailed(result) }) + test('Check quotas serves a custody fallback through the injected vault resolver', async () => { + const paths = tempPaths() + const now = Date.now() + await seedStore(paths, [ + account('custodied', { + access: '', + refresh: 'claustrum-tombstone:v1:openai', + }), + ]) + const main = { + type: 'oauth', + refresh: 'main-refresh', + access: 'main-access', + expires: now + 86_400_000, + } + const { client, getAuth } = createClient(main) + const fetchedAuthorization: string[] = [] + const fetchImpl = mock(async (_input: unknown, init?: RequestInit) => { + fetchedAuthorization.push( + new Headers(init?.headers).get('authorization') ?? '', + ) + return new Response( + JSON.stringify({ + rate_limit: { + primary_window: { + used_percent: 25, + limit_window_seconds: 18_000, + }, + }, + }), + { status: 200, headers: { 'content-type': 'application/json' } }, + ) + }) + const isFallbackRefreshInert = mock(async () => true) + const resolveFallbackAccess = mock(async () => ({ + token: 'vault-fallback-access', + provenance: { handle: 'vault-handle', recordVersion: 7 }, + })) + const reportCustodyAuthFailure = mock(async () => {}) + const methods = createAuthMethods({ + client, + getAuth, + getPaths: () => paths, + fetchImpl: fetchImpl as unknown as typeof fetch, + dependencies: { + showAuthMenu: async () => 'check-quotas', + custodyQuotaDeps: { + isFallbackRefreshInert, + resolveFallbackAccess, + reportCustodyAuthFailure, + } as never, + } as never, + }) + spyOn(console, 'log').mockImplementation(() => {}) + + const result = await oauthMethod(methods, 0).authorize({}) + + expect(isFallbackRefreshInert).toHaveBeenCalledTimes(1) + expect(resolveFallbackAccess).toHaveBeenCalledTimes(1) + expect(fetchedAuthorization).toEqual([ + 'Bearer main-access', + 'Bearer vault-fallback-access', + ]) + await expectMenuCompletionFailed(result) + }) + test('Auth doctor leaves both store files byte-unchanged', async () => { const paths = tempPaths() await seedStore(paths, [account('main', { accountId: 'chatgpt-main' })]) diff --git a/packages/opencode/src/tests/custody-main.test.ts b/packages/opencode/src/tests/custody-main.test.ts index 232903e..2028117 100644 --- a/packages/opencode/src/tests/custody-main.test.ts +++ b/packages/opencode/src/tests/custody-main.test.ts @@ -425,7 +425,7 @@ describe('main host slot', () => { ) }) - test('does not acquire refresh state or call the token endpoint for a tombstoned main slot', async () => { + test('does not acquire refresh state or send a tombstoned main slot', async () => { const originalFetch = globalThis.fetch const urls: string[] = [] globalThis.fetch = (async (url: string | URL | Request) => { @@ -465,9 +465,7 @@ describe('main host slot', () => { body: '{}', }, ) - expect(urls).toEqual([ - 'https://chatgpt.com/backend-api/codex/responses', - ]) + expect(urls).toEqual([]) expect( (await loadAccounts(getAccountPaths(configPath)))?.refresh ?.mainRefreshLeaseId, diff --git a/packages/opencode/src/tests/custody-request.test.ts b/packages/opencode/src/tests/custody-request.test.ts index 176d5ad..9dc4829 100644 --- a/packages/opencode/src/tests/custody-request.test.ts +++ b/packages/opencode/src/tests/custody-request.test.ts @@ -59,6 +59,13 @@ async function withCustodyLoader( accounts: OAuthAccount[] routing?: { mode: 'main-first' | 'fallback-first' | 'sticky-balanced' } claustrumEnabled?: boolean + manifestLabel?: string + mainAuth?: { + type: 'oauth' + access: string + refresh: string + expires: number + } credential?: { material: string; recordVersion: number } | undefined credentialForGet?: () => { material: string; recordVersion: number } now?: () => number @@ -95,7 +102,9 @@ async function withCustodyLoader( const directory = mkdtempSync(join(tmpdir(), 'custody-request-loader-')) const configPath = join(directory, 'openai-auth.json') const manifestPath = join(directory, 'handles.json') - const manifest = enrollmentManifest(options.accounts[0]?.id ?? 'custody-1') + const manifest = enrollmentManifest( + options.manifestLabel ?? options.accounts[0]?.id ?? 'custody-1', + ) if (!manifest.ok) throw new Error('expected manifest fixture') const originalFetch = globalThis.fetch const authorizations: string[] = [] @@ -195,12 +204,13 @@ async function withCustodyLoader( const loader = hooks.auth?.loader if (!loader) throw new Error('expected auth loader') const result = await loader( - async () => ({ - type: 'oauth' as const, - access: 'main-access', - refresh: 'main-refresh', - expires: Date.now() + 3_600_000, - }), + async () => + options.mainAuth ?? { + type: 'oauth' as const, + access: 'main-access', + refresh: 'main-refresh', + expires: Date.now() + 3_600_000, + }, {} as never, ) const fetchOverride = (result as { fetch?: typeof globalThis.fetch }).fetch @@ -714,6 +724,80 @@ describe('custody request resolution', () => { ) }) + it('uses the vault bearer for a tombstoned main send', async () => { + const vaultAccess = 'VAULT-MAIN-TOKEN-xyz' + await withCustodyLoader( + { + accounts: [], + manifestLabel: 'main', + mainAuth: { + type: 'oauth', + access: '', + refresh: TOMBSTONE_OPENAI, + expires: 0, + }, + credential: { material: vaultAccess, recordVersion: 71 }, + respond: () => 401, + }, + async ({ fetchOverride, authorizations, gets, reports }) => { + const [url, init] = codexRequest() + expect((await fetchOverride(url, init)).status).toBe(401) + expect(gets()).toBeGreaterThan(0) + expect(authorizations).toEqual([`Bearer ${vaultAccess}`]) + await new Promise((resolve) => setTimeout(resolve, 0)) + expect(reports).toEqual([ + { recordVersion: 71, reporterSource: 'direct' }, + ]) + }, + ) + }) + + it('refuses a tombstoned main without a cached vault bearer and serves a fallback', async () => { + const fallback = liveAccount('fallback', { accountId: 'acct-fallback' }) + await withCustodyLoader( + { + accounts: [fallback], + manifestLabel: 'main', + mainAuth: { + type: 'oauth', + access: 'stale-main-access', + refresh: TOMBSTONE_OPENAI, + expires: Date.now() + 60_000, + }, + credential: undefined, + respond: (authorization) => + authorization === `Bearer ${fallback.access}` ? 200 : 401, + }, + async ({ fetchOverride, authorizations }) => { + const [url, init] = codexRequest() + expect((await fetchOverride(url, init)).status).toBe(200) + expect(authorizations).toEqual([`Bearer ${fallback.access}`]) + }, + ) + }) + + it('serves real main material in local mode', async () => { + await withCustodyLoader( + { + accounts: [], + claustrumEnabled: false, + mainAuth: { + type: 'oauth', + access: 'local-main-access', + refresh: 'local-main-refresh', + expires: Date.now() + 60_000, + }, + credential: undefined, + respond: () => 200, + }, + async ({ fetchOverride, authorizations }) => { + const [url, init] = codexRequest() + expect((await fetchOverride(url, init)).status).toBe(200) + expect(authorizations).toEqual(['Bearer local-main-access']) + }, + ) + }) + it('reports the served vault version after a reactive fallback 401', async () => { const fallback = makeSentinelAccount({ id: 'reactive',