From 7cee12d5c0ffa4bb0b456b5fccbe52ee888c2f5d Mon Sep 17 00:00:00 2001 From: jif-oai Date: Fri, 28 Nov 2025 11:54:57 +0000 Subject: [PATCH 1/5] chore: better session recycling --- codex-rs/core/src/unified_exec/mod.rs | 2 -- codex-rs/core/src/unified_exec/session_manager.rs | 7 +++---- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs index 35df80255b..81efd7bd6c 100644 --- a/codex-rs/core/src/unified_exec/mod.rs +++ b/codex-rs/core/src/unified_exec/mod.rs @@ -22,7 +22,6 @@ //! - `session_manager.rs`: orchestration (approvals, sandboxing, reuse) and request handling. use std::collections::HashMap; -use std::collections::HashSet; use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; @@ -99,7 +98,6 @@ pub(crate) struct UnifiedExecResponse { #[derive(Default)] pub(crate) struct UnifiedExecSessionManager { sessions: Mutex>, - used_session_ids: Mutex>, } struct SessionEntry { diff --git a/codex-rs/core/src/unified_exec/session_manager.rs b/codex-rs/core/src/unified_exec/session_manager.rs index df8f543212..f4e76144e1 100644 --- a/codex-rs/core/src/unified_exec/session_manager.rs +++ b/codex-rs/core/src/unified_exec/session_manager.rs @@ -81,7 +81,7 @@ struct PreparedSessionHandles { impl UnifiedExecSessionManager { pub(crate) async fn allocate_process_id(&self) -> String { loop { - let mut store = self.used_session_ids.lock().await; + let store = self.sessions.lock().await; let process_id = if !cfg!(test) && !cfg!(feature = "deterministic_process_ids") { // production mode → random @@ -90,7 +90,7 @@ impl UnifiedExecSessionManager { // test or deterministic mode let next = store .iter() - .filter_map(|s| s.parse::().ok()) + .filter_map(|(s, _)| s.parse::().ok()) .max() .map(|m| std::cmp::max(m, 999) + 1) .unwrap_or(1000); @@ -98,11 +98,10 @@ impl UnifiedExecSessionManager { next.to_string() }; - if store.contains(&process_id) { + if store.contains_key(&process_id) { continue; } - store.insert(process_id.clone()); return process_id; } } From 946022afd9dc91f2495add4a416bfc2eb3398703 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Fri, 28 Nov 2025 13:18:44 +0000 Subject: [PATCH 2/5] Better --- codex-rs/core/src/unified_exec/mod.rs | 29 +++++++++++++-- .../core/src/unified_exec/session_manager.rs | 37 ++++++++++--------- 2 files changed, 44 insertions(+), 22 deletions(-) diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs index 81efd7bd6c..a368bf8f5c 100644 --- a/codex-rs/core/src/unified_exec/mod.rs +++ b/codex-rs/core/src/unified_exec/mod.rs @@ -21,7 +21,7 @@ //! - `session.rs`: PTY session lifecycle + output buffering. //! - `session_manager.rs`: orchestration (approvals, sandboxing, reuse) and request handling. -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; @@ -97,7 +97,26 @@ pub(crate) struct UnifiedExecResponse { #[derive(Default)] pub(crate) struct UnifiedExecSessionManager { - sessions: Mutex>, + session_store: Mutex, +} + +// Required for mutex sharing. +#[derive(Default)] +pub(crate) struct SessionStore { + sessions: HashMap, + reserved_sessions_id: HashSet +} + +impl SessionStore { + pub fn remove(&mut self, session_id: &str) { + self.sessions.remove(&session_id); + self.reserved_sessions_id.remove(&session_id); + } + + pub fn clear(&mut self) { + self.reserved_sessions_id.clear(); + self.sessions.clear(); + } } struct SessionEntry { @@ -382,9 +401,10 @@ mod tests { session .services .unified_exec_manager - .sessions + .session_store .lock() .await + .sessions .is_empty() ); @@ -423,9 +443,10 @@ mod tests { session .services .unified_exec_manager - .sessions + .session_store .lock() .await + .sessions .is_empty() ); diff --git a/codex-rs/core/src/unified_exec/session_manager.rs b/codex-rs/core/src/unified_exec/session_manager.rs index f4e76144e1..1a13fd3ca3 100644 --- a/codex-rs/core/src/unified_exec/session_manager.rs +++ b/codex-rs/core/src/unified_exec/session_manager.rs @@ -33,7 +33,7 @@ use crate::truncate::TruncationPolicy; use crate::truncate::approx_token_count; use crate::truncate::formatted_truncate_text; -use super::ExecCommandRequest; +use super::{ExecCommandRequest, SessionStore}; use super::MAX_UNIFIED_EXEC_SESSIONS; use super::SessionEntry; use super::UnifiedExecContext; @@ -81,16 +81,16 @@ struct PreparedSessionHandles { impl UnifiedExecSessionManager { pub(crate) async fn allocate_process_id(&self) -> String { loop { - let store = self.sessions.lock().await; + let mut store = self.session_store.lock().await; let process_id = if !cfg!(test) && !cfg!(feature = "deterministic_process_ids") { // production mode → random rand::rng().random_range(1_000..100_000).to_string() } else { // test or deterministic mode - let next = store + let next = store.reserved_sessions_id .iter() - .filter_map(|(s, _)| s.parse::().ok()) + .filter_map(|s| s.parse::().ok()) .max() .map(|m| std::cmp::max(m, 999) + 1) .unwrap_or(1000); @@ -98,10 +98,11 @@ impl UnifiedExecSessionManager { next.to_string() }; - if store.contains_key(&process_id) { + if store.reserved_sessions_id.contains(&process_id) { continue; } + store.reserved_sessions_id.insert(process_id.clone()); return process_id; } } @@ -330,8 +331,8 @@ impl UnifiedExecSessionManager { } async fn refresh_session_state(&self, process_id: &str) -> SessionStatus { - let mut sessions = self.sessions.lock().await; - let Some(entry) = sessions.get(process_id) else { + let mut store = self.session_store.lock().await; + let Some(entry) = store.sessions.get(process_id) else { return SessionStatus::Unknown; }; @@ -339,7 +340,7 @@ impl UnifiedExecSessionManager { let process_id = entry.process_id.clone(); if entry.session.has_exited() { - let Some(entry) = sessions.remove(&process_id) else { + let Some(entry) = store.remove(&process_id) else { return SessionStatus::Unknown; }; SessionStatus::Exited { @@ -359,8 +360,8 @@ impl UnifiedExecSessionManager { &self, process_id: &str, ) -> Result { - let mut sessions = self.sessions.lock().await; - let entry = sessions + let mut store = self.session_store.lock().await; + let entry = store.sessions .get_mut(process_id) .ok_or(UnifiedExecError::UnknownSessionId { process_id: process_id.to_string(), @@ -416,9 +417,9 @@ impl UnifiedExecSessionManager { started_at, last_used: started_at, }; - let mut sessions = self.sessions.lock().await; - Self::prune_sessions_if_needed(&mut sessions); - sessions.insert(process_id, entry); + let mut store = self.session_store.lock().await; + Self::prune_sessions_if_needed(&mut store); + store.sessions.insert(process_id, entry); } async fn emit_exec_end_from_entry( @@ -628,18 +629,18 @@ impl UnifiedExecSessionManager { collected } - fn prune_sessions_if_needed(sessions: &mut HashMap) { - if sessions.len() < MAX_UNIFIED_EXEC_SESSIONS { + fn prune_sessions_if_needed(store: &mut SessionStore) { + if store.sessions.len() < MAX_UNIFIED_EXEC_SESSIONS { return; } - let meta: Vec<(String, Instant, bool)> = sessions + let meta: Vec<(String, Instant, bool)> = store.sessions .iter() .map(|(id, entry)| (id.clone(), entry.last_used, entry.session.has_exited())) .collect(); if let Some(session_id) = Self::session_id_to_prune_from_meta(&meta) { - sessions.remove(&session_id); + store.remove(&session_id); } } @@ -673,7 +674,7 @@ impl UnifiedExecSessionManager { } pub(crate) async fn terminate_all_sessions(&self) { - let mut sessions = self.sessions.lock().await; + let mut sessions = self.session_store.lock().await; sessions.clear(); } } From cfac4347d652b6ee8badb7c62428f26efadb925e Mon Sep 17 00:00:00 2001 From: jif-oai Date: Fri, 28 Nov 2025 13:20:04 +0000 Subject: [PATCH 3/5] Nit --- codex-rs/core/src/unified_exec/mod.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs index a368bf8f5c..8d8a3f3148 100644 --- a/codex-rs/core/src/unified_exec/mod.rs +++ b/codex-rs/core/src/unified_exec/mod.rs @@ -108,9 +108,9 @@ pub(crate) struct SessionStore { } impl SessionStore { - pub fn remove(&mut self, session_id: &str) { - self.sessions.remove(&session_id); + pub fn remove(&mut self, session_id: &str) -> Option { self.reserved_sessions_id.remove(&session_id); + self.sessions.remove(&session_id) } pub fn clear(&mut self) { From 5a7547d9974cfad470d29cd71f57917eaa4ed6bf Mon Sep 17 00:00:00 2001 From: jif-oai Date: Fri, 28 Nov 2025 13:28:13 +0000 Subject: [PATCH 4/5] Better --- codex-rs/core/src/unified_exec/mod.rs | 13 ++++++------ .../core/src/unified_exec/session_manager.rs | 21 ++++++++++++------- 2 files changed, 20 insertions(+), 14 deletions(-) diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs index 8d8a3f3148..8dd3e7099f 100644 --- a/codex-rs/core/src/unified_exec/mod.rs +++ b/codex-rs/core/src/unified_exec/mod.rs @@ -21,7 +21,8 @@ //! - `session.rs`: PTY session lifecycle + output buffering. //! - `session_manager.rs`: orchestration (approvals, sandboxing, reuse) and request handling. -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; +use std::collections::HashSet; use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; @@ -104,16 +105,16 @@ pub(crate) struct UnifiedExecSessionManager { #[derive(Default)] pub(crate) struct SessionStore { sessions: HashMap, - reserved_sessions_id: HashSet + reserved_sessions_id: HashSet, } impl SessionStore { - pub fn remove(&mut self, session_id: &str) -> Option { - self.reserved_sessions_id.remove(&session_id); - self.sessions.remove(&session_id) + pub(crate) fn remove(&mut self, session_id: &str) -> Option { + self.reserved_sessions_id.remove(session_id); + self.sessions.remove(session_id) } - pub fn clear(&mut self) { + pub(crate) fn clear(&mut self) { self.reserved_sessions_id.clear(); self.sessions.clear(); } diff --git a/codex-rs/core/src/unified_exec/session_manager.rs b/codex-rs/core/src/unified_exec/session_manager.rs index 1a13fd3ca3..23044c0f78 100644 --- a/codex-rs/core/src/unified_exec/session_manager.rs +++ b/codex-rs/core/src/unified_exec/session_manager.rs @@ -33,9 +33,10 @@ use crate::truncate::TruncationPolicy; use crate::truncate::approx_token_count; use crate::truncate::formatted_truncate_text; -use super::{ExecCommandRequest, SessionStore}; +use super::ExecCommandRequest; use super::MAX_UNIFIED_EXEC_SESSIONS; use super::SessionEntry; +use super::SessionStore; use super::UnifiedExecContext; use super::UnifiedExecError; use super::UnifiedExecResponse; @@ -88,7 +89,8 @@ impl UnifiedExecSessionManager { rand::rng().random_range(1_000..100_000).to_string() } else { // test or deterministic mode - let next = store.reserved_sessions_id + let next = store + .reserved_sessions_id .iter() .filter_map(|s| s.parse::().ok()) .max() @@ -361,11 +363,13 @@ impl UnifiedExecSessionManager { process_id: &str, ) -> Result { let mut store = self.session_store.lock().await; - let entry = store.sessions - .get_mut(process_id) - .ok_or(UnifiedExecError::UnknownSessionId { - process_id: process_id.to_string(), - })?; + let entry = + store + .sessions + .get_mut(process_id) + .ok_or(UnifiedExecError::UnknownSessionId { + process_id: process_id.to_string(), + })?; entry.last_used = Instant::now(); let OutputHandles { output_buffer, @@ -634,7 +638,8 @@ impl UnifiedExecSessionManager { return; } - let meta: Vec<(String, Instant, bool)> = store.sessions + let meta: Vec<(String, Instant, bool)> = store + .sessions .iter() .map(|(id, entry)| (id.clone(), entry.last_used, entry.session.has_exited())) .collect(); From 2436f5176d95a028a82133eb2ec5a20c56cd1573 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Fri, 28 Nov 2025 14:28:06 +0000 Subject: [PATCH 5/5] clippy --- codex-rs/core/src/unified_exec/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs index 8dd3e7099f..8c33548065 100644 --- a/codex-rs/core/src/unified_exec/mod.rs +++ b/codex-rs/core/src/unified_exec/mod.rs @@ -109,7 +109,7 @@ pub(crate) struct SessionStore { } impl SessionStore { - pub(crate) fn remove(&mut self, session_id: &str) -> Option { + fn remove(&mut self, session_id: &str) -> Option { self.reserved_sessions_id.remove(session_id); self.sessions.remove(session_id) }