fix(vscode): preserve mode on first send

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

packages/kilo-vscode/tests/unit/prompt-send-contract.test.ts

@@ -213,12 +213,47 @@ describe("sendMessage / sendCommand draft id contract", () => {
     // from ":pending:<id>" to ":session:<newSessionId>". 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)")
   })
 })
 

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> = {}): 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<string, string> = {}
+    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<string, string> = {}
+    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([])
+  })
+})

packages/kilo-vscode/webview-ui/src/context/session-agent.ts

@@ -8,3 +8,33 @@ export function resolveSessionAgent(messages: Message[], names: Set<string>): st
     return name
   }
 }
+
+export function draftAgentSelection(selections: Record<string, string>, draft: string, pending: string | null) {
+  if (selections[draft]) return undefined
+  return pending ?? undefined
+}
+
+export function createDraftAgentSeed(opts: {
+  selections: () => Record<string, string>
+  pending: () => string | null
+  active: (draft: string) => boolean
+  set: (draft: string, agent: string) => void
+  drop: (draft: string) => void
+}) {
+  const seeded = new Set<string>()
+  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)
+    },
+  }
+}

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)