From 5381933eea4ed2b691e4258bbbaeba52ce7551c9 Mon Sep 17 00:00:00 2001 From: qer Date: Fri, 17 Jul 2026 19:50:18 +0800 Subject: [PATCH] fix(vscode): reliable cancel and preserved session model on attach --- apps/vscode/src/handlers/chat.handler.ts | 20 +++++- apps/vscode/src/runtime/kimi-runtime.ts | 15 ++-- apps/vscode/src/runtime/session-runtime.ts | 6 +- apps/vscode/test/bridge-handler.test.ts | 16 +++++ .../test/kimi-harness.integration.test.ts | 69 +++++++++++++++++-- apps/vscode/test/kimi-runtime.test.ts | 7 +- apps/vscode/test/session-runtime.test.ts | 8 +++ 7 files changed, 122 insertions(+), 19 deletions(-) diff --git a/apps/vscode/src/handlers/chat.handler.ts b/apps/vscode/src/handlers/chat.handler.ts index 21608381e2..c5b21f89ad 100644 --- a/apps/vscode/src/handlers/chat.handler.ts +++ b/apps/vscode/src/handlers/chat.handler.ts @@ -6,6 +6,7 @@ import type { ApprovalResponse, ContentPart } from "../../shared/legacy-sdk"; import { getUserMessage } from "../../shared/errors"; import type { ErrorPhase } from "../../shared/types"; import { VSCodeSettings } from "../config/vscode-settings"; +import { normalizeEffort } from "../runtime/kimi-runtime"; import type { SessionRuntime } from "../runtime/session-runtime"; import { isWorkspacePathContained, relativeWorkspacePath } from "../utils/workspace-path"; import { parseHostSlashCommand, runHostSlashCommand } from "./slash-command"; @@ -101,11 +102,23 @@ const streamChat: Handler = async (params, } try { + // Attach no longer overwrites session modes with the configured defaults + // (resumed sessions keep their own), so apply the model/effort that the + // composer submitted with this prompt before the turn starts. const status = await runtime.session.getStatus(); + let model = status.model; + if (params.model && model !== params.model) { + await runtime.session.setModel(params.model); + model = params.model; + } + const effort = normalizeEffort(params.effort ?? (params.thinking === true ? "on" : "off")); + if (status.thinkingEffort !== effort) { + await runtime.session.setThinking(effort); + } if (params.planMode !== undefined && status.planMode !== params.planMode) { await runtime.session.setPlanMode(params.planMode); } - runtime.announceSessionStart(status.model); + runtime.announceSessionStart(model); } catch (error) { emitCaughtError(ctx, error, "preflight", runtime.id); return { done: false }; @@ -133,7 +146,10 @@ const streamChat: Handler = async (params, const abortChat: Handler = async (_, ctx) => { const runtime = ctx.getSession(); - if (runtime !== undefined) await runtime.cancel(); + // Do not claim an abort when there is no runtime to cancel — the webview + // would otherwise show the task as stopped while the engine keeps running. + if (runtime === undefined) return { aborted: false }; + await runtime.cancel(); return { aborted: true }; }; diff --git a/apps/vscode/src/runtime/kimi-runtime.ts b/apps/vscode/src/runtime/kimi-runtime.ts index 9da8e97543..9c5a9e5444 100644 --- a/apps/vscode/src/runtime/kimi-runtime.ts +++ b/apps/vscode/src/runtime/kimi-runtime.ts @@ -256,21 +256,18 @@ async function applySessionSettings( legacyApproval: LegacyApprovalFlags, ): Promise { const status = await session.getStatus(); - if (options.model && status.model !== options.model) { - await session.setModel(options.model); - } - // Thinking effort is applied only when the session is created (see - // openSession). An existing session keeps its own effort — the global - // config value is a default for new sessions, matching CLI/TUI resume - // semantics. Effort changes made in the picker reach the active session - // through the SaveConfig handler instead. + // Model and thinking effort are applied only when the session is created + // (see openSession). An existing session keeps its own — the global config + // values are defaults for new sessions, matching CLI/TUI resume semantics. + // Changes made in the pickers reach the active session through the + // SaveConfig handler instead. const permission = corePermissionForLegacyApproval(legacyApproval); if (status.permission !== permission) { await session.setPermission(permission); } } -function normalizeEffort(effort: string): ThinkingEffort { +export function normalizeEffort(effort: string): ThinkingEffort { return (effort.trim() || "off") as ThinkingEffort; } diff --git a/apps/vscode/src/runtime/session-runtime.ts b/apps/vscode/src/runtime/session-runtime.ts index 45d60ac500..22a6536bba 100644 --- a/apps/vscode/src/runtime/session-runtime.ts +++ b/apps/vscode/src/runtime/session-runtime.ts @@ -327,7 +327,11 @@ export class SessionRuntime { } async cancel(): Promise { - if (this.closed || !this.hasActiveWork) return; + // Always reach the engine, even when the host believes nothing is active. + // The host-side bookkeeping can drift from engine truth after an abnormal + // error path; session.cancel() is a harmless no-op when the engine is + // idle, but it is the only way to recover a turn the host lost track of. + if (this.closed) return; this.reverseRpc.cancelAll("Turn cancelled"); const cancellingHostAction = this.hostActionActive; const hostActionId = this.activeHostActionId; diff --git a/apps/vscode/test/bridge-handler.test.ts b/apps/vscode/test/bridge-handler.test.ts index a877efc600..add81ea091 100644 --- a/apps/vscode/test/bridge-handler.test.ts +++ b/apps/vscode/test/bridge-handler.test.ts @@ -130,6 +130,22 @@ describe("Webview RPC boundary (validates requests before host dispatch)", () => expect(showLogs).not.toHaveBeenCalled(); }); + it("reports aborted: false when the view has no runtime to cancel", async () => { + const result = await bridge.handle({ id: "rpc-1", method: Methods.AbortChat }, "view-1"); + + expect(result).toEqual({ id: "rpc-1", result: { aborted: false } }); + }); + + it("cancels the view's runtime when aborting a chat", async () => { + const cancel = vi.fn(async () => undefined); + vi.spyOn(bridge.runtime, "getSessionForView").mockReturnValue({ cancel } as never); + + const result = await bridge.handle({ id: "rpc-1", method: Methods.AbortChat }, "view-1"); + + expect(result).toEqual({ id: "rpc-1", result: { aborted: true } }); + expect(cancel).toHaveBeenCalledOnce(); + }); + it.each(["missingMethod", "toString", "constructor", "__proto__"])( "does not dispatch the unknown or prototype method %s", async (method) => { diff --git a/apps/vscode/test/kimi-harness.integration.test.ts b/apps/vscode/test/kimi-harness.integration.test.ts index 36c9505735..d5f1e725c6 100644 --- a/apps/vscode/test/kimi-harness.integration.test.ts +++ b/apps/vscode/test/kimi-harness.integration.test.ts @@ -22,6 +22,9 @@ vi.mock("vscode", () => ({ showWarningMessage: async () => undefined, showTextDocument: async () => undefined, }, + workspace: { + getConfiguration: () => ({ get: (_key: string, fallback: unknown) => fallback }), + }, })); import { @@ -35,6 +38,7 @@ import { type UpdateMCPServerRequest, } from "../shared/legacy-sdk"; import { configHandlers } from "../src/handlers/config.handler"; +import { chatHandlers } from "../src/handlers/chat.handler"; import { mcpHandlers } from "../src/handlers/mcp.handler"; import { parseHostSlashCommand, runHostSlashCommand } from "../src/handlers/slash-command"; import type { HandlerContext } from "../src/handlers/types"; @@ -80,7 +84,7 @@ afterEach(async () => { } }); -async function createRuntimeRig(): Promise { +async function createRuntimeRig(extraAliases: readonly string[] = []): Promise { const rootDir = await mkdtemp(join(tmpdir(), "kimi-vscode-harness-")); const homeDir = join(rootDir, "home"); const workDir = join(rootDir, "workspace"); @@ -94,7 +98,7 @@ async function createRuntimeRig(): Promise { await provider.close(); }; - await writeProviderConfig(homeDir, `${provider.baseUrl}/v1`); + await writeProviderConfig(homeDir, `${provider.baseUrl}/v1`, extraAliases); const version = await readExtensionVersion(); const broadcasts: BroadcastRecord[] = []; const logs: LogRecord[] = []; @@ -184,7 +188,23 @@ async function readExtensionVersion(): Promise { return parsed.version; } -async function writeProviderConfig(homeDir: string, baseUrl: string): Promise { +async function writeProviderConfig( + homeDir: string, + baseUrl: string, + extraAliases: readonly string[] = [], +): Promise { + const extra = extraAliases + .map( + (alias) => ` +[models."${alias}"] +provider = "local" +model = "mock-model" +max_context_size = 128000 +capabilities = ["thinking"] +support_efforts = ["low", "high"] +`, + ) + .join("\n"); await writeFile( join(homeDir, "config.toml"), `default_model = "${MODEL_ALIAS}" @@ -198,7 +218,7 @@ api_key = "${PROVIDER_TOKEN}" provider = "local" model = "mock-model" max_context_size = 128000 - +${extra} [loop_control] max_retries_per_step = 1 `, @@ -270,6 +290,30 @@ async function openRuntimeSession(rig: RuntimeRig, sessionId?: string, yoloMode }); } +function streamChatContext(rig: RuntimeRig): HandlerContext { + return { + workDir: rig.workDir, + webviewId: "view-1", + broadcast: (event: string, data: unknown, webviewId?: string) => { + rig.broadcasts.push({ event, data, webviewId }); + }, + getOrCreateSession: async (model: string, effort: string, sessionId?: string) => + rig.runtime.openSession({ + webviewId: "view-1", + workDir: rig.workDir, + ...(sessionId === undefined ? {} : { sessionId }), + model, + effort, + yoloMode: false, + }), + getSession: () => rig.runtime.getSessionForView("view-1"), + saveAllDirty: async () => undefined, + logError: (message: string, error: unknown) => { + rig.logs.push({ message, error }); + }, + } as unknown as HandlerContext; +} + function streamEvents(broadcasts: readonly BroadcastRecord[]): unknown[] { return broadcasts .filter((record) => record.event === Events.StreamEvent) @@ -1018,6 +1062,23 @@ describe("VS Code Kimi harness integration (shares one in-process SDK home)", () await expect(runtime.session.getContext()).resolves.toEqual({ history: [], tokenCount: 0 }); }); + it("applies the composer-submitted model before the turn starts", async () => { + const rig = await createRuntimeRig(["vscode-alt"]); + routeSuccessfulPrompt(rig.provider); + const runtime = await openRuntimeSession(rig); + + const result = await chatHandlers[Methods.StreamChat]!( + { content: "hi", model: "vscode-alt", effort: "high" }, + streamChatContext(rig), + ); + + expect(result).toEqual({ done: true }); + await expect(runtime.session.getStatus()).resolves.toMatchObject({ + model: "vscode-alt", + thinkingEffort: "high", + }); + }); + it("toggles plan mode through the public session without calling the model", async () => { const rig = await createRuntimeRig(); const runtime = await openRuntimeSession(rig); diff --git a/apps/vscode/test/kimi-runtime.test.ts b/apps/vscode/test/kimi-runtime.test.ts index 8133c2916e..6a86f7f2d1 100644 --- a/apps/vscode/test/kimi-runtime.test.ts +++ b/apps/vscode/test/kimi-runtime.test.ts @@ -343,13 +343,14 @@ describe("Kimi runtime (owns shared SDK sessions for Webviews)", () => { expect(boundary.handlerInstallations).toEqual({ approval: 1, question: 1 }); }); - it("updates the SDK model when the resumed session has a different model", async () => { + it("preserves the resumed session's model instead of reapplying the configured default", async () => { const { runtime, sdk } = createRuntime(); const session = sdk.addSession("saved-1", "/workspace", { model: "old-model" }); - await runtime.openSession(openOptions({ sessionId: "saved-1", model: "new-model" })); + const opened = await runtime.openSession(openOptions({ sessionId: "saved-1", model: "new-model" })); - expect(session.setModels).toEqual(["new-model"]); + expect(session.setModels).toEqual([]); + await expect(opened.session.getStatus()).resolves.toMatchObject({ model: "old-model" }); }); it("preserves the resumed session's thinking effort instead of reapplying the configured default", async () => { diff --git a/apps/vscode/test/session-runtime.test.ts b/apps/vscode/test/session-runtime.test.ts index 6caadfabb1..b392204fde 100644 --- a/apps/vscode/test/session-runtime.test.ts +++ b/apps/vscode/test/session-runtime.test.ts @@ -508,6 +508,14 @@ describe("session runtime (adapts one SDK session for subscribed Webviews)", () expect(sdk.cancelCount()).toBe(1); }); + it("still reaches the SDK cancel when the host lost track of active work", async () => { + const { runtime, sdk } = createRuntime(); + + await runtime.cancel(); + + expect(sdk.cancelCount()).toBe(1); + }); + it("converts legacy media keys when steering an active response", async () => { const { runtime, sdk } = createRuntime(); void runtime.prompt("hello");