Skip to main content

xtrace_client/
lib.rs

1#[cfg(feature = "tracing")]
2pub mod layer;
3#[cfg(feature = "tracing")]
4pub use layer::current_trace_id;
5#[cfg(feature = "tracing")]
6pub use layer::XtraceLayer;
7
8use chrono::{DateTime, Utc};
9use reqwest::header::{HeaderMap, HeaderValue, AUTHORIZATION, CONTENT_TYPE};
10use serde::{Deserialize, Serialize};
11use serde_json::Value as JsonValue;
12use std::collections::HashMap;
13use std::time::Duration;
14use url::Url;
15use uuid::Uuid;
16
17#[derive(Debug, thiserror::Error)]
18pub enum Error {
19    #[error("invalid base url: {0}")]
20    InvalidBaseUrl(#[from] url::ParseError),
21
22    #[error("http error: {0}")]
23    Http(#[from] reqwest::Error),
24}
25
26#[derive(Clone)]
27pub struct Client {
28    base_url: Url,
29    http: reqwest::Client,
30}
31
32impl Client {
33    pub fn new(base_url: &str, bearer_token: &str) -> Result<Self, Error> {
34        let base_url = Url::parse(base_url)?;
35
36        let mut headers = HeaderMap::new();
37        headers.insert(
38            AUTHORIZATION,
39            HeaderValue::from_str(&format!("Bearer {}", bearer_token)).unwrap(),
40        );
41        headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
42
43        let http = reqwest::Client::builder()
44            .default_headers(headers)
45            .timeout(Duration::from_secs(30))
46            .build()?;
47
48        Ok(Self { base_url, http })
49    }
50
51    pub async fn healthz(&self) -> Result<(), Error> {
52        let url = self.base_url.join("healthz")?;
53        self.http.get(url).send().await?.error_for_status()?;
54        Ok(())
55    }
56
57    pub async fn ingest_batch(
58        &self,
59        req: &BatchIngestRequest,
60    ) -> Result<ApiResponse<JsonValue>, Error> {
61        let url = self.base_url.join("v1/l/batch")?;
62        let res = self
63            .http
64            .post(url)
65            .json(req)
66            .send()
67            .await?
68            .error_for_status()?;
69        Ok(res.json::<ApiResponse<JsonValue>>().await?)
70    }
71
72    pub async fn list_traces(&self, q: &TraceListQuery) -> Result<PagedData<TraceListItem>, Error> {
73        let mut url = self.base_url.join("api/public/traces")?;
74        {
75            let mut pairs = url.query_pairs_mut();
76            if let Some(v) = q.page {
77                pairs.append_pair("page", &v.to_string());
78            }
79            if let Some(v) = q.limit {
80                pairs.append_pair("limit", &v.to_string());
81            }
82            if let Some(v) = q.user_id.as_deref() {
83                pairs.append_pair("userId", v);
84            }
85            if let Some(v) = q.name.as_deref() {
86                pairs.append_pair("name", v);
87            }
88            if let Some(v) = q.session_id.as_deref() {
89                pairs.append_pair("sessionId", v);
90            }
91            if let Some(v) = q.from_timestamp.as_ref() {
92                pairs.append_pair("fromTimestamp", &v.to_rfc3339());
93            }
94            if let Some(v) = q.to_timestamp.as_ref() {
95                pairs.append_pair("toTimestamp", &v.to_rfc3339());
96            }
97            if let Some(v) = q.order_by.as_deref() {
98                pairs.append_pair("orderBy", v);
99            }
100            for tag in &q.tags {
101                pairs.append_pair("tags", tag);
102            }
103            if let Some(v) = q.version.as_deref() {
104                pairs.append_pair("version", v);
105            }
106            if let Some(v) = q.release.as_deref() {
107                pairs.append_pair("release", v);
108            }
109            for env in &q.environment {
110                pairs.append_pair("environment", env);
111            }
112            if let Some(v) = q.fields.as_deref() {
113                pairs.append_pair("fields", v);
114            }
115        }
116
117        let res = self.http.get(url).send().await?.error_for_status()?;
118        Ok(res.json::<PagedData<TraceListItem>>().await?)
119    }
120
121    pub async fn get_trace(&self, trace_id: Uuid) -> Result<TraceDetailDto, Error> {
122        let url = self
123            .base_url
124            .join(&format!("api/public/traces/{}", trace_id))?;
125        let res = self.http.get(url).send().await?.error_for_status()?;
126        Ok(res.json::<TraceDetailDto>().await?)
127    }
128
129    pub async fn metrics_daily(
130        &self,
131        q: &MetricsDailyQuery,
132    ) -> Result<PagedData<MetricsDailyItem>, Error> {
133        let mut url = self.base_url.join("api/public/metrics/daily")?;
134        {
135            let mut pairs = url.query_pairs_mut();
136            if let Some(v) = q.page {
137                pairs.append_pair("page", &v.to_string());
138            }
139            if let Some(v) = q.limit {
140                pairs.append_pair("limit", &v.to_string());
141            }
142            if let Some(v) = q.trace_name.as_deref() {
143                pairs.append_pair("traceName", v);
144            }
145            if let Some(v) = q.user_id.as_deref() {
146                pairs.append_pair("userId", v);
147            }
148            for tag in &q.tags {
149                pairs.append_pair("tags", tag);
150            }
151            if let Some(v) = q.from_timestamp.as_ref() {
152                pairs.append_pair("fromTimestamp", &v.to_rfc3339());
153            }
154            if let Some(v) = q.to_timestamp.as_ref() {
155                pairs.append_pair("toTimestamp", &v.to_rfc3339());
156            }
157            if let Some(v) = q.version.as_deref() {
158                pairs.append_pair("version", v);
159            }
160            if let Some(v) = q.release.as_deref() {
161                pairs.append_pair("release", v);
162            }
163        }
164
165        let res = self.http.get(url).send().await?.error_for_status()?;
166        Ok(res.json::<PagedData<MetricsDailyItem>>().await?)
167    }
168
169    pub async fn push_metrics(
170        &self,
171        metrics: &[MetricPoint],
172    ) -> Result<ApiResponse<JsonValue>, Error> {
173        let url = self.base_url.join("v1/metrics/batch")?;
174        let req = MetricsBatchRequest {
175            metrics: metrics.to_vec(),
176        };
177        let res = self
178            .http
179            .post(url)
180            .json(&req)
181            .send()
182            .await?
183            .error_for_status()?;
184        Ok(res.json::<ApiResponse<JsonValue>>().await?)
185    }
186
187    /// Query time-series metrics with optional downsampling and aggregation.
188    pub async fn query_metrics(
189        &self,
190        q: &MetricsQueryParams,
191    ) -> Result<MetricsQueryResponse, Error> {
192        let mut url = self.base_url.join("api/public/metrics/query")?;
193        {
194            let mut pairs = url.query_pairs_mut();
195            pairs.append_pair("name", &q.name);
196            if let Some(v) = q.from.as_ref() {
197                pairs.append_pair("from", &v.to_rfc3339());
198            }
199            if let Some(v) = q.to.as_ref() {
200                pairs.append_pair("to", &v.to_rfc3339());
201            }
202            if let Some(v) = q.labels.as_ref() {
203                pairs.append_pair("labels", &serde_json::to_string(v).unwrap_or_default());
204            }
205            if let Some(v) = q.step.as_deref() {
206                pairs.append_pair("step", v);
207            }
208            if let Some(v) = q.agg.as_deref() {
209                pairs.append_pair("agg", v);
210            }
211            if let Some(v) = q.group_by.as_deref() {
212                pairs.append_pair("group_by", v);
213            }
214        }
215
216        let res = self.http.get(url).send().await?.error_for_status()?;
217        Ok(res.json::<MetricsQueryResponse>().await?)
218    }
219
220    /// List all available metric names.
221    pub async fn list_metric_names(&self) -> Result<Vec<String>, Error> {
222        let url = self.base_url.join("api/public/metrics/names")?;
223        let res = self.http.get(url).send().await?.error_for_status()?;
224        let wrapper = res.json::<MetricNamesResponse>().await?;
225        Ok(wrapper.data)
226    }
227}
228
229#[derive(Debug, Serialize, Deserialize, Clone)]
230pub struct MetricPoint {
231    pub name: String,
232    #[serde(default)]
233    pub labels: HashMap<String, String>,
234    pub value: f64,
235    pub timestamp: DateTime<Utc>,
236}
237
238#[derive(Debug, Serialize)]
239struct MetricsBatchRequest {
240    metrics: Vec<MetricPoint>,
241}
242
243/// Parameters for `query_metrics`.
244#[derive(Debug, Default, Serialize, Deserialize)]
245pub struct MetricsQueryParams {
246    /// Required. Metric name (e.g. `pending_requests`, `kv_cache_usage`).
247    pub name: String,
248    /// Start time (inclusive). Defaults to 1 hour before `to`.
249    #[serde(default)]
250    pub from: Option<DateTime<Utc>>,
251    /// End time (inclusive). Defaults to now.
252    #[serde(default)]
253    pub to: Option<DateTime<Utc>>,
254    /// Label filter as a JSON object (JSONB containment).
255    #[serde(default)]
256    pub labels: Option<HashMap<String, String>>,
257    /// Downsample step: `1m`, `5m`, `1h`, `1d`. Default `1m`.
258    #[serde(default)]
259    pub step: Option<String>,
260    /// Aggregation: `avg`, `max`, `min`, `sum`, `last`, `p50`, `p90`, `p99`. Default `avg`.
261    #[serde(default)]
262    pub agg: Option<String>,
263    /// Group results by a specific label key instead of the full label set.
264    #[serde(default)]
265    pub group_by: Option<String>,
266}
267
268/// A single time-series data point.
269#[derive(Debug, Deserialize, Clone)]
270pub struct MetricValuePoint {
271    pub timestamp: String,
272    pub value: f64,
273}
274
275/// A time-series for one unique label combination.
276#[derive(Debug, Deserialize, Clone)]
277pub struct MetricsSeries {
278    pub labels: JsonValue,
279    pub values: Vec<MetricValuePoint>,
280}
281
282/// Metadata about the metrics query result.
283#[derive(Debug, Deserialize, Clone)]
284pub struct MetricsQueryMeta {
285    /// Timestamp (RFC3339 UTC) of the most recent data point. Absent when no data.
286    #[serde(default)]
287    pub latest_ts: Option<String>,
288    /// Number of distinct series returned.
289    pub series_count: usize,
290    /// `true` when results were truncated due to server limits.
291    pub truncated: bool,
292}
293
294/// Response from `GET /api/public/metrics/query`.
295#[derive(Debug, Deserialize, Clone)]
296pub struct MetricsQueryResponse {
297    pub data: Vec<MetricsSeries>,
298    pub meta: MetricsQueryMeta,
299}
300
301#[derive(Debug, Deserialize)]
302struct MetricNamesResponse {
303    data: Vec<String>,
304}
305
306#[derive(Debug, Deserialize)]
307pub struct ApiResponse<T> {
308    pub message: String,
309    /// Machine-readable error code. Present only on error responses.
310    /// Possible values: `UNAUTHORIZED`, `BAD_REQUEST`, `TOO_MANY_REQUESTS`,
311    /// `INTERNAL_ERROR`, `SERVICE_UNAVAILABLE`, `NOT_FOUND`.
312    #[serde(default)]
313    pub code: Option<String>,
314    #[serde(default)]
315    pub data: Option<T>,
316}
317
318#[derive(Debug, Deserialize)]
319#[serde(rename_all = "camelCase")]
320pub struct PageMeta {
321    pub page: i64,
322    pub limit: i64,
323    pub total_items: i64,
324    pub total_pages: i64,
325}
326
327#[derive(Debug, Deserialize)]
328pub struct PagedData<T> {
329    pub data: Vec<T>,
330    pub meta: PageMeta,
331}
332
333#[derive(Debug, Serialize, Deserialize, Default)]
334pub struct BatchIngestRequest {
335    #[serde(default)]
336    pub trace: Option<TraceIngest>,
337    #[serde(default)]
338    pub observations: Vec<ObservationIngest>,
339}
340
341#[derive(Debug, Serialize, Deserialize)]
342#[serde(rename_all = "camelCase")]
343pub struct TraceIngest {
344    pub id: Uuid,
345    #[serde(default)]
346    pub timestamp: Option<DateTime<Utc>>,
347
348    #[serde(default)]
349    pub name: Option<String>,
350    #[serde(default)]
351    pub input: Option<JsonValue>,
352    #[serde(default)]
353    pub output: Option<JsonValue>,
354    #[serde(default)]
355    pub session_id: Option<String>,
356    #[serde(default)]
357    pub release: Option<String>,
358    #[serde(default)]
359    pub version: Option<String>,
360    #[serde(default, rename = "userId")]
361    pub user_id: Option<String>,
362    #[serde(default)]
363    pub metadata: Option<JsonValue>,
364    #[serde(default)]
365    pub tags: Vec<String>,
366    #[serde(default)]
367    pub public: Option<bool>,
368    #[serde(default)]
369    pub environment: Option<String>,
370    #[serde(default)]
371    pub external_id: Option<String>,
372    #[serde(default)]
373    pub bookmarked: Option<bool>,
374
375    #[serde(default)]
376    pub latency: Option<f64>,
377    #[serde(default, rename = "totalCost")]
378    pub total_cost: Option<f64>,
379
380    #[serde(default, rename = "projectId")]
381    pub project_id: Option<String>,
382}
383
384#[derive(Debug, Serialize, Deserialize)]
385#[serde(rename_all = "camelCase")]
386pub struct ObservationIngest {
387    pub id: Uuid,
388    #[serde(rename = "traceId")]
389    pub trace_id: Uuid,
390
391    #[serde(default, rename = "type")]
392    pub r#type: Option<String>,
393    #[serde(default)]
394    pub name: Option<String>,
395
396    #[serde(default)]
397    pub start_time: Option<DateTime<Utc>>,
398    #[serde(default)]
399    pub end_time: Option<DateTime<Utc>>,
400    #[serde(default)]
401    pub completion_start_time: Option<DateTime<Utc>>,
402
403    #[serde(default)]
404    pub model: Option<String>,
405    #[serde(default)]
406    pub model_parameters: Option<JsonValue>,
407
408    #[serde(default)]
409    pub input: Option<JsonValue>,
410    #[serde(default)]
411    pub output: Option<JsonValue>,
412
413    #[serde(default)]
414    pub usage: Option<JsonValue>,
415
416    #[serde(default)]
417    pub level: Option<String>,
418    #[serde(default)]
419    pub status_message: Option<String>,
420    #[serde(default)]
421    pub parent_observation_id: Option<Uuid>,
422
423    #[serde(default)]
424    pub prompt_id: Option<String>,
425    #[serde(default)]
426    pub prompt_name: Option<String>,
427    #[serde(default)]
428    pub prompt_version: Option<String>,
429
430    #[serde(default)]
431    pub model_id: Option<String>,
432
433    #[serde(default)]
434    pub input_price: Option<f64>,
435    #[serde(default)]
436    pub output_price: Option<f64>,
437    #[serde(default)]
438    pub total_price: Option<f64>,
439
440    #[serde(default)]
441    pub calculated_input_cost: Option<f64>,
442    #[serde(default)]
443    pub calculated_output_cost: Option<f64>,
444    #[serde(default)]
445    pub calculated_total_cost: Option<f64>,
446
447    #[serde(default)]
448    pub latency: Option<f64>,
449    #[serde(default)]
450    pub time_to_first_token: Option<f64>,
451
452    #[serde(default)]
453    pub completion_tokens: Option<i64>,
454    #[serde(default)]
455    pub prompt_tokens: Option<i64>,
456    #[serde(default)]
457    pub total_tokens: Option<i64>,
458    #[serde(default)]
459    pub unit: Option<String>,
460
461    #[serde(default)]
462    pub metadata: Option<JsonValue>,
463
464    #[serde(default)]
465    pub environment: Option<String>,
466
467    #[serde(default, rename = "projectId")]
468    pub project_id: Option<String>,
469}
470
471#[derive(Debug, Default, Serialize, Deserialize)]
472#[serde(rename_all = "camelCase")]
473pub struct TraceListQuery {
474    #[serde(default)]
475    pub page: Option<i64>,
476    #[serde(default)]
477    pub limit: Option<i64>,
478
479    #[serde(default, rename = "userId")]
480    pub user_id: Option<String>,
481    #[serde(default)]
482    pub name: Option<String>,
483    #[serde(default, rename = "sessionId")]
484    pub session_id: Option<String>,
485
486    #[serde(default, rename = "fromTimestamp")]
487    pub from_timestamp: Option<DateTime<Utc>>,
488    #[serde(default, rename = "toTimestamp")]
489    pub to_timestamp: Option<DateTime<Utc>>,
490
491    #[serde(default, rename = "orderBy")]
492    pub order_by: Option<String>,
493
494    #[serde(default)]
495    pub tags: Vec<String>,
496
497    #[serde(default)]
498    pub version: Option<String>,
499    #[serde(default)]
500    pub release: Option<String>,
501    #[serde(default)]
502    pub environment: Vec<String>,
503
504    #[serde(default)]
505    pub fields: Option<String>,
506}
507
508#[derive(Debug, Deserialize)]
509#[serde(rename_all = "camelCase")]
510pub struct TraceListItem {
511    pub id: Uuid,
512    pub timestamp: DateTime<Utc>,
513    pub name: Option<String>,
514    #[serde(default)]
515    pub input: Option<JsonValue>,
516    #[serde(default)]
517    pub output: Option<JsonValue>,
518    pub session_id: Option<String>,
519    pub release: Option<String>,
520    pub version: Option<String>,
521    pub user_id: Option<String>,
522    #[serde(default)]
523    pub metadata: Option<JsonValue>,
524    pub tags: Vec<String>,
525    pub public: bool,
526    pub environment: String,
527    pub html_path: String,
528    pub latency: Option<f64>,
529    pub total_cost: Option<f64>,
530    pub observations: Vec<String>,
531    pub scores: Vec<String>,
532}
533
534#[derive(Debug, Default, Serialize, Deserialize)]
535#[serde(rename_all = "camelCase")]
536pub struct MetricsDailyQuery {
537    #[serde(default)]
538    pub page: Option<i64>,
539    #[serde(default)]
540    pub limit: Option<i64>,
541
542    #[serde(default, rename = "traceName")]
543    pub trace_name: Option<String>,
544    #[serde(default, rename = "userId")]
545    pub user_id: Option<String>,
546    #[serde(default)]
547    pub tags: Vec<String>,
548
549    #[serde(default, rename = "fromTimestamp")]
550    pub from_timestamp: Option<DateTime<Utc>>,
551    #[serde(default, rename = "toTimestamp")]
552    pub to_timestamp: Option<DateTime<Utc>>,
553
554    #[serde(default)]
555    pub version: Option<String>,
556    #[serde(default)]
557    pub release: Option<String>,
558}
559
560#[derive(Debug, Deserialize)]
561#[serde(rename_all = "camelCase")]
562pub struct MetricsDailyItem {
563    pub date: String,
564    pub count_traces: i64,
565    pub count_observations: i64,
566    pub total_cost: f64,
567    pub usage: JsonValue,
568}
569
570#[derive(Debug, Deserialize)]
571#[serde(rename_all = "camelCase")]
572pub struct TraceDetailDto {
573    pub id: Uuid,
574    pub timestamp: DateTime<Utc>,
575    pub name: Option<String>,
576    pub input: JsonValue,
577    pub output: JsonValue,
578    pub session_id: Option<String>,
579    pub release: Option<String>,
580    pub version: Option<String>,
581    pub user_id: Option<String>,
582    pub metadata: JsonValue,
583    pub tags: Vec<String>,
584    pub public: bool,
585    pub environment: String,
586    pub html_path: String,
587    pub latency: Option<f64>,
588    pub total_cost: Option<f64>,
589    pub observations: Vec<JsonValue>,
590    pub scores: Vec<JsonValue>,
591}