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::{Stream, StreamExt};
13
14use crate::Result;
15use crate::governance::Governance;
16use crate::knowledge::KnowledgeSource;
17use crate::memory::MemoryStore;
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};
23use deepstrike_core::mm::memory::{MemoryQuery, MemoryScope};
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 memory_store: Option<&'a dyn MemoryStore>,
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        durable_content: None,
83        is_error,
84        is_fatal,
85        error_kind,
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.memory_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            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/// Collect tool results from a plane stream (one result per initial call).
595pub async fn collect_tool_results(
596    mut stream: Pin<Box<dyn Stream<Item = Result<RunEvent>> + Send>>,
597    calls: &[ToolCall],
598) -> Result<Vec<ToolResult>> {
599    let mut by_id: HashMap<String, ToolResult> = HashMap::new();
600    while let Some(evt) = stream.next().await {
601        if let RunEvent::ToolResult {
602            call_id,
603            content,
604            is_error,
605            is_fatal,
606            error_kind,
607        } = evt?
608        {
609            by_id.insert(
610                call_id.clone(),
611                make_result(
612                    compact_str::CompactString::new(&call_id),
613                    content,
614                    is_error,
615                    is_fatal,
616                    error_kind,
617                ),
618            );
619        }
620    }
621    Ok(calls
622        .iter()
623        .filter_map(|c| by_id.remove(c.id.as_str()))
624        .collect())
625}
deepstrike_sdk/runtime/execution_plane.rs

deepstrike_sdk/runtime/
execution_plane.rs
