From b44eb3bb3509893d48c4aa7f029948d48c24ae10 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Wed, 11 Mar 2026 17:39:47 -0700 Subject: [PATCH] Simplify code-mode shared session state --- codex-rs/core/src/state/service.rs | 34 ++++---- codex-rs/core/src/tools/code_mode.rs | 116 +++++++++------------------ 2 files changed, 52 insertions(+), 98 deletions(-) diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index cc4791ae23..8621d7dc31 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -1,4 +1,5 @@ use std::collections::HashMap; +use std::collections::HashSet; use std::sync::Arc; use crate::AuthManager; @@ -16,7 +17,6 @@ use crate::plugins::PluginsManager; use crate::skills::SkillsManager; use crate::state_db::StateDbHandle; use crate::tools::code_mode::CodeModeProcess; -use crate::tools::code_mode::CodeModeYieldedSession; use crate::tools::network_approval::NetworkApprovalService; use crate::tools::runtimes::ExecveSessionApproval; use crate::tools::sandboxing::ApprovalStore; @@ -33,8 +33,8 @@ use tokio_util::sync::CancellationToken; pub(crate) struct CodeModeStoreService { stored_values: Mutex>, - process: Mutex>>>, - yielded_sessions: Mutex>, + process: Arc>>, + yielded_sessions: Mutex>, next_session_id: Mutex, } @@ -42,8 +42,8 @@ impl Default for CodeModeStoreService { fn default() -> Self { Self { stored_values: Mutex::new(HashMap::new()), - process: Mutex::new(None), - yielded_sessions: Mutex::new(HashMap::new()), + process: Arc::new(Mutex::new(None)), + yielded_sessions: Mutex::new(HashSet::new()), next_session_id: Mutex::new(1), } } @@ -58,12 +58,8 @@ impl CodeModeStoreService { *self.stored_values.lock().await = values; } - pub(crate) async fn store_process(&self, process: Arc>) { - *self.process.lock().await = Some(process); - } - - pub(crate) async fn process(&self) -> Option>> { - self.process.lock().await.clone() + pub(crate) fn process(&self) -> Arc>> { + self.process.clone() } pub(crate) async fn allocate_session_id(&self) -> i32 { @@ -73,18 +69,16 @@ impl CodeModeStoreService { session_id } - pub(crate) async fn store_yielded_session(&self, yielded_session: CodeModeYieldedSession) { + pub(crate) async fn store_yielded_session(&self, session_id: i32) { + self.yielded_sessions.lock().await.insert(session_id); + } + + pub(crate) async fn take_yielded_session(&self, session_id: i32) -> Option { self.yielded_sessions .lock() .await - .insert(yielded_session.session_id, yielded_session); - } - - pub(crate) async fn take_yielded_session( - &self, - session_id: i32, - ) -> Option { - self.yielded_sessions.lock().await.remove(&session_id) + .remove(&session_id) + .then_some(session_id) } } diff --git a/codex-rs/core/src/tools/code_mode.rs b/codex-rs/core/src/tools/code_mode.rs index 945310cb84..40ea237fed 100644 --- a/codex-rs/core/src/tools/code_mode.rs +++ b/codex-rs/core/src/tools/code_mode.rs @@ -78,11 +78,6 @@ impl CodeModeProcess { } } -#[derive(Clone, Debug)] -pub(crate) struct CodeModeYieldedSession { - pub(crate) session_id: i32, -} - #[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] #[serde(rename_all = "snake_case")] enum CodeModeToolKind { @@ -176,7 +171,7 @@ enum CodeModeSessionProgress { Finished(FunctionToolOutput), Yielded { output: FunctionToolOutput, - yielded_session: CodeModeYieldedSession, + session_id: i32, }, } @@ -241,14 +236,26 @@ pub(crate) async fn execute( .code_mode_store .allocate_session_id() .await; - let process = ensure_shared_code_mode_process(&exec) - .await - .map_err(FunctionCallError::RespondToModel)?; let result = { - let mut process = process.lock().await; + let process_slot = exec.session.services.code_mode_store.process(); + let mut process_slot = process_slot.lock().await; + let needs_spawn = match process_slot.as_mut() { + Some(process) => !matches!(process.has_exited(), Ok(false)), + None => true, + }; + if needs_spawn { + *process_slot = Some( + spawn_code_mode_process(&exec) + .await + .map_err(FunctionCallError::RespondToModel)?, + ); + } + let process = process_slot.as_mut().ok_or_else(|| { + FunctionCallError::RespondToModel(format!("{PUBLIC_TOOL_NAME} runner failed to start")) + })?; drive_code_mode_session( &exec, - &mut process, + process, session_id, CodeModeSessionAction::Start { enabled_tools, @@ -258,14 +265,11 @@ pub(crate) async fn execute( ) .await }; - if let Ok(CodeModeSessionProgress::Yielded { - yielded_session, .. - }) = &result - { + if let Ok(CodeModeSessionProgress::Yielded { session_id, .. }) = &result { exec.session .services .code_mode_store - .store_yielded_session(yielded_session.clone()) + .store_yielded_session(*session_id) .await; } match result { @@ -289,7 +293,7 @@ pub(crate) async fn wait( turn, tracker, }; - let yielded_session = exec + let yielded_session_id = exec .session .services .code_mode_store @@ -301,13 +305,23 @@ pub(crate) async fn wait( )) })?; - let process = existing_shared_code_mode_process(&exec).await?; let result = { - let mut process = process.lock().await; + let process_slot = exec.session.services.code_mode_store.process(); + let mut process_slot = process_slot.lock().await; + let process = process_slot.as_mut().ok_or_else(|| { + FunctionCallError::RespondToModel(format!( + "{PUBLIC_TOOL_NAME} runner is not available for {WAIT_TOOL_NAME}" + )) + })?; + if !matches!(process.has_exited(), Ok(false)) { + return Err(FunctionCallError::RespondToModel(format!( + "{PUBLIC_TOOL_NAME} runner is not available for {WAIT_TOOL_NAME}" + ))); + } drive_code_mode_session( &exec, - &mut process, - yielded_session.session_id, + process, + yielded_session_id, if terminate { CodeModeSessionAction::Terminate { max_output_tokens } } else { @@ -319,14 +333,11 @@ pub(crate) async fn wait( ) .await }; - if let Ok(CodeModeSessionProgress::Yielded { - yielded_session, .. - }) = &result - { + if let Ok(CodeModeSessionProgress::Yielded { session_id, .. }) = &result { exec.session .services .code_mode_store - .store_yielded_session(yielded_session.clone()) + .store_yielded_session(*session_id) .await; } @@ -384,57 +395,6 @@ async fn spawn_code_mode_process(exec: &ExecContext) -> Result Result>, String> { - if let Some(process) = exec.session.services.code_mode_store.process().await { - let is_running = { - let mut process_guard = process.lock().await; - matches!(process_guard.has_exited(), Ok(false)) - }; - if is_running { - return Ok(process); - } - } - - let process = Arc::new(tokio::sync::Mutex::new( - spawn_code_mode_process(exec).await?, - )); - exec.session - .services - .code_mode_store - .store_process(process.clone()) - .await; - Ok(process) -} - -async fn existing_shared_code_mode_process( - exec: &ExecContext, -) -> Result>, FunctionCallError> { - let process = exec - .session - .services - .code_mode_store - .process() - .await - .ok_or_else(|| { - FunctionCallError::RespondToModel(format!( - "{PUBLIC_TOOL_NAME} runner is not available for {WAIT_TOOL_NAME}" - )) - })?; - let is_running = { - let mut process_guard = process.lock().await; - matches!(process_guard.has_exited(), Ok(false)) - }; - if is_running { - Ok(process) - } else { - Err(FunctionCallError::RespondToModel(format!( - "{PUBLIC_TOOL_NAME} runner is not available for {WAIT_TOOL_NAME}" - ))) - } -} - async fn drive_code_mode_session( exec: &ExecContext, process: &mut CodeModeProcess, @@ -639,7 +599,7 @@ async fn handle_node_message( ); Ok(Some(CodeModeSessionProgress::Yielded { output: FunctionToolOutput::from_content(delta_items, Some(true)), - yielded_session: CodeModeYieldedSession { session_id }, + session_id, })) } NodeToHostMessage::Terminated {