Skip to main content

rs_adk/code_executors/
mod.rs

1//! Code execution infrastructure — sandboxed code execution for agents.
2
3/// Base trait and error types for code execution.
4pub mod base;
5/// Built-in code executor using Gemini's native code execution.
6pub mod built_in;
7/// Container-based code executor using Docker.
8pub mod container;
9/// Types used by code executors (input, output, files).
10pub mod types;
11/// Unsafe local code executor (no sandboxing).
12pub mod unsafe_local;
13/// Utility functions for extracting code blocks and building parts.
14pub mod utils;
15/// Vertex AI managed code executor.
16pub mod vertex_ai;
17
18pub use base::{CodeExecutor, CodeExecutorError};
19pub use built_in::BuiltInCodeExecutor;
20pub use container::{ContainerCodeExecutor, ContainerCodeExecutorConfig};
21pub use types::{CodeExecutionInput, CodeExecutionResult, CodeFile};
22pub use unsafe_local::UnsafeLocalCodeExecutor;
23pub use vertex_ai::{VertexAiCodeExecutor, VertexAiCodeExecutorConfig};
24
25#[cfg(test)]
26mod tests {
27    use super::*;
28    use async_trait::async_trait;
29
30    // --- CodeFile tests ---
31
32    #[test]
33    fn code_file_construction() {
34        let file = CodeFile {
35            name: "main.py".into(),
36            content: "print('hello')".into(),
37            mime_type: "text/x-python".into(),
38        };
39        assert_eq!(file.name, "main.py");
40        assert_eq!(file.content, "print('hello')");
41        assert_eq!(file.mime_type, "text/x-python");
42    }
43
44    #[test]
45    fn code_file_serde_roundtrip() {
46        let file = CodeFile {
47            name: "data.csv".into(),
48            content: "a,b\n1,2".into(),
49            mime_type: "text/csv".into(),
50        };
51        let json = serde_json::to_string(&file).unwrap();
52        let deserialized: CodeFile = serde_json::from_str(&json).unwrap();
53        assert_eq!(deserialized.name, file.name);
54        assert_eq!(deserialized.content, file.content);
55        assert_eq!(deserialized.mime_type, file.mime_type);
56    }
57
58    // --- CodeExecutionResult tests ---
59
60    #[test]
61    fn code_execution_result_empty() {
62        let result = CodeExecutionResult::empty();
63        assert!(result.stdout.is_empty());
64        assert!(result.stderr.is_empty());
65        assert!(result.output_files.is_empty());
66    }
67
68    // --- CodeExecutionInput tests ---
69
70    #[test]
71    fn code_execution_input_construction() {
72        let input = CodeExecutionInput {
73            code: "x = 1 + 2".into(),
74            input_files: vec![CodeFile {
75                name: "input.txt".into(),
76                content: "hello".into(),
77                mime_type: "text/plain".into(),
78            }],
79            execution_id: Some("exec-123".into()),
80        };
81        assert_eq!(input.code, "x = 1 + 2");
82        assert_eq!(input.input_files.len(), 1);
83        assert_eq!(input.execution_id.as_deref(), Some("exec-123"));
84    }
85
86    #[test]
87    fn code_execution_input_no_execution_id() {
88        let input = CodeExecutionInput {
89            code: "pass".into(),
90            input_files: Vec::new(),
91            execution_id: None,
92        };
93        assert!(input.execution_id.is_none());
94        assert!(input.input_files.is_empty());
95    }
96
97    // --- CodeExecutor trait default methods ---
98
99    /// A minimal executor to test default trait method implementations.
100    struct StubExecutor;
101
102    #[async_trait]
103    impl CodeExecutor for StubExecutor {
104        async fn execute_code(
105            &self,
106            _input: CodeExecutionInput,
107        ) -> Result<CodeExecutionResult, CodeExecutorError> {
108            Ok(CodeExecutionResult::empty())
109        }
110    }
111
112    #[test]
113    fn default_code_block_delimiters() {
114        let exec = StubExecutor;
115        let delims = exec.code_block_delimiters();
116        assert_eq!(delims.len(), 2);
117        assert_eq!(
118            delims[0],
119            ("```tool_code\n".to_string(), "\n```".to_string())
120        );
121        assert_eq!(delims[1], ("```python\n".to_string(), "\n```".to_string()));
122    }
123
124    #[test]
125    fn default_execution_result_delimiters() {
126        let exec = StubExecutor;
127        let (open, close) = exec.execution_result_delimiters();
128        assert_eq!(open, "```tool_output\n");
129        assert_eq!(close, "\n```");
130    }
131
132    #[test]
133    fn default_error_retry_attempts() {
134        let exec = StubExecutor;
135        assert_eq!(exec.error_retry_attempts(), 2);
136    }
137
138    #[test]
139    fn default_stateful() {
140        let exec = StubExecutor;
141        assert!(!exec.stateful());
142    }
143
144    #[tokio::test]
145    async fn stub_executor_execute_code() {
146        let exec = StubExecutor;
147        let input = CodeExecutionInput {
148            code: "1 + 1".into(),
149            input_files: Vec::new(),
150            execution_id: None,
151        };
152        let result = exec.execute_code(input).await.unwrap();
153        assert!(result.stdout.is_empty());
154    }
155
156    // --- Object safety ---
157
158    fn _assert_object_safe(_: &dyn CodeExecutor) {}
159
160    #[test]
161    fn code_executor_is_object_safe() {
162        let exec = StubExecutor;
163        _assert_object_safe(&exec);
164    }
165
166    // --- BuiltInCodeExecutor tests ---
167
168    #[test]
169    fn built_in_process_llm_request_adds_code_execution_for_gemini2() {
170        let executor = built_in::BuiltInCodeExecutor;
171        let mut request = crate::llm::LlmRequest::from_text("hello");
172        executor
173            .process_llm_request(&mut request, "gemini-2.5-flash")
174            .unwrap();
175        assert_eq!(request.tools.len(), 1);
176        assert!(request.tools[0].code_execution.is_some());
177    }
178
179    #[test]
180    fn built_in_process_llm_request_rejects_non_gemini2() {
181        let executor = built_in::BuiltInCodeExecutor;
182        let mut request = crate::llm::LlmRequest::from_text("hello");
183        let err = executor
184            .process_llm_request(&mut request, "gemini-1.5-pro")
185            .unwrap_err();
186        assert!(
187            err.to_string().contains("Gemini 2.0+"),
188            "expected UnsupportedModel error, got: {}",
189            err
190        );
191        assert!(request.tools.is_empty());
192    }
193
194    #[tokio::test]
195    async fn built_in_execute_code_returns_empty() {
196        let executor = built_in::BuiltInCodeExecutor;
197        let input = CodeExecutionInput {
198            code: "print('hi')".into(),
199            input_files: Vec::new(),
200            execution_id: None,
201        };
202        let result = executor.execute_code(input).await.unwrap();
203        assert!(result.stdout.is_empty());
204        assert!(result.stderr.is_empty());
205        assert!(result.output_files.is_empty());
206    }
207
208    // --- utils tests ---
209
210    #[test]
211    fn extract_code_from_tool_code_block() {
212        let text = "Some text\n```tool_code\nprint('hello')\n```\nMore text";
213        let delimiters = vec![
214            ("```tool_code\n".to_string(), "\n```".to_string()),
215            ("```python\n".to_string(), "\n```".to_string()),
216        ];
217        let (code, remaining) = utils::extract_code_from_text(text, &delimiters).unwrap();
218        assert_eq!(code, "print('hello')");
219        assert_eq!(remaining, "Some text\n\nMore text");
220    }
221
222    #[test]
223    fn extract_code_from_python_block() {
224        let text = "Intro\n```python\nx = 42\n```\nDone";
225        let delimiters = vec![
226            ("```tool_code\n".to_string(), "\n```".to_string()),
227            ("```python\n".to_string(), "\n```".to_string()),
228        ];
229        let (code, remaining) = utils::extract_code_from_text(text, &delimiters).unwrap();
230        assert_eq!(code, "x = 42");
231        assert_eq!(remaining, "Intro\n\nDone");
232    }
233
234    #[test]
235    fn extract_code_returns_none_when_no_block() {
236        let text = "Just plain text, no code blocks here.";
237        let delimiters = vec![
238            ("```tool_code\n".to_string(), "\n```".to_string()),
239            ("```python\n".to_string(), "\n```".to_string()),
240        ];
241        assert!(utils::extract_code_from_text(text, &delimiters).is_none());
242    }
243
244    #[test]
245    fn build_executable_code_part_correct() {
246        let part = utils::build_executable_code_part("x = 1");
247        match part {
248            rs_genai::prelude::Part::ExecutableCode { executable_code } => {
249                assert_eq!(executable_code.language, "PYTHON");
250                assert_eq!(executable_code.code, "x = 1");
251            }
252            other => panic!("expected ExecutableCode, got: {:?}", other),
253        }
254    }
255
256    #[test]
257    fn build_code_execution_result_part_ok_outcome() {
258        let part = utils::build_code_execution_result_part("42\n", "");
259        match part {
260            rs_genai::prelude::Part::CodeExecutionResult {
261                code_execution_result,
262            } => {
263                assert_eq!(code_execution_result.outcome, "OK");
264                assert_eq!(code_execution_result.output.as_deref(), Some("42\n"));
265            }
266            other => panic!("expected CodeExecutionResult, got: {:?}", other),
267        }
268    }
269
270    #[test]
271    fn build_code_execution_result_part_failed_outcome() {
272        let part =
273            utils::build_code_execution_result_part("partial output", "NameError: x not defined");
274        match part {
275            rs_genai::prelude::Part::CodeExecutionResult {
276                code_execution_result,
277            } => {
278                assert_eq!(code_execution_result.outcome, "FAILED");
279                let output = code_execution_result.output.unwrap();
280                assert!(output.contains("partial output"));
281                assert!(output.contains("NameError: x not defined"));
282            }
283            other => panic!("expected CodeExecutionResult, got: {:?}", other),
284        }
285    }
286}