From 92d90369108daefad443183379a55357e8448b1a Mon Sep 17 00:00:00 2001 From: mchen-oai Date: Thu, 11 Jun 2026 16:55:01 -0700 Subject: [PATCH] Add spans to turn lifecycle gaps (#27623) ## Why Codex app-server latency traces do not granularly cover turn task startup and inter-request handoffs. These spans help attribute time across task execution, startup prewarm, in-flight tool completion, and rollout persistence. ## What changed - Add `session_task.run` spans around task execution and `session_task.flush_rollout` around flushing pending conversation transcript writes to durable storage - Add `regular_task.prepare_run_turn` around regular-turn startup (Send the `TurnStarted` event, reset turn-specific reasoning state, and resolve any startup prewarm) - Add `startup_prewarm.resolve` around waiting for background session prewarming to finish, fail, time out, or be cancelled - Add a function-level trace span around draining in-flight tool calls (Wait for tool calls to complete, record tool result in conversation history, and other bookkeeping) ## Verification Trigger Codex rollout and observe new spans are included --- codex-rs/core/src/session/mod.rs | 1 + codex-rs/core/src/session/turn.rs | 1 + codex-rs/core/src/session_startup_prewarm.rs | 2 ++ codex-rs/core/src/tasks/mod.rs | 2 ++ codex-rs/core/src/tasks/regular.rs | 29 +++++++++++--------- 5 files changed, 22 insertions(+), 13 deletions(-) diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 68c05e2c4..08c691b9c 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -1106,6 +1106,7 @@ impl Session { } /// Flush rollout writes and return the final durability-barrier result. + #[instrument(name = "session.flush_rollout", level = "trace", skip_all)] pub(crate) async fn flush_rollout(&self) -> std::io::Result<()> { if let Some(live_thread) = self.live_thread() { live_thread.flush().await.map_err(std::io::Error::other) diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 7bb34aec8..da1d82531 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -1764,6 +1764,7 @@ async fn handle_assistant_item_done_in_plan_mode( false } +#[instrument(level = "trace", skip_all)] async fn drain_in_flight( in_flight: &mut FuturesOrdered>>, sess: Arc, diff --git a/codex-rs/core/src/session_startup_prewarm.rs b/codex-rs/core/src/session_startup_prewarm.rs index 4fee5e774..0474eb870 100644 --- a/codex-rs/core/src/session_startup_prewarm.rs +++ b/codex-rs/core/src/session_startup_prewarm.rs @@ -6,6 +6,7 @@ use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; use tokio_util::task::AbortOnDropHandle; use tracing::info; +use tracing::instrument; use tracing::warn; use crate::client::ModelClientSession; @@ -53,6 +54,7 @@ impl SessionStartupPrewarmHandle { let _ = self.task.await; } + #[instrument(name = "startup_prewarm.resolve", level = "trace", skip_all)] async fn resolve( self, session_telemetry: &SessionTelemetry, diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 52af0d04d..61f698015 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -19,6 +19,7 @@ use tracing::Span; use tracing::field; use tracing::info_span; use tracing::trace; +use tracing::trace_span; use tracing::warn; use crate::codex_thread::BackgroundTerminalInfo; @@ -400,6 +401,7 @@ impl Session { task_input, task_cancellation_token.child_token(), ) + .instrument(trace_span!("session_task.run")) .await; let sess = session_ctx.clone_session(); if let Err(err) = sess.flush_rollout().await { diff --git a/codex-rs/core/src/tasks/regular.rs b/codex-rs/core/src/tasks/regular.rs index c01cae914..c5cc2da31 100644 --- a/codex-rs/core/src/tasks/regular.rs +++ b/codex-rs/core/src/tasks/regular.rs @@ -45,19 +45,22 @@ impl SessionTask for RegularTask { let run_turn_span = trace_span!("run_turn"); // Regular turns emit `TurnStarted` inline so first-turn lifecycle does // not wait on startup prewarm resolution. - let event = EventMsg::TurnStarted(TurnStartedEvent { - turn_id: ctx.sub_id.clone(), - trace_id: ctx.trace_id.clone(), - started_at: ctx.turn_timing_state.started_at_unix_secs().await, - model_context_window: ctx.model_context_window(), - collaboration_mode_kind: ctx.collaboration_mode.mode, - }); - sess.send_event(ctx.as_ref(), event).await; - sess.set_server_reasoning_included(/*included*/ false).await; - let prewarmed_client_session = match sess - .consume_startup_prewarm_for_regular_turn(&cancellation_token) - .await - { + let prewarmed_client_session = async { + let event = EventMsg::TurnStarted(TurnStartedEvent { + turn_id: ctx.sub_id.clone(), + trace_id: ctx.trace_id.clone(), + started_at: ctx.turn_timing_state.started_at_unix_secs().await, + model_context_window: ctx.model_context_window(), + collaboration_mode_kind: ctx.collaboration_mode.mode, + }); + sess.send_event(ctx.as_ref(), event).await; + sess.set_server_reasoning_included(/*included*/ false).await; + sess.consume_startup_prewarm_for_regular_turn(&cancellation_token) + .await + } + .instrument(trace_span!("regular_task.prepare_run_turn")) + .await; + let prewarmed_client_session = match prewarmed_client_session { SessionStartupPrewarmResolution::Cancelled => return None, SessionStartupPrewarmResolution::Unavailable { .. } => None, SessionStartupPrewarmResolution::Ready(prewarmed_client_session) => {