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-host/tests/grpc.rs

@@ -52,6 +52,8 @@ mod large_tool_delegate;
 mod network_policy_tests;
 #[path = "support/recording_delegate.rs"]
 mod recording_delegate;
+#[path = "grpc/session_recovery_tests.rs"]
+mod session_recovery_tests;
 
 use host::HostHarness;
 use large_tool_delegate::LargeToolResultDelegate;

codex-rs/code-mode-host/tests/grpc/session_recovery_tests.rs

@@ -0,0 +1,101 @@
+//! Exercises missing-session recovery while the host's lease streams remain open.
+
+use std::sync::Arc;
+use std::sync::Mutex;
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
+
+use anyhow::Result;
+use codex_code_mode::CodeModeSessionProvider;
+use codex_code_mode::GrpcCodeModeSessionProvider;
+use codex_code_mode::NoopCodeModeSessionDelegate;
+use codex_code_mode_host::GrpcCodeModeHost;
+use codex_code_mode_protocol::grpc::code_mode_host_server::CodeModeHostServer;
+use pretty_assertions::assert_eq;
+use tokio::net::TcpListener;
+use tokio_stream::wrappers::TcpListenerStream;
+use tokio_util::task::AbortOnDropHandle;
+use tonic::Code;
+use tonic::Status;
+use tonic::transport::Server;
+
+use super::execute;
+use super::request;
+use super::text_response;
+
+#[tokio::test]
+async fn only_missing_sessions_reopen_without_replaying_execution() -> Result<()> {
+    for (code, expected_cell, expected_value) in [
+        (Code::NotFound, "g2:1", "undefined"),
+        (Code::InvalidArgument, "2", "1"),
+    ] {
+        let rejection = Arc::new(Mutex::new(/*t*/ None));
+        let next_rejection = Arc::clone(&rejection);
+        let request_count = Arc::new(AtomicUsize::new(/*v*/ 0));
+        let count = Arc::clone(&request_count);
+        let service = CodeModeHostServer::with_interceptor(
+            GrpcCodeModeHost::new(),
+            move |request: tonic::Request<()>| {
+                count.fetch_add(/*val*/ 1, Ordering::SeqCst);
+                match next_rejection.lock().unwrap().take() {
+                    Some(rejection) => Err(rejection),
+                    None => Ok(request),
+                }
+            },
+        );
+        let listener = TcpListener::bind("127.0.0.1:0").await?;
+        let endpoint = format!("http://{}", listener.local_addr()?);
+        let _server = AbortOnDropHandle::new(tokio::spawn(
+            Server::builder()
+                .add_service(service)
+                .serve_with_incoming(TcpListenerStream::new(listener)),
+        ));
+        let session = GrpcCodeModeSessionProvider::new(endpoint)
+            .create_session()
+            .await
+            .map_err(anyhow::Error::msg)?;
+        execute(
+            &session,
+            request(r#"store("before", 1);"#),
+            Arc::new(NoopCodeModeSessionDelegate),
+        )
+        .await?;
+
+        // Reject the next Execute without closing the lease or subscription streams.
+        let failure = Status::new(code, "injected RPC rejection");
+        *rejection.lock().unwrap() = Some(failure.clone());
+        let requests_before_rejection = request_count.load(Ordering::SeqCst);
+        let failed = execute(
+            &session,
+            request(r#"store("before", "replayed");"#),
+            Arc::new(NoopCodeModeSessionDelegate),
+        )
+        .await
+        .unwrap_err();
+        assert_eq!(
+            failed.to_string(),
+            format!("gRPC code-mode execution failed: {failure}")
+        );
+        assert_eq!(
+            request_count.load(Ordering::SeqCst),
+            requests_before_rejection + 1,
+            "the failed call must return without opening a session or replaying Execute"
+        );
+        let actual = execute(
+            &session,
+            request(r#"text(String(load("before")))"#),
+            Arc::new(NoopCodeModeSessionDelegate),
+        )
+        .await?;
+        assert_eq!(
+            actual,
+            text_response(
+                expected_cell,
+                expected_value,
+                actual.code_mode_host_duration()
+            )
+        );
+        session.shutdown().await.map_err(anyhow::Error::msg)?;
+    }
+    Ok(())
+}

codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto

@@ -5,6 +5,10 @@ package codex.code_mode.v1;
 // Hosts stateful JavaScript execution and delegates nested tool calls to the
 // session owner. Large tool inputs and outputs stay off the session event
 // stream so independent HTTP/2 streams can make progress concurrently.
+//
+// Execute (including its stream), Wait, and Terminate reserve NOT_FOUND for an
+// unknown session_id. Clients may retire the session binding on this status.
+// Missing cells use WaitResponse.missing_cell instead.
 service CodeModeHost {
   // Opens a session lease. The first event is always SessionOpened; dropping
   // this stream closes the session and terminates its active cells.
@@ -19,6 +23,7 @@ service CodeModeHost {
 
   // Each result receives its own HTTP/2 stream, preventing a large response
   // from blocking unrelated tool completions or session control events.
+  // NOT_FOUND may identify a missing invocation rather than a missing session.
   rpc CompleteToolCall(CompleteToolCallRequest)
       returns (CompleteToolCallResponse);
   rpc AcknowledgeNotification(AcknowledgeNotificationRequest)

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,