diff --git a/bt-daemon/Cargo.lock b/bt-daemon/Cargo.lock index 4699e9ae..9cfaf31c 100644 --- a/bt-daemon/Cargo.lock +++ b/bt-daemon/Cargo.lock @@ -84,9 +84,9 @@ checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" [[package]] name = "arc-swap" -version = "1.9.2" +version = "1.9.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c049c0be4daef0b145cb3555416b3b8ef5b7888a38aea1a3a155801fe7b0810b" +checksum = "380962b3d93b99642e50b9febf72c071e97e8b32d86eae9cf4ef582bb960ae6e" dependencies = [ "rustversion", ] @@ -126,9 +126,9 @@ checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" [[package]] name = "aws-lc-rs" -version = "1.18.2" +version = "1.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "faac5829c2b74c28f830747e7818ccfb684261b5f48a1118b1e2a13d36dfab13" +checksum = "048006790404103c488e6c5616b3e7e702a71292aadda08e8f7698a9e0d2eb53" dependencies = [ "aws-lc-sys", "zeroize", @@ -136,9 +136,9 @@ dependencies = [ [[package]] name = "aws-lc-sys" -version = "0.46.0" +version = "0.47.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f1622d8446a2d4b2ce0c7eefc73dd43a99779028d5ee5c2dd8073a658ba8a2bc" +checksum = "8bfbe32ac36b3c428696ce315464181bc4029d689fd58e3f2a4e58c252242f52" dependencies = [ "cc", "cmake", @@ -268,7 +268,7 @@ dependencies = [ [[package]] name = "braintrust-sdk-rust" version = "0.1.0-alpha.2" -source = "git+https://github.com/braintrustdata/braintrust-sdk-rust?rev=45a2c88ca9d9ccc393ec0338ac438199fc3c5df2#45a2c88ca9d9ccc393ec0338ac438199fc3c5df2" +source = "git+https://github.com/braintrustdata/braintrust-sdk-rust?rev=1e4f901fe44b028f9a3ea59a204d7651cc38acfc#1e4f901fe44b028f9a3ea59a204d7651cc38acfc" dependencies = [ "anyhow", "arc-swap", @@ -278,11 +278,12 @@ dependencies = [ "bon", "bytes", "chrono", - "crossbeam", + "crossbeam-utils", "futures", "indexmap", "regex", "reqwest 0.13.5", + "sdd", "serde", "serde_json", "serde_repr", @@ -496,61 +497,11 @@ dependencies = [ "libc", ] -[[package]] -name = "crossbeam" -version = "0.8.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1137cd7e7fc0fb5d3c5a8678be38ec56e819125d8d7907411fe24ccb943faca8" -dependencies = [ - "crossbeam-channel", - "crossbeam-deque", - "crossbeam-epoch", - "crossbeam-queue", - "crossbeam-utils", -] - -[[package]] -name = "crossbeam-channel" -version = "0.5.16" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d85363c37faeca707aef026efa9f3b34d077bce547e48f770770625c6013679e" -dependencies = [ - "crossbeam-utils", -] - -[[package]] -name = "crossbeam-deque" -version = "0.8.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" -dependencies = [ - "crossbeam-epoch", - "crossbeam-utils", -] - -[[package]] -name = "crossbeam-epoch" -version = "0.9.20" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" -dependencies = [ - "crossbeam-utils", -] - -[[package]] -name = "crossbeam-queue" -version = "0.3.13" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "803d13fb3b09d88be9f4dbc29062c66b19bf7170867ceb746d2a8689bf6c7a26" -dependencies = [ - "crossbeam-utils", -] - [[package]] name = "crossbeam-utils" -version = "0.8.22" +version = "0.8.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" +checksum = "a31eee39dddec8330830986fcd7625edb5a24ec90ea038215273bbc3adb08ac6" [[package]] name = "crypto-common" @@ -1753,6 +1704,7 @@ dependencies = [ "rustls-platform-verifier", "serde", "serde_json", + "serde_urlencoded", "sync_wrapper", "tokio", "tokio-rustls", @@ -1949,6 +1901,12 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "saa" +version = "5.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bd85992a680de788217d9c70b74aa7685b108cf1aba02f728fb27f1b27199ad1" + [[package]] name = "same-file" version = "1.0.6" @@ -1973,6 +1931,15 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "sdd" +version = "4.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a31877eee03686b905845993b5d9809a220c9dd5a10e9d6021d7ca67384c9070" +dependencies = [ + "saa", +] + [[package]] name = "security-framework" version = "3.7.0" @@ -2671,9 +2638,9 @@ dependencies = [ [[package]] name = "webpki-root-certs" -version = "1.0.10" +version = "1.0.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "58ec80e73fd8baf57bf0d9615554c75348bda990d6d5bbeab964dd068319fecb" +checksum = "e49d0e65abc00c142e8f9428da991da9f1a250c39be97459901eb12b9b49f62a" dependencies = [ "rustls-pki-types", ] diff --git a/bt-daemon/Cargo.toml b/bt-daemon/Cargo.toml index d38f914c..0817618d 100644 --- a/bt-daemon/Cargo.toml +++ b/bt-daemon/Cargo.toml @@ -23,7 +23,7 @@ expect_used = "deny" [dependencies] # Keep this pinned exactly, matching the dependency policy used by `bt`. -braintrust-sdk-rust = { git = "https://github.com/braintrustdata/braintrust-sdk-rust", rev = "45a2c88ca9d9ccc393ec0338ac438199fc3c5df2" } +braintrust-sdk-rust = { git = "https://github.com/braintrustdata/braintrust-sdk-rust", rev = "1e4f901fe44b028f9a3ea59a204d7651cc38acfc" } base64 = "0.22" anyhow = "1" muse-code-sdk = "=1.4.12" diff --git a/bt-daemon/src/pipeline/submission.rs b/bt-daemon/src/pipeline/submission.rs index 634cb26f..6fc401c2 100644 --- a/bt-daemon/src/pipeline/submission.rs +++ b/bt-daemon/src/pipeline/submission.rs @@ -54,7 +54,25 @@ impl std::error::Error for Failure { impl Failure { fn classify(cause: anyhow::Error) -> Self { - let retry = match cause.downcast_ref::() { + let sdk_error = cause + .downcast_ref::() + .or_else( + || match cause.downcast_ref::() { + Some(braintrust_sdk_rust::FlushError::Delivery(failure)) => { + failure.cause().cause().downcast_ref() + } + _ => None, + }, + ) + .or_else( + || match cause.downcast_ref::() { + Some(braintrust_sdk_rust::WriteError::Delivery(failure)) => { + failure.cause().downcast_ref() + } + _ => None, + }, + ); + let retry = match sdk_error { Some(braintrust_sdk_rust::BraintrustError::Api { status: 401, .. }) => { RetryCondition::Credentials } @@ -121,3 +139,47 @@ impl Stage for Submission { } } } + +#[cfg(test)] +mod tests { + use super::*; + use braintrust_sdk_rust::{BraintrustClient, ParentSpanInfo, SpanLog}; + use wiremock::{ + matchers::{method, path}, + Mock, MockServer, ResponseTemplate, + }; + + #[tokio::test] + async fn classifies_delivery_errors_from_both_flush_and_write() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with( + ResponseTemplate::new(400).set_delay(std::time::Duration::from_millis(50)), + ) + .mount(&server) + .await; + let client = BraintrustClient::builder() + .skip_login(true) + .api_url(server.uri()) + .app_url(server.uri()) + .build() + .await + .unwrap(); + let span = client + .span_builder_with_credentials("token", "org") + .parent_info(ParentSpanInfo::ProjectLogs { + object_id: "project".into(), + }) + .build(); + let write = span.log(SpanLog::builder().build().unwrap()); + let flush = client.flush().await.unwrap_err(); + let write = write.await.unwrap_err(); + for error in [anyhow::Error::new(flush), anyhow::Error::new(write)] { + assert!(matches!( + Failure::classify(error).retry, + RetryCondition::Configuration { recheck: false } + )); + } + } +} diff --git a/bt-daemon/src/sink/braintrust.rs b/bt-daemon/src/sink/braintrust.rs index 4f036eba..8b019dfb 100644 --- a/bt-daemon/src/sink/braintrust.rs +++ b/bt-daemon/src/sink/braintrust.rs @@ -1,6 +1,6 @@ //! The Braintrust sink: maps sink-neutral [`SpanOp`]s onto `braintrust-sdk-rust`. //! -//! Each routing branch owns its checked SDK queue. Clients are built lazily +//! Each routing branch owns its SDK queue and write completions. Clients are built lazily //! using its configured backend URLs; flush results cannot include another //! branch's rejected operations. Within a client, each //! session's token/org travel per span (`span_builder_with_credentials`) and @@ -12,8 +12,8 @@ use super::{Sink, SinkFactory}; use crate::translate::{SpanOp, SpanRow, SpanType}; use crate::wire::{SessionConfig, TraceDestination}; use braintrust_sdk_rust::{ - BraintrustClient, ParentSpanInfo, SpanComponents, SpanHandle, SpanLog, SpanObjectType, - SpanType as SdkSpanType, DEFAULT_API_URL, DEFAULT_APP_URL, + BraintrustClient, Completion, DeliveryReceipt, ParentSpanInfo, SpanComponents, SpanHandle, + SpanLog, SpanObjectType, SpanType as SdkSpanType, WriteError, DEFAULT_API_URL, DEFAULT_APP_URL, }; use serde_json::{json, Map, Value}; use std::collections::{HashMap, VecDeque}; @@ -37,6 +37,7 @@ struct ClientCache { const CLIENT_CACHE_CAPACITY: usize = 16; const QUEUE_FLUSH_BYTE_BUDGET: usize = 8 * 1024 * 1024; +const MAX_PENDING_COMPLETIONS: usize = 1024; #[derive(Default)] struct ClientCacheState { @@ -50,15 +51,10 @@ struct CachedClient { } impl CachedClient { - async fn account_and_flush(&self, bytes: usize) -> anyhow::Result<()> { + async fn account(&self, bytes: usize) -> bool { let mut queued_bytes = self.queued_bytes.lock().await; *queued_bytes = queued_bytes.saturating_add(bytes); - if *queued_bytes < QUEUE_FLUSH_BYTE_BUDGET { - return Ok(()); - } - self.client.flush().await.map_err(anyhow::Error::new)?; - *queued_bytes = 0; - Ok(()) + *queued_bytes >= QUEUE_FLUSH_BYTE_BUDGET } async fn flush(&self) -> anyhow::Result<()> { @@ -115,7 +111,6 @@ impl ClientCache { } let client = BraintrustClient::builder() .skip_login(true) - .checked_delivery(true) .api_url(api_url.to_string()) .app_url(app_url.to_string()) .build() @@ -161,6 +156,7 @@ impl SinkFactory for BraintrustSinkFactory { client: None, open: HashMap::new(), pending_roots: HashMap::new(), + pending_completions: Vec::new(), delivered_root: None, })) } @@ -200,6 +196,8 @@ struct BraintrustSink { /// Roots marked by translators as provisional, keyed by their local span /// id. Each is exported only once direct descendant work arrives. pending_roots: HashMap>, + /// Retain acknowledgements even when automatic delivery finishes before flush. + pending_completions: Vec>, /// Root span id of the trace that most recently received a row, which /// is also that root's deterministic row id in Braintrust. delivered_root: Option, @@ -285,9 +283,13 @@ impl BraintrustSink { let handle = self.open.get(&row.span_id).ok_or_else(|| { anyhow::anyhow!("span handle was not initialized for {}", row.span_id) })?; - handle.log(build_log(row, &self.source, false)?); - if let Some(end) = row.end_ms { - handle.end_with_time(ms_to_secs(end)); + let completion = handle.log(build_log(row, &self.source, false)?); + if let Err(error) = completion.submission() { + return Err(anyhow::anyhow!("braintrust span admission failed: {error}")); + } + self.pending_completions.push(completion); + if row.end_ms.is_some() { + // build_log includes the end metric in this tracked write. // SpanHandle retains the complete accumulated input/output. Once a // span is terminal, use stateless SDK merges for any unusually late // update instead of pinning that payload for the session lifetime. @@ -296,7 +298,7 @@ impl BraintrustSink { Ok(()) } - fn merge_closed(&self, client: &BraintrustClient, row: &SpanRow) -> anyhow::Result<()> { + fn merge_closed(&mut self, client: &BraintrustClient, row: &SpanRow) -> anyhow::Result<()> { let creds = self .creds .as_ref() @@ -315,14 +317,16 @@ impl BraintrustSink { context.remove("span_origin"); } } - client - .update_span_with_credentials( + let completion = client + .update_span_with_credentials_tracked( creds.token.clone(), creds.org_id.clone(), &components.to_str(), build_log(row, &self.source, true)?, ) - .map_err(|error| anyhow::anyhow!("braintrust span merge failed: {error}")) + .map_err(|error| anyhow::anyhow!("braintrust span merge failed: {error}"))?; + self.pending_completions.push(completion); + Ok(()) } fn span_components(&self, row: &SpanRow, creds: &Creds, project: &str) -> SpanComponents { @@ -361,9 +365,11 @@ impl BraintrustSink { } SpanOp::Merge(row) => self.merge_closed(&client.client, row)?, } - client - .account_and_flush(serialized_len(op).unwrap_or(0)) - .await + let over_budget = client.account(serialized_len(op).unwrap_or(0)).await; + if over_budget || self.pending_completions.len() >= MAX_PENDING_COMPLETIONS { + self.flush().await?; + } + Ok(()) } } @@ -451,10 +457,18 @@ impl Sink for BraintrustSink { } async fn flush(&mut self) -> anyhow::Result<()> { - match &self.client { - Some(client) => client.flush().await, - None => Ok(()), + let mut failure = match &self.client { + Some(client) => client.flush().await.err(), + None => None, + }; + for completion in self.pending_completions.drain(..) { + if let Err(error) = completion.await { + if failure.is_none() { + failure = Some(anyhow::Error::new(error)); + } + } } + failure.map_or(Ok(()), Err) } fn has_pending_delivery(&self) -> bool { @@ -771,6 +785,53 @@ mod tests { use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; + #[tokio::test] + async fn sink_flush_observes_failure_already_retired_by_sdk() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(401)) + .mount(&server) + .await; + let client = BraintrustClient::builder() + .skip_login(true) + .api_url(server.uri()) + .app_url(server.uri()) + .build() + .await + .unwrap(); + let handle = client + .span_builder_with_credentials("token", "org") + .parent_info(ParentSpanInfo::ProjectLogs { + object_id: "project".into(), + }) + .build(); + let completion = handle.log(SpanLog::builder().input(json!("value")).build().unwrap()); + // Complete and retire the failed batch before the sink issues its flush. + // Depending on scheduling, the automatic worker may already have + // retired the failure, in which case this barrier succeeds. + let _ = client.flush().await; + let cached = Arc::new(CachedClient { + client, + queued_bytes: AsyncMutex::new(0), + }); + let mut sink = BraintrustSink { + cache: Arc::new(ClientCache::new()), + default_api_url: None, + default_app_url: None, + source: "codex".into(), + creds: None, + urls: None, + client: Some(cached), + open: HashMap::new(), + pending_roots: HashMap::new(), + pending_completions: vec![completion], + delivered_root: None, + }; + assert!(sink.flush().await.is_err()); + assert!(sink.flush().await.is_ok()); + } + #[tokio::test] async fn root_context_retains_sdk_environment_on_insert_and_stateless_merge() { let server = MockServer::start().await; @@ -786,7 +847,6 @@ mod tests { .await; let client = BraintrustClient::builder() .skip_login(true) - .checked_delivery(true) .api_url(server.uri()) .app_url(server.uri()) .environment("test", Some("sink-test"))