openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-27 18:48:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

Capture direct-tool critical paths behind a feature flag (#57201)

Browse Source 
## What changed

Add the default-disabled `critical_path_timing` feature to capture dependency paths for root-agent turns. Connect direct tool dispatch, execution-lock handoffs, and completed tool batches to turn timing so overlapping work shares its common history without double counting.

Omit captures for unsupported dependencies, including CodeMode, agent communication, approval review, and user waits. Settle aborted tasks before freezing timing, invalidate captures on forced teardown, and apply runtime feature changes only to future turns.

## Testing

Add coverage for serial and overlapping tool dispatch, execution-lock dependency handoffs, cancellation, unsupported paths, frozen timing after completion, and configuration refresh boundaries.

GitOrigin-RevId: 71458b654b1c4a05799886928b449c948ea33109
This commit is contained in:
taohong-oai
2026-10-26 19:19:10 +00:00
committed by copyberry
parent 36021e1ad5
commit f371f9701f
31 changed files with 1436 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/Cargo.lock
Generated
+5
View File
@@ -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]]
codex-rs/Cargo.toml
+1
View File
@@ -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" }
codex-rs/core/Cargo.toml
+2
View File
@@ -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 }
codex-rs/core/config.schema.json
+6
View File
@@ -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"
},
codex-rs/core/src/guardian/decision.rs
+83 -72
View File
@@ -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<Session>,
context: impl Into<GuardianReviewContext>,
review_id: String,
request: impl Into<ReviewAction>,
reasons: ApprovalRequestReasons,
mut options: GuardianReviewOptions,
) -> Option<ReviewDecision> {
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<ReviewDecision>> {
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
}
})
}
codex-rs/core/src/guardian/decision_timing_tests.rs
+63
View File
@@ -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)
);
}
codex-rs/core/src/guardian/mod.rs
+4
View File
@@ -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;
codex-rs/core/src/session/config_refresh.rs
+11 -1
View File
@@ -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();
codex-rs/core/src/session/config_refresh_tests.rs
+91
View File
@@ -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)
);
}
codex-rs/core/src/session/elicitation_holders_tests.rs
+50 -7
View File
@@ -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<super::turn_context::TurnContext>, 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<bool>) {
pause_state
.wait_for(|paused| *paused)
@@ -32,9 +55,12 @@ async fn wait_until_released(pause_state: &mut watch::Receiver<bool>) {
.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);
}
codex-rs/core/src/session/mcp.rs
+9
View File
@@ -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 = {
codex-rs/core/src/session/mod.rs
+23
View File
@@ -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<crate::PendingAgentMessage>,
) {
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,
codex-rs/core/src/session/tests.rs
+58 -5
View File
@@ -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]
codex-rs/core/src/session/turn.rs
+3 -1
View File
@@ -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",
codex-rs/core/src/tasks/mod.rs
+6
View File
@@ -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))
codex-rs/core/src/tasks/mod_tests.rs
+70
View File
@@ -494,3 +494,73 @@ fn emit_compact_metric_records_auto_local() {
])
);
}
struct UncooperativeTimedTask {
entered: Arc<tokio::sync::Notify>,
}
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<Self>,
_session: Arc<Session>,
context: Arc<TurnContext>,
_input: Vec<TurnInput>,
_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
);
}
codex-rs/core/src/tasks/regular.rs
+8
View File
@@ -44,6 +44,14 @@ impl SessionTask for RegularTask {
input: Vec<TurnInput>,
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,
codex-rs/core/src/test_support.rs
+32
View File
@@ -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;
codex-rs/core/src/tools/parallel.rs
+146 -6
View File
@@ -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<StepContext>,
tracker: SharedTurnDiffTracker,
parallel_execution: Arc<RwLock<()>>,
timing: Option<DirectToolTiming>,
}
#[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<StepContext>,
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<ToolBatchTiming> {
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(
&registered.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<Vec<u8>> =
@@ -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()
codex-rs/core/src/turn_timing.rs
+76 -5
View File
@@ -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<TurnTimingStateInner>,
profile: StdMutex<TurnProfileState>,
}
@@ -74,6 +79,8 @@ struct TurnTimingStateInner {
#[derive(Debug, Default)]
struct TurnProfileState {
path: Option<CriticalPath>,
tool_dependency: Option<CriticalPath>,
started_at: Option<Instant>,
last_transition_at: Option<Instant>,
active_phase: Option<TurnProfilePhase>,
@@ -102,9 +109,33 @@ pub(crate) struct TurnProfileTimingGuard {
timing: Arc<TurnTimingState>,
phase: TurnProfilePhase,
active: bool,
tool_batch: Option<ToolBatchTiming>,
}
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<CriticalPath> {
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<Self>) -> TurnProfileTimingGuard {
pub(crate) fn begin_tool_blocking(
self: &Arc<Self>,
batch: Option<ToolBatchTiming>,
) -> 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<CriticalPath> {
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 {
codex-rs/core/src/turn_timing_tests.rs
+106 -1
View File
@@ -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));
}
codex-rs/core/tests/suite/auto_review.rs
+27 -3
View File
@@ -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<codex_extension_api::ApprovalDecision>> {
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()))
codex-rs/core/tests/suite/critical_path_runtime.rs
+251
View File
@@ -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<Semaphore>,
}
impl ToolContributor for GatedTool {
fn tools(
&self,
_session_store: &ExtensionData,
_thread_store: &ExtensionData,
) -> Vec<Arc<dyn for<'call> ToolExecutor<ToolCall<'call>>>> {
vec![Arc::new(self.clone())]
}
}
impl<'call> ToolExecutor<ToolCall<'call>> 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<dyn ToolOutput>,
)
})
}
}
#[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(())
}
codex-rs/core/tests/suite/mod.rs
+1
View File
@@ -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"))]
codex-rs/core/tests/suite/turn_phase_trace.rs
+47 -4
View File
@@ -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)
codex-rs/direct-tool-timing/Cargo.toml
+3
View File
@@ -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 }
codex-rs/direct-tool-timing/src/lib.rs
+4
View File
@@ -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;
codex-rs/direct-tool-timing/src/recorder.rs
+53 -1
View File
@@ -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<Mutex<State>>);
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<String, Vec<DirectToolInvocation>>,
}
struct DirectToolInvocation {
runtime_owner: Option<(String, std::sync::Weak<ToolTimingGuard>)>,
identity: Arc<()>,
path: Option<CriticalPath>,
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<CriticalPath>,
) -> Arc<ToolTimingGuard> {
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<ToolBlockingTimingGuard> {
self.owner_for_turn(call_id, turn_id)?.start_review()
}
pub fn owner_for_turn(&self, call_id: &str, turn_id: &str) -> Option<Arc<ToolTimingGuard>> {
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<CompletedWork> {
let mut state = self.state();
let calls = state.calls.get_mut(call_id)?;
codex-rs/direct-tool-timing/src/recorder/runtime.rs
+90
View File
@@ -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<Mutex<GateState>>);
#[derive(Default)]
struct GateState {
latest: Option<CompletedWork>,
writer: Option<CompletedWork>,
}
pub struct AdmittedToolTiming {
work: Arc<ToolTimingGuard>,
gate: ToolGateTiming,
admission: ToolAdmission,
}
impl ToolGateTiming {
pub fn admit(
&self,
work: Arc<ToolTimingGuard>,
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;
codex-rs/direct-tool-timing/src/recorder/runtime_tests.rs
+99
View File
@@ -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());
}
codex-rs/features/src/lib.rs
+8
View File
@@ -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",