Skip to main content

byokey_proxy/handler/
usage.rs

1//! Usage statistics endpoints — current counters and historical data.
2
3use 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/// Query parameters for the usage history endpoint.
13#[derive(Deserialize, ToSchema, IntoParams)]
14pub struct UsageHistoryQuery {
15    /// Start of the time range (unix timestamp). Defaults to 24 hours ago.
16    pub from: Option<i64>,
17    /// End of the time range (unix timestamp). Defaults to now.
18    pub to: Option<i64>,
19    /// Optional model name filter.
20    pub model: Option<String>,
21}
22
23/// Response for the usage history endpoint.
24#[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/// Returns current in-memory usage counters.
33#[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/// Returns bucketed usage history from the persistent store.
46#[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    // Auto-select bucket size based on range.
73    let range = to - from;
74    let bucket_secs = if range <= 86400 {
75        3600 // hourly
76    } else if range <= 86400 * 7 {
77        21600 // 6-hour
78    } else {
79        86400 // daily
80    };
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}