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
2026-11-05 20:10:33 +00:00 · 2026-09-11 17:00:07 +00:00
parent bc5957eac9
commit 3305c4f31d
30 changed files with 1027 additions and 637 deletions
@@ -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 {
-        Self::with_delegate_and_limits(delegate, CodeModeSessionCellExecutionLimits::default())
-    }
-
-    pub fn with_delegate_and_limits(
-        delegate: Arc<dyn CodeModeSessionDelegate>,
-        cell_execution_limits: CodeModeSessionCellExecutionLimits,
-    ) -> Self {
+    pub fn with_limits(cell_execution_limits: CodeModeSessionCellExecutionLimits) -> Self {
        Self {
-            runtime: SessionRuntime::new(Arc::new(ProtocolDelegate { delegate })),
+            runtime: SessionRuntime::new(),
            cell_execution_limits: CodeModeSessionCellExecutionLimits {
                max_heap_size_bytes: None,
                ..cell_execution_limits
@@ -57,16 +49,12 @@ impl InProcessCodeModeSession {
        }
    }

-    pub fn with_delegate_and_task_failure_handler(
-        delegate: Arc<dyn CodeModeSessionDelegate>,
+    pub fn with_task_failure_handler(
        task_failure_handler: Arc<dyn Fn(String) + Send + Sync>,
        cell_execution_limits: CodeModeSessionCellExecutionLimits,
    ) -> Self {
        Self {
-            runtime: SessionRuntime::new_with_task_failure_handler(
-                Arc::new(ProtocolDelegate { delegate }),
-                Some(task_failure_handler),
-            ),
+            runtime: SessionRuntime::new_with_task_failure_handler(Some(task_failure_handler)),
            cell_execution_limits: CodeModeSessionCellExecutionLimits {
                max_heap_size_bytes: None,
                ..cell_execution_limits
@@ -74,13 +62,18 @@ impl InProcessCodeModeSession {
        }
    }

-    pub async fn execute(&self, request: ExecuteRequest) -> Result<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> {
@@ -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 {
-            source: r#"text("before"); yield_control(); text("after");"#.to_string(),
-            yield_time_ms: Some(60_000),
-            ..execute_request("")
-        })
+        .execute(
+            ExecuteRequest {
+                source: r#"text("before"); yield_control(); text("after");"#.to_string(),
+                yield_time_ms: Some(60_000),
+                ..execute_request("")
+            },
+            Arc::new(NoopCodeModeSessionDelegate),
+        )
        .await
        .unwrap();

@@ -228,20 +302,23 @@ async fn yields_and_resumes() {
 #[tokio::test]
 async fn returns_and_resumes_from_the_pending_frontier() {
    let (delegate, mut events_rx) = BlockingDelegate::new();
-    let service = InProcessCodeModeSession::with_delegate(delegate.clone());
+    let service = InProcessCodeModeSession::new();

    assert_eq!(
        service
-            .execute_to_pending(ExecuteRequest {
-                enabled_tools: vec![blocking_tool()],
-                source: r#"
+            .execute_to_pending(
+                ExecuteRequest {
+                    enabled_tools: vec![blocking_tool()],
+                    source: r#"
 await tools.block({});
 text("after");
 "#
-                .to_string(),
-                yield_time_ms: Some(60_000),
-                ..execute_request("")
-            })
+                    .to_string(),
+                    yield_time_ms: Some(60_000),
+                    ..execute_request("")
+                },
+                delegate.clone()
+            )
            .await
            .unwrap(),
        ExecuteToPendingOutcome::Pending {
@@ -278,9 +355,10 @@ text("after");
 async fn observed_natural_completion_wins_over_termination() {
    let service = InProcessCodeModeSession::new();
    let cell = service
-        .execute(execute_request(
-            r#"yield_control(); store("finished", true); text("done");"#,
-        ))
+        .execute(
+            execute_request(r#"yield_control(); store("finished", true); text("done");"#),
+            Arc::new(NoopCodeModeSessionDelegate),
+        )
        .await
        .unwrap();

@@ -295,10 +373,13 @@ async fn observed_natural_completion_wins_over_termination() {
    tokio::time::timeout(Duration::from_secs(1), async {
        loop {
            let response = service
-                .execute(ExecuteRequest {
-                    yield_time_ms: Some(60_000),
-                    ..execute_request(r#"text(String(load("finished")));"#)
-                })
+                .execute(
+                    ExecuteRequest {
+                        yield_time_ms: Some(60_000),
+                        ..execute_request(r#"text(String(load("finished")));"#)
+                    },
+                    Arc::new(NoopCodeModeSessionDelegate),
+                )
                .await
                .unwrap()
                .initial_response()
@@ -335,11 +416,12 @@ async fn observed_natural_completion_wins_over_termination() {
 #[tokio::test]
 async fn termination_cancels_pending_callbacks_before_responding() {
    let (delegate, mut events_rx) = BlockingDelegate::new();
-    let service = InProcessCodeModeSession::with_delegate(delegate.clone());
+    let service = InProcessCodeModeSession::new();
    let cell = service
-        .execute(execute_request(
-            r#"notify("pending"); await new Promise(() => {});"#,
-        ))
+        .execute(
+            execute_request(r#"notify("pending"); await new Promise(() => {});"#),
+            delegate.clone(),
+        )
        .await
        .unwrap();

@@ -377,9 +459,9 @@ async fn termination_cancels_pending_callbacks_before_responding() {
 #[tokio::test]
 async fn shutdown_cancels_notifications_while_natural_completion_is_draining() {
    let (delegate, mut events_rx) = HeldNotificationDelegate::new();
-    let service = Arc::new(InProcessCodeModeSession::with_delegate(delegate.clone()));
+    let service = Arc::new(InProcessCodeModeSession::new());
    service
-        .execute(execute_request(r#"notify("pending");"#))
+        .execute(execute_request(r#"notify("pending");"#), delegate.clone())
        .await
        .unwrap();

@@ -407,11 +489,12 @@ async fn shutdown_cancels_notifications_while_natural_completion_is_draining() {
 #[tokio::test]
 async fn repeated_termination_is_rejected_while_callback_cleanup_is_pending() {
    let (delegate, mut events_rx) = HeldNotificationDelegate::new();
-    let service = Arc::new(InProcessCodeModeSession::with_delegate(delegate.clone()));
+    let service = Arc::new(InProcessCodeModeSession::new());
    let cell = service
-        .execute(execute_request(
-            r#"notify("pending"); await new Promise(() => {});"#,
-        ))
+        .execute(
+            execute_request(r#"notify("pending"); await new Promise(() => {});"#),
+            delegate.clone(),
+        )
        .await
        .unwrap();

@@ -461,7 +544,10 @@ async fn repeated_termination_is_rejected_while_callback_cleanup_is_pending() {
 async fn second_observer_is_rejected_without_displacing_the_first() {
    let service = InProcessCodeModeSession::new();
    let cell = service
-        .execute(execute_request("await new Promise(() => {});"))
+        .execute(
+            execute_request("await new Promise(() => {});"),
+            Arc::new(NoopCodeModeSessionDelegate),
+        )
        .await
        .unwrap();

@@ -509,14 +595,17 @@ async fn second_observer_is_rejected_without_displacing_the_first() {
 #[tokio::test]
 async fn natural_completion_cleans_up_callbacks_before_responding() {
    let (delegate, mut events_rx) = BlockingDelegate::new();
-    let service = InProcessCodeModeSession::with_delegate(delegate.clone());
+    let service = InProcessCodeModeSession::new();
    let cell = service
-        .execute(ExecuteRequest {
-            enabled_tools: vec![blocking_tool()],
-            source: r#"tools.block({}); text("done");"#.to_string(),
-            yield_time_ms: Some(60_000),
-            ..execute_request("")
-        })
+        .execute(
+            ExecuteRequest {
+                enabled_tools: vec![blocking_tool()],
+                source: r#"tools.block({}); text("done");"#.to_string(),
+                yield_time_ms: Some(60_000),
+                ..execute_request("")
+            },
+            delegate.clone(),
+        )
        .await
        .unwrap();

@@ -1,3 +1,4 @@
+use codex_code_mode_protocol::NoopCodeModeSessionDelegate;
 use std::sync::Arc;
 use std::sync::atomic::AtomicBool;
 use std::sync::atomic::Ordering;
@@ -43,13 +44,10 @@ fn resolve_yield_timeout_applies_grace_before_session_limits() {
        (Some(10_500), 5_000, Duration::from_secs(5)),
        (Some(u64::MAX), u64::MAX, Duration::from_millis(u64::MAX)),
    ] {
-        let session = InProcessCodeModeSession::with_delegate_and_limits(
-            Arc::new(ReleasableToolDelegate::default()),
-            CodeModeSessionCellExecutionLimits {
-                max_yield_time_ms,
-                max_heap_size_bytes: None,
-            },
-        );
+        let session = InProcessCodeModeSession::with_limits(CodeModeSessionCellExecutionLimits {
+            max_yield_time_ms,
+            max_heap_size_bytes: None,
+        });

        assert_eq!(
            session.resolve_yield_timeout(requested_yield_time_ms),
@@ -62,14 +60,14 @@ fn resolve_yield_timeout_applies_grace_before_session_limits() {
 #[tokio::test(start_paused = true)]
 async fn execute_waits_for_nested_tool_during_yield_grace() {
    let delegate = Arc::new(ReleasableToolDelegate::default());
-    let service = InProcessCodeModeSession::with_delegate(delegate.clone());
+    let service = InProcessCodeModeSession::new();
    let request = ExecuteRequest {
        enabled_tools: vec![echo_tool()],
        source: r#"await tools.echo({}); text("done");"#.to_string(),
        yield_time_ms: Some(10_000),
        ..execute_request("")
    };
-    let started = service.execute(request).await.unwrap();
+    let started = service.execute(request, delegate.clone()).await.unwrap();
    let response = tokio::spawn(started.initial_response());
    wait_until_tool_started(&delegate).await;
    tokio::time::advance(Duration::from_millis(10_500)).await;
@@ -93,20 +91,20 @@ async fn execute_waits_for_nested_tool_during_yield_grace() {
 #[tokio::test(start_paused = true)]
 async fn execute_and_wait_clamp_yield_grace_without_stopping_the_cell() {
    let delegate = Arc::new(ReleasableToolDelegate::default());
-    let service = InProcessCodeModeSession::with_delegate_and_limits(
-        delegate.clone(),
-        CodeModeSessionCellExecutionLimits {
-            max_yield_time_ms: Some(/*value*/ 10_000),
-            max_heap_size_bytes: None,
-        },
-    );
+    let service = InProcessCodeModeSession::with_limits(CodeModeSessionCellExecutionLimits {
+        max_yield_time_ms: Some(/*value*/ 10_000),
+        max_heap_size_bytes: None,
+    });
    let started = service
-        .execute(ExecuteRequest {
-            enabled_tools: vec![echo_tool()],
-            source: r#"await tools.echo({}); text("done");"#.to_string(),
-            yield_time_ms: None,
-            ..execute_request("")
-        })
+        .execute(
+            ExecuteRequest {
+                enabled_tools: vec![echo_tool()],
+                source: r#"await tools.echo({}); text("done");"#.to_string(),
+                yield_time_ms: None,
+                ..execute_request("")
+            },
+            delegate.clone(),
+        )
        .await
        .unwrap();
    let initial_response = tokio::spawn(started.initial_response());
@@ -169,13 +167,16 @@ async fn execute_and_wait_clamp_yield_grace_without_stopping_the_cell() {
 #[tokio::test(start_paused = true)]
 async fn wait_waits_for_nested_tool_during_yield_grace() {
    let delegate = Arc::new(ReleasableToolDelegate::default());
-    let service = InProcessCodeModeSession::with_delegate(delegate.clone());
+    let service = InProcessCodeModeSession::new();
    let initial_response = service
-        .execute_to_pending(ExecuteRequest {
-            enabled_tools: vec![echo_tool()],
-            source: r#"await tools.echo({}); text("done");"#.to_string(),
-            ..execute_request("")
-        })
+        .execute_to_pending(
+            ExecuteRequest {
+                enabled_tools: vec![echo_tool()],
+                source: r#"await tools.echo({}); text("done");"#.to_string(),
+                ..execute_request("")
+            },
+            delegate.clone(),
+        )
        .await
        .unwrap();
    assert_eq!(
@@ -215,29 +216,30 @@ async fn wait_waits_for_nested_tool_during_yield_grace() {
 #[tokio::test(start_paused = true)]
 async fn zero_yield_limit_is_immediate_and_scoped_to_its_session() {
    let zero_delegate = Arc::new(ReleasableToolDelegate::default());
-    let zero_session = InProcessCodeModeSession::with_delegate_and_limits(
-        zero_delegate.clone(),
-        CodeModeSessionCellExecutionLimits {
-            max_yield_time_ms: Some(/*value*/ 0),
-            max_heap_size_bytes: None,
-        },
-    );
+    let zero_session = InProcessCodeModeSession::with_limits(CodeModeSessionCellExecutionLimits {
+        max_yield_time_ms: Some(/*value*/ 0),
+        max_heap_size_bytes: None,
+    });
    let limited_delegate = Arc::new(ReleasableToolDelegate::default());
-    let limited_session = InProcessCodeModeSession::with_delegate_and_limits(
-        limited_delegate.clone(),
-        CodeModeSessionCellExecutionLimits {
+    let limited_session =
+        InProcessCodeModeSession::with_limits(CodeModeSessionCellExecutionLimits {
            max_yield_time_ms: Some(/*value*/ 10),
            max_heap_size_bytes: None,
-        },
-    );
+        });
    let request = ExecuteRequest {
        enabled_tools: vec![echo_tool()],
        source: "await tools.echo({});".to_string(),
        yield_time_ms: Some(/*value*/ 60_000),
        ..execute_request("")
    };
-    let zero_started = zero_session.execute(request.clone()).await.unwrap();
-    let limited_started = limited_session.execute(request).await.unwrap();
+    let zero_started = zero_session
+        .execute(request.clone(), zero_delegate.clone())
+        .await
+        .unwrap();
+    let limited_started = limited_session
+        .execute(request, limited_delegate.clone())
+        .await
+        .unwrap();
    let zero_response = tokio::spawn(zero_started.initial_response());
    let limited_response = tokio::spawn(limited_started.initial_response());
    wait_until_tool_started(&zero_delegate).await;
@@ -372,7 +374,7 @@ fn echo_tool() -> ToolDefinition {

 async fn execute(service: &InProcessCodeModeSession, request: ExecuteRequest) -> RuntimeResponse {
    service
-        .execute(request)
+        .execute(request, Arc::new(NoopCodeModeSessionDelegate))
        .await
        .unwrap()
        .initial_response()
@@ -521,10 +523,13 @@ async fn shutdown_interrupts_cpu_bound_cells() {
    let service = InProcessCodeModeSession::new();

    let cell = service
-        .execute(ExecuteRequest {
-            source: "while (true) {}".to_string(),
-            ..execute_request("")
-        })
+        .execute(
+            ExecuteRequest {
+                source: "while (true) {}".to_string(),
+                ..execute_request("")
+            },
+            Arc::new(NoopCodeModeSessionDelegate),
+        )
        .await
        .unwrap();
    assert_eq!(
@@ -548,7 +553,10 @@ async fn start_cell_rejects_new_cell_after_shutdown_begins() {
    service.shutdown().await.unwrap();

    let error = service
-        .execute(execute_request("text('late');"))
+        .execute(
+            execute_request("text('late');"),
+            Arc::new(NoopCodeModeSessionDelegate),
+        )
        .await
        .err()
        .unwrap();
@@ -561,11 +569,14 @@ async fn execute_to_pending_returns_completed_for_synchronous_results() {
    let service = InProcessCodeModeSession::new();

    let response = service
-        .execute_to_pending(ExecuteRequest {
-            source: r#"text("done");"#.to_string(),
-            yield_time_ms: Some(60_000),
-            ..execute_request("")
-        })
+        .execute_to_pending(
+            ExecuteRequest {
+                source: r#"text("done");"#.to_string(),
+                yield_time_ms: Some(60_000),
+                ..execute_request("")
+            },
+            Arc::new(NoopCodeModeSessionDelegate),
+        )
        .await
        .unwrap();

@@ -588,11 +599,14 @@ async fn execute_to_pending_returns_once_the_runtime_is_quiescent() {

    let response = tokio::time::timeout(
        Duration::from_secs(1),
-        service.execute_to_pending(ExecuteRequest {
-            source: r#"text("before"); await new Promise(() => {});"#.to_string(),
-            yield_time_ms: Some(60_000),
-            ..execute_request("")
-        }),
+        service.execute_to_pending(
+            ExecuteRequest {
+                source: r#"text("before"); await new Promise(() => {});"#.to_string(),
+                yield_time_ms: Some(60_000),
+                ..execute_request("")
+            },
+            Arc::new(NoopCodeModeSessionDelegate),
+        ),
    )
    .await
    .unwrap()
@@ -626,18 +640,21 @@ async fn execute_to_pending_identifies_tool_calls_in_paused_frontier() {
    let service = InProcessCodeModeSession::new();

    let response = service
-        .execute_to_pending(ExecuteRequest {
-            enabled_tools: vec![echo_tool()],
-            source: r#"
+        .execute_to_pending(
+            ExecuteRequest {
+                enabled_tools: vec![echo_tool()],
+                source: r#"
 await Promise.all([
  tools.echo({ value: "first" }),
  tools.echo({ value: "second" }),
 ]);
 "#
-            .to_string(),
-            yield_time_ms: Some(60_000),
-            ..execute_request("")
-        })
+                .to_string(),
+                yield_time_ms: Some(60_000),
+                ..execute_request("")
+            },
+            Arc::new(NoopCodeModeSessionDelegate),
+        )
        .await
        .unwrap();

@@ -667,9 +684,10 @@ async fn execute_to_pending_excludes_delayed_timeout_tool_calls_until_wait() {
    let service = InProcessCodeModeSession::new();

    let initial_response = service
-        .execute_to_pending(ExecuteRequest {
-            enabled_tools: vec![echo_tool()],
-            source: r#"
+        .execute_to_pending(
+            ExecuteRequest {
+                enabled_tools: vec![echo_tool()],
+                source: r#"
 setTimeout(() => {
  tools.echo({ value: "delayed" });
 }, 1000);
@@ -678,10 +696,12 @@ await Promise.all([
  tools.echo({ value: "third" }),
 ]);
 "#
-            .to_string(),
-            yield_time_ms: Some(60_000),
-            ..execute_request("")
-        })
+                .to_string(),
+                yield_time_ms: Some(60_000),
+                ..execute_request("")
+            },
+            Arc::new(NoopCodeModeSessionDelegate),
+        )
        .await
        .unwrap();

@@ -730,20 +750,23 @@ await Promise.all([
 #[tokio::test]
 async fn wait_to_pending_returns_after_resumed_runtime_becomes_quiescent_again() {
    let delegate = Arc::new(ReleasableToolDelegate::default());
-    let service = InProcessCodeModeSession::with_delegate(delegate.clone());
+    let service = InProcessCodeModeSession::new();

    let initial_response = service
-        .execute_to_pending(ExecuteRequest {
-            enabled_tools: vec![echo_tool()],
-            source: r#"
+        .execute_to_pending(
+            ExecuteRequest {
+                enabled_tools: vec![echo_tool()],
+                source: r#"
 await tools.echo({});
 text("after");
 await new Promise(() => {});
 "#
-            .to_string(),
-            yield_time_ms: Some(60_000),
-            ..execute_request("")
-        })
+                .to_string(),
+                yield_time_ms: Some(60_000),
+                ..execute_request("")
+            },
+            delegate.clone(),
+        )
        .await
        .unwrap();

@@ -794,19 +817,22 @@ await new Promise(() => {});
 #[tokio::test]
 async fn wait_to_pending_returns_completed_after_resumed_runtime_finishes() {
    let delegate = Arc::new(ReleasableToolDelegate::default());
-    let service = InProcessCodeModeSession::with_delegate(delegate.clone());
+    let service = InProcessCodeModeSession::new();

    let initial_response = service
-        .execute_to_pending(ExecuteRequest {
-            enabled_tools: vec![echo_tool()],
-            source: r#"
+        .execute_to_pending(
+            ExecuteRequest {
+                enabled_tools: vec![echo_tool()],
+                source: r#"
 await tools.echo({});
 text("done");
 "#
-            .to_string(),
-            yield_time_ms: Some(60_000),
-            ..execute_request("")
-        })
+                .to_string(),
+                yield_time_ms: Some(60_000),
+                ..execute_request("")
+            },
+            delegate.clone(),
+        )
        .await
        .unwrap();

@@ -1271,9 +1297,7 @@ async fn audio_helper_rejects_non_data_urls() {
 async fn image_helper_accepts_raw_mcp_image_block_with_original_detail() {
    let service = InProcessCodeModeSession::new();

-    let response = execute(
-            &service,
-            ExecuteRequest {
+    let response = execute(&service, ExecuteRequest {
                source: r#"
 image({
  type: "image",
@@ -1285,8 +1309,7 @@ image({
                .to_string(),
                yield_time_ms: None,
                ..execute_request("")
-            },
-        )
+            })
        .await;

    assert_eq!(
@@ -1383,9 +1406,7 @@ image(
 async fn image_helper_second_arg_overrides_raw_mcp_image_detail() {
    let service = InProcessCodeModeSession::new();

-    let response = execute(
-            &service,
-            ExecuteRequest {
+    let response = execute(&service, ExecuteRequest {
                source: r#"
 image(
  {
@@ -1400,8 +1421,7 @@ image(
                .to_string(),
                yield_time_ms: None,
                ..execute_request("")
-            },
-        )
+            })
        .await;

    assert_eq!(
@@ -1559,9 +1579,7 @@ image({
 async fn image_helper_rejects_raw_mcp_result_container() {
    let service = InProcessCodeModeSession::new();

-    let response = execute(
-            &service,
-            ExecuteRequest {
+    let response = execute(&service, ExecuteRequest {
                source: r#"
 image({
  content: [
@@ -1578,8 +1596,7 @@ image({
                .to_string(),
                yield_time_ms: None,
                ..execute_request("")
-            },
-        )
+            })
        .await;

    assert_eq!(
@@ -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> {
-    inner: Arc<Inner<D>>,
+pub(crate) struct SessionRuntime {
+    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> {
-    pub(crate) fn new(delegate: Arc<D>) -> Self {
-        Self::new_with_task_failure_handler(delegate, /*task_failure_handler*/ None)
+impl SessionRuntime {
+    pub(crate) fn new() -> Self {
+        Self::new_with_task_failure_handler(/*task_failure_handler*/ None)
    }

    pub(crate) fn new_with_task_failure_handler(
-        delegate: Arc<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
-            .delegate
+        self.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
-            .delegate
+        self.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);
    }
 }

@@ -66,16 +66,14 @@ impl SessionRuntimeDelegate for PanickingClosedDelegate {
 #[tokio::test]
 async fn reports_cell_actor_panics_to_the_owner() {
    let (failure_tx, mut failure_rx) = tokio::sync::mpsc::unbounded_channel();
-    let runtime = SessionRuntime::new_with_task_failure_handler(
-        Arc::new(PanickingClosedDelegate),
-        Some(Arc::new(move |reason| {
-            let _ = failure_tx.send(reason);
-        })),
-    );
+    let runtime = SessionRuntime::new_with_task_failure_handler(Some(Arc::new(move |reason| {
+        let _ = failure_tx.send(reason);
+    })));
    let started = runtime
        .execute(
            execute_request(r#"text("done");"#),
            ObserveMode::YieldAfter(Duration::from_secs(1)),
+            Arc::new(PanickingClosedDelegate),
        )
        .await
        .expect("start cell");
@@ -97,9 +95,10 @@ async fn reports_cell_actor_panics_to_the_owner() {

 #[tokio::test]
 async fn termination_rejects_a_waiting_store_commit_before_the_next_cell_can_load_it() {
-    let runtime = SessionRuntime::new(Arc::new(RecordingDelegate));
+    let runtime = SessionRuntime::new();
    let cell_state = Arc::new(CellState::new(CancellationToken::new()));
    let host = RuntimeCellHost {
+        delegate: Arc::new(RecordingDelegate),
        cell_id: CellId::new("terminating-writer"),
        inner: Arc::clone(&runtime.inner),
        execution_context: opentelemetry::Context::new(),
@@ -154,6 +153,7 @@ async fn termination_rejects_a_waiting_store_commit_before_the_next_cell_can_loa
                source: r#"text(String(load("candidate")));"#.to_string(),
            },
            ObserveMode::YieldAfter(Duration::from_secs(1)),
+            Arc::new(RecordingDelegate),
        )
        .await
        .unwrap();
@@ -179,7 +179,7 @@ fn execute_request(source: &str) -> CreateCellRequest {

 #[tokio::test]
 async fn cell_id_allocation_fails_before_wrapping() {
-    let runtime = SessionRuntime::new(Arc::new(RecordingDelegate));
+    let runtime = SessionRuntime::new();
    runtime
        .inner
        .next_cell_id
@@ -190,6 +190,7 @@ async fn cell_id_allocation_fails_before_wrapping() {
            .execute(
                execute_request(r#"text("unreachable");"#),
                ObserveMode::YieldAfter(Duration::from_secs(1)),
+                Arc::new(RecordingDelegate)
            )
            .await
            .err(),
@@ -203,12 +204,13 @@ async fn cell_id_allocation_fails_before_wrapping() {
    reason = "test holds the registry lock to force admission ahead of shutdown"
 )]
 async fn shutdown_rejects_cell_admission_queued_before_the_registry_lock() {
-    let runtime = Arc::new(SessionRuntime::new(Arc::new(RecordingDelegate)));
+    let runtime = Arc::new(SessionRuntime::new());
    let cells = runtime.inner.cells.lock().await;

    let execution = runtime.execute(
        execute_request("while (true) {}"),
        ObserveMode::YieldAfter(Duration::from_millis(/*millis*/ 1)),
+        Arc::new(RecordingDelegate),
    );
    tokio::pin!(execution);
    std::future::poll_fn(|context| match execution.as_mut().poll(context) {
@@ -238,11 +240,12 @@ async fn shutdown_rejects_cell_admission_queued_before_the_registry_lock() {

 #[tokio::test]
 async fn drop_terminates_cells_when_the_registry_is_locked() {
-    let runtime = SessionRuntime::new(Arc::new(RecordingDelegate));
+    let runtime = SessionRuntime::new();
    let started = runtime
        .execute(
            execute_request("while (true) {}"),
            ObserveMode::YieldAfter(Duration::from_millis(/*millis*/ 1)),
+            Arc::new(RecordingDelegate),
        )
        .await
        .unwrap();