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/Cargo.lock

@@ -2717,6 +2717,7 @@ dependencies = [
  "http-body-util",
  "pretty_assertions",
  "prost",
+ "rand 0.9.3",
  "reqwest 0.12.28",
  "serde_json",
  "tokio",

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

@@ -11,7 +11,7 @@ async fn session_admission_retries_only_transient_failures_and_is_bounded() -> R
     for (code, failures, expected_attempts, succeeds) in [
         (Code::Unavailable, 1, 2, true),
         (Code::ResourceExhausted, 1, 2, true),
-        (Code::Unavailable, usize::MAX, 3, false),
+        (Code::Unavailable, usize::MAX, 5, false),
         (Code::InvalidArgument, 1, 1, false),
         (Code::Unknown, 1, 1, false),
     ] {
@@ -106,3 +106,198 @@ async fn session_admission_recovers_after_connection_reset() -> Result<()> {
     drop(server);
     Ok(())
 }
+
+#[tokio::test]
+async fn routing_metadata_is_scoped_to_each_session_on_a_shared_provider() -> Result<()> {
+    use axum::body::Body;
+    use axum::extract::Request;
+    use axum::response::Response;
+    use tonic::codegen::Service;
+
+    let backends = Arc::new(["a", "b"].map(|_| {
+        axum::Router::new().fallback_service(grpc::code_mode_host_server::CodeModeHostServer::new(
+            codex_code_mode_host::GrpcCodeModeHost::new(),
+        ))
+    }));
+    let admissions = Arc::new(AtomicUsize::new(0));
+    let rejection = Arc::new(std::sync::Mutex::new(/*t*/ None::<tonic::Status>));
+    let rejected_executions = Arc::new(AtomicUsize::new(0));
+    let next_rejection = Arc::clone(&rejection);
+    let executions = Arc::clone(&rejected_executions);
+    let opens = Arc::clone(&admissions);
+    let app = axum::Router::new().fallback(move |request: Request| {
+        let backends = Arc::clone(&backends);
+        let admissions = Arc::clone(&admissions);
+        let rejection = Arc::clone(&next_rejection);
+        let executions = Arc::clone(&executions);
+        async move {
+            let opening = request.uri().path().ends_with("/OpenSession");
+            let route = request.headers().get("x-code-mode-route");
+            let token = if opening {
+                assert_eq!(
+                    route, None,
+                    "a new admission must not inherit any session's route"
+                );
+                match admissions.fetch_add(1, Ordering::SeqCst) {
+                    0 => {
+                        // Even a rejected admission can carry response metadata. The
+                        // next attempt must still be free to select another backend.
+                        let mut response =
+                            tonic::Status::resource_exhausted("full").into_http::<Body>();
+                        response
+                            .headers_mut()
+                            .insert("x-code-mode-route", "rejected".parse().unwrap());
+                        return response;
+                    }
+                    1 => "a",
+                    _ => "b",
+                }
+            } else {
+                route
+                    .expect("every session RPC needs its routing token")
+                    .to_str()
+                    .unwrap()
+            };
+            let index = match token {
+                "a" => 0,
+                "b" => 1,
+                _ => panic!("unexpected routing token"),
+            };
+            if token == "a"
+                && request.uri().path().ends_with("/Execute")
+                && let Some(error) = rejection.lock().unwrap().as_ref()
+            {
+                executions.fetch_add(1, Ordering::SeqCst);
+                return error.clone().into_http::<Body>();
+            }
+            let mut response: Response = backends[index].clone().call(request).await.unwrap();
+            if opening {
+                response
+                    .headers_mut()
+                    .insert("x-code-mode-route", ["a", "b"][index].parse().unwrap());
+            }
+            response
+        }
+    });
+    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
+    let address = listener.local_addr()?;
+    let _server = tokio_util::task::AbortOnDropHandle::new(tokio::spawn(async move {
+        axum::serve(listener, app)
+            .await
+            .expect("routing test server");
+    }));
+    let provider = GrpcCodeModeSessionProvider::new(format!("http://{address}"));
+    let parent = provider
+        .create_session()
+        .await
+        .map_err(anyhow::Error::msg)?;
+    let child = provider
+        .create_session()
+        .await
+        .map_err(anyhow::Error::msg)?;
+    for (session, value) in [(&parent, 11), (&child, 22), (&parent, 33)] {
+        let mut req = request(&format!(
+            "await tools.echo({{value: {value}}}); text({value});"
+        ));
+        req.enabled_tools = vec![tool("echo")];
+        let actual = execute(session, req, Arc::new(RecordingDelegate::default())).await?;
+        let RuntimeResponse::Result { error_text, .. } = actual else {
+            anyhow::bail!("session did not finish");
+        };
+        assert_eq!(error_text, None);
+        let mut req = request("await new Promise(() => {});");
+        req.yield_time_ms = Some(1);
+        let started = session
+            .execute(
+                req,
+                Arc::new(NoopCodeModeSessionDelegate),
+                /*preempt*/ None,
+            )
+            .await
+            .map_err(anyhow::Error::msg)?;
+        let id = started.cell_id.clone();
+        assert!(matches!(
+            started
+                .initial_response()
+                .await
+                .map_err(anyhow::Error::msg)?,
+            RuntimeResponse::Yielded { .. }
+        ));
+        let outcome = session
+            .wait(
+                WaitRequest {
+                    cell_id: id.clone(),
+                    yield_time_ms: 1,
+                },
+                /*preempt*/ None,
+            )
+            .await
+            .map_err(anyhow::Error::msg)?;
+        assert!(matches!(
+            outcome,
+            WaitOutcome::LiveCell(RuntimeResponse::Yielded { .. })
+        ));
+        session.terminate(id).await.map_err(anyhow::Error::msg)?;
+    }
+    for session in [&parent, &child] {
+        execute(
+            session,
+            request(r#"store("preserved", "original");"#),
+            Arc::new(NoopCodeModeSessionDelegate),
+        )
+        .await?;
+    }
+    for retired in [false, true] {
+        let mut error = tonic::Status::unavailable("no healthy upstream");
+        if retired {
+            error
+                .metadata_mut()
+                .insert("x-code-mode-route-lost", "true".parse()?);
+        }
+        *rejection.lock().unwrap() = Some(error);
+        let opens_before = opens.load(Ordering::SeqCst);
+        let failures_before = rejected_executions.load(Ordering::SeqCst);
+        execute(
+            &parent,
+            request(r#"store("preserved", "replayed");"#),
+            Arc::new(NoopCodeModeSessionDelegate),
+        )
+        .await
+        .expect_err("the rejected execution must not be replayed");
+        assert_eq!(opens.load(Ordering::SeqCst), opens_before);
+        assert_eq!(
+            rejected_executions.load(Ordering::SeqCst),
+            failures_before + 1
+        );
+        if !retired {
+            *rejection.lock().unwrap() = None;
+        }
+        // The old lease and callback streams are still open. Only the explicit
+        // route-loss response may replace the parent; the child must keep its state.
+        for (session, expected) in [
+            (&parent, if retired { "undefined" } else { "original" }),
+            (&child, "original"),
+        ] {
+            let actual = execute(
+                session,
+                request(r#"text(String(load("preserved")));"#),
+                Arc::new(NoopCodeModeSessionDelegate),
+            )
+            .await?;
+            let RuntimeResponse::Result { cell_id, .. } = &actual else {
+                anyhow::bail!("execution did not finish");
+            };
+            assert_eq!(
+                actual,
+                text_response(cell_id.as_str(), expected, actual.code_mode_host_duration())
+            );
+        }
+        assert_eq!(
+            opens.load(Ordering::SeqCst),
+            opens_before + usize::from(retired)
+        );
+    }
+    parent.shutdown().await.map_err(anyhow::Error::msg)?;
+    child.shutdown().await.map_err(anyhow::Error::msg)?;
+    Ok(())
+}

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

@@ -25,9 +25,11 @@ 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"),
+    for (code, route_lost, expected_cell, expected_value) in [
+        (Code::NotFound, false, "g2:1", "undefined"),
+        (Code::InvalidArgument, false, "2", "1"),
+        (Code::Unavailable, false, "2", "1"),
+        (Code::Unavailable, true, "2", "1"),
     ] {
         let rejection = Arc::new(Mutex::new(/*t*/ None));
         let next_rejection = Arc::clone(&rejection);
@@ -62,7 +64,13 @@ async fn only_missing_sessions_reopen_without_replaying_execution() -> Result<()
         .await?;
 
         // Reject the next Execute without closing the lease or subscription streams.
-        let failure = Status::new(code, "injected RPC rejection");
+        let mut failure = Status::new(code, "injected RPC rejection");
+        if route_lost {
+            // A direct host has no proxy route to retire.
+            failure
+                .metadata_mut()
+                .insert("x-code-mode-route-lost", "true".parse()?);
+        }
         *rejection.lock().unwrap() = Some(failure.clone());
         let requests_before_rejection = request_count.load(Ordering::SeqCst);
         let failed = execute(

codex-rs/code-mode/Cargo.toml

@@ -21,6 +21,7 @@ codex-protocol = { workspace = true }
 futures = { workspace = true }
 http-body-util = "0.1.3"
 prost = "0.14.3"
+rand = { workspace = true }
 reqwest = { workspace = true, features = ["stream"] }
 serde_json = { workspace = true }
 tokio = { workspace = true, features = ["io-util", "macros", "net", "process", "rt", "sync", "time"] }

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

@@ -1,7 +1,6 @@
 use std::panic::AssertUnwindSafe;
 use std::sync::Arc;
 use std::sync::PoisonError;
-use std::sync::atomic::Ordering;
 use std::time::Duration;
 
 use codex_code_mode_protocol::CellId;
@@ -55,15 +54,11 @@ impl SessionInner {
                         }
                     }
                     Ok(None) => {
-                        if !inner.shutdown_requested.load(Ordering::Acquire) {
-                            inner.fail(format!("gRPC code-mode {stream_name} closed unexpectedly"));
-                        }
+                        inner.fail(format!("gRPC code-mode {stream_name} closed unexpectedly"));
                         return;
                     }
                     Err(error) => {
-                        if !inner.shutdown_requested.load(Ordering::Acquire) {
-                            inner.fail(deadline::failure(stream_name, error));
-                        }
+                        inner.fail(deadline::failure(stream_name, error));
                         return;
                     }
                 }
@@ -199,7 +194,7 @@ impl SessionInner {
                 self,
                 "tool invocation completion",
                 Duration::ZERO,
-                client.complete_tool_call(request),
+                client.complete_tool_call(self.request(request)),
             ) => {
                 if let Err(error) = result
                     && !cancellation.is_cancelled()

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

@@ -1,13 +1,19 @@
-//! Keeps late tool-completion errors scoped to the canceled invocation.
+//! Keeps late callback failures from disrupting cancellation and shutdown.
 
 use std::collections::HashMap;
 use std::sync::Arc;
 use std::sync::Mutex;
 use std::sync::atomic::AtomicBool;
+use std::time::Duration;
 
+use http_body_util::Full;
 use pretty_assertions::assert_eq;
+use tokio::sync::Notify;
 use tokio_util::sync::CancellationToken;
 use tokio_util::task::TaskTracker;
+use tonic::body::Body;
+use tonic::codegen::Bytes;
+use tonic::codegen::http::Response;
 use tower::service_fn;
 use tower::util::BoxCloneSyncService;
 
@@ -15,6 +21,7 @@ use super::super::GrpcClient;
 use super::super::state::SessionState;
 use super::super::transport::SharedTransport;
 use super::SessionInner;
+use crate::remote_session::wait_for_watch;
 
 #[tokio::test]
 async fn missing_invocation_completion_racing_cancellation_preserves_session() {
@@ -25,18 +32,7 @@ async fn missing_invocation_completion_racing_cancellation_preserves_session() {
         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)),
-    };
+    let session = session(client);
 
     session
         .complete_tool_call(
@@ -48,3 +44,59 @@ async fn missing_invocation_completion_racing_cancellation_preserves_session() {
 
     assert_eq!(session.require_open(), Ok(()));
 }
+
+fn session(client: GrpcClient) -> Arc<SessionInner> {
+    Arc::new(SessionInner {
+        id: "session".to_string(),
+        client: client.clone(),
+        route: None,
+        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)),
+    })
+}
+
+#[tokio::test(start_paused = true)]
+async fn shutdown_is_ordered_with_session_failure() {
+    let close_started = Arc::new(Notify::new());
+    let finish_close = Arc::new(Notify::new());
+    let started = Arc::clone(&close_started);
+    let finish = Arc::clone(&finish_close);
+    let client = GrpcClient::new(BoxCloneSyncService::new(service_fn(move |_request| {
+        let started = Arc::clone(&started);
+        let finish = Arc::clone(&finish);
+        async move {
+            started.notify_one();
+            finish.notified().await;
+            Ok(Response::builder()
+                .header("content-type", "application/grpc")
+                .header("grpc-status", "0")
+                // An empty protobuf CloseSessionResponse in a gRPC frame.
+                .body(Body::new(Full::new(Bytes::from_static(b"\0\0\0\0\0"))))
+                .expect("close response"))
+        }
+    })));
+    let closing = session(client.clone());
+    let result = closing.request_shutdown();
+    tokio::time::timeout(Duration::from_secs(5), close_started.notified())
+        .await
+        .expect("CloseSession started");
+
+    // A callback can decide to fail just before shutdown starts, then reach
+    // fail() after CloseSession is in flight. Recheck at the state transition.
+    closing.fail("late callback failure".to_string());
+    assert!(!closing.stopped.is_cancelled());
+    finish_close.notify_one();
+    assert_eq!(wait_for_watch(result).await, Ok(()));
+
+    // If failure wins first, shutdown observes a closed session and skips the RPC.
+    let failed = session(client);
+    failed.fail("earlier callback failure".to_string());
+    assert!(failed.stopped.is_cancelled());
+    assert_eq!(wait_for_watch(failed.request_shutdown()).await, Ok(()));
+}

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

@@ -34,7 +34,21 @@ pub(super) async fn request<T>(
 
     match result {
         Ok(value) => Ok(value),
-        Err(RequestError::Failed(error)) => Err(failure(operation, error)),
+        Err(RequestError::Failed(error)) => {
+            let reason = failure(operation, &error);
+            // Only the proxy can declare a pinned route lost. Ordinary RPC and
+            // connection failures must not discard a potentially live session.
+            if session.route.is_some()
+                && error.code() == tonic::Code::Unavailable
+                && error
+                    .metadata()
+                    .get("x-code-mode-route-lost")
+                    .is_some_and(|value| value == "true")
+            {
+                session.fail(reason.clone());
+            }
+            Err(reason)
+        }
         Err(RequestError::TimedOut(reason)) => {
             session.fail(reason.clone());
             Err(reason)

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(|_| ())

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

@@ -124,7 +124,7 @@ impl SessionInner {
                 .saturating_add(Duration::from_secs(1));
         let opening = async {
             let mut client = self.client();
-            let mut request = tonic::Request::new(request);
+            let mut request = self.request(request);
             if let Some(traceparent) = trace.and_then(|trace| trace.traceparent)
                 && let Ok(traceparent) = traceparent.parse()
             {
@@ -321,8 +321,13 @@ impl SessionInner {
                 grpc::yield_observation_request::Observation::WaitId(wait_id),
             )
         });
-        let response =
-            deadline::session_request(self, "wait", runtime_timeout, client.wait(request)).await;
+        let response = deadline::session_request(
+            self,
+            "wait",
+            runtime_timeout,
+            client.wait(self.request(request)),
+        )
+        .await;
         drop(watcher);
         cancellation.disarm();
         self.prune_wait_slots();
@@ -336,10 +341,10 @@ impl SessionInner {
         observation: grpc::yield_observation_request::Observation,
     ) -> AbortOnDropHandle<()> {
         let mut client = self.client();
-        let request = grpc::YieldObservationRequest {
+        let request = self.request(grpc::YieldObservationRequest {
             session_id: self.id.clone(),
             observation: Some(observation),
-        };
+        });
         AbortOnDropHandle::new(tokio::spawn(async move {
             yield_signal.cancelled().await;
             if let Err(error) = deadline::startup("yield observation", async move {
@@ -360,10 +365,10 @@ impl SessionInner {
             self,
             "termination",
             Duration::ZERO,
-            client.terminate(grpc::TerminateRequest {
+            client.terminate(self.request(grpc::TerminateRequest {
                 session_id: self.id.clone(),
                 cell_id: cell_id.as_str().to_string(),
-            }),
+            })),
         )
         .await?
         .into_inner();
@@ -467,10 +472,10 @@ impl Drop for WaitCancellation {
                 &session,
                 "wait cancellation",
                 Duration::ZERO,
-                client.cancel_wait(grpc::CancelWaitRequest {
+                client.cancel_wait(session.request(grpc::CancelWaitRequest {
                     session_id: session.id.clone(),
                     wait_id,
-                }),
+                })),
             )
             .await;
             if let Err(error) = result