diff --git a/.changeset/footer-managed-quotas.md b/.changeset/footer-managed-quotas.md new file mode 100644 index 000000000..e73304df7 --- /dev/null +++ b/.changeset/footer-managed-quotas.md @@ -0,0 +1,5 @@ +--- +'@moonshot-ai/kimi-code': minor +--- + +Show managed usage quotas in the footer status bar. Quotas are polled every 30 seconds, cached for offline visibility, and updated live with current-turn token usage. Percentages are colored with a green-to-red gradient based on consumption. diff --git a/apps/kimi-code/src/tui/components/chrome/footer.ts b/apps/kimi-code/src/tui/components/chrome/footer.ts index 127009506..a9d4e6868 100644 --- a/apps/kimi-code/src/tui/components/chrome/footer.ts +++ b/apps/kimi-code/src/tui/components/chrome/footer.ts @@ -4,6 +4,7 @@ * Layout: * Line 1: [yolo] [plan] * Line 2: context: XX.X% (tokens/max) + * Line 3+: right-aligned quota rows: "label: XX% (reset in ...)" */ import type { Component } from '@earendil-works/pi-tui'; @@ -13,7 +14,7 @@ import chalk from 'chalk'; import { isRainbowDancing, renderDanceFooterModel } from '#/tui/easter-eggs/dance'; import { currentTheme } from '#/tui/theme'; import type { ColorPalette } from '#/tui/theme/colors'; -import type { AppState } from '#/tui/types'; +import type { AppState, QuotaInfo } from '#/tui/types'; import { createGitStatusCache, formatGitBadgeBase, @@ -198,12 +199,131 @@ function safeUsage(usage: number): number { return safeUsageRatio(usage); } -function formatContextStatus(usage: number, tokens?: number, maxTokens?: number): string { - const pct = `${(safeUsage(usage) * 100).toFixed(1)}%`; - if (maxTokens && maxTokens > 0 && tokens !== undefined) { - return `context: ${pct} (${formatTokenCount(tokens)}/${formatTokenCount(maxTokens)})`; +function hslToHex(h: number, s: number, l: number): string { + const normalizedH = h / 360; + const a = (s * Math.min(l, 100 - l)) / 100; + const f = (n: number): string => { + const k = (n + normalizedH * 12) % 12; + const color = l - a * Math.max(Math.min(k - 3, 9 - k, 1), -1); + return Math.round((color / 100) * 255) + .toString(16) + .padStart(2, '0'); + }; + return `#${f(0)}${f(8)}${f(4)}`; +} + +function formatResetHint(hint: string | undefined, labelName?: string): string { + if (hint === undefined) return ''; + if (hint === 'reset') return '(reset)'; + if (hint.startsWith('resets in ')) { + let parts = hint.slice('resets in '.length).split(' '); + // For weekly quotas, drop minutes when days are still present so the + // countdown column does not jitter. When less than a day remains, hours + // shift left and minutes take the hours slot. + if ((labelName ?? '').includes('week') && parts.some((p) => p.endsWith('d')) && parts.length > 2) { + parts = parts.slice(0, 2); + } + const duration = parts.join(', '); + return `(${duration})`; + } + return `(${hint})`; +} + +function shortenQuotaLabel(label: string): string { + const normalized = label.toLowerCase().replace(/\s+/g, ' ').trim(); + if (normalized.includes('week')) return 'week'; + return normalized.replace(/(\s*limit)$/, ''); +} + +interface StatusRow { + readonly labelName: string; + readonly percent: string; + readonly suffix: string; + readonly ratio: number; + readonly colored: boolean; +} + +function buildStatusRows(state: AppState): StatusRow[] { + const contextSuffix = + state.maxContextTokens && state.maxContextTokens > 0 && state.contextTokens !== undefined + ? `(${formatTokenCount(state.contextTokens)}/${formatTokenCount(state.maxContextTokens)})` + : ''; + const rows: StatusRow[] = [ + { + labelName: 'context', + percent: `${(safeUsage(state.contextUsage) * 100).toFixed(1)}%`, + suffix: contextSuffix, + ratio: state.contextUsage, + colored: false, + }, + ]; + + for (const quota of state.quotas ?? []) { + if (quota.limit <= 0) continue; + const ratio = Math.max(0, Math.min(quota.used / quota.limit, 1)); + rows.push({ + labelName: shortenQuotaLabel(quota.label), + percent: `${(ratio * 100).toFixed(1)}%`, + suffix: formatResetHint(quota.resetHint, shortenQuotaLabel(quota.label)), + ratio, + colored: true, + }); + } + + return rows; +} + +function formatStatusLines( + state: AppState, + width: number, + colors: ColorPalette, + transientHint: string | null, +): string[] { + const rows = buildStatusRows(state); + if (rows.length === 0) return []; + + const labelNameWidth = Math.max(...rows.map((r) => visibleWidth(r.labelName))); + // Reserve space for 100.0 % so the percentage column never shifts when a + // quota fills up. + const percentColWidth = Math.max(visibleWidth('100.0%'), ...rows.map((r) => visibleWidth(r.percent))); + const suffixColWidth = Math.max(...rows.map((r) => visibleWidth(r.suffix))); + const gap = 3; + const blockWidth = labelNameWidth + 1 + gap + percentColWidth + gap + suffixColWidth; + const leftPad = Math.max(0, width - blockWidth); + + const lines: string[] = []; + for (let i = 0; i < rows.length; i++) { + const row = rows[i]!; + // Smooth green -> yellow -> red gradient. Hue interpolates linearly from + // green (120) at 0 % to red (0) at 100 % with moderate saturation. + const numberColor = row.colored + ? chalk.hex(hslToHex(Math.round((1 - row.ratio) * 120), 70, 45)) + : chalk.hex(colors.text); + const content = + row.labelName.padStart(labelNameWidth) + + ':' + + ' '.repeat(gap) + + numberColor(row.percent.padStart(percentColWidth)) + + ' '.repeat(gap) + + chalk.hex(colors.text)(row.suffix.padEnd(suffixColWidth)); + + let line: string; + if (i === 0 && transientHint) { + const maxHintWidth = Math.max(0, width - blockWidth - 1); + const shownHint = + visibleWidth(transientHint) <= maxHintWidth + ? transientHint + : truncateToWidth(transientHint, maxHintWidth, '…'); + const hintWidth = visibleWidth(shownHint); + const pad = Math.max(0, leftPad - hintWidth); + line = + chalk.hex(colors.warning).bold(shownHint) + ' '.repeat(pad) + content; + } else { + line = ' '.repeat(leftPad) + content; + } + lines.push(truncateToWidth(line, width)); } - return `context: ${pct}`; + return lines; } export function formatFooterGitBadge(status: GitStatus, colors: ColorPalette): string { @@ -351,32 +471,14 @@ export class FooterComponent implements Component { line1 = truncateToWidth(leftLine, width, '…'); } - // ── Line 2: transient hint (bottom-left) + context (right) ── - const contextText = formatContextStatus( - state.contextUsage, - state.contextTokens, - state.maxContextTokens, - ); - const contextWidth = visibleWidth(contextText); - let line2: string; - if (this.transientHint) { - const maxHintWidth = Math.max(0, width - contextWidth - 1); - const shownHint = - visibleWidth(this.transientHint) <= maxHintWidth - ? this.transientHint - : truncateToWidth(this.transientHint, maxHintWidth, '…'); - const hintWidth = visibleWidth(shownHint); - const pad = Math.max(0, width - hintWidth - contextWidth); - line2 = - chalk.hex(colors.warning).bold(shownHint) + - ' '.repeat(pad) + - chalk.hex(colors.text)(contextText); - } else { - const leftPad = Math.max(0, width - contextWidth); - line2 = ' '.repeat(leftPad) + chalk.hex(colors.text)(contextText); + // ── Lines 2+: unified context + quota status block. Colons, percentages + // and suffixes share columns; the block is right-aligned. + const statusLines = formatStatusLines(state, width, colors, this.transientHint); + if (statusLines.length > 0) { + return [truncateToWidth(line1, width), ...statusLines]; } - return [truncateToWidth(line1, width), truncateToWidth(line2, width)]; + return [truncateToWidth(line1, width)]; } private syncGoalClock(goal: AppState['goal']): void { diff --git a/apps/kimi-code/src/tui/controllers/session-event-handler.ts b/apps/kimi-code/src/tui/controllers/session-event-handler.ts index 919568191..7c4c52182 100644 --- a/apps/kimi-code/src/tui/controllers/session-event-handler.ts +++ b/apps/kimi-code/src/tui/controllers/session-event-handler.ts @@ -38,6 +38,7 @@ import { type SwarmModeMarkerState, } from '../components/messages/swarm-markers'; import { + isManagedUsageProvider, OAUTH_LOGIN_REQUIRED_CODE, OAUTH_LOGIN_REQUIRED_STARTUP_NOTICE, } from '../constant/kimi-tui'; @@ -78,6 +79,7 @@ import { SubAgentEventHandler } from './subagent-event-handler'; import type { AppState, LivePaneState, + QuotaInfo, QueuedMessage, ToolCallBlockData, ToolResultBlockData, @@ -95,6 +97,7 @@ export interface SessionEventHost { requireSession(): Session; setAppState(patch: Partial): void; + fetchManagedQuotas(): Promise; patchLivePane(patch: Partial): void; resetLivePane(): void; showError(msg: string): void; @@ -142,6 +145,9 @@ export class SessionEventHandler { private queuedGoalPromotionPending = false; private queuedGoalPromotionInFlight = false; private queuedGoalPromotionTimer: ReturnType | undefined; + private quotaRefreshTimer: ReturnType | undefined; + private quotaRefreshInFlight = false; + private lastServerQuotas: readonly QuotaInfo[] | undefined; resetRuntimeState(): void { this.backgroundTasks.clear(); @@ -157,6 +163,8 @@ export class SessionEventHandler { this.queuedGoalPromotionPending = false; this.queuedGoalPromotionInFlight = false; this.clearQueuedGoalPromotionTimer(); + this.clearQuotaRefreshTimer(); + this.lastServerQuotas = undefined; this.stopAllMcpServerStatusSpinners(); } @@ -190,6 +198,7 @@ export class SessionEventHandler { this.handleEvent(event, sendQueued); }); void this.syncMcpServerStatusSnapshot(session); + this.scheduleQuotaRefresh(); } async syncMcpServerStatusSnapshot(session: Session): Promise { @@ -573,6 +582,7 @@ export class SessionEventHandler { } if (event.model !== undefined) patch.model = event.model; if (Object.keys(patch).length > 0) this.host.setAppState(patch); + if (event.model !== undefined) this.scheduleQuotaRefresh(); if (event.swarmMode === false) { this.host.state.swarmModeEntry = undefined; if (shouldRenderSwarmEnded) { @@ -687,6 +697,59 @@ export class SessionEventHandler { this.queuedGoalPromotionTimer = undefined; } + private scheduleQuotaRefresh(): void { + this.clearQuotaRefreshTimer(); + const providerKey = this.host.state.appState.availableModels[this.host.state.appState.model]?.provider; + if (!isManagedUsageProvider(providerKey)) return; + + void this.refreshQuota(); + this.quotaRefreshTimer = setInterval(() => { + void this.refreshQuota(); + }, 30_000); + this.quotaRefreshTimer.unref?.(); + } + + private clearQuotaRefreshTimer(): void { + if (this.quotaRefreshTimer === undefined) return; + clearInterval(this.quotaRefreshTimer); + this.quotaRefreshTimer = undefined; + } + + private async refreshQuota(): Promise { + if (this.quotaRefreshInFlight) return; + if (this.host.aborted) return; + const providerKey = this.host.state.appState.availableModels[this.host.state.appState.model]?.provider; + if (!isManagedUsageProvider(providerKey)) return; + + this.quotaRefreshInFlight = true; + try { + const quotas = await this.host.fetchManagedQuotas(); + if (this.host.aborted) return; + if (quotas !== undefined) { + this.lastServerQuotas = quotas; + } + const server = this.lastServerQuotas; + if (server === undefined) return; + const current = this.host.state.appState.quotas; + if ( + current !== undefined && + current.length === server.length && + server.every( + (q, i) => + q.label === current[i]?.label && + q.used === current[i]?.used && + q.limit === current[i]?.limit && + q.resetHint === current[i]?.resetHint, + ) + ) { + return; + } + this.host.setAppState({ quotas: server }); + } finally { + this.quotaRefreshInFlight = false; + } + } + requestQueuedGoalPromotion(): void { this.queuedGoalPromotionPending = true; this.goalCompletionTurnEnded = true; diff --git a/apps/kimi-code/src/tui/kimi-tui.ts b/apps/kimi-code/src/tui/kimi-tui.ts index 0337785f0..a1bb88053 100644 --- a/apps/kimi-code/src/tui/kimi-tui.ts +++ b/apps/kimi-code/src/tui/kimi-tui.ts @@ -23,6 +23,7 @@ import type { MigrationPlan } from '@moonshot-ai/migration-legacy'; import { resolve } from 'pathe'; import type { CLIOptions } from '#/cli/options'; +import { DEFAULT_OAUTH_PROVIDER_NAME } from '#/constant/app'; import { MigrationScreenComponent, type MigrationScreenResult } from '#/migration/index'; import { appendInputHistory, loadInputHistory } from '#/utils/history/input-history'; import { openUrl } from '#/utils/open-url'; @@ -113,6 +114,7 @@ import { type LivePaneState, type LoginProgressSpinnerHandle, type QueuedMessage, + type QuotaInfo, type TranscriptEntry, type TUIStartupOptions, type TUIStartupState, @@ -186,6 +188,7 @@ function createInitialAppState(input: KimiTUIStartupInput): AppState { goal: null, mcpServersSummary: null, banner: undefined, + quotas: undefined, }; } @@ -250,6 +253,39 @@ export class KimiTUI { this.harness.track(event, properties); } + async fetchManagedQuotas(): Promise { + const providerKey = this.state.appState.availableModels[this.state.appState.model]?.provider; + if (providerKey !== DEFAULT_OAUTH_PROVIDER_NAME) return undefined; + + let res; + try { + res = await this.harness.auth.getManagedUsage(providerKey); + } catch { + return undefined; + } + if (res.kind !== 'ok') return undefined; + + const rows: QuotaInfo[] = []; + if (res.summary !== null && res.summary.limit > 0) { + rows.push({ + label: res.summary.label, + used: Math.max(0, res.summary.used), + limit: res.summary.limit, + resetHint: res.summary.resetHint, + }); + } + for (const row of res.limits) { + if (row.limit <= 0) continue; + rows.push({ + label: row.label, + used: Math.max(0, row.used), + limit: row.limit, + resetHint: row.resetHint, + }); + } + return rows.length > 0 ? rows : undefined; + } + constructor(harness: KimiHarness, startupInput: KimiTUIStartupInput) { this.harness = harness; const tuiOptions: KimiTUIOptions = { diff --git a/apps/kimi-code/src/tui/types.ts b/apps/kimi-code/src/tui/types.ts index bbf047073..c92a8bf03 100644 --- a/apps/kimi-code/src/tui/types.ts +++ b/apps/kimi-code/src/tui/types.ts @@ -18,6 +18,13 @@ export interface BannerState { subText: string | null; } +export interface QuotaInfo { + readonly label: string; + readonly used: number; + readonly limit: number; + readonly resetHint?: string; +} + export interface AppState { model: string; workDir: string; @@ -46,6 +53,8 @@ export interface AppState { mcpServersSummary: string | null; /** Optional banner shown below the welcome panel; null means no banner to render. */ banner?: BannerState | null; + /** Managed quota rows (e.g. weekly limit, 5h limit). */ + quotas?: readonly QuotaInfo[]; } export interface ToolCallBlockData { diff --git a/apps/kimi-code/test/tui/components/chrome/footer.test.ts b/apps/kimi-code/test/tui/components/chrome/footer.test.ts index ab0878d6b..6bc711cc1 100644 --- a/apps/kimi-code/test/tui/components/chrome/footer.test.ts +++ b/apps/kimi-code/test/tui/components/chrome/footer.test.ts @@ -1,3 +1,4 @@ +import { visibleWidth } from '@earendil-works/pi-tui'; import chalk from 'chalk'; import { afterEach, beforeEach, describe, expect, it } from 'vitest'; @@ -16,6 +17,10 @@ function truecolorCodes(text: string): Set { return codes; } +function stripAnsi(text: string): string { + return text.replaceAll(/\u001B\[[0-9;]*m/g, ''); +} + // Dark dance colors the footer never uses outside of /dance. const RAINBOW_CYAN = '91,192,190'; const RAINBOW_GREEN = '78,200,126'; @@ -103,4 +108,84 @@ describe('FooterComponent', () => { currentTheme.setPalette(darkColors); } }); + + it('renders no quota rows when quotas are unset', () => { + const footer = new FooterComponent(appState); + expect(footer.render(80).length).toBe(2); + }); + + it('renders managed quota rows in aligned columns with reset hint', () => { + const state: AppState = { + ...appState, + contextUsage: 0.5, + contextTokens: 1_000, + maxContextTokens: 2_000, + quotas: [ + { label: 'Weekly limit', used: 41, limit: 100, resetHint: 'resets in 5d 3h' }, + { label: '5H LIMIT', used: 65, limit: 100, resetHint: 'resets in 1h 3m' }, + ], + }; + const footer = new FooterComponent(state); + const lines = footer.render(200); + + expect(lines.length).toBe(4); + const contextLine = stripAnsi(lines[1]!); + const weekLine = stripAnsi(lines[2]!); + const hourLine = stripAnsi(lines[3]!); + + expect(contextLine.trimStart().startsWith('context:')).toBe(true); + expect(weekLine.trimStart().startsWith('week:')).toBe(true); + expect(hourLine.trimStart().startsWith('5h:')).toBe(true); + expect(contextLine).toMatch(/context:\s+50\.0%\s+\(1\.0k\/2\.0k\)/); + expect(weekLine).toMatch(/week\s*:\s+41\.0%\s+\(5d, 3h\)/); + expect(hourLine).toMatch(/5h\s*:\s+65\.0%\s+\(1h, 3m\)/); + + // Colons, percentages and suffixes share the same columns. + const colonIdx = contextLine.indexOf(':'); + expect(weekLine.indexOf(':')).toBe(colonIdx); + expect(hourLine.indexOf(':')).toBe(colonIdx); + expect(weekLine.indexOf('41.0%')).toBe(hourLine.indexOf('65.0%')); + expect(weekLine.indexOf('(5d, 3h)')).toBe(hourLine.indexOf('(1h, 3m)')); + + // All rows share the same visible width (right-aligned block). + expect(stripAnsi(contextLine).length).toBe(stripAnsi(weekLine).length); + expect(stripAnsi(weekLine).length).toBe(stripAnsi(hourLine).length); + }); + + it('lowercases quota labels and colors the percentage', () => { + const state: AppState = { + ...appState, + contextUsage: 0, + contextTokens: 1_000_000, + maxContextTokens: 2_000_000, + quotas: [{ label: '5H LIMIT', used: 50, limit: 100, resetHint: 'reset' }], + }; + const footer = new FooterComponent(state); + const lines = footer.render(120); + const quotaLine = lines[2]!; + + expect(stripAnsi(quotaLine).trimStart().startsWith('5h:')).toBe(true); + expect(stripAnsi(quotaLine)).toMatch(/5h\s*:\s+50\.0%\s+\(reset\)/); + expect(truecolorCodes(quotaLine).size).toBeGreaterThan(0); + }); + + it('reserves column width for 100.0 % so the block does not shift', () => { + const state: AppState = { + ...appState, + contextUsage: 0, + quotas: [ + { label: 'week', used: 44.6, limit: 100, resetHint: 'resets in 5d 3h' }, + { label: '5h', used: 100, limit: 100, resetHint: 'resets in 2h 13m' }, + ], + }; + const footer = new FooterComponent(state); + const lines = footer.render(200); + const weekLine = stripAnsi(lines[2]!); + const fullLine = stripAnsi(lines[3]!); + + expect(weekLine).toMatch(/week\s*:\s+44\.6%\s+\(5d, 3h\)/); + expect(fullLine).toMatch(/5h\s*:\s+100\.0%\s+\(2h, 13m\)/); + // The right edge of the percentage column should align. + expect(weekLine.indexOf('%')).toBe(fullLine.indexOf('%')); + }); }); diff --git a/apps/kimi-code/test/tui/components/panels/footer-context.test.ts b/apps/kimi-code/test/tui/components/panels/footer-context.test.ts index fb1cfd5fe..08dbd8623 100644 --- a/apps/kimi-code/test/tui/components/panels/footer-context.test.ts +++ b/apps/kimi-code/test/tui/components/panels/footer-context.test.ts @@ -47,7 +47,7 @@ describe('FooterComponent — context NaN resilience', () => { const fc = new FooterComponent(baseState({ contextUsage: Number.NaN })); const out = strip(fc.render(120).join('')); expect(out).not.toMatch(/NaN/); - expect(out).toMatch(/context: 0\.0%/); + expect(out).toMatch(/context:\s+0\.0%/); }); it('undefined-ish (coerced) usage → renders 0.0%', () => { @@ -56,19 +56,19 @@ describe('FooterComponent — context NaN resilience', () => { ); const out = strip(fc.render(120).join('')); expect(out).not.toMatch(/NaN/); - expect(out).toMatch(/context: 0\.0%/); + expect(out).toMatch(/context:\s+0\.0%/); }); it('clamps ratios above 1.0 → renders 100.0%', () => { const fc = new FooterComponent(baseState({ contextUsage: 1.5 })); const out = strip(fc.render(120).join('')); - expect(out).toMatch(/context: 100\.0%/); + expect(out).toMatch(/context:\s+100\.0%/); }); it('ratio 0.427 → renders 42.7%', () => { const fc = new FooterComponent(baseState({ contextUsage: 0.427 })); const out = strip(fc.render(200).join('')); - expect(out).toMatch(/context: 42\.7%/); + expect(out).toMatch(/context:\s+42\.7%/); }); it('tokens provided but max=0 → falls back to percent-only, no division-by-zero artefact', () => { @@ -77,7 +77,7 @@ describe('FooterComponent — context NaN resilience', () => { ); const out = strip(fc.render(200).join('')); expect(out).not.toMatch(/Infinity|NaN/); - expect(out).toMatch(/context: 0\.0%/); + expect(out).toMatch(/context:\s+0\.0%/); // With maxTokens=0, token-count annotation is suppressed. expect(out).not.toMatch(/\(500\//); }); @@ -90,7 +90,7 @@ describe('FooterComponent — context NaN resilience', () => { const out = strip(footer.render(200).join('')); expect(out).toContain('kimi-k2-5'); expect(out).not.toContain(' k2 '); - expect(out).toMatch(/context: 50\.0%/); + expect(out).toMatch(/context:\s+50\.0%/); }); it('shows "thinking" label when thinking is enabled, hides it when disabled', () => { @@ -108,7 +108,7 @@ describe('FooterComponent — context NaN resilience', () => { const [, line2] = footer.render(120); expect(strip(line2 ?? '')).toContain('Press Ctrl-C again to exit'); - expect(strip(line2 ?? '')).toContain('context: 0.0%'); + expect(strip(line2 ?? '')).toMatch(/context:\s+0\.0%/); }); it('highlights the pull request badge separately from git status text', () => { diff --git a/apps/kimi-code/test/tui/controllers/session-event-handler-quota.test.ts b/apps/kimi-code/test/tui/controllers/session-event-handler-quota.test.ts new file mode 100644 index 000000000..cb7edddc8 --- /dev/null +++ b/apps/kimi-code/test/tui/controllers/session-event-handler-quota.test.ts @@ -0,0 +1,152 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +import { DEFAULT_OAUTH_PROVIDER_NAME } from '#/constant/app'; +import { SessionEventHandler } from '#/tui/controllers/session-event-handler'; +import { getBuiltInPalette } from '#/tui/theme'; +import type { AppState, QuotaInfo } from '#/tui/types'; + +function makeHost(options: { quotas?: QuotaInfo[]; fetchError?: boolean } = {}) { + const appState: Partial = { + sessionId: 's1', + streamingPhase: 'idle', + model: 'kimi-k2', + permissionMode: 'manual', + availableModels: { + 'kimi-k2': { + model: 'kimi-k2', + provider: DEFAULT_OAUTH_PROVIDER_NAME, + } as any, + openai: { + model: 'openai', + provider: 'openai', + } as any, + }, + quotas: undefined, + }; + + const host = { + state: { + appState, + queuedMessages: [], + theme: { palette: getBuiltInPalette('dark') }, + toolOutputExpanded: false, + todoPanel: { getTodos: vi.fn(() => []) }, + transcriptContainer: { addChild: vi.fn() }, + ui: { requestRender: vi.fn() }, + }, + session: undefined, + aborted: false, + sessionEventUnsubscribe: undefined, + streamingUI: { + setTurnId: vi.fn(), + flushNow: vi.fn(), + resetToolUi: vi.fn(), + finalizeTurn: vi.fn(), + hasThinkingDraft: vi.fn(() => false), + flushThinkingToTranscript: vi.fn(), + appendAssistantDelta: vi.fn(), + scheduleFlush: vi.fn(), + }, + requireSession: vi.fn(), + setAppState: vi.fn(), + fetchManagedQuotas: vi.fn(async () => + options.fetchError === true ? undefined : options.quotas, + ), + patchLivePane: vi.fn(), + resetLivePane: vi.fn(), + showError: vi.fn(), + showStatus: vi.fn(), + showNotice: vi.fn(), + track: vi.fn(), + mountEditorReplacement: vi.fn(), + restoreEditor: vi.fn(), + restoreInputText: vi.fn(), + appendTranscriptEntry: vi.fn(), + sendNormalUserInput: vi.fn(), + sendQueuedMessage: vi.fn(), + shiftQueuedMessage: vi.fn(), + btwPanelController: { routeEvent: vi.fn(() => false) }, + tasksBrowserController: {}, + }; + + host.setAppState.mockImplementation((patch: Record) => { + Object.assign(host.state.appState, patch); + }); + + return { host: host as any }; +} + +describe('SessionEventHandler quotas', () => { + beforeEach(() => { + vi.useFakeTimers(); + }); + + it('fetches managed quotas on session start', async () => { + const quotas: QuotaInfo[] = [{ label: 'Weekly limit', used: 10, limit: 100 }]; + const { host } = makeHost({ quotas }); + const handler = new SessionEventHandler(host); + + (handler as any).scheduleQuotaRefresh(); + await vi.runOnlyPendingTimersAsync(); + + expect(host.fetchManagedQuotas).toHaveBeenCalled(); + expect(host.state.appState.quotas).toEqual(quotas); + }); + + it('keeps last known quotas when the fetch fails', async () => { + const quotas: QuotaInfo[] = [{ label: 'Weekly limit', used: 10, limit: 100 }]; + const { host } = makeHost({ quotas }); + const handler = new SessionEventHandler(host); + + (handler as any).scheduleQuotaRefresh(); + await vi.runOnlyPendingTimersAsync(); + expect(host.state.appState.quotas).toEqual(quotas); + + host.fetchManagedQuotas = vi.fn(async () => undefined); + await (handler as any).refreshQuota(); + + expect(host.state.appState.quotas).toEqual(quotas); + }); + + it('does not add current-turn usage to server quotas', async () => { + const quotas: QuotaInfo[] = [{ label: 'Weekly limit', used: 10, limit: 100 }]; + const { host } = makeHost({ quotas }); + const handler = new SessionEventHandler(host); + + (handler as any).scheduleQuotaRefresh(); + await vi.runOnlyPendingTimersAsync(); + + handler.handleEvent( + { + type: 'agent.status.updated', + agentId: 'main', + usage: { + currentTurn: { inputOther: 5, inputCacheRead: 0, inputCacheCreation: 0, output: 3 }, + }, + } as any, + () => {}, + ); + + expect(host.state.appState.quotas).toEqual(quotas); + }); + + it('starts polling when the active model switches to a managed provider', async () => { + const quotas: QuotaInfo[] = [{ label: 'Weekly limit', used: 10, limit: 100 }]; + const { host } = makeHost({ quotas }); + host.state.appState.model = 'openai'; + const handler = new SessionEventHandler(host); + + (handler as any).scheduleQuotaRefresh(); + await vi.runOnlyPendingTimersAsync(); + expect(host.fetchManagedQuotas).not.toHaveBeenCalled(); + + handler.handleEvent( + { type: 'agent.status.updated', agentId: 'main', model: 'kimi-k2' } as any, + () => {}, + ); + await vi.runOnlyPendingTimersAsync(); + + expect(host.fetchManagedQuotas).toHaveBeenCalled(); + expect(host.state.appState.quotas).toEqual(quotas); + }); +});