diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 061bc7769c..0b3295100e 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2499,9 +2499,11 @@ dependencies = [ "deno_core_icudata", "pretty_assertions", "serde_json", + "sha2 0.10.9", "tokio", "tokio-util", "tracing", + "uuid", "v8", ] diff --git a/codex-rs/code-mode-protocol/src/session.rs b/codex-rs/code-mode-protocol/src/session.rs index a836e27943..8dffeceabd 100644 --- a/codex-rs/code-mode-protocol/src/session.rs +++ b/codex-rs/code-mode-protocol/src/session.rs @@ -61,9 +61,6 @@ pub trait CodeModeSessionDelegate: Send + Sync { text: String, cancellation_token: CancellationToken, ) -> NotificationFuture<'a>; - - /// Releases delegate state associated with a cell after it reaches a terminal state. - fn cell_closed(&self, cell_id: &CellId); } /// A stateful code-mode session owned by one Codex thread. diff --git a/codex-rs/code-mode/Cargo.toml b/codex-rs/code-mode/Cargo.toml index 09f1d7380d..c569d04f1a 100644 --- a/codex-rs/code-mode/Cargo.toml +++ b/codex-rs/code-mode/Cargo.toml @@ -20,9 +20,11 @@ codex-code-mode-protocol = { workspace = true } codex-protocol = { workspace = true } deno_core_icudata = { workspace = true } serde_json = { workspace = true } +sha2 = { workspace = true } tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] } tokio-util = { workspace = true, features = ["rt"] } tracing = { workspace = true } +uuid = { workspace = true, features = ["v4"] } v8 = { workspace = true } [dev-dependencies] diff --git a/codex-rs/code-mode/src/lib.rs b/codex-rs/code-mode/src/lib.rs index cbe8bfe6df..079e8962d6 100644 --- a/codex-rs/code-mode/src/lib.rs +++ b/codex-rs/code-mode/src/lib.rs @@ -1,9 +1,10 @@ mod cell_actor; mod runtime; mod service; -mod session_runtime; +pub mod session_runtime; pub use codex_code_mode_protocol::*; pub use service::CodeModeService; pub use service::InProcessCodeModeSessionProvider; pub use service::NoopCodeModeSessionDelegate; +pub use session_runtime::SessionRuntime; diff --git a/codex-rs/code-mode/src/service.rs b/codex-rs/code-mode/src/service.rs index 0f3181d864..ca2955e85d 100644 --- a/codex-rs/code-mode/src/service.rs +++ b/codex-rs/code-mode/src/service.rs @@ -49,8 +49,6 @@ impl CodeModeSessionDelegate for NoopCodeModeSessionDelegate { ) -> NotificationFuture<'a> { Box::pin(async { Ok(()) }) } - - fn cell_closed(&self, _cell_id: &CellId) {} } #[derive(Default)] @@ -79,9 +77,11 @@ impl CodeModeService { } pub fn with_delegate(delegate: Arc) -> Self { - Self { - runtime: SessionRuntime::new(Arc::new(ProtocolDelegate { delegate })), - } + #[cfg(not(test))] + let runtime = SessionRuntime::new(Arc::new(ProtocolDelegate { delegate })); + #[cfg(test)] + let runtime = SessionRuntime::new_for_test(Arc::new(ProtocolDelegate { delegate })); + Self { runtime } } pub async fn create_cell(&self, request: CreateCellRequest) -> Result { @@ -251,10 +251,6 @@ impl runtime::SessionRuntimeDelegate for ProtocolDelegate { ) .await } - - fn cell_closed(&self, cell_id: &runtime::CellId) { - self.delegate.cell_closed(&protocol_cell_id(cell_id)); - } } fn runtime_request(request: CreateCellRequest) -> runtime::CreateCellRequest { diff --git a/codex-rs/code-mode/src/service_contract_tests.rs b/codex-rs/code-mode/src/service_contract_tests.rs index bbeff7f346..23338c53f9 100644 --- a/codex-rs/code-mode/src/service_contract_tests.rs +++ b/codex-rs/code-mode/src/service_contract_tests.rs @@ -6,6 +6,7 @@ use std::time::Duration; use codex_protocol::ToolName; use pretty_assertions::assert_eq; use tokio::sync::Notify; +use tokio::sync::Semaphore; use tokio::sync::mpsc; use tokio_util::sync::CancellationToken; @@ -16,8 +17,8 @@ use crate::ToolDefinition; #[derive(Debug, PartialEq)] enum DelegateEvent { NotificationStarted, + NotificationFinished, ToolStarted, - CellClosed(CellId), } struct BlockingDelegate { @@ -38,6 +39,11 @@ struct NeverResolvingNotificationDelegate { events_tx: mpsc::UnboundedSender, } +struct ReleasableNotificationDelegate { + events_tx: mpsc::UnboundedSender, + notification_release: Semaphore, +} + struct NeverResolvingToolDelegate { events_tx: mpsc::UnboundedSender, } @@ -49,6 +55,23 @@ impl NeverResolvingNotificationDelegate { } } +impl ReleasableNotificationDelegate { + fn new() -> (Arc, mpsc::UnboundedReceiver) { + let (events_tx, events_rx) = mpsc::unbounded_channel(); + ( + Arc::new(Self { + events_tx, + notification_release: Semaphore::new(/*permits*/ 0), + }), + events_rx, + ) + } + + fn release_notification(&self) { + self.notification_release.add_permits(/*permits*/ 1); + } +} + impl NeverResolvingToolDelegate { fn new() -> (Arc, mpsc::UnboundedReceiver) { let (events_tx, events_rx) = mpsc::unbounded_channel(); @@ -77,11 +100,39 @@ impl CodeModeSessionDelegate for NeverResolvingNotificationDelegate { std::future::pending().await }) } +} - fn cell_closed(&self, cell_id: &CellId) { - let _ = self - .events_tx - .send(DelegateEvent::CellClosed(cell_id.clone())); +impl CodeModeSessionDelegate for ReleasableNotificationDelegate { + fn invoke_tool<'a>( + &'a self, + _invocation: CodeModeNestedToolCall, + cancellation_token: CancellationToken, + ) -> ToolInvocationFuture<'a> { + Box::pin(async move { + cancellation_token.cancelled().await; + Err("cancelled".to_string()) + }) + } + + fn notify<'a>( + &'a self, + _call_id: String, + _cell_id: CellId, + _text: String, + cancellation_token: CancellationToken, + ) -> NotificationFuture<'a> { + Box::pin(async move { + let _ = self.events_tx.send(DelegateEvent::NotificationStarted); + tokio::select! { + _ = self.notification_release.acquire() => { + let _ = self.events_tx.send(DelegateEvent::NotificationFinished); + Ok(()) + } + _ = cancellation_token.cancelled() => { + Err("cancelled".to_string()) + } + } + }) } } @@ -106,12 +157,6 @@ impl CodeModeSessionDelegate for NeverResolvingToolDelegate { ) -> NotificationFuture<'a> { Box::pin(async { Ok(()) }) } - - fn cell_closed(&self, cell_id: &CellId) { - let _ = self - .events_tx - .send(DelegateEvent::CellClosed(cell_id.clone())); - } } impl BlockingDelegate { @@ -166,12 +211,6 @@ impl CodeModeSessionDelegate for BlockingDelegate { Err("cancelled".to_string()) }) } - - fn cell_closed(&self, cell_id: &CellId) { - let _ = self - .events_tx - .send(DelegateEvent::CellClosed(cell_id.clone())); - } } fn cell_id(value: &str) -> CellId { @@ -278,6 +317,61 @@ async fn yields_and_resumes() { ); } +#[tokio::test] +async fn yield_before_first_observation_preserves_its_output_boundary() { + let (delegate, mut events_rx) = BlockingDelegate::new(); + let service = CodeModeService::with_delegate(delegate.clone()); + let cell_id = service + .create_cell(CreateCellRequest { + enabled_tools: vec![blocking_tool()], + source: r#" +text("before"); +yield_control(); +text("after"); +await tools.block({}); +"# + .to_string(), + ..execute_request("") + }) + .await + .unwrap(); + assert_eq!(next_event(&mut events_rx).await, DelegateEvent::ToolStarted); + + assert_eq!( + service + .observe(ObserveRequest { + cell_id: cell_id.clone(), + yield_time_ms: 60_000, + }) + .await + .unwrap(), + CellOutcome::LiveCell(RuntimeResponse::Yielded { + cell_id: cell_id.clone(), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "before".to_string(), + }], + }) + ); + + delegate.release_tool(); + assert_eq!( + service + .observe(ObserveRequest { + cell_id: cell_id.clone(), + yield_time_ms: 60_000, + }) + .await + .unwrap(), + CellOutcome::LiveCell(RuntimeResponse::Result { + cell_id, + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "after".to_string(), + }], + error_text: None, + }) + ); +} + #[tokio::test] async fn returns_and_resumes_from_the_pending_frontier() { let (delegate, mut events_rx) = BlockingDelegate::new(); @@ -409,10 +503,6 @@ async fn termination_discards_pending_callbacks_before_responding() { content_items: Vec::new(), }) ); - assert_eq!( - next_event(&mut events_rx).await, - DelegateEvent::CellClosed(cell_id("1")) - ); } #[tokio::test] @@ -442,10 +532,6 @@ await tools.block({}); }) ); delegate.release_tool(); - assert_eq!( - next_event(&mut events_rx).await, - DelegateEvent::CellClosed(cell_id("1")) - ); assert_eq!( execute_with_yield_time( @@ -485,10 +571,6 @@ async fn shutdown_does_not_await_notifications_during_natural_completion() { .await .expect("shutdown should not await a non-cooperative notification") .unwrap(); - assert_eq!( - next_event(&mut events_rx).await, - DelegateEvent::CellClosed(cell_id("1")) - ); } #[tokio::test] @@ -509,10 +591,6 @@ async fn shutdown_does_not_await_a_non_cooperative_nested_tool() { .await .expect("shutdown should not await a non-cooperative nested tool") .unwrap(); - assert_eq!( - next_event(&mut events_rx).await, - DelegateEvent::CellClosed(cell_id("1")) - ); } #[tokio::test] @@ -550,9 +628,68 @@ async fn termination_does_not_await_a_non_cooperative_notification() { content_items: Vec::new(), }) ); +} + +#[tokio::test] +async fn create_cell_returns_before_natural_completion() { + let (delegate, mut events_rx) = ReleasableNotificationDelegate::new(); + let service = CodeModeService::with_delegate(delegate.clone()); + let created_cell_id = service + .create_cell(execute_request(r#"notify("pending");"#)) + .await + .unwrap(); + assert_eq!(created_cell_id, cell_id("1")); + let mut observation = Box::pin(service.observe(ObserveRequest { + cell_id: created_cell_id, + yield_time_ms: 60_000, + })); + assert_eq!( next_event(&mut events_rx).await, - DelegateEvent::CellClosed(cell_id("1")) + DelegateEvent::NotificationStarted + ); + std::future::poll_fn(|context| match observation.as_mut().poll(context) { + std::task::Poll::Pending => std::task::Poll::Ready(()), + std::task::Poll::Ready(result) => { + panic!("observation returned while the notification was blocked: {result:?}") + } + }) + .await; + + delegate.release_notification(); + + assert_eq!( + observation.await.unwrap(), + CellOutcome::LiveCell(RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: Vec::new(), + error_text: None, + }) + ); + assert_eq!( + next_event(&mut events_rx).await, + DelegateEvent::NotificationFinished + ); +} + +#[tokio::test] +async fn created_cell_can_be_terminated_before_observation() { + let service = CodeModeService::new(); + let created_cell_id = service + .create_cell(CreateCellRequest { + source: "await new Promise(() => {});".to_string(), + ..execute_request("") + }) + .await + .unwrap(); + + assert_eq!(created_cell_id, cell_id("1")); + assert_eq!( + service.terminate(created_cell_id.clone()).await.unwrap(), + CellOutcome::LiveCell(RuntimeResponse::Terminated { + cell_id: created_cell_id, + content_items: Vec::new(), + }) ); } diff --git a/codex-rs/code-mode/src/service_tests.rs b/codex-rs/code-mode/src/service_tests.rs index db821481cc..e0b036e8df 100644 --- a/codex-rs/code-mode/src/service_tests.rs +++ b/codex-rs/code-mode/src/service_tests.rs @@ -1,4 +1,6 @@ +use std::collections::VecDeque; use std::sync::Arc; +use std::sync::Mutex; use std::time::Duration; use super::CellId; @@ -10,8 +12,16 @@ use super::NotificationFuture; use super::ObserveRequest; use super::ObserveToPendingOutcome; use super::ObserveToPendingRequest; +use super::ProtocolDelegate; use super::RuntimeResponse; use super::ToolInvocationFuture; +use super::missing_cell_response; +use super::pending_outcome; +use super::protocol_cell_id; +use super::runtime; +use super::runtime_cell_id; +use super::runtime_request; +use super::runtime_response; use crate::CodeModeToolKind; use crate::CreateCellRequest; use crate::FunctionCallOutputContentItem; @@ -20,6 +30,7 @@ use crate::ToolDefinition; use codex_protocol::ToolName; use pretty_assertions::assert_eq; use serde_json::Value as JsonValue; +use serde_json::json; use tokio::sync::Notify; use tokio_util::sync::CancellationToken; @@ -57,8 +68,88 @@ impl CodeModeSessionDelegate for ReleasableToolDelegate { ) -> NotificationFuture<'a> { Box::pin(async { Ok(()) }) } +} - fn cell_closed(&self, _cell_id: &CellId) {} +#[derive(Debug, PartialEq)] +enum RecordedDelegateCall { + Tool { + invocation: CodeModeNestedToolCall, + cancellation_requested: bool, + }, + Notification { + call_id: String, + cell_id: CellId, + text: String, + cancellation_requested: bool, + }, +} + +struct RecordingDelegate { + calls: Mutex>, + tool_results: Mutex>>, + notification_results: Mutex>>, +} + +impl RecordingDelegate { + fn new( + tool_results: Vec>, + notification_results: Vec>, + ) -> Self { + Self { + calls: Mutex::new(Vec::new()), + tool_results: Mutex::new(tool_results.into()), + notification_results: Mutex::new(notification_results.into()), + } + } + + fn take_calls(&self) -> Vec { + std::mem::take(&mut *self.calls.lock().unwrap()) + } +} + +impl CodeModeSessionDelegate for RecordingDelegate { + fn invoke_tool<'a>( + &'a self, + invocation: CodeModeNestedToolCall, + cancellation_token: CancellationToken, + ) -> ToolInvocationFuture<'a> { + Box::pin(async move { + self.calls.lock().unwrap().push(RecordedDelegateCall::Tool { + invocation, + cancellation_requested: cancellation_token.is_cancelled(), + }); + self.tool_results + .lock() + .unwrap() + .pop_front() + .expect("test must provide one result per tool call") + }) + } + + fn notify<'a>( + &'a self, + call_id: String, + cell_id: CellId, + text: String, + cancellation_token: CancellationToken, + ) -> NotificationFuture<'a> { + Box::pin(async move { + self.calls + .lock() + .unwrap() + .push(RecordedDelegateCall::Notification { + call_id, + cell_id, + text, + cancellation_requested: cancellation_token.is_cancelled(), + }); + self.notification_results + .lock() + .unwrap() + .pop_front() + .expect("test must provide one result per notification") + }) + } } fn execute_request(source: &str) -> CreateCellRequest { @@ -990,3 +1081,310 @@ async fn observe_reports_missing_cell_separately_from_runtime_results() { }) ); } + +#[test] +fn protocol_requests_map_to_runtime_requests_field_for_field() { + let request = CreateCellRequest { + tool_call_id: "response-call-7".to_string(), + enabled_tools: vec![ + ToolDefinition { + name: "mcp__search__query".to_string(), + tool_name: ToolName::namespaced("search", "query"), + description: "Search indexed documents".to_string(), + kind: CodeModeToolKind::Function, + input_schema: Some(json!({"type": "object"})), + output_schema: None, + }, + ToolDefinition { + name: "apply_patch".to_string(), + tool_name: ToolName::plain("apply_patch"), + description: "Apply a patch".to_string(), + kind: CodeModeToolKind::Freeform, + input_schema: None, + output_schema: Some(json!({"type": "string"})), + }, + ], + source: "await tools.mcp__search__query({ q: 'rust' });".to_string(), + }; + + assert_eq!( + runtime_request(request), + runtime::CreateCellRequest { + tool_call_id: "response-call-7".to_string(), + enabled_tools: vec![ + runtime::ToolDefinition { + name: "mcp__search__query".to_string(), + tool_name: runtime::ToolName { + name: "query".to_string(), + namespace: Some("search".to_string()), + }, + description: "Search indexed documents".to_string(), + kind: runtime::ToolKind::Function, + }, + runtime::ToolDefinition { + name: "apply_patch".to_string(), + tool_name: runtime::ToolName { + name: "apply_patch".to_string(), + namespace: None, + }, + description: "Apply a patch".to_string(), + kind: runtime::ToolKind::Freeform, + }, + ], + source: "await tools.mcp__search__query({ q: 'rust' });".to_string(), + } + ); + + let protocol_id = cell_id("cell-a7"); + assert_eq!( + protocol_cell_id(&runtime_cell_id(&protocol_id)), + protocol_id + ); +} + +#[tokio::test] +async fn protocol_delegate_maps_callbacks_cancellation_and_errors_field_for_field() { + let delegate = Arc::new(RecordingDelegate::new( + vec![ + Ok(json!({"matches": 3})), + Err("freeform tool failed".to_string()), + ], + vec![Err("notification failed".to_string())], + )); + let adapter = ProtocolDelegate { + delegate: delegate.clone(), + }; + let cancelled = CancellationToken::new(); + cancelled.cancel(); + + assert_eq!( + runtime::SessionRuntimeDelegate::invoke_tool( + &adapter, + runtime::NestedToolCall { + cell_id: runtime::CellId::new("cell-a7"), + runtime_tool_call_id: "runtime-call-1".to_string(), + tool_name: runtime::ToolName { + name: "query".to_string(), + namespace: Some("search".to_string()), + }, + tool_kind: runtime::ToolKind::Function, + input: Some(json!({"q": "rust"})), + }, + CancellationToken::new(), + ) + .await, + Ok(json!({"matches": 3})) + ); + assert_eq!( + runtime::SessionRuntimeDelegate::invoke_tool( + &adapter, + runtime::NestedToolCall { + cell_id: runtime::CellId::new("cell-b9"), + runtime_tool_call_id: "runtime-call-2".to_string(), + tool_name: runtime::ToolName { + name: "apply_patch".to_string(), + namespace: None, + }, + tool_kind: runtime::ToolKind::Freeform, + input: None, + }, + cancelled.clone(), + ) + .await, + Err("freeform tool failed".to_string()) + ); + assert_eq!( + runtime::SessionRuntimeDelegate::notify( + &adapter, + "notification-1".to_string(), + runtime::CellId::new("cell-c3"), + "progress".to_string(), + cancelled, + ) + .await, + Err("notification failed".to_string()) + ); + + assert_eq!( + delegate.take_calls(), + vec![ + RecordedDelegateCall::Tool { + invocation: CodeModeNestedToolCall { + cell_id: cell_id("cell-a7"), + runtime_tool_call_id: "runtime-call-1".to_string(), + tool_name: ToolName::namespaced("search", "query"), + tool_kind: CodeModeToolKind::Function, + input: Some(json!({"q": "rust"})), + }, + cancellation_requested: false, + }, + RecordedDelegateCall::Tool { + invocation: CodeModeNestedToolCall { + cell_id: cell_id("cell-b9"), + runtime_tool_call_id: "runtime-call-2".to_string(), + tool_name: ToolName::plain("apply_patch"), + tool_kind: CodeModeToolKind::Freeform, + input: None, + }, + cancellation_requested: true, + }, + RecordedDelegateCall::Notification { + call_id: "notification-1".to_string(), + cell_id: cell_id("cell-c3"), + text: "progress".to_string(), + cancellation_requested: true, + }, + ] + ); +} + +#[test] +fn runtime_events_map_to_protocol_outcomes_field_for_field() { + let output_items = vec![ + runtime::OutputItem::Text { + text: "before".to_string(), + }, + runtime::OutputItem::Image { + image_url: "data:image/png;base64,auto".to_string(), + detail: Some(runtime::ImageDetail::Auto), + }, + runtime::OutputItem::Image { + image_url: "data:image/png;base64,low".to_string(), + detail: Some(runtime::ImageDetail::Low), + }, + runtime::OutputItem::Image { + image_url: "data:image/png;base64,high".to_string(), + detail: Some(runtime::ImageDetail::High), + }, + runtime::OutputItem::Image { + image_url: "data:image/png;base64,original".to_string(), + detail: Some(runtime::ImageDetail::Original), + }, + runtime::OutputItem::Image { + image_url: "data:image/png;base64,default".to_string(), + detail: None, + }, + ]; + let expected_output_items = vec![ + FunctionCallOutputContentItem::InputText { + text: "before".to_string(), + }, + FunctionCallOutputContentItem::InputImage { + image_url: "data:image/png;base64,auto".to_string(), + detail: Some(crate::ImageDetail::Auto), + }, + FunctionCallOutputContentItem::InputImage { + image_url: "data:image/png;base64,low".to_string(), + detail: Some(crate::ImageDetail::Low), + }, + FunctionCallOutputContentItem::InputImage { + image_url: "data:image/png;base64,high".to_string(), + detail: Some(crate::ImageDetail::High), + }, + FunctionCallOutputContentItem::InputImage { + image_url: "data:image/png;base64,original".to_string(), + detail: Some(crate::ImageDetail::Original), + }, + FunctionCallOutputContentItem::InputImage { + image_url: "data:image/png;base64,default".to_string(), + detail: None, + }, + ]; + + assert_eq!( + runtime_response( + &cell_id("cell-a7"), + runtime::CellEvent::Yielded { + content_items: output_items, + }, + ), + Ok(RuntimeResponse::Yielded { + cell_id: cell_id("cell-a7"), + content_items: expected_output_items, + }) + ); + assert_eq!( + runtime_response( + &cell_id("cell-b9"), + runtime::CellEvent::Completed { + content_items: vec![runtime::OutputItem::Text { + text: "failed".to_string(), + }], + error_text: Some("tool failed".to_string()), + }, + ), + Ok(RuntimeResponse::Result { + cell_id: cell_id("cell-b9"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "failed".to_string(), + }], + error_text: Some("tool failed".to_string()), + }) + ); + assert_eq!( + runtime_response( + &cell_id("cell-c3"), + runtime::CellEvent::Terminated { + content_items: vec![runtime::OutputItem::Text { + text: "partial".to_string(), + }], + }, + ), + Ok(RuntimeResponse::Terminated { + cell_id: cell_id("cell-c3"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "partial".to_string(), + }], + }) + ); + assert_eq!( + pending_outcome( + &cell_id("cell-d4"), + runtime::CellEvent::Pending { + content_items: vec![runtime::OutputItem::Text { + text: "waiting".to_string(), + }], + pending_tool_call_ids: vec!["runtime-call-1".to_string()], + }, + ), + Ok(PendingOutcome::Pending { + cell_id: cell_id("cell-d4"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "waiting".to_string(), + }], + pending_tool_call_ids: vec!["runtime-call-1".to_string()], + }) + ); + assert_eq!( + pending_outcome( + &cell_id("cell-e5"), + runtime::CellEvent::Completed { + content_items: Vec::new(), + error_text: None, + }, + ), + Ok(PendingOutcome::Completed(RuntimeResponse::Result { + cell_id: cell_id("cell-e5"), + content_items: Vec::new(), + error_text: None, + })) + ); + assert_eq!( + runtime_response( + &cell_id("cell-f6"), + runtime::CellEvent::Pending { + content_items: Vec::new(), + pending_tool_call_ids: Vec::new(), + }, + ), + Err("cell returned a pending frontier unexpectedly".to_string()) + ); + assert_eq!( + missing_cell_response(cell_id("missing")), + RuntimeResponse::Result { + cell_id: cell_id("missing"), + content_items: Vec::new(), + error_text: Some("exec cell missing not found".to_string()), + } + ); +} diff --git a/codex-rs/code-mode/src/session_runtime/mod.rs b/codex-rs/code-mode/src/session_runtime/mod.rs index 229d93bcb9..3bfe91db9d 100644 --- a/codex-rs/code-mode/src/session_runtime/mod.rs +++ b/codex-rs/code-mode/src/session_runtime/mod.rs @@ -8,22 +8,24 @@ use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; use serde_json::Value as JsonValue; +use sha2::Digest; +use sha2::Sha256; use tokio::sync::Mutex; use tokio_util::sync::CancellationToken; use tokio_util::task::TaskTracker; -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; +pub use self::types::CellEvent; +pub use self::types::CellId; +pub use self::types::CreateCellRequest; +pub use self::types::Error; +pub use self::types::ImageDetail; +pub use self::types::NestedToolCall; +pub use self::types::ObserveMode; +pub use self::types::OutputItem; +pub use self::types::SessionRuntimeDelegate; +pub use self::types::ToolDefinition; +pub use self::types::ToolKind; +pub use self::types::ToolName; use crate::cell_actor::CellActor; use crate::cell_actor::CellError; use crate::cell_actor::CellEventFuture; @@ -35,8 +37,11 @@ use crate::cell_actor::CompletionCommit; type RuntimeEventFuture = Pin> + Send + 'static>>; +const CELL_ID_ALPHABET: &[u8; 32] = b"0123456789abcdefghjkmnpqrstvwxyz"; +const CELL_ID_LENGTH: usize = 16; + /// Owns all cells and shared state for one transport-neutral code-mode session. -pub(crate) struct SessionRuntime { +pub struct SessionRuntime { inner: Arc>, } @@ -46,11 +51,27 @@ struct Inner { cell_tasks: TaskTracker, shutdown_token: CancellationToken, delegate: Arc, + cell_id_namespace: CellIdNamespace, next_cell_id: AtomicU64, } +enum CellIdNamespace { + Runtime(uuid::Uuid), + #[cfg(test)] + Unscoped, +} + impl SessionRuntime { - pub(crate) fn new(delegate: Arc) -> Self { + pub fn new(delegate: Arc) -> Self { + Self::with_cell_id_namespace(delegate, CellIdNamespace::Runtime(uuid::Uuid::new_v4())) + } + + #[cfg(test)] + pub(crate) fn new_for_test(delegate: Arc) -> Self { + Self::with_cell_id_namespace(delegate, CellIdNamespace::Unscoped) + } + + fn with_cell_id_namespace(delegate: Arc, cell_id_namespace: CellIdNamespace) -> Self { Self { inner: Arc::new(Inner { stored_values: Mutex::new(HashMap::new()), @@ -58,16 +79,17 @@ impl SessionRuntime { cell_tasks: TaskTracker::new(), shutdown_token: CancellationToken::new(), delegate, + cell_id_namespace, next_cell_id: AtomicU64::new(1), }), } } - pub(crate) fn is_alive(&self) -> bool { + pub fn is_alive(&self) -> bool { !self.inner.shutdown_token.is_cancelled() } - pub(crate) async fn create_cell(&self, request: CreateCellRequest) -> Result { + pub async fn create_cell(&self, request: CreateCellRequest) -> Result { if self.inner.shutdown_token.is_cancelled() { return Err(Error::ShuttingDown); } @@ -76,15 +98,11 @@ impl SessionRuntime { Ok(cell_id) } - pub(crate) async fn observe( - &self, - cell_id: &CellId, - mode: ObserveMode, - ) -> Result { + pub 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( + pub async fn begin_observe( &self, cell_id: &CellId, mode: ObserveMode, @@ -102,7 +120,7 @@ impl SessionRuntime { }) } - pub(crate) async fn terminate(&self, cell_id: &CellId) -> Result { + pub async fn terminate(&self, cell_id: &CellId) -> Result { let handle = self .inner .cells @@ -117,7 +135,7 @@ impl SessionRuntime { .map_err(|error| actor_error(cell_id, error)) } - pub(crate) async fn shutdown(&self) -> Result<(), Error> { + pub async fn shutdown(&self) -> Result<(), Error> { self.begin_shutdown(); // Taking the registry lock ensures every cell that passed the shutdown // check has registered its actor with the tracker before we wait. @@ -129,12 +147,29 @@ impl SessionRuntime { } fn allocate_cell_id(&self) -> CellId { - CellId::new( - self.inner - .next_cell_id - .fetch_add(1, Ordering::Relaxed) - .to_string(), - ) + let sequence = self.inner.next_cell_id.fetch_add(1, Ordering::Relaxed); + let value = match &self.inner.cell_id_namespace { + CellIdNamespace::Runtime(runtime_id) => { + let digest = Sha256::new() + .chain_update(runtime_id.as_bytes()) + .chain_update(sequence.to_be_bytes()) + .finalize(); + let mut encoded = String::with_capacity(CELL_ID_LENGTH); + for chunk in digest[..10].chunks_exact(5) { + let chunk_bits = chunk + .iter() + .fold(0_u64, |value, byte| (value << 8) | u64::from(*byte)); + for shift in (0..40).step_by(5).rev() { + let index = ((chunk_bits >> shift) & 0x1f) as usize; + encoded.push(char::from(CELL_ID_ALPHABET[index])); + } + } + encoded + } + #[cfg(test)] + CellIdNamespace::Unscoped => sequence.to_string(), + }; + CellId::new(value) } async fn start_cell(&self, cell_id: CellId, request: CreateCellRequest) -> Result<(), Error> { @@ -172,12 +207,12 @@ impl Drop for SessionRuntime { } /// An admitted cell event that has not reached its requested frontier yet. -pub(crate) struct PendingEvent { +pub struct PendingEvent { event: RuntimeEventFuture, } impl PendingEvent { - pub(crate) async fn event(self) -> Result { + pub async fn event(self) -> Result { self.event.await } } @@ -242,7 +277,6 @@ impl CellHost for RuntimeCellHost { async fn closed(&self) { self.inner.cells.lock().await.remove(&self.cell_id); - self.inner.delegate.cell_closed(&self.cell_id); } } diff --git a/codex-rs/code-mode/src/session_runtime/tests.rs b/codex-rs/code-mode/src/session_runtime/tests.rs index 13c540c786..c7dfde5a6b 100644 --- a/codex-rs/code-mode/src/session_runtime/tests.rs +++ b/codex-rs/code-mode/src/session_runtime/tests.rs @@ -33,8 +33,6 @@ impl SessionRuntimeDelegate for RecordingDelegate { ) -> Result<(), String> { Ok(()) } - - fn cell_closed(&self, _cell_id: &CellId) {} } #[tokio::test] @@ -161,7 +159,6 @@ async fn drop_terminates_cells_when_the_registry_is_locked() { .create_cell(execute_request("while (true) {}")) .await .unwrap(); - assert_eq!(cell_id, CellId::new("1")); assert_eq!( runtime .observe( @@ -184,3 +181,31 @@ async fn drop_terminates_cells_when_the_registry_is_locked() { .unwrap(); assert!(inner.cell_tasks.is_empty()); } + +#[tokio::test] +async fn cell_ids_are_unique_across_runtime_instances() { + let first_runtime = SessionRuntime::new(Arc::new(RecordingDelegate)); + let second_runtime = SessionRuntime::new(Arc::new(RecordingDelegate)); + let first_cell_id = first_runtime + .create_cell(execute_request("await new Promise(() => {});")) + .await + .unwrap(); + let second_cell_id = second_runtime + .create_cell(execute_request("await new Promise(() => {});")) + .await + .unwrap(); + + assert_ne!(first_cell_id, second_cell_id); + for cell_id in [&first_cell_id, &second_cell_id] { + assert_eq!(cell_id.as_str().len(), CELL_ID_LENGTH); + assert!( + cell_id + .as_str() + .bytes() + .all(|byte| CELL_ID_ALPHABET.contains(&byte)) + ); + } + + first_runtime.shutdown().await.unwrap(); + second_runtime.shutdown().await.unwrap(); +} diff --git a/codex-rs/code-mode/src/session_runtime/types.rs b/codex-rs/code-mode/src/session_runtime/types.rs index b04bd0f236..3498b6fc52 100644 --- a/codex-rs/code-mode/src/session_runtime/types.rs +++ b/codex-rs/code-mode/src/session_runtime/types.rs @@ -7,14 +7,14 @@ use tokio_util::sync::CancellationToken; /// Identifies one execution cell within a session runtime. #[derive(Clone, Debug, Eq, Hash, PartialEq)] -pub(crate) struct CellId(String); +pub struct CellId(String); impl CellId { - pub(crate) fn new(value: impl Into) -> Self { + pub fn new(value: impl Into) -> Self { Self(value.into()) } - pub(crate) fn as_str(&self) -> &str { + pub fn as_str(&self) -> &str { &self.0 } } @@ -27,14 +27,14 @@ impl fmt::Display for CellId { /// Selects the next observable frontier for a running cell. #[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub(crate) enum ObserveMode { +pub enum ObserveMode { YieldAfter(Duration), PendingFrontier, } /// An observable cell lifecycle event. #[derive(Clone, Debug, PartialEq)] -pub(crate) enum CellEvent { +pub enum CellEvent { Yielded { content_items: Vec, }, @@ -53,7 +53,7 @@ pub(crate) enum CellEvent { /// Output emitted by a cell since its preceding observation. #[derive(Clone, Debug, PartialEq)] -pub(crate) enum OutputItem { +pub enum OutputItem { Text { text: String, }, @@ -65,7 +65,7 @@ pub(crate) enum OutputItem { /// Requested image fidelity for an output image. #[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub(crate) enum ImageDetail { +pub enum ImageDetail { Auto, Low, High, @@ -75,48 +75,51 @@ pub(crate) enum ImageDetail { /// Transport-neutral input for creating a cell. /// /// The owning session assigns the cell ID when it admits the request. -pub(crate) struct CreateCellRequest { - pub(crate) tool_call_id: String, - pub(crate) enabled_tools: Vec, - pub(crate) source: String, +#[derive(Debug, PartialEq)] +pub struct CreateCellRequest { + pub tool_call_id: String, + pub enabled_tools: Vec, + pub source: String, } /// Tool metadata exposed to code running inside a cell. -pub(crate) struct ToolDefinition { - pub(crate) name: String, - pub(crate) tool_name: ToolName, - pub(crate) description: String, - pub(crate) kind: ToolKind, +#[derive(Debug, PartialEq)] +pub struct ToolDefinition { + pub name: String, + pub tool_name: ToolName, + pub description: String, + pub kind: ToolKind, } /// A tool name with an optional namespace. #[derive(Clone, Debug, Eq, PartialEq)] -pub(crate) struct ToolName { - pub(crate) name: String, - pub(crate) namespace: Option, +pub struct ToolName { + pub name: String, + pub namespace: Option, } /// The JavaScript calling convention for a tool. #[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub(crate) enum ToolKind { +pub 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, +#[derive(Debug, PartialEq)] +pub struct NestedToolCall { + pub cell_id: CellId, + pub runtime_tool_call_id: String, + pub tool_name: ToolName, + pub tool_kind: ToolKind, + pub 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 { +/// Implementations should forward callback cancellation tokens to downstream +/// work. The runtime stops awaiting callbacks once cancellation begins. +pub trait SessionRuntimeDelegate: Send + Sync + 'static { fn invoke_tool( &self, invocation: NestedToolCall, @@ -130,13 +133,11 @@ pub(crate) trait SessionRuntimeDelegate: Send + Sync + 'static { 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 { +pub enum Error { ShuttingDown, DuplicateCell(CellId), MissingCell(CellId), diff --git a/codex-rs/core/src/tools/code_mode/delegate.rs b/codex-rs/core/src/tools/code_mode/delegate.rs index 7dd6615095..fa6fb89300 100644 --- a/codex-rs/core/src/tools/code_mode/delegate.rs +++ b/codex-rs/core/src/tools/code_mode/delegate.rs @@ -238,10 +238,6 @@ impl CodeModeSessionDelegate for CodeModeDispatchBroker { } }) } - - fn cell_closed(&self, cell_id: &CellId) { - self.close_cell(cell_id); - } } enum DispatchMessage { diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 584bb7d588..a1fd5c7162 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -1412,7 +1412,7 @@ text("phase 3"); assert_regex_match( concat!( r"(?s)\A", - r"Script running with cell ID \d+\nWall time \d+\.\d seconds\nOutput:\n\z" + r"Script running with cell ID [0-9abcdefghjkmnpqrstvwxyz]{16}\nWall time \d+\.\d seconds\nOutput:\n\z" ), text_item(&first_items, /*index*/ 0), ); @@ -1453,7 +1453,7 @@ text("phase 3"); assert_regex_match( concat!( r"(?s)\A", - r"Script running with cell ID \d+\nWall time \d+\.\d seconds\nOutput:\n\z" + r"Script running with cell ID [0-9abcdefghjkmnpqrstvwxyz]{16}\nWall time \d+\.\d seconds\nOutput:\n\z" ), text_item(&second_items, /*index*/ 0), ); @@ -1556,7 +1556,7 @@ while (true) {} assert_regex_match( concat!( r"(?s)\A", - r"Script running with cell ID \d+\nWall time \d+\.\d seconds\nOutput:\n\z" + r"Script running with cell ID [0-9abcdefghjkmnpqrstvwxyz]{16}\nWall time \d+\.\d seconds\nOutput:\n\z" ), text_item(&first_items, /*index*/ 0), ); @@ -2235,7 +2235,7 @@ text("session b done"); assert_regex_match( concat!( r"(?s)\A", - r"Script running with cell ID \d+\nWall time \d+\.\d seconds\nOutput:\n\z" + r"Script running with cell ID [0-9abcdefghjkmnpqrstvwxyz]{16}\nWall time \d+\.\d seconds\nOutput:\n\z" ), text_item(&third_items, /*index*/ 0), ); @@ -2357,7 +2357,7 @@ text("after yield"); assert_regex_match( concat!( r"(?s)\A", - r"Script running with cell ID \d+\nWall time \d+\.\d seconds\nOutput:\n\z" + r"Script running with cell ID [0-9abcdefghjkmnpqrstvwxyz]{16}\nWall time \d+\.\d seconds\nOutput:\n\z" ), text_item(&first_items, /*index*/ 0), );