diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index fa5fca382252..1cd658724997 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -32,13 +32,15 @@ use codex_utils_output_truncation::truncate_function_output_items_with_policy; use codex_utils_output_truncation::truncate_text; use std::num::NonZeroUsize; use std::ops::Deref; +use std::sync::Arc; use std::sync::LazyLock; /// Transcript of thread history #[derive(Debug, Clone, Default)] pub(crate) struct ContextManager { - /// The oldest items are at the beginning of the vector. - items: Vec, + /// The oldest items are at the beginning of the vector. Snapshots share the vector until a + /// caller needs to mutate it, avoiding deep copies for read-only history consumers. + items: Arc>, /// Bumped whenever history is rewritten, such as compaction or rollback. history_version: u64, token_info: Option, @@ -60,7 +62,7 @@ pub(crate) struct ContextManager { impl ContextManager { pub(crate) fn new() -> Self { Self { - items: Vec::new(), + items: Arc::new(Vec::new()), history_version: 0, token_info: TokenUsageInfo::new_or_append( &None, &None, /*model_context_window*/ None, @@ -130,8 +132,8 @@ impl ContextManager { continue; } - let processed = self.process_item(item_ref, policy); - self.items.push(processed); + let processed = Self::process_item(item_ref, policy); + Arc::make_mut(&mut self.items).push(processed); } } @@ -140,7 +142,7 @@ impl ContextManager { /// is stripped from messages and tool outputs according to `input_modalities`. pub(crate) fn for_prompt(mut self, input_modalities: &[InputModality]) -> Vec { self.normalize_history(input_modalities); - self.items + Arc::unwrap_or_clone(self.items) } /// Returns raw items in the history. @@ -150,7 +152,7 @@ impl ContextManager { /// Returns raw items in the history and consumes the snapshot. pub(crate) fn into_raw_items(self) -> Vec { - self.items + Arc::unwrap_or_clone(self.items) } pub(crate) fn history_version(&self) -> u64 { @@ -188,17 +190,18 @@ impl ContextManager { if !self.items.is_empty() { // Remove the oldest item (front of the list). Items are ordered from // oldest → newest, so index 0 is the first entry recorded. - let removed = self.items.remove(0); + let items = Arc::make_mut(&mut self.items); + let removed = items.remove(0); // If the removed item participates in a call/output pair, also remove // its corresponding counterpart to keep the invariants intact without // running a full normalization pass. - normalize::remove_corresponding_for(&mut self.items, &removed); + normalize::remove_corresponding_for(items, &removed); self.world_state_baseline = None; } } pub(crate) fn replace(&mut self, items: Vec) { - self.items = items; + self.items = Arc::new(items); self.history_version = self.history_version.saturating_add(1); self.world_state_baseline = None; } @@ -227,7 +230,7 @@ impl ContextManager { let snapshot = self.items.clone(); let user_positions = user_message_positions(&snapshot); let Some(&first_instruction_turn_idx) = user_positions.first() else { - self.replace(snapshot); + self.replace(Arc::unwrap_or_clone(snapshot)); return; }; @@ -323,20 +326,22 @@ impl ContextManager { /// 2. every output has a corresponding call entry /// 3. unsupported image and audio content is stripped from messages and tool outputs fn normalize_history(&mut self, input_modalities: &[InputModality]) { + let items = Arc::make_mut(&mut self.items); + // all function/tool calls must have a corresponding output - normalize::ensure_call_outputs_present(&mut self.items); + normalize::ensure_call_outputs_present(items); // all outputs must have a corresponding function/tool call - normalize::remove_orphan_outputs(&mut self.items); + normalize::remove_orphan_outputs(items); // strip images when model does not support them - normalize::strip_images_when_unsupported(input_modalities, &mut self.items); + normalize::strip_images_when_unsupported(input_modalities, items); // strip audio when model does not support it - normalize::strip_audio_when_unsupported(input_modalities, &mut self.items); + normalize::strip_audio_when_unsupported(input_modalities, items); } - fn process_item(&self, item: &ResponseItem, policy: TruncationPolicy) -> ResponseItem { + fn process_item(item: &ResponseItem, policy: TruncationPolicy) -> ResponseItem { let policy_with_serialization_budget = policy * 1.2; match item { ResponseItem::FunctionCallOutput { diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index 86f9cea51cc1..027bdee1a7ab 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -444,6 +444,31 @@ fn for_prompt_preserves_inter_agent_assistant_messages() { assert_eq!(history.for_prompt(&default_input_modalities()), vec![item]); } +#[test] +fn cloned_history_shares_items_until_mutated() { + let first = assistant_msg(&"first ".repeat(1_024)); + let second = assistant_msg("second"); + let history = create_history_with_items(vec![first.clone()]); + let mut snapshot = history.clone(); + + assert!(std::ptr::eq( + history.raw_items().as_ptr(), + snapshot.raw_items().as_ptr() + )); + + snapshot.record_items( + std::slice::from_ref(&second), + TruncationPolicy::Tokens(10_000), + ); + + assert!(!std::ptr::eq( + history.raw_items().as_ptr(), + snapshot.raw_items().as_ptr() + )); + assert_eq!(history.raw_items(), std::slice::from_ref(&first)); + assert_eq!(snapshot.raw_items(), &[first, second]); +} + #[test] fn drop_last_n_user_turns_treats_inter_agent_assistant_messages_as_instruction_turns() { let first_turn = user_input_text_msg("first"); diff --git a/codex-rs/core/src/realtime_context.rs b/codex-rs/core/src/realtime_context.rs index 0e400cb106c8..ec881bdcc963 100644 --- a/codex-rs/core/src/realtime_context.rs +++ b/codex-rs/core/src/realtime_context.rs @@ -62,8 +62,10 @@ pub(crate) async fn build_realtime_startup_context( ) -> Option { let config = sess.get_config().await; let cwd = config.cwd.clone(); - let history = sess.clone_history().await; - let current_thread_section = build_current_thread_section(history.raw_items()); + let current_thread_section = { + let history = sess.clone_history().await; + build_current_thread_section(history.raw_items()) + }; let recent_threads = load_recent_threads(sess).await; let recent_work_section = build_recent_work_section(&cwd, &recent_threads).await; let workspace_section = build_workspace_section_with_user_root(&cwd, home_dir()).await; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index b96c190e7b5f..a3c714207a2d 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -1217,8 +1217,8 @@ impl Session { &self, turn_context: &TurnContext, ) -> Option { - let state = self.state.lock().await; - state.history.estimate_token_count(turn_context) + let history = self.clone_history().await; + history.estimate_token_count(turn_context) } pub(crate) async fn get_base_instructions(&self) -> BaseInstructions {