openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-07 20:30:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Retry transient gRPC code-mode session admission failures (#51185)

Browse Source 
## Why

Transient admission errors or connection resets can prevent a code-mode session from opening before any execution starts.

## What changed

- Retry `OpenSession` and the initial lease message on `Unavailable` or `ResourceExhausted`, with at most three attempts and 100 ms and 200 ms delays within a single startup deadline.
- Classify HTTP connection and request failures as `Unavailable` so admission can recover from connection resets, while preserving error sources.
- Keep retries limited to initial admission so established executions are never replayed.

## Testing

Add gRPC integration tests covering transient status recovery, the retry limit, immediate failure for non-retryable statuses, and successful session execution after a connection reset.

GitOrigin-RevId: 3b654bcd377a5addf7c3ed58bb1c72a449b2f7f7
This commit is contained in:
cassirer-openai
2026-10-05 22:18:55 +00:00
committed by copyberry
parent 7f21b5ee9c
commit aa6635ece5
4 changed files with 165 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/code-mode/src/grpc_session/mod.rs
+32 -15
View File
@@ -112,21 +112,38 @@ 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 mut open_session_request = tonic::Request::new(grpc::OpenSessionRequest {
cell_execution_limits,
});
inject_span_traceparent(&mut open_session_request, &open_session_span);
let (lease, first) = async {
let mut lease =
deadline::startup("session opening", client.open_session(open_session_request))
.await?
.into_inner();
let first = deadline::startup("session lease opening", lease.message())
.await?
.ok_or_else(|| "gRPC code-mode session lease ended before opening".to_string())?;
Ok::<_, String>((lease, first))
}
.instrument(open_session_span)
let (lease, first) = deadline::startup("session opening", async {
let mut retries = 0;
loop {
let mut request = tonic::Request::new(grpc::OpenSessionRequest {
cell_execution_limits,
});
inject_span_traceparent(&mut request, &open_session_span);
let result = async {
let mut lease = client.open_session(request).await?.into_inner();
let first = lease.message().await?.ok_or_else(|| {
tonic::Status::unavailable("session lease ended before opening")
})?;
Ok::<_, tonic::Status>((lease, first))
}
.await;
let Err(error) = &result else { return result };
if retries == 2
|| !matches!(
error.code(),
tonic::Code::Unavailable | tonic::Code::ResourceExhausted
)
{
return result;
}
// No execution has started, and dropping a failed lease releases
// any partial admission. Never apply these retries to Execute.
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())
.await?;
let Some(grpc::session_event::Event::Opened(opened)) = first.event else {
return Err("gRPC code-mode session lease omitted its opening event".to_string());
codex-rs/code-mode/src/grpc_session/transport.rs
+23 -7
View File
@@ -1,10 +1,11 @@
use std::io;
use std::sync::Arc;
use codex_code_mode_protocol::grpc::code_mode_host_client::CodeModeHostClient;
use codex_code_mode_protocol::host::MAX_FRAME_BYTES;
use codex_http_client::ClientRouteClass;
use codex_http_client::HttpClientFactory;
use http_body_util::BodyExt;
use tonic::Status;
use tonic::body::Body;
use tonic::codegen::http::Request;
use tonic::codegen::http::Response;
@@ -17,7 +18,7 @@ use tower::util::BoxCloneSyncService;
use super::GrpcClient;
pub(super) type GrpcTransport = BoxCloneSyncService<Request<Body>, Response<Body>, io::Error>;
pub(super) type GrpcTransport = BoxCloneSyncService<Request<Body>, Response<Body>, Status>;
pub(super) enum SharedTransport {
Url {
@@ -38,7 +39,7 @@ impl SharedTransport {
}
pub(super) fn with_channel(channel: Channel) -> Self {
let transport = channel.map_err(io::Error::other);
let transport = channel.map_err(|error| Status::from_error(Box::new(error)));
Self::Connected(
CodeModeHostClient::new(BoxCloneSyncService::new(transport))
.max_decoding_message_size(MAX_FRAME_BYTES)
@@ -71,7 +72,7 @@ impl SharedTransport {
let channel = Endpoint::from_shared(endpoint.clone())
.map_err(|error| format!("invalid gRPC code-mode Unix socket endpoint: {error}"))?
.connect_lazy();
let transport = channel.map_err(io::Error::other);
let transport = channel.map_err(|error| Status::from_error(Box::new(error)));
CodeModeHostClient::new(BoxCloneSyncService::new(transport))
} else {
let target = reqwest::Url::parse(endpoint)
@@ -104,13 +105,28 @@ impl SharedTransport {
async move {
let request =
request.map(|body| reqwest::Body::wrap_stream(body.into_data_stream()));
let request = reqwest::Request::try_from(request).map_err(io::Error::other)?;
let request = reqwest::Request::try_from(request)
.map_err(|error| Status::invalid_argument(error.to_string()))?;
let response = client
.execute(request)
.await
.map_err(io::Error::other)?
.map_err(|error| {
// Sending can fail on DNS or a reset connection before
// headers arrive. Keep these distinct from RPC errors.
let mut status = if matches!(
&error,
codex_http_client::HttpError::Request(error)
if error.is_connect() || error.is_request()
) {
Status::unavailable(error.to_string())
} else {
Status::unknown(error.to_string())
};
status.set_source(Arc::new(error));
status
})?
.into_http_response();
Ok::<_, io::Error>(response.map(Body::new))
Ok::<_, Status>(response.map(Body::new))
}
});
CodeModeHostClient::with_origin(BoxCloneSyncService::new(transport), origin)