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

codex-rs/code-mode-host/src/grpc/mod.rs

@@ -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>;

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

@@ -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),
         )

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

@@ -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),

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

@@ -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.

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

@@ -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,
+        }))
+    }
+}

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

@@ -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
+        }
     }
 }
 

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

@@ -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"]

codex-rs/code-mode-host/tests/grpc/process_metrics_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));
+    }
+}

codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto

@@ -10,6 +10,10 @@ package codex.code_mode.v1;
 // unknown session_id. Clients may retire the session binding on this status.
 // Missing cells use WaitResponse.missing_cell instead.
 service CodeModeHost {
+  // Optional host-level observation. UNIMPLEMENTED means no collector. Sampling
+  // never opens/closes a session and failures affect only this RPC stream.
+  rpc GetProcessMetrics(GetProcessMetricsRequest) returns (GetProcessMetricsResponse);
+
   // Opens a session lease. The first event is always SessionOpened; dropping
   // this stream closes the session and terminates its active cells.
   rpc OpenSession(OpenSessionRequest) returns (stream SessionEvent);
@@ -43,6 +47,15 @@ service CodeModeHost {
   rpc Terminate(TerminateRequest) returns (WaitResponse);
 }
 
+message GetProcessMetricsRequest {}
+
+message GetProcessMetricsResponse {
+  uint32 process_id = 1;
+  double process_start_time_seconds = 2;
+  double cpu_seconds_total = 3;
+  uint64 resident_memory_bytes = 4;
+}
+
 message OpenSessionRequest {
   optional SessionCellExecutionLimits cell_execution_limits = 1;
   // Absent preserves the legacy client-side truncation behavior.

codex-rs/code-mode-protocol/src/lib.rs

@@ -2,6 +2,7 @@ mod description;
 pub mod grpc;
 pub mod host;
 mod json_schema_types;
+mod process_metrics;
 mod response;
 mod runtime;
 mod session;
@@ -24,6 +25,8 @@ pub use description::parse_exec_source;
 pub use description::render_code_mode_sample;
 pub use json_schema_types::DEFAULT_INPUT_SCHEMA_MAX_BYTES;
 pub use json_schema_types::render_json_schema_to_typescript;
+pub use process_metrics::CodeModeProcessMetrics;
+pub use process_metrics::CodeModeProcessMetricsSample;
 pub use response::DEFAULT_IMAGE_DETAIL;
 pub use response::FunctionCallOutputContentItem;
 pub use response::ImageDetail;

codex-rs/code-mode-protocol/src/process_metrics.rs

@@ -0,0 +1,21 @@
+//! Optional process observations, independent of sessions. Process-owned
+//! providers observe existing children; remote transports own connection policy.
+
+/// A cumulative observation of one host process, shared across its sessions.
+#[derive(Clone, Debug, PartialEq)]
+pub struct CodeModeProcessMetrics {
+    pub process_id: u32,
+    /// Unix timestamp, stable for this process across all metrics transports.
+    pub process_start_time_seconds: f64,
+    pub cpu_seconds_total: f64,
+    pub resident_memory_bytes: u64,
+}
+
+/// Unsupported capability and an absent process are distinct from a failed sample.
+#[derive(Clone, Debug, PartialEq)]
+pub enum CodeModeProcessMetricsSample {
+    Unsupported,
+    /// The process owner knows no host is running; never inferred from an RPC status.
+    NotRunning,
+    Snapshot(CodeModeProcessMetrics),
+}

codex-rs/code-mode-protocol/src/session.rs

@@ -223,6 +223,16 @@ pub trait CodeModeSessionProvider: Send + Sync {
         Ok(())
     }
 
+    /// Observes a host without starting a session. Process-owned providers must
+    /// not start or replace a host. Network providers may connect or reconnect
+    /// according to their transport policy. Sampling failures and cancellation
+    /// must not explicitly reset the shared execution connection.
+    fn sample_process_metrics(
+        &self,
+    ) -> CodeModeSessionResultFuture<'_, crate::CodeModeProcessMetricsSample> {
+        Box::pin(async { Ok(crate::CodeModeProcessMetricsSample::Unsupported) })
+    }
+
     fn create_session(&self) -> CodeModeSessionProviderFuture<'_>;
 
     /// Creates a session whose cells share the supplied execution limits.

codex-rs/code-mode/src/grpc_session/mod.rs

@@ -46,6 +46,7 @@ mod conversion;
 mod deadline;
 mod generation;
 mod operations;
+mod process_metrics;
 mod reconnect;
 mod state;
 mod transport;
@@ -65,6 +66,11 @@ fn inject_span_traceparent<T>(request: &mut tonic::Request<T>, span: &tracing::S
 }
 
 /// Creates code-mode sessions over an HTTP/2 gRPC connection.
+///
+/// Metrics requests do not open sessions, retry RPCs, or explicitly reset the
+/// connection. URL providers may initialize their client or establish a network
+/// connection; caller-supplied channels and transports retain their reconnect
+/// policy. Preventing child creation requires a process-owned transport.
 #[derive(Clone)]
 pub struct GrpcCodeModeSessionProvider {
     transport: Arc<SharedTransport>,
@@ -94,7 +100,8 @@ impl GrpcCodeModeSessionProvider {
         Self::from_transport(SharedTransport::with_channel(channel))
     }
 
-    /// Uses a caller-owned transport to control connection establishment and recovery.
+    /// Uses a caller-owned transport, which can prevent observation RPCs from
+    /// starting or replacing process-owned connections.
     pub fn with_transport<T>(transport: T) -> Self
     where
         T: tower::Service<
@@ -238,6 +245,13 @@ impl GrpcCodeModeSessionProvider {
 }
 
 impl CodeModeSessionProvider for GrpcCodeModeSessionProvider {
+    fn sample_process_metrics(
+        &self,
+    ) -> CodeModeSessionResultFuture<'_, codex_code_mode_protocol::CodeModeProcessMetricsSample>
+    {
+        Box::pin(process_metrics::sample(self))
+    }
+
     fn create_session(&self) -> CodeModeSessionProviderFuture<'_> {
         self.create_session_with_limits(CodeModeSessionCellExecutionLimits::default())
     }

codex-rs/code-mode/src/grpc_session/process_metrics.rs

@@ -0,0 +1,50 @@
+//! A bounded, non-retrying host observation. No session binding or recovery is
+//! involved: a failed RPC never explicitly retires the execution connection.
+//! Connection establishment and reconnection remain the transport's policy.
+
+use std::time::Duration;
+
+use codex_code_mode_protocol::CodeModeProcessMetrics;
+use codex_code_mode_protocol::CodeModeProcessMetricsSample;
+use codex_code_mode_protocol::grpc;
+use tonic::Code;
+
+use super::GrpcCodeModeSessionProvider;
+
+pub(super) async fn sample(
+    provider: &GrpcCodeModeSessionProvider,
+) -> Result<CodeModeProcessMetricsSample, String> {
+    let deadline = Duration::from_secs(/*secs*/ 1);
+    tokio::time::timeout(deadline, async {
+        let mut client = provider.transport.client().await?;
+        let mut request = tonic::Request::new(grpc::GetProcessMetricsRequest {});
+        request.set_timeout(deadline);
+        match client.get_process_metrics(request).await {
+            Ok(response) => {
+                let snapshot = response.into_inner();
+                if snapshot.process_id == 0
+                    || !snapshot.process_start_time_seconds.is_finite()
+                    || snapshot.process_start_time_seconds <= 0.0
+                    || !snapshot.cpu_seconds_total.is_finite()
+                    || snapshot.cpu_seconds_total < 0.0
+                {
+                    return Err("invalid code-mode process metrics snapshot".to_string());
+                }
+                Ok(CodeModeProcessMetricsSample::Snapshot(
+                    CodeModeProcessMetrics {
+                        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,
+                    },
+                ))
+            }
+            Err(status) if status.code() == Code::Unimplemented => {
+                Ok(CodeModeProcessMetricsSample::Unsupported)
+            }
+            Err(status) => Err(format!("code-mode process metrics failed: {status}")),
+        }
+    })
+    .await
+    .map_err(|_| "code-mode process metrics timed out".to_string())?
+}

codex-rs/core/src/code_mode_host.rs

@@ -1,6 +1,6 @@
 //! Owns the local gRPC code-mode client and host process. Sessions share a lazy
-//! HTTP/2 connection. Only session admission can replace a dead host; other RPCs
-//! use the current generation. Dropping its I/O kills and reaps the child.
+//! HTTP/2 connection. Only session admission can replace a dead host; metrics
+//! use the current generation and cannot reconnect. Dropping its I/O reaps the child.
 
 use std::io;
 use std::path::Path;
@@ -17,9 +17,11 @@ use std::task::Context;
 use std::task::Poll;
 use std::time::Duration;
 
+use codex_code_mode::CodeModeProcessMetricsSample;
 use codex_code_mode::CodeModeSessionCellExecutionLimits;
 use codex_code_mode::CodeModeSessionProvider;
 use codex_code_mode::CodeModeSessionProviderFuture;
+use codex_code_mode::CodeModeSessionResultFuture;
 use codex_code_mode::GrpcCodeModeSessionProvider;
 use codex_code_mode::ProcessOwnedCodeModeSessionProvider;
 use codex_protocol::shell_environment::scrub_non_inheritable_env_vars;
@@ -63,6 +65,26 @@ impl CodeModeSessionProvider for ProcessOwnedGrpcCodeModeSessionProvider {
             .availability()
     }
 
+    fn sample_process_metrics(
+        &self,
+    ) -> CodeModeSessionResultFuture<'_, CodeModeProcessMetricsSample> {
+        // Absence is local lifecycle state, not a meaning assigned to any remote
+        // gRPC error. A child dying after this observation produces a failed sample.
+        if !self
+            .current
+            .lock()
+            .unwrap_or_else(PoisonError::into_inner)
+            .as_ref()
+            .is_some_and(|host| host.alive.load(Ordering::Acquire))
+        {
+            return Box::pin(async { Ok(CodeModeProcessMetricsSample::NotRunning) });
+        }
+        match self.provider.get() {
+            Some(provider) => provider.sample_process_metrics(),
+            None => Box::pin(async { Ok(CodeModeProcessMetricsSample::NotRunning) }),
+        }
+    }
+
     fn create_session(&self) -> CodeModeSessionProviderFuture<'_> {
         self.create_session_with_limits(CodeModeSessionCellExecutionLimits::default())
     }
