From a2daa463f36e3f060757b2cd541a83dd2f71d551 Mon Sep 17 00:00:00 2001 From: qer Date: Fri, 17 Jul 2026 18:20:48 +0800 Subject: [PATCH 1/6] feat(vscode): reliable cancel, preserved session model, and conversation undo --- apps/vscode/shared/bridge.ts | 7 ++++ apps/vscode/src/handlers/chat.handler.ts | 25 +++++++++++++-- apps/vscode/src/runtime/kimi-runtime.ts | 13 +++----- apps/vscode/src/runtime/session-runtime.ts | 19 ++++++++++- apps/vscode/test/bridge-handler.test.ts | 31 ++++++++++++++++++ .../test/kimi-harness.integration.test.ts | 11 +++++++ apps/vscode/test/kimi-runtime.test.ts | 7 ++-- apps/vscode/test/session-runtime.test.ts | 32 +++++++++++++++++++ .../src/components/inputarea/InputArea.tsx | 19 +++++++++-- apps/vscode/webview-ui/src/services/bridge.ts | 4 +++ .../webview-ui/src/stores/chat.store.ts | 13 ++++++++ 11 files changed, 165 insertions(+), 16 deletions(-) diff --git a/apps/vscode/shared/bridge.ts b/apps/vscode/shared/bridge.ts index 805c1f32f6..c12c33af0a 100644 --- a/apps/vscode/shared/bridge.ts +++ b/apps/vscode/shared/bridge.ts @@ -37,6 +37,7 @@ export const Methods = { ResetSession: "resetSession", SetPlanMode: "setPlanMode", SteerChat: "steerChat", + UndoChat: "undoChat", RespondApproval: "respondApproval", GetKimiSessions: "getKimiSessions", @@ -183,6 +184,12 @@ function validateParams(method: RpcMethod, params: unknown): boolean { return hasBoolean(params, "enabled"); case Methods.SteerChat: return isPlainObject(params) && isContent(params["content"]); + case Methods.UndoChat: + return params === undefined || ( + isPlainObject(params) + && (params["count"] === undefined + || (Number.isInteger(params["count"]) && (params["count"] as number) >= 1)) + ); case Methods.GetProjectFiles: return params === undefined || ( isPlainObject(params) diff --git a/apps/vscode/src/handlers/chat.handler.ts b/apps/vscode/src/handlers/chat.handler.ts index 21608381e2..ca2b3a5839 100644 --- a/apps/vscode/src/handlers/chat.handler.ts +++ b/apps/vscode/src/handlers/chat.handler.ts @@ -4,8 +4,9 @@ import { isKimiError } from "@moonshot-ai/kimi-code-sdk"; import { Events, Methods } from "../../shared/bridge"; import type { ApprovalResponse, ContentPart } from "../../shared/legacy-sdk"; import { getUserMessage } from "../../shared/errors"; -import type { ErrorPhase } from "../../shared/types"; +import type { ErrorPhase, UIStreamEvent } from "../../shared/types"; import { VSCodeSettings } from "../config/vscode-settings"; +import { replaySessionToWebviewEvents } from "../runtime/replay-adapter"; import type { SessionRuntime } from "../runtime/session-runtime"; import { isWorkspacePathContained, relativeWorkspacePath } from "../utils/workspace-path"; import { parseHostSlashCommand, runHostSlashCommand } from "./slash-command"; @@ -133,7 +134,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 }; }; @@ -160,6 +164,22 @@ const steerChat: Handler<{ content: string | ContentPart[] }, { ok: boolean }> = return { ok: true }; }; +const undoChat: Handler<{ count?: number }, { events: UIStreamEvent[] }> = async (params, ctx) => { + const runtime = ctx.getSession(); + if (runtime === undefined) throw new Error("There is no conversation to undo."); + const requested = Number(params?.count); + const count = Number.isFinite(requested) && requested >= 1 ? Math.floor(requested) : 1; + await runtime.undoHistory(count); + + // Return a fresh replay so the webview can re-render the truncated + // conversation, mirroring how LoadKimiSessionHistory re-hydrates views. + const resumeState = runtime.session.getResumeState(); + if (resumeState?.agents["main"] === undefined) { + throw new Error("Session history is unavailable."); + } + return { events: replaySessionToWebviewEvents(resumeState, runtime.id) }; +}; + const resetSession: Handler = async (_, ctx) => { const runtime = ctx.getSession(); if (runtime !== undefined) injectedEditorContextSessions.delete(runtime.id); @@ -175,6 +195,7 @@ export const chatHandlers: Record> = { [Methods.RespondQuestion]: respondQuestion, [Methods.SetPlanMode]: setPlanMode, [Methods.SteerChat]: steerChat, + [Methods.UndoChat]: undoChat, [Methods.ResetSession]: resetSession, }; diff --git a/apps/vscode/src/runtime/kimi-runtime.ts b/apps/vscode/src/runtime/kimi-runtime.ts index 9da8e97543..aa98a2a78d 100644 --- a/apps/vscode/src/runtime/kimi-runtime.ts +++ b/apps/vscode/src/runtime/kimi-runtime.ts @@ -256,14 +256,11 @@ 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); diff --git a/apps/vscode/src/runtime/session-runtime.ts b/apps/vscode/src/runtime/session-runtime.ts index 45d60ac500..5ae7164a7a 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; @@ -347,6 +351,19 @@ export class SessionRuntime { if (failure !== undefined) throw failure.reason; } + /** + * Drop the last `count` user turns from the conversation, matching the + * CLI's /undo. Only meaningful while idle; the engine itself reports when + * there is nothing (or a compaction boundary) left to undo. + */ + async undoHistory(count: number): Promise { + this.ensureOpen(); + if (this.isBusy) { + throw new Error("Wait for the current response to finish before undoing."); + } + await this.session.undoHistory(count); + } + /** * Stop any in-flight work, wait for its terminal event, and keep new turns * out until the supplied operation finishes. Forking uses this so it reads a diff --git a/apps/vscode/test/bridge-handler.test.ts b/apps/vscode/test/bridge-handler.test.ts index a877efc600..88bbbdc597 100644 --- a/apps/vscode/test/bridge-handler.test.ts +++ b/apps/vscode/test/bridge-handler.test.ts @@ -130,6 +130,37 @@ 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("rejects undoChat when the view has no runtime", async () => { + const result = await bridge.handle({ id: "rpc-1", method: Methods.UndoChat }, "view-1"); + + expect(result).toEqual({ id: "rpc-1", error: "There is no conversation to undo." }); + }); + + it("rejects an invalid undoChat count at the boundary", async () => { + const result = await bridge.handle( + { id: "rpc-1", method: Methods.UndoChat, params: { count: 0 } }, + "view-1", + ); + + expect(result).toEqual({ id: "rpc-1", error: "Invalid bridge params for method: undoChat" }); + }); + 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..107d5517c7 100644 --- a/apps/vscode/test/kimi-harness.integration.test.ts +++ b/apps/vscode/test/kimi-harness.integration.test.ts @@ -1018,6 +1018,17 @@ describe("VS Code Kimi harness integration (shares one in-process SDK home)", () await expect(runtime.session.getContext()).resolves.toEqual({ history: [], tokenCount: 0 }); }); + it("undoes the last user turn on a real session", async () => { + const rig = await createRuntimeRig(); + routeSuccessfulPrompt(rig.provider); + const runtime = await openRuntimeSession(rig); + await expect(runtime.prompt("hello")).resolves.toEqual({ status: "finished" }); + + await runtime.undoHistory(1); + + await expect(runtime.session.getContext()).resolves.toMatchObject({ history: [] }); + }); + 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..650e634a39 100644 --- a/apps/vscode/test/session-runtime.test.ts +++ b/apps/vscode/test/session-runtime.test.ts @@ -45,6 +45,7 @@ interface FakeSessionBoundary { readonly subscriptionCount: () => number; readonly cancelCount: () => number; readonly cancelCompactionCount: () => number; + readonly undoCounts: readonly number[]; readonly closeCount: () => number; emit(event: Event): void; rejectNextPrompt(error: Error): void; @@ -69,6 +70,7 @@ function createFakeSession(): FakeSessionBoundary { let subscriptions = 0; let cancellations = 0; let compactionCancellations = 0; + const undoCounts: number[] = []; let closes = 0; let permission: PermissionMode = "manual"; @@ -115,6 +117,9 @@ function createFakeSession(): FakeSessionBoundary { async cancelCompaction() { compactionCancellations += 1; }, + async undoHistory(count: number) { + undoCounts.push(count); + }, async getStatus() { return { thinkingEffort: "off", @@ -152,6 +157,7 @@ function createFakeSession(): FakeSessionBoundary { subscriptionCount: () => subscriptions, cancelCount: () => cancellations, cancelCompactionCount: () => compactionCancellations, + undoCounts, closeCount: () => closes, emit(event) { for (const listener of listeners) listener(event); @@ -508,6 +514,32 @@ 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("forwards undoHistory to the SDK while the session is idle", async () => { + const { runtime, sdk } = createRuntime(); + + await runtime.undoHistory(2); + + expect([...sdk.undoCounts]).toEqual([2]); + }); + + it("rejects undoHistory while a response is still generating", async () => { + const { runtime, sdk } = createRuntime(); + void runtime.prompt("hello"); + + await expect(runtime.undoHistory(1)).rejects.toThrow( + "Wait for the current response to finish before undoing.", + ); + expect(sdk.undoCounts).toEqual([]); + }); + it("converts legacy media keys when steering an active response", async () => { const { runtime, sdk } = createRuntime(); void runtime.prompt("hello"); diff --git a/apps/vscode/webview-ui/src/components/inputarea/InputArea.tsx b/apps/vscode/webview-ui/src/components/inputarea/InputArea.tsx index 5d1b09cd64..629af2c79e 100644 --- a/apps/vscode/webview-ui/src/components/inputarea/InputArea.tsx +++ b/apps/vscode/webview-ui/src/components/inputarea/InputArea.tsx @@ -1,6 +1,6 @@ import { Fragment, useRef, useMemo, useState, useEffect, useCallback } from "react"; import { useMemoizedFn } from "ahooks"; -import { IconSend, IconPlayerStop, IconChevronDown, IconPlus } from "@tabler/icons-react"; +import { IconSend, IconPlayerStop, IconChevronDown, IconPlus, IconArrowBackUp } from "@tabler/icons-react"; import { Button } from "@/components/ui/button"; import { DropdownMenu, @@ -50,7 +50,7 @@ export function InputArea({ onAuthAction }: InputAreaProps) { const [cursorPos, setCursorPos] = useState(0); const [previewMedia, setPreviewMedia] = useState(null); - const { isStreaming, sendMessage, abort, draftMedia, removeDraftMedia, hasProcessingMedia, getMediaInConversation, pendingInput, planMode } = useChatStore(); + const { isStreaming, sendMessage, abort, draftMedia, removeDraftMedia, hasProcessingMedia, getMediaInConversation, pendingInput, planMode, undoLastTurn, sessionId, messages } = useChatStore(); const { currentModel, thinkingEffort, updateModel, toggleThinking, selectThinkingEffort, models, extensionConfig, getCurrentThinkingMode } = useSettingsStore(); const isProcessing = hasProcessingMedia(); @@ -463,6 +463,21 @@ export function InputArea({ onAuthAction }: InputAreaProps) {
+ + + + + Undo last turn + +