perf: speed up large session forks

.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.

packages/kilo-ui/src/components/basic-tool.tsx

@@ -11,9 +11,15 @@ export interface BasicToolProps extends BaseProps {
   partID?: string
 }
 
+type OpenProps = Pick<BasicToolProps, "tool" | "callID" | "partID" | "forceOpen" | "defaultOpen">
+
+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 (
     <Base
       {...props}

packages/kilo-vscode/tests/unit/task-tool-hydration.test.ts

@@ -0,0 +1,41 @@
+import { beforeEach, describe, expect, it } from "bun:test"
+import {
+  readToolOpen,
+  resetToolOpenState,
+  toolOpenKey,
+  writeToolOpen,
+} from "../../../kilo-ui/src/components/tool-open-state"
+import { taskResult, taskRunning, taskVisible } from "../../webview-ui/src/components/chat/task-tool-state"
+
+describe("completed task hydration", () => {
+  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<task_result>\nchild outcome\n</task_result>"
+    expect(taskResult(output, undefined)).toBe("child outcome")
+    expect(taskResult(output, "ses_child")).toBeUndefined()
+    expect(taskResult("plain output", undefined)).toBe("plain output")
+  })
+})

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<ToolProps> = (props) => {
   const i18n = useI18n()
@@ -33,12 +35,17 @@ const TaskToolRenderer: Component<ToolProps> = (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<ToolProps> = (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<ToolProps> = (props) => {
         status={props.status}
         tool={props.tool}
         partID={props.partID}
-        callID={props.callID}
         trigger={trigger()}
-        defaultOpen
+        defaultOpen={running()}
+        defer
+        onOpenChange={setOpen}
       >
         <div ref={autoScroll.scrollRef} onScroll={autoScroll.handleScroll} data-component="tool-output" data-scrollable>
           <div ref={autoScroll.contentRef} data-component="task-tools">
@@ -130,6 +140,7 @@ const TaskToolRenderer: Component<ToolProps> = (props) => {
                 <span data-slot="task-tool-title">{language.t("session.messages.taskStarting")}</span>
               </div>
             </Show>
+            <Show when={result()}>{(text) => <Markdown text={text()} />}</Show>
             <For each={childToolParts()}>
               {(item) => {
                 const info = createMemo(() => getToolInfo(item.tool, item.state?.input))

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 = /<task_result>\s*([\s\S]*?)\s*<\/task_result>/.exec(output)
+  return match?.[1] ?? output
+}

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<string, SessionID>([[ctx.params.sessionID, session.id]])
-      yield* remapChildren(session.id, remapped).pipe(Effect.orDie)
-      return session
+      return yield* fork({ params: ctx.params, payload })
     })
   }
 }

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<T extends MessageV2.Info>(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<string, unknown> | undefined) {
+  if (!value) return value
+  const copy = { ...value }
+  delete copy.sessionId
+  delete copy.sessionID
+  return copy
+}
+
+function input(value: Record<string, unknown>) {
+  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<string, SessionID>(),
-): Effect.Effect<void, Session.NotFound, Session.Service> {
-  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<string, unknown> }).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) ?? {},
+    },
+  }
 }