Skip to content

feat(kernel): add token remaining/usage-ratio methods to BudgetStatus#1652

Open
SH20RAJ wants to merge 1 commit into
mofa-org:mainfrom
SH20RAJ:feat/1475-budget-status-token-methods
Open

feat(kernel): add token remaining/usage-ratio methods to BudgetStatus#1652
SH20RAJ wants to merge 1 commit into
mofa-org:mainfrom
SH20RAJ:feat/1475-budget-status-token-methods

Conversation

@SH20RAJ
Copy link
Copy Markdown
Contributor

@SH20RAJ SH20RAJ commented Apr 20, 2026

Summary

Fixes #1475.

BudgetStatus in crates/mofa-kernel/src/budget.rs exposes three
convenience methods for cost tracking but nothing equivalent for
token tracking
, forcing callers to write manual calculations:

// callers had to write:
let remaining = status.config.max_tokens_per_session
    .map(|max| max.saturating_sub(status.session_tokens));

New API

This PR adds four methods that complete the symmetry between the cost
and token dimensions:

impl BudgetStatus {
    /// Tokens remaining in the current session (clamped to 0).
    pub fn remaining_session_tokens(&self) -> Option<u64>

    /// Tokens remaining in today's budget (clamped to 0).
    pub fn remaining_daily_tokens(&self) -> Option<u64>

    /// Session token usage ratio in [0.0, ∞). ≥ 1.0 signals overrun.
    pub fn session_token_usage_ratio(&self) -> Option<f64>

    /// Daily cost usage ratio in [0.0, ∞). (mirrors session_cost_usage_ratio)
    pub fn daily_cost_usage_ratio(&self) -> Option<f64>
}

All methods return None when the corresponding limit is not configured,
matching the existing cost-method contract exactly.

Design Notes

  • remaining_*_tokens uses saturating_sub so it always returns 0
    rather than panicking on overrun (u64 arithmetic).
  • session_token_usage_ratio returns a value > 1.0 on overrun, which
    lets callers distinguish "near limit" from "already exceeded" without a
    separate is_exceeded() call.

Tests added

Test Covers
test_budget_status_token_remaining_methods Basic remaining calculation
test_budget_status_token_usage_ratio Ratio at 25 %
test_budget_status_token_remaining_clamps_to_zero Overrun → 0 remaining, ratio > 1.0
test_budget_status_no_token_limits_returns_none No limits → None
cargo test -p mofa-kernel budget::tests

Checklist

Fixes mofa-org#1475

BudgetStatus had convenience helpers for cost tracking
(remaining_session_cost, remaining_daily_cost, session_cost_usage_ratio)
but lacked the equivalent methods for token tracking, forcing callers
to write manual calculations and creating an inconsistent API surface.

Add four new methods that mirror the cost conveniences:
- remaining_session_tokens() -> Option<u64>
- remaining_daily_tokens()   -> Option<u64>
- session_token_usage_ratio() -> Option<f64>
- daily_cost_usage_ratio()    -> Option<f64>   (completes cost side too)

All return None when the corresponding limit is not configured.
remaining_session_tokens / remaining_daily_tokens use saturating_sub to
clamp to 0 on overrun; the ratio methods return a value >= 1.0 to signal
overrun without panicking.

Tests added:
- test_budget_status_token_remaining_methods
- test_budget_status_token_usage_ratio
- test_budget_status_token_remaining_clamps_to_zero
- test_budget_status_no_token_limits_returns_none
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] BudgetStatus missing token remaining/usage-ratio methods

1 participant