Trace app-server requests sent to clients

This commit is contained in:
Adam Perry
2026-07-10 00:23:00 +00:00
parent ed116c2b2c
commit 9f614931a3
6 changed files with 186 additions and 7 deletions
@@ -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,
@@ -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<usize>,
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",
+1 -1
View File
@@ -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;
+1 -1
View File
@@ -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;
@@ -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<i64> {
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(())
}
+43 -3
View File
@@ -42,6 +42,8 @@ pub(crate) type ClientRequestResult = std::result::Result<Result, JSONRPCErrorEr
enum RpcResult {
Success,
Error,
Cancelled,
SendFailed,
Disconnected,
Replaced,
}
@@ -51,6 +53,8 @@ impl RpcResult {
match self {
Self::Success => "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<ClientRequestResult>,
thread_id: Option<ThreadId>,
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()