hehe_agent/
event.rs

1use hehe_core::Id;
2use serde::{Deserialize, Serialize};
3
4#[derive(Clone, Debug, Serialize, Deserialize)]
5#[serde(tag = "type", rename_all = "snake_case")]
6pub enum AgentEvent {
7    MessageStart {
8        session_id: Id,
9    },
10
11    TextDelta {
12        delta: String,
13    },
14
15    TextComplete {
16        text: String,
17    },
18
19    ToolUseStart {
20        id: String,
21        name: String,
22        input: serde_json::Value,
23    },
24
25    ToolUseEnd {
26        id: String,
27        output: String,
28        is_error: bool,
29    },
30
31    Thinking {
32        content: String,
33    },
34
35    MessageEnd {
36        session_id: Id,
37    },
38
39    Error {
40        message: String,
41    },
42}
43
44impl AgentEvent {
45    pub fn message_start(session_id: Id) -> Self {
46        Self::MessageStart { session_id }
47    }
48
49    pub fn text_delta(delta: impl Into<String>) -> Self {
50        Self::TextDelta {
51            delta: delta.into(),
52        }
53    }
54
55    pub fn text_complete(text: impl Into<String>) -> Self {
56        Self::TextComplete { text: text.into() }
57    }
58
59    pub fn tool_use_start(id: impl Into<String>, name: impl Into<String>, input: serde_json::Value) -> Self {
60        Self::ToolUseStart {
61            id: id.into(),
62            name: name.into(),
63            input,
64        }
65    }
66
67    pub fn tool_use_end(id: impl Into<String>, output: impl Into<String>, is_error: bool) -> Self {
68        Self::ToolUseEnd {
69            id: id.into(),
70            output: output.into(),
71            is_error,
72        }
73    }
74
75    pub fn thinking(content: impl Into<String>) -> Self {
76        Self::Thinking {
77            content: content.into(),
78        }
79    }
80
81    pub fn message_end(session_id: Id) -> Self {
82        Self::MessageEnd { session_id }
83    }
84
85    pub fn error(message: impl Into<String>) -> Self {
86        Self::Error {
87            message: message.into(),
88        }
89    }
90
91    pub fn is_error(&self) -> bool {
92        matches!(self, Self::Error { .. })
93    }
94
95    pub fn is_end(&self) -> bool {
96        matches!(self, Self::MessageEnd { .. } | Self::Error { .. })
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103
104    #[test]
105    fn test_event_serialization() {
106        let event = AgentEvent::text_delta("Hello");
107        let json = serde_json::to_string(&event).unwrap();
108        assert!(json.contains("text_delta"));
109        assert!(json.contains("Hello"));
110    }
111
112    #[test]
113    fn test_tool_use_event() {
114        let event = AgentEvent::tool_use_start("call_123", "read_file", serde_json::json!({"path": "/tmp"}));
115        
116        if let AgentEvent::ToolUseStart { id, name, input } = event {
117            assert_eq!(id, "call_123");
118            assert_eq!(name, "read_file");
119            assert!(input.get("path").is_some());
120        } else {
121            panic!("Expected ToolUseStart event");
122        }
123    }
124
125    #[test]
126    fn test_is_end() {
127        assert!(AgentEvent::message_end(Id::new()).is_end());
128        assert!(AgentEvent::error("oops").is_end());
129        assert!(!AgentEvent::text_delta("hi").is_end());
130    }
131}