openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-19 17:28:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Apply workspace routing to Guardian v2 classifier requests (#47162)

Browse Source 
## Why

Classifier requests need to preserve their owning thread's workspace routing and account binding, including when reusing pooled connections or retrying after authentication changes.

## What changed

- Share the thread's `WorkspaceRoutingContext` with Guardian v2 and resolve routing before leasing HTTP or pooled WebSocket connections.
- Reuse WebSockets only when their connection key and authentication state still match. Cache HTTP transports by URL and redirect policy, rejecting redirects when routing requires it.
- Stop classification requests and retries when the account owner changes while allowing token refreshes for the same user and workspace.

## Testing

Add regression coverage for routing failures with prewarmed sockets, shared thread routing state, account switches during HTTP and WebSocket requests, same-account token refreshes, and HTTP redirect-policy cache changes.

GitOrigin-RevId: 7b3efcadb4797b6a2de2e28e7506b297b0728d69
This commit is contained in:
acrognale-oai
2026-09-22 02:33:29 +00:00
committed by copyberry
parent b23324986d
commit 30daed37ad
11 changed files with 575 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs
+15
View File
@@ -20,6 +20,7 @@ use codex_extension_api::ExtensionMetrics;
use codex_http_client::HttpClientFactory;
use codex_login::AgentIdentityAuthPolicy;
use codex_model_provider::SharedModelProvider;
use codex_model_provider::WorkspaceRoutingContext;
use codex_protocol::ResponseItemId;
use codex_protocol::error::CodexErr;
use codex_protocol::models::ContentItem;
@@ -41,6 +42,8 @@ const MAX_CONCURRENT_REQUESTS: usize = 16;
pub struct LunaSamplerConfig {
/// Provider and credentials selected for the owning thread.
pub provider: SharedModelProvider,
/// Routing scope and retained configuration layers for the owning thread.
pub workspace_routing: WorkspaceRoutingContext,
/// Effective proxy, custom-CA, and cookie configuration.
pub http_client_factory: HttpClientFactory,
/// Agent-identity policy selected for the owning thread.
@@ -151,6 +154,17 @@ impl LunaSampler {
/// Sends one tool-less classification request using an available transport.
pub async fn sample(&self, request: LunaSamplingRequest) -> Result<String, LunaSamplerError> {
let auth_owner_generation = self
.config
.provider
.auth_manager()
.filter(|_| self.config.provider.info().auth.is_none())
.map(|manager| {
manager
.auth_change_state_receiver()
.borrow()
.owner_generation
});
if request.parent_compaction.is_some()
&& !self.supports_parent_compaction(request.parent_compaction_hash.as_deref())
{
@@ -258,6 +272,7 @@ impl LunaSampler {
});
}
execution::SamplingExecution {
auth_owner_generation,
config: Arc::clone(&self.config),
connections: Arc::clone(&self.connections),
request,
codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool.rs
+93 -18
View File
@@ -1,6 +1,8 @@
//! Keeps WebSocket establishment off the classification path.
//! One background worker fills the pool. Opening timeouts pause replenishment;
//! requests with no healthy idle socket use HTTP under the same concurrency limit.
//! Each lease resolves workspace routing before reusing a socket or starting HTTP.
//! Constrained HTTP clients reject redirects, including after a cache policy change.
use super::super::metrics::sampler_failure_reason;
use super::INITIAL_WEBSOCKET_CONNECTIONS;
@@ -26,8 +28,11 @@ use codex_login::default_client::ClientRedirectPolicy;
use codex_login::default_client::add_originator_header;
use codex_login::default_client::create_client_for_route_async;
use codex_login::default_client::default_headers;
use codex_model_provider::ACCOUNT_ROUTING_HEADER;
use codex_model_provider::AgentIdentitySessionFallback;
use codex_model_provider::ProviderAuthScope;
use codex_model_provider::ResolvedResponsesProvider;
use codex_model_provider::ResponsesConnectionKey;
use codex_protocol::ThreadId;
use http::HeaderMap;
use http::HeaderValue;
@@ -58,7 +63,13 @@ pub(super) struct ConnectionPool {
sockets: Arc<Semaphore>,
replenishing: Arc<tokio::sync::Mutex<()>>,
retry_after: Mutex<Option<Instant>>,
http_transport: Mutex<Option<(String, Arc<OnceCell<ReqwestTransport>>)>>,
http_transport: Mutex<Option<Arc<HttpTransport>>>,
}
struct HttpTransport {
url: String,
redirect_policy: ClientRedirectPolicy,
transport: OnceCell<ReqwestTransport>,
}
pub(super) struct PooledConnection {
@@ -68,9 +79,18 @@ pub(super) struct PooledConnection {
thread_id: String,
pub(super) expires_at: Instant,
auth_changes: Option<tokio::sync::watch::Receiver<u64>>,
key: ResponsesConnectionKey,
_permit: OwnedSemaphorePermit,
}
struct ClientSetup {
provider: Provider,
auth: SharedAuthProvider,
redirect_policy: ClientRedirectPolicy,
key: ResponsesConnectionKey,
auth_changes: Option<tokio::sync::watch::Receiver<u64>>,
}
enum Connection {
Websocket(PooledConnection),
Http(ResponsesClient<ReqwestTransport>),
@@ -162,6 +182,13 @@ impl ConnectionPool {
.acquire_owned()
.await
.map_err(|error| LunaSamplerError::Api(ApiError::Stream(error.to_string())))?;
let ClientSetup {
mut provider,
auth,
redirect_policy,
key,
auth_changes,
} = self.client_setup().await?;
let connection = loop {
let idle = self
.idle_connections
@@ -170,10 +197,11 @@ impl ConnectionPool {
.pop();
match idle {
Some(connection)
if connection
.auth_changes
.as_ref()
.is_none_or(|auth| !auth.has_changed().unwrap_or(true))
if connection.key == key
&& connection
.auth_changes
.as_ref()
.is_none_or(|auth| !auth.has_changed().unwrap_or(true))
&& Instant::now() < connection.expires_at
&& !connection.connection.is_closed().await =>
{
@@ -191,33 +219,43 @@ impl ConnectionPool {
}
None => {
self.replenish();
let (mut provider, auth) = self.client_setup().await?;
// Sampling owns the retry budget across both transports.
provider.retry.max_attempts = 0;
let request_kind = self.responses_request_kind().await;
let url = provider.url_for_path("/responses");
let redirect_policy = if provider.headers.contains_key(ACCOUNT_ROUTING_HEADER) {
ClientRedirectPolicy::Reject
} else {
redirect_policy
};
let transport = {
let mut cached = self
.http_transport
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if let Some((cached_url, transport)) = cached.as_ref()
&& cached_url == &url
if let Some(transport) = cached.as_ref()
&& transport.url == url
&& transport.redirect_policy == redirect_policy
{
Arc::clone(transport)
} else {
let transport = Arc::new(OnceCell::new());
*cached = Some((url.clone(), Arc::clone(&transport)));
let transport = Arc::new(HttpTransport {
url: url.clone(),
redirect_policy,
transport: OnceCell::new(),
});
*cached = Some(Arc::clone(&transport));
transport
}
};
let transport = transport
.transport
.get_or_try_init(|| async {
let client = create_client_for_route_async(
self.config.http_client_factory.clone(),
url,
ClientRouteClass::Api,
ClientRedirectPolicy::Default,
redirect_policy,
)
.await
.map_err(|error| {
@@ -237,6 +275,14 @@ impl ConnectionPool {
)
}
};
if auth_changes
.as_ref()
.is_some_and(|auth| auth.has_changed().unwrap_or(true))
{
return Err(LunaSamplerError::Api(ApiError::Stream(
"authentication changed while leasing".into(),
)));
}
Ok(ConnectionLease {
thread_id,
request_kind,
@@ -246,11 +292,20 @@ impl ConnectionPool {
})
}
async fn client_setup(&self) -> Result<(Provider, SharedAuthProvider), LunaSamplerError> {
let provider = self
async fn client_setup(&self) -> Result<ClientSetup, LunaSamplerError> {
let mut auth_changes = self
.config
.provider
.api_provider()
.auth_manager()
.map(|manager| manager.auth_change_receiver());
let revision = auth_changes.as_mut().map(|auth| *auth.borrow_and_update());
let ResolvedResponsesProvider {
provider,
redirect_policy,
} = self
.config
.provider
.responses_api_provider(&self.config.workspace_routing)
.await
.map_err(LunaSamplerError::Provider)?;
let auth = self
@@ -264,7 +319,22 @@ impl ConnectionPool {
.await
.map_err(LunaSamplerError::Provider)?
.auth;
Ok((provider, auth))
if auth_changes
.as_ref()
.is_some_and(|auth| auth.has_changed().unwrap_or(true))
{
return Err(LunaSamplerError::Api(ApiError::Stream(
"authentication changed while resolving routing".into(),
)));
}
let key = ResponsesConnectionKey::new(&provider, revision);
Ok(ClientSetup {
provider,
auth,
redirect_policy,
key,
auth_changes,
})
}
fn headers(
@@ -326,9 +396,13 @@ impl ConnectionPool {
&self,
permit: OwnedSemaphorePermit,
) -> Result<PooledConnection, LunaSamplerError> {
let auth_manager = self.config.provider.auth_manager();
let auth_changes = auth_manager.map(|manager| manager.auth_change_receiver());
let (provider, auth) = self.client_setup().await?;
let ClientSetup {
provider,
auth,
redirect_policy: _,
key,
auth_changes,
} = self.client_setup().await?;
let thread_id = ThreadId::new().to_string();
let request_kind = self.responses_request_kind().await;
let mut headers = self.headers(&thread_id, request_kind)?;
@@ -379,6 +453,7 @@ impl ConnectionPool {
thread_id,
expires_at: Instant::now() + MAX_WEBSOCKET_AGE,
auth_changes,
key,
_permit: permit,
})
}
codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_tests.rs
+91
View File
@@ -6,6 +6,10 @@ use super::super::tests::sampler_config;
use super::*;
use anyhow::Result;
use codex_login::AuthManager;
use codex_login::WorkspaceRouting;
use codex_login::WorkspaceRoutingRequest;
use codex_login::WorkspaceRoutingResolver;
use codex_model_provider::ModelProviderFuture;
use codex_model_provider::create_model_provider;
use codex_model_provider_info::ModelProviderInfo;
use core_test_support::responses;
@@ -369,3 +373,90 @@ async fn supersession_closes_http_before_headers_and_while_draining_the_body() -
}
Ok(())
}
struct RoutingPolicy(Mutex<Option<WorkspaceRouting>>);
impl WorkspaceRoutingResolver for RoutingPolicy {
fn resolve(
&self,
_request: WorkspaceRoutingRequest,
) -> ModelProviderFuture<'_, std::io::Result<Option<WorkspaceRouting>>> {
Box::pin(async { Ok(self.0.lock().unwrap().clone()) })
}
}
#[tokio::test]
async fn http_cache_preserves_no_constraint_redirect_policy_and_configured_headers() -> Result<()> {
let base_url = "https://localhost:9/backend-api/codex".to_owned();
let auth =
AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing());
let resolver = Arc::new(RoutingPolicy(Mutex::new(None)));
let routing_owner: Arc<dyn WorkspaceRoutingResolver> = resolver.clone();
auth.set_workspace_routing_resolver(Arc::downgrade(&routing_owner));
let mut config = sampler_config(base_url.clone());
config.provider = create_model_provider(
ModelProviderInfo::create_openai_provider(Some(base_url)),
Some(auth),
);
let pool = ConnectionPool::new(Arc::new(config));
// Exercise HTTP leases without starting a background WebSocket opener.
*pool.retry_after.lock().unwrap() = Some(Instant::now() + CONNECT_COOLDOWN);
let mut previous: Option<Arc<HttpTransport>> = None;
for (routing, expected_policy) in [
(None, ClientRedirectPolicy::Default),
(Some("NO_CONSTRAINT"), ClientRedirectPolicy::Reject),
(Some("NO_CONSTRAINT"), ClientRedirectPolicy::Reject),
(None, ClientRedirectPolicy::Default),
] {
*resolver.0.lock().unwrap() = routing.map(|account_routing_override| WorkspaceRouting {
chatgpt_account_id: "account_id".into(),
backend_origin: "https://localhost:9".into(),
account_routing_override: account_routing_override.into(),
});
let setup = pool.client_setup().await?;
assert!(!setup.provider.headers.contains_key(ACCOUNT_ROUTING_HEADER));
let lease = pool.lease().await?;
assert!(matches!(lease.connection, Connection::Http(_)));
let cached = Arc::clone(pool.http_transport.lock().unwrap().as_ref().unwrap());
assert_eq!(
(cached.url.as_str(), cached.redirect_policy),
(
"https://localhost:9/backend-api/codex/responses",
expected_policy
),
);
if let Some(previous) = previous {
assert_eq!(
Arc::ptr_eq(&previous, &cached),
previous.redirect_policy == expected_policy,
);
}
previous = Some(cached);
}
let mut config = sampler_config("https://localhost:9/v1".into());
let mut info = config.provider.info().clone();
info.http_headers = Some(std::collections::HashMap::from([(
ACCOUNT_ROUTING_HEADER.into(),
"us".into(),
)]));
config.provider = create_model_provider(info, config.provider.auth_manager());
let pool = ConnectionPool::new(Arc::new(config));
*pool.retry_after.lock().unwrap() = Some(Instant::now() + CONNECT_COOLDOWN);
assert_eq!(
pool.client_setup().await?.redirect_policy,
ClientRedirectPolicy::Default
);
let lease = pool.lease().await?;
assert!(matches!(lease.connection, Connection::Http(_)));
assert_eq!(
pool.http_transport
.lock()
.unwrap()
.as_ref()
.unwrap()
.redirect_policy,
ClientRedirectPolicy::Reject,
);
Ok(())
}
codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs
+42
View File
@@ -1,5 +1,6 @@
//! Executes one prepared classifier request with the existing retry and streaming rules.
//! The first output returns immediately; a detached drain preserves reuse and token accounting.
//! Requests and retries stop when the account owner that started the classification changes.
use super::CLASSIFICATION_TOKEN_USAGE_METRIC;
use super::ConnectionPool;
@@ -30,6 +31,7 @@ const RESPONSES_LITE_METADATA_KEY: &str =
const TURN_METADATA_KEY: &str = "x-codex-turn-metadata";
pub(super) struct SamplingExecution {
pub(super) auth_owner_generation: Option<u64>,
pub(super) config: Arc<LunaSamplerConfig>,
pub(super) connections: Arc<ConnectionPool>,
pub(super) request: ResponsesApiRequest,
@@ -107,6 +109,40 @@ impl SamplingExecution {
mut superseded: oneshot::Receiver<()>,
scored: Arc<AtomicBool>,
) -> Result<String, LunaSamplerError> {
let auth_changes = self
.config
.provider
.auth_manager()
.filter(|_| self.config.provider.info().auth.is_none())
.map(|manager| manager.auth_change_state_receiver());
let owner_generation = self.auth_owner_generation;
let account_changed_error = || {
LunaSamplerError::Provider(
std::io::Error::other("account changed during Luna classification").into(),
)
};
let ensure_account_owner = || {
if auth_changes
.as_ref()
.map(|changes| changes.borrow().owner_generation)
!= owner_generation
{
return Err(account_changed_error());
}
Ok(())
};
let mut owner_changes = auth_changes.clone();
let owner_changed = async move {
let Some(changes) = owner_changes.as_mut() else {
return std::future::pending::<()>().await;
};
while Some(changes.borrow_and_update().owner_generation) == owner_generation {
if changes.changed().await.is_err() {
break;
}
}
};
tokio::pin!(owner_changed);
let mut retries = 0;
let mut auth_recovery = self
.config
@@ -114,8 +150,10 @@ impl SamplingExecution {
.auth_manager()
.map(|manager| manager.unauthorized_recovery());
'retry: loop {
ensure_account_owner()?;
let lease = match tokio::select! {
biased;
_ = &mut owner_changed => return Err(account_changed_error()),
_ = &mut superseded => return Err(LunaSamplerError::Superseded),
lease = self.connections.lease() => lease,
} {
@@ -130,6 +168,7 @@ impl SamplingExecution {
return Err(error);
}
};
ensure_account_owner()?;
self.request.service_tier = if lease.request_kind == RequestMode::GuardianClassifier {
None
} else {
@@ -168,6 +207,7 @@ impl SamplingExecution {
self.request.client_metadata = Some(client_metadata);
let mut stream = match tokio::select! {
biased;
_ = &mut owner_changed => return Err(account_changed_error()),
_ = &mut superseded => return Err(LunaSamplerError::Superseded),
stream = lease.stream_request(&self.request) => stream,
} {
@@ -187,6 +227,7 @@ impl SamplingExecution {
let mut output = String::new();
while let Some(event) = tokio::select! {
biased;
_ = &mut owner_changed => return Err(account_changed_error()),
_ = &mut superseded => {
return if scored.load(Ordering::Relaxed) && !output.is_empty() {
Ok(output)
@@ -196,6 +237,7 @@ impl SamplingExecution {
}
event = stream.rx_event.recv() => event,
} {
ensure_account_owner()?;
let event = match event {
Ok(event) => event,
Err(error) => {
codex-rs/ext/guardian-v2/src/async_scorer/sampler_routing_tests.rs
+307
View File
@@ -0,0 +1,307 @@
//! Keeps classifier evidence within its owning account and thread routing binding.
use super::connect_sampler;
use super::proxy_websocket_servers;
use super::sample_request;
use super::sampler_config;
use anyhow::Result;
use codex_login::AuthManager;
use codex_login::CodexAuth;
use codex_login::ExternalAuth;
use codex_login::ExternalAuthFuture;
use codex_login::ExternalAuthRefreshContext;
use codex_login::WorkspaceRouting;
use codex_login::WorkspaceRoutingRequest;
use codex_login::WorkspaceRoutingResolver;
use codex_model_provider::ModelProviderFuture;
use codex_model_provider::create_model_provider;
use codex_model_provider_info::ModelProviderInfo;
use core_test_support::responses;
use core_test_support::responses::WebSocketConnectionConfig;
use core_test_support::responses::ev_assistant_message;
use core_test_support::responses::ev_completed;
use pretty_assertions::assert_eq;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::time::Duration;
type RoutingFuture<'a> = ModelProviderFuture<'a, std::io::Result<Option<WorkspaceRouting>>>;
struct UnavailableRouting;
impl WorkspaceRoutingResolver for UnavailableRouting {
fn resolve(&self, _request: WorkspaceRoutingRequest) -> RoutingFuture<'_> {
Box::pin(async { Err(std::io::Error::other("discovery unavailable")) })
}
}
#[tokio::test]
async fn failed_routing_blocks_a_prewarmed_classifier_connection() -> Result<()> {
core_test_support::skip_if_no_network!(Ok(()));
let response = vec![vec![ev_assistant_message("r1", "low"), ev_completed("r1")]];
let server = responses::start_websocket_server(vec![response.clone(), response]).await;
let base_url = proxy_websocket_servers(&[&server, &server])
.await?
.replace("/v1", "/backend-api/codex");
let auth =
AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing());
let mut config = sampler_config(base_url.clone());
config.provider = create_model_provider(
ModelProviderInfo::create_openai_provider(Some(base_url)),
Some(auth.clone()),
);
let sampler = connect_sampler(config).await?;
let resolver: Arc<dyn WorkspaceRoutingResolver> = Arc::new(UnavailableRouting);
auth.set_workspace_routing_resolver(Arc::downgrade(&resolver));
let error = sampler.sample(sample_request("turn-1")).await.unwrap_err();
assert!(
error.to_string().contains("discovery unavailable"),
"{error}"
);
assert!(server.connections().iter().all(Vec::is_empty));
Ok(())
}
struct AuthAfterRouting {
initial: CodexAuth,
updated: CodexAuth,
resolved: AtomicBool,
}
impl ExternalAuth for AuthAfterRouting {
fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> {
Box::pin(async {
Ok(if self.resolved.load(Ordering::SeqCst) {
self.updated.clone()
} else {
self.initial.clone()
})
})
}
fn refresh(&self, _context: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> {
ExternalAuth::resolve(self)
}
}
impl WorkspaceRoutingResolver for AuthAfterRouting {
fn resolve(&self, _request: WorkspaceRoutingRequest) -> RoutingFuture<'_> {
Box::pin(async {
self.resolved.store(/*val*/ true, Ordering::SeqCst);
Ok(None)
})
}
}
#[tokio::test]
async fn classification_retries_refreshes_but_never_switches_accounts() -> Result<()> {
core_test_support::skip_if_no_network!(Ok(()));
// Both tokens identify the same user; their signatures differ after refresh.
let token =
"header.eyJodHRwczovL2FwaS5vcGVuYWkuY29tL2F1dGgiOnsiY2hhdGdwdF91c2VyX2lkIjoidXNlci1hIn19";
for account in ["workspace-a", "workspace-b"] {
let server = responses::start_mock_server().await;
let response = responses::mount_sse_once(
&server,
responses::sse(vec![
ev_assistant_message("score", "low"),
ev_completed("score"),
]),
)
.await;
let home = tempfile::tempdir()?;
let changes = Arc::new(AuthAfterRouting {
initial: CodexAuth::from_external_chatgpt_tokens(
&format!("{token}.initial"),
"workspace-a",
/*chatgpt_plan_type*/ None,
)?,
updated: CodexAuth::from_external_chatgpt_tokens(
&format!("{token}.refreshed"),
account,
/*chatgpt_plan_type*/ None,
)?,
resolved: AtomicBool::new(/*v*/ false),
});
let auth = AuthManager::from_auth_for_testing_with_home(
changes.initial.clone(),
home.path().to_owned(),
);
auth.set_external_auth(changes.clone()).await?;
let resolver: Arc<dyn WorkspaceRoutingResolver> = changes;
auth.set_workspace_routing_resolver(Arc::downgrade(&resolver));
let mut config = sampler_config(format!("{}/backend-api/codex", server.uri()));
config.provider = create_model_provider(config.provider.info().clone(), Some(auth));
let sampler = super::LunaSampler::new(config);
let result = sampler.sample(sample_request("turn-1")).await;
if account == "workspace-a" {
assert_eq!(result?, "low");
let request = response.single_request();
assert_eq!(
(
request.header("chatgpt-account-id"),
request.header("authorization")
),
(
Some("workspace-a".into()),
Some(format!("Bearer {token}.refreshed"))
),
);
} else {
assert!(result.unwrap_err().to_string().contains("account changed"));
assert!(
response.requests().is_empty(),
"evidence must not reach the new account"
);
}
}
Ok(())
}
struct ChangedBootstrap(AtomicBool);
impl WorkspaceRoutingResolver for ChangedBootstrap {
fn resolve(&self, request: WorkspaceRoutingRequest) -> RoutingFuture<'_> {
Box::pin(async move {
if request.previously_routed {
Err(std::io::Error::other(
"bootstrap changed; start a new thread",
))
} else if self.0.swap(/*val*/ false, Ordering::SeqCst) {
Ok(Some(WorkspaceRouting {
chatgpt_account_id: "account_id".into(),
backend_origin: "https://gov.chatgpt.com".into(),
account_routing_override: "us_cr".into(),
}))
} else {
Ok(None)
}
})
}
}
#[tokio::test]
async fn idle_classifier_retains_the_owning_threads_workspace_binding() -> Result<()> {
core_test_support::skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let test = core_test_support::test_codex::test_codex()
.build_with_auto_env(&server)
.await?;
let auth =
AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing());
let resolver: Arc<dyn WorkspaceRoutingResolver> =
Arc::new(ChangedBootstrap(AtomicBool::new(/*v*/ true)));
auth.set_workspace_routing_resolver(Arc::downgrade(&resolver));
let mut config = test.config.clone();
config.model_provider = ModelProviderInfo::create_openai_provider(Some(format!(
"{}/backend-api/codex",
server.uri()
)));
let session_store = codex_extension_api::ExtensionData::new("session-1");
let input = codex_extension_api::ThreadStartInput {
config: &config,
session_source: &codex_protocol::protocol::SessionSource::Exec,
persistent_thread_state_available: false,
environments: &[],
mcp_resource_client: None,
extension_metrics: None,
session_store: &session_store,
thread_store: test.codex.thread_extension_data(),
};
let owning_context = input
.thread_store
.get::<codex_model_provider::WorkspaceRoutingContext>()
.expect("Core seeds the owning thread's routing context before extensions start");
let classifier =
crate::async_scorer::startup::sampler_config(&input, auth.clone(), /*manager*/ None).await;
classifier
.provider
.responses_api_provider(&owning_context)
.await?;
let sampler = super::LunaSampler::new(classifier);
let error = sampler.sample(sample_request("turn-1")).await.unwrap_err();
assert!(error.to_string().contains("start a new thread"), "{error}");
assert!(server.received_requests().await.unwrap().is_empty());
Ok(())
}
#[tokio::test]
async fn account_switch_cancels_request_opening_and_first_token_waits() -> Result<()> {
core_test_support::skip_if_no_network!(Ok(()));
let token = "header.eyJodHRwczovL2FwaS5vcGVuYWkuY29tL2F1dGgiOnsiY2hhdGdwdF91c2VyX2lkIjoidXNlci1hIn19.signature";
for transport in ["http", "websocket"] {
let http = responses::start_mock_server().await;
let opening = responses::mount_response_once(
&http,
responses::sse_response(String::new()).set_delay(Duration::from_secs(/*secs*/ 60)),
)
.await;
let stalled = WebSocketConnectionConfig {
requests: vec![Vec::new()],
response_headers: Vec::new(),
accept_delay: None,
close_after_requests: false,
};
let websocket = responses::start_websocket_server_with_headers(vec![stalled; 2]).await;
let changes = Arc::new(AuthAfterRouting {
initial: CodexAuth::from_external_chatgpt_tokens(
token,
"workspace-a",
/*chatgpt_plan_type*/ None,
)?,
updated: CodexAuth::from_external_chatgpt_tokens(
token,
"workspace-b",
/*chatgpt_plan_type*/ None,
)?,
resolved: AtomicBool::new(/*v*/ false),
});
let home = tempfile::tempdir()?;
let auth = AuthManager::from_auth_for_testing_with_home(
changes.initial.clone(),
home.path().to_owned(),
);
auth.set_external_auth(changes.clone()).await?;
let base_url = if transport == "http" {
format!("{}/v1", http.uri())
} else {
proxy_websocket_servers(&[&websocket, &websocket]).await?
};
let mut config = sampler_config(base_url);
config.provider = create_model_provider(config.provider.info().clone(), Some(auth.clone()));
let sampler = if transport == "http" {
super::LunaSampler::new(config)
} else {
connect_sampler(config).await?
};
let sample = sampler.sample(sample_request("turn-1"));
tokio::pin!(sample);
tokio::time::timeout(Duration::from_secs(/*secs*/ 10), async {
loop {
if !opening.requests().is_empty() || websocket.connections().iter().any(|requests| !requests.is_empty()) {
break Ok::<(), anyhow::Error>(());
}
tokio::select! {
result = &mut sample => anyhow::bail!("classification finished before account switch: {result:?}"),
() = tokio::task::yield_now() => {}
}
}
}).await??;
changes.resolved.store(/*val*/ true, Ordering::SeqCst);
auth.set_external_auth(changes).await?;
let error = tokio::time::timeout(Duration::from_secs(/*secs*/ 2), &mut sample)
.await?
.unwrap_err();
assert!(
error.to_string().contains("account changed"),
"{transport}: {error}"
);
assert_eq!(opening.requests().len(), usize::from(transport == "http"));
assert_eq!(
websocket.connections().iter().map(Vec::len).sum::<usize>(),
usize::from(transport == "websocket")
);
}
Ok(())
}
codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs
+12
View File
@@ -48,6 +48,9 @@ use super::LunaSamplerError;
use super::LunaSamplingRequest;
use super::MAX_CONCURRENT_REQUESTS;
#[path = "sampler_routing_tests.rs"]
mod routing;
impl LunaSampler {
/// Waits for warm sockets to enter the client pool, beyond the server handshake.
pub(in crate::async_scorer) async fn wait_for_prewarm(&self, timeout: Duration) -> Result<()> {
@@ -196,6 +199,9 @@ pub(in crate::async_scorer) async fn proxy_websocket_servers_with_http(
pub(super) fn sampler_config(base_url: String) -> LunaSamplerConfig {
LunaSamplerConfig {
workspace_routing: codex_model_provider::WorkspaceRoutingContext::new(
"https://chatgpt.com/backend-api".into(),
),
provider: create_model_provider(
ModelProviderInfo::create_openai_provider(Some(base_url)),
Some(AuthManager::from_auth_for_testing(CodexAuth::from_api_key(
@@ -506,6 +512,9 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_classifications
);
let sampler = connect_sampler(LunaSamplerConfig {
workspace_routing: codex_model_provider::WorkspaceRoutingContext::new(
"https://chatgpt.com/backend-api".into(),
),
provider,
http_client_factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
agent_identity_policy: AgentIdentityAuthPolicy::JwtOnly,
@@ -751,6 +760,9 @@ async fn sampler_returns_classification_token_before_terminal_response_events()
))),
);
let sampler = connect_sampler(LunaSamplerConfig {
workspace_routing: codex_model_provider::WorkspaceRoutingContext::new(
"https://chatgpt.com/backend-api".into(),
),
provider,
http_client_factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
agent_identity_policy: AgentIdentityAuthPolicy::JwtOnly,
codex-rs/ext/guardian-v2/src/async_scorer/startup.rs
+5
View File
@@ -42,6 +42,11 @@ pub(super) async fn sampler_config(
});
let luna_compaction_hash = luna_model.and_then(|model| model.comp_hash);
LunaSamplerConfig {
workspace_routing: input
.thread_store
.get_or_init(|| input.config.workspace_routing_context())
.as_ref()
.clone(),
provider: create_model_provider(input.config.model_provider.clone(), Some(auth_manager)),
http_client_factory: input.config.http_client_factory(),
agent_identity_policy: if input.config.features.enabled(Feature::UseAgentIdentity) {