fix(vscode): prevent duplicate permission prompts in Agent Manager

.changeset/permission-prompt-race.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+Prevent duplicate permission prompts and response races in Agent Manager.

packages/kilo-vscode/src/KiloProvider.ts

@@ -2510,6 +2510,7 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
       this.activeAlerts.delete(sessionID)
       this.postMessage({ type: "sessionCostAlertResolved", sessionID: sessionID, limit: deletedAlertLimit })
     }
+    this.connectionService.clearPermissionSession(sessionID)
     this.connectionService.pruneSession(sessionID)
     if (this.currentSession?.id === sessionID) {
       this.contextSessionID = undefined
@@ -4533,11 +4534,15 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
       extraDirectories: this.opts.worktreeDirectories,
       postMessage: (msg) => this.postMessage(msg),
       getWorkspaceDirectory: (sid) => this.getWorkspaceDirectory(sid),
-      recordPermissionDirectory: (id, dir) => this.connectionService.recordPermissionDirectory(id, dir),
+      recordPermissionDirectory: (id, dir, sid) => this.connectionService.recordPermissionDirectory(id, dir, sid),
       getPermissionDirectory: (id) => this.connectionService.getPermissionDirectory(id),
+      getPermissionSession: (id) => this.connectionService.getPermissionSession(id),
       clearPermissionDirectory: (id) => this.connectionService.clearPermissionDirectory(id),
       getPermissionRevision: () => this.connectionService.getPermissionRevision(),
       prunePermissionDirectories: (active, dirs) => this.connectionService.prunePermissionDirectories(active, dirs),
+      runPermissionResponse: (id, sid, action) => this.connectionService.runPermissionResponse(id, sid, action),
+      isPermissionResponseClaimed: (id) => this.connectionService.isPermissionResponseClaimed(id),
+      clearPermissionResponse: (id) => this.connectionService.clearPermissionResponse(id),
     }
   }
 

packages/kilo-vscode/src/kilo-provider/handlers/permission-handler.ts

@@ -9,6 +9,11 @@ import type { KiloClient, PermissionRequest } from "@kilocode/sdk/v2/client"
 import { isNotFoundError } from "./not-found"
 
 export type RecoverablePermission = PermissionRequest
