Skip to main content

hh_core/
step.rs

1//! Step-assignment pass (FR-3.4): 1-based step ordinals for semantic events.
2//!
3//! The pass is a pure function over a slice of [`EventRow`]s: it assigns each
4//! event's `step` in place. The store reads rows via
5//! [`crate::store::Store::list_events`], runs this pass, and writes the
6//! ordinals back in one transaction
7//! ([`crate::store::Store::assign_steps`]).
8//!
9//! ## Rules (FR-3.4)
10//! - `terminal_output` events are not steps → `step = None`.
11//! - Every other semantic event gets its own incrementing 1-based step, *except*
12//!   a `tool_result` whose `correlates` points to a `tool_call` present in the
13//!   session: that result shares the call's step (a call and its result are one
14//!   step).
15//! - A `tool_result` with `correlates = None`, or pointing to an id absent from
16//!   the session (orphan / concurrent result-before-call from another source),
17//!   gets its own step.
18//!
19//! The pass is **order-independent for correlation**: it builds an id→step map
20//! of the calls in pass 1 and resolves deferred results against it in pass 2,
21//! so a result that sorts before its call (e.g. a tool_result block emitted
22//! ahead of the tool_use block in the same record, or concurrent sources) still
23//! resolves to the call's step.
24//!
25//! ## Deviation (flagged, ADR-0002)
26//!
27//! ADR-0001 said step ordinals are "derived at read time". Storing them at
28//! finalize (authoritative `events.step` column) + self-healing on
29//! [`crate::store::Store::open`] is a deliberate deviation recorded in
30//! `docs/adr/0002-stored-step-ordinals.md`: it makes `hh list`'s step count a
31//! trivial `COUNT(DISTINCT step)`, and self-heal repairs both a crashed
32//! finalize and the attached-MCP-proxy late-event race on the next `hh`
33//! invocation.
34
35use crate::event::{EventKind, EventRow};
36use std::collections::{HashMap, HashSet};
37
38/// Assign 1-based step ordinals to `events` in place (FR-3.4). See the module
39/// docs for the rules. Idempotent: re-running on already-assigned rows yields
40/// the same ordinals.
41pub fn assign_steps(events: &mut [EventRow]) {
42    // Order by (ts_ms, id) so ordinals are stable and chronological. `id` is
43    // the rowid PK, so (ts_ms, id) is a strict total order (no ties).
44    events.sort_by_key(|e| (e.ts_ms, e.id));
45
46    // Upfront knowledge of which ids are present and which are tool calls, so a
47    // result can decide to defer even when its call sorts later in the order.
48    let mut present: HashSet<i64> = HashSet::with_capacity(events.len());
49    let mut is_call: HashSet<i64> = HashSet::new();
50    for e in events.iter() {
51        present.insert(e.id);
52        if e.kind == EventKind::ToolCall {
53            is_call.insert(e.id);
54        }
55    }
56
57    let mut call_step: HashMap<i64, i64> = HashMap::new();
58    let mut deferred: Vec<usize> = Vec::new();
59    let mut counter: i64 = 0;
60
61    // Pass 1: assign steps to terminal-excluded + non-deferred events; record
62    // each call's step for pass 2.
63    for (i, e) in events.iter_mut().enumerate() {
64        if e.kind == EventKind::TerminalOutput {
65            e.step = None;
66            continue;
67        }
68        // Defer a tool_result only if it points to a tool_call present in this
69        // session; otherwise it gets its own step now (orphan / no correlation).
70        let defer = e.kind == EventKind::ToolResult
71            && matches!(e.correlates, Some(cid) if present.contains(&cid) && is_call.contains(&cid));
72        if defer {
73            deferred.push(i);
74            continue;
75        }
76        counter += 1;
77        e.step = Some(counter);
78        if e.kind == EventKind::ToolCall {
79            call_step.insert(e.id, counter);
80        }
81    }
82
83    // Pass 2: deferred results borrow their call's step. The deferral guard
84    // guarantees the call exists and is a ToolCall (which received a step in
85    // pass 1), so the lookup succeeds; the let-else + fallback are defensive
86    // only — no unwrap/expect, per CLAUDE.md.
87    for &i in &deferred {
88        let e = &mut events[i];
89        let Some(cid) = e.correlates else {
90            counter += 1;
91            e.step = Some(counter);
92            continue;
93        };
94        if let Some(&step) = call_step.get(&cid) {
95            e.step = Some(step);
96        } else {
97            counter += 1;
98            e.step = Some(counter);
99        }
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106    use crate::event::{EventKind, EventRow};
107
108    /// Build an EventRow with the given (id, ts_ms, kind, correlates); step
109    /// starts as `None` (the pass assigns it).
110    fn row(id: i64, ts_ms: i64, kind: EventKind, correlates: Option<i64>) -> EventRow {
111        EventRow {
112            id,
113            session_id: "s".into(),
114            ts_ms,
115            kind,
116            step: None,
117            correlates,
118        }
119    }
120
121    fn steps(rows: &[EventRow]) -> Vec<Option<i64>> {
122        rows.iter().map(|r| r.step).collect()
123    }
124
125    #[test]
126    fn assign_steps_normal_call_result_share() {
127        let mut rows = vec![
128            row(1, 0, EventKind::ToolCall, None),
129            row(2, 5, EventKind::ToolResult, Some(1)),
130        ];
131        assign_steps(&mut rows);
132        // Call gets step 1; result shares step 1.
133        assert_eq!(steps(&rows), vec![Some(1), Some(1)]);
134    }
135
136    #[test]
137    fn assign_steps_result_before_call() {
138        // Concurrent source: result sorts before the call by ts_ms.
139        let mut rows = vec![
140            row(2, 0, EventKind::ToolResult, Some(1)),
141            row(1, 10, EventKind::ToolCall, None),
142        ];
143        assign_steps(&mut rows);
144        // After sort by (ts_ms, id): result(id=2,ts=0) then call(id=1,ts=10).
145        // The result defers (correlates → a present call); pass 2 gives it the
146        // call's step. The call is the first step assigned → step 1; result → 1.
147        assert_eq!(steps(&rows), vec![Some(1), Some(1)]);
148    }
149
150    #[test]
151    fn assign_steps_orphan_result_own_step() {
152        // Result points to a call id not present in the session → orphan, own step.
153        let mut rows = vec![
154            row(1, 0, EventKind::UserMessage, None),
155            row(2, 5, EventKind::ToolResult, Some(99)),
156        ];
157        assign_steps(&mut rows);
158        // UserMessage → step 1; orphan result → step 2 (not deferred).
159        assert_eq!(steps(&rows), vec![Some(1), Some(2)]);
160    }
161
162    #[test]
163    fn assign_steps_none_correlates_own_step() {
164        let mut rows = vec![
165            row(1, 0, EventKind::ToolCall, None),
166            row(2, 5, EventKind::ToolResult, None),
167        ];
168        assign_steps(&mut rows);
169        // Call → 1; result with no correlation → 2.
170        assert_eq!(steps(&rows), vec![Some(1), Some(2)]);
171    }
172
173    #[test]
174    fn assign_steps_terminal_output_null() {
175        let mut rows = vec![
176            row(1, 0, EventKind::UserMessage, None),
177            row(2, 1, EventKind::TerminalOutput, None),
178            row(3, 2, EventKind::AgentMessage, None),
179        ];
180        assign_steps(&mut rows);
181        // terminal_output is not a step; the two semantic events are 1 and 2.
182        assert_eq!(steps(&rows), vec![Some(1), None, Some(2)]);
183    }
184
185    #[test]
186    fn assign_steps_multiple_pairs() {
187        // Two independent call/result pairs → 2 distinct steps.
188        let mut rows = vec![
189            row(1, 0, EventKind::ToolCall, None),
190            row(2, 1, EventKind::ToolResult, Some(1)),
191            row(3, 10, EventKind::ToolCall, None),
192            row(4, 11, EventKind::ToolResult, Some(3)),
193        ];
194        assign_steps(&mut rows);
195        assert_eq!(steps(&rows), vec![Some(1), Some(1), Some(2), Some(2)]);
196    }
197
198    #[test]
199    fn assign_steps_thinking_own_step() {
200        let mut rows = vec![
201            row(1, 0, EventKind::Thinking, None),
202            row(2, 5, EventKind::AgentMessage, None),
203        ];
204        assign_steps(&mut rows);
205        assert_eq!(steps(&rows), vec![Some(1), Some(2)]);
206    }
207
208    #[test]
209    fn assign_steps_idempotent() {
210        let mut rows = vec![
211            row(1, 0, EventKind::ToolCall, None),
212            row(2, 5, EventKind::ToolResult, Some(1)),
213            row(3, 6, EventKind::AgentMessage, None),
214        ];
215        assign_steps(&mut rows);
216        let first = steps(&rows);
217        assign_steps(&mut rows);
218        assert_eq!(steps(&rows), first, "re-running must be idempotent");
219        assert_eq!(first, vec![Some(1), Some(1), Some(2)]);
220    }
221
222    #[test]
223    fn assign_steps_result_pointing_at_non_call_present_id_gets_own_step() {
224        // correlates points to an id that IS present but is not a ToolCall —
225        // must not defer (only tool_call targets share); gets its own step.
226        let mut rows = vec![
227            row(1, 0, EventKind::UserMessage, None),
228            row(2, 5, EventKind::ToolResult, Some(1)),
229        ];
230        assign_steps(&mut rows);
231        assert_eq!(steps(&rows), vec![Some(1), Some(2)]);
232    }
233}