Unify SDK ingestion upload orchestration

src/api/client.rs

@@ -34,18 +34,7 @@ impl Client {
         if config.api_key.trim().is_empty() {
             return Err(Error::Config("api_key must not be empty".into()));
         }
-        let mut headers = HeaderMap::new();
-        let mut token = HeaderValue::from_str(&format!("Bearer {}", config.api_key))
-            .map_err(|_| Error::Config("invalid api_key header".into()))?;
-        token.set_sensitive(true);
-        headers.insert(AUTHORIZATION, token);
-        if let Some(org) = config.org_name {
-            headers.insert(
-                "x-bt-org-name",
-                HeaderValue::from_str(&org)
-                    .map_err(|_| Error::Config("invalid organization header".into()))?,
-            );
-        }
+        let headers = scoped_headers(&config.api_key, config.org_name.as_deref())?;
         let ca = match (config.custom_ca_bundle_path, config.custom_ca_bundle_pem) {
             (Some(_), Some(_)) => {
                 return Err(Error::Config(
@@ -72,6 +61,26 @@ impl Client {
             }),
         })
     }
+    /// Reuse the SDK's configured transport without copying session state into
+    /// the API layer. This scope is immutable and used only for ingestion.
+    pub(crate) fn for_ingestion(
+        http: reqwest::Client,
+        api_url: &Url,
+        token: &str,
+        org_name: Option<&str>,
+    ) -> Result<Self, Error> {
+        let api_url = base_url(api_url.as_str())?;
+        let headers = scoped_headers(token, org_name)?;
+        Ok(Self {
+            inner: Arc::new(Inner {
+                http,
+                app_url: api_url.clone(),
+                api_url,
+                headers,
+            }),
+        })
+    }
+
     pub fn logs3(&self) -> super::logs3::Logs3 {
         super::logs3::Logs3 {
             client: self.clone(),
@@ -199,6 +208,22 @@ impl Client {
         Ok(())
     }
 }
+fn scoped_headers(token: &str, org_name: Option<&str>) -> Result<HeaderMap, Error> {
+    let mut headers = HeaderMap::new();
+    let mut token = HeaderValue::from_str(&format!("Bearer {token}"))
+        .map_err(|_| Error::Config("invalid api_key header".into()))?;
+    token.set_sensitive(true);
+    headers.insert(AUTHORIZATION, token);
+    if let Some(org) = org_name {
+        headers.insert(
+            "x-bt-org-name",
+            HeaderValue::from_str(org)
+                .map_err(|_| Error::Config("invalid organization header".into()))?,
+        );
+    }
+    Ok(headers)
+}
+
 fn base_url(value: &str) -> Result<Url, Error> {
     let mut url = Url::parse(value).map_err(|_| Error::Config("invalid endpoint URL".into()))?;
     if !matches!(url.scheme(), "http" | "https")

src/http.rs

@@ -69,6 +69,8 @@ pub(crate) fn build_http_client_builder(
     user_agent: Option<&str>,
 ) -> Result<reqwest::ClientBuilder> {
     let mut builder = Client::builder()
+        // SDK retries own request repetition; redirects must not add attempts.
+        .redirect(reqwest::redirect::Policy::none())
         .timeout(timeout)
         .user_agent(user_agent.unwrap_or(USER_AGENT));
 

src/internal/mod.rs

@@ -1,3 +1,5 @@
 pub(crate) mod batching;
 pub(crate) mod ingest;
+pub(crate) mod object_store;
 pub(crate) mod retry;
+pub(crate) mod upload;

src/internal/object_store.rs

@@ -0,0 +1,84 @@
+//! One-attempt signed object uploads, independent of Braintrust credentials.
+use crate::api::logs3::Logs3OverflowUpload;
+use crate::error::{BraintrustError, Result};
+use bytes::Bytes;
+use url::Url;
+
+/// Only signed grant information reaches object storage. Braintrust credentials
+/// are scoped to the API client, never added by this transport.
+pub(crate) struct ObjectStoreTransport {
+    pub(crate) http: reqwest::Client,
+}
+impl ObjectStoreTransport {
+    pub(crate) async fn upload(&self, grant: &Logs3OverflowUpload, payload: Bytes) -> Result<()> {
+        let url = Url::parse(&grant.signed_url)
+            .map_err(|_| BraintrustError::InvalidConfig("invalid overflow signed URL".into()))?;
+        let mut headers = reqwest::header::HeaderMap::new();
+        if let Some(pairs) = &grant.headers {
+            for (key, value) in pairs {
+                let key =
+                    reqwest::header::HeaderName::from_bytes(key.as_bytes()).map_err(|_| {
+                        BraintrustError::InvalidConfig("invalid overflow header name".into())
+                    })?;
+                let value = reqwest::header::HeaderValue::from_str(value).map_err(|_| {
+                    BraintrustError::InvalidConfig("invalid overflow header value".into())
+                })?;
+                headers.insert(key, value);
+            }
+        }
+        let request = if grant.method.eq_ignore_ascii_case("POST") {
+            let fields = grant.fields.as_ref().ok_or_else(|| {
+                BraintrustError::InvalidConfig("overflow POST upload missing fields".into())
+            })?;
+            let content_type = fields
+                .get("Content-Type")
+                .map(String::as_str)
+                .unwrap_or("application/json");
+            let mut form = reqwest::multipart::Form::new();
+            for (key, value) in fields {
+                form = form.text(key.clone(), value.clone());
+            }
+            form = form.part(
+                "file",
+                reqwest::multipart::Part::bytes(payload.to_vec())
+                    .mime_str(content_type)
+                    .map_err(|_| {
+                        BraintrustError::InvalidConfig("invalid overflow content type".into())
+                    })?,
+            );
+            headers.remove(reqwest::header::CONTENT_TYPE);
+            self.http.post(url).headers(headers).multipart(form)
+        } else if grant.method.eq_ignore_ascii_case("PUT") {
+            if url
+                .host_str()
+                .is_some_and(|host| host.ends_with(".blob.core.windows.net"))
+                && !headers.contains_key("x-ms-blob-type")
+            {
+                headers.insert(
+                    "x-ms-blob-type",
+                    reqwest::header::HeaderValue::from_static("BlockBlob"),
+                );
+            }
+            self.http
+                .put(url)
+                .header("content-type", "application/json")
+                .headers(headers)
+                .body(payload)
+        } else {
+            return Err(BraintrustError::InvalidConfig(
+                "unsupported overflow upload method".into(),
+            ));
+        };
+        let response = request
+            .send()
+            .await
+            .map_err(|error| BraintrustError::Http(error.without_url()))?;
+        if !response.status().is_success() {
+            return Err(BraintrustError::Api {
+                status: response.status().as_u16(),
+                message: response.text().await.unwrap_or_default(),
+            });
+        }
+        Ok(())
+    }
+}

src/internal/retry.rs

@@ -65,6 +65,23 @@ impl<B: Iterator<Item = Duration>, S: Sleeper> RetryExecutor<B, S> {
 
     pub(crate) async fn run<T, E, F, Fut, C, N>(
         mut self,
+        operation: F,
+        retryable: C,
+        on_retry: N,
+    ) -> RetryOutcome<T, E>
+    where
+        F: FnMut(Attempt) -> Fut,
+        Fut: Future<Output = Result<T, E>>,
+        C: FnMut(&E) -> bool,
+        N: FnMut(Attempt, &E, Duration),
+    {
+        self.run_shared(operation, retryable, on_retry).await
+    }
+
+    /// Execute one stage, sharing the remaining retry allowance with later
+    /// stages. Successful stages do not consume retries or reset the schedule.
+    pub(crate) async fn run_shared<T, E, F, Fut, C, N>(
+        &mut self,
         mut operation: F,
         mut retryable: C,
         mut on_retry: N,
@@ -75,10 +92,13 @@ impl<B: Iterator<Item = Duration>, S: Sleeper> RetryExecutor<B, S> {
         C: FnMut(&E) -> bool,
         N: FnMut(Attempt, &E, Duration),
     {
-        // Construction always provides the initial attempt, even for an empty
-        // delay schedule. A classified error is returned without cloning it.
-        while let Some(attempt) = self.next().await {
-            let attempts = attempt.index.saturating_add(1);
+        if self.next_index == Some(0) {
+            self.next().await;
+        }
+        let mut attempts: usize = 0;
+        loop {
+            let attempt = Attempt { index: attempts };
+            attempts = attempts.saturating_add(1);
             match operation(attempt).await {
                 Ok(value) => {
                     return RetryOutcome {
@@ -87,7 +107,7 @@ impl<B: Iterator<Item = Duration>, S: Sleeper> RetryExecutor<B, S> {
                     }
                 }
                 Err(error) => {
-                    if attempt.index == self.max_retries || !retryable(&error) {
+                    if !retryable(&error) || self.next_index.is_none() {
                         return RetryOutcome {
                             result: Err(error),
                             attempts,
@@ -100,10 +120,15 @@ impl<B: Iterator<Item = Duration>, S: Sleeper> RetryExecutor<B, S> {
                         };
                     };
                     on_retry(attempt, &error, delay);
+                    if self.next().await.is_none() {
+                        return RetryOutcome {
+                            result: Err(error),
+                            attempts,
+                        };
+                    }
                 }
             }
         }
-        unreachable!("initial attempt exists and a terminal error returns immediately")
     }
 }
 
@@ -146,6 +171,46 @@ mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn stages_share_retries_without_resetting_the_delay_schedule() {
+        let sleeper = FakeSleeper::default();
+        let mut executor = RetryExecutor::with_sleeper(
+            2,
+            exponential(Duration::from_secs(1), Duration::from_secs(8)),
+            sleeper.clone(),
+        );
+        for _ in 0..2 {
+            let outcome = executor
+                .run_shared(
+                    |attempt| async move {
+                        if attempt.index == 0 {
+                            Err("transient")
+                        } else {
+                            Ok(())
+                        }
+                    },
+                    |_| true,
+                    |_, _, _| {},
+                )
+                .await;
+            assert_eq!(outcome.attempts, 2);
+            assert_eq!(outcome.result, Ok(()));
+        }
+        let outcome = executor
+            .run_shared(
+                |_| async { Err::<(), _>("exhausted") },
+                |_| true,
+                |_, _, _| {},
+            )
+            .await;
+        assert_eq!(outcome.attempts, 1);
+        assert_eq!(outcome.result, Err("exhausted"));
+        assert_eq!(
+            *sleeper.delays.lock().unwrap(),
+            vec![Duration::from_secs(1), Duration::from_secs(2)]
+        );
+    }
+
     #[tokio::test]
     async fn zero_retries_still_attempts_once_without_sleeping() {
         let sleeper = FakeSleeper::default();