diff --git a/Cargo.lock b/Cargo.lock index 4637613580866d..f5784e6a657970 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3597,6 +3597,7 @@ dependencies = [ "http_client", "log", "net", + "oauth_callback_server", "parking_lot", "pollster 0.4.0", "postage", @@ -3608,7 +3609,6 @@ dependencies = [ "sha2", "slotmap", "tempfile", - "tiny_http", "url", "util", ] @@ -9683,20 +9683,26 @@ dependencies = [ "log", "menu", "mistral", + "oauth_callback_server", "ollama", "open_ai", "open_router", "opencode", + "parking_lot", "pretty_assertions", + "rand 0.9.4", "release_channel", "schemars 1.0.4", "serde", "serde_json", "settings", + "sha2", + "smol", "strum 0.27.2", "tokio", "ui", "ui_input", + "url", "util", "x_ai", ] @@ -11509,6 +11515,17 @@ dependencies = [ "tokio-util", ] +[[package]] +name = "oauth_callback_server" +version = "0.1.0" +dependencies = [ + "anyhow", + "futures 0.3.32", + "log", + "tiny_http", + "url", +] + [[package]] name = "objc" version = "0.2.7" diff --git a/Cargo.toml b/Cargo.toml index a0135052e37843..3d2701f4766d7a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -140,6 +140,7 @@ members = [ "crates/net", "crates/node_runtime", "crates/notifications", + "crates/oauth_callback_server", "crates/ollama", "crates/onboarding", "crates/opencode", @@ -399,6 +400,7 @@ nc = { path = "crates/nc" } net = { path = "crates/net" } node_runtime = { path = "crates/node_runtime" } notifications = { path = "crates/notifications" } +oauth_callback_server = { path = "crates/oauth_callback_server" } ollama = { path = "crates/ollama" } onboarding = { path = "crates/onboarding" } opencode = { path = "crates/opencode" } diff --git a/crates/context_server/Cargo.toml b/crates/context_server/Cargo.toml index 39288c5a6d87be..efea3dd250f800 100644 --- a/crates/context_server/Cargo.toml +++ b/crates/context_server/Cargo.toml @@ -26,6 +26,7 @@ gpui.workspace = true http_client = { workspace = true, features = ["test-support"] } log.workspace = true net.workspace = true +oauth_callback_server.workspace = true parking_lot.workspace = true rand.workspace = true postage.workspace = true @@ -36,7 +37,6 @@ settings.workspace = true sha2.workspace = true slotmap.workspace = true tempfile.workspace = true -tiny_http.workspace = true url = { workspace = true, features = ["serde"] } util.workspace = true diff --git a/crates/context_server/src/oauth.rs b/crates/context_server/src/oauth.rs index de6b7d50e8a020..c5afa28c9b8649 100644 --- a/crates/context_server/src/oauth.rs +++ b/crates/context_server/src/oauth.rs @@ -20,18 +20,18 @@ use anyhow::{Context as _, Result, anyhow, bail}; use async_trait::async_trait; use base64::Engine as _; use futures::AsyncReadExt as _; +use futures::FutureExt as _; use futures::channel::mpsc; +use futures::future::BoxFuture; use http_client::{AsyncBody, HttpClient, Request}; use parking_lot::Mutex as SyncMutex; use rand::Rng as _; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; -use std::str::FromStr; use std::sync::Arc; use std::time::{Duration, SystemTime}; use url::Url; -use util::ResultExt as _; /// The CIMD URL where Zed's OAuth client metadata document is hosted. pub const CIMD_URL: &str = "https://zed.dev/oauth/client-metadata.json"; @@ -992,58 +992,14 @@ impl OAuthCallback { /// Parse the query string from a callback URL like /// `http://127.0.0.1:/callback?code=...&state=...`. pub fn parse_query(query: &str) -> Result { - let mut code: Option = None; - let mut state: Option = None; - let mut error: Option = None; - let mut error_description: Option = None; - - for (key, value) in url::form_urlencoded::parse(query.as_bytes()) { - match key.as_ref() { - "code" => { - if !value.is_empty() { - code = Some(value.into_owned()); - } - } - "state" => { - if !value.is_empty() { - state = Some(value.into_owned()); - } - } - "error" => { - if !value.is_empty() { - error = Some(value.into_owned()); - } - } - "error_description" => { - if !value.is_empty() { - error_description = Some(value.into_owned()); - } - } - _ => {} - } - } - - // Check for OAuth error response (RFC 6749 Section 4.1.2.1) before - // checking for missing code/state. - if let Some(error_code) = error { - bail!( - "OAuth authorization failed: {} ({})", - error_code, - error_description.as_deref().unwrap_or("no description") - ); - } - - let code = code.ok_or_else(|| anyhow!("missing 'code' parameter in OAuth callback"))?; - let state = state.ok_or_else(|| anyhow!("missing 'state' parameter in OAuth callback"))?; - - Ok(Self { code, state }) + let params = oauth_callback_server::OAuthCallbackParams::parse_query(query)?; + Ok(Self { + code: params.code, + state: params.state, + }) } } -/// How long to wait for the browser to complete the OAuth flow before giving -/// up and releasing the loopback port. -const CALLBACK_TIMEOUT: Duration = Duration::from_secs(2 * 60); - /// Start a loopback HTTP server to receive the OAuth authorization callback. /// /// Binds to an ephemeral loopback port for each flow. @@ -1056,104 +1012,24 @@ const CALLBACK_TIMEOUT: Duration = Duration::from_secs(2 * 60); /// contains `code` and `state` query parameters, responds with a minimal /// HTML page telling the user they can close the tab, and shuts down. /// -/// The callback server shuts down when the returned oneshot receiver is dropped -/// (e.g. because the authentication task was cancelled), or after a timeout -/// ([CALLBACK_TIMEOUT]). -pub async fn start_callback_server() -> Result<( - String, - futures::channel::oneshot::Receiver>, -)> { - let server = tiny_http::Server::http("127.0.0.1:0") - .map_err(|e| anyhow!(e).context("Failed to bind loopback listener for OAuth callback"))?; - let port = server - .server_addr() - .to_ip() - .context("server not bound to a TCP address")? - .port(); - - let redirect_uri = format!("http://127.0.0.1:{}/callback", port); - - let (tx, rx) = futures::channel::oneshot::channel(); - - // `tiny_http` is blocking, so we run it on a background thread. - // The `recv_timeout` loop lets us check for cancellation (the receiver - // being dropped) and enforce an overall timeout. - std::thread::spawn(move || { - let deadline = std::time::Instant::now() + CALLBACK_TIMEOUT; - - loop { - if tx.is_canceled() { - return; - } - let remaining = deadline.saturating_duration_since(std::time::Instant::now()); - if remaining.is_zero() { - return; - } - - let timeout = remaining.min(Duration::from_millis(500)); - let Some(request) = (match server.recv_timeout(timeout) { - Ok(req) => req, - Err(_) => { - let _ = tx.send(Err(anyhow!("OAuth callback server I/O error"))); - return; - } - }) else { - // Timeout with no request — loop back and check cancellation. - continue; - }; - - let result = handle_callback_request(&request); - - let (status_code, body) = match &result { - Ok(_) => ( - 200, - "

Authorization successful

\ -

You can close this tab and return to Zed.

", - ), - Err(err) => { - log::error!("OAuth callback error: {}", err); - ( - 400, - "

Authorization failed

\ -

Something went wrong. Please try again from Zed.

", - ) - } - }; - - let response = tiny_http::Response::from_string(body) - .with_status_code(status_code) - .with_header( - tiny_http::Header::from_str("Content-Type: text/html") - .expect("failed to construct response header"), - ) - .with_header( - tiny_http::Header::from_str("Keep-Alive: timeout=0,max=0") - .expect("failed to construct response header"), - ); - request.respond(response).log_err(); - - let _ = tx.send(result); - return; +/// The callback server shuts down when the returned future is dropped (e.g. +/// because the authentication task was cancelled), or after a timeout. +pub fn start_callback_server() -> Result<(String, BoxFuture<'static, Result>)> { + let (redirect_uri, rx) = oauth_callback_server::start_oauth_callback_server()?; + let future = async move { + match rx.await { + Ok(Ok(params)) => Ok(OAuthCallback { + code: params.code, + state: params.state, + }), + Ok(Err(e)) => Err(e), + Err(_) => Err(anyhow!( + "OAuth callback server was shut down before receiving a response" + )), } - }); - - Ok((redirect_uri, rx)) -} - -/// Extract the `code` and `state` query parameters from an OAuth callback -/// request to `/callback`. -fn handle_callback_request(request: &tiny_http::Request) -> Result { - let url = Url::parse(&format!("http://localhost{}", request.url())) - .context("malformed callback request URL")?; - - if url.path() != "/callback" { - bail!("unexpected path in OAuth callback: {}", url.path()); } - - let query = url - .query() - .ok_or_else(|| anyhow!("OAuth callback has no query string"))?; - OAuthCallback::parse_query(query) + .boxed(); + Ok((redirect_uri, future)) } // -- JSON fetch helper ------------------------------------------------------- diff --git a/crates/language_models/Cargo.toml b/crates/language_models/Cargo.toml index 7afcbb8aea109b..7e2d2618ea9695 100644 --- a/crates/language_models/Cargo.toml +++ b/crates/language_models/Cargo.toml @@ -47,19 +47,24 @@ 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"] } 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"] } @@ -71,5 +76,6 @@ feature_flags.workspace = true gpui = { workspace = true, features = ["test-support"] } http_client = { workspace = true, features = ["test-support"] } language_model = { workspace = true, features = ["test-support"] } +parking_lot.workspace = true pretty_assertions.workspace = true settings = { workspace = true, features = ["test-support"] } diff --git a/crates/language_models/src/language_models.rs b/crates/language_models/src/language_models.rs index 1920815871b7c7..97ac0b2c0abb63 100644 --- a/crates/language_models/src/language_models.rs +++ b/crates/language_models/src/language_models.rs @@ -27,6 +27,7 @@ use crate::provider::ollama::OllamaLanguageModelProvider; use crate::provider::open_ai::OpenAiLanguageModelProvider; use crate::provider::open_ai_compatible::OpenAiCompatibleLanguageModelProvider; use crate::provider::open_router::OpenRouterLanguageModelProvider; +use crate::provider::openai_subscribed::OpenAiSubscribedProvider; use crate::provider::opencode::OpenCodeLanguageModelProvider; use crate::provider::vercel_ai_gateway::VercelAiGatewayLanguageModelProvider; use crate::provider::x_ai::XAiLanguageModelProvider; @@ -324,10 +325,18 @@ fn register_language_model_providers( registry.register_provider( Arc::new(OpenCodeLanguageModelProvider::new( client.http_client(), - credentials_provider, + credentials_provider.clone(), cx, )), cx, ); registry.register_provider(Arc::new(CopilotChatLanguageModelProvider::new(cx)), cx); + registry.register_provider( + Arc::new(OpenAiSubscribedProvider::new( + client.http_client(), + credentials_provider, + cx, + )), + cx, + ); } diff --git a/crates/language_models/src/provider.rs b/crates/language_models/src/provider.rs index acf6579cae6b0b..8c0b4fbfd6c8e9 100644 --- a/crates/language_models/src/provider.rs +++ b/crates/language_models/src/provider.rs @@ -10,6 +10,7 @@ pub mod ollama; pub mod open_ai; pub mod open_ai_compatible; pub mod open_router; +pub mod openai_subscribed; pub mod opencode; pub mod vercel_ai_gateway; diff --git a/crates/language_models/src/provider/open_ai.rs b/crates/language_models/src/provider/open_ai.rs index 3ecc4ba02f0919..ddf8b0e688597a 100644 --- a/crates/language_models/src/provider/open_ai.rs +++ b/crates/language_models/src/provider/open_ai.rs @@ -383,6 +383,7 @@ impl OpenAiLanguageModel { &api_url, &api_key, request, + vec![], ); let response = request.await?; Ok(response) diff --git a/crates/language_models/src/provider/open_ai_compatible.rs b/crates/language_models/src/provider/open_ai_compatible.rs index b2baf79ecdc6bb..d4e1f21f20a057 100644 --- a/crates/language_models/src/provider/open_ai_compatible.rs +++ b/crates/language_models/src/provider/open_ai_compatible.rs @@ -289,6 +289,7 @@ impl OpenAiCompatibleLanguageModel { &api_url, &api_key, request, + vec![], ); let response = request.await?; Ok(response) diff --git a/crates/language_models/src/provider/openai_subscribed.rs b/crates/language_models/src/provider/openai_subscribed.rs new file mode 100644 index 00000000000000..822db3b1f7ddc6 --- /dev/null +++ b/crates/language_models/src/provider/openai_subscribed.rs @@ -0,0 +1,1464 @@ +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::{AnyView, App, AsyncApp, Context, Entity, SharedString, Task, Window}; +use http_client::{AsyncBody, HttpClient, Method, Request as HttpRequest}; +use language_model::{ + AuthenticateError, IconOrSvg, LanguageModel, LanguageModelCompletionError, + LanguageModelCompletionEvent, LanguageModelId, LanguageModelName, LanguageModelProvider, + LanguageModelProviderId, LanguageModelProviderName, LanguageModelProviderState, + 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 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 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, +} + +impl OpenAiSubscribedProvider { + pub fn new( + http_client: Arc, + 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 sign_out(&self, cx: &mut App) -> Task> { + do_sign_out(&self.state.downgrade(), cx) + } + + 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), + }) + } +} + +impl LanguageModelProviderState for OpenAiSubscribedProvider { + type ObservableEntity = State; + + fn observable_entity(&self) -> Option> { + Some(self.state.clone()) + } +} + +impl LanguageModelProvider for OpenAiSubscribedProvider { + fn id(&self) -> LanguageModelProviderId { + PROVIDER_ID + } + + fn name(&self) -> LanguageModelProviderName { + PROVIDER_NAME + } + + fn icon(&self) -> IconOrSvg { + IconOrSvg::Icon(IconName::AiOpenAi) + } + + fn default_model(&self, _cx: &App) -> Option> { + Some(self.create_language_model(ChatGptModel::Gpt55)) + } + + fn default_fast_model(&self, _cx: &App) -> Option> { + // No GPT-5.5 Mini exists yet; per the OpenAI Codex docs, gpt-5.4-mini + // is the recommended fast/cheap default alongside gpt-5.5. + Some(self.create_language_model(ChatGptModel::Gpt54Mini)) + } + + fn provided_models(&self, _cx: &App) -> Vec> { + ChatGptModel::all() + .into_iter() + .map(|m| self.create_language_model(m)) + .collect() + } + + fn is_authenticated(&self, cx: &App) -> bool { + self.state.read(cx).is_authenticated() + } + + fn authenticate(&self, cx: &mut App) -> Task> { + if self.is_authenticated(cx) { + return Task::ready(Ok(())); + } + let load_task = self.state.read(cx).load_task.clone(); + 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()) + .unwrap_or(false); + if is_auth { + Ok(()) + } else { + Err(anyhow!( + "Sign in with your ChatGPT Plus or Pro subscription to use this provider." + ) + .into()) + } + }) + } else { + Task::ready(Err(anyhow!( + "Sign in with your ChatGPT Plus or Pro subscription to use this provider." + ) + .into())) + } + } + + fn configuration_view( + &self, + _target_agent: language_model::ConfigurationViewTargetAgent, + _window: &mut Window, + cx: &mut App, + ) -> AnyView { + let state = self.state.clone(); + let http_client = self.http_client.clone(); + cx.new(|_cx| ConfigurationView { state, http_client }) + .into() + } + + fn reset_credentials(&self, cx: &mut App) -> Task> { + self.sign_out(cx) + } +} + +// +// 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 or +// gpt-5.3-codex 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 { + Gpt55, + Gpt54, + Gpt54Mini, + Gpt53Codex, + Gpt52, +} + +impl ChatGptModel { + fn all() -> Vec { + vec![ + Self::Gpt55, + Self::Gpt54, + Self::Gpt54Mini, + Self::Gpt53Codex, + Self::Gpt52, + ] + } + + fn id(&self) -> &str { + match self { + Self::Gpt55 => "gpt-5.5", + Self::Gpt54 => "gpt-5.4", + Self::Gpt54Mini => "gpt-5.4-mini", + Self::Gpt53Codex => "gpt-5.3-codex", + Self::Gpt52 => "gpt-5.2", + } + } + + fn display_name(&self) -> &str { + match self { + Self::Gpt55 => "GPT-5.5", + Self::Gpt54 => "GPT-5.4", + Self::Gpt54Mini => "GPT-5.4 Mini", + Self::Gpt53Codex => "GPT-5.3 Codex", + Self::Gpt52 => "GPT-5.2", + } + } + + fn max_token_count(&self) -> u64 { + // All Codex-supported models share the backend's 272K input cap, even + // when the raw model exposes a larger context window via the public + // API (e.g. gpt-5.4 has max_context_window 1M, but the Codex backend + // caps it at 272K). Source: openai/codex models-manager/models.json. + 272_000 + } + + fn max_output_tokens(&self) -> Option { + Some(128_000) + } + + fn supports_images(&self) -> bool { + true + } + + fn reasoning_effort(&self) -> Option { + // Codex bundled models all default to Medium reasoning effort. + Some(ReasoningEffort::Medium) + } + + fn supports_none_reasoning_effort(&self) -> bool { + // The Codex backend's supported_reasoning_levels for every model in + // this list is low/medium/high/xhigh -- sending `none` is rejected. + false + } + + fn supports_parallel_tool_calls(&self) -> bool { + true + } + + fn supports_prompt_cache_key(&self) -> bool { + true + } +} + +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 { + self.model.reasoning_effort().is_some() + } + + 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, + request: LanguageModelRequest, + cx: &AsyncApp, + ) -> BoxFuture< + 'static, + Result< + futures::stream::BoxStream< + 'static, + Result, + >, + LanguageModelCompletionError, + >, + > { + // 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 = 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.reasoning_effort(), + self.model.supports_none_reasoning_effort(), + ); + responses_request.store = Some(false); + + // The Codex backend requires system messages to be in the top-level + // `instructions` field rather than as input items. + let mut instructions = Vec::new(); + responses_request.input.retain(|item| { + if let open_ai::responses::ResponseInputItem::Message(msg) = item { + if msg.role == open_ai::Role::System { + for part in &msg.content { + if let open_ai::responses::ResponseInputContent::Text { text } = part { + instructions.push(text.clone()); + } + } + return false; + } + } + true + }); + if !instructions.is_empty() { + responses_request.instructions = Some(instructions.join("\n\n")); + } + + 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 extra_headers: Vec<(String, String)> = vec![ + ("originator".into(), "zed".into()), + ("OpenAI-Beta".into(), "responses=experimental".into()), + ]; + if let Some(ref id) = creds.account_id { + if !id.is_empty() { + extra_headers.push(("ChatGPT-Account-Id".into(), id.clone())); + } + } + + 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(); + 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) + .append_pair( + "scope", + "openid profile email offline_access api.connectors.read api.connectors.invoke", + ) + .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, +} + +impl Render for ConfigurationView { + fn render(&mut self, _window: &mut Window, cx: &mut Context) -> impl IntoElement { + let state = self.state.read(cx); + + if state.is_authenticated() { + let label = state + .email() + .map(|e| format!("Signed in as {e}")) + .unwrap_or_else(|| "Signed in".to_string()); + + let weak_state = self.state.downgrade(); + return v_flex() + .child( + ConfiguredApiCard::new(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); + })), + ) + .into_any_element(); + } + + if state.is_signing_in() { + return v_flex() + .child(Label::new("Signing in…").color(Color::Muted)) + .into_any_element(); + } + + let last_auth_error = state.last_auth_error.clone(); + let provider_state = self.state.clone(); + let http_client = self.http_client.clone(); + + v_flex() + .gap_2() + .when_some(last_auth_error, |this, error| { + this.child(Label::new(error).color(Color::Error)) + }) + .child(Label::new( + "Sign in with your ChatGPT Plus or Pro subscription to use OpenAI models in Zed's agent.", + )) + .child( + Button::new("sign-in", "Sign in with ChatGPT") + .on_click(move |_, _window, cx| { + do_sign_in(&provider_state, &http_client, cx); + }), + ) + .into_any_element() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use gpui::TestAppContext; + use http_client::FakeHttpClient; + use parking_lot::Mutex; + use std::future::Future; + use std::pin::Pin; + use std::sync::atomic::{AtomicUsize, Ordering}; + + 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_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/language_models/src/provider/opencode.rs b/crates/language_models/src/provider/opencode.rs index 64b3a6fcdbd704..1d77c59f5d9335 100644 --- a/crates/language_models/src/provider/opencode.rs +++ b/crates/language_models/src/provider/opencode.rs @@ -479,6 +479,7 @@ impl OpenCodeLanguageModel { &api_url, &api_key, request, + vec![], ); let response = request.await?; Ok(response) diff --git a/crates/oauth_callback_server/Cargo.toml b/crates/oauth_callback_server/Cargo.toml new file mode 100644 index 00000000000000..5fcc6233af56d5 --- /dev/null +++ b/crates/oauth_callback_server/Cargo.toml @@ -0,0 +1,23 @@ +[package] +name = "oauth_callback_server" +version = "0.1.0" +edition.workspace = true +publish = false +license = "Apache-2.0" +description = "Loopback OAuth 2.0 callback server and shared HTML response page for Zed sign-in flows" + +[lints] +workspace = true + +[lib] +path = "src/oauth_callback_server.rs" +doctest = false + +[dependencies] +anyhow.workspace = true + +[target.'cfg(not(target_family = "wasm"))'.dependencies] +futures.workspace = true +log.workspace = true +tiny_http.workspace = true +url.workspace = true diff --git a/crates/oauth_callback_server/LICENSE-APACHE b/crates/oauth_callback_server/LICENSE-APACHE new file mode 120000 index 00000000000000..1cd601d0a3affa --- /dev/null +++ b/crates/oauth_callback_server/LICENSE-APACHE @@ -0,0 +1 @@ +../../LICENSE-APACHE \ No newline at end of file diff --git a/crates/oauth_callback_server/src/oauth_callback_server.rs b/crates/oauth_callback_server/src/oauth_callback_server.rs new file mode 100644 index 00000000000000..44ea22947347ef --- /dev/null +++ b/crates/oauth_callback_server/src/oauth_callback_server.rs @@ -0,0 +1,493 @@ +//! Loopback OAuth 2.0 callback server and shared HTML response page. +//! +//! Used by Zed's OAuth-based sign-in flows (e.g. MCP servers, ChatGPT +//! Subscription) to receive the authorization code redirect from the user's +//! browser. The HTML response page rendered to the browser is kept alongside +//! the server so all OAuth callback presentation lives in one place. + +/// Generate a styled HTML page for OAuth callback responses. +/// +/// Returns a complete HTML document (no HTTP headers) with a centered card +/// layout styled to match Zed's dark theme. The `title` is rendered as a +/// heading and `message` as body text below it. +/// +/// When `is_error` is true, a red X icon is shown instead of the green +/// checkmark. +pub fn oauth_callback_page(title: &str, message: &str, is_error: bool) -> String { + let title = html_escape(title); + let message = html_escape(message); + let (icon_bg, icon_svg) = if is_error { + ( + "#f38ba8", + r#""#, + ) + } else { + ( + "#a6e3a1", + r#""#, + ) + }; + format!( + r#" + + + + +{title} — Zed + + + +
+
+ {icon_svg} +
+

{title}

+

{message}

+
Zed
+
+ +"#, + title = title, + message = message, + icon_bg = icon_bg, + icon_svg = icon_svg, + ) +} + +fn html_escape(input: &str) -> String { + let mut output = String::with_capacity(input.len()); + for ch in input.chars() { + match ch { + '&' => output.push_str("&"), + '<' => output.push_str("<"), + '>' => output.push_str(">"), + '"' => output.push_str("""), + '\'' => output.push_str("'"), + _ => output.push(ch), + } + } + output +} + +#[cfg(not(target_family = "wasm"))] +mod server { + use super::oauth_callback_page; + use anyhow::{Context as _, Result, anyhow}; + use std::str::FromStr; + use std::time::Duration; + use url::Url; + + /// Parsed OAuth callback parameters from the authorization server redirect. + pub struct OAuthCallbackParams { + pub code: String, + pub state: String, + } + + /// Configuration for the loopback OAuth callback server. + /// + /// OAuth servers compare `redirect_uri` against a per-client allow-list using + /// exact string matching (RFC 6749 §3.1.2), so the `host`, `preferred_port`, + /// and `path` here must match what's registered for the OAuth client_id. + #[derive(Clone, Copy)] + pub struct OAuthCallbackServerConfig { + /// Host portion of the redirect URI (typically `127.0.0.1` or `localhost`). + pub host: &'static str, + /// Preferred port. Use `0` for an OS-assigned ephemeral port. + pub preferred_port: u16, + /// Optional fallback port if `preferred_port` is unavailable. Only used + /// when `preferred_port` is non-zero. + pub fallback_port: Option, + /// Callback path on the redirect URI (e.g. `/callback`, `/auth/callback`). + pub path: &'static str, + } + + impl Default for OAuthCallbackServerConfig { + fn default() -> Self { + Self { + host: "127.0.0.1", + preferred_port: 0, + fallback_port: None, + path: "/callback", + } + } + } + + impl OAuthCallbackParams { + /// Parse the query string from a callback URL like + /// `http://127.0.0.1:/callback?code=...&state=...`. + pub fn parse_query(query: &str) -> Result { + let mut code: Option = None; + let mut state: Option = None; + let mut error: Option = None; + let mut error_description: Option = None; + + for (key, value) in url::form_urlencoded::parse(query.as_bytes()) { + match key.as_ref() { + "code" => { + if !value.is_empty() { + code = Some(value.into_owned()); + } + } + "state" => { + if !value.is_empty() { + state = Some(value.into_owned()); + } + } + "error" => { + if !value.is_empty() { + error = Some(value.into_owned()); + } + } + "error_description" => { + if !value.is_empty() { + error_description = Some(value.into_owned()); + } + } + _ => {} + } + } + + if let Some(error_code) = error { + anyhow::bail!( + "OAuth authorization failed: {} ({})", + error_code, + error_description.as_deref().unwrap_or("no description") + ); + } + + let code = code.ok_or_else(|| anyhow!("missing 'code' parameter in OAuth callback"))?; + let state = + state.ok_or_else(|| anyhow!("missing 'state' parameter in OAuth callback"))?; + + Ok(Self { code, state }) + } + } + + /// How long to wait for the browser to complete the OAuth flow before giving + /// up and releasing the loopback port. + const OAUTH_CALLBACK_TIMEOUT: Duration = Duration::from_secs(2 * 60); + + /// Start a loopback HTTP server to receive the OAuth authorization callback. + /// + /// Binds to an ephemeral loopback port. Returns `(redirect_uri, callback_future)`. + /// The caller should use the redirect URI in the authorization request, open + /// the browser, then await the future to receive the callback. + pub fn start_oauth_callback_server() -> Result<( + String, + futures::channel::oneshot::Receiver>, + )> { + start_oauth_callback_server_with_config(OAuthCallbackServerConfig::default()) + } + + /// Start a loopback HTTP server with custom host/port/path. + /// + /// Use this when the OAuth client requires a specific redirect URI that the + /// default ephemeral-port `http://127.0.0.1:/callback` doesn't match. + pub fn start_oauth_callback_server_with_config( + config: OAuthCallbackServerConfig, + ) -> Result<( + String, + futures::channel::oneshot::Receiver>, + )> { + let server = bind_callback_server(&config)?; + let port = server + .server_addr() + .to_ip() + .ok_or_else(|| anyhow!("server not bound to a TCP address"))? + .port(); + + let redirect_uri = format!("http://{}:{}{}", config.host, port, config.path); + let expected_path = config.path; + + let (tx, rx) = futures::channel::oneshot::channel(); + + std::thread::spawn(move || { + let deadline = std::time::Instant::now() + OAUTH_CALLBACK_TIMEOUT; + + loop { + if tx.is_canceled() { + return; + } + let remaining = deadline.saturating_duration_since(std::time::Instant::now()); + if remaining.is_zero() { + return; + } + + let timeout = remaining.min(Duration::from_millis(500)); + let Some(request) = (match server.recv_timeout(timeout) { + Ok(req) => req, + Err(_) => { + let _ = tx.send(Err(anyhow!("OAuth callback server I/O error"))); + return; + } + }) else { + continue; + }; + + let raw_url = request.url().to_string(); + let raw_path = raw_url.split('?').next().unwrap_or(&raw_url); + if raw_path == CANCEL_PATH { + let response = tiny_http::Response::from_string("Cancelled") + .with_status_code(200) + .with_header( + tiny_http::Header::from_str("Content-Type: text/plain") + .expect("failed to construct response header"), + ) + .with_header( + tiny_http::Header::from_str("Connection: close") + .expect("failed to construct response header"), + ); + if let Err(err) = request.respond(response) { + log::error!("Failed to send OAuth cancel response: {}", err); + } + let _ = tx.send(Err(anyhow!( + "OAuth callback server was cancelled by another sign-in attempt" + ))); + return; + } + + let result = handle_oauth_callback_request(&request, expected_path); + + let (status_code, body) = match &result { + Ok(_) => ( + 200, + oauth_callback_page( + "Authorization Successful", + "You can close this tab and return to Zed.", + false, + ), + ), + Err(err) => { + log::error!("OAuth callback error: {}", err); + ( + 400, + oauth_callback_page( + "Authorization Failed", + "Something went wrong. Please try again from Zed.", + true, + ), + ) + } + }; + + let response = tiny_http::Response::from_string(body) + .with_status_code(status_code) + .with_header( + tiny_http::Header::from_str("Content-Type: text/html") + .expect("failed to construct response header"), + ) + .with_header( + tiny_http::Header::from_str("Keep-Alive: timeout=0,max=0") + .expect("failed to construct response header"), + ); + if let Err(err) = request.respond(response) { + log::error!("Failed to send OAuth callback response: {}", err); + } + + let _ = tx.send(result); + return; + } + }); + + Ok((redirect_uri, rx)) + } + + fn handle_oauth_callback_request( + request: &tiny_http::Request, + expected_path: &str, + ) -> Result { + let url = Url::parse(&format!("http://localhost{}", request.url())) + .context("malformed callback request URL")?; + + if url.path() != expected_path { + anyhow::bail!("unexpected path in OAuth callback: {}", url.path()); + } + + let query = url + .query() + .ok_or_else(|| anyhow!("OAuth callback has no query string"))?; + OAuthCallbackParams::parse_query(query) + } + + /// Callback path reserved for evicting a previously-running OAuth callback + /// server bound to the same port. Always handled, regardless of `config.path`. + const CANCEL_PATH: &str = "/cancel"; + + const BIND_MAX_ATTEMPTS: u32 = 10; + const BIND_RETRY_DELAY: Duration = Duration::from_millis(200); + const CANCEL_REQUEST_TIMEOUT: Duration = Duration::from_secs(2); + + fn bind_callback_server(config: &OAuthCallbackServerConfig) -> Result { + // Ephemeral ports always succeed; skip the cancel-retry dance entirely. + if config.preferred_port == 0 { + let addr = format!("{}:0", config.host); + return tiny_http::Server::http(&addr).map_err(|err| { + anyhow!(err).context(format!( + "Failed to bind loopback listener for OAuth callback on {addr}" + )) + }); + } + + match try_bind_with_cancel(config.host, config.preferred_port) { + Ok(server) => Ok(server), + Err(primary_err) => { + let Some(fallback_port) = config.fallback_port else { + return Err(primary_err.context(format!( + "Failed to bind loopback listener for OAuth callback on {}:{}", + config.host, config.preferred_port, + ))); + }; + log::warn!( + "OAuth callback port {}:{} unavailable; falling back to port {}", + config.host, + config.preferred_port, + fallback_port, + ); + try_bind_with_cancel(config.host, fallback_port).map_err(|fallback_err| { + fallback_err.context(format!( + "Failed to bind loopback listener for OAuth callback on {}:{} or {}:{}", + config.host, config.preferred_port, config.host, fallback_port, + )) + }) + } + } + } + + /// Attempts to bind to a fixed `host:port`. On `AddrInUse`, sends a single + /// `GET /cancel` to the existing listener (to evict a previous OAuth flow + /// from this or a compatible client) and retries. + fn try_bind_with_cancel(host: &'static str, port: u16) -> Result { + let addr = format!("{host}:{port}"); + let mut cancel_attempted = false; + let mut last_err: Option = None; + + for _ in 0..BIND_MAX_ATTEMPTS { + match tiny_http::Server::http(&addr) { + Ok(server) => return Ok(server), + Err(err) => { + let is_addr_in_use = err + .downcast_ref::() + .map(|io_err| io_err.kind() == std::io::ErrorKind::AddrInUse) + .unwrap_or(false); + + if !is_addr_in_use { + return Err(anyhow!(err).context(format!( + "Failed to bind loopback listener for OAuth callback on {addr}" + ))); + } + + if !cancel_attempted { + cancel_attempted = true; + if let Err(cancel_err) = send_cancel_request(host, port) { + log::warn!( + "Failed to cancel previous OAuth callback server on {addr}: {cancel_err}" + ); + } + } + + last_err = Some(anyhow!(err)); + std::thread::sleep(BIND_RETRY_DELAY); + } + } + } + + Err(last_err + .unwrap_or_else(|| anyhow!("unknown bind error")) + .context(format!( + "OAuth callback port {addr} remained in use after {BIND_MAX_ATTEMPTS} attempts" + ))) + } + + /// Sends `GET /cancel` to a listener on `host:port`, asking it to shut down. + /// + /// Best-effort: errors here are surfaced to the caller for logging but do + /// not block the subsequent rebind attempt. + fn send_cancel_request(host: &str, port: u16) -> std::io::Result<()> { + use std::io::{Read as _, Write as _}; + use std::net::{TcpStream, ToSocketAddrs as _}; + + let addr = format!("{host}:{port}") + .to_socket_addrs()? + .next() + .ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("could not resolve {host}:{port}"), + ) + })?; + let mut stream = TcpStream::connect_timeout(&addr, CANCEL_REQUEST_TIMEOUT)?; + stream.set_read_timeout(Some(CANCEL_REQUEST_TIMEOUT))?; + stream.set_write_timeout(Some(CANCEL_REQUEST_TIMEOUT))?; + + stream.write_all(b"GET /cancel HTTP/1.1\r\n")?; + stream.write_all(format!("Host: {host}:{port}\r\n").as_bytes())?; + stream.write_all(b"Connection: close\r\n\r\n")?; + + // Drain the response so the server can close cleanly. We don't care + // about the body; errors here are harmless. + let mut buf = [0u8; 64]; + let _ = stream.read(&mut buf); + Ok(()) + } +} + +#[cfg(not(target_family = "wasm"))] +pub use server::{ + OAuthCallbackParams, OAuthCallbackServerConfig, start_oauth_callback_server, + start_oauth_callback_server_with_config, +}; diff --git a/crates/open_ai/src/completion.rs b/crates/open_ai/src/completion.rs index 5be0a5d6286ca9..3b62391678e47f 100644 --- a/crates/open_ai/src/completion.rs +++ b/crates/open_ai/src/completion.rs @@ -259,8 +259,9 @@ pub fn into_open_ai_response( ResponseRequest { model: model_id.into(), + instructions: None, input: input_items, - store: false, + store: Some(false), include, stream, temperature, diff --git a/crates/open_ai/src/responses.rs b/crates/open_ai/src/responses.rs index 37335454c321c1..954f9b5ad56f57 100644 --- a/crates/open_ai/src/responses.rs +++ b/crates/open_ai/src/responses.rs @@ -9,9 +9,10 @@ use crate::{ReasoningEffort, RequestError, Role, ToolChoice}; #[derive(Serialize, Debug)] pub struct Request { pub model: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub instructions: Option, #[serde(skip_serializing_if = "Vec::is_empty")] pub input: Vec, - pub store: bool, #[serde(skip_serializing_if = "Vec::is_empty")] pub include: Vec, #[serde(default)] @@ -32,6 +33,8 @@ pub struct Request { pub prompt_cache_key: Option, #[serde(skip_serializing_if = "Option::is_none")] pub reasoning: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub store: Option, } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)] @@ -395,13 +398,17 @@ pub async fn stream_response( api_url: &str, api_key: &str, request: Request, + extra_headers: Vec<(String, String)>, ) -> Result>, RequestError> { let uri = format!("{api_url}/responses"); - let request_builder = HttpRequest::builder() + let mut request_builder = HttpRequest::builder() .method(Method::POST) .uri(uri) .header("Content-Type", "application/json") .header("Authorization", format!("Bearer {}", api_key.trim())); + for (name, value) in &extra_headers { + request_builder = request_builder.header(name.as_str(), value.as_str()); + } let is_streaming = request.stream; let request = request_builder diff --git a/crates/project/src/context_server_store.rs b/crates/project/src/context_server_store.rs index 9effe4a4638762..de2e1e3ceffdc1 100644 --- a/crates/project/src/context_server_store.rs +++ b/crates/project/src/context_server_store.rs @@ -1063,9 +1063,8 @@ impl ContextServerStore { // Start a loopback HTTP server on an ephemeral port. The redirect URI // includes this port so the browser sends the callback directly to our // process. - let (redirect_uri, callback_rx) = oauth::start_callback_server() - .await - .context("Failed to start OAuth callback server")?; + let (redirect_uri, callback_rx) = + oauth::start_callback_server().context("Failed to start OAuth callback server")?; let http_client = cx.update(|cx| cx.http_client()); let credentials_provider = cx.update(|cx| zed_credentials_provider::global(cx)); @@ -1093,9 +1092,6 @@ impl ContextServerStore { let callback = callback_rx .await - .map_err(|_| { - anyhow::anyhow!("OAuth callback server was shut down before receiving a response") - })? .context("OAuth callback server received an invalid request")?; if callback.state != state_param {