Merge pull request #10609 from Kilo-Org/neat-cyclamen

fix(vscode): add context handoff for forked sessions

.changeset/forked-session-context.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+Preserve prior context in forked sessions while recognizing the selected direction and current worktree context.

packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts

@@ -1152,6 +1152,7 @@ export class AgentManagerProvider implements Disposable {
       {
         getClient: () => this.connectionService.getClient(),
         state: this.getStateManager(),
+        directory: this.getRoot(),
         postError: (msg) => this.postToWebview({ type: "error", message: msg }),
         registerWorktreeSession: (sid, dir) => this.registerWorktreeSession(sid, dir),
         pushState: () => this.pushState(),

packages/kilo-vscode/src/agent-manager/continue-in-worktree.ts

@@ -4,6 +4,7 @@ import type { WorktreeStateManager } from "./WorktreeStateManager"
 import { capture as captureGitState, apply as applyGitState, type GitSnapshot } from "./git-transfer"
 import { getErrorMessage } from "../kilo-provider-utils"
 import { PLATFORM } from "./constants"
+import { recordForkHandoff } from "./fork-handoff"
 
 export interface ContinueContext {
   root: string
@@ -82,6 +83,9 @@ export async function forkSession(ctx: ContinueContext, sessionId: string, dir:
   }
   try {
     const { data } = await client.session.fork({ sessionID: sessionId, directory: dir }, { throwOnError: true })
+    await recordForkHandoff({ client, sessionId: data.id, directory: dir }).catch((err) => {
+      ctx.log("Failed to record fork handoff:", getErrorMessage(err))
+    })
     return { ok: true, value: data }
   } catch (err) {
     return { ok: false, error: `Failed to fork session: ${getErrorMessage(err)}` }

packages/kilo-vscode/src/agent-manager/fork-handoff.ts

@@ -0,0 +1,41 @@
+import type { KiloClient } from "@kilocode/sdk/v2/client"
+
+export interface ForkHandoffInput {
+  client: KiloClient
+  sessionId: string
+  directory?: string
+}
+
+export function forkText(input: Pick<ForkHandoffInput, "directory">): string {
+  return [
+    "<system-reminder>",
+    "This session was forked from an existing session in the current repository or worktree.",
+    ...(input.directory
+      ? [
+          `Use this as the current working directory: ${input.directory}`,
+          "For this fork, this location supersedes any earlier repository or worktree location retained in the copied context.",
+        ]
+      : []),
+    "The prior conversation context was retained intentionally.",
+    "The user may continue the same task, explore an alternative approach, or provide new instructions.",
+    "Follow the user's next instruction as the direction for this fork, using retained context when relevant.",
+    "</system-reminder>",
+  ].join("\n")
+}
+
+export async function recordForkHandoff(input: ForkHandoffInput): Promise<void> {
+  const payload = {
+    sessionID: input.sessionId,
+    ...(input.directory ? { directory: input.directory } : {}),
+    noReply: true,
+    parts: [
+      {
+        type: "text",
+        text: forkText(input),
+        synthetic: true,
+      },
+    ],
+  } as Parameters<KiloClient["session"]["promptAsync"]>[0]
+
+  await input.client.session.promptAsync(payload, { throwOnError: true })
+}

packages/kilo-vscode/src/agent-manager/fork-session.ts

@@ -3,10 +3,12 @@ import { getErrorMessage } from "../kilo-provider-utils"
 import { TelemetryProxy, TelemetryEventName } from "../services/telemetry"
 import type { WorktreeStateManager } from "./WorktreeStateManager"
 import { PLATFORM } from "./constants"
+import { recordForkHandoff } from "./fork-handoff"
 
 export interface ForkContext {
   getClient: () => KiloClient
   state: WorktreeStateManager | undefined
+  directory: string | undefined
   postError: (message: string) => void
   registerWorktreeSession: (sessionId: string, directory: string) => void
   pushState: () => void
@@ -37,8 +39,8 @@ export async function forkSession(
   }
 
   const directory = (() => {
-    if (!worktreeId || !ctx.state) return undefined
-    return ctx.state.getWorktree(worktreeId)?.path
+    if (!worktreeId || !ctx.state) return ctx.directory
+    return ctx.state.getWorktree(worktreeId)?.path ?? ctx.directory
   })()
 
   let forked: Session
@@ -63,6 +65,10 @@ export async function forkSession(
     if (directory) ctx.registerWorktreeSession(forked.id, directory)
   }
 
+  await recordForkHandoff({ client, sessionId: forked.id, directory }).catch((err) => {
+    ctx.log("forkSession: failed to record fork handoff:", getErrorMessage(err))
+  })
+
   ctx.pushState()
   ctx.notifyForked(forked, sessionId, worktreeId)
   ctx.registerSession(forked)

packages/kilo-vscode/src/kilo-provider/fork-session.ts

@@ -32,6 +32,7 @@ export async function handleForkSession(ctx: ForkContext, sessionId: string, mes
     {
       getClient: () => ctx.connection.getClient(),
       state: undefined,
+      directory: ctx.directory(sessionId),
       postError: (message) => ctx.post({ type: "error", message }),
       registerWorktreeSession: () => {},
       pushState: () => {},

packages/kilo-vscode/tests/unit/continue-in-worktree.test.ts

@@ -1,4 +1,4 @@
-import { describe, expect, it } from "bun:test"
+import { describe, expect, it, mock } from "bun:test"
 import {
   abortSession,
   captureState,
@@ -8,6 +8,7 @@ import {
   type StepResult,
 } from "../../src/agent-manager/continue-in-worktree"
 import type { CreateWorktreeResult } from "../../src/agent-manager/WorktreeManager"
+import { forkText } from "../../src/agent-manager/fork-handoff"
 import type { Session } from "@kilocode/sdk/v2/client"
 
 const noop = () => {}
@@ -100,17 +101,27 @@ describe("continue-in-worktree steps", () => {
       if (!res.ok) expect(res.error).toContain("fork failed")
     })
 
-    it("returns forked session on success", async () => {
+    it("records handoff instructions for the forked worktree session", async () => {
       const forked = session("forked-1")
+      const promptAsync = mock(async () => ({}))
       const c = ctx({
         getClient: () =>
           ({
-            session: { fork: () => Promise.resolve({ data: forked }) },
+            session: { fork: () => Promise.resolve({ data: forked }), promptAsync },
           }) as never,
       })
       const res = await forkSession(c, "session-1", "/tmp/wt")
       expect(res.ok).toBe(true)
       if (res.ok) expect(res.value.id).toBe("forked-1")
+      expect(promptAsync).toHaveBeenCalledWith(
+        {
+          sessionID: "forked-1",
+          directory: "/tmp/wt",
+          noReply: true,
+          parts: [{ type: "text", text: forkText({ directory: "/tmp/wt" }), synthetic: true }],
+        },
+        { throwOnError: true },
+      )
     })
   })
 

packages/kilo-vscode/tests/unit/fork-handoff.test.ts

@@ -0,0 +1,42 @@
+import { describe, expect, it, mock } from "bun:test"
+import { forkText, recordForkHandoff } from "../../src/agent-manager/fork-handoff"
+
+describe("fork handoff", () => {
+  it("describes retained context without assuming a new task", () => {
+    const text = forkText({ directory: "/repo/.kilo/worktrees/feature" })
+
+    expect(text).toContain("This session was forked from an existing session in the current repository or worktree.")
+    expect(text).toContain("Use this as the current working directory: /repo/.kilo/worktrees/feature")
+    expect(text).toContain("this location supersedes any earlier repository or worktree location")
+    expect(text).toContain("The prior conversation context was retained intentionally.")
+    expect(text).toContain("continue the same task, explore an alternative approach, or provide new instructions")
+    expect(text).toContain("Follow the user's next instruction as the direction for this fork")
+  })
+
+  it("records a hidden no-reply handoff in the forked session", async () => {
+    const promptAsync = mock(async () => ({}))
+    const client = { session: { promptAsync } }
+
+    await recordForkHandoff({
+      client: client as never,
+      sessionId: "session-fork",
+      directory: "/repo/.kilo/worktrees/feature",
+    })
+
+    expect(promptAsync).toHaveBeenCalledWith(
+      {
+        sessionID: "session-fork",
+        directory: "/repo/.kilo/worktrees/feature",
+        noReply: true,
+        parts: [
+          {
+            type: "text",
+            text: forkText({ directory: "/repo/.kilo/worktrees/feature" }),
+            synthetic: true,
+          },
+        ],
+      },
+      { throwOnError: true },
+    )
+  })
+})

packages/kilo-vscode/tests/unit/fork-session.test.ts

@@ -0,0 +1,87 @@
+import { describe, expect, it, mock } from "bun:test"
+import type { Session } from "@kilocode/sdk/v2/client"
+import { forkText } from "../../src/agent-manager/fork-handoff"
+import { forkSession, type ForkContext } from "../../src/agent-manager/fork-session"
+
+const noop = () => {}
+
+function session(id: string): Session {
+  return { id, title: id, createdAt: "", updatedAt: "" } as Session
+}
+
+function ctx(client: unknown, overrides: Partial<ForkContext> = {}): ForkContext {
+  return {
+    getClient: () => client as never,
+    state: undefined,
+    directory: "/repo",
+    postError: noop,
+    registerWorktreeSession: noop,
+    pushState: noop,
+    notifyForked: noop,
+    registerSession: noop,
+    log: noop,
+    ...overrides,
+  }
+}
+
+describe("agent manager fork session", () => {
+  it("records the hidden handoff in the current repository", async () => {
+    const fork = mock(async () => ({ data: session("forked") }))
+    const promptAsync = mock(async () => ({}))
+    const client = { session: { fork, promptAsync } }
+
+    await forkSession(ctx(client), "source", undefined, "message")
+
+    expect(fork).toHaveBeenCalledWith(
+      { sessionID: "source", directory: "/repo", messageID: "message" },
+      { throwOnError: true },
+    )
+    expect(promptAsync).toHaveBeenCalledWith(
+      {
+        sessionID: "forked",
+        directory: "/repo",
+        noReply: true,
+        parts: [{ type: "text", text: forkText({ directory: "/repo" }), synthetic: true }],
+      },
+      { throwOnError: true },
+    )
+  })
+
+  it("uses the selected worktree directory for the handoff", async () => {
+    const fork = mock(async () => ({ data: session("forked") }))
+    const promptAsync = mock(async () => ({}))
+    const client = { session: { fork, promptAsync } }
+    const state = {
+      getWorktree: () => ({ path: "/repo/.kilo/worktrees/feature" }),
+      addSession: mock(() => undefined),
+    }
+
+    await forkSession(ctx(client, { state: state as never }), "source", "worktree")
+
+    expect(fork).toHaveBeenCalledWith(
+      { sessionID: "source", directory: "/repo/.kilo/worktrees/feature" },
+      { throwOnError: true },
+    )
+    expect(promptAsync).toHaveBeenCalledWith(expect.objectContaining({ directory: "/repo/.kilo/worktrees/feature" }), {
+      throwOnError: true,
+    })
+  })
+
+  it("still exposes the fork when recording the handoff fails", async () => {
+    const notify = mock(() => undefined)
+    const log = mock(() => undefined)
+    const client = {
+      session: {
+        fork: mock(async () => ({ data: session("forked") })),
+        promptAsync: mock(async () => {
+          throw new Error("handoff failed")
+        }),
+      },
+    }
+
+    await forkSession(ctx(client, { notifyForked: notify, log }), "source")
+
+    expect(notify).toHaveBeenCalledWith(expect.objectContaining({ id: "forked" }), "source", undefined)
+    expect(log).toHaveBeenCalledWith("forkSession: failed to record fork handoff:", "handoff failed")
+  })
+})