Validate trace daemon against the Rust SDK refactor

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"

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::<braintrust_sdk_rust::BraintrustError>() {
+        let sdk_error = cause
+            .downcast_ref::<braintrust_sdk_rust::BraintrustError>()
+            .or_else(
+                || match cause.downcast_ref::<braintrust_sdk_rust::FlushError>() {
+                    Some(braintrust_sdk_rust::FlushError::Delivery(failure)) => {
+                        failure.cause().cause().downcast_ref()
+                    }
+                    _ => None,
+                },
+            )
+            .or_else(
+                || match cause.downcast_ref::<braintrust_sdk_rust::WriteError>() {
+                    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 }
+            ));
+        }
+    }
+}