diff --git a/src/api/btql.rs b/src/api/btql.rs index 9bee687..f4075ab 100644 --- a/src/api/btql.rs +++ b/src/api/btql.rs @@ -1,5 +1,5 @@ //! BTQL dataset readback, independent of SDK ingestion. -use super::{Client, Error, Plane}; +use super::{Client, Error, Plane, Transport}; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; @@ -63,8 +63,23 @@ pub struct Btql { pub(crate) client: Client, } impl Btql { + /// Read one dataset page without pagination loops or visibility polling. + /// Submission is immediate; dropping the token does not cancel accepted work. + pub fn query_dataset_records( + &self, + query: BTQLQuery, + ) -> crate::Completion { + self.client + .submit(move |transport| async move { transport.query_dataset_records(query).await }) + } +} + +impl Transport { /// Execute one dataset page request; no visibility polling or pagination loop. - pub async fn query_dataset_records(&self, query: BTQLQuery) -> Result { - self.client.post_json(Plane::Data, "btql", &query).await + pub(crate) async fn query_dataset_records( + &self, + query: BTQLQuery, + ) -> Result { + self.post_json(Plane::Data, "btql", &query).await } } diff --git a/src/api/client.rs b/src/api/client.rs index aecb013..fe95709 100644 --- a/src/api/client.rs +++ b/src/api/client.rs @@ -1,88 +1,48 @@ -use std::sync::Arc; - -use reqwest::{ - header::{HeaderMap, HeaderValue, AUTHORIZATION}, - RequestBuilder, Response, +use std::{ + future::Future, + sync::{Arc, Mutex}, + time::Duration, }; -use serde::{de::DeserializeOwned, Serialize}; -use url::Url; - -use super::{Config, Error}; -use crate::http::{build_http_client_builder, CaBundleConfig}; -#[derive(Clone, Copy)] -pub(crate) enum Plane { - Control, - Data, -} +use super::{Config, Error, Transport}; +use crate::internal::{ + execution::ExecutionBinding, + shared_outcome::{SharedOutcome, SharedResponder}, +}; +use crate::{Completion, ExecutionError, ShutdownError}; -/// Immutable HTTP client shared by lightweight endpoint views. +/// Eager one-attempt API operations sharing one execution binding and lifecycle. +/// Clones and endpoint views share admission and shutdown state. Construction +/// validates local configuration and starts the selected binding without login. #[derive(Clone)] pub struct Client { - inner: Arc, -} -struct Inner { - http: reqwest::Client, - api_url: Url, - app_url: Url, - headers: HeaderMap, + transport: Transport, + execution: ExecutionBinding, + shutdown: Arc>>>, } impl Client { pub fn new(config: Config) -> Result { - let api_url = base_url(&config.api_url)?; - let app_url = base_url(&config.app_url)?; - if config.api_key.trim().is_empty() { - return Err(Error::Config("api_key must not be empty".into())); - } - let headers = scoped_headers(&config.api_key, config.org_name.as_deref())?; - let ca = match (config.custom_ca_bundle_path, config.custom_ca_bundle_pem) { - (Some(_), Some(_)) => { - return Err(Error::Config( - "choose a CA file or PEM data, not both".into(), - )) - } - (Some(path), None) => Some(CaBundleConfig::File(path)), - (None, Some(pem)) => Some(CaBundleConfig::from_pem(pem, "API configuration")), - (None, None) => CaBundleConfig::from_environment(), - }; - let http = - build_http_client_builder(config.timeout, ca.as_ref(), config.user_agent.as_deref()) - .map_err(|err| Error::Config(err.to_string()))? - // Redirects would be additional HTTP requests and can change routing. - .redirect(reqwest::redirect::Policy::none()) - .build() - .map_err(|_| Error::Config("failed to build HTTP client".into()))?; + let runtime = config.runtime.clone(); + let transport = Transport::new(config)?; + let execution = ExecutionBinding::new(runtime).map_err(|error| match error { + crate::BraintrustError::Execution(error) => Error::Execution(error), + error => Error::Config(error.to_string()), + })?; Ok(Self { - inner: Arc::new(Inner { - http, - api_url, - app_url, - headers, - }), + transport, + execution, + shutdown: Arc::new(Mutex::new(None)), }) } - /// Reuse the SDK's configured transport without copying session state into - /// the API layer. Credentials and endpoint routes stay immutable. - pub(crate) fn for_scope( - http: reqwest::Client, - api_url: &Url, - app_url: &Url, - token: &str, - org_name: Option<&str>, - ) -> Result { - let api_url = base_url(api_url.as_str())?; - let app_url = base_url(app_url.as_str())?; - let headers = scoped_headers(token, org_name)?; - Ok(Self { - inner: Arc::new(Inner { - http, - app_url, - api_url, - headers, - }), - }) + pub(crate) fn submit(&self, make: F) -> Completion + where + T: Send + 'static, + F: FnOnce(Transport) -> W, + W: Future> + Send + 'static, + { + self.execution + .response_try(|| Ok(make(self.transport.clone()))) } - pub fn logs3(&self) -> super::logs3::Logs3 { super::logs3::Logs3 { client: self.clone(), @@ -108,160 +68,43 @@ impl Client { client: self.clone(), } } - /// One login request. Organization selection and client reconfiguration are explicit. - pub async fn login(&self) -> Result { - let response = self - .send( - self.request(Plane::Control, "api/apikey/login", true)? - .header("content-type", "application/json"), - "api/apikey/login", - ) - .await?; - decode(response, "api/apikey/login").await - } - fn request( - &self, - plane: Plane, - path: &'static str, - post: bool, - ) -> Result { - let base = if matches!(plane, Plane::Control) { - &self.inner.app_url - } else { - &self.inner.api_url - }; - let url = base - .join(path) - .map_err(|_| Error::Config("invalid endpoint path".into()))?; - let request = if post { - self.inner.http.post(url) - } else { - self.inner.http.get(url) - }; - Ok(request.headers(self.inner.headers.clone())) - } - async fn send( - &self, - request: RequestBuilder, - operation: &'static str, - ) -> Result { - let response = request - .send() - .await - .map_err(|source| transport(operation, source))?; - if response.status().is_success() { - return Ok(response); + /// Submit one login request without changing this client's immutable scope. + pub fn login(&self) -> Completion { + self.submit(|transport| async move { transport.login().await }) + } + /// Close admission immediately and wait for accepted requests under one + /// deadline. Repeated calls observe the first deadline and outcome. Request + /// failures remain on their own tokens and are not a shutdown failure log. + pub fn shutdown(&self, timeout: Duration) -> Completion<(), ShutdownError> { + let mut state = self + .shutdown + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(outcome) = &*state { + return outcome.observe(); } - let status = response.status().as_u16(); - let body = response - .text() - .await - .map_err(|source| transport(operation, source))?; - Err(Error::Http { - operation, - status, - body, - }) - } - pub(crate) async fn get_json( - &self, - plane: Plane, - path: &'static str, - query: &[(&str, &str)], - ) -> Result { - let response = self - .send(self.request(plane, path, false)?.query(query), path) - .await?; - decode(response, path).await + let outcome = SharedOutcome::default(); + let (sender, completion) = Completion::pair(); + let responder = SharedResponder::new(sender, Some(outcome.clone())); + *state = Some(outcome); + self.execution.close_responses(); + let execution = self.execution.clone(); + let deadline = tokio::time::Instant::now().checked_add(timeout); + self.execution.spawn(async move { + let work = execution.responses_drained(); + let result = match deadline { + Some(deadline) => tokio::time::timeout_at(deadline, work).await, + None => Ok(work.await), + }; + let result = match result { + Ok(result) => result.map_err(ShutdownError::Execution), + Err(_) => { + execution.cancel_responses(ExecutionError::TimedOut); + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + } + }; + let _ = responder.send(result); + }); + completion } - pub(crate) async fn post_json( - &self, - plane: Plane, - path: &'static str, - body: &Q, - ) -> Result { - let response = self - .send(self.request(plane, path, true)?.json(body), path) - .await?; - decode(response, path).await - } - pub(crate) async fn post_empty( - &self, - plane: Plane, - path: &'static str, - body: &Q, - ) -> Result<(), Error> { - self.send(self.request(plane, path, true)?.json(body), path) - .await?; - Ok(()) - } - pub(crate) async fn post_bytes( - &self, - path: &'static str, - payload: Vec, - ) -> Result<(), Error> { - self.send( - self.request(Plane::Data, path, true)? - .header("content-type", "application/json") - .body(payload), - path, - ) - .await?; - Ok(()) - } -} -fn scoped_headers(token: &str, org_name: Option<&str>) -> Result { - let mut headers = HeaderMap::new(); - let mut token = HeaderValue::from_str(&format!("Bearer {token}")) - .map_err(|_| Error::Config("invalid api_key header".into()))?; - token.set_sensitive(true); - headers.insert(AUTHORIZATION, token); - if let Some(org) = org_name { - headers.insert( - "x-bt-org-name", - HeaderValue::from_str(org) - .map_err(|_| Error::Config("invalid organization header".into()))?, - ); - } - Ok(headers) -} - -fn base_url(value: &str) -> Result { - let mut url = Url::parse(value).map_err(|_| Error::Config("invalid endpoint URL".into()))?; - if !matches!(url.scheme(), "http" | "https") - || url.host_str().is_none() - || !url.username().is_empty() - || url.password().is_some() - || url.query().is_some() - || url.fragment().is_some() - { - return Err(Error::Config( - "endpoint must be an HTTP(S) base URL without credentials, query or fragment".into(), - )); - } - if !url.path().ends_with('/') { - url.set_path(&format!("{}/", url.path())); - } - Ok(url) -} -fn transport(operation: &'static str, source: reqwest::Error) -> Error { - Error::Transport { - operation, - source: source.without_url(), - } -} -async fn decode( - response: Response, - operation: &'static str, -) -> Result { - let status = response.status().as_u16(); - let body = response - .bytes() - .await - .map_err(|source| transport(operation, source))?; - serde_json::from_slice(&body).map_err(|source| Error::Decode { - operation, - status, - source, - }) } diff --git a/src/api/config.rs b/src/api/config.rs index a582b75..e63f8d0 100644 --- a/src/api/config.rs +++ b/src/api/config.rs @@ -11,6 +11,10 @@ pub struct Config { pub app_url: String, pub org_name: Option, pub timeout: Duration, + /// Explicit execution context. None owns a dedicated runtime; ambient Tokio + /// contexts are never selected automatically. A supplied runtime must remain + /// alive and driven until its accepted operations and shutdown complete. + pub runtime: Option, pub user_agent: Option, /// An explicit PEM file; mutually exclusive with `custom_ca_bundle_pem`. pub custom_ca_bundle_path: Option, @@ -25,6 +29,7 @@ impl Config { app_url: DEFAULT_APP_URL.into(), org_name: None, timeout: Duration::from_secs(30), + runtime: None, user_agent: None, custom_ca_bundle_path: None, custom_ca_bundle_pem: None, diff --git a/src/api/datasets.rs b/src/api/datasets.rs index 14185d3..10c58d7 100644 --- a/src/api/datasets.rs +++ b/src/api/datasets.rs @@ -1,5 +1,5 @@ //! Dataset registration, row mutation, and summary endpoint contracts. -use super::{Client, Error, Plane}; +use super::{Client, Error, Plane, Transport}; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; @@ -129,42 +129,70 @@ pub struct Datasets { pub(crate) client: Client, } impl Datasets { + /// Register one dataset on the control plane. + /// Submission is immediate; dropping the token does not cancel accepted work. + pub fn register( + &self, + request: DatasetRegisterRequest, + ) -> crate::Completion { + self.client + .submit(move |transport| async move { transport.register_dataset(request).await }) + } + /// Submit exactly these dataset mutations once, without buffering or splitting. + /// Submission is immediate; dropping the token does not cancel accepted work. + pub fn write(&self, rows: &[DatasetWrite]) -> crate::Completion<(), Error> { + self.client.submit(move |transport| { + let rows = rows.to_vec(); + async move { transport.write_dataset(&rows).await } + }) + } + /// Fetch one dataset summary. + /// Submission is immediate; dropping the token does not cancel accepted work. + pub fn summarize(&self, dataset_id: &str) -> crate::Completion { + self.client.submit(move |transport| { + let dataset_id = dataset_id.to_owned(); + async move { transport.summarize_dataset(&dataset_id).await } + }) + } +} + +impl Transport { /// Register once on the control plane; no resolution cache or retry. - pub async fn register( + pub(crate) async fn register_dataset( &self, request: DatasetRegisterRequest, ) -> Result { - self.client - .post_json(Plane::Control, "api/dataset/register", &request) + self.post_json(Plane::Control, "api/dataset/register", &request) .await } /// Send exactly one typed dataset mutation request on the data plane. /// No retry, buffering, splitting, or overflow upload is performed. - pub async fn write(&self, rows: &[DatasetWrite]) -> Result<(), Error> { + pub(crate) async fn write_dataset(&self, rows: &[DatasetWrite]) -> Result<(), Error> { #[derive(Serialize)] struct Request<'a> { rows: &'a [DatasetWrite], api_version: u8, } - self.client - .post_empty( - Plane::Data, - "logs3", - &Request { - rows, - api_version: crate::types::LOGS_API_VERSION, - }, - ) - .await + self.post_empty( + Plane::Data, + "logs3", + &Request { + rows, + api_version: crate::types::LOGS_API_VERSION, + }, + ) + .await } - pub async fn summarize(&self, dataset_id: &str) -> Result { - self.client - .get_json( - Plane::Data, - "dataset-summary", - &[("dataset_id", dataset_id)], - ) - .await + pub(crate) async fn summarize_dataset( + &self, + dataset_id: &str, + ) -> Result { + self.get_json( + Plane::Data, + "dataset-summary", + &[("dataset_id", dataset_id)], + ) + .await } } diff --git a/src/api/error.rs b/src/api/error.rs index 055faa6..024600e 100644 --- a/src/api/error.rs +++ b/src/api/error.rs @@ -4,6 +4,10 @@ use thiserror::Error; #[derive(Debug, Error)] #[non_exhaustive] pub enum Error { + #[error(transparent)] + Admission(#[from] crate::EnqueueError), + #[error(transparent)] + Execution(#[from] crate::ExecutionError), #[error("invalid API configuration: {0}")] Config(String), #[error("{operation}: transport failed")] diff --git a/src/api/experiments.rs b/src/api/experiments.rs index 0e96bf4..65cf443 100644 --- a/src/api/experiments.rs +++ b/src/api/experiments.rs @@ -1,5 +1,5 @@ //! Existing experiment registration, base and comparison endpoint contracts. -use super::{Client, Error, Plane}; +use super::{Client, Error, Plane, Transport}; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use std::collections::HashMap; @@ -121,24 +121,60 @@ pub struct Experiments { pub(crate) client: Client, } impl Experiments { - pub async fn register( + /// Register one experiment on the control plane. + /// Submission is immediate; dropping the token does not cancel accepted work. + pub fn register( &self, request: ExperimentRegisterRequest, - ) -> Result { + ) -> crate::Completion { + self.client + .submit(move |transport| async move { transport.register_experiment(request).await }) + } + /// Fetch the base experiment, preserving all endpoint errors including HTTP 400. + /// Submission is immediate; dropping the token does not cancel accepted work. + pub fn base( + &self, + request: BaseExperimentRequest, + ) -> crate::Completion { self.client - .post_json(Plane::Control, "api/experiment/register", &request) + .submit(move |transport| async move { transport.base_experiment(request).await }) + } + /// Fetch one experiment comparison. + /// Submission is immediate; dropping the token does not cancel accepted work. + pub fn compare( + &self, + experiment_id: &str, + base_experiment_id: Option<&str>, + ) -> crate::Completion { + self.client.submit(move |transport| { + let experiment_id = experiment_id.to_owned(); + let base_experiment_id = base_experiment_id.map(str::to_owned); + async move { + transport + .compare_experiments(&experiment_id, base_experiment_id.as_deref()) + .await + } + }) + } +} + +impl Transport { + pub(crate) async fn register_experiment( + &self, + request: ExperimentRegisterRequest, + ) -> Result { + self.post_json(Plane::Control, "api/experiment/register", &request) .await } /// Return the actual endpoint result, including a 400 response; SDK interpretation belongs above this layer. - pub async fn base( + pub(crate) async fn base_experiment( &self, request: BaseExperimentRequest, ) -> Result { - self.client - .post_json(Plane::Control, "api/base_experiment/get_id", &request) + self.post_json(Plane::Control, "api/base_experiment/get_id", &request) .await } - pub async fn compare( + pub(crate) async fn compare_experiments( &self, experiment_id: &str, base_experiment_id: Option<&str>, @@ -147,8 +183,7 @@ impl Experiments { if let Some(base) = base_experiment_id { query.push(("base_experiment_id", base)); } - self.client - .get_json(Plane::Control, "experiment-comparison2", &query) + self.get_json(Plane::Control, "experiment-comparison2", &query) .await } } diff --git a/src/api/logs3.rs b/src/api/logs3.rs index b8f56ab..c63ec0e 100644 --- a/src/api/logs3.rs +++ b/src/api/logs3.rs @@ -1,5 +1,5 @@ //! Direct logs3 requests: no batching, splitting, retries or implicit overflow. -use super::{Client, Error, Plane}; +use super::{Client, Error, Plane, Transport}; use serde::{Deserialize, Serialize}; use serde_json::Value; use std::collections::HashMap; @@ -74,25 +74,64 @@ pub struct Logs3 { pub(crate) client: Client, } impl Logs3 { - pub async fn version(&self) -> Result { - self.client.get_json(Plane::Data, "version", &[]).await + /// Fetch ingestion capabilities once. + /// Submission is immediate; dropping the token does not cancel accepted work. + pub fn version(&self) -> crate::Completion { + self.client + .submit(move |transport| async move { transport.logs3_version().await }) } /// Submit exactly this body once, accepting any successful HTTP status. - pub async fn post_raw(&self, payload: Vec) -> Result<(), Error> { - self.client.post_bytes("logs3", payload).await + /// Submission is immediate; dropping the token does not cancel accepted work. + pub fn post_raw(&self, payload: Vec) -> crate::Completion<(), Error> { + self.client + .submit(move |transport| async move { transport.logs3_post_raw(payload).await }) } - pub async fn insert(&self, request: InsertRows) -> Result<(), Error> { - self.client.post_empty(Plane::Data, "logs3", &request).await + /// Submit exactly these rows once, without batching or retries. + /// Submission is immediate; dropping the token does not cancel accepted work. + pub fn insert(&self, request: InsertRows) -> crate::Completion<(), Error> { + self.client + .submit(move |transport| async move { transport.logs3_insert(request).await }) } - pub async fn request_overflow_upload( + /// Request an overflow grant without uploading its object or committing it. + /// Submission is immediate; dropping the token does not cancel accepted work. + pub fn request_overflow_upload( &self, request: Logs3OverflowUploadRequest, - ) -> Result { + ) -> crate::Completion { + self.client.submit(move |transport| async move { + transport.logs3_request_overflow_upload(request).await + }) + } + /// Submit an explicit overflow reference once. + /// Submission is immediate; dropping the token does not cancel accepted work. + pub fn commit_overflow(&self, request: Logs3OverflowRequest) -> crate::Completion<(), Error> { self.client - .post_json(Plane::Data, "logs3/overflow", &request) + .submit(move |transport| async move { transport.logs3_commit_overflow(request).await }) + } +} + +impl Transport { + pub(crate) async fn logs3_version(&self) -> Result { + self.get_json(Plane::Data, "version", &[]).await + } + /// Submit exactly this body once, accepting any successful HTTP status. + pub(crate) async fn logs3_post_raw(&self, payload: Vec) -> Result<(), Error> { + self.post_bytes("logs3", payload).await + } + pub(crate) async fn logs3_insert(&self, request: InsertRows) -> Result<(), Error> { + self.post_empty(Plane::Data, "logs3", &request).await + } + pub(crate) async fn logs3_request_overflow_upload( + &self, + request: Logs3OverflowUploadRequest, + ) -> Result { + self.post_json(Plane::Data, "logs3/overflow", &request) .await } - pub async fn commit_overflow(&self, request: Logs3OverflowRequest) -> Result<(), Error> { - self.client.post_empty(Plane::Data, "logs3", &request).await + pub(crate) async fn logs3_commit_overflow( + &self, + request: Logs3OverflowRequest, + ) -> Result<(), Error> { + self.post_empty(Plane::Data, "logs3", &request).await } } diff --git a/src/api/mod.rs b/src/api/mod.rs index 9c5faae..869c6ab 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -1,10 +1,30 @@ //! Direct, one-attempt Braintrust HTTP endpoints. //! -//! This is the transport foundation of the SDK refactor. It has no background -//! login, queue, retries, batching or SDK resource cache. Construct it with -//! explicit credentials/endpoints; a login response never changes its scope. -//! Endpoint methods are direct async calls in this incremental layer. The eager -//! completion facade is introduced separately, above the same endpoint logic. +//! Calls submit immediately and return [`crate::Completion`] tokens. Dropping +//! a token leaves accepted work running. Configure an explicit runtime or use +//! the default owned runtime, then call [`Client::shutdown`] to close admission +//! and drain accepted requests under a deadline. +//! +//! Explicit credentials and routes are immutable. Each accepted endpoint call +//! makes one HTTP attempt: no implicit login, retries, batching, overflow upload +//! or resource cache. SDK resource policy reuses the private HTTP transport +//! inside its own admitted operations. +//! +//! ```no_run +//! use braintrust_sdk_rust::api::{Client, Config, logs3::InsertRows}; +//! use serde_json::json; +//! use std::time::Duration; +//! +//! # async fn example() -> Result<(), Box> { +//! let client = Client::new(Config::new("token"))?; +//! let completion = client.logs3().insert(InsertRows::new(vec![ +//! json!({"id": "row", "project_id": "project"}), +//! ])); // request admission happens now +//! completion.await?; +//! client.shutdown(Duration::from_secs(10)).await?; +//! # Ok(()) +//! # } +//! ``` pub mod btql; mod client; @@ -17,6 +37,9 @@ pub mod logs3; pub mod projects; pub use client::Client; -use client::Plane; pub use config::{Config, DEFAULT_API_URL, DEFAULT_APP_URL}; pub use error::Error; +use transport::Plane; + +mod transport; +pub(crate) use transport::Transport; diff --git a/src/api/projects.rs b/src/api/projects.rs index 5207e92..59716e3 100644 --- a/src/api/projects.rs +++ b/src/api/projects.rs @@ -1,5 +1,5 @@ //! Existing project registration endpoint. -use super::{Client, Error, Plane}; +use super::{Client, Error, Plane, Transport}; use serde::{Deserialize, Serialize}; #[derive(Debug, Clone, Serialize)] @@ -25,12 +25,23 @@ pub struct Projects { pub(crate) client: Client, } impl Projects { - pub async fn register( + /// Register one project on the control plane. + /// Submission is immediate; dropping the token does not cancel accepted work. + pub fn register( &self, request: ProjectRegisterRequest, - ) -> Result { + ) -> crate::Completion { self.client - .post_json(Plane::Control, "api/project/register", &request) + .submit(move |transport| async move { transport.register_project(request).await }) + } +} + +impl Transport { + pub(crate) async fn register_project( + &self, + request: ProjectRegisterRequest, + ) -> Result { + self.post_json(Plane::Control, "api/project/register", &request) .await } } diff --git a/src/api/transport.rs b/src/api/transport.rs new file mode 100644 index 0000000..7dbee48 --- /dev/null +++ b/src/api/transport.rs @@ -0,0 +1,242 @@ +use std::sync::Arc; + +use reqwest::{ + header::{HeaderMap, HeaderValue, AUTHORIZATION}, + RequestBuilder, Response, +}; +use serde::{de::DeserializeOwned, Serialize}; +use url::Url; + +use super::{Config, Error}; +use crate::http::{build_http_client_builder, CaBundleConfig}; + +#[derive(Clone, Copy)] +pub(crate) enum Plane { + Control, + Data, +} + +/// One-attempt HTTP transport. It owns no scheduler or operation admission. +#[derive(Clone)] +pub(crate) struct Transport { + inner: Arc, +} +struct Inner { + http: reqwest::Client, + api_url: Url, + app_url: Url, + headers: HeaderMap, +} +impl Transport { + pub fn new(config: Config) -> Result { + let api_url = base_url(&config.api_url)?; + let app_url = base_url(&config.app_url)?; + if config.api_key.trim().is_empty() { + return Err(Error::Config("api_key must not be empty".into())); + } + let headers = scoped_headers(&config.api_key, config.org_name.as_deref())?; + let ca = match (config.custom_ca_bundle_path, config.custom_ca_bundle_pem) { + (Some(_), Some(_)) => { + return Err(Error::Config( + "choose a CA file or PEM data, not both".into(), + )) + } + (Some(path), None) => Some(CaBundleConfig::File(path)), + (None, Some(pem)) => Some(CaBundleConfig::from_pem(pem, "API configuration")), + (None, None) => CaBundleConfig::from_environment(), + }; + let http = + build_http_client_builder(config.timeout, ca.as_ref(), config.user_agent.as_deref()) + .map_err(|err| Error::Config(err.to_string()))? + // Redirects would be additional HTTP requests and can change routing. + .redirect(reqwest::redirect::Policy::none()) + .build() + .map_err(|_| Error::Config("failed to build HTTP client".into()))?; + Ok(Self { + inner: Arc::new(Inner { + http, + api_url, + app_url, + headers, + }), + }) + } + /// Reuse the SDK's configured transport without copying session state into + /// the API layer. Credentials and endpoint routes stay immutable. + pub(crate) fn for_scope( + http: reqwest::Client, + api_url: &Url, + app_url: &Url, + token: &str, + org_name: Option<&str>, + ) -> Result { + let api_url = base_url(api_url.as_str())?; + let app_url = base_url(app_url.as_str())?; + let headers = scoped_headers(token, org_name)?; + Ok(Self { + inner: Arc::new(Inner { + http, + app_url, + api_url, + headers, + }), + }) + } + + /// One login request. Organization selection and client reconfiguration are explicit. + pub(crate) async fn login(&self) -> Result { + let response = self + .send( + self.request(Plane::Control, "api/apikey/login", true)? + .header("content-type", "application/json"), + "api/apikey/login", + ) + .await?; + decode(response, "api/apikey/login").await + } + fn request( + &self, + plane: Plane, + path: &'static str, + post: bool, + ) -> Result { + let base = if matches!(plane, Plane::Control) { + &self.inner.app_url + } else { + &self.inner.api_url + }; + let url = base + .join(path) + .map_err(|_| Error::Config("invalid endpoint path".into()))?; + let request = if post { + self.inner.http.post(url) + } else { + self.inner.http.get(url) + }; + Ok(request.headers(self.inner.headers.clone())) + } + async fn send( + &self, + request: RequestBuilder, + operation: &'static str, + ) -> Result { + let response = request + .send() + .await + .map_err(|source| transport(operation, source))?; + if response.status().is_success() { + return Ok(response); + } + let status = response.status().as_u16(); + let body = response + .text() + .await + .map_err(|source| transport(operation, source))?; + Err(Error::Http { + operation, + status, + body, + }) + } + pub(crate) async fn get_json( + &self, + plane: Plane, + path: &'static str, + query: &[(&str, &str)], + ) -> Result { + let response = self + .send(self.request(plane, path, false)?.query(query), path) + .await?; + decode(response, path).await + } + pub(crate) async fn post_json( + &self, + plane: Plane, + path: &'static str, + body: &Q, + ) -> Result { + let response = self + .send(self.request(plane, path, true)?.json(body), path) + .await?; + decode(response, path).await + } + pub(crate) async fn post_empty( + &self, + plane: Plane, + path: &'static str, + body: &Q, + ) -> Result<(), Error> { + self.send(self.request(plane, path, true)?.json(body), path) + .await?; + Ok(()) + } + pub(crate) async fn post_bytes( + &self, + path: &'static str, + payload: Vec, + ) -> Result<(), Error> { + self.send( + self.request(Plane::Data, path, true)? + .header("content-type", "application/json") + .body(payload), + path, + ) + .await?; + Ok(()) + } +} +fn scoped_headers(token: &str, org_name: Option<&str>) -> Result { + let mut headers = HeaderMap::new(); + let mut token = HeaderValue::from_str(&format!("Bearer {token}")) + .map_err(|_| Error::Config("invalid api_key header".into()))?; + token.set_sensitive(true); + headers.insert(AUTHORIZATION, token); + if let Some(org) = org_name { + headers.insert( + "x-bt-org-name", + HeaderValue::from_str(org) + .map_err(|_| Error::Config("invalid organization header".into()))?, + ); + } + Ok(headers) +} + +fn base_url(value: &str) -> Result { + let mut url = Url::parse(value).map_err(|_| Error::Config("invalid endpoint URL".into()))?; + if !matches!(url.scheme(), "http" | "https") + || url.host_str().is_none() + || !url.username().is_empty() + || url.password().is_some() + || url.query().is_some() + || url.fragment().is_some() + { + return Err(Error::Config( + "endpoint must be an HTTP(S) base URL without credentials, query or fragment".into(), + )); + } + if !url.path().ends_with('/') { + url.set_path(&format!("{}/", url.path())); + } + Ok(url) +} +fn transport(operation: &'static str, source: reqwest::Error) -> Error { + Error::Transport { + operation, + source: source.without_url(), + } +} +async fn decode( + response: Response, + operation: &'static str, +) -> Result { + let status = response.status().as_u16(); + let body = response + .bytes() + .await + .map_err(|source| transport(operation, source))?; + serde_json::from_slice(&body).map_err(|source| Error::Decode { + operation, + status, + source, + }) +} diff --git a/src/error.rs b/src/error.rs index c6f1be9..de1d6aa 100644 --- a/src/error.rs +++ b/src/error.rs @@ -69,6 +69,8 @@ impl From for BraintrustError { impl From for BraintrustError { fn from(error: crate::api::Error) -> Self { match error { + crate::api::Error::Admission(error) => Self::Admission(error), + crate::api::Error::Execution(error) => Self::Execution(error), crate::api::Error::Config(message) => BraintrustError::InvalidConfig(message), crate::api::Error::Transport { source, .. } => BraintrustError::Http(source), crate::api::Error::Http { status, body, .. } => BraintrustError::Api { diff --git a/src/internal/ports.rs b/src/internal/ports.rs index ba9da91..2693c46 100644 --- a/src/internal/ports.rs +++ b/src/internal/ports.rs @@ -48,13 +48,10 @@ impl Drop for SpanOperation { } /// Shared delivery lifecycle, independent of resource HTTP operations. -#[async_trait] pub(crate) trait DeliveryControl: Send + Sync { fn flush(&self) -> crate::Completion; - async fn trigger_flush(&self) -> Result<()>; } -#[async_trait] pub(crate) trait SpanSubmitter: DeliveryControl { fn operation(&self) -> std::result::Result; diff --git a/src/internal/resources.rs b/src/internal/resources.rs index 423ec18..65938b8 100644 --- a/src/internal/resources.rs +++ b/src/internal/resources.rs @@ -25,7 +25,7 @@ impl ResourceTransport { pub fn new(http: reqwest::Client) -> Self { Self { http } } - fn resource_api(&self, scope: &CredentialScope) -> Result { + fn resource_api(&self, scope: &CredentialScope) -> Result { scope.api(&self.http) } } @@ -40,8 +40,7 @@ impl ExperimentApi for ResourceTransport { // Registration is part of admitted event preparation or an explicitly // bound metadata operation; it must remain valid during queue drain. self.resource_api(scope)? - .experiments() - .register(request) + .register_experiment(request) .await .map_err(BraintrustError::from) } @@ -55,7 +54,7 @@ impl ExperimentApi for ResourceTransport { let request = BaseExperimentRequest { id: experiment_id.to_owned(), }; - let response = match api.experiments().base(request).await { + let response = match api.base_experiment(request).await { // The SDK interprets 400 as no base; the API layer preserves the status. Err(crate::api::Error::Http { status: 400, .. }) => return Ok(None), result => result.map_err(BraintrustError::from)?, @@ -73,8 +72,7 @@ impl ExperimentApi for ResourceTransport { base_experiment_id: Option<&str>, ) -> Result { self.resource_api(scope)? - .experiments() - .compare(experiment_id, base_experiment_id) + .compare_experiments(experiment_id, base_experiment_id) .await .map_err(BraintrustError::from) } @@ -91,8 +89,7 @@ impl DatasetApi for ResourceTransport { request: DatasetRegisterRequest, ) -> Result { self.resource_api(scope)? - .datasets() - .register(request) + .register_dataset(request) .await .map_err(BraintrustError::from) } @@ -103,7 +100,6 @@ impl DatasetApi for ResourceTransport { query: BTQLQuery, ) -> Result { self.resource_api(scope)? - .btql() .query_dataset_records(query) .await .map_err(BraintrustError::from) @@ -115,16 +111,14 @@ impl DatasetApi for ResourceTransport { dataset_id: &str, ) -> Result { self.resource_api(scope)? - .datasets() - .summarize(dataset_id) + .summarize_dataset(dataset_id) .await .map_err(BraintrustError::from) } async fn write_dataset_record(&self, scope: &CredentialScope, row: DatasetWrite) -> Result<()> { self.resource_api(scope)? - .datasets() - .write(&[row]) + .write_dataset(&[row]) .await .map_err(BraintrustError::from) } diff --git a/src/internal/session.rs b/src/internal/session.rs index b82e1db..506d907 100644 --- a/src/internal/session.rs +++ b/src/internal/session.rs @@ -47,8 +47,8 @@ impl CredentialScope { identity.org_name = Some(org_name); Self(Arc::new(identity)) } - pub fn api(&self, http: &reqwest::Client) -> crate::error::Result { - crate::api::Client::for_scope( + pub fn api(&self, http: &reqwest::Client) -> crate::error::Result { + crate::api::Transport::for_scope( http.clone(), &self.0.api_url, &self.0.app_url, diff --git a/src/internal/upload.rs b/src/internal/upload.rs index 03ba089..57a6258 100644 --- a/src/internal/upload.rs +++ b/src/internal/upload.rs @@ -81,7 +81,7 @@ pub(crate) struct UploadResult { } pub(crate) struct UploadCoordinator { - logs: api::logs3::Logs3, + logs: api::Transport, objects: ObjectStoreTransport, policy: UploadPolicy, } @@ -93,9 +93,8 @@ impl UploadCoordinator { org_name: Option<&str>, policy: UploadPolicy, ) -> Result { - let logs = api::Client::for_scope(http.clone(), api_url, api_url, token, org_name) - .map_err(BraintrustError::from)? - .logs3(); + let logs = api::Transport::for_scope(http.clone(), api_url, api_url, token, org_name) + .map_err(BraintrustError::from)?; Ok(Self { logs, objects: ObjectStoreTransport { http }, @@ -108,7 +107,7 @@ impl UploadCoordinator { fallback: usize, override_bytes: Option, ) -> Capabilities { - let server_limit = match self.logs.version().await { + let server_limit = match self.logs.logs3_version().await { Ok(version) => version.logs3_payload_max_bytes.filter(|limit| *limit > 0), Err(error) => { warn!(%error, "failed to fetch ingestion capabilities"); @@ -150,7 +149,7 @@ impl UploadCoordinator { match budget .run(|| async { self.logs - .post_raw(batch.data.to_vec()) + .logs3_post_raw(batch.data.to_vec()) .await .map_err(BraintrustError::from) }) @@ -181,7 +180,7 @@ impl UploadCoordinator { let grant = budget .run(|| async { self.logs - .request_overflow_upload(Logs3OverflowUploadRequest { + .logs3_request_overflow_upload(Logs3OverflowUploadRequest { content_type: "application/json".into(), size_bytes: batch.data.len(), rows: batch.overflow_rows.clone(), @@ -196,7 +195,7 @@ impl UploadCoordinator { budget .run(|| async { self.logs - .commit_overflow(Logs3OverflowRequest { + .logs3_commit_overflow(Logs3OverflowRequest { rows: Logs3OverflowReference { reference_type: "logs3_overflow", key: grant.key.clone(), diff --git a/src/log_queue/config.rs b/src/log_queue/config.rs index b371107..e0924aa 100644 --- a/src/log_queue/config.rs +++ b/src/log_queue/config.rs @@ -66,13 +66,6 @@ fn default_flush_chunk_size() -> usize { DEFAULT_FLUSH_CHUNK_SIZE } -fn default_sync_flush() -> bool { - matches!( - std::env::var("BRAINTRUST_SYNC_FLUSH").ok().as_deref(), - Some("1") | Some("true") | Some("yes") - ) -} - fn default_failed_publish_payloads_dir() -> Option { std::env::var("BRAINTRUST_FAILED_PUBLISH_PAYLOADS_DIR") .ok() @@ -100,7 +93,8 @@ fn default_all_publish_payloads_dir() -> Option { /// ``` #[derive(bon::Builder, Debug, Clone)] pub struct LogQueueConfig { - /// Max queue size (BRAINTRUST_QUEUE_DROP_EXCEEDING_MAXSIZE, default: 15000) + /// Maximum accepted writes, including in-flight work + /// (BRAINTRUST_QUEUE_DROP_EXCEEDING_MAXSIZE, default: 15000). #[builder(default = default_queue_max_size())] queue_max_size: usize, @@ -122,15 +116,6 @@ pub struct LogQueueConfig { #[builder(default = default_num_retries())] num_retries: usize, - /// Legacy compatibility value. Delivery is always automatic; ignored. - #[builder(default = default_sync_flush())] - sync_flush: bool, - - /// Legacy compatibility value. Flush reports covered active failures and - /// delivery remains automatic; ignored. - #[builder(default = false)] - checked_delivery: bool, - /// Chunk size for flush processing (BRAINTRUST_LOG_FLUSH_CHUNK_SIZE, default: 25). /// Rows drained from the queue are merged and sent in sequential chunks of this size. /// Each chunk's batches are sent concurrently (matching TypeScript SDK behavior). @@ -142,11 +127,6 @@ pub struct LogQueueConfig { /// Debug: save all payloads (BRAINTRUST_ALL_PUBLISH_PAYLOADS_DIR) all_publish_payloads_dir: Option, - - /// Legacy compatibility value. Admission is always bounded by queue_max_size, - /// including in-flight work; ignored. - #[builder(default = true)] - enforce_queue_size_limit: bool, } impl Default for LogQueueConfig { @@ -176,14 +156,6 @@ impl LogQueueConfig { self.num_retries } - pub fn sync_flush(&self) -> bool { - self.sync_flush - } - - pub fn checked_delivery(&self) -> bool { - self.checked_delivery - } - pub fn flush_chunk_size(&self) -> usize { self.flush_chunk_size } @@ -199,10 +171,6 @@ impl LogQueueConfig { .clone() .or_else(default_all_publish_payloads_dir) } - - pub fn enforce_queue_size_limit(&self) -> bool { - self.enforce_queue_size_limit - } } #[cfg(test)] @@ -216,16 +184,6 @@ mod tests { assert_eq!(config.batch_max_items(), DEFAULT_BATCH_MAX_ITEMS); assert_eq!(config.batch_max_bytes(), DEFAULT_BATCH_MAX_BYTES); assert_eq!(config.num_retries(), DEFAULT_NUM_RETRIES); - assert!(!config.sync_flush()); - assert!(config.enforce_queue_size_limit()); - } - - #[test] - fn test_enforce_queue_size_limit_can_be_enabled() { - let config = LogQueueConfig::builder() - .enforce_queue_size_limit(true) - .build(); - assert!(config.enforce_queue_size_limit()); } #[test] @@ -234,33 +192,11 @@ mod tests { .queue_max_size(5000) .batch_max_items(50) .num_retries(5) - .sync_flush(true) .build(); assert_eq!(config.queue_max_size(), 5000); assert_eq!(config.batch_max_items(), 50); assert_eq!(config.num_retries(), 5); - assert!(config.sync_flush()); - } - - #[test] - fn test_sync_flush_env_var_accepts_truthy_strings() { - for truthy in &["1", "true", "yes"] { - unsafe { std::env::set_var("BRAINTRUST_SYNC_FLUSH", truthy) }; - assert!( - default_sync_flush(), - "BRAINTRUST_SYNC_FLUSH={truthy} should enable sync flush" - ); - } - for falsy in &["0", "false", "no", ""] { - unsafe { std::env::set_var("BRAINTRUST_SYNC_FLUSH", falsy) }; - assert!( - !default_sync_flush(), - "BRAINTRUST_SYNC_FLUSH={falsy} should not enable sync flush" - ); - } - unsafe { std::env::remove_var("BRAINTRUST_SYNC_FLUSH") }; - assert!(!default_sync_flush(), "unset should not enable sync flush"); } #[test] diff --git a/src/log_queue/queue.rs b/src/log_queue/queue.rs index b0d2226..2028949 100644 --- a/src/log_queue/queue.rs +++ b/src/log_queue/queue.rs @@ -453,8 +453,7 @@ impl LogQueueCore { } let response = scope .api(&self.client)? - .projects() - .register(ProjectRegisterRequest { + .register_project(ProjectRegisterRequest { project_name: project_name.to_owned(), org_id: (!scope.org_id().is_empty()).then(|| scope.org_id().to_owned()), org_name: scope.org_name().map(str::to_owned), @@ -639,7 +638,7 @@ impl LogQueueCore { } } -/// Compatibility adapter connecting span preparation to the bounded write service. +/// Span preparation and submission through the bounded write service. #[derive(Clone)] pub struct LogQueue { core: Arc, @@ -655,7 +654,6 @@ impl LogQueue { pub fn new( config: LogQueueConfig, client: reqwest::Client, - _worker_queue_size: usize, execution: ExecutionBinding, ) -> Self { Self { @@ -796,25 +794,13 @@ impl LogQueue { pub(crate) fn is_started(&self) -> bool { matches!(*self.state(), QueueState::Running(_)) } - pub async fn trigger_flush_command(&self) -> Result<()> { - let token = self.flush(); - if let Err(error) = &*token { - return Err(error.clone().into()); - } - Ok(()) - } } -#[async_trait::async_trait] impl DeliveryControl for LogQueue { fn flush(&self) -> Completion { LogQueue::flush(self) } - async fn trigger_flush(&self) -> Result<()> { - self.trigger_flush_command().await - } } -#[async_trait::async_trait] impl SpanSubmitter for LogQueue { fn operation(&self) -> std::result::Result { let lease = { @@ -885,7 +871,6 @@ mod tests { LogQueue::new( LogQueueConfig::builder().num_retries(0).build(), reqwest::Client::new(), - 32, ExecutionBinding::external(tokio::runtime::Handle::current()), ) } @@ -919,7 +904,6 @@ mod tests { .num_retries(0) .build(), reqwest::Client::new(), - 8, ExecutionBinding::external(tokio::runtime::Handle::current()), ); assert!(!queue.is_started()); diff --git a/src/logger.rs b/src/logger.rs index 3b22eb2..b106e44 100644 --- a/src/logger.rs +++ b/src/logger.rs @@ -2,8 +2,6 @@ use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; -use serde::Serialize; -use serde_json::{Map, Value}; use tokio::sync::Notify; use url::Url; @@ -26,7 +24,6 @@ use crate::types::{ParentSpanInfo, SpanAttributes, SpanEventData, SpanObjectType // The TypeScript SDK applies no explicit timeout. const REQUEST_TIMEOUT: Duration = Duration::from_secs(30); const LOGIN_TIMEOUT: Duration = Duration::from_secs(30); -const DEFAULT_QUEUE_SIZE: usize = 256; pub use crate::api::{DEFAULT_API_URL, DEFAULT_APP_URL}; use crate::api::login::LoginResponse; @@ -175,18 +172,14 @@ pub struct BraintrustClientBuilder { org_name: Option, default_project: Option, ca_bundle: Option, - queue_size: usize, blocking_login: bool, skip_login: bool, /// Maximum items per HTTP batch. batch_max_items: Option, /// Maximum bytes per HTTP batch. batch_max_bytes: Option, - /// Maximum queue capacity (None = unlimited). + /// Maximum accepted writes, including in-flight work (None uses environment defaults). queue_max_size: Option, - /// Whether the queue should enforce `queue_max_size`. - enforce_queue_size_limit: bool, - checked_delivery: bool, environment: Option, span_origin: Option, user_agent: Option, @@ -211,14 +204,11 @@ impl BraintrustClientBuilder { org_name: std::env::var("BRAINTRUST_ORG_NAME").ok(), default_project: std::env::var("BRAINTRUST_DEFAULT_PROJECT").ok(), ca_bundle: CaBundleConfig::from_environment(), - queue_size: DEFAULT_QUEUE_SIZE, blocking_login: false, skip_login: false, batch_max_items: None, batch_max_bytes: None, queue_max_size: None, - enforce_queue_size_limit: true, - checked_delivery: false, environment: None, span_origin: None, user_agent: None, @@ -255,29 +245,14 @@ impl BraintrustClientBuilder { self } - /// Set the maximum queue capacity. - /// When full, the newest arriving events are dropped. + /// Set the maximum accepted writes, including in-flight work. + /// When full, new writes return `EnqueueError::Full` at issuance. /// Default: 15,000 (or `BRAINTRUST_QUEUE_DROP_EXCEEDING_MAXSIZE` env var). pub fn with_queue_max_size(mut self, max_size: usize) -> Self { self.queue_max_size = Some(max_size); self } - /// Legacy compatibility setting. Delivery admission is always bounded by - /// `with_queue_max_size`, including in-flight writes. This value is ignored. - pub fn enforce_queue_size_limit(mut self, enforce: bool) -> Self { - self.enforce_queue_size_limit = enforce; - self - } - - /// Legacy compatibility setting. Delivery is always automatic and flush - /// reports failures in covered active work. This value is ignored. Await a - /// write completion to observe that write's result independently of flush. - pub fn checked_delivery(mut self, checked: bool) -> Self { - self.checked_delivery = checked; - self - } - /// Set the app URL (overrides `BRAINTRUST_APP_URL` env var). pub fn app_url(mut self, url: impl Into) -> Self { self.app_url = Some(url.into()); @@ -316,12 +291,6 @@ impl BraintrustClientBuilder { self } - /// Set the internal queue size for buffering log events. - pub fn queue_size(mut self, size: usize) -> Self { - self.queue_size = size; - self - } - /// Wait for login before initialization completes (default: false). /// /// When `false`, the build completion returns a client while login continues @@ -416,17 +385,10 @@ impl BraintrustClientBuilder { .maybe_batch_max_items(self.batch_max_items) .maybe_batch_max_bytes(self.batch_max_bytes) .maybe_queue_max_size(self.queue_max_size) - .enforce_queue_size_limit(self.enforce_queue_size_limit) - .checked_delivery(self.checked_delivery) .build(); // Commands carry their own immutable scope; the worker has no session fallback. - let queue = LogQueue::new( - log_config, - http_client.clone(), - self.queue_size, - execution.clone(), - ); + let queue = LogQueue::new(log_config, http_client.clone(), execution.clone()); let client = BraintrustClient { inner: Arc::new(ClientInner { @@ -524,7 +486,6 @@ impl SpanUpdate { struct ClientInner { lifecycle: Lifecycle, execution: ExecutionBinding, - #[allow(dead_code)] api_url: Url, app_url: Url, queue: LogQueue, @@ -1086,164 +1047,6 @@ impl BraintrustClient { pub fn flush(&self) -> crate::Completion { self.inner.queue.flush() } - - /// Trigger a non-blocking background flush. - /// Does not wait for completion - useful for streaming writes. - pub async fn trigger_flush(&self) -> Result<()> { - self.inner.queue.trigger_flush_command().await - } -} - -#[allow(dead_code)] -#[derive(Debug, Default)] -struct MutableSpanEvent { - input: Option, - output: Option, - expected: Option, - error: Option, - scores: Option>, - metadata: Option>, - metrics: Option>, - tags: Option>, - context: Option, - span_attributes: Option, - extra: HashMap, -} - -#[allow(dead_code)] -#[derive(Serialize)] -struct SerializableSpanEvent { - #[serde(skip_serializing_if = "Option::is_none")] - input: Option, - #[serde(skip_serializing_if = "Option::is_none")] - output: Option, - #[serde(skip_serializing_if = "Option::is_none")] - expected: Option, - #[serde(skip_serializing_if = "Option::is_none")] - error: Option, - #[serde(skip_serializing_if = "Option::is_none")] - scores: Option, - #[serde(skip_serializing_if = "Option::is_none")] - metadata: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - metrics: Option, - #[serde(skip_serializing_if = "Option::is_none")] - tags: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - context: Option, - #[serde(skip_serializing_if = "Option::is_none")] - span_attributes: Option, - #[serde(flatten)] - extra: HashMap, -} - -impl MutableSpanEvent { - fn to_event_map(&self) -> Map { - let serialized = SerializableSpanEvent { - input: self.input.clone(), - output: self.output.clone(), - expected: self.expected.clone(), - error: self.error.clone(), - scores: self.scores.as_ref().map(numeric_map_to_value), - metadata: self.metadata.clone(), - metrics: self.metrics.as_ref().map(numeric_map_to_value), - tags: self.tags.clone(), - context: self.context.clone(), - span_attributes: self - .span_attributes - .clone() - .and_then(|value| serde_json::to_value(value).ok()), - extra: self.extra.clone(), - }; - serde_json::to_value(serialized) - .ok() - .and_then(|value| value.as_object().cloned()) - .unwrap_or_default() - } - - fn apply_merged_event_map(&mut self, mut event: Map) { - self.input = event.remove("input"); - self.output = event.remove("output"); - self.expected = event.remove("expected"); - self.error = event.remove("error"); - self.scores = event.remove("scores").and_then(parse_numeric_map); - self.metadata = event.remove("metadata").and_then(|value| match value { - Value::Object(map) => Some(map), - _ => None, - }); - self.metrics = event.remove("metrics").and_then(parse_numeric_map); - self.tags = event.remove("tags").and_then(parse_string_array); - self.context = event.remove("context"); - self.span_attributes = event - .remove("span_attributes") - .and_then(|value| serde_json::from_value(value).ok()); - self.extra = HashMap::from_iter(event); - } -} - -#[allow(dead_code)] -fn numeric_map_to_value(map: &HashMap) -> Value { - Value::Object(Map::from_iter( - map.iter() - .map(|(key, value)| (key.clone(), Value::from(*value))), - )) -} - -#[allow(dead_code)] -fn apply_propagated_event( - propagated_event: &Map, - event_fields: &mut MutableSpanEvent, -) { - let mut event = event_fields.to_event_map(); - merge_event_maps(&mut event, propagated_event); - event_fields.apply_merged_event_map(event); -} - -#[allow(dead_code)] -fn merge_event_maps(merge_into: &mut Map, merge_from: &Map) { - for (key, merge_from_value) in merge_from { - match (merge_into.get_mut(key), merge_from_value) { - (Some(Value::Object(merge_into_object)), Value::Object(merge_from_object)) => { - merge_event_maps(merge_into_object, merge_from_object); - } - (Some(Value::Array(merge_into_array)), Value::Array(merge_from_array)) - if key == "tags" => - { - for item in merge_from_array { - if !merge_into_array.iter().any(|existing| existing == item) { - merge_into_array.push(item.clone()); - } - } - } - _ => { - merge_into.insert(key.clone(), merge_from_value.clone()); - } - } - } -} - -#[allow(dead_code)] -fn parse_numeric_map(value: Value) -> Option> { - match value { - Value::Object(map) => Some(HashMap::from_iter( - map.into_iter() - .filter_map(|(key, value)| value.as_f64().map(|number| (key, number))), - )), - _ => None, - } -} - -#[allow(dead_code)] -fn parse_string_array(value: Value) -> Option> { - match value { - Value::Array(values) => Some( - values - .into_iter() - .filter_map(|value| value.as_str().map(ToString::to_string)) - .collect(), - ), - _ => None, - } } #[cfg(test)] @@ -1619,13 +1422,13 @@ mod tests { } #[tokio::test] - async fn builder_capacity_is_enforced_even_with_the_legacy_unbounded_setting() { + async fn builder_capacity_is_always_bounded() { let client = BraintrustClient::builder() + .runtime(crate::current_runtime()) .app_url("https://example.com") .api_url("https://example.com") .skip_login(true) .with_queue_max_size(1) - .enforce_queue_size_limit(false) .build() .await .expect("client"); diff --git a/src/span.rs b/src/span.rs index 7d6158f..5ef0b19 100644 --- a/src/span.rs +++ b/src/span.rs @@ -5,8 +5,6 @@ use std::path::PathBuf; use std::sync::{Arc, Mutex}; use std::time::{SystemTime, UNIX_EPOCH}; -#[cfg(test)] -use async_trait::async_trait; use serde_json::{json, Map, Value}; use uuid::Uuid; @@ -678,12 +676,6 @@ impl SpanHandle { .end_time } - /// Trigger a non-blocking background flush. - /// Useful for streaming writes to get partial updates visible immediately. - pub async fn trigger_flush(&self) -> Result<()> { - self.submitter.trigger_flush().await - } - /// Export the span as SpanComponents for passing to other systems or SDKs. /// This creates a serializable representation that can be used as a parent span /// in another context (e.g., across service boundaries via HTTP headers). @@ -1118,7 +1110,6 @@ mod tests { reject: AtomicBool, accepted: Mutex>, } - #[async_trait] impl SpanSubmitter for Submitter { fn operation( &self, @@ -1154,16 +1145,12 @@ mod tests { completion } } - #[async_trait] impl crate::internal::ports::DeliveryControl for Submitter { fn flush(&self) -> crate::Completion { let (sender, completion) = crate::Completion::pair(); let _ = sender.send(Ok(crate::FlushReport::default())); completion } - async fn trigger_flush(&self) -> Result<()> { - Ok(()) - } } let submitter = Arc::new(Submitter { reject: AtomicBool::new(true), @@ -1199,18 +1186,13 @@ mod tests { async fn concurrent_ends_share_one_snapshot_and_one_delivery_result() { use crate::internal::shared_outcome::WriteResponder; struct Submitter(Mutex>); - #[async_trait] impl crate::internal::ports::DeliveryControl for Submitter { fn flush(&self) -> crate::Completion { let (sender, completion) = crate::Completion::pair(); let _ = sender.send(Ok(crate::FlushReport::default())); completion } - async fn trigger_flush(&self) -> Result<()> { - Ok(()) - } } - #[async_trait] impl SpanSubmitter for Submitter { fn operation( &self, diff --git a/src/stream.rs b/src/stream.rs index ea72c71..26a9bf1 100644 --- a/src/stream.rs +++ b/src/stream.rs @@ -1091,7 +1091,7 @@ where drop(agg); // Non-blocking flush for real-time updates - let _ = span.trigger_flush().await; + drop(span.flush()); *last_flush.lock().await = Instant::now(); } } diff --git a/src/test_utils.rs b/src/test_utils.rs index a6007c0..96fc203 100644 --- a/src/test_utils.rs +++ b/src/test_utils.rs @@ -1,7 +1,5 @@ use std::sync::{Arc, Mutex}; -use async_trait::async_trait; - use crate::error::Result; use crate::internal::session::CredentialScope; use crate::span::{SpanBuilder, SpanHandle}; @@ -58,7 +56,6 @@ impl MockBraintrustClient { } } -#[async_trait] impl SpanSubmitter for MockBraintrustClient { fn operation( &self, @@ -112,18 +109,12 @@ impl SpanSubmitter for MockBraintrustClient { } } -#[async_trait] impl DeliveryControl for MockBraintrustClient { fn flush(&self) -> crate::Completion { let (sender, completion) = crate::Completion::pair(); let _ = sender.send(Ok(crate::FlushReport::default())); completion } - - async fn trigger_flush(&self) -> Result<()> { - // Mock: no-op for tests - Ok(()) - } } pub(crate) fn mock_span_builder() -> (SpanBuilder, TestSpanCollector) { diff --git a/tests/eager_api.rs b/tests/eager_api.rs new file mode 100644 index 0000000..2e15533 --- /dev/null +++ b/tests/eager_api.rs @@ -0,0 +1,312 @@ +use braintrust_sdk_rust::{ + api::{ + self, + datasets::{DatasetFields, DatasetWrite}, + logs3::InsertRows, + }, + Completion, EnqueueError, ExecutionError, ShutdownError, +}; +use serde_json::json; +use std::{future::IntoFuture, time::Duration}; +use wiremock::{ + matchers::{body_json, method, path}, + Mock, MockServer, ResponseTemplate, +}; + +fn config(server: &MockServer) -> api::Config { + let mut config = api::Config::new("token"); + config.api_url = server.uri(); + config.app_url = server.uri(); + config +} +fn abandon(completion: Completion) { + assert!(completion.is_ok()); + drop(completion); +} + +#[tokio::test] +async fn all_endpoint_families_run_unobserved_from_a_plain_thread_and_drain_after_owner_drop() { + use api::{ + btql::{BTQLQuery, BTQLQueryInner}, + datasets::DatasetRegisterRequest, + experiments::{BaseExperimentRequest, ExperimentRegisterRequest}, + logs3::{Logs3OverflowReference, Logs3OverflowRequest, Logs3OverflowUploadRequest}, + projects::ProjectRegisterRequest, + }; + let server = MockServer::start().await; + let fixtures = [ + ("POST", "/api/apikey/login", json!({"org_info":[]})), + ( + "POST", + "/api/project/register", + json!({"project":{"id":"p"}}), + ), + ( + "POST", + "/api/dataset/register", + json!({"project":{"id":"p","name":"p"},"dataset":{"id":"d","name":"d"}}), + ), + ( + "GET", + "/dataset-summary", + json!({"project_name":"p","dataset_name":"d","project_id":"p","dataset_id":"d"}), + ), + ("POST", "/btql", json!({"data":[],"cursor":null})), + ( + "POST", + "/api/experiment/register", + json!({"project":{"id":"p","name":"p"},"experiment":{"id":"e","name":"e"}}), + ), + ("POST", "/api/base_experiment/get_id", json!({})), + ("GET", "/experiment-comparison2", json!({})), + ("GET", "/version", json!({})), + ( + "POST", + "/logs3/overflow", + json!({"method":"PUT","signedUrl":"http://unused.invalid","key":"k"}), + ), + ]; + for (verb, endpoint, response) in fixtures { + Mock::given(method(verb)) + .and(path(endpoint)) + .respond_with(ResponseTemplate::new(200).set_body_json(response)) + .expect(1) + .mount(&server) + .await; + } + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(204)) + .expect(4) + .mount(&server) + .await; + let config = config(&server); + let shutdown = std::thread::spawn(move || { + let client = api::Client::new(config).unwrap(); + abandon(client.login()); + abandon(client.projects().register(ProjectRegisterRequest { + project_name: "p".into(), + org_id: None, + org_name: None, + })); + abandon(client.datasets().register(DatasetRegisterRequest { + project_name: "p".into(), + org_name: "org".into(), + dataset_name: "d".into(), + description: None, + metadata: None, + })); + abandon(client.datasets().write(&[DatasetWrite::Delete { + dataset_id: "d".into(), + id: "row".into(), + }])); + abandon(client.datasets().summarize(&String::from("d"))); + abandon(client.btql().query_dataset_records(BTQLQuery { + query: BTQLQueryInner { + from: "dataset('d')".into(), + limit: None, + cursor: None, + }, + })); + abandon(client.experiments().register(ExperimentRegisterRequest { + project_name: "p".into(), + org_id: "org".into(), + experiment_name: None, + description: None, + base_experiment: None, + metadata: None, + public: None, + update: None, + repo_info: None, + })); + abandon( + client + .experiments() + .base(BaseExperimentRequest { id: "e".into() }), + ); + abandon( + client + .experiments() + .compare(&String::from("e"), Some(&String::from("base"))), + ); + abandon(client.logs3().version()); + abandon( + client + .logs3() + .insert(InsertRows::new(vec![json!({"id":"insert"})])), + ); + abandon( + client + .logs3() + .post_raw(br#"{"rows":[],"api_version":2}"#.to_vec()), + ); + abandon( + client + .logs3() + .request_overflow_upload(Logs3OverflowUploadRequest { + content_type: "application/json".into(), + size_bytes: 2, + rows: vec![], + }), + ); + abandon(client.logs3().commit_overflow(Logs3OverflowRequest { + rows: Logs3OverflowReference { + reference_type: "file", + key: "k".into(), + }, + api_version: 2, + })); + client.shutdown(Duration::from_secs(5)) + }) + .join() + .unwrap(); + shutdown.await.unwrap(); + assert_eq!(server.received_requests().await.unwrap().len(), 14); +} + +#[tokio::test] +async fn borrowed_request_data_is_captured_at_issuance_and_shutdown_closes_all_views() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .and(body_json( + json!({"rows":[{"dataset_id":"d","id":"original","input":1}],"api_version":2}), + )) + .respond_with(ResponseTemplate::new(204)) + .expect(1) + .mount(&server) + .await; + let mut config = config(&server); + config.runtime = braintrust_sdk_rust::current_runtime(); + let client = api::Client::new(config).unwrap(); + let view = client.datasets(); + let mut rows = vec![DatasetWrite::Insert { + dataset_id: "d".into(), + id: "original".into(), + fields: DatasetFields { + input: Some(json!(1)), + ..Default::default() + }, + }]; + let write = view.write(&rows); + rows.clear(); + let shutdown = client.shutdown(Duration::from_secs(5)); + let rejected = view.summarize("d"); + assert!(matches!( + &*rejected, + Err(api::Error::Admission(EnqueueError::Closed)) + )); + assert!(matches!( + rejected.await, + Err(api::Error::Admission(EnqueueError::Closed)) + )); + assert!(matches!( + client.clone().login().await, + Err(api::Error::Admission(EnqueueError::Closed)) + )); + drop(view); + drop(client); + shutdown.await.unwrap(); + write.await.unwrap(); +} + +#[test] +fn bounded_admission_and_runtime_loss_do_not_depend_on_observer_polling() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let mut config = api::Config::new("token"); + config.runtime = Some(runtime.handle().clone()); + config.api_url = "http://127.0.0.1:9".into(); + let client = api::Client::new(config).unwrap(); + let mut accepted = Vec::new(); + for _ in 0..1024 { + let token = client.logs3().version(); + assert!(token.is_ok()); + accepted.push(token); + } + drop(accepted.pop()); + let full = client.logs3().version(); + assert!(matches!( + &*full, + Err(api::Error::Admission(EnqueueError::Full)) + )); + assert!(matches!( + futures::executor::block_on(full.into_future()), + Err(api::Error::Admission(EnqueueError::Full)) + )); + let shutdown = client.shutdown(Duration::from_secs(5)); + drop(client); + drop(runtime); + for token in accepted { + assert!(matches!( + futures::executor::block_on(token.into_future()), + Err(api::Error::Execution(ExecutionError::WorkerLost)) + )); + } + assert!(matches!( + futures::executor::block_on(shutdown.into_future()), + Err(ShutdownError::Execution(ExecutionError::WorkerLost)) + )); +} + +#[tokio::test] +async fn first_shutdown_deadline_cancels_requests_and_is_shared_by_later_observers() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/version")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(json!({})) + .set_delay(Duration::from_secs(2)), + ) + .expect(1) + .mount(&server) + .await; + let client = api::Client::new(config(&server)).unwrap(); + let request = client.logs3().version(); + tokio::time::timeout(Duration::from_secs(5), async { + while server.received_requests().await.unwrap().is_empty() { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + let first = client.shutdown(Duration::from_millis(20)); + let later = client.clone().shutdown(Duration::from_secs(10)); + for token in [first, later] { + assert!(matches!( + token.await, + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + )); + } + assert!(matches!( + request.await, + Err(api::Error::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + client.shutdown(Duration::from_secs(10)).await, + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + )); +} + +#[tokio::test] +async fn request_errors_remain_on_their_tokens_and_do_not_poison_shutdown() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/version")) + .respond_with(ResponseTemplate::new(500).set_body_string("failed")) + .expect(1) + .mount(&server) + .await; + let client = api::Client::new(config(&server)).unwrap(); + let request = client.logs3().version(); + let shutdown = client.shutdown(Duration::from_secs(5)); + shutdown.await.unwrap(); + assert!(matches!( + request.await, + Err(api::Error::Http { status: 500, .. }) + )); + client.shutdown(Duration::ZERO).await.unwrap(); +}