Add ChatGPT subscription provider via OAuth 2.0 PKCE

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" }

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
 

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:<port>/callback?code=...&state=...`.
     pub fn parse_query(query: &str) -> Result<Self> {
-        let mut code: Option<String> = None;
-        let mut state: Option<String> = None;
-        let mut error: Option<String> = None;
-        let mut error_description: Option<String> = 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<Result<OAuthCallback>>,
-)> {
-    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,
-                    "<html><body><h1>Authorization successful</h1>\
-                     <p>You can close this tab and return to Zed.</p></body></html>",
-                ),
-                Err(err) => {
-                    log::error!("OAuth callback error: {}", err);
-                    (
-                        400,
-                        "<html><body><h1>Authorization failed</h1>\
-                         <p>Something went wrong. Please try again from Zed.</p></body></html>",
-                    )
-                }
-            };
-
-            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<OAuthCallback>>)> {
+    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<OAuthCallback> {
-    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 -------------------------------------------------------

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"] }

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,
+    );
 }

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;

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)

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)