From 9fc7efea18a708c7b48fd7b6db8671dbfc5e1905 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Tue, 20 Oct 2026 21:33:02 +0800 Subject: [PATCH 1/7] Share resource registration state and propagate operation failures --- src/dataset.rs | 126 +++++++++------ src/error.rs | 14 ++ src/experiments/builder.rs | 14 +- src/experiments/experiment.rs | 187 ++++++++++----------- src/experiments/mod.rs | 6 +- src/lib.rs | 2 +- src/logger.rs | 58 ++++--- tests/dataset_test.rs | 240 ++++++++++++++++++++++++++- tests/eval_test.rs | 1 + tests/experiments.rs | 295 +++++++++++++++++++++++++++++++++- 10 files changed, 748 insertions(+), 195 deletions(-) diff --git a/src/dataset.rs b/src/dataset.rs index 7eee9ed..60bb161 100644 --- a/src/dataset.rs +++ b/src/dataset.rs @@ -398,7 +398,7 @@ impl, metadata: Option>, - lazy_metadata: OnceCell>, + lazy_metadata: Arc>, pending_rows: Arc>>, } +impl Clone + for Dataset +{ + fn clone(&self) -> Self { + Self { + submitter: self.submitter.clone(), + token: self.token.clone(), + org_name: self.org_name.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(), + pending_rows: self.pending_rows.clone(), + } + } +} + impl fmt::Debug for Dataset { @@ -454,7 +472,7 @@ impl 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(), @@ -462,18 +480,16 @@ impl Ok(DatasetMetadata { - dataset_id: response.dataset.id, - project_id: response.project.id, - }), - Err(e) => Err(e.to_string()), - } + let response = self + .submitter + .register_dataset(&self.token, 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). @@ -552,7 +568,7 @@ impl) -> Result<(), BraintrustError> { +/// let mut records = dataset.fetch(Some(100)).await?; +/// while let Some(record) = records.next().await? { +/// println!("{}", record.id); +/// } +/// # Ok(()) +/// # } +/// ``` #[allow(private_bounds)] pub struct DatasetIterator { fetcher: Arc, @@ -678,53 +705,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.token, 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/error.rs b/src/error.rs index 7444f6a..779576c 100644 --- a/src/error.rs +++ b/src/error.rs @@ -37,3 +37,17 @@ impl From for BraintrustError { 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()) + } +} diff --git a/src/experiments/builder.rs b/src/experiments/builder.rs index 32ffa69..e9b9fd6 100644 --- a/src/experiments/builder.rs +++ b/src/experiments/builder.rs @@ -1,13 +1,11 @@ //! 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; @@ -33,9 +31,8 @@ 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. +/// Registration is shared by clones of the resulting handle. Explicit +/// credentials belong to the resource and do not mutate the client's session. #[allow(private_bounds)] pub struct ExperimentBuilder { submitter: Arc, @@ -182,10 +179,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..687d3e7 100644 --- a/src/experiments/experiment.rs +++ b/src/experiments/experiment.rs @@ -6,8 +6,7 @@ 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::error::Result; use crate::experiments::api::{ @@ -49,13 +48,16 @@ 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 { @@ -73,10 +75,7 @@ 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(), } } } @@ -94,23 +93,20 @@ impl< /// /// 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 +116,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(); @@ -206,7 +204,9 @@ impl< }; self.submitter - .submit(self.token.clone(), payload, Some(parent_info)); + .submit(self.token.clone(), payload, Some(parent_info)) + .await?; + Ok(()) } /// Get a summary of the experiment's performance. @@ -217,15 +217,10 @@ 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 @@ -297,21 +292,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 +315,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 and final-flush errors are returned. Individual writes expose their + /// delivery results through the completion tokens returned inside the callback. /// /// # Example /// @@ -343,24 +337,24 @@ 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, 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 + span.end(); + span.flush().await?; + 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(), @@ -393,14 +387,16 @@ impl< }; self.submitter - .submit(self.token.clone(), payload, Some(parent_info)); + .submit(self.token.clone(), payload, Some(parent_info)) + .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; + let metadata = self.ensure_registered().await?; self.submitter .fetch_base_experiment(&self.token, &metadata.experiment_id) .await @@ -410,18 +406,19 @@ 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(), @@ -433,37 +430,21 @@ impl< update: self.update, repo_info: self.repo_info.clone(), }; - - match self + let response = 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()), - } - } - } + .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( + fn create_internal_span( &self, metadata: &ExperimentMetadata, start_time: f64, diff --git a/src/experiments/mod.rs b/src/experiments/mod.rs index d2e0bc4..870e762 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,8 +20,8 @@ //! .output(json!({"answer": "4"})) //! .expected(json!({"answer": "4"})) //! .score("accuracy", 1.0) -//! .build()? -//! ).await; +//! .build() +//! ).await?; //! //! experiment.flush().await?; //! ``` diff --git a/src/lib.rs b/src/lib.rs index eb51164..a3b9857 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -21,7 +21,7 @@ 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::{ diff --git a/src/logger.rs b/src/logger.rs index 493e23e..0f61d0a 100644 --- a/src/logger.rs +++ b/src/logger.rs @@ -42,7 +42,7 @@ pub struct LoginState { inner: Arc>, } -#[derive(Debug, Clone)] +#[derive(Clone)] struct LoginStateInner { api_key: String, org_id: String, @@ -119,13 +119,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() } } @@ -550,6 +548,9 @@ impl BraintrustClient { /// /// This waits for login to complete if it hasn't already. 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() @@ -564,6 +565,16 @@ impl BraintrustClient { )) } + /// 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(Arc::new(self.clone()), token, org_name) + } + /// Create a span builder with explicit token and org_id. /// /// Use this for per-session / multi-tenant logging: the supplied credentials @@ -597,7 +608,7 @@ 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? @@ -893,8 +904,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 /// @@ -913,15 +923,6 @@ impl BraintrustClient { 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) } @@ -1404,6 +1405,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; diff --git a/tests/dataset_test.rs b/tests/dataset_test.rs index b4ccfd6..f5358af 100644 --- a/tests/dataset_test.rs +++ b/tests/dataset_test.rs @@ -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,241 @@ 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()); +} 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..6e406a2 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,285 @@ 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(_)) + )); +} From e89fbefb60120ff1e7bb49f98fef8007bf1c3995 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Tue, 20 Oct 2026 21:44:16 +0800 Subject: [PATCH 2/7] 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 bae2dc7..a9bd25c 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}; @@ -70,6 +70,76 @@ pub struct DatasetInfo { pub(crate) 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)] #[non_exhaustive] @@ -104,6 +174,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 741019b..9a0ce3c 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 0f61d0a..5479123 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 @@ -979,278 +987,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 1b1deee..f432244 100644 --- a/tests/api_client.rs +++ b/tests/api_client.rs @@ -458,3 +458,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 f5358af..567040a 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()); +} From 86c6c9bcdd56bffa43b93443894eb2f917470f1b Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Tue, 20 Oct 2026 21:53:53 +0800 Subject: [PATCH 3/7] Capture immutable credential scopes for resource operations --- src/dataset.rs | 74 ++---- src/experiments/api/traits.rs | 7 +- src/experiments/builder.rs | 25 +- src/experiments/experiment.rs | 45 ++-- src/experiments/span_builder.rs | 9 +- src/internal/mod.rs | 1 + src/internal/resources.rs | 36 ++- src/internal/session.rs | 81 +++++++ src/log_queue/queue.rs | 303 +++++++++--------------- src/logger.rs | 124 +++++----- src/span.rs | 42 ++-- src/test_utils.rs | 7 +- tests/credential_scopes.rs | 404 ++++++++++++++++++++++++++++++++ 13 files changed, 743 insertions(+), 415 deletions(-) create mode 100644 src/internal/session.rs create mode 100644 tests/credential_scopes.rs diff --git a/src/dataset.rs b/src/dataset.rs index a568369..9d80ef5 100644 --- a/src/dataset.rs +++ b/src/dataset.rs @@ -35,6 +35,7 @@ use uuid::Uuid; use crate::api::datasets::{DatasetFields, DatasetWrite}; use crate::error::Result; +use crate::internal::session::CredentialScope; use crate::span::SpanSubmitter; // ============================================================================ @@ -251,7 +252,7 @@ pub(crate) trait DatasetRegistrar: Send + Sync { /// Register a dataset with the API. async fn register_dataset( &self, - token: &str, + scope: &CredentialScope, request: DatasetRegisterRequest, ) -> Result; } @@ -262,8 +263,7 @@ pub(crate) trait DatasetFetcher: Send + Sync { /// Fetch dataset records via BTQL query. async fn fetch_dataset_records( &self, - token: &str, - org_name: &str, + scope: &CredentialScope, query: BTQLQuery, ) -> Result; } @@ -274,8 +274,7 @@ pub(crate) trait DatasetSummarizer: Send + Sync { /// Fetch dataset summary. async fn fetch_dataset_summary( &self, - token: &str, - org_name: &str, + scope: &CredentialScope, dataset_id: &str, ) -> Result; } @@ -283,12 +282,7 @@ pub(crate) trait DatasetSummarizer: Send + Sync { /// 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<()>; + async fn write_dataset_record(&self, scope: &CredentialScope, row: DatasetWrite) -> Result<()>; } // ============================================================================ @@ -330,8 +324,7 @@ pub struct DatasetBuilder< S: SpanSubmitter + DatasetRegistrar + DatasetFetcher + DatasetSummarizer + DatasetWriter, > { submitter: Arc, - token: String, - org_name: String, + scope: CredentialScope, project_name: Option, dataset_name: Option, description: Option, @@ -344,8 +337,7 @@ impl Self { Self { submitter: Arc::clone(&self.submitter), - token: self.token.clone(), - org_name: self.org_name.clone(), + scope: self.scope.clone(), project_name: self.project_name.clone(), dataset_name: self.dataset_name.clone(), description: self.description.clone(), @@ -365,15 +357,10 @@ impl< > DatasetBuilder { /// Create a new DatasetBuilder (internal use). - pub(crate) fn new( - submitter: Arc, - token: impl Into, - org_name: impl Into, - ) -> Self { + pub(crate) fn new(submitter: Arc, scope: CredentialScope) -> Self { Self { submitter, - token: token.into(), - org_name: org_name.into(), + scope, project_name: None, dataset_name: None, description: None, @@ -424,8 +411,7 @@ impl< Ok(Dataset { submitter: self.submitter, - token: self.token, - org_name: self.org_name, + scope: self.scope, project_name, dataset_name, description: self.description, @@ -447,8 +433,7 @@ pub struct Dataset< S: SpanSubmitter + DatasetRegistrar + DatasetFetcher + DatasetSummarizer + DatasetWriter, > { submitter: Arc, - token: String, - org_name: String, + scope: CredentialScope, project_name: String, dataset_name: String, description: Option, @@ -462,8 +447,7 @@ impl Self { Self { submitter: self.submitter.clone(), - token: self.token.clone(), - org_name: self.org_name.clone(), + scope: self.scope.clone(), project_name: self.project_name.clone(), dataset_name: self.dataset_name.clone(), description: self.description.clone(), @@ -501,14 +485,14 @@ impl< .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(), }; let response = self .submitter - .register_dataset(&self.token, request) + .register_dataset(&self.scope, request) .await?; Ok(DatasetMetadata { dataset_id: response.dataset.id, @@ -539,8 +523,7 @@ impl< let id = Uuid::new_v4().to_string(); self.submitter .write_dataset_record( - &self.token, - &self.org_name, + &self.scope, DatasetWrite::Insert { dataset_id, id: id.clone(), @@ -556,8 +539,7 @@ impl< let dataset_id = self.ensure_registered().await?.dataset_id.clone(); self.submitter .write_dataset_record( - &self.token, - &self.org_name, + &self.scope, DatasetWrite::Update { dataset_id, id: id.to_owned(), @@ -572,8 +554,7 @@ impl< let dataset_id = self.ensure_registered().await?.dataset_id.clone(); self.submitter .write_dataset_record( - &self.token, - &self.org_name, + &self.scope, DatasetWrite::Delete { dataset_id, id: id.to_owned(), @@ -589,8 +570,7 @@ impl< let metadata = self.ensure_registered().await?; Ok(DatasetIterator::new( Arc::clone(&self.submitter), - self.token.clone(), - self.org_name.clone(), + self.scope.clone(), metadata.dataset_id.clone(), batch_size.unwrap_or(1000).max(1), )) @@ -603,7 +583,7 @@ impl< // Fetch summary from API let response = self .submitter - .fetch_dataset_summary(&self.token, &self.org_name, &metadata.dataset_id) + .fetch_dataset_summary(&self.scope, &metadata.dataset_id) .await?; Ok(DatasetSummary { @@ -643,8 +623,7 @@ impl< #[allow(private_bounds)] pub struct DatasetIterator { fetcher: Arc, - token: String, - org_name: String, + scope: CredentialScope, dataset_id: String, batch_size: usize, cursor: Option, @@ -654,17 +633,10 @@ pub struct DatasetIterator { #[allow(private_bounds)] impl DatasetIterator { - fn new( - fetcher: Arc, - token: String, - org_name: String, - dataset_id: String, - batch_size: usize, - ) -> Self { + fn new(fetcher: Arc, scope: CredentialScope, dataset_id: String, batch_size: usize) -> Self { Self { fetcher, - token, - org_name, + scope, dataset_id, batch_size, cursor: None, @@ -693,7 +665,7 @@ impl DatasetIterator { // A failed request leaves the cursor and EOF state untouched. let response = self .fetcher - .fetch_dataset_records(&self.token, &self.org_name, query) + .fetch_dataset_records(&self.scope, query) .await?; if response.cursor.is_some() && response.cursor == previous_cursor { return Err(crate::BraintrustError::Background( diff --git a/src/experiments/api/traits.rs b/src/experiments/api/traits.rs index 54704f2..e14bd8c 100644 --- a/src/experiments/api/traits.rs +++ b/src/experiments/api/traits.rs @@ -6,13 +6,14 @@ use crate::experiments::api::registration::{ ExperimentRegisterRequest, ExperimentRegisterResponse, }; use crate::experiments::metadata::BaseExperimentInfo; +use crate::internal::session::CredentialScope; /// 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, + scope: &CredentialScope, request: ExperimentRegisterRequest, ) -> Result; } @@ -22,7 +23,7 @@ pub(crate) trait ExperimentRegistrar: Send + Sync { pub(crate) trait BaseExperimentFetcher: Send + Sync { async fn fetch_base_experiment( &self, - token: &str, + scope: &CredentialScope, experiment_id: &str, ) -> Result>; } @@ -32,7 +33,7 @@ pub(crate) trait BaseExperimentFetcher: Send + Sync { pub(crate) trait ExperimentComparisonFetcher: Send + Sync { async fn fetch_experiment_comparison( &self, - token: &str, + scope: &CredentialScope, experiment_id: &str, base_experiment_id: Option<&str>, ) -> Result; diff --git a/src/experiments/builder.rs b/src/experiments/builder.rs index e9b9fd6..9ba3376 100644 --- a/src/experiments/builder.rs +++ b/src/experiments/builder.rs @@ -7,6 +7,7 @@ use serde_json::{Map, Value}; use crate::experiments::experiment::{Experiment, ExperimentState}; use crate::experiments::metadata::RepoInfo; +use crate::internal::session::CredentialScope; use crate::span::SpanSubmitter; /// Error type for ExperimentBuilder validation. @@ -36,9 +37,7 @@ impl std::error::Error for ExperimentBuilderError {} #[allow(private_bounds)] pub struct ExperimentBuilder { submitter: Arc, - token: String, - org_id: String, - org_name: Option, + scope: CredentialScope, project_name: Option, experiment_name: Option, description: Option, @@ -53,9 +52,7 @@ 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(), + scope: self.scope.clone(), project_name: self.project_name.clone(), experiment_name: self.experiment_name.clone(), description: self.description.clone(), @@ -71,16 +68,10 @@ impl Clone for ExperimentBuilder { #[allow(private_bounds)] impl ExperimentBuilder { /// Create a new ExperimentBuilder (internal use). - pub(crate) fn new( - submitter: Arc, - token: impl Into, - org_id: impl Into, - ) -> Self { + 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, experiment_name: None, description: None, @@ -94,7 +85,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 } @@ -168,9 +159,7 @@ impl ExperimentBuilder { Ok(Experiment { 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, diff --git a/src/experiments/experiment.rs b/src/experiments/experiment.rs index 687d3e7..e8eb515 100644 --- a/src/experiments/experiment.rs +++ b/src/experiments/experiment.rs @@ -20,6 +20,7 @@ use crate::experiments::metadata::{ }; use crate::experiments::span_builder::ExperimentSpanBuilder; use crate::experiments::summary::{ExperimentSummary, MetricSummary, ScoreSummary}; +use crate::internal::session::CredentialScope; use crate::span::{SpanBuilder, SpanHandle, SpanLog, SpanSubmitter}; use crate::types::{ParentSpanInfo, SpanAttributes, SpanPayload, SpanType}; @@ -37,9 +38,7 @@ pub(crate) struct ExperimentMetadata { #[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(crate) scope: CredentialScope, pub(crate) project_name: String, pub(crate) experiment_name: Option, pub(crate) description: Option, @@ -64,9 +63,7 @@ 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(), + scope: self.scope.clone(), project_name: self.project_name.clone(), experiment_name: self.experiment_name.clone(), description: self.description.clone(), @@ -183,8 +180,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, @@ -204,7 +201,7 @@ impl< }; self.submitter - .submit(self.token.clone(), payload, Some(parent_info)) + .submit(self.scope.clone(), payload, Some(parent_info)) .await?; Ok(()) } @@ -226,7 +223,7 @@ impl< let comparison = self .submitter .fetch_experiment_comparison( - &self.token, + &self.scope, &metadata.experiment_id, base_exp_id.as_deref(), ) @@ -361,8 +358,8 @@ impl< 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, @@ -387,7 +384,7 @@ impl< }; self.submitter - .submit(self.token.clone(), payload, Some(parent_info)) + .submit(self.scope.clone(), payload, Some(parent_info)) .await?; Ok(()) } @@ -398,7 +395,7 @@ impl< pub async fn fetch_base_experiment(&self) -> Result> { let metadata = self.ensure_registered().await?; self.submitter - .fetch_base_experiment(&self.token, &metadata.experiment_id) + .fetch_base_experiment(&self.scope, &metadata.experiment_id) .await } @@ -421,7 +418,7 @@ impl< .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(), @@ -432,7 +429,7 @@ impl< }; let response = self .submitter - .register_experiment(&self.token, request) + .register_experiment(&self.scope, request) .await?; Ok(ExperimentMetadata { project_id: response.project.id, @@ -455,18 +452,10 @@ impl< // 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/span_builder.rs b/src/experiments/span_builder.rs index 4f519e5..0efb335 100644 --- a/src/experiments/span_builder.rs +++ b/src/experiments/span_builder.rs @@ -46,18 +46,13 @@ impl<'a, S: SpanSubmitter + 'static> ExperimentSpanBuilder<'a, S> { // 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( + let builder = SpanBuilder::new( Arc::clone(&self.experiment.submitter), - self.experiment.token.clone(), - self.experiment.org_id.clone(), + self.experiment.scope.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()); - } - builder.build() } } diff --git a/src/internal/mod.rs b/src/internal/mod.rs index 3556ec5..7eb3eba 100644 --- a/src/internal/mod.rs +++ b/src/internal/mod.rs @@ -4,4 +4,5 @@ pub(crate) mod ingest; pub(crate) mod object_store; pub(crate) mod resources; pub(crate) mod retry; +pub(crate) mod session; pub(crate) mod upload; diff --git a/src/internal/resources.rs b/src/internal/resources.rs index 66c4aa2..4d74ee2 100644 --- a/src/internal/resources.rs +++ b/src/internal/resources.rs @@ -1,4 +1,5 @@ //! 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::{ @@ -23,10 +24,10 @@ use crate::BraintrustClient; impl ExperimentRegistrar for BraintrustClient { async fn register_experiment( &self, - token: &str, + scope: &CredentialScope, request: ExperimentRegisterRequest, ) -> Result { - self.resource_api(token, None)? + self.resource_api(scope)? .experiments() .register(request) .await @@ -38,11 +39,11 @@ impl ExperimentRegistrar for BraintrustClient { impl BaseExperimentFetcher for BraintrustClient { async fn fetch_base_experiment( &self, - token: &str, + scope: &CredentialScope, experiment_id: &str, ) -> Result> { let response = match self - .resource_api(token, None)? + .resource_api(scope)? .experiments() .base(BaseExperimentRequest { id: experiment_id.to_owned(), @@ -64,11 +65,11 @@ impl BaseExperimentFetcher for BraintrustClient { impl ExperimentComparisonFetcher for BraintrustClient { async fn fetch_experiment_comparison( &self, - token: &str, + scope: &CredentialScope, experiment_id: &str, base_experiment_id: Option<&str>, ) -> Result { - self.resource_api(token, None)? + self.resource_api(scope)? .experiments() .compare(experiment_id, base_experiment_id) .await @@ -80,10 +81,10 @@ impl ExperimentComparisonFetcher for BraintrustClient { impl DatasetRegistrar for BraintrustClient { async fn register_dataset( &self, - token: &str, + scope: &CredentialScope, request: DatasetRegisterRequest, ) -> Result { - self.resource_api(token, Some(&request.org_name))? + self.resource_api(scope)? .datasets() .register(request) .await @@ -95,11 +96,10 @@ impl DatasetRegistrar for BraintrustClient { impl DatasetFetcher for BraintrustClient { async fn fetch_dataset_records( &self, - token: &str, - org_name: &str, + scope: &CredentialScope, query: BTQLQuery, ) -> Result { - self.resource_api(token, Some(org_name))? + self.resource_api(scope)? .btql() .query_dataset_records(query) .await @@ -111,11 +111,10 @@ impl DatasetFetcher for BraintrustClient { impl DatasetSummarizer for BraintrustClient { async fn fetch_dataset_summary( &self, - token: &str, - org_name: &str, + scope: &CredentialScope, dataset_id: &str, ) -> Result { - self.resource_api(token, Some(org_name))? + self.resource_api(scope)? .datasets() .summarize(dataset_id) .await @@ -125,13 +124,8 @@ impl DatasetSummarizer for BraintrustClient { #[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))? + async fn write_dataset_record(&self, scope: &CredentialScope, row: DatasetWrite) -> Result<()> { + self.resource_api(scope)? .datasets() .write(&[row]) .await 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/log_queue/queue.rs b/src/log_queue/queue.rs index de69ff0..92b2bb1 100644 --- a/src/log_queue/queue.rs +++ b/src/log_queue/queue.rs @@ -1,4 +1,3 @@ -use anyhow::Context; use chrono::Utc; use super::batching::batch_and_serialize_rows; @@ -8,7 +7,7 @@ use super::merge::merge_row_into; use super::row_key::RowKey; use crate::error::{BraintrustError, Result}; use crate::internal::delivery::WriteService; -use crate::logger::LoginState; +use crate::internal::session::CredentialScope; use crate::span_components::{project_logs_identifier, ProjectLogsIdentifier}; use crate::types::{LogDestination, Logs3Row, ParentSpanInfo, SpanObjectType, SpanPayload}; use crate::{ @@ -19,18 +18,27 @@ use indexmap::IndexMap; use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; -use tokio::sync::{Mutex as TokioMutex, OnceCell}; +use tokio::sync::Mutex as TokioMutex; use tracing::warn; -use url::Url; -use crate::api::projects::{ProjectRegisterRequest, ProjectRegisterResponse}; +use crate::api::projects::ProjectRegisterRequest; + +// 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. /// /// Only the span batch processor interprets these fields; the delivery worker /// treats commands as opaque values. struct SubmitCommand { - token: String, + scope: CredentialScope, payload: SpanPayload, parent_info: Option, } @@ -38,46 +46,34 @@ struct SubmitCommand { /// Preparation and protocol orchestration, independent of admission/lifecycle. struct SpanBatchProcessor { config: LogQueueConfig, - login_state: LoginState, client: reqwest::Client, - api_url: Url, - app_url: Url, - project_cache: TokioMutex>, - version_info: OnceCell<(usize, bool)>, + project_cache: TokioMutex>, + version_info: TokioMutex>, } impl SpanBatchProcessor { - fn new( - config: LogQueueConfig, - login_state: LoginState, - client: reqwest::Client, - api_url: Url, - app_url: Url, - ) -> Arc { + fn new(config: LogQueueConfig, client: reqwest::Client) -> Arc { Arc::new(Self { config, - login_state, client, - api_url, - app_url, project_cache: TokioMutex::new(IndexMap::new()), - version_info: OnceCell::new(), + version_info: TokioMutex::new(IndexMap::new()), }) } - /// 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 + /// 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; + } + 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 } /// Process one detached logical batch: prepare → merge → encode → upload. @@ -95,43 +91,18 @@ impl SpanBatchProcessor { } 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 Err(e.into()); - } - }; - - // 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"); @@ -146,21 +117,6 @@ impl SpanBatchProcessor { 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, @@ -170,7 +126,10 @@ impl SpanBatchProcessor { ), ); - 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(); @@ -220,9 +179,9 @@ impl SpanBatchProcessor { .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, @@ -251,7 +210,7 @@ impl SpanBatchProcessor { /// 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 { @@ -260,8 +219,8 @@ impl SpanBatchProcessor { span_id, is_merge, span_components, - org_id, - org_name, + org_id: _, + org_name: _, project_name, input, output, @@ -276,14 +235,12 @@ impl SpanBatchProcessor { 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 @@ -319,10 +276,7 @@ impl SpanBatchProcessor { } 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 }; @@ -346,9 +300,7 @@ impl SpanBatchProcessor { 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 }) => ( @@ -384,14 +336,8 @@ impl SpanBatchProcessor { 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 => { @@ -441,9 +387,7 @@ impl SpanBatchProcessor { 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 { @@ -457,9 +401,7 @@ impl SpanBatchProcessor { 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!( @@ -474,71 +416,32 @@ impl SpanBatchProcessor { } } - /// 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: project_name.to_owned(), - org_id: (!org_id.is_empty()).then(|| org_id.to_owned()), - org_name: org_name.map(str::to_owned), - }; - - 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) + .map_err(BraintrustError::from)?; + cache_insert( + &mut *self.project_cache.lock().await, + cache_key, + response.project.id.clone(), + ); + Ok(response.project.id) } } @@ -595,7 +498,7 @@ impl RejectionReporter { }; handle.spawn(async move { let SubmitCommand { - token: _, + scope: _, payload, parent_info, } = cmd; @@ -723,21 +626,14 @@ pub struct LogQueue { writes: WriteService, } impl LogQueue { - pub fn new( - config: LogQueueConfig, - login_state: LoginState, - client: reqwest::Client, - _worker_queue_size: usize, - api_url: Url, - app_url: Url, - ) -> Self { + pub fn new(config: LogQueueConfig, client: reqwest::Client, _worker_queue_size: usize) -> Self { let capacity = config.queue_max_size(); let rejections = Arc::new(RejectionReporter { dropped_count: AtomicUsize::new(0), last_drop_log_time: AtomicU64::new(0), config: config.clone(), }); - let processor = SpanBatchProcessor::new(config, login_state, client, api_url, app_url); + let processor = SpanBatchProcessor::new(config, client); let writes = WriteService::new( capacity, tokio::runtime::Handle::current(), @@ -751,12 +647,12 @@ impl LogQueue { pub fn submit( &self, - token: impl Into, + scope: CredentialScope, payload: SpanPayload, parent_info: Option, ) -> Completion { let command = SubmitCommand { - token: token.into(), + scope, payload, parent_info, }; @@ -800,6 +696,7 @@ impl LogQueue { mod tests { use super::*; use serde_json::{json, Value}; + use url::Url; use wiremock::{ matchers::{header, method, path}, Mock, MockServer, ResponseTemplate, @@ -827,12 +724,18 @@ mod tests { extra: std::collections::HashMap::new(), } } - fn queue(server: &MockServer) -> LogQueue { + fn queue(_server: &MockServer) -> LogQueue { LogQueue::new( LogQueueConfig::builder().num_retries(0).build(), - LoginState::new(), reqwest::Client::new(), 32, + ) + } + 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(), ) @@ -849,7 +752,7 @@ mod tests { Mock::given(method("GET")) .and(path("/version")) .respond_with(ResponseTemplate::new(200).set_body_json(json!({}))) - .expect(1) + .expect(2) .mount(&server) .await; Mock::given(method("POST")) @@ -867,9 +770,9 @@ mod tests { .mount(&server) .await; let queue = queue(&server); - let first = queue.submit("good", payload("merged", false), parent()); - let update = queue.submit("good", payload("merged", true), parent()); - let last = queue.submit("bad", payload("other", false), parent()); + let first = queue.submit(scope(&server, "good"), payload("merged", false), parent()); + let update = queue.submit(scope(&server, "good"), payload("merged", true), parent()); + let last = queue.submit(scope(&server, "bad"), payload("other", false), parent()); let Err(FlushError::Delivery(failure)) = queue.flush().await else { panic!("expected failed flush"); }; @@ -894,7 +797,7 @@ mod tests { 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("good", payload("fresh", false), parent()); + let fresh = queue.submit(scope(&server, "good"), payload("fresh", false), parent()); assert_eq!(queue.flush().await.unwrap().delivered, 1); fresh.await.unwrap(); queue @@ -928,8 +831,16 @@ mod tests { .mount(&server) .await; let queue = queue(&server); - let first = queue.submit("token", payload("committed", false), parent()); - let last = queue.submit("token", payload("late-failure", false), parent()); + let first = queue.submit( + scope(&server, "token"), + payload("committed", false), + parent(), + ); + let last = queue.submit( + scope(&server, "token"), + payload("late-failure", false), + parent(), + ); let Err(FlushError::Delivery(failure)) = queue.flush().await else { panic!("expected failure"); }; diff --git a/src/logger.rs b/src/logger.rs index 5479123..4284874 100644 --- a/src/logger.rs +++ b/src/logger.rs @@ -12,6 +12,7 @@ use crate::error::{BraintrustError, Result}; use crate::experiments::ExperimentBuilder; use crate::http::{build_http_client, CaBundleConfig}; 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_components::SpanComponents; @@ -78,6 +79,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() @@ -365,17 +385,8 @@ impl BraintrustClientBuilder { .checked_delivery(self.checked_delivery) .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 { @@ -449,20 +460,23 @@ impl std::fmt::Debug for ClientInner { } impl BraintrustClient { - /// Bind explicit credentials to the configured resource endpoints without login. - pub(crate) fn resource_api( + pub(crate) fn resource_api(&self, scope: &CredentialScope) -> Result { + scope.api(&self.inner.http_client) + } + + fn explicit_scope( &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: 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(), ) - .map_err(Into::into) } /// Create a new builder for configuring the client. @@ -533,13 +547,7 @@ impl BraintrustClient { )); } 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(Arc::new(self.clone()), state.scope()?); if let Some(environment) = &self.inner.environment { builder = builder.span_origin_environment(environment.clone()); } @@ -560,17 +568,7 @@ impl BraintrustClient { 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, - )) + Ok(DatasetBuilder::new(Arc::new(self.clone()), state.scope()?)) } /// Create a dataset builder with immutable explicit credentials, without @@ -580,14 +578,17 @@ impl BraintrustClient { token: impl Into, org_name: impl Into, ) -> DatasetBuilder { - DatasetBuilder::new(Arc::new(self.clone()), token, org_name) + DatasetBuilder::new( + Arc::new(self.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, @@ -598,7 +599,8 @@ impl BraintrustClient { let org_id = org_id.into(); let submitter = Arc::new(self.clone()); - let mut builder = crate::span::SpanBuilder::new(submitter, token, org_id); + 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()); } @@ -622,13 +624,7 @@ impl BraintrustClient { 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) .map(|_| ()) } @@ -647,7 +643,7 @@ 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) .map(|_| ()) } @@ -662,14 +658,16 @@ impl BraintrustClient { exported: &str, event: crate::span::SpanLog, ) -> Result> { - self.update_span_internal(token.into(), org_id.into(), None, exported, event) + self.update_span_internal( + self.explicit_scope(token.into(), org_id.into(), 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> { @@ -749,8 +747,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, @@ -765,7 +763,7 @@ impl BraintrustClient { extra: event_data.extra, }; - let completion = self.submit_payload(token, payload, None); + let completion = self.submit_payload(scope, payload, None); if let Err(error) = completion.submission() { return Err(BraintrustError::Background(error.to_string())); } @@ -932,17 +930,17 @@ impl BraintrustClient { let org_id = org_id.into(); let submitter = Arc::new(self.clone()); - ExperimentBuilder::new(submitter, token, org_id) + ExperimentBuilder::new(submitter, self.explicit_scope(token, org_id, None)) } /// 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, ) -> crate::Completion { - self.inner.queue.submit(token, payload, parent_info) + self.inner.queue.submit(scope, payload, parent_info) } /// Stop accepting writes, drain accepted work and stop the delivery worker. @@ -971,11 +969,11 @@ impl BraintrustClient { impl SpanSubmitter for BraintrustClient { fn submit( &self, - token: String, + scope: CredentialScope, payload: SpanPayload, parent_info: Option, ) -> crate::Completion { - self.submit_payload(token, payload, parent_info) + self.submit_payload(scope, payload, parent_info) } fn flush(&self) -> crate::Completion { diff --git a/src/span.rs b/src/span.rs index 7fad7f9..ae7417e 100644 --- a/src/span.rs +++ b/src/span.rs @@ -10,6 +10,7 @@ use serde_json::{json, Map, Value}; use uuid::Uuid; use crate::error::Result; +use crate::internal::session::CredentialScope; use crate::span_components::SpanComponents; use crate::types::{ ParentSpanInfo, SpanAttributes, SpanEventData, SpanObjectType, SpanPayload, SpanType, @@ -252,7 +253,7 @@ pub(crate) trait SpanSubmitter: Send + Sync { /// Attempt admission immediately and observe the logical batch outcome. fn submit( &self, - token: String, + scope: CredentialScope, payload: SpanPayload, parent_info: Option, ) -> crate::Completion; @@ -268,9 +269,7 @@ pub(crate) trait SpanSubmitter: Send + Sync { #[allow(private_bounds)] pub struct SpanBuilder { submitter: Arc, - token: String, - org_id: String, - org_name: Option, + scope: CredentialScope, project_name: Option, parent_info: Option, span_type: SpanType, @@ -287,9 +286,7 @@ 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, @@ -374,16 +371,10 @@ impl SpanOrigin { #[allow(private_bounds)] impl SpanBuilder { - pub(crate) fn new( - submitter: Arc, - token: impl Into, - org_id: impl Into, - ) -> Self { + 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(), @@ -450,7 +441,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 } @@ -521,15 +512,15 @@ 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, @@ -547,7 +538,7 @@ impl SpanBuilder { #[allow(private_bounds)] pub struct SpanHandle { submitter: Arc, - token: String, + scope: CredentialScope, parent_info: Option, /// Stable identifier set at creation, accessible without locking. row_id: String, @@ -558,7 +549,7 @@ 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), @@ -589,7 +580,7 @@ impl SpanHandle { let payload = current_span_payload(&mut inner); let completion = self.submitter - .submit(self.token.clone(), payload, self.parent_info.clone()); + .submit(self.scope.clone(), payload, self.parent_info.clone()); if completion.is_ok() { inner.has_flushed = true; } @@ -629,7 +620,7 @@ impl SpanHandle { let payload = current_span_payload(&mut inner); let completion = self.submitter - .submit(self.token.clone(), payload, self.parent_info.clone()); + .submit(self.scope.clone(), payload, self.parent_info.clone()); if completion.is_ok() { inner.has_flushed = true; inner.end_submitted = true; @@ -1081,7 +1072,7 @@ mod tests { impl SpanSubmitter for Submitter { fn submit( &self, - _: String, + _: CredentialScope, payload: SpanPayload, _: Option, ) -> crate::Completion { @@ -1107,7 +1098,8 @@ mod tests { reject: AtomicBool::new(true), accepted: Mutex::new(Vec::new()), }); - let span = SpanBuilder::new(submitter.clone(), "token", "org").build(); + 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()); diff --git a/src/test_utils.rs b/src/test_utils.rs index ec1dcd5..4489b43 100644 --- a/src/test_utils.rs +++ b/src/test_utils.rs @@ -3,6 +3,7 @@ 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}; @@ -60,12 +61,12 @@ impl MockBraintrustClient { impl SpanSubmitter for MockBraintrustClient { fn submit( &self, - token: String, + scope: CredentialScope, payload: SpanPayload, parent_info: Option, ) -> crate::Completion { self.collector.push(CapturedSpan { - token, + token: scope.token().to_owned(), payload, parent: parent_info, }); @@ -88,7 +89,7 @@ impl SpanSubmitter for MockBraintrustClient { 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) } 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 + ); +} From ef95ff9925fd3a871f21982c89db1e268e55a787 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Tue, 20 Oct 2026 22:02:39 +0800 Subject: [PATCH 4/7] Expose concrete resource handles behind private I/O ports --- src/completion.rs | 2 +- src/dataset.rs | 162 +++++++++------------------- src/eval/runner.rs | 2 +- src/eval/types.rs | 6 +- src/experiments/api/comparison.rs | 1 - src/experiments/api/mod.rs | 11 -- src/experiments/api/registration.rs | 2 - src/experiments/api/traits.rs | 40 ------- src/experiments/builder.rs | 26 +++-- src/experiments/experiment.rs | 48 +++------ src/experiments/mod.rs | 3 - src/experiments/span_builder.rs | 28 +++-- src/internal/mod.rs | 1 + src/internal/ports.rs | 73 +++++++++++++ src/internal/resources.rs | 39 +++---- src/lib.rs | 2 + src/log_queue/queue.rs | 22 ++++ src/logger.rs | 89 +++++++++------ src/span.rs | 57 +++++----- src/stream.rs | 21 ++-- src/test_utils.rs | 9 +- tests/dataset_test.rs | 26 ++--- tests/resource_handles.rs | 130 ++++++++++++++++++++++ 23 files changed, 446 insertions(+), 354 deletions(-) delete mode 100644 src/experiments/api/comparison.rs delete mode 100644 src/experiments/api/mod.rs delete mode 100644 src/experiments/api/registration.rs delete mode 100644 src/experiments/api/traits.rs create mode 100644 src/internal/ports.rs create mode 100644 tests/resource_handles.rs diff --git a/src/completion.rs b/src/completion.rs index b9e6e81..96f7d55 100644 --- a/src/completion.rs +++ b/src/completion.rs @@ -18,7 +18,7 @@ use tokio::sync::oneshot; /// use braintrust_sdk_rust::{BraintrustClient, SpanHandle, SpanLog}; /// use std::time::Duration; /// -/// # async fn example(client: &BraintrustClient, span: &SpanHandle) +/// # async fn example(client: &BraintrustClient, span: &SpanHandle) /// # -> Result<(), Box> { /// let completion = span.log(SpanLog::default()); // admission happens now /// if let Err(error) = completion.submission() { diff --git a/src/dataset.rs b/src/dataset.rs index 9d80ef5..74fde64 100644 --- a/src/dataset.rs +++ b/src/dataset.rs @@ -27,7 +27,6 @@ use std::collections::VecDeque; use std::fmt; use std::sync::Arc; -use async_trait::async_trait; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use tokio::sync::OnceCell; @@ -35,8 +34,8 @@ use uuid::Uuid; use crate::api::datasets::{DatasetFields, DatasetWrite}; use crate::error::Result; +use crate::internal::ports::{DatasetApi, DeliveryControl}; use crate::internal::session::CredentialScope; -use crate::span::SpanSubmitter; // ============================================================================ // Dataset Insert Types @@ -237,53 +236,8 @@ pub(crate) struct DatasetMetadata { pub project_id: String, } -pub(crate) use crate::api::btql::{BTQLQuery, BTQLQueryInner, BTQLResponse}; -pub(crate) use crate::api::datasets::{ - DatasetRegisterRequest, DatasetRegisterResponse, DatasetSummaryResponse, -}; - -// ============================================================================ -// 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, - scope: &CredentialScope, - 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, - scope: &CredentialScope, - 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, - scope: &CredentialScope, - 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, scope: &CredentialScope, row: DatasetWrite) -> Result<()>; -} +use crate::api::btql::{BTQLQuery, BTQLQueryInner}; +use crate::api::datasets::DatasetRegisterRequest; // ============================================================================ // DatasetBuilder Error Types @@ -319,11 +273,9 @@ impl std::error::Error for DatasetBuilderError {} // ============================================================================ /// Builder for creating datasets with lazy registration. -#[allow(private_bounds)] -pub struct DatasetBuilder< - S: SpanSubmitter + DatasetRegistrar + DatasetFetcher + DatasetSummarizer + DatasetWriter, -> { - submitter: Arc, +pub struct DatasetBuilder { + api: Arc, + lifecycle: Arc, scope: CredentialScope, project_name: Option, dataset_name: Option, @@ -331,12 +283,11 @@ pub struct DatasetBuilder< metadata: Option>, } -impl Clone - for DatasetBuilder -{ +impl Clone for DatasetBuilder { fn clone(&self) -> Self { Self { - submitter: Arc::clone(&self.submitter), + api: self.api.clone(), + lifecycle: self.lifecycle.clone(), scope: self.scope.clone(), project_name: self.project_name.clone(), dataset_name: self.dataset_name.clone(), @@ -346,20 +297,16 @@ impl DatasetBuilder -{ +impl DatasetBuilder { /// Create a new DatasetBuilder (internal use). - pub(crate) fn new(submitter: Arc, scope: CredentialScope) -> Self { + pub(crate) fn new( + api: Arc, + lifecycle: Arc, + scope: CredentialScope, + ) -> Self { Self { - submitter, + api, + lifecycle, scope, project_name: None, dataset_name: None, @@ -401,7 +348,7 @@ impl< } /// Build the dataset (lazy registration). - pub fn build(self) -> std::result::Result, DatasetBuilderError> { + pub fn build(self) -> std::result::Result { let project_name = self .project_name .ok_or(DatasetBuilderError::MissingProjectName)?; @@ -410,7 +357,8 @@ impl< .ok_or(DatasetBuilderError::MissingDatasetName)?; Ok(Dataset { - submitter: self.submitter, + api: self.api, + lifecycle: self.lifecycle, scope: self.scope, project_name, dataset_name, @@ -428,11 +376,9 @@ impl< /// A dataset for storing test/evaluation data. /// /// Datasets are lazily registered with the API on first use. -#[allow(private_bounds)] -pub struct Dataset< - S: SpanSubmitter + DatasetRegistrar + DatasetFetcher + DatasetSummarizer + DatasetWriter, -> { - submitter: Arc, +pub struct Dataset { + api: Arc, + lifecycle: Arc, scope: CredentialScope, project_name: String, dataset_name: String, @@ -441,12 +387,11 @@ pub struct Dataset< lazy_metadata: Arc>, } -impl Clone - for Dataset -{ +impl Clone for Dataset { fn clone(&self) -> Self { Self { - submitter: self.submitter.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(), @@ -457,9 +402,7 @@ impl - fmt::Debug for Dataset -{ +impl fmt::Debug for Dataset { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("Dataset") .field("project_name", &self.project_name) @@ -469,16 +412,7 @@ impl Dataset -{ +impl Dataset { /// Ensure the dataset is registered, returning cached metadata. async fn ensure_registered(&self) -> Result<&DatasetMetadata> { self.lazy_metadata @@ -490,10 +424,7 @@ impl< description: self.description.clone(), metadata: self.metadata.clone(), }; - let response = self - .submitter - .register_dataset(&self.scope, request) - .await?; + let response = self.api.register_dataset(&self.scope, request).await?; Ok(DatasetMetadata { dataset_id: response.dataset.id, project_id: response.project.id, @@ -521,7 +452,7 @@ impl< 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(); - self.submitter + self.api .write_dataset_record( &self.scope, DatasetWrite::Insert { @@ -537,7 +468,7 @@ impl< /// 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 + self.api .write_dataset_record( &self.scope, DatasetWrite::Update { @@ -552,7 +483,7 @@ impl< /// 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 + self.api .write_dataset_record( &self.scope, DatasetWrite::Delete { @@ -566,10 +497,10 @@ impl< /// 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.api.clone(), self.scope.clone(), metadata.dataset_id.clone(), batch_size.unwrap_or(1000).max(1), @@ -582,7 +513,7 @@ impl< // Fetch summary from API let response = self - .submitter + .api .fetch_dataset_summary(&self.scope, &metadata.dataset_id) .await?; @@ -600,7 +531,7 @@ impl< /// 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() + self.lifecycle.flush() } } @@ -611,18 +542,17 @@ impl< /// Iterator for fetching dataset records with pagination. /// /// ```no_run -/// use braintrust_sdk_rust::{BraintrustClient, BraintrustError, Dataset}; -/// # async fn read(dataset: &Dataset) -> Result<(), BraintrustError> { +/// 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); +/// println!("{}", record.id()); /// } /// # Ok(()) /// # } /// ``` -#[allow(private_bounds)] -pub struct DatasetIterator { - fetcher: Arc, +pub struct DatasetIterator { + fetcher: Arc, scope: CredentialScope, dataset_id: String, batch_size: usize, @@ -631,9 +561,13 @@ pub struct DatasetIterator { done: bool, } -#[allow(private_bounds)] -impl DatasetIterator { - fn new(fetcher: Arc, scope: CredentialScope, dataset_id: String, batch_size: usize) -> Self { +impl DatasetIterator { + fn new( + fetcher: Arc, + scope: CredentialScope, + dataset_id: String, + batch_size: usize, + ) -> Self { Self { fetcher, scope, diff --git a/src/eval/runner.rs b/src/eval/runner.rs index 96db594..052ed33 100644 --- a/src/eval/runner.rs +++ b/src/eval/runner.rs @@ -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, 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 173bffc..0000000 --- a/src/experiments/api/comparison.rs +++ /dev/null @@ -1 +0,0 @@ -pub(crate) use crate::api::experiments::ExperimentComparisonResponse; diff --git a/src/experiments/api/mod.rs b/src/experiments/api/mod.rs deleted file mode 100644 index 14d106c..0000000 --- a/src/experiments/api/mod.rs +++ /dev/null @@ -1,11 +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 registration::ExperimentRegisterRequest; -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 4dac4cb..0000000 --- a/src/experiments/api/registration.rs +++ /dev/null @@ -1,2 +0,0 @@ -//! Compatibility imports for existing experiment callers. -pub(crate) use crate::api::experiments::{ExperimentRegisterRequest, ExperimentRegisterResponse}; diff --git a/src/experiments/api/traits.rs b/src/experiments/api/traits.rs deleted file mode 100644 index e14bd8c..0000000 --- a/src/experiments/api/traits.rs +++ /dev/null @@ -1,40 +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; -use crate::internal::session::CredentialScope; - -/// Trait for registering experiments with the Braintrust API. -#[async_trait::async_trait] -pub(crate) trait ExperimentRegistrar: Send + Sync { - async fn register_experiment( - &self, - scope: &CredentialScope, - 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, - scope: &CredentialScope, - 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, - scope: &CredentialScope, - experiment_id: &str, - base_experiment_id: Option<&str>, - ) -> Result; -} diff --git a/src/experiments/builder.rs b/src/experiments/builder.rs index 9ba3376..fe97b25 100644 --- a/src/experiments/builder.rs +++ b/src/experiments/builder.rs @@ -7,8 +7,8 @@ use serde_json::{Map, Value}; use crate::experiments::experiment::{Experiment, ExperimentState}; use crate::experiments::metadata::RepoInfo; +use crate::internal::ports::{ExperimentApi, SpanSubmitter}; use crate::internal::session::CredentialScope; -use crate::span::SpanSubmitter; /// Error type for ExperimentBuilder validation. #[derive(Debug, Clone, PartialEq)] @@ -34,9 +34,9 @@ impl std::error::Error for ExperimentBuilderError {} /// /// Registration is shared by clones of the resulting handle. Explicit /// credentials belong to the resource and do not mutate the client's session. -#[allow(private_bounds)] -pub struct ExperimentBuilder { - submitter: Arc, +pub struct ExperimentBuilder { + api: Arc, + submitter: Arc, scope: CredentialScope, project_name: Option, experiment_name: Option, @@ -48,10 +48,11 @@ pub struct ExperimentBuilder { repo_info: Option, } -impl Clone for ExperimentBuilder { +impl Clone for ExperimentBuilder { fn clone(&self) -> Self { Self { - submitter: Arc::clone(&self.submitter), + api: self.api.clone(), + submitter: self.submitter.clone(), scope: self.scope.clone(), project_name: self.project_name.clone(), experiment_name: self.experiment_name.clone(), @@ -65,11 +66,15 @@ impl Clone for ExperimentBuilder { } } -#[allow(private_bounds)] -impl ExperimentBuilder { +impl ExperimentBuilder { /// Create a new ExperimentBuilder (internal use). - pub(crate) fn new(submitter: Arc, scope: CredentialScope) -> Self { + pub(crate) fn new( + api: Arc, + submitter: Arc, + scope: CredentialScope, + ) -> Self { Self { + api, submitter, scope, project_name: None, @@ -152,12 +157,13 @@ 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, scope: self.scope, project_name, diff --git a/src/experiments/experiment.rs b/src/experiments/experiment.rs index e8eb515..02c33cb 100644 --- a/src/experiments/experiment.rs +++ b/src/experiments/experiment.rs @@ -8,11 +8,8 @@ use std::time::{SystemTime, UNIX_EPOCH}; use serde_json::{Map, Value}; 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::{ @@ -20,8 +17,9 @@ use crate::experiments::metadata::{ }; use crate::experiments::span_builder::ExperimentSpanBuilder; use crate::experiments::summary::{ExperimentSummary, MetricSummary, ScoreSummary}; +use crate::internal::ports::{ExperimentApi, SpanSubmitter}; use crate::internal::session::CredentialScope; -use crate::span::{SpanBuilder, SpanHandle, SpanLog, SpanSubmitter}; +use crate::span::{SpanBuilder, SpanHandle, SpanLog}; use crate::types::{ParentSpanInfo, SpanAttributes, SpanPayload, SpanType}; /// Metadata from experiment registration, cached after lazy initialization. @@ -35,9 +33,9 @@ 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 struct Experiment { + pub(crate) api: Arc, + pub(crate) submitter: Arc, pub(crate) scope: CredentialScope, pub(crate) project_name: String, pub(crate) experiment_name: Option, @@ -59,10 +57,11 @@ pub(crate) struct ExperimentState { called_start_span: AtomicBool, } -impl Clone for Experiment { +impl Clone for Experiment { fn clone(&self) -> Self { Self { - submitter: Arc::clone(&self.submitter), + api: self.api.clone(), + submitter: self.submitter.clone(), scope: self.scope.clone(), project_name: self.project_name.clone(), experiment_name: self.experiment_name.clone(), @@ -77,15 +76,7 @@ impl Clone for Experiment { } } -#[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()`. @@ -142,7 +133,7 @@ impl< /// /// Use this for more control over span lifecycle. Once called, /// you cannot use `log()` on this experiment. - pub async fn start_span(&self) -> Result> { + 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())) @@ -221,7 +212,7 @@ impl< // Fetch comparison data from API let comparison = self - .submitter + .api .fetch_experiment_comparison( &self.scope, &metadata.experiment_id, @@ -336,7 +327,7 @@ impl< /// ``` 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(); @@ -394,7 +385,7 @@ impl< /// 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 + self.api .fetch_base_experiment(&self.scope, &metadata.experiment_id) .await } @@ -427,10 +418,7 @@ impl< update: self.update, repo_info: self.repo_info.clone(), }; - let response = self - .submitter - .register_experiment(&self.scope, request) - .await?; + let response = self.api.register_experiment(&self.scope, request).await?; Ok(ExperimentMetadata { project_id: response.project.id, experiment_id: response.experiment.id, @@ -441,11 +429,7 @@ impl< } /// Create an internal span for log() calls. - 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(), }; diff --git a/src/experiments/mod.rs b/src/experiments/mod.rs index 870e762..d38f8c2 100644 --- a/src/experiments/mod.rs +++ b/src/experiments/mod.rs @@ -26,9 +26,6 @@ //! 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 0efb335..1f04a6d 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,19 +40,16 @@ 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 builder = SpanBuilder::new( - Arc::clone(&self.experiment.submitter), - self.experiment.scope.clone(), - ) - .parent_info(parent_info) - .span_type(self.span_type); + let builder = SpanBuilder::new(self.submitter, self.scope) + .parent_info(parent_info) + .span_type(self.span_type); builder.build() } diff --git a/src/internal/mod.rs b/src/internal/mod.rs index 7eb3eba..965aab2 100644 --- a/src/internal/mod.rs +++ b/src/internal/mod.rs @@ -2,6 +2,7 @@ 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; diff --git a/src/internal/ports.rs b/src/internal/ports.rs new file mode 100644 index 0000000..a5ba29b --- /dev/null +++ b/src/internal/ports.rs @@ -0,0 +1,73 @@ +//! 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::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, + ) -> 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 index 4d74ee2..557aec6 100644 --- a/src/internal/resources.rs +++ b/src/internal/resources.rs @@ -12,16 +12,26 @@ use crate::api::{ 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; +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 ExperimentRegistrar for BraintrustClient { +impl ExperimentApi for ResourceTransport { async fn register_experiment( &self, scope: &CredentialScope, @@ -33,10 +43,7 @@ impl ExperimentRegistrar for BraintrustClient { .await .map_err(Into::into) } -} -#[async_trait] -impl BaseExperimentFetcher for BraintrustClient { async fn fetch_base_experiment( &self, scope: &CredentialScope, @@ -59,10 +66,7 @@ impl BaseExperimentFetcher for BraintrustClient { _ => None, }) } -} -#[async_trait] -impl ExperimentComparisonFetcher for BraintrustClient { async fn fetch_experiment_comparison( &self, scope: &CredentialScope, @@ -78,7 +82,7 @@ impl ExperimentComparisonFetcher for BraintrustClient { } #[async_trait] -impl DatasetRegistrar for BraintrustClient { +impl DatasetApi for ResourceTransport { async fn register_dataset( &self, scope: &CredentialScope, @@ -90,10 +94,7 @@ impl DatasetRegistrar for BraintrustClient { .await .map_err(Into::into) } -} -#[async_trait] -impl DatasetFetcher for BraintrustClient { async fn fetch_dataset_records( &self, scope: &CredentialScope, @@ -105,10 +106,7 @@ impl DatasetFetcher for BraintrustClient { .await .map_err(Into::into) } -} -#[async_trait] -impl DatasetSummarizer for BraintrustClient { async fn fetch_dataset_summary( &self, scope: &CredentialScope, @@ -120,10 +118,7 @@ impl DatasetSummarizer for BraintrustClient { .await .map_err(Into::into) } -} -#[async_trait] -impl DatasetWriter for BraintrustClient { async fn write_dataset_record(&self, scope: &CredentialScope, row: DatasetWrite) -> Result<()> { self.resource_api(scope)? .datasets() diff --git a/src/lib.rs b/src/lib.rs index a3b9857..15c7583 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,3 +1,5 @@ +#![deny(private_bounds, private_interfaces)] + pub mod api; mod completion; mod dataset; diff --git a/src/log_queue/queue.rs b/src/log_queue/queue.rs index 92b2bb1..eb303b7 100644 --- a/src/log_queue/queue.rs +++ b/src/log_queue/queue.rs @@ -7,6 +7,7 @@ use super::merge::merge_row_into; use super::row_key::RowKey; use crate::error::{BraintrustError, Result}; use crate::internal::delivery::WriteService; +use crate::internal::ports::{DeliveryControl, SpanSubmitter}; use crate::internal::session::CredentialScope; use crate::span_components::{project_logs_identifier, ProjectLogsIdentifier}; use crate::types::{LogDestination, Logs3Row, ParentSpanInfo, SpanObjectType, SpanPayload}; @@ -692,6 +693,27 @@ impl LogQueue { } } +#[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, + ) -> Completion { + LogQueue::submit(self, scope, payload, parent_info) + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/logger.rs b/src/logger.rs index 4284874..0f74159 100644 --- a/src/logger.rs +++ b/src/logger.rs @@ -11,10 +11,12 @@ use crate::dataset::DatasetBuilder; use crate::error::{BraintrustError, Result}; 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}; @@ -460,8 +462,12 @@ impl std::fmt::Debug for ClientInner { } impl BraintrustClient { - pub(crate) fn resource_api(&self, scope: &CredentialScope) -> Result { - scope.api(&self.inner.http_client) + 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( @@ -540,14 +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 mut builder = crate::span::SpanBuilder::new(Arc::new(self.clone()), state.scope()?); + 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()); } @@ -563,12 +569,16 @@ 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?; - Ok(DatasetBuilder::new(Arc::new(self.clone()), state.scope()?)) + Ok(DatasetBuilder::new( + self.resource_transport(), + Arc::new(self.inner.queue.clone()), + state.scope()?, + )) } /// Create a dataset builder with immutable explicit credentials, without @@ -577,9 +587,10 @@ impl BraintrustClient { &self, token: impl Into, org_name: impl Into, - ) -> DatasetBuilder { + ) -> DatasetBuilder { DatasetBuilder::new( - Arc::new(self.clone()), + self.resource_transport(), + Arc::new(self.inner.queue.clone()), self.explicit_scope(token.into(), String::new(), Some(org_name.into())), ) } @@ -594,11 +605,11 @@ impl BraintrustClient { &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 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 { @@ -925,12 +936,16 @@ impl BraintrustClient { &self, token: impl Into, org_id: impl Into, - ) -> ExperimentBuilder { + ) -> ExperimentBuilder { let token = token.into(); let org_id = org_id.into(); - let submitter = Arc::new(self.clone()); - ExperimentBuilder::new(submitter, self.explicit_scope(token, org_id, None)) + let submitter = self.span_submitter(); + ExperimentBuilder::new( + self.resource_transport(), + submitter, + self.explicit_scope(token, org_id, None), + ) } /// Submit immediately; the token observes admission and full batch delivery. @@ -965,26 +980,6 @@ impl BraintrustClient { } } -#[async_trait::async_trait] -impl SpanSubmitter for BraintrustClient { - fn submit( - &self, - scope: CredentialScope, - payload: SpanPayload, - parent_info: Option, - ) -> crate::Completion { - self.submit_payload(scope, payload, parent_info) - } - - fn flush(&self) -> crate::Completion { - BraintrustClient::flush(self) - } - - async fn trigger_flush(&self) -> Result<()> { - BraintrustClient::trigger_flush(self).await - } -} - #[allow(dead_code)] #[derive(Debug, Default)] struct MutableSpanEvent { @@ -1171,6 +1166,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 diff --git a/src/span.rs b/src/span.rs index ae7417e..dc6e8db 100644 --- a/src/span.rs +++ b/src/span.rs @@ -5,11 +5,13 @@ 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::span_components::SpanComponents; use crate::types::{ @@ -248,27 +250,8 @@ impl SpanLogBuilder { } } -#[async_trait] -pub(crate) trait SpanSubmitter: Send + Sync { - /// Attempt admission immediately and observe the logical batch outcome. - fn submit( - &self, - scope: CredentialScope, - payload: SpanPayload, - parent_info: Option, - ) -> crate::Completion; - - /// Flush all queued span data through to the server. - fn flush(&self) -> crate::Completion; - - /// Trigger a non-blocking background flush. - async fn trigger_flush(&self) -> Result<()>; -} - -#[derive(Debug)] -#[allow(private_bounds)] -pub struct SpanBuilder { - submitter: Arc, +pub struct SpanBuilder { + submitter: Arc, scope: CredentialScope, project_name: Option, parent_info: Option, @@ -282,7 +265,17 @@ 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), @@ -369,9 +362,8 @@ impl SpanOrigin { } } -#[allow(private_bounds)] -impl SpanBuilder { - pub(crate) fn new(submitter: Arc, scope: CredentialScope) -> Self { +impl SpanBuilder { + pub(crate) fn new(submitter: Arc, scope: CredentialScope) -> Self { Self { submitter, scope, @@ -479,7 +471,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. @@ -535,9 +527,8 @@ impl SpanBuilder { } } -#[allow(private_bounds)] -pub struct SpanHandle { - submitter: Arc, +pub struct SpanHandle { + submitter: Arc, scope: CredentialScope, parent_info: Option, /// Stable identifier set at creation, accessible without locking. @@ -545,7 +536,7 @@ pub struct SpanHandle { inner: Arc>, } -impl Clone for SpanHandle { +impl Clone for SpanHandle { fn clone(&self) -> Self { Self { submitter: Arc::clone(&self.submitter), @@ -557,8 +548,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 @@ -1085,6 +1075,9 @@ mod tests { } completion } + } + #[async_trait] + impl crate::internal::ports::DeliveryControl for Submitter { fn flush(&self) -> crate::Completion { let (sender, completion) = crate::Completion::pair(); let _ = sender.send(Ok(crate::FlushReport::default())); diff --git a/src/stream.rs b/src/stream.rs index 248050e..e4f061c 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() { @@ -1204,11 +1198,10 @@ async fn finalize_span( } } -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 4489b43..18ca211 100644 --- a/src/test_utils.rs +++ b/src/test_utils.rs @@ -7,7 +7,7 @@ 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}; #[derive(Clone, Default)] pub(crate) struct TestSpanCollector { @@ -74,7 +74,10 @@ impl SpanSubmitter for MockBraintrustClient { let _ = sender.send(Ok(crate::DeliveryReceipt::new())); completion } +} +#[async_trait] +impl DeliveryControl for MockBraintrustClient { fn flush(&self) -> crate::Completion { let (sender, completion) = crate::Completion::pair(); let _ = sender.send(Ok(crate::FlushReport::default())); @@ -87,13 +90,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, 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/tests/dataset_test.rs b/tests/dataset_test.rs index 567040a..8b16d23 100644 --- a/tests/dataset_test.rs +++ b/tests/dataset_test.rs @@ -298,7 +298,7 @@ async fn isolated_client(server: &MockServer) -> BraintrustClient { 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 { +fn dataset(client: &BraintrustClient) -> braintrust_sdk_rust::Dataset { client .dataset_builder_with_credentials("token", "tenant") .project_name("project") @@ -381,12 +381,12 @@ async fn pagination_failure_is_an_error_and_retry_preserves_the_failed_cursor() 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_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_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 @@ -469,7 +469,7 @@ async fn empty_pages_with_a_continuation_do_not_end_iteration() { .await .unwrap(); assert_eq!(records.len(), 1); - assert_eq!(records[0].id, "record"); + assert_eq!(records[0].id(), "record"); for request in server .received_requests() .await @@ -516,11 +516,11 @@ async fn a_non_advancing_cursor_is_reported_and_does_not_poison_a_later_retry() 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_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_eq!(iterator.next().await.unwrap().unwrap().id(), "second"); assert!(iterator.next().await.unwrap().is_none()); } @@ -705,7 +705,7 @@ async fn resource_operations_preserve_configured_planes_prefixes_and_transport_h let dataset = dataset(&client); dataset.insert(record(1)).await.unwrap(); assert_eq!( - dataset.fetch(None).await.unwrap().collect().await.unwrap()[0].id, + dataset.fetch(None).await.unwrap().collect().await.unwrap()[0].id(), "record" ); dataset.summarize().await.unwrap(); @@ -720,13 +720,7 @@ async fn resource_operations_preserve_configured_planes_prefixes_and_transport_h } 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() + request.headers.get(name).unwrap().to_str().unwrap() } #[tokio::test] @@ -749,9 +743,9 @@ async fn dataset_flush_observes_the_shared_trace_queue_and_retires_its_failure() .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!(matches!(flush.await, Err(FlushError::Delivery(failure)) if failure.failed() == 1)); assert!(write.await.is_err()); - assert_eq!(dataset.flush().await.unwrap().delivered, 0); + assert_eq!(dataset.flush().await.unwrap().delivered(), 0); let requests = server.received_requests().await.unwrap(); assert!(requests .iter() diff --git a/tests/resource_handles.rs b/tests/resource_handles.rs new file mode 100644 index 0000000..ee364cd --- /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(); +} From 99aa20be8e52b87cba9d5fa46422e3ebf71adf86 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Thu, 29 Oct 2026 23:20:04 +0800 Subject: [PATCH 5/7] Make resource data types extensible --- src/api/datasets.rs | 156 ++++++++++++++++++++++++------------ src/dataset.rs | 20 +---- src/error.rs | 25 ++++++ src/experiments/metadata.rs | 107 ++++++++++++++----------- src/experiments/mod.rs | 5 +- src/lib.rs | 3 +- src/span.rs | 95 ++++++++++++++-------- tests/api_client.rs | 31 +++---- tests/dataset_test.rs | 1 + tests/experiments.rs | 8 +- tests/span_lifecycle.rs | 10 +-- 11 files changed, 278 insertions(+), 183 deletions(-) diff --git a/src/api/datasets.rs b/src/api/datasets.rs index a9bd25c..d49d083 100644 --- a/src/api/datasets.rs +++ b/src/api/datasets.rs @@ -71,72 +71,122 @@ pub struct DatasetInfo { } /// Fields stored in a dataset record. Absent fields are omitted from the wire. -#[derive(Debug, Clone, Default, Serialize)] +/// +/// ``` +/// use braintrust_sdk_rust::api::datasets::DatasetFields; +/// use serde_json::json; +/// let fields = DatasetFields::builder().input(json!(null)).tags(vec!["test".into()]).build(); +/// assert_eq!(fields.input(), Some(&json!(null))); +/// assert!(fields.expected().is_none()); +/// ``` +/// +/// Construct fields through the builder rather than depending on their layout: +/// +/// ```compile_fail +/// use braintrust_sdk_rust::api::datasets::DatasetFields; +/// let mut fields = DatasetFields::default(); +/// fields.input = Some(serde_json::json!(1)); +/// ``` +#[derive(bon::Builder, Debug, Clone, Default, Serialize)] pub struct DatasetFields { #[serde(skip_serializing_if = "Option::is_none")] - pub input: Option, + #[builder(into)] + pub(crate) input: Option, #[serde(skip_serializing_if = "Option::is_none")] - pub expected: Option, + #[builder(into)] + pub(crate) expected: Option, #[serde(skip_serializing_if = "Option::is_none")] - pub metadata: Option>, + pub(crate) metadata: Option>, #[serde(skip_serializing_if = "Option::is_none")] - pub tags: Option>, + pub(crate) tags: Option>, +} +impl DatasetFields { + /// Return the supplied input, including an explicit JSON null. + pub fn input(&self) -> Option<&Value> { + self.input.as_ref() + } + /// Return the expected output, if supplied. + pub fn expected(&self) -> Option<&Value> { + self.expected.as_ref() + } + /// Return the metadata entries, if supplied. + pub fn metadata(&self) -> Option<&Map> { + self.metadata.as_ref() + } + /// Return the tags, if supplied. + pub fn tags(&self) -> Option<&[String]> { + self.tags.as_deref() + } } /// A dataset row mutation, independent of tracing/span payloads. -#[derive(Debug, Clone)] -pub enum DatasetWrite { - Insert { - dataset_id: String, - id: String, +/// +/// Use [`Self::insert`], [`Self::update`], or [`Self::delete`] to select the +/// operation. Its private fields keep wire flags and payload layout extensible. +/// +/// ```compile_fail +/// use braintrust_sdk_rust::api::datasets::DatasetWrite; +/// let mut write = DatasetWrite::delete("dataset", "row"); +/// write.id = "other".into(); +/// ``` +#[derive(Debug, Clone, Serialize)] +pub struct DatasetWrite { + dataset_id: String, + id: String, + #[serde(flatten)] + fields: Option, + #[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, +} +impl DatasetWrite { + /// Insert a record with the supplied ID and fields. + pub fn insert( + dataset_id: impl Into, + id: impl Into, fields: DatasetFields, - }, + ) -> Self { + Self { + dataset_id: dataset_id.into(), + id: id.into(), + fields: Some(fields), + is_merge: None, + is_delete: None, + } + } /// Merge only the supplied fields into an existing record. - Update { - dataset_id: String, - id: String, + pub fn update( + dataset_id: impl Into, + id: impl Into, 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, + ) -> Self { + Self { + is_merge: Some(true), + ..Self::insert(dataset_id, id, fields) } - 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, + } + /// Delete a record without sending record fields. + pub fn delete(dataset_id: impl Into, id: impl Into) -> Self { + Self { + dataset_id: dataset_id.into(), + id: id.into(), + fields: None, + is_merge: None, + is_delete: Some(true), } - .serialize(serializer) + } + /// Return the dataset receiving this mutation. + pub fn dataset_id(&self) -> &str { + &self.dataset_id + } + /// Return the record ID. + pub fn id(&self) -> &str { + &self.id + } + /// Return the supplied record fields. Deletes have no fields. + pub fn fields(&self) -> Option<&DatasetFields> { + self.fields.as_ref() } } diff --git a/src/dataset.rs b/src/dataset.rs index 74fde64..046685a 100644 --- a/src/dataset.rs +++ b/src/dataset.rs @@ -455,11 +455,7 @@ impl Dataset { self.api .write_dataset_record( &self.scope, - DatasetWrite::Insert { - dataset_id, - id: id.clone(), - fields: record.into(), - }, + DatasetWrite::insert(dataset_id, id.clone(), record.into()), ) .await?; Ok(id) @@ -471,11 +467,7 @@ impl Dataset { self.api .write_dataset_record( &self.scope, - DatasetWrite::Update { - dataset_id, - id: id.to_owned(), - fields: record.into(), - }, + DatasetWrite::update(dataset_id, id, record.into()), ) .await } @@ -484,13 +476,7 @@ impl Dataset { 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(), - }, - ) + .write_dataset_record(&self.scope, DatasetWrite::delete(dataset_id, id)) .await } diff --git a/src/error.rs b/src/error.rs index 7d2c366..cc89536 100644 --- a/src/error.rs +++ b/src/error.rs @@ -1,6 +1,20 @@ use thiserror::Error; +/// SDK failures. Match known cases with a fallback for future variants. +/// +/// ```compile_fail +/// use braintrust_sdk_rust::BraintrustError; +/// fn consume(error: BraintrustError) { +/// match error { +/// BraintrustError::InvalidConfig(_) | BraintrustError::Network(_) +/// | BraintrustError::Background(_) | BraintrustError::StreamAggregation(_) +/// | BraintrustError::Http(_) | BraintrustError::Api { .. } +/// | BraintrustError::ChannelClosed => {} +/// } +/// } +/// ``` #[derive(Debug, Error)] +#[non_exhaustive] pub enum BraintrustError { #[error("invalid configuration: {0}")] InvalidConfig(String), @@ -8,7 +22,18 @@ pub enum BraintrustError { Http(#[from] reqwest::Error), #[error("network error: {0}")] Network(String), + /// API failures may gain additional diagnostic fields. + /// + /// ```compile_fail + /// use braintrust_sdk_rust::BraintrustError; + /// fn consume(error: BraintrustError) { + /// if let BraintrustError::Api { status, message } = error { + /// let _ = (status, message); + /// } + /// } + /// ``` #[error("API error ({status}): {message}")] + #[non_exhaustive] Api { status: u16, message: String }, #[error("internal channel closed")] ChannelClosed, diff --git a/src/experiments/metadata.rs b/src/experiments/metadata.rs index 9ca9327..4a7257f 100644 --- a/src/experiments/metadata.rs +++ b/src/experiments/metadata.rs @@ -3,74 +3,89 @@ use serde::{Deserialize, Serialize}; /// Metadata about the project containing an experiment. +/// +/// Read values through [`Self::id`] and [`Self::name`]. +/// +/// ```compile_fail +/// use braintrust_sdk_rust::ProjectMetadata; +/// fn rename(project: &mut ProjectMetadata) { project.name = "other".into(); } +/// ``` #[derive(Debug, Clone)] pub struct ProjectMetadata { /// The project ID. - pub id: String, + pub(crate) id: String, /// The project name. - pub name: String, + pub(crate) name: String, } /// Base experiment identifier for comparison. +/// +/// ```compile_fail +/// use braintrust_sdk_rust::BaseExperimentInfo; +/// fn rename(base: &mut BaseExperimentInfo) { base.name = "other".into(); } +/// ``` #[derive(Debug, Clone)] pub struct BaseExperimentInfo { /// The base experiment ID. - pub id: String, + pub(crate) id: String, /// The base experiment name. - pub name: String, + pub(crate) name: String, } /// Exported experiment context for cross-process tracing. /// /// Use this to pass experiment context to another process or service. +/// +/// ``` +/// use braintrust_sdk_rust::ExportedExperiment; +/// let context: ExportedExperiment = serde_json::from_str( +/// r#"{"object_type":"experiment","object_id":"test-experiment"}"#, +/// ).unwrap(); +/// assert_eq!(context.object_type(), "experiment"); +/// assert_eq!(context.object_id(), "test-experiment"); +/// ``` +/// +/// ```compile_fail +/// use braintrust_sdk_rust::ExportedExperiment; +/// fn retarget(context: &mut ExportedExperiment) { context.object_id = "other".into(); } +/// ``` #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ExportedExperiment { /// The object type (always "experiment"). - pub object_type: String, + pub(crate) object_type: String, /// The experiment ID. - pub object_id: String, + pub(crate) object_id: String, } -/// Settings for git metadata collection. -#[derive(Debug, Clone, Default)] -pub struct GitMetadataSettings { - /// How much git metadata to collect. - pub collect: GitMetadataCollect, - /// Specific fields to collect (when collect is `Some`). - pub fields: Option>, -} +pub use crate::api::experiments::RepoInfo; -/// How much git metadata to collect. -#[derive(Debug, Clone, Default)] -pub enum GitMetadataCollect { - /// Collect all available git metadata. - #[default] - All, - /// Don't collect any git metadata. - None, - /// Collect only specific fields. - Some, +impl ProjectMetadata { + /// Return the project ID. + pub fn id(&self) -> &str { + &self.id + } + /// Return the project name. + pub fn name(&self) -> &str { + &self.name + } } - -/// Individual git metadata fields that can be collected. -#[derive(Debug, Clone)] -pub enum GitMetadataField { - /// Git commit SHA. - Commit, - /// Current branch name. - Branch, - /// Current tag (if any). - Tag, - /// Whether the working directory has uncommitted changes. - Dirty, - /// Git author name. - AuthorName, - /// Git author email. - AuthorEmail, - /// Git commit message. - CommitMessage, - /// Git commit timestamp. - CommitTime, +impl BaseExperimentInfo { + /// Return the base experiment ID. + pub fn id(&self) -> &str { + &self.id + } + /// Return the base experiment name. + pub fn name(&self) -> &str { + &self.name + } +} +impl ExportedExperiment { + /// Return the exported object type. + pub fn object_type(&self) -> &str { + &self.object_type + } + /// Return the experiment ID. + pub fn object_id(&self) -> &str { + &self.object_id + } } - -pub use crate::api::experiments::RepoInfo; diff --git a/src/experiments/mod.rs b/src/experiments/mod.rs index d38f8c2..fa361e8 100644 --- a/src/experiments/mod.rs +++ b/src/experiments/mod.rs @@ -40,9 +40,6 @@ pub use builder::{ExperimentBuilder, ExperimentBuilderError}; pub use event_log::{ExperimentLog, ExperimentLogBuilder}; pub use experiment::Experiment; pub use feedback::{Feedback, FeedbackBuilder, FeedbackBuilderError}; -pub use metadata::{ - BaseExperimentInfo, ExportedExperiment, GitMetadataCollect, GitMetadataField, - GitMetadataSettings, ProjectMetadata, RepoInfo, -}; +pub use metadata::{BaseExperimentInfo, ExportedExperiment, ProjectMetadata, RepoInfo}; pub use span_builder::ExperimentSpanBuilder; pub use summary::{ExperimentSummary, MetricSummary, ScoreSummary}; diff --git a/src/lib.rs b/src/lib.rs index 15c7583..23379a3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -34,8 +34,7 @@ pub use error::{BraintrustError, Result}; pub use experiments::{ BaseExperimentInfo, Experiment, ExperimentBuilder, ExperimentBuilderError, ExperimentLog, ExperimentLogBuilder, ExperimentSpanBuilder, ExperimentSummary, ExportedExperiment, Feedback, - FeedbackBuilder, FeedbackBuilderError, GitMetadataCollect, GitMetadataField, - GitMetadataSettings, MetricSummary, ProjectMetadata, RepoInfo, ScoreSummary, + FeedbackBuilder, FeedbackBuilderError, MetricSummary, ProjectMetadata, RepoInfo, ScoreSummary, }; pub use extractors::{extract_anthropic_usage, extract_openai_usage}; pub use log_queue::LogQueueConfig; diff --git a/src/span.rs b/src/span.rs index dc6e8db..a8133a7 100644 --- a/src/span.rs +++ b/src/span.rs @@ -294,7 +294,22 @@ impl Clone for SpanBuilder { } } +/// Small options bag identifying the environment that produced a span. +/// +/// Use the constructor so future options can be added: +/// +/// ``` +/// use braintrust_sdk_rust::SpanOriginEnvironment; +/// let environment = SpanOriginEnvironment::new("test", Some("test-environment")); +/// assert_eq!(environment.environment_type.as_deref(), Some("test")); +/// ``` +/// +/// ```compile_fail +/// use braintrust_sdk_rust::SpanOriginEnvironment; +/// let _ = SpanOriginEnvironment { environment_type: Some("test".into()), name: None }; +/// ``` #[derive(Clone, Debug, PartialEq, Eq)] +#[non_exhaustive] pub struct SpanOriginEnvironment { pub environment_type: Option, pub name: Option, @@ -319,46 +334,57 @@ impl SpanOriginEnvironment { /// Any explicit `span_origin` values already present in a span's `context` take /// precedence over these overrides, which in turn take precedence over the /// built-in defaults. -#[derive(Clone, Debug, Default, PartialEq, Eq)] +/// +/// ``` +/// use braintrust_sdk_rust::{SpanOrigin, SpanOriginEnvironment}; +/// let origin = SpanOrigin::builder() +/// .name("test-instrumentation").version("1.0") +/// .instrumentation("test-plugin") +/// .environment(SpanOriginEnvironment::new("test", Some("test-environment"))) +/// .build(); +/// assert_eq!(origin.name(), Some("test-instrumentation")); +/// assert_eq!(origin.version(), Some("1.0")); +/// assert_eq!(origin.instrumentation(), Some("test-plugin")); +/// assert!(origin.environment().is_some()); +/// assert_eq!(SpanOrigin::builder().build(), SpanOrigin::default()); +/// ``` +/// +/// ```compile_fail +/// use braintrust_sdk_rust::SpanOrigin; +/// let mut origin = SpanOrigin::default(); +/// origin.name = Some("other".into()); +/// ``` +#[derive(bon::Builder, Clone, Debug, Default, PartialEq, Eq)] pub struct SpanOrigin { /// Overrides `span_origin.name` (e.g. `"braintrust.plugin.codex"`). - pub name: Option, + #[builder(into)] + name: Option, /// Overrides `span_origin.version` (e.g. a plugin version). - pub version: Option, + #[builder(into)] + version: Option, /// Overrides `span_origin.instrumentation.name`. - pub instrumentation: Option, + #[builder(into)] + instrumentation: Option, /// Overrides `span_origin.environment`. - pub environment: Option, + environment: Option, } impl SpanOrigin { - /// Create an empty override (all defaults preserved). - pub fn new() -> Self { - Self::default() + /// Return the name override. + pub fn name(&self) -> Option<&str> { + self.name.as_deref() } - - /// Set the `span_origin.name` override. - pub fn name(mut self, name: impl Into) -> Self { - self.name = Some(name.into()); - self + /// Return the version override. + pub fn version(&self) -> Option<&str> { + self.version.as_deref() } - - /// Set the `span_origin.version` override. - pub fn version(mut self, version: impl Into) -> Self { - self.version = Some(version.into()); - self + /// Return the instrumentation name override. + pub fn instrumentation(&self) -> Option<&str> { + self.instrumentation.as_deref() } - - /// Set the `span_origin.instrumentation.name` override. - pub fn instrumentation(mut self, instrumentation: impl Into) -> Self { - self.instrumentation = Some(instrumentation.into()); - self - } - - /// Set the `span_origin.environment` override. - pub fn with_environment(mut self, environment: SpanOriginEnvironment) -> Self { - self.environment = Some(environment); - self + /// Return the environment override. + pub fn environment(&self) -> Option<&SpanOriginEnvironment> { + self.environment.as_ref() } } @@ -1500,10 +1526,11 @@ mod tests { span.log( SpanLog::builder() .span_origin( - SpanOrigin::new() + SpanOrigin::builder() .name("braintrust.plugin.codex") .version("9.9.9") - .instrumentation("codex-plugin"), + .instrumentation("codex-plugin") + .build(), ) .span_attribute("session_id", json!("sess-1")) .build() @@ -1536,7 +1563,11 @@ mod tests { // Builder-level default applies when the SpanLog does not override. let (builder, collector) = mock_span_builder(); let span = builder - .span_origin(SpanOrigin::new().name("braintrust.plugin.default")) + .span_origin( + SpanOrigin::builder() + .name("braintrust.plugin.default") + .build(), + ) .build(); span.log(SpanLog::builder().input(json!("x")).build().expect("build")); span.flush().await.expect("flush"); diff --git a/tests/api_client.rs b/tests/api_client.rs index f432244..6fe5622 100644 --- a/tests/api_client.rs +++ b/tests/api_client.rs @@ -466,26 +466,17 @@ async fn typed_dataset_mutations_use_the_data_plane_once_without_trace_fields() 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(), - }, + DatasetWrite::insert( + "dataset", + "row", + DatasetFields::builder().input(json!(null)).build(), + ), + DatasetWrite::update( + "dataset", + "row", + DatasetFields::builder().expected(json!(42)).build(), + ), + DatasetWrite::delete("dataset", "row"), ]; Mock::given(method("POST")) .and(path("/logs3")) diff --git a/tests/dataset_test.rs b/tests/dataset_test.rs index 8b16d23..12e122a 100644 --- a/tests/dataset_test.rs +++ b/tests/dataset_test.rs @@ -620,6 +620,7 @@ async fn mutation_errors_are_direct_and_do_not_replay_or_poison_later_operations Err(braintrust_sdk_rust::BraintrustError::Api { status: actual, message, + .. }) => { assert_eq!(actual, status); assert_eq!(message, "rejected"); diff --git a/tests/experiments.rs b/tests/experiments.rs index 9de77b4..bddd24d 100644 --- a/tests/experiments.rs +++ b/tests/experiments.rs @@ -292,12 +292,12 @@ async fn registration_failure_has_no_invented_id_and_clones_share_a_successful_r 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!(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.clone().export().await.unwrap().object_id(), "experiment" ); experiment @@ -468,7 +468,7 @@ async fn explicit_experiment_credentials_never_prime_the_shared_session() { 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_eq!(exported.object_id(), "experiment"); assert!(!state.is_logged_in()); } diff --git a/tests/span_lifecycle.rs b/tests/span_lifecycle.rs index 0510018..df810a8 100644 --- a/tests/span_lifecycle.rs +++ b/tests/span_lifecycle.rs @@ -220,7 +220,7 @@ async fn client_update_span_skip_span_origin_preserves_only_explicit_context() { .api_url(server.uri()) .app_url(server.uri()) .environment("ci", Some("test")) - .span_origin(SpanOrigin::new().name("client-origin")) + .span_origin(SpanOrigin::builder().name("client-origin").build()) .build() .await .expect("client"); @@ -256,7 +256,7 @@ async fn client_update_span_skip_span_origin_preserves_only_explicit_context() { .to_str(); let mut event = SpanLog::builder() .skip_span_origin() - .span_origin(SpanOrigin::new().name("event-origin")) + .span_origin(SpanOrigin::builder().name("event-origin").build()) .output(json!({"status": "updated"})); if let Some(context) = context { event = event.context(context.clone()); @@ -305,7 +305,7 @@ async fn span_handle_skip_span_origin_persists_through_logs_and_end() { .api_url(server.uri()) .app_url(server.uri()) .environment("ci", Some("test")) - .span_origin(SpanOrigin::new().name("client-origin")) + .span_origin(SpanOrigin::builder().name("client-origin").build()) .build() .await .expect("client"); @@ -319,10 +319,10 @@ async fn span_handle_skip_span_origin_persists_through_logs_and_end() { .parent_info(ParentSpanInfo::ProjectLogs { object_id: "proj-id".to_string(), }) - .span_origin(SpanOrigin::new().name("builder-origin")) + .span_origin(SpanOrigin::builder().name("builder-origin").build()) .build(); let mut event = SpanLog::builder() - .span_origin(SpanOrigin::new().name("event-origin")) + .span_origin(SpanOrigin::builder().name("event-origin").build()) .skip_span_origin() .input(json!("input")); if let Some(context) = context { From ba766a28d9b9d410c8c7c47ecbcdc946fe7b502b Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Thu, 29 Oct 2026 23:56:01 +0800 Subject: [PATCH 6/7] Allow partial dataset updates and require insert input --- src/dataset.rs | 187 ++++++++++++------------------------- src/lib.rs | 4 +- tests/credential_scopes.rs | 2 +- tests/dataset_test.rs | 124 +++++++++++++++++++++--- tests/resource_handles.rs | 2 +- 5 files changed, 173 insertions(+), 146 deletions(-) diff --git a/src/dataset.rs b/src/dataset.rs index 046685a..7a15cd7 100644 --- a/src/dataset.rs +++ b/src/dataset.rs @@ -19,7 +19,7 @@ //! DatasetInsert::builder() //! .input(json!({"question": "What is 2+2?"})) //! .expected(json!({"answer": "4"})) -//! .build()? +//! .build() //! ).await?; //! ``` @@ -32,7 +32,8 @@ use serde_json::{Map, Value}; use tokio::sync::OnceCell; use uuid::Uuid; -use crate::api::datasets::{DatasetFields, DatasetWrite}; +pub use crate::api::datasets::DatasetFields; +use crate::api::datasets::DatasetWrite; use crate::error::Result; use crate::internal::ports::{DatasetApi, DeliveryControl}; use crate::internal::session::CredentialScope; @@ -41,45 +42,61 @@ use crate::internal::session::CredentialScope; // Dataset Insert Types // ============================================================================ -/// Error type for DatasetInsert builder validation. -#[derive(Debug, Clone, PartialEq)] -#[non_exhaustive] -pub enum DatasetInsertBuilderError { - /// Input is required but was not provided. - MissingInput, -} - -impl fmt::Display for DatasetInsertBuilderError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::MissingInput => write!(f, "input is required but was not provided"), - } - } -} - -impl std::error::Error for DatasetInsertBuilderError {} - -/// Data to insert into a dataset. -#[derive(Clone, Default)] -#[non_exhaustive] +/// Data to insert into a dataset. Input is required, including when its value is JSON null. +/// +/// ``` +/// use braintrust_sdk_rust::DatasetInsert; +/// use serde_json::json; +/// let record = DatasetInsert::builder().input(json!(null)).build(); +/// assert_eq!(record.input(), &json!(null)); +/// ``` +/// +/// Missing input is rejected at compile time: +/// +/// ```compile_fail +/// use braintrust_sdk_rust::DatasetInsert; +/// let record = DatasetInsert::builder().expected(serde_json::json!(4)).build(); +/// ``` +/// +/// There is no default record that bypasses the required input: +/// +/// ```compile_fail +/// use braintrust_sdk_rust::DatasetInsert; +/// let record = DatasetInsert::default(); +/// ``` +#[derive(bon::Builder, Debug, Clone)] pub struct DatasetInsert { - pub(crate) input: Option, - pub(crate) expected: Option, - pub(crate) metadata: Option>, - pub(crate) tags: Option>, + #[builder(into)] + input: Value, + #[builder(into)] + expected: Option, + metadata: Option>, + tags: Option>, } impl DatasetInsert { - /// Create a new DatasetInsert builder. - pub fn builder() -> DatasetInsertBuilder { - DatasetInsertBuilder::new() + /// Return the input, including an explicit JSON null. + pub fn input(&self) -> &Value { + &self.input + } + /// Return the expected output, if supplied. + pub fn expected(&self) -> Option<&Value> { + self.expected.as_ref() + } + /// Return the metadata entries, if supplied. + pub fn metadata(&self) -> Option<&Map> { + self.metadata.as_ref() + } + /// Return the tags, if supplied. + pub fn tags(&self) -> Option<&[String]> { + self.tags.as_deref() } } impl From for DatasetFields { fn from(record: DatasetInsert) -> Self { Self { - input: record.input, + input: Some(record.input), expected: record.expected, metadata: record.metadata, tags: record.tags, @@ -87,69 +104,6 @@ impl From for DatasetFields { } } -/// Builder for DatasetInsert. -#[derive(Clone, Default)] -pub struct DatasetInsertBuilder { - inner: DatasetInsert, -} - -impl DatasetInsertBuilder { - /// Create a new DatasetInsertBuilder. - pub fn new() -> Self { - Self::default() - } - - /// Set the input data (required). - pub fn input(mut self, input: impl Into) -> Self { - self.inner.input = Some(input.into()); - self - } - - /// Set the expected output for evaluation. - pub fn expected(mut self, expected: impl Into) -> Self { - self.inner.expected = Some(expected.into()); - self - } - - /// Set the metadata map. - pub fn metadata(mut self, metadata: Map) -> Self { - self.inner.metadata = Some(metadata); - self - } - - /// Add a single metadata entry. - pub fn metadata_entry(mut self, key: impl Into, value: impl Into) -> Self { - self.inner - .metadata - .get_or_insert_with(Map::new) - .insert(key.into(), value.into()); - self - } - - /// Set multiple tags at once. - pub fn tags(mut self, tags: Vec) -> Self { - self.inner.tags = Some(tags); - self - } - - /// Add a single tag. - pub fn tag(mut self, tag: impl Into) -> Self { - self.inner - .tags - .get_or_insert_with(Vec::new) - .push(tag.into()); - self - } - - /// Build the DatasetInsert. - pub fn build(self) -> std::result::Result { - if self.inner.input.is_none() { - return Err(DatasetInsertBuilderError::MissingInput); - } - Ok(self.inner) - } -} - // ============================================================================ // Dataset Record // ============================================================================ @@ -462,13 +416,19 @@ impl Dataset { } /// Merge supplied fields into a record, returning after API acknowledgement. - pub async fn update(&self, id: &str, record: DatasetInsert) -> Result<()> { + /// Omitted fields remain unchanged; explicit JSON null values are sent as null. + /// + /// ```no_run + /// use braintrust_sdk_rust::{Dataset, DatasetFields}; + /// # async fn update(dataset: &Dataset) -> braintrust_sdk_rust::Result<()> { + /// dataset.update("record-id", DatasetFields::builder().expected(4).build()).await?; + /// # Ok(()) + /// # } + /// ``` + pub async fn update(&self, id: &str, fields: DatasetFields) -> Result<()> { let dataset_id = self.ensure_registered().await?.dataset_id.clone(); self.api - .write_dataset_record( - &self.scope, - DatasetWrite::update(dataset_id, id, record.into()), - ) + .write_dataset_record(&self.scope, DatasetWrite::update(dataset_id, id, fields)) .await } @@ -613,35 +573,6 @@ impl DatasetIterator { mod tests { use super::*; - #[test] - fn test_dataset_insert_builder_requires_input() { - let result = DatasetInsert::builder() - .expected(serde_json::json!({"answer": "4"})) - .build(); - - assert!(matches!( - result, - Err(DatasetInsertBuilderError::MissingInput) - )); - } - - #[test] - fn test_dataset_insert_builder_works() { - let result = DatasetInsert::builder() - .input(serde_json::json!({"question": "What is 2+2?"})) - .expected(serde_json::json!({"answer": "4"})) - .tag("math") - .tag("addition") - .metadata_entry("difficulty", serde_json::json!("easy")) - .build(); - - assert!(result.is_ok()); - let insert = result.unwrap(); - assert!(insert.input.is_some()); - assert!(insert.expected.is_some()); - assert_eq!(insert.tags.as_ref().map(|t| t.len()), Some(2)); - } - #[test] fn test_dataset_summary_getters() { let mut summary = DatasetSummary::new("project".to_string(), "dataset".to_string()); diff --git a/src/lib.rs b/src/lib.rs index 23379a3..fc8b10a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -23,8 +23,8 @@ mod types; pub use completion::{Completion, CompletionFuture}; pub use dataset::{ - Dataset, DatasetBuilder, DatasetBuilderError, DatasetInsert, DatasetInsertBuilder, - DatasetInsertBuilderError, DatasetIterator, DatasetRecord, DatasetSummary, + Dataset, DatasetBuilder, DatasetBuilderError, DatasetFields, DatasetInsert, + DatasetInsertBuilder, DatasetIterator, DatasetRecord, DatasetSummary, }; pub use delivery::{ DeliveryError, DeliveryReceipt, EnqueueError, ExecutionError, FlushError, FlushFailure, diff --git a/tests/credential_scopes.rs b/tests/credential_scopes.rs index b9a4c65..86ed7cd 100644 --- a/tests/credential_scopes.rs +++ b/tests/credential_scopes.rs @@ -11,7 +11,7 @@ fn registration() -> ResponseTemplate { ) } fn record() -> DatasetInsert { - DatasetInsert::builder().input(json!(1)).build().unwrap() + DatasetInsert::builder().input(json!(1)).build() } #[tokio::test] diff --git a/tests/dataset_test.rs b/tests/dataset_test.rs index 12e122a..4feb500 100644 --- a/tests/dataset_test.rs +++ b/tests/dataset_test.rs @@ -1,6 +1,6 @@ //! Integration tests for Datasets API. -use braintrust_sdk_rust::{BraintrustClient, DatasetInsert}; +use braintrust_sdk_rust::{BraintrustClient, DatasetFields, DatasetInsert}; use serde_json::json; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; @@ -133,8 +133,7 @@ async fn dataset_insert_is_acknowledged_without_flush() { DatasetInsert::builder() .input(json!({"question": "What is 2+2?"})) .expected(json!({"answer": "4"})) - .build() - .expect("record"), + .build(), ) .await .expect("insert"); @@ -525,10 +524,7 @@ async fn a_non_advancing_cursor_is_reported_and_does_not_poison_a_later_retry() } fn record(value: i32) -> DatasetInsert { - DatasetInsert::builder() - .input(json!(value)) - .build() - .unwrap() + DatasetInsert::builder().input(json!(value)).build() } #[tokio::test] @@ -553,14 +549,16 @@ async fn mutations_have_dataset_wire_shapes_and_share_registration_across_clones DatasetInsert::builder() .input(json!(1)) .expected(json!(2)) - .metadata_entry("source", "test") - .tag("sample") - .build() - .unwrap(), + .metadata(serde_json::Map::from_iter([( + "source".into(), + json!("test"), + )])) + .tags(vec!["sample".into()]) + .build(), ) .await .unwrap(); - dataset.clone().update(&id, record(3)).await.unwrap(); + dataset.clone().update(&id, record(3).into()).await.unwrap(); dataset.delete(&id).await.unwrap(); // Neither barrier registers or replays dataset work. dataset.flush().await.unwrap(); @@ -585,6 +583,104 @@ async fn mutations_have_dataset_wire_shapes_and_share_registration_across_clones } } +#[tokio::test] +async fn partial_updates_send_only_supplied_fields_and_preserve_null_and_empty_values() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/dataset/register")) + .respond_with(registration()) + .expect(1) + .mount(&server) + .await; + let patches = [ + ( + DatasetFields::builder().expected(4).build(), + json!({"expected":4}), + ), + ( + DatasetFields::builder() + .metadata(serde_json::Map::from_iter([( + "source".into(), + json!("reviewed"), + )])) + .build(), + json!({"metadata":{"source":"reviewed"}}), + ), + ( + DatasetFields::builder() + .tags(vec!["reviewed".into()]) + .build(), + json!({"tags":["reviewed"]}), + ), + ( + DatasetFields::builder().expected(json!(null)).build(), + json!({"expected":null}), + ), + ( + DatasetFields::builder().input(json!(null)).build(), + json!({"input":null}), + ), + ( + DatasetFields::builder().tags(vec![]).build(), + json!({"tags":[]}), + ), + ( + DatasetFields::builder() + .metadata(serde_json::Map::new()) + .build(), + json!({"metadata":{}}), + ), + ]; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(204)) + .expect(patches.len() as u64) + .mount(&server) + .await; + let client = isolated_client(&server).await; + let dataset = dataset(&client); + for (fields, expected_fields) in patches { + dataset.update("existing", fields).await.unwrap(); + let requests = server.received_requests().await.unwrap(); + let body: serde_json::Value = + serde_json::from_slice(&requests.last().unwrap().body).unwrap(); + let mut row = expected_fields; + row["dataset_id"] = json!("dataset"); + row["id"] = json!("existing"); + row["_is_merge"] = json!(true); + assert_eq!(body, json!({"api_version":2,"rows":[row]})); + } +} + +#[tokio::test] +async fn inserts_preserve_explicit_null_input() { + 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(1) + .mount(&server) + .await; + let client = isolated_client(&server).await; + let dataset = dataset(&client); + let id = dataset + .insert(DatasetInsert::builder().input(json!(null)).build()) + .await + .unwrap(); + let requests = server.received_requests().await.unwrap(); + let body: serde_json::Value = serde_json::from_slice(&requests.last().unwrap().body).unwrap(); + assert_eq!( + body, + json!({"api_version":2,"rows":[{"dataset_id":"dataset","id":id,"input":null}]}) + ); +} + #[tokio::test] async fn mutation_errors_are_direct_and_do_not_replay_or_poison_later_operations() { for status in [401, 413, 429, 503] { @@ -613,7 +709,7 @@ async fn mutation_errors_are_direct_and_do_not_replay_or_poison_later_operations 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 update = dataset.update("old", record(2).into()).await; let delete = dataset.delete("old").await; for result in [insert, update, delete] { match result { @@ -655,7 +751,7 @@ async fn registration_failure_prevents_mutation_and_flush_does_not_register_a_da 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.update("old", record(2).into()).await.is_err()); assert!(dataset.delete("old").await.is_err()); assert!(server .received_requests() diff --git a/tests/resource_handles.rs b/tests/resource_handles.rs index ee364cd..01956df 100644 --- a/tests/resource_handles.rs +++ b/tests/resource_handles.rs @@ -88,7 +88,7 @@ async fn resources_and_detached_span_builders_work_after_the_client_is_dropped() child.log(SpanLog::default()).await.unwrap(); resources .dataset - .insert(DatasetInsert::builder().input(json!(1)).build().unwrap()) + .insert(DatasetInsert::builder().input(json!(1)).build()) .await .unwrap(); let records: DatasetIterator = resources.dataset.fetch(None).await.unwrap(); From 2cde9ba6e17bc787d1d267bd3ee844d02c3004b5 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Fri, 30 Oct 2026 00:23:52 +0800 Subject: [PATCH 7/7] Simplify resource builders and transport boundaries --- src/dataset.rs | 184 ++++++-------------------------- src/experiments/builder.rs | 180 ------------------------------- src/experiments/experiment.rs | 69 ++++++++---- src/experiments/mod.rs | 6 +- src/experiments/span_builder.rs | 40 +++---- src/internal/ports.rs | 52 +-------- src/internal/resources.rs | 28 +++-- src/lib.rs | 10 +- src/logger.rs | 22 ++-- src/span.rs | 42 ++------ tests/credential_scopes.rs | 9 +- tests/dataset_test.rs | 78 +------------- tests/experiments.rs | 130 +++++++++++++++------- tests/resource_handles.rs | 8 +- 14 files changed, 234 insertions(+), 624 deletions(-) delete mode 100644 src/experiments/builder.rs diff --git a/src/dataset.rs b/src/dataset.rs index 7a15cd7..e12b329 100644 --- a/src/dataset.rs +++ b/src/dataset.rs @@ -13,7 +13,7 @@ //! .await? //! .project_name("my-project") //! .dataset_name("test-cases") -//! .build()?; +//! .build(); //! //! let id = dataset.insert( //! DatasetInsert::builder() @@ -35,7 +35,8 @@ use uuid::Uuid; pub use crate::api::datasets::DatasetFields; use crate::api::datasets::DatasetWrite; use crate::error::Result; -use crate::internal::ports::{DatasetApi, DeliveryControl}; +use crate::internal::ports::DeliveryControl; +use crate::internal::resources::ResourceTransport; use crate::internal::session::CredentialScope; // ============================================================================ @@ -193,169 +194,44 @@ pub(crate) struct DatasetMetadata { use crate::api::btql::{BTQLQuery, BTQLQueryInner}; use crate::api::datasets::DatasetRegisterRequest; -// ============================================================================ -// DatasetBuilder Error Types -// ============================================================================ - -/// Error type for DatasetBuilder validation. -#[derive(Debug, Clone, PartialEq)] -#[non_exhaustive] -pub enum DatasetBuilderError { - /// Project name is required but was not provided. - MissingProjectName, - /// Dataset name is required but was not provided. - MissingDatasetName, -} - -impl fmt::Display for DatasetBuilderError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::MissingProjectName => { - write!(f, "project_name is required but was not provided") - } - Self::MissingDatasetName => { - write!(f, "dataset_name is required but was not provided") - } - } - } -} - -impl std::error::Error for DatasetBuilderError {} - -// ============================================================================ -// DatasetBuilder -// ============================================================================ - -/// Builder for creating datasets with lazy registration. -pub struct DatasetBuilder { - api: Arc, - lifecycle: Arc, - scope: CredentialScope, - project_name: Option, - dataset_name: Option, - description: Option, - metadata: Option>, -} - -impl Clone for DatasetBuilder { - 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(), - } - } -} - -impl DatasetBuilder { - /// Create a new DatasetBuilder (internal use). - pub(crate) fn new( - api: Arc, - lifecycle: Arc, - scope: CredentialScope, - ) -> Self { - Self { - api, - lifecycle, - scope, - project_name: None, - dataset_name: None, - description: None, - metadata: None, - } - } - - /// Set the project name (required). - pub fn project_name(mut self, project_name: impl Into) -> Self { - self.project_name = Some(project_name.into()); - self - } - - /// Set the dataset name (required). - pub fn dataset_name(mut self, dataset_name: impl Into) -> Self { - self.dataset_name = Some(dataset_name.into()); - self - } - - /// Set the dataset description. - pub fn description(mut self, description: impl Into) -> Self { - self.description = Some(description.into()); - self - } - - /// Set the dataset metadata. - pub fn metadata(mut self, metadata: Map) -> Self { - self.metadata = Some(metadata); - self - } - - /// Add a single metadata entry. - pub fn metadata_entry(mut self, key: impl Into, value: impl Into) -> Self { - self.metadata - .get_or_insert_with(Map::new) - .insert(key.into(), value.into()); - self - } - - /// Build the dataset (lazy registration). - pub fn build(self) -> std::result::Result { - let project_name = self - .project_name - .ok_or(DatasetBuilderError::MissingProjectName)?; - let dataset_name = self - .dataset_name - .ok_or(DatasetBuilderError::MissingDatasetName)?; - - Ok(Dataset { - api: self.api, - lifecycle: self.lifecycle, - scope: self.scope, - project_name, - dataset_name, - description: self.description, - metadata: self.metadata, - lazy_metadata: Arc::new(OnceCell::new()), - }) - } -} - -// ============================================================================ -// Dataset -// ============================================================================ - /// A dataset for storing test/evaluation data. /// /// Datasets are lazily registered with the API on first use. +/// The builder requires both a project name and a dataset name: +/// +/// ```compile_fail +/// # fn missing_project(client: &braintrust_sdk_rust::BraintrustClient) { +/// client.dataset_builder_with_credentials("token", "tenant").dataset_name("d").build(); +/// # } +/// ``` +/// +/// ```compile_fail +/// # fn missing_dataset(client: &braintrust_sdk_rust::BraintrustClient) { +/// client.dataset_builder_with_credentials("token", "tenant").project_name("p").build(); +/// # } +/// ``` +#[derive(bon::Builder, Clone)] +#[builder(start_fn(name = new, vis = "pub(crate)"), derive(Clone))] pub struct Dataset { - api: Arc, + #[builder(start_fn)] + api: ResourceTransport, + #[builder(start_fn)] lifecycle: Arc, + #[builder(start_fn)] scope: CredentialScope, + /// Project containing the dataset (required). + #[builder(into)] project_name: String, + /// Dataset name (required). + #[builder(into)] dataset_name: String, + #[builder(into)] description: Option, metadata: Option>, + #[builder(skip)] lazy_metadata: Arc>, } -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") @@ -498,7 +374,7 @@ impl Dataset { /// # } /// ``` pub struct DatasetIterator { - fetcher: Arc, + fetcher: ResourceTransport, scope: CredentialScope, dataset_id: String, batch_size: usize, @@ -509,7 +385,7 @@ pub struct DatasetIterator { impl DatasetIterator { fn new( - fetcher: Arc, + fetcher: ResourceTransport, scope: CredentialScope, dataset_id: String, batch_size: usize, diff --git a/src/experiments/builder.rs b/src/experiments/builder.rs deleted file mode 100644 index fe97b25..0000000 --- a/src/experiments/builder.rs +++ /dev/null @@ -1,180 +0,0 @@ -//! Builder for creating experiments. - -use std::fmt; -use std::sync::Arc; - -use serde_json::{Map, Value}; - -use crate::experiments::experiment::{Experiment, ExperimentState}; -use crate::experiments::metadata::RepoInfo; -use crate::internal::ports::{ExperimentApi, SpanSubmitter}; -use crate::internal::session::CredentialScope; - -/// Error type for ExperimentBuilder validation. -#[derive(Debug, Clone, PartialEq)] -#[non_exhaustive] -pub enum ExperimentBuilderError { - /// Project name is required but was not provided. - MissingProjectName, -} - -impl fmt::Display for ExperimentBuilderError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::MissingProjectName => { - write!(f, "project_name is required but was not provided") - } - } - } -} - -impl std::error::Error for ExperimentBuilderError {} - -/// Builder for creating experiments with lazy registration. -/// -/// 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, - base_experiment: Option, - metadata: Option>, - public: Option, - update: Option, - repo_info: Option, -} - -impl Clone for ExperimentBuilder { - fn clone(&self) -> Self { - Self { - 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(), - base_experiment: self.base_experiment.clone(), - metadata: self.metadata.clone(), - public: self.public, - update: self.update, - repo_info: self.repo_info.clone(), - } - } -} - -impl ExperimentBuilder { - /// Create a new ExperimentBuilder (internal use). - pub(crate) fn new( - api: Arc, - submitter: Arc, - scope: CredentialScope, - ) -> Self { - Self { - api, - submitter, - scope, - project_name: None, - experiment_name: None, - description: None, - base_experiment: None, - metadata: None, - public: None, - update: None, - repo_info: None, - } - } - - /// Set the organization name. - pub fn org_name(mut self, org_name: impl Into) -> Self { - self.scope = self.scope.with_org_name(org_name.into()); - self - } - - /// Set the project name (required). - pub fn project_name(mut self, project_name: impl Into) -> Self { - self.project_name = Some(project_name.into()); - self - } - - /// Set the experiment name. If not provided, a name will be auto-generated. - pub fn experiment_name(mut self, experiment_name: impl Into) -> Self { - self.experiment_name = Some(experiment_name.into()); - self - } - - /// Set the experiment description. - pub fn description(mut self, description: impl Into) -> Self { - self.description = Some(description.into()); - self - } - - /// Set the base experiment name for comparison. - pub fn base_experiment(mut self, base_experiment: impl Into) -> Self { - self.base_experiment = Some(base_experiment.into()); - self - } - - /// Set the experiment metadata. - pub fn metadata(mut self, metadata: Map) -> Self { - self.metadata = Some(metadata); - self - } - - /// Add a single metadata entry. - pub fn metadata_entry(mut self, key: impl Into, value: impl Into) -> Self { - self.metadata - .get_or_insert_with(Map::new) - .insert(key.into(), value.into()); - self - } - - /// Set whether the experiment is public. - pub fn public(mut self, public: bool) -> Self { - self.public = Some(public); - self - } - - /// Continue logging to an existing experiment instead of creating a new one. - /// - /// When set to `true`, if an experiment with the same name already exists, - /// logging will continue to that experiment. If `false` (the default), - /// a new experiment will always be created. - pub fn update(mut self, update: bool) -> Self { - self.update = Some(update); - self - } - - /// Set explicit repository information. - /// - /// This overrides any automatic git metadata collection. - pub fn repo_info(mut self, repo_info: RepoInfo) -> Self { - self.repo_info = Some(repo_info); - self - } - - /// Build the experiment. Returns an error if required fields are missing. - 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, - scope: self.scope, - project_name, - experiment_name: self.experiment_name, - description: self.description, - base_experiment: self.base_experiment, - metadata: self.metadata, - public: self.public, - update: self.update, - repo_info: self.repo_info, - state: Arc::new(ExperimentState::default()), - }) - } -} diff --git a/src/experiments/experiment.rs b/src/experiments/experiment.rs index 02c33cb..3dae615 100644 --- a/src/experiments/experiment.rs +++ b/src/experiments/experiment.rs @@ -17,7 +17,8 @@ use crate::experiments::metadata::{ }; use crate::experiments::span_builder::ExperimentSpanBuilder; use crate::experiments::summary::{ExperimentSummary, MetricSummary, ScoreSummary}; -use crate::internal::ports::{ExperimentApi, SpanSubmitter}; +use crate::internal::ports::SpanSubmitter; +use crate::internal::resources::ResourceTransport; use crate::internal::session::CredentialScope; use crate::span::{SpanBuilder, SpanHandle, SpanLog}; use crate::types::{ParentSpanInfo, SpanAttributes, SpanPayload, SpanType}; @@ -33,8 +34,16 @@ pub(crate) struct ExperimentMetadata { /// Handle to a Braintrust experiment. /// /// Use this to log evaluation events, create spans, log feedback, and get summaries. +/// The builder requires a project name before it can build: +/// +/// ```compile_fail +/// # fn missing_project(client: &braintrust_sdk_rust::BraintrustClient) { +/// client.experiment_builder_with_credentials("token", "org").experiment_name("e").build(); +/// # } +/// ``` +#[derive(Clone)] pub struct Experiment { - pub(crate) api: Arc, + pub(crate) api: ResourceTransport, pub(crate) submitter: Arc, pub(crate) scope: CredentialScope, pub(crate) project_name: String, @@ -57,26 +66,48 @@ pub(crate) struct ExperimentState { called_start_span: AtomicBool, } -impl Clone for Experiment { - fn clone(&self) -> Self { +#[bon::bon] +impl Experiment { + /// Create an experiment with shared lazy registration and immutable credentials. + #[builder(start_fn(name = builder, vis = "pub(crate)"), builder_type(vis = "pub"), derive(Clone))] + pub(crate) fn new( + #[builder(start_fn)] api: ResourceTransport, + #[builder(start_fn)] submitter: Arc, + #[builder(start_fn)] scope: CredentialScope, + #[builder(into)] project_name: String, + #[builder(into)] experiment_name: Option, + #[builder(into)] description: Option, + #[builder(into)] base_experiment: Option, + metadata: Option>, + public: Option, + /// Continue an existing experiment with this name when true; otherwise create a new one. + update: Option, + /// Explicit repository metadata for this experiment. + repo_info: Option, + /// Organization header override for this resource, without changing the client session. + #[builder(into)] + org_name: Option, + ) -> Self { + let scope = match org_name { + Some(name) => scope.with_org_name(name), + None => scope, + }; Self { - 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(), - base_experiment: self.base_experiment.clone(), - update: self.update, - repo_info: self.repo_info.clone(), - metadata: self.metadata.clone(), - public: self.public, - state: self.state.clone(), + api, + submitter, + scope, + project_name, + experiment_name, + description, + base_experiment, + metadata, + public, + update, + repo_info, + state: Arc::new(ExperimentState::default()), } } -} -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()`. @@ -136,7 +167,7 @@ impl Experiment { 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())) + Ok(ExperimentSpanBuilder::new(self, &metadata.experiment_id)) } /// Log feedback on an existing span/record. diff --git a/src/experiments/mod.rs b/src/experiments/mod.rs index fa361e8..1865b2b 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() @@ -27,7 +27,6 @@ //! ``` // Public modules -mod builder; mod event_log; mod experiment; mod feedback; @@ -36,9 +35,8 @@ mod span_builder; mod summary; // Public re-exports -pub use builder::{ExperimentBuilder, ExperimentBuilderError}; pub use event_log::{ExperimentLog, ExperimentLogBuilder}; -pub use experiment::Experiment; +pub use experiment::{Experiment, ExperimentBuilder}; pub use feedback::{Feedback, FeedbackBuilder, FeedbackBuilderError}; pub use metadata::{BaseExperimentInfo, ExportedExperiment, ProjectMetadata, RepoInfo}; pub use span_builder::ExperimentSpanBuilder; diff --git a/src/experiments/span_builder.rs b/src/experiments/span_builder.rs index 1f04a6d..f17c655 100644 --- a/src/experiments/span_builder.rs +++ b/src/experiments/span_builder.rs @@ -1,56 +1,40 @@ //! Builder for creating spans within experiments. -use std::sync::Arc; - -use crate::experiments::experiment::{Experiment, ExperimentMetadata}; -use crate::internal::{ports::SpanSubmitter, session::CredentialScope}; +use crate::experiments::experiment::Experiment; use crate::span::{SpanBuilder, SpanHandle}; use crate::types::{ParentSpanInfo, SpanType}; /// Builder for creating spans within an experiment. +#[derive(Clone)] pub struct ExperimentSpanBuilder { - submitter: Arc, - scope: CredentialScope, - metadata: ExperimentMetadata, - name: Option, - span_type: SpanType, + inner: SpanBuilder, } impl ExperimentSpanBuilder { - pub(crate) fn new(experiment: &Experiment, metadata: ExperimentMetadata) -> Self { + pub(crate) fn new(experiment: &Experiment, experiment_id: &str) -> Self { Self { - submitter: experiment.submitter.clone(), - scope: experiment.scope.clone(), - metadata, - name: None, - span_type: SpanType::Eval, + inner: SpanBuilder::new(experiment.submitter.clone(), experiment.scope.clone()) + .parent_info(ParentSpanInfo::Experiment { + object_id: experiment_id.to_owned(), + }) + .span_type(SpanType::Eval), } } /// Set the span name. pub fn name(mut self, name: impl Into) -> Self { - self.name = Some(name.into()); + self.inner = self.inner.name(name); self } /// Set the span type. pub fn span_type(mut self, span_type: SpanType) -> Self { - self.span_type = span_type; + self.inner = self.inner.span_type(span_type); self } /// Build the span. 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 builder = SpanBuilder::new(self.submitter, self.scope) - .parent_info(parent_info) - .span_type(self.span_type); - - builder.build() + self.inner.build() } } diff --git a/src/internal/ports.rs b/src/internal/ports.rs index a5ba29b..02dfc62 100644 --- a/src/internal/ports.rs +++ b/src/internal/ports.rs @@ -1,17 +1,7 @@ -//! Private dependency-injection boundaries for resource I/O and delivery. +//! Private dependency-injection boundaries for span 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::types::{ParentSpanInfo, SpanPayload}; @@ -31,43 +21,3 @@ pub(crate) trait SpanSubmitter: DeliveryControl { parent_info: 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 index 557aec6..44e6652 100644 --- a/src/internal/resources.rs +++ b/src/internal/resources.rs @@ -1,6 +1,5 @@ -//! Resource ports adapt SDK policy to the shared one-attempt API contracts. +//! Concrete resource HTTP adapter over the shared one-attempt API contracts. use crate::internal::session::CredentialScope; -use async_trait::async_trait; use crate::api::{ btql::{BTQLQuery, BTQLResponse}, @@ -14,7 +13,6 @@ use crate::api::{ }; 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)] @@ -28,11 +26,8 @@ impl ResourceTransport { fn resource_api(&self, scope: &CredentialScope) -> Result { scope.api(&self.http) } -} -#[async_trait] -impl ExperimentApi for ResourceTransport { - async fn register_experiment( + pub async fn register_experiment( &self, scope: &CredentialScope, request: ExperimentRegisterRequest, @@ -44,7 +39,7 @@ impl ExperimentApi for ResourceTransport { .map_err(Into::into) } - async fn fetch_base_experiment( + pub async fn fetch_base_experiment( &self, scope: &CredentialScope, experiment_id: &str, @@ -67,7 +62,7 @@ impl ExperimentApi for ResourceTransport { }) } - async fn fetch_experiment_comparison( + pub async fn fetch_experiment_comparison( &self, scope: &CredentialScope, experiment_id: &str, @@ -79,11 +74,8 @@ impl ExperimentApi for ResourceTransport { .await .map_err(Into::into) } -} -#[async_trait] -impl DatasetApi for ResourceTransport { - async fn register_dataset( + pub async fn register_dataset( &self, scope: &CredentialScope, request: DatasetRegisterRequest, @@ -95,7 +87,7 @@ impl DatasetApi for ResourceTransport { .map_err(Into::into) } - async fn fetch_dataset_records( + pub async fn fetch_dataset_records( &self, scope: &CredentialScope, query: BTQLQuery, @@ -107,7 +99,7 @@ impl DatasetApi for ResourceTransport { .map_err(Into::into) } - async fn fetch_dataset_summary( + pub async fn fetch_dataset_summary( &self, scope: &CredentialScope, dataset_id: &str, @@ -119,7 +111,11 @@ impl DatasetApi for ResourceTransport { .map_err(Into::into) } - async fn write_dataset_record(&self, scope: &CredentialScope, row: DatasetWrite) -> Result<()> { + pub async fn write_dataset_record( + &self, + scope: &CredentialScope, + row: DatasetWrite, + ) -> Result<()> { self.resource_api(scope)? .datasets() .write(&[row]) diff --git a/src/lib.rs b/src/lib.rs index fc8b10a..44ed86b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -23,8 +23,8 @@ mod types; pub use completion::{Completion, CompletionFuture}; pub use dataset::{ - Dataset, DatasetBuilder, DatasetBuilderError, DatasetFields, DatasetInsert, - DatasetInsertBuilder, DatasetIterator, DatasetRecord, DatasetSummary, + Dataset, DatasetBuilder, DatasetFields, DatasetInsert, DatasetInsertBuilder, DatasetIterator, + DatasetRecord, DatasetSummary, }; pub use delivery::{ DeliveryError, DeliveryReceipt, EnqueueError, ExecutionError, FlushError, FlushFailure, @@ -32,9 +32,9 @@ pub use delivery::{ }; pub use error::{BraintrustError, Result}; pub use experiments::{ - BaseExperimentInfo, Experiment, ExperimentBuilder, ExperimentBuilderError, ExperimentLog, - ExperimentLogBuilder, ExperimentSpanBuilder, ExperimentSummary, ExportedExperiment, Feedback, - FeedbackBuilder, FeedbackBuilderError, MetricSummary, ProjectMetadata, RepoInfo, ScoreSummary, + BaseExperimentInfo, Experiment, ExperimentBuilder, ExperimentLog, ExperimentLogBuilder, + ExperimentSpanBuilder, ExperimentSummary, ExportedExperiment, Feedback, FeedbackBuilder, + FeedbackBuilderError, MetricSummary, ProjectMetadata, RepoInfo, ScoreSummary, }; pub use extractors::{extract_anthropic_usage, extract_openai_usage}; pub use log_queue::LogQueueConfig; diff --git a/src/logger.rs b/src/logger.rs index 0f74159..4184ad8 100644 --- a/src/logger.rs +++ b/src/logger.rs @@ -7,9 +7,9 @@ use serde_json::{Map, Value}; use tokio::sync::Notify; use url::Url; -use crate::dataset::DatasetBuilder; +use crate::dataset::{Dataset, DatasetBuilder}; use crate::error::{BraintrustError, Result}; -use crate::experiments::ExperimentBuilder; +use crate::experiments::{Experiment, ExperimentBuilder}; use crate::http::{build_http_client, CaBundleConfig}; use crate::internal::ports::SpanSubmitter; use crate::internal::resources::ResourceTransport; @@ -462,8 +462,8 @@ impl std::fmt::Debug for ClientInner { } impl BraintrustClient { - fn resource_transport(&self) -> Arc { - Arc::new(ResourceTransport::new(self.inner.http_client.clone())) + fn resource_transport(&self) -> ResourceTransport { + ResourceTransport::new(self.inner.http_client.clone()) } fn span_submitter(&self) -> Arc { @@ -574,7 +574,7 @@ impl BraintrustClient { 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?; - Ok(DatasetBuilder::new( + Ok(Dataset::new( self.resource_transport(), Arc::new(self.inner.queue.clone()), state.scope()?, @@ -588,7 +588,7 @@ impl BraintrustClient { token: impl Into, org_name: impl Into, ) -> DatasetBuilder { - DatasetBuilder::new( + Dataset::new( self.resource_transport(), Arc::new(self.inner.queue.clone()), self.explicit_scope(token.into(), String::new(), Some(org_name.into())), @@ -930,7 +930,7 @@ impl BraintrustClient { /// .experiment_builder_with_credentials(&api_key, &org_id) /// .project_name("my-project") /// .experiment_name("baseline") - /// .build()?; + /// .build(); /// ``` pub fn experiment_builder_with_credentials( &self, @@ -941,7 +941,7 @@ impl BraintrustClient { let org_id = org_id.into(); let submitter = self.span_submitter(); - ExperimentBuilder::new( + Experiment::builder( self.resource_transport(), submitter, self.explicit_scope(token, org_id, None), @@ -1167,7 +1167,7 @@ mod tests { } #[tokio::test] - async fn resource_ports_do_not_retain_the_whole_client_or_login_state() { + async fn resource_dependencies_do_not_retain_the_whole_client_or_login_state() { let client = BraintrustClient::builder() .skip_login(true) .build() @@ -1188,8 +1188,8 @@ mod tests { "resource dependencies must not retain ClientInner" ); let _span = span.build(); - let _dataset = dataset.build().unwrap(); - let _experiment = experiment.build().unwrap(); + let _dataset = dataset.build(); + let _experiment = experiment.build(); } #[tokio::test] diff --git a/src/span.rs b/src/span.rs index a8133a7..903ea9c 100644 --- a/src/span.rs +++ b/src/span.rs @@ -250,7 +250,9 @@ impl SpanLogBuilder { } } +#[derive(Clone)] pub struct SpanBuilder { + name: Option, submitter: Arc, scope: CredentialScope, project_name: Option, @@ -275,25 +277,6 @@ impl fmt::Debug for SpanBuilder { } } -impl Clone for SpanBuilder { - fn clone(&self) -> Self { - Self { - submitter: Arc::clone(&self.submitter), - 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(), - start_time_override: self.start_time_override, - environment: self.environment.clone(), - span_id_override: self.span_id_override.clone(), - row_id_override: self.row_id_override.clone(), - span_origin: self.span_origin.clone(), - span_attributes_extra: self.span_attributes_extra.clone(), - } - } -} - /// Small options bag identifying the environment that produced a span. /// /// Use the constructor so future options can be added: @@ -389,8 +372,15 @@ impl SpanOrigin { } impl SpanBuilder { + /// Set the span name. + pub fn name(mut self, name: impl Into) -> Self { + self.name = Some(name.into()); + self + } + pub(crate) fn new(submitter: Arc, scope: CredentialScope) -> Self { Self { + name: None, submitter, scope, project_name: None, @@ -534,6 +524,7 @@ impl SpanBuilder { parent_info: self.parent_info, row_id: row_id.clone(), inner: Arc::new(Mutex::new(SpanData { + name: self.name, row_id, span_id, has_flushed: false, @@ -553,6 +544,7 @@ impl SpanBuilder { } } +#[derive(Clone)] pub struct SpanHandle { submitter: Arc, scope: CredentialScope, @@ -562,18 +554,6 @@ pub struct SpanHandle { inner: Arc>, } -impl Clone for SpanHandle { - fn clone(&self) -> Self { - Self { - submitter: Arc::clone(&self.submitter), - scope: self.scope.clone(), - parent_info: self.parent_info.clone(), - row_id: self.row_id.clone(), - inner: Arc::clone(&self.inner), - } - } -} - impl SpanHandle { /// Return the stable row ID for this span, usable with `log_feedback()` and `update_span()`. pub fn row_id(&self) -> &str { diff --git a/tests/credential_scopes.rs b/tests/credential_scopes.rs index 86ed7cd..712c318 100644 --- a/tests/credential_scopes.rs +++ b/tests/credential_scopes.rs @@ -110,14 +110,12 @@ async fn logged_in_resources_capture_the_org_route_while_explicit_resources_keep .unwrap() .project_name("p") .dataset_name("d") - .build() - .unwrap(); + .build(); let explicit = client .dataset_builder_with_credentials("explicit", "explicit-tenant") .project_name("p") .dataset_name("d") - .build() - .unwrap(); + .build(); for dataset in [session.clone(), explicit] { dataset.insert(record()).await.unwrap(); assert_eq!( @@ -330,8 +328,7 @@ async fn experiment_clones_and_child_spans_preserve_the_explicit_scope() { .experiment_builder_with_credentials("explicit", "org") .org_name("tenant") .project_name("p") - .build() - .unwrap(); + .build(); let span = experiment.start_span().await.unwrap().build(); span.log(SpanLog::default()).await.unwrap(); experiment.clone().summarize().await.unwrap(); diff --git a/tests/dataset_test.rs b/tests/dataset_test.rs index 4feb500..7734478 100644 --- a/tests/dataset_test.rs +++ b/tests/dataset_test.rs @@ -14,74 +14,6 @@ fn mock_login_response(orgs: &[(&str, &str)]) -> ResponseTemplate { ResponseTemplate::new(200).set_body_json(json!({ "org_info": org_info })) } -#[tokio::test] -async fn dataset_builder_requires_project_name() { - let server = MockServer::start().await; - - Mock::given(method("POST")) - .and(path("/api/apikey/login")) - .respond_with(mock_login_response(&[("org-id", "Test Org")])) - .mount(&server) - .await; - - let client = BraintrustClient::builder() - .api_key("test-key") - .app_url(server.uri()) - .api_url(server.uri()) - .blocking_login(true) - .build() - .await - .expect("client"); - - let result = client - .dataset_builder() - .await - .unwrap() - .dataset_name("test-dataset") - .build(); - - assert!(result.is_err()); - let err = result.unwrap_err(); - assert_eq!( - err.to_string(), - "project_name is required but was not provided" - ); -} - -#[tokio::test] -async fn dataset_builder_requires_dataset_name() { - let server = MockServer::start().await; - - Mock::given(method("POST")) - .and(path("/api/apikey/login")) - .respond_with(mock_login_response(&[("org-id", "Test Org")])) - .mount(&server) - .await; - - let client = BraintrustClient::builder() - .api_key("test-key") - .app_url(server.uri()) - .api_url(server.uri()) - .blocking_login(true) - .build() - .await - .expect("client"); - - let result = client - .dataset_builder() - .await - .unwrap() - .project_name("test-project") - .build(); - - assert!(result.is_err()); - let err = result.unwrap_err(); - assert_eq!( - err.to_string(), - "dataset_name is required but was not provided" - ); -} - #[tokio::test] async fn dataset_insert_is_acknowledged_without_flush() { let server = MockServer::start().await; @@ -124,8 +56,7 @@ async fn dataset_insert_is_acknowledged_without_flush() { .unwrap() .project_name("test-project") .dataset_name("test-dataset") - .build() - .expect("dataset builder"); + .build(); // Insert a record let record_id = dataset @@ -210,8 +141,7 @@ async fn dataset_fetch_returns_records() { .unwrap() .project_name("test-project") .dataset_name("test-dataset") - .build() - .expect("dataset builder"); + .build(); // Fetch records let iter = dataset.fetch(Some(10)).await.expect("fetch"); @@ -273,8 +203,7 @@ async fn dataset_summarize_returns_summary() { .unwrap() .project_name("test-project") .dataset_name("test-dataset") - .build() - .expect("dataset builder"); + .build(); let summary = dataset.summarize().await.expect("summarize"); @@ -303,7 +232,6 @@ fn dataset(client: &BraintrustClient) -> braintrust_sdk_rust::Dataset { .project_name("project") .dataset_name("dataset") .build() - .unwrap() } #[tokio::test] diff --git a/tests/experiments.rs b/tests/experiments.rs index bddd24d..617560e 100644 --- a/tests/experiments.rs +++ b/tests/experiments.rs @@ -35,8 +35,7 @@ async fn experiment_log_flushes_to_logs_endpoint() { .experiment_builder_with_credentials("token", "org-id") .project_name("test-project") .experiment_name("test-experiment") - .build() - .expect("build experiment"); + .build(); let _id = experiment .log( @@ -71,24 +70,6 @@ async fn experiment_log_flushes_to_logs_endpoint() { assert_eq!(row["experiment_id"], "exp-id"); } -#[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() - .await - .expect("client"); - - let result = client - .experiment_builder_with_credentials("token", "org-id") - .experiment_name("test-experiment") - .build(); - - assert!(result.is_err()); -} - #[tokio::test] async fn experiment_log_feedback_sends_merge_event() { let server = MockServer::start().await; @@ -120,8 +101,7 @@ async fn experiment_log_feedback_sends_merge_event() { let experiment = client .experiment_builder_with_credentials("token", "org-id") .project_name("test-project") - .build() - .expect("build experiment"); + .build(); // Log an event let id = experiment @@ -213,8 +193,7 @@ async fn experiment_summarize_returns_summary() { .experiment_builder_with_credentials("token", "org-id") .project_name("test-project") .experiment_name("my-experiment") - .build() - .expect("build experiment"); + .build(); // Log an event first experiment @@ -285,8 +264,7 @@ async fn registration_failure_has_no_invented_id_and_clones_share_a_successful_r let experiment = client .experiment_builder_with_credentials("token", "org") .project_name("project") - .build() - .unwrap(); + .build(); let clone = experiment.clone(); assert_api_error(experiment.project().await, 503); assert_eq!(experiment.experiment_id(), None); @@ -336,8 +314,7 @@ async fn every_registration_dependent_operation_returns_the_real_failure() { let experiment = client .experiment_builder_with_credentials("token", "org") .project_name("project") - .build() - .unwrap(); + .build(); assert_api_error(experiment.start_span().await, 400); assert_api_error( experiment @@ -394,8 +371,7 @@ async fn event_feedback_and_update_methods_propagate_delivery_rejection() { let experiment = client .experiment_builder_with_credentials("token", "org") .project_name("project") - .build() - .unwrap(); + .build(); assert_api_error( experiment .log(ExperimentLog::builder().input(json!("value")).build()) @@ -441,8 +417,7 @@ async fn comparison_setup_errors_are_not_silently_treated_as_no_base_experiment( let experiment = client .experiment_builder_with_credentials("token", "org") .project_name("project") - .build() - .unwrap(); + .build(); assert_api_error(experiment.summarize().await, 401); } @@ -460,8 +435,7 @@ async fn explicit_experiment_credentials_never_prime_the_shared_session() { let experiment = client .experiment_builder_with_credentials("private-token", "tenant") .project_name("project") - .build() - .unwrap(); + .build(); assert!(!state.is_logged_in()); experiment.project().await.unwrap(); assert!(!state.is_logged_in()); @@ -491,8 +465,7 @@ async fn traced_propagates_final_flush_errors_instead_of_discarding_them() { let experiment = client .experiment_builder_with_credentials("token", "org") .project_name("project") - .build() - .unwrap(); + .build(); assert_api_error(experiment.traced(|_| async { 42 }).await, 401); } @@ -509,8 +482,7 @@ async fn clones_share_the_selected_span_recording_interface() { let experiment = client .experiment_builder_with_credentials("token", "org") .project_name("project") - .build() - .unwrap(); + .build(); let clone = experiment.clone(); experiment.start_span().await.unwrap(); assert!(matches!( @@ -557,8 +529,7 @@ async fn experiment_resource_adapters_use_the_control_plane_and_preserve_no_base let experiment = client .experiment_builder_with_credentials("explicit", "org") .project_name("project") - .build() - .unwrap(); + .build(); experiment.summarize().await.unwrap(); let requests = control.received_requests().await.unwrap(); assert_eq!(requests.len(), 3); @@ -568,3 +539,82 @@ async fn experiment_resource_adapters_use_the_control_plane_and_preserve_no_base .all(|(key, _)| key != "base_experiment_id")); assert!(data.received_requests().await.unwrap().is_empty()); } + +#[tokio::test] +async fn experiment_builders_preserve_options_and_named_span_configuration() { + use braintrust_sdk_rust::{RepoInfo, SpanLog, SpanType}; + use wiremock::matchers::{body_json, header}; + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/experiment/register")) + .and(header("x-bt-org-name", "tenant")) + .and(body_json(json!({ + "project_name":"project", "org_id":"org", "experiment_name":"named", + "description":"description", "base_experiment":"baseline", + "metadata":{"source":"test"}, "public":true, "update":true, + "repo_info":{"commit":"abc"} + }))) + .respond_with(registered_experiment()) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .and(header("x-bt-org-name", "tenant")) + .respond_with(ResponseTemplate::new(204)) + .expect(2) + .mount(&server) + .await; + let client = isolated_client(&server).await; + let builder = client + .experiment_builder_with_credentials("token", "org") + .project_name("project") + .experiment_name("named") + .description("description") + .base_experiment("baseline") + .metadata(serde_json::Map::from_iter([( + "source".into(), + json!("test"), + )])) + .public(true) + .update(true) + .repo_info(RepoInfo::builder().commit("abc").build()) + .org_name("tenant"); + let experiment = builder.clone().build(); + experiment + .start_span() + .await + .unwrap() + .name("evaluation") + .build() + .log(SpanLog::default()) + .await + .unwrap(); + experiment + .clone() + .start_span() + .await + .unwrap() + .name("custom") + .span_type(SpanType::Task) + .build() + .log(SpanLog::default()) + .await + .unwrap(); + + let requests = server.received_requests().await.unwrap(); + let rows: Vec = requests + .iter() + .filter(|r| r.url.path() == "/logs3") + .map(|r| serde_json::from_slice::(&r.body).unwrap()["rows"][0].clone()) + .collect(); + for (row, (name, span_type)) in rows + .iter() + .zip([("evaluation", "eval"), ("custom", "task")]) + { + assert_eq!(row["experiment_id"], "experiment"); + assert_eq!(row["span_attributes"]["name"], name); + assert_eq!(row["span_attributes"]["type"], span_type); + } +} diff --git a/tests/resource_handles.rs b/tests/resource_handles.rs index 01956df..8cce512 100644 --- a/tests/resource_handles.rs +++ b/tests/resource_handles.rs @@ -68,18 +68,18 @@ async fn resources_and_detached_span_builders_work_after_the_client_is_dropped() .parent_info(ParentSpanInfo::ProjectLogs { object_id: "project".into(), }); - let experiment: ExperimentBuilder = client + let experiment = client .experiment_builder_with_credentials("token", "org") .project_name("p"); - let dataset: DatasetBuilder = client + let dataset = 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(), + experiment: experiment.build(), + dataset: dataset.build(), }; resources.span.log(SpanLog::default()).await.unwrap(); let child: ExperimentSpanBuilder = resources.experiment.start_span().await.unwrap();