Skip to main content

harness_loop/
trajectory.rs

1//! Trajectory checking: did the agent take a *correct path*?
2//!
3//! [`SessionRecorder`](crate::SessionRecorder) records what a run did and
4//! `harness trace` shows it, but neither says whether the run was right. A
5//! final answer can look fine while the agent skipped reading the file it
6//! edited, ran `git add -A`, or retried the same failing command ten times.
7//! This module judges the path itself, offline, from the JSONL log:
8//!
9//! - [`Trajectory`] flattens the raw event stream into model turns and tool
10//!   calls (a `PreTool` joined to its `PostTool`), which is the shape rules
11//!   actually want to ask questions of.
12//! - [`Rule`] / [`TrajectoryCheck`] / [`RuleSet`] state what a correct path
13//!   looks like: required tools, ordering, forbidden calls, argument shapes,
14//!   budgets, loops, and a clean finish.
15//! - [`Trajectory::compare_to_golden`] measures structural drift from a known
16//!   good run. Only tool names are compared, never model text, because the
17//!   text changes on every run while the shape of a good run does not.
18//!
19//! Everything here is deterministic and needs no model, so it can gate CI.
20
21use 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/// One step of a run, in the order it happened.
30#[derive(Debug, Clone, PartialEq)]
31pub enum Step {
32    /// The model answered once.
33    Model {
34        /// Whether the turn carried non-empty text. A turn with only tool
35        /// calls has none; the final turn of a finished run must.
36        has_text: bool,
37        /// How many tool calls this turn asked for.
38        tool_calls: usize,
39        input_tokens: u32,
40        output_tokens: u32,
41    },
42    /// A tool was called.
43    Tool {
44        id: String,
45        name: String,
46        args: Value,
47        ok: bool,
48        /// Why it failed. For a call that never got a result this is
49        /// `"no result"`, so a crashed run still reads as a failed call.
50        error: Option<String>,
51    },
52}
53
54/// A run flattened into ordered [`Step`]s plus totals.
55#[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    /// An `End` event was seen. A log without one was cut off.
63    pub ended: bool,
64}
65
66impl Trajectory {
67    /// Build the normalized view from recorded events.
68    ///
69    /// A `PreTool` starts a tool step already marked failed ("no result"); the
70    /// matching `PostTool` (same `call_id`) overwrites that. Doing it in this
71    /// order means a run that died mid-call is visible as a failure instead of
72    /// the call silently vanishing.
73    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    /// Read a JSONL session log and normalize it.
122    pub fn load(path: &Path) -> std::io::Result<Self> {
123        Ok(Self::from_events(&read_session(path)?))
124    }
125
126    /// Tool names in call order: the structure of the run.
127    pub fn tool_sequence(&self) -> Vec<String> {
128        self.tool_steps().map(|(_, s)| s.name.to_string()).collect()
129    }
130
131    /// Compare this run's tool sequence with a known good one.
132    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
171/// Borrowed view of a [`Step::Tool`], so rules do not re-match the enum.
172struct ToolStep<'a> {
173    name: &'a str,
174    args: &'a Value,
175    ok: bool,
176    error: Option<&'a str>,
177}
178
179/// Best human-readable reason a tool failed: a structured `error`/`message`
180/// field first, then a bare string payload, then the trace, then a stub.
181fn 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
196/// Levenshtein distance over whole tool names (one-row dynamic programming).
197fn 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/// How far a run's tool sequence is from a golden one.
216#[derive(Debug, Clone, PartialEq)]
217pub struct GoldenDiff {
218    /// Insertions, deletions and substitutions needed to turn one sequence
219    /// into the other.
220    pub distance: usize,
221    pub golden_len: usize,
222    pub actual_len: usize,
223    /// `distance` over the longer sequence: 0.0 identical, 1.0 nothing in
224    /// common. Normalizing makes one threshold work for short and long runs.
225    pub ratio: f64,
226}
227
228impl GoldenDiff {
229    /// Drift is acceptable when the ratio does not exceed `max_ratio`.
230    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/// One statement about what a correct path looks like.
246///
247/// Deserializes from TOML with a `rule = "<name>"` discriminator, so rule files
248/// can live next to the code and be reviewed like any other config.
249#[derive(Debug, Clone, PartialEq, Deserialize)]
250#[serde(tag = "rule", rename_all = "snake_case")]
251pub enum Rule {
252    /// The tool was called successfully at least once.
253    Requires { tool: String },
254    /// The first successful `before` precedes the first `after`. Passes if
255    /// `after` is never called; fails if `after` runs with no earlier `before`.
256    Order { before: String, after: String },
257    /// The tool is never called. With `args_match`, only calls whose
258    /// JSON-serialized args match the regex are violations.
259    Forbids {
260        tool: String,
261        #[serde(default)]
262        args_match: Option<String>,
263    },
264    /// Every call to `tool` has a string at `field` (a `/json/pointer`, or a
265    /// top-level key) matching `pattern`.
266    Args {
267        tool: String,
268        field: String,
269        pattern: String,
270    },
271    /// Upper bounds on model calls, tool calls, and total tokens.
272    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    /// The same call (tool and args) failed `times` or more times: a loop.
281    NoRepeatedFailure { times: u32 },
282    /// The run ended and its last model turn produced text.
283    Ends,
284}
285
286impl Rule {
287    /// Short stable name used in reports.
288    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    /// Reject a rule whose regex does not compile. Loading a rule file calls
301    /// this so a typo fails up front instead of at the end of a long check.
302    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/// One broken rule.
471#[derive(Debug, Clone, PartialEq)]
472pub struct Violation {
473    /// The rule's [`Rule::name`].
474    pub rule: String,
475    pub message: String,
476    /// Index into [`Trajectory::steps`], when the problem has a location.
477    pub step: Option<usize>,
478}
479
480/// Result of checking a trajectory against rules.
481#[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    /// Panic with the full report if any rule failed. For tests, where the
492    /// panic message is the diagnosis.
493    #[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/// Builder for a set of rules, for checks written in Rust.
518#[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    /// Add any rule, including ones the shortcuts below do not cover.
529    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    /// Run every rule; all violations are collected, not just the first.
593    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/// Rules loaded from a TOML file (`[[rules]]` tables).
607#[derive(Debug, Clone, Default, Deserialize)]
608pub struct RuleSet {
609    #[serde(default)]
610    pub rules: Vec<Rule>,
611}
612
613impl RuleSet {
614    /// Parse and validate rules from TOML text.
615    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}