diff --git a/src/api.rs b/src/api.rs new file mode 100644 index 00000000..ed76793e --- /dev/null +++ b/src/api.rs @@ -0,0 +1,620 @@ +use std::fs; +use std::io::{self, Read, Write}; +use std::path::PathBuf; + +use anyhow::{bail, Context, Result}; +use clap::{Args, Subcommand, ValueEnum}; +use futures_util::StreamExt; +use reqwest::Method; +use serde::Serialize; +use serde_json::Value; + +use crate::args::BaseArgs; +use crate::auth::login; +use crate::http::{ + build_http_client, ApiClient, HttpError, RawRequestBody, ServiceBase, DEFAULT_HTTP_TIMEOUT, +}; + +const OPENAPI_SPEC_URL: &str = + "https://raw.githubusercontent.com/braintrustdata/braintrust-openapi/main/openapi/spec.json"; + +#[derive(Debug, Clone, Args)] +#[command(after_help = "\ +Examples: + bt api get /v1/project + bt api get /v1/project --json + bt api post /v1/project_score --body '{\"name\":\"example\"}' + bt api post /api/project_score/register --base app --body-file payload.json + bt api spec --filter project_score +")] +pub struct ApiArgs { + #[command(subcommand)] + command: ApiCommand, +} + +#[derive(Debug, Clone, Subcommand)] +enum ApiCommand { + /// Send an authenticated GET request + Get(ReadRequestArgs), + /// Send an authenticated POST request + Post(WriteRequestArgs), + /// Send an authenticated PUT request + Put(WriteRequestArgs), + /// Send an authenticated PATCH request + Patch(WriteRequestArgs), + /// Send an authenticated DELETE request + Delete(ReadRequestArgs), + /// Fetch and inspect the Braintrust OpenAPI spec + Spec(SpecArgs), +} + +#[derive(Debug, Clone, Args)] +struct ReadRequestArgs { + #[command(flatten)] + target: RequestTargetArgs, +} + +#[derive(Debug, Clone, Args)] +struct WriteRequestArgs { + #[command(flatten)] + target: RequestTargetArgs, + + #[command(flatten)] + body: RequestBodyArgs, +} + +#[derive(Debug, Clone, Args)] +struct RequestTargetArgs { + /// Relative path to request (for example: /v1/project or /api/project_score/register) + #[arg(value_name = "PATH")] + path: String, + + /// Which Braintrust base URL to target + #[arg(long, value_enum, default_value_t = ApiBase::Auto)] + base: ApiBase, + + /// Extra request header in Name: Value form + #[arg(long = "header", value_parser = parse_header_arg, value_name = "NAME:VALUE")] + headers: Vec, +} + +#[derive(Debug, Clone, Args)] +struct RequestBodyArgs { + /// Inline request body + #[arg(long, conflicts_with = "body_file", value_name = "BODY")] + body: Option, + + /// Read request body from a file, or '-' for stdin + #[arg(long, conflicts_with = "body", value_name = "FILE")] + body_file: Option, + + /// Content-Type to send when a body is present + #[arg(long, value_name = "MIME_TYPE")] + content_type: Option, +} + +#[derive(Debug, Clone, Args)] +struct SpecArgs { + /// Print the raw JSON OpenAPI document + #[arg(long)] + raw: bool, + + /// Filter operations by path, method, summary, or operation id + #[arg(long, value_name = "TEXT")] + filter: Option, + + /// Override the source URL for the spec + #[arg(long, hide = true, value_name = "URL")] + source_url: Option, +} + +#[derive(Debug, Clone, Copy, ValueEnum, Eq, PartialEq)] +enum ApiBase { + Auto, + Api, + App, +} + +#[derive(Debug, Clone, Eq, PartialEq)] +struct HeaderArg { + name: String, + value: String, +} + +pub async fn run(base: BaseArgs, args: ApiArgs) -> Result<()> { + match args.command { + ApiCommand::Spec(args) => run_spec(base.json, args).await, + command => { + let ctx = login(&base).await?; + let client = ApiClient::new(&ctx)?; + run_authenticated_command(&client, command, base.json).await + } + } +} + +async fn run_authenticated_command( + client: &ApiClient, + command: ApiCommand, + json: bool, +) -> Result<()> { + match command { + ApiCommand::Get(args) => run_request(client, Method::GET, args.target, None, json).await, + ApiCommand::Post(args) => { + run_request( + client, + Method::POST, + args.target, + load_request_body(args.body)?, + json, + ) + .await + } + ApiCommand::Put(args) => { + run_request( + client, + Method::PUT, + args.target, + load_request_body(args.body)?, + json, + ) + .await + } + ApiCommand::Patch(args) => { + run_request( + client, + Method::PATCH, + args.target, + load_request_body(args.body)?, + json, + ) + .await + } + ApiCommand::Delete(args) => { + run_request(client, Method::DELETE, args.target, None, json).await + } + ApiCommand::Spec(_) => unreachable!("spec commands are handled before auth is loaded"), + } +} + +async fn run_request( + client: &ApiClient, + method: Method, + args: RequestTargetArgs, + body: Option, + json: bool, +) -> Result<()> { + if is_absolute_url(&args.path) { + bail!("absolute URLs are not supported; pass a relative path such as /v1/project"); + } + + let service = resolve_service_base(&args.path, args.base); + let headers = args + .headers + .into_iter() + .map(|header| (header.name, header.value)) + .collect::>(); + + let response = client + .request_raw(method, service, &args.path, &headers, body) + .await?; + let status = response.status(); + if !status.is_success() { + let bytes = response + .bytes() + .await + .context("failed to read error response body")?; + let body = String::from_utf8_lossy(&bytes).into_owned(); + return Err(HttpError { status, body }.into()); + } + + if json { + let bytes = response + .bytes() + .await + .context("failed to read JSON response body")?; + let value = parse_json_response_body(&bytes)?; + println!("{}", serde_json::to_string(&value)?); + return Ok(()); + } + + let mut stdout = io::stdout().lock(); + let mut stream = response.bytes_stream(); + while let Some(chunk) = stream.next().await { + stdout + .write_all(&chunk.context("failed to read response body")?) + .context("failed to write response body")?; + } + stdout.flush().context("failed to flush response body")?; + Ok(()) +} + +fn parse_json_response_body(bytes: &[u8]) -> Result { + if bytes.is_empty() { + return Ok(Value::Null); + } + serde_json::from_slice(bytes).context("response body is not valid JSON") +} + +async fn run_spec(json: bool, args: SpecArgs) -> Result<()> { + let source_url = args + .source_url + .unwrap_or_else(|| OPENAPI_SPEC_URL.to_string()); + let spec = fetch_openapi_spec(&source_url).await?; + if args.raw { + print!("{spec}"); + return Ok(()); + } + + let operations = filter_operations(&extract_operations(&spec)?, args.filter.as_deref()); + if json { + println!( + "{}", + serde_json::to_string(&SpecOutput { + source_url, + operations, + })? + ); + return Ok(()); + } + + for operation in operations { + println!("{}", format_operation_line(&operation)); + } + Ok(()) +} + +fn load_request_body(args: RequestBodyArgs) -> Result> { + let RequestBodyArgs { + body, + body_file, + content_type, + } = args; + + let body_bytes = if let Some(body) = body { + Some(body.into_bytes()) + } else if let Some(path) = body_file { + Some(read_body_source(&path)?) + } else { + None + }; + + if body_bytes.is_none() { + if content_type.is_some() { + bail!("--content-type requires --body or --body-file"); + } + return Ok(None); + } + + Ok(body_bytes.map(|bytes| RawRequestBody { + bytes, + content_type: Some(content_type.unwrap_or_else(|| "application/json".to_string())), + })) +} + +fn read_body_source(path: &PathBuf) -> Result> { + if path.as_os_str() == "-" { + let mut bytes = Vec::new(); + io::stdin() + .read_to_end(&mut bytes) + .context("failed to read request body from stdin")?; + return Ok(bytes); + } + + fs::read(path).with_context(|| format!("failed to read request body from {}", path.display())) +} + +fn parse_header_arg(value: &str) -> Result { + let Some((name, raw_value)) = value.split_once(':') else { + return Err("header must be in Name: Value form".to_string()); + }; + + let name = name.trim(); + if name.is_empty() { + return Err("header name cannot be empty".to_string()); + } + + Ok(HeaderArg { + name: name.to_string(), + value: raw_value.trim().to_string(), + }) +} + +fn resolve_service_base(path: &str, base: ApiBase) -> ServiceBase { + match base { + ApiBase::Api => ServiceBase::Api, + ApiBase::App => ServiceBase::App, + ApiBase::Auto => { + if is_app_path(path) { + ServiceBase::App + } else { + ServiceBase::Api + } + } + } +} + +fn is_app_path(path: &str) -> bool { + let trimmed = path.trim_start_matches('/'); + trimmed == "api" || trimmed.starts_with("api/") +} + +fn is_absolute_url(path: &str) -> bool { + path.starts_with("http://") || path.starts_with("https://") +} + +#[derive(Debug, Clone, Serialize, Eq, PartialEq)] +struct SpecOperation { + method: String, + path: String, + operation_id: Option, + summary: Option, +} + +#[derive(Debug, Serialize)] +struct SpecOutput { + source_url: String, + operations: Vec, +} + +async fn fetch_openapi_spec(source_url: &str) -> Result { + let client = + build_http_client(DEFAULT_HTTP_TIMEOUT).context("failed to initialize HTTP client")?; + let response = client + .get(source_url) + .send() + .await + .with_context(|| format!("failed to fetch OpenAPI spec from {source_url}"))?; + let status = response.status(); + let body = response + .text() + .await + .context("failed to read OpenAPI spec response body")?; + if !status.is_success() { + bail!("failed to fetch OpenAPI spec ({status}): {body}"); + } + Ok(body) +} + +fn extract_operations(spec_json: &str) -> Result> { + let spec: Value = serde_json::from_str(spec_json).context("failed to parse OpenAPI spec")?; + let paths = spec + .get("paths") + .and_then(Value::as_object) + .ok_or_else(|| anyhow::anyhow!("OpenAPI spec is missing the top-level `paths` object"))?; + + let mut operations = Vec::new(); + for (path, path_item) in paths { + let Some(path_item) = path_item.as_object() else { + continue; + }; + + for method in ["get", "post", "put", "patch", "delete", "options", "head"] { + let Some(operation) = path_item.get(method).and_then(Value::as_object) else { + continue; + }; + operations.push(SpecOperation { + method: method.to_uppercase(), + path: path.clone(), + operation_id: operation + .get("operationId") + .and_then(Value::as_str) + .map(ToOwned::to_owned), + summary: operation + .get("summary") + .and_then(Value::as_str) + .map(ToOwned::to_owned), + }); + } + } + + operations.sort_by(|left, right| { + left.path + .cmp(&right.path) + .then_with(|| left.method.cmp(&right.method)) + }); + Ok(operations) +} + +fn filter_operations(operations: &[SpecOperation], filter: Option<&str>) -> Vec { + let Some(filter) = filter.map(str::trim).filter(|value| !value.is_empty()) else { + return operations.to_vec(); + }; + let needle = filter.to_ascii_lowercase(); + + operations + .iter() + .filter(|operation| operation_matches_filter(operation, &needle)) + .cloned() + .collect() +} + +fn operation_matches_filter(operation: &SpecOperation, needle: &str) -> bool { + [ + operation.method.as_str(), + operation.path.as_str(), + operation.operation_id.as_deref().unwrap_or(""), + operation.summary.as_deref().unwrap_or(""), + ] + .into_iter() + .any(|value| value.to_ascii_lowercase().contains(needle)) +} + +fn format_operation_line(operation: &SpecOperation) -> String { + let mut line = format!("{} {}", operation.method, operation.path); + let detail = operation + .summary + .as_deref() + .or(operation.operation_id.as_deref()) + .unwrap_or(""); + if !detail.is_empty() { + line.push_str(" "); + line.push_str(detail); + } + line +} + +#[cfg(test)] +mod tests { + use super::*; + + const SAMPLE_SPEC: &str = r#"{ + "paths": { + "/v1/project_score": { + "get": { + "summary": "List project scores", + "operationId": "listProjectScores" + }, + "post": { + "summary": "Create a project score", + "operationId": "createProjectScore" + } + }, + "/brainstore/automation/reset-cursors": { + "post": { + "summary": "Reset automation cursors", + "operationId": "resetAutomationCursors" + } + } + } + }"#; + + #[test] + fn parse_header_arg_splits_name_and_value() { + let header = parse_header_arg("X-Test: hello").expect("parse header"); + assert_eq!( + header, + HeaderArg { + name: "X-Test".to_string(), + value: "hello".to_string(), + } + ); + } + + #[test] + fn parse_header_arg_rejects_invalid_input() { + let err = parse_header_arg("invalid").expect_err("expected parse failure"); + assert!(err.contains("Name: Value")); + } + + #[test] + fn resolve_service_base_uses_app_for_next_api_routes() { + assert_eq!( + resolve_service_base("/api/project_score/register", ApiBase::Auto), + ServiceBase::App + ); + assert_eq!( + resolve_service_base("api/foo", ApiBase::Auto), + ServiceBase::App + ); + } + + #[test] + fn resolve_service_base_uses_api_for_non_app_routes() { + assert_eq!( + resolve_service_base("/v1/project_score", ApiBase::Auto), + ServiceBase::Api + ); + assert_eq!( + resolve_service_base("/brainstore/automation/reset-cursors", ApiBase::Auto), + ServiceBase::Api + ); + } + + #[test] + fn load_request_body_defaults_to_json_content_type() { + let body = load_request_body(RequestBodyArgs { + body: Some("{\"ok\":true}".to_string()), + body_file: None, + content_type: None, + }) + .expect("load request body") + .expect("expected request body"); + + assert_eq!(body.content_type.as_deref(), Some("application/json")); + assert_eq!(body.bytes, br#"{"ok":true}"#); + } + + #[test] + fn load_request_body_rejects_content_type_without_body() { + let err = load_request_body(RequestBodyArgs { + body: None, + body_file: None, + content_type: Some("application/json".to_string()), + }) + .expect_err("expected failure"); + + assert!(err.to_string().contains("--content-type requires")); + } + + #[test] + fn parse_json_response_body_uses_null_for_empty_responses() { + assert_eq!( + parse_json_response_body(b"").expect("parse response"), + Value::Null + ); + } + + #[test] + fn parse_json_response_body_rejects_non_json_responses() { + let err = parse_json_response_body(b"plain text").expect_err("expected parse failure"); + assert!(err.to_string().contains("not valid JSON")); + } + + #[test] + fn extract_operations_reads_methods_and_metadata() { + let operations = extract_operations(SAMPLE_SPEC).expect("extract operations"); + assert_eq!( + operations, + vec![ + SpecOperation { + method: "POST".to_string(), + path: "/brainstore/automation/reset-cursors".to_string(), + operation_id: Some("resetAutomationCursors".to_string()), + summary: Some("Reset automation cursors".to_string()), + }, + SpecOperation { + method: "GET".to_string(), + path: "/v1/project_score".to_string(), + operation_id: Some("listProjectScores".to_string()), + summary: Some("List project scores".to_string()), + }, + SpecOperation { + method: "POST".to_string(), + path: "/v1/project_score".to_string(), + operation_id: Some("createProjectScore".to_string()), + summary: Some("Create a project score".to_string()), + }, + ] + ); + } + + #[test] + fn filter_operations_matches_path_method_summary_and_operation_id() { + let operations = extract_operations(SAMPLE_SPEC).expect("extract operations"); + + let by_path = filter_operations(&operations, Some("project_score")); + assert_eq!(by_path.len(), 2); + + let by_method = filter_operations(&operations, Some("post")); + assert_eq!(by_method.len(), 2); + + let by_summary = filter_operations(&operations, Some("reset automation")); + assert_eq!(by_summary.len(), 1); + + let by_operation_id = filter_operations(&operations, Some("createProjectScore")); + assert_eq!(by_operation_id.len(), 1); + assert_eq!(by_operation_id[0].path, "/v1/project_score"); + } + + #[test] + fn format_operation_line_prefers_summary() { + let line = format_operation_line(&SpecOperation { + method: "POST".to_string(), + path: "/v1/project_score".to_string(), + operation_id: Some("createProjectScore".to_string()), + summary: Some("Create a project score".to_string()), + }); + + assert_eq!(line, "POST /v1/project_score Create a project score"); + } +} diff --git a/src/auth.rs b/src/auth.rs index a915a636..38a67eaf 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -15,7 +15,7 @@ use actix_web::{dev::ServerHandle, web, App, HttpResponse, HttpServer}; use anyhow::{bail, Context, Result}; use base64::Engine as _; use braintrust_sdk_rust::{BraintrustClient, LoginState}; -use chrono::{DateTime, Months, Utc}; +use chrono::{DateTime, Duration as ChronoDuration, Months, Utc}; use clap::{Args, Subcommand}; use crossterm::event::{self, Event, KeyCode, KeyEventKind}; use dialoguer::{Confirm, Input, Password}; @@ -37,6 +37,7 @@ use crate::{ const KEYCHAIN_SERVICE: &str = "com.braintrust.bt.cli"; const OAUTH_SCOPE: &str = "mcp"; +const TOKEN_EXCHANGE_GRANT: &str = "urn:ietf:params:oauth:grant-type:token-exchange"; const OAUTH_CALLBACK_TIMEOUT: Duration = Duration::from_secs(300); const OAUTH_REFRESH_SAFETY_WINDOW_SECONDS: u64 = 60; const AI_PROVIDER_KEY_STALENESS_CHECK_INTERVAL_SECONDS: i64 = 24 * 60 * 60; @@ -375,6 +376,25 @@ struct OAuthTokenResponse { expires_in: Option, } +#[derive(Debug, Clone, Deserialize)] +struct OrgScopedTokenResponse { + access_token: String, + token_type: String, + expires_in: u64, +} + +#[derive(Debug, Clone, Serialize)] +struct OrgScopedTokenRequest<'a> { + grant_type: &'a str, + org_id: &'a str, +} + +#[derive(Debug, Clone)] +struct ResolvedBearerToken { + token: String, + expires_at: Option, +} + #[derive(Debug, Clone, Deserialize)] struct OAuthErrorResponse { #[serde(default)] @@ -389,6 +409,7 @@ Examples: bt auth login bt auth profiles bt auth refresh + bt auth token --header bt auth logout --profile work ")] pub struct AuthArgs { @@ -402,6 +423,8 @@ enum AuthCommand { Login(AuthLoginArgs), /// Force-refresh OAuth access token for a profile Refresh, + /// Print a short-lived org-scoped bearer token for the current auth context + Token(AuthTokenArgs), /// List auth profiles and check connection status Profiles(AuthProfilesArgs), /// Log out by removing a saved profile @@ -415,6 +438,13 @@ struct AuthProfilesArgs { profile: Option, } +#[derive(Debug, Clone, Args)] +struct AuthTokenArgs { + /// Print a full Authorization header instead of only the token + #[arg(long)] + header: bool, +} + #[derive(Debug, Clone, Args)] struct AuthLoginArgs { /// Use OAuth login instead of API key login @@ -450,11 +480,150 @@ pub async fn run(base: BaseArgs, args: AuthArgs) -> Result<()> { match args.command { AuthCommand::Login(login_args) => run_login_set(&base, login_args).await, AuthCommand::Refresh => run_login_refresh(&base).await, + AuthCommand::Token(token_args) => run_token(&base, token_args).await, AuthCommand::Profiles(profile_args) => run_profiles(&base, profile_args).await, AuthCommand::Logout(logout_args) => run_login_logout(base, logout_args), } } +#[derive(Debug, Clone, Serialize)] +struct AuthTokenOutput { + token: String, + authorization_header: String, + api_url: Option, + app_url: Option, + org_name: Option, + expires_at: Option, +} + +async fn run_token(base: &BaseArgs, args: AuthTokenArgs) -> Result<()> { + let auth = resolve_auth(base).await?; + let token = resolve_short_lived_bearer_token(&auth).await?; + + if base.json { + println!( + "{}", + serde_json::to_string(&build_auth_token_output(&token, &auth))? + ); + return Ok(()); + } + + if args.header { + println!("{}", format_authorization_header(&token.token)); + } else { + println!("{}", token.token); + } + Ok(()) +} + +fn format_authorization_header(token: &str) -> String { + format!("Authorization: Bearer {token}") +} + +fn build_auth_token_output(token: &ResolvedBearerToken, auth: &ResolvedAuth) -> AuthTokenOutput { + AuthTokenOutput { + token: token.token.clone(), + authorization_header: format_authorization_header(&token.token), + api_url: auth.api_url.clone(), + app_url: auth.app_url.clone(), + org_name: auth.org_name.clone(), + expires_at: token.expires_at.clone(), + } +} + +async fn resolve_api_bearer_token(auth: &ResolvedAuth) -> Result { + let token = auth.api_key.clone().ok_or_else(|| { + anyhow::anyhow!( + "no login credentials found; set BRAINTRUST_API_KEY, pass --api-key, or run `bt auth login`" + ) + })?; + if !auth.is_oauth { + return Ok(ResolvedBearerToken { + token, + expires_at: None, + }); + } + + resolve_short_lived_bearer_token(auth).await +} + +async fn resolve_short_lived_bearer_token(auth: &ResolvedAuth) -> Result { + let token = auth.api_key.clone().ok_or_else(|| { + anyhow::anyhow!( + "no login credentials found; set BRAINTRUST_API_KEY, pass --api-key, or run `bt auth login`" + ) + })?; + let app_url = auth + .app_url + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or(DEFAULT_APP_URL); + let orgs = fetch_login_orgs(&token, app_url).await?; + let requested_org_name = auth + .org_name + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()); + let org = match requested_org_name { + Some(org_name) => find_login_org(&orgs, org_name).ok_or_else(|| { + anyhow::anyhow!("credential does not have access to organization '{org_name}'") + })?, + None if orgs.len() == 1 => &orgs[0], + None => { + bail!("an active org is required to mint a short-lived token; pass --org or select a profile with an org") + } + }; + + exchange_credential_for_org_scoped_token(&token, app_url, &org.id).await +} + +async fn exchange_credential_for_org_scoped_token( + login_token: &str, + app_url: &str, + org_id: &str, +) -> Result { + let token_url = format!("{}/api/oauth/token", app_url.trim_end_matches('/')); + let client = build_http_client(crate::http::DEFAULT_HTTP_TIMEOUT) + .context("failed to initialize HTTP client")?; + let requested_at = Utc::now(); + let response = client + .post(&token_url) + .bearer_auth(login_token) + .json(&OrgScopedTokenRequest { + grant_type: TOKEN_EXCHANGE_GRANT, + org_id, + }) + .send() + .await + .with_context(|| format!("failed to call OAuth token endpoint {token_url}"))?; + + if !response.status().is_success() { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + return Err(crate::http::HttpError { status, body }.into()); + } + + let payload: OrgScopedTokenResponse = response + .json() + .await + .context("failed to parse OAuth token response")?; + if !payload.token_type.eq_ignore_ascii_case("Bearer") { + bail!("OAuth token endpoint returned unsupported token type"); + } + let expires_in = i64::try_from(payload.expires_in) + .context("OAuth token expiration exceeds supported range")?; + let expires_at = requested_at + .checked_add_signed(ChronoDuration::seconds(expires_in)) + .context("OAuth token expiration exceeds supported range")? + .to_rfc3339(); + + Ok(ResolvedBearerToken { + token: payload.access_token, + expires_at: Some(expires_at), + }) +} + pub async fn login_read_only(base: &BaseArgs) -> Result { if !has_cached_project_id(base) { return login(base).await; @@ -473,11 +642,7 @@ pub async fn login_read_only(base: &BaseArgs) -> Result { pub async fn fast_login(base: &BaseArgs) -> Result { maybe_warn_api_key_override(base); let auth = resolve_auth(base).await?; - let api_key = auth.api_key.clone().ok_or_else(|| { - anyhow::anyhow!( - "no login credentials found; set BRAINTRUST_API_KEY, pass --api-key, or run `bt auth login`" - ) - })?; + let api_key = resolve_api_bearer_token(&auth).await?.token; let org_name = auth.org_name.clone().unwrap_or_default(); let api_url = auth .api_url @@ -507,11 +672,7 @@ pub async fn fast_login(base: &BaseArgs) -> Result { pub async fn login(base: &BaseArgs) -> Result { maybe_warn_api_key_override(base); let auth = resolve_auth(base).await?; - let api_key = auth.api_key.clone().ok_or_else(|| { - anyhow::anyhow!( - "no login credentials found; set BRAINTRUST_API_KEY, pass --api-key, or run `bt auth login`" - ) - })?; + let api_key = resolve_api_bearer_token(&auth).await?.token; let mut builder = BraintrustClient::builder() .blocking_login(true) @@ -972,8 +1133,10 @@ pub async fn resolved_auth_env(base: &BaseArgs) -> Result> let auth = resolve_auth(base).await?; let mut envs = Vec::new(); - if let Some(api_key) = auth.api_key { - envs.push(("BRAINTRUST_API_KEY".to_string(), api_key)); + let token = resolve_api_bearer_token(&auth).await?; + envs.push(("BRAINTRUST_API_KEY".to_string(), token.token)); + if let Some(expires_at) = token.expires_at { + envs.push(("BRAINTRUST_TOKEN_EXPIRES_AT".to_string(), expires_at)); } if let Some(api_url) = auth.api_url { envs.push(("BRAINTRUST_API_URL".to_string(), api_url)); @@ -3664,6 +3827,7 @@ fn auth_store_path() -> Result { #[cfg(test)] mod tests { + use actix_web::HttpRequest; use futures_util::lock::Mutex; use tempfile::TempDir; @@ -3673,10 +3837,84 @@ mod tests { ffi::OsString, fs, path::PathBuf, - sync::OnceLock, + sync::{Arc, Mutex as StdMutex, OnceLock}, time::{SystemTime, UNIX_EPOCH}, }; + #[derive(Debug, Default)] + struct MockTokenExchangeState { + login_authorization: StdMutex>, + token_authorization: StdMutex>, + token_request: StdMutex>, + } + + #[derive(Debug, Clone, Deserialize, PartialEq, Eq)] + struct OrgScopedTokenRequestBody { + grant_type: String, + org_id: String, + } + + async fn mock_login_orgs( + state: web::Data>, + request: HttpRequest, + ) -> HttpResponse { + *state + .login_authorization + .lock() + .expect("login_authorization lock") = request + .headers() + .get("authorization") + .and_then(|value| value.to_str().ok()) + .map(str::to_string); + HttpResponse::Ok().json(serde_json::json!({ + "org_info": [{ + "id": "11111111-1111-4111-8111-111111111111", + "name": "Acme", + "api_url": "https://api.example.com" + }] + })) + } + + async fn mock_token_exchange( + state: web::Data>, + request: HttpRequest, + body: web::Json, + ) -> HttpResponse { + *state + .token_authorization + .lock() + .expect("token_authorization lock") = request + .headers() + .get("authorization") + .and_then(|value| value.to_str().ok()) + .map(str::to_string); + *state.token_request.lock().expect("token_request lock") = Some(body.into_inner()); + HttpResponse::Ok().json(serde_json::json!({ + "access_token": "scoped-token", + "token_type": "Bearer", + "expires_in": 3600 + })) + } + + fn start_mock_token_exchange_server() -> (String, Arc) { + let state = Arc::new(MockTokenExchangeState::default()); + let listener = TcpListener::bind(("127.0.0.1", 0)).expect("bind mock server"); + let addr = listener.local_addr().expect("mock server addr"); + let server_state = state.clone(); + let server = HttpServer::new(move || { + App::new() + .app_data(web::Data::new(server_state.clone())) + .route("/api/apikey/login", web::post().to(mock_login_orgs)) + .route("/api/oauth/token", web::post().to(mock_token_exchange)) + }) + .listen(listener) + .expect("listen mock server") + .run(); + tokio::spawn(server); + std::thread::sleep(Duration::from_millis(25)); + (format!("http://{addr}"), state) + } + fn make_base() -> BaseArgs { BaseArgs { json: false, @@ -3714,6 +3952,81 @@ mod tests { .with_timezone(&Utc) } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn resolve_api_bearer_token_uses_oauth_token_exchange_endpoint() { + let (app_url, state) = start_mock_token_exchange_server(); + let started_at = Utc::now(); + let token = resolve_api_bearer_token(&ResolvedAuth { + api_key: Some("oauth-access-token".to_string()), + api_url: Some("https://api.example.com".to_string()), + app_url: Some(app_url), + org_name: Some("acme".to_string()), + is_oauth: true, + }) + .await + .expect("resolve bearer token"); + + assert_eq!(token.token, "scoped-token"); + let expires_at = + DateTime::parse_from_rfc3339(token.expires_at.as_deref().expect("token expiration")) + .expect("valid token expiration") + .with_timezone(&Utc); + assert!(expires_at >= started_at + ChronoDuration::seconds(3599)); + assert!(expires_at <= Utc::now() + ChronoDuration::seconds(3601)); + assert_eq!( + state + .login_authorization + .lock() + .expect("login_authorization lock") + .as_deref(), + Some("Bearer oauth-access-token") + ); + assert_eq!( + state + .token_authorization + .lock() + .expect("token_authorization lock") + .as_deref(), + Some("Bearer oauth-access-token") + ); + assert_eq!( + state + .token_request + .lock() + .expect("token_request lock") + .as_ref(), + Some(&OrgScopedTokenRequestBody { + grant_type: TOKEN_EXCHANGE_GRANT.to_string(), + org_id: "11111111-1111-4111-8111-111111111111".to_string(), + }) + ); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn resolve_short_lived_bearer_token_exchanges_api_keys() { + let (app_url, state) = start_mock_token_exchange_server(); + let token = resolve_short_lived_bearer_token(&ResolvedAuth { + api_key: Some("permanent-api-key".to_string()), + api_url: Some("https://api.example.com".to_string()), + app_url: Some(app_url), + org_name: None, + is_oauth: false, + }) + .await + .expect("resolve short-lived bearer token"); + + assert_eq!(token.token, "scoped-token"); + assert!(token.expires_at.is_some()); + assert_eq!( + state + .token_authorization + .lock() + .expect("token_authorization lock") + .as_deref(), + Some("Bearer permanent-api-key") + ); + } + fn ai_provider_secret( id: Option<&str>, name: &str, diff --git a/src/http.rs b/src/http.rs index d5a500d0..34d31a36 100644 --- a/src/http.rs +++ b/src/http.rs @@ -12,6 +12,18 @@ use crate::auth::LoginContext; pub const DEFAULT_HTTP_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); pub const BT_USER_AGENT: &str = concat!("bt-cli/", env!("CARGO_PKG_VERSION")); +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ServiceBase { + Api, + App, +} + +#[derive(Debug, Clone)] +pub struct RawRequestBody { + pub bytes: Vec, + pub content_type: Option, +} + pub fn build_http_client(timeout: std::time::Duration) -> Result { build_http_client_from_builder(Client::builder().timeout(timeout)) } @@ -30,7 +42,8 @@ pub fn build_http_client_from_builder(mut builder: ClientBuilder) -> Result String { + self.url_for_service(ServiceBase::Api, path) + } + + pub fn url_for_service(&self, service: ServiceBase, path: &str) -> String { let path = path.trim_start_matches('/'); - format!("{}/{}", self.base_url, path) + let base_url = match service { + ServiceBase::Api => &self.api_url, + ServiceBase::App => &self.app_url, + }; + format!("{}/{}", base_url, path) } pub fn api_key(&self) -> &str { @@ -175,7 +197,7 @@ impl ApiClient { } pub fn base_url(&self) -> &str { - &self.base_url + &self.api_url } pub fn org_id(&self) -> &str { @@ -186,6 +208,30 @@ impl ApiClient { &self.org_name } + pub async fn request_raw( + &self, + method: reqwest::Method, + service: ServiceBase, + path: &str, + headers: &[(String, String)], + body: Option, + ) -> Result { + let url = self.url_for_service(service, path); + let mut request = self.http.request(method, &url).bearer_auth(&self.api_key); + + for (key, value) in headers { + request = request.header(key, value); + } + if let Some(body) = body { + if let Some(content_type) = body.content_type { + request = request.header(CONTENT_TYPE, content_type); + } + request = request.body(body.bytes); + } + + request.send().await.context("request failed") + } + pub async fn get(&self, path: &str) -> Result { self.get_with_headers(path, &[]).await } diff --git a/src/main.rs b/src/main.rs index 5f12e60f..b7746668 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,6 +2,7 @@ use anyhow::{Context, Result}; use clap::{parser::ValueSource, ArgMatches, CommandFactory, FromArgMatches, Parser, Subcommand}; use std::ffi::{OsStr, OsString}; +mod api; mod args; mod auth; #[allow(dead_code)] @@ -72,6 +73,7 @@ Projects & resources experiments Manage experiments Data & evaluation + api Send authenticated HTTP requests datasets Manage datasets eval Run eval files sql Run SQL queries against Braintrust @@ -127,6 +129,8 @@ enum Commands { Setup(CLIArgs), /// Manage workflow docs for coding agents Docs(CLIArgs), + /// Send authenticated HTTP requests + Api(CLIArgs), /// Run SQL queries against Braintrust Sql(CLIArgs), /// Authenticate bt with Braintrust @@ -175,6 +179,7 @@ impl Commands { Commands::Init(cmd) => &cmd.base, Commands::Setup(cmd) => &cmd.base, Commands::Docs(cmd) => &cmd.base, + Commands::Api(cmd) => &cmd.base, Commands::Sql(cmd) => &cmd.base, Commands::Auth(cmd) => &cmd.base, Commands::View(cmd) => &cmd.base, @@ -202,6 +207,7 @@ impl Commands { Commands::Init(cmd) => &mut cmd.base, Commands::Setup(cmd) => &mut cmd.base, Commands::Docs(cmd) => &mut cmd.base, + Commands::Api(cmd) => &mut cmd.base, Commands::Sql(cmd) => &mut cmd.base, Commands::Auth(cmd) => &mut cmd.base, Commands::View(cmd) => &mut cmd.base, @@ -310,6 +316,7 @@ fn try_main() -> Result<()> { Commands::Auth(cmd) => auth::run(cmd.base, cmd.args).await?, Commands::View(cmd) => traces::run(cmd.base, cmd.args).await?, Commands::Init(cmd) => init::run(cmd.base, cmd.args).await?, + Commands::Api(cmd) => api::run(cmd.base, cmd.args).await?, Commands::Sql(cmd) => sql::run(cmd.base, cmd.args).await?, Commands::Setup(cmd) => setup::run_setup_top(cmd.base, cmd.args).await?, Commands::Docs(cmd) => setup::run_docs_top(cmd.base, cmd.args).await?,