Add opt-in gRPC over stdio for the code-mode host (#52723)

## What changed

- Add `grpc+stdio://` as a host transport and the disabled-by-default `code_mode_host_grpc` feature flag to select it.
- Share a lazy HTTP/2 channel and host process across code-mode sessions while keeping session state isolated.
- Use HTTP/2 keepalives to detect unresponsive hosts, kill and reap disconnected children, and spawn a fresh host on reconnect. Exit the host when its stdio connection closes.

## Testing

Add coverage for shared-host session isolation, host exit on stdin EOF, and recovery after an unresponsive host is reaped. Exercise nested tool calls with both legacy stdio and gRPC stdio, asserting that the feature flag selects the expected transport.

GitOrigin-RevId: 31389d8540de236b8e0571d3c24e29e117b75c4a
This commit is contained in:
Channing Conger
2026-10-09 23:26:34 +00:00
committed by copyberry
parent 07757f448f
commit 2dc0791c7d
17 changed files with 566 additions and 18 deletions
+88
View File
@@ -0,0 +1,88 @@
//! Serves one HTTP/2 gRPC connection on stdin/stdout. Stdout contains only protocol
//! bytes, and closing the connection stops the host rather than accepting another client.
use std::io;
use std::pin::Pin;
use std::task::Context;
use std::task::Poll;
use anyhow::Context as _;
use anyhow::Result;
use codex_code_mode_protocol::grpc::code_mode_host_server::CodeModeHostServer;
use codex_code_mode_protocol::host::MAX_FRAME_BYTES;
use futures::StreamExt;
use tokio::io::AsyncRead;
use tokio::io::AsyncWrite;
use tokio::io::ReadBuf;
use tokio_util::sync::CancellationToken;
use tonic::transport::Server;
use tonic::transport::server::Connected;
use crate::GrpcCodeModeHost;
pub(crate) async fn run() -> Result<()> {
let closed = CancellationToken::new();
let connection = StdioConnection {
io: tokio::io::join(tokio::io::stdin(), tokio::io::stdout()),
closed: closed.clone(),
};
// Keep the incoming stream open while the single connection is being served.
let incoming = futures::stream::once(async { Ok::<_, io::Error>(connection) })
.chain(futures::stream::pending());
let server = Server::builder()
.add_service(
CodeModeHostServer::new(GrpcCodeModeHost::new())
.max_decoding_message_size(MAX_FRAME_BYTES)
.max_encoding_message_size(MAX_FRAME_BYTES),
)
.serve_with_incoming(incoming);
tokio::select! {
result = server => result.context("code-mode gRPC stdio connection failed"),
_ = closed.cancelled() => Ok(()),
}
}
struct StdioConnection {
io: tokio::io::Join<tokio::io::Stdin, tokio::io::Stdout>,
closed: CancellationToken,
}
impl Connected for StdioConnection {
type ConnectInfo = ();
fn connect_info(&self) -> Self::ConnectInfo {}
}
impl AsyncRead for StdioConnection {
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
Pin::new(&mut self.io).poll_read(cx, buf)
}
}
impl AsyncWrite for StdioConnection {
fn poll_write(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
Pin::new(&mut self.io).poll_write(cx, buf)
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.io).poll_flush(cx)
}
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.io).poll_shutdown(cx)
}
}
impl Drop for StdioConnection {
fn drop(&mut self) {
self.closed.cancel();
}
}
+1
View File
@@ -48,6 +48,7 @@ pub use self::transport::DEFAULT_LISTEN_URL;
mod delegate;
mod grpc;
mod grpc_stdio;
mod grpc_transport;
mod peer;
mod transport;
+1 -1
View File
@@ -18,7 +18,7 @@ const OTEL_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 5);
#[derive(Debug, Parser)]
struct Cli {
/// Transport endpoint: `stdio`, `stdio://`, or `grpc://IP:PORT`.
/// Transport endpoint: `stdio`, `stdio://`, `grpc+stdio://`, or `grpc://IP:PORT`.
#[arg(
long,
value_name = "URL",
+6 -1
View File
@@ -11,17 +11,22 @@ pub const DEFAULT_LISTEN_URL: &str = "stdio";
#[derive(Debug, Clone, Eq, PartialEq)]
enum ListenTransport {
Stdio,
GrpcStdio,
Grpc(SocketAddr),
}
pub(crate) async fn run_transport(listen_url: &str) -> Result<()> {
match parse_listen_url(listen_url)? {
ListenTransport::Stdio => crate::run_stdio().await,
ListenTransport::GrpcStdio => crate::grpc_stdio::run().await,
ListenTransport::Grpc(bind_address) => grpc_transport::run_tcp_listener(bind_address).await,
}
}
fn parse_listen_url(listen_url: &str) -> Result<ListenTransport> {
if listen_url == "grpc+stdio://" {
return Ok(ListenTransport::GrpcStdio);
}
if matches!(listen_url, "stdio" | "stdio://") {
return Ok(ListenTransport::Stdio);
}
@@ -36,7 +41,7 @@ fn parse_listen_url(listen_url: &str) -> Result<ListenTransport> {
}
anyhow::bail!(
"unsupported --listen URL `{listen_url}`; expected `grpc://IP:PORT`, `stdio`, or `stdio://`"
"unsupported --listen URL `{listen_url}`; expected `grpc://IP:PORT`, `grpc+stdio://`, `stdio`, or `stdio://`"
);
}
@@ -41,6 +41,10 @@ fn parse_listen_url_accepts_stdio_transports() {
parse_listen_url("stdio://").expect("stdio URL should parse"),
ListenTransport::Stdio
);
assert_eq!(
parse_listen_url("grpc+stdio://").expect("gRPC stdio URL should parse"),
ListenTransport::GrpcStdio
);
}
#[test]