Skip to content

Commit f305deb

Browse files
feat(tui): re-evaluate quota polling when active model changes
Schedule or cancel the managed-usage refresh timer when the active model changes, so switching to/from a managed provider starts/stops polling immediately. Add unit test for the switch-to-managed case.
1 parent 66cd397 commit f305deb

2 files changed

Lines changed: 25 additions & 0 deletions

File tree

apps/kimi-code/src/tui/controllers/session-event-handler.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -587,6 +587,7 @@ export class SessionEventHandler {
587587
if (event.model !== undefined) patch.model = event.model;
588588
if (Object.keys(patch).length > 0) this.host.setAppState(patch);
589589
if (event.usage !== undefined) this.applyLiveUsage(event.usage);
590+
if (event.model !== undefined) this.scheduleQuotaRefresh();
590591
if (event.swarmMode === false) {
591592
this.host.state.swarmModeEntry = undefined;
592593
if (shouldRenderSwarmEnded) {

apps/kimi-code/test/tui/controllers/session-event-handler-quota.test.ts

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,10 @@ function makeHost(options: { quotas?: QuotaInfo[]; fetchError?: boolean } = {})
1616
model: 'kimi-k2',
1717
provider: DEFAULT_OAUTH_PROVIDER_NAME,
1818
} as any,
19+
openai: {
20+
model: 'openai',
21+
provider: 'openai',
22+
} as any,
1923
},
2024
quotas: undefined,
2125
};
@@ -142,4 +146,24 @@ describe('SessionEventHandler quotas', () => {
142146
(handler as any).applyLiveUsage({ currentTurn: undefined });
143147
expect(host.state.appState.quotas[0].used).toBe(10);
144148
});
149+
150+
it('starts polling when the active model switches to a managed provider', async () => {
151+
const quotas: QuotaInfo[] = [{ label: 'Weekly limit', used: 10, limit: 100 }];
152+
const { host } = makeHost({ quotas });
153+
host.state.appState.model = 'openai';
154+
const handler = new SessionEventHandler(host);
155+
156+
(handler as any).scheduleQuotaRefresh();
157+
await vi.runOnlyPendingTimersAsync();
158+
expect(host.fetchManagedQuotas).not.toHaveBeenCalled();
159+
160+
handler.handleEvent(
161+
{ type: 'agent.status.updated', agentId: 'main', model: 'kimi-k2' } as any,
162+
() => {},
163+
);
164+
await vi.runOnlyPendingTimersAsync();
165+
166+
expect(host.fetchManagedQuotas).toHaveBeenCalled();
167+
expect(host.state.appState.quotas).toEqual(quotas);
168+
});
145169
});

0 commit comments

Comments
 (0)