lex_store/budget.rs
1//! Per-session budget ledger (#292 slice 1).
2//!
3//! Today `OperationKind::budget_delta()` records `(from, to)`
4//! budget pairs on every op that touches a `[budget(N)]`-bearing
5//! function. `lex audit --budget` already rolls those up per
6//! signature. What's missing is the per-session aggregate: "how
7//! much budget did session X cause to be spent across all the ops
8//! it authored?"
9//!
10//! This module is the read-only ledger. Slice 2 layers a
11//! `policy.json` cap on top; slice 3 wires the apply-path gate.
12//!
13//! # Spend model
14//!
15//! For each op tagged with an `intent_id` resolving to a session:
16//!
17//! - `AddFunction` with `budget_cost = Some(n)` contributes `n`.
18//! - `ModifyBody` / `ChangeEffectSig` / `ReplaceMatchArm` /
19//! `RenameLocal` / `InlineLet` with `(from_budget, to_budget)`
20//! contribute `max(0, to - from)` (only budget *increases*
21//! count toward spend; refactor-to-cheaper doesn't refund).
22//! - Ops without `intent_id` are excluded — there's no session to
23//! attribute them to.
24//!
25//! # Cost
26//!
27//! Single walk of the op log + one `IntentLog::get` per distinct
28//! intent. For studies past ~100k ops, slice 2 will add an on-disk
29//! cache keyed by `(session_id, head_op)` so re-reads are O(1).
30
31use serde::{Deserialize, Serialize};
32use std::collections::BTreeMap;
33
34use crate::store::{Store, StoreError};
35
36/// Rollup of a single session's budget spend.
37///
38/// `cap` and `remaining` are populated from `policy.json`'s
39/// `session_budgets` (#292 slices 2 + 3). When no cap is set
40/// either by per-session override or by `default_cap`, both
41/// fields are `None`.
42#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
43pub struct SessionBudget {
44 pub session_id: String,
45 /// Sum of monotonic budget cost over all ops in this session.
46 pub spent: u64,
47 /// How many ops were attributed to this session (only those
48 /// that contributed a non-zero increment count).
49 pub op_count: usize,
50 /// Resolved cap from `policy.session_budgets`. `None` means
51 /// no enforcement.
52 #[serde(default, skip_serializing_if = "Option::is_none")]
53 pub cap: Option<u64>,
54 /// `cap - spent` when `cap` is set. Negative when over.
55 /// `None` when uncapped.
56 #[serde(default, skip_serializing_if = "Option::is_none")]
57 pub remaining: Option<i64>,
58}
59
60impl Store {
61 /// Compute the budget spent by the given `session_id` across
62 /// every op currently reachable from any branch head. Returns
63 /// `(spent: 0, op_count: 0)` for unknown sessions, with cap
64 /// populated from `policy.session_budgets`.
65 pub fn session_budget(&self, session_id: &str) -> Result<SessionBudget, StoreError> {
66 let all = self.all_session_budgets()?;
67 if let Some(b) = all.into_iter().find(|b| b.session_id == session_id) {
68 return Ok(b);
69 }
70 // Unknown session: zero spend, but the cap (if any) still
71 // applies. Useful for "show me what budget this brand-new
72 // session has" queries.
73 let cap = self.session_budget_cap(session_id)?;
74 let remaining = cap.map(|c| c as i64);
75 Ok(SessionBudget {
76 session_id: session_id.into(),
77 spent: 0,
78 op_count: 0,
79 cap,
80 remaining,
81 })
82 }
83
84 /// Resolve the budget cap configured for `session_id` from
85 /// `policy.json`'s `session_budgets` (#292 slice 2). Returns
86 /// `None` when no enforcement is configured.
87 pub fn session_budget_cap(&self, session_id: &str) -> Result<Option<u64>, StoreError> {
88 let policy = crate::policy::load(self.root())?.unwrap_or_default();
89 Ok(policy.session_budgets.cap_for(session_id))
90 }
91
92 /// Compute per-session budget rollups across every branch.
93 /// Returns one entry per distinct session that contributed at
94 /// least one budget-bearing op. Sorted by `session_id` so the
95 /// output is deterministic.
96 pub fn all_session_budgets(&self) -> Result<Vec<SessionBudget>, StoreError> {
97 let log = lex_vcs::OpLog::open(self.root())?;
98 let intent_log = lex_vcs::IntentLog::open(self.root())?;
99
100 // Collect the union of ops reachable from every branch
101 // head. Walking each branch separately and unioning by
102 // op_id avoids double-counting on diamond histories.
103 let mut visited: std::collections::BTreeSet<lex_vcs::OpId> = Default::default();
104 let mut records: Vec<lex_vcs::OperationRecord> = Vec::new();
105 for branch_name in self.list_branches()? {
106 let Some(branch) = self.get_branch(&branch_name)? else { continue };
107 let Some(head) = branch.head_op else { continue };
108 for rec in log.walk_back(&head, None)? {
109 if visited.insert(rec.op_id.clone()) {
110 records.push(rec);
111 }
112 }
113 }
114
115 // Walk records → resolve intent → resolve session → tally.
116 // Intent lookups are cached so we don't hit the IntentLog
117 // once per op when many ops share an intent.
118 let mut intent_to_session: BTreeMap<String, Option<String>> = BTreeMap::new();
119 let mut buckets: BTreeMap<String, (u64, usize)> = BTreeMap::new();
120 for rec in &records {
121 let Some(intent_id) = rec.op.intent_id.as_deref() else { continue };
122 let session = match intent_to_session.get(intent_id) {
123 Some(s) => s.clone(),
124 None => {
125 let s = intent_log.get(&intent_id.to_string())?
126 .map(|i| i.session_id);
127 intent_to_session.insert(intent_id.into(), s.clone());
128 s
129 }
130 };
131 let Some(session_id) = session else { continue };
132
133 let increment = monotonic_spend(&rec.op.kind);
134 if increment == 0 { continue; }
135 let entry = buckets.entry(session_id).or_insert((0, 0));
136 entry.0 += increment;
137 entry.1 += 1;
138 }
139
140 let policy = crate::policy::load(self.root())?.unwrap_or_default();
141 let out: Vec<SessionBudget> = buckets
142 .into_iter()
143 .map(|(session_id, (spent, op_count))| {
144 let cap = policy.session_budgets.cap_for(&session_id);
145 let remaining = cap.map(|c| (c as i64) - (spent as i64));
146 SessionBudget { session_id, spent, op_count, cap, remaining }
147 })
148 .collect();
149 Ok(out)
150 }
151}
152
153/// Convert an op's `budget_delta` into a monotonic spend amount.
154/// `AddFunction` contributes its full `budget_cost`; modify-shape
155/// ops contribute the delta only when budget *increased*.
156fn monotonic_spend(kind: &lex_vcs::OperationKind) -> u64 {
157 monotonic_spend_of(kind)
158}
159
160/// Crate-public form used by [`crate::Store::apply_operation_checked`]'s
161/// budget gate (#292 slice 3). Same semantics as the private
162/// helper; exposed under a separate name so the test-only
163/// `monotonic_spend` keeps its `#[cfg(test)]`-friendly shape.
164pub(crate) fn monotonic_spend_of(kind: &lex_vcs::OperationKind) -> u64 {
165 let (from, to) = kind.budget_delta();
166 match (from, to) {
167 (None, Some(n)) => n,
168 (Some(f), Some(t)) if t > f => t - f,
169 _ => 0,
170 }
171}
172
173#[cfg(test)]
174mod tests {
175 use super::*;
176
177 #[test]
178 fn monotonic_spend_handles_each_shape() {
179 use lex_vcs::OperationKind;
180 // AddFunction: full cost contributes.
181 let k = OperationKind::AddFunction {
182 sig_id: "f".into(),
183 stage_id: "s".into(),
184 effects: Default::default(),
185 budget_cost: Some(10),
186 };
187 assert_eq!(monotonic_spend(&k), 10);
188
189 // ModifyBody: only increases count.
190 let k = OperationKind::ModifyBody {
191 sig_id: "f".into(),
192 from_stage_id: "a".into(),
193 to_stage_id: "b".into(),
194 from_budget: Some(10),
195 to_budget: Some(15),
196 };
197 assert_eq!(monotonic_spend(&k), 5);
198
199 let k = OperationKind::ModifyBody {
200 sig_id: "f".into(),
201 from_stage_id: "a".into(),
202 to_stage_id: "b".into(),
203 from_budget: Some(15),
204 to_budget: Some(10),
205 };
206 assert_eq!(monotonic_spend(&k), 0, "decrease doesn't refund");
207
208 // ModifyBody with no budget data on either side: zero.
209 let k = OperationKind::ModifyBody {
210 sig_id: "f".into(),
211 from_stage_id: "a".into(),
212 to_stage_id: "b".into(),
213 from_budget: None,
214 to_budget: None,
215 };
216 assert_eq!(monotonic_spend(&k), 0);
217 }
218}