deepstrike_sdk/runtime/
remote_vpc_plane.rs

1use std::collections::HashSet;
2use std::pin::Pin;
3use std::sync::Arc;
4
5use async_stream::try_stream;
6use deepstrike_core::types::message::{ToolCall, ToolSchema};
7use futures::stream::Stream;
8use futures::StreamExt;
9
10use crate::run_event::RunEvent;
11use crate::runtime::credential_vault::CredentialVault;
12use crate::runtime::execution_plane::{ExecutionPlane, LocalExecutionPlane, RunContext};
13use crate::tools::RegisteredTool;
14use crate::Result;
15
16pub struct RemoteVpcOptions {
17    /// Base URL of the remote worker endpoint inside the customer VPC.
18    /// Expected route: `POST {base_url}/execute`  body: `{ name, arguments }`
19    ///                                            response: `{ output, isError }`
20    pub base_url: String,
21    pub vault: Arc<dyn CredentialVault>,
22    /// Vault key whose value is sent verbatim as the `Authorization` header.
23    /// Fetched fresh on every `execute_all` call — never stored in the session log.
24    pub auth_credential_key: Option<String>,
25    /// Static tool schemas served by this VPC worker.
26    pub schemas: Vec<ToolSchema>,
27    /// Per-call HTTP timeout in ms. Default: 30 000.
28    pub timeout_ms: u64,
29}
30
31/// `ExecutionPlane` that forwards tool calls over HTTP to a worker inside a customer VPC.
32///
33/// Credentials are fetched from a `CredentialVault` at call time and injected into the
34/// `Authorization` header — they are never forwarded to the model or stored in the session log.
35///
36/// Local tools registered via `register()` run in-process and take priority over any remote
37/// schema with the same name.
38pub struct RemoteVpcPlane {
39    base_url: String,
40    vault: Arc<dyn CredentialVault>,
41    auth_key: Option<String>,
42    remote_schemas: Vec<ToolSchema>,
43    client: reqwest::Client,
44    local: LocalExecutionPlane,
45}
46
47impl RemoteVpcPlane {
48    pub fn new(opts: RemoteVpcOptions) -> Self {
49        let client = reqwest::Client::builder()
50            .timeout(std::time::Duration::from_millis(opts.timeout_ms))
51            .build()
52            .unwrap_or_default();
53        Self {
54            base_url: opts.base_url.trim_end_matches('/').to_string(),
55            vault: opts.vault,
56            auth_key: opts.auth_credential_key,
57            remote_schemas: opts.schemas,
58            client,
59            local: LocalExecutionPlane::new(),
60        }
61    }
62
63    /// Register an in-process tool. Local tools take priority over remote schemas of the same name.
64    pub fn register(&mut self, tool: RegisteredTool) -> &mut Self {
65        self.local.register(tool);
66        self
67    }
68
69    pub fn unregister(&mut self, name: &str) -> &mut Self {
70        self.local.unregister(name);
71        self
72    }
73}
74
75impl ExecutionPlane for RemoteVpcPlane {
76    fn schemas(&self) -> Vec<ToolSchema> {
77        let local_names: HashSet<_> = self
78            .local
79            .schemas()
80            .into_iter()
81            .map(|s| s.name.clone())
82            .collect();
83        let mut result = self.local.schemas();
84        result.extend(
85            self.remote_schemas
86                .iter()
87                .filter(|s| !local_names.contains(&s.name))
88                .cloned(),
89        );
90        result
91    }
92
93    fn execute_all<'a>(
94        &'a self,
95        calls: &'a [ToolCall],
96        ctx: RunContext<'a>,
97    ) -> Pin<Box<dyn Stream<Item = Result<RunEvent>> + Send + 'a>> {
98        Box::pin(try_stream! {
99            let RunContext {
100                agent_id, memory_scope, skill_dir, dream_store, knowledge_source, governance, on_tool_suspend, on_permission_request,
101            } = ctx;
102
103            let local_names: HashSet<String> =
104                self.local.schemas().into_iter().map(|s| s.name.to_string()).collect();
105            let local_calls: Vec<_> =
106                calls.iter().filter(|c| local_names.contains(c.name.as_str())).cloned().collect();
107            let remote_calls: Vec<_> =
108                calls.iter().filter(|c| !local_names.contains(c.name.as_str())).cloned().collect();
109
110            if !local_calls.is_empty() {
111                let local_ctx = RunContext {
112                    agent_id,
113                    memory_scope,
114                    skill_dir,
115                    dream_store,
116                    knowledge_source,
117                    governance: governance.clone(),
118                    on_tool_suspend: on_tool_suspend.clone(),
119                    on_permission_request: on_permission_request.clone(),
120                };
121                let mut s = self.local.execute_all(&local_calls, local_ctx);
122                while let Some(evt) = s.next().await {
123                    yield evt?;
124                }
125            }
126
127            if !remote_calls.is_empty() {
128                let auth = if let Some(key) = &self.auth_key {
129                    self.vault.get(key).await
130                } else {
131                    None
132                };
133
134                // Fire all remote calls concurrently; yield results in dispatch order.
135                let futs: Vec<_> = remote_calls
136                    .iter()
137                    .map(|call| {
138                        let client = self.client.clone();
139                        let url = format!("{}/execute", self.base_url);
140                        let auth = auth.clone();
141                        let call = call.clone();
142                        async move { call_remote(client, url, auth, call).await }
143                    })
144                    .collect();
145
146                let results = futures::future::join_all(futs).await;
147                for (call, (content, is_error)) in remote_calls.iter().zip(results) {
148                    yield RunEvent::ToolResult {
149                        call_id: call.id.to_string(),
150                        content,
151                        is_error,
152                        is_fatal: false,
153                        error_kind: None,
154                    };
155                }
156            }
157        })
158    }
159}
160
161// ── HTTP helper ───────────────────────────────────────────────────────────────
162
163async fn call_remote(
164    client: reqwest::Client,
165    url: String,
166    auth: Option<String>,
167    call: ToolCall,
168) -> (String, bool) {
169    let mut req = client.post(&url).json(&serde_json::json!({
170        "name": call.name.as_str(),
171        "arguments": call.arguments,
172    }));
173    if let Some(token) = auth {
174        req = req.header("Authorization", token);
175    }
176
177    match req.send().await {
178        Ok(resp) => {
179            if !resp.status().is_success() {
180                let status = resp.status().as_u16();
181                let body = resp.text().await.unwrap_or_default();
182                let msg = if body.is_empty() {
183                    format!("HTTP {status}")
184                } else {
185                    format!("HTTP {status}: {body}")
186                };
187                return (msg, true);
188            }
189            match resp.json::<serde_json::Value>().await {
190                Ok(result) => {
191                    let output = result
192                        .get("output")
193                        .and_then(|v| v.as_str())
194                        .unwrap_or("")
195                        .to_string();
196                    let is_error = result
197                        .get("isError")
198                        .and_then(|v| v.as_bool())
199                        .unwrap_or(false);
200                    (output, is_error)
201                }
202                Err(e) => (e.to_string(), true),
203            }
204        }
205        Err(e) => (e.to_string(), true),
206    }
207}
deepstrike_sdk/runtime/remote_vpc_plane.rs

deepstrike_sdk/runtime/
remote_vpc_plane.rs
