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

Tie gRPC code-mode connection lifetime to live sessions (#59783)

Browse Source 
## 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
This commit is contained in:
Alex Kotliarskyi
2026-11-06 01:49:33 +00:00
committed by copyberry
parent 1840e7dead
commit f6a69b5873
10 changed files with 195 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/code-mode/src/grpc_session/callbacks.rs
+1 -1
View File
@@ -184,7 +184,6 @@ impl SessionInner {
result: Result<serde_json::Value, String>,
) {
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",
codex-rs/code-mode/src/grpc_session/callbacks_tests.rs
+2 -2
View File
@@ -49,13 +49,13 @@ async fn missing_invocation_completion_racing_cancellation_preserves_session() {
}
fn session(client: GrpcClient) -> Arc<SessionInner> {
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),
codex-rs/code-mode/src/grpc_session/mod.rs
+10 -16
View File
@@ -140,7 +140,7 @@ impl GrpcCodeModeSessionProvider {
options: &CodeModeSessionOptions,
pinned_mode: Option<CodeModeOutputTruncation>,
) -> Result<Arc<GrpcCodeModeSession>, 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<tonic::metadata::MetadataValue<tonic::metadata::Ascii>>,
@@ -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<GrpcClient, String> {
self.state
.lock()
.unwrap_or_else(PoisonError::into_inner)
.client()
}
pub(super) fn request<T>(&self, message: T) -> tonic::Request<T> {
@@ -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",
codex-rs/code-mode/src/grpc_session/operations.rs
+11 -6
View File
@@ -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<Self>,
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<WaitOutcome, String> {
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",
codex-rs/code-mode/src/grpc_session/process_metrics.rs
+2 -2
View File
@@ -16,10 +16,10 @@ pub(super) async fn sample(
) -> Result<CodeModeProcessMetricsSample, String> {
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
codex-rs/code-mode/src/grpc_session/state.rs
+24 -4
View File
@@ -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<Arc<super::GrpcClient>>,
executions: HashMap<String, ExecutionRecord>,
invocations: HashMap<String, ActiveCallback>,
notifications: usize,
@@ -96,12 +98,29 @@ pub(super) struct SessionState {
}
impl SessionState {
pub(super) fn with_client(client: Arc<super::GrpcClient>) -> Self {
Self {
client: Some(client),
..Self::default()
}
}
pub(super) fn client(&self) -> Result<super::GrpcClient, String> {
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() {
codex-rs/code-mode/src/grpc_session/transport.rs
+16 -9
View File
@@ -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<Option<(HttpClientFactory, GrpcClient)>>,
client: tokio::sync::Mutex<Option<(HttpClientFactory, Weak<GrpcClient>)>>,
},
Connected(GrpcClient),
Connected(Arc<GrpcClient>),
}
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<GrpcClient, String> {
pub(super) async fn client(&self) -> Result<Arc<GrpcClient>, 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)
}
}