1use std::future::Future;
2use std::path::Path;
3use std::pin::Pin;
4
5use crate::file_claims::{ClaimResult, FileClaimRegistry, FileStaleEvent};
6use crate::network::NetworkPolicy;
7use crate::panel::PanelSink;
8use crate::sandbox::SandboxSettings;
9use crate::swarm_hub::SwarmHub;
10use crate::types::{PermissionSet, ToolDefinition, ToolResult};
11
12#[derive(Clone)]
14pub struct ToolContext {
15 pub working_dir: String,
16 pub session_id: Option<String>,
17 pub sandbox: SandboxSettings,
19 pub network: NetworkPolicy,
21 pub file_claims: Option<FileClaimRegistry>,
23 pub agent_id: Option<String>,
25 pub agent_label: Option<String>,
27 pub file_index: Option<std::sync::Arc<whycodes_index::WorkspaceIndex>>,
30 pub panel: Option<PanelSink>,
32 pub todo_sink: Option<crate::todo::TodoSink>,
34 pub swarm_hub: Option<SwarmHub>,
36}
37
38impl std::fmt::Debug for ToolContext {
39 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
40 f.debug_struct("ToolContext")
41 .field("working_dir", &self.working_dir)
42 .field("session_id", &self.session_id)
43 .field("sandbox", &self.sandbox)
44 .field("network", &self.network)
45 .field("file_claims", &self.file_claims)
46 .field("agent_id", &self.agent_id)
47 .field("agent_label", &self.agent_label)
48 .field("file_index", &self.file_index.as_ref().map(|_| "Some"))
49 .field("panel", &self.panel.as_ref().map(|_| "Some"))
50 .field("todo_sink", &self.todo_sink.as_ref().map(|_| "Some"))
51 .field("swarm_hub", &self.swarm_hub)
52 .finish()
53 }
54}
55
56impl ToolContext {
57 pub fn new(working_dir: impl Into<String>) -> Self {
58 Self {
59 working_dir: working_dir.into(),
60 session_id: None,
61 sandbox: SandboxSettings::default(),
62 network: NetworkPolicy::unrestricted(),
63 file_claims: None,
64 agent_id: None,
65 agent_label: None,
66 file_index: None,
67 panel: None,
68 todo_sink: None,
69 swarm_hub: None,
70 }
71 }
72
73 pub fn unsandboxed(working_dir: impl Into<String>) -> Self {
74 Self {
75 working_dir: working_dir.into(),
76 session_id: None,
77 sandbox: SandboxSettings::off(),
78 network: NetworkPolicy::unrestricted(),
79 file_claims: None,
80 agent_id: None,
81 agent_label: None,
82 file_index: None,
83 panel: None,
84 todo_sink: None,
85 swarm_hub: None,
86 }
87 }
88
89 pub fn check_file_write(&self, path: &Path) -> Result<(), String> {
94 let Some(reg) = self.file_claims.as_ref() else {
95 return Ok(());
96 };
97 let Some(id) = self.agent_id.as_deref() else {
98 return Ok(());
99 };
100 let label = self.agent_label.as_deref().unwrap_or(id);
101 match reg.try_claim(id, label, path) {
102 ClaimResult::Acquired | ClaimResult::Held => Ok(()),
103 ClaimResult::Conflict {
104 owner_label,
105 owner_id: _,
106 } => {
107 let shown = path.display();
108 Err(format!(
109 "File conflict: `{shown}` is claimed by swarm agent `{owner_label}`. \
110 Choose a different file, or wait for that agent to finish."
111 ))
112 }
113 }
114 }
115
116 pub fn check_file_read(&self, path: &Path) -> Option<FileStaleEvent> {
119 let reg = self.file_claims.as_ref()?;
120 let id = self.agent_id.as_deref()?;
121 reg.note_read(id, path)
122 }
123}
124
125pub type ToolFuture<'a> = Pin<Box<dyn Future<Output = ToolResult> + Send + 'a>>;
129
130pub trait Tool: Send + Sync {
132 fn name(&self) -> &str;
134
135 fn description(&self) -> &str;
137
138 fn parameters(&self) -> serde_json::Value;
140
141 fn execute<'a>(&'a self, args: serde_json::Value, ctx: &'a ToolContext) -> ToolFuture<'a>;
143
144 fn definition(&self) -> ToolDefinition {
146 ToolDefinition {
147 name: self.name().to_string(),
148 description: self.description().to_string(),
149 parameters: self.parameters(),
150 }
151 }
152
153 fn is_allowed(&self, permissions: &PermissionSet) -> bool {
159 permissions.is_tool_allowed(self.name())
160 }
161}
162
163#[cfg(test)]
164mod tests {
165 use super::*;
166 use crate::file_claims::FileClaimRegistry;
167 use crate::types::PermissionSet;
168 use std::path::Path;
169
170 struct DummyTool;
171
172 impl Tool for DummyTool {
173 fn name(&self) -> &str {
174 "dummy"
175 }
176 fn description(&self) -> &str {
177 "d"
178 }
179 fn parameters(&self) -> serde_json::Value {
180 serde_json::json!({"type": "object"})
181 }
182 fn execute<'a>(
183 &'a self,
184 _args: serde_json::Value,
185 _ctx: &'a ToolContext,
186 ) -> ToolFuture<'a> {
187 Box::pin(async {
188 crate::types::ToolResult {
189 tool_call_id: "1".into(),
190 content: "ok".into(),
191 is_error: false,
192 }
193 })
194 }
195 }
196
197 #[tokio::test]
198 async fn context_debug_claims_and_tool_trait() {
199 let ctx = ToolContext::new("/w");
200 assert_eq!(ctx.working_dir, "/w");
201 assert!(format!("{ctx:?}").contains("working_dir"));
202 let unsandboxed = ToolContext::unsandboxed("/u");
203 assert_eq!(unsandboxed.sandbox.mode, crate::sandbox::SandboxMode::Off);
204 assert!(ctx.check_file_write(Path::new("a.rs")).is_ok());
205 assert!(ctx.check_file_read(Path::new("a.rs")).is_none());
206
207 let mut claimed = ToolContext::new("/w");
208 claimed.file_claims = Some(FileClaimRegistry::new());
209 assert!(claimed.check_file_write(Path::new("a.rs")).is_ok());
210 claimed.agent_id = Some("w0".into());
211 claimed.agent_label = Some("worker".into());
212 assert!(claimed.check_file_write(Path::new("a.rs")).is_ok());
213 assert!(claimed.check_file_read(Path::new("a.rs")).is_none());
214
215 let mut other = ToolContext::new("/w");
216 other.file_claims = claimed.file_claims.clone();
217 other.agent_id = Some("w1".into());
218 other.agent_label = Some("other".into());
219 let err = other.check_file_write(Path::new("a.rs")).unwrap_err();
220 assert!(err.contains("File conflict"));
221
222 let tool = DummyTool;
223 let def = tool.definition();
224 assert_eq!(def.name, "dummy");
225 let result = tool.execute(serde_json::json!({}), &ctx).await;
226 assert!(!result.is_error);
227 assert_eq!(result.content, "ok");
228 let perms = PermissionSet {
229 allowed_tools: None,
230 denied_tools: None,
231 allow_file_writes: true,
232 allow_network: true,
233 allow_shell: true,
234 allowed_paths: None,
235 rules: Default::default(),
236 };
237 assert!(tool.is_allowed(&perms));
238 }
239}