From 9ee855ec57e3fc2da6c0f640a21ad04d7c965de7 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Tue, 2 Dec 2025 11:56:00 +0000 Subject: [PATCH] feat: add warning message for the model (#7445) Add a warning message as a user turn to the model if the model does not behave as expected (here, for example, if the model opens too many `unified_exec` sessions) --- codex-rs/core/src/codex.rs | 50 +++++++++++++++++++ codex-rs/core/src/features.rs | 16 ++++-- codex-rs/core/src/unified_exec/mod.rs | 3 ++ .../core/src/unified_exec/session_manager.rs | 27 ++++++++-- 4 files changed, 87 insertions(+), 9 deletions(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 4e3949fa0..8ae83d003 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -1035,6 +1035,22 @@ impl Session { state.record_items(items.iter(), turn_context.truncation_policy); } + pub(crate) async fn record_model_warning(&self, message: impl Into, ctx: &TurnContext) { + if !self.enabled(Feature::ModelWarnings).await { + return; + } + + let item = ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: format!("Warning: {}", message.into()), + }], + }; + + self.record_conversation_items(ctx, &[item]).await; + } + pub(crate) async fn replace_history(&self, items: Vec) { let mut state = self.state.lock().await; state.replace_history(items); @@ -2806,6 +2822,40 @@ 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 + .state + .lock() + .await + .session_configuration + .features + .enable(Feature::ModelWarnings); + + session + .record_model_warning("too many unified exec sessions", &turn_context) + .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/features.rs b/codex-rs/core/src/features.rs index 14f0404b3..687e8b162 100644 --- a/codex-rs/core/src/features.rs +++ b/codex-rs/core/src/features.rs @@ -53,6 +53,8 @@ pub enum Feature { ParallelToolCalls, /// Experimental skills injection (CLI flag-driven). Skills, + /// Send warnings to the model to correct it on the tool usage. + ModelWarnings, } impl Feature { @@ -267,6 +269,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, @@ -323,10 +331,10 @@ 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, }, FeatureSpec { id: Feature::Skills, diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs index 8c3354806..34b62df34 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 23044c0f7..37a12bf2d 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,22 @@ 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. Reuse older sessions or close them to prevent automatic pruning of old session"), + &context.turn + ) + .await; + }; } async fn emit_exec_end_from_entry( @@ -633,9 +647,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 +660,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.