feat: keep background agents running when the main agent stops

.changeset/improve-background-agent-control.md

@@ -0,0 +1,7 @@
+---
+"kilo-code": patch
+"@kilocode/cli": patch
+"@kilocode/sdk": patch
+---
+
+Support stopping the main agent without cancelling background agents, and add Stop all to the VS Code background-agent bar.

packages/kilo-vscode/src/KiloProvider.ts

@@ -1140,7 +1140,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
         }
         case "abort":
           this.cancelRetry(message.sessionID ?? "")
-          await this.handleAbort(message.sessionID)
+          await this.handleAbort(message.sessionID, message.scope)
           break
         case "revertSession":
           this.checkpoint(message.sessionID, () =>
@@ -4200,14 +4200,18 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
     await Promise.all([...targets].map((sid) => this.stopSession(sid)))
   }
 
-  private stopSession(sid: string): Promise<boolean> {
+  private stopSession(sid: string, scope?: "session" | "tree"): Promise<boolean> {
     this.cancelRetry(sid)
     const client = this.client
     if (!client) return Promise.resolve(false)
     const dir = this.getWorkspaceDirectory(sid)
     this.aborts.preserve(sid, this.sessionStatusMap.get(sid), dir)
-    return this.aborts.stop(client, sid, dir, (dir, action) =>
-      this.connectionService.runExplicitAbort(sid, dir, action),
+    return this.aborts.stop(
+      client,
+      sid,
+      dir,
+      (dir, action) => this.connectionService.runExplicitAbort(sid, dir, action),
+      scope,
     )
   }
 
@@ -4232,9 +4236,9 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
     }
   }
 
