Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .changeset/retry-permission-transport-drops.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
"kilo-code": patch
---

Retry transient connection drops when approving permissions so auto-approve and manual approvals no longer leave agents waiting.
20 changes: 5 additions & 15 deletions packages/kilo-vscode/src/commands/toggle-auto-approve.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
import * as vscode from "vscode"
import type { Event, KiloClient } from "@kilocode/sdk/v2/client"
import { replyOnce } from "../kilo-provider/handlers/permission-handler"
import { retry } from "../services/cli-backend/retry"
import type { KiloConnectionService } from "../services/cli-backend/connection-service"

/**
Expand Down Expand Up @@ -70,15 +72,11 @@ export function registerToggleAutoApprove(
for (const dir of directories()) {
if (generation !== snapshot) break
try {
const { data: pending } = await client.permission.list({ directory: dir }, { throwOnError: true })
const { data: pending } = await retry(() => client.permission.list({ directory: dir }, { throwOnError: true }))
for (const req of pending) {
if (generation !== snapshot) break
if (req.metadata?.["sandboxEscalation"] === true) continue
await client.permission
.reply({ requestID: req.id, directory: dir, reply: "once" }, { throwOnError: true })
.catch((err) => {
console.error("[Kilo New] toggleAutoApprove: failed to drain pending:", err)
})
await replyOnce(client, req.id, dir, () => generation === snapshot)
}
} catch (err) {
console.error("[Kilo New] toggleAutoApprove: failed to list pending permissions:", err)
Expand All @@ -95,15 +93,7 @@ export function registerToggleAutoApprove(
if (event.properties.metadata?.["sandboxEscalation"] === true) return false
const dir =
directory ?? connectionService.getPermissionDirectory(event.properties.id) ?? resolve(event.properties.sessionID)
return client.permission
.reply({ requestID: event.properties.id, directory: dir, reply: "once" }, { throwOnError: true })
.then(
() => true,
(err) => {
console.error("[Kilo New] toggleAutoApprove: failed to auto-reply:", err)
return false
},
)
return replyOnce(client, event.properties.id, dir)
}

context.subscriptions.push(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

import type { KiloClient, PermissionRequest } from "@kilocode/sdk/v2/client"
import { permissionSettled, respondToPermission } from "@kilocode/sdk/permission"
import { retry } from "../../services/cli-backend/retry"
import { isNotFoundError } from "./not-found"

export type RecoverablePermission = PermissionRequest
Expand All @@ -16,6 +17,8 @@ export type PermissionResponseResult =
| { kind: "stale" }
| { kind: "error" }

const CANCELLED = "permission reply cancelled"

export interface PermissionContext {
readonly client: KiloClient | null
readonly currentSessionId: string | undefined
Expand Down Expand Up @@ -43,6 +46,33 @@ export function recoveryDirs(workspace: string, dirs: ReadonlyMap<string, string
return [...new Set([workspace, ...dirs.values(), ...extra])]
}

/**
* Reply "once" to a permission request, retrying transient transport drops.
* A dropped pooled socket must not strand the request while the agent waits.
* When shouldContinue returns false the reply stops before the next attempt, so
* disabling auto-approve cancels an in-flight retry. Its message is not
* transient on purpose, so the retry helper stops instead of looping.
*/
export async function replyOnce(
client: KiloClient,
requestID: string,
directory: string,
shouldContinue?: () => boolean,
): Promise<boolean> {
try {
await retry(() => {
if (shouldContinue && !shouldContinue()) throw new Error(CANCELLED)
return client.permission.reply({ requestID, directory, reply: "once" }, { throwOnError: true })
})
return true
} catch (error) {
if (!(error instanceof Error && error.message === CANCELLED)) {
console.error("[Kilo New] permission-handler: failed to reply once:", error)
}
return false
}
}

