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
This commit is contained in:
aakash-oai
2026-10-28 21:34:25 +00:00
committed by copyberry
parent 7ce5fa7bdc
commit d373d64967
18 changed files with 559 additions and 15 deletions
+29
View File
@@ -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>;
+2 -2
View File
@@ -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),
+17 -1
View File
@@ -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,
}))
}
}
+5 -3
View File
@@ -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
}
}
}
+2
View File
@@ -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));
}
}
@@ -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.
+3
View File
@@ -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;
@@ -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),
}
@@ -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.
+15 -1
View File
@@ -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())
}
@@ -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())?
}
+27 -5
View File
@@ -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))
@@ -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)?;
+40 -1
View File
@@ -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();
+12
View File
@@ -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,