mirror of
https://github.com/openai/codex.git
synced 2026-11-13 14:46:14 +00:00
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:
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");
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Reference in new issue
Block a user