Add OAuth Web Flow auth option for llm provider extensions

This commit is contained in:
Richard Feldman
2025-12-04 16:50:44 -05:00
parent a48bd10da0
commit 2d3a3521ba
5 changed files with 346 additions and 0 deletions

View File

@@ -249,4 +249,36 @@ world extension {
/// Read an environment variable.
import llm-get-env-var: func(name: string) -> option<string>;
// =========================================================================
// OAuth Web Auth Flow Imports
// =========================================================================
use llm-provider.{oauth-web-auth-config, oauth-web-auth-result, oauth-http-request, oauth-http-response};
/// Start an OAuth web authentication flow.
///
/// This will:
/// 1. Start a localhost server to receive the OAuth callback
/// 2. Open the auth URL in the user's default browser
/// 3. Wait for the callback (up to the timeout)
/// 4. Return the callback URL with query parameters
///
/// The extension is responsible for:
/// - Constructing the auth URL with client_id, redirect_uri, scope, state, etc.
/// - Parsing the callback URL to extract the authorization code
/// - Exchanging the code for tokens using llm-oauth-http-request
import llm-oauth-start-web-auth: func(config: oauth-web-auth-config) -> result<oauth-web-auth-result, string>;
/// Make an HTTP request for OAuth token exchange.
///
/// This is a simple HTTP client for OAuth flows, allowing the extension
/// to handle token exchange with full control over serialization.
import llm-oauth-http-request: func(request: oauth-http-request) -> result<oauth-http-response, string>;
/// Open a URL in the user's default browser.
///
/// Useful for OAuth flows that need to open a browser but handle the
/// callback differently (e.g., polling-based flows).
import llm-oauth-open-browser: func(url: string) -> result<_, string>;
}

View File

@@ -252,4 +252,51 @@ interface llm-provider {
/// Minimum token count for a message to be cached.
min-total-token-count: u64,
}
// =========================================================================
// OAuth Web Auth Flow Types
// =========================================================================
/// Configuration for starting an OAuth web authentication flow.
record oauth-web-auth-config {
/// The URL to open in the user's browser to start authentication.
/// This should include client_id, redirect_uri, scope, state, etc.
auth-url: string,
/// The path to listen on for the OAuth callback (e.g., "/callback").
/// A localhost server will be started to receive the redirect.
callback-path: string,
/// Timeout in seconds to wait for the callback (default: 300 = 5 minutes).
timeout-secs: option<u32>,
}
/// Result of an OAuth web authentication flow.
record oauth-web-auth-result {
/// The full callback URL that was received, including query parameters.
/// The extension is responsible for parsing the code, state, etc.
callback-url: string,
/// The port that was used for the localhost callback server.
port: u32,
}
/// A generic HTTP request for OAuth token exchange.
record oauth-http-request {
/// The URL to request.
url: string,
/// HTTP method (e.g., "POST", "GET").
method: string,
/// Request headers as key-value pairs.
headers: list<tuple<string, string>>,
/// Request body as a string (for form-encoded or JSON bodies).
body: string,
}
/// Response from an OAuth HTTP request.
record oauth-http-response {
/// HTTP status code.
status: u16,
/// Response headers as key-value pairs.
headers: list<tuple<string, string>>,
/// Response body as a string.
body: string,
}
}

View File

@@ -249,4 +249,36 @@ world extension {
/// Read an environment variable.
import llm-get-env-var: func(name: string) -> option<string>;
// =========================================================================
// OAuth Web Auth Flow Imports
// =========================================================================
use llm-provider.{oauth-web-auth-config, oauth-web-auth-result, oauth-http-request, oauth-http-response};
/// Start an OAuth web authentication flow.
///
/// This will:
/// 1. Start a localhost server to receive the OAuth callback
/// 2. Open the auth URL in the user's default browser
/// 3. Wait for the callback (up to the timeout)
/// 4. Return the callback URL with query parameters
///
/// The extension is responsible for:
/// - Constructing the auth URL with client_id, redirect_uri, scope, state, etc.
/// - Parsing the callback URL to extract the authorization code
/// - Exchanging the code for tokens using llm-oauth-http-request
import llm-oauth-start-web-auth: func(config: oauth-web-auth-config) -> result<oauth-web-auth-result, string>;
/// Make an HTTP request for OAuth token exchange.
///
/// This is a simple HTTP client for OAuth flows, allowing the extension
/// to handle token exchange with full control over serialization.
import llm-oauth-http-request: func(request: oauth-http-request) -> result<oauth-http-response, string>;
/// Open a URL in the user's default browser.
///
/// Useful for OAuth flows that need to open a browser but handle the
/// callback differently (e.g., polling-based flows).
import llm-oauth-open-browser: func(url: string) -> result<_, string>;
}

