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}