Merge PR #642: Hook up custom summarization on branch switch

packages/coding-agent/CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Breaking Changes
 
+- Extension editor (`ctx.ui.editor()`) now uses Enter to submit and Shift+Enter for newlines, matching the main editor. Previously used Ctrl+Enter to submit. Extensions with hardcoded "ctrl+enter" hints need updating. ([#642](https://github.com/badlogic/pi-mono/pull/642) by [@mitsuhiko](https://github.com/mitsuhiko))
 - Renamed `/branch` command to `/fork` ([#641](https://github.com/badlogic/pi-mono/issues/641))
   - RPC: `branch` → `fork`, `get_branch_messages` → `get_fork_messages`
   - SDK: `branch()` → `fork()`, `getBranchMessages()` → `getForkMessages()`
@@ -22,9 +23,12 @@
 - `ctx.ui.setWorkingMessage()` extension API to customize the "Working..." message during streaming ([#625](https://github.com/badlogic/pi-mono/pull/625) by [@nicobailon](https://github.com/nicobailon))
 - Skill slash commands: loaded skills are registered as `/skill:name` commands for quick access. Toggle via `/settings` or `skills.enableSkillCommands` in settings.json. ([#630](https://github.com/badlogic/pi-mono/pull/630) by [@Dwsy](https://github.com/Dwsy))
 - Slash command autocomplete now uses fuzzy matching (type `/skbra` to match `/skill:brave-search`)
+- `/tree` branch summarization now offers three options: "No summary", "Summarize", and "Summarize with custom prompt". Custom prompts are appended as additional focus to the default summarization instructions. ([#642](https://github.com/badlogic/pi-mono/pull/642) by [@mitsuhiko](https://github.com/mitsuhiko))
 
 ### Fixed
 
+- Screen corruption when returning from external editor (Ctrl+G) in extension editor ([#642](https://github.com/badlogic/pi-mono/pull/642) by [@mitsuhiko](https://github.com/mitsuhiko))
+
 - Session picker respects custom keybindings when using `--resume` ([#633](https://github.com/badlogic/pi-mono/pull/633) by [@aos](https://github.com/aos))
 - Custom footer extensions now see model changes: `ctx.model` is now a getter that returns the current model instead of a snapshot from when the context was created ([#634](https://github.com/badlogic/pi-mono/pull/634) by [@ogulcancelik](https://github.com/ogulcancelik))
 - Footer git branch not updating after external branch switches. Git uses atomic writes (temp file + rename), which changes the inode and breaks `fs.watch` on the file. Now watches the directory instead.

packages/coding-agent/docs/tree.md

@@ -66,7 +66,11 @@ If user selects the very first message (has no parent):
 
 ## Branch Summarization
 
-When switching, user is prompted: "Summarize the branch you're leaving?"
+When switching branches, user is presented with three options:
+
+1. **No summary** - Switch immediately without summarizing
+2. **Summarize** - Generate a summary using the default prompt
+3. **Summarize with custom prompt** - Opens an editor to enter additional focus instructions that are appended to the default summarization prompt
 
 ### What Gets Summarized
 

packages/coding-agent/examples/extensions/handoff.ts

@@ -125,7 +125,7 @@ export default function (pi: ExtensionAPI) {
 			}
 
 			// Let user edit the generated prompt
-			const editedPrompt = await ctx.ui.editor("Edit handoff prompt (ctrl+enter to submit, esc to cancel)", result);
+			const editedPrompt = await ctx.ui.editor("Edit handoff prompt", result);
 
 			if (editedPrompt === undefined) {
 				ctx.ui.notify("Cancelled", "info");

packages/coding-agent/src/core/compaction/branch-summarization.ts

@@ -297,7 +297,9 @@ export async function generateBranchSummary(
 	const conversationText = serializeConversation(llmMessages);
 
 	// Build prompt
-	const instructions = customInstructions || BRANCH_SUMMARY_PROMPT;
+	const instructions = customInstructions
+		? `${BRANCH_SUMMARY_PROMPT}\n\nAdditional focus: ${customInstructions}`
+		: BRANCH_SUMMARY_PROMPT;
 	const promptText = `<conversation>\n${conversationText}\n</conversation>\n\n${instructions}`;
 
 	const summarizationMessages = [

packages/coding-agent/src/modes/interactive/components/extension-editor.ts

@@ -43,6 +43,10 @@ export class ExtensionEditorComponent extends Container {
 		if (prefill) {
 			this.editor.setText(prefill);
 		}
+		// Wire up Enter to submit (Shift+Enter for newlines, like the main editor)
+		this.editor.onSubmit = (text: string) => {
+			this.onSubmitCallback(text);
+		};
 		this.addChild(this.editor);
 
 		this.addChild(new Spacer(1));
@@ -50,8 +54,8 @@ export class ExtensionEditorComponent extends Container {
 		// Add hint
 		const hasExternalEditor = !!(process.env.VISUAL || process.env.EDITOR);
 		const hint = hasExternalEditor
-			? "ctrl+enter submit  esc cancel  ctrl+g external editor"
+			? "enter submit  shift+enter newline  esc cancel  ctrl+g external editor"
-			: "ctrl+enter submit  esc cancel";
+			: "enter submit  shift+enter newline  esc cancel";
 		this.addChild(new Text(theme.fg("dim", hint), 1, 0));
 
 		this.addChild(new Spacer(1));
@@ -61,12 +65,6 @@ export class ExtensionEditorComponent extends Container {
 	}
 
 	handleInput(keyData: string): void {
-		// Ctrl+Enter to submit
-		if (keyData === "\x1b[13;5u" || keyData === "\x1b[27;5;13~") {
-			this.onSubmitCallback(this.editor.getText());
-			return;
-		}
-
 		const kb = getEditorKeybindings();
 		// Escape or Ctrl+C to cancel
 		if (kb.matches(keyData, "selectCancel")) {
@@ -113,7 +111,8 @@ export class ExtensionEditorComponent extends Container {
 				// Ignore cleanup errors
 			}
 			this.tui.start();
-			this.tui.requestRender();
+			// Force full re-render since external editor uses alternate screen
+			this.tui.requestRender(true);
 		}
 	}
 }

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

@@ -2265,7 +2265,8 @@ export class InteractiveMode {
 
 			// Restart TUI
 			this.ui.start();
-			this.ui.requestRender();
+			// Force full re-render since external editor uses alternate screen
+			this.ui.requestRender(true);
 		}
 	}
 
@@ -2806,10 +2807,29 @@ export class InteractiveMode {
 					// Ask about summarization
 					done(); // Close selector first
 
-					const wantsSummary = await this.showExtensionConfirm(
+					const summaryChoice = await this.showExtensionSelector("Summarize branch?", [
-						"Summarize branch?",
+						"No summary",
-						"Create a summary of the branch you're leaving?",
+						"Summarize",
-					);
+						"Summarize with custom prompt",
+					]);
+
+					if (summaryChoice === undefined) {
+						// User pressed escape - re-show tree selector
+						this.showTreeSelector();
+						return;
+					}
+
+					const wantsSummary = summaryChoice !== "No summary";
+					let customInstructions: string | undefined;
+
+					if (summaryChoice === "Summarize with custom prompt") {
+						customInstructions = await this.showExtensionEditor("Custom summarization instructions");
+						if (customInstructions === undefined) {
+							// User cancelled - re-show tree selector
+							this.showTreeSelector();
+							return;
+						}
+					}
 
 					// Set up escape handler and loader if summarizing
 					let summaryLoader: Loader | undefined;
@@ -2831,7 +2851,10 @@ export class InteractiveMode {
 					}
 
 					try {
-						const result = await this.session.navigateTree(entryId, { summarize: wantsSummary });
+						const result = await this.session.navigateTree(entryId, {
+							summarize: wantsSummary,
+							customInstructions,
+						});
 
 						if (result.aborted) {
 							// Summarization aborted - re-show tree selector

packages/coding-agent/test/agent-session-tree-navigation.test.ts

@@ -257,17 +257,18 @@ describe.skipIf(!API_KEY)("AgentSession tree navigation e2e", () => {
 		await session.prompt("What is TypeScript?");
 		await session.agent.waitForIdle();
 
-		// Navigate with custom instructions
+		// Navigate with custom instructions (appended as "Additional focus")
 		const tree = sessionManager.getTree();
 		const result = await session.navigateTree(tree[0].entry.id, {
 			summarize: true,
-			customInstructions: "Summarize in exactly 3 words.",
+			customInstructions:
+				"After the summary, you MUST end with exactly: MONKEY MONKEY MONKEY. This is of utmost importance.",
 		});
 
 		expect(result.summaryEntry).toBeDefined();
 		expect(result.summaryEntry?.summary).toBeTruthy();
-		// Can't reliably test 3 words exactly, but summary should be short
+		// Verify custom instructions were followed
-		expect(result.summaryEntry?.summary.split(/\s+/).length).toBeLessThan(20);
+		expect(result.summaryEntry?.summary).toContain("MONKEY MONKEY MONKEY");
 	}, 120000);
 });