From 30daed37ad8035f041f65a4c4615fbc590dc8552 Mon Sep 17 00:00:00 2001 From: acrognale-oai Date: Tue, 22 Sep 2026 02:25:08 +0000 Subject: [PATCH] Apply workspace routing to Guardian v2 classifier requests (#47162) ## 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 --- codex-rs/app-server/README.md | 1 + codex-rs/core/src/session/session.rs | 4 +- .../guardian-v2/src/async_scorer/sampler.rs | 15 + .../async_scorer/sampler/connection_pool.rs | 111 ++++++- .../sampler/connection_pool_tests.rs | 91 ++++++ .../src/async_scorer/sampler/execution.rs | 42 +++ .../src/async_scorer/sampler_routing_tests.rs | 307 ++++++++++++++++++ .../src/async_scorer/sampler_tests.rs | 12 + .../guardian-v2/src/async_scorer/startup.rs | 5 + .../model-provider/src/workspace_routing.rs | 8 +- .../src/workspace_routing_tests.rs | 3 +- 11 files changed, 575 insertions(+), 24 deletions(-) create mode 100644 codex-rs/ext/guardian-v2/src/async_scorer/sampler_routing_tests.rs diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index be16b5d1c5..5da9f5bdcf 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -310,6 +310,7 @@ Routing compares origins by scheme, host, and effective port, ignoring API paths configured base URL. Responses HTTP (including compaction) and WebSockets wait for discovery and preserve API paths. +Guardian v2 classifier HTTP and pooled WebSockets use the same routing. HTTP redirects are rejected. `us` and `us_cr` set `X-OpenAI-Account-Routing-Override`; `NO_CONSTRAINT` omits it. diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 4be0cb1218..b65e77d767 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1613,6 +1613,8 @@ impl Session { let mcp_resource_client = Arc::new(McpResourceClient::new(Arc::clone(&mcp_runtime))); let extension_metrics = extension_metrics::from_session_telemetry(session_telemetry.clone()); + let workspace_routing = thread_extension_data + .get_or_init(|| config.workspace_routing_context()); for contributor in extensions.thread_lifecycle_contributors() { contributor.on_thread_start(codex_extension_api::ThreadStartInput { config: config.as_ref(), @@ -1711,7 +1713,7 @@ impl Session { .enabled(Feature::ConcurrentReasoningSummaries), attestation_provider, config.http_client_factory(), - config.workspace_routing_context(), + workspace_routing.as_ref().clone(), ) .with_restored_history(matches!( &initial_history, diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs index 2861b7bea2..e46ca6b161 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs @@ -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 { + 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, diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool.rs index 8fa4227171..70817b00cd 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool.rs @@ -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, replenishing: Arc>, retry_after: Mutex>, - http_transport: Mutex>)>>, + http_transport: Mutex>>, +} + +struct HttpTransport { + url: String, + redirect_policy: ClientRedirectPolicy, + transport: OnceCell, } pub(super) struct PooledConnection { @@ -68,9 +79,18 @@ pub(super) struct PooledConnection { thread_id: String, pub(super) expires_at: Instant, auth_changes: Option>, + key: ResponsesConnectionKey, _permit: OwnedSemaphorePermit, } +struct ClientSetup { + provider: Provider, + auth: SharedAuthProvider, + redirect_policy: ClientRedirectPolicy, + key: ResponsesConnectionKey, + auth_changes: Option>, +} + enum Connection { Websocket(PooledConnection), Http(ResponsesClient), @@ -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 { + 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 { - 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, }) } diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_tests.rs index efde4f6cdf..6cf716de8d 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_tests.rs @@ -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>); + +impl WorkspaceRoutingResolver for RoutingPolicy { + fn resolve( + &self, + _request: WorkspaceRoutingRequest, + ) -> ModelProviderFuture<'_, std::io::Result>> { + 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 = 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> = 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(()) +} diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs index b8dd61c799..0e50a96e90 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs @@ -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, pub(super) config: Arc, pub(super) connections: Arc, pub(super) request: ResponsesApiRequest, @@ -107,6 +109,40 @@ impl SamplingExecution { mut superseded: oneshot::Receiver<()>, scored: Arc, ) -> Result { + 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) => { diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler_routing_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler_routing_tests.rs new file mode 100644 index 0000000000..5a94aed3ed --- /dev/null +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler_routing_tests.rs @@ -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>>; + +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 = 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 = 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 = + 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::() + .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::from(transport == "websocket") + ); + } + Ok(()) +} diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs index f63c170620..ac2a4f2539 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs @@ -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, diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/startup.rs b/codex-rs/ext/guardian-v2/src/async_scorer/startup.rs index df05e4a933..96d8c046a5 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/startup.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/startup.rs @@ -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) { diff --git a/codex-rs/model-provider/src/workspace_routing.rs b/codex-rs/model-provider/src/workspace_routing.rs index d057a665b4..267a264446 100644 --- a/codex-rs/model-provider/src/workspace_routing.rs +++ b/codex-rs/model-provider/src/workspace_routing.rs @@ -13,11 +13,11 @@ pub const ACCOUNT_ROUTING_HEADER: &str = "x-openai-account-routing-override"; /// Routing scope for one session's provider. Once routed, it cannot become /// independent merely because the account owner's bootstrap configuration changes. -/// Lookups and the first successful routing transition are serialized per session. -#[derive(Debug)] +/// Clones serialize discovery and share the first successful routing transition. +#[derive(Clone, Debug)] pub struct WorkspaceRoutingContext { pub(crate) chatgpt_base_url: String, - pub(crate) previously_routed: Mutex, + pub(crate) previously_routed: Arc>, pub(crate) session: Option>, } @@ -25,7 +25,7 @@ impl WorkspaceRoutingContext { pub fn new(chatgpt_base_url: String) -> Self { Self { chatgpt_base_url, - previously_routed: Mutex::new(/*t*/ false), + previously_routed: Arc::new(Mutex::new(/*t*/ false)), session: None, } } diff --git a/codex-rs/model-provider/src/workspace_routing_tests.rs b/codex-rs/model-provider/src/workspace_routing_tests.rs index 95f20eea52..65ddb2bc1d 100644 --- a/codex-rs/model-provider/src/workspace_routing_tests.rs +++ b/codex-rs/model-provider/src/workspace_routing_tests.rs @@ -67,7 +67,8 @@ async fn concurrent_discovery_cannot_forget_established_routing() { Some(auth), ); let mut first = provider.responses_api_provider(&context); - let mut second = provider.responses_api_provider(&context); + let classifier_context = context.clone(); + let mut second = provider.responses_api_provider(&classifier_context); let early_second = std::future::poll_fn(|cx| { assert!(first.as_mut().poll(cx).is_pending()); std::task::Poll::Ready(second.as_mut().poll(cx))