fix(cli): resume subagents after session fork

.changeset/forked-subagent-resume.md

@@ -0,0 +1,5 @@
+---
+"@kilocode/cli": patch
+---
+
+Allow subagent tasks to be resumed after their parent session is forked.

packages/opencode/src/kilocode/session/fork.ts

@@ -1,41 +1,180 @@
+import { Effect } from "effect"
+import type { Session } from "@/session/session"
 import { MessageV2 } from "@/session/message-v2"
+import { MessageID, PartID, SessionID } from "@/session/schema"
 import { KiloPartLifecycle } from "./part-lifecycle"
 
 const task = "task"
-const stale = /^[ \t]*task_id:[^\r\n]*(?:(?:\r?\n){1,2}|$)/m
+type Ops = Pick<Session.Interface, "get" | "messages" | "create" | "updateMessage" | "updatePart">
 
-// Prepare a source part for a forked transcript copy: drop transient parts (returns undefined) and detach
-// task calls into historical results. The caller assigns fresh ids and publishes via Session.updatePart.
+// Keep terminal task references so the fork can give them private child sessions.
+// In-flight jobs cannot be copied safely, so those remain historical errors.
 export function prepareForkedPart(part: MessageV2.Part): MessageV2.Part | undefined {
   if (KiloPartLifecycle.transient(part)) return undefined
-  return structuredClone(detachPart(part))
+  if (
+    part.type === "tool" &&
+    part.tool === task &&
+    (part.state.status === "pending" || part.state.status === "running")
+  ) {
+    return structuredClone(detachPart(part))
+  }
+  return structuredClone(part)
 }
 
