Codex Handbook
core/src/turn_timing.rs 391 lines
use std::sync::Arc;use std::sync::Mutex as StdMutex;use std::time::Duration;use std::time::Instant;use std::time::SystemTime;use std::time::UNIX_EPOCH;use codex_analytics::TurnProfile;use codex_otel::TURN_TTFM_DURATION_METRIC;use codex_protocol::items::TurnItem;use codex_protocol::models::ResponseItem;use tokio::sync::Mutex;use crate::ResponseEvent;use crate::session::turn_context::TurnContext;use crate::stream_events_utils::raw_assistant_output_text_from_item;pub(crate) async fn record_turn_ttft_metric(turn_context: &TurnContext, event: &ResponseEvent) {    let Some(duration) = turn_context        .turn_timing_state        .record_ttft_for_response_event(event)        .await    else {        return;    };    turn_context.session_telemetry.record_turn_ttft(duration);}pub(crate) async fn record_turn_ttfm_metric(turn_context: &TurnContext, item: &TurnItem) {    let Some(duration) = turn_context        .turn_timing_state        .record_ttfm_for_turn_item(item)        .await    else {        return;    };    turn_context        .session_telemetry        .record_duration(TURN_TTFM_DURATION_METRIC, duration, &[]);}#[derive(Debug, Default)]pub(crate) struct TurnTimingState {    state: Mutex<TurnTimingStateInner>,    profile: StdMutex<TurnProfileState>,}#[derive(Debug, Default)]struct TurnTimingStateInner {    started_at: Option<Instant>,    started_at_unix_secs: Option<i64>,    first_token_at: Option<Instant>,    first_message_at: Option<Instant>,}#[derive(Debug, Default)]struct TurnProfileState {    started_at: Option<Instant>,    last_transition_at: Option<Instant>,    active_phase: Option<TurnProfilePhase>,    seen_sampling: bool,    before_first_sampling: Duration,    sampling: Duration,    between_sampling_overhead: Duration,    tool_blocking: Duration,    pending_idle_after_sampling: Duration,    sampling_request_count: u32,    sampling_retry_count: u32,    completed_profile: Option<TurnProfile>,}#[derive(Clone, Copy, Debug, PartialEq, Eq)]enum TurnProfilePhase {    Sampling,    ToolBlocking,}#[must_use]pub(crate) struct TurnProfileTimingGuard {    timing: Arc<TurnTimingState>,    phase: TurnProfilePhase,    active: bool,}impl TurnTimingState {    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;        state.started_at = Some(started_at);        state.started_at_unix_secs = Some(started_at_unix_ms / 1000);        state.first_token_at = None;        state.first_message_at = None;        self.profile_state().start(started_at);        started_at_unix_ms    }    pub(crate) async fn started_at_unix_secs(&self) -> Option<i64> {        self.state.lock().await.started_at_unix_secs    }    pub(crate) async fn completed_at_and_duration_ms(&self) -> (Option<i64>, Option<i64>) {        let state = self.state.lock().await;        let completed_at = Some(now_unix_timestamp_secs());        let duration_ms = state            .started_at            .map(|started_at| i64::try_from(started_at.elapsed().as_millis()).unwrap_or(i64::MAX));        (completed_at, duration_ms)    }    pub(crate) async fn time_to_first_token_ms(&self) -> Option<i64> {        let state = self.state.lock().await;        state            .time_to_first_token()            .map(|duration| i64::try_from(duration.as_millis()).unwrap_or(i64::MAX))    }    pub(crate) fn complete_profile(&self) -> TurnProfile {        self.profile_state().complete(Instant::now())    }    pub(crate) fn begin_sampling(self: &Arc<Self>) -> TurnProfileTimingGuard {        let active = self.profile_state().begin_sampling(Instant::now());        TurnProfileTimingGuard {            timing: Arc::clone(self),            phase: TurnProfilePhase::Sampling,            active,        }    }    pub(crate) fn record_sampling_retry(&self) {        self.profile_state().record_sampling_retry();    }    pub(crate) fn begin_tool_blocking(self: &Arc<Self>) -> TurnProfileTimingGuard {        let active = self.profile_state().begin_tool_blocking(Instant::now());        TurnProfileTimingGuard {            timing: Arc::clone(self),            phase: TurnProfilePhase::ToolBlocking,            active,        }    }    pub(crate) async fn record_ttft_for_response_event(        &self,        event: &ResponseEvent,    ) -> Option<Duration> {        if !response_event_records_turn_ttft(event) {            return None;        }        let mut state = self.state.lock().await;        state.record_turn_ttft()    }    pub(crate) async fn record_ttfm_for_turn_item(&self, item: &TurnItem) -> Option<Duration> {        if !matches!(item, TurnItem::AgentMessage(_)) {            return None;        }        let mut state = self.state.lock().await;        state.record_turn_ttfm()    }    fn profile_state(&self) -> std::sync::MutexGuard<'_, TurnProfileState> {        self.profile            .lock()            .unwrap_or_else(std::sync::PoisonError::into_inner)    }}impl Drop for TurnProfileTimingGuard {    fn drop(&mut self) {        if self.active {            self.timing                .profile_state()                .end_phase(Instant::now(), self.phase);        }    }}fn now_unix_timestamp_secs() -> i64 {    now_unix_timestamp_ms() / 1000}pub(crate) fn now_unix_timestamp_ms() -> i64 {    let duration = SystemTime::now()        .duration_since(UNIX_EPOCH)        .unwrap_or_default();    i64::try_from(duration.as_millis()).unwrap_or(i64::MAX)}fn duration_to_u64_ms(duration: Duration) -> u64 {    u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)}impl TurnProfileState {    fn start(&mut self, started_at: Instant) {        *self = Self {            started_at: Some(started_at),            last_transition_at: Some(started_at),            ..Self::default()        };    }    fn begin_sampling(&mut self, now: Instant) -> bool {        if self.completed_profile.is_some()            || self.started_at.is_none()            || self.active_phase.is_some()        {            return false;        }        self.advance(now);        if self.seen_sampling {            self.between_sampling_overhead += std::mem::take(&mut self.pending_idle_after_sampling);        }        self.seen_sampling = true;        self.active_phase = Some(TurnProfilePhase::Sampling);        self.sampling_request_count = self.sampling_request_count.saturating_add(1);        true    }    fn record_sampling_retry(&mut self) {        if self.completed_profile.is_none() && self.started_at.is_some() {            self.sampling_retry_count = self.sampling_retry_count.saturating_add(1);        }    }    fn begin_tool_blocking(&mut self, now: Instant) -> bool {        if self.completed_profile.is_some()            || self.started_at.is_none()            || self.active_phase.is_some()        {            return false;        }        self.advance(now);        self.active_phase = Some(TurnProfilePhase::ToolBlocking);        true    }    fn end_phase(&mut self, now: Instant, phase: TurnProfilePhase) {        if self.completed_profile.is_some() || self.active_phase != Some(phase) {            return;        }        self.advance(now);        self.active_phase = None;    }    fn advance(&mut self, now: Instant) {        let Some(previous) = self.last_transition_at.replace(now) else {            return;        };        let elapsed = now.saturating_duration_since(previous);        match self.active_phase {            Some(TurnProfilePhase::Sampling) => self.sampling += elapsed,            Some(TurnProfilePhase::ToolBlocking) => self.tool_blocking += elapsed,            None if self.seen_sampling => self.pending_idle_after_sampling += elapsed,            None => self.before_first_sampling += elapsed,        }    }    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 {            std::mem::take(&mut self.pending_idle_after_sampling)        } else {            Duration::ZERO        };        let mut profile = TurnProfile {            before_first_sampling_ms: duration_to_u64_ms(self.before_first_sampling),            sampling_ms: duration_to_u64_ms(self.sampling),            between_sampling_overhead_ms: duration_to_u64_ms(self.between_sampling_overhead),            tool_blocking_ms: duration_to_u64_ms(self.tool_blocking),            after_last_sampling_ms: duration_to_u64_ms(after_last_sampling),            sampling_request_count: self.sampling_request_count,            sampling_retry_count: self.sampling_retry_count,        };        let total_ms = self            .started_at            .map(|started_at| duration_to_u64_ms(now.saturating_duration_since(started_at)))            .unwrap_or_default();        let classified_ms = profile            .before_first_sampling_ms            .saturating_add(profile.sampling_ms)            .saturating_add(profile.between_sampling_overhead_ms)            .saturating_add(profile.tool_blocking_ms)            .saturating_add(profile.after_last_sampling_ms);        let rounding_ms = total_ms.saturating_sub(classified_ms);        match final_phase {            Some(TurnProfilePhase::Sampling) => profile.sampling_ms += rounding_ms,            Some(TurnProfilePhase::ToolBlocking) => profile.tool_blocking_ms += rounding_ms,            None if self.seen_sampling => profile.after_last_sampling_ms += rounding_ms,            None => profile.before_first_sampling_ms += rounding_ms,        }        self.active_phase = None;        self.completed_profile = Some(profile.clone());        profile    }}impl TurnTimingStateInner {    fn time_to_first_token(&self) -> Option<Duration> {        Some(self.first_token_at?.duration_since(self.started_at?))    }    fn record_turn_ttft(&mut self) -> Option<Duration> {        if self.first_token_at.is_some() {            return None;        }        self.started_at?;        self.first_token_at = Some(Instant::now());        self.time_to_first_token()    }    fn record_turn_ttfm(&mut self) -> Option<Duration> {        if self.first_message_at.is_some() {            return None;        }        let started_at = self.started_at?;        let first_message_at = Instant::now();        self.first_message_at = Some(first_message_at);        Some(first_message_at.duration_since(started_at))    }}fn response_event_records_turn_ttft(event: &ResponseEvent) -> bool {    match event {        ResponseEvent::OutputItemDone(item) | ResponseEvent::OutputItemAdded(item) => {            response_item_records_turn_ttft(item)        }        ResponseEvent::OutputTextDelta(_)        | ResponseEvent::ReasoningSummaryDelta { .. }        | ResponseEvent::ReasoningContentDelta { .. } => true,        ResponseEvent::Created        | ResponseEvent::ServerModel(_)        | ResponseEvent::ModelVerifications(_)        | ResponseEvent::TurnModerationMetadata(_)        | ResponseEvent::ServerReasoningIncluded(_)        | ResponseEvent::ToolCallInputDelta { .. }        | ResponseEvent::Completed { .. }        | ResponseEvent::ReasoningSummaryPartAdded { .. }        | ResponseEvent::RateLimits(_)        | ResponseEvent::ModelsEtag(_) => false,    }}fn response_item_records_turn_ttft(item: &ResponseItem) -> bool {    match item {        ResponseItem::Message { .. } => {            raw_assistant_output_text_from_item(item).is_some_and(|text| !text.is_empty())        }        ResponseItem::Reasoning {            summary, content, ..        } => {            summary.iter().any(|entry| match entry {                codex_protocol::models::ReasoningItemReasoningSummary::SummaryText { text } => {                    !text.is_empty()                }            }) || content.as_ref().is_some_and(|entries| {                entries.iter().any(|entry| match entry {                    codex_protocol::models::ReasoningItemContent::ReasoningText { text }                    | codex_protocol::models::ReasoningItemContent::Text { text } => {                        !text.is_empty()                    }                })            })        }        ResponseItem::AgentMessage { .. } => false,        ResponseItem::LocalShellCall { .. }        | ResponseItem::FunctionCall { .. }        | ResponseItem::CustomToolCall { .. }        | ResponseItem::ToolSearchCall { .. }        | ResponseItem::WebSearchCall { .. }        | ResponseItem::ImageGenerationCall { .. }        | ResponseItem::Compaction { .. }        | ResponseItem::ContextCompaction { .. } => true,        ResponseItem::CompactionTrigger { .. } => false,        ResponseItem::FunctionCallOutput { .. }        | ResponseItem::CustomToolCallOutput { .. }        | ResponseItem::ToolSearchOutput { .. }        | ResponseItem::Other => false,    }}#[cfg(test)]#[path = "turn_timing_tests.rs"]mod tests;