vtcode_auth/openai_chatgpt_oauth/
session.rs1use 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#[derive(Clone, Serialize, Deserialize)]
28pub struct OpenAIChatGptSession {
29 pub openai_api_key: String,
32 pub id_token: String,
34 pub access_token: String,
36 pub refresh_token: String,
38 pub account_id: Option<String>,
40 pub email: Option<String>,
42 pub plan: Option<String>,
44 pub obtained_at: u64,
46 pub refreshed_at: u64,
48 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#[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#[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#[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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
269pub enum OpenAIChatGptSessionProvenance {
270 Native,
272 CodexFallback,
274}
275
276#[derive(Debug, Clone)]
281pub struct OpenAICredentialOverview {
282 pub api_key_available: bool,
283 pub chatgpt_email: Option<String>,
285 pub chatgpt_plan: Option<String>,
287 pub chatgpt_session_present: bool,
289 pub chatgpt_session_provenance: Option<OpenAIChatGptSessionProvenance>,
291 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#[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}
310pub 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}
368pub 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}