export function recoverablePermissions(
perms: RecoverablePermission[],
tracked: Set<string>,
Expand Down Expand Up @@ -148,7 +178,8 @@ export async function handlePermissionResponse(
* recovered instead of leaving the server blocked indefinitely.
*/
export async function fetchAndSendPendingPermissions(ctx: PermissionContext): Promise<void> {
if (!ctx.client) return
const client = ctx.client
if (!client) return
try {
const dirs = recoveryDirs(ctx.getWorkspaceDirectory(), ctx.sessionDirectories, ctx.extraDirectories?.() ?? [])

Expand All @@ -158,8 +189,12 @@ export async function fetchAndSendPendingPermissions(ctx: PermissionContext): Pr
const valid = new Set<string>()
const pending: Array<{ perm: RecoverablePermission; dir: string }> = []
for (const dir of dirs) {
const { data, error } = await ctx.client.permission.list({ directory: dir })
if (error) {
let data: RecoverablePermission[] | undefined
try {
data = await retry(() => client.permission.list({ directory: dir }, { throwOnError: true })).then(
(result) => result.data,
)
} catch (error) {
console.error(`[Kilo New] KiloProvider: Failed to fetch pending permissions for ${dir}:`, error)
continue
}
Expand Down
110 changes: 109 additions & 1 deletion packages/kilo-vscode/tests/unit/permission-recovery.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,17 @@ import {
handlePermissionResponse,
recoverablePermissions,
recoveryDirs,
replyOnce,
type RecoverablePermission,
type PermissionContext,
} from "../../src/kilo-provider/handlers/permission-handler"
import { KiloConnectionService } from "../../src/services/cli-backend/connection-service"

/** Transient transport failures carry no HTTP status and are safe to retry. */
function terminated() {
return new TypeError("terminated")
}

/** Minimal permission shape returned by the SDK's permission.list(). */
function pending(id: string, sessionID: string, permission = "bash"): RecoverablePermission {
return {
Expand All @@ -35,7 +41,7 @@ function permissionClient(
const dir = args?.directory ?? ""
queries.push(dir)
const error = errors?.list?.[dir]
if (error) return { data: undefined, error }
if (error) throw error
return { data: permsPerDir[dir] ?? [] }
},
saveAlwaysRules: async (args: unknown) => {
Expand Down Expand Up @@ -328,6 +334,75 @@ describe("handlePermissionResponse", () => {
})
})

describe("replyOnce", () => {
it("retries a transient transport drop", async () => {
let calls = 0
const client = {
permission: {
reply: async () => {
calls += 1
if (calls === 1) throw terminated()
return { data: true }
},
},
}
const log = spyOn(console, "error").mockImplementation(() => {})
try {
expect(await replyOnce(client as never, "p1", "/workspace")).toBe(true)
} finally {
log.mockRestore()
}
expect(calls).toBe(2)
})

it("does not retry a decisive not-found reply", async () => {
let calls = 0
const client = {
permission: {
reply: async () => {
calls += 1
throw new Error("Permission request not found: p1", { cause: { status: 404 } })
},
},
}
const log = spyOn(console, "error").mockImplementation(() => {})
try {
expect(await replyOnce(client as never, "p1", "/workspace")).toBe(false)
} finally {
log.mockRestore()
}
expect(calls).toBe(1)
})

it("stops retrying when the caller cancels", async () => {
Comment thread
marius-kilocode marked this conversation as resolved.
let calls = 0
let checks = 0
let allowed = true
const client = {
permission: {
reply: async () => {
calls += 1
allowed = false
throw terminated()
},
},
}
const log = spyOn(console, "error").mockImplementation(() => {})
try {
expect(
await replyOnce(client as never, "p1", "/workspace", () => {
checks += 1
return allowed
}),
).toBe(false)
} finally {
log.mockRestore()
}
expect(calls).toBe(1)
expect(checks).toBe(2)
})
})

describe("recoverablePermissions", () => {
it("filters out untracked permissions", () => {
const seen = new Set<string>()
Expand Down Expand Up @@ -419,6 +494,39 @@ describe("fetchAndSendPendingPermissions", () => {
expect(permDirs.get("worktree-pending")).toBe("/workspace/.kilo/worktrees/failing")
})

it("retries a transient list failure during recovery", async () => {
const messages: unknown[] = []
let calls = 0
const client = {
permission: {
list: async () => {
calls += 1
if (calls === 1) throw terminated()
return { data: [pending("p1", "s1")] }
},
},
}
const fake: PermissionContext = {
client: client as unknown as PermissionContext["client"],
currentSessionId: undefined,
trackedSessionIds: new Set(["s1"]),
sessionDirectories: new Map(),
extraDirectories: () => [],
postMessage: (msg) => messages.push(msg),
getWorkspaceDirectory: () => "/workspace",
recordPermissionDirectory: () => {},
getPermissionDirectory: () => undefined,
clearPermissionDirectory: () => {},
getPermissionRevision: () => 0,
prunePermissionDirectories: () => {},
}

await fetchAndSendPendingPermissions(fake)

expect(calls).toBe(2)
expect(messages).toHaveLength(1)
})

it("deduplicates directories", async () => {
const dirs = new Map([
["s1", "/workspace/.kilo/worktrees/alpha"],
Expand Down
86 changes: 69 additions & 17 deletions packages/sdk/js/src/kilocode/permission.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,50 @@ type Decision = {
message?: string
}

const ATTEMPTS = 3

// Mirrors the transient classifier in
// packages/kilo-vscode/src/services/cli-backend/retry.ts. The SDK cannot import
// that module, so keep the two lists in sync by hand, including the exact
// "terminated" match for undici pooled-connection drops. Only transport-level
// failures are retried. A status-less client error such as a response parse
// failure or a server-version mismatch is definitive and must not be replayed.
// Re-sending an identical rule set after a lost response is tolerable because
// duplicate patterns do not change the decision.
const TRANSIENT = [
"load failed",
"network connection was lost",
"network request failed",
"failed to fetch",
"fetch failed",
"econnreset",
"econnrefused",
"etimedout",
"socket hang up",
]

const TRANSIENT_EXACT = ["terminated"]

function transport(error: unknown): boolean {
if (!error) return false
const message = String(error instanceof Error ? error.message : error)
.toLowerCase()
.trim()
if (TRANSIENT_EXACT.includes(message)) return true
return TRANSIENT.some((entry) => message.includes(entry))
}

async function send<T>(fn: () => Promise<T>, budget: () => number): Promise<T> {
for (let attempt = 0; ; attempt++) {
try {
return await fn()
} catch (error) {
if (attempt >= ATTEMPTS - 1 || !transport(error) || budget() <= 0) throw error
Comment thread
WebReflection marked this conversation as resolved.
await new Promise((resolve) => setTimeout(resolve, Math.min(100 * 2 ** attempt, budget())))
}
}
}

/**
* Send a permission decision with one bounded wait across both requests.
* A timeout aborts only the client wait, so callers must reconcile an aborted
Expand All @@ -25,25 +69,33 @@ export async function respondToPermission(
const saved = input.approvedAlways.length > 0 || input.deniedAlways.length > 0
try {
if (saved) {
await client.permission.saveAlwaysRules(
{
requestID: input.requestID,
directory: input.directory,
approvedAlways: input.approvedAlways,
deniedAlways: input.deniedAlways,
},
{ throwOnError: true, signal: AbortSignal.timeout(budget()) },
await send(
() =>
client.permission.saveAlwaysRules(
{
requestID: input.requestID,
directory: input.directory,
approvedAlways: input.approvedAlways,
deniedAlways: input.deniedAlways,
},
{ throwOnError: true, signal: AbortSignal.timeout(budget()) },
),
budget,
)
}
await client.permission.reply(
{
requestID: input.requestID,
directory: input.directory,
reply: input.reply,
interactive: true,
...(input.message ? { message: input.message } : {}),
},
{ throwOnError: true, signal: AbortSignal.timeout(budget()) },
await send(
() =>
client.permission.reply(
{
requestID: input.requestID,
directory: input.directory,
reply: input.reply,
interactive: true,
...(input.message ? { message: input.message } : {}),
},
{ throwOnError: true, signal: AbortSignal.timeout(budget()) },
),
budget,
)
return { saved }
} catch (error) {
Expand Down
Loading
Loading