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
This commit is contained in:
Dylan Hurd authored and copyberry committed 2026-08-18 18:08:17 +00:00
1 parent a1dc95d5af
commit 8193c56a59
3 files changed
+196 -27

No files matched your search

@@ -26,6 +26,7 @@ use codex_login::ExternalAuthFuture;
use codex_login::ExternalAuthRefreshContext; use codex_login::ExternalAuthRefreshContext;
use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::ModelProviderInfo;
use codex_protocol::ResponseItemId; use codex_protocol::ResponseItemId;
use codex_protocol::ThreadId;
use codex_protocol::models::ContentItem; use codex_protocol::models::ContentItem;
use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::FunctionCallOutputContentItem;
use codex_protocol::models::FunctionCallOutputPayload; 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 session_store = ExtensionData::new("session-1");
let thread_store = test.codex.thread_extension_data(); let thread_store = test.codex.thread_extension_data();
assert_eq!(request["model"], "gpt-5.6-luna"); 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!( assert_eq!(
request["client_metadata"]["thread_id"], turn_metadata,
thread_id.to_string() 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["client_metadata"]["turn_id"], "turn-1");
assert_eq!(request["reasoning"]["effort"], "low"); assert_eq!(request["reasoning"]["effort"], "low");
+44 -13
View File
@@ -27,6 +27,7 @@ use codex_login::default_client::default_headers;
use codex_model_provider::AgentIdentitySessionFallback; use codex_model_provider::AgentIdentitySessionFallback;
use codex_model_provider::ProviderAuthScope; use codex_model_provider::ProviderAuthScope;
use codex_model_provider::SharedModelProvider; use codex_model_provider::SharedModelProvider;
use codex_protocol::ThreadId;
use codex_protocol::error::CodexErr; use codex_protocol::error::CodexErr;
use codex_protocol::models::ContentItem; use codex_protocol::models::ContentItem;
use codex_protocol::models::ResponseItem; use codex_protocol::models::ResponseItem;
@@ -35,6 +36,7 @@ use codex_protocol::protocol::SessionSource;
use http::HeaderValue; use http::HeaderValue;
use http::StatusCode; use http::StatusCode;
use serde_json::Value; use serde_json::Value;
use serde_json::json;
use thiserror::Error; use thiserror::Error;
use tokio::sync::OwnedSemaphorePermit; use tokio::sync::OwnedSemaphorePermit;
use tokio::sync::Semaphore; 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_WEBSOCKETS_BETA: &str = "responses_websockets=2026-02-06";
const RESPONSES_LITE_METADATA_KEY: &str = const RESPONSES_LITE_METADATA_KEY: &str =
"ws_request_header_x_openai_internal_codex_responses_lite"; "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. /// Host-owned provider, authentication, and attribution for one Luna connection.
pub struct LunaSamplerConfig { pub struct LunaSamplerConfig {
@@ -117,7 +120,9 @@ pub enum LunaSamplerError {
struct PooledConnection { struct PooledConnection {
connection: ResponsesWebsocketConnection, 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<tokio::sync::watch::Receiver<u64>>, auth_changes: Option<tokio::sync::watch::Receiver<u64>>,
} }
@@ -192,9 +197,19 @@ impl LunaSampler {
.await .await
.map_err(LunaSamplerError::Provider)? .map_err(LunaSamplerError::Provider)?
.auth; .auth;
let thread_id = ThreadId::new().to_string();
let mut headers = build_session_headers( let mut headers = build_session_headers(
Some(self.config.session_id.clone()), 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( headers.insert(
"openai-beta", "openai-beta",
@@ -207,7 +222,7 @@ impl LunaSampler {
if let Some(originator) = self.config.originator.as_deref() { if let Some(originator) = self.config.originator.as_deref() {
add_originator_header(&mut headers, originator); 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); headers.insert("x-client-request-id", request_id);
} }
@@ -258,7 +273,8 @@ impl LunaSampler {
Ok(PooledConnection { Ok(PooledConnection {
connection, connection,
connected_at: Instant::now(), thread_id,
expires_at: Instant::now() + MAX_WEBSOCKET_AGE,
auth_changes, auth_changes,
}) })
} }
@@ -280,7 +296,7 @@ impl LunaSampler {
.auth_changes .auth_changes
.as_ref() .as_ref()
.is_none_or(|auth| !auth.has_changed().unwrap_or(true)) .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 => && !connection.connection.is_closed().await =>
{ {
break connection; break connection;
@@ -355,12 +371,7 @@ impl LunaSampler {
/// Sends one structured, tool-less request on an exclusively leased WebSocket. /// Sends one structured, tool-less request on an exclusively leased WebSocket.
pub async fn sample(&self, request: LunaSamplingRequest) -> Result<String, LunaSamplerError> { pub async fn sample(&self, request: LunaSamplingRequest) -> Result<String, LunaSamplerError> {
let metadata = HashMap::from([ let turn_id = request.turn_id;
("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 mut input = vec![ let mut input = vec![
ResponseItem::AdditionalTools { ResponseItem::AdditionalTools {
id: None, id: None,
@@ -405,7 +416,7 @@ impl LunaSampler {
phase: None, phase: None,
internal_chat_message_metadata_passthrough: None, internal_chat_message_metadata_passthrough: None,
}); });
let request = ResponsesApiRequest { let mut request = ResponsesApiRequest {
model: MODEL.to_owned(), model: MODEL.to_owned(),
instructions: String::new(), instructions: String::new(),
input, input,
@@ -428,7 +439,7 @@ impl LunaSampler {
&Some(request.output_schema), &Some(request.output_schema),
/*output_schema_strict*/ true, /*output_schema_strict*/ true,
), ),
client_metadata: Some(metadata), client_metadata: None,
}; };
let (supersede, mut superseded) = oneshot::channel(); let (supersede, mut superseded) = oneshot::channel();
let scored = Arc::new(AtomicBool::new(false)); let scored = Arc::new(AtomicBool::new(false));
@@ -475,6 +486,26 @@ impl LunaSampler {
return Err(error); 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 let mut stream = match lease
.connection .connection
.connection .connection
+128 -12
View File
@@ -10,6 +10,7 @@ use codex_login::ExternalAuthRefreshContext;
use codex_model_provider::create_model_provider; use codex_model_provider::create_model_provider;
use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::ModelProviderInfo;
use codex_protocol::ResponseItemId; use codex_protocol::ResponseItemId;
use codex_protocol::ThreadId;
use codex_protocol::models::ResponseItem; use codex_protocol::models::ResponseItem;
use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::openai_models::ReasoningEffort;
use codex_protocol::protocol::SessionSource; 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 core_test_support::skip_if_no_network;
use pretty_assertions::assert_eq; use pretty_assertions::assert_eq;
use serde_json::json; use serde_json::json;
use std::collections::HashSet;
use std::sync::Arc; use std::sync::Arc;
use std::sync::atomic::AtomicUsize; use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering; use std::sync::atomic::Ordering;
@@ -35,6 +37,54 @@ use super::LunaSamplingRequest;
use super::MAX_SAMPLING_RETRIES; use super::MAX_SAMPLING_RETRIES;
use super::MAX_WEBSOCKET_CONNECTIONS; use super::MAX_WEBSOCKET_CONNECTIONS;
fn assert_connection_metadata(server: &responses::WebSocketTestServer) -> Result<String> {
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<String> { async fn proxy_websocket_servers(servers: &[&responses::WebSocketTestServer]) -> Result<String> {
let listener = TcpListener::bind("127.0.0.1:0").await?; let listener = TcpListener::bind("127.0.0.1:0").await?;
let address = listener.local_addr()?; let address = listener.local_addr()?;
@@ -174,7 +224,6 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_structured_requ
Some("true".to_owned()) Some("true".to_owned())
); );
assert_eq!(handshake.header("session-id"), Some("session-1".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!( assert_eq!(
handshake.header("originator"), handshake.header("originator"),
Some("guardian-v2-test".to_owned()) 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()); requests.extend(refreshed.single_connection());
assert_eq!(requests.len(), 2); 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!( assert_eq!(
requests[0].body_json()["input"][2]["content"], requests[0].body_json()["input"][2]["content"],
json!([ json!([
@@ -252,20 +306,14 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_structured_requ
assert_eq!(request["tool_choice"], "none"); assert_eq!(request["tool_choice"], "none");
assert_eq!(request["text"]["format"]["strict"], true); assert_eq!(request["text"]["format"]["strict"], true);
assert_eq!(request["prompt_cache_key"], "guardian-v2:thread-1"); 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!( assert_eq!(
request["client_metadata"]["ws_request_header_x_openai_internal_codex_responses_lite"], request["client_metadata"]["turn_id"],
"true" format!("turn-{}", index + 1)
); );
assert!(request.get("tools").is_none()); assert!(request.get("tools").is_none());
let effort = if index == 0 { "none" } else { "medium" }; let effort = if index == 0 { "none" } else { "medium" };
assert_eq!(request["reasoning"]["effort"], effort); assert_eq!(request["reasoning"]["effort"], effort);
assert_eq!(request["reasoning"]["context"], "all_turns"); assert_eq!(request["reasoning"]["context"], "all_turns");
assert_eq!(
request["client_metadata"]["turn_id"],
format!("turn-{}", index + 1)
);
} }
Ok(()) Ok(())
@@ -493,9 +541,26 @@ async fn sampler_grows_its_pool_for_overlapping_requests() -> Result<()> {
r#"{"score":0.25}"#.to_owned(), r#"{"score":0.25}"#.to_owned(),
) )
); );
let mut thread_ids = HashSet::new();
let mut turn_ids = HashSet::new();
for server in [&first, &second, &third] { for server in [&first, &second, &third] {
assert_eq!(server.single_connection().len(), 1); 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(()) Ok(())
} }
@@ -629,10 +694,61 @@ async fn sampler_retries_expired_websockets_on_another_warm_connection() -> Resu
let healthy_requests = healthy.single_connection(); let healthy_requests = healthy.single_connection();
assert_eq!(expired_requests.len(), 1); assert_eq!(expired_requests.len(), 1);
assert_eq!(healthy_requests.len(), 1); assert_eq!(healthy_requests.len(), 1);
assert_eq!( assert_ne!(
expired_requests[0].body_json(), assert_connection_metadata(&expired)?,
healthy_requests[0].body_json() 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::<Result<HashSet<_>>>()?;
assert_eq!(thread_ids.len(), 3);
assert_eq!(second.single_connection().len(), 1);
assert_eq!(replacement.single_connection().len(), 1);
Ok(()) Ok(())
} }