kcode_gemini_api/
accounting.rs

1use std::{collections::BTreeMap, fmt, sync::Mutex};
2
3use chrono::{DateTime, Datelike, TimeZone, Timelike, Utc};
4
5use crate::{CostBreakdown, Error, Money, Result, ServiceTier, TokenUsage};
6
7type WindowBounds = (Option<DateTime<Utc>>, Option<DateTime<Utc>>);
8
9/// UTC period used by one spending limit.
10#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
11pub enum LimitPeriod {
12    /// Current UTC clock hour.
13    Hourly,
14    /// Current UTC calendar day.
15    Daily,
16    /// Current UTC calendar month.
17    Monthly,
18}
19
20impl LimitPeriod {
21    /// Returns the stable name for the period.
22    pub const fn as_str(self) -> &'static str {
23        match self {
24            Self::Hourly => "hourly",
25            Self::Daily => "daily",
26            Self::Monthly => "monthly",
27        }
28    }
29}
30
31impl fmt::Display for LimitPeriod {
32    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
33        formatter.write_str(self.as_str())
34    }
35}
36
37/// Optional spending limits for the current session's UTC hour, day, and month.
38#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
39pub struct SpendingLimits {
40    /// Current UTC clock-hour limit, or `None` when disabled.
41    pub hourly: Option<Money>,
42    /// Current UTC calendar-day limit, or `None` when disabled.
43    pub daily: Option<Money>,
44    /// Current UTC calendar-month limit, or `None` when disabled.
45    pub monthly: Option<Money>,
46}
47
48impl SpendingLimits {
49    pub(crate) const fn any(self) -> bool {
50        self.hourly.is_some() || self.daily.is_some() || self.monthly.is_some()
51    }
52}
53
54/// Current-session spend and configured limits for all three UTC periods.
55#[derive(Clone, Copy, Debug, Eq, PartialEq)]
56pub struct LimitStatus {
57    /// Active limit configuration.
58    pub limits: SpendingLimits,
59    /// Session spend since the start of the current UTC hour.
60    pub hourly_spent: Money,
61    /// Session spend since the start of the current UTC day.
62    pub daily_spent: Money,
63    /// Session spend since the start of the current UTC month.
64    pub monthly_spent: Money,
65}
66
67/// Time selection for a current-session usage breakdown.
68#[derive(Clone, Debug, Eq, PartialEq)]
69pub enum UsageWindow {
70    /// Every request recorded by the live client.
71    AllTime,
72    /// Current UTC clock hour.
73    CurrentHour,
74    /// Current UTC calendar day.
75    CurrentDay,
76    /// Current UTC calendar month.
77    CurrentMonth,
78    /// Explicit half-open UTC range `[start, end)`.
79    Range {
80        /// Inclusive range start.
81        start: DateTime<Utc>,
82        /// Exclusive range end.
83        end: DateTime<Utc>,
84    },
85}
86
87/// One in-memory usage record. Prompts, media, responses, and credentials are not recorded.
88#[derive(Clone, Debug, Eq, PartialEq)]
89pub struct UsageRecord {
90    /// Monotonic identifier within the current client session.
91    pub id: u64,
92    /// Time at which the provider attempt finished.
93    pub occurred_at: DateTime<Utc>,
94    /// Stable operation name such as `infer_flash_lite` or `nano_banana_pro`.
95    pub operation: String,
96    /// Exact requested Gemini model identifier.
97    pub model: String,
98    /// Requested service tier.
99    pub service_tier: ServiceTier,
100    /// Whether the provider interaction produced a usable response.
101    pub succeeded: bool,
102    /// Gemini interaction identifier when the provider supplied one.
103    pub provider_request_id: Option<String>,
104    /// Local failure class without raw provider or request content.
105    pub failure_kind: Option<String>,
106    /// Provider-reported token and grounding usage.
107    pub usage: TokenUsage,
108    /// Locally calculated cost at the crate's compiled pricing version.
109    pub cost: CostBreakdown,
110}
111
112/// Aggregated request, token, and cost values.
113#[derive(Clone, Debug, Default, Eq, PartialEq)]
114pub struct UsageTotals {
115    /// Total provider attempts.
116    pub requests: u64,
117    /// Attempts that returned usable model output.
118    pub successful_requests: u64,
119    /// Attempts that failed locally or at the provider after admission.
120    pub failed_requests: u64,
121    /// Summed provider token and grounding usage.
122    pub usage: TokenUsage,
123    /// Summed locally calculated costs.
124    pub cost: CostBreakdown,
125}
126
127impl UsageTotals {
128    fn add_record(&mut self, record: &UsageRecord) {
129        self.requests = self.requests.saturating_add(1);
130        if record.succeeded {
131            self.successful_requests = self.successful_requests.saturating_add(1);
132        } else {
133            self.failed_requests = self.failed_requests.saturating_add(1);
134        }
135        self.usage.saturating_add(&record.usage);
136        self.cost.saturating_add(&record.cost);
137    }
138}
139
140/// Usage totals grouped by one model or operation name.
141#[derive(Clone, Debug, Eq, PartialEq)]
142pub struct GroupedUsage {
143    /// Exact grouping key.
144    pub key: String,
145    /// Aggregated usage for the key.
146    pub totals: UsageTotals,
147}
148
149/// Complete current-session aggregate report for one selected window.
150#[derive(Clone, Debug, Eq, PartialEq)]
151pub struct UsageBreakdown {
152    /// Requested time window.
153    pub window: UsageWindow,
154    /// Report generation time.
155    pub generated_at: DateTime<Utc>,
156    /// Aggregate across every matching record.
157    pub totals: UsageTotals,
158    /// Matching usage grouped by exact model identifier.
159    pub by_model: Vec<GroupedUsage>,
160    /// Matching usage grouped by stable operation name.
161    pub by_operation: Vec<GroupedUsage>,
162    /// Current UTC hour/day/month limit state, independent of the report window.
163    pub limit_status: LimitStatus,
164    /// Clarifies the scope and authority of the cost calculation.
165    pub accounting_note: String,
166}
167
168#[derive(Default)]
169struct SessionState {
170    next_id: u64,
171    records: Vec<UsageRecord>,
172    limits: SpendingLimits,
173}
174
175pub(crate) struct Accounting {
176    state: Mutex<SessionState>,
177}
178
179impl Accounting {
180    pub(crate) fn new() -> Self {
181        Self {
182            state: Mutex::new(SessionState {
183                next_id: 1,
184                ..SessionState::default()
185            }),
186        }
187    }
188
189    pub(crate) fn limits(&self) -> Result<SpendingLimits> {
190        Ok(self.lock()?.limits)
191    }
192
193    pub(crate) fn set_limits(&self, limits: SpendingLimits) -> Result<()> {
194        self.lock()?.limits = limits;
195        Ok(())
196    }
197
198    pub(crate) fn enforce_limits(&self, now: DateTime<Utc>) -> Result<()> {
199        let status = self.limit_status_at(now)?;
200        for (period, limit, spent) in [
201            (
202                LimitPeriod::Hourly,
203                status.limits.hourly,
204                status.hourly_spent,
205            ),
206            (LimitPeriod::Daily, status.limits.daily, status.daily_spent),
207            (
208                LimitPeriod::Monthly,
209                status.limits.monthly,
210                status.monthly_spent,
211            ),
212        ] {
213            if let Some(limit) = limit
214                && spent >= limit
215            {
216                return Err(Error::SpendingLimitReached {
217                    period,
218                    limit,
219                    spent,
220                });
221            }
222        }
223        Ok(())
224    }
225
226    pub(crate) fn limit_status(&self) -> Result<LimitStatus> {
227        self.limit_status_at(Utc::now())
228    }
229
230    fn limit_status_at(&self, now: DateTime<Utc>) -> Result<LimitStatus> {
231        let state = self.lock()?;
232        let (hour, day, month) = period_starts(now);
233        Ok(LimitStatus {
234            limits: state.limits,
235            hourly_spent: sum_cost_since(&state.records, hour),
236            daily_spent: sum_cost_since(&state.records, day),
237            monthly_spent: sum_cost_since(&state.records, month),
238        })
239    }
240
241    #[allow(clippy::too_many_arguments)]
242    pub(crate) fn record(
243        &self,
244        operation: &str,
245        model: &str,
246        service_tier: ServiceTier,
247        succeeded: bool,
248        provider_request_id: Option<&str>,
249        failure_kind: Option<&str>,
250        usage: &TokenUsage,
251        cost: &CostBreakdown,
252    ) -> Result<u64> {
253        let mut state = self.lock()?;
254        let id = state.next_id;
255        state.next_id = state.next_id.checked_add(1).ok_or_else(|| {
256            Error::Accounting("session usage identifier space is exhausted".into())
257        })?;
258        state.records.push(UsageRecord {
259            id,
260            occurred_at: Utc::now(),
261            operation: operation.into(),
262            model: model.into(),
263            service_tier,
264            succeeded,
265            provider_request_id: provider_request_id.map(str::to_owned),
266            failure_kind: failure_kind.map(str::to_owned),
267            usage: usage.clone(),
268            cost: cost.clone(),
269        });
270        Ok(id)
271    }
272
273    pub(crate) fn breakdown(&self, window: UsageWindow) -> Result<UsageBreakdown> {
274        let generated_at = Utc::now();
275        let records = self.records(&window, None)?;
276        let mut totals = UsageTotals::default();
277        let mut models = BTreeMap::<String, UsageTotals>::new();
278        let mut operations = BTreeMap::<String, UsageTotals>::new();
279        for record in &records {
280            totals.add_record(record);
281            models
282                .entry(record.model.clone())
283                .or_default()
284                .add_record(record);
285            operations
286                .entry(record.operation.clone())
287                .or_default()
288                .add_record(record);
289        }
290        Ok(UsageBreakdown {
291            window,
292            generated_at,
293            totals,
294            by_model: grouped(models),
295            by_operation: grouped(operations),
296            limit_status: self.limit_status_at(generated_at)?,
297            accounting_note: concat!(
298                "This report covers only requests made through this live Gemini client session. ",
299                "Costs apply compiled Gemini Developer API rates to provider-reported usage and ",
300                "are local estimates; Google AI Studio and Cloud Billing remain authoritative. ",
301                "All records and limits disappear when the session is dropped."
302            )
303            .into(),
304        })
305    }
306
307    pub(crate) fn usage_records(
308        &self,
309        window: UsageWindow,
310        maximum: usize,
311    ) -> Result<Vec<UsageRecord>> {
312        if maximum == 0 || maximum > 10_000 {
313            return Err(Error::InvalidInput(
314                "maximum usage records must be between 1 and 10000".into(),
315            ));
316        }
317        self.records(&window, Some(maximum))
318    }
319
320    fn records(&self, window: &UsageWindow, maximum: Option<usize>) -> Result<Vec<UsageRecord>> {
321        let (start, end) = window_bounds(window, Utc::now())?;
322        let state = self.lock()?;
323        let mut records = state
324            .records
325            .iter()
326            .rev()
327            .filter(|record| {
328                start.is_none_or(|start| record.occurred_at >= start)
329                    && end.is_none_or(|end| record.occurred_at < end)
330            })
331            .take(maximum.unwrap_or(usize::MAX))
332            .cloned()
333            .collect::<Vec<_>>();
334        records.sort_by(|left, right| {
335            right
336                .occurred_at
337                .cmp(&left.occurred_at)
338                .then_with(|| right.id.cmp(&left.id))
339        });
340        Ok(records)
341    }
342
343    fn lock(&self) -> Result<std::sync::MutexGuard<'_, SessionState>> {
344        self.state
345            .lock()
346            .map_err(|_| Error::Accounting("in-memory session accounting is unavailable".into()))
347    }
348}
349
350fn grouped(values: BTreeMap<String, UsageTotals>) -> Vec<GroupedUsage> {
351    values
352        .into_iter()
353        .map(|(key, totals)| GroupedUsage { key, totals })
354        .collect()
355}
356
357fn sum_cost_since(records: &[UsageRecord], start: DateTime<Utc>) -> Money {
358    records
359        .iter()
360        .filter(|record| record.occurred_at >= start)
361        .fold(Money::ZERO, |total, record| {
362            total.saturating_add(record.cost.total)
363        })
364}
365
366fn period_starts(now: DateTime<Utc>) -> (DateTime<Utc>, DateTime<Utc>, DateTime<Utc>) {
367    let hour = now
368        .with_minute(0)
369        .and_then(|value| value.with_second(0))
370        .and_then(|value| value.with_nanosecond(0))
371        .expect("valid UTC clock-hour boundary");
372    let day = Utc
373        .with_ymd_and_hms(now.year(), now.month(), now.day(), 0, 0, 0)
374        .single()
375        .expect("valid UTC calendar-day boundary");
376    let month = Utc
377        .with_ymd_and_hms(now.year(), now.month(), 1, 0, 0, 0)
378        .single()
379        .expect("valid UTC calendar-month boundary");
380    (hour, day, month)
381}
382
383fn window_bounds(window: &UsageWindow, now: DateTime<Utc>) -> Result<WindowBounds> {
384    let (hour, day, month) = period_starts(now);
385    match window {
386        UsageWindow::AllTime => Ok((None, None)),
387        UsageWindow::CurrentHour => Ok((Some(hour), None)),
388        UsageWindow::CurrentDay => Ok((Some(day), None)),
389        UsageWindow::CurrentMonth => Ok((Some(month), None)),
390        UsageWindow::Range { start, end } if start < end => Ok((Some(*start), Some(*end))),
391        UsageWindow::Range { .. } => Err(Error::InvalidInput(
392            "usage range end must be later than its start".into(),
393        )),
394    }
395}
396
397#[cfg(test)]
398mod tests {
399    use super::*;
400    use crate::{CostAccuracy, Modality, ModalityTokens};
401
402    fn sample_usage() -> TokenUsage {
403        TokenUsage {
404            input_tokens: 10,
405            output_tokens: 4,
406            total_tokens: 14,
407            input_by_modality: vec![ModalityTokens {
408                modality: Modality::Text,
409                tokens: 10,
410            }],
411            ..TokenUsage::default()
412        }
413    }
414
415    fn sample_cost(nanos: u64) -> CostBreakdown {
416        CostBreakdown {
417            input: Money::from_usd_nanos(nanos),
418            total: Money::from_usd_nanos(nanos),
419            accuracy: CostAccuracy::Exact,
420            ..CostBreakdown::default()
421        }
422    }
423
424    #[test]
425    fn limits_block_at_the_boundary() {
426        let accounting = Accounting::new();
427        accounting
428            .set_limits(SpendingLimits {
429                hourly: Some(Money::from_usd_nanos(100)),
430                daily: None,
431                monthly: None,
432            })
433            .unwrap();
434        accounting
435            .record(
436                "infer_flash_lite",
437                "gemini-3.1-flash-lite",
438                ServiceTier::Standard,
439                true,
440                Some("interaction"),
441                None,
442                &sample_usage(),
443                &sample_cost(100),
444            )
445            .unwrap();
446        let error = accounting.enforce_limits(Utc::now()).unwrap_err();
447        assert!(matches!(
448            error,
449            Error::SpendingLimitReached {
450                period: LimitPeriod::Hourly,
451                ..
452            }
453        ));
454    }
455
456    #[test]
457    fn breakdown_groups_without_storing_content() {
458        let accounting = Accounting::new();
459        accounting
460            .record(
461                "infer_pro",
462                "gemini-3.1-pro-preview",
463                ServiceTier::Standard,
464                true,
465                Some("interaction"),
466                None,
467                &sample_usage(),
468                &sample_cost(5_000),
469            )
470            .unwrap();
471        let report = accounting.breakdown(UsageWindow::AllTime).unwrap();
472        assert_eq!(report.totals.requests, 1);
473        assert_eq!(report.totals.usage.input_tokens, 10);
474        assert_eq!(report.totals.cost.total.usd_nanos(), 5_000);
475        assert_eq!(report.by_model[0].key, "gemini-3.1-pro-preview");
476        let records = accounting.usage_records(UsageWindow::AllTime, 10).unwrap();
477        assert_eq!(
478            records[0].provider_request_id.as_deref(),
479            Some("interaction")
480        );
481    }
482
483    #[test]
484    fn a_new_session_has_no_records_or_limits() {
485        let previous = Accounting::new();
486        previous
487            .set_limits(SpendingLimits {
488                hourly: Some(Money::from_usd_nanos(100)),
489                ..SpendingLimits::default()
490            })
491            .unwrap();
492        previous
493            .record(
494                "infer_pro",
495                "gemini-3.1-pro-preview",
496                ServiceTier::Standard,
497                true,
498                None,
499                None,
500                &sample_usage(),
501                &sample_cost(50),
502            )
503            .unwrap();
504
505        let next = Accounting::new();
506        assert_eq!(next.limits().unwrap(), SpendingLimits::default());
507        assert!(
508            next.usage_records(UsageWindow::AllTime, 10)
509                .unwrap()
510                .is_empty()
511        );
512    }
513
514    #[test]
515    fn invalid_range_is_rejected() {
516        let now = Utc::now();
517        assert!(
518            Accounting::new()
519                .breakdown(UsageWindow::Range {
520                    start: now,
521                    end: now,
522                })
523                .is_err()
524        );
525    }
526}
kcode_gemini_api/accounting.rs

kcode_gemini_api/
accounting.rs
