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

deepstrike_sdk/runtime/
execution_plane.rs
