import { isKimiError, type ContentPart as SdkContentPart, type Event, type PromptInput, type Session, type SessionSummary, } from "@moonshot-ai/kimi-code-sdk"; import type { ContentPart as LegacyContentPart, ApprovalResponse } from "../../shared/legacy-sdk"; import { Events } from "../../shared/bridge"; import { getUserMessage } from "../../shared/errors"; import type { ErrorPhase, UIStreamEvent } from "../../shared/types"; import { adaptSdkEvent, createEventAdapterState, type EventAdapterState, type TurnTerminalMetadata, } from "./event-adapter"; import { corePermissionForLegacyApproval, legacyApprovalMetadata, type LegacyApprovalFlags, } from "./legacy-approval"; import { ReverseRpcController } from "./reverse-rpc"; export type RuntimeBroadcast = (event: string, data: unknown, webviewId?: string) => void; export interface SessionRuntimeOptions { readonly session: Session; readonly legacyApproval: LegacyApprovalFlags; readonly broadcast: RuntimeBroadcast; readonly captureBaseline: ( session: Pick, filePath: string, webviewIds: readonly string[], ) => void; readonly log: (message: string, error?: unknown) => void; } interface ActivePrompt { readonly input: LegacyContentPart[] | string; started: boolean; settled: boolean; resolve: (result: PromptResult) => void; } const ALREADY_GENERATING_MESSAGE = "A response is already being generated for this session."; export interface PromptResult { readonly status: "finished" | "cancelled" | "failed"; } interface SuppressedError { readonly code: string; readonly message: string; } interface PendingHostCompaction { readonly actionId: number; readonly resolve: (result: "completed" | "cancelled") => void; readonly reject: (error: unknown) => void; } /** * Owns the one SDK event subscription and reverse-RPC handlers for a session. * Any number of Webviews may subscribe without replacing each other's approval * handler or duplicating streamed events. */ export class SessionRuntime { readonly session: Session; private readonly broadcast: RuntimeBroadcast; private readonly captureBaseline: SessionRuntimeOptions["captureBaseline"]; private readonly log: SessionRuntimeOptions["log"]; private readonly webviewIds = new Set(); private readonly reverseRpc: ReverseRpcController; private readonly unsubscribe: () => void; private adapterState: EventAdapterState = createEventAdapterState(); private activePrompt: ActivePrompt | undefined; private hostActionActive = false; private hostActionSequence = 0; private activeHostActionId: number | undefined; private readonly cancelledHostActions = new Set(); private pendingHostCompaction: PendingHostCompaction | undefined; private readonly activeWorkSettledWaiters = new Set<() => void>(); private exclusiveActionActive = false; private readonly terminalKeys = new Set(); private suppressedError: SuppressedError | undefined; private legacyApproval: LegacyApprovalFlags; private closed = false; constructor(options: SessionRuntimeOptions) { this.session = options.session; this.broadcast = options.broadcast; this.captureBaseline = options.captureBaseline; this.log = options.log; this.legacyApproval = options.legacyApproval; this.reverseRpc = new ReverseRpcController((event) => this.emitStreamEvent(event)); // Forward every approval request to the user. The engine permission mode // (mapped from the legacy flags) already auto-approves what yolo/auto // allow internally; anything that reaches this handler is an exception // (sensitive file, plan review, ask rule) the user must decide on. this.session.setApprovalHandler((request) => this.reverseRpc.requestApproval(request)); this.session.setQuestionHandler((request) => this.reverseRpc.requestQuestion(request)); this.unsubscribe = this.session.onEvent((event) => this.onSdkEvent(event)); } get id(): string { return this.session.id; } get summary(): SessionSummary | undefined { return this.session.summary; } get subscribers(): readonly string[] { return [...this.webviewIds]; } get isBusy(): boolean { return this.hasActiveWork || this.exclusiveActionActive; } get legacyApprovalFlags(): LegacyApprovalFlags { return this.legacyApproval; } async toggleLegacyApproval(kind: keyof LegacyApprovalFlags): Promise { const next = { ...this.legacyApproval, [kind]: !this.legacyApproval[kind] }; await this.applyLegacyApproval(next); return next; } async setLegacyYoloMode(enabled: boolean): Promise { if (this.legacyApproval.yolo === enabled) return; await this.applyLegacyApproval({ ...this.legacyApproval, yolo: enabled }); } subscribe(webviewId: string): void { this.ensureOpen(); this.webviewIds.add(webviewId); } /** * Push the session's current status to a view. Called whenever a view * opens or re-enters a session so the display (model, thinking effort, * plan mode) matches engine truth instead of the global defaults. */ async announceStatus(webviewId: string): Promise { this.ensureOpen(); const status = await this.session.getStatus(); if (this.closed || !this.webviewIds.has(webviewId)) return; this.broadcast( Events.StreamEvent, { type: "StatusUpdate", payload: { model: status.model, thinking_effort: status.thinkingEffort, plan_mode: status.planMode, context_usage: status.contextUsage, }, _sessionId: this.id, }, webviewId, ); } unsubscribeView(webviewId: string): void { this.webviewIds.delete(webviewId); } async prompt(input: string | LegacyContentPart[]): Promise { return this.runTurnAction(input, () => this.session.prompt(toSdkPromptInput(input))); } async runTurnAction( input: string | LegacyContentPart[], action: () => Promise, ): Promise { this.ensureOpen(); if (this.isBusy) { // A re-entrant turn request must never disturb the active turn — it fails // only itself. When a turn or host action is running, its later terminal // stream event unlocks every subscribed view, so a non-terminal warning // is enough. An exclusive operation (e.g. fork materialization) emits no // such terminal event, so reject terminally: the caller's composer must // unlock rather than hang until the handshake timeout. this.emitError( new Error(ALREADY_GENERATING_MESSAGE), "runtime", { terminal: this.hasActiveWork ? false : undefined }, ); return { status: "failed" }; } let resolveCompletion!: (result: PromptResult) => void; const completion = new Promise((resolve) => { resolveCompletion = resolve; }); const active: ActivePrompt = { input, started: false, settled: false, resolve: resolveCompletion, }; this.activePrompt = active; try { await action(); } catch (error) { // Only settle the prompt this call created. Once the event pipeline or a // cancel has settled it, the failure was already reported — settling it // again would misreport the active turn and could emit a duplicate error. if (!active.settled) { this.emitError(error, active.started ? "runtime" : "preflight"); this.settlePrompt({ status: "failed" }); } } return completion; } beginHostAction(input: string | LegacyContentPart[], forkable = false): number { this.ensureOpen(); if (this.isBusy) { throw new Error(ALREADY_GENERATING_MESSAGE); } const actionId = ++this.hostActionSequence; this.hostActionActive = true; this.activeHostActionId = actionId; this.emitStreamEvent({ type: "TurnBegin", payload: { user_input: input, forkable }, _sessionId: this.id, }); this.emitStreamEvent({ type: "StepBegin", payload: { n: 1 }, _sessionId: this.id, }); return actionId; } emitHostText(text: string, actionId = this.activeHostActionId): void { if (!this.hostActionActive || actionId !== this.activeHostActionId || text.length === 0) return; this.emitStreamEvent({ type: "ContentPart", payload: { type: "text", text }, _sessionId: this.id, }); } announceSessionStart(model?: string): void { this.emitStreamEvent({ type: "session_start", sessionId: this.id, ...(model === undefined ? {} : { model }), _sessionId: this.id, }); } completeHostAction( status: "finished" | "cancelled" = "finished", actionId = this.activeHostActionId, ): void { if (!this.hostActionActive || actionId !== this.activeHostActionId) return; this.hostActionActive = false; this.activeHostActionId = undefined; this.emitStreamEvent({ type: "stream_complete", result: { status }, _sessionId: this.id, }); this.notifyActiveWorkSettled(); } failHostAction(actionId: number): void { if (!this.hostActionActive || actionId !== this.activeHostActionId) return; this.hostActionActive = false; this.activeHostActionId = undefined; this.notifyActiveWorkSettled(); } wasHostActionCancelled(actionId: number): boolean { return this.cancelledHostActions.has(actionId); } releaseHostAction(actionId: number): void { this.cancelledHostActions.delete(actionId); } async compactHostAction(actionId: number, instruction?: string): Promise { if (!this.hostActionActive || actionId !== this.activeHostActionId) { throw new Error("The host action is no longer active."); } if (this.pendingHostCompaction !== undefined) { throw new Error("A context compaction is already running."); } let resolveCompletion!: (result: "completed" | "cancelled") => void; let rejectCompletion!: (error: unknown) => void; const completion = new Promise<"completed" | "cancelled">((resolve, reject) => { resolveCompletion = resolve; rejectCompletion = reject; }); this.pendingHostCompaction = { actionId, resolve: resolveCompletion, reject: rejectCompletion, }; try { await this.session.compact(instruction === undefined ? {} : { instruction }); } catch (error) { if (this.pendingHostCompaction?.actionId === actionId) { this.pendingHostCompaction = undefined; rejectCompletion(error); } } const result = await completion; if (result === "cancelled") { throw new Error("Context compaction was cancelled."); } } async cancel(): Promise { // 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; if (cancellingHostAction && hostActionId !== undefined) { this.cancelledHostActions.add(hostActionId); this.completeHostAction("cancelled", hostActionId); } // A manual compaction is not a model turn, so Session.cancel() alone does // not stop it. Calling both public cancellation surfaces is harmless when // the other operation is idle and keeps the Stop button correct for both // normal turns and host-side slash commands such as /compact and /init. const results = await Promise.allSettled([ this.session.cancel(), ...(cancellingHostAction ? [this.session.cancelCompaction()] : []), ]); const failure = results.find((result): result is PromiseRejectedResult => result.status === "rejected"); if (failure !== undefined) throw failure.reason; } /** * 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 * fully settled session instead of racing the asynchronous cancel event. */ async runExclusiveAfterCancelling(action: () => Promise): Promise { this.ensureOpen(); if (this.exclusiveActionActive) { throw new Error("Another session operation is already in progress."); } this.exclusiveActionActive = true; try { const settled = this.waitForActiveWorkToSettle(); await this.cancel(); await settled; this.ensureOpen(); return await action(); } finally { this.exclusiveActionActive = false; } } async steer(input: string | LegacyContentPart[]): Promise { this.ensureOpen(); await this.session.steer(toSdkPromptInput(input)); this.emitStreamEvent({ type: "SteerInput", payload: { user_input: input }, _sessionId: this.id, }); } respondApproval(id: string, response: ApprovalResponse): boolean { return this.reverseRpc.respondApproval(id, response); } respondQuestion(id: string, answers: Record): boolean { return this.reverseRpc.respondQuestion(id, answers); } async close(): Promise { if (this.closed) return; this.closed = true; this.pendingHostCompaction?.reject(new Error("Session closed during context compaction.")); this.pendingHostCompaction = undefined; this.reverseRpc.cancelAll("Session closed"); this.unsubscribe(); this.session.setApprovalHandler(undefined); this.session.setQuestionHandler(undefined); if (this.activePrompt !== undefined || this.hostActionActive) { try { await this.session.cancel(); } catch (error) { this.log("Failed to cancel the active turn while closing a session", error); } if (this.activePrompt !== undefined) this.settlePrompt({ status: "cancelled" }); this.hostActionActive = false; this.activeHostActionId = undefined; this.notifyActiveWorkSettled(); } this.cancelledHostActions.clear(); await this.session.close(); this.webviewIds.clear(); } private async applyLegacyApproval(flags: LegacyApprovalFlags): Promise { this.ensureOpen(); const permission = corePermissionForLegacyApproval(flags); const status = await this.session.getStatus(); const permissionChanged = status.permission !== permission; if (permissionChanged) await this.session.setPermission(permission); try { await this.session.updateMetadata(legacyApprovalMetadata(flags)); } catch (error) { if (permissionChanged) { await this.session.setPermission(status.permission).catch((rollbackError: unknown) => { this.log("Failed to restore session permission after a metadata error", rollbackError); }); } throw error; } this.legacyApproval = flags; } private onSdkEvent(event: Event): void { if (this.closed) return; if (event.type === "compaction.completed" || event.type === "compaction.cancelled") { const pending = this.pendingHostCompaction; if (pending !== undefined) { this.pendingHostCompaction = undefined; pending.resolve(event.type === "compaction.completed" ? "completed" : "cancelled"); } } if (event.type === "turn.started" && event.agentId === "main" && this.activePrompt !== undefined) { this.activePrompt.started = true; } if (event.type === "tool.call.started") { this.captureFileBaseline(event); } if (event.type === "turn.step.retrying") { this.log( `Provider retry ${event.nextAttempt}/${event.maxAttempts} in ${event.delayMs}ms`, new Error(event.errorMessage), ); } if (event.type === "error" && this.consumeSuppressedError(event.code, event.message)) { return; } const pendingInput = this.activePrompt?.input; const adapted = adaptSdkEvent(this.adapterState, event, { pendingInput, errorPhase: this.activePrompt?.started === false ? "preflight" : "runtime", }); this.adapterState = adapted.state; if (adapted.terminal !== undefined) { this.emitTerminal(adapted.terminal); return; } if (adapted.event !== undefined) { // Errors the core reports while the active turn keeps running (they are // not followed by a terminal turn.ended) must not look turn-ending to the // Webview — otherwise the UI unlocks mid-turn and the next send collides // with the still-active prompt. const wireEvent = adapted.event.type === "error" && this.activePrompt?.started === true ? { ...adapted.event, terminal: false as const } : adapted.event; this.emitStreamEvent(wireEvent); if (adapted.event.type === "error" && this.activePrompt !== undefined && !this.activePrompt.started) { this.settlePrompt({ status: "failed" }); } } } private captureFileBaseline(event: Extract): void { if (event.name !== "Write" && event.name !== "Edit") return; if (!isRecord(event.args)) return; const filePath = event.args["path"]; if (typeof filePath !== "string" || filePath.length === 0) return; const summary = this.session.summary; this.captureBaseline( { id: this.session.id, workDir: this.session.workDir, metadata: summary?.metadata, }, filePath, this.subscribers, ); } private emitTerminal(terminal: TurnTerminalMetadata): void { if (this.terminalKeys.has(terminal.key)) return; this.terminalKeys.add(terminal.key); if (terminal.reason === "completed") { this.emitStreamEvent({ type: "stream_complete", result: { status: "finished" }, _sessionId: terminal.sessionId, }); this.settlePrompt({ status: "finished" }); return; } if (terminal.reason === "cancelled") { this.reverseRpc.cancelAll("Turn cancelled"); this.emitStreamEvent({ type: "stream_complete", result: { status: "cancelled" }, _sessionId: terminal.sessionId, }); this.settlePrompt({ status: "cancelled" }); return; } const code = terminal.error?.code ?? `turn.${terminal.reason}`; this.reverseRpc.cancelAll("Turn ended"); const detail = terminal.error?.message ?? `Turn ended with reason: ${terminal.reason}`; const message = getUserMessage(code, detail); this.log("Session turn failed", new Error(`${code}: ${detail}`)); this.emitStreamEvent({ type: "error", code, message, detail, phase: "runtime", _sessionId: terminal.sessionId, }); if (terminal.error !== undefined) { this.suppressedError = { code: terminal.error.code, message: terminal.error.message }; } this.settlePrompt({ status: "failed" }); } private consumeSuppressedError(code: string, message: string): boolean { const suppressed = this.suppressedError; if (suppressed === undefined) return false; this.suppressedError = undefined; return suppressed.code === code && suppressed.message === message; } private emitError(error: unknown, phase: ErrorPhase, options?: { readonly terminal?: boolean }): void { const code = isKimiError(error) ? error.code : "internal"; const detail = error instanceof Error ? error.message : String(error); this.log(`Session ${phase} error`, error); this.emitStreamEvent({ type: "error", code, message: getUserMessage(code, detail), detail, phase, _sessionId: this.session.id, terminal: options?.terminal, }); } private emitStreamEvent(event: UIStreamEvent | { type: string; payload: unknown }): void { for (const webviewId of this.webviewIds) { this.broadcast(Events.StreamEvent, event, webviewId); } } private settlePrompt(result: PromptResult): void { const active = this.activePrompt; if (active === undefined || active.settled) return; active.settled = true; this.activePrompt = undefined; active.resolve(result); this.notifyActiveWorkSettled(); } private get hasActiveWork(): boolean { return this.activePrompt !== undefined || this.hostActionActive; } private waitForActiveWorkToSettle(): Promise { if (!this.hasActiveWork) return Promise.resolve(); return new Promise((resolve) => { this.activeWorkSettledWaiters.add(resolve); }); } private notifyActiveWorkSettled(): void { if (this.hasActiveWork) return; for (const resolve of this.activeWorkSettledWaiters) resolve(); this.activeWorkSettledWaiters.clear(); } private ensureOpen(): void { if (this.closed) throw new Error("Session is closed."); } } export function toSdkPromptInput(input: string | LegacyContentPart[]): string | PromptInput { if (typeof input === "string") return input; const parts: SdkContentPart[] = []; for (const part of input) { switch (part.type) { case "text": parts.push({ type: "text", text: part.text }); break; case "image_url": parts.push({ type: "image_url", imageUrl: { url: part.image_url.url, ...(part.image_url.id === null || part.image_url.id === undefined ? {} : { id: part.image_url.id }), }, }); break; case "video_url": parts.push({ type: "video_url", videoUrl: { url: part.video_url.url, ...(part.video_url.id === null || part.video_url.id === undefined ? {} : { id: part.video_url.id }), }, }); break; case "audio_url": case "think": // PromptInput intentionally accepts user text/images/videos only. break; } } return parts as PromptInput; } function isRecord(value: unknown): value is Record { return typeof value === "object" && value !== null && !Array.isArray(value); }