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() + ); +}