-function metadata(value: Record<string, unknown> | undefined) {
+function childID(part: MessageV2.Part) {
+  if (part.type !== "tool" || part.tool !== task) return undefined
+  const state = part.state
+  const metadata = state.status === "pending" ? undefined : state.metadata
+  const values = [
+    metadata?.sessionId,
+    metadata?.sessionID,
+    part.metadata?.sessionId,
+    part.metadata?.sessionID,
+    state.input.task_id,
+  ]
+  return values.find((value): value is string => typeof value === "string")
+}
+
+function mapRecord(value: Record<string, unknown> | undefined, map: Map<string, SessionID>, keys: string[]) {
   if (!value) return value
   const copy = { ...value }
-  delete copy.sessionId
-  delete copy.sessionID
+  for (const key of keys) {
+    const id = copy[key]
+    if (typeof id !== "string") continue
+    const replacement = map.get(id)
+    if (replacement) copy[key] = replacement
+  }
   return copy
 }
 
-function input(value: Record<string, unknown>) {
-  const copy = { ...value }
-  delete copy.task_id
-  return copy
+function output(value: string, map: Map<string, SessionID>) {
+  return [...map].reduce(
+    (text, [source, target]) =>
+      source === target
+        ? text
+        : text
+            .replaceAll(source, target)
+            .replaceAll(`task id="${source}"`, `task id="${target}"`)
+            .replaceAll(`task_id="${source}"`, `task_id="${target}"`)
+            .replaceAll(`task_id: ${source}`, `task_id: ${target}`),
+    value,
+  )
+}
+
+function resumeHint(sessionID: string) {
+  return [
+    `This subagent session can be resumed: call the task tool again with task_id="${sessionID}"`,
+    "and a prompt describing how to continue or recover. Its prior context is preserved.",
+  ].join(" ")
+}
+
+function remapPart(part: MessageV2.Part, map: Map<string, SessionID>) {
+  if (part.type === "text") {
+    const text = output(part.text, map)
+    return text === part.text ? part : { ...part, text }
+  }
+  if (part.type !== "tool" || part.tool !== task) return part
+  const next = structuredClone(part)
+  next.metadata = mapRecord(next.metadata, map, ["sessionId", "sessionID", "parentSessionId", "parentSessionID"])
+  const input = mapRecord(next.state.input, map, ["task_id"])
+  if (input) next.state.input = input
+  if (next.state.status !== "pending") {
+    const metadata = mapRecord(next.state.metadata, map, [
+      "sessionId",
+      "sessionID",
+      "parentSessionId",
+      "parentSessionID",
+    ])
+    if (metadata) next.state.metadata = metadata
+  }
+  if (next.state.status === "completed") next.state.output = output(next.state.output, map)
+  if (next.state.status === "error") next.state.error = output(next.state.error, map)
+  return next
+}
+
+function copy(input: { source: Session.Info; parentID: SessionID; ops: Ops }) {
+  return Effect.gen(function* () {
+    const target = yield* input.ops.create({
+      parentID: input.parentID,
+      title: input.source.title,
+      agent: input.source.agent,
+      model: input.source.model,
+      metadata: structuredClone(input.source.metadata),
+      permission: input.source.permission ? [...input.source.permission] : undefined,
+      workspaceID: input.source.workspaceID,
+    })
+    const msgs = yield* input.ops.messages({ sessionID: input.source.id })
+    const ids = new Map<string, MessageID>()
+
+    for (const msg of msgs) {
+      const id = MessageID.ascending()
+      ids.set(msg.info.id, id)
+      const parentID = msg.info.role === "assistant" ? ids.get(msg.info.parentID) : undefined
+      const cloned = yield* input.ops.updateMessage({
+        ...msg.info,
+        id,
+        sessionID: target.id,
+        ...(msg.info.role === "assistant" && { cost: 0 }),
+        ...(parentID && { parentID }),
+      })
+
+      for (const part of msg.parts) {
+        const prepared = prepareForkedPart(part)
+        if (!prepared) continue
+        const next: MessageV2.Part = {
+          ...prepared,
+          id: PartID.ascending(),
+          messageID: cloned.id,
+          sessionID: target.id,
+          ...(prepared.type === "step-finish" && { cost: 0 }),
+        }
+        if (next.type === "compaction" && next.tail_start_id) next.tail_start_id = ids.get(next.tail_start_id)
+        yield* input.ops.updatePart(next)
+      }
+    }
+
+    return target
+  })
+}
+
+export function remapChildren(input: {
+  sessionID: SessionID
+  ops: Ops
+  remapped?: Map<string, SessionID>
+}): Effect.Effect<void, Session.NotFound> {
+  return Effect.gen(function* () {
+    const map = input.remapped ?? new Map<string, SessionID>()
+    const msgs = yield* input.ops.messages({ sessionID: input.sessionID })
+    const refs = msgs.flatMap((msg) =>
+      msg.parts.flatMap((part) => {
+        const child = childID(part)
+        return child ? [{ part, child }] : []
+      }),
+    )
+
+    for (const ref of refs) {
+      if (map.has(ref.child)) continue
+      const source = yield* input.ops.get(SessionID.make(ref.child)).pipe(Effect.orElseSucceed(() => undefined))
+      if (!source) continue
+      const target = yield* copy({ source, parentID: input.sessionID, ops: input.ops })
+      map.set(ref.child, target.id)
+      yield* remapChildren({ sessionID: target.id, ops: input.ops, remapped: map })
+    }
+
+    for (const msg of msgs) {
+      for (const part of msg.parts) {
+        const next = remapPart(part, map)
+        if (next !== part) yield* input.ops.updatePart(next)
+      }
+    }
+  })
 }
 
-/**
- * Turns copied task calls into detached historical results.
- *
- * Child sessions are execution state, not conversation context. Their final
- * result is already embedded in the parent task part, so a fork keeps that
- * result while dropping references that could resume, stream, or route prompts
- * to a child owned by the source session.
- */
 function detachPart(part: MessageV2.Part): MessageV2.Part {
   if (part.type !== "tool" || part.tool !== task) return part
 
+  const child = childID(part)
+  const hint = child ? `\n${resumeHint(child)}` : ""
   const top = metadata(part.metadata)
   const state = part.state
   if (state.status === "pending") {
@@ -46,46 +185,30 @@ function detachPart(part: MessageV2.Part): MessageV2.Part {
       state: {
         status: "error",
         input: input(state.input),
-        error: "Task was still pending when this session was forked.",
+        error: `Task was still pending when this session was forked.${hint}`,
         time: { start: now, end: now },
       },
     }
   }
 
-  if (state.status === "running") {
-    return {
-      ...part,
-      metadata: top,
-      state: {
-        status: "error",
-        input: input(state.input),
-        error: "Task was still running when this session was forked.",
-        metadata: metadata(state.metadata),
-        time: { start: state.time.start, end: Date.now() },
-      },
-    }
-  }
-
-  if (state.status === "error") {
-    return {
-      ...part,
-      metadata: top,
-      state: {
-        ...state,
-        input: input(state.input),
-        metadata: metadata(state.metadata),
-      },
-    }
-  }
-
   return {
     ...part,
     metadata: top,
     state: {
-      ...state,
+      status: "error",
       input: input(state.input),
-      output: state.output.replace(stale, ""),
-      metadata: metadata(state.metadata) ?? {},
+      error: `Task was still running when this session was forked.${hint}`,
+      metadata: metadata(state.metadata),
+      time: { start: state.time.start, end: Date.now() },
     },
   }
 }
+
+function metadata(value: Record<string, unknown> | undefined) {
+  if (!value) return value
+  return { ...value }
+}
+
+function input(value: Record<string, unknown>) {
+  return { ...value }
+}

packages/opencode/src/kilocode/session/index.ts

@@ -1,4 +1,4 @@
-import { prepareForkedPart as _prepareForkedPart } from "./fork"
+import { prepareForkedPart as _prepareForkedPart, remapChildren as _remapChildren } from "./fork"
 import z from "zod"
 import { Cause, Effect, Schema } from "effect"
 import { Bus } from "@/bus"
@@ -437,6 +437,7 @@ export namespace KiloSession {
   }
 
   export const prepareForkedPart = _prepareForkedPart
+  export const remapChildren = _remapChildren
 }
 
 export { kiloSessionFork } from "./fork-command"

packages/opencode/src/session/session.ts

@@ -906,6 +906,13 @@ export const layer: Layer.Layer<
       }
       // kilocode_change - preserve imported/cumulative diffs when forking (self-contained Storage runtime keeps this shared file off the legacy Storage layer)
       yield* carryForkDiff(input.sessionID, session.id)
