From 8e1e99ca0522e2e727ab995b7a4f35d6ae3d0617 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Wed, 31 Dec 2025 13:47:34 +0100 Subject: [PATCH] Change branch() to use entryId instead of entryIndex - AgentSession.branch(entryId: string) now takes entry ID - SessionBeforeBranchEvent.entryId replaces entryIndex - getUserMessagesForBranching() returns entryId - Update RPC types and client - Update UserMessageSelectorComponent - Update hook examples and tests - Update docs (hooks.md, sdk.md) --- packages/coding-agent/docs/hooks.md | 18 ++++++++++----- packages/coding-agent/docs/sdk.md | 17 +++++++------- .../examples/hooks/confirm-destructive.ts | 2 +- .../examples/hooks/git-checkpoint.ts | 17 +++++++++----- .../coding-agent/src/core/agent-session.ts | 22 +++++++++---------- packages/coding-agent/src/core/hooks/types.ts | 4 ++-- .../components/user-message-selector.ts | 8 +++---- .../src/modes/interactive/interactive-mode.ts | 6 ++--- .../coding-agent/src/modes/rpc/rpc-client.ts | 8 +++---- .../coding-agent/src/modes/rpc/rpc-mode.ts | 2 +- .../coding-agent/src/modes/rpc/rpc-types.ts | 4 ++-- .../test/agent-session-branching.test.ts | 6 ++--- 12 files changed, 64 insertions(+), 50 deletions(-) diff --git a/packages/coding-agent/docs/hooks.md b/packages/coding-agent/docs/hooks.md index 2cde1a0c..dba42e0b 100644 --- a/packages/coding-agent/docs/hooks.md +++ b/packages/coding-agent/docs/hooks.md @@ -195,7 +195,7 @@ Fired when branching via `/branch`. ```typescript pi.on("session_before_branch", async (event, ctx) => { - // event.entryIndex - entry index being branched from + // event.entryId - ID of the entry being branched from return { cancel: true }; // Cancel branch // OR @@ -634,15 +634,23 @@ export default function (pi: HookAPI) { import type { HookAPI } from "@mariozechner/pi-coding-agent/hooks"; export default function (pi: HookAPI) { - const checkpoints = new Map(); + const checkpoints = new Map(); + let currentEntryId: string | undefined; - pi.on("turn_start", async (event) => { + pi.on("tool_result", async (_event, ctx) => { + const leaf = ctx.sessionManager.getLeafEntry(); + if (leaf) currentEntryId = leaf.id; + }); + + pi.on("turn_start", async () => { const { stdout } = await pi.exec("git", ["stash", "create"]); - if (stdout.trim()) checkpoints.set(event.turnIndex, stdout.trim()); + if (stdout.trim() && currentEntryId) { + checkpoints.set(currentEntryId, stdout.trim()); + } }); pi.on("session_before_branch", async (event, ctx) => { - const ref = checkpoints.get(event.entryIndex); + const ref = checkpoints.get(event.entryId); if (!ref || !ctx.hasUI) return; const ok = await ctx.ui.confirm("Restore?", "Restore code to checkpoint?"); diff --git a/packages/coding-agent/docs/sdk.md b/packages/coding-agent/docs/sdk.md index 0a9eccdc..6363fa1e 100644 --- a/packages/coding-agent/docs/sdk.md +++ b/packages/coding-agent/docs/sdk.md @@ -99,11 +99,12 @@ interface AgentSession { isStreaming: boolean; // Session management - newSession(): Promise; // Returns false if cancelled by hook + reset(): Promise; // Returns false if cancelled by hook switchSession(sessionPath: string): Promise; - // Branching (tree-based) - branch(entryId: string): Promise<{ cancelled: boolean }>; + // Branching + branch(entryId: string): Promise<{ selectedText: string; cancelled: boolean }>; // Creates new session file + navigateTree(targetId: string, options?: { summarize?: boolean }): Promise<{ editorText?: string; cancelled: boolean }>; // In-place navigation // Hook message injection sendHookMessage(message: HookMessage, triggerTurn?: boolean): void; @@ -400,10 +401,10 @@ const { session } = await createAgentSession({ ```typescript import { Type } from "@sinclair/typebox"; -import { createAgentSession, discoverCustomTools, type CustomAgentTool } from "@mariozechner/pi-coding-agent"; +import { createAgentSession, discoverCustomTools, type CustomTool } from "@mariozechner/pi-coding-agent"; // Inline custom tool -const myTool: CustomAgentTool = { +const myTool: CustomTool = { name: "my_tool", label: "My Tool", description: "Does something useful", @@ -793,7 +794,7 @@ import { readTool, bashTool, type HookFactory, - type CustomAgentTool, + type CustomTool, } from "@mariozechner/pi-coding-agent"; // Set up auth storage (custom location) @@ -816,7 +817,7 @@ const auditHook: HookFactory = (api) => { }; // Inline tool -const statusTool: CustomAgentTool = { +const statusTool: CustomTool = { name: "status", label: "Status", description: "Get system status", @@ -932,7 +933,7 @@ createGrepTool, createFindTool, createLsTool // Types type CreateAgentSessionOptions type CreateAgentSessionResult -type CustomAgentTool +type CustomTool type HookFactory type Skill type FileSlashCommand diff --git a/packages/coding-agent/examples/hooks/confirm-destructive.ts b/packages/coding-agent/examples/hooks/confirm-destructive.ts index 75c5ee0c..ef189b23 100644 --- a/packages/coding-agent/examples/hooks/confirm-destructive.ts +++ b/packages/coding-agent/examples/hooks/confirm-destructive.ts @@ -45,7 +45,7 @@ export default function (pi: HookAPI) { pi.on("session_before_branch", async (event, ctx) => { if (!ctx.hasUI) return; - const choice = await ctx.ui.select(`Branch from turn ${event.entryIndex}?`, [ + const choice = await ctx.ui.select(`Branch from entry ${event.entryId.slice(0, 8)}?`, [ "Yes, create branch", "No, stay in current session", ]); diff --git a/packages/coding-agent/examples/hooks/git-checkpoint.ts b/packages/coding-agent/examples/hooks/git-checkpoint.ts index 87c8f0b5..6190be0d 100644 --- a/packages/coding-agent/examples/hooks/git-checkpoint.ts +++ b/packages/coding-agent/examples/hooks/git-checkpoint.ts @@ -8,19 +8,26 @@ import type { HookAPI } from "@mariozechner/pi-coding-agent/hooks"; export default function (pi: HookAPI) { - const checkpoints = new Map(); + const checkpoints = new Map(); + let currentEntryId: string | undefined; - pi.on("turn_start", async (event) => { + // Track the current entry ID when user messages are saved + pi.on("tool_result", async (_event, ctx) => { + const leaf = ctx.sessionManager.getLeafEntry(); + if (leaf) currentEntryId = leaf.id; + }); + + pi.on("turn_start", async () => { // Create a git stash entry before LLM makes changes const { stdout } = await pi.exec("git", ["stash", "create"]); const ref = stdout.trim(); - if (ref) { - checkpoints.set(event.turnIndex, ref); + if (ref && currentEntryId) { + checkpoints.set(currentEntryId, ref); } }); pi.on("session_before_branch", async (event, ctx) => { - const ref = checkpoints.get(event.entryIndex); + const ref = checkpoints.get(event.entryId); if (!ref) return; if (!ctx.hasUI) { diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index e2d63e09..3e5db9aa 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -1498,21 +1498,20 @@ export class AgentSession { } /** - * Create a branch from a specific entry index. + * Create a branch from a specific entry. * Emits before_branch/branch session events to hooks. * - * @param entryIndex Index into session entries to branch from + * @param entryId ID of the entry to branch from * @returns Object with: * - selectedText: The text of the selected user message (for editor pre-fill) * - cancelled: True if a hook cancelled the branch */ - async branch(entryIndex: number): Promise<{ selectedText: string; cancelled: boolean }> { + async branch(entryId: string): Promise<{ selectedText: string; cancelled: boolean }> { const previousSessionFile = this.sessionFile; - const entries = this.sessionManager.getEntries(); - const selectedEntry = entries[entryIndex]; + const selectedEntry = this.sessionManager.getEntry(entryId); if (!selectedEntry || selectedEntry.type !== "message" || selectedEntry.message.role !== "user") { - throw new Error("Invalid entry index for branching"); + throw new Error("Invalid entry ID for branching"); } const selectedText = this._extractUserMessageText(selectedEntry.message.content); @@ -1523,7 +1522,7 @@ export class AgentSession { if (this._hookRunner?.hasHandlers("session_before_branch")) { const result = (await this._hookRunner.emit({ type: "session_before_branch", - entryIndex: entryIndex, + entryId, })) as SessionBeforeBranchResult | undefined; if (result?.cancel) { @@ -1729,18 +1728,17 @@ export class AgentSession { /** * Get all user messages from session for branch selector. */ - getUserMessagesForBranching(): Array<{ entryIndex: number; text: string }> { + getUserMessagesForBranching(): Array<{ entryId: string; text: string }> { const entries = this.sessionManager.getEntries(); - const result: Array<{ entryIndex: number; text: string }> = []; + const result: Array<{ entryId: string; text: string }> = []; - for (let i = 0; i < entries.length; i++) { - const entry = entries[i]; + for (const entry of entries) { if (entry.type !== "message") continue; if (entry.message.role !== "user") continue; const text = this._extractUserMessageText(entry.message.content); if (text) { - result.push({ entryIndex: i, text }); + result.push({ entryId: entry.id, text }); } } diff --git a/packages/coding-agent/src/core/hooks/types.ts b/packages/coding-agent/src/core/hooks/types.ts index 879b111c..6acc843a 100644 --- a/packages/coding-agent/src/core/hooks/types.ts +++ b/packages/coding-agent/src/core/hooks/types.ts @@ -121,8 +121,8 @@ export interface SessionNewEvent { /** Fired before branching a session (can be cancelled) */ export interface SessionBeforeBranchEvent { type: "session_before_branch"; - /** Index of the entry in the session (SessionManager.getEntries()) to branch from */ - entryIndex: number; + /** ID of the entry to branch from */ + entryId: string; } /** Fired after branching a session */ diff --git a/packages/coding-agent/src/modes/interactive/components/user-message-selector.ts b/packages/coding-agent/src/modes/interactive/components/user-message-selector.ts index 18cd769c..8a8f2152 100644 --- a/packages/coding-agent/src/modes/interactive/components/user-message-selector.ts +++ b/packages/coding-agent/src/modes/interactive/components/user-message-selector.ts @@ -14,7 +14,7 @@ import { theme } from "../theme/theme.js"; import { DynamicBorder } from "./dynamic-border.js"; interface UserMessageItem { - index: number; // Index in the full messages array + id: string; // Entry ID in the session text: string; // The message text timestamp?: string; // Optional timestamp if available } @@ -25,7 +25,7 @@ interface UserMessageItem { class UserMessageList implements Component { private messages: UserMessageItem[] = []; private selectedIndex: number = 0; - public onSelect?: (messageIndex: number) => void; + public onSelect?: (entryId: string) => void; public onCancel?: () => void; private maxVisible: number = 10; // Max messages visible @@ -101,7 +101,7 @@ class UserMessageList implements Component { else if (isEnter(keyData)) { const selected = this.messages[this.selectedIndex]; if (selected && this.onSelect) { - this.onSelect(selected.index); + this.onSelect(selected.id); } } // Escape - cancel @@ -125,7 +125,7 @@ class UserMessageList implements Component { export class UserMessageSelectorComponent extends Container { private messageList: UserMessageList; - constructor(messages: UserMessageItem[], onSelect: (messageIndex: number) => void, onCancel: () => void) { + constructor(messages: UserMessageItem[], onSelect: (entryId: string) => void, onCancel: () => void) { super(); // Add header diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index 809eca88..42ef8be6 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -1570,9 +1570,9 @@ export class InteractiveMode { this.showSelector((done) => { const selector = new UserMessageSelectorComponent( - userMessages.map((m) => ({ index: m.entryIndex, text: m.text })), - async (entryIndex) => { - const result = await this.session.branch(entryIndex); + userMessages.map((m) => ({ id: m.entryId, text: m.text })), + async (entryId) => { + const result = await this.session.branch(entryId); if (result.cancelled) { // Hook cancelled the branch done(); diff --git a/packages/coding-agent/src/modes/rpc/rpc-client.ts b/packages/coding-agent/src/modes/rpc/rpc-client.ts index 0249ca11..7877ab36 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-client.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-client.ts @@ -326,17 +326,17 @@ export class RpcClient { * Branch from a specific message. * @returns Object with `text` (the message text) and `cancelled` (if hook cancelled) */ - async branch(entryIndex: number): Promise<{ text: string; cancelled: boolean }> { - const response = await this.send({ type: "branch", entryIndex }); + async branch(entryId: string): Promise<{ text: string; cancelled: boolean }> { + const response = await this.send({ type: "branch", entryId }); return this.getData(response); } /** * Get messages available for branching. */ - async getBranchMessages(): Promise> { + async getBranchMessages(): Promise> { const response = await this.send({ type: "get_branch_messages" }); - return this.getData<{ messages: Array<{ entryIndex: number; text: string }> }>(response).messages; + return this.getData<{ messages: Array<{ entryId: string; text: string }> }>(response).messages; } /** diff --git a/packages/coding-agent/src/modes/rpc/rpc-mode.ts b/packages/coding-agent/src/modes/rpc/rpc-mode.ts index e9fabf2a..09a0fde6 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-mode.ts @@ -347,7 +347,7 @@ export async function runRpcMode(session: AgentSession): Promise { } case "branch": { - const result = await session.branch(command.entryIndex); + const result = await session.branch(command.entryId); return success(id, "branch", { text: result.selectedText, cancelled: result.cancelled }); } diff --git a/packages/coding-agent/src/modes/rpc/rpc-types.ts b/packages/coding-agent/src/modes/rpc/rpc-types.ts index 5feead90..aa525687 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-types.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-types.ts @@ -53,7 +53,7 @@ export type RpcCommand = | { id?: string; type: "get_session_stats" } | { id?: string; type: "export_html"; outputPath?: string } | { id?: string; type: "switch_session"; sessionPath: string } - | { id?: string; type: "branch"; entryIndex: number } + | { id?: string; type: "branch"; entryId: string } | { id?: string; type: "get_branch_messages" } | { id?: string; type: "get_last_assistant_text" } @@ -150,7 +150,7 @@ export type RpcResponse = type: "response"; command: "get_branch_messages"; success: true; - data: { messages: Array<{ entryIndex: number; text: string }> }; + data: { messages: Array<{ entryId: string; text: string }> }; } | { id?: string; diff --git a/packages/coding-agent/test/agent-session-branching.test.ts b/packages/coding-agent/test/agent-session-branching.test.ts index 71b78ab7..33f70853 100644 --- a/packages/coding-agent/test/agent-session-branching.test.ts +++ b/packages/coding-agent/test/agent-session-branching.test.ts @@ -83,7 +83,7 @@ describe.skipIf(!API_KEY)("AgentSession branching", () => { expect(userMessages[0].text).toBe("Say hello"); // Branch from the first message - const result = await session.branch(userMessages[0].entryIndex); + const result = await session.branch(userMessages[0].entryId); expect(result.selectedText).toBe("Say hello"); expect(result.cancelled).toBe(false); @@ -113,7 +113,7 @@ describe.skipIf(!API_KEY)("AgentSession branching", () => { expect(session.messages.length).toBeGreaterThan(0); // Branch from the first message - const result = await session.branch(userMessages[0].entryIndex); + const result = await session.branch(userMessages[0].entryId); expect(result.selectedText).toBe("Say hi"); expect(result.cancelled).toBe(false); @@ -143,7 +143,7 @@ describe.skipIf(!API_KEY)("AgentSession branching", () => { // Branch from second message (keeps first message + response) const secondMessage = userMessages[1]; - const result = await session.branch(secondMessage.entryIndex); + const result = await session.branch(secondMessage.entryId); expect(result.selectedText).toBe("Say two"); // After branching, should have first user message + assistant response