Serve the vault bearer for a tombstoned main

packages/opencode/src/auth/methods.ts

@@ -60,6 +60,12 @@ export interface AuthMethodDependencies {
   readStoreIds: typeof readStoreIds
   openBrowser(url: string): boolean | undefined | Promise<boolean | undefined>
   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<AuthDetails> =>
@@ -345,6 +352,7 @@ export function createAuthMethods({
       whamFn: whamUsageFn,
       respectBackoff: false,
       readSidebarState: async () => ({ main: {}, fallbacks: [] }),
+      ...deps.custodyQuotaDeps,
     })
     printQuotaResults(results)
   }

packages/opencode/src/index.ts

@@ -1302,21 +1302,30 @@ export async function CodexAuthPlugin(
   let loaderGetAuth:
     | Parameters<NonNullable<NonNullable<Hooks['auth']>['loader']>>[0]
     | undefined
+  const custodyQuotaDepsForAuthMenu: Pick<
+    Parameters<typeof refreshAllQuota>[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<typeof loadAccounts>>,
+        ): ReturnType<typeof resolveFallbackAccess> {
+          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<StickyRouteCandidate[]> {
           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,
+                    )
               }
             }
 

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' })])

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,