+      // kilocode_change start - fork terminal task children under the new parent and remap their references
+      yield* KiloSession.remapChildren({
+        sessionID: session.id,
+        remapped: new Map([[input.sessionID, session.id]]),
+        ops: { get, messages, create, updateMessage, updatePart },
+      })
+      // kilocode_change end
       return session
     })
 

packages/opencode/test/kilocode/session-fork-remap.test.ts

@@ -49,15 +49,12 @@ afterAll(async () => {
 })
 
 const sessions = {
-  create: (input?: Parameters<Session.Interface["create"]>[0]) =>
-    runtime.runPromise((svc) => svc.create(input)),
+  create: (input?: Parameters<Session.Interface["create"]>[0]) => runtime.runPromise((svc) => svc.create(input)),
+  get: (id: SessionID) => runtime.runPromise((svc) => svc.get(id)),
   list: () => runtime.runPromise((svc) => svc.list()),
-  messages: (input: Parameters<Session.Interface["messages"]>[0]) =>
-    runtime.runPromise((svc) => svc.messages(input)),
-  updateMessage: <T extends MessageV2.Info>(msg: T) =>
-    runtime.runPromise((svc) => svc.updateMessage(msg)),
-  updatePart: <T extends MessageV2.Part>(part: T) =>
-    runtime.runPromise((svc) => svc.updatePart(part)),
+  messages: (input: Parameters<Session.Interface["messages"]>[0]) => runtime.runPromise((svc) => svc.messages(input)),
+  updateMessage: <T extends MessageV2.Info>(msg: T) => runtime.runPromise((svc) => svc.updateMessage(msg)),
+  updatePart: <T extends MessageV2.Part>(part: T) => runtime.runPromise((svc) => svc.updatePart(part)),
 }
 
 afterEach(async () => {
@@ -190,9 +187,9 @@ describe("Session.fork cost accounting", () => {
   )
 })
 
-describe("Session.fork task detachment", () => {
+describe("Session.fork task children", () => {
   test(
-    "keeps completed task outcomes without cloning child sessions",
+    "clones completed task children under the forked parent",
     async () => {
       await using tmp = await tmpdir({ git: true })
       await instance({
@@ -212,6 +209,13 @@ describe("Session.fork task detachment", () => {
           const user = await userMsg(parent.id)
           const assistant = await asstMsg(parent.id, user)
           await sessions.updatePart(taskPart({ messageID: assistant, sessionID: parent.id, childSessionID: child.id }))
+          await sessions.updatePart({
+            id: PartID.ascending(),
+            messageID: assistant,
+            sessionID: parent.id,
+            type: "text",
+            text: `Subagent task ID: ${child.id}`,
+          } as MessageV2.TextPart)
           const before = await sessions.list()
 
           const server = HttpRouter.toWebHandler(HttpApiApp.routes, { disableLogger: true })
@@ -220,23 +224,36 @@ describe("Session.fork task detachment", () => {
             directory: tmp.path,
             fetch: ((request: Request) => server.handler(request, HttpApiApp.context)) as unknown as typeof fetch,
           })
-          const { data: forked } = await client.session.fork(
-            { sessionID: parent.id, directory: tmp.path },
-            { throwOnError: true },
-          ).finally(() => server.dispose())
+          const { data: forked } = await client.session
+            .fork({ sessionID: parent.id, directory: tmp.path }, { throwOnError: true })
+            .finally(() => server.dispose())
 
           const after = await sessions.list()
-          expect(after).toHaveLength(before.length + 1)
+          expect(after).toHaveLength(before.length + 2)
 
           const msgs = await sessions.messages({ sessionID: SessionID.make(forked.id) })
           const tool = msgs.flatMap((msg) => msg.parts).find((part) => part.type === "tool") as MessageV2.ToolPart
           expect(tool.state.status).toBe("completed")
           if (tool.state.status !== "completed") throw new Error("expected completed task")
-          expect(tool.metadata).toEqual({ trace: "keep" })
-          expect(tool.state.metadata).toEqual({ model: { modelID: "test", providerID: "test" } })
-          expect(tool.state.input.task_id).toBeUndefined()
+          const clonedID = tool.state.input.task_id
+          expect(clonedID).not.toBe(child.id)
+          expect(tool.metadata).toEqual({ sessionId: clonedID, trace: "keep" })
+          expect(tool.state.metadata).toEqual({
+            sessionId: clonedID,
+            model: { modelID: "test", providerID: "test" },
+          })
+          if (typeof clonedID !== "string") throw new Error("expected a cloned task ID")
           expect(tool.state.output).toBe(
-            "Background task completed: test task\r\n<task_result>\r\nchild outcome\r\n</task_result>",
+            `Background task completed: test task\r\n\ttask_id: ${clonedID} (for resuming to continue this task if needed)\r\n\r\n<task_result>\r\nchild outcome\r\n</task_result>`,
+          )
+
+          const clone = await sessions.get(SessionID.descending(clonedID))
+          expect(clone.parentID).toBe(SessionID.descending(forked.id))
+          expect((await sessions.messages({ sessionID: clone.id }))[0]?.parts).toContainEqual(
+            expect.objectContaining({ text: "child message content" }),
+          )
+          expect(msgs.flatMap((msg) => msg.parts)).toContainEqual(
+            expect.objectContaining({ text: `Subagent task ID: ${clonedID}` }),
           )
 
           const source = await sessions.messages({ sessionID: parent.id })
@@ -252,7 +269,7 @@ describe("Session.fork task detachment", () => {
   )
 
   test(
-    "turns copied running tasks into terminal historical errors",
+    "turns copied running tasks into resumable historical errors",
     async () => {
       await using tmp = await tmpdir({ git: true })
       await instance({
@@ -278,15 +295,22 @@ describe("Session.fork task detachment", () => {
             },
           } as MessageV2.ToolPart)
 
+          const before = await sessions.list()
           const forked = await Session.fork({ sessionID: parent.id })
+          const after = await sessions.list()
           const msgs = await sessions.messages({ sessionID: forked.id })
           const tool = msgs.flatMap((msg) => msg.parts).find((part) => part.type === "tool") as MessageV2.ToolPart
           expect(tool.state.status).toBe("error")
           if (tool.state.status !== "error") throw new Error("expected detached task error")
           expect(tool.state.error).toContain("still running")
-          expect(tool.state.input.task_id).toBeUndefined()
-          expect(tool.state.metadata).toEqual({ variant: "high" })
-          expect(tool.metadata).toEqual({})
+          const clonedID = tool.state.input.task_id
+          if (typeof clonedID !== "string") throw new Error("expected a cloned task ID")
+          expect(clonedID).not.toBe(child.id)
+          expect(tool.state.error).toContain(`task_id="${clonedID}"`)
+          expect(tool.state.metadata).toEqual({ sessionId: clonedID, variant: "high" })
+          expect(tool.metadata).toEqual({ sessionId: clonedID })
+          expect(after).toHaveLength(before.length + 2)
+          expect((await sessions.get(SessionID.descending(clonedID))).parentID).toBe(forked.id)
         },
       })
     },
@@ -294,7 +318,7 @@ describe("Session.fork task detachment", () => {
   )
 
   test(
-    "detaches pending and errored task references",
+    "detaches in-flight tasks and remaps errored task references",
     async () => {
       await using tmp = await tmpdir({ git: true })
       await instance({
@@ -329,7 +353,7 @@ describe("Session.fork task detachment", () => {
             state: {
               status: "error",
               input: { task_id: child.id },
-              error: "original error",
+              error: `original error; task_id="${child.id}"`,
               metadata: { sessionID: child.id, detail: "keep" },
               time: { start: Date.now(), end: Date.now() },
             },
@@ -344,15 +368,20 @@ describe("Session.fork task detachment", () => {
           expect(pending?.state.status).toBe("error")
           if (!pending || pending.state.status !== "error") throw new Error("expected detached pending task")
           expect(pending.state.error).toContain("still pending")
-          expect(pending.state.input.task_id).toBeUndefined()
-          expect(pending.metadata).toEqual({})
+          const pendingID = pending.state.input.task_id
+          if (typeof pendingID !== "string") throw new Error("expected a cloned pending task ID")
+          expect(pendingID).not.toBe(child.id)
+          expect(pending.state.error).toContain(`task_id="${pendingID}"`)
+          expect(pending.metadata).toEqual({ sessionID: pendingID })
 
           expect(errored?.state.status).toBe("error")
           if (!errored || errored.state.status !== "error") throw new Error("expected detached errored task")
-          expect(errored.state.error).toBe("original error")
-          expect(errored.state.input.task_id).toBeUndefined()
-          expect(errored.state.metadata).toEqual({ detail: "keep" })
-          expect(errored.metadata).toEqual({})
+          const clonedID = errored.state.input.task_id
+          if (typeof clonedID !== "string") throw new Error("expected a cloned task ID")
+          expect(clonedID).toBe(pendingID)
+          expect(errored.state.error).toBe(`original error; task_id="${clonedID}"`)
+          expect(errored.state.metadata).toEqual({ sessionID: clonedID, detail: "keep" })
+          expect(errored.metadata).toEqual({ sessionId: clonedID })
         },
       })
     },

packages/opencode/test/tool/task.test.ts

@@ -269,6 +269,69 @@ describe("tool.task", () => {
     }),
   )
 
+  // kilocode_change start - verify forked task children remain resumable
+  it.instance("execute resumes a cloned task session after the parent is forked", () =>
+    Effect.gen(function* () {
+      const sessions = yield* Session.Service
+      const { chat, assistant } = yield* seed()
+      const child = yield* sessions.create({ parentID: chat.id, title: "Existing child" })
+      yield* sessions.updatePart({
+        id: PartID.ascending(),
+        messageID: assistant.id,
+        sessionID: chat.id,
+        type: "tool",
+        callID: "call_1",
+        tool: "task",
+        metadata: { sessionId: child.id },
+        state: {
+          status: "completed",
+          input: { description: "inspect bug", prompt: "continue", task_id: child.id },
+          output: `<task id="${child.id}"><task_result>done</task_result></task>`,
+          title: "inspect bug",
+          metadata: { sessionId: child.id },
+          time: { start: Date.now(), end: Date.now() },
+        },
+      } as MessageV2.ToolPart)
+
+      const forked = yield* sessions.fork({ sessionID: chat.id })
+      const msgs = yield* sessions.messages({ sessionID: forked.id })
+      const part = msgs.flatMap((msg) => msg.parts).find((item) => item.type === "tool" && item.tool === "task") as
+        | MessageV2.ToolPart
+        | undefined
+      if (!part || part.state.status !== "completed") throw new Error("expected a completed task part")
+      const id = part.state.input.task_id
+      if (typeof id !== "string") throw new Error("expected a cloned task ID")
+      const parent = msgs.find((msg) => msg.info.role === "assistant")
+      if (!parent || parent.info.role !== "assistant") throw new Error("expected a forked assistant message")
+
+      const tool = yield* TaskTool
+      const def = yield* tool.init()
+      let seen: SessionPrompt.PromptInput | undefined
+      yield* def.execute(
+        {
+          description: "inspect bug",
+          prompt: "continue from the fork",
+          subagent_type: "general",
+          task_id: id,
+        },
+        {
+          sessionID: forked.id,
+          messageID: parent.info.id,
+          agent: "build",
+          abort: new AbortController().signal,
+          extra: { promptOps: stubOps({ onPrompt: (input) => (seen = input) }) },
+          messages: [],
+          metadata: () => Effect.void,
+          ask: () => Effect.void,
+        },
+      )
+
+      expect(seen?.sessionID).toBe(SessionID.descending(id))
+      expect((yield* sessions.get(SessionID.descending(id))).parentID).toBe(forked.id)
+    }),
+  )
+  // kilocode_change end
+
   // kilocode_change start - resumed children rebuild parent platform attribution after restart
   it.instance("execute preserves platform attribution when resuming a task", () =>
     Effect.gen(function* () {