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