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}
47
48/// A directed edge in the agent graph.
49#[derive(Debug, Clone, Serialize, Deserialize)]
50pub struct AgentEdge {
51    pub from_instance_id: String,
52    pub to_instance_id: String,
53    pub edge_type: AgentEdgeType,
54    pub timestamp: String,
55    #[serde(default, skip_serializing_if = "Vec::is_empty")]
56    pub artifacts: Vec<String>,
57}
58
59/// The complete agent collaboration graph for a session.
60#[derive(Debug, Clone, Default, Serialize, Deserialize)]
61pub struct AgentGraph {
62    pub nodes: Vec<AgentNode>,
63    pub edges: Vec<AgentEdge>,
64}
65
66impl AgentGraph {
67    /// Build an agent graph from a sequence of session events.
68    pub fn from_events(events: &[SessionEvent]) -> Self {
69        let mut nodes_map: BTreeMap<String, AgentNode> = BTreeMap::new();
70        let mut edges: Vec<AgentEdge> = Vec::new();
71        let mut parent_map: BTreeMap<String, String> = BTreeMap::new(); // child -> parent instance
72
73        for event in events {
74            let instance_id = &event.agent_instance_id;
75
76            // Ensure node exists
77            let node = nodes_map.entry(instance_id.clone()).or_insert_with(|| AgentNode {
78                agent_id: event.agent_id.clone(),
79                agent_instance_id: instance_id.clone(),
80                agent_name: event.agent_name.clone(),
81                agent_role: event.agent_role.clone(),
82                host_id: event.host_id.clone(),
83                started_at: None,
84                completed_at: None,
85                status: None,
86                depth: 0,
87                tool_calls: 0,
88            });
89
90            match &event.event_type {
91                EventType::AgentStarted { parent_agent_instance_id } => {
92                    node.started_at = Some(event.timestamp.clone());
93                    if let Some(parent_id) = parent_agent_instance_id {
94                        parent_map.insert(instance_id.clone(), parent_id.clone());
95                    }
96                }
97
98                EventType::AgentSpawned { spawned_by_agent_instance_id, .. } => {
99                    node.started_at = Some(event.timestamp.clone());
100                    parent_map.insert(instance_id.clone(), spawned_by_agent_instance_id.clone());
101                    edges.push(AgentEdge {
102                        from_instance_id: spawned_by_agent_instance_id.clone(),
103                        to_instance_id: instance_id.clone(),
104                        edge_type: AgentEdgeType::ParentChild,
105                        timestamp: event.timestamp.clone(),
106                        artifacts: Vec::new(),
107                    });
108                }
109
110                EventType::AgentHandoff { from_agent_instance_id, to_agent_instance_id, artifacts } => {
111                    edges.push(AgentEdge {
112                        from_instance_id: from_agent_instance_id.clone(),
113                        to_instance_id: to_agent_instance_id.clone(),
114                        edge_type: AgentEdgeType::Handoff,
115                        timestamp: event.timestamp.clone(),
116                        artifacts: artifacts.clone(),
117                    });
118                    // Ensure the target node exists
119                    nodes_map.entry(to_agent_instance_id.clone()).or_insert_with(|| AgentNode {
120                        agent_id: String::new(),
121                        agent_instance_id: to_agent_instance_id.clone(),
122                        agent_name: String::new(),
123                        agent_role: None,
124                        host_id: event.host_id.clone(),
125                        started_at: None,
126                        completed_at: None,
127                        status: None,
128                        depth: 0,
129                        tool_calls: 0,
130                    });
131                }
132
133                EventType::AgentCollaborated { collaborator_agent_instance_ids } => {
134                    for collab_id in collaborator_agent_instance_ids {
135                        edges.push(AgentEdge {
136                            from_instance_id: instance_id.clone(),
137                            to_instance_id: collab_id.clone(),
138                            edge_type: AgentEdgeType::Collaboration,
139                            timestamp: event.timestamp.clone(),
140                            artifacts: Vec::new(),
141                        });
142                    }
143                }
144
145                EventType::AgentReturned { returned_to_agent_instance_id } => {
146                    edges.push(AgentEdge {
147                        from_instance_id: instance_id.clone(),
148                        to_instance_id: returned_to_agent_instance_id.clone(),
149                        edge_type: AgentEdgeType::Return,
150                        timestamp: event.timestamp.clone(),
151                        artifacts: Vec::new(),
152                    });
153                }
154
155                EventType::AgentCompleted { .. } => {
156                    node.completed_at = Some(event.timestamp.clone());
157                    node.status = Some("completed".into());
158                }
159
160                EventType::AgentFailed { .. } => {
161                    node.completed_at = Some(event.timestamp.clone());
162                    node.status = Some("failed".into());
163                }
164
165                EventType::AgentCalledTool { .. } => {
166                    node.tool_calls += 1;
167                }
168
169                _ => {}
170            }
171        }
172
173        // Compute depths from parent map
174        let mut depth_cache: BTreeMap<String, u32> = BTreeMap::new();
175        let instances: Vec<String> = nodes_map.keys().cloned().collect();
176        for inst in &instances {
177            let depth = compute_depth(inst, &parent_map, &mut depth_cache);
178            if let Some(node) = nodes_map.get_mut(inst) {
179                node.depth = depth;
180            }
181        }
182
183        let nodes: Vec<AgentNode> = nodes_map.into_values().collect();
184
185        AgentGraph { nodes, edges }
186    }
187
188    /// Return the maximum depth in the graph.
189    pub fn max_depth(&self) -> u32 {
190        self.nodes.iter().map(|n| n.depth).max().unwrap_or(0)
191    }
192
193    /// Return the set of unique host IDs across all agents.
194    pub fn host_ids(&self) -> BTreeSet<String> {
195        self.nodes.iter().map(|n| n.host_id.clone()).collect()
196    }
197
198    /// Total number of handoff edges.
199    pub fn handoff_count(&self) -> u32 {
200        self.edges.iter()
201            .filter(|e| e.edge_type == AgentEdgeType::Handoff)
202            .count() as u32
203    }
204
205    /// Total number of spawn (parent-child) edges.
206    pub fn spawn_count(&self) -> u32 {
207        self.edges.iter()
208            .filter(|e| e.edge_type == AgentEdgeType::ParentChild)
209            .count() as u32
210    }
211}
212
213fn compute_depth(
214    instance_id: &str,
215    parent_map: &BTreeMap<String, String>,
216    cache: &mut BTreeMap<String, u32>,
217) -> u32 {
218    if let Some(&d) = cache.get(instance_id) {
219        return d;
220    }
221    let depth = match parent_map.get(instance_id) {
222        Some(parent) => 1 + compute_depth(parent, parent_map, cache),
223        None => 0,
224    };
225    cache.insert(instance_id.to_string(), depth);
226    depth
227}
228
229#[cfg(test)]
230mod tests {
231    use super::*;
232    use crate::session::event::*;
233
234    fn evt(instance_id: &str, host: &str, event_type: EventType) -> SessionEvent {
235        SessionEvent {
236            session_id: "ssn_001".into(),
237            event_id: generate_event_id(),
238            timestamp: "2026-04-05T08:00:00Z".into(),
239            sequence_no: 0,
240            trace_id: "trace_1".into(),
241            span_id: generate_span_id(),
242            parent_span_id: None,
243            agent_id: format!("agent://{instance_id}"),
244            agent_instance_id: instance_id.into(),
245            agent_name: instance_id.into(),
246            agent_role: None,
247            host_id: host.into(),
248            tool_runtime_id: None,
249            event_type,
250            artifact_ref: None,
251            meta: None,
252        }
253    }
254
255    #[test]
256    fn builds_graph_from_spawn_and_handoff() {
257        let events = vec![
258            evt("root", "host_a", EventType::AgentStarted {
259                parent_agent_instance_id: None,
260            }),
261            evt("child1", "host_a", EventType::AgentSpawned {
262                spawned_by_agent_instance_id: "root".into(),
263                reason: Some("review code".into()),
264            }),
265            evt("child2", "host_b", EventType::AgentSpawned {
266                spawned_by_agent_instance_id: "root".into(),
267                reason: None,
268            }),
269            evt("root", "host_a", EventType::AgentHandoff {
270                from_agent_instance_id: "root".into(),
271                to_agent_instance_id: "child1".into(),
272                artifacts: vec!["art_001".into()],
273            }),
274            evt("child1", "host_a", EventType::AgentCompleted {
275                termination_reason: None,
276            }),
277        ];
278
279        let graph = AgentGraph::from_events(&events);
280        assert_eq!(graph.nodes.len(), 3);
281        assert_eq!(graph.max_depth(), 1);
282        assert_eq!(graph.handoff_count(), 1);
283        assert_eq!(graph.spawn_count(), 2);
284        assert_eq!(graph.host_ids().len(), 2);
285    }
286
287    #[test]
288    fn nested_depth() {
289        let events = vec![
290            evt("root", "h", EventType::AgentStarted { parent_agent_instance_id: None }),
291            evt("l1", "h", EventType::AgentSpawned { spawned_by_agent_instance_id: "root".into(), reason: None }),
292            evt("l2", "h", EventType::AgentSpawned { spawned_by_agent_instance_id: "l1".into(), reason: None }),
293            evt("l3", "h", EventType::AgentSpawned { spawned_by_agent_instance_id: "l2".into(), reason: None }),
294        ];
295
296        let graph = AgentGraph::from_events(&events);
297        assert_eq!(graph.max_depth(), 3);
298        let l3 = graph.nodes.iter().find(|n| n.agent_instance_id == "l3").unwrap();
299        assert_eq!(l3.depth, 3);
300    }
301}
treeship_core/session/graph.rs

treeship_core/session/
graph.rs
