fix(tui): restored image attachments on /tree and esc-esc branch

- branch() and navigateTree() now return the selected user message's image
  parts (selectedImages/editorImages) alongside the text, extracted in marker
  order by #extractUserMessageImages.
- CustomEditor.setDraft() replaces the composer draft with text plus its
  pending images, so restored [Image #N] markers resolve on resubmit instead
  of degrading to literal text.
- Wired all six restore call sites (selector-controller, extension-ui-controller)
  through setDraft; updated rpc-subagents mocks for the new branch shape.
- Added offline regression tests for branch/navigateTree image restitution,
  multi-image marker order, and text-only prompts.

packages/coding-agent/CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+### Fixed
+
+- Restoring a prompt with image attachments via esc-esc branch or `/tree` now re-attaches the images to the composer draft: previously only the text (with its `[Image #N]` markers) was restored, so resubmitting sent the literal marker with no image.
+
 ## [17.1.7] - 2026-07-27
 
 ### Added

packages/coding-agent/src/modes/components/custom-editor.ts

@@ -427,6 +427,17 @@ export class CustomEditor extends Editor {
 		this.pendingImageLinks = [];
 	}
 
+	/** Replace the composer draft with a restored historical prompt: sets the text and
+	 *  re-attaches the message's images so positional `[Image #N]` markers resolve on
+	 *  resubmit instead of degrading to literal text (esc-esc branch, `/tree`). Source
+	 *  links are unknown for restored drafts, so every link slot is `undefined`. */
+	setDraft(text: string, images?: readonly ImageContent[]): void {
+		this.setText(text);
+		this.imageLinks = undefined;
+		this.pendingImages = images ? [...images] : [];
+		this.pendingImageLinks = images ? images.map(() => undefined) : [];
+	}
+
 	/** Treat image/paste markers as indivisible: a stray backspace deletes the whole token
 	 *  instead of corrupting `[Paste #1, +30 lines]` into plain text. */
 	override atomicTokenPattern = PLACEHOLDER_REGEX;

packages/coding-agent/src/modes/controllers/extension-ui-controller.ts

@@ -246,7 +246,7 @@ export class ExtensionUiController {
 				// Update UI
 				this.ctx.renderInitialMessages({ clearTerminalHistory: true });
 				await this.ctx.reloadTodos();
-				this.ctx.editor.setText(result.selectedText);
+				this.ctx.editor.setDraft(result.selectedText, result.selectedImages);
 				this.ctx.showStatus("Branched to new session");
 
 				return { cancelled: false };
@@ -261,7 +261,7 @@ export class ExtensionUiController {
 				this.ctx.renderInitialMessages({ clearTerminalHistory: true });
 				await this.ctx.reloadTodos();
 				if (result.editorText && !this.ctx.editor.getText().trim()) {
-					this.ctx.editor.setText(result.editorText);
+					this.ctx.editor.setDraft(result.editorText, result.editorImages);
 				}
 				this.ctx.showStatus("Navigated to selected point");
 
@@ -476,7 +476,7 @@ export class ExtensionUiController {
 				// Update UI
 				this.ctx.renderInitialMessages({ clearTerminalHistory: true });
 				await this.ctx.reloadTodos();
-				this.ctx.editor.setText(result.selectedText);
+				this.ctx.editor.setDraft(result.selectedText, result.selectedImages);
 				this.ctx.showStatus("Branched to new session");
 
 				return { cancelled: false };
@@ -491,7 +491,7 @@ export class ExtensionUiController {
 				this.ctx.renderInitialMessages({ clearTerminalHistory: true });
 				await this.ctx.reloadTodos();
 				if (result.editorText && !this.ctx.editor.getText().trim()) {
-					this.ctx.editor.setText(result.editorText);
+					this.ctx.editor.setDraft(result.editorText, result.editorImages);
 				}
 				this.ctx.showStatus("Navigated to selected point");
 

packages/coding-agent/src/modes/controllers/selector-controller.ts

@@ -1102,7 +1102,7 @@ export class SelectorController {
 					}
 
 					this.ctx.renderInitialMessages({ clearTerminalHistory: true });
-					this.ctx.editor.setText(result.selectedText);
+					this.ctx.editor.setDraft(result.selectedText, result.selectedImages);
 					done();
 					this.ctx.showStatus("Branched to new session");
 				},
@@ -1277,7 +1277,7 @@ export class SelectorController {
 						this.ctx.renderInitialMessages({ clearTerminalHistory: true });
 						await this.ctx.reloadTodos();
 						if (result.editorText && !this.ctx.editor.getText().trim()) {
-							this.ctx.editor.setText(result.editorText);
+							this.ctx.editor.setDraft(result.editorText, result.editorImages);
 						}
 						this.ctx.showStatus("Navigated to selected point");
 

packages/coding-agent/src/session/agent-session.ts

@@ -72,6 +72,7 @@ import type {
 	ToolChoice,
 	ToolResultMessage,
 	UsageReport,
+	UserMessage,
 } from "@oh-my-pi/pi-ai";
 import { type Effort, streamSimple } from "@oh-my-pi/pi-ai";
 import * as AIError from "@oh-my-pi/pi-ai/error";
@@ -7206,10 +7207,12 @@ export class AgentSession {
 	 * @param entryId ID of the entry to branch from
 	 * @returns Object with:
 	 *   - selectedText: The text of the selected user message (for editor pre-fill)
+	 *   - selectedImages: Image attachments of the selected user message (for editor draft restore)
 	 *   - cancelled: True if a hook cancelled the branch
 	 */
 	async branch(entryId: string): Promise<{
 		selectedText: string;
+		selectedImages: ImageContent[];
 		cancelled: boolean;
 	}> {
 		const previousSessionFile = this.sessionFile;
@@ -7220,6 +7223,7 @@ export class AgentSession {
 		}
 
 		const selectedText = this.#extractUserMessageText(selectedEntry.message.content);
+		const selectedImages = this.#extractUserMessageImages(selectedEntry.message.content);
 
 		let skipConversationRestore = false;
 
@@ -7231,7 +7235,7 @@ export class AgentSession {
 			})) as SessionBeforeBranchResult | undefined;
 
 			if (result?.cancel) {
-				return { selectedText, cancelled: true };
+				return { selectedText, selectedImages, cancelled: true };
 			}
 			skipConversationRestore = result?.skipConversationRestore ?? false;
 		}
@@ -7287,7 +7291,7 @@ export class AgentSession {
 			this.#closeCodexProviderSessionsForHistoryRewrite();
 		}
 
-		return { selectedText, cancelled: false };
+		return { selectedText, selectedImages, cancelled: false };
 	}
 
 	async branchFromBtw(
@@ -7402,7 +7406,7 @@ 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
-	 * @returns Result with editorText (if user message) and cancelled status
+	 * @returns Result with editorText/editorImages (if user message) and cancelled status
 	 */
 	async navigateTree(
 		targetId: string,
@@ -7432,6 +7436,8 @@ export class AgentSession {
 		} = {},
 	): Promise<{
 		editorText?: string;
+		/** Image attachments of the target user message, parallel to the positional `[Image #N]` markers in {@link editorText}. */
+		editorImages?: ImageContent[];
 		cancelled: boolean;
 		aborted?: boolean;
 		summaryEntry?: BranchSummaryEntry;
@@ -7604,6 +7610,7 @@ export class AgentSession {
 		// Determine the new leaf position based on target type
 		let newLeafId: string | null;
 		let editorText: string | undefined;
+		let editorImages: ImageContent[] | undefined;
 		// Set when the second-pass `ask` re-answer branch below actually commits a
 		// new sibling answer — the trigger for resuming the agent afterwards so the
 		// model consumes it, mirroring a live `ask` completion (issue #6483).
@@ -7613,6 +7620,8 @@ export class AgentSession {
 			// User message: leaf = parent (null if root), text goes to editor
 			newLeafId = targetEntry.parentId;
 			editorText = this.#extractUserMessageText(targetEntry.message.content);
+			const targetImages = this.#extractUserMessageImages(targetEntry.message.content);
+			if (targetImages.length > 0) editorImages = targetImages;
 		} else if (targetEntry.type === "custom_message" && targetEntry.customType !== SKILL_PROMPT_MESSAGE_TYPE) {
 			// Custom message: leaf = parent (null if root), text goes to editor
 			newLeafId = targetEntry.parentId;
@@ -7712,6 +7721,7 @@ export class AgentSession {
 			const rawContext = this.sessionManager.buildSessionContext();
 			return {
 				editorText,
+				editorImages,
 				cancelled: false,
 				summaryEntry,
 				sessionContext: rawContext,
@@ -7720,6 +7730,7 @@ export class AgentSession {
 		}
 		return {
 			editorText,
+			editorImages,
 			cancelled: false,
 			summaryEntry,
 			sessionContext: stateContext,
@@ -7838,6 +7849,14 @@ export class AgentSession {
 		return "";
 	}
 
+	/** Image parts of a stored user message, in submission order — index N-1 backs the
+	 *  `[Image #N]` marker in the message text, so restoring them alongside the text keeps
+	 *  positional markers resolvable on resubmit. */
+	#extractUserMessageImages(content: UserMessage["content"]): ImageContent[] {
+		if (!Array.isArray(content)) return [];
+		return content.filter((c): c is ImageContent => c.type === "image");
+	}
+
 	/**
 	 * Get session statistics.
 	 */

packages/coding-agent/test/agent-session-branching.test.ts

@@ -12,6 +12,7 @@ import * as fs from "node:fs";
 import * as os from "node:os";
 import * as path from "node:path";
 import { Agent } from "@oh-my-pi/pi-agent-core";
+import type { ImageContent, UserMessage } from "@oh-my-pi/pi-ai";
 import { getBundledModel } from "@oh-my-pi/pi-catalog/models";
 import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
 import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
@@ -20,7 +21,7 @@ import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage";
 import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager";
 import { createTools, type ToolSession } from "@oh-my-pi/pi-coding-agent/tools";
 import { removeSyncWithRetries, Snowflake } from "@oh-my-pi/pi-utils";
-import { e2eApiKey } from "./utilities";
+import { assistantMsg, createTestSession, e2eApiKey } from "./utilities";
 
 describe.skipIf(!e2eApiKey("ANTHROPIC_API_KEY"))("AgentSession branching", () => {
 	let session: AgentSession;
@@ -165,3 +166,96 @@ describe.skipIf(!e2eApiKey("ANTHROPIC_API_KEY"))("AgentSession branching", () =>
 		expect(session.messages[1].role).toBe("assistant");
 	});
 });
+
+const HISTORICAL_IMAGE: ImageContent = {
+	type: "image",
+	data: "aW1hZ2U=",
+	mimeType: "image/png",
+};
+
+function historicalImagePrompt(text: string): UserMessage {
+	return {
+		role: "user",
+		content: [{ type: "text", text }, HISTORICAL_IMAGE],
+		timestamp: Date.now(),
+	};
+}
+
+describe("AgentSession historical image prompts", () => {
+	it("returns the selected images when branching from a user prompt", async () => {
+		const ctx = await createTestSession({ inMemory: true });
+		try {
+			const text = "Inspect [Image #1, 1x1]";
+			const entryId = ctx.sessionManager.appendMessage(historicalImagePrompt(text));
+
+			const result = await ctx.session.branch(entryId);
+
+			expect(result).toEqual({
+				selectedText: text,
+				selectedImages: [HISTORICAL_IMAGE],
+				cancelled: false,
+			});
+		} finally {
+			await ctx.cleanup();
+		}
+	});
+
+	it("returns the target images when navigating to a user prompt", async () => {
+		const ctx = await createTestSession({ inMemory: true });
+		try {
+			const text = "Compare [Image #1, 1x1]";
+			const entryId = ctx.sessionManager.appendMessage(historicalImagePrompt(text));
+			ctx.sessionManager.appendMessage(assistantMsg("Compared."));
+
+			const result = await ctx.session.navigateTree(entryId);
+
+			expect(result).toMatchObject({
+				editorText: text,
+				editorImages: [HISTORICAL_IMAGE],
+				cancelled: false,
+			});
+		} finally {
+			await ctx.cleanup();
+		}
+	});
+
+	it("preserves multi-image order so positional markers stay aligned", async () => {
+		const ctx = await createTestSession({ inMemory: true });
+		try {
+			const second: ImageContent = { type: "image", data: "Qg==", mimeType: "image/jpeg" };
+			const text = "compare [Image #1, 1x1] with [Image #2, 2x2]";
+			const entryId = ctx.sessionManager.appendMessage({
+				role: "user",
+				content: [{ type: "text", text }, HISTORICAL_IMAGE, second],
+				timestamp: Date.now(),
+			} satisfies UserMessage);
+			ctx.sessionManager.appendMessage(assistantMsg("Compared."));
+
+			const result = await ctx.session.navigateTree(entryId);
+
+			expect(result.editorText).toBe(text);
+			expect(result.editorImages).toEqual([HISTORICAL_IMAGE, second]);
+		} finally {
+			await ctx.cleanup();
+		}
+	});
+
+	it("leaves editorImages unset when the prompt has no images", async () => {
+		const ctx = await createTestSession({ inMemory: true });
+		try {
+			const entryId = ctx.sessionManager.appendMessage({
+				role: "user",
+				content: "plain text turn",
+				timestamp: Date.now(),
+			} satisfies UserMessage);
+			ctx.sessionManager.appendMessage(assistantMsg("ok"));
+
+			const result = await ctx.session.navigateTree(entryId);
+
+			expect(result.editorText).toBe("plain text turn");
+			expect(result.editorImages).toBeUndefined();
+		} finally {
+			await ctx.cleanup();
+		}
+	});
+});

packages/coding-agent/test/modes/components/custom-editor.test.ts

@@ -1,4 +1,5 @@
 import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "bun:test";
+import type { ImageContent } from "@oh-my-pi/pi-ai";
 import { CURSOR_MARKER } from "@oh-my-pi/pi-tui";
 import { setKittyProtocolActive } from "@oh-my-pi/pi-tui/keys";
 import { $ } from "bun";
@@ -80,6 +81,33 @@ describe("CustomEditor placeholder decoration", () => {
 	});
 });
 
+describe("CustomEditor restored image drafts", () => {
+	beforeAll(async () => {
+		await initTheme();
+	});
+
+	it("submits restored images with their historical prompt", () => {
+		const editor = new CustomEditor(getEditorTheme());
+		const image: ImageContent = {
+			type: "image",
+			data: "aW1hZ2U=",
+			mimeType: "image/png",
+		};
+		let submitted: { text: string; images: ImageContent[] } | undefined;
+		editor.onSubmit = text => {
+			submitted = { text, images: [...editor.pendingImages] };
+		};
+
+		editor.setDraft("Inspect [Image #1, 1x1]", [image]);
+		editor.submit();
+
+		expect(submitted).toEqual({
+			text: "Inspect [Image #1, 1x1]",
+			images: [image],
+		});
+	});
+});
+
 describe("CustomEditor queue shorthand decoration", () => {
 	beforeAll(async () => {
 		await initTheme();

packages/coding-agent/test/rpc-subagents.test.ts

@@ -2,6 +2,7 @@ import { afterEach, describe, expect, test } from "bun:test";
 import * as fs from "node:fs";
 import * as os from "node:os";
 import * as path from "node:path";
+import type { ImageContent } from "@oh-my-pi/pi-ai";
 import { RpcClient } from "@oh-my-pi/pi-coding-agent/modes/rpc/rpc-client";
 import {
 	handleRpcSessionChange,
@@ -70,14 +71,15 @@ function createRegistryWithSnapshot(): RpcSubagentRegistry {
 type SessionChangeStubOptions = {
 	newSession?: boolean;
 	switchSession?: boolean;
-	branch?: { selectedText: string; cancelled: boolean };
+	branch?: { selectedText: string; selectedImages: ImageContent[]; cancelled: boolean };
 };
 
 function createSessionChangeSession(options: SessionChangeStubOptions): RpcSessionChangeSession {
 	return {
 		newSession: async (_options?: unknown) => options.newSession ?? true,
 		switchSession: async (_sessionPath: string) => options.switchSession ?? true,
-		branch: async (_entryId: string) => options.branch ?? { selectedText: "branched text", cancelled: false },
+		branch: async (_entryId: string) =>
+			options.branch ?? { selectedText: "branched text", selectedImages: [], cancelled: false },
 	};
 }
 
@@ -208,7 +210,9 @@ describe("RPC subagent registry", () => {
 			},
 			{
 				command: { type: "branch", entryId: "entry-1" },
-				session: createSessionChangeSession({ branch: { selectedText: "Branch text", cancelled: false } }),
+				session: createSessionChangeSession({
+					branch: { selectedText: "Branch text", selectedImages: [], cancelled: false },
+				}),
 				expected: { type: "branch", data: { text: "Branch text", cancelled: false } },
 			},
 		];
@@ -247,7 +251,7 @@ describe("RPC subagent registry", () => {
 			},
 			{
 				command: { type: "branch", entryId: "entry-1" },
-				session: createSessionChangeSession({ branch: { selectedText: "", cancelled: true } }),
+				session: createSessionChangeSession({ branch: { selectedText: "", selectedImages: [], cancelled: true } }),
 				expected: { type: "branch", data: { text: "", cancelled: true } },
 			},
 		];