Reland "refactor transcript view to handle HistoryCells" (#3753)

Reland of #3538

Author: nornagon-openai

codex-rs/core/src/codex.rs

@@ -130,7 +130,7 @@ use codex_protocol::models::ResponseItem;
 use codex_protocol::models::ShellToolCallParams;
 use codex_protocol::protocol::InitialHistory;
 
-mod compact;
+pub mod compact;
 use self::compact::build_compacted_history;
 use self::compact::collect_user_messages;
 
@@ -710,7 +710,7 @@ impl Session {
         self.persist_rollout_items(&rollout_items).await;
     }
 
-    fn build_initial_context(&self, turn_context: &TurnContext) -> Vec<ResponseItem> {
+    pub(crate) fn build_initial_context(&self, turn_context: &TurnContext) -> Vec<ResponseItem> {
         let mut items = Vec::<ResponseItem>::with_capacity(2);
         if let Some(user_instructions) = turn_context.user_instructions.as_deref() {
             items.push(UserInstructions::new(user_instructions.to_string()).into());
@@ -3325,6 +3325,9 @@ async fn exit_review_mode(
         .await;
 }
 
+#[cfg(test)]
+pub(crate) use tests::make_session_and_context;
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -3565,7 +3568,7 @@ mod tests {
         })
     }
 
-    fn make_session_and_context() -> (Session, TurnContext) {
+    pub(crate) fn make_session_and_context() -> (Session, TurnContext) {
         let (tx_event, _rx_event) = async_channel::unbounded();
         let codex_home = tempfile::tempdir().expect("create temp dir");
         let config = Config::load_from_base_config_with_overrides(

codex-rs/core/src/codex/compact.rs

@@ -203,7 +203,7 @@ async fn run_compact_task_inner(
     sess.send_event(event).await;
 }
 
-fn content_items_to_text(content: &[ContentItem]) -> Option<String> {
+pub fn content_items_to_text(content: &[ContentItem]) -> Option<String> {
     let mut pieces = Vec::new();
     for item in content {
         match item {
@@ -235,7 +235,7 @@ pub(crate) fn collect_user_messages(items: &[ResponseItem]) -> Vec<String> {
         .collect()
 }
 
-fn is_session_prefix_message(text: &str) -> bool {
+pub fn is_session_prefix_message(text: &str) -> bool {
     matches!(
         InputMessageKind::from(("user", text)),
         InputMessageKind::UserInstructions | InputMessageKind::EnvironmentContext

codex-rs/core/src/conversation_manager.rs

@@ -3,6 +3,8 @@ use crate::CodexAuth;
 use crate::codex::Codex;
 use crate::codex::CodexSpawnOk;
 use crate::codex::INITIAL_SUBMIT_ID;
+use crate::codex::compact::content_items_to_text;
+use crate::codex::compact::is_session_prefix_message;
 use crate::codex_conversation::CodexConversation;
 use crate::config::Config;
 use crate::error::CodexErr;
@@ -134,19 +136,19 @@ impl ConversationManager {
         self.conversations.write().await.remove(conversation_id)
     }
 
-    /// Fork an existing conversation by dropping the last `drop_last_messages`
-    /// user/assistant messages from its transcript and starting a new
+    /// Fork an existing conversation by taking messages up to the given position
+    /// (not including the message at the given position) and starting a new
     /// conversation with identical configuration (unless overridden by the
     /// caller's `config`). The new conversation will have a fresh id.
     pub async fn fork_conversation(
         &self,
-        num_messages_to_drop: usize,
+        nth_user_message: usize,
         config: Config,
         path: PathBuf,
     ) -> CodexResult<NewConversation> {
         // Compute the prefix up to the cut point.
         let history = RolloutRecorder::get_rollout_history(&path).await?;
-        let history = truncate_after_dropping_last_messages(history, num_messages_to_drop);
+        let history = truncate_before_nth_user_message(history, nth_user_message);
 
         // Spawn a new conversation with the computed initial history.
         let auth_manager = self.auth_manager.clone();
@@ -159,33 +161,30 @@ impl ConversationManager {
     }
 }
 
-/// Return a prefix of `items` obtained by dropping the last `n` user messages
-/// and all items that follow them.
-fn truncate_after_dropping_last_messages(history: InitialHistory, n: usize) -> InitialHistory {
-    if n == 0 {
-        return InitialHistory::Forked(history.get_rollout_items());
-    }
-
-    // Work directly on rollout items, and cut the vector at the nth-from-last user message input.
+/// Return a prefix of `items` obtained by cutting strictly before the nth user message
+/// (0-based) and all items that follow it.
+fn truncate_before_nth_user_message(history: InitialHistory, n: usize) -> InitialHistory {
+    // Work directly on rollout items, and cut the vector at the nth user message input.
     let items: Vec<RolloutItem> = history.get_rollout_items();
 
     // Find indices of user message inputs in rollout order.
     let mut user_positions: Vec<usize> = Vec::new();
     for (idx, item) in items.iter().enumerate() {
-        if let RolloutItem::ResponseItem(ResponseItem::Message { role, .. }) = item
+        if let RolloutItem::ResponseItem(ResponseItem::Message { role, content, .. }) = item
             && role == "user"
+            && content_items_to_text(content).is_some_and(|text| !is_session_prefix_message(&text))
         {
             user_positions.push(idx);
         }
     }
 
-    // If fewer than n user messages exist, treat as empty.
-    if user_positions.len() < n {
+    // If fewer than or equal to n user messages exist, treat as empty (out of range).
+    if user_positions.len() <= n {
         return InitialHistory::New;
     }
 
-    // Cut strictly before the nth-from-last user message (do not keep the nth itself).
-    let cut_idx = user_positions[user_positions.len() - n];
+    // Cut strictly before the nth user message (do not keep the nth itself).
+    let cut_idx = user_positions[n];
     let rolled: Vec<RolloutItem> = items.into_iter().take(cut_idx).collect();
 
     if rolled.is_empty() {
@@ -198,9 +197,11 @@ fn truncate_after_dropping_last_messages(history: InitialHistory, n: usize) -> I
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::codex::make_session_and_context;
     use codex_protocol::models::ContentItem;
     use codex_protocol::models::ReasoningItemReasoningSummary;
     use codex_protocol::models::ResponseItem;
+    use pretty_assertions::assert_eq;
 
     fn user_msg(text: &str) -> ResponseItem {
         ResponseItem::Message {
@@ -252,7 +253,7 @@ mod tests {
             .cloned()
             .map(RolloutItem::ResponseItem)
             .collect();
-        let truncated = truncate_after_dropping_last_messages(InitialHistory::Forked(initial), 1);
+        let truncated = truncate_before_nth_user_message(InitialHistory::Forked(initial), 1);
         let got_items = truncated.get_rollout_items();
         let expected_items = vec![
             RolloutItem::ResponseItem(items[0].clone()),
@@ -269,7 +270,37 @@ mod tests {
             .cloned()
             .map(RolloutItem::ResponseItem)
             .collect();
-        let truncated2 = truncate_after_dropping_last_messages(InitialHistory::Forked(initial2), 2);
+        let truncated2 = truncate_before_nth_user_message(InitialHistory::Forked(initial2), 2);
         assert!(matches!(truncated2, InitialHistory::New));
     }
+
+    #[test]
+    fn ignores_session_prefix_messages_when_truncating() {
+        let (session, turn_context) = make_session_and_context();
+        let mut items = session.build_initial_context(&turn_context);
+        items.push(user_msg("feature request"));
+        items.push(assistant_msg("ack"));
+        items.push(user_msg("second question"));
+        items.push(assistant_msg("answer"));
+
+        let rollout_items: Vec<RolloutItem> = items
+            .iter()
+            .cloned()
+            .map(RolloutItem::ResponseItem)
+            .collect();
+
+        let truncated = truncate_before_nth_user_message(InitialHistory::Forked(rollout_items), 1);
+        let got_items = truncated.get_rollout_items();
+
+        let expected: Vec<RolloutItem> = vec![
+            RolloutItem::ResponseItem(items[0].clone()),
+            RolloutItem::ResponseItem(items[1].clone()),
+            RolloutItem::ResponseItem(items[2].clone()),
+        ];
+
+        assert_eq!(
+            serde_json::to_value(&got_items).unwrap(),
+            serde_json::to_value(&expected).unwrap()
+        );
+    }
 }

codex-rs/core/src/lib.rs

@@ -92,6 +92,8 @@ pub use client_common::Prompt;
 pub use client_common::REVIEW_PROMPT;
 pub use client_common::ResponseEvent;
 pub use client_common::ResponseStream;
+pub use codex::compact::content_items_to_text;
+pub use codex::compact::is_session_prefix_message;
 pub use codex_protocol::models::ContentItem;
 pub use codex_protocol::models::LocalShellAction;
 pub use codex_protocol::models::LocalShellExecAction;

codex-rs/core/tests/suite/compact_resume_fork.rs

@@ -74,7 +74,7 @@ async fn compact_resume_and_fork_preserve_model_history_view() {
         "compact+resume test expects resumed path {resumed_path:?} to exist",
     );
 
-    let forked = fork_conversation(&manager, &config, resumed_path, 1).await;
+    let forked = fork_conversation(&manager, &config, resumed_path, 2).await;
     user_turn(&forked, "AFTER_FORK").await;
 
     // 3. Capture the requests to the model and validate the history slices.
@@ -100,17 +100,15 @@ async fn compact_resume_and_fork_preserve_model_history_view() {
         "after-resume input should have at least as many items as after-compact",
     );
     assert_eq!(compact_arr.as_slice(), &resume_arr[..compact_arr.len()]);
-    eprint!(
-        "len of compact: {}, len of fork: {}",
-        compact_arr.len(),
-        fork_arr.len()
-    );
-    eprintln!("input_after_fork:{}", json!(input_after_fork));
+
     assert!(
         compact_arr.len() <= fork_arr.len(),
         "after-fork input should have at least as many items as after-compact",
     );
-    assert_eq!(compact_arr.as_slice(), &fork_arr[..compact_arr.len()]);
+    assert_eq!(
+        &compact_arr.as_slice()[..compact_arr.len()],
+        &fork_arr[..compact_arr.len()]
+    );
 
     let prompt = requests[0]["instructions"]
         .as_str()
@@ -824,14 +822,15 @@ async fn resume_conversation(
     conversation
 }
 
+#[cfg(test)]
 async fn fork_conversation(
     manager: &ConversationManager,
     config: &Config,
     path: std::path::PathBuf,
-    back_steps: usize,
+    nth_user_message: usize,
 ) -> Arc<CodexConversation> {
     let NewConversation { conversation, .. } = manager
-        .fork_conversation(back_steps, config.clone(), path)
+        .fork_conversation(nth_user_message, config.clone(), path)
         .await
         .expect("fork conversation");
     conversation

codex-rs/core/tests/suite/fork_conversation.rs

@@ -5,6 +5,8 @@ use codex_core::ModelProviderInfo;
 use codex_core::NewConversation;
 use codex_core::ResponseItem;
 use codex_core::built_in_model_providers;
+use codex_core::content_items_to_text;
+use codex_core::is_session_prefix_message;
 use codex_core::protocol::ConversationPathResponseEvent;
 use codex_core::protocol::EventMsg;
 use codex_core::protocol::InputItem;
@@ -104,13 +106,16 @@ async fn fork_conversation_twice_drops_to_first_message() {
         items
     };
 
-    // Compute expected prefixes after each fork by truncating base rollout at nth-from-last user input.
+    // Compute expected prefixes after each fork by truncating base rollout
+    // strictly before the nth user input (0-based).
     let base_items = read_items(&base_path);
     let find_user_input_positions = |items: &[RolloutItem]| -> Vec<usize> {
         let mut pos = Vec::new();
         for (i, it) in items.iter().enumerate() {
             if let RolloutItem::ResponseItem(ResponseItem::Message { role, content, .. }) = it
                 && role == "user"
+                && content_items_to_text(content)
+                    .is_some_and(|text| !is_session_prefix_message(&text))
             {
                 // Consider any user message as an input boundary; recorder stores both EventMsg and ResponseItem.
                 // We specifically look for input items, which are represented as ContentItem::InputText.
@@ -126,11 +131,8 @@ async fn fork_conversation_twice_drops_to_first_message() {
     };
     let user_inputs = find_user_input_positions(&base_items);
 
-    // After dropping last user input (n=1), cut strictly before that input if present, else empty.
-    let cut1 = user_inputs
-        .get(user_inputs.len().saturating_sub(1))
-        .copied()
-        .unwrap_or(0);
+    // After cutting at nth user input (n=1 → second user message), cut strictly before that input.
+    let cut1 = user_inputs.get(1).copied().unwrap_or(0);
     let expected_after_first: Vec<RolloutItem> = base_items[..cut1].to_vec();
 
     // After dropping again (n=1 on fork1), compute expected relative to fork1's rollout.
@@ -161,12 +163,12 @@ async fn fork_conversation_twice_drops_to_first_message() {
         serde_json::to_value(&expected_after_first).unwrap()
     );
 
-    // Fork again with n=1 → drops the (new) last user message, leaving only the first.
+    // Fork again with n=0 → drops the (new) last user message, leaving only the first.
     let NewConversation {
         conversation: codex_fork2,
         ..
     } = conversation_manager
-        .fork_conversation(1, config_for_fork.clone(), fork1_path.clone())
+        .fork_conversation(0, config_for_fork.clone(), fork1_path.clone())
         .await
         .expect("fork 2");
 

codex-rs/tui/src/app.rs

@@ -3,6 +3,7 @@ use crate::app_event::AppEvent;
 use crate::app_event_sender::AppEventSender;
 use crate::chatwidget::ChatWidget;
 use crate::file_search::FileSearchManager;
+use crate::history_cell::HistoryCell;
 use crate::pager_overlay::Overlay;
 use crate::resume_picker::ResumeSelection;
 use crate::tui;
@@ -52,7 +53,7 @@ pub(crate) struct App {
 
     pub(crate) file_search: FileSearchManager,
 
-    pub(crate) transcript_lines: Vec<Line<'static>>,
+    pub(crate) transcript_cells: Vec<Arc<dyn HistoryCell>>,
 
     // Pager overlay state (Transcript or Static like Diff)
     pub(crate) overlay: Option<Overlay>,
@@ -138,7 +139,7 @@ impl App {
             active_profile,
             file_search,
             enhanced_keys_supported,
-            transcript_lines: Vec::new(),
+            transcript_cells: Vec::new(),
             overlay: None,
             deferred_history_lines: Vec::new(),
             has_emitted_history_lines: false,
@@ -225,15 +226,12 @@ impl App {
                 tui.frame_requester().schedule_frame();
             }
             AppEvent::InsertHistoryCell(cell) => {
-                let mut cell_transcript = cell.transcript_lines();
-                if !cell.is_stream_continuation() && !self.transcript_lines.is_empty() {
-                    cell_transcript.insert(0, Line::from(""));
-                }
+                let cell: Arc<dyn HistoryCell> = cell.into();
                 if let Some(Overlay::Transcript(t)) = &mut self.overlay {
-                    t.insert_lines(cell_transcript.clone());
+                    t.insert_cell(cell.clone());
                     tui.frame_requester().schedule_frame();
                 }
-                self.transcript_lines.extend(cell_transcript.clone());
+                self.transcript_cells.push(cell.clone());
                 let mut display = cell.display_lines(tui.terminal.last_known_screen_size.width);
                 if !display.is_empty() {
                     // Only insert a separating blank line for new cells that are not
@@ -380,7 +378,7 @@ impl App {
             } => {
                 // Enter alternate screen and set viewport to full size.
                 let _ = tui.enter_alt_screen();
-                self.overlay = Some(Overlay::new_transcript(self.transcript_lines.clone()));
+                self.overlay = Some(Overlay::new_transcript(self.transcript_cells.clone()));
                 tui.frame_requester().schedule_frame();
             }
             // Esc primes/advances backtracking only in normal (not working) mode
@@ -405,7 +403,7 @@ impl App {
                 kind: KeyEventKind::Press,
                 ..
             } if self.backtrack.primed
-                && self.backtrack.count > 0
+                && self.backtrack.nth_user_message != usize::MAX
                 && self.chat_widget.composer_is_empty() =>
             {
                 // Delegate to helper for clarity; preserves behavior.
@@ -439,7 +437,6 @@ mod tests {
     use codex_core::AuthManager;
     use codex_core::CodexAuth;
     use codex_core::ConversationManager;
-    use ratatui::text::Line;
     use std::sync::Arc;
     use std::sync::atomic::AtomicBool;
 
@@ -462,7 +459,7 @@ mod tests {
             config,
             active_profile: None,
             file_search,
-            transcript_lines: Vec::<Line<'static>>::new(),
+            transcript_cells: Vec::new(),
             overlay: None,
             deferred_history_lines: Vec::new(),
             has_emitted_history_lines: false,