diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index b9ab8a1d90..dfa360b42a 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2954,6 +2954,7 @@ dependencies = [ "codex-context-fragments", "codex-core-plugins", "codex-diagnostics", + "codex-direct-tool-timing", "codex-exec-server", "codex-exec-server-test-support", "codex-execpolicy", @@ -3003,6 +3004,7 @@ dependencies = [ "codex-utils-audio", "codex-utils-cache", "codex-utils-cargo-bin", + "codex-utils-critical-path", "codex-utils-git-discovery", "codex-utils-head-tail-buffer", "codex-utils-home-dir", @@ -3178,6 +3180,9 @@ version = "0.0.0" dependencies = [ "codex-utils-critical-path", "pretty_assertions", + "test-case", + "tokio", + "tokio-util", ] [[package]] diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 1602f7625c..134c3f920b 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -225,6 +225,7 @@ codex-core-api = { path = "core-api" } codex-core-plugin-common = { path = "core-plugin-common" } codex-core-plugins = { path = "core-plugins" } codex-diagnostics = { path = "diagnostics" } +codex-direct-tool-timing = { path = "direct-tool-timing" } codex-exec = { path = "exec" } codex-file-system = { path = "file-system" } codex-exec-server-protocol = { path = "exec-server-protocol" } diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 9bb640698e..542abf90d1 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -34,6 +34,7 @@ codex-context-fragments = { workspace = true } codex-config = { workspace = true } codex-core-plugins = { workspace = true } codex-diagnostics = { workspace = true } +codex-direct-tool-timing = { workspace = true } codex-exec-server = { workspace = true } codex-extension-api = { workspace = true } codex-extension-items = { workspace = true } @@ -78,6 +79,7 @@ codex-tools = { workspace = true } codex-utils-absolute-path = { workspace = true } codex-utils-audio = { workspace = true } codex-utils-cache = { workspace = true } +codex-utils-critical-path = { workspace = true } codex-utils-git-discovery = { workspace = true } codex-utils-head-tail-buffer = { workspace = true } codex-utils-image = { workspace = true } diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index bbfed66c0d..0f83fd98aa 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -978,6 +978,9 @@ "credential_masking": { "type": "boolean" }, + "critical_path_timing": { + "type": "boolean" + }, "current_time_reminder": { "$ref": "#/definitions/FeatureToml_for_CurrentTimeReminderConfigToml" }, @@ -7580,6 +7583,9 @@ "credential_masking": { "type": "boolean" }, + "critical_path_timing": { + "type": "boolean" + }, "current_time_reminder": { "$ref": "#/definitions/FeatureToml_for_CurrentTimeReminderConfigToml" }, diff --git a/codex-rs/core/src/guardian/decision.rs b/codex-rs/core/src/guardian/decision.rs index 79981f57e0..3e35b4ad59 100644 --- a/codex-rs/core/src/guardian/decision.rs +++ b/codex-rs/core/src/guardian/decision.rs @@ -26,13 +26,12 @@ pub(crate) fn spawn_approval_decision( let request: ReviewAction = request.into(); let (tx, rx) = oneshot::channel(); let runtime = session.services.runtime_handle.clone(); + let decision = decide_approval(session, context, review_id, request, reasons, options); let spawn_result = std::thread::Builder::new() .name("codex-approval-review".to_string()) .stack_size(THREAD_STACK_SIZE_BYTES) .spawn(move || { - let decision = runtime.block_on(decide_approval( - session, context, review_id, request, reasons, options, - )); + let decision = runtime.block_on(decision); let _ = tx.send(decision); }); if let Err(err) = spawn_result { @@ -43,80 +42,92 @@ pub(crate) fn spawn_approval_decision( /// `None` requests the existing user flow. No contributor is never an implicit allow. /// Both metric sinks use the originating turn retained by the review context. -pub(crate) async fn decide_approval( +pub(crate) fn decide_approval( session: Arc, context: impl Into, review_id: String, request: impl Into, reasons: ApprovalRequestReasons, mut options: GuardianReviewOptions, -) -> Option { - let context = context.into(); - let request = request.into(); - let (_, history_reset) = session.history_reset().await; - let cancellation = match options.external_cancel.take() { - Some(external) => crate::exec::cancel_when_either(external, history_reset.clone()), - None => history_reset.child_token(), - }; - let _cancel_on_drop = cancellation.clone().drop_guard(); - let turn = context.turn(); - let live_config = session.get_config().await; - let requirements = live_config.config_layer_stack.requirements(); - let model_requires_review = - requirements.auto_review_required_for_model(&context.model_info.slug); - let require_guardian = options.require_guardian - || model_requires_review - || requirements - .approvals_reviewer - .can_set(&codex_protocol::config_types::ApprovalsReviewer::User) - .is_err(); - let full_access = context.environments().has_full_access( - context.approval_policy, - &turn.config.permissions.effective_permission_profile(), - ); - let require_synchronous_review = options.require_synchronous_review; - let retried = reasons.retry.is_some(); - let decision = codex_guardian_reviewer::ReviewRequest { - host: ReviewRuntime { - session: Arc::clone(&session), - history_reset: history_reset.clone(), - context: context.clone(), - request: request.clone(), - reasons, - options, - }, - approval_id: &review_id, - tool_call_id: request.tool_call_id.as_deref(), - action: request.action.as_ref().ok(), - thread_id: session.thread_id, - thread_store: &session.services.thread_extension_data, - category: request.category, - approval_policy: context.approval_policy, - approvals_reviewer: context.approvals_reviewer, - require_guardian, - require_synchronous_review, - model_requires_review, - async_enabled: turn.config.features.enabled(codex_features::Feature::GuardianV2), - retried, - escalated_exec: matches!(&request.request, Ok(GuardianApprovalRequest::ExecCommand { sandbox_permissions, .. }) if sandbox_permissions.requires_escalated_permissions()), - full_access, - cancellation: cancellation.clone(), - timeout: turn.config.guardian_review_timeout, - // Keep the existing turn-level reporting policy during this ownership move. - model: turn.model_info(), - telemetry: &turn.session_telemetry, - log_assessments: turn.config.otel.guardian_assessment_logging_enabled(), - analytics: &session.services.analytics_events_client, - metrics: Some(crate::session::extension_metrics::from_session_telemetry( - turn.session_telemetry.clone(), - )), - } - .decide(&session.services.extensions) - .await; - // Enforce cancellation after extension callbacks too, including cached decisions. - if history_reset.is_cancelled() || cancellation.is_cancelled() { - Some(ReviewDecision::Abort) - } else { - decision +) -> futures::future::BoxFuture<'static, Option> { + let context: GuardianReviewContext = context.into(); + let request: ReviewAction = request.into(); + // Invalidate before scheduling: cancellation may finalize the turn before + // the review future is first polled. Replace this boundary with attribution. + if context + .turn() + .config + .features + .enabled(codex_features::Feature::CriticalPathTiming) + { + context.turn().turn_timing_state.invalidate_critical_path(); } + Box::pin(async move { + let (_, history_reset) = session.history_reset().await; + let cancellation = match options.external_cancel.take() { + Some(external) => crate::exec::cancel_when_either(external, history_reset.clone()), + None => history_reset.child_token(), + }; + let _cancel_on_drop = cancellation.clone().drop_guard(); + let turn = context.turn(); + let live_config = session.get_config().await; + let requirements = live_config.config_layer_stack.requirements(); + let model_requires_review = + requirements.auto_review_required_for_model(&context.model_info.slug); + let require_guardian = options.require_guardian + || model_requires_review + || requirements + .approvals_reviewer + .can_set(&codex_protocol::config_types::ApprovalsReviewer::User) + .is_err(); + let full_access = context.environments().has_full_access( + context.approval_policy, + &turn.config.permissions.effective_permission_profile(), + ); + let require_synchronous_review = options.require_synchronous_review; + let retried = reasons.retry.is_some(); + let decision = codex_guardian_reviewer::ReviewRequest { + host: ReviewRuntime { + session: Arc::clone(&session), + history_reset: history_reset.clone(), + context: context.clone(), + request: request.clone(), + reasons, + options, + }, + approval_id: &review_id, + tool_call_id: request.tool_call_id.as_deref(), + action: request.action.as_ref().ok(), + thread_id: session.thread_id, + thread_store: &session.services.thread_extension_data, + category: request.category, + approval_policy: context.approval_policy, + approvals_reviewer: context.approvals_reviewer, + require_guardian, + require_synchronous_review, + model_requires_review, + async_enabled: turn.config.features.enabled(codex_features::Feature::GuardianV2), + retried, + escalated_exec: matches!(&request.request, Ok(GuardianApprovalRequest::ExecCommand { sandbox_permissions, .. }) if sandbox_permissions.requires_escalated_permissions()), + full_access, + cancellation: cancellation.clone(), + timeout: turn.config.guardian_review_timeout, + // Keep the existing turn-level reporting policy during this ownership move. + model: turn.model_info(), + telemetry: &turn.session_telemetry, + log_assessments: turn.config.otel.guardian_assessment_logging_enabled(), + analytics: &session.services.analytics_events_client, + metrics: Some(crate::session::extension_metrics::from_session_telemetry( + turn.session_telemetry.clone(), + )), + } + .decide(&session.services.extensions) + .await; + // Enforce cancellation after extension callbacks too, including cached decisions. + if history_reset.is_cancelled() || cancellation.is_cancelled() { + Some(ReviewDecision::Abort) + } else { + decision + } + }) } diff --git a/codex-rs/core/src/guardian/decision_timing_tests.rs b/codex-rs/core/src/guardian/decision_timing_tests.rs new file mode 100644 index 0000000000..712b595d8b --- /dev/null +++ b/codex-rs/core/src/guardian/decision_timing_tests.rs @@ -0,0 +1,63 @@ +#![cfg(test)] +//! Keeps uninstrumented review work out of complete critical-path captures. + +use pretty_assertions::assert_eq; +use std::sync::Arc; +use std::time::Instant; + +#[test_case::test_case(true; "capture_enabled")] +#[test_case::test_case(false; "capture_disabled")] +#[tokio::test] +async fn uncaptured_review_invalidates_before_polling(critical_path_timing: bool) { + let (session, mut turn) = crate::session::tests::make_session_and_context().await; + Arc::make_mut(&mut turn.config) + .features + .set_enabled( + codex_features::Feature::CriticalPathTiming, + critical_path_timing, + ) + .unwrap(); + let turn = Arc::new(turn); + turn.turn_timing_state.enable_critical_path(); + turn.turn_timing_state + .mark_turn_started(Instant::now()) + .await; + assert!( + turn.turn_timing_state + .snapshot_path(Instant::now()) + .is_some() + ); + let cancellation = tokio_util::sync::CancellationToken::new(); + cancellation.cancel(); + let decision = crate::guardian::decide_approval( + Arc::new(session), + &turn, + "review".to_string(), + crate::guardian::GuardianApprovalRequest::ApplyPatch { + id: "call".to_string(), + environment_id: "local".to_string(), + cwd: codex_utils_path_uri::PathUri::parse("file:///tmp").unwrap(), + files: Vec::new(), + patch: String::new(), + }, + crate::tools::sandboxing::ApprovalRequestReasons::default(), + crate::guardian::GuardianReviewOptions { + require_guardian: true, + plugin_attribution_override: None, + approval_request_source: codex_analytics::GuardianApprovalRequestSource::MainTurn, + external_cancel: Some(cancellation), + require_synchronous_review: false, + }, + ); + // No polling or worker scheduling is needed to protect a finalizing turn. + assert_eq!( + turn.turn_timing_state + .snapshot_path(Instant::now()) + .is_none(), + critical_path_timing + ); + assert_eq!( + decision.await, + Some(codex_protocol::protocol::ReviewDecision::Abort) + ); +} diff --git a/codex-rs/core/src/guardian/mod.rs b/codex-rs/core/src/guardian/mod.rs index 7602d6721e..652a43d958 100644 --- a/codex-rs/core/src/guardian/mod.rs +++ b/codex-rs/core/src/guardian/mod.rs @@ -224,3 +224,7 @@ use review_session::build_guardian_review_session_config as build_guardian_revie #[cfg(test)] mod tests; + +#[cfg(test)] +#[path = "decision_timing_tests.rs"] +mod decision_timing_tests; diff --git a/codex-rs/core/src/session/config_refresh.rs b/codex-rs/core/src/session/config_refresh.rs index 08a99f81dc..016eae3978 100644 --- a/codex-rs/core/src/session/config_refresh.rs +++ b/codex-rs/core/src/session/config_refresh.rs @@ -1,4 +1,4 @@ -//! Publishes resolved MCP configuration while preserving session-static settings. +//! Publishes resolved MCP configuration and host timing rollouts for future turns. use std::sync::Arc; @@ -107,6 +107,16 @@ impl Session { } } } + // Publish host rollout changes only to future turn snapshots. Active turns + // must retain one capture policy across tool dispatch and Guardian review. + if matches!(scope, RuntimeConfigRefresh::User) + && let Err(err) = config.features.set_enabled( + Feature::CriticalPathTiming, + next_config.features.enabled(Feature::CriticalPathTiming), + ) + { + warn!("failed to refresh critical path timing config: {err}"); + } let mut config = Arc::new(config); let notify_contributors = !matches!(scope, RuntimeConfigRefresh::Mcp) && !self.services.extensions.config_contributors().is_empty(); diff --git a/codex-rs/core/src/session/config_refresh_tests.rs b/codex-rs/core/src/session/config_refresh_tests.rs index 927a286309..b584558364 100644 --- a/codex-rs/core/src/session/config_refresh_tests.rs +++ b/codex-rs/core/src/session/config_refresh_tests.rs @@ -202,3 +202,94 @@ async fn hook_refresh_retries_current_plugin_policy_after_owner_replacement() -> assert!(session.hooks().preview_session_start(&request).is_empty()); Ok(()) } + +#[test_case::test_case(RuntimeConfigRefresh::User, false; "host enables")] +#[test_case::test_case(RuntimeConfigRefresh::User, true; "host disables")] +#[test_case::test_case(RuntimeConfigRefresh::Mcp, false; "MCP retains disabled")] +#[test_case::test_case(RuntimeConfigRefresh::Mcp, true; "MCP retains enabled")] +#[test_case::test_case(RuntimeConfigRefresh::UserFiles, false; "file retains disabled")] +#[test_case::test_case(RuntimeConfigRefresh::UserFiles, true; "file retains enabled")] +#[tokio::test] +async fn critical_path_refresh_changes_only_future_turns( + scope: RuntimeConfigRefresh, + initially_enabled: bool, +) { + let (session, _) = make_session_and_context().await; + let mut initial = session.get_config().await.as_ref().clone(); + initial + .features + .set_enabled(Feature::CriticalPathTiming, initially_enabled) + .unwrap(); + let original = Arc::new(initial); + session + .state + .lock() + .await + .session_configuration + .original_config_do_not_use = Arc::clone(&original); + let active_turn = session.new_default_turn().await; + let mut next_config = original.as_ref().clone(); + next_config + .features + .set_enabled(Feature::CriticalPathTiming, !initially_enabled) + .unwrap(); + + assert_eq!( + session.refresh_config(original, next_config, scope).await, + crate::ConfigRefreshOutcome::Published + ); + + let next_turn = session.new_default_turn().await; + assert_eq!( + active_turn + .config + .features + .enabled(Feature::CriticalPathTiming), + initially_enabled + ); + assert_eq!( + next_turn + .config + .features + .enabled(Feature::CriticalPathTiming), + initially_enabled != matches!(scope, RuntimeConfigRefresh::User) + ); +} + +#[tokio::test] +async fn critical_path_refresh_respects_managed_disable() { + let (session, _) = make_session_and_context().await; + let mut initial = session.get_config().await.as_ref().clone(); + initial.features = crate::config::ManagedFeatures::from_configured( + initial.features.get().clone(), + Some(Sourced::new( + serde_json::from_value(json!({"critical_path_timing": false})).unwrap(), + RequirementSource::LegacyManagedConfigTomlFromMdm, + )), + ) + .unwrap(); + let original = Arc::new(initial); + session + .state + .lock() + .await + .session_configuration + .original_config_do_not_use = Arc::clone(&original); + let mut next_config = original.as_ref().clone(); + let mut incoming_features = original.features.get().clone(); + incoming_features.enable(Feature::CriticalPathTiming); + next_config.features = incoming_features.into(); + + assert_eq!( + session.refresh_runtime_config(original, next_config).await, + crate::ConfigRefreshOutcome::Published + ); + assert!( + !session + .new_default_turn() + .await + .config + .features + .enabled(Feature::CriticalPathTiming) + ); +} diff --git a/codex-rs/core/src/session/elicitation_holders_tests.rs b/codex-rs/core/src/session/elicitation_holders_tests.rs index c1d8ab53f0..2d86b74c62 100644 --- a/codex-rs/core/src/session/elicitation_holders_tests.rs +++ b/codex-rs/core/src/session/elicitation_holders_tests.rs @@ -1,5 +1,7 @@ +use codex_features::Feature; use std::collections::HashMap; use std::sync::Arc; +use std::time::Instant; use codex_core_plugins::PluginCommandAttribution; use codex_plugin::PluginId; @@ -18,6 +20,27 @@ use super::tests::make_session_and_context_with_rx; use crate::session::step_context::StepContext; use crate::state::ActiveTurn; +async fn start_timing(turn: &mut Arc, enabled: bool) { + Arc::make_mut(&mut Arc::get_mut(turn).expect("unique test turn").config) + .features + .set_enabled(Feature::CriticalPathTiming, enabled) + .unwrap(); + // Populate state independently to verify disabled capture does not invalidate it. + turn.turn_timing_state.enable_critical_path(); + turn.turn_timing_state + .mark_turn_started(Instant::now()) + .await; +} + +fn assert_uncaptured_wait(turn: &super::turn_context::TurnContext, enabled: bool) { + assert_eq!( + turn.turn_timing_state + .snapshot_path(Instant::now()) + .is_none(), + enabled + ); +} + async fn wait_until_held(pause_state: &mut watch::Receiver) { pause_state .wait_for(|paused| *paused) @@ -32,9 +55,12 @@ async fn wait_until_released(pause_state: &mut watch::Receiver) { .expect("elicitation service should remain available"); } +#[test_case::test_case(true; "capture_on")] +#[test_case::test_case(false; "capture_off")] #[tokio::test] -async fn command_approval_holds_an_elicitation_until_response() { - let (session, turn_context, events) = make_session_and_context_with_rx().await; +async fn command_approval_holds_an_elicitation_until_response(capture: bool) { + let (session, mut turn_context, events) = make_session_and_context_with_rx().await; + start_timing(&mut turn_context, capture).await; *session.active_turn.lock().await = Some(ActiveTurn::default()); let mut pause_state = session.subscribe_elicitation_pause_state(); #[allow(deprecated)] @@ -82,11 +108,13 @@ async fn command_approval_holds_an_elicitation_until_response() { .await; request.await.expect("approval task"); wait_until_released(&mut pause_state).await; + assert_uncaptured_wait(&turn_context, capture); } #[tokio::test] async fn patch_approval_holds_an_elicitation_until_response() { - let (session, turn_context, events) = make_session_and_context_with_rx().await; + let (session, mut turn_context, events) = make_session_and_context_with_rx().await; + start_timing(&mut turn_context, /*enabled*/ true).await; *session.active_turn.lock().await = Some(ActiveTurn::default()); let mut pause_state = session.subscribe_elicitation_pause_state(); @@ -113,11 +141,13 @@ async fn patch_approval_holds_an_elicitation_until_response() { .await; request.await.expect("approval task"); wait_until_released(&mut pause_state).await; + assert_uncaptured_wait(&turn_context, /*enabled*/ true); } #[tokio::test] async fn permission_request_holds_an_elicitation_until_response() { - let (session, turn_context, events) = make_session_and_context_with_rx().await; + let (session, mut turn_context, events) = make_session_and_context_with_rx().await; + start_timing(&mut turn_context, /*enabled*/ true).await; *session.active_turn.lock().await = Some(ActiveTurn::default()); let mut pause_state = session.subscribe_elicitation_pause_state(); @@ -160,11 +190,16 @@ async fn permission_request_holds_an_elicitation_until_response() { .await; request.await.expect("permission request task"); wait_until_released(&mut pause_state).await; + assert_uncaptured_wait(&turn_context, /*enabled*/ true); } +#[test_case::test_case(true, false; "blocking")] +#[test_case::test_case(false, false; "nonblocking_ui_label")] +#[test_case::test_case(true, true; "cancelled")] #[tokio::test] -async fn request_user_input_holds_an_elicitation_until_response() { - let (session, turn_context, events) = make_session_and_context_with_rx().await; +async fn request_user_input_holds_an_elicitation_until_response(blocking: bool, cancel: bool) { + let (session, mut turn_context, events) = make_session_and_context_with_rx().await; + start_timing(&mut turn_context, /*enabled*/ true).await; *session.active_turn.lock().await = Some(ActiveTurn::default()); let mut pause_state = session.subscribe_elicitation_pause_state(); @@ -178,7 +213,7 @@ async fn request_user_input_holds_an_elicitation_until_response() { "call-1".to_string(), RequestUserInputArgs { questions: Vec::new(), - is_blocking: true, + is_blocking: blocking, auto_resolution_ms: None, }, ) @@ -189,6 +224,13 @@ async fn request_user_input_holds_an_elicitation_until_response() { events.recv().await.expect("request user input event"); wait_until_held(&mut pause_state).await; + if cancel { + request.abort(); + assert!(matches!(request.await, Err(error) if error.is_cancelled())); + wait_until_released(&mut pause_state).await; + assert_uncaptured_wait(&turn_context, /*enabled*/ true); + return; + } let response = RequestUserInputResponse { answers: HashMap::new(), }; @@ -198,4 +240,5 @@ async fn request_user_input_holds_an_elicitation_until_response() { request.await.expect("request user input task"); wait_until_released(&mut pause_state).await; + assert_uncaptured_wait(&turn_context, /*enabled*/ true); } diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index b312cc7345..895770147c 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -732,6 +732,15 @@ impl Session { }; } + // This callback carries a request ID, not a reliable tool invocation owner. + // Do not report its human wait as ordinary tool execution. + if turn_context + .config + .features + .enabled(Feature::CriticalPathTiming) + { + turn_context.turn_timing_state.invalidate_critical_path(); + } let _elicitation = self.services.elicitations.register(); let (tx_response, rx_response) = oneshot::channel(); let prev_entry = { diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index b6dccabe02..1e70c9a7e4 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -2998,6 +2998,17 @@ impl Session { .await; } + fn invalidate_uncaptured_user_wait(&self, turn_context: &TurnContext) { + // Replace this boundary with owner-scoped timing when user waits are connected. + if turn_context + .config + .features + .enabled(Feature::CriticalPathTiming) + { + turn_context.turn_timing_state.invalidate_critical_path(); + } + } + /// Emit an exec approval request event and await the user's decision. /// /// The request is keyed by `call_id` + `approval_id` so matching responses @@ -3100,6 +3111,7 @@ impl Session { available_decisions: Some(available_decisions), parsed_cmd, }); + self.invalidate_uncaptured_user_wait(turn_context); self.send_event(turn_context, event).await; rx_approve.await.unwrap_or(ReviewDecision::Abort) } @@ -3142,6 +3154,7 @@ impl Session { reason, grant_root, }); + self.invalidate_uncaptured_user_wait(turn_context); self.send_event(turn_context, event).await; rx_approve.await.unwrap_or(ReviewDecision::Abort) } @@ -3307,6 +3320,7 @@ impl Session { permissions: requested_permissions, cwd: Some(cwd.clone().into()), }); + self.invalidate_uncaptured_user_wait(turn_context); self.send_event(turn_context.as_ref(), event).await; tokio::select! { biased; @@ -3360,6 +3374,7 @@ impl Session { turn_context .turn_metadata_state .mark_user_input_requested_during_turn(); + self.invalidate_uncaptured_user_wait(turn_context); self.send_event(turn_context, event).await; rx_response.await.ok() } @@ -4189,6 +4204,14 @@ impl Session { communication: impl Into, ) { let communication = communication.into(); + // Agent-message dependencies are not captured by direct-tool timing. + if turn_context + .config + .features + .enabled(Feature::CriticalPathTiming) + { + turn_context.turn_timing_state.invalidate_critical_path(); + } let mut response_item = communication.to_model_input_item(); if let ResponseItem::AgentMessage { content, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 241ead83e8..e85587b72e 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -222,6 +222,7 @@ use opentelemetry_sdk::metrics::data::MetricData; use opentelemetry_sdk::metrics::data::ResourceMetrics; use std::path::Path; use std::time::Duration; +use std::time::Instant; use test_case::test_case; use tokio::sync::Notify; use tokio::sync::Semaphore; @@ -715,15 +716,35 @@ async fn request_mcp_server_elicitation_auto_accepts_when_auto_deny_is_enabled( assert!(rx.try_recv().is_err()); } -#[test_case(SessionSource::Exec; "root")] -#[test_case(SessionSource::SubAgent(SubAgentSource::Review); "subagent")] -#[test_case(SessionSource::Internal(InternalSessionSource::Guardian); "guardian")] +#[test_case(SessionSource::Exec, true; "root")] +#[test_case(SessionSource::Exec, false; "capture_off")] +#[test_case(SessionSource::SubAgent(SubAgentSource::Review), true; "subagent")] +#[test_case(SessionSource::Internal(InternalSessionSource::Guardian), true; "guardian")] #[tokio::test] -async fn request_mcp_server_elicitation_waits_for_user_response(source: SessionSource) { +async fn request_mcp_server_elicitation_waits_for_user_response( + source: SessionSource, + capture: bool, +) { let (session, mut turn_context, rx) = make_session_and_context_with_rx().await; Arc::get_mut(&mut turn_context) .expect("turn context should not be shared") .session_source = source; + Arc::make_mut(&mut Arc::get_mut(&mut turn_context).unwrap().config) + .features + .set_enabled(Feature::CriticalPathTiming, capture) + .unwrap(); + // Enable state independently to prove a disabled feature does not invalidate it. + turn_context.turn_timing_state.enable_critical_path(); + turn_context + .turn_timing_state + .mark_turn_started(Instant::now()) + .await; + assert!( + turn_context + .turn_timing_state + .snapshot_path(Instant::now()) + .is_some() + ); *session.active_turn.lock().await = Some(ActiveTurn::default()); let paused = session.subscribe_elicitation_pause_state(); let request = ElicitationRequest::Url { @@ -764,6 +785,13 @@ async fn request_mcp_server_elicitation_waits_for_user_response(source: SessionS } ); assert!(*paused.borrow()); + assert_eq!( + turn_context + .turn_timing_state + .snapshot_path(Instant::now()) + .is_none(), + capture + ); assert!(!pending.is_finished()); let response = ElicitationResponse { @@ -2358,6 +2386,10 @@ async fn refresh_config_rejects_stale_owner_without_effects(scope: RuntimeConfig !stale_owner.features.enabled(Feature::Mcp20260728), ) .expect("set MCP protocol rollout"); + stale_config + .features + .enable(Feature::CriticalPathTiming) + .unwrap(); assert_eq!( session .refresh_runtime_config(Arc::clone(&stale_owner), stale_owner.as_ref().clone()) @@ -2782,7 +2814,22 @@ async fn inter_agent_communication_waits_for_confirmed_delivery_persistence() { #[tokio::test] async fn record_inter_agent_communication_sets_turn_id_in_rollout_and_resume() { - let (mut session, turn_context) = make_session_and_context().await; + let (mut session, mut turn_context) = make_session_and_context().await; + Arc::make_mut(&mut turn_context.config) + .features + .enable(Feature::CriticalPathTiming) + .unwrap(); + turn_context.turn_timing_state.enable_critical_path(); + turn_context + .turn_timing_state + .mark_turn_started(std::time::Instant::now()) + .await; + assert!( + turn_context + .turn_timing_state + .snapshot_path(std::time::Instant::now()) + .is_some() + ); let rollout_path = attach_thread_persistence(&mut session).await; let communication = InterAgentCommunication::new( AgentPath::root().join("worker").expect("worker path"), @@ -2803,6 +2850,12 @@ async fn record_inter_agent_communication_sets_turn_id_in_rollout_and_resume() { .record_inter_agent_communication(&turn_context, turn_context.model_info(), communication) .await; + assert!( + turn_context + .turn_timing_state + .snapshot_path(std::time::Instant::now()) + .is_none() + ); let recorded_history = session.clone_history().await; let recorded_items = raw_history_items(&recorded_history); let create_time = recorded_items[0] diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 6f2e622367..ed445e0dae 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -3429,7 +3429,9 @@ async fn try_run_sampling_request( .await; if !in_flight.is_empty() { - let _tool_blocking_timing_guard = turn_context.turn_timing_state.begin_tool_blocking(); + let _tool_blocking_timing_guard = turn_context + .turn_timing_state + .begin_tool_blocking(tool_runtime.batch_timing()); let _tool_blocking_span = trace_span!( "codex.tool_blocking", codex.turn.phase = "tool_blocking", diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 2ebfcc173d..2d2f7e66e3 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -997,10 +997,16 @@ impl Session { }, _ = tokio::time::sleep(Duration::from_millis(GRACEFULL_INTERRUPTION_TIMEOUT_MS)) => { warn!("task {sub_id} didn't complete gracefully after {}ms", GRACEFULL_INTERRUPTION_TIMEOUT_MS); + // Forced teardown can abort independently spawned tool dispatches + // whose guards are destroyed later. Do not publish complete totals + // for a dependency graph that has not settled. + task.turn_context.turn_timing_state.invalidate_critical_path(); } } task.handle.abort(); + // Aborting schedules destruction; join before freezing the turn's timing. + let _ = task.handle.await; session_task .abort(Arc::clone(self), Arc::clone(&task.turn_context)) diff --git a/codex-rs/core/src/tasks/mod_tests.rs b/codex-rs/core/src/tasks/mod_tests.rs index beeda10229..ba03d83a15 100644 --- a/codex-rs/core/src/tasks/mod_tests.rs +++ b/codex-rs/core/src/tasks/mod_tests.rs @@ -494,3 +494,73 @@ fn emit_compact_metric_records_auto_local() { ]) ); } + +struct UncooperativeTimedTask { + entered: Arc, +} + +impl SessionTask for UncooperativeTimedTask { + fn kind(&self) -> TaskKind { + TaskKind::Regular + } + fn span_name(&self) -> &'static str { + "session_task.uncooperative_timing" + } + async fn run( + self: Arc, + _session: Arc, + context: Arc, + _input: Vec, + _cancellation_token: CancellationToken, + ) -> SessionTaskResult { + let timing = &context.turn_timing_state; + timing.enable_critical_path(); + let batch = codex_direct_tool_timing::ToolBatchTiming::default(); + let now = std::time::Instant::now(); + let _blocking; + let _work = timing.direct_tool_timings.start_turn_direct_call( + "uncooperative", + &context.sub_id, + batch.clone(), + now, + timing.snapshot_path(now), + ); + _blocking = timing.begin_tool_blocking(Some(batch)); + let _review = _work.start_review(); + self.entered.notify_one(); + std::future::pending().await + } +} + +#[tokio::test(flavor = "current_thread")] +async fn hard_abort_drops_timing_guards_before_profile_completion() { + let (session, context, _receiver) = make_session_and_context_with_rx().await; + let entered = Arc::new(tokio::sync::Notify::new()); + session + .spawn_task( + Arc::clone(&context), + Vec::new(), + UncooperativeTimedTask { + entered: Arc::clone(&entered), + }, + ) + .await; + entered.notified().await; + session + .abort_all_tasks(codex_protocol::protocol::TurnAbortReason::Interrupted) + .await; + // The task's review and tool guards are gone before abort completion. + assert!( + context + .turn_timing_state + .direct_tool_timings + .owner_for_turn("uncooperative", &context.sub_id) + .is_none() + ); + assert_eq!( + context + .turn_timing_state + .snapshot_path(std::time::Instant::now()), + None + ); +} diff --git a/codex-rs/core/src/tasks/regular.rs b/codex-rs/core/src/tasks/regular.rs index 3472d95953..6c6f3d10da 100644 --- a/codex-rs/core/src/tasks/regular.rs +++ b/codex-rs/core/src/tasks/regular.rs @@ -44,6 +44,14 @@ impl SessionTask for RegularTask { input: Vec, cancellation_token: CancellationToken, ) -> SessionTaskResult { + if ctx + .config + .features + .enabled(codex_features::Feature::CriticalPathTiming) + && !ctx.session_source.is_non_root_agent() + { + ctx.turn_timing_state.enable_critical_path(); + } let run_turn_span = trace_span!( "run_turn", conversation.id = %sess.thread_id, diff --git a/codex-rs/core/src/test_support.rs b/codex-rs/core/src/test_support.rs index 359ca8a0ea..b9f94b7eb1 100644 --- a/codex-rs/core/src/test_support.rs +++ b/codex-rs/core/src/test_support.rs @@ -96,6 +96,38 @@ pub fn mcp_attribution_snapshot( .mcp_attribution_snapshot() } +/// Retains the active turn's timing after completion without exporting telemetry. +/// Returns `None` when there is no running task. The observer reports exclusive +/// durations for the caller's interval in inference, automatic review, local +/// tools, overhead, and unknown order. +pub async fn observe_turn_timing( + thread: &crate::CodexThread, +) -> Option< + impl Fn(std::time::Instant, std::time::Instant) -> Option<[std::time::Duration; 5]> + use<>, +> { + let active = thread.session.active_turn.lock().await; + let timing = Arc::clone( + &active + .as_ref()? + .task + .as_ref()? + .turn_context + .turn_timing_state, + ); + Some(move |start, end| { + timing.snapshot_path(end).map(|path| { + let totals = path.durations(start, end); + [ + totals.inference, + totals.auto_review, + totals.tools, + totals.overhead, + totals.unknown, + ] + }) + }) +} + /// Test-only provider that supplies no user instructions. #[derive(Debug, Default)] pub struct EmptyUserInstructionsProvider; diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index 4dea6384f3..7d11aff463 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -27,6 +27,9 @@ use crate::tools::registry::AnyToolResult; use crate::tools::registry::ToolArgumentDiffConsumer; use crate::tools::router::ToolCall; use crate::tools::router::ToolCallSource; +use codex_direct_tool_timing::ToolAdmission; +use codex_direct_tool_timing::ToolBatchTiming; +use codex_direct_tool_timing::ToolGateTiming; use codex_history::ResponseItemEnvelope; use codex_protocol::error::CodexErr; use codex_protocol::models::ResponseInputItem; @@ -49,6 +52,25 @@ pub(crate) struct ToolCallRuntime { step_context: Arc, tracker: SharedTurnDiffTracker, parallel_execution: Arc>, + timing: Option, +} + +#[derive(Clone, Default)] +struct DirectToolTiming { + gate: ToolGateTiming, + batch: ToolBatchTiming, +} + +// Remove this boundary when CodeMode and subagent dependencies have runtime capture. +fn direct_path_supported(tool_name: &codex_tools::ToolName, source: &ToolCallSource) -> bool { + !(matches!(source, ToolCallSource::CodeMode { .. }) + || (tool_name.is_default_namespace() + && matches!(tool_name.name.as_str(), + crate::tools::code_mode::PUBLIC_TOOL_NAME | crate::tools::code_mode::WAIT_TOOL_NAME)) + // Multi-agent tool namespaces can be overridden by the tool catalog. + || matches!(tool_name.name.as_str(), "spawn_agent" | "send_input" + | "close_agent" | "resume_agent" | "wait_agent" | "list_agents" + | "followup_task" | "interrupt_agent" | "send_message")) } impl ToolCallRuntime { @@ -57,11 +79,18 @@ impl ToolCallRuntime { step_context: Arc, tracker: SharedTurnDiffTracker, ) -> Self { + let timing = step_context + .turn + .config + .features + .enabled(codex_features::Feature::CriticalPathTiming) + .then(DirectToolTiming::default); Self { session, step_context, tracker, parallel_execution: Arc::new(RwLock::new(())), + timing, } } @@ -152,6 +181,35 @@ impl ToolCallRuntime { let lock = Arc::clone(&self.parallel_execution); let invocation_cancellation_token = cancellation_token.clone(); let started = Instant::now(); + if turn + .config + .features + .enabled(codex_features::Feature::CriticalPathTiming) + && !direct_path_supported( + &tool_runtime + .as_ref() + .map_or_else(|| call.tool_name.clone(), |runtime| runtime.tool_name()), + &source, + ) + { + turn.turn_timing_state.invalidate_critical_path(); + } + let work_timing = self.timing.and_then(|timing| match &source { + ToolCallSource::Direct | ToolCallSource::DirectPlaintextMessage => { + let work = turn + .turn_timing_state + .direct_tool_timings + .start_turn_direct_call( + &call.call_id, + &turn.sub_id, + timing.batch, + started, + turn.turn_timing_state.snapshot_path(started), + ); + Some((timing.gate, work)) + } + ToolCallSource::CodeMode { .. } => None, + }); let mut tool_call_timing_guard = ToolCallTimingGuard::capture( started, &session.thread_id, @@ -208,7 +266,7 @@ impl ToolCallRuntime { async move { let _message_admission = message_admission?; // Readiness and admission remain cancellable even when the handler must finish. - let guard = async { + let (guard, admission_started) = async { if let Some(tool_runtime) = tool_runtime && let Some(readiness) = tool_runtime.wait_until_ready(&session) { @@ -217,7 +275,8 @@ impl ToolCallRuntime { .await; } - if supports_parallel { + let admission_started = Instant::now(); + let guard = if supports_parallel { Either::Left( lock.read() .instrument(trace_span!("tool_dispatch.parallel_lock")) @@ -229,10 +288,24 @@ impl ToolCallRuntime { .instrument(trace_span!("tool_dispatch.exclusive_lock")) .await, ) - } + }; + (guard, admission_started) } .or_cancel_tool(&invocation_cancellation_token) .await?; + // Publish the completed path before releasing the admission lock, + // including when cancellation drops the handler future. + let _admitted_timing = work_timing.map(|(gate, work)| { + gate.admit( + work, + if supports_parallel { + ToolAdmission::Parallel + } else { + ToolAdmission::Exclusive + }, + admission_started, + ) + }); // 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 { @@ -250,6 +323,7 @@ impl ToolCallRuntime { dispatch_call_state, ) .await; + drop(_admitted_timing); drop(guard); // The sampling loop collects results in order only after its stream ends. // Record readiness here, before either caller encodes or collects the result. @@ -347,6 +421,10 @@ impl ToolCallRuntime { } impl ToolCallRuntime { + pub(crate) fn batch_timing(&self) -> Option { + self.timing.as_ref().map(|timing| timing.batch.clone()) + } + fn tool_task_join_error(err: JoinError) -> FunctionCallError { FunctionCallError::Fatal(format!("tool task failed to receive: {err:?}")) } @@ -529,6 +607,40 @@ mod tests { use tokio::sync::oneshot; use tracing_test::internal::MockWriter; + #[test] + fn unsupported_handlers_remain_incomplete_after_namespace_resolution() { + let handler = crate::tools::multi_agent_tool::multi_agent_v2_handler( + crate::tools::handlers::multi_agents_v2::WaitAgentHandler::default(), + Some("custom_collaboration"), + /*description_override*/ None, + /*parameters_override*/ None, + ); + let registry = ToolRegistry::from_tools([handler]); + let registered = registry + .registered_tool(&codex_tools::ToolName::namespaced( + "custom_collaboration", + "wait_agent", + )) + .unwrap(); + assert!(!direct_path_supported( + ®istered.runtime.tool_name(), + &ToolCallSource::Direct + )); + for name in [ + crate::tools::code_mode::PUBLIC_TOOL_NAME, + crate::tools::code_mode::WAIT_TOOL_NAME, + ] { + assert!(!direct_path_supported( + &codex_tools::ToolName::plain(name).with_default_namespace(), + &ToolCallSource::Direct + )); + } + assert!(direct_path_supported( + &codex_tools::ToolName::namespaced("mcp", "exec"), + &ToolCallSource::Direct + )); + } + #[test] fn tool_call_timing_guard_ignores_code_mode_source() { let buffer: &'static std::sync::Mutex> = @@ -622,13 +734,34 @@ mod tests { }); } - #[test_case::test_case(false; "ordinary_handler")] - #[test_case::test_case(true; "dynamic_handler")] + #[test_case::test_case(false, false; "ordinary_capture_disabled")] + #[test_case::test_case(false, true; "ordinary_capture_enabled")] + #[test_case::test_case(true, false; "dynamic_capture_disabled")] + #[test_case::test_case(true, true; "dynamic_capture_enabled")] #[tokio::test] async fn cancellation_before_dispatch_admission_logs_dispatch_only_timing( dynamic: bool, + critical_path_timing: bool, ) -> anyhow::Result<()> { - let (session, turn_context) = crate::session::tests::make_session_and_context().await; + let (session, mut turn_context) = crate::session::tests::make_session_and_context().await; + Arc::make_mut(&mut turn_context.config) + .features + .enable(codex_features::Feature::RuntimeMetrics) + .unwrap(); + Arc::make_mut(&mut turn_context.config) + .features + .set_enabled( + codex_features::Feature::CriticalPathTiming, + critical_path_timing, + ) + .unwrap(); + if critical_path_timing { + turn_context.turn_timing_state.enable_critical_path(); + } + turn_context + .turn_timing_state + .mark_turn_started(Instant::now()) + .await; let session = Arc::new(session); let turn_context = Arc::new(turn_context); let tool_name = codex_tools::ToolName::plain("test_tool"); @@ -661,6 +794,8 @@ mod tests { let step_context = step_context.with_tool_router_for_test(router); let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); let runtime = ToolCallRuntime::new(session, step_context, tracker); + let batch = runtime.batch_timing(); + assert_eq!(batch.is_some(), critical_path_timing); let execution_gate = Arc::clone(&runtime.parallel_execution); let execution_gate_guard = execution_gate .try_write_owned() @@ -700,6 +835,11 @@ mod tests { .expect("cancelled tool response task should join") .expect("cancelled tool call should produce a response"); + assert_eq!( + batch.and_then(|batch| batch.completed()).is_some(), + critical_path_timing, + "capture requires its own gate and must finish before cancellation returns" + ); let logs = String::from_utf8( buffer .lock() diff --git a/codex-rs/core/src/turn_timing.rs b/codex-rs/core/src/turn_timing.rs index 37d238b7ec..84d1d05ef7 100644 --- a/codex-rs/core/src/turn_timing.rs +++ b/codex-rs/core/src/turn_timing.rs @@ -7,6 +7,7 @@ use std::time::SystemTime; use std::time::UNIX_EPOCH; use codex_analytics::TurnProfile; +use codex_direct_tool_timing::DirectToolTimings; use codex_otel::TURN_TTFM_DURATION_METRIC; use codex_protocol::ThreadId; use codex_protocol::items::TurnItem; @@ -16,6 +17,9 @@ use tokio::sync::Mutex; use crate::ResponseEvent; use crate::session::turn_context::TurnContext; use crate::stream_events_utils::raw_assistant_output_text_from_item; +use codex_direct_tool_timing::ToolBatchTiming; +use codex_utils_critical_path::CriticalPath; +use codex_utils_critical_path::WorkKind; // Reuse the finalized analytics counters instead of counting sampled child spans. // Recovery can complete the same turn ID again, so consumers must retain the @@ -59,6 +63,7 @@ pub(crate) async fn record_turn_ttfm_metric(turn_context: &TurnContext, item: &T #[derive(Debug, Default)] pub(crate) struct TurnTimingState { + pub(crate) direct_tool_timings: DirectToolTimings, state: Mutex, profile: StdMutex, } @@ -74,6 +79,8 @@ struct TurnTimingStateInner { #[derive(Debug, Default)] struct TurnProfileState { + path: Option, + tool_dependency: Option, started_at: Option, last_transition_at: Option, active_phase: Option, @@ -102,9 +109,33 @@ pub(crate) struct TurnProfileTimingGuard { timing: Arc, phase: TurnProfilePhase, active: bool, + tool_batch: Option, } impl TurnTimingState { + pub(crate) fn enable_critical_path(&self) { + let mut state = self.profile_state(); + if state.completed_profile.is_none() { + state.path = Some(CriticalPath::default()); + } + } + + pub(crate) fn invalidate_critical_path(&self) { + let mut state = self.profile_state(); + if state.completed_profile.is_none() { + state.path = None; + } + } + + pub(crate) fn snapshot_path(&self, now: Instant) -> Option { + let state = self.profile_state(); + if state.completed_profile.is_some() { + state.path.clone() + } else { + state.path_at(now) + } + } + pub(crate) async fn mark_turn_started(&self, started_at: Instant) -> i64 { let started_at_unix_ms = now_unix_timestamp_ms(); let mut state = self.state.lock().await; @@ -166,6 +197,7 @@ impl TurnTimingState { timing: Arc::clone(self), phase: TurnProfilePhase::Sampling, active, + tool_batch: None, } } @@ -186,15 +218,20 @@ impl TurnTimingState { timing: Arc::clone(self), phase: TurnProfilePhase::Compaction, active, + tool_batch: None, } } - pub(crate) fn begin_tool_blocking(self: &Arc) -> TurnProfileTimingGuard { + pub(crate) fn begin_tool_blocking( + self: &Arc, + batch: Option, + ) -> TurnProfileTimingGuard { let active = self.profile_state().begin_tool_blocking(Instant::now()); TurnProfileTimingGuard { timing: Arc::clone(self), phase: TurnProfilePhase::ToolBlocking, active, + tool_batch: batch, } } @@ -227,9 +264,14 @@ impl TurnTimingState { impl Drop for TurnProfileTimingGuard { fn drop(&mut self) { if self.active { - self.timing - .profile_state() - .end_phase(Instant::now(), self.phase); + let mut state = self.timing.profile_state(); + if state.completed_profile.is_some() { + return; + } + if let Some(batch) = &self.tool_batch { + state.tool_dependency = batch.completed().and_then(|work| work.path); + } + state.end_phase(Instant::now(), self.phase); } } } @@ -251,7 +293,9 @@ fn duration_to_u64_ms(duration: Duration) -> u64 { impl TurnProfileState { fn start(&mut self, started_at: Instant) { + let path = self.path.clone(); *self = Self { + path, started_at: Some(started_at), last_transition_at: Some(started_at), ..Self::default() @@ -318,6 +362,11 @@ impl TurnProfileState { return; }; let elapsed = now.saturating_duration_since(previous); + self.path = self + .path + .as_ref() + .map(|path| self.extend_path(path.clone(), previous, now)); + self.tool_dependency = None; match self.active_phase { Some(TurnProfilePhase::Sampling) => self.sampling += elapsed, Some(TurnProfilePhase::Compaction) => self.compaction += elapsed, @@ -327,11 +376,33 @@ impl TurnProfileState { } } + fn path_at(&self, now: Instant) -> Option { + let path = self.path.clone()?; + Some(self.extend_path(path, self.last_transition_at?, now)) + } + + fn extend_path(&self, path: CriticalPath, previous: Instant, now: Instant) -> CriticalPath { + match self.active_phase { + // Compaction is model-related elapsed work, including local bookkeeping; + // the existing profile still exposes its separate compaction_ms duration. + Some(TurnProfilePhase::Sampling | TurnProfilePhase::Compaction) => { + path.work(previous, now, WorkKind::Inference) + } + Some(TurnProfilePhase::ToolBlocking) => { + let dependency = self + .tool_dependency + .clone() + .unwrap_or_else(|| path.clone().work(previous, now, WorkKind::Unknown)); + path.join(previous, now, dependency) + } + None => path.work(previous, now, WorkKind::Overhead), + } + } + fn complete(&mut self, now: Instant) -> TurnProfile { if let Some(profile) = self.completed_profile.as_ref() { return profile.clone(); } - let final_phase = self.active_phase; self.advance(now); let after_last_sampling = if self.seen_sampling { diff --git a/codex-rs/core/src/turn_timing_tests.rs b/codex-rs/core/src/turn_timing_tests.rs index 7aa7d3d177..1045ad16af 100644 --- a/codex-rs/core/src/turn_timing_tests.rs +++ b/codex-rs/core/src/turn_timing_tests.rs @@ -1,10 +1,15 @@ +use super::ToolBatchTiming; use codex_analytics::TurnProfile; use codex_protocol::items::AgentMessageItem; use codex_protocol::items::TurnItem; use codex_protocol::models::ContentItem; use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ResponseItem; +use codex_utils_critical_path::CriticalPath; +use codex_utils_critical_path::PathDurations; +use codex_utils_critical_path::WorkKind; use pretty_assertions::assert_eq; +use std::sync::Arc; use std::time::Duration; use std::time::Instant; use std::time::SystemTime; @@ -256,7 +261,10 @@ fn turn_profile_breaks_down_sampling_blocking_and_retry_overhead() { #[test] fn turn_profile_counts_compaction_as_an_exclusive_phase() { let started_at = Instant::now(); - let mut state = TurnProfileState::default(); + let mut state = TurnProfileState { + path: Some(CriticalPath::default()), + ..Default::default() + }; state.start(started_at); let _ = state.begin_compaction(started_at + Duration::from_millis(100)); @@ -285,6 +293,14 @@ fn turn_profile_counts_compaction_as_an_exclusive_phase() { tools_change_count: 0, } ); + let totals = state + .path + .as_ref() + .unwrap() + .durations(started_at, started_at + Duration::from_millis(1_100)); + assert_eq!(totals.inference, Duration::from_millis(800)); + assert_eq!(totals.overhead, Duration::from_millis(300)); + assert_eq!(totals.user_wait, Duration::ZERO); } #[tokio::test] @@ -307,3 +323,92 @@ async fn turn_profile_and_duration_share_a_completion_instant() { Some(i64::try_from(classified_ms).expect("profile duration should fit i64")) ); } + +#[test] +fn capture_preserves_phases_and_freezes_missing_dependencies() { + let started_at = Instant::now(); + let time_after_turn_start = |millis| started_at + Duration::from_millis(millis); + let mut state = TurnProfileState::default(); + state.start(started_at); + assert!(state.path_at(time_after_turn_start(10)).is_none()); + state.path = Some(CriticalPath::default()); + assert!(state.begin_sampling(time_after_turn_start(10))); + state.end_phase(time_after_turn_start(20), TurnProfilePhase::Sampling); + assert!(state.begin_tool_blocking(time_after_turn_start(20))); + state.end_phase(time_after_turn_start(50), TurnProfilePhase::ToolBlocking); + let profile = state.complete(time_after_turn_start(60)); + assert_eq!( + state + .path + .as_ref() + .unwrap() + .durations(started_at, time_after_turn_start(60)), + PathDurations { + inference: Duration::from_millis(10), + unknown: Duration::from_millis(30), + overhead: Duration::from_millis(20), + ..PathDurations::default() + } + ); + let frozen_path = state.path.clone(); + state.end_phase(time_after_turn_start(70), TurnProfilePhase::ToolBlocking); + assert_eq!(state.complete(time_after_turn_start(80)), profile); + assert_eq!(state.path, frozen_path); +} + +#[tokio::test] +async fn invalidation_and_late_guard_drop_cannot_restore_a_frozen_path() { + let timing = Arc::new(TurnTimingState::default()); + timing.enable_critical_path(); + timing.mark_turn_started(Instant::now()).await; + let guard = timing.begin_tool_blocking(Some(ToolBatchTiming::default())); + timing.invalidate_critical_path(); + timing.complete_profile_and_duration_ms().await; + timing.enable_critical_path(); + drop(guard); + assert!(timing.snapshot_path(Instant::now()).is_none()); +} + +#[tokio::test] +async fn completed_batch_guard_retains_its_path_across_turns() { + let started_at = Instant::now() - Duration::from_secs(2); + let previous_turn = TurnTimingState::default(); + let previous_batch = ToolBatchTiming::default(); + let previous_call = previous_turn.direct_tool_timings.start_turn_direct_call( + "reviewed", + "turn", + previous_batch.clone(), + started_at, + Some(CriticalPath::default()), + ); + drop(previous_turn); + let timing = Arc::new(TurnTimingState::default()); + timing.enable_critical_path(); + timing.mark_turn_started(started_at).await; + let batch = ToolBatchTiming::default(); + let review_path = CriticalPath::default().work( + started_at, + started_at + Duration::from_secs(1), + WorkKind::AutoReview, + ); + let call = timing.direct_tool_timings.start_turn_direct_call( + "reviewed", + "turn", + batch.clone(), + started_at + Duration::from_secs(1), + Some(review_path), + ); + // A late guard from an earlier turn cannot finish or invalidate the reused ID. + drop(previous_call); + assert!(previous_batch.completed().unwrap().path.is_some()); + assert!(batch.completed().is_none()); + let guard = timing.begin_tool_blocking(Some(batch)); + drop(call); + drop(guard); + timing.complete_profile_and_duration_ms().await; + let path = timing.snapshot_path(Instant::now()).unwrap(); + let totals = path.durations(started_at, path.end().unwrap()); + assert_eq!(totals.auto_review, Duration::from_secs(1)); + assert_eq!(totals.unknown, Duration::ZERO); + assert!(totals.tools >= Duration::from_secs(1)); +} diff --git a/codex-rs/core/tests/suite/auto_review.rs b/codex-rs/core/tests/suite/auto_review.rs index 13f0e8d778..e53d538054 100644 --- a/codex-rs/core/tests/suite/auto_review.rs +++ b/codex-rs/core/tests/suite/auto_review.rs @@ -59,7 +59,10 @@ use wiremock::ResponseTemplate; use wiremock::matchers::method; use wiremock::matchers::path; -struct ApprovedReviewContributor; +struct ApprovedReviewContributor { + entered: tokio::sync::Notify, + release: tokio::sync::Semaphore, +} impl ApprovalReviewContributor for ApprovedReviewContributor { fn decide<'a>( @@ -67,6 +70,8 @@ impl ApprovalReviewContributor for ApprovedReviewContributor { input: &'a codex_extension_api::ApprovalDecisionInput<'_>, ) -> ExtensionFuture<'a, Option> { Box::pin(async move { + self.entered.notify_one(); + self.release.acquire().await.expect("review gate").forget(); assert!(input.metrics.is_some()); assert_eq!(input.action["tool"], "request_permissions"); Some(codex_extension_api::ApprovalDecision::Allow) @@ -125,11 +130,19 @@ async fn approval_review_contributor_skips_existing_guardian_model_call() -> Res .await; let mut extensions = ExtensionRegistryBuilder::new(); - extensions.approval_review_contributor(Arc::new(ApprovedReviewContributor)); + let reviewer = Arc::new(ApprovedReviewContributor { + entered: tokio::sync::Notify::new(), + release: tokio::sync::Semaphore::new(/* permits */ 0), + }); + extensions.approval_review_contributor(reviewer.clone()); let mut builder = test_codex() .with_extensions(Arc::new(extensions.build())) .with_config(|config| { config.approvals_reviewer = ApprovalsReviewer::AutoReview; + config + .features + .enable(Feature::CriticalPathTiming) + .expect("critical-path capture is configurable"); config .features .enable(Feature::RequestPermissionsTool) @@ -143,6 +156,13 @@ async fn approval_review_contributor_skips_existing_guardian_model_call() -> Res }])) .await?; + reviewer.entered.notified().await; + let observe = codex_core::test_support::observe_turn_timing(&test.codex) + .await + .expect("active turn"); + let now = std::time::Instant::now(); + reviewer.release.add_permits(/* n */ 1); + loop { match wait_for_event(&test.codex, |_| true).await { EventMsg::GuardianAssessment(event) => { @@ -156,6 +176,7 @@ async fn approval_review_contributor_skips_existing_guardian_model_call() -> Res } } + assert_eq!(observe(now, now), None); let requests = responses.requests(); assert_eq!(requests.len(), 2); assert!( @@ -332,7 +353,10 @@ async fn required_model_bypasses_extension_approval_when_guardian_v2_is_disabled let requirements = format!("{reviewer_requirements}[auto_review]\nrequired_on_models = [\"{model}\"]\n"); let mut extensions = ExtensionRegistryBuilder::new(); - extensions.approval_review_contributor(Arc::new(ApprovedReviewContributor)); + extensions.approval_review_contributor(Arc::new(ApprovedReviewContributor { + entered: tokio::sync::Notify::new(), + release: tokio::sync::Semaphore::new(/* permits */ 1), + })); let mut builder = test_codex() .with_model(model) .with_extensions(Arc::new(extensions.build())) diff --git a/codex-rs/core/tests/suite/critical_path_runtime.rs b/codex-rs/core/tests/suite/critical_path_runtime.rs new file mode 100644 index 0000000000..49d3506ea8 --- /dev/null +++ b/codex-rs/core/tests/suite/critical_path_runtime.rs @@ -0,0 +1,251 @@ +//! Exercises dependency capture through real tool dispatch and admission gates. +//! Channels hold work open and establish overlap without scheduling sleeps. + +use std::sync::Arc; +use std::time::Duration; +use std::time::Instant; + +use codex_core::TurnInputRequest; +use codex_extension_api::ExtensionData; +use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::ToolContributor; +use codex_features::Feature; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::Op; +use codex_protocol::user_input::UserInput; +use codex_tools::JsonToolOutput; +use codex_tools::ToolCall; +use codex_tools::ToolExecutor; +use codex_tools::ToolExecutorFuture; +use codex_tools::ToolName; +use codex_tools::ToolOutput; +use codex_tools::ToolSpec; +use core_test_support::responses; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; +use pretty_assertions::assert_eq; +use tokio::sync::Semaphore; +use tokio::sync::mpsc; + +#[derive(Clone, Copy)] +enum Admission { + Serial, + Parallel, + Unsupported, +} + +#[derive(Clone)] +struct GatedTool { + admission: Admission, + entered: mpsc::UnboundedSender<(String, Instant)>, + release: Arc, +} + +impl ToolContributor for GatedTool { + fn tools( + &self, + _session_store: &ExtensionData, + _thread_store: &ExtensionData, + ) -> Vec ToolExecutor>>> { + vec![Arc::new(self.clone())] + } +} + +impl<'call> ToolExecutor> for GatedTool { + fn tool_name(&self) -> ToolName { + ToolName::plain(if matches!(self.admission, Admission::Unsupported) { + "wait_agent" + } else { + "timing_gate" + }) + } + + fn spec(&self) -> ToolSpec { + ToolSpec::Function(codex_tools::ResponsesApiTool { + name: self.tool_name().name, + description: "Wait for the host.".to_string(), + strict: false, + parameters: codex_tools::JsonSchema::default(), + output_schema: None, + defer_loading: None, + }) + } + + fn supports_parallel_tool_calls(&self) -> bool { + matches!(self.admission, Admission::Parallel) + } + + fn handle<'a>(&'a self, call: ToolCall<'call>) -> ToolExecutorFuture<'a> + where + 'call: 'a, + { + Box::pin(async move { + self.entered + .send((call.call_id.to_string(), Instant::now())) + .expect("test receiver"); + self.release.acquire().await.expect("test gate").forget(); + Ok( + Box::new(JsonToolOutput::new(serde_json::json!({"done": true}))) + as Box, + ) + }) + } +} + +#[test_case::test_case(Admission::Serial, false; "default_disabled")] +#[test_case::test_case(Admission::Unsupported, true; "unsupported_enabled")] +#[test_case::test_case(Admission::Serial, true; "serial_enabled")] +#[test_case::test_case(Admission::Parallel, true; "overlapping_enabled")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn critical_path_follows_runtime_admission( + admission: Admission, + critical_path_timing: bool, +) -> anyhow::Result<()> { + let server = responses::start_mock_server().await; + let (entered, mut entries) = mpsc::unbounded_channel(); + let gate = GatedTool { + admission, + entered, + release: Arc::new(Semaphore::new(/* permits */ 0)), + }; + let name = gate.tool_name().name; + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.tool_contributor(Arc::new(gate.clone())); + let test = test_codex() + .with_extensions(Arc::new(extensions.build())) + .with_config(move |config| { + if critical_path_timing { + config + .features + .enable(Feature::CriticalPathTiming) + .expect("critical-path capture is configurable"); + } + }) + .build_with_auto_env(&server) + .await?; + let first = responses::mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("tools"), + responses::ev_function_call("a", &name, "{}"), + responses::ev_function_call("b", &name, "{}"), + responses::ev_completed("tools"), + ]), + ) + .await; + let start = Instant::now(); + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "Run both tools".to_string(), + text_elements: Vec::new(), + }])) + .await?; + let (first_id, first_start) = entries.recv().await.expect("first admitted tool"); + let observe = codex_core::test_support::observe_turn_timing(&test.codex) + .await + .expect("active turn"); + let reply = responses::mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("reply"), + responses::ev_assistant_message("answer", "Done"), + responses::ev_completed("reply"), + ]), + ) + .await; + let (second_start, held_start, held_end) = if matches!(admission, Admission::Parallel) { + let (second_id, second_start) = entries.recv().await.expect("second overlapping tool"); + assert_ne!(first_id, second_id); + // Both handlers are inside their gate before either can finish. + let held_end = Instant::now(); + gate.release.add_permits(/* n */ 2); + (second_start, second_start, held_end) + } else { + assert!( + entries.try_recv().is_err(), + "exclusive handler still owns admission" + ); + let held_end = Instant::now(); + gate.release.add_permits(/* n */ 1); + let (_, second_start) = entries.recv().await.expect("second serial tool"); + gate.release.add_permits(/* n */ 1); + (second_start, first_start, held_end) + }; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + first.single_request(); + let request = reply.single_request(); + for id in ["a", "b"] { + assert!( + request + .function_call_output(id) + .to_string() + .contains("done") + ); + } + let end = Instant::now(); + if matches!(admission, Admission::Unsupported) || !critical_path_timing { + assert_eq!(observe(start, end), None); + } else { + let [inference, review, tools, overhead, unknown] = + observe(start, end).expect("complete direct-tool path"); + assert_eq!( + observe(held_start, held_end), + Some([ + Duration::ZERO, + Duration::ZERO, + held_end.duration_since(held_start), + Duration::ZERO, + Duration::ZERO, + ]) + ); + assert!(inference > Duration::ZERO); + assert!(tools >= second_start.duration_since(first_start)); + assert_eq!((review, unknown), (Duration::ZERO, Duration::ZERO)); + assert_eq!( + inference + review + tools + overhead + unknown, + end.duration_since(start) + ); + // Retaining both branches would count their common prefix twice. + assert!(tools < end.duration_since(start)); + } + Ok(()) +} + +#[tokio::test] +async fn standalone_user_shell_omits_uncaptured_path() -> anyhow::Result<()> { + core_test_support::skip_if_target_windows!(Ok(()), "Unix shell workload"); + let server = responses::start_mock_server().await; + let test = test_codex() + .with_config(|config| { + config + .features + .enable(Feature::CriticalPathTiming) + .expect("critical-path capture is configurable"); + }) + .build_with_auto_env(&server) + .await?; + test.codex + .submit(Op::RunUserShellCommand { + command: "sleep 30".to_string(), + timeout_ms: None, + }) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::ExecCommandBegin(_)) + }) + .await; + let observe = codex_core::test_support::observe_turn_timing(&test.codex) + .await + .expect("active turn"); + test.codex.submit(Op::Interrupt).await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnAborted(_)) + }) + .await; + let now = Instant::now(); + assert_eq!(observe(now, now), None); + Ok(()) +} diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index 9deb12935c..9bb0c274f2 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -99,6 +99,7 @@ mod guardian_cached_score; mod guardian_checkpoint_migration; // Uses the same command-approval harness as guardian_review below. mod canonical_plugin_connectors; +mod critical_path_runtime; mod gateway_auth; mod guardian_connector_trust; #[cfg(not(target_os = "windows"))] diff --git a/codex-rs/core/tests/suite/turn_phase_trace.rs b/codex-rs/core/tests/suite/turn_phase_trace.rs index 0d1d2f18f7..aa4fac477b 100644 --- a/codex-rs/core/tests/suite/turn_phase_trace.rs +++ b/codex-rs/core/tests/suite/turn_phase_trace.rs @@ -63,19 +63,27 @@ enum Finish { Interrupt, } -#[test_case(Finish::Reply; "reply")] -#[test_case(Finish::Interrupt; "interrupt")] +#[test_case(Finish::Reply, false; "reply_disabled")] +#[test_case(Finish::Interrupt, false; "interrupt_disabled")] +#[test_case(Finish::Reply, true; "reply_enabled")] +#[test_case(Finish::Interrupt, true; "interrupt_enabled")] #[tokio::test(flavor = "current_thread")] async fn steer_joins_existing_turn_without_extending_sampling( finish: Finish, + critical_path_timing: bool, ) -> anyhow::Result<()> { if std::env::var_os(SUBPROCESS_ENV_VAR).is_none() { let case = match finish { Finish::Reply => "reply", Finish::Interrupt => "interrupt", }; + let capture = if critical_path_timing { + "enabled" + } else { + "disabled" + }; let test_name = format!( - "suite::turn_phase_trace::steer_joins_existing_turn_without_extending_sampling::{case}" + "suite::turn_phase_trace::steer_joins_existing_turn_without_extending_sampling::{case}_{capture}" ); let output = Command::new(std::env::current_exe()?) .arg("--exact") @@ -117,7 +125,22 @@ async fn steer_joins_existing_turn_without_extending_sampling( ]), ) .await; - let mut test = test_codex().build_with_auto_env(&server).await?; + let mut test = test_codex() + .with_config(move |config| { + config + .features + .enable(codex_features::Feature::RuntimeMetrics) + .expect("runtime metrics feature can be configured"); + config + .features + .set_enabled( + codex_features::Feature::CriticalPathTiming, + critical_path_timing, + ) + .expect("critical path capture can be configured independently"); + }) + .build_with_auto_env(&server) + .await?; let thread = test .thread_manager .start_thread(StartThreadOptions { @@ -152,6 +175,10 @@ async fn steer_joins_existing_turn_without_extending_sampling( unreachable!(); }; first.single_request(); + let timing_start = std::time::Instant::now(); + let observe_timing = codex_core::test_support::observe_turn_timing(&test.codex) + .await + .expect("active turn"); // The tool is still blocked. Sampling must already be exportable, even though // the dispatch future and its trace remain alive until the host responds. let sampling = wait_for_span(&exported, |span| span.name == "codex.sampling").await?; @@ -167,6 +194,7 @@ async fn steer_joins_existing_turn_without_extending_sampling( turn_id: call.turn_id.clone() } ); + let mut interrupted_capture = None; match finish { Finish::Reply => { let reply = mount_sse_once( @@ -204,6 +232,8 @@ async fn steer_joins_existing_turn_without_extending_sampling( matches!(event, EventMsg::TurnAborted(_)) }) .await; + let closed_at = std::time::Instant::now(); + interrupted_capture = Some((closed_at, observe_timing(timing_start, closed_at))); let reply = mount_sse_once( &server, sse(vec![ @@ -253,6 +283,19 @@ async fn steer_joins_existing_turn_without_extending_sampling( ); } } + if let Some((closed_at, captured)) = interrupted_capture { + // Recovery reuses a turn ID, but must not mutate the closed capture. + assert_eq!(observe_timing(timing_start, closed_at), captured); + } + let timing_end = std::time::Instant::now(); + let totals = observe_timing(timing_start, timing_end); + if critical_path_timing { + let [_, review, tools, _, unknown] = totals.expect("captured direct-tool path"); + assert!(tools > Duration::ZERO); + assert_eq!((review, unknown), (Duration::ZERO, Duration::ZERO)); + } else { + assert_eq!(totals, None); + } test.codex.submit(Op::Shutdown).await?; wait_for_event(&test.codex, |event| { matches!(event, EventMsg::ShutdownComplete) diff --git a/codex-rs/direct-tool-timing/Cargo.toml b/codex-rs/direct-tool-timing/Cargo.toml index 82a8aadac3..4b5306b1ca 100644 --- a/codex-rs/direct-tool-timing/Cargo.toml +++ b/codex-rs/direct-tool-timing/Cargo.toml @@ -15,3 +15,6 @@ codex-utils-critical-path = { workspace = true } [dev-dependencies] pretty_assertions = { workspace = true } +test-case = "3.3.1" +tokio = { workspace = true, features = ["macros", "rt", "sync"] } +tokio-util = { workspace = true } diff --git a/codex-rs/direct-tool-timing/src/lib.rs b/codex-rs/direct-tool-timing/src/lib.rs index 4b007deefb..193ff6a66a 100644 --- a/codex-rs/direct-tool-timing/src/lib.rs +++ b/codex-rs/direct-tool-timing/src/lib.rs @@ -5,3 +5,7 @@ pub use recorder::DirectToolTimings; pub use recorder::ToolBatchTiming; pub use recorder::ToolBlockingTimingGuard; pub use recorder::ToolTimingGuard; + +pub use recorder::AdmittedToolTiming; +pub use recorder::ToolAdmission; +pub use recorder::ToolGateTiming; diff --git a/codex-rs/direct-tool-timing/src/recorder.rs b/codex-rs/direct-tool-timing/src/recorder.rs index 5ae32e768a..22af5b8d9e 100644 --- a/codex-rs/direct-tool-timing/src/recorder.rs +++ b/codex-rs/direct-tool-timing/src/recorder.rs @@ -3,13 +3,19 @@ //! Prefixes and dependency paths can contain any work kind; this recorder adds tool and //! review and human-wait intervals, or Unknown for a dependency wait whose history is unavailable. //! Review callbacks must retain their invocation guard and bracket blocking automated -//! reviews with the review guard's lifetime. Execution-lock admission timing is separate. +//! reviews with the review guard's lifetime. Runtime lookup retains weak owners; +//! execution-lock admission timing is separate. use std::collections::HashMap; use std::sync::Arc; use std::sync::Mutex; use std::time::Instant; +mod runtime; +pub use runtime::AdmittedToolTiming; +pub use runtime::ToolAdmission; +pub use runtime::ToolGateTiming; + use codex_utils_critical_path::CriticalPath; use codex_utils_critical_path::WorkKind; @@ -67,12 +73,19 @@ impl ToolBatchTiming { #[derive(Clone, Default)] pub struct DirectToolTimings(Arc>); +impl std::fmt::Debug for DirectToolTimings { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DirectToolTimings").finish_non_exhaustive() + } +} + #[derive(Default)] struct State { calls: HashMap>, } struct DirectToolInvocation { + runtime_owner: Option<(String, std::sync::Weak)>, identity: Arc<()>, path: Option, segment_started_at: Instant, @@ -164,6 +177,7 @@ impl DirectToolTimings { None }; calls.push(DirectToolInvocation { + runtime_owner: None, identity: Arc::clone(&identity), path, segment_started_at: started_at, @@ -191,6 +205,44 @@ impl DirectToolTimings { } } + pub fn start_turn_direct_call( + &self, + call_id: &str, + turn_id: &str, + batch: ToolBatchTiming, + started_at: Instant, + prefix: Option, + ) -> Arc { + let work = Arc::new(self.start_direct_call(call_id, batch, started_at, prefix)); + let mut state = self.state(); + if let Some(invocation) = state.calls.get_mut(call_id).and_then(|calls| { + calls + .iter_mut() + .find(|invocation| Arc::ptr_eq(&invocation.identity, &work.identity)) + }) { + invocation.runtime_owner = Some((turn_id.to_owned(), Arc::downgrade(&work))); + } + work + } + + pub fn start_review_for_turn( + &self, + call_id: &str, + turn_id: &str, + ) -> Option { + self.owner_for_turn(call_id, turn_id)?.start_review() + } + + pub fn owner_for_turn(&self, call_id: &str, turn_id: &str) -> Option> { + let mut state = self.state(); + let (owner_turn, owner) = state.unique_work(call_id)?.runtime_owner.as_ref()?; + if owner_turn != turn_id { + return None; + } + // Retain the invocation identity across registry unlock and later ID reuse. + owner.upgrade() + } + fn finish(&self, call_id: &str, identity: &Arc<()>) -> Option { let mut state = self.state(); let calls = state.calls.get_mut(call_id)?; diff --git a/codex-rs/direct-tool-timing/src/recorder/runtime.rs b/codex-rs/direct-tool-timing/src/recorder/runtime.rs new file mode 100644 index 0000000000..db63a47451 --- /dev/null +++ b/codex-rs/direct-tool-timing/src/recorder/runtime.rs @@ -0,0 +1,90 @@ +//! Records runtime execution-lock handoffs. +//! +//! Admission guards must finish before their execution lock is released, so the +//! next admitted invocation can inherit the completed dependency path. + +use super::CompletedWork; +use super::ToolTimingGuard; +use std::sync::Arc; +use std::sync::Mutex; +use std::time::Instant; + +#[derive(Clone, Copy)] +pub enum ToolAdmission { + Parallel, + Exclusive, +} + +#[derive(Clone, Default)] +pub struct ToolGateTiming(Arc>); + +#[derive(Default)] +struct GateState { + latest: Option, + writer: Option, +} + +pub struct AdmittedToolTiming { + work: Arc, + gate: ToolGateTiming, + admission: ToolAdmission, +} + +impl ToolGateTiming { + pub fn admit( + &self, + work: Arc, + admission: ToolAdmission, + started_at: Instant, + ) -> AdmittedToolTiming { + let state = self + .0 + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let dependency = match admission { + ToolAdmission::Parallel => &state.writer, + ToolAdmission::Exclusive => &state.latest, + }; + if let Some(dependency) = dependency + && dependency.ended_at > started_at + { + work.timings.join( + &work.call_id, + started_at, + Instant::now(), + dependency.path.clone(), + ); + } + AdmittedToolTiming { + work, + gate: self.clone(), + admission, + } + } +} + +impl Drop for AdmittedToolTiming { + fn drop(&mut self) { + if let Some(completed) = self.work.finish() { + let mut state = self + .gate + .0 + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if matches!(self.admission, ToolAdmission::Exclusive) { + state.writer = Some(completed.clone()); + } + if state + .latest + .as_ref() + .is_none_or(|previous| previous.ended_at < completed.ended_at) + { + state.latest = Some(completed); + } + } + } +} + +#[cfg(test)] +#[path = "runtime_tests.rs"] +mod tests; diff --git a/codex-rs/direct-tool-timing/src/recorder/runtime_tests.rs b/codex-rs/direct-tool-timing/src/recorder/runtime_tests.rs new file mode 100644 index 0000000000..e4571318b9 --- /dev/null +++ b/codex-rs/direct-tool-timing/src/recorder/runtime_tests.rs @@ -0,0 +1,99 @@ +#![cfg(test)] +//! Verifies dependency publication at real execution-lock handoffs. + +use super::super::DirectToolTimings; +use super::super::ToolBatchTiming; +use super::*; +use codex_utils_critical_path::CriticalPath; +use codex_utils_critical_path::PathDurations; +use codex_utils_critical_path::WorkKind; +use pretty_assertions::assert_eq; +use std::future::Future; +use std::task::Context; +use std::task::Waker; +use std::time::Duration; +use tokio::sync::RwLock; +use tokio_util::either::Either; + +#[test_case::test_case(ToolAdmission::Parallel; "parallel_after_exclusive")] +#[test_case::test_case(ToolAdmission::Exclusive; "exclusive_after_parallel")] +#[tokio::test] +async fn lock_handoff_publishes_review_history_before_admitting_waiter(admission: ToolAdmission) { + let start = Instant::now() - Duration::from_secs(10); + let tool_start = start + Duration::from_secs(1); + let shared_initial_path = CriticalPath::default().work(start, tool_start, WorkKind::Inference); + let timings = DirectToolTimings::default(); + let batch = ToolBatchTiming::default(); + let gate = ToolGateTiming::default(); + let lock = RwLock::new(()); + let (releasing_lock, releasing_admission) = match admission { + ToolAdmission::Parallel => (Either::Left(lock.write().await), ToolAdmission::Exclusive), + ToolAdmission::Exclusive => (Either::Right(lock.read().await), ToolAdmission::Parallel), + }; + let releasing_work = timings.start_turn_direct_call( + "releasing", + "turn", + batch.clone(), + tool_start, + Some(shared_initial_path.clone()), + ); + let releasing_timing = gate.admit(releasing_work, releasing_admission, tool_start); + // Fixed historical boundaries keep attribution exact without sleeping for a review. + { + let mut state = timings.state(); + let work = state.unique_work("releasing").unwrap(); + work.record_until(start + Duration::from_secs(2)); + work.active_auto_reviews = 1; + work.record_until(start + Duration::from_secs(5)); + work.active_auto_reviews = 0; + } + let waiting_work = timings.start_turn_direct_call( + "waiting", + "turn", + batch.clone(), + tool_start, + Some(shared_initial_path), + ); + let acquire_waiting_lock = async { + match admission { + ToolAdmission::Parallel => Either::Left(lock.read().await), + ToolAdmission::Exclusive => Either::Right(lock.write().await), + } + }; + tokio::pin!(acquire_waiting_lock); + let mut context = Context::from_waker(Waker::noop()); + assert!( + acquire_waiting_lock + .as_mut() + .poll(&mut context) + .is_pending() + ); + + // Runtime cleanup must publish the dependency while it still holds the execution lock. + drop(releasing_timing); + assert!(batch.completed().is_none()); + assert!( + acquire_waiting_lock + .as_mut() + .poll(&mut context) + .is_pending() + ); + drop(releasing_lock); + let waiting_lock = acquire_waiting_lock.await; + let waiting_timing = gate.admit(waiting_work, admission, tool_start); + assert!(batch.completed().is_none()); + drop(waiting_timing); + drop(waiting_lock); + + let completed = batch.completed().unwrap(); + assert_eq!( + completed.path.unwrap().durations(start, completed.ended_at), + PathDurations { + inference: Duration::from_secs(1), + auto_review: Duration::from_secs(3), + tools: completed.ended_at.duration_since(start) - Duration::from_secs(4), + ..PathDurations::default() + } + ); + assert!(timings.state().calls.is_empty()); +} diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index 294f2994d4..19cc2c1c7d 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -203,6 +203,8 @@ pub enum Feature { CwdRelativeTurnDiffs, /// Enable runtime metrics snapshots via a manual reader. RuntimeMetrics, + /// Capture internal direct-tool dependency paths for turn attribution. + CriticalPathTiming, /// Enable startup memory extraction and file-backed memory consolidation. MemoryTool, /// Enable importing project-scoped memory from external agents. @@ -1326,6 +1328,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::UnderDevelopment, default_enabled: false, }, + FeatureSpec { + id: Feature::CriticalPathTiming, + key: "critical_path_timing", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::Sqlite, key: "sqlite",