hehe_core/event/
mod.rs

1use crate::types::{AgentId, EventId, MessageId, SessionId, Timestamp, ToolCallId};
2use serde::{Deserialize, Serialize};
3
4#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
5#[serde(rename_all = "snake_case")]
6pub enum EventKind {
7    AgentStarted,
8    AgentStopped,
9    AgentPaused,
10    AgentResumed,
11    AgentError,
12    SessionCreated,
13    SessionEnded,
14    MessageReceived,
15    MessageSent,
16    MessageStreaming,
17    MessageStreamEnd,
18    ToolCallStarted,
19    ToolCallCompleted,
20    ToolCallFailed,
21    ToolCallCancelled,
22    LlmRequestStarted,
23    LlmRequestCompleted,
24    LlmRequestFailed,
25    LlmTokenUsage,
26    StorageWrite,
27    StorageDelete,
28    ConfigReloaded,
29    PluginLoaded,
30    PluginUnloaded,
31    Custom(String),
32}
33
34#[derive(Clone, Debug, Default, Serialize, Deserialize)]
35pub struct TokenUsage {
36    pub input_tokens: u32,
37    pub output_tokens: u32,
38    #[serde(skip_serializing_if = "Option::is_none")]
39    pub cache_read_tokens: Option<u32>,
40    #[serde(skip_serializing_if = "Option::is_none")]
41    pub cache_write_tokens: Option<u32>,
42}
43
44impl TokenUsage {
45    pub fn new(input: u32, output: u32) -> Self {
46        Self {
47            input_tokens: input,
48            output_tokens: output,
49            cache_read_tokens: None,
50            cache_write_tokens: None,
51        }
52    }
53
54    pub fn total(&self) -> u32 {
55        self.input_tokens + self.output_tokens
56    }
57}
58
59#[derive(Clone, Debug, Serialize, Deserialize)]
60#[serde(untagged)]
61pub enum EventPayload {
62    None,
63    Agent {
64        agent_id: AgentId,
65    },
66    Session {
67        session_id: SessionId,
68        agent_id: AgentId,
69    },
70    Message {
71        message_id: MessageId,
72        session_id: SessionId,
73    },
74    ToolCall {
75        tool_call_id: ToolCallId,
76        tool_name: String,
77    },
78    Llm {
79        provider: String,
80        model: String,
81        #[serde(skip_serializing_if = "Option::is_none")]
82        usage: Option<TokenUsage>,
83    },
84    Error {
85        code: String,
86        message: String,
87    },
88    Custom(serde_json::Value),
89}
90
91impl Default for EventPayload {
92    fn default() -> Self {
93        Self::None
94    }
95}
96
97#[derive(Clone, Debug, Serialize, Deserialize)]
98pub struct Event {
99    pub id: EventId,
100    pub kind: EventKind,
101    pub payload: EventPayload,
102    pub timestamp: Timestamp,
103    #[serde(skip_serializing_if = "Option::is_none")]
104    pub source: Option<String>,
105    #[serde(skip_serializing_if = "Option::is_none")]
106    pub trace_id: Option<String>,
107}
108
109impl Event {
110    pub fn new(kind: EventKind) -> Self {
111        Self {
112            id: EventId::new(),
113            kind,
114            payload: EventPayload::None,
115            timestamp: Timestamp::now(),
116            source: None,
117            trace_id: None,
118        }
119    }
120
121    pub fn with_payload(mut self, payload: EventPayload) -> Self {
122        self.payload = payload;
123        self
124    }
125
126    pub fn with_source(mut self, source: impl Into<String>) -> Self {
127        self.source = Some(source.into());
128        self
129    }
130
131    pub fn with_trace_id(mut self, trace_id: impl Into<String>) -> Self {
132        self.trace_id = Some(trace_id.into());
133        self
134    }
135
136    pub fn agent_started(agent_id: AgentId) -> Self {
137        Self::new(EventKind::AgentStarted).with_payload(EventPayload::Agent { agent_id })
138    }
139
140    pub fn agent_stopped(agent_id: AgentId) -> Self {
141        Self::new(EventKind::AgentStopped).with_payload(EventPayload::Agent { agent_id })
142    }
143
144    pub fn tool_call_started(tool_call_id: ToolCallId, tool_name: impl Into<String>) -> Self {
145        Self::new(EventKind::ToolCallStarted).with_payload(EventPayload::ToolCall {
146            tool_call_id,
147            tool_name: tool_name.into(),
148        })
149    }
150
151    pub fn tool_call_completed(tool_call_id: ToolCallId, tool_name: impl Into<String>) -> Self {
152        Self::new(EventKind::ToolCallCompleted).with_payload(EventPayload::ToolCall {
153            tool_call_id,
154            tool_name: tool_name.into(),
155        })
156    }
157
158    pub fn llm_completed(
159        provider: impl Into<String>,
160        model: impl Into<String>,
161        usage: Option<TokenUsage>,
162    ) -> Self {
163        Self::new(EventKind::LlmRequestCompleted).with_payload(EventPayload::Llm {
164            provider: provider.into(),
165            model: model.into(),
166            usage,
167        })
168    }
169
170    pub fn error(code: impl Into<String>, message: impl Into<String>) -> Self {
171        Self::new(EventKind::AgentError).with_payload(EventPayload::Error {
172            code: code.into(),
173            message: message.into(),
174        })
175    }
176}
177
178#[async_trait::async_trait]
179pub trait EventEmitter: Send + Sync {
180    async fn emit(&self, event: Event);
181}
182
183#[async_trait::async_trait]
184pub trait EventSubscriber: Send + Sync {
185    fn event_kinds(&self) -> Vec<EventKind>;
186    async fn on_event(&self, event: &Event);
187}