synth_ai_core/
trace_upload.rs

1//! Trace upload helpers for large traces via presigned URLs.
2
3use serde::{Deserialize, Serialize};
4use serde_json::Value;
5use std::time::Duration;
6
7use crate::http::{HttpClient, HttpError};
8use crate::CoreError;
9
10/// Response from creating an upload URL.
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct UploadUrlResponse {
13    pub trace_id: String,
14    pub trace_ref: String,
15    pub upload_url: String,
16    pub expires_in_seconds: i64,
17    pub storage_key: String,
18    #[serde(default)]
19    pub max_size_bytes: Option<i64>,
20}
21
22/// Client for trace upload endpoints.
23pub struct TraceUploadClient {
24    http: HttpClient,
25    base_url: String,
26    timeout_secs: u64,
27}
28
29impl TraceUploadClient {
30    pub fn new(base_url: &str, api_key: &str, timeout_secs: u64) -> Result<Self, CoreError> {
31        let http = HttpClient::new(base_url, api_key, timeout_secs)
32            .map_err(|e| CoreError::Internal(format!("failed to create http client: {}", e)))?;
33        Ok(Self {
34            http,
35            base_url: base_url.trim_end_matches('/').to_string(),
36            timeout_secs,
37        })
38    }
39
40    pub fn trace_size(trace: &Value) -> Result<usize, CoreError> {
41        let payload = serde_json::to_string(trace)
42            .map_err(|e| CoreError::Validation(format!("failed to serialize trace: {}", e)))?;
43        Ok(payload.len())
44    }
45
46    pub fn should_upload(trace: &Value, threshold_bytes: usize) -> Result<bool, CoreError> {
47        Ok(Self::trace_size(trace)? > threshold_bytes)
48    }
49
50    pub async fn create_upload_url(
51        &self,
52        content_type: Option<&str>,
53        expires_in_seconds: Option<i64>,
54    ) -> Result<UploadUrlResponse, CoreError> {
55        let mut payload = serde_json::Map::new();
56        payload.insert(
57            "content_type".to_string(),
58            Value::String(content_type.unwrap_or("application/json").to_string()),
59        );
60        if let Some(expires) = expires_in_seconds {
61            payload.insert(
62                "expires_in_seconds".to_string(),
63                Value::Number(expires.into()),
64            );
65        }
66
67        let path = "/v1/traces/upload-url";
68        let response: UploadUrlResponse = self
69            .http
70            .post_json(path, &Value::Object(payload))
71            .await
72            .map_err(map_http_error)?;
73
74        Ok(response)
75    }
76
77    pub async fn upload_trace(
78        &self,
79        trace: &Value,
80        content_type: Option<&str>,
81        expires_in_seconds: Option<i64>,
82    ) -> Result<String, CoreError> {
83        let content_type = content_type.unwrap_or("application/json");
84        let payload = serde_json::to_vec(trace)
85            .map_err(|e| CoreError::Validation(format!("failed to serialize trace: {}", e)))?;
86
87        let upload = self
88            .create_upload_url(Some(content_type), expires_in_seconds)
89            .await?;
90
91        if let Some(max_size) = upload.max_size_bytes {
92            if payload.len() as i64 > max_size {
93                return Err(CoreError::Validation(format!(
94                    "trace exceeds max size {} bytes",
95                    max_size
96                )));
97            }
98        }
99
100        let client = reqwest::Client::builder()
101            .timeout(Duration::from_secs(self.timeout_secs))
102            .build()
103            .map_err(|e| CoreError::Internal(format!("failed to build upload client: {}", e)))?;
104
105        let resp = client
106            .put(&upload.upload_url)
107            .header("Content-Type", content_type)
108            .body(payload)
109            .send()
110            .await
111            .map_err(|e| CoreError::Http(e))?;
112
113        if !resp.status().is_success() {
114            let status = resp.status().as_u16();
115            let body = resp.text().await.unwrap_or_default();
116            return Err(CoreError::HttpResponse(crate::HttpErrorInfo {
117                status,
118                url: upload.upload_url,
119                message: "trace_upload_failed".to_string(),
120                body_snippet: Some(body.chars().take(200).collect()),
121            }));
122        }
123
124        Ok(upload.trace_ref)
125    }
126
127    pub fn base_url(&self) -> &str {
128        &self.base_url
129    }
130}
131
132fn map_http_error(err: HttpError) -> CoreError {
133    match err {
134        HttpError::Response(detail) => CoreError::HttpResponse(crate::HttpErrorInfo {
135            status: detail.status,
136            url: detail.url,
137            message: detail.message,
138            body_snippet: detail.body_snippet,
139        }),
140        HttpError::Request(e) => CoreError::Http(e),
141        _ => CoreError::Internal(format!("{}", err)),
142    }
143}
synth_ai_core/trace_upload.rs

synth_ai_core/
trace_upload.rs
