From f6a69b587344c187d4c12de947ee258b205f142c Mon Sep 17 00:00:00 2001 From: Alex Kotliarskyi Date: Fri, 6 Nov 2026 01:10:02 +0000 Subject: [PATCH] Tie gRPC code-mode connection lifetime to live sessions (#59783) ## Why Some hosts assign a worker per connection. Idle connection eviction can prevent a quiet session from reaching the worker holding its lease, while retaining a connection after the last session closes unnecessarily keeps that worker allocated. ## What changed - Disable the HTTP connection pool idle timeout for gRPC code-mode connections. - Cache URL-based clients weakly and let live sessions own the shared client, releasing their ownership on closure or failure even while the provider remains alive. - Acquire clients for deferred RPCs only when needed and return the session error if the client has already been released. ## Testing Add regression tests that verify connection reuse and stored state after 120 seconds of inactivity, and connection closure after the last shared session fails while the provider and session handles remain alive. GitOrigin-RevId: bd1ff00c4003c03ed173c23dd2ed373ed354ab65 --- codex-rs/code-mode-host/Cargo.toml | 1 + .../tests/grpc/session_recovery_tests.rs | 116 ++++++++++++++++++ .../code-mode/src/grpc_session/callbacks.rs | 2 +- .../src/grpc_session/callbacks_tests.rs | 4 +- codex-rs/code-mode/src/grpc_session/mod.rs | 26 ++-- .../code-mode/src/grpc_session/operations.rs | 17 ++- .../src/grpc_session/process_metrics.rs | 4 +- codex-rs/code-mode/src/grpc_session/state.rs | 28 ++++- .../code-mode/src/grpc_session/transport.rs | 25 ++-- codex-rs/http-client/src/client_builder.rs | 12 ++ 10 files changed, 195 insertions(+), 40 deletions(-) diff --git a/codex-rs/code-mode-host/Cargo.toml b/codex-rs/code-mode-host/Cargo.toml index 98e45c5e0b..7668df1157 100644 --- a/codex-rs/code-mode-host/Cargo.toml +++ b/codex-rs/code-mode-host/Cargo.toml @@ -38,6 +38,7 @@ tracing-subscriber = { workspace = true } uuid = { workspace = true, features = ["v4"] } [dev-dependencies] +tokio = { workspace = true, features = ["test-util"] } codex-plugin-host = { workspace = true } codex-code-mode = { workspace = true } codex-utils-cargo-bin = { workspace = true } diff --git a/codex-rs/code-mode-host/tests/grpc/session_recovery_tests.rs b/codex-rs/code-mode-host/tests/grpc/session_recovery_tests.rs index 822ebab886..6d675f5c6b 100644 --- a/codex-rs/code-mode-host/tests/grpc/session_recovery_tests.rs +++ b/codex-rs/code-mode-host/tests/grpc/session_recovery_tests.rs @@ -2,6 +2,7 @@ use std::sync::Arc; use std::sync::Mutex; +use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; @@ -15,8 +16,10 @@ use codex_code_mode::NoopCodeModeSessionDelegate; use codex_code_mode::RuntimeResponse; use codex_code_mode_host::GrpcCodeModeHost; use codex_code_mode_protocol::grpc::code_mode_host_server::CodeModeHostServer; +use futures::StreamExt; use pretty_assertions::assert_eq; use tokio::net::TcpListener; +use tokio::net::TcpStream; use tokio_stream::wrappers::TcpListenerStream; use tokio_util::task::AbortOnDropHandle; use tonic::Code; @@ -27,6 +30,119 @@ use super::execute; use super::request; use super::text_response; +#[tokio::test] +async fn quiet_session_reuses_its_worker_connection() -> Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let endpoint = format!("http://{}", listener.local_addr()?); + let connections = Arc::new(AtomicUsize::new(0)); + let accepted = Arc::clone(&connections); + let incoming = TcpListenerStream::new(listener).inspect(move |_| { + accepted.fetch_add(1, Ordering::SeqCst); + }); + let _server = AbortOnDropHandle::new(tokio::spawn( + Server::builder() + .add_service(CodeModeHostServer::new(GrpcCodeModeHost::new())) + .serve_with_incoming(incoming), + )); + let session = GrpcCodeModeSessionProvider::new(endpoint) + .create_session() + .await + .map_err(anyhow::Error::msg)?; + execute( + &session, + request(r#"store("key", "persisted");"#), + Arc::new(NoopCodeModeSessionDelegate), + ) + .await?; + + tokio::time::pause(); + tokio::time::advance(std::time::Duration::from_secs(120)).await; + tokio::time::resume(); + + let result = execute( + &session, + request(r#"text(load("key"));"#), + Arc::new(NoopCodeModeSessionDelegate), + ) + .await?; + assert_eq!( + result, + text_response("2", "persisted", result.code_mode_host_duration()) + ); + assert_eq!(connections.load(Ordering::SeqCst), 1); + Ok(()) +} + +#[tokio::test] +async fn last_failed_session_releases_the_shared_connection() -> Result<()> { + let reject = Arc::new(AtomicBool::new(false)); + let reject_next_rpc = Arc::clone(&reject); + let service = CodeModeHostServer::with_interceptor(GrpcCodeModeHost::new(), move |req| { + if reject_next_rpc.swap(false, Ordering::SeqCst) { + Err(Status::not_found("session was lost")) + } else { + Ok(req) + } + }); + let host = TcpListener::bind("127.0.0.1:0").await?; + let host_addr = host.local_addr()?; + let _server = AbortOnDropHandle::new(tokio::spawn( + Server::builder() + .add_service(service) + .serve_with_incoming(TcpListenerStream::new(host)), + )); + + // Forward exactly one TCP connection. Its EOF is the worker-release signal. + let listener = TcpListener::bind("127.0.0.1:0").await?; + let provider = GrpcCodeModeSessionProvider::new(format!("http://{}", listener.local_addr()?)); + let connection = AbortOnDropHandle::new(tokio::spawn(async move { + let (mut client, _) = listener.accept().await?; + drop(listener); + let mut host = TcpStream::connect(host_addr).await?; + tokio::io::copy_bidirectional(&mut client, &mut host).await + })); + let first = provider + .create_session() + .await + .map_err(anyhow::Error::msg)?; + let last = provider + .create_session() + .await + .map_err(anyhow::Error::msg)?; + first.shutdown().await.map_err(anyhow::Error::msg)?; + let result = execute( + &last, + request("text('still live')"), + Arc::new(NoopCodeModeSessionDelegate), + ) + .await?; + assert_eq!( + result, + text_response("g2:1", "still live", result.code_mode_host_duration()) + ); + + reject.store(true, Ordering::SeqCst); + let error = execute( + &last, + request("text('must not run')"), + Arc::new(NoopCodeModeSessionDelegate), + ) + .await + .unwrap_err(); + assert_eq!( + error.to_string(), + format!( + "gRPC code-mode execution failed: {}", + Status::not_found("session was lost") + ) + ); + tokio::time::timeout(super::TEST_TIMEOUT, connection) + .await + .expect("last failed session kept its worker connection open")??; + drop((first, last, provider)); + Ok(()) +} + #[tokio::test] async fn only_missing_sessions_reopen_without_replaying_execution() -> Result<()> { for (code, route_lost, expected_cell, expected_value) in [ diff --git a/codex-rs/code-mode/src/grpc_session/callbacks.rs b/codex-rs/code-mode/src/grpc_session/callbacks.rs index 7800b5fc41..d423e4cdfc 100644 --- a/codex-rs/code-mode/src/grpc_session/callbacks.rs +++ b/codex-rs/code-mode/src/grpc_session/callbacks.rs @@ -184,7 +184,6 @@ impl SessionInner { result: Result, ) { let request = completion::request(&self.id, &invocation_id, result); - let mut client = self.client(); tokio::select! { biased; _ = cancellation.cancelled() => {} @@ -193,6 +192,7 @@ impl SessionInner { result = async { let _permit = self.completion_permits.acquire().await .map_err(|_| "code-mode completion queue is closed".to_string())?; + let mut client = self.client()?; deadline::request( self, "tool invocation completion", diff --git a/codex-rs/code-mode/src/grpc_session/callbacks_tests.rs b/codex-rs/code-mode/src/grpc_session/callbacks_tests.rs index a30c4ee724..1ebfc416fc 100644 --- a/codex-rs/code-mode/src/grpc_session/callbacks_tests.rs +++ b/codex-rs/code-mode/src/grpc_session/callbacks_tests.rs @@ -49,13 +49,13 @@ async fn missing_invocation_completion_racing_cancellation_preserves_session() { } fn session(client: GrpcClient) -> Arc { + let client = Arc::new(client); Arc::new(SessionInner { id: "session".to_string(), output_truncation: CodeModeOutputTruncation::Client, - client: client.clone(), route: None, runtime: tokio::runtime::Handle::current(), - state: Mutex::new(SessionState::default()), + state: Mutex::new(SessionState::with_client(Arc::clone(&client))), wait_slots: Mutex::new(HashMap::new()), shutdown_requested: AtomicBool::new(/*v*/ false), shutdown_result: Mutex::new(/*t*/ None), diff --git a/codex-rs/code-mode/src/grpc_session/mod.rs b/codex-rs/code-mode/src/grpc_session/mod.rs index 61f8b3324d..eec382d9a6 100644 --- a/codex-rs/code-mode/src/grpc_session/mod.rs +++ b/codex-rs/code-mode/src/grpc_session/mod.rs @@ -140,7 +140,7 @@ impl GrpcCodeModeSessionProvider { options: &CodeModeSessionOptions, pinned_mode: Option, ) -> Result, String> { - let mut client = deadline::startup("transport connection", self.transport.client()).await?; + 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; let limits = grpc::SessionCellExecutionLimits { @@ -172,7 +172,7 @@ impl GrpcCodeModeSessionProvider { }); inject_span_traceparent(&mut request, &open_session_span); let result = async { - let response = client.open_session(request).await?; + let response = client.as_ref().clone().open_session(request).await?; let route = response.metadata().get(ROUTE_HEADER).cloned(); if route .as_ref() @@ -243,10 +243,9 @@ impl GrpcCodeModeSessionProvider { let inner = Arc::new(SessionInner { id: opened.session_id, output_truncation, - client, route, runtime: tokio::runtime::Handle::current(), - state: Mutex::new(SessionState::default()), + state: Mutex::new(SessionState::with_client(client)), wait_slots: Mutex::new(HashMap::new()), shutdown_requested: AtomicBool::new(false), shutdown_result: Mutex::new(None), @@ -269,7 +268,7 @@ impl GrpcCodeModeSessionProvider { tool_names: Vec::new(), }); inject_span_traceparent(&mut request, &subscribe_span); - let mut client = inner.client(); + let mut client = inner.client()?; let response = match deadline::startup("tool subscription", client.subscribe_to_tool_calls(request)) .instrument(subscribe_span) @@ -386,7 +385,6 @@ impl Drop for GrpcCodeModeSession { pub(super) struct SessionInner { pub(super) id: String, output_truncation: CodeModeOutputTruncation, - 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>, @@ -402,8 +400,11 @@ pub(super) struct SessionInner { } impl SessionInner { - pub(super) fn client(&self) -> GrpcClient { - self.client.clone() + pub(super) fn client(&self) -> Result { + self.state + .lock() + .unwrap_or_else(PoisonError::into_inner) + .client() } pub(super) fn request(&self, message: T) -> tonic::Request { @@ -481,14 +482,7 @@ impl SessionInner { } async fn drive_shutdown(&self) -> Result<(), String> { - let is_open = self - .state - .lock() - .unwrap_or_else(PoisonError::into_inner) - .require_open() - .is_ok(); - let result = if is_open { - let mut client = self.client(); + let result = if let Ok(mut client) = self.client() { deadline::request( self, "session shutdown", diff --git a/codex-rs/code-mode/src/grpc_session/operations.rs b/codex-rs/code-mode/src/grpc_session/operations.rs index 877590198e..5b60d0434a 100644 --- a/codex-rs/code-mode/src/grpc_session/operations.rs +++ b/codex-rs/code-mode/src/grpc_session/operations.rs @@ -124,7 +124,7 @@ impl SessionInner { Duration::from_millis(request.yield_time_ms.unwrap_or(DEFAULT_EXEC_YIELD_TIME_MS)) .saturating_add(Duration::from_secs(1)); let opening = async { - let mut client = self.client(); + let mut client = self.client()?; let mut request = self.request(request); if let Some(traceparent) = trace.and_then(|trace| trace.traceparent) && let Ok(traceparent) = traceparent.parse() @@ -319,7 +319,7 @@ impl SessionInner { request.max_output_tokens, )?, }; - let mut client = self.client(); + let mut client = self.client()?; let watcher = preempt.map(|signal| { self.watch_yield( signal, @@ -341,17 +341,20 @@ impl SessionInner { } fn watch_yield( - &self, + self: &Arc, yield_signal: CancellationToken, observation: grpc::yield_observation_request::Observation, ) -> AbortOnDropHandle<()> { - let mut client = self.client(); + let session = Arc::clone(self); let request = self.request(grpc::YieldObservationRequest { session_id: self.id.clone(), observation: Some(observation), }); AbortOnDropHandle::new(tokio::spawn(async move { yield_signal.cancelled().await; + let Ok(mut client) = session.client() else { + return; + }; if let Err(error) = deadline::startup("yield observation", async move { client.yield_observation(request).await.map(|_| ()) }) @@ -365,7 +368,7 @@ impl SessionInner { pub(super) async fn terminate(&self, request: TerminateRequest) -> Result { self.require_open()?; - let mut client = self.client(); + let mut client = self.client()?; let response = deadline::session_request( self, "termination", @@ -475,7 +478,9 @@ impl Drop for WaitCancellation { } let session = Arc::clone(&self.session); self.session.runtime.spawn(async move { - let mut client = session.client(); + let Ok(mut client) = session.client() else { + return; + }; let result = deadline::request( &session, "wait cancellation", diff --git a/codex-rs/code-mode/src/grpc_session/process_metrics.rs b/codex-rs/code-mode/src/grpc_session/process_metrics.rs index 5c381520cc..2f269f291b 100644 --- a/codex-rs/code-mode/src/grpc_session/process_metrics.rs +++ b/codex-rs/code-mode/src/grpc_session/process_metrics.rs @@ -16,10 +16,10 @@ pub(super) async fn sample( ) -> Result { let deadline = Duration::from_secs(/*secs*/ 1); tokio::time::timeout(deadline, async { - let mut client = provider.transport.client().await?; + let client = provider.transport.client().await?; let mut request = tonic::Request::new(grpc::GetProcessMetricsRequest {}); request.set_timeout(deadline); - match client.get_process_metrics(request).await { + match client.as_ref().clone().get_process_metrics(request).await { Ok(response) => { let snapshot = response.into_inner(); if snapshot.process_id == 0 diff --git a/codex-rs/code-mode/src/grpc_session/state.rs b/codex-rs/code-mode/src/grpc_session/state.rs index 0fbb82ace0..8e2d13efdc 100644 --- a/codex-rs/code-mode/src/grpc_session/state.rs +++ b/codex-rs/code-mode/src/grpc_session/state.rs @@ -86,6 +86,8 @@ impl RecentIds { #[derive(Default)] pub(super) struct SessionState { + // Only a live session owns a share of the worker connection. + client: Option>, executions: HashMap, invocations: HashMap, notifications: usize, @@ -96,12 +98,29 @@ pub(super) struct SessionState { } impl SessionState { + pub(super) fn with_client(client: Arc) -> Self { + Self { + client: Some(client), + ..Self::default() + } + } + + pub(super) fn client(&self) -> Result { + self.client + .as_deref() + .cloned() + .ok_or_else(|| self.closed_error()) + } + + fn closed_error(&self) -> String { + self.failure + .clone() + .unwrap_or_else(|| "code-mode gRPC session is closed".to_string()) + } + pub(super) fn require_open(&self) -> Result<(), String> { if self.closed { - return Err(self - .failure - .clone() - .unwrap_or_else(|| "code-mode gRPC session is closed".to_string())); + return Err(self.closed_error()); } Ok(()) } @@ -341,6 +360,7 @@ impl SessionState { return Vec::new(); } self.closed = true; + self.client = None; self.failure = failure; self.notifications = 0; for (_, callback) in self.invocations.drain() { diff --git a/codex-rs/code-mode/src/grpc_session/transport.rs b/codex-rs/code-mode/src/grpc_session/transport.rs index 286a122d5f..e34cfdb084 100644 --- a/codex-rs/code-mode/src/grpc_session/transport.rs +++ b/codex-rs/code-mode/src/grpc_session/transport.rs @@ -1,4 +1,5 @@ use std::sync::Arc; +use std::sync::Weak; use codex_code_mode_protocol::grpc::code_mode_host_client::CodeModeHostClient; use codex_code_mode_protocol::host::MAX_FRAME_BYTES; @@ -24,9 +25,9 @@ pub(super) enum SharedTransport { Url { endpoint: String, http_client_factory: HttpClientFactory, - client: tokio::sync::Mutex>, + client: tokio::sync::Mutex)>>, }, - Connected(GrpcClient), + Connected(Arc), } impl SharedTransport { @@ -40,24 +41,24 @@ impl SharedTransport { pub(super) fn with_channel(channel: Channel) -> Self { let transport = channel.map_err(|error| Status::from_error(Box::new(error))); - Self::Connected( + Self::Connected(Arc::new( CodeModeHostClient::new(BoxCloneSyncService::new(transport)) .accept_compressed(tonic::codec::CompressionEncoding::Zstd) .max_decoding_message_size(MAX_FRAME_BYTES) .max_encoding_message_size(MAX_FRAME_BYTES), - ) + )) } pub(super) fn with_transport(transport: GrpcTransport) -> Self { - Self::Connected( + Self::Connected(Arc::new( CodeModeHostClient::new(transport) .accept_compressed(tonic::codec::CompressionEncoding::Zstd) .max_decoding_message_size(MAX_FRAME_BYTES) .max_encoding_message_size(MAX_FRAME_BYTES), - ) + )) } - pub(super) async fn client(&self) -> Result { + pub(super) async fn client(&self) -> Result, String> { let (endpoint, http_client_factory, client) = match self { Self::Url { endpoint, @@ -75,8 +76,9 @@ impl SharedTransport { let mut cached = client.lock().await; if let Some((factory, client)) = cached.as_ref() && factory == &bound_factory + && let Some(client) = client.upgrade() { - return Ok(client.clone()); + return Ok(client); } let client = if endpoint.starts_with("unix:") { let channel = Endpoint::from_shared(endpoint.clone()) @@ -107,6 +109,9 @@ impl SharedTransport { ClientRouteClass::Other, codex_http_client::HttpClientBuilder::new() .without_request_logging() + // Some hosts assign a worker per connection. A quiet session + // must still reach the worker holding its open lease. + .without_pool_idle_timeout() .http2_prior_knowledge() .without_redirects(), ); @@ -145,7 +150,9 @@ impl SharedTransport { .accept_compressed(tonic::codec::CompressionEncoding::Zstd) .max_decoding_message_size(MAX_FRAME_BYTES) .max_encoding_message_size(MAX_FRAME_BYTES); - *cached = Some((bound_factory, client.clone())); + // Let the last session release its worker even if the provider stays alive. + let client = Arc::new(client); + *cached = Some((bound_factory, Arc::downgrade(&client))); Ok(client) } } diff --git a/codex-rs/http-client/src/client_builder.rs b/codex-rs/http-client/src/client_builder.rs index 7de37b90ff..cdcc0bb58f 100644 --- a/codex-rs/http-client/src/client_builder.rs +++ b/codex-rs/http-client/src/client_builder.rs @@ -35,6 +35,7 @@ use crate::with_chatgpt_cloudflare_cookie_store; #[derive(Clone)] pub struct HttpClientBuilder { http2_prior_knowledge: bool, + disable_pool_idle_timeout: bool, pub(crate) default_headers: Option, follow_redirects: bool, pub(crate) redirect_observed: Option>, @@ -108,6 +109,13 @@ impl HttpClientBuilder { self.http2_prior_knowledge = true; self } + + /// Keep pooled connections reusable during long gaps between requests. + pub fn without_pool_idle_timeout(mut self) -> Self { + self.disable_pool_idle_timeout = true; + self + } + pub fn new() -> Self { Self::default() } @@ -386,6 +394,9 @@ impl HttpClientBuilder { if let Some(path) = self.windows_proxy_pipe_path { builder = builder.windows_named_pipe(path); } + if self.disable_pool_idle_timeout { + builder = builder.pool_idle_timeout(None); + } if self.http2_prior_knowledge { builder = builder.http2_prior_knowledge(); } @@ -426,6 +437,7 @@ impl Default for HttpClientBuilder { fn default() -> Self { Self { http2_prior_knowledge: false, + disable_pool_idle_timeout: false, default_headers: None, follow_redirects: true, redirect_observed: None,