From 659bc99845154b8bb55b8cf03b33ce49988a83ca Mon Sep 17 00:00:00 2001 From: marius-kilocode Date: Wed, 15 Apr 2026 10:33:19 +0200 Subject: [PATCH 1/3] fix(cli): remap child session references on fork to prevent subagent state leaks When forking a session, task tool parts referencing child sessions were copied verbatim, causing both original and forked sessions to share the same subagent references. This led to SSE events and permission prompts leaking across sessions in the Agent Manager. --- .../opencode/src/kilocode/session/fork.ts | 65 +++++ .../opencode/src/kilocode/session/index.ts | 3 + packages/opencode/src/session/index.ts | 10 +- .../test/kilocode/session-fork-remap.test.ts | 251 ++++++++++++++++++ 4 files changed, 326 insertions(+), 3 deletions(-) create mode 100644 packages/opencode/src/kilocode/session/fork.ts create mode 100644 packages/opencode/test/kilocode/session-fork-remap.test.ts diff --git a/packages/opencode/src/kilocode/session/fork.ts b/packages/opencode/src/kilocode/session/fork.ts new file mode 100644 index 000000000000..2eee0242795a --- /dev/null +++ b/packages/opencode/src/kilocode/session/fork.ts @@ -0,0 +1,65 @@ +import { Session } from "@/session" +import { MessageV2 } from "@/session/message-v2" +import { SessionID, PartID } from "@/session/schema" +import { Log } from "@/util/log" + +const log = Log.create({ service: "session.fork" }) + +/** + * Extracts the child session ID from a task tool part. + */ +function childID(part: MessageV2.Part): string | undefined { + if (part.type !== "tool" || part.tool !== "task") return undefined + return (part.state as { metadata?: { sessionId?: string } }).metadata?.sessionId +} + +/** + * Recursively fork all child (subagent) sessions referenced by task tool parts + * in the given session, then update the parts to point at the forked copies. + * + * This prevents subagent state from leaking between forked sessions in the + * same worktree: without remapping, two forked sessions would share the same + * child session references, causing SSE events and permission prompts to bleed + * across sessions. + */ +export async function remapChildren(sid: SessionID): Promise { + const msgs = await Session.messages({ sessionID: sid }) + const refs: { part: MessageV2.ToolPart; child: string }[] = [] + for (const msg of msgs) { + for (const part of msg.parts) { + const child = childID(part) + if (child) refs.push({ part: part as MessageV2.ToolPart, child }) + } + } + if (refs.length === 0) return + + const remapped = new Map() + for (const ref of refs) { + if (remapped.has(ref.child)) continue + const exists = await Session.get(SessionID.make(ref.child)).catch(() => undefined) + if (!exists) continue + const forked = await Session.fork({ sessionID: SessionID.make(ref.child) }) + remapped.set(ref.child, forked.id) + await remapChildren(forked.id) + } + + if (remapped.size === 0) return + + for (const ref of refs) { + const replacement = remapped.get(ref.child) + if (!replacement) continue + const meta = (ref.part.state as { metadata?: Record }).metadata + if (!meta) continue + await Session.updatePart({ + ...ref.part, + id: PartID.make(ref.part.id), + sessionID: SessionID.make(ref.part.sessionID), + state: { + ...ref.part.state, + metadata: { ...meta, sessionId: replacement }, + }, + } as MessageV2.ToolPart) + } + + log.info("remapped child sessions", { session: sid, count: remapped.size }) +} diff --git a/packages/opencode/src/kilocode/session/index.ts b/packages/opencode/src/kilocode/session/index.ts index 2b76e3b5a1d8..3ba6f44e41c5 100644 --- a/packages/opencode/src/kilocode/session/index.ts +++ b/packages/opencode/src/kilocode/session/index.ts @@ -1,4 +1,5 @@ // kilocode_change - new file +import { remapChildren as _remapChildren } from "./fork" import z from "zod" import { BusEvent } from "@/bus/bus-event" import { Database, eq, and, gte, isNull, desc, like, inArray, lt, or } from "@/storage/db" @@ -282,4 +283,6 @@ export namespace KiloSession { yield { ...input.fromRow(row), project } as T & { project: ProjectInfo | null } } } + + export const remapChildren = _remapChildren } diff --git a/packages/opencode/src/session/index.ts b/packages/opencode/src/session/index.ts index 644e52f29038..91f6b86d131a 100644 --- a/packages/opencode/src/session/index.ts +++ b/packages/opencode/src/session/index.ts @@ -745,9 +745,13 @@ export namespace Session { (input) => runPromise((svc) => svc.create(input)), ) - export const fork = fn(z.object({ sessionID: SessionID.zod, messageID: MessageID.zod.optional() }), (input) => - runPromise((svc) => svc.fork(input)), - ) + // kilocode_change start - remap child session references after fork to prevent subagent state leaks + export const fork = fn(z.object({ sessionID: SessionID.zod, messageID: MessageID.zod.optional() }), async (input) => { + const session = await runPromise((svc) => svc.fork(input)) + await KiloSession.remapChildren(session.id) + return session + }) + // kilocode_change end export const get = fn(SessionID.zod, (id) => runPromise((svc) => svc.get(id))) export const share = fn(SessionID.zod, (id) => runPromise((svc) => svc.share(id))) diff --git a/packages/opencode/test/kilocode/session-fork-remap.test.ts b/packages/opencode/test/kilocode/session-fork-remap.test.ts new file mode 100644 index 000000000000..eb148f9f7111 --- /dev/null +++ b/packages/opencode/test/kilocode/session-fork-remap.test.ts @@ -0,0 +1,251 @@ +import { afterEach, describe, expect, test } from "bun:test" +import { Instance } from "../../src/project/instance" +import { Session } from "../../src/session" +import { MessageV2 } from "../../src/session/message-v2" +import { MessageID, PartID, SessionID } from "../../src/session/schema" +import { Log } from "../../src/util/log" +import { tmpdir } from "../fixture/fixture" + +Log.init({ print: false }) + +afterEach(async () => { + await Instance.disposeAll() +}) + +function taskPart(input: { messageID: string; sessionID: string; childSessionID: string }): MessageV2.ToolPart { + return { + id: PartID.ascending(), + messageID: MessageID.make(input.messageID), + sessionID: SessionID.make(input.sessionID), + type: "tool", + callID: "call_1", + tool: "task", + state: { + status: "completed", + input: { description: "test task", prompt: "do something" }, + output: `task_id: ${input.childSessionID}`, + title: "test task", + metadata: { + sessionId: input.childSessionID, + model: { modelID: "test", providerID: "test" }, + }, + time: { start: Date.now(), end: Date.now() }, + }, + } +} + +async function userMsg(sid: string) { + const id = MessageID.ascending() + await Session.updateMessage({ + id, + sessionID: SessionID.make(sid), + role: "user", + time: { created: Date.now() }, + agent: "test", + model: { providerID: "test", modelID: "test" }, + tools: {}, + } as MessageV2.User) + return id +} + +async function asstMsg(sid: string, parent: string) { + const id = MessageID.ascending() + await Session.updateMessage({ + id, + sessionID: SessionID.make(sid), + role: "assistant", + time: { created: Date.now() }, + parentID: MessageID.make(parent), + modelID: "test", + providerID: "test", + mode: "", + agent: "test", + path: { cwd: "/tmp", root: "/tmp" }, + cost: 0, + tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } }, + } as MessageV2.Assistant) + return id +} + +describe("Session.fork child session remapping", () => { + test( + "forked session gets its own copy of child sessions", + async () => { + await using tmp = await tmpdir({ git: true }) + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const parent = await Session.create({ title: "parent" }) + const child = await Session.create({ parentID: parent.id, title: "child subagent" }) + + // Add a user message to the child so it has content + const childMsgId = await userMsg(child.id) + await Session.updatePart({ + id: PartID.ascending(), + messageID: childMsgId, + sessionID: child.id, + type: "text", + text: "child message content", + } as MessageV2.TextPart) + + // Add a user message then an assistant message with a task tool part referencing the child + const parentUserMsg = await userMsg(parent.id) + await Session.updatePart({ + id: PartID.ascending(), + messageID: parentUserMsg, + sessionID: parent.id, + type: "text", + text: "do something", + } as MessageV2.TextPart) + + const parentAsstMsg = await asstMsg(parent.id, parentUserMsg) + await Session.updatePart( + taskPart({ + messageID: parentAsstMsg, + sessionID: parent.id, + childSessionID: child.id, + }), + ) + + // Fork the parent session + const forked = await Session.fork({ sessionID: parent.id }) + expect(forked.id).not.toBe(parent.id) + + // Check that the forked session's task part references a DIFFERENT child session + const forkedMsgs = await Session.messages({ sessionID: forked.id }) + const parts = forkedMsgs.flatMap((m) => m.parts) + const tools = parts.filter((p) => p.type === "tool" && p.tool === "task") as MessageV2.ToolPart[] + + expect(tools).toHaveLength(1) + const meta = (tools[0].state as unknown as { metadata: { sessionId: string } }).metadata + expect(meta.sessionId).not.toBe(child.id) + + // Verify the forked child session actually exists and has content + const forkedChild = await Session.get(SessionID.make(meta.sessionId)) + expect(forkedChild).toBeDefined() + expect(forkedChild.id).not.toBe(child.id) + + const forkedChildMsgs = await Session.messages({ sessionID: forkedChild.id }) + expect(forkedChildMsgs).toHaveLength(1) + expect(forkedChildMsgs[0].parts[0].type).toBe("text") + }, + }) + }, + { timeout: 30000 }, + ) + + test( + "nested child sessions are also remapped", + async () => { + await using tmp = await tmpdir({ git: true }) + await Instance.provide({ + directory: tmp.path, + fn: async () => { + // grandchild -> child -> parent + const parent = await Session.create({ title: "parent" }) + const child = await Session.create({ parentID: parent.id, title: "child" }) + const grandchild = await Session.create({ parentID: child.id, title: "grandchild" }) + + // grandchild has a text message + const gcMsgId = await userMsg(grandchild.id) + await Session.updatePart({ + id: PartID.ascending(), + messageID: gcMsgId, + sessionID: grandchild.id, + type: "text", + text: "grandchild content", + } as MessageV2.TextPart) + + // child references grandchild via task part + const childUserMsg = await userMsg(child.id) + await Session.updatePart({ + id: PartID.ascending(), + messageID: childUserMsg, + sessionID: child.id, + type: "text", + text: "question", + } as MessageV2.TextPart) + const childAsstMsg = await asstMsg(child.id, childUserMsg) + await Session.updatePart( + taskPart({ + messageID: childAsstMsg, + sessionID: child.id, + childSessionID: grandchild.id, + }), + ) + + // parent references child via task part + const parentUserMsg = await userMsg(parent.id) + await Session.updatePart({ + id: PartID.ascending(), + messageID: parentUserMsg, + sessionID: parent.id, + type: "text", + text: "request", + } as MessageV2.TextPart) + const parentAsstMsg = await asstMsg(parent.id, parentUserMsg) + await Session.updatePart( + taskPart({ + messageID: parentAsstMsg, + sessionID: parent.id, + childSessionID: child.id, + }), + ) + + const forked = await Session.fork({ sessionID: parent.id }) + + // Verify parent-level remap + const forkedMsgs = await Session.messages({ sessionID: forked.id }) + const tools = forkedMsgs + .flatMap((m) => m.parts) + .filter((p) => p.type === "tool" && p.tool === "task") as MessageV2.ToolPart[] + const forkedChildID = (tools[0].state as unknown as { metadata: { sessionId: string } }).metadata.sessionId + expect(forkedChildID).not.toBe(child.id) + + // Verify child-level remap (grandchild) + const forkedChildMsgs = await Session.messages({ sessionID: SessionID.make(forkedChildID) }) + const childTools = forkedChildMsgs + .flatMap((m) => m.parts) + .filter((p) => p.type === "tool" && p.tool === "task") as MessageV2.ToolPart[] + expect(childTools).toHaveLength(1) + const forkedGrandchildID = (childTools[0].state as unknown as { metadata: { sessionId: string } }).metadata + .sessionId + expect(forkedGrandchildID).not.toBe(grandchild.id) + + // Verify grandchild content was copied + const gcMsgs = await Session.messages({ sessionID: SessionID.make(forkedGrandchildID) }) + expect(gcMsgs).toHaveLength(1) + }, + }) + }, + { timeout: 30000 }, + ) + + test( + "non-task tool parts are not affected", + async () => { + await using tmp = await tmpdir({ git: true }) + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const parent = await Session.create({ title: "parent" }) + const parentUserMsg = await userMsg(parent.id) + await Session.updatePart({ + id: PartID.ascending(), + messageID: parentUserMsg, + sessionID: parent.id, + type: "text", + text: "hello", + } as MessageV2.TextPart) + + const forked = await Session.fork({ sessionID: parent.id }) + const forkedMsgs = await Session.messages({ sessionID: forked.id }) + expect(forkedMsgs).toHaveLength(1) + expect(forkedMsgs[0].parts[0].type).toBe("text") + expect((forkedMsgs[0].parts[0] as MessageV2.TextPart).text).toBe("hello") + }, + }) + }, + { timeout: 30000 }, + ) +}) From bb3e76203835af96880d80b100c5330cadd5ba41 Mon Sep 17 00:00:00 2001 From: marius-kilocode Date: Wed, 15 Apr 2026 10:45:17 +0200 Subject: [PATCH 2/3] fix: remove redundant recursive remapChildren call Session.fork() already calls remapChildren via the hook, so the explicit recursive call inside remapChildren was double-forking nested descendants. --- packages/opencode/src/kilocode/session/fork.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/opencode/src/kilocode/session/fork.ts b/packages/opencode/src/kilocode/session/fork.ts index 2eee0242795a..bfbc0f24e7f1 100644 --- a/packages/opencode/src/kilocode/session/fork.ts +++ b/packages/opencode/src/kilocode/session/fork.ts @@ -38,9 +38,10 @@ export async function remapChildren(sid: SessionID): Promise { if (remapped.has(ref.child)) continue const exists = await Session.get(SessionID.make(ref.child)).catch(() => undefined) if (!exists) continue + // Session.fork() already calls remapChildren on the forked child, + // so nested subagents are handled recursively without an explicit call here. const forked = await Session.fork({ sessionID: SessionID.make(ref.child) }) remapped.set(ref.child, forked.id) - await remapChildren(forked.id) } if (remapped.size === 0) return From d295f65749a71e31d8623881c7172e7732d6194b Mon Sep 17 00:00:00 2001 From: marius-kilocode Date: Wed, 15 Apr 2026 14:50:53 +0200 Subject: [PATCH 3/3] refactor(cli): move fork remapping wrapper into kilocode session --- packages/opencode/src/kilocode/session/index.ts | 14 ++++++++++++++ packages/opencode/src/session/index.ts | 10 ++-------- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/packages/opencode/src/kilocode/session/index.ts b/packages/opencode/src/kilocode/session/index.ts index 3ba6f44e41c5..b723fc307efb 100644 --- a/packages/opencode/src/kilocode/session/index.ts +++ b/packages/opencode/src/kilocode/session/index.ts @@ -2,6 +2,10 @@ import { remapChildren as _remapChildren } from "./fork" import z from "zod" import { BusEvent } from "@/bus/bus-event" +import { Session } from "@/session" +import { MessageID, SessionID } from "@/session/schema" +import { makeRuntime } from "@/effect/run-service" +import { fn } from "@/util/fn" import { Database, eq, and, gte, isNull, desc, like, inArray, lt, or } from "@/storage/db" import type { SQL } from "@/storage/db" import { ProjectTable } from "@/project/project.sql" @@ -286,3 +290,13 @@ export namespace KiloSession { export const remapChildren = _remapChildren } + +export const kiloSessionFork = fn( + z.object({ sessionID: SessionID.zod, messageID: MessageID.zod.optional() }), + async (input) => { + const { runPromise } = makeRuntime(Session.Service, Session.defaultLayer) + const session = await runPromise((svc) => svc.fork(input)) + await KiloSession.remapChildren(session.id) + return session + }, +) diff --git a/packages/opencode/src/session/index.ts b/packages/opencode/src/session/index.ts index 91f6b86d131a..57fb2331b909 100644 --- a/packages/opencode/src/session/index.ts +++ b/packages/opencode/src/session/index.ts @@ -25,7 +25,7 @@ import { Snapshot } from "@/snapshot" import { ProjectID } from "../project/schema" import { WorkspaceID } from "../control-plane/schema" import { SessionID, MessageID, PartID } from "./schema" -import { KiloSession } from "@/kilocode/session" // kilocode_change +import { KiloSession, kiloSessionFork } from "@/kilocode/session" // kilocode_change import type { Provider } from "@/provider/provider" import { Permission } from "@/permission" @@ -745,13 +745,7 @@ export namespace Session { (input) => runPromise((svc) => svc.create(input)), ) - // kilocode_change start - remap child session references after fork to prevent subagent state leaks - export const fork = fn(z.object({ sessionID: SessionID.zod, messageID: MessageID.zod.optional() }), async (input) => { - const session = await runPromise((svc) => svc.fork(input)) - await KiloSession.remapChildren(session.id) - return session - }) - // kilocode_change end + export const fork = kiloSessionFork // kilocode_change export const get = fn(SessionID.zod, (id) => runPromise((svc) => svc.get(id))) export const share = fn(SessionID.zod, (id) => runPromise((svc) => svc.share(id)))