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,