From fd1096ae73363da56b298dada784a44580f780d8 Mon Sep 17 00:00:00 2001 From: GLENNA ZHANG Date: Wed, 10 Jun 2026 14:44:35 -0700 Subject: [PATCH] propagate response execution model identity --- .../tests/suite/v2/safety_check_downgrade.rs | 107 +++++++++ codex-rs/codex-api/src/common.rs | 8 + .../src/endpoint/responses_websocket.rs | 14 ++ codex-rs/codex-api/src/lib.rs | 1 + codex-rs/codex-api/src/sse/responses.rs | 211 +++++++++++++++--- codex-rs/core/src/session/review.rs | 3 + codex-rs/core/src/session/turn.rs | 15 ++ codex-rs/core/src/session/turn_context.rs | 9 + codex-rs/core/src/tasks/mod.rs | 5 +- codex-rs/core/src/turn_timing.rs | 1 + .../core/tests/suite/client_websockets.rs | 104 +++++++++ .../tests/suite/safety_check_downgrade.rs | 74 +++++- codex-rs/otel/src/events/session_telemetry.rs | 1 + 13 files changed, 520 insertions(+), 33 deletions(-) diff --git a/codex-rs/app-server/tests/suite/v2/safety_check_downgrade.rs b/codex-rs/app-server/tests/suite/v2/safety_check_downgrade.rs index ea3cd8dada..e4c5efe11b 100644 --- a/codex-rs/app-server/tests/suite/v2/safety_check_downgrade.rs +++ b/codex-rs/app-server/tests/suite/v2/safety_check_downgrade.rs @@ -15,9 +15,11 @@ use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ThreadItem; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; +use codex_app_server_protocol::TurnCompletedNotification; use codex_app_server_protocol::TurnModerationMetadataNotification; use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; +use codex_app_server_protocol::TurnStatus; use codex_app_server_protocol::UserInput; use core_test_support::responses; use core_test_support::skip_if_no_network; @@ -29,10 +31,59 @@ use wiremock::ResponseTemplate; const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); const REQUESTED_MODEL: &str = "gpt-5.4"; const SERVER_MODEL: &str = "gpt-5.3-codex"; +const MODEL_SNAPSHOT: &str = "test-model-snapshot"; const TRUSTED_ACCESS_FOR_CYBER_VERIFICATION: &str = "trusted_access_for_cyber"; const CYBER_POLICY_MESSAGE: &str = "This request has been flagged for potentially high-risk cyber activity."; +#[tokio::test] +async fn response_identity_headers_propagate_to_turn_completed_v2() -> Result<()> { + skip_if_no_network!(Ok(())); + + let body = responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_assistant_message("msg-1", "Done"), + responses::ev_completed("resp-1"), + ]); + let response = responses::sse_response(body) + .insert_header("OpenAI-Model", SERVER_MODEL) + .insert_header("OpenAI-Model-Snapshot", MODEL_SNAPSHOT); + + let (thread_id, turn_id, completed, _) = run_turn_and_read_completion(response).await?; + + assert_eq!(completed.thread_id, thread_id); + assert_eq!(completed.turn.id, turn_id); + assert_eq!(completed.turn.status, TurnStatus::Completed); + assert_eq!(completed.final_model.as_deref(), Some(SERVER_MODEL)); + assert_eq!(completed.model_snapshot.as_deref(), Some(MODEL_SNAPSHOT)); + + Ok(()) +} + +#[tokio::test] +async fn missing_response_identity_headers_leave_turn_completed_fields_null() -> Result<()> { + skip_if_no_network!(Ok(())); + + let body = responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_assistant_message("msg-1", "Done"), + responses::ev_completed("resp-1"), + ]); + + let (thread_id, turn_id, completed, params) = + run_turn_and_read_completion(responses::sse_response(body)).await?; + + assert_eq!(completed.thread_id, thread_id); + assert_eq!(completed.turn.id, turn_id); + assert_eq!(completed.turn.status, TurnStatus::Completed); + assert_eq!(completed.final_model, None); + assert_eq!(completed.model_snapshot, None); + assert_eq!(params.get("finalModel"), Some(&serde_json::Value::Null)); + assert_eq!(params.get("modelSnapshot"), Some(&serde_json::Value::Null)); + + Ok(()) +} + #[tokio::test] async fn openai_model_header_mismatch_emits_model_rerouted_notification_v2() -> Result<()> { skip_if_no_network!(Ok(())); @@ -433,6 +484,62 @@ async fn collect_turn_notifications_and_validate_no_warning_item( } } +async fn run_turn_and_read_completion( + response: ResponseTemplate, +) -> Result<(String, String, TurnCompletedNotification, serde_json::Value)> { + let server = responses::start_mock_server().await; + let _response_mock = responses::mount_response_once(&server, response).await; + + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), &server.uri())?; + + let mut mcp = TestAppServer::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let thread_req = mcp + .send_thread_start_request(ThreadStartParams { + model: Some(REQUESTED_MODEL.to_string()), + ..Default::default() + }) + .await?; + let thread_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; + + let turn_req = mcp + .send_turn_start_request(TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![UserInput::Text { + text: "capture response identity".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }) + .await?; + let turn_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), + ) + .await??; + let TurnStartResponse { turn } = to_response::(turn_resp)?; + + let notification = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("turn/completed"), + ) + .await??; + let params = notification + .params + .ok_or_else(|| anyhow::anyhow!("turn/completed notification must include params"))?; + let completed = serde_json::from_value(params.clone())?; + + Ok((thread.id, turn.id, completed, params)) +} + async fn collect_model_verification_notifications_and_validate_no_warning_item( mcp: &mut TestAppServer, ) -> Result { diff --git a/codex-rs/codex-api/src/common.rs b/codex-rs/codex-api/src/common.rs index 17753e799a..4ffbec126d 100644 --- a/codex-rs/codex-api/src/common.rs +++ b/codex-rs/codex-api/src/common.rs @@ -69,6 +69,12 @@ pub struct MemorySummarizeOutput { pub memory_summary: String, } +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct ResponseModelIdentity { + pub final_model: Option, + pub model_snapshot: Option, +} + #[derive(Debug)] pub enum ResponseEvent { Created, @@ -77,6 +83,8 @@ pub enum ResponseEvent { /// Emitted when the server includes `OpenAI-Model` on the stream response. /// This can differ from the requested model when backend safety routing applies. ServerModel(String), + /// Authoritative execution identity reported for this response. + ServerModelIdentity(ResponseModelIdentity), /// Emitted when the server recommends additional account verification. ModelVerifications(Vec), /// Emitted when the server includes moderation metadata for first-party turn presentation. diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket.rs b/codex-rs/codex-api/src/endpoint/responses_websocket.rs index 44cb6a544b..8dc884ffce 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket.rs +++ b/codex-rs/codex-api/src/endpoint/responses_websocket.rs @@ -1,5 +1,6 @@ use crate::auth::SharedAuthProvider; use crate::common::ResponseEvent; +use crate::common::ResponseModelIdentity; use crate::common::ResponseStream; use crate::common::ResponsesWsRequest; use crate::error::ApiError; @@ -633,6 +634,7 @@ async fn run_websocket_response_stream( connection_reused: bool, ) -> Result<(), ApiError> { let mut last_server_model: Option = None; + let mut last_server_model_identity: Option = None; send_websocket_request( ws_stream, request_body, @@ -698,6 +700,18 @@ async fn run_websocket_response_stream( .await; last_server_model = Some(model); } + let server_model_identity = event.response_model_identity(); + if (server_model_identity.final_model.is_some() + || server_model_identity.model_snapshot.is_some()) + && last_server_model_identity.as_ref() != Some(&server_model_identity) + { + let _ = tx_event + .send(Ok(ResponseEvent::ServerModelIdentity( + server_model_identity.clone(), + ))) + .await; + last_server_model_identity = Some(server_model_identity); + } if let Some(verifications) = model_verifications && tx_event .send(Ok(ResponseEvent::ModelVerifications(verifications))) diff --git a/codex-rs/codex-api/src/lib.rs b/codex-rs/codex-api/src/lib.rs index d2d7679406..17d0d958c2 100644 --- a/codex-rs/codex-api/src/lib.rs +++ b/codex-rs/codex-api/src/lib.rs @@ -33,6 +33,7 @@ pub use crate::common::Reasoning; pub use crate::common::ReasoningContext; pub use crate::common::ResponseCreateWsRequest; pub use crate::common::ResponseEvent; +pub use crate::common::ResponseModelIdentity; pub use crate::common::ResponseStream; pub use crate::common::ResponsesApiRequest; pub use crate::common::ResponsesWsRequest; diff --git a/codex-rs/codex-api/src/sse/responses.rs b/codex-rs/codex-api/src/sse/responses.rs index ab1be640df..0ea3dd64c6 100644 --- a/codex-rs/codex-api/src/sse/responses.rs +++ b/codex-rs/codex-api/src/sse/responses.rs @@ -1,4 +1,5 @@ use crate::common::ResponseEvent; +use crate::common::ResponseModelIdentity; use crate::common::ResponseStream; use crate::error::ApiError; use crate::rate_limits::parse_all_rate_limits; @@ -24,6 +25,7 @@ use tracing::trace; const X_REASONING_INCLUDED_HEADER: &str = "x-reasoning-included"; const OPENAI_MODEL_HEADER: &str = "openai-model"; +const OPENAI_MODEL_SNAPSHOT_HEADER: &str = "openai-model-snapshot"; const REQUEST_ID_HEADER: &str = "x-request-id"; const TRUSTED_ACCESS_FOR_CYBER_VERIFICATION: &str = "trusted_access_for_cyber"; @@ -44,6 +46,11 @@ pub fn spawn_response_stream( .get(OPENAI_MODEL_HEADER) .and_then(|v| v.to_str().ok()) .map(ToString::to_string); + let server_model_snapshot = stream_response + .headers + .get(OPENAI_MODEL_SNAPSHOT_HEADER) + .and_then(|v| v.to_str().ok()) + .map(ToString::to_string); let reasoning_included = stream_response .headers .get(X_REASONING_INCLUDED_HEADER) @@ -63,8 +70,20 @@ pub fn spawn_response_stream( } let (tx_event, rx_event) = mpsc::channel::>(1600); tokio::spawn(async move { - if let Some(model) = server_model { - let _ = tx_event.send(Ok(ResponseEvent::ServerModel(model))).await; + if let Some(model) = server_model.as_ref() { + let _ = tx_event + .send(Ok(ResponseEvent::ServerModel(model.clone()))) + .await; + } + if server_model.is_some() || server_model_snapshot.is_some() { + let _ = tx_event + .send(Ok(ResponseEvent::ServerModelIdentity( + ResponseModelIdentity { + final_model: server_model, + model_snapshot: server_model_snapshot, + }, + ))) + .await; } for snapshot in rate_limit_snapshots { let _ = tx_event.send(Ok(ResponseEvent::RateLimits(snapshot))).await; @@ -163,27 +182,39 @@ impl ResponsesStreamEvent { &self.kind } - /// Returns the effective model reported by the server, if present. - /// - /// Precedence: - /// 1. `response.headers` for standard Responses stream events. - /// 2. top-level `headers` for websocket metadata events. - pub fn response_model(&self) -> Option { - let response_headers_model = self - .response + fn response_headers(&self) -> Option<&Value> { + self.response .as_ref() .and_then(|response| response.get("headers")) - .and_then(header_openai_model_value_from_json); + .filter(|headers| headers.is_object()) + .or_else(|| self.headers.as_ref().filter(|headers| headers.is_object())) + } - match response_headers_model { - Some(model) => Some(model), - None => self - .headers - .as_ref() - .and_then(header_openai_model_value_from_json), + /// Returns the execution-model identity reported by the server, if present. + /// + /// Both fields are read from one container: `response.headers` takes + /// precedence over the top-level `headers` used by websocket metadata. + pub fn response_model_identity(&self) -> ResponseModelIdentity { + let headers = self.response_headers(); + ResponseModelIdentity { + final_model: headers.and_then(header_openai_model_value_from_json), + model_snapshot: headers.and_then(header_openai_model_snapshot_value_from_json), } } + /// Preserves the legacy model-warning fallback across metadata containers. + pub(crate) fn response_model(&self) -> Option { + self.response + .as_ref() + .and_then(|response| response.get("headers")) + .and_then(header_openai_model_value_from_json) + .or_else(|| { + self.headers + .as_ref() + .and_then(header_openai_model_value_from_json) + }) + } + pub(crate) fn model_verifications(&self) -> Option> { if self.kind() != "response.metadata" { return None; @@ -220,6 +251,17 @@ fn header_openai_model_value_from_json(value: &Value) -> Option { }) } +fn header_openai_model_snapshot_value_from_json(value: &Value) -> Option { + let headers = value.as_object()?; + headers.iter().find_map(|(name, value)| { + if name.eq_ignore_ascii_case(OPENAI_MODEL_SNAPSHOT_HEADER) { + json_value_as_string(value) + } else { + None + } + }) +} + fn model_verifications_from_json_value(value: &Value) -> Option> { let verifications = value .as_array() @@ -418,6 +460,7 @@ pub async fn process_sse( let mut stream = stream.eventsource(); let mut response_error: Option = None; let mut last_server_model: Option = None; + let mut last_server_model_identity: Option = None; loop { let start = Instant::now(); @@ -471,6 +514,22 @@ pub async fn process_sse( } last_server_model = Some(model); } + let server_model_identity = event.response_model_identity(); + if (server_model_identity.final_model.is_some() + || server_model_identity.model_snapshot.is_some()) + && last_server_model_identity.as_ref() != Some(&server_model_identity) + { + if tx_event + .send(Ok(ResponseEvent::ServerModelIdentity( + server_model_identity.clone(), + ))) + .await + .is_err() + { + return; + } + last_server_model_identity = Some(server_model_identity); + } if let Some(verifications) = model_verifications && tx_event .send(Ok(ResponseEvent::ModelVerifications(verifications))) @@ -1068,6 +1127,10 @@ mod tests { OPENAI_MODEL_HEADER, HeaderValue::from_static(CYBER_RESTRICTED_MODEL_FOR_TESTS), ); + headers.insert( + OPENAI_MODEL_SNAPSHOT_HEADER, + HeaderValue::from_static(MODEL_SNAPSHOT_FOR_TESTS), + ); let bytes = stream::iter(Vec::>::new()); let stream_response = StreamResponse { status: StatusCode::OK, @@ -1094,6 +1157,20 @@ mod tests { } other => panic!("expected server model event, got {other:?}"), } + let event = stream + .rx_event + .recv() + .await + .expect("expected server model identity event") + .expect("expected ok event"); + assert_matches!( + event, + ResponseEvent::ServerModelIdentity(ResponseModelIdentity { + final_model: Some(final_model), + model_snapshot: Some(model_snapshot), + }) if final_model == CYBER_RESTRICTED_MODEL_FOR_TESTS + && model_snapshot == MODEL_SNAPSHOT_FOR_TESTS + ); } #[tokio::test] @@ -1131,6 +1208,18 @@ mod tests { .iter() .any(|event| matches!(event, ResponseEvent::ModelVerifications(_))) ); + assert!( + !events + .iter() + .any(|event| matches!(event, ResponseEvent::ServerModelIdentity(_))) + ); + assert!(events.iter().any(|event| matches!( + event, + ResponseEvent::Completed { + response_id, + .. + } if response_id == "resp-1" + ))); } #[tokio::test] @@ -1173,7 +1262,8 @@ mod tests { "response": { "id": "resp-1", "headers": { - "OpenAI-Model": CYBER_RESTRICTED_MODEL_FOR_TESTS + "OpenAI-Model": CYBER_RESTRICTED_MODEL_FOR_TESTS, + "OpenAI-Model-Snapshot": MODEL_SNAPSHOT_FOR_TESTS } } }), @@ -1186,14 +1276,22 @@ mod tests { ]) .await; - assert_eq!(events.len(), 3); + assert_eq!(events.len(), 4); assert_matches!( &events[0], ResponseEvent::ServerModel(model) if model == CYBER_RESTRICTED_MODEL_FOR_TESTS ); - assert_matches!(&events[1], ResponseEvent::Created); assert_matches!( - &events[2], + &events[1], + ResponseEvent::ServerModelIdentity(ResponseModelIdentity { + final_model: Some(final_model), + model_snapshot: Some(model_snapshot), + }) if final_model == CYBER_RESTRICTED_MODEL_FOR_TESTS + && model_snapshot == MODEL_SNAPSHOT_FOR_TESTS + ); + assert_matches!(&events[2], ResponseEvent::Created); + assert_matches!( + &events[3], ResponseEvent::Completed { response_id, token_usage: None, @@ -1278,13 +1376,17 @@ mod tests { "type": "response.metadata", "headers": { "openai-model": CYBER_RESTRICTED_MODEL_FOR_TESTS, + "OPENAI-MODEL-SNAPSHOT": MODEL_SNAPSHOT_FOR_TESTS, } })) .expect("expected event to deserialize"); assert_eq!( - ev.response_model().as_deref(), - Some(CYBER_RESTRICTED_MODEL_FOR_TESTS) + ev.response_model_identity(), + ResponseModelIdentity { + final_model: Some(CYBER_RESTRICTED_MODEL_FOR_TESTS.to_string()), + model_snapshot: Some(MODEL_SNAPSHOT_FOR_TESTS.to_string()), + } ); } @@ -1293,7 +1395,34 @@ mod tests { let ev: ResponsesStreamEvent = serde_json::from_value(json!({ "type": "response.created", "headers": { - "openai-model": "top-level-model" + "openai-model": "top-level-model", + "openai-model-snapshot": "top-level-snapshot" + }, + "response": { + "id": "resp-1", + "headers": { + "openai-model": CYBER_RESTRICTED_MODEL_FOR_TESTS, + "openai-model-snapshot": MODEL_SNAPSHOT_FOR_TESTS + } + } + })) + .expect("expected event to deserialize"); + + assert_eq!( + ev.response_model_identity(), + ResponseModelIdentity { + final_model: Some(CYBER_RESTRICTED_MODEL_FOR_TESTS.to_string()), + model_snapshot: Some(MODEL_SNAPSHOT_FOR_TESTS.to_string()), + } + ); + } + + #[test] + fn responses_stream_event_does_not_mix_header_containers() { + let ev: ResponsesStreamEvent = serde_json::from_value(json!({ + "type": "response.created", + "headers": { + "openai-model-snapshot": "top-level-snapshot" }, "response": { "id": "resp-1", @@ -1305,8 +1434,37 @@ mod tests { .expect("expected event to deserialize"); assert_eq!( - ev.response_model().as_deref(), - Some(CYBER_RESTRICTED_MODEL_FOR_TESTS) + ev.response_model_identity(), + ResponseModelIdentity { + final_model: Some(CYBER_RESTRICTED_MODEL_FOR_TESTS.to_string()), + model_snapshot: None, + } + ); + } + + #[test] + fn responses_stream_event_preserves_legacy_model_fallback() { + let ev: ResponsesStreamEvent = serde_json::from_value(json!({ + "type": "response.created", + "headers": { + "openai-model": "top-level-model" + }, + "response": { + "id": "resp-1", + "headers": { + "openai-model-snapshot": MODEL_SNAPSHOT_FOR_TESTS + } + } + })) + .expect("expected event to deserialize"); + + assert_eq!(ev.response_model().as_deref(), Some("top-level-model")); + assert_eq!( + ev.response_model_identity(), + ResponseModelIdentity { + final_model: None, + model_snapshot: Some(MODEL_SNAPSHOT_FOR_TESTS.to_string()), + } ); } @@ -1398,4 +1556,5 @@ mod tests { } const CYBER_RESTRICTED_MODEL_FOR_TESTS: &str = "gpt-5.3-codex"; + const MODEL_SNAPSHOT_FOR_TESTS: &str = "test-model-snapshot"; } diff --git a/codex-rs/core/src/session/review.rs b/codex-rs/core/src/session/review.rs index 00a89b2b5c..578d33da69 100644 --- a/codex-rs/core/src/session/review.rs +++ b/codex-rs/core/src/session/review.rs @@ -3,6 +3,8 @@ use codex_core_skills::HostLoadedSkills; use codex_protocol::openai_models::ToolMode; use std::sync::atomic::AtomicBool; +use crate::session::turn_context::ModelExecutionIdentity; + /// Spawn a review thread using the given prompt. pub(super) async fn spawn_review_thread( sess: Arc, @@ -152,6 +154,7 @@ pub(super) async fn spawn_review_thread( extension_data, turn_skills: TurnSkillsContext::new(parent_turn_context.turn_skills.outcome.clone()), turn_timing_state: Arc::new(TurnTimingState::default()), + model_execution_identity: Arc::new(Mutex::new(ModelExecutionIdentity::default())), server_model_warning_emitted: AtomicBool::new(false), model_verification_emitted: AtomicBool::new(false), }; diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index f65ada6112..ea2d30d861 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -40,6 +40,7 @@ use crate::responses_retry::handle_retryable_response_stream_error; use crate::session::PreviousTurnSettings; use crate::session::TurnInput; use crate::session::session::Session; +use crate::session::turn_context::ModelExecutionIdentity; use crate::session::turn_context::TurnContext; use crate::stream_events_utils::HandleOutputCtx; use crate::stream_events_utils::TurnItemContributorPolicy; @@ -140,6 +141,7 @@ pub(crate) async fn run_turn( prewarmed_client_session: Option, cancellation_token: CancellationToken, ) -> Option { + *turn_context.model_execution_identity.lock().await = ModelExecutionIdentity::default(); let mut client_session = prewarmed_client_session.unwrap_or_else(|| sess.services.model_client.new_session()); // TODO(ccunningham): Pre-turn compaction runs before context updates and the @@ -242,6 +244,7 @@ pub(crate) async fn run_turn( let SamplingRequestResult { needs_follow_up: model_needs_follow_up, last_agent_message: sampling_request_last_agent_message, + model_execution_identity: sampling_request_model_execution_identity, } = sampling_request_output; can_drain_pending_input = true; let (has_pending_input, token_status, estimated_token_count) = async { @@ -326,6 +329,8 @@ pub(crate) async fn run_turn( .await; } } + *turn_context.model_execution_identity.lock().await = + sampling_request_model_execution_identity; if stop_outcome.should_stop { break; } @@ -1183,6 +1188,7 @@ pub(crate) async fn built_tools( struct SamplingRequestResult { needs_follow_up: bool, last_agent_message: Option, + model_execution_identity: ModelExecutionIdentity, } /// Ephemeral per-response state for streaming a single proposed plan. @@ -1797,6 +1803,7 @@ async fn try_run_sampling_request( FuturesOrdered::new(); let mut needs_follow_up = false; let mut last_agent_message: Option = None; + let mut model_execution_identity = ModelExecutionIdentity::default(); let mut active_item: Option = None; let mut active_tool_argument_diff_consumer: Option<( String, @@ -1945,6 +1952,7 @@ async fn try_run_sampling_request( break Ok(SamplingRequestResult { needs_follow_up: true, last_agent_message, + model_execution_identity, }); } } @@ -2033,6 +2041,12 @@ async fn try_run_sampling_request( .store(true, Ordering::Relaxed); } } + ResponseEvent::ServerModelIdentity(identity) => { + model_execution_identity = ModelExecutionIdentity { + final_model: identity.final_model, + model_snapshot: identity.model_snapshot, + }; + } ResponseEvent::ModelVerifications(verifications) => { if !turn_context .model_verification_emitted @@ -2081,6 +2095,7 @@ async fn try_run_sampling_request( break Ok(SamplingRequestResult { needs_follow_up, last_agent_message, + model_execution_identity, }); } ResponseEvent::OutputTextDelta(delta) => { diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index 10eff3f3b4..127891b17e 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -43,6 +43,12 @@ pub(crate) struct TurnEnvironment { pub(crate) shell: Option, } +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub(crate) struct ModelExecutionIdentity { + pub(crate) final_model: Option, + pub(crate) model_snapshot: Option, +} + impl TurnEnvironment { pub(crate) fn selection(&self) -> TurnEnvironmentSelection { TurnEnvironmentSelection { @@ -102,6 +108,7 @@ pub struct TurnContext { pub(crate) extension_data: Arc, pub(crate) turn_skills: TurnSkillsContext, pub(crate) turn_timing_state: Arc, + pub(crate) model_execution_identity: Arc>, pub(crate) server_model_warning_emitted: AtomicBool, pub(crate) model_verification_emitted: AtomicBool, } @@ -270,6 +277,7 @@ impl TurnContext { extension_data: Arc::clone(&self.extension_data), turn_skills: self.turn_skills.clone(), turn_timing_state: Arc::clone(&self.turn_timing_state), + model_execution_identity: Arc::clone(&self.model_execution_identity), server_model_warning_emitted: AtomicBool::new( self.server_model_warning_emitted.load(Ordering::Relaxed), ), @@ -572,6 +580,7 @@ impl Session { extension_data, turn_skills: TurnSkillsContext::new(skills_outcome), turn_timing_state: Arc::new(TurnTimingState::default()), + model_execution_identity: Arc::new(Mutex::new(ModelExecutionIdentity::default())), server_model_warning_emitted: AtomicBool::new(false), model_verification_emitted: AtomicBool::new(false), } diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index e33d781cf8..24391490ec 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -739,11 +739,12 @@ impl Session { }); self.emit_turn_stop_lifecycle(turn_context.extension_data.as_ref()) .await; + let model_execution_identity = turn_context.model_execution_identity.lock().await.clone(); let event = EventMsg::TurnComplete(TurnCompleteEvent { turn_id: turn_context.sub_id.clone(), last_agent_message, - final_model: None, - model_snapshot: None, + final_model: model_execution_identity.final_model, + model_snapshot: model_execution_identity.model_snapshot, completed_at, duration_ms, time_to_first_token_ms, diff --git a/codex-rs/core/src/turn_timing.rs b/codex-rs/core/src/turn_timing.rs index 380e246fae..1a375b4cfb 100644 --- a/codex-rs/core/src/turn_timing.rs +++ b/codex-rs/core/src/turn_timing.rs @@ -337,6 +337,7 @@ fn response_event_records_turn_ttft(event: &ResponseEvent) -> bool { | ResponseEvent::ReasoningContentDelta { .. } => true, ResponseEvent::Created | ResponseEvent::ServerModel(_) + | ResponseEvent::ServerModelIdentity(_) | ResponseEvent::ModelVerifications(_) | ResponseEvent::TurnModerationMetadata(_) | ResponseEvent::ServerReasoningIncluded(_) diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index 8b135b6b75..c16a58d3d2 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -1103,6 +1103,73 @@ async fn responses_websocket_emits_reasoning_included_event() { server.shutdown().await; } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn responses_websocket_uses_per_request_model_identity_metadata() { + skip_if_no_network!(); + + let first_identity = json!({ + "type": "codex.response.metadata", + "headers": { + "OpenAI-Model": "route-first", + "OpenAI-Model-Snapshot": "snapshot-first" + } + }); + let second_identity = json!({ + "type": "codex.response.metadata", + "headers": { + "OpenAI-Model": "route-terminal", + "OpenAI-Model-Snapshot": "snapshot-terminal" + } + }); + let server = start_websocket_server_with_headers(vec![WebSocketConnectionConfig { + requests: vec![ + vec![ + first_identity, + ev_response_created("resp-1"), + ev_completed("resp-1"), + ], + vec![ + second_identity, + ev_response_created("resp-2"), + ev_completed("resp-2"), + ], + ], + response_headers: vec![("OpenAI-Model".to_string(), "handshake-model".to_string())], + accept_delay: None, + close_after_requests: true, + }]) + .await; + + let harness = websocket_harness(&server).await; + let mut client_session = harness.client.new_session(); + let prompt = prompt_with_input(vec![message_item("hello")]); + + let first = + stream_until_complete_and_capture_model_identity(&mut client_session, &harness, &prompt) + .await; + let second = + stream_until_complete_and_capture_model_identity(&mut client_session, &harness, &prompt) + .await; + + assert_eq!( + first, + ( + Some("route-first".to_string()), + Some("snapshot-first".to_string()) + ) + ); + assert_eq!( + second, + ( + Some("route-terminal".to_string()), + Some("snapshot-terminal".to_string()) + ) + ); + assert_eq!(server.handshakes().len(), 1); + + server.shutdown().await; +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn responses_websocket_emits_rate_limit_events() { skip_if_no_network!(); @@ -2171,6 +2238,43 @@ async fn stream_until_complete_with_service_tier( .await; } +async fn stream_until_complete_and_capture_model_identity( + client_session: &mut ModelClientSession, + harness: &WebsocketTestHarness, + prompt: &Prompt, +) -> (Option, Option) { + let mut stream = client_session + .stream( + TEST_WINDOW_ID, + prompt, + &harness.model_info, + &harness.session_telemetry, + harness.effort.clone(), + harness.summary, + /*service_tier*/ None, + /*turn_metadata_header*/ None, + &InferenceTraceContext::disabled(), + ) + .await + .expect("websocket stream failed"); + let mut identity = (None, None); + + while let Some(event) = stream.next().await { + match event.expect("websocket response event should parse") { + ResponseEvent::ServerModelIdentity(response_identity) => { + identity = ( + response_identity.final_model, + response_identity.model_snapshot, + ); + } + ResponseEvent::Completed { .. } => return identity, + _ => {} + } + } + + panic!("websocket stream ended before completion"); +} + async fn stream_until_complete_with_turn_metadata( client_session: &mut ModelClientSession, harness: &WebsocketTestHarness, diff --git a/codex-rs/core/tests/suite/safety_check_downgrade.rs b/codex-rs/core/tests/suite/safety_check_downgrade.rs index a271702c9c..290367721d 100644 --- a/codex-rs/core/tests/suite/safety_check_downgrade.rs +++ b/codex-rs/core/tests/suite/safety_check_downgrade.rs @@ -29,7 +29,10 @@ use pretty_assertions::assert_eq; use wiremock::ResponseTemplate; const SERVER_MODEL: &str = "gpt-5.2"; +const TERMINAL_SERVER_MODEL: &str = "gpt-5.1-codex"; const REQUESTED_MODEL: &str = "gpt-5.3-codex"; +const FIRST_MODEL_SNAPSHOT: &str = "test-first-model-snapshot"; +const TERMINAL_MODEL_SNAPSHOT: &str = "test-terminal-model-snapshot"; const TRUSTED_ACCESS_FOR_CYBER_VERIFICATION: &str = "trusted_access_for_cyber"; const CYBER_POLICY_MESSAGE: &str = @@ -221,13 +224,15 @@ async fn openai_model_header_mismatch_only_emits_one_warning_per_turn() -> Resul ), core_test_support::responses::ev_completed("resp-1"), ])) - .insert_header("OpenAI-Model", SERVER_MODEL); + .insert_header("OpenAI-Model", SERVER_MODEL) + .insert_header("OpenAI-Model-Snapshot", FIRST_MODEL_SNAPSHOT); let second_response = sse_response(sse(vec![ ev_response_created("resp-2"), ev_assistant_message("msg-1", "done"), core_test_support::responses::ev_completed("resp-2"), ])) - .insert_header("OpenAI-Model", SERVER_MODEL); + .insert_header("OpenAI-Model", TERMINAL_SERVER_MODEL) + .insert_header("OpenAI-Model-Snapshot", TERMINAL_MODEL_SNAPSHOT); let _mock = mount_response_sequence(&server, vec![first_response, second_response]).await; let mut builder = test_codex().with_model(REQUESTED_MODEL); @@ -238,18 +243,77 @@ async fn openai_model_header_mismatch_only_emits_one_warning_per_turn() -> Resul .await?; let mut warning_count = 0; - loop { + let turn_complete = loop { let event = wait_for_event(&test.codex, |_| true).await; match event { EventMsg::Warning(warning) if warning.message.contains(REQUESTED_MODEL) => { warning_count += 1; } - EventMsg::TurnComplete(_) => break, + EventMsg::TurnComplete(turn_complete) => break turn_complete, _ => {} } - } + }; assert_eq!(warning_count, 1); + assert_eq!( + turn_complete.final_model.as_deref(), + Some(TERMINAL_SERVER_MODEL) + ); + assert_eq!( + turn_complete.model_snapshot.as_deref(), + Some(TERMINAL_MODEL_SNAPSHOT) + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn nonterminal_response_identity_is_not_reported_when_follow_up_fails() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let tool_args = serde_json::json!({ + "command": "echo hello", + "timeout_ms": 1_000 + }); + let first_response = sse_response(sse(vec![ + ev_response_created("resp-1"), + ev_function_call( + "call-1", + "shell_command", + &serde_json::to_string(&tool_args)?, + ), + core_test_support::responses::ev_completed("resp-1"), + ])) + .insert_header("OpenAI-Model", SERVER_MODEL) + .insert_header("OpenAI-Model-Snapshot", FIRST_MODEL_SNAPSHOT); + let failed_follow_up = ResponseTemplate::new(400).set_body_json(serde_json::json!({ + "error": { + "message": "synthetic follow-up failure", + "type": "invalid_request_error", + "param": null, + "code": "invalid_prompt" + } + })); + let _mock = mount_response_sequence(&server, vec![first_response, failed_follow_up]).await; + + let mut builder = test_codex().with_model(REQUESTED_MODEL); + let test = builder.build(&server).await?; + + test.codex + .submit(disabled_text_turn(&test, "trigger failed follow-up")) + .await?; + + let turn_complete = wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + let EventMsg::TurnComplete(turn_complete) = turn_complete else { + panic!("expected turn complete event"); + }; + + assert_eq!(turn_complete.final_model, None); + assert_eq!(turn_complete.model_snapshot, None); Ok(()) } diff --git a/codex-rs/otel/src/events/session_telemetry.rs b/codex-rs/otel/src/events/session_telemetry.rs index 00c13046dd..08e30a4996 100644 --- a/codex-rs/otel/src/events/session_telemetry.rs +++ b/codex-rs/otel/src/events/session_telemetry.rs @@ -1207,6 +1207,7 @@ impl SessionTelemetry { "reasoning_summary_part_added".into() } ResponseEvent::ServerModel(_) => "server_model".into(), + ResponseEvent::ServerModelIdentity(_) => "server_model_identity".into(), ResponseEvent::ModelVerifications(_) => "model_verifications".into(), ResponseEvent::TurnModerationMetadata(_) => "turn_moderation_metadata".into(), ResponseEvent::ServerReasoningIncluded(_) => "server_reasoning_included".into(),