byokey_proxy/handler/
usage.rs1use crate::AppState;
4use axum::{
5 Json,
6 extract::{Query, State},
7};
8use serde::{Deserialize, Serialize};
9use std::sync::Arc;
10use utoipa::{IntoParams, ToSchema};
11
12#[derive(Deserialize, ToSchema, IntoParams)]
14pub struct UsageHistoryQuery {
15 pub from: Option<i64>,
17 pub to: Option<i64>,
19 pub model: Option<String>,
21}
22
23#[derive(Serialize, ToSchema)]
25pub struct UsageHistoryResponse {
26 pub from: i64,
27 pub to: i64,
28 pub bucket_seconds: i64,
29 pub buckets: Vec<byokey_types::UsageBucket>,
30}
31
32#[utoipa::path(
34 get,
35 path = "/v0/management/usage",
36 responses((status = 200, body = crate::usage::UsageSnapshot)),
37 tag = "management"
38)]
39pub async fn usage_handler(
40 State(state): State<Arc<AppState>>,
41) -> Json<crate::usage::UsageSnapshot> {
42 Json(state.usage.snapshot())
43}
44
45#[utoipa::path(
47 get,
48 path = "/v0/management/usage/history",
49 params(UsageHistoryQuery),
50 responses(
51 (status = 200, body = UsageHistoryResponse),
52 ),
53 tag = "management"
54)]
55pub async fn usage_history_handler(
56 State(state): State<Arc<AppState>>,
57 Query(q): Query<UsageHistoryQuery>,
58) -> Json<serde_json::Value> {
59 let Some(store) = state.usage.store() else {
60 return Json(serde_json::json!({ "error": "no persistent usage store configured" }));
61 };
62
63 #[allow(clippy::cast_possible_wrap)]
64 let now = std::time::SystemTime::now()
65 .duration_since(std::time::UNIX_EPOCH)
66 .unwrap_or_default()
67 .as_secs() as i64;
68
69 let to = q.to.unwrap_or(now);
70 let from = q.from.unwrap_or(to - 86400);
71
72 let range = to - from;
74 let bucket_secs = if range <= 86400 {
75 3600 } else if range <= 86400 * 7 {
77 21600 } else {
79 86400 };
81
82 match store.query(from, to, q.model.as_deref(), bucket_secs).await {
83 Ok(buckets) => Json(serde_json::json!({
84 "from": from,
85 "to": to,
86 "bucket_seconds": bucket_secs,
87 "buckets": buckets,
88 })),
89 Err(e) => Json(serde_json::json!({ "error": e.to_string() })),
90 }
91}