Preserve local image attachments when restoring killed text (#57502)

## Why

Restoring killed text should restore its local image attachments along with their placeholders.

## What changed

Keep image paths in the composer's kill buffer, including across composer replacement, and restore attachments with fresh labels when text is pasted back. Preserve plain-text insertion in Vim Replace mode so per-character recovery continues to work.

## Testing

Add regression coverage for restoring images after submission, transferring the kill buffer between composers, and preserving existing images and Backspace recovery in Vim Replace mode.

GitOrigin-RevId: 9bc052a1b372d5a1245ddc54c7ffafee4a637494

codex-rs/tui/src/bottom_pane/chat_composer.rs

@@ -149,7 +149,9 @@
 //! dispatch, they intentionally preserve the textarea kill buffer. That lets users `Ctrl+K` part
 //! of a draft, perform a composer action such as changing reasoning level, and then `Ctrl+Y` the
 //! killed text back into the now-empty draft. Replacing the chat widget carries that buffer into
-//! the fresh composer so Vim yanks survive `/new` and thread switches.
+//! the fresh composer so Vim yanks survive `/new` and thread switches. Local images in the killed
+//! text retain their attachment paths and receive fresh labels when restored outside Vim Replace
+//! mode, whose per-character recovery requires plain text.
 //!
 //! The numeric auto-submit path used by the slash popup performs the same pending-paste expansion
 //! and attachment pruning, and clears pending paste state on success.
@@ -426,7 +428,7 @@ use crate::app_event::ConnectorsSnapshot;
 use crate::app_event_sender::AppEventSender;
 use crate::bottom_pane::LocalImageAttachment;
 use crate::bottom_pane::MentionBinding;
-use crate::bottom_pane::textarea::KillBufferSnapshot;
+use crate::bottom_pane::textarea::KillBufferEdit;
 use crate::bottom_pane::textarea::TextArea;
 use crate::clipboard_paste::normalize_pasted_path;
 use crate::clipboard_paste::pasted_image_format;
@@ -631,6 +633,7 @@ pub(crate) struct ChatComposer {
     effort_observed: bool,
     luna_reserve_active: bool,
     attachments: AttachmentState,
+    kill_buffer_images: Vec<(Range<usize>, PathBuf)>,
     placeholder_text: String,
     blocks_direct_input: bool,
     is_task_running: bool,
@@ -668,6 +671,11 @@ pub(crate) struct ChatComposer {
     vim_normal_keymap: VimNormalKeymap,
 }
 
+pub(crate) struct KillBufferSnapshot {
+    textarea: crate::bottom_pane::textarea::KillBufferSnapshot,
+    images: Vec<(Range<usize>, PathBuf)>,
+}
+
 /// A resolved legacy `$` target plus any catalog built while disambiguating shell syntax.
 struct MentionCompletionTarget {
     range: Range<usize>,
@@ -809,6 +817,7 @@ impl ChatComposer {
             effort_observed: false,
             luna_reserve_active: false,
             attachments: AttachmentState::default(),
+            kill_buffer_images: Vec::new(),
             placeholder_text,
             blocks_direct_input: false,
             is_task_running: false,
@@ -1390,11 +1399,17 @@ impl ChatComposer {
     }
 
     pub(crate) fn take_kill_buffer_snapshot(&mut self) -> KillBufferSnapshot {
-        self.draft.textarea.take_kill_buffer_snapshot()
+        KillBufferSnapshot {
+            textarea: self.draft.textarea.take_kill_buffer_snapshot(),
+            images: std::mem::take(&mut self.kill_buffer_images),
+        }
     }
 
     pub(crate) fn restore_kill_buffer_snapshot(&mut self, snapshot: KillBufferSnapshot) {
-        self.draft.textarea.restore_kill_buffer_snapshot(snapshot);
+        self.draft
+            .textarea
+            .restore_kill_buffer_snapshot(snapshot.textarea);
+        self.kill_buffer_images = snapshot.images;
     }
 
     /// Restore draft history transferred from the startup composer.
@@ -3763,8 +3778,27 @@ impl ChatComposer {
 
         self.begin_vim_edit(input);
         self.draft.textarea.input(input);
+        match self.draft.textarea.take_kill_buffer_edit() {
+            Some(KillBufferEdit::Stored(elements)) => {
+                self.kill_buffer_images = elements
+                    .into_iter()
+                    .filter_map(|(range, text)| {
+                        self.attachments
+                            .local_images
+                            .iter()
+                            .find(|image| image.placeholder == text)
+                            .map(|image| (range, image.path.clone()))
+                    })
+                    .collect();
+            }
+            Some(KillBufferEdit::Inserted(start)) => self.attachments.restore_killed_images(
+                &mut self.draft.textarea,
+                start,
+                &self.kill_buffer_images,
+            ),
+            None => {}
+        }
         self.sync_bash_mode_from_text();
-
         if let Some(elements_before) = elements_before {
             self.reconcile_deleted_elements(elements_before);
         }
@@ -9327,7 +9361,11 @@ mod tests {
             "image_placeholder_single",
             /*enhanced_keys_supported*/ false,
             |composer| {
-                composer.attach_image(PathBuf::from("/tmp/image1.png"));
+                let (mut previous, _rx) = new_test_composer();
+                previous.attach_image(PathBuf::from("/tmp/image1.png"));
+                previous.handle_key_event(KeyEvent::new(KeyCode::Char('u'), KeyModifiers::CONTROL));
+                composer.restore_kill_buffer_snapshot(previous.take_kill_buffer_snapshot());
+                composer.handle_key_event(KeyEvent::new(KeyCode::Char('y'), KeyModifiers::CONTROL));
             },
         );
 
@@ -9568,7 +9606,7 @@ mod tests {
     }
 
     #[test]
-    fn kill_buffer_persists_after_submit() {
+    fn kill_buffer_preserves_image_after_submit() {
         use crossterm::event::KeyCode;
         use crossterm::event::KeyEvent;
         use crossterm::event::KeyModifiers;
@@ -9583,11 +9621,12 @@ mod tests {
             /*disable_paste_burst*/ false,
         );
         composer.set_steer_enabled(true);
-        composer.draft.textarea.insert_str("restore me");
-        composer.draft.textarea.set_cursor(/*pos*/ 0);
+        composer.draft.textarea.insert_str("restore me ");
+        let image_path = PathBuf::from("/tmp/killed-image.png");
+        composer.attach_image(image_path.clone());
 
         let (_result, _needs_redraw) =
-            composer.handle_key_event(KeyEvent::new(KeyCode::Char('k'), KeyModifiers::CONTROL));
+            composer.handle_key_event(KeyEvent::new(KeyCode::Char('u'), KeyModifiers::CONTROL));
         assert!(composer.draft.textarea.is_empty());
 
         composer.draft.textarea.insert_str("hello");
@@ -9596,9 +9635,27 @@ mod tests {
         assert!(matches!(result, InputResult::Submitted { .. }));
         assert!(composer.draft.textarea.is_empty());
 
+        let other_image_path = PathBuf::from("/tmp/other-image.png");
+        composer.attach_image(other_image_path.clone());
+        composer.handle_paste(" ".into());
         let (_result, _needs_redraw) =
             composer.handle_key_event(KeyEvent::new(KeyCode::Char('y'), KeyModifiers::CONTROL));
-        assert_eq!(composer.draft.textarea.text(), "restore me");
+        let (result, _) =
+            composer.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE));
+        assert_eq!(
+            result,
+            InputResult::Submitted {
+                text: "[Image #1] restore me [Image #2]".into(),
+                text_elements: vec![
+                    TextElement::new((0..10).into(), Some("[Image #1]".into())),
+                    TextElement::new((22..32).into(), Some("[Image #2]".into())),
+                ],
+            }
+        );
+        assert_eq!(
+            composer.take_recent_submission_images(),
+            vec![other_image_path, image_path]
+        );
     }
 
     #[test]

codex-rs/tui/src/bottom_pane/chat_composer/attachment_state.rs

@@ -2,6 +2,7 @@
 //! remote-image rows, and keyboard selection over those rows.
 
 use std::collections::HashSet;
+use std::ops::Range;
 use std::path::PathBuf;
 
 use crossterm::event::KeyCode;
@@ -100,6 +101,33 @@ impl AttachmentState {
         self.local_images.push(AttachedImage { placeholder, path });
     }
 
+    pub(super) fn restore_killed_images(
+        &mut self,
+        textarea: &mut TextArea,
+        start: usize,
+        images: &[(Range<usize>, PathBuf)],
+    ) {
+        let mut position = start;
+        let mut previous_end = 0;
+        for (range, path) in images {
+            position += range.start - previous_end;
+            let placeholder =
+                local_image_label_text(self.remote_image_urls.len() + self.local_images.len() + 1);
+            let cursor_at_start = textarea.cursor() == position;
+            textarea.replace_range(position..position + range.len(), &placeholder);
+            let _ = textarea.add_element_range(position..position + placeholder.len());
+            if cursor_at_start {
+                textarea.set_cursor(position);
+            }
+            position += placeholder.len();
+            self.local_images.push(AttachedImage {
+                placeholder,
+                path: path.clone(),
+            });
+            previous_end = range.end;
+        }
+    }
+
     pub(super) fn prune_local_images_for_submission(
         &mut self,
         text: &str,

codex-rs/tui/src/bottom_pane/chat_composer/vim_history_tests.rs

@@ -603,6 +603,27 @@ fn undo_and_redo_restore_deleted_image_attachment() {
     assert!(composer.local_image_paths().is_empty());
 }
 
+#[test]
+fn vim_replace_yank_preserves_existing_image_and_recovery() {
+    let mut composer = vim_composer("");
+    keys(&mut composer, "i");
+    composer.attach_image(PathBuf::from("killed.png"));
+    composer.handle_key_event(KeyEvent::new(KeyCode::Char('u'), KeyModifiers::CONTROL));
+    composer.attach_image(PathBuf::from("existing.png"));
+    keys(&mut composer, "abcdefghij");
+    let before = (composer.current_text(), composer.local_image_paths());
+    escape(&mut composer);
+    keys(&mut composer, "0R");
+    composer.handle_key_event(KeyEvent::new(KeyCode::Char('y'), KeyModifiers::CONTROL));
+    for _ in 0.."[Image #1]".len() {
+        composer.handle_key_event(KeyEvent::from(KeyCode::Backspace));
+    }
+    assert_eq!(
+        (composer.current_text(), composer.local_image_paths()),
+        before
+    );
+}
+
 #[test]
 fn find_jump_and_search_navigation_never_snapshot_the_draft() {
     let mut composer = vim_composer("alpha beta\nalpha gamma");

codex-rs/tui/src/bottom_pane/mod.rs

@@ -266,10 +266,10 @@ pub(crate) use chat_composer::ChatComposerConfig;
 pub(crate) use chat_composer::ClipboardPasteTarget;
 pub(crate) use chat_composer::ComposerDraftSnapshot;
 pub(crate) use chat_composer::InputResult;
+pub(crate) use chat_composer::KillBufferSnapshot;
 pub(crate) use chat_composer::QueuedInputAction;
 pub(crate) use chat_composer::RestrictedInputMode;
 pub(crate) use chat_composer_history::HistoryEntry;
-pub(crate) use textarea::KillBufferSnapshot;
 
 use crate::status_indicator_widget::StatusDetailsCapitalization;
 use crate::status_indicator_widget::StatusIndicatorWidget;

codex-rs/tui/src/bottom_pane/textarea.rs

@@ -153,6 +153,7 @@ pub(crate) struct TextArea {
     next_element_id: u64,
     kill_buffer: String,
     kill_buffer_kind: KillBufferKind,
+    kill_buffer_edit: Option<KillBufferEdit>,
     vim_enabled: bool,
     vim_mode: VimMode,
     vim_pending: VimPending,
@@ -201,6 +202,14 @@ pub(crate) struct KillBufferSnapshot {
     kind: KillBufferKind,
 }
 
+#[derive(Debug)]
+pub(super) enum KillBufferEdit {
+    /// Atomic elements with ranges relative to the newly stored kill buffer.
+    Stored(Vec<(Range<usize>, String)>),
+    /// Start of the kill buffer's contents in the visible draft.
+    Inserted(usize),
+}
+
 impl Default for TextArea {
     fn default() -> Self {
         Self::new()
@@ -223,6 +232,7 @@ impl TextArea {
             next_element_id: 1,
             kill_buffer: String::new(),
             kill_buffer_kind: KillBufferKind::Characterwise,
+            kill_buffer_edit: None,
             vim_enabled: false,
             vim_mode: VimMode::Insert,
             vim_pending: VimPending::None,
@@ -632,6 +642,7 @@ impl TextArea {
     }
 
     pub fn input(&mut self, event: KeyEvent) {
+        self.kill_buffer_edit = None;
         // Only process key presses or repeats; ignore releases to avoid inserting
         // characters on key-up events when modifiers are no longer reported.
         if !matches!(event.kind, KeyEventKind::Press | KeyEventKind::Repeat) {
@@ -1322,7 +1333,17 @@ impl TextArea {
             return;
         }
         let text = self.kill_buffer.clone();
-        self.insert_str(&text);
+        let target = self.edit_target();
+        let start = target.start();
+        self.insert_str_at_target(target, &text);
+        // Rehydrating elements would invalidate Vim Replace's per-character recovery.
+        if !self.is_vim_replace_mode() {
+            self.kill_buffer_edit = Some(KillBufferEdit::Inserted(start));
+        }
+    }
+
+    pub(super) fn take_kill_buffer_edit(&mut self) -> Option<KillBufferEdit> {
+        self.kill_buffer_edit.take()
     }
 
     pub(crate) fn take_kill_buffer_snapshot(&mut self) -> KillBufferSnapshot {
@@ -1351,12 +1372,7 @@ impl TextArea {
             return;
         }
 
-        let removed = self.text[range.clone()].to_string();
-        if removed.is_empty() {
-            return;
-        }
-
-        self.store_kill_buffer(removed, kind);
+        self.store_kill_buffer(range.clone(), kind);
         self.replace_range_raw(range, "");
     }
 
@@ -1373,16 +1389,22 @@ impl TextArea {
         if range.start >= range.end {
             return;
         }
-        let removed = self.text[range].to_string();
-        if removed.is_empty() {
-            return;
-        }
-        self.store_kill_buffer(removed, kind);
+        self.store_kill_buffer(range, kind);
     }
 
-    fn store_kill_buffer(&mut self, text: String, kind: KillBufferKind) {
-        self.kill_buffer = text;
+    fn store_kill_buffer(&mut self, range: Range<usize>, kind: KillBufferKind) {
+        let elements = self
+            .text_element_ranges_overlapping(range.clone())
+            .map(|element| {
+                (
+                    element.start - range.start..element.end - range.start,
+                    self.text[element.clone()].to_string(),
+                )
+            })
+            .collect();
+        self.kill_buffer = self.text[range].to_string();
         self.kill_buffer_kind = kind;
+        self.kill_buffer_edit = Some(KillBufferEdit::Stored(elements));
     }
 
     fn paste_after_cursor(&mut self) {
@@ -1395,8 +1417,7 @@ impl TextArea {
         }
         let insert_at = self.next_atomic_boundary(self.cursor_pos);
         self.set_cursor(insert_at);
-        let text = self.kill_buffer.clone();
-        self.insert_str(&text);
+        self.yank();
     }
 
     fn paste_line_after_current_line(&mut self) {
@@ -1417,6 +1438,7 @@ impl TextArea {
             format!("\n{}", self.kill_buffer.trim_end_matches('\n'))
         };
         self.insert_str_at(insert_at, &text);
+        self.kill_buffer_edit = Some(KillBufferEdit::Inserted(cursor));
         self.set_cursor(cursor.min(self.text.len()));
     }