Report exec runtime resets before accepting further requests (#52825)

## Why

Replacing a gRPC code-mode runtime loses stored values and running cells. Silently continuing after reconnection hides that loss, while prior external tool operations may already have completed.

## What changed

- Capture the runtime generation when `execute` and `wait` calls are created. If recovery replaces it, reject those calls with an explicit reset error, including concurrent calls and calls first polled after recovery.
- Explain the lost state and advise inspecting prior tool results before retrying. Allow a fresh `exec` call to continue on the replacement runtime.
- Keep `terminate` from reopening the runtime so cleanup cannot hide a reset from the next request.

## Testing

Extend gRPC recovery tests to cover reset errors after host restarts, missing sessions, and route loss, rejection of queued calls, and successful fresh execution with cleared stored state.

GitOrigin-RevId: 914cab3bb45e30a3f3b2afacb5631e42b18e21b0
This commit is contained in:
Jennifer Zhao
2026-10-10 08:51:21 +00:00
committed by copyberry
parent c3d3b142d1
commit 806d9732c9
4 changed files with 100 additions and 15 deletions
@@ -75,8 +75,10 @@ impl CodeModeSession for ReconnectableSession {
delegate: Arc<dyn CodeModeSessionDelegate>,
preempt: Option<CancellationToken>,
) -> CodeModeSessionResultFuture<'a, StartedCell> {
// Capture before polling so queued calls cannot silently cross a runtime reset.
let generation = self.inner.binding().map(|binding| binding.generation);
Box::pin(async move {
let binding = self.inner.get_or_open_binding().await?;
let binding = self.inner.binding_for_request(generation).await?;
let delegate = Arc::new(GenerationDelegate {
delegate,
generation: binding.generation,
@@ -91,8 +93,9 @@ impl CodeModeSession for ReconnectableSession {
request: WaitRequest,
preempt: Option<CancellationToken>,
) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
let generation = self.inner.binding().map(|binding| binding.generation);
Box::pin(async move {
let binding = self.inner.get_or_open_binding().await?;
let binding = self.inner.binding_for_request(generation).await?;
let request = WaitRequest {
cell_id: generation::remote_cell_id(binding.generation, &request.cell_id)?,
yield_time_ms: request.yield_time_ms,
@@ -104,7 +107,11 @@ impl CodeModeSession for ReconnectableSession {
fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
Box::pin(async move {
let binding = self.inner.get_or_open_binding().await?;
// Cleanup must not replace the runtime and hide its loss from the next request.
let binding = self
.inner
.binding()
.ok_or_else(|| SHUTDOWN_ERROR.to_string())?;
let cell_id = generation::remote_cell_id(binding.generation, &cell_id)?;
let outcome = binding.session.terminate(cell_id).await?;
Ok(generation::public_wait_outcome(binding.generation, outcome))
@@ -125,6 +132,18 @@ impl Drop for ReconnectableSession {
}
impl ReconnectInner {
async fn binding_for_request(&self, generation: Option<u64>) -> Result<SessionBinding, String> {
let binding = self.get_or_open_binding().await?;
if generation != Some(binding.generation) {
return Err(concat!(
"The exec runtime was replaced. Stored values and running cells were lost. ",
"Prior external tool operations may have completed; inspect their results before retrying. ",
"You may continue with a fresh exec call.",
).into());
}
Ok(binding)
}
async fn get_or_open_binding(&self) -> Result<SessionBinding, String> {
if self.shutdown_requested.is_cancelled() {
return Err(SHUTDOWN_ERROR.to_string());
@@ -148,12 +167,7 @@ impl ReconnectInner {
return Ok(binding);
}
let previous_binding = self
.binding
.lock()
.unwrap_or_else(PoisonError::into_inner)
.clone();
if let Some(binding) = previous_binding {
if let Some(binding) = self.binding() {
wait_for_watch(binding.session.inner.request_shutdown()).await?;
}
@@ -186,12 +200,15 @@ impl ReconnectInner {
}
fn live_binding(&self) -> Option<SessionBinding> {
self.binding()
.filter(|binding| !binding.session.inner.stopped.is_cancelled())
}
fn binding(&self) -> Option<SessionBinding> {
self.binding
.lock()
.unwrap_or_else(PoisonError::into_inner)
.as_ref()
.filter(|binding| !binding.session.inner.stopped.is_cancelled())
.cloned()
.clone()
}
fn request_shutdown(self: &Arc<Self>) -> ShutdownResultReceiver {