Skip to main content

vtcode_auth/openai_chatgpt_oauth/
session.rs

1//! ChatGPT session model, auth handle, and OAuth flow helpers.
2
3use anyhow::{Context, Result, anyhow, bail};
4use async_trait::async_trait;
5use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
6use reqwest::Client;
7use ring::rand::{SecureRandom, SystemRandom};
8use serde::{Deserialize, Serialize};
9use std::fmt;
10use std::sync::{Arc, Mutex};
11use tokio::sync::Mutex as AsyncMutex;
12
13use crate::{OpenAIAuthConfig, OpenAIPreferredMethod};
14
15pub use super::super::credentials::AuthCredentialsStoreMode;
16use super::super::pkce::PkceChallenge;
17use super::refresh::refresh_openai_chatgpt_session_from_snapshot;
18use super::{
19    OPENAI_AUTH_URL, OPENAI_CALLBACK_PATH, OPENAI_TOKEN_URL, REFRESH_INTERVAL_SECS, REFRESH_SKEW_SECS, now_secs,
20    resolve_oauth_client_identity,
21};
22
23/// Stored OpenAI ChatGPT subscription session.
24///
25/// Custom `Debug` redacts all token fields to prevent credential leakage
26/// through `tracing::debug!(?session)` or error wrappers.
27#[derive(Clone, Serialize, Deserialize)]
28pub struct OpenAIChatGptSession {
29    /// Exchanged OpenAI bearer token used for normal API calls when available.
30    /// If unavailable, VT Code falls back to the OAuth access token.
31    pub openai_api_key: String,
32    /// OAuth ID token from the sign-in flow.
33    pub id_token: String,
34    /// OAuth access token from the sign-in flow.
35    pub access_token: String,
36    /// Refresh token used to renew the session.
37    pub refresh_token: String,
38    /// ChatGPT workspace/account identifier, if present.
39    pub account_id: Option<String>,
40    /// Account email, if present.
41    pub email: Option<String>,
42    /// ChatGPT plan type, if present.
43    pub plan: Option<String>,
44    /// When the session was originally created.
45    pub obtained_at: u64,
46    /// When the OAuth/API-key exchange was last refreshed.
47    pub refreshed_at: u64,
48    /// Access-token expiry, if supplied by the authority.
49    pub expires_at: Option<u64>,
50}
51
52impl fmt::Debug for OpenAIChatGptSession {
53    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
54        f.debug_struct("OpenAIChatGptSession")
55            .field("openai_api_key", &"<redacted>")
56            .field("id_token", &"<redacted>")
57            .field("access_token", &"<redacted>")
58            .field("refresh_token", &"<redacted>")
59            .field("account_id", &self.account_id)
60            .field("email", &self.email)
61            .field("plan", &self.plan)
62            .field("obtained_at", &self.obtained_at)
63            .field("refreshed_at", &self.refreshed_at)
64            .field("expires_at", &self.expires_at)
65            .finish()
66    }
67}
68
69impl OpenAIChatGptSession {
70    pub(super) fn is_refresh_due(&self) -> bool {
71        let now = now_secs();
72        if let Some(expires_at) = self.expires_at
73            && now.saturating_add(REFRESH_SKEW_SECS) >= expires_at
74        {
75            return true;
76        }
77        now.saturating_sub(self.refreshed_at) >= REFRESH_INTERVAL_SECS
78    }
79}
80
81/// Host-provided refresher for externally managed ChatGPT auth tokens.
82#[async_trait]
83pub trait OpenAIChatGptSessionRefresher: Send + Sync {
84    async fn refresh_session(&self, current: &OpenAIChatGptSession) -> Result<OpenAIChatGptSession>;
85}
86
87#[derive(Clone)]
88enum OpenAIChatGptAuthRefreshStrategy {
89    Stored {
90        storage_mode: AuthCredentialsStoreMode,
91    },
92    External {
93        refresher: Arc<dyn OpenAIChatGptSessionRefresher>,
94    },
95}
96
97impl fmt::Debug for OpenAIChatGptAuthRefreshStrategy {
98    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
99        match self {
100            Self::Stored { storage_mode } => f.debug_struct("Stored").field("storage_mode", storage_mode).finish(),
101            Self::External { .. } => f.debug_struct("External").finish_non_exhaustive(),
102        }
103    }
104}
105
106/// Runtime auth state shared by OpenAI provider instances.
107#[derive(Clone)]
108pub struct OpenAIChatGptAuthHandle {
109    session: Arc<Mutex<OpenAIChatGptSession>>,
110    refresh_gate: Arc<AsyncMutex<()>>,
111    auto_refresh: bool,
112    refresh_strategy: OpenAIChatGptAuthRefreshStrategy,
113}
114
115impl fmt::Debug for OpenAIChatGptAuthHandle {
116    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
117        f.debug_struct("OpenAIChatGptAuthHandle")
118            .field("auto_refresh", &self.auto_refresh)
119            .field("refresh_strategy", &self.refresh_strategy)
120            .finish()
121    }
122}
123
124impl OpenAIChatGptAuthHandle {
125    pub fn new(
126        session: OpenAIChatGptSession,
127        auth_config: OpenAIAuthConfig,
128        storage_mode: AuthCredentialsStoreMode,
129    ) -> Self {
130        Self {
131            session: Arc::new(Mutex::new(session)),
132            refresh_gate: Arc::new(AsyncMutex::new(())),
133            auto_refresh: auth_config.auto_refresh,
134            refresh_strategy: OpenAIChatGptAuthRefreshStrategy::Stored { storage_mode },
135        }
136    }
137
138    pub fn new_external(
139        session: OpenAIChatGptSession,
140        auto_refresh: bool,
141        refresher: Arc<dyn OpenAIChatGptSessionRefresher>,
142    ) -> Self {
143        Self {
144            session: Arc::new(Mutex::new(session)),
145            refresh_gate: Arc::new(AsyncMutex::new(())),
146            auto_refresh,
147            refresh_strategy: OpenAIChatGptAuthRefreshStrategy::External { refresher },
148        }
149    }
150
151    pub fn snapshot(&self) -> Result<OpenAIChatGptSession> {
152        self.session
153            .lock()
154            .map(|guard| guard.clone())
155            .map_err(|_| anyhow!("openai chatgpt auth mutex poisoned"))
156    }
157
158    pub fn current_api_key(&self) -> Result<String> {
159        self.snapshot().map(|session| active_api_bearer_token(&session).to_string())
160    }
161
162    pub fn provider_label(&self) -> &'static str {
163        "OpenAI (ChatGPT)"
164    }
165
166    pub async fn refresh_if_needed(&self) -> Result<()> {
167        if !self.auto_refresh {
168            return Ok(());
169        }
170
171        self.refresh_when(|session| session.is_refresh_due()).await
172    }
173
174    pub async fn force_refresh(&self) -> Result<()> {
175        self.refresh_when(|_| true).await
176    }
177
178    async fn refresh_when<P>(&self, should_refresh: P) -> Result<()>
179    where
180        P: FnOnce(&OpenAIChatGptSession) -> bool,
181    {
182        let _refresh_guard = self.refresh_gate.lock().await;
183        let session = self.snapshot()?;
184        if !should_refresh(&session) {
185            return Ok(());
186        }
187
188        let refreshed = match &self.refresh_strategy {
189            OpenAIChatGptAuthRefreshStrategy::Stored { storage_mode } => {
190                refresh_openai_chatgpt_session_from_snapshot(&session, *storage_mode).await?
191            }
192            OpenAIChatGptAuthRefreshStrategy::External { refresher } => refresher.refresh_session(&session).await?,
193        };
194        self.replace_session(refreshed)
195    }
196
197    #[must_use]
198    pub(super) fn using_external_tokens(&self) -> bool {
199        matches!(self.refresh_strategy, OpenAIChatGptAuthRefreshStrategy::External { .. })
200    }
201
202    fn replace_session(&self, session: OpenAIChatGptSession) -> Result<()> {
203        let mut guard = self.session.lock().map_err(|_| anyhow!("openai chatgpt auth mutex poisoned"))?;
204        *guard = session;
205        Ok(())
206    }
207}
208
209/// OpenAI auth resolution chosen for the current runtime.
210///
211/// Custom `Debug` redacts the bearer `api_key` to prevent credential leakage.
212#[derive(Clone)]
213pub enum OpenAIResolvedAuth {
214    ApiKey {
215        api_key: String,
216    },
217    ChatGpt {
218        api_key: String,
219        handle: OpenAIChatGptAuthHandle,
220    },
221}
222
223impl fmt::Debug for OpenAIResolvedAuth {
224    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
225        match self {
226            Self::ApiKey { .. } => f.debug_struct("OpenAIResolvedAuth::ApiKey").finish(),
227            Self::ChatGpt { .. } => f.debug_struct("OpenAIResolvedAuth::ChatGpt").finish(),
228        }
229    }
230}
231
232impl OpenAIResolvedAuth {
233    pub fn api_key(&self) -> &str {
234        match self {
235            Self::ApiKey { api_key } => api_key,
236            Self::ChatGpt { api_key, .. } => api_key,
237        }
238    }
239
240    pub fn handle(&self) -> Option<OpenAIChatGptAuthHandle> {
241        match self {
242            Self::ApiKey { .. } => None,
243            Self::ChatGpt { handle, .. } => Some(handle.clone()),
244        }
245    }
246
247    pub(super) fn using_chatgpt(&self) -> bool {
248        matches!(self, Self::ChatGpt { .. })
249    }
250}
251
252pub(crate) fn active_api_bearer_token(session: &OpenAIChatGptSession) -> &str {
253    if session.openai_api_key.trim().is_empty() {
254        session.access_token.as_str()
255    } else {
256        session.openai_api_key.as_str()
257    }
258}
259
260#[derive(Debug, Clone, Copy, PartialEq, Eq)]
261pub enum OpenAIResolvedAuthSource {
262    ApiKey,
263    ChatGpt,
264}
265
266/// Where the ChatGPT session originated — used by CLI/TUI to render accurate
267/// status without directly inspecting the filesystem.
268#[derive(Debug, Clone, Copy, PartialEq, Eq)]
269pub enum OpenAIChatGptSessionProvenance {
270    /// Session stored in VT Code's own credential storage (full auto-refresh).
271    Native,
272    /// Session loaded from Codex CLI's `~/.codex/auth.json` (managed by Codex).
273    CodexFallback,
274}
275
276/// Redacted summary of available OpenAI credentials for CLI/TUI display.
277///
278/// Does NOT carry token data — only metadata (email, plan, provenance, expiry)
279/// so credential values can never leak through `Debug` or logging.
280#[derive(Debug, Clone)]
281pub struct OpenAICredentialOverview {
282    pub api_key_available: bool,
283    /// Email from the ChatGPT session's ID token, if available.
284    pub chatgpt_email: Option<String>,
285    /// Plan type from the ChatGPT session's ID token, if available.
286    pub chatgpt_plan: Option<String>,
287    /// `true` when a ChatGPT session (native or Codex fallback) is available.
288    pub chatgpt_session_present: bool,
289    /// Provenance of the ChatGPT session — `None` when no session is available.
290    pub chatgpt_session_provenance: Option<OpenAIChatGptSessionProvenance>,
291    /// `true` only when Codex's auth.json was **successfully parsed** into a
292    /// usable session (not merely that the file exists on disk).
293    pub codex_fallback_available: bool,
294    pub active_source: Option<OpenAIResolvedAuthSource>,
295    pub preferred_method: OpenAIPreferredMethod,
296    pub notice: Option<String>,
297    pub recommendation: Option<String>,
298}
299
300/// Generic auth status reused by slash auth/status output.
301#[derive(Debug, Clone)]
302pub enum OpenAIChatGptAuthStatus {
303    Authenticated {
304        label: Option<String>,
305        age_seconds: u64,
306        expires_in: Option<u64>,
307    },
308    NotAuthenticated,
309}
310/// Build the OpenAI ChatGPT OAuth authorization URL.
311pub fn get_openai_chatgpt_auth_url(challenge: &PkceChallenge, callback_port: u16, state: &str) -> Result<String> {
312    let redirect_uri = format!("http://localhost:{callback_port}{OPENAI_CALLBACK_PATH}");
313    let identity = resolve_oauth_client_identity()?;
314    let query = [
315        ("response_type", "code".to_string()),
316        ("client_id", identity.client_id.clone()),
317        ("redirect_uri", redirect_uri),
318        ("scope", "openid profile email offline_access api.connectors.read api.connectors.invoke".to_string()),
319        ("code_challenge", challenge.code_challenge.clone()),
320        ("code_challenge_method", challenge.code_challenge_method.clone()),
321        ("id_token_add_organizations", "true".to_string()),
322        ("codex_cli_simplified_flow", "true".to_string()),
323        ("state", state.to_string()),
324        ("originator", identity.originator),
325    ];
326
327    let encoded = query
328        .iter()
329        .map(|(key, value)| format!("{key}={}", urlencoding::encode(value)))
330        .collect::<Vec<_>>()
331        .join("&");
332    Ok(format!("{OPENAI_AUTH_URL}?{encoded}"))
333}
334
335pub fn generate_openai_oauth_state() -> Result<String> {
336    let mut state_bytes = [0_u8; 32];
337    SystemRandom::new()
338        .fill(&mut state_bytes)
339        .map_err(|_| anyhow!("failed to generate openai oauth state"))?;
340    Ok(URL_SAFE_NO_PAD.encode(state_bytes))
341}
342
343pub fn parse_openai_chatgpt_manual_callback_input(input: &str, expected_state: &str) -> Result<String> {
344    let trimmed = input.trim();
345    if trimmed.is_empty() {
346        bail!("missing authorization callback input");
347    }
348
349    let query = if trimmed.contains("://") {
350        let url = reqwest::Url::parse(trimmed).context("invalid callback url")?;
351        url.query()
352            .ok_or_else(|| anyhow!("callback url did not include a query string"))?
353            .to_string()
354    } else if trimmed.contains('=') {
355        trimmed.trim_start_matches('?').to_string()
356    } else {
357        bail!("paste the full redirect url or query string containing code and state");
358    };
359
360    let code = extract_query_value(&query, "code")
361        .ok_or_else(|| anyhow!("callback input did not include an authorization code"))?;
362    let state = extract_query_value(&query, "state").ok_or_else(|| anyhow!("callback input did not include state"))?;
363    if state != expected_state {
364        bail!("OAuth error: state mismatch");
365    }
366    Ok(code)
367}
368/// Resolve the active OpenAI auth source for the current configuration.
369pub fn resolve_openai_auth(
370    auth_config: &OpenAIAuthConfig,
371    storage_mode: AuthCredentialsStoreMode,
372    api_key: Option<String>,
373) -> Result<OpenAIResolvedAuth> {
374    crate::auth_service::OpenAIAccountAuthService::new(auth_config.clone(), storage_mode).resolve_runtime_auth(api_key)
375}
376
377pub fn summarize_openai_credentials(
378    auth_config: &OpenAIAuthConfig,
379    storage_mode: AuthCredentialsStoreMode,
380    api_key: Option<String>,
381) -> Result<OpenAICredentialOverview> {
382    crate::auth_service::OpenAIAccountAuthService::new(auth_config.clone(), storage_mode).summarize_credentials(api_key)
383}
384fn extract_query_value(query: &str, key: &str) -> Option<String> {
385    query
386        .trim_start_matches('?')
387        .split('&')
388        .filter_map(|pair| {
389            let (pair_key, pair_value) = pair.split_once('=')?;
390            (pair_key == key)
391                .then(|| urlencoding::decode(pair_value).ok().map(|value| value.into_owned()))
392                .flatten()
393        })
394        .find(|value| !value.is_empty())
395}