fix(tui): preserve mode attachments across async paths

This commit is contained in:
Fatih Al-Aziz
2026-08-11 16:25:20 +07:00
parent bb42aaaced
commit de9ee6d412
6 changed files with 89 additions and 32 deletions
@@ -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;
@@ -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<void> {
@@ -3695,7 +3696,16 @@ export class InteractiveMode implements InteractiveModeContext {
): Promise<void> {
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 }));
}
}
@@ -27,7 +27,7 @@ async function runWithDetachedModeDraft(
run: () => Promise<void>,
): Promise<void> {
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<SlashCommandSpec> = [
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),
);
},
},
{
@@ -85,6 +85,8 @@ export interface TuiSlashCommandRuntime {
ctx: InteractiveModeContext;
/** Post-extension-hook attachments belonging to the submitted slash draft. */
input?: Pick<SubmittedUserInput, "images" | "imageLinks">;
/** The editor snapshot was cleared before asynchronous input hooks ran. */
draftDetached?: boolean;
}
/** Unified slash-command spec consumed by both TUI and ACP dispatchers. */
@@ -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" }];
@@ -5,7 +5,7 @@ import type { InteractiveModeContext, SubmittedUserInput } from "@oh-my-pi/pi-co
type Attachments = Pick<SubmittedUserInput, "images" | "imageLinks">;
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({});