Skip to main content

aether_evals/agents/
fake_agent.rs

1use super::agent::{Agent, AgentConfig, RunError};
2use aether_core::events::AgentMessage;
3use llm::ToolCallResult;
4use std::path::PathBuf;
5use tokio::sync::mpsc::Sender;
6
7#[derive(Clone)]
8pub struct FakeAgent {
9    messages: Vec<AgentMessage>,
10    file_writes: Vec<(PathBuf, String)>,
11}
12
13impl FakeAgent {
14    pub fn new(messages: Vec<AgentMessage>) -> Self {
15        Self { messages, file_writes: Vec::new() }
16    }
17
18    pub fn success() -> Self {
19        Self::new(vec![AgentMessage::text("fake_1", "Task completed successfully", true, "fake"), AgentMessage::Done])
20    }
21
22    pub fn with_tool_call(tool_name: impl Into<String>, result: impl Into<String>) -> Self {
23        let tool_name = tool_name.into();
24        Self::new(vec![
25            AgentMessage::ToolResult {
26                result: ToolCallResult {
27                    id: "fake_call_1".to_string(),
28                    name: tool_name,
29                    arguments: "{}".to_string(),
30                    result: result.into(),
31                },
32                result_meta: None,
33                model_name: "fake".to_string(),
34            },
35            AgentMessage::text("fake_2", "Task completed using tools", true, "fake"),
36            AgentMessage::Done,
37        ])
38    }
39
40    pub fn writes_file(path: impl Into<PathBuf>, contents: impl Into<String>) -> Self {
41        Self::success().with_file_write(path, contents)
42    }
43
44    pub fn with_file_write(mut self, path: impl Into<PathBuf>, contents: impl Into<String>) -> Self {
45        self.file_writes.push((path.into(), contents.into()));
46        self
47    }
48}
49
50impl Agent for FakeAgent {
51    async fn run(&self, config: AgentConfig<'_>, tx: Sender<AgentMessage>) -> Result<(), RunError> {
52        for (path, contents) in &self.file_writes {
53            let file_path = config.workspace().join(path);
54            if let Some(parent) = file_path.parent() {
55                tokio::fs::create_dir_all(parent).await.map_err(|error| {
56                    RunError::ExecutionFailed(format!("failed to create fake file parent: {error}"))
57                })?;
58            }
59            tokio::fs::write(&file_path, contents)
60                .await
61                .map_err(|error| RunError::ExecutionFailed(format!("failed to write fake file: {error}")))?;
62        }
63
64        for message in &self.messages {
65            tx.send(message.clone()).await.map_err(|error| RunError::ChannelSendFailed(error.to_string()))?;
66        }
67
68        Ok(())
69    }
70}
71
72#[cfg(test)]
73mod tests {
74    use super::*;
75    use tokio::sync::mpsc;
76
77    #[tokio::test]
78    async fn fake_agent_success_sends_done() {
79        let agent = FakeAgent::success();
80        let (tx, mut rx) = mpsc::channel(10);
81        let temp_dir = tempfile::tempdir().unwrap();
82        let config = AgentConfig { workspace_root: temp_dir.path(), relative_cwd: None, task_prompt: "test task" };
83        let result = agent.run(config, tx).await;
84        assert!(result.is_ok());
85
86        let message = rx.recv().await.unwrap();
87        assert!(matches!(message, AgentMessage::Text { .. }));
88
89        let message = rx.recv().await.unwrap();
90        assert!(matches!(message, AgentMessage::Done));
91    }
92
93    #[tokio::test]
94    async fn fake_agent_with_tool_call_sends_tool_messages() {
95        let agent = FakeAgent::with_tool_call("bash", "success");
96        let (tx, mut rx) = mpsc::channel(10);
97        let temp_dir = tempfile::tempdir().unwrap();
98        let config = AgentConfig { workspace_root: temp_dir.path(), relative_cwd: None, task_prompt: "test task" };
99        let result = agent.run(config, tx).await;
100        assert!(result.is_ok());
101
102        let mut count = 0;
103        while let Some(message) = rx.recv().await {
104            count += 1;
105            if matches!(message, AgentMessage::Done) {
106                break;
107            }
108        }
109        assert_eq!(count, 3);
110    }
111
112    #[tokio::test]
113    async fn fake_agent_writes_file_in_workspace() {
114        let agent = FakeAgent::writes_file("nested/hello.txt", "hello");
115        let (tx, _rx) = mpsc::channel(10);
116        let temp_dir = tempfile::tempdir().unwrap();
117        let config = AgentConfig { workspace_root: temp_dir.path(), relative_cwd: None, task_prompt: "test task" };
118        agent.run(config, tx).await.unwrap();
119        assert_eq!(std::fs::read_to_string(temp_dir.path().join("nested/hello.txt")).unwrap(), "hello");
120    }
121}