From eb8334db016e7942ce3f99a55353e2e10cfadf39 Mon Sep 17 00:00:00 2001 From: Marius Date: Fri, 8 May 2026 13:45:47 +0200 Subject: [PATCH] fix(agent-manager): preserve promoted session worktree context (#10058) * fix(agent-manager): preserve promoted session worktree context * fix(cli): guard worktree context outside instance scope --- .../painted-toaster-worktree-awareness.md | 5 +++ packages/kilo-vscode/src/KiloProvider.ts | 16 ++++---- .../src/agent-manager/AgentManagerProvider.ts | 15 +++++++ .../src/agent-manager/promotion-handoff.ts | 35 ++++++++++++++++ .../tests/unit/promotion-handoff.test.ts | 40 +++++++++++++++++++ .../opencode/src/kilocode/editor-context.ts | 8 ++++ .../opencode/src/kilocode/session/prompt.ts | 12 +++++- .../opencode/src/kilocode/system-prompt.ts | 25 ++++++++++++ packages/opencode/src/session/system.ts | 23 ++--------- .../test/kilocode/system-prompt.test.ts | 15 +++++++ 10 files changed, 166 insertions(+), 28 deletions(-) create mode 100644 .changeset/painted-toaster-worktree-awareness.md create mode 100644 packages/kilo-vscode/src/agent-manager/promotion-handoff.ts create mode 100644 packages/kilo-vscode/tests/unit/promotion-handoff.test.ts create mode 100644 packages/opencode/src/kilocode/system-prompt.ts diff --git a/.changeset/painted-toaster-worktree-awareness.md b/.changeset/painted-toaster-worktree-awareness.md new file mode 100644 index 00000000000..e62c883cdbd --- /dev/null +++ b/.changeset/painted-toaster-worktree-awareness.md @@ -0,0 +1,5 @@ +--- +"kilo-code": patch +--- + +Keep Agent Manager promoted sessions aware of their worktree location. diff --git a/packages/kilo-vscode/src/KiloProvider.ts b/packages/kilo-vscode/src/KiloProvider.ts index 88314baa945..cf3c33479e2 100644 --- a/packages/kilo-vscode/src/KiloProvider.ts +++ b/packages/kilo-vscode/src/KiloProvider.ts @@ -2560,14 +2560,14 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper } parts.push({ type: "text", text }) - const editorContext = await this.gatherEditorContext() - - if (messageID) { - this.connectionService.recordMessageSessionId(messageID, resolved!.sid) - } - const sid = resolved!.sid const dir = resolved!.dir + const editorContext = await this.gatherEditorContext(dir) + + if (messageID) { + this.connectionService.recordMessageSessionId(messageID, sid) + } + await runWithMessageConfirmation(this.confirmations, messageID, "KiloProvider: Message request", () => this.withRetry( () => @@ -3253,8 +3253,8 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper return controller } - private async gatherEditorContext(): Promise { - const workspaceDir = this.getWorkspaceDirectory() + private async gatherEditorContext(dir?: string): Promise { + const workspaceDir = dir ?? this.getWorkspaceDirectory() const controller = await this.getIgnoreController(workspaceDir) const toRelative = (fsPath: string): string | undefined => { diff --git a/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts b/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts index c4d875c5420..4890788eaeb 100644 --- a/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts +++ b/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts @@ -29,6 +29,7 @@ import { forkSession } from "./fork-session" import { continueInWorktree } from "./continue-in-worktree" import { WorktreeDiffController } from "./worktree-diff-controller" import { WorktreeImporter } from "./worktree-importer" +import { recordPromotionHandoff } from "./promotion-handoff" import { restoreWorktrees } from "./state-recovery" import { diffSummary as localDiffSummary, diffFile as localDiffFile } from "./local-diff" import { parseToolRequest, startFromTool, type ToolRequest } from "./tool-start" @@ -1017,11 +1018,25 @@ export class AgentManagerProvider implements Disposable { } this.registerWorktreeSession(sessionId, created.result.path) + await this.recordPromotionHandoff(sessionId, created.result.path, created.result.branch) this.notifyWorktreeReady(sessionId, created.result, created.worktree.id) this.log(`Promoted session ${sessionId} to worktree ${created.worktree.id}`) return null } + private async recordPromotionHandoff(sessionId: string, dir: string, branch: string): Promise { + try { + await recordPromotionHandoff({ + client: this.connectionService.getClient(), + sessionId, + directory: dir, + branch, + }) + } catch (err) { + this.log("Failed to record worktree promotion handoff:", getErrorMessage(err)) + } + } + /** Add a new session to an existing worktree. */ private async onAddSessionToWorktree(worktreeId: string): Promise { let client: KiloClient diff --git a/packages/kilo-vscode/src/agent-manager/promotion-handoff.ts b/packages/kilo-vscode/src/agent-manager/promotion-handoff.ts new file mode 100644 index 00000000000..31d20870b55 --- /dev/null +++ b/packages/kilo-vscode/src/agent-manager/promotion-handoff.ts @@ -0,0 +1,35 @@ +import type { KiloClient } from "@kilocode/sdk/v2/client" + +export interface PromoteHandoffInput { + client: KiloClient + sessionId: string + directory: string + branch: string +} + +export function handoffText(input: Omit): string { + return [ + "", + "This session was moved to a git worktree.", + `Use this as the current working directory: ${input.directory}`, + `The worktree branch is: ${input.branch}`, + "", + ].join("\n") +} + +export async function recordPromotionHandoff(input: PromoteHandoffInput): Promise { + const payload = { + sessionID: input.sessionId, + directory: input.directory, + noReply: true, + parts: [ + { + type: "text", + text: handoffText(input), + synthetic: true, + }, + ], + } as Parameters[0] + + await input.client.session.promptAsync(payload, { throwOnError: true }) +} diff --git a/packages/kilo-vscode/tests/unit/promotion-handoff.test.ts b/packages/kilo-vscode/tests/unit/promotion-handoff.test.ts new file mode 100644 index 00000000000..401b3418e34 --- /dev/null +++ b/packages/kilo-vscode/tests/unit/promotion-handoff.test.ts @@ -0,0 +1,40 @@ +import { describe, expect, it, mock } from "bun:test" +import { handoffText, recordPromotionHandoff } from "../../src/agent-manager/promotion-handoff" + +describe("promotion handoff", () => { + it("describes the new worktree location", () => { + const text = handoffText({ directory: "/repo/.kilo/worktrees/feature", branch: "feature/test" }) + + expect(text).toContain("This session was moved to a git worktree.") + expect(text).toContain("Use this as the current working directory: /repo/.kilo/worktrees/feature") + expect(text).toContain("The worktree branch is: feature/test") + }) + + it("records a hidden no-reply handoff in the worktree instance", async () => { + const promptAsync = mock(async () => ({})) + const client = { session: { promptAsync } } + + await recordPromotionHandoff({ + client: client as never, + sessionId: "session-1", + directory: "/repo/.kilo/worktrees/feature", + branch: "feature/test", + }) + + expect(promptAsync).toHaveBeenCalledWith( + { + sessionID: "session-1", + directory: "/repo/.kilo/worktrees/feature", + noReply: true, + parts: [ + { + type: "text", + text: handoffText({ directory: "/repo/.kilo/worktrees/feature", branch: "feature/test" }), + synthetic: true, + }, + ], + }, + { throwOnError: true }, + ) + }) +}) diff --git a/packages/opencode/src/kilocode/editor-context.ts b/packages/opencode/src/kilocode/editor-context.ts index 3a087bc82af..1b90c019808 100644 --- a/packages/opencode/src/kilocode/editor-context.ts +++ b/packages/opencode/src/kilocode/editor-context.ts @@ -1,4 +1,6 @@ export interface EditorContext { + directory?: string + worktree?: string visibleFiles?: string[] openTabs?: string[] activeFile?: string @@ -38,6 +40,12 @@ function timestamp(): string { export function environmentDetails(ctx?: EditorContext): string { const lines: string[] = [`Current time: ${timestamp()}`] + if (ctx?.directory) { + lines.push(`Working directory: ${ctx.directory}`) + } + if (ctx?.worktree) { + lines.push(`Workspace root folder: ${ctx.worktree}`) + } if (ctx?.activeFile) { lines.push(`Active file: ${ctx.activeFile}`) } diff --git a/packages/opencode/src/kilocode/session/prompt.ts b/packages/opencode/src/kilocode/session/prompt.ts index 37bfab8d568..be01a7c0bd1 100644 --- a/packages/opencode/src/kilocode/session/prompt.ts +++ b/packages/opencode/src/kilocode/session/prompt.ts @@ -139,7 +139,17 @@ export namespace KiloSessionPrompt { cache: EnvCache }) { if (input.cache.user !== input.lastUser.id) { - input.cache.block = environmentDetails(input.lastUser.editorContext) + const ctx = (() => { + try { + return Instance.current + } catch { + return undefined + } + })() + input.cache.block = environmentDetails({ + ...input.lastUser.editorContext, + ...(ctx ? { directory: ctx.directory, worktree: ctx.worktree } : {}), + }) input.cache.user = input.lastUser.id } if (!input.cache.block) return diff --git a/packages/opencode/src/kilocode/system-prompt.ts b/packages/opencode/src/kilocode/system-prompt.ts new file mode 100644 index 00000000000..02743776743 --- /dev/null +++ b/packages/opencode/src/kilocode/system-prompt.ts @@ -0,0 +1,25 @@ +// kilocode_change - new file + +import { Global } from "@opencode-ai/core/global" +import { staticEnvLines, type EditorContext } from "@/kilocode/editor-context" +import type { Provider } from "@/provider/provider" +import type { InstanceContext } from "@/project/instance" + +export namespace KilocodeSystemPrompt { + export function environment(input: { ctx: InstanceContext; model: Provider.Model; editor?: EditorContext }) { + return [ + [ + `You are powered by the model named ${input.model.api.id}. The exact model ID is ${input.model.providerID}/${input.model.api.id}`, + `Here is some useful information about the environment you are running in:`, + ``, + ` Is directory a git repo: ${input.ctx.project.vcs === "git" ? "yes" : "no"}`, + ` Platform: ${process.platform}`, + ` Today's date: ${new Date().toDateString()}`, + ` Project config: .kilo/command/*.md, .kilo/agent/*.md, kilo.json, AGENTS.md. Put new commands and agents in .kilo/. Do not use .kilocode/ or .opencode/.`, + ` Global config: ${Global.Path.config}/ (same structure)`, + ...staticEnvLines(input.editor), + ``, + ].join("\n"), + ] + } +} diff --git a/packages/opencode/src/session/system.ts b/packages/opencode/src/session/system.ts index f41e8980b54..cbe1b1aa00c 100644 --- a/packages/opencode/src/session/system.ts +++ b/packages/opencode/src/session/system.ts @@ -1,6 +1,5 @@ import { Context, Effect, Layer } from "effect" -import { Global } from "@opencode-ai/core/global" // kilocode_change import { InstanceState } from "@/effect/instance-state" import PROMPT_ANTHROPIC from "./prompt/anthropic.txt" @@ -21,7 +20,8 @@ import { Skill } from "@/skill" // kilocode_change start import SOUL from "../kilocode/soul.txt" -import { staticEnvLines, type EditorContext } from "../kilocode/editor-context" +import type { EditorContext } from "../kilocode/editor-context" +import { KilocodeSystemPrompt } from "../kilocode/system-prompt" import { isLing } from "../kilocode/model-match" // kilocode_change end @@ -97,25 +97,10 @@ export const layer = Layer.effect( model: Provider.Model, editorContext?: EditorContext, ) { - // kilocode_change end const ctx = yield* InstanceState.context - return [ - [ - `You are powered by the model named ${model.api.id}. The exact model ID is ${model.providerID}/${model.api.id}`, - `Here is some useful information about the environment you are running in:`, - ``, - ` Working directory: ${ctx.directory}`, - ` Workspace root folder: ${ctx.worktree}`, - ` Is directory a git repo: ${ctx.project.vcs === "git" ? "yes" : "no"}`, - ` Platform: ${process.platform}`, - ` Today's date: ${new Date().toDateString()}`, - ` Project config: .kilo/command/*.md, .kilo/agent/*.md, kilo.json, AGENTS.md. Put new commands and agents in .kilo/. Do not use .kilocode/ or .opencode/.`, // kilocode_change - ` Global config: ${Global.Path.config}/ (same structure)`, // kilocode_change - ...staticEnvLines(editorContext), // kilocode_change - ``, - ].join("\n"), - ] + return KilocodeSystemPrompt.environment({ ctx, model, editor: editorContext }) }), + // kilocode_change end skills: Effect.fn("SystemPrompt.skills")(function* (agent: Agent.Info) { if (Permission.disabled(["skill"], agent.permission).has("skill")) return diff --git a/packages/opencode/test/kilocode/system-prompt.test.ts b/packages/opencode/test/kilocode/system-prompt.test.ts index 4945204e622..221ac3f8ca5 100644 --- a/packages/opencode/test/kilocode/system-prompt.test.ts +++ b/packages/opencode/test/kilocode/system-prompt.test.ts @@ -1,5 +1,6 @@ import { describe, expect, test } from "bun:test" import { SystemPrompt } from "../../src/session/system" +import { environmentDetails } from "../../src/kilocode/editor-context" import { ProviderTest } from "../fake/provider" import PROMPT_ANTHROPIC from "../../src/session/prompt/anthropic.txt" @@ -125,3 +126,17 @@ describe("SystemPrompt.provider", () => { }) }) }) + +describe("environmentDetails", () => { + test("includes cwd and worktree in dynamic context", () => { + const result = environmentDetails({ + directory: "/repo/.kilo/worktrees/feature", + worktree: "/repo/.kilo/worktrees/feature", + activeFile: "src/app.ts", + }) + + expect(result).toContain("Working directory: /repo/.kilo/worktrees/feature") + expect(result).toContain("Workspace root folder: /repo/.kilo/worktrees/feature") + expect(result).toContain("Active file: src/app.ts") + }) +})