@@ -89,8 +111,8 @@ impl CodeModeSessionProvider for ProcessOwnedGrpcCodeModeSessionProvider {
 }
 
 // Every channel gets exactly one child. Tonic's internal reconnect loop cannot
-// spawn a process, including if the child dies between channel selection and
-// dispatch. Session admission selects a new generation during recovery.
+// spawn a process, including if the child dies between a metrics liveness check
+// and dispatch. Session admission selects a new generation during recovery.
 struct HostTransport {
     host_program: PathBuf,
     current: Arc<Mutex<Option<HostChannel>>>,
@@ -133,7 +155,7 @@ impl HostTransport {
                     .map_err(|error| Status::unavailable(error.to_string()))?;
                 let alive = connection.alive.clone();
                 let connection = Arc::new(Mutex::new(Some(connection)));
-                // Detect a live but unresponsive child even while the connection is idle.
+                // Preserve execution liveness detection independently of metrics cadence.
                 let channel = Endpoint::from_static("http://code-mode-host")
                     .http2_keep_alive_interval(Duration::from_secs(/*secs*/ 30))
                     .keep_alive_timeout(Duration::from_secs(/*secs*/ 10))

codex-rs/core/src/code_mode_host_lifecycle_tests.rs

@@ -4,11 +4,13 @@
 use std::sync::Arc;
 use std::sync::Mutex;
 use std::sync::atomic::AtomicBool;
+use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering;
 use std::time::Duration;
 
 use anyhow::Context;
 use anyhow::Result;
+use codex_code_mode::CodeModeProcessMetricsSample;
 use codex_code_mode::CodeModeSessionProvider;
 use codex_code_mode::FunctionCallOutputContentItem;
 use codex_code_mode::GrpcCodeModeSessionProvider;
@@ -18,6 +20,7 @@ use pretty_assertions::assert_eq;
 use tokio::sync::Notify;
 use tokio::sync::Semaphore;
 use tokio::time::timeout;
+use tokio_util::sync::CancellationToken;
 use tonic::Status;
 use tonic::body::Body;
 use tower::ServiceExt;
@@ -105,6 +108,25 @@ async fn child_exit_after_channel_selection_does_not_replay_execution() -> Resul
     )
     .await?
     .map_err(anyhow::Error::msg)?;
+    let samples = Arc::new(AtomicUsize::new(0));
+    let stop = CancellationToken::new();
+    let sampler = tokio::spawn({
+        let provider = provider.clone();
+        let samples = samples.clone();
+        let stop = stop.clone();
+        async move {
+            let mut interval = tokio::time::interval(Duration::from_millis(/*millis*/ 5));
+            loop {
+                tokio::select! {
+                    _ = stop.cancelled() => break,
+                    _ = interval.tick() => {
+                        let _ = provider.sample_process_metrics().await;
+                        samples.fetch_add(1, Ordering::SeqCst);
+                    }
+                }
+            }
+        }
+    });
     armed.store(true, Ordering::SeqCst);
     let pending = tokio::spawn({
         let session = sessions[0].clone();
@@ -123,6 +145,8 @@ async fn child_exit_after_channel_selection_does_not_replay_execution() -> Resul
     })
     .await
     .context("child was not reaped")?;
+    let absent = provider.sample_process_metrics().await.unwrap_err();
+    assert!(absent.contains("host is not running"), "{absent}");
     assert_eq!(recorder.invocations()?.len(), 1);
     dispatch.add_permits(1);
     assert!(timeout(TEST_TIMEOUT, pending).await??.is_err());
@@ -147,6 +171,16 @@ async fn child_exit_after_channel_selection_does_not_replay_execution() -> Resul
         Ok::<_, anyhow::Error>(())
     }));
     timeout(TEST_TIMEOUT, recovery).await??;
