fix: fence quota wrapper disposal and prove lock post-rename verification

Author: iceteaSA

packages/core/src/file-lock.test.ts

@@ -951,6 +951,13 @@ describe('acquireFencedFileLock — renewal TOCTOU', () => {
     const target = join(root, 'state.json')
     const lockPath = `${target}.accounts.lock`
     const clearIntervalSpy = spyOn(globalThis, 'clearInterval')
+    let renewalTick: (() => void) | null = null
+    const setIntervalSpy = spyOn(globalThis, 'setInterval').mockImplementation(
+      ((callback: () => void) => {
+        renewalTick = callback
+        return 0 as unknown as ReturnType<typeof setInterval>
+      }) as typeof setInterval,
+    )
 
     let pausedResolve!: () => void
     const paused = new Promise<void>((resolve) => {
@@ -967,7 +974,7 @@ describe('acquireFencedFileLock — renewal TOCTOU', () => {
         path: target,
         name: 'accounts',
         ttlMs: 60_000,
-        renewIntervalMs: 10,
+        renewIntervalMs: 60_000,
         onStep: async (step) => {
           if (step === 'renew-committed' && !hookFired) {
             hookFired = true
@@ -977,15 +984,25 @@ describe('acquireFencedFileLock — renewal TOCTOU', () => {
         },
       })
       expect(lock).not.toBeNull()
+      expect(renewalTick).not.toBeNull()
 
-      // A's renewal has just renamed its refreshed payload onto the
-      // lock file; owner B's acquisition lands in the post-commit
-      // window, before A's verify re-read.
+      // Drive exactly one renewal. With no later timer callback available,
+      // only the immediate post-rename re-read can observe B's takeover.
+      renewalTick!()
       await paused
       await writeLock(lockPath, 'owner-B', Date.now() + 60_000)
       gateResolve()
 
-      await lock!.whenLost()
+      let timeoutHandle: ReturnType<typeof setTimeout> | null = null
+      const loss = await Promise.race([
+        lock!.whenLost(),
+        new Promise<'timeout'>((resolve) => {
+          timeoutHandle = setTimeout(() => resolve('timeout'), 100)
+        }),
+      ]).finally(() => {
+        if (timeoutHandle !== null) clearTimeout(timeoutHandle)
+      })
+      expect(loss).toBeUndefined()
       expect(lock!.hasLost()).toBe(true)
       expect(clearIntervalSpy).toHaveBeenCalledTimes(1)
 
@@ -1000,6 +1017,7 @@ describe('acquireFencedFileLock — renewal TOCTOU', () => {
       await rm(lockPath, { force: true })
     } finally {
       clearIntervalSpy.mockRestore()
+      setIntervalSpy.mockRestore()
     }
   })
 

packages/core/src/file-lock.ts

@@ -396,13 +396,10 @@ function buildLock(
         }
         if (observed.expiresAt > now()) {
           if (released || lost) return
-          // TOCTOU protection: the lock file may have been swapped out
-          // from under us between the read above and the write below.
-          // Compare-and-swap via a sibling temp file + atomic rename so
-          // we never overwrite a fresh owner's lock content directly.
-          // The remaining race window — between the re-read and the
-          // rename — collapses to a single inode replace that POSIX
-          // rename(2) makes atomic at the filesystem level.
+          // Rename cannot compare-and-swap ownership. Re-read after the
+          // rename and mark this lock lost if another owner replaced it;
+          // writers call assertOwned() before mutating shared state, so a
+          // hijack in this seam is detected and stops future renewals.
           const tempPath = `${lockPath}.${ownerId}.tmp`
           let shouldCommit = false
           try {

packages/opencode/src/plugin/index.ts

@@ -32,15 +32,15 @@ import {
 import { createEventHandler } from './event-handler'
 import { createFetchInterceptor } from './fetch-interceptor'
 import { createGoogleSearchTool } from './google-search-tool'
-import { createPluginLifecycle } from './lifecycle'
+import { createPluginLifecycle, type PluginLifecycle } from './lifecycle'
 import { createLogger, initLogger, setRuntimeLogLevel } from './logger'
 import { createOAuthMethods, openBrowserWithSystem } from './oauth-methods'
 import {
   createOperatorSettingsController,
   type OperatorSettingsController,
 } from './operator-settings'
 import { persistAccountPool } from './persist-account-pool'
-import { createOpenCodeQuotaManager } from './quota'
+import { createOpenCodeQuotaManager, type QuotaManager } from './quota'
 import { createSessionRecoveryHook } from './recovery'
 import { initHealthTracker, initTokenTracker } from './rotation'
 import { AgySessionRegistry } from './session-context'
@@ -72,6 +72,13 @@ export interface CreateAntigravityPluginOptions {
   dependencies?: PluginDependencyOverrides
 }
 
+export function registerQuotaManagerProducer(
+  lifecycle: PluginLifecycle,
+  quotaManager: QuotaManager,
+): void {
+  lifecycle.register({ dispose: () => quotaManager.dispose() }, 'producer')
+}
+
 export const createAntigravityPlugin =
   (providerId: string, options: CreateAntigravityPluginOptions = {}) =>
   async (input: PluginInput): Promise<PluginResult> => {
@@ -146,7 +153,7 @@ export const createAntigravityPlugin =
     // final post-refresh write is enqueued before the drain flushes —
     // a consumer-phase registration could let a refresh enqueue a write
     // after drainSidebarWrites() already asserted the queue was empty.
-    lifecycle.register({ dispose: () => quotaManager.dispose() }, 'producer')
+    registerQuotaManagerProducer(lifecycle, quotaManager)
 
     // Operator settings controller backs the /antigravity-* slash commands.
     // The controller loads existing persisted settings at first read, mutates

packages/opencode/src/plugin/quota.test.ts

@@ -1,4 +1,12 @@
-import { afterEach, beforeEach, describe, expect, it } from 'bun:test'
+import {
+  afterEach,
+  beforeEach,
+  describe,
+  expect,
+  it,
+  mock,
+  spyOn,
+} from 'bun:test'
 import { mkdtempSync, rmSync } from 'node:fs'
 import { tmpdir } from 'node:os'
 import { join } from 'node:path'
@@ -7,12 +15,21 @@ import type { AccountMetadataV3 } from '@cortexkit/antigravity-auth-core'
 
 import {
   DEFAULT_SIDEBAR_STATE,
+  drainSidebarWrites,
   readSidebarState,
   SIDEBAR_STATE_ENV,
   SIDEBAR_STATE_VERSION,
   type SidebarStateV1,
+  setSidebarMergeHooks,
 } from '../sidebar-state'
-import { classifyQuotaGroup, pushSidebarQuotaSnapshot } from './quota.ts'
+import { registerQuotaManagerProducer } from './index.ts'
+import { createPluginLifecycle } from './lifecycle.ts'
+import {
+  classifyQuotaGroup,
+  createOpenCodeQuotaManager,
+  pushSidebarQuotaSnapshot,
+} from './quota.ts'
+import type { PluginClient } from './types.ts'
 
 interface QuotaSnapshotAccount {
   index: number
@@ -154,4 +171,98 @@ describe('pushSidebarQuotaSnapshot', () => {
     const state = read()
     expect(state.accounts).toEqual([])
   })
+
+  it('fences the real quota wrapper sidebar enqueue before the lifecycle drain', async () => {
+    const events: string[] = []
+    let releaseFetch!: () => void
+    const fetchGate = new Promise<void>((resolve) => {
+      releaseFetch = resolve
+    })
+    let fetchStartedResolve!: () => void
+    const fetchStarted = new Promise<void>((resolve) => {
+      fetchStartedResolve = resolve
+    })
+    const fetchSpy = spyOn(globalThis, 'fetch').mockImplementation(
+      (async () =>
+        new Response(
+          JSON.stringify({ access_token: 'access-token', expires_in: 3600 }),
+          { status: 200 },
+        )) as unknown as typeof fetch,
+    )
+    const client = {
+      auth: { set: mock(async () => {}) },
+    } as unknown as PluginClient
+    const account: AccountMetadataV3 = {
+      refreshToken: 'refresh-token',
+      managedProjectId: 'managed-project',
+      addedAt: 0,
+      lastUsed: 0,
+    }
+    const manager = createOpenCodeQuotaManager(client, 'google', {
+      getAccountsForSidebar: () => [
+        {
+          index: 0,
+          email: 'primary@example.test',
+          cachedQuota: {
+            claude: { remainingFraction: 0.42, modelCount: 1 },
+          },
+        },
+      ],
+      fetchVia: async () => {
+        events.push('fetch:start')
+        fetchStartedResolve()
+        await fetchGate
+        return new Response('unavailable', { status: 503 })
+      },
+    })
+    const lifecycle = createPluginLifecycle({
+      sessionRegistry: { clear: () => {} },
+      shutdownDiskSignatureCache: async () => {},
+      clearFetchState: () => {},
+      drainSidebarWrites: async () => {
+        events.push('lifecycle:drain')
+        await drainSidebarWrites()
+        events.push(
+          readSidebarState(stateFile).accounts.length === 1
+            ? 'drain:sees-sidebar-write'
+            : 'drain:misses-sidebar-write',
+        )
+      },
+    })
+    registerQuotaManagerProducer(lifecycle, manager)
+    setSidebarMergeHooks({
+      onStep: async (step) => {
+        if (step === 'await-lock') events.push('sidebar:write-start')
+      },
+    })
+
+    const refresh = manager.refreshAccounts([account], {
+      indexFor: () => 0,
+      force: true,
+    })
+    await fetchStarted
+    const dispose = lifecycle.dispose()
+    releaseFetch()
+
+    try {
+      await dispose
+      await manager.refreshAccounts([account], {
+        indexFor: () => 0,
+        force: true,
+      })
+      await drainSidebarWrites()
+      expect(events).toEqual([
+        'fetch:start',
+        'fetch:start',
+        'sidebar:write-start',
+        'lifecycle:drain',
+        'drain:sees-sidebar-write',
+      ])
+    } finally {
+      await refresh
+      await drainSidebarWrites()
+      setSidebarMergeHooks(null)
+      fetchSpy.mockRestore()
+    }
+  })
 })

packages/opencode/src/plugin/quota.ts

@@ -140,35 +140,68 @@ export function createOpenCodeQuotaManager(
   const originalRefreshAccount = manager.refreshAccount
   const originalRefreshAccounts = manager.refreshAccounts
   const getAccountsForSidebar = options.getAccountsForSidebar
+  let disposed = false
+  const inFlight = new Set<Promise<unknown>>()
 
-  const pushAfterRefresh = (account: AccountMetadataV3): void => {
+  const pushAfterRefresh = async (
+    account: AccountMetadataV3,
+  ): Promise<void> => {
     if (!getAccountsForSidebar) return
-    void pushSidebarQuotaSnapshot(
+    await pushSidebarQuotaSnapshot(
       getAccountsForSidebar,
       manager.getBackoffUntil(account),
     ).catch(() => {
-      // pushSidebarQuotaSnapshot already swallows and logs; the .catch
-      // here is just to keep the Promise from surfacing as an unhandled
-      // rejection when the lock was retried past the 2s budget.
+      // Sidebar persistence remains best-effort when lock contention
+      // outlives its retry budget.
     })
   }
 
+  const track = <T>(operation: Promise<T>): Promise<T> => {
+    inFlight.add(operation)
+    void operation.then(
+      () => inFlight.delete(operation),
+      () => inFlight.delete(operation),
+    )
+    return operation
+  }
+
+  const dispose = async (): Promise<void> => {
+    if (disposed) return
+    disposed = true
+    await manager.dispose()
+    await Promise.allSettled(inFlight)
+  }
+
   return {
     ...manager,
     async refreshAccount(account, refreshOptions) {
-      const result = await originalRefreshAccount(account, refreshOptions)
-      pushAfterRefresh(account)
-      return result
+      const shouldPush = !disposed
+      return track(
+        (async () => {
+          const result = await originalRefreshAccount(account, refreshOptions)
+          if (shouldPush) await pushAfterRefresh(account)
+          return result
+        })(),
+      )
     },
     async refreshAccounts(accounts, refreshOptions) {
-      const results = await originalRefreshAccounts(accounts, refreshOptions)
-      // Push one snapshot per batch — the AccountManager's view is updated
-      // by the caller (oauth-methods / fetch-interceptor) BEFORE we read
-      // here, so a single post-batch snapshot captures the full diff.
-      const lastAccount = accounts[accounts.length - 1]
-      if (lastAccount) pushAfterRefresh(lastAccount)
-      return results
+      const shouldPush = !disposed
+      return track(
+        (async () => {
+          const results = await originalRefreshAccounts(
+            accounts,
+            refreshOptions,
+          )
+          // Push one snapshot per batch — the AccountManager's view is updated
+          // by the caller (oauth-methods / fetch-interceptor) BEFORE we read
+          // here, so a single post-batch snapshot captures the full diff.
+          const lastAccount = accounts[accounts.length - 1]
+          if (shouldPush && lastAccount) await pushAfterRefresh(lastAccount)
+          return results
+        })(),
+      )
     },
+    dispose,
   }
 }