From 4ea52f2dd17d56ba6c7a1ac0896b17ff020314ba Mon Sep 17 00:00:00 2001 From: marius-kilocode Date: Thu, 6 Aug 2026 11:16:17 +0200 Subject: [PATCH] fix(cli): resume subagents after session fork --- .changeset/forked-subagent-resume.md | 5 + .../opencode/src/kilocode/session/fork.ts | 221 ++++++++++++++---- .../opencode/src/kilocode/session/index.ts | 3 +- packages/opencode/src/session/session.ts | 7 + .../test/kilocode/session-fork-remap.test.ts | 91 +++++--- packages/opencode/test/tool/task.test.ts | 63 +++++ 6 files changed, 309 insertions(+), 81 deletions(-) create mode 100644 .changeset/forked-subagent-resume.md diff --git a/.changeset/forked-subagent-resume.md b/.changeset/forked-subagent-resume.md new file mode 100644 index 00000000000..2a4f8cd4837 --- /dev/null +++ b/.changeset/forked-subagent-resume.md @@ -0,0 +1,5 @@ +--- +"@kilocode/cli": patch +--- + +Allow subagent tasks to be resumed after their parent session is forked. diff --git a/packages/opencode/src/kilocode/session/fork.ts b/packages/opencode/src/kilocode/session/fork.ts index b95374bf035..9c64afd28c5 100644 --- a/packages/opencode/src/kilocode/session/fork.ts +++ b/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 -// 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 | 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 | undefined, map: Map, 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) { - const copy = { ...value } - delete copy.task_id - return copy +function output(value: string, map: Map) { + 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) { + 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() + + 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 +}): Effect.Effect { + return Effect.gen(function* () { + const map = input.remapped ?? new Map() + 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 | undefined) { + if (!value) return value + return { ...value } +} + +function input(value: Record) { + return { ...value } +} diff --git a/packages/opencode/src/kilocode/session/index.ts b/packages/opencode/src/kilocode/session/index.ts index 3810eafaf36..679b8440724 100644 --- a/packages/opencode/src/kilocode/session/index.ts +++ b/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" diff --git a/packages/opencode/src/session/session.ts b/packages/opencode/src/session/session.ts index 87c5e2fb7a3..4ad5fb566af 100644 --- a/packages/opencode/src/session/session.ts +++ b/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 }) diff --git a/packages/opencode/test/kilocode/session-fork-remap.test.ts b/packages/opencode/test/kilocode/session-fork-remap.test.ts index 5d0586988c6..6e3ebe6a577 100644 --- a/packages/opencode/test/kilocode/session-fork-remap.test.ts +++ b/packages/opencode/test/kilocode/session-fork-remap.test.ts @@ -49,15 +49,12 @@ afterAll(async () => { }) const sessions = { - create: (input?: Parameters[0]) => - runtime.runPromise((svc) => svc.create(input)), + create: (input?: Parameters[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[0]) => - runtime.runPromise((svc) => svc.messages(input)), - updateMessage: (msg: T) => - runtime.runPromise((svc) => svc.updateMessage(msg)), - updatePart: (part: T) => - runtime.runPromise((svc) => svc.updatePart(part)), + messages: (input: Parameters[0]) => runtime.runPromise((svc) => svc.messages(input)), + updateMessage: (msg: T) => runtime.runPromise((svc) => svc.updateMessage(msg)), + updatePart: (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\r\nchild outcome\r\n", + `Background task completed: test task\r\n\ttask_id: ${clonedID} (for resuming to continue this task if needed)\r\n\r\n\r\nchild outcome\r\n`, + ) + + 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 }) }, }) }, diff --git a/packages/opencode/test/tool/task.test.ts b/packages/opencode/test/tool/task.test.ts index 5f05b95df7e..8a03ea08f7e 100644 --- a/packages/opencode/test/tool/task.test.ts +++ b/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: `done`, + 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* () {