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

codex-rs/code-mode-host/src/grpc_transport.rs

@@ -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;
@@ -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(
+        .accept_http1(/*accept_http1*/ true)
-            CodeModeHostServer::new(GrpcCodeModeHost::new())
+        .add_routes(routes.into())
-                .max_decoding_message_size(MAX_FRAME_BYTES)
-                .max_encoding_message_size(MAX_FRAME_BYTES),
-        )
         .serve_with_incoming(listener)
         .await
         .context("code-mode gRPC TCP listener failed")

codex-rs/code-mode-host/tests/grpc_tcp.rs

@@ -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")?)