fix(vscode): stop moved active sessions

.changeset/stop-moved-agent-sessions.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+Keep the Stop button working when an active session moves between the workspace and an Agent Manager worktree.

packages/kilo-vscode/src/KiloProvider.ts

@@ -72,7 +72,7 @@ import { fetchMessagePage, MESSAGE_PAGE_LIMIT } from "./kilo-provider/message-pa
 import { childID } from "./kilo-provider/task-session"
 import { VisibleTaskStreams } from "./kilo-provider/visible-task-streams"
 import { handleNetworkEvent, clearNetworkWaits } from "./kilo-provider/network"
-import { abortSession } from "./kilo-provider/abort"
+import { SessionAbort } from "./kilo-provider/abort"
 import {
   buildAutocompleteSettingsMessage,
   validAutocompleteSetting,
@@ -310,6 +310,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
   private syncedChildSessions: Set<string> = new Set()
   private sessionStatusMap = new Map<string, SessionStatus["type"]>() // Latest status used for destructive config warnings.
   private sessionDirectories = new Map<string, string>() // Per-session directory overrides, such as Agent Manager worktrees.
+  private readonly aborts = new SessionAbort()
   private projectID: string | undefined // Current workspace project ID used to filter sessions.
   private loadMessagesAbort: AbortController | null = null // Current load request cancellation.
   private lastReconciledAt = new Map<string, number>() // Per-session focus-mode reconcile timestamp.
@@ -657,10 +658,12 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
    * When set, all operations for this session use this directory instead of the workspace root.
    */
   public setSessionDirectory(sessionId: string, directory: string): void {
+    this.aborts.preserve(sessionId, this.sessionStatusMap.get(sessionId), this.getWorkspaceDirectory(sessionId))
     this.sessionDirectories.set(sessionId, directory)
   }
 
   public clearSessionDirectory(sessionId: string): void {
+    this.aborts.preserve(sessionId, this.sessionStatusMap.get(sessionId), this.getWorkspaceDirectory(sessionId))
     this.sessionDirectories.delete(sessionId)
   }
 
@@ -1755,6 +1758,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
       this.visibleTaskStreams.delete(sessionID)
       this.syncedChildSessions.delete(sessionID)
       this.sessionDirectories.delete(sessionID)
+      this.aborts.delete(sessionID)
       this.lastReconciledAt.delete(sessionID)
       this.connectionService.pruneSession(sessionID)
       if (this.currentSession?.id === sessionID) {
@@ -2754,24 +2758,11 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
   }
 
   private async handleAbort(sessionID?: string): Promise<void> {
-    if (!this.client) {
-      return
-    }
-
-    const targetSessionID = sessionID || this.currentSession?.id
-    if (!targetSessionID) {
-      return
-    }
-
-    try {
-      await abortSession({
-        client: this.client,
-        sessionID: targetSessionID,
-        dir: this.getWorkspaceDirectory(targetSessionID),
-      })
-    } catch (error) {
-      console.error("[Kilo New] KiloProvider: Failed to abort session:", error)
-    }
+    const sid = sessionID || this.currentSession?.id
+    if (!this.client || !sid || !(await this.aborts.stop(this.client, sid, this.getWorkspaceDirectory(sid)))) return
+    this.sessionStatusMap.set(sid, "idle")
+    this.streams.flush(sid)
+    this.postMessage({ type: "sessionStatus", sessionID: sid, status: "idle" })
   }
 
   private async handleRevertSession(sessionID: string, messageID: string, partID?: string): Promise<void> {
@@ -3156,6 +3147,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
     if (event.type === "session.status") {
       const sid = event.properties.sessionID
       this.sessionStatusMap.set(sid, event.properties.status.type)
+      this.aborts.observe(sid, event.properties.status.type, directory)
       const msg = mapSSEEventToWebviewMessage(event, sid)
       if (msg) {
         this.streams.flush(sid)
@@ -3188,6 +3180,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
     if (event.type === "server.instance.disposed") {
       const props = event.properties as Record<string, unknown> | null
       const dir = typeof props?.directory === "string" ? props.directory : undefined
+      if (dir) for (const sid of this.aborts.dispose(dir)) this.sessionStatusMap.set(sid, "idle")
       if (dir && !sameDirectory(dir, this.getWorkspaceDirectory())) return
       void this.reloadAfterAuthChange()
       return
@@ -3600,6 +3593,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
     this.trackedSessionIds.clear()
     this.syncedChildSessions.clear()
     this.sessionDirectories.clear()
+    this.aborts.clear()
     this.sessionStatusMap.clear()
     this.ignoreController?.dispose()
     this.chatAutocomplete?.dispose()

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

@@ -1,4 +1,69 @@
-import type { KiloClient } from "@kilocode/sdk/v2/client"
+import type { KiloClient, SessionStatus } from "@kilocode/sdk/v2/client"
+import { sameDirectory } from "../kilo-provider-utils"
+
+export class SessionAbort {
+  private active = new Map<string, Set<string>>()
+
+  observe(sessionID: string, status: SessionStatus["type"], dir?: string) {
+    if (!dir) return
+    const dirs = this.active.get(sessionID)
+    if (status === "idle") {
+      if (!dirs) return
+      for (const entry of dirs) {
+        if (sameDirectory(entry, dir)) dirs.delete(entry)
+      }
+      if (dirs.size === 0) this.active.delete(sessionID)
+      return
+    }
+    if (!dirs) {
+      this.active.set(sessionID, new Set([dir]))
+      return
+    }
+    if (![...dirs].some((entry) => sameDirectory(entry, dir))) dirs.add(dir)
+  }
+
+  preserve(sessionID: string, status: SessionStatus["type"] | undefined, dir: string) {
+    if (!status || status === "idle" || this.active.has(sessionID)) return
+    this.observe(sessionID, status, dir)
+  }
+
+  async stop(client: KiloClient, sessionID: string, fallback: string) {
+    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) => abortSession({ client, sessionID, dir })))
+    const failures = results.flatMap((result, index) =>
+      result.status === "rejected" ? [{ dir: dirs[index], error: result.reason }] : [],
+    )
+    if (failures.length > 0) {
+      console.error("[Kilo New] KiloProvider: Failed to abort session in one or more directories:", failures)
+      return false
+    }
+    if (known) this.active.delete(sessionID)
+    return known
+  }
+
+  dispose(dir: string) {
+    const idle: string[] = []
+    for (const [sessionID, dirs] of this.active) {
+      for (const entry of dirs) {
+        if (sameDirectory(entry, dir)) dirs.delete(entry)
+      }
+      if (dirs.size > 0) continue
+      this.active.delete(sessionID)
+      idle.push(sessionID)
+    }
+    return idle
+  }
+
+  delete(sessionID: string) {
+    this.active.delete(sessionID)
+  }
+
+  clear() {
+    this.active.clear()
+  }
+}
 
 export async function abortSession(input: { client: KiloClient; sessionID: string; dir: string }) {
   await input.client.session.abort({ sessionID: input.sessionID, directory: input.dir }, { throwOnError: true })

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

@@ -1,6 +1,6 @@
 import { describe, expect, it } from "bun:test"
 import type { KiloClient } from "@kilocode/sdk/v2/client"
-import { abortSession } from "../../src/kilo-provider/abort"
+import { abortSession, SessionAbort } from "../../src/kilo-provider/abort"
 
 function client(calls: unknown[], fail = false) {
   return {
@@ -14,6 +14,53 @@ function client(calls: unknown[], fail = false) {
   } as unknown as KiloClient
 }
 
+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")
+
+    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 },
+      },
+    ])
+  })
+
+  it("forgets an owner when its instance becomes idle", async () => {
+    const calls: unknown[] = []
+    const aborts = new SessionAbort()
+    aborts.observe("session_1", "busy", "/repo")
+    aborts.observe("session_1", "idle", "/repo")
+
+    expect(await aborts.stop(client(calls), "session_1", "/repo/worktree")).toBe(false)
+    expect(calls).toEqual([
+      {
+        type: "abort",
+        params: { sessionID: "session_1", directory: "/repo/worktree" },
+        opts: { throwOnError: true },
+      },
+    ])
+  })
+
+  it("deduplicates equivalent directory paths", async () => {
+    const calls: unknown[] = []
+    const aborts = new SessionAbort()
+    aborts.observe("session_1", "busy", "/repo/worktree")
+
+    expect(await aborts.stop(client(calls), "session_1", "/repo/worktree/.")).toBe(true)
+    expect(calls).toHaveLength(1)
+  })
+})
+
 describe("abortSession", () => {
   it("calls session.abort with the session id and directory", async () => {
     const calls: unknown[] = []

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

@@ -1,4 +1,4 @@
-import { describe, it, expect } from "bun:test"
+import { describe, it, expect, spyOn } from "bun:test"
 import type { PartUpdate } from "../../src/shared/stream-messages"
 
 // vscode mock is provided by the shared preload (tests/setup/vscode-mock.ts)
@@ -44,19 +44,27 @@ function createClient(options?: {
   deleteDeferred?: Deferred<unknown>
   sessionData?: unknown
   sessionGet?: (params: { sessionID: string; directory?: string }) => Promise<{ data: unknown }>
+  abortFailures?: string[]
 }) {
   const calls: { before?: string; limit?: number }[] = []
   const stopped: { sessionID: string; directory?: string }[] = []
+  const aborted: { sessionID: string; directory?: string }[] = []
   return {
     calls,
     stopped,
+    aborted,
     session: {
       list: async () => ({ data: [] }),
       get: async (params: { sessionID: string; directory?: string }) => {
         if (options?.sessionGet) return options.sessionGet(params)
         return { data: options?.sessionData ?? null }
       },
       status: async () => ({ data: {} }),
+      abort: async (params: { sessionID: string; directory?: string }) => {
+        aborted.push(params)
+        if (params.directory && options?.abortFailures?.includes(params.directory)) throw new Error("abort failed")
+        return { data: true }
+      },
       messages: async (params: { before?: string; limit?: number }) => {
         calls.push({ before: params.before, limit: params.limit })
         if (options?.messagesDeferred) return options.messagesDeferred.promise
@@ -118,7 +126,8 @@ type ProviderInternals = {
   trackedSessionIds: Set<string>
   streams: { push: (msg: PartUpdate) => void }
   stopCurrentSessionProcesses: (next?: string) => void
-  handleEvent: (event: unknown) => void
+  handleEvent: (event: unknown, directory?: string) => void
+  handleAbort: (sid?: string) => Promise<void>
   handleLoadMessages: (sid: string, opts?: { mode?: string; before?: string; limit?: number }) => Promise<void>
   handleDeleteSession: (sid: string) => Promise<void>
 }
@@ -137,6 +146,70 @@ function makeProvider(client: ReturnType<typeof createClient>) {
   return { provider, internal, sent }
 }
 
+describe("KiloProvider.handleAbort", () => {
+  it("aborts the original owner after a running session moves to a worktree", async () => {
+    const client = createClient()
+    const { provider, internal, sent } = makeProvider(client)
+    internal.handleEvent(
+      {
+        type: "session.status",
+        properties: { sessionID: "s1", status: { type: "busy" } },
+      },
+      "/repo",
+    )
+    provider.setSessionDirectory("s1", "/repo/worktree")
+
+    await internal.handleAbort("s1")
+
+    expect(client.aborted).toEqual([
+      { sessionID: "s1", directory: "/repo" },
+      { sessionID: "s1", directory: "/repo/worktree" },
+    ])
+    expect(sent.at(-1)).toMatchObject({ type: "sessionStatus", sessionID: "s1", status: "idle" })
+  })
+
+  it("preserves the original owner when the status event lacks a directory", async () => {
+    const client = createClient()
+    const { provider, internal } = makeProvider(client)
+    internal.handleEvent({
+      type: "session.status",
+      properties: { sessionID: "s1", status: { type: "busy" } },
+    })
+    provider.setSessionDirectory("s1", "/repo/worktree")
+
+    await internal.handleAbort("s1")
+
+    expect(client.aborted).toEqual([
+      { sessionID: "s1", directory: "/repo" },
+      { sessionID: "s1", directory: "/repo/worktree" },
+    ])
+  })
+
+  it("attempts every owner and stays busy when one abort fails", async () => {
+    const error = spyOn(console, "error").mockImplementation(() => {})
+    const client = createClient({ abortFailures: ["/repo"] })
+    const { provider, internal, sent } = makeProvider(client)
+    internal.handleEvent(
+      {
+        type: "session.status",
+        properties: { sessionID: "s1", status: { type: "busy" } },
+      },
+      "/repo",
+    )
+    provider.setSessionDirectory("s1", "/repo/worktree")
+
+    await internal.handleAbort("s1")
+
+    expect(client.aborted).toEqual([
+      { sessionID: "s1", directory: "/repo" },
+      { sessionID: "s1", directory: "/repo/worktree" },
+    ])
+    expect(sent.at(-1)).toMatchObject({ type: "sessionStatus", sessionID: "s1", status: "busy" })
+    expect(error).toHaveBeenCalledTimes(1)
+    error.mockRestore()
+  })
+})
+
 describe("KiloProvider.handleLoadMessages / focus mode freshness", () => {
   it("stops background processes for the previous session when switching sessions", async () => {
     const client = createClient({