fix: abort Agent Manager prompts during startup

.changeset/abort-starting-sessions.md

@@ -0,0 +1,6 @@
+---
+"kilo-code": patch
+"@kilocode/cli": patch
+---
+
+Allow Escape to stop Agent Manager prompts while their sessions are still starting.

packages/kilo-vscode/tests/unit/abort-state.test.ts

@@ -0,0 +1,64 @@
+import { describe, expect, it } from "bun:test"
+import { createAbortState } from "../../webview-ui/src/context/abort-state"
+
+describe("pending prompt abort state", () => {
+  it("waits for the same submission to become cancellable", () => {
+    const aborts = createAbortState()
+
+    expect(aborts.request("draft", "idle", "message")).toBe(false)
+    expect(aborts.update("draft", "busy")).toBe(true)
+    expect(aborts.update("draft", "busy")).toBe(false)
+    expect(aborts.update("draft", "idle")).toBe(false)
+    expect(aborts.update("draft", "busy")).toBe(false)
+  })
+
+  it("moves pending cancellation to the created session", () => {
+    const aborts = createAbortState()
+
+    expect(aborts.request("draft", "idle", "message")).toBe(false)
+    aborts.move("draft", "session")
+
+    expect(aborts.update("draft", "busy")).toBe(false)
+    expect(aborts.update("session", "busy")).toBe(true)
+  })
+
+  it("does not retain cancellation after an idle terminal status", () => {
+    const aborts = createAbortState()
+
+    expect(aborts.request("session", "idle", "message")).toBe(false)
+    expect(aborts.update("session", "idle")).toBe(false)
+    expect(aborts.update("session", "busy")).toBe(false)
+  })
+
+  it("allows retrying an abort while the session remains active", () => {
+    const aborts = createAbortState()
+
+    expect(aborts.request("session", "busy")).toBe(true)
+    expect(aborts.request("session", "busy")).toBe(true)
+    expect(aborts.update("session", "idle")).toBe(false)
+    expect(aborts.request("session", "busy")).toBe(true)
+  })
+
+  it("clears cancellation when the matching submission finishes", () => {
+    const aborts = createAbortState()
+
+    expect(aborts.request("session", "idle", "message")).toBe(false)
+    aborts.finish("other")
+    expect(aborts.update("session", "busy")).toBe(true)
+
+    expect(aborts.update("session", "idle")).toBe(false)
+    expect(aborts.request("session", "idle", "message")).toBe(false)
+    aborts.finish("message")
+    expect(aborts.update("session", "busy")).toBe(false)
+  })
+
+  it("preserves active destination state during duplicate draft migration", () => {
+    const aborts = createAbortState()
+
+    expect(aborts.request("draft", "idle", "message")).toBe(false)
+    expect(aborts.request("session", "busy")).toBe(true)
+    aborts.move("draft", "session")
+
+    expect(aborts.request("session", "busy")).toBe(true)
+  })
+})

packages/kilo-vscode/tests/unit/prompt-input-utils.test.ts

