Add OAuth Web Flow auth option for llm provider extensions
This commit is contained in:
@@ -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>;
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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>;
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
|
||||
Reference in New Issue
Block a user