fix(next): retain ordered prompt attachments through send and retry

cmux-tui/crates/acpmux/src/hub/prompt_images.rs

@@ -186,6 +186,10 @@ pub(super) fn inline_parts(
     let kept_images: HashSet<&str> =
         images.iter().filter_map(|image| image.get("id").and_then(Value::as_str)).collect();
     let display = validated_display_parts(display_parts, blocks);
+    let companion_images: HashSet<u64> = display
+        .values()
+        .filter_map(|part| part.get("companionOf").and_then(Value::as_u64))
+        .collect();
     blocks
         .iter()
         .enumerate()
@@ -222,6 +226,9 @@ pub(super) fn inline_parts(
                     let persisted = kept_images.contains(id);
                     let mut part = named;
                     part["kind"] = Value::String("image".into());
+                    if companion_images.contains(&(block_index as u64)) {
+                        part["blockIndex"] = json!(block_index);
+                    }
                     if !persisted {
                         part["persisted"] = Value::Bool(false);
                     }
@@ -317,6 +324,20 @@ pub(crate) fn validated_display_parts(
         }
         let mut out =
             serde_json::Map::from_iter([(String::from("kind"), Value::String(kind.into()))]);
+        if let Some(raw) = object.get("companionOf") {
+            let Some(index) = raw.as_u64().and_then(|n| usize::try_from(n).ok()) else { continue };
+            if kind != "file"
+                || block_type != "resource_link"
+                || index.checked_add(1) != Some(block_index)
+                || !blocks.get(index).is_some_and(|target| {
+                    target.get("type").and_then(Value::as_str) == Some("image")
+                        && decode(target).is_some()
+                })
+            {
+                continue;
+            }
+            out.insert("companionOf".into(), json!(index));
+        }
         if let Some(name) = object.get("name").and_then(Value::as_str) {
             if name.is_empty() || name.chars().count() > MAX_DISPLAY_NAME {
                 continue;
@@ -362,7 +383,7 @@ pub(crate) fn validated_display_parts(
 
 fn copy_display_fields(part: &mut Value, display: &Value) {
     let Some(object) = display.as_object() else { return };
-    for key in ["name", "mimeType", "size", "pageCount", "start", "end"] {
+    for key in ["name", "mimeType", "size", "pageCount", "start", "end", "companionOf"] {
         if let Some(value) = object.get(key) {
             part[key] = value.clone();
         }
@@ -659,4 +680,33 @@ mod tests {
             vec![json!({"type": "resource_link", "uri": "", "name": "archive.zip"})];
         assert!(validated_display_parts(Some(&empty_resource), &resource_block).is_empty());
     }
+    #[test]
+    fn inline_parts_validate_companion_relation() {
+        let image = json!({"type": "image", "mimeType": "image/png", "data": "AQID"});
+        let file =
+            json!({"type": "resource_link", "name": "one.png", "uri": "file:///tmp/one.png"});
+        let blocks = vec![
+            image.clone(),
+            file.clone(),
+            json!({"type": "text", "text": "between"}),
+            image,
+            file,
+        ];
+        let display = json!([
+            {"kind": "file", "blockIndex": 1, "companionOf": 0},
+            {"kind": "file", "blockIndex": 4, "companionOf": 3}
+        ]);
+        let parts = inline_parts(&blocks, &[], Some(&display));
+        assert_eq!(parts[0]["id"], parts[3]["id"], "same bytes, distinct block identities");
+        assert_eq!(parts[0]["blockIndex"], 0);
+        assert_eq!(parts[1]["companionOf"], 0);
+        assert_eq!(parts[3]["blockIndex"], 3);
+        assert_eq!(parts[4]["companionOf"], 3);
+        assert_eq!(parts[1]["uri"], "file:///tmp/one.png", "the file is still delivered");
+        for (block, target) in [(1, 2), (1, 99), (1, -1), (0, 3), (4, 0)] {
+            let invalid = json!([{"kind": if block == 0 { "image" } else { "file" }, "blockIndex": block, "companionOf": target}]);
+            let parts = inline_parts(&blocks, &[], Some(&invalid));
+            assert!(parts.iter().all(|part| part.get("companionOf").is_none()));
+        }
+    }
 }

webviews/src/agent-session/acpmux/App.tsx

@@ -5434,11 +5434,10 @@ function AcpmuxPane() {
       }),
       ...(connected && {
         retry: (rowId: string, prompt: string) =>
-          void (async () =>
-            callNative("chat.send", {
-              text: prompt,
-              attachments: (await directClient.current?.promptAttachments(rowId)) ?? [],
-            }))().catch(() => undefined),
+          void (async () => {
+            const draft = await directClient.current?.resendPrompt(rowId, prompt);
+            return callNative("chat.send", { text: draft?.input ?? prompt, attachments: draft?.attachments ?? [] });
+          })().catch(() => undefined),
       }),
       // Continue after a stop (Codex's) sends its prompt as the composer would: it shows as the
       // reader's prompt and starts the next turn, or queues behind one already sent. Like Restart
@@ -5480,6 +5479,8 @@ function AcpmuxPane() {
           takeBackSteer(row, {
             withdraw: (promptId) => editQueued(promptId, { op: "withdraw" }),
             attachments: (rowId) => directClient.current?.promptAttachments(rowId) ?? Promise.resolve([]),
+            draft: (rowId, text) =>
+              directClient.current?.resendPrompt(rowId, text) ?? Promise.reject(new Error(translate("error.notOpen"))),
             restore: (text, attachments) => composerHandle.current?.restore(text, attachments),
           }),
       }),
@@ -7320,7 +7321,8 @@ function AcpmuxPane() {
               const live = summary?.sessionId ? { sessionId: summary.sessionId, model: summary.model } : undefined;
               if (!(await harnessSwitch.pickModel(next, live))) return;
             }
-            await send(String(prompt ?? ""), await client.promptAttachments(String(rowId ?? "")));
+            const draft = await client.resendPrompt(String(rowId ?? ""), String(prompt ?? ""));
+            await send(draft.input, draft.attachments);
           },
           // Edit and Resend (the message menu): the prompt goes back into the composer, before
           // anything typed since, with the caret in it.
@@ -7396,7 +7398,7 @@ function AcpmuxPane() {
           // edited text with the prompt's attachments; the original chat stays as it was.
           "chat.editResend": async ({ rowId, throughSeq, text }) => {
             harnessSwitch.cancel();
-            const attachments = await client.promptAttachments(String(rowId ?? ""));
+            const draft = await client.resendPrompt(String(rowId ?? ""), String(text ?? ""));
             // A fork's own edited prompt edits again from the chat it was forked from, at the same
             // turn, so every version of the prompt is a sibling (conversation/branches.ts).
             const current = snapshotRef.current;
@@ -7411,7 +7413,7 @@ function AcpmuxPane() {
             await persistSession(forked);
             // The reader may open another chat while the fork is saved: the edit belongs to the fork.
             if (client.selectedSession !== forked) return false;
-            await send(String(text ?? ""), attachments);
+            await send(draft.input, draft.attachments);
             return true;
           },
           // `/btw <question>`: acpmux starts the side session; the composer then opens it.

webviews/src/agent-session/acpmux/conversation/PendingSteer.tsx

@@ -124,15 +124,20 @@ export async function takeBackSteer(
   io: {
     withdraw: (promptId: string) => Promise<boolean>;
     attachments: (rowId: string) => Promise<ComposerAttachment[]>;
+    draft?: (rowId: string, text: string) => Promise<{ input: string; attachments: ComposerAttachment[] }>;
     restore: (text: string, attachments: ComposerAttachment[]) => void;
   },
 ): Promise<boolean> {
   const promptId = row.steerPromptId;
   if (!promptId || row.steer !== "waiting") return false;
-  const attachments = await io.attachments(row.id).catch(() => undefined);
+  const text = row.resend ?? row.text ?? "";
+  const draft = await (
+    io.draft ? io.draft(row.id, text) : io.attachments(row.id).then((attachments) => ({ input: text, attachments }))
+  ).catch(() => undefined);
+  const attachments = draft?.attachments;
   // An image acpmux could not give back is left out of what it read: the steer stays then too.
   const images = attachments?.filter((attachment) => attachment.kind === "image").length ?? 0;
-  if (!attachments || images < (row.images?.length ?? 0) || !(await io.withdraw(promptId))) return false;
-  io.restore(row.resend ?? row.text ?? "", attachments);
+  if (!draft || images < (row.images?.length ?? 0) || !(await io.withdraw(promptId))) return false;
+  io.restore(draft.input, draft.attachments);
   return true;
 }