Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 18 additions & 1 deletion Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,7 @@ members = [
"crates/net",
"crates/node_runtime",
"crates/notifications",
"crates/oauth_callback_server",
"crates/ollama",
"crates/onboarding",
"crates/opencode",
Expand Down Expand Up @@ -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" }
Expand Down
2 changes: 1 addition & 1 deletion crates/context_server/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand Down
170 changes: 23 additions & 147 deletions crates/context_server/src/oauth.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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.
Expand All @@ -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 -------------------------------------------------------
Expand Down
6 changes: 6 additions & 0 deletions crates/language_models/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"] }

Expand All @@ -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"] }
11 changes: 10 additions & 1 deletion crates/language_models/src/language_models.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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,
);
}
1 change: 1 addition & 0 deletions crates/language_models/src/provider.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
1 change: 1 addition & 0 deletions crates/language_models/src/provider/open_ai.rs
Original file line number Diff line number Diff line change
Expand Up @@ -383,6 +383,7 @@ impl OpenAiLanguageModel {
&api_url,
&api_key,
request,
vec![],
);
let response = request.await?;
Ok(response)
Expand Down
1 change: 1 addition & 0 deletions crates/language_models/src/provider/open_ai_compatible.rs
Original file line number Diff line number Diff line change
Expand Up @@ -289,6 +289,7 @@ impl OpenAiCompatibleLanguageModel {
&api_url,
&api_key,
request,
vec![],
);
let response = request.await?;
Ok(response)
Expand Down
Loading
Loading