Add optional process metrics sampling for code-mode hosts (#57841)

## What changed - Add `GetProcessMetrics` and provider sampling APIs for process ID, start time, cumulative CPU time, and resident memory, exposed through `ThreadManager` without creating a thread or session. - Distinguish unsupported collection, a locally absent host, and sampling failures. Process-owned providers never start or replace a host when sampling. - Bound gRPC sampling to one second without retries or explicit connection resets. Run optional collectors on a detached worker with at most one collection in flight, so stalled collection cannot block the HTTP/2 driver or runtime shutdown. ## Testing Add coverage for snapshots, collector errors and panics, deadlines, overload, and shutdown with a stuck collector. Extend lifecycle tests to check that sampling preserves session state and keepalive behavior without starting or restarting local hosts. GitOrigin-RevId: eeafe275891d0464636f28a29345a87aed1a12dc
2026-10-29 19:08:43 +00:00 · 2026-10-28 21:34:25 +00:00
parent 7ce5fa7bdc
commit d373d64967
18 changed files with 559 additions and 15 deletions
@@ -47,6 +47,7 @@ fn trace_context_from_request<T>(request: &Request<T>) -> Option<W3cTraceContext
 #[derive(Clone)]
 pub struct GrpcCodeModeHost {
    state: Arc<GrpcHostState>,
+    process_metrics: Option<Arc<crate::process_metrics::ProcessMetrics>>,
 }

 impl GrpcCodeModeHost {
@@ -54,6 +55,17 @@ impl GrpcCodeModeHost {
    pub fn new() -> Self {
        Self {
            state: Arc::new(GrpcHostState::new()),
+            process_metrics: None,
+        }
+    }
+
+    /// Uses a collector whose resources were prepared before sandbox entry.
+    pub fn with_process_metrics(collector: crate::CodeModeProcessMetricsCollector) -> Self {
+        Self {
+            process_metrics: Some(Arc::new(crate::process_metrics::ProcessMetrics::new(
+                collector,
+            ))),
+            ..Self::new()
        }
    }

@@ -302,6 +314,23 @@ impl Default for GrpcCodeModeHost {
 }

 impl CodeModeHost for GrpcCodeModeHost {
+    fn get_process_metrics<'a, 'async_trait>(
+        &'a self,
+        _request: Request<proto::GetProcessMetricsRequest>,
+    ) -> GrpcFuture<'async_trait, proto::GetProcessMetricsResponse>
+    where
+        'a: 'async_trait,
+        Self: 'async_trait,
+    {
+        Box::pin(async move {
+            self.process_metrics
+                .as_ref()
+                .ok_or_else(|| Status::unimplemented("process metrics are unsupported"))?
+                .sample()
+                .await
+        })
+    }
+
    type OpenSessionStream = GrpcStream<proto::SessionEvent>;
    type SubscribeToToolCallsStream = GrpcStream<proto::ToolCall>;
    type ExecuteStream = GrpcStream<proto::ExecuteEvent>;
@@ -20,7 +20,7 @@ use tonic::transport::server::Connected;

 use crate::GrpcCodeModeHost;

-pub(crate) async fn run() -> Result<()> {
+pub(crate) async fn run(host: GrpcCodeModeHost) -> Result<()> {
    let closed = CancellationToken::new();
    let connection = StdioConnection {
        io: tokio::io::join(tokio::io::stdin(), tokio::io::stdout()),
@@ -31,7 +31,7 @@ pub(crate) async fn run() -> Result<()> {
        .chain(futures::stream::pending());
    let server = Server::builder()
        .add_service(
-            CodeModeHostServer::new(GrpcCodeModeHost::new())
+            CodeModeHostServer::new(host)
                .max_decoding_message_size(MAX_FRAME_BYTES)
                .max_encoding_message_size(MAX_FRAME_BYTES),
        )
@@ -21,7 +21,10 @@ use tracing::info;

 use crate::GrpcCodeModeHost;

-pub(super) async fn run_tcp_listener(bind_address: SocketAddr) -> Result<()> {
+pub(super) async fn run_tcp_listener(
+    bind_address: SocketAddr,
+    host: GrpcCodeModeHost,
+) -> Result<()> {
    let listener = bind_tcp_listener(bind_address).await?;
    let local_address = listener
        .local_addr()
@@ -33,7 +36,7 @@ pub(super) async fn run_tcp_listener(bind_address: SocketAddr) -> Result<()> {
        .context("failed to publish code-mode gRPC listen address")?;

    let routes = Routes::new(
-        CodeModeHostServer::new(GrpcCodeModeHost::new())
+        CodeModeHostServer::new(host)
            .send_compressed(tonic::codec::CompressionEncoding::Zstd)
            .max_decoding_message_size(MAX_FRAME_BYTES)
            .max_encoding_message_size(MAX_FRAME_BYTES),
@@ -44,13 +44,16 @@ use self::delegate::RemoteDelegate;
 use self::peer::HostPeer;

 pub use self::grpc::GrpcCodeModeHost;
+pub use self::process_metrics::CodeModeProcessMetricsCollector;
 pub use self::transport::DEFAULT_LISTEN_URL;
+pub use codex_code_mode_protocol::CodeModeProcessMetrics;

 mod delegate;
 mod grpc;
 mod grpc_stdio;
 mod grpc_transport;
 mod peer;
+mod process_metrics;
 mod transport;

 const MAX_IN_FLIGHT_REQUESTS: usize = 256;
@@ -95,7 +98,20 @@ impl HostLimits {
    fields(otel.name = "code_mode_host.run_main")
 )]
 pub async fn run_main(listen_url: &str) -> Result<()> {
-    transport::run_transport(listen_url).await
+    transport::run_transport(listen_url, GrpcCodeModeHost::new()).await
+}
+
+/// Runs with an optional gRPC metrics collector initialized before sandbox entry.
+/// The legacy stdio transport remains unchanged.
+pub async fn run_main_with_process_metrics(
+    listen_url: &str,
+    collector: CodeModeProcessMetricsCollector,
+) -> Result<()> {
+    transport::run_transport(
+        listen_url,
+        GrpcCodeModeHost::with_process_metrics(collector),
+    )
+    .await
 }

 /// Runs one code-mode host connection over the process standard streams.
@@ -0,0 +1,69 @@
+//! Isolates optional synchronous metrics collection from the HTTP/2 driver and
+//! session admission. A stuck collector occupies at most one detached thread,
+//! which is never joined during host or Tokio runtime shutdown.
+
+use std::panic::AssertUnwindSafe;
+use std::panic::catch_unwind;
+use std::sync::Arc;
+
+use codex_code_mode_protocol::CodeModeProcessMetrics;
+use codex_code_mode_protocol::grpc;
+use tokio::sync::Semaphore;
+use tokio::sync::oneshot;
+use tonic::Response;
+use tonic::Status;
+
+/// Collectors capture any resources they need before the host enters its sandbox.
+/// They must be bounded, read-only observations that do not acquire execution
+/// locks. Isolation covers unwinding panics and stalled collection, not process
+/// aborts, resource exhaustion, or interference through shared state. A collector
+/// that never returns disables further sampling for this host; it is not retried
+/// on another thread, and its worker is not joined during Tokio runtime shutdown.
+pub type CodeModeProcessMetricsCollector =
+    Arc<dyn Fn() -> Result<CodeModeProcessMetrics, String> + Send + Sync>;
+
+pub(crate) struct ProcessMetrics {
+    collector: CodeModeProcessMetricsCollector,
+    permits: Arc<Semaphore>,
+}
+
+impl ProcessMetrics {
+    pub(crate) fn new(collector: CodeModeProcessMetricsCollector) -> Self {
+        Self {
+            collector,
+            permits: Arc::new(Semaphore::new(/*permits*/ 1)),
+        }
+    }
+
+    pub(crate) async fn sample(&self) -> Result<Response<grpc::GetProcessMetricsResponse>, Status> {
+        let permit = self
+            .permits
+            .clone()
+            .try_acquire_owned()
+            .map_err(|_| Status::resource_exhausted("process metrics collector is busy"))?;
+        let collector = self.collector.clone();
+        let (sender, receiver) = oneshot::channel();
+        // Do not use Tokio's blocking pool: runtime drop joins running blocking
+        // tasks even after the request times out. Dropping this thread's handle
+        // detaches it; the permit still bounds abandoned collection to one thread.
+        std::thread::Builder::new()
+            .name("code-mode-metrics".to_string())
+            .spawn(move || {
+                let result = catch_unwind(AssertUnwindSafe(|| collector()));
+                drop(permit);
+                let _ = sender.send(result);
+            })
+            .map_err(|_| Status::internal("process metrics worker could not start"))?;
+        let snapshot = receiver
+            .await
+            .map_err(|_| Status::internal("process metrics worker ended"))?
+            .map_err(|_| Status::internal("process metrics collector panicked"))?
+            .map_err(|_| Status::internal("process metrics collection failed"))?;
+        Ok(Response::new(grpc::GetProcessMetricsResponse {
+            process_id: snapshot.process_id,
+            process_start_time_seconds: snapshot.process_start_time_seconds,
+            cpu_seconds_total: snapshot.cpu_seconds_total,
+            resident_memory_bytes: snapshot.resident_memory_bytes,
+        }))
+    }
+}
@@ -15,11 +15,13 @@ enum ListenTransport {
    Grpc(SocketAddr),
 }

-pub(crate) async fn run_transport(listen_url: &str) -> Result<()> {
+pub(crate) async fn run_transport(listen_url: &str, host: crate::GrpcCodeModeHost) -> 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,
+        ListenTransport::GrpcStdio => crate::grpc_stdio::run(host).await,
+        ListenTransport::Grpc(bind_address) => {
+            grpc_transport::run_tcp_listener(bind_address, host).await
+        }
    }
 }

@@ -52,6 +52,8 @@ mod image_cancellation_tests;
 mod large_tool_delegate;
 #[path = "grpc/network_policy_tests.rs"]
 mod network_policy_tests;
+#[path = "grpc/process_metrics_tests.rs"]
+mod process_metrics_tests;
 #[path = "support/recording_delegate.rs"]
 mod recording_delegate;
 #[path = "grpc/session_recovery_tests.rs"]
@@ -0,0 +1,205 @@
+//! Process observations are independent of session state and HTTP/2 keepalive.
+
+use std::sync::Mutex;
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
+
+use codex_code_mode::CodeModeProcessMetrics;
+use codex_code_mode::CodeModeProcessMetricsSample;
+use futures::StreamExt;
+use pretty_assertions::assert_eq;
+use tokio_stream::wrappers::TcpListenerStream;
+use tonic::transport::Endpoint;
+
+use super::*;
+
+#[tokio::test]
+async fn metrics_failures_and_deadlines_preserve_sessions_and_keepalive() -> Result<()> {
+    let snapshot = CodeModeProcessMetrics {
+        process_id: 42,
+        process_start_time_seconds: 1_700_000_000.0,
+        cpu_seconds_total: 1.5,
+        resident_memory_bytes: 4096,
+    };
+    let expected = snapshot.clone();
+    let calls = Arc::new(AtomicUsize::new(0));
+    let collector_calls = calls.clone();
+    let (release, blocked) = std::sync::mpsc::channel::<()>();
+    let blocked = Mutex::new(blocked);
+    let host = codex_code_mode_host::GrpcCodeModeHost::with_process_metrics(Arc::new(move || {
+        match collector_calls.fetch_add(1, Ordering::SeqCst) {
+            0 => Ok(snapshot.clone()),
+            1 => Err("metrics unavailable".to_string()),
+            2 => panic!("collector panic"),
+            _ => {
+                blocked
+                    .lock()
+                    .unwrap()
+                    .recv()
+                    .map_err(|error| error.to_string())?;
+                Ok(snapshot.clone())
+            }
+        }
+    }));
+    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
+    let address = listener.local_addr()?;
+    let connections = Arc::new(AtomicUsize::new(0));
+    let accepted = connections.clone();
+    let incoming = TcpListenerStream::new(listener).inspect(move |connection| {
+        if connection.is_ok() {
+            accepted.fetch_add(1, Ordering::SeqCst);
+        }
+    });
+    let server = tokio_util::task::AbortOnDropHandle::new(tokio::spawn(async move {
+        tonic::transport::Server::builder()
+            .add_service(grpc::code_mode_host_server::CodeModeHostServer::new(host))
+            .serve_with_incoming(incoming)
+            .await
+    }));
+    // Multiple keepalive intervals elapse while the collector is blocked, even
+    // on this single-threaded runtime. A stalled HTTP/2 driver would lose state.
+    let channel = Endpoint::from_shared(format!("http://{address}"))?
+        .http2_keep_alive_interval(Duration::from_millis(/*millis*/ 50))
+        .keep_alive_timeout(Duration::from_millis(/*millis*/ 100))
+        .keep_alive_while_idle(/*enabled*/ true)
+        .connect()
+        .await?;
+    let provider = GrpcCodeModeSessionProvider::with_channel(channel);
+    let session = provider
+        .create_session()
+        .await
+        .map_err(anyhow::Error::msg)?;
+    execute(
+        &session,
+        request("store('survives', 42);"),
+        Arc::new(NoopCodeModeSessionDelegate),
+    )
+    .await?;
+    assert_eq!(
+        provider
+            .sample_process_metrics()
+            .await
+            .map_err(anyhow::Error::msg)?,
+        CodeModeProcessMetricsSample::Snapshot(expected)
+    );
+    for cell in 2..=4 {
+        assert!(provider.sample_process_metrics().await.is_err());
+        let response = execute(
+            &session,
+            request("text(load('survives'));"),
+            Arc::new(NoopCodeModeSessionDelegate),
+        )
+        .await?;
+        assert_eq!(
+            response,
+            text_response(&cell.to_string(), "42", response.code_mode_host_duration())
+        );
+    }
+    // A canceled blocked RPC retains its collector quota: no unbounded worker
+    // threads are added, and neither the sample nor overload opens another channel.
+    assert!(provider.sample_process_metrics().await.is_err());
+    assert_eq!(calls.load(Ordering::SeqCst), 4);
+    assert_eq!(connections.load(Ordering::SeqCst), 1);
+    release.send(())?;
+    session.shutdown().await.map_err(anyhow::Error::msg)?;
+    drop(server);
+    Ok(())
+}
+
+/// A running server's collection precondition failure is not local host absence.
+#[tokio::test]
+async fn remote_metrics_precondition_is_a_failure() -> Result<()> {
+    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
+    let address = listener.local_addr()?;
+    let server = tokio_util::task::AbortOnDropHandle::new(tokio::spawn(async move {
+        tonic::transport::Server::builder()
+            .add_service(
+                grpc::code_mode_host_server::CodeModeHostServer::with_interceptor(
+                    codex_code_mode_host::GrpcCodeModeHost::new(),
+                    |_request| {
+                        Err(tonic::Status::failed_precondition(
+                            "collector is misconfigured",
+                        ))
+                    },
+                ),
+            )
+            .serve_with_incoming(TcpListenerStream::new(listener))
+            .await
+    }));
+    let channel = Endpoint::from_shared(format!("http://{address}"))?
+        .connect()
+        .await?;
+    let provider = GrpcCodeModeSessionProvider::with_channel(channel);
+    let error = provider.sample_process_metrics().await.unwrap_err();
+    assert!(error.contains("collector is misconfigured"), "{error}");
+    drop(server);
+    Ok(())
+}
+
+/// A permanently stuck collector is bounded and cannot hang normal runtime drop.
+#[test]
+fn stuck_collector_does_not_block_runtime_shutdown() -> Result<()> {
+    const CHILD_ENV: &str = "CODE_MODE_METRICS_SHUTDOWN_TEST_CHILD";
+    if let Some(shutdown_marker) = std::env::var_os(CHILD_ENV) {
+        let runtime = tokio::runtime::Builder::new_current_thread()
+            .enable_all()
+            .build()?;
+        runtime.block_on(async {
+            let calls = Arc::new(AtomicUsize::new(0));
+            let collector_calls = calls.clone();
+            let host =
+                codex_code_mode_host::GrpcCodeModeHost::with_process_metrics(Arc::new(move || {
+                    collector_calls.fetch_add(1, Ordering::SeqCst);
+                    loop {
+                        std::thread::park();
+                    }
+                }));
+            let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
+            let address = listener.local_addr()?;
+            let server = tokio_util::task::AbortOnDropHandle::new(tokio::spawn(async move {
+                tonic::transport::Server::builder()
+                    .add_service(grpc::code_mode_host_server::CodeModeHostServer::new(host))
+                    .serve_with_incoming(TcpListenerStream::new(listener))
+                    .await
+            }));
+            let channel = Endpoint::from_shared(format!("http://{address}"))?
+                .connect()
+                .await?;
+            let provider = GrpcCodeModeSessionProvider::with_channel(channel);
+            assert!(provider.sample_process_metrics().await.is_err());
+            let busy = provider.sample_process_metrics().await.unwrap_err();
+            assert!(busy.contains("collector is busy"), "{busy}");
+            assert_eq!(calls.load(Ordering::SeqCst), 1);
+            drop(server);
+            Ok::<_, anyhow::Error>(())
+        })?;
+        drop(runtime);
+        std::fs::write(shutdown_marker, "runtime dropped")?;
+        return Ok(());
+    }
+
+    let directory = tempfile::tempdir()?;
+    let shutdown_marker = directory.path().join("shutdown-completed");
+    let mut child = std::process::Command::new(std::env::current_exe()?)
+        .args([
+            "--exact",
+            "process_metrics_tests::stuck_collector_does_not_block_runtime_shutdown",
+            "--nocapture",
+        ])
+        .env(CHILD_ENV, &shutdown_marker)
+        .spawn()?;
+    let deadline = std::time::Instant::now() + Duration::from_secs(/*secs*/ 10);
+    loop {
+        if let Some(status) = child.try_wait()? {
+            assert!(status.success(), "metrics subprocess failed: {status}");
+            assert_eq!(std::fs::read_to_string(shutdown_marker)?, "runtime dropped");
+            return Ok(());
+        }
+        if std::time::Instant::now() >= deadline {
+            child.kill()?;
+            child.wait()?;
+            anyhow::bail!("a stuck metrics collector blocked runtime shutdown");
+        }
+        std::thread::sleep(Duration::from_millis(/*millis*/ 10));
+    }
+}