Reuse decoded WebSocket event kinds in telemetry (#59769)

## Why

WebSocket telemetry reparses text frames to extract event kinds that the transport already decodes.

## What changed

Pass decoded event kinds to telemetry and report receive failures separately. Keep JSON parsing for timing metrics and as a fallback when stream event decoding fails. Remove the `tokio-tungstenite` dependency from `codex-otel` and restrict it to tests in `codex-core`.

Limit wrapped error parsing to error events or failed stream decoding, preserving HTTP error forwarding when unrelated stream fields are malformed.

## Testing

Update runtime summary coverage to include receive failures and extend the WebSocket error forwarding test with a malformed stream field.

GitOrigin-RevId: e1f9d200fd6d30ec124a2354bdbc01f3082ddde3
This commit is contained in:
Charlie Marsh
2026-11-06 00:24:03 +00:00
committed by copyberry
parent 08e4f43c9a
commit 74417e9c6b
9 changed files with 94 additions and 89 deletions
-1
View File
@@ -4546,7 +4546,6 @@ dependencies = [
"tempfile",
"thiserror 2.0.18",
"tokio",
"tokio-tungstenite",
"tracing",
"tracing-opentelemetry",
"tracing-subscriber",
@@ -768,6 +768,18 @@ struct InterruptAcknowledgement {
response_id: String,
}
fn is_websocket_receive_failure(
response: &Result<Option<Result<Message, WsError>>, ApiError>,
) -> bool {
match response {
Ok(Some(Ok(message))) => match message {
Message::Text(_) | Message::Ping(_) | Message::Pong(_) => false,
Message::Binary(_) | Message::Close(_) | Message::Frame(_) => true,
},
Ok(Some(Err(_))) | Ok(None) | Err(_) => true,
}
}
#[allow(clippy::too_many_arguments)]
async fn run_websocket_response_stream(
ws_stream: &mut WsStream,
@@ -827,8 +839,11 @@ async fn run_websocket_response_stream(
continue;
}
};
if let Some(t) = telemetry.as_ref() {
t.on_ws_event(&response, poll_start.elapsed());
let poll_duration = poll_start.elapsed();
if let Some(telemetry) = &telemetry
&& is_websocket_receive_failure(&response)
{
telemetry.on_ws_receive_failure(poll_duration);
}
let message = match response {
Ok(Some(Ok(msg))) => msg,
@@ -847,14 +862,21 @@ async fn run_websocket_response_stream(
match message {
Message::Text(text) => {
if let Some(wrapped_error) = parse_wrapped_websocket_error_event(&text)
let event = serde_json::from_str::<ResponsesStreamEvent>(&text);
let kind = event.as_ref().ok().map(ResponsesStreamEvent::kind);
if let Some(t) = telemetry.as_ref() {
t.on_ws_text_event(&text, kind, poll_duration);
}
// A wrapped error can still be valid when an unrelated stream field is malformed.
if matches!(kind, Some("error") | None)
&& let Some(wrapped_error) = parse_wrapped_websocket_error_event(&text)
&& let Some(error) =
map_wrapped_websocket_error_event(wrapped_error, text.to_string())
{
return Err(error);
}
let event = match serde_json::from_str::<ResponsesStreamEvent>(&text) {
let event = match event {
Ok(event) => event,
Err(err) => {
debug!("failed to parse websocket event: {err}, data: {text}");
+7 -7
View File
@@ -11,8 +11,6 @@ use std::future::Future;
use std::sync::Arc;
use std::time::Duration;
use tokio::time::Instant;
use tokio_tungstenite::tungstenite::Error;
use tokio_tungstenite::tungstenite::Message;
/// Generic telemetry.
pub trait SseTelemetry: Send + Sync {
@@ -33,13 +31,15 @@ pub trait SseTelemetry: Send + Sync {
/// Telemetry for Responses WebSocket transport.
pub trait WebsocketTelemetry: Send + Sync {
/// Reports how long it took to send the request and whether the send failed.
fn on_ws_request(&self, duration: Duration, error: Option<&ApiError>, connection_reused: bool);
fn on_ws_event(
&self,
result: &Result<Option<Result<Message, Error>>, ApiError>,
duration: Duration,
);
/// Reports the poll duration for timeouts, stream errors, end of stream,
/// and binary, close, or raw frames.
fn on_ws_receive_failure(&self, duration: Duration);
/// Reports each text frame and its poll duration. `kind` is `None` if decoding the stream event failed.
fn on_ws_text_event(&self, text: &str, kind: Option<&str>, duration: Duration);
}
pub(crate) trait WithStatus {
+1 -1
View File
@@ -123,7 +123,6 @@ tokio = { workspace = true, features = [
"signal",
] }
tokio-util = { workspace = true, features = ["rt"] }
tokio-tungstenite = { workspace = true }
toml = { workspace = true }
toml_edit = { workspace = true }
tonic = { workspace = true }
@@ -177,6 +176,7 @@ serial_test = { workspace = true }
tempfile = { workspace = true }
test-log = { workspace = true }
tokio-rustls = { workspace = true }
tokio-tungstenite = { workspace = true }
tracing-opentelemetry = { workspace = true }
tracing-subscriber = { workspace = true }
tracing-test = { workspace = true, features = ["no-env-filter"] }
+7 -8
View File
@@ -113,8 +113,6 @@ use std::time::Instant;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio::sync::oneshot::error::TryRecvError;
use tokio_tungstenite::tungstenite::Error;
use tokio_tungstenite::tungstenite::Message;
use tokio_util::sync::CancellationToken;
use tracing::instrument;
use tracing::trace;
@@ -3204,13 +3202,14 @@ impl WebsocketTelemetry for ApiTelemetry {
);
}
fn on_ws_event(
&self,
result: &std::result::Result<Option<std::result::Result<Message, Error>>, ApiError>,
duration: Duration,
) {
fn on_ws_receive_failure(&self, duration: Duration) {
self.session_telemetry
.record_websocket_event(result, duration);
.record_websocket_receive_failure(duration);
}
fn on_ws_text_event(&self, text: &str, kind: Option<&str>, duration: Duration) {
self.session_telemetry
.record_websocket_text_event(text, kind, duration);
}
}
@@ -1885,6 +1885,8 @@ async fn responses_websocket_invalid_request_error_with_status_is_forwarded() {
let invalid_request_error = json!({
"type": "error",
// An unexpected stream field must not hide the wrapped HTTP error.
"delta": {"unexpected": true},
"status": 400,
"error": {
"type": "invalid_request_error",
-1
View File
@@ -56,7 +56,6 @@ serde_json = { workspace = true }
strum_macros = { workspace = true }
thiserror = { workspace = true }
tokio = { workspace = true }
tokio-tungstenite = { workspace = true }
tracing = { workspace = true }
tracing-opentelemetry = { workspace = true }
tracing-subscriber = { workspace = true }
+38 -49
View File
@@ -40,7 +40,6 @@ use crate::sanitize_metric_tag_value;
use crate::tool_result::ToolResultEvent;
use crate::tool_result::emit_tool_result;
use codex_api::AgentIdentityTelemetry;
use codex_api::ApiError;
use codex_api::ResponseEvent;
use codex_protocol::AgentPath;
use codex_protocol::ThreadId;
@@ -969,62 +968,52 @@ impl SessionTelemetry {
);
}
pub fn record_websocket_event(
pub fn record_websocket_receive_failure(&self, duration: Duration) {
self.record_websocket_event_metrics(
WEBSOCKET_UNKNOWN_KIND,
/*success*/ false,
duration,
);
}
/// Records the transport's decoded event kind. Timing events also carry metric values.
pub fn record_websocket_text_event(
&self,
result: &Result<
Option<
Result<
tokio_tungstenite::tungstenite::Message,
tokio_tungstenite::tungstenite::Error,
>,
>,
ApiError,
>,
text: &str,
parsed_kind: Option<&str>,
duration: Duration,
) {
let mut kind = None;
let mut success = true;
match result {
Ok(Some(Ok(message))) => match message {
tokio_tungstenite::tungstenite::Message::Text(text) => {
match serde_json::from_str::<serde_json::Value>(text) {
Ok(value) => {
kind = value
.get("type")
.and_then(|value| value.as_str())
.map(std::string::ToString::to_string);
if kind.as_deref() == Some(RESPONSES_WEBSOCKET_TIMING_KIND) {
self.record_responses_websocket_timing_metrics(&value);
}
if kind.as_deref() == Some("response.failed") {
success = false;
}
}
Err(_) => {
kind = Some("parse_error".to_string());
success = false;
}
let value;
let kind = match parsed_kind {
Some(kind) if kind != RESPONSES_WEBSOCKET_TIMING_KIND => kind,
_ => {
// Valid JSON with an unknown or malformed stream shape still contributes its kind.
value = match serde_json::from_str::<serde_json::Value>(text) {
Ok(value) => value,
Err(_) => {
return self.record_websocket_event_metrics(
"parse_error",
/*success*/ false,
duration,
);
}
};
let kind = value
.get("type")
.and_then(serde_json::Value::as_str)
.unwrap_or(WEBSOCKET_UNKNOWN_KIND);
if kind == RESPONSES_WEBSOCKET_TIMING_KIND {
self.record_responses_websocket_timing_metrics(&value);
}
tokio_tungstenite::tungstenite::Message::Ping(_)
| tokio_tungstenite::tungstenite::Message::Pong(_) => {
return;
}
tokio_tungstenite::tungstenite::Message::Binary(_)
| tokio_tungstenite::tungstenite::Message::Close(_)
| tokio_tungstenite::tungstenite::Message::Frame(_) => {
success = false;
}
},
Ok(Some(Err(_))) | Ok(None) | Err(_) => {
success = false;
kind
}
}
};
self.record_websocket_event_metrics(kind, kind != "response.failed", duration);
}
let kind_str = kind.as_deref().unwrap_or(WEBSOCKET_UNKNOWN_KIND);
fn record_websocket_event_metrics(&self, kind: &str, success: bool, duration: Duration) {
let success_str = if success { "true" } else { "false" };
let tags = [("kind", kind_str), ("success", success_str)];
let tags = [("kind", kind), ("success", success_str)];
self.counter(WEBSOCKET_EVENT_COUNT_METRIC, /*inc*/ 1, &tags);
self.record_duration(WEBSOCKET_EVENT_DURATION_METRIC, duration, &tags);
}
+13 -18
View File
@@ -12,7 +12,6 @@ use eventsource_stream::Event as StreamEvent;
use opentelemetry_sdk::metrics::InMemoryMetricExporter;
use pretty_assertions::assert_eq;
use std::time::Duration;
use tokio_tungstenite::tungstenite::Message;
#[test]
fn runtime_metrics_summary_collects_tool_api_and_streaming_metrics() -> Result<()> {
@@ -80,21 +79,17 @@ fn runtime_metrics_summary_collects_tool_api_and_streaming_metrics() -> Result<(
retry: None,
})));
manager.log_sse_event(&sse_response, Duration::from_millis(120));
let ws_response: std::result::Result<
Option<std::result::Result<Message, tokio_tungstenite::tungstenite::Error>>,
codex_api::ApiError,
> = Ok(Some(Ok(Message::Text(
r#"{"type":"response.created"}"#.into(),
))));
manager.record_websocket_event(&ws_response, Duration::from_millis(80));
let ws_timing_response: std::result::Result<
Option<std::result::Result<Message, tokio_tungstenite::tungstenite::Error>>,
codex_api::ApiError,
> = Ok(Some(Ok(Message::Text(
r#"{"type":"responsesapi.websocket_timing","timing_metrics":{"responses_duration_excl_engine_and_client_tool_time_ms":124.25,"engine_service_total_ms":457,"engine_iapi_ttft_total_ms":211,"engine_service_ttft_total_ms":233,"engine_iapi_tbt_across_engine_calls_ms":2.450638,"engine_service_tbt_across_engine_calls_ms":5.267279}}"#
.into(),
))));
manager.record_websocket_event(&ws_timing_response, Duration::from_millis(20));
manager.record_websocket_text_event(
r#"{"type":"response.created"}"#,
Some("response.created"),
Duration::from_millis(80),
);
manager.record_websocket_text_event(
r#"{"type":"responsesapi.websocket_timing","timing_metrics":{"responses_duration_excl_engine_and_client_tool_time_ms":124.25,"engine_service_total_ms":457,"engine_iapi_ttft_total_ms":211,"engine_service_ttft_total_ms":233,"engine_iapi_tbt_across_engine_calls_ms":2.450638,"engine_service_tbt_across_engine_calls_ms":5.267279}}"#,
Some("responsesapi.websocket_timing"),
Duration::from_millis(20),
);
manager.record_websocket_receive_failure(Duration::from_millis(10));
manager.record_duration(
"codex.turn.ttft.duration_ms",
Duration::from_millis(95),
@@ -127,8 +122,8 @@ fn runtime_metrics_summary_collects_tool_api_and_streaming_metrics() -> Result<(
duration_ms: 400,
},
websocket_events: RuntimeMetricTotals {
count: 2,
duration_ms: 100,
count: 3,
duration_ms: 110,
},
responses_api_overhead_ms: 124,
responses_api_inference_time_ms: 457,