This commit is contained in:
Ahmed Ibrahim
2026-01-14 09:44:10 -08:00
parent b594def567
commit e90eeee6b2
5 changed files with 55 additions and 32 deletions

View File

@@ -453,6 +453,10 @@ pub(crate) struct SessionConfiguration {
}
impl SessionConfiguration {
pub(crate) fn codex_home(&self) -> &PathBuf {
&self.original_config_do_not_use.codex_home
}
pub(crate) fn apply(&self, updates: &SessionSettingsUpdate) -> ConstraintResult<Self> {
let mut next_configuration = self.clone();
if let Some(model) = updates.model.clone() {
@@ -686,7 +690,7 @@ impl Session {
.await
.map(Arc::new);
}
let state = SessionState::new(session_configuration.clone());
let state = SessionState::new(session_configuration.clone(), conversation_id);
let services = SessionServices {
mcp_connection_manager: Arc::new(RwLock::new(McpConnectionManager::default())),
@@ -3347,7 +3351,8 @@ mod tests {
session_source: SessionSource::Exec,
};
let mut state = SessionState::new(session_configuration);
let conversation_id = ThreadId::default();
let mut state = SessionState::new(session_configuration, conversation_id);
let initial = RateLimitSnapshot {
primary: Some(RateLimitWindow {
used_percent: 10.0,
@@ -3413,7 +3418,8 @@ mod tests {
session_source: SessionSource::Exec,
};
let mut state = SessionState::new(session_configuration);
let conversation_id = ThreadId::default();
let mut state = SessionState::new(session_configuration, conversation_id);
let initial = RateLimitSnapshot {
primary: Some(RateLimitWindow {
used_percent: 15.0,
@@ -3674,7 +3680,7 @@ mod tests {
session_configuration.session_source.clone(),
);
let state = SessionState::new(session_configuration.clone());
let state = SessionState::new(session_configuration.clone(), conversation_id);
let skills_manager = Arc::new(SkillsManager::new(config.codex_home.clone()));
let services = SessionServices {
@@ -3769,7 +3775,7 @@ mod tests {
session_configuration.session_source.clone(),
);
let state = SessionState::new(session_configuration.clone());
let state = SessionState::new(session_configuration.clone(), conversation_id);
let skills_manager = Arc::new(SkillsManager::new(config.codex_home.clone()));
let services = SessionServices {

View File

@@ -15,8 +15,10 @@ use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::TokenUsage;
use codex_protocol::protocol::TokenUsageInfo;
use std::ops::Deref;
use std::path::Path;
use std::path::PathBuf;
use tracing::warn;
use uuid::Uuid;
/// Transcript of thread history
#[derive(Debug, Clone, Default)]
@@ -24,6 +26,7 @@ pub(crate) struct ContextManager {
/// The oldest items are at the beginning of the vector.
items: Vec<ResponseItem>,
token_info: Option<TokenUsageInfo>,
user_message_dir: Option<PathBuf>,
}
impl ContextManager {
@@ -31,6 +34,7 @@ impl ContextManager {
Self {
items: Vec::new(),
token_info: TokenUsageInfo::new_or_append(&None, &None, None),
user_message_dir: None,
}
}
@@ -42,6 +46,10 @@ impl ContextManager {
self.token_info = info;
}
pub(crate) fn set_user_message_dir(&mut self, dir: Option<PathBuf>) {
self.user_message_dir = dir;
}
pub(crate) fn set_token_usage_full(&mut self, context_window: i64) {
match &mut self.token_info {
Some(info) => info.fill_to_context_window(context_window),
@@ -318,7 +326,10 @@ impl ContextManager {
return item.clone();
}
let Some(path) = write_message_to_temp_file(&text) else {
let Some(user_message_dir) = self.user_message_dir.as_ref() else {
return item.clone();
};
let Some(path) = write_message_to_user_dir(&text, user_message_dir) else {
return item.clone();
};
@@ -368,31 +379,18 @@ fn message_text(content: &[ContentItem]) -> String {
pieces.join("\n")
}
fn write_message_to_temp_file(text: &str) -> Option<PathBuf> {
let mut temp = match tempfile::Builder::new()
.prefix("codex-user-input-")
.suffix(".txt")
.tempfile()
{
Ok(temp) => temp,
Err(err) => {
warn!(error = %err, "failed to create temp file for user message");
return None;
}
};
if let Err(err) = std::io::Write::write_all(&mut temp, text.as_bytes()) {
warn!(error = %err, "failed to write user message to temp file");
fn write_message_to_user_dir(text: &str, user_message_dir: &Path) -> Option<PathBuf> {
if let Err(err) = std::fs::create_dir_all(user_message_dir) {
warn!(error = %err, "failed to create user message directory");
return None;
}
let (_, path) = match temp.keep() {
Ok(kept) => kept,
Err(err) => {
warn!(error = %err, "failed to persist user message temp file");
return None;
}
};
let id = Uuid::new_v4();
let path = user_message_dir.join(format!("user-message-{id}.txt"));
if let Err(err) = std::fs::write(&path, text.as_bytes()) {
warn!(error = %err, "failed to write user message to file");
return None;
}
Some(path)
}

View File

@@ -15,6 +15,7 @@ use pretty_assertions::assert_eq;
use regex_lite::Regex;
use std::fs;
use std::path::PathBuf;
use tempfile::tempdir;
const EXEC_FORMAT_MAX_BYTES: usize = 10_000;
const EXEC_FORMAT_MAX_TOKENS: usize = 2_500;
@@ -86,8 +87,16 @@ fn truncate_exec_output(content: &str) -> String {
}
#[test]
fn stores_oversized_user_message_in_temp_file() {
fn stores_oversized_user_message_in_user_message_dir() {
let mut history = ContextManager::new();
let temp_dir = tempdir().expect("create temp dir");
history.set_user_message_dir(Some(
temp_dir
.path()
.join("context")
.join("usermsgs")
.join("test-session"),
));
history.set_token_info(Some(TokenUsageInfo {
total_token_usage: TokenUsage::default(),
last_token_usage: TokenUsage::default(),

View File

@@ -1,5 +1,6 @@
//! Session-wide mutable state.
use codex_protocol::ThreadId;
use codex_protocol::models::ResponseItem;
use crate::codex::SessionConfiguration;
@@ -18,8 +19,17 @@ 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,
conversation_id: ThreadId,
) -> Self {
let mut history = ContextManager::new();
let user_message_dir = session_configuration
.codex_home()
.join("context")
.join("usermsgs")
.join(conversation_id.to_string());
history.set_user_message_dir(Some(user_message_dir));
Self {
session_configuration,
history,

View File

@@ -244,7 +244,7 @@ async fn tool_call_output_exceeds_limit_truncated_chars_limit() -> Result<()> {
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn oversized_user_message_is_written_to_temp_file() -> Result<()> {
async fn oversized_user_message_is_written_to_user_message_dir() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;