From 63f009e9dad2e70454b7ed6434d8aa28dfb52b51 Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Sun, 21 Jun 2026 11:18:36 -0700 Subject: [PATCH] code-mode: move session ownership into runtime (#29285) ## Summary - Move code-mode cell ownership and shared stored values from `CodeModeService` into `SessionRuntime`. - Keep the protocol-facing execute/wait behavior behind the existing service adapter. - Add runtime-level ownership and isolation coverage. ## Why This establishes a transport-neutral session boundary before later lifecycle and create/observe changes. ## Impact No intended model-facing behavior change. This is an ownership and layering refactor. ## Validation - Stack-tip validation: `just test -p codex-code-mode -p codex-code-mode-protocol` (70 passed). - Parent branch: `cconger/code-mode-runtime-compact-03a-runtime-types`. --- codex-rs/code-mode/src/service.rs | 346 +++++++----------- .../code-mode/src/service_contract_tests.rs | 16 +- codex-rs/code-mode/src/service_tests.rs | 11 +- codex-rs/code-mode/src/session_runtime/mod.rs | 268 ++++++++++++++ .../code-mode/src/session_runtime/types.rs | 93 +++++ 5 files changed, 510 insertions(+), 224 deletions(-) diff --git a/codex-rs/code-mode/src/service.rs b/codex-rs/code-mode/src/service.rs index bf32156a0..bb7aa842e 100644 --- a/codex-rs/code-mode/src/service.rs +++ b/codex-rs/code-mode/src/service.rs @@ -1,8 +1,4 @@ -use std::collections::HashMap; use std::sync::Arc; -use std::sync::atomic::AtomicBool; -use std::sync::atomic::AtomicU64; -use std::sync::atomic::Ordering; use std::time::Duration; use codex_code_mode_protocol::CellId; @@ -27,24 +23,11 @@ use codex_code_mode_protocol::WaitRequest; use codex_code_mode_protocol::WaitToPendingOutcome; use codex_code_mode_protocol::WaitToPendingRequest; use serde_json::Value as JsonValue; -use tokio::sync::Mutex; use tokio::sync::oneshot; use tokio_util::sync::CancellationToken; -use crate::cell_actor::CellActor; -use crate::cell_actor::CellError; -use crate::cell_actor::CellEventFuture; -use crate::cell_actor::CellHandle; -use crate::cell_actor::CellHost; -use crate::cell_actor::CellToolCall; -use crate::session_runtime::CellEvent; -use crate::session_runtime::CreateCellRequest as CellRequest; -use crate::session_runtime::ImageDetail as CellImageDetail; -use crate::session_runtime::ObserveMode; -use crate::session_runtime::OutputItem as CellOutputItem; -use crate::session_runtime::ToolDefinition as CellToolDefinition; -use crate::session_runtime::ToolKind as CellToolKind; -use crate::session_runtime::ToolName as CellToolName; +use crate::session_runtime as runtime; +use crate::session_runtime::SessionRuntime; pub struct NoopCodeModeSessionDelegate; @@ -89,16 +72,8 @@ impl CodeModeSessionProvider for InProcessCodeModeSessionProvider { } } -struct Inner { - stored_values: Mutex>, - cells: Mutex>, - delegate: Arc, - shutting_down: AtomicBool, - next_cell_id: AtomicU64, -} - pub struct CodeModeService { - inner: Arc, + runtime: SessionRuntime, } impl CodeModeService { @@ -108,45 +83,31 @@ impl CodeModeService { pub fn with_delegate(delegate: Arc) -> Self { Self { - inner: Arc::new(Inner { - stored_values: Mutex::new(HashMap::new()), - cells: Mutex::new(HashMap::new()), - delegate, - shutting_down: AtomicBool::new(false), - next_cell_id: AtomicU64::new(1), - }), + runtime: SessionRuntime::new(Arc::new(ProtocolDelegate { delegate })), } } - fn allocate_cell_id(&self) -> CellId { - CellId::new( - self.inner - .next_cell_id - .fetch_add(1, Ordering::Relaxed) - .to_string(), - ) - } - pub async fn execute(&self, request: ExecuteRequest) -> Result { let yield_time_ms = request.yield_time_ms.unwrap_or(DEFAULT_EXEC_YIELD_TIME_MS); - let cell_id = self.allocate_cell_id(); - let initial_event = self - .start_cell( - cell_id.clone(), - request, - ObserveMode::YieldAfter(Duration::from_millis(yield_time_ms)), + let started = self + .runtime + .execute( + runtime_request(request), + runtime::ObserveMode::YieldAfter(Duration::from_millis(yield_time_ms)), ) - .await?; + .await + .map_err(|error| error.to_string())?; + let cell_id = protocol_cell_id(&started.cell_id); let response_cell_id = cell_id.clone(); let (response_tx, response_rx) = oneshot::channel(); tokio::spawn(async move { - let response = initial_event + let response = started + .initial_event() .await - .map_err(|error| cell_error_text(&response_cell_id, error)) + .map_err(|error| error.to_string()) .and_then(|event| runtime_response(&response_cell_id, event)); let _ = response_tx.send(response); }); - Ok(StartedCell::from_result_receiver(cell_id, response_rx)) } @@ -154,45 +115,22 @@ impl CodeModeService { &self, request: ExecuteRequest, ) -> Result { - let cell_id = self.allocate_cell_id(); - let event = self - .start_cell(cell_id.clone(), request, ObserveMode::PendingFrontier) - .await? + let started = self + .runtime + .execute( + runtime_request(request), + runtime::ObserveMode::PendingFrontier, + ) .await - .map_err(|error| cell_error_text(&cell_id, error))?; + .map_err(|error| error.to_string())?; + let cell_id = protocol_cell_id(&started.cell_id); + let event = started + .initial_event() + .await + .map_err(|error| error.to_string())?; pending_outcome(&cell_id, event) } - async fn start_cell( - &self, - cell_id: CellId, - request: ExecuteRequest, - initial_observe_mode: ObserveMode, - ) -> Result { - let stored_values = self.inner.stored_values.lock().await.clone(); - let host = Arc::new(ServiceCellHost { - cell_id: cell_id.clone(), - inner: Arc::clone(&self.inner), - }); - let mut cells = self.inner.cells.lock().await; - if self.inner.shutting_down.load(Ordering::Acquire) { - return Err("code mode session is shutting down".to_string()); - } - if cells.contains_key(&cell_id) { - return Err(format!("exec cell {cell_id} already exists")); - } - let (handle, initial_event, task) = CellActor::prepare( - cell_request(request), - stored_values, - host, - initial_observe_mode, - )?; - cells.insert(cell_id, handle); - drop(cells); - tokio::spawn(task); - Ok(initial_event) - } - pub async fn wait(&self, request: WaitRequest) -> Result { self.begin_wait(request).await.await } @@ -205,24 +143,39 @@ impl CodeModeService { cell_id, yield_time_ms, } = request; - let handle = self.inner.cells.lock().await.get(&cell_id).cloned(); - let Some(handle) = handle else { - return missing_wait(cell_id); - }; - wait_for_event( - cell_id, - handle.observe(ObserveMode::YieldAfter(Duration::from_millis( - yield_time_ms, - ))), - ) + let runtime_cell_id = runtime_cell_id(&cell_id); + match self + .runtime + .begin_observe( + &runtime_cell_id, + runtime::ObserveMode::YieldAfter(Duration::from_millis(yield_time_ms)), + ) + .await + { + Ok(pending_event) => Box::pin(async move { + match pending_event.event().await { + Ok(event) => Ok(WaitOutcome::LiveCell(runtime_response(&cell_id, event)?)), + Err(runtime::Error::MissingCell(_) | runtime::Error::ClosedCell(_)) => { + Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))) + } + Err(error) => Err(error.to_string()), + } + }), + Err(runtime::Error::MissingCell(_) | runtime::Error::ClosedCell(_)) => { + missing_wait(cell_id) + } + Err(error) => Box::pin(async move { Err(error.to_string()) }), + } } pub async fn terminate(&self, cell_id: CellId) -> Result { - let handle = self.inner.cells.lock().await.get(&cell_id).cloned(); - let Some(handle) = handle else { - return Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))); - }; - wait_for_event(cell_id, handle.terminate()).await + match self.runtime.terminate(&runtime_cell_id(&cell_id)).await { + Ok(event) => Ok(WaitOutcome::LiveCell(runtime_response(&cell_id, event)?)), + Err(runtime::Error::MissingCell(_) | runtime::Error::ClosedCell(_)) => { + Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))) + } + Err(error) => Err(error.to_string()), + } } pub async fn wait_to_pending( @@ -230,40 +183,29 @@ impl CodeModeService { request: WaitToPendingRequest, ) -> Result { let cell_id = request.cell_id; - let handle = self.inner.cells.lock().await.get(&cell_id).cloned(); - let Some(handle) = handle else { - return Ok(WaitToPendingOutcome::MissingCell(missing_cell_response( - cell_id, - ))); - }; - match handle.observe(ObserveMode::PendingFrontier).await { + match self + .runtime + .observe( + &runtime_cell_id(&cell_id), + runtime::ObserveMode::PendingFrontier, + ) + .await + { Ok(event) => Ok(WaitToPendingOutcome::LiveCell(pending_outcome( &cell_id, event, )?)), - Err(CellError::Closed) => Ok(WaitToPendingOutcome::MissingCell(missing_cell_response( - cell_id, - ))), - Err(error) => Err(cell_error_text(&cell_id, error)), + Err(runtime::Error::MissingCell(_) | runtime::Error::ClosedCell(_)) => Ok( + WaitToPendingOutcome::MissingCell(missing_cell_response(cell_id)), + ), + Err(error) => Err(error.to_string()), } } pub async fn shutdown(&self) -> Result<(), String> { - self.inner.shutting_down.store(true, Ordering::Release); - let handles = self - .inner - .cells - .lock() + self.runtime + .shutdown() .await - .values() - .cloned() - .collect::>(); - for handle in handles { - handle.shutdown(); - } - while !self.inner.cells.lock().await.is_empty() { - tokio::task::yield_now().await; - } - Ok(()) + .map_err(|error| error.to_string()) } } @@ -273,20 +215,9 @@ impl Default for CodeModeService { } } -impl Drop for CodeModeService { - fn drop(&mut self) { - self.inner.shutting_down.store(true, Ordering::Release); - if let Ok(cells) = self.inner.cells.try_lock() { - for handle in cells.values() { - handle.shutdown(); - } - } - } -} - impl CodeModeSession for CodeModeService { fn is_alive(&self) -> bool { - !self.inner.shutting_down.load(Ordering::Acquire) + self.runtime.is_alive() } fn execute<'a>( @@ -309,30 +240,28 @@ impl CodeModeSession for CodeModeService { } } -struct ServiceCellHost { - cell_id: CellId, - inner: Arc, +struct ProtocolDelegate { + delegate: Arc, } -impl CellHost for ServiceCellHost { +impl runtime::SessionRuntimeDelegate for ProtocolDelegate { async fn invoke_tool( &self, - invocation: CellToolCall, + invocation: runtime::NestedToolCall, cancellation_token: CancellationToken, ) -> Result { - self.inner - .delegate + self.delegate .invoke_tool( CodeModeNestedToolCall { - cell_id: self.cell_id.clone(), - runtime_tool_call_id: invocation.id, + cell_id: protocol_cell_id(&invocation.cell_id), + runtime_tool_call_id: invocation.runtime_tool_call_id, tool_name: codex_protocol::ToolName { - name: invocation.name.name, - namespace: invocation.name.namespace, + name: invocation.tool_name.name, + namespace: invocation.tool_name.namespace, }, - tool_kind: match invocation.kind { - CellToolKind::Function => CodeModeToolKind::Function, - CellToolKind::Freeform => CodeModeToolKind::Freeform, + tool_kind: match invocation.tool_kind { + runtime::ToolKind::Function => CodeModeToolKind::Function, + runtime::ToolKind::Freeform => CodeModeToolKind::Freeform, }, input: invocation.input, }, @@ -344,45 +273,41 @@ impl CellHost for ServiceCellHost { async fn notify( &self, call_id: String, + cell_id: runtime::CellId, text: String, cancellation_token: CancellationToken, ) -> Result<(), String> { - self.inner - .delegate - .notify(call_id, self.cell_id.clone(), text, cancellation_token) + self.delegate + .notify( + call_id, + protocol_cell_id(&cell_id), + text, + cancellation_token, + ) .await } - async fn commit_stored_values(&self, stored_value_writes: HashMap) { - self.inner - .stored_values - .lock() - .await - .extend(stored_value_writes); - } - - async fn closed(&self) { - self.inner.cells.lock().await.remove(&self.cell_id); - self.inner.delegate.cell_closed(&self.cell_id); + fn cell_closed(&self, cell_id: &runtime::CellId) { + self.delegate.cell_closed(&protocol_cell_id(cell_id)); } } -fn cell_request(request: ExecuteRequest) -> CellRequest { - CellRequest { +fn runtime_request(request: ExecuteRequest) -> runtime::CreateCellRequest { + runtime::CreateCellRequest { tool_call_id: request.tool_call_id, enabled_tools: request .enabled_tools .into_iter() - .map(|definition| CellToolDefinition { + .map(|definition| runtime::ToolDefinition { name: definition.name, - tool_name: CellToolName { + tool_name: runtime::ToolName { name: definition.tool_name.name, namespace: definition.tool_name.namespace, }, description: definition.description, kind: match definition.kind { - CodeModeToolKind::Function => CellToolKind::Function, - CodeModeToolKind::Freeform => CellToolKind::Freeform, + CodeModeToolKind::Function => runtime::ToolKind::Function, + CodeModeToolKind::Freeform => runtime::ToolKind::Freeform, }, }) .collect(), @@ -390,22 +315,20 @@ fn cell_request(request: ExecuteRequest) -> CellRequest { } } -fn wait_for_event( - cell_id: CellId, - event: CellEventFuture, -) -> CodeModeSessionResultFuture<'static, WaitOutcome> { - Box::pin(async move { - match event.await { - Ok(event) => Ok(WaitOutcome::LiveCell(runtime_response(&cell_id, event)?)), - Err(CellError::Closed) => Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))), - Err(error) => Err(cell_error_text(&cell_id, error)), - } - }) +fn runtime_cell_id(cell_id: &CellId) -> runtime::CellId { + runtime::CellId::new(cell_id.as_str()) } -fn pending_outcome(cell_id: &CellId, event: CellEvent) -> Result { +fn protocol_cell_id(cell_id: &runtime::CellId) -> CellId { + CellId::new(cell_id.as_str().to_string()) +} + +fn pending_outcome( + cell_id: &CellId, + event: runtime::CellEvent, +) -> Result { match event { - CellEvent::Pending { + runtime::CellEvent::Pending { content_items, pending_tool_call_ids, } => Ok(ExecuteToPendingOutcome::Pending { @@ -419,13 +342,16 @@ fn pending_outcome(cell_id: &CellId, event: CellEvent) -> Result Result { +fn runtime_response( + cell_id: &CellId, + event: runtime::CellEvent, +) -> Result { match event { - CellEvent::Yielded { content_items } => Ok(RuntimeResponse::Yielded { + runtime::CellEvent::Yielded { content_items } => Ok(RuntimeResponse::Yielded { cell_id: cell_id.clone(), content_items: content_items.into_iter().map(output_item).collect(), }), - CellEvent::Completed { + runtime::CellEvent::Completed { content_items, error_text, } => Ok(RuntimeResponse::Result { @@ -433,36 +359,30 @@ fn runtime_response(cell_id: &CellId, event: CellEvent) -> Result Ok(RuntimeResponse::Terminated { + runtime::CellEvent::Terminated { content_items } => Ok(RuntimeResponse::Terminated { cell_id: cell_id.clone(), content_items: content_items.into_iter().map(output_item).collect(), }), - CellEvent::Pending { .. } => { + runtime::CellEvent::Pending { .. } => { Err("cell returned a pending frontier unexpectedly".to_string()) } } } -fn output_item(item: CellOutputItem) -> FunctionCallOutputContentItem { +fn output_item(item: runtime::OutputItem) -> FunctionCallOutputContentItem { match item { - CellOutputItem::Text { text } => FunctionCallOutputContentItem::InputText { text }, - CellOutputItem::Image { image_url, detail } => FunctionCallOutputContentItem::InputImage { - image_url, - detail: detail.map(|detail| match detail { - CellImageDetail::Auto => ImageDetail::Auto, - CellImageDetail::Low => ImageDetail::Low, - CellImageDetail::High => ImageDetail::High, - CellImageDetail::Original => ImageDetail::Original, - }), - }, - } -} - -fn cell_error_text(cell_id: &CellId, error: CellError) -> String { - match error { - CellError::Busy => format!("exec cell {cell_id} already has an active observer"), - CellError::AlreadyTerminating => format!("exec cell {cell_id} is already terminating"), - CellError::Closed => format!("exec cell {cell_id} closed unexpectedly"), + runtime::OutputItem::Text { text } => FunctionCallOutputContentItem::InputText { text }, + runtime::OutputItem::Image { image_url, detail } => { + FunctionCallOutputContentItem::InputImage { + image_url, + detail: detail.map(|detail| match detail { + runtime::ImageDetail::Auto => ImageDetail::Auto, + runtime::ImageDetail::Low => ImageDetail::Low, + runtime::ImageDetail::High => ImageDetail::High, + runtime::ImageDetail::Original => ImageDetail::Original, + }), + } + } } } diff --git a/codex-rs/code-mode/src/service_contract_tests.rs b/codex-rs/code-mode/src/service_contract_tests.rs index 3d814e98c..948431032 100644 --- a/codex-rs/code-mode/src/service_contract_tests.rs +++ b/codex-rs/code-mode/src/service_contract_tests.rs @@ -288,8 +288,20 @@ async fn observed_natural_completion_wins_over_termination() { ); tokio::time::timeout(Duration::from_secs(1), async { loop { - if service.inner.stored_values.lock().await.get("finished") - == Some(&serde_json::json!(true)) + let response = service + .execute(execute_request(r#"text(String(load("finished")));"#)) + .await + .unwrap() + .initial_response() + .await + .unwrap(); + let RuntimeResponse::Result { content_items, .. } = response else { + panic!("expected stored-value probe to complete"); + }; + if content_items + == vec![FunctionCallOutputContentItem::InputText { + text: "true".to_string(), + }] { break; } diff --git a/codex-rs/code-mode/src/service_tests.rs b/codex-rs/code-mode/src/service_tests.rs index 688bd89cd..d90d4b6af 100644 --- a/codex-rs/code-mode/src/service_tests.rs +++ b/codex-rs/code-mode/src/service_tests.rs @@ -1,5 +1,4 @@ use std::sync::Arc; -use std::sync::atomic::Ordering; use std::time::Duration; use super::CellId; @@ -7,7 +6,6 @@ use super::CodeModeNestedToolCall; use super::CodeModeService; use super::CodeModeSessionDelegate; use super::NotificationFuture; -use super::ObserveMode; use super::RuntimeResponse; use super::ToolInvocationFuture; use super::WaitOutcome; @@ -216,20 +214,15 @@ async fn shutdown_interrupts_cpu_bound_cells() { #[tokio::test] async fn start_cell_rejects_new_cell_after_shutdown_begins() { let service = CodeModeService::new(); - service.inner.shutting_down.store(true, Ordering::Release); + service.shutdown().await.unwrap(); let error = service - .start_cell( - cell_id("late-cell"), - execute_request(""), - ObserveMode::YieldAfter(Duration::from_millis(1)), - ) + .execute(execute_request("text('late');")) .await .err() .unwrap(); assert_eq!(error, "code mode session is shutting down".to_string()); - assert!(service.inner.cells.lock().await.is_empty()); } #[tokio::test] diff --git a/codex-rs/code-mode/src/session_runtime/mod.rs b/codex-rs/code-mode/src/session_runtime/mod.rs index bb09e0e43..f72aadb10 100644 --- a/codex-rs/code-mode/src/session_runtime/mod.rs +++ b/codex-rs/code-mode/src/session_runtime/mod.rs @@ -1,10 +1,278 @@ mod types; +use std::collections::HashMap; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; + +use serde_json::Value as JsonValue; +use tokio::sync::Mutex; +use tokio_util::sync::CancellationToken; + pub(crate) use self::types::CellEvent; +pub(crate) use self::types::CellId; pub(crate) use self::types::CreateCellRequest; +pub(crate) use self::types::Error; pub(crate) use self::types::ImageDetail; +pub(crate) use self::types::NestedToolCall; pub(crate) use self::types::ObserveMode; pub(crate) use self::types::OutputItem; +pub(crate) use self::types::SessionRuntimeDelegate; pub(crate) use self::types::ToolDefinition; pub(crate) use self::types::ToolKind; pub(crate) use self::types::ToolName; +use crate::cell_actor::CellActor; +use crate::cell_actor::CellError; +use crate::cell_actor::CellEventFuture; +use crate::cell_actor::CellHandle; +use crate::cell_actor::CellHost; +use crate::cell_actor::CellToolCall; + +type RuntimeEventFuture = Pin> + Send + 'static>>; + +/// Owns all cells and shared state for one transport-neutral code-mode session. +pub(crate) struct SessionRuntime { + inner: Arc>, +} + +struct Inner { + stored_values: Mutex>, + cells: Mutex>, + delegate: Arc, + shutting_down: AtomicBool, + next_cell_id: AtomicU64, +} + +impl SessionRuntime { + pub(crate) fn new(delegate: Arc) -> Self { + Self { + inner: Arc::new(Inner { + stored_values: Mutex::new(HashMap::new()), + cells: Mutex::new(HashMap::new()), + delegate, + shutting_down: AtomicBool::new(false), + next_cell_id: AtomicU64::new(1), + }), + } + } + + pub(crate) fn is_alive(&self) -> bool { + !self.inner.shutting_down.load(Ordering::Acquire) + } + + pub(crate) async fn execute( + &self, + request: CreateCellRequest, + initial_observe_mode: ObserveMode, + ) -> Result { + let cell_id = self.allocate_cell_id(); + let initial_event = self + .start_cell(cell_id.clone(), request, initial_observe_mode) + .await?; + Ok(StartedCell { + cell_id, + initial_event, + }) + } + + pub(crate) async fn observe( + &self, + cell_id: &CellId, + mode: ObserveMode, + ) -> Result { + self.begin_observe(cell_id, mode).await?.event().await + } + + pub(crate) async fn begin_observe( + &self, + cell_id: &CellId, + mode: ObserveMode, + ) -> Result { + let handle = self + .inner + .cells + .lock() + .await + .get(cell_id) + .cloned() + .ok_or_else(|| Error::MissingCell(cell_id.clone()))?; + Ok(PendingEvent { + event: map_actor_event(cell_id.clone(), handle.observe(mode)), + }) + } + + pub(crate) async fn terminate(&self, cell_id: &CellId) -> Result { + let handle = self + .inner + .cells + .lock() + .await + .get(cell_id) + .cloned() + .ok_or_else(|| Error::MissingCell(cell_id.clone()))?; + handle + .terminate() + .await + .map_err(|error| actor_error(cell_id, error)) + } + + pub(crate) async fn shutdown(&self) -> Result<(), Error> { + self.inner.shutting_down.store(true, Ordering::Release); + let handles = self + .inner + .cells + .lock() + .await + .values() + .cloned() + .collect::>(); + for handle in handles { + handle.shutdown(); + } + while !self.inner.cells.lock().await.is_empty() { + tokio::task::yield_now().await; + } + Ok(()) + } + + fn allocate_cell_id(&self) -> CellId { + CellId::new( + self.inner + .next_cell_id + .fetch_add(1, Ordering::Relaxed) + .to_string(), + ) + } + + async fn start_cell( + &self, + cell_id: CellId, + request: CreateCellRequest, + initial_observe_mode: ObserveMode, + ) -> Result { + let stored_values = self.inner.stored_values.lock().await.clone(); + let host = Arc::new(RuntimeCellHost { + cell_id: cell_id.clone(), + inner: Arc::clone(&self.inner), + }); + let mut cells = self.inner.cells.lock().await; + if self.inner.shutting_down.load(Ordering::Acquire) { + return Err(Error::ShuttingDown); + } + if cells.contains_key(&cell_id) { + return Err(Error::DuplicateCell(cell_id)); + } + let (handle, initial_event, task) = + CellActor::prepare(request, stored_values, host, initial_observe_mode) + .map_err(Error::Runtime)?; + cells.insert(cell_id.clone(), handle); + drop(cells); + tokio::spawn(task); + Ok(map_actor_event(cell_id, initial_event)) + } + + fn begin_shutdown(&self) { + self.inner.shutting_down.store(true, Ordering::Release); + if let Ok(cells) = self.inner.cells.try_lock() { + for handle in cells.values() { + handle.shutdown(); + } + } + } +} + +impl Drop for SessionRuntime { + fn drop(&mut self) { + self.begin_shutdown(); + } +} + +/// A cell admitted by [`SessionRuntime::execute`]. +pub(crate) struct StartedCell { + pub(crate) cell_id: CellId, + initial_event: RuntimeEventFuture, +} + +impl StartedCell { + pub(crate) async fn initial_event(self) -> Result { + self.initial_event.await + } +} + +/// An admitted observation that has not reached its requested frontier yet. +pub(crate) struct PendingEvent { + event: RuntimeEventFuture, +} + +impl PendingEvent { + pub(crate) async fn event(self) -> Result { + self.event.await + } +} + +struct RuntimeCellHost { + cell_id: CellId, + inner: Arc>, +} + +impl CellHost for RuntimeCellHost { + async fn invoke_tool( + &self, + invocation: CellToolCall, + cancellation_token: CancellationToken, + ) -> Result { + self.inner + .delegate + .invoke_tool( + NestedToolCall { + cell_id: self.cell_id.clone(), + runtime_tool_call_id: invocation.id, + tool_name: invocation.name, + tool_kind: invocation.kind, + input: invocation.input, + }, + cancellation_token, + ) + .await + } + + async fn notify( + &self, + call_id: String, + text: String, + cancellation_token: CancellationToken, + ) -> Result<(), String> { + self.inner + .delegate + .notify(call_id, self.cell_id.clone(), text, cancellation_token) + .await + } + + async fn commit_stored_values(&self, stored_value_writes: HashMap) { + self.inner + .stored_values + .lock() + .await + .extend(stored_value_writes); + } + + async fn closed(&self) { + self.inner.cells.lock().await.remove(&self.cell_id); + self.inner.delegate.cell_closed(&self.cell_id); + } +} + +fn map_actor_event(cell_id: CellId, event: CellEventFuture) -> RuntimeEventFuture { + Box::pin(async move { event.await.map_err(|error| actor_error(&cell_id, error)) }) +} + +fn actor_error(cell_id: &CellId, error: CellError) -> Error { + match error { + CellError::Busy => Error::BusyObserver(cell_id.clone()), + CellError::AlreadyTerminating => Error::AlreadyTerminating(cell_id.clone()), + CellError::Closed => Error::ClosedCell(cell_id.clone()), + } +} diff --git a/codex-rs/code-mode/src/session_runtime/types.rs b/codex-rs/code-mode/src/session_runtime/types.rs index f3719c96c..b04bd0f23 100644 --- a/codex-rs/code-mode/src/session_runtime/types.rs +++ b/codex-rs/code-mode/src/session_runtime/types.rs @@ -1,5 +1,30 @@ +use std::fmt; +use std::future::Future; use std::time::Duration; +use serde_json::Value as JsonValue; +use tokio_util::sync::CancellationToken; + +/// Identifies one execution cell within a session runtime. +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +pub(crate) struct CellId(String); + +impl CellId { + pub(crate) fn new(value: impl Into) -> Self { + Self(value.into()) + } + + pub(crate) fn as_str(&self) -> &str { + &self.0 + } +} + +impl fmt::Display for CellId { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(self.as_str()) + } +} + /// Selects the next observable frontier for a running cell. #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub(crate) enum ObserveMode { @@ -77,3 +102,71 @@ pub(crate) enum ToolKind { Function, Freeform, } + +/// A nested tool request emitted by a running cell. +pub(crate) struct NestedToolCall { + pub(crate) cell_id: CellId, + pub(crate) runtime_tool_call_id: String, + pub(crate) tool_name: ToolName, + pub(crate) tool_kind: ToolKind, + pub(crate) input: Option, +} + +/// Host callbacks used by cells owned by a [`super::SessionRuntime`]. +/// +/// Implementations must honor cancellation tokens. `cell_closed` is called +/// after the runtime has stopped routing requests to the cell. +pub(crate) trait SessionRuntimeDelegate: Send + Sync + 'static { + fn invoke_tool( + &self, + invocation: NestedToolCall, + cancellation_token: CancellationToken, + ) -> impl Future> + Send; + + fn notify( + &self, + call_id: String, + cell_id: CellId, + text: String, + cancellation_token: CancellationToken, + ) -> impl Future> + Send; + + fn cell_closed(&self, cell_id: &CellId); +} + +/// A failure reported by a session runtime operation. +#[derive(Clone, Debug, Eq, PartialEq)] +pub(crate) enum Error { + ShuttingDown, + DuplicateCell(CellId), + MissingCell(CellId), + BusyObserver(CellId), + AlreadyTerminating(CellId), + ClosedCell(CellId), + Runtime(String), +} + +impl fmt::Display for Error { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::ShuttingDown => formatter.write_str("code mode session is shutting down"), + Self::DuplicateCell(cell_id) => write!(formatter, "exec cell {cell_id} already exists"), + Self::MissingCell(cell_id) => write!(formatter, "exec cell {cell_id} not found"), + Self::BusyObserver(cell_id) => { + write!( + formatter, + "exec cell {cell_id} already has an active observer" + ) + } + Self::AlreadyTerminating(cell_id) => { + write!(formatter, "exec cell {cell_id} is already terminating") + } + Self::ClosedCell(cell_id) => { + write!(formatter, "exec cell {cell_id} closed unexpectedly") + } + Self::Runtime(error_text) => formatter.write_str(error_text), + } + } +} + +impl std::error::Error for Error {}