From 85f331772f54a9d518fefa6b1596cbde8cd78fd1 Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Tue, 11 Aug 2026 23:32:58 +0000 Subject: [PATCH] Route gRPC code-mode sessions through the shared HTTP client (#38087) ## What changed - Build URL-based gRPC code-mode connections with `HttpClientFactory` so they support the application's outbound proxy and custom CA configuration. - Accept `http` and `https` origins while rejecting endpoints with unsupported schemes, paths, queries, or fragments. - Preserve custom tonic channel injection and gRPC frame-size limits through the new transport adapter. GitOrigin-RevId: 142f0b572b3ab0154e8fe860cb304752a9af5784 --- codex-rs/Cargo.lock | 3 + codex-rs/code-mode/Cargo.toml | 3 + codex-rs/code-mode/src/grpc_session/mod.rs | 30 +++-- .../code-mode/src/grpc_session/transport.rs | 110 +++++++++++++++--- codex-rs/deny.toml | 2 + 5 files changed, 119 insertions(+), 29 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 0e91214228..a67d1f3d7b 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2563,13 +2563,16 @@ dependencies = [ "codex-protocol", "codex-websocket-client", "futures", + "http-body-util", "pretty_assertions", "prost", + "reqwest 0.12.28", "serde_json", "tokio", "tokio-tungstenite", "tokio-util", "tonic", + "tower", "tracing", "uuid", ] diff --git a/codex-rs/code-mode/Cargo.toml b/codex-rs/code-mode/Cargo.toml index e7d28dc788..0ef5079beb 100644 --- a/codex-rs/code-mode/Cargo.toml +++ b/codex-rs/code-mode/Cargo.toml @@ -19,12 +19,15 @@ codex-install-context = { workspace = true } codex-protocol = { workspace = true } codex-websocket-client = { workspace = true } futures = { workspace = true } +http-body-util = "0.1.3" prost = "0.14.3" +reqwest = { workspace = true, features = ["stream"] } serde_json = { workspace = true } tokio = { workspace = true, features = ["io-util", "macros", "net", "process", "rt", "sync", "time"] } tokio-tungstenite = { workspace = true } tokio-util = { workspace = true, features = ["rt"] } tonic = { workspace = true } +tower = { version = "0.5.3", features = ["util"] } tracing = { workspace = true } uuid = { workspace = true, features = ["v4"] } diff --git a/codex-rs/code-mode/src/grpc_session/mod.rs b/codex-rs/code-mode/src/grpc_session/mod.rs index 76f9e9452f..43a325f5f0 100644 --- a/codex-rs/code-mode/src/grpc_session/mod.rs +++ b/codex-rs/code-mode/src/grpc_session/mod.rs @@ -21,7 +21,8 @@ use codex_code_mode_protocol::WaitOutcome; use codex_code_mode_protocol::WaitRequest; use codex_code_mode_protocol::grpc; use codex_code_mode_protocol::grpc::code_mode_host_client::CodeModeHostClient; -use codex_code_mode_protocol::host::MAX_FRAME_BYTES; +use codex_http_client::HttpClientFactory; +use codex_http_client::OutboundProxyPolicy; use tokio::sync::watch; use tokio_util::sync::CancellationToken; use tokio_util::task::TaskTracker; @@ -29,6 +30,7 @@ use tonic::transport::Channel; use self::operations::WaitSlot; use self::state::SessionState; +use self::transport::GrpcTransport; use self::transport::SharedTransport; use crate::remote_session::ShutdownResultReceiver; use crate::remote_session::wait_for_watch; @@ -41,7 +43,7 @@ mod operations; mod state; mod transport; -type GrpcClient = CodeModeHostClient; +type GrpcClient = CodeModeHostClient; /// Creates code-mode sessions over an HTTP/2 gRPC connection. #[derive(Clone)] @@ -50,9 +52,20 @@ pub struct GrpcCodeModeSessionProvider { } impl GrpcCodeModeSessionProvider { - /// Connects lazily to an `http://` gRPC endpoint. + /// Connects lazily to an `http://` or `https://` gRPC endpoint. pub fn new(endpoint: impl Into) -> Self { - Self::from_transport(SharedTransport::new(endpoint.into())) + Self::with_http_client_factory( + endpoint, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ) + } + + /// Connects using the application's resolved outbound proxy and custom CA policy. + pub fn with_http_client_factory( + endpoint: impl Into, + http_client_factory: HttpClientFactory, + ) -> Self { + Self::from_transport(SharedTransport::new(endpoint.into(), http_client_factory)) } /// Uses an existing channel, including channels backed by custom transports. @@ -71,8 +84,7 @@ impl GrpcCodeModeSessionProvider { delegate: Arc, limits: CodeModeSessionCellExecutionLimits, ) -> Result, String> { - let channel = deadline::startup("transport connection", self.transport.channel()).await?; - let mut client = grpc_client(channel); + let mut client = deadline::startup("transport connection", self.transport.client()).await?; let limits = grpc::SessionCellExecutionLimits { max_yield_time_ms: limits.max_yield_time_ms, max_heap_size_bytes: limits @@ -327,9 +339,3 @@ fn validate_identifier(value: &str, field: &str) -> Result<(), String> { } Ok(()) } - -fn grpc_client(channel: Channel) -> GrpcClient { - CodeModeHostClient::new(channel) - .max_decoding_message_size(MAX_FRAME_BYTES) - .max_encoding_message_size(MAX_FRAME_BYTES) -} diff --git a/codex-rs/code-mode/src/grpc_session/transport.rs b/codex-rs/code-mode/src/grpc_session/transport.rs index 23ef21408b..ce98f8e188 100644 --- a/codex-rs/code-mode/src/grpc_session/transport.rs +++ b/codex-rs/code-mode/src/grpc_session/transport.rs @@ -1,44 +1,120 @@ +use std::io; + +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::body::Body; +use tonic::codegen::http::Request; +use tonic::codegen::http::Response; +use tonic::codegen::http::Uri; use tonic::transport::Channel; -use tonic::transport::Endpoint; +use tower::ServiceExt; +use tower::service_fn; +use tower::util::BoxCloneSyncService; + +use super::GrpcClient; + +pub(super) type GrpcTransport = BoxCloneSyncService, Response, io::Error>; pub(super) struct SharedTransport { endpoint: TransportEndpoint, - channel: tokio::sync::OnceCell, + client: tokio::sync::OnceCell, } enum TransportEndpoint { - Url(String), + Url { + endpoint: String, + http_client_factory: HttpClientFactory, + }, Connected(Channel), } impl SharedTransport { - pub(super) fn new(endpoint: String) -> Self { + pub(super) fn new(endpoint: String, http_client_factory: HttpClientFactory) -> Self { Self { - endpoint: TransportEndpoint::Url(endpoint), - channel: tokio::sync::OnceCell::new(), + endpoint: TransportEndpoint::Url { + endpoint, + http_client_factory, + }, + client: tokio::sync::OnceCell::new(), } } pub(super) fn with_channel(channel: Channel) -> Self { Self { endpoint: TransportEndpoint::Connected(channel), - channel: tokio::sync::OnceCell::new(), + client: tokio::sync::OnceCell::new(), } } - pub(super) async fn channel(&self) -> Result { - self.channel + pub(super) async fn client(&self) -> Result { + self.client .get_or_try_init(|| async { - match &self.endpoint { - TransportEndpoint::Url(endpoint) => Endpoint::from_shared(endpoint.clone()) - .map_err(|error| format!("invalid gRPC code-mode host endpoint: {error}"))? - .connect() + let client = match &self.endpoint { + TransportEndpoint::Url { + endpoint, + http_client_factory, + } => { + let target = reqwest::Url::parse(endpoint) + .map_err(|error| format!("invalid gRPC code-mode host URL: {error}"))?; + if !matches!(target.scheme(), "http" | "https") { + return Err("gRPC code-mode host URL must use http or https".to_string()); + } + if target.path() != "/" + || target.query().is_some() + || target.fragment().is_some() + { + return Err("gRPC code-mode host URL must not include a path, query, or fragment".to_string()); + } + let origin: Uri = endpoint + .parse() + .map_err(|error| format!("invalid gRPC code-mode host origin: {error}"))?; + let endpoint = endpoint.clone(); + let http_client_factory = http_client_factory.clone(); + let client = tokio::task::spawn_blocking(move || { + http_client_factory + .build_reqwest_client( + reqwest::Client::builder() + .http2_prior_knowledge() + .redirect(reqwest::redirect::Policy::none()), + &endpoint, + ClientRouteClass::Other, + ) + .map_err(|error| { + format!( + "failed to configure gRPC code-mode host transport: {error}" + ) + }) + }) .await .map_err(|error| { - format!("failed to connect to gRPC code-mode host: {error}") - }), - TransportEndpoint::Connected(channel) => Ok(channel.clone()), - } + format!("gRPC code-mode host transport task failed: {error}") + })??; + let transport = service_fn(move |request: Request| { + let client = client.clone(); + 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 response: Response = + client.execute(request).await.map_err(io::Error::other)?.into(); + Ok::<_, io::Error>(response.map(Body::new)) + } + }); + CodeModeHostClient::with_origin(BoxCloneSyncService::new(transport), origin) + } + TransportEndpoint::Connected(channel) => { + let transport = channel.clone().map_err(io::Error::other); + CodeModeHostClient::new(BoxCloneSyncService::new(transport)) + } + }; + Ok(client + .max_decoding_message_size(MAX_FRAME_BYTES) + .max_encoding_message_size(MAX_FRAME_BYTES)) }) .await .cloned() diff --git a/codex-rs/deny.toml b/codex-rs/deny.toml index d785e1f2f3..d6ccfcc642 100644 --- a/codex-rs/deny.toml +++ b/codex-rs/deny.toml @@ -238,6 +238,8 @@ deny = [ { crate = "reqwest", wrappers = [ # Intended owner. "codex-http-client", + # Intentional exception: gRPC uses reqwest directly as tonic's HTTP/2 transport. + "codex-code-mode", # Intentional exception: this standalone, privileged Responses API proxy owns its blocking # upstream HTTP transport independently of Codex. "codex-responses-api-proxy",