Recover gRPC code-mode sessions after missing-session errors (#52235)

## Why

A host can reject a request for an unknown session while its lease streams remain open. Keeping that session binding prevents subsequent execution from recovering.

## What changed

Invalidate the session when `Execute` (including its response stream), `Wait`, or `Terminate` returns `NOT_FOUND`, allowing the next execution to open a fresh session. Return the failed call's error without replaying execution.

Document that these RPCs reserve `NOT_FOUND` for unknown sessions, while `CompleteToolCall` can use it for a missing invocation.

## Testing

Add regression coverage for recovery with lease streams still open, no replay of failed execution, and preservation of session state after `INVALID_ARGUMENT`. Add a cancellation-race test confirming that a late tool-completion `NOT_FOUND` leaves the session open.

GitOrigin-RevId: 85b38fdbd6d59268447b6aa6a85ea7c60ec869a7

codex-rs/code-mode/src/grpc_session/callbacks.rs

@@ -193,6 +193,8 @@ impl SessionInner {
         tokio::select! {
             biased;
             _ = cancellation.cancelled() => {}
+            // NotFound can mean an invocation in a live session was canceled.
+            // Check cancellation before treating a completion error as fatal.
             result = deadline::request(
                 self,
                 "tool invocation completion",
@@ -266,3 +268,7 @@ impl SessionInner {
         Ok(())
     }
 }
+
+#[cfg(test)]
+#[path = "callbacks_tests.rs"]
+mod tests;

codex-rs/code-mode/src/grpc_session/callbacks_tests.rs

@@ -0,0 +1,50 @@
+//! Keeps late tool-completion errors scoped to the canceled invocation.
+
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::sync::Mutex;
+use std::sync::atomic::AtomicBool;
+
+use pretty_assertions::assert_eq;
+use tokio_util::sync::CancellationToken;
+use tokio_util::task::TaskTracker;
+use tower::service_fn;
+use tower::util::BoxCloneSyncService;
+
+use super::super::GrpcClient;
+use super::super::state::SessionState;
+use super::super::transport::SharedTransport;
+use super::SessionInner;
+
+#[tokio::test]
+async fn missing_invocation_completion_racing_cancellation_preserves_session() {
+    let cancellation = CancellationToken::new();
+    let completion_cancellation = cancellation.clone();
+    let client = GrpcClient::new(BoxCloneSyncService::new(service_fn(move |_request| {
+        // Cancel after the biased select has polled cancellation, before the RPC fails.
+        completion_cancellation.cancel();
+        async { Ok(tonic::Status::not_found("unknown code-mode tool invocation").into_http()) }
+    })));
+    let session = SessionInner {
+        id: "session".to_string(),
+        client: client.clone(),
+        runtime: tokio::runtime::Handle::current(),
+        state: Mutex::new(SessionState::default()),
+        wait_slots: Mutex::new(HashMap::new()),
+        shutdown_requested: AtomicBool::new(/*v*/ false),
+        shutdown_result: Mutex::new(/*t*/ None),
+        stopped: CancellationToken::new(),
+        stream_tasks: TaskTracker::new(),
+        _transport: Arc::new(SharedTransport::Connected(client)),
+    };
+
+    session
+        .complete_tool_call(
+            "invocation".to_string(),
+            cancellation,
+            Ok(serde_json::Value::Null),
+        )
+        .await;
+
+    assert_eq!(session.require_open(), Ok(()));
+}

codex-rs/code-mode/src/grpc_session/deadline.rs

@@ -42,6 +42,25 @@ pub(super) async fn request<T>(
     }
 }
 
+/// Invalidates the session on RPCs where `NotFound` means the session is missing.
+pub(super) async fn session_request<T>(
+    session: &SessionInner,
+    operation: &str,
+    runtime_timeout: Duration,
+    rpc: impl Future<Output = Result<T, tonic::Status>>,
+) -> Result<T, String> {
+    request(session, operation, runtime_timeout, async {
+        let result = rpc.await;
+        if let Err(error) = &result
+            && error.code() == tonic::Code::NotFound
+        {
+            session.fail(failure(operation, error));
+        }
+        result
+    })
+    .await
+}
+
 pub(super) fn failure(operation: &str, error: impl Display) -> String {
     let mut message = format!("gRPC code-mode {operation} failed: {error}");
     if message.len() > MAX_ERROR_BYTES {

codex-rs/code-mode/src/grpc_session/operations.rs

@@ -130,11 +130,15 @@ impl SessionInner {
             {
                 request.metadata_mut().insert("traceparent", traceparent);
             }
-            let mut stream =
-                deadline::request(&self, "execution", Duration::ZERO, client.execute(request))
-                    .await?
-                    .into_inner();
-            let first = deadline::request(
+            let mut stream = deadline::session_request(
+                &self,
+                "execution",
+                Duration::ZERO,
+                client.execute(request),
+            )
+            .await?
+            .into_inner();
+            let first = deadline::session_request(
                 &self,
                 "execution starting event",
                 Duration::ZERO,
@@ -217,7 +221,7 @@ impl SessionInner {
         let outcome = tokio::select! {
             biased;
             _ = response_tx.closed() => return,
-            outcome = deadline::request(
+            outcome = deadline::session_request(
                 &self,
                 "execution outcome",
                 runtime_timeout,
@@ -317,7 +321,8 @@ impl SessionInner {
                 grpc::yield_observation_request::Observation::WaitId(wait_id),
             )
         });
-        let response = deadline::request(self, "wait", runtime_timeout, client.wait(request)).await;
+        let response =
+            deadline::session_request(self, "wait", runtime_timeout, client.wait(request)).await;
         drop(watcher);
         cancellation.disarm();
         self.prune_wait_slots();
@@ -351,7 +356,7 @@ impl SessionInner {
     pub(super) async fn terminate(&self, cell_id: CellId) -> Result<WaitOutcome, String> {
         self.require_open()?;
         let mut client = self.client();
-        let response = deadline::request(
+        let response = deadline::session_request(
             self,
             "termination",
             Duration::ZERO,