From 4a64cc688ed03d7c62ffd2bf48341f9197abf30b Mon Sep 17 00:00:00 2001 From: Dhruv Date: Tue, 2 Jun 2026 04:00:43 +0530 Subject: [PATCH] Exchange workspace token when switching account sessions --- codex-rs/app-server/README.md | 2 +- codex-rs/app-server/src/account_sessions.rs | 37 ++++++++++++++++--- .../request_processors/account_processor.rs | 1 + codex-rs/backend-client/src/client.rs | 21 +++++++++++ codex-rs/backend-client/src/lib.rs | 1 + codex-rs/backend-client/src/types.rs | 7 ++++ 6 files changed, 63 insertions(+), 6 deletions(-) diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 5778a56261..e72db9aa98 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -1866,7 +1866,7 @@ Field notes: ### 6) Save and switch ChatGPT account sessions -Desktop clients can save the current managed ChatGPT auth payload after login, list saved sessions, and switch the active workspace. These methods return the full saved-session snapshot because the number of interactive account sessions is intentionally small. +Desktop clients can save the current managed ChatGPT auth payload after login, list saved sessions, and switch the active workspace. Switching a workspace exchanges the saved access token for a newly signed workspace-scoped token without opening a browser. These methods return the full saved-session snapshot because the number of interactive account sessions is intentionally small. ```json { "method": "accountSession/add", "id": 6, "params": { "switchToAddedAccount": true } } diff --git a/codex-rs/app-server/src/account_sessions.rs b/codex-rs/app-server/src/account_sessions.rs index 73485394c1..a541e8e6bb 100644 --- a/codex-rs/app-server/src/account_sessions.rs +++ b/codex-rs/app-server/src/account_sessions.rs @@ -17,6 +17,7 @@ use codex_login::logout; use codex_login::revoke_account_session_auth; use codex_login::save_account_session_auth; use codex_login::save_auth; +use codex_login::token_data::parse_chatgpt_jwt_claims; use serde::Deserialize; use serde::Serialize; use std::fs::OpenOptions; @@ -202,26 +203,52 @@ impl<'a> AccountSessionsStore<'a> { Ok(Self::response(stored)) } - pub(crate) fn switch( + pub(crate) async fn switch( &self, session_id: &str, account_id: &str, ) -> std::io::Result { self.sync_active_auth()?; let mut stored = self.load()?; - let session = stored + let index = stored .sessions - .iter_mut() - .find(|session| session.session_id == session_id) + .iter() + .position(|session| session.session_id == session_id) .ok_or_else(|| std::io::Error::other("Saved ChatGPT account session not found"))?; let mut auth_json = self - .load_session_auth(&session.session_id)? + .load_session_auth(&stored.sessions[index].session_id)? .ok_or_else(|| std::io::Error::other("Saved ChatGPT account session has no tokens"))?; + let auth = CodexAuth::from_account_session_auth_dot_json( + self.codex_home, + &stored.sessions[index].session_id, + auth_json.clone(), + self.auth_credentials_store_mode, + Some(self.chatgpt_base_url), + ) + .await?; + let client = BackendClient::from_auth(self.chatgpt_base_url, &auth) + .map_err(std::io::Error::other)?; + let replacement = client + .switch_workspace_token(account_id) + .await + .map_err(std::io::Error::other)?; + let plan = parse_chatgpt_jwt_claims(&replacement.access_token) + .ok() + .and_then(|claims| claims.get_chatgpt_plan_type_raw()); let tokens = auth_json .tokens .as_mut() .ok_or_else(|| std::io::Error::other("Saved ChatGPT account session has no tokens"))?; + tokens.access_token = replacement.access_token; + if let Some(refresh_token) = replacement.refresh_token { + tokens.refresh_token = refresh_token; + } tokens.account_id = Some(account_id.to_string()); + auth_json.last_refresh = Some(Utc::now()); + let session = &mut stored.sessions[index]; + if let Some(plan) = plan { + session.plan = Some(plan); + } session.selected_workspace_account_id = Some(account_id.to_string()); session.last_used_at = Utc::now().timestamp(); stored.active_session_id = Some(session_id.to_string()); 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 f1631ac949..61f1580a10 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -149,6 +149,7 @@ impl AccountRequestProcessor { let response = self .account_sessions_store() .switch(¶ms.session_id, ¶ms.account_id) + .await .map_err(|err| internal_error(format!("failed to switch account session: {err}")))?; self.sync_auth_after_account_session_change().await; Ok(Some(response.into())) diff --git a/codex-rs/backend-client/src/client.rs b/codex-rs/backend-client/src/client.rs index 24d2776590..545c29eeac 100644 --- a/codex-rs/backend-client/src/client.rs +++ b/codex-rs/backend-client/src/client.rs @@ -4,6 +4,7 @@ use crate::types::ConfigBundleResponse; use crate::types::PaginatedListTaskListItem; use crate::types::RateLimitReachedKind as BackendRateLimitReachedKind; use crate::types::RateLimitStatusPayload; +use crate::types::SwitchWorkspaceTokenResponse; use crate::types::TurnAttemptsSiblingTurnsResponse; use anyhow::Result; use codex_api::SharedAuthProvider; @@ -97,6 +98,11 @@ struct SendAddCreditsNudgeEmailRequest { credit_type: AddCreditsNudgeCreditType, } +#[derive(Serialize)] +struct SwitchWorkspaceTokenRequest<'a> { + workspace_id: &'a str, +} + #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum PathStyle { /// /api/codex/… @@ -313,6 +319,21 @@ impl Client { self.decode_json(&url, &ct, &body) } + pub async fn switch_workspace_token( + &self, + workspace_id: &str, + ) -> Result { + let url = format!("{}/accounts/switch-workspace-token", self.base_url); + let req = self + .http + .post(&url) + .headers(self.headers()) + .header(CONTENT_TYPE, HeaderValue::from_static("application/json")) + .json(&SwitchWorkspaceTokenRequest { workspace_id }); + let (body, ct) = self.exec_request(req, "POST", &url).await?; + self.decode_json(&url, &ct, &body) + } + pub async fn send_add_credits_nudge_email( &self, credit_type: AddCreditsNudgeCreditType, diff --git a/codex-rs/backend-client/src/lib.rs b/codex-rs/backend-client/src/lib.rs index dfd7e816f6..8078727ae2 100644 --- a/codex-rs/backend-client/src/lib.rs +++ b/codex-rs/backend-client/src/lib.rs @@ -13,5 +13,6 @@ pub use types::DeliveredConfigToml; pub use types::DeliveredRequirementsToml; pub use types::DeliveredTomlFragment; pub use types::PaginatedListTaskListItem; +pub use types::SwitchWorkspaceTokenResponse; pub use types::TaskListItem; pub use types::TurnAttemptsSiblingTurnsResponse; diff --git a/codex-rs/backend-client/src/types.rs b/codex-rs/backend-client/src/types.rs index 06989241fa..6a9350b759 100644 --- a/codex-rs/backend-client/src/types.rs +++ b/codex-rs/backend-client/src/types.rs @@ -35,6 +35,13 @@ pub struct AccountEntry { pub structure: String, } +#[derive(Clone, Debug, Deserialize)] +pub struct SwitchWorkspaceTokenResponse { + pub access_token: String, + #[serde(default)] + pub refresh_token: Option, +} + #[derive(Deserialize)] struct RawAccountsCheckResponse { #[serde(default)]