mirror of
https://github.com/openai/codex.git
synced 2026-11-11 21:10:39 +00:00
Preserve typed gRPC code-mode session admission errors (#60496)
## Why Session-opening failures were returned as strings, preventing callers from inspecting the terminal gRPC admission status directly. ## What changed - Add `GrpcCodeModeSessionProvider::open_session` returning `GrpcSessionOpenError`, distinguishing admission status codes, admission timeouts, and other setup failures. - Keep string errors for the existing `CodeModeSessionProvider` trait API. - Re-export `ProcessOwnedGrpcCodeModeSessionProvider` from `codex-core`. ## Testing Extend admission tests to assert terminal status codes, cover exhausted `ResourceExhausted` retries, and verify that `DeadlineExceeded` is not retried. GitOrigin-RevId: 9d7f2e733e12b8bb7def368fa5fc5a09ea2a374e
This commit is contained in:
@@ -1,4 +1,5 @@
|
||||
use std::fmt::Display;
|
||||
use std::fmt::Formatter;
|
||||
use std::future::Future;
|
||||
use std::time::Duration;
|
||||
|
||||
@@ -7,6 +8,35 @@ use super::SessionInner;
|
||||
const TRANSPORT_TIMEOUT: Duration = Duration::from_secs(60);
|
||||
const MAX_ERROR_BYTES: usize = 512;
|
||||
|
||||
/// Terminal failure while establishing a new gRPC code-mode session.
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
pub enum GrpcSessionOpenError {
|
||||
/// Final status from OpenSession or the first read of its lease, with the bounded diagnostic.
|
||||
AdmissionStatus { code: tonic::Code, message: String },
|
||||
/// The client deadline elapsed while waiting for admission, including retries.
|
||||
AdmissionTimeout(String),
|
||||
/// Setup failed outside admission, for example during transport setup or tool subscription.
|
||||
Other(String),
|
||||
}
|
||||
|
||||
impl Display for GrpcSessionOpenError {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::AdmissionStatus { message, .. }
|
||||
| Self::AdmissionTimeout(message)
|
||||
| Self::Other(message) => f.write_str(message),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for GrpcSessionOpenError {}
|
||||
|
||||
impl<T: Into<String>> From<T> for GrpcSessionOpenError {
|
||||
fn from(message: T) -> Self {
|
||||
Self::Other(message.into())
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn startup<T, E: Display>(
|
||||
operation: &str,
|
||||
request: impl Future<Output = Result<T, E>>,
|
||||
@@ -85,7 +115,7 @@ pub(super) fn failure(operation: &str, error: impl Display) -> String {
|
||||
message
|
||||
}
|
||||
|
||||
async fn enforce<T, E>(
|
||||
pub(super) async fn enforce<T, E>(
|
||||
operation: &str,
|
||||
runtime_timeout: Duration,
|
||||
request: impl Future<Output = Result<T, E>>,
|
||||
@@ -100,7 +130,7 @@ async fn enforce<T, E>(
|
||||
}
|
||||
}
|
||||
|
||||
enum RequestError<E> {
|
||||
pub(super) enum RequestError<E> {
|
||||
Failed(E),
|
||||
TimedOut(String),
|
||||
}
|
||||
|
||||
@@ -27,6 +27,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 futures::TryFutureExt;
|
||||
use rand::Rng;
|
||||
use tokio::sync::Semaphore;
|
||||
use tokio::sync::watch;
|
||||
@@ -35,6 +36,8 @@ use tokio_util::task::TaskTracker;
|
||||
use tonic::transport::Channel;
|
||||
use tracing::Instrument;
|
||||
|
||||
pub use self::deadline::GrpcSessionOpenError;
|
||||
use self::deadline::RequestError;
|
||||
use self::operations::WaitSlot;
|
||||
use self::state::ClosedCell;
|
||||
use self::state::SessionState;
|
||||
@@ -134,12 +137,22 @@ impl GrpcCodeModeSessionProvider {
|
||||
}
|
||||
}
|
||||
|
||||
/// Opens a new session and preserves the terminal admission status for the caller.
|
||||
pub async fn open_session(
|
||||
&self,
|
||||
options: CodeModeSessionOptions,
|
||||
) -> Result<Arc<dyn CodeModeSession>, GrpcSessionOpenError> {
|
||||
let session = Arc::new(reconnect::ReconnectableSession::new(self.clone(), options));
|
||||
session.initialize().await?;
|
||||
Ok(session)
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "code_mode.grpc.open_binding", level = "info", skip_all)]
|
||||
async fn open_binding(
|
||||
&self,
|
||||
options: &CodeModeSessionOptions,
|
||||
pinned_mode: Option<CodeModeOutputTruncation>,
|
||||
) -> Result<Arc<GrpcCodeModeSession>, String> {
|
||||
) -> Result<Arc<GrpcCodeModeSession>, GrpcSessionOpenError> {
|
||||
let client = deadline::startup("transport connection", self.transport.client()).await?;
|
||||
let requested_mode = pinned_mode.unwrap_or(options.output_truncation);
|
||||
let limits = &options.cell_execution_limits;
|
||||
@@ -155,7 +168,8 @@ 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, route) = deadline::startup("session opening", async {
|
||||
let operation = "session opening";
|
||||
let (lease, first, route) = deadline::enforce(operation, Duration::ZERO, async {
|
||||
let mut retries = 0;
|
||||
loop {
|
||||
let mut request = tonic::Request::new(grpc::OpenSessionRequest {
|
||||
@@ -211,9 +225,16 @@ impl GrpcCodeModeSessionProvider {
|
||||
}
|
||||
})
|
||||
.instrument(open_session_span.clone())
|
||||
.await?;
|
||||
.await
|
||||
.map_err(|error| match error {
|
||||
RequestError::Failed(status) => GrpcSessionOpenError::AdmissionStatus {
|
||||
code: status.code(),
|
||||
message: deadline::failure(operation, status),
|
||||
},
|
||||
RequestError::TimedOut(reason) => GrpcSessionOpenError::AdmissionTimeout(reason),
|
||||
})?;
|
||||
let Some(grpc::session_event::Event::Opened(opened)) = first.event else {
|
||||
return Err("gRPC code-mode session lease omitted its opening event".to_string());
|
||||
return Err("gRPC code-mode session lease omitted its opening event".into());
|
||||
};
|
||||
validate_identifier(&opened.session_id, "session ID")?;
|
||||
let host_mode =
|
||||
@@ -230,13 +251,13 @@ impl GrpcCodeModeSessionProvider {
|
||||
{
|
||||
return Err(format!(
|
||||
"gRPC code-mode host selected {output_truncation:?} output truncation, requested {requested_mode:?}"
|
||||
));
|
||||
).into());
|
||||
}
|
||||
if pinned_mode.is_some_and(|previous| previous != output_truncation) {
|
||||
// Dropping the opening lease releases this incompatible session before subscribing.
|
||||
return Err(
|
||||
"gRPC code-mode host changed the session output truncation mode on reconnect"
|
||||
.to_string(),
|
||||
.into(),
|
||||
);
|
||||
}
|
||||
|
||||
@@ -277,7 +298,7 @@ impl GrpcCodeModeSessionProvider {
|
||||
Ok(response) => response,
|
||||
Err(error) => {
|
||||
let _ = wait_for_watch(inner.request_shutdown()).await;
|
||||
return Err(error);
|
||||
return Err(error.into());
|
||||
}
|
||||
};
|
||||
inner.spawn_tool_subscription(response.into_inner());
|
||||
@@ -313,11 +334,10 @@ impl CodeModeSessionProvider for GrpcCodeModeSessionProvider {
|
||||
&'a self,
|
||||
options: CodeModeSessionOptions,
|
||||
) -> CodeModeSessionProviderFuture<'a> {
|
||||
Box::pin(async move {
|
||||
let session = Arc::new(reconnect::ReconnectableSession::new(self.clone(), options));
|
||||
session.initialize().await?;
|
||||
Ok(session as _)
|
||||
})
|
||||
Box::pin(
|
||||
self.open_session(options)
|
||||
.map_err(|error| error.to_string()),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -20,6 +20,7 @@ use tokio_util::sync::CancellationToken;
|
||||
|
||||
use super::GrpcCodeModeSession;
|
||||
use super::GrpcCodeModeSessionProvider;
|
||||
use super::GrpcSessionOpenError;
|
||||
use super::SHUTDOWN_ERROR;
|
||||
use super::generation;
|
||||
use super::generation::GenerationDelegate;
|
||||
@@ -66,7 +67,7 @@ impl ReconnectableSession {
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn initialize(&self) -> Result<(), String> {
|
||||
pub(super) async fn initialize(&self) -> Result<(), GrpcSessionOpenError> {
|
||||
self.inner.get_or_open_binding().await.map(|_| ())
|
||||
}
|
||||
}
|
||||
@@ -151,7 +152,10 @@ impl Drop for ReconnectableSession {
|
||||
|
||||
impl ReconnectInner {
|
||||
async fn binding_for_request(&self, generation: Option<u64>) -> Result<SessionBinding, String> {
|
||||
let binding = self.get_or_open_binding().await?;
|
||||
let binding = self
|
||||
.get_or_open_binding()
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
if generation != Some(binding.generation) {
|
||||
return Err(concat!(
|
||||
"The exec runtime was replaced. Stored values and running cells were lost. ",
|
||||
@@ -162,9 +166,9 @@ impl ReconnectInner {
|
||||
Ok(binding)
|
||||
}
|
||||
|
||||
async fn get_or_open_binding(&self) -> Result<SessionBinding, String> {
|
||||
async fn get_or_open_binding(&self) -> Result<SessionBinding, GrpcSessionOpenError> {
|
||||
if self.shutdown_requested.is_cancelled() {
|
||||
return Err(SHUTDOWN_ERROR.to_string());
|
||||
return Err(SHUTDOWN_ERROR.into());
|
||||
}
|
||||
if let Some(binding) = self.live_binding() {
|
||||
return Ok(binding);
|
||||
@@ -173,13 +177,13 @@ impl ReconnectInner {
|
||||
let _opening_permit = tokio::select! {
|
||||
biased;
|
||||
_ = self.shutdown_requested.cancelled() => {
|
||||
return Err(SHUTDOWN_ERROR.to_string());
|
||||
return Err(SHUTDOWN_ERROR.into());
|
||||
}
|
||||
permit = self.opening_permit.acquire() => permit
|
||||
.map_err(|_| "gRPC code-mode session opening coordinator closed".to_string())?,
|
||||
};
|
||||
if self.shutdown_requested.is_cancelled() {
|
||||
return Err(SHUTDOWN_ERROR.to_string());
|
||||
return Err(SHUTDOWN_ERROR.into());
|
||||
}
|
||||
if let Some(binding) = self.live_binding() {
|
||||
return Ok(binding);
|
||||
@@ -196,7 +200,7 @@ impl ReconnectInner {
|
||||
let session = tokio::select! {
|
||||
biased;
|
||||
_ = self.shutdown_requested.cancelled() => {
|
||||
return Err(SHUTDOWN_ERROR.to_string());
|
||||
return Err(SHUTDOWN_ERROR.into());
|
||||
}
|
||||
session = self.provider.open_binding(&self.options, self.output_truncation.get().copied()) => session?,
|
||||
};
|
||||
@@ -217,7 +221,7 @@ impl ReconnectInner {
|
||||
};
|
||||
if !published {
|
||||
let _ = wait_for_watch(binding.session.inner.request_shutdown()).await;
|
||||
return Err(SHUTDOWN_ERROR.to_string());
|
||||
return Err(SHUTDOWN_ERROR.into());
|
||||
}
|
||||
Ok(binding)
|
||||
}
|
||||
|
||||
@@ -3,6 +3,7 @@ mod remote_session;
|
||||
|
||||
pub use codex_code_mode_protocol::*;
|
||||
pub use grpc_session::GrpcCodeModeSessionProvider;
|
||||
pub use grpc_session::GrpcSessionOpenError;
|
||||
pub use remote_session::DisabledCodeModeSessionProvider;
|
||||
pub use remote_session::ProcessOwnedCodeModeSession;
|
||||
pub use remote_session::ProcessOwnedCodeModeSessionProvider;
|
||||
|
||||
Reference in New Issue
Block a user