From 039c93be60467ce97cdf9862473ce5bf938cb7c9 Mon Sep 17 00:00:00 2001 From: roboomp Date: Thu, 25 Jun 2026 13:53:00 +0000 Subject: [PATCH] fix(tui): restore streaming steer image draft on prompt error Wrap the streaming Enter steer dispatch in try/catch so prompt failures restore text plus pendingImages / pendingImageLinks / imageLinks and surface showError instead of losing an image-only draft.\n\nAlso remove a forbidden ReturnType<> from the follow-up image regression helper.\n\nFixes #3467 --- .../src/modes/controllers/input-controller.ts | 24 ++++++++++--- .../input-controller-followup-image.test.ts | 5 +-- ...interrupt-and-flush-empty-messages.test.ts | 36 +++++++++++++++++-- 3 files changed, 56 insertions(+), 9 deletions(-) diff --git a/packages/coding-agent/src/modes/controllers/input-controller.ts b/packages/coding-agent/src/modes/controllers/input-controller.ts index 2fb841324..bbebe009a 100644 --- a/packages/coding-agent/src/modes/controllers/input-controller.ts +++ b/packages/coding-agent/src/modes/controllers/input-controller.ts @@ -719,11 +719,25 @@ export class InputController { // (a user-role `message_start` event) leaves any draft the user has // typed since queuing intact. Same protection as #783, applied to // the streaming/queue path. - await this.ctx.withLocalSubmission( - text, - () => this.ctx.session.prompt(text, { streamingBehavior: "steer", images }), - { imageCount: images?.length ?? 0 }, - ); + try { + await this.ctx.withLocalSubmission( + text, + () => this.ctx.session.prompt(text, { streamingBehavior: "steer", images }), + { imageCount: images?.length ?? 0 }, + ); + } catch (error) { + // Don't lose the queued steer draft: restore text and images so + // the user can retry after dispatch validation/queue failures. + this.ctx.editor.setText(text); + if (images && images.length > 0) { + this.ctx.editor.pendingImages = [...images]; + this.ctx.editor.pendingImageLinks = inputImageLinks + ? [...inputImageLinks] + : images.map(() => undefined); + this.ctx.editor.imageLinks = this.ctx.editor.pendingImageLinks; + } + this.ctx.showError(error instanceof Error ? error.message : String(error)); + } this.ctx.updatePendingMessagesDisplay(); this.ctx.ui.requestRender(); return; diff --git a/packages/coding-agent/test/input-controller-followup-image.test.ts b/packages/coding-agent/test/input-controller-followup-image.test.ts index dc45d2acf..4982e104d 100644 --- a/packages/coding-agent/test/input-controller-followup-image.test.ts +++ b/packages/coding-agent/test/input-controller-followup-image.test.ts @@ -51,6 +51,7 @@ function createContext(opts: { const prompt = vi.fn(async (_text: string, _options?: PromptOptionsLike) => {}); const updatePendingMessagesDisplay = vi.fn(); const requestRender = vi.fn(); + const showError = vi.fn(); const ctx = { editor, @@ -68,11 +69,11 @@ function createContext(opts: { compactionQueuedMessages: [], locallySubmittedUserSignatures: new Set(), updatePendingMessagesDisplay, - showError: vi.fn(), + showError, withLocalSubmission: async (_text: string, fn: () => unknown) => fn(), } as unknown as InteractiveModeContext; - return { ctx, editor, prompt, showError: ctx.showError as ReturnType }; + return { ctx, editor, prompt, showError }; } describe("InputController.handleFollowUp image forwarding", () => { diff --git a/packages/coding-agent/test/issue-interrupt-and-flush-empty-messages.test.ts b/packages/coding-agent/test/issue-interrupt-and-flush-empty-messages.test.ts index 0cd4325c2..875bb99a3 100644 --- a/packages/coding-agent/test/issue-interrupt-and-flush-empty-messages.test.ts +++ b/packages/coding-agent/test/issue-interrupt-and-flush-empty-messages.test.ts @@ -4,7 +4,11 @@ import { InputController } from "@oh-my-pi/pi-coding-agent/modes/controllers/inp import type { InteractiveModeContext } from "@oh-my-pi/pi-coding-agent/modes/types"; import { USER_INTERRUPT_LABEL } from "@oh-my-pi/pi-coding-agent/session/messages"; -function createContext(options?: { queuedMessageCount?: number; pendingImages?: ImageContent[] }) { +function createContext(options?: { + queuedMessageCount?: number; + pendingImages?: ImageContent[]; + pendingImageLinks?: (string | undefined)[]; +}) { let editorText = ""; const abort = vi.fn(async () => {}); const prompt = vi.fn(async () => {}); @@ -22,7 +26,10 @@ function createContext(options?: { queuedMessageCount?: number; pendingImages?: }, addToHistory: vi.fn(), pendingImages: options?.pendingImages ? [...options.pendingImages] : ([] as ImageContent[]), - pendingImageLinks: options?.pendingImages?.map(() => undefined) ?? ([] as (string | undefined)[]), + pendingImageLinks: + options?.pendingImageLinks ?? + options?.pendingImages?.map(() => undefined) ?? + ([] as (string | undefined)[]), }, ui: { requestRender }, session: { @@ -86,6 +93,31 @@ describe("empty submit with queued messages", () => { expect(requestRender).toHaveBeenCalledTimes(1); }); + it("restores an image-only steer when streaming dispatch rejects", async () => { + const image: ImageContent = { type: "image", mimeType: "image/png", data: "aW1hZ2U=" }; + const { ctx, abort, prompt, showError, updatePendingMessagesDisplay, requestRender } = createContext({ + queuedMessageCount: 0, + pendingImages: [image], + pendingImageLinks: ["local://draft.png"], + }); + prompt.mockImplementationOnce(async () => { + throw new Error("queue rejected"); + }); + const controller = new InputController(ctx); + controller.setupEditorSubmitHandler(); + + await ctx.editor.onSubmit?.(""); + + expect(abort).not.toHaveBeenCalled(); + expect(showError).toHaveBeenCalledWith("queue rejected"); + expect(ctx.editor.getText()).toBe(""); + expect(ctx.editor.pendingImages).toEqual([image]); + expect(ctx.editor.pendingImageLinks).toEqual(["local://draft.png"]); + expect(ctx.editor.imageLinks).toEqual(["local://draft.png"]); + expect(updatePendingMessagesDisplay).toHaveBeenCalledTimes(1); + expect(requestRender).toHaveBeenCalledTimes(1); + }); + it("queues an image-only steer instead of aborting when messages are already queued", async () => { const image: ImageContent = { type: "image", mimeType: "image/png", data: "aW1hZ2U=" }; const { ctx, abort, prompt } = createContext({ queuedMessageCount: 1, pendingImages: [image] });