+export type PermissionResponse = "once" | "always" | "reject"
+export type PermissionResponseResult =
+  | { kind: "resolved"; sessionID: string; response: PermissionResponse }
+  | { kind: "stale" }
+  | { kind: "error" }
 
 export interface PermissionContext {
   readonly client: KiloClient | null
@@ -18,21 +23,35 @@ export interface PermissionContext {
   readonly extraDirectories?: () => string[]
   postMessage(msg: unknown): void
   getWorkspaceDirectory(sessionId?: string): string
-  recordPermissionDirectory(requestID: string, directory: string): void
+  recordPermissionDirectory(requestID: string, directory: string, sessionID?: string): void
   getPermissionDirectory(requestID: string): string | undefined
+  getPermissionSession?(requestID: string): string | undefined
   clearPermissionDirectory(requestID: string): void
   getPermissionRevision(): number
   prunePermissionDirectories(active: Set<string>, dirs?: Set<string>): void
+  runPermissionResponse?: (
+    requestID: string,
+    sessionID: string,
+    action: () => Promise<PermissionResponseResult>,
+  ) => Promise<PermissionResponseResult>
+  isPermissionResponseClaimed?: (requestID: string) => boolean
+  clearPermissionResponse?: (requestID: string) => void
 }
 
 export function recoveryDirs(workspace: string, dirs: ReadonlyMap<string, string>, extra: string[] = []) {
   return [...new Set([workspace, ...dirs.values(), ...extra])]
 }
 
-export function recoverablePermissions(perms: RecoverablePermission[], tracked: Set<string>, seen: Set<string>) {
+export function recoverablePermissions(
+  perms: RecoverablePermission[],
+  tracked: Set<string>,
+  seen: Set<string>,
+  claimed: (requestID: string) => boolean = () => false,
+) {
   return perms.filter((perm) => {
     if (seen.has(perm.id)) return false
     seen.add(perm.id)
+    if (claimed(perm.id)) return false
     return tracked.has(perm.sessionID)
   })
 }
@@ -45,68 +64,93 @@ export async function handlePermissionResponse(
   ctx: PermissionContext,
   permissionId: string,
   sessionID: string,
-  response: "once" | "always" | "reject",
+  response: PermissionResponse,
   approvedAlways: string[],
   deniedAlways: string[],
 ): Promise<void> {
-  if (!ctx.client) {
+  const client = ctx.client
+  if (!client) {
     ctx.postMessage({ type: "permissionError", permissionID: permissionId })
     return
   }
 
-  const target = sessionID || ctx.currentSessionId
-  if (!target) {
-    console.error("[Kilo New] KiloProvider: No sessionID for permission response")
+  const dir = ctx.getPermissionDirectory(permissionId)
+  const target = ctx.getPermissionSession?.(permissionId) ?? sessionID
+  const claimed = ctx.isPermissionResponseClaimed?.(permissionId) ?? false
+  if (!target || (!dir && !claimed) || (ctx.getPermissionSession?.(permissionId) && target !== sessionID)) {
+    console.error("[Kilo New] KiloProvider: Unknown permission route")
     ctx.postMessage({ type: "permissionError", permissionID: permissionId })
     return
   }
 
-  const dir = ctx.getPermissionDirectory(permissionId) ?? ctx.getWorkspaceDirectory(target)
+  const run =
+    ctx.runPermissionResponse ??
+    ((_requestID: string, _sessionID: string, action: () => Promise<PermissionResponseResult>) => action())
+  const action = async (): Promise<PermissionResponseResult> => {
+    if (!dir) return { kind: "error" }
 
-  const staleCleanup = () => {
-    ctx.clearPermissionDirectory(permissionId)
-    ctx.postMessage({ type: "permissionError", permissionID: permissionId, stale: true })
-    void fetchAndSendPendingPermissions(ctx)
-  }
+    if (approvedAlways.length > 0 || deniedAlways.length > 0) {
+      const saveResult = await client.permission
+        .saveAlwaysRules(
+          {
+            requestID: permissionId,
+            directory: dir,
+            approvedAlways,
+            deniedAlways,
+          },
+          { throwOnError: true },
+        )
+        .then(() => "ok" as const)
+        .catch((error: unknown) => {
+          if (isNotFoundError(error)) return "stale" as const
+          console.error("[Kilo New] KiloProvider: Failed to save always-rules:", error)
+          return "error" as const
+        })
+      if (saveResult === "stale") {
+        ctx.clearPermissionDirectory(permissionId)
+        void fetchAndSendPendingPermissions(ctx)
+        return { kind: "stale" }
+      }
+      if (saveResult === "error") return { kind: "error" }
+    }
 
-  if (approvedAlways.length > 0 || deniedAlways.length > 0) {
-    const saveResult = await ctx.client.permission
-      .saveAlwaysRules(
-        {
-          requestID: permissionId,
-          directory: dir,
-          approvedAlways,
-          deniedAlways,
-        },
-        { throwOnError: true },
-      )
+    const replyResult = await client.permission
+      .reply({ requestID: permissionId, reply: response, directory: dir, interactive: true }, { throwOnError: true })
       .then(() => "ok" as const)
       .catch((error: unknown) => {
         if (isNotFoundError(error)) return "stale" as const
-        console.error("[Kilo New] KiloProvider: Failed to save always-rules:", error)
-        ctx.postMessage({ type: "permissionError", permissionID: permissionId })
+        console.error("[Kilo New] KiloProvider: Failed to respond to permission:", error)
         return "error" as const
       })
-    if (saveResult === "stale") {
-      staleCleanup()
-      return
+    if (replyResult === "stale") {
+      ctx.clearPermissionDirectory(permissionId)
+      void fetchAndSendPendingPermissions(ctx)
+      return { kind: "stale" }
     }
-    if (saveResult === "error") return
+    if (replyResult !== "ok") return { kind: "error" }
+    ctx.clearPermissionDirectory(permissionId)
+    return { kind: "resolved", sessionID: target, response }
   }
 
-  const replyResult = await ctx.client.permission
-    .reply({ requestID: permissionId, reply: response, directory: dir, interactive: true }, { throwOnError: true })
-    .then(() => "ok" as const)
-    .catch((error: unknown) => {
-      if (isNotFoundError(error)) return "stale" as const
-      console.error("[Kilo New] KiloProvider: Failed to respond to permission:", error)
-      ctx.postMessage({ type: "permissionError", permissionID: permissionId })
-      return "error" as const
-    })
-  if (replyResult === "stale") staleCleanup()
-  if (replyResult !== "ok") return
-  ctx.clearPermissionDirectory(permissionId)
-  ctx.postMessage({ type: "permissionResolved", permissionID: permissionId, sessionID: target, response })
+  const result = await run(permissionId, target, action).catch((error: unknown) => {
+    console.error("[Kilo New] KiloProvider: Failed to process permission response:", error)
+    return { kind: "error" } as const
+  })
+  if (result.kind === "error") {
+    ctx.clearPermissionResponse?.(permissionId)
+    ctx.postMessage({ type: "permissionError", permissionID: permissionId })
+    return
+  }
+  if (result.kind === "stale") {
+    ctx.postMessage({ type: "permissionError", permissionID: permissionId, stale: true })
+    return
+  }
+  ctx.postMessage({
+    type: "permissionResolved",
+    permissionID: permissionId,
+    sessionID: result.sessionID,
+    response: result.response,
+  })
 }
 
 /**
@@ -133,11 +177,17 @@ export async function fetchAndSendPendingPermissions(ctx: PermissionContext): Pr
         }
         valid.add(dir)
         if (!data) continue
-        for (const perm of recoverablePermissions(data, ctx.trackedSessionIds, seen)) pending.push({ perm, dir })
+        for (const perm of recoverablePermissions(
+          data,
+          ctx.trackedSessionIds,
+          seen,
+          (id) => ctx.isPermissionResponseClaimed?.(id) ?? false,
+        ))
+          pending.push({ perm, dir })
       }
       if (ctx.getPermissionRevision() !== revision) continue
       for (const { perm, dir } of pending) {
-        ctx.recordPermissionDirectory(perm.id, dir)
+        ctx.recordPermissionDirectory(perm.id, dir, perm.sessionID)
         ctx.postMessage({
           type: "permissionRequest",
           permission: {