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
This commit is contained in:
Abhinav authored and copyberry committed 2026-09-11 17:00:07 +00:00
1 parent bc5957eac9
commit 3305c4f31d
30 files changed
+1027 -637

No files matched your search

@@ -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!(