diff --git a/.changeset/fast-session-forks.md b/.changeset/fast-session-forks.md new file mode 100644 index 00000000000..c42452d0feb --- /dev/null +++ b/.changeset/fast-session-forks.md @@ -0,0 +1,7 @@ +--- +"@kilocode/cli": patch +"@kilocode/kilo-ui": patch +"kilo-code": patch +--- + +Speed up large session forks by retaining final task outcomes instead of duplicating resumable subagent histories, and load completed task details only when expanded. diff --git a/packages/kilo-ui/src/components/basic-tool.tsx b/packages/kilo-ui/src/components/basic-tool.tsx index 1cf5a25bf1f..0f4234ad1a4 100644 --- a/packages/kilo-ui/src/components/basic-tool.tsx +++ b/packages/kilo-ui/src/components/basic-tool.tsx @@ -11,9 +11,15 @@ export interface BasicToolProps extends BaseProps { partID?: string } +type OpenProps = Pick + +export function initialOpen(props: OpenProps) { + return props.forceOpen ? true : readToolOpen(toolOpenKey(props), props.defaultOpen) +} + export function BasicTool(props: BasicToolProps) { const key = () => toolOpenKey(props) - const initial = () => (props.forceOpen ? true : readToolOpen(key(), props.defaultOpen)) + const initial = () => initialOpen(props) return ( { + beforeEach(() => resetToolOpenState()) + + it("opens running tasks and collapses completed tasks by default", () => { + expect(taskRunning("pending")).toBe(true) + expect(taskRunning("running")).toBe(true) + expect(taskRunning("completed")).toBe(false) + expect(readToolOpen(toolOpenKey({ tool: "task", partID: "part-new" }), taskRunning("completed"))).toBe(false) + }) + + it("keeps expansion state isolated by copied part ID", () => { + const source = { tool: "task", partID: "part-source", defaultOpen: false } + const fork = { tool: "task", partID: "part-fork", defaultOpen: false } + writeToolOpen(toolOpenKey(source), true) + + expect(readToolOpen(toolOpenKey(source), source.defaultOpen)).toBe(true) + expect(readToolOpen(toolOpenKey(fork), fork.defaultOpen)).toBe(false) + }) + + it("hydrates and streams a child only while expanded", () => { + expect(taskVisible(false, "ses_child")).toBeUndefined() + expect(taskVisible(true, "ses_child")).toBe("ses_child") + expect(taskVisible(true, undefined)).toBeUndefined() + }) + + it("renders the retained result when a fork has no child session", () => { + const output = "task_id: stale\n\n\nchild outcome\n" + expect(taskResult(output, undefined)).toBe("child outcome") + expect(taskResult(output, "ses_child")).toBeUndefined() + expect(taskResult("plain output", undefined)).toBe("plain output") + }) +}) diff --git a/packages/kilo-vscode/webview-ui/src/components/chat/TaskToolExpanded.tsx b/packages/kilo-vscode/webview-ui/src/components/chat/TaskToolExpanded.tsx index e672cbae562..62162fda12e 100644 --- a/packages/kilo-vscode/webview-ui/src/components/chat/TaskToolExpanded.tsx +++ b/packages/kilo-vscode/webview-ui/src/components/chat/TaskToolExpanded.tsx @@ -1,23 +1,25 @@ /** * TaskToolExpanded component - * Registers a custom "task" tool renderer that matches the v1.0.25 layout: - * a BasicTool open by default with a compact scrollable list of child tool calls, - * each shown as: icon + title + subtitle. + * Registers a custom "task" tool renderer with a compact scrollable list of + * child tool calls. Running tasks open immediately; completed tasks load their + * child details only when expanded. * * Call registerExpandedTaskTool() once at app startup to activate. */ -import { Component, createEffect, createMemo, For, Show, onCleanup } from "solid-js" +import { Component, createEffect, createMemo, createSignal, For, Show, onCleanup } from "solid-js" import { ToolRegistry, ToolProps, getToolInfo } from "@kilocode/kilo-ui/message-part" -import { BasicTool } from "@kilocode/kilo-ui/basic-tool" +import { BasicTool, initialOpen } from "@kilocode/kilo-ui/basic-tool" import { Icon } from "@kilocode/kilo-ui/icon" import { IconButton } from "@kilocode/kilo-ui/icon-button" +import { Markdown } from "@kilocode/kilo-ui/markdown" import { useLanguage } from "../../context/language" import { useI18n } from "@kilocode/kilo-ui/context/i18n" import { createAutoScroll } from "@kilocode/kilo-ui/hooks" import { useSession } from "../../context/session" import { useVSCode } from "../../context/vscode" import { childID } from "../../context/session-utils" +import { taskResult, taskRunning, taskVisible } from "./task-tool-state" const TaskToolRenderer: Component = (props) => { const i18n = useI18n() @@ -33,12 +35,17 @@ const TaskToolRenderer: Component = (props) => { state: { metadata: props.metadata as { sessionId?: string } }, }) - const running = createMemo(() => props.status === "pending" || props.status === "running") + const running = createMemo(() => taskRunning(props.status)) + const [open, setOpen] = createSignal( + initialOpen({ + tool: props.tool, + partID: props.partID, + defaultOpen: running(), + }), + ) - // Warm child session data immediately so completed task tools already have - // their compact child tool list available when the user expands them. createEffect(() => { - const id = childSessionId() + const id = taskVisible(open(), childSessionId()) if (!id) return session.syncSession(id) }) @@ -62,15 +69,17 @@ const TaskToolRenderer: Component = (props) => { return id ? session.getSessionToolCount(id) : 0 }) + const result = createMemo(() => taskResult(props.output, childSessionId())) + createEffect((prev: string | undefined) => { - const id = childSessionId() + const id = taskVisible(open(), childSessionId()) if (prev && prev !== id) vscode.postMessage({ type: "streamSessionVisible", sessionID: prev, visible: false }) if (id && id !== prev) vscode.postMessage({ type: "streamSessionVisible", sessionID: id, visible: true }) return id }) onCleanup(() => { - const id = childSessionId() + const id = taskVisible(open(), childSessionId()) if (id) vscode.postMessage({ type: "streamSessionVisible", sessionID: id, visible: false }) }) @@ -119,9 +128,10 @@ const TaskToolRenderer: Component = (props) => { status={props.status} tool={props.tool} partID={props.partID} - callID={props.callID} trigger={trigger()} - defaultOpen + defaultOpen={running()} + defer + onOpenChange={setOpen} >
@@ -130,6 +140,7 @@ const TaskToolRenderer: Component = (props) => { {language.t("session.messages.taskStarting")}
+ {(text) => } {(item) => { const info = createMemo(() => getToolInfo(item.tool, item.state?.input)) diff --git a/packages/kilo-vscode/webview-ui/src/components/chat/task-tool-state.ts b/packages/kilo-vscode/webview-ui/src/components/chat/task-tool-state.ts new file mode 100644 index 00000000000..f4079c5ef6f --- /dev/null +++ b/packages/kilo-vscode/webview-ui/src/components/chat/task-tool-state.ts @@ -0,0 +1,13 @@ +export function taskRunning(status: string | undefined) { + return status === "pending" || status === "running" +} + +export function taskVisible(open: boolean | undefined, id: string | undefined) { + return open ? id : undefined +} + +export function taskResult(output: string | undefined, id: string | undefined) { + if (id || typeof output !== "string") return + const match = /\s*([\s\S]*?)\s*<\/task_result>/.exec(output) + return match?.[1] ?? output +} diff --git a/packages/opencode/src/kilocode/server/httpapi/session-fork.ts b/packages/opencode/src/kilocode/server/httpapi/session-fork.ts index 1d3e0bd5e6d..82d33d2c3c6 100644 --- a/packages/opencode/src/kilocode/server/httpapi/session-fork.ts +++ b/packages/opencode/src/kilocode/server/httpapi/session-fork.ts @@ -1,4 +1,3 @@ -import { remapChildren } from "@/kilocode/session/fork" import { SessionID } from "@/session/schema" import { ForkPayload } from "@/server/routes/instance/httpapi/groups/session" import { Effect, Schema } from "effect" @@ -30,10 +29,7 @@ export namespace KiloSessionHttpApi { Effect.mapError(() => new HttpApiError.BadRequest({})), ) }) - const session = yield* fork({ params: ctx.params, payload }) - const remapped = new Map([[ctx.params.sessionID, session.id]]) - yield* remapChildren(session.id, remapped).pipe(Effect.orDie) - return session + return yield* fork({ params: ctx.params, payload }) }) } } diff --git a/packages/opencode/src/kilocode/session/fork.ts b/packages/opencode/src/kilocode/session/fork.ts index ce3fab18000..3e804f9d39e 100644 --- a/packages/opencode/src/kilocode/session/fork.ts +++ b/packages/opencode/src/kilocode/session/fork.ts @@ -1,71 +1,125 @@ -import { Effect } from "effect" -import { Session } from "@/session/session" import { MessageV2 } from "@/session/message-v2" -import { SessionID, PartID } from "@/session/schema" -import * as Log from "@opencode-ai/core/util/log" +import { SessionID } from "@/session/schema" +import { Database } from "@/storage/db" +import { SyncEvent } from "@/sync" +import { Effect } from "effect" -const log = Log.create({ service: "session.fork" }) +const task = "task" +const stale = /^[ \t]*task_id:[^\r\n]*(?:(?:\r?\n){1,2}|$)/m -/** - * 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 +type Item = { type: "message"; info: MessageV2.Info } | { type: "part"; part: MessageV2.Part; time: number } + +export function writer(sessionID: SessionID, sync: SyncEvent.Interface) { + const items: Item[] = [] + return { + message(info: T) { + items.push({ type: "message", info }) + return info + }, + part(part: MessageV2.Part) { + items.push({ type: "part", part: structuredClone(detachPart(part)), time: Date.now() }) + }, + commit() { + return Effect.sync(() => + Database.transaction( + () => { + // sync.run stays synchronous with publishing disabled, and its nested transaction reuses this active transaction. + for (const item of items) { + if (item.type === "message") { + Effect.runSync(sync.run(MessageV2.Event.Updated, { sessionID, info: item.info }, { publish: false })) + continue + } + Effect.runSync( + sync.run( + MessageV2.Event.PartUpdated, + { sessionID, part: item.part, time: item.time }, + { publish: false }, + ), + ) + } + }, + { behavior: "immediate" }, + ), + ) + }, + } +} + +function metadata(value: Record | undefined) { + if (!value) return value + const copy = { ...value } + delete copy.sessionId + delete copy.sessionID + return copy +} + +function input(value: Record) { + const copy = { ...value } + delete copy.task_id + return copy } /** - * 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. + * Turns copied task calls into detached historical results. * - * 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. + * 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. */ -export function remapChildren( - sid: SessionID, - remapped = new Map(), -): Effect.Effect { - return Effect.gen(function* () { - const sessions = yield* Session.Service - const msgs = yield* sessions.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 +function detachPart(part: MessageV2.Part): MessageV2.Part { + if (part.type !== "tool" || part.tool !== task) return part - for (const ref of refs) { - if (remapped.has(ref.child)) continue - const exists = yield* sessions.get(SessionID.make(ref.child)).pipe(Effect.orElseSucceed(() => undefined)) - if (!exists) continue - const forked = yield* sessions.fork({ sessionID: SessionID.make(ref.child) }) - remapped.set(ref.child, forked.id) - yield* remapChildren(forked.id, remapped) + const top = metadata(part.metadata) + const state = part.state + if (state.status === "pending") { + const now = Date.now() + return { + ...part, + metadata: top, + state: { + status: "error", + input: input(state.input), + error: "Task was still pending when this session was forked.", + time: { start: now, end: now }, + }, } + } - if (remapped.size === 0) return + 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() }, + }, + } + } - 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 - yield* sessions.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) + if (state.status === "error") { + return { + ...part, + metadata: top, + state: { + ...state, + input: input(state.input), + metadata: metadata(state.metadata), + }, } + } - log.info("remapped child sessions", { session: sid, count: remapped.size }) - }) + return { + ...part, + metadata: top, + state: { + ...state, + input: input(state.input), + output: state.output.replace(stale, ""), + metadata: metadata(state.metadata) ?? {}, + }, + } } diff --git a/packages/opencode/src/kilocode/session/index.ts b/packages/opencode/src/kilocode/session/index.ts index 3db9fd383e4..086db38c0a8 100644 --- a/packages/opencode/src/kilocode/session/index.ts +++ b/packages/opencode/src/kilocode/session/index.ts @@ -1,4 +1,4 @@ -import { remapChildren as _remapChildren } from "./fork" +import { writer as _writer } from "./fork" import z from "zod" import { Cause, Effect, Schema } from "effect" import { BusEvent } from "@/bus/bus-event" @@ -407,21 +407,13 @@ export namespace KiloSession { } } - export const remapChildren = _remapChildren + export const writer = _writer } export const kiloSessionFork = fn( z.object({ sessionID: toZod(SessionID), messageID: toZod(MessageID).optional() }), async (input) => { const { AppRuntime } = await import("@/effect/app-runtime") - return AppRuntime.runPromise( - Effect.gen(function* () { - const sessions = yield* Session.Service - const session = yield* sessions.fork(input) - const remapped = new Map([[input.sessionID, session.id]]) - yield* KiloSession.remapChildren(session.id, remapped) - return session - }), - ) + return AppRuntime.runPromise(Session.Service.use((sessions) => sessions.fork(input))) }, ) diff --git a/packages/opencode/src/session/session.ts b/packages/opencode/src/session/session.ts index e564acfd598..ced818f542b 100644 --- a/packages/opencode/src/session/session.ts +++ b/packages/opencode/src/session/session.ts @@ -744,6 +744,7 @@ export const layer: Layer.Layer< }) const msgs = yield* messages({ sessionID: input.sessionID }) const idMap = new Map() + const writer = KiloSession.writer(session.id, sync) // kilocode_change - commit copied transcript in one transaction for (const msg of msgs) { if (input.messageID && msg.info.id >= input.messageID) break @@ -751,13 +752,15 @@ export const layer: Layer.Layer< idMap.set(msg.info.id, newID) const parentID = msg.info.role === "assistant" && msg.info.parentID ? idMap.get(msg.info.parentID) : undefined - const cloned = yield* updateMessage({ + // kilocode_change start - queue copied messages for the atomic transcript commit + const cloned = writer.message({ ...msg.info, sessionID: session.id, id: newID, - ...(msg.info.role === "assistant" && { cost: 0 }), // kilocode_change - count only spend incurred after the fork + ...(msg.info.role === "assistant" && { cost: 0 }), // count only spend incurred after the fork ...(parentID && { parentID }), }) + // kilocode_change end for (const part of msg.parts) { const p: MessageV2.Part = { @@ -770,9 +773,10 @@ export const layer: Layer.Layer< if (p.type === "compaction" && p.tail_start_id) { p.tail_start_id = idMap.get(p.tail_start_id) } - yield* updatePart(p) + writer.part(p) // kilocode_change - queue copied parts for the atomic transcript commit } } + yield* writer.commit() // kilocode_change - the caller hydrates after commit; copied-row events stay silent // kilocode_change start - preserve imported/cumulative diffs when forking sessions const local = yield* storage .read(["session_diff", input.sessionID]) @@ -1033,7 +1037,7 @@ export function* listGlobal(input?: { } // kilocode_change end -// kilocode_change - preserve Kilo recursive fork/remap behavior without a Session service-local Promise runtime +// kilocode_change - delegate the exported Promise facade to the Kilo session runtime export const fork = kiloSessionFork export * as Session from "./session" diff --git a/packages/opencode/test/kilocode/session-fork-remap.test.ts b/packages/opencode/test/kilocode/session-fork-remap.test.ts index 342de6a8285..77c0618a402 100644 --- a/packages/opencode/test/kilocode/session-fork-remap.test.ts +++ b/packages/opencode/test/kilocode/session-fork-remap.test.ts @@ -8,14 +8,16 @@ import { MessageV2 } from "../../src/session/message-v2" import { MessageID, PartID, SessionID } from "../../src/session/schema" import * as Log from "@opencode-ai/core/util/log" import { disposeAllInstances, tmpdir } from "../fixture/fixture" +import { Database, eq } from "../../src/storage/db" +import { EventSequenceTable, EventTable } from "../../src/sync/event.sql" +import { Flag } from "@opencode-ai/core/flag/flag" Log.init({ print: false }) const sessions = { create: (input?: Parameters[0]) => Effect.runPromise(Session.Service.use((svc) => svc.create(input)).pipe(Effect.provide(Session.defaultLayer))), - get: (id: SessionID) => - Effect.runPromise(Session.Service.use((svc) => svc.get(id)).pipe(Effect.provide(Session.defaultLayer))), + list: () => Effect.runPromise(Session.Service.use((svc) => svc.list()).pipe(Effect.provide(Session.defaultLayer))), messages: (input: Parameters[0]) => Effect.runPromise(Session.Service.use((svc) => svc.messages(input)).pipe(Effect.provide(Session.defaultLayer))), updateMessage: (msg: T) => @@ -36,10 +38,18 @@ function taskPart(input: { messageID: string; sessionID: string; childSessionID: type: "tool", callID: "call_1", tool: "task", + metadata: { sessionId: input.childSessionID, trace: "keep" }, state: { status: "completed", - input: { description: "test task", prompt: "do something" }, - output: `task_id: ${input.childSessionID}`, + input: { description: "test task", prompt: "do something", task_id: input.childSessionID }, + output: [ + "Background task completed: test task", + `\ttask_id: ${input.childSessionID} (for resuming to continue this task if needed)`, + "", + "", + "child outcome", + "", + ].join("\r\n"), title: "test task", metadata: { sessionId: input.childSessionID, @@ -125,9 +135,9 @@ describe("Session.fork cost accounting", () => { ) }) -describe("Session.fork child session remapping", () => { +describe("Session.fork task detachment", () => { test( - "forked session gets its own copy of child sessions", + "keeps completed task outcomes without cloning child sessions", async () => { await using tmp = await tmpdir({ git: true }) await WithInstance.provide({ @@ -135,37 +145,20 @@ describe("Session.fork child session remapping", () => { fn: async () => { const parent = await sessions.create({ title: "parent" }) const child = await sessions.create({ parentID: parent.id, title: "child subagent" }) - - // Add a user message to the child so it has content - const childMsgId = await userMsg(child.id) + const childMsg = await userMsg(child.id) await sessions.updatePart({ id: PartID.ascending(), - messageID: childMsgId, + messageID: childMsg, 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 sessions.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 sessions.updatePart( - taskPart({ - messageID: parentAsstMsg, - sessionID: parent.id, - childSessionID: child.id, - }), - ) + 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 })) + const before = await sessions.list() - // Exercise the SDK and HTTP route used by Agent Manager. const client = createKiloClient({ baseUrl: "http://localhost", directory: tmp.path, @@ -175,25 +168,27 @@ describe("Session.fork child session remapping", () => { { sessionID: parent.id, directory: tmp.path }, { throwOnError: true }, ) - expect(forked.id).not.toBe(parent.id) - - // Check that the forked session's task part references a DIFFERENT child session - const forkedMsgs = await sessions.messages({ sessionID: SessionID.make(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 sessions.get(SessionID.make(meta.sessionId)) - expect(forkedChild).toBeDefined() - expect(forkedChild.id).not.toBe(child.id) + const after = await sessions.list() + expect(after).toHaveLength(before.length + 1) + + 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() + expect(tool.state.output).toBe( + "Background task completed: test task\r\n\r\nchild outcome\r\n", + ) - const forkedChildMsgs = await sessions.messages({ sessionID: forkedChild.id }) - expect(forkedChildMsgs).toHaveLength(1) - expect(forkedChildMsgs[0].parts[0].type).toBe("text") + const source = await sessions.messages({ sessionID: parent.id }) + const original = source.flatMap((msg) => msg.parts).find((part) => part.type === "tool") as MessageV2.ToolPart + expect(original.state.status).toBe("completed") + if (original.state.status !== "completed") throw new Error("expected completed source task") + expect(original.state.metadata.sessionId).toBe(child.id) + expect(original.state.input.task_id).toBe(child.id) }, }) }, @@ -201,86 +196,41 @@ describe("Session.fork child session remapping", () => { ) test( - "nested child sessions are also remapped", + "turns copied running tasks into terminal historical errors", async () => { await using tmp = await tmpdir({ git: true }) await WithInstance.provide({ directory: tmp.path, fn: async () => { - // grandchild -> child -> parent const parent = await sessions.create({ title: "parent" }) const child = await sessions.create({ parentID: parent.id, title: "child" }) - const grandchild = await sessions.create({ parentID: child.id, title: "grandchild" }) - - // grandchild has a text message - const gcMsgId = await userMsg(grandchild.id) - await sessions.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) + const user = await userMsg(parent.id) + const assistant = await asstMsg(parent.id, user) await sessions.updatePart({ id: PartID.ascending(), - messageID: childUserMsg, - sessionID: child.id, - type: "text", - text: "question", - } as MessageV2.TextPart) - const childAsstMsg = await asstMsg(child.id, childUserMsg) - await sessions.updatePart( - taskPart({ - messageID: childAsstMsg, - sessionID: child.id, - childSessionID: grandchild.id, - }), - ) - - // parent references child via task part - const parentUserMsg = await userMsg(parent.id) - await sessions.updatePart({ - id: PartID.ascending(), - messageID: parentUserMsg, + messageID: assistant, sessionID: parent.id, - type: "text", - text: "request", - } as MessageV2.TextPart) - const parentAsstMsg = await asstMsg(parent.id, parentUserMsg) - await sessions.updatePart( - taskPart({ - messageID: parentAsstMsg, - sessionID: parent.id, - childSessionID: child.id, - }), - ) + type: "tool", + callID: "call_running", + tool: "task", + metadata: { sessionId: child.id }, + state: { + status: "running", + input: { description: "running", task_id: child.id }, + metadata: { sessionId: child.id, variant: "high" }, + time: { start: Date.now() }, + }, + } as MessageV2.ToolPart) const forked = await Session.fork({ sessionID: parent.id }) - - // Verify parent-level remap - const forkedMsgs = await sessions.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 sessions.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 sessions.messages({ sessionID: SessionID.make(forkedGrandchildID) }) - expect(gcMsgs).toHaveLength(1) + 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({}) }, }) }, @@ -288,39 +238,65 @@ describe("Session.fork child session remapping", () => { ) test( - "self-referential task metadata remaps to the forked session", + "detaches pending and errored task references", async () => { await using tmp = await tmpdir({ git: true }) await WithInstance.provide({ directory: tmp.path, fn: async () => { const parent = await sessions.create({ title: "parent" }) - const msg = await userMsg(parent.id) + const child = await sessions.create({ parentID: parent.id, title: "child" }) + const user = await userMsg(parent.id) + const assistant = await asstMsg(parent.id, user) await sessions.updatePart({ id: PartID.ascending(), - messageID: msg, + messageID: assistant, sessionID: parent.id, - type: "text", - text: "self task", - } as MessageV2.TextPart) - const asst = await asstMsg(parent.id, msg) - await sessions.updatePart( - taskPart({ - messageID: asst, - sessionID: parent.id, - childSessionID: parent.id, - }), - ) + type: "tool", + callID: "call_pending", + tool: "task", + metadata: { sessionID: child.id }, + state: { + status: "pending", + input: { task_id: child.id }, + raw: "pending", + }, + } as MessageV2.ToolPart) + await sessions.updatePart({ + id: PartID.ascending(), + messageID: assistant, + sessionID: parent.id, + type: "tool", + callID: "call_error", + tool: "task", + metadata: { sessionId: child.id }, + state: { + status: "error", + input: { task_id: child.id }, + error: "original error", + metadata: { sessionID: child.id, detail: "keep" }, + time: { start: Date.now(), end: Date.now() }, + }, + } as MessageV2.ToolPart) const forked = await Session.fork({ sessionID: parent.id }) const msgs = await sessions.messages({ sessionID: forked.id }) - const tools = msgs - .flatMap((m) => m.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).toBe(forked.id) + const tools = msgs.flatMap((msg) => msg.parts).filter((part) => part.type === "tool") + const pending = tools.find((part) => part.callID === "call_pending") + const errored = tools.find((part) => part.callID === "call_error") + + 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({}) + + 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({}) }, }) }, @@ -328,95 +304,78 @@ describe("Session.fork child session remapping", () => { ) test( - "cyclic task metadata remaps each session once", + "preserves workspace sync event sequencing in the atomic copy", async () => { - await using tmp = await tmpdir({ git: true }) - await WithInstance.provide({ - directory: tmp.path, - fn: async () => { - const parent = await sessions.create({ title: "parent" }) - const child = await sessions.create({ parentID: parent.id, title: "child" }) - - const parentMsg = await userMsg(parent.id) - await sessions.updatePart({ - id: PartID.ascending(), - messageID: parentMsg, - sessionID: parent.id, - type: "text", - text: "call child", - } as MessageV2.TextPart) - const parentAsst = await asstMsg(parent.id, parentMsg) - await sessions.updatePart( - taskPart({ - messageID: parentAsst, + const flag = Flag.KILO_EXPERIMENTAL_WORKSPACES + Flag.KILO_EXPERIMENTAL_WORKSPACES = true + try { + await using tmp = await tmpdir({ git: true }) + await WithInstance.provide({ + directory: tmp.path, + fn: async () => { + const parent = await sessions.create({ title: "parent" }) + const user = await userMsg(parent.id) + await sessions.updatePart({ + id: PartID.ascending(), + messageID: user, sessionID: parent.id, - childSessionID: child.id, - }), - ) - - const childMsg = await userMsg(child.id) - await sessions.updatePart({ - id: PartID.ascending(), - messageID: childMsg, - sessionID: child.id, - type: "text", - text: "call parent", - } as MessageV2.TextPart) - const childAsst = await asstMsg(child.id, childMsg) - await sessions.updatePart( - taskPart({ - messageID: childAsst, - sessionID: child.id, - childSessionID: parent.id, - }), - ) - - const forked = await Session.fork({ sessionID: parent.id }) - const msgs = await sessions.messages({ sessionID: forked.id }) - const tools = msgs - .flatMap((m) => m.parts) - .filter((p) => p.type === "tool" && p.tool === "task") as MessageV2.ToolPart[] - const id = (tools[0].state as unknown as { metadata: { sessionId: SessionID } }).metadata.sessionId - - expect(id).not.toBe(child.id) - const copy = await sessions.get(SessionID.make(id)) - expect(copy.id).toBe(id) - - const childMsgs = await sessions.messages({ sessionID: copy.id }) - const childTools = childMsgs - .flatMap((m) => m.parts) - .filter((p) => p.type === "tool" && p.tool === "task") as MessageV2.ToolPart[] - const back = (childTools[0].state as unknown as { metadata: { sessionId: string } }).metadata.sessionId - - expect(back).toBe(forked.id) - }, - }) + type: "text", + text: "hello", + } as MessageV2.TextPart) + + const forked = await Session.fork({ sessionID: parent.id }) + const rows = Database.use((db) => + db + .select({ seq: EventTable.seq, type: EventTable.type }) + .from(EventTable) + .where(eq(EventTable.aggregate_id, forked.id)) + .orderBy(EventTable.seq) + .all(), + ) + const sequence = Database.use((db) => + db + .select({ seq: EventSequenceTable.seq }) + .from(EventSequenceTable) + .where(eq(EventSequenceTable.aggregate_id, forked.id)) + .get(), + ) + + expect(rows).toEqual([ + { seq: 0, type: "session.created.1" }, + { seq: 1, type: "message.updated.1" }, + { seq: 2, type: "message.part.updated.1" }, + ]) + expect(sequence?.seq).toBe(2) + }, + }) + } finally { + Flag.KILO_EXPERIMENTAL_WORKSPACES = flag + } }, { timeout: 30000 }, ) test( - "non-task tool parts are not affected", + "does not alter non-task parts", async () => { await using tmp = await tmpdir({ git: true }) await WithInstance.provide({ directory: tmp.path, fn: async () => { const parent = await sessions.create({ title: "parent" }) - const parentUserMsg = await userMsg(parent.id) + const user = await userMsg(parent.id) await sessions.updatePart({ id: PartID.ascending(), - messageID: parentUserMsg, + messageID: user, sessionID: parent.id, type: "text", text: "hello", } as MessageV2.TextPart) const forked = await Session.fork({ sessionID: parent.id }) - const forkedMsgs = await sessions.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") + const msgs = await sessions.messages({ sessionID: forked.id }) + expect(msgs).toHaveLength(1) + expect(msgs[0].parts[0]).toMatchObject({ type: "text", text: "hello" }) }, }) },