View File

@@ -252,4 +252,51 @@ interface llm-provider {
/// Minimum token count for a message to be cached.
min-total-token-count: u64,
}
// =========================================================================
// OAuth Web Auth Flow Types
// =========================================================================
/// Configuration for starting an OAuth web authentication flow.
record oauth-web-auth-config {
/// The URL to open in the user's browser to start authentication.
/// This should include client_id, redirect_uri, scope, state, etc.
auth-url: string,
/// The path to listen on for the OAuth callback (e.g., "/callback").
/// A localhost server will be started to receive the redirect.
callback-path: string,
/// Timeout in seconds to wait for the callback (default: 300 = 5 minutes).
timeout-secs: option<u32>,
}
/// Result of an OAuth web authentication flow.
record oauth-web-auth-result {
/// The full callback URL that was received, including query parameters.
/// The extension is responsible for parsing the code, state, etc.
callback-url: string,
/// The port that was used for the localhost callback server.
port: u32,
}
/// A generic HTTP request for OAuth token exchange.
record oauth-http-request {
/// The URL to request.
url: string,
/// HTTP method (e.g., "POST", "GET").
method: string,
/// Request headers as key-value pairs.
headers: list<tuple<string, string>>,
/// Request body as a string (for form-encoded or JSON bodies).
body: string,
}
/// Response from an OAuth HTTP request.
record oauth-http-response {
/// HTTP status code.
status: u16,
/// Response headers as key-value pairs.
headers: list<tuple<string, string>>,
/// Response body as a string.
body: string,
}
}

View File