@@ -171,30 +171,34 @@ describe("isPromptBlocked", () => {
 
 describe("isPromptBusy", () => {
   it("returns true when busy and neither suggesting nor questioning", () => {
-    expect(isPromptBusy("busy", false, false)).toBe(true)
+    expect(isPromptBusy("busy", false, false, false)).toBe(true)
+  })
+
+  it("returns true while submitting before the backend reports busy", () => {
+    expect(isPromptBusy("idle", false, false, true)).toBe(true)
   })
 
   it("returns false when idle regardless of suggesting/questioning", () => {
-    expect(isPromptBusy("idle", false, false)).toBe(false)
-    expect(isPromptBusy("idle", true, false)).toBe(false)
-    expect(isPromptBusy("idle", false, true)).toBe(false)
-    expect(isPromptBusy("idle", true, true)).toBe(false)
+    expect(isPromptBusy("idle", false, false, false)).toBe(false)
+    expect(isPromptBusy("idle", true, false, false)).toBe(false)
+    expect(isPromptBusy("idle", false, true, false)).toBe(false)
+    expect(isPromptBusy("idle", true, true, false)).toBe(false)
   })
 
   it("returns false when busy but suggesting is true (suggestion decoupling)", () => {
-    expect(isPromptBusy("busy", true, false)).toBe(false)
+    expect(isPromptBusy("busy", true, false, false)).toBe(false)
   })
 
   it("returns false when busy but questioning is true (question decoupling)", () => {
-    expect(isPromptBusy("busy", false, true)).toBe(false)
+    expect(isPromptBusy("busy", false, true, false)).toBe(false)
   })
 
   it("returns false when busy and both suggesting and questioning", () => {
-    expect(isPromptBusy("busy", true, true)).toBe(false)
+    expect(isPromptBusy("busy", true, true, false)).toBe(false)
   })
 
   it("returns true for non-idle non-busy status when not suggesting/questioning", () => {
-    expect(isPromptBusy("retry", false, false)).toBe(true)
+    expect(isPromptBusy("retry", false, false, false)).toBe(true)
   })
 })
 

packages/kilo-vscode/webview-ui/src/components/chat/ChatView.tsx

@@ -82,7 +82,7 @@ export const ChatView: Component<ChatViewProps> = (props) => {
   onMount(() => {
     if (props.readonly) return
     const handler = (e: KeyboardEvent) => {
-      if (e.key !== "Escape" || session.status() === "idle" || e.defaultPrevented) return
+      if (e.key !== "Escape" || (!session.submitting() && session.status() === "idle") || e.defaultPrevented) return
       e.preventDefault()
       session.abort()
     }

packages/kilo-vscode/webview-ui/src/components/chat/PromptInput.tsx

@@ -344,7 +344,8 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
   window.addEventListener("exportSessionTranscript", onExport)
   onCleanup(() => window.removeEventListener("exportSessionTranscript", onExport))
 
-  const isBusy = () => isPromptBusy(session.status(), !!props.suggesting?.(), !!props.questioning?.())
+  const isBusy = () =>
+    isPromptBusy(session.status(), !!props.suggesting?.(), !!props.questioning?.(), session.submitting())
   const isDisabled = () => !server.isConnected()
   const canUseSpeech = () => canUseSpeechToText(config(), provider.connected(), server.profileData())
   const speechModel = () => selectedSpeechToTextModel(config())

packages/kilo-vscode/webview-ui/src/components/chat/prompt-input-utils.ts

@@ -87,8 +87,8 @@ export function isPromptBlocked(permissions: number): boolean {
  * Returns false (idle-like) when the session is busy only because
  * a suggestion or question tool call is pending.
  */
-export function isPromptBusy(status: string, suggesting: boolean, questioning: boolean): boolean {
-  return status !== "idle" && !suggesting && !questioning
+export function isPromptBusy(status: string, suggesting: boolean, questioning: boolean, submitting: boolean): boolean {
+  return submitting || (status !== "idle" && !suggesting && !questioning)
 }
 
 /**

packages/kilo-vscode/webview-ui/src/context/abort-state.ts

@@ -0,0 +1,50 @@
+type Entry = {
+  phase: "pending" | "active"
+  messageID?: string
+}
+
+export function createAbortState() {
+  const entries = new Map<string, Entry>()
+
+  return {
+    request(id: string, status: string, messageID?: string) {
+      const entry = entries.get(id)
+      if (entry?.phase === "active") return status !== "idle"
+      if (entry) return false
+      if (status === "idle") {
+        if (!messageID) return false
+        entries.set(id, { phase: "pending", messageID })
+        return false
+      }
+      entries.set(id, { phase: "active" })
+      return true
+    },
+    move(from: string, to: string) {
+      const source = entries.get(from)
+      if (!source) return
+      entries.delete(from)
+      const target = entries.get(to)
+      if (target?.phase === "active") return
+      entries.set(to, source)
+    },
+    update(id: string, status: string) {
+      const entry = entries.get(id)
+      if (!entry) return false
+      if (status === "idle") {
+        entries.delete(id)
+        return false
+      }
+      if (entry.phase === "active") return false
+      entries.set(id, { phase: "active" })
+      return true
+    },
+    finish(messageID: string) {
+      for (const [id, entry] of entries) {
+        if (entry.phase === "pending" && entry.messageID === messageID) entries.delete(id)
+      }
+    },
+    clear(id: string) {
+      entries.delete(id)
+    },
+  }
+}

packages/kilo-vscode/webview-ui/src/context/session.tsx

@@ -71,6 +71,7 @@ import { state as todoState } from "./todo-revert"
 import { getVariant, sessionVariantKeys, transferVariants, variantKey } from "./session-variant-store"
 import { KILO_AUTO, parseModelString } from "../../../src/shared/provider-model"
 import { visibleMessages as filterVisibleMessages } from "./session-queue"
+import { createAbortState } from "./abort-state"
 
 const RECENT_LIMIT = 5
 const MESSAGE_PAGE_LIMIT = 80
@@ -303,6 +304,7 @@ export const SessionProvider: ParentComponent = (props) => {
   const [busySinceMap, setBusySinceMap] = createStore<Record<string, number>>({})
   const [submissionMap, setSubmissionMap] = createStore<Record<string, number>>({})
   const pendingSubmissions = new Map<string, string>()
+  const aborts = createAbortState()
 
   const idle: SessionStatusInfo = { type: "idle" }
 
@@ -442,6 +444,7 @@ export const SessionProvider: ParentComponent = (props) => {
     if (!busySinceMap[sid]) setBusySinceMap(sid, Date.now())
   }
   const finishSubmission = (messageID: string) => {
+    aborts.finish(messageID)
     const sid = pendingSubmissions.get(messageID)
     if (!sid) return
     pendingSubmissions.delete(messageID)
@@ -464,7 +467,9 @@ export const SessionProvider: ParentComponent = (props) => {
   }
   const confirmSubmissions = (sid: string) => {
     for (const [id, scope] of pendingSubmissions) {
-      if (scope === sid) pendingSubmissions.delete(id)
+      if (scope !== sid) continue
+      aborts.finish(id)
+      pendingSubmissions.delete(id)
     }
     setSubmissionMap(
       produce((map) => {
@@ -1070,6 +1075,7 @@ export const SessionProvider: ParentComponent = (props) => {
 
   // Event handlers
   function handleSessionCreated(session: SessionInfo, draftID?: string) {
+    if (draftID) aborts.move(draftID, session.id)
     batch(() => {
       setStore("sessions", session.id, session)
 
@@ -1496,6 +1502,7 @@ export const SessionProvider: ParentComponent = (props) => {
     message?: string,
     next?: number,
   ) {
+    const shouldAbort = aborts.update(sessionID, newStatus)
     confirmSubmissions(sessionID)
     const prev = statusMap[sessionID] ?? { type: "idle" }
     const info: SessionStatusInfo =
@@ -1522,6 +1529,7 @@ export const SessionProvider: ParentComponent = (props) => {
       // messages themselves will be reconciled on the next messagesLoaded.
       pendingOptimistic.delete(sessionID)
     }
+    if (shouldAbort) vscode.postMessage({ type: "abort", sessionID })
   }
 
   function handlePermissionRequest(permission: PermissionRequest) {
@@ -1640,6 +1648,7 @@ export const SessionProvider: ParentComponent = (props) => {
   function handleSendMessageFailed(message: SendMessageFailedMessage) {
     const sid = message.sessionID ?? message.draftID
     if (message.messageID) finishSubmission(message.messageID)
+    if (!message.messageID && sid) aborts.clear(sid)
     if (sid && message.messageID) {
       pendingOptimistic.get(sid)?.delete(message.messageID)
       stash.remove(message.messageID)
@@ -1777,6 +1786,7 @@ export const SessionProvider: ParentComponent = (props) => {
 
   function handleSessionDeleted(sessionID: string) {
     pendingOptimistic.delete(sessionID)
+    aborts.clear(sessionID)
     confirmSubmissions(sessionID)
     batch(() => {
       // Collect message IDs so we can clean up their parts (store + stash)
@@ -2197,10 +2207,13 @@ export const SessionProvider: ParentComponent = (props) => {
 
   function abort() {
     const sessionID = currentSessionID()
-    if (!sessionID) {
-      console.warn("[Kilo New] Cannot abort: no current session")
+    const scope = sessionID ?? draftSessionID()
+    if (!scope) {
+      console.warn("[Kilo New] Cannot abort: no current or pending session")
       return
     }
+    const messageID = [...pendingSubmissions].reverse().find(([, sid]) => sid === scope)?.[0]
+    if (!aborts.request(scope, status(), messageID) || !sessionID) return
 
     vscode.postMessage({
       type: "abort",

packages/kilo-vscode/webview-ui/src/stories/StoryProviders.tsx

@@ -180,6 +180,7 @@ export function mockSessionValue(overrides?: {
     setCurrentSessionID: noop,
     sessions: () => [],
     status: () => status,
+    submitting: () => false,
     statusInfo: () => ({ type: status }),
     closeReason: () => overrides?.closeReason,
     statusText: () => (status === "idle" ? undefined : "Thinking…"),

packages/opencode/src/effect/runner.ts

@@ -1,4 +1,5 @@
 import { Cause, Deferred, Effect, Exit, Fiber, Latch, Schema, Scope, SynchronizedRef } from "effect"
+import { KiloRunner } from "@/kilocode/effect/runner" // kilocode_change
 
 export interface Runner<A, E = never> {
   readonly state: State<A, E>
@@ -80,16 +81,19 @@ export const make = <A, E = never>(
         ] as const,
     ).pipe(Effect.flatten)
 
-  const startRun = (work: Effect.Effect<A, E>, done: Deferred.Deferred<A, E | Cancelled>) =>
-    Effect.gen(function* () {
-      const id = next()
-      const fiber = yield* work.pipe(
-        Effect.onExit((exit) => finishRun(id, done, exit)),
-        Effect.forkIn(scope),
-      )
-      return { id, done, fiber } satisfies RunHandle<A, E>
+  // kilocode_change start - do not let work publish busy before the Running state is committed
+  const startRun = (work: Effect.Effect<A, E>, done: Deferred.Deferred<A, E | Cancelled>) => {
+    const id = next()
+    return KiloRunner.start({
+      work,
+      scope,
+      finish: (exit) => finishRun(id, done, exit),
+      handle: (fiber) => ({ id, done, fiber }) satisfies RunHandle<A, E>,
     })
+  }
+  // kilocode_change end
 
+  // kilocode_change start - open work only after the Running state is committed
   const finishShell = (id: number) =>
     SynchronizedRef.modifyEffect(
       ref,
@@ -98,12 +102,12 @@ export const make = <A, E = never>(
           return [idle, { _tag: "Idle" }] as const
         }
         if (st._tag === "ShellThenRun" && st.shell.id === id) {
-          const run = yield* startRun(st.run.work, st.run.done)
-          return [Effect.void, { _tag: "Running", run }] as const
+          return yield* KiloRunner.commit(startRun(st.run.work, st.run.done), Effect.void)
         }
         return [Effect.void, st] as const
       }),
     ).pipe(Effect.flatten)
+  // kilocode_change end
 
   const stopShell = (shell: ShellHandle<A, E>) =>
     Effect.gen(function* () {
@@ -112,6 +116,7 @@ export const make = <A, E = never>(
       yield* Fiber.interrupt(shell.fiber)
     })
 
+  // kilocode_change start - open work only after the Running state is committed
   const ensureRunning = (work: Effect.Effect<A, E>) =>
     SynchronizedRef.modifyEffect(
       ref,
@@ -130,8 +135,7 @@ export const make = <A, E = never>(
           }
           case "Idle": {
             const done = yield* Deferred.make<A, E | Cancelled>()
-            const run = yield* startRun(work, done)
-            return [awaitDone(done), { _tag: "Running", run }] as const
+            return yield* KiloRunner.commit(startRun(work, done), awaitDone(done))
           }
         }
       }),
@@ -140,6 +144,7 @@ export const make = <A, E = never>(
   const startShell = (work: Effect.Effect<A, E>, ready?: Latch.Latch): Effect.Effect<A, E | Busy> =>
     SynchronizedRef.modifyEffect(
       ref,
+      // kilocode_change end
       Effect.fnUntraced(function* (st) {
         if (st._tag !== "Idle") {
           const reject: Effect.Effect<A, E | Busy> = Effect.fail(new Busy())