diff --git a/codex-rs/app-server-client/src/lib.rs b/codex-rs/app-server-client/src/lib.rs index c5840e4d4c..588e5b1566 100644 --- a/codex-rs/app-server-client/src/lib.rs +++ b/codex-rs/app-server-client/src/lib.rs @@ -1354,6 +1354,7 @@ mod tests { JSONRPCMessage::Response(JSONRPCResponse { id: request.id, result: serde_json::to_value(GetAccountResponse { + workspace_routing: None, account: None, requires_openai_auth: false, }) @@ -1408,6 +1409,7 @@ mod tests { JSONRPCMessage::Response(JSONRPCResponse { id: request.id, result: serde_json::to_value(GetAccountResponse { + workspace_routing: None, account: None, requires_openai_auth: false, }) @@ -1484,6 +1486,7 @@ mod tests { assert_eq!( response, GetAccountResponse { + workspace_routing: None, account: None, requires_openai_auth: false, } @@ -1587,6 +1590,7 @@ mod tests { JSONRPCMessage::Response(JSONRPCResponse { id: request.id, result: serde_json::to_value(GetAccountResponse { + workspace_routing: None, account: None, requires_openai_auth: false, }) @@ -1640,6 +1644,7 @@ mod tests { assert_eq!( first_response, GetAccountResponse { + workspace_routing: None, account: None, requires_openai_auth: false, } diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index fec88ff951..ac76ffbb2d 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -6705,6 +6705,15 @@ "title": "AccountRateLimitsUpdatedNotification", "type": "object" }, + "AccountRoutingOverride": { + "description": "Backend routing policy. Wire values match the accounts/check contract.", + "enum": [ + "NO_CONSTRAINT", + "us", + "us_cr" + ], + "type": "string" + }, "AccountTokenUsageDailyBucket": { "properties": { "startDate": { @@ -25918,6 +25927,25 @@ ], "type": "string" }, + "WorkspaceRouting": { + "properties": { + "accountRoutingOverride": { + "$ref": "#/definitions/v2/AccountRoutingOverride" + }, + "backendOrigin": { + "type": "string" + }, + "chatgptAccountId": { + "type": "string" + } + }, + "required": [ + "accountRoutingOverride", + "backendOrigin", + "chatgptAccountId" + ], + "type": "object" + }, "WriteStatus": { "enum": [ "ok", diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json index fcb6755496..7a95a06a2a 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json @@ -121,6 +121,15 @@ "title": "AccountRateLimitsUpdatedNotification", "type": "object" }, + "AccountRoutingOverride": { + "description": "Backend routing policy. Wire values match the accounts/check contract.", + "enum": [ + "NO_CONSTRAINT", + "us", + "us_cr" + ], + "type": "string" + }, "AccountTokenUsageDailyBucket": { "properties": { "startDate": { @@ -23632,6 +23641,25 @@ ], "type": "string" }, + "WorkspaceRouting": { + "properties": { + "accountRoutingOverride": { + "$ref": "#/definitions/AccountRoutingOverride" + }, + "backendOrigin": { + "type": "string" + }, + "chatgptAccountId": { + "type": "string" + } + }, + "required": [ + "accountRoutingOverride", + "backendOrigin", + "chatgptAccountId" + ], + "type": "object" + }, "WriteStatus": { "enum": [ "ok", diff --git a/codex-rs/app-server-protocol/schema/json/v2/GetAccountResponse.json b/codex-rs/app-server-protocol/schema/json/v2/GetAccountResponse.json index e5537a3e4f..4613d90f4f 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/GetAccountResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/GetAccountResponse.json @@ -68,6 +68,15 @@ } ] }, + "AccountRoutingOverride": { + "description": "Backend routing policy. Wire values match the accounts/check contract.", + "enum": [ + "NO_CONSTRAINT", + "us", + "us_cr" + ], + "type": "string" + }, "PlanType": { "enum": [ "free", @@ -89,6 +98,25 @@ "unknown" ], "type": "string" + }, + "WorkspaceRouting": { + "properties": { + "accountRoutingOverride": { + "$ref": "#/definitions/AccountRoutingOverride" + }, + "backendOrigin": { + "type": "string" + }, + "chatgptAccountId": { + "type": "string" + } + }, + "required": [ + "accountRoutingOverride", + "backendOrigin", + "chatgptAccountId" + ], + "type": "object" } }, "properties": { diff --git a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-experimental.json.zst b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-experimental.json.zst index 7d613c2a4f..b4f9045ce0 100644 Binary files a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-experimental.json.zst and b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-experimental.json.zst differ diff --git a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst index 6fd1937421..8222eefe7f 100644 Binary files a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst and b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst differ diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/AccountRoutingOverride.ts b/codex-rs/app-server-protocol/schema/typescript/v2/AccountRoutingOverride.ts new file mode 100644 index 0000000000..30b8c8e5bc --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/AccountRoutingOverride.ts @@ -0,0 +1,8 @@ +// GENERATED CODE! DO NOT MODIFY BY HAND! + +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. + +/** + * Backend routing policy. Wire values match the accounts/check contract. + */ +export type AccountRoutingOverride = "NO_CONSTRAINT" | "us" | "us_cr"; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/GetAccountResponse.ts b/codex-rs/app-server-protocol/schema/typescript/v2/GetAccountResponse.ts index 83da4f4e5e..e88cc6abe3 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/GetAccountResponse.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/GetAccountResponse.ts @@ -3,4 +3,4 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. import type { Account } from "./Account"; -export type GetAccountResponse = { account: Account | null, requiresOpenaiAuth: boolean, }; +export type GetAccountResponse = {account: Account | null, requiresOpenaiAuth: boolean}; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/WorkspaceRouting.ts b/codex-rs/app-server-protocol/schema/typescript/v2/WorkspaceRouting.ts new file mode 100644 index 0000000000..f40eec6e33 --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/WorkspaceRouting.ts @@ -0,0 +1,6 @@ +// GENERATED CODE! DO NOT MODIFY BY HAND! + +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { AccountRoutingOverride } from "./AccountRoutingOverride"; + +export type WorkspaceRouting = { chatgptAccountId: string, backendOrigin: string, accountRoutingOverride: AccountRoutingOverride, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/index.ts b/codex-rs/app-server-protocol/schema/typescript/v2/index.ts index 242c9fc3e2..e5b2a0c7a3 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/index.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/index.ts @@ -3,6 +3,7 @@ export type { Account } from "./Account"; export type { AccountLoginCompletedNotification } from "./AccountLoginCompletedNotification"; export type { AccountRateLimitsUpdatedNotification } from "./AccountRateLimitsUpdatedNotification"; +export type { AccountRoutingOverride } from "./AccountRoutingOverride"; export type { AccountTokenUsageDailyBucket } from "./AccountTokenUsageDailyBucket"; export type { AccountTokenUsageSummary } from "./AccountTokenUsageSummary"; export type { AccountUpdatedNotification } from "./AccountUpdatedNotification"; @@ -624,4 +625,5 @@ export type { WindowsSandboxSetupStartResponse } from "./WindowsSandboxSetupStar export type { WindowsWorldWritableWarningNotification } from "./WindowsWorldWritableWarningNotification"; export type { WorkspaceMessage } from "./WorkspaceMessage"; export type { WorkspaceMessageType } from "./WorkspaceMessageType"; +export type { WorkspaceRouting } from "./WorkspaceRouting"; export type { WriteStatus } from "./WriteStatus"; diff --git a/codex-rs/app-server-protocol/src/protocol/v2/account.rs b/codex-rs/app-server-protocol/src/protocol/v2/account.rs index 40965ed3be..35c2b6eb8a 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/account.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/account.rs @@ -551,12 +551,35 @@ pub struct GetAccountParams { pub refresh_token: bool, } -#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS, ExperimentalApi)] #[serde(rename_all = "camelCase")] #[ts(export_to = "v2/")] pub struct GetAccountResponse { pub account: Option, pub requires_openai_auth: bool, + #[experimental("account/read.workspaceRouting")] + pub workspace_routing: Option, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct WorkspaceRouting { + pub chatgpt_account_id: String, + pub backend_origin: String, + pub account_routing_override: AccountRoutingOverride, +} + +/// Backend routing policy. Wire values match the accounts/check contract. +#[derive(Serialize, Deserialize, Debug, Clone, Copy, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "snake_case")] +#[ts(export_to = "v2/", rename_all = "snake_case")] +pub enum AccountRoutingOverride { + #[serde(rename = "NO_CONSTRAINT")] + #[ts(rename = "NO_CONSTRAINT")] + NoConstraint, + Us, + UsCr, } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 6a15fa2dd8..358ae55943 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -235,3 +235,11 @@ Existing rollouts may contain historical `ThreadRolledBack` events. Their replay and migration remain supported so resuming, reading, and forking those threads preserves the surviving history. This disk compatibility does not require restoring support for new `thread/rollback` requests. + +# Selected workspace routing + +The experimental `account/read.workspaceRouting` response field returns the selected ChatGPT workspace's `chatgptAccountId`, resolved HTTPS `backendOrigin`, and backend-provided `accountRoutingOverride`. The routing value is `us`, `us_cr`, or the explicit `NO_CONSTRAINT` value. API-only and signed-out accounts return `null` and do not need `accounts/check`. + +App-server discovers routing for saved ChatGPT logins at startup and for new logins or workspace switches. After requirements and routing are ready, it sends the existing `account/updated` notification. Newly initialized connections also receive this notification once saved-workspace routing is ready, including when discovery finished before the connection initialized. Clients then reread `configRequirements/read` and `account/read`. Saved ChatGPT credentials without a selected workspace ID retain their account information and return `workspaceRouting: null`; app-server does not guess a workspace from the backend's default account. Discovery failures for a selected workspace, including missing or null fields from older backends, return an `account/read` error. They never produce a successful unrestricted result. A later read retries failed discovery. Logout clears the cached routing, and results from earlier authentication owners are discarded. Token refreshes for the same known user and workspace invalidate cached routing without cancelling discovery or failing sign-in. Configuration is reloaded after discovery; a changed backend, model provider, or required backend rejects the result so the next read discovers against current configuration. Account notifications recheck the auth owner generation after waiting for outbound queue capacity. Superseded sign-in attempts emit a failed `account/login/completed` event instead of silently dropping completion. Notifications remain snapshots: clients reread current account and requirements state rather than treating a queued notification as authorization. + +The origin of a required `chatgpt_base_url` must match the discovered origin by scheme, host, and effective port. The base URL's API path is not part of this comparison. Either origin alone is sufficient. If requirements specify no base URL and discovery explicitly returns `NO_CONSTRAINT`, the effective `chatgpt_base_url` supplies the origin, including its existing default. `backendOrigin` is always a resolved origin; `accountRoutingOverride` preserves `NO_CONSTRAINT` when the backend explicitly returns it. Discovering an origin does not change API paths or apply routing headers to requests. diff --git a/codex-rs/app-server/src/account_notifications.rs b/codex-rs/app-server/src/account_notifications.rs new file mode 100644 index 0000000000..7e64caa226 --- /dev/null +++ b/codex-rs/app-server/src/account_notifications.rs @@ -0,0 +1,57 @@ +//! Publishes account snapshots only while their authentication owner is current. +//! Superseded login attempts still receive a terminal completion notification. + +use super::*; +use codex_app_server_protocol::AccountLoginCompletedNotification; +use codex_app_server_protocol::AccountUpdatedNotification; +use codex_login::AuthChangeState; +use tokio::sync::watch; + +pub(crate) enum AccountNotification { + LoginCompleted(AccountLoginCompletedNotification), + Updated(AccountUpdatedNotification), +} + +impl OutgoingMessageSender { + pub(crate) async fn send_account_notification( + &self, + connection_id: Option, + auth_changes: &watch::Receiver, + owner_generation: u64, + notification: AccountNotification, + ) { + let Ok(permit) = self.sender.reserve().await else { + return; + }; + // Keep the revision read lock through enqueue; never wait while holding it. + let current_auth = auth_changes.borrow(); + let notification = match notification { + AccountNotification::LoginCompleted(mut payload) => { + if payload.success && current_auth.owner_generation != owner_generation { + payload.success = false; + payload.error = Some("account changed before sign-in completed".into()); + } + ServerNotification::AccountLoginCompleted(payload) + } + AccountNotification::Updated(payload) => { + if current_auth.owner_generation != owner_generation { + return; + } + ServerNotification::AccountUpdated(payload) + } + }; + let message = timestamped_server_notification(notification); + permit.send(match connection_id { + Some(connection_id) => OutgoingEnvelope::ToConnection { + connection_id, + message, + write_complete_tx: None, + }, + None => OutgoingEnvelope::Broadcast { message }, + }); + } +} + +#[cfg(test)] +#[path = "account_notifications_tests.rs"] +mod tests; diff --git a/codex-rs/app-server/src/account_notifications_tests.rs b/codex-rs/app-server/src/account_notifications_tests.rs new file mode 100644 index 0000000000..fba40deacc --- /dev/null +++ b/codex-rs/app-server/src/account_notifications_tests.rs @@ -0,0 +1,99 @@ +//! Exercises auth changes while notification delivery waits for queue capacity. + +use super::*; +use pretty_assertions::assert_eq; +use test_case::test_case; + +#[test_case(0; "same_owner_refresh")] +#[test_case(1; "owner_change")] +#[test_case(2; "coalesced_owner_changes")] +#[tokio::test] +async fn queued_notifications_follow_auth_owner_changes(owner_generation: u64) { + for login in [false, true] { + let (tx, mut rx) = mpsc::channel(/*buffer*/ 1); + let outgoing = OutgoingMessageSender::new( + tx.clone(), + codex_analytics::AnalyticsEventsClient::disabled(), + ); + let (changes, auth_changes) = watch::channel(AuthChangeState::default()); + let updated = AccountUpdatedNotification { + auth_mode: None, + plan_type: None, + }; + tx.send(OutgoingEnvelope::Broadcast { + message: timestamped_server_notification(ServerNotification::AccountUpdated( + updated.clone(), + )), + }) + .await + .unwrap(); + let payload = AccountLoginCompletedNotification { + login_id: Some("queued-login".into()), + success: true, + error: None, + onboarding_entrypoint: None, + }; + let notification = if login { + AccountNotification::LoginCompleted(payload.clone()) + } else { + AccountNotification::Updated(updated.clone()) + }; + let send = outgoing.send_account_notification( + Some(ConnectionId(7)), + &auth_changes, + /*owner_generation*/ 0, + notification, + ); + tokio::pin!(send); + // Poll delivery with a full queue before changing authentication. + assert!(futures::poll!(&mut send).is_pending()); + changes.send_replace(AuthChangeState { + generation: 2, + owner_generation, + }); + rx.recv().await.unwrap(); + send.await; + if login { + let OutgoingEnvelope::ToConnection { + connection_id, + message: OutgoingMessage::AppServerNotification(envelope), + .. + } = rx.try_recv().unwrap() + else { + panic!("expected targeted login completion"); + }; + assert_eq!(connection_id, ConnectionId(7)); + let ServerNotification::AccountLoginCompleted(actual) = envelope.notification else { + panic!("expected login completion"); + }; + assert_eq!( + actual, + if owner_generation == 0 { + payload + } else { + AccountLoginCompletedNotification { + success: false, + error: Some("account changed before sign-in completed".into()), + ..payload + } + } + ); + } else if owner_generation == 0 { + let OutgoingEnvelope::ToConnection { + connection_id, + message: OutgoingMessage::AppServerNotification(envelope), + .. + } = rx.try_recv().unwrap() + else { + panic!("expected targeted account update"); + }; + assert_eq!(connection_id, ConnectionId(7)); + let ServerNotification::AccountUpdated(actual) = envelope.notification else { + panic!("expected account update"); + }; + assert_eq!(actual, updated); + } else { + assert!(rx.try_recv().is_err()); + } + } +} diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index f08d34ee44..9c7c11281f 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -803,6 +803,8 @@ impl MessageProcessor { connection_id: ConnectionId, request_attestation: bool, ) { + self.account_processor + .notify_workspace_routing_to_connection(connection_id); self.thread_processor .connection_initialized( connection_id, @@ -932,13 +934,7 @@ impl MessageProcessor { ) .await?; if connection_initialized { - self.thread_processor - .connection_initialized( - connection_id, - ConnectionCapabilities { - request_attestation: session.request_attestation(), - }, - ) + self.connection_initialized(connection_id, session.request_attestation()) .await; } return Ok(()); diff --git a/codex-rs/app-server/src/outgoing_message.rs b/codex-rs/app-server/src/outgoing_message.rs index e6803f0b4d..e6a109ceb3 100644 --- a/codex-rs/app-server/src/outgoing_message.rs +++ b/codex-rs/app-server/src/outgoing_message.rs @@ -46,6 +46,10 @@ pub(crate) type ClientRequestResult = std::result::Result, config_manager: ConfigManager, active_login: Arc>>, + workspace_routing: Arc>>, + workspace_routing_fetch: Arc, + workspace_routing_shutdown: CancellationToken, } impl AccountRequestProcessor { @@ -101,14 +106,26 @@ impl AccountRequestProcessor { config: Arc, config_manager: ConfigManager, ) -> Self { - Self { + let processor = Self { auth_manager, thread_manager, outgoing, config, config_manager, active_login: Arc::new(Mutex::new(None)), - } + workspace_routing: Arc::new(Mutex::new(None)), + workspace_routing_fetch: Arc::new(Semaphore::new(/*permits*/ 1)), + workspace_routing_shutdown: CancellationToken::new(), + }; + let startup = processor.clone(); + tokio::spawn(async move { + let _ = startup + .get_account_response(GetAccountParams { + refresh_token: false, + }) + .await; + }); + processor } pub(crate) async fn login_account( @@ -196,6 +213,7 @@ impl AccountRequestProcessor { } pub(crate) fn clear_external_auth(&self) { + self.workspace_routing_shutdown.cancel(); self.auth_manager.clear_external_auth(); } @@ -638,10 +656,7 @@ impl AccountRequestProcessor { }); } - let outgoing_clone = self.outgoing.clone(); - let config_manager = self.config_manager.clone(); - let thread_manager = Arc::clone(&self.thread_manager); - let config = Arc::clone(&self.config); + let processor = self.clone(); let active_login = self.active_login.clone(); let auth_url = server.auth_url.clone(); tokio::spawn(async move { @@ -667,19 +682,14 @@ impl AccountRequestProcessor { } }; - Self::send_chatgpt_login_completion_notifications( - &outgoing_clone, - config_manager, - thread_manager, - config, - AccountLoginCompletedNotification { + processor + .send_chatgpt_login_completion_notifications(AccountLoginCompletedNotification { login_id: Some(login_id.to_string()), success, error: error_msg, onboarding_entrypoint, - }, - ) - .await; + }) + .await; // Clear the active login if it matches this attempt. It may have been replaced or cancelled. let mut guard = active_login.lock().await; @@ -728,10 +738,7 @@ impl AccountRequestProcessor { let verification_url = device_code.verification_url.clone(); let user_code = device_code.user_code.clone(); - let outgoing_clone = self.outgoing.clone(); - let config_manager = self.config_manager.clone(); - let thread_manager = Arc::clone(&self.thread_manager); - let config = Arc::clone(&self.config); + let processor = self.clone(); let active_login = self.active_login.clone(); tokio::spawn(async move { let (success, error_msg) = tokio::select! { @@ -746,19 +753,14 @@ impl AccountRequestProcessor { } }; - Self::send_chatgpt_login_completion_notifications( - &outgoing_clone, - config_manager, - thread_manager, - config, - AccountLoginCompletedNotification { + processor + .send_chatgpt_login_completion_notifications(AccountLoginCompletedNotification { login_id: Some(login_id.to_string()), success, error: error_msg, onboarding_entrypoint: None, - }, - ) - .await; + }) + .await; let mut guard = active_login.lock().await; if guard.as_ref().map(ActiveLogin::login_id) == Some(login_id) { @@ -878,76 +880,87 @@ impl AccountRequestProcessor { } async fn send_login_success_notifications(&self, login_id: Option) { - Self::maybe_refresh_plugin_caches_for_current_config( - &self.config_manager, - &self.thread_manager, - self.auth_manager.auth_cached(), - ) - .await; - - let payload_login_completed = AccountLoginCompletedNotification { + self.send_account_login_notifications(AccountLoginCompletedNotification { login_id: login_id.map(|id| id.to_string()), success: true, error: None, onboarding_entrypoint: None, - }; - self.outgoing - .send_server_notification(ServerNotification::AccountLoginCompleted( - payload_login_completed, - )) - .await; - - self.outgoing - .send_server_notification(ServerNotification::AccountUpdated( - self.current_account_updated_notification(), - )) - .await; + }) + .await; } - async fn send_chatgpt_login_completion_notifications( - outgoing: &OutgoingMessageSender, - config_manager: ConfigManager, - thread_manager: Arc, - config: Arc, - payload_v2: AccountLoginCompletedNotification, + async fn send_account_login_notifications( + &self, + mut payload: AccountLoginCompletedNotification, ) { - let success = payload_v2.success; - outgoing - .send_server_notification(ServerNotification::AccountLoginCompleted(payload_v2)) + let auth_changes = self.auth_manager.auth_change_state_receiver(); + let owner_generation = auth_changes.borrow().owner_generation; + if payload.success + && let Err(error) = self + .get_account_response(GetAccountParams { + refresh_token: false, + }) + .await + { + payload.success = false; + payload.error = Some(error.message); + } + if payload.success && auth_changes.borrow().owner_generation == owner_generation { + Self::maybe_refresh_plugin_caches_for_current_config( + &self.config_manager, + &self.thread_manager, + self.auth_manager.auth_cached(), + ) + .await; + } + + let success = payload.success; + self.outgoing + .send_account_notification( + /*connection_id*/ None, + &auth_changes, + owner_generation, + AccountNotification::LoginCompleted(payload), + ) .await; if success { - let auth_manager = thread_manager.auth_manager(); - auth_manager.reload().await; - config_manager.replace_cloud_config_bundle_loader( - auth_manager.clone(), - config.chatgpt_base_url.clone(), - config.http_client_factory(), - ); - config_manager - .sync_default_client_residency_requirement() - .await; - - let auth = auth_manager.auth_cached(); - Self::maybe_refresh_plugin_caches_for_current_config( - &config_manager, - &thread_manager, - auth.clone(), - ) - .await; - let payload_v2 = AccountUpdatedNotification { - auth_mode: auth - .as_ref() - .map(CodexAuth::api_auth_mode) - .map(auth_mode_to_api), - plan_type: auth.as_ref().and_then(CodexAuth::account_plan_type), - }; - outgoing - .send_server_notification(ServerNotification::AccountUpdated(payload_v2)) + let notification = self.current_account_updated_notification(); + self.outgoing + .send_account_notification( + /*connection_id*/ None, + &auth_changes, + owner_generation, + AccountNotification::Updated(notification), + ) .await; } } + async fn send_chatgpt_login_completion_notifications( + &self, + mut payload_v2: AccountLoginCompletedNotification, + ) { + if payload_v2.success { + self.auth_manager.reload().await; + let auth_changes = self.auth_manager.auth_change_state_receiver(); + let owner_generation = auth_changes.borrow().owner_generation; + self.config_manager.replace_cloud_config_bundle_loader( + self.auth_manager.clone(), + self.config.chatgpt_base_url.clone(), + self.config.http_client_factory(), + ); + self.config_manager + .sync_default_client_residency_requirement() + .await; + if auth_changes.borrow().owner_generation != owner_generation { + payload_v2.success = false; + payload_v2.error = Some("account changed before sign-in completed".into()); + } + } + self.send_account_login_notifications(payload_v2).await; + } + async fn logout_common(&self) -> std::result::Result, JSONRPCErrorError> { if self.auth_manager.is_workload_identity_selected() { return Err(self.configured_auth_owned_by_host_error()); @@ -974,6 +987,7 @@ impl AccountRequestProcessor { } self.config_manager.clear_cloud_config_bundle_loader(); + *self.workspace_routing.lock().await = None; Self::maybe_refresh_plugin_caches_for_current_config( &self.config_manager, @@ -1103,29 +1117,6 @@ impl AccountRequestProcessor { Ok(response) } - async fn get_account_response( - &self, - params: GetAccountParams, - ) -> Result { - let do_refresh = params.refresh_token; - - self.refresh_token_if_requested(do_refresh).await; - - let config = self.load_latest_config().await; - let provider = - create_model_provider(config.model_provider, Some(self.auth_manager.clone())); - let account_state = match provider.account_state() { - Ok(account_state) => account_state, - Err(err) => return Err(invalid_request(err.to_string())), - }; - let account = account_state.account.map(Account::from); - - Ok(GetAccountResponse { - account, - requires_openai_auth: account_state.requires_openai_auth, - }) - } - async fn get_account_rate_limits_response( &self, params: GetAccountRateLimitsParams, diff --git a/codex-rs/app-server/src/request_processors/account_processor/workspace_routing.rs b/codex-rs/app-server/src/request_processors/account_processor/workspace_routing.rs new file mode 100644 index 0000000000..ea1d1322f3 --- /dev/null +++ b/codex-rs/app-server/src/request_processors/account_processor/workspace_routing.rs @@ -0,0 +1,276 @@ +//! Discovers routing for the selected ChatGPT workspace without crossing auth owners. +//! Credential refreshes invalidate the cache, but do not cancel same-owner discovery. + +use super::*; +use codex_app_server_protocol::AccountRoutingOverride; +use codex_app_server_protocol::WorkspaceRouting; +use codex_backend_client::AccountEntry; +use url::Url; + +pub(super) struct CachedWorkspaceRouting { + auth_generation: u64, + effective_chatgpt_base_url: String, + required_chatgpt_base_url: Option, + routing: WorkspaceRouting, +} + +impl AccountRequestProcessor { + pub(crate) fn notify_workspace_routing_to_connection(&self, connection_id: ConnectionId) { + let processor = self.clone(); + let auth_changes = self.auth_manager.auth_change_state_receiver(); + let owner_generation = auth_changes.borrow().owner_generation; + tokio::spawn(async move { + if auth_changes.borrow().owner_generation != owner_generation { + return; + } + if let Ok(response) = processor + .get_account_response(GetAccountParams { + refresh_token: false, + }) + .await + && response.workspace_routing.is_some() + && auth_changes.borrow().owner_generation == owner_generation + { + let notification = processor.current_account_updated_notification(); + processor + .outgoing + .send_account_notification( + Some(connection_id), + &auth_changes, + owner_generation, + AccountNotification::Updated(notification), + ) + .await; + } + }); + } + + pub(super) async fn get_account_response( + &self, + params: GetAccountParams, + ) -> Result { + self.refresh_token_if_requested(params.refresh_token).await; + let mut auth_changes = self.auth_manager.auth_change_state_receiver(); + let auth_state = *auth_changes.borrow_and_update(); + let current_auth_changes = auth_changes.clone(); + let read = Box::pin(async { + let _fetch_permit = self + .workspace_routing_fetch + .acquire() + .await + .map_err(|_| internal_error("workspace routing discovery cancelled"))?; + let config = match self + .config_manager + .load_latest_config(/*fallback_cwd*/ None) + .await + { + Ok(config) => config, + Err(_) + if !self + .auth_manager + .auth_cached() + .as_ref() + .is_some_and(CodexAuth::is_chatgpt_auth) => + { + self.config.as_ref().clone() + } + Err(_) => return Err(internal_error("failed to load workspace requirements")), + }; + let auth = self.auth_manager.auth_cached(); + let provider = create_model_provider( + config.model_provider.clone(), + Some(self.auth_manager.clone()), + ); + let account_state = provider + .account_state() + .map_err(|err| invalid_request(err.to_string()))?; + let account = account_state.account.map(Account::from); + let workspace_routing = if let Some((auth, account_id)) = auth + .as_ref() + .filter(|auth| { + auth.is_chatgpt_auth() && matches!(account, Some(Account::Chatgpt { .. })) + }) + .and_then(|auth| auth.get_account_id().map(|account_id| (auth, account_id))) + { + if account_id.is_empty() { + return Err(internal_error( + "workspace routing requires a ChatGPT account id", + )); + } + let required_chatgpt_base_url = config + .config_layer_stack + .requirements_toml() + .chatgpt_base_url + .clone(); + let cached = self + .workspace_routing + .lock() + .await + .as_ref() + .filter(|cached| { + cached.auth_generation == auth_state.generation + && cached.effective_chatgpt_base_url == config.chatgpt_base_url + && cached.required_chatgpt_base_url == required_chatgpt_base_url + && cached.routing.chatgpt_account_id == account_id + }) + .map(|cached| cached.routing.clone()); + if let Some(cached) = cached { + Some(cached) + } else { + *self.workspace_routing.lock().await = None; + let client = BackendClient::from_auth( + &config.chatgpt_base_url, + auth, + config.http_client_factory(), + ); + let response = client + .get_accounts_check() + .await + .map_err(|_| internal_error("workspace routing discovery failed"))?; + let mut accounts = response + .accounts + .into_iter() + .filter(|account| account.id == account_id); + let entry = accounts.next().ok_or_else(|| { + internal_error("selected workspace missing from routing discovery") + })?; + if accounts.next().is_some() { + return Err(internal_error("duplicate workspace in routing discovery")); + } + let latest_config = self + .config_manager + .load_latest_config(/*fallback_cwd*/ None) + .await + .map_err(|_| internal_error("failed to reload workspace requirements"))?; + if latest_config.chatgpt_base_url != config.chatgpt_base_url + || latest_config.model_provider != config.model_provider + || latest_config + .config_layer_stack + .requirements_toml() + .chatgpt_base_url + != required_chatgpt_base_url + { + return Err(internal_error( + "configuration changed during workspace routing discovery; retry account/read", + )); + } + let routing = resolve_routing( + entry, + required_chatgpt_base_url.as_deref(), + &config.chatgpt_base_url, + )?; + let mut cached = self.workspace_routing.lock().await; + if current_auth_changes.borrow().owner_generation != auth_state.owner_generation + { + return Err(internal_error( + "account changed during workspace routing discovery", + )); + } + *cached = Some(CachedWorkspaceRouting { + auth_generation: auth_state.generation, + effective_chatgpt_base_url: config.chatgpt_base_url.clone(), + required_chatgpt_base_url, + routing: routing.clone(), + }); + Some(routing) + } + } else { + *self.workspace_routing.lock().await = None; + None + }; + Ok(GetAccountResponse { + account, + requires_openai_auth: account_state.requires_openai_auth, + workspace_routing, + }) + }); + let result = tokio::select! { + biased; + _ = self.workspace_routing_shutdown.cancelled() => { + return Err(internal_error("workspace routing discovery cancelled during shutdown")); + } + _ = auth_changes.wait_for(|state| state.owner_generation != auth_state.owner_generation) => { + return Err(internal_error("account changed during workspace routing discovery")); + } + result = tokio::time::timeout(Duration::from_secs(/*secs*/ 10), read) => { + result.map_err(|_| internal_error("workspace routing discovery timed out"))? + } + }; + if auth_changes.borrow().owner_generation != auth_state.owner_generation { + return Err(internal_error( + "account changed during workspace routing discovery", + )); + } + result + } +} + +fn resolve_routing( + entry: AccountEntry, + required_chatgpt_base_url: Option<&str>, + effective_base_url: &str, +) -> Result { + let backend = entry + .workspace_backend_origin + .ok_or_else(|| internal_error("workspace routing discovery missing backend origin"))?; + let account_routing_override = match entry.account_routing_override.as_deref() { + Some("NO_CONSTRAINT") => AccountRoutingOverride::NoConstraint, + Some("us") => AccountRoutingOverride::Us, + Some("us_cr") => AccountRoutingOverride::UsCr, + _ => { + return Err(internal_error( + "workspace routing discovery has invalid account routing override", + )); + } + }; + let required = required_chatgpt_base_url + .map(parse_backend_url) + .transpose()?; + let discovered = if backend == "NO_CONSTRAINT" { + None + } else { + let url = parse_backend_url(&backend)?; + if url.path() != "/" || url.query().is_some() || url.fragment().is_some() { + return Err(internal_error( + "workspace routing discovery must return an origin", + )); + } + Some(url) + }; + let origin = match (required, discovered) { + (Some(required), Some(discovered)) => { + if required.origin() != discovered.origin() { + return Err(internal_error( + "required ChatGPT backend conflicts with workspace routing", + )); + } + required.origin() + } + (Some(url), None) | (None, Some(url)) => url.origin(), + (None, None) => parse_backend_url(effective_base_url)?.origin(), + }; + Ok(WorkspaceRouting { + chatgpt_account_id: entry.id, + backend_origin: origin.ascii_serialization(), + account_routing_override, + }) +} + +fn parse_backend_url(value: &str) -> Result { + let url = Url::parse(value).map_err(|_| internal_error("invalid workspace backend URL"))?; + if url.scheme() != "https" + || url.host_str().is_none() + || !url.username().is_empty() + || url.password().is_some() + || value.trim() != value + { + return Err(internal_error( + "workspace backend must use an HTTPS origin without credentials", + )); + } + Ok(url) +} + +#[cfg(test)] +#[path = "workspace_routing_tests.rs"] +mod tests; diff --git a/codex-rs/app-server/src/request_processors/account_processor/workspace_routing_tests.rs b/codex-rs/app-server/src/request_processors/account_processor/workspace_routing_tests.rs new file mode 100644 index 0000000000..4532d4c106 --- /dev/null +++ b/codex-rs/app-server/src/request_processors/account_processor/workspace_routing_tests.rs @@ -0,0 +1,83 @@ +//! Origin resolution and strict discovery validation. + +use super::*; +use pretty_assertions::assert_eq; +use test_case::test_case; + +#[test_case(Some("https://gov.chatgpt.com/backend-api/"), "NO_CONSTRAINT", "https://gov.chatgpt.com"; "configured_only")] +#[test_case(None, "https://gov.chatgpt.com", "https://gov.chatgpt.com"; "discovered_only")] +#[test_case(Some("https://GOV.chatgpt.com:443/backend-api/"), "https://gov.chatgpt.com", "https://gov.chatgpt.com"; "matching_effective_port")] +#[test_case(Some("https://example.com:8443/backend-api/"), "https://example.com:8443", "https://example.com:8443"; "custom_port")] +#[test_case(None, "NO_CONSTRAINT", "https://chatgpt.com"; "existing_default")] +fn resolves_origins(required: Option<&str>, discovered: &str, expected: &str) { + let entry = serde_json::from_value(serde_json::json!({ + "id": "workspace", "workspace_backend_origin": discovered, + "account_routing_override": "NO_CONSTRAINT", + })) + .unwrap(); + assert_eq!( + resolve_routing(entry, required, "https://chatgpt.com/backend-api/").unwrap(), + WorkspaceRouting { + chatgpt_account_id: "workspace".into(), + backend_origin: expected.into(), + account_routing_override: AccountRoutingOverride::NoConstraint, + } + ); +} + +#[test_case("https://other.example/backend-api/", "https://gov.chatgpt.com"; "host_conflict")] +#[test_case("http://gov.chatgpt.com/backend-api/", "https://gov.chatgpt.com"; "scheme_conflict")] +#[test_case("https://gov.chatgpt.com:444/backend-api/", "https://gov.chatgpt.com"; "port_conflict")] +fn rejects_conflicting_origins(required: &str, discovered: &str) { + let entry = serde_json::from_value(serde_json::json!({ + "id": "workspace", "workspace_backend_origin": discovered, + "account_routing_override": "us_cr", + })) + .unwrap(); + assert!(resolve_routing(entry, Some(required), "https://chatgpt.com").is_err()); +} + +#[test_case(serde_json::json!(null), serde_json::json!("us_cr"); "null_origin")] +#[test_case(serde_json::json!("NO_CONSTRAINT"), serde_json::json!(null); "null_routing")] +#[test_case(serde_json::json!(""), serde_json::json!("us_cr"); "empty_origin")] +#[test_case(serde_json::json!("https://example.com/backend-api/"), serde_json::json!("us_cr"); "origin_with_path")] +#[test_case(serde_json::json!("https://example.com?query"), serde_json::json!("us_cr"); "origin_with_query")] +#[test_case(serde_json::json!("https://example.com#fragment"), serde_json::json!("us_cr"); "origin_with_fragment")] +#[test_case(serde_json::json!("https://user:pass@example.com"), serde_json::json!("us_cr"); "credentials")] +#[test_case(serde_json::json!("http://example.com"), serde_json::json!("us_cr"); "insecure_origin")] +#[test_case(serde_json::json!("NO_CONSTRAINT"), serde_json::json!("unknown"); "unknown_routing")] +#[test_case(serde_json::json!("NO_CONSTRAINT"), serde_json::json!(""); "empty_routing")] +fn rejects_invalid_discovery(backend: serde_json::Value, routing: serde_json::Value) { + let entry = serde_json::from_value(serde_json::json!({ + "id": "workspace", "workspace_backend_origin": backend, "account_routing_override": routing, + })) + .unwrap(); + assert!( + resolve_routing( + entry, + /*required_chatgpt_base_url*/ None, + "https://chatgpt.com" + ) + .is_err() + ); +} + +#[test] +fn unrestricted_discovery_uses_effective_custom_backend() { + let entry = serde_json::from_value(serde_json::json!({ + "id": "workspace", "workspace_backend_origin": "NO_CONSTRAINT", "account_routing_override": "us", + })).unwrap(); + assert_eq!( + resolve_routing( + entry, + /*required_chatgpt_base_url*/ None, + "https://custom.example:8443/backend-api/" + ) + .unwrap(), + WorkspaceRouting { + chatgpt_account_id: "workspace".into(), + backend_origin: "https://custom.example:8443".into(), + account_routing_override: AccountRoutingOverride::Us, + } + ); +} diff --git a/codex-rs/app-server/tests/common/auth_fixtures.rs b/codex-rs/app-server/tests/common/auth_fixtures.rs index 24fbd1c6b5..52a830eb6d 100644 --- a/codex-rs/app-server/tests/common/auth_fixtures.rs +++ b/codex-rs/app-server/tests/common/auth_fixtures.rs @@ -15,6 +15,24 @@ use codex_login::token_data::parse_chatgpt_jwt_claims; use codex_protocol::auth::AuthMode; use serde_json::json; +pub async fn mount_workspace_routing(server: &wiremock::MockServer) { + wiremock::Mock::given(wiremock::matchers::method("GET")) + .and(wiremock::matchers::path("/backend-api/wham/accounts/check")) + .respond_with(|request: &wiremock::Request| { + let account_id = request + .headers + .get("chatgpt-account-id") + .and_then(|value| value.to_str().ok()) + .unwrap_or("account-123"); + wiremock::ResponseTemplate::new(200).set_body_json(json!({ + "accounts": [{"id": account_id, "workspace_backend_origin": "https://chatgpt.com", + "account_routing_override": "NO_CONSTRAINT"}], + })) + }) + .mount(server) + .await; +} + /// Builder for writing a fake ChatGPT auth.json in tests. #[derive(Debug, Clone)] pub struct ChatGptAuthFixture { diff --git a/codex-rs/app-server/tests/common/lib.rs b/codex-rs/app-server/tests/common/lib.rs index 7b545e94f3..f7ebde5147 100644 --- a/codex-rs/app-server/tests/common/lib.rs +++ b/codex-rs/app-server/tests/common/lib.rs @@ -16,6 +16,7 @@ pub use analytics_server::start_analytics_events_server; pub use auth_fixtures::ChatGptAuthFixture; pub use auth_fixtures::ChatGptIdTokenClaims; pub use auth_fixtures::encode_id_token; +pub use auth_fixtures::mount_workspace_routing; pub use auth_fixtures::write_chatgpt_auth; use codex_app_server_protocol::JSONRPCResponse; pub use config::MockResponsesConfig; diff --git a/codex-rs/app-server/tests/common/test_app_server.rs b/codex-rs/app-server/tests/common/test_app_server.rs index 8a0a2fefe3..39fbef236c 100644 --- a/codex-rs/app-server/tests/common/test_app_server.rs +++ b/codex-rs/app-server/tests/common/test_app_server.rs @@ -184,6 +184,7 @@ impl TestAppServer { env_overrides: Vec::new(), args: vec![DISABLE_PLUGIN_STARTUP_TASKS_ARG.to_string()], exec_server_delay: None, + mock_chatgpt_backend: false, } } @@ -1856,6 +1857,7 @@ pub struct TestAppServerBuilder { env_overrides: Vec<(String, Option)>, args: Vec, exec_server_delay: Option, + mock_chatgpt_backend: bool, } enum TestAppServerEnvironment { @@ -1864,6 +1866,11 @@ enum TestAppServerEnvironment { } impl TestAppServerBuilder { + pub fn with_mock_chatgpt_backend(mut self) -> Self { + self.mock_chatgpt_backend = true; + self + } + /// Uses this existing CODEX_HOME instead of a temporary one. pub fn with_codex_home(mut self, codex_home: &Path) -> Self { self.codex_home = Some(codex_home.to_path_buf()); @@ -1958,6 +1965,7 @@ impl TestAppServerBuilder { mut env_overrides, args, exec_server_delay, + mock_chatgpt_backend, } = self; let (codex_home, owned_codex_home) = match codex_home { Some(codex_home) => (codex_home, None), @@ -1969,7 +1977,9 @@ impl TestAppServerBuilder { ) } }; - let attribution_settings_server = if codex_home.join("auth.json").is_file() { + let attribution_settings_server = if mock_chatgpt_backend + || codex_home.join("auth.json").is_file() + { let config_path = codex_home.join("config.toml"); let config = std::fs::read_to_string(&config_path)?; if config @@ -1979,6 +1989,12 @@ impl TestAppServerBuilder { None } else { let settings_server = MockServer::start().await; + crate::mount_workspace_routing(&settings_server).await; + Mock::given(method("GET")) + .and(path("/backend-api/wham/config/bundle")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({}))) + .mount(&settings_server) + .await; Mock::given(method("GET")) .and(path("/backend-api/wham/settings/user")) .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ diff --git a/codex-rs/app-server/tests/suite/auth.rs b/codex-rs/app-server/tests/suite/auth.rs index 504c48fef8..79bd14f003 100644 --- a/codex-rs/app-server/tests/suite/auth.rs +++ b/codex-rs/app-server/tests/suite/auth.rs @@ -163,6 +163,7 @@ async fn personal_access_token_without_email_supports_auth_status_and_account_re let authapi_base_url = server.uri(); let mut mcp = TestAppServer::builder() + .with_mock_chatgpt_backend() .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[ @@ -219,6 +220,12 @@ async fn personal_access_token_without_email_supports_auth_status_and_account_re assert_eq!( to_response::(response)?, GetAccountResponse { + workspace_routing: Some(codex_app_server_protocol::WorkspaceRouting { + chatgpt_account_id: "account-123".to_string(), + backend_origin: "https://chatgpt.com".to_string(), + account_routing_override: + codex_app_server_protocol::AccountRoutingOverride::NoConstraint, + }), account: Some(Account::Chatgpt { email: None, plan_type: AccountPlanType::EnterpriseCbpAutomation, diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index 6db61e21fe..f2966c6850 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -57,14 +57,17 @@ use pretty_assertions::assert_eq; use serde_json::json; use serial_test::serial; use std::path::Path; +use std::sync::Arc; use std::time::Duration; use tempfile::TempDir; use test_case::test_case; +use tokio::sync::Notify; use tokio::time::timeout; use url::Url; use wiremock::Mock; use wiremock::MockServer; use wiremock::ResponseTemplate; +use wiremock::matchers::header; use wiremock::matchers::method; use wiremock::matchers::path; @@ -80,6 +83,16 @@ const WORKSPACE_ID_REFRESHED: &str = "123e4567-e89b-42d3-a456-426614174012"; const WORKSPACE_ID_DEVICE: &str = "123e4567-e89b-42d3-a456-426614174013"; const WORKSPACE_ID_STALE: &str = "123e4567-e89b-42d3-a456-426614174014"; +fn expected_workspace_routing( + account_id: &str, +) -> Option { + Some(codex_app_server_protocol::WorkspaceRouting { + chatgpt_account_id: account_id.to_string(), + backend_origin: "https://chatgpt.com".to_string(), + account_routing_override: codex_app_server_protocol::AccountRoutingOverride::NoConstraint, + }) +} + // Helper to create a minimal config.toml for the app server #[derive(Default)] struct CreateConfigTomlParams { @@ -426,6 +439,7 @@ async fn set_auth_token_updates_account_and_notifies() -> Result<()> { )?; let mut mcp = TestAppServer::builder() + .with_mock_chatgpt_backend() .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) @@ -465,6 +479,7 @@ async fn set_auth_token_updates_account_and_notifies() -> Result<()> { assert_eq!( account, GetAccountResponse { + workspace_routing: expected_workspace_routing(WORKSPACE_ID_EMBEDDED), account: Some(Account::Chatgpt { email: Some("embedded@example.com".to_string()), plan_type: AccountPlanType::Pro, @@ -509,6 +524,7 @@ async fn account_read_refresh_token_is_noop_in_external_mode() -> Result<()> { )?; let mut mcp = TestAppServer::builder() + .with_mock_chatgpt_backend() .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[("OPENAI_API_KEY", None)]) @@ -541,6 +557,7 @@ async fn account_read_refresh_token_is_noop_in_external_mode() -> Result<()> { assert_eq!( account, GetAccountResponse { + workspace_routing: expected_workspace_routing(WORKSPACE_ID_EMBEDDED), account: Some(Account::Chatgpt { email: Some("embedded@example.com".to_string()), plan_type: AccountPlanType::Pro, @@ -588,6 +605,7 @@ async fn respond_to_refresh_request( } async fn mount_disabled_attribution_settings(mock_server: &MockServer) { + app_test_support::mount_workspace_routing(mock_server).await; Mock::given(method("GET")) .and(path("/backend-api/wham/settings/user")) .respond_with(ResponseTemplate::new(200).set_body_json(json!({ @@ -642,6 +660,7 @@ async fn external_auth_refreshes_on_unauthorized() -> Result<()> { )?; let mut mcp = TestAppServer::builder() + .with_mock_chatgpt_backend() .with_codex_home(codex_home.path()) .with_env_overrides(&[("OPENAI_API_KEY", None)]) .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) @@ -743,6 +762,7 @@ async fn external_auth_refresh_error_fails_turn() -> Result<()> { )?; let mut mcp = TestAppServer::builder() + .with_mock_chatgpt_backend() .with_codex_home(codex_home.path()) .with_env_overrides(&[("OPENAI_API_KEY", None)]) .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) @@ -860,6 +880,7 @@ async fn external_auth_refresh_mismatched_workspace_fails_turn() -> Result<()> { )?; let mut mcp = TestAppServer::builder() + .with_mock_chatgpt_backend() .with_codex_home(codex_home.path()) .with_env_overrides(&[("OPENAI_API_KEY", None)]) .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) @@ -970,6 +991,7 @@ async fn external_auth_refresh_invalid_access_token_fails_turn() -> Result<()> { )?; let mut mcp = TestAppServer::builder() + .with_mock_chatgpt_backend() .with_codex_home(codex_home.path()) .with_env_overrides(&[("OPENAI_API_KEY", None)]) .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) @@ -1225,6 +1247,7 @@ async fn login_amazon_bedrock_replaces_primary_auth_and_persists_provider( assert_eq!( read_account(&mut mcp).await?, GetAccountResponse { + workspace_routing: None, account: Some(Account::AmazonBedrock { uses_codex_managed_credentials: true, }), @@ -1255,6 +1278,7 @@ async fn login_amazon_bedrock_replaces_primary_auth_and_persists_provider( assert_eq!( read_account(&mut mcp).await?, GetAccountResponse { + workspace_routing: None, account: None, requires_openai_auth: true, } @@ -1468,6 +1492,7 @@ async fn logout_managed_bedrock_restores_default_account( assert_eq!( read_account(&mut mcp).await?, GetAccountResponse { + workspace_routing: None, account: Some(Account::AmazonBedrock { uses_codex_managed_credentials: true, }), @@ -1518,6 +1543,7 @@ async fn logout_managed_bedrock_restores_default_account( assert_eq!( read_account(&mut mcp).await?, GetAccountResponse { + workspace_routing: None, account: None, requires_openai_auth: true, } @@ -1582,6 +1608,7 @@ async fn logout_aws_managed_bedrock_clears_provider_and_restores_default_account assert_eq!( read_account(&mut mcp).await?, GetAccountResponse { + workspace_routing: None, account: Some(Account::AmazonBedrock { uses_codex_managed_credentials: false, }), @@ -1620,6 +1647,7 @@ async fn logout_aws_managed_bedrock_clears_provider_and_restores_default_account assert_eq!( read_account(&mut mcp).await?, GetAccountResponse { + workspace_routing: None, account: None, requires_openai_auth: true, } @@ -1688,6 +1716,7 @@ async fn logout_managed_bedrock_preserves_changed_provider_without_experimental_ assert_eq!( read_account(&mut mcp).await?, GetAccountResponse { + workspace_routing: None, account: None, requires_openai_auth: false, } @@ -1767,6 +1796,7 @@ async fn login_managed_bedrock_updates_active_bedrock_account() -> Result<()> { assert_eq!( read_account(&mut mcp).await?, GetAccountResponse { + workspace_routing: None, account: Some(Account::AmazonBedrock { uses_codex_managed_credentials: true, }), @@ -1885,6 +1915,7 @@ async fn login_account_amazon_bedrock_rejected_with_external_chatgpt_auth() -> R )?; let mut mcp = TestAppServer::builder() + .with_mock_chatgpt_backend() .with_codex_home(codex_home.path()) .without_auto_env() .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) @@ -1939,6 +1970,7 @@ async fn login_account_api_key_rejected_when_forced_chatgpt() -> Result<()> { )?; let mut mcp = TestAppServer::builder() + .with_mock_chatgpt_backend() .with_codex_home(codex_home.path()) .without_auto_env() .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) @@ -2008,6 +2040,7 @@ async fn login_account_chatgpt_device_code_returns_error_when_disabled() -> Resu let issuer = mock_server.uri(); let mut mcp = TestAppServer::builder() + .with_mock_chatgpt_backend() .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[ @@ -2073,6 +2106,7 @@ async fn login_account_chatgpt_device_code_succeeds_and_notifies() -> Result<()> let issuer = mock_server.uri(); let mut mcp = TestAppServer::builder() + .with_mock_chatgpt_backend() .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[ @@ -2127,6 +2161,198 @@ async fn login_account_chatgpt_device_code_succeeds_and_notifies() -> Result<()> Ok(()) } +#[derive(Clone, Copy)] +enum LoginRefreshTrigger { + AuthStatus, + UnauthorizedConfig, +} + +#[test_case("/backend-api/wham/config/bundle", LoginRefreshTrigger::AuthStatus; "refresh_during_requirements")] +#[test_case("/backend-api/wham/accounts/check", LoginRefreshTrigger::AuthStatus; "refresh_during_routing")] +#[test_case("/backend-api/wham/config/bundle", LoginRefreshTrigger::UnauthorizedConfig; "refresh_during_account_read_requirements")] +#[tokio::test] +async fn login_survives_same_owner_token_refresh( + delayed_path: &str, + refresh_trigger: LoginRefreshTrigger, +) -> Result<()> { + let codex_home = TempDir::new()?; + let backend = MockServer::start().await; + create_config_toml( + codex_home.path(), + CreateConfigTomlParams { + requires_openai_auth: Some(true), + chatgpt_base_url: Some(format!("{}/backend-api", backend.uri())), + ..Default::default() + }, + )?; + write_models_cache(codex_home.path()).await?; + mock_device_code_usercode(&backend, /*interval_seconds*/ 0).await; + mock_device_code_token_success(&backend).await; + let id_token = encode_id_token( + &ChatGptIdTokenClaims::new() + .email("device@example.com") + .plan_type("enterprise") + .chatgpt_user_id("device-user") + .chatgpt_account_id(WORKSPACE_ID_DEVICE), + )?; + mock_oauth_token(&backend, &id_token).await; + Mock::given(method("POST")) + .and(path("/oauth/refresh")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id_token": id_token, + "access_token": "refreshed-access-token", + "refresh_token": "refreshed-refresh-token", + }))) + .expect(1) + .mount(&backend) + .await; + let started = Arc::new(Notify::new()); + for (request_path, response) in [ + ("/backend-api/wham/config/bundle", json!({})), + ( + "/backend-api/wham/accounts/check", + json!({"accounts": [{ + "id": WORKSPACE_ID_DEVICE, "workspace_backend_origin": "https://chatgpt.com", + "account_routing_override": "NO_CONSTRAINT", + }]}), + ), + ] { + let mock = Mock::given(method("GET")).and(path(request_path)); + let mock = if request_path == "/backend-api/wham/accounts/check" { + mock.and(header("authorization", "Bearer refreshed-access-token")) + } else { + mock + }; + mock.respond_with(ResponseTemplate::new(200).set_body_json(response.clone())) + .mount(&backend) + .await; + if request_path == delayed_path { + let request_started = Arc::clone(&started); + Mock::given(method("GET")) + .and(path(request_path)) + .and(header("authorization", "Bearer access-token-123")) + .respond_with(move |_: &wiremock::Request| { + request_started.notify_one(); + ResponseTemplate::new(match refresh_trigger { + LoginRefreshTrigger::AuthStatus => 200, + LoginRefreshTrigger::UnauthorizedConfig => 401, + }) + .set_delay(Duration::from_secs(/*secs*/ 2)) + .set_body_json(response.clone()) + }) + .with_priority(1) + .expect(1..=2) + .mount(&backend) + .await; + } + } + let issuer = backend.uri(); + let refresh_url = format!("{issuer}/oauth/refresh"); + let mut server = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .with_env_overrides(&[ + ("OPENAI_API_KEY", None), + (LOGIN_ISSUER_ENV_VAR, Some(issuer.as_str())), + ( + REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR, + Some(refresh_url.as_str()), + ), + ]) + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) + .await?; + let request = server + .send_login_account_chatgpt_device_code_request() + .await?; + let login: LoginAccountResponse = + timeout(DEFAULT_READ_TIMEOUT, server.read_response(request)).await??; + let LoginAccountResponse::ChatgptDeviceCode { login_id, .. } = login else { + bail!("unexpected login response: {login:?}"); + }; + timeout(DEFAULT_READ_TIMEOUT, started.notified()).await?; + let expected_account = GetAccountResponse { + account: Some(Account::Chatgpt { + email: Some("device@example.com".into()), + plan_type: AccountPlanType::Enterprise, + }), + requires_openai_auth: true, + workspace_routing: expected_workspace_routing(WORKSPACE_ID_DEVICE), + }; + match refresh_trigger { + LoginRefreshTrigger::AuthStatus => { + let request = server + .send_get_auth_status_request(GetAuthStatusParams { + include_token: Some(true), + refresh_token: Some(true), + }) + .await?; + let refreshed: GetAuthStatusResponse = + timeout(DEFAULT_READ_TIMEOUT, server.read_response(request)).await??; + assert_eq!( + refreshed.auth_token.as_deref(), + Some("refreshed-access-token") + ); + } + LoginRefreshTrigger::UnauthorizedConfig => { + assert_eq!(read_account(&mut server).await?, expected_account); + } + } + let completed: AccountLoginCompletedNotification = timeout( + DEFAULT_READ_TIMEOUT, + server.read_notification("account/login/completed"), + ) + .await??; + assert_eq!( + completed, + AccountLoginCompletedNotification { + login_id: Some(login_id), + success: true, + error: None, + onboarding_entrypoint: None, + } + ); + let updated: AccountUpdatedNotification = timeout( + DEFAULT_READ_TIMEOUT, + server.read_notification("account/updated"), + ) + .await??; + assert_eq!( + updated, + AccountUpdatedNotification { + auth_mode: Some(AuthMode::Chatgpt), + plan_type: Some(AccountPlanType::Enterprise), + } + ); + assert_eq!(read_account(&mut server).await?, expected_account); + let requests = backend + .received_requests() + .await + .expect("recorded requests"); + let routing_tokens = requests + .iter() + .filter(|request| request.url.path() == "/backend-api/wham/accounts/check") + .map(|request| { + request.headers["authorization"] + .to_str() + .expect("routing authorization header") + }) + .collect::>(); + assert_eq!( + routing_tokens, + if delayed_path == "/backend-api/wham/accounts/check" { + vec!["Bearer access-token-123", "Bearer refreshed-access-token"] + } else if matches!(refresh_trigger, LoginRefreshTrigger::UnauthorizedConfig) { + vec![ + "Bearer refreshed-access-token", + "Bearer refreshed-access-token", + ] + } else { + vec!["Bearer refreshed-access-token"] + } + ); + backend.verify().await; + Ok(()) +} + #[tokio::test] async fn login_account_chatgpt_device_code_failure_notifies_without_account_update() -> Result<()> { let codex_home = TempDir::new()?; @@ -2146,6 +2372,7 @@ async fn login_account_chatgpt_device_code_failure_notifies_without_account_upda let issuer = mock_server.uri(); let mut mcp = TestAppServer::builder() + .with_mock_chatgpt_backend() .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[ @@ -2217,6 +2444,7 @@ async fn login_account_chatgpt_device_code_can_be_cancelled() -> Result<()> { let issuer = mock_server.uri(); let mut mcp = TestAppServer::builder() + .with_mock_chatgpt_backend() .with_codex_home(codex_home.path()) .without_auto_env() .with_env_overrides(&[ @@ -2282,6 +2510,7 @@ async fn login_account_chatgpt_start_can_be_cancelled() -> Result<()> { create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; let mut mcp = TestAppServer::builder() + .with_mock_chatgpt_backend() .with_codex_home(codex_home.path()) .without_auto_env() .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) @@ -2351,6 +2580,7 @@ async fn login_account_chatgpt_uses_oauth_overrides() -> Result<()> { let issuer = mock_server.uri(); let mut mcp = TestAppServer::builder() + .with_mock_chatgpt_backend() .with_codex_home(codex_home.path()) .without_auto_env() // Exercise packaged-build behavior without the debug-only startup flag. @@ -2550,6 +2780,7 @@ async fn set_auth_token_cancels_active_chatgpt_login() -> Result<()> { create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; let mut mcp = TestAppServer::builder() + .with_mock_chatgpt_backend() .with_codex_home(codex_home.path()) .without_auto_env() .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) @@ -2615,6 +2846,7 @@ async fn login_account_chatgpt_includes_forced_workspace_query_param() -> Result )?; let mut mcp = TestAppServer::builder() + .with_mock_chatgpt_backend() .with_codex_home(codex_home.path()) .without_auto_env() .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) @@ -2650,6 +2882,7 @@ async fn login_account_chatgpt_includes_forced_workspace_allowlist_query_param() )?; let mut mcp = TestAppServer::builder() + .with_mock_chatgpt_backend() .with_codex_home(codex_home.path()) .without_auto_env() .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) @@ -2738,6 +2971,7 @@ async fn get_account_with_api_key() -> Result<()> { timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let expected = GetAccountResponse { + workspace_routing: None, account: Some(Account::ApiKey {}), requires_openai_auth: true, }; @@ -2771,6 +3005,7 @@ async fn get_account_when_auth_not_required() -> Result<()> { timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let expected = GetAccountResponse { + workspace_routing: None, account: None, requires_openai_auth: false, }; @@ -2811,6 +3046,7 @@ region = "us-west-2" timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let expected = GetAccountResponse { + workspace_routing: None, account: Some(Account::AmazonBedrock { uses_codex_managed_credentials: false, }), @@ -2849,6 +3085,7 @@ command = "print-token" assert_eq!( read_account(&mut mcp).await?, GetAccountResponse { + workspace_routing: None, account: Some(Account::AmazonBedrock { uses_codex_managed_credentials: false, }), @@ -2887,6 +3124,7 @@ region = "us-west-2" assert_eq!( read_account(&mut mcp).await?, GetAccountResponse { + workspace_routing: None, account: Some(Account::AmazonBedrock { uses_codex_managed_credentials: false, }), @@ -2952,6 +3190,7 @@ async fn get_account_with_managed_bedrock_provider() -> Result<()> { assert_eq!( received, GetAccountResponse { + workspace_routing: None, account: Some(Account::AmazonBedrock { uses_codex_managed_credentials: true, }), @@ -2974,6 +3213,7 @@ async fn get_account_with_chatgpt() -> Result<()> { write_chatgpt_auth( codex_home.path(), ChatGptAuthFixture::new("access-chatgpt") + .account_id("account-123") .email("user@example.com") .plan_type("pro"), AuthCredentialsStoreMode::File, @@ -2995,6 +3235,7 @@ async fn get_account_with_chatgpt() -> Result<()> { timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let expected = GetAccountResponse { + workspace_routing: expected_workspace_routing("account-123"), account: Some(Account::Chatgpt { email: Some("user@example.com".to_string()), plan_type: AccountPlanType::Pro, @@ -3024,6 +3265,7 @@ async fn get_account_with_chatgpt_plan_variants_returns_plan_type( write_chatgpt_auth( codex_home.path(), ChatGptAuthFixture::new("access-chatgpt") + .account_id("account-123") .email("user@example.com") .plan_type(plan_type), AuthCredentialsStoreMode::File, @@ -3047,6 +3289,7 @@ async fn get_account_with_chatgpt_plan_variants_returns_plan_type( assert_eq!( received, GetAccountResponse { + workspace_routing: expected_workspace_routing("account-123"), account: Some(Account::Chatgpt { email: Some("user@example.com".to_string()), plan_type: expected_plan, @@ -3069,7 +3312,9 @@ async fn get_account_with_chatgpt_without_email() -> Result<()> { )?; write_chatgpt_auth( codex_home.path(), - ChatGptAuthFixture::new("access-chatgpt").plan_type("pro"), + ChatGptAuthFixture::new("access-chatgpt") + .account_id("account-123") + .plan_type("pro"), AuthCredentialsStoreMode::File, )?; @@ -3091,6 +3336,7 @@ async fn get_account_with_chatgpt_without_email() -> Result<()> { assert_eq!( received, GetAccountResponse { + workspace_routing: expected_workspace_routing("account-123"), account: Some(Account::Chatgpt { email: None, plan_type: AccountPlanType::Pro, @@ -3172,6 +3418,7 @@ async fn get_account_omits_chatgpt_after_permanent_refresh_failure() -> Result<( assert_eq!( received, GetAccountResponse { + workspace_routing: None, account: None, requires_openai_auth: true, } @@ -3192,7 +3439,9 @@ async fn get_account_with_chatgpt_missing_plan_claim_returns_unknown() -> Result )?; write_chatgpt_auth( codex_home.path(), - ChatGptAuthFixture::new("access-chatgpt").email("user@example.com"), + ChatGptAuthFixture::new("access-chatgpt") + .account_id("account-123") + .email("user@example.com"), AuthCredentialsStoreMode::File, )?; @@ -3212,6 +3461,7 @@ async fn get_account_with_chatgpt_missing_plan_claim_returns_unknown() -> Result timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; let expected = GetAccountResponse { + workspace_routing: expected_workspace_routing("account-123"), account: Some(Account::Chatgpt { email: Some("user@example.com".to_string()), plan_type: AccountPlanType::Unknown, diff --git a/codex-rs/app-server/tests/suite/v2/config_requirements_login.rs b/codex-rs/app-server/tests/suite/v2/config_requirements_login.rs index 829a4ad8d6..679271935a 100644 --- a/codex-rs/app-server/tests/suite/v2/config_requirements_login.rs +++ b/codex-rs/app-server/tests/suite/v2/config_requirements_login.rs @@ -170,6 +170,7 @@ command = "print-token" uses_codex_managed_credentials: false }), requires_openai_auth: false, + workspace_routing: None, } ); assert!( diff --git a/codex-rs/app-server/tests/suite/v2/mod.rs b/codex-rs/app-server/tests/suite/v2/mod.rs index 049e561d91..405d1bfbdc 100644 --- a/codex-rs/app-server/tests/suite/v2/mod.rs +++ b/codex-rs/app-server/tests/suite/v2/mod.rs @@ -146,6 +146,7 @@ mod turn_steer; mod view_image; mod web_search; mod windows_sandbox_setup; +mod workspace_routing; mod user_verification; mod user_verification_mcp; diff --git a/codex-rs/app-server/tests/suite/v2/plugin_share.rs b/codex-rs/app-server/tests/suite/v2/plugin_share.rs index d64c52ca23..06ca0a5874 100644 --- a/codex-rs/app-server/tests/suite/v2/plugin_share.rs +++ b/codex-rs/app-server/tests/suite/v2/plugin_share.rs @@ -389,7 +389,8 @@ plugin_sharing = false .received_requests() .await .expect("wiremock should record requests") - .is_empty() + .iter() + .all(|request| request.url.path() == "/backend-api/wham/accounts/check") ); Ok(()) } diff --git a/codex-rs/app-server/tests/suite/v2/rate_limit_reset_credits.rs b/codex-rs/app-server/tests/suite/v2/rate_limit_reset_credits.rs index 288a6154f6..96ba01d67a 100644 --- a/codex-rs/app-server/tests/suite/v2/rate_limit_reset_credits.rs +++ b/codex-rs/app-server/tests/suite/v2/rate_limit_reset_credits.rs @@ -228,6 +228,16 @@ async fn consume_account_rate_limit_reset_credit_surfaces_backend_failure() -> R #[tokio::test] async fn consume_timeout_releases_account_auth_queue() -> Result<()> { let (codex_home, server) = chatgpt_test_context().await?; + Mock::given(method("GET")) + .and(path("/api/codex/accounts/check")) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"accounts": [{ + "id": "account-123", "workspace_backend_origin": "https://chatgpt.com", + "account_routing_override": "NO_CONSTRAINT" + }]})), + ) + .mount(&server) + .await; Mock::given(method("POST")) .and(path("/api/codex/rate-limit-reset-credits/consume")) .respond_with( diff --git a/codex-rs/app-server/tests/suite/v2/workspace_routing.rs b/codex-rs/app-server/tests/suite/v2/workspace_routing.rs new file mode 100644 index 0000000000..d7a22d40ed --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/workspace_routing.rs @@ -0,0 +1,458 @@ +//! Selected-workspace discovery through the public account API. + +use anyhow::Result; +use app_test_support::ChatGptAuthFixture; +use app_test_support::ChatGptIdTokenClaims; +use app_test_support::TestAppServer; +use app_test_support::encode_id_token; +use app_test_support::write_chatgpt_auth; +use codex_app_server_protocol::ClientInfo; +use codex_app_server_protocol::GetAccountParams; +use codex_app_server_protocol::LoginAccountResponse; +use codex_app_server_protocol::LogoutAccountResponse; +use codex_app_server_protocol::RequestId; +use codex_config::types::AuthCredentialsStoreMode; +use pretty_assertions::assert_eq; +use serde_json::Value; +use serde_json::json; +use std::sync::Arc; +use std::time::Duration; +use tempfile::TempDir; +use tokio::sync::Notify; +use tokio::time::timeout; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::header; +use wiremock::matchers::method; +use wiremock::matchers::path; + +const READ_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 60); + +async fn config(home: &TempDir, backend: &MockServer) -> Result<()> { + std::fs::write( + home.path().join("config.toml"), + format!( + "chatgpt_base_url = '{}/backend-api/'\ncli_auth_credentials_store = 'file'\n", + backend.uri() + ), + )?; + Mock::given(method("GET")).and(path("/backend-api/wham/config/bundle")) + .respond_with(|request: &wiremock::Request| { + let account_id = request.headers.get("chatgpt-account-id").expect("workspace header") + .to_str().expect("workspace header text"); + ResponseTemplate::new(200).set_body_json(json!({"requirements_toml": {"enterprise_managed": [{ + "id": "workspace-policy", "name": "Workspace policy", + "contents": format!("[application.network.domains]\n'{account_id}.example' = 'allow'"), + }]}})) + }).mount(backend).await; + Ok(()) +} + +async fn start(home: &TempDir) -> Result { + TestAppServer::builder() + .with_codex_home(home.path()) + .with_env_overrides(&[("OPENAI_API_KEY", None)]) + .build_initialized_with_timeout(READ_TIMEOUT) + .await +} + +async fn read(server: &mut TestAppServer) -> Result { + let request = server + .send_get_account_request(GetAccountParams { + refresh_token: false, + }) + .await?; + timeout(READ_TIMEOUT, server.read_response(request)).await? +} + +async fn login(server: &mut TestAppServer, account: &str) -> Result<()> { + let token = encode_id_token( + &ChatGptIdTokenClaims::new() + .email("user@example.com") + .plan_type("enterprise") + .chatgpt_account_id(account), + )?; + let request = server + .send_chatgpt_auth_tokens_login_request(token, account.into(), Some("enterprise".into())) + .await?; + let _: LoginAccountResponse = timeout(READ_TIMEOUT, server.read_response(request)).await??; + Ok(()) +} + +fn routing(account: &str, origin: &str, routing: &str) -> Value { + json!({"chatgptAccountId": account, "backendOrigin": origin, "accountRoutingOverride": routing}) +} + +#[tokio::test] +async fn saved_workspace_is_discovered_once_and_not_the_default_account() -> Result<()> { + let backend = MockServer::start().await; + Mock::given(method("GET")).and(path("/backend-api/wham/accounts/check")) + .and(header("chatgpt-account-id", "selected")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "accounts": [ + {"id": "other", "workspace_backend_origin": "https://other.example", "account_routing_override": "us"}, + {"id": "selected", "workspace_backend_origin": "https://gov.chatgpt.com", "account_routing_override": "NO_CONSTRAINT"} + ], "default_account_id": "other" + }))).expect(1).mount(&backend).await; + let home = TempDir::new()?; + config(&home, &backend).await?; + write_chatgpt_auth( + home.path(), + ChatGptAuthFixture::new("token") + .account_id("selected") + .email("user@example.com") + .plan_type("pro"), + AuthCredentialsStoreMode::File, + )?; + let mut server = start(&home).await?; + let notification = timeout( + READ_TIMEOUT, + server.read_stream_until_notification_message("account/updated"), + ) + .await??; + assert_eq!( + notification.params, + Some(json!({"authMode": "chatgpt", "planType": "pro"})) + ); + let expected = json!({ + "account": {"type": "chatgpt", "email": "user@example.com", "planType": "pro"}, + "requiresOpenaiAuth": true, "workspaceRouting": routing("selected", "https://gov.chatgpt.com", "NO_CONSTRAINT") + }); + assert_eq!(read(&mut server).await?, expected); + assert_eq!(read(&mut server).await?, expected); + backend.verify().await; + Ok(()) +} + +#[tokio::test] +async fn saved_chatgpt_login_without_selected_workspace_preserves_account() -> Result<()> { + let backend = MockServer::start().await; + Mock::given(path("/backend-api/wham/accounts/check")) + .respond_with(ResponseTemplate::new(500)) + .expect(0) + .mount(&backend) + .await; + let home = TempDir::new()?; + config(&home, &backend).await?; + write_chatgpt_auth( + home.path(), + ChatGptAuthFixture::new("token").plan_type("pro"), + AuthCredentialsStoreMode::File, + )?; + let mut server = start(&home).await?; + assert_eq!( + read(&mut server).await?, + json!({"account": {"type": "chatgpt", "email": null, "planType": "pro"}, + "requiresOpenaiAuth": true, "workspaceRouting": null}) + ); + backend.verify().await; + Ok(()) +} + +#[tokio::test] +async fn stable_clients_do_not_treat_failed_workspace_discovery_as_unrestricted() -> Result<()> { + let backend = MockServer::start().await; + Mock::given(path("/backend-api/wham/accounts/check")) + .respond_with(ResponseTemplate::new(503)) + .mount(&backend) + .await; + let home = TempDir::new()?; + config(&home, &backend).await?; + write_chatgpt_auth( + home.path(), + ChatGptAuthFixture::new("token") + .account_id("selected") + .plan_type("pro"), + AuthCredentialsStoreMode::File, + )?; + let mut server = TestAppServer::builder() + .with_codex_home(home.path()) + .with_env_overrides(&[("OPENAI_API_KEY", None)]) + .build() + .await?; + timeout( + READ_TIMEOUT, + server.initialize_with_capabilities( + ClientInfo { + name: "routing-stable-client".into(), + title: None, + version: "1.0.0".into(), + }, + /*capabilities*/ None, + ), + ) + .await??; + let request = server + .send_get_account_request(GetAccountParams { + refresh_token: false, + }) + .await?; + let error = timeout( + READ_TIMEOUT, + server.read_stream_until_error_message(RequestId::Integer(request)), + ) + .await??; + assert_eq!(error.error.message, "workspace routing discovery failed"); + backend.reset().await; + Mock::given(path("/backend-api/wham/accounts/check")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"accounts": [{ + "id": "selected", "workspace_backend_origin": "https://chatgpt.com", "account_routing_override": "NO_CONSTRAINT" + }]}))) + .mount(&backend) + .await; + assert_eq!( + read(&mut server).await?, + json!({"account": {"type": "chatgpt", "email": null, "planType": "pro"}, + "requiresOpenaiAuth": true, "workspaceRouting": routing("selected", "https://chatgpt.com", "NO_CONSTRAINT")}) + ); + Ok(()) +} + +#[tokio::test] +async fn login_and_workspace_switch_notify_after_routing_is_ready_then_logout_clears_it() +-> Result<()> { + let backend = MockServer::start().await; + Mock::given(method("GET")).and(path("/backend-api/wham/accounts/check")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"accounts": [ + {"id": "first", "workspace_backend_origin": "https://chatgpt.com", "account_routing_override": "us_cr"}, + {"id": "second", "workspace_backend_origin": "https://gov.chatgpt.com", "account_routing_override": "us"} + ]}))).expect(2).mount(&backend).await; + let home = TempDir::new()?; + config(&home, &backend).await?; + let mut server = start(&home).await?; + for (account, origin, value) in [ + ("first", "https://chatgpt.com", "us_cr"), + ("second", "https://gov.chatgpt.com", "us"), + ] { + login(&mut server, account).await?; + let notification = timeout( + READ_TIMEOUT, + server.read_stream_until_notification_message("account/updated"), + ) + .await??; + assert_eq!( + notification.params, + Some(json!({"authMode": "chatgptAuthTokens", "planType": "enterprise"})) + ); + assert_eq!( + read(&mut server).await?["workspaceRouting"], + routing(account, origin, value) + ); + let request = server.send_config_requirements_read_request().await?; + let requirements: Value = timeout(READ_TIMEOUT, server.read_response(request)).await??; + assert_eq!( + requirements["requirements"]["application"], + json!({"network": { + "enabled": true, "domains": {format!("{account}.example"): "allow"} + }}) + ); + } + let request = server.send_logout_account_request().await?; + let _: LogoutAccountResponse = timeout(READ_TIMEOUT, server.read_response(request)).await??; + assert_eq!( + read(&mut server).await?, + json!({"account": null, "requiresOpenaiAuth": true, "workspaceRouting": null}) + ); + backend.verify().await; + Ok(()) +} + +#[tokio::test] +async fn unavailable_or_malformed_discovery_never_returns_unrestricted_success() -> Result<()> { + for response in [ + ResponseTemplate::new(503), + ResponseTemplate::new(200).set_body_json(json!({"accounts": [{"id": "selected"}]})), + ResponseTemplate::new(200).set_body_json(json!({"accounts": [{"id": "selected", "account_routing_override": "us_cr"}]})), + ResponseTemplate::new(200).set_body_json(json!({"accounts": [{"id": "selected", "workspace_backend_origin": "NO_CONSTRAINT"}]})), + ResponseTemplate::new(200).set_body_json(json!({"accounts": [{"id": "selected", "workspace_backend_origin": null, "account_routing_override": null}]})), + ResponseTemplate::new(200).set_body_json(json!({"accounts": [{"id": "selected", "workspace_backend_origin": 7, "account_routing_override": "us_cr"}]})), + ResponseTemplate::new(200).set_body_json(json!({"accounts": [{"id": "selected", "workspace_backend_origin": "NO_CONSTRAINT", "account_routing_override": "unknown"}]})), + ResponseTemplate::new(200).set_body_json(json!({"accounts": []})), + ] { + let backend = MockServer::start().await; + Mock::given(method("GET")).and(path("/backend-api/wham/accounts/check")) + .respond_with(response).mount(&backend).await; + let home = TempDir::new()?; + config(&home, &backend).await?; + write_chatgpt_auth(home.path(), ChatGptAuthFixture::new("token").account_id("selected"), AuthCredentialsStoreMode::File)?; + let mut server = start(&home).await?; + let request = server.send_get_account_request(GetAccountParams { refresh_token: false }).await?; + let error = timeout(READ_TIMEOUT, server.read_stream_until_error_message(RequestId::Integer(request))).await??; + assert!(error.error.message.contains("routing"), "{error:?}"); + backend.reset().await; + Mock::given(method("GET")).and(path("/backend-api/wham/accounts/check")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"accounts": [{ + "id": "selected", "workspace_backend_origin": "https://chatgpt.com", "account_routing_override": "NO_CONSTRAINT" + }]}))).mount(&backend).await; + assert_eq!(read(&mut server).await?["workspaceRouting"], routing("selected", "https://chatgpt.com", "NO_CONSTRAINT")); + } + Ok(()) +} + +#[tokio::test] +async fn late_startup_discovery_is_discarded_on_workspace_switch_and_logout() -> Result<()> { + for next_account in [Some("second"), Some("first"), None] { + let backend = MockServer::start().await; + let started = Arc::new(Notify::new()); + let request_started = Arc::clone(&started); + Mock::given(method("GET")).and(path("/backend-api/wham/accounts/check")) + .and(header("chatgpt-account-id", "first")) + .and(header("authorization", "Bearer token")) + .respond_with(move |_: &wiremock::Request| { + request_started.notify_one(); + ResponseTemplate::new(200).set_delay(Duration::from_secs(/*secs*/ 2)) + .set_body_json(json!({"accounts": [{"id": "first", "workspace_backend_origin": "https://old.example", "account_routing_override": "us_cr"}]})) + }).with_priority(1).mount(&backend).await; + Mock::given(method("GET")).and(path("/backend-api/wham/accounts/check")) + .and(header("chatgpt-account-id", next_account.unwrap_or("second"))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"accounts": [ + {"id": next_account.unwrap_or("second"), "workspace_backend_origin": "https://new.example", "account_routing_override": "us"} + ]}))).mount(&backend).await; + let home = TempDir::new()?; + config(&home, &backend).await?; + write_chatgpt_auth( + home.path(), + ChatGptAuthFixture::new("token").account_id("first"), + AuthCredentialsStoreMode::File, + )?; + let mut server = start(&home).await?; + timeout(READ_TIMEOUT, started.notified()).await?; + if let Some(account) = next_account { + login(&mut server, account).await?; + } else { + let request = server.send_logout_account_request().await?; + let _: LogoutAccountResponse = + timeout(READ_TIMEOUT, server.read_response(request)).await??; + } + let expected = next_account + .map(|account| routing(account, "https://new.example", "us")) + .unwrap_or(Value::Null); + assert_eq!(read(&mut server).await?["workspaceRouting"], expected); + tokio::time::sleep(Duration::from_secs(/*secs*/ 2)).await; + assert_eq!(read(&mut server).await?["workspaceRouting"], expected); + } + Ok(()) +} + +#[tokio::test] +async fn api_only_login_does_not_discover_chatgpt_routing() -> Result<()> { + let backend = MockServer::start().await; + Mock::given(path("/backend-api/wham/accounts/check")) + .respond_with(ResponseTemplate::new(500)) + .expect(0) + .mount(&backend) + .await; + let home = TempDir::new()?; + config(&home, &backend).await?; + std::fs::write( + home.path().join("requirements.toml"), + "allowed_login_methods = ['api']", + )?; + let mut server = start(&home).await?; + let request = server.send_login_account_api_key_request("sk-test").await?; + let _: LoginAccountResponse = timeout(READ_TIMEOUT, server.read_response(request)).await??; + assert_eq!( + read(&mut server).await?, + json!({"account": {"type": "apiKey"}, "requiresOpenaiAuth": true, "workspaceRouting": null}) + ); + backend.verify().await; + Ok(()) +} + +#[tokio::test] +async fn failed_workspace_requirements_do_not_fall_back_to_startup_config() -> Result<()> { + let backend = MockServer::start().await; + let home = TempDir::new()?; + config(&home, &backend).await?; + let mut server = start(&home).await?; + backend.reset().await; + Mock::given(path("/backend-api/wham/config/bundle")) + .respond_with(ResponseTemplate::new(200).set_body_json( + json!({"requirements_toml": {"enterprise_managed": [{ + "id": "invalid", "name": "Invalid requirements", "contents": "application = true", + }]}}), + )) + .mount(&backend) + .await; + Mock::given(path("/backend-api/wham/accounts/check")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&backend) + .await; + login(&mut server, "selected").await?; + let notification = timeout( + READ_TIMEOUT, + server.read_stream_until_notification_message("account/login/completed"), + ) + .await??; + assert_eq!(notification.params.unwrap()["success"], false); + let request = server + .send_get_account_request(GetAccountParams { + refresh_token: false, + }) + .await?; + let error = timeout( + READ_TIMEOUT, + server.read_stream_until_error_message(RequestId::Integer(request)), + ) + .await??; + assert_eq!(error.error.message, "failed to load workspace requirements"); + backend.verify().await; + Ok(()) +} + +#[tokio::test] +async fn backend_config_changed_during_discovery_is_retried_with_fresh_config() -> Result<()> { + let old_backend = MockServer::start().await; + let new_backend = MockServer::start().await; + let started = Arc::new(Notify::new()); + let request_started = Arc::clone(&started); + Mock::given(path("/backend-api/wham/accounts/check")) + .respond_with(move |_: &wiremock::Request| { + request_started.notify_one(); + ResponseTemplate::new(200) + .set_delay(Duration::from_secs(/*secs*/ 2)) + .set_body_json(json!({"accounts": [{ + "id": "selected", "workspace_backend_origin": "https://old.example", + "account_routing_override": "us_cr" + }]})) + }) + .mount(&old_backend) + .await; + Mock::given(path("/backend-api/wham/accounts/check")) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"accounts": [{ + "id": "selected", "workspace_backend_origin": "https://new.example", + "account_routing_override": "us" + }]})), + ) + .expect(1) + .mount(&new_backend) + .await; + let home = TempDir::new()?; + config(&home, &old_backend).await?; + let mut server = start(&home).await?; + login(&mut server, "selected").await?; + timeout(READ_TIMEOUT, started.notified()).await?; + config(&home, &new_backend).await?; + let completed = timeout( + READ_TIMEOUT, + server.read_stream_until_notification_message("account/login/completed"), + ) + .await??; + assert_eq!( + completed.params, + Some(json!({ + "loginId": null, "success": false, + "error": "configuration changed during workspace routing discovery; retry account/read", + "onboardingEntrypoint": null, + })) + ); + assert_eq!( + read(&mut server).await?["workspaceRouting"], + routing("selected", "https://new.example", "us") + ); + new_backend.verify().await; + Ok(()) +} diff --git a/codex-rs/backend-client/src/types.rs b/codex-rs/backend-client/src/types.rs index 791f9c4ace..6b45c6e20c 100644 --- a/codex-rs/backend-client/src/types.rs +++ b/codex-rs/backend-client/src/types.rs @@ -136,6 +136,8 @@ pub struct AccountsCheckResponse { #[derive(Clone, Debug, Deserialize)] pub struct AccountEntry { pub id: String, + pub workspace_backend_origin: Option, + pub account_routing_override: Option, #[serde(default)] pub name: Option, #[serde(default)] @@ -198,6 +200,8 @@ impl<'de> Deserialize<'de> for AccountsCheckResponse { let account = accounts.remove(account_id)?.account; Some(AccountEntry { id: account.account_id?, + workspace_backend_origin: None, + account_routing_override: None, name: account.name, profile_picture_url: account.profile_picture_url, structure: account.structure, diff --git a/codex-rs/tui/src/app_server_session.rs b/codex-rs/tui/src/app_server_session.rs index cc576eb791..a50000f024 100644 --- a/codex-rs/tui/src/app_server_session.rs +++ b/codex-rs/tui/src/app_server_session.rs @@ -2562,6 +2562,7 @@ mod tests { let mut app_server = crate::start_embedded_app_server_for_picker(&config).await?; let next_request_id = app_server.next_request_id; let account = GetAccountResponse { + workspace_routing: None, account: Some(Account::Chatgpt { email: Some("teammate@openai.com".to_string()), plan_type: codex_protocol::account::PlanType::Plus, diff --git a/sdk/python/src/openai_codex/generated/v2_all.py b/sdk/python/src/openai_codex/generated/v2_all.py index c0ca73d30e..621298ce6c 100644 --- a/sdk/python/src/openai_codex/generated/v2_all.py +++ b/sdk/python/src/openai_codex/generated/v2_all.py @@ -43,6 +43,12 @@ class AmazonBedrockAccount(BaseModel): ] = False +class AccountRoutingOverride(Enum): + no_constraint = "NO_CONSTRAINT" + us = "us" + us_cr = "us_cr" + + class AccountTokenUsageDailyBucket(BaseModel): model_config = ConfigDict( populate_by_name=True, @@ -6465,6 +6471,17 @@ class WorkspaceMessageType(Enum): unknown = "unknown" +class WorkspaceRouting(BaseModel): + model_config = ConfigDict( + populate_by_name=True, + ) + account_routing_override: Annotated[ + AccountRoutingOverride, Field(alias="accountRoutingOverride") + ] + backend_origin: Annotated[str, Field(alias="backendOrigin")] + chatgpt_account_id: Annotated[str, Field(alias="chatgptAccountId")] + + class WriteStatus(Enum): ok = "ok" ok_overridden = "okOverridden"