diff --git a/codex-rs/app-server-protocol/src/protocol/common.rs b/codex-rs/app-server-protocol/src/protocol/common.rs index bf56f7f28a..17bef19b26 100644 --- a/codex-rs/app-server-protocol/src/protocol/common.rs +++ b/codex-rs/app-server-protocol/src/protocol/common.rs @@ -1212,13 +1212,14 @@ macro_rules! server_request_definitions { ),* $(,)? ) => { /// Request initiated from the server and sent to the client. - #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] + #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Display, JsonSchema, TS)] #[allow(clippy::large_enum_variant)] #[serde(tag = "method", rename_all = "camelCase")] + #[strum(serialize_all = "camelCase")] pub enum ServerRequest { $( $(#[doc = $variant_doc])* - $(#[serde(rename = $wire)] #[ts(rename = $wire)])? + $(#[serde(rename = $wire)] #[ts(rename = $wire)] #[strum(serialize = $wire)])? $variant { #[serde(rename = "id")] request_id: RequestId, @@ -1234,6 +1235,10 @@ macro_rules! server_request_definitions { } } + pub fn method(&self) -> String { + self.to_string() + } + pub fn response_from_result( &self, result: crate::Result, diff --git a/codex-rs/app-server/src/app_server_tracing.rs b/codex-rs/app-server/src/app_server_tracing.rs index 119fc281d1..f8c17a3acb 100644 --- a/codex-rs/app-server/src/app_server_tracing.rs +++ b/codex-rs/app-server/src/app_server_tracing.rs @@ -13,9 +13,11 @@ use crate::transport::AppServerTransport; use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::InitializeParams; use codex_app_server_protocol::JSONRPCRequest; +use codex_app_server_protocol::RequestId; use codex_otel::set_parent_from_context; use codex_otel::set_parent_from_w3c_trace_context; use codex_otel::traceparent_context_from_env; +use codex_protocol::ThreadId; use codex_protocol::protocol::W3cTraceContext; use tracing::Span; use tracing::field; @@ -82,6 +84,34 @@ pub(crate) fn typed_request_span( span } +pub(crate) fn server_request_span( + method: &str, + request_id: &RequestId, + target_count: Option, + thread_id: Option<&ThreadId>, +) -> Span { + let thread_id = thread_id.map(ToString::to_string); + let span = info_span!( + "app_server.server_request", + otel.kind = "client", + otel.name = method, + rpc.system = "jsonrpc", + rpc.method = method, + rpc.request_id = %request_id, + rpc.result = field::Empty, + rpc.error_code = field::Empty, + app_server.target_count = field::Empty, + codex.thread.id = thread_id.as_deref().unwrap_or(""), + ); + if let Some(target_count) = target_count { + span.record( + "app_server.target_count", + i64::try_from(target_count).unwrap_or(i64::MAX), + ); + } + span +} + fn transport_name(transport: &AppServerTransport) -> &'static str { match transport { AppServerTransport::Stdio => "stdio", diff --git a/codex-rs/app-server/src/attestation.rs b/codex-rs/app-server/src/attestation.rs index 206c38cee1..64349a3a4c 100644 --- a/codex-rs/app-server/src/attestation.rs +++ b/codex-rs/app-server/src/attestation.rs @@ -75,7 +75,7 @@ async fn request_attestation_header_value_with_timeout( .send_request_to_connections( Some(&connection_ids), ServerRequestPayload::AttestationGenerate(AttestationGenerateParams {}), - /*thread_id*/ None, + Some(thread_id), ) .await; diff --git a/codex-rs/app-server/src/current_time.rs b/codex-rs/app-server/src/current_time.rs index 57be618af3..a2184a13aa 100644 --- a/codex-rs/app-server/src/current_time.rs +++ b/codex-rs/app-server/src/current_time.rs @@ -110,7 +110,7 @@ async fn request_current_time( ServerRequestPayload::CurrentTimeRead(CurrentTimeReadParams { thread_id: thread_id.to_string(), }), - /*thread_id*/ None, + Some(thread_id), ) .await; diff --git a/codex-rs/app-server/src/message_processor_tracing_tests.rs b/codex-rs/app-server/src/message_processor_tracing_tests.rs index 08486c0f89..6f755e6c38 100644 --- a/codex-rs/app-server/src/message_processor_tracing_tests.rs +++ b/codex-rs/app-server/src/message_processor_tracing_tests.rs @@ -12,11 +12,15 @@ use app_test_support::write_mock_responses_config_toml; use codex_analytics::AppServerRpcTransport; use codex_app_server_protocol::ClientInfo; use codex_app_server_protocol::ClientRequest; +use codex_app_server_protocol::CurrentTimeReadParams; use codex_app_server_protocol::InitializeCapabilities; use codex_app_server_protocol::InitializeParams; use codex_app_server_protocol::InitializeResponse; +use codex_app_server_protocol::JSONRPCError; +use codex_app_server_protocol::JSONRPCErrorError; use codex_app_server_protocol::JSONRPCRequest; use codex_app_server_protocol::RequestId; +use codex_app_server_protocol::ServerRequestPayload; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::TurnStartParams; @@ -30,6 +34,7 @@ use codex_core::config::ConfigBuilder; use codex_exec_server::EnvironmentManager; use codex_feedback::CodexFeedback; use codex_login::AuthManager; +use codex_protocol::ThreadId; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::W3cTraceContext; use opentelemetry::global; @@ -303,6 +308,27 @@ fn span_attr<'a>(span: &'a SpanData, key: &str) -> Option<&'a str> { }) } +fn span_attr_i64(span: &SpanData, key: &str) -> Option { + span.attributes + .iter() + .find(|kv| kv.key.as_str() == key) + .and_then(|kv| match &kv.value { + opentelemetry::Value::I64(value) => Some(*value), + _ => None, + }) +} + +fn find_rpc_span<'a>(spans: &'a [SpanData], kind: SpanKind, method: &str) -> &'a SpanData { + spans + .iter() + .find(|span| { + span.span_kind == kind + && span_attr(span, "rpc.system") == Some("jsonrpc") + && span_attr(span, "rpc.method") == Some(method) + }) + .unwrap_or_else(|| panic!("missing {kind:?} span for rpc.method={method}")) +} + fn find_rpc_span_with_trace<'a>( spans: &'a [SpanData], kind: SpanKind, @@ -718,3 +744,81 @@ async fn turn_start_jsonrpc_span_parents_core_turn_spans() -> Result<()> { Ok(()) } + +#[tokio::test(flavor = "current_thread")] +#[serial(app_server_tracing)] +async fn server_request_span_records_terminal_outcomes() -> Result<()> { + let harness = TracingHarness::new().await?; + let outgoing = Arc::clone(&harness.processor.outgoing); + let connection_ids = [TEST_CONNECTION_ID]; + let thread_id = ThreadId::new(); + + harness.reset_tracing(); + let (error_id, _error_rx) = outgoing + .send_request_to_connections( + Some(&connection_ids), + ServerRequestPayload::CurrentTimeRead(CurrentTimeReadParams { + thread_id: thread_id.to_string(), + }), + Some(thread_id), + ) + .await; + let expected_error = JSONRPCErrorError { + code: -32_042, + message: "clock unavailable".to_string(), + data: None, + }; + harness + .processor + .process_error(JSONRPCError { + error: expected_error.clone(), + id: error_id.clone(), + }) + .await; + let error_spans = wait_for_exported_spans(harness.tracing, |spans| { + spans.iter().any(|span| { + span.span_kind == SpanKind::Client + && span_attr(span, "rpc.method") == Some("currentTime/read") + && span_attr(span, "rpc.result") == Some("error") + }) + }) + .await; + let error_span = find_rpc_span(&error_spans, SpanKind::Client, "currentTime/read"); + assert_eq!(error_span.name.as_ref(), "currentTime/read"); + assert_eq!(span_attr(error_span, "rpc.result"), Some("error")); + let thread_id_string = thread_id.to_string(); + assert_eq!( + span_attr_i64(error_span, "app_server.target_count"), + Some(1) + ); + assert_eq!( + span_attr(error_span, "codex.thread.id"), + Some(thread_id_string.as_str()) + ); + assert_eq!( + span_attr_i64(error_span, "rpc.error_code"), + Some(expected_error.code) + ); + + harness.reset_tracing(); + let (cancelled_id, _cancelled_rx) = outgoing + .send_request_to_connections( + Some(&connection_ids), + ServerRequestPayload::CurrentTimeRead(CurrentTimeReadParams { + thread_id: thread_id.to_string(), + }), + Some(thread_id), + ) + .await; + assert!(outgoing.cancel_request(&cancelled_id).await); + let _ = wait_for_exported_spans(harness.tracing, |spans| { + spans + .iter() + .any(|span| span_attr(span, "rpc.result") == Some("cancelled")) + }) + .await; + + harness.shutdown().await; + + Ok(()) +} diff --git a/codex-rs/app-server/src/outgoing_message.rs b/codex-rs/app-server/src/outgoing_message.rs index 760b77f018..ddd206a6c2 100644 --- a/codex-rs/app-server/src/outgoing_message.rs +++ b/codex-rs/app-server/src/outgoing_message.rs @@ -42,6 +42,8 @@ pub(crate) type ClientRequestResult = std::result::Result "success", Self::Error => "error", + Self::Cancelled => "cancelled", + Self::SendFailed => "send_failed", Self::Disconnected => "disconnected", Self::Replaced => "replaced", } @@ -142,6 +146,25 @@ struct PendingCallbackEntry { callback: oneshot::Sender, thread_id: Option, request: ServerRequest, + span: Span, +} + +impl PendingCallbackEntry { + fn record_result(&self, result: RpcResult) { + self.span.record("rpc.result", result.as_str()); + } + + fn record_error(&self, error: &JSONRPCErrorError) { + self.record_result(RpcResult::Error); + self.span.record("rpc.error_code", error.code); + } + + fn record_cancelled(&self, error: Option<&JSONRPCErrorError>) { + self.record_result(RpcResult::Cancelled); + if let Some(error) = error { + self.span.record("rpc.error_code", error.code); + } + } } impl ThreadScopedOutgoingMessageSender { @@ -324,18 +347,29 @@ impl OutgoingMessageSender { let id = self.next_request_id(); let outgoing_message_id = id.clone(); let request = request.request_with_id(outgoing_message_id.clone()); + let method = request.method(); + let request_span = crate::app_server_tracing::server_request_span( + method.as_str(), + &outgoing_message_id, + connection_ids.map(<[ConnectionId]>::len), + thread_id.as_ref(), + ); let (tx_approve, rx_approve) = oneshot::channel(); { let mut request_id_to_callback = self.request_id_to_callback.lock().await; - request_id_to_callback.insert( + if let Some(replaced) = request_id_to_callback.insert( id, PendingCallbackEntry { callback: tx_approve, thread_id, request: request.clone(), + span: request_span.clone(), }, - ); + ) { + replaced.record_result(RpcResult::Replaced); + warn!("replaced unresolved server request callback"); + } } let outgoing_message = OutgoingMessage::Request(request.clone()); @@ -374,6 +408,7 @@ impl OutgoingMessageSender { }; if let Err(err) = send_result { + request_span.record("rpc.result", RpcResult::SendFailed.as_str()); warn!("failed to send request {outgoing_message_id:?} to client: {err:?}"); let mut request_id_to_callback = self.request_id_to_callback.lock().await; request_id_to_callback.remove(&outgoing_message_id); @@ -407,6 +442,7 @@ impl OutgoingMessageSender { match entry { Some((id, entry)) => { + entry.record_result(RpcResult::Success); let completed_at_ms = now_unix_timestamp_ms(); if let Ok(response) = entry.request.response_from_result(result.clone()) && !matches!(response, ServerResponse::PermissionsRequestApproval { .. }) @@ -429,6 +465,7 @@ impl OutgoingMessageSender { match entry { Some((id, entry)) => { + entry.record_error(&error); warn!("client responded with error for {id:?}: {error:?}"); self.analytics_events_client .track_server_request_aborted(now_unix_timestamp_ms(), id.clone()); @@ -444,7 +481,8 @@ impl OutgoingMessageSender { pub(crate) async fn cancel_request(&self, id: &RequestId) -> bool { let entry = self.take_request_callback(id).await; - if let Some((request_id, _entry)) = entry { + if let Some((request_id, entry)) = entry { + entry.record_cancelled(/*error*/ None); self.analytics_events_client .track_server_request_aborted(now_unix_timestamp_ms(), request_id); true @@ -463,6 +501,7 @@ impl OutgoingMessageSender { }; for entry in entries { + entry.record_cancelled(error.as_ref()); self.analytics_events_client .track_server_request_aborted(now_unix_timestamp_ms(), entry.request.id().clone()); if let Some(error) = error.as_ref() @@ -521,6 +560,7 @@ impl OutgoingMessageSender { }; for entry in entries { + entry.record_cancelled(error.as_ref()); self.analytics_events_client .track_server_request_aborted(now_unix_timestamp_ms(), entry.request.id().clone()); if let Some(error) = error.as_ref()