a_agent/tools/
runner.rs

1use std::collections::HashMap;
2use std::path::PathBuf;
3use std::sync::Arc;
4use std::time::Duration;
5
6use async_trait::async_trait;
7use futures_util::future::join_all;
8use serde::Deserialize;
9use tokio::sync::{Mutex, Semaphore};
10use tokio_util::sync::CancellationToken;
11
12use crate::model::{StreamEvent, ToolCall, ToolResult};
13use crate::provider::EventSink;
14
15use super::bash::{BashArgs, BashOptions, OutputSink, execute_bash_cancellable};
16use super::patch::{FileSnapshot, affected_paths, apply_patch_with_snapshots};
17use super::read::{ReadArgs, read_text_file_bounded};
18
19/// A tool result plus anything the harness needs that must not be shown to the
20/// model. `ToolResult::output` goes into the conversation, so file snapshots
21/// travel beside it rather than inside it.
22#[derive(Debug, Clone)]
23pub struct ToolOutcome {
24    pub result: ToolResult,
25    pub snapshots: Vec<FileSnapshot>,
26}
27
28impl From<ToolResult> for ToolOutcome {
29    fn from(result: ToolResult) -> Self {
30        Self {
31            result,
32            snapshots: Vec::new(),
33        }
34    }
35}
36
37#[async_trait]
38pub trait ToolExecutor: Send + Sync {
39    async fn execute(&self, call: ToolCall) -> ToolOutcome;
40
41    async fn execute_with(
42        &self,
43        call: ToolCall,
44        _events: EventSink,
45        _cancel: CancellationToken,
46    ) -> ToolOutcome {
47        self.execute(call).await
48    }
49}
50
51pub struct CoreToolExecutor {
52    cwd: PathBuf,
53    read_max_lines: usize,
54    max_output_bytes: usize,
55    snapshot_max_bytes: usize,
56    bash_options: BashOptions,
57}
58
59impl CoreToolExecutor {
60    pub fn new(
61        cwd: PathBuf,
62        read_max_lines: usize,
63        bash_timeout: Duration,
64        max_output_bytes: usize,
65    ) -> Self {
66        Self::with_snapshot_limit(
67            cwd,
68            read_max_lines,
69            BashOptions {
70                timeout: bash_timeout,
71                max_timeout: bash_timeout,
72                max_output_bytes,
73            },
74            max_output_bytes,
75            1024 * 1024,
76        )
77    }
78
79    pub fn with_snapshot_limit(
80        cwd: PathBuf,
81        read_max_lines: usize,
82        bash_options: BashOptions,
83        max_output_bytes: usize,
84        snapshot_max_bytes: usize,
85    ) -> Self {
86        Self {
87            cwd,
88            read_max_lines,
89            max_output_bytes,
90            snapshot_max_bytes,
91            bash_options,
92        }
93    }
94}
95
96#[derive(Deserialize)]
97struct PatchArgs {
98    patch: String,
99}
100
101#[async_trait]
102impl ToolExecutor for CoreToolExecutor {
103    async fn execute(&self, call: ToolCall) -> ToolOutcome {
104        self.execute_with(call, EventSink::default(), CancellationToken::new())
105            .await
106    }
107
108    async fn execute_with(
109        &self,
110        call: ToolCall,
111        events: EventSink,
112        cancel: CancellationToken,
113    ) -> ToolOutcome {
114        if cancel.is_cancelled() {
115            return ToolResult::error(call.id, "tool execution cancelled").into();
116        }
117        let call_id = call.id.clone();
118        let mut snapshots = Vec::new();
119        let result = match call.name.as_str() {
120            "read" => match serde_json::from_str::<ReadArgs>(&call.arguments) {
121                Ok(args) => {
122                    read_text_file_bounded(
123                        &self.cwd,
124                        &args,
125                        self.read_max_lines,
126                        self.max_output_bytes,
127                    )
128                    .await
129                }
130                Err(error) => Err(error.into()),
131            },
132            "apply_patch" => match patch_text(&call.arguments) {
133                Ok(patch) => apply_patch_with_snapshots(&self.cwd, &patch, self.snapshot_max_bytes)
134                    .await
135                    .map(|summary| {
136                        snapshots = summary.snapshots;
137                        summary
138                            .files
139                            .iter()
140                            .map(|file| {
141                                format!("{} (+{} -{})", file.path, file.added, file.removed)
142                            })
143                            .collect::<Vec<_>>()
144                            .join("\n")
145                    }),
146                Err(error) => Err(error),
147            },
148            "bash" => match serde_json::from_str::<BashArgs>(&call.arguments) {
149                Ok(args) => {
150                    let output_events = events.clone();
151                    let output_id = call_id.clone();
152                    let output_sink: OutputSink = Arc::new(move |delta| {
153                        output_events.emit(StreamEvent::ToolExecutionOutput {
154                            id: output_id.clone(),
155                            delta,
156                        });
157                    });
158                    execute_bash_cancellable(
159                        &self.cwd,
160                        &args,
161                        &self.bash_options,
162                        Some(output_sink),
163                        cancel,
164                    )
165                    .await
166                    .map(|result| {
167                        format!(
168                            "{}\n[exit code: {}]",
169                            result.output,
170                            result
171                                .exit_code
172                                .map_or_else(|| "signal".into(), |code| code.to_string())
173                        )
174                    })
175                }
176                Err(error) => Err(error.into()),
177            },
178            name => Err(anyhow::anyhow!(
179                "unknown tool '{name}'; available tools: read, apply_patch, bash"
180            )),
181        };
182        let result = match result {
183            Ok(output)
184                if call.name == "bash"
185                    && (output.contains("[bash cancelled]")
186                        || output.contains("[bash timed out after ")) =>
187            {
188                ToolResult::error(call.id, output)
189            }
190            Ok(output) => ToolResult::success(call.id, output),
191            Err(error) => ToolResult::error(call.id, error.to_string()),
192        };
193        ToolOutcome { result, snapshots }
194    }
195}
196
197pub struct ToolRunner {
198    executor: Arc<dyn ToolExecutor>,
199    semaphore: Arc<Semaphore>,
200    path_locks: Arc<Mutex<HashMap<String, Arc<Mutex<()>>>>>,
201}
202
203impl ToolRunner {
204    pub fn new(executor: Arc<dyn ToolExecutor>, max_parallel: usize) -> Self {
205        Self {
206            executor,
207            semaphore: Arc::new(Semaphore::new(max_parallel.max(1))),
208            path_locks: Arc::new(Mutex::new(HashMap::new())),
209        }
210    }
211
212    pub async fn execute(&self, calls: Vec<ToolCall>) -> Vec<ToolOutcome> {
213        self.execute_with(calls, EventSink::default(), CancellationToken::new())
214            .await
215    }
216
217    pub async fn execute_with(
218        &self,
219        calls: Vec<ToolCall>,
220        events: EventSink,
221        cancel: CancellationToken,
222    ) -> Vec<ToolOutcome> {
223        let futures = calls.into_iter().map(|call| {
224            let executor = self.executor.clone();
225            let semaphore = self.semaphore.clone();
226            let path_locks = self.path_locks.clone();
227            let events = events.clone();
228            let cancel = cancel.clone();
229            async move {
230                let _permit = semaphore.acquire_owned().await.expect("semaphore closed");
231                let paths = scheduling_paths(&call);
232                let locks = {
233                    let mut registry = path_locks.lock().await;
234                    paths
235                        .into_iter()
236                        .map(|path| registry.entry(path).or_default().clone())
237                        .collect::<Vec<_>>()
238                };
239                let mut guards = Vec::with_capacity(locks.len());
240                for lock in locks {
241                    guards.push(lock.lock_owned().await);
242                }
243                let result = executor.execute_with(call, events, cancel).await;
244                drop(guards);
245                result
246            }
247        });
248        join_all(futures).await
249    }
250}
251
252fn scheduling_paths(call: &ToolCall) -> Vec<String> {
253    if call.name != "apply_patch" {
254        return Vec::new();
255    }
256    patch_text(&call.arguments)
257        .and_then(|patch| affected_paths(&patch))
258        .unwrap_or_default()
259}
260
261fn patch_text(arguments: &str) -> anyhow::Result<String> {
262    if arguments.trim_start().starts_with("*** Begin Patch") {
263        return Ok(arguments.to_owned());
264    }
265    Ok(serde_json::from_str::<PatchArgs>(arguments)?.patch)
266}
a_agent/tools/runner.rs

a_agent/tools/
runner.rs
