Move timing start bookkeeping into helpers

codex-rs/analytics/src/client.rs

@@ -25,6 +25,7 @@ use crate::facts::TrackEventsContext;
 use crate::facts::TurnResolvedConfigFact;
 use crate::facts::TurnTimingBreakdownFact;
 use crate::facts::TurnTokenUsageFact;
+use crate::now_unix_seconds;
 use crate::reducer::AnalyticsReducer;
 use codex_app_server_protocol::ClientRequest;
 use codex_app_server_protocol::ClientResponsePayload;
@@ -43,12 +44,36 @@ use std::collections::HashSet;
 use std::sync::Arc;
 use std::sync::Mutex;
 use std::time::Duration;
+use std::time::Instant;
 use tokio::sync::mpsc;
 
 const ANALYTICS_EVENTS_QUEUE_SIZE: usize = 256;
 const ANALYTICS_EVENTS_TIMEOUT: Duration = Duration::from_secs(10);
 const ANALYTICS_EVENT_DEDUPE_MAX_KEYS: usize = 4096;
 
+#[derive(Clone, Copy, Debug)]
+pub struct StartedTimer {
+    started_at: Instant,
+}
+
+impl StartedTimer {
+    #[must_use]
+    pub fn start() -> Self {
+        Self {
+            started_at: Instant::now(),
+        }
+    }
+
+    #[must_use]
+    pub fn elapsed(self) -> Duration {
+        self.started_at.elapsed()
+    }
+
+    fn elapsed_ms(self) -> u64 {
+        self.elapsed().as_millis().try_into().unwrap_or(u64::MAX)
+    }
+}
+
 #[derive(Clone)]
 pub(crate) struct AnalyticsEventsQueue {
     pub(crate) sender: mpsc::Sender<AnalyticsFact>,
@@ -171,15 +196,14 @@ impl AnalyticsEventsClient {
     pub fn track_app_server_started(
         &self,
         rpc_transport: AppServerRpcTransport,
-        startup_duration_ms: u64,
-        completed_at: u64,
+        timer: StartedTimer,
     ) {
         self.record_fact(AnalyticsFact::Custom(
             CustomAnalyticsFact::AppServerStarted(AppServerStartedInput {
                 runtime: current_runtime_metadata(),
                 rpc_transport,
-                startup_duration_ms,
-                completed_at,
+                startup_duration_ms: timer.elapsed_ms(),
+                completed_at: now_unix_seconds(),
             }),
         ));
     }
@@ -190,11 +214,11 @@ impl AnalyticsEventsClient {
         ));
     }
 
-    pub fn track_thread_start_timing(&self, thread_id: String, duration_ms: u64) {
+    pub fn track_thread_start_timing(&self, thread_id: String, timer: StartedTimer) {
         self.record_fact(AnalyticsFact::Custom(
             CustomAnalyticsFact::ThreadStartTiming(ThreadStartTimingFact {
                 thread_id,
-                duration_ms,
+                duration_ms: timer.elapsed_ms(),
             }),
         ));
     }

codex-rs/analytics/src/lib.rs

@@ -10,6 +10,7 @@ use std::time::UNIX_EPOCH;
 pub use accepted_lines::accepted_line_fingerprints_from_unified_diff;
 pub use accepted_lines::fingerprint_hash;
 pub use client::AnalyticsEventsClient;
+pub use client::StartedTimer;
 pub use events::AppServerRpcTransport;
 pub use events::GuardianApprovalRequestSource;
 pub use events::GuardianReviewAnalyticsResult;

codex-rs/app-server/src/lib.rs

@@ -17,7 +17,6 @@ use std::io::Result as IoResult;
 use std::sync::Arc;
 use std::sync::RwLock;
 use std::sync::atomic::AtomicBool;
-use std::time::Instant;
 
 use crate::analytics_utils::analytics_events_client_from_config;
 use crate::config_manager::ConfigManager;
@@ -42,7 +41,7 @@ use crate::transport::start_remote_control;
 use crate::transport::start_stdio_connection;
 use crate::transport::start_websocket_acceptor;
 use codex_analytics::AppServerRpcTransport;
-use codex_analytics::now_unix_seconds;
+use codex_analytics::StartedTimer;
 use codex_app_server_protocol::ConfigLayerSource;
 use codex_app_server_protocol::ConfigWarningNotification;
 use codex_app_server_protocol::JSONRPCMessage;
@@ -430,7 +429,7 @@ pub async fn run_main_with_transport_options(
     auth: AppServerWebsocketAuthSettings,
     runtime_options: AppServerRuntimeOptions,
 ) -> IoResult<()> {
-    let app_server_start_started_at = Instant::now();
+    let app_server_start_timer = StartedTimer::start();
     let (transport_event_tx, mut transport_event_rx) =
         mpsc::channel::<TransportEvent>(CHANNEL_CAPACITY);
     let (outgoing_tx, mut outgoing_rx) = mpsc::channel::<OutgoingEnvelope>(CHANNEL_CAPACITY);
@@ -1070,15 +1069,7 @@ pub async fn run_main_with_transport_options(
             info!("processor task exited (channel closed)");
         }
     });
