feat: wire truncate_before_message_id through frontend for edit/retry

Thread truncation metadata from edit/retry UI through the WebSocket API:
- acpApi.ts: prompt() accepts optional _meta parameter
- acp.ts: acpSendMessage passes truncateBeforeMessageId via _meta
- useChat.ts: sendMessage accepts and forwards truncation param,
  retryMessage passes the user message ID for backend truncation
- ChatView.tsx: handleSaveEdit passes truncation ID, uses deferred
  pattern for edit-during-streaming (BLOCKER 2), removes flushSync
  hack and dead 'edit' reason from pendingSend

Signed-off-by: Taylor Ho <taylorkmho@gmail.com>

ui/goose2/src/features/chat/hooks/useChat.ts

@@ -129,6 +129,7 @@ export function useChat(
       text: string,
       overridePersona?: { id: string; name?: string },
       attachments?: ChatAttachmentDraft[],
+      truncateBeforeMessageId?: string,
     ) => {
       const images = buildAcpImages(attachments);
       const hasAttachments = (attachments?.length ?? 0) > 0;
@@ -242,6 +243,7 @@ export function useChat(
           images: images?.map(
             (img) => [img.base64, img.mimeType] as [string, string],
           ),
+          truncateBeforeMessageId,
         });
 
         store.setChatState(sessionId, "idle");
@@ -371,7 +373,9 @@ export function useChat(
       // Bail if there's nothing to re-send — don't truncate
       if (!hasContent) return;
 
-      // Truncate from the user message onward (removes user msg + all responses)
+      // Truncate locally for immediate UI feedback; the backend truncates
+      // via _meta.truncate_before_message_id on the PromptRequest.
+      const truncateMessageId = userMessage.id;
       store.setMessages(sessionId, sessionMessages.slice(0, truncateFromIndex));
 
       const targetPersonaId = userMessage.metadata?.targetPersonaId;
@@ -382,6 +386,7 @@ export function useChat(
           ? { id: targetPersonaId, name: targetPersonaName }
           : undefined,
         attachmentDrafts.length > 0 ? attachmentDrafts : undefined,
+        truncateMessageId,
       );
     },
     [sessionId, store, sendMessage, chatState],

ui/goose2/src/features/chat/ui/ChatView.tsx

@@ -1,5 +1,4 @@
 import { useState, useEffect, useRef, useCallback, useMemo } from "react";
-import { flushSync } from "react-dom";
 import { useTranslation } from "react-i18next";
 import { AnimatePresence } from "motion/react";
 import { MessageTimeline } from "./MessageTimeline";
