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()) }) }