/** * Scenario: untrusted Webview RPC messages cross into the VS Code extension host. * Responsibilities: validate requests, preserve public model metadata, omit private paths, and recover visibly from persisted state errors. * Wiring: the real BridgeHandler and handlers; VS Code and the public Node SDK harness boundary are replaced. * Run: pnpm --filter kimi-code exec vitest run --config vitest.config.ts test/bridge-handler.test.ts */ import { mkdtemp, readdir, rm, writeFile } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { afterEach, beforeEach, describe, expect, it, vi, type Mock } from "vitest"; import type * as vscode from "vscode"; import { Methods } from "../shared/bridge"; import { BridgeHandler } from "../src/bridge-handler"; const host = vi.hoisted(() => { const watcher = { onDidChange: vi.fn(), onDidCreate: vi.fn(), onDidDelete: vi.fn(), dispose: vi.fn(), }; const harness = { homeDir: "/tmp/kimi-code-test-home", close: vi.fn(async () => undefined), getConfig: vi.fn(), setConfig: vi.fn(async () => undefined), listSessions: vi.fn(async () => []), resumeSession: vi.fn(), forkSession: vi.fn(), deleteSession: vi.fn(async () => undefined), }; const showWarningMessage = vi.fn(async () => undefined as string | undefined); class Uri { readonly scheme = "file"; readonly authority = ""; readonly path: string; constructor(readonly fsPath: string) { this.path = fsPath; } static joinPath(base: Uri, ...segments: string[]): Uri { return new Uri(join(base.fsPath, ...segments)); } toString(): string { return `file://${this.path}`; } } return { Uri, watcher, harness, showWarningMessage, workspaceFolders: [] as Array<{ uri: Uri }>, }; }); vi.mock("vscode", () => ({ Uri: host.Uri, workspace: { get workspaceFolders() { return host.workspaceFolders; }, getConfiguration: () => ({ get: (_key: string, fallback: unknown) => fallback }), createFileSystemWatcher: () => host.watcher, textDocuments: [], }, window: { showWarningMessage: host.showWarningMessage }, })); vi.mock("@moonshot-ai/kimi-code-sdk", async (importOriginal) => { const original = await importOriginal(); return { ...original, createKimiHarness: () => host.harness }; }); let bridge: BridgeHandler; let root: string; let showLogs: Mock<() => void>; let writeLog: Mock<(message: string) => void>; let workspaceState: { get: ReturnType; update: ReturnType }; beforeEach(async () => { root = await mkdtemp(join(tmpdir(), "kimi-vscode-bridge-")); host.workspaceFolders.splice(0, host.workspaceFolders.length, { uri: new host.Uri(root) }); showLogs = vi.fn(); writeLog = vi.fn(); host.harness.resumeSession.mockReset(); host.harness.getConfig.mockReset(); host.harness.getConfig.mockResolvedValue({ models: {} }); host.showWarningMessage.mockReset(); host.showWarningMessage.mockResolvedValue(undefined); workspaceState = { get: vi.fn((_key, fallback) => fallback), update: vi.fn() }; bridge = new BridgeHandler( vi.fn(), workspaceState as unknown as vscode.Memento, join(root, "global-storage"), vi.fn(), showLogs, writeLog, ); }); afterEach(async () => { await bridge.dispose(); vi.clearAllMocks(); await rm(root, { recursive: true, force: true }); }); describe("Webview RPC boundary (validates requests before host dispatch)", () => { it("returns a readable error when the envelope is not a plain object", async () => { const result = await bridge.handle([], "view-1"); expect(result).toEqual({ id: "", error: "Invalid bridge request: expected a plain object.", }); }); it("does not execute a known handler when the request id is blank", async () => { const result = await bridge.handle({ id: " ", method: Methods.ShowLogs }, "view-1"); expect(result).toEqual({ id: "", error: "Invalid bridge request: id must be a non-empty string.", }); 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) => { const result = await bridge.handle({ id: "rpc-1", method }, "view-1"); expect(result).toEqual({ id: "rpc-1", error: `Unknown bridge method: ${method}` }); expect(showLogs).not.toHaveBeenCalled(); }, ); it("does not execute a no-params handler when a payload is supplied", async () => { const result = await bridge.handle( { id: "rpc-1", method: Methods.ShowLogs, params: {} }, "view-1", ); expect(result).toEqual({ id: "rpc-1", error: "Invalid bridge params for method: showLogs", }); expect(showLogs).not.toHaveBeenCalled(); }); it("does not execute an object-payload handler when a required field has the wrong type", async () => { const result = await bridge.handle( { id: "rpc-1", method: Methods.AddInputHistory, params: { text: 42 } }, "view-1", ); expect(result).toEqual({ id: "rpc-1", error: "Invalid bridge params for method: addInputHistory", }); expect(workspaceState.update).not.toHaveBeenCalled(); }); it("dispatches a valid request through the existing bridge surface", async () => { const result = await bridge.handle({ id: "rpc-1", method: Methods.ShowLogs }, "view-1"); expect(result).toEqual({ id: "rpc-1", result: { ok: true } }); expect(showLogs).toHaveBeenCalledOnce(); }); it("keeps provider identity when configured models share a display name", async () => { host.harness.getConfig.mockResolvedValueOnce({ defaultModel: "openai/shared", models: { "openai/shared": { provider: "openai", model: "shared", displayName: "Shared", maxContextSize: 128_000, }, "proxy/shared": { provider: "company-proxy", model: "shared", displayName: "Shared", maxContextSize: 128_000, }, }, }); const result = await bridge.handle({ id: "rpc-models", method: Methods.GetModels }, "view-1"); expect(result).toMatchObject({ id: "rpc-models", result: { defaultModel: "openai/shared", models: [ { id: "openai/shared", name: "Shared", provider: "openai" }, { id: "proxy/shared", name: "Shared", provider: "company-proxy" }, ], }, }); }); it("preserves adaptive thinking metadata in the Webview model list", async () => { host.harness.getConfig.mockResolvedValueOnce({ defaultModel: "anthropic/claude", models: { "anthropic/claude": { provider: "anthropic", model: "claude-sonnet", maxContextSize: 200_000, adaptiveThinking: true, }, }, }); const result = await bridge.handle({ id: "rpc-models", method: Methods.GetModels }, "view-1"); expect(result).toMatchObject({ result: { models: [{ id: "anthropic/claude", name: "claude-sonnet", provider: "anthropic", adaptive_thinking: true, }], }, }); }); it("does not expose the session storage path when listing sessions", async () => { host.harness.listSessions.mockResolvedValueOnce([ { id: "session-1", workDir: root, sessionDir: "/private/kimi/sessions/session-1", updatedAt: 123, title: "Visible title", }, ] as never); const result = await bridge.handle( { id: "rpc-1", method: Methods.GetKimiSessions }, "view-1", ); expect(result).toEqual({ id: "rpc-1", result: [{ id: "session-1", workDir: root, updatedAt: 123, brief: "Visible title" }], }); expect(JSON.stringify(result)).not.toContain("/private/kimi/sessions"); }); it("does not expose the session storage path when forking a session", async () => { const source = { id: "session-1", workDir: root, sessionDir: "/private/kimi/sessions/session-1", updatedAt: 123, }; const target = { id: "session-2", workDir: root, sessionDir: "/private/kimi/sessions/session-2", updatedAt: 124, }; host.harness.listSessions.mockResolvedValueOnce([source] as never); host.harness.forkSession.mockResolvedValueOnce({ summary: target, close: vi.fn() }); const result = await bridge.handle( { id: "rpc-1", method: Methods.ForkKimiSession, params: { sessionId: "session-1", turnIndex: 0 }, }, "view-1", ); expect(result).toEqual({ id: "rpc-1", result: { sessionId: "session-2" } }); expect(JSON.stringify(result)).not.toContain("/private/kimi/sessions"); }); it("runs a fork through the active session cancellation boundary", async () => { const source = { id: "session-1", workDir: root, sessionDir: "/private/kimi/sessions/session-1", updatedAt: 123, }; const target = { id: "session-2", workDir: root, sessionDir: "/private/kimi/sessions/session-2", updatedAt: 124, }; const runExclusiveAfterCancelling = vi.fn(async (action: () => Promise) => action()); vi.spyOn(bridge.runtime, "getSession").mockReturnValue({ runExclusiveAfterCancelling, } as never); host.harness.listSessions.mockResolvedValueOnce([source] as never); host.harness.forkSession.mockResolvedValueOnce({ summary: target, close: vi.fn() }); const result = await bridge.handle( { id: "rpc-1", method: Methods.ForkKimiSession, params: { sessionId: "session-1", turnIndex: 0 }, }, "view-1", ); expect(result).toEqual({ id: "rpc-1", result: { sessionId: "session-2" } }); expect(runExclusiveAfterCancelling).toHaveBeenCalledOnce(); expect(host.harness.forkSession).toHaveBeenCalledOnce(); }); it("closes and removes a fork when its baseline cannot be materialized", async () => { const source = { id: "session-1", workDir: root, updatedAt: 123 }; const target = { id: "session-2", workDir: root, updatedAt: 124 }; const close = vi.fn(async () => undefined); host.harness.listSessions.mockResolvedValueOnce([source] as never); host.harness.forkSession.mockResolvedValueOnce({ summary: target, close }); vi.spyOn(bridge.baselineManager, "materializeToFork").mockRejectedValueOnce( new Error("baseline unavailable"), ); const deleteBaseline = vi.spyOn(bridge.baselineManager, "deleteSession"); const result = await bridge.handle( { id: "rpc-1", method: Methods.ForkKimiSession, params: { sessionId: "session-1", turnIndex: 0 }, }, "view-1", ); expect(result).toEqual({ id: "rpc-1", error: "baseline unavailable" }); expect(close).toHaveBeenCalledOnce(); expect(host.harness.deleteSession).toHaveBeenCalledWith("session-2"); expect(deleteBaseline).toHaveBeenCalledWith("session-2"); }); it("keeps conversation history available when its baseline snapshot disappears", async () => { const session = createResumedSession("session-1", root); host.harness.resumeSession.mockResolvedValueOnce(session as never); host.showWarningMessage.mockResolvedValueOnce("Show Logs"); const sourcePath = join(root, "app.ts"); await writeFile(sourcePath, "original\n", "utf-8"); await bridge.baselineManager.capture(session.summary, sourcePath); const baselinesRoot = join(root, "global-storage", "baselines"); const [homeDirectory] = await readdir(baselinesRoot); const [sessionDirectory] = await readdir(join(baselinesRoot, homeDirectory!)); const snapshotsDirectory = join( baselinesRoot, homeDirectory!, sessionDirectory!, "snapshots", ); const [snapshot] = await readdir(snapshotsDirectory); await rm(join(snapshotsDirectory, snapshot!)); const result = await bridge.handle( { id: "rpc-1", method: Methods.LoadKimiSessionHistory, params: { kimiSessionId: "session-1" }, }, "view-1", ); expect(result).toEqual({ id: "rpc-1", result: expect.arrayContaining([ expect.objectContaining({ type: "StatusUpdate", _sessionId: "session-1" }), ]), }); expect(writeLog).toHaveBeenCalledWith( expect.stringMatching(/Unable to restore session file changes.*Unable to read baseline snapshot/), ); await vi.waitFor(() => expect(showLogs).toHaveBeenCalledOnce()); }); it("returns a readable error when persisted session state is corrupt without wedging the bridge", async () => { host.harness.resumeSession.mockRejectedValueOnce( new Error("Session state is invalid JSON at line 4"), ); const failed = await bridge.handle( { id: "rpc-1", method: Methods.LoadKimiSessionHistory, params: { kimiSessionId: "session-1" }, }, "view-1", ); const next = await bridge.handle({ id: "rpc-2", method: Methods.ShowLogs }, "view-1"); expect(failed).toEqual({ id: "rpc-1", error: "Session state is invalid JSON at line 4", }); expect(writeLog).toHaveBeenCalledWith( expect.stringContaining("Session state is invalid JSON at line 4"), ); expect(next).toEqual({ id: "rpc-2", result: { ok: true } }); }); }); describe("Webview config saves (thinking effort persistence parity with the TUI)", () => { const effortModel = { provider: "managed:kimi-code", model: "reasoning", supportEfforts: ["low", "high", "max"], defaultEffort: "high", }; function mockConfig(thinking?: { enabled: boolean; effort?: string }) { host.harness.getConfig.mockResolvedValue({ defaultModel: "kimi/reasoning", thinking, models: { "kimi/reasoning": effortModel }, } as never); } it("persists a non-top effort as the global default", async () => { mockConfig(); const result = await bridge.handle( { id: "rpc-1", method: Methods.SaveConfig, params: { model: "kimi/reasoning", thinking: true, effort: "high" } }, "view-1", ); expect(result).toEqual({ id: "rpc-1", result: { ok: true } }); expect(host.harness.setConfig).toHaveBeenCalledWith({ defaultModel: "kimi/reasoning", thinking: { enabled: true, effort: "high" }, }); }); it("keeps the model's top declared tier session-only", async () => { mockConfig(); await bridge.handle( { id: "rpc-1", method: Methods.SaveConfig, params: { model: "kimi/reasoning", thinking: true, effort: "max" } }, "view-1", ); expect(host.harness.setConfig).toHaveBeenCalledWith({ defaultModel: "kimi/reasoning", thinking: { enabled: true }, }); }); it("persists the concrete effort when the model's levels are unknown", async () => { host.harness.getConfig.mockResolvedValue({ defaultModel: "other/model", models: {} }); await bridge.handle( { id: "rpc-1", method: Methods.SaveConfig, params: { model: "custom/model", thinking: true, effort: "max" } }, "view-1", ); expect(host.harness.setConfig).toHaveBeenCalledWith({ defaultModel: "custom/model", thinking: { enabled: true, effort: "max" }, }); }); it("leaves the stored effort alone when the pick re-confirms the active effort", async () => { mockConfig({ enabled: false, effort: "low" }); await bridge.handle( { id: "rpc-1", method: Methods.SaveConfig, params: { model: "kimi/reasoning", thinking: true, effort: "high", effortChanged: false } }, "view-1", ); expect(host.harness.setConfig).toHaveBeenCalledWith({ defaultModel: "kimi/reasoning", thinking: { enabled: true }, }); }); it("skips the config write entirely when nothing changed", async () => { mockConfig({ enabled: true, effort: "high" }); await bridge.handle( { id: "rpc-1", method: Methods.SaveConfig, params: { model: "kimi/reasoning", thinking: true, effort: "high" } }, "view-1", ); expect(host.harness.setConfig).not.toHaveBeenCalled(); }); }); function createResumedSession(id: string, workDir: string) { const close = vi.fn(async () => undefined); const summary = { id, workDir, sessionDir: join("/private/kimi/sessions", id), createdAt: 1, updatedAt: 2, metadata: { vscode_legacy_approval: { yolo: false, afk: false } }, }; return { id, workDir, summary, close, getResumeState: () => ({ sessionMetadata: { agents: {} }, agents: { main: { type: "main", config: { cwd: workDir, modelAlias: "test-model", modelCapabilities: { image_in: false, video_in: false, audio_in: false, thinking: false, tool_use: true, max_context_tokens: 128_000, }, thinkingEffort: "off", systemPrompt: "", }, context: { history: [], tokenCount: 0 }, replay: [], permission: { mode: "manual", rules: [] }, plan: null, usage: {}, tools: [], background: [], }, }, }), getStatus: async () => ({ permission: "manual" }), setPermission: async () => undefined, updateMetadata: async () => undefined, setApprovalHandler: () => undefined, setQuestionHandler: () => undefined, onEvent: () => () => undefined, }; }