Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 50 additions & 0 deletions codex-rs/core/src/codex.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<String>, 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<ResponseItem>) {
let mut state = self.state.lock().await;
state.replace_history(items);
Expand Down Expand Up @@ -2787,6 +2803,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,
Expand Down
16 changes: 12 additions & 4 deletions codex-rs/core/src/features.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
3 changes: 3 additions & 0 deletions codex-rs/core/src/unified_exec/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Shouldn't we use MAX_UNIFIED_EXEC_SESSIONS? or something relative to it?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We want something slightly lower on purpose. To avoid having to prune

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's use it relative to it at least so we have only one magic number?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Then we need a relative factor so same thing. I kind of like the idea of soft and hard limits

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I need to merge to move on but happy to discuss this mater further and I can do a follow-up PR if you don't agree


pub(crate) struct UnifiedExecContext {
pub session: Arc<Session>,
pub turn: Arc<TurnContext>,
Expand Down
27 changes: 22 additions & 5 deletions codex-rs/core/src/unified_exec/session_manager.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand Down
Loading