Skip to main content

av_sandbox/
sandbox.rs

1//! The sandbox pipeline: parse → schema → policy → budget, with timing.
2
3use crate::denial_code::DenialCode;
4use crate::policy::{PolicyDecision, PolicyEngine};
5use crate::rpc::{denial_error, parse_tool_call, RpcError};
6use av_state::{ActionBudget, BudgetDecision, BudgetSpec, StateStore};
7use serde_json::Value;
8use std::collections::HashMap;
9
10/// Sandbox configuration.
11pub struct SandboxConfig {
12    /// Per-tool JSON Schemas for argument validation.
13    pub schemas: HashMap<String, Value>,
14    /// Action budget spec applied per session.
15    pub budget: BudgetSpec,
16    /// Argument field carrying a payout amount in USD (e.g. `amount_usd`).
17    /// When present on a call, it is charged against `max_payout_usd_micros`.
18    pub payout_field: String,
19    /// Reject tools without a configured schema.
20    pub require_schema: bool,
21}
22
23impl Default for SandboxConfig {
24    fn default() -> Self {
25        Self {
26            schemas: HashMap::new(),
27            budget: BudgetSpec::default(),
28            payout_field: "amount_usd".into(),
29            require_schema: false,
30        }
31    }
32}
33
34/// Verdict for one intercepted tool call.
35#[derive(Debug)]
36#[non_exhaustive]
37pub enum ToolVerdict {
38    /// Forward to the downstream tool server.
39    Allowed {
40        /// Tool name.
41        tool: String,
42        /// Remaining headroom under the binding budget dimension.
43        budget_remaining: u64,
44        /// Decision latency in microseconds (SLA surface, R23).
45        elapsed_us: u64,
46        /// Payout amount (USD micros) that
47        /// `ActionBudget::try_tool_call` debited on this call. Threaded
48        /// out so a lost `execution.claim()` race in the harness can
49        /// call [`ActionBudget::refund_tool_call`] with the exact same
50        /// amount, closing the concurrent-MCP budget
51        /// double-spend.
52        payout_micros: u64,
53        /// Principal id whose ledger was ALSO debited (when a
54        /// principal-scoped budget is bound). Threaded out so every
55        /// harness failure path that refunds the session ledger can
56        /// mirror the refund on the principal ledger — without this,
57        /// each lost claim race / saturated worker permanently consumed
58        /// principal quota with zero tools executed.
59        principal_id: Option<String>,
60    },
61    /// Block; respond to the agent with `response`.
62    Blocked {
63        /// Tool name (`<unparsed>` when the payload never parsed).
64        tool: String,
65        /// Which gate blocked: `parse` | `schema` | `policy` | `budget`.
66        stage: &'static str,
67        /// Name of the policy engine that denied, when the block came
68        /// from the named policy chain (`stage: "policy"`). Parse,
69        /// schema, and budget gates are configured limits rather than
70        /// named policies, so they stay `None` — the console's
71        /// per-policy attribution must never invent a policy that the
72        /// operator can't find in their policy inventory.
73        policy: Option<String>,
74        /// Structured denial code for machine consumption.
75        denial_code: DenialCode,
76        /// Human/machine-readable reason.
77        reason: String,
78        /// Ready-to-send JSON-RPC authorization error.
79        response: Value,
80        /// Decision latency in microseconds.
81        elapsed_us: u64,
82    },
83}
84
85impl ToolVerdict {
86    /// True when the call may proceed.
87    pub fn is_allowed(&self) -> bool {
88        matches!(self, Self::Allowed { .. })
89    }
90
91    /// Decision latency in microseconds.
92    pub fn elapsed_us(&self) -> u64 {
93        match self {
94            Self::Allowed { elapsed_us, .. } | Self::Blocked { elapsed_us, .. } => *elapsed_us,
95        }
96    }
97}
98
99/// The sandbox: compiled schemas + policy chain + budget spec.
100pub struct Sandbox {
101    validators: HashMap<String, jsonschema::Validator>,
102    policies: Vec<Box<dyn PolicyEngine>>,
103    config: SandboxConfig,
104}
105
106impl Sandbox {
107    /// Build a sandbox, compiling every tool schema up front (a schema that
108    /// fails to compile is a configuration error surfaced at boot, not a
109    /// silently-skipped validation at request time).
110    pub fn new(config: SandboxConfig, policies: Vec<Box<dyn PolicyEngine>>) -> Result<Self, String> {
111        let mut validators = HashMap::new();
112        for (tool, schema) in &config.schemas {
113            // Enforce `format` keywords. In
114            // jsonschema 0.49 / draft 2020-12, `format` is
115            // annotation-only unless opted in — an operator writing
116            // `"format": "uri"` into a tool-argument schema expected an
117            // enforced gate, not a no-op.
118            let v = jsonschema::options()
119                .should_validate_formats(true)
120                .build(schema)
121                .map_err(|e| format!("schema for tool {tool:?} does not compile: {e}"))?;
122            validators.insert(tool.clone(), v);
123        }
124        Ok(Self {
125            validators,
126            policies,
127            config,
128        })
129    }
130
131    /// Evaluate a raw MCP payload for `session`.
132    pub fn check(&self, store: &dyn StateStore, session: &str, raw: &[u8]) -> ToolVerdict {
133        self.check_with_principal(store, session, None, raw)
134    }
135
136    /// Same as [`Self::check`] but layers a principal-scoped budget check
137    /// on top of the session-scoped one. When `principal` is `Some(spec)`,
138    /// the tool call is refused unless BOTH ledgers admit the spend; the
139    /// principal ledger persists across every session belonging to the
140    /// same identity, which is what makes header-rotation attacks
141    /// (resetting the session bucket by minting a new session id)
142    /// visible.
143    ///
144    /// The two ledgers commit independently: if the principal check
145    /// passes and the session check fails, the principal debit is
146    /// refunded on the same dimensions before returning. Both budgets
147    /// must use identical `BudgetSpec` shapes so the refund is exact.
148    pub fn check_with_principal(
149        &self,
150        store: &dyn StateStore,
151        session: &str,
152        principal: Option<(&str, &BudgetSpec)>,
153        raw: &[u8],
154    ) -> ToolVerdict {
155        let started = std::time::Instant::now();
156        let elapsed = |s: std::time::Instant| u64::try_from(s.elapsed().as_micros()).unwrap_or(u64::MAX);
157
158        // Gate 1: parse.
159        let req = match parse_tool_call(raw) {
160            Ok(r) => r,
161            Err(e) => {
162                let (stage, reason, denial_code) = match &e {
163                    // Passthrough refusal parsed fine as JSON-RPC — a
164                    // deliberate policy choice, distinct HTTP class
165                    // (403) from true protocol failures (400).
166                    RpcError::NotToolCall(_) => (
167                        "passthrough",
168                        format!("passthrough refused: {e}"),
169                        DenialCode::PassthroughRefused,
170                    ),
171                    _ => ("parse", e.to_string(), DenialCode::ParseError),
172                };
173                // Protocol-level failures
174                // get the reserved JSON-RPC codes (-32700/-32600/
175                // -32602), not the -32001 policy code — the request
176                // never reached an authorization decision. The id is
177                // echoed whenever the envelope made it detectable
178                // (JSON-RPC 2.0 §5; see `detectable_error_id`).
179                let id = crate::rpc::detectable_error_id(raw, &e);
180                return ToolVerdict::Blocked {
181                    tool: "<unparsed>".into(),
182                    stage,
183                    policy: None,
184                    denial_code,
185                    reason,
186                    response: crate::rpc::protocol_error(id.as_ref(), &e),
187                    elapsed_us: elapsed(started),
188                };
189            }
190        };
191
192        // Gate 2: schema.
193        if let Some(validator) = self.validators.get(&req.tool) {
194            // Bounded work: a hostile client cannot force unbounded String
195            // allocation by sending pathologically-invalid arguments.
196            let errors: Vec<String> = validator
197                .iter_errors(&req.arguments)
198                .take(3)
199                .map(|e| e.to_string())
200                .collect();
201            if !errors.is_empty() {
202                let reason = format!("schema validation failed: {}", errors.join("; "));
203                return ToolVerdict::Blocked {
204                    tool: req.tool.clone(),
205                    stage: "schema",
206                    policy: None,
207                    denial_code: DenialCode::SchemaViolation,
208                    reason: reason.clone(),
209                    response: denial_error(req.id.as_ref(), DenialCode::SchemaViolation, &reason),
210                    elapsed_us: elapsed(started),
211                };
212            }
213        } else if self.config.require_schema {
214            let reason = format!("no argument schema configured for tool {:?}", req.tool);
215            return ToolVerdict::Blocked {
216                tool: req.tool.clone(),
217                stage: "schema",
218                policy: None,
219                denial_code: DenialCode::SchemaMissing,
220                reason: reason.clone(),
221                response: denial_error(req.id.as_ref(), DenialCode::SchemaMissing, &reason),
222                elapsed_us: elapsed(started),
223            };
224        }
225
226        // Gate 3: policy chain (first deny wins).
227        for policy in &self.policies {
228            if let PolicyDecision::Deny { reason } = policy.evaluate(&req.tool, &req.arguments) {
229                let reason = format!("policy {:?}: {reason}", policy.name());
230                return ToolVerdict::Blocked {
231                    tool: req.tool.clone(),
232                    stage: "policy",
233                    policy: Some(policy.name().to_owned()),
234                    denial_code: DenialCode::PolicyDenied,
235                    reason: reason.clone(),
236                    response: denial_error(req.id.as_ref(), DenialCode::PolicyDenied, &reason),
237                    elapsed_us: elapsed(started),
238                };
239            }
240        }
241
242        // Gate 4: budget (atomic check-and-spend).
243        let payout_micros = match extract_payout_micros(&req.arguments, &self.config.payout_field) {
244            Ok(m) => m,
245            Err(reason) => {
246                return ToolVerdict::Blocked {
247                    tool: req.tool.clone(),
248                    stage: "budget",
249                    policy: None,
250                    denial_code: DenialCode::PayoutInvalid,
251                    reason: reason.clone(),
252                    response: denial_error(req.id.as_ref(), DenialCode::PayoutInvalid, &reason),
253                    elapsed_us: elapsed(started),
254                }
255            }
256        };
257        let budget = ActionBudget::new(store, session, &self.config.budget);
258        // Principal-scoped pre-check. Debit the principal ledger
259        // first — a header-rotation attack that resets the session bucket
260        // still charges the same principal. If the principal ledger admits
261        // the spend but the session ledger refuses, we refund the principal
262        // debit before returning so the failed call consumed no quota
263        // anywhere.
264        if let Some((principal_id, principal_spec)) = principal {
265            let principal_budget = ActionBudget::for_principal(store, principal_id, principal_spec);
266            match principal_budget.try_tool_call(&req.tool, payout_micros) {
267                Ok(BudgetDecision::Allowed { .. }) => {}
268                Ok(BudgetDecision::Refused { limit, cap }) => {
269                    let reason = format!("principal action budget exceeded: {limit} (cap {cap})");
270                    return ToolVerdict::Blocked {
271                        tool: req.tool.clone(),
272                        stage: "budget",
273                        policy: None,
274                        denial_code: DenialCode::BudgetExceeded,
275                        reason: reason.clone(),
276                        response: denial_error(req.id.as_ref(), DenialCode::BudgetExceeded, &reason),
277                        elapsed_us: elapsed(started),
278                    };
279                }
280                Err(e) => {
281                    let reason = format!("principal budget check failed closed: {e}");
282                    return ToolVerdict::Blocked {
283                        tool: req.tool.clone(),
284                        stage: "budget",
285                        policy: None,
286                        denial_code: DenialCode::BudgetError,
287                        reason: reason.clone(),
288                        response: denial_error(req.id.as_ref(), DenialCode::BudgetError, &reason),
289                        elapsed_us: elapsed(started),
290                    };
291                }
292            }
293            match budget.try_tool_call(&req.tool, payout_micros) {
294                Ok(BudgetDecision::Allowed { remaining }) => {
295                    return ToolVerdict::Allowed {
296                        tool: req.tool,
297                        budget_remaining: remaining,
298                        elapsed_us: elapsed(started),
299                        payout_micros,
300                        principal_id: Some(principal_id.to_owned()),
301                    };
302                }
303                Ok(BudgetDecision::Refused { limit, cap }) => {
304                    // Session refused; unwind the principal debit so a
305                    // failed call is fully non-consumptive.
306                    principal_budget.refund_tool_call(&req.tool, payout_micros);
307                    let reason = format!("action budget exceeded: {limit} (cap {cap})");
308                    return ToolVerdict::Blocked {
309                        tool: req.tool.clone(),
310                        stage: "budget",
311                        policy: None,
312                        denial_code: DenialCode::BudgetExceeded,
313                        reason: reason.clone(),
314                        response: denial_error(req.id.as_ref(), DenialCode::BudgetExceeded, &reason),
315                        elapsed_us: elapsed(started),
316                    };
317                }
318                Err(e) => {
319                    principal_budget.refund_tool_call(&req.tool, payout_micros);
320                    let reason = format!("budget check failed closed: {e}");
321                    return ToolVerdict::Blocked {
322                        tool: req.tool.clone(),
323                        stage: "budget",
324                        policy: None,
325                        denial_code: DenialCode::BudgetError,
326                        reason: reason.clone(),
327                        response: denial_error(req.id.as_ref(), DenialCode::BudgetError, &reason),
328                        elapsed_us: elapsed(started),
329                    };
330                }
331            }
332        }
333        match budget.try_tool_call(&req.tool, payout_micros) {
334            Ok(BudgetDecision::Allowed { remaining }) => ToolVerdict::Allowed {
335                tool: req.tool,
336                budget_remaining: remaining,
337                elapsed_us: elapsed(started),
338                // Thread the actual debited amount so the
339                // harness can refund exactly this much on a lost
340                // execution.claim() race.
341                payout_micros,
342                principal_id: None,
343            },
344            Ok(BudgetDecision::Refused { limit, cap }) => {
345                let reason = format!("action budget exceeded: {limit} (cap {cap})");
346                ToolVerdict::Blocked {
347                    tool: req.tool.clone(),
348                    stage: "budget",
349                    policy: None,
350                    denial_code: DenialCode::BudgetExceeded,
351                    reason: reason.clone(),
352                    response: denial_error(req.id.as_ref(), DenialCode::BudgetExceeded, &reason),
353                    elapsed_us: elapsed(started),
354                }
355            }
356            Err(e) => {
357                // State-store failure fails closed.
358                let reason = format!("budget check failed closed: {e}");
359                ToolVerdict::Blocked {
360                    tool: req.tool.clone(),
361                    stage: "budget",
362                    policy: None,
363                    denial_code: DenialCode::BudgetError,
364                    reason: reason.clone(),
365                    response: denial_error(req.id.as_ref(), DenialCode::BudgetError, &reason),
366                    elapsed_us: elapsed(started),
367                }
368            }
369        }
370    }
371
372    /// Evaluate an arbitrary operation payload against the configured native
373    /// and WASM policy chain. Used for chat sanitization before compression.
374    pub fn sanitize(&self, operation: &str, payload: &Value) -> Result<(), String> {
375        for policy in &self.policies {
376            if let PolicyDecision::Deny { reason } = policy.evaluate(operation, payload) {
377                return Err(format!("policy {:?}: {reason}", policy.name()));
378            }
379        }
380        Ok(())
381    }
382}
383
384/// Extract a payout amount (USD float or integer) as micro-USD. Rejects
385/// negatives, NaN/Inf, and absurd magnitudes rather than truncating silently.
386fn extract_payout_micros(arguments: &Value, field: &str) -> Result<u64, String> {
387    let Some(v) = arguments.get(field) else {
388        return Ok(0);
389    };
390    let usd = v.as_f64().ok_or_else(|| format!("{field} must be a number"))?;
391    if !usd.is_finite() || usd < 0.0 {
392        return Err(format!("{field} must be a finite non-negative number, got {usd}"));
393    }
394    if usd > 1.0e12 {
395        return Err(format!("{field} of {usd} exceeds sanity bounds"));
396    }
397    // Ceil, not round: a positive sub-micro amount (e.g. 0.0000004 USD)
398    // rounded to ZERO micros, and a zero spend skips payout enforcement
399    // entirely (`ActionBudget` only gates payout > 0) — including the
400    // fail-closed "payout with NO cap configured is refused" rule. A
401    // hostile tool argument could drip unaccounted positive payouts
402    // forever. Any positive amount now costs at least one micro-USD,
403    // so the cap (or its absence) always gets a say; exact zero stays
404    // zero.
405    #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
406    Ok((usd * av_core::units::USD_MICROS_PER_DOLLAR as f64).ceil() as u64)
407}
408
409#[cfg(test)]
410mod tests {
411    #![allow(
412        clippy::unwrap_used,
413        clippy::expect_used,
414        clippy::panic,
415        clippy::indexing_slicing
416    )]
417
418    use super::*;
419    use crate::policy::NativePolicy;
420    use av_state::InMemoryStore;
421    use serde_json::json;
422
423    fn raw_call(tool: &str, args: Value) -> Vec<u8> {
424        serde_json::to_vec(&json!({
425            "jsonrpc": "2.0", "id": 1, "method": "tools/call",
426            "params": {"name": tool, "arguments": args}
427        }))
428        .unwrap()
429    }
430
431    fn sandbox() -> Sandbox {
432        let mut schemas = HashMap::new();
433        schemas.insert(
434            "db_write".to_owned(),
435            json!({
436                "type": "object",
437                "required": ["table", "row"],
438                "properties": {
439                    "table": {"type": "string", "minLength": 1},
440                    "row": {"type": "object"}
441                },
442                "additionalProperties": false
443            }),
444        );
445        let budget = BudgetSpec {
446            max_tool_calls: [("db_write".to_owned(), 3)].into_iter().collect(),
447            max_payout_usd_micros: Some(50_000_000),
448            max_total_tool_calls: Some(100),
449            max_tokens: None,
450        };
451        Sandbox::new(
452            SandboxConfig {
453                schemas,
454                budget,
455                payout_field: "amount_usd".into(),
456                require_schema: false,
457            },
458            vec![Box::new(NativePolicy::deny_tools(&["rm_rf"]))],
459        )
460        .unwrap()
461    }
462
463    #[test]
464    fn valid_call_allowed_with_latency() {
465        let store = InMemoryStore::new();
466        let sandbox = sandbox();
467        let v = sandbox.check(
468            &store,
469            "s",
470            &raw_call("db_write", json!({"table": "t", "row": {}})),
471        );
472        assert!(v.is_allowed(), "{v:?}");
473        // R23: block/allow decision < 5ms = 5000µs. A single debug-profile
474        // sample flakes under full-suite load (scheduler preemption between
475        // the two Instant reads), so take the minimum over several runs:
476        // one clean time slice is enough for a genuinely fast decision,
477        // while a real regression stays slow on every sample. Distinct
478        // sessions keep the per-session db_write budget (3) out of play.
479        // The release SLA gate and criterion bench own the authoritative
480        // measurement.
481        let min_elapsed_us = (0..10)
482            .map(|i| {
483                let verdict = sandbox.check(
484                    &store,
485                    &format!("s-latency-{i}"),
486                    &raw_call("db_write", json!({"table": "t", "row": {}})),
487                );
488                assert!(verdict.is_allowed(), "{verdict:?}");
489                verdict.elapsed_us()
490            })
491            .min()
492            .unwrap_or(u64::MAX);
493        assert!(min_elapsed_us < 5_000, "fastest decision took {min_elapsed_us}µs");
494    }
495
496    #[test]
497    fn schema_invalid_blocked() {
498        let store = InMemoryStore::new();
499        let s = sandbox();
500        // Missing required field.
501        let v = s.check(&store, "s", &raw_call("db_write", json!({"table": "t"})));
502        match &v {
503            ToolVerdict::Blocked { stage, response, .. } => {
504                assert_eq!(*stage, "schema");
505                assert_eq!(response["error"]["code"], -32001);
506            }
507            other => panic!("{other:?}"),
508        }
509        // Type confusion.
510        let v = s.check(
511            &store,
512            "s",
513            &raw_call("db_write", json!({"table": 42, "row": {}})),
514        );
515        assert!(!v.is_allowed());
516        // Extra field (additionalProperties: false).
517        let v = s.check(
518            &store,
519            "s",
520            &raw_call("db_write", json!({"table": "t", "row": {}, "backdoor": true})),
521        );
522        assert!(!v.is_allowed());
523    }
524
525    #[test]
526    fn policy_denied_blocked() {
527        let store = InMemoryStore::new();
528        let v = sandbox().check(&store, "s", &raw_call("rm_rf", json!({})));
529        match v {
530            ToolVerdict::Blocked {
531                stage,
532                policy,
533                reason,
534                ..
535            } => {
536                assert_eq!(stage, "policy");
537                assert!(reason.contains("deny-listed"), "{reason}");
538                // Named-policy denials carry the engine name so the
539                // console can attribute the block to a policy row.
540                assert_eq!(policy.as_deref(), Some("deny_tools"));
541            }
542            other => panic!("{other:?}"),
543        }
544    }
545
546    #[test]
547    fn non_policy_gates_carry_no_policy_attribution() {
548        let store = InMemoryStore::new();
549        // Parse gate: invalid JSON-RPC.
550        let parse = sandbox().check(&store, "s", b"not json");
551        match parse {
552            ToolVerdict::Blocked { stage, policy, .. } => {
553                assert_eq!(stage, "parse");
554                assert!(policy.is_none(), "parse gate must not invent a policy name");
555            }
556            other => panic!("{other:?}"),
557        }
558    }
559
560    /// JSON-RPC 2.0 §5: parse-gate refusals of a WELL-FORMED envelope
561    /// must echo the request id in the error response. Pre-fix the
562    /// blocked-response body always carried `id: null`, so SDK
563    /// correlation tables never matched and the pending call hung.
564    #[test]
565    fn blocked_parse_verdict_echoes_the_request_id() {
566        let store = InMemoryStore::new();
567        let raw = br#"{"jsonrpc":"2.0","id":42,"method":"initialize","params":{}}"#;
568        match sandbox().check(&store, "s", raw) {
569            ToolVerdict::Blocked {
570                response,
571                stage,
572                reason,
573                ..
574            } => {
575                assert_eq!(response["id"], json!(42), "{response}");
576                // Mutation-run hardening (round 8): the passthrough
577                // stage label was unpinned — deleting the NotToolCall
578                // match arm reclassified passthrough refusals as
579                // "parse" failures, corrupting the stage-keyed metrics
580                // and events operators alert on.
581                assert_eq!(stage, "passthrough");
582                assert!(reason.contains("passthrough refused"), "{reason}");
583            }
584            other => panic!("{other:?}"),
585        }
586        // Unparseable JSON: the id is undetectable — null stands.
587        match sandbox().check(&store, "s", b"not json") {
588            ToolVerdict::Blocked { response, stage, .. } => {
589                assert_eq!(response["id"], serde_json::Value::Null, "{response}");
590                assert_eq!(stage, "parse");
591            }
592            other => panic!("{other:?}"),
593        }
594    }
595
596    #[test]
597    fn budget_enforced_across_calls() {
598        let store = InMemoryStore::new();
599        let s = sandbox();
600        let args = json!({"table": "t", "row": {}});
601        for _ in 0..3 {
602            assert!(s
603                .check(&store, "sess", &raw_call("db_write", args.clone()))
604                .is_allowed());
605        }
606        let v = s.check(&store, "sess", &raw_call("db_write", args));
607        match v {
608            ToolVerdict::Blocked { stage, reason, .. } => {
609                assert_eq!(stage, "budget");
610                assert!(reason.contains("db_write"), "{reason}");
611            }
612            other => panic!("{other:?}"),
613        }
614    }
615
616    /// An `Allowed` verdict must thread the bound principal id out so the
617    /// harness's failure paths (lost claim race, saturated worker, close
618    /// race) can mirror the session-ledger refund on the principal ledger.
619    /// Without it, every failed claim permanently consumed principal quota
620    /// with zero tools executed — and the principal ledger persists across
621    /// sessions, so that leak never healed. The refund round-trip below is
622    /// exactly the compensation the harness performs.
623    #[test]
624    fn allowed_verdict_threads_principal_id_and_refund_restores_the_ledger() {
625        let store = InMemoryStore::new();
626        let s = sandbox();
627        let args = json!({"amount_usd": 1.0});
628        let principal_spec = BudgetSpec {
629            max_tool_calls: std::collections::BTreeMap::new(),
630            max_payout_usd_micros: Some(50_000_000),
631            max_total_tool_calls: Some(1),
632            max_tokens: None,
633        };
634        // No principal bound: field is None.
635        match s.check(&store, "sess-a", &raw_call("payout", args.clone())) {
636            ToolVerdict::Allowed { principal_id, .. } => assert_eq!(principal_id, None),
637            other => panic!("{other:?}"),
638        }
639        // Principal bound: the verdict carries the id and payout the
640        // harness needs for compensation.
641        let verdict = s.check_with_principal(
642            &store,
643            "sess-b",
644            Some(("principal-1", &principal_spec)),
645            &raw_call("payout", args.clone()),
646        );
647        let (tool, payout_micros, principal_id) = match verdict {
648            ToolVerdict::Allowed {
649                tool,
650                payout_micros,
651                principal_id,
652                ..
653            } => (tool, payout_micros, principal_id),
654            other => panic!("{other:?}"),
655        };
656        assert_eq!(principal_id.as_deref(), Some("principal-1"));
657        assert_eq!(payout_micros, 1_000_000);
658        // Simulate the harness failure path: refund BOTH ledgers with the
659        // threaded values. With max_total_tool_calls = 1, the next call
660        // only admits if the principal refund actually landed.
661        ActionBudget::new(&store, "sess-b", &s.config.budget).refund_tool_call(&tool, payout_micros);
662        ActionBudget::for_principal(&store, "principal-1", &principal_spec)
663            .refund_tool_call(&tool, payout_micros);
664        let retry = s.check_with_principal(
665            &store,
666            "sess-b",
667            Some(("principal-1", &principal_spec)),
668            &raw_call("payout", args),
669        );
670        assert!(
671            retry.is_allowed(),
672            "principal ledger did not heal after the compensating refund: {retry:?}"
673        );
674    }
675
676    #[test]
677    fn payout_cap_enforced_and_hostile_amounts_rejected() {
678        let store = InMemoryStore::new();
679        let s = sandbox();
680        assert!(s
681            .check(&store, "p", &raw_call("payout", json!({"amount_usd": 49.5})))
682            .is_allowed());
683        // Would cross $50 total.
684        let v = s.check(&store, "p", &raw_call("payout", json!({"amount_usd": 1.0})));
685        assert!(!v.is_allowed(), "{v:?}");
686        // Hostile numerics.
687        for bad in [
688            json!({"amount_usd": -5}),
689            json!({"amount_usd": "1e9"}),
690            json!({"amount_usd": 1e15}),
691        ] {
692            let v = s.check(&store, "p2", &raw_call("payout", bad));
693            assert!(!v.is_allowed(), "hostile payout accepted: {v:?}");
694        }
695    }
696
697    #[test]
698    fn unparsable_payload_blocked_at_parse() {
699        let store = InMemoryStore::new();
700        let v = sandbox().check(&store, "s", b"\xff\xfe not json");
701        match v {
702            ToolVerdict::Blocked { stage, tool, .. } => {
703                assert_eq!(stage, "parse");
704                assert_eq!(tool, "<unparsed>");
705            }
706            other => panic!("{other:?}"),
707        }
708    }
709
710    #[test]
711    fn tools_without_schema_pass_schema_gate() {
712        let store = InMemoryStore::new();
713        let v = sandbox().check(&store, "s", &raw_call("search", json!({"anything": [1, 2]})));
714        assert!(v.is_allowed(), "{v:?}");
715    }
716
717    #[test]
718    fn required_schema_mode_blocks_unknown_tools() {
719        let sandbox = Sandbox::new(
720            SandboxConfig {
721                require_schema: true,
722                ..SandboxConfig::default()
723            },
724            Vec::new(),
725        )
726        .unwrap();
727        let verdict = sandbox.check(&InMemoryStore::new(), "session", &raw_call("unknown", json!({})));
728        assert!(matches!(verdict, ToolVerdict::Blocked { stage: "schema", .. }));
729    }
730
731    #[test]
732    fn bad_boot_schema_is_a_boot_error() {
733        let mut schemas = HashMap::new();
734        schemas.insert("t".to_owned(), json!({"type": "not-a-type"}));
735        let err = Sandbox::new(
736            SandboxConfig {
737                schemas,
738                ..SandboxConfig::default()
739            },
740            vec![],
741        );
742        assert!(err.is_err(), "invalid schema must fail at boot, not be skipped");
743    }
744
745    #[test]
746    fn chat_payload_uses_the_same_policy_chain() {
747        let sandbox = Sandbox::new(
748            SandboxConfig::default(),
749            vec![Box::new(NativePolicy::new("no-secret", |operation, payload| {
750                if operation == "chat/completions" && payload.to_string().contains("secret") {
751                    PolicyDecision::Deny {
752                        reason: "secret content".into(),
753                    }
754                } else {
755                    PolicyDecision::Allow
756                }
757            }))],
758        )
759        .unwrap();
760        assert!(sandbox
761            .sanitize("chat/completions", &json!({"messages": ["safe"]}))
762            .is_ok());
763        assert!(sandbox
764            .sanitize("chat/completions", &json!({"messages": ["secret"]}))
765            .is_err());
766    }
767}
768
769#[cfg(test)]
770mod payout_boundary_tests {
771    #![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
772
773    use super::*;
774    use serde_json::json;
775
776    /// Mutation-run hardening: pin the payout parser's exact
777    /// bounds. `usd < 0.0` -> `<=` would refuse a legitimate zero payout;
778    /// `usd > 1.0e12` -> `>=`/`==` shift or invert the sanity ceiling.
779    #[test]
780    fn payout_bounds_are_exact() {
781        let extract = |v: serde_json::Value| extract_payout_micros(&json!({"amount_usd": v}), "amount_usd");
782        assert_eq!(extract(json!(0.0)).unwrap(), 0, "zero payout is legitimate");
783        assert_eq!(
784            extract(json!(1.0e12)).unwrap(),
785            1_000_000_000_000_u64 * av_core::units::USD_MICROS_PER_DOLLAR,
786            "the sanity ceiling itself is accepted"
787        );
788        let over = extract(json!(2.0e12));
789        assert!(
790            matches!(over, Err(ref m) if m.contains("sanity")),
791            "past the ceiling must be refused, got {over:?}"
792        );
793        let negative = extract(json!(-0.25));
794        assert!(matches!(negative, Err(ref m) if m.contains("non-negative")));
795        // Sub-micro positive amounts must CEIL to one micro, never round
796        // to zero: a zero result skips payout enforcement entirely —
797        // including the fail-closed no-cap refusal — so a hostile tool
798        // argument of 0.0000004 USD per call could drip unaccounted
799        // positive payouts forever.
800        assert_eq!(
801            extract(json!(0.0000004)).unwrap(),
802            1,
803            "positive sub-micro payouts must cost at least one micro-USD"
804        );
805        assert_eq!(extract(json!(0.0000019)).unwrap(), 2, "ceil, not round");
806    }
807}