feat(auth): add OAuth authorization code + PKCE for Codex provider (#2790)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Amit Singh <amitksingh1490@gmail.com>
Co-authored-by: ForgeCode <noreply@forgecode.dev>

Author: 4 people

Cargo.lock

@@ -89,6 +89,7 @@ syn = { version = "2.0.117", features = ["derive", "parsing"] }
 sysinfo = "0.38.3"
 tempfile = "3.27.0"
 termimad = "0.34.1"
+tiny_http = "0.12.0"
 syntect = { version = "5", default-features = false, features = ["default-syntaxes", "default-themes", "regex-onig"] }
 thiserror = "2.0.18"
 toml_edit = { version = "0.22", features = ["serde"] }

Cargo.toml

@@ -26,6 +26,10 @@ pub struct OAuthTokenResponse {
     /// OAuth scopes granted
     #[serde(skip_serializing_if = "Option::is_none")]
     pub scope: Option<String>,
+
+    /// ID token containing user identity claims (OpenID Connect)
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub id_token: Option<String>,
 }
 
 fn default_token_type() -> String {

crates/forge_domain/src/auth/auth_token_response.rs

@@ -57,6 +57,68 @@ impl AuthStrategy for ApiKeyStrategy {
     }
 }
 
+/// Extract the ChatGPT account ID from a JWT token's claims.
+///
+/// Checks `chatgpt_account_id`, `https://api.openai.com/auth.chatgpt_account_id`,
+/// and `organizations[0].id` in that order, matching the opencode
+/// implementation.
+fn extract_chatgpt_account_id(token: &str) -> Option<String> {
+    let parts: Vec<&str> = token.split('.').collect();
+    if parts.len() != 3 {
+        return None;
+    }
+    use base64::Engine;
+    let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD
+        .decode(parts[1])
+        .ok()?;
+    let claims: serde_json::Value = serde_json::from_slice(&payload).ok()?;
+
+    // Try chatgpt_account_id first
+    if let Some(id) = claims["chatgpt_account_id"].as_str() {
+        return Some(id.to_string());
+    }
+    // Try nested auth claim
+    if let Some(id) = claims["https://api.openai.com/auth"]["chatgpt_account_id"].as_str() {
+        return Some(id.to_string());
+    }
+    // Fall back to organizations[0].id
+    if let Some(id) = claims["organizations"]
+        .as_array()
+        .and_then(|orgs| orgs.first())
+        .and_then(|org| org["id"].as_str())
+    {
+        return Some(id.to_string());
+    }
+    None
+}
+
+/// Adds Codex-specific credential metadata derived from OAuth tokens.
+///
+/// Tries to extract the account ID from the `id_token` first (which typically
+/// contains the user identity claims in OpenID Connect flows), then falls back
+/// to the `access_token` if needed.
+fn enrich_codex_oauth_credential(
+    provider_id: &ProviderId,
+    credential: &mut AuthCredential,
+    id_token: Option<&str>,
+    access_token: &str,
+) {
+    if *provider_id != ProviderId::CODEX {
+        return;
+    }
+
+    // Try id_token first (preferred for user identity claims)
+    let account_id = id_token
+        .and_then(extract_chatgpt_account_id)
+        .or_else(|| extract_chatgpt_account_id(access_token));
+
+    if let Some(account_id) = account_id {
+        credential
+            .url_params
+            .insert("chatgpt_account_id".to_string().into(), account_id.into());
+    }
+}
+
 /// OAuth Code Strategy - Browser redirect flow
 pub struct OAuthCodeStrategy<T> {
     provider_id: ProviderId,
@@ -96,7 +158,7 @@ impl<T: OAuthHttpProvider> AuthStrategy for OAuthCodeStrategy<T> {
                 let token_response = self
                     .adapter
                     .exchange_code(
-                        &self.config,
+                        &ctx.request.oauth_config,
                         ctx.response.code.as_str(),
                         ctx.request.pkce_verifier.as_ref().map(|v| v.as_str()),
                     )
@@ -107,12 +169,21 @@ impl<T: OAuthHttpProvider> AuthStrategy for OAuthCodeStrategy<T> {
                         ))
                     })?;
 
