From dec38fc9bb8304f3e34d00d7cee77f5c0488584f Mon Sep 17 00:00:00 2001 From: Ben Romano Date: Tue, 3 Nov 2026 00:36:49 +0000 Subject: [PATCH] Add WebSocket connection metrics with reauthentication labels (#59018) ## What changed - Record connection attempts and durations as `codex.websocket.connect` and `codex.websocket.connect.duration_ms`, labeled with `success`, `reauth`, and the configured `product_sku` when present. - Detect reauthentication by comparing the current auth revision with the last successful connection, preserving that revision across connection resets. Include `auth.reauth` in connection logs and traces. - Populate session telemetry with the configured product SKU. ## Testing Extend tests to check reauthentication labels after account switches, product SKU labels for restored sessions, and failed-handshake counts, durations, and labels. GitOrigin-RevId: 53a38b4460ab9021058dfe2fe104d0441993b18a --- codex-rs/core/src/client.rs | 13 ++- codex-rs/core/src/session/session.rs | 1 + .../core/tests/suite/client_websockets.rs | 86 ++++++++++++++----- codex-rs/otel/src/events/session_telemetry.rs | 16 +++- codex-rs/otel/src/metrics/names.rs | 2 + .../tests/suite/otel_export_routing_policy.rs | 30 ++++++- 6 files changed, 123 insertions(+), 25 deletions(-) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 889a62680f..c1f1027388 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -362,6 +362,9 @@ struct WebsocketContinuation { #[derive(Debug, Default)] struct WebsocketSession { connection: Option, + // Outer None means no successful handshake yet. The inner revision is optional for + // providers without managed auth. Keep it across resets so retries compare to the last socket. + last_connected_auth_revision: Option>, responses_headers: ApiHeaderMap, parent_routing_affinity: Option, runtime_id: Option, @@ -520,6 +523,7 @@ impl WebsocketSession { *self = Self { auth_owner_generation: self.auth_owner_generation, continuation_reset_reason, + last_connected_auth_revision: self.last_connected_auth_revision, ..Default::default() }; } @@ -1415,6 +1419,7 @@ impl ModelClient { request_route_telemetry: RequestRouteTelemetry, responses_headers: &ApiHeaderMap, parent_routing_affinity: Option<&codex_http_client::ChatGptRoutingAffinity>, + reauth: bool, ) -> std::result::Result { let mut headers = self.build_websocket_headers(responses_metadata).await; headers.extend(responses_headers.clone()); @@ -1474,7 +1479,7 @@ impl ModelClient { auth_context.recovery_mode, auth_context.recovery_phase, request_route_telemetry.endpoint, - /*connection_reused*/ false, + reauth, response_debug.request_id.as_deref(), response_debug.cf_ray.as_deref(), response_debug.auth_error.as_deref(), @@ -1837,6 +1842,10 @@ impl ModelClientSession { } else { reset_reason }); + let reauth = self + .websocket_session + .last_connected_auth_revision + .is_some_and(|previous| previous != auth_revision); let new_conn = self .client .connect_websocket( @@ -1849,9 +1858,11 @@ impl ModelClientSession { request_route_telemetry, responses_headers, parent_routing_affinity.as_ref(), + reauth, ) .await?; self.websocket_session.connection = Some(new_conn); + self.websocket_session.last_connected_auth_revision = Some(auth_revision); self.websocket_session.parent_routing_affinity = parent_routing_affinity; self.websocket_session.responses_headers = responses_headers.clone(); self.websocket_session.runtime_id = runtime_id; diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index ae17b9a009..160190b29a 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1396,6 +1396,7 @@ impl Session { session_configuration.session_source.clone(), ) .with_auth_env(auth_env_telemetry.to_otel_metadata()) + .with_product_sku(config.apps_mcp_product_sku.as_deref()) .with_user_id(telemetry_auth.and_then(CodexAuth::get_chatgpt_user_id)) .with_tool_result_log_config(config.otel.tool_result); if !config.otel.log_sse_deltas { diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index c48ae76402..eea77b19ee 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -69,6 +69,7 @@ use core_test_support::tracing::install_test_tracing; use core_test_support::wait_for_event; use futures::StreamExt; use opentelemetry_sdk::metrics::InMemoryMetricExporter; +use opentelemetry_sdk::metrics::data::ResourceMetrics; use pretty_assertions::assert_eq; use serde_json::json; use std::sync::Arc; @@ -1138,12 +1139,30 @@ async fn responses_websocket_reconnects_after_account_switch() { ], "new_turn_on_switch={new_turn_on_switch}", ); - assert_continuation_metrics( - &harness.session_telemetry, + let snapshot = harness + .session_telemetry + .snapshot_metrics() + .expect("metrics snapshot"); + assert_counter( + &snapshot, + codex_otel::WEBSOCKET_CONNECT_COUNT_METRIC, + ["reauth", "success"], + &[(["false", "true"], 1), (["true", "true"], 1)], + ); + assert_counter( + &snapshot, + codex_otel::WEBSOCKET_CONTINUATION_COUNT_METRIC, + CONTINUATION_TAGS, &[ - (["full", "other", "generation"], 1), - (["full", "no_previous_request", "generation"], 1), - (["incremental", "incremental", "generation"], 3), + (["full", "other", "generation", "none", "none"], 1), + ( + ["full", "no_previous_request", "generation", "none", "none"], + 1, + ), + ( + ["incremental", "incremental", "generation", "none", "none"], + 3, + ), ], ); server.shutdown().await; @@ -2936,6 +2955,7 @@ async fn responses_websocket_restored_history_metric(fork: bool) -> anyhow::Resu let manager = &initial.thread_manager; let mut options = codex_core::StartThreadOptions::new(initial.config.clone()); + options.config.apps_mcp_product_sku = Some("codex".to_string()); options.thread_extension_init.insert(metrics); let restored = if fork { manager @@ -2948,11 +2968,23 @@ async fn responses_websocket_restored_history_metric(fork: bool) -> anyhow::Resu }; initial.codex = restored.thread; initial.submit_text_turn("continue").await?; - assert_continuation_metrics( - &initial.codex.session_telemetry(), + let snapshot = initial.codex.session_telemetry().snapshot_metrics()?; + assert_counter( + &snapshot, + codex_otel::WEBSOCKET_CONNECT_COUNT_METRIC, + ["product_sku", "reauth", "success"], + &[(["codex", "false", "true"], 1)], + ); + assert_counter( + &snapshot, + codex_otel::WEBSOCKET_CONTINUATION_COUNT_METRIC, + CONTINUATION_TAGS, &[ - (["full", "restored_history", "warmup"], 1), - (["incremental", "incremental", "generation"], 1), + (["full", "restored_history", "warmup", "none", "none"], 1), + ( + ["incremental", "incremental", "generation", "none", "none"], + 1, + ), ], ); assert!(server.connections()[1].iter().any(|request| { @@ -2973,19 +3005,40 @@ fn assert_continuation_metrics(telemetry: &SessionTelemetry, expected: &[([&str; assert_continuation_metrics_with_item_types(telemetry, &expected); } +const CONTINUATION_TAGS: [&str; 5] = [ + "mode", + "reason", + "phase", + "previous_item_type", + "current_item_type", +]; + fn assert_continuation_metrics_with_item_types( telemetry: &SessionTelemetry, expected: &[([&str; 5], u64)], +) { + assert_counter( + &telemetry.snapshot_metrics().expect("metrics snapshot"), + codex_otel::WEBSOCKET_CONTINUATION_COUNT_METRIC, + CONTINUATION_TAGS, + expected, + ); +} + +fn assert_counter( + snapshot: &ResourceMetrics, + name: &str, + tag_names: [&str; N], + expected: &[([&str; N], u64)], ) { use opentelemetry_sdk::metrics::data::AggregatedMetrics; use opentelemetry_sdk::metrics::data::MetricData; - let snapshot = telemetry.snapshot_metrics().expect("metrics snapshot"); let metric = snapshot .scope_metrics() .flat_map(opentelemetry_sdk::metrics::data::ScopeMetrics::metrics) - .find(|metric| metric.name() == codex_otel::WEBSOCKET_CONTINUATION_COUNT_METRIC) - .expect("continuation counter"); + .find(|metric| metric.name() == name) + .unwrap_or_else(|| panic!("missing counter {name}")); let AggregatedMetrics::U64(MetricData::Sum(sum)) = metric.data() else { panic!("expected counter"); }; @@ -2996,14 +3049,7 @@ fn assert_continuation_metrics_with_item_types( let mut actual: Vec<_> = sum .data_points() .map(|point| { - let tags = [ - "mode", - "reason", - "phase", - "previous_item_type", - "current_item_type", - ] - .map(|key| { + let tags = tag_names.map(|key| { point .attributes() .find(|attr| attr.key.as_str() == key) diff --git a/codex-rs/otel/src/events/session_telemetry.rs b/codex-rs/otel/src/events/session_telemetry.rs index 43affe902d..9b68177003 100644 --- a/codex-rs/otel/src/events/session_telemetry.rs +++ b/codex-rs/otel/src/events/session_telemetry.rs @@ -27,6 +27,8 @@ use crate::metrics::TOOL_CALL_COUNT_METRIC; use crate::metrics::TOOL_CALL_DURATION_METRIC; use crate::metrics::TURN_COST_MICROUSD_METRIC; use crate::metrics::TURN_TTFT_DURATION_METRIC; +use crate::metrics::WEBSOCKET_CONNECT_COUNT_METRIC; +use crate::metrics::WEBSOCKET_CONNECT_DURATION_METRIC; use crate::metrics::WEBSOCKET_EVENT_COUNT_METRIC; use crate::metrics::WEBSOCKET_EVENT_DURATION_METRIC; use crate::metrics::WEBSOCKET_REQUEST_COUNT_METRIC; @@ -834,7 +836,7 @@ impl SessionTelemetry { recovery_mode: Option<&str>, recovery_phase: Option<&str>, endpoint: &str, - connection_reused: bool, + reauth: bool, request_id: Option<&str>, cf_ray: Option<&str>, auth_error: Option<&str>, @@ -846,6 +848,15 @@ impl SessionTelemetry { .map(|code| (200..=299).contains(&code)) .unwrap_or(true); let success_str = if success { "true" } else { "false" }; + let mut tags = vec![ + ("success", success_str), + ("reauth", if reauth { "true" } else { "false" }), + ]; + if let Some(product_sku) = self.metadata.product_sku { + tags.push(("product_sku", product_sku)); + } + self.counter(WEBSOCKET_CONNECT_COUNT_METRIC, /*inc*/ 1, &tags); + self.record_duration(WEBSOCKET_CONNECT_DURATION_METRIC, duration, &tags); log_and_trace_event!( self, common: { @@ -866,7 +877,8 @@ impl SessionTelemetry { auth.env_provider_key_name = self.metadata.auth_env.provider_env_key_name.as_deref(), auth.env_provider_key_present = self.metadata.auth_env.provider_env_key_present, auth.env_refresh_token_url_override_present = self.metadata.auth_env.refresh_token_url_override_present, - auth.connection_reused = connection_reused, + auth.connection_reused = false, + auth.reauth = reauth, auth.request_id = request_id, auth.cf_ray = cf_ray, auth.error = auth_error, diff --git a/codex-rs/otel/src/metrics/names.rs b/codex-rs/otel/src/metrics/names.rs index e0d60d99f0..999f89eee8 100644 --- a/codex-rs/otel/src/metrics/names.rs +++ b/codex-rs/otel/src/metrics/names.rs @@ -16,6 +16,8 @@ pub const API_CALL_COUNT_METRIC: &str = "codex.api_request"; pub const API_CALL_DURATION_METRIC: &str = "codex.api_request.duration_ms"; pub const SSE_EVENT_COUNT_METRIC: &str = "codex.sse_event"; pub const SSE_EVENT_DURATION_METRIC: &str = "codex.sse_event.duration_ms"; +pub const WEBSOCKET_CONNECT_COUNT_METRIC: &str = "codex.websocket.connect"; +pub const WEBSOCKET_CONNECT_DURATION_METRIC: &str = "codex.websocket.connect.duration_ms"; pub const WEBSOCKET_REQUEST_COUNT_METRIC: &str = "codex.websocket.request"; pub const WEBSOCKET_CONTINUATION_COUNT_METRIC: &str = "codex.websocket.continuation"; pub const WEBSOCKET_REQUEST_DURATION_METRIC: &str = "codex.websocket.request.duration_ms"; diff --git a/codex-rs/otel/tests/suite/otel_export_routing_policy.rs b/codex-rs/otel/tests/suite/otel_export_routing_policy.rs index 897e3f420c..92ca1afc3f 100644 --- a/codex-rs/otel/tests/suite/otel_export_routing_policy.rs +++ b/codex-rs/otel/tests/suite/otel_export_routing_policy.rs @@ -1062,6 +1062,7 @@ fn otel_export_routing_policy_routes_websocket_connect_auth_observability() { .with_filter(filter_fn(OtelProvider::trace_export_filter)), ); + let (metrics, metric_exporter) = build_metrics_with_defaults(&[]).unwrap(); tracing::subscriber::with_default(subscriber, || { tracing::callsite::rebuild_interest_cache(); let manager = SessionTelemetry::new( @@ -1076,7 +1077,9 @@ fn otel_export_routing_policy_routes_websocket_connect_auth_observability() { "tty".to_string(), SessionSource::Cli, ) - .with_auth_env(auth_env_metadata()); + .with_auth_env(auth_env_metadata()) + .with_product_sku(Some("codex")) + .with_metrics_without_metadata_tags(metrics); let root_span = tracing::info_span!("root"); let _root_guard = root_span.enter(); let agent_identity_telemetry = AgentIdentityTelemetry { @@ -1093,15 +1096,38 @@ fn otel_export_routing_policy_routes_websocket_connect_auth_observability() { Some("managed"), Some("reload"), "/responses", - /*connection_reused*/ false, + /*reauth*/ true, Some("req-ws-401"), Some("ray-ws-401"), Some("missing_authorization_header"), Some("token_expired"), Some(&agent_identity_telemetry), ); + manager.shutdown_metrics().unwrap(); }); + let exported = latest_metrics(&metric_exporter); + let expected_labels = BTreeMap::from([ + ("product_sku".to_string(), "codex".to_string()), + ("reauth".to_string(), "true".to_string()), + ("success".to_string(), "false".to_string()), + ]); + let counter = find_metric(&exported, codex_otel::WEBSOCKET_CONNECT_COUNT_METRIC).unwrap(); + let AggregatedMetrics::U64(MetricData::Sum(sum)) = counter.data() else { + panic!("expected connect counter"); + }; + let point = sum.data_points().next().expect("failed handshake"); + assert_eq!(point.value(), 1); + assert_eq!(attributes_to_map(point.attributes()), expected_labels); + + let duration = find_metric(&exported, codex_otel::WEBSOCKET_CONNECT_DURATION_METRIC).unwrap(); + let AggregatedMetrics::F64(MetricData::Histogram(histogram)) = duration.data() else { + panic!("expected connect duration histogram"); + }; + let point = histogram.data_points().next().expect("failed handshake"); + assert_eq!((point.count(), point.sum()), (1, 17.0)); + assert_eq!(attributes_to_map(point.attributes()), expected_labels); + logger_provider.force_flush().expect("flush logs"); tracer_provider.force_flush().expect("flush traces");