@@ -24,12 +24,15 @@ use gpui::{BackgroundExecutor, SharedString};
use language::{BinaryStatus, LanguageName, language_settings::AllLanguageSettings};
use project::project_settings::ProjectSettings;
use semver::Version;
use smol::net::TcpListener;
use std::{
env,
io::{BufRead, Write},
net::Ipv4Addr,
path::{Path, PathBuf},
str::FromStr,
sync::{Arc, OnceLock},
time::Duration,
};
use task::{SpawnInTerminal, ZedDebugConfig};
use url::Url;
@@ -1244,6 +1247,191 @@ impl ExtensionImports for WasmState {
Ok(env::var(&name).ok())
}
async fn llm_oauth_start_web_auth(
&mut self,
config: llm_provider::OauthWebAuthConfig,
) -> wasmtime::Result<Result<llm_provider::OauthWebAuthResult, String>> {
let auth_url = config.auth_url;
let callback_path = config.callback_path;
let timeout_secs = config.timeout_secs.unwrap_or(300);
self.on_main_thread(move |cx| {
async move {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.map_err(|e| anyhow::anyhow!("Failed to bind localhost server: {}", e))?;
let port = listener
.local_addr()
.map_err(|e| anyhow::anyhow!("Failed to get local address: {}", e))?
.port();
cx.update(|cx| {
cx.open_url(&auth_url);
})?;
let accept_future = async {
let (stream, _) = listener
.accept()
.await
.map_err(|e| anyhow::anyhow!("Failed to accept connection: {}", e))?;
let mut reader = smol::io::BufReader::new(&stream);
let mut request_line = String::new();
smol::io::AsyncBufReadExt::read_line(&mut reader, &mut request_line)
.await
.map_err(|e| anyhow::anyhow!("Failed to read request: {}", e))?;
let callback_url = if let Some(path_start) = request_line.find(' ') {
if let Some(path_end) = request_line[path_start + 1..].find(' ') {
let path = &request_line[path_start + 1..path_start + 1 + path_end];
if path.starts_with(&callback_path) || path.starts_with(&format!("/{}", callback_path.trim_start_matches('/'))) {
format!("http://localhost:{}{}", port, path)
} else {
return Err(anyhow::anyhow!(
"Unexpected callback path: {}",
path
));
}
} else {
return Err(anyhow::anyhow!("Malformed HTTP request"));
}
} else {
return Err(anyhow::anyhow!("Malformed HTTP request"));
};
let response = "HTTP/1.1 200 OK\r\n\
Content-Type: text/html\r\n\
Connection: close\r\n\
\r\n\
<!DOCTYPE html>\
<html><head><title>Authentication Complete</title></head>\
<body style=\"font-family: system-ui, sans-serif; display: flex; justify-content: center; align-items: center; height: 100vh; margin: 0;\">\
<div style=\"text-align: center;\">\
<h1>Authentication Complete</h1>\
<p>You can close this window and return to Zed.</p>\
</div></body></html>";
let mut writer = &stream;
smol::io::AsyncWriteExt::write_all(&mut writer, response.as_bytes())
.await
.ok();
smol::io::AsyncWriteExt::flush(&mut writer).await.ok();
Ok(callback_url)
};
let timeout_duration = Duration::from_secs(timeout_secs as u64);
let callback_url = smol::future::or(
accept_future,
async {
smol::Timer::after(timeout_duration).await;
Err(anyhow::anyhow!(
"OAuth callback timed out after {} seconds",
timeout_secs
))
},
)
.await?;
Ok(llm_provider::OauthWebAuthResult {
callback_url,
port: port as u32,
})
}
.boxed_local()
})
.await
.to_wasmtime_result()
}
async fn llm_oauth_http_request(
&mut self,
request: llm_provider::OauthHttpRequest,
) -> wasmtime::Result<Result<llm_provider::OauthHttpResponse, String>> {
let http_client = self.http_client.clone();
self.on_main_thread(move |_cx| {
async move {
let method = match request.method.to_uppercase().as_str() {
"GET" => ::http_client::Method::GET,
"POST" => ::http_client::Method::POST,
"PUT" => ::http_client::Method::PUT,
"DELETE" => ::http_client::Method::DELETE,
"PATCH" => ::http_client::Method::PATCH,
_ => {
return Err(anyhow::anyhow!(
"Unsupported HTTP method: {}",
request.method
));
}
};
let mut builder = ::http_client::HttpRequest::builder()
.method(method)
.uri(&request.url);
for (key, value) in &request.headers {
builder = builder.header(key.as_str(), value.as_str());
}
let body = if request.body.is_empty() {
AsyncBody::empty()
} else {
AsyncBody::from(request.body.into_bytes())
};
let http_request = builder
.body(body)
.map_err(|e| anyhow::anyhow!("Failed to build request: {}", e))?;
let mut response = http_client
.send(http_request)
.await
.map_err(|e| anyhow::anyhow!("HTTP request failed: {}", e))?;
let status = response.status().as_u16();
let headers: Vec<(String, String)> = response
.headers()
.iter()
.map(|(k, v)| (k.to_string(), v.to_str().unwrap_or("").to_string()))
.collect();
let mut body_bytes = Vec::new();
futures::AsyncReadExt::read_to_end(response.body_mut(), &mut body_bytes)
.await
.map_err(|e| anyhow::anyhow!("Failed to read response body: {}", e))?;
let body = String::from_utf8_lossy(&body_bytes).to_string();
Ok(llm_provider::OauthHttpResponse {
status,
headers,
body,
})
}
.boxed_local()
})
.await
.to_wasmtime_result()
}
async fn llm_oauth_open_browser(
&mut self,
url: String,
) -> wasmtime::Result<Result<(), String>> {
self.on_main_thread(move |cx| {
async move {
cx.update(|cx| {
cx.open_url(&url);
})?;
Ok(())
}
.boxed_local()
})
.await
.to_wasmtime_result()
}
}
// =============================================================================