From 82eb14467d2d3e4655bd10675ea976085bd1644a Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Tue, 20 Oct 2026 21:44:16 +0800 Subject: [PATCH] Replace buffered dataset spans with typed resource writes --- src/api/client.rs | 8 +- src/api/datasets.rs | 92 ++++++++- src/dataset.rs | 267 +++++++++++-------------- src/error.rs | 21 ++ src/experiments/api/mod.rs | 6 +- src/experiments/api/registration.rs | 5 +- src/internal/mod.rs | 1 + src/internal/resources.rs | 140 +++++++++++++ src/internal/upload.rs | 28 +-- src/logger.rs | 300 ++-------------------------- tests/api_client.rs | 47 +++++ tests/dataset_test.rs | 266 ++++++++++++++++++++++-- tests/experiments.rs | 48 +++++ 13 files changed, 746 insertions(+), 483 deletions(-) create mode 100644 src/internal/resources.rs diff --git a/src/api/client.rs b/src/api/client.rs index 1d088d8..aecb013 100644 --- a/src/api/client.rs +++ b/src/api/client.rs @@ -62,19 +62,21 @@ impl Client { }) } /// Reuse the SDK's configured transport without copying session state into - /// the API layer. This scope is immutable and used only for ingestion. - pub(crate) fn for_ingestion( + /// 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.clone(), + app_url, api_url, headers, }), diff --git a/src/api/datasets.rs b/src/api/datasets.rs index 7778efa..14185d3 100644 --- a/src/api/datasets.rs +++ b/src/api/datasets.rs @@ -1,4 +1,4 @@ -//! Dataset registration and summary endpoint contracts. +//! Dataset registration, row mutation, and summary endpoint contracts. use super::{Client, Error, Plane}; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; @@ -36,6 +36,76 @@ pub struct DatasetInfo { 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 { @@ -68,6 +138,26 @@ impl Datasets { .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( diff --git a/src/dataset.rs b/src/dataset.rs index 60bb161..a568369 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}; // ============================================================================ // 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 { @@ -252,7 +260,12 @@ pub(crate) trait DatasetRegistrar: Send + Sync { #[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; + async fn fetch_dataset_records( + &self, + token: &str, + org_name: &str, + query: BTQLQuery, + ) -> Result; } /// Trait for fetching dataset summary. @@ -262,10 +275,22 @@ pub(crate) trait DatasetSummarizer: Send + Sync { async fn fetch_dataset_summary( &self, token: &str, + org_name: &str, dataset_id: &str, ) -> Result; } +/// Direct dataset mutation port; writes are acknowledged by the endpoint. +#[async_trait] +pub(crate) trait DatasetWriter: Send + Sync { + async fn write_dataset_record( + &self, + token: &str, + org_name: &str, + row: DatasetWrite, + ) -> Result<()>; +} + // ============================================================================ // DatasetBuilder Error Types // ============================================================================ @@ -301,8 +326,9 @@ impl std::error::Error for DatasetBuilderError {} /// Builder for creating datasets with lazy registration. #[allow(private_bounds)] -pub struct DatasetBuilder -{ +pub struct DatasetBuilder< + S: SpanSubmitter + DatasetRegistrar + DatasetFetcher + DatasetSummarizer + DatasetWriter, +> { submitter: Arc, token: String, org_name: String, @@ -312,7 +338,7 @@ pub struct DatasetBuilder>, } -impl Clone +impl Clone for DatasetBuilder { fn clone(&self) -> Self { @@ -329,8 +355,14 @@ impl C } #[allow(private_bounds)] -impl - DatasetBuilder +impl< + S: SpanSubmitter + + DatasetRegistrar + + DatasetFetcher + + DatasetSummarizer + + DatasetWriter + + 'static, + > DatasetBuilder { /// Create a new DatasetBuilder (internal use). pub(crate) fn new( @@ -399,7 +431,6 @@ impl { +pub struct Dataset< + S: SpanSubmitter + DatasetRegistrar + DatasetFetcher + DatasetSummarizer + DatasetWriter, +> { submitter: Arc, token: String, org_name: String, @@ -421,10 +454,9 @@ pub struct Dataset, metadata: Option>, lazy_metadata: Arc>, - pending_rows: Arc>>, } -impl Clone +impl Clone for Dataset { fn clone(&self) -> Self { @@ -437,13 +469,12 @@ impl C description: self.description.clone(), metadata: self.metadata.clone(), lazy_metadata: self.lazy_metadata.clone(), - pending_rows: self.pending_rows.clone(), } } } -impl fmt::Debug - for Dataset +impl + fmt::Debug for Dataset { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("Dataset") @@ -454,20 +485,15 @@ 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< + S: SpanSubmitter + + DatasetRegistrar + + DatasetFetcher + + DatasetSummarizer + + DatasetWriter + + 'static, + > Dataset { /// Ensure the dataset is registered, returning cached metadata. async fn ensure_registered(&self) -> Result<&DatasetMetadata> { @@ -507,56 +533,53 @@ 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.submitter + .write_dataset_record( + &self.token, + &self.org_name, + 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.submitter + .write_dataset_record( + &self.token, + &self.org_name, + 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.submitter + .write_dataset_record( + &self.token, + &self.org_name, + DatasetWrite::Delete { + dataset_id, + id: id.to_owned(), + }, + ) + .await } /// Fetch records with pagination. @@ -567,6 +590,7 @@ impl 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) + .fetch_dataset_summary(&self.token, &self.org_name, &metadata.dataset_id) .await?; Ok(DatasetSummary { @@ -596,70 +617,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.submitter.flush() } } @@ -683,6 +644,7 @@ impl { fetcher: Arc, token: String, + org_name: String, dataset_id: String, batch_size: usize, cursor: Option, @@ -692,10 +654,17 @@ pub struct DatasetIterator { #[allow(private_bounds)] impl DatasetIterator { - fn new(fetcher: Arc, token: String, dataset_id: String, batch_size: usize) -> Self { + fn new( + fetcher: Arc, + token: String, + org_name: String, + dataset_id: String, + batch_size: usize, + ) -> Self { Self { fetcher, token, + org_name, dataset_id, batch_size, cursor: None, @@ -724,7 +693,7 @@ impl DatasetIterator { // A failed request leaves the cursor and EOF state untouched. let response = self .fetcher - .fetch_dataset_records(&self.token, query) + .fetch_dataset_records(&self.token, &self.org_name, query) .await?; if response.cursor.is_some() && response.cursor == previous_cursor { return Err(crate::BraintrustError::Background( diff --git a/src/error.rs b/src/error.rs index 779576c..7d2c366 100644 --- a/src/error.rs +++ b/src/error.rs @@ -51,3 +51,24 @@ impl From for BraintrustError { 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/experiments/api/mod.rs b/src/experiments/api/mod.rs index 92d5af2..14d106c 100644 --- a/src/experiments/api/mod.rs +++ b/src/experiments/api/mod.rs @@ -7,9 +7,5 @@ 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 registration::ExperimentRegisterRequest; pub(crate) use traits::{BaseExperimentFetcher, ExperimentComparisonFetcher, ExperimentRegistrar}; diff --git a/src/experiments/api/registration.rs b/src/experiments/api/registration.rs index 7a898fd..4dac4cb 100644 --- a/src/experiments/api/registration.rs +++ b/src/experiments/api/registration.rs @@ -1,5 +1,2 @@ //! Compatibility imports for existing experiment callers. -pub(crate) use crate::api::experiments::{ - BaseExperimentRequest, BaseExperimentResponse, ExperimentRegisterRequest, - ExperimentRegisterResponse, -}; +pub(crate) use crate::api::experiments::{ExperimentRegisterRequest, ExperimentRegisterResponse}; diff --git a/src/internal/mod.rs b/src/internal/mod.rs index a10d05f..3556ec5 100644 --- a/src/internal/mod.rs +++ b/src/internal/mod.rs @@ -2,5 +2,6 @@ pub(crate) mod batching; pub(crate) mod delivery; pub(crate) mod ingest; pub(crate) mod object_store; +pub(crate) mod resources; pub(crate) mod retry; pub(crate) mod upload; diff --git a/src/internal/resources.rs b/src/internal/resources.rs new file mode 100644 index 0000000..66c4aa2 --- /dev/null +++ b/src/internal/resources.rs @@ -0,0 +1,140 @@ +//! Resource ports adapt SDK policy to the shared one-attempt API contracts. +use async_trait::async_trait; + +use crate::api::{ + btql::{BTQLQuery, BTQLResponse}, + datasets::{ + DatasetRegisterRequest, DatasetRegisterResponse, DatasetSummaryResponse, DatasetWrite, + }, + experiments::{ + BaseExperimentRequest, ExperimentComparisonResponse, ExperimentRegisterRequest, + ExperimentRegisterResponse, + }, +}; +use crate::dataset::{DatasetFetcher, DatasetRegistrar, DatasetSummarizer, DatasetWriter}; +use crate::error::{BraintrustError, Result}; +use crate::experiments::api::{ + BaseExperimentFetcher, ExperimentComparisonFetcher, ExperimentRegistrar, +}; +use crate::experiments::BaseExperimentInfo; +use crate::BraintrustClient; + +#[async_trait] +impl ExperimentRegistrar for BraintrustClient { + async fn register_experiment( + &self, + token: &str, + request: ExperimentRegisterRequest, + ) -> Result { + self.resource_api(token, None)? + .experiments() + .register(request) + .await + .map_err(Into::into) + } +} + +#[async_trait] +impl BaseExperimentFetcher for BraintrustClient { + async fn fetch_base_experiment( + &self, + token: &str, + experiment_id: &str, + ) -> Result> { + let response = match self + .resource_api(token, None)? + .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_trait] +impl ExperimentComparisonFetcher for BraintrustClient { + async fn fetch_experiment_comparison( + &self, + token: &str, + experiment_id: &str, + base_experiment_id: Option<&str>, + ) -> Result { + self.resource_api(token, None)? + .experiments() + .compare(experiment_id, base_experiment_id) + .await + .map_err(Into::into) + } +} + +#[async_trait] +impl DatasetRegistrar for BraintrustClient { + async fn register_dataset( + &self, + token: &str, + request: DatasetRegisterRequest, + ) -> Result { + self.resource_api(token, Some(&request.org_name))? + .datasets() + .register(request) + .await + .map_err(Into::into) + } +} + +#[async_trait] +impl DatasetFetcher for BraintrustClient { + async fn fetch_dataset_records( + &self, + token: &str, + org_name: &str, + query: BTQLQuery, + ) -> Result { + self.resource_api(token, Some(org_name))? + .btql() + .query_dataset_records(query) + .await + .map_err(Into::into) + } +} + +#[async_trait] +impl DatasetSummarizer for BraintrustClient { + async fn fetch_dataset_summary( + &self, + token: &str, + org_name: &str, + dataset_id: &str, + ) -> Result { + self.resource_api(token, Some(org_name))? + .datasets() + .summarize(dataset_id) + .await + .map_err(Into::into) + } +} + +#[async_trait] +impl DatasetWriter for BraintrustClient { + async fn write_dataset_record( + &self, + token: &str, + org_name: &str, + row: DatasetWrite, + ) -> Result<()> { + self.resource_api(token, Some(org_name))? + .datasets() + .write(&[row]) + .await + .map_err(Into::into) + } +} diff --git a/src/internal/upload.rs b/src/internal/upload.rs index 3028a2f..03ba089 100644 --- a/src/internal/upload.rs +++ b/src/internal/upload.rs @@ -72,24 +72,6 @@ fn retryable(error: &BraintrustError) -> bool { _ => false, } } -pub(crate) fn api_error(error: api::Error) -> BraintrustError { - match error { - api::Error::Config(message) => BraintrustError::InvalidConfig(message), - api::Error::Transport { source, .. } => BraintrustError::Http(source), - api::Error::Http { status, body, .. } => BraintrustError::Api { - status, - message: body, - }, - api::Error::Decode { - operation, - status, - source, - } => BraintrustError::Api { - status, - message: format!("{operation}: invalid response: {source}"), - }, - } -} #[derive(Default)] pub(crate) struct UploadResult { @@ -111,8 +93,8 @@ impl UploadCoordinator { org_name: Option<&str>, policy: UploadPolicy, ) -> Result { - let logs = api::Client::for_ingestion(http.clone(), api_url, token, org_name) - .map_err(api_error)? + let logs = api::Client::for_scope(http.clone(), api_url, api_url, token, org_name) + .map_err(BraintrustError::from)? .logs3(); Ok(Self { logs, @@ -170,7 +152,7 @@ impl UploadCoordinator { self.logs .post_raw(batch.data.to_vec()) .await - .map_err(api_error) + .map_err(BraintrustError::from) }) .await { @@ -205,7 +187,7 @@ impl UploadCoordinator { rows: batch.overflow_rows.clone(), }) .await - .map_err(api_error) + .map_err(BraintrustError::from) }) .await?; budget @@ -222,7 +204,7 @@ impl UploadCoordinator { api_version: LOGS_API_VERSION, }) .await - .map_err(api_error) + .map_err(BraintrustError::from) }) .await } diff --git a/src/logger.rs b/src/logger.rs index 675c47e..9a0058d 100644 --- a/src/logger.rs +++ b/src/logger.rs @@ -7,17 +7,9 @@ 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::retry::{exponential, RetryExecutor}; use crate::log_queue::{LogQueue, LogQueueConfig}; @@ -457,6 +449,22 @@ impl std::fmt::Debug for ClientInner { } impl BraintrustClient { + /// Bind explicit credentials to the configured resource endpoints without login. + pub(crate) fn resource_api( + &self, + token: &str, + org_name: Option<&str>, + ) -> Result { + crate::api::Client::for_scope( + self.inner.http_client.clone(), + &self.inner.api_url, + &self.inner.app_url, + token, + org_name, + ) + .map_err(Into::into) + } + /// Create a new builder for configuring the client. /// /// # Example @@ -963,278 +971,6 @@ impl SpanSubmitter for BraintrustClient { } } -// 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 { diff --git a/tests/api_client.rs b/tests/api_client.rs index dfc05a1..450d79a 100644 --- a/tests/api_client.rs +++ b/tests/api_client.rs @@ -397,3 +397,50 @@ async fn base_path_is_preserved_without_a_trailing_slash() { .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/dataset_test.rs b/tests/dataset_test.rs index c335a80..f982236 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] @@ -523,3 +523,237 @@ async fn a_non_advancing_cursor_is_reported_and_does_not_poison_a_later_retry() 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/experiments.rs b/tests/experiments.rs index 6e406a2..9de77b4 100644 --- a/tests/experiments.rs +++ b/tests/experiments.rs @@ -520,3 +520,51 @@ async fn clones_share_the_selected_span_recording_interface() { 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()); +}