diff --git a/codex-rs/app-server-protocol/src/protocol/thread_history.rs b/codex-rs/app-server-protocol/src/protocol/thread_history.rs index 32cfe744e8..7a1d70c85f 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history.rs @@ -138,17 +138,6 @@ impl ThreadHistoryTurnChange { duration_ms: turn.duration_ms, } } - - fn from_turn(turn: &Turn) -> Self { - Self { - turn_id: turn.id.clone(), - status: turn.status.clone(), - error: turn.error.clone(), - started_at: turn.started_at, - completed_at: turn.completed_at, - duration_ms: turn.duration_ms, - } - } } /// Coalesces per-rollout-item changes into an end-of-batch view. It preserves @@ -234,7 +223,9 @@ impl ThreadHistoryChangeAccumulator { } pub struct ThreadHistoryBuilder { - turns: Vec, + // Retain the builder representation so late completions can reuse each + // finished turn's item index without adding it to the public Turn type. + turns: Vec, current_turn: Option, next_item_index: i64, current_rollout_index: usize, @@ -266,14 +257,14 @@ impl ThreadHistoryBuilder { pub fn finish(mut self) -> Vec { self.finish_current_turn(); - self.turns + self.turns.into_iter().map(Turn::from).collect() } pub fn active_turn_snapshot(&self) -> Option { self.current_turn .as_ref() .map(Turn::from) - .or_else(|| self.turns.last().cloned()) + .or_else(|| self.turns.last().map(Turn::from)) } /// Returns the id of the active turn without materializing its items. @@ -289,7 +280,12 @@ impl ThreadHistoryBuilder { .as_ref() .filter(|turn| turn.id == turn_id) .map(Turn::from) - .or_else(|| self.turns.iter().find(|turn| turn.id == turn_id).cloned()) + .or_else(|| { + self.turns + .iter() + .find(|turn| turn.id == turn_id) + .map(Turn::from) + }) } /// Returns the index of the active turn snapshot within the finished turn list. @@ -1226,7 +1222,7 @@ impl ThreadHistoryBuilder { turn.status = TurnStatus::Interrupted; turn.completed_at = payload.completed_at; turn.duration_ms = payload.duration_ms; - let changed_turn = ThreadHistoryTurnChange::from_turn(turn); + let changed_turn = ThreadHistoryTurnChange::from_pending_turn(turn); self.record_changed_turn(changed_turn); return; } @@ -1294,7 +1290,7 @@ impl ThreadHistoryBuilder { } turn.completed_at = payload.completed_at; turn.duration_ms = payload.duration_ms; - let changed_turn = ThreadHistoryTurnChange::from_turn(turn); + let changed_turn = ThreadHistoryTurnChange::from_pending_turn(turn); self.record_changed_turn(changed_turn); return; } @@ -1347,7 +1343,7 @@ impl ThreadHistoryBuilder { if turn.items.is_empty() && !turn.opened_explicitly && !turn.saw_compaction { return; } - self.turns.push(Turn::from(turn)); + self.turns.push(turn); } } @@ -1362,6 +1358,7 @@ impl ThreadHistoryBuilder { PendingTurn { id, items: Vec::new(), + item_index: TurnItemIndex::default(), error: None, status: TurnStatus::Completed, started_at: None, @@ -1392,7 +1389,7 @@ impl ThreadHistoryBuilder { let changed_item = { let turn = self.ensure_turn(); let changed_item = tracking_changes.then(|| (turn.id.clone(), item.clone())); - turn.items.push(item); + turn.item_index.push(&mut turn.items, item); changed_item }; if let Some((turn_id, item)) = changed_item { @@ -1406,7 +1403,7 @@ impl ThreadHistoryBuilder { && turn.id == turn_id { let changed_item = { - let item = upsert_turn_item(&mut turn.items, item); + let item = turn.item_index.upsert(&mut turn.items, item); tracking_changes.then(|| (turn.id.clone(), item.clone())) }; if let Some((turn_id, item)) = changed_item { @@ -1417,7 +1414,7 @@ impl ThreadHistoryBuilder { if let Some(turn) = self.turns.iter_mut().find(|turn| turn.id == turn_id) { let changed_item = { - let item = upsert_turn_item(&mut turn.items, item); + let item = turn.item_index.upsert(&mut turn.items, item); tracking_changes.then(|| (turn.id.clone(), item.clone())) }; if let Some((turn_id, item)) = changed_item { @@ -1436,7 +1433,7 @@ impl ThreadHistoryBuilder { let tracking_changes = self.is_tracking_changes(); let changed_item = { let turn = self.ensure_turn(); - let item = upsert_turn_item(&mut turn.items, item); + let item = turn.item_index.upsert(&mut turn.items, item); tracking_changes.then(|| (turn.id.clone(), item.clone())) }; if let Some((turn_id, item)) = changed_item { @@ -1545,22 +1542,54 @@ fn convert_dynamic_tool_content_items( .collect() } -fn upsert_turn_item(items: &mut Vec, item: ThreadItem) -> &ThreadItem { - if let Some(existing_item_index) = items - .iter() - .position(|existing_item| existing_item.id() == item.id()) - { - items[existing_item_index] = item; - return &items[existing_item_index]; +const TURN_ITEM_INDEX_THRESHOLD: usize = 32; + +/// Lazily indexes a turn's append-only item list. Replacements keep the same ID +/// and position; duplicate IDs continue to resolve to their first occurrence. +/// Keep this index with its turn, including after completion and across rollback. +#[derive(Default)] +struct TurnItemIndex { + positions: Option>, +} + +impl TurnItemIndex { + fn push(&mut self, items: &mut Vec, item: ThreadItem) { + if let Some(positions) = &mut self.positions { + positions + .entry(item.id().to_string()) + .or_insert(items.len()); + } + items.push(item); + } + + fn upsert<'a>(&mut self, items: &'a mut Vec, item: ThreadItem) -> &'a ThreadItem { + if self.positions.is_none() && items.len() >= TURN_ITEM_INDEX_THRESHOLD { + let mut positions = HashMap::with_capacity(items.len()); + for (index, existing) in items.iter().enumerate() { + positions.entry(existing.id().to_string()).or_insert(index); + } + self.positions = Some(positions); + } + + let existing_index = match &self.positions { + Some(positions) => positions.get(item.id()).copied(), + None => items.iter().position(|existing| existing.id() == item.id()), + }; + if let Some(index) = existing_index { + items[index] = item; + &items[index] + } else { + let index = items.len(); + self.push(items, item); + &items[index] + } } - let inserted_item_index = items.len(); - items.push(item); - &items[inserted_item_index] } struct PendingTurn { id: String, items: Vec, + item_index: TurnItemIndex, error: Option, status: TurnStatus, started_at: Option, @@ -1678,6 +1707,41 @@ mod tests { use std::time::Duration; use uuid::Uuid; + fn indexed_sleep_item(id: &str, duration_ms: u64) -> ThreadItem { + ThreadItem::Sleep(CoreSleepItem { + id: id.to_string(), + duration_ms, + }) + } + + #[test] + fn pushed_duplicates_resolve_to_first_occurrence_before_and_after_indexing() { + for count in [2, TURN_ITEM_INDEX_THRESHOLD, TURN_ITEM_INDEX_THRESHOLD * 2] { + let mut index = TurnItemIndex::default(); + let mut items = Vec::new(); + let original = indexed_sleep_item("duplicate", /*duration_ms*/ 1); + for _ in 0..count { + index.push(&mut items, original.clone()); + } + let updated = indexed_sleep_item("duplicate", /*duration_ms*/ 2); + index.upsert(&mut items, updated.clone()); + // Appends after index activation must also keep the first position. + index.push(&mut items, original.clone()); + index.upsert(&mut items, updated.clone()); + let appended = indexed_sleep_item("appended", /*duration_ms*/ 3); + index.push( + &mut items, + indexed_sleep_item("appended", /*duration_ms*/ 1), + ); + index.upsert(&mut items, appended.clone()); + + let mut expected = vec![original; count + 1]; + expected[0] = updated; + expected.push(appended); + assert_eq!(items, expected); + } + } + #[test] fn builds_multiple_turns_with_reasoning_items() { let events = vec![ @@ -2063,6 +2127,27 @@ mod tests { ); } + #[test] + fn indexed_items_stay_with_their_turn_after_rollback() { + let mut builder = ThreadHistoryBuilder::new(); + let mut expected: Vec<_> = (0..64) + .map(|i| indexed_sleep_item(&i.to_string(), /*duration_ms*/ 1)) + .collect(); + for turn_id in ["retained", "discarded"] { + builder.ensure_turn().id = turn_id.into(); + for item in &expected { + builder.upsert_item_in_current_turn(item.clone()); + } + builder.finish_current_turn(); + // Reuse IDs at different positions in the other turn. + expected.reverse(); + } + builder.handle_thread_rollback(&ThreadRolledBackEvent { num_turns: 1 }); + expected[0] = indexed_sleep_item("0", /*duration_ms*/ 2); + builder.upsert_item_in_turn_id("retained", expected[0].clone()); + assert_eq!(builder.finish()[0].items, expected); + } + #[test] fn rebuilds_extension_image_generation_item_from_persisted_completion() { let turn_id = "turn-1"; diff --git a/codex-rs/tui/src/bottom_pane/chat_composer/history_search.rs b/codex-rs/tui/src/bottom_pane/chat_composer/history_search.rs index 5217f5aac9..e1141c7768 100644 --- a/codex-rs/tui/src/bottom_pane/chat_composer/history_search.rs +++ b/codex-rs/tui/src/bottom_pane/chat_composer/history_search.rs @@ -419,24 +419,21 @@ impl ChatComposer { let mut ranges = Vec::new(); let mut search_from = 0; + // Use two-pointer method to find matches in linear time. + let mut start_span = 0; + let mut end_span = 0; while search_from <= folded.len() && let Some(relative_start) = folded[search_from..].find(&query_lower) { let folded_start = search_from + relative_start; let folded_end = folded_start + query_lower.len(); - if let Some((_, first_original)) = folded_spans.iter().find(|(folded_range, _)| { - folded_range.end > folded_start && folded_range.start < folded_end - }) { - let original_end = folded_spans - .iter() - .rev() - .find(|(folded_range, _)| { - folded_range.end > folded_start && folded_range.start < folded_end - }) - .map(|(_, original_range)| original_range.end) - .unwrap_or(first_original.end); - ranges.push(first_original.start..original_end); + while folded_spans[start_span].0.end <= folded_start { + start_span += 1; } + while folded_spans[end_span].0.end < folded_end { + end_span += 1; + } + ranges.push(folded_spans[start_span].1.start..folded_spans[end_span].1.end); search_from = folded_end; } ranges @@ -569,6 +566,25 @@ mod tests { assert!(ChatComposer::case_insensitive_match_ranges("git", "").is_empty()); } + #[test] + fn history_search_match_ranges_preserve_unicode_boundaries() { + for (text, query, expected) in [ + ("İİ", "i", vec![0..2, 2..4]), + ("İİ", "\u{307}", vec![0..2, 2..4]), + ("İİİ", "\u{307}i", vec![0..4, 2..6]), + ("éÉ é", "É", vec![0..2, 2..4, 5..7]), + ("aaaaa", "aa", vec![0..2, 2..4]), + ("", "x", vec![]), + ("abc", "z", vec![]), + ] { + assert_eq!( + ChatComposer::case_insensitive_match_ranges(text, query), + expected, + "text: {text:?}, query: {query:?}" + ); + } + } + #[test] fn history_search_accepts_matching_entry() { let (tx, _rx) = unbounded_channel::();