From 82fe09d23e21c642a59e217cbacf17ca9feca66a Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Fri, 22 May 2026 15:45:19 -0700 Subject: [PATCH] codex: simplify thread start timing state --- codex-rs/core/src/thread_manager.rs | 29 +++---- codex-rs/core/src/thread_start_timing.rs | 103 ++++++++++++++++------- 2 files changed, 83 insertions(+), 49 deletions(-) diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 7cadeb95ce..1de537f3b3 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -1189,11 +1189,10 @@ impl ThreadManagerState { user_shell_override: Option, ) -> CodexResult { let is_resumed_thread = matches!(&initial_history, InitialHistory::Resumed(_)); - let mut thread_start_timing = matches!( - &initial_history, - InitialHistory::New | InitialHistory::Cleared - ) - .then(ThreadStartTiming::start); + let mut thread_start_timing = match &initial_history { + InitialHistory::New | InitialHistory::Cleared => ThreadStartTiming::start(), + InitialHistory::Forked(_) | InitialHistory::Resumed(_) => ThreadStartTiming::Disabled, + }; if let InitialHistory::Resumed(resumed) = &initial_history { let mut threads = self.threads.write().await; if let Some(thread) = threads.get(&resumed.conversation_id).cloned() { @@ -1220,9 +1219,7 @@ impl ThreadManagerState { let parent_rollout_thread_trace = self .parent_rollout_thread_trace_for_source(&session_source, &initial_history) .await; - if let Some(thread_start_timing) = thread_start_timing.as_mut() { - thread_start_timing.mark_prepare_completed(); - } + thread_start_timing.mark_prepare_completed(); let tracked_session_source = session_source.clone(); let CodexSpawnOk { codex, thread_id, .. @@ -1254,19 +1251,15 @@ impl ThreadManagerState { attestation_provider: self.attestation_provider.clone(), }) .await?; - if let Some(thread_start_timing) = thread_start_timing.as_mut() { - thread_start_timing.mark_spawn_completed(); - } + thread_start_timing.mark_spawn_completed(); let new_thread = self .finalize_thread_spawn(codex, thread_id, tracked_session_source) .await?; - if let Some(mut thread_start_timing) = thread_start_timing { - thread_start_timing.mark_finalize_completed(); - if let Some(analytics_events_client) = self.analytics_events_client.as_ref() { - analytics_events_client.track_thread_start_timing( - thread_start_timing.into_fact(new_thread.thread_id.to_string()), - ); - } + thread_start_timing.mark_finalize_completed(); + if let Some(analytics_events_client) = self.analytics_events_client.as_ref() + && let Some(fact) = thread_start_timing.into_fact(new_thread.thread_id.to_string()) + { + analytics_events_client.track_thread_start_timing(fact); } if is_resumed_thread { new_thread.thread.emit_thread_resume_lifecycle().await; diff --git a/codex-rs/core/src/thread_start_timing.rs b/codex-rs/core/src/thread_start_timing.rs index 0cf5457d5a..d97f14bd92 100644 --- a/codex-rs/core/src/thread_start_timing.rs +++ b/codex-rs/core/src/thread_start_timing.rs @@ -4,16 +4,19 @@ use std::time::Instant; use codex_analytics::ThreadStartTimingFact; #[derive(Debug)] -pub(crate) struct ThreadStartTiming { - phase_started_at: Instant, - prepare_duration: Option, - spawn_duration: Option, - finalize_duration: Option, +pub(crate) enum ThreadStartTiming { + Enabled { + phase_started_at: Instant, + prepare_duration: Option, + spawn_duration: Option, + finalize_duration: Option, + }, + Disabled, } impl ThreadStartTiming { pub(crate) fn start() -> Self { - Self { + Self::Enabled { phase_started_at: Instant::now(), prepare_duration: None, spawn_duration: None, @@ -22,34 +25,72 @@ impl ThreadStartTiming { } pub(crate) fn mark_prepare_completed(&mut self) { - self.prepare_duration = Some(self.finish_phase()); - } - - pub(crate) fn mark_spawn_completed(&mut self) { - self.spawn_duration = Some(self.finish_phase()); - } - - pub(crate) fn mark_finalize_completed(&mut self) { - self.finalize_duration = Some(self.finish_phase()); - } - - pub(crate) fn into_fact(self, thread_id: String) -> ThreadStartTimingFact { - let prepare_duration = self.prepare_duration.unwrap_or_default(); - let spawn_duration = self.spawn_duration.unwrap_or_default(); - let finalize_duration = self.finalize_duration.unwrap_or_default(); - ThreadStartTimingFact { - thread_id, - duration_ms: duration_ms(prepare_duration + spawn_duration + finalize_duration), - prepare_duration_ms: duration_ms(prepare_duration), - spawn_duration_ms: duration_ms(spawn_duration), - finalize_duration_ms: duration_ms(finalize_duration), + let Some(duration) = self.finish_phase() else { + return; + }; + if let Self::Enabled { + prepare_duration, .. + } = self + { + *prepare_duration = Some(duration); } } - fn finish_phase(&mut self) -> Duration { - let duration = self.phase_started_at.elapsed(); - self.phase_started_at = Instant::now(); - duration + pub(crate) fn mark_spawn_completed(&mut self) { + let Some(duration) = self.finish_phase() else { + return; + }; + if let Self::Enabled { spawn_duration, .. } = self { + *spawn_duration = Some(duration); + } + } + + pub(crate) fn mark_finalize_completed(&mut self) { + let Some(duration) = self.finish_phase() else { + return; + }; + if let Self::Enabled { + finalize_duration, .. + } = self + { + *finalize_duration = Some(duration); + } + } + + pub(crate) fn into_fact(self, thread_id: String) -> Option { + match self { + Self::Enabled { + prepare_duration, + spawn_duration, + finalize_duration, + .. + } => { + let prepare_duration = prepare_duration.unwrap_or_default(); + let spawn_duration = spawn_duration.unwrap_or_default(); + let finalize_duration = finalize_duration.unwrap_or_default(); + Some(ThreadStartTimingFact { + thread_id, + duration_ms: duration_ms(prepare_duration + spawn_duration + finalize_duration), + prepare_duration_ms: duration_ms(prepare_duration), + spawn_duration_ms: duration_ms(spawn_duration), + finalize_duration_ms: duration_ms(finalize_duration), + }) + } + Self::Disabled => None, + } + } + + fn finish_phase(&mut self) -> Option { + match self { + Self::Enabled { + phase_started_at, .. + } => { + let duration = phase_started_at.elapsed(); + *phase_started_at = Instant::now(); + Some(duration) + } + Self::Disabled => None, + } } }