From 58e2e8cf3c2a261394a96065510d8ef7cc2ba939 Mon Sep 17 00:00:00 2001 From: acrognale-oai Date: Wed, 16 Sep 2026 00:04:34 +0000 Subject: [PATCH] Add workspace routing support for Responses requests (#45812) ## What changed - Add an optional `WorkspaceRoutingResolver` to `AuthManager` and pass session configuration into routing lookups. Callers without a resolver retain existing routing behavior. - Apply resolved backend origins and account routing overrides to eligible ChatGPT Responses HTTP requests and WebSocket handshakes. Preserve endpoint paths, validate routing values, and reject HTTP redirects for routed requests, including `NO_CONSTRAINT` routes. - Key cached WebSocket connections by destination, routing header, and auth revision. Rebuild request setup after credential refreshes and reject account changes during setup or routing discovery. - Serialize routing lookups per session and retain whether the session has previously been routed. ## Testing Add tests for path preservation, provider exclusions, concurrent routing lookups, unavailable resolvers, workspace mismatches, redirect rejection, and credential refresh or account changes during request setup. GitOrigin-RevId: 2b6609019e84a315745c94858955f3f16426fe7f --- codex-rs/Cargo.lock | 1 + codex-rs/app-server/src/config_manager.rs | 18 +- .../src/config_manager_service_tests.rs | 4 +- codex-rs/app-server/src/mcp_refresh.rs | 5 +- .../src/request_processors/apps_processor.rs | 5 +- .../request_processors/catalog_processor.rs | 5 +- .../request_processors/config_processor.rs | 5 +- .../src/request_processors/mcp_processor.rs | 5 +- codex-rs/core/src/client.rs | 178 ++++++++--- codex-rs/core/src/client_tests.rs | 290 ++++++++++++++++++ codex-rs/core/src/config/config_tests.rs | 23 +- codex-rs/core/src/config/mod.rs | 25 +- codex-rs/core/src/session/session.rs | 1 + codex-rs/core/src/session/tests.rs | 5 + codex-rs/core/tests/responses_headers.rs | 3 + codex-rs/core/tests/suite/client.rs | 2 + .../core/tests/suite/client_websockets.rs | 1 + .../model_provider_requirements_tests.rs | 10 +- codex-rs/login/src/auth/default_client.rs | 32 +- codex-rs/login/src/auth/manager.rs | 14 + .../src/auth/manager/workspace_routing.rs | 88 ++++++ codex-rs/login/src/lib.rs | 4 + codex-rs/memories/write/src/runtime.rs | 1 + codex-rs/model-provider/Cargo.toml | 1 + codex-rs/model-provider/src/lib.rs | 9 + codex-rs/model-provider/src/provider.rs | 46 +++ .../model-provider/src/workspace_routing.rs | 103 +++++++ .../src/workspace_routing_tests.rs | 199 ++++++++++++ 28 files changed, 1003 insertions(+), 80 deletions(-) create mode 100644 codex-rs/login/src/auth/manager/workspace_routing.rs create mode 100644 codex-rs/model-provider/src/workspace_routing.rs create mode 100644 codex-rs/model-provider/src/workspace_routing_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 3a4550fd90..cdb51f237f 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -4036,6 +4036,7 @@ dependencies = [ "tempfile", "tokio", "tracing", + "url", "wiremock", ] diff --git a/codex-rs/app-server/src/config_manager.rs b/codex-rs/app-server/src/config_manager.rs index 3dca9dd143..c54c844516 100644 --- a/codex-rs/app-server/src/config_manager.rs +++ b/codex-rs/app-server/src/config_manager.rs @@ -163,16 +163,18 @@ impl ConfigManager { manager.load_latest_config(/*fallback_cwd*/ None).await } - pub(crate) async fn load_latest_config_for_thread( + pub(crate) async fn load_latest_config_with_session_layers( &self, - thread_config: &Config, + session_layers: &ConfigLayerStack, + cwd: &Path, ) -> std::io::Result { - let refreshed_config = self - .load_latest_config(Some(thread_config.cwd.to_path_buf())) - .await?; - let mut config = thread_config - .rebuild_preserving_session_layers(&refreshed_config) - .await?; + let refreshed_config = self.load_latest_config(Some(cwd.to_path_buf())).await?; + let mut config = Config::rebuild_with_session_layers( + session_layers, + cwd.to_path_buf(), + &refreshed_config, + ) + .await?; self.apply_runtime_feature_enablement(&mut config); self.apply_arg0_paths(&mut config); Ok(config) diff --git a/codex-rs/app-server/src/config_manager_service_tests.rs b/codex-rs/app-server/src/config_manager_service_tests.rs index a418909299..dae2c0c28e 100644 --- a/codex-rs/app-server/src/config_manager_service_tests.rs +++ b/codex-rs/app-server/src/config_manager_service_tests.rs @@ -1344,7 +1344,9 @@ async fn managed_auth_policy_survives_unusable_requirements_file_changes() -> Re )?; for refreshed in [ service.load_latest_config(/*fallback_cwd*/ None).await?, - service.load_latest_config_for_thread(&startup).await?, + service + .load_latest_config_with_session_layers(&startup.config_layer_stack, &startup.cwd) + .await?, ] { assert_eq!(refreshed.forced_login_method, None); assert_eq!(refreshed.forced_chatgpt_workspace_id, None); diff --git a/codex-rs/app-server/src/mcp_refresh.rs b/codex-rs/app-server/src/mcp_refresh.rs index 2953580378..b654ad6a28 100644 --- a/codex-rs/app-server/src/mcp_refresh.rs +++ b/codex-rs/app-server/src/mcp_refresh.rs @@ -57,7 +57,10 @@ async fn load_refresh_config( ) -> io::Result { let thread_config = thread.config().await; config_manager - .load_latest_config_for_thread(thread_config.as_ref()) + .load_latest_config_with_session_layers( + &thread_config.config_layer_stack, + &thread_config.cwd, + ) .await } diff --git a/codex-rs/app-server/src/request_processors/apps_processor.rs b/codex-rs/app-server/src/request_processors/apps_processor.rs index 28add513f4..be74a8fb61 100644 --- a/codex-rs/app-server/src/request_processors/apps_processor.rs +++ b/codex-rs/app-server/src/request_processors/apps_processor.rs @@ -361,7 +361,10 @@ impl AppsRequestProcessor { let (_, thread) = self.load_thread(thread_id).await?; let thread_config = thread.config().await; self.config_manager - .load_latest_config_for_thread(thread_config.as_ref()) + .load_latest_config_with_session_layers( + &thread_config.config_layer_stack, + &thread_config.cwd, + ) .await .map_err(|err| internal_error(format!("failed to reload config: {err}"))) } diff --git a/codex-rs/app-server/src/request_processors/catalog_processor.rs b/codex-rs/app-server/src/request_processors/catalog_processor.rs index ac2ed55997..45b8f132ff 100644 --- a/codex-rs/app-server/src/request_processors/catalog_processor.rs +++ b/codex-rs/app-server/src/request_processors/catalog_processor.rs @@ -331,7 +331,10 @@ impl CatalogRequestProcessor { .map_err(|_| invalid_request(format!("thread not found: {thread_id}")))?; let thread_config = thread.config().await; self.config_manager - .load_latest_config_for_thread(thread_config.as_ref()) + .load_latest_config_with_session_layers( + &thread_config.config_layer_stack, + &thread_config.cwd, + ) .await .map_err(|err| internal_error(format!("failed to reload config: {err}")))? } diff --git a/codex-rs/app-server/src/request_processors/config_processor.rs b/codex-rs/app-server/src/request_processors/config_processor.rs index b6b5348ede..adee7e1f65 100644 --- a/codex-rs/app-server/src/request_processors/config_processor.rs +++ b/codex-rs/app-server/src/request_processors/config_processor.rs @@ -371,7 +371,10 @@ pub(super) async fn reload_user_config( }; let current_config = thread.config().await; let next_config = match config_manager - .load_latest_config_for_thread(current_config.as_ref()) + .load_latest_config_with_session_layers( + ¤t_config.config_layer_stack, + ¤t_config.cwd, + ) .await { Ok(config) => config, diff --git a/codex-rs/app-server/src/request_processors/mcp_processor.rs b/codex-rs/app-server/src/request_processors/mcp_processor.rs index f6e910392e..9c22739de9 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -279,7 +279,10 @@ impl McpRequestProcessor { let thread_config = thread.config().await; let config = self .config_manager - .load_latest_config_for_thread(thread_config.as_ref()) + .load_latest_config_with_session_layers( + &thread_config.config_layer_stack, + &thread_config.cwd, + ) .await .map_err(|err| internal_error(format!("failed to reload config: {err}")))?; (config, Some(thread)) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index b60f810b99..5957c192bb 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -71,6 +71,7 @@ use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::RefreshTokenError; use codex_login::UnauthorizedRecovery; +use codex_login::default_client::ClientRedirectPolicy; use codex_login::default_client::add_originator_header; use codex_login::default_client::create_client_for_route; use codex_otel::SessionTelemetry; @@ -135,7 +136,9 @@ use codex_login::auth_env_telemetry::collect_auth_env_telemetry; use codex_model_provider::AgentIdentitySessionFallback; use codex_model_provider::ProviderAuthScope; use codex_model_provider::ProviderUnauthorizedRecovery; +use codex_model_provider::ResponsesConnectionKey; use codex_model_provider::SharedModelProvider; +use codex_model_provider::WorkspaceRoutingContext; use codex_model_provider::create_model_provider; #[cfg(test)] use codex_model_provider_info::DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS; @@ -193,6 +196,7 @@ fn session_telemetry_for_request( struct ModelClientState { thread_id: ThreadId, provider: SharedModelProvider, + workspace_routing: WorkspaceRoutingContext, auth_env_telemetry: AuthEnvTelemetry, session_source: SessionSource, originator: String, @@ -209,6 +213,11 @@ struct ModelClientState { cached_websocket_session: StdMutex, } +enum ClientRouting { + Workspace, + ConfiguredProvider, +} + /// Resolved API client setup for a single request attempt. /// /// Keeping this as a single bundle ensures prewarm and normal request paths @@ -216,7 +225,9 @@ struct ModelClientState { struct CurrentClientSetup { auth: Option, auth_owner_generation: Option, + auth_revision: Option, api_provider: ApiProvider, + redirect_policy: ClientRedirectPolicy, api_auth: SharedAuthProvider, agent_identity_telemetry: Option, } @@ -295,6 +306,7 @@ struct WebsocketSession { responses_headers: ApiHeaderMap, /// Owner of the cached state, including before a connection is opened. auth_owner_generation: Option, + connection_key: Option, last_request: Option, last_response_rx: Option>, last_response_from_untraced_warmup: bool, @@ -443,6 +455,7 @@ impl ModelClient { concurrent_reasoning_summaries_enabled: bool, attestation_provider: Option>, http_client_factory: HttpClientFactory, + workspace_routing: WorkspaceRoutingContext, ) -> Self { let model_provider = create_model_provider(provider_info, auth_manager); let codex_api_key_env_enabled = model_provider @@ -456,6 +469,7 @@ impl ModelClient { state: Arc::new(ModelClientState { thread_id, provider: model_provider, + workspace_routing, auth_env_telemetry, session_source, originator, @@ -596,7 +610,9 @@ impl ModelClient { ) -> Result { // Create the media call over HTTP first, then retain matching auth so realtime can attach // the server-side control WebSocket to the call id from that HTTP response. - let client_setup = self.current_client_setup().await?; + let client_setup = self + .current_client_setup(ClientRouting::ConfiguredProvider) + .await?; if let Some(header_value) = self.generate_attestation_header_for().await { extra_headers.insert(X_OAI_ATTESTATION_HEADER, header_value); } @@ -605,7 +621,11 @@ impl ModelClient { client_setup.api_auth.as_ref(), )); let api_provider = api_provider_override.unwrap_or(client_setup.api_provider); - let transport = self.build_api_transport(&api_provider, REALTIME_CALLS_ENDPOINT)?; + let transport = self.build_api_transport( + &api_provider, + REALTIME_CALLS_ENDPOINT, + client_setup.redirect_policy, + )?; let response = ApiRealtimeCallClient::new(transport, api_provider, client_setup.api_auth) .create_with_session_and_headers(sdp, session_config, extra_headers) .await @@ -621,7 +641,9 @@ impl ModelClient { &self, mut extra_headers: ApiHeaderMap, ) -> Result { - let client_setup = self.current_client_setup().await?; + let client_setup = self + .current_client_setup(ClientRouting::ConfiguredProvider) + .await?; if let Some(header_value) = self.generate_attestation_header_for().await { extra_headers.insert(X_OAI_ATTESTATION_HEADER, header_value); } @@ -648,9 +670,14 @@ impl ModelClient { return Ok(Vec::new()); } - let client_setup = self.current_client_setup().await?; - let transport = - self.build_api_transport(&client_setup.api_provider, MEMORIES_SUMMARIZE_ENDPOINT)?; + let client_setup = self + .current_client_setup(ClientRouting::ConfiguredProvider) + .await?; + let transport = self.build_api_transport( + &client_setup.api_provider, + MEMORIES_SUMMARIZE_ENDPOINT, + client_setup.redirect_policy, + )?; let request_telemetry = Self::build_request_telemetry( session_telemetry, AuthRequestTelemetryContext::new( @@ -944,28 +971,68 @@ impl ModelClient { /// /// This centralizes setup used by both prewarm and normal request paths so they stay in /// lockstep when auth/provider resolution changes. - async fn current_client_setup(&self) -> Result { + async fn current_client_setup(&self, routing: ClientRouting) -> Result { // Capture before resolving credentials so an account switch during setup cannot label // an old connection with the new owner's revision. let auth_owner_generation = self.auth_owner_generation(); - let auth = self.state.provider.auth().await; - let api_provider = self.state.provider.api_provider().await?; - let resolved_auth = self + let auth_changes = self .state .provider - .api_auth_for_scope(ProviderAuthScope { - agent_identity_policy: self.agent_identity_policy, - session_source: self.state.session_source.clone(), - agent_identity_session_fallback: self.state.agent_identity_session_fallback.clone(), - }) - .await?; - Ok(CurrentClientSetup { - auth, - auth_owner_generation, - api_provider, - api_auth: resolved_auth.auth, - agent_identity_telemetry: resolved_auth.agent_identity_telemetry, - }) + .auth_manager() + .map(|manager| manager.auth_change_receiver()); + loop { + let revision = auth_changes.as_ref().map(|changes| *changes.borrow()); + let auth = self.state.provider.auth().await; + let (api_provider, redirect_policy) = match routing { + ClientRouting::Workspace => { + let resolved = self + .state + .provider + .responses_api_provider(&self.state.workspace_routing) + .await?; + (resolved.provider, resolved.redirect_policy) + } + ClientRouting::ConfiguredProvider => ( + self.state.provider.api_provider().await?, + ClientRedirectPolicy::Default, + ), + }; + let resolved_auth = self + .state + .provider + .api_auth_for_scope(ProviderAuthScope { + agent_identity_policy: self.agent_identity_policy, + session_source: self.state.session_source.clone(), + agent_identity_session_fallback: self + .state + .agent_identity_session_fallback + .clone(), + }) + .await?; + // Command-backed bearer refreshes cannot change workspace routing. Keep the captured + // revisions so a refresh during setup still invalidates cached WebSocket state. + if self.state.provider.info().auth.is_none() { + if self.auth_owner_generation() != auth_owner_generation { + return Err(std::io::Error::other( + "account changed while preparing model request", + ) + .into()); + } + // Rebuild the entire bundle after a credential-only refresh, including routing. + if auth_changes.as_ref().map(|changes| *changes.borrow()) != revision { + continue; + } + } + return Ok(CurrentClientSetup { + auth, + auth_owner_generation, + auth_revision: revision, + api_provider, + redirect_policy, + api_auth: resolved_auth.auth, + agent_identity_telemetry: resolved_auth.agent_identity_telemetry, + }); + } } fn responses_headers(&self, auth: Option<&CodexAuth>, model: &str) -> ApiHeaderMap { @@ -1047,19 +1114,30 @@ impl ModelClient { &self, api_provider: &ApiProvider, endpoint: &str, + redirect_policy: ClientRedirectPolicy, ) -> Result { - let request_url = api_provider.url_for_path(endpoint); + let redirect_policy = if api_provider + .headers + .contains_key(codex_model_provider::ACCOUNT_ROUTING_HEADER) + { + ClientRedirectPolicy::Reject + } else { + redirect_policy + }; let client = create_client_for_route( &self.http_client_factory, - &request_url, + &api_provider.url_for_path(endpoint), ClientRouteClass::Api, + redirect_policy, ) .map_err(std::io::Error::from)?; Ok(ReqwestTransport::from_http_client(client)) } pub(crate) async fn prewarm_auth(&self) -> Result<()> { - self.current_client_setup().await.map(|_| ()) + self.current_client_setup(ClientRouting::Workspace) + .await + .map(|_| ()) } /// Opens a websocket connection using the same header and telemetry wiring as normal turns. @@ -1203,6 +1281,7 @@ impl ModelClientSession { fn reset_websocket_session(&mut self) { self.websocket_session.connection = None; self.websocket_session.responses_headers.clear(); + self.websocket_session.connection_key = None; self.websocket_session.last_request = None; self.websocket_session.last_response_rx = None; self.websocket_session.last_response_from_untraced_warmup = false; @@ -1340,15 +1419,23 @@ impl ModelClientSession { if !self.client.responses_websocket_enabled() { return Ok(()); } - if self.websocket_session.connection.is_some() { + let client_setup = self + .client + .current_client_setup(ClientRouting::Workspace) + .await + .map_err(|err| { + ApiError::Stream(format!( + "failed to build websocket prewarm client setup: {err}" + )) + })?; + let connection_key = + ResponsesConnectionKey::new(&client_setup.api_provider, client_setup.auth_revision); + if self.websocket_session.connection.is_some() + && self.websocket_session.connection_key.as_ref() == Some(&connection_key) + { return Ok(()); } - - let client_setup = self.client.current_client_setup().await.map_err(|err| { - ApiError::Stream(format!( - "failed to build websocket prewarm client setup: {err}" - )) - })?; + self.reset_websocket_session(); let auth_context = AuthRequestTelemetryContext::new( client_setup.auth.as_ref().map(CodexAuth::auth_mode), client_setup.api_auth.as_ref(), @@ -1361,6 +1448,7 @@ impl ModelClientSession { self.websocket_connection(WebsocketConnectParams { session_telemetry, api_provider: client_setup.api_provider, + auth_revision: client_setup.auth_revision, api_auth: client_setup.api_auth, auth_owner_generation: client_setup.auth_owner_generation, responses_metadata, @@ -1391,6 +1479,7 @@ impl ModelClientSession { let WebsocketConnectParams { session_telemetry, api_provider, + auth_revision, api_auth, auth_owner_generation, responses_metadata, @@ -1398,9 +1487,11 @@ impl ModelClientSession { request_route_telemetry, responses_headers, } = params; + let connection_key = ResponsesConnectionKey::new(&api_provider, auth_revision); let needs_new = match self.websocket_session.connection.as_ref() { Some(conn) => { self.websocket_session.responses_headers != *responses_headers + || self.websocket_session.connection_key.as_ref() != Some(&connection_key) || conn.is_closed().await } None => true, @@ -1438,6 +1529,7 @@ impl ModelClientSession { self.websocket_session.connection = Some(new_conn); self.websocket_session.responses_headers = responses_headers.clone(); self.websocket_session.auth_owner_generation = auth_owner_generation; + self.websocket_session.connection_key = Some(connection_key); self.websocket_session .set_connection_reused(/*connection_reused*/ false); } else { @@ -1499,14 +1591,19 @@ impl ModelClientSession { let mut provider_auth_recovery_attempted = false; let mut pending_retry = PendingUnauthorizedRetry::default(); loop { - let client_setup = self.client.current_client_setup().await?; + let client_setup = self + .client + .current_client_setup(ClientRouting::Workspace) + .await?; let responses_headers = self .client .responses_headers(client_setup.auth.as_ref(), &model_info.slug); tracing::Span::current().record("api.path", "/responses"); - let transport = self - .client - .build_api_transport(&client_setup.api_provider, "/responses")?; + let transport = self.client.build_api_transport( + &client_setup.api_provider, + "/responses", + client_setup.redirect_policy, + )?; let request_auth_context = AuthRequestTelemetryContext::new( client_setup.auth.as_ref().map(CodexAuth::auth_mode), client_setup.api_auth.as_ref(), @@ -1676,7 +1773,10 @@ impl ModelClientSession { let mut provider_auth_recovery_attempted = false; let mut pending_retry = PendingUnauthorizedRetry::default(); loop { - let client_setup = self.client.current_client_setup().await?; + let client_setup = self + .client + .current_client_setup(ClientRouting::Workspace) + .await?; let responses_headers = self .client .responses_headers(client_setup.auth.as_ref(), &model_info.slug); @@ -1729,6 +1829,7 @@ impl ModelClientSession { .websocket_connection(WebsocketConnectParams { session_telemetry, api_provider: client_setup.api_provider, + auth_revision: client_setup.auth_revision, api_auth: client_setup.api_auth, auth_owner_generation: client_setup.auth_owner_generation, responses_metadata: &websocket_metadata, @@ -2342,6 +2443,7 @@ impl AuthRequestTelemetryContext { struct WebsocketConnectParams<'a> { session_telemetry: &'a SessionTelemetry, api_provider: codex_api::Provider, + auth_revision: Option, api_auth: SharedAuthProvider, auth_owner_generation: Option, responses_metadata: &'a CodexResponsesMetadata, diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 36d54c9cc4..b129556af1 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -14,6 +14,7 @@ use crate::GenerateAttestationFuture; use crate::responses_metadata::CodexResponsesMetadata; use crate::test_support::TestCodexResponsesRequestKind; use crate::test_support::responses_metadata as test_responses_metadata; +use base64::Engine; use codex_api::AgentIdentityTelemetry; use codex_api::ApiError; use codex_api::ResponseEvent; @@ -120,6 +121,9 @@ fn test_model_client_with_thread_id( /*concurrent_reasoning_summaries_enabled*/ false, /*attestation_provider*/ None, HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + codex_model_provider::WorkspaceRoutingContext::new( + "https://chatgpt.com/backend-api".into(), + ), ) } @@ -127,6 +131,289 @@ fn test_model_provider() -> SharedModelProvider { test_model_client(SessionSource::Cli).state.provider.clone() } +#[tokio::test] +async fn workspace_routed_http_rejects_redirects_without_a_routing_header() { + use codex_client::HttpTransport; + use codex_login::WorkspaceRouting; + use codex_login::WorkspaceRoutingRequest; + use codex_login::WorkspaceRoutingResolver; + + struct Routing(Option<&'static str>); + impl WorkspaceRoutingResolver for Routing { + fn resolve( + &self, + _request: WorkspaceRoutingRequest, + ) -> Pin< + Box< + dyn std::future::Future>> + + Send + + '_, + >, + > { + Box::pin(async move { + Ok(self.0.map(|override_value| WorkspaceRouting { + chatgpt_account_id: "account_id".into(), + backend_origin: "https://gov.chatgpt.com".into(), + account_routing_override: override_value.into(), + })) + }) + } + } + + for routing_override in [Some("NO_CONSTRAINT"), Some("us_cr"), None] { + let origin = MockServer::start().await; + let destination = MockServer::start().await; + Mock::given(method("POST")) + .respond_with( + ResponseTemplate::new(/*status*/ 307) + .insert_header("location", format!("{}/responses", destination.uri())), + ) + .mount(&origin) + .await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(/*status*/ 200)) + .mount(&destination) + .await; + let manager = + AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let resolver: Arc = Arc::new(Routing(routing_override)); + manager.set_workspace_routing_resolver(Arc::downgrade(&resolver)); + let mut client = test_model_client(SessionSource::Exec); + Arc::get_mut(&mut client.state).unwrap().provider = create_model_provider( + ModelProviderInfo::create_openai_provider(/*base_url*/ None), + Some(manager), + ); + let mut setup = client + .current_client_setup(super::ClientRouting::Workspace) + .await + .unwrap(); + if routing_override.is_some() { + assert_eq!( + setup.api_provider.base_url, + "https://gov.chatgpt.com/backend-api/codex" + ); + } + // Exercise the resolved route's redirect policy against loopback HTTP servers. + setup.api_provider.base_url = origin.uri(); + let transport = client + .build_api_transport(&setup.api_provider, "/responses", setup.redirect_policy) + .unwrap(); + let request = setup + .api_provider + .build_request(http::Method::POST, "/responses") + .with_json(&json!({"input": "workspace content"})); + let result = transport.execute(request).await; + if routing_override.is_some() { + assert!( + matches!( + result, + Err(TransportError::Http { + status: http::StatusCode::TEMPORARY_REDIRECT, + .. + }) + ), + "workspace redirect must be rejected: {routing_override:?}" + ); + } else { + assert_eq!(result.unwrap().status, http::StatusCode::OK); + } + assert_eq!( + destination.received_requests().await.unwrap().len(), + usize::from(routing_override.is_none()) + ); + } +} + +#[derive(Debug)] +enum SetupRefresh { + Command(PathBuf), + ChatGpt { + home: PathBuf, + token: String, + workspace: String, + }, +} + +#[derive(Debug)] +struct SetupRefreshProvider { + inner: SharedModelProvider, + refresh: SetupRefresh, + setup_calls: AtomicUsize, +} + +impl ModelProvider for SetupRefreshProvider { + fn info(&self) -> &ModelProviderInfo { + self.inner.info() + } + + fn auth_manager(&self) -> Option> { + self.inner.auth_manager() + } + + fn auth(&self) -> ModelProviderFuture<'_, Option> { + self.inner.auth() + } + + fn account_state(&self) -> ProviderAccountResult { + self.inner.account_state() + } + + fn api_provider( + &self, + ) -> ModelProviderFuture<'_, codex_protocol::error::Result> { + Box::pin(async move { + self.setup_calls.fetch_add(1, Ordering::SeqCst); + let manager = self.inner.auth_manager().expect("auth manager"); + match &self.refresh { + SetupRefresh::Command(token_path) => { + std::fs::write(token_path, "refreshed-token")?; + manager + .refresh_token_from_authority() + .await + .expect("refresh command token"); + } + SetupRefresh::ChatGpt { + home, + token, + workspace, + } => { + codex_login::auth::login_with_chatgpt_auth_tokens( + home, token, workspace, /*chatgpt_plan_type*/ None, + )?; + manager.reload().await; + } + } + self.inner.api_provider().await + }) + } + + fn models_manager( + &self, + codex_home: PathBuf, + config_model_catalog: Option, + ) -> SharedModelsManager { + self.inner.models_manager(codex_home, config_model_catalog) + } +} + +#[tokio::test] +async fn client_setup_accepts_command_credential_refresh() { + for routing in [ + super::ClientRouting::Workspace, + super::ClientRouting::ConfiguredProvider, + ] { + let tempdir = TempDir::new().unwrap(); + let token_path = tempdir.path().join("token.txt"); + std::fs::write(&token_path, "initial-token").unwrap(); + let mut info = test_model_provider().info().clone(); + info.auth = Some(codex_protocol::config_types::ModelProviderAuthInfo { + command: if cfg!(windows) { "cmd.exe" } else { "cat" }.into(), + args: if cfg!(windows) { + vec!["/D", "/C", "type", "token.txt"] + } else { + vec!["token.txt"] + } + .into_iter() + .map(Into::into) + .collect(), + timeout_ms: std::num::NonZeroU64::new(/*n*/ 5_000).unwrap(), + refresh_interval_ms: 60_000, + cwd: tempdir.path().try_into().unwrap(), + }); + let provider = Arc::new(SetupRefreshProvider { + inner: create_model_provider(info, /*auth_manager*/ None), + refresh: SetupRefresh::Command(token_path), + setup_calls: AtomicUsize::new(/*v*/ 0), + }); + let manager = provider.auth_manager().unwrap(); + let mut client = test_model_client(SessionSource::Exec); + Arc::get_mut(&mut client.state).unwrap().provider = provider; + + let setup = client.current_client_setup(routing).await.unwrap(); + let mut headers = http::HeaderMap::new(); + setup.api_auth.add_auth_headers(&mut headers); + assert_eq!( + headers.get(http::header::AUTHORIZATION).unwrap(), + "Bearer refreshed-token" + ); + let refreshed_revision = Some(*manager.auth_change_receiver().borrow()); + assert_ne!( + codex_model_provider::ResponsesConnectionKey::new( + &setup.api_provider, + setup.auth_revision + ), + codex_model_provider::ResponsesConnectionKey::new( + &setup.api_provider, + refreshed_revision + ), + ); + assert_ne!(setup.auth_owner_generation, client.auth_owner_generation()); + } +} + +#[tokio::test] +async fn client_setup_rebuilds_chatgpt_refresh_but_rejects_account_switches() { + for (user, workspace, expected_calls) in [ + ("user-a", "workspace-a", 2), + ("user-b", "workspace-a", 1), + ("user-a", "workspace-b", 1), + ] { + let token = |user: &str, revision: &str| { + let claims = + json!({"jti": revision, "https://api.openai.com/auth": {"chatgpt_user_id": user}}); + let payload = + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(claims.to_string()); + format!("header.{payload}.signature") + }; + let home = TempDir::new().unwrap(); + let initial = CodexAuth::from_external_chatgpt_tokens( + &token("user-a", "initial"), + "workspace-a", + /*chatgpt_plan_type*/ None, + ) + .unwrap(); + let manager = + AuthManager::from_auth_for_testing_with_home(initial, home.path().to_path_buf()); + let refreshed_token = token(user, "refreshed"); + let mut info = test_model_provider().info().clone(); + info.requires_openai_auth = true; + let provider = Arc::new(SetupRefreshProvider { + inner: create_model_provider(info, Some(manager.clone())), + refresh: SetupRefresh::ChatGpt { + home: home.path().to_path_buf(), + token: refreshed_token.clone(), + workspace: workspace.into(), + }, + setup_calls: AtomicUsize::new(/*v*/ 0), + }); + let mut client = test_model_client(SessionSource::Exec); + Arc::get_mut(&mut client.state).unwrap().provider = provider.clone(); + let result = client + .current_client_setup(super::ClientRouting::ConfiguredProvider) + .await; + if expected_calls == 2 { + let setup = result.unwrap(); + let mut headers = http::HeaderMap::new(); + setup.api_auth.add_auth_headers(&mut headers); + assert_eq!( + headers.get(http::header::AUTHORIZATION).unwrap(), + &format!("Bearer {refreshed_token}") + ); + assert_eq!(setup.auth.unwrap().get_token().unwrap(), refreshed_token); + assert_eq!( + (setup.auth_revision, setup.auth_owner_generation), + (Some(*manager.auth_change_receiver().borrow()), Some(0)) + ); + } else { + assert_eq!( + result.err().expect("account switch must fail").to_string(), + "account changed while preparing model request" + ); + } + assert_eq!(provider.setup_calls.load(Ordering::SeqCst), expected_calls); + } +} + fn test_responses_metadata_for_client( client: &ModelClient, turn_id: Option<&str>, @@ -1197,6 +1484,9 @@ fn model_client_with_counting_attestation( calls: attestation_calls.clone(), })), HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + codex_model_provider::WorkspaceRoutingContext::new( + "https://chatgpt.com/backend-api".into(), + ), ); (model_client, attestation_calls) } diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index 748ecc8410..d4e1adf178 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -5307,7 +5307,7 @@ fn filter_plugin_mcp_servers_by_matchers_enforces_name_and_invocation() { } #[tokio::test] -async fn rebuild_preserving_session_layers_refreshes_requirements() -> std::io::Result<()> { +async fn rebuild_with_session_layers_refreshes_requirements() -> std::io::Result<()> { let codex_home = TempDir::new()?; let user_file = AbsolutePathBuf::resolve_path_against_base(CONFIG_TOML_FILE, codex_home.path()); let project_dot_codex = @@ -5476,9 +5476,12 @@ async fn rebuild_preserving_session_layers_refreshes_requirements() -> std::io:: thread_layer_stack, ) .await?; - let config = thread_config - .rebuild_preserving_session_layers(&refreshed_config) - .await?; + let config = Config::rebuild_with_session_layers( + &thread_config.config_layer_stack, + thread_config.cwd.to_path_buf(), + &refreshed_config, + ) + .await?; assert_eq!( config.mcp_servers.get(), @@ -5516,8 +5519,7 @@ async fn rebuild_preserving_session_layers_refreshes_requirements() -> std::io:: } #[tokio::test] -async fn rebuild_preserving_session_layers_refreshes_plugin_derived_mcp_config() --> anyhow::Result<()> { +async fn rebuild_with_session_layers_refreshes_plugin_derived_mcp_config() -> anyhow::Result<()> { let codex_home = TempDir::new()?; let plugin_root = codex_home .path() @@ -5599,9 +5601,12 @@ async fn rebuild_preserving_session_layers_refreshes_plugin_derived_mcp_config() thread_layer_stack, ) .await?; - let config = thread_config - .rebuild_preserving_session_layers(&refreshed_config) - .await?; + let config = Config::rebuild_with_session_layers( + &thread_config.config_layer_stack, + thread_config.cwd.to_path_buf(), + &refreshed_config, + ) + .await?; let plugins_manager = plugins_manager_for_config(&config, auth_manager_from_optional_auth(/*auth*/ None)); let mcp_config = config.to_mcp_config(&plugins_manager).await; diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index ef972143cf..789b207da3 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -1826,12 +1826,23 @@ impl Config { } } - pub async fn rebuild_preserving_session_layers( - &self, + pub fn workspace_routing_context(&self) -> codex_model_provider::WorkspaceRoutingContext { + codex_model_provider::WorkspaceRoutingContext::new(self.chatgpt_base_url.clone()) + .with_session(codex_login::WorkspaceRoutingSession { + cwd: self.cwd.to_path_buf(), + config_layer_stack: self.config_layer_stack.clone(), + }) + } + + pub async fn rebuild_with_session_layers( + session_layers: &ConfigLayerStack, + cwd: PathBuf, refreshed_config: &Config, ) -> std::io::Result { - let config_layer_stack = - self.layer_stack_preserving_session(&refreshed_config.config_layer_stack)?; + let config_layer_stack = Self::layer_stack_preserving_session( + session_layers, + &refreshed_config.config_layer_stack, + )?; let cfg: ConfigToml = config_layer_stack .effective_config() .try_into() @@ -1846,7 +1857,7 @@ impl Config { LOCAL_FS.as_ref(), cfg, ConfigOverrides { - cwd: Some(self.cwd.to_path_buf()), + cwd: Some(cwd), default_zsh_path, ..Default::default() }, @@ -1857,7 +1868,7 @@ impl Config { } fn layer_stack_preserving_session( - &self, + session_layers: &ConfigLayerStack, refreshed_layers: &ConfigLayerStack, ) -> std::io::Result { let mut layers = refreshed_layers @@ -1866,7 +1877,7 @@ impl Config { .cloned() .collect::>(); layers.extend( - self.config_layer_stack + session_layers .all_layers_low_to_high() .filter(|layer| is_session_layer(&layer.name)) .cloned(), diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 83e6024ce6..97e067c759 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1673,6 +1673,7 @@ impl Session { .enabled(Feature::ConcurrentReasoningSummaries), attestation_provider, config.http_client_factory(), + config.workspace_routing_context(), ) .with_session_context( crate::guardian::prompt_cache_key_override_for_review_session( diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 00e7e935ca..745756d48a 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -743,6 +743,9 @@ fn test_model_client_session() -> crate::client::ModelClientSession { /*concurrent_reasoning_summaries_enabled*/ false, /*attestation_provider*/ None, HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + codex_model_provider::WorkspaceRoutingContext::new( + "https://chatgpt.com/backend-api".into(), + ), ) .new_session() } @@ -6119,6 +6122,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { .enabled(Feature::ConcurrentReasoningSummaries), /*attestation_provider*/ None, config.http_client_factory(), + config.workspace_routing_context(), ), executed_tool_calls: executed_tool_calls.clone(), code_mode_service: crate::tools::code_mode::CodeModeService::new( @@ -8304,6 +8308,7 @@ where .enabled(Feature::ConcurrentReasoningSummaries), /*attestation_provider*/ None, config.http_client_factory(), + config.workspace_routing_context(), ), executed_tool_calls: executed_tool_calls.clone(), code_mode_service: crate::tools::code_mode::CodeModeService::new( diff --git a/codex-rs/core/tests/responses_headers.rs b/codex-rs/core/tests/responses_headers.rs index be184753e7..b3b1e86e91 100644 --- a/codex-rs/core/tests/responses_headers.rs +++ b/codex-rs/core/tests/responses_headers.rs @@ -136,6 +136,7 @@ async fn responses_stream_includes_subagent_header_on_review() { /*concurrent_reasoning_summaries_enabled*/ false, /*attestation_provider*/ None, config.http_client_factory(), + config.workspace_routing_context(), ); let responses_metadata = test_turn_responses_metadata(&client, thread_id, &session_source); let mut client_session = client.new_session(); @@ -273,6 +274,7 @@ async fn responses_stream_includes_subagent_header_on_other() { /*concurrent_reasoning_summaries_enabled*/ false, /*attestation_provider*/ None, config.http_client_factory(), + config.workspace_routing_context(), ); let responses_metadata = test_turn_responses_metadata(&client, thread_id, &session_source); let mut client_session = client.new_session(); @@ -395,6 +397,7 @@ async fn responses_respects_model_info_overrides_from_config() { /*concurrent_reasoning_summaries_enabled*/ false, /*attestation_provider*/ None, config.http_client_factory(), + config.workspace_routing_context(), ); let responses_metadata = test_turn_responses_metadata(&client, thread_id, &session_source); let mut client_session = client.new_session(); diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index 79d0eb4f6d..502f3ccf38 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -1579,6 +1579,7 @@ async fn send_request_with_provider(provider: ModelProviderInfo) { .enabled(Feature::ConcurrentReasoningSummaries), /*attestation_provider*/ None, config.http_client_factory(), + config.workspace_routing_context(), ); let responses_metadata = test_turn_responses_metadata(&client, thread_id); let mut client_session = client.new_session(); @@ -3070,6 +3071,7 @@ async fn azure_responses_request_does_not_store_and_preserves_prefixed_item_ids( /*concurrent_reasoning_summaries_enabled*/ false, /*attestation_provider*/ None, config.http_client_factory(), + config.workspace_routing_context(), ); let responses_metadata = test_turn_responses_metadata(&client, thread_id); let mut client_session = client.new_session(); diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index 4d2ba3c926..6a72bffb3b 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -2836,6 +2836,7 @@ async fn websocket_harness_with_provider_options_and_auth( .enabled(Feature::ConcurrentReasoningSummaries), /*attestation_provider*/ None, http_client_factory, + config.workspace_routing_context(), ); WebsocketTestHarness { diff --git a/codex-rs/core/tests/suite/model_provider_requirements_tests.rs b/codex-rs/core/tests/suite/model_provider_requirements_tests.rs index c096bf8fb1..a6090f1de5 100644 --- a/codex-rs/core/tests/suite/model_provider_requirements_tests.rs +++ b/codex-rs/core/tests/suite/model_provider_requirements_tests.rs @@ -4,6 +4,7 @@ use anyhow::Result; use codex_config::LoaderOverrides; use codex_config::config_toml::ConfigToml; use codex_config::test_support::CloudConfigBundleFixture; +use codex_core::config::Config; use codex_core::config::ConfigBuilder; use codex_core::config::ConfigOverrides; use codex_login::CodexAuth; @@ -177,9 +178,12 @@ X-Managed = "required" insta::assert_debug_snapshot!("provider_selection_override_warning", warnings); } assert_eq!(&overridden.model_provider, &test.config.model_provider); - let rebuilt = overridden - .rebuild_preserving_session_layers(&overridden) - .await?; + let rebuilt = Config::rebuild_with_session_layers( + &overridden.config_layer_stack, + overridden.cwd.to_path_buf(), + &overridden, + ) + .await?; assert_eq!(rebuilt.model_provider, test.config.model_provider); } } diff --git a/codex-rs/login/src/auth/default_client.rs b/codex-rs/login/src/auth/default_client.rs index cd4cbd5431..0d6a0fcc99 100644 --- a/codex-rs/login/src/auth/default_client.rs +++ b/codex-rs/login/src/auth/default_client.rs @@ -253,24 +253,32 @@ pub fn create_client_for_route( http_client_factory: &HttpClientFactory, request_url: &str, route_class: ClientRouteClass, + redirect_policy: ClientRedirectPolicy, ) -> Result { + let builder = match redirect_policy { + ClientRedirectPolicy::Default => default_http_client_builder(), + ClientRedirectPolicy::Reject => default_http_client_builder().without_redirects(), + }; if matches!( http_client_factory.outbound_proxy_policy(), OutboundProxyPolicy::ReqwestDefault ) { - return Ok(create_client()); + return Ok(build_default_client(builder)); } if is_sandboxed() { // Preserve the sandbox's existing no-proxy policy; sandboxed command egress is routed // separately through network-proxy. - return Ok(create_client()); + return Ok(build_default_client(builder)); } - default_http_client_builder().build_respecting_outbound_proxy_policy( - http_client_factory, - request_url, - route_class, - ) + builder.build_respecting_outbound_proxy_policy(http_client_factory, request_url, route_class) +} + +/// Redirect handling for the shared HTTP client builders. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ClientRedirectPolicy { + Default, + Reject, } /// Builds the default Codex HTTP client for a concrete outbound route without blocking the @@ -286,8 +294,13 @@ pub async fn create_client_for_route_async( .map_err(std::io::Error::other)?; tokio::task::spawn_blocking(move || { let _permit = permit; - create_client_for_route(&http_client_factory, &request_url, route_class) - .map_err(std::io::Error::from) + create_client_for_route( + &http_client_factory, + &request_url, + route_class, + ClientRedirectPolicy::Default, + ) + .map_err(std::io::Error::from) }) .await .map_err(std::io::Error::other)? @@ -329,6 +342,7 @@ pub(crate) fn create_default_auth_client( auth_route_config.http_client_factory(), endpoint, ClientRouteClass::Auth, + ClientRedirectPolicy::Default, ) } diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index baf243faf6..db0f2e2057 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -1,3 +1,5 @@ +mod workspace_routing; + use chrono::Utc; use http::StatusCode; use serde::Deserialize; @@ -12,7 +14,9 @@ use std::path::PathBuf; use std::pin::Pin; use std::sync::Arc; use std::sync::Mutex; +use std::sync::OnceLock; use std::sync::RwLock; +use std::sync::Weak; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; use std::time::Duration; @@ -74,6 +78,10 @@ use codex_protocol::auth::RefreshTokenFailedReason; use codex_protocol::protocol::SessionSource; use serde_json::Value; use thiserror::Error; +pub use workspace_routing::WorkspaceRouting; +pub use workspace_routing::WorkspaceRoutingRequest; +pub use workspace_routing::WorkspaceRoutingResolver; +pub use workspace_routing::WorkspaceRoutingSession; /// Authentication mechanism used by the current user. #[derive(Debug, Clone)] @@ -2051,6 +2059,7 @@ pub struct AuthManager { inner: RwLock, auth_change_tx: watch::Sender, auth_change_state_tx: watch::Sender, + workspace_routing_resolver: OnceLock>, enable_codex_api_key_env: bool, auth_credentials_store_mode: AuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, @@ -2188,6 +2197,7 @@ impl AuthManager { }), auth_change_tx, auth_change_state_tx: watch::channel(AuthChangeState::default()).0, + workspace_routing_resolver: OnceLock::new(), enable_codex_api_key_env, auth_credentials_store_mode, keyring_backend_kind, @@ -2223,6 +2233,7 @@ impl AuthManager { inner: RwLock::new(cached), auth_change_tx, auth_change_state_tx: watch::channel(AuthChangeState::default()).0, + workspace_routing_resolver: OnceLock::new(), enable_codex_api_key_env: false, auth_credentials_store_mode: AuthCredentialsStoreMode::File, keyring_backend_kind: AuthKeyringBackendKind::default(), @@ -2252,6 +2263,7 @@ impl AuthManager { inner: RwLock::new(cached), auth_change_tx, auth_change_state_tx: watch::channel(AuthChangeState::default()).0, + workspace_routing_resolver: OnceLock::new(), enable_codex_api_key_env: false, auth_credentials_store_mode: AuthCredentialsStoreMode::File, keyring_backend_kind: AuthKeyringBackendKind::default(), @@ -2285,6 +2297,7 @@ impl AuthManager { inner: RwLock::new(cached), auth_change_tx, auth_change_state_tx: watch::channel(AuthChangeState::default()).0, + workspace_routing_resolver: OnceLock::new(), enable_codex_api_key_env: false, auth_credentials_store_mode: AuthCredentialsStoreMode::File, keyring_backend_kind: AuthKeyringBackendKind::default(), @@ -2316,6 +2329,7 @@ impl AuthManager { }), auth_change_tx, auth_change_state_tx: watch::channel(AuthChangeState::default()).0, + workspace_routing_resolver: OnceLock::new(), enable_codex_api_key_env: false, auth_credentials_store_mode: AuthCredentialsStoreMode::File, keyring_backend_kind: AuthKeyringBackendKind::default(), diff --git a/codex-rs/login/src/auth/manager/workspace_routing.rs b/codex-rs/login/src/auth/manager/workspace_routing.rs new file mode 100644 index 0000000000..0ceb56a35d --- /dev/null +++ b/codex-rs/login/src/auth/manager/workspace_routing.rs @@ -0,0 +1,88 @@ +//! Connects account-owned routing discovery to ChatGPT request builders. +//! The resolver is weakly held so it cannot keep its app-server owner alive. + +use codex_config::ConfigLayerStack; +use std::future::Future; +use std::io; +use std::path::PathBuf; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::Weak; + +use crate::AuthManager; +use crate::CodexAuth; + +/// A successful discovery for one selected workspace. +#[derive(Clone, Debug)] +pub struct WorkspaceRouting { + pub chatgpt_account_id: String, + pub backend_origin: String, + pub account_routing_override: String, +} + +/// Retained configuration layers and project location for one model session. +/// The account owner refreshes global settings without discarding session overrides. +pub struct WorkspaceRoutingSession { + pub cwd: PathBuf, + pub config_layer_stack: ConfigLayerStack, +} + +impl std::fmt::Debug for WorkspaceRoutingSession { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("WorkspaceRoutingSession") + .finish_non_exhaustive() + } +} + +/// The provider and bootstrap selected for one model session. +pub struct WorkspaceRoutingRequest { + pub provider_base_url: String, + pub chatgpt_base_url: String, + pub previously_routed: bool, + pub session: Option>, +} + +/// Resolves routing through the account owner's cache and requirements loader. +/// Returns no routing for independent destinations or no selected workspace. +/// Destination scope is decided here, before making a discovery request. +/// Implementations must reject failed discovery and changes to the selected account. +pub trait WorkspaceRoutingResolver: Send + Sync { + fn resolve( + &self, + request: WorkspaceRoutingRequest, + ) -> Pin>> + Send + '_>>; +} + +impl AuthManager { + /// Installs the app-server's routing owner before accepting model requests. + pub fn set_workspace_routing_resolver(&self, resolver: Weak) { + assert!(self.workspace_routing_resolver.set(resolver).is_ok()); + } + + /// CLI callers without a discovery owner retain their existing routing behavior. + pub async fn workspace_routing( + &self, + auth: &CodexAuth, + request: WorkspaceRoutingRequest, + ) -> io::Result> { + let Some(resolver) = self.workspace_routing_resolver.get() else { + return Ok(None); + }; + let auth_changes = self.auth_change_receiver(); + let revision = *auth_changes.borrow(); + let resolver = resolver + .upgrade() + .ok_or_else(|| io::Error::other("workspace routing owner is unavailable"))?; + let routing = resolver.resolve(request).await?; + if *auth_changes.borrow() != revision + || routing.as_ref().is_some_and(|routing| { + auth.get_account_id().as_deref() != Some(routing.chatgpt_account_id.as_str()) + }) + { + return Err(io::Error::other( + "account changed during workspace routing discovery", + )); + } + Ok(routing) + } +} diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index e33e587121..6661a78d8e 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -2,6 +2,10 @@ pub mod auth; pub mod auth_env_telemetry; pub mod test_support; pub mod token_data; +pub use auth::WorkspaceRouting; +pub use auth::WorkspaceRoutingRequest; +pub use auth::WorkspaceRoutingResolver; +pub use auth::WorkspaceRoutingSession; mod callback_params; mod device_code_auth; diff --git a/codex-rs/memories/write/src/runtime.rs b/codex-rs/memories/write/src/runtime.rs index d9079c48bb..3efdfb21c5 100644 --- a/codex-rs/memories/write/src/runtime.rs +++ b/codex-rs/memories/write/src/runtime.rs @@ -304,6 +304,7 @@ impl MemoryStartupContext { /*concurrent_reasoning_summaries_enabled*/ false, /*attestation_provider*/ None, config.http_client_factory(), + config.workspace_routing_context(), ); let mut client_session = model_client.new_session(); diff --git a/codex-rs/model-provider/Cargo.toml b/codex-rs/model-provider/Cargo.toml index 9ba32d881a..c9ad5683da 100644 --- a/codex-rs/model-provider/Cargo.toml +++ b/codex-rs/model-provider/Cargo.toml @@ -31,6 +31,7 @@ serde_json = { workspace = true } sha2 = { workspace = true } tokio = { workspace = true, features = ["io-util", "process", "sync", "time"] } tracing = { workspace = true, features = ["log"] } +url = { workspace = true } [dev-dependencies] base64 = { workspace = true } diff --git a/codex-rs/model-provider/src/lib.rs b/codex-rs/model-provider/src/lib.rs index b3ca480d8c..0abe111a40 100644 --- a/codex-rs/model-provider/src/lib.rs +++ b/codex-rs/model-provider/src/lib.rs @@ -6,6 +6,11 @@ mod models_identity; mod provider; mod shared_state; pub mod test_support; +mod workspace_routing; +pub use workspace_routing::ACCOUNT_ROUTING_HEADER; +pub use workspace_routing::ResolvedResponsesProvider; +pub use workspace_routing::ResponsesConnectionKey; +pub use workspace_routing::WorkspaceRoutingContext; pub use amazon_bedrock::is_supported_amazon_bedrock_region; pub use auth::AgentIdentitySessionFallback; @@ -31,3 +36,7 @@ pub use provider::ProviderUnauthorizedRecovery; pub use provider::RemoteCompactionSupport; pub use provider::SharedModelProvider; pub use provider::create_model_provider; + +#[cfg(test)] +#[path = "workspace_routing_tests.rs"] +mod workspace_routing_tests; diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index 6239bf0cb5..58b51d1bb3 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -11,6 +11,8 @@ use codex_api::TransportError; use codex_api::is_azure_responses_provider; use codex_login::AuthManager; use codex_login::CodexAuth; +use codex_login::WorkspaceRoutingRequest; +use codex_login::default_client::ClientRedirectPolicy; use codex_login::default_client::RESIDENCY_HEADER_NAME; use codex_login::default_client::ResidencyRequirement; use codex_login::default_client::read_default_client_residency_requirement; @@ -24,6 +26,7 @@ use codex_protocol::error::CodexErr; use codex_protocol::openai_models::ModelsResponse; use http::HeaderValue; +use crate::ResolvedResponsesProvider; use crate::amazon_bedrock::AmazonBedrockModelProvider; use crate::auth::ProviderAuthScope; use crate::auth::ResolvedProviderAuth; @@ -31,6 +34,7 @@ use crate::auth::auth_manager_for_provider; use crate::auth::resolve_provider_auth; use crate::auth::resolve_provider_auth_for_scope; use crate::models_endpoint::OpenAiModelsEndpoint; +use crate::workspace_routing::WorkspaceRoutingContext; pub(crate) fn enforce_managed_residency(provider: &mut Provider) { if let Some(requirement) = read_default_client_residency_requirement() { @@ -234,6 +238,48 @@ pub trait ModelProvider: fmt::Debug + Send + Sync { }) } + /// Resolves routing for Responses HTTP, compaction, and WebSocket handshakes. + #[expect( + clippy::await_holding_invalid_type, + reason = "serialize discovery and the session's first successful routing transition" + )] + fn responses_api_provider<'a>( + &'a self, + routing_context: &'a WorkspaceRoutingContext, + ) -> ModelProviderFuture<'a, codex_protocol::error::Result> { + Box::pin(async move { + let mut provider = self.api_provider().await?; + let mut redirect_policy = ClientRedirectPolicy::Default; + if provider_uses_first_party_auth_path(self.info()) + && self.info().supports_codex_backend_routes() + && let Some(auth) = self.auth().await.filter(CodexAuth::is_chatgpt_auth) + && let Some(auth_manager) = self.auth_manager() + { + let mut previously_routed = routing_context.previously_routed.lock().await; + if let Some(routing) = auth_manager + .workspace_routing( + &auth, + WorkspaceRoutingRequest { + provider_base_url: provider.base_url.clone(), + chatgpt_base_url: routing_context.chatgpt_base_url.clone(), + previously_routed: *previously_routed, + session: routing_context.session.clone(), + }, + ) + .await? + { + crate::workspace_routing::apply_workspace_routing(&mut provider, routing)?; + redirect_policy = ClientRedirectPolicy::Reject; + *previously_routed = true; + } + } + Ok(ResolvedResponsesProvider { + provider, + redirect_policy, + }) + }) + } + /// Returns the provider base URL that will be used at request time. fn runtime_base_url( &self, diff --git a/codex-rs/model-provider/src/workspace_routing.rs b/codex-rs/model-provider/src/workspace_routing.rs new file mode 100644 index 0000000000..d057a665b4 --- /dev/null +++ b/codex-rs/model-provider/src/workspace_routing.rs @@ -0,0 +1,103 @@ +//! Applies discovered routing and keeps requests within the selected ChatGPT backend. + +use codex_login::WorkspaceRouting; +use codex_login::WorkspaceRoutingSession; +use codex_login::default_client::ClientRedirectPolicy; +use http::HeaderValue; +use std::io; +use std::sync::Arc; +use tokio::sync::Mutex; +use url::Url; + +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)] +pub struct WorkspaceRoutingContext { + pub(crate) chatgpt_base_url: String, + pub(crate) previously_routed: Mutex, + pub(crate) session: Option>, +} + +impl WorkspaceRoutingContext { + pub fn new(chatgpt_base_url: String) -> Self { + Self { + chatgpt_base_url, + previously_routed: Mutex::new(/*t*/ false), + session: None, + } + } + + pub fn with_session(mut self, session: WorkspaceRoutingSession) -> Self { + self.session = Some(Arc::new(session)); + self + } +} + +/// Changes only the origin of requests to the selected ChatGPT backend. +pub(crate) fn apply_workspace_routing( + provider: &mut codex_api::Provider, + routing: WorkspaceRouting, +) -> io::Result<()> { + let base_url = &mut provider.base_url; + let headers = &mut provider.headers; + let parse = |value: &str| Url::parse(value).map_err(io::Error::other); + let mut url = parse(base_url)?; + let backend = parse(&routing.backend_origin)?; + if backend.scheme() != "https" + || backend.host_str().is_none() + || !backend.username().is_empty() + || backend.password().is_some() + || backend.path() != "/" + || backend.query().is_some() + || backend.fragment().is_some() + { + return Err(io::Error::other("invalid workspace backend origin")); + } + let header = match routing.account_routing_override.as_str() { + "NO_CONSTRAINT" => None, + "us" | "us_cr" => Some( + HeaderValue::from_str(&routing.account_routing_override).map_err(io::Error::other)?, + ), + _ => return Err(io::Error::other("invalid workspace routing override")), + }; + url.set_scheme(backend.scheme()) + .map_err(|()| io::Error::other("invalid backend scheme"))?; + url.set_host(backend.host_str()).map_err(io::Error::other)?; + url.set_port(backend.port()) + .map_err(|()| io::Error::other("invalid backend port"))?; + *base_url = url.into(); + headers.remove(ACCOUNT_ROUTING_HEADER); + if let Some(header) = header { + headers.insert(ACCOUNT_ROUTING_HEADER, header); + } + Ok(()) +} + +/// API deployment and redirect policy resolved for one Responses request. +/// Workspace routes reject redirects even when their routing override supplies no header. +#[derive(Debug)] +pub struct ResolvedResponsesProvider { + pub provider: codex_api::Provider, + pub redirect_policy: ClientRedirectPolicy, +} + +/// Effective destination and credentials that permit reusing a Responses socket. +#[derive(Debug, PartialEq, Eq)] +pub struct ResponsesConnectionKey { + base_url: String, + routing_header: Option, + auth_revision: Option, +} + +impl ResponsesConnectionKey { + pub fn new(provider: &codex_api::Provider, auth_revision: Option) -> Self { + Self { + base_url: provider.base_url.clone(), + routing_header: provider.headers.get(ACCOUNT_ROUTING_HEADER).cloned(), + auth_revision, + } + } +} diff --git a/codex-rs/model-provider/src/workspace_routing_tests.rs b/codex-rs/model-provider/src/workspace_routing_tests.rs new file mode 100644 index 0000000000..95f20eea52 --- /dev/null +++ b/codex-rs/model-provider/src/workspace_routing_tests.rs @@ -0,0 +1,199 @@ +//! Routing is limited to the selected ChatGPT backend and successful discovery. + +use std::future::Future; +use std::io; +use std::pin::Pin; +use std::sync::Arc; + +use crate::ACCOUNT_ROUTING_HEADER; +use crate::WorkspaceRoutingContext; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_login::WorkspaceRouting; +use codex_login::WorkspaceRoutingRequest; +use codex_login::WorkspaceRoutingResolver; +use codex_login::default_client::ClientRedirectPolicy; +use codex_model_provider_info::ModelProviderInfo; +use pretty_assertions::assert_eq; + +use crate::create_model_provider; + +struct Routing(WorkspaceRouting); + +struct ChangedBootstrap { + first_lookup: std::sync::atomic::AtomicBool, + release_first: tokio::sync::Notify, +} + +impl WorkspaceRoutingResolver for ChangedBootstrap { + fn resolve( + &self, + request: WorkspaceRoutingRequest, + ) -> Pin>> + Send + '_>> { + Box::pin(async move { + if self + .first_lookup + .swap(false, std::sync::atomic::Ordering::Relaxed) + { + self.release_first.notified().await; + Ok(Some(WorkspaceRouting { + chatgpt_account_id: "account_id".into(), + backend_origin: "https://gov.chatgpt.com".into(), + account_routing_override: "us_cr".into(), + })) + } else if request.previously_routed { + Err(io::Error::other("bootstrap changed")) + } else { + Ok(None) + } + }) + } +} + +#[tokio::test] +async fn concurrent_discovery_cannot_forget_established_routing() { + let context = WorkspaceRoutingContext::new("https://chatgpt.com/backend-api".into()); + let auth = + AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let resolver = Arc::new(ChangedBootstrap { + first_lookup: std::sync::atomic::AtomicBool::new(true), + release_first: tokio::sync::Notify::new(), + }); + auth.set_workspace_routing_resolver(Arc::downgrade( + &(resolver.clone() as Arc), + )); + let provider = create_model_provider( + ModelProviderInfo::create_openai_provider(/*base_url*/ None), + Some(auth), + ); + let mut first = provider.responses_api_provider(&context); + let mut second = provider.responses_api_provider(&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)) + }) + .await; + resolver.release_first.notify_one(); + assert_eq!( + first.await.unwrap().provider.base_url, + "https://gov.chatgpt.com/backend-api/codex" + ); + let second = match early_second { + std::task::Poll::Ready(result) => result, + std::task::Poll::Pending => second.await, + }; + assert!( + second.is_err(), + "a concurrent lookup must retain the first lookup's routing scope" + ); +} + +impl WorkspaceRoutingResolver for Routing { + fn resolve( + &self, + request: WorkspaceRoutingRequest, + ) -> Pin>> + Send + '_>> { + Box::pin(async move { + let origin = url::Url::parse(&request.provider_base_url) + .unwrap() + .origin(); + let bootstrap = url::Url::parse(&request.chatgpt_base_url).unwrap().origin(); + Ok((origin == bootstrap).then(|| self.0.clone())) + }) + } +} + +#[tokio::test] +async fn workspace_routing_preserves_paths_and_excludes_other_providers() { + let routing_context = WorkspaceRoutingContext::new("https://chatgpt.com/backend-api".into()); + let auth = + AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let resolver: Arc = Arc::new(Routing(WorkspaceRouting { + chatgpt_account_id: "account_id".into(), + backend_origin: "https://gov.chatgpt.com:8443".into(), + account_routing_override: "us_cr".into(), + })); + auth.set_workspace_routing_resolver(Arc::downgrade(&resolver)); + let info = ModelProviderInfo::create_openai_provider(/*base_url*/ None); + let provider = create_model_provider(info.clone(), Some(auth.clone())); + let routed = provider + .responses_api_provider(&routing_context) + .await + .unwrap(); + assert_eq!( + ( + routed.provider.base_url.as_str(), + routed.provider.headers[ACCOUNT_ROUTING_HEADER] + .to_str() + .unwrap(), + routed.redirect_policy, + ), + ( + "https://gov.chatgpt.com:8443/backend-api/codex", + "us_cr", + ClientRedirectPolicy::Reject + ) + ); + let mut custom_auth = info.clone(); + custom_auth.experimental_bearer_token = Some("custom-token".into()); + let mut unrelated_path = info.clone(); + unrelated_path.base_url = Some("https://chatgpt.com/v1".into()); + for info in [custom_auth, unrelated_path] { + let provider = create_model_provider(info, Some(auth.clone())); + let ordinary = provider.api_provider().await.unwrap(); + let responses = provider + .responses_api_provider(&routing_context) + .await + .unwrap(); + assert_eq!( + ( + responses.provider.base_url, + responses.provider.headers, + responses.redirect_policy + ), + ( + ordinary.base_url, + ordinary.headers, + ClientRedirectPolicy::Default + ) + ); + } + let api_auth = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("api-key")); + api_auth.set_workspace_routing_resolver(Arc::downgrade(&resolver)); + let api = create_model_provider(info, Some(api_auth)) + .responses_api_provider(&routing_context) + .await + .unwrap(); + assert_eq!(api.provider.base_url, "https://api.openai.com/v1"); + assert!(!api.provider.headers.contains_key(ACCOUNT_ROUTING_HEADER)); +} + +#[tokio::test] +async fn missing_routing_owner_and_wrong_workspace_fail_closed() { + let routing_context = WorkspaceRoutingContext::new("https://chatgpt.com/backend-api".into()); + let auth = + AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let resolver: Arc = Arc::new(Routing(WorkspaceRouting { + chatgpt_account_id: "other-account".into(), + backend_origin: "https://chatgpt.com".into(), + account_routing_override: "NO_CONSTRAINT".into(), + })); + auth.set_workspace_routing_resolver(Arc::downgrade(&resolver)); + let provider = create_model_provider( + ModelProviderInfo::create_openai_provider(/*base_url*/ None), + Some(auth), + ); + assert!( + provider + .responses_api_provider(&routing_context) + .await + .is_err() + ); + drop(resolver); + assert!( + provider + .responses_api_provider(&routing_context) + .await + .is_err() + ); +}