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                EventType::AgentCalledTool { .. } => {
186                    node.tool_calls += 1;
187                }
188
189                EventType::AgentCompletedProcess { .. } => {
190                    node.tool_calls += 1;
191                }
192
193                EventType::AgentDecision { ref model, tokens_in, tokens_out, ref provider, .. } => {
194                    if let Some(ref m) = model {
195                        node.model = Some(m.clone());
196                    }
197                    if let Some(ref p) = provider {
198                        node.provider = Some(p.clone());
199                    }
200                    if let Some(t) = tokens_in { node.tokens_in += t; }
201                    if let Some(t) = tokens_out { node.tokens_out += t; }
202                }
203
204                _ => {}
205            }
206        }
207
208        // Compute depths from parent map
209        let mut depth_cache: BTreeMap<String, u32> = BTreeMap::new();
210        let instances: Vec<String> = nodes_map.keys().cloned().collect();
211        for inst in &instances {
212            let depth = compute_depth(inst, &parent_map, &mut depth_cache);
213            if let Some(node) = nodes_map.get_mut(inst) {
214                node.depth = depth;
215            }
216        }
217
218        let nodes: Vec<AgentNode> = nodes_map.into_values().collect();
219
220        AgentGraph { nodes, edges }
221    }
222
223    /// Return the maximum depth in the graph.
224    pub fn max_depth(&self) -> u32 {
225        self.nodes.iter().map(|n| n.depth).max().unwrap_or(0)
226    }
227
228    /// Return the set of unique host IDs across all agents.
229    pub fn host_ids(&self) -> BTreeSet<String> {
230        self.nodes.iter().map(|n| n.host_id.clone()).collect()
231    }
232
233    /// Total number of handoff edges.
234    pub fn handoff_count(&self) -> u32 {
235        self.edges.iter()
236            .filter(|e| e.edge_type == AgentEdgeType::Handoff)
237            .count() as u32
238    }
239
240    /// Total number of spawn (parent-child) edges.
241    pub fn spawn_count(&self) -> u32 {
242        self.edges.iter()
243            .filter(|e| e.edge_type == AgentEdgeType::ParentChild)
244            .count() as u32
245    }
246}
247
248fn compute_depth(
249    instance_id: &str,
250    parent_map: &BTreeMap<String, String>,
251    cache: &mut BTreeMap<String, u32>,
252) -> u32 {
253    if let Some(&d) = cache.get(instance_id) {
254        return d;
255    }
256    let depth = match parent_map.get(instance_id) {
257        Some(parent) => 1 + compute_depth(parent, parent_map, cache),
258        None => 0,
259    };
260    cache.insert(instance_id.to_string(), depth);
261    depth
262}
263
264#[cfg(test)]
265mod tests {
266    use super::*;
267    use crate::session::event::*;
268
269    fn evt(instance_id: &str, host: &str, event_type: EventType) -> SessionEvent {
270        SessionEvent {
271            session_id: "ssn_001".into(),
272            event_id: generate_event_id(),
273            timestamp: "2026-04-05T08:00:00Z".into(),
274            sequence_no: 0,
275            trace_id: "trace_1".into(),
276            span_id: generate_span_id(),
277            parent_span_id: None,
278            agent_id: format!("agent://{instance_id}"),
279            agent_instance_id: instance_id.into(),
280            agent_name: instance_id.into(),
281            agent_role: None,
282            host_id: host.into(),
283            tool_runtime_id: None,
284            event_type,
285            artifact_ref: None,
286            meta: None,
287        }
288    }
289
290    #[test]
291    fn builds_graph_from_spawn_and_handoff() {
292        let events = vec![
293            evt("root", "host_a", EventType::AgentStarted {
294                parent_agent_instance_id: None,
295            }),
296            evt("child1", "host_a", EventType::AgentSpawned {
297                spawned_by_agent_instance_id: "root".into(),
298                reason: Some("review code".into()),
299            }),
300            evt("child2", "host_b", EventType::AgentSpawned {
301                spawned_by_agent_instance_id: "root".into(),
302                reason: None,
303            }),
304            evt("root", "host_a", EventType::AgentHandoff {
305                from_agent_instance_id: "root".into(),
306                to_agent_instance_id: "child1".into(),
307                artifacts: vec!["art_001".into()],
308            }),
309            evt("child1", "host_a", EventType::AgentCompleted {
310                termination_reason: None,
311            }),
312        ];
313
314        let graph = AgentGraph::from_events(&events);
315        assert_eq!(graph.nodes.len(), 3);
316        assert_eq!(graph.max_depth(), 1);
317        assert_eq!(graph.handoff_count(), 1);
318        assert_eq!(graph.spawn_count(), 2);
319        assert_eq!(graph.host_ids().len(), 2);
320    }
321
322    #[test]
323    fn nested_depth() {
324        let events = vec![
325            evt("root", "h", EventType::AgentStarted { parent_agent_instance_id: None }),
326            evt("l1", "h", EventType::AgentSpawned { spawned_by_agent_instance_id: "root".into(), reason: None }),
327            evt("l2", "h", EventType::AgentSpawned { spawned_by_agent_instance_id: "l1".into(), reason: None }),
328            evt("l3", "h", EventType::AgentSpawned { spawned_by_agent_instance_id: "l2".into(), reason: None }),
329        ];
330
331        let graph = AgentGraph::from_events(&events);
332        assert_eq!(graph.max_depth(), 3);
333        let l3 = graph.nodes.iter().find(|n| n.agent_instance_id == "l3").unwrap();
334        assert_eq!(l3.depth, 3);
335    }
336}
treeship_core/session/graph.rs

treeship_core/session/
graph.rs
