Add message-level Agent Manager session forks

.changeset/agent-manager-message-fork.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": minor
+---
+
+Support forking Agent Manager sessions from a specific user message.

packages/kilo-ui/src/components/message-part.tsx

@@ -706,6 +706,7 @@ export function UserMessageDisplay(props: {
   interrupted?: boolean
   animate?: boolean
   queued?: boolean
+  onFork?: () => void
   onRevert?: () => void
 }) {
   const data = useData()
@@ -847,6 +848,21 @@ export function UserMessageDisplay(props: {
                   </Show>
                 </span>
               </Show>
+              <Show when={props.onFork}>
+                <Tooltip value={i18n.t("ui.message.forkMessage")} placement="right" gutter={4}>
+                  <IconButton
+                    icon="fork"
+                    size="normal"
+                    variant="ghost"
+                    onMouseDown={(e) => e.preventDefault()}
+                    onClick={(event) => {
+                      event.stopPropagation()
+                      props.onFork?.()
+                    }}
+                    aria-label={i18n.t("ui.message.forkMessage")}
+                  />
+                </Tooltip>
+              </Show>
               <Show when={props.onRevert}>
                 <Tooltip value={i18n.t("ui.message.revertMessage")} placement="right" gutter={4}>
                   <IconButton

packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts

@@ -297,7 +297,7 @@ export class AgentManagerProvider implements Disposable {
     if (m.type === "agentManager.removeStaleWorktree") return this.onRemoveStaleWorktree(m.worktreeId)
     if (m.type === "agentManager.promoteSession") return this.onPromoteSession(m.sessionId)
     if (m.type === "agentManager.addSessionToWorktree") return this.onAddSessionToWorktree(m.worktreeId)
-    if (m.type === "agentManager.forkSession") return this.onForkSession(m.sessionId, m.worktreeId)
+    if (m.type === "agentManager.forkSession") return this.onForkSession(m.sessionId, m.worktreeId, m.messageId)
     if (m.type === "agentManager.closeSession") return this.onCloseSession(m.sessionId)
   }
 
@@ -931,7 +931,7 @@ export class AgentManagerProvider implements Disposable {
     return null
   }
 
-  private onForkSession(sessionId: string, worktreeId?: string) {
+  private onForkSession(sessionId: string, worktreeId?: string, messageId?: string) {
     return forkSession(
       {
         getClient: () => this.connectionService.getClient(),
@@ -951,6 +951,7 @@ export class AgentManagerProvider implements Disposable {
       },
       sessionId,
       worktreeId,
+      messageId,
     )
   }
 

packages/kilo-vscode/src/agent-manager/fork-session.ts

@@ -21,7 +21,12 @@ export interface ForkContext {
  *
  * Pure orchestration — no vscode imports.
  */
-export async function forkSession(ctx: ForkContext, sessionId: string, worktreeId?: string): Promise<null> {
+export async function forkSession(
+  ctx: ForkContext,
+  sessionId: string,
+  worktreeId?: string,
+  messageId?: string,
+): Promise<null> {
   let client: KiloClient
   try {
     client = ctx.getClient()
@@ -38,7 +43,8 @@ export async function forkSession(ctx: ForkContext, sessionId: string, worktreeI
 
   let forked: Session
   try {
-    const { data } = await client.session.fork({ sessionID: sessionId, directory }, { throwOnError: true })
+    const input = { sessionID: sessionId, directory, ...(messageId ? { messageID: messageId } : {}) }
+    const { data } = await client.session.fork(input, { throwOnError: true })
     forked = data
   } catch (error) {
     const err = getErrorMessage(error)

packages/kilo-vscode/src/agent-manager/types.ts

@@ -590,6 +590,7 @@ interface ForkSessionIn {
   type: "agentManager.forkSession"
   sessionId: string
   worktreeId?: string
+  messageId?: string
 }
 
 interface AbortIn {

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

@@ -1888,13 +1888,12 @@ const AgentManagerContent: Component = () => {
     if (sel === LOCAL) addPendingTab()
     else if (sel) vscode.postMessage({ type: "agentManager.addSessionToWorktree", worktreeId: sel })
   }
-
-  const handleForkSession = (sessionId: string) => {
+  const handleForkSession = (sessionId: string, messageId?: string) => {
     const sel = selection()
-    if (sel === LOCAL) vscode.postMessage({ type: "agentManager.forkSession", sessionId })
-    else if (sel) vscode.postMessage({ type: "agentManager.forkSession", sessionId, worktreeId: sel })
+    const msg = { type: "agentManager.forkSession" as const, sessionId, ...(messageId ? { messageId } : {}) }
+    if (!sel || sel === LOCAL) return vscode.postMessage(msg)
+    vscode.postMessage({ ...msg, worktreeId: sel })
   }
-
   const handleCloseTab = (sessionId: string) => {
     const pending = isPending(sessionId)
     const isActive = pending ? sessionId === activePendingId() : session.currentSessionID() === sessionId
@@ -2948,6 +2947,7 @@ const AgentManagerContent: Component = () => {
                   openLocally(id)
                 }}
                 onShowHistory={() => setHistory(true)}
+                onForkMessage={readOnly() ? undefined : handleForkSession}
                 readonly={readOnly()}
                 continueInWorktree={selection() === LOCAL}
                 promptBoxId={`agent-manager:${selection() ?? "unassigned"}`}

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

@@ -24,6 +24,7 @@ import { isPromptBlocked, isSuggesting, isQuestioning } from "./prompt-input-uti
 interface ChatViewProps {
   onSelectSession?: (id: string) => void
   onShowHistory?: () => void
+  onForkMessage?: (sessionId: string, messageId: string) => void
   readonly?: boolean
   /** When true, show the "Continue in Worktree" button. Defaults to true in the sidebar. */
   continueInWorktree?: boolean
@@ -136,6 +137,7 @@ export const ChatView: Component<ChatViewProps> = (props) => {
           <MessageList
             onSelectSession={props.onSelectSession}
             onShowHistory={props.onShowHistory}
+            onForkMessage={props.onForkMessage}
             questions={standaloneQuestions}
             suggestions={standaloneSuggestions}
             readonly={props.readonly}

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

@@ -48,6 +48,7 @@ const KiloLogo = (): JSX.Element => {
 interface MessageListProps {
   onSelectSession?: (id: string) => void
   onShowHistory?: () => void
+  onForkMessage?: (sessionId: string, messageId: string) => void
   /** Non-tool question requests to render inline at the bottom of the message list */
   questions?: () => QuestionRequest[]
   /** Non-tool suggestion requests to render inline at the bottom of the message list */
@@ -238,7 +239,7 @@ export const MessageList: Component<MessageListProps> = (props) => {
                     return index() > active
                   })
 
-                  return <VscodeSessionTurn turn={turn} queued={queued()} />
+                  return <VscodeSessionTurn turn={turn} queued={queued()} onForkMessage={props.onForkMessage} />
                 }}
               </Virtualizer>
             </Show>

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

@@ -55,6 +55,7 @@ export interface VscodeTurn {
 interface VscodeSessionTurnProps {
   turn: VscodeTurn
   queued?: boolean
+  onForkMessage?: (sessionId: string, messageId: string) => void
 }
 
 export const VscodeSessionTurn: Component<VscodeSessionTurnProps> = (props) => {
@@ -153,6 +154,7 @@ export const VscodeSessionTurn: Component<VscodeSessionTurnProps> = (props) => {
               parts={parts() as unknown as Parameters<typeof UserMessageDisplay>[0]["parts"]}
               interrupted={interrupted()}
               queued={props.queued}
+              onFork={props.onForkMessage ? () => props.onForkMessage?.(msg().sessionID, msg().id) : undefined}
               onRevert={
                 assistantMessages().length > 0 && !session.revert()
                   ? () => {

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

@@ -12,15 +12,18 @@ import { ToolRegistry } from "@kilocode/kilo-ui/message-part"
 
 /** Tools that should be open by default in the VS Code sidebar. */
 const DEFAULT_OPEN_TOOLS = ["bash"]
+const registered = new Set<string>()
 
 export function registerVscodeToolOverrides() {
   for (const name of DEFAULT_OPEN_TOOLS) {
+    if (registered.has(name)) continue
     const upstream = ToolRegistry.render(name)
     if (!upstream) continue
 
     ToolRegistry.register({
       name,
       render: (props) => <Dynamic component={upstream} {...props} defaultOpen />,
     })
+    registered.add(name)
   }
 }

packages/kilo-vscode/webview-ui/src/types/messages.ts

@@ -2054,6 +2054,7 @@ export interface ForkSessionRequest {
   type: "agentManager.forkSession"
   sessionId: string
   worktreeId?: string
+  messageId?: string
 }
 
 // Close (remove) a session from its worktree