From 316637aa3410fc6e72c0abd3df06d0ba802cb661 Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Sun, 8 Feb 2026 21:37:07 -0800 Subject: [PATCH] fix --- codex-rs/core/src/codex.rs | 42 ++++- .../context_manager/context_discoverable.rs | 123 ++++++++++++- codex-rs/core/src/context_manager/history.rs | 49 ++++-- .../core/src/context_manager/history_tests.rs | 161 +++++++++++++++++- codex-rs/core/src/state/session.rs | 10 +- codex-rs/core/tests/suite/mod.rs | 1 + 6 files changed, 359 insertions(+), 27 deletions(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 11fbeb6498..bf2aeb4161 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -1044,7 +1044,11 @@ impl Session { } }; session_configuration.thread_name = thread_name.clone(); - let state = SessionState::new(session_configuration.clone()); + let state = SessionState::new( + session_configuration.clone(), + conversation_id.to_string(), + config.codex_home.clone(), + ); let services = SessionServices { mcp_connection_manager: Arc::new(RwLock::new(McpConnectionManager::default())), @@ -2029,7 +2033,10 @@ impl Session { turn_context: &TurnContext, rollout_items: &[RolloutItem], ) -> Vec { - let mut history = ContextManager::new(); + let mut history = ContextManager::new( + Arc::new(self.conversation_id.to_string()), + Arc::new(turn_context.config.codex_home.clone()), + ); for item in rollout_items { match item { RolloutItem::ResponseItem(response_item) => { @@ -5483,6 +5490,7 @@ mod tests { let codex_home = tempfile::tempdir().expect("create temp dir"); let config = build_test_config(codex_home.path()).await; let config = Arc::new(config); + let conversation_id = ThreadId::default(); let model = ModelsManager::get_model_offline(config.model.as_deref()); let model_info = ModelsManager::construct_model_info_offline(model.as_str(), &config); let reasoning_effort = config.model_reasoning_effort; @@ -5517,7 +5525,11 @@ mod tests { dynamic_tools: Vec::new(), }; - let mut state = SessionState::new(session_configuration); + let mut state = SessionState::new( + session_configuration, + conversation_id.to_string(), + config.codex_home.clone(), + ); let initial = RateLimitSnapshot { primary: Some(RateLimitWindow { used_percent: 10.0, @@ -5566,6 +5578,7 @@ mod tests { let codex_home = tempfile::tempdir().expect("create temp dir"); let config = build_test_config(codex_home.path()).await; let config = Arc::new(config); + let conversation_id = ThreadId::default(); let model = ModelsManager::get_model_offline(config.model.as_deref()); let model_info = ModelsManager::construct_model_info_offline(model.as_str(), &config); let reasoning_effort = config.model_reasoning_effort; @@ -5600,7 +5613,11 @@ mod tests { dynamic_tools: Vec::new(), }; - let mut state = SessionState::new(session_configuration); + let mut state = SessionState::new( + session_configuration, + conversation_id.to_string(), + config.codex_home.clone(), + ); let initial = RateLimitSnapshot { primary: Some(RateLimitWindow { used_percent: 15.0, @@ -5885,7 +5902,11 @@ mod tests { session_configuration.session_source.clone(), ); - let mut state = SessionState::new(session_configuration.clone()); + let mut state = SessionState::new( + session_configuration.clone(), + conversation_id.to_string(), + config.codex_home.clone(), + ); mark_state_initial_context_seeded(&mut state); let skills_manager = Arc::new(SkillsManager::new(config.codex_home.clone())); @@ -6017,7 +6038,11 @@ mod tests { session_configuration.session_source.clone(), ); - let mut state = SessionState::new(session_configuration.clone()); + let mut state = SessionState::new( + session_configuration.clone(), + conversation_id.to_string(), + config.codex_home.clone(), + ); mark_state_initial_context_seeded(&mut state); let skills_manager = Arc::new(SkillsManager::new(config.codex_home.clone())); @@ -6516,7 +6541,10 @@ mod tests { turn_context: &TurnContext, ) -> (Vec, Vec) { let mut rollout_items = Vec::new(); - let mut live_history = ContextManager::new(); + let mut live_history = ContextManager::new( + Arc::new(session.conversation_id.to_string()), + Arc::new(turn_context.config.codex_home.clone()), + ); let initial_context = session.build_initial_context(turn_context).await; for item in &initial_context { diff --git a/codex-rs/core/src/context_manager/context_discoverable.rs b/codex-rs/core/src/context_manager/context_discoverable.rs index 3689cd614f..0a5f490cf1 100644 --- a/codex-rs/core/src/context_manager/context_discoverable.rs +++ b/codex-rs/core/src/context_manager/context_discoverable.rs @@ -1,18 +1,38 @@ use crate::truncate::TruncationPolicy; +use crate::truncate::approx_token_count; use crate::truncate::truncate_function_output_items_with_policy; use crate::truncate::truncate_text; +use codex_protocol::models::ContentItem; use codex_protocol::models::CustomToolCallOutput; use codex_protocol::models::FunctionCallOutput; use codex_protocol::models::FunctionCallOutputBody; use codex_protocol::models::FunctionCallOutputPayload; +use codex_protocol::models::Message; use codex_protocol::models::ResponseItem; +use sha2::Digest; +use sha2::Sha256; +use std::path::Path; pub(super) trait ContextDiscoverable { - fn discoverable_history_item(&self, policy: TruncationPolicy) -> ResponseItem; + fn discoverable_history_item( + &self, + policy: TruncationPolicy, + codex_home: Option<&Path>, + thread_id: Option<&str>, + context_window_tokens: Option, + current_usage_tokens: i64, + ) -> ResponseItem; } impl ContextDiscoverable for FunctionCallOutput { - fn discoverable_history_item(&self, policy: TruncationPolicy) -> ResponseItem { + fn discoverable_history_item( + &self, + policy: TruncationPolicy, + _codex_home: Option<&Path>, + _thread_id: Option<&str>, + _context_window_tokens: Option, + _current_usage_tokens: i64, + ) -> ResponseItem { let body = match &self.output.body { FunctionCallOutputBody::Text(content) => { FunctionCallOutputBody::Text(truncate_text(content, policy)) @@ -33,10 +53,107 @@ impl ContextDiscoverable for FunctionCallOutput { } impl ContextDiscoverable for CustomToolCallOutput { - fn discoverable_history_item(&self, policy: TruncationPolicy) -> ResponseItem { + fn discoverable_history_item( + &self, + policy: TruncationPolicy, + _codex_home: Option<&Path>, + _thread_id: Option<&str>, + _context_window_tokens: Option, + _current_usage_tokens: i64, + ) -> ResponseItem { ResponseItem::CustomToolCallOutput(CustomToolCallOutput { call_id: self.call_id.clone(), output: truncate_text(&self.output, policy), }) } } + +impl ContextDiscoverable for Message { + fn discoverable_history_item( + &self, + _policy: TruncationPolicy, + codex_home: Option<&Path>, + thread_id: Option<&str>, + context_window_tokens: Option, + current_usage_tokens: i64, + ) -> ResponseItem { + let (Some(codex_home), Some(thread_id), Some(context_window_tokens)) = + (codex_home, thread_id, context_window_tokens) + else { + return ResponseItem::Message(self.clone()); + }; + + if self.role != "user" { + return ResponseItem::Message(self.clone()); + } + + let Some(user_text) = extract_message_text(&self.content) else { + return ResponseItem::Message(self.clone()); + }; + + let estimated_tokens = i64::try_from(approx_token_count(&user_text)).unwrap_or(i64::MAX); + let offload_threshold = context_window_tokens.saturating_mul(95) / 100; + if current_usage_tokens.saturating_add(estimated_tokens) <= offload_threshold { + return ResponseItem::Message(self.clone()); + } + + let checksum = checksum_hex(&user_text); + let output_path = codex_home + .join("discovarable_items") + .join(thread_id) + .join("user_message") + .join(checksum); + + if let Some(parent) = output_path.parent() + && let Err(err) = std::fs::create_dir_all(parent) + { + tracing::warn!( + path = %parent.display(), + "failed to create discoverable message directory: {err}" + ); + return ResponseItem::Message(self.clone()); + } + + if let Err(err) = std::fs::write(&output_path, &user_text) { + tracing::warn!( + path = %output_path.display(), + "failed to write discoverable user message file: {err}" + ); + return ResponseItem::Message(self.clone()); + } + + let replacement = format!( + "User message was too large. Read it from <{}>", + output_path.display() + ); + let mut rewritten = self.clone(); + rewritten.content = vec![ContentItem::InputText { text: replacement }]; + ResponseItem::Message(rewritten) + } +} + +fn extract_message_text(content: &[ContentItem]) -> Option { + let parts: Vec<&str> = content + .iter() + .filter_map(|item| match item { + ContentItem::InputText { text } | ContentItem::OutputText { text } + if !text.trim().is_empty() => + { + Some(text.as_str()) + } + ContentItem::InputText { .. } | ContentItem::OutputText { .. } => None, + ContentItem::InputImage { .. } => None, + }) + .collect(); + if parts.is_empty() { + None + } else { + Some(parts.join("\n\n")) + } +} + +fn checksum_hex(content: &str) -> String { + let mut hasher = Sha256::new(); + hasher.update(content.as_bytes()); + format!("{:x}", hasher.finalize()) +} diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index 0b637f1fd5..93ad67e87a 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -15,20 +15,26 @@ use codex_protocol::models::ResponseItem; use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TokenUsageInfo; use std::ops::Deref; +use std::path::PathBuf; +use std::sync::Arc; /// Transcript of thread history -#[derive(Debug, Clone, Default)] +#[derive(Debug, Clone)] pub(crate) struct ContextManager { /// The oldest items are at the beginning of the vector. items: Vec, token_info: Option, + thread_id: Arc, + codex_home: Arc, } impl ContextManager { - pub(crate) fn new() -> Self { + pub(crate) fn new(thread_id: Arc, codex_home: Arc) -> Self { Self { items: Vec::new(), token_info: TokenUsageInfo::new_or_append(&None, &None, None), + thread_id, + codex_home, } } @@ -288,10 +294,18 @@ impl ContextManager { } fn process_item(&self, item: &ResponseItem, policy: TruncationPolicy) -> ResponseItem { + let used_tokens = self + .token_info + .as_ref() + .map(|info| info.last_token_usage.total_tokens) + .unwrap_or(0); + let context_window = self + .token_info + .as_ref() + .and_then(|info| info.model_context_window); let policy_with_serialization_budget = policy * 1.2; match item { - ResponseItem::Message(_) - | ResponseItem::Reasoning(_) + ResponseItem::Reasoning(_) | ResponseItem::LocalShellCall(_) | ResponseItem::FunctionCall(_) | ResponseItem::CustomToolCall(_) @@ -299,12 +313,27 @@ impl ContextManager { | ResponseItem::Compaction(_) | ResponseItem::GhostSnapshot(_) | ResponseItem::Other => item.clone(), - ResponseItem::FunctionCallOutput(item) => { - item.discoverable_history_item(policy_with_serialization_budget) - } - ResponseItem::CustomToolCallOutput(item) => { - item.discoverable_history_item(policy_with_serialization_budget) - } + ResponseItem::FunctionCallOutput(item) => item.discoverable_history_item( + policy_with_serialization_budget, + Some(self.codex_home.as_path()), + Some(self.thread_id.as_str()), + context_window, + used_tokens, + ), + ResponseItem::CustomToolCallOutput(item) => item.discoverable_history_item( + policy_with_serialization_budget, + Some(self.codex_home.as_path()), + Some(self.thread_id.as_str()), + context_window, + used_tokens, + ), + ResponseItem::Message(item) => item.discoverable_history_item( + policy_with_serialization_budget, + Some(self.codex_home.as_path()), + Some(self.thread_id.as_str()), + context_window, + used_tokens, + ), } } } diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index b3ca57cf97..64174cd56c 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -12,8 +12,13 @@ use codex_protocol::models::LocalShellExecAction; use codex_protocol::models::LocalShellStatus; use codex_protocol::models::ReasoningItemContent; use codex_protocol::models::ReasoningItemReasoningSummary; +use codex_protocol::protocol::TokenUsage; +use codex_protocol::protocol::TokenUsageInfo; use pretty_assertions::assert_eq; use regex_lite::Regex; +use std::path::PathBuf; +use std::sync::Arc; +use tempfile::tempdir; const EXEC_FORMAT_MAX_BYTES: usize = 10_000; const EXEC_FORMAT_MAX_TOKENS: usize = 2_500; @@ -31,7 +36,10 @@ fn assistant_msg(text: &str) -> ResponseItem { } fn create_history_with_items(items: Vec) -> ContextManager { - let mut h = ContextManager::new(); + let mut h = ContextManager::new( + Arc::new("test-thread".to_string()), + Arc::new(PathBuf::from("/tmp/test-codex-home")), + ); // Use a generous but fixed token budget; tests only rely on truncation // behavior, not on a specific model's token limit. h.record_items(items.iter(), TruncationPolicy::Tokens(10_000)); @@ -62,6 +70,24 @@ fn user_input_text_msg(text: &str) -> ResponseItem { }) } +fn set_context_window_and_usage( + history: &mut ContextManager, + context_window: i64, + total_usage: i64, +) { + history.set_token_info(Some(TokenUsageInfo { + total_token_usage: TokenUsage { + total_tokens: total_usage, + ..TokenUsage::default() + }, + last_token_usage: TokenUsage { + total_tokens: total_usage, + ..TokenUsage::default() + }, + model_context_window: Some(context_window), + })); +} + fn function_call_output(call_id: &str, content: &str) -> ResponseItem { ResponseItem::FunctionCallOutput(codex_protocol::models::FunctionCallOutput { call_id: call_id.to_string(), @@ -110,7 +136,10 @@ fn approx_token_count_for_text(text: &str) -> i64 { #[test] fn filters_non_api_messages() { - let mut h = ContextManager::default(); + let mut h = ContextManager::new( + Arc::new("test-thread".to_string()), + Arc::new(PathBuf::from("/tmp/test-codex-home")), + ); let policy = TruncationPolicy::Tokens(10_000); // System message is not API messages; Other is ignored. let system = ResponseItem::Message(codex_protocol::models::Message { @@ -590,7 +619,10 @@ fn normalization_retains_local_shell_outputs() { #[test] fn record_items_truncates_function_call_output_content() { - let mut history = ContextManager::new(); + let mut history = ContextManager::new( + Arc::new("test-thread".to_string()), + Arc::new(PathBuf::from("/tmp/test-codex-home")), + ); // Any reasonably small token budget works; the test only cares that // truncation happens and the marker is present. let policy = TruncationPolicy::Tokens(1_000); @@ -629,7 +661,10 @@ fn record_items_truncates_function_call_output_content() { #[test] fn record_items_truncates_custom_tool_call_output_content() { - let mut history = ContextManager::new(); + let mut history = ContextManager::new( + Arc::new("test-thread".to_string()), + Arc::new(PathBuf::from("/tmp/test-codex-home")), + ); let policy = TruncationPolicy::Tokens(1_000); let line = "custom output that is very long\n"; let long_output = line.repeat(2_500); @@ -662,7 +697,10 @@ fn record_items_truncates_custom_tool_call_output_content() { #[test] fn record_items_respects_custom_token_limit() { - let mut history = ContextManager::new(); + let mut history = ContextManager::new( + Arc::new("test-thread".to_string()), + Arc::new(PathBuf::from("/tmp/test-codex-home")), + ); let policy = TruncationPolicy::Tokens(10); let long_output = "tokenized content repeated many times ".repeat(200); let item = ResponseItem::FunctionCallOutput(codex_protocol::models::FunctionCallOutput { @@ -689,6 +727,119 @@ fn record_items_respects_custom_token_limit() { ); } +#[test] +fn record_items_with_discoverability_offloads_large_user_message() { + let codex_home = tempdir().expect("create tempdir"); + let mut history = ContextManager::new( + Arc::new("thread-1".to_string()), + Arc::new(codex_home.path().to_path_buf()), + ); + let policy = TruncationPolicy::Tokens(10_000); + let large_text = "large user prompt ".repeat(200); + let item = user_input_text_msg(&large_text); + set_context_window_and_usage(&mut history, 1, 1); + + history.record_items([&item], policy); + + assert_eq!(history.raw_items().len(), 1); + let stored = &history.raw_items()[0]; + let ResponseItem::Message(message) = stored else { + panic!("expected message item, got {stored:?}"); + }; + let [ContentItem::InputText { text }] = message.content.as_slice() else { + panic!("expected a single discoverable pointer text, got {message:?}"); + }; + assert!(text.starts_with("User message was too large. Read it from <")); + let path = text + .split('<') + .nth(1) + .and_then(|s| s.split('>').next()) + .expect("discoverable path enclosed in angle brackets"); + assert!(path.contains("/discovarable_items/thread-1/user_message/")); + let persisted = std::fs::read_to_string(path).expect("discoverable file exists"); + assert_eq!(persisted, large_text); +} + +#[test] +fn record_items_with_discoverability_keeps_small_user_message_in_history() { + let codex_home = tempdir().expect("create tempdir"); + let mut history = ContextManager::new( + Arc::new("thread-2".to_string()), + Arc::new(codex_home.path().to_path_buf()), + ); + let policy = TruncationPolicy::Tokens(10_000); + let item = user_input_text_msg("small prompt"); + set_context_window_and_usage(&mut history, 10_000, 0); + + history.record_items([&item], policy); + + assert_eq!(history.raw_items(), vec![item]); +} + +#[test] +fn record_items_with_discoverability_does_not_change_non_user_messages() { + let codex_home = tempdir().expect("create tempdir"); + let mut history = ContextManager::new( + Arc::new("thread-3".to_string()), + Arc::new(codex_home.path().to_path_buf()), + ); + let policy = TruncationPolicy::Tokens(10_000); + let item = assistant_msg(&"assistant reply ".repeat(300)); + set_context_window_and_usage(&mut history, 1, 1); + + history.record_items([&item], policy); + + assert_eq!(history.raw_items(), vec![item]); +} + +#[test] +fn record_items_with_discoverability_ignores_image_only_user_messages() { + let codex_home = tempdir().expect("create tempdir"); + let mut history = ContextManager::new( + Arc::new("thread-4".to_string()), + Arc::new(codex_home.path().to_path_buf()), + ); + let policy = TruncationPolicy::Tokens(10_000); + let item = ResponseItem::Message(codex_protocol::models::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputImage { + image_url: "data:image/png;base64,AAA".to_string(), + }], + end_turn: None, + phase: None, + }); + set_context_window_and_usage(&mut history, 1, 1); + + history.record_items([&item], policy); + + assert_eq!(history.raw_items(), vec![item]); +} + +#[test] +fn record_items_with_discoverability_offloads_once_above_ninety_five_percent_window() { + let codex_home = tempdir().expect("create tempdir"); + let mut history = ContextManager::new( + Arc::new("thread-95".to_string()), + Arc::new(codex_home.path().to_path_buf()), + ); + let policy = TruncationPolicy::Tokens(10_000); + let text = "near context threshold ".repeat(100); + let item = user_input_text_msg(&text); + set_context_window_and_usage(&mut history, 1_000, 900); + + history.record_items([&item], policy); + + let stored = &history.raw_items()[0]; + let ResponseItem::Message(message) = stored else { + panic!("expected message item, got {stored:?}"); + }; + let [ContentItem::InputText { text }] = message.content.as_slice() else { + panic!("expected discoverable pointer text, got {message:?}"); + }; + assert!(text.starts_with("User message was too large. Read it from <")); +} + fn assert_truncated_message_matches(message: &str, line: &str, expected_removed: usize) { let pattern = truncated_message_pattern(line); let regex = Regex::new(&pattern).unwrap_or_else(|err| { diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 29bc14d1cc..d6c3ee4eaf 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -3,6 +3,8 @@ use codex_protocol::models::ResponseItem; use std::collections::HashMap; use std::collections::HashSet; +use std::path::PathBuf; +use std::sync::Arc; use crate::codex::SessionConfiguration; use crate::context_manager::ContextManager; @@ -30,8 +32,12 @@ pub(crate) struct SessionState { impl SessionState { /// Create a new session state mirroring previous `State::default()` semantics. - pub(crate) fn new(session_configuration: SessionConfiguration) -> Self { - let history = ContextManager::new(); + pub(crate) fn new( + session_configuration: SessionConfiguration, + thread_id: String, + codex_home: PathBuf, + ) -> Self { + let history = ContextManager::new(Arc::new(thread_id), Arc::new(codex_home)); Self { session_configuration, history, diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index 379f521682..2dbbda2c77 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -70,6 +70,7 @@ mod compact; mod compact_remote; mod compact_resume_fork; mod deprecation_notice; +mod discoverable_history; mod exec; mod exec_policy; mod fork_thread;