memra_server/metering.rs
1//! The admission/accounting seam (lane engine-billing-extraction-20260829).
2//!
3//! The server's job at this boundary is to ADMIT, DENY, and REPORT COUNTS. What
4//! admission means — budgets, prices, tenancy policy — is the deployment's business,
5//! supplied behind these traits. The stock binary wires the in-repo reference
6//! implementation (`ledger::ReferenceMetering`); a deployment-owned binary can wire its own.
7//! Everything here speaks tokens and verdicts, never money: the vocabulary is the
8//! boundary, and it is what lets the policy half live outside this crate.
9//!
10//! The seam is `pub`: a deployment-owned binary wires its implementation through
11//! [`crate::ServerWiring`] and runs the same server.
12
13use std::any::Any;
14
15/// An opaque reservation handle minted by [`Metering::reserve`] and consumed by
16/// [`Metering::open`] of the SAME implementation. The server carries it between the
17/// two calls without looking inside; dropping it un-consumed must release whatever
18/// it holds (the reference implementation refunds on drop).
19pub type Permit = Box<dyn Any + Send>;
20
21/// The request identity a receipt is opened with. All borrowed: this is a view of
22/// state the handler already owns, taken for the duration of one `open` call.
23pub struct RequestMeta<'a> {
24 pub request_id: &'a str,
25 pub tenant: &'a str,
26 /// The authenticated key's non-secret identification prefix, when the request
27 /// carries per-key identity (multi-key ring). What per-PRINCIPAL policy means —
28 /// spend caps, quotas — is the implementation's business.
29 pub principal: Option<&'a str>,
30 pub model: &'a str,
31 pub route: &'static str,
32 pub lane: &'static str,
33 pub stream: bool,
34 /// The request's EFFECTIVE completion-token bound at receipt time: the caller's
35 /// `max_tokens`, or the model-default/`max_output` it resolved to
36 /// (`apply_model_request_limits`); `None` when no bound resolved (context-bounded
37 /// generation). D2 gap G4 (darklanes `research/d2-shadow-20260831/RESULTS.md`):
38 /// the ledger row carries this so live-shadow replay reads the real bound instead
39 /// of recovering it by inverting the billing reservation. An implementation with a
40 /// positional row format appends this as a NEW trailing column (column order in an
41 /// existing ledger is an API, so the discipline is append-only).
42 pub max_tokens: Option<u64>,
43 /// The context reservation the budget admission charged for this request:
44 /// `prompt_tokens + completion bound`, in tokens (the same quantities handed to
45 /// [`Metering::reserve`]). `None` when no reservation ran (limits not enforced, or
46 /// the receipt records a pre-reservation rejection). Same G4 append-only column
47 /// rule as `max_tokens`.
48 pub reserved_ctx: Option<u64>,
49}
50
51/// Token counts as the worker measured them. The one usage shape that crosses the
52/// seam; anything priced is derived from these on the implementation's side.
53#[derive(Debug, Clone, Copy, Default)]
54pub struct UsageCounts {
55 pub prompt_tokens: u64,
56 pub cached_prompt_tokens: u64,
57 pub completion_tokens: u64,
58}
59
60/// Why admission said no. Mirrors the HTTP contract the handlers already speak:
61/// `Insufficient`/`Blocked` answer 402 (one recovery action for callers),
62/// `Unenrolled` answers 402 with its own code, `Unavailable` is the fail-closed 500.
63#[derive(Debug, PartialEq, Eq)]
64pub enum AdmitError {
65 Insufficient,
66 Blocked,
67 Unenrolled,
68 /// The PRINCIPAL (per-key) spend ceiling is reached while the tenant itself may
69 /// still have balance. Its own 402 code: the caller's recovery is raising or
70 /// clearing the key's cap, not adding credit.
71 PrincipalCapped,
72 Unavailable(String),
73}
74
75/// Limits-source health for the operator metrics surface. Counts only.
76#[derive(Debug, Clone, Copy)]
77pub struct LimitsHealth {
78 pub source_reload_failed: u64,
79 pub source_reload_consecutive: u32,
80 pub source_available: bool,
81}
82
83/// Per-deployment admission + usage accounting. One object, present iff the
84/// deployment configured accounting at all (`AppState.metering: Option<Arc<dyn ..>>`
85/// mirrors the old `request_ledger: Option<Ledger>` exactly).
86pub trait Metering: Send + Sync {
87 /// Whether per-tenant admission limits are configured at all. `false` = every
88 /// authenticated tenant is admitted without reservation (counting may still run).
89 fn enforces_limits(&self) -> bool;
90
91 /// Whether this tenant is subject to limits. With limits enforced, an unknown
92 /// tenant is NOT a free pass — the caller rejects it as unenrolled.
93 fn is_limited(&self, tenant: &str) -> Result<bool, AdmitError>;
94
95 /// Reserve headroom for a request's worst case, in tokens. `Ok(Some(permit))`
96 /// rides the receipt and is settled to worker-truth usage; `Ok(None)` means the
97 /// implementation needs no per-request hold. `principal` is the authenticated
98 /// key's non-secret prefix when the request carries one — the hook for per-key
99 /// policy (spend caps) on the implementation's side.
100 fn reserve(
101 &self,
102 tenant: &str,
103 principal: Option<&str>,
104 model: &str,
105 prompt_tokens: u64,
106 completion_bound: u64,
107 ) -> Result<Option<Permit>, AdmitError>;
108
109 /// Open the request's usage receipt. Every terminal outcome settles it through
110 /// one of the [`Receipt`] methods; dropping it unfinalized is the abandoned-client
111 /// path and must stay safe (the reference implementation prices the partial).
112 fn open(&self, meta: &RequestMeta<'_>, permit: Option<Permit>) -> Box<dyn Receipt>;
113
114 /// Whether this tenant's requests are captured (a retention policy the
115 /// implementation owns). The one pre-receipt check handlers make so an unmarked
116 /// tenant never pays for a prompt copy; the receipt's own [`Receipt::wants_capture`]
117 /// is the post-open gate and the implementation's settle-time re-check stays
118 /// authoritative.
119 fn captures(&self, _tenant: &str) -> bool {
120 false
121 }
122
123 /// Limits-source health for the operator metrics surface, when limits exist.
124 fn limits_health(&self) -> Option<LimitsHealth>;
125
126 /// The graceful-drain deadline expired with requests still in flight: everything
127 /// dropped from this moment on was killed by OUR shutdown, not abandoned by its
128 /// client. Fault attribution (owner ruling 2026-08-23): the implementation must
129 /// settle those drops without billing the caller. Latched — the process is exiting.
130 fn drain_kill(&self) {}
131}
132
133/// One request's accounting record, admission row to terminal row. Method names
134/// deliberately match the reference implementation's inherent methods so the
135/// handler code reads identically through the seam.
136pub trait Receipt: Send {
137 /// Whether this receipt was opened captured. Gates the caller's lazy prompt build;
138 /// `false` makes `arm_capture` a no-op.
139 fn wants_capture(&self) -> bool {
140 false
141 }
142 /// Attach the prompt payload to a captured request. Where it goes and the
143 /// settle-time consent re-check belong to the implementation; the seam never
144 /// names a storage type.
145 fn arm_capture(&mut self, prompt: serde_json::Value);
146 fn capture_completion_delta(&mut self, text: &str);
147 fn record_prompt_usage(
148 &mut self,
149 prompt_tokens: u64,
150 cached_prompt_tokens: u64,
151 ) -> Result<(), String>;
152 fn record_completion_token(&mut self) -> Result<(), String>;
153 fn complete(&mut self, usage: UsageCounts, worker_elapsed_s: f64) -> Result<(), String>;
154 /// Deadline-partial: billed like `complete` but census-distinct.
155 fn complete_deadline_partial(
156 &mut self,
157 usage: UsageCounts,
158 worker_elapsed_s: f64,
159 ) -> Result<(), String>;
160 fn reject(&mut self, status: u16, error_code: &str) -> Result<(), String>;
161 /// Terminal rows with a NAMED zero-debit outcome (`deadline_exceeded`,
162 /// `shed_deadline`, `shed_queue`) — `reject`'s twin for outcomes the census
163 /// distinguishes. Never bills.
164 fn settle_unbilled(
165 &mut self,
166 outcome: &'static str,
167 status: u16,
168 error_code: &str,
169 ) -> Result<(), String>;
170}
171
172/// What a metering factory gets to see at construction time: the server's loaded
173/// model roster. Deliberately small — an implementation brings its own prices,
174/// policies, and storage; the server only vouches for what it serves.
175pub struct MeteringInit<'a> {
176 pub models: &'a [String],
177}
178
179/// Deployment hook: build the metering implementation once models are loaded.
180/// `Ok(None)` = no accounting (the stock no-ledger shape). An `Err` is a startup
181/// FATAL — accounting configuration never fails open.
182pub type MeteringFactory = Box<
183 dyn FnOnce(&MeteringInit<'_>) -> Result<Option<std::sync::Arc<dyn Metering>>, String> + Send,
184>;