From 6b6707f30cbfdddaccce490e7cc4a3e7a39250b0 Mon Sep 17 00:00:00 2001 From: Armin Ronacher Date: Fri, 16 Jan 2026 19:52:54 +0100 Subject: [PATCH 1/2] Improve navigateTree API --- .../coding-agent/src/core/agent-session.ts | 37 +++++++++++++++++-- .../core/compaction/branch-summarization.ts | 15 ++++++-- .../src/core/extensions/runner.ts | 2 +- .../coding-agent/src/core/extensions/types.ts | 17 ++++++++- .../src/modes/interactive/interactive-mode.ts | 7 +++- 5 files changed, 68 insertions(+), 10 deletions(-) diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 61f9eba9..63d6b6db 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -1972,11 +1972,13 @@ export class AgentSession { * @param targetId The entry ID to navigate to * @param options.summarize Whether user wants to summarize abandoned branch * @param options.customInstructions Custom instructions for summarizer + * @param options.replaceInstructions If true, customInstructions replaces the default prompt + * @param options.label Label to attach to the branch summary entry * @returns Result with editorText (if user message) and cancelled status */ async navigateTree( targetId: string, - options: { summarize?: boolean; customInstructions?: string } = {}, + options: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string } = {}, ): Promise<{ editorText?: string; cancelled: boolean; aborted?: boolean; summaryEntry?: BranchSummaryEntry }> { const oldLeafId = this.sessionManager.getLeafId(); @@ -2002,13 +2004,20 @@ export class AgentSession { targetId, ); - // Prepare event data + // Prepare event data - mutable so extensions can override + let customInstructions = options.customInstructions; + let replaceInstructions = options.replaceInstructions; + let label = options.label; + const preparation: TreePreparation = { targetId, oldLeafId, commonAncestorId, entriesToSummarize, userWantsSummary: options.summarize ?? false, + customInstructions, + replaceInstructions, + label, }; // Set up abort controller for summarization @@ -2032,6 +2041,17 @@ export class AgentSession { extensionSummary = result.summary; fromExtension = true; } + + // Allow extensions to override instructions and label + if (result?.customInstructions !== undefined) { + customInstructions = result.customInstructions; + } + if (result?.replaceInstructions !== undefined) { + replaceInstructions = result.replaceInstructions; + } + if (result?.label !== undefined) { + label = result.label; + } } // Run default summarizer if needed @@ -2048,7 +2068,8 @@ export class AgentSession { model, apiKey, signal: this._branchSummaryAbortController.signal, - customInstructions: options.customInstructions, + customInstructions, + replaceInstructions, reserveTokens: branchSummarySettings.reserveTokens, }); this._branchSummaryAbortController = undefined; @@ -2098,6 +2119,11 @@ export class AgentSession { // Create summary at target position (can be null for root) const summaryId = this.sessionManager.branchWithSummary(newLeafId, summaryText, summaryDetails, fromExtension); summaryEntry = this.sessionManager.getEntry(summaryId) as BranchSummaryEntry; + + // Attach label to the summary entry + if (label) { + this.sessionManager.appendLabelChange(summaryId, label); + } } else if (newLeafId === null) { // No summary, navigating to root - reset leaf this.sessionManager.resetLeaf(); @@ -2106,6 +2132,11 @@ export class AgentSession { this.sessionManager.branch(newLeafId); } + // Attach label to target entry when not summarizing (no summary entry to label) + if (label && !summaryText) { + this.sessionManager.appendLabelChange(targetId, label); + } + // Update agent state const sessionContext = this.sessionManager.buildSessionContext(); this.agent.replaceMessages(sessionContext.messages); diff --git a/packages/coding-agent/src/core/compaction/branch-summarization.ts b/packages/coding-agent/src/core/compaction/branch-summarization.ts index f66cb80c..5e3b1896 100644 --- a/packages/coding-agent/src/core/compaction/branch-summarization.ts +++ b/packages/coding-agent/src/core/compaction/branch-summarization.ts @@ -71,6 +71,8 @@ export interface GenerateBranchSummaryOptions { signal: AbortSignal; /** Optional custom instructions for summarization */ customInstructions?: string; + /** If true, customInstructions replaces the default prompt instead of being appended */ + replaceInstructions?: boolean; /** Tokens reserved for prompt + LLM response (default 16384) */ reserveTokens?: number; } @@ -279,7 +281,7 @@ export async function generateBranchSummary( entries: SessionEntry[], options: GenerateBranchSummaryOptions, ): Promise { - const { model, apiKey, signal, customInstructions, reserveTokens = 16384 } = options; + const { model, apiKey, signal, customInstructions, replaceInstructions, reserveTokens = 16384 } = options; // Token budget = context window minus reserved space for prompt + response const contextWindow = model.contextWindow || 128000; @@ -297,9 +299,14 @@ export async function generateBranchSummary( const conversationText = serializeConversation(llmMessages); // Build prompt - const instructions = customInstructions - ? `${BRANCH_SUMMARY_PROMPT}\n\nAdditional focus: ${customInstructions}` - : BRANCH_SUMMARY_PROMPT; + let instructions: string; + if (replaceInstructions && customInstructions) { + instructions = customInstructions; + } else if (customInstructions) { + instructions = `${BRANCH_SUMMARY_PROMPT}\n\nAdditional focus: ${customInstructions}`; + } else { + instructions = BRANCH_SUMMARY_PROMPT; + } const promptText = `\n${conversationText}\n\n\n${instructions}`; const summarizationMessages = [ diff --git a/packages/coding-agent/src/core/extensions/runner.ts b/packages/coding-agent/src/core/extensions/runner.ts index 117f89c5..4fa869ae 100644 --- a/packages/coding-agent/src/core/extensions/runner.ts +++ b/packages/coding-agent/src/core/extensions/runner.ts @@ -57,7 +57,7 @@ export type ForkHandler = (entryId: string) => Promise<{ cancelled: boolean }>; export type NavigateTreeHandler = ( targetId: string, - options?: { summarize?: boolean }, + options?: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string }, ) => Promise<{ cancelled: boolean }>; export type ShutdownHandler = () => void; diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index 0c2f656c..b122465c 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -344,6 +344,12 @@ export interface TreePreparation { commonAncestorId: string | null; entriesToSummarize: SessionEntry[]; userWantsSummary: boolean; + /** Custom instructions for summarization */ + customInstructions?: string; + /** If true, customInstructions replaces the default prompt instead of being appended */ + replaceInstructions?: boolean; + /** Label to attach to the branch summary entry */ + label?: string; } /** Fired before navigating in the session tree (can be cancelled) */ @@ -633,6 +639,12 @@ export interface SessionBeforeTreeResult { summary: string; details?: unknown; }; + /** Override custom instructions for summarization */ + customInstructions?: string; + /** Override whether customInstructions replaces the default prompt */ + replaceInstructions?: boolean; + /** Override label to attach to the branch summary entry */ + label?: string; } // ============================================================================ @@ -917,7 +929,10 @@ export interface ExtensionCommandContextActions { setup?: (sessionManager: SessionManager) => Promise; }) => Promise<{ cancelled: boolean }>; fork: (entryId: string) => Promise<{ cancelled: boolean }>; - navigateTree: (targetId: string, options?: { summarize?: boolean }) => Promise<{ cancelled: boolean }>; + navigateTree: ( + targetId: string, + options?: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string }, + ) => Promise<{ cancelled: boolean }>; } /** diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index 4a4cdc5b..a02a969c 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -717,7 +717,12 @@ export class InteractiveMode { return { cancelled: false }; }, navigateTree: async (targetId, options) => { - const result = await this.session.navigateTree(targetId, { summarize: options?.summarize }); + const result = await this.session.navigateTree(targetId, { + summarize: options?.summarize, + customInstructions: options?.customInstructions, + replaceInstructions: options?.replaceInstructions, + label: options?.label, + }); if (result.cancelled) { return { cancelled: true }; } From 572ec64dbdba60c05010109c0bd9e9a787085ea9 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Fri, 16 Jan 2026 21:47:01 +0100 Subject: [PATCH 2/2] Fix navigateTree API: add missing type updates, handler passthrough, and docs - Update ExtensionCommandContext.navigateTree type signature - Pass new options through in print-mode and rpc-mode handlers - Update docs/extensions.md, docs/sdk.md, docs/tree.md - Add changelog entry --- packages/coding-agent/CHANGELOG.md | 1 + packages/coding-agent/docs/extensions.md | 9 ++++++++ packages/coding-agent/docs/sdk.md | 2 +- packages/coding-agent/docs/tree.md | 22 +++++++++++++++++-- .../coding-agent/src/core/extensions/types.ts | 5 ++++- packages/coding-agent/src/modes/print-mode.ts | 7 +++++- .../coding-agent/src/modes/rpc/rpc-mode.ts | 7 +++++- 7 files changed, 47 insertions(+), 6 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index bca2129a..cca30e6d 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -9,6 +9,7 @@ - Bash tool now displays the timeout value in the UI when a timeout is set ([#780](https://github.com/badlogic/pi-mono/pull/780) by [@dannote](https://github.com/dannote)) - Export `getShellConfig` for extensions to detect user's shell environment ([#766](https://github.com/badlogic/pi-mono/pull/766) by [@dannote](https://github.com/dannote)) - Added `thinkingText` and `selectedBg` to theme schema ([#763](https://github.com/badlogic/pi-mono/pull/763) by [@scutifer](https://github.com/scutifer)) +- `navigateTree()` now supports `replaceInstructions` option to replace the default summarization prompt entirely, and `label` option to attach a label to the branch summary entry ([#787](https://github.com/badlogic/pi-mono/pull/787) by [@mitsuhiko](https://github.com/mitsuhiko)) ### Fixed diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index 5c80694b..7ef92582 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -733,9 +733,18 @@ Navigate to a different point in the session tree: ```typescript const result = await ctx.navigateTree("entry-id-456", { summarize: true, + customInstructions: "Focus on error handling changes", + replaceInstructions: false, // true = replace default prompt entirely + label: "review-checkpoint", }); ``` +Options: +- `summarize`: Whether to generate a summary of the abandoned branch +- `customInstructions`: Custom instructions for the summarizer +- `replaceInstructions`: If true, `customInstructions` replaces the default prompt instead of being appended +- `label`: Label to attach to the branch summary entry (or target entry if not summarizing) + ## ExtensionAPI Methods ### pi.on(event, handler) diff --git a/packages/coding-agent/docs/sdk.md b/packages/coding-agent/docs/sdk.md index d231b291..a2052783 100644 --- a/packages/coding-agent/docs/sdk.md +++ b/packages/coding-agent/docs/sdk.md @@ -110,7 +110,7 @@ interface AgentSession { // Forking fork(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 + navigateTree(targetId: string, options?: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string }): Promise<{ editorText?: string; cancelled: boolean }>; // In-place navigation // Hook message injection sendHookMessage(message: HookMessage, triggerTurn?: boolean): Promise; diff --git a/packages/coding-agent/docs/tree.md b/packages/coding-agent/docs/tree.md index 3d2604d2..23a0b12c 100644 --- a/packages/coding-agent/docs/tree.md +++ b/packages/coding-agent/docs/tree.md @@ -110,10 +110,21 @@ interface BranchSummaryEntry { ```typescript async navigateTree( targetId: string, - options?: { summarize?: boolean; customInstructions?: string } + options?: { + summarize?: boolean; + customInstructions?: string; + replaceInstructions?: boolean; + label?: string; + } ): Promise<{ editorText?: string; cancelled: boolean }> ``` +Options: +- `summarize`: Whether to generate a summary of the abandoned branch +- `customInstructions`: Custom instructions for the summarizer +- `replaceInstructions`: If true, `customInstructions` replaces the default prompt instead of being appended +- `label`: Label to attach to the branch summary entry (or target entry if not summarizing) + Flow: 1. Validate target, check no-op (target === current leaf) 2. Find common ancestor between old leaf and target @@ -153,21 +164,28 @@ interface TreePreparation { commonAncestorId: string | null; entriesToSummarize: SessionEntry[]; userWantsSummary: boolean; + customInstructions?: string; + replaceInstructions?: boolean; + label?: string; } interface SessionBeforeTreeEvent { type: "session_before_tree"; preparation: TreePreparation; - model: Model; signal: AbortSignal; } interface SessionBeforeTreeResult { cancel?: boolean; summary?: { summary: string; details?: unknown }; + customInstructions?: string; // Override custom instructions + replaceInstructions?: boolean; // Override replace mode + label?: string; // Override label } ``` +Extensions can override `customInstructions`, `replaceInstructions`, and `label` by returning them from the `session_before_tree` handler. + ### `session_tree` ```typescript diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index b122465c..1753576a 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -237,7 +237,10 @@ export interface ExtensionCommandContext extends ExtensionContext { fork(entryId: string): Promise<{ cancelled: boolean }>; /** Navigate to a different point in the session tree. */ - navigateTree(targetId: string, options?: { summarize?: boolean }): Promise<{ cancelled: boolean }>; + navigateTree( + targetId: string, + options?: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string }, + ): Promise<{ cancelled: boolean }>; } // ============================================================================ diff --git a/packages/coding-agent/src/modes/print-mode.ts b/packages/coding-agent/src/modes/print-mode.ts index a2b0081e..2a9a5277 100644 --- a/packages/coding-agent/src/modes/print-mode.ts +++ b/packages/coding-agent/src/modes/print-mode.ts @@ -95,7 +95,12 @@ export async function runPrintMode(session: AgentSession, options: PrintModeOpti return { cancelled: result.cancelled }; }, navigateTree: async (targetId, options) => { - const result = await session.navigateTree(targetId, { summarize: options?.summarize }); + const result = await session.navigateTree(targetId, { + summarize: options?.summarize, + customInstructions: options?.customInstructions, + replaceInstructions: options?.replaceInstructions, + label: options?.label, + }); return { cancelled: result.cancelled }; }, }, diff --git a/packages/coding-agent/src/modes/rpc/rpc-mode.ts b/packages/coding-agent/src/modes/rpc/rpc-mode.ts index dc8a7a4a..50c4ed3c 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-mode.ts @@ -311,7 +311,12 @@ export async function runRpcMode(session: AgentSession): Promise { return { cancelled: result.cancelled }; }, navigateTree: async (targetId, options) => { - const result = await session.navigateTree(targetId, { summarize: options?.summarize }); + const result = await session.navigateTree(targetId, { + summarize: options?.summarize, + customInstructions: options?.customInstructions, + replaceInstructions: options?.replaceInstructions, + label: options?.label, + }); return { cancelled: result.cancelled }; }, },