fix(tui): restore early-aborted prompt to composer (#325)

* fix(tui): restore early-aborted prompt to composer

* fix(tui): preserve restored submission edits and metadata

Author: hetaoBackend

README.md

@@ -145,7 +145,7 @@ Inside the TUI, use `/sessions` to find previous sessions and `/help` to see all
 | Reference a workspace file or directory | `@` |
 | Toggle Plan Mode | `Shift+Tab` |
 | Switch permission modes | `Alt+M` |
-| Close a panel or interrupt a running task | `Esc` |
+| Close a panel or interrupt a running task; interrupting before the model replies returns the message to the composer | `Esc` |
 
 ## Uninstall
 

README_ZH.md

@@ -145,7 +145,7 @@ mcode --session
 | 引用工作区文件或目录 | `@` |
 | 切换 Plan Mode | `Shift+Tab` |
 | 切换权限模式 | `Alt+M` |
-| 关闭面板或中断正在运行的任务 | `Esc` |
+| 关闭面板或中断正在运行的任务；在模型回复之前中断会把消息放回输入框 | `Esc` |
 
 ## 卸载
 

packages/tui/src/application/run-coordinator.ts

@@ -87,7 +87,7 @@ export interface TuiRunCoordinatorOptions {
   nowMs?: () => number;
 }
 
-const DEFAULT_CANCELLATION_SETTLEMENT_TIMEOUT_MS = 1_000;
+export const DEFAULT_CANCELLATION_SETTLEMENT_TIMEOUT_MS = 1_000;
 
 export class TuiRunCoordinator {
   private activeRun?: ActiveRun;

packages/tui/src/tui/app.ts

@@ -17,6 +17,10 @@ import { TranscriptView } from './transcript/view.js';
 import { TuiChatController, type TuiChatSnapshot } from './controller/chat-controller.js';
 import { TuiChromeFlow } from './controller/product/chrome-flow.js';
 import type { TuiDraftLifecycle } from './features/composer/draft-lifecycle.js';
+import { retainAvailableAttachmentPlaceholders } from './features/composer/draft-recovery.js';
+import type { TuiSubmissionSnapshot } from './features/composer/submission.js';
+import { statSync } from 'node:fs';
+import type { TuiTransportAttachment } from '../types/invocation.js';
 import { TuiRunProjection } from './state/run-projection.js';
 import { createTuiState, TuiEffectRunner, TuiStateStore } from './state/index.js';
 import { FeedbackFlow as Feedback } from './controller/product/feedback-flow.js';
@@ -358,6 +362,9 @@ export function createTuiApp(options: CreateTuiAppOptions): TuiApp {
     keybindings: options.keybindings,
     isStopped: () => stopped,
   });
+  // Turns whose prompt was already returned to the composer by the abort
+  // restore; keys are runtime turn ids, unique for the app's lifetime.
+  const restoredAbortTurnIds = new Set<string>();
   abortLiveTurn = createTuiAbortLiveTurn({
     controller,
     runProjection,
@@ -371,6 +378,125 @@ export function createTuiApp(options: CreateTuiAppOptions): TuiApp {
       if (!stopped) tui.requestRender();
     },
     append: appendLocalCell,
+    onLiveTurnAborted: ({ turnId }) => {
+      // Return the aborted prompt to the composer only when the turn produced
+      // no user-visible output: any assistant text, tool call, or steer
+      // message means the user was interacting with a live response, not
+      // regretting a fresh submission. Thinking is internal model process and
+      // does not disqualify: it is the most common regret moment, and nothing
+      // of value is discarded by resending.
+      if (sessionFlow.isSideModeActive()) return;
+      const cells = transcript.snapshot().filter((cell) => cell.turnId === turnId);
+      const userCell = transcript.get(`user:${turnId}`);
+      if (!userCell) return; // runtime-owned or retry-continuation turns
+      // A second Esc in the settle window can reach the runtime-owned branch
+      // with the same turnId; track restored turns so a repeat concatenates
+      // nothing. (Cell status cannot key this: abort-time markTurn also
+      // cancels still-pending user cells before the first restore.)
+      if (restoredAbortTurnIds.has(turnId)) return;
+      const hasIrreversibleActivity = cells.some((cell) => {
+        if (cell.id === `user:${turnId}`) return false;
+        if (cell.kind === 'thinking' || cell.kind === 'turn-duration') return false;
+        // markTurn('cancelled') synthesizes an empty assistant placeholder for
+        // turns with no assistant output before this callback can run.
+        if ((cell.kind === 'assistant' || cell.kind === 'assistant-preamble') && !cell.content) {
+          return false;
+        }
+        return true;
+      });
+      if (hasIrreversibleActivity) return;
+      // Preferred path: replay the ORIGINAL submission retained for this turn
+      // (complete attachment set, transport content, client intent, editor
+      // state) instead of reconstructing a lossy copy from the display cell.
+      const retained = commandFlow.getRetainedSubmission(turnId);
+      const retainedFiltered = retained
+        ? filterRetainedSubmissionForRestore(retained)
+        : undefined;
+      if (retained && retainedFiltered) {
+        commandFlow.restoreSubmission(retainedFiltered);
+        commandFlow.markAbortSubmissionRestored(retained.submissionId);
+        commandFlow.dropRetainedSubmission(turnId);
+        restoredAbortTurnIds.add(turnId);
+        markAbortedUserRowCancelled(transcript, turnId);
+        draftLifecycle?.recordRestoredSubmission(retainedFiltered);
+        chromeFlow?.setHint('Stopped · message restored to the Composer.');
+        updateChrome(controller.snapshot());
+        tui.requestRender();
+        return;
+      }
+      const text = userCell.content;
+      const draftAttachments = (userCell.attachments ?? []).flatMap((attachment) =>
+        attachment.filePath && typeof attachment.sizeBytes === 'number'
+          ? [
+              {
+                type: attachment.type,
+                fileName: attachment.fileName,
+                mimeType: attachment.mimeType,
+                sizeBytes: attachment.sizeBytes,
+                filePath: attachment.filePath,
+              },
+            ]
+          : [],
+      );
+      // Fallback path: reconstruct from the display cell. The transport list
+      // must carry BOTH halves of a mixed row (asset-backed AND file-backed
+      // entries): submission-time selection is `transportAttachments ??
+      // attachments`, so a partial transport list would silently drop the
+      // file-backed attachments on resubmit.
+      const transportAttachments: TuiTransportAttachment[] = (userCell.attachments ?? []).flatMap(
+        (attachment): TuiTransportAttachment[] =>
+          attachment.assetId
+            ? [
+                {
+                  type: attachment.type,
+                  fileName: attachment.fileName,
+                  mimeType: attachment.mimeType,
+                  ...(attachment.filePath ? { filePath: attachment.filePath } : {}),
+                  assetId: attachment.assetId,
+                },
+              ]
+            : attachment.filePath && typeof attachment.sizeBytes === 'number'
+              ? [
+                  {
+                    type: attachment.type,
+                    fileName: attachment.fileName,
+                    mimeType: attachment.mimeType,
+                    filePath: attachment.filePath,
+                  },
+                ]
+              : [],
+      );
+      if (!text.trim() && draftAttachments.length === 0 && transportAttachments.length === 0) {
+        return;
+      }
+      const restoredSubmission: TuiSubmissionSnapshot = {
+        submissionId: `abort-restore:${turnId}`,
+        sessionId: controller.snapshot().session?.sessionId,
+        content: text,
+        attachments: draftAttachments,
+        ...(transportAttachments.length > 0 ? { transportAttachments } : {}),
+        createdAtMs: Date.now(),
+        editor: {
+          schemaVersion: 1,
+          text,
+          cursor: text.length,
+          pastes: [],
+          pasteCounter: 0,
+        },
+      };
+      commandFlow.restoreSubmission(restoredSubmission);
+      restoredAbortTurnIds.add(turnId);
+      markAbortedUserRowCancelled(transcript, turnId);
+      // The restored text already carries the submission's content; a pending
+      // retry for it would merge the same text again on the next hydrate.
+      draftLifecycle?.recordRestoredSubmission(restoredSubmission);
+      chromeFlow?.setHint('Stopped · message restored to the Composer.');
+      // The funnel's updateChrome tail ran during the settle wait, before this
+      // hint was set; push the chrome once more so the hint is not left
+      // unpresented.
+      updateChrome(controller.snapshot());
+      tui.requestRender();
+    },
   });
   const planModeFlow = new TuiPlanModeFlow({
     currentSession: () => controller.snapshot().session,
@@ -447,8 +573,11 @@ export function createTuiApp(options: CreateTuiAppOptions): TuiApp {
     },
     followBottom: () => layout.forceFollowBottom(),
     requestWelcomeRebuild: () => tui.requestImmediateRender(),
-    switchComposerDraft: (sessionKey) =>
-      draftLifecycle?.switchSession(sessionKey) ?? Promise.resolve(),
+    switchComposerDraft: (sessionKey) => {
+      // Retained turn submissions belong to the previous session's turns.
+      commandFlow.clearRetainedSubmissions();
+      return draftLifecycle?.switchSession(sessionKey) ?? Promise.resolve();
+    },
     detachForegroundObserver: () => runtimeEventFlow?.detachForegroundObserver(),
     adoptForegroundRun: () =>
       sessionLifecycle.adoptForegroundRun(
@@ -730,8 +859,8 @@ export function createTuiApp(options: CreateTuiAppOptions): TuiApp {
     activity.dispose();
     widgets.imagePreview.dispose();
     renderer.prepareTranscriptExit();
-    disposeComponents(surfaceHost, editor, transcriptView, status, goal);
     draftStopPromise ??= draftLifecycle?.stop() ?? Promise.resolve();
+    disposeComponents(surfaceHost, editor, transcriptView, status, goal);
     if (stopOptions.abortActiveTurn !== false) {
       const snapshot = controller.snapshot();
       if (snapshot.activeTurnId) void controller.abort().catch(() => undefined);
@@ -808,8 +937,70 @@ export function createTuiApp(options: CreateTuiAppOptions): TuiApp {
       updateChromeAndRequestRender();
     },
     submit: (input) => commandFlow.submit(input).then(() => undefined),
+    /** Exposed for integration tests that submit with an explicit seed. */
+    commandFlow,
     abortTurn: abortLiveTurn,
     leaveUi,
     stop,
   };
 }
+
+function fileExists(filePath: string): boolean {
+  try {
+    return statSync(filePath).isFile();
+  } catch {
+    return false;
+  }
+}
+
+/**
+ * Filters a retained submission snapshot for abort-time restore: drops
+ * attachments whose backing file disappeared (asset-backed entries stay),
+ * and prunes editor placeholders accordingly. Returns undefined when nothing
+ * restorable remains.
+ */
+function filterRetainedSubmissionForRestore(
+  snapshot: TuiSubmissionSnapshot,
+): TuiSubmissionSnapshot | undefined {
+  const attachments = snapshot.attachments.filter((attachment) =>
+    fileExists(attachment.filePath),
+  );
+  const transportAttachments = (snapshot.transportAttachments ?? []).filter((attachment) =>
+    attachment.assetId ? true : attachment.filePath ? fileExists(attachment.filePath) : false,
+  );
+  if (!snapshot.content.trim() && attachments.length === 0 && transportAttachments.length === 0) {
+    return undefined;
+  }
+  return {
+    ...snapshot,
+    submissionId: `abort-restore:${snapshot.submissionId}`,
+    attachments,
+    ...(snapshot.transportAttachments ? { transportAttachments } : {}),
+    editor: retainAvailableAttachmentPlaceholders(
+      // Some submission paths (automation results, test seeds) submit with a
+      // payload but an empty editor draft; restoring that verbatim would
+      // leave the composer blank, so fall back to the visible content.
+      snapshot.editor.text.trim()
+        ? snapshot.editor
+        : { ...snapshot.editor, text: snapshot.content, cursor: snapshot.content.length },
+      attachments,
+    ),
+  };
+}
+
+/** Marks the aborted turn's user row cancelled so it stays visible in history. */
+function markAbortedUserRowCancelled(transcript: TranscriptStore, turnId: string): void {
+  const userCell = transcript.get(`user:${turnId}`);
+  if (!userCell) return;
+  // markTurn does not rewrite the user cell once the runtime echo flipped it
+  // to 'succeeded', so upsert explicitly. createdAtMs is force-preserved by
+  // the store's merge for existing cells and is intentionally omitted.
+  transcript.upsert({
+    id: `user:${turnId}`,
+    kind: 'user',
+    status: 'cancelled',
+    ...(userCell.content ? { content: userCell.content } : {}),
+    updatedAtMs: Date.now(),
+    ...(userCell.attachments ? { attachments: userCell.attachments } : {}),
+  });
+}

packages/tui/src/tui/controller/chat-controller-types.ts

@@ -70,4 +70,10 @@ export interface TuiSubmitOptions {
   beforeTurnAdmission?: (sessionId: string) => void | Promise<void>;
   onRuntimeAccepted?: (sessionId: string) => void;
   optimisticRequestId?: string;
+  /**
+   * Fired once a new turn has been minted and projected (before the first
+   * await), so the caller can retain per-turn state such as the original
+   * submission snapshot. Not fired for steer paths or retry continuations.
+   */
+  onTurnStarted?: (turnId: string) => void;
 }

packages/tui/src/tui/controller/chat-controller.ts

@@ -389,6 +389,9 @@ export class TuiChatController {
       errorRetryable: undefined,
     });
     if (!optimisticCell && !isRetryContinuation) this.onUserSubmissionProjected?.();
+    // Before the first await: the turn id is final, so the caller can retain
+    // the original submission snapshot for this turn.
+    if (!isRetryContinuation) options.onTurnStarted?.(turnId);
 
     try {
       await this.requireLoginForAgentAction();