diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 61edfedb72..b1c6a2ac96 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2572,6 +2572,7 @@ dependencies = [ "codex-code-mode-protocol", "codex-http-client", "codex-install-context", + "codex-otel", "codex-protocol", "futures", "http-body-util", diff --git a/codex-rs/code-mode-host/src/grpc/mod.rs b/codex-rs/code-mode-host/src/grpc/mod.rs index 4647f6cdce..7a90a7469c 100644 --- a/codex-rs/code-mode-host/src/grpc/mod.rs +++ b/codex-rs/code-mode-host/src/grpc/mod.rs @@ -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 = Pin> + Send + 'static>>; type GrpcFuture<'a, T> = Pin, Status>> + Send + 'a>>; +fn trace_context_from_request(request: &Request) -> Option { + 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, ) -> Result>, 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>( diff --git a/codex-rs/code-mode-host/src/grpc/robustness_tests.rs b/codex-rs/code-mode-host/src/grpc/robustness_tests.rs index 31c5eb1095..45aac60556 100644 --- a/codex-rs/code-mode-host/src/grpc/robustness_tests.rs +++ b/codex-rs/code-mode-host/src/grpc/robustness_tests.rs @@ -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); diff --git a/codex-rs/code-mode-host/src/grpc/routing.rs b/codex-rs/code-mode-host/src/grpc/routing.rs index c1879a19fd..4a7ddfeecd 100644 --- a/codex-rs/code-mode-host/src/grpc/routing.rs +++ b/codex-rs/code-mode-host/src/grpc/routing.rs @@ -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::>(); - (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!( diff --git a/codex-rs/code-mode-host/src/grpc/session.rs b/codex-rs/code-mode-host/src/grpc/session.rs index 3a215726e3..a8ce05bf6f 100644 --- a/codex-rs/code-mode-host/src/grpc/session.rs +++ b/codex-rs/code-mode-host/src/grpc/session.rs @@ -71,6 +71,7 @@ pub(super) struct SessionState { pub(super) struct ExecutionState { pub(super) execution_id: String, + pub(super) traceparent: Option, pub(super) tool_call_sequence: u64, permit: OwnedSemaphorePermit, } @@ -296,6 +297,7 @@ impl GrpcSession { execution_id: String, cell_id: String, permit: OwnedSemaphorePermit, + traceparent: Option, ) -> 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, }); diff --git a/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto b/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto index ece24730cc..f70f547e1b 100644 --- a/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto +++ b/codex-rs/code-mode-protocol/src/grpc/codex.code_mode.v1.proto @@ -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 { diff --git a/codex-rs/code-mode/Cargo.toml b/codex-rs/code-mode/Cargo.toml index 8bc8b84933..b6c0edf251 100644 --- a/codex-rs/code-mode/Cargo.toml +++ b/codex-rs/code-mode/Cargo.toml @@ -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" diff --git a/codex-rs/code-mode/src/grpc_session/callbacks.rs b/codex-rs/code-mode/src/grpc_session/callbacks.rs index 27714a836a..8f84fbbd64 100644 --- a/codex-rs/code-mode/src/grpc_session/callbacks.rs +++ b/codex-rs/code-mode/src/grpc_session/callbacks.rs @@ -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(()) } diff --git a/codex-rs/code-mode/src/grpc_session/conversion_tests.rs b/codex-rs/code-mode/src/grpc_session/conversion_tests.rs index ef7847c969..ed0e66e2a9 100644 --- a/codex-rs/code-mode/src/grpc_session/conversion_tests.rs +++ b/codex-rs/code-mode/src/grpc_session/conversion_tests.rs @@ -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!( diff --git a/codex-rs/code-mode/src/grpc_session/mod.rs b/codex-rs/code-mode/src/grpc_session/mod.rs index a7932bbd2f..38b3635818 100644 --- a/codex-rs/code-mode/src/grpc_session/mod.rs +++ b/codex-rs/code-mode/src/grpc_session/mod.rs @@ -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; const SHUTDOWN_ERROR: &str = "code mode session is shutting down"; +fn inject_span_traceparent(request: &mut tonic::Request, 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, @@ -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, diff --git a/codex-rs/code-mode/src/grpc_session/operations.rs b/codex-rs/code-mode/src/grpc_session/operations.rs index edf4cd540b..ea29cc4788 100644 --- a/codex-rs/code-mode/src/grpc_session/operations.rs +++ b/codex-rs/code-mode/src/grpc_session/operations.rs @@ -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 { 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>, + trace: Option, ) { 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? diff --git a/codex-rs/code-mode/src/grpc_session/state_tests.rs b/codex-rs/code-mode/src/grpc_session/state_tests.rs index b91f50311e..5a146304c4 100644 --- a/codex-rs/code-mode/src/grpc_session/state_tests.rs +++ b/codex-rs/code-mode/src/grpc_session/state_tests.rs @@ -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, } } diff --git a/codex-rs/core/src/tools/code_mode/delegate.rs b/codex-rs/core/src/tools/code_mode/delegate.rs index bfa6a1c0bb..56389b0f3b 100644 --- a/codex-rs/core/src/tools/code_mode/delegate.rs +++ b/codex-rs/core/src/tools/code_mode/delegate.rs @@ -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>, + span: tracing::Span, }, Notify { call_id: String,