-                build_oauth_credential(
+                let access_token = token_response.access_token.clone();
+                let id_token = token_response.id_token.clone();
+                let mut credential = build_oauth_credential(
                     self.provider_id.clone(),
                     token_response,
-                    &self.config,
+                    &ctx.request.oauth_config,
                     chrono::Duration::hours(1), // Code flow default
-                )
+                )?;
+                enrich_codex_oauth_credential(
+                    &self.provider_id,
+                    &mut credential,
+                    id_token.as_deref(),
+                    &access_token,
+                );
+                Ok(credential)
             }
             _ => Err(AuthError::InvalidContext("Expected Code context".to_string()).into()),
         }
@@ -479,41 +550,6 @@ struct CodexDeviceTokenResponse {
     code_verifier: String,
 }
 
-/// Extract the ChatGPT account ID from a JWT token's claims.
-///
-/// Checks `chatgpt_account_id`, `https://api.openai.com/auth.chatgpt_account_id`,
-/// and `organizations[0].id` in that order, matching the opencode
-/// implementation.
-fn extract_chatgpt_account_id(token: &str) -> Option<String> {
-    let parts: Vec<&str> = token.split('.').collect();
-    if parts.len() != 3 {
-        return None;
-    }
-    use base64::Engine;
-    let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD
-        .decode(parts[1])
-        .ok()?;
-    let claims: serde_json::Value = serde_json::from_slice(&payload).ok()?;
-
-    // Try chatgpt_account_id first
-    if let Some(id) = claims["chatgpt_account_id"].as_str() {
-        return Some(id.to_string());
-    }
-    // Try nested auth claim
-    if let Some(id) = claims["https://api.openai.com/auth"]["chatgpt_account_id"].as_str() {
-        return Some(id.to_string());
-    }
-    // Fall back to organizations[0].id
-    if let Some(id) = claims["organizations"]
-        .as_array()
-        .and_then(|orgs| orgs.first())
-        .and_then(|org| org["id"].as_str())
-    {
-        return Some(id.to_string());
-    }
-    None
-}
-
 #[async_trait::async_trait]
 impl AuthStrategy for CodexDeviceStrategy {
     async fn init(&self) -> anyhow::Result<AuthContextRequest> {
@@ -570,11 +606,8 @@ impl AuthStrategy for CodexDeviceStrategy {
                 // Poll for authorization code using the custom OpenAI endpoint
                 let token_response = codex_poll_for_tokens(&ctx.request, &self.config).await?;
 
-                // Extract ChatGPT account ID from the access token JWT.
-                // This is used for the optional `ChatGPT-Account-Id` request
-                // header when available.
-                let account_id = extract_chatgpt_account_id(&token_response.access_token);
-
+                let access_token = token_response.access_token.clone();
+                let id_token = token_response.id_token.clone();
                 let mut credential = build_oauth_credential(
                     self.provider_id.clone(),
                     token_response,
@@ -583,12 +616,13 @@ impl AuthStrategy for CodexDeviceStrategy {
                 )?;
 
                 // Store account_id in url_params so it's persisted and available
-                // for chat request headers
-                if let Some(id) = account_id {
-                    credential
-                        .url_params
-                        .insert("chatgpt_account_id".to_string().into(), id.into());
-                }
+                // for chat request headers.
+                enrich_codex_oauth_credential(
+                    &self.provider_id,
+                    &mut credential,
+                    id_token.as_deref(),
+                    &access_token,
+                );
 
                 Ok(credential)
             }

crates/forge_infra/src/auth/strategy.rs

@@ -36,6 +36,7 @@ pub(crate) fn into_domain<T: oauth2::TokenResponse>(token: T) -> OAuthTokenRespo
                 .collect::<Vec<_>>()
                 .join(" ")
         }),
+        id_token: None, // oauth2 crate doesn't provide id_token directly
     }
 }
 
@@ -98,6 +99,7 @@ pub(crate) fn build_token_response(
         expires_at: None,
         token_type: "Bearer".to_string(),
         scope: None,
+        id_token: None,
     }
 }
 

crates/forge_infra/src/auth/util.rs

@@ -66,6 +66,7 @@ strip-ansi-escapes.workspace = true
 terminal_size = "0.4"
 rustls.workspace = true
 tempfile.workspace = true
+tiny_http.workspace = true
 
 [target.'cfg(windows)'.dependencies]
 enable-ansi-support.workspace = true

crates/forge_main/Cargo.toml

@@ -7,6 +7,7 @@ mod editor;
 mod info;
 mod input;
 mod model;
+mod oauth_callback;
 mod porcelain;
 mod prompt;
 mod sandbox;