github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-20 00:18:49 +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/Cargo.lock
Generated
+1
View File
@@ -2572,6 +2572,7 @@ dependencies = [
"codex-code-mode-protocol",
"codex-http-client",
"codex-install-context",
"codex-otel",
"codex-protocol",
"futures",
"http-body-util",
codex-rs/code-mode-host/src/grpc/mod.rs
+66 -34
View File
@@ -14,6 +14,7 @@ use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::WaitRequest;
use codex_code_mode_protocol::grpc as proto;
use codex_code_mode_protocol::grpc::code_mode_host_server::CodeModeHost;
use codex_protocol::protocol::W3cTraceContext;
use futures::Stream;
use futures::StreamExt;
use tokio::sync::mpsc;
@@ -21,6 +22,7 @@ use tokio_stream::wrappers::ReceiverStream;
use tonic::Request;
use tonic::Response;
use tonic::Status;
use tracing::Instrument;
use self::session::GrpcHostState;
use self::session::GrpcSession;
@@ -29,6 +31,17 @@ use self::waits::WaitRegistration;
type GrpcStream<T> = Pin<Box<dyn Stream<Item = Result<T, Status>> + Send + 'static>>;
type GrpcFuture<'a, T> = Pin<Box<dyn Future<Output = Result<Response<T>, Status>> + Send + 'a>>;
fn trace_context_from_request<T>(request: &Request<T>) -> Option<W3cTraceContext> {
request
.metadata()
.get("traceparent")
.and_then(|value| value.to_str().ok())
.map(|traceparent| W3cTraceContext {
traceparent: Some(traceparent.to_string()),
tracestate: None,
})
}
/// Serves transport-independent, leased code-mode sessions over gRPC.
#[derive(Clone)]
pub struct GrpcCodeModeHost {
@@ -43,12 +56,6 @@ impl GrpcCodeModeHost {
}
}
#[tracing::instrument(
name = "code_mode_host.grpc.open_session",
level = "info",
skip_all,
fields(otel.name = "code_mode_host.grpc.open_session")
)]
async fn open_session_request(
&self,
request: proto::OpenSessionRequest,
@@ -119,20 +126,10 @@ impl GrpcCodeModeHost {
Ok(Response::new(proto::AcknowledgeNotificationResponse {}))
}
#[tracing::instrument(
name = "code_mode_host.grpc.execute",
level = "info",
skip_all,
fields(
otel.name = "code_mode_host.grpc.execute",
session.id = %request.session_id,
execution.id = %request.execution_id,
call_id = %request.tool_call_id,
)
)]
async fn execute_request(
&self,
request: proto::ExecuteRequest,
callback_traceparent: Option<String>,
) -> Result<Response<GrpcStream<proto::ExecuteEvent>>, Status> {
let session = self.state.session(&request.session_id)?;
validation::identifier(&request.execution_id, "execution ID")?;
@@ -154,7 +151,12 @@ impl GrpcCodeModeHost {
}
};
let cell_id = started.cell_id.clone();
session.admit_execution(execution_id.clone(), cell_id.to_string(), cell_permit)?;
session.admit_execution(
execution_id.clone(),
cell_id.to_string(),
cell_permit,
callback_traceparent,
)?;
let (sender, receiver) = mpsc::channel(/*buffer*/ 2);
sender
@@ -167,22 +169,26 @@ impl GrpcCodeModeHost {
)),
}))
.map_err(|_| Status::internal("failed to publish code-mode execution admission"))?;
tokio::spawn(async move {
let _request_permit = request_permit;
tokio::select! {
biased;
_ = sender.closed() => {}
response = started.initial_response() => {
let event = response.map(|response| proto::ExecuteEvent {
event: Some(proto::execute_event::Event::Outcome(
conversions::execution_outcome(response),
)),
}).map_err(Status::internal);
let _ = sender.send(event).await;
let outcome_span = tracing::Span::current();
tokio::spawn(
async move {
let _request_permit = request_permit;
tokio::select! {
biased;
_ = sender.closed() => {}
response = started.initial_response() => {
let event = response.map(|response| proto::ExecuteEvent {
event: Some(proto::execute_event::Event::Outcome(
conversions::execution_outcome(response),
)),
}).map_err(Status::internal);
let _ = sender.send(event).await;
}
_ = session.closed.cancelled() => {}
}
_ = session.closed.cancelled() => {}
}
});
.instrument(outcome_span),
);
let stream = ReceiverStream::new(receiver).inspect(move |event| {
if matches!(
@@ -278,7 +284,16 @@ impl CodeModeHost for GrpcCodeModeHost {
'a: 'async_trait,
Self: 'async_trait,
{
Box::pin(self.open_session_request(request.into_inner()))
let trace = trace_context_from_request(&request);
let request = request.into_inner();
let open_session_span = tracing::info_span!("code_mode_host.grpc.open_session");
if let Some(trace) = trace.as_ref() {
codex_otel::set_parent_from_w3c_trace_context(&open_session_span, trace);
}
Box::pin(
self.open_session_request(request)
.instrument(open_session_span),
)
}
fn close_session<'a, 'async_trait>(
@@ -333,7 +348,24 @@ impl CodeModeHost for GrpcCodeModeHost {
'a: 'async_trait,
Self: 'async_trait,
{
Box::pin(self.execute_request(request.into_inner()))
let trace = trace_context_from_request(&request);
let request = request.into_inner();
let execute_span = tracing::info_span!(
"code_mode_host.grpc.execute",
otel.name = "code_mode_host.grpc.execute",
session.id = %request.session_id,
execution.id = %request.execution_id,
call_id = %request.tool_call_id,
);
if let Some(trace) = trace.as_ref() {
codex_otel::set_parent_from_w3c_trace_context(&execute_span, trace);
}
let callback_traceparent =
codex_otel::span_w3c_trace_context(&execute_span).and_then(|trace| trace.traceparent);
Box::pin(
self.execute_request(request, callback_traceparent)
.instrument(execute_span),
)
}
fn wait<'a, 'async_trait>(
codex-rs/code-mode-host/src/grpc/robustness_tests.rs
+1
View File
@@ -139,6 +139,7 @@ async fn dropping_execution_before_admission_releases_its_reservation() {
execution_id,
"cell".to_string(),
host.state.cell_permit().expect("reserve cell permit"),
/*traceparent*/ None,
)
.expect_err("abandoned execution must not admit a runtime cell");
assert_eq!(error.code(), Code::Cancelled);
codex-rs/code-mode-host/src/grpc/routing.rs
+3 -2
View File
@@ -103,7 +103,7 @@ impl GrpcSession {
let canonical_tool_name =
ToolName::new(tool_name.namespace.clone(), tool_name.name.clone())
.with_default_namespace();
let (sequence, subscriptions) = {
let (sequence, traceparent, subscriptions) = {
let state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
let Some(execution) = state.cells.get(&cell_id) else {
return Err("code-mode cell closed before dispatching its tool call".to_string());
@@ -124,7 +124,7 @@ impl GrpcSession {
})
.map(|subscription| (subscription.id, subscription.sender.clone()))
.collect::<Vec<_>>();
(sequence, subscriptions)
(sequence, execution.traceparent.clone(), subscriptions)
};
if subscriptions.is_empty() {
return Err("no code-mode tool subscription matches the requested tool".to_string());
@@ -139,6 +139,7 @@ impl GrpcSession {
tool_kind: conversions::tool_kind(invocation.tool_kind),
input_json,
sequence,
traceparent,
};
if message.encoded_len() > MAX_FRAME_BYTES {
return Err(format!(
codex-rs/code-mode-host/src/grpc/session.rs
+3
View File
@@ -71,6 +71,7 @@ pub(super) struct SessionState {
pub(super) struct ExecutionState {
pub(super) execution_id: String,
pub(super) traceparent: Option<String>,
pub(super) tool_call_sequence: u64,
permit: OwnedSemaphorePermit,
}
@@ -296,6 +297,7 @@ impl GrpcSession {
execution_id: String,
cell_id: String,
permit: OwnedSemaphorePermit,
traceparent: Option<String>,
) -> Result<(), Status> {
let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner);
if !state.pending_executions.remove(&execution_id) {
@@ -308,6 +310,7 @@ impl GrpcSession {
};
entry.insert(ExecutionState {
execution_id,
traceparent,
tool_call_sequence: 0,
permit,
});
codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto
+5
View File
@@ -82,6 +82,11 @@ message ToolCall {
// Starts at one and increases independently for each execution.
uint64 sequence = 9;
// The host execution span's W3C parent context for this streamed callback.
// gRPC metadata is fixed when the subscription stream opens, so each call
// carries its own context in the message.
optional string traceparent = 10;
}
message CompleteToolCallRequest {
codex-rs/code-mode/Cargo.toml
+1
View File
@@ -16,6 +16,7 @@ workspace = true
codex-code-mode-protocol = { workspace = true }
codex-http-client = { workspace = true }
codex-install-context = { workspace = true }
codex-otel = { workspace = true }
codex-protocol = { workspace = true }
futures = { workspace = true }
http-body-util = "0.1.3"
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(())
}
codex-rs/code-mode/src/grpc_session/conversion_tests.rs
+1
View File
@@ -73,6 +73,7 @@ fn tool_call_decodes_structured_input_and_namespace() {
tool_kind: grpc::ToolKind::Function as i32,
input_json: Some(br#"{"query":"hello"}"#.to_vec()),
sequence: 1,
traceparent: None,
};
assert_eq!(
codex-rs/code-mode/src/grpc_session/mod.rs
+36 -13
View File
@@ -27,6 +27,7 @@ use tokio::sync::watch;
use tokio_util::sync::CancellationToken;
use tokio_util::task::TaskTracker;
use tonic::transport::Channel;
use tracing::Instrument;
use self::operations::WaitSlot;
use self::state::SessionState;
@@ -49,6 +50,15 @@ type GrpcClient = CodeModeHostClient<GrpcTransport>;
const SHUTDOWN_ERROR: &str = "code mode session is shutting down";
fn inject_span_traceparent<T>(request: &mut tonic::Request<T>, span: &tracing::Span) {
if let Some(traceparent) =
codex_otel::span_w3c_trace_context(span).and_then(|trace| trace.traceparent)
&& let Ok(traceparent) = traceparent.parse()
{
request.metadata_mut().insert("traceparent", traceparent);
}
}
/// Creates code-mode sessions over an HTTP/2 gRPC connection.
#[derive(Clone)]
pub struct GrpcCodeModeSessionProvider {
@@ -83,6 +93,7 @@ impl GrpcCodeModeSessionProvider {
}
}
#[tracing::instrument(name = "code_mode.grpc.open_binding", level = "info", skip_all)]
async fn open_binding(
&self,
delegate: Arc<dyn CodeModeSessionDelegate>,
@@ -100,17 +111,23 @@ impl GrpcCodeModeSessionProvider {
let cell_execution_limits = (limits.max_yield_time_ms.is_some()
|| limits.max_heap_size_bytes.is_some())
.then_some(limits);
let mut lease = deadline::startup(
"session opening",
client.open_session(grpc::OpenSessionRequest {
cell_execution_limits,
}),
)
.await?
.into_inner();
let first = deadline::startup("session lease opening", lease.message())
.await?
.ok_or_else(|| "gRPC code-mode session lease ended before opening".to_string())?;
let open_session_span = tracing::info_span!("code_mode.grpc.open_session");
let mut open_session_request = tonic::Request::new(grpc::OpenSessionRequest {
cell_execution_limits,
});
inject_span_traceparent(&mut open_session_request, &open_session_span);
let (lease, first) = async {
let mut lease =
deadline::startup("session opening", client.open_session(open_session_request))
.await?
.into_inner();
let first = deadline::startup("session lease opening", lease.message())
.await?
.ok_or_else(|| "gRPC code-mode session lease ended before opening".to_string())?;
Ok::<_, String>((lease, first))
}
.instrument(open_session_span)
.await?;
let Some(grpc::session_event::Event::Opened(opened)) = first.event else {
return Err("gRPC code-mode session lease omitted its opening event".to_string());
};
@@ -134,13 +151,19 @@ impl GrpcCodeModeSessionProvider {
};
inner.spawn_session_events(lease);
let request = grpc::SubscribeToToolCallsRequest {
let subscribe_span = tracing::info_span!(
"code_mode.grpc.subscribe_to_tool_calls",
session.id = %inner.id,
);
let mut request = tonic::Request::new(grpc::SubscribeToToolCallsRequest {
session_id: inner.id.clone(),
tool_names: Vec::new(),
};
});
inject_span_traceparent(&mut request, &subscribe_span);
let mut client = inner.client();
let response =
match deadline::startup("tool subscription", client.subscribe_to_tool_calls(request))
.instrument(subscribe_span)
.await
{
Ok(response) => response,
codex-rs/code-mode/src/grpc_session/operations.rs
+20 -1
View File
@@ -12,8 +12,10 @@ use codex_code_mode_protocol::StartedCell;
use codex_code_mode_protocol::WaitOutcome;
use codex_code_mode_protocol::WaitRequest;
use codex_code_mode_protocol::grpc;
use codex_protocol::protocol::W3cTraceContext;
use tokio::sync::OwnedMutexGuard;
use tokio::sync::oneshot;
use tracing::Instrument;
use tracing::debug;
use uuid::Uuid;
@@ -68,6 +70,13 @@ impl SessionInner {
) -> Result<StartedCell, String> {
self.require_open()?;
let execution_id = Uuid::new_v4().to_string();
let execute_span = tracing::info_span!(
"code_mode.grpc.execute",
otel.name = "code_mode.grpc.execute",
execution.id = %execution_id,
call_id = %request.tool_call_id,
);
let trace = codex_otel::span_w3c_trace_context(&execute_span);
let request = conversion::execute_request(&self.id, execution_id.clone(), request)?;
self.state
.lock()
@@ -81,7 +90,10 @@ impl SessionInner {
let (started_tx, started_rx) = oneshot::channel();
let inner = Arc::clone(self);
self.stream_tasks.spawn(async move {
inner.drive_execution(request, ownership, started_tx).await;
inner
.drive_execution(request, ownership, started_tx, trace)
.instrument(execute_span)
.await;
});
started_rx
.await
@@ -93,12 +105,19 @@ impl SessionInner {
request: grpc::ExecuteRequest,
ownership: ExecutionOwnership,
started_tx: oneshot::Sender<Result<StartedCell, String>>,
trace: Option<W3cTraceContext>,
) {
let runtime_timeout =
Duration::from_millis(request.yield_time_ms.unwrap_or(DEFAULT_EXEC_YIELD_TIME_MS))
.saturating_add(Duration::from_secs(1));
let opening = async {
let mut client = self.client();
let mut request = tonic::Request::new(request);
if let Some(traceparent) = trace.and_then(|trace| trace.traceparent)
&& let Ok(traceparent) = traceparent.parse()
{
request.metadata_mut().insert("traceparent", traceparent);
}
let mut stream =
deadline::request(&self, "execution", Duration::ZERO, client.execute(request))
.await?
codex-rs/code-mode/src/grpc_session/state_tests.rs
+1
View File
@@ -44,6 +44,7 @@ fn tool_call(execution_id: &str, invocation_id: u128) -> grpc::ToolCall {
tool_kind: grpc::ToolKind::Function as i32,
input_json: None,
sequence: 1,
traceparent: None,
}
}
codex-rs/core/src/tools/code_mode/delegate.rs
+7 -2
View File
@@ -14,6 +14,7 @@ use serde_json::Value as JsonValue;
use tokio::sync::oneshot;
use tokio::sync::watch;
use tokio_util::sync::CancellationToken;
use tracing::Instrument;
use super::ExecContext;
use super::PUBLIC_TOOL_NAME;
@@ -133,6 +134,7 @@ impl CodeModeDispatchBroker {
invocation,
cancellation_token,
response_tx,
span,
} => {
let cell_id = invocation.cell_id.clone();
if !wait_until_cell_ready_for_dispatch(
@@ -147,8 +149,9 @@ impl CodeModeDispatchBroker {
}
let host = Arc::clone(&host);
tokio::spawn(async move {
let invocation =
host.invoke_tool(invocation, cancellation_token.clone());
let invocation = host
.invoke_tool(invocation, cancellation_token.clone())
.instrument(span);
tokio::pin!(invocation);
let response = tokio::select! {
biased;
@@ -236,6 +239,7 @@ impl CodeModeSessionDelegate for CodeModeDispatchBroker {
invocation,
cancellation_token: cancellation_token.clone(),
response_tx,
span: tracing::Span::current(),
})
.await
.map_err(|_| "code mode nested tool dispatcher is unavailable".to_string())?;
@@ -291,6 +295,7 @@ enum DispatchMessage {
invocation: CodeModeNestedToolCall,
cancellation_token: CancellationToken,
response_tx: oneshot::Sender<Result<JsonValue, String>>,
span: tracing::Span,
},
Notify {
call_id: String,