diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index bde4e48f00d0..671c2a2f8d25 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -233,8 +233,13 @@ where .map(|(&method, route)| (method, route)) } - pub(crate) fn notification_route(&self, method: &str) -> Option<&NotificationRoute> { - self.notification_routes.get(method) + pub(crate) fn notification_route( + &self, + method: &str, + ) -> Option<(&'static str, &NotificationRoute)> { + self.notification_routes + .get_key_value(method) + .map(|(&method, route)| (method, route)) } } diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index 15343754c9ee..9c8d3e80db02 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -185,25 +185,50 @@ async fn run_connection( } } codex_exec_server_protocol::JSONRPCMessage::Notification(notification) => { - let Some(route) = router.notification_route(notification.method.as_str()) - else { - warn!( - "closing exec-server connection after unexpected notification: {}", - notification.method - ); - break; - }; - let result = tokio::select! { - result = route(Arc::clone(&handler), notification) => result, - _ = disconnected_rx.changed() => { - debug!( - "exec-server transport disconnected while handling notification" + let method = notification.method.clone(); + let notification_span = tracing::info_span!( + "codex.exec_server.notification", + otel.kind = "server", + otel.name = tracing::field::Empty, + rpc.system = "jsonrpc", + rpc.method = method.as_str(), + rpc.transport = transport.metric_tag(), + method = method.as_str(), + result = tracing::field::Empty, + ); + let continue_connection = async { + let Some((method, route)) = router.notification_route(method.as_str()) + else { + tracing::Span::current().record("otel.name", "unknown"); + tracing::Span::current().record("result", "error"); + warn!( + "closing exec-server connection after unexpected notification: {}", + notification.method ); - break; + return false; + }; + tracing::Span::current().record("otel.name", method); + let result = tokio::select! { + result = route(Arc::clone(&handler), notification) => result, + _ = disconnected_rx.changed() => { + tracing::Span::current().record("result", "disconnected"); + debug!( + "exec-server transport disconnected while handling notification" + ); + return false; + } + }; + tracing::Span::current() + .record("result", if result.is_ok() { "success" } else { "error" }); + if let Err(err) = result { + warn!("closing exec-server connection after protocol error: {err}"); + return false; } - }; - if let Err(err) = result { - warn!("closing exec-server connection after protocol error: {err}"); + true + } + .instrument(notification_span) + .await; + if !continue_connection { break; } }