openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-15 12:08:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix

Browse Source
This commit is contained in:
Ahmed Ibrahim
2026-02-08 21:37:07 -08:00
parent 0fd69f1b37
commit 316637aa34
6 changed files with 359 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
codex-rs/core/src/codex.rs
View File

@@ -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<ResponseItem> {
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<RolloutItem>, Vec<ResponseItem>) {
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 {

123
codex-rs/core/src/context_manager/context_discoverable.rs
View File

@@ -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<i64>,
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<i64>,
_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<i64>,
_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<i64>,
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<String> {
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())
}

49
codex-rs/core/src/context_manager/history.rs
View File

@@ -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<ResponseItem>,
token_info: Option<TokenUsageInfo>,
thread_id: Arc<String>,
codex_home: Arc<PathBuf>,
}
impl ContextManager {
pub(crate) fn new() -> Self {
pub(crate) fn new(thread_id: Arc<String>, codex_home: Arc<PathBuf>) -> 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,
),
}
}
}

161
codex-rs/core/src/context_manager/history_tests.rs
View File

@@ -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<ResponseItem>) -> 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| {

10
codex-rs/core/src/state/session.rs
View File

@@ -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,

1
codex-rs/core/tests/suite/mod.rs
View File

@@ -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;