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

Propagate trace context through gRPC code mode (#41017)

Browse Source 
## What changed

- Inject W3C `traceparent` metadata into code-mode session and execution requests.
- Carry each execution's trace context on streamed tool callbacks so callback and nested-tool spans remain connected across the gRPC boundary.
- Preserve the relevant spans when execution, callback, and dispatch work moves into spawned tasks.

GitOrigin-RevId: 6482963cb9b2e2f0ea27f1345bba208b70a80d0b
This commit is contained in:
yqz-openai
2026-08-27 01:55:53 +00:00
committed by copyberry
parent 7c3747941a
commit eed1dee69e
13 changed files with 191 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/code-mode/src/grpc_session/callbacks.rs
+46 -24
View File
@@ -6,8 +6,10 @@ use std::time::Duration;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::grpc;
use codex_protocol::protocol::W3cTraceContext;
use futures::FutureExt;
use tokio_util::sync::CancellationToken;
use tracing::Instrument;
use tracing::warn;
use super::SessionInner;
@@ -112,6 +114,23 @@ impl SessionInner {
.lock()
.unwrap_or_else(PoisonError::into_inner)
.admit_invocation(&call)?;
let callback_span = tracing::info_span!(
"code_mode.grpc.callback",
otel.name = "code_mode.grpc.callback",
session.id = %call.session_id,
execution.id = %call.execution_id,
cell.id = %call.cell_id,
invocation.id = %call.invocation_id,
);
if let Some(traceparent) = call.traceparent.as_ref() {
codex_otel::set_parent_from_w3c_trace_context(
&callback_span,
&W3cTraceContext {
traceparent: Some(traceparent.clone()),
tracestate: None,
},
);
}
let invocation_id = call.invocation_id.clone();
let cancellation = match admission {
CallbackAdmission::Active(cancellation) => Ok(cancellation),
@@ -133,31 +152,34 @@ impl SessionInner {
};
let invocation = conversion::tool_call(call);
let inner = Arc::clone(self);
tokio::spawn(async move {
let result = match invocation {
Ok(invocation) => {
let callback = AssertUnwindSafe(async {
inner
.delegate
.invoke_tool(invocation, cancellation.child_token())
.await
})
.catch_unwind();
tokio::select! {
biased;
_ = cancellation.cancelled() => return,
result = callback => match result {
Ok(result) => result,
Err(_) => Err("code-mode tool delegate panicked".to_string()),
},
tokio::spawn(
async move {
let result = match invocation {
Ok(invocation) => {
let callback = AssertUnwindSafe(async {
inner
.delegate
.invoke_tool(invocation, cancellation.child_token())
.await
})
.catch_unwind();
tokio::select! {
biased;
_ = cancellation.cancelled() => return,
result = callback => match result {
Ok(result) => result,
Err(_) => Err("code-mode tool delegate panicked".to_string()),
},
}
}
}
Err(error) => Err(error),
};
inner
.complete_tool_call(invocation_id, cancellation, result)
.await;
});
Err(error) => Err(error),
};
inner
.complete_tool_call(invocation_id, cancellation, result)
.await;
}
.instrument(callback_span),
);
Ok(())
}