openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-07 20:30:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Simplify code-mode shared session state

Browse Source
This commit is contained in:
pakrym-oai
2026-03-11 17:39:47 -07:00
parent e1cc893f95
commit b44eb3bb35
2 changed files with 52 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/state/service.rs
+14 -20
View File
@@ -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<HashMap<String, JsonValue>>,
process: Mutex<Option<Arc<Mutex<CodeModeProcess>>>>,
yielded_sessions: Mutex<HashMap<i32, CodeModeYieldedSession>>,
process: Arc<Mutex<Option<CodeModeProcess>>>,
yielded_sessions: Mutex<HashSet<i32>>,
next_session_id: Mutex<i32>,
}
@@ -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<Mutex<CodeModeProcess>>) {
*self.process.lock().await = Some(process);
}
pub(crate) async fn process(&self) -> Option<Arc<Mutex<CodeModeProcess>>> {
self.process.lock().await.clone()
pub(crate) fn process(&self) -> Arc<Mutex<Option<CodeModeProcess>>> {
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<i32> {
self.yielded_sessions
.lock()
.await
.insert(yielded_session.session_id, yielded_session);
}
pub(crate) async fn take_yielded_session(
&self,
session_id: i32,
) -> Option<CodeModeYieldedSession> {
self.yielded_sessions.lock().await.remove(&session_id)
.remove(&session_id)
.then_some(session_id)
}
}
codex-rs/core/src/tools/code_mode.rs
+38 -78
View File
@@ -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<CodeModeProcess,
})
}
async fn ensure_shared_code_mode_process(
exec: &ExecContext,
) -> Result<Arc<tokio::sync::Mutex<CodeModeProcess>>, 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<Arc<tokio::sync::Mutex<CodeModeProcess>>, 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 {