1use std::collections::HashSet;
19
20use pv_state::BudgetSpec;
21use serde_json::Value;
22
23#[derive(Debug, Clone, PartialEq, Eq)]
27pub struct ToolPattern {
28 parts: Vec<String>,
29}
30
31impl ToolPattern {
32 pub fn from_scope(scope: &str) -> Option<Self> {
37 let scope = scope.trim();
38 let rest = if scope == "tool" {
39 "*"
40 } else {
41 let rest = scope.strip_prefix("tool.")?;
42 if rest.is_empty() {
43 "*"
44 } else {
45 rest
46 }
47 };
48 Some(Self {
49 parts: rest.split('*').map(str::to_owned).collect(),
50 })
51 }
52
53 pub fn matches(&self, tool: &str) -> bool {
55 let Some((first, rest)) = self.parts.split_first() else {
56 return false;
57 };
58 let Some(mut remaining) = tool.strip_prefix(first.as_str()) else {
59 return false;
60 };
61 let Some((last, middle)) = rest.split_last() else {
62 return remaining.is_empty();
64 };
65 for part in middle {
66 match remaining.find(part.as_str()) {
67 Some(index) => remaining = remaining.get(index + part.len()..).unwrap_or_default(),
68 None => return false,
69 }
70 }
71 remaining.ends_with(last.as_str())
72 }
73}
74
75#[derive(Debug, Clone)]
77pub struct PushedAllowlist {
78 pub policy: String,
80 pub scope: ToolPattern,
82 pub allow_only: Option<HashSet<String>>,
84 pub deny: HashSet<String>,
86}
87
88impl PushedAllowlist {
89 pub fn refuses(&self, tool: &str) -> Option<String> {
91 if !self.scope.matches(tool) {
92 return None;
93 }
94 if self.deny.contains(tool) {
95 return Some(format!("tool {tool:?} is deny-listed"));
96 }
97 match &self.allow_only {
98 Some(allowed) if !allowed.contains(tool) => Some(format!("tool {tool:?} is not allow-listed")),
99 _ => None,
100 }
101 }
102}
103
104pub struct PushedSchema {
106 pub policy: String,
108 pub tool: String,
110 validator: jsonschema::Validator,
111}
112
113impl std::fmt::Debug for PushedSchema {
114 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
115 f.debug_struct("PushedSchema")
116 .field("policy", &self.policy)
117 .field("tool", &self.tool)
118 .finish_non_exhaustive()
119 }
120}
121
122struct NoExternalRefs;
130
131impl jsonschema::Retrieve for NoExternalRefs {
132 fn retrieve(
133 &self,
134 uri: &jsonschema::Uri<String>,
135 ) -> Result<Value, Box<dyn std::error::Error + Send + Sync>> {
136 Err(format!(
137 "pushed schemas may not reference external documents ({})",
138 uri.scheme().as_str()
139 )
140 .into())
141 }
142}
143
144impl PushedSchema {
145 pub fn compile(policy: &str, tool: &str, schema: &Value) -> Result<Self, String> {
149 let validator = jsonschema::options()
150 .should_validate_formats(true)
151 .with_retriever(NoExternalRefs)
152 .build(schema)
153 .map_err(|e| format!("json_schema does not compile: {e}"))?;
154 Ok(Self {
155 policy: policy.to_owned(),
156 tool: tool.to_owned(),
157 validator,
158 })
159 }
160
161 pub fn errors(&self, arguments: &Value) -> Vec<String> {
163 self.validator
164 .iter_errors(arguments)
165 .take(3)
166 .map(|e| e.to_string())
167 .collect()
168 }
169}
170
171#[derive(Debug, Default)]
173pub struct PushedRules {
174 pub etag: Option<String>,
176 pub budget: BudgetSpec,
178 pub budget_policies: Vec<String>,
180 pub allowlists: Vec<PushedAllowlist>,
182 pub schemas: Vec<PushedSchema>,
184}
185
186impl PushedRules {
187 pub fn budget_policy_for(&self, local: &BudgetSpec, limit: &str, cap: u64) -> Option<String> {
192 let (pushed, local_cap) = match limit {
193 "max_tokens" => (self.budget.max_tokens, local.max_tokens),
194 "max_payout_usd_micros" => (self.budget.max_payout_usd_micros, local.max_payout_usd_micros),
195 "max_total_tool_calls" => (self.budget.max_total_tool_calls, local.max_total_tool_calls),
196 other => {
197 let tool = other.strip_prefix("max_tool_calls.")?;
198 (
199 self.budget.max_tool_calls.get(tool).copied(),
200 local.max_tool_calls.get(tool).copied(),
201 )
202 }
203 };
204 let from_console = pushed == Some(cap) && local_cap.is_none_or(|local_cap| local_cap > cap);
205 if from_console && !self.budget_policies.is_empty() {
206 Some(self.budget_policies.join(","))
207 } else {
208 None
209 }
210 }
211}
212
213#[cfg(test)]
214mod tests {
215 #![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
216
217 use super::*;
218 use serde_json::json;
219
220 #[test]
221 fn pushed_schemas_never_fetch_external_references() {
222 let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
224 listener.set_nonblocking(true).unwrap();
225 let url = format!("http://{}/schema.json", listener.local_addr().unwrap());
226 let err = PushedSchema::compile("p", "t", &json!({"$ref": url})).unwrap_err();
227 assert!(err.contains("may not reference external documents"), "{err}");
228 assert!(listener.accept().is_err(), "the daemon must not connect anywhere");
229
230 let meta = format!("http://{}/meta.json", listener.local_addr().unwrap());
232 assert!(PushedSchema::compile("p", "t", &json!({"$schema": meta})).is_err());
233 assert!(PushedSchema::compile("p", "t", &json!({"$dynamicRef": url})).is_err());
234 for draft in [
236 "https://json-schema.org/draft/2020-12/schema",
237 "https://json-schema.org/draft/2019-09/schema",
238 "http://json-schema.org/draft-07/schema#",
239 "http://json-schema.org/draft-07/schema",
240 "http://json-schema.org/draft-06/schema#",
241 "http://json-schema.org/draft-04/schema#",
242 ] {
243 PushedSchema::compile("p", "t", &json!({"$schema": draft, "type": "object"})).unwrap();
244 }
245 assert!(listener.accept().is_err(), "the daemon must not connect anywhere");
246
247 let file = std::env::temp_dir().join(format!("pv-pushed-ref-{}.json", std::process::id()));
248 std::fs::write(&file, r#"{"type":"object"}"#).unwrap();
249 let path = file.display().to_string().replace('\\', "/");
251 let file_ref = if path.starts_with('/') {
252 format!("file://{path}")
253 } else {
254 format!("file:///{path}")
255 };
256 let err = PushedSchema::compile("p", "t", &json!({"$ref": file_ref})).unwrap_err();
257 let _ = std::fs::remove_file(&file);
258 assert!(err.contains("may not reference external documents"), "{err}");
259
260 let local = PushedSchema::compile(
262 "p",
263 "t",
264 &json!({"$defs": {"amount": {"type": "integer", "maximum": 10}}, "properties": {"n": {"$ref": "#/$defs/amount"}}}),
265 )
266 .unwrap();
267 assert!(local.errors(&json!({"n": 5})).is_empty());
268 assert!(!local.errors(&json!({"n": 50})).is_empty());
269 }
270
271 #[test]
272 fn scopes_become_tool_patterns() {
273 let all = ToolPattern::from_scope("tool.*").unwrap();
274 assert!(all.matches("anything") && all.matches(""));
275 assert!(ToolPattern::from_scope("tool").unwrap().matches("x"));
276 let exact = ToolPattern::from_scope("tool.create_payment").unwrap();
277 assert!(exact.matches("create_payment"));
278 assert!(!exact.matches("create_payment_v2") && !exact.matches("create"));
279 let family = ToolPattern::from_scope("tool.payments.*").unwrap();
280 assert!(family.matches("payments.refund") && family.matches("payments."));
281 assert!(!family.matches("payment.refund"));
282 let middle = ToolPattern::from_scope("tool.db_*_write").unwrap();
283 assert!(middle.matches("db_orders_write") && middle.matches("db__write"));
284 assert!(!middle.matches("db_orders_read"));
285 for scope in ["llm.request", "llm.egress", "*", "", "tools.x"] {
286 assert!(ToolPattern::from_scope(scope).is_none(), "{scope}");
287 }
288 }
289
290 #[test]
291 fn allowlists_apply_only_inside_their_scope() {
292 let rule = PushedAllowlist {
293 policy: "runtime.tools".into(),
294 scope: ToolPattern::from_scope("tool.payments.*").unwrap(),
295 allow_only: Some(["payments.quote".to_owned()].into_iter().collect()),
296 deny: ["payments.wire".to_owned()].into_iter().collect(),
297 };
298 assert!(rule.refuses("search").is_none(), "outside the scope");
299 assert!(rule.refuses("payments.quote").is_none());
300 assert!(rule.refuses("payments.wire").unwrap().contains("deny-listed"));
301 assert!(rule
302 .refuses("payments.refund")
303 .unwrap()
304 .contains("not allow-listed"));
305 }
306
307 #[test]
308 fn pushed_schema_reports_bounded_errors() {
309 let schema = PushedSchema::compile(
310 "procurement.vendors",
311 "create_purchase_order",
312 &json!({"type": "object", "required": ["vendor"], "properties": {"vendor": {"enum": ["Contoso"]}}}),
313 )
314 .unwrap();
315 assert!(schema.errors(&json!({"vendor": "Contoso"})).is_empty());
316 assert!(!schema.errors(&json!({"vendor": "Initech"})).is_empty());
317 assert!(PushedSchema::compile("p", "t", &json!({"type": 12})).is_err());
318 }
319
320 #[test]
321 fn budget_refusals_name_the_console_policy_only_when_it_set_the_cap() {
322 let mut rules = PushedRules {
323 budget_policies: vec!["finance.cap".into()],
324 ..PushedRules::default()
325 };
326 rules.budget.max_total_tool_calls = Some(2);
327 rules.budget.max_tool_calls.insert("db_write".into(), 1);
328 let local = BudgetSpec {
329 max_total_tool_calls: Some(10),
330 ..BudgetSpec::default()
331 };
332 assert_eq!(
333 rules
334 .budget_policy_for(&local, "max_total_tool_calls", 2)
335 .as_deref(),
336 Some("finance.cap")
337 );
338 assert_eq!(
339 rules
340 .budget_policy_for(&local, "max_tool_calls.db_write", 1)
341 .as_deref(),
342 Some("finance.cap")
343 );
344 let tighter_local = BudgetSpec {
345 max_total_tool_calls: Some(1),
346 ..BudgetSpec::default()
347 };
348 assert!(rules
349 .budget_policy_for(&tighter_local, "max_total_tool_calls", 1)
350 .is_none());
351 assert!(rules.budget_policy_for(&local, "max_tokens", 5).is_none());
352 }
353}