tsk/agent/
mod.rs

1use async_trait::async_trait;
2use std::sync::Arc;
3
4mod claude_code;
5mod log_processor;
6mod no_op;
7mod no_op_log_processor;
8mod provider;
9
10pub use self::log_processor::LogProcessor;
11pub use claude_code::{ClaudeCodeAgent, TaskResult};
12pub use no_op::NoOpAgent;
13pub use provider::AgentProvider;
14
15/// Trait defining the interface for AI agents that can execute tasks
16#[async_trait]
17pub trait Agent: Send + Sync {
18    /// Returns the command to execute the agent with the given instruction file
19    fn build_command(&self, instruction_path: &str) -> Vec<String>;
20
21    /// Returns the volumes to mount for this agent
22    /// Format: Vec<(host_path, container_path, options)> where options is like ":ro" for read-only
23    fn volumes(&self) -> Vec<(String, String, String)>;
24
25    /// Returns environment variables for this agent
26    fn environment(&self) -> Vec<(String, String)>;
27
28    /// Creates a log processor for this agent's output
29    fn create_log_processor(
30        &self,
31        file_system: Arc<dyn crate::context::file_system::FileSystemOperations>,
32    ) -> Box<dyn LogProcessor>;
33
34    /// Returns the agent's unique identifier
35    fn name(&self) -> &str;
36
37    /// Validates that this agent is properly configured
38    async fn validate(&self) -> Result<(), String> {
39        Ok(())
40    }
41
42    /// Performs any necessary warmup steps before launching the Docker container
43    ///
44    /// This method is called after validation but before container creation.
45    /// It can be used to execute host-side setup commands, refresh credentials,
46    /// or perform any other preparatory work needed by the agent.
47    ///
48    /// The default implementation does nothing, allowing backward compatibility.
49    async fn warmup(&self) -> Result<(), String> {
50        Ok(())
51    }
52}
53
54#[cfg(test)]
55mod tests {
56    use super::*;
57    use std::path::Path;
58
59    /// Test agent for testing purposes
60    struct TestAgent {
61        name: String,
62    }
63
64    impl TestAgent {
65        fn new(name: &str) -> Self {
66            Self {
67                name: name.to_string(),
68            }
69        }
70    }
71
72    #[async_trait]
73    impl Agent for TestAgent {
74        fn build_command(&self, instruction_path: &str) -> Vec<String> {
75            vec!["test".to_string(), instruction_path.to_string()]
76        }
77
78        fn volumes(&self) -> Vec<(String, String, String)> {
79            vec![("/test".to_string(), "/test".to_string(), ":ro".to_string())]
80        }
81
82        fn environment(&self) -> Vec<(String, String)> {
83            vec![("TEST_VAR".to_string(), "test_value".to_string())]
84        }
85
86        fn create_log_processor(
87            &self,
88            _file_system: Arc<dyn crate::context::file_system::FileSystemOperations>,
89        ) -> Box<dyn LogProcessor> {
90            struct TestLogProcessor;
91
92            #[async_trait]
93            impl LogProcessor for TestLogProcessor {
94                fn process_line(&mut self, _line: &str) -> Option<String> {
95                    Some("test".to_string())
96                }
97
98                fn get_full_log(&self) -> String {
99                    "test log".to_string()
100                }
101
102                async fn save_full_log(&self, _path: &Path) -> Result<(), String> {
103                    Ok(())
104                }
105
106                fn get_final_result(&self) -> Option<&super::TaskResult> {
107                    None
108                }
109            }
110
111            Box::new(TestLogProcessor)
112        }
113
114        fn name(&self) -> &str {
115            &self.name
116        }
117    }
118
119    #[test]
120    fn test_agent_trait_is_object_safe() {
121        // This test ensures that the Agent trait can be used as a trait object
122        fn _assert_object_safe(_: &dyn Agent) {}
123
124        let agent = TestAgent::new("test");
125        _assert_object_safe(&agent);
126    }
127
128    #[tokio::test]
129    async fn test_no_op_agent() {
130        let agent = NoOpAgent;
131
132        // Test agent name
133        assert_eq!(agent.name(), "no-op");
134
135        // Test build_command
136        let command = agent.build_command("/instructions/test.md");
137        assert_eq!(command.len(), 3);
138        assert_eq!(command[0], "sh");
139        assert_eq!(command[1], "-c");
140        assert!(command[2].contains("cat '/instructions/test.md'"));
141
142        // Test volumes
143        let volumes = agent.volumes();
144        assert!(volumes.is_empty());
145
146        // Test environment
147        let env = agent.environment();
148        assert!(env.is_empty());
149
150        // Test validation (should always succeed)
151        assert!(agent.validate().await.is_ok());
152
153        // Test warmup (should always succeed)
154        assert!(agent.warmup().await.is_ok());
155    }
156
157    #[test]
158    fn test_no_op_log_processor() {
159        use super::no_op_log_processor::NoOpLogProcessor;
160
161        let mut processor = NoOpLogProcessor::new();
162
163        // Test process_line passes through
164        let line = "test output line";
165        let result = processor.process_line(line);
166        assert_eq!(result, Some(line.to_string()));
167
168        // Test get_full_log
169        processor.process_line("line 1");
170        processor.process_line("line 2");
171        let full_log = processor.get_full_log();
172        assert!(full_log.contains("line 1"));
173        assert!(full_log.contains("line 2"));
174
175        // Test get_final_result returns None
176        assert!(processor.get_final_result().is_none());
177    }
178}