openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-27 18:48:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

Prewarm more Guardian sampling WebSockets (#40403)

Browse Source 
## What changed

- Increase the Guardian sampler's initial WebSocket pool from 2 connections to 8.
- Keep pool growth and recovery tests focused on the two-connection scenarios they exercise while making the remaining tests derive their expectations from the initial pool size.

GitOrigin-RevId: e2eabcdde5828ac7df333271f3253e74d696f396
This commit is contained in:
jif
2026-08-24 13:23:10 +00:00
committed by copyberry
parent 80cce09d05
commit 77b30a21e1
3 changed files with 113 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs
+32 -18
View File
@@ -65,6 +65,7 @@ use crate::async_scorer::config::CLASSIFICATION_OUTPUT_INSTRUCTIONS;
use crate::async_scorer::config::DEFAULT_MODEL_CONTEXT_ITEM_TOKENS;
use crate::async_scorer::config::DEFAULT_PARENT_COMPACTION_TOKENS;
use crate::async_scorer::sampler::CLASSIFICATION_TOKEN_USAGE_METRIC;
use crate::async_scorer::sampler::INITIAL_WEBSOCKET_CONNECTIONS;
use crate::async_scorer::sampler::MODEL;
use crate::async_scorer::transcript::truncate_entry;
@@ -98,12 +99,10 @@ async fn installed_extension_reconnects_after_auth_refresh() -> Result<()> {
];
// Keep the sampled connection open for another request so only auth
// invalidation, not a server close, forces the next handshake.
let server = responses::start_websocket_server(vec![
Vec::new(),
vec![events.clone(), events.clone()],
vec![events],
])
.await;
let mut connections = vec![Vec::new(); INITIAL_WEBSOCKET_CONNECTIONS - 1];
connections.push(vec![events.clone(), events.clone()]);
connections.push(vec![events]);
let server = responses::start_websocket_server(connections).await;
let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("original"));
auth_manager
.set_external_auth(Arc::new(RefreshableAuth(std::sync::Mutex::new("original"))))
@@ -180,25 +179,27 @@ async fn installed_extension_reconnects_after_auth_refresh() -> Result<()> {
);
}
let mut expected_authorizations =
vec![Some("Bearer original".to_owned()); INITIAL_WEBSOCKET_CONNECTIONS];
expected_authorizations.push(Some("Bearer refreshed".to_owned()));
assert_eq!(
server
.handshakes()
.iter()
.map(|handshake| handshake.header("authorization"))
.collect::<Vec<_>>(),
vec![
Some("Bearer original".to_owned()),
Some("Bearer original".to_owned()),
Some("Bearer refreshed".to_owned()),
]
expected_authorizations
);
let mut expected_requests = vec![0; INITIAL_WEBSOCKET_CONNECTIONS - 1];
expected_requests.extend([1, 1]);
assert_eq!(
server
.connections()
.iter()
.map(Vec::len)
.collect::<Vec<_>>(),
vec![0, 1, 1]
expected_requests
);
Ok(())
}
@@ -587,7 +588,9 @@ async fn sample_configured_conversation_history(
"total_tokens": 150,
});
let events = vec![ev_assistant_message("sample", "high"), completed];
let server = responses::start_websocket_server(vec![Vec::new(), vec![events]]).await;
let mut connections = vec![Vec::new(); INITIAL_WEBSOCKET_CONNECTIONS - 1];
connections.push(vec![events]);
let server = responses::start_websocket_server(connections).await;
let provider_info = ModelProviderInfo::create_openai_provider(Some(format!(
"http://{}/v1",
server.uri().trim_start_matches("ws://")
@@ -661,7 +664,10 @@ async fn sample_configured_conversation_history(
let request = tokio::time::timeout(
Duration::from_secs(5),
server.wait_for_request(/*connection_index*/ 1, /*request_index*/ 0),
server.wait_for_request(
/*connection_index*/ INITIAL_WEBSOCKET_CONNECTIONS - 1,
/*request_index*/ 0,
),
)
.await?;
Ok((request.body_json(), test, registry))
@@ -823,7 +829,9 @@ async fn contributor_fails_closed_when_luna_classification_fails() -> Result<()>
ev_assistant_message("sample", "invalid"),
ev_completed("response-invalid"),
];
let server = responses::start_websocket_server(vec![Vec::new(), vec![invalid_score]]).await;
let mut connections = vec![Vec::new(); INITIAL_WEBSOCKET_CONNECTIONS - 1];
connections.push(vec![invalid_score]);
let server = responses::start_websocket_server(connections).await;
let mut config = fixture.test.config.clone();
config.model_provider = ModelProviderInfo::create_openai_provider(Some(format!(
"http://{}/v1",
@@ -1647,7 +1655,8 @@ async fn contributor_skips_models_requiring_managed_guardian_review() -> Result<
.build_with_auto_env(&thread_server)
.await?;
let server = responses::start_websocket_server(vec![Vec::new(), Vec::new()]).await;
let server =
responses::start_websocket_server(vec![Vec::new(); INITIAL_WEBSOCKET_CONNECTIONS]).await;
let provider_info = ModelProviderInfo::create_openai_provider(Some(format!(
"http://{}/v1",
server.uri().trim_start_matches("ws://")
@@ -2028,7 +2037,9 @@ async fn contributor_reuses_the_latest_compatible_parent_compaction() -> Result<
ev_assistant_message("sample", "low"),
ev_completed("response-1"),
];
let server = responses::start_websocket_server(vec![Vec::new(), vec![events]]).await;
let mut connections = vec![Vec::new(); INITIAL_WEBSOCKET_CONNECTIONS - 1];
connections.push(vec![events]);
let server = responses::start_websocket_server(connections).await;
let provider_info = ModelProviderInfo::create_openai_provider(Some(format!(
"http://{}/v1",
server.uri().trim_start_matches("ws://")
@@ -2109,7 +2120,10 @@ async fn contributor_reuses_the_latest_compatible_parent_compaction() -> Result<
let request = tokio::time::timeout(
Duration::from_secs(5),
server.wait_for_request(/*connection_index*/ 1, /*request_index*/ 0),
server.wait_for_request(
/*connection_index*/ INITIAL_WEBSOCKET_CONNECTIONS - 1,
/*request_index*/ 0,
),
)
.await?
.body_json();
codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs
+1 -1
View File
@@ -46,7 +46,7 @@ pub(crate) const MODEL: &str = "gpt-5.6-luna";
pub(crate) const CLASSIFICATION_TOKEN_USAGE_METRIC: &str =
"codex.guardian_v2.classification.token_usage";
const MAX_OUTPUT_BYTES: usize = 8 * 1024;
const INITIAL_WEBSOCKET_CONNECTIONS: usize = 2;
pub(super) const INITIAL_WEBSOCKET_CONNECTIONS: usize = 8;
const MAX_WEBSOCKET_CONNECTIONS: usize = 16;
const MAX_SAMPLING_RETRIES: usize = 2;
const MAX_WEBSOCKET_AGE: Duration = Duration::from_secs(55 * 60);
codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs
+80 -13
View File
@@ -35,6 +35,7 @@ use tokio::net::TcpListener;
use tokio::net::TcpStream;
use super::CLASSIFICATION_TOKEN_USAGE_METRIC;
use super::INITIAL_WEBSOCKET_CONNECTIONS;
use super::LunaSampler;
use super::LunaSamplerConfig;
use super::LunaSamplingRequest;
@@ -88,7 +89,20 @@ fn assert_connection_metadata(server: &responses::WebSocketTestServer) -> Result
Ok(thread_id)
}
#[derive(Clone, Copy)]
enum ProxyPrewarmLimit {
AllConnections,
StopAfter { ready_connections: usize },
}
async fn proxy_websocket_servers(servers: &[&responses::WebSocketTestServer]) -> Result<String> {
proxy_websocket_servers_with_prewarm_limit(servers, ProxyPrewarmLimit::AllConnections).await
}
async fn proxy_websocket_servers_with_prewarm_limit(
servers: &[&responses::WebSocketTestServer],
prewarm_limit: ProxyPrewarmLimit,
) -> Result<String> {
let listener = TcpListener::bind("127.0.0.1:0").await?;
let address = listener.local_addr()?;
let targets = servers
@@ -96,7 +110,15 @@ async fn proxy_websocket_servers(servers: &[&responses::WebSocketTestServer]) ->
.map(|server| server.uri().trim_start_matches("ws://").to_owned())
.collect::<Vec<_>>();
tokio::spawn(async move {
for target in targets {
for (index, target) in targets.into_iter().enumerate() {
if let ProxyPrewarmLimit::StopAfter { ready_connections } = prewarm_limit
&& index == ready_connections
{
let Ok((connection, _)) = listener.accept().await else {
return;
};
drop(connection);
}
let Ok((mut incoming, _)) = listener.accept().await else {
return;
};
@@ -171,7 +193,9 @@ async fn sampler_records_token_usage_after_returning_an_early_classification() -
ev_output_text_delta("low"),
ev_completed_with_tokens("response-1", /*total_tokens*/ 37),
];
let server = responses::start_websocket_server(vec![Vec::new(), vec![events]]).await;
let mut connections = vec![Vec::new(); INITIAL_WEBSOCKET_CONNECTIONS - 1];
connections.push(vec![events]);
let server = responses::start_websocket_server(connections).await;
let metrics = Arc::new(RecordingMetrics::default());
let mut config = sampler_config(format!(
"http://{}/v1",
@@ -240,7 +264,13 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_classifications
let refreshed =
responses::start_websocket_server(vec![vec![scripted_requests[1].clone()]]).await;
let server = responses::start_websocket_server(vec![scripted_requests]).await;
let base_url = proxy_websocket_servers(&[&idle_server, &server, &refreshed]).await?;
let base_url = proxy_websocket_servers_with_prewarm_limit(
&[&idle_server, &server, &refreshed],
ProxyPrewarmLimit::StopAfter {
ready_connections: 2,
},
)
.await?;
let manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("test-api-key"));
manager
.set_external_auth(Arc::new(RefreshableAuth(std::sync::Mutex::new(
@@ -391,8 +421,9 @@ async fn sampler_reuses_parent_compaction_only_for_matching_model_hashes() -> Re
ev_assistant_message("sample", "low"),
ev_completed("response-1"),
];
let server =
responses::start_websocket_server(vec![Vec::new(), vec![events.clone(), events]]).await;
let mut connections = vec![Vec::new(); INITIAL_WEBSOCKET_CONNECTIONS - 1];
connections.push(vec![events.clone(), events]);
let server = responses::start_websocket_server(connections).await;
let mut config = sampler_config(format!(
"http://{}/v1",
server.uri().trim_start_matches("ws://")
@@ -412,7 +443,10 @@ async fn sampler_reuses_parent_compaction_only_for_matching_model_hashes() -> Re
assert_eq!(sampler.sample(request).await?, "low");
let request = server
.wait_for_request(/*connection_index*/ 1, /*request_index*/ 0)
.wait_for_request(
/*connection_index*/ INITIAL_WEBSOCKET_CONNECTIONS - 1,
/*request_index*/ 0,
)
.await
.body_json();
let input = request["input"].as_array().expect("input items");
@@ -430,7 +464,10 @@ async fn sampler_reuses_parent_compaction_only_for_matching_model_hashes() -> Re
switched_request.parent_compaction_hash = Some("incompatible".to_owned());
assert_eq!(sampler.sample(switched_request).await?, "low");
let switched_request = server
.wait_for_request(/*connection_index*/ 1, /*request_index*/ 1)
.wait_for_request(
/*connection_index*/ INITIAL_WEBSOCKET_CONNECTIONS - 1,
/*request_index*/ 1,
)
.await
.body_json();
assert_eq!(switched_request["input"][2]["role"], "user");
@@ -509,7 +546,9 @@ async fn sampler_keeps_first_classification_token_when_later_output_disagrees()
ev_assistant_message("sample", "lowhigh"),
ev_completed("response-1"),
];
let server = responses::start_websocket_server(vec![Vec::new(), vec![events]]).await;
let mut connections = vec![Vec::new(); INITIAL_WEBSOCKET_CONNECTIONS - 1];
connections.push(vec![events]);
let server = responses::start_websocket_server(connections).await;
let sampler = LunaSampler::connect(sampler_config(format!(
"http://{}/v1",
server.uri().trim_start_matches("ws://")
@@ -589,7 +628,13 @@ async fn sampler_grows_its_pool_for_overlapping_requests() -> Result<()> {
let second = responses::start_websocket_server(vec![response("response-2")]).await;
let third = responses::start_websocket_server(vec![response("response-3")]).await;
let sampler = LunaSampler::connect(sampler_config(
proxy_websocket_servers(&[&first, &second, &third]).await?,
proxy_websocket_servers_with_prewarm_limit(
&[&first, &second, &third],
ProxyPrewarmLimit::StopAfter {
ready_connections: 2,
},
)
.await?,
))
.await?;
@@ -661,7 +706,11 @@ async fn sampler_replaces_scored_drains_before_unfinished_classifications() -> R
]]])
.await,
);
let server_refs = servers.iter().collect::<Vec<_>>();
let server_refs = servers[2..INITIAL_WEBSOCKET_CONNECTIONS]
.iter()
.chain(servers[..2].iter())
.chain(servers[INITIAL_WEBSOCKET_CONNECTIONS..].iter())
.collect::<Vec<_>>();
let sampler = Arc::new(
LunaSampler::connect(sampler_config(proxy_websocket_servers(&server_refs).await?)).await?,
);
@@ -785,7 +834,13 @@ async fn sampler_assigns_a_fresh_identity_when_replacing_aged_connections() -> R
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?,
proxy_websocket_servers_with_prewarm_limit(
&[&first, &second, &replacement],
ProxyPrewarmLimit::StopAfter {
ready_connections: 2,
},
)
.await?,
))
.await?;
@@ -830,7 +885,13 @@ async fn sampler_reconnects_after_transient_service_failures() -> Result<()> {
]]])
.await;
let sampler = LunaSampler::connect(sampler_config(
proxy_websocket_servers(&[&first, &second, &recovered]).await?,
proxy_websocket_servers_with_prewarm_limit(
&[&first, &second, &recovered],
ProxyPrewarmLimit::StopAfter {
ready_connections: 2,
},
)
.await?,
))
.await?;
@@ -865,7 +926,13 @@ async fn sampler_limits_transient_recovery_attempts() -> Result<()> {
let third = responses::start_websocket_server(unavailable()).await;
let unused = responses::start_websocket_server(unavailable()).await;
let sampler = LunaSampler::connect(sampler_config(
proxy_websocket_servers(&[&first, &second, &third, &unused]).await?,
proxy_websocket_servers_with_prewarm_limit(
&[&first, &second, &third, &unused],
ProxyPrewarmLimit::StopAfter {
ready_connections: 2,
},
)
.await?,
))
.await?;