-    analytics_events_client.track_app_server_started(
-        analytics_transport,
-        app_server_start_started_at
-            .elapsed()
-            .as_millis()
-            .try_into()
-            .unwrap_or(u64::MAX),
-        now_unix_seconds(),
-    );
+    analytics_events_client.track_app_server_started(analytics_transport, app_server_start_timer);
 
     drop(transport_event_tx);
 

codex-rs/app-server/src/request_processors/thread_processor.rs

@@ -1,5 +1,6 @@
 use super::*;
 use crate::error_code::method_not_found;
+use codex_analytics::StartedTimer;
 use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_DANGER_FULL_ACCESS;
 use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_WORKSPACE;
 
@@ -989,7 +990,7 @@ impl ThreadRequestProcessor {
         experimental_raw_events: bool,
         request_trace: Option<W3cTraceContext>,
     ) -> Result<(), JSONRPCErrorError> {
-        let thread_start_started_at = std::time::Instant::now();
+        let thread_start_timer = StartedTimer::start();
         let requested_cwd = typesafe_overrides.cwd.clone();
         let mut config = config_manager
             .load_with_overrides(config_overrides.clone(), typesafe_overrides.clone())
@@ -1208,14 +1209,7 @@ impl ThreadRequestProcessor {
             active_permission_profile,
             reasoning_effort: config_snapshot.reasoning_effort,
         };
-        analytics_events_client.track_thread_start_timing(
-            thread.id.clone(),
-            thread_start_started_at
-                .elapsed()
-                .as_millis()
-                .try_into()
-                .unwrap_or(u64::MAX),
-        );
+        analytics_events_client.track_thread_start_timing(thread.id.clone(), thread_start_timer);
         let notif = thread_started_notification(thread);
         listener_task_context
             .outgoing
@@ -1236,7 +1230,7 @@ impl ThreadRequestProcessor {
             .await;
         session_telemetry.record_startup_phase(
             "thread_start_total",
-            thread_start_started_at.elapsed(),
+            thread_start_timer.elapsed(),
             Some("ready"),
         );
         Ok(())

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

@@ -2,7 +2,6 @@ use std::collections::HashMap;
 use std::collections::HashSet;
 use std::sync::Arc;
 use std::sync::atomic::Ordering;
-use std::time::Instant;
 
 use crate::SkillInjections;
 use crate::build_skill_injections;
@@ -1738,7 +1737,7 @@ async fn try_run_sampling_request(
         .turn_timing_state
         .mark_model_request_started()
         .await;
-    let sampling_started_at = Instant::now();
+    turn_context.turn_timing_state.mark_sampling_started().await;
     let stream_result = client_session
         .stream(
             prompt,
@@ -1758,14 +1757,14 @@ async fn try_run_sampling_request(
         Ok(Err(err)) => {
             turn_context
                 .turn_timing_state
-                .record_sampling_duration(sampling_started_at.elapsed())
+                .mark_sampling_completed()
                 .await;
             return Err(err);
         }
         Err(codex_async_utils::CancelErr::Cancelled) => {
             turn_context
                 .turn_timing_state
-                .record_sampling_duration(sampling_started_at.elapsed())
+                .mark_sampling_completed()
                 .await;
             return Err(CodexErr::TurnAborted);
         }
@@ -2176,7 +2175,7 @@ async fn try_run_sampling_request(
     .await;
     turn_context
         .turn_timing_state
-        .record_sampling_duration(sampling_started_at.elapsed())
+        .mark_sampling_completed()
         .await;
 
     if sess
@@ -2188,12 +2187,17 @@ async fn try_run_sampling_request(
         client_session.send_response_processed(response_id).await;
     }
 
-    let blocking_tool_started_at = Instant::now();
-    drain_in_flight(&mut in_flight, sess.clone(), turn_context.clone()).await?;
     turn_context
         .turn_timing_state
-        .record_blocking_tool_critical_path_duration(blocking_tool_started_at.elapsed())
+        .mark_blocking_tool_critical_path_started()
         .await;
+    let drain_in_flight_result =
+        drain_in_flight(&mut in_flight, sess.clone(), turn_context.clone()).await;
+    turn_context
+        .turn_timing_state
+        .mark_blocking_tool_critical_path_completed()
+        .await;
+    drain_in_flight_result?;
 
     if should_emit_token_count {
         // A tool call such as request_user_input can intentionally pause the turn. Emit token

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

@@ -6,7 +6,6 @@ mod user_shell;
 
 use std::sync::Arc;
 use std::time::Duration;
-use std::time::Instant;
 
 use codex_extension_api::ExtensionData;
 use futures::future::BoxFuture;
@@ -319,11 +318,7 @@ impl Session {
         let task: Arc<dyn AnySessionTask> = Arc::new(task);
         let task_kind = task.kind();
         let span_name = task.span_name();
-        let started_at = Instant::now();
-        let turn_started_at_unix_ms = turn_context
-            .turn_timing_state
-            .mark_turn_started(started_at)
-            .await;
+        let turn_started_at_unix_ms = turn_context.turn_timing_state.mark_turn_started().await;
         turn_context
             .turn_metadata_state
             .set_turn_started_at_unix_ms(turn_started_at_unix_ms);

codex-rs/core/src/turn_timing.rs

@@ -55,12 +55,15 @@ struct TurnTimingStateInner {
     first_token_at: Option<Instant>,
     first_message_at: Option<Instant>,
     first_request_started_at: Option<Instant>,
+    sampling_started_at: Option<Instant>,
     sampling_duration: Duration,
+    blocking_tool_critical_path_started_at: Option<Instant>,
     blocking_tool_critical_path_duration: Duration,
 }
 
 impl TurnTimingState {
-    pub(crate) async fn mark_turn_started(&self, started_at: Instant) -> i64 {
+    pub(crate) async fn mark_turn_started(&self) -> i64 {
+        let started_at = Instant::now();
         let started_at_unix_ms = now_unix_timestamp_ms();
         let mut state = self.state.lock().await;
         state.started_at = Some(started_at);
@@ -68,7 +71,9 @@ impl TurnTimingState {
         state.first_token_at = None;
         state.first_message_at = None;
         state.first_request_started_at = None;
+        state.sampling_started_at = None;
         state.sampling_duration = Duration::default();
+        state.blocking_tool_critical_path_started_at = None;
         state.blocking_tool_critical_path_duration = Duration::default();
         started_at_unix_ms
     }
@@ -119,16 +124,34 @@ impl TurnTimingState {
         }
     }
 
-    pub(crate) async fn record_sampling_duration(&self, duration: Duration) {
+    pub(crate) async fn mark_sampling_started(&self) {
         let mut state = self.state.lock().await;
-        state.sampling_duration = state.sampling_duration.saturating_add(duration);
+        if state.sampling_started_at.is_none() {
+            state.sampling_started_at = Some(Instant::now());
+        }
     }
 
-    pub(crate) async fn record_blocking_tool_critical_path_duration(&self, duration: Duration) {
+    pub(crate) async fn mark_sampling_completed(&self) {
         let mut state = self.state.lock().await;
-        state.blocking_tool_critical_path_duration = state
-            .blocking_tool_critical_path_duration
-            .saturating_add(duration);
+        if let Some(started_at) = state.sampling_started_at.take() {
+            state.sampling_duration = state.sampling_duration.saturating_add(started_at.elapsed());
+        }
+    }
+
+    pub(crate) async fn mark_blocking_tool_critical_path_started(&self) {
+        let mut state = self.state.lock().await;
+        if state.blocking_tool_critical_path_started_at.is_none() {
+            state.blocking_tool_critical_path_started_at = Some(Instant::now());
+        }
+    }
+
+    pub(crate) async fn mark_blocking_tool_critical_path_completed(&self) {
+        let mut state = self.state.lock().await;
+        if let Some(started_at) = state.blocking_tool_critical_path_started_at.take() {
+            state.blocking_tool_critical_path_duration = state
+                .blocking_tool_critical_path_duration
+                .saturating_add(started_at.elapsed());
+        }
     }
 
     pub(crate) async fn timing_breakdown(&self) -> TurnTimingBreakdown {

codex-rs/core/src/turn_timing_tests.rs

@@ -4,9 +4,10 @@ use codex_protocol::models::ContentItem;
 use codex_protocol::models::FunctionCallOutputPayload;
 use codex_protocol::models::ResponseItem;
 use pretty_assertions::assert_eq;
-use std::time::Instant;
+use std::time::Duration;
 use std::time::SystemTime;
 use std::time::UNIX_EPOCH;
+use tokio::time::sleep;
 
 use super::TurnTimingState;
 use super::response_item_records_turn_ttft;
@@ -22,7 +23,7 @@ async fn turn_timing_state_records_ttft_only_once_per_turn() {
         None
     );
 
-    state.mark_turn_started(Instant::now()).await;
+    state.mark_turn_started().await;
     assert_eq!(
         state
             .record_ttft_for_response_event(&ResponseEvent::Created)
@@ -46,7 +47,7 @@ async fn turn_timing_state_records_ttft_only_once_per_turn() {
 #[tokio::test]
 async fn turn_timing_state_records_ttfm_independently_of_ttft() {
     let state = TurnTimingState::default();
-    state.mark_turn_started(Instant::now()).await;
+    state.mark_turn_started().await;
 
     assert!(
         state
@@ -86,7 +87,7 @@ async fn turn_timing_state_records_turn_started_epoch_millis() {
         .expect("system time should be after unix epoch")
         .as_millis();
 
-    let started_at_unix_ms = state.mark_turn_started(Instant::now()).await;
+    let started_at_unix_ms = state.mark_turn_started().await;
 
     let after = SystemTime::now()
         .duration_since(UNIX_EPOCH)
@@ -102,22 +103,22 @@ async fn turn_timing_state_records_turn_started_epoch_millis() {
 #[tokio::test]
 async fn turn_timing_state_tracks_request_start_and_duration_breakdown() {
     let state = TurnTimingState::default();
-    state.mark_turn_started(Instant::now()).await;
+    state.mark_turn_started().await;
     state.mark_model_request_started().await;
-    state
-        .record_sampling_duration(std::time::Duration::from_millis(120))
-        .await;
-    state
-        .record_sampling_duration(std::time::Duration::from_millis(30))
-        .await;
-    state
-        .record_blocking_tool_critical_path_duration(std::time::Duration::from_millis(45))
-        .await;
+    state.mark_sampling_started().await;
+    sleep(Duration::from_millis(10)).await;
+    state.mark_sampling_completed().await;
+    state.mark_sampling_started().await;
+    sleep(Duration::from_millis(5)).await;
+    state.mark_sampling_completed().await;
+    state.mark_blocking_tool_critical_path_started().await;
+    sleep(Duration::from_millis(5)).await;
+    state.mark_blocking_tool_critical_path_completed().await;
 
     let breakdown = state.timing_breakdown().await;
 
-    assert_eq!(breakdown.sampling_duration_ms, 150);
-    assert_eq!(breakdown.blocking_tool_critical_path_duration_ms, 45);
+    assert!(breakdown.sampling_duration_ms >= 15);
+    assert!(breakdown.blocking_tool_critical_path_duration_ms >= 5);
     assert!(breakdown.request_start_delay_ms.is_some());
 }