diff --git a/crates/buzz-agent/src/catalog.rs b/crates/buzz-agent/src/catalog.rs index 1386f7dcd7..a4f55261ee 100644 --- a/crates/buzz-agent/src/catalog.rs +++ b/crates/buzz-agent/src/catalog.rs @@ -143,6 +143,7 @@ pub async fn discover_databricks_models_with_cache_dir( ) -> Result, AgentError> { let token_source = if matches!(cfg.provider, Provider::Databricks | Provider::DatabricksV2) && cfg.api_key.is_empty() + && std::env::var_os(crate::sandbox_runtime::BROKER_ENV).is_none() { crate::auth::PkceOAuthTokenSource::new(crate::llm::databricks_pkce_config( &cfg.base_url, diff --git a/crates/buzz-agent/src/llm.rs b/crates/buzz-agent/src/llm.rs index 1bf35bdf23..edf1af7d62 100644 --- a/crates/buzz-agent/src/llm.rs +++ b/crates/buzz-agent/src/llm.rs @@ -2096,6 +2096,9 @@ pub(crate) fn databricks_pkce_config( /// discovery URL. First request without a cached token triggers a browser /// flow; subsequent requests use the cache + refresh transparently. pub(crate) fn build_token_source(cfg: &Config) -> Result, AgentError> { + if let Some(source) = crate::sandbox_runtime::token_source(cfg)? { + return Ok(source); + } match cfg.provider { Provider::Anthropic | Provider::OpenAi | Provider::OpenRouter => { Ok(Arc::new(StaticTokenSource::new(cfg.api_key.clone()))) diff --git a/crates/buzz-agent/src/sandbox_runtime.rs b/crates/buzz-agent/src/sandbox_runtime.rs index 712c9e2f50..100276c6fd 100644 --- a/crates/buzz-agent/src/sandbox_runtime.rs +++ b/crates/buzz-agent/src/sandbox_runtime.rs @@ -1,11 +1,50 @@ -//! Local TLS verification for agents started by trusted external launchers. -use crate::AgentError; +//! Trusted-launcher services for a confined native agent. No arbitrary URLs or file operations. +use crate::{ + auth::TokenSource, + config::{Config, Provider}, + AgentError, +}; +use async_trait::async_trait; +use axum::{ + extract::{DefaultBodyLimit, State}, + http::{HeaderMap, StatusCode}, + routing::post, + Json, Router, +}; +use serde::{Deserialize, Serialize}; #[cfg(test)] -use std::{path::Path, sync::Arc}; +use std::path::Path; +use std::{sync::Arc, time::Duration}; +/// Launcher-supplied JSON capability for a fixed loopback credential broker. +pub const BROKER_ENV: &str = "BUZZ_SANDBOX_AUTH_BROKER"; /// Path to the immutable certificate snapshot supplied by the launcher. pub const ROOTS_ENV: &str = "BUZZ_SANDBOX_TLS_ROOTS"; +/// Connection capability passed only to the protected agent, never its tools. +#[derive(Clone, Serialize, Deserialize)] +pub struct BrokerConfig { + /// Loopback TCP port selected by the trusted launcher. + pub port: u16, + /// Random bearer capability authenticating requests to the launcher. + pub secret: String, + /// Exact configured model service URL; prevents cross-provider reuse. + pub host: String, +} +#[derive(Serialize, Deserialize)] +struct Request { + rejected: Option, +} +#[derive(Serialize, Deserialize)] +struct Response { + token: String, +} +#[derive(Clone)] +struct BrokerState { + secret: String, + source: Arc, +} + /// Native trust is captured before confinement. Verification remains enabled, /// using rustls locally instead of macOS trustd (which can bypass network policy). pub fn http_builder() -> Result { @@ -20,10 +59,149 @@ fn builder_with_roots( .map_err(|e| AgentError::Llm(e.to_string())) } +pub(crate) fn token_source(cfg: &Config) -> Result>, AgentError> { + let Some(raw) = std::env::var_os(BROKER_ENV) else { + return Ok(None); + }; + if !matches!(cfg.provider, Provider::Databricks | Provider::DatabricksV2) + || !cfg.api_key.is_empty() + { + return Ok(None); + } + let broker: BrokerConfig = serde_json::from_str(&raw.to_string_lossy()) + .map_err(|_| AgentError::Llm("invalid sandbox auth broker".into()))?; + if broker.host != cfg.base_url || broker.port == 0 || broker.secret.len() < 32 { + return Err(AgentError::Llm( + "sandbox auth broker provider mismatch".into(), + )); + } + let client = reqwest::Client::builder() + .no_proxy() + .timeout(Duration::from_secs(60)) + .build() + .map_err(|_| AgentError::Llm("sandbox auth client unavailable".into()))?; + Ok(Some(Arc::new(BrokerClient { broker, client }))) +} +struct BrokerClient { + broker: BrokerConfig, + client: reqwest::Client, +} +impl BrokerClient { + async fn request(&self, rejected: Option<&str>) -> Result { + let response = self + .client + .post(format!("http://127.0.0.1:{}/token", self.broker.port)) + .bearer_auth(&self.broker.secret) + .json(&Request { + rejected: rejected.map(str::to_owned), + }) + .send() + .await + .map_err(|_| AgentError::Llm("sandbox sign-in service unavailable".into()))?; + if !response.status().is_success() { + return Err(AgentError::LlmAuth( + "Refresh sign-in from the desktop model picker".into(), + )); + } + Ok(response + .json::() + .await + .map_err(|_| AgentError::Llm("invalid sandbox sign-in response".into()))? + .token) + } +} +#[async_trait] +impl TokenSource for BrokerClient { + async fn bearer(&self) -> Result { + self.request(None).await + } + async fn refresh_now(&self, rejected: &str) -> Result { + self.request(Some(rejected)).await + } +} + +/// The launcher fixes the workspace and cache. The child can only request a +/// bearer or refresh a rejected bearer, never choose a destination or open a browser. +pub fn serve( + listener: tokio::net::TcpListener, + host: &str, + secret: String, +) -> Result, AgentError> { + let source = + crate::auth::PkceOAuthTokenSource::new(crate::llm::databricks_pkce_config(host, None))?; + Ok(serve_source(listener, secret, source)) +} +fn serve_source( + listener: tokio::net::TcpListener, + secret: String, + source: Arc, +) -> tokio::task::JoinHandle<()> { + let app = Router::new() + .route("/token", post(token)) + .layer(DefaultBodyLimit::max(16 * 1024)) + .with_state(BrokerState { secret, source }); + tokio::spawn(async move { + if let Err(error) = axum::serve(listener, app).await { + tracing::warn!(%error, "sandbox auth broker stopped"); + } + }) +} +async fn token( + State(state): State, + headers: HeaderMap, + Json(request): Json, +) -> Result, StatusCode> { + if headers.get("authorization").and_then(|v| v.to_str().ok()) + != Some(format!("Bearer {}", state.secret).as_str()) + { + return Err(StatusCode::UNAUTHORIZED); + } + let result = tokio::time::timeout(Duration::from_secs(55), async { + match request.rejected { + Some(value) => state.source.refresh_now(&value).await, + None => state.source.bearer_no_browser().await, + } + }) + .await + .map_err(|_| StatusCode::GATEWAY_TIMEOUT)?; + let token = result.map_err(|_| StatusCode::UNAUTHORIZED)?; + Ok(Json(Response { token })) +} + #[cfg(test)] mod tests { use super::*; - use reqwest::StatusCode; + #[tokio::test] + async fn broker_rejects_unauthorized_callers_and_serves_fixed_source() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let task = serve_source( + listener, + "test-broker-capability".into(), + Arc::new(crate::auth::StaticTokenSource::new("synthetic-bearer")), + ); + let client = reqwest::Client::builder().no_proxy().build().unwrap(); + let url = format!("http://127.0.0.1:{port}/token"); + let denied = client + .post(&url) + .json(&Request { rejected: None }) + .send() + .await + .unwrap(); + assert_eq!(denied.status(), StatusCode::UNAUTHORIZED); + let allowed = client + .post(&url) + .bearer_auth("test-broker-capability") + .json(&Request { rejected: None }) + .send() + .await + .unwrap(); + assert_eq!( + allowed.json::().await.unwrap().token, + "synthetic-bearer" + ); + task.abort(); + } #[test] fn missing_trust_snapshot_fails_closed() { assert!(builder_with_roots(