Add experimental ChatGPT workspace switching to app-server (#57514)

## What changed

- Add `accountSession/list` and `accountSession/switch` for the active managed ChatGPT login, with optional workspace membership refresh.
- Validate membership and workspace policy, pause new turns, and unload running threads before switching credentials. Notify clients with `thread/closed` and `account/updated`; selecting the current workspace leaves threads running.
- Support browser sign-in for SSO through `allowBrowserLogin`, binding the callback to the requesting login and workspace. Require a separate sign-in for government workspaces.
- Preserve the selected workspace across token refresh and conditionally restore the previous workspace if routing validation fails. Guard credential writes against concurrent login or logout changes.

## Testing

Add app-server and auth tests covering workspace switching, credential preservation, policy and identity rejection, conditional rollback, stale sessions, and SSO callback validation.

GitOrigin-RevId: 1234abc33b71e69b66318da3e201fccc08867906

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())

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<crate::turn_admission::TurnAdmissionPause>,
+    ) -> 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) {

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<AccountSessionsResponse, JSONRPCErrorError> {
+        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::<Vec<_>>();
+            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<codex_backend_client::AccountsCheckResponse, JSONRPCErrorError> {
+        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<AccountSessionsSwitchResponse, JSONRPCErrorError> {
+        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(&params.session_id, &params.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(&params.account_id)
+        {
+            return self.account_session_switch_response(&params).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(&params.session_id, &params.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(
+                    &params.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(&params.session_id, &params.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(&params).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<AccountSessionsSwitchResponse, JSONRPCErrorError> {
+        let response = self
+            .list_account_sessions(AccountSessionsListParams {
+                refresh_workspace_metadata: false,
+            })
+            .await?;
+        if response.active_session_id.as_deref() != Some(&params.session_id)
+            || !response.sessions.iter().any(|session| {
+                session.session_id == params.session_id
+                    && session.selected_workspace_account_id.as_deref() == Some(&params.account_id)
+            })
+        {
+            return Err(invalid_request(
+                "Authentication changed during workspace switching",
+            ));
+        }
+        Ok(AccountSessionsSwitchResponse::Switched { sessions: response })
+    }
+}

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;

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

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<TurnAdmissionPause, JSONRPCErrorError> {
+        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<TurnPermit, JSONRPCErrorError> {
-        self.try_admit().ok_or_else(server_draining_error)
-    }
-
-    fn try_admit(&self) -> Option<TurnPermit> {
         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<Box<dyn Send>> {
-        self.try_admit()
+        self.admit()
+            .ok()
             .map(|permit| Box::new(permit) as Box<dyn Send>)
     }
 }

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