fix(agent-manager): start new tasks in selected worktree (#9614)

* fix(agent-manager): start new tasks in selected worktree

* test(cli): skip flaky bash truncation cancel case

* fix(agent-manager): preserve drafts for worktree new tasks

* fix(cli): serialize model selection writes

* fix(agent-manager): clear stale captured drafts

* fix(agent-manager): keep drafts through successful apply

* test(cli): re-enable truncation cancel test

* chore(agent-manager): remove worktree task plan

* test(cli): drop redundant truncation markers

* revert(cli): remove model selection write fix

.changeset/worktree-new-task.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+Start Agent Manager new tasks in the selected worktree instead of the workspace root.

packages/kilo-vscode/webview-ui/agent-manager/AgentManagerApp.tsx

@@ -100,6 +100,7 @@ import { BranchSelect } from "./BranchSelect"
 import { WorktreeItem } from "./WorktreeItem"
 import SectionHeader from "./SectionHeader"
 import { randomColor } from "./section-colors"
+import { createNewTaskDrafts } from "./new-task-drafts"
 import {
   buildTopLevelItems,
   buildSidebarOrder,
@@ -1164,6 +1165,18 @@ const AgentManagerContent: Component = () => {
     }
     window.addEventListener("focus", onWindowFocus)
 
+    const drafts = createNewTaskDrafts()
+    const newTaskHandler = (e: Event) => {
+      const sel = selection()
+      if (!sel || sel === LOCAL) return
+      e.stopImmediatePropagation()
+      const draft = drafts.create(sel)
+      window.dispatchEvent(new CustomEvent("agentManagerCaptureDraft", { detail: { id: draft.id } }))
+      terms.setActiveId(undefined)
+      vscode.postMessage({ type: "agentManager.addSessionToWorktree", worktreeId: sel })
+    }
+    window.addEventListener("newTaskRequest", newTaskHandler, true)
+
     // Add created sessions as local tabs (both direct from the prompt and
     // backend follow-ups). Dedups HTTP + SSE firing together.
     const unsubCreate = vscode.onMessage((msg) => {
@@ -1259,6 +1272,7 @@ const AgentManagerContent: Component = () => {
         saveTabMemory()
         appendToTabOrder(ev.worktreeId, ev.sessionId)
         setSelection(ev.worktreeId)
+        drafts.apply(ev.worktreeId, ev.sessionId)
         session.selectSession(ev.sessionId)
       }
 
@@ -1502,6 +1516,8 @@ const AgentManagerContent: Component = () => {
       window.removeEventListener("keydown", preventDefaults, true)
       window.removeEventListener("keydown", deleteKeyHandler)
       window.removeEventListener("focus", onWindowFocus)
+      window.removeEventListener("newTaskRequest", newTaskHandler, true)
+      drafts.cleanup()
       unsubCreate()
       unsubSessions()
       unsubRun()

packages/kilo-vscode/webview-ui/agent-manager/new-task-drafts.ts

@@ -0,0 +1,63 @@
+export interface NewTaskDraft {
+  id: string
+  worktreeId: string
+}
+
+export function createNewTaskDrafts(timeout = 30_000) {
+  let seq = 0
+  const tasks = new Map<string, string[]>()
+  const timers = new Map<string, ReturnType<typeof setTimeout>>()
+
+  const remove = (task: NewTaskDraft, discard = false) => {
+    const ids = tasks.get(task.worktreeId) ?? []
+    const next = ids.filter((id) => id !== task.id)
+    if (next.length === 0) tasks.delete(task.worktreeId)
+    else tasks.set(task.worktreeId, next)
+    const timer = timers.get(task.id)
+    if (timer) clearTimeout(timer)
+    timers.delete(task.id)
+    if (!discard) return
+    window.dispatchEvent(new CustomEvent("agentManagerDiscardDraft", { detail: { id: task.id } }))
+  }
+
+  const create = (worktreeId: string) => {
+    const task = { id: `task:${++seq}`, worktreeId }
+    tasks.set(worktreeId, [...(tasks.get(worktreeId) ?? []), task.id])
+    timers.set(
+      task.id,
+      setTimeout(() => remove(task, true), timeout),
+    )
+    return task
+  }
+
+  const take = (worktreeId: string) => {
+    const id = tasks.get(worktreeId)?.[0]
+    if (!id) return undefined
+    const task = { id, worktreeId }
+    remove(task)
+    return task
+  }
+
+  const cleanup = () => {
+    for (const ids of tasks.values()) {
+      for (const id of ids) {
+        window.dispatchEvent(new CustomEvent("agentManagerDiscardDraft", { detail: { id } }))
+      }
+    }
+    for (const timer of timers.values()) clearTimeout(timer)
+    timers.clear()
+    tasks.clear()
+  }
+
+  const apply = (worktreeId: string, sessionId: string) => {
+    const task = take(worktreeId)
+    if (!task) return
+    window.dispatchEvent(
+      new CustomEvent("agentManagerApplyDraft", {
+        detail: { id: task.id, sessionId, boxId: `agent-manager:${worktreeId}` },
+      }),
+    )
+  }
+
+  return { create, apply, cleanup }
+}

packages/kilo-vscode/webview-ui/src/components/chat/PromptInput.tsx

@@ -119,6 +119,11 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
     if (imgs.length > 0) imageDrafts.set(key, imgs)
     else imageDrafts.delete(key)
   }
+  const readDraft = () => ({
+    text: text().trim(),
+    comments: reviewComments(),
+    images: imageAttach.images(),
+  })
 
   const [text, setText] = createSignal("")
   const [reviewComments, setReviewComments] = createSignal<ReviewComment[]>([])
@@ -276,6 +281,35 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
   window.addEventListener("newTaskRequest", onNewTaskRequest)
   onCleanup(() => window.removeEventListener("newTaskRequest", onNewTaskRequest))
 
+  const captured = new Map<string, ReturnType<typeof readDraft>>()
+  const onAgentManagerCaptureDraft = (event: Event) => {
+    if (!(event instanceof CustomEvent) || typeof event.detail?.id !== "string") return
+    captured.set(event.detail.id, readDraft())
+  }
+  window.addEventListener("agentManagerCaptureDraft", onAgentManagerCaptureDraft)
+  onCleanup(() => window.removeEventListener("agentManagerCaptureDraft", onAgentManagerCaptureDraft))
+
+  const onAgentManagerApplyDraft = (event: Event) => {
+    if (!(event instanceof CustomEvent)) return
+    const id = event.detail?.id
+    const sid = event.detail?.sessionId
+    const box = event.detail?.boxId
+    if (typeof id !== "string" || typeof sid !== "string" || typeof box !== "string") return
+    const draft = captured.get(id)
+    captured.delete(id)
+    if (!draft) return
+    saveDraft(scopeDraftKey(box, sessionDraftKey(sid)), draft.text, draft.comments, draft.images)
+  }
+  window.addEventListener("agentManagerApplyDraft", onAgentManagerApplyDraft)
+  onCleanup(() => window.removeEventListener("agentManagerApplyDraft", onAgentManagerApplyDraft))
+
+  const onAgentManagerDiscardDraft = (event: Event) => {
+    if (!(event instanceof CustomEvent) || typeof event.detail?.id !== "string") return
+    captured.delete(event.detail.id)
+  }
+  window.addEventListener("agentManagerDiscardDraft", onAgentManagerDiscardDraft)
+  onCleanup(() => window.removeEventListener("agentManagerDiscardDraft", onAgentManagerDiscardDraft))
+
   // Compact/summarize the current session (mirrors canCompact guards in TaskHeader)
   const onCompact = () => {
     if (session.status() === "busy") return