github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-20 00:18:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

Track WebSocket continuation modes and full-input send reasons (#46051)

Browse Source 
## Why

Per-socket metrics label resends after reconnect as initial requests, obscuring why full input was sent. Client telemetry can preserve that reason across reconnect failures and turn boundaries.

## What changed

- Add `codex.websocket.continuation` to count `response.create` send attempts, including failed sends, with `mode`, `reason`, and `phase` tags.
- Distinguish incremental continuations, fresh requests, restored history, closed connections, and other full-input sends. Track warmup and generation separately.
- Preserve the first continuation reset reason until the next send and identify resumed or forked history.
- Allow session telemetry to use a `MetricsClient` supplied through thread extensions and document the counter's semantics.

## Testing

Add metric assertions for incremental reuse, warmup, changed requests, account switches, reconnects, and resumed or forked history. Add a unit test for preserving reset reasons across repeated resets and turn boundaries.

GitOrigin-RevId: 29324f5e9d5677e0c6068f89b2256a92a5a0d157
This commit is contained in:
vkg-oai
2026-09-16 22:43:57 +00:00
committed by copyberry
parent 821ad43f9d
commit c56dda711c
7 changed files with 295 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/client.rs
+90 -64
View File
@@ -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<Arc<CodexResponsesHeaders>>,
event_sender: Option<Sender<ProtocolEvent>>,
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<ResponseItem>,
}
struct WebsocketContinuation {
response_id: String,
items: Vec<ResponseItem>,
from_untraced_warmup: bool,
}
#[derive(Debug, Default)]
struct WebsocketSession {
connection: Option<ApiWebSocketConnection>,
@@ -311,6 +319,7 @@ struct WebsocketSession {
last_response_rx: Option<oneshot::Receiver<LastResponse>>,
last_response_from_untraced_warmup: bool,
connection_reused: StdMutex<bool>,
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<String>,
@@ -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<Vec<ResponseItem>> {
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<ResponseItem>)>, 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<WebsocketContinuation> {
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,
codex-rs/core/src/client_tests.rs
+44 -4
View File
@@ -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(&current);
let continuation = session.prepare_websocket_request(&current);
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(),
codex-rs/core/src/session/session.rs
+7
View File
@@ -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::<codex_otel::MetricsClient>() {
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,
codex-rs/core/tests/suite/client_websockets.rs
+131
View File
@@ -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,
codex-rs/otel/README.md
+20
View File
@@ -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:
codex-rs/otel/src/metrics/names.rs
+1
View File
@@ -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";
codex-rs/tui/src/app/tests/snapshots/codex_tui__app__tests__approvals_reviewer_error_tests__approvals_reviewer_config_error.snap
+2 -2
View File
@@ -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: <CODEX_HOME>/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: <CODEX_HOME>/config.toml:1:23: unclosed array, expected `]` (code -32603)