diff --git a/codex-rs/app-server-protocol/schema/json/ChatgptAuthTokensRefreshParams.json b/codex-rs/app-server-protocol/schema/json/ChatgptAuthTokensRefreshParams.json index b3e828e80..8b320fd67 100644 --- a/codex-rs/app-server-protocol/schema/json/ChatgptAuthTokensRefreshParams.json +++ b/codex-rs/app-server-protocol/schema/json/ChatgptAuthTokensRefreshParams.json @@ -15,7 +15,7 @@ }, "properties": { "previousAccountId": { - "description": "Workspace/account identifier that Codex was previously using.\n\nClients that manage multiple accounts/workspaces can use this as a hint to refresh the token for the correct workspace.\n\nThis may be `null` when the prior ID token did not include a workspace identifier (`chatgpt_account_id`) or when the token could not be parsed.", + "description": "Workspace/account identifier that Codex was previously using.\n\nClients that manage multiple accounts/workspaces can use this as a hint to refresh the token for the correct workspace.\n\nThis may be `null` when the prior auth state did not include a workspace identifier (`chatgpt_account_id`).", "type": [ "string", "null" diff --git a/codex-rs/app-server-protocol/schema/json/ChatgptAuthTokensRefreshResponse.json b/codex-rs/app-server-protocol/schema/json/ChatgptAuthTokensRefreshResponse.json index 48d149492..6d88e784c 100644 --- a/codex-rs/app-server-protocol/schema/json/ChatgptAuthTokensRefreshResponse.json +++ b/codex-rs/app-server-protocol/schema/json/ChatgptAuthTokensRefreshResponse.json @@ -4,13 +4,19 @@ "accessToken": { "type": "string" }, - "idToken": { + "chatgptAccountId": { "type": "string" + }, + "chatgptPlanType": { + "type": [ + "string", + "null" + ] } }, "required": [ "accessToken", - "idToken" + "chatgptAccountId" ], "title": "ChatgptAuthTokensRefreshResponse", "type": "object" diff --git a/codex-rs/app-server-protocol/schema/json/ClientRequest.json b/codex-rs/app-server-protocol/schema/json/ClientRequest.json index c981dfda6..2797dc366 100644 --- a/codex-rs/app-server-protocol/schema/json/ClientRequest.json +++ b/codex-rs/app-server-protocol/schema/json/ClientRequest.json @@ -998,13 +998,20 @@ "description": "[UNSTABLE] FOR OPENAI INTERNAL USE ONLY - DO NOT USE. The access token must contain the same scopes that Codex-managed ChatGPT auth tokens have.", "properties": { "accessToken": { - "description": "Access token (JWT) supplied by the client. This token is used for backend API requests.", + "description": "Access token (JWT) supplied by the client. This token is used for backend API requests and email extraction.", "type": "string" }, - "idToken": { - "description": "ID token (JWT) supplied by the client.\n\nThis token is used for identity and account metadata (email, plan type, workspace id).", + "chatgptAccountId": { + "description": "Workspace/account identifier supplied by the client.", "type": "string" }, + "chatgptPlanType": { + "description": "Optional plan type supplied by the client.\n\nWhen `null`, Codex attempts to derive the plan type from access-token claims. If unavailable, the plan defaults to `unknown`.", + "type": [ + "string", + "null" + ] + }, "type": { "enum": [ "chatgptAuthTokens" @@ -1015,7 +1022,7 @@ }, "required": [ "accessToken", - "idToken", + "chatgptAccountId", "type" ], "title": "ChatgptAuthTokensLoginAccountParams", diff --git a/codex-rs/app-server-protocol/schema/json/ServerRequest.json b/codex-rs/app-server-protocol/schema/json/ServerRequest.json index ad0c2e354..3fa487385 100644 --- a/codex-rs/app-server-protocol/schema/json/ServerRequest.json +++ b/codex-rs/app-server-protocol/schema/json/ServerRequest.json @@ -41,7 +41,7 @@ "ChatgptAuthTokensRefreshParams": { "properties": { "previousAccountId": { - "description": "Workspace/account identifier that Codex was previously using.\n\nClients that manage multiple accounts/workspaces can use this as a hint to refresh the token for the correct workspace.\n\nThis may be `null` when the prior ID token did not include a workspace identifier (`chatgpt_account_id`) or when the token could not be parsed.", + "description": "Workspace/account identifier that Codex was previously using.\n\nClients that manage multiple accounts/workspaces can use this as a hint to refresh the token for the correct workspace.\n\nThis may be `null` when the prior auth state did not include a workspace identifier (`chatgpt_account_id`).", "type": [ "string", "null" 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 11e8fe9e4..cdb6407b5 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 @@ -338,7 +338,7 @@ "$schema": "http://json-schema.org/draft-07/schema#", "properties": { "previousAccountId": { - "description": "Workspace/account identifier that Codex was previously using.\n\nClients that manage multiple accounts/workspaces can use this as a hint to refresh the token for the correct workspace.\n\nThis may be `null` when the prior ID token did not include a workspace identifier (`chatgpt_account_id`) or when the token could not be parsed.", + "description": "Workspace/account identifier that Codex was previously using.\n\nClients that manage multiple accounts/workspaces can use this as a hint to refresh the token for the correct workspace.\n\nThis may be `null` when the prior auth state did not include a workspace identifier (`chatgpt_account_id`).", "type": [ "string", "null" @@ -371,13 +371,19 @@ "accessToken": { "type": "string" }, - "idToken": { + "chatgptAccountId": { "type": "string" + }, + "chatgptPlanType": { + "type": [ + "string", + "null" + ] } }, "required": [ "accessToken", - "idToken" + "chatgptAccountId" ], "title": "ChatgptAuthTokensRefreshResponse", "type": "object" @@ -12071,13 +12077,20 @@ "description": "[UNSTABLE] FOR OPENAI INTERNAL USE ONLY - DO NOT USE. The access token must contain the same scopes that Codex-managed ChatGPT auth tokens have.", "properties": { "accessToken": { - "description": "Access token (JWT) supplied by the client. This token is used for backend API requests.", + "description": "Access token (JWT) supplied by the client. This token is used for backend API requests and email extraction.", "type": "string" }, - "idToken": { - "description": "ID token (JWT) supplied by the client.\n\nThis token is used for identity and account metadata (email, plan type, workspace id).", + "chatgptAccountId": { + "description": "Workspace/account identifier supplied by the client.", "type": "string" }, + "chatgptPlanType": { + "description": "Optional plan type supplied by the client.\n\nWhen `null`, Codex attempts to derive the plan type from access-token claims. If unavailable, the plan defaults to `unknown`.", + "type": [ + "string", + "null" + ] + }, "type": { "enum": [ "chatgptAuthTokens" @@ -12088,7 +12101,7 @@ }, "required": [ "accessToken", - "idToken", + "chatgptAccountId", "type" ], "title": "ChatgptAuthTokensv2::LoginAccountParams", diff --git a/codex-rs/app-server-protocol/schema/json/v2/LoginAccountParams.json b/codex-rs/app-server-protocol/schema/json/v2/LoginAccountParams.json index 66df09435..ce6bdd4a3 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/LoginAccountParams.json +++ b/codex-rs/app-server-protocol/schema/json/v2/LoginAccountParams.json @@ -41,13 +41,20 @@ "description": "[UNSTABLE] FOR OPENAI INTERNAL USE ONLY - DO NOT USE. The access token must contain the same scopes that Codex-managed ChatGPT auth tokens have.", "properties": { "accessToken": { - "description": "Access token (JWT) supplied by the client. This token is used for backend API requests.", + "description": "Access token (JWT) supplied by the client. This token is used for backend API requests and email extraction.", "type": "string" }, - "idToken": { - "description": "ID token (JWT) supplied by the client.\n\nThis token is used for identity and account metadata (email, plan type, workspace id).", + "chatgptAccountId": { + "description": "Workspace/account identifier supplied by the client.", "type": "string" }, + "chatgptPlanType": { + "description": "Optional plan type supplied by the client.\n\nWhen `null`, Codex attempts to derive the plan type from access-token claims. If unavailable, the plan defaults to `unknown`.", + "type": [ + "string", + "null" + ] + }, "type": { "enum": [ "chatgptAuthTokens" @@ -58,7 +65,7 @@ }, "required": [ "accessToken", - "idToken", + "chatgptAccountId", "type" ], "title": "ChatgptAuthTokensv2::LoginAccountParams", diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/ChatgptAuthTokensRefreshParams.ts b/codex-rs/app-server-protocol/schema/typescript/v2/ChatgptAuthTokensRefreshParams.ts index 4393c7f7a..7358d8034 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/ChatgptAuthTokensRefreshParams.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/ChatgptAuthTokensRefreshParams.ts @@ -10,7 +10,7 @@ export type ChatgptAuthTokensRefreshParams = { reason: ChatgptAuthTokensRefreshR * Clients that manage multiple accounts/workspaces can use this as a hint * to refresh the token for the correct workspace. * - * This may be `null` when the prior ID token did not include a workspace - * identifier (`chatgpt_account_id`) or when the token could not be parsed. + * This may be `null` when the prior auth state did not include a workspace + * identifier (`chatgpt_account_id`). */ previousAccountId?: string | null, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/ChatgptAuthTokensRefreshResponse.ts b/codex-rs/app-server-protocol/schema/typescript/v2/ChatgptAuthTokensRefreshResponse.ts index f7f7ecba8..30bf03e83 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/ChatgptAuthTokensRefreshResponse.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/ChatgptAuthTokensRefreshResponse.ts @@ -2,4 +2,4 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. -export type ChatgptAuthTokensRefreshResponse = { idToken: string, accessToken: string, }; +export type ChatgptAuthTokensRefreshResponse = { accessToken: string, chatgptAccountId: string, chatgptPlanType: string | null, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/LoginAccountParams.ts b/codex-rs/app-server-protocol/schema/typescript/v2/LoginAccountParams.ts index 5c1f4c02a..ef668f9c1 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/LoginAccountParams.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/LoginAccountParams.ts @@ -3,15 +3,19 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. export type LoginAccountParams = { "type": "apiKey", apiKey: string, } | { "type": "chatgpt" } | { "type": "chatgptAuthTokens", -/** - * ID token (JWT) supplied by the client. - * - * This token is used for identity and account metadata (email, plan type, - * workspace id). - */ -idToken: string, /** * Access token (JWT) supplied by the client. - * This token is used for backend API requests. + * This token is used for backend API requests and email extraction. */ -accessToken: string, }; +accessToken: string, +/** + * Workspace/account identifier supplied by the client. + */ +chatgptAccountId: string, +/** + * Optional plan type supplied by the client. + * + * When `null`, Codex attempts to derive the plan type from access-token + * claims. If unavailable, the plan defaults to `unknown`. + */ +chatgptPlanType?: string | null, }; diff --git a/codex-rs/app-server-protocol/src/protocol/common.rs b/codex-rs/app-server-protocol/src/protocol/common.rs index 7b4db0d6d..bc7738924 100644 --- a/codex-rs/app-server-protocol/src/protocol/common.rs +++ b/codex-rs/app-server-protocol/src/protocol/common.rs @@ -1003,7 +1003,8 @@ mod tests { request_id: RequestId::Integer(5), params: v2::LoginAccountParams::ChatgptAuthTokens { access_token: "access-token".to_string(), - id_token: "id-token".to_string(), + chatgpt_account_id: "org-123".to_string(), + chatgpt_plan_type: Some("business".to_string()), }, }; assert_eq!( @@ -1013,7 +1014,8 @@ mod tests { "params": { "type": "chatgptAuthTokens", "accessToken": "access-token", - "idToken": "id-token" + "chatgptAccountId": "org-123", + "chatgptPlanType": "business" } }), serde_json::to_value(&request)?, diff --git a/codex-rs/app-server-protocol/src/protocol/v2.rs b/codex-rs/app-server-protocol/src/protocol/v2.rs index a75bdaab3..c3cfd6335 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2.rs @@ -905,21 +905,20 @@ pub enum LoginAccountParams { /// [UNSTABLE] FOR OPENAI INTERNAL USE ONLY - DO NOT USE. /// The access token must contain the same scopes that Codex-managed ChatGPT auth tokens have. #[experimental("account/login/start.chatgptAuthTokens")] - #[serde(rename = "chatgptAuthTokens")] - #[ts(rename = "chatgptAuthTokens")] + #[serde(rename = "chatgptAuthTokens", rename_all = "camelCase")] + #[ts(rename = "chatgptAuthTokens", rename_all = "camelCase")] ChatgptAuthTokens { - /// ID token (JWT) supplied by the client. - /// - /// This token is used for identity and account metadata (email, plan type, - /// workspace id). - #[serde(rename = "idToken")] - #[ts(rename = "idToken")] - id_token: String, /// Access token (JWT) supplied by the client. - /// This token is used for backend API requests. - #[serde(rename = "accessToken")] - #[ts(rename = "accessToken")] + /// This token is used for backend API requests and email extraction. access_token: String, + /// Workspace/account identifier supplied by the client. + chatgpt_account_id: String, + /// Optional plan type supplied by the client. + /// + /// When `null`, Codex attempts to derive the plan type from access-token + /// claims. If unavailable, the plan defaults to `unknown`. + #[ts(optional = nullable)] + chatgpt_plan_type: Option, }, } @@ -991,8 +990,8 @@ pub struct ChatgptAuthTokensRefreshParams { /// Clients that manage multiple accounts/workspaces can use this as a hint /// to refresh the token for the correct workspace. /// - /// This may be `null` when the prior ID token did not include a workspace - /// identifier (`chatgpt_account_id`) or when the token could not be parsed. + /// This may be `null` when the prior auth state did not include a workspace + /// identifier (`chatgpt_account_id`). #[ts(optional = nullable)] pub previous_account_id: Option, } @@ -1001,8 +1000,9 @@ pub struct ChatgptAuthTokensRefreshParams { #[serde(rename_all = "camelCase")] #[ts(export_to = "v2/")] pub struct ChatgptAuthTokensRefreshResponse { - pub id_token: String, pub access_token: String, + pub chatgpt_account_id: String, + pub chatgpt_plan_type: Option, } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 0c24ed30c..3dfabc3c0 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -203,7 +203,6 @@ use codex_core::skills::remote::download_remote_skill; use codex_core::skills::remote::list_remote_skills; use codex_core::state_db::StateDbHandle; use codex_core::state_db::open_if_present; -use codex_core::token_data::parse_id_token; use codex_core::windows_sandbox::WindowsSandboxLevelExt; use codex_feedback::CodexFeedback; use codex_login::ServerOptions as LoginServerOptions; @@ -789,11 +788,17 @@ impl CodexMessageProcessor { self.login_chatgpt_v2(request_id).await; } LoginAccountParams::ChatgptAuthTokens { - id_token, access_token, + chatgpt_account_id, + chatgpt_plan_type, } => { - self.login_chatgpt_auth_tokens(request_id, id_token, access_token) - .await; + self.login_chatgpt_auth_tokens( + request_id, + access_token, + chatgpt_account_id, + chatgpt_plan_type, + ) + .await; } } } @@ -1224,8 +1229,9 @@ impl CodexMessageProcessor { async fn login_chatgpt_auth_tokens( &mut self, request_id: ConnectionRequestId, - id_token: String, access_token: String, + chatgpt_account_id: String, + chatgpt_plan_type: Option, ) { if matches!( self.config.forced_login_method, @@ -1249,27 +1255,13 @@ impl CodexMessageProcessor { } } - let id_token_info = match parse_id_token(&id_token) { - Ok(info) => info, - Err(err) => { - let error = JSONRPCErrorError { - code: INVALID_REQUEST_ERROR_CODE, - message: format!("invalid id token: {err}"), - data: None, - }; - self.outgoing.send_error(request_id, error).await; - return; - } - }; - if let Some(expected_workspace) = self.config.forced_chatgpt_workspace_id.as_deref() - && id_token_info.chatgpt_account_id.as_deref() != Some(expected_workspace) + && chatgpt_account_id != expected_workspace { - let account_id = id_token_info.chatgpt_account_id; let error = JSONRPCErrorError { code: INVALID_REQUEST_ERROR_CODE, message: format!( - "External auth must use workspace {expected_workspace}, but received {account_id:?}." + "External auth must use workspace {expected_workspace}, but received {chatgpt_account_id:?}." ), data: None, }; @@ -1277,9 +1269,12 @@ impl CodexMessageProcessor { return; } - if let Err(err) = - login_with_chatgpt_auth_tokens(&self.config.codex_home, &id_token, &access_token) - { + if let Err(err) = login_with_chatgpt_auth_tokens( + &self.config.codex_home, + &access_token, + &chatgpt_account_id, + chatgpt_plan_type.as_deref(), + ) { let error = JSONRPCErrorError { code: INTERNAL_ERROR_CODE, message: format!("failed to set external auth: {err}"), diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 2646e5f0a..26da44df3 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -100,7 +100,8 @@ impl ExternalAuthRefresher for ExternalAuthRefreshBridge { Ok(ExternalAuthTokens { access_token: response.access_token, - id_token: response.id_token, + chatgpt_account_id: response.chatgpt_account_id, + chatgpt_plan_type: response.chatgpt_plan_type, }) } } diff --git a/codex-rs/app-server/tests/common/auth_fixtures.rs b/codex-rs/app-server/tests/common/auth_fixtures.rs index e689e4183..c52109ac5 100644 --- a/codex-rs/app-server/tests/common/auth_fixtures.rs +++ b/codex-rs/app-server/tests/common/auth_fixtures.rs @@ -11,7 +11,7 @@ use codex_core::auth::AuthCredentialsStoreMode; use codex_core::auth::AuthDotJson; use codex_core::auth::save_auth; use codex_core::token_data::TokenData; -use codex_core::token_data::parse_id_token; +use codex_core::token_data::parse_chatgpt_jwt_claims; use serde_json::json; /// Builder for writing a fake ChatGPT auth.json in tests. @@ -148,7 +148,7 @@ pub fn write_chatgpt_auth( cli_auth_credentials_store_mode: AuthCredentialsStoreMode, ) -> Result<()> { let id_token_raw = encode_id_token(&fixture.claims)?; - let id_token = parse_id_token(&id_token_raw).context("parse id token")?; + let id_token = parse_chatgpt_jwt_claims(&id_token_raw).context("parse id token")?; let tokens = TokenData { id_token, access_token: fixture.access_token, diff --git a/codex-rs/app-server/tests/common/mcp_process.rs b/codex-rs/app-server/tests/common/mcp_process.rs index d4541e6f4..c1d004d40 100644 --- a/codex-rs/app-server/tests/common/mcp_process.rs +++ b/codex-rs/app-server/tests/common/mcp_process.rs @@ -334,12 +334,14 @@ impl McpProcess { /// Send an `account/login/start` JSON-RPC request with ChatGPT auth tokens. pub async fn send_chatgpt_auth_tokens_login_request( &mut self, - id_token: String, access_token: String, + chatgpt_account_id: String, + chatgpt_plan_type: Option, ) -> anyhow::Result { let params = LoginAccountParams::ChatgptAuthTokens { - id_token, access_token, + chatgpt_account_id, + chatgpt_plan_type, }; let params = Some(serde_json::to_value(params)?); self.send_request("account/login/start", params).await diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index d3145345e..0c9680095 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -166,19 +166,22 @@ async fn set_auth_token_updates_account_and_notifies() -> Result<()> { )?; write_models_cache(codex_home.path())?; - let id_token = encode_id_token( + let access_token = encode_id_token( &ChatGptIdTokenClaims::new() .email("embedded@example.com") .plan_type("pro") .chatgpt_account_id("org-embedded"), )?; - let access_token = "access-embedded".to_string(); let mut mcp = McpProcess::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let set_id = mcp - .send_chatgpt_auth_tokens_login_request(id_token.clone(), access_token) + .send_chatgpt_auth_tokens_login_request( + access_token, + "org-embedded".to_string(), + Some("pro".to_string()), + ) .await?; let set_resp: JSONRPCResponse = timeout( DEFAULT_READ_TIMEOUT, @@ -236,7 +239,7 @@ async fn account_read_refresh_token_is_noop_in_external_mode() -> Result<()> { )?; write_models_cache(codex_home.path())?; - let id_token = encode_id_token( + let access_token = encode_id_token( &ChatGptIdTokenClaims::new() .email("embedded@example.com") .plan_type("pro") @@ -247,7 +250,11 @@ async fn account_read_refresh_token_is_noop_in_external_mode() -> Result<()> { timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let set_id = mcp - .send_chatgpt_auth_tokens_login_request(id_token, "access-embedded".to_string()) + .send_chatgpt_auth_tokens_login_request( + access_token, + "org-embedded".to_string(), + Some("pro".to_string()), + ) .await?; let set_resp: JSONRPCResponse = timeout( DEFAULT_READ_TIMEOUT, @@ -300,7 +307,8 @@ async fn account_read_refresh_token_is_noop_in_external_mode() -> Result<()> { async fn respond_to_refresh_request( mcp: &mut McpProcess, access_token: &str, - id_token: &str, + chatgpt_account_id: &str, + chatgpt_plan_type: Option<&str>, ) -> Result<()> { let refresh_req: ServerRequest = timeout( DEFAULT_READ_TIMEOUT, @@ -313,7 +321,8 @@ async fn respond_to_refresh_request( assert_eq!(params.reason, ChatgptAuthTokensRefreshReason::Unauthorized); let response = ChatgptAuthTokensRefreshResponse { access_token: access_token.to_string(), - id_token: id_token.to_string(), + chatgpt_account_id: chatgpt_account_id.to_string(), + chatgpt_plan_type: chatgpt_plan_type.map(str::to_string), }; mcp.send_response(request_id, serde_json::to_value(response)?) .await?; @@ -349,28 +358,27 @@ async fn external_auth_refreshes_on_unauthorized() -> Result<()> { ) .await; - let initial_id_token = encode_id_token( + let initial_access_token = encode_id_token( &ChatGptIdTokenClaims::new() .email("initial@example.com") .plan_type("pro") .chatgpt_account_id("org-initial"), )?; - let refreshed_id_token = encode_id_token( + let refreshed_access_token = encode_id_token( &ChatGptIdTokenClaims::new() .email("refreshed@example.com") .plan_type("pro") .chatgpt_account_id("org-refreshed"), )?; - let initial_access_token = "access-initial".to_string(); - let refreshed_access_token = "access-refreshed".to_string(); let mut mcp = McpProcess::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let set_id = mcp .send_chatgpt_auth_tokens_login_request( - initial_id_token.clone(), initial_access_token.clone(), + "org-initial".to_string(), + Some("pro".to_string()), ) .await?; let set_resp: JSONRPCResponse = timeout( @@ -409,7 +417,13 @@ async fn external_auth_refreshes_on_unauthorized() -> Result<()> { ..Default::default() }) .await?; - respond_to_refresh_request(&mut mcp, &refreshed_access_token, &refreshed_id_token).await?; + respond_to_refresh_request( + &mut mcp, + &refreshed_access_token, + "org-refreshed", + Some("pro"), + ) + .await?; let _turn_resp: JSONRPCResponse = timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), @@ -456,7 +470,7 @@ async fn external_auth_refresh_error_fails_turn() -> Result<()> { let _responses_mock = responses::mount_response_sequence(&mock_server, vec![unauthorized]).await; - let initial_id_token = encode_id_token( + let initial_access_token = encode_id_token( &ChatGptIdTokenClaims::new() .email("initial@example.com") .plan_type("pro") @@ -467,7 +481,11 @@ async fn external_auth_refresh_error_fails_turn() -> Result<()> { timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let set_id = mcp - .send_chatgpt_auth_tokens_login_request(initial_id_token, "access-initial".to_string()) + .send_chatgpt_auth_tokens_login_request( + initial_access_token, + "org-initial".to_string(), + Some("pro".to_string()), + ) .await?; let set_resp: JSONRPCResponse = timeout( DEFAULT_READ_TIMEOUT, @@ -568,13 +586,13 @@ async fn external_auth_refresh_mismatched_workspace_fails_turn() -> Result<()> { let _responses_mock = responses::mount_response_sequence(&mock_server, vec![unauthorized]).await; - let initial_id_token = encode_id_token( + let initial_access_token = encode_id_token( &ChatGptIdTokenClaims::new() .email("initial@example.com") .plan_type("pro") .chatgpt_account_id("org-expected"), )?; - let refreshed_id_token = encode_id_token( + let refreshed_access_token = encode_id_token( &ChatGptIdTokenClaims::new() .email("refreshed@example.com") .plan_type("pro") @@ -585,7 +603,11 @@ async fn external_auth_refresh_mismatched_workspace_fails_turn() -> Result<()> { timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let set_id = mcp - .send_chatgpt_auth_tokens_login_request(initial_id_token, "access-initial".to_string()) + .send_chatgpt_auth_tokens_login_request( + initial_access_token, + "org-expected".to_string(), + Some("pro".to_string()), + ) .await?; let set_resp: JSONRPCResponse = timeout( DEFAULT_READ_TIMEOUT, @@ -636,8 +658,9 @@ async fn external_auth_refresh_mismatched_workspace_fails_turn() -> Result<()> { mcp.send_response( request_id, serde_json::to_value(ChatgptAuthTokensRefreshResponse { - access_token: "access-refreshed".to_string(), - id_token: refreshed_id_token, + access_token: refreshed_access_token, + chatgpt_account_id: "org-other".to_string(), + chatgpt_plan_type: Some("pro".to_string()), })?, ) .await?; @@ -664,8 +687,8 @@ async fn external_auth_refresh_mismatched_workspace_fails_turn() -> Result<()> { } #[tokio::test] -// Refresh returns a malformed id_token; turn fails. -async fn external_auth_refresh_invalid_id_token_fails_turn() -> Result<()> { +// Refresh returns a malformed access token; turn fails. +async fn external_auth_refresh_invalid_access_token_fails_turn() -> Result<()> { let codex_home = TempDir::new()?; let mock_server = MockServer::start().await; create_config_toml( @@ -684,7 +707,7 @@ async fn external_auth_refresh_invalid_id_token_fails_turn() -> Result<()> { let _responses_mock = responses::mount_response_sequence(&mock_server, vec![unauthorized]).await; - let initial_id_token = encode_id_token( + let initial_access_token = encode_id_token( &ChatGptIdTokenClaims::new() .email("initial@example.com") .plan_type("pro") @@ -695,7 +718,11 @@ async fn external_auth_refresh_invalid_id_token_fails_turn() -> Result<()> { timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let set_id = mcp - .send_chatgpt_auth_tokens_login_request(initial_id_token, "access-initial".to_string()) + .send_chatgpt_auth_tokens_login_request( + initial_access_token, + "org-initial".to_string(), + Some("pro".to_string()), + ) .await?; let set_resp: JSONRPCResponse = timeout( DEFAULT_READ_TIMEOUT, @@ -746,8 +773,9 @@ async fn external_auth_refresh_invalid_id_token_fails_turn() -> Result<()> { mcp.send_response( request_id, serde_json::to_value(ChatgptAuthTokensRefreshResponse { - access_token: "access-refreshed".to_string(), - id_token: "not-a-jwt".to_string(), + access_token: "not-a-jwt".to_string(), + chatgpt_account_id: "org-initial".to_string(), + chatgpt_plan_type: Some("pro".to_string()), })?, ) .await?; @@ -967,7 +995,7 @@ async fn set_auth_token_cancels_active_chatgpt_login() -> Result<()> { bail!("unexpected login response: {login:?}"); }; - let id_token = encode_id_token( + let access_token = encode_id_token( &ChatGptIdTokenClaims::new() .email("embedded@example.com") .plan_type("pro") @@ -976,7 +1004,11 @@ async fn set_auth_token_cancels_active_chatgpt_login() -> Result<()> { // Set an external auth token instead of completing the ChatGPT login flow. // This should cancel the active login attempt. let set_id = mcp - .send_chatgpt_auth_tokens_login_request(id_token, "access-embedded".to_string()) + .send_chatgpt_auth_tokens_login_request( + access_token, + "org-embedded".to_string(), + Some("pro".to_string()), + ) .await?; let set_resp: JSONRPCResponse = timeout( DEFAULT_READ_TIMEOUT, diff --git a/codex-rs/core/src/auth.rs b/codex-rs/core/src/auth.rs index 7e58ff125..a6282305e 100644 --- a/codex-rs/core/src/auth.rs +++ b/codex-rs/core/src/auth.rs @@ -26,11 +26,10 @@ use crate::auth::storage::create_auth_storage; use crate::config::Config; use crate::error::RefreshTokenFailedError; use crate::error::RefreshTokenFailedReason; -use crate::token_data::IdTokenInfo; use crate::token_data::KnownPlan as InternalKnownPlan; use crate::token_data::PlanType as InternalPlanType; use crate::token_data::TokenData; -use crate::token_data::parse_id_token; +use crate::token_data::parse_chatgpt_jwt_claims; use crate::util::try_parse_error_message; use codex_client::CodexHttpClient; use codex_protocol::account::PlanType as AccountPlanType; @@ -115,7 +114,8 @@ pub enum RefreshTokenError { #[derive(Clone, Debug, PartialEq, Eq)] pub struct ExternalAuthTokens { pub access_token: String, - pub id_token: String, + pub chatgpt_account_id: String, + pub chatgpt_plan_type: Option, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -399,10 +399,15 @@ pub fn login_with_api_key( /// Writes an in-memory auth payload for externally managed ChatGPT tokens. pub fn login_with_chatgpt_auth_tokens( codex_home: &Path, - id_token: &str, access_token: &str, + chatgpt_account_id: &str, + chatgpt_plan_type: Option<&str>, ) -> std::io::Result<()> { - let auth_dot_json = AuthDotJson::from_external_token_strings(id_token, access_token)?; + let auth_dot_json = AuthDotJson::from_external_access_token( + access_token, + chatgpt_account_id, + chatgpt_plan_type, + )?; save_auth( codex_home, &auth_dot_json, @@ -591,7 +596,7 @@ fn update_tokens( let tokens = auth_dot_json.tokens.get_or_insert_with(TokenData::default); if let Some(id_token) = id_token { - tokens.id_token = parse_id_token(&id_token).map_err(std::io::Error::other)?; + tokens.id_token = parse_chatgpt_jwt_claims(&id_token).map_err(std::io::Error::other)?; } if let Some(access_token) = access_token { tokens.access_token = access_token; @@ -727,30 +732,42 @@ fn refresh_token_endpoint() -> String { } impl AuthDotJson { - fn from_external_tokens(external: &ExternalAuthTokens, id_token: IdTokenInfo) -> Self { - let account_id = id_token.chatgpt_account_id.clone(); + fn from_external_tokens(external: &ExternalAuthTokens) -> std::io::Result { + let mut token_info = + parse_chatgpt_jwt_claims(&external.access_token).map_err(std::io::Error::other)?; + token_info.chatgpt_account_id = Some(external.chatgpt_account_id.clone()); + token_info.chatgpt_plan_type = external + .chatgpt_plan_type + .as_deref() + .map(InternalPlanType::from_raw_value) + .or(token_info.chatgpt_plan_type) + .or(Some(InternalPlanType::Unknown("unknown".to_string()))); let tokens = TokenData { - id_token, + id_token: token_info, access_token: external.access_token.clone(), refresh_token: String::new(), - account_id, + account_id: Some(external.chatgpt_account_id.clone()), }; - Self { + Ok(Self { auth_mode: Some(ApiAuthMode::ChatgptAuthTokens), openai_api_key: None, tokens: Some(tokens), last_refresh: Some(Utc::now()), - } + }) } - fn from_external_token_strings(id_token: &str, access_token: &str) -> std::io::Result { - let id_token_info = parse_id_token(id_token).map_err(std::io::Error::other)?; + fn from_external_access_token( + access_token: &str, + chatgpt_account_id: &str, + chatgpt_plan_type: Option<&str>, + ) -> std::io::Result { let external = ExternalAuthTokens { access_token: access_token.to_string(), - id_token: id_token.to_string(), + chatgpt_account_id: chatgpt_account_id.to_string(), + chatgpt_plan_type: chatgpt_plan_type.map(str::to_string), }; - Ok(Self::from_external_tokens(&external, id_token_info)) + Self::from_external_tokens(&external) } fn resolved_mode(&self) -> ApiAuthMode { @@ -1233,19 +1250,18 @@ impl AuthManager { }; let refreshed = refresher.refresh(context).await?; - let id_token = parse_id_token(&refreshed.id_token) - .map_err(|err| RefreshTokenError::Transient(std::io::Error::other(err)))?; - if let Some(expected_workspace_id) = forced_chatgpt_workspace_id.as_deref() { - let actual_workspace_id = id_token.chatgpt_account_id.as_deref(); - if actual_workspace_id != Some(expected_workspace_id) { - return Err(RefreshTokenError::Transient(std::io::Error::other( - format!( - "external auth refresh returned workspace {actual_workspace_id:?}, expected {expected_workspace_id:?}", - ), - ))); - } + if let Some(expected_workspace_id) = forced_chatgpt_workspace_id.as_deref() + && refreshed.chatgpt_account_id != expected_workspace_id + { + return Err(RefreshTokenError::Transient(std::io::Error::other( + format!( + "external auth refresh returned workspace {:?}, expected {expected_workspace_id:?}", + refreshed.chatgpt_account_id, + ), + ))); } - let auth_dot_json = AuthDotJson::from_external_tokens(&refreshed, id_token); + let auth_dot_json = + AuthDotJson::from_external_tokens(&refreshed).map_err(RefreshTokenError::Transient)?; save_auth( &self.codex_home, &auth_dot_json, diff --git a/codex-rs/core/src/auth/storage.rs b/codex-rs/core/src/auth/storage.rs index 1ac1b2ee1..81d17e4e1 100644 --- a/codex-rs/core/src/auth/storage.rs +++ b/codex-rs/core/src/auth/storage.rs @@ -502,7 +502,7 @@ mod tests { let signature_b64 = encode(b"sig"); let fake_jwt = format!("{header_b64}.{payload_b64}.{signature_b64}"); - crate::token_data::parse_id_token(&fake_jwt).expect("fake JWT should parse") + crate::token_data::parse_chatgpt_jwt_claims(&fake_jwt).expect("fake JWT should parse") } fn auth_with_prefix(prefix: &str) -> AuthDotJson { diff --git a/codex-rs/core/src/token_data.rs b/codex-rs/core/src/token_data.rs index e38660ff5..85babd851 100644 --- a/codex-rs/core/src/token_data.rs +++ b/codex-rs/core/src/token_data.rs @@ -60,6 +60,22 @@ pub(crate) enum PlanType { Unknown(String), } +impl PlanType { + pub(crate) fn from_raw_value(raw: &str) -> Self { + match raw.to_ascii_lowercase().as_str() { + "free" => Self::Known(KnownPlan::Free), + "go" => Self::Known(KnownPlan::Go), + "plus" => Self::Known(KnownPlan::Plus), + "pro" => Self::Known(KnownPlan::Pro), + "team" => Self::Known(KnownPlan::Team), + "business" => Self::Known(KnownPlan::Business), + "enterprise" => Self::Known(KnownPlan::Enterprise), + "education" | "edu" => Self::Known(KnownPlan::Edu), + _ => Self::Unknown(raw.to_string()), + } + } +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] pub(crate) enum KnownPlan { @@ -111,9 +127,9 @@ pub enum IdTokenInfoError { Json(#[from] serde_json::Error), } -pub fn parse_id_token(id_token: &str) -> Result { +pub fn parse_chatgpt_jwt_claims(jwt: &str) -> Result { // JWT format: header.payload.signature - let mut parts = id_token.split('.'); + let mut parts = jwt.split('.'); let (_header_b64, payload_b64, _sig_b64) = match (parts.next(), parts.next(), parts.next()) { (Some(h), Some(p), Some(s)) if !h.is_empty() && !p.is_empty() && !s.is_empty() => (h, p, s), _ => return Err(IdTokenInfoError::InvalidFormat), @@ -128,14 +144,14 @@ pub fn parse_id_token(id_token: &str) -> Result { match claims.auth { Some(auth) => Ok(IdTokenInfo { email, - raw_jwt: id_token.to_string(), + raw_jwt: jwt.to_string(), chatgpt_plan_type: auth.chatgpt_plan_type, chatgpt_user_id: auth.chatgpt_user_id.or(auth.user_id), chatgpt_account_id: auth.chatgpt_account_id, }), None => Ok(IdTokenInfo { email, - raw_jwt: id_token.to_string(), + raw_jwt: jwt.to_string(), chatgpt_plan_type: None, chatgpt_user_id: None, chatgpt_account_id: None, @@ -148,7 +164,7 @@ where D: serde::Deserializer<'de>, { let s = String::deserialize(deserializer)?; - parse_id_token(&s).map_err(serde::de::Error::custom) + parse_chatgpt_jwt_claims(&s).map_err(serde::de::Error::custom) } fn serialize_id_token(id_token: &IdTokenInfo, serializer: S) -> Result @@ -191,7 +207,7 @@ mod tests { let signature_b64 = b64url_no_pad(b"sig"); let fake_jwt = format!("{header_b64}.{payload_b64}.{signature_b64}"); - let info = parse_id_token(&fake_jwt).expect("should parse"); + let info = parse_chatgpt_jwt_claims(&fake_jwt).expect("should parse"); assert_eq!(info.email.as_deref(), Some("user@example.com")); assert_eq!(info.get_chatgpt_plan_type().as_deref(), Some("Pro")); } @@ -223,7 +239,7 @@ mod tests { let signature_b64 = b64url_no_pad(b"sig"); let fake_jwt = format!("{header_b64}.{payload_b64}.{signature_b64}"); - let info = parse_id_token(&fake_jwt).expect("should parse"); + let info = parse_chatgpt_jwt_claims(&fake_jwt).expect("should parse"); assert_eq!(info.email.as_deref(), Some("user@example.com")); assert_eq!(info.get_chatgpt_plan_type().as_deref(), Some("Go")); } @@ -250,7 +266,7 @@ mod tests { let signature_b64 = b64url_no_pad(b"sig"); let fake_jwt = format!("{header_b64}.{payload_b64}.{signature_b64}"); - let info = parse_id_token(&fake_jwt).expect("should parse"); + let info = parse_chatgpt_jwt_claims(&fake_jwt).expect("should parse"); assert!(info.email.is_none()); assert!(info.get_chatgpt_plan_type().is_none()); } diff --git a/codex-rs/login/src/server.rs b/codex-rs/login/src/server.rs index e571d9b82..d97de008e 100644 --- a/codex-rs/login/src/server.rs +++ b/codex-rs/login/src/server.rs @@ -20,7 +20,7 @@ use codex_core::auth::AuthDotJson; use codex_core::auth::save_auth; use codex_core::default_client::originator; use codex_core::token_data::TokenData; -use codex_core::token_data::parse_id_token; +use codex_core::token_data::parse_chatgpt_jwt_claims; use rand::RngCore; use serde_json::Value as JsonValue; use tiny_http::Header; @@ -548,7 +548,7 @@ pub(crate) async fn persist_tokens_async( let codex_home = codex_home.to_path_buf(); tokio::task::spawn_blocking(move || { let mut tokens = TokenData { - id_token: parse_id_token(&id_token).map_err(io::Error::other)?, + id_token: parse_chatgpt_jwt_claims(&id_token).map_err(io::Error::other)?, access_token, refresh_token, account_id: None,