Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Add a health endpoint to the code-mode gRPC listener (#38806)
## What changed

- Serve `GET /healthz` with a `200 OK` response over HTTP/1.1 and HTTP/2.
- Continue requiring HTTP/2 for all other requests so gRPC methods are not exposed over HTTP/1.1.

## Testing

- Add TCP listener integration coverage for HTTP/1.1 and HTTP/2 health checks and rejection of HTTP/1.1 gRPC requests.

GitOrigin-RevId: ae7bbf56323fbc76769375a6d8e90653e8adc860
  • Loading branch information
cconger authored and copyberry committed Aug 15, 2026
commit b3cc21737803549679e2009193c04205f7d8d19c
34 changes: 29 additions & 5 deletions codex-rs/code-mode-host/src/grpc_transport.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,17 @@ use std::net::SocketAddr;

use anyhow::Context;
use anyhow::Result;
use axum::extract::Request;
use axum::http::Method;
use axum::http::StatusCode;
use axum::http::Version;
use axum::middleware;
use axum::middleware::Next;
use axum::routing::get;
use codex_code_mode_protocol::grpc::code_mode_host_server::CodeModeHostServer;
use codex_code_mode_protocol::host::MAX_FRAME_BYTES;
use tokio::net::TcpListener;
use tonic::service::Routes;
use tonic::transport::Server;
use tonic::transport::server::TcpIncoming;
use tracing::info;
Expand All @@ -24,12 +32,28 @@ pub(super) async fn run_tcp_listener(bind_address: SocketAddr) -> Result<()> {
.flush()
.context("failed to publish code-mode gRPC listen address")?;

let routes = Routes::new(
CodeModeHostServer::new(GrpcCodeModeHost::new())
.max_decoding_message_size(MAX_FRAME_BYTES)
.max_encoding_message_size(MAX_FRAME_BYTES),
)
.into_axum_router()
.route("/healthz", get(|| async { StatusCode::OK }))
.layer(middleware::from_fn(
|request: Request, next: Next| async move {
if request.version() != Version::HTTP_2
&& (request.method() != Method::GET || request.uri().path() != "/healthz")
{
return Err(StatusCode::HTTP_VERSION_NOT_SUPPORTED);
}

Ok(next.run(request).await)
},
));

Server::builder()
.add_service(
CodeModeHostServer::new(GrpcCodeModeHost::new())
.max_decoding_message_size(MAX_FRAME_BYTES)
.max_encoding_message_size(MAX_FRAME_BYTES),
)
.accept_http1(/*accept_http1*/ true)
.add_routes(routes.into())
.serve_with_incoming(listener)
.await
.context("code-mode gRPC TCP listener failed")
Expand Down
104 changes: 104 additions & 0 deletions codex-rs/code-mode-host/tests/grpc_tcp.rs
Original file line number Diff line number Diff line change
@@ -1,18 +1,122 @@
use std::future::poll_fn;
use std::process::Stdio;
use std::time::Duration;

use anyhow::Context;
use anyhow::Result;
use axum::http::Request;
use axum::http::StatusCode;
use axum::http::Version;
use codex_code_mode_protocol::grpc;
use codex_code_mode_protocol::grpc::code_mode_host_client::CodeModeHostClient;
use pretty_assertions::assert_eq;
use tokio::io::AsyncBufReadExt;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncWriteExt;
use tokio::io::BufReader;
use tokio::net::TcpStream;
use tokio::process::Command;
use tokio::time::timeout;
use tonic::body::Body;
use tonic::codegen::Service;
use tonic::transport::Endpoint;

#[path = "support/host.rs"]
mod host;

use host::HostHarness;

const TEST_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 10);

#[tokio::test]
async fn tcp_listener_serves_http1_healthz() -> Result<()> {
let host = HostHarness::start("grpc://127.0.0.1:0").await?;
let request = "GET /healthz HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n";
let response = http1_response(&host.endpoint, request.as_bytes()).await?;

assert_eq!(response.lines().next(), Some("HTTP/1.1 200 OK"));
Ok(())
}

#[tokio::test]
async fn tcp_listener_rejects_http1_grpc_requests() -> Result<()> {
let host = HostHarness::start("grpc://127.0.0.1:0").await?;
let mut request = concat!(
"POST /codex.code_mode.v1.CodeModeHost/OpenSession HTTP/1.1\r\n",
"Host: localhost\r\n",
"Content-Type: application/grpc\r\n",
"Content-Length: 5\r\n",
"Connection: close\r\n\r\n"
)
.as_bytes()
.to_vec();
request.extend_from_slice(&[0; 5]);
let response = http1_response(&host.endpoint, &request).await?;

assert_eq!(
response.lines().next(),
Some("HTTP/1.1 505 HTTP Version Not Supported")
);
Ok(())
}

#[tokio::test]
async fn tcp_listener_serves_http2_healthz() -> Result<()> {
let host = HostHarness::start("grpc://127.0.0.1:0").await?;
let mut channel = Endpoint::from_shared(host.endpoint.clone())
.context("gRPC code-mode host published an invalid endpoint")?
.connect_timeout(TEST_TIMEOUT)
.timeout(TEST_TIMEOUT)
.connect()
.await
.context("failed to connect to gRPC code-mode host health endpoint")?;
let request = Request::builder()
.uri(format!("{}/healthz", host.endpoint))
.body(Body::empty())
.context("failed to build gRPC code-mode host health request")?;

let response = timeout(TEST_TIMEOUT, async {
poll_fn(|context| channel.poll_ready(context))
.await
.context("gRPC code-mode host health channel is unavailable")?;
channel
.call(request)
.await
.context("failed to request gRPC code-mode host health")
})
.await
.context("timed out requesting gRPC code-mode host health")??;

assert_eq!(response.version(), Version::HTTP_2);
assert_eq!(response.status(), StatusCode::OK);
Ok(())
}

async fn http1_response(endpoint: &str, request: &[u8]) -> Result<String> {
let address = endpoint
.strip_prefix("http://")
.context("gRPC code-mode host URL should use http://")?;

timeout(TEST_TIMEOUT, async {
let mut stream = TcpStream::connect(address)
.await
.context("failed to connect to gRPC code-mode host")?;
stream
.write_all(request)
.await
.context("failed to send HTTP/1.1 request to gRPC code-mode host")?;

let mut response = String::new();
stream
.read_to_string(&mut response)
.await
.context("failed to read HTTP/1.1 response from gRPC code-mode host")?;
Ok(response)
})
.await
.context("timed out waiting for HTTP/1.1 response from gRPC code-mode host")?
}

#[tokio::test]
async fn tcp_listener_opens_a_grpc_session() -> Result<()> {
let mut host = Command::new(codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")?)
Expand Down
Loading