diff --git a/hub/src/sync/usageService.test.ts b/hub/src/sync/usageService.test.ts index 83bd79dc..23cf31d6 100644 --- a/hub/src/sync/usageService.test.ts +++ b/hub/src/sync/usageService.test.ts @@ -188,14 +188,14 @@ describe('usage service', () => { store.close() }) - it('accepts Codex events that only contain last_token_usage', () => { + it('falls back to the session model for Codex events without a model', () => { const store = new Store(':memory:') const session = store.sessions.getOrCreateSession( 'codex-last-usage-test', { path: '/tmp', host: 'test', flavor: 'codex' }, null, 'default', - 'test-model' + 'deepseek-v4-flash[1m]' ) addAgentMessage(store, session.id, { @@ -214,7 +214,45 @@ describe('usage service', () => { expect(result.totals.requests).toBe(1) expect(result.totals.totalTokens).toBe(110) expect(result.totals.uncachedTokens).toBe(30) - expect(result.byModel).toEqual([expect.objectContaining({ key: 'unknown' })]) + expect(result.byModel).toEqual([ + expect.objectContaining({ key: 'deepseek-v4-flash[1m]', totalTokens: 110 }) + ]) + store.close() + }) + + it('preserves an indexed fallback model across session model changes', () => { + const store = new Store(':memory:') + const session = store.sessions.getOrCreateSession( + 'codex-model-fallback-rebuild-test', + { path: '/tmp', host: 'test', flavor: 'codex' }, + null, + 'default', + 'deepseek-v4-flash' + ) + + addAgentMessage(store, session.id, { + type: 'codex', + data: { + type: 'token_count', + thread_id: 'thread-1', + turn_id: 'turn-1', + info: { + total_token_usage: { input_tokens: 100, output_tokens: 10 }, + last_token_usage: { input_tokens: 100, output_tokens: 10 } + } + } + }) + + expect(getUsageSummary(store, 'default', 'all').byModel).toEqual([ + expect.objectContaining({ key: 'deepseek-v4-flash', totalTokens: 110 }) + ]) + + store.sessions.setSessionModel(session.id, 'gpt-5.6-sol', 'default') + store.messages.bumpMessageEpoch(session.id) + + expect(getUsageSummary(store, 'default', 'all').byModel).toEqual([ + expect.objectContaining({ key: 'deepseek-v4-flash', totalTokens: 110 }) + ]) store.close() }) diff --git a/hub/src/sync/usageService.ts b/hub/src/sync/usageService.ts index 3700753d..6979667c 100644 --- a/hub/src/sync/usageService.ts +++ b/hub/src/sync/usageService.ts @@ -31,6 +31,10 @@ function sessionAgent(session: StoredSession): string { return typeof flavor === 'string' && flavor.trim() ? flavor.trim() : 'unknown' } +function sessionModel(session: StoredSession): string | null { + return typeof session.model === 'string' && session.model.trim() ? session.model.trim() : null +} + function parseUsageEvent(session: StoredSession, message: StoredMessage): UsageEvent | null { const envelope = asRecord(message.content) if (envelope?.role !== 'agent') return null @@ -169,11 +173,32 @@ function collectUsageEvents(store: Store, sessions: StoredSession[]): void { const afterSeq = replaceEvents ? 0 : scanState.lastSeq const messages = store.messages.getMessagesAfterSeq(session.id, afterSeq) const events = new Map() + let indexedModels: Map | null = null + const getIndexedModel = (sourceKey: string): string | null => { + if (indexedModels === null) { + indexedModels = new Map( + store.usage.getEvents([session.id]) + .filter((event): event is UsageEvent & { model: string } => event.model !== null) + .map((event) => [event.sourceKey, event.model]) + ) + } + return indexedModels.get(sourceKey) ?? null + } + const fallbackModel = sessionModel(session) for (const message of messages) { const event = parseUsageEvent(session, message) if (!event) continue - if (event.kind === 'delta' || !events.has(event.sourceKey)) { + const existingEvent = events.get(event.sourceKey) + const explicitModel = event.model + event.model = explicitModel + ?? existingEvent?.model + ?? getIndexedModel(event.sourceKey) + ?? fallbackModel + if (event.kind === 'delta' || !existingEvent) { events.set(event.sourceKey, event) + } else if (explicitModel !== null) { + // A replay may add model metadata missing from the original snapshot. + existingEvent.model = explicitModel } } const lastSeq = messages.at(-1)?.seq ?? afterSeq