From cfd335d2b1241f44d0ce1b40b29f2716e15ac6df Mon Sep 17 00:00:00 2001 From: can1357 Date: Tue, 28 Jul 2026 00:57:34 +0200 Subject: [PATCH] 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 | 4 + .../src/modes/components/custom-editor.ts | 11 +++ .../controllers/extension-ui-controller.ts | 8 +- .../modes/controllers/selector-controller.ts | 4 +- .../coding-agent/src/session/agent-session.ts | 25 ++++- .../test/agent-session-branching.test.ts | 96 ++++++++++++++++++- .../modes/components/custom-editor.test.ts | 28 ++++++ .../coding-agent/test/rpc-subagents.test.ts | 12 ++- 8 files changed, 174 insertions(+), 14 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index a507a8391..ec219e9cb 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/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 diff --git a/packages/coding-agent/src/modes/components/custom-editor.ts b/packages/coding-agent/src/modes/components/custom-editor.ts index 2c2aa2132..96c53f5d2 100644 --- a/packages/coding-agent/src/modes/components/custom-editor.ts +++ b/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; diff --git a/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts b/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts index 21014ced4..750bdd193 100644 --- a/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts +++ b/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"); diff --git a/packages/coding-agent/src/modes/controllers/selector-controller.ts b/packages/coding-agent/src/modes/controllers/selector-controller.ts index b3b7ab8e9..b4b64dbc5 100644 --- a/packages/coding-agent/src/modes/controllers/selector-controller.ts +++ b/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"); diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index 385a10f9c..ff6376d8c 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/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. */ diff --git a/packages/coding-agent/test/agent-session-branching.test.ts b/packages/coding-agent/test/agent-session-branching.test.ts index 7b915b7e4..ea9f25598 100644 --- a/packages/coding-agent/test/agent-session-branching.test.ts +++ b/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(); + } + }); +}); diff --git a/packages/coding-agent/test/modes/components/custom-editor.test.ts b/packages/coding-agent/test/modes/components/custom-editor.test.ts index 742031718..f77ffe850 100644 --- a/packages/coding-agent/test/modes/components/custom-editor.test.ts +++ b/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(); diff --git a/packages/coding-agent/test/rpc-subagents.test.ts b/packages/coding-agent/test/rpc-subagents.test.ts index b021107f8..81b272885 100644 --- a/packages/coding-agent/test/rpc-subagents.test.ts +++ b/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 } }, }, ];