diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index a0db3b4027..ebac1d6879 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -75,6 +75,7 @@ use codex_login::default_client::ClientRedirectPolicy; use codex_login::default_client::add_originator_header; use codex_login::default_client::create_client_for_route; use codex_otel::SessionTelemetry; +use codex_otel::WEBSOCKET_CONTINUATION_COUNT_METRIC; use codex_otel::current_span_w3c_trace_context; use codex_protocol::ResponseItemId; use codex_protocol::auth::AuthMode; @@ -262,6 +263,7 @@ pub struct ModelClient { codex_responses_headers: Option>, event_sender: Option>, http_client_factory: HttpClientFactory, + restored_history: bool, } /// A turn-scoped streaming session created from a [`ModelClient`]. @@ -300,6 +302,12 @@ struct LastResponse { items_added: Vec, } +struct WebsocketContinuation { + response_id: String, + items: Vec, + from_untraced_warmup: bool, +} + #[derive(Debug, Default)] struct WebsocketSession { connection: Option, @@ -311,6 +319,7 @@ struct WebsocketSession { last_response_rx: Option>, last_response_from_untraced_warmup: bool, connection_reused: StdMutex, + continuation_reset_reason: Option<&'static str>, } // This is intentionally not a `PartialEq` implementation: request equality includes `input` and @@ -395,6 +404,19 @@ fn response_items_equal_ignoring_internal_metadata( } impl WebsocketSession { + fn reset(&mut self, reason: Option<&'static str>) { + // Per-socket backend metrics call a resend after reconnect "initial". + // Retain the loss reason across reconnects/turns until the next send. + let continuation_reset_reason = self + .continuation_reset_reason + .or_else(|| self.last_request.as_ref().and(reason)); + *self = Self { + auth_owner_generation: self.auth_owner_generation, + continuation_reset_reason, + ..Default::default() + }; + } + fn set_connection_reused(&self, connection_reused: bool) { *self .connection_reused @@ -495,9 +517,15 @@ impl ModelClient { codex_responses_headers: None, event_sender: None, http_client_factory, + restored_history: false, } } + pub(crate) fn with_restored_history(mut self, restored_history: bool) -> Self { + self.restored_history = restored_history; + self + } + pub(crate) fn with_session_context( mut self, prompt_cache_key_override: Option, @@ -543,10 +571,8 @@ impl ModelClient { let mut websocket_session = self.take_cached_websocket_session(); if websocket_session.auth_owner_generation != auth_owner_generation { // Drop the old owner's cache before this turn can establish fresh routing state. - websocket_session = WebsocketSession { - auth_owner_generation, - ..Default::default() - }; + websocket_session.reset(Some("other")); + websocket_session.auth_owner_generation = auth_owner_generation; } ModelClientSession { client: self.clone(), @@ -1283,17 +1309,6 @@ impl Drop for ModelClientSession { } impl ModelClientSession { - fn reset_websocket_session(&mut self) { - self.websocket_session.connection = None; - self.websocket_session.responses_headers.clear(); - self.websocket_session.connection_key = None; - self.websocket_session.last_request = None; - self.websocket_session.last_response_rx = None; - self.websocket_session.last_response_from_untraced_warmup = false; - self.websocket_session - .set_connection_reused(/*connection_reused*/ false); - } - #[allow(clippy::too_many_arguments)] /// Builds shared Responses API transport options and request-body options. /// @@ -1332,13 +1347,12 @@ impl ModelClientSession { /// Checks whether the current request is an incremental extension of the previous request. /// We only reuse an incremental input delta when non-input request fields are unchanged and - /// `input` is a strict extension of the previous known input. Server-returned output items + /// `input` extends or equals the previous known input. Server-returned output items /// are treated as part of the baseline so we do not resend them. fn get_incremental_items( &self, request: &ResponsesApiRequest, - last_response: Option<&LastResponse>, - allow_empty_delta: bool, + last_response: &LastResponse, ) -> Option> { let previous_request = self.websocket_session.last_request.as_ref()?; if !responses_request_properties_match(previous_request, request) { @@ -1346,8 +1360,7 @@ impl ModelClientSession { return None; } - let response_items = - last_response.map_or(&[][..], |response| response.items_added.as_slice()); + let response_items = &last_response.items_added; let previous_items_len = previous_request .input .len() @@ -1368,9 +1381,6 @@ impl ModelClientSession { trace!("incremental request failed, items didn't match"); return None; } - if !allow_empty_delta && incremental_items.is_empty() { - return None; - } Some(incremental_items.to_vec()) } @@ -1387,29 +1397,20 @@ impl ModelClientSession { fn prepare_websocket_request( &mut self, request: &ResponsesApiRequest, - ) -> (Option<(String, Vec)>, bool) { - let Some(last_response) = self.get_last_response() else { - return (None, false); - }; - let previous_response_id_from_untraced_warmup = - self.websocket_session.last_response_from_untraced_warmup; - let Some(incremental_items) = self.get_incremental_items( - request, - Some(&last_response), - /*allow_empty_delta*/ true, - ) else { - return (None, false); - }; + ) -> Option { + let last_response = self.get_last_response()?; + let items = self.get_incremental_items(request, &last_response)?; if last_response.response_id.is_empty() { trace!("incremental request failed, no previous response id"); - return (None, false); + return None; } - ( - Some((last_response.response_id, incremental_items)), - previous_response_id_from_untraced_warmup, - ) + Some(WebsocketContinuation { + response_id: last_response.response_id, + items, + from_untraced_warmup: self.websocket_session.last_response_from_untraced_warmup, + }) } /// Opportunistically preconnects a websocket for this turn-scoped client session. @@ -1440,7 +1441,7 @@ impl ModelClientSession { { return Ok(()); } - self.reset_websocket_session(); + self.websocket_session.reset(Some("other")); let auth_context = AuthRequestTelemetryContext::new( client_setup.auth.as_ref().map(CodexAuth::auth_mode), client_setup.api_auth.as_ref(), @@ -1493,14 +1494,17 @@ impl ModelClientSession { responses_headers, } = params; let connection_key = ResponsesConnectionKey::new(&api_provider, auth_revision); - let needs_new = match self.websocket_session.connection.as_ref() { - Some(conn) => { - self.websocket_session.responses_headers != *responses_headers - || self.websocket_session.connection_key.as_ref() != Some(&connection_key) - || conn.is_closed().await + let reset_reason = match self.websocket_session.connection.as_ref() { + Some(_) + if self.websocket_session.responses_headers != *responses_headers + || self.websocket_session.connection_key.as_ref() != Some(&connection_key) => + { + Some("other") } - None => true, + Some(conn) if conn.is_closed().await => Some("connection_closed"), + Some(_) | None => None, }; + let needs_new = self.websocket_session.connection.is_none() || reset_reason.is_some(); // Resolving an external auth provider can change ownership during client setup. let owner_changed = self.websocket_session.auth_owner_generation != auth_owner_generation || self.client.auth_owner_generation() != auth_owner_generation; @@ -1509,8 +1513,12 @@ impl ModelClientSession { } if needs_new || owner_changed { - self.reset_websocket_session(); - let new_conn = match self + self.websocket_session.reset(if owner_changed { + Some("other") + } else { + reset_reason + }); + let new_conn = self .client .connect_websocket( session_telemetry, @@ -1521,16 +1529,7 @@ impl ModelClientSession { request_route_telemetry, responses_headers, ) - .await - { - Ok(new_conn) => new_conn, - Err(err) => { - if matches!(err, ApiError::Transport(TransportError::Timeout)) { - self.reset_websocket_session(); - } - return Err(err); - } - }; + .await?; self.websocket_session.connection = Some(new_conn); self.websocket_session.responses_headers = responses_headers.clone(); self.websocket_session.auth_owner_generation = auth_owner_generation; @@ -1882,8 +1881,26 @@ impl ModelClientSession { if let Some(turn_state) = self.turn_state.get() { client_metadata.insert(X_CODEX_TURN_STATE_HEADER.to_string(), turn_state.clone()); } - let (incremental_request, previous_response_id_from_untraced_warmup) = - self.prepare_websocket_request(&request); + let continuation = self.prepare_websocket_request(&request); + let (mode, reason) = if continuation.is_some() { + ("incremental", "incremental") + } else { + let reason = self + .websocket_session + .continuation_reset_reason + .take() + .unwrap_or(if self.websocket_session.last_request.is_some() { + "other" + } else if self.client.restored_history { + "restored_history" + } else { + "no_previous_request" + }); + ("full", reason) + }; + let previous_response_id_from_untraced_warmup = continuation + .as_ref() + .is_some_and(|c| c.from_untraced_warmup); let inference_trace_attempt = if warmup { // Prewarm sends `generate=false`; it is connection setup, not a // model inference attempt that should appear in rollout traces. @@ -1898,8 +1915,8 @@ impl ModelClientSession { inference_trace_attempt.record_started(&request); } - let (previous_response_id, mut incremental_items) = match incremental_request { - Some((response_id, items)) => (Some(response_id), Some(items)), + let (previous_response_id, mut incremental_items) = match continuation { + Some(continuation) => (Some(continuation.response_id), Some(continuation.items)), None => (None, None), }; let original_item_ids = if let Some(incremental_items) = &mut incremental_items { @@ -1944,6 +1961,15 @@ impl ModelClientSession { "websocket connection is unavailable".to_string(), )) })?; + request_session_telemetry.counter( + WEBSOCKET_CONTINUATION_COUNT_METRIC, + /*inc*/ 1, + &[ + ("mode", mode), + ("reason", reason), + ("phase", if warmup { "warmup" } else { "generation" }), + ], + ); let stream_result = websocket_connection .stream_request( ws_request, diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 76895776ee..b36b23a095 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -645,13 +645,16 @@ fn websocket_incremental_reuse_tracks_raw_result_metadata() -> anyhow::Result<() }) .unwrap(); session.websocket_session.last_response_rx = Some(receiver); - let (incremental, from_warmup) = session.prepare_websocket_request(¤t); + let continuation = session.prepare_websocket_request(¤t); assert_eq!( - incremental, - expect_incremental.then_some(("previous-response".to_string(), vec![follow_up])), + continuation.map(|continuation| ( + continuation.response_id, + continuation.items, + continuation.from_untraced_warmup, + )), + expect_incremental.then_some(("previous-response".to_string(), vec![follow_up], false)), "{scenario}", ); - assert!(!from_warmup); } Ok(()) } @@ -809,6 +812,43 @@ fn test_session_telemetry() -> SessionTelemetry { ) } +#[test] +fn websocket_continuation_reset_reason_survives_failed_reconnect_and_turn_boundary() { + for (reason, later_reason) in [ + ("connection_closed", "other"), + ("other", "connection_closed"), + ] { + let client = test_model_client(SessionSource::Cli); + let request = client + .build_responses_request( + &Prompt::default(), + &test_model_info(), + /*effort*/ None, + codex_protocol::config_types::ReasoningSummary::None, + /*service_tier*/ None, + &test_responses_metadata_for_client( + &client, + /*turn_id*/ None, + format!("{}:0", client.state.thread_id), + /*parent_thread_id*/ None, + TestCodexResponsesRequestKind::Turn, + ), + ) + .expect("build continuation request"); + let mut session = client.new_session(); + session.websocket_session.last_request = Some(request); + session.websocket_session.reset(Some(reason)); + session.websocket_session.reset(/*reason*/ None); + session.websocket_session.reset(Some(later_reason)); + drop(session); + let session = client.new_session(); + assert_eq!( + session.websocket_session.continuation_reset_reason, + Some(reason) + ); + } +} + fn spawned_session_source() -> SessionSource { SessionSource::SubAgent(SubAgentSource::ThreadSpawn { parent_thread_id: ThreadId::new(), diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 8a20bd7fc5..9f6d02e1bc 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1244,6 +1244,9 @@ impl Session { ) .with_auth_env(auth_env_telemetry.to_otel_metadata()) .with_tool_result_log_config(config.otel.tool_result); + if let Some(metrics) = thread_extension_data.get::() { + session_telemetry = session_telemetry.with_metrics(metrics.as_ref().clone()); + } if let Some(service_name) = session_configuration.metrics_service_name.as_deref() { session_telemetry = session_telemetry.with_metrics_service_name(service_name); } @@ -1683,6 +1686,10 @@ impl Session { config.http_client_factory(), config.workspace_routing_context(), ) + .with_restored_history(matches!( + &initial_history, + InitialHistory::Resumed(_) | InitialHistory::Forked(_) + )) .with_session_context( crate::guardian::prompt_cache_key_override_for_review_session( &session_configuration.session_source, diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index 6a72bffb3b..9996efbaaa 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -697,6 +697,14 @@ async fn responses_websocket_request_prewarm_reuses_connection() { Some("sequential_cutoff") ); + assert_continuation_metrics( + &harness.session_telemetry, + &[ + (["full", "no_previous_request", "warmup"], 1), + (["incremental", "incremental", "generation"], 1), + ], + ); + server.shutdown().await; } @@ -1000,6 +1008,14 @@ async fn responses_websocket_reconnects_after_account_switch() { ], "new_turn_on_switch={new_turn_on_switch}", ); + assert_continuation_metrics( + &harness.session_telemetry, + &[ + (["full", "other", "generation"], 1), + (["full", "no_previous_request", "generation"], 1), + (["incremental", "incremental", "generation"], 3), + ], + ); server.shutdown().await; } } @@ -2063,6 +2079,13 @@ async fn responses_websocket_uses_incremental_create_on_prefix() { serde_json::to_value(&prompt_two.input[2..]).expect("serialize incremental items") ); + assert_continuation_metrics( + &harness.session_telemetry, + &[ + (["full", "no_previous_request", "generation"], 1), + (["incremental", "incremental", "generation"], 1), + ], + ); server.shutdown().await; } @@ -2329,6 +2352,13 @@ async fn responses_websocket_creates_on_non_prefix() { serde_json::to_value(&prompt_two.input).unwrap() ); + assert_continuation_metrics( + &harness.session_telemetry, + &[ + (["full", "other", "generation"], 1), + (["full", "no_previous_request", "generation"], 1), + ], + ); server.shutdown().await; } @@ -2365,6 +2395,13 @@ async fn responses_websocket_creates_when_non_input_request_fields_change() { serde_json::to_value(&prompt_two.input).expect("serialize full input") ); + assert_continuation_metrics( + &harness.session_telemetry, + &[ + (["full", "no_previous_request", "generation"], 1), + (["full", "other", "generation"], 1), + ], + ); server.shutdown().await; } @@ -2536,6 +2573,15 @@ async fn responses_websocket_v2_after_error_uses_full_create_without_previous_re serde_json::to_value(&prompt_three.input).unwrap() ); + assert_continuation_metrics( + &harness.session_telemetry, + &[ + (["full", "connection_closed", "generation"], 1), + (["full", "no_previous_request", "generation"], 1), + (["incremental", "incremental", "generation"], 1), + ], + ); + server.shutdown().await; } @@ -2853,6 +2899,91 @@ async fn websocket_harness_with_provider_options_and_auth( } } +#[test_case::test_case(false; "resume")] +#[test_case::test_case(true; "fork")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn responses_websocket_restored_history_metric(fork: bool) -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let exporter = InMemoryMetricExporter::default(); + let config = + MetricsConfig::in_memory("test", "codex-core", "test", exporter).with_runtime_reader(); + let metrics = MetricsClient::new(config)?; + let warmup = vec![ev_response_created("warmup"), ev_completed("warmup")]; + let turn = vec![ev_response_created("turn"), ev_completed("turn")]; + let server = start_websocket_server(vec![vec![warmup, turn]; 2]).await; + let mut initial = test_codex().build_with_websocket_server(&server).await?; + initial.submit_text_turn("original history marker").await?; + let rollout_path = initial.codex.rollout_path().expect("persisted history"); + initial.codex.shutdown_and_wait().await?; + + let manager = &initial.thread_manager; + let mut options = codex_core::StartThreadOptions::new(initial.config.clone()); + options.thread_extension_init.insert(metrics); + let restored = if fork { + manager + .fork_thread(codex_core::ForkSnapshot::Interrupted, options, rollout_path) + .await? + } else { + options.initial_history = + codex_rollout::RolloutRecorder::get_rollout_history(&rollout_path).await?; + manager.start_thread(options).await? + }; + initial.codex = restored.thread; + initial.submit_text_turn("continue").await?; + assert_continuation_metrics( + &initial.codex.session_telemetry(), + &[ + (["full", "restored_history", "warmup"], 1), + (["incremental", "incremental", "generation"], 1), + ], + ); + assert!(server.connections()[1].iter().any(|request| { + request.body_json()["input"] + .to_string() + .contains("original history marker") + })); + initial.codex.shutdown_and_wait().await?; + server.shutdown().await; + Ok(()) +} + +fn assert_continuation_metrics(telemetry: &SessionTelemetry, expected: &[([&str; 3], 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"); + let AggregatedMetrics::U64(MetricData::Sum(sum)) = metric.data() else { + panic!("expected counter"); + }; + let mut expected: Vec<_> = expected + .iter() + .map(|(tags, count)| (tags.map(str::to_owned), *count)) + .collect(); + let mut actual: Vec<_> = sum + .data_points() + .map(|point| { + let tags = ["mode", "reason", "phase"].map(|key| { + point + .attributes() + .find(|attr| attr.key.as_str() == key) + .unwrap() + .value + .to_string() + }); + (tags, point.value()) + }) + .collect(); + actual.sort(); + expected.sort(); + assert_eq!(actual, expected); +} + async fn stream_until_complete( client_session: &mut ModelClientSession, harness: &WebsocketTestHarness, diff --git a/codex-rs/otel/README.md b/codex-rs/otel/README.md index 9deab4d511..68d05d0b61 100644 --- a/codex-rs/otel/README.md +++ b/codex-rs/otel/README.md @@ -144,6 +144,26 @@ metrics.counter("codex.turns", 1, &[("model", "gpt-5.1")])?; metrics.shutdown()?; // flushes in-memory exporter ``` +## WebSocket continuation + +`codex.websocket.continuation` counts `response.create` send attempts, including +failed sends. It carries existing session tags plus `mode` (`incremental`/`full`), +`phase` (`warmup`/`generation`), and `reason`: + +| Reason | Meaning | +| --- | --- | +| `incremental` | Send the previous response ID and new input. | +| `no_previous_request` | First request from a fresh client. | +| `restored_history` | First request after loading resumed or forked history. | +| `connection_closed` | Full input after observing the previous socket closed. | +| `other` | Full input for another reason, such as changed input/settings or unavailable response state. | + +Per-socket backend metrics label a resend after reconnect as `initial`; this client +metric retains the first reset reason through reconnect failures and turn boundaries. +Include warmups (`generate=false`), which can send full input before an incremental +generation. Closes may be intentional; restored history includes manual resumes/forks. +This measures client send attempts, not disconnect rates or engine cache reuse. + ## Trace context Trace propagation helpers remain separate from the session event emitter: diff --git a/codex-rs/otel/src/metrics/names.rs b/codex-rs/otel/src/metrics/names.rs index 1d4994e918..31b2aad09c 100644 --- a/codex-rs/otel/src/metrics/names.rs +++ b/codex-rs/otel/src/metrics/names.rs @@ -12,6 +12,7 @@ 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_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"; pub const WEBSOCKET_EVENT_COUNT_METRIC: &str = "codex.websocket.event"; pub const WEBSOCKET_EVENT_DURATION_METRIC: &str = "codex.websocket.event.duration_ms"; diff --git a/codex-rs/tui/src/app/tests/snapshots/codex_tui__app__tests__approvals_reviewer_error_tests__approvals_reviewer_config_error.snap b/codex-rs/tui/src/app/tests/snapshots/codex_tui__app__tests__approvals_reviewer_error_tests__approvals_reviewer_config_error.snap index ed4e880e4d..dcd2ad82fa 100644 --- a/codex-rs/tui/src/app/tests/snapshots/codex_tui__app__tests__approvals_reviewer_error_tests__approvals_reviewer_config_error.snap +++ b/codex-rs/tui/src/app/tests/snapshots/codex_tui__app__tests__approvals_reviewer_error_tests__approvals_reviewer_config_error.snap @@ -1,6 +1,6 @@ --- source: tui/src/app/tests/approvals_reviewer_error_tests.rs -assertion_line: 34 +assertion_line: 36 expression: rendered --- -■ Failed to save approvals reviewer: config/batchWrite failed in TUI: config/batchWrite failed: failed to load configuration: /config.toml:1:24: unclosed array, expected `]` (code -32603) +■ Failed to save approvals reviewer: config/batchWrite failed in TUI: config/batchWrite failed: failed to load configuration: /config.toml:1:23: unclosed array, expected `]` (code -32603)