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 3179934928..3383756a7c 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 @@ -6919,6 +6919,103 @@ ], "type": "string" }, + "AccountSession": { + "properties": { + "displayName": { + "type": [ + "string", + "null" + ] + }, + "email": { + "type": [ + "string", + "null" + ] + }, + "imageUrl": { + "type": [ + "string", + "null" + ] + }, + "isActive": { + "type": "boolean" + }, + "lastUsedAt": { + "format": "int64", + "type": "integer" + }, + "selectedWorkspaceAccountId": { + "type": [ + "string", + "null" + ] + }, + "sessionId": { + "type": "string" + }, + "userId": { + "type": [ + "string", + "null" + ] + }, + "workspaces": { + "items": { + "$ref": "#/definitions/v2/AccountSessionWorkspace" + }, + "type": "array" + } + }, + "required": [ + "isActive", + "lastUsedAt", + "sessionId", + "workspaces" + ], + "type": "object" + }, + "AccountSessionWorkspace": { + "properties": { + "accountId": { + "type": "string" + }, + "imageUrl": { + "type": [ + "string", + "null" + ] + }, + "kind": { + "anyOf": [ + { + "$ref": "#/definitions/v2/AccountSessionWorkspaceKind" + }, + { + "type": "null" + } + ] + }, + "name": { + "type": [ + "string", + "null" + ] + } + }, + "required": [ + "accountId" + ], + "type": "object" + }, + "AccountSessionWorkspaceKind": { + "enum": [ + "personal", + "workspace" + ], + "type": "string" + }, "AccountTokenUsageDailyBucket": { "properties": { "startDate": { 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 bce1d13765..e1355e2bf6 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 @@ -154,6 +154,103 @@ ], "type": "string" }, + "AccountSession": { + "properties": { + "displayName": { + "type": [ + "string", + "null" + ] + }, + "email": { + "type": [ + "string", + "null" + ] + }, + "imageUrl": { + "type": [ + "string", + "null" + ] + }, + "isActive": { + "type": "boolean" + }, + "lastUsedAt": { + "format": "int64", + "type": "integer" + }, + "selectedWorkspaceAccountId": { + "type": [ + "string", + "null" + ] + }, + "sessionId": { + "type": "string" + }, + "userId": { + "type": [ + "string", + "null" + ] + }, + "workspaces": { + "items": { + "$ref": "#/definitions/AccountSessionWorkspace" + }, + "type": "array" + } + }, + "required": [ + "isActive", + "lastUsedAt", + "sessionId", + "workspaces" + ], + "type": "object" + }, + "AccountSessionWorkspace": { + "properties": { + "accountId": { + "type": "string" + }, + "imageUrl": { + "type": [ + "string", + "null" + ] + }, + "kind": { + "anyOf": [ + { + "$ref": "#/definitions/AccountSessionWorkspaceKind" + }, + { + "type": "null" + } + ] + }, + "name": { + "type": [ + "string", + "null" + ] + } + }, + "required": [ + "accountId" + ], + "type": "object" + }, + "AccountSessionWorkspaceKind": { + "enum": [ + "personal", + "workspace" + ], + "type": "string" + }, "AccountTokenUsageDailyBucket": { "properties": { "startDate": { 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 fec1d795b5..eb8ee98152 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 0ff4c59a10..2b8e25fe9e 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/AccountSession.ts b/codex-rs/app-server-protocol/schema/typescript/v2/AccountSession.ts new file mode 100644 index 0000000000..71a36b0d75 --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/AccountSession.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 { AccountSessionWorkspace } from "./AccountSessionWorkspace"; + +export type AccountSession = { sessionId: string, email: string | null, userId: string | null, displayName: string | null, imageUrl: string | null, lastUsedAt: bigint, isActive: boolean, selectedWorkspaceAccountId: string | null, workspaces: Array, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/AccountSessionWorkspace.ts b/codex-rs/app-server-protocol/schema/typescript/v2/AccountSessionWorkspace.ts new file mode 100644 index 0000000000..b55d2a9ea7 --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/AccountSessionWorkspace.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 { AccountSessionWorkspaceKind } from "./AccountSessionWorkspaceKind"; + +export type AccountSessionWorkspace = { accountId: string, name: string | null, imageUrl: string | null, kind: AccountSessionWorkspaceKind | null, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/AccountSessionWorkspaceKind.ts b/codex-rs/app-server-protocol/schema/typescript/v2/AccountSessionWorkspaceKind.ts new file mode 100644 index 0000000000..1ece70db35 --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/AccountSessionWorkspaceKind.ts @@ -0,0 +1,5 @@ +// 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. + +export type AccountSessionWorkspaceKind = "personal" | "workspace"; 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 e8b3866f29..c36b59f3e3 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/index.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/index.ts @@ -4,6 +4,9 @@ export type { Account } from "./Account"; export type { AccountLoginCompletedNotification } from "./AccountLoginCompletedNotification"; export type { AccountRateLimitsUpdatedNotification } from "./AccountRateLimitsUpdatedNotification"; export type { AccountRoutingOverride } from "./AccountRoutingOverride"; +export type { AccountSession } from "./AccountSession"; +export type { AccountSessionWorkspace } from "./AccountSessionWorkspace"; +export type { AccountSessionWorkspaceKind } from "./AccountSessionWorkspaceKind"; export type { AccountTokenUsageDailyBucket } from "./AccountTokenUsageDailyBucket"; export type { AccountTokenUsageSummary } from "./AccountTokenUsageSummary"; export type { AccountUpdatedNotification } from "./AccountUpdatedNotification"; diff --git a/codex-rs/app-server-protocol/src/protocol/common.rs b/codex-rs/app-server-protocol/src/protocol/common.rs index 1afa3c6b42..00ac5b6a4c 100644 --- a/codex-rs/app-server-protocol/src/protocol/common.rs +++ b/codex-rs/app-server-protocol/src/protocol/common.rs @@ -1358,6 +1358,22 @@ client_request_definitions! { response: v2::CancelLoginAccountResponse, }, + #[experimental("accountSession/list")] + AccountSessionsList => "accountSession/list" { + params: v2::AccountSessionsListParams, + serialization: global("account-auth"), + manual_payload_conversion: manual, + response: v2::AccountSessionsResponse, + }, + + #[experimental("accountSession/switch")] + AccountSessionsSwitch => "accountSession/switch" { + params: v2::AccountSessionsSwitchParams, + serialization: global("account-auth"), + manual_payload_conversion: manual, + response: v2::AccountSessionsSwitchResponse, + }, + LogoutAccount => "account/logout" { params: #[ts(type = "undefined")] #[serde(skip_serializing_if = "Option::is_none")] Option<()>, serialization: global("account-auth"), 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 f745f1f985..a5a5009d12 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/account.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/account.rs @@ -227,8 +227,26 @@ pub struct AccountSessionsLogoutParams { #[serde(rename_all = "camelCase")] #[ts(export_to = "v2/")] pub struct AccountSessionsSwitchParams { + /// Opaque active-login identifier returned by accountSession/list. pub session_id: String, + /// Workspace to select within that login; this is not a user identifier. pub account_id: String, + /// Allow Core to start browser sign-in if this session needs SSO for the workspace. + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + #[ts(optional, as = "Option")] + pub allow_browser_login: bool, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema, TS)] +#[serde(untagged, rename_all_fields = "camelCase")] +#[ts(untagged, rename_all_fields = "camelCase", export_to = "v2/")] +pub enum AccountSessionsSwitchResponse { + Switched { + #[serde(flatten)] + sessions: AccountSessionsResponse, + }, + /// Open the URL and await account/login/completed for this login ID. + BrowserLogin { login_id: String, auth_url: String }, } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema, TS)] diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 26a385cec5..a557a984d9 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -647,6 +647,29 @@ Omitted or null cursors preserve normal first-page behavior. Continue anchored pages with the returned opaque string `nextCursor`; response fields and `backwardsCursor` semantics are unchanged. +## Workspace switching (experimental) + +For a managed ChatGPT login, `accountSession/list` returns an opaque +`sessionId` for that login and its selected workspace. Set +`refreshWorkspaceMetadata` to fetch the user's memberships. + +Call `accountSession/switch` with that `sessionId` and the target workspace's +`accountId`. Core checks the membership, stops loaded threads, exchanges the +session token and validates the new workspace's routing. Clients receive +`thread/closed` for unloaded threads and `account/updated` after a successful +switch. Selecting the current workspace leaves its threads running. + +The server remembers the selected workspace on refresh. If routing validation +fails, Core attempts to restore the previous workspace while the same credentials +are still current. + +Clients that can open a browser may pass `allowBrowserLogin: true`. If the +workspace requires SSO, the response contains `loginId` and `authUrl`. Open +that URL and wait for `account/login/completed`. The browser can reuse an +existing SSO session. Core accepts the callback only for the requesting login +and target workspace. Cancel through `account/login/cancel` to keep the current +workspace. Government workspaces require a separate sign-in. + # GCP provider account The experimental `account/gcp/setup` method configures the built-in Vertex AI provider. diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index c664fac4b2..81d3de09ef 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -1842,6 +1842,16 @@ impl MessageProcessor { .await .map(|response| Some(response.into())) } + ClientRequest::AccountSessionsList { params, .. } => self + .account_processor + .list_account_sessions(params) + .await + .map(|response| Some(ClientResponsePayload::AccountSessionsList(response))), + ClientRequest::AccountSessionsSwitch { params, .. } => self + .account_processor + .switch_account_session(params, &self.thread_processor, &self.turn_admission) + .await + .map(|response| Some(ClientResponsePayload::AccountSessionsSwitch(response))), ClientRequest::LogoutAccount { .. } => { self.account_processor .logout_account(request_id.clone()) diff --git a/codex-rs/app-server/src/request_processors/account_processor.rs b/codex-rs/app-server/src/request_processors/account_processor.rs index 2aaa283f78..80f88efb7f 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -24,6 +24,7 @@ mod gateway_oauth; mod gcp_setup; mod rate_limit_resets; mod workspace_routing; +mod workspace_switching; pub(super) use enterprise_login::EnterpriseLoginCompletion; pub(super) use enterprise_login::EnterpriseLoginTarget; @@ -690,6 +691,15 @@ impl AccountRequestProcessor { let opts = self .login_chatgpt_common(codex_streamlined_login, login_success_page) .await?; + let (login_id, auth_url) = self.start_browser_login(opts, /*pause*/ None).await?; + Ok(LoginAccountResponse::Chatgpt { login_id, auth_url }) + } + + async fn start_browser_login( + &self, + opts: LoginServerOptions, + pause: Option, + ) -> Result<(String, String), JSONRPCErrorError> { let server = run_login_server(opts) .map_err(|err| internal_error(format!("failed to start login server: {err}")))?; let login_id = Uuid::new_v4(); @@ -711,6 +721,8 @@ impl AccountRequestProcessor { let active_login = self.active_login.clone(); let auth_url = server.auth_url.clone(); let task = async move { + // Workspace credentials may change in the callback. Resume turns after caches and clients are updated. + let _pause = pause; let (success, error_msg, onboarding_entrypoint) = match tokio::time::timeout( LOGIN_CHATGPT_TIMEOUT, server.block_until_done_with_callback_result(), @@ -750,10 +762,7 @@ impl AccountRequestProcessor { }; tokio::spawn(AuthStorageOriginator::current().scope(task)); - Ok(LoginAccountResponse::Chatgpt { - login_id: login_id.to_string(), - auth_url, - }) + Ok((login_id.to_string(), auth_url)) } async fn login_chatgpt_device_code_v2(&self, request_id: ConnectionRequestId) { diff --git a/codex-rs/app-server/src/request_processors/account_processor/workspace_switching.rs b/codex-rs/app-server/src/request_processors/account_processor/workspace_switching.rs new file mode 100644 index 0000000000..c518cf371b --- /dev/null +++ b/codex-rs/app-server/src/request_processors/account_processor/workspace_switching.rs @@ -0,0 +1,255 @@ +//! Lists the active managed login and switches its workspace through the account-session API. + +use super::*; +use codex_app_server_protocol::AccountSession; +use codex_app_server_protocol::AccountSessionWorkspace; +use codex_app_server_protocol::AccountSessionWorkspaceKind; +use codex_app_server_protocol::AccountSessionsListParams; +use codex_app_server_protocol::AccountSessionsResponse; +use codex_app_server_protocol::AccountSessionsSwitchParams; +use codex_app_server_protocol::AccountSessionsSwitchResponse; + +impl AccountRequestProcessor { + pub(crate) async fn list_account_sessions( + &self, + params: AccountSessionsListParams, + ) -> Result { + let auth = if params.refresh_workspace_metadata { + self.auth_manager.auth().await + } else { + self.auth_manager.auth_cached() + } + .filter(|_| !self.auth_manager.has_external_auth()); + let session_id = auth + .as_ref() + .and_then(CodexAuth::managed_account_session_id); + let (Some(auth), Some(session_id)) = (auth, session_id) else { + return Ok(AccountSessionsResponse { + active_session_id: None, + sessions: vec![], + }); + }; + let account_id = auth.get_account_id(); + let mut workspaces = account_id + .iter() + .map(|id| AccountSessionWorkspace { + account_id: id.clone(), + name: None, + image_url: None, + kind: None, + }) + .collect(); + if params.refresh_workspace_metadata { + let accounts = self.workspace_memberships(&auth).await?; + workspaces = accounts + .accounts + .into_iter() + .map(|account| AccountSessionWorkspace { + account_id: account.id, + name: account.name, + image_url: account.profile_picture_url, + kind: match account.structure.as_str() { + "personal" => Some(AccountSessionWorkspaceKind::Personal), + "workspace" => Some(AccountSessionWorkspaceKind::Workspace), + _ => None, + }, + }) + .collect::>(); + workspaces.sort_by_key(|workspace| { + accounts + .account_ordering + .iter() + .position(|id| id == &workspace.account_id) + .unwrap_or(usize::MAX) + }); + let current = self.auth_manager.auth_cached(); + if self.auth_manager.has_external_auth() + || current + .as_ref() + .and_then(CodexAuth::managed_account_session_id) + .as_deref() + != Some(&session_id) + || current.and_then(|auth| auth.get_account_id()) != account_id + { + return Err(invalid_request( + "Authentication changed while listing account sessions", + )); + } + } + Ok(AccountSessionsResponse { + active_session_id: Some(session_id.clone()), + sessions: vec![AccountSession { + session_id, + email: auth.get_account_email(), + user_id: auth.get_chatgpt_user_id(), + display_name: None, + image_url: None, + last_used_at: chrono::Utc::now().timestamp(), + is_active: true, + selected_workspace_account_id: account_id, + workspaces, + }], + }) + } + + async fn workspace_memberships( + &self, + auth: &CodexAuth, + ) -> Result { + let client = BackendClient::from_auth( + self.config.chatgpt_base_url.clone(), + auth, + self.config.http_client_factory(), + ); + tokio::time::timeout( + Duration::from_secs(/*secs*/ 30), + client.get_accounts_check(), + ) + .await + .map_err(|_| internal_error("Workspace metadata request timed out"))? + .map_err(|_| internal_error("Failed to fetch workspace metadata")) + } + + pub(crate) async fn switch_account_session( + &self, + params: AccountSessionsSwitchParams, + threads: &crate::request_processors::ThreadRequestProcessor, + admission: &crate::turn_admission::TurnAdmission, + ) -> Result { + if self.active_login.lock().await.is_some() { + return Err(invalid_request( + "Finish or cancel the current login before switching workspaces", + )); + } + self.auth_manager + .validate_workspace_switch(¶ms.session_id, ¶ms.account_id) + .map_err(|error| invalid_request(error.to_string()))?; + if self + .auth_manager + .auth_cached() + .and_then(|auth| auth.get_account_id()) + .as_deref() + == Some(¶ms.account_id) + { + return self.account_session_switch_response(¶ms).await; + } + let auth = self + .auth_manager + .auth_cached() + .ok_or_else(|| invalid_request("Authentication is unavailable"))?; + let memberships = self.workspace_memberships(&auth).await?; + let mut matches = memberships + .accounts + .iter() + .filter(|account| account.id == params.account_id); + let Some(target) = matches.next() else { + return Err(invalid_request("The selected workspace is unavailable")); + }; + if matches.next().is_some() || target.is_deactivated { + return Err(invalid_request("The selected workspace is unavailable")); + } + // The exchange strips the FedRAMP hint from the server session's token templates. + if auth.is_fedramp_account() + || target.is_fedramp_compliant_workspace + || target.workspace_backend_origin.as_deref() == Some("https://gov.chatgpt.com") + { + return Err(invalid_request("Sign in to switch government workspaces")); + } + self.auth_manager + .validate_workspace_switch(¶ms.session_id, ¶ms.account_id) + .map_err(|error| invalid_request(error.to_string()))?; + if target.can_access_with_session == Some(false) { + if !params.allow_browser_login { + return Err(invalid_request( + "The selected workspace requires a new sign-in", + )); + } + let mut opts = self + .login_chatgpt_common( + /*codex_streamlined_login*/ true, + LoginSuccessPage::default(), + ) + .await?; + opts.workspace_login = Some( + auth.workspace_login( + ¶ms.session_id, + params.account_id.clone(), + target.sso_connection_name.clone(), + ) + .map_err(|error| invalid_request(error.to_string()))?, + ); + opts.forced_chatgpt_workspace_id = Some(vec![params.account_id]); + let pause = admission.pause().await?; + threads.unload_for_workspace_switch().await?; + let (login_id, auth_url) = self.start_browser_login(opts, Some(pause)).await?; + return Ok(AccountSessionsSwitchResponse::BrowserLogin { login_id, auth_url }); + } + let _pause = admission.pause().await?; + threads.unload_for_workspace_switch().await?; + let switched = self + .auth_manager + .switch_workspace(¶ms.session_id, ¶ms.account_id) + .await + .map_err(|error| invalid_request(error.to_string()))?; + let auth_changes = self.auth_manager.auth_change_state_receiver(); + let owner_generation = auth_changes.borrow().owner_generation; + self.reload_workspace_config().await; + // Routing includes fatal authorization and policy checks. Restore the + // previous login on any failure, unless another auth operation replaced it. + if let Err(error) = self.read_account(/*request*/ None).await { + let rollback = self.auth_manager.rollback_workspace_switch(switched).await; + self.reload_workspace_config().await; + rollback.map_err(|rollback| internal_error(format!("Workspace validation failed: {error}; could not restore the previous workspace: {rollback}")))?; + return Err(internal_error(format!( + "Workspace validation failed: {error}" + ))); + } + self.thread_manager.invalidate_mcp_runtimes().await; + self.model_catalog.invalidate().await; + Self::maybe_refresh_plugin_caches_for_current_config( + &self.config_manager, + &self.thread_manager, + self.auth_manager.auth_cached(), + ) + .await; + self.outgoing + .send_account_notification( + /*connection_id*/ None, + &auth_changes, + owner_generation, + AccountNotification::Updated(self.current_account_updated_notification()), + ) + .await; + self.account_session_switch_response(¶ms).await + } + + async fn reload_workspace_config(&self) { + self.config_manager + .update_managed_config_loader(self.auth_manager.clone(), &self.config); + self.config_manager + .sync_default_client_residency_requirement() + .await; + } + + async fn account_session_switch_response( + &self, + params: &AccountSessionsSwitchParams, + ) -> Result { + let response = self + .list_account_sessions(AccountSessionsListParams { + refresh_workspace_metadata: false, + }) + .await?; + if response.active_session_id.as_deref() != Some(¶ms.session_id) + || !response.sessions.iter().any(|session| { + session.session_id == params.session_id + && session.selected_workspace_account_id.as_deref() == Some(¶ms.account_id) + }) + { + return Err(invalid_request( + "Authentication changed during workspace switching", + )); + } + Ok(AccountSessionsSwitchResponse::Switched { sessions: response }) + } +} diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 64f8bd39e8..0e7c82e51c 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -4,6 +4,8 @@ mod daemon_continuation; #[path = "daemon_snapshot.rs"] mod daemon_snapshot; +mod workspace_switching; + use super::persisted_resume_settings::PersistedResumeSettings; use super::persisted_resume_settings::latest_owned_thread_settings; use super::persisted_resume_settings::latest_persisted_model_settings; diff --git a/codex-rs/app-server/src/request_processors/thread_processor/workspace_switching.rs b/codex-rs/app-server/src/request_processors/thread_processor/workspace_switching.rs new file mode 100644 index 0000000000..32533cb5ae --- /dev/null +++ b/codex-rs/app-server/src/request_processors/thread_processor/workspace_switching.rs @@ -0,0 +1,34 @@ +//! Unloads old-workspace runtimes while turn admission is paused by the account switch. +//! Resuming a thread builds fresh policy and transport state for the new workspace. + +use super::*; + +impl ThreadRequestProcessor { + pub(crate) async fn unload_for_workspace_switch(&self) -> Result<(), JSONRPCErrorError> { + let report = self + .thread_manager + .shutdown_all_threads_bounded(Duration::from_secs(/*secs*/ 10)) + .await; + for thread_id in report.completed { + self.finalize_thread_teardown(thread_id).await; + self.outgoing + .send_server_notification(ServerNotification::ThreadClosed( + codex_app_server_protocol::ThreadClosedNotification { + thread_id: thread_id.to_string(), + }, + )) + .await; + } + // A delegated child may have started while its parent was shutting down. + // Never replace credentials until every runtime has stopped. + if !report.submit_failed.is_empty() + || !report.timed_out.is_empty() + || !self.thread_manager.list_thread_ids().await.is_empty() + { + return Err(internal_error( + "Some threads are still running; retry the workspace switch", + )); + } + Ok(()) + } +} diff --git a/codex-rs/app-server/src/turn_admission.rs b/codex-rs/app-server/src/turn_admission.rs index 6c40044596..552875473b 100644 --- a/codex-rs/app-server/src/turn_admission.rs +++ b/codex-rs/app-server/src/turn_admission.rs @@ -6,11 +6,13 @@ use std::sync::Arc; use std::sync::Mutex; use tokio::sync::watch; +use crate::error_code::invalid_request; use crate::error_code::server_draining_error; #[derive(Debug, Default)] struct AdmissionState { closed: bool, + paused: bool, active: usize, } @@ -30,8 +32,32 @@ impl Default for TurnAdmission { } pub(crate) struct TurnPermit(TurnAdmission); +pub(crate) struct TurnAdmissionPause(TurnAdmission); impl TurnAdmission { + pub(crate) async fn pause(&self) -> Result { + let mut active = self.subscribe_active(); + { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state.closed || state.paused { + return Err(invalid_request("Runtime changes are already paused")); + } + state.paused = true; + } + let pause = TurnAdmissionPause(self.clone()); + tokio::time::timeout( + std::time::Duration::from_secs(/*secs*/ 10), + active.wait_for(|active| *active == 0), + ) + .await + .map_err(|_| invalid_request("Work is still starting; retry the workspace switch"))? + .map_err(|_| invalid_request("Turn admission is unavailable"))?; + Ok(pause) + } + pub(crate) fn begin_drain(&self) { self.state .lock() @@ -46,26 +72,38 @@ impl TurnAdmission { // Admit and close take the same short lock. The permit keeps shutdown from // finishing while an earlier request is still preparing or submitting work. pub(crate) fn admit(&self) -> Result { - self.try_admit().ok_or_else(server_draining_error) - } - - fn try_admit(&self) -> Option { let mut state = self .state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); if state.closed { - return None; + return Err(server_draining_error()); + } + if state.paused { + return Err(invalid_request( + "Workspace change in progress; retry when it completes", + )); } state.active += 1; self.active_tx.send_replace(state.active); - Some(TurnPermit(self.clone())) + Ok(TurnPermit(self.clone())) + } +} + +impl Drop for TurnAdmissionPause { + fn drop(&mut self) { + self.0 + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .paused = false; } } impl TurnStartAdmission for TurnAdmission { fn admit_turn_start(&self) -> Option> { - self.try_admit() + self.admit() + .ok() .map(|permit| Box::new(permit) as Box) } } diff --git a/codex-rs/app-server/src/turn_admission_tests.rs b/codex-rs/app-server/src/turn_admission_tests.rs index 472954a946..d36f15eea6 100644 --- a/codex-rs/app-server/src/turn_admission_tests.rs +++ b/codex-rs/app-server/src/turn_admission_tests.rs @@ -26,3 +26,26 @@ fn drain_rejects_new_work_without_waiting_for_admitted_work() { drop(automatic); assert_eq!(*active.borrow(), 0); } + +#[tokio::test] +async fn workspace_pause_waits_for_admitted_work_and_preserves_shutdown_drain() { + let admission = TurnAdmission::default(); + let in_flight = admission.admit().unwrap(); + let pending_pause = admission.pause(); + tokio::pin!(pending_pause); + tokio::select! { + biased; + _ = &mut pending_pause => panic!("pause completed before the existing request"), + () = async {} => {} + } + assert!(admission.admit().is_err()); + assert!(admission.admit_turn_start().is_none()); + drop(in_flight); + let pause = pending_pause.await.unwrap(); + drop(pause); + drop(admission.admit().unwrap()); + let pause = admission.pause().await.unwrap(); + admission.begin_drain(); + drop(pause); + assert!(admission.admit().is_err()); +} diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index 259438bee8..339438a5fd 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -3729,3 +3729,6 @@ async fn get_account_with_chatgpt_missing_plan_claim_returns_unknown() -> Result assert_eq!(received, expected); Ok(()) } + +#[path = "account_workspace_switching.rs"] +mod workspace_switching; diff --git a/codex-rs/app-server/tests/suite/v2/account_workspace_switching.rs b/codex-rs/app-server/tests/suite/v2/account_workspace_switching.rs new file mode 100644 index 0000000000..78d9a56230 --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/account_workspace_switching.rs @@ -0,0 +1,520 @@ +//! Exercises workspace switching through the app-server's public account API. + +use super::*; +use codex_app_server_protocol::AccountSessionsResponse; +use pretty_assertions::assert_eq; +use test_case::test_case; +use wiremock::matchers::body_json; + +#[derive(Clone, Copy, PartialEq)] +enum SwitchScenario { + Success, + MissingWorkspace, + DuplicateWorkspace, + InvalidRouting, + Unauthorized, + PendingLogin, + RequiresSignIn, + BrowserSignIn, + GovernmentWorkspace, +} + +#[test_case(SwitchScenario::Success; "success_and_refresh")] +#[test_case(SwitchScenario::MissingWorkspace; "missing_workspace")] +#[test_case(SwitchScenario::DuplicateWorkspace; "duplicate_workspace")] +#[test_case(SwitchScenario::InvalidRouting; "invalid_routing")] +#[test_case(SwitchScenario::Unauthorized; "unauthorized")] +#[test_case(SwitchScenario::PendingLogin; "pending_login")] +#[test_case(SwitchScenario::RequiresSignIn; "requires_sign_in")] +#[test_case(SwitchScenario::BrowserSignIn; "sso_sign_in_and_cancel")] +#[test_case(SwitchScenario::GovernmentWorkspace; "government_workspace")] +#[tokio::test] +async fn workspace_switch_validates_and_preserves_managed_login( + scenario: SwitchScenario, +) -> Result<()> { + let home = TempDir::new()?; + let backend = MockServer::start().await; + Mock::given(path("/backend-api/wham/settings/user")) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"commit_attribution_enabled": false})), + ) + .mount(&backend) + .await; + Mock::given(path("/backend-api/wham/config/bundle")) + .respond_with(ResponseTemplate::new(/*s*/ 200).set_body_json(json!({}))) + .mount(&backend) + .await; + create_config_toml( + home.path(), + CreateConfigTomlParams { + requires_openai_auth: Some(true), + base_url: Some(format!("{}/v1", backend.uri())), + chatgpt_base_url: Some(format!("{}/backend-api", backend.uri())), + ..Default::default() + }, + )?; + write_models_cache(home.path()).await?; + write_chatgpt_auth( + home.path(), + ChatGptAuthFixture::new("original-token") + .email("person@example.com") + .plan_type("plus") + .chatgpt_user_id("user-123") + .chatgpt_account_id(WORKSPACE_ID_INITIAL) + .account_id(WORKSPACE_ID_INITIAL), + AuthCredentialsStoreMode::File, + )?; + let original = load_file_auth(home.path())?.expect("managed auth exists"); + let inference = if scenario == SwitchScenario::Success { + let command = json!({ + "cmd": "echo workspace-switch", + "sandbox_permissions": "require_escalated", + "justification": "test workspace switching during approval", + }); + let pending = responses::sse(vec![ + responses::ev_function_call("pending-command", "exec_command", &command.to_string()), + responses::ev_completed("old-workspace-turn"), + ]); + Some( + responses::mount_response_sequence( + &backend, + vec![ + responses::sse_response(pending) + .insert_header("x-codex-turn-state", "old-workspace-state"), + responses::sse_response(responses::sse_completed("new-workspace-turn")), + ], + ) + .await, + ) + } else { + None + }; + let access_token = encode_id_token( + &ChatGptIdTokenClaims::new() + .chatgpt_user_id("user-123") + .chatgpt_account_id(WORKSPACE_ID_SECOND_ALLOWED) + .plan_type("business"), + )?; + let rejected_before_exchange = matches!( + scenario, + SwitchScenario::PendingLogin + | SwitchScenario::RequiresSignIn + | SwitchScenario::BrowserSignIn + | SwitchScenario::GovernmentWorkspace + ); + Mock::given(method("POST")) + .and(path("/backend-api/accounts/switch-workspace-token")) + .and(header("authorization", "Bearer original-token")) + .and(body_json( + json!({"workspace_id": WORKSPACE_ID_SECOND_ALLOWED}), + )) + .respond_with( + ResponseTemplate::new(/*s*/ 200).set_body_json(json!({"access_token": access_token})), + ) + .expect(u64::from(!rejected_before_exchange)) + .mount(&backend) + .await; + let refresh_url = format!("{}/oauth/token", backend.uri()); + let source = json!({ + "id": WORKSPACE_ID_INITIAL, + "workspace_backend_origin": "https://chatgpt.com", + "account_routing_override": "NO_CONSTRAINT", + }); + let mut target = json!({ + "id": WORKSPACE_ID_SECOND_ALLOWED, + "workspace_backend_origin": "https://chatgpt.com", + "account_routing_override": "NO_CONSTRAINT", + "can_access_with_session": !matches!(scenario, SwitchScenario::RequiresSignIn | SwitchScenario::BrowserSignIn), + "sso_connection_name": "test-enterprise", + "is_fedramp_compliant_workspace": scenario == SwitchScenario::GovernmentWorkspace, + }); + Mock::given(path("/backend-api/wham/accounts/check")) + .and(header("chatgpt-account-id", WORKSPACE_ID_INITIAL)) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"accounts": [source, target]})), + ) + .mount(&backend) + .await; + let target_memberships = match scenario { + SwitchScenario::MissingWorkspace => vec![], + SwitchScenario::DuplicateWorkspace => vec![target.clone(), target], + SwitchScenario::InvalidRouting => { + target["workspace_backend_origin"] = json!("invalid"); + vec![target] + } + _ => vec![target], + }; + let response = if scenario == SwitchScenario::Unauthorized { + ResponseTemplate::new(401) + } else { + ResponseTemplate::new(200).set_body_json(json!({"accounts": target_memberships})) + }; + Mock::given(path("/backend-api/wham/accounts/check")) + .and(header("chatgpt-account-id", WORKSPACE_ID_SECOND_ALLOWED)) + .respond_with(response) + .mount(&backend) + .await; + let original_token = encode_id_token( + &ChatGptIdTokenClaims::new() + .chatgpt_user_id("user-123") + .chatgpt_account_id(WORKSPACE_ID_INITIAL) + .plan_type("plus"), + )?; + Mock::given(path("/backend-api/accounts/switch-workspace-token")) + .and(body_json(json!({"workspace_id": WORKSPACE_ID_INITIAL}))) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"access_token":original_token})), + ) + .expect(u64::from( + !rejected_before_exchange && scenario != SwitchScenario::Success, + )) + .mount(&backend) + .await; + Mock::given(path("/oauth/token")) + .respond_with(if scenario == SwitchScenario::Unauthorized { + ResponseTemplate::new(401).set_body_json(json!({"error": "invalid_grant"})) + } else { + ResponseTemplate::new(200).set_body_json( + json!({"access_token": access_token, "refresh_token": "rotated-refresh-token"}), + ) + }) + .mount(&backend) + .await; + mock_device_code_usercode(&backend, /*interval_seconds*/ 1).await; + mock_device_code_token_failure(&backend, /*status*/ 404).await; + let issuer = backend.uri(); + let mut app = TestAppServer::builder() + .with_codex_home(home.path()) + .with_env_overrides(&[ + ("OPENAI_API_KEY", None), + ( + codex_login::REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR, + Some(refresh_url.as_str()), + ), + (LOGIN_ISSUER_ENV_VAR, Some(issuer.as_str())), + ]) + .build_initialized() + .await?; + let id = app + .send_request( + "accountSession/list", + Some(json!({ + "refreshWorkspaceMetadata": scenario == SwitchScenario::Success, + })), + ) + .await?; + let listed: AccountSessionsResponse = app.read_response(id).await?; + let session_id = listed.active_session_id.expect("active managed session"); + assert_eq!(session_id, "user-123"); + assert_eq!(listed.sessions.len(), 1); + let session = &listed.sessions[0]; + assert_eq!(session.session_id, session_id); + assert_eq!(session.user_id.as_deref(), Some("user-123")); + assert_eq!( + session.selected_workspace_account_id.as_deref(), + Some(WORKSPACE_ID_INITIAL) + ); + assert!( + session + .workspaces + .iter() + .any(|workspace| workspace.account_id == WORKSPACE_ID_INITIAL) + ); + let login_id = if scenario == SwitchScenario::PendingLogin { + let id = app.send_login_account_chatgpt_device_code_request().await?; + let LoginAccountResponse::ChatgptDeviceCode { login_id, .. } = + app.read_response(id).await? + else { + bail!("expected device login"); + }; + Some(login_id) + } else { + None + }; + let thread_id = if scenario == SwitchScenario::Success { + let id = app + .send_thread_start_request_with_auto_env(codex_app_server_protocol::ThreadStartParams { + model: Some("mock-model".into()), + approval_policy: Some(codex_app_server_protocol::AskForApproval::OnRequest), + approvals_reviewer: Some(codex_app_server_protocol::ApprovalsReviewer::User), + sandbox: Some(codex_app_server_protocol::SandboxMode::WorkspaceWrite), + ..Default::default() + }) + .await?; + let thread: codex_app_server_protocol::ThreadStartResponse = app.read_response(id).await?; + let id = app + .send_request( + "turn/start", + Some(json!({ + "threadId": thread.thread.id, + "input": [{"type": "text", "text": "Run the command"}], + })), + ) + .await?; + let _: codex_app_server_protocol::TurnStartResponse = app.read_response(id).await?; + let request = timeout( + DEFAULT_READ_TIMEOUT, + app.read_stream_until_request_message(), + ) + .await??; + assert!(matches!( + request, + ServerRequest::CommandExecutionRequestApproval { .. } + )); + let _: AccountUpdatedNotification = app.read_notification("account/updated").await?; + let id = app + .send_request( + "accountSession/switch", + Some(json!({ + "sessionId": "another-session", "accountId": WORKSPACE_ID_SECOND_ALLOWED, + })), + ) + .await?; + let error = app + .read_stream_until_error_message(RequestId::Integer(id)) + .await?; + assert!( + error + .error + .message + .contains("active managed account session") + ); + let id = app + .send_request("thread/loaded/list", Some(json!({}))) + .await?; + let loaded: codex_app_server_protocol::ThreadLoadedListResponse = + app.read_response(id).await?; + assert_eq!(loaded.data, vec![thread.thread.id.clone()]); + assert_eq!(load_file_auth(home.path())?, Some(original.clone())); + Some(thread.thread.id) + } else { + None + }; + let id = app + .send_request( + "accountSession/switch", + Some(json!({ + "sessionId": session_id, "accountId": WORKSPACE_ID_SECOND_ALLOWED, + "allowBrowserLogin": scenario == SwitchScenario::BrowserSignIn, + })), + ) + .await?; + if scenario == SwitchScenario::BrowserSignIn { + let response: codex_app_server_protocol::AccountSessionsSwitchResponse = + app.read_response(id).await?; + let codex_app_server_protocol::AccountSessionsSwitchResponse::BrowserLogin { + login_id, + auth_url, + } = response + else { + bail!("expected browser SSO"); + }; + let url = url::Url::parse(&auth_url)?; + let query: std::collections::HashMap<_, _> = url.query_pairs().collect(); + for (name, value) in [ + ("current_workspace_id", WORKSPACE_ID_SECOND_ALLOWED), + ("allowed_workspace_id", WORKSPACE_ID_SECOND_ALLOWED), + ("login_hint", "person@example.com"), + ("connection", "test-enterprise"), + ] { + assert_eq!(query.get(name).map(AsRef::as_ref), Some(value)); + } + assert!(!query.contains_key("prompt")); + assert!(!query.contains_key("reauth")); + let id = app + .send_cancel_login_account_request(CancelLoginAccountParams { login_id }) + .await?; + let _: CancelLoginAccountResponse = app.read_response(id).await?; + assert_eq!(load_file_auth(home.path())?, Some(original)); + return Ok(()); + } + if scenario != SwitchScenario::Success { + let error = app + .read_stream_until_error_message(RequestId::Integer(id)) + .await?; + let expected = match scenario { + SwitchScenario::PendingLogin => "current login", + SwitchScenario::RequiresSignIn => "requires a new sign-in", + SwitchScenario::GovernmentWorkspace => "government workspaces", + _ => "Workspace validation failed", + }; + assert!(error.error.message.contains(expected), "{error:?}"); + assert_eq!(load_file_auth(home.path())?, Some(original)); + if let Some(login_id) = login_id { + let id = app + .send_cancel_login_account_request(CancelLoginAccountParams { login_id }) + .await?; + let cancelled: CancelLoginAccountResponse = app.read_response(id).await?; + assert_eq!(cancelled.status, CancelLoginAccountStatus::Canceled); + } + return Ok(()); + } + let thread_id = thread_id.expect("success scenario starts a thread"); + let closed = timeout( + DEFAULT_READ_TIMEOUT, + app.read_stream_until_notification_message("thread/closed"), + ) + .await??; + assert_eq!( + closed.params.expect("thread/closed params")["threadId"], + thread_id.as_str() + ); + let interrupted: TurnCompletedNotification = app.read_notification("turn/completed").await?; + assert_eq!(interrupted.turn.status, TurnStatus::Interrupted); + let switched: AccountSessionsResponse = + timeout(DEFAULT_READ_TIMEOUT, app.read_response(id)).await??; + assert_eq!( + switched.active_session_id.as_deref(), + Some(session_id.as_str()) + ); + assert_eq!( + switched.sessions[0] + .selected_workspace_account_id + .as_deref(), + Some(WORKSPACE_ID_SECOND_ALLOWED) + ); + let notification = timeout( + DEFAULT_READ_TIMEOUT, + app.read_stream_until_notification_message("account/updated"), + ) + .await??; + assert_eq!( + serde_json::from_value::( + notification.params.expect("account/updated params") + )?, + AccountUpdatedNotification { + auth_mode: Some(AuthMode::Chatgpt), + plan_type: Some(AccountPlanType::Business), + } + ); + assert_eq!( + read_account(&mut app).await?.account, + Some(Account::Chatgpt { + email: Some("person@example.com".into()), + workspace_name: None, + plan_type: AccountPlanType::Business, + }) + ); + let mut expected = original; + let tokens = expected.tokens.as_mut().expect("managed tokens exist"); + tokens.account_id = Some(WORKSPACE_ID_SECOND_ALLOWED.into()); + tokens.access_token = access_token; + assert_eq!(load_file_auth(home.path())?, Some(expected.clone())); + let id = app + .send_request("account/read", Some(json!({"refreshToken": true}))) + .await?; + let account: GetAccountResponse = app.read_response(id).await?; + assert_eq!( + account.account, + Some(Account::Chatgpt { + email: Some("person@example.com".into()), + workspace_name: None, + plan_type: AccountPlanType::Business + }) + ); + let refreshed = load_file_auth(home.path())?.expect("managed auth exists"); + expected + .tokens + .as_mut() + .expect("managed tokens exist") + .refresh_token = "rotated-refresh-token".into(); + expected.last_refresh = refreshed.last_refresh; + assert_eq!(refreshed, expected); + let id = app + .send_thread_resume_request(codex_app_server_protocol::ThreadResumeParams { + thread_id: thread_id.clone(), + exclude_turns: true, + ..Default::default() + }) + .await?; + let _: codex_app_server_protocol::ThreadResumeResponse = app.read_response(id).await?; + let id = app + .send_request( + "accountSession/switch", + Some(json!({"sessionId": session_id, "accountId": WORKSPACE_ID_SECOND_ALLOWED})), + ) + .await?; + let unchanged: AccountSessionsResponse = app.read_response(id).await?; + assert_eq!(unchanged.active_session_id, Some(session_id)); + assert_eq!( + unchanged.sessions[0] + .selected_workspace_account_id + .as_deref(), + Some(WORKSPACE_ID_SECOND_ALLOWED) + ); + let id = app + .send_request("thread/loaded/list", Some(json!({}))) + .await?; + let loaded: codex_app_server_protocol::ThreadLoadedListResponse = app.read_response(id).await?; + assert_eq!(loaded.data, vec![thread_id]); + let id = app + .send_request( + "turn/start", + Some(json!({ + "threadId": loaded.data[0], + "input": [{"type": "text", "text": "Continue in the selected workspace"}], + })), + ) + .await?; + let _: codex_app_server_protocol::TurnStartResponse = app.read_response(id).await?; + let completed: TurnCompletedNotification = timeout( + DEFAULT_READ_TIMEOUT, + app.read_notification("turn/completed"), + ) + .await??; + assert_eq!(completed.turn.status, TurnStatus::Completed); + let requests = inference + .expect("success scenario mounts inference") + .requests(); + assert_eq!(requests.len(), 2); + assert_eq!( + requests[0].header("chatgpt-account-id"), + Some(WORKSPACE_ID_INITIAL.into()) + ); + assert_eq!( + requests[1].header("chatgpt-account-id"), + Some(WORKSPACE_ID_SECOND_ALLOWED.into()) + ); + assert_eq!(requests[1].header("x-codex-turn-state"), None); + Ok(()) +} + +#[tokio::test] +async fn account_session_list_is_empty_without_a_managed_chatgpt_login() -> Result<()> { + let home = TempDir::new()?; + let mut app = TestAppServer::builder() + .with_codex_home(home.path()) + .with_env_overrides(&[("OPENAI_API_KEY", None)]) + .build_initialized() + .await?; + for api_key in [None, Some("test-api-key")] { + if let Some(api_key) = api_key { + let id = app + .send_request( + "account/login/start", + Some(json!({ + "type": "apiKey", "apiKey": api_key, + })), + ) + .await?; + let _: LoginAccountResponse = app.read_response(id).await?; + } + let id = app + .send_request( + "accountSession/list", + Some(json!({ + "refreshWorkspaceMetadata": true, + })), + ) + .await?; + let response: AccountSessionsResponse = app.read_response(id).await?; + assert_eq!( + response, + AccountSessionsResponse { + active_session_id: None, + sessions: vec![] + } + ); + } + Ok(()) +} diff --git a/codex-rs/backend-client/src/types.rs b/codex-rs/backend-client/src/types.rs index 925716d67e..65a705c5bc 100644 --- a/codex-rs/backend-client/src/types.rs +++ b/codex-rs/backend-client/src/types.rs @@ -136,6 +136,12 @@ pub struct AccountsCheckResponse { #[derive(Clone, Debug, Deserialize)] pub struct AccountEntry { pub id: String, + pub can_access_with_session: Option, + pub sso_connection_name: Option, + #[serde(default)] + pub is_deactivated: bool, + #[serde(default)] + pub is_fedramp_compliant_workspace: bool, /// Current subscription reported by the accounts endpoint, independent of token claims. #[serde(default)] pub plan_type: Option, @@ -175,12 +181,18 @@ impl Default for RawAccounts { #[derive(Deserialize)] struct ChatGptAccountEntry { account: ChatGptAccountInfo, + can_access_with_session: Option, + sso_connection_name: Option, } #[derive(Deserialize)] struct ChatGptAccountInfo { account_id: Option, #[serde(default)] + is_deactivated: bool, + #[serde(default)] + is_fedramp_compliant_workspace: bool, + #[serde(default)] plan_type: Option, #[serde(default)] name: Option, @@ -202,9 +214,14 @@ impl<'de> Deserialize<'de> for AccountsCheckResponse { .account_ordering .iter() .filter_map(|account_id| { - let account = accounts.remove(account_id)?.account; + let entry = accounts.remove(account_id)?; + let account = entry.account; Some(AccountEntry { id: account.account_id?, + can_access_with_session: entry.can_access_with_session, + sso_connection_name: entry.sso_connection_name, + is_deactivated: account.is_deactivated, + is_fedramp_compliant_workspace: account.is_fedramp_compliant_workspace, plan_type: account.plan_type, workspace_backend_origin: None, account_routing_override: None, diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index c52d4f62fb..48cbeb984a 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -1187,6 +1187,19 @@ async fn loads_api_key_from_auth_json() { } } +#[test] +fn logout_does_not_create_a_missing_file_store() -> std::io::Result<()> { + let root = tempdir()?; + let missing_home = root.path().join("missing"); + assert!(!logout( + &missing_home, + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default() + )?); + assert!(!missing_home.exists()); + Ok(()) +} + #[test] fn logout_removes_auth_file() -> Result<(), std::io::Error> { let dir = tempdir()?; diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 9da4328ec7..15864c40a8 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -1,4 +1,7 @@ mod workspace_routing; +mod workspace_switching; +pub use workspace_switching::WorkspaceLogin; +pub use workspace_switching::WorkspaceSwitch; #[path = "account_grant.rs"] mod account_grant; @@ -604,7 +607,7 @@ impl CodexAuth { tokens: Some(tokens), last_refresh: Some(_), .. - }) => Ok(tokens), + }) => Ok(workspace_switching::selected_workspace_tokens(tokens)), _ => Err(std::io::Error::other("Token data is not available.")), } } @@ -738,7 +741,9 @@ impl CodexAuth { /// Returns `None` if token-backed ChatGPT auth is unavailable. fn get_current_token_data(&self) -> Option { - self.get_current_auth_json().and_then(|t| t.tokens) + self.get_current_auth_json() + .and_then(|t| t.tokens) + .map(workspace_switching::selected_workspace_tokens) } fn stored_managed_chatgpt_agent_identity_record( @@ -960,13 +965,12 @@ fn persist_agent_identity_record( storage: &Arc, record: AgentIdentityAuthRecord, ) -> std::io::Result<()> { - let _commit = auth_storage_commit_guard()?; + let _commit = storage.lock_for_write()?; let mut guard = auth_dot_json .lock() .map_err(|_| std::io::Error::other("failed to lock auth state"))?; let stored = storage.load()?; if let Some(previous) = guard.as_ref().and_then(|auth| auth.tokens.as_ref()) - && previous.login_owner_id.is_some() && !stored .as_ref() .and_then(|auth| auth.tokens.as_ref()) @@ -976,9 +980,7 @@ fn persist_agent_identity_record( "Account changed during agent identity registration.", )); } - let mut auth = stored - .or_else(|| guard.clone()) - .ok_or_else(|| std::io::Error::other("auth data is not available"))?; + let mut auth = stored.ok_or_else(|| std::io::Error::other("auth data is not available"))?; auth.agent_identity = Some(AgentIdentityStorage::Record(record)); storage.save(&auth)?; // Publishing credentials belongs to AuthManager. Persisting identity metadata must @@ -1022,12 +1024,15 @@ pub fn logout( auth_credentials_store_mode: AuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, ) -> std::io::Result { - let _commit = auth_storage_commit_guard()?; + if auth_credentials_store_mode == AuthCredentialsStoreMode::File && !codex_home.try_exists()? { + return Ok(false); + } let storage = create_auth_storage( codex_home.to_path_buf(), auth_credentials_store_mode, keyring_backend_kind, ); + let _commit = storage.lock_for_write()?; storage.delete() } @@ -1193,12 +1198,12 @@ pub fn save_auth( auth_credentials_store_mode: AuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, ) -> std::io::Result<()> { - let _commit = auth_storage_commit_guard()?; let storage = create_auth_storage( codex_home.to_path_buf(), auth_credentials_store_mode, keyring_backend_kind, ); + let _commit = storage.lock_for_write()?; storage.save(auth) } @@ -1658,7 +1663,7 @@ pub(super) fn persist_tokens( access_token: Option, refresh_token: Option, ) -> std::io::Result { - let _commit = auth_storage_commit_guard()?; + let _commit = storage.lock_for_write()?; let stored = storage.load()?; let mut auth_dot_json = stored.ok_or(std::io::Error::other("Token data is not available."))?; if auth_dot_json.resolved_mode() != AuthMode::Chatgpt { diff --git a/codex-rs/login/src/auth/manager/workspace_switching.rs b/codex-rs/login/src/auth/manager/workspace_switching.rs new file mode 100644 index 0000000000..bec42f3b12 --- /dev/null +++ b/codex-rs/login/src/auth/manager/workspace_switching.rs @@ -0,0 +1,370 @@ +//! Switches the workspace of managed ChatGPT auth without a new login. +//! Refresh credentials remain owned by the existing auth store. + +use super::*; +use codex_http_client::ClientRouteClass; +use codex_http_client::RouteAwareClientPool; +use serde_json::json; +use std::io; + +#[derive(Deserialize)] +struct WorkspaceTokenResponse { + access_token: String, +} + +/// Allows the caller to restore credentials if target workspace validation fails. +/// Restoration is conditional so it cannot overwrite a newer login or logout. +pub struct WorkspaceSwitch { + storage: Arc, + previous: AuthDotJson, + selected: AuthDotJson, +} + +impl std::fmt::Debug for WorkspaceSwitch { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("WorkspaceSwitch").finish_non_exhaustive() + } +} + +/// Binds a browser callback to the managed login that requested the workspace. +#[derive(Clone)] +pub struct WorkspaceLogin { + storage: Arc, + original: AuthDotJson, + pub(crate) workspace_id: String, + pub(crate) email: Option, + pub(crate) connection: Option, +} + +impl std::fmt::Debug for WorkspaceLogin { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("WorkspaceLogin").finish_non_exhaustive() + } +} + +impl WorkspaceLogin { + pub(crate) fn save(&self, mut updated: AuthDotJson) -> io::Result<()> { + let previous = self + .original + .tokens + .as_ref() + .ok_or_else(|| io::Error::other("Missing managed login"))?; + let tokens = updated + .tokens + .as_mut() + .ok_or_else(|| io::Error::other("Missing sign-in credentials"))?; + let access = parse_chatgpt_jwt_claims(&tokens.access_token).map_err(io::Error::other)?; + if previous.id_token.chatgpt_user_id.is_none() + || tokens.id_token.chatgpt_user_id != previous.id_token.chatgpt_user_id + || access.chatgpt_user_id != previous.id_token.chatgpt_user_id + || tokens.account_id.as_deref() != Some(&self.workspace_id) + || access.chatgpt_account_id != tokens.account_id + { + return Err(io::Error::new( + io::ErrorKind::PermissionDenied, + "Sign-in returned a different user or workspace", + )); + } + tokens.login_owner_id.clone_from(&previous.login_owner_id); + replace_if_current(self.storage.as_ref(), &self.original, &updated) + } +} + +impl CodexAuth { + pub fn workspace_login( + &self, + session_id: &str, + workspace_id: String, + connection: Option, + ) -> io::Result { + let Self::Chatgpt(auth) = self else { + return Err(io::Error::other( + "Workspace sign-in requires a managed login", + )); + }; + let original = auth + .current_auth_json() + .ok_or_else(|| io::Error::other("Missing managed login"))?; + if original.tokens.as_ref().and_then(managed_session_id) != Some(session_id) { + return Err(io::Error::other("The active login changed")); + } + Ok(WorkspaceLogin { + storage: auth.storage().clone(), + original, + workspace_id, + email: self.get_account_email(), + connection, + }) + } + + /// Stable across token rotation and workspace changes for the managed login. + pub fn managed_account_session_id(&self) -> Option { + if !matches!(self, Self::Chatgpt(_)) { + return None; + } + let tokens = self.get_current_token_data()?; + managed_session_id(&tokens).map(str::to_owned) + } +} + +impl AuthManager { + /// Checks local eligibility without changing credentials or an in-progress login. + pub fn validate_workspace_switch( + &self, + session_id: &str, + workspace_id: &str, + ) -> io::Result<()> { + if workspace_id.trim().is_empty() { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "accountId must not be empty", + )); + } + ensure_auth_workspace_allowed( + self.effective_chatgpt_workspaces().as_deref(), + workspace_id, + )?; + let auth = self.auth_cached(); + if !matches!(auth, Some(CodexAuth::Chatgpt(_))) { + return Err(io::Error::new( + io::ErrorKind::Unsupported, + "Workspace switching requires a managed ChatGPT login", + )); + } + if self.has_external_auth() { + return Err(io::Error::new( + io::ErrorKind::Unsupported, + "Workspace switching is owned by the external authentication provider", + )); + } + if auth + .and_then(|auth| auth.managed_account_session_id()) + .as_deref() + != Some(session_id) + { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "Only the active managed account session can be switched", + )); + } + Ok(()) + } + + /// Exchanges the current session into another workspace without exporting credentials. + pub async fn switch_workspace( + &self, + session_id: &str, + workspace_id: &str, + ) -> io::Result { + self.validate_workspace_switch(session_id, workspace_id)?; + self.auth().await; + let _guard = self + .refresh_lock + .acquire() + .await + .map_err(io::Error::other)?; + self.validate_workspace_switch(session_id, workspace_id)?; + let Some(CodexAuth::Chatgpt(auth)) = self.auth_cached() else { + return Err(io::Error::other("Authentication is unavailable")); + }; + let snapshot = auth + .current_auth_json() + .ok_or_else(|| io::Error::other("Authentication is unavailable"))?; + let tokens = snapshot + .tokens + .as_ref() + .ok_or_else(|| io::Error::other("Authentication is unavailable"))?; + let user_id = tokens + .id_token + .chatgpt_user_id + .as_deref() + .filter(|id| !id.is_empty()) + .ok_or_else(|| io::Error::other("Authentication is missing a user identity"))?; + // reload() may replace cached auth without the refresh lock. Bind the + // request to the exact snapshot used by both the exchange and conditional save. + if managed_session_id(tokens) != Some(session_id) { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "Only the active managed account session can be switched", + )); + } + if tokens.account_id.as_deref() == Some(workspace_id) { + return Ok(WorkspaceSwitch { + storage: auth.storage().clone(), + previous: snapshot.clone(), + selected: snapshot, + }); + } + let access_token = self + .exchange_workspace_token(&tokens.access_token, workspace_id, user_id) + .await?; + ensure_auth_workspace_allowed( + self.effective_chatgpt_workspaces().as_deref(), + workspace_id, + )?; + if self.has_external_auth() { + return Err(io::Error::new( + io::ErrorKind::Interrupted, + "Authentication changed during workspace switching", + )); + } + let mut updated = snapshot.clone(); + let tokens = updated + .tokens + .as_mut() + .ok_or_else(|| io::Error::other("Authentication is unavailable"))?; + tokens.access_token = access_token; + tokens.account_id = Some(workspace_id.to_owned()); + // The exchange does not renew the session or issue refresh/ID tokens. + updated.agent_identity = None; + replace_if_current(auth.storage().as_ref(), &snapshot, &updated)?; + self.reload().await; + Ok(WorkspaceSwitch { + storage: auth.storage().clone(), + previous: snapshot, + selected: updated, + }) + } + + /// Restores a rejected workspace change only while its credentials are still current. + pub async fn rollback_workspace_switch(&self, switched: WorkspaceSwitch) -> io::Result<()> { + let _guard = self + .refresh_lock + .acquire() + .await + .map_err(io::Error::other)?; + if self.has_external_auth() + || switched.storage.load_account_grant()?.as_ref() != Some(&switched.selected) + { + return Err(io::Error::new( + io::ErrorKind::Interrupted, + "Account credentials changed", + )); + } + let selected = switched + .selected + .tokens + .as_ref() + .ok_or_else(|| io::Error::other("Missing selected credentials"))?; + let previous_workspace = switched + .previous + .tokens + .as_ref() + .and_then(|tokens| tokens.account_id.as_deref()) + .ok_or_else(|| io::Error::other("Missing previous workspace"))?; + let user = selected + .id_token + .chatgpt_user_id + .as_deref() + .ok_or_else(|| io::Error::other("Missing selected user"))?; + // The exchange also selects the workspace used by the server on the next refresh. + self.exchange_workspace_token(&selected.access_token, previous_workspace, user) + .await?; + let result = replace_if_current( + switched.storage.as_ref(), + &switched.selected, + &switched.previous, + ); + self.reload().await; + result + } + + pub(super) async fn exchange_workspace_token( + &self, + access_token: &str, + workspace_id: &str, + user_id: &str, + ) -> io::Result { + let base_url = self + .chatgpt_base_url + .as_deref() + .unwrap_or(ChatGptEnvironment::default().chatgpt_base_url()) + .trim_end_matches('/'); + let endpoint = format!("{base_url}/accounts/switch-workspace-token"); + let client = RouteAwareClientPool::with_chatgpt_cloudflare_cookies_without_redirects_or_request_logging( + self.auth_route_config.http_client_factory().clone(), ClientRouteClass::Auth, + ); + let response = client + .post(&endpoint) + .headers(crate::default_client::default_headers()) + .header( + http::header::AUTHORIZATION, + format!("Bearer {access_token}"), + ) + .json(&json!({ "workspace_id": workspace_id })) + .timeout(Duration::from_secs(/*secs*/ 30)) + .send() + .await + .map_err(io::Error::other)?; + if !response.status().is_success() { + let kind = if response.status().is_client_error() { + io::ErrorKind::PermissionDenied + } else { + io::ErrorKind::Other + }; + return Err(io::Error::new( + kind, + format!("Workspace switch was rejected (HTTP {})", response.status()), + )); + } + let response: WorkspaceTokenResponse = response.json().await.map_err(io::Error::other)?; + let claims = parse_chatgpt_jwt_claims(&response.access_token).map_err(io::Error::other)?; + if claims.chatgpt_account_id.as_deref() != Some(workspace_id) + || claims.chatgpt_user_id.as_deref() != Some(user_id) + { + return Err(io::Error::other( + "Workspace exchange returned a different user or workspace", + )); + } + Ok(response.access_token) + } +} + +// Older credentials predate login_owner_id. Keep them usable without migrating the store. +fn managed_session_id(tokens: &TokenData) -> Option<&str> { + let user_id = tokens + .id_token + .chatgpt_user_id + .as_deref() + .filter(|id| !id.is_empty())?; + Some( + tokens + .login_owner_id + .as_deref() + .filter(|id| !id.is_empty()) + .unwrap_or(user_id), + ) +} + +fn replace_if_current( + storage: &dyn AuthStorageBackend, + expected: &AuthDotJson, + updated: &AuthDotJson, +) -> io::Result<()> { + let _commit = storage.lock_for_write()?; + if storage.load_account_grant()?.as_ref() != Some(expected) { + return Err(io::Error::new( + io::ErrorKind::Interrupted, + "Account credentials changed", + )); + } + storage.save(updated) +} + +/// The retained ID token can still describe the previous workspace until refresh. +pub(super) fn selected_workspace_tokens(mut tokens: TokenData) -> TokenData { + if tokens.account_id != tokens.id_token.chatgpt_account_id + && let Ok(claims) = parse_chatgpt_jwt_claims(&tokens.access_token) + && claims.chatgpt_account_id == tokens.account_id + && claims.chatgpt_user_id == tokens.id_token.chatgpt_user_id + { + tokens.id_token.chatgpt_account_id = claims.chatgpt_account_id; + tokens.id_token.chatgpt_plan_type = claims.chatgpt_plan_type; + tokens.id_token.chatgpt_account_is_fedramp = claims.chatgpt_account_is_fedramp; + } + tokens +} + +#[cfg(test)] +#[path = "workspace_switching_tests.rs"] +mod tests; diff --git a/codex-rs/login/src/auth/manager/workspace_switching_tests.rs b/codex-rs/login/src/auth/manager/workspace_switching_tests.rs new file mode 100644 index 0000000000..774dd2f66a --- /dev/null +++ b/codex-rs/login/src/auth/manager/workspace_switching_tests.rs @@ -0,0 +1,354 @@ +#![cfg(test)] + +//! Verifies workspace exchange persistence and failed-switch isolation. + +use super::*; +use base64::Engine; +use pretty_assertions::assert_eq; +use tempfile::TempDir; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::body_json; +use wiremock::matchers::header; +use wiremock::matchers::method; +use wiremock::matchers::path; + +fn token(workspace: &str, user: &str, plan: &str) -> String { + let claims = json!({ + "exp": (Utc::now() + chrono::Duration::hours(/*hours*/ 1)).timestamp(), + "email": "person@example.com", + "https://api.openai.com/auth": { + "chatgpt_account_id": workspace, + "chatgpt_user_id": user, + "chatgpt_plan_type": plan, + }, + }); + let encoded = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(claims.to_string()); + format!("e30.{encoded}.signature") +} + +async fn session(server: &MockServer) -> anyhow::Result<(TempDir, AuthManager, AuthDotJson)> { + let home = tempfile::tempdir()?; + let jwt = token("personal", "user-1", "plus"); + let snapshot = AuthDotJson { + auth_mode: Some(AuthMode::Chatgpt), + openai_api_key: None, + tokens: Some(TokenData { + id_token: parse_chatgpt_jwt_claims(&jwt)?, + access_token: jwt, + refresh_token: "retained-refresh-token".into(), + account_id: Some("personal".into()), + login_owner_id: Some("login-1".into()), + oauth_provenance: None, + }), + last_refresh: Some(Utc::now()), + agent_identity: None, + personal_access_token: None, + bedrock_api_key: None, + bedrock_access_keys: None, + }; + create_auth_storage( + home.path().to_path_buf(), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) + .save(&snapshot)?; + let manager = AuthManager::new( + home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, + Some(server.uri()), + AuthKeyringBackendKind::default(), + crate::test_support::transport_default_auth_route_config(), + ) + .await; + Ok((home, manager, snapshot)) +} + +#[tokio::test] +async fn workspace_switch_preserves_login_and_survives_reload() -> anyhow::Result<()> { + let server = MockServer::start().await; + let (home, manager, snapshot) = session(&server).await?; + let access_token = token("business", "user-1", "business"); + Mock::given(method("POST")) + .and(path("/accounts/switch-workspace-token")) + .and(header( + "authorization", + format!("Bearer {}", snapshot.tokens.as_ref().unwrap().access_token), + )) + .and(body_json(json!({"workspace_id": "business"}))) + .respond_with( + ResponseTemplate::new(/*s*/ 200).set_body_json(json!({"access_token": access_token})), + ) + .expect(/*r*/ 1) + .mount(&server) + .await; + let before = *manager.auth_change_state_receiver().borrow(); + manager.switch_workspace("login-1", "business").await?; + enforce_login_restrictions(&AuthConfig { + codex_home: home.path().to_path_buf(), + auth_credentials_store_mode: AuthCredentialsStoreMode::File, + keyring_backend_kind: AuthKeyringBackendKind::default(), + forced_login_method: None, + forced_chatgpt_workspace_id: Some(vec!["business".into()]), + managed_auth_policy: ManagedAuthPolicy::default(), + chatgpt_base_url: Some(server.uri()), + auth_route_config: crate::test_support::transport_default_auth_route_config(), + }) + .await?; + manager.reload().await; + let auth = manager.auth_cached().unwrap(); + let mut expected = snapshot.clone(); + let tokens = expected.tokens.as_mut().unwrap(); + tokens.access_token = access_token; + tokens.account_id = Some("business".into()); + assert_eq!(auth.get_current_auth_json(), Some(expected)); + assert_eq!( + ( + auth.get_account_id(), + auth.get_chatgpt_user_id(), + auth.account_plan_type() + ), + ( + Some("business".into()), + Some("user-1".into()), + Some(AccountPlanType::Business) + ), + ); + assert!( + manager + .auth_change_state_receiver() + .borrow() + .owner_generation + > before.owner_generation + ); + manager.switch_workspace("login-1", "business").await?; + Ok(()) +} + +#[tokio::test] +async fn workspace_switch_rejects_stale_agent_identity_and_rolls_back_conditionally() +-> anyhow::Result<()> { + let server = MockServer::start().await; + let (_home, manager, snapshot) = session(&server).await?; + let original_auth = manager.auth_cached().unwrap(); + Mock::given(method("POST")) + .respond_with( + ResponseTemplate::new(/*s*/ 200) + .set_body_json(json!({"access_token": token("business", "user-1", "business")})), + ) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(body_json(json!({"workspace_id": "personal"}))) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(json!({"access_token": token("personal", "user-1", "plus")})), + ) + .with_priority(1) + .expect(1) + .mount(&server) + .await; + let switched = manager.switch_workspace("login-1", "business").await?; + let selected = manager.auth_cached().unwrap().get_current_auth_json(); + let error = original_auth + .persist_managed_chatgpt_agent_identity_record(AgentIdentityAuthRecord { + agent_runtime_id: "runtime-personal".into(), + agent_private_key: "private-key-personal".into(), + account_id: "personal".into(), + chatgpt_user_id: "user-1".into(), + email: None, + plan_type: AccountPlanType::Plus, + chatgpt_account_is_fedramp: false, + task_id: None, + }) + .unwrap_err(); + assert!(error.to_string().contains("Account changed")); + manager.reload().await; + assert_eq!( + manager.auth_cached().unwrap().get_current_auth_json(), + selected + ); + manager.rollback_workspace_switch(switched).await?; + assert_eq!( + manager.auth_cached().unwrap().get_current_auth_json(), + Some(snapshot) + ); + let switched = manager.switch_workspace("login-1", "business").await?; + manager.logout().await?; + assert_eq!( + manager + .rollback_workspace_switch(switched) + .await + .unwrap_err() + .kind(), + io::ErrorKind::Interrupted + ); + assert!(manager.auth_cached().is_none()); + Ok(()) +} + +#[tokio::test] +async fn workspace_switch_rejection_preserves_current_login() -> anyhow::Result<()> { + for (status, response) in [ + ( + 403, + json!({"error": {"code": "invalid_workspace_selected"}}), + ), + ( + 200, + json!({"access_token": token("business", "another-user", "business")}), + ), + ( + 200, + json!({"access_token": token("wrong-workspace", "user-1", "business")}), + ), + ] { + let server = MockServer::start().await; + let (_home, manager, snapshot) = session(&server).await?; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(status).set_body_json(response)) + .expect(/*r*/ 1) + .mount(&server) + .await; + assert!( + manager + .switch_workspace("login-1", "business") + .await + .is_err() + ); + manager.reload().await; + assert_eq!( + manager.auth_cached().unwrap().get_current_auth_json(), + Some(snapshot) + ); + } + Ok(()) +} + +#[tokio::test] +async fn workspace_switch_obeys_workspace_policy() -> anyhow::Result<()> { + let server = MockServer::start().await; + let (_home, manager, snapshot) = session(&server).await?; + manager.set_forced_chatgpt_workspace_id(Some(vec!["personal".into()])); + assert_eq!( + manager + .switch_workspace("login-1", "business") + .await + .unwrap_err() + .kind(), + io::ErrorKind::PermissionDenied + ); + assert!(server.received_requests().await.unwrap().is_empty()); + assert_eq!( + manager.auth_cached().unwrap().get_current_auth_json(), + Some(snapshot) + ); + Ok(()) +} + +#[tokio::test] +async fn workspace_switch_does_not_restore_a_logged_out_session() -> anyhow::Result<()> { + let server = MockServer::start().await; + let (home, manager, _) = session(&server).await?; + let received = Arc::new(tokio::sync::Notify::new()); + let notify = received.clone(); + Mock::given(method("POST")) + .respond_with(move |_: &wiremock::Request| { + notify.notify_one(); + ResponseTemplate::new(/*s*/ 200) + .set_body_json(json!({"access_token": token("business", "user-1", "business")})) + .set_delay(Duration::from_millis(/*millis*/ 100)) + }) + .expect(/*r*/ 1) + .mount(&server) + .await; + let (switched, logged_out) = + tokio::join!(manager.switch_workspace("login-1", "business"), async { + received.notified().await; + manager.logout().await + },); + logged_out?; + assert_eq!(switched.unwrap_err().kind(), io::ErrorKind::Interrupted); + assert!(manager.auth_cached().is_none()); + assert!( + create_auth_storage( + home.path().to_path_buf(), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default() + ) + .load()? + .is_none() + ); + Ok(()) +} + +#[tokio::test] +async fn workspace_switch_rejects_a_session_from_a_previous_login() -> anyhow::Result<()> { + let server = MockServer::start().await; + let (home, manager, mut snapshot) = session(&server).await?; + let session_id = manager + .auth_cached() + .unwrap() + .managed_account_session_id() + .unwrap(); + let jwt = token("personal", "user-1", "plus"); + let tokens = snapshot.tokens.as_mut().unwrap(); + tokens.login_owner_id = Some("new-login".into()); + tokens.id_token = parse_chatgpt_jwt_claims(&jwt)?; + tokens.access_token = jwt; + save_auth( + home.path(), + &snapshot, + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + manager.reload().await; + assert_eq!( + manager + .switch_workspace(&session_id, "business") + .await + .unwrap_err() + .kind(), + io::ErrorKind::InvalidInput + ); + assert!(server.received_requests().await.unwrap().is_empty()); + assert_eq!( + manager.auth_cached().unwrap().get_current_auth_json(), + Some(snapshot) + ); + Ok(()) +} + +#[tokio::test] +async fn workspace_sso_only_commits_for_the_requesting_login_and_workspace() -> anyhow::Result<()> { + let backend = MockServer::start().await; + let (_home, manager, original) = session(&backend).await?; + let login = manager.auth_cached().unwrap().workspace_login( + "login-1", + "team".into(), + /*connection*/ None, + )?; + let result = |workspace: &str, user: &str| -> anyhow::Result { + let mut updated = original.clone(); + let tokens = updated.tokens.as_mut().unwrap(); + tokens.access_token = token(workspace, user, "business"); + tokens.id_token = parse_chatgpt_jwt_claims(&tokens.access_token)?; + tokens.account_id = Some(workspace.into()); + tokens.login_owner_id = Some("new-login".into()); + Ok(updated) + }; + assert!(login.save(result("team", "different-user")?).is_err()); + assert!(login.save(result("other-workspace", "user-1")?).is_err()); + assert_eq!(login.storage.load_account_grant()?, Some(original.clone())); + login.save(result("team", "user-1")?)?; + let selected = login.storage.load_account_grant()?.unwrap().tokens.unwrap(); + assert_eq!(selected.account_id.as_deref(), Some("team")); + assert_eq!(selected.login_owner_id.as_deref(), Some("login-1")); + manager.logout().await?; + assert!(login.save(result("team", "user-1")?).is_err()); + assert!(login.storage.load_account_grant()?.is_none()); + Ok(()) +} diff --git a/codex-rs/login/src/auth/storage.rs b/codex-rs/login/src/auth/storage.rs index d990d609b7..60adea09ca 100644 --- a/codex-rs/login/src/auth/storage.rs +++ b/codex-rs/login/src/auth/storage.rs @@ -182,6 +182,31 @@ pub(super) fn auth_storage_commit_guard() -> std::io::Result Option<&Path> { + None + } + + /// Serialize the whole read/modify/write across managers and processes. + fn lock_for_write( + &self, + ) -> std::io::Result<(std::sync::MutexGuard<'static, ()>, Option)> { + let guard = auth_storage_commit_guard()?; + let file = if let Some(home) = self.persistent_home() { + std::fs::create_dir_all(home)?; + let file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(home.join("auth.lock"))?; + file.lock()?; + Some(file) + } else { + None + }; + Ok((guard, file)) + } + fn load(&self) -> std::io::Result>; /// An unavailable keyring cannot authorize using an older fallback record. fn load_account_grant(&self) -> std::io::Result> { @@ -214,6 +239,10 @@ impl FileAuthStorage { } impl AuthStorageBackend for FileAuthStorage { + fn persistent_home(&self) -> Option<&Path> { + Some(&self.codex_home) + } + fn load(&self) -> std::io::Result> { let auth_file = get_auth_file(&self.codex_home); let auth_dot_json = match self.try_read_auth_json(&auth_file) { @@ -320,6 +349,10 @@ impl DirectKeyringAuthStorage { } impl AuthStorageBackend for DirectKeyringAuthStorage { + fn persistent_home(&self) -> Option<&Path> { + Some(&self.codex_home) + } + fn load(&self) -> std::io::Result> { let key = compute_store_key(&self.codex_home)?; self.load_from_keyring(&key) @@ -398,6 +431,10 @@ impl SecretsKeyringAuthStorage { } impl AuthStorageBackend for SecretsKeyringAuthStorage { + fn persistent_home(&self) -> Option<&Path> { + Some(&self.codex_home) + } + fn load(&self) -> std::io::Result> { match self .secrets_manager @@ -489,6 +526,10 @@ impl AutoAuthStorage { } impl AuthStorageBackend for AutoAuthStorage { + fn persistent_home(&self) -> Option<&Path> { + Some(&self.file_storage.codex_home) + } + fn load_account_grant(&self) -> std::io::Result> { match self.keyring_storage.load()? { Some(auth) => Ok(Some(auth)), diff --git a/codex-rs/login/src/auth/storage_telemetry.rs b/codex-rs/login/src/auth/storage_telemetry.rs index 9fb4639981..ac2aa19bf4 100644 --- a/codex-rs/login/src/auth/storage_telemetry.rs +++ b/codex-rs/login/src/auth/storage_telemetry.rs @@ -71,6 +71,10 @@ impl ObservedStorage { } impl AuthStorageBackend for ObservedStorage { + fn persistent_home(&self) -> Option<&std::path::Path> { + self.inner.persistent_home() + } + fn load(&self) -> std::io::Result> { let mut telemetry = telemetry(self.mode, self.kind, Operation::Load); let result = self.inner.load(); diff --git a/codex-rs/login/src/server.rs b/codex-rs/login/src/server.rs index d4f65c1fd7..98295d8cf8 100644 --- a/codex-rs/login/src/server.rs +++ b/codex-rs/login/src/server.rs @@ -27,6 +27,7 @@ use std::time::Duration; use crate::auth::AuthDotJson; use crate::auth::AuthKeyringBackendKind; +use crate::auth::WorkspaceLogin; use crate::auth::save_auth; use crate::callback_params::LIFE_SCIENCES_OAUTH_STATE_SUFFIX; use crate::callback_params::LoginCallbackResult; @@ -92,6 +93,7 @@ pub struct ServerOptions { pub port: u16, pub open_browser: bool, pub force_state: Option, + pub workspace_login: Option, pub forced_chatgpt_workspace_id: Option>, pub codex_streamlined_login: bool, pub login_success_page: LoginSuccessPage, @@ -117,6 +119,7 @@ impl ServerOptions { port: DEFAULT_PORT, open_browser: true, force_state: None, + workspace_login: None, forced_chatgpt_workspace_id, codex_streamlined_login: false, login_success_page: LoginSuccessPage::default(), @@ -199,6 +202,7 @@ pub fn run_login_server(opts: ServerOptions) -> io::Result { &pkce, &state, opts.forced_chatgpt_workspace_id.as_deref(), + opts.workspace_login.as_ref(), )?; if opts.open_browser { @@ -580,6 +584,7 @@ fn build_authorize_url( pkce: &PkceCodes, state: &str, forced_chatgpt_workspace_ids: Option<&[String]>, + workspace_login: Option<&WorkspaceLogin>, ) -> io::Result { let originator = originator().value; let workspace_ids = forced_chatgpt_workspace_ids.map(|ids| ids.join(",")); @@ -591,6 +596,15 @@ fn build_authorize_url( if let Some(workspace_ids) = workspace_ids.as_deref() { extra_parameters.push(("allowed_workspace_id", workspace_ids)); } + if let Some(login) = workspace_login { + extra_parameters.push(("current_workspace_id", &login.workspace_id)); + if let Some(email) = &login.email { + extra_parameters.push(("login_hint", email)); + } + if let Some(connection) = &login.connection { + extra_parameters.push(("connection", connection)); + } + } build_authorization_url(AuthorizationRequest { endpoint: &format!("{issuer}/oauth/authorize"), client_id, @@ -830,6 +844,7 @@ pub(crate) async fn persist_tokens_async( let id_token = exchanged.id_token.clone(); let access_token = exchanged.access_token.clone(); let refresh_token = exchanged.refresh_token.clone(); + let workspace_login = opts.workspace_login.clone(); let provenance = AccountOAuthProvenance { issuer: opts.issuer.clone(), client_id: opts.client_id.clone(), @@ -864,6 +879,9 @@ pub(crate) async fn persist_tokens_async( bedrock_api_key: None, bedrock_access_keys: None, }; + if let Some(login) = workspace_login { + return login.save(auth); + } save_auth( &codex_home, &auth, diff --git a/codex-rs/login/tests/suite/login_server_e2e.rs b/codex-rs/login/tests/suite/login_server_e2e.rs index 03d48769de..d53d066136 100644 --- a/codex-rs/login/tests/suite/login_server_e2e.rs +++ b/codex-rs/login/tests/suite/login_server_e2e.rs @@ -147,6 +147,7 @@ async fn end_to_end_login_flow_persists_auth_json() -> Result<()> { port: 0, open_browser: false, force_state: Some(state), + workspace_login: None, forced_chatgpt_workspace_id: Some(vec![chatgpt_account_id.to_string()]), codex_streamlined_login: false, auth_keyring_backend_kind: AuthKeyringBackendKind::Direct, @@ -258,6 +259,7 @@ async fn hosted_login_redirects_to_configured_open_app_url() -> Result<()> { port: 0, open_browser: false, force_state: Some("streamlined_state".to_string()), + workspace_login: None, forced_chatgpt_workspace_id: None, codex_streamlined_login: false, login_success_page: LoginSuccessPage::Hosted { @@ -311,6 +313,7 @@ async fn creates_missing_codex_home_dir() -> Result<()> { port: 0, open_browser: false, force_state: Some(state), + workspace_login: None, forced_chatgpt_workspace_id: None, codex_streamlined_login: false, auth_keyring_backend_kind: AuthKeyringBackendKind::Direct, @@ -354,6 +357,7 @@ async fn login_server_includes_forced_workspaces_as_one_query_param() -> Result< port: 0, open_browser: false, force_state: Some(state), + workspace_login: None, forced_chatgpt_workspace_id: Some(vec![ WORKSPACE_ID_ALLOWED.to_string(), WORKSPACE_ID_SECOND_ALLOWED.to_string(), @@ -398,6 +402,7 @@ async fn forced_chatgpt_workspace_id_mismatch_blocks_login() -> Result<()> { port: 0, open_browser: false, force_state: Some(state.clone()), + workspace_login: None, forced_chatgpt_workspace_id: Some(vec![WORKSPACE_ID_ALLOWED.to_string()]), codex_streamlined_login: false, auth_keyring_backend_kind: AuthKeyringBackendKind::Direct, @@ -461,6 +466,7 @@ async fn oauth_access_denied_missing_entitlement_blocks_login_with_clear_error() port: 0, open_browser: false, force_state: Some(state.clone()), + workspace_login: None, forced_chatgpt_workspace_id: None, codex_streamlined_login: false, auth_keyring_backend_kind: AuthKeyringBackendKind::Direct, @@ -532,6 +538,7 @@ async fn oauth_access_denied_unknown_reason_uses_generic_error_page() -> Result< port: 0, open_browser: false, force_state: Some(state.clone()), + workspace_login: None, forced_chatgpt_workspace_id: None, codex_streamlined_login: false, auth_keyring_backend_kind: AuthKeyringBackendKind::Direct, @@ -682,6 +689,7 @@ async fn cancels_previous_login_server_when_port_is_in_use() -> Result<()> { port: 0, open_browser: false, force_state: Some("cancel_state".to_string()), + workspace_login: None, forced_chatgpt_workspace_id: None, codex_streamlined_login: false, auth_keyring_backend_kind: AuthKeyringBackendKind::Direct, @@ -706,6 +714,7 @@ async fn cancels_previous_login_server_when_port_is_in_use() -> Result<()> { port: login_port, open_browser: false, force_state: Some("cancel_state_2".to_string()), + workspace_login: None, forced_chatgpt_workspace_id: None, codex_streamlined_login: false, auth_keyring_backend_kind: AuthKeyringBackendKind::Direct, diff --git a/sdk/python/src/openai_codex/generated/v2_all.py b/sdk/python/src/openai_codex/generated/v2_all.py index d58f29ae71..2afe8d7cc8 100644 --- a/sdk/python/src/openai_codex/generated/v2_all.py +++ b/sdk/python/src/openai_codex/generated/v2_all.py @@ -52,6 +52,11 @@ class AccountRoutingOverride(Enum): us_cr = "us_cr" +class AccountSessionWorkspaceKind(Enum): + personal = "personal" + workspace = "workspace" + + class AccountTokenUsageDailyBucket(BaseModel): model_config = ConfigDict( populate_by_name=True, @@ -6892,6 +6897,16 @@ class AccountLoginCompletedNotification(BaseModel): success: bool +class AccountSessionWorkspace(BaseModel): + model_config = ConfigDict( + populate_by_name=True, + ) + account_id: Annotated[str, Field(alias="accountId")] + image_url: Annotated[str | None, Field(alias="imageUrl")] = None + kind: AccountSessionWorkspaceKind | None = None + name: str | None = None + + class AccountUpdatedNotification(BaseModel): model_config = ConfigDict( populate_by_name=True, @@ -10893,6 +10908,23 @@ class AccountRateLimitsUpdatedNotification(BaseModel): rate_limits: Annotated[RateLimitSnapshot, Field(alias="rateLimits")] +class AccountSession(BaseModel): + model_config = ConfigDict( + populate_by_name=True, + ) + display_name: Annotated[str | None, Field(alias="displayName")] = None + email: str | None = None + image_url: Annotated[str | None, Field(alias="imageUrl")] = None + is_active: Annotated[bool, Field(alias="isActive")] + last_used_at: Annotated[int, Field(alias="lastUsedAt")] + selected_workspace_account_id: Annotated[ + str | None, Field(alias="selectedWorkspaceAccountId") + ] = None + session_id: Annotated[str, Field(alias="sessionId")] + user_id: Annotated[str | None, Field(alias="userId")] = None + workspaces: list[AccountSessionWorkspace] + + class AppInfo(BaseModel): model_config = ConfigDict( populate_by_name=True,