From bb7c083fff7cafa2f2c7a99dce7b0cc7a7c73f45 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Thu, 2 Jul 2026 11:28:09 -0700 Subject: [PATCH] telemetry: log structured tool and inference timing events --- .../app-server/tests/common/json_logging.rs | 106 ++++++++++-- .../tests/common/test_app_server.rs | 57 ++++++- codex-rs/app-server/tests/suite/logging.rs | 145 +++++++++++++++++ codex-rs/core/src/tools/parallel.rs | 154 +++++++++++++++++- codex-rs/core/tests/suite/otel.rs | 2 +- 5 files changed, 442 insertions(+), 22 deletions(-) diff --git a/codex-rs/app-server/tests/common/json_logging.rs b/codex-rs/app-server/tests/common/json_logging.rs index c6c78bd0a4..8683d1bdaf 100644 --- a/codex-rs/app-server/tests/common/json_logging.rs +++ b/codex-rs/app-server/tests/common/json_logging.rs @@ -1,11 +1,79 @@ use std::path::Path; use std::process::Command; use std::process::Stdio; +use std::sync::Arc; +use std::sync::Mutex; +use std::time::Duration; use anyhow::Context; use anyhow::Result; use serde_json::Value; use serde_json::json; +use tokio::sync::Notify; + +#[derive(Clone, Default)] +pub(crate) struct JsonLogCapture { + lines: Arc>>, + updated: Arc, +} + +impl JsonLogCapture { + pub(crate) fn record(&self, line: String) { + self.lines + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(line); + self.updated.notify_one(); + } + + pub(crate) async fn wait_for_event(&self, event_name: &str) -> Result { + let mut events = self.wait_for_events(event_name, /*count*/ 1).await?; + Ok(events.remove(0)) + } + + pub(crate) async fn wait_for_events( + &self, + event_name: &str, + count: usize, + ) -> Result> { + let result = tokio::time::timeout(Duration::from_secs(10), async { + loop { + let updated = self.updated.notified(); + let events = self + .events()? + .into_iter() + .filter(|event| event["fields"]["event.name"].as_str() == Some(event_name)) + .collect::>(); + if events.len() >= count { + return Ok(events); + } + updated.await; + } + }) + .await; + match result { + Ok(result) => result, + Err(_) => { + let lines = self + .lines + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .join("\n"); + anyhow::bail!( + "timed out waiting for {count} JSON log event(s) named `{event_name}`; captured stderr:\n{lines}" + ) + } + } + } + + pub(crate) fn events(&self) -> Result> { + let lines = self + .lines + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + json_log_events(lines.iter().map(String::as_str)) + } +} pub fn app_server_json_shutdown_event( binary: &str, @@ -31,25 +99,39 @@ pub fn app_server_json_shutdown_event( let stderr = String::from_utf8(output.stderr)?; anyhow::ensure!(output.status.success(), "app-server failed: {stderr}"); - let events = stderr - .lines() - .filter(|line| !line.is_empty()) - .map(serde_json::from_str::) - .collect::>>() - .with_context(|| format!("app-server stderr was not JSONL: {stderr}"))?; + let events = json_log_events(stderr.lines()) + .with_context(|| format!("app-server stderr was not valid JSONL: {stderr}"))?; let event = events .iter() .find(|event| event["fields"]["message"] == "processor task exited") .context("missing INFO shutdown event in app-server JSON logs")?; - let timestamp = event["timestamp"] - .as_str() - .context("shutdown event did not include a timestamp")?; - chrono::DateTime::parse_from_rfc3339(timestamp) - .with_context(|| format!("shutdown event timestamp was not RFC 3339: {timestamp}"))?; - Ok(json!({ "level": event["level"], "fields": event["fields"], "target": event["target"], })) } + +fn json_log_events<'a>(lines: impl IntoIterator) -> Result> { + lines + .into_iter() + .filter(|line| !line.is_empty()) + .map(|line| { + let event = serde_json::from_str::(line) + .with_context(|| format!("log line was not JSON: {line}"))?; + anyhow::ensure!( + event["level"].is_string() + && event["fields"].is_object() + && event["target"].is_string(), + "JSON log event did not include level, fields, and target: {line}" + ); + let timestamp = event["timestamp"] + .as_str() + .with_context(|| format!("JSON log event did not include a timestamp: {line}"))?; + chrono::DateTime::parse_from_rfc3339(timestamp).with_context(|| { + format!("JSON log event timestamp was not RFC 3339: {timestamp}") + })?; + Ok(event) + }) + .collect() +} diff --git a/codex-rs/app-server/tests/common/test_app_server.rs b/codex-rs/app-server/tests/common/test_app_server.rs index fbe43e7689..cbb8ff161e 100644 --- a/codex-rs/app-server/tests/common/test_app_server.rs +++ b/codex-rs/app-server/tests/common/test_app_server.rs @@ -123,6 +123,8 @@ use core_test_support::test_codex::TestEnv; use core_test_support::test_codex::test_env; use tokio::process::Command; +use crate::json_logging::JsonLogCapture; + pub struct TestAppServer { next_request_id: AtomicI64, /// Retain this child process until the client is dropped. The Tokio runtime @@ -134,6 +136,7 @@ pub struct TestAppServer { stdout: BufReader, pending_messages: VecDeque, auto_env: Option, + json_logs: JsonLogCapture, } pub const DEFAULT_CLIENT_NAME: &str = "codex-app-server-tests"; @@ -160,6 +163,31 @@ impl TestAppServer { /// URL-based configuration, this helper rejects a `codex_home` containing /// that file. pub async fn new_with_auto_env(codex_home: &Path) -> anyhow::Result { + Self::new_with_auto_env_and_env(codex_home, &[]).await + } + + /// Starts an auto-environment app server that emits JSON logs. + /// + /// `rust_log` is the value to use for the `RUST_LOG` environment variable. + pub async fn new_with_auto_env_and_json_logging( + codex_home: &Path, + rust_log: impl Into, + ) -> anyhow::Result { + let rust_log = rust_log.into(); + Self::new_with_auto_env_and_env( + codex_home, + &[ + ("LOG_FORMAT", Some("json")), + ("RUST_LOG", Some(rust_log.as_str())), + ], + ) + .await + } + + async fn new_with_auto_env_and_env( + codex_home: &Path, + extra_env_overrides: &[(&str, Option<&str>)], + ) -> anyhow::Result { let environments_toml = codex_home.join("environments.toml"); ensure!( !environments_toml @@ -172,7 +200,7 @@ impl TestAppServer { let auto_env = test_env().await?; // Noise registry configuration takes precedence over the URL-based // provider, so clear inherited values to keep the selection hermetic. - let env_overrides = [ + let mut env_overrides = vec![ ( CODEX_EXEC_SERVER_URL_ENV_VAR, auto_env.environment().exec_server_url(), @@ -182,6 +210,7 @@ impl TestAppServer { (CODEX_EXEC_SERVER_NOISE_AUTH_TOKEN_ENV_VAR, None), (CODEX_EXEC_SERVER_NOISE_CHATGPT_ACCOUNT_ID_ENV_VAR, None), ]; + env_overrides.extend_from_slice(extra_env_overrides); let mut app_server = Self::new_with_env(codex_home, &env_overrides).await?; app_server.auto_env = Some(auto_env); Ok(app_server) @@ -202,6 +231,28 @@ impl TestAppServer { }) } + /// Waits for a JSON stderr event whose structured `event.name` field matches. + pub async fn wait_for_json_log_event( + &self, + event_name: &str, + ) -> anyhow::Result { + self.json_logs.wait_for_event(event_name).await + } + + /// Waits for the requested number of JSON stderr events with the same `event.name` field. + pub async fn wait_for_json_log_events( + &self, + event_name: &str, + count: usize, + ) -> anyhow::Result> { + self.json_logs.wait_for_events(event_name, count).await + } + + /// Returns every stderr line parsed and validated as a JSON log event. + pub fn json_log_events(&self) -> anyhow::Result> { + self.json_logs.events() + } + pub async fn new_without_managed_config(codex_home: &Path) -> anyhow::Result { Self::new_with_env(codex_home, &[(DISABLE_MANAGED_CONFIG_ENV_VAR, Some("1"))]).await } @@ -322,10 +373,13 @@ impl TestAppServer { // Forward child's stderr to our stderr so failures are visible even // when stdout/stderr are captured by the test harness. + let json_logs = JsonLogCapture::default(); if let Some(stderr) = process.stderr.take() { + let json_logs = json_logs.clone(); let mut stderr_reader = BufReader::new(stderr).lines(); tokio::spawn(async move { while let Ok(Some(line)) = stderr_reader.next_line().await { + json_logs.record(line.clone()); eprintln!("[mcp stderr] {line}"); } }); @@ -337,6 +391,7 @@ impl TestAppServer { stdout, pending_messages: VecDeque::new(), auto_env: None, + json_logs, }) } diff --git a/codex-rs/app-server/tests/suite/logging.rs b/codex-rs/app-server/tests/suite/logging.rs index ea2c31a766..8d39a485d1 100644 --- a/codex-rs/app-server/tests/suite/logging.rs +++ b/codex-rs/app-server/tests/suite/logging.rs @@ -1,8 +1,30 @@ +use anyhow::Context; use anyhow::Result; +use app_test_support::TestAppServer; use app_test_support::app_server_json_shutdown_event; +use app_test_support::create_exec_command_sse_response; +use app_test_support::create_final_assistant_message_sse_response; +use app_test_support::create_mock_responses_server_sequence; +use app_test_support::to_response; +use app_test_support::write_mock_responses_config_toml; +use codex_app_server_protocol::JSONRPCResponse; +use codex_app_server_protocol::RequestId; +use codex_app_server_protocol::ThreadStartParams; +use codex_app_server_protocol::ThreadStartResponse; +use codex_app_server_protocol::TurnStartParams; +use codex_app_server_protocol::TurnStartResponse; +use codex_app_server_protocol::UserInput; +use codex_features::Feature; +use core_test_support::skip_if_no_network; use pretty_assertions::assert_eq; +use serde_json::Value; use serde_json::json; +use std::collections::BTreeMap; use tempfile::TempDir; +use tokio::time::Duration; +use tokio::time::timeout; + +const READ_TIMEOUT: Duration = Duration::from_secs(10); #[test] fn standalone_app_server_emits_json_info_events() -> Result<()> { @@ -25,3 +47,126 @@ fn standalone_app_server_emits_json_info_events() -> Result<()> { Ok(()) } + +#[tokio::test] +async fn app_server_emits_structured_tool_call_timing_event() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = create_mock_responses_server_sequence(vec![ + create_exec_command_sse_response("exec-call-1")?, + create_final_assistant_message_sse_response("done")?, + ]) + .await; + let codex_home = TempDir::new()?; + write_mock_responses_config_toml( + codex_home.path(), + &server.uri(), + &BTreeMap::from([(Feature::UnifiedExec, true)]), + /*auto_compact_limit*/ 100_000, + /*requires_openai_auth*/ None, + "mock_provider", + "compact", + )?; + + let mut app_server = TestAppServer::new_with_auto_env_and_json_logging( + codex_home.path(), + "warn,codex_core::tools::parallel=info", + ) + .await?; + timeout(READ_TIMEOUT, app_server.initialize()).await??; + + let thread_start_id = app_server + .send_thread_start_request_with_auto_env(ThreadStartParams { + model: Some("mock-model".to_string()), + ..Default::default() + }) + .await?; + let thread_start_response: JSONRPCResponse = timeout( + READ_TIMEOUT, + app_server.read_stream_until_response_message(RequestId::Integer(thread_start_id)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response(thread_start_response)?; + + let turn_start_id = app_server + .send_turn_start_request(TurnStartParams { + thread_id: thread.id.clone(), + input: vec![UserInput::Text { + text: "run a command".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }) + .await?; + let turn_start_response: JSONRPCResponse = timeout( + READ_TIMEOUT, + app_server.read_stream_until_response_message(RequestId::Integer(turn_start_id)), + ) + .await??; + let TurnStartResponse { turn } = to_response(turn_start_response)?; + + timeout( + READ_TIMEOUT, + app_server.read_stream_until_notification_message("turn/completed"), + ) + .await??; + + let mut tool_call = app_server + .wait_for_json_log_event("codex.tool_call") + .await?; + let tool_call_object = tool_call + .as_object_mut() + .context("tool call log event must be an object")?; + // JsonLogCapture already validates the timestamp as RFC 3339. + tool_call_object + .remove("timestamp") + .context("tool call log event must include a timestamp")?; + let fields = tool_call_object + .get_mut("fields") + .and_then(Value::as_object_mut) + .context("tool call log event fields must be an object")?; + let trace_id = fields + .remove("trace_id") + .context("tool call log event must include trace_id")?; + anyhow::ensure!(trace_id.is_string(), "trace_id must be a string"); + let dispatch_duration_ms = fields + .remove("dispatch_duration_ms") + .and_then(|duration| duration.as_u64()) + .context("dispatch_duration_ms must be a nonnegative integer")?; + let handler_duration_ms = fields + .remove("handler_duration_ms") + .and_then(|duration| duration.as_u64()) + .context("handler_duration_ms must be a nonnegative integer")?; + let total_duration_ms = fields + .remove("total_duration_ms") + .and_then(|duration| duration.as_u64()) + .context("total_duration_ms must be a nonnegative integer")?; + let accounted_duration_ms = dispatch_duration_ms + .checked_add(handler_duration_ms) + .context("dispatch and handler durations must not overflow")?; + anyhow::ensure!( + total_duration_ms >= accounted_duration_ms + && total_duration_ms - accounted_duration_ms <= 1, + "dispatch and handler durations must account for total duration within integer truncation" + ); + + assert_eq!( + tool_call, + json!({ + "level": "INFO", + "fields": { + "message": "tool call completed", + "event.name": "codex.tool_call", + "conversation.id": thread.id, + "turn_id": turn.id, + "tool_name": "exec_command", + "call_id": "exec-call-1", + "tool_source": "direct", + "execution_started": true, + }, + "target": "codex_core::tools::parallel", + }) + ); + + Ok(()) +} diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index a2ee89ab83..61697899a5 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -1,4 +1,5 @@ use std::sync::Arc; +use std::sync::OnceLock; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use std::time::Instant; @@ -9,6 +10,7 @@ use tokio_util::either::Either; use tokio_util::sync::CancellationToken; use tokio_util::task::AbortOnDropHandle; use tracing::Instrument; +use tracing::info; use tracing::instrument; use tracing::trace_span; @@ -27,6 +29,15 @@ use crate::tools::router::ToolRouter; use codex_protocol::error::CodexErr; use codex_protocol::models::ResponseInputItem; +struct ToolCallTimingGuard { + started_at: Instant, + execution_started_at: Arc>, + conversation_id: String, + turn_id: String, + call_id: String, + tool_name: codex_tools::ToolName, +} + #[derive(Clone)] pub(crate) struct ToolCallRuntime { router: Arc, @@ -96,6 +107,11 @@ impl ToolCallRuntime { let invocation_cancellation_token = cancellation_token.clone(); let wait_for_runtime_cancellation = self.router.tool_waits_for_runtime_cancellation(&call); let started = Instant::now(); + let tool_call_timing_guard = + ToolCallTimingGuard::capture(started, &session.thread_id, &turn.sub_id, &call, &source); + let execution_started_at = tool_call_timing_guard + .as_ref() + .map(|timing| Arc::clone(&timing.execution_started_at)); let abort_session = Arc::clone(&session); let abort_source = source.clone(); let abort_turn = Arc::clone(&turn); @@ -112,13 +128,18 @@ impl ToolCallRuntime { ); let abort_dispatch_span = dispatch_span.clone(); - let mut handle: AbortOnDropHandle> = + let mut dispatch_handle: AbortOnDropHandle> = AbortOnDropHandle::new(tokio::spawn(async move { let _guard = if supports_parallel { Either::Left(lock.read().await) } else { Either::Right(lock.write().await) }; + // Admission through the parallel-execution gate marks the end + // of dispatch waiting and the start of handler execution. + if let Some(execution_started_at) = execution_started_at { + let _ = execution_started_at.set(Instant::now()); + } router .dispatch_tool_call_with_terminal_outcome( @@ -135,28 +156,29 @@ impl ToolCallRuntime { })); async move { + let _tool_call_timing_guard = tool_call_timing_guard; tokio::select! { - res = &mut handle => res.map_err(Self::tool_task_join_error)?, + res = &mut dispatch_handle => res.map_err(Self::tool_task_join_error)?, _ = cancellation_token.cancelled() => { - if terminal_outcome_reached.load(Ordering::Acquire) || handle.is_finished() { - handle.await.map_err(Self::tool_task_join_error)? + if terminal_outcome_reached.load(Ordering::Acquire) || dispatch_handle.is_finished() { + dispatch_handle.await.map_err(Self::tool_task_join_error)? } else { let secs = started.elapsed().as_secs_f32().max(0.1); abort_dispatch_span.record("aborted", true); if wait_for_runtime_cancellation { if terminal_outcome_reached.swap(true, Ordering::AcqRel) { - return handle.await.map_err(Self::tool_task_join_error)?; + return dispatch_handle.await.map_err(Self::tool_task_join_error)?; } // The abort owns the terminal outcome; await only so // the runtime can finish process teardown. - match handle.await { + match dispatch_handle.await { Ok(_) => {} Err(err) if err.is_cancelled() => {} Err(err) => return Err(Self::tool_task_join_error(err)), } } else { - handle.abort(); - match handle.await { + dispatch_handle.abort(); + match dispatch_handle.await { Ok(result) => return result, Err(err) if err.is_cancelled() => {} Err(err) => return Err(Self::tool_task_join_error(err)), @@ -237,6 +259,85 @@ impl ToolCallRuntime { } } +impl ToolCallTimingGuard { + fn capture( + started_at: Instant, + conversation_id: &impl std::fmt::Display, + turn_id: &str, + call: &ToolCall, + source: &ToolCallSource, + ) -> Option { + if !matches!(source, ToolCallSource::Direct) || !tracing::enabled!(tracing::Level::INFO) { + return None; + } + + Some(Self { + started_at, + execution_started_at: Arc::new(OnceLock::new()), + conversation_id: conversation_id.to_string(), + turn_id: turn_id.to_string(), + call_id: call.call_id.clone(), + tool_name: call.tool_name.clone(), + }) + } +} + +impl Drop for ToolCallTimingGuard { + fn drop(&mut self) { + let completed_at = Instant::now(); + // Snapshot once so a concurrently-starting dispatch cannot make one + // event internally inconsistent. + let execution_started_at = self + .execution_started_at + .get() + .copied() + .filter(|execution_started_at| *execution_started_at <= completed_at); + let duration_ms = |duration: std::time::Duration| u64::try_from(duration.as_millis()).ok(); + let total_duration_ms = duration_ms(completed_at.duration_since(self.started_at)); + let dispatch_duration_ms = execution_started_at.map_or_else( + || total_duration_ms, + |execution_started_at| { + duration_ms(execution_started_at.duration_since(self.started_at)) + }, + ); + let handler_duration_ms = execution_started_at.map_or(Some(0), |execution_started_at| { + duration_ms(completed_at.duration_since(execution_started_at)) + }); + + macro_rules! log_tool_call { + ($dispatch_duration_ms:expr, $handler_duration_ms:expr, $total_duration_ms:expr) => { + info!( + event.name = "codex.tool_call", + trace_id = %codex_otel::current_span_trace_id().unwrap_or_default(), + conversation.id = %self.conversation_id, + turn_id = %self.turn_id, + tool_name = %self.tool_name, + call_id = %self.call_id, + tool_source = "direct", + execution_started = execution_started_at.is_some(), + dispatch_duration_ms = $dispatch_duration_ms, + handler_duration_ms = $handler_duration_ms, + total_duration_ms = $total_duration_ms, + "tool call completed" + ); + }; + } + + match (dispatch_duration_ms, handler_duration_ms, total_duration_ms) { + (Some(dispatch_duration_ms), Some(handler_duration_ms), Some(total_duration_ms)) => { + log_tool_call!(dispatch_duration_ms, handler_duration_ms, total_duration_ms); + } + _ => { + log_tool_call!( + tracing::field::Empty, + tracing::field::Empty, + tracing::field::Empty + ); + } + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -256,6 +357,43 @@ mod tests { use tokio::sync::Notify; use tokio::sync::oneshot; + #[test] + fn tool_call_timing_guard_ignores_code_mode_source() { + let subscriber = tracing_subscriber::fmt() + .with_max_level(tracing::Level::INFO) + .finish(); + tracing::subscriber::with_default(subscriber, || { + let call = ToolCall { + tool_name: codex_tools::ToolName::plain("test_tool"), + call_id: "call-1".to_string(), + payload: ToolPayload::Function { + arguments: "{}".to_string(), + }, + }; + let direct_guard = ToolCallTimingGuard::capture( + Instant::now(), + &"conversation-id", + "turn-id", + &call, + &ToolCallSource::Direct, + ); + assert!(direct_guard.is_some()); + drop(direct_guard); + + let code_mode_guard = ToolCallTimingGuard::capture( + Instant::now(), + &"conversation-id", + "turn-id", + &call, + &ToolCallSource::CodeMode { + cell_id: "cell-1".to_string(), + runtime_tool_call_id: "runtime-call-1".to_string(), + }, + ); + assert!(code_mode_guard.is_none()); + }); + } + struct ImmediateHandler { tool_name: codex_tools::ToolName, } diff --git a/codex-rs/core/tests/suite/otel.rs b/codex-rs/core/tests/suite/otel.rs index 4a2b2993ad..e8a4af523d 100644 --- a/codex-rs/core/tests/suite/otel.rs +++ b/codex-rs/core/tests/suite/otel.rs @@ -716,7 +716,7 @@ async fn turn_and_completed_response_spans_record_token_usage() { ); } -#[tokio::test] +#[tokio::test(flavor = "current_thread")] async fn handle_responses_span_records_response_kind_and_tool_name() { let buffer: &'static Mutex> = Box::leak(Box::new(Mutex::new(Vec::new()))); let subscriber = tracing_subscriber::fmt()