Attribute approval waits to tool critical-path timing (#57207)

## Why

Automatic reviews and user prompts previously invalidated critical-path captures even when their owning tool invocation was known.

## What changed

Record automatic review time as `auto_review` and waits for approvals, permission requests, and user input as `user_wait` when `CriticalPathTiming` is enabled. Bind timing to the owning invocation before asynchronous work begins, and use the triggering tool's identity for network approvals. Continue invalidating captures when a wait cannot be attributed.

## Testing

Add coverage for cancellation, reused call IDs, network ownership, and disabled capture. Update the automatic review integration test to verify attributed review duration.

GitOrigin-RevId: 0a1f5feebd0c4f5ed67dc058e832ec5f8818db02

codex-rs/core/src/session/elicitation_holders_tests.rs

@@ -1,6 +1,10 @@
+use codex_direct_tool_timing::ToolBatchTiming;
+use codex_direct_tool_timing::ToolTimingGuard;
 use codex_features::Feature;
+use codex_utils_critical_path::CriticalPath;
 use std::collections::HashMap;
 use std::sync::Arc;
+use std::time::Duration;
 use std::time::Instant;
 
 use codex_core_plugins::PluginCommandAttribution;
@@ -20,25 +24,32 @@ 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) {
+fn start_timing(
+    turn: &mut Arc<super::turn_context::TurnContext>,
+    enabled: bool,
+) -> (Arc<ToolTimingGuard>, Instant) {
     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;
+    let start = Instant::now();
+    let owner = turn
+        .turn_timing_state
+        .direct_tool_timings
+        .start_turn_direct_call(
+            "call-1",
+            &turn.sub_id,
+            ToolBatchTiming::default(),
+            start,
+            Some(CriticalPath::default()),
+        );
+    (owner, start)
 }
 
-fn assert_uncaptured_wait(turn: &super::turn_context::TurnContext, enabled: bool) {
-    assert_eq!(
-        turn.turn_timing_state
-            .snapshot_path(Instant::now())
-            .is_none(),
-        enabled
-    );
+fn assert_user_wait(owner: &ToolTimingGuard, start: Instant, enabled: bool) {
+    let completed = owner.finish().unwrap();
+    let totals = completed.path.unwrap().durations(start, completed.ended_at);
+    assert_eq!(totals.user_wait > Duration::ZERO, enabled);
 }
 
 async fn wait_until_held(pause_state: &mut watch::Receiver<bool>) {
@@ -55,12 +66,13 @@ 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")]
+#[test_case::test_case(true, "call-1"; "capture_on")]
+#[test_case::test_case(false, "call-1"; "capture_off")]
+#[test_case::test_case(true, "network-prompt"; "distinct_network_prompt_identity")]
 #[tokio::test]
-async fn command_approval_holds_an_elicitation_until_response(capture: bool) {
+async fn command_approval_holds_an_elicitation_until_response(capture: bool, prompt_id: &str) {
     let (session, mut turn_context, events) = make_session_and_context_with_rx().await;
-    start_timing(&mut turn_context, capture).await;
+    let (timing, timing_start) = start_timing(&mut turn_context, capture);
     *session.active_turn.lock().await = Some(ActiveTurn::default());
     let mut pause_state = session.subscribe_elicitation_pause_state();
     #[allow(deprecated)]
@@ -71,6 +83,7 @@ async fn command_approval_holds_an_elicitation_until_response(capture: bool) {
     };
 
     let request = tokio::spawn({
+        let prompt_id = prompt_id.to_owned();
         let session = session.clone();
         let turn_context = turn_context.clone();
         async move {
@@ -80,7 +93,8 @@ async fn command_approval_holds_an_elicitation_until_response(capture: bool) {
                     ExecApprovalKind::Command,
                     crate::guardian::GuardianReviewContext::from(turn_context.clone())
                         .model_context(),
-                    "call-1".to_string(),
+                    prompt_id,
+                    Some("call-1"),
                     /*approval_id*/ None,
                     /*environment_id*/ None,
                     vec!["echo".to_string()],
@@ -104,17 +118,19 @@ async fn command_approval_holds_an_elicitation_until_response(capture: bool) {
     assert_eq!(event.script_path.as_deref(), Some("scripts/run.py"));
     wait_until_held(&mut pause_state).await;
     session
-        .notify_approval("call-1", ReviewDecision::Approved)
+        .notify_approval(prompt_id, ReviewDecision::Approved)
         .await;
     request.await.expect("approval task");
     wait_until_released(&mut pause_state).await;
-    assert_uncaptured_wait(&turn_context, capture);
+    assert_user_wait(&timing, timing_start, capture);
 }
 
+#[test_case::test_case(true; "capture_on")]
+#[test_case::test_case(false; "capture_off")]
 #[tokio::test]
-async fn patch_approval_holds_an_elicitation_until_response() {
+async fn patch_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, /*enabled*/ true).await;
+    let (timing, timing_start) = start_timing(&mut turn_context, capture);
     *session.active_turn.lock().await = Some(ActiveTurn::default());
     let mut pause_state = session.subscribe_elicitation_pause_state();
 
@@ -141,13 +157,15 @@ 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);
+    assert_user_wait(&timing, timing_start, capture);
 }
 
+#[test_case::test_case(true; "capture_on")]
+#[test_case::test_case(false; "capture_off")]
 #[tokio::test]
-async fn permission_request_holds_an_elicitation_until_response() {
+async fn permission_request_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, /*enabled*/ true).await;
+    let (timing, timing_start) = start_timing(&mut turn_context, capture);
     *session.active_turn.lock().await = Some(ActiveTurn::default());
     let mut pause_state = session.subscribe_elicitation_pause_state();
 
@@ -190,16 +208,21 @@ 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);
+    assert_user_wait(&timing, timing_start, capture);
 }
 
-#[test_case::test_case(true, false; "blocking")]
-#[test_case::test_case(false, false; "nonblocking_ui_label")]
-#[test_case::test_case(true, true; "cancelled")]
+#[test_case::test_case(true, true, false; "blocking")]
+#[test_case::test_case(true, false, false; "nonblocking_ui_label")]
+#[test_case::test_case(true, true, true; "cancelled")]
+#[test_case::test_case(false, true, false; "capture_off")]
 #[tokio::test]
-async fn request_user_input_holds_an_elicitation_until_response(blocking: bool, cancel: bool) {
+async fn request_user_input_holds_an_elicitation_until_response(
+    capture: bool,
+    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;
+    let (timing, timing_start) = start_timing(&mut turn_context, capture);
     *session.active_turn.lock().await = Some(ActiveTurn::default());
     let mut pause_state = session.subscribe_elicitation_pause_state();
 
@@ -228,7 +251,7 @@ async fn request_user_input_holds_an_elicitation_until_response(blocking: bool,
         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);
+        assert_user_wait(&timing, timing_start, capture);
         return;
     }
     let response = RequestUserInputResponse {
@@ -240,5 +263,41 @@ async fn request_user_input_holds_an_elicitation_until_response(blocking: bool,
 
     request.await.expect("request user input task");
     wait_until_released(&mut pause_state).await;
-    assert_uncaptured_wait(&turn_context, /*enabled*/ true);
+    assert_user_wait(&timing, timing_start, capture);
+}
+
+#[test_case::test_case(true; "enabled_unowned_wait_is_incomplete")]
+#[test_case::test_case(false; "disabled_unowned_wait_is_unchanged")]
+#[tokio::test]
+async fn unowned_human_wait_does_not_claim_a_complete_path(capture: bool) {
+    let (session, mut turn_context, _rx) = make_session_and_context_with_rx().await;
+    let (_tool_timing, _) = start_timing(&mut turn_context, capture);
+    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()
+    );
+    assert!(
+        session
+            .user_wait_timing_owner(&turn_context, Some("unknown-call"))
+            .is_none()
+    );
+    assert!(
+        session
+            .start_user_wait_timing(&turn_context, /*owner*/ None)
+            .is_none()
+    );
+    assert_eq!(
+        turn_context
+            .turn_timing_state
+            .snapshot_path(Instant::now())
+            .is_none(),
+        capture
+    );
 }

codex-rs/core/src/session/mod.rs

@@ -2998,15 +2998,42 @@ 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
+    fn user_wait_timing_owner(
+        &self,
+        turn_context: &TurnContext,
+        call_id: Option<&str>,
+    ) -> Option<Arc<codex_direct_tool_timing::ToolTimingGuard>> {
+        if !turn_context
             .config
             .features
             .enabled(Feature::CriticalPathTiming)
+        {
+            return None;
+        }
+        // Retain the invocation identity before request preparation can suspend.
+        call_id.and_then(|call_id| {
+            turn_context
+                .turn_timing_state
+                .direct_tool_timings
+                .owner_for_turn(call_id, &turn_context.sub_id)
+        })
+    }
+
+    fn start_user_wait_timing(
+        &self,
+        turn_context: &TurnContext,
+        owner: Option<&codex_direct_tool_timing::ToolTimingGuard>,
+    ) -> Option<codex_direct_tool_timing::ToolBlockingTimingGuard> {
+        let timing = owner.and_then(codex_direct_tool_timing::ToolTimingGuard::start_user_wait);
+        if timing.is_none()
+            && turn_context
+                .config
+                .features
+                .enabled(Feature::CriticalPathTiming)
         {
             turn_context.turn_timing_state.invalidate_critical_path();
         }
+        timing
     }
 
     /// Emit an exec approval request event and await the user's decision.
@@ -3030,6 +3057,7 @@ impl Session {
         kind: ExecApprovalKind,
         model_context: ModelInvocationContext,
         call_id: String,
+        timing_call_id: Option<&str>,
         approval_id: Option<String>,
         environment_id: Option<String>,
         command: Vec<String>,
@@ -3041,6 +3069,7 @@ impl Session {
         available_decisions: Option<Vec<ReviewDecision>>,
         plugin_attribution_override: Option<PluginCommandAttribution>,
     ) -> ReviewDecision {
+        let user_wait_owner = self.user_wait_timing_owner(turn_context, timing_call_id);
         let _elicitation = self.services.elicitations.register();
         //  command-level approvals use `call_id`.
         // `approval_id` identifies subcommand callbacks and stdin writes.
@@ -3111,7 +3140,7 @@ impl Session {
             available_decisions: Some(available_decisions),
             parsed_cmd,
         });
-        self.invalidate_uncaptured_user_wait(turn_context);
+        let _user_wait = self.start_user_wait_timing(turn_context, user_wait_owner.as_deref());
         self.send_event(turn_context, event).await;
         rx_approve.await.unwrap_or(ReviewDecision::Abort)
     }
@@ -3128,6 +3157,7 @@ impl Session {
         reason: Option<String>,
         grant_root: Option<PathBuf>,
     ) -> ReviewDecision {
+        let user_wait_owner = self.user_wait_timing_owner(turn_context, Some(&call_id));
         let _elicitation = self.services.elicitations.register();
         // Add the tx_approve callback to the map before sending the request.
         let (tx_approve, rx_approve) = oneshot::channel();
@@ -3154,7 +3184,7 @@ impl Session {
             reason,
             grant_root,
         });
-        self.invalidate_uncaptured_user_wait(turn_context);
+        let _user_wait = self.start_user_wait_timing(turn_context, user_wait_owner.as_deref());
         self.send_event(turn_context, event).await;
         rx_approve.await.unwrap_or(ReviewDecision::Abort)
     }
@@ -3172,6 +3202,7 @@ impl Session {
         cancellation_token: CancellationToken,
     ) -> Option<RequestPermissionsResponse> {
         let turn_context = &step_context.turn;
+        let user_wait_owner = self.user_wait_timing_owner(turn_context, Some(&call_id));
         let approval_policy = step_context.settings.approval_policy();
         let Some(environment) = step_context
             .environments
@@ -3320,11 +3351,12 @@ impl Session {
             permissions: requested_permissions,
             cwd: Some(cwd.clone().into()),
         });
-        self.invalidate_uncaptured_user_wait(turn_context);
+        let _user_wait = self.start_user_wait_timing(turn_context, user_wait_owner.as_deref());
         self.send_event(turn_context.as_ref(), event).await;
         tokio::select! {
             biased;
             _ = cancellation_token.cancelled() => {
+                drop(_user_wait);
                 let mut active = self.active_turn.lock().await;
                 if let Some(at) = active.as_mut() {
                     let mut ts = at.turn_state.lock().await;
@@ -3346,6 +3378,7 @@ impl Session {
         call_id: String,
         args: RequestUserInputArgs,
     ) -> Option<AcceptedUserInputResponse> {
+        let user_wait_owner = self.user_wait_timing_owner(turn_context, Some(&call_id));
         let _elicitation = self.services.elicitations.register();
         let sub_id = turn_context.sub_id.clone();
         let (tx_response, rx_response) = oneshot::channel();
@@ -3374,7 +3407,7 @@ impl Session {
         turn_context
             .turn_metadata_state
             .mark_user_input_requested_during_turn();
-        self.invalidate_uncaptured_user_wait(turn_context);
+        let _user_wait = self.start_user_wait_timing(turn_context, user_wait_owner.as_deref());
         self.send_event(turn_context, event).await;
         rx_response.await.ok()
     }