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

deepstrike_sdk/runtime/
execution_plane.rs
