fix(cli): wait for background continuations before headless exit

.changeset/headless-session-drain.md

@@ -0,0 +1,6 @@
+---
+"@kilocode/cli": patch
+"@kilocode/sdk": patch
+---
+
+Wait for background subagents and their parent continuations before completing headless runs, including sessions resumed from another directory. Report cancelled headless runs as failures.

packages/opencode/src/cli/cmd/run.ts

@@ -279,12 +279,12 @@ export const RunCommand = effectCmd({
     const { InstanceRef } = yield* Effect.promise(() => import("@/effect/instance-ref"))
     const { ServerAuth } = yield* Effect.promise(() => import("@/server/auth"))
     // kilocode_change start - lazy Kilo implementations (see top-of-file note)
-    const { createKiloClient } = yield* Effect.promise(() => import("@kilocode/sdk/v2"))
     const { buildRunMessage } = yield* Effect.promise(() => import("@/kilocode/cli/cmd/run-message"))
     const { importCloudSession, validateCloudFork, reportCloudImportError } = yield* Effect.promise(
       () => import("@/kilocode/cloud-session"),
     )
     const { KiloRunAuto } = yield* Effect.promise(() => import("@/kilocode/cli/run-auto"))
+    const { KiloRunDrain } = yield* Effect.promise(() => import("@/kilocode/cli/run-drain"))
     const { KiloHeadless } = yield* Effect.promise(() => import("@/kilocode/permission/headless"))
     const { KiloRun, KiloRunDaemon } = yield* Effect.promise(() => import("@/kilocode/cli/cmd/run"))
     // kilocode_change end
@@ -368,7 +368,8 @@ export const RunCommand = effectCmd({
         ? ServerAuth.headers({ password: args.password, username: args.username })
         : undefined
       const attachSDK = (dir?: string) => {
-        return createKiloClient({
+        return KiloRunDrain.client({
+          // kilocode_change
           baseUrl: args.attach!,
           directory: dir,
           headers: attachHeaders,
@@ -769,6 +770,7 @@ export const RunCommand = effectCmd({
         const sessionID = sess.id
         // kilocode_change start - track Task children; plain headless runs deny subagent asks instead of hanging (#11903)
         const tracked = KiloRunAuto.create(sessionID) // kilocode_change - named to avoid shadowing the `auto` flag
+        const drain = KiloRunDrain.create(sessionID)
         if (!args.attach && !args.auto && !skipPermissions) KiloHeadless.mark(sessionID) // kilocode_change - --yolo skips too
         // kilocode_change end
 
@@ -800,6 +802,7 @@ export const RunCommand = effectCmd({
 
           // kilocode_change start - revert to upstream: consume native events without normalizing sync copies
           for await (const event of events.stream) {
+            if (drain.event(event)) break // kilocode_change
             // kilocode_change end
 
             if (
@@ -903,14 +906,6 @@ export const RunCommand = effectCmd({
             }
             // kilocode_change end
 
-            if (
-              event.type === "session.status" &&
-              event.properties.sessionID === sessionID &&
-              event.properties.status.type === "idle"
-            ) {
-              break
-            }
-
             // kilocode_change start - non-interactive runs dismiss suggestions so they don't block
             if (event.type === "suggestion.shown") {
               const suggestion = event.properties
@@ -923,6 +918,7 @@ export const RunCommand = effectCmd({
 
             if (event.type === "permission.asked") {
               const permission = event.properties
+              if (!KiloRunAuto.allowed(tracked, permission.sessionID)) continue // kilocode_change
               // kilocode_change start - skill shell batches need an interactive human decision. The server ignores
               // non-interactive approvals, so headless runs must reject explicitly rather than leave them pending.
               if (permission.metadata?.["skillShell"] === true || permission.metadata?.["sandboxEscalation"] === true) {
@@ -991,7 +987,7 @@ export const RunCommand = effectCmd({
             // kilocode_change start - bounded network retry handling
             if (event.type === "session.network.asked") {
               const request = event.properties
-              if (request.sessionID !== sessionID) continue
+              if (!KiloRunAuto.allowed(tracked, request.sessionID)) continue
               retries++
               if (retries > MAX_RETRIES) {
                 UI.println(
@@ -1002,7 +998,7 @@ export const RunCommand = effectCmd({
                 continue
               }
               const delay = Math.min(5000 * Math.pow(2, retries - 1), 60000)
-              await new Promise((resolve) => setTimeout(resolve, delay))
+              await drain.pause(delay)
               await client.network.reply({ requestID: request.id })
             }
             // kilocode_change end
@@ -1017,8 +1013,8 @@ export const RunCommand = effectCmd({
           // kilocode_change end
           return error
         }
-        const cwd = args.attach ? (directory ?? sess.directory ?? (await current(sdk))) : (directory ?? root)
-        const client = args.attach ? attachSDK(cwd) : sdk
+        const cwd = sess.directory ?? directory ?? (await current(sdk)) // kilocode_change
+        const client = KiloRunDrain.scope(sdk, cwd, interactive ? undefined : drain.signal) // kilocode_change
         // kilocode_change start - classify deferred attach commands in the session directory
         const builtin = deferred ? await KiloRun.resolveBuiltin(client, args.command, cwd) : initial
         if (deferred) {
@@ -1032,63 +1028,64 @@ export const RunCommand = effectCmd({
 
         await share(client, sessionID)
 
+        // kilocode_change start
         if (!interactive) {
-          const events = await client.event.subscribe()
-          const completed = loop(client, events).catch((e) => {
-            console.error(e)
-            process.exitCode = 1
+          const events = await client.event.subscribe(undefined, {
+            signal: drain.signal,
+            sseMaxRetryAttempts: 1,
+            onSseError: (error) => drain.end(error),
           })
-          async function finish() {
-            if (args.attach) return
-            const error = await completed
-            if (error) process.exitCode = 1
-          }
-
-          // kilocode_change start - handle built-in session commands
-          if (builtin) {
-            const result = await KiloRun.runBuiltin(client, sessionID, builtin, args.model, sess.model, cwd)
-            if (result.error) {
-              if (!emit("error", { error: result.error })) UI.error(formatRunError(result.error))
-              process.exitCode = 1
-            }
-            return
-          }
-          // kilocode_change end
-
-          if (args.command) {
-            const result = await client.session.command({
-              sessionID,
-              agent,
-              model: args.model,
-              command: args.command,
-              arguments: message,
-              variant: args.variant,
-            })
+          const completed = loop(client, events).then(
+            (error) => {
+              drain.end()
+              return error
+            },
+            (error) => {
+              drain.end(error)
+              return undefined
+            },
+          )
+          try {
+            await drain.race(KiloRunDrain.check(client, drain.signal))
+            await drain.ready()
+            const result = await drain.race(
+              builtin
+                ? KiloRun.runBuiltin(client, sessionID, builtin, args.model, sess.model, cwd)
+                : args.command
+                  ? client.session.command({
+                      sessionID,
+                      agent,
+                      model: args.model,
+                      command: args.command,
+                      arguments: message,
+                      variant: args.variant,
+                    })
+                  : client.session.prompt({
+                      sessionID,
+                      agent,
+                      model: pick(args.model),
+                      variant: args.variant,
+                      parts: [...files, { type: "text", text: message }],
+                    }),
+            )
             if (result.error) {
               if (!emit("error", { error: result.error })) UI.error(formatRunError(result.error))
               process.exitCode = 1
-              return
             }
-            await finish()
-            return
-          }
-
-          const model = pick(args.model)
-          const result = await client.session.prompt({
-            sessionID,
-            agent,
-            model,
-            variant: args.variant,
-            parts: [...files, { type: "text", text: message }],
-          })
-          if (result.error) {
-            if (!emit("error", { error: result.error })) UI.error(formatRunError(result.error))
+            await drain.wait(client, cwd)
+            if (await completed) process.exitCode = 1
+          } catch (error) {
+            const text = error instanceof Error ? error.message : String(error)
+            if (!emit("error", { error: text })) UI.error(text)
             process.exitCode = 1
-            return
+          } finally {
+            drain.close()
+            await completed
+            await KiloRunDrain.flush()
           }
-          await finish()
           return
         }
+        // kilocode_change end
 
         const model = pick(args.model)
         const { runInteractiveMode } = await import("./run/runtime")
@@ -1169,7 +1166,8 @@ export const RunCommand = effectCmd({
         if (auth) headers.set("Authorization", auth)
         return Server.Default().app.fetch(new Request(request, { headers }))
       }) as typeof globalThis.fetch
-      const sdk = createKiloClient({
+      const sdk = KiloRunDrain.client({
+        // kilocode_change
         baseUrl: "http://kilo.internal",
         fetch: fetchFn,
         directory,

packages/opencode/src/effect/app-runtime.ts

@@ -55,6 +55,7 @@ import { BackgroundJob } from "@/background/job"
 import { RuntimeFlags } from "@/effect/runtime-flags"
 // kilocode_change start
 import { Notebook } from "@/kilocode/notebook/service"
+import { SessionDrain } from "@/kilocode/session/drain"
 import { AgentManager } from "@/kilocode/agent-manager/service"
 // kilocode_change end
 import { EventV2Bridge } from "@/event-v2-bridge"
@@ -70,7 +71,14 @@ import { Pty } from "@opencode-ai/core/pty" // kilocode_change
 
 // kilocode_change start - retain Kilo runtime services in the upstream node graph
 const memory = LayerNode.make({ service: MemoryService.Service, layer: MemoryService.layer, deps: [] })
-const kilo = LayerNode.group([Credential.node, ModelCache.node, AgentManager.node, Notebook.node, memory])
+const kilo = LayerNode.group([
+  Credential.node,
+  ModelCache.node,
+  AgentManager.node,
+  Notebook.node,
+  SessionDrain.node,
+  memory,
+])
 // kilocode_change end
 
 export const AppLayer = AppNodeBuilderV1.build(

packages/opencode/src/effect/instance-registry.ts

@@ -1,4 +1,5 @@
 import type { WorkspaceV2 } from "@opencode-ai/core/workspace" // kilocode_change
+import { dispose } from "@/kilocode/effect/instance-registry" // kilocode_change
 
 const disposers = new Set<(directory: string, workspaceID?: WorkspaceV2.ID) => Promise<void>>() // kilocode_change
 
@@ -13,6 +14,8 @@ export function registerDisposer(
 }
 
 export async function disposeInstance(directory: string, workspaceID?: WorkspaceV2.ID) {
-  await Promise.allSettled([...disposers].map((disposer) => disposer(directory, workspaceID)))
+  await dispose(directory, workspaceID, () =>
+    Promise.allSettled([...disposers].map((disposer) => disposer(directory, workspaceID))),
+  )
 }
 // kilocode_change end

packages/opencode/src/effect/runner.ts

@@ -43,6 +43,7 @@ export const make = <A, E = never>(
     onIdle?: Effect.Effect<void>
     onBusy?: Effect.Effect<void>
     onInterrupt?: Effect.Effect<A, E>
+    lease?: Effect.Effect<() => void> // kilocode_change
   },
 ): Runner<A, E> => {
   const ref = SynchronizedRef.makeUnsafe<State<A, E>>({ _tag: "Idle" })
@@ -82,11 +83,22 @@ export const make = <A, E = never>(
     ).pipe(Effect.flatten)
 
   // kilocode_change start - do not let work publish busy before the Running state is committed
-  const startRun = (work: Effect.Effect<A, E>, done: Deferred.Deferred<A, E | Cancelled>) => {
+  const current = (fiber: Fiber.Fiber<unknown, unknown>) => {
+    const st = state()
+    return st._tag === "Running" && st.run.fiber === fiber
+  }
+
+  const startRun = (
+    work: Effect.Effect<A, E>,
+    done: Deferred.Deferred<A, E | Cancelled>,
+    record: (started: KiloRunner.Started) => void,
+  ) => {
     const id = next()
     return KiloRunner.start({
       work,
       scope,
+      lease: opts?.lease,
+      record,
       finish: (exit) => finishRun(id, done, exit),
       handle: (fiber) => ({ id, done, fiber }) satisfies RunHandle<A, E>,
     })
@@ -95,17 +107,19 @@ export const make = <A, E = never>(
 
   // kilocode_change start - open work only after the Running state is committed
   const finishShell = (id: number) =>
-    SynchronizedRef.modifyEffect(
-      ref,
-      Effect.fnUntraced(function* (st) {
-        if (st._tag === "Shell" && st.shell.id === id) {
-          return [idle, { _tag: "Idle" }] as const
-        }
-        if (st._tag === "ShellThenRun" && st.shell.id === id) {
-          return yield* KiloRunner.commit(startRun(st.run.work, st.run.done), Effect.void)
-        }
-        return [Effect.void, st] as const
-      }),
+    KiloRunner.guard(current, (record) =>
+      SynchronizedRef.modifyEffect(
+        ref,
+        Effect.fnUntraced(function* (st) {
+          if (st._tag === "Shell" && st.shell.id === id) {
+            return [idle, { _tag: "Idle" }] as const
+          }
+          if (st._tag === "ShellThenRun" && st.shell.id === id) {
+            return yield* KiloRunner.commit(startRun(st.run.work, st.run.done, record), Effect.void)
+          }
+          return [Effect.void, st] as const
+        }),
+      ),
     ).pipe(Effect.flatten)
   // kilocode_change end
 
@@ -118,28 +132,30 @@ export const make = <A, E = never>(
 
   // kilocode_change start - open work only after the Running state is committed
   const ensureRunning = (work: Effect.Effect<A, E>, valid?: () => boolean) =>
-    SynchronizedRef.modifyEffect(
-      ref,
-      Effect.fnUntraced(function* (st) {
-        if (valid && !valid()) return [onInterrupt ?? Effect.die(new Cancelled()), st] as const
-        switch (st._tag) {
-          case "Running":
-          case "ShellThenRun":
-            return [awaitDone(st.run.done), st] as const
-          case "Shell": {
-            const run = {
-              id: next(),
-              done: yield* Deferred.make<A, E | Cancelled>(),
-              work,
-            } satisfies PendingHandle<A, E>
-            return [awaitDone(run.done), { _tag: "ShellThenRun", shell: st.shell, run }] as const
-          }
-          case "Idle": {
-            const done = yield* Deferred.make<A, E | Cancelled>()
-            return yield* KiloRunner.commit(startRun(work, done), awaitDone(done))
+    KiloRunner.guard(current, (record) =>
+      SynchronizedRef.modifyEffect(
+        ref,
+        Effect.fnUntraced(function* (st) {
+          if (valid && !valid()) return [onInterrupt ?? Effect.die(new Cancelled()), st] as const
+          switch (st._tag) {
+            case "Running":
+            case "ShellThenRun":
+              return [awaitDone(st.run.done), st] as const
+            case "Shell": {
+              const run = {
+                id: next(),
+                done: yield* Deferred.make<A, E | Cancelled>(),
+                work,
+              } satisfies PendingHandle<A, E>
+              return [awaitDone(run.done), { _tag: "ShellThenRun", shell: st.shell, run }] as const
+            }
+            case "Idle": {
+              const done = yield* Deferred.make<A, E | Cancelled>()
+              return yield* KiloRunner.commit(startRun(work, done, record), awaitDone(done))
+            }
           }
-        }
-      }),
+        }),
+      ),
     ).pipe(Effect.flatten)
 
   const startShell = (work: Effect.Effect<A, E>, ready?: Latch.Latch): Effect.Effect<A, E | Busy> =>
@@ -154,7 +170,12 @@ export const make = <A, E = never>(
         yield* onBusy
         const id = next()
         const cancelled = yield* Deferred.make<void>()
-        const fiber = yield* work.pipe(Effect.ensuring(finishShell(id)), Effect.forkChild)
+        // kilocode_change start
+        const fiber = yield* KiloRunner.fork(
+          work.pipe(Effect.ensuring(finishShell(id)), Effect.forkChild({ uninterruptible: false })),
+          opts?.lease,
+        )
+        // kilocode_change end
         const shell = { id, cancelled, ready, fiber } satisfies ShellHandle<A, E>
         return [
           Effect.gen(function* () {