braintrust_sdk_rust/
logger.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3use std::time::Duration;
4
5use anyhow::Context;
6use async_trait::async_trait;
7use chrono::Utc;
8use serde::{Deserialize, Serialize};
9use tokio::sync::{mpsc, oneshot, Notify, RwLock};
10use tokio::task::JoinHandle;
11use tracing::warn;
12use url::Url;
13
14use crate::error::{BraintrustError, Result};
15use crate::span::SpanSubmitter;
16use crate::types::{
17    LogDestination, Logs3Request, Logs3Row, ParentSpanInfo, SpanObjectType, SpanPayload,
18    LOGS_API_VERSION,
19};
20
21const DEFAULT_QUEUE_SIZE: usize = 256;
22const REQUEST_TIMEOUT: Duration = Duration::from_secs(10);
23const LOGIN_TIMEOUT: Duration = Duration::from_secs(30);
24
25/// Organization info returned from login.
26#[derive(Debug, Clone, Deserialize)]
27pub struct OrgInfo {
28    pub id: String,
29    pub name: String,
30    #[serde(default)]
31    pub api_url: Option<String>,
32}
33
34/// Response from the login endpoint.
35#[derive(Debug, Deserialize)]
36struct LoginResponse {
37    org_info: Vec<OrgInfo>,
38}
39
40/// Logged-in state containing API key and org info.
41#[derive(Debug, Clone)]
42pub struct LoginState {
43    pub api_key: String,
44    pub org_id: String,
45    pub org_name: String,
46    pub api_url: Option<String>,
47}
48
49/// Builder for creating a BraintrustClient with configuration.
50///
51/// Configuration is loaded from environment variables by default,
52/// and can be overridden using builder methods.
53///
54/// # Example
55///
56/// ```no_run
57/// use braintrust_sdk_rust::BraintrustClient;
58///
59/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
60/// // Using environment variables (BRAINTRUST_API_KEY, etc.)
61/// let client = BraintrustClient::builder().build().await?;
62///
63/// // With explicit configuration
64/// let client = BraintrustClient::builder()
65///     .api_key("sk-...")
66///     .org_name("my-org")
67///     .default_project("my-project")
68///     .blocking_login(true)
69///     .build()
70///     .await?;
71/// # Ok(())
72/// # }
73/// ```
74pub struct BraintrustClientBuilder {
75    api_key: Option<String>,
76    app_url: Option<String>,
77    api_url: Option<String>,
78    org_name: Option<String>,
79    default_project: Option<String>,
80    queue_size: usize,
81    blocking_login: bool,
82}
83
84impl BraintrustClientBuilder {
85    /// Create a new builder with defaults from environment variables.
86    ///
87    /// Supported environment variables:
88    /// - `BRAINTRUST_API_KEY`: API key for authentication (required)
89    /// - `BRAINTRUST_APP_URL`: Braintrust app URL (default: `https://www.braintrust.dev`)
90    /// - `BRAINTRUST_API_URL`: API endpoint URL (default: `https://api.braintrust.dev`)
91    /// - `BRAINTRUST_ORG_NAME`: Organization name (default: first org from login)
92    /// - `BRAINTRUST_DEFAULT_PROJECT`: Default project name
93    pub fn new() -> Self {
94        Self {
95            api_key: std::env::var("BRAINTRUST_API_KEY").ok(),
96            app_url: std::env::var("BRAINTRUST_APP_URL").ok(),
97            api_url: std::env::var("BRAINTRUST_API_URL").ok(),
98            org_name: std::env::var("BRAINTRUST_ORG_NAME").ok(),
99            default_project: std::env::var("BRAINTRUST_DEFAULT_PROJECT").ok(),
100            queue_size: DEFAULT_QUEUE_SIZE,
101            blocking_login: false,
102        }
103    }
104
105    /// Set the API key (overrides `BRAINTRUST_API_KEY` env var).
106    pub fn api_key(mut self, api_key: impl Into<String>) -> Self {
107        self.api_key = Some(api_key.into());
108        self
109    }
110
111    /// Set the app URL (overrides `BRAINTRUST_APP_URL` env var).
112    pub fn app_url(mut self, url: impl Into<String>) -> Self {
113        self.app_url = Some(url.into());
114        self
115    }
116
117    /// Set the API URL (overrides `BRAINTRUST_API_URL` env var).
118    pub fn api_url(mut self, url: impl Into<String>) -> Self {
119        self.api_url = Some(url.into());
120        self
121    }
122
123    /// Set the organization name (overrides `BRAINTRUST_ORG_NAME` env var).
124    pub fn org_name(mut self, name: impl Into<String>) -> Self {
125        self.org_name = Some(name.into());
126        self
127    }
128
129    /// Set the default project name (overrides `BRAINTRUST_DEFAULT_PROJECT` env var).
130    pub fn default_project(mut self, name: impl Into<String>) -> Self {
131        self.default_project = Some(name.into());
132        self
133    }
134
135    /// Set the internal queue size for buffering log events.
136    pub fn queue_size(mut self, size: usize) -> Self {
137        self.queue_size = size;
138        self
139    }
140
141    /// Block until login completes (default: false, login happens in background).
142    ///
143    /// When `false` (default), login happens asynchronously in a background task.
144    /// When `true`, the `build()` method waits for login to complete before returning.
145    pub fn blocking_login(mut self, blocking: bool) -> Self {
146        self.blocking_login = blocking;
147        self
148    }
149
150    /// Build the client, performing login.
151    ///
152    /// If `blocking_login` is true, waits for login to complete.
153    /// Otherwise, login happens in the background with retry logic.
154    pub async fn build(self) -> Result<BraintrustClient> {
155        // Validate required fields
156        let api_key = self.api_key.ok_or_else(|| {
157            BraintrustError::InvalidConfig(
158                "API key required: set BRAINTRUST_API_KEY or call .api_key()".into(),
159            )
160        })?;
161
162        let app_url_str = self
163            .app_url
164            .unwrap_or_else(|| "https://www.braintrust.dev".into());
165        let api_url_str = self
166            .api_url
167            .unwrap_or_else(|| "https://api.braintrust.dev".into());
168
169        let app_url = Url::parse(&app_url_str)
170            .map_err(|e| BraintrustError::InvalidConfig(format!("invalid app_url: {}", e)))?;
171        let api_url = Url::parse(&api_url_str)
172            .map_err(|e| BraintrustError::InvalidConfig(format!("invalid api_url: {}", e)))?;
173
174        let http_client = reqwest::Client::builder()
175            .timeout(REQUEST_TIMEOUT)
176            .build()
177            .map_err(|e| BraintrustError::InvalidConfig(e.to_string()))?;
178
179        let (sender, receiver) = mpsc::channel(self.queue_size.max(32));
180        let worker = tokio::spawn(run_worker(api_url.clone(), app_url.clone(), receiver));
181
182        let client = BraintrustClient {
183            inner: Arc::new(ClientInner {
184                api_url,
185                app_url,
186                sender,
187                worker,
188                login_state: RwLock::new(None),
189                login_notify: Notify::new(),
190                http_client,
191                default_project: self.default_project,
192            }),
193        };
194
195        // Perform login
196        if self.blocking_login {
197            client
198                .perform_login(&api_key, self.org_name.as_deref())
199                .await?;
200        } else {
201            client.start_background_login(api_key, self.org_name);
202        }
203
204        Ok(client)
205    }
206}
207
208impl Default for BraintrustClientBuilder {
209    fn default() -> Self {
210        Self::new()
211    }
212}
213
214#[derive(Clone, Debug)]
215pub struct BraintrustClient {
216    inner: Arc<ClientInner>,
217}
218
219struct ClientInner {
220    #[allow(dead_code)]
221    api_url: Url,
222    app_url: Url,
223    sender: mpsc::Sender<LogCommand>,
224    #[allow(dead_code)]
225    worker: JoinHandle<()>,
226    login_state: RwLock<Option<LoginState>>,
227    login_notify: Notify,
228    http_client: reqwest::Client,
229    default_project: Option<String>,
230}
231
232impl std::fmt::Debug for ClientInner {
233    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
234        f.debug_struct("ClientInner")
235            .field("api_url", &self.api_url)
236            .field("app_url", &self.app_url)
237            .field("default_project", &self.default_project)
238            .finish_non_exhaustive()
239    }
240}
241
242impl BraintrustClient {
243    /// Create a new builder for configuring the client.
244    ///
245    /// # Example
246    ///
247    /// ```no_run
248    /// use braintrust_sdk_rust::BraintrustClient;
249    ///
250    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
251    /// let client = BraintrustClient::builder()
252    ///     .api_key("sk-...")
253    ///     .build()
254    ///     .await?;
255    /// # Ok(())
256    /// # }
257    /// ```
258    pub fn builder() -> BraintrustClientBuilder {
259        BraintrustClientBuilder::new()
260    }
261
262    /// Check if the client is logged in.
263    pub async fn is_logged_in(&self) -> bool {
264        self.inner.login_state.read().await.is_some()
265    }
266
267    /// Wait for login to complete (useful if using background login).
268    ///
269    /// Returns the login state once available, or an error if login times out.
270    pub async fn wait_for_login(&self) -> Result<LoginState> {
271        self.wait_for_login_state().await
272    }
273
274    /// Get the current login state, if logged in.
275    pub async fn login_state(&self) -> Option<LoginState> {
276        self.inner.login_state.read().await.clone()
277    }
278
279    /// Create a span builder using the logged-in state and default project.
280    ///
281    /// This waits for login to complete if it hasn't already.
282    ///
283    /// # Example
284    ///
285    /// ```no_run
286    /// use braintrust_sdk_rust::BraintrustClient;
287    ///
288    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
289    /// let client = BraintrustClient::builder()
290    ///     .api_key("sk-...")
291    ///     .default_project("my-project")
292    ///     .build()
293    ///     .await?;
294    ///
295    /// let span = client.span_builder().await?.build();
296    /// # Ok(())
297    /// # }
298    /// ```
299    pub async fn span_builder(&self) -> Result<crate::span::SpanBuilder<Self>> {
300        let state = self.wait_for_login_state().await?;
301        let mut builder =
302            crate::span::SpanBuilder::new(Arc::new(self.clone()), &state.api_key, &state.org_id);
303        if let Some(ref project) = self.inner.default_project {
304            builder = builder.project_name(project);
305        }
306        Ok(builder)
307    }
308
309    /// Create a span builder with explicit token and org_id.
310    ///
311    /// Use this if you already have the org_id and don't want to use the login state.
312    pub fn span_builder_with_credentials(
313        &self,
314        token: impl Into<String>,
315        org_id: impl Into<String>,
316    ) -> crate::span::SpanBuilder<Self> {
317        let submitter = Arc::new(self.clone());
318        crate::span::SpanBuilder::new(submitter, token, org_id)
319    }
320
321    /// Perform login synchronously.
322    async fn perform_login(&self, api_key: &str, org_name: Option<&str>) -> Result<()> {
323        let login_url = self
324            .inner
325            .app_url
326            .join("api/apikey/login")
327            .map_err(|e| BraintrustError::InvalidConfig(e.to_string()))?;
328
329        let response = self
330            .inner
331            .http_client
332            .post(login_url)
333            .bearer_auth(api_key)
334            .header("Content-Type", "application/json")
335            .send()
336            .await
337            .map_err(|e| BraintrustError::Network(e.to_string()))?;
338
339        if !response.status().is_success() {
340            let status = response.status();
341            let body = response.text().await.unwrap_or_default();
342            return Err(BraintrustError::Api {
343                status: status.as_u16(),
344                message: body,
345            });
346        }
347
348        let login_response: LoginResponse =
349            response.json().await.map_err(|e| BraintrustError::Api {
350                status: 200,
351                message: format!("Failed to parse login response: {}", e),
352            })?;
353
354        // Find matching org or use first
355        let org = if let Some(name) = org_name {
356            login_response
357                .org_info
358                .into_iter()
359                .find(|o| o.name == name)
360                .ok_or_else(|| {
361                    BraintrustError::InvalidConfig(format!("Organization '{}' not found", name))
362                })?
363        } else {
364            login_response.org_info.into_iter().next().ok_or_else(|| {
365                BraintrustError::InvalidConfig(
366                    "No organizations found for this API key".to_string(),
367                )
368            })?
369        };
370
371        let state = LoginState {
372            api_key: api_key.to_string(),
373            org_id: org.id,
374            org_name: org.name,
375            api_url: org.api_url,
376        };
377
378        *self.inner.login_state.write().await = Some(state);
379        self.inner.login_notify.notify_waiters();
380        Ok(())
381    }
382
383    /// Start login in a background task with retry logic.
384    fn start_background_login(&self, api_key: String, org_name: Option<String>) {
385        let client = self.clone();
386        tokio::spawn(async move {
387            // Retry with exponential backoff
388            let mut delay = Duration::from_millis(100);
389            let max_delay = Duration::from_secs(5);
390
391            loop {
392                match client.perform_login(&api_key, org_name.as_deref()).await {
393                    Ok(()) => {
394                        tracing::debug!("Background login completed successfully");
395                        break;
396                    }
397                    Err(e) => {
398                        tracing::warn!("Background login failed: {}, retrying in {:?}", e, delay);
399                        tokio::time::sleep(delay).await;
400                        delay = (delay * 2).min(max_delay);
401                    }
402                }
403            }
404        });
405    }
406
407    /// Wait for login state to be available.
408    async fn wait_for_login_state(&self) -> Result<LoginState> {
409        // Check if already logged in
410        if let Some(state) = self.inner.login_state.read().await.clone() {
411            return Ok(state);
412        }
413
414        // Get notification future BEFORE checking state to avoid race condition
415        let notified = self.inner.login_notify.notified();
416
417        // Check state again (may have been set between our first check and now)
418        if let Some(state) = self.inner.login_state.read().await.clone() {
419            return Ok(state);
420        }
421
422        // Wait for notification or timeout
423        tokio::select! {
424            _ = notified => {
425                // Login completed - state is guaranteed to be set since
426                // notify_waiters() is only called after setting state
427                self.inner.login_state.read().await.clone().ok_or_else(|| {
428                    BraintrustError::InvalidConfig(
429                        "Login notification received but state not set".into(),
430                    )
431                })
432            }
433            _ = tokio::time::sleep(LOGIN_TIMEOUT) => {
434                Err(BraintrustError::InvalidConfig(
435                    "Timeout waiting for login to complete".into(),
436                ))
437            }
438        }
439    }
440
441    /// Submit a span payload for logging (fire-and-forget).
442    ///
443    /// Returns immediately after queuing. HTTP submission happens in the background.
444    /// Errors are logged as warnings but not propagated to callers.
445    pub(crate) async fn submit_payload(
446        &self,
447        token: impl Into<String>,
448        payload: SpanPayload,
449        parent_info: Option<ParentSpanInfo>,
450    ) -> Result<()> {
451        let cmd = LogCommand::Submit(Box::new(SubmitCommand {
452            token: token.into(),
453            payload,
454            parent_info,
455        }));
456        self.inner
457            .sender
458            .send(cmd)
459            .await
460            .map_err(|_| BraintrustError::ChannelClosed)?;
461        Ok(())
462    }
463
464    /// Flush all pending log events.
465    pub async fn flush(&self) -> Result<()> {
466        let (tx, rx) = oneshot::channel();
467        self.inner
468            .sender
469            .send(LogCommand::Flush(tx))
470            .await
471            .map_err(|_| BraintrustError::ChannelClosed)?;
472        rx.await
473            .map_err(|_| BraintrustError::ChannelClosed)?
474            .map_err(|e| BraintrustError::Background(e.to_string()))
475    }
476}
477
478#[async_trait]
479impl SpanSubmitter for BraintrustClient {
480    async fn submit(
481        &self,
482        token: impl Into<String> + Send,
483        payload: SpanPayload,
484        parent_info: Option<ParentSpanInfo>,
485    ) -> Result<()> {
486        self.submit_payload(token, payload, parent_info).await
487    }
488}
489
490enum LogCommand {
491    Submit(Box<SubmitCommand>),
492    Flush(oneshot::Sender<std::result::Result<(), anyhow::Error>>),
493}
494
495struct SubmitCommand {
496    token: String,
497    payload: SpanPayload,
498    parent_info: Option<ParentSpanInfo>,
499}
500
501/// Request body for project registration.
502#[derive(Serialize)]
503struct ProjectRegisterRequest<'a> {
504    project_name: &'a str,
505    #[serde(skip_serializing_if = "Option::is_none")]
506    org_id: Option<&'a str>,
507    #[serde(skip_serializing_if = "Option::is_none")]
508    org_name: Option<&'a str>,
509}
510
511/// Response from project registration.
512#[derive(Deserialize)]
513struct ProjectRegisterResponse {
514    project: ProjectInfo,
515}
516
517/// Project info in registration response.
518#[derive(Deserialize)]
519struct ProjectInfo {
520    id: String,
521}
522
523async fn run_worker(api_url: Url, app_url: Url, mut receiver: mpsc::Receiver<LogCommand>) {
524    let mut state = WorkerState::new(api_url, app_url);
525    while let Some(cmd) = receiver.recv().await {
526        match cmd {
527            LogCommand::Submit(cmd) => {
528                let SubmitCommand {
529                    token,
530                    payload,
531                    parent_info,
532                } = *cmd;
533                // Fire-and-forget: log errors but don't propagate
534                if let Err(e) = state.submit_payload(&token, payload, parent_info).await {
535                    warn!(error = %e, "failed to submit span to Braintrust");
536                }
537            }
538            LogCommand::Flush(response) => {
539                let _ = response.send(Ok(()));
540            }
541        }
542    }
543}
544
545struct WorkerState {
546    /// URL for data plane requests (e.g., /logs3)
547    api_url: Url,
548    /// URL for control plane requests (e.g., /api/project/register)
549    app_url: Url,
550    client: reqwest::Client,
551    project_cache: HashMap<String, String>,
552}
553
554impl WorkerState {
555    fn new(api_url: Url, app_url: Url) -> Self {
556        let client = reqwest::Client::builder()
557            .timeout(REQUEST_TIMEOUT)
558            .build()
559            .expect("reqwest client");
560        Self {
561            api_url,
562            app_url,
563            client,
564            project_cache: HashMap::new(),
565        }
566    }
567
568    async fn submit_payload(
569        &mut self,
570        token: &str,
571        payload: SpanPayload,
572        parent_info: Option<ParentSpanInfo>,
573    ) -> std::result::Result<(), anyhow::Error> {
574        let SpanPayload {
575            row_id,
576            span_id,
577            is_merge,
578            org_id,
579            org_name,
580            project_name,
581            input,
582            output,
583            metadata,
584            metrics,
585            span_attributes,
586        } = payload;
587
588        let project_id = if let Some(ref project_name) = project_name {
589            Some(
590                self.ensure_project_id(token, &org_id, org_name.as_deref(), project_name)
591                    .await?,
592            )
593        } else {
594            None
595        };
596
597        let logs_url = self
598            .api_url
599            .join("logs3")
600            .map_err(|e| anyhow::anyhow!("invalid logs url: {e}"))?;
601
602        // row_id and span_id come from payload - generated once at span creation, reused on every flush
603
604        // Determine destination and span hierarchy based on parent info
605        let (root_span_id, span_parents, destination) = match parent_info {
606            None => {
607                // No parent - use project_id if available, otherwise fail
608                let dest = match project_id {
609                    Some(pid) => LogDestination::project_logs(pid),
610                    None => {
611                        anyhow::bail!("no destination: either parent_info or project_name required")
612                    }
613                };
614                (span_id.clone(), None, dest)
615            }
616            Some(ParentSpanInfo::Experiment { object_id }) => {
617                (span_id.clone(), None, LogDestination::experiment(object_id))
618            }
619            Some(ParentSpanInfo::ProjectLogs { object_id }) => (
620                span_id.clone(),
621                None,
622                LogDestination::project_logs(object_id),
623            ),
624            Some(ParentSpanInfo::ProjectName { project_name }) => {
625                let proj_id = self
626                    .ensure_project_id(token, &org_id, org_name.as_deref(), &project_name)
627                    .await?;
628                (span_id.clone(), None, LogDestination::project_logs(proj_id))
629            }
630            Some(ParentSpanInfo::PlaygroundLogs { object_id }) => (
631                span_id.clone(),
632                None,
633                LogDestination::playground_logs(object_id),
634            ),
635            Some(ParentSpanInfo::FullSpan {
636                object_type,
637                object_id,
638                span_id: parent_span_id,
639                root_span_id: parent_root_span_id,
640            }) => {
641                let span_parents = Some(vec![parent_span_id]);
642                let dest = match object_type {
643                    SpanObjectType::Experiment => LogDestination::experiment(object_id),
644                    SpanObjectType::ProjectLogs => LogDestination::project_logs(object_id),
645                    SpanObjectType::PlaygroundLogs => LogDestination::playground_logs(object_id),
646                };
647                (parent_root_span_id, span_parents, dest)
648            }
649        };
650
651        let row = Logs3Row {
652            id: row_id,
653            is_merge: if is_merge { Some(true) } else { None },
654            span_id,
655            root_span_id,
656            span_parents,
657            destination,
658            org_id,
659            org_name,
660            input,
661            output,
662            metadata,
663            metrics,
664            span_attributes,
665            created: Utc::now(),
666        };
667
668        let request = Logs3Request {
669            rows: vec![row],
670            api_version: LOGS_API_VERSION,
671        };
672
673        let json_bytes = serde_json::to_vec(&request)
674            .map_err(|e| anyhow::anyhow!("JSON serialization failed: {e}"))?;
675
676        let response = self
677            .client
678            .post(logs_url)
679            .bearer_auth(token)
680            .header("content-type", "application/json")
681            .body(json_bytes)
682            .send()
683            .await?;
684
685        if !response.status().is_success() {
686            let status = response.status();
687            let body = response
688                .text()
689                .await
690                .unwrap_or_else(|_| "<unavailable>".to_string());
691            tracing::warn!("failed to submit span: [{status}] {body}");
692        }
693
694        Ok(())
695    }
696
697    async fn ensure_project_id(
698        &mut self,
699        token: &str,
700        org_id: &str,
701        org_name: Option<&str>,
702        project_name: &str,
703    ) -> std::result::Result<String, anyhow::Error> {
704        let cache_key = format!("{org_id}:{project_name}");
705        if let Some(project_id) = self.project_cache.get(&cache_key) {
706            return Ok(project_id.clone());
707        }
708
709        let request = ProjectRegisterRequest {
710            project_name,
711            org_id: (!org_id.is_empty()).then_some(org_id),
712            org_name,
713        };
714
715        let url = self
716            .app_url
717            .join("api/project/register")
718            .map_err(|e| anyhow::anyhow!("invalid project register url: {e}"))?;
719        let response = self
720            .client
721            .post(url)
722            .bearer_auth(token)
723            .json(&request)
724            .send()
725            .await?;
726        let status = response.status();
727        if !status.is_success() {
728            let text = response.text().await.unwrap_or_default();
729            anyhow::bail!("register project failed: [{status}] {text}");
730        }
731
732        let register_response: ProjectRegisterResponse = response
733            .json()
734            .await
735            .context("failed to parse project registration response")?;
736
737        self.project_cache
738            .insert(cache_key, register_response.project.id.clone());
739        Ok(register_response.project.id)
740    }
741}
742
743#[cfg(test)]
744mod tests {
745    use super::*;
746    use crate::span::SpanLog;
747    use serde_json::Value;
748    use wiremock::matchers::{method, path};
749    use wiremock::{Mock, MockServer, ResponseTemplate};
750
751    // Helper to create a mock login response
752    fn mock_login_response(orgs: &[(&str, &str)]) -> ResponseTemplate {
753        let org_info: Vec<_> = orgs
754            .iter()
755            .map(|(id, name)| serde_json::json!({ "id": id, "name": name }))
756            .collect();
757        ResponseTemplate::new(200).set_body_json(serde_json::json!({ "org_info": org_info }))
758    }
759
760    #[tokio::test]
761    async fn builder_rejects_missing_api_key() {
762        // Clear any env vars that might be set
763        std::env::remove_var("BRAINTRUST_API_KEY");
764
765        let result = BraintrustClient::builder()
766            .app_url("https://example.com")
767            .build()
768            .await;
769
770        assert!(result.is_err());
771        let err = result.unwrap_err();
772        assert!(matches!(err, BraintrustError::InvalidConfig(_)));
773    }
774
775    #[tokio::test]
776    async fn builder_rejects_invalid_app_url() {
777        let result = BraintrustClient::builder()
778            .api_key("test-key")
779            .app_url("::not a url::")
780            .build()
781            .await;
782
783        assert!(result.is_err());
784        let err = result.unwrap_err();
785        assert!(matches!(err, BraintrustError::InvalidConfig(_)));
786    }
787
788    #[tokio::test]
789    async fn builder_rejects_invalid_api_url() {
790        let result = BraintrustClient::builder()
791            .api_key("test-key")
792            .api_url("::not a url::")
793            .build()
794            .await;
795
796        assert!(result.is_err());
797        let err = result.unwrap_err();
798        assert!(matches!(err, BraintrustError::InvalidConfig(_)));
799    }
800
801    #[tokio::test]
802    async fn project_registration_is_cached() {
803        let server = MockServer::start().await;
804
805        Mock::given(method("POST"))
806            .and(path("/api/project/register"))
807            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
808                "project": { "id": "test-project-id" }
809            })))
810            .expect(1)
811            .mount(&server)
812            .await;
813
814        Mock::given(method("POST"))
815            .and(path("/logs3"))
816            .respond_with(ResponseTemplate::new(200).set_body_string("{}"))
817            .mount(&server)
818            .await;
819
820        Mock::given(method("POST"))
821            .and(path("/api/apikey/login"))
822            .respond_with(mock_login_response(&[("org-id", "Test Org")]))
823            .mount(&server)
824            .await;
825
826        let client = BraintrustClient::builder()
827            .api_key("token")
828            .app_url(server.uri())
829            .api_url(server.uri())
830            .blocking_login(true)
831            .build()
832            .await
833            .expect("client");
834
835        for _ in 0..2 {
836            let span = client
837                .span_builder()
838                .await
839                .expect("span_builder")
840                .project_name("demo-project")
841                .build();
842            span.log(SpanLog {
843                input: Some(Value::String("hello".into())),
844                ..Default::default()
845            })
846            .await;
847            span.flush().await.expect("flush");
848            client.flush().await.expect("client flush");
849        }
850
851        let register_calls = server
852            .received_requests()
853            .await
854            .unwrap()
855            .into_iter()
856            .filter(|request| request.url.path() == "/api/project/register")
857            .count();
858
859        assert_eq!(register_calls, 1);
860    }
861
862    #[tokio::test]
863    async fn logs_request_contains_span_rows() {
864        let server = MockServer::start().await;
865
866        Mock::given(method("POST"))
867            .and(path("/api/project/register"))
868            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
869                "project": { "id": "proj-id" }
870            })))
871            .mount(&server)
872            .await;
873
874        Mock::given(method("POST"))
875            .and(path("/logs3"))
876            .respond_with(ResponseTemplate::new(200).set_body_string("{}"))
877            .mount(&server)
878            .await;
879
880        Mock::given(method("POST"))
881            .and(path("/api/apikey/login"))
882            .respond_with(mock_login_response(&[("org-id", "Test Org")]))
883            .mount(&server)
884            .await;
885
886        let client = BraintrustClient::builder()
887            .api_key("token")
888            .app_url(server.uri())
889            .api_url(server.uri())
890            .blocking_login(true)
891            .build()
892            .await
893            .expect("client");
894
895        let span = client
896            .span_builder()
897            .await
898            .expect("span_builder")
899            .project_name("demo-project")
900            .build();
901        span.log(SpanLog {
902            input: Some(Value::String("input".into())),
903            ..Default::default()
904        })
905        .await;
906        span.flush().await.expect("flush");
907        client.flush().await.expect("client flush");
908
909        let logs_request = server
910            .received_requests()
911            .await
912            .unwrap()
913            .into_iter()
914            .find(|request| request.url.path() == "/logs3")
915            .expect("logs request present");
916        let body: Value = serde_json::from_slice(&logs_request.body).expect("json");
917        assert!(body.get("rows").is_some());
918    }
919
920    #[tokio::test]
921    async fn blocking_login_returns_first_org_by_default() {
922        let server = MockServer::start().await;
923
924        Mock::given(method("POST"))
925            .and(path("/api/apikey/login"))
926            .respond_with(mock_login_response(&[
927                ("org-1", "First Org"),
928                ("org-2", "Second Org"),
929            ]))
930            .mount(&server)
931            .await;
932
933        let client = BraintrustClient::builder()
934            .api_key("test-api-key")
935            .app_url(server.uri())
936            .api_url(server.uri())
937            .blocking_login(true)
938            .build()
939            .await
940            .expect("client");
941
942        let login_state = client.login_state().await.expect("should be logged in");
943
944        assert_eq!(login_state.org_id, "org-1");
945        assert_eq!(login_state.org_name, "First Org");
946        assert_eq!(login_state.api_key, "test-api-key");
947    }
948
949    #[tokio::test]
950    async fn blocking_login_selects_org_by_name() {
951        let server = MockServer::start().await;
952
953        Mock::given(method("POST"))
954            .and(path("/api/apikey/login"))
955            .respond_with(mock_login_response(&[
956                ("org-1", "First Org"),
957                ("org-2", "Second Org"),
958            ]))
959            .mount(&server)
960            .await;
961
962        let client = BraintrustClient::builder()
963            .api_key("test-api-key")
964            .app_url(server.uri())
965            .api_url(server.uri())
966            .org_name("Second Org")
967            .blocking_login(true)
968            .build()
969            .await
970            .expect("client");
971
972        let login_state = client.login_state().await.expect("should be logged in");
973
974        assert_eq!(login_state.org_id, "org-2");
975        assert_eq!(login_state.org_name, "Second Org");
976    }
977
978    #[tokio::test]
979    async fn blocking_login_errors_when_org_not_found() {
980        let server = MockServer::start().await;
981
982        Mock::given(method("POST"))
983            .and(path("/api/apikey/login"))
984            .respond_with(mock_login_response(&[("org-1", "First Org")]))
985            .mount(&server)
986            .await;
987
988        let result = BraintrustClient::builder()
989            .api_key("test-api-key")
990            .app_url(server.uri())
991            .api_url(server.uri())
992            .org_name("Nonexistent Org")
993            .blocking_login(true)
994            .build()
995            .await;
996
997        assert!(result.is_err());
998        let err = result.unwrap_err();
999        assert!(matches!(err, BraintrustError::InvalidConfig(_)));
1000    }
1001
1002    #[tokio::test]
1003    async fn blocking_login_errors_when_no_orgs_returned() {
1004        let server = MockServer::start().await;
1005
1006        Mock::given(method("POST"))
1007            .and(path("/api/apikey/login"))
1008            .respond_with(
1009                ResponseTemplate::new(200).set_body_json(serde_json::json!({ "org_info": [] })),
1010            )
1011            .mount(&server)
1012            .await;
1013
1014        let result = BraintrustClient::builder()
1015            .api_key("test-api-key")
1016            .app_url(server.uri())
1017            .api_url(server.uri())
1018            .blocking_login(true)
1019            .build()
1020            .await;
1021
1022        assert!(result.is_err());
1023        let err = result.unwrap_err();
1024        assert!(matches!(err, BraintrustError::InvalidConfig(_)));
1025    }
1026
1027    #[tokio::test]
1028    async fn blocking_login_errors_on_api_failure() {
1029        let server = MockServer::start().await;
1030
1031        Mock::given(method("POST"))
1032            .and(path("/api/apikey/login"))
1033            .respond_with(ResponseTemplate::new(401).set_body_string("Unauthorized"))
1034            .mount(&server)
1035            .await;
1036
1037        let result = BraintrustClient::builder()
1038            .api_key("bad-api-key")
1039            .app_url(server.uri())
1040            .api_url(server.uri())
1041            .blocking_login(true)
1042            .build()
1043            .await;
1044
1045        assert!(result.is_err());
1046        let err = result.unwrap_err();
1047        assert!(matches!(err, BraintrustError::Api { status: 401, .. }));
1048    }
1049
1050    #[tokio::test]
1051    async fn is_logged_in_returns_false_initially_with_background_login() {
1052        let server = MockServer::start().await;
1053
1054        // Don't mount any mock - login will fail and retry in background
1055        let client = BraintrustClient::builder()
1056            .api_key("test-api-key")
1057            .app_url(server.uri())
1058            .api_url(server.uri())
1059            // blocking_login defaults to false
1060            .build()
1061            .await
1062            .expect("client");
1063
1064        // Initially should not be logged in (background login hasn't completed)
1065        let is_logged_in = client.is_logged_in().await;
1066        assert!(!is_logged_in);
1067    }
1068
1069    #[tokio::test]
1070    async fn wait_for_login_succeeds_after_background_login() {
1071        let server = MockServer::start().await;
1072
1073        Mock::given(method("POST"))
1074            .and(path("/api/apikey/login"))
1075            .respond_with(mock_login_response(&[("org-1", "Test Org")]))
1076            .mount(&server)
1077            .await;
1078
1079        let client = BraintrustClient::builder()
1080            .api_key("test-api-key")
1081            .app_url(server.uri())
1082            .api_url(server.uri())
1083            .build()
1084            .await
1085            .expect("client");
1086
1087        // Wait for background login to complete
1088        let login_state = client.wait_for_login().await.expect("login should succeed");
1089        assert_eq!(login_state.org_id, "org-1");
1090        assert!(client.is_logged_in().await);
1091    }
1092
1093    #[tokio::test]
1094    async fn span_builder_uses_login_state_and_default_project() {
1095        let server = MockServer::start().await;
1096
1097        Mock::given(method("POST"))
1098            .and(path("/api/apikey/login"))
1099            .respond_with(mock_login_response(&[("org-123", "Test Org")]))
1100            .mount(&server)
1101            .await;
1102
1103        Mock::given(method("POST"))
1104            .and(path("/api/project/register"))
1105            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
1106                "project": { "id": "proj-id" }
1107            })))
1108            .mount(&server)
1109            .await;
1110
1111        Mock::given(method("POST"))
1112            .and(path("/logs3"))
1113            .respond_with(ResponseTemplate::new(200).set_body_string("{}"))
1114            .mount(&server)
1115            .await;
1116
1117        let client = BraintrustClient::builder()
1118            .api_key("test-api-key")
1119            .app_url(server.uri())
1120            .api_url(server.uri())
1121            .default_project("test-project")
1122            .blocking_login(true)
1123            .build()
1124            .await
1125            .expect("client");
1126
1127        // span_builder() should use login state and default project
1128        let span = client.span_builder().await.expect("span_builder").build();
1129
1130        span.log(SpanLog {
1131            input: Some(Value::String("test".into())),
1132            ..Default::default()
1133        })
1134        .await;
1135        span.flush().await.expect("flush");
1136        client.flush().await.expect("client flush");
1137
1138        // Verify the logs request was made with the correct org_id
1139        let logs_request = server
1140            .received_requests()
1141            .await
1142            .unwrap()
1143            .into_iter()
1144            .find(|request| request.url.path() == "/logs3")
1145            .expect("logs request present");
1146        let body: Value = serde_json::from_slice(&logs_request.body).expect("json");
1147        let rows = body.get("rows").and_then(|r| r.as_array()).unwrap();
1148        assert_eq!(rows.len(), 1);
1149        assert_eq!(
1150            rows[0].get("org_id").and_then(|v| v.as_str()),
1151            Some("org-123")
1152        );
1153    }
1154
1155    #[tokio::test]
1156    async fn span_builder_with_credentials_bypasses_login() {
1157        let server = MockServer::start().await;
1158
1159        Mock::given(method("POST"))
1160            .and(path("/api/project/register"))
1161            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
1162                "project": { "id": "proj-id" }
1163            })))
1164            .mount(&server)
1165            .await;
1166
1167        Mock::given(method("POST"))
1168            .and(path("/logs3"))
1169            .respond_with(ResponseTemplate::new(200).set_body_string("{}"))
1170            .mount(&server)
1171            .await;
1172
1173        // Create client without any login mock - background login will fail
1174        let client = BraintrustClient::builder()
1175            .api_key("test-api-key")
1176            .app_url(server.uri())
1177            .api_url(server.uri())
1178            .build()
1179            .await
1180            .expect("client");
1181
1182        // Use span_builder_with_credentials to bypass login
1183        let span = client
1184            .span_builder_with_credentials("explicit-token", "explicit-org-id")
1185            .project_name("demo-project")
1186            .build();
1187
1188        span.log(SpanLog {
1189            input: Some(Value::String("test".into())),
1190            ..Default::default()
1191        })
1192        .await;
1193        span.flush().await.expect("flush");
1194        client.flush().await.expect("client flush");
1195
1196        // Verify the logs request was made with the explicit org_id
1197        let logs_request = server
1198            .received_requests()
1199            .await
1200            .unwrap()
1201            .into_iter()
1202            .find(|request| request.url.path() == "/logs3")
1203            .expect("logs request present");
1204        let body: Value = serde_json::from_slice(&logs_request.body).expect("json");
1205        let rows = body.get("rows").and_then(|r| r.as_array()).unwrap();
1206        assert_eq!(
1207            rows[0].get("org_id").and_then(|v| v.as_str()),
1208            Some("explicit-org-id")
1209        );
1210    }
1211
1212    #[tokio::test]
1213    async fn separate_urls_for_data_and_control_plane() {
1214        // Use two separate mock servers to verify requests go to the correct URLs
1215        let api_server = MockServer::start().await; // Data plane: /logs3
1216        let app_server = MockServer::start().await; // Control plane: /api/project/register, /api/apikey/login
1217
1218        // Mount login mock on app_server
1219        Mock::given(method("POST"))
1220            .and(path("/api/apikey/login"))
1221            .respond_with(mock_login_response(&[("org-id", "Test Org")]))
1222            .mount(&app_server)
1223            .await;
1224
1225        // Mount project registration on app_server (control plane)
1226        Mock::given(method("POST"))
1227            .and(path("/api/project/register"))
1228            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
1229                "project": { "id": "proj-id" }
1230            })))
1231            .expect(1)
1232            .mount(&app_server)
1233            .await;
1234
1235        // Mount logs endpoint on api_server (data plane)
1236        Mock::given(method("POST"))
1237            .and(path("/logs3"))
1238            .respond_with(ResponseTemplate::new(200).set_body_string("{}"))
1239            .expect(1)
1240            .mount(&api_server)
1241            .await;
1242
1243        let client = BraintrustClient::builder()
1244            .api_key("test-api-key")
1245            .app_url(app_server.uri()) // Control plane
1246            .api_url(api_server.uri()) // Data plane
1247            .blocking_login(true)
1248            .build()
1249            .await
1250            .expect("client");
1251
1252        let span = client
1253            .span_builder()
1254            .await
1255            .expect("span_builder")
1256            .project_name("test-project")
1257            .build();
1258
1259        span.log(SpanLog {
1260            input: Some(Value::String("test".into())),
1261            ..Default::default()
1262        })
1263        .await;
1264        span.flush().await.expect("flush");
1265        client.flush().await.expect("client flush");
1266
1267        // Verify api_server received the /logs3 request
1268        let api_requests = api_server.received_requests().await.unwrap();
1269        assert_eq!(api_requests.len(), 1);
1270        assert_eq!(api_requests[0].url.path(), "/logs3");
1271
1272        // Verify app_server received the /api/project/register request (and login)
1273        let app_requests = app_server.received_requests().await.unwrap();
1274        let register_requests: Vec<_> = app_requests
1275            .iter()
1276            .filter(|r| r.url.path() == "/api/project/register")
1277            .collect();
1278        assert_eq!(register_requests.len(), 1);
1279
1280        // Verify app_server did NOT receive /logs3
1281        let logs_on_app: Vec<_> = app_requests
1282            .iter()
1283            .filter(|r| r.url.path() == "/logs3")
1284            .collect();
1285        assert!(logs_on_app.is_empty(), "/logs3 should not go to app_server");
1286
1287        // Verify api_server did NOT receive /api/project/register
1288        let register_on_api: Vec<_> = api_requests
1289            .iter()
1290            .filter(|r| r.url.path() == "/api/project/register")
1291            .collect();
1292        assert!(
1293            register_on_api.is_empty(),
1294            "/api/project/register should not go to api_server"
1295        );
1296    }
1297}
braintrust_sdk_rust/logger.rs

braintrust_sdk_rust/
logger.rs
