llm/providers/codex/
oauth.rs

1use crate::LlmError;
2use aether_auth::{
3    BrowserOAuthHandler, OAuthCredential, OAuthCredentialStorage, OAuthError, OAuthHandler, oauth_http_client,
4};
5use base64::Engine;
6use base64::engine::general_purpose::URL_SAFE_NO_PAD;
7use oauth2::basic::BasicClient;
8use oauth2::{AuthUrl, AuthorizationCode, ClientId, PkceCodeChallenge, RedirectUrl, TokenUrl};
9use std::sync::Arc;
10use tokio::sync::Mutex;
11use url::Url;
12
13const CLIENT_ID: &str = "app_EMoamEEZ73f0CkXaXp7hrann";
14const AUTHORIZE_URL: &str = "https://auth.openai.com/oauth/authorize";
15const TOKEN_URL: &str = "https://auth.openai.com/oauth/token";
16const REDIRECT_URI: &str = "http://localhost:1455/auth/callback";
17const SCOPE: &str = "openid profile email offline_access";
18
19/// Run the full Codex OAuth flow: open browser, capture callback, exchange token, save credentials.
20///
21pub async fn perform_codex_oauth_flow(store: &dyn OAuthCredentialStorage) -> Result<(), LlmError> {
22    let (pkce_challenge, pkce_verifier) = PkceCodeChallenge::new_random_sha256();
23    let state = generate_random_state();
24
25    let auth_url = Url::parse_with_params(
26        AUTHORIZE_URL,
27        &[
28            ("response_type", "code"),
29            ("client_id", CLIENT_ID),
30            ("redirect_uri", REDIRECT_URI),
31            ("scope", SCOPE),
32            ("code_challenge", pkce_challenge.as_str()),
33            ("code_challenge_method", "S256"),
34            ("state", &state),
35            ("id_token_add_organizations", "true"),
36            ("codex_cli_simplified_flow", "true"),
37            ("originator", "codex_cli_rs"),
38        ],
39    )
40    .map_err(|e| OAuthError::TokenExchange(format!("Failed to build auth URL: {e}")))?;
41
42    // Port 1455 is hardcoded because the Codex API has a fixed redirect URI
43    // (http://localhost:1455/auth/callback) registered with OpenAI's OAuth server.
44    let handler = BrowserOAuthHandler::with_redirect_uri(REDIRECT_URI, 1455)?;
45    let callback_url = handler.authorize(auth_url.as_str()).await?;
46    let callback = Url::parse(&callback_url)
47        .map_err(|error| OAuthError::InvalidCallback(format!("Invalid callback URL: {error}")))?;
48    let mut code = None;
49    let mut callback_state = None;
50    for (name, value) in callback.query_pairs() {
51        match name.as_ref() {
52            "code" => code = Some(value.into_owned()),
53            "state" => callback_state = Some(value.into_owned()),
54            _ => {}
55        }
56    }
57    if callback_state.as_deref() != Some(&state) {
58        return Err(OAuthError::StateMismatch.into());
59    }
60    let code = code.ok_or_else(|| OAuthError::InvalidCallback("No authorization code in callback".to_string()))?;
61
62    let oauth_client = BasicClient::new(ClientId::new(CLIENT_ID.to_string()))
63        .set_auth_uri(
64            AuthUrl::new(AUTHORIZE_URL.to_string())
65                .map_err(|e| OAuthError::TokenExchange(format!("invalid auth URL: {e}")))?,
66        )
67        .set_token_uri(
68            TokenUrl::new(TOKEN_URL.to_string())
69                .map_err(|e| OAuthError::TokenExchange(format!("invalid token URL: {e}")))?,
70        )
71        .set_redirect_uri(
72            RedirectUrl::new(REDIRECT_URI.to_string())
73                .map_err(|e| OAuthError::TokenExchange(format!("invalid redirect URI: {e}")))?,
74        );
75
76    let http_client = oauth_http_client()?;
77
78    let token_response = oauth_client
79        .exchange_code(AuthorizationCode::new(code))
80        .set_pkce_verifier(pkce_verifier)
81        .request_async(&http_client)
82        .await
83        .map_err(|e| OAuthError::TokenExchange(e.to_string()))?;
84
85    let credential = OAuthCredential::from_token_response(CLIENT_ID.to_string(), &token_response);
86    store.save_credential(super::PROVIDER_ID, credential).await?;
87
88    Ok(())
89}
90
91/// In-memory cache of the most recently validated credential and its derived account ID.
92struct CachedToken {
93    credential: OAuthCredential,
94    account_id: String,
95}
96
97/// Manages OAuth tokens for the Codex backend API.
98///
99/// Holds an `Arc<dyn OAuthCredentialStorage>` so callers can swap in keyring-backed,
100/// file-backed, or in-memory stores without changing this type.
101pub struct CodexTokenManager {
102    store: Arc<dyn OAuthCredentialStorage>,
103    credential_key: String,
104    token_url: TokenUrl,
105    cached: Mutex<Option<CachedToken>>,
106}
107
108impl CodexTokenManager {
109    pub fn new(store: Arc<dyn OAuthCredentialStorage>, credential_key: &str) -> Self {
110        Self::new_with_token_url(
111            store,
112            credential_key,
113            TokenUrl::new(TOKEN_URL.to_string()).expect("hardcoded Codex token URL is valid"),
114        )
115    }
116
117    fn new_with_token_url(store: Arc<dyn OAuthCredentialStorage>, credential_key: &str, token_url: TokenUrl) -> Self {
118        Self { store, credential_key: credential_key.to_string(), token_url, cached: Mutex::new(None) }
119    }
120
121    /// Get a valid access token and account ID.
122    ///
123    /// Returns `(access_token, account_id)`. The account ID is extracted from
124    /// the JWT's `https://api.openai.com/auth` claim field `chatgpt_account_id`.
125    pub async fn get_valid_token(&self) -> Result<(String, String), LlmError> {
126        let mut cache = self.cached.lock().await;
127        if let Some(cached) = cache.as_ref()
128            && !cached.credential.needs_refresh()
129        {
130            return Ok((cached.credential.access_token.clone(), cached.account_id.clone()));
131        }
132
133        let credential = self.load_or_refresh().await?;
134        let account_id = extract_account_id(&credential.access_token)?;
135        let access_token = credential.access_token.clone();
136        *cache = Some(CachedToken { credential, account_id: account_id.clone() });
137        Ok((access_token, account_id))
138    }
139
140    async fn load_or_refresh(&self) -> Result<OAuthCredential, LlmError> {
141        let stored = self.store.load_credential(&self.credential_key).await?.ok_or_else(|| {
142            OAuthError::NoCredentials(
143                "No Codex OAuth credentials found. Run `aether` and select a codex model to trigger OAuth login."
144                    .to_string(),
145            )
146        })?;
147
148        if !stored.needs_refresh() {
149            return Ok(stored);
150        }
151
152        let refreshed = stored.refresh(&self.token_url).await?;
153        self.store.save_credential(&self.credential_key, refreshed.clone()).await?;
154        Ok(refreshed)
155    }
156
157    /// Clear the cached token (e.g. after a 401 response)
158    pub async fn clear_cache(&self) {
159        *self.cached.lock().await = None;
160    }
161}
162
163/// Extract the account ID from a JWT access token.
164///
165/// The JWT payload contains a claim at `https://api.openai.com/auth`
166/// with a `chatgpt_account_id` field.
167pub fn extract_account_id(access_token: &str) -> Result<String, LlmError> {
168    let parts: Vec<&str> = access_token.split('.').collect();
169    if parts.len() != 3 {
170        return Err(OAuthError::InvalidJwt("expected 3 dot-separated parts".to_string()).into());
171    }
172
173    let decoded = URL_SAFE_NO_PAD
174        .decode(parts[1])
175        .map_err(|e| OAuthError::InvalidJwt(format!("failed to decode payload: {e}")))?;
176
177    let payload: serde_json::Value = serde_json::from_slice(&decoded)
178        .map_err(|e| OAuthError::InvalidJwt(format!("failed to parse payload: {e}")))?;
179
180    let account_id = payload
181        .get("https://api.openai.com/auth")
182        .and_then(|auth| auth.get("chatgpt_account_id"))
183        .and_then(|v| v.as_str())
184        .ok_or_else(|| OAuthError::InvalidJwt("missing chatgpt_account_id in token".to_string()))?;
185
186    Ok(account_id.to_string())
187}
188
189fn generate_random_state() -> String {
190    uuid::Uuid::new_v4().to_string()
191}
192
193#[cfg(test)]
194mod tests {
195    use super::*;
196    use aether_auth::{FakeOAuthCredentialStore, OAuthCredential};
197    use axum::Router;
198    use axum::body::{Body, to_bytes};
199    use axum::extract::State;
200    use axum::http::{HeaderMap, Method, Request, StatusCode};
201    use axum::response::IntoResponse;
202    use axum::routing::post;
203    use std::collections::HashMap;
204    use tokio::net::TcpListener;
205    use tokio::sync::{Mutex as TokioMutex, oneshot};
206
207    /// Create a test JWT with a given payload
208    fn make_test_jwt(payload: &serde_json::Value) -> String {
209        let header = URL_SAFE_NO_PAD.encode(r#"{"alg":"RS256","typ":"JWT"}"#);
210        let payload_json = serde_json::to_string(payload).unwrap();
211        let payload_b64url = URL_SAFE_NO_PAD.encode(payload_json.as_bytes());
212        format!("{header}.{payload_b64url}.fake_signature")
213    }
214
215    #[test]
216    fn extract_account_id_from_valid_jwt() {
217        let payload = serde_json::json!({
218            "sub": "user_123",
219            "https://api.openai.com/auth": {
220                "chatgpt_account_id": "acct_abc123"
221            }
222        });
223
224        let jwt = make_test_jwt(&payload);
225        let account_id = extract_account_id(&jwt).unwrap();
226        assert_eq!(account_id, "acct_abc123");
227    }
228
229    #[test]
230    fn extract_account_id_missing_claim() {
231        let payload = serde_json::json!({
232            "sub": "user_123"
233        });
234
235        let jwt = make_test_jwt(&payload);
236        let result = extract_account_id(&jwt);
237        assert!(result.is_err());
238        assert!(result.unwrap_err().to_string().contains("chatgpt_account_id"));
239    }
240
241    #[test]
242    fn extract_account_id_invalid_jwt_format() {
243        let result = extract_account_id("not.a.valid.jwt.too.many.parts");
244        assert!(result.is_err());
245
246        let result = extract_account_id("toofewparts");
247        assert!(result.is_err());
248    }
249
250    #[test]
251    fn extract_account_id_invalid_base64() {
252        let result = extract_account_id("header.!!!invalid!!!.signature");
253        assert!(result.is_err());
254    }
255
256    #[test]
257    fn auth_url_is_well_formed() {
258        let (pkce_challenge, _) = PkceCodeChallenge::new_random_sha256();
259        let state = "test-state";
260
261        let auth_url = Url::parse_with_params(
262            AUTHORIZE_URL,
263            &[
264                ("response_type", "code"),
265                ("client_id", CLIENT_ID),
266                ("redirect_uri", REDIRECT_URI),
267                ("scope", SCOPE),
268                ("code_challenge", pkce_challenge.as_str()),
269                ("code_challenge_method", "S256"),
270                ("state", state),
271                ("id_token_add_organizations", "true"),
272                ("codex_cli_simplified_flow", "true"),
273                ("originator", "codex_cli_rs"),
274            ],
275        )
276        .unwrap();
277
278        let url_str = auth_url.as_str();
279        assert!(url_str.starts_with(AUTHORIZE_URL));
280        assert!(url_str.contains("client_id="));
281        assert!(url_str.contains("redirect_uri="));
282        assert!(url_str.contains("scope="));
283        assert!(url_str.contains("code_challenge="));
284        assert!(url_str.contains("state=test-state"));
285    }
286
287    #[test]
288    fn generate_random_state_is_valid_uuid() {
289        let state = generate_random_state();
290        assert!(!state.is_empty());
291        assert!(uuid::Uuid::parse_str(&state).is_ok());
292    }
293
294    #[test]
295    fn oauth_constants_are_valid() {
296        assert!(AUTHORIZE_URL.starts_with("https://"));
297        assert!(TOKEN_URL.starts_with("https://"));
298        assert!(REDIRECT_URI.starts_with("http://localhost:"));
299        assert!(SCOPE.contains("openid"));
300    }
301
302    #[tokio::test]
303    async fn codex_token_manager_refreshes_expired_credential() {
304        let new_access_token = test_jwt_for_account("acct_new");
305        let endpoint = FakeTokenEndpoint::start(TokenEndpointResponse::success(&new_access_token, None)).await;
306        let store = Arc::new(
307            FakeOAuthCredentialStore::new()
308                .with_credential("codex", expired_credential("old-access", Some("refresh-old"))),
309        );
310
311        let manager = CodexTokenManager::new_with_token_url(store.clone(), "codex", endpoint.url.clone());
312        let (access_token, account_id) = manager.get_valid_token().await.unwrap();
313        let request = endpoint.request.await.expect("token endpoint request");
314        let saved = store.load_credential("codex").await.unwrap().unwrap();
315
316        assert_eq!(access_token, new_access_token);
317        assert_eq!(account_id, "acct_new");
318        assert_eq!(saved.access_token, new_access_token);
319        assert_eq!(saved.refresh_token.as_deref(), Some("refresh-old"));
320        assert_eq!(request.method, Method::POST);
321        assert_eq!(request.path, "/oauth/token");
322        assert_eq!(request.form.get("grant_type").map(String::as_str), Some("refresh_token"));
323        assert_eq!(request.form.get("refresh_token").map(String::as_str), Some("refresh-old"));
324        assert_eq!(request.form.get("client_id").map(String::as_str), Some(CLIENT_ID));
325        assert!(request.headers.get("accept").is_some());
326    }
327
328    #[tokio::test]
329    async fn codex_token_manager_saves_rotated_refresh_token() {
330        let new_access_token = test_jwt_for_account("acct_new");
331        let endpoint =
332            FakeTokenEndpoint::start(TokenEndpointResponse::success(&new_access_token, Some("refresh-new"))).await;
333        let store = Arc::new(
334            FakeOAuthCredentialStore::new()
335                .with_credential("codex", expired_credential("old-access", Some("refresh-old"))),
336        );
337        let manager = CodexTokenManager::new_with_token_url(store.clone(), "codex", endpoint.url.clone());
338        manager.get_valid_token().await.unwrap();
339        let saved = store.load_credential("codex").await.unwrap().unwrap();
340
341        assert_eq!(saved.access_token, new_access_token);
342        assert_eq!(saved.refresh_token.as_deref(), Some("refresh-new"));
343    }
344
345    #[tokio::test]
346    async fn codex_token_manager_uses_unexpired_credential_without_refresh() {
347        let access_token = test_jwt_for_account("acct_existing");
348        let store = Arc::new(FakeOAuthCredentialStore::new().with_credential(
349            "codex",
350            OAuthCredential {
351                client_id: CLIENT_ID.to_string(),
352                access_token: access_token.clone(),
353                refresh_token: Some("refresh-old".to_string()),
354                expires_at: Some(u64::MAX),
355            },
356        ));
357
358        let manager = CodexTokenManager::new_with_token_url(
359            store,
360            "codex",
361            TokenUrl::new("http://127.0.0.1:9/oauth/token".to_string()).unwrap(),
362        );
363
364        let (returned_token, account_id) = manager.get_valid_token().await.unwrap();
365        assert_eq!(returned_token, access_token);
366        assert_eq!(account_id, "acct_existing");
367    }
368
369    #[tokio::test]
370    async fn codex_token_manager_errors_when_credential_is_missing() {
371        let store = Arc::new(FakeOAuthCredentialStore::new());
372        let manager = CodexTokenManager::new_with_token_url(
373            store,
374            "codex",
375            TokenUrl::new("http://127.0.0.1:9/oauth/token".to_string()).unwrap(),
376        );
377
378        let error = manager.get_valid_token().await.unwrap_err();
379        assert!(error.to_string().contains("No Codex OAuth credentials found"));
380        assert!(error.to_string().contains("select a codex model"));
381    }
382
383    #[tokio::test]
384    async fn codex_token_manager_errors_when_expired_without_refresh_token() {
385        let original = expired_credential("old-access", None);
386        let store = Arc::new(FakeOAuthCredentialStore::new().with_credential("codex", original.clone()));
387        let manager = CodexTokenManager::new_with_token_url(
388            store.clone(),
389            "codex",
390            TokenUrl::new("http://127.0.0.1:9/oauth/token".to_string()).unwrap(),
391        );
392
393        let error = manager.get_valid_token().await.unwrap_err();
394        let saved = store.load_credential("codex").await.unwrap().unwrap();
395
396        assert!(error.to_string().contains("Re-run OAuth login"));
397        assert_eq!(saved.access_token, original.access_token);
398        assert_eq!(saved.refresh_token, original.refresh_token);
399    }
400
401    #[tokio::test]
402    async fn codex_token_manager_does_not_overwrite_credential_when_refresh_fails() {
403        let endpoint = FakeTokenEndpoint::start(TokenEndpointResponse::failure()).await;
404        let original = expired_credential("old-access", Some("refresh-old"));
405        let store = Arc::new(FakeOAuthCredentialStore::new().with_credential("codex", original.clone()));
406        let manager = CodexTokenManager::new_with_token_url(store.clone(), "codex", endpoint.url.clone());
407
408        let result = manager.get_valid_token().await;
409        let saved = store.load_credential("codex").await.unwrap().unwrap();
410
411        assert!(result.is_err());
412        assert_eq!(saved.access_token, original.access_token);
413        assert_eq!(saved.refresh_token, original.refresh_token);
414    }
415
416    struct FakeTokenEndpoint {
417        url: TokenUrl,
418        request: oneshot::Receiver<CapturedTokenRequest>,
419    }
420
421    struct CapturedTokenRequest {
422        method: Method,
423        path: String,
424        headers: HeaderMap,
425        form: HashMap<String, String>,
426    }
427
428    #[derive(Clone)]
429    struct FakeTokenState {
430        response: TokenEndpointResponse,
431        request_tx: Arc<TokioMutex<Option<oneshot::Sender<CapturedTokenRequest>>>>,
432        shutdown_tx: Arc<TokioMutex<Option<oneshot::Sender<()>>>>,
433    }
434
435    #[derive(Clone)]
436    struct TokenEndpointResponse {
437        status: StatusCode,
438        body: serde_json::Value,
439    }
440
441    impl TokenEndpointResponse {
442        fn success(access_token: &str, refresh_token: Option<&str>) -> Self {
443            let mut body = serde_json::json!({
444                "access_token": access_token,
445                "token_type": "Bearer",
446                "expires_in": 3600
447            });
448            if let Some(refresh_token) = refresh_token {
449                body["refresh_token"] = serde_json::Value::String(refresh_token.to_string());
450            }
451            Self { status: StatusCode::OK, body }
452        }
453
454        fn failure() -> Self {
455            Self { status: StatusCode::BAD_REQUEST, body: serde_json::json!({ "error": "invalid_grant" }) }
456        }
457    }
458
459    impl FakeTokenEndpoint {
460        async fn start(response: TokenEndpointResponse) -> Self {
461            let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind fake token endpoint");
462            let url = TokenUrl::new(format!(
463                "http://{}/oauth/token",
464                listener.local_addr().expect("fake token endpoint address")
465            ))
466            .expect("fake token endpoint URL is valid");
467            let (request_tx, request) = oneshot::channel();
468            let (shutdown_tx, shutdown) = oneshot::channel();
469            let state = FakeTokenState {
470                response,
471                request_tx: Arc::new(TokioMutex::new(Some(request_tx))),
472                shutdown_tx: Arc::new(TokioMutex::new(Some(shutdown_tx))),
473            };
474            let app = Router::new().route("/oauth/token", post(capture_token_request)).with_state(state);
475            tokio::spawn(async move {
476                axum::serve(listener, app)
477                    .with_graceful_shutdown(async {
478                        let _ = shutdown.await;
479                    })
480                    .await
481                    .expect("serve fake token endpoint");
482            });
483            Self { url, request }
484        }
485    }
486
487    async fn capture_token_request(State(state): State<FakeTokenState>, request: Request<Body>) -> impl IntoResponse {
488        let (parts, body) = request.into_parts();
489        let body = to_bytes(body, usize::MAX).await.expect("read token request body");
490        let form = url::form_urlencoded::parse(&body).into_owned().collect();
491        if let Some(tx) = state.request_tx.lock().await.take() {
492            let _ = tx.send(CapturedTokenRequest {
493                method: parts.method,
494                path: parts.uri.path().to_string(),
495                headers: parts.headers,
496                form,
497            });
498        }
499        if let Some(tx) = state.shutdown_tx.lock().await.take() {
500            let _ = tx.send(());
501        }
502        (state.response.status, axum::Json(state.response.body))
503    }
504
505    fn expired_credential(access_token: &str, refresh_token: Option<&str>) -> OAuthCredential {
506        OAuthCredential {
507            client_id: CLIENT_ID.to_string(),
508            access_token: access_token.to_string(),
509            refresh_token: refresh_token.map(str::to_string),
510            expires_at: Some(0),
511        }
512    }
513
514    fn test_jwt_for_account(account_id: &str) -> String {
515        make_test_jwt(&serde_json::json!({
516            "https://api.openai.com/auth": {
517                "chatgpt_account_id": account_id
518            }
519        }))
520    }
521}
llm/providers/codex/oauth.rs

llm/providers/codex/
oauth.rs
