1use crate::replay::{SessionEvent, read_session};
22use regex::Regex;
23use serde::Deserialize;
24use serde_json::Value;
25use std::collections::HashMap;
26use std::fmt;
27use std::path::Path;
28
29#[derive(Debug, Clone, PartialEq)]
31pub enum Step {
32 Model {
34 has_text: bool,
37 tool_calls: usize,
39 input_tokens: u32,
40 output_tokens: u32,
41 },
42 Tool {
44 id: String,
45 name: String,
46 args: Value,
47 ok: bool,
48 error: Option<String>,
51 },
52}
53
54#[derive(Debug, Clone, Default)]
56pub struct Trajectory {
57 pub steps: Vec<Step>,
58 pub model_calls: usize,
59 pub tool_calls: usize,
60 pub input_tokens: u32,
61 pub output_tokens: u32,
62 pub ended: bool,
64}
65
66impl Trajectory {
67 pub fn from_events(events: &[SessionEvent]) -> Self {
74 let mut t = Trajectory::default();
75 let mut pending: HashMap<String, usize> = HashMap::new();
76 for e in events {
77 match e {
78 SessionEvent::PostModel { output, .. } => {
79 t.model_calls += 1;
80 t.input_tokens += output.usage.input_tokens;
81 t.output_tokens += output.usage.output_tokens;
82 t.steps.push(Step::Model {
83 has_text: output.text.as_deref().is_some_and(|s| !s.trim().is_empty()),
84 tool_calls: output.tool_calls.len(),
85 input_tokens: output.usage.input_tokens,
86 output_tokens: output.usage.output_tokens,
87 });
88 }
89 SessionEvent::PreTool { action, .. } => {
90 t.tool_calls += 1;
91 pending.insert(action.call_id.clone(), t.steps.len());
92 t.steps.push(Step::Tool {
93 id: action.call_id.clone(),
94 name: action.tool.clone(),
95 args: action.args.clone(),
96 ok: false,
97 error: Some("no result".into()),
98 });
99 }
100 SessionEvent::PostTool {
101 call_id, result, ..
102 } => {
103 if let Some(&i) = pending.get(call_id)
104 && let Step::Tool { ok, error, .. } = &mut t.steps[i]
105 {
106 *ok = result.ok;
107 *error = if result.ok {
108 None
109 } else {
110 Some(failure_reason(result))
111 };
112 }
113 }
114 SessionEvent::End { .. } => t.ended = true,
115 _ => {}
116 }
117 }
118 t
119 }
120
121 pub fn load(path: &Path) -> std::io::Result<Self> {
123 Ok(Self::from_events(&read_session(path)?))
124 }
125
126 pub fn tool_sequence(&self) -> Vec<String> {
128 self.tool_steps().map(|(_, s)| s.name.to_string()).collect()
129 }
130
131 pub fn compare_to_golden(&self, golden: &Trajectory) -> GoldenDiff {
133 let a = self.tool_sequence();
134 let g = golden.tool_sequence();
135 let distance = edit_distance(&a, &g);
136 let longest = a.len().max(g.len());
137 GoldenDiff {
138 distance,
139 golden_len: g.len(),
140 actual_len: a.len(),
141 ratio: if longest == 0 {
142 0.0
143 } else {
144 distance as f64 / longest as f64
145 },
146 }
147 }
148
149 fn tool_steps(&self) -> impl Iterator<Item = (usize, ToolStep<'_>)> {
150 self.steps.iter().enumerate().filter_map(|(i, s)| match s {
151 Step::Tool {
152 name,
153 args,
154 ok,
155 error,
156 ..
157 } => Some((
158 i,
159 ToolStep {
160 name,
161 args,
162 ok: *ok,
163 error: error.as_deref(),
164 },
165 )),
166 _ => None,
167 })
168 }
169}
170
171struct ToolStep<'a> {
173 name: &'a str,
174 args: &'a Value,
175 ok: bool,
176 error: Option<&'a str>,
177}
178
179fn failure_reason(r: &harness_core::ToolResult) -> String {
182 for key in ["error", "message"] {
183 if let Some(s) = r.content.get(key).and_then(Value::as_str) {
184 return s.to_string();
185 }
186 }
187 if let Some(s) = r.content.as_str().filter(|s| !s.is_empty()) {
188 return s.to_string();
189 }
190 r.trace
191 .clone()
192 .filter(|s| !s.is_empty())
193 .unwrap_or_else(|| "tool failed".into())
194}
195
196fn edit_distance(a: &[String], b: &[String]) -> usize {
198 let mut row: Vec<usize> = (0..=b.len()).collect();
199 for (i, x) in a.iter().enumerate() {
200 let mut prev = row[0];
201 row[0] = i + 1;
202 for (j, y) in b.iter().enumerate() {
203 let cur = row[j + 1];
204 row[j + 1] = if x == y {
205 prev
206 } else {
207 1 + prev.min(row[j]).min(cur)
208 };
209 prev = cur;
210 }
211 }
212 row[b.len()]
213}
214
215#[derive(Debug, Clone, PartialEq)]
217pub struct GoldenDiff {
218 pub distance: usize,
221 pub golden_len: usize,
222 pub actual_len: usize,
223 pub ratio: f64,
226}
227
228impl GoldenDiff {
229 pub fn within(&self, max_ratio: f64) -> bool {
231 self.ratio <= max_ratio
232 }
233}
234
235impl fmt::Display for GoldenDiff {
236 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
237 write!(
238 f,
239 "tool sequence drift: edit distance {} (golden {} calls, actual {} calls, ratio {:.2})",
240 self.distance, self.golden_len, self.actual_len, self.ratio
241 )
242 }
243}
244
245#[derive(Debug, Clone, PartialEq, Deserialize)]
250#[serde(tag = "rule", rename_all = "snake_case")]
251pub enum Rule {
252 Requires { tool: String },
254 Order { before: String, after: String },
257 Forbids {
260 tool: String,
261 #[serde(default)]
262 args_match: Option<String>,
263 },
264 Args {
267 tool: String,
268 field: String,
269 pattern: String,
270 },
271 MaxSteps {
273 #[serde(default)]
274 model_calls: Option<u32>,
275 #[serde(default)]
276 tool_calls: Option<u32>,
277 #[serde(default)]
278 tokens: Option<u32>,
279 },
280 NoRepeatedFailure { times: u32 },
282 Ends,
284}
285
286impl Rule {
287 pub fn name(&self) -> &'static str {
289 match self {
290 Rule::Requires { .. } => "requires",
291 Rule::Order { .. } => "order",
292 Rule::Forbids { .. } => "forbids",
293 Rule::Args { .. } => "args",
294 Rule::MaxSteps { .. } => "max_steps",
295 Rule::NoRepeatedFailure { .. } => "no_repeated_failure",
296 Rule::Ends => "ends",
297 }
298 }
299
300 pub fn validate(&self) -> Result<(), String> {
303 let pat = match self {
304 Rule::Forbids {
305 args_match: Some(p),
306 ..
307 } => p,
308 Rule::Args { pattern, .. } => pattern,
309 _ => return Ok(()),
310 };
311 Regex::new(pat)
312 .map(|_| ())
313 .map_err(|e| format!("{}: invalid regex {pat:?}: {e}", self.name()))
314 }
315
316 fn check(&self, t: &Trajectory, out: &mut Vec<Violation>) {
317 let name = self.name();
318 let mut fail = |step: Option<usize>, message: String| {
319 out.push(Violation {
320 rule: name.into(),
321 message,
322 step,
323 })
324 };
325 match self {
326 Rule::Requires { tool } => {
327 if !t.tool_steps().any(|(_, s)| s.name == tool && s.ok) {
328 fail(None, format!("`{tool}` was never called successfully"));
329 }
330 }
331 Rule::Order { before, after } => {
332 let first_after = t.tool_steps().find(|(_, s)| s.name == after);
333 let first_before = t.tool_steps().find(|(_, s)| s.name == before && s.ok);
334 if let Some((ai, _)) = first_after {
335 match first_before {
336 Some((bi, _)) if bi < ai => {}
337 Some(_) => fail(
338 Some(ai),
339 format!("`{after}` was called before the first successful `{before}`"),
340 ),
341 None => fail(
342 Some(ai),
343 format!("`{after}` was called without a prior successful `{before}`"),
344 ),
345 }
346 }
347 }
348 Rule::Forbids { tool, args_match } => {
349 let re = args_match.as_deref().map(Regex::new);
350 let re = match re {
351 Some(Ok(r)) => Some(r),
352 Some(Err(e)) => {
353 fail(None, format!("invalid regex: {e}"));
354 return;
355 }
356 None => None,
357 };
358 for (i, s) in t.tool_steps().filter(|(_, s)| s.name == tool) {
359 let args = s.args.to_string();
360 match &re {
361 None => fail(Some(i), format!("forbidden tool `{tool}` was called")),
362 Some(r) if r.is_match(&args) => fail(
363 Some(i),
364 format!("forbidden call to `{tool}` with args {args}"),
365 ),
366 Some(_) => {}
367 }
368 }
369 }
370 Rule::Args {
371 tool,
372 field,
373 pattern,
374 } => {
375 let re = match Regex::new(pattern) {
376 Ok(r) => r,
377 Err(e) => {
378 fail(None, format!("invalid regex: {e}"));
379 return;
380 }
381 };
382 for (i, s) in t.tool_steps().filter(|(_, s)| s.name == tool) {
383 let v = if field.starts_with('/') {
384 s.args.pointer(field)
385 } else {
386 s.args.get(field.as_str())
387 };
388 match v.and_then(Value::as_str) {
389 None => fail(
390 Some(i),
391 format!("`{tool}` has no string at `{field}` (args {})", s.args),
392 ),
393 Some(text) if !re.is_match(text) => fail(
394 Some(i),
395 format!("`{tool}` `{field}` = {text:?} does not match /{pattern}/"),
396 ),
397 Some(_) => {}
398 }
399 }
400 }
401 Rule::MaxSteps {
402 model_calls,
403 tool_calls,
404 tokens,
405 } => {
406 if let Some(max) = model_calls
407 && t.model_calls > *max as usize
408 {
409 fail(
410 None,
411 format!("{} model calls exceeds the limit of {max}", t.model_calls),
412 );
413 }
414 if let Some(max) = tool_calls
415 && t.tool_calls > *max as usize
416 {
417 fail(
418 None,
419 format!("{} tool calls exceeds the limit of {max}", t.tool_calls),
420 );
421 }
422 if let Some(max) = tokens {
423 let total = u64::from(t.input_tokens) + u64::from(t.output_tokens);
424 if total > u64::from(*max) {
425 fail(None, format!("{total} tokens exceeds the limit of {max}"));
426 }
427 }
428 }
429 Rule::NoRepeatedFailure { times } => {
430 let limit = (*times).max(1) as usize;
431 let mut seen: HashMap<(String, String), usize> = HashMap::new();
432 for (i, s) in t.tool_steps().filter(|(_, s)| !s.ok) {
433 let n = seen
434 .entry((s.name.to_string(), s.args.to_string()))
435 .or_insert(0);
436 *n += 1;
437 if *n == limit {
438 fail(
439 Some(i),
440 format!(
441 "`{}` failed {limit} times with the same args {} ({})",
442 s.name,
443 s.args,
444 s.error.unwrap_or("no error")
445 ),
446 );
447 }
448 }
449 }
450 Rule::Ends => {
451 if !t.ended {
452 fail(None, "no End event: the run was cut off".into());
453 }
454 let last_model = t
455 .steps
456 .iter()
457 .enumerate()
458 .rev()
459 .find(|(_, s)| matches!(s, Step::Model { .. }));
460 match last_model {
461 Some((_, Step::Model { has_text: true, .. })) => {}
462 Some((i, _)) => fail(Some(i), "the last model turn produced no text".into()),
463 None => fail(None, "the run has no model turn".into()),
464 }
465 }
466 }
467 }
468}
469
470#[derive(Debug, Clone, PartialEq)]
472pub struct Violation {
473 pub rule: String,
475 pub message: String,
476 pub step: Option<usize>,
478}
479
480#[derive(Debug, Clone, Default, PartialEq)]
482pub struct Report {
483 pub violations: Vec<Violation>,
484}
485
486impl Report {
487 pub fn is_ok(&self) -> bool {
488 self.violations.is_empty()
489 }
490
491 #[track_caller]
494 pub fn assert_ok(&self) {
495 assert!(self.is_ok(), "trajectory check failed:\n{self}");
496 }
497}
498
499impl fmt::Display for Report {
500 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
501 if self.violations.is_empty() {
502 return write!(f, "trajectory ok");
503 }
504 for (i, v) in self.violations.iter().enumerate() {
505 if i > 0 {
506 writeln!(f)?;
507 }
508 match v.step {
509 Some(s) => write!(f, "[{}] step {}: {}", v.rule, s, v.message)?,
510 None => write!(f, "[{}] {}", v.rule, v.message)?,
511 }
512 }
513 Ok(())
514 }
515}
516
517#[derive(Debug, Clone, Default)]
519pub struct TrajectoryCheck {
520 rules: Vec<Rule>,
521}
522
523impl TrajectoryCheck {
524 pub fn new() -> Self {
525 Self::default()
526 }
527
528 pub fn rule(mut self, rule: Rule) -> Self {
530 self.rules.push(rule);
531 self
532 }
533
534 pub fn requires(self, tool: &str) -> Self {
535 self.rule(Rule::Requires { tool: tool.into() })
536 }
537
538 pub fn order(self, before: &str, after: &str) -> Self {
539 self.rule(Rule::Order {
540 before: before.into(),
541 after: after.into(),
542 })
543 }
544
545 pub fn forbids(self, tool: &str, args_match: Option<&str>) -> Self {
546 self.rule(Rule::Forbids {
547 tool: tool.into(),
548 args_match: args_match.map(Into::into),
549 })
550 }
551
552 pub fn args(self, tool: &str, field: &str, pattern: &str) -> Self {
553 self.rule(Rule::Args {
554 tool: tool.into(),
555 field: field.into(),
556 pattern: pattern.into(),
557 })
558 }
559
560 pub fn max_model_calls(self, n: u32) -> Self {
561 self.rule(Rule::MaxSteps {
562 model_calls: Some(n),
563 tool_calls: None,
564 tokens: None,
565 })
566 }
567
568 pub fn max_tool_calls(self, n: u32) -> Self {
569 self.rule(Rule::MaxSteps {
570 model_calls: None,
571 tool_calls: Some(n),
572 tokens: None,
573 })
574 }
575
576 pub fn max_tokens(self, n: u32) -> Self {
577 self.rule(Rule::MaxSteps {
578 model_calls: None,
579 tool_calls: None,
580 tokens: Some(n),
581 })
582 }
583
584 pub fn no_repeated_failure(self, times: u32) -> Self {
585 self.rule(Rule::NoRepeatedFailure { times })
586 }
587
588 pub fn ends(self) -> Self {
589 self.rule(Rule::Ends)
590 }
591
592 pub fn check(&self, t: &Trajectory) -> Report {
594 check_rules(&self.rules, t)
595 }
596}
597
598fn check_rules(rules: &[Rule], t: &Trajectory) -> Report {
599 let mut violations = Vec::new();
600 for r in rules {
601 r.check(t, &mut violations);
602 }
603 Report { violations }
604}
605
606#[derive(Debug, Clone, Default, Deserialize)]
608pub struct RuleSet {
609 #[serde(default)]
610 pub rules: Vec<Rule>,
611}
612
613impl RuleSet {
614 pub fn from_toml_str(s: &str) -> anyhow::Result<Self> {
616 let set: RuleSet = toml::from_str(s)?;
617 for r in &set.rules {
618 r.validate().map_err(anyhow::Error::msg)?;
619 }
620 Ok(set)
621 }
622
623 pub fn from_path(path: &Path) -> anyhow::Result<Self> {
624 let text = std::fs::read_to_string(path)
625 .map_err(|e| anyhow::anyhow!("read {}: {e}", path.display()))?;
626 Self::from_toml_str(&text).map_err(|e| anyhow::anyhow!("{}: {e}", path.display()))
627 }
628
629 pub fn check(&self, t: &Trajectory) -> Report {
630 check_rules(&self.rules, t)
631 }
632}