diff --git a/packages/opencode/src/kilocode/session/fork.ts b/packages/opencode/src/kilocode/session/fork.ts new file mode 100644 index 000000000000..bfbc0f24e7f1 --- /dev/null +++ b/packages/opencode/src/kilocode/session/fork.ts @@ -0,0 +1,66 @@ +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 + // 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) + } + + 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..b723fc307efb 100644 --- a/packages/opencode/src/kilocode/session/index.ts +++ b/packages/opencode/src/kilocode/session/index.ts @@ -1,6 +1,11 @@ // kilocode_change - new file +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" @@ -282,4 +287,16 @@ export namespace KiloSession { yield { ...input.fromRow(row), project } as T & { project: ProjectInfo | null } } } + + 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 644e52f29038..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,9 +745,7 @@ 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)), - ) + 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))) 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 }, + ) +})