mirror of
https://github.com/openai/codex.git
synced 2026-11-08 20:40:33 +00:00
Trace app-server requests sent to clients
This commit is contained in:
@@ -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",
|
||||
|
||||
@@ -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;
|
||||
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user