kaizen/store/
span_tree.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2//! In-memory span tree assembled from flat `ToolSpanView` rows.
3
4use crate::metrics::types::ToolSpanView;
5use serde::{Deserialize, Serialize};
6use std::collections::HashMap;
7
8#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct SpanNode {
10    pub span: ToolSpanView,
11    pub children: Vec<SpanNode>,
12    pub subtree_cost_usd_e6: i64,
13    pub subtree_token_count: u64,
14}
15
16/// Assemble a forest of `SpanNode` from a flat ordered list.
17pub fn build_tree(spans: Vec<ToolSpanView>) -> Vec<SpanNode> {
18    let ids: Vec<String> = spans.iter().map(|s| s.span_id.clone()).collect();
19    let mut nodes: HashMap<String, SpanNode> = spans
20        .into_iter()
21        .map(|s| {
22            let cost = s.subtree_cost_usd_e6.unwrap_or(0);
23            let tokens = s.subtree_token_count.unwrap_or(0) as u64;
24            (
25                s.span_id.clone(),
26                SpanNode {
27                    span: s,
28                    children: vec![],
29                    subtree_cost_usd_e6: cost,
30                    subtree_token_count: tokens,
31                },
32            )
33        })
34        .collect();
35    let mut roots: Vec<String> = Vec::new();
36    for id in &ids {
37        let pid = nodes[id].span.parent_span_id.clone();
38        match pid {
39            Some(p) if nodes.contains_key(&p) => {
40                let child = nodes.remove(id).expect("id present");
41                nodes
42                    .get_mut(&p)
43                    .expect("parent present")
44                    .children
45                    .push(child);
46            }
47            _ => roots.push(id.clone()),
48        }
49    }
50    roots
51        .into_iter()
52        .filter_map(|id| nodes.remove(&id))
53        .collect()
54}
kaizen/store/span_tree.rs

kaizen/store/
span_tree.rs
