diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 018d02a..5a9a35e 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -4,7 +4,6 @@ on: push: branches: [main] pull_request: - branches: [main] env: CARGO_TERM_COLOR: always diff --git a/Cargo.toml b/Cargo.toml index 264b265..22e77af 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,13 +14,11 @@ categories = ["api-bindings", "asynchronous", "development-tools"] anyhow = "1" base64 = "0.22" bon = "3" -arc-swap = "1" async-trait = "0.1" backoff = { version = "0.4", features = ["tokio"] } bytes = "1" futures = "0.3" chrono = { version = "0.4", features = ["serde"] } -crossbeam = "0.8" indexmap = "2" reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "stream", "rustls-tls-native-roots"] } serde = { version = "1", features = ["derive"] } diff --git a/src/api/btql.rs b/src/api/btql.rs new file mode 100644 index 0000000..9bee687 --- /dev/null +++ b/src/api/btql.rs @@ -0,0 +1,70 @@ +//! BTQL dataset readback, independent of SDK ingestion. +use super::{Client, Error, Plane}; +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; + +/// A record within a dataset. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[non_exhaustive] +pub struct DatasetRecord { + /// Unique record ID. + pub id: String, + /// Input data for the operation. + #[serde(skip_serializing_if = "Option::is_none")] + pub input: Option, + /// Expected output for evaluation. + #[serde(skip_serializing_if = "Option::is_none")] + pub expected: Option, + /// Key-value metadata. + #[serde(skip_serializing_if = "Option::is_none")] + pub metadata: Option>, + /// String tags for filtering. + #[serde(skip_serializing_if = "Option::is_none")] + pub tags: Option>, + /// Transaction ID (for versioning). + #[serde(rename = "_xact_id", skip_serializing_if = "Option::is_none")] + pub xact_id: Option, + /// ISO 8601 timestamp. + #[serde(skip_serializing_if = "Option::is_none")] + pub created: Option, +} + +/// BTQL query for fetching dataset records. +#[derive(Debug, Clone, Serialize)] +pub struct BTQLQuery { + pub query: BTQLQueryInner, +} + +/// Inner BTQL query structure. +#[derive(Debug, Clone, Serialize)] +pub struct BTQLQueryInner { + /// Dataset reference: dataset('id') + pub from: String, + /// Maximum records per page. + #[serde(skip_serializing_if = "Option::is_none")] + pub limit: Option, + /// Pagination cursor token. + #[serde(skip_serializing_if = "Option::is_none")] + pub cursor: Option, +} + +/// Response from BTQL query. +#[derive(Debug, Clone, Deserialize)] +pub struct BTQLResponse { + /// Records in JSONL format. + #[serde(default)] + pub data: Vec, + /// Cursor for next page (None if last page). + pub cursor: Option, +} + +#[derive(Clone)] +pub struct Btql { + pub(crate) client: Client, +} +impl Btql { + /// 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 + } +} diff --git a/src/api/client.rs b/src/api/client.rs new file mode 100644 index 0000000..aecb013 --- /dev/null +++ b/src/api/client.rs @@ -0,0 +1,267 @@ +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, +} + +/// Immutable HTTP client shared by lightweight endpoint views. +#[derive(Clone)] +pub struct Client { + inner: Arc, +} +struct Inner { + http: reqwest::Client, + api_url: Url, + app_url: Url, + headers: HeaderMap, +} +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()))?; + 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, + }), + }) + } + + pub fn logs3(&self) -> super::logs3::Logs3 { + super::logs3::Logs3 { + client: self.clone(), + } + } + pub fn btql(&self) -> super::btql::Btql { + super::btql::Btql { + client: self.clone(), + } + } + pub fn datasets(&self) -> super::datasets::Datasets { + super::datasets::Datasets { + client: self.clone(), + } + } + pub fn experiments(&self) -> super::experiments::Experiments { + super::experiments::Experiments { + client: self.clone(), + } + } + pub fn projects(&self) -> super::projects::Projects { + super::projects::Projects { + 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); + } + 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/api/config.rs b/src/api/config.rs new file mode 100644 index 0000000..a582b75 --- /dev/null +++ b/src/api/config.rs @@ -0,0 +1,42 @@ +use std::{fmt, path::PathBuf, time::Duration}; + +pub const DEFAULT_APP_URL: &str = "https://www.braintrust.dev"; +pub const DEFAULT_API_URL: &str = "https://api.braintrust.dev"; + +/// Explicit configuration for direct API calls. Construction does not log in. +#[derive(Clone)] +pub struct Config { + pub(crate) api_key: String, + pub api_url: String, + pub app_url: String, + pub org_name: Option, + pub timeout: Duration, + pub user_agent: Option, + /// An explicit PEM file; mutually exclusive with `custom_ca_bundle_pem`. + pub custom_ca_bundle_path: Option, + /// Raw PEM data; otherwise the existing environment CA convention applies. + pub custom_ca_bundle_pem: Option>, +} +impl Config { + pub fn new(api_key: impl Into) -> Self { + Self { + api_key: api_key.into(), + api_url: DEFAULT_API_URL.into(), + app_url: DEFAULT_APP_URL.into(), + org_name: None, + timeout: Duration::from_secs(30), + user_agent: None, + custom_ca_bundle_path: None, + custom_ca_bundle_pem: None, + } + } +} +impl fmt::Debug for Config { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + // Endpoints may be invalid and contain credentials before validation. + f.debug_struct("Config") + .field("api_key", &"[redacted]") + .field("timeout", &self.timeout) + .finish_non_exhaustive() + } +} diff --git a/src/api/datasets.rs b/src/api/datasets.rs new file mode 100644 index 0000000..14185d3 --- /dev/null +++ b/src/api/datasets.rs @@ -0,0 +1,170 @@ +//! Dataset registration, row mutation, and summary endpoint contracts. +use super::{Client, Error, Plane}; +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; + +/// Request body for dataset registration. +#[derive(Debug, Clone, Serialize)] +pub struct DatasetRegisterRequest { + pub project_name: String, + pub org_name: String, + pub dataset_name: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub description: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub metadata: Option>, +} + +/// Response from dataset registration. +#[derive(Debug, Clone, Deserialize)] +pub struct DatasetRegisterResponse { + pub project: DatasetProjectInfo, + pub dataset: DatasetInfo, +} + +/// Project info from registration response. +#[derive(Debug, Clone, Deserialize)] +pub struct DatasetProjectInfo { + pub id: String, + pub name: String, +} + +/// Dataset info from registration response. +#[derive(Debug, Clone, Deserialize)] +pub struct DatasetInfo { + pub id: String, + pub name: String, +} + +/// Fields stored in a dataset record. Absent fields are omitted from the wire. +#[derive(Debug, Clone, Default, Serialize)] +pub struct DatasetFields { + #[serde(skip_serializing_if = "Option::is_none")] + pub input: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub expected: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub metadata: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub tags: Option>, +} + +/// A dataset row mutation, independent of tracing/span payloads. +#[derive(Debug, Clone)] +pub enum DatasetWrite { + Insert { + dataset_id: String, + id: String, + fields: DatasetFields, + }, + /// Merge only the supplied fields into an existing record. + Update { + dataset_id: String, + id: String, + fields: DatasetFields, + }, + Delete { + dataset_id: String, + id: String, + }, +} + +impl Serialize for DatasetWrite { + fn serialize(&self, serializer: S) -> Result { + #[derive(Serialize)] + struct Row<'a> { + dataset_id: &'a str, + id: &'a str, + #[serde(flatten)] + fields: Option<&'a DatasetFields>, + #[serde(rename = "_is_merge", skip_serializing_if = "Option::is_none")] + is_merge: Option, + #[serde(rename = "_object_delete", skip_serializing_if = "Option::is_none")] + is_delete: Option, + } + let (dataset_id, id, fields, is_merge, is_delete) = match self { + Self::Insert { + dataset_id, + id, + fields, + } => (dataset_id, id, Some(fields), None, None), + Self::Update { + dataset_id, + id, + fields, + } => (dataset_id, id, Some(fields), Some(true), None), + Self::Delete { dataset_id, id } => (dataset_id, id, None, None, Some(true)), + }; + Row { + dataset_id, + id, + fields, + is_merge, + is_delete, + } + .serialize(serializer) + } +} + +/// Response from dataset summary API. +#[derive(Debug, Clone, Deserialize)] +pub struct DatasetSummaryResponse { + pub project_name: String, + pub dataset_name: String, + pub project_id: String, + pub dataset_id: String, + pub project_url: Option, + pub dataset_url: Option, + pub data_summary: Option, +} + +/// Data summary statistics. +#[derive(Debug, Clone, Deserialize)] +pub struct DataSummaryStats { + pub total_records: Option, +} + +#[derive(Clone)] +pub struct Datasets { + pub(crate) client: Client, +} +impl Datasets { + /// Register once on the control plane; no resolution cache or retry. + pub async fn register( + &self, + request: DatasetRegisterRequest, + ) -> Result { + self.client + .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> { + #[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 + } + + pub async fn summarize(&self, dataset_id: &str) -> Result { + self.client + .get_json( + Plane::Data, + "dataset-summary", + &[("dataset_id", dataset_id)], + ) + .await + } +} diff --git a/src/api/error.rs b/src/api/error.rs new file mode 100644 index 0000000..055faa6 --- /dev/null +++ b/src/api/error.rs @@ -0,0 +1,29 @@ +use thiserror::Error; + +/// Direct request failures retain their operation, status and original source. +#[derive(Debug, Error)] +#[non_exhaustive] +pub enum Error { + #[error("invalid API configuration: {0}")] + Config(String), + #[error("{operation}: transport failed")] + Transport { + operation: &'static str, + #[source] + source: reqwest::Error, + }, + #[error("{operation}: HTTP {status}")] + Http { + operation: &'static str, + status: u16, + /// Server response retained for programmatic inspection, not Display. + body: String, + }, + #[error("{operation}: invalid JSON response (HTTP {status})")] + Decode { + operation: &'static str, + status: u16, + #[source] + source: serde_json::Error, + }, +} diff --git a/src/api/experiments.rs b/src/api/experiments.rs new file mode 100644 index 0000000..0e96bf4 --- /dev/null +++ b/src/api/experiments.rs @@ -0,0 +1,154 @@ +//! Existing experiment registration, base and comparison endpoint contracts. +use super::{Client, Error, Plane}; +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; +use std::collections::HashMap; + +/// Repository information for experiment tracking. +#[derive(Debug, Clone, Default, Serialize)] +pub struct RepoInfo { + /// Git commit SHA. + #[serde(skip_serializing_if = "Option::is_none")] + pub commit: Option, + /// Current branch name. + #[serde(skip_serializing_if = "Option::is_none")] + pub branch: Option, + /// Current tag (if any). + #[serde(skip_serializing_if = "Option::is_none")] + pub tag: Option, + /// Whether the working directory has uncommitted changes. + #[serde(skip_serializing_if = "Option::is_none")] + pub dirty: Option, + /// Git author name. + #[serde(skip_serializing_if = "Option::is_none")] + pub author_name: Option, + /// Git author email. + #[serde(skip_serializing_if = "Option::is_none")] + pub author_email: Option, + /// Git commit message. + #[serde(skip_serializing_if = "Option::is_none")] + pub commit_message: Option, + /// Git commit timestamp. + #[serde(skip_serializing_if = "Option::is_none")] + pub commit_time: Option, +} + +/// Request body for experiment registration. +#[derive(Debug, Clone, Serialize)] +pub struct ExperimentRegisterRequest { + pub project_name: String, + pub org_id: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub experiment_name: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub description: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub base_experiment: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub metadata: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub public: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub update: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub repo_info: Option, +} + +/// Response from experiment registration. +#[derive(Debug, Clone, Deserialize)] +pub struct ExperimentRegisterResponse { + pub project: ProjectInfo, + pub experiment: ExperimentInfo, +} + +/// Project information from registration response. +#[derive(Debug, Clone, Deserialize)] +pub struct ProjectInfo { + pub id: String, + pub name: String, +} + +/// Experiment information from registration response. +#[derive(Debug, Clone, Deserialize)] +pub struct ExperimentInfo { + pub id: String, + pub name: String, +} + +/// Request body for fetching base experiment. +#[derive(Debug, Clone, Serialize)] +pub struct BaseExperimentRequest { + pub id: String, +} + +/// Response from base experiment fetch. +#[derive(Debug, Clone, Deserialize)] +pub struct BaseExperimentResponse { + pub base_exp_id: Option, + pub base_exp_name: Option, +} +/// Response from experiment comparison API. +#[derive(Debug, Clone, Deserialize)] +pub struct ExperimentComparisonResponse { + #[serde(default)] + pub scores: HashMap, + #[serde(default)] + pub metrics: HashMap, + pub base_experiment_name: Option, +} + +/// Score data from comparison API. +#[derive(Debug, Clone, Deserialize)] +pub struct ComparisonScoreData { + pub score: f64, + pub diff: Option, + pub improvements: Option, + pub regressions: Option, +} + +/// Metric data from comparison API. +#[derive(Debug, Clone, Deserialize)] +pub struct ComparisonMetricData { + pub metric: f64, + pub unit: Option, + pub diff: Option, + pub improvements: Option, + pub regressions: Option, +} + +#[derive(Clone)] +pub struct Experiments { + pub(crate) client: Client, +} +impl Experiments { + pub async fn register( + &self, + request: ExperimentRegisterRequest, + ) -> Result { + self.client + .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( + &self, + request: BaseExperimentRequest, + ) -> Result { + self.client + .post_json(Plane::Control, "api/base_experiment/get_id", &request) + .await + } + pub async fn compare( + &self, + experiment_id: &str, + base_experiment_id: Option<&str>, + ) -> Result { + let mut query = vec![("experiment_id", experiment_id)]; + if let Some(base) = base_experiment_id { + query.push(("base_experiment_id", base)); + } + self.client + .get_json(Plane::Control, "experiment-comparison2", &query) + .await + } +} diff --git a/src/api/login.rs b/src/api/login.rs new file mode 100644 index 0000000..158762d --- /dev/null +++ b/src/api/login.rs @@ -0,0 +1,34 @@ +//! Login response values. Login never mutates SDK session state. +use serde::Deserialize; + +/// Organization info returned from login. +#[derive(Debug, Clone, Deserialize)] +pub struct OrgInfo { + pub(crate) id: String, + pub(crate) name: String, + #[serde(default)] + pub(crate) api_url: Option, +} + +impl OrgInfo { + /// The organization's unique ID. + pub fn id(&self) -> &str { + &self.id + } + + /// The organization's name. + pub fn name(&self) -> &str { + &self.name + } + + /// The organization's API URL, if set. + pub fn api_url(&self) -> Option<&str> { + self.api_url.as_deref() + } +} + +/// Response from the login endpoint. +#[derive(Debug, Deserialize)] +pub struct LoginResponse { + pub org_info: Vec, +} diff --git a/src/api/logs3.rs b/src/api/logs3.rs new file mode 100644 index 0000000..b8f56ab --- /dev/null +++ b/src/api/logs3.rs @@ -0,0 +1,98 @@ +//! Direct logs3 requests: no batching, splitting, retries or implicit overflow. +use super::{Client, Error, Plane}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::HashMap; + +/// Response from POST logs3/overflow — provides a signed URL to upload the payload. +#[derive(Debug, Clone, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct Logs3OverflowUpload { + pub method: String, + pub signed_url: String, + pub headers: Option>, + pub fields: Option>, + pub key: String, +} + +/// A single row's overflow metadata, sent when requesting the overflow upload URL. +#[derive(Debug, Clone, Serialize)] +pub struct Logs3OverflowInputRow { + /// Key identifying fields extracted from the row (experiment_id, dataset_id, etc.) + pub object_ids: serde_json::Map, + #[serde(skip_serializing_if = "Option::is_none")] + pub is_delete: Option, + pub input_row: Logs3OverflowInputRowMeta, +} + +#[derive(Debug, Clone, Serialize)] +pub struct Logs3OverflowInputRowMeta { + pub byte_size: usize, +} + +/// Reference sent to POST logs3 after a successful overflow upload. +#[derive(Debug, Clone, Serialize)] +pub struct Logs3OverflowRequest { + pub rows: Logs3OverflowReference, + pub api_version: u8, +} + +#[derive(Debug, Clone, Serialize)] +pub struct Logs3OverflowReference { + #[serde(rename = "type")] + pub reference_type: &'static str, + pub key: String, +} + +#[derive(Debug, Clone, Deserialize)] +pub struct VersionInfo { + pub logs3_payload_max_bytes: Option, +} +/// Explicit request for an overflow grant. Uploading the object is a separate operation. +#[derive(Debug, Clone, Serialize)] +pub struct Logs3OverflowUploadRequest { + pub content_type: String, + pub size_bytes: usize, + pub rows: Vec, +} +/// Input for a single logs3 insertion. Rows use the existing ingestion wire format. +#[derive(Debug, Clone, Serialize)] +pub struct InsertRows { + pub rows: Vec, + pub api_version: u8, +} +impl InsertRows { + pub fn new(rows: Vec) -> Self { + Self { + rows, + api_version: 2, + } + } +} +#[derive(Clone)] +pub struct Logs3 { + pub(crate) client: Client, +} +impl Logs3 { + pub async fn version(&self) -> Result { + self.client.get_json(Plane::Data, "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 + } + pub async fn insert(&self, request: InsertRows) -> Result<(), Error> { + self.client.post_empty(Plane::Data, "logs3", &request).await + } + pub async fn request_overflow_upload( + &self, + request: Logs3OverflowUploadRequest, + ) -> Result { + self.client + .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 + } +} diff --git a/src/api/mod.rs b/src/api/mod.rs new file mode 100644 index 0000000..9c5faae --- /dev/null +++ b/src/api/mod.rs @@ -0,0 +1,22 @@ +//! 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. + +pub mod btql; +mod client; +mod config; +pub mod datasets; +mod error; +pub mod experiments; +pub mod login; +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; diff --git a/src/api/projects.rs b/src/api/projects.rs new file mode 100644 index 0000000..5207e92 --- /dev/null +++ b/src/api/projects.rs @@ -0,0 +1,36 @@ +//! Existing project registration endpoint. +use super::{Client, Error, Plane}; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Serialize)] +pub struct ProjectRegisterRequest { + pub project_name: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub org_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub org_name: Option, +} +#[derive(Debug, Clone, Deserialize)] +pub struct ProjectRegisterResponse { + pub project: ProjectInfo, +} +#[derive(Debug, Clone, Deserialize)] +pub struct ProjectInfo { + pub id: String, + #[serde(default)] + pub name: Option, +} +#[derive(Clone)] +pub struct Projects { + pub(crate) client: Client, +} +impl Projects { + pub async fn register( + &self, + request: ProjectRegisterRequest, + ) -> Result { + self.client + .post_json(Plane::Control, "api/project/register", &request) + .await + } +} diff --git a/src/completion.rs b/src/completion.rs new file mode 100644 index 0000000..9a35ef8 --- /dev/null +++ b/src/completion.rs @@ -0,0 +1,131 @@ +//! Observation of work submitted when an SDK method is called. +use crate::ExecutionError; +use std::{ + future::{Future, IntoFuture}, + ops::Deref, + pin::Pin, + task::{Context, Poll}, +}; +use tokio::sync::oneshot; + +/// A synchronously recorded issuance result and an eventual operation result. +/// +/// Inspect `&*completion` to check admission without waiting. Awaiting returns +/// that error, or the accepted operation's result. Dropping this token does not +/// cancel accepted work. No work is started by polling the token. +/// +/// ```no_run +/// use braintrust_sdk_rust::{BraintrustClient, SpanHandle, SpanLog}; +/// use std::time::Duration; +/// +/// # async fn example(client: &BraintrustClient, span: &SpanHandle) +/// # -> Result<(), Box> { +/// let completion = span.log(SpanLog::default()); // admission happens now +/// if let Err(error) = &*completion { +/// eprintln!("submission rejected: {error}"); +/// } +/// completion.await?; // observes admission failure or full batch delivery +/// let report = client.flush().await?; +/// client.shutdown(Duration::from_secs(10)).await?; +/// # Ok(()) +/// # } +/// ``` +pub struct Completion { + issued: Result<(), E>, + receiver: oneshot::Receiver>, +} +impl Completion { + pub(crate) fn pair() -> (oneshot::Sender>, Self) { + let (sender, receiver) = oneshot::channel(); + ( + sender, + Self { + issued: Ok(()), + receiver, + }, + ) + } + pub(crate) fn record_issuance(&mut self, result: Result<(), E>) { + self.issued = result; + } +} +impl Deref for Completion { + type Target = Result<(), E>; + fn deref(&self) -> &Self::Target { + &self.issued + } +} +impl std::fmt::Debug for Completion +where + E: std::fmt::Debug, +{ + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Completion") + .field("issued", &self.issued) + .finish_non_exhaustive() + } +} +/// The future produced by consuming a [`Completion`]. +pub struct CompletionFuture { + issued: Option>, + receiver: oneshot::Receiver>, +} +// Neither field is structurally pinned; the channel receiver is Unpin. +impl Unpin for CompletionFuture {} +impl> IntoFuture for Completion { + type Output = Result; + type IntoFuture = CompletionFuture; + fn into_future(self) -> Self::IntoFuture { + CompletionFuture { + issued: Some(self.issued), + receiver: self.receiver, + } + } +} +impl> Future for CompletionFuture { + type Output = Result; + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.get_mut(); + if let Some(Err(error)) = this.issued.take() { + return Poll::Ready(Err(error)); + } + match Pin::new(&mut this.receiver).poll(cx) { + Poll::Pending => Poll::Pending, + Poll::Ready(Ok(result)) => Poll::Ready(result), + Poll::Ready(Err(_)) => Poll::Ready(Err(ExecutionError::WorkerLost.into())), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + #[tokio::test] + async fn issuance_error_needs_neither_clone_nor_producer() { + #[derive(Debug, PartialEq)] + enum Error { + Admission, + Execution, + } + impl From for Error { + fn from(_: ExecutionError) -> Self { + Self::Execution + } + } + let (sender, mut token) = Completion::<(), Error>::pair(); + token.record_issuance(Err(Error::Admission)); + assert_eq!(&*token, &Err(Error::Admission)); + drop(sender); + assert_eq!(token.await, Err(Error::Admission)); + } + #[tokio::test] + async fn successful_issuance_stays_successful_after_operation_failure() { + let (sender, token) = Completion::<(), ExecutionError>::pair(); + sender.send(Err(ExecutionError::TimedOut)).unwrap(); + assert!(token.is_ok()); + assert_eq!(token.await, Err(ExecutionError::TimedOut)); + let (sender, token) = Completion::<(), ExecutionError>::pair(); + drop(sender); + assert_eq!(token.await, Err(ExecutionError::WorkerLost)); + } +} diff --git a/src/dataset.rs b/src/dataset.rs index 0ada6a8..5bcacf7 100644 --- a/src/dataset.rs +++ b/src/dataset.rs @@ -20,25 +20,22 @@ //! .input(json!({"question": "What is 2+2?"})) //! .expected(json!({"answer": "4"})) //! .build()? -//! ).await; -//! -//! dataset.flush().await?; +//! ).await?; //! ``` use std::collections::VecDeque; use std::fmt; use std::sync::Arc; -use std::time::{SystemTime, UNIX_EPOCH}; -use async_trait::async_trait; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; -use tokio::sync::{Mutex, OnceCell}; +use tokio::sync::OnceCell; use uuid::Uuid; +use crate::api::datasets::{DatasetFields, DatasetWrite}; use crate::error::Result; -use crate::span::SpanSubmitter; -use crate::types::{ParentSpanInfo, SpanPayload}; +use crate::internal::ports::{DatasetApi, DeliveryControl}; +use crate::internal::session::CredentialScope; // ============================================================================ // Dataset Insert Types @@ -79,6 +76,17 @@ impl DatasetInsert { } } +impl From for DatasetFields { + fn from(record: DatasetInsert) -> Self { + Self { + input: record.input, + expected: record.expected, + metadata: record.metadata, + tags: record.tags, + } + } +} + /// Builder for DatasetInsert. #[derive(Clone, Default)] pub struct DatasetInsertBuilder { @@ -146,31 +154,7 @@ impl DatasetInsertBuilder { // Dataset Record // ============================================================================ -/// A record within a dataset. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[non_exhaustive] -pub struct DatasetRecord { - /// Unique record ID. - pub id: String, - /// Input data for the operation. - #[serde(skip_serializing_if = "Option::is_none")] - pub input: Option, - /// Expected output for evaluation. - #[serde(skip_serializing_if = "Option::is_none")] - pub expected: Option, - /// Key-value metadata. - #[serde(skip_serializing_if = "Option::is_none")] - pub metadata: Option>, - /// String tags for filtering. - #[serde(skip_serializing_if = "Option::is_none")] - pub tags: Option>, - /// Transaction ID (for versioning). - #[serde(rename = "_xact_id", skip_serializing_if = "Option::is_none")] - pub xact_id: Option, - /// ISO 8601 timestamp. - #[serde(skip_serializing_if = "Option::is_none")] - pub created: Option, -} +pub use crate::api::btql::DatasetRecord; // ============================================================================ // Dataset Summary @@ -252,128 +236,8 @@ pub(crate) struct DatasetMetadata { pub project_id: String, } -/// Request body for dataset registration. -#[derive(Debug, Clone, Serialize)] -pub(crate) struct DatasetRegisterRequest { - pub project_name: String, - pub org_name: String, - pub dataset_name: String, - #[serde(skip_serializing_if = "Option::is_none")] - pub description: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub metadata: Option>, -} - -/// Response from dataset registration. -#[derive(Debug, Clone, Deserialize)] -pub(crate) struct DatasetRegisterResponse { - pub project: DatasetProjectInfo, - pub dataset: DatasetInfo, -} - -/// Project info from registration response. -#[derive(Debug, Clone, Deserialize)] -pub(crate) struct DatasetProjectInfo { - pub id: String, - #[allow(dead_code)] - pub name: String, -} - -/// Dataset info from registration response. -#[derive(Debug, Clone, Deserialize)] -pub(crate) struct DatasetInfo { - pub id: String, - #[allow(dead_code)] - pub name: String, -} - -// ============================================================================ -// BTQL Query Types (Internal) -// ============================================================================ - -/// BTQL query for fetching dataset records. -#[derive(Debug, Clone, Serialize)] -pub(crate) struct BTQLQuery { - pub query: BTQLQueryInner, -} - -/// Inner BTQL query structure. -#[derive(Debug, Clone, Serialize)] -pub(crate) struct BTQLQueryInner { - /// Dataset reference: dataset('id') - pub from: String, - /// Maximum records per page. - #[serde(skip_serializing_if = "Option::is_none")] - pub limit: Option, - /// Pagination cursor token. - #[serde(skip_serializing_if = "Option::is_none")] - pub cursor: Option, -} - -/// Response from BTQL query. -#[derive(Debug, Clone, Deserialize)] -pub(crate) struct BTQLResponse { - /// Records in JSONL format. - #[serde(default)] - pub data: Vec, - /// Cursor for next page (None if last page). - pub cursor: Option, -} - -// ============================================================================ -// Dataset Summary API Types (Internal) -// ============================================================================ - -/// Response from dataset summary API. -#[derive(Debug, Clone, Deserialize)] -pub(crate) struct DatasetSummaryResponse { - pub project_name: String, - pub dataset_name: String, - pub project_id: String, - pub dataset_id: String, - pub project_url: Option, - pub dataset_url: Option, - pub data_summary: Option, -} - -/// Data summary statistics. -#[derive(Debug, Clone, Deserialize)] -pub(crate) struct DataSummaryStats { - pub total_records: Option, -} - -// ============================================================================ -// Traits for API Operations -// ============================================================================ - -/// Trait for registering datasets. -#[async_trait] -pub(crate) trait DatasetRegistrar: Send + Sync { - /// Register a dataset with the API. - async fn register_dataset( - &self, - token: &str, - request: DatasetRegisterRequest, - ) -> Result; -} - -/// Trait for fetching dataset records via BTQL. -#[async_trait] -pub(crate) trait DatasetFetcher: Send + Sync { - /// Fetch dataset records via BTQL query. - async fn fetch_dataset_records(&self, token: &str, query: BTQLQuery) -> Result; -} - -/// Trait for fetching dataset summary. -#[async_trait] -pub(crate) trait DatasetSummarizer: Send + Sync { - /// Fetch dataset summary. - async fn fetch_dataset_summary( - &self, - token: &str, - dataset_id: &str, - ) -> Result; -} +use crate::api::btql::{BTQLQuery, BTQLQueryInner}; +use crate::api::datasets::DatasetRegisterRequest; // ============================================================================ // DatasetBuilder Error Types @@ -409,26 +273,22 @@ impl std::error::Error for DatasetBuilderError {} // ============================================================================ /// Builder for creating datasets with lazy registration. -#[allow(private_bounds)] -pub struct DatasetBuilder -{ - submitter: Arc, - token: String, - org_name: String, +pub struct DatasetBuilder { + api: Arc, + lifecycle: Arc, + scope: CredentialScope, project_name: Option, dataset_name: Option, description: Option, metadata: Option>, } -impl Clone - for DatasetBuilder -{ +impl Clone for DatasetBuilder { fn clone(&self) -> Self { Self { - submitter: Arc::clone(&self.submitter), - token: self.token.clone(), - org_name: self.org_name.clone(), + api: self.api.clone(), + lifecycle: self.lifecycle.clone(), + scope: self.scope.clone(), project_name: self.project_name.clone(), dataset_name: self.dataset_name.clone(), description: self.description.clone(), @@ -437,20 +297,17 @@ impl C } } -#[allow(private_bounds)] -impl - DatasetBuilder -{ +impl DatasetBuilder { /// Create a new DatasetBuilder (internal use). pub(crate) fn new( - submitter: Arc, - token: impl Into, - org_name: impl Into, + api: Arc, + lifecycle: Arc, + scope: CredentialScope, ) -> Self { Self { - submitter, - token: token.into(), - org_name: org_name.into(), + api, + lifecycle, + scope, project_name: None, dataset_name: None, description: None, @@ -491,7 +348,7 @@ impl std::result::Result, DatasetBuilderError> { + pub fn build(self) -> std::result::Result { let project_name = self .project_name .ok_or(DatasetBuilderError::MissingProjectName)?; @@ -500,15 +357,14 @@ impl { - submitter: Arc, - token: String, - org_name: String, +pub struct Dataset { + api: Arc, + lifecycle: Arc, + scope: CredentialScope, project_name: String, dataset_name: String, description: Option, metadata: Option>, - lazy_metadata: OnceCell>, - pending_rows: Arc>>, + lazy_metadata: Arc>, } -impl fmt::Debug - for Dataset -{ +impl Clone for Dataset { + fn clone(&self) -> Self { + Self { + api: self.api.clone(), + lifecycle: self.lifecycle.clone(), + scope: self.scope.clone(), + project_name: self.project_name.clone(), + dataset_name: self.dataset_name.clone(), + description: self.description.clone(), + metadata: self.metadata.clone(), + lazy_metadata: self.lazy_metadata.clone(), + } + } +} + +impl fmt::Debug for Dataset { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("Dataset") .field("project_name", &self.project_name) @@ -545,44 +412,25 @@ impl f } } -/// A pending dataset row waiting to be submitted. -struct PendingDatasetRow { - id: String, - input: Option, - expected: Option, - metadata: Option>, - tags: Option>, - is_merge: bool, - is_delete: bool, -} - -#[allow(private_bounds)] -impl - Dataset -{ +impl Dataset { /// Ensure the dataset is registered, returning cached metadata. async fn ensure_registered(&self) -> Result<&DatasetMetadata> { self.lazy_metadata - .get_or_init(|| async { + .get_or_try_init(|| async { let request = DatasetRegisterRequest { project_name: self.project_name.clone(), - org_name: self.org_name.clone(), + org_name: self.scope.org_name().unwrap_or_default().to_owned(), dataset_name: self.dataset_name.clone(), description: self.description.clone(), metadata: self.metadata.clone(), }; - - match self.submitter.register_dataset(&self.token, request).await { - Ok(response) => Ok(DatasetMetadata { - dataset_id: response.dataset.id, - project_id: response.project.id, - }), - Err(e) => Err(e.to_string()), - } + let response = self.api.register_dataset(&self.scope, request).await?; + Ok(DatasetMetadata { + dataset_id: response.dataset.id, + project_id: response.project.id, + }) }) .await - .as_ref() - .map_err(|e| crate::error::BraintrustError::InvalidConfig(e.clone())) } /// Get the dataset ID (triggers registration if needed). @@ -600,82 +448,73 @@ impl String { + /// Insert a record and return its generated ID after the API acknowledges it. + pub async fn insert(&self, record: DatasetInsert) -> Result { + let dataset_id = self.ensure_registered().await?.dataset_id.clone(); let id = Uuid::new_v4().to_string(); - - let row = PendingDatasetRow { - id: id.clone(), - input: record.input, - expected: record.expected, - metadata: record.metadata, - tags: record.tags, - is_merge: false, - is_delete: false, - }; - - self.pending_rows.lock().await.push(row); - id - } - - /// Update an existing record by ID (merge semantics). - /// - /// Unspecified fields are preserved. - pub async fn update(&self, id: &str, record: DatasetInsert) { - let row = PendingDatasetRow { - id: id.to_string(), - input: record.input, - expected: record.expected, - metadata: record.metadata, - tags: record.tags, - is_merge: true, - is_delete: false, - }; - - self.pending_rows.lock().await.push(row); + self.api + .write_dataset_record( + &self.scope, + DatasetWrite::Insert { + dataset_id, + id: id.clone(), + fields: record.into(), + }, + ) + .await?; + Ok(id) + } + + /// Merge supplied fields into a record, returning after API acknowledgement. + pub async fn update(&self, id: &str, record: DatasetInsert) -> Result<()> { + let dataset_id = self.ensure_registered().await?.dataset_id.clone(); + self.api + .write_dataset_record( + &self.scope, + DatasetWrite::Update { + dataset_id, + id: id.to_owned(), + fields: record.into(), + }, + ) + .await } - /// Delete a record by ID. - pub async fn delete(&self, id: &str) { - let row = PendingDatasetRow { - id: id.to_string(), - input: None, - expected: None, - metadata: None, - tags: None, - is_merge: false, - is_delete: true, - }; - - self.pending_rows.lock().await.push(row); + /// Delete a record, returning after API acknowledgement. + pub async fn delete(&self, id: &str) -> Result<()> { + let dataset_id = self.ensure_registered().await?.dataset_id.clone(); + self.api + .write_dataset_record( + &self.scope, + DatasetWrite::Delete { + dataset_id, + id: id.to_owned(), + }, + ) + .await } /// Fetch records with pagination. /// /// Returns an iterator that fetches records in batches. - pub async fn fetch(&self, batch_size: Option) -> Result> { + pub async fn fetch(&self, batch_size: Option) -> Result { let metadata = self.ensure_registered().await?; Ok(DatasetIterator::new( - Arc::clone(&self.submitter), - self.token.clone(), + self.api.clone(), + self.scope.clone(), metadata.dataset_id.clone(), - batch_size.unwrap_or(1000), + batch_size.unwrap_or(1000).max(1), )) } /// Get dataset summary statistics. pub async fn summarize(&self) -> Result { - // Flush pending operations first - self.flush().await?; - let metadata = self.ensure_registered().await?; // Fetch summary from API let response = self - .submitter - .fetch_dataset_summary(&self.token, &metadata.dataset_id) + .api + .fetch_dataset_summary(&self.scope, &metadata.dataset_id) .await?; Ok(DatasetSummary { @@ -689,70 +528,10 @@ impl Result<()> { - let metadata = self.ensure_registered().await?; - let rows = { - let mut pending = self.pending_rows.lock().await; - std::mem::take(&mut *pending) - }; - - if rows.is_empty() { - return Ok(()); - } - - // Submit each row - for row in rows { - let now = SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_secs_f64(); - - let payload = SpanPayload { - row_id: row.id.clone(), - span_id: row.id, - is_merge: row.is_merge, - span_components: None, - org_id: String::new(), - org_name: Some(self.org_name.clone()), - project_name: Some(self.project_name.clone()), - input: row.input, - output: None, - expected: row.expected, - error: None, - scores: None, - metadata: row.metadata, - metrics: None, - tags: row.tags, - context: Some(serde_json::json!({ - "calendar_time": now, - })), - span_attributes: if row.is_delete { - Some(crate::types::SpanAttributes { - name: None, - span_type: None, - purpose: None, - extra: { - let mut map = std::collections::HashMap::new(); - map.insert("_object_delete".to_string(), Value::Bool(true)); - map - }, - }) - } else { - None - }, - extra: std::collections::HashMap::new(), - }; - - let parent_info = ParentSpanInfo::Dataset { - object_id: metadata.dataset_id.clone(), - }; - - self.submitter - .submit(self.token.clone(), payload, Some(parent_info)); - } - - Ok(()) + /// Flush the client's shared trace queue. Dataset mutations are already + /// acknowledged when their individual operation completes. + pub fn flush(&self) -> crate::Completion { + self.lifecycle.flush() } } @@ -761,10 +540,20 @@ impl { - fetcher: Arc, - token: String, +/// +/// ```no_run +/// use braintrust_sdk_rust::{BraintrustError, Dataset}; +/// # async fn read(dataset: &Dataset) -> Result<(), BraintrustError> { +/// let mut records = dataset.fetch(Some(100)).await?; +/// while let Some(record) = records.next().await? { +/// println!("{}", record.id); +/// } +/// # Ok(()) +/// # } +/// ``` +pub struct DatasetIterator { + fetcher: Arc, + scope: CredentialScope, dataset_id: String, batch_size: usize, cursor: Option, @@ -772,12 +561,16 @@ pub struct DatasetIterator { done: bool, } -#[allow(private_bounds)] -impl DatasetIterator { - fn new(fetcher: Arc, token: String, dataset_id: String, batch_size: usize) -> Self { +impl DatasetIterator { + fn new( + fetcher: Arc, + scope: CredentialScope, + dataset_id: String, + batch_size: usize, + ) -> Self { Self { fetcher, - token, + scope, dataset_id, batch_size, cursor: None, @@ -787,53 +580,46 @@ impl DatasetIterator { } /// Get the next record, fetching a new batch if needed. - pub async fn next(&mut self) -> Option { - // Return buffered record if available - if let Some(record) = self.buffer.pop_front() { - return Some(record); - } - - // If we've exhausted all pages, return None - if self.done { - return None; - } - - // Fetch next batch - let query = BTQLQuery { - query: BTQLQueryInner { - from: format!("dataset('{}')", self.dataset_id), - limit: Some(self.batch_size), - cursor: self.cursor.take(), - }, - }; - - match self.fetcher.fetch_dataset_records(&self.token, query).await { - Ok(response) => { - // Update cursor for next page - self.cursor = response.cursor; - self.done = self.cursor.is_none(); - - // Buffer the records - self.buffer.extend(response.data); - - // Return first record - self.buffer.pop_front() + pub async fn next(&mut self) -> Result> { + loop { + if let Some(record) = self.buffer.pop_front() { + return Ok(Some(record)); } - Err(e) => { - tracing::warn!("failed to fetch dataset records: {}", e); - self.done = true; - None + if self.done { + return Ok(None); + } + let previous_cursor = self.cursor.clone(); + let query = BTQLQuery { + query: BTQLQueryInner { + from: format!("dataset('{}')", self.dataset_id), + limit: Some(self.batch_size), + cursor: previous_cursor.clone(), + }, + }; + // A failed request leaves the cursor and EOF state untouched. + let response = self + .fetcher + .fetch_dataset_records(&self.scope, query) + .await?; + if response.cursor.is_some() && response.cursor == previous_cursor { + return Err(crate::BraintrustError::Background( + "dataset pagination cursor did not advance".into(), + )); } + self.cursor = response.cursor; + self.done = self.cursor.is_none(); + self.buffer.extend(response.data); + // Empty pages with a continuation are not end-of-data. } } - /// Collect all remaining records into a Vec. - pub async fn collect(mut self) -> Vec { + /// Collect all remaining records, preserving any page failure as an error. + pub async fn collect(mut self) -> Result> { let mut records = Vec::new(); - while let Some(record) = self.next().await { + while let Some(record) = self.next().await? { records.push(record); } - records + Ok(records) } } diff --git a/src/delivery.rs b/src/delivery.rs new file mode 100644 index 0000000..25a0172 --- /dev/null +++ b/src/delivery.rs @@ -0,0 +1,87 @@ +//! Outcomes of queued ingestion and its lifecycle operations. +use std::sync::Arc; +use thiserror::Error; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)] +pub enum EnqueueError { + #[error("delivery admission capacity is full")] + Full, + #[error("delivery queue is closed")] + Closed, +} +#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)] +pub enum ExecutionError { + #[error("delivery worker stopped before completing accepted work")] + WorkerLost, + #[error("delivery shutdown timed out")] + TimedOut, +} +/// A terminal preparation, encoding or upload failure shared by a logical batch. +#[derive(Debug, Error)] +#[error("log delivery failed: {source}")] +pub struct DeliveryError { + #[source] + pub(crate) source: anyhow::Error, +} +impl DeliveryError { + /// Inspect the original preparation, encoding or transport error. + pub fn cause(&self) -> &anyhow::Error { + &self.source + } +} +#[derive(Debug, Clone, Error)] +pub enum WriteError { + #[error(transparent)] + Admission(#[from] EnqueueError), + #[error(transparent)] + Execution(#[from] ExecutionError), + #[error(transparent)] + Delivery(Arc), +} +/// Acknowledgment that an operation completed successfully. +/// For admitted writes, the entire logical batch succeeded. +/// Backend query visibility is separate from delivery acknowledgment. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DeliveryReceipt; +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct FlushReport { + /// Accepted operations delivered through this flush's cutoff, before merging. + pub delivered: u64, +} +#[derive(Debug, Clone, Error)] +#[error("{failed} accepted writes failed: {cause}")] +pub struct FlushFailure { + pub failed: u64, + pub cause: Arc, +} +#[derive(Debug, Clone, Error)] +pub enum FlushError { + #[error(transparent)] + Admission(#[from] EnqueueError), + #[error(transparent)] + Execution(#[from] ExecutionError), + #[error(transparent)] + Delivery(FlushFailure), +} +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct ShutdownReport { + pub delivered: u64, +} +#[derive(Debug, Clone, Error)] +pub enum ShutdownError { + #[error(transparent)] + Admission(#[from] EnqueueError), + #[error(transparent)] + Execution(#[from] ExecutionError), + #[error(transparent)] + Delivery(FlushFailure), +} +impl From for ShutdownError { + fn from(error: FlushError) -> Self { + match error { + FlushError::Admission(error) => Self::Admission(error), + FlushError::Execution(error) => Self::Execution(error), + FlushError::Delivery(error) => Self::Delivery(error), + } + } +} diff --git a/src/error.rs b/src/error.rs index 9b6729e..7d2c366 100644 --- a/src/error.rs +++ b/src/error.rs @@ -19,3 +19,56 @@ pub enum BraintrustError { } pub type Result = std::result::Result; + +// Compatibility for resource APIs that still return BraintrustError while the +// public operation surfaces migrate to typed completion errors. +impl From for BraintrustError { + fn from(error: crate::FlushError) -> Self { + if let crate::FlushError::Delivery(failure) = &error { + if let Some(Self::Api { status, message }) = + failure.cause.cause().downcast_ref::() + { + return Self::Api { + status: *status, + message: message.clone(), + }; + } + } + Self::Background(error.to_string()) + } +} + +impl From for BraintrustError { + fn from(error: crate::WriteError) -> Self { + if let crate::WriteError::Delivery(cause) = &error { + if let Some(Self::Api { status, message }) = cause.cause().downcast_ref::() { + return Self::Api { + status: *status, + message: message.clone(), + }; + } + } + Self::Background(error.to_string()) + } +} + +impl From for BraintrustError { + fn from(error: crate::api::Error) -> Self { + match 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 { + status, + message: body, + }, + crate::api::Error::Decode { + operation, + status, + source, + } => BraintrustError::Api { + status, + message: format!("{operation}: invalid response: {source}"), + }, + } + } +} diff --git a/src/eval/runner.rs b/src/eval/runner.rs index 96db594..cecb373 100644 --- a/src/eval/runner.rs +++ b/src/eval/runner.rs @@ -163,7 +163,7 @@ where } // End root span - root_span.end(); + root_span.end().await?; // Compute statistics let score_stats = EvalSummary::::compute_stats(&results); @@ -200,7 +200,7 @@ where case: Case, task: &dyn Task, scorers: &[Box>], - _root_span: &SpanHandle, + _root_span: &SpanHandle, client: &BraintrustClient, case_num: usize, quiet: bool, @@ -248,7 +248,7 @@ where .build() .unwrap(), ); - case_span.end(); + case_span.end().await?; return Ok(EvalResult { input: case.input.clone(), @@ -322,7 +322,7 @@ where } case_span.log(final_log.build().unwrap()); - case_span.end(); + case_span.end().await?; if !quiet { let score_str = all_scores diff --git a/src/eval/types.rs b/src/eval/types.rs index ff0ed05..acbaf0e 100644 --- a/src/eval/types.rs +++ b/src/eval/types.rs @@ -2,7 +2,7 @@ use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use std::sync::{Arc, Mutex}; -use crate::{BraintrustClient, SpanHandle}; +use crate::SpanHandle; /// A single test case for evaluation #[derive(Debug, Clone, Serialize, Deserialize)] @@ -203,7 +203,7 @@ impl EvalSummary { /// Hooks provided to tasks for instrumentation and metadata pub struct TaskHooks<'a, I, O> { /// Span handle for logging - pub span: &'a SpanHandle, + pub span: &'a SpanHandle, /// Expected output (if available) pub expected: Option<&'a O>, /// Mutable metadata map for the task to add context @@ -215,7 +215,7 @@ pub struct TaskHooks<'a, I, O> { } impl<'a, I, O> TaskHooks<'a, I, O> { - pub fn new(span: &'a SpanHandle, expected: Option<&'a O>) -> Self { + pub fn new(span: &'a SpanHandle, expected: Option<&'a O>) -> Self { Self { span, expected, diff --git a/src/experiments/api/comparison.rs b/src/experiments/api/comparison.rs deleted file mode 100644 index 5317b91..0000000 --- a/src/experiments/api/comparison.rs +++ /dev/null @@ -1,37 +0,0 @@ -//! Internal API types for experiment comparison. - -use std::collections::HashMap; - -use serde::Deserialize; - -/// Response from experiment comparison API. -#[derive(Debug, Clone, Deserialize)] -#[allow(dead_code)] -pub(crate) struct ExperimentComparisonResponse { - #[serde(default)] - pub scores: HashMap, - #[serde(default)] - pub metrics: HashMap, - pub base_experiment_name: Option, -} - -/// Score data from comparison API. -#[derive(Debug, Clone, Deserialize)] -#[allow(dead_code)] -pub(crate) struct ComparisonScoreData { - pub score: f64, - pub diff: Option, - pub improvements: Option, - pub regressions: Option, -} - -/// Metric data from comparison API. -#[derive(Debug, Clone, Deserialize)] -#[allow(dead_code)] -pub(crate) struct ComparisonMetricData { - pub metric: f64, - pub unit: Option, - pub diff: Option, - pub improvements: Option, - pub regressions: Option, -} diff --git a/src/experiments/api/mod.rs b/src/experiments/api/mod.rs deleted file mode 100644 index 92d5af2..0000000 --- a/src/experiments/api/mod.rs +++ /dev/null @@ -1,15 +0,0 @@ -//! Internal API contracts and traits. -//! -//! This module contains internal API types and traits used for communicating -//! with the Braintrust backend. These are not part of the public API. - -pub(crate) mod comparison; -pub(crate) mod registration; -pub(crate) mod traits; - -pub(crate) use comparison::ExperimentComparisonResponse; -pub(crate) use registration::{ - BaseExperimentRequest, BaseExperimentResponse, ExperimentRegisterRequest, - ExperimentRegisterResponse, -}; -pub(crate) use traits::{BaseExperimentFetcher, ExperimentComparisonFetcher, ExperimentRegistrar}; diff --git a/src/experiments/api/registration.rs b/src/experiments/api/registration.rs deleted file mode 100644 index 28d71f9..0000000 --- a/src/experiments/api/registration.rs +++ /dev/null @@ -1,62 +0,0 @@ -//! Internal API types for experiment registration. - -use serde::{Deserialize, Serialize}; -use serde_json::{Map, Value}; - -use crate::experiments::metadata::RepoInfo; - -/// Request body for experiment registration. -#[derive(Debug, Clone, Serialize)] -pub(crate) struct ExperimentRegisterRequest { - pub project_name: String, - pub org_id: String, - #[serde(skip_serializing_if = "Option::is_none")] - pub experiment_name: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub description: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub base_experiment: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub metadata: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub public: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub update: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub repo_info: Option, -} - -/// Response from experiment registration. -#[derive(Debug, Clone, Deserialize)] -pub(crate) struct ExperimentRegisterResponse { - pub project: ProjectInfo, - pub experiment: ExperimentInfo, -} - -/// Project information from registration response. -#[derive(Debug, Clone, Deserialize)] -#[allow(dead_code)] -pub(crate) struct ProjectInfo { - pub id: String, - pub name: String, -} - -/// Experiment information from registration response. -#[derive(Debug, Clone, Deserialize)] -pub(crate) struct ExperimentInfo { - pub id: String, - pub name: String, -} - -/// Request body for fetching base experiment. -#[derive(Debug, Clone, Serialize)] -pub(crate) struct BaseExperimentRequest { - pub id: String, -} - -/// Response from base experiment fetch. -#[derive(Debug, Clone, Deserialize)] -pub(crate) struct BaseExperimentResponse { - pub base_exp_id: Option, - pub base_exp_name: Option, -} diff --git a/src/experiments/api/traits.rs b/src/experiments/api/traits.rs deleted file mode 100644 index 54704f2..0000000 --- a/src/experiments/api/traits.rs +++ /dev/null @@ -1,39 +0,0 @@ -//! Internal API trait definitions for experiments. - -use crate::error::Result; -use crate::experiments::api::comparison::ExperimentComparisonResponse; -use crate::experiments::api::registration::{ - ExperimentRegisterRequest, ExperimentRegisterResponse, -}; -use crate::experiments::metadata::BaseExperimentInfo; - -/// Trait for registering experiments with the Braintrust API. -#[async_trait::async_trait] -pub(crate) trait ExperimentRegistrar: Send + Sync { - async fn register_experiment( - &self, - token: &str, - request: ExperimentRegisterRequest, - ) -> Result; -} - -/// Trait for fetching base experiment information. -#[async_trait::async_trait] -pub(crate) trait BaseExperimentFetcher: Send + Sync { - async fn fetch_base_experiment( - &self, - token: &str, - experiment_id: &str, - ) -> Result>; -} - -/// Trait for fetching experiment comparison data. -#[async_trait::async_trait] -pub(crate) trait ExperimentComparisonFetcher: Send + Sync { - async fn fetch_experiment_comparison( - &self, - token: &str, - experiment_id: &str, - base_experiment_id: Option<&str>, - ) -> Result; -} diff --git a/src/experiments/builder.rs b/src/experiments/builder.rs index 32ffa69..fe97b25 100644 --- a/src/experiments/builder.rs +++ b/src/experiments/builder.rs @@ -1,15 +1,14 @@ //! Builder for creating experiments. use std::fmt; -use std::sync::atomic::{AtomicBool, AtomicU64}; use std::sync::Arc; use serde_json::{Map, Value}; -use tokio::sync::{Mutex, OnceCell}; -use crate::experiments::experiment::Experiment; +use crate::experiments::experiment::{Experiment, ExperimentState}; use crate::experiments::metadata::RepoInfo; -use crate::span::SpanSubmitter; +use crate::internal::ports::{ExperimentApi, SpanSubmitter}; +use crate::internal::session::CredentialScope; /// Error type for ExperimentBuilder validation. #[derive(Debug, Clone, PartialEq)] @@ -33,15 +32,12 @@ impl std::error::Error for ExperimentBuilderError {} /// Builder for creating experiments with lazy registration. /// -/// Note: This builder uses `experiment_builder_with_credentials` naming to indicate -/// that API credentials must be passed explicitly. This will be deprecated in favor -/// of `experiment_builder()` once the login-flow branch lands. -#[allow(private_bounds)] -pub struct ExperimentBuilder { - submitter: Arc, - token: String, - org_id: String, - org_name: Option, +/// Registration is shared by clones of the resulting handle. Explicit +/// credentials belong to the resource and do not mutate the client's session. +pub struct ExperimentBuilder { + api: Arc, + submitter: Arc, + scope: CredentialScope, project_name: Option, experiment_name: Option, description: Option, @@ -52,13 +48,12 @@ pub struct ExperimentBuilder { repo_info: Option, } -impl Clone for ExperimentBuilder { +impl Clone for ExperimentBuilder { fn clone(&self) -> Self { Self { - submitter: Arc::clone(&self.submitter), - token: self.token.clone(), - org_id: self.org_id.clone(), - org_name: self.org_name.clone(), + api: self.api.clone(), + submitter: self.submitter.clone(), + scope: self.scope.clone(), project_name: self.project_name.clone(), experiment_name: self.experiment_name.clone(), description: self.description.clone(), @@ -71,19 +66,17 @@ impl Clone for ExperimentBuilder { } } -#[allow(private_bounds)] -impl ExperimentBuilder { +impl ExperimentBuilder { /// Create a new ExperimentBuilder (internal use). pub(crate) fn new( - submitter: Arc, - token: impl Into, - org_id: impl Into, + api: Arc, + submitter: Arc, + scope: CredentialScope, ) -> Self { Self { + api, submitter, - token: token.into(), - org_id: org_id.into(), - org_name: None, + scope, project_name: None, experiment_name: None, description: None, @@ -97,7 +90,7 @@ impl ExperimentBuilder { /// Set the organization name. pub fn org_name(mut self, org_name: impl Into) -> Self { - self.org_name = Some(org_name.into()); + self.scope = self.scope.with_org_name(org_name.into()); self } @@ -164,16 +157,15 @@ impl ExperimentBuilder { } /// Build the experiment. Returns an error if required fields are missing. - pub fn build(self) -> std::result::Result, ExperimentBuilderError> { + pub fn build(self) -> std::result::Result { let project_name = self .project_name .ok_or(ExperimentBuilderError::MissingProjectName)?; Ok(Experiment { + api: self.api, submitter: self.submitter, - token: self.token, - org_id: self.org_id, - org_name: self.org_name, + scope: self.scope, project_name, experiment_name: self.experiment_name, description: self.description, @@ -182,10 +174,7 @@ impl ExperimentBuilder { public: self.public, update: self.update, repo_info: self.repo_info, - lazy_metadata: OnceCell::new(), - last_start_time: AtomicU64::new(0), - called_start_span: AtomicBool::new(false), - pending_flushes: Arc::new(Mutex::new(Vec::new())), + state: Arc::new(ExperimentState::default()), }) } } diff --git a/src/experiments/experiment.rs b/src/experiments/experiment.rs index b5142fd..7bdcf2d 100644 --- a/src/experiments/experiment.rs +++ b/src/experiments/experiment.rs @@ -6,14 +6,10 @@ use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; use serde_json::{Map, Value}; -use tokio::sync::{Mutex, OnceCell}; -use uuid::Uuid; +use tokio::sync::OnceCell; +use crate::api::experiments::ExperimentRegisterRequest; use crate::error::Result; -use crate::experiments::api::{ - BaseExperimentFetcher, ExperimentComparisonFetcher, ExperimentRegisterRequest, - ExperimentRegistrar, -}; use crate::experiments::event_log::ExperimentLog; use crate::experiments::feedback::Feedback; use crate::experiments::metadata::{ @@ -21,7 +17,9 @@ use crate::experiments::metadata::{ }; use crate::experiments::span_builder::ExperimentSpanBuilder; use crate::experiments::summary::{ExperimentSummary, MetricSummary, ScoreSummary}; -use crate::span::{SpanBuilder, SpanHandle, SpanLog, SpanSubmitter}; +use crate::internal::ports::{ExperimentApi, SpanSubmitter}; +use crate::internal::session::CredentialScope; +use crate::span::{SpanBuilder, SpanHandle, SpanLog}; use crate::types::{ParentSpanInfo, SpanAttributes, SpanPayload, SpanType}; /// Metadata from experiment registration, cached after lazy initialization. @@ -35,12 +33,10 @@ pub(crate) struct ExperimentMetadata { /// Handle to a Braintrust experiment. /// /// Use this to log evaluation events, create spans, log feedback, and get summaries. -#[allow(private_bounds)] -pub struct Experiment { - pub(crate) submitter: Arc, - pub(crate) token: String, - pub(crate) org_id: String, - pub(crate) org_name: Option, +pub struct Experiment { + pub(crate) api: Arc, + pub(crate) submitter: Arc, + pub(crate) scope: CredentialScope, pub(crate) project_name: String, pub(crate) experiment_name: Option, pub(crate) description: Option, @@ -49,22 +45,24 @@ pub struct Experiment { pub(crate) public: Option, pub(crate) update: Option, pub(crate) repo_info: Option, - pub(crate) lazy_metadata: OnceCell, - /// Last start time for sequential log() calls (stored as bits). - pub(crate) last_start_time: AtomicU64, - /// Whether start_span() has been called (prevents mixing with log()). - pub(crate) called_start_span: AtomicBool, - /// Pending span handles that need to be flushed. - pub(crate) pending_flushes: Arc>>>, + pub(crate) state: Arc, +} + +/// Registration and live experiment state have the same lifetime as the handle +/// identity. Cloning a handle never resets initialization or recording state. +#[derive(Default)] +pub(crate) struct ExperimentState { + metadata: OnceCell, + last_start_time: AtomicU64, + called_start_span: AtomicBool, } -impl Clone for Experiment { +impl Clone for Experiment { fn clone(&self) -> Self { Self { - submitter: Arc::clone(&self.submitter), - token: self.token.clone(), - org_id: self.org_id.clone(), - org_name: self.org_name.clone(), + api: self.api.clone(), + submitter: self.submitter.clone(), + scope: self.scope.clone(), project_name: self.project_name.clone(), experiment_name: self.experiment_name.clone(), description: self.description.clone(), @@ -73,44 +71,30 @@ impl Clone for Experiment { repo_info: self.repo_info.clone(), metadata: self.metadata.clone(), public: self.public, - lazy_metadata: OnceCell::new(), - last_start_time: AtomicU64::new(self.last_start_time.load(Ordering::Acquire)), - called_start_span: AtomicBool::new(self.called_start_span.load(Ordering::Acquire)), - pending_flushes: Arc::clone(&self.pending_flushes), + state: self.state.clone(), } } } -#[allow(private_bounds)] -impl< - S: SpanSubmitter - + ExperimentRegistrar - + BaseExperimentFetcher - + ExperimentComparisonFetcher - + 'static, - > Experiment -{ +impl Experiment { /// Log an event to the experiment, creating and ending a span internally. /// /// Returns the span ID which can be used with `log_feedback()`. /// - /// # Panics - /// - /// Panics if `start_span()` has already been called on this experiment, - /// as mixing `log()` and `start_span()` is not supported. - pub async fn log(&self, event: ExperimentLog) -> String { - if self.called_start_span.load(Ordering::Acquire) { - panic!( - "Cannot use log() after start_span() has been called. \ - Use start_span() consistently or log() consistently, but not both." - ); + /// Returns an error if registration/delivery fails, or if start_span() has + /// already selected the span-recording interface for this experiment. + pub async fn log(&self, event: ExperimentLog) -> Result { + if self.state.called_start_span.load(Ordering::Acquire) { + return Err(crate::BraintrustError::InvalidConfig( + "cannot use log() after start_span(); use the span to log events".into(), + )); } // Get or create experiment registration - let metadata = self.ensure_registered().await; + let metadata = self.ensure_registered().await?; // Create a span for this log event - let start_time = f64::from_bits(self.last_start_time.load(Ordering::Acquire)); + let start_time = f64::from_bits(self.state.last_start_time.load(Ordering::Acquire)); let start_time = if start_time == 0.0 { SystemTime::now() .duration_since(UNIX_EPOCH) @@ -120,44 +104,46 @@ impl< start_time }; - let span = self.create_internal_span(&metadata, start_time).await; + let span = self.create_internal_span(metadata, start_time); // Get the row_id before logging - it's stable from span creation let row_id = span.row_id().to_string(); - // Convert ExperimentLog to SpanLog and log - let span_log = event.into_span_log(); - span.log(span_log); - - // Record end time for next log() call let end_time = SystemTime::now() .duration_since(UNIX_EPOCH) - .map(|d| d.as_secs_f64()) + .map(|duration| duration.as_secs_f64()) .unwrap_or(0.0); - self.last_start_time + let mut span_log = event.into_span_log(); + span_log + .metrics + .get_or_insert_with(HashMap::new) + .entry("end".into()) + .or_insert(end_time); + // A single complete event snapshot owns the delivery obligation. + span.log(span_log).await?; + + self.state + .last_start_time .store(end_time.to_bits(), Ordering::Release); - // Track for flush - self.pending_flushes.lock().await.push(span); - - row_id + Ok(row_id) } /// Start a new root span for this experiment. /// /// Use this for more control over span lifecycle. Once called, /// you cannot use `log()` on this experiment. - pub async fn start_span(&self) -> ExperimentSpanBuilder<'_, S> { - self.called_start_span.store(true, Ordering::Release); - let metadata = self.ensure_registered().await; - ExperimentSpanBuilder::new(self, metadata) + pub async fn start_span(&self) -> Result { + let metadata = self.ensure_registered().await?; + self.state.called_start_span.store(true, Ordering::Release); + Ok(ExperimentSpanBuilder::new(self, metadata.clone())) } /// Log feedback on an existing span/record. /// /// Use this to add scores, comments, or other feedback after the initial log. - pub async fn log_feedback(&self, id: &str, feedback: Feedback) { - let metadata = self.ensure_registered().await; + pub async fn log_feedback(&self, id: &str, feedback: Feedback) -> Result<()> { + let metadata = self.ensure_registered().await?; // Create a merge payload for the existing span let mut scores = HashMap::new(); @@ -185,8 +171,8 @@ impl< span_id: id.to_string(), // Use same ID for merge is_merge: true, span_components: None, - org_id: self.org_id.clone(), - org_name: self.org_name.clone(), + org_id: self.scope.org_id().to_owned(), + org_name: self.scope.org_name().map(str::to_owned), project_name: None, input: None, output: None, @@ -206,7 +192,9 @@ impl< }; self.submitter - .submit(self.token.clone(), payload, Some(parent_info)); + .submit(self.scope.clone(), payload, Some(parent_info), None) + .await?; + Ok(()) } /// Get a summary of the experiment's performance. @@ -217,21 +205,16 @@ impl< // Flush all pending data first self.flush().await?; - let metadata = self.ensure_registered().await; + let metadata = self.ensure_registered().await?; // Try to get base experiment for comparison - let base_exp_id = self - .fetch_base_experiment() - .await - .ok() - .flatten() - .map(|b| b.id); + let base_exp_id = self.fetch_base_experiment().await?.map(|base| base.id); // Fetch comparison data from API let comparison = self - .submitter + .api .fetch_experiment_comparison( - &self.token, + &self.scope, &metadata.experiment_id, base_exp_id.as_deref(), ) @@ -297,21 +280,19 @@ impl< /// Flush all pending experiment data to Braintrust. pub async fn flush(&self) -> Result<()> { - let mut pending = self.pending_flushes.lock().await; - for span in pending.drain(..) { - span.flush().await?; - } + self.submitter.flush().await?; Ok(()) } /// Get the experiment ID (after registration). pub fn experiment_id(&self) -> Option { - self.lazy_metadata.get().map(|m| m.experiment_id.clone()) + self.state.metadata.get().map(|m| m.experiment_id.clone()) } /// Get the experiment name (after registration). pub fn experiment_name(&self) -> String { - self.lazy_metadata + self.state + .metadata .get() .map(|m| m.experiment_name.clone()) .unwrap_or_else(|| { @@ -322,18 +303,19 @@ impl< } /// Get the project containing this experiment. - pub async fn project(&self) -> ProjectMetadata { - let metadata = self.ensure_registered().await; - ProjectMetadata { + pub async fn project(&self) -> Result { + let metadata = self.ensure_registered().await?; + Ok(ProjectMetadata { id: metadata.project_id.clone(), name: self.project_name.clone(), - } + }) } /// Execute a callback within a traced span. /// - /// The span is automatically flushed when the callback completes (success or error). - /// This ensures all span data is sent to Braintrust before the method returns. + /// Finalizes the span and flushes active work when the callback completes. + /// Setup, final-write, and final-flush errors are returned. Individual writes expose their + /// delivery results through the completion tokens returned inside the callback. /// /// # Example /// @@ -343,32 +325,36 @@ impl< /// compute_result() /// }).await; /// ``` - pub async fn traced(&self, callback: F) -> R + pub async fn traced(&self, callback: F) -> Result where - F: FnOnce(SpanHandle) -> Fut, + F: FnOnce(SpanHandle) -> Fut, Fut: std::future::Future, { - let span = self.start_span().await.name("root").build(); + let span = self.start_span().await?.name("root").build(); let result = callback(span.clone()).await; - // Flush span data - ignore errors since this is best-effort - let _ = span.flush().await; - result + let end = span.end(); + let flush = span.flush(); + let ended = end.await; + let flushed = flush.await; + ended?; + flushed?; + Ok(result) } /// Update an existing span by ID, merging new data with existing data. /// /// This is useful for adding information to a span after it has been created, /// such as adding scores or metadata computed later. - pub async fn update_span(&self, id: &str, event: SpanLog) { - let metadata = self.ensure_registered().await; + pub async fn update_span(&self, id: &str, event: SpanLog) -> Result<()> { + let metadata = self.ensure_registered().await?; let payload = SpanPayload { row_id: id.to_string(), span_id: id.to_string(), is_merge: true, span_components: None, - org_id: self.org_id.clone(), - org_name: self.org_name.clone(), + org_id: self.scope.org_id().to_owned(), + org_name: self.scope.org_name().map(str::to_owned), project_name: None, input: event.input, output: event.output, @@ -393,16 +379,18 @@ impl< }; self.submitter - .submit(self.token.clone(), payload, Some(parent_info)); + .submit(self.scope.clone(), payload, Some(parent_info), None) + .await?; + Ok(()) } /// Fetch the base experiment used for comparison. /// /// Returns `None` if no base experiment has been configured or detected. pub async fn fetch_base_experiment(&self) -> Result> { - let metadata = self.ensure_registered().await; - self.submitter - .fetch_base_experiment(&self.token, &metadata.experiment_id) + let metadata = self.ensure_registered().await?; + self.api + .fetch_base_experiment(&self.scope, &metadata.experiment_id) .await } @@ -410,21 +398,22 @@ impl< /// /// Use this to pass experiment context to another process or service, /// allowing spans created there to be linked to this experiment. - pub async fn export(&self) -> ExportedExperiment { - let metadata = self.ensure_registered().await; - ExportedExperiment { + pub async fn export(&self) -> Result { + let metadata = self.ensure_registered().await?; + Ok(ExportedExperiment { object_type: "experiment".to_string(), object_id: metadata.experiment_id.clone(), - } + }) } /// Ensure the experiment is registered, returning cached metadata. - pub(crate) async fn ensure_registered(&self) -> ExperimentMetadata { - self.lazy_metadata - .get_or_init(|| async { + pub(crate) async fn ensure_registered(&self) -> Result<&ExperimentMetadata> { + self.state + .metadata + .get_or_try_init(|| async { let request = ExperimentRegisterRequest { project_name: self.project_name.clone(), - org_id: self.org_id.clone(), + org_id: self.scope.org_id().to_owned(), experiment_name: self.experiment_name.clone(), description: self.description.clone(), base_experiment: self.base_experiment.clone(), @@ -433,59 +422,28 @@ impl< update: self.update, repo_info: self.repo_info.clone(), }; - - match self - .submitter - .register_experiment(&self.token, request) - .await - { - Ok(response) => ExperimentMetadata { - project_id: response.project.id, - experiment_id: response.experiment.id, - experiment_name: response.experiment.name, - }, - Err(e) => { - // Log error but return placeholder to avoid blocking - tracing::warn!(error = %e, "experiment registration failed"); - ExperimentMetadata { - project_id: String::new(), - experiment_id: Uuid::new_v4().to_string(), - experiment_name: self - .experiment_name - .clone() - .unwrap_or_else(|| "unnamed".to_string()), - } - } - } + let response = self.api.register_experiment(&self.scope, request).await?; + Ok(ExperimentMetadata { + project_id: response.project.id, + experiment_id: response.experiment.id, + experiment_name: response.experiment.name, + }) }) .await - .clone() } /// Create an internal span for log() calls. - async fn create_internal_span( - &self, - metadata: &ExperimentMetadata, - start_time: f64, - ) -> SpanHandle { + fn create_internal_span(&self, metadata: &ExperimentMetadata, start_time: f64) -> SpanHandle { let parent_info = ParentSpanInfo::Experiment { object_id: metadata.experiment_id.clone(), }; // Note: When using ParentSpanInfo::Experiment, we don't need project_name // because the destination is determined by the experiment_id. - let mut builder = SpanBuilder::new( - Arc::clone(&self.submitter), - self.token.clone(), - self.org_id.clone(), - ) - .parent_info(parent_info) - .span_type(SpanType::Eval) - .start_time(start_time); - - if let Some(org_name) = &self.org_name { - builder = builder.org_name(org_name.clone()); - } + let builder = SpanBuilder::new(Arc::clone(&self.submitter), self.scope.clone()) + .parent_info(parent_info) + .span_type(SpanType::Eval) + .start_time(start_time); builder.build() } diff --git a/src/experiments/metadata.rs b/src/experiments/metadata.rs index c510dfb..9ca9327 100644 --- a/src/experiments/metadata.rs +++ b/src/experiments/metadata.rs @@ -73,31 +73,4 @@ pub enum GitMetadataField { CommitTime, } -/// Repository information for experiment tracking. -#[derive(Debug, Clone, Default, Serialize)] -pub struct RepoInfo { - /// Git commit SHA. - #[serde(skip_serializing_if = "Option::is_none")] - pub commit: Option, - /// Current branch name. - #[serde(skip_serializing_if = "Option::is_none")] - pub branch: Option, - /// Current tag (if any). - #[serde(skip_serializing_if = "Option::is_none")] - pub tag: Option, - /// Whether the working directory has uncommitted changes. - #[serde(skip_serializing_if = "Option::is_none")] - pub dirty: Option, - /// Git author name. - #[serde(skip_serializing_if = "Option::is_none")] - pub author_name: Option, - /// Git author email. - #[serde(skip_serializing_if = "Option::is_none")] - pub author_email: Option, - /// Git commit message. - #[serde(skip_serializing_if = "Option::is_none")] - pub commit_message: Option, - /// Git commit timestamp. - #[serde(skip_serializing_if = "Option::is_none")] - pub commit_time: Option, -} +pub use crate::api::experiments::RepoInfo; diff --git a/src/experiments/mod.rs b/src/experiments/mod.rs index d2e0bc4..d38f8c2 100644 --- a/src/experiments/mod.rs +++ b/src/experiments/mod.rs @@ -12,7 +12,7 @@ //! .experiment_builder_with_credentials(&api_key, &org_id) //! .project_name("my-project") //! .experiment_name("gpt-4-baseline") -//! .build(); +//! .build()?; //! //! let id = experiment.log( //! ExperimentLog::builder() @@ -20,15 +20,12 @@ //! .output(json!({"answer": "4"})) //! .expected(json!({"answer": "4"})) //! .score("accuracy", 1.0) -//! .build()? -//! ).await; +//! .build() +//! ).await?; //! //! experiment.flush().await?; //! ``` -// Internal API module (not public) -pub(crate) mod api; - // Public modules mod builder; mod event_log; diff --git a/src/experiments/span_builder.rs b/src/experiments/span_builder.rs index 4f519e5..f556afc 100644 --- a/src/experiments/span_builder.rs +++ b/src/experiments/span_builder.rs @@ -3,23 +3,24 @@ use std::sync::Arc; use crate::experiments::experiment::{Experiment, ExperimentMetadata}; -use crate::span::{SpanBuilder, SpanHandle, SpanSubmitter}; +use crate::internal::{ports::SpanSubmitter, session::CredentialScope}; +use crate::span::{SpanBuilder, SpanHandle}; use crate::types::{ParentSpanInfo, SpanType}; /// Builder for creating spans within an experiment. -#[allow(private_bounds)] -pub struct ExperimentSpanBuilder<'a, S: SpanSubmitter> { - experiment: &'a Experiment, +pub struct ExperimentSpanBuilder { + submitter: Arc, + scope: CredentialScope, metadata: ExperimentMetadata, name: Option, span_type: SpanType, } -#[allow(private_bounds)] -impl<'a, S: SpanSubmitter + 'static> ExperimentSpanBuilder<'a, S> { - pub(crate) fn new(experiment: &'a Experiment, metadata: ExperimentMetadata) -> Self { +impl ExperimentSpanBuilder { + pub(crate) fn new(experiment: &Experiment, metadata: ExperimentMetadata) -> Self { Self { - experiment, + submitter: experiment.submitter.clone(), + scope: experiment.scope.clone(), metadata, name: None, span_type: SpanType::Eval, @@ -39,25 +40,20 @@ impl<'a, S: SpanSubmitter + 'static> ExperimentSpanBuilder<'a, S> { } /// Build the span. - pub fn build(self) -> SpanHandle { + pub fn build(self) -> SpanHandle { let parent_info = ParentSpanInfo::Experiment { object_id: self.metadata.experiment_id.clone(), }; // Note: When using ParentSpanInfo::Experiment, we don't need project_name // because the destination is determined by the experiment_id. - let mut builder = SpanBuilder::new( - Arc::clone(&self.experiment.submitter), - self.experiment.token.clone(), - self.experiment.org_id.clone(), - ) - .parent_info(parent_info) - .span_type(self.span_type); - - if let Some(org_name) = &self.experiment.org_name { - builder = builder.org_name(org_name.clone()); - } + let mut builder = SpanBuilder::new(self.submitter, self.scope) + .parent_info(parent_info) + .span_type(self.span_type); + if let Some(name) = self.name { + builder = builder.name(name); + } builder.build() } } diff --git a/src/http.rs b/src/http.rs index 140c964..ec76f9a 100644 --- a/src/http.rs +++ b/src/http.rs @@ -58,7 +58,19 @@ pub(crate) fn build_http_client( ca_bundle: Option<&CaBundleConfig>, user_agent: Option<&str>, ) -> Result { + build_http_client_builder(timeout, ca_bundle, user_agent)? + .build() + .map_err(|err| BraintrustError::InvalidConfig(err.to_string())) +} + +pub(crate) fn build_http_client_builder( + timeout: std::time::Duration, + ca_bundle: Option<&CaBundleConfig>, + user_agent: Option<&str>, +) -> Result { let mut builder = Client::builder() + // SDK retries own request repetition; redirects must not add attempts. + .redirect(reqwest::redirect::Policy::none()) .timeout(timeout) .user_agent(user_agent.unwrap_or(USER_AGENT)); @@ -81,9 +93,7 @@ pub(crate) fn build_http_client( } } - builder - .build() - .map_err(|err| BraintrustError::InvalidConfig(err.to_string())) + Ok(builder) } #[cfg(test)] @@ -290,6 +300,20 @@ cb0bfZ7fHHs= server.await.expect("TLS test server task"); } + #[tokio::test] + async fn direct_api_client_uses_shared_custom_tls_configuration() { + let (url, certificate_pem, server) = start_tls_server(1).await; + let mut config = crate::api::Config::new("token"); + config.api_url = url; + config.custom_ca_bundle_pem = Some(certificate_pem.into_bytes()); + let api = crate::api::Client::new(config).expect("direct API client"); + api.logs3() + .post_raw(b"{}".to_vec()) + .await + .expect("private TLS request"); + server.await.expect("TLS test server task"); + } + #[tokio::test] async fn environment_ca_bundle_enables_private_tls_server() { let _lock = ENV_LOCK.lock().await; diff --git a/src/internal/batching.rs b/src/internal/batching.rs new file mode 100644 index 0000000..9991d14 --- /dev/null +++ b/src/internal/batching.rs @@ -0,0 +1,110 @@ +/// Limits include the container and separators supplied by the protocol encoder. +#[derive(Clone, Copy)] +pub(crate) struct BatchLimits { + pub max_items: usize, + pub max_bytes: usize, + pub container_bytes: usize, + pub separator_bytes: usize, +} + +/// Ordered groups: oversized items remain in their original position. +pub(crate) enum Partition { + Batch(Vec), + Oversized(T), +} + +pub(crate) struct Batcher; + +impl Batcher { + pub fn partition( + items: Vec, + limits: BatchLimits, + weight: impl Fn(&T) -> usize, + ) -> Vec> { + let mut output = Vec::new(); + let mut batch = Vec::new(); + let mut bytes = limits.container_bytes; + for item in items { + let size = weight(&item); + let singleton = limits.container_bytes.checked_add(size); + let next = bytes + .checked_add(if batch.is_empty() { + 0 + } else { + limits.separator_bytes + }) + .and_then(|bytes| bytes.checked_add(size)); + let fits = next.is_some_and(|bytes| bytes <= limits.max_bytes) + && batch.len() < limits.max_items.max(1); + if !batch.is_empty() && !fits { + output.push(Partition::Batch(std::mem::take(&mut batch))); + bytes = limits.container_bytes; + } + if singleton.is_none_or(|bytes| bytes > limits.max_bytes) { + output.push(Partition::Oversized(item)); + } else { + bytes += if batch.is_empty() { + 0 + } else { + limits.separator_bytes + }; + bytes += size; + batch.push(item); + } + } + if !batch.is_empty() { + output.push(Partition::Batch(batch)); + } + output + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn partitions_in_order_with_explicit_oversized_items() { + let parts = Batcher::partition( + vec![2, 3, 50, 1, 1, 1], + BatchLimits { + max_items: 2, + max_bytes: 8, + container_bytes: 2, + separator_bytes: 1, + }, + |size| *size, + ); + let parts: Vec<_> = parts + .into_iter() + .map(|part| match part { + Partition::Batch(items) => (false, items), + Partition::Oversized(item) => (true, vec![item]), + }) + .collect(); + assert_eq!( + parts, + vec![ + (false, vec![2, 3]), + (true, vec![50]), + (false, vec![1, 1]), + (false, vec![1]) + ] + ); + } + + #[test] + fn weights_cannot_wrap_and_empty_input_has_no_batches() { + let limits = BatchLimits { + max_items: 0, + max_bytes: usize::MAX, + container_bytes: 2, + separator_bytes: 1, + }; + assert!(Batcher::partition(Vec::::new(), limits, |size| *size).is_empty()); + assert!(matches!( + Batcher::partition(vec![usize::MAX], limits, |size| *size)[0], + Partition::Oversized(_) + )); + } +} diff --git a/src/internal/delivery.rs b/src/internal/delivery.rs new file mode 100644 index 0000000..87cb815 --- /dev/null +++ b/src/internal/delivery.rs @@ -0,0 +1,903 @@ +//! Atomic admission, fixed flush cutoffs and active-batch outcome ownership. +use crate::{ + Completion, DeliveryError, DeliveryReceipt, EnqueueError, ExecutionError, FlushError, + FlushFailure, FlushReport, ShutdownError, ShutdownReport, WriteError, +}; +use futures::{ + future::{BoxFuture, Shared}, + FutureExt, +}; +use std::{ + collections::VecDeque, + sync::{Arc, Mutex, MutexGuard}, + time::Duration, +}; +use tokio::{ + runtime::Handle, + sync::{oneshot, Notify, Semaphore}, + task::AbortHandle, +}; +use tracing::warn; + +use crate::internal::write_outcome::{WriteOutcome, WriteResponder}; +type WriteSender = WriteResponder; +type FlushSender = oneshot::Sender>; +#[derive(Clone)] +enum BatchError { + Delivery(Arc), + Execution(ExecutionError), +} +struct Accepted { + command: C, + sender: WriteSender, +} +struct Batch { + commands: Vec, + completion: Arc, +} +struct BatchCompletion { + count: usize, + senders: Mutex>>, +} +fn lock(value: &Mutex) -> MutexGuard<'_, T> { + value + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} +impl BatchCompletion { + fn finish(&self, result: &Result<(), BatchError>) { + if let Some(writes) = lock(&self.senders).take() { + let write_result = match result { + Ok(()) => Ok(DeliveryReceipt), + Err(BatchError::Delivery(error)) => Err(WriteError::Delivery(error.clone())), + Err(BatchError::Execution(error)) => Err(WriteError::Execution(*error)), + }; + for sender in writes { + let _ = sender.send(write_result.clone()); + } + } + } +} +/// One bounded summary per active flush, rather than cloning every covered +/// batch into each observer. Its cutoff is the last batch present at admission. +struct Barrier { + through: Arc, + sender: Option, + delivered: u64, + failure: Option, + shutdown: bool, +} +impl Barrier { + fn record(&mut self, count: u64, result: &Result<(), BatchError>) { + match result { + Ok(()) => self.delivered += count, + Err(BatchError::Delivery(error)) => match &mut self.failure { + Some(failure) => failure.failed += count, + None => { + self.failure = Some(FlushFailure { + failed: count, + cause: error.clone(), + }) + } + }, + Err(BatchError::Execution(error)) => { + if let Some(sender) = self.sender.take() { + let _ = sender.send(Err((*error).into())); + } + } + } + } + fn finish(&mut self) { + if let Some(sender) = self.sender.take() { + let result = match self.failure.take() { + Some(failure) => Err(FlushError::Delivery(failure)), + None => Ok(FlushReport { + delivered: self.delivered, + }), + }; + let _ = sender.send(result); + } + } +} +struct State { + pending: Vec>, + ready: VecDeque>, + active: Option>, + barriers: Vec, + outstanding: usize, + accepting: bool, + terminal: Option, + depth_warned: bool, +} +impl State { + fn seal(&mut self) { + if self.pending.is_empty() { + return; + } + let (commands, writes): (Vec<_>, Vec<_>) = std::mem::take(&mut self.pending) + .into_iter() + .map(|item| (item.command, item.sender)) + .unzip(); + let count = commands.len(); + self.ready.push_back(Batch { + commands, + completion: Arc::new(BatchCompletion { + count, + senders: Mutex::new(Some(writes)), + }), + }); + } + fn barrier(&mut self, sender: FlushSender, shutdown: bool) { + self.seal(); + if let Some(through) = self + .ready + .back() + .map(|batch| batch.completion.clone()) + .or_else(|| self.active.clone()) + { + self.barriers.push(Barrier { + through, + sender: Some(sender), + delivered: 0, + failure: None, + shutdown, + }); + } else { + let _ = sender.send(Ok(FlushReport::default())); + } + } + fn finish(&mut self, batch: &Arc, result: Result<(), BatchError>) { + if self + .active + .as_ref() + .is_some_and(|active| Arc::ptr_eq(active, batch)) + { + self.outstanding -= batch.count; + self.active = None; + if self.outstanding == 0 { + self.depth_warned = false; + } + batch.finish(&result); + // Every registered barrier covers this batch: completed cutoffs + // are removed here, before admitting any subsequent flush. + self.barriers.retain_mut(|barrier| { + barrier.record(batch.count as u64, &result); + if Arc::ptr_eq(&barrier.through, batch) { + barrier.finish(); + false + } else { + true + } + }); + } + } +} +struct Inner { + state: Mutex>, + wake: Notify, + capacity: usize, + control_limit: usize, + shutdown_controls: Arc, + abort: Mutex>, +} +impl Inner { + fn stop(&self, error: ExecutionError) { + let mut state = lock(&self.state); + state.accepting = false; + if state.terminal.is_some() { + return; + } + state.terminal = Some(error); + for accepted in state.pending.drain(..) { + let _ = accepted.sender.send(Err(error.into())); + } + let result = Err(BatchError::Execution(error)); + if let Some(active) = state.active.take() { + active.finish(&result); + } + for batch in state.ready.drain(..) { + batch.completion.finish(&result); + } + for mut barrier in state.barriers.drain(..) { + barrier.record(0, &result); + } + state.outstanding = 0; + self.wake.notify_one(); + } + fn close(&self) { + lock(&self.state).accepting = false; + self.wake.notify_one(); + } +} + +/// Producers own no completed-outcome registry. A batch owns its write senders +/// until all its physical requests finish; barriers summarize only active work. +pub(crate) struct WriteService { + inner: Arc>, + owners: Arc>, + runtime: Handle, + exit: Shared>>, +} +impl Clone for WriteService { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone(), + owners: self.owners.clone(), + runtime: self.runtime.clone(), + exit: self.exit.clone(), + } + } +} +struct ProducerOwner { + inner: std::sync::Weak>, +} +impl Drop for ProducerOwner { + fn drop(&mut self) { + if let Some(inner) = self.inner.upgrade() { + inner.close(); + } + } +} +impl WriteService { + pub fn new(capacity: usize, runtime: Handle, process: F) -> Self + where + F: Fn(Vec) -> BoxFuture<'static, anyhow::Result<()>> + Send + 'static, + { + let control_limit = capacity.clamp(1, 1024); + let inner = Arc::new(Inner { + state: Mutex::new(State { + pending: Vec::new(), + ready: VecDeque::new(), + active: None, + barriers: Vec::new(), + outstanding: 0, + accepting: true, + terminal: None, + depth_warned: false, + }), + wake: Notify::new(), + capacity, + control_limit, + shutdown_controls: Arc::new(Semaphore::new(control_limit)), + abort: Mutex::new(None), + }); + let (sender, receiver) = oneshot::channel(); + let exit = async move { receiver.await.unwrap_or(Err(ExecutionError::WorkerLost)) } + .boxed() + .shared(); + // Construct outside the future so cancellation before its first poll + // still rejects pending work and notifies every accepted observer. + let guard = WorkerGuard { + inner: inner.clone(), + exit: Some(sender), + }; + let task = runtime.spawn(async move { run(guard, process).await }); + *lock(&inner.abort) = Some(task.abort_handle()); + let owners = Arc::new(ProducerOwner { + inner: Arc::downgrade(&inner), + }); + Self { + inner, + owners, + runtime, + exit, + } + } + + #[cfg(test)] + pub fn enqueue(&self, command: C) -> Completion { + self.enqueue_with_rejection(command, None, |_, _| {}) + } + + pub fn enqueue_with_rejection( + &self, + command: C, + outcome: Option, + rejected: impl FnOnce(C, EnqueueError), + ) -> Completion { + let (sender, mut completion) = Completion::pair(); + let sender = WriteResponder::new(sender, outcome); + let mut state = lock(&self.inner.state); + let error = if !state.accepting { + Some(EnqueueError::Closed) + } else if state.outstanding >= self.inner.capacity { + Some(EnqueueError::Full) + } else { + None + }; + if let Some(error) = error { + completion.record_issuance(Err(error.into())); + let _ = sender.send(Err(error.into())); + drop(state); + rejected(command, error); + } else { + // Sender registration and publication have one linearization point. + state.pending.push(Accepted { command, sender }); + state.outstanding += 1; + if !state.depth_warned + && state.outstanding >= self.inner.capacity - self.inner.capacity / 5 + { + state.depth_warned = true; + warn!( + outstanding = state.outstanding, + capacity = self.inner.capacity, + "delivery queue has reached 80% of capacity" + ); + } + self.inner.wake.notify_one(); + } + completion + } + + pub fn flush(&self) -> Completion { + let (sender, mut completion) = Completion::pair(); + let mut state = lock(&self.inner.state); + if !state.accepting { + completion.record_issuance(Err(EnqueueError::Closed.into())); + } else if state + .barriers + .iter() + .filter(|barrier| !barrier.shutdown) + .count() + >= self.inner.control_limit + { + completion.record_issuance(Err(EnqueueError::Full.into())); + } else { + state.barrier(sender, false); + self.inner.wake.notify_one(); + } + completion + } + + pub fn shutdown(&self, timeout: Duration) -> Completion { + let (sender, mut completion) = Completion::pair(); + let Ok(permit) = self.inner.shutdown_controls.clone().try_acquire_owned() else { + completion.record_issuance(Err(EnqueueError::Full.into())); + return completion; + }; + let (barrier_sender, barrier) = Completion::pair(); + { + let mut state = lock(&self.inner.state); + if let Some(error) = state.terminal { + completion.record_issuance(Err(error.into())); + return completion; + } + state.accepting = false; + state.barrier(barrier_sender, true); + } + let inner = self.inner.clone(); + let exit = self.exit.clone(); + self.inner.wake.notify_one(); + self.runtime.spawn(async move { + let _permit = permit; + let work = async { + let report = barrier.await; + exit.await.map_err(ShutdownError::Execution)?; + report + .map(|report| ShutdownReport { + delivered: report.delivered, + }) + .map_err(Into::into) + }; + let result = match tokio::time::timeout(timeout, work).await { + Ok(result) => result, + Err(_) => { + inner.stop(ExecutionError::TimedOut); + if let Some(abort) = lock(&inner.abort).as_ref() { + abort.abort(); + } + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + } + }; + let _ = sender.send(result); + }); + completion + } + + #[cfg(test)] + pub fn is_empty(&self) -> bool { + lock(&self.inner.state).outstanding == 0 + } + #[cfg(test)] + fn outstanding(&self) -> usize { + lock(&self.inner.state).outstanding + } +} +struct WorkerGuard { + inner: Arc>, + exit: Option>>, +} +impl Drop for WorkerGuard { + fn drop(&mut self) { + if let Some(exit) = self.exit.take() { + let error = lock(&self.inner.state) + .terminal + .unwrap_or(ExecutionError::WorkerLost); + self.inner.stop(error); + let _ = exit.send(Err(error)); + } + } +} +async fn run(mut guard: WorkerGuard, process: F) +where + C: Send + 'static, + F: Fn(Vec) -> BoxFuture<'static, anyhow::Result<()>>, +{ + loop { + let wake = guard.inner.wake.notified(); + let next = { + let mut state = lock(&guard.inner.state); + state.seal(); + match state.ready.pop_front() { + Some(batch) => { + state.active = Some(batch.completion.clone()); + Some(batch) + } + None if !state.accepting => break, + None => None, + } + }; + if let Some(batch) = next { + let result = process(batch.commands) + .await + .map_err(|source| BatchError::Delivery(Arc::new(DeliveryError { source }))); + if let Err(BatchError::Delivery(error)) = &result { + warn!(%error, "delivery batch failed"); + } + lock(&guard.inner.state).finish(&batch.completion, result); + } else { + wake.await; + } + } + if let Some(exit) = guard.exit.take() { + let _ = exit.send(Ok(())); + } +} +#[cfg(test)] +mod tests { + use super::*; + use std::{ + future::IntoFuture, + sync::atomic::{AtomicBool, Ordering}, + }; + use tokio::sync::mpsc; + + fn gated( + capacity: usize, + ) -> ( + WriteService, + mpsc::UnboundedReceiver>, + Arc, + ) { + let (entered, receiver) = mpsc::unbounded_channel(); + let gate = Arc::new(Semaphore::new(0)); + let process_gate = gate.clone(); + let service = WriteService::new(capacity, Handle::current(), move |commands| { + let gate = process_gate.clone(); + let entered = entered.clone(); + Box::pin(async move { + entered.send(commands).unwrap(); + gate.acquire().await.unwrap().forget(); + Ok(()) + }) + }); + (service, receiver, gate) + } + + #[tokio::test] + async fn concurrent_flushes_observe_fixed_cutoffs_including_active_work() { + let (service, mut entered, gate) = gated(8); + let one = service.enqueue(1); + let first = service.flush(); + assert_eq!(entered.recv().await.unwrap(), vec![1]); + let concurrent = service.flush(); + let two = service.enqueue(2); + let mut second = Box::pin(service.flush().into_future()); + let three = service.enqueue(3); // beyond second's cutoff + gate.add_permits(1); + assert_eq!(first.await.unwrap().delivered, 1); + assert_eq!(concurrent.await.unwrap().delivered, 1); + one.await.unwrap(); + assert_eq!(entered.recv().await.unwrap(), vec![2]); + assert!(futures::poll!(second.as_mut()).is_pending()); + gate.add_permits(1); + assert_eq!(second.await.unwrap().delivered, 2); + two.await.unwrap(); + assert_eq!(entered.recv().await.unwrap(), vec![3]); + // The second flush finished without waiting for the later arrival. + gate.add_permits(1); + three.await.unwrap(); + assert_eq!(service.flush().await.unwrap().delivered, 0); + service.shutdown(Duration::from_secs(1)).await.unwrap(); + } + + #[tokio::test] + async fn a_new_flush_does_not_inherit_a_failed_prefix_from_an_older_active_flush() { + let (entered, mut receiver) = mpsc::unbounded_channel(); + let gate = Arc::new(Semaphore::new(0)); + let process_gate = gate.clone(); + let service = WriteService::new(4, Handle::current(), move |commands: Vec| { + let entered = entered.clone(); + let gate = process_gate.clone(); + Box::pin(async move { + entered.send(commands.clone()).unwrap(); + gate.acquire().await.unwrap().forget(); + if commands == [1] { + anyhow::bail!("first batch failed"); + } + Ok(()) + }) + }); + let first = service.enqueue(1); + assert_eq!(receiver.recv().await.unwrap(), vec![1]); + let second = service.enqueue(2); + let covering = service.flush(); + gate.add_permits(1); + assert!(matches!(first.await, Err(WriteError::Delivery(_)))); + assert_eq!(receiver.recv().await.unwrap(), vec![2]); + let fresh = service.flush(); + gate.add_permits(1); + second.await.unwrap(); + let Err(FlushError::Delivery(failure)) = covering.await else { + panic!("expected captured failure"); + }; + assert_eq!(failure.failed, 1); + assert_eq!(fresh.await.unwrap().delivered, 1); + service.shutdown(Duration::from_secs(1)).await.unwrap(); + } + + #[tokio::test] + async fn graceful_shutdown_reports_its_failed_batch_after_notifying_all_writes() { + let service = WriteService::new(4, Handle::current(), |_: Vec| { + Box::pin(async { + anyhow::bail!("terminal failure"); + }) + }); + let first = service.enqueue(1); + let second = service.enqueue(2); + let shutdown = service.shutdown(Duration::from_secs(1)); + let Err(ShutdownError::Delivery(failure)) = shutdown.await else { + panic!("expected delivery failure"); + }; + assert_eq!(failure.failed, 2); + for completion in [first, second] { + assert!(matches!(completion.await, Err(WriteError::Delivery(_)))); + } + assert!(service.is_empty()); + } + + #[tokio::test] + async fn capacity_includes_in_flight_work_and_rejection_is_observable_twice() { + let (service, mut entered, gate) = gated(1); + let first = service.enqueue(1); + entered.recv().await.unwrap(); + let rejected = service.enqueue(2); + assert!(matches!( + &*rejected, + Err(WriteError::Admission(EnqueueError::Full)) + )); + assert!(matches!( + rejected.await, + Err(WriteError::Admission(EnqueueError::Full)) + )); + assert_eq!(service.outstanding(), 1); + gate.add_permits(1); + first.await.unwrap(); + assert!(service.is_empty()); + gate.add_permits(1); + service.enqueue(3).await.unwrap(); + service.shutdown(Duration::from_secs(1)).await.unwrap(); + } + + #[tokio::test] + async fn dropped_observer_does_not_cancel_automatic_delivery_or_fill_the_queue() { + let (service, mut entered, gate) = gated(2); + drop(service.enqueue(1)); + assert_eq!(entered.recv().await.unwrap(), vec![1]); + let flush = service.flush(); + drop(flush); // lifecycle observers also do not cancel accepted work + gate.add_permits(1); + service.shutdown(Duration::from_secs(1)).await.unwrap(); + assert!(service.is_empty()); + } + + #[tokio::test] + async fn every_failed_batch_notifies_all_members_releases_errors_and_allows_fresh_work() { + let fail = Arc::new(AtomicBool::new(true)); + let process_fail = fail.clone(); + let service = WriteService::new(4, Handle::current(), move |_: Vec| { + let fail = process_fail.load(Ordering::Relaxed); + Box::pin(async move { + if fail { + anyhow::bail!("terminal failure"); + } + Ok(()) + }) + }); + for _ in 0..100 { + let first = service.enqueue(1); + let second = service.enqueue(2); + let flush = service.flush(); + let Err(FlushError::Delivery(failure)) = flush.await else { + panic!("expected failure"); + }; + assert_eq!(failure.failed, 2); + let weak = Arc::downgrade(&failure.cause); + for completion in [first, second] { + let Err(WriteError::Delivery(error)) = completion.await else { + panic!("expected failure"); + }; + assert!(Arc::ptr_eq(&error, &failure.cause)); + } + drop(failure); + assert!(weak.upgrade().is_none(), "SDK must release terminal errors"); + assert_eq!(service.outstanding(), 0); + assert_eq!(service.flush().await.unwrap().delivered, 0); + } + fail.store(false, Ordering::Relaxed); + let fresh = service.enqueue(3); + assert_eq!(service.flush().await.unwrap().delivered, 1); + fresh.await.unwrap(); + service.shutdown(Duration::from_secs(1)).await.unwrap(); + } + + #[tokio::test] + async fn shutdown_stops_admission_immediately_drains_and_is_shared_by_clones() { + let (service, mut entered, gate) = gated(4); + let write = service.enqueue(1); + let shutdown = service.shutdown(Duration::from_secs(1)); + let concurrent = service.clone().shutdown(Duration::from_secs(1)); + assert!(shutdown.is_ok()); + assert!(matches!( + service.enqueue(2).await, + Err(WriteError::Admission(EnqueueError::Closed)) + )); + entered.recv().await.unwrap(); + gate.add_permits(1); + write.await.unwrap(); + assert_eq!(shutdown.await.unwrap().delivered, 1); + assert_eq!(concurrent.await.unwrap().delivered, 1); + assert!(service.is_empty()); + } + + #[tokio::test] + async fn shutdown_timeout_fails_active_work_even_when_flush_slots_are_full() { + let (service, mut entered, _) = gated(1); + let write = service.enqueue(1); + entered.recv().await.unwrap(); + let flush = service.flush(); + let rejected = service.flush(); + assert!(matches!( + &*rejected, + Err(FlushError::Admission(EnqueueError::Full)) + )); + let shutdown = service.shutdown(Duration::from_millis(10)); + assert!( + shutdown.is_ok(), + "shutdown must not compete with flush slots" + ); + assert!(matches!( + shutdown.await, + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + write.await, + Err(WriteError::Execution(ExecutionError::TimedOut)) + )); + assert!(matches!( + flush.await, + Err(FlushError::Execution(ExecutionError::TimedOut)) + )); + assert!(service.is_empty()); + assert!(matches!( + service.enqueue(2).await, + Err(WriteError::Admission(EnqueueError::Closed)) + )); + } + + #[tokio::test] + async fn shutdown_timeout_also_retires_batches_queued_behind_active_work() { + let (service, mut entered, _) = gated(4); + let first = service.enqueue(1); + entered.recv().await.unwrap(); + let second = service.enqueue(2); + let third = service.enqueue(3); + let flush = service.flush(); + let shutdown = service.shutdown(Duration::from_millis(10)); + assert!(matches!( + shutdown.await, + Err(ShutdownError::Execution(ExecutionError::TimedOut)) + )); + for completion in [first, second, third] { + assert!(matches!( + completion.await, + Err(WriteError::Execution(ExecutionError::TimedOut)) + )); + } + assert!(matches!( + flush.await, + Err(FlushError::Execution(ExecutionError::TimedOut)) + )); + assert_eq!(service.outstanding(), 0); + assert!(lock(&service.inner.state).barriers.is_empty()); + assert!(lock(&service.inner.state).ready.is_empty()); + } + + #[tokio::test] + async fn very_large_write_capacity_does_not_overflow_lifecycle_control_limits() { + let (service, _, _) = gated(usize::MAX); + assert_eq!(service.inner.control_limit, 1024); + service.shutdown(Duration::from_secs(1)).await.unwrap(); + } + + #[tokio::test] + async fn worker_cancellation_before_first_poll_completes_every_accepted_sender() { + let (service, _, _) = gated(4); + let outcome = WriteOutcome::default(); + let first = service.enqueue_with_rejection(1, Some(outcome.clone()), |_, _| {}); + let repeated_end = outcome.observe(); + let second = service.enqueue(2); + let flush = service.flush(); + lock(&service.inner.abort).as_ref().unwrap().abort(); + for completion in [first, second, repeated_end] { + assert!(matches!( + completion.await, + Err(WriteError::Execution(ExecutionError::WorkerLost)) + )); + } + assert!(matches!( + flush.await, + Err(FlushError::Execution(ExecutionError::WorkerLost)) + )); + assert!(matches!( + outcome.observe().await, + Err(WriteError::Execution(ExecutionError::WorkerLost)) + )); + assert!(service.is_empty()); + assert!(matches!( + service.enqueue(3).await, + Err(WriteError::Admission(EnqueueError::Closed)) + )); + } + + #[tokio::test] + async fn worker_panic_notifies_active_and_pending_members() { + let (entered, mut receiver) = mpsc::unbounded_channel(); + let gate = Arc::new(Semaphore::new(0)); + let process_gate = gate.clone(); + let service = WriteService::new(4, Handle::current(), move |_: Vec| { + let entered = entered.clone(); + let gate = process_gate.clone(); + Box::pin(async move { + entered.send(()).unwrap(); + gate.acquire().await.unwrap().forget(); + panic!("test processor failed"); + }) + }); + let active = service.enqueue(1); + receiver.recv().await.unwrap(); + let pending = service.enqueue(2); + let flush = service.flush(); + gate.add_permits(1); + for completion in [active, pending] { + assert!(matches!( + completion.await, + Err(WriteError::Execution(ExecutionError::WorkerLost)) + )); + } + assert!(matches!( + flush.await, + Err(FlushError::Execution(ExecutionError::WorkerLost)) + )); + assert!(service.is_empty()); + } + + #[tokio::test] + async fn dropping_the_last_producers_drains_work_and_stops_the_worker() { + let (service, mut entered, gate) = gated(4); + let write = service.enqueue(1); + let clone = service.clone(); + let exit = service.exit.clone(); + std::thread::scope(|scope| { + scope.spawn(move || drop(service)); + scope.spawn(move || drop(clone)); + }); + entered.recv().await.unwrap(); + gate.add_permits(1); + write.await.unwrap(); + tokio::time::timeout(Duration::from_secs(1), exit) + .await + .unwrap() + .unwrap(); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn fast_workers_and_racing_flushes_cannot_lose_registered_completions() { + let seen = Arc::new(Mutex::new(std::collections::HashSet::new())); + let processed = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let process_seen = seen.clone(); + let process_count = processed.clone(); + let service = WriteService::new(1024, Handle::current(), move |commands: Vec| { + lock(&process_seen).extend(commands.iter().copied()); + process_count.fetch_add(commands.len(), Ordering::Relaxed); + Box::pin(async { Ok(()) }) + }); + let (writes, flushes) = std::thread::scope(|scope| { + let producers: Vec<_> = (0..4) + .map(|producer| { + let service = &service; + scope.spawn(move || { + (0..50) + .map(|item| service.enqueue(producer * 50 + item)) + .collect::>() + }) + }) + .collect(); + let flushers: Vec<_> = (0..2) + .map(|_| { + let service = &service; + scope.spawn(move || (0..50).map(|_| service.flush()).collect::>()) + }) + .collect(); + ( + producers + .into_iter() + .flat_map(|producer| producer.join().unwrap()) + .collect::>(), + flushers + .into_iter() + .flat_map(|flusher| flusher.join().unwrap()) + .collect::>(), + ) + }); + tokio::time::timeout(Duration::from_secs(5), async { + for write in writes { + write.await.unwrap(); + } + for flush in flushes { + flush.await.unwrap(); + } + }) + .await + .unwrap(); + assert_eq!(lock(&seen).len(), 200); + assert_eq!(processed.load(Ordering::Relaxed), 200); + assert!(service.is_empty()); + service.shutdown(Duration::from_secs(1)).await.unwrap(); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn racing_producers_cannot_overbook_capacity() { + let (service, _entered, gate) = gated(8); + let completions = std::thread::scope(|scope| { + let threads: Vec<_> = (0..32) + .map(|number| { + let service = &service; + scope.spawn(move || service.enqueue(number)) + }) + .collect(); + threads + .into_iter() + .map(|thread| thread.join().unwrap()) + .collect::>() + }); + assert_eq!( + completions + .iter() + .filter(|completion| completion.is_ok()) + .count(), + 8 + ); + assert_eq!(service.outstanding(), 8); + gate.add_permits(8); + let mut delivered = 0; + for completion in completions { + if completion.await.is_ok() { + delivered += 1; + } + } + assert_eq!(delivered, 8); + service.shutdown(Duration::from_secs(1)).await.unwrap(); + assert!(service.is_empty()); + } +} diff --git a/src/internal/ingest.rs b/src/internal/ingest.rs new file mode 100644 index 0000000..27e2ec5 --- /dev/null +++ b/src/internal/ingest.rs @@ -0,0 +1,247 @@ +//! Pure logs3 encoding and upload decisions. No transport or retry policy. +use super::batching::{BatchLimits, Batcher, Partition}; +use crate::types::{Logs3OverflowInputRow, Logs3OverflowInputRowMeta, LOGS_API_VERSION}; +use bytes::Bytes; +use serde::Serialize; +use serde_json::{Map, Value}; + +#[derive(Debug, Clone)] +pub(crate) struct EncodedRow { + data: Bytes, + overflow: Logs3OverflowInputRow, +} + +impl EncodedRow { + pub fn encode( + row: &impl Serialize, + mut overflow: Logs3OverflowInputRow, + ) -> Result { + let data = serde_json::to_vec(row)?; + overflow.input_row.byte_size = data.len(); + Ok(Self { + data: data.into(), + overflow, + }) + } + + pub fn from_map(row: &Map) -> Result { + let object_ids = [ + "experiment_id", + "dataset_id", + "prompt_session_id", + "project_id", + "log_id", + "function_data", + ] + .into_iter() + .filter_map(|key| row.get(key).map(|value| (key.to_owned(), value.clone()))) + .collect(); + Self::encode( + row, + Logs3OverflowInputRow { + object_ids, + is_delete: row.get("_object_delete").and_then(Value::as_bool), + input_row: Logs3OverflowInputRowMeta { byte_size: 0 }, + }, + ) + } +} + +#[derive(Debug, Clone)] +pub(crate) struct EncodedBatch { + pub data: Bytes, + pub overflow_rows: Vec, + rows: Vec, +} + +impl EncodedBatch { + fn encode(rows: Vec) -> Self { + let mut data = Vec::new(); + data.extend_from_slice(b"{\"rows\":["); + for (index, row) in rows.iter().enumerate() { + if index > 0 { + data.push(b','); + } + data.extend_from_slice(&row.data); + } + data.extend_from_slice(format!("],\"api_version\":{LOGS_API_VERSION}}}").as_bytes()); + Self { + data: data.into(), + overflow_rows: rows.iter().map(|row| row.overflow.clone()).collect(), + rows, + } + } + + pub fn partition(rows: Vec, max_items: usize, max_bytes: usize) -> Vec { + let container_bytes = Self::encode(Vec::new()).data.len(); + Batcher::partition( + rows, + BatchLimits { + max_items, + max_bytes, + container_bytes, + separator_bytes: 1, + }, + |row| row.data.len(), + ) + .into_iter() + .map(|part| { + Self::encode(match part { + Partition::Batch(rows) => rows, + // Keep the oversized row as a singleton for the planner to handle. + Partition::Oversized(row) => vec![row], + }) + }) + .collect() + } + + pub fn len(&self) -> usize { + self.rows.len() + } + + fn split(mut self) -> (Self, Self) { + let right = self.rows.split_off(self.rows.len() / 2); + (Self::encode(self.rows), Self::encode(right)) + } +} + +pub(crate) enum UploadPlan { + Direct(EncodedBatch), + Overflow(EncodedBatch), + Split(EncodedBatch, EncodedBatch), +} + +#[derive(Clone, Copy)] +pub(crate) struct UploadPlanner { + pub max_request_bytes: usize, + pub overflow: bool, +} + +impl UploadPlanner { + pub fn plan(&self, batch: EncodedBatch) -> Result { + self.decide(batch, false) + } + + /// A rejected direct request must change shape, never retry identical bytes. + pub fn rejected( + &self, + batch: EncodedBatch, + ) -> Result { + self.decide(batch, true) + } + + fn decide( + &self, + batch: EncodedBatch, + rejected: bool, + ) -> Result { + if !rejected && batch.data.len() <= self.max_request_bytes { + Ok(UploadPlan::Direct(batch)) + } else if self.overflow { + Ok(UploadPlan::Overflow(batch)) + } else if batch.len() > 1 { + let (left, right) = batch.split(); + Ok(UploadPlan::Split(left, right)) + } else { + Err(crate::error::BraintrustError::Api { + status: 413, + message: "single row exceeds server payload limit and overflow is unavailable" + .into(), + }) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn rows() -> Vec { + ["é\\\"", "second", "third"] + .into_iter() + .map(|id| { + EncodedRow::from_map( + json!({"id":id,"dataset_id":"ds","_object_delete":true}) + .as_object() + .unwrap(), + ) + .unwrap() + }) + .collect() + } + + #[test] + fn encoded_limits_include_envelope_separators_and_utf8() { + let expected = EncodedBatch::encode(rows()); + let exact = EncodedBatch::partition(rows(), 3, expected.data.len()); + assert_eq!(exact.len(), 1); + assert_eq!(exact[0].data, expected.data); + let smaller = EncodedBatch::partition(rows(), 3, expected.data.len() - 1); + assert_eq!(smaller.len(), 2); + for batch in smaller { + assert!(batch.data.len() < expected.data.len()); + let value: Value = serde_json::from_slice(&batch.data).unwrap(); + for (row, meta) in value["rows"] + .as_array() + .unwrap() + .iter() + .zip(batch.overflow_rows) + { + assert_eq!( + meta.input_row.byte_size, + serde_json::to_vec(row).unwrap().len() + ); + assert_eq!(meta.object_ids["dataset_id"], "ds"); + assert_eq!(meta.is_delete, Some(true)); + } + } + } + + #[test] + fn planner_handles_oversized_and_rejected_batches_without_reordering() { + let batch = EncodedBatch::encode(rows()); + let planner = UploadPlanner { + max_request_bytes: batch.data.len(), + overflow: false, + }; + assert!(matches!( + planner.plan(batch.clone()).unwrap(), + UploadPlan::Direct(_) + )); + let UploadPlan::Split(left, right) = planner.rejected(batch).unwrap() else { + panic!("expected split") + }; + let ids: Vec<_> = [left.clone(), right] + .iter() + .flat_map(|batch| { + serde_json::from_slice::(&batch.data).unwrap()["rows"] + .as_array() + .unwrap() + .iter() + .map(|row| row["id"].clone()) + .collect::>() + }) + .collect(); + assert_eq!(ids, vec![json!("é\\\""), json!("second"), json!("third")]); + assert!(planner.rejected(left.clone()).is_err()); + assert!(matches!( + UploadPlanner { + overflow: true, + ..planner + } + .rejected(left) + .unwrap(), + UploadPlan::Overflow(_) + )); + assert!(matches!( + UploadPlanner { + max_request_bytes: 1, + overflow: true + } + .plan(EncodedBatch::encode(rows())) + .unwrap(), + UploadPlan::Overflow(_) + )); + } +} diff --git a/src/internal/mod.rs b/src/internal/mod.rs new file mode 100644 index 0000000..a55e1dc --- /dev/null +++ b/src/internal/mod.rs @@ -0,0 +1,10 @@ +pub(crate) mod batching; +pub(crate) mod delivery; +pub(crate) mod ingest; +pub(crate) mod object_store; +pub(crate) mod ports; +pub(crate) mod resources; +pub(crate) mod retry; +pub(crate) mod session; +pub(crate) mod upload; +pub(crate) mod write_outcome; diff --git a/src/internal/object_store.rs b/src/internal/object_store.rs new file mode 100644 index 0000000..51ff4af --- /dev/null +++ b/src/internal/object_store.rs @@ -0,0 +1,84 @@ +//! One-attempt signed object uploads, independent of Braintrust credentials. +use crate::api::logs3::Logs3OverflowUpload; +use crate::error::{BraintrustError, Result}; +use bytes::Bytes; +use url::Url; + +/// Only signed grant information reaches object storage. Braintrust credentials +/// are scoped to the API client, never added by this transport. +pub(crate) struct ObjectStoreTransport { + pub(crate) http: reqwest::Client, +} +impl ObjectStoreTransport { + pub(crate) async fn upload(&self, grant: &Logs3OverflowUpload, payload: Bytes) -> Result<()> { + let url = Url::parse(&grant.signed_url) + .map_err(|_| BraintrustError::InvalidConfig("invalid overflow signed URL".into()))?; + let mut headers = reqwest::header::HeaderMap::new(); + if let Some(pairs) = &grant.headers { + for (key, value) in pairs { + let key = + reqwest::header::HeaderName::from_bytes(key.as_bytes()).map_err(|_| { + BraintrustError::InvalidConfig("invalid overflow header name".into()) + })?; + let value = reqwest::header::HeaderValue::from_str(value).map_err(|_| { + BraintrustError::InvalidConfig("invalid overflow header value".into()) + })?; + headers.insert(key, value); + } + } + let request = if grant.method.eq_ignore_ascii_case("POST") { + let fields = grant.fields.as_ref().ok_or_else(|| { + BraintrustError::InvalidConfig("overflow POST upload missing fields".into()) + })?; + let content_type = fields + .get("Content-Type") + .map(String::as_str) + .unwrap_or("application/json"); + let mut form = reqwest::multipart::Form::new(); + for (key, value) in fields { + form = form.text(key.clone(), value.clone()); + } + form = form.part( + "file", + reqwest::multipart::Part::bytes(payload.to_vec()) + .mime_str(content_type) + .map_err(|_| { + BraintrustError::InvalidConfig("invalid overflow content type".into()) + })?, + ); + headers.remove(reqwest::header::CONTENT_TYPE); + self.http.post(url).headers(headers).multipart(form) + } else if grant.method.eq_ignore_ascii_case("PUT") { + if url + .host_str() + .is_some_and(|host| host.ends_with(".blob.core.windows.net")) + && !headers.contains_key("x-ms-blob-type") + { + headers.insert( + "x-ms-blob-type", + reqwest::header::HeaderValue::from_static("BlockBlob"), + ); + } + self.http + .put(url) + .header("content-type", "application/json") + .headers(headers) + .body(payload) + } else { + return Err(BraintrustError::InvalidConfig( + "unsupported overflow upload method".into(), + )); + }; + let response = request + .send() + .await + .map_err(|error| BraintrustError::Http(error.without_url()))?; + if !response.status().is_success() { + return Err(BraintrustError::Api { + status: response.status().as_u16(), + message: response.text().await.unwrap_or_default(), + }); + } + Ok(()) + } +} diff --git a/src/internal/ports.rs b/src/internal/ports.rs new file mode 100644 index 0000000..b6a6a34 --- /dev/null +++ b/src/internal/ports.rs @@ -0,0 +1,75 @@ +//! Private dependency-injection boundaries for resource I/O and delivery. +use async_trait::async_trait; + +use crate::api::{ + btql::{BTQLQuery, BTQLResponse}, + datasets::{ + DatasetRegisterRequest, DatasetRegisterResponse, DatasetSummaryResponse, DatasetWrite, + }, + experiments::{ + ExperimentComparisonResponse, ExperimentRegisterRequest, ExperimentRegisterResponse, + }, +}; +use crate::error::Result; +use crate::experiments::BaseExperimentInfo; +use crate::internal::session::CredentialScope; +use crate::internal::write_outcome::WriteOutcome; +use crate::types::{ParentSpanInfo, SpanPayload}; + +/// 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 submit( + &self, + scope: CredentialScope, + payload: SpanPayload, + parent_info: Option, + outcome: Option, + ) -> crate::Completion; +} + +#[async_trait] +pub(crate) trait DatasetApi: Send + Sync { + async fn register_dataset( + &self, + scope: &CredentialScope, + request: DatasetRegisterRequest, + ) -> Result; + async fn fetch_dataset_records( + &self, + scope: &CredentialScope, + query: BTQLQuery, + ) -> Result; + async fn fetch_dataset_summary( + &self, + scope: &CredentialScope, + dataset_id: &str, + ) -> Result; + async fn write_dataset_record(&self, scope: &CredentialScope, row: DatasetWrite) -> Result<()>; +} + +#[async_trait] +pub(crate) trait ExperimentApi: Send + Sync { + async fn register_experiment( + &self, + scope: &CredentialScope, + request: ExperimentRegisterRequest, + ) -> Result; + async fn fetch_base_experiment( + &self, + scope: &CredentialScope, + experiment_id: &str, + ) -> Result>; + async fn fetch_experiment_comparison( + &self, + scope: &CredentialScope, + experiment_id: &str, + base_experiment_id: Option<&str>, + ) -> Result; +} diff --git a/src/internal/resources.rs b/src/internal/resources.rs new file mode 100644 index 0000000..557aec6 --- /dev/null +++ b/src/internal/resources.rs @@ -0,0 +1,129 @@ +//! Resource ports adapt SDK policy to the shared one-attempt API contracts. +use crate::internal::session::CredentialScope; +use async_trait::async_trait; + +use crate::api::{ + btql::{BTQLQuery, BTQLResponse}, + datasets::{ + DatasetRegisterRequest, DatasetRegisterResponse, DatasetSummaryResponse, DatasetWrite, + }, + experiments::{ + BaseExperimentRequest, ExperimentComparisonResponse, ExperimentRegisterRequest, + ExperimentRegisterResponse, + }, +}; +use crate::error::{BraintrustError, Result}; +use crate::experiments::BaseExperimentInfo; +use crate::internal::ports::{DatasetApi, ExperimentApi}; + +/// Shared HTTP pool without retaining client/session/delivery state. +#[derive(Clone)] +pub(crate) struct ResourceTransport { + http: reqwest::Client, +} +impl ResourceTransport { + pub fn new(http: reqwest::Client) -> Self { + Self { http } + } + fn resource_api(&self, scope: &CredentialScope) -> Result { + scope.api(&self.http) + } +} + +#[async_trait] +impl ExperimentApi for ResourceTransport { + async fn register_experiment( + &self, + scope: &CredentialScope, + request: ExperimentRegisterRequest, + ) -> Result { + self.resource_api(scope)? + .experiments() + .register(request) + .await + .map_err(Into::into) + } + + async fn fetch_base_experiment( + &self, + scope: &CredentialScope, + experiment_id: &str, + ) -> Result> { + let response = match self + .resource_api(scope)? + .experiments() + .base(BaseExperimentRequest { + id: experiment_id.to_owned(), + }) + .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)?, + }; + Ok(match (response.base_exp_id, response.base_exp_name) { + (Some(id), Some(name)) => Some(BaseExperimentInfo { id, name }), + _ => None, + }) + } + + async fn fetch_experiment_comparison( + &self, + scope: &CredentialScope, + experiment_id: &str, + base_experiment_id: Option<&str>, + ) -> Result { + self.resource_api(scope)? + .experiments() + .compare(experiment_id, base_experiment_id) + .await + .map_err(Into::into) + } +} + +#[async_trait] +impl DatasetApi for ResourceTransport { + async fn register_dataset( + &self, + scope: &CredentialScope, + request: DatasetRegisterRequest, + ) -> Result { + self.resource_api(scope)? + .datasets() + .register(request) + .await + .map_err(Into::into) + } + + async fn fetch_dataset_records( + &self, + scope: &CredentialScope, + query: BTQLQuery, + ) -> Result { + self.resource_api(scope)? + .btql() + .query_dataset_records(query) + .await + .map_err(Into::into) + } + + async fn fetch_dataset_summary( + &self, + scope: &CredentialScope, + dataset_id: &str, + ) -> Result { + self.resource_api(scope)? + .datasets() + .summarize(dataset_id) + .await + .map_err(Into::into) + } + + async fn write_dataset_record(&self, scope: &CredentialScope, row: DatasetWrite) -> Result<()> { + self.resource_api(scope)? + .datasets() + .write(&[row]) + .await + .map_err(Into::into) + } +} diff --git a/src/internal/retry.rs b/src/internal/retry.rs new file mode 100644 index 0000000..6c6653f --- /dev/null +++ b/src/internal/retry.rs @@ -0,0 +1,340 @@ +//! Retry mechanics shared by protocol-specific callers. +//! +//! Callers decide replay safety, classification and diagnostics. This module +//! only bounds attempts and advances a supplied delay schedule; it retains no +//! failed operations once the run ends. + +use std::{future::Future, iter::Peekable, time::Duration}; + +use backoff::{backoff::Backoff, ExponentialBackoffBuilder}; +use futures::future::BoxFuture; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct Attempt { + /// Zero-based index; total attempts are configured retries plus one. + pub(crate) index: usize, +} + +pub(crate) trait Sleeper { + fn sleep(&self, delay: Duration) -> BoxFuture<'_, ()>; +} +pub(crate) struct TokioSleeper; +impl Sleeper for TokioSleeper { + fn sleep(&self, delay: Duration) -> BoxFuture<'_, ()> { + Box::pin(tokio::time::sleep(delay)) + } +} + +pub(crate) struct RetryOutcome { + pub(crate) result: Result, + pub(crate) attempts: usize, +} + +pub(crate) struct RetryExecutor, S = TokioSleeper> { + delays: Peekable, + sleeper: S, + max_retries: usize, + next_index: Option, +} +impl> RetryExecutor { + pub(crate) fn new(max_retries: usize, delays: B) -> Self { + Self::with_sleeper(max_retries, delays, TokioSleeper) + } +} +impl, S: Sleeper> RetryExecutor { + fn with_sleeper(max_retries: usize, delays: B, sleeper: S) -> Self { + Self { + delays: delays.peekable(), + sleeper, + max_retries, + next_index: Some(0), + } + } + + /// Advance a stateful protocol orchestrator without moving its state into + /// an operation factory. Breaking/dropping this cursor stops further retries. + pub(crate) async fn next(&mut self) -> Option { + let index = self.next_index.take()?; + if index > 0 { + let delay = self.delays.next()?; + self.sleeper.sleep(delay).await; + } + self.next_index = (index < self.max_retries).then(|| index + 1); + Some(Attempt { index }) + } + + pub(crate) async fn run( + mut self, + operation: F, + retryable: C, + on_retry: N, + ) -> RetryOutcome + where + F: FnMut(Attempt) -> Fut, + Fut: Future>, + C: FnMut(&E) -> bool, + N: FnMut(Attempt, &E, Duration), + { + self.run_shared(operation, retryable, on_retry).await + } + + /// Execute one stage, sharing the remaining retry allowance with later + /// stages. Successful stages do not consume retries or reset the schedule. + pub(crate) async fn run_shared( + &mut self, + mut operation: F, + mut retryable: C, + mut on_retry: N, + ) -> RetryOutcome + where + F: FnMut(Attempt) -> Fut, + Fut: Future>, + C: FnMut(&E) -> bool, + N: FnMut(Attempt, &E, Duration), + { + if self.next_index == Some(0) { + self.next().await; + } + let mut attempts: usize = 0; + loop { + let attempt = Attempt { index: attempts }; + attempts = attempts.saturating_add(1); + match operation(attempt).await { + Ok(value) => { + return RetryOutcome { + result: Ok(value), + attempts, + } + } + Err(error) => { + if !retryable(&error) || self.next_index.is_none() { + return RetryOutcome { + result: Err(error), + attempts, + }; + } + let Some(delay) = self.delays.peek().copied() else { + return RetryOutcome { + result: Err(error), + attempts, + }; + }; + on_retry(attempt, &error, delay); + if self.next().await.is_none() { + return RetryOutcome { + result: Err(error), + attempts, + }; + } + } + } + } + } +} + +/// Deterministic capped doubling; the first delay follows the initial attempt. +pub(crate) fn exponential(initial: Duration, max: Duration) -> impl Iterator { + std::iter::successors(Some(initial.min(max)), move |delay| { + Some(delay.saturating_mul(2).min(max)) + }) +} + +/// Existing jittered upload policy, separate from executing HTTP operations. +pub(crate) fn jittered_exponential( + initial: Duration, + max: Duration, + randomization: f64, +) -> impl Iterator { + let mut backoff = ExponentialBackoffBuilder::new() + .with_initial_interval(initial) + .with_multiplier(2.0) + .with_randomization_factor(randomization) + .with_max_interval(max) + .with_max_elapsed_time(None) + .build(); + std::iter::from_fn(move || backoff.next_backoff()) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::{Arc, Mutex}; + + #[derive(Default, Clone)] + struct FakeSleeper { + delays: Arc>>, + } + impl Sleeper for FakeSleeper { + fn sleep(&self, delay: Duration) -> BoxFuture<'_, ()> { + self.delays.lock().unwrap().push(delay); + Box::pin(async {}) + } + } + + #[tokio::test] + async fn stages_share_retries_without_resetting_the_delay_schedule() { + let sleeper = FakeSleeper::default(); + let mut executor = RetryExecutor::with_sleeper( + 2, + exponential(Duration::from_secs(1), Duration::from_secs(8)), + sleeper.clone(), + ); + for _ in 0..2 { + let outcome = executor + .run_shared( + |attempt| async move { + if attempt.index == 0 { + Err("transient") + } else { + Ok(()) + } + }, + |_| true, + |_, _, _| {}, + ) + .await; + assert_eq!(outcome.attempts, 2); + assert_eq!(outcome.result, Ok(())); + } + let outcome = executor + .run_shared( + |_| async { Err::<(), _>("exhausted") }, + |_| true, + |_, _, _| {}, + ) + .await; + assert_eq!(outcome.attempts, 1); + assert_eq!(outcome.result, Err("exhausted")); + assert_eq!( + *sleeper.delays.lock().unwrap(), + vec![Duration::from_secs(1), Duration::from_secs(2)] + ); + } + + #[tokio::test] + async fn zero_retries_still_attempts_once_without_sleeping() { + let sleeper = FakeSleeper::default(); + let seen = sleeper.clone(); + let outcome = RetryExecutor::with_sleeper( + 0, + exponential(Duration::from_secs(1), Duration::from_secs(30)), + sleeper, + ) + .run( + |a| async move { + assert_eq!(a.index, 0); + Err::<(), _>("original") + }, + |_| true, + |_, _, _| panic!("no retry"), + ) + .await; + assert_eq!(outcome.result, Err("original")); + assert_eq!(outcome.attempts, 1); + assert!(seen.delays.lock().unwrap().is_empty()); + } + + #[tokio::test] + async fn retries_are_additional_attempts_with_capped_delays_and_notifications() { + let sleeper = FakeSleeper::default(); + let seen = sleeper.clone(); + let mut notifications = Vec::new(); + let outcome = RetryExecutor::with_sleeper( + 4, + exponential(Duration::from_secs(1), Duration::from_secs(3)), + sleeper, + ) + .run( + |a| async move { + if a.index == 4 { + Ok(a.index) + } else { + Err(a.index) + } + }, + |_| true, + |attempt, error, delay| notifications.push((attempt.index, *error, delay)), + ) + .await; + assert_eq!(outcome.result, Ok(4)); + assert_eq!(outcome.attempts, 5); + let delays = vec![ + Duration::from_secs(1), + Duration::from_secs(2), + Duration::from_secs(3), + Duration::from_secs(3), + ]; + assert_eq!(*seen.delays.lock().unwrap(), delays); + assert_eq!( + notifications.iter().map(|(_, _, d)| *d).collect::>(), + delays + ); + } + + #[tokio::test] + async fn permanent_errors_stop_without_sleeping_or_requiring_clone() { + struct OwnedError(Box); + let sleeper = FakeSleeper::default(); + let seen = sleeper.clone(); + let outcome = RetryExecutor::with_sleeper( + 10, + exponential(Duration::from_secs(1), Duration::from_secs(30)), + sleeper, + ) + .run( + |_| async { Err::<(), _>(OwnedError(Box::new(42))) }, + |_| false, + |_, _, _| panic!("permanent"), + ) + .await; + assert_eq!(*outcome.result.err().unwrap().0, 42); + assert_eq!(outcome.attempts, 1); + assert!(seen.delays.lock().unwrap().is_empty()); + } + + #[tokio::test] + async fn finite_schedule_returns_the_last_error_without_an_extra_attempt() { + let sleeper = FakeSleeper::default(); + let seen = sleeper.clone(); + let outcome = + RetryExecutor::with_sleeper(20, [Duration::from_secs(7)].into_iter(), sleeper) + .run( + |a| async move { Err::<(), _>(a.index) }, + |_| true, + |_, _, _| {}, + ) + .await; + assert_eq!(outcome.result, Err(1)); + assert_eq!(outcome.attempts, 2); + assert_eq!(*seen.delays.lock().unwrap(), [Duration::from_secs(7)]); + } + + #[tokio::test] + async fn dropping_a_retry_wait_does_not_schedule_another_attempt() { + struct PendingSleeper; + impl Sleeper for PendingSleeper { + fn sleep(&self, _: Duration) -> BoxFuture<'_, ()> { + Box::pin(std::future::pending()) + } + } + let attempts = std::sync::atomic::AtomicUsize::new(0); + let mut run = Box::pin( + RetryExecutor::with_sleeper( + 3, + exponential(Duration::from_secs(1), Duration::from_secs(3)), + PendingSleeper, + ) + .run( + |_| { + attempts.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + async { Err::<(), _>(()) } + }, + |_| true, + |_, _, _| {}, + ), + ); + assert!(futures::poll!(run.as_mut()).is_pending()); + drop(run); + assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 1); + } +} diff --git a/src/internal/session.rs b/src/internal/session.rs new file mode 100644 index 0000000..b82e1db --- /dev/null +++ b/src/internal/session.rs @@ -0,0 +1,81 @@ +//! Immutable credentials and endpoint routes captured by resource handles. +use std::{fmt, sync::Arc}; +use url::Url; + +/// Full request identity. Equality includes credentials, both routes and the org +/// header, so caches and upload groups cannot cross authentication boundaries. +#[derive(Clone, PartialEq, Eq, Hash)] +pub(crate) struct CredentialScope(Arc); +#[derive(Clone, PartialEq, Eq, Hash)] +struct Identity { + token: String, + org_id: String, + org_name: Option, + api_url: Url, + app_url: Url, +} +impl CredentialScope { + pub fn new( + token: String, + org_id: String, + org_name: Option, + api_url: Url, + app_url: Url, + ) -> Self { + Self(Arc::new(Identity { + token, + org_id, + org_name, + api_url, + app_url, + })) + } + pub fn token(&self) -> &str { + &self.0.token + } + pub fn org_id(&self) -> &str { + &self.0.org_id + } + pub fn org_name(&self) -> Option<&str> { + self.0.org_name.as_deref() + } + pub fn api_url(&self) -> &Url { + &self.0.api_url + } + pub fn with_org_name(self, org_name: String) -> Self { + let mut identity = (*self.0).clone(); + 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( + http.clone(), + &self.0.api_url, + &self.0.app_url, + self.token(), + self.org_name(), + ) + .map_err(Into::into) + } +} +impl fmt::Debug for CredentialScope { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("CredentialScope") + .field("org_id", &self.org_id()) + .field("org_name", &self.org_name()) + .finish_non_exhaustive() + } +} + +#[cfg(test)] +impl CredentialScope { + pub fn test(token: &str, org_id: &str) -> Self { + Self::new( + token.into(), + org_id.into(), + None, + Url::parse("http://localhost").unwrap(), + Url::parse("http://localhost").unwrap(), + ) + } +} diff --git a/src/internal/upload.rs b/src/internal/upload.rs new file mode 100644 index 0000000..03ba089 --- /dev/null +++ b/src/internal/upload.rs @@ -0,0 +1,496 @@ +//! SDK ingestion orchestration, built on the one-attempt API. +use super::ingest::{EncodedBatch, UploadPlan, UploadPlanner}; +use super::object_store::ObjectStoreTransport; +use super::retry::{exponential, jittered_exponential, RetryExecutor}; +use crate::api::{ + self, + logs3::{Logs3OverflowReference, Logs3OverflowRequest, Logs3OverflowUploadRequest}, +}; +use crate::error::{BraintrustError, Result}; +use crate::types::LOGS_API_VERSION; +use std::{future::Future, time::Duration}; +use tracing::warn; +use url::Url; + +#[derive(Clone, Copy, Debug)] +pub(crate) struct Capabilities { + pub max_request_size: usize, + pub overflow: bool, +} + +/// Existing schedules remain caller policy; orchestration and classification +/// are shared. Retries are shared across all stages and splits of one batch. +#[derive(Clone, Copy)] +pub(crate) enum UploadPolicy { + Queue { retries: usize }, + Standalone, +} +impl UploadPolicy { + fn budget(self) -> RetryBudget { + let (retries, delays): (_, Box + Send>) = match self { + Self::Queue { retries } => ( + retries, + Box::new(exponential(Duration::from_secs(1), Duration::from_secs(30))), + ), + Self::Standalone => ( + 4, + Box::new(jittered_exponential( + Duration::from_millis(300), + Duration::from_secs(8), + 0.2, + )), + ), + }; + RetryBudget { + executor: RetryExecutor::new(retries, delays), + requests: 0, + } + } +} +struct RetryBudget { + executor: RetryExecutor + Send>>, + requests: usize, +} +impl RetryBudget { + async fn run(&mut self, mut operation: F) -> Result + where + F: FnMut() -> Fut, + Fut: Future>, + { + let outcome = self.executor.run_shared( + |_| operation(), retryable, + |attempt, error, _| warn!(%error, retry = attempt.index + 1, "ingestion request failed, retrying"), + ).await; + self.requests += outcome.attempts; + outcome.result + } +} +fn retryable(error: &BraintrustError) -> bool { + match error { + BraintrustError::Http(error) => !error.is_builder(), + BraintrustError::Api { status, .. } => *status == 429 || (500..600).contains(status), + _ => false, + } +} + +#[derive(Default)] +pub(crate) struct UploadResult { + pub rows_uploaded: usize, + pub bytes_processed: usize, + pub requests_sent: usize, +} + +pub(crate) struct UploadCoordinator { + logs: api::logs3::Logs3, + objects: ObjectStoreTransport, + policy: UploadPolicy, +} +impl UploadCoordinator { + pub fn new( + http: reqwest::Client, + api_url: &Url, + token: &str, + 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(); + Ok(Self { + logs, + objects: ObjectStoreTransport { http }, + policy, + }) + } + + pub async fn capabilities( + &self, + fallback: usize, + override_bytes: Option, + ) -> Capabilities { + let server_limit = match self.logs.version().await { + Ok(version) => version.logs3_payload_max_bytes.filter(|limit| *limit > 0), + Err(error) => { + warn!(%error, "failed to fetch ingestion capabilities"); + None + } + }; + let max_request_size = match (server_limit, override_bytes) { + (Some(server), Some(local)) => server.min(local), + (Some(server), None) => server, + (None, Some(local)) => local, + (None, None) => fallback, + } + .max(1); + Capabilities { + max_request_size, + overflow: server_limit.is_some(), + } + } + + pub async fn upload( + &self, + batch: EncodedBatch, + capabilities: Capabilities, + ) -> Result { + let planner = UploadPlanner { + max_request_bytes: capabilities.max_request_size, + overflow: capabilities.overflow, + }; + let mut pending = vec![planner.plan(batch)?]; + let mut budget = self.policy.budget(); + let mut result = UploadResult::default(); + while let Some(plan) = pending.pop() { + match plan { + UploadPlan::Split(left, right) => { + pending.push(planner.plan(right)?); + pending.push(planner.plan(left)?); + } + UploadPlan::Direct(batch) => { + match budget + .run(|| async { + self.logs + .post_raw(batch.data.to_vec()) + .await + .map_err(BraintrustError::from) + }) + .await + { + Ok(()) => { + result.rows_uploaded += batch.len(); + result.bytes_processed += batch.data.len(); + } + Err(BraintrustError::Api { status: 413, .. }) => { + pending.push(planner.rejected(batch)?) + } + Err(error) => return Err(error), + } + } + UploadPlan::Overflow(batch) => { + self.overflow(&batch, &mut budget).await?; + result.rows_uploaded += batch.len(); + result.bytes_processed += batch.data.len(); + } + } + } + result.requests_sent = budget.requests; + Ok(result) + } + + async fn overflow(&self, batch: &EncodedBatch, budget: &mut RetryBudget) -> Result<()> { + let grant = budget + .run(|| async { + self.logs + .request_overflow_upload(Logs3OverflowUploadRequest { + content_type: "application/json".into(), + size_bytes: batch.data.len(), + rows: batch.overflow_rows.clone(), + }) + .await + .map_err(BraintrustError::from) + }) + .await?; + budget + .run(|| self.objects.upload(&grant, batch.data.clone())) + .await?; + budget + .run(|| async { + self.logs + .commit_overflow(Logs3OverflowRequest { + rows: Logs3OverflowReference { + reference_type: "logs3_overflow", + key: grant.key.clone(), + }, + api_version: LOGS_API_VERSION, + }) + .await + .map_err(BraintrustError::from) + }) + .await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::api::logs3::Logs3OverflowUpload; + use crate::internal::ingest::EncodedRow; + use serde_json::{json, Value}; + use wiremock::{ + matchers::{header, method, path}, + Mock, MockServer, ResponseTemplate, + }; + + fn batch() -> EncodedBatch { + EncodedBatch::partition( + vec![EncodedRow::from_map( + json!({"id":"row", "project_id":"project", "log_id":"g", "input":"é"}) + .as_object() + .unwrap(), + ) + .unwrap()], + 100, + usize::MAX, + ) + .pop() + .unwrap() + } + fn coordinator(server: &MockServer, policy: UploadPolicy) -> UploadCoordinator { + UploadCoordinator::new( + crate::http::build_http_client(Duration::from_secs(3), None, None).unwrap(), + &Url::parse(&server.uri()).unwrap(), + "secret", + Some("tenant"), + policy, + ) + .unwrap() + } + fn capabilities() -> Capabilities { + Capabilities { + max_request_size: 1, + overflow: true, + } + } + + #[tokio::test] + async fn both_callers_reuse_successful_overflow_stages_and_preserve_scopes() { + for policy in [UploadPolicy::Queue { retries: 2 }, UploadPolicy::Standalone] { + let server = MockServer::start().await; + Mock::given(method("POST")).and(path("/logs3/overflow")) + .and(header("authorization", "Bearer secret")).and(header("x-bt-org-name", "tenant")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"method":"PUT", "signedUrl":format!("{}/object", server.uri()), "key":"key"}))) + .expect(1).mount(&server).await; + Mock::given(method("PUT")) + .and(path("/object")) + .respond_with(ResponseTemplate::new(503)) + .with_priority(1) + .up_to_n_times(1) + .expect(1) + .mount(&server) + .await; + Mock::given(method("PUT")) + .and(path("/object")) + .respond_with(ResponseTemplate::new(200)) + .with_priority(2) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(503)) + .with_priority(1) + .up_to_n_times(1) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(200)) + .with_priority(2) + .expect(1) + .mount(&server) + .await; + let batch = batch(); + let data = batch.data.clone(); + let result = coordinator(&server, policy) + .upload(batch, capabilities()) + .await + .unwrap(); + assert_eq!(result.requests_sent, 5); + assert_eq!(result.rows_uploaded, 1); + assert_eq!(result.bytes_processed, data.len()); + let requests = server.received_requests().await.unwrap(); + for request in requests + .iter() + .filter(|request| request.url.path() == "/object") + { + assert!(!request + .headers + .keys() + .any(|key| key.as_str() == "authorization")); + assert!(!request + .headers + .keys() + .any(|key| key.as_str() == "x-bt-org-name")); + assert_eq!(request.body, data); + } + let grant: Value = serde_json::from_slice(&requests[0].body).unwrap(); + assert_eq!(grant["size_bytes"], data.len()); + assert_eq!(grant["rows"][0]["object_ids"]["project_id"], "project"); + } + } + + #[tokio::test] + async fn retry_budget_is_shared_across_grant_object_and_reference() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3/overflow")) + .respond_with(ResponseTemplate::new(503)) + .with_priority(1) + .up_to_n_times(1) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3/overflow")) + .respond_with(ResponseTemplate::new(200).set_body_json( + json!({"method":"PUT", "signedUrl":format!("{}/object",server.uri()), "key":"key"}), + )) + .with_priority(2) + .expect(1) + .mount(&server) + .await; + Mock::given(method("PUT")) + .and(path("/object")) + .respond_with(ResponseTemplate::new(503)) + .with_priority(1) + .up_to_n_times(1) + .expect(1) + .mount(&server) + .await; + Mock::given(method("PUT")) + .and(path("/object")) + .respond_with(ResponseTemplate::new(200)) + .with_priority(2) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(503)) + .expect(1) + .mount(&server) + .await; + let error = coordinator(&server, UploadPolicy::Queue { retries: 2 }) + .upload(batch(), capabilities()) + .await + .err() + .unwrap(); + assert!(matches!(error, BraintrustError::Api { status: 503, .. })); + assert_eq!(server.received_requests().await.unwrap().len(), 5); + } + + #[tokio::test] + async fn permanent_errors_stop_at_each_overflow_stage() { + for stage in ["grant", "object", "reference"] { + let server = MockServer::start().await; + let grant = if stage == "grant" { + ResponseTemplate::new(400) + } else { + ResponseTemplate::new(200).set_body_json(json!({"method":"PUT","signedUrl":format!("{}/object", server.uri()),"key":"key"})) + }; + Mock::given(method("POST")) + .and(path("/logs3/overflow")) + .respond_with(grant) + .expect(1) + .mount(&server) + .await; + Mock::given(method("PUT")) + .and(path("/object")) + .respond_with(ResponseTemplate::new(if stage == "object" { + 400 + } else { + 200 + })) + .expect(if stage == "grant" { 0 } else { 1 }) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(400)) + .expect(if stage == "reference" { 1 } else { 0 }) + .mount(&server) + .await; + assert!(matches!( + coordinator(&server, UploadPolicy::Queue { retries: 3 }) + .upload(batch(), capabilities()) + .await, + Err(BraintrustError::Api { status: 400, .. }) + )); + } + } + + #[tokio::test] + async fn signed_post_uses_multipart_boundary_and_preserves_grant_headers() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/object")) + .and(header("x-signed", "value")) + .respond_with(ResponseTemplate::new(204)) + .expect(1) + .mount(&server) + .await; + let payload = batch().data; + ObjectStoreTransport { + http: reqwest::Client::new(), + } + .upload( + &Logs3OverflowUpload { + method: "POST".into(), + signed_url: format!("{}/object", server.uri()), + headers: Some( + [ + ("content-type".into(), "must-be-replaced".into()), + ("x-signed".into(), "value".into()), + ] + .into(), + ), + fields: Some( + [ + ("Content-Type".into(), "application/json".into()), + ("key".into(), "object-key".into()), + ] + .into(), + ), + key: "key".into(), + }, + payload.clone(), + ) + .await + .unwrap(); + let requests = server.received_requests().await.unwrap(); + let content_type = requests[0] + .headers + .iter() + .find(|(key, _)| key.as_str() == "content-type") + .unwrap() + .1 + .get(0) + .unwrap() + .as_str(); + assert!(content_type.starts_with("multipart/form-data; boundary=")); + let body = String::from_utf8(requests[0].body.clone()).unwrap(); + assert!(body.contains("name=\"file\"")); + assert!(body.contains(std::str::from_utf8(&payload).unwrap())); + assert!(body.contains("object-key")); + } + + #[tokio::test] + async fn redirects_do_not_create_hidden_attempts_or_forward_credentials() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(307).insert_header("location", "/redirected")) + .expect(1) + .mount(&server) + .await; + Mock::given(path("/redirected")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&server) + .await; + let result = coordinator(&server, UploadPolicy::Queue { retries: 2 }) + .upload( + batch(), + Capabilities { + max_request_size: usize::MAX, + overflow: false, + }, + ) + .await; + assert!(matches!( + result, + Err(BraintrustError::Api { status: 307, .. }) + )); + } +} diff --git a/src/internal/write_outcome.rs b/src/internal/write_outcome.rs new file mode 100644 index 0000000..12415f1 --- /dev/null +++ b/src/internal/write_outcome.rs @@ -0,0 +1,133 @@ +//! One final write outcome shared by repeated span-end observers. +use std::sync::{Arc, Mutex, MutexGuard}; +use tokio::sync::oneshot; + +use crate::{Completion, DeliveryReceipt, EnqueueError, ExecutionError, WriteError}; + +type Outcome = Result; +type Sender = oneshot::Sender; +const MAX_PENDING_OBSERVERS: usize = 1024; + +/// Lives with a finished span, not with the queue. Holds no command or payload. +#[derive(Clone, Default)] +pub(crate) struct WriteOutcome(Arc>); +#[derive(Default)] +struct State { + result: Option, + observers: Vec, +} +fn lock(value: &Mutex) -> MutexGuard<'_, T> { + value + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} +impl WriteOutcome { + pub fn observe(&self) -> Completion { + let (sender, mut completion) = Completion::pair(); + let mut state = lock(&self.0); + if let Some(result) = &state.result { + let _ = sender.send(result.clone()); + } else { + state.observers.retain(|observer| !observer.is_closed()); + if state.observers.len() >= MAX_PENDING_OBSERVERS { + completion.record_issuance(Err(EnqueueError::Full.into())); + } else { + state.observers.push(sender); + } + } + completion + } + + fn finish(&self, result: Outcome) { + let mut state = lock(&self.0); + if state.result.is_some() { + return; + } + for sender in state.observers.drain(..) { + let _ = sender.send(result.clone()); + } + state.result = Some(result); + } +} + +/// Connect the existing batch sender to a span's optional shared final outcome. +/// No relay task or polling of a completion token is needed. +pub(crate) struct WriteResponder { + sender: Option, + outcome: Option, +} +impl WriteResponder { + pub fn new(sender: Sender, outcome: Option) -> Self { + Self { + sender: Some(sender), + outcome, + } + } + pub fn send(mut self, result: Outcome) -> Result<(), Outcome> { + if let Some(outcome) = self.outcome.take() { + outcome.finish(result.clone()); + } + match self.sender.take() { + Some(sender) => sender.send(result), + None => Err(result), + } + } +} +impl Drop for WriteResponder { + fn drop(&mut self) { + if let Some(outcome) = self.outcome.take() { + outcome.finish(Err(ExecutionError::WorkerLost.into())); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn pending_observers_are_bounded_and_dropped_tokens_release_slots() { + let outcome = WriteOutcome::default(); + let mut observers: Vec<_> = (0..MAX_PENDING_OBSERVERS) + .map(|_| outcome.observe()) + .collect(); + let full = outcome.observe(); + assert!(matches!( + &*full, + Err(WriteError::Admission(EnqueueError::Full)) + )); + assert!(matches!( + full.await, + Err(WriteError::Admission(EnqueueError::Full)) + )); + drop(observers.pop()); + let replacement = outcome.observe(); + assert!(replacement.is_ok()); + outcome.finish(Ok(DeliveryReceipt)); + for observer in observers { + observer.await.unwrap(); + } + replacement.await.unwrap(); + outcome.observe().await.unwrap(); + } + + #[tokio::test] + async fn producer_loss_completes_present_and_future_observers_without_a_relay_task() { + let outcome = WriteOutcome::default(); + let before = outcome.observe(); + let (sender, initial) = Completion::pair(); + drop(WriteResponder::new(sender, Some(outcome.clone()))); + assert!(matches!( + initial.await, + Err(WriteError::Execution(ExecutionError::WorkerLost)) + )); + assert!(matches!( + before.await, + Err(WriteError::Execution(ExecutionError::WorkerLost)) + )); + assert!(matches!( + outcome.observe().await, + Err(WriteError::Execution(ExecutionError::WorkerLost)) + )); + } +} diff --git a/src/lib.rs b/src/lib.rs index b1fd6a8..15c7583 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,9 +1,15 @@ +#![deny(private_bounds, private_interfaces)] + +pub mod api; +mod completion; mod dataset; +mod delivery; mod error; pub mod eval; mod experiments; mod extractors; mod http; +mod internal; mod json_merge; mod log_queue; mod logger; @@ -15,10 +21,15 @@ mod stream; pub(crate) mod test_utils; mod types; +pub use completion::{Completion, CompletionFuture}; pub use dataset::{ - DatasetBuilder, DatasetBuilderError, DatasetInsert, DatasetInsertBuilder, + Dataset, DatasetBuilder, DatasetBuilderError, DatasetInsert, DatasetInsertBuilder, DatasetInsertBuilderError, DatasetIterator, DatasetRecord, DatasetSummary, }; +pub use delivery::{ + DeliveryError, DeliveryReceipt, EnqueueError, ExecutionError, FlushError, FlushFailure, + FlushReport, ShutdownError, ShutdownReport, WriteError, +}; pub use error::{BraintrustError, Result}; pub use experiments::{ BaseExperimentInfo, Experiment, ExperimentBuilder, ExperimentBuilderError, ExperimentLog, diff --git a/src/log_queue/batching.rs b/src/log_queue/batching.rs index 0aaae99..338e6f6 100644 --- a/src/log_queue/batching.rs +++ b/src/log_queue/batching.rs @@ -1,134 +1,30 @@ use super::config::LogQueueConfig; -use crate::types::{Logs3OverflowInputRow, Logs3OverflowInputRowMeta, Logs3Row, LOGS_API_VERSION}; +use crate::types::{Logs3OverflowInputRow, Logs3OverflowInputRowMeta, Logs3Row}; use tracing::warn; -/// Generic batching function — matches TypeScript SDK's `batchItems()`. -/// -/// Splits `items` into batches respecting `max_items` and `max_bytes` limits. -/// `byte_size` is called once per item to determine its size. -/// -/// Flush condition (matches TypeScript): a new batch is started when the current -/// batch is non-empty AND (`batch.len() >= max_items` OR `batch_bytes + item_size >= max_bytes`). -/// This means a single item that exceeds `max_bytes` on its own is always placed in its own batch. -/// -/// Note: the byte comparison uses `>=` (not `>`), so two items that sum to exactly `max_bytes` -/// are split into separate batches. This matches the TypeScript SDK's strict `<` check in the -/// negated condition (`!(size < max && len < max)`). The test -/// `test_byte_boundary_flushes_at_exact_limit` verifies this invariant explicitly. -pub(crate) fn batch_items( - items: Vec, - max_items: usize, - max_bytes: usize, - byte_size: impl Fn(&T) -> usize, -) -> Vec> { - let max_items = max_items.max(1); - let max_bytes = max_bytes.max(1); - let mut output = Vec::new(); - let mut batch: Vec = Vec::new(); - let mut batch_bytes = 0usize; - - for item in items { - let item_size = byte_size(&item); - if !batch.is_empty() && (batch.len() >= max_items || batch_bytes + item_size >= max_bytes) { - output.push(batch); - batch = Vec::new(); - batch_bytes = 0; - } - batch_bytes += item_size; - batch.push(item); - } - if !batch.is_empty() { - output.push(batch); - } - output -} - -/// A serialized batch with per-row overflow metadata. -pub(crate) struct SerializedBatch { - /// The full batch serialized as a Logs3Request JSON payload. - pub data: Vec, - /// Per-row overflow metadata (for `POST logs3/overflow` requests). - pub overflow_rows: Vec, -} +pub(crate) use crate::internal::ingest::EncodedBatch as SerializedBatch; +use crate::internal::ingest::EncodedRow; -/// Batch rows respecting size limits, then serialize each batch to JSON. -/// -/// Each row is serialized once directly to bytes, and overflow metadata is -/// extracted from the `Logs3Row` struct without intermediate `Value` representation. -/// -/// `batch_max_bytes` is the effective byte limit for this flush, which may be -/// lower than `config.batch_max_bytes()` when the server advertises a tighter -/// limit via `GET /version` (callers should use `min(config, server_limit / 2)`). +/// Encode rows once, then partition using the complete request size. pub(crate) fn batch_and_serialize_rows( rows: Vec, config: &LogQueueConfig, batch_max_bytes: usize, ) -> Result, anyhow::Error> { - // Serialize each row once directly to bytes, extract overflow metadata from struct - let prepared: Vec<(Vec, Logs3OverflowInputRow)> = rows + let rows = rows .into_iter() .map(|row| { - let overflow_row = build_overflow_row_from_logs3row(&row); - let row_bytes = match serde_json::to_vec(&row) { - Ok(bytes) => bytes, - Err(error) => { - warn!(error = %error, "failed to serialize row"); - return Err(error); - } - }; - let overflow_row = Logs3OverflowInputRow { - input_row: Logs3OverflowInputRowMeta { - byte_size: row_bytes.len(), - }, - ..overflow_row - }; - Ok((row_bytes, overflow_row)) - }) - .collect::>()?; - - // Split into batches using the shared batch_items logic. - let batches = batch_items( - prepared, - config.batch_max_items(), - batch_max_bytes, - |(row_bytes, _)| row_bytes.len(), - ); - - // Assemble each batch by concatenating the pre-serialized row bytes. - batches - .into_iter() - .map(|batch| { - let (row_bytes, overflow_rows): (Vec<_>, Vec<_>) = batch.into_iter().unzip(); - assemble_logs3_request(&row_bytes, overflow_rows).map_err(|error| { - warn!(error = %error, "failed to assemble batch"); + EncodedRow::encode(&row, build_overflow_row_from_logs3row(&row)).map_err(|error| { + warn!(%error, "failed to serialize row"); error }) }) - .collect() -} - -/// Assemble a `Logs3Request` payload by concatenating pre-serialized row bytes. -/// -/// Produces `{"rows":[,,...],"api_version":}` without re-serializing rows. -fn assemble_logs3_request( - row_bytes: &[Vec], - overflow_rows: Vec, -) -> Result { - let mut data = Vec::new(); - data.extend_from_slice(b"{\"rows\":["); - for (i, row) in row_bytes.iter().enumerate() { - if i > 0 { - data.push(b','); - } - data.extend_from_slice(row); - } - data.extend_from_slice(b"],\"api_version\":"); - data.extend_from_slice(LOGS_API_VERSION.to_string().as_bytes()); - data.push(b'}'); - Ok(SerializedBatch { - data, - overflow_rows, - }) + .collect::, _>>()?; + Ok(SerializedBatch::partition( + rows, + config.batch_max_items(), + batch_max_bytes, + )) } /// Build overflow metadata directly from a Logs3Row struct. @@ -182,6 +78,7 @@ mod tests { use super::*; use crate::log_queue::config::LogQueueConfig; use crate::types::LogDestination; + use crate::types::LOGS_API_VERSION; use chrono::Utc; use std::collections::HashMap; @@ -277,21 +174,6 @@ mod tests { assert!(overflow_row.input_row.byte_size > 0); } - #[test] - fn test_byte_boundary_flushes_at_exact_limit() { - // A batch should flush when adding the next item would reach the limit (>=), - // not just exceed it (>). With max_bytes=100, two items of 50 bytes each - // should NOT be in the same batch (50 + 50 >= 100). - let items: Vec> = vec![vec![0u8; 50], vec![0u8; 50]]; - let batches = batch_items(items, 1000, 100, |b| b.len()); - - assert_eq!( - batches.len(), - 2, - "Items summing to exactly max_bytes should be split" - ); - } - #[test] fn test_overflow_metadata_is_delete_set_from_object_delete_field() { let config = LogQueueConfig::default(); @@ -322,15 +204,4 @@ mod tests { "is_delete should be None when _object_delete is not set" ); } - - #[test] - fn test_single_oversized_item_gets_its_own_batch() { - // An item larger than max_bytes should still be placed in a batch alone. - let items: Vec> = vec![vec![0u8; 200], vec![0u8; 10]]; - let batches = batch_items(items, 1000, 100, |b| b.len()); - - assert_eq!(batches.len(), 2); - assert_eq!(batches[0].len(), 1); // oversized item alone - assert_eq!(batches[1].len(), 1); - } } diff --git a/src/log_queue/config.rs b/src/log_queue/config.rs index a5d1a7b..b371107 100644 --- a/src/log_queue/config.rs +++ b/src/log_queue/config.rs @@ -122,13 +122,12 @@ pub struct LogQueueConfig { #[builder(default = default_num_retries())] num_retries: usize, - /// Sync flush mode (BRAINTRUST_SYNC_FLUSH, default: false) + /// Legacy compatibility value. Delivery is always automatic; ignored. #[builder(default = default_sync_flush())] sync_flush: bool, - /// Return preparation and HTTP failures to an explicit flush caller. - /// Checked clients also disable automatic flushing so no failure can be - /// consumed before that caller observes it. + /// Legacy compatibility value. Flush reports covered active failures and + /// delivery remains automatic; ignored. #[builder(default = false)] checked_delivery: bool, @@ -144,11 +143,9 @@ pub struct LogQueueConfig { /// Debug: save all payloads (BRAINTRUST_ALL_PUBLISH_PAYLOADS_DIR) all_publish_payloads_dir: Option, - /// When true, the queue drops new items once it reaches `queue_max_size`. - /// Defaults to false (unbounded), matching the TypeScript SDK's default behavior where - /// data loss via queue overflow is opt-in. Set to true to apply a hard cap and protect - /// against runaway memory growth in high-throughput scenarios. - #[builder(default = false)] + /// Legacy compatibility value. Admission is always bounded by queue_max_size, + /// including in-flight work; ignored. + #[builder(default = true)] enforce_queue_size_limit: bool, } @@ -220,7 +217,7 @@ mod tests { 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()); // unbounded by default (matches TS SDK) + assert!(config.enforce_queue_size_limit()); } #[test] diff --git a/src/log_queue/http.rs b/src/log_queue/http.rs index e2dcc25..806355b 100644 --- a/src/log_queue/http.rs +++ b/src/log_queue/http.rs @@ -1,8 +1,8 @@ use super::batching::SerializedBatch; use super::config::LogQueueConfig; -use crate::types::{Logs3OverflowReference, Logs3OverflowRequest, Logs3OverflowUpload}; +use crate::internal::upload::{Capabilities, UploadCoordinator, UploadPolicy}; use std::path::PathBuf; -use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; +use std::time::{SystemTime, UNIX_EPOCH}; use tracing::warn; use url::Url; @@ -26,73 +26,32 @@ pub(crate) fn save_payload_debug( Ok(()) } -/// Fetch version info from the server and return (max_request_size, can_use_overflow). -/// -/// `can_use_overflow` is true when the server returns a positive `logs3_payload_max_bytes`, -/// indicating that the overflow upload path is supported. +/// Compatibility adapter for the queue; protocol behavior lives in one coordinator. pub(crate) async fn fetch_version_info( client: &reqwest::Client, api_url: &Url, token: &str, org_name: Option<&str>, ) -> (usize, bool) { - let version_url = match api_url.join("version") { - Ok(url) => url, - Err(e) => { - warn!(error = %e, "Invalid API URL for version fetch"); - return (crate::log_queue::config::DEFAULT_BATCH_MAX_BYTES * 2, false); + let fallback = crate::log_queue::config::DEFAULT_BATCH_MAX_BYTES * 2; + match UploadCoordinator::new( + client.clone(), + api_url, + token, + org_name, + UploadPolicy::Queue { retries: 0 }, + ) { + Ok(coordinator) => { + let capabilities = coordinator.capabilities(fallback, None).await; + (capabilities.max_request_size, capabilities.overflow) } - }; - - let mut request = client.get(version_url).bearer_auth(token); - if let Some(org) = org_name { - request = request.header("x-bt-org-name", org); - } - let response = request.send().await; - - match response { - Ok(resp) if resp.status().is_success() => match resp.json::().await { - Ok(json) => { - let server_limit = json - .get("logs3_payload_max_bytes") - .and_then(|v| v.as_u64()) - .filter(|&v| v > 0) - .map(|v| v as usize); - - match server_limit { - Some(limit) => (limit, true), - None => (crate::log_queue::config::DEFAULT_BATCH_MAX_BYTES * 2, false), - } - } - Err(e) => { - warn!(error = %e, "Failed to parse version response"); - (crate::log_queue::config::DEFAULT_BATCH_MAX_BYTES * 2, false) - } - }, - Ok(resp) => { - warn!(status = %resp.status(), "Version fetch returned non-success status"); - (crate::log_queue::config::DEFAULT_BATCH_MAX_BYTES * 2, false) - } - Err(e) => { - warn!(error = %e, "Failed to fetch version info"); - (crate::log_queue::config::DEFAULT_BATCH_MAX_BYTES * 2, false) + Err(error) => { + warn!(%error, "invalid ingestion scope"); + (fallback, false) } } } -/// Send a batch with retry logic (matching TypeScript SDK behavior). -/// -/// If the batch exceeds `max_request_size` and `can_use_overflow` is true, -/// uses the overflow upload path (POST logs3/overflow → signed URL → POST logs3). -/// -/// Retry semantics match the TypeScript SDK: -/// - `config.num_retries()` is the number of *retries* (not total attempts) -/// - Total attempts = num_retries + 1 -/// - Retry delay starts at 1 second, doubles each attempt: `1s * 2^attempt` -/// -/// Note: the TypeScript SDK retries on all errors including 4xx. This implementation -/// intentionally diverges: only 429 and 5xx are retried since retrying on 400 Bad -/// Request (malformed payload) cannot succeed and wastes time. #[allow(clippy::too_many_arguments)] pub(crate) async fn send_batch_with_retry( client: &reqwest::Client, @@ -104,290 +63,30 @@ pub(crate) async fn send_batch_with_retry( max_request_size: usize, can_use_overflow: bool, ) -> Result<(), anyhow::Error> { - let json_bytes = batch.data; - let overflow_rows = batch.overflow_rows; - - // Save to debug directories if configured - save_payload_debug(&json_bytes, config.all_publish_payloads_dir(), "all").ok(); - - let logs_url = api_url - .join("logs3") - .map_err(|e| anyhow::anyhow!("invalid logs url: {e}"))?; - - let use_overflow = can_use_overflow && json_bytes.len() > max_request_size; - - // If overflow is needed, the signed URL is requested once and reused across retries - // (matches TypeScript SDK's `overflowUpload` caching pattern). - let mut overflow_upload: Option = None; - - // Total attempts = num_retries + 1 (matches TypeScript SDK: numTries = env + 1) - for attempt in 0..=config.num_retries() { - // Sleep before each retry (exponential backoff: 1s, 2s, 4s, ..., capped at 30s) - if attempt > 0 { - let delay_secs = 1u64.checked_shl((attempt - 1) as u32).unwrap_or(30).min(30); - tokio::time::sleep(Duration::from_secs(delay_secs)).await; - } - - if use_overflow { - // Get or create the overflow upload (cached across retries) - if overflow_upload.is_none() { - let upload_result = - request_overflow_upload(client, api_url, token, &json_bytes, &overflow_rows) - .await; - - match upload_result { - Ok(upload) => { - match upload_overflow_payload(client, &upload, &json_bytes).await { - Ok(()) => overflow_upload = Some(upload), - Err(e) => { - if attempt < config.num_retries() { - warn!(error = %e, attempt, "overflow upload failed, retrying"); - continue; - } - save_payload_debug( - &json_bytes, - config.failed_publish_payloads_dir(), - "failed", - ) - .ok(); - return Err(e); - } - } - } - Err(e) => { - if attempt < config.num_retries() { - warn!(error = %e, attempt, "overflow URL request failed, retrying"); - continue; - } - save_payload_debug( - &json_bytes, - config.failed_publish_payloads_dir(), - "failed", - ) - .ok(); - return Err(e); - } - } - } - - // Send the overflow reference to logs3 - let key = overflow_upload.as_ref().unwrap().key.clone(); - let overflow_ref = Logs3OverflowRequest { - rows: Logs3OverflowReference { - reference_type: "logs3_overflow", - key, - }, - api_version: crate::types::LOGS_API_VERSION, - }; - let overflow_bytes = serde_json::to_vec(&overflow_ref) - .map_err(|e| anyhow::anyhow!("failed to serialize overflow ref: {e}"))?; - - let mut overflow_req = client - .post(logs_url.clone()) - .bearer_auth(token) - .header("content-type", "application/json") - .body(overflow_bytes); - if let Some(org) = org_name { - overflow_req = overflow_req.header("x-bt-org-name", org); - } - let req_start = Instant::now(); - match overflow_req.send().await { - Ok(resp) if resp.status().is_success() => { - tracing::debug!( - bytes = json_bytes.len(), - elapsed_ms = req_start.elapsed().as_millis(), - "overflow batch sent" - ); - return Ok(()); - } - Ok(resp) => { - let status = resp.status(); - let body = resp.text().await.unwrap_or_default(); - let is_retryable = status == reqwest::StatusCode::TOO_MANY_REQUESTS - || status.is_server_error(); - if is_retryable && attempt < config.num_retries() { - warn!(status = %status, attempt, "overflow ref send failed, retrying"); - continue; - } - save_payload_debug(&json_bytes, config.failed_publish_payloads_dir(), "failed") - .ok(); - return Err(crate::error::BraintrustError::Api { - status: status.as_u16(), - message: body, - } - .into()); - } - Err(e) => { - if attempt < config.num_retries() { - warn!(error = %e, attempt, "overflow ref send network error, retrying"); - continue; - } - save_payload_debug(&json_bytes, config.failed_publish_payloads_dir(), "failed") - .ok(); - return Err(e.into()); - } - } - } else { - // Normal path: POST directly to logs3 - let mut normal_req = client - .post(logs_url.clone()) - .bearer_auth(token) - .header("content-type", "application/json") - .body(json_bytes.clone()); - if let Some(org) = org_name { - normal_req = normal_req.header("x-bt-org-name", org); - } - let req_start = Instant::now(); - match normal_req.send().await { - Ok(resp) if resp.status().is_success() => { - tracing::debug!( - bytes = json_bytes.len(), - elapsed_ms = req_start.elapsed().as_millis(), - "batch sent" - ); - return Ok(()); - } - Ok(resp) => { - let status = resp.status(); - let body = resp.text().await.unwrap_or_default(); - let is_retryable = status == reqwest::StatusCode::TOO_MANY_REQUESTS - || status.is_server_error(); - if is_retryable && attempt < config.num_retries() { - warn!(status = %status, attempt, "batch send failed with retryable error, retrying"); - continue; - } - save_payload_debug(&json_bytes, config.failed_publish_payloads_dir(), "failed") - .ok(); - return Err(crate::error::BraintrustError::Api { - status: status.as_u16(), - message: body, - } - .into()); - } - Err(e) => { - if attempt < config.num_retries() { - warn!(error = %e, attempt, "batch send network error, retrying"); - continue; - } - save_payload_debug(&json_bytes, config.failed_publish_payloads_dir(), "failed") - .ok(); - return Err(e.into()); - } - } - } - } - - unreachable!("retry loop should always return or continue") -} - -/// Request an overflow upload URL from the server. -pub(crate) async fn request_overflow_upload( - client: &reqwest::Client, - api_url: &Url, - token: &str, - json_bytes: &[u8], - overflow_rows: &[crate::types::Logs3OverflowInputRow], -) -> Result { - let overflow_url = api_url - .join("logs3/overflow") - .map_err(|e| anyhow::anyhow!("invalid overflow url: {e}"))?; - - let request_body = serde_json::json!({ - "content_type": "application/json", - "size_bytes": json_bytes.len(), - "rows": overflow_rows, - }); - - let resp = client - .post(overflow_url) - .bearer_auth(token) - .json(&request_body) - .send() - .await - .map_err(|e| anyhow::anyhow!("overflow URL request failed: {e}"))?; - - if !resp.status().is_success() { - let status = resp.status(); - let body = resp.text().await.unwrap_or_default(); - return Err(crate::error::BraintrustError::Api { - status: status.as_u16(), - message: body, - } - .into()); - } - - let upload: Logs3OverflowUpload = resp - .json() - .await - .map_err(|e| anyhow::anyhow!("failed to parse overflow response: {e}"))?; - - Ok(upload) -} - -/// Upload the payload to the signed URL provided by the overflow response. -pub(crate) async fn upload_overflow_payload( - client: &reqwest::Client, - upload: &Logs3OverflowUpload, - payload: &[u8], -) -> Result<(), anyhow::Error> { - let resp = if upload.method.eq_ignore_ascii_case("POST") { - // Multipart POST (e.g. S3) - let fields = upload - .fields - .as_ref() - .ok_or_else(|| anyhow::anyhow!("overflow POST upload missing fields"))?; - - let mut form = reqwest::multipart::Form::new(); - for (key, value) in fields { - form = form.text(key.clone(), value.clone()); - } - let content_type = fields - .get("Content-Type") - .map(|s| s.as_str()) - .unwrap_or("application/json"); - let part = reqwest::multipart::Part::bytes(payload.to_vec()) - .mime_str(content_type) - .map_err(|e| anyhow::anyhow!("invalid content-type: {e}"))?; - form = form.part("file", part); - - let mut req = client.post(&upload.signed_url).multipart(form); - if let Some(headers) = &upload.headers { - for (key, value) in headers { - if key.to_lowercase() != "content-type" { - req = req.header(key, value); - } - } - } - req.send().await - } else { - // PUT (e.g. Azure, GCS) - let mut req = client - .put(&upload.signed_url) - .header("content-type", "application/json") - .body(payload.to_vec()); - if let Some(headers) = &upload.headers { - for (key, value) in headers { - req = req.header(key, value); - } - } - // Azure Blob Storage requires `x-ms-blob-type: BlockBlob` on SAS-URL PUT - // requests; without it the server returns 400. The TypeScript SDK uses the - // same hostname-based check (`addAzureBlobHeaders`). A more robust signal - // would be an `upload_type` field in the server's overflow response, but - // that is a server-side change outside this SDK's scope for now. - if upload.signed_url.contains("blob.core.windows.net") { - req = req.header("x-ms-blob-type", "BlockBlob"); - } - req.send().await - }; - - let resp = resp.map_err(|e| anyhow::anyhow!("overflow upload failed: {e}"))?; - if !resp.status().is_success() { - let status = resp.status(); - let body = resp.text().await.unwrap_or_default(); - anyhow::bail!("overflow upload failed: [{status}] {body}"); + save_payload_debug(&batch.data, config.all_publish_payloads_dir(), "all").ok(); + let data = batch.data.clone(); + let coordinator = UploadCoordinator::new( + client.clone(), + api_url, + token, + org_name, + UploadPolicy::Queue { + retries: config.num_retries(), + }, + )?; + let result = coordinator + .upload( + batch, + Capabilities { + max_request_size, + overflow: can_use_overflow, + }, + ) + .await; + if result.is_err() { + save_payload_debug(&data, config.failed_publish_payloads_dir(), "failed").ok(); } - Ok(()) + result.map(|_| ()).map_err(Into::into) } #[cfg(test)] @@ -395,11 +94,14 @@ mod tests { use super::*; use crate::log_queue::config::LogQueueConfig; - fn make_batch(data: &[u8]) -> SerializedBatch { - SerializedBatch { - data: data.to_vec(), - overflow_rows: vec![], - } + fn make_batch() -> SerializedBatch { + SerializedBatch::partition( + vec![crate::internal::ingest::EncodedRow::from_map(&serde_json::Map::new()).unwrap()], + 1, + usize::MAX, + ) + .pop() + .unwrap() } #[test] @@ -437,6 +139,203 @@ mod tests { assert!(result.is_ok()); } + #[tokio::test] + async fn rejected_direct_batch_splits_in_order_without_retrying_identical_bytes() { + use wiremock::{ + matchers::{method, path}, + Mock, MockServer, ResponseTemplate, + }; + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(|request: &wiremock::Request| { + let body: serde_json::Value = serde_json::from_slice(&request.body).unwrap(); + ResponseTemplate::new(if body["rows"].as_array().unwrap().len() > 1 { + 413 + } else { + 200 + }) + }) + .expect(5) + .mount(&server) + .await; + let rows = ["first", "second", "third"] + .into_iter() + .map(|id| { + crate::internal::ingest::EncodedRow::from_map( + serde_json::json!({"id":id}).as_object().unwrap(), + ) + .unwrap() + }) + .collect(); + let batch = SerializedBatch::partition(rows, 100, usize::MAX) + .pop() + .unwrap(); + send_batch_with_retry( + &reqwest::Client::new(), + &Url::parse(&server.uri()).unwrap(), + "token", + None, + batch, + &LogQueueConfig::builder().num_retries(2).build(), + usize::MAX, + false, + ) + .await + .unwrap(); + let requests = server.received_requests().await.unwrap(); + let delivered: Vec<_> = requests + .iter() + .filter_map(|request| { + let body: serde_json::Value = serde_json::from_slice(&request.body).unwrap(); + let rows = body["rows"].as_array().unwrap(); + (rows.len() == 1).then(|| rows[0]["id"].clone()) + }) + .collect(); + assert_eq!( + delivered, + vec![ + serde_json::json!("first"), + serde_json::json!("second"), + serde_json::json!("third") + ] + ); + } + + #[tokio::test] + async fn rejected_singleton_uses_overflow_once_when_supported() { + use wiremock::{ + matchers::{method, path}, + Mock, MockServer, ResponseTemplate, + }; + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(413)) + .up_to_n_times(1) + .with_priority(1) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(200)) + .with_priority(2) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")).and(path("/logs3/overflow")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({"method":"PUT","signedUrl":format!("{}/object",server.uri()),"key":"key"}))) + .expect(1).mount(&server).await; + Mock::given(method("PUT")) + .and(path("/object")) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + send_batch_with_retry( + &reqwest::Client::new(), + &Url::parse(&server.uri()).unwrap(), + "token", + None, + make_batch(), + &LogQueueConfig::builder().num_retries(0).build(), + usize::MAX, + true, + ) + .await + .unwrap(); + } + + #[tokio::test] + async fn exhausted_queue_retries_send_exactly_the_configured_attempts() { + use wiremock::{ + matchers::{method, path}, + Mock, MockServer, ResponseTemplate, + }; + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(503)) + .expect(3) + .mount(&server) + .await; + let result = send_batch_with_retry( + &reqwest::Client::new(), + &Url::parse(&server.uri()).unwrap(), + "token", + None, + make_batch(), + &LogQueueConfig::builder().num_retries(2).build(), + usize::MAX, + false, + ) + .await; + assert!(result.is_err()); + assert_eq!(server.received_requests().await.unwrap().len(), 3); + } + + #[tokio::test] + async fn overflow_reference_retry_reuses_the_successful_grant_and_object_upload() { + use wiremock::{ + matchers::{method, path}, + Mock, MockServer, ResponseTemplate, + }; + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3/overflow")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "method":"PUT", "signedUrl":format!("{}/object", server.uri()), "key":"saved-grant" + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("PUT")) + .and(path("/object")) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(503)) + .with_priority(1) + .up_to_n_times(1) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(200)) + .with_priority(2) + .expect(1) + .mount(&server) + .await; + send_batch_with_retry( + &reqwest::Client::new(), + &Url::parse(&server.uri()).unwrap(), + "token", + None, + make_batch(), + &LogQueueConfig::builder().num_retries(1).build(), + 1, + true, + ) + .await + .unwrap(); + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests.len(), 4); + let object = requests.iter().find(|r| r.url.path() == "/object").unwrap(); + assert!(!object + .headers + .keys() + .any(|name| name.as_str().eq_ignore_ascii_case("authorization"))); + for request in requests.iter().filter(|r| r.url.path() == "/logs3") { + let body: serde_json::Value = serde_json::from_slice(&request.body).unwrap(); + assert_eq!(body["rows"]["key"], "saved-grant"); + } + } + #[tokio::test] async fn test_send_batch_with_retry_success() { use wiremock::{matchers::method, Mock, MockServer, ResponseTemplate}; @@ -456,7 +355,7 @@ mod tests { &api_url, "token", None, - make_batch(b"{}"), + make_batch(), &config, usize::MAX, false, @@ -492,7 +391,7 @@ mod tests { &api_url, "token", None, - make_batch(b"{}"), + make_batch(), &config, usize::MAX, false, @@ -527,7 +426,7 @@ mod tests { &api_url, "token", None, - make_batch(b"{}"), + make_batch(), &config, usize::MAX, false, @@ -556,7 +455,7 @@ mod tests { &api_url, "token", None, - make_batch(b"{}"), + make_batch(), &config, usize::MAX, false, @@ -586,7 +485,7 @@ mod tests { &api_url, "token", None, - make_batch(b"{}"), + make_batch(), &config, usize::MAX, false, @@ -621,7 +520,7 @@ mod tests { &api_url, "token", None, - make_batch(b"{}"), + make_batch(), &config, usize::MAX, false, diff --git a/src/log_queue/mod.rs b/src/log_queue/mod.rs index 32ffd50..48a4b3b 100644 --- a/src/log_queue/mod.rs +++ b/src/log_queue/mod.rs @@ -4,9 +4,7 @@ mod http; mod merge; mod queue; mod row_key; -pub(crate) mod worker; // Public API exports -pub(crate) use batching::batch_items; pub use config::LogQueueConfig; pub use queue::LogQueue; diff --git a/src/log_queue/queue.rs b/src/log_queue/queue.rs index 2b46720..24ee715 100644 --- a/src/log_queue/queue.rs +++ b/src/log_queue/queue.rs @@ -1,49 +1,38 @@ -use anyhow::Context; use chrono::Utc; -use serde::{Deserialize, Serialize}; use super::batching::batch_and_serialize_rows; use super::config::LogQueueConfig; use super::http::{fetch_version_info, save_payload_debug, send_batch_with_retry}; use super::merge::merge_row_into; use super::row_key::RowKey; -use super::worker::LogCommand; use crate::error::{BraintrustError, Result}; -use crate::logger::LoginState; +use crate::internal::delivery::WriteService; +use crate::internal::ports::{DeliveryControl, SpanSubmitter}; +use crate::internal::session::CredentialScope; +use crate::internal::write_outcome::WriteOutcome; use crate::span_components::{project_logs_identifier, ProjectLogsIdentifier}; use crate::types::{LogDestination, Logs3Row, ParentSpanInfo, SpanObjectType, SpanPayload}; -use arc_swap::ArcSwap; -use crossbeam::channel::{unbounded, Receiver, Sender, TrySendError}; +use crate::{ + Completion, DeliveryReceipt, FlushError, FlushReport, ShutdownError, ShutdownReport, WriteError, +}; use futures::future::join_all; use indexmap::IndexMap; -use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; -use tokio::sync::{mpsc, oneshot, Mutex as TokioMutex, OnceCell}; -use tokio::task::JoinHandle; +use tokio::sync::Mutex as TokioMutex; use tracing::warn; -use url::Url; -/// Request body for project registration. -#[derive(Serialize)] -struct ProjectRegisterRequest<'a> { - project_name: &'a str, - #[serde(skip_serializing_if = "Option::is_none")] - org_id: Option<&'a str>, - #[serde(skip_serializing_if = "Option::is_none")] - org_name: Option<&'a str>, -} - -/// Response from project registration. -#[derive(Deserialize)] -struct ProjectRegisterResponse { - project: ProjectInfo, -} +use crate::api::projects::ProjectRegisterRequest; -/// Project info in registration response. -#[derive(Deserialize)] -struct ProjectInfo { - id: String, +// Bound successful lookup caches too, so one shared client cannot retain an +// unlimited number of per-tenant credentials. Eviction only repeats a lookup. +const LOOKUP_CACHE_CAPACITY: usize = 1024; +fn cache_insert(cache: &mut IndexMap, key: K, value: V) { + if cache.len() >= LOOKUP_CACHE_CAPACITY && !cache.contains_key(&key) { + cache.shift_remove_index(0); + } + cache.insert(key, value); } /// A raw span payload queued for deferred preparation and flushing. @@ -51,285 +40,75 @@ struct ProjectInfo { /// Fields are private — construction and destructuring happen only within this /// module. `worker.rs` never touches the fields directly. struct SubmitCommand { - token: String, + scope: CredentialScope, payload: SpanPayload, parent_info: Option, } -struct CmdChannel { - sender: Sender, - receiver: Receiver, - queued_count: AtomicUsize, -} - -fn make_cmd_channel() -> CmdChannel { - let (sender, receiver) = unbounded(); - CmdChannel { - sender, - receiver, - queued_count: AtomicUsize::new(0), - } -} - -fn reserve_queue_slot( - queued_count: &AtomicUsize, - queue_max_size: usize, - enforce_queue_size_limit: bool, -) -> bool { - let mut current = queued_count.load(Ordering::Relaxed); - loop { - if enforce_queue_size_limit && current >= queue_max_size { - return false; - } - - match queued_count.compare_exchange_weak( - current, - current + 1, - Ordering::Relaxed, - Ordering::Relaxed, - ) { - Ok(_) => return true, - Err(next) => current = next, - } - } -} - -/// All state shared between `LogQueue` and the background worker. -/// -/// Always accessed through `Arc`. Fields no longer need their own `Arc` -/// wrappers — the outer `Arc` provides shared ownership and the field types provide -/// their own interior mutability where needed. -pub(super) struct LogQueueCore { - /// Lock-free queue channel (sender, receiver) pair, atomically swappable. - /// Holds unprocessed `SubmitCommand`s; preparation (project registration, row - /// building) happens during `flush_internal`. - channel: ArcSwap, +/// Preparation and protocol orchestration, independent of admission/lifecycle. +struct LogQueueCore { dropped_count: AtomicUsize, last_drop_log_time: AtomicU64, - queue_depth_warning_emitted: AtomicBool, config: LogQueueConfig, - login_state: LoginState, client: reqwest::Client, - /// Client-configured data-plane URL. Used as a fallback for flushing when the - /// shared login state has no URL (e.g. per-request-credential callers that - /// never log in). - api_url: Url, - /// Client-configured control-plane URL. Fallback for project registration - /// when the shared login state has no URL. - app_url: Url, - /// Cache of (org_id:project_name) → project_id to avoid redundant HTTP calls. - project_cache: TokioMutex>, - /// Handle to the currently in-progress flush task, if any. - /// Using JoinHandle allows checking completion status and awaiting the result. - flush_handle: TokioMutex>>>, - /// Lazily cached (max_request_size, can_use_overflow) from GET /version. - /// Fetched once on first flush and reused thereafter. - version_info: OnceCell<(usize, bool)>, + project_cache: TokioMutex>, + version_info: TokioMutex>, } - impl LogQueueCore { - fn new( - config: LogQueueConfig, - login_state: LoginState, - client: reqwest::Client, - api_url: Url, - app_url: Url, - ) -> Arc { - let channel = make_cmd_channel(); + fn new(config: LogQueueConfig, client: reqwest::Client) -> Arc { Arc::new(Self { - channel: ArcSwap::from_pointee(channel), dropped_count: AtomicUsize::new(0), last_drop_log_time: AtomicU64::new(0), - queue_depth_warning_emitted: AtomicBool::new(false), config, - login_state, client, - api_url, - app_url, project_cache: TokioMutex::new(IndexMap::new()), - flush_handle: TokioMutex::new(None), - version_info: OnceCell::new(), + version_info: TokioMutex::new(IndexMap::new()), }) } - /// Push a submit command directly to the lock-free queue. - /// - /// This is called by `LogQueue::submit` so that ALL drops (queue-full and - /// channel-closed) are handled uniformly here — with drop count tracking, - /// throttled warnings, and optional debug dumps. - fn push(self: &Arc, cmd: SubmitCommand) { - let channel = self.channel.load(); - if !reserve_queue_slot( - &channel.queued_count, - self.config.queue_max_size(), - self.config.enforce_queue_size_limit(), - ) { - self.dropped_count.fetch_add(1, Ordering::Relaxed); - self.log_drop_warning(); - self.dump_dropped_row_if_configured(cmd); - return; - } - - match channel.sender.try_send(cmd) { - Ok(()) => self.log_queue_depth_warning(channel.queued_count.load(Ordering::Relaxed)), - Err(TrySendError::Full(_)) => unreachable!("the log queue is unbounded"), - Err(TrySendError::Disconnected(_)) => { - channel.queued_count.fetch_sub(1, Ordering::Relaxed); - self.dropped_count.fetch_add(1, Ordering::Relaxed); - self.log_drop_warning(); - } - } - } - - /// Trigger a background flush if none is in progress. - /// Non-blocking — spawns a tokio task if needed. - /// No-op if no tokio runtime is available (e.g. in sync test contexts). - async fn trigger_flush_background(self: &Arc) { - let mut guard = self.flush_handle.lock().await; - - // Check if there's an existing flush in progress - if let Some(handle) = guard.as_ref() { - if !handle.is_finished() { - // Flush already in progress - return; - } + /// Capabilities are cached within the exact authentication and routing scope. + async fn get_version_info(&self, scope: &CredentialScope) -> (usize, bool) { + if let Some(info) = self.version_info.lock().await.get(scope).copied() { + return info; } - - // Consume the finished handle to surface any panic or error from the task. - if let Some(handle) = guard.take() { - if let Err(err) = handle.await { - warn!("background flush task failed: {:?}", err); - } - } - - // No flush in progress, spawn a new one - let Ok(runtime_handle) = tokio::runtime::Handle::try_current() else { - // No runtime available - return; - }; - - let core = self.clone(); - let join_handle = runtime_handle.spawn(async move { core.flush_internal().await }); - - *guard = Some(join_handle); - } - - /// Flush all rows in the lock-free queue and wait for completion. - /// - /// If a flush is already in progress, waits for it to complete and then - /// flushes any remaining items (matches TypeScript SDK behavior). - /// - /// Called by the background worker in response to `LogCommand::Flush` / - /// `LogCommand::TriggerFlush`. Does NOT drain the worker's mpsc channel — - /// callers that need that should use `LogQueue::flush_all()`. - pub(super) async fn flush(self: &Arc) -> std::result::Result<(), anyhow::Error> { - loop { - // Trigger a flush (no-op if one already running) - self.trigger_flush_background().await; - - // Wait for the current flush to complete, if any - let handle = { - let mut guard = self.flush_handle.lock().await; - guard.take() - }; - - if let Some(handle) = handle { - // Wait for the flush task to complete and propagate any errors - match handle.await { - Ok(Ok(())) => { - // Flush succeeded - } - Ok(Err(e)) => { - // Flush failed - return Err(e); - } - Err(e) => { - // Task panicked or was cancelled - return Err(anyhow::anyhow!("flush task failed: {}", e)); - } - } - } - - // Check if queue is now empty - if self.is_empty() { - return Ok(()); - } - - // Queue is non-empty — items arrived while the flush was running. - // Loop back to the top, which calls trigger_flush_background() again. - } - } - - /// Get the cached (max_request_size, can_use_overflow) from GET /version. - /// Fetches lazily on first call and caches the result. - async fn get_version_info( - &self, - api_url: &Url, - token: &str, - org_name: Option<&str>, - ) -> (usize, bool) { - *self - .version_info - .get_or_init(|| async { - fetch_version_info(&self.client, api_url, token, org_name).await - }) - .await + let info = fetch_version_info( + &self.client, + scope.api_url(), + scope.token(), + scope.org_name(), + ) + .await; + cache_insert(&mut *self.version_info.lock().await, scope.clone(), info); + info } - /// Internal flush implementation: drain → prepare → chunk → merge → batch → send. + /// Process one detached logical batch: prepare → merge → encode → upload. /// - /// Commands are drained atomically. Each command is prepared (project - /// registration, row building) before batching. Rows are then processed in + /// The write service owns membership and the fixed cutoff. Each command is + /// prepared (project registration, row building) before batching. Rows are then processed in /// sequential chunks of `min(batch_max_items, flush_chunk_size)`. /// Within each chunk, batches are sent concurrently. - async fn flush_internal(self: &Arc) -> std::result::Result<(), anyhow::Error> { - let cmds = self.drain_all(); + async fn process( + self: &Arc, + cmds: Vec, + ) -> std::result::Result<(), anyhow::Error> { if cmds.is_empty() { return Ok(()); } let mut first_error = None; - // Resolve the data-plane URL. Prefer the shared login state (covers the - // logged-in path, including an org-specific api_url returned by login) and - // fall back to the client-configured URL so per-request-credential callers - // work without ever populating the shared login state. - let api_url_str = self - .login_state - .api_url() - .unwrap_or_else(|| self.api_url.to_string()); - let api_url = match Url::parse(&api_url_str) { - Ok(url) => url, - Err(e) => { - warn!(error = %e, "Invalid API URL"); - return if self.config.checked_delivery() { - Err(e.into()) - } else { - Ok(()) - }; - } - }; - - // Prepare rows and bucket them by their upload credentials. The batch - // "destination" is (token, org_id): rows uploaded with different tokens or - // orgs must never be merged or batched into the same request under one - // token. `org_name` is carried alongside for the `x-bt-org-name` header. - type CredKey = (String, String); - let mut groups: IndexMap, Vec)> = IndexMap::new(); + // No live session lookup: registration, merging and upload share the + // exact scope captured by each producer when its handle was created. + let mut groups: IndexMap> = IndexMap::new(); for cmd in cmds { let SubmitCommand { - token, + scope, payload, parent_info, } = cmd; - match self.prepare_row(&token, payload, parent_info).await { + match self.prepare_row(&scope, payload, parent_info).await { Ok(row) => { - let key = (token, row.org_id.clone()); - let org_name = row.org_name.clone().or_else(|| self.login_state.org_name()); - groups - .entry(key) - .or_insert_with(|| (org_name, Vec::new())) - .1 - .push(row); + groups.entry(scope).or_default().push(row); } Err(e) => { warn!(error = %e, "failed to prepare span, dropping"); @@ -341,28 +120,9 @@ impl LogQueueCore { } if groups.is_empty() { - return if self.config.checked_delivery() { - first_error.map_or(Ok(()), Err) - } else { - Ok(()) - }; + return first_error.map_or(Ok(()), Err); } - // Version info reflects a server-wide limit, so it is fetched once and - // reused across all credential groups (using any group's credentials). - let (probe_token, probe_org) = { - let ((token, _org_id), (org_name, _rows)) = groups.get_index(0).unwrap(); - (token.clone(), org_name.clone()) - }; - let (max_request_size, can_use_overflow) = self - .get_version_info(&api_url, &probe_token, probe_org.as_deref()) - .await; - - // Effective batch byte limit: the smaller of the locally configured limit and half the - // server-advertised max request size. This mirrors the TypeScript SDK which uses - // `maxRequestSize / 2` as the batching limit, where maxRequestSize comes from /version. - let effective_batch_bytes = self.config.batch_max_bytes().min(max_request_size / 2); - // Chunk size matches TypeScript SDK: max(1, min(batchMaxItems, flushChunkSize)) let chunk_size = std::cmp::max( 1, @@ -372,7 +132,10 @@ impl LogQueueCore { ), ); - for ((token, _org_id), (org_name, rows)) in groups { + for (scope, rows) in groups { + let (max_request_size, can_use_overflow) = self.get_version_info(&scope).await; + // Preserve the queue policy of half the advertised request limit. + let effective_batch_bytes = self.config.batch_max_bytes().min(max_request_size / 2); // Process chunks sequentially (matches TypeScript SDK) let chunks: Vec> = rows.chunks(chunk_size).map(<[_]>::to_vec).collect(); @@ -422,9 +185,9 @@ impl LogQueueCore { .map(|batch| { send_batch_with_retry( &self.client, - &api_url, - &token, - org_name.as_deref(), + scope.api_url(), + scope.token(), + scope.org_name(), batch, &self.config, max_request_size, @@ -444,11 +207,7 @@ impl LogQueueCore { } } } - if self.config.checked_delivery() { - first_error.map_or(Ok(()), Err) - } else { - Ok(()) - } + first_error.map_or(Ok(()), Err) } /// Prepare a `Logs3Row` from a submit command's token, payload, and parent info. @@ -457,7 +216,7 @@ impl LogQueueCore { /// if needed) and builds the complete row for the flush pipeline. async fn prepare_row( &self, - token: &str, + scope: &CredentialScope, payload: SpanPayload, parent_info: Option, ) -> std::result::Result { @@ -466,8 +225,8 @@ impl LogQueueCore { span_id, is_merge, span_components, - org_id, - org_name, + org_id: _, + org_name: _, project_name, input, output, @@ -482,14 +241,12 @@ impl LogQueueCore { extra, } = payload; + let org_id = scope.org_id().to_owned(); + let org_name = scope.org_name().map(str::to_owned); + if let Some(span_components) = span_components { let destination = self - .destination_from_span_components( - token, - &org_id, - org_name.as_deref(), - &span_components, - ) + .destination_from_span_components(scope, &span_components) .await?; let root_span_id = span_components .root_span_id @@ -525,10 +282,7 @@ impl LogQueueCore { } let project_id = if let Some(ref pn) = project_name { - Some( - self.ensure_project_id(token, &org_id, org_name.as_deref(), pn) - .await?, - ) + Some(self.ensure_project_id(scope, pn).await?) } else { None }; @@ -552,9 +306,7 @@ impl LogQueueCore { LogDestination::project_logs(object_id), ), Some(ParentSpanInfo::ProjectName { project_name }) => { - let proj_id = self - .ensure_project_id(token, &org_id, org_name.as_deref(), &project_name) - .await?; + let proj_id = self.ensure_project_id(scope, &project_name).await?; (span_id.clone(), None, LogDestination::project_logs(proj_id)) } Some(ParentSpanInfo::PlaygroundLogs { object_id }) => ( @@ -590,14 +342,8 @@ impl LogQueueCore { LogDestination::project_logs(object_id) } Some(ProjectLogsIdentifier::ProjectName(project_name)) => { - let project_id = self - .ensure_project_id( - token, - &org_id, - org_name.as_deref(), - &project_name, - ) - .await?; + let project_id = + self.ensure_project_id(scope, &project_name).await?; LogDestination::project_logs(project_id) } None => { @@ -647,9 +393,7 @@ impl LogQueueCore { async fn destination_from_span_components( &self, - token: &str, - org_id: &str, - org_name: Option<&str>, + scope: &CredentialScope, span_components: &crate::span_components::SpanComponents, ) -> std::result::Result { match span_components.object_type { @@ -663,9 +407,7 @@ impl LogQueueCore { Ok(LogDestination::project_logs(object_id)) } Some(ProjectLogsIdentifier::ProjectName(project_name)) => { - let project_id = self - .ensure_project_id(token, org_id, org_name, &project_name) - .await?; + let project_id = self.ensure_project_id(scope, &project_name).await?; Ok(LogDestination::project_logs(project_id)) } None => anyhow::bail!( @@ -680,104 +422,32 @@ impl LogQueueCore { } } - /// Ensure a project ID is available for the given project name, registering - /// it via the API if it is not yet cached. - /// - /// Uses `login_state.app_url()` for the registration endpoint so the URL - /// doesn't need to be passed separately. + /// Resolve project names only within the producer's immutable request scope. async fn ensure_project_id( &self, - token: &str, - org_id: &str, - org_name: Option<&str>, + scope: &CredentialScope, project_name: &str, ) -> std::result::Result { - let cache_key = format!("{org_id}:{project_name}"); - { - let cache = self.project_cache.lock().await; - if let Some(project_id) = cache.get(&cache_key) { - return Ok(project_id.clone()); - } - } - - // Prefer the login-state control-plane URL (logged-in path), falling back - // to the client-configured URL for per-request-credential callers. - let app_url = match self.login_state.app_url() { - Some(app_url_str) => { - Url::parse(&app_url_str).map_err(|e| anyhow::anyhow!("invalid app URL: {e}"))? - } - None => self.app_url.clone(), - }; - let url = app_url - .join("api/project/register") - .map_err(|e| anyhow::anyhow!("invalid project register url: {e}"))?; - - let request = ProjectRegisterRequest { - project_name, - org_id: (!org_id.is_empty()).then_some(org_id), - org_name, - }; - - let response = self - .client - .post(url) - .bearer_auth(token) - .json(&request) - .send() - .await?; - let status = response.status(); - if !status.is_success() { - let text = response.text().await.unwrap_or_default(); - return Err(BraintrustError::Api { - status: status.as_u16(), - message: text, - } - .into()); + let cache_key = (scope.clone(), project_name.to_owned()); + if let Some(id) = self.project_cache.lock().await.get(&cache_key).cloned() { + return Ok(id); } - - let register_response: ProjectRegisterResponse = response - .json() + let response = scope + .api(&self.client)? + .projects() + .register(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), + }) .await - .context("failed to parse project registration response")?; - - { - let mut cache = self.project_cache.lock().await; - cache.insert(cache_key, register_response.project.id.clone()); - } - Ok(register_response.project.id) - } - - /// Drain all commands from the queue with precise flush boundaries. - /// - /// Atomically swaps in a new channel using lock-free CAS, then drains the old channel. - fn drain_all(&self) -> Vec { - // A new channel starts a new fill cycle, so it may emit the high-water warning. - self.queue_depth_warning_emitted - .store(false, Ordering::Relaxed); - let old_channel = self.channel.swap(Arc::new(make_cmd_channel())); - let mut items = Vec::new(); - while let Ok(item) = old_channel.receiver.try_recv() { - old_channel.queued_count.fetch_sub(1, Ordering::Relaxed); - items.push(item); - } - items - } - - /// Check if the queue is approximately empty. - pub fn is_empty(&self) -> bool { - self.channel.load().receiver.is_empty() - } - - /// Get the approximate number of items in the queue (tests only). - #[cfg(test)] - pub fn len(&self) -> usize { - self.channel.load().receiver.len() - } - - /// Get the total number of dropped items (tests only). - #[cfg(test)] - pub fn dropped_count(&self) -> usize { - self.dropped_count.load(Ordering::Relaxed) + .map_err(BraintrustError::from)?; + cache_insert( + &mut *self.project_cache.lock().await, + cache_key, + response.project.id.clone(), + ); + Ok(response.project.id) } /// Log a drop warning with throttling. @@ -811,23 +481,6 @@ impl LogQueueCore { } } - /// Warn once per fill cycle when the queue reaches 80% of its configured maximum. - fn log_queue_depth_warning(&self, queue_size: usize) { - let queue_max_size = self.config.queue_max_size(); - let warning_threshold = queue_max_size.saturating_mul(4).div_ceil(5); - - if queue_size >= warning_threshold - && !self - .queue_depth_warning_emitted - .swap(true, Ordering::Relaxed) - { - warn!( - queue_size, - queue_max_size, "Log queue has reached 80% of the configured maximum depth" - ); - } - } - /// If debug directories are configured, serialize the dropped command's payload /// and write it to disk. Spawns a background task — best-effort, non-blocking. fn dump_dropped_row_if_configured(&self, cmd: SubmitCommand) { @@ -842,7 +495,7 @@ impl LogQueueCore { }; handle.spawn(async move { let SubmitCommand { - token: _, + scope: _, payload, parent_info, } = cmd; @@ -963,405 +616,279 @@ impl LogQueueCore { } } -/// Lock-free log queue with batching and HTTP dispatch. -/// -/// This is the complete log processing system: -/// - Owns the background worker for flush coordination -/// - Lock-free queue for accumulating submit commands (bounded or unbounded) -/// - Row preparation (project registration, destination resolution) deferred to flush time -/// - Automatic batching and serialization via spawn_blocking -/// - HTTP dispatch with retry logic (concurrent per-chunk batches) -/// - Atomic drain for precise flush boundaries -/// - Matches TypeScript SDK `activeFlush` pattern: only one flush runs at a time -/// - Lazy version fetch for dynamic maxRequestSize and overflow support +/// Compatibility adapter connecting span preparation to the bounded write service. #[derive(Clone)] pub struct LogQueue { - /// All shared state, also held by the background worker via Arc. core: Arc, - /// Channel to the background worker (for flush commands only). - worker_sender: mpsc::Sender, - /// Handle to the background worker task. - #[allow(dead_code)] - worker_handle: Arc>, + writes: WriteService, } - impl LogQueue { - /// Create a new log queue that spawns and owns the background worker. - /// - /// The worker handles flush coordination; row preparation (project registration, - /// destination resolution) happens inside `flush_internal` at flush time. - #[allow(clippy::too_many_arguments)] - pub fn new( - config: LogQueueConfig, - login_state: LoginState, - client: reqwest::Client, - worker_queue_size: usize, - api_url: Url, - app_url: Url, - ) -> Self { - let core = LogQueueCore::new( - config, - login_state.clone(), - client.clone(), - api_url, - app_url, + pub fn new(config: LogQueueConfig, client: reqwest::Client, _worker_queue_size: usize) -> Self { + let capacity = config.queue_max_size(); + let core = LogQueueCore::new(config, client); + let processor = core.clone(); + let writes = WriteService::new( + capacity, + tokio::runtime::Handle::current(), + move |commands| { + let processor = processor.clone(); + Box::pin(async move { processor.process(commands).await }) + }, ); - - let (worker_sender, worker_receiver) = mpsc::channel(worker_queue_size.max(32)); - - let worker_handle = Arc::new(tokio::spawn(super::worker::run_worker( - worker_receiver, - core.clone(), - ))); - - Self { - core, - worker_sender, - worker_handle, - } - } - - #[cfg(test)] - pub(crate) fn enforces_queue_size_limit(&self) -> bool { - self.core.config.enforce_queue_size_limit() + Self { core, writes } } - /// Submit a span payload for processing. - /// - /// The command is pushed directly into the lock-free queue. All drops - /// (queue-full and channel-closed) go through `LogQueueCore::push`, which - /// handles drop counting, throttled warnings, and optional debug dumps - /// uniformly. Row preparation (project registration, destination resolution) - /// is deferred to flush time. pub fn submit( &self, - token: impl Into, + scope: CredentialScope, payload: SpanPayload, parent_info: Option, - ) { - self.core.push(SubmitCommand { - token: token.into(), + outcome: Option, + ) -> Completion { + let command = SubmitCommand { + scope, payload, parent_info, - }); - // Signal the worker to flush. try_send is non-blocking and best-effort: - // if the channel is full the worker already has a flush pending. - if !self.core.config.sync_flush() { - let _ = self.worker_sender.try_send(LogCommand::TriggerFlush); - } - } - - /// Flush all pending data through the worker, waiting for completion. - /// - /// Sends a `Flush` command through the worker channel, which ensures any - /// previously submitted rows (still in the worker's mpsc queue) are processed - /// before the flush begins. This is the correct flush path for external callers. - pub async fn flush_all(&self) -> Result<()> { - let (tx, rx) = oneshot::channel(); - self.worker_sender - .send(LogCommand::Flush(tx)) - .await - .map_err(|_| BraintrustError::ChannelClosed)?; - rx.await - .map_err(|_| BraintrustError::ChannelClosed)? - .map_err(|error| { - error.downcast::().unwrap_or_else(|error| { - error - .downcast::() - .map(BraintrustError::Http) - .unwrap_or_else(|error| BraintrustError::Background(error.to_string())) - }) + }; + self.writes + .enqueue_with_rejection(command, outcome, |command, error| { + self.core.dropped_count.fetch_add(1, Ordering::Relaxed); + if matches!(error, crate::EnqueueError::Full) { + self.core.log_drop_warning(); + } else { + warn!(%error, "span submission rejected"); + } + self.core.dump_dropped_row_if_configured(command); }) } - /// Return true if the lock-free row queue is approximately empty. - /// - /// Note: this inspects only the lock-free queue, not the worker's mpsc backlog, - /// so it may briefly report empty while a submit is in flight to the worker. + pub fn flush(&self) -> Completion { + self.writes.flush() + } + pub fn shutdown( + &self, + timeout: std::time::Duration, + ) -> Completion { + self.writes.shutdown(timeout) + } + #[cfg(test)] pub(crate) fn is_empty(&self) -> bool { - self.core.is_empty() + self.writes.is_empty() } - - /// Trigger a non-blocking background flush via the worker. pub async fn trigger_flush_command(&self) -> Result<()> { - self.worker_sender - .send(LogCommand::TriggerFlush) - .await - .map_err(|_| BraintrustError::ChannelClosed) + let token = self.flush(); + if let Err(error) = &*token { + return Err(error.clone().into()); + } + Ok(()) } } -impl Drop for LogQueue { - fn drop(&mut self) { - // Skip if the lock-free queue is visibly empty. Credentials travel with each - // queued row, so a flush works regardless of shared login state. - // - // Note: `is_empty()` only inspects the lock-free row queue, not the worker's - // mpsc channel. Items submitted but not yet processed by the worker would be - // missed by this guard. In practice the window is very small (a race between - // Drop and the worker task), and when it does occur the worker's own "channel - // closed" handler performs a final flush. Callers that need a hard guarantee - // should call `flush_all()` before dropping. - if self.core.is_empty() { - return; - } - - // Use flush_all() rather than core.flush() directly. Routing through the - // worker ensures any Submit commands ahead of this Drop in the mpsc queue are - // processed before the flush begins. - // - // block_in_place panics on current-thread runtimes (e.g. #[tokio::test] - // default). In that case we skip the flush here — callers on a current-thread - // runtime should call flush_all() explicitly before dropping. - if let Ok(handle) = tokio::runtime::Handle::try_current() { - // `block_in_place` panics on a current-thread runtime. In that case, - // skip the synchronous flush; dropping `worker_sender` closes the - // command channel and the worker performs a final flush. - if handle.runtime_flavor() != tokio::runtime::RuntimeFlavor::MultiThread { - return; - } - tokio::task::block_in_place(|| { - handle.block_on(async { - let _ = self.flush_all().await; - }); - }); - } else { - let rt = tokio::runtime::Runtime::new().unwrap(); - rt.block_on(async { - let _ = self.flush_all().await; - }); - } +#[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 submit( + &self, + scope: CredentialScope, + payload: SpanPayload, + parent_info: Option, + outcome: Option, + ) -> Completion { + LogQueue::submit(self, scope, payload, parent_info, outcome) } } #[cfg(test)] mod tests { use super::*; - use crate::log_queue::config::LogQueueConfig; - - fn make_test_cmd(id: &str) -> SubmitCommand { - SubmitCommand { - token: "test-token".to_string(), - payload: SpanPayload { - row_id: id.to_string(), - span_id: id.to_string(), - is_merge: false, - span_components: None, - org_id: "org".to_string(), - org_name: Some("test-org".to_string()), - project_name: None, - input: None, - output: None, - expected: None, - error: None, - scores: None, - metadata: None, - metrics: None, - tags: None, - context: None, - span_attributes: None, - extra: std::collections::HashMap::new(), - }, - parent_info: Some(ParentSpanInfo::Experiment { - object_id: "exp-test".to_string(), - }), + use serde_json::{json, Value}; + use url::Url; + use wiremock::{ + matchers::{header, method, path}, + Mock, MockServer, ResponseTemplate, + }; + + fn payload(id: &str, merge: bool) -> SpanPayload { + SpanPayload { + row_id: id.into(), + span_id: id.into(), + is_merge: merge, + span_components: None, + org_id: "org".into(), + org_name: Some("tenant".into()), + project_name: None, + input: Some(json!({"value": if merge { 2 } else { 1 }})), + output: None, + expected: None, + error: None, + scores: None, + metadata: None, + metrics: None, + tags: None, + context: None, + span_attributes: None, + extra: std::collections::HashMap::new(), } } - - fn make_queue(max_size: usize) -> LogQueue { - // No login state set - prevents background flushes from making real HTTP calls - // which would hang tests waiting for connections to non-existent servers. - let config = LogQueueConfig::builder() - .queue_max_size(max_size) - .enforce_queue_size_limit(true) // Enable bounding so drop tests work correctly - .sync_flush(true) // Disable auto-flush to avoid background HTTP calls - .build(); - let login_state = LoginState::new(); + fn queue(_server: &MockServer) -> LogQueue { LogQueue::new( - config, - login_state, + LogQueueConfig::builder().num_retries(0).build(), reqwest::Client::new(), - 256, - Url::parse("http://localhost").unwrap(), - Url::parse("http://localhost").unwrap(), + 32, ) } - - #[tokio::test] - async fn test_drop_does_not_panic_on_current_thread_runtime() { - let config = LogQueueConfig::builder() - .queue_max_size(10) - .enforce_queue_size_limit(true) - .sync_flush(true) - .build(); - let login_state = LoginState::new(); - login_state.set( - "test-key".to_string(), - "org".to_string(), - "test-org".to_string(), - "http://127.0.0.1:9".to_string(), - "http://127.0.0.1:9".to_string(), - ); - let queue = LogQueue::new( - config, - login_state, - reqwest::Client::new(), - 256, - Url::parse("http://localhost").unwrap(), - Url::parse("http://localhost").unwrap(), - ); - queue.core.push(make_test_cmd("1")); - // Logged in + non-empty queue: Drop takes the flush path. `#[tokio::test]` - // runs on a current-thread runtime, where this previously panicked in - // `block_in_place`. - drop(queue); + fn scope(server: &MockServer, token: &str) -> CredentialScope { + CredentialScope::new( + token.into(), + "org".into(), + Some("tenant".into()), + Url::parse(&server.uri()).unwrap(), + Url::parse(&server.uri()).unwrap(), + ) } - - #[tokio::test] - async fn test_push_queues_item() { - let queue = make_queue(10); - assert!(queue.core.is_empty()); - queue.core.push(make_test_cmd("1")); - assert_eq!(queue.core.len(), 1); + fn parent() -> Option { + Some(ParentSpanInfo::Experiment { + object_id: "experiment".into(), + }) } #[tokio::test] - async fn test_queue_depth_warning_emits_at_eighty_percent_once_per_fill_cycle() { - let queue = make_queue(10); - - for id in 1..8 { - queue.core.push(make_test_cmd(&id.to_string())); + async fn late_route_failure_fails_merged_members_and_successful_siblings_then_retires() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/version")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({}))) + .expect(2) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .and(header("authorization", "Bearer good")) + .respond_with(ResponseTemplate::new(200)) + .expect(2) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .and(header("authorization", "Bearer bad")) + .respond_with(ResponseTemplate::new(400)) + .expect(1) + .mount(&server) + .await; + let queue = queue(&server); + let first = queue.submit( + scope(&server, "good"), + payload("merged", false), + parent(), + None, + ); + let update = queue.submit( + scope(&server, "good"), + payload("merged", true), + parent(), + None, + ); + let last = queue.submit( + scope(&server, "bad"), + payload("other", false), + parent(), + None, + ); + let Err(FlushError::Delivery(failure)) = queue.flush().await else { + panic!("expected failed flush"); + }; + assert_eq!( + failure.failed, 3, + "count accepted operations before merging" + ); + for token in [first, update, last] { + let Err(WriteError::Delivery(error)) = token.await else { + panic!("expected batch error"); + }; + assert!(Arc::ptr_eq(&error, &failure.cause)); } - assert!(!queue - .core - .queue_depth_warning_emitted - .load(Ordering::Relaxed)); - - queue.core.push(make_test_cmd("8")); - assert!(queue - .core - .queue_depth_warning_emitted - .load(Ordering::Relaxed)); - - queue.core.push(make_test_cmd("9")); - assert!(queue - .core - .queue_depth_warning_emitted - .load(Ordering::Relaxed)); - - queue.core.drain_all(); - assert!(!queue - .core - .queue_depth_warning_emitted - .load(Ordering::Relaxed)); - } - - #[tokio::test] - async fn test_push_drops_when_full() { - let queue = make_queue(2); - queue.core.push(make_test_cmd("1")); - queue.core.push(make_test_cmd("2")); - assert_eq!(queue.core.len(), 2); - - // This push should be dropped since the queue is full - queue.core.push(make_test_cmd("3")); - assert_eq!(queue.core.len(), 2); - // The throttled warning fires immediately on the first drop (last_drop_log_time - // starts at 0, which is well beyond the throttle window), resetting the counter. - assert_eq!(queue.core.dropped_count(), 0); - } - - #[tokio::test] - async fn test_drop_count_accumulates_within_throttle_window() { - // With a very long logging period the counter accumulates across multiple drops - // without resetting (the throttle hasn't expired yet after the first warning fires). - let config = LogQueueConfig::builder() - .queue_max_size(1) - .enforce_queue_size_limit(true) - .sync_flush(true) - .build(); - let login_state = LoginState::new(); - let queue = LogQueue::new( - config, - login_state, - reqwest::Client::new(), - 256, - Url::parse("http://localhost").unwrap(), - Url::parse("http://localhost").unwrap(), + assert!(queue.is_empty()); + assert_eq!(queue.flush().await.unwrap().delivered, 0); + let requests = server.received_requests().await.unwrap(); + let uploads: Vec<_> = requests + .iter() + .filter(|req| req.url.path() == "/logs3") + .collect(); + assert_eq!(uploads.len(), 2); + let first_body: Value = serde_json::from_slice(&uploads[0].body).unwrap(); + assert_eq!(first_body["rows"].as_array().unwrap().len(), 1); + assert_eq!(first_body["rows"][0]["input"]["value"], 2); + let fresh = queue.submit( + scope(&server, "good"), + payload("fresh", false), + parent(), + None, ); - - queue.core.push(make_test_cmd("1")); // fills queue - - // First drop: warning fires (last_drop_log_time=0 is always past throttle), - // resets counter to 0 and records timestamp. - queue.core.push(make_test_cmd("2")); // dropped, count reset to 0 after warning - assert_eq!(queue.core.dropped_count(), 0); - - // Subsequent drops within the same throttle window accumulate without reset. - queue.core.push(make_test_cmd("3")); // dropped, counter now 1 - queue.core.push(make_test_cmd("4")); // dropped, counter now 2 - assert_eq!(queue.core.dropped_count(), 2); + assert_eq!(queue.flush().await.unwrap().delivered, 1); + fresh.await.unwrap(); + queue + .shutdown(std::time::Duration::from_secs(1)) + .await + .unwrap(); } #[tokio::test] - async fn test_unbounded_queue_never_drops() { - let config = LogQueueConfig::builder() - .enforce_queue_size_limit(false) - .sync_flush(true) - .build(); - let login_state = LoginState::new(); - let queue = LogQueue::new( - config, - login_state, - reqwest::Client::new(), - 256, - Url::parse("http://localhost").unwrap(), - Url::parse("http://localhost").unwrap(), + async fn rejected_split_sibling_fails_every_member_despite_an_earlier_committed_request() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/version")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({}))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(|request: &wiremock::Request| { + let body: Value = serde_json::from_slice(&request.body).unwrap(); + let rows = body["rows"].as_array().unwrap(); + ResponseTemplate::new(if rows.len() > 1 { + 413 + } else if rows[0]["id"] == "late-failure" { + 400 + } else { + 200 + }) + }) + .expect(3) + .mount(&server) + .await; + let queue = queue(&server); + let first = queue.submit( + scope(&server, "token"), + payload("committed", false), + parent(), + None, ); - - // Push many more items than the default queue_max_size - for i in 0..20_000 { - queue.core.push(make_test_cmd(&i.to_string())); + let last = queue.submit( + scope(&server, "token"), + payload("late-failure", false), + parent(), + None, + ); + let Err(FlushError::Delivery(failure)) = queue.flush().await else { + panic!("expected failure"); + }; + assert_eq!(failure.failed, 2); + for token in [first, last] { + assert!(matches!(token.await, Err(WriteError::Delivery(_)))); } - assert_eq!(queue.core.len(), 20_000); - assert_eq!(queue.core.dropped_count(), 0); - } - - #[tokio::test] - async fn test_drain_all_returns_all_items() { - let queue = make_queue(10); - queue.core.push(make_test_cmd("1")); - queue.core.push(make_test_cmd("2")); - queue.core.push(make_test_cmd("3")); - - let drained = queue.core.drain_all(); - assert_eq!(drained.len(), 3); - assert!(queue.core.is_empty()); - } - - #[tokio::test] - async fn test_drain_all_atomic_boundary() { - // Items pushed after drain starts should go to the new channel, not the drained batch - let queue = make_queue(10); - queue.core.push(make_test_cmd("1")); - queue.core.push(make_test_cmd("2")); - - let drained = queue.core.drain_all(); - assert_eq!(drained.len(), 2); - - // Push after drain goes to the new channel - queue.core.push(make_test_cmd("3")); - assert_eq!(queue.core.len(), 1); - } - - #[tokio::test] - async fn test_push_no_auto_flush_in_sync_mode() { - // make_queue uses sync_flush=true, so push() should not trigger background flush - let queue = make_queue(10); - - queue.core.push(make_test_cmd("1")); - // No flush should be in progress - in sync mode user calls flush() explicitly - assert!(queue.core.flush_handle.try_lock().unwrap().is_none()); - assert_eq!(queue.core.len(), 1); + assert!(queue.is_empty()); + queue.flush().await.unwrap(); + queue + .shutdown(std::time::Duration::from_secs(1)) + .await + .unwrap(); } } diff --git a/src/log_queue/worker.rs b/src/log_queue/worker.rs deleted file mode 100644 index 45b281e..0000000 --- a/src/log_queue/worker.rs +++ /dev/null @@ -1,34 +0,0 @@ -use std::sync::Arc; -use tokio::sync::{mpsc, oneshot}; -use tracing::warn; - -use super::queue::LogQueueCore; - -/// Commands sent through the worker channel (flush coordination only). -pub(super) enum LogCommand { - Flush(oneshot::Sender>), - TriggerFlush, -} - -pub(super) async fn run_worker(mut receiver: mpsc::Receiver, core: Arc) { - loop { - match receiver.recv().await { - Some(LogCommand::Flush(response)) => { - let result = core.flush().await; - let _ = response.send(result); - } - Some(LogCommand::TriggerFlush) => { - if let Err(e) = core.flush().await { - warn!(error = %e, "background flush failed"); - } - } - None => { - // Channel closed — flush any remaining items and exit. - if let Err(e) = core.flush().await { - warn!(error = %e, "final flush failed"); - } - break; - } - } - } -} diff --git a/src/logger.rs b/src/logger.rs index 5c69ebf..3699e2f 100644 --- a/src/logger.rs +++ b/src/logger.rs @@ -2,25 +2,21 @@ use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; -use serde::{Deserialize, Serialize}; +use serde::Serialize; use serde_json::{Map, Value}; use tokio::sync::Notify; use url::Url; -use crate::dataset::{ - BTQLQuery, BTQLResponse, DatasetBuilder, DatasetFetcher, DatasetRegisterRequest, - DatasetRegisterResponse, DatasetRegistrar, DatasetSummarizer, DatasetSummaryResponse, -}; +use crate::dataset::DatasetBuilder; use crate::error::{BraintrustError, Result}; -use crate::experiments::api::{ - BaseExperimentFetcher, BaseExperimentRequest, BaseExperimentResponse, - ExperimentComparisonFetcher, ExperimentComparisonResponse, ExperimentRegisterRequest, - ExperimentRegisterResponse, ExperimentRegistrar, -}; -use crate::experiments::{BaseExperimentInfo, ExperimentBuilder}; +use crate::experiments::ExperimentBuilder; use crate::http::{build_http_client, CaBundleConfig}; +use crate::internal::ports::SpanSubmitter; +use crate::internal::resources::ResourceTransport; +use crate::internal::retry::{exponential, RetryExecutor}; +use crate::internal::session::CredentialScope; use crate::log_queue::{LogQueue, LogQueueConfig}; -use crate::span::{merge_span_origin_context, SpanOrigin, SpanOriginEnvironment, SpanSubmitter}; +use crate::span::{merge_span_origin_context, SpanOrigin, SpanOriginEnvironment}; use crate::span_components::SpanComponents; use crate::types::{ParentSpanInfo, SpanAttributes, SpanEventData, SpanObjectType, SpanPayload}; @@ -29,40 +25,10 @@ use crate::types::{ParentSpanInfo, SpanAttributes, SpanEventData, SpanObjectType const REQUEST_TIMEOUT: Duration = Duration::from_secs(30); const LOGIN_TIMEOUT: Duration = Duration::from_secs(30); const DEFAULT_QUEUE_SIZE: usize = 256; -pub const DEFAULT_APP_URL: &str = "https://www.braintrust.dev"; -pub const DEFAULT_API_URL: &str = "https://api.braintrust.dev"; - -/// Organization info returned from login. -#[derive(Debug, Clone, Deserialize)] -pub struct OrgInfo { - id: String, - name: String, - #[serde(default)] - api_url: Option, -} - -impl OrgInfo { - /// The organization's unique ID. - pub fn id(&self) -> &str { - &self.id - } - - /// The organization's name. - pub fn name(&self) -> &str { - &self.name - } - - /// The organization's API URL, if set. - pub fn api_url(&self) -> Option<&str> { - self.api_url.as_deref() - } -} +pub use crate::api::{DEFAULT_API_URL, DEFAULT_APP_URL}; -/// Response from the login endpoint. -#[derive(Debug, Deserialize)] -struct LoginResponse { - org_info: Vec, -} +use crate::api::login::LoginResponse; +pub use crate::api::login::OrgInfo; /// Logged-in state containing API key and org info. /// Handles internal synchronization for the transition from not-logged-in to logged-in. @@ -71,7 +37,7 @@ pub struct LoginState { inner: Arc>, } -#[derive(Debug, Clone)] +#[derive(Clone)] struct LoginStateInner { api_key: String, org_id: String, @@ -115,6 +81,25 @@ impl LoginState { .is_ok() } + /// Capture the selected organization and routes together, without fallback. + fn scope(&self) -> Result { + let state = self + .inner + .get() + .ok_or_else(|| BraintrustError::InvalidConfig("Not logged in".into()))?; + let api_url = Url::parse(&state.api_url) + .map_err(|_| BraintrustError::InvalidConfig("invalid login API URL".into()))?; + let app_url = Url::parse(&state.app_url) + .map_err(|_| BraintrustError::InvalidConfig("invalid login app URL".into()))?; + Ok(CredentialScope::new( + state.api_key.clone(), + state.org_id.clone(), + Some(state.org_name.clone()), + api_url, + app_url, + )) + } + /// Check if logged in. pub fn is_logged_in(&self) -> bool { self.inner.get().is_some() @@ -148,13 +133,11 @@ impl LoginState { impl std::fmt::Debug for LoginState { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self.inner.get() { - Some(inner) => f.debug_struct("LoginState").field("inner", inner).finish(), - None => f - .debug_struct("LoginState") - .field("inner", &"") - .finish(), - } + f.debug_struct("LoginState") + .field("is_logged_in", &self.is_logged_in()) + .field("org_id", &self.org_id()) + .field("org_name", &self.org_name()) + .finish_non_exhaustive() } } @@ -231,7 +214,7 @@ impl BraintrustClientBuilder { batch_max_items: None, batch_max_bytes: None, queue_max_size: None, - enforce_queue_size_limit: false, + enforce_queue_size_limit: true, checked_delivery: false, environment: None, span_origin: None, @@ -276,17 +259,16 @@ impl BraintrustClientBuilder { self } - /// Enable or disable enforcement of the queue size limit. - /// - /// When enabled, the queue drops newly arriving events once it reaches - /// `queue_max_size`. Defaults to false, leaving the queue unbounded. + /// 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 } - /// Make `flush` report preparation and HTTP delivery failures. Automatic - /// flushing is disabled so the explicit caller observes every result. + /// 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 @@ -403,20 +385,10 @@ impl BraintrustClientBuilder { .maybe_queue_max_size(self.queue_max_size) .enforce_queue_size_limit(self.enforce_queue_size_limit) .checked_delivery(self.checked_delivery) - .maybe_sync_flush(self.checked_delivery.then_some(true)) .build(); - // LogQueue owns the background worker. It is given the client-configured - // URLs so that flushes with per-request credentials work without relying on - // the shared (set-once) login state to hold URLs. - let queue = LogQueue::new( - log_config, - login_state.clone(), - http_client.clone(), - self.queue_size, - api_url.clone(), - app_url.clone(), - ); + // Commands carry their own immutable scope; the worker has no session fallback. + let queue = LogQueue::new(log_config, http_client.clone(), self.queue_size); let client = BraintrustClient { inner: Arc::new(ClientInner { @@ -490,6 +462,29 @@ impl std::fmt::Debug for ClientInner { } impl BraintrustClient { + fn resource_transport(&self) -> Arc { + Arc::new(ResourceTransport::new(self.inner.http_client.clone())) + } + + fn span_submitter(&self) -> Arc { + Arc::new(self.inner.queue.clone()) + } + + fn explicit_scope( + &self, + token: String, + org_id: String, + org_name: Option, + ) -> CredentialScope { + CredentialScope::new( + token, + org_id, + org_name, + self.inner.api_url.clone(), + self.inner.app_url.clone(), + ) + } + /// Create a new builder for configuring the client. /// /// # Example @@ -551,20 +546,14 @@ impl BraintrustClient { /// # Ok(()) /// # } /// ``` - pub async fn span_builder(&self) -> Result> { + pub async fn span_builder(&self) -> Result { if self.inner.login_skipped { return Err(BraintrustError::InvalidConfig( "span_builder() requires login; use span_builder_with_credentials() when skip_login is enabled".into(), )); } let state = self.wait_for_login_state().await?; - let api_key = state - .api_key() - .ok_or_else(|| BraintrustError::InvalidConfig("Not logged in".into()))?; - let org_id = state - .org_id() - .ok_or_else(|| BraintrustError::InvalidConfig("Not logged in".into()))?; - let mut builder = crate::span::SpanBuilder::new(Arc::new(self.clone()), &api_key, &org_id); + let mut builder = crate::span::SpanBuilder::new(self.span_submitter(), state.scope()?); if let Some(environment) = &self.inner.environment { builder = builder.span_origin_environment(environment.clone()); } @@ -580,37 +569,49 @@ impl BraintrustClient { /// Create a dataset builder using the logged-in state. /// /// This waits for login to complete if it hasn't already. - pub async fn dataset_builder(&self) -> Result> { + pub async fn dataset_builder(&self) -> Result { + if self.inner.login_skipped { + return Err(BraintrustError::InvalidConfig("dataset_builder() requires login; use dataset_builder_with_credentials() when skip_login is enabled".into())); + } let state = self.wait_for_login_state().await?; - let api_key = state - .api_key() - .ok_or_else(|| BraintrustError::InvalidConfig("Not logged in".into()))?; - let org_name = state - .org_name() - .ok_or_else(|| BraintrustError::InvalidConfig("Not logged in".into()))?; Ok(DatasetBuilder::new( - Arc::new(self.clone()), - api_key, - org_name, + self.resource_transport(), + Arc::new(self.inner.queue.clone()), + state.scope()?, )) } + /// Create a dataset builder with immutable explicit credentials, without + /// logging in or changing the client's shared session. + pub fn dataset_builder_with_credentials( + &self, + token: impl Into, + org_name: impl Into, + ) -> DatasetBuilder { + DatasetBuilder::new( + self.resource_transport(), + Arc::new(self.inner.queue.clone()), + self.explicit_scope(token.into(), String::new(), Some(org_name.into())), + ) + } + /// Create a span builder with explicit token and org_id. /// /// Use this for per-session / multi-tenant logging: the supplied credentials - /// travel with each span through the upload pipeline and are NOT written into - /// the client's shared login state. Spans created from the same client with + /// and configured endpoint routes are captured by the builder and travel with + /// each span. They do not change the client's shared login state. Spans created from the same client with /// different `(token, org_id)` pairs each upload with their own credentials. pub fn span_builder_with_credentials( &self, token: impl Into, org_id: impl Into, - ) -> crate::span::SpanBuilder { + ) -> crate::span::SpanBuilder { let token = token.into(); let org_id = org_id.into(); - let submitter = Arc::new(self.clone()); - let mut builder = crate::span::SpanBuilder::new(submitter, token, org_id); + let submitter = self.span_submitter(); + let mut builder = + crate::span::SpanBuilder::new(submitter, self.explicit_scope(token, org_id, None)); if let Some(environment) = &self.inner.environment { builder = builder.span_origin_environment(environment.clone()); } @@ -628,19 +629,13 @@ impl BraintrustClient { self.inner.login_state.clone() } else if self.inner.login_skipped { return Err(BraintrustError::InvalidConfig( - "update_span() requires credentials; call span_builder_with_credentials() or experiment_builder_with_credentials() first when skip_login is enabled".into(), + "update_span() requires a logged-in client; use update_span_with_credentials() when skip_login is enabled".into(), )); } else { self.wait_for_login_state().await? }; - let api_key = login_state - .api_key() - .ok_or_else(|| BraintrustError::InvalidConfig("Not logged in".into()))?; - let org_id = login_state - .org_id() - .ok_or_else(|| BraintrustError::InvalidConfig("Not logged in".into()))?; - self.update_span_internal(api_key, org_id, login_state.org_name(), exported, event) + self.update_span_internal(login_state.scope()?, exported, event) } /// Update an existing span using explicit credentials instead of shared login state. @@ -658,14 +653,12 @@ impl BraintrustClient { let token = token.into(); let org_id = org_id.into(); - self.update_span_internal(token, org_id, None, exported, event) + self.update_span_internal(self.explicit_scope(token, org_id, None), exported, event) } fn update_span_internal( &self, - token: String, - org_id: String, - org_name: Option, + scope: CredentialScope, exported: &str, event: crate::span::SpanLog, ) -> Result<()> { @@ -745,8 +738,8 @@ impl BraintrustClient { span_id, is_merge: true, span_components: Some(components), - org_id, - org_name, + org_id: scope.org_id().to_owned(), + org_name: scope.org_name().map(str::to_owned), project_name: None, input: event_data.input, output: event_data.output, @@ -761,7 +754,10 @@ impl BraintrustClient { extra: event_data.extra, }; - self.submit_payload(token, payload, None); + let completion = self.submit_payload(scope, payload, None); + if let Err(error) = &*completion { + return Err(BraintrustError::Background(error.to_string())); + } Ok(()) } @@ -837,37 +833,33 @@ impl BraintrustClient { fn start_background_login(&self, api_key: String, org_name: Option) { let client = self.clone(); tokio::spawn(async move { - // Retry with exponential backoff up to a fixed maximum number of attempts. - const MAX_ATTEMPTS: u32 = 10; - let mut delay = Duration::from_millis(100); - let max_delay = Duration::from_secs(5); - - for attempt in 1..=MAX_ATTEMPTS { - match client.perform_login(&api_key, org_name.as_deref()).await { - Ok(()) => { - tracing::debug!("Background login completed successfully"); - return; - } - Err(e) if attempt < MAX_ATTEMPTS => { - tracing::warn!( - attempt, - max = MAX_ATTEMPTS, - "Background login failed: {}, retrying in {:?}", - e, - delay - ); - tokio::time::sleep(delay).await; - delay = (delay * 2).min(max_delay); - } - Err(e) => { - tracing::error!( - "Background login failed after {} attempts: {}. \ - Spans will not be uploaded until the client is re-created.", - MAX_ATTEMPTS, - e - ); - } - } + const MAX_ATTEMPTS: usize = 10; + let outcome = RetryExecutor::new( + MAX_ATTEMPTS - 1, + exponential(Duration::from_millis(100), Duration::from_secs(5)), + ) + .run( + |_| client.perform_login(&api_key, org_name.as_deref()), + |_| true, // Preserve the existing background-login policy. + |attempt, error, delay| { + tracing::warn!( + attempt = attempt.index + 1, + max = MAX_ATTEMPTS, + "Background login failed: {}, retrying in {:?}", + error, + delay + ); + }, + ) + .await; + match outcome.result { + Ok(()) => tracing::debug!("Background login completed successfully"), + Err(error) => tracing::error!( + "Background login failed after {} attempts: {}. \ + Spans will not be uploaded until the client is re-created.", + outcome.attempts, + error + ), } }); } @@ -909,8 +901,7 @@ impl BraintrustClient { /// Create an experiment builder with explicit credentials. /// - /// Note: This will be deprecated in favor of `experiment_builder()` once - /// the login-flow branch lands and credentials are stored in the client. + /// Credentials belong to this resource and never populate shared login state. /// /// # Example /// @@ -925,41 +916,41 @@ impl BraintrustClient { &self, token: impl Into, org_id: impl Into, - ) -> ExperimentBuilder { + ) -> ExperimentBuilder { let token = token.into(); let org_id = org_id.into(); - // Populate login state with explicit credentials so flush_internal() can send data. - let _ = self.inner.login_state.set( - token.clone(), - org_id.clone(), - String::new(), // org_name unknown without login - self.inner.api_url.to_string(), - self.inner.app_url.to_string(), - ); - - let submitter = Arc::new(self.clone()); - ExperimentBuilder::new(submitter, token, org_id) + let submitter = self.span_submitter(); + ExperimentBuilder::new( + self.resource_transport(), + submitter, + self.explicit_scope(token, org_id, None), + ) } - /// Submit a span payload for logging (fire-and-forget). - /// - /// Returns immediately after queuing. HTTP submission happens in the background. - /// Errors are logged as warnings but not propagated to callers. - /// Submit a span payload synchronously (fire-and-forget). - /// The payload is queued internally and processed by a background worker. + /// Submit immediately; the token observes admission and full batch delivery. pub(crate) fn submit_payload( &self, - token: impl Into, + scope: CredentialScope, payload: SpanPayload, parent_info: Option, - ) { - self.inner.queue.submit(token, payload, parent_info) + ) -> crate::Completion { + self.inner.queue.submit(scope, payload, parent_info, None) + } + + /// Stop accepting writes, drain accepted work and stop the delivery worker. + /// A timeout fails remaining completions and discards active SDK state. + pub fn shutdown( + &self, + timeout: std::time::Duration, + ) -> crate::Completion { + self.inner.queue.shutdown(timeout) } - /// Flush all pending log events. - pub async fn flush(&self) -> Result<()> { - self.inner.queue.flush_all().await + /// Capture a cutoff immediately and observe delivery of the covered batch. + /// New writes belong to later work; completed failures are never replayed. + pub fn flush(&self) -> crate::Completion { + self.inner.queue.flush() } /// Trigger a non-blocking background flush. @@ -969,332 +960,6 @@ impl BraintrustClient { } } -impl Drop for BraintrustClient { - fn drop(&mut self) { - // Only flush on the last client reference, and only when there is pending - // data. `strong_count == 1` means no other `BraintrustClient` clones are - // alive. We check queue emptiness rather than login state so that - // per-request-credential (skip_login / multi-tenant) callers, which never - // populate the shared login state, still get their pending spans flushed. - if Arc::strong_count(&self.inner) != 1 || self.inner.queue.is_empty() { - return; - } - - // Flush via flush_all(), which routes through the background worker's mpsc channel. - // This ensures any Submit commands still sitting in the worker's backlog are processed - // before the flush begins — resolving the race in LogQueue::Drop where is_empty() only - // inspects the lock-free crossbeam queue, not the worker channel. - // After this flush completes, LogQueue::Drop will find an empty queue and exit early. - // - // block_in_place panics on current-thread runtimes (e.g. #[tokio::test] default). - // In that case we skip the flush here; LogQueue::Drop provides a best-effort fallback - // for the lock-free queue, and tests should always call flush() explicitly. - let inner = self.inner.clone(); - if let Ok(handle) = tokio::runtime::Handle::try_current() { - if handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread { - tokio::task::block_in_place(|| { - handle.block_on(async move { - let _ = inner.queue.flush_all().await; - }); - }); - } - } else { - // No async runtime available — spin up a temporary one. - let rt = tokio::runtime::Runtime::new().unwrap(); - rt.block_on(async move { - let _ = inner.queue.flush_all().await; - }); - } - } -} - -#[async_trait::async_trait] -impl SpanSubmitter for BraintrustClient { - fn submit(&self, token: String, payload: SpanPayload, parent_info: Option) { - self.submit_payload(token, payload, parent_info) - } - - async fn flush(&self) -> Result<()> { - BraintrustClient::flush(self).await - } - - async fn trigger_flush(&self) -> Result<()> { - BraintrustClient::trigger_flush(self).await - } -} - -// Experiment trait implementations for BraintrustClient -#[async_trait::async_trait] -impl ExperimentRegistrar for BraintrustClient { - async fn register_experiment( - &self, - token: &str, - request: ExperimentRegisterRequest, - ) -> Result { - let url = self - .inner - .app_url - .join("api/experiment/register") - .map_err(|e| { - BraintrustError::InvalidConfig(format!("invalid experiment register url: {e}")) - })?; - - let response = self - .inner - .http_client - .post(url) - .bearer_auth(token) - .json(&request) - .send() - .await?; - - let status = response.status(); - if !status.is_success() { - let text = response.text().await.unwrap_or_default(); - return Err(BraintrustError::Api { - status: status.as_u16(), - message: format!("register experiment failed: {text}"), - }); - } - - let register_response: ExperimentRegisterResponse = response.json().await.map_err(|e| { - BraintrustError::Background(format!( - "failed to parse experiment registration response: {e}" - )) - })?; - - Ok(register_response) - } -} - -#[async_trait::async_trait] -impl BaseExperimentFetcher for BraintrustClient { - async fn fetch_base_experiment( - &self, - token: &str, - experiment_id: &str, - ) -> Result> { - let url = self - .inner - .app_url - .join("api/base_experiment/get_id") - .map_err(|e| { - BraintrustError::InvalidConfig(format!("invalid base experiment url: {e}")) - })?; - - let request = BaseExperimentRequest { - id: experiment_id.to_string(), - }; - - let response = self - .inner - .http_client - .post(url) - .bearer_auth(token) - .json(&request) - .send() - .await?; - - let status = response.status(); - - // 400 means no base experiment - return None - if status.as_u16() == 400 { - return Ok(None); - } - - if !status.is_success() { - let text = response.text().await.unwrap_or_default(); - return Err(BraintrustError::Api { - status: status.as_u16(), - message: format!("fetch base experiment failed: {text}"), - }); - } - - let base_response: BaseExperimentResponse = response.json().await.map_err(|e| { - BraintrustError::Background(format!("failed to parse base experiment response: {e}")) - })?; - - match (base_response.base_exp_id, base_response.base_exp_name) { - (Some(id), Some(name)) => Ok(Some(BaseExperimentInfo { id, name })), - _ => Ok(None), - } - } -} - -#[async_trait::async_trait] -impl ExperimentComparisonFetcher for BraintrustClient { - async fn fetch_experiment_comparison( - &self, - token: &str, - experiment_id: &str, - base_experiment_id: Option<&str>, - ) -> Result { - let mut url = self - .inner - .app_url - .join("experiment-comparison2") - .map_err(|e| { - BraintrustError::InvalidConfig(format!("invalid experiment comparison url: {e}")) - })?; - - url.query_pairs_mut() - .append_pair("experiment_id", experiment_id); - - if let Some(base_id) = base_experiment_id { - url.query_pairs_mut() - .append_pair("base_experiment_id", base_id); - } - - let response = self - .inner - .http_client - .get(url) - .bearer_auth(token) - .send() - .await?; - - let status = response.status(); - if !status.is_success() { - let text = response.text().await.unwrap_or_default(); - return Err(BraintrustError::Api { - status: status.as_u16(), - message: format!("fetch experiment comparison failed: {text}"), - }); - } - - let comparison: ExperimentComparisonResponse = response.json().await.map_err(|e| { - BraintrustError::Background(format!( - "failed to parse experiment comparison response: {e}" - )) - })?; - - Ok(comparison) - } -} - -// Dataset trait implementations for BraintrustClient -#[async_trait::async_trait] -impl DatasetRegistrar for BraintrustClient { - async fn register_dataset( - &self, - token: &str, - request: DatasetRegisterRequest, - ) -> Result { - let url = self - .inner - .app_url - .join("api/dataset/register") - .map_err(|e| BraintrustError::InvalidConfig(e.to_string()))?; - - let response = self - .inner - .http_client - .post(url) - .bearer_auth(token) - .json(&request) - .send() - .await - .map_err(|e| BraintrustError::Network(e.to_string()))?; - - if !response.status().is_success() { - let status = response.status(); - let body = response.text().await.unwrap_or_default(); - return Err(BraintrustError::Api { - status: status.as_u16(), - message: body, - }); - } - - response - .json::() - .await - .map_err(|e| BraintrustError::Api { - status: 200, - message: format!("Failed to parse dataset register response: {}", e), - }) - } -} - -#[async_trait::async_trait] -impl DatasetFetcher for BraintrustClient { - async fn fetch_dataset_records(&self, token: &str, query: BTQLQuery) -> Result { - let url = self - .inner - .api_url - .join("btql") - .map_err(|e| BraintrustError::InvalidConfig(e.to_string()))?; - - let response = self - .inner - .http_client - .post(url) - .bearer_auth(token) - .json(&query) - .send() - .await - .map_err(|e| BraintrustError::Network(e.to_string()))?; - - if !response.status().is_success() { - let status = response.status(); - let body = response.text().await.unwrap_or_default(); - return Err(BraintrustError::Api { - status: status.as_u16(), - message: body, - }); - } - - response - .json::() - .await - .map_err(|e| BraintrustError::Api { - status: 200, - message: format!("Failed to parse BTQL response: {}", e), - }) - } -} - -#[async_trait::async_trait] -impl DatasetSummarizer for BraintrustClient { - async fn fetch_dataset_summary( - &self, - token: &str, - dataset_id: &str, - ) -> Result { - let mut url = self - .inner - .api_url - .join("dataset-summary") - .map_err(|e| BraintrustError::InvalidConfig(e.to_string()))?; - url.query_pairs_mut().append_pair("dataset_id", dataset_id); - - let response = self - .inner - .http_client - .get(url) - .bearer_auth(token) - .send() - .await - .map_err(|e| BraintrustError::Network(e.to_string()))?; - - if !response.status().is_success() { - let status = response.status(); - let body = response.text().await.unwrap_or_default(); - return Err(BraintrustError::Api { - status: status.as_u16(), - message: body, - }); - } - - response - .json::() - .await - .map_err(|e| BraintrustError::Api { - status: 200, - message: format!("Failed to parse dataset summary response: {}", e), - }) - } -} - #[allow(dead_code)] #[derive(Debug, Default)] struct MutableSpanEvent { @@ -1449,6 +1114,23 @@ fn parse_string_array(value: Value) -> Option> { #[cfg(test)] mod tests { + #[test] + fn login_state_debug_never_exposes_tokens_or_unvalidated_endpoints() { + let state = super::LoginState::new(); + state.set( + "private-token".into(), + "org".into(), + "tenant".into(), + "https://user:password@example.com/?api_key=secret".into(), + "https://example.com/?token=secret".into(), + ); + let debug = format!("{state:?}"); + assert!(debug.contains("tenant")); + for secret in ["private-token", "password", "secret"] { + assert!(!debug.contains(secret)); + } + } + use super::*; use crate::span::SpanLog; use serde_json::Value; @@ -1464,6 +1146,32 @@ mod tests { ResponseTemplate::new(200).set_body_json(serde_json::json!({ "org_info": org_info })) } + #[tokio::test] + async fn resource_ports_do_not_retain_the_whole_client_or_login_state() { + let client = BraintrustClient::builder() + .skip_login(true) + .build() + .await + .unwrap(); + let weak = Arc::downgrade(&client.inner); + let span = client.span_builder_with_credentials("token", "org"); + let dataset = client + .dataset_builder_with_credentials("token", "tenant") + .project_name("p") + .dataset_name("d"); + let experiment = client + .experiment_builder_with_credentials("token", "org") + .project_name("p"); + drop(client); + assert!( + weak.upgrade().is_none(), + "resource dependencies must not retain ClientInner" + ); + let _span = span.build(); + let _dataset = dataset.build().unwrap(); + let _experiment = experiment.build().unwrap(); + } + #[tokio::test] async fn builder_rejects_missing_api_key() { // Clear any env vars that might be set @@ -1563,17 +1271,23 @@ mod tests { } #[tokio::test] - async fn builder_configures_queue_size_limit_enforcement() { + async fn builder_capacity_is_enforced_even_with_the_legacy_unbounded_setting() { let client = BraintrustClient::builder() .app_url("https://example.com") .api_url("https://example.com") .skip_login(true) - .enforce_queue_size_limit(true) + .with_queue_max_size(1) + .enforce_queue_size_limit(false) .build() .await .expect("client"); - assert!(client.inner.queue.enforces_queue_size_limit()); + let span = client.span_builder_with_credentials("token", "org").build(); + assert!(span.log(crate::SpanLog::default()).is_ok()); + assert!(matches!( + &*span.log(crate::SpanLog::default()), + Err(crate::WriteError::Admission(crate::EnqueueError::Full)) + )); } #[tokio::test] diff --git a/src/logs3.rs b/src/logs3.rs index b5966b2..1582f31 100644 --- a/src/logs3.rs +++ b/src/logs3.rs @@ -1,32 +1,12 @@ -use std::collections::HashMap; -use std::time::Duration; - -use backoff::backoff::Backoff; -use backoff::ExponentialBackoffBuilder; -use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; -use reqwest::{multipart, Client, Url}; -use serde::{Deserialize, Serialize}; -use serde_json::{json, Map, Value}; -use tokio::time::sleep; - use crate::error::{BraintrustError, Result}; use crate::http::{build_http_client, CaBundleConfig}; -use crate::log_queue::batch_items; +use crate::internal::ingest::{EncodedBatch, EncodedRow}; +use crate::internal::upload::{Capabilities, UploadCoordinator, UploadPolicy}; +use reqwest::{Client, Url}; +use serde_json::{Map, Value}; +use std::time::Duration; const DEFAULT_MAX_REQUEST_SIZE: usize = 6 * 1024 * 1024; -const LOGS_API_VERSION: u8 = 2; -const LOGS3_OVERFLOW_REFERENCE_TYPE: &str = "logs3_overflow"; -const MAX_ATTEMPTS: usize = 5; -const BASE_BACKOFF_MS: u64 = 300; -const MAX_BACKOFF_MS: u64 = 8_000; -const OBJECT_ID_KEYS: [&str; 6] = [ - "experiment_id", - "dataset_id", - "prompt_session_id", - "project_id", - "log_id", - "function_data", -]; #[derive(Debug, Clone, Default)] pub struct Logs3UploadResult { @@ -42,56 +22,7 @@ pub struct Logs3BatchUploader { api_key: String, org_name: Option, max_request_size_override: Option, - payload_limits: Option, -} - -#[derive(Debug, Clone)] -struct PayloadLimits { - max_request_size: usize, - can_use_overflow: bool, -} - -#[derive(Debug, Clone)] -struct RowPayload { - row_json: String, - row_bytes: usize, - overflow_meta: Logs3OverflowInputRow, -} - -#[derive(Debug, Clone, Serialize)] -struct Logs3OverflowInputRow { - object_ids: Map, - #[serde(skip_serializing_if = "Option::is_none")] - is_delete: Option, - input_row: OverflowInputRowInfo, -} - -#[derive(Debug, Clone, Serialize)] -struct OverflowInputRowInfo { - byte_size: usize, -} - -#[derive(Debug, Deserialize)] -struct VersionInfo { - #[serde(default)] - logs3_payload_max_bytes: Option, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "UPPERCASE")] -enum OverflowMethod { - Put, - Post, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -struct Logs3OverflowUpload { - method: OverflowMethod, - signed_url: String, - headers: Option>, - fields: Option>, - key: String, + payload_limits: Option, } impl Logs3BatchUploader { @@ -158,398 +89,45 @@ impl Logs3BatchUploader { return Ok(Logs3UploadResult::default()); } - let limits = self.get_payload_limits().await; - let mut row_payloads = Vec::with_capacity(rows.len()); - for row in rows { - row_payloads.push(RowPayload::from_row(row)?); - } - - let batch_limit_rows = batch_max_num_items.max(1); - let batch_limit_bytes = (limits.max_request_size / 2).max(1); - let batches = batch_items(row_payloads, batch_limit_rows, batch_limit_bytes, |item| { - item.row_bytes - }); + let coordinator = UploadCoordinator::new( + self.client.clone(), + &self.api_url, + &self.api_key, + self.org_name.as_deref(), + UploadPolicy::Standalone, + )?; + let limits = match self.payload_limits { + Some(limits) => limits, + None => { + let limits = coordinator + .capabilities(DEFAULT_MAX_REQUEST_SIZE, self.max_request_size_override) + .await; + self.payload_limits = Some(limits); + limits + } + }; + let row_payloads = rows + .iter() + .map(EncodedRow::from_map) + .collect::, _>>() + .map_err(|error| { + BraintrustError::InvalidConfig(format!("failed to serialize logs3 row: {error}")) + })?; + let batches = EncodedBatch::partition( + row_payloads, + batch_max_num_items.max(1), + (limits.max_request_size / 2).max(1), + ); let mut result = Logs3UploadResult::default(); for batch in batches { - let batch_result = self.submit_row_batch(&batch, &limits).await?; + let batch_result = coordinator.upload(batch, limits).await?; result.bytes_processed += batch_result.bytes_processed; result.requests_sent += batch_result.requests_sent; result.rows_uploaded += batch_result.rows_uploaded; } Ok(result) } - - async fn get_payload_limits(&mut self) -> PayloadLimits { - if let Some(existing) = &self.payload_limits { - return existing.clone(); - } - let fetched = self.fetch_payload_limits().await; - self.payload_limits = Some(fetched.clone()); - fetched - } - - async fn fetch_payload_limits(&self) -> PayloadLimits { - let mut server_limit: Option = None; - if let Ok(url) = self.api_url.join("version") { - let mut request = self.client.get(url).bearer_auth(&self.api_key); - if let Some(org_name) = &self.org_name { - request = request.header("x-bt-org-name", org_name); - } - if let Ok(response) = request.send().await { - if response.status().is_success() { - if let Ok(version) = response.json::().await { - server_limit = version.logs3_payload_max_bytes.filter(|v| *v > 0); - } - } - } - } - - let can_use_overflow = server_limit.is_some(); - let mut max_request_size = DEFAULT_MAX_REQUEST_SIZE; - if let Some(override_bytes) = self.max_request_size_override { - max_request_size = if let Some(server_limit) = server_limit { - override_bytes.min(server_limit) - } else { - override_bytes - }; - } else if let Some(server_limit) = server_limit { - max_request_size = server_limit; - } - PayloadLimits { - max_request_size: max_request_size.max(1), - can_use_overflow, - } - } - - async fn submit_row_batch( - &self, - items: &[RowPayload], - limits: &PayloadLimits, - ) -> Result { - if items.is_empty() { - return Ok(Logs3UploadResult::default()); - } - let mut pending: Vec> = vec![items.to_vec()]; - let mut result = Logs3UploadResult::default(); - - while let Some(chunk) = pending.pop() { - if chunk.is_empty() { - continue; - } - let payload = construct_logs3_payload(&chunk); - let payload_bytes = payload.len(); - - if limits.can_use_overflow && payload_bytes > limits.max_request_size { - let overflow_result = self.submit_overflow(&chunk, payload, payload_bytes).await?; - result.rows_uploaded += overflow_result.rows_uploaded; - result.bytes_processed += overflow_result.bytes_processed; - result.requests_sent += overflow_result.requests_sent; - continue; - } - - let (post_result, post_attempts) = - run_with_backoff(|| self.post_logs3_raw(&payload)).await; - result.requests_sent += post_attempts; - match post_result { - Ok(()) => { - result.rows_uploaded += chunk.len(); - result.bytes_processed += payload_bytes; - } - Err(BraintrustError::Api { status, .. }) if status == 413 => { - if limits.can_use_overflow { - let overflow_result = self - .submit_overflow(&chunk, payload.clone(), payload_bytes) - .await?; - result.rows_uploaded += overflow_result.rows_uploaded; - result.bytes_processed += overflow_result.bytes_processed; - result.requests_sent += overflow_result.requests_sent; - } else if chunk.len() > 1 { - let mid = chunk.len() / 2; - pending.push(chunk[mid..].to_vec()); - pending.push(chunk[..mid].to_vec()); - } else { - return Err(BraintrustError::Api { - status, - message: - "single row exceeds server payload limit and overflow is unavailable" - .to_string(), - }); - } - } - Err(err) => return Err(err), - } - } - - Ok(result) - } - - async fn submit_overflow( - &self, - items: &[RowPayload], - payload: String, - payload_bytes: usize, - ) -> Result { - let overflow_rows: Vec = items - .iter() - .map(|item| item.overflow_meta.clone()) - .collect(); - let mut requests_sent = 0usize; - let (upload_result, upload_attempts) = - run_with_backoff(|| self.request_overflow_upload(&overflow_rows, payload_bytes)).await; - requests_sent += upload_attempts; - let upload = upload_result?; - - let (upload_payload_result, upload_payload_attempts) = - run_with_backoff(|| self.upload_overflow_payload(&upload, &payload)).await; - requests_sent += upload_payload_attempts; - upload_payload_result?; - - let overflow_reference = json!({ - "rows": { - "type": LOGS3_OVERFLOW_REFERENCE_TYPE, - "key": upload.key, - }, - "api_version": LOGS_API_VERSION, - }) - .to_string(); - let (post_result, post_attempts) = - run_with_backoff(|| self.post_logs3_raw(&overflow_reference)).await; - requests_sent += post_attempts; - post_result?; - - Ok(Logs3UploadResult { - rows_uploaded: items.len(), - bytes_processed: payload_bytes, - requests_sent, - }) - } - - async fn request_overflow_upload( - &self, - rows: &[Logs3OverflowInputRow], - payload_bytes: usize, - ) -> Result { - let url = self - .api_url - .join("logs3/overflow") - .map_err(|e| BraintrustError::InvalidConfig(format!("invalid overflow URL: {e}")))?; - - let request_body = json!({ - "content_type": "application/json", - "size_bytes": payload_bytes, - "rows": rows, - }); - let mut request = self - .client - .post(url) - .bearer_auth(&self.api_key) - .header("content-type", "application/json") - .json(&request_body); - if let Some(org_name) = &self.org_name { - request = request.header("x-bt-org-name", org_name); - } - let response = request.send().await?; - if !response.status().is_success() { - let status = response.status().as_u16(); - let message = response.text().await.unwrap_or_default(); - return Err(BraintrustError::Api { status, message }); - } - response - .json::() - .await - .map_err(|e| BraintrustError::Api { - status: 200, - message: format!("failed to parse logs3 overflow response: {e}"), - }) - } - - async fn upload_overflow_payload( - &self, - upload: &Logs3OverflowUpload, - payload: &str, - ) -> Result<()> { - let signed_url = Url::parse(&upload.signed_url).map_err(|e| { - BraintrustError::InvalidConfig(format!("invalid logs3 overflow signed URL: {e}")) - })?; - match upload.method { - OverflowMethod::Post => { - let fields = upload.fields.clone().ok_or_else(|| { - BraintrustError::InvalidConfig( - "logs3 overflow POST upload missing form fields".to_string(), - ) - })?; - let content_type = fields - .get("Content-Type") - .cloned() - .unwrap_or_else(|| "application/json".to_string()); - let mut form = multipart::Form::new(); - for (key, value) in &fields { - form = form.text(key.clone(), value.clone()); - } - let file_part = multipart::Part::text(payload.to_string()) - .mime_str(&content_type) - .map_err(|e| { - BraintrustError::InvalidConfig(format!( - "invalid overflow content-type '{content_type}': {e}" - )) - })?; - form = form.part("file", file_part); - - let mut request = self.client.post(signed_url).multipart(form); - let mut headers = header_map_from_pairs(upload.headers.as_ref())?; - headers.remove("content-type"); - if !headers.is_empty() { - request = request.headers(headers); - } - let response = request.send().await?; - if !response.status().is_success() { - let status = response.status().as_u16(); - let message = response.text().await.unwrap_or_default(); - return Err(BraintrustError::Api { status, message }); - } - } - OverflowMethod::Put => { - let mut headers = header_map_from_pairs(upload.headers.as_ref())?; - if upload.signed_url.contains("blob.core.windows.net") - && !headers.contains_key("x-ms-blob-type") - { - headers.insert("x-ms-blob-type", HeaderValue::from_static("BlockBlob")); - } - - let mut request = self.client.put(signed_url).body(payload.to_string()); - if !headers.is_empty() { - request = request.headers(headers); - } - let response = request.send().await?; - if !response.status().is_success() { - let status = response.status().as_u16(); - let message = response.text().await.unwrap_or_default(); - return Err(BraintrustError::Api { status, message }); - } - } - } - Ok(()) - } - - async fn post_logs3_raw(&self, payload: &str) -> Result<()> { - let url = self - .api_url - .join("logs3") - .map_err(|e| BraintrustError::InvalidConfig(format!("invalid logs3 URL: {e}")))?; - let mut request = self - .client - .post(url) - .bearer_auth(&self.api_key) - .header("content-type", "application/json") - .body(payload.to_string()); - if let Some(org_name) = &self.org_name { - request = request.header("x-bt-org-name", org_name); - } - let response = request.send().await?; - if !response.status().is_success() { - let status = response.status().as_u16(); - let message = response.text().await.unwrap_or_default(); - return Err(BraintrustError::Api { status, message }); - } - Ok(()) - } -} - -impl RowPayload { - fn from_row(row: &Map) -> Result { - let row_json = serde_json::to_string(row).map_err(|e| { - BraintrustError::InvalidConfig(format!("failed to serialize logs3 row: {e}")) - })?; - let row_bytes = row_json.len(); - let mut object_ids = Map::new(); - for key in OBJECT_ID_KEYS { - if let Some(value) = row.get(key) { - object_ids.insert(key.to_string(), value.clone()); - } - } - let is_delete = row.get("_object_delete").and_then(Value::as_bool); - Ok(Self { - row_json, - row_bytes, - overflow_meta: Logs3OverflowInputRow { - object_ids, - is_delete, - input_row: OverflowInputRowInfo { - byte_size: row_bytes, - }, - }, - }) - } -} - -fn construct_logs3_payload(items: &[RowPayload]) -> String { - let rows = items - .iter() - .map(|item| item.row_json.as_str()) - .collect::>() - .join(","); - format!("{{\"rows\":[{rows}],\"api_version\":{LOGS_API_VERSION}}}") -} - -fn header_map_from_pairs(headers: Option<&HashMap>) -> Result { - let mut out = HeaderMap::new(); - if let Some(headers) = headers { - for (key, value) in headers { - let name = HeaderName::from_bytes(key.as_bytes()).map_err(|e| { - BraintrustError::InvalidConfig(format!("invalid HTTP header name '{key}': {e}")) - })?; - let header_value = HeaderValue::from_str(value).map_err(|e| { - BraintrustError::InvalidConfig(format!( - "invalid HTTP header value for '{key}': {e}" - )) - })?; - out.insert(name, header_value); - } - } - Ok(out) -} - -fn should_retry(err: &BraintrustError) -> bool { - match err { - BraintrustError::Http(http_err) => http_err.is_timeout() || http_err.is_connect(), - BraintrustError::Network(_) - | BraintrustError::ChannelClosed - | BraintrustError::Background(_) => true, - BraintrustError::Api { status, .. } => *status == 429 || (*status >= 500 && *status <= 599), - BraintrustError::InvalidConfig(_) | BraintrustError::StreamAggregation(_) => false, - } -} - -async fn run_with_backoff(mut operation: F) -> (Result, usize) -where - F: FnMut() -> Fut, - Fut: std::future::Future>, -{ - let mut attempts = 0usize; - let mut backoff = ExponentialBackoffBuilder::new() - .with_initial_interval(Duration::from_millis(BASE_BACKOFF_MS)) - .with_multiplier(2.0) - .with_randomization_factor(0.2) - .with_max_interval(Duration::from_millis(MAX_BACKOFF_MS)) - .with_max_elapsed_time(None) - .build(); - - loop { - attempts += 1; - match operation().await { - Ok(value) => return (Ok(value), attempts), - Err(err) => { - if !should_retry(&err) || attempts >= MAX_ATTEMPTS { - return (Err(err), attempts); - } - match backoff.next_backoff() { - Some(delay) => sleep(delay).await, - None => return (Err(err), attempts), - } - } - } - } } #[cfg(test)] diff --git a/src/span.rs b/src/span.rs index 30a88dc..7d8e7e4 100644 --- a/src/span.rs +++ b/src/span.rs @@ -5,11 +5,15 @@ 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; use crate::error::Result; +use crate::internal::ports::SpanSubmitter; +use crate::internal::session::CredentialScope; +use crate::internal::write_outcome::WriteOutcome; use crate::span_components::SpanComponents; use crate::types::{ ParentSpanInfo, SpanAttributes, SpanEventData, SpanObjectType, SpanPayload, SpanType, @@ -247,31 +251,14 @@ impl SpanLogBuilder { } } -#[async_trait] -pub(crate) trait SpanSubmitter: Send + Sync { - /// Submit a span payload for queuing (fire-and-forget). - /// The payload is queued internally and processed by a background worker. - /// Dropping is handled internally if the queue is full. - fn submit(&self, token: String, payload: SpanPayload, parent_info: Option); - - /// Flush all queued span data through to the server. - async fn flush(&self) -> Result<()>; - - /// Trigger a non-blocking background flush. - async fn trigger_flush(&self) -> Result<()>; -} - -#[derive(Debug)] -#[allow(private_bounds)] -pub struct SpanBuilder { - submitter: Arc, - token: String, - org_id: String, - org_name: Option, +pub struct SpanBuilder { + submitter: Arc, + scope: CredentialScope, project_name: Option, parent_info: Option, span_type: SpanType, purpose: Option, + name: Option, start_time_override: Option, environment: Option, span_id_override: Option, @@ -280,17 +267,26 @@ pub struct SpanBuilder { span_attributes_extra: HashMap, } -impl Clone for SpanBuilder { +impl fmt::Debug for SpanBuilder { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("SpanBuilder") + .field("scope", &self.scope) + .field("project_name", &self.project_name) + .field("parent_info", &self.parent_info) + .finish_non_exhaustive() + } +} + +impl Clone for SpanBuilder { fn clone(&self) -> Self { Self { submitter: Arc::clone(&self.submitter), - token: self.token.clone(), - org_id: self.org_id.clone(), - org_name: self.org_name.clone(), + scope: self.scope.clone(), project_name: self.project_name.clone(), parent_info: self.parent_info.clone(), span_type: self.span_type, purpose: self.purpose.clone(), + name: self.name.clone(), start_time_override: self.start_time_override, environment: self.environment.clone(), span_id_override: self.span_id_override.clone(), @@ -369,22 +365,16 @@ impl SpanOrigin { } } -#[allow(private_bounds)] -impl SpanBuilder { - pub(crate) fn new( - submitter: Arc, - token: impl Into, - org_id: impl Into, - ) -> Self { +impl SpanBuilder { + pub(crate) fn new(submitter: Arc, scope: CredentialScope) -> Self { Self { submitter, - token: token.into(), - org_id: org_id.into(), - org_name: None, + scope, project_name: None, parent_info: None, span_type: SpanType::default(), purpose: None, + name: None, start_time_override: None, environment: detected_environment(), span_id_override: None, @@ -394,6 +384,12 @@ impl SpanBuilder { } } + /// Set the span name before its first snapshot is submitted. + pub fn name(mut self, name: impl Into) -> Self { + self.name = Some(name.into()); + self + } + pub fn start_time(mut self, start_time: f64) -> Self { self.start_time_override = Some(start_time); self @@ -447,7 +443,7 @@ impl SpanBuilder { } pub fn org_name(mut self, org_name: impl Into) -> Self { - self.org_name = Some(org_name.into()); + self.scope = self.scope.with_org_name(org_name.into()); self } @@ -485,7 +481,7 @@ impl SpanBuilder { self } - pub fn build(self) -> SpanHandle { + pub fn build(self) -> SpanHandle { use std::time::{SystemTime, UNIX_EPOCH}; // Generate both IDs ONCE at span creation - reused for all flushes. @@ -518,19 +514,20 @@ impl SpanBuilder { SpanHandle { submitter: Arc::clone(&self.submitter), - token: self.token, + scope: self.scope.clone(), parent_info: self.parent_info, row_id: row_id.clone(), inner: Arc::new(Mutex::new(SpanData { row_id, span_id, has_flushed: false, - org_id: self.org_id, - org_name: self.org_name, + org_id: self.scope.org_id().to_owned(), + org_name: self.scope.org_name().map(str::to_owned), project_name: self.project_name, start_time, span_type: self.span_type, purpose: self.purpose, + name: self.name, environment: self.environment, span_origin: self.span_origin, span_attributes_extra: self.span_attributes_extra, @@ -541,21 +538,20 @@ impl SpanBuilder { } } -#[allow(private_bounds)] -pub struct SpanHandle { - submitter: Arc, - token: String, +pub struct SpanHandle { + submitter: Arc, + scope: CredentialScope, parent_info: Option, /// Stable identifier set at creation, accessible without locking. row_id: String, inner: Arc>, } -impl Clone for SpanHandle { +impl Clone for SpanHandle { fn clone(&self) -> Self { Self { submitter: Arc::clone(&self.submitter), - token: self.token.clone(), + scope: self.scope.clone(), parent_info: self.parent_info.clone(), row_id: self.row_id.clone(), inner: Arc::clone(&self.inner), @@ -563,8 +559,7 @@ impl Clone for SpanHandle { } } -#[allow(private_bounds)] -impl SpanHandle { +impl SpanHandle { /// Return the stable row ID for this span, usable with `log_feedback()` and `update_span()`. pub fn row_id(&self) -> &str { &self.row_id @@ -572,19 +567,33 @@ impl SpanHandle { /// Log event data to this span. All fields are optional. /// Multiple calls will merge data (later values overwrite earlier ones). - /// Each call synchronously queues the current span snapshot for background processing. - pub fn log(&self, event: SpanLog) { - let payload = { - let mut inner = self - .inner - .lock() - .expect("span mutex should not be poisoned"); - apply_span_log_to_data(&mut inner, event); - current_span_payload(&mut inner) - }; - - self.submitter - .submit(self.token.clone(), payload, self.parent_info.clone()); + /// Each call attempts admission immediately. Await the returned token for full + /// delivery; ignoring it leaves accepted work running automatically. After + /// end is called, later mutations are ignored with a diagnostic and return + /// an immediately successful completion without submitting a write. + pub fn log( + &self, + event: SpanLog, + ) -> crate::Completion { + let mut inner = self + .inner + .lock() + .expect("span mutex should not be poisoned"); + if inner.end_time.is_some() { + tracing::warn!(row_id = %self.row_id, "ignoring mutation of a finished span"); + let (sender, completion) = crate::Completion::pair(); + let _ = sender.send(Ok(crate::DeliveryReceipt)); + return completion; + } + apply_span_log_to_data(&mut inner, event); + let payload = current_span_payload(&mut inner); + let completion = + self.submitter + .submit(self.scope.clone(), payload, self.parent_info.clone(), None); + if completion.is_ok() { + inner.has_flushed = true; + } + completion } /// Flush all queued span data through to Braintrust. @@ -592,38 +601,54 @@ impl SpanHandle { /// This does not create a new queued span update. Call [`SpanHandle::log`] or /// [`SpanHandle::end`] first to submit changes, then call `flush()` when you need /// to ensure the background queue has been processed. - pub async fn flush(&self) -> Result<()> { - self.submitter.flush().await + pub fn flush(&self) -> crate::Completion { + self.submitter.flush() } - /// Mark the span as ended with the current timestamp. - /// - /// Calling `end()` multiple times is idempotent: once an end time is set, - /// subsequent calls return the same value without overwriting it. - pub fn end(&self) -> f64 { + /// Finish this span and observe delivery of its final snapshot. + /// Repeated calls observe the same admitted write, without resubmission. + /// Pending observers are bounded; admission can be full while the existing + /// final write continues. Completed outcomes remain observable on this handle. + pub fn end(&self) -> crate::Completion { self.end_with_time(epoch_secs()) } - /// Mark the span as ended with an explicit timestamp (seconds since Unix epoch). - /// - /// Calling this multiple times is idempotent: once an end time is set, - /// subsequent calls return the previously-set value. - pub fn end_with_time(&self, end_time: f64) -> f64 { - let (end_time, payload) = { - let mut inner = self - .inner - .lock() - .expect("span mutex should not be poisoned"); - if let Some(existing) = inner.end_time { - return existing; - } - inner.end_time = Some(end_time); - (end_time, current_span_payload(&mut inner)) - }; + /// Finish with a timestamp chosen once (seconds since Unix epoch). + /// If admission is rejected, a later end call retries the same final snapshot + /// and timestamp. An accepted write's terminal failure is never replayed. + pub fn end_with_time( + &self, + end_time: f64, + ) -> crate::Completion { + let mut inner = self + .inner + .lock() + .expect("span mutex should not be poisoned"); + inner.end_time.get_or_insert(end_time); + if let Some(outcome) = &inner.end_outcome { + return outcome.observe(); + } + let outcome = WriteOutcome::default(); + let payload = current_span_payload(&mut inner); + let completion = self.submitter.submit( + self.scope.clone(), + payload, + self.parent_info.clone(), + Some(outcome.clone()), + ); + if completion.is_ok() { + inner.has_flushed = true; + inner.end_outcome = Some(outcome); + } + completion + } - self.submitter - .submit(self.token.clone(), payload, self.parent_info.clone()); - end_time + /// The chosen end timestamp, including when final admission needs retrying. + pub fn end_time(&self) -> Option { + self.inner + .lock() + .expect("span mutex should not be poisoned") + .end_time } /// Trigger a non-blocking background flush. @@ -758,9 +783,7 @@ fn current_span_payload(inner: &mut SpanData) -> SpanPayload { inner.metrics.insert("end".to_string(), end); } - let payload: SpanPayload = inner.clone().into(); - inner.has_flushed = true; - payload + inner.clone().into() } #[derive(Clone, Default)] @@ -768,6 +791,7 @@ struct SpanData { row_id: String, span_id: String, has_flushed: bool, + end_outcome: Option, org_id: String, org_name: Option, project_name: Option, @@ -1060,6 +1084,137 @@ mod tests { use crate::ProjectLogsIdentifier; use serde_json::json; + #[test] + fn rejected_admission_preserves_first_write_and_retryable_final_timestamp() { + use std::sync::atomic::{AtomicBool, Ordering}; + struct Submitter { + reject: AtomicBool, + accepted: Mutex>, + } + #[async_trait] + impl SpanSubmitter for Submitter { + fn submit( + &self, + _: CredentialScope, + payload: SpanPayload, + _: Option, + outcome: Option, + ) -> crate::Completion { + let (sender, mut completion) = crate::Completion::pair(); + let sender = crate::internal::write_outcome::WriteResponder::new(sender, outcome); + if self.reject.load(Ordering::Relaxed) { + completion.record_issuance(Err(crate::EnqueueError::Full.into())); + } else { + self.accepted.lock().unwrap().push(payload); + let _ = sender.send(Ok(crate::DeliveryReceipt)); + } + 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), + accepted: Mutex::new(Vec::new()), + }); + let span = + SpanBuilder::new(submitter.clone(), CredentialScope::test("token", "org")).build(); + assert!(span.log(SpanLog::default()).is_err()); + submitter.reject.store(false, Ordering::Relaxed); + assert!(span.log(SpanLog::default()).is_ok()); + assert!( + !submitter.accepted.lock().unwrap()[0].is_merge, + "rejected writes do not count as a first admission" + ); + submitter.reject.store(true, Ordering::Relaxed); + assert!(span.end_with_time(10.0).is_err()); + assert_eq!(span.end_time(), Some(10.0)); + assert_eq!(submitter.accepted.lock().unwrap().len(), 1); + assert!(span + .log(SpanLog::builder().input(json!("too late")).build().unwrap()) + .is_ok()); + submitter.reject.store(false, Ordering::Relaxed); + assert!(span.end_with_time(20.0).is_ok()); + assert!(span.end_with_time(30.0).is_ok()); + assert_eq!(span.end_time(), Some(10.0)); + let accepted = submitter.accepted.lock().unwrap(); + assert_eq!(accepted.len(), 2); + assert_eq!(accepted[1].metrics.as_ref().unwrap()["end"], 10.0); + assert!(accepted[1].input.is_none()); + } + + #[tokio::test] + async fn concurrent_ends_share_one_snapshot_and_one_delivery_result() { + use crate::internal::write_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 submit( + &self, + _: CredentialScope, + payload: SpanPayload, + _: Option, + outcome: Option, + ) -> crate::Completion { + let (sender, completion) = crate::Completion::pair(); + self.0 + .lock() + .unwrap() + .push((payload, WriteResponder::new(sender, outcome))); + completion + } + } + let submitter = Arc::new(Submitter(Mutex::new(Vec::new()))); + let span = + SpanBuilder::new(submitter.clone(), CredentialScope::test("token", "org")).build(); + drop(span.end_with_time(123.0)); // Dropping the first observer does not cancel the write. + let observers = std::thread::scope(|threads| { + (0..32) + .map(|_| { + let span = span.clone(); + threads.spawn(move || span.end_with_time(456.0)) + }) + .collect::>() + .into_iter() + .map(|thread| thread.join().unwrap()) + .collect::>() + }); + assert!(observers.iter().all(|observer| observer.is_ok())); + let (payload, response) = { + let mut submitted = submitter.0.lock().unwrap(); + assert_eq!(submitted.len(), 1); + submitted.pop().unwrap() + }; + assert_eq!(payload.metrics.unwrap()["end"], 123.0); + let _ = response.send(Ok(crate::DeliveryReceipt)); + for observer in observers { + observer.await.unwrap(); + } + span.end_with_time(999.0).await.unwrap(); + assert_eq!(span.end_time(), Some(123.0)); + assert!(submitter.0.lock().unwrap().is_empty()); + } + #[tokio::test] async fn span_logs_input_and_output() { let (span, collector) = build_test_span(); @@ -1256,8 +1411,9 @@ mod tests { let (span, _collector) = build_test_span(); let first = span.end_with_time(123.0); let second = span.end_with_time(456.0); - assert_eq!(first, 123.0); - assert_eq!(second, 123.0); + first.await.unwrap(); + second.await.unwrap(); + assert_eq!(span.end_time(), Some(123.0)); } #[tokio::test] diff --git a/src/stream.rs b/src/stream.rs index 248050e..ea72c71 100644 --- a/src/stream.rs +++ b/src/stream.rs @@ -66,7 +66,7 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; use tokio::sync::Mutex; -use crate::span::{SpanHandle, SpanLog, SpanSubmitter}; +use crate::span::{SpanHandle, SpanLog}; use crate::types::{usage_metrics_to_map, UsageMetrics}; // ============================================================================= @@ -1021,16 +1021,13 @@ impl BraintrustStream { /// # Type Parameters /// - `S`: The stream type yielding `Result` /// - `E`: The error type (allows use with any error type) -/// - `Sub`: The span submitter type -#[allow(private_bounds)] -pub fn wrap_stream_with_span( +pub fn wrap_stream_with_span( stream: S, - span: SpanHandle, + span: SpanHandle, ) -> Pin> + Send>> where S: Stream> + Send + Unpin + 'static, E: Send + 'static, - Sub: SpanSubmitter + 'static, { use futures::StreamExt; @@ -1150,18 +1147,15 @@ enum FinalizeState { } /// A wrapper stream that logs aggregated output when the stream is exhausted. -struct SpanCompleteWrapper { +struct SpanCompleteWrapper { inner: S, - span: Option>, + span: Option, aggregator: Option>>, finalize_state: FinalizeState, } /// Finalize the stream by logging and flushing the span. -async fn finalize_span( - span: SpanHandle, - aggregator: Arc>, -) { +async fn finalize_span(span: SpanHandle, aggregator: Arc>) { let mut agg = aggregator.lock().await; if !agg.is_empty() { match agg.final_value() { @@ -1197,18 +1191,21 @@ async fn finalize_span( } } } - span.end(); + let end = span.end(); + let flush = span.flush(); + if let Err(error) = end.await { + tracing::warn!("Failed to deliver final span: {}", error); + } // Flush span with aggregated output - if let Err(e) = span.flush().await { + if let Err(e) = flush.await { tracing::warn!("Failed to flush span: {}", e); } } -impl Stream for SpanCompleteWrapper +impl Stream for SpanCompleteWrapper where S: Stream> + Unpin, E: Send + 'static, - Sub: SpanSubmitter + 'static, { type Item = std::result::Result; diff --git a/src/test_utils.rs b/src/test_utils.rs index 0075a6b..dc527db 100644 --- a/src/test_utils.rs +++ b/src/test_utils.rs @@ -3,10 +3,12 @@ use std::sync::{Arc, Mutex}; use async_trait::async_trait; use crate::error::Result; +use crate::internal::session::CredentialScope; use crate::span::{SpanBuilder, SpanHandle}; use crate::types::{ParentSpanInfo, SpanPayload}; -use crate::span::SpanSubmitter; +use crate::internal::ports::{DeliveryControl, SpanSubmitter}; +use crate::internal::write_outcome::{WriteOutcome, WriteResponder}; #[derive(Clone, Default)] pub(crate) struct TestSpanCollector { @@ -58,16 +60,30 @@ impl MockBraintrustClient { #[async_trait] impl SpanSubmitter for MockBraintrustClient { - fn submit(&self, token: String, payload: SpanPayload, parent_info: Option) { + fn submit( + &self, + scope: CredentialScope, + payload: SpanPayload, + parent_info: Option, + outcome: Option, + ) -> crate::Completion { self.collector.push(CapturedSpan { - token, + token: scope.token().to_owned(), payload, parent: parent_info, }); + let (sender, completion) = crate::Completion::pair(); + let _ = WriteResponder::new(sender, outcome).send(Ok(crate::DeliveryReceipt)); + completion } +} - async fn flush(&self) -> Result<()> { - Ok(()) +#[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<()> { @@ -76,13 +92,13 @@ impl SpanSubmitter for MockBraintrustClient { } } -pub(crate) fn mock_span_builder() -> (SpanBuilder, TestSpanCollector) { +pub(crate) fn mock_span_builder() -> (SpanBuilder, TestSpanCollector) { let (submitter, collector) = MockBraintrustClient::new(); - let builder = SpanBuilder::new(submitter, "test-token", "org-id"); + let builder = SpanBuilder::new(submitter, CredentialScope::test("test-token", "org-id")); (builder, collector) } -pub(crate) fn build_test_span() -> (SpanHandle, TestSpanCollector) { +pub(crate) fn build_test_span() -> (SpanHandle, TestSpanCollector) { let (builder, collector) = mock_span_builder(); (builder.build(), collector) } diff --git a/src/types.rs b/src/types.rs index 5246a02..d7d3b23 100644 --- a/src/types.rs +++ b/src/types.rs @@ -340,45 +340,7 @@ impl LogDestination { } } -/// Response from POST logs3/overflow — provides a signed URL to upload the payload. -#[derive(Debug, Clone, Deserialize)] -#[serde(rename_all = "camelCase")] -pub(crate) struct Logs3OverflowUpload { - pub method: String, - pub signed_url: String, - pub headers: Option>, - pub fields: Option>, - pub key: String, -} - -/// A single row's overflow metadata, sent when requesting the overflow upload URL. -#[derive(Debug, Clone, Serialize)] -pub(crate) struct Logs3OverflowInputRow { - /// Key identifying fields extracted from the row (experiment_id, dataset_id, etc.) - pub object_ids: serde_json::Map, - #[serde(skip_serializing_if = "Option::is_none")] - pub is_delete: Option, - pub input_row: Logs3OverflowInputRowMeta, -} - -#[derive(Debug, Clone, Serialize)] -pub(crate) struct Logs3OverflowInputRowMeta { - pub byte_size: usize, -} - -/// Reference sent to POST logs3 after a successful overflow upload. -#[derive(Debug, Clone, Serialize)] -pub(crate) struct Logs3OverflowRequest { - pub rows: Logs3OverflowReference, - pub api_version: u8, -} - -#[derive(Debug, Clone, Serialize)] -pub(crate) struct Logs3OverflowReference { - #[serde(rename = "type")] - pub reference_type: &'static str, - pub key: String, -} +pub(crate) use crate::api::logs3::{Logs3OverflowInputRow, Logs3OverflowInputRowMeta}; /// The object ID keys that identify a row's destination, used for overflow metadata. /// Matches the TypeScript SDK's OBJECT_ID_KEYS. diff --git a/tests/api_client.rs b/tests/api_client.rs new file mode 100644 index 0000000..450d79a --- /dev/null +++ b/tests/api_client.rs @@ -0,0 +1,446 @@ +use std::{error::Error as _, time::Duration}; + +use braintrust_sdk_rust::api::{ + self, + btql::{BTQLQuery, BTQLQueryInner}, + datasets::DatasetRegisterRequest, + experiments::{BaseExperimentRequest, ExperimentRegisterRequest}, + logs3::{InsertRows, Logs3OverflowUploadRequest}, + projects::ProjectRegisterRequest, +}; +use serde_json::json; +use wiremock::{ + matchers::{body_json, header, method, path}, + Mock, MockServer, ResponseTemplate, +}; + +fn config(api: &MockServer, app: &MockServer, token: &str) -> api::Config { + let mut config = api::Config::new(token); + config.api_url = api.uri(); + config.app_url = app.uri(); + config.org_name = Some("isolated-org".into()); + config.timeout = Duration::from_secs(3); + config +} + +#[tokio::test] +async fn explicit_scope_works_without_login_and_login_does_not_mutate_it() { + let api = MockServer::start().await; + let app = MockServer::start().await; + let mut settings = config(&api, &app, "explicit-token"); + settings.user_agent = Some("manual-ingest/1".into()); + let client = api::Client::new(settings).unwrap(); + assert!(api.received_requests().await.unwrap().is_empty()); + assert!(app.received_requests().await.unwrap().is_empty()); + + Mock::given(method("POST")) + .and(path("/api/apikey/login")) + .and(header("authorization", "Bearer explicit-token")) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"org_info":[{ + "id":"other-org", "name":"other", "api_url":"http://must-not-be-used.invalid" + }]})), + ) + .expect(1) + .mount(&app) + .await; + assert_eq!(client.login().await.unwrap().org_info[0].id(), "other-org"); + Mock::given(method("POST")) + .and(path("/logs3")) + .and(header("authorization", "Bearer explicit-token")) + .and(header("x-bt-org-name", "isolated-org")) + .and(header("user-agent", "manual-ingest/1")) + .and(body_json( + json!({"rows":[{"id":"row-1","project_id":"project"}],"api_version":2}), + )) + .respond_with(ResponseTemplate::new(204)) + .expect(1) + .mount(&api) + .await; + client + .clone() + .logs3() + .insert(InsertRows::new(vec![ + json!({"id":"row-1","project_id":"project"}), + ])) + .await + .unwrap(); + assert_eq!(app.received_requests().await.unwrap().len(), 1); + assert_eq!(api.received_requests().await.unwrap().len(), 1); +} + +#[tokio::test] +async fn failures_make_one_attempt_without_retry_or_overflow() { + for status in [400, 401, 413, 429, 500, 503] { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(status).set_body_string("rejected")) + .expect(1) + .mount(&server) + .await; + let client = api::Client::new(config(&server, &server, "token")).unwrap(); + let error = client + .logs3() + .post_raw(br#"{"rows":[],"api_version":2}"#.to_vec()) + .await + .unwrap_err(); + match error { + api::Error::Http { + status: actual, + operation, + body, + } => { + assert_eq!(actual, status); + assert_eq!(operation, "logs3"); + assert_eq!(body, "rejected"); + } + other => panic!("unexpected {other:?}"), + } + assert_eq!(server.received_requests().await.unwrap().len(), 1); + } +} + +#[tokio::test] +async fn redirects_are_returned_without_a_followup_request() { + let source = MockServer::start().await; + let target = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with( + ResponseTemplate::new(307) + .insert_header("location", format!("{}/logs3", target.uri()).as_str()), + ) + .expect(1) + .mount(&source) + .await; + let client = api::Client::new(config(&source, &source, "token")).unwrap(); + assert!(matches!( + client.logs3().post_raw(vec![]).await, + Err(api::Error::Http { status: 307, .. }) + )); + assert!(target.received_requests().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn btql_readback_preserves_cursor_and_request_errors() { + let data = MockServer::start().await; + let control = MockServer::start().await; + let client = api::Client::new(config(&data, &control, "token")).unwrap(); + Mock::given(method("POST")) + .and(path("/btql")) + .and(body_json( + json!({"query":{"from":"dataset('id')","limit":5,"cursor":"next"}}), + )) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(json!({"data":[{"id":"record","input":1}],"cursor":"last"})), + ) + .expect(1) + .mount(&data) + .await; + let page = client + .btql() + .query_dataset_records(BTQLQuery { + query: BTQLQueryInner { + from: "dataset('id')".into(), + limit: Some(5), + cursor: Some("next".into()), + }, + }) + .await + .unwrap(); + assert_eq!(page.cursor.as_deref(), Some("last")); + assert_eq!(page.data.len(), 1); + assert!(control.received_requests().await.unwrap().is_empty()); + + data.reset().await; + Mock::given(method("POST")) + .and(path("/btql")) + .respond_with(ResponseTemplate::new(500)) + .expect(1) + .mount(&data) + .await; + assert!(matches!( + client + .btql() + .query_dataset_records(BTQLQuery { + query: BTQLQueryInner { + from: "dataset('id')".into(), + limit: None, + cursor: None, + } + }) + .await, + Err(api::Error::Http { status: 500, .. }) + )); +} + +#[tokio::test] +async fn malformed_success_is_a_decode_error_with_status_and_source() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/version")) + .respond_with(ResponseTemplate::new(202).set_body_string("not json")) + .expect(1) + .mount(&server) + .await; + let client = api::Client::new(config(&server, &server, "token")).unwrap(); + let error = client.logs3().version().await.unwrap_err(); + assert!(error.source().is_some()); + assert!(matches!( + error, + api::Error::Decode { + status: 202, + operation: "version", + .. + } + )); +} + +#[tokio::test] +async fn selected_resource_endpoints_keep_their_wire_contracts_and_planes() { + let data = MockServer::start().await; + let control = MockServer::start().await; + let client = api::Client::new(config(&data, &control, "token")).unwrap(); + for (route, response) in [ + ("/api/project/register", json!({"project":{"id":"project"}})), + ( + "/api/dataset/register", + json!({"project":{"id":"project","name":"p"},"dataset":{"id":"dataset","name":"d"}}), + ), + ( + "/api/experiment/register", + json!({"project":{"id":"project","name":"p"},"experiment":{"id":"experiment","name":"e"}}), + ), + ( + "/api/base_experiment/get_id", + json!({"base_exp_id":"base","base_exp_name":"baseline"}), + ), + ] { + Mock::given(method("POST")) + .and(path(route)) + .and(header("authorization", "Bearer token")) + .respond_with(ResponseTemplate::new(200).set_body_json(response)) + .expect(1) + .mount(&control) + .await; + } + assert_eq!( + client + .projects() + .register(ProjectRegisterRequest { + project_name: "p".into(), + org_id: Some("org".into()), + org_name: None + }) + .await + .unwrap() + .project + .id, + "project" + ); + assert_eq!( + client + .datasets() + .register(DatasetRegisterRequest { + project_name: "p".into(), + org_name: "org".into(), + dataset_name: "d".into(), + description: None, + metadata: None + }) + .await + .unwrap() + .dataset + .id, + "dataset" + ); + assert_eq!( + client + .experiments() + .register(ExperimentRegisterRequest { + project_name: "p".into(), + org_id: "org".into(), + experiment_name: Some("e".into()), + description: None, + base_experiment: None, + metadata: None, + public: None, + update: None, + repo_info: None, + }) + .await + .unwrap() + .experiment + .id, + "experiment" + ); + assert_eq!( + client + .experiments() + .base(BaseExperimentRequest { + id: "experiment".into() + }) + .await + .unwrap() + .base_exp_id + .as_deref(), + Some("base") + ); + Mock::given(method("GET")) + .and(path("/experiment-comparison2")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"scores":{},"metrics":{}}))) + .expect(1) + .mount(&control) + .await; + client + .experiments() + .compare("experiment & name", Some("base/id")) + .await + .unwrap(); + let requests = control.received_requests().await.unwrap(); + let comparison = requests + .iter() + .find(|r| r.url.path() == "/experiment-comparison2") + .unwrap(); + assert_eq!( + comparison + .url + .query_pairs() + .find(|(k, _)| k == "experiment_id") + .unwrap() + .1, + "experiment & name" + ); + Mock::given(method("GET")).and(path("/dataset-summary")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"project_name":"p","project_id":"project","dataset_name":"d","dataset_id":"dataset"}))) + .expect(1).mount(&data).await; + assert_eq!( + client + .datasets() + .summarize("dataset") + .await + .unwrap() + .dataset_id, + "dataset" + ); + assert_eq!(data.received_requests().await.unwrap().len(), 1); +} + +#[tokio::test] +async fn overflow_grant_is_explicit_and_does_not_upload_or_commit() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3/overflow")) + .and(body_json( + json!({"content_type":"application/json","size_bytes":123,"rows":[]}), + )) + .respond_with(ResponseTemplate::new(200).set_body_json( + json!({"method":"PUT","signedUrl":"https://signed.invalid/upload","key":"key"}), + )) + .expect(1) + .mount(&server) + .await; + let client = api::Client::new(config(&server, &server, "token")).unwrap(); + let grant = client + .logs3() + .request_overflow_upload(Logs3OverflowUploadRequest { + content_type: "application/json".into(), + size_bytes: 123, + rows: vec![], + }) + .await + .unwrap(); + assert_eq!(grant.key, "key"); + assert_eq!(server.received_requests().await.unwrap().len(), 1); +} + +#[test] +fn config_rejects_ambiguous_routes_and_redacts_credentials() { + let mut settings = api::Config::new("do-not-print-token"); + settings.api_url = "https://user:secret@example.test".into(); + assert!(!format!("{settings:?}").contains("secret")); + assert!(!format!("{settings:?}").contains("do-not-print-token")); + for url in [ + "file:///tmp/foo", + "https://u:p@example.test", + "https://example.test/?key=secret", + "https://example.test/#fragment", + ] { + settings.api_url = url.into(); + assert!(matches!( + api::Client::new(settings.clone()), + Err(api::Error::Config(_)) + )); + } + assert!(api::Client::new(api::Config::new("\ninvalid-token")).is_err()); +} + +#[tokio::test] +async fn base_path_is_preserved_without_a_trailing_slash() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/prefix/version")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({}))) + .expect(1) + .mount(&server) + .await; + let mut settings = config(&server, &server, "token"); + settings.api_url = format!("{}/prefix", server.uri()); + assert!(api::Client::new(settings) + .unwrap() + .logs3() + .version() + .await + .unwrap() + .logs3_payload_max_bytes + .is_none()); +} + +#[tokio::test] +async fn typed_dataset_mutations_use_the_data_plane_once_without_trace_fields() { + use api::datasets::{DatasetFields, DatasetWrite}; + let data = MockServer::start().await; + let control = MockServer::start().await; + let client = api::Client::new(config(&data, &control, "token")).unwrap(); + let writes = [ + DatasetWrite::Insert { + dataset_id: "dataset".into(), + id: "row".into(), + fields: DatasetFields { + input: Some(json!(null)), + ..Default::default() + }, + }, + DatasetWrite::Update { + dataset_id: "dataset".into(), + id: "row".into(), + fields: DatasetFields { + expected: Some(json!(42)), + ..Default::default() + }, + }, + DatasetWrite::Delete { + dataset_id: "dataset".into(), + id: "row".into(), + }, + ]; + Mock::given(method("POST")) + .and(path("/logs3")) + .and(body_json(json!({"api_version":2,"rows":[ + {"dataset_id":"dataset","id":"row","input":null}, + {"dataset_id":"dataset","id":"row","expected":42,"_is_merge":true}, + {"dataset_id":"dataset","id":"row","_object_delete":true} + ]}))) + .respond_with(ResponseTemplate::new(413).set_body_string("too large")) + .expect(1) + .mount(&data) + .await; + assert!(matches!( + client.datasets().write(&writes).await, + Err(api::Error::Http { status: 413, .. }) + )); + assert_eq!(data.received_requests().await.unwrap().len(), 1); + assert!(control.received_requests().await.unwrap().is_empty()); +} diff --git a/tests/credential_scopes.rs b/tests/credential_scopes.rs new file mode 100644 index 0000000..b9a4c65 --- /dev/null +++ b/tests/credential_scopes.rs @@ -0,0 +1,404 @@ +use braintrust_sdk_rust::{BraintrustClient, DatasetInsert, SpanLog}; +use serde_json::json; +use wiremock::{ + matchers::{header, method, path}, + Match, Mock, MockServer, ResponseTemplate, +}; + +fn registration() -> ResponseTemplate { + ResponseTemplate::new(200).set_body_json( + json!({"project":{"id":"project","name":"p"},"dataset":{"id":"dataset","name":"d"}}), + ) +} +fn record() -> DatasetInsert { + DatasetInsert::builder().input(json!(1)).build().unwrap() +} + +#[tokio::test] +async fn explicit_handles_keep_their_routes_after_shared_login_state_is_populated() { + let configured = MockServer::start().await; + let session = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/project/register")) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"project":{"id":"explicit-project"}})), + ) + .expect(1) + .mount(&configured) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .and(header("authorization", "Bearer explicit")) + .respond_with(ResponseTemplate::new(204)) + .expect(1) + .mount(&configured) + .await; + let client = BraintrustClient::builder() + .skip_login(true) + .api_url(configured.uri()) + .app_url(configured.uri()) + .build() + .await + .unwrap(); + let builder = client + .span_builder_with_credentials("explicit", "org") + .project_name("p"); + assert!(client.login_state().await.set( + "session-secret".into(), + "other-org".into(), + "other".into(), + session.uri(), + session.uri() + )); + let span = builder.build(); + let write = span.log(SpanLog::builder().input(json!("explicit")).build().unwrap()); + client.flush().await.unwrap(); + write.await.unwrap(); + assert!(session.received_requests().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn logged_in_resources_capture_the_org_route_while_explicit_resources_keep_configured_routes() +{ + let control = MockServer::start().await; + let configured = MockServer::start().await; + let org_data = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/apikey/login")) + .respond_with(ResponseTemplate::new(200).set_body_json( + json!({"org_info":[{"id":"org", "name":"tenant", "api_url":org_data.uri()}]}), + )) + .expect(1) + .mount(&control) + .await; + Mock::given(method("POST")) + .and(path("/api/dataset/register")) + .respond_with(registration()) + .expect(2) + .mount(&control) + .await; + for server in [&configured, &org_data] { + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(204)) + .expect(1) + .mount(server) + .await; + Mock::given(method("POST")) + .and(path("/btql")) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"data":[{"id":"record"}]})), + ) + .expect(1) + .mount(server) + .await; + Mock::given(method("GET")).and(path("/dataset-summary")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"project_name":"p","project_id":"project","dataset_name":"d","dataset_id":"dataset"}))) + .expect(1).mount(server).await; + } + let client = BraintrustClient::builder() + .api_key("session") + .blocking_login(true) + .app_url(control.uri()) + .api_url(configured.uri()) + .build() + .await + .unwrap(); + let session = client + .dataset_builder() + .await + .unwrap() + .project_name("p") + .dataset_name("d") + .build() + .unwrap(); + let explicit = client + .dataset_builder_with_credentials("explicit", "explicit-tenant") + .project_name("p") + .dataset_name("d") + .build() + .unwrap(); + for dataset in [session.clone(), explicit] { + dataset.insert(record()).await.unwrap(); + assert_eq!( + dataset + .fetch(None) + .await + .unwrap() + .collect() + .await + .unwrap() + .len(), + 1 + ); + dataset.summarize().await.unwrap(); + } + for (server, token, org) in [ + (&configured, "explicit", "explicit-tenant"), + (&org_data, "session", "tenant"), + ] { + for request in server.received_requests().await.unwrap() { + assert!(header("authorization", format!("Bearer {token}").as_str()).matches(&request)); + assert!(header("x-bt-org-name", org).matches(&request)); + } + } +} + +#[tokio::test] +async fn organization_headers_partition_registration_capabilities_and_uploads() { + let server = MockServer::start().await; + for org in ["tenant-a", "tenant-b"] { + Mock::given(method("POST")) + .and(path("/api/project/register")) + .and(header("x-bt-org-name", org)) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"project":{"id":org}}))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/version")) + .and(header("x-bt-org-name", org)) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({}))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .and(header("x-bt-org-name", org)) + .respond_with(ResponseTemplate::new(204)) + .expect(1) + .mount(&server) + .await; + } + let client = BraintrustClient::builder() + .skip_login(true) + .app_url(server.uri()) + .api_url(server.uri()) + .build() + .await + .unwrap(); + let builder = client + .span_builder_with_credentials("shared-token", "shared-org") + .project_name("same-name"); + let a = builder.clone().org_name("tenant-a").build(); + let b = builder.org_name("tenant-b").build(); + let first = a.log(SpanLog::default()); + let second = b.log(SpanLog::default()); + client.flush().await.unwrap(); + first.await.unwrap(); + second.await.unwrap(); + for request in server + .received_requests() + .await + .unwrap() + .iter() + .filter(|r| r.url.path() == "/logs3") + { + let body: serde_json::Value = serde_json::from_slice(&request.body).unwrap(); + let rows = body["rows"].as_array().unwrap(); + assert_eq!(rows.len(), 1); + let org = rows[0]["project_id"].as_str().unwrap(); + assert!(header("x-bt-org-name", org).matches(request)); + } +} + +#[tokio::test] +async fn mixed_trace_scopes_keep_separate_routes_and_project_cache_entries() { + let control = MockServer::start().await; + let configured = MockServer::start().await; + let selected = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/apikey/login")) + .respond_with(ResponseTemplate::new(200).set_body_json( + json!({"org_info":[{"id":"org","name":"tenant","api_url":selected.uri()}]}), + )) + .mount(&control) + .await; + Mock::given(method("POST")) + .and(path("/api/project/register")) + .and(header("authorization", "Bearer shared")) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"project":{"id":"shared-project"}})), + ) + .expect(2) + .mount(&control) + .await; + for server in [&selected, &configured] { + Mock::given(method("GET")) + .and(path("/version")) + .and(header("authorization", "Bearer shared")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({}))) + .expect(1) + .mount(server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .and(header("authorization", "Bearer shared")) + .respond_with(ResponseTemplate::new(204)) + .expect(2) + .mount(server) + .await; + } + let client = BraintrustClient::builder() + .api_key("shared") + .blocking_login(true) + .app_url(control.uri()) + .api_url(configured.uri()) + .build() + .await + .unwrap(); + // Identical credentials and selectors still have separate data routes. + let logged = client + .span_builder() + .await + .unwrap() + .project_name("p") + .build(); + let explicit = client + .span_builder_with_credentials("shared", "org") + .org_name("tenant") + .project_name("p") + .build(); + for _ in 0..2 { + let a = logged.log(SpanLog::default()); + let b = explicit.clone().log(SpanLog::default()); + client.flush().await.unwrap(); + a.await.unwrap(); + b.await.unwrap(); + } + for server in [&selected, &configured] { + for request in server + .received_requests() + .await + .unwrap() + .iter() + .filter(|r| r.url.path() == "/logs3") + { + let body: serde_json::Value = serde_json::from_slice(&request.body).unwrap(); + assert_eq!(body["rows"].as_array().unwrap().len(), 1); + assert_eq!(body["rows"][0]["project_id"], "shared-project"); + assert!(header("x-bt-org-name", "tenant").matches(request)); + } + } +} + +#[tokio::test] +async fn experiment_clones_and_child_spans_preserve_the_explicit_scope() { + let server = MockServer::start().await; + for (route, response) in [ + ( + "/api/experiment/register", + json!({"project":{"id":"project","name":"p"},"experiment":{"id":"experiment","name":"e"}}), + ), + ( + "/api/base_experiment/get_id", + json!({"base_exp_id":"base","base_exp_name":"base"}), + ), + ] { + Mock::given(method("POST")) + .and(path(route)) + .and(header("authorization", "Bearer explicit")) + .and(header("x-bt-org-name", "tenant")) + .respond_with(ResponseTemplate::new(200).set_body_json(response)) + .expect(1) + .mount(&server) + .await; + } + Mock::given(method("GET")) + .and(path("/experiment-comparison2")) + .and(header("x-bt-org-name", "tenant")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"scores":{},"metrics":{}}))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .and(header("authorization", "Bearer explicit")) + .and(header("x-bt-org-name", "tenant")) + .respond_with(ResponseTemplate::new(204)) + .expect(1) + .mount(&server) + .await; + let client = BraintrustClient::builder() + .skip_login(true) + .app_url(server.uri()) + .api_url(server.uri()) + .build() + .await + .unwrap(); + let experiment = client + .experiment_builder_with_credentials("explicit", "org") + .org_name("tenant") + .project_name("p") + .build() + .unwrap(); + let span = experiment.start_span().await.unwrap().build(); + span.log(SpanLog::default()).await.unwrap(); + experiment.clone().summarize().await.unwrap(); + assert!(!client.is_logged_in().await); +} + +#[tokio::test] +async fn span_builder_debug_redacts_resource_credentials() { + let client = BraintrustClient::builder() + .skip_login(true) + .build() + .await + .unwrap(); + let builder = client.span_builder_with_credentials("private-resource-token", "org"); + assert!(!format!("{builder:?}").contains("private-resource-token")); +} + +#[tokio::test] +async fn project_registration_cannot_reuse_another_tokens_authorization() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/project/register")) + .and(header("authorization", "Bearer allowed")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"project":{"id":"project"}}))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/api/project/register")) + .and(header("authorization", "Bearer rejected")) + .respond_with(ResponseTemplate::new(403)) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .and(header("authorization", "Bearer allowed")) + .respond_with(ResponseTemplate::new(204)) + .expect(1) + .mount(&server) + .await; + let client = BraintrustClient::builder() + .skip_login(true) + .api_url(server.uri()) + .app_url(server.uri()) + .build() + .await + .unwrap(); + let allowed = client + .span_builder_with_credentials("allowed", "org") + .project_name("p") + .build(); + allowed.log(SpanLog::default()).await.unwrap(); + let rejected = client + .span_builder_with_credentials("rejected", "org") + .project_name("p") + .build(); + let error = rejected.log(SpanLog::default()).await.unwrap_err(); + assert!(error.to_string().contains("403")); + assert_eq!( + server + .received_requests() + .await + .unwrap() + .iter() + .filter(|r| r.url.path() == "/logs3") + .count(), + 1 + ); +} diff --git a/tests/dataset_test.rs b/tests/dataset_test.rs index 928d26b..6de0cd4 100644 --- a/tests/dataset_test.rs +++ b/tests/dataset_test.rs @@ -83,7 +83,7 @@ async fn dataset_builder_requires_dataset_name() { } #[tokio::test] -async fn dataset_insert_flushes_to_logs_endpoint() { +async fn dataset_insert_is_acknowledged_without_flush() { let server = MockServer::start().await; Mock::given(method("POST")) @@ -102,14 +102,6 @@ async fn dataset_insert_flushes_to_logs_endpoint() { .mount(&server) .await; - Mock::given(method("POST")) - .and(path("/api/project/register")) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "project": { "id": "proj-123" } - }))) - .mount(&server) - .await; - Mock::given(method("POST")) .and(path("/logs3")) .respond_with(ResponseTemplate::new(200).set_body_json(json!({}))) @@ -144,16 +136,24 @@ async fn dataset_insert_flushes_to_logs_endpoint() { .build() .expect("record"), ) - .await; + .await + .expect("insert"); assert!(!record_id.is_empty()); - // Flush to send to API - dataset.flush().await.expect("flush"); - client.flush().await.expect("client flush"); - - // Give the background worker time to process - tokio::time::sleep(std::time::Duration::from_millis(100)).await; + let requests = server.received_requests().await.unwrap(); + let writes: Vec<_> = requests + .iter() + .filter(|request| request.url.path() == "/logs3") + .collect(); + assert_eq!(writes.len(), 1); + assert_eq!( + serde_json::from_slice::(&writes[0].body).unwrap(), + json!({ + "api_version": 2, "rows": [{"id":record_id, "dataset_id":"ds-456", + "input":{"question":"What is 2+2?"}, "expected":{"answer":"4"}}] + }) + ); } #[tokio::test] @@ -216,7 +216,7 @@ async fn dataset_fetch_returns_records() { // Fetch records let iter = dataset.fetch(Some(10)).await.expect("fetch"); - let records = iter.collect().await; + let records = iter.collect().await.expect("collect records"); assert_eq!(records.len(), 2); assert_eq!(records[0].id, "rec-1"); @@ -285,3 +285,475 @@ async fn dataset_summarize_returns_summary() { assert_eq!(summary.dataset_id(), Some("ds-456")); assert_eq!(summary.total_records(), Some(42)); } + +async fn isolated_client(server: &MockServer) -> BraintrustClient { + BraintrustClient::builder() + .skip_login(true) + .app_url(server.uri()) + .api_url(server.uri()) + .build() + .await + .unwrap() +} +fn registration() -> ResponseTemplate { + ResponseTemplate::new(200).set_body_json(json!({"project":{"id":"project","name":"project"},"dataset":{"id":"dataset","name":"dataset"}})) +} +fn dataset(client: &BraintrustClient) -> braintrust_sdk_rust::Dataset { + client + .dataset_builder_with_credentials("token", "tenant") + .project_name("project") + .dataset_name("dataset") + .build() + .unwrap() +} + +#[tokio::test] +async fn failed_dataset_registration_is_retryable_and_clone_registration_is_shared() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/dataset/register")) + .respond_with(ResponseTemplate::new(503)) + .with_priority(1) + .up_to_n_times(1) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/api/dataset/register")) + .respond_with(registration()) + .with_priority(2) + .expect(1) + .mount(&server) + .await; + let client = isolated_client(&server).await; + let dataset = dataset(&client); + let clone = dataset.clone(); + assert!(matches!( + dataset.dataset_id().await, + Err(braintrust_sdk_rust::BraintrustError::Api { status: 503, .. }) + )); + let (first, second) = tokio::join!(dataset.dataset_id(), clone.dataset_id()); + assert_eq!(first.unwrap(), "dataset"); + assert_eq!(second.unwrap(), "dataset"); + assert_eq!(dataset.clone().dataset_id().await.unwrap(), "dataset"); + assert!(!client.is_logged_in().await); +} + +#[tokio::test] +async fn pagination_failure_is_an_error_and_retry_preserves_the_failed_cursor() { + use wiremock::matchers::body_partial_json; + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/dataset/register")) + .respond_with(registration()) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/btql")) + .and(body_partial_json(json!({"query":{"cursor":"next"}}))) + .respond_with(ResponseTemplate::new(503)) + .with_priority(1) + .up_to_n_times(1) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/btql")) + .and(body_partial_json(json!({"query":{"cursor":"next"}}))) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(json!({"data":[{"id":"second"}],"cursor":null})), + ) + .with_priority(2) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/btql")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(json!({"data":[{"id":"first"}],"cursor":"next"})), + ) + .with_priority(3) + .expect(1) + .mount(&server) + .await; + let client = isolated_client(&server).await; + let dataset = dataset(&client); + let mut iterator = dataset.fetch(Some(1)).await.unwrap(); + assert_eq!(iterator.next().await.unwrap().unwrap().id, "first"); + assert!(matches!( + iterator.next().await, + Err(braintrust_sdk_rust::BraintrustError::Api { status: 503, .. }) + )); + assert_eq!(iterator.next().await.unwrap().unwrap().id, "second"); + assert!(iterator.next().await.unwrap().is_none()); + let requests = server.received_requests().await.unwrap(); + let cursors: Vec<_> = requests + .iter() + .filter(|request| request.url.path() == "/btql") + .map(|request| { + serde_json::from_slice::(&request.body).unwrap()["query"]["cursor"] + .clone() + }) + .collect(); + assert_eq!(cursors, vec![json!(null), json!("next"), json!("next")]); +} + +#[tokio::test] +async fn collecting_records_propagates_page_errors_instead_of_returning_partial_success() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/dataset/register")) + .respond_with(registration()) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/btql")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(json!({"data":[{"id":"first"}],"cursor":"next"})), + ) + .with_priority(1) + .up_to_n_times(1) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/btql")) + .respond_with(ResponseTemplate::new(401)) + .with_priority(2) + .expect(1) + .mount(&server) + .await; + let client = isolated_client(&server).await; + let records = dataset(&client).fetch(None).await.unwrap().collect().await; + assert!(matches!( + records, + Err(braintrust_sdk_rust::BraintrustError::Api { status: 401, .. }) + )); +} + +#[tokio::test] +async fn empty_pages_with_a_continuation_do_not_end_iteration() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/dataset/register")) + .respond_with(registration()) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/btql")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"data":[],"cursor":"next"}))) + .with_priority(1) + .up_to_n_times(1) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/btql")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(json!({"data":[{"id":"record"}],"cursor":null})), + ) + .with_priority(2) + .expect(1) + .mount(&server) + .await; + let client = isolated_client(&server).await; + let records = dataset(&client) + .fetch(Some(0)) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(records.len(), 1); + assert_eq!(records[0].id, "record"); + for request in server + .received_requests() + .await + .unwrap() + .iter() + .filter(|request| request.url.path() == "/btql") + { + assert_eq!( + serde_json::from_slice::(&request.body).unwrap()["query"]["limit"], + 1 + ); + } +} + +#[tokio::test] +async fn a_non_advancing_cursor_is_reported_and_does_not_poison_a_later_retry() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/dataset/register")) + .respond_with(registration()) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/btql")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(json!({"data":[{"id":"first"}],"cursor":"stuck"})), + ) + .with_priority(1) + .up_to_n_times(2) + .expect(2) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/btql")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(json!({"data":[{"id":"second"}],"cursor":null})), + ) + .with_priority(2) + .expect(1) + .mount(&server) + .await; + let client = isolated_client(&server).await; + let dataset = dataset(&client); + let mut iterator = dataset.fetch(None).await.unwrap(); + assert_eq!(iterator.next().await.unwrap().unwrap().id, "first"); + assert!( + matches!(iterator.next().await, Err(braintrust_sdk_rust::BraintrustError::Background(message)) if message.contains("did not advance")) + ); + assert_eq!(iterator.next().await.unwrap().unwrap().id, "second"); + assert!(iterator.next().await.unwrap().is_none()); +} + +fn record(value: i32) -> DatasetInsert { + DatasetInsert::builder() + .input(json!(value)) + .build() + .unwrap() +} + +#[tokio::test] +async fn mutations_have_dataset_wire_shapes_and_share_registration_across_clones() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/dataset/register")) + .respond_with(registration()) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(204)) + .expect(3) + .mount(&server) + .await; + let client = isolated_client(&server).await; + let dataset = dataset(&client); + let id = dataset + .insert( + DatasetInsert::builder() + .input(json!(1)) + .expected(json!(2)) + .metadata_entry("source", "test") + .tag("sample") + .build() + .unwrap(), + ) + .await + .unwrap(); + dataset.clone().update(&id, record(3)).await.unwrap(); + dataset.delete(&id).await.unwrap(); + // Neither barrier registers or replays dataset work. + dataset.flush().await.unwrap(); + client.flush().await.unwrap(); + let requests = server.received_requests().await.unwrap(); + let bodies: Vec = requests + .iter() + .filter(|r| r.url.path() == "/logs3") + .map(|r| serde_json::from_slice(&r.body).unwrap()) + .collect(); + assert_eq!( + bodies, + vec![ + json!({"api_version":2,"rows":[{"dataset_id":"dataset","id":id,"input":1,"expected":2,"metadata":{"source":"test"},"tags":["sample"]}]}), + json!({"api_version":2,"rows":[{"dataset_id":"dataset","id":id,"input":3,"_is_merge":true}]}), + json!({"api_version":2,"rows":[{"dataset_id":"dataset","id":id,"_object_delete":true}]}), + ] + ); + for request in requests { + assert_eq!(request_header(&request, "authorization"), "Bearer token"); + assert_eq!(request_header(&request, "x-bt-org-name"), "tenant"); + } +} + +#[tokio::test] +async fn mutation_errors_are_direct_and_do_not_replay_or_poison_later_operations() { + for status in [401, 413, 429, 503] { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/dataset/register")) + .respond_with(registration()) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(status).set_body_string("rejected")) + .with_priority(1) + .up_to_n_times(3) + .expect(3) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(204)) + .with_priority(2) + .expect(1) + .mount(&server) + .await; + let client = isolated_client(&server).await; + let dataset = dataset(&client); + let insert = dataset.insert(record(1)).await.map(|_| ()); + let update = dataset.update("old", record(2)).await; + let delete = dataset.delete("old").await; + for result in [insert, update, delete] { + match result { + Err(braintrust_sdk_rust::BraintrustError::Api { + status: actual, + message, + }) => { + assert_eq!(actual, status); + assert_eq!(message, "rejected"); + } + other => panic!("expected API error, got {other:?}"), + } + } + dataset.flush().await.unwrap(); + let id = dataset.insert(record(4)).await.unwrap(); + dataset.flush().await.unwrap(); + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests.len(), 5); // One registration and four one-attempt writes. + assert_eq!( + serde_json::from_slice::(&requests[4].body).unwrap()["rows"][0] + ["id"], + id + ); + } +} + +#[tokio::test] +async fn registration_failure_prevents_mutation_and_flush_does_not_register_a_dataset() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/dataset/register")) + .respond_with(ResponseTemplate::new(401)) + .expect(3) + .mount(&server) + .await; + let client = isolated_client(&server).await; + let dataset = dataset(&client); + dataset.flush().await.unwrap(); + assert!(server.received_requests().await.unwrap().is_empty()); + assert!(dataset.insert(record(1)).await.is_err()); + assert!(dataset.update("old", record(2)).await.is_err()); + assert!(dataset.delete("old").await.is_err()); + assert!(server + .received_requests() + .await + .unwrap() + .iter() + .all(|r| r.url.path() == "/api/dataset/register")); +} + +#[tokio::test] +async fn resource_operations_preserve_configured_planes_prefixes_and_transport_headers() { + use wiremock::matchers::{header, query_param}; + let control = MockServer::start().await; + let data = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/control/api/dataset/register")) + .and(header("authorization", "Bearer token")) + .and(header("x-bt-org-name", "tenant")) + .respond_with(registration()) + .expect(1) + .mount(&control) + .await; + Mock::given(method("POST")) + .and(path("/data/logs3")) + .respond_with(ResponseTemplate::new(204)) + .expect(1) + .mount(&data) + .await; + Mock::given(method("POST")) + .and(path("/data/btql")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"data":[{"id":"record"}]}))) + .expect(1) + .mount(&data) + .await; + Mock::given(method("GET")).and(path("/data/dataset-summary")) + .and(query_param("dataset_id", "dataset")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "project_name":"project","dataset_name":"dataset","project_id":"project","dataset_id":"dataset" + }))).expect(1).mount(&data).await; + let client = BraintrustClient::builder() + .skip_login(true) + .app_url(format!("{}/control", control.uri())) + .api_url(format!("{}/data", data.uri())) + .user_agent("resource-client/1") + .build() + .await + .unwrap(); + let dataset = dataset(&client); + dataset.insert(record(1)).await.unwrap(); + assert_eq!( + dataset.fetch(None).await.unwrap().collect().await.unwrap()[0].id, + "record" + ); + dataset.summarize().await.unwrap(); + let requests = data.received_requests().await.unwrap(); + assert_eq!(requests.len(), 3); + for request in requests { + assert_eq!(request_header(&request, "authorization"), "Bearer token"); + assert_eq!(request_header(&request, "x-bt-org-name"), "tenant"); + assert_eq!(request_header(&request, "user-agent"), "resource-client/1"); + } + assert_eq!(control.received_requests().await.unwrap().len(), 1); +} + +fn request_header<'a>(request: &'a wiremock::Request, name: &str) -> &'a str { + request + .headers + .iter() + .find(|(key, _)| key.as_str() == name) + .unwrap() + .1[0] + .as_str() +} + +#[tokio::test] +async fn dataset_flush_observes_the_shared_trace_queue_and_retires_its_failure() { + use braintrust_sdk_rust::{FlushError, ParentSpanInfo, SpanLog}; + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(401)) + .expect(1) + .mount(&server) + .await; + let client = isolated_client(&server).await; + let dataset = dataset(&client); + let span = client + .span_builder_with_credentials("token", "org") + .parent_info(ParentSpanInfo::Experiment { + object_id: "experiment".into(), + }) + .build(); + let write = span.log(SpanLog::builder().input(json!(1)).build().unwrap()); + let flush = dataset.flush(); + assert!(matches!(flush.await, Err(FlushError::Delivery(failure)) if failure.failed == 1)); + assert!(write.await.is_err()); + assert_eq!(dataset.flush().await.unwrap().delivered, 0); + let requests = server.received_requests().await.unwrap(); + assert!(requests + .iter() + .all(|r| r.url.path() != "/api/dataset/register")); +} diff --git a/tests/eval_test.rs b/tests/eval_test.rs index 93555be..b07eb0d 100644 --- a/tests/eval_test.rs +++ b/tests/eval_test.rs @@ -1,3 +1,4 @@ +use braintrust_sdk_rust::eval::Dataset; use braintrust_sdk_rust::eval::*; use braintrust_sdk_rust::*; diff --git a/tests/experiments.rs b/tests/experiments.rs index ff4852b..9de77b4 100644 --- a/tests/experiments.rs +++ b/tests/experiments.rs @@ -47,7 +47,8 @@ async fn experiment_log_flushes_to_logs_endpoint() { .score("accuracy", 1.0) .build(), ) - .await; + .await + .expect("log event"); experiment.flush().await.expect("flush"); client.flush().await.expect("client flush"); @@ -73,6 +74,7 @@ async fn experiment_log_flushes_to_logs_endpoint() { #[tokio::test] async fn experiment_builder_requires_project_name() { let client = BraintrustClient::builder() + .skip_login(true) .api_key("token") .api_url("https://api.braintrust.dev") .build() @@ -124,7 +126,8 @@ async fn experiment_log_feedback_sends_merge_event() { // Log an event let id = experiment .log(ExperimentLog::builder().input(json!("test input")).build()) - .await; + .await + .expect("log event"); // Log feedback on the event experiment @@ -136,7 +139,8 @@ async fn experiment_log_feedback_sends_merge_event() { .build() .expect("build feedback"), ) - .await; + .await + .expect("log event"); experiment.flush().await.expect("flush"); client.flush().await.expect("client flush"); @@ -215,7 +219,8 @@ async fn experiment_summarize_returns_summary() { // Log an event first experiment .log(ExperimentLog::builder().input(json!("test")).build()) - .await; + .await + .expect("log event"); let summary = experiment.summarize().await.expect("summarize"); @@ -233,3 +238,333 @@ async fn experiment_summarize_returns_summary() { assert_eq!(accuracy_score.improvements(), Some(10)); assert_eq!(accuracy_score.regressions(), Some(2)); } + +async fn isolated_client(server: &MockServer) -> BraintrustClient { + BraintrustClient::builder() + .skip_login(true) + .app_url(server.uri()) + .api_url(server.uri()) + .build() + .await + .unwrap() +} +fn registered_experiment() -> ResponseTemplate { + ResponseTemplate::new(200).set_body_json(json!({"project":{"id":"project","name":"project"},"experiment":{"id":"experiment","name":"experiment"}})) +} +fn assert_api_error(result: Result, expected: u16) { + assert!( + matches!(result, Err(braintrust_sdk_rust::BraintrustError::Api {status, ..}) if status == expected) + ); +} + +#[tokio::test] +async fn registration_failure_has_no_invented_id_and_clones_share_a_successful_retry() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/experiment/register")) + .respond_with(ResponseTemplate::new(503)) + .with_priority(1) + .up_to_n_times(1) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/api/experiment/register")) + .respond_with(registered_experiment()) + .with_priority(2) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + let client = isolated_client(&server).await; + let experiment = client + .experiment_builder_with_credentials("token", "org") + .project_name("project") + .build() + .unwrap(); + let clone = experiment.clone(); + assert_api_error(experiment.project().await, 503); + assert_eq!(experiment.experiment_id(), None); + assert_eq!(clone.experiment_id(), None); + let (first, second) = tokio::join!(experiment.project(), clone.project()); + assert_eq!(first.unwrap().id, "project"); + assert_eq!(second.unwrap().id, "project"); + assert_eq!(experiment.experiment_id().as_deref(), Some("experiment")); + assert_eq!(clone.experiment_id().as_deref(), Some("experiment")); + assert_eq!( + experiment.clone().export().await.unwrap().object_id, + "experiment" + ); + experiment + .log(ExperimentLog::builder().input(json!("value")).build()) + .await + .unwrap(); + let requests = server.received_requests().await.unwrap(); + let body: Value = serde_json::from_slice( + &requests + .iter() + .find(|req| req.url.path() == "/logs3") + .unwrap() + .body, + ) + .unwrap(); + assert_eq!(body["rows"][0]["experiment_id"], "experiment"); + assert!(body["rows"][0]["metrics"]["end"].is_number()); +} + +#[tokio::test] +async fn every_registration_dependent_operation_returns_the_real_failure() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/experiment/register")) + .respond_with(ResponseTemplate::new(400)) + .expect(9) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&server) + .await; + let client = isolated_client(&server).await; + let experiment = client + .experiment_builder_with_credentials("token", "org") + .project_name("project") + .build() + .unwrap(); + assert_api_error(experiment.start_span().await, 400); + assert_api_error( + experiment + .log(ExperimentLog::builder().input(json!("value")).build()) + .await, + 400, + ); + assert_api_error(experiment.project().await, 400); + assert_api_error(experiment.export().await, 400); + assert_api_error( + experiment + .log_feedback("row", Feedback::builder().build().unwrap()) + .await, + 400, + ); + assert_api_error( + experiment + .update_span("row", braintrust_sdk_rust::SpanLog::default()) + .await, + 400, + ); + assert_api_error(experiment.fetch_base_experiment().await, 400); + assert_api_error(experiment.summarize().await, 400); + let called = std::sync::atomic::AtomicBool::new(false); + assert_api_error( + experiment + .traced(|_| async { + called.store(true, std::sync::atomic::Ordering::Relaxed); + 42 + }) + .await, + 400, + ); + assert!(!called.load(std::sync::atomic::Ordering::Relaxed)); + assert_eq!(experiment.experiment_id(), None); +} + +#[tokio::test] +async fn event_feedback_and_update_methods_propagate_delivery_rejection() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/experiment/register")) + .respond_with(registered_experiment()) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(401)) + .expect(3) + .mount(&server) + .await; + let client = isolated_client(&server).await; + let experiment = client + .experiment_builder_with_credentials("token", "org") + .project_name("project") + .build() + .unwrap(); + assert_api_error( + experiment + .log(ExperimentLog::builder().input(json!("value")).build()) + .await, + 401, + ); + assert_api_error( + experiment + .log_feedback("row", Feedback::builder().build().unwrap()) + .await, + 401, + ); + assert_api_error( + experiment + .update_span("row", braintrust_sdk_rust::SpanLog::default()) + .await, + 401, + ); + experiment.flush().await.unwrap(); // completed failures were already observed and retired +} + +#[tokio::test] +async fn comparison_setup_errors_are_not_silently_treated_as_no_base_experiment() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/experiment/register")) + .respond_with(registered_experiment()) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/api/base_experiment/get_id")) + .respond_with(ResponseTemplate::new(401)) + .expect(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/experiment-comparison2")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&server) + .await; + let client = isolated_client(&server).await; + let experiment = client + .experiment_builder_with_credentials("token", "org") + .project_name("project") + .build() + .unwrap(); + assert_api_error(experiment.summarize().await, 401); +} + +#[tokio::test] +async fn explicit_experiment_credentials_never_prime_the_shared_session() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/experiment/register")) + .respond_with(registered_experiment()) + .expect(1) + .mount(&server) + .await; + let client = isolated_client(&server).await; + let state = client.login_state().await; + let experiment = client + .experiment_builder_with_credentials("private-token", "tenant") + .project_name("project") + .build() + .unwrap(); + assert!(!state.is_logged_in()); + experiment.project().await.unwrap(); + assert!(!state.is_logged_in()); + assert!(!client.is_logged_in().await); + assert!(client.dataset_builder().await.is_err()); + let exported = experiment.export().await.unwrap(); + assert_eq!(exported.object_id, "experiment"); + assert!(!state.is_logged_in()); +} + +#[tokio::test] +async fn traced_propagates_final_flush_errors_instead_of_discarding_them() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/experiment/register")) + .respond_with(registered_experiment()) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(401)) + .expect(1) + .mount(&server) + .await; + let client = isolated_client(&server).await; + let experiment = client + .experiment_builder_with_credentials("token", "org") + .project_name("project") + .build() + .unwrap(); + assert_api_error(experiment.traced(|_| async { 42 }).await, 401); +} + +#[tokio::test] +async fn clones_share_the_selected_span_recording_interface() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/experiment/register")) + .respond_with(registered_experiment()) + .expect(1) + .mount(&server) + .await; + let client = isolated_client(&server).await; + let experiment = client + .experiment_builder_with_credentials("token", "org") + .project_name("project") + .build() + .unwrap(); + let clone = experiment.clone(); + experiment.start_span().await.unwrap(); + assert!(matches!( + clone + .log(ExperimentLog::builder().input(json!("value")).build()) + .await, + Err(braintrust_sdk_rust::BraintrustError::InvalidConfig(_)) + )); +} + +#[tokio::test] +async fn experiment_resource_adapters_use_the_control_plane_and_preserve_no_base_policy() { + use wiremock::matchers::{body_json, header, query_param}; + let control = MockServer::start().await; + let data = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/control/api/experiment/register")) + .and(header("authorization", "Bearer explicit")) + .respond_with(registered_experiment()) + .expect(1) + .mount(&control) + .await; + Mock::given(method("POST")) + .and(path("/control/api/base_experiment/get_id")) + .and(body_json(json!({"id":"experiment"}))) + .respond_with(ResponseTemplate::new(400)) + .expect(1) + .mount(&control) + .await; + Mock::given(method("GET")) + .and(path("/control/experiment-comparison2")) + .and(query_param("experiment_id", "experiment")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"scores":{},"metrics":{}}))) + .expect(1) + .mount(&control) + .await; + let client = BraintrustClient::builder() + .skip_login(true) + .app_url(format!("{}/control", control.uri())) + .api_url(format!("{}/data", data.uri())) + .build() + .await + .unwrap(); + let experiment = client + .experiment_builder_with_credentials("explicit", "org") + .project_name("project") + .build() + .unwrap(); + experiment.summarize().await.unwrap(); + let requests = control.received_requests().await.unwrap(); + assert_eq!(requests.len(), 3); + assert!(requests[2] + .url + .query_pairs() + .all(|(key, _)| key != "base_experiment_id")); + assert!(data.received_requests().await.unwrap().is_empty()); +} diff --git a/tests/finished_spans.rs b/tests/finished_spans.rs new file mode 100644 index 0000000..a2fa11f --- /dev/null +++ b/tests/finished_spans.rs @@ -0,0 +1,215 @@ +use std::sync::Arc; + +use braintrust_sdk_rust::{ + BraintrustClient, DeliveryError, EnqueueError, FlushError, ParentSpanInfo, SpanHandle, SpanLog, + WriteError, +}; +use serde_json::{json, Value}; +use wiremock::{ + matchers::{method, path}, + Mock, MockServer, ResponseTemplate, +}; + +async fn client(server: &MockServer) -> BraintrustClient { + BraintrustClient::builder() + .skip_login(true) + .api_url(server.uri()) + .app_url(server.uri()) + .build() + .await + .unwrap() +} +fn span(client: &BraintrustClient) -> SpanHandle { + client + .span_builder_with_credentials("token", "org") + .parent_info(ParentSpanInfo::ProjectLogs { + object_id: "project".into(), + }) + .build() +} +async fn delivery_error( + completion: braintrust_sdk_rust::Completion, +) -> Arc { + assert!(completion.is_ok(), "observation was issued successfully"); + match completion.await { + Err(WriteError::Delivery(error)) => error, + other => panic!("expected delivery error, got {other:?}"), + } +} + +#[tokio::test] +async fn repeated_ends_share_a_retired_failure_without_replaying_or_poisoning_fresh_work() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(401)) + .with_priority(1) + .up_to_n_times(1) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(204)) + .with_priority(2) + .expect(1) + .mount(&server) + .await; + let client = client(&server).await; + let finished = span(&client); + let first = finished.end_with_time(123.0); + let second = finished.clone().end_with_time(456.0); + let flush = client.flush(); + let error = delivery_error(first).await; + assert!(Arc::ptr_eq(&error, &delivery_error(second).await)); + match flush.await { + Err(FlushError::Delivery(failure)) => { + assert_eq!(failure.failed, 1); + assert!(Arc::ptr_eq(&error, &failure.cause)); + } + other => panic!("expected failed flush, got {other:?}"), + } + assert_eq!(client.flush().await.unwrap().delivered, 0); + assert!(Arc::ptr_eq(&error, &delivery_error(finished.end()).await)); + assert_eq!(finished.end_time(), Some(123.0)); + let fresh = span(&client); + fresh.end().await.unwrap(); + assert_eq!(client.flush().await.unwrap().delivered, 0); + let weak = Arc::downgrade(&error); + drop(error); + drop(finished); + assert!( + weak.upgrade().is_none(), + "only the live finished span may retain its outcome" + ); + assert_eq!( + server + .received_requests() + .await + .unwrap() + .iter() + .filter(|r| r.url.path() == "/logs3") + .count(), + 2 + ); +} + +#[tokio::test] +async fn finishing_a_parent_does_not_finish_its_child_and_explicit_updates_remain_independent() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(204)) + .expect(4) + .mount(&server) + .await; + let client = client(&server).await; + let parent = span(&client); + let context = parent.export().unwrap(); + let exported = context.to_str(); + let child = client + .span_builder_with_credentials("token", "org") + .parent_info(context.to_parent_span_info().unwrap()) + .build(); + parent.end_with_time(10.0).await.unwrap(); + assert_eq!(child.end_time(), None); + child + .log( + SpanLog::builder() + .input(json!("child still active")) + .build() + .unwrap(), + ) + .await + .unwrap(); + child.end_with_time(20.0).await.unwrap(); + client + .update_span_with_credentials( + "token", + "org", + &exported, + SpanLog::builder() + .output(json!("explicit patch")) + .build() + .unwrap(), + ) + .unwrap(); + client.flush().await.unwrap(); + let bodies: Vec = server + .received_requests() + .await + .unwrap() + .iter() + .filter(|r| r.url.path() == "/logs3") + .map(|r| serde_json::from_slice(&r.body).unwrap()) + .collect(); + assert_eq!(bodies[1]["rows"][0]["input"], "child still active"); + assert!(bodies[1]["rows"][0]["metrics"].get("end").is_none()); + assert_eq!(bodies[2]["rows"][0]["metrics"]["end"], 20.0); + assert_eq!(bodies[3]["rows"][0]["id"], parent.row_id()); + assert_eq!(bodies[3]["rows"][0]["output"], "explicit patch"); + assert_eq!(parent.end_time(), Some(10.0)); +} + +#[tokio::test] +async fn final_admission_rejection_is_visible_and_keeps_the_chosen_timestamp() { + let server = MockServer::start().await; + let client = client(&server).await; + let span = span(&client); + client + .shutdown(std::time::Duration::from_secs(1)) + .await + .unwrap(); + let end = span.end_with_time(123.0); + assert!(matches!( + &*end, + Err(WriteError::Admission(EnqueueError::Closed)) + )); + assert!(matches!( + end.await, + Err(WriteError::Admission(EnqueueError::Closed)) + )); + assert!(matches!( + span.end_with_time(999.0).await, + Err(WriteError::Admission(EnqueueError::Closed)) + )); + assert_eq!(span.end_time(), Some(123.0)); + span.log(SpanLog::builder().input(json!("ignored")).build().unwrap()) + .await + .unwrap(); + assert!(server.received_requests().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn named_experiment_span_keeps_its_initial_name_in_the_final_snapshot() { + let server = MockServer::start().await; + Mock::given(method("POST")).and(path("/api/experiment/register")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"project":{"id":"project","name":"p"},"experiment":{"id":"experiment","name":"e"}}))) + .expect(1).mount(&server).await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(204)) + .expect(1) + .mount(&server) + .await; + let client = client(&server).await; + let experiment = client + .experiment_builder_with_credentials("token", "org") + .project_name("p") + .build() + .unwrap(); + let span = experiment + .start_span() + .await + .unwrap() + .name("original") + .build(); + span.end_with_time(123.0).await.unwrap(); + span.log(SpanLog::builder().name("too late").build().unwrap()) + .await + .unwrap(); + let requests = server.received_requests().await.unwrap(); + let request = requests.iter().find(|r| r.url.path() == "/logs3").unwrap(); + let body: Value = serde_json::from_slice(&request.body).unwrap(); + assert_eq!(body["rows"][0]["span_attributes"]["name"], "original"); +} diff --git a/tests/http_flow.rs b/tests/http_flow.rs index 28e539b..1721184 100644 --- a/tests/http_flow.rs +++ b/tests/http_flow.rs @@ -64,9 +64,8 @@ async fn submits_with_bearer_token() { } #[tokio::test] -async fn flush_is_fire_and_forget() { - // flush() should succeed even when project registration would fail. - // Errors are logged but not propagated (resilient design). +async fn preparation_failure_reaches_write_and_covering_flush() { + // Preparation failure terminates the logical batch, then is forgotten. let server = MockServer::start().await; Mock::given(method("POST")) @@ -88,28 +87,26 @@ async fn flush_is_fire_and_forget() { .project_name("demo") .build(); - span.log( + let completion = span.log( SpanLog::builder() .input(Value::String("input".into())) .build() .expect("build"), ); - // flush() waits for queued work to be processed. Background submission errors - // are still swallowed by the queue and should not bubble out here. - let result = span.flush().await; - assert!( - result.is_ok(), - "flush() should not surface background submission errors" - ); - - // Drain pending items before drop so the Drop impl's synchronous flush path - // is a no-op (queue is empty). The 500 from project registration is swallowed. - let _ = client.flush().await; + assert!(span.flush().await.is_err()); + assert!(matches!( + completion.await, + Err(braintrust_sdk_rust::WriteError::Delivery(_)) + )); + client + .flush() + .await + .expect("old failures must not poison a later flush"); } #[tokio::test] -async fn checked_flush_reports_backend_rejection() { +async fn automatic_delivery_reports_backend_rejection_through_completion() { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/api/project/register")) @@ -126,7 +123,6 @@ async fn checked_flush_reports_backend_rejection() { let client = BraintrustClient::builder() .skip_login(true) - .checked_delivery(true) .api_url(server.uri()) .app_url(server.uri()) .build() @@ -136,9 +132,52 @@ async fn checked_flush_reports_backend_rejection() { .span_builder_with_credentials("secret-token", "org") .project_name("demo") .build(); - span.log(SpanLog::builder().input(json!("input")).build().unwrap()); - assert!(matches!( - client.flush().await, - Err(BraintrustError::Api { status: 401, .. }) - )); + let completion = span.log(SpanLog::builder().input(json!("input")).build().unwrap()); + assert!( + matches!(completion.await, Err(braintrust_sdk_rust::WriteError::Delivery(error)) + if matches!(error.cause().downcast_ref::(), Some(BraintrustError::Api { status: 401, .. }))) + ); + client + .flush() + .await + .expect("completed background failure is retired"); +} + +#[tokio::test] +async fn submission_and_flush_can_be_issued_without_a_call_site_runtime() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/project/register")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"project":{"id":"proj"}}))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + let client = BraintrustClient::builder() + .skip_login(true) + .api_url(server.uri()) + .app_url(server.uri()) + .build() + .await + .unwrap(); + let span = client + .span_builder_with_credentials("token", "org") + .project_name("demo") + .build(); + let (write, flush) = std::thread::spawn(move || { + let write = span.log(SpanLog::builder().input(json!("input")).build().unwrap()); + let flush = client.flush(); + assert!(write.is_ok()); + assert!(flush.is_ok()); + (write, flush) + }) + .join() + .unwrap(); + write.await.unwrap(); + assert_eq!(flush.await.unwrap().delivered, 1); } diff --git a/tests/resource_handles.rs b/tests/resource_handles.rs new file mode 100644 index 0000000..7e6812c --- /dev/null +++ b/tests/resource_handles.rs @@ -0,0 +1,130 @@ +//! Consumer coverage for the concrete public resource model. +use braintrust_sdk_rust::{ + BraintrustClient, Dataset, DatasetBuilder, DatasetInsert, DatasetIterator, Experiment, + ExperimentBuilder, ExperimentSpanBuilder, ParentSpanInfo, SpanBuilder, SpanHandle, SpanLog, +}; +use serde_json::json; +use wiremock::{ + matchers::{method, path}, + Mock, MockServer, ResponseTemplate, +}; + +// Applications can store these types without naming an SDK implementation or +// carrying its private bounds through their own structs and functions. +struct Resources { + span: SpanHandle, + experiment: Experiment, + dataset: Dataset, +} + +#[test] +fn resource_handles_and_builders_are_send_sync_and_owned() { + fn assert_owned() {} + assert_owned::(); + assert_owned::(); + assert_owned::(); + assert_owned::(); + assert_owned::(); + assert_owned::(); + assert_owned::(); + assert_owned::(); +} + +#[tokio::test] +async fn resources_and_detached_span_builders_work_after_the_client_is_dropped() { + let server = MockServer::start().await; + Mock::given(method("POST")).and(path("/api/experiment/register")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"project":{"id":"project","name":"p"},"experiment":{"id":"experiment","name":"e"}}))) + .expect(1).mount(&server).await; + Mock::given(method("POST")) + .and(path("/api/dataset/register")) + .respond_with(ResponseTemplate::new(200).set_body_json( + json!({"project":{"id":"project","name":"p"},"dataset":{"id":"dataset","name":"d"}}), + )) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(204)) + .expect(3) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/btql")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"data":[{"id":"record"}]}))) + .expect(1) + .mount(&server) + .await; + let client = BraintrustClient::builder() + .skip_login(true) + .app_url(server.uri()) + .api_url(server.uri()) + .build() + .await + .unwrap(); + let span: SpanBuilder = client + .span_builder_with_credentials("token", "org") + .parent_info(ParentSpanInfo::ProjectLogs { + object_id: "project".into(), + }); + let experiment: ExperimentBuilder = client + .experiment_builder_with_credentials("token", "org") + .project_name("p"); + let dataset: DatasetBuilder = client + .dataset_builder_with_credentials("token", "tenant") + .project_name("p") + .dataset_name("d"); + drop(client); + let resources = Resources { + span: span.build(), + experiment: experiment.build().unwrap(), + dataset: dataset.build().unwrap(), + }; + resources.span.log(SpanLog::default()).await.unwrap(); + let child: ExperimentSpanBuilder = resources.experiment.start_span().await.unwrap(); + drop(resources.experiment); + let child: SpanHandle = child.build(); + child.log(SpanLog::default()).await.unwrap(); + resources + .dataset + .insert(DatasetInsert::builder().input(json!(1)).build().unwrap()) + .await + .unwrap(); + let records: DatasetIterator = resources.dataset.fetch(None).await.unwrap(); + drop(resources.dataset); + assert_eq!(records.collect().await.unwrap()[0].id, "record"); + resources.span.flush().await.unwrap(); +} + +#[tokio::test] +async fn dropping_the_last_resource_and_client_still_drains_an_accepted_write() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(204)) + .expect(1) + .mount(&server) + .await; + let client = BraintrustClient::builder() + .skip_login(true) + .app_url(server.uri()) + .api_url(server.uri()) + .build() + .await + .unwrap(); + let span: SpanHandle = client + .span_builder_with_credentials("token", "org") + .parent_info(ParentSpanInfo::ProjectLogs { + object_id: "project".into(), + }) + .build(); + let write = span.log(SpanLog::default()); + assert!(write.is_ok()); + drop(span); + drop(client); + tokio::time::timeout(std::time::Duration::from_secs(5), async { write.await }) + .await + .unwrap() + .unwrap(); +} diff --git a/tests/span_lifecycle.rs b/tests/span_lifecycle.rs index 0510018..b7e876e 100644 --- a/tests/span_lifecycle.rs +++ b/tests/span_lifecycle.rs @@ -820,3 +820,44 @@ fn usage_extractors_return_expected_metrics() { assert_eq!(anthropic_usage.completion_tokens(), Some(7)); assert_eq!(anthropic_usage.total_tokens(), Some(10)); } + +#[tokio::test] +async fn finished_span_ignores_later_mutations() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(204)) + .expect(1) + .mount(&server) + .await; + let client = BraintrustClient::builder() + .skip_login(true) + .api_url(server.uri()) + .app_url(server.uri()) + .build() + .await + .unwrap(); + let span = client + .span_builder_with_credentials("token", "org") + .parent_info(ParentSpanInfo::ProjectLogs { + object_id: "project".into(), + }) + .build(); + span.end_with_time(123.0); + let covering = client.flush(); + span.clone() + .log(SpanLog::builder().input(json!("too late")).build().unwrap()) + .await + .unwrap(); + covering.await.unwrap(); + client.flush().await.unwrap(); + let requests = server.received_requests().await.unwrap(); + let uploads: Vec<_> = requests + .iter() + .filter(|r| r.url.path() == "/logs3") + .collect(); + assert_eq!(uploads.len(), 1); + let body: Value = serde_json::from_slice(&uploads[0].body).unwrap(); + assert!(body["rows"][0].get("input").is_none()); + assert_eq!(body["rows"][0]["metrics"]["end"], 123.0); +}