From 2efa216ee5bfffa6e01f51ae5add7c5b9034833c Mon Sep 17 00:00:00 2001 From: marius-kilocode Date: Fri, 29 May 2026 11:41:13 +0200 Subject: [PATCH] fix(agent-manager): keep slow snapshot initialization non-blocking --- .changeset/agent-manager-snapshot-wait.md | 6 ++++ packages/kilo-vscode/src/KiloProvider.ts | 2 ++ .../src/agent-manager/constants.ts | 3 ++ .../src/agent-manager/tool-start.ts | 3 +- .../src/agent-manager/vscode-host.ts | 3 +- .../kilo-vscode/src/kilo-provider/options.ts | 1 + .../unit/agent-manager-tool-start.test.ts | 1 + .../opencode/src/kilocode/snapshot/track.ts | 23 ++++++++++++- packages/opencode/src/session/processor.ts | 15 +++++--- packages/opencode/src/session/prompt.ts | 24 ++++++++++--- packages/opencode/src/snapshot/index.ts | 34 ++++++++++--------- .../kilocode/snapshot-track-timeout.test.ts | 23 +++++++++++++ packages/sdk/js/src/v2/gen/sdk.gen.ts | 6 ++++ packages/sdk/js/src/v2/gen/types.gen.ts | 3 ++ packages/sdk/openapi.json | 12 +++++++ 15 files changed, 132 insertions(+), 27 deletions(-) create mode 100644 .changeset/agent-manager-snapshot-wait.md diff --git a/.changeset/agent-manager-snapshot-wait.md b/.changeset/agent-manager-snapshot-wait.md new file mode 100644 index 00000000000..4f3228ca08c --- /dev/null +++ b/.changeset/agent-manager-snapshot-wait.md @@ -0,0 +1,6 @@ +--- +"@kilocode/cli": patch +"kilo-code": patch +--- + +Keep Agent Manager turns running while slow snapshot baselines initialize instead of stopping for an interactive question. diff --git a/packages/kilo-vscode/src/KiloProvider.ts b/packages/kilo-vscode/src/KiloProvider.ts index 8657c062cf2..02d5afea78b 100644 --- a/packages/kilo-vscode/src/KiloProvider.ts +++ b/packages/kilo-vscode/src/KiloProvider.ts @@ -2655,6 +2655,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper agent, variant, editorContext, + snapshotInitialization: this.opts.snapshotInitialization, }), sid, messageID, @@ -2732,6 +2733,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper agent, variant, parts, + snapshotInitialization: this.opts.snapshotInitialization, }), sid, messageID, diff --git a/packages/kilo-vscode/src/agent-manager/constants.ts b/packages/kilo-vscode/src/agent-manager/constants.ts index 771ce4172f2..a9592acc85a 100644 --- a/packages/kilo-vscode/src/agent-manager/constants.ts +++ b/packages/kilo-vscode/src/agent-manager/constants.ts @@ -13,6 +13,9 @@ export const MAX_MULTI_VERSIONS = 4 /** Telemetry source identifier for all Agent Manager events. */ export const PLATFORM = "agent-manager" as const +/** Keep baseline snapshots without interrupting concurrently started agents. */ +export const SNAPSHOT_INITIALIZATION = "wait" as const + /** Kilo config directory name (project-level and inside worktrees). */ export const KILO_DIR = ".kilo" diff --git a/packages/kilo-vscode/src/agent-manager/tool-start.ts b/packages/kilo-vscode/src/agent-manager/tool-start.ts index 51b9b0bdcb4..da261902f69 100644 --- a/packages/kilo-vscode/src/agent-manager/tool-start.ts +++ b/packages/kilo-vscode/src/agent-manager/tool-start.ts @@ -3,7 +3,7 @@ import { sanitizeBranchName, versionedName } from "./branch-name" import type { CreateWorktreeResult } from "./WorktreeManager" import type { WorktreeStateManager } from "./WorktreeStateManager" import type { PanelContext } from "./host" -import { PLATFORM } from "./constants" +import { PLATFORM, SNAPSHOT_INITIALIZATION } from "./constants" const LABEL_MAX = 28 const PREFIX = new Set(["feat", "fix", "chore", "bug", "issue", "task", "branch"]) @@ -103,6 +103,7 @@ async function prompt(client: KiloClient, sid: string, dir: string, task: ToolTa sessionID: sid, directory: dir, parts: [{ type: "text", text: body }], + snapshotInitialization: SNAPSHOT_INITIALIZATION, }, { throwOnError: true }, ) diff --git a/packages/kilo-vscode/src/agent-manager/vscode-host.ts b/packages/kilo-vscode/src/agent-manager/vscode-host.ts index 94f28cba3b1..8d5f9ed9aca 100644 --- a/packages/kilo-vscode/src/agent-manager/vscode-host.ts +++ b/packages/kilo-vscode/src/agent-manager/vscode-host.ts @@ -9,7 +9,7 @@ import * as vscode from "vscode" import type { Host, PanelContext, OutputHandle, SessionProvider, Disposable } from "./host" import type { KiloConnectionService } from "../services/cli-backend" import { KiloProvider } from "../KiloProvider" -import { PLATFORM } from "./constants" +import { PLATFORM, SNAPSHOT_INITIALIZATION } from "./constants" import { DiffVirtualProvider } from "../DiffVirtualProvider" import { buildWebviewHtml } from "../utils" import { openFileInEditor, getWorkspaceRoot } from "../review-utils" @@ -90,6 +90,7 @@ export class VscodeHost implements Host { const provider = new KiloProvider(this.extensionUri, this.connectionService, this.context, { platform: PLATFORM, + snapshotInitialization: SNAPSHOT_INITIALIZATION, slimEditMetadata: true, worktreeDirectories: () => opts.worktreeDirectories?.() ?? [], }) diff --git a/packages/kilo-vscode/src/kilo-provider/options.ts b/packages/kilo-vscode/src/kilo-provider/options.ts index 8fad1c57c48..e8f747ff8f2 100644 --- a/packages/kilo-vscode/src/kilo-provider/options.ts +++ b/packages/kilo-vscode/src/kilo-provider/options.ts @@ -1,6 +1,7 @@ export type KiloProviderOptions = { projectDirectory?: string | null platform?: string + snapshotInitialization?: "wait" slimEditMetadata?: boolean tabTitle?: (title: string) => void onSidebarVisibilityChange?: (visible: boolean) => void diff --git a/packages/kilo-vscode/tests/unit/agent-manager-tool-start.test.ts b/packages/kilo-vscode/tests/unit/agent-manager-tool-start.test.ts index 8742ab2776a..7d1c158bb73 100644 --- a/packages/kilo-vscode/tests/unit/agent-manager-tool-start.test.ts +++ b/packages/kilo-vscode/tests/unit/agent-manager-tool-start.test.ts @@ -87,6 +87,7 @@ describe("agent manager tool start", () => { sessionID: "s-local", directory: "/repo", parts: [{ type: "text", text: "Do work" }], + snapshotInitialization: "wait", }), { throwOnError: true }, ) diff --git a/packages/opencode/src/kilocode/snapshot/track.ts b/packages/opencode/src/kilocode/snapshot/track.ts index 724eb75d749..dff8ebefdc0 100644 --- a/packages/opencode/src/kilocode/snapshot/track.ts +++ b/packages/opencode/src/kilocode/snapshot/track.ts @@ -9,7 +9,8 @@ // // 1. Runs the real `track()` in a forked fiber. // 2. Waits up to `TIMEOUT_MS` for it to complete. -// 3. If it times out AND we have a sessionID to target, asks the user: +// 3. If it times out AND we have a sessionID to target, either waits +// silently when the caller selected that product policy, or asks the user: // - "Continue with snapshots": keep waiting on this turn; snapshot // finishes eventually and undo/redo stays functional. Future turns // are fast because the snapshot index is built. @@ -139,6 +140,9 @@ export namespace KiloSnapshotTrack { /** Answer shape returned by `askUser`. Three-valued because dismiss !== disable. */ export type Answer = "continue" | "disable" | "dismissed" + /** Managed callers can retain snapshots without blocking on the interactive warning. */ + export type SnapshotInitialization = "wait" + /** * Hooks injected by the snapshot layer. Split out so the unit tests can * substitute fakes without reaching for the real Question/filesystem stack. @@ -180,6 +184,7 @@ export namespace KiloSnapshotTrack { export interface WrapInput { readonly inner: Effect.Effect readonly state: State + readonly snapshotInitialization?: SnapshotInitialization readonly sessionID?: SessionID /** * When provided, the wrapper injects an in-message "initializing snapshot…" @@ -302,6 +307,22 @@ export namespace KiloSnapshotTrack { // dialog is visible. The progress fiber is only torn down once we've // decided whether to keep waiting, disable, or skip below. + // Managed products such as Agent Manager expect concurrent snapshot + // initialization and cannot stop started turns for an inline question. + // Retain the snapshot baseline, but wait silently after the threshold. + if (input.snapshotInitialization === "wait") { + log.info("snapshot track slow; waiting without question") + const finished = yield* Fiber.join(fiber).pipe( + Effect.catch((err) => { + log.warn("snapshot track failed while waiting without question", { err }) + return Effect.succeed(undefined as string | undefined) + }), + ) + if (progressFiber) yield* Fiber.interrupt(progressFiber) + yield* clearProgress() + return finished + } + // Slow path. No target session to prompt against, or we've already // prompted on this instance — skip silently. if (!input.sessionID || input.state.asked) { diff --git a/packages/opencode/src/session/processor.ts b/packages/opencode/src/session/processor.ts index 9723f489339..31808d9f164 100644 --- a/packages/opencode/src/session/processor.ts +++ b/packages/opencode/src/session/processor.ts @@ -59,6 +59,7 @@ type Input = { sessionID: SessionID model: Provider.Model telemetry?: ReviewTelemetry // kilocode_change + snapshotInitialization?: "wait" // kilocode_change - managed clients wait silently for slow baseline creation } export interface Interface { @@ -121,10 +122,11 @@ export const layer: Layer.Layer< // Pre-capture snapshot before the LLM stream starts. The AI SDK // may execute tools internally before emitting start-step events, // so capturing inside the event handler can be too late. - // kilocode_change start - pass sessionID + messageID so the slow-repo prompt/progress indicator can attach + // kilocode_change start - pass turn context for slow-snapshot UI/policy handling const initialSnapshot = yield* snapshot.track({ sessionID: input.sessionID, messageID: input.assistantMessage.id, + snapshotInitialization: input.snapshotInitialization, }) // kilocode_change end const ctx: ProcessorContext = { @@ -505,9 +507,13 @@ export const layer: Layer.Layer< case "start-step": ctx.stepStart = performance.now() // kilocode_change ctx.step = { reasoning: false, text: false, tool: false } // kilocode_change - // kilocode_change start - pass sessionID + messageID so the slow-repo prompt/progress indicator can attach + // kilocode_change start - pass turn context for slow-snapshot UI/policy handling if (!ctx.snapshot) - ctx.snapshot = yield* snapshot.track({ sessionID: ctx.sessionID, messageID: ctx.assistantMessage.id }) + ctx.snapshot = yield* snapshot.track({ + sessionID: ctx.sessionID, + messageID: ctx.assistantMessage.id, + snapshotInitialization: input.snapshotInitialization, + }) // kilocode_change end if (!ctx.assistantMessage.summary) { // TODO(v2): Temporary dual-write while migrating session messages to v2 events. @@ -533,10 +539,11 @@ export const layer: Layer.Layer< return case "finish-step": { - // kilocode_change start - pass sessionID + messageID + // kilocode_change start - pass turn context for slow-snapshot UI/policy handling const completedSnapshot = yield* snapshot.track({ sessionID: ctx.sessionID, messageID: ctx.assistantMessage.id, + snapshotInitialization: input.snapshotInitialization, }) // kilocode_change end const usage = Session.getUsage({ diff --git a/packages/opencode/src/session/prompt.ts b/packages/opencode/src/session/prompt.ts index 17748fa93f7..d411db4e3ad 100644 --- a/packages/opencode/src/session/prompt.ts +++ b/packages/opencode/src/session/prompt.ts @@ -1446,7 +1446,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the return yield* KiloSessionPromptQueue.enqueue( input.sessionID, message.info.id, - loop({ sessionID: input.sessionID }), + loop({ sessionID: input.sessionID, snapshotInitialization: input.snapshotInitialization }), // kilocode_change lastAssistant(input.sessionID), ) // kilocode_change end @@ -1469,8 +1469,11 @@ NOTE: At any point in time through this workflow you should feel free to ask the // kilocode_change — mutable close-reason per session, set by runLoop and read by loop const closeReasons = new Map() - const runLoop: (sessionID: SessionID) => Effect.Effect = Effect.fn("SessionPrompt.run")( - function* (sessionID: SessionID) { + // kilocode_change start - retain request-scoped snapshot initialization policy + const runLoop: (input: LoopInput) => Effect.Effect = Effect.fn("SessionPrompt.run")( + function* (input: LoopInput) { + const sessionID = input.sessionID + // kilocode_change end // kilocode_change — cache environment details per turn (prompt caching) const envCache: KiloSessionPrompt.EnvCache = {} closeReasons.delete(sessionID) // kilocode_change @@ -1652,6 +1655,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the sessionID, model, telemetry, // kilocode_change + snapshotInitialization: input.snapshotInitialization, // kilocode_change }) const outcome: "break" | "continue" = yield* Effect.gen(function* () { @@ -1846,7 +1850,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the yield* KiloSessionPrompt.recoverProviderFinishError({ sessionID: input.sessionID, status, sessions }) yield* bus.publish(KiloSession.Event.TurnOpen, { sessionID: input.sessionID }) return yield* Effect.onExit( - state.ensureRunning(input.sessionID, lastAssistant(input.sessionID), runLoop(input.sessionID)), + state.ensureRunning(input.sessionID, lastAssistant(input.sessionID), runLoop(input)), // kilocode_change Effect.fnUntraced(function* (exit) { yield* bus.publish(KiloSession.Event.TurnClose, { sessionID: input.sessionID, @@ -1976,6 +1980,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the agent: userAgent, parts, variant: input.variant, + snapshotInitialization: input.snapshotInitialization, // kilocode_change }) yield* bus.publish(Command.Event.Executed, { name: input.command, @@ -2046,6 +2051,11 @@ export const PromptInput = Schema.Struct({ format: Schema.optional(MessageV2.Format), system: Schema.optional(Schema.String), variant: Schema.optional(Schema.String), + // kilocode_change start - managed product slow-snapshot policy + snapshotInitialization: Schema.optional(Schema.Literal("wait")).annotate({ + description: "Wait silently if snapshot initialization is slow instead of asking the user.", + }), + // kilocode_change end // kilocode_change start - reuse shared editor context schema editorContext: Schema.optional(MessageV2.EditorContext), // kilocode_change end @@ -2074,6 +2084,7 @@ export type PromptInput = Omit, "parts" | export class LoopInput extends Schema.Class("SessionPrompt.LoopInput")({ sessionID: SessionID, + snapshotInitialization: Schema.optional(Schema.Literal("wait")), // kilocode_change }) { static readonly zod = zod(this) } @@ -2095,6 +2106,11 @@ export const CommandInput = Schema.Struct({ arguments: Schema.String, command: Schema.String, variant: Schema.optional(Schema.String), + // kilocode_change start - managed product slow-snapshot policy + snapshotInitialization: Schema.optional(Schema.Literal("wait")).annotate({ + description: "Wait silently if snapshot initialization is slow instead of asking the user.", + }), + // kilocode_change end // Inlined (no identifier annotation) to keep the original SDK output — the // PromptInput call site below references FilePartInput by ref via the // Schema export in message-v2.ts. diff --git a/packages/opencode/src/snapshot/index.ts b/packages/opencode/src/snapshot/index.ts index 0ddbd6f8b48..cb7d649b582 100644 --- a/packages/opencode/src/snapshot/index.ts +++ b/packages/opencode/src/snapshot/index.ts @@ -59,14 +59,17 @@ const cache = new Map>() const max = 100 // kilocode_change end -type State = Omit +type State = Omit & { readonly trackState: KiloSnapshotTrack.State } // kilocode_change export interface Interface { readonly init: () => Effect.Effect readonly cleanup: () => Effect.Effect - // kilocode_change start - accept optional sessionID/messageID so the slow-repo prompt can target - // a client and the in-message "initializing snapshot" indicator can attach to the live turn - readonly track: (opts?: { sessionID?: SessionID; messageID?: MessageID }) => Effect.Effect + // kilocode_change start - accept prompt context so slow snapshots can target UI or honor managed caller policy + readonly track: (opts?: { + sessionID?: SessionID + messageID?: MessageID + snapshotInitialization?: KiloSnapshotTrack.SnapshotInitialization + }) => Effect.Effect // kilocode_change end readonly patch: (hash: string) => Effect.Effect readonly restore: (snapshot: string) => Effect.Effect @@ -782,14 +785,10 @@ export const layer: Layer.Layer< Effect.forkScoped, ) - return { cleanup, track, patch, restore, revert, diff, diffFull } + return { cleanup, track, patch, restore, revert, diff, diffFull, trackState: KiloSnapshotTrack.makeState() } // kilocode_change - scope slow prompt state by worktree }), ) - // kilocode_change start - per-instance state for the slow-repo track wrapper - const trackState = KiloSnapshotTrack.makeState() - // kilocode_change end - return Service.of({ init: Effect.fn("Snapshot.init")(function* () { yield* InstanceState.get(state) @@ -797,14 +796,17 @@ export const layer: Layer.Layer< cleanup: Effect.fn("Snapshot.cleanup")(function* () { return yield* InstanceState.useEffect(state, (s) => s.cleanup()) }), - // kilocode_change start - timeout + interactive "disable for this project" prompt + // kilocode_change start - timeout guard with interactive and managed wait policies track: Effect.fn("Snapshot.track")(function* (opts) { - return yield* KiloSnapshotTrack.wrap({ - inner: InstanceState.useEffect(state, (s) => s.track()), - state: trackState, - sessionID: opts?.sessionID, - messageID: opts?.messageID, - }) + return yield* InstanceState.useEffect(state, (s) => + KiloSnapshotTrack.wrap({ + inner: s.track(), + state: s.trackState, + snapshotInitialization: opts?.snapshotInitialization, + sessionID: opts?.sessionID, + messageID: opts?.messageID, + }), + ) // kilocode_change end }), patch: Effect.fn("Snapshot.patch")(function* (hash: string) { diff --git a/packages/opencode/test/kilocode/snapshot-track-timeout.test.ts b/packages/opencode/test/kilocode/snapshot-track-timeout.test.ts index 896a336e449..6e5488cfd8a 100644 --- a/packages/opencode/test/kilocode/snapshot-track-timeout.test.ts +++ b/packages/opencode/test/kilocode/snapshot-track-timeout.test.ts @@ -289,6 +289,29 @@ describe("KiloSnapshotTrack.wrap", () => { expect(state.disabledForSession).toBe(false) }) + test('timeout + snapshot initialization "wait" keeps waiting without asking', async () => { + const state = KiloSnapshotTrack.makeState() + const { hooks, calls } = makeHooks("disable") + + const result = await Effect.runPromise( + KiloSnapshotTrack.wrap({ + inner: slowInner(50, "managed-hash"), + state, + snapshotInitialization: "wait", + sessionID: SESSION, + messageID: MESSAGE, + hooks, + timeoutMs: 10, + progressDelayMs: 2, + }), + ) + + expect(result).toBe("managed-hash") + expect(calls.ask).toBe(0) + expect(calls.persist).toBe(0) + expect(state.disabledForSession).toBe(false) + }) + test("continue path keeps `asked` sticky when the fiber finished with no hash", async () => { const state = KiloSnapshotTrack.makeState() const { hooks, calls } = makeHooks("continue") diff --git a/packages/sdk/js/src/v2/gen/sdk.gen.ts b/packages/sdk/js/src/v2/gen/sdk.gen.ts index 31297a92fdc..e6e0cf5b9a9 100644 --- a/packages/sdk/js/src/v2/gen/sdk.gen.ts +++ b/packages/sdk/js/src/v2/gen/sdk.gen.ts @@ -3599,6 +3599,7 @@ export class Session2 extends HeyApiClient { format?: OutputFormat system?: string variant?: string + snapshotInitialization?: "wait" editorContext?: { visibleFiles?: Array openTabs?: Array @@ -3625,6 +3626,7 @@ export class Session2 extends HeyApiClient { { in: "body", key: "format" }, { in: "body", key: "system" }, { in: "body", key: "variant" }, + { in: "body", key: "snapshotInitialization" }, { in: "body", key: "editorContext" }, { in: "body", key: "parts" }, ], @@ -3959,6 +3961,7 @@ export class Session2 extends HeyApiClient { format?: OutputFormat system?: string variant?: string + snapshotInitialization?: "wait" editorContext?: { visibleFiles?: Array openTabs?: Array @@ -3985,6 +3988,7 @@ export class Session2 extends HeyApiClient { { in: "body", key: "format" }, { in: "body", key: "system" }, { in: "body", key: "variant" }, + { in: "body", key: "snapshotInitialization" }, { in: "body", key: "editorContext" }, { in: "body", key: "parts" }, ], @@ -4019,6 +4023,7 @@ export class Session2 extends HeyApiClient { arguments?: string command?: string variant?: string + snapshotInitialization?: "wait" parts?: Array<{ id?: string type: "file" @@ -4044,6 +4049,7 @@ export class Session2 extends HeyApiClient { { in: "body", key: "arguments" }, { in: "body", key: "command" }, { in: "body", key: "variant" }, + { in: "body", key: "snapshotInitialization" }, { in: "body", key: "parts" }, ], }, diff --git a/packages/sdk/js/src/v2/gen/types.gen.ts b/packages/sdk/js/src/v2/gen/types.gen.ts index fdebfac185c..5311a337028 100644 --- a/packages/sdk/js/src/v2/gen/types.gen.ts +++ b/packages/sdk/js/src/v2/gen/types.gen.ts @@ -5967,6 +5967,7 @@ export type SessionPromptData = { format?: OutputFormat system?: string variant?: string + snapshotInitialization?: "wait" editorContext?: { visibleFiles?: Array openTabs?: Array @@ -6308,6 +6309,7 @@ export type SessionPromptAsyncData = { format?: OutputFormat system?: string variant?: string + snapshotInitialization?: "wait" editorContext?: { visibleFiles?: Array openTabs?: Array @@ -6356,6 +6358,7 @@ export type SessionCommandData = { arguments: string command: string variant?: string + snapshotInitialization?: "wait" parts?: Array<{ id?: string type: "file" diff --git a/packages/sdk/openapi.json b/packages/sdk/openapi.json index 8f41c9639c3..62c5858362a 100644 --- a/packages/sdk/openapi.json +++ b/packages/sdk/openapi.json @@ -5723,6 +5723,10 @@ "variant": { "type": "string" }, + "snapshotInitialization": { + "type": "string", + "enum": ["wait"] + }, "editorContext": { "type": "object", "properties": { @@ -6547,6 +6551,10 @@ "variant": { "type": "string" }, + "snapshotInitialization": { + "type": "string", + "enum": ["wait"] + }, "editorContext": { "type": "object", "properties": { @@ -6706,6 +6714,10 @@ "variant": { "type": "string" }, + "snapshotInitialization": { + "type": "string", + "enum": ["wait"] + }, "parts": { "type": "array", "items": {