diff --git a/Cargo.lock b/Cargo.lock index 2a8077b622..ff1eba29fe 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9887,10 +9887,10 @@ dependencies = [ "log", "menu", "mistral", - "oauth_callback_server", "ollama", "open_ai", "open_router", + "openai_subscribed", "opencode", "parking_lot", "pretty_assertions", @@ -9900,13 +9900,10 @@ dependencies = [ "serde", "serde_json", "settings", - "sha2", - "smol", "strum 0.27.2", "tokio", "ui", "ui_input", - "url", "util", "x_ai", ] @@ -12464,6 +12461,30 @@ dependencies = [ "thiserror 2.0.17", ] +[[package]] +name = "openai_subscribed" +version = "0.1.0" +dependencies = [ + "anyhow", + "base64 0.22.1", + "credentials_provider", + "futures 0.3.32", + "gpui", + "http_client", + "language_model", + "log", + "oauth_callback_server", + "open_ai", + "parking_lot", + "rand 0.9.4", + "serde", + "serde_json", + "sha2", + "smol", + "url", + "util", +] + [[package]] name = "opencode" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index 6c1a3251ab..a1a2e71087 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -145,6 +145,7 @@ members = [ "crates/oauth_callback_server", "crates/ollama", "crates/onboarding", + "crates/openai_subscribed", "crates/opencode", "crates/open_ai", "crates/open_path_prompt", @@ -410,6 +411,7 @@ notifications = { path = "crates/notifications" } oauth_callback_server = { path = "crates/oauth_callback_server" } ollama = { path = "crates/ollama" } onboarding = { path = "crates/onboarding" } +openai_subscribed = { path = "crates/openai_subscribed" } opencode = { path = "crates/opencode" } open_ai = { path = "crates/open_ai" } open_path_prompt = { path = "crates/open_path_prompt" } diff --git a/crates/language_models/Cargo.toml b/crates/language_models/Cargo.toml index 8da9eb332e..dbefcd9a76 100644 --- a/crates/language_models/Cargo.toml +++ b/crates/language_models/Cargo.toml @@ -49,24 +49,21 @@ lmstudio = { workspace = true, features = ["schemars"] } log.workspace = true menu.workspace = true mistral = { workspace = true, features = ["schemars"] } -oauth_callback_server.workspace = true ollama = { workspace = true, features = ["schemars"] } open_ai = { workspace = true, features = ["schemars"] } +openai_subscribed.workspace = true opencode = { workspace = true, features = ["schemars"] } open_router = { workspace = true, features = ["schemars"] } rand.workspace = true release_channel.workspace = true schemars.workspace = true -sha2.workspace = true serde.workspace = true serde_json.workspace = true settings.workspace = true -smol.workspace = true strum.workspace = true tokio = { workspace = true, features = ["rt", "rt-multi-thread"] } ui.workspace = true ui_input.workspace = true -url.workspace = true util.workspace = true x_ai = { workspace = true, features = ["schemars"] } diff --git a/crates/language_models/src/provider/openai_subscribed.rs b/crates/language_models/src/provider/openai_subscribed.rs index e67c9f3b65..2b57c537ce 100644 --- a/crates/language_models/src/provider/openai_subscribed.rs +++ b/crates/language_models/src/provider/openai_subscribed.rs @@ -1,104 +1,21 @@ -use anyhow::{Context as _, Result, anyhow}; -use base64::Engine as _; -use base64::engine::general_purpose::URL_SAFE_NO_PAD; +use anyhow::{Result, anyhow}; use credentials_provider::CredentialsProvider; -use futures::{FutureExt, StreamExt, future::BoxFuture, future::Shared}; -use gpui::{App, AsyncApp, Context, Entity, SharedString, Task, Window}; -use http_client::{ - AsyncBody, CustomHeaders, HttpClient, Method, Request as HttpRequest, - http::{HeaderName, HeaderValue}, -}; +use futures::future::Shared; +use gpui::{App, Context, Entity, SharedString, Task, Window}; +use http_client::HttpClient; use language_model::{ AuthenticateError, FastModeConfirmation, IconOrSvg, InlineDescription, LanguageModel, - LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelEffortLevel, - LanguageModelId, LanguageModelName, LanguageModelProvider, LanguageModelProviderId, - LanguageModelProviderName, LanguageModelProviderState, LanguageModelRequest, - LanguageModelToolChoice, ProviderSettingsView, RateLimiter, + LanguageModelProvider, LanguageModelProviderId, LanguageModelProviderName, + LanguageModelProviderState, ProviderSettingsView, }; -use open_ai::{ReasoningEffort, responses::stream_response}; -use rand::RngCore as _; -use serde::{Deserialize, Serialize}; -use sha2::{Digest, Sha256}; +use openai_subscribed::{ChatGptModel, PROVIDER_ID, PROVIDER_NAME, State, create_language_model}; use std::sync::Arc; -use std::time::{SystemTime, UNIX_EPOCH}; use ui::{ConfiguredApiCard, prelude::*}; -use url::form_urlencoded; -use util::ResultExt as _; - -use crate::provider::open_ai::{OpenAiResponseEventMapper, into_open_ai_response}; - -const PROVIDER_ID: LanguageModelProviderId = LanguageModelProviderId::new("openai-subscribed"); -const PROVIDER_NAME: LanguageModelProviderName = - LanguageModelProviderName::new("ChatGPT Subscription"); const SUBSCRIPTION_DESCRIPTION: &str = "Sign in with your ChatGPT Plus or Pro subscription to use OpenAI models in Zed's agent."; -const CODEX_BASE_URL: &str = "https://chatgpt.com/backend-api/codex"; -const OPENAI_TOKEN_URL: &str = "https://auth.openai.com/oauth/token"; -const OPENAI_AUTHORIZE_URL: &str = "https://auth.openai.com/oauth/authorize"; -const CLIENT_ID: &str = "app_EMoamEEZ73f0CkXaXp7hrann"; - -const CREDENTIALS_KEY: &str = "https://chatgpt.com/backend-api/codex"; -const TOKEN_REFRESH_BUFFER_MS: u64 = 5 * 60 * 1000; - -#[derive(Serialize, Deserialize, Clone, Debug)] -struct CodexCredentials { - access_token: String, - refresh_token: String, - expires_at_ms: u64, - account_id: Option, - email: Option, -} - -impl CodexCredentials { - fn is_expired(&self) -> bool { - let now = now_ms(); - now + TOKEN_REFRESH_BUFFER_MS >= self.expires_at_ms - } -} - -pub struct State { - credentials: Option, - sign_in_task: Option>>, - refresh_task: Option>>>>, - load_task: Option>>>>, - credentials_provider: Arc, - auth_generation: u64, - last_auth_error: Option, -} - -#[derive(Debug)] -enum RefreshError { - Fatal(anyhow::Error), - Transient(anyhow::Error), -} - -impl std::fmt::Display for RefreshError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - RefreshError::Fatal(e) => write!(f, "{e}"), - RefreshError::Transient(e) => write!(f, "{e}"), - } - } -} - -impl State { - fn is_authenticated(&self) -> bool { - self.credentials.is_some() - } - - fn email(&self) -> Option<&str> { - self.credentials.as_ref().and_then(|c| c.email.as_deref()) - } - - fn is_signing_in(&self) -> bool { - self.sign_in_task.is_some() - } -} - pub struct OpenAiSubscribedProvider { - http_client: Arc, state: Entity, } @@ -108,63 +25,8 @@ impl OpenAiSubscribedProvider { credentials_provider: Arc, cx: &mut App, ) -> Self { - let state = cx.new(|_cx| State { - credentials: None, - sign_in_task: None, - refresh_task: None, - load_task: None, - credentials_provider, - auth_generation: 0, - last_auth_error: None, - }); - - let provider = Self { http_client, state }; - - provider.load_credentials(cx); - - provider - } - - fn load_credentials(&self, cx: &mut App) { - let state = self.state.downgrade(); - let load_task = cx - .spawn(async move |cx| { - let credentials_provider = - state.read_with(&*cx, |s, _| s.credentials_provider.clone())?; - let result = credentials_provider - .read_credentials(CREDENTIALS_KEY, &*cx) - .await; - state.update(cx, |s, cx| { - if let Ok(Some((_, bytes))) = result { - match serde_json::from_slice::(&bytes) { - Ok(creds) => s.credentials = Some(creds), - Err(err) => { - log::warn!( - "Failed to deserialize ChatGPT subscription credentials: {err}" - ); - } - } - } - s.load_task = None; - cx.notify(); - })?; - Ok::<(), Arc>(()) - }) - .shared(); - - self.state.update(cx, |s, _| { - s.load_task = Some(load_task); - }); - } - - fn create_language_model(&self, model: ChatGptModel) -> Arc { - Arc::new(OpenAiSubscribedLanguageModel { - id: LanguageModelId::from(model.id().to_string()), - model, - state: self.state.clone(), - http_client: self.http_client.clone(), - request_limiter: RateLimiter::new(4), - }) + let state = cx.new(|cx| State::new(http_client, credentials_provider, cx)); + Self { state } } } @@ -189,18 +51,22 @@ impl LanguageModelProvider for OpenAiSubscribedProvider { IconOrSvg::Icon(IconName::AiOpenAi) } - fn default_model(&self, _cx: &App) -> Option> { - Some(self.create_language_model(ChatGptModel::Gpt55)) + fn default_model(&self, cx: &App) -> Option> { + Some(create_language_model(ChatGptModel::Gpt55, &self.state, cx)) } - fn default_fast_model(&self, _cx: &App) -> Option> { - Some(self.create_language_model(ChatGptModel::Gpt56Luna)) + fn default_fast_model(&self, cx: &App) -> Option> { + Some(create_language_model( + ChatGptModel::Gpt56Luna, + &self.state, + cx, + )) } - fn provided_models(&self, _cx: &App) -> Vec> { + fn provided_models(&self, cx: &App) -> Vec> { ChatGptModel::all() .into_iter() - .map(|m| self.create_language_model(m)) + .map(|model| create_language_model(model, &self.state, cx)) .collect() } @@ -212,13 +78,13 @@ impl LanguageModelProvider for OpenAiSubscribedProvider { if self.is_authenticated(cx) { return Task::ready(Ok(())); } - let load_task = self.state.read(cx).load_task.clone(); + let load_task: Option> = self.state.read(cx).load_task(); if let Some(load_task) = load_task { let weak_state = self.state.downgrade(); cx.spawn(async move |cx| { let _ = load_task.await; let is_auth = weak_state - .read_with(&*cx, |s, _| s.is_authenticated()) + .read_with(&*cx, |state, _| state.is_authenticated()) .unwrap_or(false); if is_auth { Ok(()) @@ -256,11 +122,9 @@ impl LanguageModelProvider for OpenAiSubscribedProvider { description, create_view: Arc::new({ let state = self.state.clone(); - let http_client = self.http_client.clone(); move |_window, cx| { cx.new(|_cx| ConfigurationView { state: state.clone(), - http_client: http_client.clone(), compact: true, }) .into() @@ -293,791 +157,8 @@ impl LanguageModelProvider for OpenAiSubscribedProvider { } } -// -// The ChatGPT Subscription provider routes requests to chatgpt.com/backend-api/codex, -// which only supports a subset of OpenAI models. This list is maintained separately -// from the standard OpenAI API model list (open_ai::Model). -// -// TODO: The Codex CLI fetches this list dynamically from -// `GET /models?client_version=...` (see -// codex-rs/codex-api/src/endpoint/models.rs in openai/codex) and falls back to -// a bundled models.json. Beyond going stale, the static approach also can't -// model per-account access (e.g. free accounts cannot use gpt-5.4 even though -// paid accounts can), so the backend still rejects some requests. The bundled -// list at -// codex-rs/models-manager/models.json (openai/codex) is the closest -// approximation; the entries below mirror that file's picker-visible models. -#[derive(Clone, Debug, PartialEq)] -enum ChatGptModel { - Gpt56Sol, - Gpt56Terra, - Gpt56Luna, - Gpt55, - Gpt54, - Gpt54Mini, -} - -impl ChatGptModel { - fn all() -> Vec { - vec![ - Self::Gpt56Sol, - Self::Gpt56Terra, - Self::Gpt56Luna, - Self::Gpt55, - Self::Gpt54, - Self::Gpt54Mini, - ] - } - - fn id(&self) -> &str { - match self { - Self::Gpt56Sol => "gpt-5.6-sol", - Self::Gpt56Terra => "gpt-5.6-terra", - Self::Gpt56Luna => "gpt-5.6-luna", - Self::Gpt55 => "gpt-5.5", - Self::Gpt54 => "gpt-5.4", - Self::Gpt54Mini => "gpt-5.4-mini", - } - } - - fn display_name(&self) -> &str { - match self { - Self::Gpt56Sol => "GPT-5.6 Sol", - Self::Gpt56Terra => "GPT-5.6 Terra", - Self::Gpt56Luna => "GPT-5.6 Luna", - Self::Gpt55 => "GPT-5.5", - Self::Gpt54 => "GPT-5.4", - Self::Gpt54Mini => "GPT-5.4 Mini", - } - } - - fn max_token_count(&self) -> u64 { - match self { - Self::Gpt56Sol | Self::Gpt56Terra | Self::Gpt56Luna => 372_000, - Self::Gpt55 | Self::Gpt54 | Self::Gpt54Mini => 272_000, - } - } - - fn max_output_tokens(&self) -> Option { - // Codex model metadata does not expose a max output token cap for these - // models. Source: openai/codex models-manager/models.json. - None - } - - fn supports_images(&self) -> bool { - true - } - - fn default_reasoning_effort(&self) -> Option { - match self { - Self::Gpt56Sol => Some(ReasoningEffort::Low), - Self::Gpt56Terra | Self::Gpt56Luna | Self::Gpt55 | Self::Gpt54 | Self::Gpt54Mini => { - Some(ReasoningEffort::Medium) - } - } - } - - fn supported_reasoning_efforts(&self) -> &'static [ReasoningEffort] { - match self { - Self::Gpt56Sol | Self::Gpt56Terra | Self::Gpt56Luna => &[ - ReasoningEffort::Low, - ReasoningEffort::Medium, - ReasoningEffort::High, - ReasoningEffort::XHigh, - ReasoningEffort::Max, - ], - Self::Gpt55 | Self::Gpt54 | Self::Gpt54Mini => &[ - ReasoningEffort::Low, - ReasoningEffort::Medium, - ReasoningEffort::High, - ReasoningEffort::XHigh, - ], - } - } - - fn supports_parallel_tool_calls(&self) -> bool { - true - } - - fn supports_prompt_cache_key(&self) -> bool { - true - } - - fn supports_priority(&self) -> bool { - match self { - Self::Gpt56Sol | Self::Gpt56Terra | Self::Gpt56Luna | Self::Gpt55 | Self::Gpt54 => true, - Self::Gpt54Mini => false, - } - } -} - -struct OpenAiSubscribedLanguageModel { - id: LanguageModelId, - model: ChatGptModel, - state: Entity, - http_client: Arc, - request_limiter: RateLimiter, -} - -impl LanguageModel for OpenAiSubscribedLanguageModel { - fn id(&self) -> LanguageModelId { - self.id.clone() - } - - fn name(&self) -> LanguageModelName { - LanguageModelName::from(self.model.display_name().to_string()) - } - - fn provider_id(&self) -> LanguageModelProviderId { - PROVIDER_ID - } - - fn provider_name(&self) -> LanguageModelProviderName { - PROVIDER_NAME - } - - fn supports_tools(&self) -> bool { - true - } - - fn supports_images(&self) -> bool { - self.model.supports_images() - } - - fn supports_tool_choice(&self, _choice: LanguageModelToolChoice) -> bool { - true - } - - fn supports_streaming_tools(&self) -> bool { - true - } - - fn supports_thinking(&self) -> bool { - true - } - - fn supports_fast_mode(&self) -> bool { - self.model.supports_priority() - } - - fn supported_effort_levels(&self) -> Vec { - let default_effort = self.model.default_reasoning_effort(); - self.model - .supported_reasoning_efforts() - .iter() - .copied() - .filter_map(|effort| { - let (name, value) = match effort { - ReasoningEffort::None => return None, - ReasoningEffort::Minimal => ("Minimal", "minimal"), - ReasoningEffort::Low => ("Low", "low"), - ReasoningEffort::Medium => ("Medium", "medium"), - ReasoningEffort::High => ("High", "high"), - ReasoningEffort::XHigh => ("Extra High", "xhigh"), - ReasoningEffort::Max => ("Max", "max"), - }; - - Some(LanguageModelEffortLevel { - name: name.into(), - value: value.into(), - is_default: Some(effort) == default_effort, - }) - }) - .collect() - } - - fn telemetry_id(&self) -> String { - format!("openai-subscribed/{}", self.model.id()) - } - - fn max_token_count(&self) -> u64 { - self.model.max_token_count() - } - - fn max_output_tokens(&self) -> Option { - self.model.max_output_tokens() - } - - fn stream_completion( - &self, - mut request: LanguageModelRequest, - cx: &AsyncApp, - ) -> BoxFuture< - 'static, - Result< - futures::stream::BoxStream< - 'static, - Result, - >, - LanguageModelCompletionError, - >, - > { - if !self.model.supports_priority() { - request.speed = None; - } - - // The Codex backend rejects `max_output_tokens` (`Unsupported parameter`), - // unlike the public OpenAI Responses API. Pass `None` so the field is - // omitted from the serialized request body entirely. - let mut responses_request = match into_open_ai_response( - request, - self.model.id(), - self.model.supports_parallel_tool_calls(), - self.model.supports_prompt_cache_key(), - /*max_output_tokens*/ None, - self.model.default_reasoning_effort(), - self.model - .supported_reasoning_efforts() - .contains(&ReasoningEffort::None), - &PROVIDER_ID, - ) { - Ok(request) => request, - Err(error) => return async move { Err(error.into()) }.boxed(), - }; - responses_request.store = Some(false); - - // `into_open_ai_response` already hoists system messages into - // `instructions`, which is the only form the Codex backend accepts. - // Codex has only ever been sent requests with the field present - // (possibly empty), so keep sending it even without system messages. - responses_request.instructions.get_or_insert_default(); - - let state = self.state.downgrade(); - let http_client = self.http_client.clone(); - let request_limiter = self.request_limiter.clone(); - - let future = cx.spawn(async move |cx| { - let creds = get_fresh_credentials(&state, &http_client, cx).await?; - - let mut header_pairs: Vec<(HeaderName, HeaderValue)> = vec![ - ( - HeaderName::from_static("originator"), - HeaderValue::from_static("zed"), - ), - ( - HeaderName::from_static("openai-beta"), - HeaderValue::from_static("responses=experimental"), - ), - ]; - if let Some(ref id) = creds.account_id { - if !id.is_empty() { - if let Ok(value) = HeaderValue::from_str(id) { - header_pairs.push((HeaderName::from_static("chatgpt-account-id"), value)); - } - } - } - let extra_headers = CustomHeaders::new(header_pairs); - - let access_token = creds.access_token.clone(); - request_limiter - .stream(async move { - stream_response( - http_client.as_ref(), - PROVIDER_NAME.0.as_str(), - CODEX_BASE_URL, - &access_token, - responses_request, - &extra_headers, - ) - .await - .map_err(LanguageModelCompletionError::from) - }) - .await - }); - - async move { - let mapper = OpenAiResponseEventMapper::new(PROVIDER_ID); - Ok(mapper.map_stream(future.await?.boxed()).boxed()) - } - .boxed() - } -} - -async fn get_fresh_credentials( - state: &gpui::WeakEntity, - http_client: &Arc, - cx: &mut AsyncApp, -) -> Result { - let (creds, existing_task) = state - .read_with(&*cx, |s, _| (s.credentials.clone(), s.refresh_task.clone())) - .map_err(LanguageModelCompletionError::Other)?; - - let creds = creds.ok_or(LanguageModelCompletionError::NoApiKey { - provider: PROVIDER_NAME, - })?; - - if !creds.is_expired() { - return Ok(creds); - } - - // If another caller is already refreshing, await their result. - if let Some(shared_task) = existing_task { - return shared_task - .await - .map_err(|e| LanguageModelCompletionError::Other(anyhow::anyhow!("{e}"))); - } - - // We are the first caller to notice expiry — spawn the refresh task. - let http_client_clone = http_client.clone(); - let state_clone = state.clone(); - let refresh_token_value = creds.refresh_token.clone(); - - // Capture the generation so we can detect sign-outs that happened during refresh. - let generation = state - .read_with(&*cx, |s, _| s.auth_generation) - .map_err(LanguageModelCompletionError::Other)?; - - let shared_task = cx - .spawn(async move |cx| { - let result = refresh_token(&http_client_clone, &refresh_token_value).await; - - match result { - Ok(refreshed) => { - let persist_result: Result> = async { - // Check if auth_generation changed (sign-out during refresh). - let current_generation = state_clone - .read_with(&*cx, |s, _| s.auth_generation) - .map_err(|e| Arc::new(e))?; - if current_generation != generation { - return Err(Arc::new(anyhow!( - "Sign-out occurred during token refresh" - ))); - } - - let credentials_provider = state_clone - .read_with(&*cx, |s, _| s.credentials_provider.clone()) - .map_err(|e| Arc::new(e))?; - - let json = - serde_json::to_vec(&refreshed).map_err(|e| Arc::new(e.into()))?; - - credentials_provider - .write_credentials(CREDENTIALS_KEY, "Bearer", &json, &*cx) - .await - .map_err(|e| Arc::new(e))?; - - state_clone - .update(cx, |s, _| { - s.credentials = Some(refreshed.clone()); - s.refresh_task = None; - }) - .map_err(|e| Arc::new(e))?; - - Ok(refreshed) - } - .await; - - // Clear refresh_task on failure too. - if persist_result.is_err() { - let _ = state_clone.update(cx, |s, _| { - s.refresh_task = None; - }); - } - - persist_result - } - Err(RefreshError::Fatal(e)) => { - log::error!("ChatGPT subscription token refresh failed fatally: {e:?}"); - let _ = state_clone.update(cx, |s, cx| { - s.refresh_task = None; - s.credentials = None; - s.last_auth_error = - Some("Your session has expired. Please sign in again.".into()); - cx.notify(); - }); - // Also clear the keychain so stale credentials aren't loaded next time. - if let Ok(credentials_provider) = - state_clone.read_with(&*cx, |s, _| s.credentials_provider.clone()) - { - credentials_provider - .delete_credentials(CREDENTIALS_KEY, &*cx) - .await - .log_err(); - } - Err(Arc::new(e)) - } - Err(RefreshError::Transient(e)) => { - log::warn!("ChatGPT subscription token refresh failed transiently: {e:?}"); - let _ = state_clone.update(cx, |s, _| { - s.refresh_task = None; - }); - Err(Arc::new(e)) - } - } - }) - .shared(); - - // Store the shared task so concurrent callers can join on it. - state - .update(cx, |s, _| { - s.refresh_task = Some(shared_task.clone()); - }) - .map_err(LanguageModelCompletionError::Other)?; - - shared_task - .await - .map_err(|e| LanguageModelCompletionError::Other(anyhow::anyhow!("{e}"))) -} - -#[derive(Deserialize)] -struct TokenResponse { - access_token: String, - refresh_token: String, - #[serde(default)] - id_token: Option, - expires_in: u64, - #[serde(default)] - email: Option, -} - -// The OAuth client registered for `CLIENT_ID` (the Codex CLI's client) only allows -// `http://localhost:1455/auth/callback` and `http://localhost:1457/auth/callback` -// as redirect URIs; using anything else (different host, port, or path) causes -// auth.openai.com to reject the authorize request with a generic `unknown_error` -// before redirecting back. Keep these in sync with the Codex CLI's redirect URI -// allow-list (see codex-rs/login/src/server.rs in openai/codex). -const CODEX_CALLBACK_HOST: &str = "localhost"; -const CODEX_CALLBACK_PORT: u16 = 1455; -const CODEX_CALLBACK_FALLBACK_PORT: u16 = 1457; -const CODEX_CALLBACK_PATH: &str = "/auth/callback"; - -async fn do_oauth_flow( - http_client: Arc, - cx: &AsyncApp, -) -> Result { - // Start the callback server FIRST so the redirect URI is ready - let (redirect_uri, callback_rx) = - oauth_callback_server::start_oauth_callback_server_with_config( - oauth_callback_server::OAuthCallbackServerConfig { - host: CODEX_CALLBACK_HOST, - preferred_port: CODEX_CALLBACK_PORT, - fallback_port: Some(CODEX_CALLBACK_FALLBACK_PORT), - path: CODEX_CALLBACK_PATH, - }, - ) - .context("Failed to start OAuth callback server")?; - - // PKCE verifier: 32 random bytes → base64url (no padding) - let mut verifier_bytes = [0u8; 32]; - rand::rng().fill_bytes(&mut verifier_bytes); - let verifier = URL_SAFE_NO_PAD.encode(verifier_bytes); - - // PKCE challenge: SHA-256(verifier) → base64url - let mut hasher = Sha256::new(); - hasher.update(verifier.as_bytes()); - let challenge = URL_SAFE_NO_PAD.encode(hasher.finalize().as_slice()); - - // CSRF state: 16 random bytes → hex string - let mut state_bytes = [0u8; 16]; - rand::rng().fill_bytes(&mut state_bytes); - let oauth_state: String = state_bytes.iter().map(|b| format!("{b:02x}")).collect(); - - let mut auth_url = url::Url::parse(OPENAI_AUTHORIZE_URL).expect("valid base URL"); - auth_url - .query_pairs_mut() - .append_pair("client_id", CLIENT_ID) - .append_pair("redirect_uri", &redirect_uri) - // Deliberately excludes `api.connectors.read api.connectors.invoke` - // (which Codex CLI requests): extra scopes inflate the - // access-token JWT, and the serialized credentials must fit within - // Windows Credential Manager's 2560-byte blob limit - // (CRED_MAX_CREDENTIAL_BLOB_SIZE). See #58541. - .append_pair("scope", "openid profile email offline_access") - .append_pair("response_type", "code") - .append_pair("code_challenge", &challenge) - .append_pair("code_challenge_method", "S256") - .append_pair("id_token_add_organizations", "true") - .append_pair("state", &oauth_state) - .append_pair("codex_cli_simplified_flow", "true") - .append_pair("originator", "zed"); - - // Open browser AFTER the listener is ready - cx.update(|cx| cx.open_url(auth_url.as_str())); - - // Await the callback - let callback = callback_rx - .await - .map_err(|_| anyhow!("OAuth callback was cancelled"))? - .context("OAuth callback failed")?; - - // Validate CSRF state - if callback.state != oauth_state { - return Err(anyhow!("OAuth state mismatch")); - } - - let tokens = exchange_code(&http_client, &callback.code, &verifier, &redirect_uri) - .await - .context("Token exchange failed")?; - - let jwt = tokens - .id_token - .as_deref() - .unwrap_or(tokens.access_token.as_str()); - let claims = extract_jwt_claims(jwt); - - Ok(CodexCredentials { - access_token: tokens.access_token, - refresh_token: tokens.refresh_token, - expires_at_ms: now_ms() + tokens.expires_in * 1000, - account_id: claims.account_id, - email: claims.email.or(tokens.email), - }) -} - -async fn exchange_code( - client: &Arc, - code: &str, - verifier: &str, - redirect_uri: &str, -) -> Result { - let body = form_urlencoded::Serializer::new(String::new()) - .append_pair("grant_type", "authorization_code") - .append_pair("client_id", CLIENT_ID) - .append_pair("code", code) - .append_pair("redirect_uri", redirect_uri) - .append_pair("code_verifier", verifier) - .finish(); - - let request = HttpRequest::builder() - .method(Method::POST) - .uri(OPENAI_TOKEN_URL) - .header("Content-Type", "application/x-www-form-urlencoded") - .body(AsyncBody::from(body))?; - - let mut response = client.send(request).await?; - let mut body = String::new(); - smol::io::AsyncReadExt::read_to_string(response.body_mut(), &mut body).await?; - - if !response.status().is_success() { - return Err(anyhow!( - "Token exchange failed (HTTP {}): {body}", - response.status() - )); - } - - serde_json::from_str::(&body).context("Failed to parse token response") -} - -async fn refresh_token( - client: &Arc, - refresh_token: &str, -) -> Result { - let body = form_urlencoded::Serializer::new(String::new()) - .append_pair("grant_type", "refresh_token") - .append_pair("client_id", CLIENT_ID) - .append_pair("refresh_token", refresh_token) - .finish(); - - let request = HttpRequest::builder() - .method(Method::POST) - .uri(OPENAI_TOKEN_URL) - .header("Content-Type", "application/x-www-form-urlencoded") - .body(AsyncBody::from(body)) - .map_err(|e| RefreshError::Transient(e.into()))?; - - let mut response = client - .send(request) - .await - .map_err(|e| RefreshError::Transient(e))?; - let status = response.status(); - let mut body = String::new(); - smol::io::AsyncReadExt::read_to_string(response.body_mut(), &mut body) - .await - .map_err(|e| RefreshError::Transient(e.into()))?; - - if !status.is_success() { - let err = anyhow!("Token refresh failed (HTTP {}): {body}", status); - // 400/401/403 indicate a revoked or invalid refresh token. - // 5xx and other errors are treated as transient. - if status == http_client::StatusCode::BAD_REQUEST - || status == http_client::StatusCode::UNAUTHORIZED - || status == http_client::StatusCode::FORBIDDEN - { - return Err(RefreshError::Fatal(err)); - } - return Err(RefreshError::Transient(err)); - } - - let tokens: TokenResponse = - serde_json::from_str(&body).map_err(|e| RefreshError::Transient(e.into()))?; - let jwt = tokens - .id_token - .as_deref() - .unwrap_or(tokens.access_token.as_str()); - let claims = extract_jwt_claims(jwt); - - Ok(CodexCredentials { - access_token: tokens.access_token, - refresh_token: tokens.refresh_token, - expires_at_ms: now_ms() + tokens.expires_in * 1000, - account_id: claims.account_id, - email: claims.email.or(tokens.email), - }) -} - -struct JwtClaims { - account_id: Option, - email: Option, -} - -/// Extract claims from a JWT payload (base64url middle segment). -/// Extracts `chatgpt_account_id` from three possible locations (matching Roo Code's -/// implementation) and the `email` claim. -fn extract_jwt_claims(jwt: &str) -> JwtClaims { - let Some(payload_b64) = jwt.split('.').nth(1) else { - return JwtClaims { - account_id: None, - email: None, - }; - }; - let Ok(payload) = URL_SAFE_NO_PAD.decode(payload_b64) else { - return JwtClaims { - account_id: None, - email: None, - }; - }; - let Ok(claims) = serde_json::from_slice::(&payload) else { - return JwtClaims { - account_id: None, - email: None, - }; - }; - - let account_id = claims - .get("chatgpt_account_id") - .and_then(|v| v.as_str()) - .or_else(|| { - claims - .get("https://api.openai.com/auth") - .and_then(|v| v.get("chatgpt_account_id")) - .and_then(|v| v.as_str()) - }) - .or_else(|| { - claims - .get("organizations") - .and_then(|v| v.as_array()) - .and_then(|arr| arr.first()) - .and_then(|org| org.get("id")) - .and_then(|v| v.as_str()) - }) - .map(|s| s.to_owned()); - - let email = claims - .get("email") - .and_then(|v| v.as_str()) - .map(|s| s.to_owned()); - - JwtClaims { account_id, email } -} - -fn now_ms() -> u64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|d| d.as_millis() as u64) - .unwrap_or_else(|err| { - log::error!("System clock is before UNIX epoch: {err}"); - 0 - }) -} - -fn do_sign_in(state: &Entity, http_client: &Arc, cx: &mut App) { - if state.read(cx).is_signing_in() { - return; - } - - let weak_state = state.downgrade(); - let http_client = http_client.clone(); - - let task = cx.spawn(async move |cx| { - match do_oauth_flow(http_client, &*cx).await { - Ok(creds) => { - let persist_result = async { - let credentials_provider = - weak_state.read_with(&*cx, |s, _| s.credentials_provider.clone())?; - let json = serde_json::to_vec(&creds)?; - credentials_provider - .write_credentials(CREDENTIALS_KEY, "Bearer", &json, &*cx) - .await?; - anyhow::Ok(()) - } - .await; - - match persist_result { - Ok(()) => { - weak_state - .update(cx, |s, cx| { - s.credentials = Some(creds); - s.sign_in_task = None; - s.last_auth_error = None; - cx.notify(); - }) - .log_err(); - } - Err(err) => { - log::error!( - "ChatGPT subscription sign-in failed to persist credentials: {err:?}" - ); - weak_state - .update(cx, |s, cx| { - s.sign_in_task = None; - s.last_auth_error = - Some("Failed to save credentials. Please try again.".into()); - cx.notify(); - }) - .log_err(); - } - } - } - Err(err) => { - log::error!("ChatGPT subscription sign-in failed: {err:?}"); - weak_state - .update(cx, |s, cx| { - s.sign_in_task = None; - s.last_auth_error = Some("Sign-in failed. Please try again.".into()); - cx.notify(); - }) - .log_err(); - } - } - anyhow::Ok(()) - }); - - state.update(cx, |s, cx| { - s.last_auth_error = None; - s.sign_in_task = Some(task); - cx.notify(); - }); -} - -fn do_sign_out(state: &gpui::WeakEntity, cx: &mut App) -> Task> { - let weak_state = state.clone(); - // Clear credentials and cancel in-flight work immediately so the UI - // reflects the sign-out right away. - weak_state - .update(cx, |s, cx| { - s.auth_generation += 1; - s.credentials = None; - s.sign_in_task = None; - s.refresh_task = None; - s.last_auth_error = None; - cx.notify(); - }) - .log_err(); - - cx.spawn(async move |cx| { - let credentials_provider = - weak_state.read_with(&*cx, |s, _| s.credentials_provider.clone())?; - credentials_provider - .delete_credentials(CREDENTIALS_KEY, &*cx) - .await - .context("Failed to delete ChatGPT subscription credentials from keychain")?; - anyhow::Ok(()) - }) -} - struct ConfigurationView { state: Entity, - http_client: Arc, /// When `true`, the description is rendered elsewhere (the settings row's /// left column), so it's omitted here to avoid duplication. compact: bool, @@ -1093,22 +174,23 @@ impl Render for ConfigurationView { .map(|e| format!("Signed in as {e}")) .unwrap_or_else(|| "Signed in".to_string()); - let weak_state = self.state.downgrade(); + let state_entity = self.state.clone(); return v_flex() .child( ConfiguredApiCard::new("openai-subscribed-sign-out", SharedString::from(label)) .button_label("Sign Out") .on_click(cx.listener(move |_this, _, _window, cx| { - do_sign_out(&weak_state, cx).detach_and_log_err(cx); + state_entity + .update(cx, |state, cx| state.sign_out(cx)) + .detach_and_log_err(cx); })), ) .into_any_element(); } - let last_auth_error = state.last_auth_error.clone(); + let last_auth_error = state.last_auth_error(); let provider_state = self.state.clone(); - let http_client = self.http_client.clone(); let is_signing_in = state.is_signing_in(); let button_label = if is_signing_in { @@ -1130,7 +212,7 @@ impl Render for ConfigurationView { .loading(is_signing_in) .disabled(is_signing_in) .on_click(move |_, _window, cx| { - do_sign_in(&provider_state, &http_client, cx); + provider_state.update(cx, |state, cx| state.sign_in(cx)); }), ) .when_some(last_auth_error, |this, error| { @@ -1153,12 +235,48 @@ impl Render for ConfigurationView { #[cfg(test)] mod tests { use super::*; - use gpui::TestAppContext; + use gpui::{AsyncApp, TestAppContext}; use http_client::FakeHttpClient; use parking_lot::Mutex; use std::future::Future; use std::pin::Pin; - use std::sync::atomic::{AtomicUsize, Ordering}; + + #[gpui::test] + async fn test_authenticate_awaits_initial_load(cx: &mut TestAppContext) { + let creds_json = serde_json::json!({ + "access_token": "fresh_access", + "refresh_token": "fresh_refresh", + "expires_at_ms": u64::MAX, + "account_id": null, + "email": null, + }); + let creds_provider = Arc::new(FakeCredentialsProvider::new()); + creds_provider.storage.lock().replace(( + "Bearer".to_string(), + serde_json::to_vec(&creds_json).unwrap(), + )); + + let http_client = FakeHttpClient::create(|_| async { + Ok(http_client::Response::builder() + .status(200) + .body(http_client::AsyncBody::default())?) + }); + + let provider = + cx.update(|cx| OpenAiSubscribedProvider::new(http_client, creds_provider, cx)); + + // Before load completes, authenticate should still await the load. + let auth_task = cx.update(|cx| provider.authenticate(cx)); + + // Drive the load to completion. + cx.run_until_parked(); + + let result = auth_task.await; + assert!( + result.is_ok(), + "authenticate should succeed after load completes with valid credentials" + ); + } struct FakeCredentialsProvider { storage: Mutex)>>, @@ -1203,387 +321,4 @@ mod tests { Box::pin(async { Ok(()) }) } } - - fn make_expired_credentials() -> CodexCredentials { - CodexCredentials { - access_token: "old_access".to_string(), - refresh_token: "old_refresh".to_string(), - expires_at_ms: 0, - account_id: None, - email: None, - } - } - - fn make_fresh_credentials() -> CodexCredentials { - CodexCredentials { - access_token: "fresh_access".to_string(), - refresh_token: "fresh_refresh".to_string(), - expires_at_ms: now_ms() + 3_600_000, - account_id: None, - email: None, - } - } - - fn fake_token_response() -> String { - serde_json::json!({ - "access_token": "fresh_access", - "refresh_token": "fresh_refresh", - "expires_in": 3600 - }) - .to_string() - } - - #[gpui::test] - async fn test_concurrent_refresh_deduplicates(cx: &mut TestAppContext) { - let refresh_count = Arc::new(AtomicUsize::new(0)); - let refresh_count_clone = refresh_count.clone(); - - let http_client = FakeHttpClient::create(move |_request| { - let refresh_count = refresh_count_clone.clone(); - async move { - refresh_count.fetch_add(1, Ordering::SeqCst); - let body = fake_token_response(); - Ok(http_client::Response::builder() - .status(200) - .body(http_client::AsyncBody::from(body))?) - } - }); - - let state = cx.new(|_cx| State { - credentials: Some(make_expired_credentials()), - sign_in_task: None, - refresh_task: None, - load_task: None, - credentials_provider: Arc::new(FakeCredentialsProvider::new()), - auth_generation: 0, - last_auth_error: None, - }); - - let weak_state = cx.read(|_cx| state.downgrade()); - let http: Arc = http_client; - - // Spawn two concurrent refresh attempts. - let weak1 = weak_state.clone(); - let http1 = http.clone(); - let task1 = - cx.spawn(async move |mut cx| get_fresh_credentials(&weak1, &http1, &mut cx).await); - - let weak2 = weak_state.clone(); - let http2 = http.clone(); - let task2 = - cx.spawn(async move |mut cx| get_fresh_credentials(&weak2, &http2, &mut cx).await); - - // Drive both to completion. - cx.run_until_parked(); - let result1 = task1.await; - let result2 = task2.await; - - assert!(result1.is_ok(), "first refresh should succeed"); - assert!(result2.is_ok(), "second refresh should succeed"); - assert_eq!(result1.unwrap().access_token, "fresh_access"); - assert_eq!(result2.unwrap().access_token, "fresh_access"); - assert_eq!( - refresh_count.load(Ordering::SeqCst), - 1, - "refresh_token should only be called once despite two concurrent callers" - ); - } - - #[gpui::test] - async fn test_fresh_credentials_skip_refresh(cx: &mut TestAppContext) { - let refresh_count = Arc::new(AtomicUsize::new(0)); - let refresh_count_clone = refresh_count.clone(); - - let http_client = FakeHttpClient::create(move |_request| { - let refresh_count = refresh_count_clone.clone(); - async move { - refresh_count.fetch_add(1, Ordering::SeqCst); - let body = fake_token_response(); - Ok(http_client::Response::builder() - .status(200) - .body(http_client::AsyncBody::from(body))?) - } - }); - - let state = cx.new(|_cx| State { - credentials: Some(make_fresh_credentials()), - sign_in_task: None, - refresh_task: None, - load_task: None, - credentials_provider: Arc::new(FakeCredentialsProvider::new()), - auth_generation: 0, - last_auth_error: None, - }); - - let weak_state = cx.read(|_cx| state.downgrade()); - let http: Arc = http_client; - - let weak = weak_state.clone(); - let http_clone = http.clone(); - let result = cx - .spawn(async move |mut cx| get_fresh_credentials(&weak, &http_clone, &mut cx).await) - .await; - - assert!(result.is_ok()); - assert_eq!(result.unwrap().access_token, "fresh_access"); - assert_eq!( - refresh_count.load(Ordering::SeqCst), - 0, - "no refresh should happen when credentials are fresh" - ); - } - - #[gpui::test] - async fn test_no_credentials_returns_no_api_key(cx: &mut TestAppContext) { - let http_client = FakeHttpClient::create(|_| async { - Ok(http_client::Response::builder() - .status(200) - .body(http_client::AsyncBody::default())?) - }); - - let state = cx.new(|_cx| State { - credentials: None, - sign_in_task: None, - refresh_task: None, - load_task: None, - credentials_provider: Arc::new(FakeCredentialsProvider::new()), - auth_generation: 0, - last_auth_error: None, - }); - - let weak_state = cx.read(|_cx| state.downgrade()); - let http: Arc = http_client; - - let weak = weak_state.clone(); - let http_clone = http.clone(); - let result = cx - .spawn(async move |mut cx| get_fresh_credentials(&weak, &http_clone, &mut cx).await) - .await; - - assert!(matches!( - result, - Err(LanguageModelCompletionError::NoApiKey { .. }) - )); - } - - #[gpui::test] - async fn test_fatal_refresh_clears_auth_state(cx: &mut TestAppContext) { - let http_client = FakeHttpClient::create(move |_request| async move { - Ok(http_client::Response::builder() - .status(401) - .body(http_client::AsyncBody::from(r#"{"error":"invalid_grant"}"#))?) - }); - - let creds_provider = Arc::new(FakeCredentialsProvider::new()); - let state = cx.new(|_cx| State { - credentials: Some(make_expired_credentials()), - sign_in_task: None, - refresh_task: None, - load_task: None, - credentials_provider: creds_provider.clone(), - auth_generation: 0, - last_auth_error: None, - }); - - let weak_state = cx.read(|_cx| state.downgrade()); - let http: Arc = http_client; - - let weak = weak_state.clone(); - let http_clone = http.clone(); - let result = cx - .spawn(async move |mut cx| get_fresh_credentials(&weak, &http_clone, &mut cx).await) - .await; - - cx.run_until_parked(); - - assert!(result.is_err(), "fatal refresh should return an error"); - cx.read(|cx| { - let s = state.read(cx); - assert!( - s.credentials.is_none(), - "credentials should be cleared on fatal refresh failure" - ); - assert!( - s.last_auth_error.is_some(), - "last_auth_error should be set on fatal refresh failure" - ); - }); - } - - #[gpui::test] - async fn test_transient_refresh_keeps_credentials(cx: &mut TestAppContext) { - let http_client = FakeHttpClient::create(move |_request| async move { - Ok(http_client::Response::builder() - .status(500) - .body(http_client::AsyncBody::from("Internal Server Error"))?) - }); - - let state = cx.new(|_cx| State { - credentials: Some(make_expired_credentials()), - sign_in_task: None, - refresh_task: None, - load_task: None, - credentials_provider: Arc::new(FakeCredentialsProvider::new()), - auth_generation: 0, - last_auth_error: None, - }); - - let weak_state = cx.read(|_cx| state.downgrade()); - let http: Arc = http_client; - - let weak = weak_state.clone(); - let http_clone = http.clone(); - let result = cx - .spawn(async move |mut cx| get_fresh_credentials(&weak, &http_clone, &mut cx).await) - .await; - - cx.run_until_parked(); - - assert!(result.is_err(), "transient refresh should return an error"); - cx.read(|cx| { - let s = state.read(cx); - assert!( - s.credentials.is_some(), - "credentials should be kept on transient refresh failure" - ); - assert!( - s.last_auth_error.is_none(), - "last_auth_error should not be set on transient refresh failure" - ); - }); - } - - #[gpui::test] - async fn test_sign_out_during_refresh_discards_result(cx: &mut TestAppContext) { - let (gate_tx, gate_rx) = futures::channel::oneshot::channel::<()>(); - let gate_rx = Arc::new(Mutex::new(Some(gate_rx))); - let gate_rx_clone = gate_rx.clone(); - - let http_client = FakeHttpClient::create(move |_request| { - let gate_rx = gate_rx_clone.clone(); - async move { - // Wait until the gate is opened, simulating a slow network. - let rx = gate_rx.lock().take(); - if let Some(rx) = rx { - let _ = rx.await; - } - let body = fake_token_response(); - Ok(http_client::Response::builder() - .status(200) - .body(http_client::AsyncBody::from(body))?) - } - }); - - let creds_provider = Arc::new(FakeCredentialsProvider::new()); - let state = cx.new(|_cx| State { - credentials: Some(make_expired_credentials()), - sign_in_task: None, - refresh_task: None, - load_task: None, - credentials_provider: creds_provider.clone(), - auth_generation: 0, - last_auth_error: None, - }); - - let weak_state = cx.read(|_cx| state.downgrade()); - let http: Arc = http_client; - - // Start a refresh - let weak = weak_state.clone(); - let http_clone = http.clone(); - let refresh_task = - cx.spawn(async move |mut cx| get_fresh_credentials(&weak, &http_clone, &mut cx).await); - - cx.run_until_parked(); - - // Sign out while the refresh is in-flight - cx.update(|cx| { - do_sign_out(&weak_state, cx).detach(); - }); - cx.run_until_parked(); - - // Now let the refresh respond by opening the gate - let _ = gate_tx.send(()); - cx.run_until_parked(); - - let result = refresh_task.await; - assert!(result.is_err(), "refresh should fail after sign-out"); - - cx.read(|cx| { - let s = state.read(cx); - assert!( - s.credentials.is_none(), - "sign-out should have cleared credentials" - ); - }); - } - - #[gpui::test] - async fn test_sign_out_completes_fully(cx: &mut TestAppContext) { - let creds_provider = Arc::new(FakeCredentialsProvider::new()); - // Pre-populate the credential store - creds_provider - .storage - .lock() - .replace(("Bearer".to_string(), b"some-creds".to_vec())); - - let state = cx.new(|_cx| State { - credentials: Some(make_fresh_credentials()), - sign_in_task: None, - refresh_task: None, - load_task: None, - credentials_provider: creds_provider.clone(), - auth_generation: 0, - last_auth_error: None, - }); - - let weak_state = cx.read(|_cx| state.downgrade()); - let sign_out_task = cx.update(|cx| do_sign_out(&weak_state, cx)); - - cx.run_until_parked(); - sign_out_task.await.expect("sign-out should succeed"); - - assert!( - creds_provider.storage.lock().is_none(), - "credential store should be empty after sign-out" - ); - cx.read(|cx| { - assert!( - !state.read(cx).is_authenticated(), - "state should show not authenticated" - ); - }); - } - - #[gpui::test] - async fn test_authenticate_awaits_initial_load(cx: &mut TestAppContext) { - let creds = make_fresh_credentials(); - let creds_json = serde_json::to_vec(&creds).unwrap(); - let creds_provider = Arc::new(FakeCredentialsProvider::new()); - creds_provider - .storage - .lock() - .replace(("Bearer".to_string(), creds_json)); - - let http_client = FakeHttpClient::create(|_| async { - Ok(http_client::Response::builder() - .status(200) - .body(http_client::AsyncBody::default())?) - }); - - let provider = - cx.update(|cx| OpenAiSubscribedProvider::new(http_client, creds_provider, cx)); - - // Before load completes, authenticate should still await the load. - let auth_task = cx.update(|cx| provider.authenticate(cx)); - - // Drive the load to completion. - cx.run_until_parked(); - - let result = auth_task.await; - assert!( - result.is_ok(), - "authenticate should succeed after load completes with valid credentials" - ); - } } diff --git a/crates/openai_subscribed/Cargo.toml b/crates/openai_subscribed/Cargo.toml new file mode 100644 index 0000000000..4ce7c8db7f --- /dev/null +++ b/crates/openai_subscribed/Cargo.toml @@ -0,0 +1,36 @@ +[package] +name = "openai_subscribed" +version = "0.1.0" +edition.workspace = true +publish.workspace = true +license = "GPL-3.0-or-later" + +[lints] +workspace = true + +[lib] +path = "src/openai_subscribed.rs" + +[dependencies] +anyhow.workspace = true +base64.workspace = true +credentials_provider.workspace = true +futures.workspace = true +gpui.workspace = true +http_client.workspace = true +language_model.workspace = true +log.workspace = true +oauth_callback_server.workspace = true +open_ai.workspace = true +rand.workspace = true +serde.workspace = true +serde_json.workspace = true +sha2.workspace = true +smol.workspace = true +url.workspace = true +util.workspace = true + +[dev-dependencies] +gpui = { workspace = true, features = ["test-support"] } +http_client = { workspace = true, features = ["test-support"] } +parking_lot.workspace = true diff --git a/crates/openai_subscribed/LICENSE-GPL b/crates/openai_subscribed/LICENSE-GPL new file mode 120000 index 0000000000..89e542f750 --- /dev/null +++ b/crates/openai_subscribed/LICENSE-GPL @@ -0,0 +1 @@ +../../LICENSE-GPL \ No newline at end of file diff --git a/crates/openai_subscribed/src/openai_subscribed.rs b/crates/openai_subscribed/src/openai_subscribed.rs new file mode 100644 index 0000000000..21473f05b0 --- /dev/null +++ b/crates/openai_subscribed/src/openai_subscribed.rs @@ -0,0 +1,1362 @@ +use anyhow::{Context as _, Result, anyhow}; +use base64::Engine as _; +use base64::engine::general_purpose::URL_SAFE_NO_PAD; +use credentials_provider::CredentialsProvider; +use futures::{FutureExt, StreamExt, future::BoxFuture, future::Shared}; +use gpui::{App, AsyncApp, Context, Entity, SharedString, Task, WeakEntity}; +use http_client::{ + AsyncBody, CustomHeaders, HttpClient, Method, Request as HttpRequest, + http::{HeaderName, HeaderValue}, +}; +use language_model::{ + LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent, + LanguageModelEffortLevel, LanguageModelId, LanguageModelName, LanguageModelProviderId, + LanguageModelProviderName, LanguageModelRequest, LanguageModelToolChoice, RateLimiter, +}; +use open_ai::{ReasoningEffort, responses::stream_response}; +use rand::RngCore as _; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use std::sync::Arc; +use std::time::{SystemTime, UNIX_EPOCH}; +use url::form_urlencoded; +use util::ResultExt as _; + +use open_ai::completion::{OpenAiResponseEventMapper, into_open_ai_response}; + +pub const PROVIDER_ID: LanguageModelProviderId = LanguageModelProviderId::new("openai-subscribed"); +pub const PROVIDER_NAME: LanguageModelProviderName = + LanguageModelProviderName::new("ChatGPT Subscription"); + +const CODEX_BASE_URL: &str = "https://chatgpt.com/backend-api/codex"; +const OPENAI_TOKEN_URL: &str = "https://auth.openai.com/oauth/token"; +const OPENAI_AUTHORIZE_URL: &str = "https://auth.openai.com/oauth/authorize"; +const CLIENT_ID: &str = "app_EMoamEEZ73f0CkXaXp7hrann"; + +const CREDENTIALS_KEY: &str = "https://chatgpt.com/backend-api/codex"; +const TOKEN_REFRESH_BUFFER_MS: u64 = 5 * 60 * 1000; + +#[derive(Serialize, Deserialize, Clone, Debug)] +struct CodexCredentials { + access_token: String, + refresh_token: String, + expires_at_ms: u64, + account_id: Option, + email: Option, +} + +impl CodexCredentials { + fn is_expired(&self) -> bool { + let now = now_ms(); + now + TOKEN_REFRESH_BUFFER_MS >= self.expires_at_ms + } +} + +pub struct State { + credentials: Option, + sign_in_task: Option>>, + refresh_task: Option>>>>, + load_task: Option>>>>, + credentials_provider: Arc, + http_client: Arc, + auth_generation: u64, + last_auth_error: Option, +} + +#[derive(Debug)] +enum RefreshError { + Fatal(anyhow::Error), + Transient(anyhow::Error), +} + +impl std::fmt::Display for RefreshError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + RefreshError::Fatal(e) => write!(f, "{e}"), + RefreshError::Transient(e) => write!(f, "{e}"), + } + } +} + +impl State { + /// Creates the state and starts loading any persisted credentials. + /// [`State::load_task`] resolves once the load finishes. + pub fn new( + http_client: Arc, + credentials_provider: Arc, + cx: &mut Context, + ) -> Self { + let load_task = cx + .spawn({ + let credentials_provider = credentials_provider.clone(); + async move |this, cx| { + let result = credentials_provider + .read_credentials(CREDENTIALS_KEY, cx) + .await; + this.update(cx, |state, cx| { + if let Ok(Some((_, bytes))) = result { + match serde_json::from_slice::(&bytes) { + Ok(creds) => state.credentials = Some(creds), + Err(err) => { + log::warn!( + "Failed to deserialize ChatGPT subscription credentials: {err}" + ); + } + } + } + state.load_task = None; + cx.notify(); + })?; + Ok::<(), Arc>(()) + } + }) + .shared(); + + Self { + credentials: None, + sign_in_task: None, + refresh_task: None, + load_task: Some(load_task), + credentials_provider, + http_client, + auth_generation: 0, + last_auth_error: None, + } + } + + pub fn is_authenticated(&self) -> bool { + self.credentials.is_some() + } + + pub fn email(&self) -> Option<&str> { + self.credentials.as_ref().and_then(|c| c.email.as_deref()) + } + + pub fn is_signing_in(&self) -> bool { + self.sign_in_task.is_some() + } + + pub fn last_auth_error(&self) -> Option { + self.last_auth_error.clone() + } + + /// The in-flight task loading persisted credentials, or `None` once the + /// initial load has finished. + pub fn load_task(&self) -> Option>>>> { + self.load_task.clone() + } + + /// Starts the browser-based OAuth sign-in flow. No-op while a sign-in is + /// already in progress; observe the entity to react to the outcome. + pub fn sign_in(&mut self, cx: &mut Context) { + if self.is_signing_in() { + return; + } + + let http_client = self.http_client.clone(); + let task = cx.spawn(async move |this, cx| { + match do_oauth_flow(http_client, cx).await { + Ok(creds) => { + let persist_result = async { + let credentials_provider = + this.read_with(cx, |state, _| state.credentials_provider.clone())?; + let json = serde_json::to_vec(&creds)?; + credentials_provider + .write_credentials(CREDENTIALS_KEY, "Bearer", &json, cx) + .await?; + anyhow::Ok(()) + } + .await; + + match persist_result { + Ok(()) => { + this.update(cx, |state, cx| { + state.credentials = Some(creds); + state.sign_in_task = None; + state.last_auth_error = None; + cx.notify(); + }) + .log_err(); + } + Err(err) => { + log::error!( + "ChatGPT subscription sign-in failed to persist credentials: {err:?}" + ); + this.update(cx, |state, cx| { + state.sign_in_task = None; + state.last_auth_error = + Some("Failed to save credentials. Please try again.".into()); + cx.notify(); + }) + .log_err(); + } + } + } + Err(err) => { + log::error!("ChatGPT subscription sign-in failed: {err:?}"); + this.update(cx, |state, cx| { + state.sign_in_task = None; + state.last_auth_error = Some("Sign-in failed. Please try again.".into()); + cx.notify(); + }) + .log_err(); + } + } + anyhow::Ok(()) + }); + + self.last_auth_error = None; + self.sign_in_task = Some(task); + cx.notify(); + } + + /// Clears credentials and in-flight work immediately (so observers see the + /// sign-out right away); the returned task deletes the persisted + /// credentials. + pub fn sign_out(&mut self, cx: &mut Context) -> Task> { + self.auth_generation += 1; + self.credentials = None; + self.sign_in_task = None; + self.refresh_task = None; + self.last_auth_error = None; + cx.notify(); + + let credentials_provider = self.credentials_provider.clone(); + cx.spawn(async move |_this, cx| { + credentials_provider + .delete_credentials(CREDENTIALS_KEY, cx) + .await + .context("Failed to delete ChatGPT subscription credentials from keychain")?; + anyhow::Ok(()) + }) + } +} + +// +// The ChatGPT Subscription provider routes requests to chatgpt.com/backend-api/codex, +// which only supports a subset of OpenAI models. This list is maintained separately +// from the standard OpenAI API model list (open_ai::Model). +// +// TODO: The Codex CLI fetches this list dynamically from +// `GET /models?client_version=...` (see +// codex-rs/codex-api/src/endpoint/models.rs in openai/codex) and falls back to +// a bundled models.json. Beyond going stale, the static approach also can't +// model per-account access (e.g. free accounts cannot use gpt-5.4 even though +// paid accounts can), so the backend still rejects some requests. The bundled +// list at +// codex-rs/models-manager/models.json (openai/codex) is the closest +// approximation; the entries below mirror that file's picker-visible models. +#[derive(Clone, Debug, PartialEq)] +pub enum ChatGptModel { + Gpt56Sol, + Gpt56Terra, + Gpt56Luna, + Gpt55, + Gpt54, + Gpt54Mini, +} + +impl ChatGptModel { + pub fn all() -> Vec { + vec![ + Self::Gpt56Sol, + Self::Gpt56Terra, + Self::Gpt56Luna, + Self::Gpt55, + Self::Gpt54, + Self::Gpt54Mini, + ] + } + + pub fn id(&self) -> &str { + match self { + Self::Gpt56Sol => "gpt-5.6-sol", + Self::Gpt56Terra => "gpt-5.6-terra", + Self::Gpt56Luna => "gpt-5.6-luna", + Self::Gpt55 => "gpt-5.5", + Self::Gpt54 => "gpt-5.4", + Self::Gpt54Mini => "gpt-5.4-mini", + } + } + + pub fn display_name(&self) -> &str { + match self { + Self::Gpt56Sol => "GPT-5.6 Sol", + Self::Gpt56Terra => "GPT-5.6 Terra", + Self::Gpt56Luna => "GPT-5.6 Luna", + Self::Gpt55 => "GPT-5.5", + Self::Gpt54 => "GPT-5.4", + Self::Gpt54Mini => "GPT-5.4 Mini", + } + } + + fn max_token_count(&self) -> u64 { + match self { + Self::Gpt56Sol | Self::Gpt56Terra | Self::Gpt56Luna => 372_000, + Self::Gpt55 | Self::Gpt54 | Self::Gpt54Mini => 272_000, + } + } + + fn max_output_tokens(&self) -> Option { + // Codex model metadata does not expose a max output token cap for these + // models. Source: openai/codex models-manager/models.json. + None + } + + fn supports_images(&self) -> bool { + true + } + + fn default_reasoning_effort(&self) -> Option { + match self { + Self::Gpt56Sol => Some(ReasoningEffort::Low), + Self::Gpt56Terra | Self::Gpt56Luna | Self::Gpt55 | Self::Gpt54 | Self::Gpt54Mini => { + Some(ReasoningEffort::Medium) + } + } + } + + fn supported_reasoning_efforts(&self) -> &'static [ReasoningEffort] { + match self { + Self::Gpt56Sol | Self::Gpt56Terra | Self::Gpt56Luna => &[ + ReasoningEffort::Low, + ReasoningEffort::Medium, + ReasoningEffort::High, + ReasoningEffort::XHigh, + ReasoningEffort::Max, + ], + Self::Gpt55 | Self::Gpt54 | Self::Gpt54Mini => &[ + ReasoningEffort::Low, + ReasoningEffort::Medium, + ReasoningEffort::High, + ReasoningEffort::XHigh, + ], + } + } + + fn supports_parallel_tool_calls(&self) -> bool { + true + } + + fn supports_prompt_cache_key(&self) -> bool { + true + } + + pub fn supports_priority(&self) -> bool { + match self { + Self::Gpt56Sol | Self::Gpt56Terra | Self::Gpt56Luna | Self::Gpt55 | Self::Gpt54 => true, + Self::Gpt54Mini => false, + } + } +} + +/// Creates a [`LanguageModel`] for `model` that authenticates through +/// `state`'s credentials, refreshing them as needed. +pub fn create_language_model( + model: ChatGptModel, + state: &Entity, + cx: &App, +) -> Arc { + Arc::new(OpenAiSubscribedLanguageModel { + id: LanguageModelId::from(model.id().to_string()), + http_client: state.read(cx).http_client.clone(), + model, + state: state.clone(), + request_limiter: RateLimiter::new(4), + }) +} + +struct OpenAiSubscribedLanguageModel { + id: LanguageModelId, + model: ChatGptModel, + state: Entity, + http_client: Arc, + request_limiter: RateLimiter, +} + +impl LanguageModel for OpenAiSubscribedLanguageModel { + fn id(&self) -> LanguageModelId { + self.id.clone() + } + + fn name(&self) -> LanguageModelName { + LanguageModelName::from(self.model.display_name().to_string()) + } + + fn provider_id(&self) -> LanguageModelProviderId { + PROVIDER_ID + } + + fn provider_name(&self) -> LanguageModelProviderName { + PROVIDER_NAME + } + + fn supports_tools(&self) -> bool { + true + } + + fn supports_images(&self) -> bool { + self.model.supports_images() + } + + fn supports_tool_choice(&self, _choice: LanguageModelToolChoice) -> bool { + true + } + + fn supports_streaming_tools(&self) -> bool { + true + } + + fn supports_thinking(&self) -> bool { + true + } + + fn supports_fast_mode(&self) -> bool { + self.model.supports_priority() + } + + fn supported_effort_levels(&self) -> Vec { + let default_effort = self.model.default_reasoning_effort(); + self.model + .supported_reasoning_efforts() + .iter() + .copied() + .filter_map(|effort| { + let (name, value) = match effort { + ReasoningEffort::None => return None, + ReasoningEffort::Minimal => ("Minimal", "minimal"), + ReasoningEffort::Low => ("Low", "low"), + ReasoningEffort::Medium => ("Medium", "medium"), + ReasoningEffort::High => ("High", "high"), + ReasoningEffort::XHigh => ("Extra High", "xhigh"), + ReasoningEffort::Max => ("Max", "max"), + }; + + Some(LanguageModelEffortLevel { + name: name.into(), + value: value.into(), + is_default: Some(effort) == default_effort, + }) + }) + .collect() + } + + fn telemetry_id(&self) -> String { + format!("openai-subscribed/{}", self.model.id()) + } + + fn max_token_count(&self) -> u64 { + self.model.max_token_count() + } + + fn max_output_tokens(&self) -> Option { + self.model.max_output_tokens() + } + + fn stream_completion( + &self, + mut request: LanguageModelRequest, + cx: &AsyncApp, + ) -> BoxFuture< + 'static, + Result< + futures::stream::BoxStream< + 'static, + Result, + >, + LanguageModelCompletionError, + >, + > { + if !self.model.supports_priority() { + request.speed = None; + } + + // The Codex backend rejects `max_output_tokens` (`Unsupported parameter`), + // unlike the public OpenAI Responses API. Pass `None` so the field is + // omitted from the serialized request body entirely. + let mut responses_request = match into_open_ai_response( + request, + self.model.id(), + self.model.supports_parallel_tool_calls(), + self.model.supports_prompt_cache_key(), + /*max_output_tokens*/ None, + self.model.default_reasoning_effort(), + self.model + .supported_reasoning_efforts() + .contains(&ReasoningEffort::None), + &PROVIDER_ID, + ) { + Ok(request) => request, + Err(error) => return async move { Err(error.into()) }.boxed(), + }; + responses_request.store = Some(false); + + // `into_open_ai_response` already hoists system messages into + // `instructions`, which is the only form the Codex backend accepts. + // Codex has only ever been sent requests with the field present + // (possibly empty), so keep sending it even without system messages. + responses_request.instructions.get_or_insert_default(); + + let state = self.state.downgrade(); + let http_client = self.http_client.clone(); + let request_limiter = self.request_limiter.clone(); + + let future = cx.spawn(async move |cx| { + let creds = get_fresh_credentials(&state, &http_client, cx).await?; + + let mut header_pairs: Vec<(HeaderName, HeaderValue)> = vec![ + ( + HeaderName::from_static("originator"), + HeaderValue::from_static("zed"), + ), + ( + HeaderName::from_static("openai-beta"), + HeaderValue::from_static("responses=experimental"), + ), + ]; + if let Some(ref id) = creds.account_id { + if !id.is_empty() { + if let Ok(value) = HeaderValue::from_str(id) { + header_pairs.push((HeaderName::from_static("chatgpt-account-id"), value)); + } + } + } + let extra_headers = CustomHeaders::new(header_pairs); + + let access_token = creds.access_token.clone(); + request_limiter + .stream(async move { + stream_response( + http_client.as_ref(), + PROVIDER_NAME.0.as_str(), + CODEX_BASE_URL, + &access_token, + responses_request, + &extra_headers, + ) + .await + .map_err(LanguageModelCompletionError::from) + }) + .await + }); + + async move { + let mapper = OpenAiResponseEventMapper::new(PROVIDER_ID); + Ok(mapper.map_stream(future.await?.boxed()).boxed()) + } + .boxed() + } +} + +async fn get_fresh_credentials( + state: &WeakEntity, + http_client: &Arc, + cx: &mut AsyncApp, +) -> Result { + let (creds, existing_task) = state + .read_with(&*cx, |s, _| (s.credentials.clone(), s.refresh_task.clone())) + .map_err(LanguageModelCompletionError::Other)?; + + let creds = creds.ok_or(LanguageModelCompletionError::NoApiKey { + provider: PROVIDER_NAME, + })?; + + if !creds.is_expired() { + return Ok(creds); + } + + // If another caller is already refreshing, await their result. + if let Some(shared_task) = existing_task { + return shared_task + .await + .map_err(|e| LanguageModelCompletionError::Other(anyhow::anyhow!("{e}"))); + } + + // We are the first caller to notice expiry — spawn the refresh task. + let http_client_clone = http_client.clone(); + let state_clone = state.clone(); + let refresh_token_value = creds.refresh_token.clone(); + + // Capture the generation so we can detect sign-outs that happened during refresh. + let generation = state + .read_with(&*cx, |s, _| s.auth_generation) + .map_err(LanguageModelCompletionError::Other)?; + + let shared_task = cx + .spawn(async move |cx| { + let result = refresh_token(&http_client_clone, &refresh_token_value).await; + + match result { + Ok(refreshed) => { + let persist_result: Result> = async { + // Check if auth_generation changed (sign-out during refresh). + let current_generation = state_clone + .read_with(&*cx, |s, _| s.auth_generation) + .map_err(|e| Arc::new(e))?; + if current_generation != generation { + return Err(Arc::new(anyhow!( + "Sign-out occurred during token refresh" + ))); + } + + let credentials_provider = state_clone + .read_with(&*cx, |s, _| s.credentials_provider.clone()) + .map_err(|e| Arc::new(e))?; + + let json = + serde_json::to_vec(&refreshed).map_err(|e| Arc::new(e.into()))?; + + credentials_provider + .write_credentials(CREDENTIALS_KEY, "Bearer", &json, &*cx) + .await + .map_err(|e| Arc::new(e))?; + + state_clone + .update(cx, |s, _| { + s.credentials = Some(refreshed.clone()); + s.refresh_task = None; + }) + .map_err(|e| Arc::new(e))?; + + Ok(refreshed) + } + .await; + + // Clear refresh_task on failure too. + if persist_result.is_err() { + state_clone + .update(cx, |s, _| { + s.refresh_task = None; + }) + .ok(); + } + + persist_result + } + Err(RefreshError::Fatal(e)) => { + log::error!("ChatGPT subscription token refresh failed fatally: {e:?}"); + state_clone + .update(cx, |s, cx| { + s.refresh_task = None; + s.credentials = None; + s.last_auth_error = + Some("Your session has expired. Please sign in again.".into()); + cx.notify(); + }) + .ok(); + // Also clear the keychain so stale credentials aren't loaded next time. + if let Ok(credentials_provider) = + state_clone.read_with(&*cx, |s, _| s.credentials_provider.clone()) + { + credentials_provider + .delete_credentials(CREDENTIALS_KEY, &*cx) + .await + .log_err(); + } + Err(Arc::new(e)) + } + Err(RefreshError::Transient(e)) => { + log::warn!("ChatGPT subscription token refresh failed transiently: {e:?}"); + state_clone + .update(cx, |s, _| { + s.refresh_task = None; + }) + .ok(); + Err(Arc::new(e)) + } + } + }) + .shared(); + + // Store the shared task so concurrent callers can join on it. + state + .update(cx, |s, _| { + s.refresh_task = Some(shared_task.clone()); + }) + .map_err(LanguageModelCompletionError::Other)?; + + shared_task + .await + .map_err(|e| LanguageModelCompletionError::Other(anyhow::anyhow!("{e}"))) +} + +#[derive(Deserialize)] +struct TokenResponse { + access_token: String, + refresh_token: String, + #[serde(default)] + id_token: Option, + expires_in: u64, + #[serde(default)] + email: Option, +} + +// The OAuth client registered for `CLIENT_ID` (the Codex CLI's client) only allows +// `http://localhost:1455/auth/callback` and `http://localhost:1457/auth/callback` +// as redirect URIs; using anything else (different host, port, or path) causes +// auth.openai.com to reject the authorize request with a generic `unknown_error` +// before redirecting back. Keep these in sync with the Codex CLI's redirect URI +// allow-list (see codex-rs/login/src/server.rs in openai/codex). +const CODEX_CALLBACK_HOST: &str = "localhost"; +const CODEX_CALLBACK_PORT: u16 = 1455; +const CODEX_CALLBACK_FALLBACK_PORT: u16 = 1457; +const CODEX_CALLBACK_PATH: &str = "/auth/callback"; + +async fn do_oauth_flow( + http_client: Arc, + cx: &AsyncApp, +) -> Result { + // Start the callback server FIRST so the redirect URI is ready + let (redirect_uri, callback_rx) = + oauth_callback_server::start_oauth_callback_server_with_config( + oauth_callback_server::OAuthCallbackServerConfig { + host: CODEX_CALLBACK_HOST, + preferred_port: CODEX_CALLBACK_PORT, + fallback_port: Some(CODEX_CALLBACK_FALLBACK_PORT), + path: CODEX_CALLBACK_PATH, + }, + ) + .context("Failed to start OAuth callback server")?; + + // PKCE verifier: 32 random bytes → base64url (no padding) + let mut verifier_bytes = [0u8; 32]; + rand::rng().fill_bytes(&mut verifier_bytes); + let verifier = URL_SAFE_NO_PAD.encode(verifier_bytes); + + // PKCE challenge: SHA-256(verifier) → base64url + let mut hasher = Sha256::new(); + hasher.update(verifier.as_bytes()); + let challenge = URL_SAFE_NO_PAD.encode(hasher.finalize().as_slice()); + + // CSRF state: 16 random bytes → hex string + let mut state_bytes = [0u8; 16]; + rand::rng().fill_bytes(&mut state_bytes); + let oauth_state: String = state_bytes.iter().map(|b| format!("{b:02x}")).collect(); + + let mut auth_url = url::Url::parse(OPENAI_AUTHORIZE_URL).expect("valid base URL"); + auth_url + .query_pairs_mut() + .append_pair("client_id", CLIENT_ID) + .append_pair("redirect_uri", &redirect_uri) + // Deliberately excludes `api.connectors.read api.connectors.invoke` + // (which Codex CLI requests): extra scopes inflate the + // access-token JWT, and the serialized credentials must fit within + // Windows Credential Manager's 2560-byte blob limit + // (CRED_MAX_CREDENTIAL_BLOB_SIZE). See #58541. + .append_pair("scope", "openid profile email offline_access") + .append_pair("response_type", "code") + .append_pair("code_challenge", &challenge) + .append_pair("code_challenge_method", "S256") + .append_pair("id_token_add_organizations", "true") + .append_pair("state", &oauth_state) + .append_pair("codex_cli_simplified_flow", "true") + .append_pair("originator", "zed"); + + // Open browser AFTER the listener is ready + cx.update(|cx| cx.open_url(auth_url.as_str())); + + // Await the callback + let callback = callback_rx + .await + .map_err(|_| anyhow!("OAuth callback was cancelled"))? + .context("OAuth callback failed")?; + + // Validate CSRF state + if callback.state != oauth_state { + return Err(anyhow!("OAuth state mismatch")); + } + + let tokens = exchange_code(&http_client, &callback.code, &verifier, &redirect_uri) + .await + .context("Token exchange failed")?; + + let jwt = tokens + .id_token + .as_deref() + .unwrap_or(tokens.access_token.as_str()); + let claims = extract_jwt_claims(jwt); + + Ok(CodexCredentials { + access_token: tokens.access_token, + refresh_token: tokens.refresh_token, + expires_at_ms: now_ms() + tokens.expires_in * 1000, + account_id: claims.account_id, + email: claims.email.or(tokens.email), + }) +} + +async fn exchange_code( + client: &Arc, + code: &str, + verifier: &str, + redirect_uri: &str, +) -> Result { + let body = form_urlencoded::Serializer::new(String::new()) + .append_pair("grant_type", "authorization_code") + .append_pair("client_id", CLIENT_ID) + .append_pair("code", code) + .append_pair("redirect_uri", redirect_uri) + .append_pair("code_verifier", verifier) + .finish(); + + let request = HttpRequest::builder() + .method(Method::POST) + .uri(OPENAI_TOKEN_URL) + .header("Content-Type", "application/x-www-form-urlencoded") + .body(AsyncBody::from(body))?; + + let mut response = client.send(request).await?; + let mut body = String::new(); + smol::io::AsyncReadExt::read_to_string(response.body_mut(), &mut body).await?; + + if !response.status().is_success() { + return Err(anyhow!( + "Token exchange failed (HTTP {}): {body}", + response.status() + )); + } + + serde_json::from_str::(&body).context("Failed to parse token response") +} + +async fn refresh_token( + client: &Arc, + refresh_token: &str, +) -> Result { + let body = form_urlencoded::Serializer::new(String::new()) + .append_pair("grant_type", "refresh_token") + .append_pair("client_id", CLIENT_ID) + .append_pair("refresh_token", refresh_token) + .finish(); + + let request = HttpRequest::builder() + .method(Method::POST) + .uri(OPENAI_TOKEN_URL) + .header("Content-Type", "application/x-www-form-urlencoded") + .body(AsyncBody::from(body)) + .map_err(|e| RefreshError::Transient(e.into()))?; + + let mut response = client + .send(request) + .await + .map_err(|e| RefreshError::Transient(e))?; + let status = response.status(); + let mut body = String::new(); + smol::io::AsyncReadExt::read_to_string(response.body_mut(), &mut body) + .await + .map_err(|e| RefreshError::Transient(e.into()))?; + + if !status.is_success() { + let err = anyhow!("Token refresh failed (HTTP {}): {body}", status); + // 400/401/403 indicate a revoked or invalid refresh token. + // 5xx and other errors are treated as transient. + if status == http_client::StatusCode::BAD_REQUEST + || status == http_client::StatusCode::UNAUTHORIZED + || status == http_client::StatusCode::FORBIDDEN + { + return Err(RefreshError::Fatal(err)); + } + return Err(RefreshError::Transient(err)); + } + + let tokens: TokenResponse = + serde_json::from_str(&body).map_err(|e| RefreshError::Transient(e.into()))?; + let jwt = tokens + .id_token + .as_deref() + .unwrap_or(tokens.access_token.as_str()); + let claims = extract_jwt_claims(jwt); + + Ok(CodexCredentials { + access_token: tokens.access_token, + refresh_token: tokens.refresh_token, + expires_at_ms: now_ms() + tokens.expires_in * 1000, + account_id: claims.account_id, + email: claims.email.or(tokens.email), + }) +} + +struct JwtClaims { + account_id: Option, + email: Option, +} + +/// Extract claims from a JWT payload (base64url middle segment). +/// Extracts `chatgpt_account_id` from three possible locations (matching Roo Code's +/// implementation) and the `email` claim. +fn extract_jwt_claims(jwt: &str) -> JwtClaims { + let Some(payload_b64) = jwt.split('.').nth(1) else { + return JwtClaims { + account_id: None, + email: None, + }; + }; + let Ok(payload) = URL_SAFE_NO_PAD.decode(payload_b64) else { + return JwtClaims { + account_id: None, + email: None, + }; + }; + let Ok(claims) = serde_json::from_slice::(&payload) else { + return JwtClaims { + account_id: None, + email: None, + }; + }; + + let account_id = claims + .get("chatgpt_account_id") + .and_then(|v| v.as_str()) + .or_else(|| { + claims + .get("https://api.openai.com/auth") + .and_then(|v| v.get("chatgpt_account_id")) + .and_then(|v| v.as_str()) + }) + .or_else(|| { + claims + .get("organizations") + .and_then(|v| v.as_array()) + .and_then(|arr| arr.first()) + .and_then(|org| org.get("id")) + .and_then(|v| v.as_str()) + }) + .map(|s| s.to_owned()); + + let email = claims + .get("email") + .and_then(|v| v.as_str()) + .map(|s| s.to_owned()); + + JwtClaims { account_id, email } +} + +fn now_ms() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_millis() as u64) + .unwrap_or_else(|err| { + log::error!("System clock is before UNIX epoch: {err}"); + 0 + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use gpui::{AppContext as _, TestAppContext}; + use http_client::FakeHttpClient; + use parking_lot::Mutex; + use std::future::Future; + use std::pin::Pin; + use std::sync::atomic::{AtomicUsize, Ordering}; + + #[gpui::test] + async fn test_concurrent_refresh_deduplicates(cx: &mut TestAppContext) { + let refresh_count = Arc::new(AtomicUsize::new(0)); + let refresh_count_clone = refresh_count.clone(); + + let http_client = FakeHttpClient::create(move |_request| { + let refresh_count = refresh_count_clone.clone(); + async move { + refresh_count.fetch_add(1, Ordering::SeqCst); + let body = fake_token_response(); + Ok(http_client::Response::builder() + .status(200) + .body(http_client::AsyncBody::from(body))?) + } + }); + + let http: Arc = http_client; + let state = make_state(http.clone(), Some(make_expired_credentials()), cx); + + let weak_state = cx.read(|_cx| state.downgrade()); + + // Spawn two concurrent refresh attempts. + let weak1 = weak_state.clone(); + let http1 = http.clone(); + let task1 = + cx.spawn(async move |mut cx| get_fresh_credentials(&weak1, &http1, &mut cx).await); + + let weak2 = weak_state.clone(); + let http2 = http.clone(); + let task2 = + cx.spawn(async move |mut cx| get_fresh_credentials(&weak2, &http2, &mut cx).await); + + // Drive both to completion. + cx.run_until_parked(); + let result1 = task1.await; + let result2 = task2.await; + + assert!(result1.is_ok(), "first refresh should succeed"); + assert!(result2.is_ok(), "second refresh should succeed"); + assert_eq!(result1.unwrap().access_token, "fresh_access"); + assert_eq!(result2.unwrap().access_token, "fresh_access"); + assert_eq!( + refresh_count.load(Ordering::SeqCst), + 1, + "refresh_token should only be called once despite two concurrent callers" + ); + } + + #[gpui::test] + async fn test_fresh_credentials_skip_refresh(cx: &mut TestAppContext) { + let refresh_count = Arc::new(AtomicUsize::new(0)); + let refresh_count_clone = refresh_count.clone(); + + let http_client = FakeHttpClient::create(move |_request| { + let refresh_count = refresh_count_clone.clone(); + async move { + refresh_count.fetch_add(1, Ordering::SeqCst); + let body = fake_token_response(); + Ok(http_client::Response::builder() + .status(200) + .body(http_client::AsyncBody::from(body))?) + } + }); + + let http: Arc = http_client; + let state = make_state(http.clone(), Some(make_fresh_credentials()), cx); + + let weak_state = cx.read(|_cx| state.downgrade()); + + let weak = weak_state.clone(); + let http_clone = http.clone(); + let result = cx + .spawn(async move |mut cx| get_fresh_credentials(&weak, &http_clone, &mut cx).await) + .await; + + assert!(result.is_ok()); + assert_eq!(result.unwrap().access_token, "fresh_access"); + assert_eq!( + refresh_count.load(Ordering::SeqCst), + 0, + "no refresh should happen when credentials are fresh" + ); + } + + #[gpui::test] + async fn test_no_credentials_returns_no_api_key(cx: &mut TestAppContext) { + let http_client = FakeHttpClient::create(|_| async { + Ok(http_client::Response::builder() + .status(200) + .body(http_client::AsyncBody::default())?) + }); + + let http: Arc = http_client; + let state = make_state(http.clone(), None, cx); + + let weak_state = cx.read(|_cx| state.downgrade()); + + let weak = weak_state.clone(); + let http_clone = http.clone(); + let result = cx + .spawn(async move |mut cx| get_fresh_credentials(&weak, &http_clone, &mut cx).await) + .await; + + assert!(matches!( + result, + Err(LanguageModelCompletionError::NoApiKey { .. }) + )); + } + + #[gpui::test] + async fn test_fatal_refresh_clears_auth_state(cx: &mut TestAppContext) { + let http_client = FakeHttpClient::create(move |_request| async move { + Ok(http_client::Response::builder() + .status(401) + .body(http_client::AsyncBody::from(r#"{"error":"invalid_grant"}"#))?) + }); + + let http: Arc = http_client; + let state = make_state(http.clone(), Some(make_expired_credentials()), cx); + + let weak_state = cx.read(|_cx| state.downgrade()); + + let weak = weak_state.clone(); + let http_clone = http.clone(); + let result = cx + .spawn(async move |mut cx| get_fresh_credentials(&weak, &http_clone, &mut cx).await) + .await; + + cx.run_until_parked(); + + assert!(result.is_err(), "fatal refresh should return an error"); + cx.read(|cx| { + let s = state.read(cx); + assert!( + s.credentials.is_none(), + "credentials should be cleared on fatal refresh failure" + ); + assert!( + s.last_auth_error.is_some(), + "last_auth_error should be set on fatal refresh failure" + ); + }); + } + + #[gpui::test] + async fn test_transient_refresh_keeps_credentials(cx: &mut TestAppContext) { + let http_client = FakeHttpClient::create(move |_request| async move { + Ok(http_client::Response::builder() + .status(500) + .body(http_client::AsyncBody::from("Internal Server Error"))?) + }); + + let http: Arc = http_client; + let state = make_state(http.clone(), Some(make_expired_credentials()), cx); + + let weak_state = cx.read(|_cx| state.downgrade()); + + let weak = weak_state.clone(); + let http_clone = http.clone(); + let result = cx + .spawn(async move |mut cx| get_fresh_credentials(&weak, &http_clone, &mut cx).await) + .await; + + cx.run_until_parked(); + + assert!(result.is_err(), "transient refresh should return an error"); + cx.read(|cx| { + let s = state.read(cx); + assert!( + s.credentials.is_some(), + "credentials should be kept on transient refresh failure" + ); + assert!( + s.last_auth_error.is_none(), + "last_auth_error should not be set on transient refresh failure" + ); + }); + } + + #[gpui::test] + async fn test_sign_out_during_refresh_discards_result(cx: &mut TestAppContext) { + let (gate_tx, gate_rx) = futures::channel::oneshot::channel::<()>(); + let gate_rx = Arc::new(Mutex::new(Some(gate_rx))); + let gate_rx_clone = gate_rx.clone(); + + let http_client = FakeHttpClient::create(move |_request| { + let gate_rx = gate_rx_clone.clone(); + async move { + // Wait until the gate is opened, simulating a slow network. + let rx = gate_rx.lock().take(); + if let Some(rx) = rx { + let _ = rx.await; + } + let body = fake_token_response(); + Ok(http_client::Response::builder() + .status(200) + .body(http_client::AsyncBody::from(body))?) + } + }); + + let http: Arc = http_client; + let state = make_state(http.clone(), Some(make_expired_credentials()), cx); + + let weak_state = cx.read(|_cx| state.downgrade()); + + // Start a refresh + let weak = weak_state.clone(); + let http_clone = http.clone(); + let refresh_task = + cx.spawn(async move |mut cx| get_fresh_credentials(&weak, &http_clone, &mut cx).await); + + cx.run_until_parked(); + + // Sign out while the refresh is in-flight + state.update(cx, |state, cx| { + state.sign_out(cx).detach(); + }); + cx.run_until_parked(); + + // Now let the refresh respond by opening the gate + let _ = gate_tx.send(()); + cx.run_until_parked(); + + let result = refresh_task.await; + assert!(result.is_err(), "refresh should fail after sign-out"); + + cx.read(|cx| { + let s = state.read(cx); + assert!( + s.credentials.is_none(), + "sign-out should have cleared credentials" + ); + }); + } + + #[gpui::test] + async fn test_sign_out_completes_fully(cx: &mut TestAppContext) { + let creds_provider = Arc::new(FakeCredentialsProvider::new()); + // Pre-populate the credential store + creds_provider + .storage + .lock() + .replace(("Bearer".to_string(), b"some-creds".to_vec())); + + let http: Arc = FakeHttpClient::create(|_| async { + Ok(http_client::Response::builder() + .status(200) + .body(http_client::AsyncBody::default())?) + }); + let state = make_state_with_credentials_provider( + http, + Some(make_fresh_credentials()), + creds_provider.clone(), + cx, + ); + + let sign_out_task = state.update(cx, |state, cx| state.sign_out(cx)); + + cx.run_until_parked(); + sign_out_task.await.expect("sign-out should succeed"); + + assert!( + creds_provider.storage.lock().is_none(), + "credential store should be empty after sign-out" + ); + cx.read(|cx| { + assert!( + !state.read(cx).is_authenticated(), + "state should show not authenticated" + ); + }); + } + + #[gpui::test] + async fn test_initial_load_restores_persisted_credentials(cx: &mut TestAppContext) { + let creds = make_fresh_credentials(); + let creds_json = serde_json::to_vec(&creds).unwrap(); + let creds_provider = Arc::new(FakeCredentialsProvider::new()); + creds_provider + .storage + .lock() + .replace(("Bearer".to_string(), creds_json)); + + let http: Arc = FakeHttpClient::create(|_| async { + Ok(http_client::Response::builder() + .status(200) + .body(http_client::AsyncBody::default())?) + }); + + let state = cx.new(|cx| State::new(http, creds_provider, cx)); + + let load_task = cx + .read(|cx| state.read(cx).load_task()) + .expect("constructor should start the credentials load"); + + cx.run_until_parked(); + load_task.await.expect("load should succeed"); + + cx.read(|cx| { + let state = state.read(cx); + assert!(state.is_authenticated()); + assert!(state.load_task().is_none()); + }); + } + + struct FakeCredentialsProvider { + storage: Mutex)>>, + } + + impl FakeCredentialsProvider { + fn new() -> Self { + Self { + storage: Mutex::new(None), + } + } + } + + impl CredentialsProvider for FakeCredentialsProvider { + fn read_credentials<'a>( + &'a self, + _url: &'a str, + _cx: &'a AsyncApp, + ) -> Pin)>>> + 'a>> { + Box::pin(async { Ok(self.storage.lock().clone()) }) + } + + fn write_credentials<'a>( + &'a self, + _url: &'a str, + username: &'a str, + password: &'a [u8], + _cx: &'a AsyncApp, + ) -> Pin> + 'a>> { + self.storage + .lock() + .replace((username.to_string(), password.to_vec())); + Box::pin(async { Ok(()) }) + } + + fn delete_credentials<'a>( + &'a self, + _url: &'a str, + _cx: &'a AsyncApp, + ) -> Pin> + 'a>> { + *self.storage.lock() = None; + Box::pin(async { Ok(()) }) + } + } + + fn make_state( + http_client: Arc, + credentials: Option, + cx: &mut TestAppContext, + ) -> Entity { + make_state_with_credentials_provider( + http_client, + credentials, + Arc::new(FakeCredentialsProvider::new()), + cx, + ) + } + + fn make_state_with_credentials_provider( + http_client: Arc, + credentials: Option, + credentials_provider: Arc, + cx: &mut TestAppContext, + ) -> Entity { + cx.new(|_cx| State { + credentials, + sign_in_task: None, + refresh_task: None, + load_task: None, + credentials_provider, + http_client, + auth_generation: 0, + last_auth_error: None, + }) + } + + fn make_expired_credentials() -> CodexCredentials { + CodexCredentials { + access_token: "old_access".to_string(), + refresh_token: "old_refresh".to_string(), + expires_at_ms: 0, + account_id: None, + email: None, + } + } + + fn make_fresh_credentials() -> CodexCredentials { + CodexCredentials { + access_token: "fresh_access".to_string(), + refresh_token: "fresh_refresh".to_string(), + expires_at_ms: now_ms() + 3_600_000, + account_id: None, + email: None, + } + } + + fn fake_token_response() -> String { + serde_json::json!({ + "access_token": "fresh_access", + "refresh_token": "fresh_refresh", + "expires_in": 3600 + }) + .to_string() + } +}