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
2026-09-11 17:00:07 +00:00
committed by copyberry
parent bc5957eac9
commit 3305c4f31d
30 changed files with 1027 additions and 637 deletions
+7 -7
View File
@@ -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>(
&'a self,
delegate: Arc<dyn CodeModeSessionDelegate>,
) -> CodeModeSessionProviderFuture<'a>;
fn create_session(&self) -> CodeModeSessionProviderFuture<'_>;
/// Creates a session whose cells share the supplied execution limits.
///
@@ -182,11 +183,10 @@ pub trait CodeModeSessionProvider: Send + Sync {
/// explicitly implement this method before accepting non-default limits.
fn create_session_with_limits<'a>(
&'a self,
delegate: Arc<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())