Scope code mode callback delegates to individual executions (#44865)

## Why

A delegate bound at session creation cannot provide different callbacks for cells sharing that session. Each execution needs to retain its own delegate across yields and release it when the cell is cleaned up.

## What changed

- Pass `CodeModeSessionDelegate` to `execute` instead of session creation.
- Route tool calls, notifications, and cell closure callbacks through the execution's delegate in the in-process runtime and the gRPC and stdio transports.
- Retain delegates with pending executions and live cells, releasing them through closure and cancellation cleanup.

## Testing

Add coverage for distinct delegates across yielded cells, gRPC callbacks before cell admission, and delegate release after completion or abandoned execution cleanup. Update transport tests to verify callbacks reach the owning cell's delegate.

GitOrigin-RevId: 7469f52104b993e790a40c65fb800ad02b7fd606

codex-rs/code-mode-host/src/grpc/mod.rs

@@ -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)?
             }
         };

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<Self> {
         Arc::new_cyclic(|weak: &Weak<Self>| {
-            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(
+                runtime: Arc::new(InProcessCodeModeSession::with_task_failure_handler(
-                    InProcessCodeModeSession::with_delegate_and_task_failure_handler(
+                    failure_handler,
-                        delegate,
+                    limits,
-                        failure_handler,
+                )),
-                        limits,
-                    ),
-                ),
                 closed,
                 state: Mutex::new(SessionState::default()),
                 events,

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(
+            Arc::new(InProcessCodeModeSession::with_task_failure_handler(
-                InProcessCodeModeSession::with_delegate_and_task_failure_handler(
+                task_failure_handler,
-                    delegate,
+                cell_execution_limits,
-                    task_failure_handler,
+            )),
-                    cell_execution_limits,
-                ),
-            ),
         );
         Ok(())
     }

codex-rs/code-mode-host/tests/grpc.rs

@@ -144,10 +144,11 @@ fn text_response(
 async fn execute(
     session: &Arc<dyn CodeModeSession>,
     request: ExecuteRequest,
+    delegate: Arc<dyn CodeModeSessionDelegate>,
 ) -> Result<RuntimeResponse> {
     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(
+        .create_session_with_limits(CodeModeSessionCellExecutionLimits {
-            Arc::new(NoopCodeModeSessionDelegate),
+            max_yield_time_ms: Some(/*value*/ 1),
-            CodeModeSessionCellExecutionLimits {
+            max_heap_size_bytes: Some(/*value*/ 16 * 1024 * 1024),
-                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(
+        .create_session_with_limits(CodeModeSessionCellExecutionLimits {
-            Arc::new(NoopCodeModeSessionDelegate),
+            max_yield_time_ms: Some(/*value*/ 1_000),
-            CodeModeSessionCellExecutionLimits {
+            max_heap_size_bytes: None,
-                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, callback, delegate.clone()),
-        execute(&session, request(r#"text("concurrent")"#)),
+        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())

codex-rs/code-mode-host/tests/grpc_notifications.rs

@@ -120,10 +120,11 @@ fn text_response(
 async fn execute(
     session: &Arc<dyn CodeModeSession>,
     request: ExecuteRequest,
+    delegate: Arc<dyn CodeModeSessionDelegate>,
 ) -> Result<RuntimeResponse> {
     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!(

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 {
+            .execute(
-                tool_call_id: format!("call-{index}"),
+                ExecuteRequest {
-                source: format!(
+                    tool_call_id: format!("call-{index}"),
-                    "const until = Date.now() + 150; while (Date.now() < until) {{}} \
+                    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(),
+                    enabled_tools: Vec::new(),
-                yield_time_ms: Some(/*value*/ 5_000),
+                    yield_time_ms: Some(/*value*/ 5_000),
-                max_output_tokens: Some(/*value*/ 1_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 {
+            .execute(
-                tool_call_id: "call-1".to_string(),
+                ExecuteRequest {
-                source: "await new Promise(() => {});".to_string(),
+                    tool_call_id: "call-1".to_string(),
-                enabled_tools: Vec::new(),
+                    source: "await new Promise(() => {});".to_string(),
-                yield_time_ms: Some(/*value*/ 200),
+                    enabled_tools: Vec::new(),
-                max_output_tokens: Some(/*value*/ 1_000),
+                    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();

codex-rs/code-mode-host/tests/stdio.rs

@@ -174,9 +174,13 @@ fn execute_request(source: &str) -> ExecuteRequest {
     }
 }
 
-async fn execute(session: &Arc<dyn CodeModeSession>, request: ExecuteRequest) -> RuntimeResponse {
+async fn execute(
+    session: &Arc<dyn CodeModeSession>,
+    request: ExecuteRequest,
+    delegate: Arc<dyn CodeModeSessionDelegate>,
+) -> RuntimeResponse {
     session
-        .execute(request)
+        .execute(request, delegate.clone())
         .await
         .expect("start execution")
         .initial_response()
@@ -187,8 +191,12 @@ async fn execute(session: &Arc<dyn CodeModeSession>, request: ExecuteRequest) ->
 async fn execute_to_terminal(
     session: &Arc<dyn CodeModeSession>,
     request: ExecuteRequest,
+    delegate: Arc<dyn CodeModeSessionDelegate>,
 ) -> 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(
+        .create_session_with_limits(CodeModeSessionCellExecutionLimits {
-            Arc::new(RecordingDelegate::default()),
+            max_yield_time_ms: Some(1),
-            CodeModeSessionCellExecutionLimits {
+            max_heap_size_bytes: None,
-                max_yield_time_ms: Some(1),
+        })
-                max_heap_size_bytes: None,
-            },
-        )
         .await
         .expect("create limited session");
     let other = provider
-        .create_session_with_limits(
+        .create_session_with_limits(CodeModeSessionCellExecutionLimits {
-            Arc::new(RecordingDelegate::default()),
+            max_yield_time_ms: Some(1_000),
-            CodeModeSessionCellExecutionLimits {
+            max_heap_size_bytes: None,
-                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!(
+        .execute(
-            r#"text("x".repeat({MAX_FRAME_BYTES}));"#
+            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 {

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<dyn CodeModeSessionDelegate>,
     ) -> 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>(
+    fn create_session(&self) -> CodeModeSessionProviderFuture<'_>;
-        &'a self,
-        delegate: Arc<dyn CodeModeSessionDelegate>,
-    ) -> CodeModeSessionProviderFuture<'a>;
 
     /// 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<dyn CodeModeSessionDelegate>,
         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())

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<ProtocolDelegate>,
+    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<dyn CodeModeSessionDelegate>) -> Self {
+    pub fn with_limits(cell_execution_limits: CodeModeSessionCellExecutionLimits) -> Self {
-        Self::with_delegate_and_limits(delegate, CodeModeSessionCellExecutionLimits::default())
-    }
-
-    pub fn with_delegate_and_limits(
-        delegate: Arc<dyn CodeModeSessionDelegate>,
-        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(
+    pub fn with_task_failure_handler(
-        delegate: Arc<dyn CodeModeSessionDelegate>,
         task_failure_handler: Arc<dyn Fn(String) + Send + Sync>,
         cell_execution_limits: CodeModeSessionCellExecutionLimits,
     ) -> Self {
         Self {
-            runtime: SessionRuntime::new_with_task_failure_handler(
+            runtime: SessionRuntime::new_with_task_failure_handler(Some(task_failure_handler)),
-                Arc::new(ProtocolDelegate { delegate }),
-                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<StartedCell, String> {
+    pub async fn execute(
+        &self,
+        request: ExecuteRequest,
+        delegate: Arc<dyn CodeModeSessionDelegate>,
+    ) -> Result<StartedCell, String> {
         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<dyn CodeModeSessionDelegate>,
     ) -> Result<ExecuteToPendingOutcome, String> {
         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<dyn CodeModeSessionDelegate>,
     ) -> 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> {

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<DelegateEvent>) -> 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 {
+        .execute(
-            source: r#"text("before"); yield_control(); text("after");"#.to_string(),
+            ExecuteRequest {
-            yield_time_ms: Some(60_000),
+                source: r#"text("before"); yield_control(); text("after");"#.to_string(),
-            ..execute_request("")
+                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 {
+            .execute_to_pending(
-                enabled_tools: vec![blocking_tool()],
+                ExecuteRequest {
-                source: r#"
+                    enabled_tools: vec![blocking_tool()],
+                    source: r#"
 await tools.block({});
 text("after");
 "#
-                .to_string(),
+                    .to_string(),
-                yield_time_ms: Some(60_000),
+                    yield_time_ms: Some(60_000),
-                ..execute_request("")
+                    ..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(
+        .execute(
-            r#"yield_control(); store("finished", true); text("done");"#,
+            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 {
+                .execute(
-                    yield_time_ms: Some(60_000),
+                    ExecuteRequest {
-                    ..execute_request(r#"text(String(load("finished")));"#)
+                        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(
+        .execute(
-            r#"notify("pending"); await new Promise(() => {});"#,
+            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(
+        .execute(
-            r#"notify("pending"); await new Promise(() => {});"#,
+            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 {
+        .execute(
-            enabled_tools: vec![blocking_tool()],
+            ExecuteRequest {
-            source: r#"tools.block({}); text("done");"#.to_string(),
+                enabled_tools: vec![blocking_tool()],
-            yield_time_ms: Some(60_000),
+                source: r#"tools.block({}); text("done");"#.to_string(),
-            ..execute_request("")
+                yield_time_ms: Some(60_000),
-        })
+                ..execute_request("")
+            },
+            delegate.clone(),
+        )
         .await
         .unwrap();
 

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(
+        let session = InProcessCodeModeSession::with_limits(CodeModeSessionCellExecutionLimits {
-            Arc::new(ReleasableToolDelegate::default()),
+            max_yield_time_ms,
-            CodeModeSessionCellExecutionLimits {
+            max_heap_size_bytes: None,
-                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(
+    let service = InProcessCodeModeSession::with_limits(CodeModeSessionCellExecutionLimits {
-        delegate.clone(),
+        max_yield_time_ms: Some(/*value*/ 10_000),
-        CodeModeSessionCellExecutionLimits {
+        max_heap_size_bytes: None,
-            max_yield_time_ms: Some(/*value*/ 10_000),
+    });
-            max_heap_size_bytes: None,
-        },
-    );
     let started = service
-        .execute(ExecuteRequest {
+        .execute(
-            enabled_tools: vec![echo_tool()],
+            ExecuteRequest {
-            source: r#"await tools.echo({}); text("done");"#.to_string(),
+                enabled_tools: vec![echo_tool()],
-            yield_time_ms: None,
+                source: r#"await tools.echo({}); text("done");"#.to_string(),
-            ..execute_request("")
+                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 {
+        .execute_to_pending(
-            enabled_tools: vec![echo_tool()],
+            ExecuteRequest {
-            source: r#"await tools.echo({}); text("done");"#.to_string(),
+                enabled_tools: vec![echo_tool()],
-            ..execute_request("")
+                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(
+    let zero_session = InProcessCodeModeSession::with_limits(CodeModeSessionCellExecutionLimits {
-        zero_delegate.clone(),
+        max_yield_time_ms: Some(/*value*/ 0),
-        CodeModeSessionCellExecutionLimits {
+        max_heap_size_bytes: None,
-            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(
+    let limited_session =
-        limited_delegate.clone(),
+        InProcessCodeModeSession::with_limits(CodeModeSessionCellExecutionLimits {
-        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 zero_started = zero_session
-    let limited_started = limited_session.execute(request).await.unwrap();
+        .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 {
+        .execute(
-            source: "while (true) {}".to_string(),
+            ExecuteRequest {
-            ..execute_request("")
+                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 {
+        .execute_to_pending(
-            source: r#"text("done");"#.to_string(),
+            ExecuteRequest {
-            yield_time_ms: Some(60_000),
+                source: r#"text("done");"#.to_string(),
-            ..execute_request("")
+                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 {
+        service.execute_to_pending(
-            source: r#"text("before"); await new Promise(() => {});"#.to_string(),
+            ExecuteRequest {
-            yield_time_ms: Some(60_000),
+                source: r#"text("before"); await new Promise(() => {});"#.to_string(),
-            ..execute_request("")
+                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 {
+        .execute_to_pending(
-            enabled_tools: vec![echo_tool()],
+            ExecuteRequest {
-            source: r#"
+                enabled_tools: vec![echo_tool()],
+                source: r#"
 await Promise.all([
   tools.echo({ value: "first" }),
   tools.echo({ value: "second" }),
 ]);
 "#
-            .to_string(),
+                .to_string(),
-            yield_time_ms: Some(60_000),
+                yield_time_ms: Some(60_000),
-            ..execute_request("")
+                ..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 {
+        .execute_to_pending(
-            enabled_tools: vec![echo_tool()],
+            ExecuteRequest {
-            source: r#"
+                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(),
+                .to_string(),
-            yield_time_ms: Some(60_000),
+                yield_time_ms: Some(60_000),
-            ..execute_request("")
+                ..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 {
+        .execute_to_pending(
-            enabled_tools: vec![echo_tool()],
+            ExecuteRequest {
-            source: r#"
+                enabled_tools: vec![echo_tool()],
+                source: r#"
 await tools.echo({});
 text("after");
 await new Promise(() => {});
 "#
-            .to_string(),
+                .to_string(),
-            yield_time_ms: Some(60_000),
+                yield_time_ms: Some(60_000),
-            ..execute_request("")
+                ..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 {
+        .execute_to_pending(
-            enabled_tools: vec![echo_tool()],
+            ExecuteRequest {
-            source: r#"
+                enabled_tools: vec![echo_tool()],
+                source: r#"
 await tools.echo({});
 text("done");
 "#
-            .to_string(),
+                .to_string(),
-            yield_time_ms: Some(60_000),
+                yield_time_ms: Some(60_000),
-            ..execute_request("")
+                ..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(
+    let response = execute(&service, ExecuteRequest {
-            &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(
+    let response = execute(&service, ExecuteRequest {
-            &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(
+    let response = execute(&service, ExecuteRequest {
-            &service,
-            ExecuteRequest {
                 source: r#"
 image({
   content: [
@@ -1578,8 +1596,7 @@ image({
                 .to_string(),
                 yield_time_ms: None,
                 ..execute_request("")
-            },
+            })
-        )
         .await;
 
     assert_eq!(

codex-rs/code-mode-runtime/src/session_runtime/mod.rs

@@ -38,27 +38,25 @@ use crate::cell_actor::CompletionCommit;
 type RuntimeEventFuture = Pin<Box<dyn Future<Output = Result<CellEvent, Error>> + Send + 'static>>;
 
 /// Owns all cells and shared state for one transport-neutral code-mode session.
-pub(crate) struct SessionRuntime<D: SessionRuntimeDelegate> {
+pub(crate) struct SessionRuntime {
-    inner: Arc<Inner<D>>,
+    inner: Arc<Inner>,
 }
 
-struct Inner<D: SessionRuntimeDelegate> {
+struct Inner {
     stored_values: Mutex<HashMap<String, JsonValue>>,
     cells: Mutex<HashMap<CellId, CellHandle>>,
     cell_tasks: TaskTracker,
     shutdown_token: CancellationToken,
-    delegate: Arc<D>,
     task_failure_handler: Option<TaskFailureHandler>,
     next_cell_id: AtomicU64,
 }
 
-impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
+impl SessionRuntime {
-    pub(crate) fn new(delegate: Arc<D>) -> Self {
+    pub(crate) fn new() -> Self {
-        Self::new_with_task_failure_handler(delegate, /*task_failure_handler*/ None)
+        Self::new_with_task_failure_handler(/*task_failure_handler*/ None)
     }
 
     pub(crate) fn new_with_task_failure_handler(
-        delegate: Arc<D>,
         task_failure_handler: Option<TaskFailureHandler>,
     ) -> Self {
         Self {
@@ -67,24 +65,24 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
                 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<D: SessionRuntimeDelegate>(
         &self,
         request: CreateCellRequest,
         initial_observe_mode: ObserveMode,
+        delegate: Arc<D>,
     ) -> Result<StartedCell, Error> {
         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<D: SessionRuntimeDelegate> SessionRuntime<D> {
             .map_err(|_| Error::CellIdSpaceExhausted)
     }
 
-    async fn start_cell(
+    async fn start_cell<D: SessionRuntimeDelegate>(
         &self,
         cell_id: CellId,
         request: CreateCellRequest,
         initial_observe_mode: ObserveMode,
+        delegate: Arc<D>,
     ) -> Result<RuntimeEventFuture, Error> {
         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<D: SessionRuntimeDelegate> SessionRuntime<D> {
     }
 }
 
-impl<D: SessionRuntimeDelegate> Drop for SessionRuntime<D> {
+impl Drop for SessionRuntime {
     fn drop(&mut self) {
         self.begin_shutdown();
     }
@@ -235,8 +235,9 @@ impl PendingEvent {
 }
 
 struct RuntimeCellHost<D: SessionRuntimeDelegate> {
+    delegate: Arc<D>,
     cell_id: CellId,
-    inner: Arc<Inner<D>>,
+    inner: Arc<Inner>,
     // 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<D: SessionRuntimeDelegate> CellHost for RuntimeCellHost<D> {
         invocation: CellToolCall,
         cancellation_token: CancellationToken,
     ) -> Result<JsonValue, String> {
-        self.inner
+        self.delegate
-            .delegate
             .invoke_tool(
                 NestedToolCall {
                     cell_id: self.cell_id.clone(),
@@ -270,8 +270,7 @@ impl<D: SessionRuntimeDelegate> CellHost for RuntimeCellHost<D> {
         text: String,
         cancellation_token: CancellationToken,
     ) -> Result<(), String> {
-        self.inner
+        self.delegate
-            .delegate
             .notify(call_id, self.cell_id.clone(), text, cancellation_token)
             .await
     }
@@ -298,7 +297,7 @@ impl<D: SessionRuntimeDelegate> CellHost for RuntimeCellHost<D> {
 
     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);
     }
 }
 

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(
+    let runtime = SessionRuntime::new_with_task_failure_handler(Some(Arc::new(move |reason| {
-        Arc::new(PanickingClosedDelegate),
+        let _ = failure_tx.send(reason);
-        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();

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 {
+        .execute(
-            tool_call_id: "call_1".to_string(),
+            ExecuteRequest {
-            enabled_tools: Vec::new(),
+                tool_call_id: "call_1".to_string(),
-            source: "21 * 2;".to_string(),
+                enabled_tools: Vec::new(),
-            yield_time_ms: None,
+                source: "21 * 2;".to_string(),
-            max_output_tokens: None,
+                yield_time_ms: None,
-        })
+                max_output_tokens: None,
+            },
+            Arc::new(NoopCodeModeSessionDelegate),
+        )
         .await
         .expect("start code-mode cell");
     let cell_id = started.cell_id.clone();

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;
+                                biased;
-                            _ = cancellation.cancelled() => return,
+                                _ = cancellation.cancelled() => return,
-                            result = callback => match result {
+                                result = callback => match result {
-                                Ok(result) => result,
+                                    Ok(result) => result,
-                                Err(_) => Err("code-mode tool delegate panicked".to_string()),
+                                    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(&notification)?;
-        let cancellation = match admission {
+        let (cancellation, delegate) = match admission {
-            CallbackAdmission::Active(cancellation) => cancellation,
+            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),

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<dyn CodeModeSessionDelegate>,
         limits: CodeModeSessionCellExecutionLimits,
     ) -> Result<Arc<GrpcCodeModeSession>, 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>(
+    fn create_session(&self) -> CodeModeSessionProviderFuture<'_> {
-        &'a self,
+        self.create_session_with_limits(CodeModeSessionCellExecutionLimits::default())
-        delegate: Arc<dyn CodeModeSessionDelegate>,
-    ) -> CodeModeSessionProviderFuture<'a> {
-        self.create_session_with_limits(delegate, CodeModeSessionCellExecutionLimits::default())
     }
 
     fn create_session_with_limits<'a>(
         &'a self,
-        delegate: Arc<dyn CodeModeSessionDelegate>,
         limits: CodeModeSessionCellExecutionLimits,
     ) -> CodeModeSessionProviderFuture<'a> {
         Box::pin(async move {
-            let session = Arc::new(reconnect::ReconnectableSession::new(
+            let session = Arc::new(reconnect::ReconnectableSession::new(self.clone(), limits));
-                self.clone(),
-                delegate,
-                limits,
-            ));
             session.initialize().await?;
             Ok(session as _)
         })
@@ -229,8 +220,9 @@ impl CodeModeSession for GrpcCodeModeSession {
     fn execute<'a>(
         &'a self,
         request: ExecuteRequest,
+        delegate: Arc<dyn CodeModeSessionDelegate>,
     ) -> 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<dyn CodeModeSessionDelegate>,
     runtime: tokio::runtime::Handle,
     state: Mutex<SessionState>,
     wait_slots: Mutex<HashMap<CellId, Weak<WaitSlot>>>,
@@ -282,14 +273,14 @@ impl SessionInner {
         Ok(())
     }
 
-    pub(super) fn report_closed_cell(&self, cell_id: Option<CellId>) {
+    pub(super) fn report_closed_cell(&self, cell: Option<ClosedCell>) {
-        if let Some(cell_id) = cell_id {
+        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);
             }));
         }
     }

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<Self>,
         request: ExecuteRequest,
+        delegate: Arc<dyn CodeModeSessionDelegate>,
     ) -> Result<StartedCell, String> {
         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,

codex-rs/code-mode/src/grpc_session/reconnect.rs

@@ -31,7 +31,6 @@ pub(super) struct ReconnectableSession {
 
 struct ReconnectInner {
     provider: GrpcCodeModeSessionProvider,
-    delegate: Arc<dyn CodeModeSessionDelegate>,
     limits: CodeModeSessionCellExecutionLimits,
     binding: Mutex<Option<SessionBinding>>,
     opening_permit: Semaphore,
@@ -49,13 +48,11 @@ struct SessionBinding {
 impl ReconnectableSession {
     pub(super) fn new(
         provider: GrpcCodeModeSessionProvider,
-        delegate: Arc<dyn CodeModeSessionDelegate>,
         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<dyn CodeModeSessionDelegate>,
     ) -> 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,

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<dyn CodeModeSessionDelegate>);
+
 pub(super) enum CallbackAdmission {
-    Active(CancellationToken),
+    Active(CancellationToken, Arc<dyn CodeModeSessionDelegate>),
     Cancelled,
     Closed,
     Rejected(String),
 }
 
-#[derive(Default)]
 struct ExecutionRecord {
+    delegate: Arc<dyn CodeModeSessionDelegate>,
     cell_id: Option<CellId>,
     tool_call_id: String,
     enabled_tools: HashMap<ToolName, i32>,
@@ -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<dyn CodeModeSessionDelegate>,
+    ) -> 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<Option<CellId>, String> {
+    ) -> Result<Option<ClosedCell>, 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<CellId> {
+    pub(super) fn finish_notification(&mut self, execution_id: &str) -> Option<ClosedCell> {
         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<Option<CellId>, String> {
+    ) -> Result<Option<ClosedCell>, 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<String>) -> Vec<CellId> {
+    pub(super) fn close(&mut self, failure: Option<String>) -> Vec<ClosedCell> {
         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<CellId> {
+    fn close_execution_if_ready(&mut self, execution_id: &str) -> Option<ClosedCell> {
         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<CellId> {
+    pub(super) fn remove_execution(&mut self, execution_id: &str) -> Option<ClosedCell> {
         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> {

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<Arc<dyn codex_code_mode_protocol::CodeModeSessionDelegate>> = 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(&notice).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(&notification("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"),
+            .expect("record cell closure")
-        None
+            .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"),
+            .expect("claim started cell")
-        None
+            .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"),
+            .expect("record early cell closure")
-        None
+            .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!(state.remove_execution("execution").is_none());
-    assert_eq!(
+    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"),
+            .expect("ignore closure for abandoned execution")
-        None
+            .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"),
+            .expect("ignore delayed cell closure")
-        None
+            .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(&notification("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(&notification("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(&notification("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(&notification("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<_>>(),
         vec![CellId::new("cell".to_string())]
     );
     assert!(cancellation.is_cancelled());

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>(
+    fn create_session(&self) -> CodeModeSessionProviderFuture<'_> {
-        &'a self,
+        self.create_session_with_limits(CodeModeSessionCellExecutionLimits::default())
-        delegate: Arc<dyn CodeModeSessionDelegate>,
-    ) -> CodeModeSessionProviderFuture<'a> {
-        self.create_session_with_limits(delegate, CodeModeSessionCellExecutionLimits::default())
     }
 
     fn create_session_with_limits<'a>(
         &'a self,
-        delegate: Arc<dyn CodeModeSessionDelegate>,
         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>(
+    fn create_session(&self) -> CodeModeSessionProviderFuture<'_> {
-        &'a self,
-        _delegate: Arc<dyn CodeModeSessionDelegate>,
-    ) -> CodeModeSessionProviderFuture<'a> {
         Box::pin(async { Err("code-mode host is disabled".to_string()) })
     }
 
     fn create_session_with_limits<'a>(
         &'a self,
-        delegate: Arc<dyn CodeModeSessionDelegate>,
         _limits: CodeModeSessionCellExecutionLimits,
     ) -> CodeModeSessionProviderFuture<'a> {
-        self.create_session(delegate)
+        self.create_session()
     }
 }
 
 async fn create_host_session(
-    delegate: Arc<dyn CodeModeSessionDelegate>,
     host: Arc<OwnedCodeModeHost>,
     limits: CodeModeSessionCellExecutionLimits,
 ) -> Result<Arc<dyn CodeModeSession>, 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<OwnedCodeModeHost>,
-    delegate: Arc<dyn CodeModeSessionDelegate>,
     limits: CodeModeSessionCellExecutionLimits,
     state: StdMutex<SessionState>,
     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(
+    fn with_host(host: Arc<OwnedCodeModeHost>, limits: CodeModeSessionCellExecutionLimits) -> Self {
-        delegate: Arc<dyn CodeModeSessionDelegate>,
-        host: Arc<OwnedCodeModeHost>,
-        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<StartedCell, String> {
+    pub async fn execute(
+        &self,
+        request: ExecuteRequest,
+        delegate: Arc<dyn CodeModeSessionDelegate>,
+    ) -> Result<StartedCell, String> {
         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<WaitOutcome, String> {
@@ -348,11 +338,7 @@ impl SessionInner {
         let result = match self.host.connection().await {
             Ok(connection) => {
                 let cleanup = connection
-                    .open_session(
+                    .open_session(remote.clone(), self.limits.clone())
-                        remote.clone(),
-                        Arc::clone(&self.delegate),
-                        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<dyn CodeModeSessionDelegate>,
     ) -> 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> {

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<dyn CodeModeSessionDelegate>,
         limits: CodeModeSessionCellExecutionLimits,
     ) -> Result<SessionCleanup, String> {
         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<dyn CodeModeSessionDelegate>,
     ) -> Result<StartedCell, String> {
         let cancellation = CallerCancellation::new();
         let (response_tx, response_rx) = oneshot::channel();
         self.send(DriverCommand::Execute {
             session,
             request,
+            delegate,
             caller_cancellation: cancellation.token(),
             response_tx,
         })

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(
+            } => self.open_session(session, limits, cleanup, caller_cancellation, response_tx),
-                session,
-                delegate,
-                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<dyn CodeModeSessionDelegate>,
         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<dyn CodeModeSessionDelegate>,
         caller_cancellation: CancellationToken,
         response_tx: oneshot::Sender<Result<DeliveredExecute, String>>,
     ) -> bool {
@@ -179,6 +171,7 @@ impl ConnectionDriver {
             request_id,
             PendingRequest::Execute {
                 session,
+                delegate,
                 response_tx,
                 initial_response_tx,
                 initial_response_rx,

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
+                    self.sessions.insert_ready(session.clone(), cleanup);
-                        .insert_ready(session.clone(), delegate, 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

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<dyn CodeModeSessionDelegate>,
     cleanup: SessionCleanup,
     phase: SessionPhase,
-    cells: HashMap<WireCellId, CellId>,
+    cells: HashMap<WireCellId, CellOwner>,
 }
 
 #[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(
+    pub(super) fn insert_ready(&mut self, session: RemoteSession, cleanup: SessionCleanup) {
-        &mut self,
-        session: RemoteSession,
-        delegate: Arc<dyn CodeModeSessionDelegate>,
-        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<dyn CodeModeSessionDelegate>,
     ) -> Result<CellId, CellAdmissionError> {
         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,
+            cell_id: owner.cell_id.clone(),
-            delegate: Arc::clone(&session.delegate),
+            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_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),
-        })
     }
 
     pub(super) fn remove_session(&mut self, session_id: &SessionId) -> Vec<CellOwner> {
         let Some(session) = self.records.remove(session_id) else {
             return Vec::new();
         };
-        session
+        session.cells.into_values().collect()
-            .cells
-            .into_values()
-            .map(|cell_id| CellOwner {
-                session_id: session_id.clone(),
-                cell_id,
-                delegate: Arc::clone(&session.delegate),
-            })
-            .collect()
     }
 
     pub(super) fn drain(&mut self) -> Vec<FailedSession> {
         let sessions = std::mem::take(&mut self.records);
         sessions
-            .into_iter()
+            .into_values()
-            .map(|(session_id, session)| {
+            .map(|session| {
-                let cells = session
+                let cells = session.cells.into_values().collect();
-                    .cells
-                    .into_values()
-                    .map(|cell_id| CellOwner {
-                        session_id: session_id.clone(),
-                        cell_id,
-                        delegate: Arc::clone(&session.delegate),
-                    })
-                    .collect();
                 FailedSession {
                     cleanup: session.cleanup,
                     cells,

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<dyn CodeModeSessionDelegate>,
         limits: CodeModeSessionCellExecutionLimits,
         cleanup: SessionCleanup,
         caller_cancellation: CancellationToken,
@@ -38,6 +37,7 @@ pub(in crate::remote_session::connection) enum DriverCommand {
     },
     Execute {
         session: RemoteSession,
+        delegate: Arc<dyn CodeModeSessionDelegate>,
         request: ExecuteRequest,
         caller_cancellation: CancellationToken,
         response_tx: oneshot::Sender<Result<DeliveredExecute, String>>,
@@ -136,13 +136,13 @@ pub(super) struct UnclaimedExecute {
 pub(super) enum PendingRequest {
     OpenSession {
         session: RemoteSession,
-        delegate: Arc<dyn CodeModeSessionDelegate>,
         cleanup: SessionCleanup,
         cancellation: CancellableRequest,
         response_tx: oneshot::Sender<Result<(), String>>,
     },
     Execute {
         session: RemoteSession,
+        delegate: Arc<dyn CodeModeSessionDelegate>,
         response_tx: oneshot::Sender<Result<DeliveredExecute, String>>,
         initial_response_tx: oneshot::Sender<Result<RuntimeResponse, String>>,
         initial_response_rx: oneshot::Receiver<Result<RuntimeResponse, String>>,

codex-rs/code-mode/src/remote_session/connection/driver_tests.rs

@@ -87,17 +87,13 @@ impl DriverHarness {
         }
     }
 
-    async fn open(
+    async fn open(&mut self, session: RemoteSession) -> SessionCleanup {
-        &mut self,
-        session: RemoteSession,
-        delegate: Arc<dyn CodeModeSessionDelegate>,
-    ) -> 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<dyn CodeModeSessionDelegate>,
     ) -> 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());
+    let first_delegate = Arc::new(RecordingDelegate::default());
-    harness.open(session.clone(), delegate.clone()).await;
+    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 [
+    for (tool_call_id, response_tx, delegate) in [
-        ("first-cell", first_response_tx),
+        ("first-cell", first_response_tx, first_delegate.clone()),
-        ("second-cell", second_response_tx),
+        ("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
+    harness.open(session.clone()).await;
-        .open(session.clone(), Arc::new(PanickingDelegate))
-        .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;
+    harness.open(session.clone()).await;
-    let started = harness.start_cell(session, /*request_id*/ 2, "1").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
+    harness.open(session.clone()).await;
-        .open(session.clone(), Arc::new(RecordingDelegate::default()))
-        .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
+    harness.open(session.clone()).await;
-        .open(session.clone(), Arc::new(RecordingDelegate::default()))
-        .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
+    harness.open(session.clone()).await;
-        .open(session.clone(), Arc::new(RecordingDelegate::default()))
-        .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
+    harness.open(session.clone()).await;
-        .open(session.clone(), Arc::new(RecordingDelegate::default()))
-        .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
+    harness.open(session.clone()).await;
-        .open(session.clone(), Arc::new(RecordingDelegate::default()))
-        .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(),

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 {
+        .execute(
-            tool_call_id: "call-1".to_string(),
+            codex_code_mode_protocol::ExecuteRequest {
-            enabled_tools: Vec::new(),
+                tool_call_id: "call-1".to_string(),
-            source: "text('unreachable')".to_string(),
+                enabled_tools: Vec::new(),
-            yield_time_ms: None,
+                source: "text('unreachable')".to_string(),
-            max_output_tokens: None,
+                yield_time_ms: None,
-        })
+                max_output_tokens: None,
+            },
+            Arc::new(NoopCodeModeSessionDelegate),
+        )
         .await
         .err()
         .expect("shutdown session should reject execution");

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;

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>(
+    fn create_session(&self) -> codex_code_mode::CodeModeSessionProviderFuture<'_> {
-        &'a self,
-        _delegate: Arc<dyn codex_code_mode::CodeModeSessionDelegate>,
-    ) -> codex_code_mode::CodeModeSessionProviderFuture<'a> {
         Box::pin(async {
             Ok(Arc::new(MissingCellCodeModeSession) as Arc<dyn codex_code_mode::CodeModeSession>)
         })
@@ -82,6 +79,7 @@ impl codex_code_mode::CodeModeSession for MissingCellCodeModeSession {
     fn execute<'a>(
         &'a self,
         _request: codex_code_mode::ExecuteRequest,
+        _delegate: Arc<dyn codex_code_mode::CodeModeSessionDelegate>,
     ) -> codex_code_mode::CodeModeSessionResultFuture<'a, codex_code_mode::StartedCell> {
         Box::pin(async { Err("test session cannot execute cells".to_string()) })
     }