fix(vscode): abort moved active sessions

.changeset/abort-moved-sessions.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+Stop active Agent Manager turns after moving or continuing their sessions in another worktree.

packages/kilo-vscode/src/KiloProvider.ts

@@ -71,7 +71,12 @@ 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 {
+  abortSession,
+  resolveAbortDirectories,
+  updateActiveSessionDirectory,
+  type ActiveSessionDirectories,
+} from "./kilo-provider/abort"
 import {
   buildAutocompleteSettingsMessage,
   validAutocompleteSetting,
@@ -281,6 +286,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 activeSessionDirectories: ActiveSessionDirectories = new Map()
   private permissionDirectories = new Map<string, string>()
   private projectID: string | undefined // Current workspace project ID used to filter sessions.
   private loadMessagesAbort: AbortController | null = null // Current load request cancellation.
@@ -628,10 +634,28 @@ 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 {
+    const status = this.sessionStatusMap.get(sessionId)
+    if (status && status !== "idle" && !this.activeSessionDirectories.has(sessionId)) {
+      updateActiveSessionDirectory({
+        active: this.activeSessionDirectories,
+        sessionID: sessionId,
+        status,
+        dir: this.getWorkspaceDirectory(sessionId),
+      })
+    }
     this.sessionDirectories.set(sessionId, directory)
   }
 
   public clearSessionDirectory(sessionId: string): void {
+    const status = this.sessionStatusMap.get(sessionId)
+    if (status && status !== "idle" && !this.activeSessionDirectories.has(sessionId)) {
+      updateActiveSessionDirectory({
+        active: this.activeSessionDirectories,
+        sessionID: sessionId,
+        status,
+        dir: this.getWorkspaceDirectory(sessionId),
+      })
+    }
     this.sessionDirectories.delete(sessionId)
   }
 
@@ -1708,6 +1732,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
       this.visibleTaskStreams.delete(sessionID)
       this.syncedChildSessions.delete(sessionID)
       this.sessionDirectories.delete(sessionID)
+      this.activeSessionDirectories.delete(sessionID)
       this.lastReconciledAt.delete(sessionID)
       this.connectionService.pruneSession(sessionID)
       if (this.currentSession?.id === sessionID) {
@@ -2705,23 +2730,33 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
   }
 
   private async handleAbort(sessionID?: string): Promise<void> {
-    if (!this.client) {
-      return
-    }
+    const client = this.client
+    if (!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 dirs = resolveAbortDirectories(
+      this.activeSessionDirectories,
+      targetSessionID,
+      this.getWorkspaceDirectory(targetSessionID),
+    )
+    const results = await Promise.allSettled(
+      dirs.map((dir) =>
+        abortSession({
+          client,
+          sessionID: targetSessionID,
+          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)
     }
   }
 
@@ -3096,6 +3131,12 @@ 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)
+      updateActiveSessionDirectory({
+        active: this.activeSessionDirectories,
+        sessionID: sid,
+        status: event.properties.status.type,
+        dir: directory,
+      })
       const msg = mapSSEEventToWebviewMessage(event, sid)
       if (msg) {
         this.streams.flush(sid)
@@ -3128,6 +3169,17 @@ 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 sessionID of [...this.activeSessionDirectories.keys()]) {
+          updateActiveSessionDirectory({
+            active: this.activeSessionDirectories,
+            sessionID,
+            status: "idle",
+            dir,
+          })
+          if (!this.activeSessionDirectories.has(sessionID)) this.sessionStatusMap.set(sessionID, "idle")
+        }
+      }
       if (dir && !sameDirectory(dir, this.getWorkspaceDirectory())) return
       void this.reloadAfterAuthChange()
       return
@@ -3527,6 +3579,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
     this.trackedSessionIds.clear()
     this.syncedChildSessions.clear()
     this.sessionDirectories.clear()
+    this.activeSessionDirectories.clear()
     this.permissionDirectories.clear()
     this.sessionStatusMap.clear()
     this.ignoreController?.dispose()

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

@@ -1,4 +1,37 @@
-import type { KiloClient } from "@kilocode/sdk/v2/client"
+import type { KiloClient, SessionStatus } from "@kilocode/sdk/v2/client"
+import { sameDirectory } from "../kilo-provider-utils"
+
+export type ActiveSessionDirectories = Map<string, Set<string>>
+
+export function updateActiveSessionDirectory(input: {
+  active: ActiveSessionDirectories
+  sessionID: string
+  status: SessionStatus["type"]
+  dir?: string
+}) {
+  const target = input.dir
+  if (!target) return
+  const dirs = input.active.get(input.sessionID)
+  if (input.status === "idle") {
+    if (!dirs) return
+    for (const dir of dirs) {
+      if (sameDirectory(dir, target)) dirs.delete(dir)
+    }
+    if (dirs.size === 0) input.active.delete(input.sessionID)
+    return
+  }
+  if (!dirs) {
+    input.active.set(input.sessionID, new Set([target]))
+    return
+  }
+  if (![...dirs].some((dir) => sameDirectory(dir, target))) dirs.add(target)
+}
+
+export function resolveAbortDirectories(active: ActiveSessionDirectories, sessionID: string, fallback: string) {
+  const dirs = [...(active.get(sessionID) ?? [])]
+  if (!dirs.some((dir) => sameDirectory(dir, fallback))) dirs.push(fallback)
+  return dirs
+}
 
 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,11 @@
 import { describe, expect, it } from "bun:test"
 import type { KiloClient } from "@kilocode/sdk/v2/client"
-import { abortSession } from "../../src/kilo-provider/abort"
+import {
+  abortSession,
+  resolveAbortDirectories,
+  updateActiveSessionDirectory,
+  type ActiveSessionDirectories,
+} from "../../src/kilo-provider/abort"
 
 function client(calls: unknown[], fail = false) {
   return {
@@ -14,6 +19,58 @@ function client(calls: unknown[], fail = false) {
   } as unknown as KiloClient
 }
 
+describe("active session directories", () => {
+  function update(active: ActiveSessionDirectories, status: "busy" | "retry" | "idle", dir?: string) {
+    updateActiveSessionDirectory({ active, sessionID: "session_1", status, dir })
+  }
+
+  it("includes the active owner and current session directory", () => {
+    const active: ActiveSessionDirectories = new Map()
+    update(active, "busy", "/repo/source")
+
+    expect(resolveAbortDirectories(active, "session_1", "/repo/worktree")).toEqual(["/repo/source", "/repo/worktree"])
+  })
+
+  it("falls back to the current session directory after the active turn becomes idle", () => {
+    const active: ActiveSessionDirectories = new Map()
+    update(active, "busy", "/repo/source")
+    update(active, "idle", "/repo/source")
+
+    expect(resolveAbortDirectories(active, "session_1", "/repo/worktree")).toEqual(["/repo/worktree"])
+  })
+
+  it("retains active directories when an unrelated instance reports idle", () => {
+    const active: ActiveSessionDirectories = new Map()
+    update(active, "retry", "/repo/source")
+    update(active, "idle", "/repo/worktree")
+    update(active, "idle")
+
+    expect(resolveAbortDirectories(active, "session_1", "/repo/worktree")).toEqual(["/repo/source", "/repo/worktree"])
+  })
+
+  it("tracks concurrent instances and ignores delayed idle events from the old one", () => {
+    const active: ActiveSessionDirectories = new Map()
+    update(active, "busy", "/repo/source")
+    update(active, "busy", "/repo/worktree")
+
+    expect(resolveAbortDirectories(active, "session_1", "/repo/fallback")).toEqual([
+      "/repo/source",
+      "/repo/worktree",
+      "/repo/fallback",
+    ])
+
+    update(active, "idle", "/repo/source")
+    expect(resolveAbortDirectories(active, "session_1", "/repo/fallback")).toEqual(["/repo/worktree", "/repo/fallback"])
+  })
+
+  it("deduplicates the current directory when it is already active", () => {
+    const active: ActiveSessionDirectories = new Map()
+    update(active, "busy", "/repo/worktree")
+
+    expect(resolveAbortDirectories(active, "session_1", "/repo/worktree/.")).toEqual(["/repo/worktree"])
+  })
+})
+
 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"
 
 // vscode mock is provided by the shared preload (tests/setup/vscode-mock.ts)
 const { KiloProvider } = await import("../../src/KiloProvider")
@@ -43,12 +43,15 @@ 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 }) => {
@@ -56,6 +59,11 @@ function createClient(options?: {
         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
@@ -114,9 +122,11 @@ type ProviderInternals = {
   currentSession: { id: string; directory?: string } | null
   contextSessionID: string | undefined
   sessionDirectories: Map<string, string>
+  activeSessionDirectories: Map<string, Set<string>>
   trackedSessionIds: Set<string>
   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>
 }
@@ -135,6 +145,108 @@ function makeProvider(client: ReturnType<typeof createClient>) {
   return { provider, internal, sent }
 }
 
+describe("KiloProvider.handleAbort", () => {
+  it("keeps the active owner when a running session moves to a worktree", 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("preserves a directory-qualified owner instead of inferring the current directory", async () => {
+    const client = createClient()
+    const { provider, internal } = makeProvider(client)
+    internal.handleEvent(
+      {
+        type: "session.status",
+        properties: { sessionID: "s1", status: { type: "busy" } },
+      },
+      "/repo/source",
+    )
+    provider.setSessionDirectory("s1", "/repo/worktree")
+
+    await internal.handleAbort("s1")
+
+    expect(client.aborted).toEqual([
+      { sessionID: "s1", directory: "/repo/source" },
+      { sessionID: "s1", directory: "/repo/worktree" },
+    ])
+  })
+
+  it("falls back to the current directory after the old owner becomes idle", 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")
+    internal.handleEvent(
+      {
+        type: "session.status",
+        properties: { sessionID: "s1", status: { type: "idle" } },
+      },
+      "/repo",
+    )
+
+    await internal.handleAbort("s1")
+
+    expect(client.aborted).toEqual([{ sessionID: "s1", directory: "/repo/worktree" }])
+  })
+
+  it("deduplicates the current directory when it is also an active owner", async () => {
+    const client = createClient()
+    const { provider, internal } = makeProvider(client)
+    internal.trackedSessionIds.add("s1")
+    provider.setSessionDirectory("s1", "/repo/worktree")
+    internal.handleEvent(
+      {
+        type: "session.status",
+        properties: { sessionID: "s1", status: { type: "busy" } },
+      },
+      "/repo/worktree",
+    )
+
+    await internal.handleAbort("s1")
+
+    expect(client.aborted).toEqual([{ sessionID: "s1", directory: "/repo/worktree" }])
+  })
+
+  it("attempts every owner when one abort request fails", async () => {
+    const error = spyOn(console, "error").mockImplementation(() => {})
+    const client = createClient({ abortFailures: ["/repo"] })
+    const { provider, internal } = makeProvider(client)
+    internal.trackedSessionIds.add("s1")
+    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(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({