Add early yielding for code-mode observations (#48123)

## What changed

Add an optional `preempt` signal to `CodeModeSession::execute` and `CodeModeSession::wait` so callers can end an observation early while the cell continues running and remains available for later waits.

Support the signal in the in-process runtime and across both remote transports: `operation/yield` with the negotiated `yield-observation` capability for stdio, and `YieldObservation` for gRPC. Preserve signals received before gRPC observation registration, and retain normal timeout behavior for older hosts that lack support.

## Testing

Add runtime and transport tests covering early execute and wait yields followed by successful cell completion, yields before observation registration, and suppression of yield frames for older stdio hosts.

GitOrigin-RevId: 3f8049d8f0582be6b8f0832cc03f2d974a0b8f18
This commit is contained in:
pakrym-oai
2026-09-25 15:42:35 +00:00
committed by copyberry
parent 66c4e5a20d
commit 55543d8772
40 changed files with 1110 additions and 227 deletions
@@ -29,6 +29,10 @@ service CodeModeHost {
rpc Execute(ExecuteRequest) returns (stream ExecuteEvent);
rpc Wait(WaitRequest) returns (WaitResponse);
// Ends the current observation early without stopping its cell. The
// corresponding Execute or Wait still returns its ordinary outcome.
rpc YieldObservation(YieldObservationRequest) returns (YieldObservationResponse);
// Acknowledges that a canceled wait has retired before another wait starts.
rpc CancelWait(CancelWaitRequest) returns (CancelWaitResponse);
rpc Terminate(TerminateRequest) returns (WaitResponse);
@@ -189,6 +193,16 @@ message CancelWaitRequest {
message CancelWaitResponse {}
message YieldObservationRequest {
string session_id = 1;
oneof observation {
string execution_id = 2;
string wait_id = 3;
}
}
message YieldObservationResponse {}
message TerminateRequest {
string session_id = 1;
string cell_id = 2;
@@ -135,6 +135,9 @@ pub enum ClientToHost {
Request { id: RequestId, request: HostRequest },
#[serde(rename = "operation/cancel")]
CancelRequest { id: RequestId },
/// Yield the active observation without cancelling its cell.
#[serde(rename = "operation/yield")]
YieldRequest { id: RequestId },
#[serde(rename = "delegate/response")]
DelegateResponse {
id: DelegateRequestId,
@@ -15,6 +15,9 @@ pub const MAX_PENDING_DELEGATE_CALLS: usize = 1_024;
/// Negotiated support for cell execution resource limits on `session/open`.
pub const SESSION_RESOURCE_LIMITS_CAPABILITY: &str = "session-cell-execution-resource-limits";
/// Support for yielding an in-flight execute or wait observation.
pub const YIELD_OBSERVATION_CAPABILITY: &str = "yield-observation";
pub use codec::EncodedFrame;
pub use codec::FramedReader;
pub use codec::FramedWriter;
+9 -1
View File
@@ -152,13 +152,21 @@ impl CodeModeSessionDelegate for NoopCodeModeSessionDelegate {
/// must keep those values isolated. Implementations may execute cells
/// in-process or remotely.
pub trait CodeModeSession: Send + Sync {
/// Executes a cell, yielding its foreground observation if `preempt` is signaled.
/// The cell continues running after its observation yields.
fn execute<'a>(
&'a self,
request: ExecuteRequest,
delegate: Arc<dyn CodeModeSessionDelegate>,
preempt: Option<CancellationToken>,
) -> CodeModeSessionResultFuture<'a, StartedCell>;
fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome>;
/// Waits for a cell, yielding the observation if `preempt` is signaled while the cell continues.
fn wait<'a>(
&'a self,
request: WaitRequest,
preempt: Option<CancellationToken>,
) -> CodeModeSessionResultFuture<'a, WaitOutcome>;
fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome>;