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 fc6b59b..994cce0 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. /// /// Fields are private — construction and destructuring happen only within this /// module. `worker.rs` never touches the fields directly. struct SubmitCommand { - token: String, + scope: CredentialScope, payload: SpanPayload, parent_info: Option, } @@ -40,48 +48,36 @@ struct LogQueueCore { dropped_count: AtomicUsize, last_drop_log_time: AtomicU64, 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 LogQueueCore { - 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 { dropped_count: AtomicUsize::new(0), last_drop_log_time: AtomicU64::new(0), 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. @@ -99,43 +95,18 @@ impl LogQueueCore { } 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"); @@ -150,21 +121,6 @@ impl LogQueueCore { 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, @@ -174,7 +130,10 @@ impl LogQueueCore { ), ); - for ((token, _org_id), (org_name, rows)) in groups { + for (scope, rows) in groups { + let (max_request_size, can_use_overflow) = self.get_version_info(&scope).await; + // Preserve the queue policy of half the advertised request limit. + let effective_batch_bytes = self.config.batch_max_bytes().min(max_request_size / 2); // Process chunks sequentially (matches TypeScript SDK) let chunks: Vec> = rows.chunks(chunk_size).map(<[_]>::to_vec).collect(); @@ -224,9 +183,9 @@ impl LogQueueCore { .map(|batch| { send_batch_with_retry( &self.client, - &api_url, - &token, - org_name.as_deref(), + scope.api_url(), + scope.token(), + scope.org_name(), batch, &self.config, max_request_size, @@ -255,7 +214,7 @@ impl LogQueueCore { /// if needed) and builds the complete row for the flush pipeline. async fn prepare_row( &self, - token: &str, + scope: &CredentialScope, payload: SpanPayload, parent_info: Option, ) -> std::result::Result { @@ -264,8 +223,8 @@ impl LogQueueCore { span_id, is_merge, span_components, - org_id, - org_name, + org_id: _, + org_name: _, project_name, input, output, @@ -280,14 +239,12 @@ impl LogQueueCore { extra, } = payload; + let org_id = scope.org_id().to_owned(); + let org_name = scope.org_name().map(str::to_owned); + if let Some(span_components) = span_components { let destination = self - .destination_from_span_components( - token, - &org_id, - org_name.as_deref(), - &span_components, - ) + .destination_from_span_components(scope, &span_components) .await?; let root_span_id = span_components .root_span_id @@ -323,10 +280,7 @@ impl LogQueueCore { } let project_id = if let Some(ref pn) = project_name { - Some( - self.ensure_project_id(token, &org_id, org_name.as_deref(), pn) - .await?, - ) + Some(self.ensure_project_id(scope, pn).await?) } else { None }; @@ -350,9 +304,7 @@ impl LogQueueCore { LogDestination::project_logs(object_id), ), Some(ParentSpanInfo::ProjectName { project_name }) => { - let proj_id = self - .ensure_project_id(token, &org_id, org_name.as_deref(), &project_name) - .await?; + let proj_id = self.ensure_project_id(scope, &project_name).await?; (span_id.clone(), None, LogDestination::project_logs(proj_id)) } Some(ParentSpanInfo::PlaygroundLogs { object_id }) => ( @@ -388,14 +340,8 @@ impl LogQueueCore { LogDestination::project_logs(object_id) } Some(ProjectLogsIdentifier::ProjectName(project_name)) => { - let project_id = self - .ensure_project_id( - token, - &org_id, - org_name.as_deref(), - &project_name, - ) - .await?; + let project_id = + self.ensure_project_id(scope, &project_name).await?; LogDestination::project_logs(project_id) } None => { @@ -445,9 +391,7 @@ impl LogQueueCore { async fn destination_from_span_components( &self, - token: &str, - org_id: &str, - org_name: Option<&str>, + scope: &CredentialScope, span_components: &crate::span_components::SpanComponents, ) -> std::result::Result { match span_components.object_type { @@ -461,9 +405,7 @@ impl LogQueueCore { Ok(LogDestination::project_logs(object_id)) } Some(ProjectLogsIdentifier::ProjectName(project_name)) => { - let project_id = self - .ensure_project_id(token, org_id, org_name, &project_name) - .await?; + let project_id = self.ensure_project_id(scope, &project_name).await?; Ok(LogDestination::project_logs(project_id)) } None => anyhow::bail!( @@ -478,71 +420,32 @@ impl LogQueueCore { } } - /// Ensure a project ID is available for the given project name, registering - /// it via the API if it is not yet cached. - /// - /// Uses `login_state.app_url()` for the registration endpoint so the URL - /// doesn't need to be passed separately. + /// Resolve project names only within the producer's immutable request scope. async fn ensure_project_id( &self, - token: &str, - org_id: &str, - org_name: Option<&str>, + scope: &CredentialScope, project_name: &str, ) -> std::result::Result { - let cache_key = format!("{org_id}:{project_name}"); - { - let cache = self.project_cache.lock().await; - if let Some(project_id) = cache.get(&cache_key) { - return Ok(project_id.clone()); - } - } - - // Prefer the login-state control-plane URL (logged-in path), falling back - // to the client-configured URL for per-request-credential callers. - let app_url = match self.login_state.app_url() { - Some(app_url_str) => { - Url::parse(&app_url_str).map_err(|e| anyhow::anyhow!("invalid app URL: {e}"))? - } - None => self.app_url.clone(), - }; - let url = app_url - .join("api/project/register") - .map_err(|e| anyhow::anyhow!("invalid project register url: {e}"))?; - - let request = ProjectRegisterRequest { - project_name: 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) } /// Log a drop warning with throttling. @@ -590,7 +493,7 @@ impl LogQueueCore { }; handle.spawn(async move { let SubmitCommand { - token: _, + scope: _, payload, parent_info, } = cmd; @@ -718,16 +621,9 @@ 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 core = LogQueueCore::new(config, login_state, client, api_url, app_url); + let core = LogQueueCore::new(config, client); let processor = core.clone(); let writes = WriteService::new( capacity, @@ -742,12 +638,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, }; @@ -789,6 +685,7 @@ impl LogQueue { mod tests { use super::*; use serde_json::{json, Value}; + use url::Url; use wiremock::{ matchers::{header, method, path}, Mock, MockServer, ResponseTemplate, @@ -816,12 +713,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(), ) @@ -838,7 +741,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")) @@ -856,9 +759,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"); }; @@ -883,7 +786,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 @@ -917,8 +820,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 9a0058d..7deaa7d 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) } /// Update an existing span using explicit credentials instead of shared login state. @@ -646,14 +642,12 @@ impl BraintrustClient { let token = token.into(); let org_id = org_id.into(); - self.update_span_internal(token, org_id, None, exported, event) + self.update_span_internal(self.explicit_scope(token, org_id, None), exported, event) } fn update_span_internal( &self, - token: String, - org_id: String, - org_name: Option, + scope: CredentialScope, exported: &str, event: crate::span::SpanLog, ) -> Result<()> { @@ -733,8 +727,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, @@ -749,7 +743,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 { return Err(BraintrustError::Background(error.to_string())); } @@ -916,17 +910,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. @@ -955,11 +949,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 beb4022..f0aed07 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 b9718d4..c2ab142 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 + ); +}