diff --git a/.changeset/logout-keeps-session.md b/.changeset/logout-keeps-session.md new file mode 100644 index 0000000000..5365b1868a --- /dev/null +++ b/.changeset/logout-keeps-session.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +Logging out or removing the active provider no longer closes the current session. diff --git a/apps/kimi-code/src/tui/commands/auth.ts b/apps/kimi-code/src/tui/commands/auth.ts index 0c573a4acc..3e2534dd69 100644 --- a/apps/kimi-code/src/tui/commands/auth.ts +++ b/apps/kimi-code/src/tui/commands/auth.ts @@ -240,8 +240,10 @@ export async function handleLogoutCommand(host: SlashCommandHost): Promise } if (target === currentProvider) { + // Keep the session: only the provider credential is gone. The next turn + // fails with model.not_configured until the user logs in again or picks + // another model. await host.authFlow.refreshConfigAfterLogout(); - await host.authFlow.clearActiveSessionAfterLogout(); } else { const updated = await host.harness.getConfig({ reload: true }); host.setAppState({ @@ -253,5 +255,12 @@ export async function handleLogoutCommand(host: SlashCommandHost): Promise host.track('logout', { provider: target }); const label = target === DEFAULT_OAUTH_PROVIDER_NAME ? PRODUCT_NAME : target; + if (target === currentProvider) { + host.showStatus( + `Logged out from ${label}. Current model is unavailable — /login or /model to continue.`, + 'warning', + ); + return; + } host.showStatus(`Logged out from ${label}.`); } diff --git a/apps/kimi-code/src/tui/commands/config.ts b/apps/kimi-code/src/tui/commands/config.ts index 18f7edb5d8..46068bd186 100644 --- a/apps/kimi-code/src/tui/commands/config.ts +++ b/apps/kimi-code/src/tui/commands/config.ts @@ -523,6 +523,14 @@ async function performModelSwitch( const status = await session.getStatus(); effectiveAlias = status.model ?? alias; effectiveEffort = status.thinkingEffort; + // A logout that retained the session zeroed the footer's context + // counters; the model switch is what heals the session here, so + // restore the counters from the live status too. + host.setAppState({ + contextTokens: status.contextTokens, + maxContextTokens: status.maxContextTokens, + contextUsage: status.contextUsage, + }); } } catch (error) { const msg = formatErrorMessage(error); diff --git a/apps/kimi-code/src/tui/commands/provider.ts b/apps/kimi-code/src/tui/commands/provider.ts index 5e0d84b0e6..fb2fa24ace 100644 --- a/apps/kimi-code/src/tui/commands/provider.ts +++ b/apps/kimi-code/src/tui/commands/provider.ts @@ -88,29 +88,36 @@ async function handleProviderManagerDeleteSource( } async function handleProviderDelete(host: SlashCommandHost, providerId: string): Promise { + const activeProvider = + host.state.appState.availableModels[host.state.appState.model]?.provider; + if (providerId === DEFAULT_OAUTH_PROVIDER_NAME) { await host.harness.auth.logout(DEFAULT_OAUTH_PROVIDER_NAME); // Drop the process-wide region cache with the credential: derived // endpoints (updates, marketplace, site links, telemetry) must fall back // to the marker/default profile, not the logged-out region. refreshKimiRegion(); - await host.authFlow.refreshConfigAfterLogout(); - await host.authFlow.clearActiveSessionAfterLogout(); - return; + } else { + await host.harness.removeProvider(providerId); } - const activeProvider = - host.state.appState.availableModels[host.state.appState.model]?.provider; - const config = await host.harness.removeProvider(providerId); if (activeProvider === providerId) { + // Keep the session, mirroring /logout: only the model display is cleared; + // the next turn fails with model.not_configured until the user logs in + // again or picks another model. await host.authFlow.refreshConfigAfterLogout(); - await host.authFlow.clearActiveSessionAfterLogout(); - } else { - host.setAppState({ - availableProviders: config.providers ?? {}, - availableModels: config.models ?? {}, - }); + host.showStatus( + `Provider "${providerId}" was used by the current model — /login or /model to continue.`, + 'warning', + ); + return; } + + const updated = await host.harness.getConfig({ reload: true }); + host.setAppState({ + availableProviders: updated.providers ?? {}, + availableModels: updated.models ?? {}, + }); } async function handleProviderAdd(host: SlashCommandHost): Promise { diff --git a/apps/kimi-code/src/tui/controllers/auth-flow.ts b/apps/kimi-code/src/tui/controllers/auth-flow.ts index 67fac913c2..6bab494c48 100644 --- a/apps/kimi-code/src/tui/controllers/auth-flow.ts +++ b/apps/kimi-code/src/tui/controllers/auth-flow.ts @@ -40,7 +40,6 @@ export interface AuthFlowHost { resetSessionRuntime(): void; setSession(session: Session): Promise; syncRuntimeState(session?: Session): Promise; - closeSession(reason: string): Promise; appendStartupNotice(extra: string): void; hydrateLazyConfigDefaults(): Promise; readonly sessionEventHandler: SessionEventHandler; @@ -83,6 +82,10 @@ export class AuthFlowController { if (effort !== undefined) { await host.session.setThinking(effort); } + // Logging out with the session retained zeroed the footer's context + // counters; resync from the live session even when setModel was a + // no-op (same alias), so contextTokens/contextUsage are accurate again. + await host.syncRuntimeState(host.session); return; } @@ -134,18 +137,6 @@ export class AuthFlowController { void host.refreshPluginCommands(host.session); } - async clearActiveSessionAfterLogout(): Promise { - await this.host.closeSession('logged out'); - this.host.resetSessionRuntime(); - this.host.setAppState({ - sessionId: '', - model: '', - sessionTitle: null, - }); - await this.host.refreshSkillCommands(); - await this.host.refreshPluginCommands(); - } - async refreshConfigAfterLogin(): Promise { const { host } = this; const config = await host.harness.getConfig({ reload: true }); diff --git a/apps/kimi-code/test/tui/commands/model-switch.test.ts b/apps/kimi-code/test/tui/commands/model-switch.test.ts new file mode 100644 index 0000000000..34b9f34645 --- /dev/null +++ b/apps/kimi-code/test/tui/commands/model-switch.test.ts @@ -0,0 +1,117 @@ +/** + * Scenario: /model switching on a session retained across a provider logout. + * Responsibilities: the switch must restore the context counters that logout + * zeroed (footer + cache-expiry hint read them), alongside the model itself. + * Wiring: real command and selector with the SDK/session boundary stubbed by a small host rig. + * Run: pnpm -C apps/kimi-code exec vitest run test/tui/commands/model-switch.test.ts + */ +import type { ModelAlias } from '@moonshot-ai/kimi-code-sdk'; +import { describe, expect, it, vi } from 'vitest'; + +import type { SlashCommandHost } from '#/tui/commands'; +import { handleModelCommand } from '#/tui/commands/config'; +import { TabbedModelSelectorComponent } from '#/tui/components/dialogs/tabbed-model-selector'; + +interface PickerOptions { + readonly models: Record; + readonly currentValue: string; + readonly onSelect: (selection: { alias: string; thinking: 'off' }) => void; + readonly onSessionOnlySelect: (selection: { alias: string; thinking: 'off' }) => void; +} + +function model(name: string): ModelAlias { + return { + provider: 'test', + model: name, + maxContextSize: 200_000, + displayName: name, + } as unknown as ModelAlias; +} + +function makeHost() { + const appState = { + availableModels: { + k2: model('k2'), + g1: model('g1'), + } as Record, + availableProviders: {}, + // Post-logout state: the model display and the context counters were + // cleared while the session was retained. + model: '', + thinkingEffort: 'off' as const, + contextTokens: 0, + maxContextTokens: 0, + contextUsage: 0, + streamingPhase: 'idle' as const, + transcriptEntries: [], + }; + const session = { + id: 'ses-1', + setModel: vi.fn(async () => {}), + setThinking: vi.fn(async () => {}), + getStatus: vi.fn(async () => ({ + model: 'g1', + thinkingEffort: 'off', + permission: 'manual', + planMode: false, + contextTokens: 10, + maxContextTokens: 100, + contextUsage: 0.1, + })), + }; + const host = { + state: { + appState, + transcriptEntries: [], + }, + session, + engineV2: true, + authFlow: { + refreshOAuthProviderModels: vi.fn(async () => undefined), + }, + harness: { + getConfig: vi.fn(async () => ({})), + }, + setAppState: vi.fn((patch) => Object.assign(appState, patch)), + mountEditorReplacement: vi.fn(), + restoreEditor: vi.fn(), + showStatus: vi.fn(), + showError: vi.fn(), + showNotice: vi.fn(), + track: vi.fn(), + } as unknown as SlashCommandHost & { + mountEditorReplacement: ReturnType; + showStatus: ReturnType; + showError: ReturnType; + }; + return { host, session, appState }; +} + +function mountedPicker(host: { mountEditorReplacement: ReturnType }): PickerOptions { + expect(host.mountEditorReplacement).toHaveBeenCalledOnce(); + const component = host.mountEditorReplacement.mock.calls[0]![0]; + expect(component).toBeInstanceOf(TabbedModelSelectorComponent); + return (component as unknown as { opts: PickerOptions }).opts; +} + +describe('handleModelCommand', () => { + it('restores the context counters zeroed by a provider logout', async () => { + const { host, session, appState } = makeHost(); + + await handleModelCommand(host, ''); + mountedPicker(host).onSessionOnlySelect({ alias: 'g1', thinking: 'off' }); + + await vi.waitFor(() => { + expect(host.showStatus).toHaveBeenCalled(); + }); + expect(session.setModel).toHaveBeenCalledWith('g1'); + expect(appState).toMatchObject({ + model: 'g1', + thinkingEffort: 'off', + contextTokens: 10, + maxContextTokens: 100, + contextUsage: 0.1, + }); + expect(host.showError).not.toHaveBeenCalled(); + }); +}); diff --git a/apps/kimi-code/test/tui/kimi-tui-startup.test.ts b/apps/kimi-code/test/tui/kimi-tui-startup.test.ts index b2184ee9b6..312340ed2a 100644 --- a/apps/kimi-code/test/tui/kimi-tui-startup.test.ts +++ b/apps/kimi-code/test/tui/kimi-tui-startup.test.ts @@ -1904,7 +1904,7 @@ describe('KimiTUI startup', () => { } }); - it('tracks logout after managed credentials and session state are cleared', async () => { + it('keeps the session and clears model state when logging out the current provider', async () => { const session = makeSession(); const harness = makeHarness(session, { getConfig: vi.fn(async () => ({ @@ -1926,17 +1926,21 @@ describe('KimiTUI startup', () => { await expect(driver.init()).resolves.toBe(false); harness.track.mockClear(); + const showStatus = vi.spyOn(driver as any, 'showStatus').mockImplementation(() => {}); vi.mocked(promptLogoutProviderSelection).mockResolvedValue('managed:kimi-code'); await handleLogoutCommand(driver as any); expect(harness.auth.logout).toHaveBeenCalledWith('managed:kimi-code'); - expect(session.close).toHaveBeenCalledOnce(); + expect(session.close).not.toHaveBeenCalled(); expect(driver.state.appState).toMatchObject({ - sessionId: '', + sessionId: 'ses-1', model: '', - sessionTitle: null, }); + expect(showStatus).toHaveBeenCalledWith( + 'Logged out from Kimi Code. Current model is unavailable — /login or /model to continue.', + 'warning', + ); expect(harness.track).toHaveBeenCalledWith('logout', { provider: 'managed:kimi-code' }); }); @@ -1981,6 +1985,46 @@ describe('KimiTUI startup', () => { expect(harness.track).toHaveBeenCalledWith('logout', { provider: 'openai' }); }); + it('restores the retained session counters when logging back in after logout', async () => { + const session = makeSession(); + const harness = makeHarness(session, { + getConfig: vi.fn(async () => ({ + defaultModel: 'k2', + models: { + k2: { provider: 'managed:kimi-code', model: 'moonshot-v1', maxContextSize: 100 }, + }, + providers: { 'managed:kimi-code': { type: 'kimi' } }, + })), + auth: { + status: vi.fn(async () => ({ + providers: [{ providerName: 'managed:kimi-code', hasToken: true }], + })), + login: vi.fn(async () => {}), + logout: vi.fn(), + getManagedUsage: vi.fn(), + }, + }); + const driver = makeDriver(harness, makeStartupInput()); + + await expect(driver.init()).resolves.toBe(false); + + vi.mocked(promptLogoutProviderSelection).mockResolvedValue('managed:kimi-code'); + await handleLogoutCommand(driver as any); + expect(driver.state.appState).toMatchObject({ model: '', contextTokens: 0 }); + + vi.mocked(promptPlatformSelection).mockResolvedValue('kimi-code'); + await handleLoginCommand(driver as any); + + expect(session.setModel).toHaveBeenCalledWith('k2'); + expect(driver.state.appState).toMatchObject({ + sessionId: 'ses-1', + model: 'k2', + contextTokens: 10, + maxContextTokens: 100, + contextUsage: 0.1, + }); + }); + it('can log out a stale managed entry even after the OAuth token is gone', async () => { const session = makeSession(); const harness = makeHarness(session, {