From f00edcd5135d233c3b86770589683beb8ff3881b Mon Sep 17 00:00:00 2001 From: jif-oai Date: Mon, 1 Dec 2025 16:27:26 +0000 Subject: [PATCH 1/8] feat: add warning message for the model --- codex-rs/core/src/codex.rs | 39 +++++++++++++++++++ codex-rs/core/src/truncate.rs | 4 ++ codex-rs/core/src/unified_exec/mod.rs | 3 ++ .../core/src/unified_exec/session_manager.rs | 26 ++++++++++--- 4 files changed, 67 insertions(+), 5 deletions(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index a3d21ba8ca..f2da21d323 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -1035,6 +1035,19 @@ impl Session { state.record_items(items.iter(), turn_context.truncation_policy); } + pub(crate) async fn record_model_warning(&self, message: impl Into) { + let item = ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: format!("Warning: {}", message.into()), + }], + }; + + let mut state = self.state.lock().await; + state.record_items([item].iter(), TruncationPolicy::no_truncation()); + } + pub(crate) async fn replace_history(&self, items: Vec) { let mut state = self.state.lock().await; state.replace_history(items); @@ -2787,6 +2800,32 @@ mod tests { (session, turn_context, rx_event) } + #[tokio::test] + async fn record_model_warning_appends_user_message() { + let (session, _turn_context) = make_session_and_context(); + + session + .record_model_warning("too many unified exec sessions") + .await; + + let mut history = session.clone_history().await; + let history_items = history.get_history(); + let last = history_items.last().expect("warning recorded"); + + match last { + ResponseItem::Message { role, content, .. } => { + assert_eq!(role, "user"); + assert_eq!( + content, + &vec![ContentItem::InputText { + text: "Warning: too many unified exec sessions".to_string(), + }] + ); + } + other => panic!("expected user message, got {other:?}"), + } + } + #[derive(Clone, Copy)] struct NeverEndingTask { kind: TaskKind, diff --git a/codex-rs/core/src/truncate.rs b/codex-rs/core/src/truncate.rs index 6e38ef6986..14bf81f18f 100644 --- a/codex-rs/core/src/truncate.rs +++ b/codex-rs/core/src/truncate.rs @@ -72,6 +72,10 @@ impl TruncationPolicy { TruncationPolicy::Tokens(tokens) => approx_bytes_for_tokens(*tokens), } } + + pub fn no_truncation() -> Self { + TruncationPolicy::Tokens(999_999) + } } pub(crate) fn formatted_truncate_text(content: &str, policy: TruncationPolicy) -> String { diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs index 8c33548065..34b62df342 100644 --- a/codex-rs/core/src/unified_exec/mod.rs +++ b/codex-rs/core/src/unified_exec/mod.rs @@ -48,6 +48,9 @@ pub(crate) const UNIFIED_EXEC_OUTPUT_MAX_BYTES: usize = 1024 * 1024; // 1 MiB pub(crate) const UNIFIED_EXEC_OUTPUT_MAX_TOKENS: usize = UNIFIED_EXEC_OUTPUT_MAX_BYTES / 4; pub(crate) const MAX_UNIFIED_EXEC_SESSIONS: usize = 64; +// Send a warning message to the models when it reaches this number of sessions. +pub(crate) const WARNING_UNIFIED_EXEC_SESSIONS: usize = 60; + pub(crate) struct UnifiedExecContext { pub session: Arc, pub turn: Arc, diff --git a/codex-rs/core/src/unified_exec/session_manager.rs b/codex-rs/core/src/unified_exec/session_manager.rs index 23044c0f78..373e89da6e 100644 --- a/codex-rs/core/src/unified_exec/session_manager.rs +++ b/codex-rs/core/src/unified_exec/session_manager.rs @@ -41,6 +41,7 @@ use super::UnifiedExecContext; use super::UnifiedExecError; use super::UnifiedExecResponse; use super::UnifiedExecSessionManager; +use super::WARNING_UNIFIED_EXEC_SESSIONS; use super::WriteStdinRequest; use super::clamp_yield_time; use super::generate_chunk_id; @@ -421,9 +422,21 @@ impl UnifiedExecSessionManager { started_at, last_used: started_at, }; - let mut store = self.session_store.lock().await; - Self::prune_sessions_if_needed(&mut store); - store.sessions.insert(process_id, entry); + let number_sessions = { + let mut store = self.session_store.lock().await; + Self::prune_sessions_if_needed(&mut store); + store.sessions.insert(process_id, entry); + store.sessions.len() + }; + + if number_sessions >= WARNING_UNIFIED_EXEC_SESSIONS { + context + .session + .record_model_warning( + format!("The maximum number of unified exec sessions you can keep open is {WARNING_UNIFIED_EXEC_SESSIONS} and you currently have {number_sessions} sessions open. Re-use older sessions or close them to prevent automatic pruning of old session"), + ) + .await; + }; } async fn emit_exec_end_from_entry( @@ -633,9 +646,9 @@ impl UnifiedExecSessionManager { collected } - fn prune_sessions_if_needed(store: &mut SessionStore) { + fn prune_sessions_if_needed(store: &mut SessionStore) -> bool { if store.sessions.len() < MAX_UNIFIED_EXEC_SESSIONS { - return; + return false; } let meta: Vec<(String, Instant, bool)> = store @@ -646,7 +659,10 @@ impl UnifiedExecSessionManager { if let Some(session_id) = Self::session_id_to_prune_from_meta(&meta) { store.remove(&session_id); + return true; } + + false } // Centralized pruning policy so we can easily swap strategies later. From aa50471d50e67866eeee7fe0921672ed1fc3313e Mon Sep 17 00:00:00 2001 From: jif-oai Date: Mon, 1 Dec 2025 16:31:05 +0000 Subject: [PATCH 2/8] Update session_manager.rs --- codex-rs/core/src/unified_exec/session_manager.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codex-rs/core/src/unified_exec/session_manager.rs b/codex-rs/core/src/unified_exec/session_manager.rs index 373e89da6e..75832bfacb 100644 --- a/codex-rs/core/src/unified_exec/session_manager.rs +++ b/codex-rs/core/src/unified_exec/session_manager.rs @@ -433,7 +433,7 @@ impl UnifiedExecSessionManager { context .session .record_model_warning( - format!("The maximum number of unified exec sessions you can keep open is {WARNING_UNIFIED_EXEC_SESSIONS} and you currently have {number_sessions} sessions open. Re-use older sessions or close them to prevent automatic pruning of old session"), + format!("The maximum number of unified exec sessions you can keep open is {WARNING_UNIFIED_EXEC_SESSIONS} and you currently have {number_sessions} sessions open. Reuse older sessions or close them to prevent automatic pruning of old session"), ) .await; }; From 1762db5fa437f93b9fa3256c4fd9aa15afb8dd3c Mon Sep 17 00:00:00 2001 From: jif-oai Date: Mon, 1 Dec 2025 16:47:02 +0000 Subject: [PATCH 3/8] Add feature flag --- codex-rs/core/src/codex.rs | 4 ++++ codex-rs/core/src/features.rs | 16 ++++++++++++---- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index f2da21d323..a43708ebf8 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -1036,6 +1036,10 @@ impl Session { } pub(crate) async fn record_model_warning(&self, message: impl Into) { + if !self.enabled(Feature::ModelWarnings) { + return; + } + let item = ResponseItem::Message { id: None, role: "user".to_string(), diff --git a/codex-rs/core/src/features.rs b/codex-rs/core/src/features.rs index 0c67c9ff5c..88319b052f 100644 --- a/codex-rs/core/src/features.rs +++ b/codex-rs/core/src/features.rs @@ -51,6 +51,8 @@ pub enum Feature { ShellTool, /// Allow model to call multiple tools in parallel (only for models supporting it). ParallelToolCalls, + /// Send warnings to the model to correct it on the tool usage. + ModelWarnings, } impl Feature { @@ -265,6 +267,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::Stable, default_enabled: true, }, + FeatureSpec { + id: Feature::ShellTool, + key: "shell_tool", + stage: Stage::Stable, + default_enabled: true, + }, // Unstable features. FeatureSpec { id: Feature::UnifiedExec, @@ -321,9 +329,9 @@ pub const FEATURES: &[FeatureSpec] = &[ default_enabled: false, }, FeatureSpec { - id: Feature::ShellTool, - key: "shell_tool", - stage: Stage::Stable, - default_enabled: true, + id: Feature::ModelWarnings, + key: "warnings", + stage: Stage::Experimental, + default_enabled: false, }, ]; From ad35ef128519789d2c02c6ef6de5529fae2cbca1 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Mon, 1 Dec 2025 16:52:40 +0000 Subject: [PATCH 4/8] Update codex.rs --- codex-rs/core/src/codex.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index a43708ebf8..bc4cf37551 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -1036,7 +1036,7 @@ impl Session { } pub(crate) async fn record_model_warning(&self, message: impl Into) { - if !self.enabled(Feature::ModelWarnings) { + if !self.enabled(Feature::ModelWarnings).await { return; } From 2cd209e177676317125a5008423a0b3bb4cf4f10 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Mon, 1 Dec 2025 17:07:08 +0000 Subject: [PATCH 5/8] Fix test --- codex-rs/core/src/codex.rs | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index bc4cf37551..6f7d857d19 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -2808,6 +2808,14 @@ mod tests { async fn record_model_warning_appends_user_message() { let (session, _turn_context) = make_session_and_context(); + session + .state + .lock() + .await + .session_configuration + .features + .enable(Feature::ModelWarnings); + session .record_model_warning("too many unified exec sessions") .await; From 9ca7247f74599dac993875140e491fbdadc296e4 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Tue, 2 Dec 2025 11:08:35 +0000 Subject: [PATCH 6/8] Persist to rollout --- codex-rs/core/src/codex.rs | 9 ++++----- codex-rs/core/src/unified_exec/session_manager.rs | 1 + 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 6f7d857d19..6950f1fb89 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -1035,7 +1035,7 @@ impl Session { state.record_items(items.iter(), turn_context.truncation_policy); } - pub(crate) async fn record_model_warning(&self, message: impl Into) { + pub(crate) async fn record_model_warning(&self, message: impl Into, ctx: &TurnContext) { if !self.enabled(Feature::ModelWarnings).await { return; } @@ -1048,8 +1048,7 @@ impl Session { }], }; - let mut state = self.state.lock().await; - state.record_items([item].iter(), TruncationPolicy::no_truncation()); + self.record_conversation_items(ctx, &[item]).await; } pub(crate) async fn replace_history(&self, items: Vec) { @@ -2806,7 +2805,7 @@ mod tests { #[tokio::test] async fn record_model_warning_appends_user_message() { - let (session, _turn_context) = make_session_and_context(); + let (session, turn_context) = make_session_and_context(); session .state @@ -2817,7 +2816,7 @@ mod tests { .enable(Feature::ModelWarnings); session - .record_model_warning("too many unified exec sessions") + .record_model_warning("too many unified exec sessions", turn_context) .await; let mut history = session.clone_history().await; diff --git a/codex-rs/core/src/unified_exec/session_manager.rs b/codex-rs/core/src/unified_exec/session_manager.rs index 75832bfacb..37a12bf2df 100644 --- a/codex-rs/core/src/unified_exec/session_manager.rs +++ b/codex-rs/core/src/unified_exec/session_manager.rs @@ -434,6 +434,7 @@ impl UnifiedExecSessionManager { .session .record_model_warning( format!("The maximum number of unified exec sessions you can keep open is {WARNING_UNIFIED_EXEC_SESSIONS} and you currently have {number_sessions} sessions open. Reuse older sessions or close them to prevent automatic pruning of old session"), + &context.turn ) .await; }; From 01c4c1e273c5c77845d38d1142927daedb00e0b6 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Tue, 2 Dec 2025 11:09:06 +0000 Subject: [PATCH 7/8] NIT --- codex-rs/core/src/codex.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 6950f1fb89..92b325a6a4 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -2816,7 +2816,7 @@ mod tests { .enable(Feature::ModelWarnings); session - .record_model_warning("too many unified exec sessions", turn_context) + .record_model_warning("too many unified exec sessions", &turn_context) .await; let mut history = session.clone_history().await; From 5d29df5818ace31ff10cb4447d2ed376b3d66ee4 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Tue, 2 Dec 2025 11:41:57 +0000 Subject: [PATCH 8/8] nit --- codex-rs/core/src/truncate.rs | 4 ---- 1 file changed, 4 deletions(-) diff --git a/codex-rs/core/src/truncate.rs b/codex-rs/core/src/truncate.rs index 14bf81f18f..6e38ef6986 100644 --- a/codex-rs/core/src/truncate.rs +++ b/codex-rs/core/src/truncate.rs @@ -72,10 +72,6 @@ impl TruncationPolicy { TruncationPolicy::Tokens(tokens) => approx_bytes_for_tokens(*tokens), } } - - pub fn no_truncation() -> Self { - TruncationPolicy::Tokens(999_999) - } } pub(crate) fn formatted_truncate_text(content: &str, policy: TruncationPolicy) -> String {