1use std::collections::HashMap;
2use std::path::Path;
3use std::pin::Pin;
4use std::sync::Arc;
5use tokio::sync::Mutex;
6
7use async_stream::try_stream;
8use deepstrike_core::context::manager::{KNOWLEDGE_TOOL_NAME, MEMORY_TOOL_NAME};
9use deepstrike_core::context::skill_catalog::SKILL_TOOL_NAME;
10use deepstrike_core::types::message::{Content, ToolCall, ToolResult, ToolSchema};
11use futures::stream::FuturesUnordered;
12use futures::stream::{self, Stream, StreamExt};
13
14use crate::Result;
15use crate::governance::Governance;
16use crate::knowledge::KnowledgeSource;
17use crate::memory::DreamStore;
18use crate::run_event::RunEvent;
19use crate::runtime::sandboxed_skill::{
20 PythonSkillPolicy, SkillKind, execute_json_skill, execute_python_skill, resolve_skill_path,
21};
22use crate::tools::{RegisteredTool, ToolChunk, ToolStep, validate_tool_arguments};
23
24#[derive(Clone)]
25pub struct ToolSuspendRequest {
26 pub call_id: String,
27 pub name: String,
28 pub suspension_id: String,
29 pub payload: Option<serde_json::Value>,
30}
31
32pub type ToolSuspendHandler = std::sync::Arc<
33 dyn Fn(ToolSuspendRequest) -> futures::future::BoxFuture<'static, Result<serde_json::Value>>
34 + Send
35 + Sync,
36>;
37
38#[derive(Clone)]
39pub struct PermissionRequest {
40 pub call_id: String,
41 pub tool_name: String,
42 pub arguments: String,
43 pub reason: String,
44}
45
46#[derive(Clone)]
47pub struct PermissionResponse {
48 pub approved: bool,
49 pub responder: String,
50 pub reason: Option<String>,
51}
52
53pub type PermissionRequestHandler = std::sync::Arc<
54 dyn Fn(PermissionRequest) -> futures::future::BoxFuture<'static, Result<PermissionResponse>>
55 + Send
56 + Sync,
57>;
58
59pub struct RunContext<'a> {
61 pub agent_id: Option<&'a str>,
62 pub skill_dir: Option<&'a Path>,
63 pub dream_store: Option<&'a dyn DreamStore>,
64 pub knowledge_source: Option<&'a dyn KnowledgeSource>,
65 pub governance: Option<Arc<Mutex<Governance>>>,
66 pub on_tool_suspend: Option<ToolSuspendHandler>,
67 pub on_permission_request: Option<PermissionRequestHandler>,
68}
69
70fn make_result(
71 call_id: compact_str::CompactString,
72 output: String,
73 is_error: bool,
74 is_fatal: bool,
75 error_kind: Option<deepstrike_core::types::message::ToolErrorKind>,
76) -> ToolResult {
77 ToolResult {
78 call_id,
79 output: Content::Text(output),
80 is_error,
81 is_fatal,
82 error_kind,
83 token_count: None,
84 }
85}
86
87pub trait ExecutionPlane: Send + Sync {
89 fn schemas(&self) -> Vec<ToolSchema>;
90
91 fn execute_all<'a>(
93 &'a self,
94 calls: &'a [ToolCall],
95 ctx: RunContext<'a>,
96 ) -> Pin<Box<dyn Stream<Item = Result<RunEvent>> + Send + 'a>>;
97}
98
99pub struct LocalExecutionPlane {
101 tools: HashMap<String, Arc<RegisteredTool>>,
102}
103
104impl LocalExecutionPlane {
105 pub fn new() -> Self {
106 Self {
107 tools: HashMap::new(),
108 }
109 }
110
111 pub fn register(&mut self, tool: RegisteredTool) -> &mut Self {
112 self.tools
113 .insert(tool.schema.name.to_string(), Arc::new(tool));
114 self
115 }
116
117 pub fn unregister(&mut self, name: &str) -> &mut Self {
118 self.tools.remove(name);
119 self
120 }
121}
122
123impl Default for LocalExecutionPlane {
124 fn default() -> Self {
125 Self::new()
126 }
127}
128
129impl ExecutionPlane for LocalExecutionPlane {
130 fn schemas(&self) -> Vec<ToolSchema> {
131 self.tools.values().map(|t| t.schema.clone()).collect()
132 }
133
134 fn execute_all<'a>(
135 &'a self,
136 calls: &'a [ToolCall],
137 ctx: RunContext<'a>,
138 ) -> Pin<Box<dyn Stream<Item = Result<RunEvent>> + Send + 'a>> {
139 Box::pin(execute_all_local(self, calls, ctx))
140 }
141}
142
143fn execute_all_local<'a>(
144 plane: &'a LocalExecutionPlane,
145 calls: &'a [ToolCall],
146 ctx: RunContext<'a>,
147) -> Pin<Box<dyn Stream<Item = Result<RunEvent>> + Send + 'a>> {
148 Box::pin(try_stream! {
149 let mut permitted = Vec::new();
150 for c in calls {
151 if let Some(gov) = &ctx.governance {
152 let mut g = gov.lock().await;
153 let now_ms = std::time::SystemTime::now()
154 .duration_since(std::time::UNIX_EPOCH)
155 .unwrap_or_default()
156 .as_millis() as u64;
157 g.set_time(now_ms);
158 let args_str =
159 serde_json::to_string(&c.arguments).unwrap_or_else(|_| "{}".to_string());
160 let verdict = g.evaluate(c.name.as_str(), &args_str);
161 match verdict.kind.as_str() {
162 "deny" => {
163 let reason = verdict.reason.unwrap_or_default();
164 yield RunEvent::ToolDenied {
165 call_id: c.id.to_string(),
166 tool_name: c.name.to_string(),
167 reason: reason.clone(),
168 };
169 yield RunEvent::ToolResult {
170 call_id: c.id.to_string(),
171 content: format!("permission denied: {reason}"),
172 is_error: true,
173 is_fatal: false,
174 error_kind: Some(deepstrike_core::types::message::ToolErrorKind::GovernanceDenied),
175 };
176 continue;
177 }
178 "rate_limited" => {
179 let reason = "rate limited".to_string();
180 yield RunEvent::ToolDenied {
181 call_id: c.id.to_string(),
182 tool_name: c.name.to_string(),
183 reason: reason.clone(),
184 };
185 yield RunEvent::ToolResult {
186 call_id: c.id.to_string(),
187 content: reason,
188 is_error: true,
189 is_fatal: false,
190 error_kind: Some(deepstrike_core::types::message::ToolErrorKind::Recoverable),
191 };
192 continue;
193 }
194 "ask_user" => {
195 let reason = verdict.reason.unwrap_or_default();
196 let args_str = serde_json::to_string(&c.arguments)
197 .unwrap_or_else(|_| "{}".to_string());
198 let request = PermissionRequest {
199 call_id: c.id.to_string(),
200 tool_name: c.name.to_string(),
201 arguments: args_str.clone(),
202 reason: reason.clone(),
203 };
204 yield RunEvent::PermissionRequest {
205 call_id: request.call_id.clone(),
206 tool_name: request.tool_name.clone(),
207 arguments: args_str,
208 reason: reason.clone(),
209 };
210
211 let decision = resolve_permission_request(request, &ctx).await;
212 yield RunEvent::PermissionResolved {
213 call_id: c.id.to_string(),
214 tool_name: c.name.to_string(),
215 approved: decision.approved,
216 responder: decision.responder.clone(),
217 reason: decision.reason.clone(),
218 };
219 if decision.approved {
220 permitted.push(c.clone());
221 continue;
222 }
223
224 let denied_reason = decision.reason.unwrap_or_else(|| {
225 if reason.is_empty() { "permission denied".to_string() } else { reason.clone() }
226 });
227 yield RunEvent::ToolDenied {
228 call_id: c.id.to_string(),
229 tool_name: c.name.to_string(),
230 reason: denied_reason.clone(),
231 };
232 yield RunEvent::ToolResult {
233 call_id: c.id.to_string(),
234 content: format!("permission denied: {denied_reason}"),
235 is_error: true,
236 is_fatal: false,
237 error_kind: Some(deepstrike_core::types::message::ToolErrorKind::GovernanceDenied),
238 };
239 continue;
240 }
241 _ => {}
242 }
243 }
244 permitted.push(c.clone());
245 }
246
247 let skill_calls: Vec<_> = permitted
248 .iter()
249 .filter(|c| c.name.as_str() == SKILL_TOOL_NAME)
250 .cloned()
251 .collect();
252 let memory_calls: Vec<_> = permitted
253 .iter()
254 .filter(|c| c.name.as_str() == MEMORY_TOOL_NAME)
255 .cloned()
256 .collect();
257 let knowledge_calls: Vec<_> = permitted
258 .iter()
259 .filter(|c| c.name.as_str() == KNOWLEDGE_TOOL_NAME)
260 .cloned()
261 .collect();
262 let regular_calls: Vec<_> = permitted
263 .iter()
264 .filter(|c| {
265 !matches!(
266 c.name.as_str(),
267 SKILL_TOOL_NAME | MEMORY_TOOL_NAME | KNOWLEDGE_TOOL_NAME
268 )
269 })
270 .cloned()
271 .collect();
272
273 for c in skill_calls {
274 let name = c.arguments.get("name").and_then(|v| v.as_str()).unwrap_or("").to_string();
275 let args: std::collections::HashMap<String, serde_json::Value> = c
276 .arguments
277 .as_object()
278 .map(|m| m.iter().map(|(k, v)| (k.clone(), v.clone())).collect())
279 .unwrap_or_default();
280
281 let (content, is_error) = if let Some(dir) = ctx.skill_dir {
282 match resolve_skill_path(dir, &name) {
283 Some((path, SkillKind::Prompt)) => {
284 match tokio::fs::read_to_string(&path).await {
285 Ok(content) => (strip_frontmatter(&content).to_string(), false),
286 Err(e) => (format!("error reading skill \"{name}\": {e}"), true),
287 }
288 }
289 Some((path, SkillKind::ComputeJson)) => execute_json_skill(&path, &args),
290 Some((path, SkillKind::PythonScript)) => {
291 execute_python_skill(&path, &args, None, &PythonSkillPolicy::default()).await
292 }
293 None => (format!("Skill \"{name}\" not found."), true),
294 }
295 } else {
296 ("No skill directory configured.".into(), true)
297 };
298 let error_kind = if is_error {
299 Some(deepstrike_core::types::message::ToolErrorKind::Recoverable)
300 } else {
301 None
302 };
303 yield RunEvent::ToolResult {
304 call_id: c.id.to_string(),
305 content,
306 is_error,
307 is_fatal: false,
308 error_kind,
309 };
310 }
311
312 for c in memory_calls {
313 let query = c.arguments.get("query").and_then(|v| v.as_str()).unwrap_or("").to_string();
314 let top_k = c.arguments.get("top_k").and_then(|v| v.as_u64()).unwrap_or(5) as usize;
315 let (content, is_error) = match (ctx.dream_store, ctx.agent_id) {
316 (Some(store), Some(agent_id)) => match store.search(agent_id, &query, top_k).await {
317 Ok(entries) if !entries.is_empty() => {
318 let text = entries
319 .iter()
320 .map(|e| format!("[score={:.3}] {}", e.score, e.text))
321 .collect::<Vec<_>>()
322 .join("\n---\n");
323 (text, false)
324 }
325 Ok(_) => ("No relevant memories found.".into(), false),
326 Err(e) => (format!("Memory search error: {e}"), true),
327 },
328 _ => ("Memory retrieval not configured.".into(), true),
329 };
330 let error_kind = if is_error {
331 Some(deepstrike_core::types::message::ToolErrorKind::Recoverable)
332 } else {
333 None
334 };
335 yield RunEvent::ToolResult {
336 call_id: c.id.to_string(),
337 content,
338 is_error,
339 is_fatal: false,
340 error_kind,
341 };
342 }
343
344 for c in knowledge_calls {
345 let query = c.arguments.get("query").and_then(|v| v.as_str()).unwrap_or("").to_string();
346 let top_k = c.arguments.get("top_k").and_then(|v| v.as_u64()).unwrap_or(5) as usize;
347 let (content, is_error) = if let Some(ks) = ctx.knowledge_source {
348 match ks.retrieve(&query, top_k).await {
349 Ok(snippets) if !snippets.is_empty() => (snippets.join("\n---\n"), false),
350 Ok(_) => ("No relevant knowledge found.".into(), false),
351 Err(e) => (format!("Knowledge retrieval error: {e}"), true),
352 }
353 } else {
354 ("Knowledge source not configured.".into(), true)
355 };
356 let error_kind = if is_error {
357 Some(deepstrike_core::types::message::ToolErrorKind::Recoverable)
358 } else {
359 None
360 };
361 yield RunEvent::ToolResult {
362 call_id: c.id.to_string(),
363 content,
364 is_error,
365 is_fatal: false,
366 error_kind,
367 };
368 }
369
370 if regular_calls.is_empty() {
371 return;
372 }
373
374 struct ActiveTool {
375 call_id: compact_str::CompactString,
376 name: String,
377 session: Box<dyn crate::tools::ToolSession>,
378 resume_input: Option<serde_json::Value>,
379 combined: String,
380 }
381
382 let mut active: FuturesUnordered<
383 futures::future::BoxFuture<'_, (ActiveTool, crate::Result<ToolStep>)>,
384 > = FuturesUnordered::new();
385
386 for mut call in regular_calls {
387 if let Some(content) = try_read_spooled_argument(&call).await {
388 yield RunEvent::ToolResult {
389 call_id: call.id.to_string(),
390 content,
391 is_error: false,
392 is_fatal: false,
393 error_kind: None,
394 };
395 continue;
396 }
397
398 let Some(tool) = plane.tools.get(call.name.as_str()) else {
399 let content = format!("unknown tool: {}", call.name);
400 yield RunEvent::ToolResult {
401 call_id: call.id.to_string(),
402 content: content.clone(),
403 is_error: true,
404 is_fatal: false,
405 error_kind: Some(deepstrike_core::types::message::ToolErrorKind::Recoverable),
406 };
407 continue;
408 };
409 let original_args_str = serde_json::to_string(&call.arguments).unwrap_or_default();
410 match validate_tool_arguments(&tool.schema.parameters, &mut call.arguments) {
411 Ok(repaired) => {
412 if repaired {
413 let repaired_args_str = serde_json::to_string(&call.arguments).unwrap_or_default();
414 yield RunEvent::ToolArgumentRepaired {
415 call_id: call.id.to_string(),
416 name: call.name.to_string(),
417 original_arguments: original_args_str,
418 repaired_arguments: repaired_args_str,
419 };
420 }
421 }
422 Err(e) => {
423 let content = format!("invalid arguments: {e}");
424 yield RunEvent::ToolResult {
425 call_id: call.id.to_string(),
426 content: content.clone(),
427 is_error: true,
428 is_fatal: false,
429 error_kind: Some(deepstrike_core::types::message::ToolErrorKind::Recoverable),
430 };
431 continue;
432 }
433 }
434 let start = Arc::clone(&tool.start);
435 let args = call.arguments.clone();
436 let call_id = call.id.clone();
437 let name = call.name.to_string();
438 match start(args).await {
439 Ok(session) => {
440 let active_tool = ActiveTool {
441 call_id: call_id.clone(),
442 name: name.clone(),
443 session,
444 resume_input: None,
445 combined: String::new(),
446 };
447 active.push(Box::pin(async move {
448 let mut t = active_tool;
449 let step = t.session.next(t.resume_input.take()).await;
450 (t, step)
451 }));
452 }
453 Err(e) => {
454 let (is_fatal, error_kind) = match &e {
455 crate::Error::ToolExecutionFailed { is_fatal, error_kind, .. } => (*is_fatal, *error_kind),
456 crate::Error::ToolFail { is_fatal, error_kind, .. } => (*is_fatal, *error_kind),
457 _ => (false, Some(deepstrike_core::types::message::ToolErrorKind::Recoverable)),
458 };
459 yield RunEvent::ToolResult {
460 call_id: call_id.to_string(),
461 content: crate::format_tool_error(&e),
462 is_error: true,
463 is_fatal,
464 error_kind,
465 };
466 }
467 }
468 }
469
470 while let Some((mut tool, step)) = active.next().await {
471 match step {
472 Ok(ToolStep::Chunk(chunk)) => {
473 match &chunk {
474 ToolChunk::Suspend { suspension_id, payload } => {
475 yield RunEvent::ToolSuspend {
476 call_id: tool.call_id.to_string(),
477 name: tool.name.clone(),
478 suspension_id: suspension_id.clone(),
479 payload: payload.clone(),
480 };
481 match &ctx.on_tool_suspend {
482 Some(handler) => {
483 tool.resume_input = Some(
484 handler(ToolSuspendRequest {
485 call_id: tool.call_id.to_string(),
486 name: tool.name.clone(),
487 suspension_id: suspension_id.clone(),
488 payload: payload.clone(),
489 })
490 .await?,
491 );
492 }
493 None => {
494 let content = format!(
495 "tool suspended without resume handler: {suspension_id}"
496 );
497 yield RunEvent::ToolResult {
498 call_id: tool.call_id.to_string(),
499 content: content.clone(),
500 is_error: true,
501 is_fatal: false,
502 error_kind: Some(deepstrike_core::types::message::ToolErrorKind::Recoverable),
503 };
504 continue;
505 }
506 }
507 }
508 _ => {
509 tool.combined.push_str(chunk.text_projection());
510 yield RunEvent::ToolDelta {
511 call_id: tool.call_id.to_string(),
512 name: tool.name.clone(),
513 chunk,
514 };
515 }
516 }
517 active.push(Box::pin(async move {
518 let step = tool.session.next(tool.resume_input.take()).await;
519 (tool, step)
520 }));
521 }
522 Ok(ToolStep::Done(text)) => {
523 tool.combined.push_str(&text);
524 yield RunEvent::ToolResult {
525 call_id: tool.call_id.to_string(),
526 content: tool.combined.clone(),
527 is_error: false,
528 is_fatal: false,
529 error_kind: None,
530 };
531 }
532 Err(e) => {
533 let (is_fatal, error_kind) = match &e {
534 crate::Error::ToolExecutionFailed { is_fatal, error_kind, .. } => (*is_fatal, *error_kind),
535 crate::Error::ToolFail { is_fatal, error_kind, .. } => (*is_fatal, *error_kind),
536 _ => (false, Some(deepstrike_core::types::message::ToolErrorKind::Recoverable)),
537 };
538 yield RunEvent::ToolResult {
539 call_id: tool.call_id.to_string(),
540 content: crate::format_tool_error(&e),
541 is_error: true,
542 is_fatal,
543 error_kind,
544 };
545 }
546 }
547 }
548 })
549}
550
551fn strip_frontmatter(content: &str) -> &str {
552 let s = content.trim_start();
553 if !s.starts_with("---") {
554 return s;
555 }
556 let rest = &s[3..];
557 if let Some(end) = rest.find("\n---") {
558 rest[end + 4..].trim_start_matches('\n')
559 } else {
560 s
561 }
562}
563
564async fn resolve_permission_request(
565 request: PermissionRequest,
566 ctx: &RunContext<'_>,
567) -> PermissionResponse {
568 let Some(handler) = &ctx.on_permission_request else {
569 return PermissionResponse {
570 approved: false,
571 responder: "policy_gate".to_string(),
572 reason: Some("no permission handler configured".to_string()),
573 };
574 };
575
576 match handler(request).await {
577 Ok(response) => PermissionResponse {
578 approved: response.approved,
579 responder: if response.responder.is_empty() {
580 "host".to_string()
581 } else {
582 response.responder
583 },
584 reason: response.reason,
585 },
586 Err(err) => PermissionResponse {
587 approved: false,
588 responder: "permission_handler".to_string(),
589 reason: Some(format!("permission handler failed: {err}")),
590 },
591 }
592}
593
594
595async fn try_read_spooled_argument(call: &ToolCall) -> Option<String> {
596 let is_read_tool = matches!(
597 call.name.as_str(),
598 "read" | "read_file" | "view_file" | "read_spooled_result"
599 );
600 if !is_read_tool {
601 return None;
602 }
603 let obj = call.arguments.as_object()?;
604 for val in obj.values() {
605 if let Some(s) = val.as_str() {
606 if s.starts_with(".spool/") || s.contains("/.spool/") {
607 if let Ok(content) = tokio::fs::read_to_string(s).await {
608 return Some(content);
609 }
610 }
611 }
612 }
613 None
614}
615
616pub async fn collect_tool_results(
618 mut stream: Pin<Box<dyn Stream<Item = Result<RunEvent>> + Send>>,
619 calls: &[ToolCall],
620) -> Result<Vec<ToolResult>> {
621 let mut by_id: HashMap<String, ToolResult> = HashMap::new();
622 while let Some(evt) = stream.next().await {
623 if let RunEvent::ToolResult {
624 call_id,
625 content,
626 is_error,
627 is_fatal,
628 error_kind,
629 } = evt?
630 {
631 by_id.insert(
632 call_id.clone(),
633 make_result(
634 compact_str::CompactString::new(&call_id),
635 content,
636 is_error,
637 is_fatal,
638 error_kind,
639 ),
640 );
641 }
642 }
643 Ok(calls
644 .iter()
645 .filter_map(|c| by_id.remove(c.id.as_str()))
646 .collect())
647}