Preserve per-session routing for gRPC code mode (#52381)

## Why

Sessions sharing a transport can be admitted to different hosts. Each session needs its own routing token so subsequent RPCs reach the host holding its state.

## What changed

- Capture the optional `x-code-mode-route` metadata from successful `OpenSession` responses and attach it to subsequent session RPCs. Reject tokens longer than 256 bytes and keep new admission attempts free of existing session routes.
- Invalidate a routed session when an `Unavailable` response explicitly sets `x-code-mode-route-lost: true`, allowing a later operation to reopen it without replaying the failed execution. Ordinary `Unavailable` responses preserve the session.
- Allow up to five admission attempts for `Unavailable` and `ResourceExhausted`, with an immediate first retry followed by jittered exponential backoff.
- Order session failure with shutdown so late callback errors cannot cancel an in-flight `CloseSession` RPC.

## Testing

Add regression coverage for independent routing on a shared provider, route-loss recovery without execution replay, preservation of unaffected session state, and callback failures racing shutdown. Extend admission tests to verify the five-attempt limit and recovery tests to distinguish explicit route loss from ordinary unavailability.

GitOrigin-RevId: 5d60ef1afa56beb8c70ba7ab27865128225b0d3d

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

@@ -23,6 +23,7 @@ use codex_code_mode_protocol::grpc;
 use codex_code_mode_protocol::grpc::code_mode_host_client::CodeModeHostClient;
 use codex_http_client::HttpClientFactory;
 use codex_http_client::OutboundProxyPolicy;
+use rand::Rng;
 use tokio::sync::watch;
 use tokio_util::sync::CancellationToken;
 use tokio_util::task::TaskTracker;
@@ -50,6 +51,7 @@ mod transport;
 type GrpcClient = CodeModeHostClient<GrpcTransport>;
 
 const SHUTDOWN_ERROR: &str = "code mode session is shutting down";
+const ROUTE_HEADER: &str = "x-code-mode-route";
 
 fn inject_span_traceparent<T>(request: &mut tonic::Request<T>, span: &tracing::Span) {
     if let Some(traceparent) =
@@ -112,7 +114,7 @@ impl GrpcCodeModeSessionProvider {
             || limits.max_heap_size_bytes.is_some())
         .then_some(limits);
         let open_session_span = tracing::info_span!("code_mode.grpc.open_session");
-        let (lease, first) = deadline::startup("session opening", async {
+        let (lease, first, route) = deadline::startup("session opening", async {
             let mut retries = 0;
             loop {
                 let mut request = tonic::Request::new(grpc::OpenSessionRequest {
@@ -120,15 +122,25 @@ impl GrpcCodeModeSessionProvider {
                 });
                 inject_span_traceparent(&mut request, &open_session_span);
                 let result = async {
-                    let mut lease = client.open_session(request).await?.into_inner();
+                    let response = client.open_session(request).await?;
+                    let route = response.metadata().get(ROUTE_HEADER).cloned();
+                    if route
+                        .as_ref()
+                        .is_some_and(|value| value.as_bytes().len() > 256)
+                    {
+                        return Err(tonic::Status::invalid_argument(
+                            "session routing token is too long",
+                        ));
+                    }
+                    let mut lease = response.into_inner();
                     let first = lease.message().await?.ok_or_else(|| {
                         tonic::Status::unavailable("session lease ended before opening")
                     })?;
-                    Ok::<_, tonic::Status>((lease, first))
+                    Ok::<_, tonic::Status>((lease, first, route))
                 }
                 .await;
                 let Err(error) = &result else { return result };
-                if retries == 2
+                if retries == 4
                     || !matches!(
                         error.code(),
                         tonic::Code::Unavailable | tonic::Code::ResourceExhausted
@@ -138,9 +150,14 @@ impl GrpcCodeModeSessionProvider {
                 }
                 // No execution has started, and dropping a failed lease releases
                 // any partial admission. Never apply these retries to Execute.
+                tracing::warn!(%error, retry = retries + 1, "retrying code-mode session admission");
+                if retries > 0 {
+                    let delay_ms = 50_u64 << (retries - 1);
+                    let jitter = rand::rng().random_range(0.9..1.1);
+                    let delay = Duration::from_millis((delay_ms as f64 * jitter) as u64);
+                    tokio::time::sleep(delay).await;
+                }
                 retries += 1;
-                tracing::warn!(%error, retries, "retrying code-mode session admission");
-                tokio::time::sleep(Duration::from_millis(100 * retries)).await;
             }
         })
         .instrument(open_session_span.clone())
@@ -153,6 +170,7 @@ impl GrpcCodeModeSessionProvider {
         let inner = Arc::new(SessionInner {
             id: opened.session_id,
             client,
+            route,
             runtime: tokio::runtime::Handle::current(),
             state: Mutex::new(SessionState::default()),
             wait_slots: Mutex::new(HashMap::new()),
@@ -171,7 +189,7 @@ impl GrpcCodeModeSessionProvider {
             "code_mode.grpc.subscribe_to_tool_calls",
             session.id = %inner.id,
         );
-        let mut request = tonic::Request::new(grpc::SubscribeToToolCallsRequest {
+        let mut request = inner.request(grpc::SubscribeToToolCallsRequest {
             session_id: inner.id.clone(),
             tool_names: Vec::new(),
         });
@@ -269,6 +287,9 @@ impl Drop for GrpcCodeModeSession {
 pub(super) struct SessionInner {
     pub(super) id: String,
     pub(super) client: GrpcClient,
+    // Affinity belongs to the session: parents and children may share a transport
+    // while being admitted to different hosts. The token is opaque to Codex.
+    route: Option<tonic::metadata::MetadataValue<tonic::metadata::Ascii>>,
     runtime: tokio::runtime::Handle,
     state: Mutex<SessionState>,
     wait_slots: Mutex<HashMap<CellId, Weak<WaitSlot>>>,
@@ -284,6 +305,14 @@ impl SessionInner {
         self.client.clone()
     }
 
+    pub(super) fn request<T>(&self, message: T) -> tonic::Request<T> {
+        let mut request = tonic::Request::new(message);
+        if let Some(route) = &self.route {
+            request.metadata_mut().insert(ROUTE_HEADER, route.clone());
+        }
+        request
+    }
+
     pub(super) fn require_open(&self) -> Result<(), String> {
         self.state
             .lock()
@@ -312,11 +341,15 @@ impl SessionInner {
     }
 
     fn close_state(&self, failure: Option<String>) {
-        let cells = self
-            .state
-            .lock()
-            .unwrap_or_else(PoisonError::into_inner)
-            .close(failure);
+        let cells = {
+            let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
+            // Order failure with request_shutdown. Once shutdown starts,
+            // CloseSession owns cleanup; a late error must not cancel its RPC.
+            if failure.is_some() && self.shutdown_requested.load(Ordering::Acquire) {
+                return;
+            }
+            state.close(failure)
+        };
         self.stopped.cancel();
         self.stream_tasks.close();
         for cell_id in cells {
@@ -325,7 +358,10 @@ impl SessionInner {
     }
 
     fn request_shutdown(self: &Arc<Self>) -> ShutdownResultReceiver {
-        self.shutdown_requested.store(true, Ordering::Release);
+        {
+            let _state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
+            self.shutdown_requested.store(true, Ordering::Release);
+        }
         let mut result = self
             .shutdown_result
             .lock()
@@ -356,9 +392,9 @@ impl SessionInner {
                 self,
                 "session shutdown",
                 Duration::ZERO,
-                client.close_session(grpc::CloseSessionRequest {
+                client.close_session(self.request(grpc::CloseSessionRequest {
                     session_id: self.id.clone(),
-                }),
+                })),
             )
             .await
             .map(|_| ())