diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 810380641e..f324db1723 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -4546,7 +4546,6 @@ dependencies = [ "tempfile", "thiserror 2.0.18", "tokio", - "tokio-tungstenite", "tracing", "tracing-opentelemetry", "tracing-subscriber", diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket.rs b/codex-rs/codex-api/src/endpoint/responses_websocket.rs index 412fb1f0ab..2275978862 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket.rs +++ b/codex-rs/codex-api/src/endpoint/responses_websocket.rs @@ -768,6 +768,18 @@ struct InterruptAcknowledgement { response_id: String, } +fn is_websocket_receive_failure( + response: &Result>, 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::(&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::(&text) { + let event = match event { Ok(event) => event, Err(err) => { debug!("failed to parse websocket event: {err}, data: {text}"); diff --git a/codex-rs/codex-api/src/telemetry.rs b/codex-rs/codex-api/src/telemetry.rs index 91918a65b9..f7589d7256 100644 --- a/codex-rs/codex-api/src/telemetry.rs +++ b/codex-rs/codex-api/src/telemetry.rs @@ -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>, 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 { diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index e1d80193d3..a9acbcdfd9 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -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"] } diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 86c41c9617..c43378900b 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -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>, 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); } } diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index 8a48eb40cc..1ff75b3cf4 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -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", diff --git a/codex-rs/otel/Cargo.toml b/codex-rs/otel/Cargo.toml index f1124348bd..5d60385356 100644 --- a/codex-rs/otel/Cargo.toml +++ b/codex-rs/otel/Cargo.toml @@ -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 } diff --git a/codex-rs/otel/src/events/session_telemetry.rs b/codex-rs/otel/src/events/session_telemetry.rs index 91b9b29314..74381ece77 100644 --- a/codex-rs/otel/src/events/session_telemetry.rs +++ b/codex-rs/otel/src/events/session_telemetry.rs @@ -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::(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::(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); } diff --git a/codex-rs/otel/tests/suite/runtime_summary.rs b/codex-rs/otel/tests/suite/runtime_summary.rs index 18722e278f..049a443111 100644 --- a/codex-rs/otel/tests/suite/runtime_summary.rs +++ b/codex-rs/otel/tests/suite/runtime_summary.rs @@ -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>, - 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>, - 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,