External/buzz
1
0
Fork 0
mirror of https://github.com/block/buzz.git synced 2026-08-18 06:50:31 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat(sprout-agent): Databricks provider with OAuth 2.0 PKCE auth (#698)

Browse Source 
Signed-off-by: tlongwell-block <109685178+tlongwell-block@users.noreply.github.com>
Co-authored-by: Dawn (sprout agent) <c6237ef84fa537c78dcee78efd2d4e59f728859c7f194da42ac51ededfa0be05@sprout-oss.stage.blox.sqprod.co>
This commit is contained in:
tlongwell-block
2026-05-21 12:08:24 -04:00
committed by GitHub
co-authored by Dawn
parent 300b51b02a
commit fa9e26fa5c
7 changed files with 1234 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+69 -7
View File
@@ -1336,8 +1336,8 @@ dependencies = [
"libc",
"log",
"rustversion",
"windows-link 0.1.3",
"windows-result 0.3.4",
"windows-link 0.2.1",
"windows-result 0.4.1",
]
[[package]]
@@ -1691,7 +1691,7 @@ dependencies = [
"libc",
"percent-encoding",
"pin-project-lite",
"socket2 0.5.10",
"socket2 0.6.3",
"tokio",
"tower-service",
"tracing",
@@ -1709,7 +1709,7 @@ dependencies = [
"js-sys",
"log",
"wasm-bindgen",
"windows-core 0.61.2",
"windows-core 0.62.2",
]
[[package]]
@@ -2318,6 +2318,12 @@ dependencies = [
"thiserror 1.0.69",
]
[[package]]
name = "ndk-context"
version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "27b02d87554356db9e9a873add8782d4ea6e3e58ea071a9adb9a2e8ddb884a8b"
[[package]]
name = "negentropy"
version = "0.3.1"
@@ -2476,6 +2482,15 @@ dependencies = [
"libc",
]
[[package]]
name = "objc2"
version = "0.6.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3a12a8ed07aefc768292f076dc3ac8c48f3781c8f2d5851dd3d98950e8c5a89f"
dependencies = [
"objc2-encode",
]
[[package]]
name = "objc2-core-foundation"
version = "0.3.2"
@@ -2485,6 +2500,22 @@ dependencies = [
"bitflags",
]
[[package]]
name = "objc2-encode"
version = "4.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ef25abbcd74fb2609453eb695bd2f860d389e457f67dc17cafc8b8cbc89d0c33"
[[package]]
name = "objc2-foundation"
version = "0.3.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e3e0adef53c21f888deb4fa59fc59f7eb17404926ee8a6f59f5df0fd7f9f3272"
dependencies = [
"bitflags",
"objc2",
]
[[package]]
name = "objc2-io-kit"
version = "0.3.2"
@@ -2818,7 +2849,7 @@ dependencies = [
"quinn-udp",
"rustc-hash",
"rustls",
"socket2 0.5.10",
"socket2 0.6.3",
"thiserror 2.0.18",
"tokio",
"tracing",
@@ -2856,7 +2887,7 @@ dependencies = [
"cfg_aliases",
"libc",
"once_cell",
"socket2 0.5.10",
"socket2 0.6.3",
"tracing",
"windows-sys 0.60.2",
]
@@ -3132,6 +3163,7 @@ dependencies = [
"rustls-platform-verifier",
"serde",
"serde_json",
"serde_urlencoded",
"sync_wrapper",
"tokio",
"tokio-rustls",
@@ -3872,15 +3904,23 @@ name = "sprout-agent"
version = "0.1.0"
dependencies = [
"arc-swap",
"async-trait",
"axum",
"base64",
"getrandom 0.2.17",
"hex",
"nix",
"reqwest 0.13.3",
"rmcp",
"serde",
"serde_json",
"sha2 0.11.0",
"tempfile",
"tokio",
"tracing",
"tracing-subscriber",
"urlencoding",
"webbrowser",
]
[[package]]
@@ -5042,6 +5082,12 @@ dependencies = [
"serde_derive",
]
[[package]]
name = "urlencoding"
version = "2.1.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "daf8dba3b7eb870caf1ddeed7bc9d2a049f3cfdfae7cb521b087cc33ae4c49da"
[[package]]
name = "utf8_iter"
version = "1.0.4"
@@ -5255,6 +5301,22 @@ dependencies = [
"wasm-bindgen",
]
[[package]]
name = "webbrowser"
version = "1.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0fc95580916af1e68ff6a7be07446fc5db73ebf71cf092de939bbf5f7e189f72"
dependencies = [
"core-foundation",
"jni",
"log",
"ndk-context",
"objc2",
"objc2-foundation",
"url",
"web-sys",
]
[[package]]
name = "webpki-root-certs"
version = "1.0.7"
@@ -5320,7 +5382,7 @@ version = "0.1.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22"
dependencies = [
"windows-sys 0.48.0",
"windows-sys 0.61.2",
]
[[package]]
crates/sprout-agent/Cargo.toml
+15 -2
View File
@@ -25,15 +25,23 @@ name = "fake-mcp"
path = "tests/bin/fake_mcp.rs"
[dependencies]
tokio = { workspace = true, features = ["rt-multi-thread", "macros", "io-std", "io-util", "sync", "process", "time"] }
tokio = { workspace = true, features = ["rt-multi-thread", "macros", "io-std", "io-util", "sync", "process", "time", "net"] }
serde = { workspace = true }
serde_json = { workspace = true }
reqwest = { workspace = true }
reqwest = { workspace = true, features = ["json", "rustls", "form"] }
rmcp = { version = "1", default-features = false, features = ["client", "transport-child-process"] }
arc-swap = "1"
getrandom = "0.2"
tracing = { workspace = true }
tracing-subscriber = { workspace = true }
# OAuth 2.0 PKCE for Databricks (and future browser-auth providers).
async-trait = "0.1"
axum = { workspace = true }
base64 = "0.22"
hex = { workspace = true }
sha2 = { workspace = true }
urlencoding = "2"
webbrowser = "1"
[target.'cfg(unix)'.dependencies]
nix = { version = "0.31", default-features = false, features = ["signal", "process"] }
@@ -41,3 +49,8 @@ nix = { version = "0.31", default-features = false, features = ["signal", "proce
[dev-dependencies]
tokio = { workspace = true, features = ["test-util", "rt-multi-thread", "macros", "io-std", "io-util", "sync", "process", "time", "net"] }
nix = { version = "0.31", default-features = false, features = ["signal", "process"] }
axum = { workspace = true }
hex = { workspace = true }
serde = { workspace = true }
sha2 = { workspace = true }
tempfile = "3"
crates/sprout-agent/src/auth.rs
+552
View File
@@ -0,0 +1,552 @@
//! Token sources for the LLM transport layer.
//!
//! [`TokenSource`] decouples request auth from `Config::api_key`: providers
//! can supply a static string ([`StaticTokenSource`]) or a refreshable OAuth
//! 2.0 PKCE engine ([`PkceOAuthTokenSource`]). Engines own their own cache
//! and refresh logic; the [`Llm`] just asks for a bearer per request.
//!
//! The PKCE engine implements RFC 6749 + RFC 7636 with on-disk token
//! caching keyed by `sha256(discovery_url|client_id|scopes)`. It's the
//! same shape goose uses for Databricks, but we own the wire format and
//! cache directory so the two are independently upgradable.
//!
//! First-use (cache empty) requires a browser: the engine opens
//! `authorization_endpoint` in `webbrowser`, listens on `127.0.0.1:0`,
//! captures the redirect, and exchanges the code for a token. Subsequent
//! calls hit the cache and silently refresh when expired.
use std::fs;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use async_trait::async_trait;
use base64::Engine;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use sha2::Digest;
use tokio::sync::Mutex;
use crate::types::AgentError;
/// Buffer before `expires_at` to consider a cached token "still good".
/// Keeps us off the cliff if the clock or the server's clock drifts.
const TOKEN_REFRESH_LEEWAY: Duration = Duration::from_secs(60);
/// Wall-clock budget for the interactive browser dance. Goose uses 60s.
/// We match: any longer and the user has gone to lunch.
const BROWSER_AUTH_TIMEOUT: Duration = Duration::from_secs(60);
/// Asynchronous source of a bearer token. The [`Llm`] calls this per
/// request, so impls are expected to be cheap on the cache-hit path.
#[async_trait]
pub trait TokenSource: Send + Sync {
async fn bearer(&self) -> Result<String, AgentError>;
}
/// A token that never changes for the life of the process.
pub struct StaticTokenSource(String);
impl StaticTokenSource {
pub fn new(token: impl Into<String>) -> Self {
Self(token.into())
}
}
#[async_trait]
impl TokenSource for StaticTokenSource {
async fn bearer(&self) -> Result<String, AgentError> {
Ok(self.0.clone())
}
}
/// Static config for an OAuth 2.0 Authorization Code + PKCE provider.
///
/// The `discovery_url` must return a JSON document with at least
/// `authorization_endpoint` and `token_endpoint` (RFC 8414). The
/// `cache_namespace` is the directory under `~/.config/sprout-agent/oauth/`
/// the token JSON lives in — separates providers' caches cleanly.
#[derive(Debug, Clone)]
pub struct PkceOAuthConfig {
pub discovery_url: String,
pub client_id: String,
pub scopes: Vec<String>,
pub cache_namespace: String,
/// When `Some`, the engine writes tokens here instead of
/// `~/.config/sprout-agent/oauth/<cache_namespace>/`. Production code
/// leaves this `None`. Integration tests use it to avoid stomping on
/// a shared `$HOME` when running in parallel.
pub cache_dir_override: Option<PathBuf>,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
struct CachedToken {
access_token: String,
refresh_token: Option<String>,
/// Unix seconds. `None` means the server didn't advertise an expiry;
/// we use it without checking and rely on refresh on 401.
expires_at: Option<u64>,
}
#[derive(Debug, Clone)]
struct OidcEndpoints {
authorization_endpoint: String,
token_endpoint: String,
}
/// PKCE OAuth token source with on-disk refresh cache.
///
/// First call:
/// 1. Loads from cache if present and unexpired.
/// 2. Otherwise tries `refresh_token` if cached.
/// 3. Otherwise runs the full browser flow.
///
/// Subsequent calls hit an in-memory copy of the cached token and only
/// touch disk/network if the access token is past `expires_at`.
pub struct PkceOAuthTokenSource {
cfg: PkceOAuthConfig,
http: Client,
cache_path: PathBuf,
/// Single-flight guard: only one refresh/browser flow at a time, even
/// if many tool calls land concurrently.
state: Mutex<Option<CachedToken>>,
}
impl PkceOAuthTokenSource {
pub fn new(cfg: PkceOAuthConfig) -> Result<Arc<Self>, AgentError> {
let cache_path = cache_path_for(&cfg)?;
if let Some(parent) = cache_path.parent() {
fs::create_dir_all(parent)
.map_err(|e| AgentError::Llm(format!("oauth cache dir {parent:?}: {e}")))?;
}
let initial = read_cache(&cache_path);
Ok(Arc::new(Self {
cfg,
http: Client::new(),
cache_path,
state: Mutex::new(initial),
}))
}
/// Discover authorization + token endpoints from the well-known URL.
async fn endpoints(&self) -> Result<OidcEndpoints, AgentError> {
let v: Value = self
.http
.get(&self.cfg.discovery_url)
.send()
.await
.map_err(|e| AgentError::Llm(format!("oauth discovery: {e}")))?
.error_for_status()
.map_err(|e| AgentError::Llm(format!("oauth discovery status: {e}")))?
.json()
.await
.map_err(|e| AgentError::Llm(format!("oauth discovery json: {e}")))?;
let auth = v
.get("authorization_endpoint")
.and_then(Value::as_str)
.ok_or_else(|| {
AgentError::Llm("oauth discovery: authorization_endpoint missing".into())
})?
.to_string();
let token = v
.get("token_endpoint")
.and_then(Value::as_str)
.ok_or_else(|| AgentError::Llm("oauth discovery: token_endpoint missing".into()))?
.to_string();
Ok(OidcEndpoints {
authorization_endpoint: auth,
token_endpoint: token,
})
}
/// Persist a token to disk and the in-memory cell.
fn save(&self, state: &mut Option<CachedToken>, token: CachedToken) -> Result<(), AgentError> {
let body = serde_json::to_vec_pretty(&token)
.map_err(|e| AgentError::Llm(format!("oauth cache serialize: {e}")))?;
// Atomic rename so a concurrent reader never sees a partial write.
let tmp = self.cache_path.with_extension("json.tmp");
fs::write(&tmp, &body)
.map_err(|e| AgentError::Llm(format!("oauth cache write {tmp:?}: {e}")))?;
fs::rename(&tmp, &self.cache_path)
.map_err(|e| AgentError::Llm(format!("oauth cache rename: {e}")))?;
*state = Some(token);
Ok(())
}
/// Exchange a refresh token for a fresh access token.
async fn refresh(
&self,
endpoints: &OidcEndpoints,
refresh_token: &str,
) -> Result<CachedToken, AgentError> {
let params = [
("grant_type", "refresh_token"),
("refresh_token", refresh_token),
("client_id", &self.cfg.client_id),
];
let resp = self
.http
.post(&endpoints.token_endpoint)
.form(&params)
.send()
.await
.map_err(|e| AgentError::Llm(format!("oauth refresh: {e}")))?;
if !resp.status().is_success() {
let body = resp.text().await.unwrap_or_default();
return Err(AgentError::Llm(format!("oauth refresh failed: {body}")));
}
let v: Value = resp
.json()
.await
.map_err(|e| AgentError::Llm(format!("oauth refresh json: {e}")))?;
token_from_response(&v, Some(refresh_token))
}
/// Run the full browser-mediated Authorization Code + PKCE flow.
/// Caller must hold a TTY/browser: this opens a window and blocks.
pub async fn interactive_login(&self) -> Result<(), AgentError> {
let endpoints = self.endpoints().await?;
let token = browser_pkce_flow(&self.http, &self.cfg, &endpoints).await?;
let mut state = self.state.lock().await;
self.save(&mut state, token)?;
Ok(())
}
}
#[async_trait]
impl TokenSource for PkceOAuthTokenSource {
async fn bearer(&self) -> Result<String, AgentError> {
let mut state = self.state.lock().await;
// 1. Cache hit, still fresh.
if let Some(tok) = state.as_ref() {
if !is_expired(tok) {
return Ok(tok.access_token.clone());
}
}
// 2. Cache hit, expired, but we have a refresh token.
let refresh = state.as_ref().and_then(|t| t.refresh_token.clone());
if let Some(rt) = refresh {
let endpoints = self.endpoints().await?;
match self.refresh(&endpoints, &rt).await {
Ok(fresh) => {
let bearer = fresh.access_token.clone();
self.save(&mut state, fresh)?;
return Ok(bearer);
}
Err(e) => {
tracing::warn!(error = %e, "oauth refresh failed; falling back to browser flow");
}
}
}
// 3. No usable cache: full browser dance.
let endpoints = self.endpoints().await?;
let fresh = browser_pkce_flow(&self.http, &self.cfg, &endpoints).await?;
let bearer = fresh.access_token.clone();
self.save(&mut state, fresh)?;
Ok(bearer)
}
}
// ---- helpers -------------------------------------------------------------
/// Aborts a spawned task when dropped. Used to guarantee the localhost
/// callback server doesn't outlive a failed/abandoned PKCE attempt.
struct AbortOnDrop(tokio::task::JoinHandle<()>);
impl Drop for AbortOnDrop {
fn drop(&mut self) {
self.0.abort();
}
}
fn is_expired(t: &CachedToken) -> bool {
let Some(exp) = t.expires_at else {
return false;
};
let now = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0);
now + TOKEN_REFRESH_LEEWAY.as_secs() >= exp
}
fn cache_path_for(cfg: &PkceOAuthConfig) -> Result<PathBuf, AgentError> {
let mut h = sha2::Sha256::new();
h.update(cfg.discovery_url.as_bytes());
h.update(b"|");
h.update(cfg.client_id.as_bytes());
h.update(b"|");
h.update(cfg.scopes.join(",").as_bytes());
let hash = hex::encode(h.finalize());
let dir = match &cfg.cache_dir_override {
Some(p) => p.join(&cfg.cache_namespace),
None => {
let home = std::env::var("HOME")
.map_err(|_| AgentError::Llm("oauth cache: $HOME not set".into()))?;
PathBuf::from(home)
.join(".config")
.join("sprout-agent")
.join("oauth")
.join(&cfg.cache_namespace)
}
};
Ok(dir.join(format!("{hash}.json")))
}
fn read_cache(path: &PathBuf) -> Option<CachedToken> {
let body = fs::read(path).ok()?;
serde_json::from_slice(&body).ok()
}
/// Parse a token-endpoint JSON response. Fails loudly when `access_token`
/// is missing or empty — without this, a malformed server response would
/// be cached and `bearer()` would silently return `""` until the entry
/// expires or is deleted by hand.
fn token_from_response(
v: &Value,
fallback_refresh: Option<&str>,
) -> Result<CachedToken, AgentError> {
let access_token = v
.get("access_token")
.and_then(Value::as_str)
.filter(|s| !s.is_empty())
.ok_or_else(|| AgentError::Llm("oauth: token response missing/empty access_token".into()))?
.to_string();
let refresh_token = v
.get("refresh_token")
.and_then(Value::as_str)
.map(str::to_string)
.or_else(|| fallback_refresh.map(str::to_string));
let expires_at = v.get("expires_in").and_then(Value::as_u64).map(|secs| {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0)
+ secs
});
Ok(CachedToken {
access_token,
refresh_token,
expires_at,
})
}
/// PKCE pieces: URL-safe random verifier (~64 chars) and its SHA-256
/// challenge (RFC 7636 §4.2).
fn pkce_pair() -> Result<(String, String), AgentError> {
let mut bytes = [0u8; 48];
getrandom::getrandom(&mut bytes).map_err(|e| AgentError::Llm(format!("pkce rng: {e}")))?;
let verifier = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes);
let challenge = base64::engine::general_purpose::URL_SAFE_NO_PAD
.encode(sha2::Sha256::digest(verifier.as_bytes()));
Ok((verifier, challenge))
}
fn random_state() -> Result<String, AgentError> {
let mut bytes = [0u8; 16];
getrandom::getrandom(&mut bytes).map_err(|e| AgentError::Llm(format!("state rng: {e}")))?;
Ok(base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes))
}
/// Spin up a localhost callback server, open the authorize URL in a
/// browser, wait up to [`BROWSER_AUTH_TIMEOUT`] for the redirect, then
/// exchange the code for a token.
async fn browser_pkce_flow(
http: &Client,
cfg: &PkceOAuthConfig,
endpoints: &OidcEndpoints,
) -> Result<CachedToken, AgentError> {
use axum::{extract::Query, response::Html, routing::get, Router};
use std::collections::HashMap;
use std::net::SocketAddr;
use tokio::sync::oneshot;
let (verifier, challenge) = pkce_pair()?;
let state = random_state()?;
let (tx, rx) = oneshot::channel::<Result<String, String>>();
let tx = Arc::new(Mutex::new(Some(tx)));
let expected_state = state.clone();
let app = Router::new().route(
"/",
get(move |Query(params): Query<HashMap<String, String>>| {
let tx = Arc::clone(&tx);
let expected = expected_state.clone();
async move {
let result = match (params.get("code"), params.get("state")) {
(Some(code), Some(st)) if st == &expected => Ok(code.clone()),
(Some(_), Some(_)) => Err("state mismatch".to_string()),
_ => Err(params
.get("error")
.cloned()
.unwrap_or_else(|| "missing code".into())),
};
if let Some(sender) = tx.lock().await.take() {
let _ = sender.send(result.clone());
}
match result {
Ok(_) => Html(
"<h2>Sprout: signed in</h2><p>You can close this window.</p>".to_string(),
),
Err(e) => Html(format!("<h2>Sprout auth failed</h2><pre>{e}</pre>")),
}
}
}),
);
let listener = tokio::net::TcpListener::bind(SocketAddr::from(([127, 0, 0, 1], 0)))
.await
.map_err(|e| AgentError::Llm(format!("oauth callback bind: {e}")))?;
let port = listener
.local_addr()
.map_err(|e| AgentError::Llm(format!("oauth callback addr: {e}")))?
.port();
let redirect_uri = format!("http://localhost:{port}");
// `_server` is held until this function returns; the drop guard aborts
// the axum task on every exit path (timeout, callback error, token
// exchange failure, or success), so we never leak a listener bound to
// 127.0.0.1 past the auth attempt.
let _server = AbortOnDrop(tokio::spawn(async move {
let _ = axum::serve(listener, app).await;
}));
let auth_url = format!(
"{}?response_type=code&client_id={}&redirect_uri={}&scope={}&state={}&code_challenge={}&code_challenge_method=S256",
endpoints.authorization_endpoint,
urlencoding::encode(&cfg.client_id),
urlencoding::encode(&redirect_uri),
urlencoding::encode(&cfg.scopes.join(" ")),
urlencoding::encode(&state),
urlencoding::encode(&challenge),
);
eprintln!("Opening browser for authentication. If it doesn't open, visit:\n {auth_url}");
let _ = webbrowser::open(&auth_url);
let code = tokio::time::timeout(BROWSER_AUTH_TIMEOUT, rx)
.await
.map_err(|_| AgentError::Llm("oauth: browser auth timed out".into()))?
.map_err(|_| AgentError::Llm("oauth: callback sender dropped".into()))?
.map_err(|e| AgentError::Llm(format!("oauth callback: {e}")))?;
// Exchange code for token.
let params = [
("grant_type", "authorization_code"),
("code", &code),
("redirect_uri", &redirect_uri),
("code_verifier", &verifier),
("client_id", &cfg.client_id),
];
let resp = http
.post(&endpoints.token_endpoint)
.form(&params)
.send()
.await
.map_err(|e| AgentError::Llm(format!("oauth exchange: {e}")))?;
if !resp.status().is_success() {
let body = resp.text().await.unwrap_or_default();
return Err(AgentError::Llm(format!("oauth exchange failed: {body}")));
}
let v: Value = resp
.json()
.await
.map_err(|e| AgentError::Llm(format!("oauth exchange json: {e}")))?;
token_from_response(&v, None)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn pkce_pair_produces_valid_challenge() {
let (verifier, challenge) = pkce_pair().unwrap();
assert!(verifier.len() >= 43);
let expected = base64::engine::general_purpose::URL_SAFE_NO_PAD
.encode(sha2::Sha256::digest(verifier.as_bytes()));
assert_eq!(expected, challenge);
}
#[test]
fn cached_token_no_expiry_is_not_expired() {
let t = CachedToken {
access_token: "x".into(),
refresh_token: None,
expires_at: None,
};
assert!(!is_expired(&t));
}
#[test]
fn cached_token_far_future_is_not_expired() {
let future = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs()
+ 3600;
let t = CachedToken {
access_token: "x".into(),
refresh_token: None,
expires_at: Some(future),
};
assert!(!is_expired(&t));
}
#[test]
fn cached_token_within_leeway_is_expired() {
let near = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs()
+ 10; // 10s away, leeway is 60s → counts as expired
let t = CachedToken {
access_token: "x".into(),
refresh_token: None,
expires_at: Some(near),
};
assert!(is_expired(&t));
}
#[test]
fn cache_path_includes_namespace_and_hash() {
// HOME is required; cargo test runs set it.
let cfg = PkceOAuthConfig {
discovery_url: "https://example.com/.well-known".into(),
client_id: "abc".into(),
scopes: vec!["a".into(), "b".into()],
cache_namespace: "demo".into(),
cache_dir_override: None,
};
let p = cache_path_for(&cfg).unwrap();
assert!(p.to_string_lossy().contains("/sprout-agent/oauth/demo/"));
assert!(p.extension().and_then(|s| s.to_str()) == Some("json"));
}
#[test]
fn token_from_response_uses_fallback_refresh() {
let v: Value = serde_json::from_str(r#"{"access_token":"abc","expires_in":3600}"#).unwrap();
let t = token_from_response(&v, Some("old-refresh")).unwrap();
assert_eq!(t.access_token, "abc");
assert_eq!(t.refresh_token.as_deref(), Some("old-refresh"));
assert!(t.expires_at.is_some());
}
#[test]
fn token_from_response_rejects_missing_access_token() {
let v: Value = serde_json::from_str(r#"{"expires_in":3600}"#).unwrap();
assert!(token_from_response(&v, None).is_err());
}
#[test]
fn token_from_response_rejects_empty_access_token() {
let v: Value = serde_json::from_str(r#"{"access_token":""}"#).unwrap();
assert!(token_from_response(&v, None).is_err());
}
}
crates/sprout-agent/src/config.rs
+15
View File
@@ -26,6 +26,10 @@ const DEFAULT_SYSTEM_PROMPT: &str =
pub enum Provider {
Anthropic,
OpenAi,
/// Databricks model serving. Routes to `{base_url}/serving-endpoints/{model}/invocations`
/// with a dynamically-acquired bearer (OAuth 2.0 PKCE, or static `DATABRICKS_TOKEN`).
/// Wire format is OpenAI-chat-compatible — reuses the same body builder and parser.
Databricks,
}
/// Which OpenAI-family HTTP API to call. Set via `OPENAI_COMPAT_API`
@@ -78,10 +82,15 @@ impl Config {
let provider = match req("SPROUT_AGENT_PROVIDER")?.to_ascii_lowercase().as_str() {
"anthropic" => Provider::Anthropic,
"openai" | "openai-compat" => Provider::OpenAi,
"databricks" => Provider::Databricks,
o => return Err(format!("config: SPROUT_AGENT_PROVIDER={o} not supported")),
};
// OPENAI_COMPAT_API is only read when provider=openai, so a stray
// bad value can't break an Anthropic-only deployment.
//
// Databricks borrows api_key as the *optional* `DATABRICKS_TOKEN` escape
// hatch — empty means "use OAuth PKCE." The model lives in the URL path,
// not the request body (see `EndpointStrategy::DatabricksServing`).
let (api_key, model, base_url, openai_api) = match provider {
Provider::Anthropic => (
req("ANTHROPIC_API_KEY")?,
@@ -95,6 +104,12 @@ impl Config {
env_or("OPENAI_COMPAT_BASE_URL", "https://api.openai.com/v1"),
parse_openai_api(env("OPENAI_COMPAT_API").as_deref())?,
),
Provider::Databricks => (
env("DATABRICKS_TOKEN").unwrap_or_default(),
req("DATABRICKS_MODEL")?,
req("DATABRICKS_HOST")?,
OpenAiApi::Chat, // Databricks invocations is chat-shaped
),
};
let system_prompt = match (env("SPROUT_AGENT_SYSTEM_PROMPT"), env("SPROUT_AGENT_SYSTEM_PROMPT_FILE")) {
(Some(_), Some(_)) => return Err(
crates/sprout-agent/src/lib.rs
+40
View File
@@ -1,5 +1,6 @@
#![forbid(unsafe_code)]
mod agent;
pub mod auth;
mod config;
mod handoff;
mod llm;
@@ -48,6 +49,13 @@ fn die(msg: String) -> ! {
}
pub fn run() -> Result<(), Box<dyn std::error::Error>> {
let args: Vec<String> = std::env::args().collect();
if matches!(args.get(1).map(String::as_str), Some("auth")) {
return tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()?
.block_on(auth_subcommand(&args[2..]));
}
tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()?
@@ -55,6 +63,38 @@ pub fn run() -> Result<(), Box<dyn std::error::Error>> {
Ok(())
}
/// `sprout-agent auth <provider>` — run the interactive auth flow for a
/// provider and persist the result, then exit. Today the only provider is
/// `databricks` (OAuth 2.0 PKCE). Reads `DATABRICKS_HOST` from env; needs
/// a browser on the machine.
async fn auth_subcommand(args: &[String]) -> Result<(), Box<dyn std::error::Error>> {
let provider = args.first().map(String::as_str);
match provider {
Some("databricks") => {
let host = std::env::var("DATABRICKS_HOST")
.map_err(|_| "auth databricks: DATABRICKS_HOST required")?;
let pkce = auth::PkceOAuthConfig {
discovery_url: format!(
"{}/oidc/.well-known/oauth-authorization-server",
host.trim_end_matches('/')
),
client_id: "databricks-cli".into(),
scopes: vec!["all-apis".into(), "offline_access".into()],
cache_namespace: "databricks".into(),
cache_dir_override: None,
};
let src = auth::PkceOAuthTokenSource::new(pkce)?;
src.interactive_login().await?;
eprintln!(
"Authenticated. Token cached under ~/.config/sprout-agent/oauth/databricks/."
);
Ok(())
}
Some(other) => Err(format!("auth: unknown provider {other:?}").into()),
None => Err("auth: provider required (try: sprout-agent auth databricks)".into()),
}
}
async fn async_main() {
tracing_subscriber::fmt()
.with_writer(std::io::stderr)
crates/sprout-agent/src/llm.rs
+93 -4
View File
@@ -1,13 +1,21 @@
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use reqwest::Client;
use serde_json::{json, Value};
use crate::auth::{PkceOAuthConfig, PkceOAuthTokenSource, StaticTokenSource, TokenSource};
use crate::config::{is_openai_host, Config, OpenAiApi, Provider};
use crate::types::{
AgentError, HistoryItem, LlmResponse, ProviderStop, ToolCall, ToolDef, ToolResultContent,
};
/// Databricks OAuth client_id — the public Databricks-published CLI client.
/// PKCE-only, no secret. Same identifier goose uses, so a user's browser
/// consent for `databricks-cli` covers sprout-agent too.
const DATABRICKS_CLIENT_ID: &str = "databricks-cli";
const DATABRICKS_OAUTH_SCOPES: &[&str] = &["all-apis", "offline_access"];
const MAX_LLM_RESPONSE_BYTES: usize = 16 * 1024 * 1024;
const MAX_LLM_ERROR_BODY_BYTES: usize = 4 * 1024;
@@ -22,6 +30,12 @@ pub struct Llm {
/// == Auto`. Subsequent OpenAI calls then go straight to Responses
/// for the lifetime of the process.
auto_upgraded: AtomicBool,
/// Bearer-token source for OpenAI-family requests. Static for OpenAI
/// (the `OPENAI_COMPAT_API_KEY` env var) and Databricks-with-token
/// (the `DATABRICKS_TOKEN` env var); a refreshable PKCE engine for
/// Databricks otherwise. Anthropic doesn't use this — it always
/// reads `cfg.api_key` directly because the API expects `x-api-key`.
auth: Arc<dyn TokenSource>,
}
impl Llm {
@@ -31,9 +45,11 @@ impl Llm {
.timeout(cfg.llm_timeout)
.build()
.map_err(|e| AgentError::Llm(format!("http: {e}")))?;
let auth = build_token_source(cfg)?;
Ok(Self {
http,
auto_upgraded: AtomicBool::new(false),
auth,
})
}
@@ -50,7 +66,7 @@ impl Llm {
.await?;
parse_anthropic(v)
}
Provider::OpenAi => {
Provider::OpenAi | Provider::Databricks => {
self.openai_request(cfg, |use_responses| {
if use_responses {
(
@@ -89,7 +105,7 @@ impl Llm {
});
Ok(parse_anthropic(self.post_anthropic(cfg, &body).await?)?.text)
}
Provider::OpenAi => {
Provider::OpenAi | Provider::Databricks => {
let r = self
.openai_request(cfg, |use_responses| {
if use_responses {
@@ -159,14 +175,35 @@ impl Llm {
}
}
/// POST to an OpenAI-family endpoint. For OpenAI-compat this is just
/// `{base_url}{path}` with the body untouched. For Databricks the URL
/// becomes `{base_url}/serving-endpoints/{model}/invocations` and the
/// `model` field is stripped from the body (Databricks rejects it —
/// the endpoint path already names the model).
async fn post_openai(
&self,
cfg: &Config,
path: &str,
body: &Value,
) -> Result<Value, AgentError> {
let url = format!("{}{}", cfg.base_url.trim_end_matches('/'), path);
post(&self.http, &url, body, |r| r.bearer_auth(&cfg.api_key)).await
let bearer = self.auth.bearer().await?;
let (url, body_owned);
let body_ref: &Value = match cfg.provider {
Provider::Databricks => {
url = format!(
"{}/serving-endpoints/{}/invocations",
cfg.base_url.trim_end_matches('/'),
cfg.model
);
body_owned = strip_model(body);
&body_owned
}
_ => {
url = format!("{}{}", cfg.base_url.trim_end_matches('/'), path);
body
}
};
post(&self.http, &url, body_ref, |r| r.bearer_auth(&bearer)).await
}
/// If `err` names `/v1/responses` / "use the Responses API", latch a
@@ -704,6 +741,58 @@ where
Err(AgentError::Llm("exhausted retries".into()))
}
/// Build the `TokenSource` for the configured provider.
///
/// - `Provider::Anthropic`: a static source seeded from `cfg.api_key`. It's
/// never read for Anthropic requests (those go through `post_anthropic` with
/// `x-api-key`), but Llm holds one to keep the field non-`Option`.
/// - `Provider::OpenAi`: a static source over `OPENAI_COMPAT_API_KEY`.
/// - `Provider::Databricks`: if `DATABRICKS_TOKEN` is set, a static source.
/// Otherwise a `PkceOAuthTokenSource` pointed at the workspace's OIDC
/// discovery URL. First request without a cached token triggers a browser
/// flow; subsequent requests use the cache + refresh transparently.
fn build_token_source(cfg: &Config) -> Result<Arc<dyn TokenSource>, AgentError> {
match cfg.provider {
Provider::Anthropic | Provider::OpenAi => {
Ok(Arc::new(StaticTokenSource::new(cfg.api_key.clone())))
}
Provider::Databricks => {
if !cfg.api_key.is_empty() {
return Ok(Arc::new(StaticTokenSource::new(cfg.api_key.clone())));
}
let discovery_url = format!(
"{}/oidc/.well-known/oauth-authorization-server",
cfg.base_url.trim_end_matches('/')
);
let pkce = PkceOAuthConfig {
discovery_url,
client_id: DATABRICKS_CLIENT_ID.into(),
scopes: DATABRICKS_OAUTH_SCOPES
.iter()
.map(|s| (*s).into())
.collect(),
cache_namespace: "databricks".into(),
cache_dir_override: None,
};
Ok(PkceOAuthTokenSource::new(pkce)?)
}
}
}
/// Return a clone of `body` with any top-level `"model"` field removed.
/// Used for Databricks model-serving, which encodes the model in the URL
/// path and rejects the field in the body.
fn strip_model(body: &Value) -> Value {
match body {
Value::Object(map) => {
let mut m = map.clone();
m.remove("model");
Value::Object(m)
}
other => other.clone(),
}
}
#[cfg(test)]
mod tests {
use super::*;
crates/sprout-agent/tests/databricks_oauth.rs
+450
View File
@@ -0,0 +1,450 @@
//! Integration tests for the PKCE OAuth token source.
//!
//! No browser dance — we cover the silent-refresh and cache-hit paths
//! against a stubbed OIDC server (axum). The interactive browser flow is
//! exercised manually via the `sprout-agent auth databricks` subcommand
//! (see `lib.rs::auth_subcommand`).
//!
//! The second test module (further down) is an ACP-level envelope
//! regression: it spawns the real `sprout-agent` binary with
//! `DATABRICKS_TOKEN` set and a stub HTTP server, then asserts the wire
//! shape we send to Databricks.
use std::net::SocketAddr;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use axum::extract::Form;
use axum::{routing::get, routing::post, Json, Router};
use serde::Deserialize;
use serde_json::json;
use sprout_agent::auth::{PkceOAuthConfig, PkceOAuthTokenSource, TokenSource};
use tempfile::TempDir;
#[derive(Deserialize)]
struct TokenForm {
grant_type: String,
#[allow(dead_code)]
refresh_token: Option<String>,
}
/// Boot a stub OIDC server that:
/// - serves discovery at `/.well-known/oauth-authorization-server`
/// - issues a fresh access token for every `refresh_token` request
/// - counts how many refresh hits it gets
async fn spawn_oidc() -> (String, Arc<AtomicU64>) {
let counter = Arc::new(AtomicU64::new(0));
let counter_for_token = counter.clone();
// Bind first so we know our own base URL before building the router.
let listener = tokio::net::TcpListener::bind(SocketAddr::from(([127, 0, 0, 1], 0)))
.await
.unwrap();
let addr = listener.local_addr().unwrap();
let base = format!("http://{addr}");
let base_for_discovery = base.clone();
let app = Router::new()
.route(
"/.well-known/oauth-authorization-server",
get(move || {
let base = base_for_discovery.clone();
async move {
Json(json!({
"authorization_endpoint": format!("{base}/authorize"),
"token_endpoint": format!("{base}/token"),
}))
}
}),
)
.route(
"/token",
post(move |Form(form): Form<TokenForm>| {
let counter = counter_for_token.clone();
async move {
let n = counter.fetch_add(1, Ordering::SeqCst) + 1;
assert_eq!(form.grant_type, "refresh_token");
Json(json!({
"access_token": format!("fresh-token-{n}"),
"refresh_token": "rotated-refresh",
"expires_in": 3600,
}))
}
}),
);
tokio::spawn(async move {
let _ = axum::serve(listener, app).await;
});
(base, counter)
}
/// Cache key construction matches the auth module: sha256(discovery|client|scopes).
fn cache_path_for(cache_dir: &std::path::Path, cfg: &PkceOAuthConfig) -> std::path::PathBuf {
use sha2::Digest;
let mut h = sha2::Sha256::new();
h.update(cfg.discovery_url.as_bytes());
h.update(b"|");
h.update(cfg.client_id.as_bytes());
h.update(b"|");
h.update(cfg.scopes.join(",").as_bytes());
let hash = hex::encode(h.finalize());
cache_dir
.join(&cfg.cache_namespace)
.join(format!("{hash}.json"))
}
/// Write a token file the engine should pick up on construction.
fn seed_cache(path: &std::path::Path, body: serde_json::Value) {
std::fs::create_dir_all(path.parent().unwrap()).unwrap();
std::fs::write(path, serde_json::to_vec(&body).unwrap()).unwrap();
}
#[tokio::test]
async fn cache_hit_short_circuits_network() {
let tmp = TempDir::new().unwrap();
let (base, refresh_counter) = spawn_oidc().await;
let cfg = PkceOAuthConfig {
discovery_url: format!("{base}/.well-known/oauth-authorization-server"),
client_id: "test-client".into(),
scopes: vec!["a".into(), "b".into()],
cache_namespace: "databricks".into(),
cache_dir_override: Some(tmp.path().to_path_buf()),
};
// Seed an unexpired token in the cache.
let future = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs()
+ 3600;
let path = cache_path_for(tmp.path(), &cfg);
seed_cache(
&path,
json!({
"access_token": "cached-token",
"refresh_token": "rt",
"expires_at": future,
}),
);
let src = PkceOAuthTokenSource::new(cfg).unwrap();
let bearer = src.bearer().await.unwrap();
assert_eq!(bearer, "cached-token");
assert_eq!(
refresh_counter.load(Ordering::SeqCst),
0,
"no refresh should fire"
);
}
#[tokio::test]
async fn expired_cache_silently_refreshes() {
let tmp = TempDir::new().unwrap();
let (base, refresh_counter) = spawn_oidc().await;
let cfg = PkceOAuthConfig {
discovery_url: format!("{base}/.well-known/oauth-authorization-server"),
client_id: "test-client".into(),
scopes: vec!["a".into()],
cache_namespace: "databricks".into(),
cache_dir_override: Some(tmp.path().to_path_buf()),
};
// Seed an already-expired token with a refresh_token.
let path = cache_path_for(tmp.path(), &cfg);
seed_cache(
&path,
json!({
"access_token": "stale",
"refresh_token": "valid-refresh",
"expires_at": 1u64, // way in the past
}),
);
let src = PkceOAuthTokenSource::new(cfg).unwrap();
let bearer = src.bearer().await.unwrap();
assert_eq!(bearer, "fresh-token-1");
assert_eq!(refresh_counter.load(Ordering::SeqCst), 1);
// A second call should hit the in-memory cache and skip the network.
let bearer2 = src.bearer().await.unwrap();
assert_eq!(bearer2, "fresh-token-1");
assert_eq!(refresh_counter.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn refreshed_token_is_persisted_to_disk() {
let tmp = TempDir::new().unwrap();
let (base, _) = spawn_oidc().await;
let cfg = PkceOAuthConfig {
discovery_url: format!("{base}/.well-known/oauth-authorization-server"),
client_id: "test-client".into(),
scopes: vec!["a".into()],
cache_namespace: "databricks".into(),
cache_dir_override: Some(tmp.path().to_path_buf()),
};
let path = cache_path_for(tmp.path(), &cfg);
seed_cache(
&path,
json!({
"access_token": "stale",
"refresh_token": "valid-refresh",
"expires_at": 1u64,
}),
);
let src = PkceOAuthTokenSource::new(cfg).unwrap();
let _ = src.bearer().await.unwrap();
let on_disk: serde_json::Value =
serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
assert_eq!(on_disk["access_token"], "fresh-token-1");
assert_eq!(on_disk["refresh_token"], "rotated-refresh");
assert!(on_disk["expires_at"].is_u64());
}
// ────────────────────────────────────────────────────────────────────────────
// ACP-level envelope regression test.
//
// Boots the real sprout-agent binary with `DATABRICKS_TOKEN` set (so the
// OAuth dance is skipped) pointed at a stub HTTP server that captures every
// inbound request. Asserts the wire-level shape Databricks model serving
// requires: path is `/serving-endpoints/<model>/invocations`, Authorization
// is `Bearer <token>`, and the JSON body has *no* top-level `"model"`. This
// locks in the DRY envelope behavior so a refactor of `post_openai` can't
// silently break Databricks.
// ────────────────────────────────────────────────────────────────────────────
use std::collections::VecDeque;
use std::process::Stdio;
use std::time::Duration;
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
use tokio::net::TcpListener;
use tokio::sync::Mutex;
#[derive(Debug)]
struct CapturedRequest {
path: String,
authorization: Option<String>,
body: serde_json::Value,
}
async fn spawn_capturing_server(
responses: Vec<serde_json::Value>,
) -> (String, Arc<Mutex<Vec<CapturedRequest>>>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let url = format!("http://{}", listener.local_addr().unwrap());
let queue = Arc::new(Mutex::new(VecDeque::from(responses)));
let captured = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
let cap_for_task = captured.clone();
tokio::spawn(async move {
loop {
let (mut sock, _) = match listener.accept().await {
Ok(p) => p,
Err(_) => return,
};
let queue = queue.clone();
let captured = cap_for_task.clone();
tokio::spawn(async move {
let mut buf = Vec::new();
let mut tmp = [0u8; 8192];
while !buf.windows(4).any(|w| w == b"\r\n\r\n") {
match sock.read(&mut tmp).await {
Ok(0) | Err(_) => return,
Ok(n) => buf.extend_from_slice(&tmp[..n]),
}
if buf.len() > 4_000_000 {
return;
}
}
let header_end = buf.windows(4).position(|w| w == b"\r\n\r\n").unwrap() + 4;
let header_str = String::from_utf8_lossy(&buf[..header_end]).to_string();
let (request_line, rest) = header_str.split_once('\n').unwrap_or(("", ""));
let path = request_line
.split_whitespace()
.nth(1)
.unwrap_or("")
.to_string();
let mut authorization = None;
let mut body_len = 0usize;
for line in rest.lines() {
// Split case-insensitively on the colon but keep the value's case intact.
let Some((name, value)) = line.split_once(':') else {
continue;
};
let value = value.trim().trim_end_matches('\r').to_string();
match name.trim().to_ascii_lowercase().as_str() {
"authorization" => authorization = Some(value),
"content-length" => body_len = value.parse().unwrap_or(0),
_ => {}
}
}
while buf.len() < header_end + body_len {
match sock.read(&mut tmp).await {
Ok(0) | Err(_) => return,
Ok(n) => buf.extend_from_slice(&tmp[..n]),
}
}
let body: serde_json::Value =
serde_json::from_slice(&buf[header_end..header_end + body_len])
.unwrap_or(json!(null));
captured.lock().await.push(CapturedRequest {
path,
authorization,
body,
});
let body = queue
.lock()
.await
.pop_front()
.unwrap_or_else(|| json!({ "error": "no canned response" }));
let body_s = serde_json::to_string(&body).unwrap();
let resp = format!(
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
body_s.len(),
body_s,
);
let _ = sock.write_all(resp.as_bytes()).await;
let _ = sock.shutdown().await;
});
}
});
(url, captured)
}
struct AgentHarness {
child: tokio::process::Child,
stdin: tokio::process::ChildStdin,
stdout: BufReader<tokio::process::ChildStdout>,
next_id: i64,
}
impl Drop for AgentHarness {
fn drop(&mut self) {
let _ = self.child.start_kill();
}
}
impl AgentHarness {
async fn spawn_databricks(base_url: &str, model: &str) -> Self {
let bin = env!("CARGO_BIN_EXE_sprout-agent");
let mut cmd = tokio::process::Command::new(bin);
cmd.env("SPROUT_AGENT_PROVIDER", "databricks")
.env("DATABRICKS_HOST", base_url)
.env("DATABRICKS_MODEL", model)
.env("DATABRICKS_TOKEN", "test-bearer")
.env("SPROUT_AGENT_LLM_TIMEOUT_SECS", "5")
.env("SPROUT_AGENT_TOOL_TIMEOUT_SECS", "5")
.env("SPROUT_AGENT_MAX_ROUNDS", "2")
.env("SPROUT_AGENT_MCP_INIT_TIMEOUT_SECS", "2")
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::null())
.kill_on_drop(true);
let mut child = cmd.spawn().expect("spawn sprout-agent");
let stdin = child.stdin.take().unwrap();
let stdout = BufReader::new(child.stdout.take().unwrap());
Self {
child,
stdin,
stdout,
next_id: 1,
}
}
async fn send(&mut self, method: &str, params: serde_json::Value) -> i64 {
let id = self.next_id;
self.next_id += 1;
let mut s = serde_json::to_string(
&json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params }),
)
.unwrap();
s.push('\n');
self.stdin.write_all(s.as_bytes()).await.unwrap();
self.stdin.flush().await.unwrap();
id
}
async fn recv_for(&mut self, want_id: i64) -> serde_json::Value {
loop {
let mut line = String::new();
let n = tokio::time::timeout(Duration::from_secs(15), self.stdout.read_line(&mut line))
.await
.expect("recv timeout")
.expect("read line");
assert!(n > 0, "agent EOF");
let v: serde_json::Value = serde_json::from_str(&line).expect("non-JSON line");
if v.get("id") == Some(&json!(want_id)) {
return v;
}
}
}
}
#[tokio::test]
async fn databricks_envelope_routes_through_serving_endpoints_and_strips_model() {
// One canned chat-completions-shaped response → assistant says "ok"
// with end_turn so the agent loop exits cleanly.
let canned = vec![json!({
"id": "x",
"object": "chat.completion",
"choices": [{
"index": 0,
"message": { "role": "assistant", "content": "ok" },
"finish_reason": "stop"
}]
})];
let (base, captured) = spawn_capturing_server(canned).await;
let model = "goose-claude-4-6-sonnet";
let mut h = AgentHarness::spawn_databricks(&base, model).await;
h.send(
"initialize",
json!({ "protocolVersion": 1, "clientCapabilities": {} }),
)
.await;
h.recv_for(1).await;
h.send("session/new", json!({ "cwd": "/tmp", "mcpServers": [] }))
.await;
let r = h.recv_for(2).await;
let sid = r["result"]["sessionId"].as_str().unwrap().to_string();
h.send(
"session/prompt",
json!({ "sessionId": sid, "prompt": [{ "type": "text", "text": "say ok" }] }),
)
.await;
let _ = h.recv_for(3).await;
let reqs = captured.lock().await;
assert_eq!(reqs.len(), 1, "expected exactly one LLM request");
let req = &reqs[0];
assert_eq!(
req.path,
format!("/serving-endpoints/{model}/invocations"),
"Databricks must route to serving-endpoints/{{model}}/invocations"
);
assert_eq!(
req.authorization.as_deref(),
Some("Bearer test-bearer"),
"Authorization must be the static DATABRICKS_TOKEN as a Bearer"
);
assert!(
req.body.get("model").is_none(),
"request body must NOT include `model` (Databricks rejects it): {:?}",
req.body
);
// Sanity: the rest of the chat envelope should still be there.
assert!(
req.body
.get("messages")
.and_then(|v| v.as_array())
.is_some(),
"request body should keep the chat `messages` field"
);
}