Struct SessionBudget 

pub struct SessionBudget { /* private fields */ }

Thread-safe per-session token budget tracker.

All methods accept the session ID as a &str to avoid requiring a SessionId wrapper at the session-management layer.

Thread Safety

Internally uses a Mutex<HashMap>. For very high-throughput workloads consider sharding by session prefix.

Implementations

impl SessionBudget

pub fn new() -> Self

Create a new, empty budget store.

Panics

This function does not panic.

pub fn set_limits(&self, session_id: impl Into<String>, limits: SessionLimits)

Register or update limits for a session.

If the session already exists its accumulated spend is preserved; only the limits are updated.

Panics

This function does not panic.

pub fn remove_session(&self, session_id: &str)

Remove a session from the budget store entirely.

Useful when a session is logically closed and its budget data no longer needs to be retained.

Panics

This function does not panic.

pub fn record_tokens( &self, session_id: &str, tokens: u64, ) -> Result<BudgetOutcome, BudgetError>

Record tokens spent by session_id and check against limits.

Hard limit enforcement

If adding tokens would cause accumulated spend to reach or exceed hard_limit_tokens, the spend is not recorded and Err(BudgetError::HardLimitExceeded) is returned. The caller should reject the underlying inference request with a clear user-facing error before dispatching to a worker.

Soft limit warning

If the new accumulated spend crosses soft_limit_tokens the function returns Ok(BudgetOutcome::SoftLimitWarning { .. }). The request proceeds normally; callers may choose to surface a warning to the user.

Unregistered sessions

Sessions with no registered limits are silently accepted — all spend calls return Ok(BudgetOutcome::Ok) and no tracking is performed.

Errors

Returns BudgetError::HardLimitExceeded when the hard limit is hit.

Panics

This function does not panic.

pub fn reset_session(&self, session_id: &str)

Reset accumulated spend to zero without changing the session limits.

This is the operation behind POST /v1/sessions/{id}/budget/reset.

Does nothing if the session is not registered (no-op, not an error).

Panics

This function does not panic.

pub fn daily_reset(&self) -> usize

Perform a daily reset: zero the spend of every session whose last_reset was more than 24 hours ago (or was never set).

Intended to be called from a background task, e.g.:

use tokio::time::{interval, Duration};
use tokio_prompt_orchestrator::session::budget::SessionBudget;
use std::sync::Arc;

async fn daily_reset_task(budget: Arc<SessionBudget>) {
    let mut ticker = interval(Duration::from_secs(60 * 60)); // check hourly
    loop {
        ticker.tick().await;
        let count = budget.daily_reset();
        if count > 0 {
            tracing::info!(count, "daily budget reset applied to sessions");
        }
    }
}

Returns

Number of sessions whose budget was reset.

Panics

This function does not panic.

pub fn snapshot(&self, session_id: &str) -> Option<SessionBudgetSnapshot>

Return a snapshot of the given session’s budget.

This is the backing data for GET /v1/sessions/{id}/budget.

Returns None if the session has not been registered.

Panics

This function does not panic.

pub fn all_snapshots(&self) -> Vec<SessionBudgetSnapshot>

Return snapshots for all tracked sessions.

Panics

This function does not panic.

pub fn session_count(&self) -> usize

Return the number of sessions currently tracked.

Panics

This function does not panic.

Trait Implementations

impl Debug for SessionBudget

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Default for SessionBudget

fn default() -> SessionBudget

Returns the “default value” for a type.