fix(web): restore reverted messages to composer

apps/web/src/components/ChatView.logic.test.ts

@@ -18,8 +18,11 @@ import {
   buildThreadTurnInterruptInput,
   createLocalDispatchSnapshot,
   deriveComposerSendState,
+  deriveRevertedMessageContent,
+  deriveRevertedMessagePrompt,
   dismissBranchMismatchForSession,
   ENVIRONMENT_RECONNECT_WARNING_GRACE_MS,
+  fetchRevertedMessageAttachmentBlob,
   getStartedThreadModelChangeBlockReason,
   hasEnvironmentReconnectWarningGraceElapsed,
   hasServerAcknowledgedLocalDispatch,
@@ -34,6 +37,92 @@ import {
   shouldWriteThreadErrorToCurrentServerThread,
 } from "./ChatView.logic";
 
+describe("reverted message prompt", () => {
+  it("restores editable text without sent-only context blocks", () => {
+    const prompt = [
+      "Please fix this",
+      "",
+      "<terminal_context>",
+      "- Terminal 1:",
+      "  npm test",
+      "</terminal_context>",
+      "",
+      "<element_context>",
+      "- <button>:",
+      "  url: http://localhost:5173",
+      "</element_context>",
+      "",
+      "<preview_annotation>",
+      '{"id":"preview-1","title":"Button","x":1,"y":2}',
+      "</preview_annotation>",
+      "",
+      '<review_comment sectionId="file:a.ts" filePath="a.ts" startIndex="1" endIndex="1">',
+      "Check this line",
+      "```diff",
+      "+const value = 1;",
+      "```",
+      "</review_comment>",
+    ].join("\n");
+
+    expect(deriveRevertedMessagePrompt(prompt)).toBe("Please fix this");
+  });
+
+  it("only excludes screenshots tied to parsed preview annotations", () => {
+    const prompt = [
+      "Please fix this",
+      "",
+      "<preview_annotation>",
+      "Preview annotation:",
+      "Id: preview-1",
+      "Page: Button",
+      "The attached screenshot is the annotated preview crop.",
+      "</preview_annotation>",
+    ].join("\n");
+
+    const content = deriveRevertedMessageContent(prompt);
+
+    expect(content.prompt).toBe("Please fix this");
+    expect([...content.previewAnnotationImageNames]).toEqual(["preview-annotation-preview-1.png"]);
+    expect(content.previewAnnotationImageNames.has("preview-annotation-design.png")).toBe(false);
+  });
+});
+
+describe("reverted message attachment fetch", () => {
+  afterEach(() => {
+    vi.useRealTimers();
+    vi.restoreAllMocks();
+  });
+
+  it("aborts a stalled attachment request after the timeout", async () => {
+    vi.useFakeTimers();
+    vi.spyOn(globalThis, "fetch").mockImplementation((_input, init) => {
+      const signal = init?.signal;
+      if (!signal) throw new Error("Expected an abort signal.");
+      return new Promise<Response>((_resolve, reject) => {
+        signal.addEventListener("abort", () => reject(signal.reason), { once: true });
+      });
+    });
+
+    const pendingFetch = fetchRevertedMessageAttachmentBlob(
+      "https://example.test/image.png",
+      1_000,
+    );
+    const rejection = expect(pendingFetch).rejects.toMatchObject({ name: "AbortError" });
+
+    await vi.advanceTimersByTimeAsync(1_000);
+    await rejection;
+  });
+
+  it("clears the timeout after an attachment request completes", async () => {
+    vi.useFakeTimers();
+    vi.spyOn(globalThis, "fetch").mockResolvedValue(new Response(null, { status: 200 }));
+
+    await fetchRevertedMessageAttachmentBlob("https://example.test/image.png", 1_000);
+
+    expect(vi.getTimerCount()).toBe(0);
+  });
+});
+
 const environmentId = EnvironmentId.make("environment-local");
 const projectId = ProjectId.make("project-1");
 const threadId = ThreadId.make("thread-1");

apps/web/src/components/ChatView.logic.ts

@@ -16,19 +16,70 @@ import * as Schema from "effect/Schema";
 import { appAtomRegistry } from "../rpc/atomRegistry";
 import { environmentThreadDetails } from "../state/threads";
 import {
+  deriveDisplayedUserMessageState,
   filterTerminalContextsWithText,
   stripInlineTerminalContextPlaceholders,
   type TerminalContextDraft,
 } from "../lib/terminalContext";
+import { extractTrailingPreviewAnnotation } from "../lib/previewAnnotation";
+import { parseReviewCommentMessageSegments } from "../reviewCommentContext";
 import type { DraftThreadEnvMode } from "../composerDraftStore";
 
 export const LAST_INVOKED_SCRIPT_BY_PROJECT_KEY = "t3code:last-invoked-script-by-project";
 export const MAX_HIDDEN_MOUNTED_TERMINAL_THREADS = 10;
 export const MAX_HIDDEN_MOUNTED_PREVIEW_THREADS = 3;
 export const ENVIRONMENT_RECONNECT_WARNING_GRACE_MS = 2_000;
+export const REVERTED_MESSAGE_ATTACHMENT_TIMEOUT_MS = 10_000;
 
 export const LastInvokedScriptByProjectSchema = Schema.Record(ProjectId, Schema.String);
 
+export function deriveRevertedMessageContent(prompt: string): {
+  prompt: string;
+  previewAnnotationImageNames: ReadonlySet<string>;
+} {
+  const withoutReviewComments = parseReviewCommentMessageSegments(prompt)
+    .filter((segment) => segment.kind === "text")
+    .map((segment) => segment.text)
+    .join("");
+
+  let withoutPreviewAnnotations = withoutReviewComments;
+  const previewAnnotationImageNames = new Set<string>();
+  while (true) {
+    const extracted = extractTrailingPreviewAnnotation(withoutPreviewAnnotations);
+    if (!extracted.annotation) break;
+    if (extracted.annotation.hasScreenshot) {
+      previewAnnotationImageNames.add(`preview-annotation-${extracted.annotation.id}.png`);
+    }
+    withoutPreviewAnnotations = extracted.promptText;
+  }
+
+  return {
+    prompt: deriveDisplayedUserMessageState(withoutPreviewAnnotations).visibleText,
+    previewAnnotationImageNames,
+  };
+}
+
+export function deriveRevertedMessagePrompt(prompt: string): string {
+  return deriveRevertedMessageContent(prompt).prompt;
+}
+
+export async function fetchRevertedMessageAttachmentBlob(
+  previewUrl: string,
+  timeoutMs = REVERTED_MESSAGE_ATTACHMENT_TIMEOUT_MS,
+): Promise<Blob> {
+  const controller = new AbortController();
+  const timeoutId = globalThis.setTimeout(() => controller.abort(), timeoutMs);
+  try {
+    // The timeline may have already cached this URL from an <img> request,
+    // which carries no Origin header and therefore no CORS response header.
+    const response = await fetch(previewUrl, { cache: "reload", signal: controller.signal });
+    if (!response.ok) throw new Error("Could not load reverted message attachment.");
+    return await response.blob();
+  } finally {
+    globalThis.clearTimeout(timeoutId);
+  }
+}
+
 export function scheduleEnvironmentReconnectWarning(showWarning: () => void): () => void {
   const timeoutId = globalThis.setTimeout(showWarning, ENVIRONMENT_RECONNECT_WARNING_GRACE_MS);
   return () => globalThis.clearTimeout(timeoutId);