diff --git a/codex-rs/analytics/src/analytics_client_tests.rs b/codex-rs/analytics/src/analytics_client_tests.rs index b25d1f5c83..f6968e3101 100644 --- a/codex-rs/analytics/src/analytics_client_tests.rs +++ b/codex-rs/analytics/src/analytics_client_tests.rs @@ -530,7 +530,6 @@ async fn ingest_rejected_turn_steer( response: Box::new(sample_thread_resume_response( "thread-2", /*ephemeral*/ false, "gpt-5", )), - thread_initialization: None, }, out, ) @@ -604,7 +603,6 @@ async fn ingest_turn_prerequisites( response: Box::new(sample_thread_start_response( "thread-2", /*ephemeral*/ false, "gpt-5", )), - thread_initialization: None, }, out, ) @@ -628,7 +626,6 @@ async fn ingest_turn_prerequisites( connection_id: 7, request_id: RequestId::Integer(3), response: Box::new(sample_turn_start_response("turn-2")), - thread_initialization: None, }, out, ) @@ -695,7 +692,6 @@ async fn ingest_review_prerequisites( response: Box::new(sample_thread_start_response( "thread-1", /*ephemeral*/ false, "gpt-5", )), - thread_initialization: None, }, events, ) @@ -1608,7 +1604,6 @@ async fn initialize_caches_client_and_thread_lifecycle_publishes_once_initialize /*ephemeral*/ false, "gpt-5", )), - thread_initialization: None, }, &mut events, ) @@ -1653,7 +1648,6 @@ async fn initialize_caches_client_and_thread_lifecycle_publishes_once_initialize response: Box::new(sample_thread_resume_response( "thread-1", /*ephemeral*/ true, "gpt-5", )), - thread_initialization: None, }, &mut events, ) @@ -1724,7 +1718,6 @@ async fn unrelated_client_requests_are_ignored_by_reducer() { connection_id: 7, request_id: RequestId::Integer(3), response: Box::new(sample_turn_start_response("turn-2")), - thread_initialization: None, }, &mut events, ) @@ -1750,7 +1743,6 @@ async fn unrelated_client_responses_are_ignored_by_reducer() { response: Box::new(ClientResponsePayload::ThreadArchive( ThreadArchiveResponse {}, )), - thread_initialization: None, }, &mut events, ) @@ -1809,7 +1801,6 @@ async fn compaction_event_ingests_custom_fact() { Some(AppServerThreadSource::Subagent), Some(parent_thread_id.to_string()), )), - thread_initialization: None, }, &mut events, ) @@ -1917,7 +1908,6 @@ async fn guardian_review_event_ingests_custom_fact_with_optional_target_item() { /*ephemeral*/ false, "gpt-5", )), - thread_initialization: None, }, &mut events, ) @@ -2419,7 +2409,6 @@ async fn item_review_summaries_do_not_cross_threads_with_reused_item_ids() { response: Box::new(sample_thread_start_response( "thread-2", /*ephemeral*/ false, "gpt-5", )), - thread_initialization: None, }, &mut events, ) @@ -2707,7 +2696,6 @@ async fn subagent_thread_started_inherits_parent_connection_for_new_thread() { /*ephemeral*/ false, "gpt-5", )), - thread_initialization: None, }, &mut events, ) @@ -3462,7 +3450,6 @@ async fn accepted_turn_steer_emits_expected_event() { connection_id: 7, request_id: RequestId::Integer(4), response: Box::new(sample_turn_steer_response("turn-2")), - thread_initialization: None, }, &mut out, ) @@ -3634,7 +3621,6 @@ async fn turn_start_error_response_discards_pending_start_request() { connection_id: 7, request_id: RequestId::Integer(3), response: Box::new(sample_turn_start_response("turn-2")), - thread_initialization: None, }, &mut out, ) @@ -3934,7 +3920,6 @@ async fn accepted_steers_increment_turn_steer_count() { connection_id: 7, request_id: RequestId::Integer(4), response: Box::new(sample_turn_steer_response("turn-2")), - thread_initialization: None, }, &mut out, ) @@ -3982,7 +3967,6 @@ async fn accepted_steers_increment_turn_steer_count() { connection_id: 7, request_id: RequestId::Integer(6), response: Box::new(sample_turn_steer_response("turn-2")), - thread_initialization: None, }, &mut out, ) diff --git a/codex-rs/analytics/src/client.rs b/codex-rs/analytics/src/client.rs index 31500611bf..f33cde076b 100644 --- a/codex-rs/analytics/src/client.rs +++ b/codex-rs/analytics/src/client.rs @@ -17,7 +17,7 @@ use crate::facts::PluginStateChangedInput; use crate::facts::SkillInvocation; use crate::facts::SkillInvokedInput; use crate::facts::SubAgentThreadStartedInput; -use crate::facts::ThreadInitializationFact; +use crate::facts::ThreadInitializationResponseFact; use crate::facts::TrackEventsContext; use crate::facts::TurnCodexErrorFact; use crate::facts::TurnProfileFact; @@ -318,7 +318,6 @@ impl AnalyticsEventsClient { connection_id: u64, request_id: RequestId, response: ClientResponsePayload, - thread_initialization: Option, ) { if !matches!( response, @@ -334,10 +333,13 @@ impl AnalyticsEventsClient { connection_id, request_id, response: Box::new(response), - thread_initialization, }); } + pub fn track_thread_initialization_response(&self, fact: ThreadInitializationResponseFact) { + self.record_fact(AnalyticsFact::ThreadInitializationResponse(Box::new(fact))); + } + pub fn track_error_response( &self, connection_id: u64, diff --git a/codex-rs/analytics/src/client_tests.rs b/codex-rs/analytics/src/client_tests.rs index bfccecdcfc..b7aa9f1c97 100644 --- a/codex-rs/analytics/src/client_tests.rs +++ b/codex-rs/analytics/src/client_tests.rs @@ -252,9 +252,7 @@ fn track_response_only_enqueues_analytics_relevant_responses() { (RequestId::Integer(4), sample_turn_start_response()), (RequestId::Integer(5), sample_turn_steer_response()), ] { - client.track_response( - /*connection_id*/ 7, request_id, response, /*thread_initialization*/ None, - ); + client.track_response(/*connection_id*/ 7, request_id, response); assert!(matches!( receiver.try_recv(), Ok(AnalyticsFact::ClientResponse { .. }) @@ -265,7 +263,6 @@ fn track_response_only_enqueues_analytics_relevant_responses() { /*connection_id*/ 7, RequestId::Integer(6), ClientResponsePayload::ThreadArchive(ThreadArchiveResponse {}), - /*thread_initialization*/ None, ); assert!(matches!(receiver.try_recv(), Err(TryRecvError::Empty))); } diff --git a/codex-rs/analytics/src/facts.rs b/codex-rs/analytics/src/facts.rs index 64fda3f6ea..4e8899caab 100644 --- a/codex-rs/analytics/src/facts.rs +++ b/codex-rs/analytics/src/facts.rs @@ -27,6 +27,7 @@ use codex_protocol::protocol::SkillScope; use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::TokenUsage; use codex_protocol::request_permissions::RequestPermissionsResponse; +use serde::Deserialize; use serde::Serialize; use std::path::PathBuf; @@ -98,7 +99,7 @@ pub enum ThreadInitializationMode { Resumed, } -#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize)] +#[derive(Clone, Copy, Debug, Default, Deserialize, PartialEq, Eq, Serialize)] pub struct ThreadInitializationProfile { pub duration_ms: u64, pub app_server_duration_ms: u64, @@ -117,11 +118,18 @@ pub struct ThreadInitializationProfile { } #[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct ThreadInitializationFact { +pub struct CompletedThreadInitialization { pub initialization_mode: ThreadInitializationMode, pub profile: ThreadInitializationProfile, } +#[derive(Clone)] +pub struct ThreadInitializationResponseFact { + pub connection_id: u64, + pub response: ClientResponsePayload, + pub initialization: CompletedThreadInitialization, +} + #[derive(Clone)] pub struct TurnTokenUsageFact { pub turn_id: String, @@ -482,8 +490,8 @@ pub(crate) enum AnalyticsFact { connection_id: u64, request_id: RequestId, response: Box, - thread_initialization: Option, }, + ThreadInitializationResponse(Box), ErrorResponse { connection_id: u64, request_id: RequestId, diff --git a/codex-rs/analytics/src/lib.rs b/codex-rs/analytics/src/lib.rs index d24089ceca..4f0cbf2d6e 100644 --- a/codex-rs/analytics/src/lib.rs +++ b/codex-rs/analytics/src/lib.rs @@ -31,14 +31,15 @@ pub use facts::CompactionReason; pub use facts::CompactionStatus; pub use facts::CompactionStrategy; pub use facts::CompactionTrigger; +pub use facts::CompletedThreadInitialization; pub use facts::HookRunFact; pub use facts::InputError; pub use facts::InvocationType; pub use facts::SkillInvocation; pub use facts::SubAgentThreadStartedInput; -pub use facts::ThreadInitializationFact; pub use facts::ThreadInitializationMode; pub use facts::ThreadInitializationProfile; +pub use facts::ThreadInitializationResponseFact; pub use facts::TrackEventsContext; pub use facts::TurnCodexErrorFact; pub use facts::TurnProfile; diff --git a/codex-rs/analytics/src/reducer.rs b/codex-rs/analytics/src/reducer.rs index 363c0b5886..b6bbdcea99 100644 --- a/codex-rs/analytics/src/reducer.rs +++ b/codex-rs/analytics/src/reducer.rs @@ -62,6 +62,7 @@ use crate::facts::AnalyticsJsonRpcError; use crate::facts::AppMentionedInput; use crate::facts::AppUsedInput; use crate::facts::CodexCompactionEvent; +use crate::facts::CompletedThreadInitialization; use crate::facts::CustomAnalyticsFact; use crate::facts::HookRunInput; use crate::facts::PluginState; @@ -69,9 +70,9 @@ use crate::facts::PluginStateChangedInput; use crate::facts::PluginUsedInput; use crate::facts::SkillInvokedInput; use crate::facts::SubAgentThreadStartedInput; -use crate::facts::ThreadInitializationFact; use crate::facts::ThreadInitializationMode; use crate::facts::ThreadInitializationProfile; +use crate::facts::ThreadInitializationResponseFact; use crate::facts::TurnCodexError; use crate::facts::TurnCodexErrorFact; use crate::facts::TurnProfile; @@ -143,6 +144,30 @@ pub(crate) struct AnalyticsReducer { item_review_summaries: HashMap, } +#[derive(Clone, Copy)] +struct ThreadInitializationEventDetails { + mode: ThreadInitializationMode, + profile: Option, +} + +impl ThreadInitializationEventDetails { + fn without_profile(mode: ThreadInitializationMode) -> Self { + Self { + mode, + profile: None, + } + } +} + +impl From for ThreadInitializationEventDetails { + fn from(initialization: CompletedThreadInitialization) -> Self { + Self { + mode: initialization.initialization_mode, + profile: Some(initialization.profile), + } + } +} + struct ConnectionState { app_server_client: CodexAppServerClientMetadata, runtime: CodexRuntimeMetadata, @@ -408,13 +433,14 @@ impl AnalyticsReducer { connection_id, request_id, response, - thread_initialization, } => { if let Some(response) = response.into_client_response(request_id) { - self.ingest_response(connection_id, response, thread_initialization, out) - .await; + self.ingest_response(connection_id, response, out).await; } } + AnalyticsFact::ThreadInitializationResponse(input) => { + self.ingest_thread_initialization_response(*input, out); + } AnalyticsFact::ErrorResponse { connection_id, request_id, @@ -764,7 +790,6 @@ impl AnalyticsReducer { &mut self, connection_id: u64, response: ClientResponse, - thread_initialization: Option, out: &mut Vec, ) { match response { @@ -773,10 +798,9 @@ impl AnalyticsReducer { connection_id, response.thread, response.model, - thread_initialization - .map(|initialization| initialization.initialization_mode) - .unwrap_or(ThreadInitializationMode::New), - thread_initialization.map(|initialization| initialization.profile), + ThreadInitializationEventDetails::without_profile( + ThreadInitializationMode::New, + ), out, ); } @@ -785,8 +809,9 @@ impl AnalyticsReducer { connection_id, response.thread, response.model, - ThreadInitializationMode::Resumed, - thread_initialization.map(|initialization| initialization.profile), + ThreadInitializationEventDetails::without_profile( + ThreadInitializationMode::Resumed, + ), out, ); } @@ -795,8 +820,9 @@ impl AnalyticsReducer { connection_id, response.thread, response.model, - ThreadInitializationMode::Forked, - thread_initialization.map(|initialization| initialization.profile), + ThreadInitializationEventDetails::without_profile( + ThreadInitializationMode::Forked, + ), out, ); } @@ -826,6 +852,38 @@ impl AnalyticsReducer { } } + fn ingest_thread_initialization_response( + &mut self, + input: ThreadInitializationResponseFact, + out: &mut Vec, + ) { + let initialization = input.initialization.into(); + match input.response { + ClientResponsePayload::ThreadStart(response) => self.emit_thread_initialized( + input.connection_id, + response.thread, + response.model, + initialization, + out, + ), + ClientResponsePayload::ThreadResume(response) => self.emit_thread_initialized( + input.connection_id, + response.thread, + response.model, + initialization, + out, + ), + ClientResponsePayload::ThreadFork(response) => self.emit_thread_initialized( + input.connection_id, + response.thread, + response.model, + initialization, + out, + ), + _ => {} + } + } + fn ingest_server_request(&mut self, _connection_id: u64, request: ServerRequest) { match request { ServerRequest::CommandExecutionRequestApproval { request_id, params } => { @@ -1199,8 +1257,7 @@ impl AnalyticsReducer { connection_id: u64, thread: codex_app_server_protocol::Thread, model: String, - initialization_mode: ThreadInitializationMode, - profile: Option, + initialization: ThreadInitializationEventDetails, out: &mut Vec, ) { let session_source: SessionSource = thread.source.into(); @@ -1236,12 +1293,12 @@ impl AnalyticsReducer { model, ephemeral: thread.ephemeral, thread_source: thread_metadata.thread_source, - initialization_mode, + initialization_mode: initialization.mode, subagent_source: thread_metadata.subagent_source.clone(), parent_thread_id: thread_metadata.parent_thread_id, forked_from_thread_id, created_at: u64::try_from(thread.created_at).unwrap_or_default(), - profile, + profile: initialization.profile, }, }, )); diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 1bf7b66b11..6ec80b1843 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -655,10 +655,7 @@ impl MessageProcessor { outgoing .register_request_context(request_context.clone()) .await; - request_context - .thread_initialization_timing() - .scope(request_fut.instrument(request_context.span())) - .await; + request_fut.instrument(request_context.span()).await; } pub(crate) fn thread_created_receiver(&self) -> broadcast::Receiver { @@ -831,24 +828,26 @@ impl MessageProcessor { let rpc_gate = Arc::clone(&session.rpc_gate); let processor = Arc::clone(self); let span = request_context.span(); + let timing_context = request_context.clone(); let request = QueuedInitializedRequest::new( rpc_gate, - async move { - let processor_for_request = Arc::clone(&processor); - let result = processor_for_request - .handle_initialized_client_request( - connection_request_id, - codex_request, - request_context, - app_server_client_name, - client_version, - ) - .await; - if let Err(error) = result { - processor.outgoing.send_error(error_request_id, error).await; - } - } - .instrument(span), + timing_context + .scope_thread_initialization(async move { + let processor_for_request = Arc::clone(&processor); + let result = processor_for_request + .handle_initialized_client_request( + connection_request_id, + codex_request, + request_context, + app_server_client_name, + client_version, + ) + .await; + if let Err(error) = result { + processor.outgoing.send_error(error_request_id, error).await; + } + }) + .instrument(span), ); if let Some(scope) = serialization_scope { diff --git a/codex-rs/app-server/src/outgoing_message.rs b/codex-rs/app-server/src/outgoing_message.rs index af20c8a2c3..a47e723391 100644 --- a/codex-rs/app-server/src/outgoing_message.rs +++ b/codex-rs/app-server/src/outgoing_message.rs @@ -1,4 +1,5 @@ use std::collections::HashMap; +use std::future::Future; use std::sync::Arc; use std::sync::atomic::AtomicI64; use std::sync::atomic::Ordering; @@ -6,6 +7,8 @@ use std::time::SystemTime; use std::time::UNIX_EPOCH; use codex_analytics::AnalyticsEventsClient; +use codex_analytics::CompletedThreadInitialization; +use codex_analytics::ThreadInitializationResponseFact; use codex_app_server_protocol::ClientResponsePayload; use codex_app_server_protocol::JSONRPCErrorError; use codex_app_server_protocol::RequestId; @@ -46,8 +49,7 @@ pub(crate) struct ConnectionRequestId { pub(crate) request_id: RequestId, } -/// Trace data we keep for an incoming request until we send its final -/// response or error. +/// Request-scoped state kept until the final response or error is sent. #[derive(Clone)] pub(crate) struct RequestContext { request_id: ConnectionRequestId, @@ -78,8 +80,19 @@ impl RequestContext { self.span.clone() } - pub(crate) fn thread_initialization_timing(&self) -> ThreadInitializationTiming { - self.thread_initialization_timing.clone() + pub(crate) fn scope_thread_initialization( + &self, + future: F, + ) -> impl Future + Send + 'static + where + F: Future + Send + 'static, + { + let timing = self.thread_initialization_timing.clone(); + async move { timing.scope(future).await } + } + + fn complete_thread_initialization(&self) -> Option { + self.thread_initialization_timing.complete_request() } fn record_turn_id(&self, turn_id: &str) { @@ -526,15 +539,25 @@ impl OutgoingMessageSender { .into_jsonrpc_parts_and_payload(request_id.request_id.clone()) .map(|(id, result, response)| { if let Some(response) = response { - let thread_initialization = request_context.as_ref().and_then(|context| { - context.thread_initialization_timing.complete_request() - }); - self.analytics_events_client.track_response( - connection_id.0, - request_id_for_analytics, - response, - thread_initialization, - ); + if let Some(initialization) = request_context + .as_ref() + .and_then(RequestContext::complete_thread_initialization) + { + self.analytics_events_client + .track_thread_initialization_response( + ThreadInitializationResponseFact { + connection_id: connection_id.0, + response, + initialization, + }, + ); + } else { + self.analytics_events_client.track_response( + connection_id.0, + request_id_for_analytics, + response, + ); + } } (id, result) }); diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 78197040dc..1a1d6734e0 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -2428,7 +2428,7 @@ impl ThreadRequestProcessor { .await { Ok(true) => return Ok(()), - Ok(false) => ThreadInitializationTiming::start_cold_resume_configuration(), + Ok(false) => {} Err(error) => { self.outgoing.send_error(request_id, error).await; return Ok(()); @@ -2707,7 +2707,7 @@ impl ThreadRequestProcessor { app_server_client_name: Option, app_server_client_version: Option, ) -> Result { - ThreadInitializationTiming::start_loaded_resume(); + ThreadInitializationTiming::resume_lookup_started(); let running_thread = if params.history.is_some() { if let Ok(existing_thread_id) = ThreadId::from_string(¶ms.thread_id) && self @@ -2876,7 +2876,7 @@ impl ThreadRequestProcessor { "failed to enqueue running thread resume for thread {existing_thread_id}: thread listener command channel is closed" ))); } - ThreadInitializationTiming::complete_core(); + ThreadInitializationTiming::core_completed(); return Ok(true); } Ok(false) diff --git a/codex-rs/app-server/tests/suite/v2/analytics.rs b/codex-rs/app-server/tests/suite/v2/analytics.rs index 4681960599..324ad0da6b 100644 --- a/codex-rs/app-server/tests/suite/v2/analytics.rs +++ b/codex-rs/app-server/tests/suite/v2/analytics.rs @@ -2,6 +2,7 @@ use anyhow::Result; use app_test_support::ChatGptAuthFixture; use app_test_support::DEFAULT_CLIENT_NAME; use app_test_support::write_chatgpt_auth; +use codex_analytics::ThreadInitializationProfile; use codex_config::types::AuthCredentialsStoreMode; use codex_config::types::OtelExporterKind; use codex_config::types::OtelHttpProtocol; @@ -210,54 +211,44 @@ pub(crate) fn assert_basic_thread_initialized_event( } pub(crate) fn assert_thread_initialization_profile(event: &Value) { - let profile = &event["event_params"]; - let value = |name: &str| { - profile[name] - .as_u64() - .unwrap_or_else(|| panic!("thread initialization profile should include {name}")) - }; + let profile = thread_initialization_profile(event); assert_eq!( - (value("duration_ms"), value("core_duration_ms")), + (profile.duration_ms, profile.core_duration_ms), ( - value("app_server_duration_ms").saturating_add(value("core_duration_ms")), - value("existing_thread_lookup_ms") - .saturating_add(value("configuration_resolution_ms")) - .saturating_add(value("session_dependency_loading_ms")) - .saturating_add(value("session_construction_ms")) - .saturating_add(value("mcp_startup_ms")) - .saturating_add(value("session_activation_ms")) - .saturating_add(value("thread_registration_ms")), + profile + .app_server_duration_ms + .saturating_add(profile.core_duration_ms), + profile + .existing_thread_lookup_ms + .saturating_add(profile.configuration_resolution_ms) + .saturating_add(profile.session_dependency_loading_ms) + .saturating_add(profile.session_construction_ms) + .saturating_add(profile.mcp_startup_ms) + .saturating_add(profile.session_activation_ms) + .saturating_add(profile.thread_registration_ms), ) ); } pub(crate) fn assert_loaded_resume_thread_initialization_profile(event: &Value) { - let profile = &event["event_params"]; - let lookup_ms = profile["existing_thread_lookup_ms"] - .as_u64() - .expect("loaded resume should include existing thread lookup timing"); + let profile = thread_initialization_profile(event); assert_eq!( - ( - profile["core_duration_ms"].as_u64(), - profile["configuration_resolution_ms"].as_u64(), - profile["session_dependency_loading_ms"].as_u64(), - profile["session_construction_ms"].as_u64(), - profile["mcp_startup_ms"].as_u64(), - profile["session_activation_ms"].as_u64(), - profile["thread_registration_ms"].as_u64(), - ), - ( - Some(lookup_ms), - Some(0), - Some(0), - Some(0), - Some(0), - Some(0), - Some(0) - ) + profile, + ThreadInitializationProfile { + duration_ms: profile.duration_ms, + app_server_duration_ms: profile.app_server_duration_ms, + core_duration_ms: profile.existing_thread_lookup_ms, + existing_thread_lookup_ms: profile.existing_thread_lookup_ms, + ..Default::default() + } ); } +fn thread_initialization_profile(event: &Value) -> ThreadInitializationProfile { + serde_json::from_value(event["event_params"].clone()) + .expect("thread initialization event should include the complete profile") +} + pub(crate) async fn wait_for_thread_initialized_events( server: &MockServer, read_timeout: Duration, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 16c9790e16..5a00b42ff3 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -4,8 +4,6 @@ use crate::agents_md::LoadedAgentsMd; use crate::config::ConstraintError; use crate::skills::SkillError; use crate::state::ActiveTurn; -use crate::thread_initialization_timing::SessionDependencyBranch; -use crate::thread_initialization_timing::ThreadInitializationPhase; use crate::thread_initialization_timing::ThreadInitializationTiming; use codex_protocol::SessionId; use codex_protocol::config_types::SERVICE_TIER_DEFAULT_REQUEST_VALUE; @@ -532,13 +530,9 @@ impl Session { // - initialize thread persistence with new or resumed session info // - perform default shell discovery // - load history metadata (skipped for subagents) - ThreadInitializationTiming::transition_to( - ThreadInitializationPhase::SessionDependencyLoading, - ); + ThreadInitializationTiming::session_dependency_loading_started(); let thread_persistence_fut = async { - let _timing_guard = ThreadInitializationTiming::begin_session_dependency( - SessionDependencyBranch::ThreadPersistence, - ); + let _timing_guard = ThreadInitializationTiming::begin_thread_persistence(); if config.ephemeral { Ok::<_, anyhow::Error>(None) } else { @@ -601,9 +595,7 @@ impl Session { session_init.ephemeral = config.ephemeral, )); let state_db_fut = async { - let _timing_guard = ThreadInitializationTiming::begin_session_dependency( - SessionDependencyBranch::StateDbLoading, - ); + let _timing_guard = ThreadInitializationTiming::begin_state_db_loading(); if config.ephemeral { None } else if let Some(local_store) = @@ -624,9 +616,7 @@ impl Session { let config_for_mcp = Arc::clone(&config); let mcp_manager_for_mcp = Arc::clone(&mcp_manager); let auth_and_mcp_fut = async move { - let _timing_guard = ThreadInitializationTiming::begin_session_dependency( - SessionDependencyBranch::AuthAndMcpDiscovery, - ); + let _timing_guard = ThreadInitializationTiming::begin_auth_and_mcp_discovery(); let auth = auth_manager_clone.auth().await; let mcp_servers = mcp_manager_for_mcp .effective_servers(&config_for_mcp, auth.as_ref()) @@ -645,9 +635,7 @@ impl Session { )); let plugin_and_skill_warmup_fut = async { - let _timing_guard = ThreadInitializationTiming::begin_session_dependency( - SessionDependencyBranch::PluginAndSkillWarmup, - ); + let _timing_guard = ThreadInitializationTiming::begin_plugin_and_skill_warmup(); warm_plugins_and_skills_for_session_init( Arc::clone(&config), Arc::clone(&environment_manager), @@ -674,7 +662,7 @@ impl Session { auth_and_mcp_fut, plugin_and_skill_warmup_fut ); - ThreadInitializationTiming::transition_to(ThreadInitializationPhase::SessionConstruction); + ThreadInitializationTiming::session_construction_started(); for err in &plugin_skill_errors { error!( @@ -1129,7 +1117,7 @@ impl Session { sess.send_event_raw(event).await; } - ThreadInitializationTiming::transition_to(ThreadInitializationPhase::McpStartup); + ThreadInitializationTiming::mcp_startup_started(); let mut required_mcp_servers: Vec = mcp_servers .iter() .filter(|(_, server)| server.enabled() && server.required()) @@ -1236,7 +1224,7 @@ impl Session { anyhow::bail!("required MCP servers failed to initialize: {details}"); } } - ThreadInitializationTiming::transition_to(ThreadInitializationPhase::SessionActivation); + ThreadInitializationTiming::session_activation_started(); sess.schedule_startup_prewarm(session_configuration.base_instructions.clone()) .await; let session_start_source = match &initial_history { @@ -1260,9 +1248,7 @@ impl Session { match session_result { Ok(sess) => { live_thread_init.commit(); - ThreadInitializationTiming::transition_to( - ThreadInitializationPhase::ThreadRegistration, - ); + ThreadInitializationTiming::thread_registration_started(); Ok(sess) } Err(err) => { diff --git a/codex-rs/core/src/thread_initialization_timing.rs b/codex-rs/core/src/thread_initialization_timing.rs index 8a05b20588..12a0b406f2 100644 --- a/codex-rs/core/src/thread_initialization_timing.rs +++ b/codex-rs/core/src/thread_initialization_timing.rs @@ -6,12 +6,12 @@ use std::sync::PoisonError; use std::time::Duration; use std::time::Instant; -use codex_analytics::ThreadInitializationFact; +use codex_analytics::CompletedThreadInitialization; use codex_analytics::ThreadInitializationMode; use codex_analytics::ThreadInitializationProfile; #[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub(crate) enum ThreadInitializationPhase { +enum ThreadInitializationPhase { ExistingThreadLookup, ConfigurationResolution, SessionDependencyLoading, @@ -22,7 +22,7 @@ pub(crate) enum ThreadInitializationPhase { } #[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub(crate) enum SessionDependencyBranch { +enum SessionDependencyBranch { ThreadPersistence, StateDbLoading, AuthAndMcpDiscovery, @@ -68,41 +68,86 @@ impl ThreadInitializationTiming { .await } - pub fn start_loaded_resume() { + pub fn resume_lookup_started() { Self::start_current( ThreadInitializationMode::Resumed, ThreadInitializationPhase::ExistingThreadLookup, ); } - pub fn start_cold_resume_configuration() { - Self::transition_to(ThreadInitializationPhase::ConfigurationResolution); - } - - pub fn complete_core() { + pub fn core_completed() { Self::with_current(|timing| timing.complete_core_at(Instant::now())); } - pub fn complete_request(&self) -> Option { + pub fn complete_request(&self) -> Option { self.complete_request_at(Instant::now()) } - pub(crate) fn start_configuration_resolution(mode: ThreadInitializationMode) { - Self::start_current(mode, ThreadInitializationPhase::ConfigurationResolution); + pub(crate) fn new_thread_started() { + Self::start_configuration_resolution(ThreadInitializationMode::New); } - pub(crate) fn start_existing_thread_lookup() { - Self::start_current( - ThreadInitializationMode::Resumed, - ThreadInitializationPhase::ExistingThreadLookup, - ); + pub(crate) fn cleared_thread_started() { + Self::start_configuration_resolution(ThreadInitializationMode::Cleared); } - pub(crate) fn transition_to(phase: ThreadInitializationPhase) { + pub(crate) fn forked_thread_started() { + Self::start_configuration_resolution(ThreadInitializationMode::Forked); + } + + pub(crate) fn resumed_thread_started() { + Self::start_configuration_resolution(ThreadInitializationMode::Resumed); + } + + pub(crate) fn configuration_resolution_started() { + Self::transition_to(ThreadInitializationPhase::ConfigurationResolution); + } + + pub(crate) fn session_dependency_loading_started() { + Self::transition_to(ThreadInitializationPhase::SessionDependencyLoading); + } + + pub(crate) fn session_construction_started() { + Self::transition_to(ThreadInitializationPhase::SessionConstruction); + } + + pub(crate) fn mcp_startup_started() { + Self::transition_to(ThreadInitializationPhase::McpStartup); + } + + pub(crate) fn session_activation_started() { + Self::transition_to(ThreadInitializationPhase::SessionActivation); + } + + pub(crate) fn thread_registration_started() { + Self::transition_to(ThreadInitializationPhase::ThreadRegistration); + } + + pub(crate) fn begin_thread_persistence() -> Option { + Self::begin_session_dependency(SessionDependencyBranch::ThreadPersistence) + } + + pub(crate) fn begin_state_db_loading() -> Option { + Self::begin_session_dependency(SessionDependencyBranch::StateDbLoading) + } + + pub(crate) fn begin_auth_and_mcp_discovery() -> Option { + Self::begin_session_dependency(SessionDependencyBranch::AuthAndMcpDiscovery) + } + + pub(crate) fn begin_plugin_and_skill_warmup() -> Option { + Self::begin_session_dependency(SessionDependencyBranch::PluginAndSkillWarmup) + } + + fn transition_to(phase: ThreadInitializationPhase) { Self::with_current(|timing| timing.transition_at(Instant::now(), phase)); } - pub(crate) fn begin_session_dependency( + fn start_configuration_resolution(mode: ThreadInitializationMode) { + Self::start_current(mode, ThreadInitializationPhase::ConfigurationResolution); + } + + fn begin_session_dependency( branch: SessionDependencyBranch, ) -> Option { ACTIVE_THREAD_INITIALIZATION_TIMING @@ -168,7 +213,7 @@ impl ThreadInitializationTiming { state.core_complete = true; } - fn complete_request_at(&self, now: Instant) -> Option { + fn complete_request_at(&self, now: Instant) -> Option { let state = self.state(); if !state.core_complete { return None; @@ -178,7 +223,7 @@ impl ThreadInitializationTiming { duration_to_u64_ms(now.saturating_duration_since(state.request_started_at)); profile.app_server_duration_ms = profile.duration_ms.saturating_sub(profile.core_duration_ms); - Some(ThreadInitializationFact { + Some(CompletedThreadInitialization { initialization_mode: state.initialization_mode?, profile, }) diff --git a/codex-rs/core/src/thread_initialization_timing_tests.rs b/codex-rs/core/src/thread_initialization_timing_tests.rs index 160d41f10e..7dd10d5bd7 100644 --- a/codex-rs/core/src/thread_initialization_timing_tests.rs +++ b/codex-rs/core/src/thread_initialization_timing_tests.rs @@ -1,7 +1,7 @@ use std::time::Duration; use std::time::Instant; -use codex_analytics::ThreadInitializationFact; +use codex_analytics::CompletedThreadInitialization; use codex_analytics::ThreadInitializationMode; use codex_analytics::ThreadInitializationProfile; use pretty_assertions::assert_eq; @@ -21,54 +21,29 @@ fn thread_initialization_profile_breaks_down_request_and_core_time() { ThreadInitializationMode::Cleared, ThreadInitializationPhase::ConfigurationResolution, ); - timing.transition_at( - core_started_at + Duration::from_millis(10), - ThreadInitializationPhase::ExistingThreadLookup, - ); - timing.transition_at( - core_started_at + Duration::from_millis(15), - ThreadInitializationPhase::ConfigurationResolution, - ); - timing.transition_at( - core_started_at + Duration::from_millis(30), - ThreadInitializationPhase::SessionDependencyLoading, - ); - timing.record_dependency( - SessionDependencyBranch::ThreadPersistence, - Duration::from_millis(12), - ); - timing.record_dependency( - SessionDependencyBranch::StateDbLoading, - Duration::from_millis(8), - ); - timing.record_dependency( - SessionDependencyBranch::AuthAndMcpDiscovery, - Duration::from_millis(15), - ); - timing.record_dependency( - SessionDependencyBranch::PluginAndSkillWarmup, - Duration::from_millis(18), - ); - timing.transition_at( - core_started_at + Duration::from_millis(50), - ThreadInitializationPhase::SessionConstruction, - ); - timing.transition_at( - core_started_at + Duration::from_millis(80), - ThreadInitializationPhase::McpStartup, - ); - timing.transition_at( - core_started_at + Duration::from_millis(120), - ThreadInitializationPhase::SessionActivation, - ); - timing.transition_at( - core_started_at + Duration::from_millis(170), - ThreadInitializationPhase::ThreadRegistration, - ); + for (elapsed_ms, phase) in [ + (10, ThreadInitializationPhase::ExistingThreadLookup), + (15, ThreadInitializationPhase::ConfigurationResolution), + (30, ThreadInitializationPhase::SessionDependencyLoading), + (50, ThreadInitializationPhase::SessionConstruction), + (80, ThreadInitializationPhase::McpStartup), + (120, ThreadInitializationPhase::SessionActivation), + (170, ThreadInitializationPhase::ThreadRegistration), + ] { + timing.transition_at(core_started_at + Duration::from_millis(elapsed_ms), phase); + } + for (branch, duration_ms) in [ + (SessionDependencyBranch::ThreadPersistence, 12), + (SessionDependencyBranch::StateDbLoading, 8), + (SessionDependencyBranch::AuthAndMcpDiscovery, 15), + (SessionDependencyBranch::PluginAndSkillWarmup, 18), + ] { + timing.record_dependency(branch, Duration::from_millis(duration_ms)); + } timing.complete_core_at(core_started_at + Duration::from_millis(230)); assert_eq!( timing.complete_request_at(request_started_at + Duration::from_millis(275)), - Some(ThreadInitializationFact { + Some(CompletedThreadInitialization { initialization_mode: ThreadInitializationMode::Cleared, profile: ThreadInitializationProfile { duration_ms: 275, @@ -104,7 +79,7 @@ fn loaded_resume_records_only_existing_thread_lookup() { assert_eq!( timing.complete_request_at(request_started_at + Duration::from_millis(25)), - Some(ThreadInitializationFact { + Some(CompletedThreadInitialization { initialization_mode: ThreadInitializationMode::Resumed, profile: ThreadInitializationProfile { duration_ms: 25, diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 081af24d6f..da45778c59 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -16,10 +16,8 @@ use crate::session::resolve_multi_agent_version; use crate::shell_snapshot::ShellSnapshot; use crate::tasks::InterruptedTurnHistoryMarker; use crate::tasks::interrupted_turn_history_marker; -use crate::thread_initialization_timing::ThreadInitializationPhase; use crate::thread_initialization_timing::ThreadInitializationTiming; use codex_analytics::AnalyticsEventsClient; -use codex_analytics::ThreadInitializationMode; use codex_app_server_protocol::ThreadHistoryBuilder; use codex_app_server_protocol::TurnStatus; use codex_core_plugins::PluginsManager; @@ -587,13 +585,12 @@ impl ThreadManager { &self, options: StartThreadOptions, ) -> CodexResult { - let initialization_mode = match &options.initial_history { - InitialHistory::Cleared => ThreadInitializationMode::Cleared, - InitialHistory::Forked(_) => ThreadInitializationMode::Forked, - InitialHistory::New => ThreadInitializationMode::New, - InitialHistory::Resumed(_) => ThreadInitializationMode::Resumed, - }; - ThreadInitializationTiming::start_configuration_resolution(initialization_mode); + match &options.initial_history { + InitialHistory::Cleared => ThreadInitializationTiming::cleared_thread_started(), + InitialHistory::Forked(_) => ThreadInitializationTiming::forked_thread_started(), + InitialHistory::New => ThreadInitializationTiming::new_thread_started(), + InitialHistory::Resumed(_) => ThreadInitializationTiming::resumed_thread_started(), + } self.start_thread_with_options_and_fork_source(options, /*forked_from_thread_id*/ None) .await } @@ -673,7 +670,7 @@ impl ThreadManager { auth_manager: Arc, parent_trace: Option, ) -> CodexResult { - ThreadInitializationTiming::start_existing_thread_lookup(); + ThreadInitializationTiming::resume_lookup_started(); let initial_history = self.initial_history_from_rollout_path(rollout_path).await?; Box::pin(self.resume_thread_with_history( config, @@ -691,7 +688,7 @@ impl ThreadManager { auth_manager: Arc, parent_trace: Option, ) -> CodexResult { - ThreadInitializationTiming::start_existing_thread_lookup(); + ThreadInitializationTiming::resume_lookup_started(); let environments = default_thread_environment_selections( self.state.environment_manager.as_ref(), &config.cwd, @@ -853,9 +850,7 @@ impl ThreadManager { where S: Into, { - ThreadInitializationTiming::start_configuration_resolution( - ThreadInitializationMode::Forked, - ); + ThreadInitializationTiming::forked_thread_started(); let snapshot = snapshot.into(); let history = self.initial_history_from_rollout_path(path).await?; self.fork_thread_from_history(snapshot, config, history, thread_source, parent_trace) @@ -892,9 +887,7 @@ impl ThreadManager { where S: Into, { - ThreadInitializationTiming::start_configuration_resolution( - ThreadInitializationMode::Forked, - ); + ThreadInitializationTiming::forked_thread_started(); self.fork_thread_with_initial_history( snapshot.into(), config, @@ -1306,15 +1299,13 @@ impl ThreadManagerState { session_configured: thread.session_configured(), thread, }; - ThreadInitializationTiming::complete_core(); + ThreadInitializationTiming::core_completed(); return Ok(new_thread); } threads.remove(&resumed.conversation_id); } } - ThreadInitializationTiming::transition_to( - ThreadInitializationPhase::ConfigurationResolution, - ); + ThreadInitializationTiming::configuration_resolution_started(); let environment_selections = resolve_environment_selections(self.environment_manager.as_ref(), &environments)?; let parent_rollout_thread_trace = self @@ -1367,7 +1358,7 @@ impl ThreadManagerState { if is_resumed_thread { new_thread.thread.emit_thread_resume_lifecycle().await; } - ThreadInitializationTiming::complete_core(); + ThreadInitializationTiming::core_completed(); Ok(new_thread) }