From a2a217277cbc955564c4c757310aadd073ddf494 Mon Sep 17 00:00:00 2001 From: roboomp Date: Thu, 25 Jun 2026 13:40:14 +0000 Subject: [PATCH] fix(tui): restore focused submit image draft on prompt error Snapshot pendingImageLinks and re-seed editor.pendingImages / pendingImageLinks / imageLinks in the focused-session submit catch block so an image-only or text+image draft survives a viewSession.prompt rejection, mirroring the main controller error path.\n\nFixes #3467 --- .../src/modes/controllers/input-controller.ts | 11 +- ...-controller-focused-submit-restore.test.ts | 100 ++++++++++++++++++ 2 files changed, 110 insertions(+), 1 deletion(-) create mode 100644 packages/coding-agent/test/input-controller-focused-submit-restore.test.ts diff --git a/packages/coding-agent/src/modes/controllers/input-controller.ts b/packages/coding-agent/src/modes/controllers/input-controller.ts index 80b925abc..a0301b3ae 100644 --- a/packages/coding-agent/src/modes/controllers/input-controller.ts +++ b/packages/coding-agent/src/modes/controllers/input-controller.ts @@ -839,6 +839,8 @@ export class InputController { async #submitToFocusedSession(text: string, streamingBehavior: "steer" | "followUp"): Promise { const target = this.ctx.viewSession; const images = this.ctx.editor.pendingImages.length > 0 ? [...this.ctx.editor.pendingImages] : undefined; + const imageLinks = + images && this.ctx.editor.pendingImageLinks.length > 0 ? [...this.ctx.editor.pendingImageLinks] : undefined; if (!text && !images) { if (target.isStreaming && target.queuedMessageCount > 0) { const aborting = target.abort({ reason: USER_INTERRUPT_LABEL }); @@ -859,7 +861,14 @@ export class InputController { imageCount: images?.length ?? 0, }); } catch (error) { - this.ctx.editor.setText(text); // hand the message back, mirroring the main submit error path + // Hand the message back, mirroring the main submit error path: restore + // pasted images so the user can retry an image-only or text+image draft. + this.ctx.editor.setText(text); + if (images && images.length > 0) { + this.ctx.editor.pendingImages = [...images]; + this.ctx.editor.pendingImageLinks = imageLinks ? [...imageLinks] : images.map(() => undefined); + this.ctx.editor.imageLinks = this.ctx.editor.pendingImageLinks; + } this.ctx.showError(error instanceof Error ? error.message : String(error)); } this.ctx.updatePendingMessagesDisplay(); diff --git a/packages/coding-agent/test/input-controller-focused-submit-restore.test.ts b/packages/coding-agent/test/input-controller-focused-submit-restore.test.ts new file mode 100644 index 000000000..464e396fa --- /dev/null +++ b/packages/coding-agent/test/input-controller-focused-submit-restore.test.ts @@ -0,0 +1,100 @@ +/** + * Regression: when an image-only or text+image submission is delivered to a + * focused subagent (`#submitToFocusedSession`) and `viewSession.prompt` + * rejects, the controller must restore both `text` AND `pendingImages` / + * `pendingImageLinks`. Previously only `text` was handed back, so the pasted + * image silently disappeared from the composer on retry. + */ +import { afterEach, describe, expect, it, vi } from "bun:test"; +import type { ImageContent } from "@oh-my-pi/pi-ai"; +import { InputController } from "@oh-my-pi/pi-coding-agent/modes/controllers/input-controller"; +import type { InteractiveModeContext } from "@oh-my-pi/pi-coding-agent/modes/types"; + +function createContext(opts: { pendingImages: ImageContent[]; pendingImageLinks?: (string | undefined)[] }) { + let editorText = ""; + const prompt = vi.fn(async () => { + throw new Error("focused dispatch rejected"); + }); + const showError = vi.fn(); + const updatePendingMessagesDisplay = vi.fn(); + const requestRender = vi.fn(); + + const editor = { + setText(text: string) { + editorText = text; + }, + getText() { + return editorText; + }, + addToHistory: vi.fn(), + imageLinks: undefined as (string | undefined)[] | undefined, + pendingImages: [...opts.pendingImages], + pendingImageLinks: + opts.pendingImageLinks !== undefined ? [...opts.pendingImageLinks] : opts.pendingImages.map(() => undefined), + clearDraft(historyText?: string) { + if (historyText !== undefined) this.addToHistory(historyText); + editorText = ""; + this.imageLinks = undefined; + this.pendingImages = []; + this.pendingImageLinks = []; + }, + }; + + const ctx = { + editor, + ui: { requestRender }, + session: { isStreaming: true, isCompacting: false, extensionRunner: undefined, queuedMessageCount: 0 }, + viewSession: { isStreaming: true, queuedMessageCount: 0, prompt, abort: vi.fn(async () => {}) }, + focusedAgentId: "Worker", + compactionQueuedMessages: [], + locallySubmittedUserSignatures: new Set(), + showError, + updatePendingMessagesDisplay, + withLocalSubmission: async (_text: string, fn: () => Promise) => fn(), + } as unknown as InteractiveModeContext; + + return { ctx, editor, prompt, showError }; +} + +describe("InputController focused submit restore-on-error", () => { + afterEach(() => { + vi.restoreAllMocks(); + }); + + it("restores text and pending images when focused prompt rejects", async () => { + const image: ImageContent = { type: "image", mimeType: "image/png", data: "aW1hZ2U=" }; + const { ctx, editor, prompt, showError } = createContext({ + pendingImages: [image], + pendingImageLinks: ["local://draft.png"], + }); + editor.setText("look at this"); + + const controller = new InputController(ctx); + controller.setupEditorSubmitHandler(); + await ctx.editor.onSubmit?.("look at this"); + + expect(prompt).toHaveBeenCalledTimes(1); + expect(showError).toHaveBeenCalledWith("focused dispatch rejected"); + expect(editor.getText()).toBe("look at this"); + expect(ctx.editor.pendingImages).toEqual([image]); + expect(ctx.editor.pendingImageLinks).toEqual(["local://draft.png"]); + expect(ctx.editor.imageLinks).toEqual(["local://draft.png"]); + }); + + it("restores image-only drafts when focused prompt rejects", async () => { + const image: ImageContent = { type: "image", mimeType: "image/png", data: "aW1hZ2U=" }; + const { ctx, editor, prompt, showError } = createContext({ pendingImages: [image] }); + editor.setText(""); + + const controller = new InputController(ctx); + controller.setupEditorSubmitHandler(); + await ctx.editor.onSubmit?.(""); + + expect(prompt).toHaveBeenCalledTimes(1); + expect(showError).toHaveBeenCalledWith("focused dispatch rejected"); + expect(editor.getText()).toBe(""); + expect(ctx.editor.pendingImages).toEqual([image]); + expect(ctx.editor.pendingImageLinks).toEqual([undefined]); + expect(ctx.editor.imageLinks).toEqual([undefined]); + }); +});