Merge PR #787: Improve navigateTree API

Adds replaceInstructions and label options to navigateTree() for custom summarization prompts and branch labeling.

closes #787

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
 

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)

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<void>;

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

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);

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<BranchSummaryResult> {
-	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 = `<conversation>\n${conversationText}\n</conversation>\n\n${instructions}`;
 
 	const summarizationMessages = [

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;

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 }>;
 }
 
 // ============================================================================
@@ -344,6 +347,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 +642,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 +932,10 @@ export interface ExtensionCommandContextActions {
 		setup?: (sessionManager: SessionManager) => Promise<void>;
 	}) => 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 }>;
 }
 
 /**

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 };
 					}

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 };
 				},
 			},

packages/coding-agent/src/modes/rpc/rpc-mode.ts

@@ -311,7 +311,12 @@ export async function runRpcMode(session: AgentSession): Promise<never> {
 					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 };
 				},
 			},