diff --git a/.changeset/vscode-first-send-agent-scope.md b/.changeset/vscode-first-send-agent-scope.md new file mode 100644 index 00000000000..9cf8af712b6 --- /dev/null +++ b/.changeset/vscode-first-send-agent-scope.md @@ -0,0 +1,5 @@ +--- +"kilo-code": patch +--- + +Preserve the selected mode when sending the first message in a new VS Code task so the chosen model is paired with the correct agent instructions. diff --git a/packages/kilo-vscode/tests/unit/prompt-send-contract.test.ts b/packages/kilo-vscode/tests/unit/prompt-send-contract.test.ts index 0aee37e70bf..8163cf4aed4 100644 --- a/packages/kilo-vscode/tests/unit/prompt-send-contract.test.ts +++ b/packages/kilo-vscode/tests/unit/prompt-send-contract.test.ts @@ -213,12 +213,47 @@ describe("sendMessage / sendCommand draft id contract", () => { // from ":pending:" to ":session:". The user loses the // typed message and the new session starts empty. const body = extractFunctionBody(source, "sendMessage") - expect(body).toMatch(/!sid && !draftID \? crypto\.randomUUID\(\) : draftID/) + expect(body).toContain("const fresh = !sid && !draftID") + expect(body).toMatch(/const effectiveDraftID = fresh \? crypto\.randomUUID\(\) : draftID/) }) it("sendCommand mints a draftID when there is no current session and none was supplied", () => { const body = extractFunctionBody(source, "sendCommand") - expect(body).toMatch(/!sid && !draftID \? crypto\.randomUUID\(\) : draftID/) + expect(body).toContain("const fresh = !sid && !draftID") + expect(body).toMatch(/const effectiveDraftID = fresh \? crypto\.randomUUID\(\) : draftID/) + }) + + it("sendMessage seeds the pending agent before resolving the draft-scoped agent", () => { + // Fresh draft IDs are created after ModeSwitcher stored the selected mode in + // pendingAgentSelection(). The draft scope must inherit that pending agent + // before promptAgent(scope) runs, otherwise the first send pairs the selected + // model with the default agent's system prompt. + const body = extractFunctionBody(source, "sendMessage") + expect(body).toMatch( + /if \(fresh && effectiveDraftID\) agentDrafts\.seed\(effectiveDraftID\)[\s\S]*const agent = promptAgent\(scope\)/, + ) + }) + + it("sendCommand seeds the pending agent before resolving the draft-scoped agent", () => { + const body = extractFunctionBody(source, "sendCommand") + expect(body).toMatch( + /if \(fresh && effectiveDraftID\) agentDrafts\.seed\(effectiveDraftID\)[\s\S]*const agent = promptAgent\(scope\)/, + ) + }) + + it("does not clear a newer pending agent when a seeded draft is promoted", () => { + const body = extractFunctionBody(source, "handleSessionCreated") + const draftBlock = body.match(/if \(draftID\) \{([\s\S]*?)\} else if/) + expect(draftBlock).not.toBeNull() + expect(draftBlock![1]).not.toContain("setPendingAgentSelection(null)") + }) + + it("prunes seeded draft agents only after the draft is abandoned", () => { + const failed = extractFunctionBody(source, "handleSendMessageFailed") + expect(source).toMatch(/const agentDrafts = createDraftAgentSeed/) + expect(source).toContain("active: (draft) => !!submissionMap[draft]") + expect(failed).toContain("draftSessionID() !== message.draftID") + expect(failed).toContain("agentDrafts.prune(message.draftID)") }) }) diff --git a/packages/kilo-vscode/tests/unit/session-agent.test.ts b/packages/kilo-vscode/tests/unit/session-agent.test.ts index 9535e0f0951..d72839a6f80 100644 --- a/packages/kilo-vscode/tests/unit/session-agent.test.ts +++ b/packages/kilo-vscode/tests/unit/session-agent.test.ts @@ -1,5 +1,9 @@ import { describe, it, expect } from "bun:test" -import { resolveSessionAgent } from "../../webview-ui/src/context/session-agent" +import { + createDraftAgentSeed, + draftAgentSelection, + resolveSessionAgent, +} from "../../webview-ui/src/context/session-agent" import type { Message } from "../../webview-ui/src/types/messages" function makeMessage(overrides: Partial = {}): Message { @@ -76,3 +80,83 @@ describe("resolveSessionAgent", () => { expect(result).toBeUndefined() }) }) + +describe("draftAgentSelection", () => { + it("carries a pending agent into a new draft scope", () => { + const result = draftAgentSelection({}, "draft-1", "plan") + + expect(result).toBe("plan") + }) + + it("does not overwrite an existing draft agent", () => { + const result = draftAgentSelection({ "draft-1": "code" }, "draft-1", "plan") + + expect(result).toBeUndefined() + }) + + it("ignores missing pending agents", () => { + const result = draftAgentSelection({}, "draft-1", null) + + expect(result).toBeUndefined() + }) +}) + +describe("createDraftAgentSeed", () => { + it("seeds and prunes abandoned draft agents", () => { + const selections: Record = {} + const seed = createDraftAgentSeed({ + selections: () => selections, + pending: () => "plan", + active: () => false, + set: (draft, agent) => { + selections[draft] = agent + }, + drop: (draft) => { + delete selections[draft] + }, + }) + + seed.seed("draft-1") + expect(selections["draft-1"]).toBe("plan") + + seed.prune("draft-1") + expect(selections["draft-1"]).toBeUndefined() + }) + + it("keeps active drafts available for retry", () => { + const selections: Record = {} + const seed = createDraftAgentSeed({ + selections: () => selections, + pending: () => "code", + active: () => true, + set: (draft, agent) => { + selections[draft] = agent + }, + drop: (draft) => { + delete selections[draft] + }, + }) + + seed.seed("draft-1") + seed.prune("draft-1") + + expect(selections["draft-1"]).toBe("code") + }) + + it("promotes drafts without dropping the migrated agent", () => { + const dropped: string[] = [] + const seed = createDraftAgentSeed({ + selections: () => ({}), + pending: () => "ask", + active: () => false, + set: () => {}, + drop: (draft) => dropped.push(draft), + }) + + seed.seed("draft-1") + seed.promote("draft-1") + seed.prune("draft-1") + + expect(dropped).toEqual([]) + }) +}) diff --git a/packages/kilo-vscode/webview-ui/src/context/session-agent.ts b/packages/kilo-vscode/webview-ui/src/context/session-agent.ts index 2c7dcf4bfd3..86764e307a2 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session-agent.ts +++ b/packages/kilo-vscode/webview-ui/src/context/session-agent.ts @@ -8,3 +8,33 @@ export function resolveSessionAgent(messages: Message[], names: Set): st return name } } + +export function draftAgentSelection(selections: Record, draft: string, pending: string | null) { + if (selections[draft]) return undefined + return pending ?? undefined +} + +export function createDraftAgentSeed(opts: { + selections: () => Record + pending: () => string | null + active: (draft: string) => boolean + set: (draft: string, agent: string) => void + drop: (draft: string) => void +}) { + const seeded = new Set() + return { + seed(draft: string) { + const agent = draftAgentSelection(opts.selections(), draft, opts.pending()) + if (!agent) return + opts.set(draft, agent) + seeded.add(draft) + }, + promote(draft: string) { + seeded.delete(draft) + }, + prune(draft?: string) { + if (!draft || opts.active(draft) || !seeded.delete(draft)) return + opts.drop(draft) + }, + } +} diff --git a/packages/kilo-vscode/webview-ui/src/context/session.tsx b/packages/kilo-vscode/webview-ui/src/context/session.tsx index 4336c85a2c3..81b6f4f8ece 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session.tsx +++ b/packages/kilo-vscode/webview-ui/src/context/session.tsx @@ -80,6 +80,7 @@ import { deleteDraftsForSession } from "../utils/draft-store" import { createAbortState } from "./abort-state" import { clearIfOn, createCloudPrune } from "./session-cloud-prune" import { isSameSessionTree } from "./model-usage" +import { createDraftAgentSeed } from "./session-agent" const RECENT_LIMIT = 5 const MESSAGE_PAGE_LIMIT = 80 @@ -555,7 +556,17 @@ export const SessionProvider: ParentComponent = (props) => { if (sessionID) return store.agentSelections[sessionID] ?? defaultAgent() return selectedAgentName() } - + const agentDrafts = createDraftAgentSeed({ + selections: () => store.agentSelections, + pending: pendingAgentSelection, + active: (draft) => !!submissionMap[draft], + set: (draft, agent) => setStore("agentSelections", draft, agent), + drop: (draft) => + setStore( + "agentSelections", + produce((agents) => void delete agents[draft]), + ), + }) const agentNames = createMemo(() => new Set(agents().map((agent) => agent.name))) const { pendingCloudPrune, prune: pruneCloudOrphans } = createCloudPrune((m) => setStore("parts", produce(m)), stash) @@ -1319,6 +1330,7 @@ export const SessionProvider: ParentComponent = (props) => { for (const key of sessionVariantKeys(variants, draftID)) delete variants[key] }), ) + agentDrafts.promote(draftID) } else if (pendingAgent && !store.agentSelections[session.id]) { setStore("agentSelections", session.id, pendingAgent) setPendingAgentSelection(null) @@ -1835,6 +1847,7 @@ export const SessionProvider: ParentComponent = (props) => { }) if (!message.sessionID && message.draftID) { + if (draftSessionID() !== message.draftID) agentDrafts.prune(message.draftID) setDraftSessionID(message.draftID) } } @@ -2256,8 +2269,10 @@ export const SessionProvider: ParentComponent = (props) => { dismissQuestion(q.id) } - const effectiveDraftID = !sid && !draftID ? crypto.randomUUID() : draftID + const fresh = !sid && !draftID + const effectiveDraftID = fresh ? crypto.randomUUID() : draftID const scope = effectiveDraftID ?? sid + if (fresh && effectiveDraftID) agentDrafts.seed(effectiveDraftID) if (scope) { clearClose(scope) addOptimistic(scope, messageID, text, files, review) @@ -2328,8 +2343,10 @@ export const SessionProvider: ParentComponent = (props) => { dismissQuestion(q.id) } - const effectiveDraftID = !sid && !draftID ? crypto.randomUUID() : draftID + const fresh = !sid && !draftID + const effectiveDraftID = fresh ? crypto.randomUUID() : draftID const scope = effectiveDraftID ?? sid + if (fresh && effectiveDraftID) agentDrafts.seed(effectiveDraftID) if (scope) { clearClose(scope) addOptimistic(scope, messageID, `/${command} ${args}`.trim(), files) @@ -2500,11 +2517,13 @@ export const SessionProvider: ParentComponent = (props) => { } // Reset agent selection to default for the new session (model overrides persist) + agentDrafts.prune(draftSessionID()) setPendingAgentSelection(defaultAgent()) vscode.postMessage({ type: "createSession" }) } function clearCurrentSession() { + agentDrafts.prune(draftSessionID()) setUserClearedSession(true) setCurrentSessionID(undefined) setDraftSessionID(undefined) @@ -2553,6 +2572,7 @@ export const SessionProvider: ParentComponent = (props) => { // they update even while disconnected. Bailing out here when not connected // froze the chat on the previous session while the side diff (resolved from // the worktree selection) still moved (the reported "only the diff changes"). + agentDrafts.prune(draftSessionID()) setCurrentSessionID(id) setDraftSessionID(id) setUserClearedSession(false) @@ -2599,6 +2619,7 @@ export const SessionProvider: ParentComponent = (props) => { return } const key = `cloud:${cloudSessionId}` + agentDrafts.prune(draftSessionID()) setCloudPreviewId(cloudSessionId) setCurrentSessionID(key) setDraftSessionID(key)