openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-06 20:20:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

Decouple thread initialization profiling

Browse Source
This commit is contained in:
Ahmed Ibrahim
2026-06-08 12:21:43 -07:00
parent e58edbbef6
commit 133689b20f
14 changed files with 289 additions and 230 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/analytics/src/analytics_client_tests.rs
-16
View File
@@ -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,
)
codex-rs/analytics/src/client.rs
+5 -3
View File
@@ -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<ThreadInitializationFact>,
) {
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,
codex-rs/analytics/src/client_tests.rs
+1 -4
View File
@@ -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)));
}
codex-rs/analytics/src/facts.rs
+11 -3
View File
@@ -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<ClientResponsePayload>,
thread_initialization: Option<ThreadInitializationFact>,
},
ThreadInitializationResponse(Box<ThreadInitializationResponseFact>),
ErrorResponse {
connection_id: u64,
request_id: RequestId,
codex-rs/analytics/src/lib.rs
+2 -1
View File
@@ -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;
codex-rs/analytics/src/reducer.rs
+74 -17
View File
@@ -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<ToolItemKey, ItemReviewSummary>,
}
#[derive(Clone, Copy)]
struct ThreadInitializationEventDetails {
mode: ThreadInitializationMode,
profile: Option<ThreadInitializationProfile>,
}
impl ThreadInitializationEventDetails {
fn without_profile(mode: ThreadInitializationMode) -> Self {
Self {
mode,
profile: None,
}
}
}
impl From<CompletedThreadInitialization> 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<ThreadInitializationFact>,
out: &mut Vec<TrackEventRequest>,
) {
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<TrackEventRequest>,
) {
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<ThreadInitializationProfile>,
initialization: ThreadInitializationEventDetails,
out: &mut Vec<TrackEventRequest>,
) {
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,
},
},
));
codex-rs/app-server/src/message_processor.rs
+19 -20
View File
@@ -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<ThreadId> {
@@ -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 {
codex-rs/app-server/src/outgoing_message.rs
+36 -13
View File
@@ -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<F>(
&self,
future: F,
) -> impl Future<Output = F::Output> + 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<CompletedThreadInitialization> {
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)
});
codex-rs/app-server/src/request_processors/thread_processor.rs
+3 -3
View File
@@ -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<String>,
app_server_client_version: Option<String>,
) -> Result<bool, JSONRPCErrorError> {
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(&params.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)
codex-rs/app-server/tests/suite/v2/analytics.rs
+28 -37
View File
@@ -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,
codex-rs/core/src/session/session.rs
+9 -23
View File
@@ -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<String> = 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) => {
codex-rs/core/src/thread_initialization_timing.rs
+66 -21
View File
@@ -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<ThreadInitializationFact> {
pub fn complete_request(&self) -> Option<CompletedThreadInitialization> {
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<SessionDependencyTimingGuard> {
Self::begin_session_dependency(SessionDependencyBranch::ThreadPersistence)
}
pub(crate) fn begin_state_db_loading() -> Option<SessionDependencyTimingGuard> {
Self::begin_session_dependency(SessionDependencyBranch::StateDbLoading)
}
pub(crate) fn begin_auth_and_mcp_discovery() -> Option<SessionDependencyTimingGuard> {
Self::begin_session_dependency(SessionDependencyBranch::AuthAndMcpDiscovery)
}
pub(crate) fn begin_plugin_and_skill_warmup() -> Option<SessionDependencyTimingGuard> {
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<SessionDependencyTimingGuard> {
ACTIVE_THREAD_INITIALIZATION_TIMING
@@ -168,7 +213,7 @@ impl ThreadInitializationTiming {
state.core_complete = true;
}
fn complete_request_at(&self, now: Instant) -> Option<ThreadInitializationFact> {
fn complete_request_at(&self, now: Instant) -> Option<CompletedThreadInitialization> {
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,
})
codex-rs/core/src/thread_initialization_timing_tests.rs
+22 -47
View File
@@ -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,
codex-rs/core/src/thread_manager.rs
+13 -22
View File
@@ -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<NewThread> {
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<AuthManager>,
parent_trace: Option<W3cTraceContext>,
) -> CodexResult<NewThread> {
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<AuthManager>,
parent_trace: Option<W3cTraceContext>,
) -> CodexResult<NewThread> {
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<ForkSnapshot>,
{
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<ForkSnapshot>,
{
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)
}