From 4e6f366a75c71b6c5a2e3499e116b61c21355fbe Mon Sep 17 00:00:00 2001 From: marius-kilocode Date: Wed, 27 May 2026 13:49:03 +0200 Subject: [PATCH] fix(vscode): add context handoff for forked sessions --- .changeset/forked-session-context.md | 5 ++ .../src/agent-manager/AgentManagerProvider.ts | 1 + .../src/agent-manager/continue-in-worktree.ts | 4 + .../src/agent-manager/fork-handoff.ts | 41 +++++++++ .../src/agent-manager/fork-session.ts | 10 ++- .../src/kilo-provider/fork-session.ts | 1 + .../tests/unit/continue-in-worktree.test.ts | 17 +++- .../tests/unit/fork-handoff.test.ts | 42 +++++++++ .../tests/unit/fork-session.test.ts | 87 +++++++++++++++++++ 9 files changed, 203 insertions(+), 5 deletions(-) create mode 100644 .changeset/forked-session-context.md create mode 100644 packages/kilo-vscode/src/agent-manager/fork-handoff.ts create mode 100644 packages/kilo-vscode/tests/unit/fork-handoff.test.ts create mode 100644 packages/kilo-vscode/tests/unit/fork-session.test.ts diff --git a/.changeset/forked-session-context.md b/.changeset/forked-session-context.md new file mode 100644 index 00000000000..b3061ce95e7 --- /dev/null +++ b/.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. diff --git a/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts b/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts index d14dcbaf8ec..dd16ecc6098 100644 --- a/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts +++ b/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(), diff --git a/packages/kilo-vscode/src/agent-manager/continue-in-worktree.ts b/packages/kilo-vscode/src/agent-manager/continue-in-worktree.ts index e25cb355588..cd49f2d2064 100644 --- a/packages/kilo-vscode/src/agent-manager/continue-in-worktree.ts +++ b/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)}` } diff --git a/packages/kilo-vscode/src/agent-manager/fork-handoff.ts b/packages/kilo-vscode/src/agent-manager/fork-handoff.ts new file mode 100644 index 00000000000..efbf36361c9 --- /dev/null +++ b/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): string { + return [ + "", + "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.", + "", + ].join("\n") +} + +export async function recordForkHandoff(input: ForkHandoffInput): Promise { + const payload = { + sessionID: input.sessionId, + ...(input.directory ? { directory: input.directory } : {}), + noReply: true, + parts: [ + { + type: "text", + text: forkText(input), + synthetic: true, + }, + ], + } as Parameters[0] + + await input.client.session.promptAsync(payload, { throwOnError: true }) +} diff --git a/packages/kilo-vscode/src/agent-manager/fork-session.ts b/packages/kilo-vscode/src/agent-manager/fork-session.ts index c78e8086b14..d6d4118c6aa 100644 --- a/packages/kilo-vscode/src/agent-manager/fork-session.ts +++ b/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) diff --git a/packages/kilo-vscode/src/kilo-provider/fork-session.ts b/packages/kilo-vscode/src/kilo-provider/fork-session.ts index 1416320581c..d97fc213a0b 100644 --- a/packages/kilo-vscode/src/kilo-provider/fork-session.ts +++ b/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: () => {}, diff --git a/packages/kilo-vscode/tests/unit/continue-in-worktree.test.ts b/packages/kilo-vscode/tests/unit/continue-in-worktree.test.ts index c2dab4a182a..4b72acdcac0 100644 --- a/packages/kilo-vscode/tests/unit/continue-in-worktree.test.ts +++ b/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 }, + ) }) }) diff --git a/packages/kilo-vscode/tests/unit/fork-handoff.test.ts b/packages/kilo-vscode/tests/unit/fork-handoff.test.ts new file mode 100644 index 00000000000..9fb775cd220 --- /dev/null +++ b/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 }, + ) + }) +}) diff --git a/packages/kilo-vscode/tests/unit/fork-session.test.ts b/packages/kilo-vscode/tests/unit/fork-session.test.ts new file mode 100644 index 00000000000..b5526c2688c --- /dev/null +++ b/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 { + 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") + }) +})