-  private async handleAbort(sessionID?: string): Promise<void> {
+  private async handleAbort(sessionID?: string, scope?: "session" | "tree"): Promise<void> {
     const sid = sessionID || this.currentSession?.id
-    if (!sid || !(await this.stopSession(sid))) return
+    if (!sid || !(await this.stopSession(sid, scope))) return
     this.mark(sid)
     this.publish(sid, { type: "idle" })
   }

packages/kilo-vscode/src/agent-manager/types.ts

@@ -1030,6 +1030,7 @@ interface ForkSessionIn {
 interface AbortIn {
   type: "abort"
   sessionID: string
+  scope?: "session" | "tree"
 }
 
 interface ContinueInWorktreeIn {

packages/kilo-vscode/src/kilo-provider/abort.ts

@@ -32,13 +32,14 @@ export class SessionAbort {
     sessionID: string,
     fallback: string,
     run?: (dir: string, action: () => Promise<void>) => Promise<void>,
+    scope?: "session" | "tree",
   ) {
     const known = this.active.has(sessionID)
     const dirs = [...(this.active.get(sessionID) ?? [])]
     if (!dirs.some((dir) => sameDirectory(dir, fallback))) dirs.push(fallback)
     const results = await Promise.allSettled(
       dirs.map((dir) => {
-        const action = () => abortSession({ client, sessionID, dir })
+        const action = () => abortSession({ client, sessionID, dir, scope })
         return known && run ? run(dir, action) : action()
       }),
     )
@@ -75,6 +76,14 @@ export class SessionAbort {
   }
 }
 
-export async function abortSession(input: { client: KiloClient; sessionID: string; dir: string }) {
-  await input.client.session.abort({ sessionID: input.sessionID, directory: input.dir }, { throwOnError: true })
+export async function abortSession(input: {
+  client: KiloClient
+  sessionID: string
+  dir: string
+  scope?: "session" | "tree"
+}) {
+  await input.client.session.abort(
+    { sessionID: input.sessionID, directory: input.dir, scope: input.scope },
+    { throwOnError: true },
+  )
 }

packages/kilo-vscode/tests/fixtures/session-provider-activity.tsx

@@ -451,7 +451,7 @@ try {
 
   await emit({ type: "sessionStatus", sessionID: "root", status: "busy" })
   value.abort()
-  assert.deepEqual(sent.at(-1), { type: "abort", sessionID: "root" })
+  assert.deepEqual(sent.at(-1), { type: "abort", sessionID: "root", scope: "session" })
   await emit({ type: "sessionTurnClosed", sessionID: "root", reason: "interrupted" })
   await emit({ type: "sessionStatus", sessionID: "root", status: "idle" })
   await check("root", "idle")

packages/kilo-vscode/tests/prompt-background-agents.spec.ts

@@ -0,0 +1,65 @@
+import { expect, test } from "@playwright/test"
+import type { BackgroundJobInfo, WebviewMessage } from "../webview-ui/src/types/messages"
+
+const globals = "colorScheme:dark;theme:kilo-vscode;vscodeTheme:dark-modern"
+
+test("Stop all only cancels this session's running background agents", async ({ page }) => {
+  const calls: WebviewMessage[] = []
+  await page.exposeFunction("record", (message: WebviewMessage) => calls.push(message))
+  await page.addInitScript(() => {
+    const record = (window as unknown as { record: (message: WebviewMessage) => void }).record
+    const parent = "story-session-chat-001"
+    let jobs: BackgroundJobInfo[] = ["first", "second", "finished", "other"].map((id) => ({
+      id,
+      type: "task",
+      title: id,
+      status: id === "finished" ? "completed" : "running",
+      started_at: 1,
+      metadata: {
+        parentSessionId: id === "other" ? "other-session" : parent,
+        sessionId: `child-${id}`,
+        background: true,
+      },
+    }))
+    Object.defineProperty(window, "acquireVsCodeApi", {
+      value: () => ({
+        getState: () => undefined,
+        setState: () => {},
+        postMessage: (message: WebviewMessage) => {
+          if (message.type === "abort" || message.type === "cancelBackgroundJob") {
+            record(message)
+          }
+          if (message.type !== "requestBackgroundJobs" && message.type !== "cancelBackgroundJob") return
+          if (message.type === "cancelBackgroundJob") {
+            jobs = jobs.map((job) => (job.id === message.jobID ? { ...job, status: "cancelled" } : job))
+          }
+          window.postMessage(
+            { type: "backgroundJobsLoaded", sessionID: message.sessionID, requestID: message.requestID, jobs },
+            "*",
+          )
+        },
+      }),
+    })
+  })
+  await page.goto(`/iframe.html?id=chat--task-header-background-agents-420&viewMode=story&globals=${globals}`)
+  const bar = page.locator('[data-component="task-header-agents"]')
+  const stop = bar.getByRole("button", { name: "Stop all (2)", exact: true })
+  await expect(stop).toBeVisible()
+  await expect(bar.locator('[data-slot="task-header-todos-list"]')).toHaveCount(0)
+  await stop.click()
+  await expect(page.getByRole("dialog")).toHaveCount(0)
+  await expect
+    .poll(() => calls)
+    .toEqual(
+      ["first", "second"].map((jobID) => ({
+        type: "cancelBackgroundJob",
+        sessionID: "story-session-chat-001",
+        jobID,
+        requestID: expect.any(String),
+      })),
+    )
+  await expect(bar.getByRole("button", { name: /^Stop all/ })).toHaveCount(0)
+  await bar.locator('[data-slot="task-header-agents-toggle"]').click()
+  await expect(bar.locator('[data-slot="task-header-agent"][data-status="cancelled"]')).toHaveCount(2)
+  await expect(bar.locator('[data-slot="task-header-agent"][data-status="completed"]')).toHaveCount(1)
+})

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

@@ -15,25 +15,28 @@ function client(calls: unknown[], fail = false) {
 }
 
 describe("SessionAbort", () => {
-  it("stops the active owner and current mapped directory", async () => {
-    const calls: unknown[] = []
-    const aborts = new SessionAbort()
-    aborts.observe("session_1", "busy", "/repo")
+  it.each([undefined, "session", "tree"] as const)(
+    "stops the active owner and current mapped directory with %s scope",
+    async (scope) => {
+      const calls: unknown[] = []
+      const aborts = new SessionAbort()
+      aborts.observe("session_1", "busy", "/repo")
 
-    expect(await aborts.stop(client(calls), "session_1", "/repo/worktree")).toBe(true)
-    expect(calls).toEqual([
-      {
-        type: "abort",
-        params: { sessionID: "session_1", directory: "/repo" },
-        opts: { throwOnError: true },
-      },
-      {
-        type: "abort",
-        params: { sessionID: "session_1", directory: "/repo/worktree" },
-        opts: { throwOnError: true },
-      },
-    ])
-  })
+      expect(await aborts.stop(client(calls), "session_1", "/repo/worktree", undefined, scope)).toBe(true)
+      expect(calls).toEqual([
+        {
+          type: "abort",
+          params: { sessionID: "session_1", directory: "/repo", scope },
+          opts: { throwOnError: true },
+        },
+        {
+          type: "abort",
+          params: { sessionID: "session_1", directory: "/repo/worktree", scope },
+          opts: { throwOnError: true },
+        },
+      ])
+    },
+  )
 
   it("forgets an owner when its instance becomes idle", async () => {
     const calls: unknown[] = []

packages/kilo-vscode/tests/unit/kilo-provider-load-messages.test.ts

@@ -2,6 +2,7 @@ import { describe, it, expect, spyOn } from "bun:test"
 import type { SessionStatus } from "@kilocode/sdk/v2/client"
 import * as vscode from "vscode"
 import type { PartUpdate } from "../../src/shared/stream-messages"
+import type { AbortRequest } from "../../webview-ui/src/types/messages/webview-messages"
 
 // vscode mock is provided by the shared preload (tests/setup/vscode-mock.ts)
 const { KiloProvider, unwrapSyncEvent } = await import("../../src/KiloProvider")
@@ -79,7 +80,7 @@ function createClient(options?: {
 }) {
   const calls: { before?: string; limit?: number }[] = []
   const stopped: { sessionID: string; directory?: string }[] = []
-  const aborted: { sessionID: string; directory?: string }[] = []
+  const aborted: { sessionID: string; directory?: string; scope?: "session" | "tree" }[] = []
   const deleted: { sessionID: string; directory?: string }[] = []
   const deletedMessages: Array<{
     sessionID: string
@@ -126,7 +127,7 @@ function createClient(options?: {
         prompted.push(params)
         return { data: undefined }
       },
-      abort: async (params: { sessionID: string; directory?: string }) => {
+      abort: async (params: { sessionID: string; directory?: string; scope?: "session" | "tree" }) => {
         aborted.push(params)
         if (params.directory && options?.abortFailures?.includes(params.directory)) throw new Error("abort failed")
         await options?.abortDeferred?.promise
@@ -251,7 +252,8 @@ type ProviderInternals = {
   seedSessionStatusMap: (reconcile?: boolean) => Promise<void>
   stopCurrentSessionProcesses: (next?: string) => void
   handleEvent: (event: unknown, directory?: string) => void
-  handleAbort: (sid?: string) => Promise<void>
+  setupWebviewMessageHandler: (webview: unknown) => void
+  handleAbort: (sid?: string, scope?: "session" | "tree") => Promise<void>
   resolveSession: (sid?: string, draft?: string, context?: string, dir?: string) => Promise<unknown>
   handleCostAlertResponse: (sid: string, limit: number, response: "continue" | "stop") => Promise<void>
   setMaxCost: (value: unknown) => void
@@ -290,6 +292,32 @@ function mockMaxCost(internal: ProviderInternals, value: number) {
 }
 
 describe("KiloProvider.handleAbort", () => {
+  it.each([undefined, "session", "tree"] as const)(
+    "forwards %s abort scope without extra child or process stops",
+    async (scope) => {
+      const client = createClient()
+      const { provider, internal } = makeProvider(client)
+      const listener = Promise.withResolvers<(message: AbortRequest) => Promise<void>>()
+      internal.setupWebviewMessageHandler({
+        onDidReceiveMessage: (handler: (message: AbortRequest) => Promise<void>) => {
+          listener.resolve(handler)
+          return { dispose: () => {} }
+        },
+      })
+      const receive = await listener.promise
+      status(internal, "busy")
+      status(internal, "busy", "/repo", "child")
+
+      await receive({ type: "abort", sessionID: "s1", scope })
+
+      expect(client.aborted).toEqual([{ sessionID: "s1", directory: "/repo", scope }])
+      expect(client.stopped).toEqual([])
+      expect(internal.sessionStatusMap.get("s1")).toBe("idle")
+      expect(internal.sessionStatusMap.get("child")).toBe("busy")
+      provider.dispose()
+    },
+  )
+
   it("aborts a session whose busy status was seeded without an SSE event", async () => {
     const client = createClient()
     const { internal, sent } = makeProvider(client)

packages/kilo-vscode/tests/unit/prompt-continue.test.ts

@@ -28,6 +28,32 @@ describe("empty prompt continuation", () => {
     expect(continuation({ ...input, messages: [user, { ...assistant, finish: "stop" }] })).toBe(assistant.id)
   })
 
+  it("keeps a stopped turn resumable while background results wait for user input", () => {
+    const job = { ...user, id: "background" }
+    const part: Part = {
+      id: "result",
+      sessionID: "session",
+      messageID: job.id,
+      type: "text",
+      text: "Background task completed",
+      synthetic: true,
+      metadata: { background: true },
+    }
+    const next = {
+      ...input,
+      messages: [user, assistant, job],
+      parts: (id: string) => (id === job.id ? [part] : []),
+    }
+    expect(continuation(next)).toBe(assistant.id)
+    expect(continuation({ ...next, parts: () => [{ ...part, metadata: {} }] })).toBeUndefined()
+    expect(
+      continuation({
+        ...next,
+        messages: [...next.messages, { ...assistant, id: "reply", parentID: job.id, error: undefined, finish: "stop" }],
+      }),
+    ).toBeUndefined()
+  })
+
   it("accepts unfinished assistant and tool-call turns", () => {
     for (const finish of [undefined, "tool-calls"]) {
       expect(continuation({ ...input, messages: [user, { ...assistant, error: undefined, finish }] })).toBe(

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

@@ -99,6 +99,7 @@ export const BackgroundAgents: Component<{ readonly?: boolean }> = (props) => {
     const hidden = session.dismissedBackgroundJobs(id)
     return agents().filter((agent) => showBackgroundAgent(agent, hidden))
   })
+  const running = createMemo(() => visible().filter((agent) => agent.status === "running"))
   const summary = createMemo(() => {
     const running = visible().filter((agent) => agent.status === "running").length
     const total = visible().length
@@ -275,6 +276,18 @@ export const BackgroundAgents: Component<{ readonly?: boolean }> = (props) => {
           >
             <Icon name={open() ? "chevron-up" : "chevron-down"} size="small" />
           </Button>
+          <Show when={!props.readonly && running().length > 0}>
+            <Button
+              icon="stop"
+              variant="ghost"
+              size="small"
+              onClick={(event: MouseEvent) => {
+                for (const agent of running()) cancelAgent(event, agent)
+              }}
+            >
+              {language.t("task.backgroundAgents.stopAll", { count: String(running().length) })}
+            </Button>
+          </Show>
           <Show when={!props.readonly && visible().some((agent) => agent.status !== "running")}>
             <Button
               icon="close-small"

packages/kilo-vscode/webview-ui/src/components/shared/WorkingIndicator.tsx

@@ -89,7 +89,7 @@ export const WorkingIndicator: Component = () => {
   const handleCancelRetry = () => {
     const sid = session.currentSessionID()
     if (sid) {
-      vscode.postMessage({ type: "abort", sessionID: sid })
+      vscode.postMessage({ type: "abort", sessionID: sid, scope: "session" })
     }
   }
 

packages/kilo-vscode/webview-ui/src/context/session-continuation.ts

@@ -12,14 +12,25 @@ export function continuation(input: {
 }) {
   if (!input.id || input.id.startsWith("cloud:") || input.status !== "idle") return
   if (input.submitting || input.blocked || input.loading || input.reverted) return
-  const user = input.messages.findLast((message) => message.role === "user")
-  const assistant = input.messages.at(-1)
+  const answered = new Set(
+    input.messages.flatMap((message) => (message.role === "assistant" ? [message.parentID] : [])),
+  )
+  const messages = input.messages.filter(
+    (message) =>
+      message.role !== "user" ||
+      answered.has(message.id) ||
+      !input
+        .parts(message.id)
+        .some((part) => part.type === "text" && part.synthetic && part.metadata?.background === true),
+  )
+  const user = messages.findLast((message) => message.role === "user")
+  const assistant = messages.at(-1)
   if (!user || !assistant || assistant.role !== "assistant" || assistant.summary === true) return
   if (assistant.parentID !== user.id) return
   if (input.parts(user.id).some((part) => part.type === "compaction")) return
   if (
-    input.messages
-      .slice(input.messages.indexOf(user) + 1)
+    messages
+      .slice(messages.indexOf(user) + 1)
       .some((message) =>
         input
           .parts(message.id)