Skip to main content

everruns_core/
budget.rs

1// Budget domain types
2//
3// Extensible budgeting system for controlling resource consumption.
4// Supports multiple currencies (USD, tokens, credits), pluggable meters,
5// pluggable rules, and soft enforcement (pause/warn/stop).
6//
7// See specs/budgeting.md for the full specification.
8
9use chrono::{DateTime, Utc};
10use serde::{Deserialize, Serialize};
11
12use crate::typed_id::{BudgetId, SessionId};
13use crate::user_facing_error::UserFacingErrorFields;
14
15#[cfg(feature = "openapi")]
16use utoipa::ToSchema;
17
18// ============================================================================
19// Budget
20// ============================================================================
21
22/// Budget status.
23#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
24#[cfg_attr(feature = "openapi", derive(ToSchema))]
25#[serde(rename_all = "snake_case")]
26pub enum BudgetStatus {
27    Active,
28    Paused,
29    Exhausted,
30    Disabled,
31}
32
33impl std::fmt::Display for BudgetStatus {
34    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
35        match self {
36            BudgetStatus::Active => write!(f, "active"),
37            BudgetStatus::Paused => write!(f, "paused"),
38            BudgetStatus::Exhausted => write!(f, "exhausted"),
39            BudgetStatus::Disabled => write!(f, "disabled"),
40        }
41    }
42}
43
44impl From<&str> for BudgetStatus {
45    fn from(s: &str) -> Self {
46        match s {
47            "active" => BudgetStatus::Active,
48            "paused" => BudgetStatus::Paused,
49            "exhausted" => BudgetStatus::Exhausted,
50            "disabled" => BudgetStatus::Disabled,
51            _ => BudgetStatus::Active,
52        }
53    }
54}
55
56/// Subject type: what entity this budget constrains.
57#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
58#[cfg_attr(feature = "openapi", derive(ToSchema))]
59#[serde(rename_all = "snake_case")]
60pub enum BudgetSubjectType {
61    Session,
62    Agent,
63    User,
64    Organization,
65    /// Bound to an `App` (every session created for the app counts).
66    App,
67    /// Bound to a single `AppChannel` (only sessions for that channel count).
68    AppChannel,
69}
70
71impl BudgetSubjectType {
72    /// Wire string used in storage and the API.
73    pub fn as_wire(&self) -> &'static str {
74        match self {
75            BudgetSubjectType::Session => "session",
76            BudgetSubjectType::Agent => "agent",
77            BudgetSubjectType::User => "user",
78            BudgetSubjectType::Organization => "org",
79            BudgetSubjectType::App => "app",
80            BudgetSubjectType::AppChannel => "app_channel",
81        }
82    }
83}
84
85impl std::fmt::Display for BudgetSubjectType {
86    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
87        f.write_str(self.as_wire())
88    }
89}
90
91impl From<&str> for BudgetSubjectType {
92    fn from(s: &str) -> Self {
93        match s {
94            "session" => BudgetSubjectType::Session,
95            "agent" => BudgetSubjectType::Agent,
96            "user" => BudgetSubjectType::User,
97            "org" | "organization" => BudgetSubjectType::Organization,
98            "app" => BudgetSubjectType::App,
99            "app_channel" => BudgetSubjectType::AppChannel,
100            _ => BudgetSubjectType::Session,
101        }
102    }
103}
104
105/// Budget period configuration for recurring budgets.
106///
107/// Periods drive automatic balance reset:
108/// - `Duration` is a fixed-length sliding window (e.g. last 5 hours, last 30 days)
109///   measured from `Budget::period_started_at`. When the window elapses the
110///   balance is reset to `limit` and the window restarts.
111/// - `Calendar` aligns to a calendar boundary (`hour | day | week | month | year`)
112///   in UTC. The balance resets when the next boundary is crossed.
113/// - `Rolling` is preserved for backwards compatibility and parses common
114///   shorthand (`24h`, `5h`, `7d`, `30d`) into a `Duration`-equivalent reset
115///   policy.
116#[derive(Debug, Clone, Serialize, Deserialize)]
117#[cfg_attr(feature = "openapi", derive(ToSchema))]
118#[serde(tag = "type", rename_all = "snake_case")]
119pub enum BudgetPeriod {
120    /// Sliding window of a configurable number of seconds.
121    Duration { seconds: u64 },
122    /// Rolling window described as a humanized string ("5h", "24h", "30d").
123    Rolling { window: String },
124    /// Calendar-aligned (`hour`, `day`, `week`, `month`, `year`).
125    Calendar { unit: String },
126}
127
128impl BudgetPeriod {
129    /// Length of the period in seconds, if it can be expressed as a fixed
130    /// duration. Calendar periods return `None` (handled separately).
131    pub fn duration_seconds(&self) -> Option<u64> {
132        match self {
133            BudgetPeriod::Duration { seconds } => Some(*seconds),
134            BudgetPeriod::Rolling { window } => parse_rolling_window(window),
135            BudgetPeriod::Calendar { .. } => None,
136        }
137    }
138}
139
140/// Parse a rolling window shorthand like "5h", "30m", "7d" into seconds.
141fn parse_rolling_window(window: &str) -> Option<u64> {
142    let trimmed = window.trim();
143    if trimmed.is_empty() {
144        return None;
145    }
146    let (digits, suffix) = trimmed.split_at(
147        trimmed
148            .find(|c: char| !c.is_ascii_digit())
149            .unwrap_or(trimmed.len()),
150    );
151    let value: u64 = digits.parse().ok()?;
152    let multiplier: u64 = match suffix.trim().to_ascii_lowercase().as_str() {
153        "" | "s" | "sec" | "secs" | "second" | "seconds" => 1,
154        "m" | "min" | "mins" | "minute" | "minutes" => 60,
155        "h" | "hr" | "hrs" | "hour" | "hours" => 3_600,
156        "d" | "day" | "days" => 86_400,
157        "w" | "wk" | "wks" | "week" | "weeks" => 604_800,
158        _ => return None,
159    };
160    value.checked_mul(multiplier)
161}
162
163/// Budget — a spending cap for a subject in a currency.
164/// API response DTO.
165#[derive(Debug, Clone, Serialize, Deserialize)]
166#[cfg_attr(feature = "openapi", derive(ToSchema))]
167pub struct Budget {
168    #[cfg_attr(feature = "openapi", schema(value_type = String, example = "bdgt_01933b5a00007000800000000000001"))]
169    pub id: BudgetId,
170    pub organization_id: String,
171    pub subject_type: BudgetSubjectType,
172    /// Public ID of the subject entity.
173    pub subject_id: String,
174    /// Currency: "usd", "tokens", "credits", or custom.
175    pub currency: String,
176    /// Hard limit — budget ceiling.
177    pub limit: f64,
178    /// Soft limit — triggers pause/warn when balance drops below this.
179    #[serde(skip_serializing_if = "Option::is_none")]
180    pub soft_limit: Option<f64>,
181    /// Current remaining balance (limit minus consumed).
182    pub balance: f64,
183    /// Optional period for recurring budgets.
184    #[serde(skip_serializing_if = "Option::is_none")]
185    pub period: Option<BudgetPeriod>,
186    /// When the current period started (used to detect period rollover for
187    /// `Duration` / `Rolling` periods, and to display "resets at" in the UI).
188    /// `None` for budgets without a period.
189    #[serde(skip_serializing_if = "Option::is_none")]
190    pub period_started_at: Option<DateTime<Utc>>,
191    /// Arbitrary metadata.
192    #[serde(default, skip_serializing_if = "Option::is_none")]
193    pub metadata: Option<serde_json::Value>,
194    pub status: BudgetStatus,
195    pub created_at: DateTime<Utc>,
196    pub updated_at: DateTime<Utc>,
197}
198
199// ============================================================================
200// Ledger Entry
201// ============================================================================
202
203/// Immutable ledger entry recording resource consumption or credit against a budget.
204#[derive(Debug, Clone, Serialize, Deserialize)]
205#[cfg_attr(feature = "openapi", derive(ToSchema))]
206pub struct LedgerEntry {
207    pub id: String,
208    #[cfg_attr(feature = "openapi", schema(value_type = String))]
209    pub budget_id: BudgetId,
210    /// Positive = debit (consumption), negative = credit (top-up/refund).
211    pub amount: f64,
212    /// Which meter produced this: "llm_tokens", "tool_calls", etc.
213    pub meter_source: String,
214    /// Reference entity type: "llm_generation", "tool_execution", "manual".
215    #[serde(skip_serializing_if = "Option::is_none")]
216    pub ref_type: Option<String>,
217    /// Reference entity ID.
218    #[serde(skip_serializing_if = "Option::is_none")]
219    pub ref_id: Option<String>,
220    /// Session context for this entry.
221    #[serde(skip_serializing_if = "Option::is_none")]
222    #[cfg_attr(feature = "openapi", schema(value_type = Option<String>))]
223    pub session_id: Option<SessionId>,
224    #[serde(skip_serializing_if = "Option::is_none")]
225    pub description: Option<String>,
226    pub created_at: DateTime<Utc>,
227}
228
229// ============================================================================
230// Budget Rule Actions
231// ============================================================================
232
233/// Action returned by a budget rule after evaluation.
234#[derive(Debug, Clone, PartialEq)]
235pub enum BudgetAction {
236    /// No action needed, continue execution.
237    Continue,
238    /// Emit a warning event but keep running.
239    Warn { message: String },
240    /// Pause the session — requires user input to resume.
241    Pause { message: String },
242    /// Hard stop — terminate the current turn.
243    Stop { message: String },
244}
245
246// ============================================================================
247// Budget check result (used by worker to decide what to do)
248// ============================================================================
249
250/// Result of checking all budgets for a session.
251#[derive(Debug, Clone, Serialize, Deserialize)]
252#[cfg_attr(feature = "openapi", derive(ToSchema))]
253pub struct BudgetCheckResult {
254    /// Most restrictive action across all budgets.
255    pub action: String, // "continue", "warn", "pause", "stop"
256    /// Human-readable message (set when action != "continue").
257    #[serde(skip_serializing_if = "Option::is_none")]
258    pub message: Option<String>,
259    /// Budget that triggered the action.
260    #[serde(skip_serializing_if = "Option::is_none")]
261    #[cfg_attr(feature = "openapi", schema(value_type = Option<String>))]
262    pub budget_id: Option<BudgetId>,
263    /// Remaining balance on the most restrictive budget.
264    #[serde(skip_serializing_if = "Option::is_none")]
265    pub balance: Option<f64>,
266    /// Currency of the most restrictive budget.
267    #[serde(skip_serializing_if = "Option::is_none")]
268    pub currency: Option<String>,
269    /// Stable error code for user-facing budget failures.
270    #[serde(default, skip_serializing_if = "Option::is_none")]
271    pub error_code: Option<String>,
272    /// Structured interpolation fields for localized error rendering.
273    #[serde(default, skip_serializing_if = "Option::is_none")]
274    #[cfg_attr(feature = "openapi", schema(value_type = Option<Object>))]
275    pub error_fields: Option<UserFacingErrorFields>,
276}
277
278impl BudgetCheckResult {
279    pub fn ok() -> Self {
280        Self {
281            action: "continue".into(),
282            message: None,
283            budget_id: None,
284            balance: None,
285            currency: None,
286            error_code: None,
287            error_fields: None,
288        }
289    }
290
291    pub fn should_stop(&self) -> bool {
292        self.action == "stop"
293    }
294
295    pub fn should_pause(&self) -> bool {
296        self.action == "pause"
297    }
298}
299
300// ============================================================================
301// Budget tool response (returned by check_budget tool)
302// ============================================================================
303
304/// Summary of a single budget for the check_budget tool response.
305#[derive(Debug, Clone, Serialize, Deserialize)]
306pub struct BudgetSummary {
307    pub currency: String,
308    pub limit: f64,
309    pub balance: f64,
310    #[serde(skip_serializing_if = "Option::is_none")]
311    pub soft_limit: Option<f64>,
312    pub percent_remaining: f64,
313    pub status: String,
314}
315
316/// Full response from the check_budget tool.
317#[derive(Debug, Clone, Serialize, Deserialize)]
318pub struct BudgetToolResponse {
319    /// Overall status: "active", "warning", "paused", "exhausted", "no_budgets"
320    pub status: String,
321    /// Per-budget summaries
322    pub budgets: Vec<BudgetSummary>,
323    /// Human-readable hint for the agent
324    #[serde(skip_serializing_if = "Option::is_none")]
325    pub hint: Option<String>,
326}