Skip to main content

deepstrike_sdk/runtime/
execution_plane.rs

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
59/// Per-run context passed into `ExecutionPlane::execute_all`.
60pub 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
87/// Guarantees exactly one `tool_result` event per dispatched `ToolCall`.
88pub trait ExecutionPlane: Send + Sync {
89    fn schemas(&self) -> Vec<ToolSchema>;
90
91    /// Execute a batch of calls. Yields intermediate events; ends with one `ToolResult` per call.
92    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
99/// Executes tools in-process from a registry of `RegisteredTool`s.
100pub 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
616/// Collect tool results from a plane stream (one result per initial call).
617pub 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}