diff --git a/packages/opencode/src/index.ts b/packages/opencode/src/index.ts index ed3c1f3..bf87b38 100644 --- a/packages/opencode/src/index.ts +++ b/packages/opencode/src/index.ts @@ -2242,6 +2242,12 @@ export async function CodexAuthPlugin( candidate.quotaCheckedAt, ]), ) + const wireAccountIdByAccount = Object.fromEntries( + eligibleCandidates.map((candidate) => [ + candidate.accountId, + candidate.wireAccountId, + ]), + ) const excluded = new Set(input.excludeAccountIds) let placement: | { @@ -2258,6 +2264,7 @@ export async function CodexAuthPlugin( validPinnedAccountIds: [...candidatesById.keys()], excludeAccountIds: input.excludeAccountIds, quotaCheckedAtByAccount, + wireAccountIdByAccount, choose: (pendingBytes) => { const eligible = eligibleCandidates.filter( (candidate) => !excluded.has(candidate.accountId), diff --git a/packages/opencode/src/sidebar-state.ts b/packages/opencode/src/sidebar-state.ts index 466d578..7c13015 100644 --- a/packages/opencode/src/sidebar-state.ts +++ b/packages/opencode/src/sidebar-state.ts @@ -100,6 +100,7 @@ export type ActiveRoutingMap = Record export interface StickyAssignment { accountId: string + wireAccountId?: string assignedAt: number lastSeenAt: number inputBytes: number @@ -120,6 +121,7 @@ export interface ResolveStickyAssignmentInput { validPinnedAccountIds: readonly string[] excludeAccountIds?: readonly string[] quotaCheckedAtByAccount: Readonly> + wireAccountIdByAccount?: Readonly> choose: ( pendingBytes: ReadonlyMap, ) => StickyAssignmentChoice | undefined @@ -256,12 +258,14 @@ function normalizeStickyAssignments( ) { continue } + const wireAccountId = assignment.wireAccountId normalized[sessionHash] = { accountId: assignment.accountId, assignedAt: assignment.assignedAt, lastSeenAt: assignment.lastSeenAt, inputBytes: assignment.inputBytes, ...(quotaCheckedAt === undefined ? {} : { quotaCheckedAt }), + ...(typeof wireAccountId === 'string' ? { wireAccountId } : {}), } } @@ -652,10 +656,24 @@ function stickyAssignmentNeedsMetadataUpdate( assignment: StickyAssignment, requestBytes: number, now: number, + wireAccountId: string | undefined, ): boolean { return ( requestBytes > assignment.inputBytes || - now - assignment.lastSeenAt >= STICKY_ASSIGNMENT_LAST_SEEN_TOUCH_MS + now - assignment.lastSeenAt >= STICKY_ASSIGNMENT_LAST_SEEN_TOUCH_MS || + (assignment.wireAccountId === undefined && wireAccountId !== undefined) + ) +} + +function hasStickyIdentityMismatch( + assignment: StickyAssignment, + wireAccountId: string | undefined, +): boolean { + // Missing identity metadata must retain the cache-warm pin until a known change proves it stale. + return ( + typeof assignment.wireAccountId === 'string' && + typeof wireAccountId === 'string' && + assignment.wireAccountId !== wireAccountId ) } @@ -692,9 +710,11 @@ function readonlyPendingBytes( function pendingBytesForAssignments( assignments: StickyAssignmentMap | undefined, quotaCheckedAtByAccount: Readonly>, + excludedSessionHash?: string, ): ReadonlyMap { const pendingBytes = new Map() - for (const assignment of Object.values(assignments ?? {})) { + for (const [sessionHash, assignment] of Object.entries(assignments ?? {})) { + if (sessionHash === excludedSessionHash) continue if ( assignment.quotaCheckedAt !== quotaCheckedAtByAccount[assignment.accountId] @@ -1178,10 +1198,15 @@ export async function resolveSidebarStickyAssignment( excludedAccountIds, input.now, ) && + !hasStickyIdentityMismatch( + existing, + input.wireAccountIdByAccount?.[existing.accountId], + ) && !stickyAssignmentNeedsMetadataUpdate( existing, input.requestBytes, input.now, + input.wireAccountIdByAccount?.[existing.accountId], ) ) { return existing @@ -1202,23 +1227,38 @@ export async function resolveSidebarStickyAssignment( stickyAssignments, ) const current = stickyAssignments?.[sessionHash] + const currentIdentityMismatch = + current !== undefined && + hasStickyIdentityMismatch( + current, + input.wireAccountIdByAccount?.[current.accountId], + ) if ( isValidStickyAssignment( current, validPinnedAccountIds, excludedAccountIds, input.now, - ) + ) && + !currentIdentityMismatch ) { const metadataNeedsUpdate = stickyAssignmentNeedsMetadataUpdate( current, input.requestBytes, input.now, + input.wireAccountIdByAccount?.[current.accountId], ) const assignment = metadataNeedsUpdate ? { ...current, inputBytes: Math.max(current.inputBytes, input.requestBytes), + ...(current.wireAccountId === undefined && + input.wireAccountIdByAccount?.[current.accountId] !== undefined + ? { + wireAccountId: + input.wireAccountIdByAccount?.[current.accountId], + } + : {}), ...(input.now - current.lastSeenAt >= STICKY_ASSIGNMENT_LAST_SEEN_TOUCH_MS ? { lastSeenAt: input.now } @@ -1241,6 +1281,8 @@ export async function resolveSidebarStickyAssignment( pendingBytesForAssignments( stickyAssignments, input.quotaCheckedAtByAccount, + // A mismatched pin belongs to a prior identity and cannot steer its replacement. + currentIdentityMismatch ? sessionHash : undefined, ), ) if (!choice) { @@ -1261,6 +1303,11 @@ export async function resolveSidebarStickyAssignment( ...(choice.quotaCheckedAt === undefined ? {} : { quotaCheckedAt: choice.quotaCheckedAt }), + ...(input.wireAccountIdByAccount?.[choice.accountId] === undefined + ? {} + : { + wireAccountId: input.wireAccountIdByAccount?.[choice.accountId], + }), } return { ...latest, diff --git a/packages/opencode/src/tests/integration.test.ts b/packages/opencode/src/tests/integration.test.ts index 53e6185..ba97733 100644 --- a/packages/opencode/src/tests/integration.test.ts +++ b/packages/opencode/src/tests/integration.test.ts @@ -2977,6 +2977,7 @@ describe('integration: active fallback routing', () => { const firstAssignment = firstState.stickyAssignments?.[hashSidebarSessionId('sticky-session')] expect(firstAssignment?.accountId).toBe('fallback-2') + expect(firstAssignment?.wireAccountId).toBe('acc-fallback-2') const changed = JSON.parse(readFileSync(sidebarFile, 'utf8')) changed.fallbacks[0].quota = stickyQuota(100, Date.now()) @@ -3005,6 +3006,10 @@ describe('integration: active fallback routing', () => { finalState.stickyAssignments?.[hashSidebarSessionId('cold-session')] ?.accountId, ).toBe('fallback-1') + expect( + finalState.stickyAssignments?.[hashSidebarSessionId('cold-session')] + ?.wireAccountId, + ).toBe('acc-fallback-1') } finally { globalThis.fetch = originalFetch await hooks?.dispose?.() diff --git a/packages/opencode/src/tests/sidebar-state.test.ts b/packages/opencode/src/tests/sidebar-state.test.ts index e51e656..1094c93 100644 --- a/packages/opencode/src/tests/sidebar-state.test.ts +++ b/packages/opencode/src/tests/sidebar-state.test.ts @@ -408,6 +408,46 @@ describe('normalizeSidebarState', () => { }) }) + test('normalizes a malformed sticky wire identity away without dropping the pin', () => { + const sessionHash = hashSidebarSessionId('malformed-sticky-identity') + + expect(() => + normalizeSidebarState({ + ...DEFAULT_SIDEBAR_STATE, + stickyAssignments: { + [sessionHash]: { + accountId: 'fallback-1', + wireAccountId: 42, + assignedAt: 100, + lastSeenAt: 200, + inputBytes: 300, + }, + }, + }), + ).not.toThrow() + expect( + normalizeSidebarState({ + ...DEFAULT_SIDEBAR_STATE, + stickyAssignments: { + [sessionHash]: { + accountId: 'fallback-1', + wireAccountId: 42, + assignedAt: 100, + lastSeenAt: 200, + inputBytes: 300, + }, + }, + }).stickyAssignments, + ).toEqual({ + [sessionHash]: { + accountId: 'fallback-1', + assignedAt: 100, + lastSeenAt: 200, + inputBytes: 300, + }, + }) + }) + test('old files without sticky assignments remain valid', () => { expect( normalizeSidebarState(DEFAULT_SIDEBAR_STATE).stickyAssignments, @@ -596,6 +636,274 @@ describe('sticky assignments', () => { expect(result).toEqual(assignment) }) + test('replaces a pin when its known wire identity changes', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'oai-sb-sticky-identity-')) + const file = join(tempDir, 'sidebar-state.json') + const sessionId = 'identity-mismatch-session' + await setSidebarState( + make({ + stickyAssignments: { + [hashSidebarSessionId(sessionId)]: { + accountId: 'account-a', + wireAccountId: 'chatgpt-a-old', + assignedAt: now - 1, + lastSeenAt: now, + inputBytes: 128, + quotaCheckedAt: 10, + }, + }, + }), + file, + ) + let chooseCalls = 0 + + const result = await resolveSidebarStickyAssignment( + { + sessionId, + requestBytes: 256, + now, + validPinnedAccountIds: ['account-a', 'account-b'], + quotaCheckedAtByAccount: { 'account-a': 10, 'account-b': 20 }, + wireAccountIdByAccount: { + 'account-a': 'chatgpt-a-new', + 'account-b': 'chatgpt-b-new', + }, + choose: () => { + chooseCalls += 1 + return { accountId: 'account-b', quotaCheckedAt: 20 } + }, + }, + file, + ) + + expect(chooseCalls).toBe(1) + expect(result).toEqual({ + accountId: 'account-b', + wireAccountId: 'chatgpt-b-new', + assignedAt: now, + lastSeenAt: now, + inputBytes: 256, + quotaCheckedAt: 20, + }) + }) + + test('excludes a mismatched pin from replacement pending bytes while retaining other valid pins', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'oai-sb-sticky-identity-')) + const file = join(tempDir, 'sidebar-state.json') + const sessionId = 'identity-mismatch-pending-session' + await setSidebarState( + make({ + stickyAssignments: { + [hashSidebarSessionId(sessionId)]: { + accountId: 'account-a', + wireAccountId: 'chatgpt-a-old', + assignedAt: now - 1, + lastSeenAt: now, + inputBytes: 100, + quotaCheckedAt: 10, + }, + [hashSidebarSessionId('other-valid-session')]: { + accountId: 'account-b', + wireAccountId: 'chatgpt-b', + assignedAt: now - 1, + lastSeenAt: now, + inputBytes: 20, + quotaCheckedAt: 20, + }, + }, + }), + file, + ) + let pendingBytes: ReadonlyMap | undefined + + const result = await resolveSidebarStickyAssignment( + { + sessionId, + requestBytes: 50, + now, + validPinnedAccountIds: ['account-a', 'account-b'], + quotaCheckedAtByAccount: { 'account-a': 10, 'account-b': 20 }, + wireAccountIdByAccount: { + 'account-a': 'chatgpt-a-new', + 'account-b': 'chatgpt-b', + }, + choose: (pending) => { + pendingBytes = pending + return pending.get('account-a') === undefined && + pending.get('account-b') === 20 + ? { accountId: 'account-a', quotaCheckedAt: 10 } + : { accountId: 'account-b', quotaCheckedAt: 20 } + }, + }, + file, + ) + + expect(pendingBytes?.get('account-a')).toBeUndefined() + expect(pendingBytes?.get('account-b')).toBe(20) + expect(result?.accountId).toBe('account-a') + }) + + test('keeps a pre-upgrade pin in replacement pending bytes', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'oai-sb-sticky-identity-')) + const file = join(tempDir, 'sidebar-state.json') + const sessionId = 'pre-upgrade-pending-session' + await setSidebarState( + make({ + stickyAssignments: { + [hashSidebarSessionId(sessionId)]: { + accountId: 'account-a', + assignedAt: now - 1, + lastSeenAt: now, + inputBytes: 100, + quotaCheckedAt: 10, + }, + }, + }), + file, + ) + let pendingBytes: ReadonlyMap | undefined + + await resolveSidebarStickyAssignment( + { + sessionId, + requestBytes: 50, + now, + validPinnedAccountIds: ['account-a', 'account-b'], + excludeAccountIds: ['account-a'], + quotaCheckedAtByAccount: { 'account-a': 10, 'account-b': 20 }, + wireAccountIdByAccount: { 'account-a': 'chatgpt-a' }, + choose: (pending) => { + pendingBytes = pending + return { accountId: 'account-b', quotaCheckedAt: 20 } + }, + }, + file, + ) + + expect(pendingBytes?.get('account-a')).toBe(100) + }) + + test('retains a pre-upgrade pin and stamps its known current wire identity', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'oai-sb-sticky-identity-')) + const file = join(tempDir, 'sidebar-state.json') + const sessionId = 'identity-migration-session' + const assignment = { + accountId: 'account-a', + assignedAt: now - 1, + lastSeenAt: now, + inputBytes: 128, + quotaCheckedAt: 10, + } + await setSidebarState( + make({ + stickyAssignments: { [hashSidebarSessionId(sessionId)]: assignment }, + }), + file, + ) + + const result = await resolveSidebarStickyAssignment( + { + sessionId, + requestBytes: 128, + now, + validPinnedAccountIds: ['account-a'], + quotaCheckedAtByAccount: { 'account-a': 10 }, + wireAccountIdByAccount: { 'account-a': 'chatgpt-a' }, + choose: () => { + throw new Error('choose must not run for a pre-upgrade pin') + }, + }, + file, + ) + + expect(result).toEqual({ ...assignment, wireAccountId: 'chatgpt-a' }) + expect( + (await getSidebarState(file)).stickyAssignments?.[ + hashSidebarSessionId(sessionId) + ], + ).toEqual({ ...assignment, wireAccountId: 'chatgpt-a' }) + }) + + test('retains a pin when its current wire identity is unknown', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'oai-sb-sticky-identity-')) + const file = join(tempDir, 'sidebar-state.json') + const sessionId = 'identity-current-unknown-session' + const assignment = { + accountId: 'account-a', + wireAccountId: 'chatgpt-a', + assignedAt: now - 1, + lastSeenAt: now, + inputBytes: 128, + quotaCheckedAt: 10, + } + await setSidebarState( + make({ + stickyAssignments: { [hashSidebarSessionId(sessionId)]: assignment }, + }), + file, + ) + const before = readFileSync(file, 'utf8') + + const result = await resolveSidebarStickyAssignment( + { + sessionId, + requestBytes: 128, + now, + validPinnedAccountIds: ['account-a'], + quotaCheckedAtByAccount: { 'account-a': 10 }, + wireAccountIdByAccount: {}, + choose: () => { + throw new Error( + 'choose must not run when current identity is unknown', + ) + }, + }, + file, + ) + + expect(result).toEqual(assignment) + expect(readFileSync(file, 'utf8')).toBe(before) + }) + + test('retains a pin when its known wire identity matches without rewrite churn', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'oai-sb-sticky-identity-')) + const file = join(tempDir, 'sidebar-state.json') + const sessionId = 'identity-match-session' + const assignment = { + accountId: 'account-a', + wireAccountId: 'chatgpt-a', + assignedAt: now - 1, + lastSeenAt: now, + inputBytes: 128, + quotaCheckedAt: 10, + } + await setSidebarState( + make({ + stickyAssignments: { [hashSidebarSessionId(sessionId)]: assignment }, + }), + file, + ) + const before = readFileSync(file, 'utf8') + + const result = await resolveSidebarStickyAssignment( + { + sessionId, + requestBytes: 128, + now, + validPinnedAccountIds: ['account-a'], + quotaCheckedAtByAccount: { 'account-a': 10 }, + wireAccountIdByAccount: { 'account-a': 'chatgpt-a' }, + choose: () => { + throw new Error('choose must not run when identities match') + }, + }, + file, + ) + + expect(result).toEqual(assignment) + expect(readFileSync(file, 'utf8')).toBe(before) + }) + test('replaces excluded, expired, and disabled sticky assignments', async () => { const cases = [ {