openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-20 17:38:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
## 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
This commit is contained in:
cassirer-openai
2026-10-09 02:50:59 +00:00
committed by copyberry
parent 03b761dca9
commit 5b3c1522f8
9 changed files with 358 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/code-mode/src/grpc_session/callbacks.rs
+3 -8
View File
@@ -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
+65 -13
View File
@@ -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
+15 -1
View File
@@ -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
+51 -15
View File
@@ -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
+14 -9
View File
@@ -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