aspect_std/

tokenbudget.rs

//! Token-budget aspect for agentic LLM systems.
//!
//! Tracks cumulative LLM token spend per session and rejects calls
//! once a declared budget is exhausted. The aspect is novel relative
//! to classical rate limiters because it caps a *cost* dimension
//! (tokens) rather than a request count, and because budgets are
//! attributed per session rather than per process.

use aspect_core::{Aspect, AspectError, ProceedingJoinPoint};
use parking_lot::RwLock;
use std::any::Any;
use std::collections::HashMap;
use std::sync::Arc;

/// Identifier for the session whose token spend is being capped.
pub type SessionId = String;

/// Per-call context the host populates before the guarded join point.
#[derive(Clone, Debug)]
pub struct TokenCharge {
    /// Session whose budget is debited.
    pub session: SessionId,
    /// Tokens the call is expected to consume (input + output estimate).
    pub estimated: u64,
}

/// Token-budget aspect.
///
/// The host system calls `set_budget(session, max)` once and `bind(charge)`
/// before each guarded invocation. The around advice debits the budget
/// if the call would still fit; otherwise it short-circuits with an error.
#[derive(Clone, Default)]
pub struct TokenBudgetAspect {
    inner: Arc<Inner>,
}

#[derive(Default)]
struct Inner {
    budgets: RwLock<HashMap<SessionId, Budget>>,
    pending: RwLock<HashMap<u64, TokenCharge>>,
    next_key: RwLock<u64>,
}

#[derive(Clone, Debug)]
struct Budget {
    cap: u64,
    spent: u64,
}

impl TokenBudgetAspect {
    /// Create an empty budget tracker.
    pub fn new() -> Self {
        Self::default()
    }

    /// Set or reset the budget cap for a session.
    pub fn set_budget(&self, session: impl Into<SessionId>, cap: u64) {
        self.inner
            .budgets
            .write()
            .insert(session.into(), Budget { cap, spent: 0 });
    }

    /// Read remaining tokens for a session. Returns `None` if no budget
    /// has been declared.
    pub fn remaining(&self, session: &str) -> Option<u64> {
        self.inner
            .budgets
            .read()
            .get(session)
            .map(|b| b.cap.saturating_sub(b.spent))
    }

    /// Bind a charge to the next guarded invocation.
    pub fn bind(&self, charge: TokenCharge) -> u64 {
        let mut next = self.inner.next_key.write();
        *next += 1;
        let key = *next;
        self.inner.pending.write().insert(key, charge);
        key
    }

    /// Clear a previously bound charge. Idempotent.
    pub fn release(&self, key: u64) {
        self.inner.pending.write().remove(&key);
    }

    fn most_recent_charge(&self) -> Option<TokenCharge> {
        self.inner
            .pending
            .read()
            .iter()
            .max_by_key(|(k, _)| *k)
            .map(|(_, v)| v.clone())
    }

    fn try_debit(&self, charge: &TokenCharge) -> Result<(), AspectError> {
        let mut budgets = self.inner.budgets.write();
        let entry = budgets.get_mut(&charge.session).ok_or_else(|| {
            AspectError::execution(format!(
                "no budget declared for session '{}'",
                charge.session
            ))
        })?;
        let projected = entry.spent.saturating_add(charge.estimated);
        if projected > entry.cap {
            return Err(AspectError::execution(format!(
                "token budget exceeded for session '{}' ({} + {} > {})",
                charge.session, entry.spent, charge.estimated, entry.cap
            )));
        }
        entry.spent = projected;
        Ok(())
    }
}

impl Aspect for TokenBudgetAspect {
    fn around(&self, pjp: ProceedingJoinPoint) -> Result<Box<dyn Any>, AspectError> {
        if let Some(charge) = self.most_recent_charge() {
            self.try_debit(&charge)?;
        }
        pjp.proceed()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use aspect_core::{JoinPoint, Location};

    fn jp() -> JoinPoint {
        JoinPoint {
            function_name: "execute_tool_call",
            module_path: "agent::dispatch",
            location: Location {
                file: "agent.rs",
                line: 1,
            },
        }
    }

    #[test]
    fn budget_admits_calls_within_cap() {
        let budget = TokenBudgetAspect::new();
        budget.set_budget("s1", 1000);
        let key = budget.bind(TokenCharge {
            session: "s1".into(),
            estimated: 300,
        });

        let pjp = ProceedingJoinPoint::new(|| Ok(Box::new(0_i32) as Box<dyn Any>), jp());
        let result = budget.around(pjp);
        budget.release(key);

        assert!(result.is_ok());
        assert_eq!(budget.remaining("s1"), Some(700));
    }

    #[test]
    fn budget_denies_call_that_would_overspend() {
        let budget = TokenBudgetAspect::new();
        budget.set_budget("s1", 500);
        let key = budget.bind(TokenCharge {
            session: "s1".into(),
            estimated: 600,
        });

        let pjp = ProceedingJoinPoint::new(|| Ok(Box::new(0_i32) as Box<dyn Any>), jp());
        let result = budget.around(pjp);
        budget.release(key);

        assert!(result.is_err());
        assert_eq!(budget.remaining("s1"), Some(500));
    }

    #[test]
    fn unknown_session_is_rejected() {
        let budget = TokenBudgetAspect::new();
        let key = budget.bind(TokenCharge {
            session: "anon".into(),
            estimated: 1,
        });
        let pjp = ProceedingJoinPoint::new(|| Ok(Box::new(0_i32) as Box<dyn Any>), jp());
        let result = budget.around(pjp);
        budget.release(key);
        assert!(result.is_err());
    }

    #[test]
    fn cumulative_debits_eventually_exhaust_budget() {
        let budget = TokenBudgetAspect::new();
        budget.set_budget("s1", 100);
        for _ in 0..4 {
            let key = budget.bind(TokenCharge {
                session: "s1".into(),
                estimated: 25,
            });
            let pjp = ProceedingJoinPoint::new(|| Ok(Box::new(0_i32) as Box<dyn Any>), jp());
            assert!(budget.around(pjp).is_ok());
            budget.release(key);
        }
        let key = budget.bind(TokenCharge {
            session: "s1".into(),
            estimated: 1,
        });
        let pjp = ProceedingJoinPoint::new(|| Ok(Box::new(0_i32) as Box<dyn Any>), jp());
        assert!(budget.around(pjp).is_err());
        budget.release(key);
    }
}