treeship_core/session/
graph.rs

1//! Agent collaboration graph built from session events.
2//!
3//! Captures the full topology of agent relationships: parent-child spawning,
4//! handoffs, and collaboration edges.
5
6use std::collections::{BTreeMap, BTreeSet};
7
8use serde::{Deserialize, Serialize};
9
10use super::event::{EventType, SessionEvent};
11
12/// Type of relationship between two agents.
13#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
14#[serde(rename_all = "snake_case")]
15pub enum AgentEdgeType {
16    /// Parent spawned a child agent.
17    ParentChild,
18    /// Work was handed off from one agent to another.
19    Handoff,
20    /// Agents collaborated on a shared task.
21    Collaboration,
22    /// Agent returned control to a parent.
23    Return,
24}
25
26/// A node in the agent graph representing one agent instance.
27#[derive(Debug, Clone, Serialize, Deserialize)]
28pub struct AgentNode {
29    pub agent_id: String,
30    pub agent_instance_id: String,
31    pub agent_name: String,
32    #[serde(skip_serializing_if = "Option::is_none")]
33    pub agent_role: Option<String>,
34    pub host_id: String,
35    #[serde(skip_serializing_if = "Option::is_none")]
36    pub started_at: Option<String>,
37    #[serde(skip_serializing_if = "Option::is_none")]
38    pub completed_at: Option<String>,
39    #[serde(skip_serializing_if = "Option::is_none")]
40    pub status: Option<String>,
41    #[serde(default)]
42    pub depth: u32,
43    /// Number of tool calls made by this agent.
44    #[serde(default)]
45    pub tool_calls: u32,
46    /// Model identifier (e.g. "claude-opus-4-6"). Populated from decision events.
47    #[serde(default, skip_serializing_if = "Option::is_none")]
48    pub model: Option<String>,
49    /// Cumulative input tokens across all decisions by this agent.
50    #[serde(default)]
51    pub tokens_in: u64,
52    /// Cumulative output tokens across all decisions by this agent.
53    #[serde(default)]
54    pub tokens_out: u64,
55    /// Provider e.g. "anthropic", "openrouter", "bedrock"
56    #[serde(default, skip_serializing_if = "Option::is_none")]
57    pub provider: Option<String>,
58}
59
60/// A directed edge in the agent graph.
61#[derive(Debug, Clone, Serialize, Deserialize)]
62pub struct AgentEdge {
63    pub from_instance_id: String,
64    pub to_instance_id: String,
65    pub edge_type: AgentEdgeType,
66    pub timestamp: String,
67    #[serde(default, skip_serializing_if = "Vec::is_empty")]
68    pub artifacts: Vec<String>,
69}
70
71/// The complete agent collaboration graph for a session.
72#[derive(Debug, Clone, Default, Serialize, Deserialize)]
73pub struct AgentGraph {
74    pub nodes: Vec<AgentNode>,
75    pub edges: Vec<AgentEdge>,
76}
77
78impl AgentGraph {
79    /// Build an agent graph from a sequence of session events.
80    pub fn from_events(events: &[SessionEvent]) -> Self {
81        let mut nodes_map: BTreeMap<String, AgentNode> = BTreeMap::new();
82        let mut edges: Vec<AgentEdge> = Vec::new();
83        let mut parent_map: BTreeMap<String, String> = BTreeMap::new(); // child -> parent instance
84
85        for event in events {
86            let instance_id = &event.agent_instance_id;
87
88            // Ensure node exists
89            let node = nodes_map.entry(instance_id.clone()).or_insert_with(|| AgentNode {
90                agent_id: event.agent_id.clone(),
91                agent_instance_id: instance_id.clone(),
92                agent_name: event.agent_name.clone(),
93                agent_role: event.agent_role.clone(),
94                host_id: event.host_id.clone(),
95                started_at: None,
96                completed_at: None,
97                status: None,
98                depth: 0,
99                tool_calls: 0,
100                model: None,
101                tokens_in: 0,
102                tokens_out: 0,
103                provider: None,
104            });
105
106            match &event.event_type {
107                EventType::AgentStarted { parent_agent_instance_id } => {
108                    node.started_at = Some(event.timestamp.clone());
109                    if let Some(parent_id) = parent_agent_instance_id {
110                        parent_map.insert(instance_id.clone(), parent_id.clone());
111                    }
112                }
113
114                EventType::AgentSpawned { spawned_by_agent_instance_id, .. } => {
115                    node.started_at = Some(event.timestamp.clone());
116                    parent_map.insert(instance_id.clone(), spawned_by_agent_instance_id.clone());
117                    edges.push(AgentEdge {
118                        from_instance_id: spawned_by_agent_instance_id.clone(),
119                        to_instance_id: instance_id.clone(),
120                        edge_type: AgentEdgeType::ParentChild,
121                        timestamp: event.timestamp.clone(),
122                        artifacts: Vec::new(),
123                    });
124                }
125
126                EventType::AgentHandoff { from_agent_instance_id, to_agent_instance_id, artifacts } => {
127                    edges.push(AgentEdge {
128                        from_instance_id: from_agent_instance_id.clone(),
129                        to_instance_id: to_agent_instance_id.clone(),
130                        edge_type: AgentEdgeType::Handoff,
131                        timestamp: event.timestamp.clone(),
132                        artifacts: artifacts.clone(),
133                    });
134                    // Ensure the target node exists
135                    nodes_map.entry(to_agent_instance_id.clone()).or_insert_with(|| AgentNode {
136                        agent_id: String::new(),
137                        agent_instance_id: to_agent_instance_id.clone(),
138                        agent_name: String::new(),
139                        agent_role: None,
140                        host_id: event.host_id.clone(),
141                        started_at: None,
142                        completed_at: None,
143                        status: None,
144                        depth: 0,
145                        tool_calls: 0,
146                        model: None,
147                        tokens_in: 0,
148                        tokens_out: 0,
149                        provider: None,
150                    });
151                }
152
153                EventType::AgentCollaborated { collaborator_agent_instance_ids } => {
154                    for collab_id in collaborator_agent_instance_ids {
155                        edges.push(AgentEdge {
156                            from_instance_id: instance_id.clone(),
157                            to_instance_id: collab_id.clone(),
158                            edge_type: AgentEdgeType::Collaboration,
159                            timestamp: event.timestamp.clone(),
160                            artifacts: Vec::new(),
161                        });
162                    }
163                }
164
165                EventType::AgentReturned { returned_to_agent_instance_id } => {
166                    edges.push(AgentEdge {
167                        from_instance_id: instance_id.clone(),
168                        to_instance_id: returned_to_agent_instance_id.clone(),
169                        edge_type: AgentEdgeType::Return,
170                        timestamp: event.timestamp.clone(),
171                        artifacts: Vec::new(),
172                    });
173                }
174
175                EventType::AgentCompleted { .. } => {
176                    node.completed_at = Some(event.timestamp.clone());
177                    node.status = Some("completed".into());
178                }
179
180                EventType::AgentFailed { .. } => {
181                    node.completed_at = Some(event.timestamp.clone());
182                    node.status = Some("failed".into());
183                }
184
185                // node.tool_calls counts every action the agent took. The
186                // side-effects ledger then groups those actions by category
187                // (files_read, files_written, processes, network_connections,
188                // ports_opened, tool_invocations) -- but the per-agent total
189                // here is the cardinal count.
190                //
191                // History note: prior to v0.9.5 only AgentCalledTool and
192                // AgentCompletedProcess were counted, which made the per-agent
193                // count drop to near-zero when the Claude Code plugin started
194                // emitting specialized event types (agent.read_file, etc.).
195                // The sum of node.tool_calls across all agents (used as
196                // "Actions" in the local preview and as the headline tool
197                // count on treeship.dev/receipt/<id>) was undercounting the
198                // agent's actual activity. Adding the four file/network/port
199                // event types here fixes the counter for all consumers in
200                // one place; renderers don't need to compute the total
201                // themselves.
202                EventType::AgentCalledTool { .. } => {
203                    node.tool_calls += 1;
204                }
205
206                EventType::AgentCompletedProcess { .. } => {
207                    node.tool_calls += 1;
208                }
209
210                EventType::AgentReadFile { .. } => {
211                    node.tool_calls += 1;
212                }
213
214                EventType::AgentWroteFile { .. } => {
215                    node.tool_calls += 1;
216                }
217
218                EventType::AgentConnectedNetwork { .. } => {
219                    node.tool_calls += 1;
220                }
221
222                EventType::AgentOpenedPort { .. } => {
223                    node.tool_calls += 1;
224                }
225
226                EventType::AgentDecision { ref model, tokens_in, tokens_out, ref provider, .. } => {
227                    if let Some(ref m) = model {
228                        node.model = Some(m.clone());
229                    }
230                    if let Some(ref p) = provider {
231                        node.provider = Some(p.clone());
232                    }
233                    if let Some(t) = tokens_in { node.tokens_in += t; }
234                    if let Some(t) = tokens_out { node.tokens_out += t; }
235                }
236
237                _ => {}
238            }
239        }
240
241        // Compute depths from parent map
242        let mut depth_cache: BTreeMap<String, u32> = BTreeMap::new();
243        let instances: Vec<String> = nodes_map.keys().cloned().collect();
244        for inst in &instances {
245            let depth = compute_depth(inst, &parent_map, &mut depth_cache);
246            if let Some(node) = nodes_map.get_mut(inst) {
247                node.depth = depth;
248            }
249        }
250
251        let nodes: Vec<AgentNode> = nodes_map.into_values().collect();
252
253        AgentGraph { nodes, edges }
254    }
255
256    /// Return the maximum depth in the graph.
257    pub fn max_depth(&self) -> u32 {
258        self.nodes.iter().map(|n| n.depth).max().unwrap_or(0)
259    }
260
261    /// Return the set of unique host IDs across all agents.
262    pub fn host_ids(&self) -> BTreeSet<String> {
263        self.nodes.iter().map(|n| n.host_id.clone()).collect()
264    }
265
266    /// Total number of handoff edges.
267    pub fn handoff_count(&self) -> u32 {
268        self.edges.iter()
269            .filter(|e| e.edge_type == AgentEdgeType::Handoff)
270            .count() as u32
271    }
272
273    /// Total number of spawn (parent-child) edges.
274    pub fn spawn_count(&self) -> u32 {
275        self.edges.iter()
276            .filter(|e| e.edge_type == AgentEdgeType::ParentChild)
277            .count() as u32
278    }
279}
280
281fn compute_depth(
282    instance_id: &str,
283    parent_map: &BTreeMap<String, String>,
284    cache: &mut BTreeMap<String, u32>,
285) -> u32 {
286    if let Some(&d) = cache.get(instance_id) {
287        return d;
288    }
289    let depth = match parent_map.get(instance_id) {
290        Some(parent) => 1 + compute_depth(parent, parent_map, cache),
291        None => 0,
292    };
293    cache.insert(instance_id.to_string(), depth);
294    depth
295}
296
297#[cfg(test)]
298mod tests {
299    use super::*;
300    use crate::session::event::*;
301
302    fn evt(instance_id: &str, host: &str, event_type: EventType) -> SessionEvent {
303        SessionEvent {
304            session_id: "ssn_001".into(),
305            event_id: generate_event_id(),
306            timestamp: "2026-04-05T08:00:00Z".into(),
307            sequence_no: 0,
308            trace_id: "trace_1".into(),
309            span_id: generate_span_id(),
310            parent_span_id: None,
311            agent_id: format!("agent://{instance_id}"),
312            agent_instance_id: instance_id.into(),
313            agent_name: instance_id.into(),
314            agent_role: None,
315            host_id: host.into(),
316            tool_runtime_id: None,
317            event_type,
318            artifact_ref: None,
319            meta: None,
320        }
321    }
322
323    #[test]
324    fn builds_graph_from_spawn_and_handoff() {
325        let events = vec![
326            evt("root", "host_a", EventType::AgentStarted {
327                parent_agent_instance_id: None,
328            }),
329            evt("child1", "host_a", EventType::AgentSpawned {
330                spawned_by_agent_instance_id: "root".into(),
331                reason: Some("review code".into()),
332            }),
333            evt("child2", "host_b", EventType::AgentSpawned {
334                spawned_by_agent_instance_id: "root".into(),
335                reason: None,
336            }),
337            evt("root", "host_a", EventType::AgentHandoff {
338                from_agent_instance_id: "root".into(),
339                to_agent_instance_id: "child1".into(),
340                artifacts: vec!["art_001".into()],
341            }),
342            evt("child1", "host_a", EventType::AgentCompleted {
343                termination_reason: None,
344            }),
345        ];
346
347        let graph = AgentGraph::from_events(&events);
348        assert_eq!(graph.nodes.len(), 3);
349        assert_eq!(graph.max_depth(), 1);
350        assert_eq!(graph.handoff_count(), 1);
351        assert_eq!(graph.spawn_count(), 2);
352        assert_eq!(graph.host_ids().len(), 2);
353    }
354
355    #[test]
356    fn nested_depth() {
357        let events = vec![
358            evt("root", "h", EventType::AgentStarted { parent_agent_instance_id: None }),
359            evt("l1", "h", EventType::AgentSpawned { spawned_by_agent_instance_id: "root".into(), reason: None }),
360            evt("l2", "h", EventType::AgentSpawned { spawned_by_agent_instance_id: "l1".into(), reason: None }),
361            evt("l3", "h", EventType::AgentSpawned { spawned_by_agent_instance_id: "l2".into(), reason: None }),
362        ];
363
364        let graph = AgentGraph::from_events(&events);
365        assert_eq!(graph.max_depth(), 3);
366        let l3 = graph.nodes.iter().find(|n| n.agent_instance_id == "l3").unwrap();
367        assert_eq!(l3.depth, 3);
368    }
369
370    /// Regression test: per-agent `tool_calls` must count every action event
371    /// type the agent emits, not just `AgentCalledTool` and `AgentCompletedProcess`.
372    ///
373    /// Pre-v0.9.5, this counter ignored AgentReadFile, AgentWroteFile,
374    /// AgentConnectedNetwork, and AgentOpenedPort. As soon as the Claude Code
375    /// plugin started emitting those specialized event types (also v0.9.5),
376    /// the per-agent count -- and the `nodes.reduce(...)` total used by the
377    /// receipt renderer -- collapsed to near-zero even when the agent had
378    /// done substantial work. The fix lives in the EventType match arm above.
379    #[test]
380    fn tool_calls_counts_every_action_event_type() {
381        let events = vec![
382            evt("a", "h", EventType::AgentStarted { parent_agent_instance_id: None }),
383            evt("a", "h", EventType::AgentCalledTool {
384                tool_name: "Glob".into(),
385                tool_input_digest: None,
386                tool_output_digest: None,
387                duration_ms: None,
388            }),
389            evt("a", "h", EventType::AgentReadFile {
390                file_path: "src/foo.rs".into(),
391                digest: None,
392            }),
393            evt("a", "h", EventType::AgentWroteFile {
394                file_path: "src/bar.rs".into(),
395                digest: None,
396                operation: None,
397                additions: None,
398                deletions: None,
399            }),
400            evt("a", "h", EventType::AgentCompletedProcess {
401                process_name: "npm test".into(),
402                exit_code: Some(0),
403                duration_ms: Some(2_500),
404                command: None,
405            }),
406            evt("a", "h", EventType::AgentConnectedNetwork {
407                destination: "api.github.com".into(),
408                port: None,
409            }),
410            evt("a", "h", EventType::AgentOpenedPort {
411                port: 3000,
412                protocol: Some("tcp".into()),
413            }),
414        ];
415
416        let graph = AgentGraph::from_events(&events);
417        let agent_a = graph.nodes.iter().find(|n| n.agent_instance_id == "a").unwrap();
418
419        // 6 action events (Glob, ReadFile, WroteFile, CompletedProcess,
420        // ConnectedNetwork, OpenedPort). AgentStarted is not an action.
421        assert_eq!(
422            agent_a.tool_calls, 6,
423            "tool_calls must count all action event types (was {}, expected 6)",
424            agent_a.tool_calls
425        );
426    }
427}
treeship_core/session/graph.rs

treeship_core/session/
graph.rs
