From d373d649673ca8a5bca71f54be0ae47dbd79b466 Mon Sep 17 00:00:00 2001 From: aakash-oai Date: Wed, 28 Oct 2026 21:21:29 +0000 Subject: [PATCH] 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 | 29 +++ codex-rs/code-mode-host/src/grpc_stdio.rs | 4 +- codex-rs/code-mode-host/src/grpc_transport.rs | 7 +- codex-rs/code-mode-host/src/lib.rs | 18 +- .../code-mode-host/src/process_metrics.rs | 69 ++++++ codex-rs/code-mode-host/src/transport.rs | 8 +- codex-rs/code-mode-host/tests/grpc.rs | 2 + .../tests/grpc/process_metrics_tests.rs | 205 ++++++++++++++++++ .../src/grpc/codex.code_mode.v1.proto | 13 ++ codex-rs/code-mode-protocol/src/lib.rs | 3 + .../code-mode-protocol/src/process_metrics.rs | 21 ++ codex-rs/code-mode-protocol/src/session.rs | 10 + codex-rs/code-mode/src/grpc_session/mod.rs | 16 +- .../src/grpc_session/process_metrics.rs | 50 +++++ codex-rs/core/src/code_mode_host.rs | 32 ++- .../src/code_mode_host_lifecycle_tests.rs | 34 +++ codex-rs/core/src/code_mode_host_tests.rs | 41 +++- codex-rs/core/src/thread_manager.rs | 12 + 18 files changed, 559 insertions(+), 15 deletions(-) create mode 100644 codex-rs/code-mode-host/src/process_metrics.rs create mode 100644 codex-rs/code-mode-host/tests/grpc/process_metrics_tests.rs create mode 100644 codex-rs/code-mode-protocol/src/process_metrics.rs create mode 100644 codex-rs/code-mode/src/grpc_session/process_metrics.rs diff --git a/codex-rs/code-mode-host/src/grpc/mod.rs b/codex-rs/code-mode-host/src/grpc/mod.rs index 13fb48d3c5..3901da32c0 100644 --- a/codex-rs/code-mode-host/src/grpc/mod.rs +++ b/codex-rs/code-mode-host/src/grpc/mod.rs @@ -47,6 +47,7 @@ fn trace_context_from_request(request: &Request) -> Option, + process_metrics: Option>, } 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, + ) -> 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; type SubscribeToToolCallsStream = GrpcStream; type ExecuteStream = GrpcStream; diff --git a/codex-rs/code-mode-host/src/grpc_stdio.rs b/codex-rs/code-mode-host/src/grpc_stdio.rs index c05bf8e135..b51b9fd28d 100644 --- a/codex-rs/code-mode-host/src/grpc_stdio.rs +++ b/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), ) diff --git a/codex-rs/code-mode-host/src/grpc_transport.rs b/codex-rs/code-mode-host/src/grpc_transport.rs index 27e49456f0..da655f4cc0 100644 --- a/codex-rs/code-mode-host/src/grpc_transport.rs +++ b/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), diff --git a/codex-rs/code-mode-host/src/lib.rs b/codex-rs/code-mode-host/src/lib.rs index 77c80e80b2..b534418123 100644 --- a/codex-rs/code-mode-host/src/lib.rs +++ b/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. diff --git a/codex-rs/code-mode-host/src/process_metrics.rs b/codex-rs/code-mode-host/src/process_metrics.rs new file mode 100644 index 0000000000..422317f5d0 --- /dev/null +++ b/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 Result + Send + Sync>; + +pub(crate) struct ProcessMetrics { + collector: CodeModeProcessMetricsCollector, + permits: Arc, +} + +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, 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, + })) + } +} diff --git a/codex-rs/code-mode-host/src/transport.rs b/codex-rs/code-mode-host/src/transport.rs index 095d0d78e0..caea7c5b50 100644 --- a/codex-rs/code-mode-host/src/transport.rs +++ b/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 + } } } diff --git a/codex-rs/code-mode-host/tests/grpc.rs b/codex-rs/code-mode-host/tests/grpc.rs index a78651d59c..69aef1441f 100644 --- a/codex-rs/code-mode-host/tests/grpc.rs +++ b/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"] diff --git a/codex-rs/code-mode-host/tests/grpc/process_metrics_tests.rs b/codex-rs/code-mode-host/tests/grpc/process_metrics_tests.rs new file mode 100644 index 0000000000..1ee27d17c8 --- /dev/null +++ b/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)); + } +} diff --git a/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto b/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto index 31db1f631b..043b7f8cd1 100644 --- a/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto +++ b/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. diff --git a/codex-rs/code-mode-protocol/src/lib.rs b/codex-rs/code-mode-protocol/src/lib.rs index 6ab492f39c..0d60218234 100644 --- a/codex-rs/code-mode-protocol/src/lib.rs +++ b/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; diff --git a/codex-rs/code-mode-protocol/src/process_metrics.rs b/codex-rs/code-mode-protocol/src/process_metrics.rs new file mode 100644 index 0000000000..9d82b476d1 --- /dev/null +++ b/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), +} diff --git a/codex-rs/code-mode-protocol/src/session.rs b/codex-rs/code-mode-protocol/src/session.rs index 38655bc41e..e5e72c571c 100644 --- a/codex-rs/code-mode-protocol/src/session.rs +++ b/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. diff --git a/codex-rs/code-mode/src/grpc_session/mod.rs b/codex-rs/code-mode/src/grpc_session/mod.rs index a2c09487e6..602f2fe2d0 100644 --- a/codex-rs/code-mode/src/grpc_session/mod.rs +++ b/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(request: &mut tonic::Request, 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, @@ -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(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()) } diff --git a/codex-rs/code-mode/src/grpc_session/process_metrics.rs b/codex-rs/code-mode/src/grpc_session/process_metrics.rs new file mode 100644 index 0000000000..5c381520cc --- /dev/null +++ b/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 { + 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())? +} diff --git a/codex-rs/core/src/code_mode_host.rs b/codex-rs/core/src/code_mode_host.rs index 83c39a9a24..ffad6e2da3 100644 --- a/codex-rs/core/src/code_mode_host.rs +++ b/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>>, @@ -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)) diff --git a/codex-rs/core/src/code_mode_host_lifecycle_tests.rs b/codex-rs/core/src/code_mode_host_lifecycle_tests.rs index fff3e7ebc0..cc50cac57e 100644 --- a/codex-rs/core/src/code_mode_host_lifecycle_tests.rs +++ b/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)?; diff --git a/codex-rs/core/src/code_mode_host_tests.rs b/codex-rs/core/src/code_mode_host_tests.rs index 2088380bb7..db60c29bb5 100644 --- a/codex-rs/core/src/code_mode_host_tests.rs +++ b/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(); diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index cda2cf6f62..b56485d741 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/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 { + 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,