diff --git a/packages/coding-agent/src/modes/controllers/input-controller.ts b/packages/coding-agent/src/modes/controllers/input-controller.ts index 89170c6d6..7628c3ab4 100644 --- a/packages/coding-agent/src/modes/controllers/input-controller.ts +++ b/packages/coding-agent/src/modes/controllers/input-controller.ts @@ -23,6 +23,7 @@ import type { InteractiveModeContext } from "../../modes/types"; import manualContinuePrompt from "../../prompts/system/manual-continue.md" with { type: "text" }; import { USER_INTERRUPT_LABEL } from "../../session/messages"; import { executeBuiltinSlashCommand } from "../../slash-commands/builtin-registry"; +import { parseSlashCommand } from "../../slash-commands/helpers/parse"; import { isTinyTitleLocalModelKey } from "../../tiny/models"; import { tinyTitleClient } from "../../tiny/title-client"; import type { TinyTitleProgressEvent } from "../../tiny/title-protocol"; @@ -678,11 +679,14 @@ export class InputController { let inputImageLinks = this.ctx.editor.pendingImageLinks.length > 0 ? [...this.ctx.editor.pendingImageLinks] : undefined; let hasInputImages = (inputImages?.length ?? 0) > 0; + const submittedMode = parseSlashCommand(text)?.name; + const draftDetached = submittedMode === "plan" || submittedMode === "vibe" || submittedMode === "goal"; + if (draftDetached) this.ctx.editor.clearDraft(); if (runner?.hasHandlers("input")) { const result = await runner.emitInput(text, inputImages, "interactive"); if (result?.handled) { - this.ctx.editor.clearDraft(); + if (!draftDetached) this.ctx.editor.clearDraft(); return; } if (result?.text !== undefined) { @@ -716,7 +720,7 @@ export class InputController { (inputImages?.length ?? 0) > 0 || (inputImageLinks?.length ?? 0) > 0 ? { images: inputImages, imageLinks: inputImageLinks } : undefined; - const slashResult = await executeBuiltinSlashCommand(text, { ctx: this.ctx, input }); + const slashResult = await executeBuiltinSlashCommand(text, { ctx: this.ctx, input, draftDetached }); if (slashResult === true) { if (!shouldSkipHistory(text)) this.ctx.editor.addToHistory(text); return; diff --git a/packages/coding-agent/src/modes/interactive-mode.ts b/packages/coding-agent/src/modes/interactive-mode.ts index 50abb8d29..f34657f77 100644 --- a/packages/coding-agent/src/modes/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive-mode.ts @@ -3501,6 +3501,7 @@ export class InteractiveMode implements InteractiveModeContext { await this.#startGoalFromObjective(objective, input); } catch (error) { this.showError(error instanceof Error ? error.message : String(error)); + throw error; } } async handleGuidedGoalCommand(rest?: string): Promise { @@ -3695,7 +3696,16 @@ export class InteractiveMode implements InteractiveModeContext { ): Promise { await this.#enterGoalMode({ objective, silent: true }); this.#resetGoalContinuationSuppression(); - if (!this.session.isStreaming && this.onInputCallback) { + if (this.session.isStreaming) { + const images = input?.images?.length ? input.images : undefined; + await this.withLocalSubmission( + objective, + () => this.session.prompt(objective, { streamingBehavior: "steer", images }), + { imageCount: images?.length ?? 0 }, + ); + return; + } + if (this.onInputCallback) { this.onInputCallback(this.startPendingSubmission({ text: objective, ...input }, { preserveDraft: true })); } } @@ -3712,8 +3722,15 @@ export class InteractiveMode implements InteractiveModeContext { this.#updateGoalModeStatus(); if (this.session.isStreaming) { await this.session.sendGoalModeContext({ deliverAs: "steer" }); + const images = input?.images?.length ? input.images : undefined; + await this.withLocalSubmission( + objective, + () => this.session.prompt(objective, { streamingBehavior: "steer", images }), + { imageCount: images?.length ?? 0 }, + ); + return; } - if (!this.session.isStreaming && this.onInputCallback) { + if (this.onInputCallback) { this.onInputCallback(this.startPendingSubmission({ text: objective, ...input }, { preserveDraft: true })); } } diff --git a/packages/coding-agent/src/slash-commands/builtin-modes.ts b/packages/coding-agent/src/slash-commands/builtin-modes.ts index c02a06610..8a97481ca 100644 --- a/packages/coding-agent/src/slash-commands/builtin-modes.ts +++ b/packages/coding-agent/src/slash-commands/builtin-modes.ts @@ -27,7 +27,7 @@ async function runWithDetachedModeDraft( run: () => Promise, ): Promise { const { editor } = runtime.ctx; - editor.clearDraft(); + if (!runtime.draftDetached) editor.clearDraft(); try { await run(); } catch (error) { @@ -254,8 +254,9 @@ export const BUILTIN_MODE_SLASH_COMMANDS: ReadonlyArray = [ return state ? `Goal: ${state.goal.status} (${shortDetail(state.goal.objective)})` : "Goal: off"; }, handleTui: async (command, runtime) => { - runtime.ctx.editor.clearDraft(); - await runtime.ctx.handleGoalModeCommand(command.args || undefined, runtime.input); + await runWithDetachedModeDraft(command, runtime, () => + runtime.ctx.handleGoalModeCommand(command.args || undefined, runtime.input), + ); }, }, { diff --git a/packages/coding-agent/src/slash-commands/types.ts b/packages/coding-agent/src/slash-commands/types.ts index 8f776eef1..1b1713035 100644 --- a/packages/coding-agent/src/slash-commands/types.ts +++ b/packages/coding-agent/src/slash-commands/types.ts @@ -85,6 +85,8 @@ export interface TuiSlashCommandRuntime { ctx: InteractiveModeContext; /** Post-extension-hook attachments belonging to the submitted slash draft. */ input?: Pick; + /** The editor snapshot was cleared before asynchronous input hooks ran. */ + draftDetached?: boolean; } /** Unified slash-command spec consumed by both TUI and ACP dispatchers. */ diff --git a/packages/coding-agent/test/goals/goal-mode-integration.test.ts b/packages/coding-agent/test/goals/goal-mode-integration.test.ts index 4cb8ef8bf..05be58be7 100644 --- a/packages/coding-agent/test/goals/goal-mode-integration.test.ts +++ b/packages/coding-agent/test/goals/goal-mode-integration.test.ts @@ -206,41 +206,33 @@ describe("InteractiveMode goal mode integration", () => { expect(await toolNamesFor(harness)).toContain("goal"); }); - it("defers initial goal objective submission while streaming", async () => { - let streaming = true; - Object.defineProperty(harness.session, "isStreaming", { configurable: true, get: () => streaming }); + it("steers initial goal objective attachments while streaming", async () => { + Object.defineProperty(harness.session, "isStreaming", { configurable: true, get: () => true }); const sendGoalModeContext = vi.spyOn(harness.session, "sendGoalModeContext").mockResolvedValue(); - const waiter = await armInputWaiter(harness.mode); + const promptSpy = vi.spyOn(harness.session, "prompt").mockResolvedValue(true); + const images: ImageContent[] = [{ type: "image", data: "aW1hZ2U=", mimeType: "image/png" }]; + const objective = "[Image #1, 10x10] Ship the release"; - await harness.mode.handleGoalModeCommand("Ship the release"); - await waitForMicrotasks(); + await harness.mode.handleGoalModeCommand(objective, { images, imageLinks: ["file:///shot.png"] }); - expect(harness.session.getGoalModeState()?.goal.objective).toBe("Ship the release"); + expect(harness.session.getGoalModeState()?.goal.objective).toBe(objective); expect(sendGoalModeContext).toHaveBeenCalledWith({ deliverAs: "steer" }); - expect(waiter.getResolvedInput()).toBeUndefined(); - - streaming = false; - harness.mode.onInputCallback?.(harness.mode.startPendingSubmission({ text: "cleanup" })); - await waiter.inputPromise; + expect(promptSpy).toHaveBeenCalledWith(objective, { streamingBehavior: "steer", images }); }); - it("defers replacement goal objective submission while streaming", async () => { + it("steers replacement goal objective attachments while streaming", async () => { await harness.mode.handleGoalModeCommand("Ship the release"); - let streaming = true; - Object.defineProperty(harness.session, "isStreaming", { configurable: true, get: () => streaming }); + Object.defineProperty(harness.session, "isStreaming", { configurable: true, get: () => true }); const sendGoalModeContext = vi.spyOn(harness.session, "sendGoalModeContext").mockResolvedValue(); - const waiter = await armInputWaiter(harness.mode); + const promptSpy = vi.spyOn(harness.session, "prompt").mockResolvedValue(true); + const images: ImageContent[] = [{ type: "image", data: "aW1hZ2U=", mimeType: "image/png" }]; + const objective = "[Image #1, 10x10] Replace the objective"; - await harness.mode.handleGoalModeCommand("set Replace the objective"); - await waitForMicrotasks(); + await harness.mode.handleGoalModeCommand(`set ${objective}`, { images, imageLinks: ["file:///shot.png"] }); - expect(harness.session.getGoalModeState()?.goal.objective).toBe("Replace the objective"); + expect(harness.session.getGoalModeState()?.goal.objective).toBe(objective); expect(sendGoalModeContext).toHaveBeenCalledWith({ deliverAs: "steer" }); - expect(waiter.getResolvedInput()).toBeUndefined(); - - streaming = false; - harness.mode.onInputCallback?.(harness.mode.startPendingSubmission({ text: "cleanup" })); - await waiter.inputPromise; + expect(promptSpy).toHaveBeenCalledWith(objective, { streamingBehavior: "steer", images }); }); const attachmentCases: Array<{ @@ -295,6 +287,28 @@ describe("InteractiveMode goal mode integration", () => { expect(input?.imageLinks).toBe(imageLinks); }); } + it("restores the goal draft when setup fails", async () => { + const images: ImageContent[] = [{ type: "image", data: "aW1hZ2U=", mimeType: "image/png" }]; + const imageLinks = ["file:///shot.png"]; + const commandText = "/goal [Image #1, 10x10] fix this"; + harness.mode.editor.setText(commandText); + harness.mode.editor.pendingImages = images; + harness.mode.editor.pendingImageLinks = imageLinks; + vi.spyOn(harness.session.goalRuntime, "createGoal").mockRejectedValueOnce(new Error("goal setup failed")); + const showError = vi.spyOn(harness.mode, "showError"); + + await expect( + executeBuiltinSlashCommand(commandText, { + ctx: harness.mode, + input: { images, imageLinks }, + }), + ).rejects.toThrow("goal setup failed"); + + expect(showError).toHaveBeenCalledWith("goal setup failed"); + expect(harness.mode.editor.getText()).toBe(commandText); + expect(harness.mode.editor.pendingImages).toEqual(images); + expect(harness.mode.editor.pendingImageLinks).toEqual(imageLinks); + }); it("keeps images pasted while delayed plan setup completes in the later draft", async () => { const submittedImages: ImageContent[] = [{ type: "image", data: "b2xk", mimeType: "image/png" }]; diff --git a/packages/coding-agent/test/slash-commands/mode-attachments.test.ts b/packages/coding-agent/test/slash-commands/mode-attachments.test.ts index f7d6dfc27..a96eccb43 100644 --- a/packages/coding-agent/test/slash-commands/mode-attachments.test.ts +++ b/packages/coding-agent/test/slash-commands/mode-attachments.test.ts @@ -5,7 +5,7 @@ import type { InteractiveModeContext, SubmittedUserInput } from "@oh-my-pi/pi-co type Attachments = Pick; -function createHarness(inputResult: { images?: ImageContent[] }) { +function createHarness(inputResult: { images?: ImageContent[] } | Promise<{ images?: ImageContent[] }>) { const oldImage: ImageContent = { type: "image", data: "b2xk", mimeType: "image/png" }; const handlePlanModeCommand = vi.fn(async (_prompt?: string, _input?: Attachments) => {}); const handleVibeModeCommand = vi.fn(async (_prompt?: string, _input?: Attachments) => {}); @@ -101,6 +101,25 @@ describe("mode command attachments", () => { expect(harness.editor.pendingImages).toEqual([]); expect(harness.editor.pendingImageLinks).toEqual([]); }); + it("detaches submitted images before awaiting input extensions", async () => { + const inputResult = Promise.withResolvers<{ images?: ImageContent[] }>(); + const harness = createHarness(inputResult.promise); + const submission = harness.editor.onSubmit?.("/plan inspect this"); + if (!submission) throw new Error("expected editor submit handler"); + + expect(harness.editor.pendingImages).toEqual([]); + const laterImage: ImageContent = { type: "image", data: "bmV3", mimeType: "image/png" }; + harness.editor.setText("later draft"); + harness.editor.pendingImages = [laterImage]; + harness.editor.pendingImageLinks = ["file:///later.png"]; + inputResult.resolve({}); + await submission; + + expect(harness.handlePlanModeCommand.mock.calls[0]?.[1]?.images).toHaveLength(1); + expect(harness.editor.getText()).toBe("later draft"); + expect(harness.editor.pendingImages).toEqual([laterImage]); + expect(harness.editor.pendingImageLinks).toEqual(["file:///later.png"]); + }); it("restores a failed mode command without overwriting a later draft", async () => { const failedPlan = createHarness({});