aether_core/testing/
agent_event_builder.rs

1use crate::events::{AgentEvent, StreamState, ToolEvent};
2use crate::mcp::tool_bridge::encode_structured;
3use llm::{ToolCallError, ToolCallRequest, ToolCallResult};
4use serde::Serialize;
5
6pub fn agent_event(message_id: &str) -> AgentEventBuilder {
7    AgentEventBuilder::new(message_id)
8}
9
10pub struct AgentEventBuilder {
11    message_id: String,
12    chunks: Vec<AgentEvent>,
13    full_text: String,
14}
15
16impl AgentEventBuilder {
17    pub fn new(message_id: &str) -> Self {
18        Self { message_id: message_id.to_string(), chunks: Vec::new(), full_text: String::new() }
19    }
20
21    pub fn text(mut self, chunks: &[&str]) -> Self {
22        for chunk in chunks {
23            self.chunks.push(AgentEvent::text(&self.message_id, chunk, StreamState::Partial));
24            self.full_text.push_str(chunk);
25        }
26        self
27    }
28
29    pub fn tool_call<T: Serialize, U: Serialize>(
30        mut self,
31        tool_call_id: &str,
32        name: &str,
33        request: &T,
34        result: &U,
35    ) -> Self {
36        let request_json = serde_json::to_string(request).expect("Failed to serialize request");
37        let result_value = serde_json::to_value(result).expect("Failed to serialize result");
38
39        self.push_tool_call(tool_call_id, name, &request_json);
40
41        self.chunks.push(AgentEvent::Tool(ToolEvent::Result {
42            result: ToolCallResult {
43                id: tool_call_id.to_string(),
44                name: name.to_string(),
45                arguments: request_json,
46                result: encode_structured(&result_value),
47            },
48            result_meta: None,
49        }));
50
51        self
52    }
53
54    pub fn tool_call_with_error<T: Serialize>(
55        mut self,
56        tool_call_id: &str,
57        name: &str,
58        request: &T,
59        error_message: &str,
60    ) -> Self {
61        let request_json = serde_json::to_string(request).expect("Failed to serialize request");
62
63        let error_result = format!("Tool execution error: {error_message}");
64
65        self.push_tool_call(tool_call_id, name, &request_json);
66
67        self.chunks.push(AgentEvent::Tool(ToolEvent::Error {
68            error: ToolCallError {
69                id: tool_call_id.to_string(),
70                name: name.to_string(),
71                arguments: Some(request_json),
72                error: error_result,
73            },
74        }));
75
76        self
77    }
78
79    pub fn build(mut self) -> Vec<AgentEvent> {
80        self.chunks.push(AgentEvent::text(&self.message_id, &self.full_text, StreamState::Complete));
81
82        self.chunks
83    }
84
85    fn push_tool_call(&mut self, tool_call_id: &str, name: &str, arguments: &str) {
86        self.chunks.extend([
87            AgentEvent::Tool(ToolEvent::InputStarted { id: tool_call_id.to_string(), name: name.to_string() }),
88            AgentEvent::Tool(ToolEvent::InputDelta { id: tool_call_id.to_string(), chunk: arguments.to_string() }),
89            AgentEvent::Tool(ToolEvent::Call {
90                request: ToolCallRequest {
91                    id: tool_call_id.to_string(),
92                    name: name.to_string(),
93                    arguments: arguments.to_string(),
94                },
95            }),
96        ]);
97    }
98}
aether_core/testing/agent_event_builder.rs

aether_core/testing/
agent_event_builder.rs
