Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 11 additions & 1 deletion packages/core/src/accounts.ts
Original file line number Diff line number Diff line change
Expand Up @@ -129,6 +129,11 @@ export type AccountManagerOptions = {
fetchImpl?: typeof fetch
configPath?: string
quotaManager?: import('./quota-manager.ts').QuotaManager
// Invoked after a background quota pass persists at least one fallback storage
// change (token refresh, quota update, or error recording), so consumers
// (e.g. the OpenCode sidebar) can re-render without a request flowing through
// the fetch handler.
onFallbackStorageChanged?: () => void
}

export type AccountRefreshError = {
Expand Down Expand Up @@ -978,12 +983,14 @@ export class FallbackAccountManager {
private refreshTimer: ReturnType<typeof setInterval> | null = null
private quotaTimer: ReturnType<typeof setInterval> | null = null
readonly quotaManager: import('./quota-manager.ts').QuotaManager | null
private readonly onFallbackStorageChanged: (() => void) | undefined

constructor(options: AccountManagerOptions = {}) {
this.now = options.now ?? Date.now
this.fetchImpl = options.fetchImpl ?? fetch
this.configPath = options.configPath ?? getAccountStoragePath()
this.quotaManager = options.quotaManager ?? null
this.onFallbackStorageChanged = options.onFallbackStorageChanged
}

/**
Expand Down Expand Up @@ -1233,7 +1240,10 @@ export class FallbackAccountManager {
// Quota probes are advisory; failed probes fail closed at selection time.
}
}
if (changed) await this.save(storage)
if (changed) {
await this.save(storage)
this.onFallbackStorageChanged?.()
}
}

async refreshQuotaForAllAccounts(options: { force?: boolean } = {}) {
Expand Down
46 changes: 38 additions & 8 deletions packages/opencode/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -306,6 +306,9 @@ export const AnthropicAuthPlugin: Plugin = async (ctx) => {
})
const fallbackManager = new FallbackAccountManager({
quotaManager,
onFallbackStorageChanged: () => {
void refreshSidebarQuota()
},
})
fallbackManager.startBackgroundRefresh()
let latestRefreshMainAccessToken: (() => Promise<string>) | null = null
Expand Down Expand Up @@ -355,6 +358,13 @@ export const AnthropicAuthPlugin: Plugin = async (ctx) => {
setDumpEnabled(isDumpPersistentlyEnabled(initialStorage))
setFastModeEnabled(isFastModePersistentlyEnabled(initialStorage))

// Remembers the last explicit routing decision so quota-only sidebar refreshes
// (background main/fallback quota landing) do not reset the active account.
let lastSidebarRouting: { activeId: string | undefined; route: string } = {
activeId: 'main',
route: 'main',
}

function writeSidebarState(
storage: Awaited<ReturnType<typeof loadAccounts>>,
options: {
Expand All @@ -364,6 +374,7 @@ export const AnthropicAuthPlugin: Plugin = async (ctx) => {
mainRefreshToken?: string
},
) {
lastSidebarRouting = { activeId: options.activeId, route: options.route }
const mainEntry = quotaManager.getMain(options.mainAccessToken)
const lastApiError = quotaManager.getLastApiError()
const mainRefreshError = storage?.refresh?.mainLastRefreshError
Expand Down Expand Up @@ -425,6 +436,30 @@ export const AnthropicAuthPlugin: Plugin = async (ctx) => {
)
}

// Re-write the sidebar using the LAST known routing decision, refreshing only
// the quota numbers. Used by async quota refreshes (main + background fallback)
// so they never clobber the active account back to 'main'.
async function refreshSidebarQuota() {
const storage = await loadAccounts(accountStoragePath)
let access: string | undefined
let refresh: string | undefined
if (latestGetAuth) {
try {
const auth = await latestGetAuth()
access = auth.access
refresh = auth.refresh
} catch {
// best-effort
}
}
writeSidebarState(storage, {
activeId: lastSidebarRouting.activeId,
route: lastSidebarRouting.route,
mainAccessToken: access,
mainRefreshToken: refresh,
})
}

let latestGetAuth:
| (() => Promise<{
type: string
Expand Down Expand Up @@ -1690,14 +1725,9 @@ export const AnthropicAuthPlugin: Plugin = async (ctx) => {
// sidebar again once the new main quota lands.
void quotaManager
.refreshMain(auth.access)
.then(() =>
writeSidebarState(storage, {
activeId: 'main',
route: 'main',
mainAccessToken: auth.access,
mainRefreshToken: auth.refresh,
}),
)
.then(() => {
void refreshSidebarQuota()
})
.catch(() => {})
}
// Update the sidebar every replayable request so fallback
Expand Down
43 changes: 43 additions & 0 deletions packages/opencode/src/tests/accounts.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -707,6 +707,49 @@ describe('FallbackAccountManager', () => {
expect(fetchImpl).toHaveBeenCalledTimes(1)
})

test('refreshQuotaForDueAccounts fires onFallbackStorageChanged when storage changes', async () => {
const storage = baseStorage()
storage.accounts.push({
id: 'idle-stale',
type: 'oauth',
access: 'idle-access',
refresh: 'idle-refresh',
expires: 20_000_000,
quota: {
// checkedAt far in the past → stale → will be refreshed this pass.
five_hour: { usedPercent: 5, remainingPercent: 95, checkedAt: 1 },
seven_day: { usedPercent: 5, remainingPercent: 95, checkedAt: 1 },
},
})
await saveAccounts(storage)

const fetchImpl = mock(() =>
Promise.resolve(
new Response(
JSON.stringify({
five_hour: { utilization: 40 },
seven_day: { utilization: 30 },
}),
{ status: 200 },
),
),
) as unknown as typeof fetch

let fired = 0
const manager = new FallbackAccountManager({
fetchImpl,
now: () => 50_000_000, // well past checkedAt → stale
onFallbackStorageChanged: () => {
fired += 1
},
})

await manager.refreshQuotaForDueAccounts()

expect(fetchImpl).toHaveBeenCalled()
expect(fired).toBe(1)
})

test('refreshes fallback token and retries quota check after stale access token 401', async () => {
const storage = baseStorage()
storage.accounts.push({
Expand Down
150 changes: 150 additions & 0 deletions packages/opencode/src/tests/index.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import {
resetDumpState,
resetFastModeState,
saveAccounts,
tokenFingerprint,
} from '@cortexkit/anthropic-auth-core'
import { AnthropicAuthPlugin } from '../index'
import { getSidebarState } from '../sidebar-state'
Expand Down Expand Up @@ -2538,6 +2539,89 @@ describe('auth.loader', () => {
}
})

test('async main refresh does not clobber the active fallback in the sidebar', async () => {
const staleCheckedAt = Date.now() - 100 * 60_000 // far past → main quota is stale
await useTempAccountFile(
createFallbackStorage({
quota: {
enabled: true,
checkIntervalMinutes: 5,
minimumRemaining: { five_hour: 10, seven_day: 20 },
failClosedOnUnknownQuota: true,
// Cached, stale, and FAILING five_hour policy (used 95% → remaining 5% < 10%).
mainQuota: {
five_hour: {
usedPercent: 95,
remainingPercent: 5,
checkedAt: staleCheckedAt,
},
seven_day: {
usedPercent: 10,
remainingPercent: 90,
checkedAt: staleCheckedAt,
},
},
mainQuotaCheckedAt: staleCheckedAt,
// Bind the cached main quota to the access token the loader will use.
mainQuotaToken: tokenFingerprint('main-access'),
} as AccountStorage['quota'],
}),
)

globalThis.fetch = mock((input: any, init: any) => {
const url = extractUrl(input)
if (url.includes('/api/oauth/usage')) {
// Delay the background main refresh so it settles AFTER the fallback
// write — this is the race that causes the clobber in production.
return Bun.sleep(40).then(
() =>
new Response(
JSON.stringify({
five_hour: { utilization: 0.95 },
seven_day: { utilization: 0.1 },
}),
{ status: 200 },
),
)
}
// Anthropic messages call — fallback serves 200, main would not be reached.
return Promise.resolve(new Response('ok', { status: 200 }))
}) as unknown as typeof fetch

const plugin = await getPlugin()
const result = await plugin.auth.loader(
() =>
Promise.resolve({
type: 'oauth',
access: 'main-access',
refresh: 'main-refresh',
expires: Date.now() + 100000,
}),
{ models: {} },
)

await result.fetch(MESSAGES_URL, {
method: 'POST',
body: JSON.stringify({
model: 'claude-opus-4-8',
messages: [{ role: 'user', content: 'hello' }],
}),
})

// Fallback served → active id should be the fallback.
const state = await waitForSidebarState(
(candidate) => candidate.activeId === 'fallback-1',
)
expect(state.route).toBe('fallback')

// Let the fire-and-forget refreshMain().then(...) settle.
await Bun.sleep(80)

// REGRESSION: pre-fix the async callback rewrites activeId to 'main'.
const after = await getSidebarState()
expect(after.activeId).toBe('fallback-1')
})

test('fetch wrapper retries with fallback when main streaming body reports rate limit', async () => {
await useTempAccountFile(createFallbackStorage())
const authorizations: string[] = []
Expand Down Expand Up @@ -2690,4 +2774,70 @@ describe('auth.loader', () => {
expect(response.status).toBe(429)
expect(calls).toBe(1)
})

test('background fallback refresh updates the sidebar without a request', async () => {
await useTempAccountFile(
createFallbackStorage({
accounts: [
{
id: 'fallback-1',
type: 'oauth',
access: 'fallback-access',
refresh: 'fallback-refresh',
expires: Date.now() + 5 * 60 * 60 * 1000,
quota: {
// Stale (old checkedAt) → background pass will refresh it.
five_hour: {
usedPercent: 0,
remainingPercent: 100,
checkedAt: 1,
},
seven_day: {
usedPercent: 0,
remainingPercent: 100,
checkedAt: 1,
},
},
},
],
}),
)

globalThis.fetch = mock((input: any) => {
if (extractUrl(input).includes('/api/oauth/usage')) {
return Promise.resolve(
new Response(
JSON.stringify({
five_hour: { utilization: 0.42 },
seven_day: { utilization: 0.1 },
}),
{ status: 200 },
),
)
}
return Promise.resolve(new Response('ok', { status: 200 }))
}) as unknown as typeof fetch

const plugin = await getPlugin()
// Running the loader starts the background refresh (immediate first pass).
await plugin.auth.loader(
() =>
Promise.resolve({
type: 'oauth',
access: 'main-access',
refresh: 'main-refresh',
expires: Date.now() + 100000,
}),
{ models: {} },
)

// The background pass refreshes the stale fallback and the hook re-writes the
// sidebar — without any request to the messages endpoint.
// utilization: 0.42 → usedPercent: 0.42 (stored as-is, not multiplied by 100)
const state = await waitForSidebarState(
(candidate) =>
candidate.fallbacks[0]?.quota?.five_hour?.usedPercent === 0.42,
)
expect(state.fallbacks[0]?.id).toBe('fallback-1')
})
})