@@ -360,9 +359,7 @@ export function ChatView({
   // Unified deferred send — a single ref + effect drains any pending send
   // that had to wait for a state update (persona switch or edit truncation).
   const pendingSend = useRef<{
-    reason: "persona" | "edit";
     text: string;
-    persona?: { id: string; name?: string };
     attachments?: ChatAttachmentDraft[];
   } | null>(null);
   const queue = useMessageQueue(activeSessionId, chatState, sendMessage);
@@ -389,7 +386,7 @@ export function ChatView({
         }
         handlePersonaChange(personaId);
         // Defer the send until after persona state updates
-        pendingSend.current = { reason: "persona", text, attachments };
+        pendingSend.current = { text, attachments };
         return;
       }
       // Queue if agent is busy and no message already queued
@@ -413,52 +410,77 @@ export function ChatView({
   );
 
   /** Save an inline edit: truncate history from the edited message onward, then send the new text. */
+  /** Save an inline edit: truncate locally for instant UI feedback, then
+   *  send with truncation metadata so the backend also rewinds both the
+   *  display log and LLM context.
+   *
+   *  When the user edits while streaming, stopStreaming() sets chatState
+   *  to "idle" in the store, but sendMessage's closure still captures the
+   *  stale "streaming" value and silently bails. We defer the send into a
+   *  ref so the next render — which carries the fresh "idle" chatState and
+   *  a new sendMessage — picks it up via the useEffect below. */
+  const pendingEditSend = useRef<{
+    text: string;
+    truncateMessageId: string;
+    persona?: { id: string; name?: string };
+    attachments?: ChatAttachmentDraft[];
+  } | null>(null);
   const handleSaveEdit = useCallback(
     (messageId: string, text: string) => {
-      if (chatState !== "idle") {
+      const wasIdle = chatState === "idle";
+      if (!wasIdle) {
         stopStreaming();
       }
       const store = useChatStore.getState();
       const allMessages = store.messagesBySession[activeSessionId] ?? [];
       const editIndex = allMessages.findIndex((m) => m.id === messageId);
       if (editIndex === -1) {
-        // Message vanished — bail to avoid appending to untruncated history
         store.setEditingMessageId(activeSessionId, null);
         return;
       }
-      // Preserve persona and attachments from the original message
       const originalMessage = allMessages[editIndex];
       const targetPersonaId = originalMessage.metadata?.targetPersonaId;
       const targetPersonaName = originalMessage.metadata?.targetPersonaName;
       const originalAttachments = rebuildAttachmentDrafts(originalMessage);
+      const persona = targetPersonaId
+        ? { id: targetPersonaId, name: targetPersonaName }
+        : undefined;
+      const attachments =
+        originalAttachments.length > 0 ? originalAttachments : undefined;
 
-      // Use flushSync to commit state updates synchronously so we can call
-      // sendMessage immediately — no timing gap, no race condition.
-      flushSync(() => {
-        store.setMessages(activeSessionId, allMessages.slice(0, editIndex));
-        store.setEditingMessageId(activeSessionId, null);
-        store.setChatState(activeSessionId, "idle");
-      });
+      // Local truncation for immediate UI feedback.
+      store.setMessages(activeSessionId, allMessages.slice(0, editIndex));
+      store.setEditingMessageId(activeSessionId, null);
+      store.setChatState(activeSessionId, "idle");
 
-      sendMessage(
-        text,
-        targetPersonaId
-          ? { id: targetPersonaId, name: targetPersonaName }
-          : undefined,
-        originalAttachments.length > 0 ? originalAttachments : undefined,
-      );
+      if (wasIdle) {
+        // chatState was already "idle" — sendMessage's closure is fresh.
+        sendMessage(text, persona, attachments, messageId);
+      } else {
+        // Defer until React re-renders with fresh chatState / sendMessage.
+        pendingEditSend.current = {
+          text,
+          truncateMessageId: messageId,
+          persona,
+          attachments,
+        };
+      }
     },
     [activeSessionId, chatState, stopStreaming, sendMessage],
   );
 
-  // Single drain effect for deferred sends (persona switch path only now —
-  // edit path uses flushSync and calls sendMessage synchronously).
   useEffect(() => {
-    if (
-      pendingSend.current &&
-      pendingSend.current.reason === "persona" &&
-      selectedPersona
-    ) {
+    if (pendingEditSend.current && chatState === "idle") {
+      const { text, truncateMessageId, persona, attachments } =
+        pendingEditSend.current;
+      pendingEditSend.current = null;
+      sendMessage(text, persona, attachments, truncateMessageId);
+    }
+  }, [chatState, sendMessage]);
+
+  // Drain deferred sends (persona switch only).
+  useEffect(() => {
+    if (pendingSend.current && selectedPersona) {
       const { text, attachments } = pendingSend.current;
       pendingSend.current = null;
       sendMessage(text, undefined, attachments);

ui/goose2/src/shared/api/acp.ts

@@ -18,6 +18,10 @@ export interface AcpSendMessageOptions {
   personaName?: string;
   /** Image attachments as [base64Data, mimeType] pairs. */
   images?: [string, string][];
+  /** When set, the backend truncates this message and all subsequent messages
+   *  from both the display log and LLM context before appending the new prompt.
+   *  Used by edit and retry to rewind the conversation. */
+  truncateBeforeMessageId?: string;
 }
 
 export interface AcpPrepareSessionOptions {
@@ -36,7 +40,7 @@ export async function acpSendMessage(
   prompt: string,
   options: AcpSendMessageOptions = {},
 ): Promise<void> {
-  const { systemPrompt, personaId, images } = options;
+  const { systemPrompt, personaId, images, truncateBeforeMessageId } = options;
 
   const gooseSessionId = sessionTracker.getGooseSessionId(sessionId, personaId);
   if (!gooseSessionId) {
@@ -55,10 +59,14 @@ export async function acpSendMessage(
     }
   }
 
+  const meta = truncateBeforeMessageId
+    ? { truncate_before_message_id: truncateBeforeMessageId }
+    : undefined;
+
   const messageId = crypto.randomUUID();
   setActiveMessageId(gooseSessionId, messageId);
 
-  await directAcp.prompt(gooseSessionId, content);
+  await directAcp.prompt(gooseSessionId, content, meta);
 
   clearActiveMessageId(gooseSessionId);
 }

ui/goose2/src/shared/api/acpApi.ts

@@ -135,7 +135,8 @@ export async function loadSession(
 export async function prompt(
   sessionId: string,
   content: ContentBlock[],
+  meta?: Record<string, unknown>,
 ): Promise<PromptResponse> {
   const client = await getClient();
-  return client.prompt({ sessionId, prompt: content });
+  return client.prompt({ sessionId, prompt: content, _meta: meta });
 }