Skip to main content

a3s_code_core/tools/
program_tool.rs

1//! Tool wrapper for programmatic tool calling.
2
3use crate::program::{
4    program_verification_hints, ProgramCatalog, ProgramExecutor, ProgramResult, ProgramStepResult,
5    ProgramTrace, ProgramTraceArtifact, ProgramTraceStep, ProgramVerificationHint,
6};
7use crate::text::truncate_utf8;
8use crate::tools::types::{Tool, ToolContext, ToolOutput};
9use crate::tools::{tool_output_artifact, ToolArtifact, ToolRegistry};
10use crate::verification::VerificationReport;
11use anyhow::Result;
12use async_trait::async_trait;
13use std::sync::Arc;
14
15const MAX_PROGRAM_STEP_OUTPUT_BYTES: usize = 4 * 1024;
16
17pub struct ProgramTool {
18    registry: Arc<ToolRegistry>,
19    catalog: ProgramCatalog,
20}
21
22impl ProgramTool {
23    pub fn new(registry: Arc<ToolRegistry>) -> Self {
24        Self::with_catalog(registry, ProgramCatalog::with_builtin_programs())
25    }
26
27    pub fn with_catalog(registry: Arc<ToolRegistry>, catalog: ProgramCatalog) -> Self {
28        Self { registry, catalog }
29    }
30}
31
32#[async_trait]
33impl Tool for ProgramTool {
34    fn name(&self) -> &str {
35        "program"
36    }
37
38    fn description(&self) -> &str {
39        "Run a named harness program such as program_code_search or program_repo_map. Programs execute bounded tool chains and return summarized step results."
40    }
41
42    fn parameters(&self) -> serde_json::Value {
43        serde_json::json!({
44            "type": "object",
45            "additionalProperties": false,
46            "properties": {
47                "name": {
48                    "type": "string",
49                    "description": "Required. Program name to run.",
50                    "enum": self.catalog.list().iter().map(|program| program.name.clone()).collect::<Vec<_>>()
51                },
52                "inputs": {
53                    "type": "object",
54                    "description": "Optional. Program-specific inputs such as query, path, and glob."
55                }
56            },
57            "required": ["name"]
58        })
59    }
60
61    async fn execute(&self, args: &serde_json::Value, ctx: &ToolContext) -> Result<ToolOutput> {
62        let Some(name) = args.get("name").and_then(|value| value.as_str()) else {
63            return Ok(ToolOutput::error("name parameter is required"));
64        };
65        let inputs = args
66            .get("inputs")
67            .cloned()
68            .unwrap_or_else(|| serde_json::json!({}));
69
70        let program = match self.catalog.instantiate(name, &inputs) {
71            Ok(program) => program,
72            Err(err) => return Ok(ToolOutput::error(err.to_string())),
73        };
74
75        let executor = ProgramExecutor::new(Arc::clone(&self.registry), ctx.clone());
76        let result = executor.execute(&program).await?;
77        let rendered = render_program_result(&result, &self.registry);
78        let verification_hints = program_verification_hints(&result, Some(&rendered.trace));
79        let verification_report =
80            VerificationReport::from_program_hints(&result.program_name, &verification_hints);
81        Ok(
82            ToolOutput::success(rendered.output).with_metadata(serde_json::json!({
83                "program": {
84                    "name": result.program_name,
85                    "success": result.success,
86                    "summary": result.summary,
87                    "steps": result.steps.iter().map(|step| {
88                        serde_json::json!({
89                            "tool_name": step.tool_name,
90                            "label": step.label,
91                            "success": step.success,
92                            "metadata": step.metadata,
93                        })
94                    }).collect::<Vec<_>>(),
95                },
96                "trace": rendered.trace.to_value(),
97                "verification_hints": ProgramVerificationHint::to_values(&verification_hints),
98                "verification_report": verification_report.to_value(),
99            })),
100        )
101    }
102}
103
104#[derive(Debug)]
105struct RenderedProgram {
106    output: String,
107    trace: ProgramTrace,
108}
109
110#[derive(Debug)]
111struct RenderedStep {
112    output: String,
113    trace: ProgramTraceStep,
114}
115
116fn render_program_result(result: &ProgramResult, registry: &ToolRegistry) -> RenderedProgram {
117    let mut output = String::new();
118    output.push_str(&result.summary);
119    if let Some(summary) = program_specific_summary(result) {
120        output.push('\n');
121        output.push_str(&summary);
122    }
123
124    let mut trace_steps = Vec::with_capacity(result.steps.len());
125    for (index, step) in result.steps.iter().enumerate() {
126        let rendered_step = render_step(&result.program_name, index, step, registry);
127        let label = step.label.as_deref().unwrap_or(&step.tool_name);
128        output.push_str(&format!(
129            "\n\n## Step {}: {} [{}] ({})\n{}",
130            index + 1,
131            label,
132            step.tool_name,
133            if step.success { "ok" } else { "failed" },
134            rendered_step.output
135        ));
136        trace_steps.push(rendered_step.trace);
137    }
138
139    RenderedProgram {
140        output,
141        trace: ProgramTrace::from_result(result, trace_steps),
142    }
143}
144
145fn program_specific_summary(result: &ProgramResult) -> Option<String> {
146    match result.program_name.as_str() {
147        "program_code_search" => summarize_code_search(result),
148        "program_repo_map" => summarize_repo_map(result),
149        _ => None,
150    }
151}
152
153fn summarize_code_search(result: &ProgramResult) -> Option<String> {
154    let step = result.steps.first()?;
155    if step.output.contains("No matches found") {
156        return Some("Search summary: no matches found.".to_string());
157    }
158
159    step.output
160        .lines()
161        .rev()
162        .find(|line| line.contains("match(es) in") && line.contains("file(s)"))
163        .map(|line| format!("Search summary: {}.", line.trim()))
164}
165
166fn summarize_repo_map(result: &ProgramResult) -> Option<String> {
167    let mut files = Vec::new();
168    for step in &result.steps {
169        if step.tool_name != "glob" || !step.success || step.output.contains("No files found") {
170            continue;
171        }
172        for line in step.output.lines() {
173            let trimmed = line.trim();
174            if trimmed.is_empty() || trimmed.ends_with("file(s) found") {
175                continue;
176            }
177            files.push(trimmed.to_string());
178        }
179    }
180
181    files.sort();
182    files.dedup();
183    if files.is_empty() {
184        Some("Repo map summary: no known project files found.".to_string())
185    } else {
186        Some(format!(
187            "Repo map summary: found project files: {}.",
188            files.join(", ")
189        ))
190    }
191}
192
193fn render_step(
194    program_name: &str,
195    step_index: usize,
196    step: &ProgramStepResult,
197    registry: &ToolRegistry,
198) -> RenderedStep {
199    let base_trace = |compacted: bool, artifact: Option<ProgramTraceArtifact>| {
200        ProgramTraceStep::from_result(step_index, step, compacted, artifact)
201    };
202
203    if step.output.len() <= MAX_PROGRAM_STEP_OUTPUT_BYTES {
204        return RenderedStep {
205            output: step.output.clone(),
206            trace: base_trace(false, None),
207        };
208    }
209
210    let shown = truncate_utf8(&step.output, MAX_PROGRAM_STEP_OUTPUT_BYTES);
211    let artifact = tool_output_artifact(
212        &format!(
213            "program-step-{program_name}-{}-{step_index}",
214            step.tool_name
215        ),
216        &step.output,
217        shown.len(),
218    );
219    registry.artifact_store().put(ToolArtifact {
220        artifact_id: artifact.artifact_id.clone(),
221        artifact_uri: artifact.artifact_uri.clone(),
222        tool_name: format!("program:{program_name}:{}", step.tool_name),
223        content: step.output.clone(),
224        original_bytes: artifact.original_bytes,
225        shown_bytes: artifact.shown_bytes,
226    });
227
228    let artifact_id = artifact.artifact_id.clone();
229    let artifact_uri = artifact.artifact_uri.clone();
230    let artifact_trace = ProgramTraceArtifact {
231        artifact_id,
232        artifact_uri,
233        original_bytes: artifact.original_bytes,
234        shown_bytes: artifact.shown_bytes,
235    };
236
237    RenderedStep {
238        output: format!(
239            "{}\n\n[program step output compacted: showing the first {} of {} bytes. Full step artifact: {}.]",
240            shown, artifact.shown_bytes, artifact.original_bytes, artifact.artifact_uri
241        ),
242        trace: base_trace(true, Some(artifact_trace)),
243    }
244}
245
246#[cfg(test)]
247mod tests {
248    use super::*;
249    use crate::program::PROGRAM_TRACE_SCHEMA;
250    use crate::tools::{Tool, ToolOutput};
251    use anyhow::Result;
252    use async_trait::async_trait;
253    use std::path::PathBuf;
254
255    struct EchoGrepTool;
256
257    #[async_trait]
258    impl Tool for EchoGrepTool {
259        fn name(&self) -> &str {
260            "grep"
261        }
262
263        fn description(&self) -> &str {
264            "Echo grep args"
265        }
266
267        fn parameters(&self) -> serde_json::Value {
268            serde_json::json!({"type": "object"})
269        }
270
271        async fn execute(
272            &self,
273            args: &serde_json::Value,
274            _ctx: &ToolContext,
275        ) -> Result<ToolOutput> {
276            if args["pattern"].as_str() == Some("large") {
277                return Ok(ToolOutput::success(
278                    "x".repeat(MAX_PROGRAM_STEP_OUTPUT_BYTES + 1),
279                ));
280            }
281            if args["pattern"].as_str() == Some("missing") {
282                return Ok(ToolOutput::success("No matches found for pattern: missing"));
283            }
284
285            Ok(ToolOutput::success(format!(
286                ">src/lib.rs:1: {} in {}\n\n1 match(es) in 1 file(s)",
287                args["pattern"].as_str().unwrap_or_default(),
288                args["path"].as_str().unwrap_or_default()
289            )))
290        }
291    }
292
293    struct RepoMapTool;
294
295    #[async_trait]
296    impl Tool for RepoMapTool {
297        fn name(&self) -> &str {
298            "glob"
299        }
300
301        fn description(&self) -> &str {
302            "Return selected repo files"
303        }
304
305        fn parameters(&self) -> serde_json::Value {
306            serde_json::json!({"type": "object"})
307        }
308
309        async fn execute(
310            &self,
311            args: &serde_json::Value,
312            _ctx: &ToolContext,
313        ) -> Result<ToolOutput> {
314            match args["pattern"].as_str().unwrap_or_default() {
315                "Cargo.toml" => Ok(ToolOutput::success("Cargo.toml\n\n1 file(s) found")),
316                "README.md" => Ok(ToolOutput::success("README.md\n\n1 file(s) found")),
317                pattern => Ok(ToolOutput::success(format!(
318                    "No files found matching pattern: {pattern}"
319                ))),
320            }
321        }
322    }
323
324    struct LsTool;
325
326    #[async_trait]
327    impl Tool for LsTool {
328        fn name(&self) -> &str {
329            "ls"
330        }
331
332        fn description(&self) -> &str {
333            "List files"
334        }
335
336        fn parameters(&self) -> serde_json::Value {
337            serde_json::json!({"type": "object"})
338        }
339
340        async fn execute(
341            &self,
342            _args: &serde_json::Value,
343            _ctx: &ToolContext,
344        ) -> Result<ToolOutput> {
345            Ok(ToolOutput::success("Directory: /tmp\n\nfile Cargo.toml"))
346        }
347    }
348
349    #[tokio::test]
350    async fn program_tool_runs_catalog_program() {
351        let registry = Arc::new(ToolRegistry::new(PathBuf::from("/tmp")));
352        registry.register(Arc::new(EchoGrepTool));
353        let tool = ProgramTool::new(Arc::clone(&registry));
354        let output = tool
355            .execute(
356                &serde_json::json!({
357                    "name": "program_code_search",
358                    "inputs": {
359                        "query": "AgentLoop",
360                        "path": "core/src"
361                    }
362                }),
363                &ToolContext::new(PathBuf::from("/tmp")),
364            )
365            .await
366            .unwrap();
367
368        assert!(output.success);
369        assert!(output.content.contains("program_code_search"));
370        assert!(output.content.contains("Step 1: search_code [grep]"));
371        assert!(output
372            .content
373            .contains("Search summary: 1 match(es) in 1 file(s)."));
374        assert!(output.content.contains("AgentLoop in core/src"));
375        let metadata = output.metadata.as_ref().expect("metadata");
376        assert_eq!(metadata["program"]["name"], "program_code_search");
377        assert_eq!(metadata["trace"]["schema"], PROGRAM_TRACE_SCHEMA);
378        assert_eq!(metadata["trace"]["type"], "program_execution");
379        assert_eq!(metadata["trace"]["step_count"], 1);
380        assert_eq!(metadata["trace"]["steps"][0]["label"], "search_code");
381        assert_eq!(metadata["verification_hints"][0]["kind"], "inspect_matches");
382        assert_eq!(metadata["verification_hints"][0]["required"], true);
383        assert_eq!(metadata["verification_report"]["status"], "needs_review");
384        assert_eq!(
385            metadata["verification_report"]["checks"][0]["kind"],
386            "inspect_matches"
387        );
388    }
389
390    #[tokio::test]
391    async fn program_tool_summarizes_code_search_misses() {
392        let registry = Arc::new(ToolRegistry::new(PathBuf::from("/tmp")));
393        registry.register(Arc::new(EchoGrepTool));
394        let tool = ProgramTool::new(Arc::clone(&registry));
395        let output = tool
396            .execute(
397                &serde_json::json!({
398                    "name": "program_code_search",
399                    "inputs": {
400                        "query": "missing"
401                    }
402                }),
403                &ToolContext::new(PathBuf::from("/tmp")),
404            )
405            .await
406            .unwrap();
407
408        assert!(output.success);
409        assert!(output.content.contains("Search summary: no matches found."));
410    }
411
412    #[tokio::test]
413    async fn program_tool_summarizes_repo_map_files() {
414        let registry = Arc::new(ToolRegistry::new(PathBuf::from("/tmp")));
415        registry.register(Arc::new(LsTool));
416        registry.register(Arc::new(RepoMapTool));
417        let tool = ProgramTool::new(Arc::clone(&registry));
418        let output = tool
419            .execute(
420                &serde_json::json!({
421                    "name": "program_repo_map",
422                    "inputs": {
423                        "path": "."
424                    }
425                }),
426                &ToolContext::new(PathBuf::from("/tmp")),
427            )
428            .await
429            .unwrap();
430
431        assert!(output.success);
432        assert!(output
433            .content
434            .contains("Repo map summary: found project files: Cargo.toml, README.md."));
435    }
436
437    #[tokio::test]
438    async fn program_tool_rejects_unknown_program() {
439        let registry = Arc::new(ToolRegistry::new(PathBuf::from("/tmp")));
440        let tool = ProgramTool::new(registry);
441        let output = tool
442            .execute(
443                &serde_json::json!({ "name": "missing" }),
444                &ToolContext::new(PathBuf::from("/tmp")),
445            )
446            .await
447            .unwrap();
448
449        assert!(!output.success);
450        assert!(output.content.contains("Unknown program"));
451    }
452
453    #[tokio::test]
454    async fn program_tool_compacts_large_step_output_into_artifact() {
455        let registry = Arc::new(ToolRegistry::new(PathBuf::from("/tmp")));
456        registry.register(Arc::new(EchoGrepTool));
457        let tool = ProgramTool::new(Arc::clone(&registry));
458        let output = tool
459            .execute(
460                &serde_json::json!({
461                    "name": "program_code_search",
462                    "inputs": {
463                        "query": "large"
464                    }
465                }),
466                &ToolContext::new(PathBuf::from("/tmp")),
467            )
468            .await
469            .unwrap();
470
471        assert!(output.success);
472        assert!(output.content.contains("[program step output compacted:"));
473        let metadata = output.metadata.as_ref().expect("metadata");
474        assert_eq!(metadata["trace"]["steps"][0]["compacted"], true);
475        assert_eq!(
476            metadata["verification_hints"][1]["kind"],
477            "inspect_artifacts"
478        );
479        let trace_artifact_uri = metadata["trace"]["steps"][0]["artifact"]["artifact_uri"]
480            .as_str()
481            .expect("trace artifact uri");
482        assert_eq!(
483            metadata["verification_hints"][1]["evidence_uris"][0],
484            trace_artifact_uri
485        );
486        assert_eq!(
487            metadata["verification_report"]["checks"][1]["evidence_uris"][0],
488            trace_artifact_uri
489        );
490        let artifact_uri = output
491            .content
492            .split("Full step artifact: ")
493            .nth(1)
494            .and_then(|tail| tail.split('.').next())
495            .expect("artifact uri");
496        assert_eq!(trace_artifact_uri, artifact_uri);
497        let artifact = registry
498            .get_artifact(artifact_uri)
499            .expect("stored step artifact");
500        assert_eq!(artifact.content.len(), MAX_PROGRAM_STEP_OUTPUT_BYTES + 1);
501    }
502}