Skip to main content

lago_core/
tool_span.rs

1use crate::event::SpanStatus;
2use serde::{Deserialize, Serialize};
3
4/// Represents a tool execution span (invoke -> result).
5#[derive(Debug, Clone, Serialize, Deserialize)]
6pub struct ToolSpan {
7    pub call_id: String,
8    pub tool_name: String,
9    pub arguments: serde_json::Value,
10    pub category: Option<String>,
11    pub status: Option<SpanStatus>,
12    pub result: Option<serde_json::Value>,
13    pub started_at: u64,
14    pub ended_at: Option<u64>,
15    pub duration_ms: Option<u64>,
16}
17
18impl ToolSpan {
19    pub fn new(call_id: String, tool_name: String, arguments: serde_json::Value) -> Self {
20        Self {
21            call_id,
22            tool_name,
23            arguments,
24            category: None,
25            status: None,
26            result: None,
27            started_at: crate::event::EventEnvelope::now_micros(),
28            ended_at: None,
29            duration_ms: None,
30        }
31    }
32
33    pub fn is_complete(&self) -> bool {
34        self.status.is_some()
35    }
36}
37
38#[cfg(test)]
39mod tests {
40    use super::*;
41    use crate::event::SpanStatus;
42
43    #[test]
44    fn new_span_is_incomplete() {
45        let span = ToolSpan::new(
46            "call-1".into(),
47            "read_file".into(),
48            serde_json::json!({"path": "/foo"}),
49        );
50        assert!(!span.is_complete());
51        assert_eq!(span.call_id, "call-1");
52        assert_eq!(span.tool_name, "read_file");
53        assert!(span.started_at > 0);
54        assert!(span.ended_at.is_none());
55        assert!(span.result.is_none());
56        assert!(span.category.is_none());
57    }
58
59    #[test]
60    fn completed_span() {
61        let mut span = ToolSpan::new("call-2".into(), "write".into(), serde_json::json!({}));
62        span.status = Some(SpanStatus::Ok);
63        span.result = Some(serde_json::json!({"written": true}));
64        span.ended_at = Some(span.started_at + 100_000);
65        span.duration_ms = Some(100);
66        assert!(span.is_complete());
67    }
68
69    #[test]
70    fn span_serde_roundtrip() {
71        let span = ToolSpan::new(
72            "call-3".into(),
73            "exec".into(),
74            serde_json::json!({"cmd": "ls"}),
75        );
76        let json = serde_json::to_string(&span).unwrap();
77        let back: ToolSpan = serde_json::from_str(&json).unwrap();
78        assert_eq!(back.call_id, "call-3");
79        assert_eq!(back.tool_name, "exec");
80    }
81}