From 935d88b455ba51d7103df85305fcf3b418592751 Mon Sep 17 00:00:00 2001 From: charley-oai Date: Fri, 23 Jan 2026 10:18:19 -0800 Subject: [PATCH] Persist text element ranges and attached images across history/resume (#9116) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit **Summary** - Backtrack selection now rehydrates `text_elements` and `local_image_paths` from the chosen user history cell so Esc‑Esc history edits preserve image placeholders and attachments. - Composer prefill uses the preserved elements/attachments in both `tui` and `tui2`. - Extended backtrack selection tests to cover image placeholder elements and local image paths. **Changes** - `tui/src/app_backtrack.rs`: Backtrack selection now carries text elements + local image paths; composer prefill uses them (removes TODO). - `tui2/src/app_backtrack.rs`: Same as above. - `tui/src/app.rs`: Updated backtrack test to assert restored elements/paths. - `tui2/src/app.rs`: Same test updates. ### The original scope of this PR (threading text elements and image attachments through the codex harness thoroughly/persistently) was broken into the following PRs other than this one: The diff of this PR was reduced by changing types in a starter PR: https://github.com/openai/codex/pull/9235 Then text element metadata was added to protocol, app server, and core in this PR: https://github.com/openai/codex/pull/9331 Then the end-to-end flow was completed by wiring TUI/TUI2 input, history, and restore behavior in https://github.com/openai/codex/pull/9393 Prompt expansion was supported in this PR: https://github.com/openai/codex/pull/9518 TextElement optional placeholder field was protected in https://github.com/openai/codex/pull/9545 --- codex-rs/tui/src/app.rs | 32 +++++++++++++----- codex-rs/tui/src/app_backtrack.rs | 33 ++++++++++++++----- codex-rs/tui/src/bottom_pane/chat_composer.rs | 14 +++++--- 3 files changed, 57 insertions(+), 22 deletions(-) diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index fa24e110f..6b6b0188c 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -2203,6 +2203,7 @@ mod tests { use codex_core::protocol::SessionSource; use codex_otel::OtelManager; use codex_protocol::ThreadId; + use codex_protocol::user_input::TextElement; use insta::assert_snapshot; use pretty_assertions::assert_eq; use ratatui::prelude::Line; @@ -2503,11 +2504,14 @@ mod tests { async fn backtrack_selection_with_duplicate_history_targets_unique_turn() { let (mut app, _app_event_rx, mut op_rx) = make_test_app_with_channels().await; - let user_cell = |text: &str| -> Arc { + let user_cell = |text: &str, + text_elements: Vec, + local_image_paths: Vec| + -> Arc { Arc::new(UserHistoryCell { message: text.to_string(), - text_elements: Vec::new(), - local_image_paths: Vec::new(), + text_elements, + local_image_paths, }) as Arc }; let agent_cell = |text: &str| -> Arc { @@ -2545,18 +2549,28 @@ mod tests { )) as Arc }; + let placeholder = "[Image #1]"; + let edited_text = format!("follow-up (edited) {placeholder}"); + let edited_range = edited_text.len().saturating_sub(placeholder.len())..edited_text.len(); + let edited_text_elements = vec![TextElement::new(edited_range.into(), None)]; + let edited_local_image_paths = vec![PathBuf::from("/tmp/fake-image.png")]; + // Simulate a transcript with duplicated history (e.g., from prior backtracks) // and an edited turn appended after a session header boundary. app.transcript_cells = vec![ make_header(true), - user_cell("first question"), + user_cell("first question", Vec::new(), Vec::new()), agent_cell("answer first"), - user_cell("follow-up"), + user_cell("follow-up", Vec::new(), Vec::new()), agent_cell("answer follow-up"), make_header(false), - user_cell("first question"), + user_cell("first question", Vec::new(), Vec::new()), agent_cell("answer first"), - user_cell("follow-up (edited)"), + user_cell( + &edited_text, + edited_text_elements.clone(), + edited_local_image_paths.clone(), + ), agent_cell("answer edited"), ]; @@ -2589,7 +2603,9 @@ mod tests { .confirm_backtrack_from_main() .expect("backtrack selection"); assert_eq!(selection.nth_user_message, 1); - assert_eq!(selection.prefill, "follow-up (edited)"); + assert_eq!(selection.prefill, edited_text); + assert_eq!(selection.text_elements, edited_text_elements); + assert_eq!(selection.local_image_paths, edited_local_image_paths); app.apply_backtrack_rollback(selection); diff --git a/codex-rs/tui/src/app_backtrack.rs b/codex-rs/tui/src/app_backtrack.rs index 9a77b6bed..256355dd0 100644 --- a/codex-rs/tui/src/app_backtrack.rs +++ b/codex-rs/tui/src/app_backtrack.rs @@ -24,6 +24,7 @@ //! both committed history and in-flight activity without changing flush or coalescing behavior. use std::any::TypeId; +use std::path::PathBuf; use std::sync::Arc; use crate::app::App; @@ -37,6 +38,7 @@ use codex_core::protocol::ErrorEvent; use codex_core::protocol::EventMsg; use codex_core::protocol::Op; use codex_protocol::ThreadId; +use codex_protocol::user_input::TextElement; use color_eyre::eyre::Result; use crossterm::event::KeyCode; use crossterm::event::KeyEvent; @@ -78,6 +80,10 @@ pub(crate) struct BacktrackSelection { /// This is applied immediately on selection confirmation; if the rollback fails, the prefill /// remains as a convenience so the user can retry or edit. pub(crate) prefill: String, + /// Text elements associated with the selected user message. + pub(crate) text_elements: Vec, + /// Local image paths associated with the selected user message. + pub(crate) local_image_paths: Vec, } /// An in-flight rollback requested from core. @@ -198,16 +204,16 @@ impl App { } let prefill = selection.prefill.clone(); + let text_elements = selection.text_elements.clone(); + let local_image_paths = selection.local_image_paths.clone(); self.backtrack.pending_rollback = Some(PendingBacktrackRollback { selection, thread_id: self.chat_widget.thread_id(), }); self.chat_widget.submit_op(Op::ThreadRollback { num_turns }); - if !prefill.is_empty() { - // TODO: Rehydrate text_elements/local_image_paths from the selected user cell so - // backtrack preserves image placeholders and attachments. + if !prefill.is_empty() || !text_elements.is_empty() || !local_image_paths.is_empty() { self.chat_widget - .set_composer_text(prefill, Vec::new(), Vec::new()); + .set_composer_text(prefill, text_elements, local_image_paths); } } @@ -481,15 +487,24 @@ impl App { return None; } - let prefill = nth_user_position(&self.transcript_cells, nth_user_message) - .and_then(|idx| self.transcript_cells.get(idx)) - .and_then(|cell| cell.as_any().downcast_ref::()) - .map(|c| c.message.clone()) - .unwrap_or_default(); + let (prefill, text_elements, local_image_paths) = + nth_user_position(&self.transcript_cells, nth_user_message) + .and_then(|idx| self.transcript_cells.get(idx)) + .and_then(|cell| cell.as_any().downcast_ref::()) + .map(|cell| { + ( + cell.message.clone(), + cell.text_elements.clone(), + cell.local_image_paths.clone(), + ) + }) + .unwrap_or_else(|| (String::new(), Vec::new(), Vec::new())); Some(BacktrackSelection { nth_user_message, prefill, + text_elements, + local_image_paths, }) } diff --git a/codex-rs/tui/src/bottom_pane/chat_composer.rs b/codex-rs/tui/src/bottom_pane/chat_composer.rs index 06c2808f2..47c2c4625 100644 --- a/codex-rs/tui/src/bottom_pane/chat_composer.rs +++ b/codex-rs/tui/src/bottom_pane/chat_composer.rs @@ -212,6 +212,7 @@ pub(crate) struct ChatComposer { pending_pastes: Vec<(String, String)>, large_paste_counters: HashMap, has_focus: bool, + /// Invariant: attached images are labeled `[Image #1]..[Image #N]` in vec order. attached_images: Vec, placeholder_text: String, is_task_running: bool, @@ -475,8 +476,9 @@ impl ChatComposer { /// Replace the composer content with text from an external editor. /// Clears pending paste placeholders and keeps only attachments whose - /// placeholder labels still appear in the new text. Cursor is placed at - /// the end after rebuilding elements. + /// placeholder labels still appear in the new text. Image placeholders + /// are renumbered to `[Image #1]..[Image #N]`. Cursor is placed at the end + /// after rebuilding elements. pub(crate) fn apply_external_edit(&mut self, text: String) { self.pending_pastes.clear(); @@ -538,6 +540,8 @@ impl ChatComposer { self.textarea.insert_str(&text[idx..]); } + // Keep image placeholders normalized to [Image #1].. in attachment order. + self.relabel_attached_images_and_update_placeholders(); self.textarea.set_cursor(self.textarea.text().len()); self.sync_popups(); } @@ -6081,7 +6085,7 @@ mod tests { false, ); - let placeholder = "[image 10x10]".to_string(); + let placeholder = local_image_label_text(1); composer.textarea.insert_element(&placeholder); composer.attached_images.push(AttachedImage { placeholder: placeholder.clone(), @@ -6115,7 +6119,7 @@ mod tests { false, ); - let placeholder = "[image 10x10]".to_string(); + let placeholder = local_image_label_text(1); composer.textarea.insert_element(&placeholder); composer.attached_images.push(AttachedImage { placeholder: placeholder.clone(), @@ -6165,7 +6169,7 @@ mod tests { false, ); - let placeholder = "[image 10x10]".to_string(); + let placeholder = local_image_label_text(1); composer.textarea.insert_element(&placeholder); composer.attached_images.push(AttachedImage { placeholder: placeholder.clone(),