From 8193c56a595f66eb0f77f18d7434765eb0179d20 Mon Sep 17 00:00:00 2001 From: Dylan Hurd Date: Tue, 18 Aug 2026 17:55:05 +0000 Subject: [PATCH] Give Guardian classifier connections distinct thread identities (#39246) ## Why Guardian sampling uses a pool of WebSockets, while the bridge routes each connection by thread ID. Each socket therefore needs its own identity instead of reusing the source thread ID across the pool. ## What changed - Generate a classifier thread ID for every pooled connection and use it for request, window, and subagent attribution. - Preserve the originating thread ID in turn metadata alongside the Guardian classifier request kind. - Build request metadata after leasing a connection so retries use the identity of the socket that carries the request. ## Testing - Cover unique identities across pooled and replacement connections. - Verify classifier attribution metadata for extension requests and retries. GitOrigin-RevId: 1c36207125157ff05cab9487a2fa37bf85e5fdd8 --- .../ext/guardian-v2/src/extension_tests.rs | 26 +++- codex-rs/ext/guardian-v2/src/sampler.rs | 57 +++++-- codex-rs/ext/guardian-v2/src/sampler_tests.rs | 140 ++++++++++++++++-- 3 files changed, 196 insertions(+), 27 deletions(-) diff --git a/codex-rs/ext/guardian-v2/src/extension_tests.rs b/codex-rs/ext/guardian-v2/src/extension_tests.rs index 6389f62602..af3947819f 100644 --- a/codex-rs/ext/guardian-v2/src/extension_tests.rs +++ b/codex-rs/ext/guardian-v2/src/extension_tests.rs @@ -26,6 +26,7 @@ use codex_login::ExternalAuthFuture; use codex_login::ExternalAuthRefreshContext; use codex_model_provider_info::ModelProviderInfo; use codex_protocol::ResponseItemId; +use codex_protocol::ThreadId; use codex_protocol::models::ContentItem; use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::FunctionCallOutputPayload; @@ -1042,9 +1043,30 @@ async fn contributor_samples_tool_calls_with_the_existing_luna_pool() -> Result< let session_store = ExtensionData::new("session-1"); let thread_store = test.codex.thread_extension_data(); assert_eq!(request["model"], "gpt-5.6-luna"); + let classifier_thread_id = request["client_metadata"]["thread_id"] + .as_str() + .expect("classifier thread ID"); + assert_ne!(ThreadId::from_string(classifier_thread_id)?, thread_id); + let turn_metadata: serde_json::Value = serde_json::from_str( + request["client_metadata"]["x-codex-turn-metadata"] + .as_str() + .expect("serialized turn metadata"), + )?; assert_eq!( - request["client_metadata"]["thread_id"], - thread_id.to_string() + turn_metadata, + json!({ + "session_id": request["client_metadata"]["session_id"], + "thread_id": classifier_thread_id, + "guardian_classifier_source_thread_id": thread_id.to_string(), + "turn_id": "turn-1", + "request_kind": "guardian_classifier", + "is_guardian_mode": true, + }) + ); + assert_eq!(request["client_metadata"]["x-openai-subagent"], "guardian"); + assert_eq!( + request["client_metadata"]["x-codex-window-id"], + format!("{classifier_thread_id}:0") ); assert_eq!(request["client_metadata"]["turn_id"], "turn-1"); assert_eq!(request["reasoning"]["effort"], "low"); diff --git a/codex-rs/ext/guardian-v2/src/sampler.rs b/codex-rs/ext/guardian-v2/src/sampler.rs index 1953a770f7..7fbf96e772 100644 --- a/codex-rs/ext/guardian-v2/src/sampler.rs +++ b/codex-rs/ext/guardian-v2/src/sampler.rs @@ -27,6 +27,7 @@ use codex_login::default_client::default_headers; use codex_model_provider::AgentIdentitySessionFallback; use codex_model_provider::ProviderAuthScope; use codex_model_provider::SharedModelProvider; +use codex_protocol::ThreadId; use codex_protocol::error::CodexErr; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; @@ -35,6 +36,7 @@ use codex_protocol::protocol::SessionSource; use http::HeaderValue; use http::StatusCode; use serde_json::Value; +use serde_json::json; use thiserror::Error; use tokio::sync::OwnedSemaphorePermit; use tokio::sync::Semaphore; @@ -49,6 +51,7 @@ const MAX_WEBSOCKET_AGE: Duration = Duration::from_secs(55 * 60); const RESPONSES_WEBSOCKETS_BETA: &str = "responses_websockets=2026-02-06"; const RESPONSES_LITE_METADATA_KEY: &str = "ws_request_header_x_openai_internal_codex_responses_lite"; +const TURN_METADATA_KEY: &str = "x-codex-turn-metadata"; /// Host-owned provider, authentication, and attribution for one Luna connection. pub struct LunaSamplerConfig { @@ -117,7 +120,9 @@ pub enum LunaSamplerError { struct PooledConnection { connection: ResponsesWebsocketConnection, - connected_at: Instant, + // The bridge routes by thread ID, so each socket needs its own identity. + thread_id: String, + expires_at: Instant, auth_changes: Option>, } @@ -192,9 +197,19 @@ impl LunaSampler { .await .map_err(LunaSamplerError::Provider)? .auth; + let thread_id = ThreadId::new().to_string(); let mut headers = build_session_headers( Some(self.config.session_id.clone()), - Some(self.config.thread_id.clone()), + Some(thread_id.clone()), + ); + headers.insert("x-openai-subagent", HeaderValue::from_static("guardian")); + headers.insert( + "x-codex-window-id", + HeaderValue::from_str(&format!("{thread_id}:0")).map_err(|error| { + LunaSamplerError::Api(ApiError::Stream(format!( + "invalid classifier window ID: {error}" + ))) + })?, ); headers.insert( "openai-beta", @@ -207,7 +222,7 @@ impl LunaSampler { if let Some(originator) = self.config.originator.as_deref() { add_originator_header(&mut headers, originator); } - if let Ok(request_id) = HeaderValue::from_str(&self.config.thread_id) { + if let Ok(request_id) = HeaderValue::from_str(&thread_id) { headers.insert("x-client-request-id", request_id); } @@ -258,7 +273,8 @@ impl LunaSampler { Ok(PooledConnection { connection, - connected_at: Instant::now(), + thread_id, + expires_at: Instant::now() + MAX_WEBSOCKET_AGE, auth_changes, }) } @@ -280,7 +296,7 @@ impl LunaSampler { .auth_changes .as_ref() .is_none_or(|auth| !auth.has_changed().unwrap_or(true)) - && connection.connected_at.elapsed() < MAX_WEBSOCKET_AGE + && Instant::now() < connection.expires_at && !connection.connection.is_closed().await => { break connection; @@ -355,12 +371,7 @@ impl LunaSampler { /// Sends one structured, tool-less request on an exclusively leased WebSocket. pub async fn sample(&self, request: LunaSamplingRequest) -> Result { - let metadata = HashMap::from([ - ("session_id".to_owned(), self.config.session_id.clone()), - ("thread_id".to_owned(), self.config.thread_id.clone()), - ("turn_id".to_owned(), request.turn_id), - (RESPONSES_LITE_METADATA_KEY.to_owned(), "true".to_owned()), - ]); + let turn_id = request.turn_id; let mut input = vec![ ResponseItem::AdditionalTools { id: None, @@ -405,7 +416,7 @@ impl LunaSampler { phase: None, internal_chat_message_metadata_passthrough: None, }); - let request = ResponsesApiRequest { + let mut request = ResponsesApiRequest { model: MODEL.to_owned(), instructions: String::new(), input, @@ -428,7 +439,7 @@ impl LunaSampler { &Some(request.output_schema), /*output_schema_strict*/ true, ), - client_metadata: Some(metadata), + client_metadata: None, }; let (supersede, mut superseded) = oneshot::channel(); let scored = Arc::new(AtomicBool::new(false)); @@ -475,6 +486,26 @@ impl LunaSampler { return Err(error); } }; + let thread_id = &lease.connection.thread_id; + let turn_metadata = json!({ + "session_id": self.config.session_id, + "thread_id": thread_id, + "guardian_classifier_source_thread_id": self.config.thread_id, + "turn_id": turn_id, + "request_kind": "guardian_classifier", + "is_guardian_mode": true, + }) + .to_string(); + request.client_metadata = Some(HashMap::from([ + ("session_id".to_owned(), self.config.session_id.clone()), + ("thread_id".to_owned(), thread_id.clone()), + ("turn_id".to_owned(), turn_id.clone()), + ("x-openai-subagent".to_owned(), "guardian".to_owned()), + // Classifier requests do not advance their own context window. + ("x-codex-window-id".to_owned(), format!("{thread_id}:0")), + (RESPONSES_LITE_METADATA_KEY.to_owned(), "true".to_owned()), + (TURN_METADATA_KEY.to_owned(), turn_metadata), + ])); let mut stream = match lease .connection .connection diff --git a/codex-rs/ext/guardian-v2/src/sampler_tests.rs b/codex-rs/ext/guardian-v2/src/sampler_tests.rs index 08fb38c719..6470ec2668 100644 --- a/codex-rs/ext/guardian-v2/src/sampler_tests.rs +++ b/codex-rs/ext/guardian-v2/src/sampler_tests.rs @@ -10,6 +10,7 @@ use codex_login::ExternalAuthRefreshContext; use codex_model_provider::create_model_provider; use codex_model_provider_info::ModelProviderInfo; use codex_protocol::ResponseItemId; +use codex_protocol::ThreadId; use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::SessionSource; @@ -22,6 +23,7 @@ use core_test_support::responses::ev_output_text_delta; use core_test_support::skip_if_no_network; use pretty_assertions::assert_eq; use serde_json::json; +use std::collections::HashSet; use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; @@ -35,6 +37,54 @@ use super::LunaSamplingRequest; use super::MAX_SAMPLING_RETRIES; use super::MAX_WEBSOCKET_CONNECTIONS; +fn assert_connection_metadata(server: &responses::WebSocketTestServer) -> Result { + let handshake = server.single_handshake(); + let thread_id = handshake.header("thread-id").expect("classifier thread ID"); + ThreadId::from_string(&thread_id)?; + assert_ne!(thread_id, "thread-1"); + assert_eq!( + [ + handshake.header("x-client-request-id"), + handshake.header("x-openai-subagent"), + handshake.header("x-codex-window-id"), + ], + [ + Some(thread_id.clone()), + Some("guardian".to_owned()), + Some(format!("{thread_id}:0")), + ] + ); + for request in server.single_connection() { + let mut metadata = request.body_json()["client_metadata"].clone(); + let turn_id = metadata["turn_id"].clone(); + metadata["x-codex-turn-metadata"] = serde_json::from_str( + metadata["x-codex-turn-metadata"] + .as_str() + .expect("serialized turn metadata"), + )?; + assert_eq!( + metadata, + json!({ + "session_id": "session-1", + "thread_id": thread_id, + "turn_id": turn_id, + "x-openai-subagent": "guardian", + "x-codex-window-id": format!("{thread_id}:0"), + "ws_request_header_x_openai_internal_codex_responses_lite": "true", + "x-codex-turn-metadata": { + "session_id": "session-1", + "thread_id": thread_id, + "guardian_classifier_source_thread_id": "thread-1", + "turn_id": turn_id, + "request_kind": "guardian_classifier", + "is_guardian_mode": true, + }, + }) + ); + } + Ok(thread_id) +} + async fn proxy_websocket_servers(servers: &[&responses::WebSocketTestServer]) -> Result { let listener = TcpListener::bind("127.0.0.1:0").await?; let address = listener.local_addr()?; @@ -174,7 +224,6 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_structured_requ Some("true".to_owned()) ); assert_eq!(handshake.header("session-id"), Some("session-1".to_owned())); - assert_eq!(handshake.header("thread-id"), Some("thread-1".to_owned())); assert_eq!( handshake.header("originator"), Some("guardian-v2-test".to_owned()) @@ -237,6 +286,11 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_structured_requ ); requests.extend(refreshed.single_connection()); assert_eq!(requests.len(), 2); + let thread_id = assert_connection_metadata(&server)?; + assert_ne!( + Some(thread_id), + idle_server.single_handshake().header("thread-id") + ); assert_eq!( requests[0].body_json()["input"][2]["content"], json!([ @@ -252,20 +306,14 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_structured_requ assert_eq!(request["tool_choice"], "none"); assert_eq!(request["text"]["format"]["strict"], true); assert_eq!(request["prompt_cache_key"], "guardian-v2:thread-1"); - assert_eq!(request["client_metadata"]["session_id"], "session-1"); - assert_eq!(request["client_metadata"]["thread_id"], "thread-1"); assert_eq!( - request["client_metadata"]["ws_request_header_x_openai_internal_codex_responses_lite"], - "true" + request["client_metadata"]["turn_id"], + format!("turn-{}", index + 1) ); assert!(request.get("tools").is_none()); let effort = if index == 0 { "none" } else { "medium" }; assert_eq!(request["reasoning"]["effort"], effort); assert_eq!(request["reasoning"]["context"], "all_turns"); - assert_eq!( - request["client_metadata"]["turn_id"], - format!("turn-{}", index + 1) - ); } Ok(()) @@ -493,9 +541,26 @@ async fn sampler_grows_its_pool_for_overlapping_requests() -> Result<()> { r#"{"score":0.25}"#.to_owned(), ) ); + let mut thread_ids = HashSet::new(); + let mut turn_ids = HashSet::new(); for server in [&first, &second, &third] { assert_eq!(server.single_connection().len(), 1); + assert!(thread_ids.insert(assert_connection_metadata(server)?)); + turn_ids.insert( + server.single_connection()[0].body_json()["client_metadata"]["turn_id"] + .as_str() + .expect("turn ID") + .to_owned(), + ); } + assert_eq!( + turn_ids, + HashSet::from_iter([ + "turn-1".to_owned(), + "turn-2".to_owned(), + "turn-3".to_owned() + ]) + ); Ok(()) } @@ -629,10 +694,61 @@ async fn sampler_retries_expired_websockets_on_another_warm_connection() -> Resu let healthy_requests = healthy.single_connection(); assert_eq!(expired_requests.len(), 1); assert_eq!(healthy_requests.len(), 1); - assert_eq!( - expired_requests[0].body_json(), - healthy_requests[0].body_json() + assert_ne!( + assert_connection_metadata(&expired)?, + assert_connection_metadata(&healthy)? ); + let mut expired_request = expired_requests[0].body_json(); + let mut healthy_request = healthy_requests[0].body_json(); + for request in [&mut expired_request, &mut healthy_request] { + assert_eq!(request["client_metadata"]["turn_id"], "turn-1"); + request + .as_object_mut() + .expect("request object") + .remove("client_metadata"); + } + assert_eq!(expired_request, healthy_request); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn sampler_assigns_a_fresh_identity_when_replacing_aged_connections() -> Result<()> { + skip_if_no_network!(Ok(())); + + let response = vec![vec![ + ev_assistant_message("response-1", r#"{"score":0.25}"#), + ev_completed("response-1"), + ]]; + let first = responses::start_websocket_server(vec![response.clone()]).await; + let second = responses::start_websocket_server(vec![response.clone()]).await; + let replacement = responses::start_websocket_server(vec![response]).await; + let sampler = LunaSampler::connect(sampler_config( + proxy_websocket_servers(&[&first, &second, &replacement]).await?, + )) + .await?; + + assert_eq!( + sampler.sample(sample_request("turn-1")).await?, + r#"{"score":0.25}"# + ); + { + let mut connections = sampler.idle_connections.lock().unwrap(); + for connection in connections.iter_mut() { + connection.expires_at = std::time::Instant::now(); + } + } + assert_eq!( + sampler.sample(sample_request("turn-2")).await?, + r#"{"score":0.25}"# + ); + + let thread_ids = [&first, &second, &replacement] + .into_iter() + .map(assert_connection_metadata) + .collect::>>()?; + assert_eq!(thread_ids.len(), 3); + assert_eq!(second.single_connection().len(), 1); + assert_eq!(replacement.single_connection().len(), 1); Ok(()) }