mirror of
https://github.com/openai/codex.git
synced 2026-10-26 07:58:43 +00:00
Owen Lin
52521a5e40
### Summary Propagate trace context originating at app-server RPC method handlers -> codex core submission loop (so this includes spans such as `run_turn`!). This implements PR 2 of the app-server tracing rollout. This also removes the old lower-level env-based reparenting in core so explicit request/submission ancestry wins instead of being overridden by ambient `TRACEPARENT` state. ### What changed - Added `trace: Option<W3cTraceContext>` to codex_protocol::Submission - Taught `Codex::submit()` / `submit_with_id()` to automatically capture the current span context when constructing or forwarding a submission - Wrapped the core submission loop in a submission_dispatch span parented from Submission.trace - Warn on invalid submission trace carriers and ignore them cleanly - Removed the old env-based downstream reparenting path in core task execution - Stopped OTEL provider init from implicitly attaching env trace context process-wide - Updated mcp-server Submission call sites for the new field Added focused unit tests for: - capturing trace context into Submission - preferring `Submission.trace` when building the core dispatch span ### Why PR 1 gave us consistent inbound request spans in app-server, but that only covered the transport boundary. For long-running work like turns and reviews, the important missing piece was preserving ancestry after the request handler returns and core continues work on a different async path. This change makes that handoff explicit and keeps the parentage rules simple: - app-server request span sets the current context - `Submission.trace` snapshots that context - core restores it once, at the submission boundary - deeper core spans inherit naturally That also lets us stop relying on env-based reparenting for this path, which was too ambient and could override explicit ancestry.
93 lines
2.5 KiB
Rust
93 lines
2.5 KiB
Rust
use std::sync::Arc;
|
|
use std::sync::Mutex;
|
|
|
|
use crate::client::ModelClient;
|
|
use crate::client::ModelClientSession;
|
|
use crate::client_common::Prompt;
|
|
use crate::codex::TurnContext;
|
|
use crate::codex::run_turn;
|
|
use crate::error::Result as CodexResult;
|
|
use crate::state::TaskKind;
|
|
use async_trait::async_trait;
|
|
use codex_protocol::user_input::UserInput;
|
|
use tokio_util::sync::CancellationToken;
|
|
use tracing::Instrument;
|
|
use tracing::trace_span;
|
|
|
|
use super::SessionTask;
|
|
use super::SessionTaskContext;
|
|
|
|
pub(crate) struct RegularTask {
|
|
prewarmed_session: Mutex<Option<ModelClientSession>>,
|
|
}
|
|
|
|
impl Default for RegularTask {
|
|
fn default() -> Self {
|
|
Self {
|
|
prewarmed_session: Mutex::new(None),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl RegularTask {
|
|
pub(crate) async fn with_startup_prewarm(
|
|
model_client: ModelClient,
|
|
prompt: Prompt,
|
|
turn_context: Arc<TurnContext>,
|
|
turn_metadata_header: Option<String>,
|
|
) -> CodexResult<Self> {
|
|
let mut client_session = model_client.new_session();
|
|
client_session
|
|
.prewarm_websocket(
|
|
&prompt,
|
|
&turn_context.model_info,
|
|
&turn_context.otel_manager,
|
|
turn_context.reasoning_effort,
|
|
turn_context.reasoning_summary,
|
|
turn_context.config.service_tier,
|
|
turn_metadata_header.as_deref(),
|
|
)
|
|
.await?;
|
|
|
|
Ok(Self {
|
|
prewarmed_session: Mutex::new(Some(client_session)),
|
|
})
|
|
}
|
|
|
|
async fn take_prewarmed_session(&self) -> Option<ModelClientSession> {
|
|
self.prewarmed_session
|
|
.lock()
|
|
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
|
.take()
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl SessionTask for RegularTask {
|
|
fn kind(&self) -> TaskKind {
|
|
TaskKind::Regular
|
|
}
|
|
|
|
async fn run(
|
|
self: Arc<Self>,
|
|
session: Arc<SessionTaskContext>,
|
|
ctx: Arc<TurnContext>,
|
|
input: Vec<UserInput>,
|
|
cancellation_token: CancellationToken,
|
|
) -> Option<String> {
|
|
let sess = session.clone_session();
|
|
let run_turn_span = trace_span!("run_turn");
|
|
sess.set_server_reasoning_included(false).await;
|
|
let prewarmed_client_session = self.take_prewarmed_session().await;
|
|
run_turn(
|
|
sess,
|
|
ctx,
|
|
input,
|
|
prewarmed_client_session,
|
|
cancellation_token,
|
|
)
|
|
.instrument(run_turn_span)
|
|
.await
|
|
}
|
|
}
|