diff --git a/bt-daemon/Cargo.lock b/bt-daemon/Cargo.lock index b061e11d..6767554f 100644 --- a/bt-daemon/Cargo.lock +++ b/bt-daemon/Cargo.lock @@ -82,15 +82,6 @@ version = "1.0.104" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" -[[package]] -name = "arc-swap" -version = "1.9.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c049c0be4daef0b145cb3555416b3b8ef5b7888a38aea1a3a155801fe7b0810b" -dependencies = [ - "rustversion", -] - [[package]] name = "assert-json-diff" version = "2.0.2" @@ -124,29 +115,6 @@ version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" -[[package]] -name = "aws-lc-rs" -version = "1.18.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "faac5829c2b74c28f830747e7818ccfb684261b5f48a1118b1e2a13d36dfab13" -dependencies = [ - "aws-lc-sys", - "zeroize", -] - -[[package]] -name = "aws-lc-sys" -version = "0.46.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f1622d8446a2d4b2ce0c7eefc73dd43a99779028d5ee5c2dd8073a658ba8a2bc" -dependencies = [ - "cc", - "cmake", - "dunce", - "fs_extra", - "pkg-config", -] - [[package]] name = "axum" version = "0.8.9" @@ -219,12 +187,6 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" -[[package]] -name = "base64" -version = "0.23.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5" - [[package]] name = "bitflags" version = "2.13.1" @@ -268,25 +230,23 @@ 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=76fe249e7103876567c61bb72cc08505fef0a68e#76fe249e7103876567c61bb72cc08505fef0a68e" dependencies = [ "anyhow", - "arc-swap", "async-trait", "backoff", - "base64 0.23.1", + "base64", "bon", "bytes", "chrono", - "crossbeam", "futures", "indexmap", "regex", - "reqwest 0.13.5", + "reqwest", "serde", "serde_json", "serde_repr", - "thiserror", + "thiserror 1.0.69", "tokio", "tracing", "url", @@ -300,7 +260,7 @@ dependencies = [ "anyhow", "async-trait", "axum", - "base64 0.22.1", + "base64", "braintrust-sdk-rust", "bytes", "chrono", @@ -309,7 +269,7 @@ dependencies = [ "libc", "parking_lot", "regex", - "reqwest 0.12.28", + "reqwest", "rquickjs", "rquickjs-serde", "rusqlite", @@ -319,7 +279,7 @@ dependencies = [ "sha2", "sysinfo", "tempfile", - "thiserror", + "thiserror 2.0.19", "tokio", "tracing", "tracing-subscriber", @@ -430,31 +390,12 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" -[[package]] -name = "cmake" -version = "0.1.58" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c0f78a02292a74a88ac736019ab962ece0bc380e3f977bf72e376c5d78ff0678" -dependencies = [ - "cc", -] - [[package]] name = "colorchoice" version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" -[[package]] -name = "combine" -version = "4.6.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4356eba6976435498d05a311ee9fc486da3bc81c200d3f0d7dd9703b41913f92" -dependencies = [ - "bytes", - "memchr", -] - [[package]] name = "core-foundation" version = "0.10.1" @@ -489,62 +430,6 @@ 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" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" - [[package]] name = "crypto-common" version = "0.1.7" @@ -628,12 +513,6 @@ dependencies = [ "syn 2.0.119", ] -[[package]] -name = "dunce" -version = "1.0.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" - [[package]] name = "equivalent" version = "1.0.2" @@ -711,12 +590,6 @@ dependencies = [ "winapi", ] -[[package]] -name = "fs_extra" -version = "1.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" - [[package]] name = "futures" version = "0.3.33" @@ -979,6 +852,7 @@ dependencies = [ "hyper", "hyper-util", "rustls", + "rustls-native-certs", "tokio", "tokio-rustls", "tower-service", @@ -990,7 +864,7 @@ version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" dependencies = [ - "base64 0.22.1", + "base64", "bytes", "futures-channel", "futures-util", @@ -1177,55 +1051,6 @@ version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" -[[package]] -name = "jni" -version = "0.22.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5efd9a482cf3a427f00d6b35f14332adc7902ce91efb778580e180ff90fa3498" -dependencies = [ - "cfg-if", - "combine", - "jni-macros", - "jni-sys", - "log", - "simd_cesu8", - "thiserror", - "walkdir", - "windows-link", -] - -[[package]] -name = "jni-macros" -version = "0.22.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a00109accc170f0bdb141fed3e393c565b6f5e072365c3bd58f5b062591560a3" -dependencies = [ - "proc-macro2", - "quote", - "rustc_version", - "simd_cesu8", - "syn 2.0.119", -] - -[[package]] -name = "jni-sys" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c6377a88cb3910bee9b0fa88d4f42e1d2da8e79915598f65fb0c7ee14c878af2" -dependencies = [ - "jni-sys-macros", -] - -[[package]] -name = "jni-sys-macros" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "38c0b942f458fe50cdac086d2f946512305e5631e720728f2a61aabcd47a6264" -dependencies = [ - "quote", - "syn 2.0.119", -] - [[package]] name = "jobserver" version = "0.1.35" @@ -1498,7 +1323,7 @@ dependencies = [ "rustc-hash", "rustls", "socket2", - "thiserror", + "thiserror 2.0.19", "tokio", "tracing", "web-time", @@ -1510,7 +1335,6 @@ version = "0.11.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2f4bfc015262b9df63c8845072ce59068853ff5872180c2ce2f13038b970e560" dependencies = [ - "aws-lc-rs", "bytes", "getrandom 0.4.3", "lru-slab", @@ -1521,7 +1345,7 @@ dependencies = [ "rustls", "rustls-pki-types", "slab", - "thiserror", + "thiserror 2.0.19", "tinyvec", "tracing", "web-time", @@ -1665,42 +1489,7 @@ version = "0.12.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" dependencies = [ - "base64 0.22.1", - "bytes", - "futures-core", - "futures-util", - "http", - "http-body", - "http-body-util", - "hyper", - "hyper-util", - "js-sys", - "log", - "percent-encoding", - "pin-project-lite", - "serde", - "serde_json", - "serde_urlencoded", - "sync_wrapper", - "tokio", - "tokio-util", - "tower", - "tower-http", - "tower-service", - "url", - "wasm-bindgen", - "wasm-bindgen-futures", - "wasm-streams 0.4.2", - "web-sys", -] - -[[package]] -name = "reqwest" -version = "0.13.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "16a1cfa75cc186dd73d5818e510e042e40927bccc9c236b061cea97e1eb08029" -dependencies = [ - "base64 0.23.1", + "base64", "bytes", "futures-core", "futures-util", @@ -1717,10 +1506,11 @@ dependencies = [ "pin-project-lite", "quinn", "rustls", + "rustls-native-certs", "rustls-pki-types", - "rustls-platform-verifier", "serde", "serde_json", + "serde_urlencoded", "sync_wrapper", "tokio", "tokio-rustls", @@ -1731,7 +1521,7 @@ dependencies = [ "url", "wasm-bindgen", "wasm-bindgen-futures", - "wasm-streams 0.5.0", + "wasm-streams", "web-sys", ] @@ -1808,15 +1598,6 @@ version = "2.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" -[[package]] -name = "rustc_version" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" -dependencies = [ - "semver", -] - [[package]] name = "rustix" version = "1.1.4" @@ -1836,8 +1617,8 @@ version = "0.23.45" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0d41d731c7d2f962d1ccc364cec258de3c0e93b38c2fb3ba97ac74513048d634" dependencies = [ - "aws-lc-rs", "once_cell", + "ring", "rustls-pki-types", "rustls-webpki", "subtle", @@ -1866,40 +1647,12 @@ dependencies = [ "zeroize", ] -[[package]] -name = "rustls-platform-verifier" -version = "0.7.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1167586491e2b18b8bfbb293e8180ec17c201c4f076d7cb3070ca964e7598f98" -dependencies = [ - "core-foundation", - "core-foundation-sys", - "jni", - "log", - "once_cell", - "rustls", - "rustls-native-certs", - "rustls-platform-verifier-android", - "rustls-webpki", - "security-framework", - "security-framework-sys", - "webpki-root-certs", - "windows-sys 0.61.2", -] - -[[package]] -name = "rustls-platform-verifier-android" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eec689c0bc40ff2458a5977b6619cb718087084a18e02a131c599b62d05e1a5f" - [[package]] name = "rustls-webpki" version = "0.103.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4443b1b6d39d7c663139942cc4558262562d2c84123bb6282ef177009ca8b281" dependencies = [ - "aws-lc-rs", "ring", "rustls-pki-types", "untrusted", @@ -1917,15 +1670,6 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" -[[package]] -name = "same-file" -version = "1.0.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" -dependencies = [ - "winapi-util", -] - [[package]] name = "schannel" version = "0.1.29" @@ -2089,22 +1833,6 @@ dependencies = [ "libc", ] -[[package]] -name = "simd_cesu8" -version = "1.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11031e251abf8611c80f460e19dbdeb54a66db918e49c65a7065b46ac7aec520" -dependencies = [ - "rustc_version", - "simdutf8", -] - -[[package]] -name = "simdutf8" -version = "0.1.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" - [[package]] name = "slab" version = "0.4.12" @@ -2213,13 +1941,33 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + [[package]] name = "thiserror" version = "2.0.19" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09a43598840e33d5b0331f38c5e30d13bb11c11210a4b58f0d9b18a5a5eefcd9" dependencies = [ - "thiserror-impl", + "thiserror-impl 2.0.19", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", ] [[package]] @@ -2511,16 +2259,6 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" -[[package]] -name = "walkdir" -version = "2.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" -dependencies = [ - "same-file", - "winapi-util", -] - [[package]] name = "want" version = "0.3.1" @@ -2604,19 +2342,6 @@ dependencies = [ "web-sys", ] -[[package]] -name = "wasm-streams" -version = "0.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9d1ec4f6517c9e11ae630e200b2b65d193279042e28edd4a2cda233e46670bbb" -dependencies = [ - "futures-util", - "js-sys", - "wasm-bindgen", - "wasm-bindgen-futures", - "web-sys", -] - [[package]] name = "web-sys" version = "0.3.103" @@ -2637,15 +2362,6 @@ dependencies = [ "wasm-bindgen", ] -[[package]] -name = "webpki-root-certs" -version = "1.0.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "58ec80e73fd8baf57bf0d9615554c75348bda990d6d5bbeab964dd068319fecb" -dependencies = [ - "rustls-pki-types", -] - [[package]] name = "winapi" version = "0.3.9" @@ -2662,15 +2378,6 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" -[[package]] -name = "winapi-util" -version = "0.1.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" -dependencies = [ - "windows-sys 0.61.2", -] - [[package]] name = "winapi-x86_64-pc-windows-gnu" version = "0.4.0" @@ -2825,7 +2532,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "08db1edfb05d9b3c1542e521aea074442088292f00b5f28e435c714a98f85031" dependencies = [ "assert-json-diff", - "base64 0.22.1", + "base64", "deadpool", "futures", "http", diff --git a/bt-daemon/Cargo.toml b/bt-daemon/Cargo.toml index b185ac05..27a069a4 100644 --- a/bt-daemon/Cargo.toml +++ b/bt-daemon/Cargo.toml @@ -22,7 +22,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 = "76fe249e7103876567c61bb72cc08505fef0a68e" } base64 = "0.22" anyhow = "1" parking_lot = "0.12" diff --git a/bt-daemon/src/pipeline/consumer.rs b/bt-daemon/src/pipeline/consumer.rs index 64810ec0..6d85d16b 100644 --- a/bt-daemon/src/pipeline/consumer.rs +++ b/bt-daemon/src/pipeline/consumer.rs @@ -813,7 +813,7 @@ mod tests { let rejection = Mock::given(method("POST")) .and(path("/logs3")) .respond_with(ResponseTemplate::new(400).set_body_string("invalid row")) - .expect(1) + .expect(1..) .mount_as_scoped(&server) .await; let tmp = tempfile::tempdir().unwrap(); @@ -844,7 +844,7 @@ mod tests { Mock::given(method("POST")) .and(path("/logs3")) .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({}))) - .expect(1) + .expect(1..) .mount(&server) .await; // A new credential cannot cure a permanently rejected span body. diff --git a/bt-daemon/src/pipeline/submission.rs b/bt-daemon/src/pipeline/submission.rs index 634cb26f..e5deec17 100644 --- a/bt-daemon/src/pipeline/submission.rs +++ b/bt-daemon/src/pipeline/submission.rs @@ -54,7 +54,10 @@ impl std::error::Error for Failure { impl Failure { fn classify(cause: anyhow::Error) -> Self { - let retry = match cause.downcast_ref::() { + let retry = match cause + .chain() + .find_map(|source| source.downcast_ref::()) + { Some(braintrust_sdk_rust::BraintrustError::Api { status: 401, .. }) => { RetryCondition::Credentials } @@ -68,6 +71,8 @@ impl Failure { | Some(braintrust_sdk_rust::BraintrustError::Http(_)) | Some(braintrust_sdk_rust::BraintrustError::Network(_)) | Some(braintrust_sdk_rust::BraintrustError::ChannelClosed) + | Some(braintrust_sdk_rust::BraintrustError::Admission(_)) + | Some(braintrust_sdk_rust::BraintrustError::Execution(_)) | None => RetryCondition::Transient, Some(_) => RetryCondition::Configuration { recheck: false }, }; diff --git a/bt-daemon/src/sink/braintrust.rs b/bt-daemon/src/sink/braintrust.rs index 445cbbfd..dce2f3b7 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 observes its own SDK 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 @@ -46,26 +46,55 @@ struct ClientCacheState { struct CachedClient { client: BraintrustClient, - queued_bytes: AsyncMutex, + pending: AsyncMutex, } +type PendingWrite = std::pin::Pin> + Send>>; +const MAX_PENDING_WRITES: usize = 256; +#[derive(Default)] +struct PendingDelivery { + bytes: usize, + writes: Vec, +} +impl PendingDelivery { + async fn drain(&mut self, client: &BraintrustClient) -> anyhow::Result<()> { + let flush = client.flush(); + let writes = std::mem::take(&mut self.writes); + self.bytes = 0; + let mut failure = None; + // Observe every accepted token, including failures that retired before + // this flush's cutoff. Journal checkpoints must never rely on a queue + // retaining historical failures. The journal owns retries on failure. + for write in writes { + if let Err(error) = write.await { + failure.get_or_insert(error); + } + } + if let Err(error) = flush.await { + failure.get_or_insert_with(|| anyhow::Error::new(error)); + } + match failure { + Some(error) => Err(error), + None => Ok(()), + } + } +} impl CachedClient { - async fn account_and_flush(&self, bytes: usize) -> anyhow::Result<()> { - 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(()); + async fn account_and_flush( + &self, + bytes: usize, + writes: Vec, + ) -> anyhow::Result<()> { + let mut pending = self.pending.lock().await; + pending.bytes = pending.bytes.saturating_add(bytes); + pending.writes.extend(writes); + if pending.bytes >= QUEUE_FLUSH_BYTE_BUDGET || pending.writes.len() >= MAX_PENDING_WRITES { + pending.drain(&self.client).await?; } - self.client.flush().await.map_err(anyhow::Error::new)?; - *queued_bytes = 0; Ok(()) } - async fn flush(&self) -> anyhow::Result<()> { - let mut queued_bytes = self.queued_bytes.lock().await; - self.client.flush().await.map_err(anyhow::Error::new)?; - *queued_bytes = 0; - Ok(()) + self.pending.lock().await.drain(&self.client).await } } @@ -78,9 +107,13 @@ impl ClientCacheState { fn trim_inactive(&mut self) { while self.clients.len() >= CLIENT_CACHE_CAPACITY { let Some(index) = self.order.iter().position(|key| { - self.clients - .get(key) - .is_some_and(|client| Arc::strong_count(client) == 1) + self.clients.get(key).is_some_and(|client| { + Arc::strong_count(client) == 1 + && client + .pending + .try_lock() + .is_ok_and(|pending| pending.writes.is_empty()) + }) }) else { // Every client is still referenced by a live session; those are // active working state and cannot be evicted safely. @@ -100,6 +133,20 @@ impl ClientCache { } } + async fn flush(&self) -> anyhow::Result<()> { + let clients: Vec<_> = self.state.lock().await.clients.values().cloned().collect(); + let mut failure = None; + for client in clients { + if let Err(error) = client.flush().await { + failure.get_or_insert(error); + } + } + match failure { + Some(error) => Err(error), + None => Ok(()), + } + } + async fn get_or_build( &self, api_url: &str, @@ -115,7 +162,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() @@ -123,7 +169,7 @@ impl ClientCache { .map_err(|e| anyhow::anyhow!("braintrust client build failed: {e}"))?; let arc = Arc::new(CachedClient { client, - queued_bytes: AsyncMutex::new(0), + pending: AsyncMutex::new(PendingDelivery::default()), }); state.trim_inactive(); state.clients.insert(key.clone(), arc.clone()); @@ -195,7 +241,7 @@ struct BraintrustSink { client: Option>, /// Live span handles keyed by deterministic span id, so a later op (e.g. /// setting `end`) merges onto the same row the SDK already knows. - open: HashMap>, + open: HashMap, /// Roots marked by translators as provisional, keyed by their local span /// id. Each is exported only once direct descendant work arrives. pending_roots: HashMap>, @@ -276,23 +322,37 @@ impl BraintrustSink { Ok(()) } - fn update_open(&mut self, client: &BraintrustClient, row: &SpanRow) -> anyhow::Result<()> { + fn update_open( + &mut self, + client: &BraintrustClient, + row: &SpanRow, + ) -> anyhow::Result> { self.ensure_handle(client, row)?; 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)?); + let write = handle.log(build_log(row, &self.source, false)?); + let mut writes: Vec = vec![Box::pin(async move { + write.await.map(|_| ()).map_err(anyhow::Error::new) + })]; if let Some(end) = row.end_ms { - handle.end_with_time(ms_to_secs(end)); + let end = handle.end_with_time(ms_to_secs(end)); + writes.push(Box::pin(async move { + end.await.map(|_| ()).map_err(anyhow::Error::new) + })); // 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. self.open.remove(&row.span_id); } - Ok(()) + Ok(writes) } - fn merge_closed(&self, client: &BraintrustClient, row: &SpanRow) -> anyhow::Result<()> { + fn merge_closed( + &self, + client: &BraintrustClient, + row: &SpanRow, + ) -> anyhow::Result> { let creds = self .creds .as_ref() @@ -311,14 +371,15 @@ impl BraintrustSink { context.remove("span_origin"); } } - client - .update_span_with_credentials( - 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}")) + let write = client.update_span_with_credentials( + creds.token.clone(), + creds.org_id.clone(), + &components.to_str(), + build_log(row, &self.source, true)?, + ); + Ok(vec![Box::pin(async move { + write.await.map_err(anyhow::Error::new) + })]) } fn span_components(&self, row: &SpanRow, creds: &Creds, project: &str) -> SpanComponents { @@ -350,15 +411,15 @@ impl BraintrustSink { async fn emit_op(&mut self, op: &SpanOp) -> anyhow::Result<()> { let client = self.ensure_client().await?; - match op { + let writes = match op { SpanOp::Insert(row) => self.update_open(&client.client, row)?, SpanOp::Merge(row) if self.open.contains_key(&row.span_id) => { self.update_open(&client.client, row)? } SpanOp::Merge(row) => self.merge_closed(&client.client, row)?, - } + }; client - .account_and_flush(serialized_len(op).unwrap_or(0)) + .account_and_flush(serialized_len(op).unwrap_or(0), writes) .await } } @@ -443,10 +504,9 @@ impl Sink for BraintrustSink { } async fn flush(&mut self) -> anyhow::Result<()> { - match &self.client { - Some(client) => client.flush().await, - None => Ok(()), - } + // Configuration changes can leave accepted work on a previous URL. + // Observe every client owned by this branch before its checkpoint. + self.cache.flush().await } fn has_pending_delivery(&self) -> bool { @@ -729,7 +789,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")) @@ -773,6 +832,7 @@ mod tests { &components.to_str(), build_log(&row, "codex", true).unwrap(), ) + .await .unwrap(); client.flush().await.unwrap(); @@ -799,4 +859,70 @@ mod tests { ); } } + #[tokio::test] + async fn branch_observes_retired_write_failure_then_replays_fresh_work() { + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(400)) + .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("test-token", "test-org") + .parent_info(ParentSpanInfo::ProjectLogs { + object_id: "test-project".into(), + }) + .build(); + let write = span.log(SpanLog::default()); + assert!(client.flush().await.is_err()); + // The SDK has retired the failed batch. Its next flush cannot report it. + client.flush().await.unwrap(); + let cached = CachedClient { + client, + pending: AsyncMutex::new(PendingDelivery::default()), + }; + cached + .account_and_flush( + 1, + vec![Box::pin(async move { + write.await.map(|_| ()).map_err(anyhow::Error::new) + })], + ) + .await + .unwrap(); + assert!(cached.flush().await.is_err()); + assert!(cached.pending.lock().await.writes.is_empty()); + cached.flush().await.unwrap(); + server.reset().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(200)) + .mount(&server) + .await; + let write = span.log(SpanLog::default()); + cached + .account_and_flush( + 1, + vec![Box::pin(async move { + write.await.map(|_| ()).map_err(anyhow::Error::new) + })], + ) + .await + .unwrap(); + cached.flush().await.unwrap(); + cached + .client + .shutdown(std::time::Duration::from_secs(5)) + .await + .unwrap(); + } } diff --git a/bt-daemon/tests/pipeline.rs b/bt-daemon/tests/pipeline.rs index 58f1bd5c..461f8313 100644 --- a/bt-daemon/tests/pipeline.rs +++ b/bt-daemon/tests/pipeline.rs @@ -3947,10 +3947,12 @@ async fn changed_settings_resume_the_original_rejected_route_through_ingress() { .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({}))) .mount(&server) .await; + // Automatic delivery may send the root and event snapshots in separate + // physical requests; the checkpoint/recovery contract is batch-independent. let rejection = Mock::given(method("POST")) .and(path("/logs3")) .respond_with(ResponseTemplate::new(400).set_body_string("invalid row")) - .expect(1) + .expect(1..) .mount_as_scoped(&server) .await; let tmp = tempfile::tempdir().unwrap();