diff --git a/.changeset/agent-manager-message-fork.md b/.changeset/agent-manager-message-fork.md new file mode 100644 index 00000000000..23fc7bdf7f3 --- /dev/null +++ b/.changeset/agent-manager-message-fork.md @@ -0,0 +1,5 @@ +--- +"kilo-code": minor +--- + +Support forking Agent Manager sessions from a specific user message. diff --git a/packages/kilo-ui/src/components/message-part.tsx b/packages/kilo-ui/src/components/message-part.tsx index 1bbbdb1eb7a..e7392935d2a 100644 --- a/packages/kilo-ui/src/components/message-part.tsx +++ b/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: { + + + e.preventDefault()} + onClick={(event) => { + event.stopPropagation() + props.onFork?.() + }} + aria-label={i18n.t("ui.message.forkMessage")} + /> + + this.connectionService.getClient(), @@ -951,6 +951,7 @@ export class AgentManagerProvider implements Disposable { }, sessionId, worktreeId, + messageId, ) } diff --git a/packages/kilo-vscode/src/agent-manager/fork-session.ts b/packages/kilo-vscode/src/agent-manager/fork-session.ts index 2f688fdbabd..c78e8086b14 100644 --- a/packages/kilo-vscode/src/agent-manager/fork-session.ts +++ b/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 { +export async function forkSession( + ctx: ForkContext, + sessionId: string, + worktreeId?: string, + messageId?: string, +): Promise { 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) diff --git a/packages/kilo-vscode/src/agent-manager/types.ts b/packages/kilo-vscode/src/agent-manager/types.ts index 8392b332014..86807ba4ce9 100644 --- a/packages/kilo-vscode/src/agent-manager/types.ts +++ b/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 { diff --git a/packages/kilo-vscode/webview-ui/agent-manager/AgentManagerApp.tsx b/packages/kilo-vscode/webview-ui/agent-manager/AgentManagerApp.tsx index ffdd0c1031b..61bc4d6d782 100644 --- a/packages/kilo-vscode/webview-ui/agent-manager/AgentManagerApp.tsx +++ b/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"}`} diff --git a/packages/kilo-vscode/webview-ui/src/components/chat/ChatView.tsx b/packages/kilo-vscode/webview-ui/src/components/chat/ChatView.tsx index 6e6969b60f5..7da3fa1dce4 100644 --- a/packages/kilo-vscode/webview-ui/src/components/chat/ChatView.tsx +++ b/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 = (props) => { { 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 = (props) => { return index() > active }) - return + return }} diff --git a/packages/kilo-vscode/webview-ui/src/components/chat/VscodeSessionTurn.tsx b/packages/kilo-vscode/webview-ui/src/components/chat/VscodeSessionTurn.tsx index 4daaaa6f318..01e84785c70 100644 --- a/packages/kilo-vscode/webview-ui/src/components/chat/VscodeSessionTurn.tsx +++ b/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 = (props) => { @@ -153,6 +154,7 @@ export const VscodeSessionTurn: Component = (props) => { parts={parts() as unknown as Parameters[0]["parts"]} interrupted={interrupted()} queued={props.queued} + onFork={props.onForkMessage ? () => props.onForkMessage?.(msg().sessionID, msg().id) : undefined} onRevert={ assistantMessages().length > 0 && !session.revert() ? () => { diff --git a/packages/kilo-vscode/webview-ui/src/components/chat/VscodeToolOverrides.tsx b/packages/kilo-vscode/webview-ui/src/components/chat/VscodeToolOverrides.tsx index ba84289e659..55ddd9d102a 100644 --- a/packages/kilo-vscode/webview-ui/src/components/chat/VscodeToolOverrides.tsx +++ b/packages/kilo-vscode/webview-ui/src/components/chat/VscodeToolOverrides.tsx @@ -12,9 +12,11 @@ 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() export function registerVscodeToolOverrides() { for (const name of DEFAULT_OPEN_TOOLS) { + if (registered.has(name)) continue const upstream = ToolRegistry.render(name) if (!upstream) continue @@ -22,5 +24,6 @@ export function registerVscodeToolOverrides() { name, render: (props) => , }) + registered.add(name) } } diff --git a/packages/kilo-vscode/webview-ui/src/types/messages.ts b/packages/kilo-vscode/webview-ui/src/types/messages.ts index 3b82ebc252e..69fda551c73 100644 --- a/packages/kilo-vscode/webview-ui/src/types/messages.ts +++ b/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