1use std::collections::BTreeMap;
2use std::fmt;
3use std::path::PathBuf;
4use std::sync::Arc;
5
6use serde::{Deserialize, Serialize};
7
8use crate::artifacts::ContextArtifactAccess;
9use crate::discovery::{
10 DiscoveryAuthState, DiscoveryCacheStatus, DiscoveryCatalogItem, DiscoveryCatalogSource,
11 DiscoveryItemStatus, DiscoveryLifecycleState, DiscoveryPromotionState, DiscoveryRedaction,
12 DiscoverySchemaFormat, DiscoverySchemaReference, DiscoverySourceKind,
13};
14use crate::events::{ThreadId, TurnId};
15use crate::extension::ToolProviderId;
16use crate::goals::ThreadGoalController;
17use crate::inference::ModelSchemaPolicy;
18use crate::policy_mode::PolicyMode;
19use crate::remote_runner::RemoteWorkspace;
20use crate::trace::SubagentTraceSink;
21use crate::{ToolSchemaPolicy, normalize_tool_schema};
22
23#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
24pub struct ToolSpec {
25 pub name: String,
26 pub description: String,
27 pub parameters: serde_json::Value,
28}
29
30impl ToolSpec {
31 pub fn freeform_input_field(&self) -> Option<&'static str> {
42 Self::freeform_field_for_name(&self.name)
43 }
44
45 pub fn freeform_field_for_name(name: &str) -> Option<&'static str> {
46 match name {
47 "apply_patch" => Some("patch"),
48 _ => None,
49 }
50 }
51
52 pub fn normalized_for_model(&self, policy: ToolSchemaPolicy) -> Self {
53 let mut spec = self.clone();
54 spec.parameters = normalize_tool_schema(&spec.name, &spec.parameters, policy).schema;
55 spec
56 }
57
58 pub fn normalized_for_model_profile(&self, policy: ModelSchemaPolicy) -> Self {
59 match policy {
60 ModelSchemaPolicy::StandardRequiredFirst => {
61 self.normalized_for_model(ToolSchemaPolicy::warning())
62 }
63 ModelSchemaPolicy::RequiredFirstFlat => {
64 self.normalized_for_model(ToolSchemaPolicy::strict())
65 }
66 }
67 }
68
69 pub fn discovery_item(
70 &self,
71 provider_id: impl Into<String>,
72 schema_uri: impl Into<String>,
73 ) -> DiscoveryCatalogItem {
74 let provider_id = provider_id.into();
75 DiscoveryCatalogItem {
76 id: format!("tool:{provider_id}/{}", self.name),
77 group_id: format!("tools:{provider_id}"),
78 source: DiscoveryCatalogSource {
79 kind: DiscoverySourceKind::InternalTools,
80 id: provider_id.clone(),
81 display_name: provider_id,
82 origin: None,
83 auth_state: DiscoveryAuthState::NotRequired,
84 redaction: DiscoveryRedaction::none(),
85 },
86 name: self.name.clone(),
87 title: self.name.clone(),
88 description: Some(self.description.clone()),
89 status: DiscoveryItemStatus::Available,
90 lifecycle: DiscoveryLifecycleState::Discovered,
91 promotion: DiscoveryPromotionState::NotPromoted,
92 cache_status: DiscoveryCacheStatus::Cold,
93 schema: Some(DiscoverySchemaReference {
94 format: DiscoverySchemaFormat::JsonSchema,
95 uri: schema_uri.into(),
96 content_hash: None,
97 byte_count: None,
98 redaction: DiscoveryRedaction::none(),
99 }),
100 tags: vec!["tool".to_string()],
101 hints: Vec::new(),
102 redaction: DiscoveryRedaction::none(),
103 last_refreshed_at: None,
104 }
105 }
106}
107
108#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
109pub enum ToolChoice {
110 Auto,
111 Any,
112 None,
113 Specific(String),
114}
115
116#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
117pub struct ToolCall {
118 pub id: String,
119 pub name: String,
120 pub arguments: serde_json::Value,
121 pub raw_arguments: String,
122 pub thread_id: ThreadId,
123 pub turn_id: TurnId,
124}
125
126#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
127pub struct ToolResult {
128 pub id: String,
129 pub name: String,
130 pub text: String,
131 pub data: serde_json::Value,
132 pub is_error: bool,
133}
134
135#[derive(Clone, Default)]
136pub struct ToolExecutionHandles {
137 pub workspace: Option<Arc<dyn ScopedWorkspaceHandle>>,
138 pub remote_workspace: Option<Arc<RemoteWorkspace>>,
144 pub process_runner: Option<Arc<dyn ScopedProcessRunner>>,
145 pub subagent_trace_sink: Option<Arc<dyn SubagentTraceSink>>,
146 pub context_artifacts: Option<Arc<dyn ContextArtifactAccess>>,
147 pub goal_controller: Option<Arc<dyn ThreadGoalController>>,
148 pub swarm_progress_sink: Option<Arc<dyn crate::subagents::AgentSwarmProgressSink>>,
150 pub parent_model_selection: Option<crate::inference::ModelSelection>,
154}
155
156impl fmt::Debug for ToolExecutionHandles {
157 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
158 f.debug_struct("ToolExecutionHandles")
159 .field("workspace", &self.workspace.is_some())
160 .field("remote_workspace", &self.remote_workspace.is_some())
161 .field("process_runner", &self.process_runner.is_some())
162 .field("subagent_trace_sink", &self.subagent_trace_sink.is_some())
163 .field("context_artifacts", &self.context_artifacts.is_some())
164 .field("goal_controller", &self.goal_controller.is_some())
165 .field("swarm_progress_sink", &self.swarm_progress_sink.is_some())
166 .field("parent_model_selection", &self.parent_model_selection)
167 .finish()
168 }
169}
170
171pub trait ScopedWorkspaceHandle: Send + Sync + 'static {
172 fn workspace_root(&self) -> Option<PathBuf>;
173}
174
175pub trait ScopedProcessRunner: Send + Sync + 'static {
176 fn runner_name(&self) -> &str;
177}
178
179#[derive(Debug, Clone)]
180pub struct LocalWorkspaceHandle {
181 root: PathBuf,
182}
183
184impl LocalWorkspaceHandle {
185 pub fn new(root: impl Into<PathBuf>) -> Self {
186 Self { root: root.into() }
187 }
188}
189
190impl ScopedWorkspaceHandle for LocalWorkspaceHandle {
191 fn workspace_root(&self) -> Option<PathBuf> {
192 Some(self.root.clone())
193 }
194}
195
196#[derive(Debug, Clone, Default)]
197pub struct LocalProcessRunnerHandle;
198
199impl ScopedProcessRunner for LocalProcessRunnerHandle {
200 fn runner_name(&self) -> &str {
201 "local-process"
202 }
203}
204
205#[derive(Debug, Clone)]
206pub struct ToolExecutionContext {
207 pub thread_id: ThreadId,
208 pub turn_id: TurnId,
209 pub effective_mode: PolicyMode,
210 pub command_shell: Option<String>,
211 pub deadline_remaining_seconds: Option<u64>,
212 pub handles: ToolExecutionHandles,
213}
214
215impl ToolExecutionContext {
216 pub fn new(
217 thread_id: impl Into<ThreadId>,
218 turn_id: impl Into<TurnId>,
219 effective_mode: PolicyMode,
220 ) -> Self {
221 Self {
222 thread_id: thread_id.into(),
223 turn_id: turn_id.into(),
224 effective_mode,
225 command_shell: None,
226 deadline_remaining_seconds: None,
227 handles: ToolExecutionHandles::default(),
228 }
229 }
230
231 pub fn with_command_shell(mut self, shell: impl Into<String>) -> Self {
232 let shell = shell.into();
233 if !shell.trim().is_empty() {
234 self.command_shell = Some(shell);
235 }
236 self
237 }
238
239 pub fn with_deadline_remaining_seconds(mut self, seconds: u64) -> Self {
240 self.deadline_remaining_seconds = Some(seconds);
241 self
242 }
243
244 pub fn with_workspace_handle(mut self, handle: Arc<dyn ScopedWorkspaceHandle>) -> Self {
245 self.handles.workspace = Some(handle);
246 self
247 }
248
249 pub fn with_remote_workspace(mut self, remote: Arc<RemoteWorkspace>) -> Self {
250 self.handles.remote_workspace = Some(remote);
251 self
252 }
253
254 pub fn with_process_runner(mut self, runner: Arc<dyn ScopedProcessRunner>) -> Self {
255 self.handles.process_runner = Some(runner);
256 self
257 }
258
259 pub fn with_subagent_trace_sink(mut self, sink: Arc<dyn SubagentTraceSink>) -> Self {
260 self.handles.subagent_trace_sink = Some(sink);
261 self
262 }
263
264 pub fn with_context_artifacts(mut self, store: Arc<dyn ContextArtifactAccess>) -> Self {
265 self.handles.context_artifacts = Some(store);
266 self
267 }
268
269 pub fn with_goal_controller(mut self, controller: Arc<dyn ThreadGoalController>) -> Self {
270 self.handles.goal_controller = Some(controller);
271 self
272 }
273
274 pub fn with_parent_model_selection(
275 mut self,
276 selection: crate::inference::ModelSelection,
277 ) -> Self {
278 self.handles.parent_model_selection = Some(selection);
279 self
280 }
281
282 pub fn with_swarm_progress_sink(
283 mut self,
284 sink: Arc<dyn crate::subagents::AgentSwarmProgressSink>,
285 ) -> Self {
286 self.handles.swarm_progress_sink = Some(sink);
287 self
288 }
289
290 pub fn require_workspace(&self) -> anyhow::Result<Arc<dyn ScopedWorkspaceHandle>> {
291 self.handles
292 .workspace
293 .clone()
294 .ok_or_else(|| anyhow::anyhow!("workspace handle is not available"))
295 }
296
297 pub fn require_process_runner(&self) -> anyhow::Result<Arc<dyn ScopedProcessRunner>> {
298 self.handles
299 .process_runner
300 .clone()
301 .ok_or_else(|| anyhow::anyhow!("process runner is not available"))
302 }
303
304 pub fn require_context_artifacts(&self) -> anyhow::Result<Arc<dyn ContextArtifactAccess>> {
305 self.handles
306 .context_artifacts
307 .clone()
308 .ok_or_else(|| anyhow::anyhow!("context artifact store is not available"))
309 }
310
311 pub fn require_goal_controller(&self) -> anyhow::Result<Arc<dyn ThreadGoalController>> {
312 self.handles
313 .goal_controller
314 .clone()
315 .ok_or_else(|| anyhow::anyhow!("goal controller is not available"))
316 }
317}
318
319#[async_trait::async_trait]
320pub trait ToolExecutor: Send + Sync + 'static {
321 fn spec(&self) -> ToolSpec;
322
323 fn supports_eager_execution(&self) -> bool {
327 false
328 }
329
330 async fn execute(
331 &self,
332 ctx: ToolExecutionContext,
333 call: ToolCall,
334 ) -> anyhow::Result<ToolResult>;
335}
336
337#[derive(Default, Clone)]
338pub struct ToolRegistry {
339 tools: BTreeMap<String, Arc<dyn ToolExecutor>>,
340}
341
342impl ToolRegistry {
343 pub fn register(&mut self, tool: Arc<dyn ToolExecutor>) -> anyhow::Result<()> {
344 let name = tool.spec().name;
345 if self.tools.contains_key(&name) {
346 anyhow::bail!("tool {name:?} is already registered");
347 }
348 self.tools.insert(name, tool);
349 Ok(())
350 }
351
352 pub fn replace(&mut self, tool: Arc<dyn ToolExecutor>) {
356 self.tools.insert(tool.spec().name, tool);
357 }
358
359 pub fn specs(&self) -> Vec<ToolSpec> {
360 self.tools
361 .values()
362 .map(|tool| {
363 tool.spec()
364 .normalized_for_model(ToolSchemaPolicy::warning())
365 })
366 .collect()
367 }
368
369 pub fn specs_for_edit_tool(&self, edit_tool: Option<&str>) -> Vec<ToolSpec> {
370 self.specs_for_edit_tool_with_schema_policy(edit_tool, ModelSchemaPolicy::RequiredFirstFlat)
371 }
372
373 pub fn specs_for_edit_tool_with_schema_policy(
374 &self,
375 edit_tool: Option<&str>,
376 schema_policy: ModelSchemaPolicy,
377 ) -> Vec<ToolSpec> {
378 self.tools
379 .values()
380 .map(|tool| tool.spec())
381 .filter(|spec| keep_tool_for_edit_tool(&spec.name, edit_tool))
382 .map(|spec| spec.normalized_for_model_profile(schema_policy))
383 .collect()
384 }
385
386 pub fn get(&self, name: &str) -> Option<Arc<dyn ToolExecutor>> {
387 self.tools.get(name).cloned()
388 }
389
390 pub fn is_empty(&self) -> bool {
391 self.tools.is_empty()
392 }
393}
394
395fn keep_tool_for_edit_tool(name: &str, edit_tool: Option<&str>) -> bool {
396 match name {
397 "apply_patch" => true,
398 "write_file" | "edit" | "multi_edit" => !matches!(edit_tool, Some("patch")),
399 _ => true,
400 }
401}
402
403pub trait ToolContributor: Send + Sync + 'static {
404 fn id(&self) -> ToolProviderId;
405 fn contribute(&self, registry: &mut ToolRegistry) -> anyhow::Result<()>;
406}
407
408#[cfg(test)]
409mod tests {
410 use super::*;
411
412 #[test]
413 fn tool_spec_can_be_represented_as_discovery_item() {
414 let spec = ToolSpec {
415 name: "grep".to_string(),
416 description: "Search files".to_string(),
417 parameters: serde_json::json!({
418 "type": "object",
419 "properties": {
420 "query": { "type": "string" }
421 },
422 "required": ["query"]
423 }),
424 };
425
426 let item = spec.discovery_item(
427 "builtin-coding-tools",
428 "discovery/tools/builtin-coding-tools/grep.schema.json",
429 );
430 assert_eq!(item.id, "tool:builtin-coding-tools/grep");
431 assert_eq!(item.group_id, "tools:builtin-coding-tools");
432 assert_eq!(item.source.kind, DiscoverySourceKind::InternalTools);
433 assert_eq!(item.source.auth_state, DiscoveryAuthState::NotRequired);
434 assert_eq!(item.status, DiscoveryItemStatus::Available);
435 assert_eq!(item.lifecycle, DiscoveryLifecycleState::Discovered);
436 assert_eq!(
437 item.schema.as_ref().map(|schema| schema.format.clone()),
438 Some(DiscoverySchemaFormat::JsonSchema)
439 );
440 }
441
442 #[test]
443 fn apply_patch_is_kept_for_all_edit_tool_profiles() {
444 assert!(keep_tool_for_edit_tool("apply_patch", None));
445 assert!(keep_tool_for_edit_tool("apply_patch", Some("edit")));
446 assert!(keep_tool_for_edit_tool("apply_patch", Some("patch")));
447
448 assert!(keep_tool_for_edit_tool("edit", None));
449 assert!(keep_tool_for_edit_tool("edit", Some("edit")));
450 assert!(!keep_tool_for_edit_tool("edit", Some("patch")));
451 assert!(!keep_tool_for_edit_tool("multi_edit", Some("patch")));
452 assert!(!keep_tool_for_edit_tool("write_file", Some("patch")));
453 }
454}