submilli_engine/runtime/llm.rs
1//! The embedder-supplied seam for `submilli:llm` — gated model calls from inside
2//! a Submilli program.
3//!
4//! `call`, `batch`, and `models()` all dispatch through the single
5//! embedder-provided [`LlmProvider`]: the actual HTTP, key resolution, retry, and
6//! provider-SDK dependencies live in the embedder (submilli-server), keeping this
7//! crate free of them — the same split [`crate::runtime::mcp::McpTransport`] makes.
8//!
9//! **Two taxonomies, answering different questions.** [`LlmCallError`] is the
10//! *dispatch-level* `Err`: the request never reached the model at all, or the
11//! transport died before any element ran. [`FailureReason`] is *per-element*: the
12//! request reached the model and that one element came back unusable or
13//! incomplete. Neither is a subset of the other, and a per-element failure must
14//! never be raised as an `Err` — a whole-batch `Err` would discard every element
15//! outcome including successes that were already billed, and would leave
16//! per-element budget reconciliation with nothing to reconcile against.
17//!
18//! **No prompt or completion text appears here.** Not in an error `Display`, not
19//! in a log line, not in a permission-check context, and not in an
20//! [`LlmFailure::message`]. The taxonomies carry the model name, the numbers, a
21//! status code, and a fixed classification string — never the payload. Provider
22//! response bodies are dropped at this boundary specifically: they carry
23//! completion fragments, request echoes, and account identifiers, so unlike
24//! [`McpCallError::Upstream`](crate::runtime::mcp::McpCallError::Upstream) no
25//! variant here interpolates one.
26
27use std::future::Future;
28use std::pin::Pin;
29
30pub use super::token_ledger::SharedTokenBudget;
31use super::token_ledger::{LedgerLimits, LedgerRefusal, TokenLedger};
32
33/// The single internal host module every `submilli:llm` call dispatches through.
34pub const LLM_MODULE_NAME: &str = "submilli:llm";
35
36/// Tokens one execution may spend across every model call it makes.
37pub const DEFAULT_MAX_EXECUTION_TOKENS: u64 = 1_000_000;
38/// Tokens every live execution may spend between them.
39pub const DEFAULT_MAX_ALL_EXECUTIONS_TOKENS: u64 = 20_000_000;
40/// Tokens one execution may hold as indeterminate (null-usage) reserve before
41/// further dispatch is refused. See [`LlmLimitKind::IndeterminateSpend`].
42pub const DEFAULT_MAX_HELD_TOKENS: u64 = 200_000;
43/// Output tokens reserved per prompt when the model declares no `outputReserve`.
44pub const DEFAULT_OUTPUT_CAP: u64 = 64_000;
45/// Elements in one `batch` slice.
46pub const DEFAULT_MAX_PROMPT_COUNT: u64 = 128;
47/// UTF-8 bytes in a single prompt element.
48pub const DEFAULT_MAX_PROMPT_BYTES: u64 = 256 * 1024;
49
50/// The ceilings a dispatch is checked against.
51///
52/// Token ceilings and prompt bounds live in one struct deliberately: an operator
53/// raising how much a program may spend almost always means to raise how much it
54/// may send, and splitting them across two config surfaces makes that two edits
55/// with a silent failure between them.
56#[derive(Debug, Clone, Copy, PartialEq, Eq)]
57pub struct LlmLimits {
58 /// Tokens one execution may spend in total.
59 pub per_execution_tokens: u64,
60 /// Tokens one execution may hold as indeterminate reserve before further
61 /// dispatch is refused.
62 pub max_held_tokens: u64,
63 /// Output tokens reserved per prompt when the model declares no
64 /// `outputReserve` of its own.
65 pub default_output_cap: u64,
66 /// Elements in one `batch` slice. Enforced at the stdlib boundary.
67 pub max_prompt_count: u64,
68 /// UTF-8 bytes in a single prompt element. Enforced at the stdlib boundary.
69 pub max_prompt_bytes: u64,
70}
71
72impl Default for LlmLimits {
73 fn default() -> Self {
74 Self {
75 per_execution_tokens: DEFAULT_MAX_EXECUTION_TOKENS,
76 max_held_tokens: DEFAULT_MAX_HELD_TOKENS,
77 default_output_cap: DEFAULT_OUTPUT_CAP,
78 max_prompt_count: DEFAULT_MAX_PROMPT_COUNT,
79 max_prompt_bytes: DEFAULT_MAX_PROMPT_BYTES,
80 }
81 }
82}
83
84/// Which token ceiling a refused dispatch ran into. Carries the numbers, never
85/// the prompts — the same rule
86/// [`SessionKvLimitKind`](crate::runtime::session_kv::SessionKvLimitKind)
87/// follows for values.
88#[derive(Debug, Clone, Copy, PartialEq, Eq)]
89pub enum LlmLimitKind {
90 /// This execution's own token ceiling.
91 PerExecutionTokens { requested: u64, limit: u64 },
92 /// The server-wide budget summed across every live execution. Distinct from
93 /// [`Self::PerExecutionTokens`] because it wants the opposite response: the
94 /// tokens in the way belong to other executions, so this execution making
95 /// fewer or shorter calls need not help.
96 AllExecutionsTokens { requested: u64, limit: u64 },
97 /// Cumulative reserve held for elements whose usage the provider never
98 /// reported. Token spend is monotonic — unlike session-KV bytes, a held
99 /// estimate is never superseded by a real number — so without this bound a
100 /// long-lived execution ratchets toward zero with no recovery. Distinct from
101 /// a genuine overspend so an operator can tell a null-usage storm from real
102 /// consumption.
103 IndeterminateSpend { held: u64, limit: u64 },
104}
105
106impl std::fmt::Display for LlmLimitKind {
107 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
108 match self {
109 Self::PerExecutionTokens { requested, limit } => write!(
110 f,
111 "execution token budget: {requested} tokens exceeds the {limit} this execution \
112 may spend — use fewer or shorter calls, or split the work across executions"
113 ),
114 Self::AllExecutionsTokens { requested, limit } => write!(
115 f,
116 "server token budget: {requested} tokens exceeds the {limit} allowed across all \
117 live executions — this execution's own spend is not what is in the way, so \
118 using fewer or shorter calls need not help; the operator raises the budget \
119 with `--max-llm-tokens`"
120 ),
121 // No CLI rung names this ceiling: `max_held_tokens` is set by the
122 // embedder, unlike the two above. Saying otherwise would send an
123 // operator looking for a flag that does not exist, so the advice is
124 // the action they can actually take.
125 Self::IndeterminateSpend { held, limit } => write!(
126 f,
127 "indeterminate spend ceiling reached: {held} tokens are held for calls the \
128 provider never reported usage for, exceeding the {limit} allowed — this is \
129 unreported spend, not measured consumption; start a new execution, and if it \
130 recurs the provider is not reporting usage and the operator should raise the \
131 held-token ceiling in the server configuration"
132 ),
133 }
134 }
135}
136
137/// Why a model *dispatch* failed — the request never reached the model, or the
138/// transport died before any element ran.
139///
140/// Reserved for dispatch-level failure only. A model that answered and produced
141/// an unusable or incomplete element reports that in [`LlmOutcome`]'s failure
142/// arm, never here: returning `Err` for one bad element would throw away the
143/// sibling elements that did succeed.
144///
145/// Every variant names the model it was dispatched at and the fix the caller can
146/// act on. None carries prompt or completion text.
147#[derive(Debug, Clone, PartialEq, Eq)]
148pub enum LlmCallError {
149 /// No [`LlmProvider`] is wired into the runtime, so nothing can dispatch.
150 /// The pure-interpreter path leaves the provider `None`, and this is the
151 /// catchable configuration error a script sees there.
152 NotConfigured { model: String },
153 /// The model name is not one the provider serves. `available` lists what it
154 /// does serve, so the message can name the alternatives rather than only the
155 /// rejection.
156 UnknownModel {
157 model: String,
158 available: Vec<String>,
159 },
160 /// A token ceiling cannot cover this dispatch. Charged before any element
161 /// ran, so nothing was billed.
162 ///
163 /// `limit_kind` carries which ceiling refused, because the three want
164 /// different responses and a message that reads the same for all of them
165 /// sends a caller to shrink work that was never what filled the budget.
166 BudgetExceeded {
167 model: String,
168 limit_kind: LlmLimitKind,
169 },
170 /// The prompt slice is outside the bounds the host accepts — too many
171 /// elements, or one element too large. Counted in elements and bytes, never
172 /// quoted.
173 PromptBoundsExceeded {
174 model: String,
175 limit_kind: PromptBoundKind,
176 actual: u64,
177 limit: u64,
178 },
179 /// The provider rejected the credential, or none was resolvable when the
180 /// provider was constructed. One variant covers both because the fix is the
181 /// same key — this is the `LoadAPIKeyError`-vs-`401` split collapsed at the
182 /// point where the caller can act on it.
183 Unauthorized { model: String },
184 /// The request never completed, and died before any element ran. A transport
185 /// failure *after* elements started belongs in [`FailureReason::Transport`]
186 /// on the affected elements instead.
187 ///
188 /// `detail` is a **fixed classification string** — "connection reset",
189 /// "dns lookup failed", "tls handshake failed". It is the one free-form field
190 /// in this taxonomy, and the no-payload rule binds it: never a provider response body, and
191 /// never a request echo. The precedent this trait otherwise mirrors does
192 /// interpolate a body
193 /// ([`McpCallError::Upstream`](crate::runtime::mcp::McpCallError::Upstream)
194 /// renders `HTTP {status}: {body}`); this one deliberately does not, because a
195 /// model provider's body carries completion fragments.
196 Transport { model: String, detail: String },
197}
198
199/// Which prompt bound a dispatch ran into. Carries the numbers, never the
200/// prompts — the same rule
201/// [`SessionKvLimitKind`](crate::runtime::session_kv::SessionKvLimitKind) follows
202/// for values.
203#[derive(Debug, Clone, Copy, PartialEq, Eq)]
204pub enum PromptBoundKind {
205 /// Elements in one `call` / `batch` slice.
206 PromptCount,
207 /// UTF-8 bytes of a single prompt element.
208 PromptBytes,
209 /// UTF-8 bytes of the schema sent with every prompt, held to the
210 /// per-prompt bound.
211 SchemaBytes,
212}
213
214impl std::fmt::Display for PromptBoundKind {
215 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
216 match self {
217 Self::PromptCount => write!(f, "prompts in one batch"),
218 Self::PromptBytes => write!(f, "bytes in a single prompt"),
219 Self::SchemaBytes => write!(f, "bytes in the schema"),
220 }
221 }
222}
223
224/// The model the call named (none for a listing), then the failure and its
225/// fix. The function the program called is not known here; the caller that
226/// knows it prefixes it.
227impl std::fmt::Display for LlmCallError {
228 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
229 if !self.model().is_empty() {
230 write!(f, "model \"{}\": ", self.model())?;
231 }
232 match self {
233 Self::NotConfigured { .. } => write!(
234 f,
235 "no model provider is configured for this runtime — the operator wires one \
236 before model calls can dispatch"
237 ),
238 Self::UnknownModel { available, .. } => {
239 write!(f, "the provider does not serve that model — ")?;
240 if available.is_empty() {
241 write!(
242 f,
243 "it serves none, so declare a model in the blueprint before calling"
244 )
245 } else {
246 write!(
247 f,
248 "call llm.models() and use one of: {}",
249 available.join(", ")
250 )
251 }
252 }
253 Self::BudgetExceeded { limit_kind, .. } => write!(f, "exceeded the {limit_kind}"),
254 Self::PromptBoundsExceeded {
255 limit_kind: PromptBoundKind::SchemaBytes,
256 actual,
257 limit,
258 ..
259 } => write!(
260 f,
261 "prompt bound exceeded: {actual} {} exceeds the {limit} allowed — shorten \
262 the schema",
263 PromptBoundKind::SchemaBytes
264 ),
265 Self::PromptBoundsExceeded {
266 limit_kind,
267 actual,
268 limit,
269 ..
270 } => write!(
271 f,
272 "prompt bound exceeded: {actual} {limit_kind} exceeds the {limit} allowed — \
273 send fewer prompts per call, or shorten each one"
274 ),
275 Self::Unauthorized { .. } => write!(
276 f,
277 "the provider rejected the API credential, or none was available — the \
278 operator sets a valid key for this provider"
279 ),
280 Self::Transport { detail, .. } => write!(
281 f,
282 "transport error before any prompt ran: {detail} — no tokens were billed; \
283 retry the call"
284 ),
285 }
286 }
287}
288
289impl std::error::Error for LlmCallError {}
290
291impl LlmCallError {
292 /// The model this dispatch was aimed at. Every variant names one, so a caller
293 /// need not match to attribute the failure.
294 pub fn model(&self) -> &str {
295 match self {
296 Self::NotConfigured { model }
297 | Self::UnknownModel { model, .. }
298 | Self::BudgetExceeded { model, .. }
299 | Self::PromptBoundsExceeded { model, .. }
300 | Self::Unauthorized { model }
301 | Self::Transport { model, .. } => model,
302 }
303 }
304
305 /// Token-budget refusals become catchable `QuotaExceededError`s at the
306 /// stdlib boundary. Prompt size/count bounds remain `RangeError`s.
307 pub fn is_budget_exceeded(&self) -> bool {
308 matches!(self, Self::BudgetExceeded { .. })
309 }
310}
311
312/// Why one *element* of a dispatch is not a clean completion.
313///
314/// Per-element, and orthogonal to [`LlmCallError`]: the request did reach the
315/// model. Several of these still carry usable text on the [`LlmOutcome`] — a
316/// truncated or content-filtered completion is `ok: false` and has a partial
317/// answer worth reading.
318///
319/// A closed set of nine. It is what a guest branches on, so a provider that
320/// meets something it cannot place reports [`Self::Incomplete`] and puts the raw
321/// provider spelling in [`LlmFailure::finish_reason`] rather than widening this.
322#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
323pub enum FailureReason {
324 /// Hit the output-token ceiling mid-answer (`finishReason` of `length`). The
325 /// partial text is retained — that is the whole reason the failure arm has a
326 /// `text` field.
327 Truncated,
328 /// The provider's safety filter stopped or redacted the completion. Reached
329 /// from both directions: `generateText` resolves with a filtered result while
330 /// `generateObject` throws on the same input, and both normalize to here.
331 ContentFiltered,
332 /// The provider answered, but not with something usable: malformed output,
333 /// or output that ignored the requested schema. Distinct from
334 /// [`Self::Truncated`] — nothing was cut off, what arrived was wrong.
335 InvalidOutput,
336 /// The provider throttled this element (structurally, a 429). Honor any
337 /// reported retry-after before re-dispatching.
338 RateLimited,
339 /// The provider refused the request itself rather than failing to answer it:
340 /// a rejected credential, a malformed or oversized request, or a prompt that
341 /// does not fit the model's context window. These share an arm because they
342 /// share a shape — the request as sent will never succeed, so retrying it
343 /// unchanged cannot help.
344 RequestRejected,
345 /// The provider was reachable but could not serve the request (structurally,
346 /// a 5xx). Unlike [`Self::RequestRejected`] the request itself may be fine,
347 /// so the same bytes may succeed later.
348 ProviderUnavailable,
349 /// The connection failed for this element after dispatch had begun — no HTTP
350 /// status was ever observed. Sibling elements may still have succeeded.
351 Transport,
352 /// The element was abandoned before the model finished — deadline, shutdown,
353 /// or an aborted request. Classified separately from [`Self::Transport`]
354 /// because the cause is on this side, not the wire.
355 Cancelled,
356 /// The provider stopped for a reason this runtime does not classify. The raw
357 /// provider spelling travels in [`LlmFailure::finish_reason`] for diagnosis.
358 Incomplete,
359}
360
361impl FailureReason {
362 /// The stable kebab-case spelling. This is the wire form: it is what any
363 /// guest-facing rendering uses, so it must not drift.
364 pub fn as_str(&self) -> &'static str {
365 match self {
366 Self::Truncated => "truncated",
367 Self::ContentFiltered => "content-filtered",
368 Self::InvalidOutput => "invalid-output",
369 Self::RateLimited => "rate-limited",
370 Self::RequestRejected => "request-rejected",
371 Self::ProviderUnavailable => "provider-unavailable",
372 Self::Transport => "transport",
373 Self::Cancelled => "cancelled",
374 Self::Incomplete => "incomplete",
375 }
376 }
377
378 /// A fixed classification string for this reason, suitable as an
379 /// [`LlmFailure::message`] when a provider has nothing more specific that is
380 /// safe to say. Drawn from a closed vocabulary precisely so it cannot carry a
381 /// response body or completion fragment.
382 pub fn default_message(&self) -> &'static str {
383 match self {
384 Self::Truncated => "the model stopped at the output token cap",
385 Self::ContentFiltered => "the provider's safety filter stopped the completion",
386 Self::InvalidOutput => "the model returned output that did not match the request",
387 Self::RateLimited => "the provider throttled this request",
388 Self::RequestRejected => "the provider refused the request as sent",
389 Self::ProviderUnavailable => "the provider could not serve this request",
390 Self::Transport => "the connection failed before a response arrived",
391 Self::Cancelled => "the request was abandoned before the model finished",
392 Self::Incomplete => "the model stopped for an unclassified reason",
393 }
394 }
395
396 /// Whether re-dispatching the identical request could plausibly succeed.
397 ///
398 /// The default a provider starts from, not the last word: it may override
399 /// per element (see [`LlmFailure::retryable`]) when it knows better — a 429
400 /// carrying a retry-after is retryable, while one on an exhausted quota is
401 /// not.
402 pub fn retryable_by_default(&self) -> bool {
403 match self {
404 Self::RateLimited | Self::ProviderUnavailable | Self::Transport => true,
405 // The request as sent will not start succeeding, and a cancellation
406 // was someone's decision — retrying either is a guess, so the guest
407 // opts in rather than the runtime assuming.
408 Self::Truncated
409 | Self::ContentFiltered
410 | Self::InvalidOutput
411 | Self::RequestRejected
412 | Self::Cancelled
413 | Self::Incomplete => false,
414 }
415 }
416
417 /// Every reason, in the plan's declaration order. Kept so a mapper or a
418 /// guest-facing enumeration cannot silently omit one.
419 pub const ALL: [FailureReason; 9] = [
420 Self::Truncated,
421 Self::ContentFiltered,
422 Self::InvalidOutput,
423 Self::RateLimited,
424 Self::RequestRejected,
425 Self::ProviderUnavailable,
426 Self::Transport,
427 Self::Cancelled,
428 Self::Incomplete,
429 ];
430}
431
432impl std::fmt::Display for FailureReason {
433 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
434 f.write_str(self.as_str())
435 }
436}
437
438/// The failure-arm payload of an [`LlmOutcome`]: everything that exists only
439/// when `ok` is false.
440///
441/// Held in an `Option` on the outcome rather than flattened into it, so the
442/// success arm *cannot* carry a `message`, a `retryable`, or a `status` — those
443/// fields are meaningless on a clean completion, and a flattened struct would
444/// have to invent values for them.
445#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
446pub struct LlmFailure {
447 /// The closed-set classification a guest branches on.
448 pub reason: FailureReason,
449 /// The degraded error string.
450 ///
451 /// This is the AI SDK's rule that an `error: unknown` becomes an
452 /// `errorText: string` at any serialization boundary — a guest cannot receive
453 /// a Rust error object, so it receives this. **The no-payload rule binds it absolutely: it is
454 /// a classification, never an echo.** No prompt, no completion fragment, and
455 /// no provider response body may reach it; provider bodies carry request
456 /// echoes and account identifiers, so they are dropped at this boundary.
457 /// [`FailureReason::default_message`] supplies a safe one.
458 pub message: String,
459 /// Whether re-dispatching the identical request could plausibly succeed.
460 /// Defaulted from [`FailureReason::retryable_by_default`] and overridable per
461 /// element, because a provider sometimes knows more than the class does.
462 pub retryable: bool,
463 /// The HTTP status, when one was observed. `None` is meaningful rather than
464 /// missing: a transport death never got a status, which is exactly what
465 /// distinguishes it from a provider that answered with an error.
466 pub status: Option<u16>,
467 /// The provider's own raw stop reason, verbatim, for diagnosis only.
468 ///
469 /// **Deliberately separate from [`Self::reason`].** `reason` is the closed
470 /// set a guest branches on; this is an open string that varies per provider
471 /// and per SDK version. Folding them together would either force the closed
472 /// set to grow with every provider quirk or silently discard the detail that
473 /// makes an unclassified stop diagnosable.
474 pub finish_reason: Option<String>,
475 /// Refused by this host's own configuration or network policy before the provider
476 /// was reached, rather than answered by it. Not part of what a guest sees: the call
477 /// log records such a failure as kind `local`, so a replay does not serve it.
478 pub local: bool,
479}
480
481impl LlmFailure {
482 /// A failure of this class carrying this classification message. `retryable`
483 /// starts from the reason's default; `status` and `finish_reason` start
484 /// absent.
485 pub fn new(reason: FailureReason, message: impl Into<String>) -> Self {
486 Self {
487 reason,
488 message: message.into(),
489 retryable: reason.retryable_by_default(),
490 status: None,
491 finish_reason: None,
492 local: false,
493 }
494 }
495
496 /// Override the class default when the provider knows better.
497 pub fn retryable(mut self, retryable: bool) -> Self {
498 self.retryable = retryable;
499 self
500 }
501
502 /// Record the observed HTTP status. Left absent when none was seen.
503 pub fn with_status(mut self, status: u16) -> Self {
504 self.status = Some(status);
505 self
506 }
507
508 /// Mark the failure as this host's own refusal; see [`Self::local`].
509 pub fn refused_locally(mut self) -> Self {
510 self.local = true;
511 self
512 }
513
514 /// Record the provider's raw stop reason for diagnosis.
515 pub fn with_finish_reason(mut self, finish_reason: impl Into<String>) -> Self {
516 self.finish_reason = Some(finish_reason.into());
517 self
518 }
519}
520
521/// One element's result, positionally matched to the prompt that produced it.
522///
523/// A struct rather than an enum because success and failure are not disjoint
524/// here: a truncated or content-filtered completion is `ok: false` and *still
525/// carries usable text*, and a caller that wants the partial answer should not
526/// have to reach through an error arm to get it.
527///
528/// `ok: true` means the provider reported a **natural stop** — not merely that no
529/// error was raised. A completion that returned cleanly but stopped on `length`
530/// is `ok: false` with [`FailureReason::Truncated`]. That choice trades a class
531/// of silent wrong-answer bugs for one honest data-loss risk, which is why the
532/// failure arm keeps its text rather than discarding it.
533///
534/// The arms differ in more than a flag: `text` is always present on success and
535/// nullable on failure, and the [`LlmFailure`] payload exists only on failure.
536#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
537pub struct LlmOutcome {
538 /// The provider reported a natural stop for this element.
539 pub ok: bool,
540 /// The completion text.
541 ///
542 /// Always `Some` on the success arm. On the failure arm it is `string | null`:
543 /// `Some` — including `Some("")` — when the model produced text, and `None`
544 /// when it produced none at all. The two must not collapse, since telling
545 /// "empty because filtered" from "empty because the model said nothing" is
546 /// the distinction the `ok` rule exists to preserve.
547 pub text: Option<String>,
548 /// The failure payload. `Some` exactly when `!ok`.
549 pub failure: Option<LlmFailure>,
550 /// Tokens the provider reported for the prompt. `None` means
551 /// **indeterminate, not free** — the provider reported nothing, or reported
552 /// something non-finite. Never `Some(0)` as a stand-in for "unknown": a
553 /// throttled call may still have been billed, so a reconciler holds a
554 /// conservative reserve for `None` rather than releasing it.
555 ///
556 /// Nullable on *both* arms: a provider can resolve successfully with no usage
557 /// reported, so this cannot be confined to the failure arm.
558 pub input_tokens: Option<u64>,
559 /// Tokens the provider reported for the completion. Same
560 /// `None`-means-indeterminate rule as [`Self::input_tokens`].
561 pub output_tokens: Option<u64>,
562}
563
564impl LlmOutcome {
565 /// A natural stop with the given text. Usage is left indeterminate; a
566 /// provider that reports counts sets them explicitly.
567 pub fn success(text: impl Into<String>) -> Self {
568 Self {
569 ok: true,
570 text: Some(text.into()),
571 failure: None,
572 input_tokens: None,
573 output_tokens: None,
574 }
575 }
576
577 /// A non-clean element. `text` is whatever the model did produce: `Some` when
578 /// it produced any, `None` when it produced none — retained rather than
579 /// discarded, so a partial answer survives.
580 pub fn failed(failure: LlmFailure, text: Option<impl Into<String>>) -> Self {
581 Self {
582 ok: false,
583 text: text.map(Into::into),
584 failure: Some(failure),
585 input_tokens: None,
586 output_tokens: None,
587 }
588 }
589
590 /// Attach reported usage. A non-finite or unreported count stays `None`.
591 pub fn with_usage(mut self, input_tokens: Option<u64>, output_tokens: Option<u64>) -> Self {
592 self.input_tokens = input_tokens;
593 self.output_tokens = output_tokens;
594 self
595 }
596}
597
598/// A model the provider serves.
599///
600/// Every optional may legitimately be `None`: models are blueprint-declared, and
601/// an operator may declare a name without a description, a context window, or an
602/// output reserve. Absent means unknown, not zero.
603#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
604pub struct LlmModel {
605 pub name: String,
606 pub description: Option<String>,
607 pub context_window: Option<u64>,
608}
609
610impl LlmModel {
611 pub fn new(name: impl Into<String>) -> Self {
612 Self {
613 name: name.into(),
614 description: None,
615 context_window: None,
616 }
617 }
618}
619
620/// The embedder-provided model provider: resolves the credential, performs the
621/// completions, and reports what each one cost.
622///
623/// Object-safe boxed-future trait (rather than `async fn`) so the interpreter
624/// stays free of the embedder's async stack while still `await`-ing the call from
625/// within the async host function — the shape
626/// [`McpTransport`](crate::runtime::mcp::McpTransport) uses.
627///
628/// # Obligations on every implementor
629///
630/// These are *contractual*, not enforced by the interpreter. A second
631/// implementor inherits them:
632///
633/// 1. **Positional ordering.** The returned `Vec<LlmOutcome>` has exactly one
634/// element per input prompt, in input order. Element `i` is the outcome of
635/// `prompts[i]`, including when it failed.
636/// 2. **Bounded concurrency.** `call` takes a *slice* rather than a single
637/// prompt because the guest is single-threaded Wasm holding one host call at a
638/// time: fan-out cannot live above this boundary, so it lives below it. The
639/// implementor is responsible for bounding that concurrency — the interpreter
640/// does not limit it.
641/// 3. **Per-element failure never becomes `Err`.** One bad element is an
642/// `LlmOutcome` with `ok: false`. `Err(LlmCallError)` is only for dispatch
643/// that never ran, or a transport that died before any element started; an
644/// `Err` discards the siblings that did succeed and were billed.
645/// 4. **No prompt or completion text in errors.** Nothing an implementor puts in
646/// an [`LlmCallError`] may quote the payload.
647pub trait LlmProvider: Send + Sync {
648 fn call<'a>(
649 &'a self,
650 model: &'a str,
651 prompts: &'a [String],
652 schema_json: Option<&'a str>,
653 ) -> Pin<Box<dyn Future<Output = Result<Vec<LlmOutcome>, LlmCallError>> + Send + 'a>>;
654
655 /// Returns the locally declared model catalog without network I/O or other
656 /// side effects. The caller must inspect these names before it can check
657 /// policy for each model, so this runs before per-model authorization.
658 fn models<'a>(
659 &'a self,
660 ) -> Pin<Box<dyn Future<Output = Result<Vec<LlmModel>, LlmCallError>> + Send + 'a>>;
661
662 /// The output cap declared for one model, consulted *before* dispatch.
663 ///
664 /// Separate from [`Self::models`] because the reservation needs one model's
665 /// reserve on the hot path of every `call`; using the whole catalog here
666 /// would add unnecessary work to each dispatch. An implementor that already
667 /// holds its declarations (the blueprint case) answers from memory, and one
668 /// that does not should return `None` rather than block.
669 ///
670 /// `None` means the model declared no reserve, or the implementor cannot
671 /// answer cheaply; the caller then applies
672 /// [`LlmLimits::default_output_cap`]. Returning `None` is always safe, which
673 /// is why this defaults rather than being required of every implementor.
674 fn output_reserve(&self, _model: &str) -> Option<u64> {
675 None
676 }
677}
678
679/// One execution's token budget, reserving against its own ceiling and the
680/// server-wide one together.
681///
682/// Both reservations are taken together and both unwound if either fails, so a
683/// refused dispatch leaves neither counter charged.
684pub struct ExecutionTokenBudget {
685 limits: LlmLimits,
686 ledger: TokenLedger,
687}
688
689impl ExecutionTokenBudget {
690 pub fn new(limits: LlmLimits, aggregate: SharedTokenBudget) -> Self {
691 let ledger = TokenLedger::new(
692 LedgerLimits {
693 per_execution_tokens: limits.per_execution_tokens,
694 max_held_tokens: limits.max_held_tokens,
695 },
696 aggregate,
697 );
698 Self { limits, ledger }
699 }
700
701 pub fn limits(&self) -> LlmLimits {
702 self.limits
703 }
704
705 /// Tokens this execution currently holds against its ceiling.
706 pub fn used(&self) -> u64 {
707 self.ledger.used()
708 }
709
710 /// The part of [`Self::used`] held for unreported usage.
711 pub fn held(&self) -> u64 {
712 self.ledger.held()
713 }
714 /// Apply committed per-run spend during replay, without a provider reservation.
715 pub fn restore_spend(&self, used: u64, held: u64) {
716 self.ledger.restore_spend(used, held);
717 }
718
719 /// Reserve `tokens` against this execution and the aggregate, before
720 /// dispatch.
721 ///
722 /// The caller passes the full reservation — `estimated_input_tokens +
723 /// (output_cap × prompt_count)`, via [`Self::reservation_for`] — because
724 /// reserving only input enforces the ceiling retroactively: output tokens
725 /// are typically the expensive half, and a small prompt can legitimately
726 /// produce a very large completion with nothing standing in the way.
727 pub fn reserve(&self, model: &str, tokens: u64) -> Result<(), LlmCallError> {
728 self.ledger.reserve(tokens).map_err(|refusal| {
729 let limit_kind = match refusal {
730 LedgerRefusal::HeldReserve { held, limit } => {
731 LlmLimitKind::IndeterminateSpend { held, limit }
732 }
733 LedgerRefusal::PerExecution { requested, limit } => {
734 LlmLimitKind::PerExecutionTokens { requested, limit }
735 }
736 LedgerRefusal::AllExecutions { requested, limit } => {
737 LlmLimitKind::AllExecutionsTokens { requested, limit }
738 }
739 };
740 LlmCallError::BudgetExceeded {
741 model: model.to_string(),
742 limit_kind,
743 }
744 })
745 }
746
747 /// The reservation a dispatch of `prompt_count` prompts needs:
748 /// `estimated_input_tokens + (output_cap × prompt_count)`.
749 ///
750 /// `output_cap` is the resolved per-model `outputReserve` when the model
751 /// declares one, and [`LlmLimits::default_output_cap`] otherwise. The same
752 /// cap is passed to the provider as the request's `maxOutputTokens`, which
753 /// makes the reservation an actual upper bound rather than an estimate.
754 pub fn reservation_for(
755 &self,
756 estimated_input_tokens: u64,
757 prompt_count: u64,
758 output_cap: Option<u64>,
759 ) -> u64 {
760 let cap = output_cap.unwrap_or(self.limits.default_output_cap);
761 estimated_input_tokens.saturating_add(cap.saturating_mul(prompt_count))
762 }
763
764 /// Reconcile a completed dispatch down from its `reserved` estimate to what
765 /// the provider actually reported.
766 ///
767 /// `reported` is the summed usage of the elements the provider gave counts
768 /// for; `indeterminate` is the reserve belonging to elements it reported
769 /// nothing for. The indeterminate part is **not** released — `None` means
770 /// indeterminate, not free, since a throttled call may still have been
771 /// billed — and is instead carried as held reserve, bounded by
772 /// [`LlmLimits::max_held_tokens`].
773 pub fn reconcile(&self, reserved: u64, reported: u64, indeterminate: u64) {
774 self.ledger.reconcile(reserved, reported, indeterminate);
775 }
776
777 /// Release a reservation for prompts that never dispatched — the remainder
778 /// of a batch abandoned partway.
779 pub fn release(&self, tokens: u64) {
780 self.reconcile(tokens, 0, 0);
781 }
782}
783
784#[cfg(test)]
785mod tests {
786 use std::sync::Arc;
787 use std::sync::atomic::{AtomicU64, Ordering};
788
789 use super::{
790 ExecutionTokenBudget, FailureReason, LlmCallError, LlmFailure, LlmLimitKind, LlmLimits,
791 LlmModel, LlmOutcome, LlmProvider, PromptBoundKind, SharedTokenBudget,
792 };
793
794 const MODEL: &str = "claude-sonnet-4";
795
796 /// An execution ends by having its budget holder dropped, not by reconciling
797 /// every reservation first. Without the release on drop the reservation of
798 /// every ended execution would be held forever and the server-wide budget
799 /// would ratchet to its cap, refusing every later execution.
800 #[test]
801 fn dropping_an_execution_budget_returns_its_tokens_to_the_aggregate() {
802 let aggregate = SharedTokenBudget::new(1_000);
803 {
804 let held = ExecutionTokenBudget::new(limits_with_cap(1_000), aggregate.clone());
805 held.reserve(MODEL, 1_000)
806 .expect("the first execution fits");
807 assert_eq!(aggregate.used(), 1_000);
808 }
809 assert_eq!(
810 aggregate.used(),
811 0,
812 "the ended execution released its tokens"
813 );
814
815 let next = ExecutionTokenBudget::new(limits_with_cap(1_000), aggregate);
816 next.reserve(MODEL, 1_000)
817 .expect("a later execution reuses the capacity");
818 }
819
820 fn limits_with_cap(per_execution_tokens: u64) -> LlmLimits {
821 LlmLimits {
822 per_execution_tokens,
823 ..LlmLimits::default()
824 }
825 }
826
827 /// A panicking execution must still return its tokens: an unrelated panic
828 /// while the lock was held would otherwise strand the whole reservation, and
829 /// the aggregate would ratchet exactly as it does without any release.
830 #[test]
831 fn a_poisoned_lock_still_releases_on_drop() {
832 let aggregate = SharedTokenBudget::new(1_000);
833 {
834 let held = ExecutionTokenBudget::new(limits_with_cap(1_000), aggregate.clone());
835 held.reserve(MODEL, 600).expect("reservation fits");
836
837 let poisoner = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
838 let _guard = held.ledger.lock_state().expect("lock");
839 panic!("poison the budget's lock");
840 }));
841 assert!(poisoner.is_err(), "the panic must have unwound");
842 assert!(held.ledger.state_is_poisoned(), "the lock must be poisoned");
843 }
844 assert_eq!(
845 aggregate.used(),
846 0,
847 "a poisoned execution still returns its tokens"
848 );
849 }
850
851 /// Two executions reserving at once must not both observe the same headroom.
852 /// With `fetch_add`-and-compensate both can read the pre-add value and both
853 /// conclude they fit; the CAS loop makes exactly one win.
854 #[test]
855 fn concurrent_reservations_cannot_both_claim_the_same_headroom() {
856 // Many rounds with many threads: a race that resolves correctly once may
857 // only have been serialized by luck, and two threads rarely collide
858 // inside the few instructions between the load and the store.
859 const THREADS: u64 = 16;
860 const ROUNDS: usize = 500;
861
862 let mut overspends = 0;
863 for _ in 0..ROUNDS {
864 // Headroom for exactly one of the racing reservations.
865 let aggregate = SharedTokenBudget::new(1_000);
866 let start = Arc::new(std::sync::Barrier::new(THREADS as usize));
867 let winners = Arc::new(AtomicU64::new(0));
868
869 let threads: Vec<_> = (0..THREADS)
870 .map(|_| {
871 let aggregate = aggregate.clone();
872 let start = Arc::clone(&start);
873 let winners = Arc::clone(&winners);
874 std::thread::spawn(move || {
875 let budget =
876 ExecutionTokenBudget::new(limits_with_cap(1_000), aggregate.clone());
877 // Release every thread into the reservation together.
878 start.wait();
879 if budget.reserve(MODEL, 1_000).is_ok() {
880 winners.fetch_add(1, Ordering::Relaxed);
881 }
882 // Hold every reservation until the round is over, so a
883 // win is never handed back before it is contested.
884 std::mem::forget(budget);
885 })
886 })
887 .collect();
888 for t in threads {
889 t.join().expect("reserving thread");
890 }
891
892 let winners = winners.load(Ordering::Relaxed);
893 assert_eq!(
894 winners, 1,
895 "exactly one of {THREADS} racing reservations may claim the only headroom"
896 );
897 if aggregate.used() > aggregate.cap() {
898 overspends += 1;
899 }
900 }
901 assert_eq!(
902 overspends, 0,
903 "the aggregate was overdrawn in {overspends}/{ROUNDS} rounds"
904 );
905 }
906
907 /// The two token ceilings refuse for opposite reasons and want opposite
908 /// responses: make fewer or shorter calls, versus wait or ask an operator. A
909 /// message that reads the same for both sends a program — or the model
910 /// writing it — to shrink work that was never what filled the budget.
911 #[test]
912 fn the_per_execution_and_server_wide_token_limits_name_different_fixes() {
913 let roomy = SharedTokenBudget::new(u64::MAX);
914 let own = ExecutionTokenBudget::new(limits_with_cap(100), roomy)
915 .reserve(MODEL, 500)
916 .expect_err("over this execution's own ceiling")
917 .to_string();
918
919 // An aggregate with no headroom left by another execution, so this one is
920 // refused without ever approaching its own per-execution ceiling.
921 let aggregate = SharedTokenBudget::new(100);
922 let neighbour = ExecutionTokenBudget::new(limits_with_cap(u64::MAX), aggregate.clone());
923 neighbour.reserve(MODEL, 100).expect("neighbour fits");
924 let shared = ExecutionTokenBudget::new(limits_with_cap(u64::MAX), aggregate);
925 let server_wide = shared
926 .reserve(MODEL, 100)
927 .expect_err("over the server-wide budget")
928 .to_string();
929
930 assert_ne!(own, server_wide);
931 assert!(own.contains("this execution may spend"), "{own}");
932 assert!(own.contains("use fewer or shorter calls"), "{own}");
933 assert!(
934 !own.contains("--max-llm-tokens"),
935 "the per-execution ceiling is not the operator's flag: {own}"
936 );
937 assert!(
938 server_wide.contains("across all live executions"),
939 "{server_wide}"
940 );
941 assert!(
942 server_wide.contains("is not what is in the way"),
943 "{server_wide}"
944 );
945 assert!(
946 server_wide.contains("--max-llm-tokens"),
947 "the aggregate refusal must name the flag that raises it: {server_wide}"
948 );
949 }
950
951 /// Each refusal must name the limit it hit and the ceiling, so an operator
952 /// reading a log knows which knob is in play.
953 #[test]
954 fn exceeding_either_ceiling_names_the_limit_and_the_ceiling_hit() {
955 let aggregate = SharedTokenBudget::new(u64::MAX);
956 let budget = ExecutionTokenBudget::new(limits_with_cap(64), aggregate);
957 let err = budget.reserve(MODEL, 4_096).expect_err("over the ceiling");
958 assert!(
959 matches!(
960 err,
961 LlmCallError::BudgetExceeded {
962 limit_kind: LlmLimitKind::PerExecutionTokens { .. },
963 ..
964 }
965 ),
966 "{err:?}"
967 );
968 let rendered = err.to_string();
969 assert!(rendered.contains(MODEL), "{rendered}");
970 assert!(rendered.contains("4096"), "the request: {rendered}");
971 assert!(rendered.contains("64"), "the ceiling: {rendered}");
972 }
973
974 /// a budget refusal carries the numbers, never the payload.
975 #[test]
976 fn no_budget_error_display_leaks_prompt_text() {
977 const SECRET: &str = "the patient's diagnosis is";
978 let aggregate = SharedTokenBudget::new(10);
979 let neighbour = ExecutionTokenBudget::new(limits_with_cap(u64::MAX), aggregate.clone());
980 neighbour.reserve(MODEL, 10).expect("neighbour fits");
981
982 let rendered: Vec<String> = [
983 ExecutionTokenBudget::new(limits_with_cap(1), SharedTokenBudget::new(u64::MAX))
984 .reserve(MODEL, 900)
985 .expect_err("per-execution"),
986 ExecutionTokenBudget::new(limits_with_cap(u64::MAX), aggregate)
987 .reserve(MODEL, 900)
988 .expect_err("server-wide"),
989 held_budget_at_its_ceiling()
990 .reserve(MODEL, 1)
991 .expect_err("indeterminate"),
992 ]
993 .iter()
994 .map(ToString::to_string)
995 .collect();
996
997 for message in &rendered {
998 assert!(!message.contains(SECRET), "{message}");
999 assert!(!message.contains("patient"), "{message}");
1000 }
1001 }
1002
1003 /// An execution whose held reserve is already over the bound.
1004 fn held_budget_at_its_ceiling() -> ExecutionTokenBudget {
1005 let budget = ExecutionTokenBudget::new(
1006 LlmLimits {
1007 per_execution_tokens: u64::MAX,
1008 max_held_tokens: 100,
1009 ..LlmLimits::default()
1010 },
1011 SharedTokenBudget::new(u64::MAX),
1012 );
1013 budget.reserve(MODEL, 500).expect("the dispatch fits");
1014 // Every element came back with no usage reported, so the whole
1015 // reservation is carried as held rather than released.
1016 budget.reconcile(500, 0, 500);
1017 budget
1018 }
1019
1020 /// A batch abandoned partway must not keep charging for the prompts that
1021 /// never dispatched.
1022 #[test]
1023 fn an_undispatched_remainder_is_released_after_a_partial_batch() {
1024 let aggregate = SharedTokenBudget::new(1_000);
1025 let budget = ExecutionTokenBudget::new(limits_with_cap(1_000), aggregate.clone());
1026 budget.reserve(MODEL, 900).expect("the whole batch fits");
1027
1028 // Four of ten prompts ran and reported 120 tokens; the rest never
1029 // dispatched, so their share of the reservation is released.
1030 budget.reconcile(900, 120, 0);
1031
1032 assert_eq!(budget.used(), 120);
1033 assert_eq!(aggregate.used(), 120, "the remainder returned");
1034 budget
1035 .reserve(MODEL, 800)
1036 .expect("the released capacity is reusable");
1037 }
1038
1039 /// reserving only input enforces the ceiling retroactively. A
1040 /// 100-token prompt can legitimately produce 100k output tokens, so the
1041 /// output cap must be part of the reservation and the refusal must land
1042 /// before dispatch, not after the provider has billed.
1043 #[test]
1044 fn the_reservation_covers_output_tokens_and_refuses_before_dispatch() {
1045 let aggregate = SharedTokenBudget::new(u64::MAX);
1046 let budget = ExecutionTokenBudget::new(
1047 LlmLimits {
1048 per_execution_tokens: 10_000,
1049 default_output_cap: 64_000,
1050 ..LlmLimits::default()
1051 },
1052 aggregate.clone(),
1053 );
1054
1055 // Input alone is trivially under the ceiling.
1056 let input_only = 100;
1057 assert!(input_only < budget.limits().per_execution_tokens);
1058
1059 let reservation = budget.reservation_for(input_only, 1, None);
1060 assert_eq!(
1061 reservation,
1062 100 + 64_000,
1063 "output cap is in the reservation"
1064 );
1065
1066 let err = budget
1067 .reserve(MODEL, reservation)
1068 .expect_err("refused before dispatch, not after the provider billed");
1069 assert!(
1070 matches!(
1071 err,
1072 LlmCallError::BudgetExceeded {
1073 limit_kind: LlmLimitKind::PerExecutionTokens { .. },
1074 ..
1075 }
1076 ),
1077 "{err:?}"
1078 );
1079 assert_eq!(budget.used(), 0, "a refused dispatch charges nothing");
1080 assert_eq!(aggregate.used(), 0);
1081
1082 // The cap scales with the prompt count, and a model's own reserve wins
1083 // over the default.
1084 assert_eq!(budget.reservation_for(100, 4, None), 100 + 64_000 * 4);
1085 assert_eq!(budget.reservation_for(100, 4, Some(1_000)), 100 + 4_000);
1086 }
1087
1088 /// held reserve releases from the per-execution budget on teardown —
1089 /// that execution is over — but stays charged against the aggregate.
1090 /// Releasing it from the aggregate would let indeterminate spend escape the
1091 /// server-wide ceiling: N executions each under the per-execution ceiling
1092 /// could exceed the server-wide one without ever tripping it.
1093 #[test]
1094 fn a_held_reserve_commits_against_the_aggregate_and_releases_only_per_execution() {
1095 let aggregate = SharedTokenBudget::new(1_500);
1096 let limits = LlmLimits {
1097 per_execution_tokens: 400,
1098 max_held_tokens: 10_000,
1099 ..LlmLimits::default()
1100 };
1101
1102 // Three executions that each spend 400 tokens the provider never reports.
1103 for _ in 0..3 {
1104 let budget = ExecutionTokenBudget::new(limits, aggregate.clone());
1105 budget.reserve(MODEL, 400).expect("within its own ceiling");
1106 budget.reconcile(400, 0, 400);
1107 assert_eq!(budget.held(), 400);
1108 }
1109
1110 assert_eq!(
1111 aggregate.used(),
1112 1_200,
1113 "indeterminate spend stays charged against the aggregate after teardown"
1114 );
1115
1116 // A fourth execution, itself well under the per-execution ceiling, is
1117 // now refused by the server-wide ceiling it would otherwise have walked
1118 // straight past.
1119 let next = ExecutionTokenBudget::new(limits, aggregate);
1120 let err = next.reserve(MODEL, 400).expect_err("the aggregate refuses");
1121 assert!(
1122 matches!(
1123 err,
1124 LlmCallError::BudgetExceeded {
1125 limit_kind: LlmLimitKind::AllExecutionsTokens { .. },
1126 ..
1127 }
1128 ),
1129 "{err:?}"
1130 );
1131 }
1132
1133 /// Token spend is monotonic, so a held estimate is never superseded by a real
1134 /// number and a long-lived execution accumulating null usage would ratchet to
1135 /// zero with no recovery. The bound names *indeterminate spend*, distinctly
1136 /// from a genuine overspend, so an operator can tell a null-usage storm from
1137 /// real consumption.
1138 #[test]
1139 fn crossing_the_held_reserve_bound_is_distinguishable_from_a_genuine_overspend() {
1140 let held = held_budget_at_its_ceiling();
1141 let indeterminate = held
1142 .reserve(MODEL, 1)
1143 .expect_err("the held bound is crossed");
1144 assert!(
1145 matches!(
1146 indeterminate,
1147 LlmCallError::BudgetExceeded {
1148 limit_kind: LlmLimitKind::IndeterminateSpend { .. },
1149 ..
1150 }
1151 ),
1152 "{indeterminate:?}"
1153 );
1154
1155 let overspend =
1156 ExecutionTokenBudget::new(limits_with_cap(10), SharedTokenBudget::new(1_000))
1157 .reserve(MODEL, 900)
1158 .expect_err("a genuine overspend");
1159
1160 let indeterminate = indeterminate.to_string();
1161 let overspend = overspend.to_string();
1162 assert_ne!(indeterminate, overspend);
1163 assert!(
1164 indeterminate.contains("indeterminate spend ceiling reached"),
1165 "{indeterminate}"
1166 );
1167 assert!(
1168 indeterminate.contains("never reported usage"),
1169 "the operator must be able to tell unreported spend from measured: {indeterminate}"
1170 );
1171 assert!(
1172 !overspend.contains("indeterminate"),
1173 "a real overspend must not read as a null-usage storm: {overspend}"
1174 );
1175 }
1176
1177 /// A ceiling is something a program can catch and adapt to, so it reaches the
1178 /// guest as a catchable `QuotaExceededError`. A dispatch that failed for
1179 /// a reason retrying smaller cannot fix
1180 /// must not masquerade as one.
1181 #[test]
1182 fn only_quota_refusals_are_classified_as_quota_exceeded_errors() {
1183 for error in every_variant() {
1184 let expected = matches!(error, LlmCallError::BudgetExceeded { .. });
1185 assert_eq!(
1186 error.is_budget_exceeded(),
1187 expected,
1188 "{error:?} is misclassified for the guest boundary"
1189 );
1190 }
1191
1192 // Every token ceiling, not just the one in `every_variant`.
1193 for limit_kind in [
1194 LlmLimitKind::PerExecutionTokens {
1195 requested: 2,
1196 limit: 1,
1197 },
1198 LlmLimitKind::AllExecutionsTokens {
1199 requested: 2,
1200 limit: 1,
1201 },
1202 LlmLimitKind::IndeterminateSpend { held: 2, limit: 1 },
1203 ] {
1204 let error = LlmCallError::BudgetExceeded {
1205 model: MODEL.to_string(),
1206 limit_kind,
1207 };
1208 assert!(error.is_budget_exceeded(), "{error:?}");
1209 }
1210 }
1211
1212 /// The prompt bounds U4 enforces live alongside the token ceilings so an
1213 /// operator raising one raises the other in one place.
1214 #[test]
1215 fn default_limits_carry_both_token_ceilings_and_prompt_bounds() {
1216 let limits = LlmLimits::default();
1217 assert_eq!(limits.max_prompt_count, 128);
1218 assert_eq!(limits.max_prompt_bytes, 256 * 1024);
1219 assert_eq!(
1220 limits.per_execution_tokens,
1221 super::DEFAULT_MAX_EXECUTION_TOKENS
1222 );
1223 assert_eq!(limits.default_output_cap, super::DEFAULT_OUTPUT_CAP);
1224 }
1225
1226 /// One of every dispatch-level variant, so a new variant fails to compile
1227 /// here until it is added.
1228 fn every_variant() -> Vec<LlmCallError> {
1229 vec![
1230 LlmCallError::NotConfigured {
1231 model: MODEL.to_string(),
1232 },
1233 LlmCallError::UnknownModel {
1234 model: MODEL.to_string(),
1235 available: vec!["claude-haiku-4".to_string()],
1236 },
1237 LlmCallError::BudgetExceeded {
1238 model: MODEL.to_string(),
1239 limit_kind: LlmLimitKind::PerExecutionTokens {
1240 requested: 4_000,
1241 limit: 120,
1242 },
1243 },
1244 LlmCallError::PromptBoundsExceeded {
1245 model: MODEL.to_string(),
1246 limit_kind: PromptBoundKind::PromptCount,
1247 actual: 64,
1248 limit: 32,
1249 },
1250 LlmCallError::Unauthorized {
1251 model: MODEL.to_string(),
1252 },
1253 LlmCallError::Transport {
1254 model: MODEL.to_string(),
1255 detail: "connection reset".to_string(),
1256 },
1257 ]
1258 }
1259
1260 #[test]
1261 fn six_dispatch_variants_each_name_the_model_and_a_distinct_fix() {
1262 let variants = every_variant();
1263 assert_eq!(variants.len(), 6, "LlmCallError has six dispatch variants");
1264
1265 let mut fixes: Vec<String> = Vec::new();
1266 for err in &variants {
1267 let rendered = err.to_string();
1268 assert!(
1269 rendered.contains(MODEL),
1270 "variant must name the model: {rendered}"
1271 );
1272 assert_eq!(err.model(), MODEL);
1273 let (_, fix) = rendered.split_once(" — ").unwrap_or_else(|| {
1274 panic!("variant must suggest a fix after an em dash: {rendered}")
1275 });
1276 assert!(!fix.trim().is_empty(), "empty fix: {rendered}");
1277 fixes.push(fix.to_string());
1278 }
1279
1280 // Distinct fixes: a shared message would tell the caller nothing about
1281 // which of the six they hit.
1282 let mut sorted = fixes.clone();
1283 sorted.sort();
1284 sorted.dedup();
1285 assert_eq!(
1286 sorted.len(),
1287 fixes.len(),
1288 "each variant needs its own fix: {fixes:?}"
1289 );
1290 }
1291
1292 /// no prompt or completion text may reach an error `Display`.
1293 #[test]
1294 fn no_dispatch_error_display_leaks_prompt_text() {
1295 const SECRET_PROMPT: &str = "patient SSN is 000-00-0000";
1296 const SECRET_COMPLETION: &str = "the diagnosis is confidential";
1297
1298 for err in every_variant() {
1299 let rendered = err.to_string();
1300 assert!(
1301 !rendered.contains(SECRET_PROMPT),
1302 "prompt text leaked into Display: {rendered}"
1303 );
1304 assert!(
1305 !rendered.contains(SECRET_COMPLETION),
1306 "completion text leaked into Display: {rendered}"
1307 );
1308 }
1309
1310 // The only free-form field any variant carries is `Transport { detail }`,
1311 // which is a transport diagnostic and never built from a payload. Prove
1312 // no *other* variant has a construction path that could carry one by
1313 // showing the rendered set is otherwise fully determined by model name
1314 // and numbers.
1315 let bounded = LlmCallError::PromptBoundsExceeded {
1316 model: MODEL.to_string(),
1317 limit_kind: PromptBoundKind::PromptBytes,
1318 actual: 900_000,
1319 limit: 65_536,
1320 };
1321 let rendered = bounded.to_string();
1322 assert!(rendered.contains("900000") && rendered.contains("65536"));
1323 assert!(
1324 rendered.contains("bytes in a single prompt"),
1325 "bound kind must be named: {rendered}"
1326 );
1327 }
1328
1329 /// nine, in the plan's declaration order. The spellings are the wire
1330 /// form a guest branches on, so drift here is a guest-visible break.
1331 #[test]
1332 fn nine_failure_reasons_have_stable_kebab_case_spellings() {
1333 let spellings: Vec<&str> = FailureReason::ALL
1334 .iter()
1335 .map(FailureReason::as_str)
1336 .collect();
1337 assert_eq!(
1338 spellings,
1339 [
1340 "truncated",
1341 "content-filtered",
1342 "invalid-output",
1343 "rate-limited",
1344 "request-rejected",
1345 "provider-unavailable",
1346 "transport",
1347 "cancelled",
1348 "incomplete",
1349 ]
1350 );
1351 for reason in FailureReason::ALL {
1352 assert_eq!(reason.to_string(), reason.as_str());
1353 }
1354 }
1355
1356 /// `ok` keys off a natural stop, not off "didn't throw"; absent usage is
1357 /// `None`, never `Some(0)`; and the failure arm keeps partial text.
1358 #[test]
1359 fn truncated_outcome_is_not_ok_but_keeps_its_text_and_indeterminate_usage() {
1360 let partial = LlmOutcome::failed(
1361 LlmFailure::new(FailureReason::Truncated, "output token ceiling reached"),
1362 Some("the answer begins"),
1363 );
1364 assert!(!partial.ok, "a length stop is not a natural stop");
1365 assert_eq!(
1366 partial.text.as_deref(),
1367 Some("the answer begins"),
1368 "partial text is retained"
1369 );
1370 let failure = partial.failure.as_ref().expect("failure arm present");
1371 assert_eq!(failure.reason, FailureReason::Truncated);
1372 assert_eq!(
1373 partial.input_tokens, None,
1374 "unreported usage is indeterminate"
1375 );
1376 assert_eq!(partial.output_tokens, None);
1377
1378 let clean = LlmOutcome::success("done").with_usage(Some(12), Some(0));
1379 assert!(clean.ok);
1380 assert!(clean.failure.is_none(), "a natural stop carries no reason");
1381 assert_eq!(
1382 clean.output_tokens,
1383 Some(0),
1384 "a real zero is Some(0); only absence is None"
1385 );
1386 }
1387
1388 /// `text` is `string | null` on the failure arm specifically — a model
1389 /// that produced nothing is distinguishable from one that produced "".
1390 #[test]
1391 fn failure_text_is_nullable_and_distinct_from_empty() {
1392 let nothing = LlmOutcome::failed(
1393 LlmFailure::new(FailureReason::ContentFiltered, "filtered before any output"),
1394 None::<String>,
1395 );
1396 assert_eq!(nothing.text, None, "produced nothing at all");
1397
1398 let empty = LlmOutcome::failed(
1399 LlmFailure::new(
1400 FailureReason::ContentFiltered,
1401 "filtered after an empty turn",
1402 ),
1403 Some(""),
1404 );
1405 assert_eq!(empty.text.as_deref(), Some(""), "produced an empty string");
1406 assert_ne!(
1407 nothing.text, empty.text,
1408 "null and empty must not collapse — the whole point is telling \
1409 'empty because filtered' from 'empty because it said nothing'"
1410 );
1411 }
1412
1413 /// the failure arm carries `message`, `retryable`, and `status`, and the
1414 /// success arm structurally cannot.
1415 #[test]
1416 fn failure_arm_carries_message_retryable_and_nullable_status() {
1417 let throttled =
1418 LlmFailure::new(FailureReason::RateLimited, "provider throttled the request")
1419 .retryable(true)
1420 .with_status(429);
1421 assert_eq!(throttled.message, "provider throttled the request");
1422 assert!(throttled.retryable);
1423 assert_eq!(throttled.status, Some(429));
1424
1425 // `status: number | null` — a transport death has no HTTP status at all.
1426 // That absence is the classification: it is exactly what distinguishes a
1427 // dead connection from a provider that answered with an error.
1428 let dead = LlmFailure::new(
1429 FailureReason::Transport,
1430 "connection closed before response",
1431 );
1432 assert_eq!(dead.status, None);
1433
1434 // `retryable` defaults per class rather than uniformly: a dead connection
1435 // may succeed on the same bytes, while a refused request will not.
1436 assert!(dead.retryable, "transport is retryable by class");
1437 let refused = LlmFailure::new(FailureReason::RequestRejected, "request refused as sent");
1438 assert!(
1439 !refused.retryable,
1440 "a request the provider refused will not start succeeding unchanged"
1441 );
1442 // ...and the class default is a starting point, not the last word.
1443 assert!(
1444 !LlmFailure::new(FailureReason::RateLimited, "quota exhausted")
1445 .retryable(false)
1446 .retryable,
1447 "a provider that knows better overrides the class default"
1448 );
1449
1450 // Raw `finishReason` is a *diagnostic* field, separate from `reason`,
1451 // which is what guests branch on. They must not fold together.
1452 let truncated = LlmFailure::new(FailureReason::Truncated, "stopped at the output cap")
1453 .with_finish_reason("length");
1454 assert_eq!(truncated.reason, FailureReason::Truncated);
1455 assert_eq!(truncated.finish_reason.as_deref(), Some("length"));
1456
1457 // A provider reporting a stop reason we do not classify still surfaces the
1458 // raw string for diagnosis while `reason` stays in the closed set.
1459 let odd = LlmFailure::new(FailureReason::Incomplete, "unclassified stop")
1460 .with_finish_reason("surprise-reason");
1461 assert_eq!(odd.reason, FailureReason::Incomplete);
1462 assert_eq!(odd.finish_reason.as_deref(), Some("surprise-reason"));
1463 }
1464
1465 /// The retryable split follows the spike's structural classification: 429 and
1466 /// 5xx and a dead connection may succeed on the same bytes; everything else
1467 /// describes a request that will not start succeeding unchanged.
1468 #[test]
1469 fn retryable_defaults_split_transient_classes_from_terminal_ones() {
1470 let retryable: Vec<&str> = FailureReason::ALL
1471 .iter()
1472 .filter(|r| r.retryable_by_default())
1473 .map(FailureReason::as_str)
1474 .collect();
1475 assert_eq!(
1476 retryable,
1477 ["rate-limited", "provider-unavailable", "transport"]
1478 );
1479 }
1480
1481 /// unlike `McpCallError::Upstream`, which renders `HTTP {status}:
1482 /// {body}`, no variant here interpolates a provider response body — those
1483 /// carry completion fragments and request echoes.
1484 #[test]
1485 fn no_dispatch_error_display_leaks_a_provider_response_body() {
1486 const BODY: &str = r#"{"error":{"message":"your prompt was: summarize the memo"}}"#;
1487 // The only free-form field in the taxonomy. A provider impl is
1488 // contracted to pass a fixed classification string; prove that when it
1489 // does, nothing body-shaped can reach the rendered message.
1490 let classified = LlmCallError::Transport {
1491 model: MODEL.to_string(),
1492 detail: "connection reset".to_string(),
1493 };
1494 let rendered = classified.to_string();
1495 assert!(!rendered.contains(BODY));
1496 assert!(!rendered.contains("summarize the memo"));
1497 assert!(!rendered.contains('{'), "no JSON body shape: {rendered}");
1498 }
1499
1500 /// The degraded-error string: `message` follows the AI SDK's
1501 /// `error: unknown` → `errorText: string` rule, and carries no payload.
1502 #[test]
1503 fn failure_message_is_a_classification_not_a_payload_echo() {
1504 // Every message a provider impl is expected to produce is drawn from a
1505 // fixed classification vocabulary, so none can carry a response body.
1506 for reason in FailureReason::ALL {
1507 let failure = LlmFailure::new(reason, reason.default_message());
1508 assert!(
1509 !failure.message.is_empty(),
1510 "{reason} needs a message: it is the degraded error string"
1511 );
1512 assert!(
1513 !failure.message.contains("patient SSN"),
1514 "classification must not echo a payload: {}",
1515 failure.message
1516 );
1517 }
1518 }
1519
1520 #[test]
1521 fn llm_model_optionals_may_both_be_absent() {
1522 let declared = LlmModel::new("operator-declared");
1523 assert_eq!(declared.description, None);
1524 assert_eq!(declared.context_window, None);
1525 }
1526
1527 struct NoopProvider;
1528
1529 impl LlmProvider for NoopProvider {
1530 fn call<'a>(
1531 &'a self,
1532 model: &'a str,
1533 prompts: &'a [String],
1534 _schema_json: Option<&'a str>,
1535 ) -> std::pin::Pin<
1536 Box<
1537 dyn std::future::Future<Output = Result<Vec<LlmOutcome>, LlmCallError>> + Send + 'a,
1538 >,
1539 > {
1540 let n = prompts.len();
1541 let model = model.to_string();
1542 Box::pin(async move {
1543 if model.is_empty() {
1544 return Err(LlmCallError::NotConfigured { model });
1545 }
1546 // Obligation 1: one outcome per prompt, in order.
1547 Ok((0..n).map(|i| LlmOutcome::success(i.to_string())).collect())
1548 })
1549 }
1550
1551 fn models<'a>(
1552 &'a self,
1553 ) -> std::pin::Pin<
1554 Box<dyn std::future::Future<Output = Result<Vec<LlmModel>, LlmCallError>> + Send + 'a>,
1555 > {
1556 Box::pin(async { Ok(vec![LlmModel::new("noop")]) })
1557 }
1558 }
1559
1560 /// The trait must stay object-safe: the store holds it as `Arc<dyn ...>`.
1561 #[tokio::test]
1562 async fn provider_trait_is_object_safe_behind_arc() {
1563 let provider: Arc<dyn LlmProvider> = Arc::new(NoopProvider);
1564 let prompts = vec!["a".to_string(), "b".to_string()];
1565 let outcomes = provider
1566 .call(MODEL, &prompts, None)
1567 .await
1568 .expect("dispatch succeeds");
1569 assert_eq!(outcomes.len(), prompts.len(), "positional: one per prompt");
1570 assert_eq!(outcomes[1].text.as_deref(), Some("1"));
1571 assert_eq!(provider.models().await.unwrap()[0].name, "noop");
1572 }
1573
1574 #[tokio::test]
1575 async fn unconfigured_store_data_has_no_provider() {
1576 // The default runtime leaves the provider absent so an unconfigured
1577 // call reports a catchable configuration error rather than silently
1578 // succeeding — the `session_kv` rule.
1579 let data = crate::runtime::StoreData::with_tempdir().unwrap();
1580 assert!(data.llm_provider.is_none());
1581 }
1582}