1use 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
10pub struct SandboxConfig {
12 pub schemas: HashMap<String, Value>,
14 pub budget: BudgetSpec,
16 pub payout_field: String,
19 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#[derive(Debug)]
36#[non_exhaustive]
37pub enum ToolVerdict {
38 Allowed {
40 tool: String,
42 budget_remaining: u64,
44 elapsed_us: u64,
46 payout_micros: u64,
53 principal_id: Option<String>,
60 },
61 Blocked {
63 tool: String,
65 stage: &'static str,
67 policy: Option<String>,
74 denial_code: DenialCode,
76 reason: String,
78 response: Value,
80 elapsed_us: u64,
82 },
83}
84
85impl ToolVerdict {
86 pub fn is_allowed(&self) -> bool {
88 matches!(self, Self::Allowed { .. })
89 }
90
91 pub fn elapsed_us(&self) -> u64 {
93 match self {
94 Self::Allowed { elapsed_us, .. } | Self::Blocked { elapsed_us, .. } => *elapsed_us,
95 }
96 }
97}
98
99pub struct Sandbox {
101 validators: HashMap<String, jsonschema::Validator>,
102 policies: Vec<Box<dyn PolicyEngine>>,
103 config: SandboxConfig,
104}
105
106impl Sandbox {
107 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 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 pub fn check(&self, store: &dyn StateStore, session: &str, raw: &[u8]) -> ToolVerdict {
133 self.check_with_principal(store, session, None, raw)
134 }
135
136 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 let req = match parse_tool_call(raw) {
160 Ok(r) => r,
161 Err(e) => {
162 let (stage, reason, denial_code) = match &e {
163 RpcError::NotToolCall(_) => (
167 "passthrough",
168 format!("passthrough refused: {e}"),
169 DenialCode::PassthroughRefused,
170 ),
171 _ => ("parse", e.to_string(), DenialCode::ParseError),
172 };
173 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 if let Some(validator) = self.validators.get(&req.tool) {
194 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 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 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 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 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 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 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 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
384fn 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 #[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 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 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 let v = s.check(
511 &store,
512 "s",
513 &raw_call("db_write", json!({"table": 42, "row": {}})),
514 );
515 assert!(!v.is_allowed());
516 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 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 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 #[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 assert_eq!(stage, "passthrough");
582 assert!(reason.contains("passthrough refused"), "{reason}");
583 }
584 other => panic!("{other:?}"),
585 }
586 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 #[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 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 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 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 let v = s.check(&store, "p", &raw_call("payout", json!({"amount_usd": 1.0})));
685 assert!(!v.is_allowed(), "{v:?}");
686 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 #[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 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}