From 3305c4f31d4ab6542afa05b91e112f1d67e9744f Mon Sep 17 00:00:00 2001 From: Abhinav Date: Fri, 11 Sep 2026 16:49:41 +0000 Subject: [PATCH] Scope code mode callback delegates to individual executions (#44865) ## Why A delegate bound at session creation cannot provide different callbacks for cells sharing that session. Each execution needs to retain its own delegate across yields and release it when the cell is cleaned up. ## What changed - Pass `CodeModeSessionDelegate` to `execute` instead of session creation. - Route tool calls, notifications, and cell closure callbacks through the execution's delegate in the in-process runtime and the gRPC and stdio transports. - Retain delegates with pending executions and live cells, releasing them through closure and cancellation cleanup. ## Testing Add coverage for distinct delegates across yielded cells, gRPC callbacks before cell admission, and delegate release after completion or abandoned execution cleanup. Update transport tests to verify callbacks reach the owning cell's delegate. GitOrigin-RevId: 7469f52104b993e790a40c65fb800ad02b7fd606 --- codex-rs/code-mode-host/src/grpc/mod.rs | 2 +- codex-rs/code-mode-host/src/grpc/session.rs | 13 +- codex-rs/code-mode-host/src/lib.rs | 21 +- codex-rs/code-mode-host/tests/grpc.rs | 206 +++++++++++----- .../tests/grpc_notifications.rs | 37 ++- codex-rs/code-mode-host/tests/grpc_timing.rs | 42 ++-- codex-rs/code-mode-host/tests/stdio.rs | 147 +++++++---- codex-rs/code-mode-protocol/src/session.rs | 14 +- codex-rs/code-mode-runtime/src/service.rs | 36 ++- .../src/service_contract_tests.rs | 165 ++++++++++--- .../code-mode-runtime/src/service_tests.rs | 223 +++++++++-------- .../src/session_runtime/mod.rs | 37 ++- .../src/session_runtime/tests.rs | 23 +- codex-rs/code-mode-runtime/tests/jit.rs | 19 +- .../code-mode/src/grpc_session/callbacks.rs | 26 +- codex-rs/code-mode/src/grpc_session/mod.rs | 27 +- .../code-mode/src/grpc_session/operations.rs | 9 +- .../code-mode/src/grpc_session/reconnect.rs | 16 +- codex-rs/code-mode/src/grpc_session/state.rs | 49 +++- .../code-mode/src/grpc_session/state_tests.rs | 145 +++++++---- codex-rs/code-mode/src/remote_session.rs | 53 ++-- .../src/remote_session/connection.rs | 4 +- .../connection/driver/commands.rs | 17 +- .../connection/driver/responses.rs | 7 +- .../connection/driver/session_registry.rs | 60 ++--- .../remote_session/connection/driver/types.rs | 4 +- .../remote_session/connection/driver_tests.rs | 230 +++++++++++++----- .../code-mode/src/remote_session_tests.rs | 19 +- codex-rs/core/src/tools/code_mode/mod.rs | 7 +- .../src/tools/tool_dispatch_trace_tests.rs | 6 +- 30 files changed, 1027 insertions(+), 637 deletions(-) diff --git a/codex-rs/code-mode-host/src/grpc/mod.rs b/codex-rs/code-mode-host/src/grpc/mod.rs index 88c7dbadff..513eb3aaab 100644 --- a/codex-rs/code-mode-host/src/grpc/mod.rs +++ b/codex-rs/code-mode-host/src/grpc/mod.rs @@ -148,7 +148,7 @@ impl GrpcCodeModeHost { _ = session.closed.cancelled() => { return Err(Status::cancelled("code-mode session is closed")); } - result = session.runtime.execute(request) => { + result = session.runtime.execute(request, Arc::new(delegate::GrpcDelegate::new(Arc::downgrade(&session)))) => { result.map_err(Status::failed_precondition)? } }; diff --git a/codex-rs/code-mode-host/src/grpc/session.rs b/codex-rs/code-mode-host/src/grpc/session.rs index a8ce05bf6f..d576e6700d 100644 --- a/codex-rs/code-mode-host/src/grpc/session.rs +++ b/codex-rs/code-mode-host/src/grpc/session.rs @@ -27,7 +27,6 @@ use tonic::Status; use uuid::Uuid; use super::GrpcStream; -use super::delegate::GrpcDelegate; use super::events::EventSender; use super::validation; use super::waits::ActiveWait; @@ -213,7 +212,6 @@ impl GrpcSession { limits: CodeModeSessionCellExecutionLimits, ) -> Arc { Arc::new_cyclic(|weak: &Weak| { - let delegate = Arc::new(GrpcDelegate::new(weak.clone())); let failure_session = weak.clone(); let failure_handler = Arc::new(move |reason: String| { if let Some(session) = failure_session.upgrade() { @@ -223,13 +221,10 @@ impl GrpcSession { }); Self { id, - runtime: Arc::new( - InProcessCodeModeSession::with_delegate_and_task_failure_handler( - delegate, - failure_handler, - limits, - ), - ), + runtime: Arc::new(InProcessCodeModeSession::with_task_failure_handler( + failure_handler, + limits, + )), closed, state: Mutex::new(SessionState::default()), events, diff --git a/codex-rs/code-mode-host/src/lib.rs b/codex-rs/code-mode-host/src/lib.rs index b147983415..f128fa1513 100644 --- a/codex-rs/code-mode-host/src/lib.rs +++ b/codex-rs/code-mode-host/src/lib.rs @@ -412,7 +412,11 @@ impl HostState { ); return; }; - let result = session.execute(request).await; + let delegate = Arc::new(RemoteDelegate::new( + session_id.clone(), + Arc::clone(&self.peer), + )); + let result = session.execute(request, delegate).await; match result { Ok(started) => { let cell_id = started.cell_id.clone(); @@ -517,10 +521,6 @@ impl HostState { { return Err(format!("code-mode session ID `{session_id}` was reused")); } - let delegate = Arc::new(RemoteDelegate::new( - session_id.clone(), - Arc::clone(&self.peer), - )); let peer = Arc::downgrade(&self.peer); let task_failure_handler = Arc::new(move |reason| { if let Some(peer) = peer.upgrade() { @@ -529,13 +529,10 @@ impl HostState { }); sessions.insert( session_id, - Arc::new( - InProcessCodeModeSession::with_delegate_and_task_failure_handler( - delegate, - task_failure_handler, - cell_execution_limits, - ), - ), + Arc::new(InProcessCodeModeSession::with_task_failure_handler( + task_failure_handler, + cell_execution_limits, + )), ); Ok(()) } diff --git a/codex-rs/code-mode-host/tests/grpc.rs b/codex-rs/code-mode-host/tests/grpc.rs index a1f0963d75..0e204ff989 100644 --- a/codex-rs/code-mode-host/tests/grpc.rs +++ b/codex-rs/code-mode-host/tests/grpc.rs @@ -144,10 +144,11 @@ fn text_response( async fn execute( session: &Arc, request: ExecuteRequest, + delegate: Arc, ) -> Result { timeout(TEST_TIMEOUT, async { session - .execute(request) + .execute(request, delegate.clone()) .await .map_err(anyhow::Error::msg)? .initial_response() @@ -190,7 +191,7 @@ async fn grpc_endpoints_reject_credentials_without_disclosing_them() { ] { let provider = GrpcCodeModeSessionProvider::new(endpoint); let error = provider - .create_session(Arc::new(NoopCodeModeSessionDelegate)) + .create_session() .await .err() .expect("gRPC credentials should be rejected"); @@ -207,12 +208,18 @@ async fn tcp_session_persists_values_and_forwards_tools_notifications_and_closur assert!(host.endpoint.starts_with("http://127.0.0.1:")); let provider = GrpcCodeModeSessionProvider::new(host.endpoint); let delegate = Arc::new(RecordingDelegate::default()); + let first_delegate = Arc::new(RecordingDelegate::default()); let session = provider - .create_session(delegate.clone()) + .create_session() .await .map_err(anyhow::Error::msg)?; - let actual = execute(&session, request(r#"store("key", "persisted");"#)).await?; + let actual = execute( + &session, + request(r#"store("key", "persisted");"#), + first_delegate.clone(), + ) + .await?; assert_eq!( actual, RuntimeResponse::Result { @@ -228,7 +235,7 @@ async fn tcp_session_persists_values_and_forwards_tools_notifications_and_closur ); callback.tool_call_id = "call-2".to_string(); callback.enabled_tools = vec![tool("echo")]; - let actual = execute(&session, callback).await?; + let actual = execute(&session, callback, delegate.clone()).await?; assert_eq!( actual, text_response("2", "output", actual.code_mode_host_duration()) @@ -263,7 +270,14 @@ async fn tcp_session_persists_values_and_forwards_tools_notifications_and_closur .closed_cells .lock() .unwrap_or_else(PoisonError::into_inner), - vec![cell_id("1"), cell_id("2")] + vec![cell_id("2")] + ); + assert_eq!( + *first_delegate + .closed_cells + .lock() + .unwrap_or_else(PoisonError::into_inner), + vec![cell_id("1")] ); Ok(()) } @@ -273,14 +287,20 @@ async fn shutdown_immediately_rejects_new_operations() -> Result<()> { let host = HostHarness::start("grpc://127.0.0.1:0").await?; let provider = GrpcCodeModeSessionProvider::new(host.endpoint); let session = provider - .create_session(Arc::new(NoopCodeModeSessionDelegate)) + .create_session() .await .map_err(anyhow::Error::msg)?; let shutdown = session.shutdown(); let expected = "code mode session is shutting down".to_string(); assert_eq!( - session.execute(request("text('too late');")).await.err(), + session + .execute( + request("text('too late');"), + Arc::new(NoopCodeModeSessionDelegate) + ) + .await + .err(), Some(expected.clone()) ); assert_eq!( @@ -304,13 +324,18 @@ async fn cancelling_execution_before_admission_keeps_the_session_usable() -> Res let provider = GrpcCodeModeSessionProvider::new(host.endpoint); let delegate = Arc::new(RecordingDelegate::default()); let session = provider - .create_session(delegate.clone()) + .create_session() .await .map_err(anyhow::Error::msg)?; let mut pending = request("await new Promise(() => {});"); pending.yield_time_ms = Some(/*value*/ 1); - assert!(session.execute(pending).now_or_never().is_none()); + assert!( + session + .execute(pending, delegate.clone()) + .now_or_never() + .is_none() + ); let abandoned_cell = cell_id("1"); timeout(TEST_TIMEOUT, async { @@ -354,7 +379,12 @@ async fn cancelling_execution_before_admission_keeps_the_session_usable() -> Res .await .context("cancelled execution leaked its remote cell")??; - let actual = execute(&session, request(r#"text("still alive");"#)).await?; + let actual = execute( + &session, + request(r#"text("still alive");"#), + delegate.clone(), + ) + .await?; assert_eq!( actual, text_response("2", "still alive", actual.code_mode_host_duration()) @@ -370,12 +400,15 @@ async fn dropping_a_started_cell_off_runtime_terminates_its_buffered_remote_exec let provider = GrpcCodeModeSessionProvider::new(host.endpoint); let delegate = Arc::new(RecordingDelegate::default()); let session = provider - .create_session(delegate.clone()) + .create_session() .await .map_err(anyhow::Error::msg)?; let mut pending = request("await new Promise(() => {});"); pending.yield_time_ms = Some(/*value*/ 1); - let started = session.execute(pending).await.map_err(anyhow::Error::msg)?; + let started = session + .execute(pending, delegate.clone()) + .await + .map_err(anyhow::Error::msg)?; let abandoned_cell = started.cell_id.clone(); timeout(TEST_TIMEOUT, async { @@ -436,7 +469,12 @@ async fn dropping_a_started_cell_off_runtime_terminates_its_buffered_remote_exec .contains(&abandoned_cell) ); - let actual = execute(&session, request(r#"text("still alive");"#)).await?; + let actual = execute( + &session, + request(r#"text("still alive");"#), + delegate.clone(), + ) + .await?; assert_eq!( actual, text_response("2", "still alive", actual.code_mode_host_duration()) @@ -451,12 +489,15 @@ async fn dropping_an_initial_response_terminates_its_pending_remote_execution() let provider = GrpcCodeModeSessionProvider::new(host.endpoint); let delegate = Arc::new(RecordingDelegate::default()); let session = provider - .create_session(delegate.clone()) + .create_session() .await .map_err(anyhow::Error::msg)?; let mut pending = request("await new Promise(() => {});"); pending.yield_time_ms = Some(/*value*/ 60_000); - let started = session.execute(pending).await.map_err(anyhow::Error::msg)?; + let started = session + .execute(pending, delegate.clone()) + .await + .map_err(anyhow::Error::msg)?; let abandoned_cell = started.cell_id.clone(); let initial_response = tokio::spawn(started.initial_response()); tokio::task::yield_now().await; @@ -479,7 +520,12 @@ async fn dropping_an_initial_response_terminates_its_pending_remote_execution() .await .context("dropping an initial response did not terminate its pending remote execution")?; - let actual = execute(&session, request(r#"text("still alive");"#)).await?; + let actual = execute( + &session, + request(r#"text("still alive");"#), + delegate.clone(), + ) + .await?; assert_eq!( actual, text_response("2", "still alive", actual.code_mode_host_duration()) @@ -493,7 +539,7 @@ async fn synchronous_delegate_panics_do_not_orphan_callbacks_or_close_the_sessio let host = HostHarness::start("grpc://127.0.0.1:0").await?; let provider = GrpcCodeModeSessionProvider::new(host.endpoint); let session = provider - .create_session(Arc::new(PanickingDelegate)) + .create_session() .await .map_err(anyhow::Error::msg)?; @@ -501,7 +547,7 @@ async fn synchronous_delegate_panics_do_not_orphan_callbacks_or_close_the_sessio r#"try { await tools.echo({}); text("unexpected"); } catch (_) { text("tool recovered"); }"#, ); tool_panic.enabled_tools = vec![tool("echo")]; - let actual = execute(&session, tool_panic).await?; + let actual = execute(&session, tool_panic, Arc::new(PanickingDelegate)).await?; assert_eq!( actual, text_response("1", "tool recovered", actual.code_mode_host_duration()) @@ -510,6 +556,7 @@ async fn synchronous_delegate_panics_do_not_orphan_callbacks_or_close_the_sessio let actual = execute( &session, request(r#"notify("panic"); text("notification recovered");"#), + Arc::new(PanickingDelegate), ) .await?; assert_eq!( @@ -520,7 +567,12 @@ async fn synchronous_delegate_panics_do_not_orphan_callbacks_or_close_the_sessio actual.code_mode_host_duration() ) ); - let actual = execute(&session, request(r#"text("still alive");"#)).await?; + let actual = execute( + &session, + request(r#"text("still alive");"#), + Arc::new(PanickingDelegate), + ) + .await?; assert_eq!( actual, text_response("3", "still alive", actual.code_mode_host_duration()) @@ -535,19 +587,24 @@ async fn tool_delegate_self_cancellation_returns_an_error_without_hanging() -> R let host = HostHarness::start("grpc://127.0.0.1:0").await?; let provider = GrpcCodeModeSessionProvider::new(host.endpoint); let session = provider - .create_session(Arc::new(SelfCancellingToolDelegate)) + .create_session() .await .map_err(anyhow::Error::msg)?; let mut callback = request(r#"try { await tools.echo({}); } catch (_) { text("tool recovered"); }"#); callback.enabled_tools = vec![tool("echo")]; - let actual = execute(&session, callback).await?; + let actual = execute(&session, callback, Arc::new(SelfCancellingToolDelegate)).await?; assert_eq!( actual, text_response("1", "tool recovered", actual.code_mode_host_duration()) ); - let actual = execute(&session, request(r#"text("still alive");"#)).await?; + let actual = execute( + &session, + request(r#"text("still alive");"#), + Arc::new(SelfCancellingToolDelegate), + ) + .await?; assert_eq!( actual, text_response("2", "still alive", actual.code_mode_host_duration()) @@ -562,12 +619,15 @@ async fn concurrent_wait_rejects_without_displacing_the_active_observer() -> Res let host = HostHarness::start("grpc://127.0.0.1:0").await?; let provider = GrpcCodeModeSessionProvider::new(host.endpoint); let session = provider - .create_session(Arc::new(NoopCodeModeSessionDelegate)) + .create_session() .await .map_err(anyhow::Error::msg)?; let mut pending = request("await new Promise(() => {});"); pending.yield_time_ms = Some(/*value*/ 1); - let started = session.execute(pending).await.map_err(anyhow::Error::msg)?; + let started = session + .execute(pending, Arc::new(NoopCodeModeSessionDelegate)) + .await + .map_err(anyhow::Error::msg)?; let running_cell = started.cell_id.clone(); let actual = started .initial_response() @@ -639,12 +699,15 @@ async fn dropping_a_wait_retires_its_observer_before_the_next_wait() -> Result<( let host = HostHarness::start("grpc://127.0.0.1:0").await?; let provider = GrpcCodeModeSessionProvider::new(host.endpoint); let session = provider - .create_session(Arc::new(NoopCodeModeSessionDelegate)) + .create_session() .await .map_err(anyhow::Error::msg)?; let mut pending = request("await new Promise(() => {});"); pending.yield_time_ms = Some(/*value*/ 1); - let started = session.execute(pending).await.map_err(anyhow::Error::msg)?; + let started = session + .execute(pending, Arc::new(NoopCodeModeSessionDelegate)) + .await + .map_err(anyhow::Error::msg)?; let running_cell = started.cell_id.clone(); let actual = started .initial_response() @@ -710,12 +773,12 @@ async fn dropping_a_session_off_runtime_retires_its_active_cells() -> Result<()> let provider = GrpcCodeModeSessionProvider::new(host.endpoint); let delegate = Arc::new(RecordingDelegate::default()); let session = provider - .create_session(delegate.clone()) + .create_session() .await .map_err(anyhow::Error::msg)?; let mut pending = request("await new Promise(() => {});"); pending.yield_time_ms = Some(/*value*/ 1); - let actual = execute(&session, pending).await?; + let actual = execute(&session, pending, delegate.clone()).await?; assert_eq!( actual, RuntimeResponse::Yielded { @@ -757,11 +820,11 @@ async fn large_unary_tool_completion_does_not_block_an_independent_session() -> release: Semaphore::new(/*permits*/ 0), }); let slow_session = provider - .create_session(delegate.clone()) + .create_session() .await .map_err(anyhow::Error::msg)?; let fast_session = provider - .create_session(Arc::new(NoopCodeModeSessionDelegate)) + .create_session() .await .map_err(anyhow::Error::msg)?; let mut slow_request = request( @@ -770,7 +833,7 @@ async fn large_unary_tool_completion_does_not_block_an_independent_session() -> slow_request.enabled_tools = vec![tool("large")]; slow_request.yield_time_ms = Some(/*value*/ 20_000); let slow_cell = slow_session - .execute(slow_request) + .execute(slow_request, delegate.clone()) .await .map_err(anyhow::Error::msg)?; timeout(TEST_TIMEOUT, delegate.started.acquire()) @@ -778,7 +841,12 @@ async fn large_unary_tool_completion_does_not_block_an_independent_session() -> .context("large tool callback did not start")?? .forget(); - let actual = execute(&fast_session, request(r#"text("fast-before");"#)).await?; + let actual = execute( + &fast_session, + request(r#"text("fast-before");"#), + Arc::new(NoopCodeModeSessionDelegate), + ) + .await?; assert_eq!( actual, text_response("1", "fast-before", actual.code_mode_host_duration()) @@ -786,7 +854,11 @@ async fn large_unary_tool_completion_does_not_block_an_independent_session() -> delegate.release.add_permits(/*n*/ 1); let slow_response = timeout(TEST_TIMEOUT, slow_cell.initial_response()); - let fast_response = execute(&fast_session, request(r#"text("fast-during");"#)); + let fast_response = execute( + &fast_session, + request(r#"text("fast-during");"#), + Arc::new(NoopCodeModeSessionDelegate), + ); let (slow_response, fast_response) = tokio::join!(slow_response, fast_response); let actual = fast_response?; @@ -822,7 +894,7 @@ async fn single_subscription_processes_slow_and_fast_tools_concurrently() -> Res release: Semaphore::new(/*permits*/ 0), }); let session = provider - .create_session(delegate.clone()) + .create_session() .await .map_err(anyhow::Error::msg)?; @@ -830,7 +902,10 @@ async fn single_subscription_processes_slow_and_fast_tools_concurrently() -> Res request(r#"const result = await tools.large({}); text(String(result.value.length));"#); slow.enabled_tools = vec![tool("large")]; slow.yield_time_ms = Some(/*value*/ 20_000); - let slow_cell = session.execute(slow).await.map_err(anyhow::Error::msg)?; + let slow_cell = session + .execute(slow, delegate.clone()) + .await + .map_err(anyhow::Error::msg)?; timeout(TEST_TIMEOUT, delegate.started.acquire()) .await .context("slow tool did not start")?? @@ -838,7 +913,7 @@ async fn single_subscription_processes_slow_and_fast_tools_concurrently() -> Res let mut fast = request(r#"const result = await tools.fast({}); text(result.value);"#); fast.enabled_tools = vec![tool("fast")]; - let actual = execute(&session, fast).await?; + let actual = execute(&session, fast, delegate.clone()).await?; assert_eq!( actual, text_response("2", "isolated", actual.code_mode_host_duration()) @@ -869,29 +944,23 @@ async fn sessions_enforce_independent_yield_limits() -> Result<()> { let host = HostHarness::start("grpc://127.0.0.1:0").await?; let provider = GrpcCodeModeSessionProvider::new(host.endpoint); let limited = provider - .create_session_with_limits( - Arc::new(NoopCodeModeSessionDelegate), - CodeModeSessionCellExecutionLimits { - max_yield_time_ms: Some(/*value*/ 1), - max_heap_size_bytes: Some(/*value*/ 16 * 1024 * 1024), - }, - ) + .create_session_with_limits(CodeModeSessionCellExecutionLimits { + max_yield_time_ms: Some(/*value*/ 1), + max_heap_size_bytes: Some(/*value*/ 16 * 1024 * 1024), + }) .await .map_err(anyhow::Error::msg)?; let other = provider - .create_session_with_limits( - Arc::new(NoopCodeModeSessionDelegate), - CodeModeSessionCellExecutionLimits { - max_yield_time_ms: Some(/*value*/ 1_000), - max_heap_size_bytes: None, - }, - ) + .create_session_with_limits(CodeModeSessionCellExecutionLimits { + max_yield_time_ms: Some(/*value*/ 1_000), + max_heap_size_bytes: None, + }) .await .map_err(anyhow::Error::msg)?; let mut pending = request("await new Promise(() => {});"); pending.yield_time_ms = Some(/*value*/ 60_000); - let actual = execute(&limited, pending).await?; + let actual = execute(&limited, pending, Arc::new(NoopCodeModeSessionDelegate)).await?; assert_eq!( actual, RuntimeResponse::Yielded { @@ -921,6 +990,7 @@ async fn sessions_enforce_independent_yield_limits() -> Result<()> { let actual = execute( &other, request(r#"await new Promise(resolve => setTimeout(resolve, 25)); text("isolated");"#), + Arc::new(NoopCodeModeSessionDelegate), ) .await?; assert_eq!( @@ -939,7 +1009,12 @@ async fn sessions_enforce_independent_yield_limits() -> Result<()> { content_items: Vec::new(), }) ); - let actual = execute(&limited, request("await new Promise(() => {});")).await?; + let actual = execute( + &limited, + request("await new Promise(() => {});"), + Arc::new(NoopCodeModeSessionDelegate), + ) + .await?; assert_eq!( actual, RuntimeResponse::Yielded { @@ -1008,14 +1083,17 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> { let provider = GrpcCodeModeSessionProvider::new(original.endpoint.clone()); let delegate = Arc::new(RecordingDelegate::default()); let session = provider - .create_session(delegate.clone()) + .create_session() .await .map_err(anyhow::Error::msg)?; let mut pending = request("await tools.echo({generation: 1}); await new Promise(() => {});"); pending.enabled_tools = vec![tool("echo")]; pending.yield_time_ms = Some(/*value*/ 1); - let started = session.execute(pending).await.map_err(anyhow::Error::msg)?; + let started = session + .execute(pending, delegate.clone()) + .await + .map_err(anyhow::Error::msg)?; let old_cell_id = started.cell_id.clone(); assert_eq!(old_cell_id, cell_id("1")); assert!(matches!( @@ -1078,8 +1156,8 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> { callback.tool_call_id = "reconnected-call".to_string(); callback.enabled_tools = vec![tool("echo")]; let (callback_response, concurrent_response) = tokio::join!( - execute(&session, callback), - execute(&session, request(r#"text("concurrent")"#)), + execute(&session, callback, delegate.clone()), + execute(&session, request(r#"text("concurrent")"#), delegate.clone()), ); let callback_response = callback_response?; let RuntimeResponse::Result { @@ -1145,7 +1223,10 @@ async fn cached_session_recovers_after_a_remote_host_restarts() -> Result<()> { let mut pending = request("await new Promise(() => {});"); pending.yield_time_ms = Some(/*value*/ 1); - let started = session.execute(pending).await.map_err(anyhow::Error::msg)?; + let started = session + .execute(pending, delegate.clone()) + .await + .map_err(anyhow::Error::msg)?; let replacement_cell_id = started.cell_id.clone(); assert_eq!(replacement_cell_id, cell_id("g2:3")); let actual = started @@ -1219,10 +1300,15 @@ async fn unix_socket_endpoints_execute_code_mode_cells() -> Result<()> { format!("unix:{}", socket_path.display()), ] { let session = GrpcCodeModeSessionProvider::new(endpoint) - .create_session(Arc::new(NoopCodeModeSessionDelegate)) + .create_session() .await .map_err(anyhow::Error::msg)?; - let actual = execute(&session, request(r#"text("unix socket")"#)).await?; + let actual = execute( + &session, + request(r#"text("unix socket")"#), + Arc::new(NoopCodeModeSessionDelegate), + ) + .await?; assert_eq!( actual, text_response("1", "unix socket", actual.code_mode_host_duration()) diff --git a/codex-rs/code-mode-host/tests/grpc_notifications.rs b/codex-rs/code-mode-host/tests/grpc_notifications.rs index dedcde7028..29533cc5b9 100644 --- a/codex-rs/code-mode-host/tests/grpc_notifications.rs +++ b/codex-rs/code-mode-host/tests/grpc_notifications.rs @@ -120,10 +120,11 @@ fn text_response( async fn execute( session: &Arc, request: ExecuteRequest, + delegate: Arc, ) -> Result { timeout(TEST_TIMEOUT, async { session - .execute(request) + .execute(request, delegate.clone()) .await .map_err(anyhow::Error::msg)? .initial_response() @@ -140,13 +141,19 @@ async fn completed_cells_drain_pending_notifications_before_completion() -> Resu let provider = GrpcCodeModeSessionProvider::new(host.endpoint); let delegate = Arc::new(BlockingNotificationDelegate::new()); let session = provider - .create_session(delegate.clone()) + .create_session() .await .map_err(anyhow::Error::msg)?; let executing = Arc::clone(&session); + let execution_delegate = delegate.clone(); let completion = tokio::spawn(async move { - execute(&executing, request(r#"notify("notice"); text("done");"#)).await + execute( + &executing, + request(r#"notify("notice"); text("done");"#), + execution_delegate, + ) + .await }); timeout(TEST_TIMEOUT, delegate.started.acquire()) .await @@ -171,7 +178,12 @@ async fn completed_cells_drain_pending_notifications_before_completion() -> Resu .context("completed cell was not retired")?? .forget(); - let actual = execute(&session, request(r#"text("still alive");"#)).await?; + let actual = execute( + &session, + request(r#"text("still alive");"#), + delegate.clone(), + ) + .await?; assert_eq!( actual, text_response("2", "still alive", actual.code_mode_host_duration()) @@ -187,11 +199,14 @@ async fn completed_waits_drain_pending_notifications_before_returning() -> Resul let provider = GrpcCodeModeSessionProvider::new(host.endpoint); let delegate = Arc::new(BlockingNotificationDelegate::new()); let session = provider - .create_session(delegate.clone()) + .create_session() .await .map_err(anyhow::Error::msg)?; let pending = request(r#"yield_control(); notify("notice"); text("done");"#); - let cell = session.execute(pending).await.map_err(anyhow::Error::msg)?; + let cell = session + .execute(pending, delegate.clone()) + .await + .map_err(anyhow::Error::msg)?; let actual = cell.initial_response().await.map_err(anyhow::Error::msg)?; assert_eq!( actual, @@ -248,12 +263,15 @@ async fn termination_cancels_pending_notifications() -> Result<()> { let provider = GrpcCodeModeSessionProvider::new(host.endpoint); let delegate = Arc::new(BlockingNotificationDelegate::new()); let session = provider - .create_session(delegate.clone()) + .create_session() .await .map_err(anyhow::Error::msg)?; let mut pending = request(r#"notify("notice"); await new Promise(() => {});"#); pending.yield_time_ms = Some(/*value*/ 1); - let cell = session.execute(pending).await.map_err(anyhow::Error::msg)?; + let cell = session + .execute(pending, delegate.clone()) + .await + .map_err(anyhow::Error::msg)?; timeout(TEST_TIMEOUT, delegate.started.acquire()) .await @@ -299,13 +317,14 @@ async fn oversized_notification_text_is_delivered_unchanged() -> Result<()> { let provider = GrpcCodeModeSessionProvider::new(host.endpoint); let delegate = Arc::new(RecordingDelegate::default()); let session = provider - .create_session(delegate.clone()) + .create_session() .await .map_err(anyhow::Error::msg)?; let actual = execute( &session, request(r#"notify("🦀".repeat(512)); text("done");"#), + delegate.clone(), ) .await?; assert_eq!( diff --git a/codex-rs/code-mode-host/tests/grpc_timing.rs b/codex-rs/code-mode-host/tests/grpc_timing.rs index 677e86ee37..980e2ee697 100644 --- a/codex-rs/code-mode-host/tests/grpc_timing.rs +++ b/codex-rs/code-mode-host/tests/grpc_timing.rs @@ -119,7 +119,7 @@ async fn stdio_execution_timing_includes_javascript_but_excludes_delayed_reads() ); let delegate = Arc::new(RecordingDelegate::default()); let session = provider - .create_session(delegate.clone()) + .create_session() .await .map_err(anyhow::Error::msg)?; @@ -127,17 +127,20 @@ async fn stdio_execution_timing_includes_javascript_but_excludes_delayed_reads() for (index, suffix) in endings.into_iter().enumerate() { let observed = Instant::now(); let started = session - .execute(ExecuteRequest { - tool_call_id: format!("call-{index}"), - source: format!( - "const until = Date.now() + 150; while (Date.now() < until) {{}} \ + .execute( + ExecuteRequest { + tool_call_id: format!("call-{index}"), + source: format!( + "const until = Date.now() + 150; while (Date.now() < until) {{}} \ await new Promise(resolve => setTimeout(resolve, 150)); \ notify('finished'); {suffix}" - ), - enabled_tools: Vec::new(), - yield_time_ms: Some(/*value*/ 5_000), - max_output_tokens: Some(/*value*/ 1_000), - }) + ), + enabled_tools: Vec::new(), + yield_time_ms: Some(/*value*/ 5_000), + max_output_tokens: Some(/*value*/ 1_000), + }, + delegate.clone(), + ) .await .map_err(anyhow::Error::msg)?; timeout(TEST_TIMEOUT, delegate.notification_delivered.notified()) @@ -184,17 +187,20 @@ async fn observation_timing_excludes_previous_requests_and_background_time() -> ]; for provider in providers { let session = provider - .create_session(Arc::new(NoopCodeModeSessionDelegate)) + .create_session() .await .map_err(anyhow::Error::msg)?; let started = session - .execute(ExecuteRequest { - tool_call_id: "call-1".to_string(), - source: "await new Promise(() => {});".to_string(), - enabled_tools: Vec::new(), - yield_time_ms: Some(/*value*/ 200), - max_output_tokens: Some(/*value*/ 1_000), - }) + .execute( + ExecuteRequest { + tool_call_id: "call-1".to_string(), + source: "await new Promise(() => {});".to_string(), + enabled_tools: Vec::new(), + yield_time_ms: Some(/*value*/ 200), + max_output_tokens: Some(/*value*/ 1_000), + }, + Arc::new(NoopCodeModeSessionDelegate), + ) .await .map_err(anyhow::Error::msg)?; let cell_id = started.cell_id.clone(); diff --git a/codex-rs/code-mode-host/tests/stdio.rs b/codex-rs/code-mode-host/tests/stdio.rs index ff5fef18dc..0f37d454b8 100644 --- a/codex-rs/code-mode-host/tests/stdio.rs +++ b/codex-rs/code-mode-host/tests/stdio.rs @@ -174,9 +174,13 @@ fn execute_request(source: &str) -> ExecuteRequest { } } -async fn execute(session: &Arc, request: ExecuteRequest) -> RuntimeResponse { +async fn execute( + session: &Arc, + request: ExecuteRequest, + delegate: Arc, +) -> RuntimeResponse { session - .execute(request) + .execute(request, delegate.clone()) .await .expect("start execution") .initial_response() @@ -187,8 +191,12 @@ async fn execute(session: &Arc, request: ExecuteRequest) -> async fn execute_to_terminal( session: &Arc, request: ExecuteRequest, + delegate: Arc, ) -> RuntimeResponse { - let started = session.execute(request).await.expect("start execution"); + let started = session + .execute(request, delegate.clone()) + .await + .expect("start execution"); let mut response = started.initial_response().await.expect("initial response"); loop { match response { @@ -227,29 +235,27 @@ async fn session_execution_limits_are_isolated_on_a_shared_process_host() { codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"), )); let limited = provider - .create_session_with_limits( - Arc::new(RecordingDelegate::default()), - CodeModeSessionCellExecutionLimits { - max_yield_time_ms: Some(1), - max_heap_size_bytes: None, - }, - ) + .create_session_with_limits(CodeModeSessionCellExecutionLimits { + max_yield_time_ms: Some(1), + max_heap_size_bytes: None, + }) .await .expect("create limited session"); let other = provider - .create_session_with_limits( - Arc::new(RecordingDelegate::default()), - CodeModeSessionCellExecutionLimits { - max_yield_time_ms: Some(1_000), - max_heap_size_bytes: None, - }, - ) + .create_session_with_limits(CodeModeSessionCellExecutionLimits { + max_yield_time_ms: Some(1_000), + max_heap_size_bytes: None, + }) .await .expect("create independently limited session"); let response = tokio::time::timeout( Duration::from_secs(5), - execute(&limited, execute_request("await new Promise(() => {});")), + execute( + &limited, + execute_request("await new Promise(() => {});"), + Arc::new(RecordingDelegate::default()), + ), ) .await .expect("session limit should bound the default execution wait"); @@ -297,12 +303,18 @@ async fn remote_session_persists_values_forwards_delegates_and_controls_cells() codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"), ); let delegate = Arc::new(RecordingDelegate::default()); + let first_delegate = Arc::new(RecordingDelegate::default()); let session = provider - .create_session(delegate.clone()) + .create_session() .await .expect("create remote session"); - let actual = execute(&session, execute_request(r#"store("key", "persisted");"#)).await; + let actual = execute( + &session, + execute_request(r#"store("key", "persisted");"#), + first_delegate.clone(), + ) + .await; assert_eq!( actual, RuntimeResponse::Result { @@ -329,7 +341,7 @@ text(result.value); input_schema: None, output_schema: None, }]; - let actual = execute(&session, callback_request).await; + let actual = execute(&session, callback_request, delegate.clone()).await; assert_eq!( actual, RuntimeResponse::Result { @@ -359,7 +371,7 @@ text(result.value); let mut pending_request = execute_request("await new Promise(() => {});"); pending_request.tool_call_id = "call-3".to_string(); pending_request.yield_time_ms = Some(1); - let actual = execute(&session, pending_request).await; + let actual = execute(&session, pending_request, delegate.clone()).await; assert_eq!( actual, RuntimeResponse::Yielded { @@ -399,7 +411,14 @@ text(result.value); session.shutdown().await.expect("shutdown remote session"); assert_eq!( *delegate.closed_cells.lock().expect("closed cells lock"), - vec![cell_id("1"), cell_id("2"), cell_id("3")] + vec![cell_id("2"), cell_id("3")] + ); + assert_eq!( + *first_delegate + .closed_cells + .lock() + .expect("closed cells lock"), + vec![cell_id("1")] ); } @@ -409,12 +428,15 @@ async fn dropping_long_wait_releases_observer_before_next_wait() { codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"), ); let session = provider - .create_session(Arc::new(RecordingDelegate::default())) + .create_session() .await .expect("create remote session"); let mut request = execute_request("await new Promise(() => {});"); request.yield_time_ms = Some(1); - let started = session.execute(request).await.expect("start execution"); + let started = session + .execute(request, Arc::new(RecordingDelegate::default())) + .await + .expect("start execution"); let running_cell_id = started.cell_id.clone(); let actual = started.initial_response().await.expect("initial response"); assert_eq!( @@ -472,7 +494,7 @@ async fn unawaited_slow_tool_is_cancelled_after_parallel_tools_complete() { ); let (delegate, mut events_rx) = CancellationDelegate::new(); let session = provider - .create_session(delegate.clone()) + .create_session() .await .expect("create remote session"); let mut request = execute_request( @@ -508,7 +530,10 @@ return; }) .collect(); - let started = session.execute(request).await.expect("start execution"); + let started = session + .execute(request, delegate.clone()) + .await + .expect("start execution"); let running_cell_id = started.cell_id.clone(); let actual = started.initial_response().await.expect("initial response"); assert_eq!( @@ -587,11 +612,14 @@ async fn oversized_execute_request_does_not_close_the_shared_host() { codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"), ); let session = provider - .create_session(Arc::new(RecordingDelegate::default())) + .create_session() .await .expect("create remote session"); let error = session - .execute(execute_request(&"x".repeat(MAX_FRAME_BYTES))) + .execute( + execute_request(&"x".repeat(MAX_FRAME_BYTES)), + Arc::new(RecordingDelegate::default()), + ) .await .err() .expect("oversized execute should fail"); @@ -600,7 +628,12 @@ async fn oversized_execute_request_does_not_close_the_shared_host() { "unexpected error: {error}" ); - let actual = execute(&session, execute_request(r#"text("still alive");"#)).await; + let actual = execute( + &session, + execute_request(r#"text("still alive");"#), + Arc::new(RecordingDelegate::default()), + ) + .await; assert_eq!( actual, RuntimeResponse::Result { @@ -621,7 +654,7 @@ async fn oversized_delegate_payloads_fail_only_the_tool_call() { codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"), ); let session = provider - .create_session(Arc::new(OversizedResultDelegate)) + .create_session() .await .expect("create remote session"); let tool = |name: &str| ToolDefinition { @@ -644,7 +677,12 @@ try {{ )); oversized_argument.enabled_tools = vec![tool("big_argument")]; oversized_argument.yield_time_ms = Some(60_000); - let actual = execute_to_terminal(&session, oversized_argument).await; + let actual = execute_to_terminal( + &session, + oversized_argument, + Arc::new(OversizedResultDelegate), + ) + .await; assert_eq!( actual, RuntimeResponse::Result { @@ -668,7 +706,12 @@ try { ); oversized_result.enabled_tools = vec![tool("big_result")]; oversized_result.yield_time_ms = Some(60_000); - let actual = execute_to_terminal(&session, oversized_result).await; + let actual = execute_to_terminal( + &session, + oversized_result, + Arc::new(OversizedResultDelegate), + ) + .await; assert_eq!( actual, RuntimeResponse::Result { @@ -681,7 +724,12 @@ try { } ); - let actual = execute(&session, execute_request(r#"text("still alive");"#)).await; + let actual = execute( + &session, + execute_request(r#"text("still alive");"#), + Arc::new(OversizedResultDelegate), + ) + .await; assert_eq!( actual, RuntimeResponse::Result { @@ -702,13 +750,14 @@ async fn oversized_initial_response_does_not_close_the_shared_host() { codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"), ); let session = provider - .create_session(Arc::new(RecordingDelegate::default())) + .create_session() .await .expect("create remote session"); let started = session - .execute(execute_request(&format!( - r#"text("x".repeat({MAX_FRAME_BYTES}));"# - ))) + .execute( + execute_request(&format!(r#"text("x".repeat({MAX_FRAME_BYTES}));"#)), + Arc::new(RecordingDelegate::default()), + ) .await .expect("start oversized response"); let error = started @@ -720,7 +769,12 @@ async fn oversized_initial_response_does_not_close_the_shared_host() { "unexpected error: {error}" ); - let actual = execute(&session, execute_request(r#"text("still alive");"#)).await; + let actual = execute( + &session, + execute_request(r#"text("still alive");"#), + Arc::new(RecordingDelegate::default()), + ) + .await; assert_eq!( actual, RuntimeResponse::Result { @@ -763,11 +817,11 @@ async fn child_process_loss_cleans_up_and_rebuilds_the_shared_host() { delegate_a.hold_slow_cleanup(); let delegate_b = Arc::new(RecordingDelegate::default()); let session_a = provider - .create_session(delegate_a.clone()) + .create_session() .await .expect("create first remote session"); let session_b = provider - .create_session(delegate_b.clone()) + .create_session() .await .expect("create second remote session"); @@ -782,7 +836,7 @@ async fn child_process_loss_cleans_up_and_rebuilds_the_shared_host() { output_schema: None, }]; let started_a = session_a - .execute(request_a) + .execute(request_a, delegate_a.clone()) .await .expect("start first cell"); let cell_a = started_a.cell_id.clone(); @@ -806,7 +860,7 @@ async fn child_process_loss_cleans_up_and_rebuilds_the_shared_host() { let mut request_b = execute_request("await new Promise(() => {});"); request_b.yield_time_ms = Some(1); let started_b = session_b - .execute(request_b) + .execute(request_b, delegate_b.clone()) .await .expect("start second cell"); let cell_b = started_b.cell_id.clone(); @@ -899,7 +953,12 @@ async fn child_process_loss_cleans_up_and_rebuilds_the_shared_host() { .await .expect("unrelated session cleanup timeout"); - let actual = execute(&session_b, execute_request(r#"text("replacement");"#)).await; + let actual = execute( + &session_b, + execute_request(r#"text("replacement");"#), + delegate_b.clone(), + ) + .await; assert_eq!( actual, RuntimeResponse::Result { diff --git a/codex-rs/code-mode-protocol/src/session.rs b/codex-rs/code-mode-protocol/src/session.rs index ffba25005c..36b88b916f 100644 --- a/codex-rs/code-mode-protocol/src/session.rs +++ b/codex-rs/code-mode-protocol/src/session.rs @@ -95,7 +95,10 @@ impl StartedCell { } } -/// Host callbacks used by a code-mode session while cells are executing. +/// Host callbacks owned by one code-mode execution. +/// +/// The session retains the supplied delegate while starting and running the cell, +/// including across yields, and releases it through its existing close/cancel paths. pub trait CodeModeSessionDelegate: Send + Sync { fn invoke_tool<'a>( &'a self, @@ -152,6 +155,7 @@ pub trait CodeModeSession: Send + Sync { fn execute<'a>( &'a self, request: ExecuteRequest, + delegate: Arc, ) -> CodeModeSessionResultFuture<'a, StartedCell>; fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome>; @@ -171,10 +175,7 @@ pub trait CodeModeSessionProvider: Send + Sync { Ok(()) } - fn create_session<'a>( - &'a self, - delegate: Arc, - ) -> CodeModeSessionProviderFuture<'a>; + fn create_session(&self) -> CodeModeSessionProviderFuture<'_>; /// Creates a session whose cells share the supplied execution limits. /// @@ -182,11 +183,10 @@ pub trait CodeModeSessionProvider: Send + Sync { /// explicitly implement this method before accepting non-default limits. fn create_session_with_limits<'a>( &'a self, - delegate: Arc, limits: CodeModeSessionCellExecutionLimits, ) -> CodeModeSessionProviderFuture<'a> { if limits == CodeModeSessionCellExecutionLimits::default() { - self.create_session(delegate) + self.create_session() } else { Box::pin(async { Err("code-mode session provider does not support resource limits".to_string()) diff --git a/codex-rs/code-mode-runtime/src/service.rs b/codex-rs/code-mode-runtime/src/service.rs index 8b8d71082a..002dd8a2b2 100644 --- a/codex-rs/code-mode-runtime/src/service.rs +++ b/codex-rs/code-mode-runtime/src/service.rs @@ -13,7 +13,6 @@ use codex_code_mode_protocol::ExecuteRequest; use codex_code_mode_protocol::ExecuteToPendingOutcome; use codex_code_mode_protocol::FunctionCallOutputContentItem; use codex_code_mode_protocol::ImageDetail; -use codex_code_mode_protocol::NoopCodeModeSessionDelegate; use codex_code_mode_protocol::RuntimeResponse; use codex_code_mode_protocol::StartedCell; use codex_code_mode_protocol::WaitOutcome; @@ -31,25 +30,18 @@ const YIELD_GRACE_PERIOD: Duration = Duration::from_secs(1); const MIN_YIELD_TIME_FOR_GRACE: Duration = Duration::from_secs(10); pub struct InProcessCodeModeSession { - runtime: SessionRuntime, + runtime: SessionRuntime, cell_execution_limits: CodeModeSessionCellExecutionLimits, } impl InProcessCodeModeSession { pub fn new() -> Self { - Self::with_delegate(Arc::new(NoopCodeModeSessionDelegate)) + Self::with_limits(CodeModeSessionCellExecutionLimits::default()) } - pub fn with_delegate(delegate: Arc) -> Self { - Self::with_delegate_and_limits(delegate, CodeModeSessionCellExecutionLimits::default()) - } - - pub fn with_delegate_and_limits( - delegate: Arc, - cell_execution_limits: CodeModeSessionCellExecutionLimits, - ) -> Self { + pub fn with_limits(cell_execution_limits: CodeModeSessionCellExecutionLimits) -> Self { Self { - runtime: SessionRuntime::new(Arc::new(ProtocolDelegate { delegate })), + runtime: SessionRuntime::new(), cell_execution_limits: CodeModeSessionCellExecutionLimits { max_heap_size_bytes: None, ..cell_execution_limits @@ -57,16 +49,12 @@ impl InProcessCodeModeSession { } } - pub fn with_delegate_and_task_failure_handler( - delegate: Arc, + pub fn with_task_failure_handler( task_failure_handler: Arc, cell_execution_limits: CodeModeSessionCellExecutionLimits, ) -> Self { Self { - runtime: SessionRuntime::new_with_task_failure_handler( - Arc::new(ProtocolDelegate { delegate }), - Some(task_failure_handler), - ), + runtime: SessionRuntime::new_with_task_failure_handler(Some(task_failure_handler)), cell_execution_limits: CodeModeSessionCellExecutionLimits { max_heap_size_bytes: None, ..cell_execution_limits @@ -74,13 +62,18 @@ impl InProcessCodeModeSession { } } - pub async fn execute(&self, request: ExecuteRequest) -> Result { + pub async fn execute( + &self, + request: ExecuteRequest, + delegate: Arc, + ) -> Result { let yield_time_ms = request.yield_time_ms.unwrap_or(DEFAULT_EXEC_YIELD_TIME_MS); let started = self .runtime .execute( runtime_request(request), runtime::ObserveMode::YieldAfter(self.resolve_yield_timeout(yield_time_ms)), + Arc::new(ProtocolDelegate { delegate }), ) .await .map_err(|error| error.to_string())?; @@ -101,12 +94,14 @@ impl InProcessCodeModeSession { pub async fn execute_to_pending( &self, request: ExecuteRequest, + delegate: Arc, ) -> Result { let started = self .runtime .execute( runtime_request(request), runtime::ObserveMode::PendingFrontier, + Arc::new(ProtocolDelegate { delegate }), ) .await .map_err(|error| error.to_string())?; @@ -220,8 +215,9 @@ impl CodeModeSession for InProcessCodeModeSession { fn execute<'a>( &'a self, request: ExecuteRequest, + delegate: Arc, ) -> CodeModeSessionResultFuture<'a, StartedCell> { - Box::pin(InProcessCodeModeSession::execute(self, request)) + Box::pin(InProcessCodeModeSession::execute(self, request, delegate)) } fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome> { diff --git a/codex-rs/code-mode-runtime/src/service_contract_tests.rs b/codex-rs/code-mode-runtime/src/service_contract_tests.rs index 1141c079c0..973440a775 100644 --- a/codex-rs/code-mode-runtime/src/service_contract_tests.rs +++ b/codex-rs/code-mode-runtime/src/service_contract_tests.rs @@ -1,3 +1,4 @@ +use codex_code_mode_protocol::NoopCodeModeSessionDelegate; use std::sync::Arc; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; @@ -184,15 +185,88 @@ async fn next_event(events_rx: &mut mpsc::UnboundedReceiver) -> D .expect("delegate event channel closed") } +#[tokio::test] +async fn yielded_cells_retain_their_own_delegate_until_closed() { + let service = InProcessCodeModeSession::new(); + let (delegate_a, mut events_a) = BlockingDelegate::new(); + let (delegate_b, mut events_b) = BlockingDelegate::new(); + let weak_a = Arc::downgrade(&delegate_a); + let request = ExecuteRequest { + enabled_tools: vec![blocking_tool()], + ..execute_request("await tools.block({}); await tools.block({});") + }; + let a = service + .execute(request.clone(), delegate_a.clone()) + .await + .unwrap(); + assert_eq!(next_event(&mut events_a).await, DelegateEvent::ToolStarted); + assert!(matches!( + a.initial_response().await.unwrap(), + RuntimeResponse::Yielded { .. } + )); + + let b = service.execute(request, delegate_b.clone()).await.unwrap(); + assert_eq!(next_event(&mut events_b).await, DelegateEvent::ToolStarted); + assert!(matches!( + b.initial_response().await.unwrap(), + RuntimeResponse::Yielded { .. } + )); + + // A's next callback must still reach A after B has started in the same session. + delegate_a.release_tool(); + assert_eq!(next_event(&mut events_a).await, DelegateEvent::ToolStarted); + delegate_a.release_tool(); + assert_eq!( + service + .wait(WaitRequest { + cell_id: cell_id("1"), + yield_time_ms: 60_000 + }) + .await + .unwrap(), + WaitOutcome::LiveCell(RuntimeResponse::Result { + code_mode_host_duration: None, + cell_id: cell_id("1"), + content_items: Vec::new(), + error_text: None + }) + ); + assert_eq!( + next_event(&mut events_a).await, + DelegateEvent::CellClosed(cell_id("1")) + ); + drop(delegate_a); + assert_eq!( + tokio::time::timeout(Duration::from_secs(2), events_a.recv()) + .await + .unwrap(), + None + ); + assert!(weak_a.upgrade().is_none()); + + service.terminate(cell_id("2")).await.unwrap(); + assert_eq!( + next_event(&mut events_b).await, + DelegateEvent::ToolCancelled + ); + assert_eq!( + next_event(&mut events_b).await, + DelegateEvent::CellClosed(cell_id("2")) + ); +} + #[tokio::test] async fn yields_and_resumes() { let service = InProcessCodeModeSession::new(); let cell = service - .execute(ExecuteRequest { - source: r#"text("before"); yield_control(); text("after");"#.to_string(), - yield_time_ms: Some(60_000), - ..execute_request("") - }) + .execute( + ExecuteRequest { + source: r#"text("before"); yield_control(); text("after");"#.to_string(), + yield_time_ms: Some(60_000), + ..execute_request("") + }, + Arc::new(NoopCodeModeSessionDelegate), + ) .await .unwrap(); @@ -228,20 +302,23 @@ async fn yields_and_resumes() { #[tokio::test] async fn returns_and_resumes_from_the_pending_frontier() { let (delegate, mut events_rx) = BlockingDelegate::new(); - let service = InProcessCodeModeSession::with_delegate(delegate.clone()); + let service = InProcessCodeModeSession::new(); assert_eq!( service - .execute_to_pending(ExecuteRequest { - enabled_tools: vec![blocking_tool()], - source: r#" + .execute_to_pending( + ExecuteRequest { + enabled_tools: vec![blocking_tool()], + source: r#" await tools.block({}); text("after"); "# - .to_string(), - yield_time_ms: Some(60_000), - ..execute_request("") - }) + .to_string(), + yield_time_ms: Some(60_000), + ..execute_request("") + }, + delegate.clone() + ) .await .unwrap(), ExecuteToPendingOutcome::Pending { @@ -278,9 +355,10 @@ text("after"); async fn observed_natural_completion_wins_over_termination() { let service = InProcessCodeModeSession::new(); let cell = service - .execute(execute_request( - r#"yield_control(); store("finished", true); text("done");"#, - )) + .execute( + execute_request(r#"yield_control(); store("finished", true); text("done");"#), + Arc::new(NoopCodeModeSessionDelegate), + ) .await .unwrap(); @@ -295,10 +373,13 @@ async fn observed_natural_completion_wins_over_termination() { tokio::time::timeout(Duration::from_secs(1), async { loop { let response = service - .execute(ExecuteRequest { - yield_time_ms: Some(60_000), - ..execute_request(r#"text(String(load("finished")));"#) - }) + .execute( + ExecuteRequest { + yield_time_ms: Some(60_000), + ..execute_request(r#"text(String(load("finished")));"#) + }, + Arc::new(NoopCodeModeSessionDelegate), + ) .await .unwrap() .initial_response() @@ -335,11 +416,12 @@ async fn observed_natural_completion_wins_over_termination() { #[tokio::test] async fn termination_cancels_pending_callbacks_before_responding() { let (delegate, mut events_rx) = BlockingDelegate::new(); - let service = InProcessCodeModeSession::with_delegate(delegate.clone()); + let service = InProcessCodeModeSession::new(); let cell = service - .execute(execute_request( - r#"notify("pending"); await new Promise(() => {});"#, - )) + .execute( + execute_request(r#"notify("pending"); await new Promise(() => {});"#), + delegate.clone(), + ) .await .unwrap(); @@ -377,9 +459,9 @@ async fn termination_cancels_pending_callbacks_before_responding() { #[tokio::test] async fn shutdown_cancels_notifications_while_natural_completion_is_draining() { let (delegate, mut events_rx) = HeldNotificationDelegate::new(); - let service = Arc::new(InProcessCodeModeSession::with_delegate(delegate.clone())); + let service = Arc::new(InProcessCodeModeSession::new()); service - .execute(execute_request(r#"notify("pending");"#)) + .execute(execute_request(r#"notify("pending");"#), delegate.clone()) .await .unwrap(); @@ -407,11 +489,12 @@ async fn shutdown_cancels_notifications_while_natural_completion_is_draining() { #[tokio::test] async fn repeated_termination_is_rejected_while_callback_cleanup_is_pending() { let (delegate, mut events_rx) = HeldNotificationDelegate::new(); - let service = Arc::new(InProcessCodeModeSession::with_delegate(delegate.clone())); + let service = Arc::new(InProcessCodeModeSession::new()); let cell = service - .execute(execute_request( - r#"notify("pending"); await new Promise(() => {});"#, - )) + .execute( + execute_request(r#"notify("pending"); await new Promise(() => {});"#), + delegate.clone(), + ) .await .unwrap(); @@ -461,7 +544,10 @@ async fn repeated_termination_is_rejected_while_callback_cleanup_is_pending() { async fn second_observer_is_rejected_without_displacing_the_first() { let service = InProcessCodeModeSession::new(); let cell = service - .execute(execute_request("await new Promise(() => {});")) + .execute( + execute_request("await new Promise(() => {});"), + Arc::new(NoopCodeModeSessionDelegate), + ) .await .unwrap(); @@ -509,14 +595,17 @@ async fn second_observer_is_rejected_without_displacing_the_first() { #[tokio::test] async fn natural_completion_cleans_up_callbacks_before_responding() { let (delegate, mut events_rx) = BlockingDelegate::new(); - let service = InProcessCodeModeSession::with_delegate(delegate.clone()); + let service = InProcessCodeModeSession::new(); let cell = service - .execute(ExecuteRequest { - enabled_tools: vec![blocking_tool()], - source: r#"tools.block({}); text("done");"#.to_string(), - yield_time_ms: Some(60_000), - ..execute_request("") - }) + .execute( + ExecuteRequest { + enabled_tools: vec![blocking_tool()], + source: r#"tools.block({}); text("done");"#.to_string(), + yield_time_ms: Some(60_000), + ..execute_request("") + }, + delegate.clone(), + ) .await .unwrap(); diff --git a/codex-rs/code-mode-runtime/src/service_tests.rs b/codex-rs/code-mode-runtime/src/service_tests.rs index 66b8cf46b6..fafcb69e8c 100644 --- a/codex-rs/code-mode-runtime/src/service_tests.rs +++ b/codex-rs/code-mode-runtime/src/service_tests.rs @@ -1,3 +1,4 @@ +use codex_code_mode_protocol::NoopCodeModeSessionDelegate; use std::sync::Arc; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; @@ -43,13 +44,10 @@ fn resolve_yield_timeout_applies_grace_before_session_limits() { (Some(10_500), 5_000, Duration::from_secs(5)), (Some(u64::MAX), u64::MAX, Duration::from_millis(u64::MAX)), ] { - let session = InProcessCodeModeSession::with_delegate_and_limits( - Arc::new(ReleasableToolDelegate::default()), - CodeModeSessionCellExecutionLimits { - max_yield_time_ms, - max_heap_size_bytes: None, - }, - ); + let session = InProcessCodeModeSession::with_limits(CodeModeSessionCellExecutionLimits { + max_yield_time_ms, + max_heap_size_bytes: None, + }); assert_eq!( session.resolve_yield_timeout(requested_yield_time_ms), @@ -62,14 +60,14 @@ fn resolve_yield_timeout_applies_grace_before_session_limits() { #[tokio::test(start_paused = true)] async fn execute_waits_for_nested_tool_during_yield_grace() { let delegate = Arc::new(ReleasableToolDelegate::default()); - let service = InProcessCodeModeSession::with_delegate(delegate.clone()); + let service = InProcessCodeModeSession::new(); let request = ExecuteRequest { enabled_tools: vec![echo_tool()], source: r#"await tools.echo({}); text("done");"#.to_string(), yield_time_ms: Some(10_000), ..execute_request("") }; - let started = service.execute(request).await.unwrap(); + let started = service.execute(request, delegate.clone()).await.unwrap(); let response = tokio::spawn(started.initial_response()); wait_until_tool_started(&delegate).await; tokio::time::advance(Duration::from_millis(10_500)).await; @@ -93,20 +91,20 @@ async fn execute_waits_for_nested_tool_during_yield_grace() { #[tokio::test(start_paused = true)] async fn execute_and_wait_clamp_yield_grace_without_stopping_the_cell() { let delegate = Arc::new(ReleasableToolDelegate::default()); - let service = InProcessCodeModeSession::with_delegate_and_limits( - delegate.clone(), - CodeModeSessionCellExecutionLimits { - max_yield_time_ms: Some(/*value*/ 10_000), - max_heap_size_bytes: None, - }, - ); + let service = InProcessCodeModeSession::with_limits(CodeModeSessionCellExecutionLimits { + max_yield_time_ms: Some(/*value*/ 10_000), + max_heap_size_bytes: None, + }); let started = service - .execute(ExecuteRequest { - enabled_tools: vec![echo_tool()], - source: r#"await tools.echo({}); text("done");"#.to_string(), - yield_time_ms: None, - ..execute_request("") - }) + .execute( + ExecuteRequest { + enabled_tools: vec![echo_tool()], + source: r#"await tools.echo({}); text("done");"#.to_string(), + yield_time_ms: None, + ..execute_request("") + }, + delegate.clone(), + ) .await .unwrap(); let initial_response = tokio::spawn(started.initial_response()); @@ -169,13 +167,16 @@ async fn execute_and_wait_clamp_yield_grace_without_stopping_the_cell() { #[tokio::test(start_paused = true)] async fn wait_waits_for_nested_tool_during_yield_grace() { let delegate = Arc::new(ReleasableToolDelegate::default()); - let service = InProcessCodeModeSession::with_delegate(delegate.clone()); + let service = InProcessCodeModeSession::new(); let initial_response = service - .execute_to_pending(ExecuteRequest { - enabled_tools: vec![echo_tool()], - source: r#"await tools.echo({}); text("done");"#.to_string(), - ..execute_request("") - }) + .execute_to_pending( + ExecuteRequest { + enabled_tools: vec![echo_tool()], + source: r#"await tools.echo({}); text("done");"#.to_string(), + ..execute_request("") + }, + delegate.clone(), + ) .await .unwrap(); assert_eq!( @@ -215,29 +216,30 @@ async fn wait_waits_for_nested_tool_during_yield_grace() { #[tokio::test(start_paused = true)] async fn zero_yield_limit_is_immediate_and_scoped_to_its_session() { let zero_delegate = Arc::new(ReleasableToolDelegate::default()); - let zero_session = InProcessCodeModeSession::with_delegate_and_limits( - zero_delegate.clone(), - CodeModeSessionCellExecutionLimits { - max_yield_time_ms: Some(/*value*/ 0), - max_heap_size_bytes: None, - }, - ); + let zero_session = InProcessCodeModeSession::with_limits(CodeModeSessionCellExecutionLimits { + max_yield_time_ms: Some(/*value*/ 0), + max_heap_size_bytes: None, + }); let limited_delegate = Arc::new(ReleasableToolDelegate::default()); - let limited_session = InProcessCodeModeSession::with_delegate_and_limits( - limited_delegate.clone(), - CodeModeSessionCellExecutionLimits { + let limited_session = + InProcessCodeModeSession::with_limits(CodeModeSessionCellExecutionLimits { max_yield_time_ms: Some(/*value*/ 10), max_heap_size_bytes: None, - }, - ); + }); let request = ExecuteRequest { enabled_tools: vec![echo_tool()], source: "await tools.echo({});".to_string(), yield_time_ms: Some(/*value*/ 60_000), ..execute_request("") }; - let zero_started = zero_session.execute(request.clone()).await.unwrap(); - let limited_started = limited_session.execute(request).await.unwrap(); + let zero_started = zero_session + .execute(request.clone(), zero_delegate.clone()) + .await + .unwrap(); + let limited_started = limited_session + .execute(request, limited_delegate.clone()) + .await + .unwrap(); let zero_response = tokio::spawn(zero_started.initial_response()); let limited_response = tokio::spawn(limited_started.initial_response()); wait_until_tool_started(&zero_delegate).await; @@ -372,7 +374,7 @@ fn echo_tool() -> ToolDefinition { async fn execute(service: &InProcessCodeModeSession, request: ExecuteRequest) -> RuntimeResponse { service - .execute(request) + .execute(request, Arc::new(NoopCodeModeSessionDelegate)) .await .unwrap() .initial_response() @@ -521,10 +523,13 @@ async fn shutdown_interrupts_cpu_bound_cells() { let service = InProcessCodeModeSession::new(); let cell = service - .execute(ExecuteRequest { - source: "while (true) {}".to_string(), - ..execute_request("") - }) + .execute( + ExecuteRequest { + source: "while (true) {}".to_string(), + ..execute_request("") + }, + Arc::new(NoopCodeModeSessionDelegate), + ) .await .unwrap(); assert_eq!( @@ -548,7 +553,10 @@ async fn start_cell_rejects_new_cell_after_shutdown_begins() { service.shutdown().await.unwrap(); let error = service - .execute(execute_request("text('late');")) + .execute( + execute_request("text('late');"), + Arc::new(NoopCodeModeSessionDelegate), + ) .await .err() .unwrap(); @@ -561,11 +569,14 @@ async fn execute_to_pending_returns_completed_for_synchronous_results() { let service = InProcessCodeModeSession::new(); let response = service - .execute_to_pending(ExecuteRequest { - source: r#"text("done");"#.to_string(), - yield_time_ms: Some(60_000), - ..execute_request("") - }) + .execute_to_pending( + ExecuteRequest { + source: r#"text("done");"#.to_string(), + yield_time_ms: Some(60_000), + ..execute_request("") + }, + Arc::new(NoopCodeModeSessionDelegate), + ) .await .unwrap(); @@ -588,11 +599,14 @@ async fn execute_to_pending_returns_once_the_runtime_is_quiescent() { let response = tokio::time::timeout( Duration::from_secs(1), - service.execute_to_pending(ExecuteRequest { - source: r#"text("before"); await new Promise(() => {});"#.to_string(), - yield_time_ms: Some(60_000), - ..execute_request("") - }), + service.execute_to_pending( + ExecuteRequest { + source: r#"text("before"); await new Promise(() => {});"#.to_string(), + yield_time_ms: Some(60_000), + ..execute_request("") + }, + Arc::new(NoopCodeModeSessionDelegate), + ), ) .await .unwrap() @@ -626,18 +640,21 @@ async fn execute_to_pending_identifies_tool_calls_in_paused_frontier() { let service = InProcessCodeModeSession::new(); let response = service - .execute_to_pending(ExecuteRequest { - enabled_tools: vec![echo_tool()], - source: r#" + .execute_to_pending( + ExecuteRequest { + enabled_tools: vec![echo_tool()], + source: r#" await Promise.all([ tools.echo({ value: "first" }), tools.echo({ value: "second" }), ]); "# - .to_string(), - yield_time_ms: Some(60_000), - ..execute_request("") - }) + .to_string(), + yield_time_ms: Some(60_000), + ..execute_request("") + }, + Arc::new(NoopCodeModeSessionDelegate), + ) .await .unwrap(); @@ -667,9 +684,10 @@ async fn execute_to_pending_excludes_delayed_timeout_tool_calls_until_wait() { let service = InProcessCodeModeSession::new(); let initial_response = service - .execute_to_pending(ExecuteRequest { - enabled_tools: vec![echo_tool()], - source: r#" + .execute_to_pending( + ExecuteRequest { + enabled_tools: vec![echo_tool()], + source: r#" setTimeout(() => { tools.echo({ value: "delayed" }); }, 1000); @@ -678,10 +696,12 @@ await Promise.all([ tools.echo({ value: "third" }), ]); "# - .to_string(), - yield_time_ms: Some(60_000), - ..execute_request("") - }) + .to_string(), + yield_time_ms: Some(60_000), + ..execute_request("") + }, + Arc::new(NoopCodeModeSessionDelegate), + ) .await .unwrap(); @@ -730,20 +750,23 @@ await Promise.all([ #[tokio::test] async fn wait_to_pending_returns_after_resumed_runtime_becomes_quiescent_again() { let delegate = Arc::new(ReleasableToolDelegate::default()); - let service = InProcessCodeModeSession::with_delegate(delegate.clone()); + let service = InProcessCodeModeSession::new(); let initial_response = service - .execute_to_pending(ExecuteRequest { - enabled_tools: vec![echo_tool()], - source: r#" + .execute_to_pending( + ExecuteRequest { + enabled_tools: vec![echo_tool()], + source: r#" await tools.echo({}); text("after"); await new Promise(() => {}); "# - .to_string(), - yield_time_ms: Some(60_000), - ..execute_request("") - }) + .to_string(), + yield_time_ms: Some(60_000), + ..execute_request("") + }, + delegate.clone(), + ) .await .unwrap(); @@ -794,19 +817,22 @@ await new Promise(() => {}); #[tokio::test] async fn wait_to_pending_returns_completed_after_resumed_runtime_finishes() { let delegate = Arc::new(ReleasableToolDelegate::default()); - let service = InProcessCodeModeSession::with_delegate(delegate.clone()); + let service = InProcessCodeModeSession::new(); let initial_response = service - .execute_to_pending(ExecuteRequest { - enabled_tools: vec![echo_tool()], - source: r#" + .execute_to_pending( + ExecuteRequest { + enabled_tools: vec![echo_tool()], + source: r#" await tools.echo({}); text("done"); "# - .to_string(), - yield_time_ms: Some(60_000), - ..execute_request("") - }) + .to_string(), + yield_time_ms: Some(60_000), + ..execute_request("") + }, + delegate.clone(), + ) .await .unwrap(); @@ -1271,9 +1297,7 @@ async fn audio_helper_rejects_non_data_urls() { async fn image_helper_accepts_raw_mcp_image_block_with_original_detail() { let service = InProcessCodeModeSession::new(); - let response = execute( - &service, - ExecuteRequest { + let response = execute(&service, ExecuteRequest { source: r#" image({ type: "image", @@ -1285,8 +1309,7 @@ image({ .to_string(), yield_time_ms: None, ..execute_request("") - }, - ) + }) .await; assert_eq!( @@ -1383,9 +1406,7 @@ image( async fn image_helper_second_arg_overrides_raw_mcp_image_detail() { let service = InProcessCodeModeSession::new(); - let response = execute( - &service, - ExecuteRequest { + let response = execute(&service, ExecuteRequest { source: r#" image( { @@ -1400,8 +1421,7 @@ image( .to_string(), yield_time_ms: None, ..execute_request("") - }, - ) + }) .await; assert_eq!( @@ -1559,9 +1579,7 @@ image({ async fn image_helper_rejects_raw_mcp_result_container() { let service = InProcessCodeModeSession::new(); - let response = execute( - &service, - ExecuteRequest { + let response = execute(&service, ExecuteRequest { source: r#" image({ content: [ @@ -1578,8 +1596,7 @@ image({ .to_string(), yield_time_ms: None, ..execute_request("") - }, - ) + }) .await; assert_eq!( diff --git a/codex-rs/code-mode-runtime/src/session_runtime/mod.rs b/codex-rs/code-mode-runtime/src/session_runtime/mod.rs index d428112f2b..f069b8af0f 100644 --- a/codex-rs/code-mode-runtime/src/session_runtime/mod.rs +++ b/codex-rs/code-mode-runtime/src/session_runtime/mod.rs @@ -38,27 +38,25 @@ use crate::cell_actor::CompletionCommit; type RuntimeEventFuture = Pin> + Send + 'static>>; /// Owns all cells and shared state for one transport-neutral code-mode session. -pub(crate) struct SessionRuntime { - inner: Arc>, +pub(crate) struct SessionRuntime { + inner: Arc, } -struct Inner { +struct Inner { stored_values: Mutex>, cells: Mutex>, cell_tasks: TaskTracker, shutdown_token: CancellationToken, - delegate: Arc, task_failure_handler: Option, next_cell_id: AtomicU64, } -impl SessionRuntime { - pub(crate) fn new(delegate: Arc) -> Self { - Self::new_with_task_failure_handler(delegate, /*task_failure_handler*/ None) +impl SessionRuntime { + pub(crate) fn new() -> Self { + Self::new_with_task_failure_handler(/*task_failure_handler*/ None) } pub(crate) fn new_with_task_failure_handler( - delegate: Arc, task_failure_handler: Option, ) -> Self { Self { @@ -67,24 +65,24 @@ impl SessionRuntime { cells: Mutex::new(HashMap::new()), cell_tasks: TaskTracker::new(), shutdown_token: CancellationToken::new(), - delegate, task_failure_handler, next_cell_id: AtomicU64::new(1), }), } } - pub(crate) async fn execute( + pub(crate) async fn execute( &self, request: CreateCellRequest, initial_observe_mode: ObserveMode, + delegate: Arc, ) -> Result { if self.inner.shutdown_token.is_cancelled() { return Err(Error::ShuttingDown); } let cell_id = self.allocate_cell_id()?; let initial_event = self - .start_cell(cell_id.clone(), request, initial_observe_mode) + .start_cell(cell_id.clone(), request, initial_observe_mode, delegate) .await?; Ok(StartedCell { cell_id, @@ -154,14 +152,16 @@ impl SessionRuntime { .map_err(|_| Error::CellIdSpaceExhausted) } - async fn start_cell( + async fn start_cell( &self, cell_id: CellId, request: CreateCellRequest, initial_observe_mode: ObserveMode, + delegate: Arc, ) -> Result { let stored_values = self.inner.stored_values.lock().await.clone(); let host = Arc::new(RuntimeCellHost { + delegate, cell_id: cell_id.clone(), inner: Arc::clone(&self.inner), execution_context: opentelemetry::Context::current(), @@ -205,7 +205,7 @@ impl SessionRuntime { } } -impl Drop for SessionRuntime { +impl Drop for SessionRuntime { fn drop(&mut self) { self.begin_shutdown(); } @@ -235,8 +235,9 @@ impl PendingEvent { } struct RuntimeCellHost { + delegate: Arc, cell_id: CellId, - inner: Arc>, + inner: Arc, // Callbacks outlive the initial request and run in separate tasks. Preserve // their trace parent without retaining the request's tracing span. execution_context: opentelemetry::Context, @@ -248,8 +249,7 @@ impl CellHost for RuntimeCellHost { invocation: CellToolCall, cancellation_token: CancellationToken, ) -> Result { - self.inner - .delegate + self.delegate .invoke_tool( NestedToolCall { cell_id: self.cell_id.clone(), @@ -270,8 +270,7 @@ impl CellHost for RuntimeCellHost { text: String, cancellation_token: CancellationToken, ) -> Result<(), String> { - self.inner - .delegate + self.delegate .notify(call_id, self.cell_id.clone(), text, cancellation_token) .await } @@ -298,7 +297,7 @@ 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); + self.delegate.cell_closed(&self.cell_id); } } diff --git a/codex-rs/code-mode-runtime/src/session_runtime/tests.rs b/codex-rs/code-mode-runtime/src/session_runtime/tests.rs index 9085e9c9bf..797cd00999 100644 --- a/codex-rs/code-mode-runtime/src/session_runtime/tests.rs +++ b/codex-rs/code-mode-runtime/src/session_runtime/tests.rs @@ -66,16 +66,14 @@ impl SessionRuntimeDelegate for PanickingClosedDelegate { #[tokio::test] async fn reports_cell_actor_panics_to_the_owner() { let (failure_tx, mut failure_rx) = tokio::sync::mpsc::unbounded_channel(); - let runtime = SessionRuntime::new_with_task_failure_handler( - Arc::new(PanickingClosedDelegate), - Some(Arc::new(move |reason| { - let _ = failure_tx.send(reason); - })), - ); + let runtime = SessionRuntime::new_with_task_failure_handler(Some(Arc::new(move |reason| { + let _ = failure_tx.send(reason); + }))); let started = runtime .execute( execute_request(r#"text("done");"#), ObserveMode::YieldAfter(Duration::from_secs(1)), + Arc::new(PanickingClosedDelegate), ) .await .expect("start cell"); @@ -97,9 +95,10 @@ async fn reports_cell_actor_panics_to_the_owner() { #[tokio::test] async fn termination_rejects_a_waiting_store_commit_before_the_next_cell_can_load_it() { - let runtime = SessionRuntime::new(Arc::new(RecordingDelegate)); + let runtime = SessionRuntime::new(); let cell_state = Arc::new(CellState::new(CancellationToken::new())); let host = RuntimeCellHost { + delegate: Arc::new(RecordingDelegate), cell_id: CellId::new("terminating-writer"), inner: Arc::clone(&runtime.inner), execution_context: opentelemetry::Context::new(), @@ -154,6 +153,7 @@ async fn termination_rejects_a_waiting_store_commit_before_the_next_cell_can_loa source: r#"text(String(load("candidate")));"#.to_string(), }, ObserveMode::YieldAfter(Duration::from_secs(1)), + Arc::new(RecordingDelegate), ) .await .unwrap(); @@ -179,7 +179,7 @@ fn execute_request(source: &str) -> CreateCellRequest { #[tokio::test] async fn cell_id_allocation_fails_before_wrapping() { - let runtime = SessionRuntime::new(Arc::new(RecordingDelegate)); + let runtime = SessionRuntime::new(); runtime .inner .next_cell_id @@ -190,6 +190,7 @@ async fn cell_id_allocation_fails_before_wrapping() { .execute( execute_request(r#"text("unreachable");"#), ObserveMode::YieldAfter(Duration::from_secs(1)), + Arc::new(RecordingDelegate) ) .await .err(), @@ -203,12 +204,13 @@ async fn cell_id_allocation_fails_before_wrapping() { reason = "test holds the registry lock to force admission ahead of shutdown" )] async fn shutdown_rejects_cell_admission_queued_before_the_registry_lock() { - let runtime = Arc::new(SessionRuntime::new(Arc::new(RecordingDelegate))); + let runtime = Arc::new(SessionRuntime::new()); let cells = runtime.inner.cells.lock().await; let execution = runtime.execute( execute_request("while (true) {}"), ObserveMode::YieldAfter(Duration::from_millis(/*millis*/ 1)), + Arc::new(RecordingDelegate), ); tokio::pin!(execution); std::future::poll_fn(|context| match execution.as_mut().poll(context) { @@ -238,11 +240,12 @@ async fn shutdown_rejects_cell_admission_queued_before_the_registry_lock() { #[tokio::test] async fn drop_terminates_cells_when_the_registry_is_locked() { - let runtime = SessionRuntime::new(Arc::new(RecordingDelegate)); + let runtime = SessionRuntime::new(); let started = runtime .execute( execute_request("while (true) {}"), ObserveMode::YieldAfter(Duration::from_millis(/*millis*/ 1)), + Arc::new(RecordingDelegate), ) .await .unwrap(); diff --git a/codex-rs/code-mode-runtime/tests/jit.rs b/codex-rs/code-mode-runtime/tests/jit.rs index 35c22bfa8c..2ab5be527e 100644 --- a/codex-rs/code-mode-runtime/tests/jit.rs +++ b/codex-rs/code-mode-runtime/tests/jit.rs @@ -1,9 +1,11 @@ +use codex_code_mode_protocol::NoopCodeModeSessionDelegate; use codex_code_mode_runtime::ExecuteRequest; use codex_code_mode_runtime::InProcessCodeModeSession; use codex_code_mode_runtime::RuntimeResponse; use codex_code_mode_runtime::V8JitMode; use codex_code_mode_runtime::initialize_v8; use pretty_assertions::assert_eq; +use std::sync::Arc; #[tokio::test] async fn code_mode_runs_with_jit_disabled() { @@ -11,13 +13,16 @@ async fn code_mode_runs_with_jit_disabled() { let service = InProcessCodeModeSession::new(); let started = service - .execute(ExecuteRequest { - tool_call_id: "call_1".to_string(), - enabled_tools: Vec::new(), - source: "21 * 2;".to_string(), - yield_time_ms: None, - max_output_tokens: None, - }) + .execute( + ExecuteRequest { + tool_call_id: "call_1".to_string(), + enabled_tools: Vec::new(), + source: "21 * 2;".to_string(), + yield_time_ms: None, + max_output_tokens: None, + }, + Arc::new(NoopCodeModeSessionDelegate), + ) .await .expect("start code-mode cell"); let cell_id = started.cell_id.clone(); diff --git a/codex-rs/code-mode/src/grpc_session/callbacks.rs b/codex-rs/code-mode/src/grpc_session/callbacks.rs index 8f84fbbd64..21f74cec0c 100644 --- a/codex-rs/code-mode/src/grpc_session/callbacks.rs +++ b/codex-rs/code-mode/src/grpc_session/callbacks.rs @@ -133,12 +133,12 @@ impl SessionInner { } let invocation_id = call.invocation_id.clone(); let cancellation = match admission { - CallbackAdmission::Active(cancellation) => Ok(cancellation), + CallbackAdmission::Active(cancellation, delegate) => Ok((cancellation, delegate)), CallbackAdmission::Cancelled => return Ok(()), CallbackAdmission::Closed => Err(format!("code-mode cell {} is closed", call.cell_id)), CallbackAdmission::Rejected(error) => Err(error), }; - let cancellation = match cancellation { + let (cancellation, delegate) = match cancellation { Ok(cancellation) => cancellation, Err(error) => { let inner = Arc::clone(self); @@ -157,19 +157,18 @@ impl SessionInner { let result = match invocation { Ok(invocation) => { let callback = AssertUnwindSafe(async { - inner - .delegate + delegate .invoke_tool(invocation, cancellation.child_token()) .await }) .catch_unwind(); tokio::select! { - biased; - _ = cancellation.cancelled() => return, - result = callback => match result { - Ok(result) => result, - Err(_) => Err("code-mode tool delegate panicked".to_string()), - }, + biased; + _ = cancellation.cancelled() => return, + result = callback => match result { + Ok(result) => result, + Err(_) => Err("code-mode tool delegate panicked".to_string()), + }, } } Err(error) => Err(error), @@ -223,8 +222,8 @@ impl SessionInner { .lock() .unwrap_or_else(PoisonError::into_inner) .admit_notification(¬ification)?; - let cancellation = match admission { - CallbackAdmission::Active(cancellation) => cancellation, + let (cancellation, delegate) = match admission { + CallbackAdmission::Active(cancellation, delegate) => (cancellation, delegate), CallbackAdmission::Cancelled | CallbackAdmission::Closed => return Ok(()), CallbackAdmission::Rejected(error) => { warn!("code-mode notification was dropped: {error}"); @@ -237,8 +236,7 @@ impl SessionInner { // them without waiting for arbitrary delegate work to complete. tokio::spawn(async move { let callback = AssertUnwindSafe(async { - inner - .delegate + delegate .notify( notification.call_id, CellId::new(notification.cell_id), diff --git a/codex-rs/code-mode/src/grpc_session/mod.rs b/codex-rs/code-mode/src/grpc_session/mod.rs index 38b3635818..1905434f3d 100644 --- a/codex-rs/code-mode/src/grpc_session/mod.rs +++ b/codex-rs/code-mode/src/grpc_session/mod.rs @@ -30,6 +30,7 @@ use tonic::transport::Channel; use tracing::Instrument; use self::operations::WaitSlot; +use self::state::ClosedCell; use self::state::SessionState; use self::transport::GrpcTransport; use self::transport::SharedTransport; @@ -96,7 +97,6 @@ impl GrpcCodeModeSessionProvider { #[tracing::instrument(name = "code_mode.grpc.open_binding", level = "info", skip_all)] async fn open_binding( &self, - delegate: Arc, limits: CodeModeSessionCellExecutionLimits, ) -> Result, String> { let mut client = deadline::startup("transport connection", self.transport.client()).await?; @@ -136,7 +136,6 @@ impl GrpcCodeModeSessionProvider { let inner = Arc::new(SessionInner { id: opened.session_id, client, - delegate, runtime: tokio::runtime::Handle::current(), state: Mutex::new(SessionState::default()), wait_slots: Mutex::new(HashMap::new()), @@ -180,24 +179,16 @@ impl GrpcCodeModeSessionProvider { } impl CodeModeSessionProvider for GrpcCodeModeSessionProvider { - fn create_session<'a>( - &'a self, - delegate: Arc, - ) -> CodeModeSessionProviderFuture<'a> { - self.create_session_with_limits(delegate, CodeModeSessionCellExecutionLimits::default()) + fn create_session(&self) -> CodeModeSessionProviderFuture<'_> { + self.create_session_with_limits(CodeModeSessionCellExecutionLimits::default()) } fn create_session_with_limits<'a>( &'a self, - delegate: Arc, limits: CodeModeSessionCellExecutionLimits, ) -> CodeModeSessionProviderFuture<'a> { Box::pin(async move { - let session = Arc::new(reconnect::ReconnectableSession::new( - self.clone(), - delegate, - limits, - )); + let session = Arc::new(reconnect::ReconnectableSession::new(self.clone(), limits)); session.initialize().await?; Ok(session as _) }) @@ -229,8 +220,9 @@ impl CodeModeSession for GrpcCodeModeSession { fn execute<'a>( &'a self, request: ExecuteRequest, + delegate: Arc, ) -> CodeModeSessionResultFuture<'a, StartedCell> { - Box::pin(self.inner.execute(request)) + Box::pin(self.inner.execute(request, delegate)) } fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome> { @@ -255,7 +247,6 @@ impl Drop for GrpcCodeModeSession { pub(super) struct SessionInner { pub(super) id: String, pub(super) client: GrpcClient, - pub(super) delegate: Arc, runtime: tokio::runtime::Handle, state: Mutex, wait_slots: Mutex>>, @@ -282,14 +273,14 @@ impl SessionInner { Ok(()) } - pub(super) fn report_closed_cell(&self, cell_id: Option) { - if let Some(cell_id) = cell_id { + pub(super) fn report_closed_cell(&self, cell: Option) { + if let Some((cell_id, delegate)) = cell { self.wait_slots .lock() .unwrap_or_else(PoisonError::into_inner) .remove(&cell_id); let _ = std::panic::catch_unwind(AssertUnwindSafe(|| { - self.delegate.cell_closed(&cell_id); + delegate.cell_closed(&cell_id); })); } } diff --git a/codex-rs/code-mode/src/grpc_session/operations.rs b/codex-rs/code-mode/src/grpc_session/operations.rs index ea29cc4788..727eb6ad3a 100644 --- a/codex-rs/code-mode/src/grpc_session/operations.rs +++ b/codex-rs/code-mode/src/grpc_session/operations.rs @@ -5,6 +5,7 @@ use std::sync::atomic::Ordering; use std::time::Duration; use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CodeModeSessionDelegate; use codex_code_mode_protocol::DEFAULT_EXEC_YIELD_TIME_MS; use codex_code_mode_protocol::ExecuteRequest; use codex_code_mode_protocol::RuntimeResponse; @@ -45,10 +46,11 @@ impl Drop for ExecutionOwnership { .lock() .unwrap_or_else(PoisonError::into_inner) .remove_execution(&self.execution_id); - let Some(cell_id) = cell else { + let Some((cell_id, delegate)) = cell else { return; }; - self.session.report_closed_cell(Some(cell_id.clone())); + self.session + .report_closed_cell(Some((cell_id.clone(), delegate))); if self.session.stopped.is_cancelled() { return; } @@ -67,6 +69,7 @@ impl SessionInner { pub(super) async fn execute( self: &Arc, request: ExecuteRequest, + delegate: Arc, ) -> Result { self.require_open()?; let execution_id = Uuid::new_v4().to_string(); @@ -81,7 +84,7 @@ impl SessionInner { self.state .lock() .unwrap_or_else(PoisonError::into_inner) - .begin_execution(&request)?; + .begin_execution(&request, delegate)?; let ownership = ExecutionOwnership { session: Arc::clone(self), execution_id, diff --git a/codex-rs/code-mode/src/grpc_session/reconnect.rs b/codex-rs/code-mode/src/grpc_session/reconnect.rs index e901fde6f4..dee623d265 100644 --- a/codex-rs/code-mode/src/grpc_session/reconnect.rs +++ b/codex-rs/code-mode/src/grpc_session/reconnect.rs @@ -31,7 +31,6 @@ pub(super) struct ReconnectableSession { struct ReconnectInner { provider: GrpcCodeModeSessionProvider, - delegate: Arc, limits: CodeModeSessionCellExecutionLimits, binding: Mutex>, opening_permit: Semaphore, @@ -49,13 +48,11 @@ struct SessionBinding { impl ReconnectableSession { pub(super) fn new( provider: GrpcCodeModeSessionProvider, - delegate: Arc, limits: CodeModeSessionCellExecutionLimits, ) -> Self { Self { inner: Arc::new(ReconnectInner { provider, - delegate, limits, binding: Mutex::new(None), opening_permit: Semaphore::new(/*permits*/ 1), @@ -75,10 +72,15 @@ impl CodeModeSession for ReconnectableSession { fn execute<'a>( &'a self, request: ExecuteRequest, + delegate: Arc, ) -> CodeModeSessionResultFuture<'a, StartedCell> { Box::pin(async move { let binding = self.inner.get_or_open_binding().await?; - let started = binding.session.execute(request).await?; + let delegate = Arc::new(GenerationDelegate { + delegate, + generation: binding.generation, + }); + let started = binding.session.execute(request, delegate).await?; Ok(generation::public_started_cell(binding.generation, started)) }) } @@ -151,16 +153,12 @@ impl ReconnectInner { } let generation = self.next_generation.fetch_add(1, Ordering::Relaxed); - let delegate = Arc::new(GenerationDelegate { - delegate: Arc::clone(&self.delegate), - generation, - }); let session = tokio::select! { biased; _ = self.shutdown_requested.cancelled() => { return Err(SHUTDOWN_ERROR.to_string()); } - session = self.provider.open_binding(delegate, self.limits.clone()) => session?, + session = self.provider.open_binding(self.limits.clone()) => session?, }; let binding = SessionBinding { session, diff --git a/codex-rs/code-mode/src/grpc_session/state.rs b/codex-rs/code-mode/src/grpc_session/state.rs index 69e1d7ad58..0fbb82ace0 100644 --- a/codex-rs/code-mode/src/grpc_session/state.rs +++ b/codex-rs/code-mode/src/grpc_session/state.rs @@ -1,6 +1,9 @@ use std::collections::HashMap; use std::collections::HashSet; use std::collections::VecDeque; +use std::sync::Arc; + +use codex_code_mode_protocol::CodeModeSessionDelegate; use codex_code_mode_protocol::CellId; use codex_code_mode_protocol::grpc; @@ -16,15 +19,17 @@ struct ActiveCallback { cancellation: CancellationToken, } +pub(super) type ClosedCell = (CellId, Arc); + pub(super) enum CallbackAdmission { - Active(CancellationToken), + Active(CancellationToken, Arc), Cancelled, Closed, Rejected(String), } -#[derive(Default)] struct ExecutionRecord { + delegate: Arc, cell_id: Option, tool_call_id: String, enabled_tools: HashMap, @@ -101,7 +106,11 @@ impl SessionState { Ok(()) } - pub(super) fn begin_execution(&mut self, request: &grpc::ExecuteRequest) -> Result<(), String> { + pub(super) fn begin_execution( + &mut self, + request: &grpc::ExecuteRequest, + delegate: Arc, + ) -> Result<(), String> { self.require_open()?; if request.execution_id.is_empty() || self.executions.contains_key(&request.execution_id) { return Err("code-mode execution ID was empty or reused".to_string()); @@ -127,7 +136,13 @@ impl SessionState { ExecutionRecord { tool_call_id: request.tool_call_id.clone(), enabled_tools, - ..ExecutionRecord::default() + delegate, + cell_id: None, + started: false, + ready: false, + closed: false, + notifications: 0, + cancellation: CancellationToken::new(), }, ); Ok(()) @@ -155,7 +170,7 @@ impl SessionState { pub(super) fn mark_execution_ready( &mut self, execution_id: &str, - ) -> Result, String> { + ) -> Result, String> { self.require_open()?; let execution = self .executions @@ -224,7 +239,10 @@ impl SessionState { cancellation: cancellation.clone(), }, ); - Ok(CallbackAdmission::Active(cancellation)) + Ok(CallbackAdmission::Active( + cancellation, + Arc::clone(&execution.delegate), + )) } pub(super) fn admit_notification( @@ -255,10 +273,11 @@ impl SessionState { self.notifications += 1; Ok(CallbackAdmission::Active( execution.cancellation.child_token(), + Arc::clone(&execution.delegate), )) } - pub(super) fn finish_notification(&mut self, execution_id: &str) -> Option { + pub(super) fn finish_notification(&mut self, execution_id: &str) -> Option { let execution = self.executions.get_mut(execution_id)?; execution.notifications = execution.notifications.checked_sub(1)?; self.notifications -= 1; @@ -296,7 +315,7 @@ impl SessionState { pub(super) fn close_cell( &mut self, closed: grpc::CellClosed, - ) -> Result, String> { + ) -> Result, String> { self.require_open()?; self.check_cell_ownership(&closed.execution_id, &closed.cell_id)?; let Some(execution) = self.executions.get_mut(&closed.execution_id) else { @@ -317,7 +336,7 @@ impl SessionState { Ok(self.close_execution_if_ready(&closed.execution_id)) } - pub(super) fn close(&mut self, failure: Option) -> Vec { + pub(super) fn close(&mut self, failure: Option) -> Vec { if self.closed { return Vec::new(); } @@ -331,12 +350,14 @@ impl SessionState { .drain() .filter_map(|(_, execution)| { execution.cancellation.cancel(); - execution.cell_id + execution + .cell_id + .map(|cell_id| (cell_id, execution.delegate)) }) .collect() } - fn close_execution_if_ready(&mut self, execution_id: &str) -> Option { + fn close_execution_if_ready(&mut self, execution_id: &str) -> Option { self.executions .get(execution_id) .is_some_and(|execution| { @@ -349,12 +370,14 @@ impl SessionState { .flatten() } - pub(super) fn remove_execution(&mut self, execution_id: &str) -> Option { + pub(super) fn remove_execution(&mut self, execution_id: &str) -> Option { let execution = self.executions.remove(execution_id)?; self.notifications -= execution.notifications; execution.cancellation.cancel(); self.revoke_execution_callbacks(execution_id); - execution.cell_id + execution + .cell_id + .map(|cell_id| (cell_id, execution.delegate)) } fn check_cell_ownership(&self, execution_id: &str, cell_id: &str) -> Result<(), String> { diff --git a/codex-rs/code-mode/src/grpc_session/state_tests.rs b/codex-rs/code-mode/src/grpc_session/state_tests.rs index 5a146304c4..1b7ab62dc8 100644 --- a/codex-rs/code-mode/src/grpc_session/state_tests.rs +++ b/codex-rs/code-mode/src/grpc_session/state_tests.rs @@ -1,7 +1,9 @@ use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::NoopCodeModeSessionDelegate; use codex_code_mode_protocol::grpc; use codex_code_mode_protocol::host::MAX_PENDING_DELEGATE_CALLS; use pretty_assertions::assert_eq; +use std::sync::Arc; use uuid::Uuid; use super::CallbackAdmission; @@ -58,33 +60,66 @@ fn notification(execution_id: &str, notification_id: u128) -> grpc::Notification } } +#[test] +fn callbacks_use_their_execution_delegate_before_cell_admission() { + let mut state = SessionState::default(); + let delegates: Vec> = vec![ + Arc::new(NoopCodeModeSessionDelegate), + Arc::new(NoopCodeModeSessionDelegate), + ]; + for (index, delegate) in delegates.iter().enumerate() { + state + .begin_execution(&request(&index.to_string()), delegate.clone()) + .unwrap(); + } + for (index, delegate) in delegates.iter().enumerate() { + let mut call = tool_call(&index.to_string(), index as u128 + 1); + call.cell_id = format!("cell-{index}"); + let CallbackAdmission::Active(_, actual) = state.admit_invocation(&call).unwrap() else { + panic!("callback should be admitted before the execute response"); + }; + assert!(Arc::ptr_eq(&actual, delegate)); + let mut notice = notification(&index.to_string(), index as u128 + 1); + notice.cell_id = call.cell_id; + let CallbackAdmission::Active(_, actual) = state.admit_notification(¬ice).unwrap() + else { + panic!("notification should be admitted before the execute response"); + }; + assert!(Arc::ptr_eq(&actual, delegate)); + } + drop(state.close(/*failure*/ None)); + for delegate in delegates { + assert_eq!(Arc::strong_count(&delegate), 1); + } +} + #[test] fn cell_closure_drains_notifications_and_cancels_tool_callbacks() { let mut state = SessionState::default(); state - .begin_execution(&request("execution")) + .begin_execution(&request("execution"), Arc::new(NoopCodeModeSessionDelegate)) .expect("register execution"); let first = state .admit_invocation(&tool_call("execution", /*invocation_id*/ 1)) .expect("accept early invocation"); - let CallbackAdmission::Active(first_cancellation) = first else { + let CallbackAdmission::Active(first_cancellation, _) = first else { panic!("first callback was not admitted"); }; - let CallbackAdmission::Active(notification_cancellation) = state + let CallbackAdmission::Active(notification_cancellation, _) = state .admit_notification(¬ification("execution", /*notification_id*/ 1)) .expect("admit notification") else { panic!("notification was not admitted"); }; - assert_eq!( + assert!( state .close_cell(grpc::CellClosed { execution_id: "execution".to_string(), cell_id: "cell".to_string(), final_tool_call_sequence: 3, }) - .expect("record cell closure"), - None + .expect("record cell closure") + .is_none() ); assert!(first_cancellation.is_cancelled()); assert!(!notification_cancellation.is_cancelled()); @@ -97,14 +132,16 @@ fn cell_closure_drains_notifications_and_cancels_tool_callbacks() { .expect("reject invocation for a closed cell"), CallbackAdmission::Closed )); - assert_eq!( + assert!( state .mark_execution_ready("execution") - .expect("claim started cell"), - None + .expect("claim started cell") + .is_none() ); assert_eq!( - state.finish_notification("execution"), + state + .finish_notification("execution") + .map(|(cell_id, _)| cell_id), Some(CellId::new("cell".to_string())) ); assert!(notification_cancellation.is_cancelled()); @@ -114,18 +151,18 @@ fn cell_closure_drains_notifications_and_cancels_tool_callbacks() { fn cell_closure_waits_until_the_started_cell_is_claimed() { let mut state = SessionState::default(); state - .begin_execution(&request("execution")) + .begin_execution(&request("execution"), Arc::new(NoopCodeModeSessionDelegate)) .expect("register execution"); - assert_eq!( + assert!( state .close_cell(grpc::CellClosed { execution_id: "execution".to_string(), cell_id: "cell".to_string(), final_tool_call_sequence: 0, }) - .expect("record early cell closure"), - None + .expect("record early cell closure") + .is_none() ); state .admit_execution("execution", "cell") @@ -133,6 +170,7 @@ fn cell_closure_waits_until_the_started_cell_is_claimed() { assert_eq!( state .mark_execution_ready("execution") + .map(|cell| cell.map(|(cell_id, _)| cell_id)) .expect("claim started cell"), Some(CellId::new("cell".to_string())) ); @@ -142,7 +180,7 @@ fn cell_closure_waits_until_the_started_cell_is_claimed() { fn oversized_cell_ids_are_rejected_before_admission() { let mut state = SessionState::default(); state - .begin_execution(&request("execution")) + .begin_execution(&request("execution"), Arc::new(NoopCodeModeSessionDelegate)) .expect("register execution"); assert_eq!( @@ -152,17 +190,17 @@ fn oversized_cell_ids_are_rejected_before_admission() { grpc::MAX_IDENTIFIER_BYTES )) ); - assert_eq!(state.remove_execution("execution"), None); + assert!(state.remove_execution("execution").is_none()); } #[test] fn callbacks_cannot_claim_another_executions_cell() { let mut state = SessionState::default(); state - .begin_execution(&request("first")) + .begin_execution(&request("first"), Arc::new(NoopCodeModeSessionDelegate)) .expect("register first execution"); state - .begin_execution(&request("second")) + .begin_execution(&request("second"), Arc::new(NoopCodeModeSessionDelegate)) .expect("register second execution"); state .admit_invocation(&tool_call("first", /*invocation_id*/ 1)) @@ -198,19 +236,19 @@ fn callbacks_cannot_claim_another_executions_cell() { fn abandonment_before_start_ignores_later_cell_closure() { let mut state = SessionState::default(); state - .begin_execution(&request("execution")) + .begin_execution(&request("execution"), Arc::new(NoopCodeModeSessionDelegate)) .expect("register execution"); - assert_eq!(state.remove_execution("execution"), None); - assert_eq!( + assert!(state.remove_execution("execution").is_none()); + assert!( state .close_cell(grpc::CellClosed { execution_id: "execution".to_string(), cell_id: "cell".to_string(), final_tool_call_sequence: 0, }) - .expect("ignore closure for abandoned execution"), - None + .expect("ignore closure for abandoned execution") + .is_none() ); assert!(state.close(/*failure*/ None).is_empty()); } @@ -219,16 +257,18 @@ fn abandonment_before_start_ignores_later_cell_closure() { fn abandonment_revokes_callbacks_and_ignores_late_events() { let mut state = SessionState::default(); state - .begin_execution(&request("execution")) + .begin_execution(&request("execution"), Arc::new(NoopCodeModeSessionDelegate)) .expect("register execution"); - let CallbackAdmission::Active(invocation) = state + let CallbackAdmission::Active(invocation, _) = state .admit_invocation(&tool_call("execution", /*invocation_id*/ 1)) .expect("admit callback before execution starts") else { panic!("invocation was not admitted"); }; assert_eq!( - state.remove_execution("execution"), + state + .remove_execution("execution") + .map(|(cell_id, _)| cell_id), Some(CellId::new("cell".to_string())) ); assert!(invocation.is_cancelled()); @@ -238,15 +278,15 @@ fn abandonment_revokes_callbacks_and_ignores_late_events() { .expect("reject delayed tool invocation"), CallbackAdmission::Closed )); - assert_eq!( + assert!( state .close_cell(grpc::CellClosed { execution_id: "execution".to_string(), cell_id: "cell".to_string(), final_tool_call_sequence: 1, }) - .expect("ignore delayed cell closure"), - None + .expect("ignore delayed cell closure") + .is_none() ); assert!(state.close(/*failure*/ None).is_empty()); } @@ -255,7 +295,7 @@ fn abandonment_revokes_callbacks_and_ignores_late_events() { fn invocation_cancellation_revokes_delegate_and_late_completion() { let mut state = SessionState::default(); state - .begin_execution(&request("execution")) + .begin_execution(&request("execution"), Arc::new(NoopCodeModeSessionDelegate)) .expect("register execution"); state .admit_execution("execution", "cell") @@ -264,7 +304,7 @@ fn invocation_cancellation_revokes_delegate_and_late_completion() { .mark_execution_ready("execution") .expect("claim execution"); let invocation = tool_call("execution", /*invocation_id*/ 1); - let CallbackAdmission::Active(cancellation) = state + let CallbackAdmission::Active(cancellation, _) = state .admit_invocation(&invocation) .expect("accept invocation") else { @@ -284,6 +324,7 @@ fn invocation_cancellation_revokes_delegate_and_late_completion() { cell_id: "cell".to_string(), final_tool_call_sequence: 1, }) + .map(|cell| cell.map(|(cell_id, _)| cell_id)) .expect("close cell"), Some(CellId::new("cell".to_string())) ); @@ -293,7 +334,7 @@ fn invocation_cancellation_revokes_delegate_and_late_completion() { fn duplicate_invocation_ids_are_rejected() { let mut state = SessionState::default(); state - .begin_execution(&request("execution")) + .begin_execution(&request("execution"), Arc::new(NoopCodeModeSessionDelegate)) .expect("register execution"); state .admit_execution("execution", "cell") @@ -311,7 +352,7 @@ fn duplicate_invocation_ids_are_rejected() { fn tool_callbacks_must_match_the_executions_enabled_tools() { let mut state = SessionState::default(); state - .begin_execution(&request("execution")) + .begin_execution(&request("execution"), Arc::new(NoopCodeModeSessionDelegate)) .expect("register execution"); let mut disabled = tool_call("execution", /*invocation_id*/ 1); @@ -349,7 +390,7 @@ fn tool_callbacks_must_match_the_executions_enabled_tools() { }); assert!(matches!( state.admit_invocation(&explicit_default_namespace), - Ok(CallbackAdmission::Active(_)) + Ok(CallbackAdmission::Active(_, _)) )); assert_eq!(state.require_open(), Ok(())); } @@ -361,20 +402,20 @@ fn execution_call_ids_must_be_bounded() { oversized.tool_call_id = "x".repeat(grpc::MAX_IDENTIFIER_BYTES + 1); assert_eq!( - state.begin_execution(&oversized), + state.begin_execution(&oversized, Arc::new(NoopCodeModeSessionDelegate)), Err(format!( "gRPC code-mode host returned tool call ID exceeding {} bytes", grpc::MAX_IDENTIFIER_BYTES )) ); - assert_eq!(state.remove_execution("execution"), None); + assert!(state.remove_execution("execution").is_none()); } #[test] fn notification_call_ids_must_match_their_execution() { let mut state = SessionState::default(); state - .begin_execution(&request("execution")) + .begin_execution(&request("execution"), Arc::new(NoopCodeModeSessionDelegate)) .expect("register execution"); let mut oversized = notification("execution", /*notification_id*/ 1); @@ -395,7 +436,7 @@ fn notification_call_ids_must_match_their_execution() { ); assert!(matches!( state.admit_notification(¬ification("execution", /*notification_id*/ 3)), - Ok(CallbackAdmission::Active(_)) + Ok(CallbackAdmission::Active(_, _)) )); } @@ -403,7 +444,7 @@ fn notification_call_ids_must_match_their_execution() { fn malformed_callback_ids_are_rejected_before_retention() { let mut state = SessionState::default(); state - .begin_execution(&request("execution")) + .begin_execution(&request("execution"), Arc::new(NoopCodeModeSessionDelegate)) .expect("register execution"); let mut invalid_invocation = tool_call("execution", /*invocation_id*/ 1); invalid_invocation.invocation_id = "not-a-uuid".to_string(); @@ -442,13 +483,13 @@ fn malformed_callback_ids_are_rejected_before_retention() { fn notifications_and_tools_share_the_pending_delegate_limit() { let mut state = SessionState::default(); state - .begin_execution(&request("execution")) + .begin_execution(&request("execution"), Arc::new(NoopCodeModeSessionDelegate)) .expect("register execution"); for index in 0..MAX_PENDING_DELEGATE_CALLS { assert!(matches!( state.admit_notification(¬ification("execution", index as u128 + 1)), - Ok(CallbackAdmission::Active(_)) + Ok(CallbackAdmission::Active(_, _)) )); } @@ -463,10 +504,10 @@ fn notifications_and_tools_share_the_pending_delegate_limit() { if error == "code-mode host exceeded its pending delegate callback limit" )); assert_eq!(state.require_open(), Ok(())); - assert_eq!(state.finish_notification("execution"), None); + assert!(state.finish_notification("execution").is_none()); assert!(matches!( state.admit_invocation(&tool_call("execution", /*invocation_id*/ 2)), - Ok(CallbackAdmission::Active(_)) + Ok(CallbackAdmission::Active(_, _)) )); } @@ -474,9 +515,9 @@ fn notifications_and_tools_share_the_pending_delegate_limit() { fn terminated_cells_cancel_pending_notifications() { let mut state = SessionState::default(); state - .begin_execution(&request("execution")) + .begin_execution(&request("execution"), Arc::new(NoopCodeModeSessionDelegate)) .expect("register execution"); - let CallbackAdmission::Active(cancellation) = state + let CallbackAdmission::Active(cancellation, _) = state .admit_notification(¬ification("execution", /*notification_id*/ 1)) .expect("admit notification") else { @@ -486,25 +527,25 @@ fn terminated_cells_cancel_pending_notifications() { state.cancel_notifications(&CellId::new("cell".to_string())); assert!(cancellation.is_cancelled()); - assert_eq!(state.finish_notification("execution"), None); + assert!(state.finish_notification("execution").is_none()); } #[test] fn disconnect_revokes_callbacks_and_returns_each_live_cell_once() { let mut state = SessionState::default(); state - .begin_execution(&request("execution")) + .begin_execution(&request("execution"), Arc::new(NoopCodeModeSessionDelegate)) .expect("register execution"); state .admit_execution("execution", "cell") .expect("admit execution"); - let CallbackAdmission::Active(cancellation) = state + let CallbackAdmission::Active(cancellation, _) = state .admit_invocation(&tool_call("execution", /*invocation_id*/ 1)) .expect("accept invocation") else { panic!("invocation was not admitted"); }; - let CallbackAdmission::Active(notification_cancellation) = state + let CallbackAdmission::Active(notification_cancellation, _) = state .admit_notification(¬ification("execution", /*notification_id*/ 1)) .expect("admit notification") else { @@ -512,7 +553,11 @@ fn disconnect_revokes_callbacks_and_returns_each_live_cell_once() { }; assert_eq!( - state.close(Some("lease closed".to_string())), + state + .close(Some("lease closed".to_string())) + .into_iter() + .map(|(cell_id, _)| cell_id) + .collect::>(), vec![CellId::new("cell".to_string())] ); assert!(cancellation.is_cancelled()); diff --git a/codex-rs/code-mode/src/remote_session.rs b/codex-rs/code-mode/src/remote_session.rs index fb38e2899e..0aff7acbd0 100644 --- a/codex-rs/code-mode/src/remote_session.rs +++ b/codex-rs/code-mode/src/remote_session.rs @@ -26,7 +26,6 @@ use self::connection::Connection; use self::connection::ConnectionError; use self::connection::RemoteSession; use self::connection::SessionCleanup; -use crate::NoopCodeModeSessionDelegate; mod connection; @@ -73,19 +72,15 @@ impl CodeModeSessionProvider for ProcessOwnedCodeModeSessionProvider { } } - fn create_session<'a>( - &'a self, - delegate: Arc, - ) -> CodeModeSessionProviderFuture<'a> { - self.create_session_with_limits(delegate, CodeModeSessionCellExecutionLimits::default()) + fn create_session(&self) -> CodeModeSessionProviderFuture<'_> { + self.create_session_with_limits(CodeModeSessionCellExecutionLimits::default()) } fn create_session_with_limits<'a>( &'a self, - delegate: Arc, limits: CodeModeSessionCellExecutionLimits, ) -> CodeModeSessionProviderFuture<'a> { - Box::pin(create_host_session(delegate, self.process_host(), limits)) + Box::pin(create_host_session(self.process_host(), limits)) } } @@ -94,28 +89,23 @@ impl CodeModeSessionProvider for DisabledCodeModeSessionProvider { Err("code-mode host is disabled".to_string()) } - fn create_session<'a>( - &'a self, - _delegate: Arc, - ) -> CodeModeSessionProviderFuture<'a> { + fn create_session(&self) -> CodeModeSessionProviderFuture<'_> { Box::pin(async { Err("code-mode host is disabled".to_string()) }) } fn create_session_with_limits<'a>( &'a self, - delegate: Arc, _limits: CodeModeSessionCellExecutionLimits, ) -> CodeModeSessionProviderFuture<'a> { - self.create_session(delegate) + self.create_session() } } async fn create_host_session( - delegate: Arc, host: Arc, limits: CodeModeSessionCellExecutionLimits, ) -> Result, String> { - let session = ProcessOwnedCodeModeSession::with_host(delegate, host, limits); + let session = ProcessOwnedCodeModeSession::with_host(host, limits); session.connection().await?; Ok(Arc::new(session)) } @@ -223,7 +213,6 @@ struct SessionBinding { struct SessionInner { host: Arc, - delegate: Arc, limits: CodeModeSessionCellExecutionLimits, state: StdMutex, next_generation: AtomicU64, @@ -240,7 +229,6 @@ pub struct ProcessOwnedCodeModeSession { impl ProcessOwnedCodeModeSession { pub fn new() -> Self { Self::with_host( - Arc::new(NoopCodeModeSessionDelegate), Arc::new(OwnedCodeModeHost::new( InstallContext::current().code_mode_host_program(), )), @@ -248,15 +236,10 @@ impl ProcessOwnedCodeModeSession { ) } - fn with_host( - delegate: Arc, - host: Arc, - limits: CodeModeSessionCellExecutionLimits, - ) -> Self { + fn with_host(host: Arc, limits: CodeModeSessionCellExecutionLimits) -> Self { Self { inner: Arc::new(SessionInner { host, - delegate, limits, state: StdMutex::new(SessionState::New), next_generation: AtomicU64::new(1), @@ -271,9 +254,16 @@ impl ProcessOwnedCodeModeSession { self.inner.connection().await } - pub async fn execute(&self, request: ExecuteRequest) -> Result { + pub async fn execute( + &self, + request: ExecuteRequest, + delegate: Arc, + ) -> Result { let binding = self.connection().await?; - binding.connection.execute(binding.remote, request).await + binding + .connection + .execute(binding.remote, request, delegate) + .await } pub async fn wait(&self, request: WaitRequest) -> Result { @@ -348,11 +338,7 @@ impl SessionInner { let result = match self.host.connection().await { Ok(connection) => { let cleanup = connection - .open_session( - remote.clone(), - Arc::clone(&self.delegate), - self.limits.clone(), - ) + .open_session(remote.clone(), self.limits.clone()) .await; cleanup.map(|cleanup| SessionBinding { connection, @@ -532,8 +518,11 @@ impl CodeModeSession for ProcessOwnedCodeModeSession { fn execute<'a>( &'a self, request: ExecuteRequest, + delegate: Arc, ) -> CodeModeSessionResultFuture<'a, StartedCell> { - Box::pin(ProcessOwnedCodeModeSession::execute(self, request)) + Box::pin(ProcessOwnedCodeModeSession::execute( + self, request, delegate, + )) } fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome> { diff --git a/codex-rs/code-mode/src/remote_session/connection.rs b/codex-rs/code-mode/src/remote_session/connection.rs index 248b3b0171..eecf9c64ee 100644 --- a/codex-rs/code-mode/src/remote_session/connection.rs +++ b/codex-rs/code-mode/src/remote_session/connection.rs @@ -323,7 +323,6 @@ impl Connection { pub(super) async fn open_session( &self, session: RemoteSession, - delegate: Arc, limits: CodeModeSessionCellExecutionLimits, ) -> Result { if limits != CodeModeSessionCellExecutionLimits::default() @@ -341,7 +340,6 @@ impl Connection { let (response_tx, response_rx) = oneshot::channel(); self.send(DriverCommand::OpenSession { session, - delegate, limits, cleanup: cleanup.clone(), caller_cancellation: cancellation.token(), @@ -358,12 +356,14 @@ impl Connection { &self, session: RemoteSession, request: ExecuteRequest, + delegate: Arc, ) -> Result { let cancellation = CallerCancellation::new(); let (response_tx, response_rx) = oneshot::channel(); self.send(DriverCommand::Execute { session, request, + delegate, caller_cancellation: cancellation.token(), response_tx, }) diff --git a/codex-rs/code-mode/src/remote_session/connection/driver/commands.rs b/codex-rs/code-mode/src/remote_session/connection/driver/commands.rs index ce50e06881..ca6b5bd235 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver/commands.rs +++ b/codex-rs/code-mode/src/remote_session/connection/driver/commands.rs @@ -29,25 +29,18 @@ impl ConnectionDriver { match command { DriverCommand::OpenSession { session, - delegate, limits, cleanup, caller_cancellation, response_tx, - } => self.open_session( - session, - delegate, - limits, - cleanup, - caller_cancellation, - response_tx, - ), + } => self.open_session(session, limits, cleanup, caller_cancellation, response_tx), DriverCommand::Execute { session, request, + delegate, caller_cancellation, response_tx, - } => self.execute(session, request, caller_cancellation, response_tx), + } => self.execute(session, request, delegate, caller_cancellation, response_tx), DriverCommand::Wait { session, request, @@ -69,7 +62,6 @@ impl ConnectionDriver { fn open_session( &mut self, session: RemoteSession, - delegate: Arc, limits: CodeModeSessionCellExecutionLimits, cleanup: super::cleanup::SessionCleanup, caller_cancellation: CancellationToken, @@ -120,7 +112,6 @@ impl ConnectionDriver { request_id, PendingRequest::OpenSession { session, - delegate, cleanup, cancellation, response_tx, @@ -134,6 +125,7 @@ impl ConnectionDriver { &mut self, session: RemoteSession, request: ExecuteRequest, + delegate: Arc, caller_cancellation: CancellationToken, response_tx: oneshot::Sender>, ) -> bool { @@ -179,6 +171,7 @@ impl ConnectionDriver { request_id, PendingRequest::Execute { session, + delegate, response_tx, initial_response_tx, initial_response_rx, diff --git a/codex-rs/code-mode/src/remote_session/connection/driver/responses.rs b/codex-rs/code-mode/src/remote_session/connection/driver/responses.rs index 7f0067e5ed..aa48b7bada 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver/responses.rs +++ b/codex-rs/code-mode/src/remote_session/connection/driver/responses.rs @@ -153,15 +153,13 @@ impl ConnectionDriver { match pending { PendingRequest::OpenSession { session, - delegate, cleanup, cancellation, response_tx, } => match result { Ok(HostResponse::SessionReady { session_id }) if session_id == session.id => { let abandoned = cancellation.is_cancelled() || response_tx.is_closed(); - self.sessions - .insert_ready(session.clone(), delegate, cleanup); + self.sessions.insert_ready(session.clone(), cleanup); if abandoned || response_tx.send(Ok(())).is_err() { return self.shutdown_abandoned_session(session); } @@ -179,6 +177,7 @@ impl ConnectionDriver { }, PendingRequest::Execute { session, + delegate, response_tx, initial_response_tx, initial_response_rx, @@ -188,7 +187,7 @@ impl ConnectionDriver { // The host owns a checked, never-reused ID sequence. Retain only live // IDs so client memory scales with concurrency, not session lifetime. let remote_cell_id = cell_id.clone(); - let public_id = match self.sessions.admit_cell(&session, cell_id) { + let public_id = match self.sessions.admit_cell(&session, cell_id, delegate) { Ok(public_id) => public_id, Err(CellAdmissionError::MissingSession) => { let _ = response_tx diff --git a/codex-rs/code-mode/src/remote_session/connection/driver/session_registry.rs b/codex-rs/code-mode/src/remote_session/connection/driver/session_registry.rs index 8668baa27e..d3dda96b85 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver/session_registry.rs +++ b/codex-rs/code-mode/src/remote_session/connection/driver/session_registry.rs @@ -34,10 +34,9 @@ pub(super) enum CellAdmissionError { struct SessionRecord { remote: RemoteSession, - delegate: Arc, cleanup: SessionCleanup, phase: SessionPhase, - cells: HashMap, + cells: HashMap, } #[derive(Clone, Copy, Eq, PartialEq)] @@ -67,17 +66,11 @@ impl SessionRegistry { .is_some_and(|session| session.cells.contains_key(cell_id)) } - pub(super) fn insert_ready( - &mut self, - session: RemoteSession, - delegate: Arc, - cleanup: SessionCleanup, - ) { + pub(super) fn insert_ready(&mut self, session: RemoteSession, cleanup: SessionCleanup) { self.records.insert( session.id.clone(), SessionRecord { remote: session, - delegate, cleanup, phase: SessionPhase::Ready, cells: HashMap::new(), @@ -133,6 +126,7 @@ impl SessionRegistry { &mut self, session: &RemoteSession, cell_id: WireCellId, + delegate: Arc, ) -> Result { let Some(record) = self.records.get_mut(&session.id) else { return Err(CellAdmissionError::MissingSession); @@ -141,7 +135,14 @@ impl SessionRegistry { return Err(CellAdmissionError::DuplicateCell); } let public_id = public_cell_id(session.generation, &cell_id); - record.cells.insert(cell_id, public_id.clone()); + record.cells.insert( + cell_id, + CellOwner { + session_id: session.id.clone(), + cell_id: public_id.clone(), + delegate, + }, + ); Ok(public_id) } @@ -154,7 +155,7 @@ impl SessionRegistry { .records .get(session_id) .ok_or_else(|| format!("code-mode host delegated for unknown session {session_id}"))?; - let public_id = session.cells.get(cell_id).cloned().ok_or_else(|| { + let owner = session.cells.get(cell_id).ok_or_else(|| { format!( "code-mode host delegated for unknown cell {} in session {session_id}", cell_id.as_str() @@ -162,8 +163,8 @@ impl SessionRegistry { })?; Ok(DelegateTarget { session_id: session_id.clone(), - cell_id: public_id, - delegate: Arc::clone(&session.delegate), + cell_id: owner.cell_id.clone(), + delegate: Arc::clone(&owner.delegate), }) } @@ -178,46 +179,25 @@ impl SessionRegistry { cell_id.as_str() ) })?; - let public_id = session + session .cells .remove(cell_id) - .ok_or_else(|| format!("code-mode host closed unknown cell in session {session_id}"))?; - Ok(CellOwner { - session_id: session_id.clone(), - cell_id: public_id, - delegate: Arc::clone(&session.delegate), - }) + .ok_or_else(|| format!("code-mode host closed unknown cell in session {session_id}")) } pub(super) fn remove_session(&mut self, session_id: &SessionId) -> Vec { let Some(session) = self.records.remove(session_id) else { return Vec::new(); }; - session - .cells - .into_values() - .map(|cell_id| CellOwner { - session_id: session_id.clone(), - cell_id, - delegate: Arc::clone(&session.delegate), - }) - .collect() + session.cells.into_values().collect() } pub(super) fn drain(&mut self) -> Vec { let sessions = std::mem::take(&mut self.records); sessions - .into_iter() - .map(|(session_id, session)| { - let cells = session - .cells - .into_values() - .map(|cell_id| CellOwner { - session_id: session_id.clone(), - cell_id, - delegate: Arc::clone(&session.delegate), - }) - .collect(); + .into_values() + .map(|session| { + let cells = session.cells.into_values().collect(); FailedSession { cleanup: session.cleanup, cells, diff --git a/codex-rs/code-mode/src/remote_session/connection/driver/types.rs b/codex-rs/code-mode/src/remote_session/connection/driver/types.rs index 7b561b330a..19595786db 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver/types.rs +++ b/codex-rs/code-mode/src/remote_session/connection/driver/types.rs @@ -30,7 +30,6 @@ pub(in crate::remote_session) struct RemoteSession { pub(in crate::remote_session::connection) enum DriverCommand { OpenSession { session: RemoteSession, - delegate: Arc, limits: CodeModeSessionCellExecutionLimits, cleanup: SessionCleanup, caller_cancellation: CancellationToken, @@ -38,6 +37,7 @@ pub(in crate::remote_session::connection) enum DriverCommand { }, Execute { session: RemoteSession, + delegate: Arc, request: ExecuteRequest, caller_cancellation: CancellationToken, response_tx: oneshot::Sender>, @@ -136,13 +136,13 @@ pub(super) struct UnclaimedExecute { pub(super) enum PendingRequest { OpenSession { session: RemoteSession, - delegate: Arc, cleanup: SessionCleanup, cancellation: CancellableRequest, response_tx: oneshot::Sender>, }, Execute { session: RemoteSession, + delegate: Arc, response_tx: oneshot::Sender>, initial_response_tx: oneshot::Sender>, initial_response_rx: oneshot::Receiver>, diff --git a/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs b/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs index 01d474b0d6..6b3428b1af 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs +++ b/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs @@ -87,17 +87,13 @@ impl DriverHarness { } } - async fn open( - &mut self, - session: RemoteSession, - delegate: Arc, - ) -> SessionCleanup { + async fn open(&mut self, session: RemoteSession) -> SessionCleanup { let cleanup = SessionCleanup::new(); let (response_tx, response_rx) = oneshot::channel(); self.command_tx .send(DriverCommand::OpenSession { session: session.clone(), - delegate, + limits: Default::default(), cleanup: cleanup.clone(), caller_cancellation: CancellationToken::new(), @@ -129,11 +125,13 @@ impl DriverHarness { session: RemoteSession, request_id: i64, cell_id: &str, + delegate: Arc, ) -> codex_code_mode_protocol::StartedCell { let (response_tx, response_rx) = oneshot::channel(); self.command_tx .send(DriverCommand::Execute { session, + delegate, request: ExecuteRequest { tool_call_id: format!("call-{request_id}"), enabled_tools: Vec::new(), @@ -208,7 +206,7 @@ async fn open_session_includes_nondefault_cell_execution_limits() { .command_tx .send(DriverCommand::OpenSession { session: session.clone(), - delegate: Arc::new(RecordingDelegate::default()), + limits, cleanup: SessionCleanup::new(), caller_cancellation: CancellationToken::new(), @@ -388,18 +386,20 @@ async fn next_held_delegate_event( async fn deferred_delegates_follow_cell_readiness_and_cancellation() { let mut harness = DriverHarness::start(); let session = remote_session(); - let delegate = Arc::new(RecordingDelegate::default()); - harness.open(session.clone(), delegate.clone()).await; + let first_delegate = Arc::new(RecordingDelegate::default()); + let second_delegate = Arc::new(RecordingDelegate::default()); + harness.open(session.clone()).await; let (first_response_tx, first_response_rx) = oneshot::channel(); let (second_response_tx, second_response_rx) = oneshot::channel(); - for (tool_call_id, response_tx) in [ - ("first-cell", first_response_tx), - ("second-cell", second_response_tx), + for (tool_call_id, response_tx, delegate) in [ + ("first-cell", first_response_tx, first_delegate.clone()), + ("second-cell", second_response_tx, second_delegate.clone()), ] { harness .command_tx .send(DriverCommand::Execute { + delegate: delegate.clone(), session: session.clone(), request: ExecuteRequest { tool_call_id: tool_call_id.to_string(), @@ -481,7 +481,13 @@ async fn deferred_delegates_follow_cell_readiness_and_cancellation() { }, } ); - assert_eq!(delegate.notifications.load(Ordering::Relaxed), 1); + assert_eq!( + ( + first_delegate.notifications.load(Ordering::Relaxed), + second_delegate.notifications.load(Ordering::Relaxed) + ), + (0, 1) + ); harness .event_tx @@ -513,7 +519,13 @@ async fn deferred_delegates_follow_cell_readiness_and_cancellation() { }, } ); - assert_eq!(delegate.notifications.load(Ordering::Relaxed), 2); + assert_eq!( + ( + first_delegate.notifications.load(Ordering::Relaxed), + second_delegate.notifications.load(Ordering::Relaxed) + ), + (1, 1) + ); assert!(matches!( harness.outgoing_rx.try_recv(), Err(mpsc::error::TryRecvError::Empty) @@ -531,7 +543,7 @@ async fn dropped_open_waiter_shuts_down_committed_session() { .command_tx .send(DriverCommand::OpenSession { session: session.clone(), - delegate: Arc::new(RecordingDelegate::default()), + limits: Default::default(), cleanup, caller_cancellation: CancellationToken::new(), @@ -575,6 +587,7 @@ async fn dropped_open_waiter_shuts_down_committed_session() { harness .command_tx .send(DriverCommand::Execute { + delegate: Arc::new(RecordingDelegate::default()), session: session.clone(), request: ExecuteRequest { tool_call_id: "call-1".to_string(), @@ -603,9 +616,14 @@ async fn delegate_cancel_is_best_effort_and_sends_no_late_response() { let mut harness = DriverHarness::start(); let session = remote_session(); let delegate = Arc::new(RecordingDelegate::default()); - harness.open(session.clone(), delegate.clone()).await; + harness.open(session.clone()).await; let _started = harness - .start_cell(session.clone(), /*request_id*/ 2, "1") + .start_cell( + session.clone(), + /*request_id*/ 2, + "1", + delegate.clone(), + ) .await; let request_id = DelegateRequestId::new(/*value*/ 7); harness @@ -662,9 +680,14 @@ async fn delegate_limit_returns_an_error_without_disconnecting() { let mut harness = DriverHarness::start(); let session = remote_session(); let delegate = Arc::new(RecordingDelegate::default()); - harness.open(session.clone(), delegate.clone()).await; + harness.open(session.clone()).await; let _started = harness - .start_cell(session.clone(), /*request_id*/ 2, "1") + .start_cell( + session.clone(), + /*request_id*/ 2, + "1", + delegate.clone(), + ) .await; for value in 1..=MAX_PENDING_DELEGATE_CALLS { @@ -729,9 +752,14 @@ async fn terminate_closes_cell_without_waiting_for_delegate_cleanup() { let mut harness = DriverHarness::start(); let session = remote_session(); let (delegate, mut events_rx, release) = HeldDelegate::new(); - harness.open(session.clone(), delegate).await; + harness.open(session.clone()).await; let _started = harness - .start_cell(session.clone(), /*request_id*/ 2, "1") + .start_cell( + session.clone(), + /*request_id*/ 2, + "1", + delegate.clone(), + ) .await; let delegate_id = DelegateRequestId::new(/*value*/ 7); harness.start_tool_delegate(&session, delegate_id).await; @@ -825,9 +853,14 @@ async fn shutdown_closes_cell_without_waiting_for_delegate_cleanup() { let mut harness = DriverHarness::start(); let session = remote_session(); let (delegate, mut events_rx, release) = HeldDelegate::new(); - harness.open(session.clone(), delegate).await; + harness.open(session.clone()).await; let _started = harness - .start_cell(session.clone(), /*request_id*/ 2, "1") + .start_cell( + session.clone(), + /*request_id*/ 2, + "1", + delegate.clone(), + ) .await; let delegate_id = DelegateRequestId::new(/*value*/ 7); harness.start_tool_delegate(&session, delegate_id).await; @@ -907,9 +940,14 @@ async fn completed_delegate_request_id_cannot_be_reused() { let mut harness = DriverHarness::start(); let session = remote_session(); let delegate = Arc::new(RecordingDelegate::default()); - harness.open(session.clone(), delegate.clone()).await; + harness.open(session.clone()).await; let _started = harness - .start_cell(session.clone(), /*request_id*/ 2, "1") + .start_cell( + session.clone(), + /*request_id*/ 2, + "1", + delegate.clone(), + ) .await; let request_id = DelegateRequestId::new(/*value*/ 7); let request = || DelegateRequest::Notify { @@ -950,11 +988,14 @@ async fn completed_delegate_request_id_cannot_be_reused() { async fn delegate_task_panic_becomes_tool_error_without_killing_connection() { let mut harness = DriverHarness::start(); let session = remote_session(); - harness - .open(session.clone(), Arc::new(PanickingDelegate)) - .await; + harness.open(session.clone()).await; let _started = harness - .start_cell(session.clone(), /*request_id*/ 2, "1") + .start_cell( + session.clone(), + /*request_id*/ 2, + "1", + Arc::new(PanickingDelegate), + ) .await; harness .event_tx @@ -994,7 +1035,15 @@ async fn delegate_for_unknown_cell_returns_error_without_invocation() { let mut harness = DriverHarness::start(); let session = remote_session(); let delegate = Arc::new(RecordingDelegate::default()); - harness.open(session.clone(), delegate.clone()).await; + harness.open(session.clone()).await; + let _started = harness + .start_cell( + session.clone(), + /*request_id*/ 2, + "known", + delegate.clone(), + ) + .await; let id = DelegateRequestId::new(/*value*/ 7); harness @@ -1039,9 +1088,14 @@ async fn delegate_after_cell_close_returns_error_without_invocation() { let mut harness = DriverHarness::start(); let session = remote_session(); let delegate = Arc::new(RecordingDelegate::default()); - harness.open(session.clone(), delegate.clone()).await; + harness.open(session.clone()).await; let _started = harness - .start_cell(session.clone(), /*request_id*/ 2, "1") + .start_cell( + session.clone(), + /*request_id*/ 2, + "1", + delegate.clone(), + ) .await; harness .event_tx @@ -1090,8 +1144,10 @@ async fn mismatched_initial_response_fails_connection_and_closes_cell_once() { let mut harness = DriverHarness::start(); let session = remote_session(); let delegate = Arc::new(RecordingDelegate::default()); - harness.open(session.clone(), delegate.clone()).await; - let started = harness.start_cell(session, /*request_id*/ 2, "1").await; + harness.open(session.clone()).await; + let started = harness + .start_cell(session, /*request_id*/ 2, "1", delegate.clone()) + .await; harness .event_tx .send(DriverEvent::HostMessage(HostToClient::InitialResponse { @@ -1120,9 +1176,14 @@ async fn mismatched_wait_response_fails_connection() { let mut harness = DriverHarness::start(); let session = remote_session(); let delegate = Arc::new(RecordingDelegate::default()); - harness.open(session.clone(), delegate.clone()).await; + harness.open(session.clone()).await; let _started = harness - .start_cell(session.clone(), /*request_id*/ 2, "1") + .start_cell( + session.clone(), + /*request_id*/ 2, + "1", + delegate.clone(), + ) .await; let (response_tx, response_rx) = oneshot::channel(); harness @@ -1169,9 +1230,14 @@ async fn mismatched_terminate_response_fails_connection() { let mut harness = DriverHarness::start(); let session = remote_session(); let delegate = Arc::new(RecordingDelegate::default()); - harness.open(session.clone(), delegate.clone()).await; + harness.open(session.clone()).await; let _started = harness - .start_cell(session.clone(), /*request_id*/ 2, "1") + .start_cell( + session.clone(), + /*request_id*/ 2, + "1", + delegate.clone(), + ) .await; let (response_tx, response_rx) = oneshot::channel(); harness @@ -1213,11 +1279,14 @@ async fn mismatched_terminate_response_fails_connection() { async fn remote_wait_accepts_durations_longer_than_five_minutes() { let mut harness = DriverHarness::start(); let session = remote_session(); - harness - .open(session.clone(), Arc::new(RecordingDelegate::default())) - .await; + harness.open(session.clone()).await; let _started = harness - .start_cell(session.clone(), /*request_id*/ 2, "1") + .start_cell( + session.clone(), + /*request_id*/ 2, + "1", + Arc::new(RecordingDelegate::default()), + ) .await; let (response_tx, response_rx) = oneshot::channel(); harness @@ -1271,11 +1340,14 @@ async fn remote_wait_accepts_durations_longer_than_five_minutes() { async fn queued_remote_wait_times_out_and_invalidates_the_connection() { let mut harness = DriverHarness::start(); let session = remote_session(); - harness - .open(session.clone(), Arc::new(RecordingDelegate::default())) - .await; + harness.open(session.clone()).await; let _started = harness - .start_cell(session.clone(), /*request_id*/ 2, "1") + .start_cell( + session.clone(), + /*request_id*/ 2, + "1", + Arc::new(RecordingDelegate::default()), + ) .await; let connection = Connection { command_tx: harness.command_tx.clone(), @@ -1317,11 +1389,14 @@ async fn queued_remote_wait_times_out_and_invalidates_the_connection() { async fn queued_remote_termination_times_out_and_invalidates_the_connection() { let mut harness = DriverHarness::start(); let session = remote_session(); - harness - .open(session.clone(), Arc::new(RecordingDelegate::default())) - .await; + harness.open(session.clone()).await; let _started = harness - .start_cell(session.clone(), /*request_id*/ 2, "1") + .start_cell( + session.clone(), + /*request_id*/ 2, + "1", + Arc::new(RecordingDelegate::default()), + ) .await; let connection = Connection { command_tx: harness.command_tx.clone(), @@ -1356,11 +1431,14 @@ async fn queued_remote_termination_times_out_and_invalidates_the_connection() { async fn cancelled_wait_is_retired_before_next_wait_is_sent() { let mut harness = DriverHarness::start(); let session = remote_session(); - harness - .open(session.clone(), Arc::new(RecordingDelegate::default())) - .await; + harness.open(session.clone()).await; let _started = harness - .start_cell(session.clone(), /*request_id*/ 2, "1") + .start_cell( + session.clone(), + /*request_id*/ 2, + "1", + Arc::new(RecordingDelegate::default()), + ) .await; let first_cancellation = CancellationToken::new(); let (first_tx, first_rx) = oneshot::channel(); @@ -1450,12 +1528,13 @@ async fn abandoned_execute_is_tracked_and_terminated_after_admission() { let mut harness = DriverHarness::start(); let session = remote_session(); let delegate = Arc::new(RecordingDelegate::default()); - harness.open(session.clone(), delegate.clone()).await; + harness.open(session.clone()).await; let cancellation = CancellationToken::new(); let (execute_tx, execute_rx) = oneshot::channel(); harness .command_tx .send(DriverCommand::Execute { + delegate: delegate.clone(), session: session.clone(), request: ExecuteRequest { tool_call_id: "call-1".to_string(), @@ -1547,12 +1626,13 @@ async fn delivered_but_unclaimed_execute_is_terminated_when_the_caller_is_cancel let mut harness = DriverHarness::start(); let session = remote_session(); let delegate = Arc::new(RecordingDelegate::default()); - harness.open(session.clone(), delegate.clone()).await; + harness.open(session.clone()).await; let cancellation = CancellationToken::new(); let (execute_tx, execute_rx) = oneshot::channel(); harness .command_tx .send(DriverCommand::Execute { + delegate: delegate.clone(), session: session.clone(), request: ExecuteRequest { tool_call_id: "call-1".to_string(), @@ -1635,7 +1715,13 @@ async fn delivered_but_unclaimed_execute_is_terminated_when_the_caller_is_cancel })) .await .expect("cell close"); - tokio::task::yield_now().await; + tokio::time::timeout(Duration::from_secs(2), async { + while Arc::strong_count(&delegate) > 1 { + tokio::task::yield_now().await; + } + }) + .await + .expect("abandoned execution should release its delegate"); assert!(harness.alive.load(Ordering::Acquire)); assert_eq!( @@ -1651,13 +1737,13 @@ async fn session_accepts_more_than_4096_cells_without_growing_a_tombstone_set() let mut harness = DriverHarness::start(); let session = remote_session(); let delegate = Arc::new(RecordingDelegate::default()); - harness.open(session.clone(), delegate.clone()).await; + harness.open(session.clone()).await; for sequence in 1..=CELL_COUNT { let request_id = i64::try_from(sequence).expect("cell sequence fits in i64") + 1; let cell_id = sequence.to_string(); let started = harness - .start_cell(session.clone(), request_id, &cell_id) + .start_cell(session.clone(), request_id, &cell_id, delegate.clone()) .await; harness .event_tx @@ -1705,11 +1791,12 @@ async fn connection_failure_closes_every_live_cell_once() { let mut harness = DriverHarness::start(); let session = remote_session(); let delegate = Arc::new(RecordingDelegate::default()); - let cleanup = harness.open(session.clone(), delegate.clone()).await; + let cleanup = harness.open(session.clone()).await; let (execute_tx, execute_rx) = oneshot::channel(); harness .command_tx .send(DriverCommand::Execute { + delegate: delegate.clone(), session, request: ExecuteRequest { tool_call_id: "call-1".to_string(), @@ -1759,9 +1846,14 @@ async fn session_cleanup_does_not_wait_for_delegate_completion() { let mut harness = DriverHarness::start(); let session = remote_session(); let (delegate, mut events_rx, release) = HeldDelegate::new(); - let cleanup = harness.open(session.clone(), delegate).await; + let cleanup = harness.open(session.clone()).await; let _started = harness - .start_cell(session.clone(), /*request_id*/ 2, "1") + .start_cell( + session.clone(), + /*request_id*/ 2, + "1", + delegate.clone(), + ) .await; harness .start_tool_delegate(&session, DelegateRequestId::new(/*value*/ 7)) @@ -1806,9 +1898,14 @@ async fn aborting_driver_marks_connection_dead_and_closes_cells() { let mut harness = DriverHarness::start(); let session = remote_session(); let delegate = Arc::new(RecordingDelegate::default()); - harness.open(session.clone(), delegate.clone()).await; + harness.open(session.clone()).await; let _started = harness - .start_cell(session.clone(), /*request_id*/ 2, "1") + .start_cell( + session.clone(), + /*request_id*/ 2, + "1", + delegate.clone(), + ) .await; let (wait_tx, wait_rx) = oneshot::channel(); harness @@ -1847,9 +1944,7 @@ async fn aborting_driver_marks_connection_dead_and_closes_cells() { async fn dropped_shutdown_waiter_does_not_abort_remote_cleanup() { let mut harness = DriverHarness::start(); let session = remote_session(); - harness - .open(session.clone(), Arc::new(RecordingDelegate::default())) - .await; + harness.open(session.clone()).await; let (shutdown_tx, shutdown_rx) = oneshot::channel(); harness .command_tx @@ -1878,6 +1973,7 @@ async fn dropped_shutdown_waiter_does_not_abort_remote_cleanup() { harness .command_tx .send(DriverCommand::Execute { + delegate: Arc::new(RecordingDelegate::default()), session, request: ExecuteRequest { tool_call_id: "call-2".to_string(), diff --git a/codex-rs/code-mode/src/remote_session_tests.rs b/codex-rs/code-mode/src/remote_session_tests.rs index a2e0a08ff8..e3f86bd2f1 100644 --- a/codex-rs/code-mode/src/remote_session_tests.rs +++ b/codex-rs/code-mode/src/remote_session_tests.rs @@ -62,7 +62,7 @@ async fn provider_returns_missing_host_error() { ); let error = provider - .create_session(Arc::new(NoopCodeModeSessionDelegate)) + .create_session() .await .err() .expect("missing host should fail"); @@ -76,13 +76,16 @@ async fn shutdown_before_open_does_not_spawn_the_host() { session.shutdown().await.expect("shutdown session"); let error = session - .execute(codex_code_mode_protocol::ExecuteRequest { - tool_call_id: "call-1".to_string(), - enabled_tools: Vec::new(), - source: "text('unreachable')".to_string(), - yield_time_ms: None, - max_output_tokens: None, - }) + .execute( + codex_code_mode_protocol::ExecuteRequest { + tool_call_id: "call-1".to_string(), + enabled_tools: Vec::new(), + source: "text('unreachable')".to_string(), + yield_time_ms: None, + max_output_tokens: None, + }, + Arc::new(NoopCodeModeSessionDelegate), + ) .await .err() .expect("shutdown session should reject execution"); diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index c9355ed1b8..355f0a195c 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -129,7 +129,10 @@ impl CodeModeService { request .yield_time_ms .get_or_insert(self.default_exec_yield_time_ms); - self.session().await?.execute(request).await + self.session() + .await? + .execute(request, self.dispatch_broker.clone()) + .await } pub(crate) async fn wait( @@ -237,7 +240,7 @@ impl CodeModeService { } session = self .session_provider - .create_session(self.dispatch_broker.clone()) => session?, + .create_session() => session?, }; if self.shutdown_token.is_cancelled() { let _ = session.shutdown().await; diff --git a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs index 0755fbbb3f..692d1704e8 100644 --- a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs +++ b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs @@ -66,10 +66,7 @@ impl CoreToolRuntime for TestHandler {} struct MissingCellCodeModeSessionProvider; impl codex_code_mode::CodeModeSessionProvider for MissingCellCodeModeSessionProvider { - fn create_session<'a>( - &'a self, - _delegate: Arc, - ) -> codex_code_mode::CodeModeSessionProviderFuture<'a> { + fn create_session(&self) -> codex_code_mode::CodeModeSessionProviderFuture<'_> { Box::pin(async { Ok(Arc::new(MissingCellCodeModeSession) as Arc) }) @@ -82,6 +79,7 @@ impl codex_code_mode::CodeModeSession for MissingCellCodeModeSession { fn execute<'a>( &'a self, _request: codex_code_mode::ExecuteRequest, + _delegate: Arc, ) -> codex_code_mode::CodeModeSessionResultFuture<'a, codex_code_mode::StartedCell> { Box::pin(async { Err("test session cannot execute cells".to_string()) }) }