+    assert_eq!(
+        provider
+            .sample_process_metrics()
+            .await
+            .map_err(anyhow::Error::msg)?,
+        CodeModeProcessMetricsSample::Unsupported
+    );
+    stop.cancel();
+    timeout(TEST_TIMEOUT, sampler).await??;
+    assert!(samples.load(Ordering::SeqCst) > 0);
     assert_eq!(recorder.invocations()?.len(), 2);
     for session in sessions {
         session.shutdown().await.map_err(anyhow::Error::msg)?;

codex-rs/core/src/code_mode_host_tests.rs

@@ -5,6 +5,7 @@ use std::time::Duration;
 
 use anyhow::Context;
 use anyhow::Result;
+use codex_code_mode::CodeModeProcessMetricsSample;
 use codex_code_mode::CodeModeSession;
 use codex_code_mode::CodeModeSessionProvider;
 use codex_code_mode::ExecuteRequest;
@@ -68,8 +69,24 @@ async fn grpc_stdio_sessions_share_a_host_and_isolate_state() -> Result<()> {
     #[cfg(not(unix))]
     let host_program = codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")?;
     let provider = ProcessOwnedGrpcCodeModeSessionProvider::with_host_program(host_program);
+    assert_eq!(
+        provider
+            .sample_process_metrics()
+            .await
+            .map_err(anyhow::Error::msg)?,
+        CodeModeProcessMetricsSample::NotRunning
+    );
+    #[cfg(unix)]
+    assert!(recorder.invocations()?.is_empty());
     let (first, second) = tokio::try_join!(provider.create_session(), provider.create_session())
         .map_err(anyhow::Error::msg)?;
+    assert_eq!(
+        provider
+            .sample_process_metrics()
+            .await
+            .map_err(anyhow::Error::msg)?,
+        CodeModeProcessMetricsSample::Unsupported
+    );
     assert_eq!(
         execute(&first, "store(\"stored\", 42); text(load(\"stored\"));").await?,
         vec![FunctionCallOutputContentItem::InputText {
@@ -111,7 +128,7 @@ async fn grpc_stdio_sessions_share_a_host_and_isolate_state() -> Result<()> {
             }
         })
         .await
-        .context("host was not reaped after provider teardown")?;
+        .context("metrics retained the host after provider teardown")?;
     }
     Ok(())
 }
@@ -143,6 +160,16 @@ async fn hung_grpc_stdio_host_is_reaped_and_the_session_recovers() -> Result<()>
     if unsafe { libc::kill(pid as libc::pid_t, libc::SIGSTOP) } != 0 {
         return Err(std::io::Error::last_os_error().into());
     }
+    // A metrics deadline cancels only its stream. It must neither reap the
+    // suspended child nor advance the independent 30s/10s keepalive deadline.
+    assert!(provider.sample_process_metrics().await.is_err());
+    assert_eq!(
+        unsafe {
+            libc::kill(pid as libc::pid_t, /*sig*/ 0)
+        },
+        0
+    );
+    assert_eq!(recorder.invocations()?.len(), 1);
     timeout(Duration::from_secs(/*secs*/ 55), async {
         while unsafe {
             libc::kill(pid as libc::pid_t, /*sig*/ 0)
@@ -153,6 +180,18 @@ async fn hung_grpc_stdio_host_is_reaped_and_the_session_recovers() -> Result<()>
     })
     .await
     .context("unresponsive host was not killed and reaped")?;
+    assert_eq!(
+        provider
+            .sample_process_metrics()
+            .await
+            .map_err(anyhow::Error::msg)?,
+        CodeModeProcessMetricsSample::NotRunning
+    );
+    assert_eq!(
+        recorder.invocations()?.len(),
+        1,
+        "metrics must not restart a dead host"
+    );
     let reset = execute(&session, r#"store("recovery", "must not run");"#)
         .await
         .unwrap_err();

codex-rs/core/src/thread_manager.rs

@@ -715,6 +715,18 @@ impl ThreadManager {
         self
     }
 
+    /// Samples host metrics without creating a Codex thread or code-mode session.
+    /// Process-owned providers observe existing children; custom providers retain
+    /// their transport's connection policy.
+    pub async fn sample_code_mode_process_metrics(
+        &self,
+    ) -> Result<codex_code_mode::CodeModeProcessMetricsSample, String> {
+        self.state
+            .code_mode_session_provider
+            .sample_process_metrics()
+            .await
+    }
+
     /// Replaces the process-wide provider before this manager is shared with threads.
     pub fn with_code_mode_session_provider(
         mut self,