Skip to main content

soma_provider_adapters/
ai_sdk.rs

1//! The generic AI-SDK / sandboxed-TypeScript-handler provider kind: runs a
2//! drop-in `.ts` provider's `call(input)` export in a bounded Node sidecar.
3//! Ported from `soma-service::providers::ai_sdk` with product types swapped
4//! for their `soma-provider-core` equivalents.
5
6use std::{path::PathBuf, sync::Arc};
7
8use async_trait::async_trait;
9use serde_json::Value;
10use soma_provider_core::{
11    Provider, ProviderCall, ProviderCatalog, ProviderError, ProviderOutput, ProviderTool,
12};
13use tokio::time::Instant;
14
15use crate::{
16    error::{redact_public, SidecarError},
17    sidecar::{
18        collect_provider_env, execution_payload, output_exceeded_message, run_bounded_sidecar,
19    },
20};
21
22#[derive(Clone)]
23pub struct AiSdkProvider {
24    path: PathBuf,
25    catalog: ProviderCatalog,
26    /// Product env-namespace prefix (e.g. `"SOMA"`) applied to
27    /// `tool.env`/`provider.env` requirements — this crate has no product
28    /// identity, so the host supplies it.
29    env_prefix: String,
30}
31
32impl AiSdkProvider {
33    pub fn new(path: PathBuf, catalog: ProviderCatalog, env_prefix: impl Into<String>) -> Self {
34        Self {
35            path,
36            catalog,
37            env_prefix: env_prefix.into(),
38        }
39    }
40
41    pub fn arc(
42        path: PathBuf,
43        catalog: ProviderCatalog,
44        env_prefix: impl Into<String>,
45    ) -> Arc<Self> {
46        Arc::new(Self::new(path, catalog, env_prefix))
47    }
48}
49
50#[async_trait]
51impl Provider for AiSdkProvider {
52    fn catalog(&self) -> ProviderCatalog {
53        self.catalog.clone()
54    }
55
56    async fn call(&self, call: ProviderCall) -> Result<ProviderOutput, ProviderError> {
57        let tool = self.tool(&call)?;
58        let runtime = SidecarRuntime::from_tool(&self.catalog, tool, &call, &self.env_prefix)?;
59        let source = self.path.display().to_string();
60        let input = execution_payload(&call).map_err(|error| {
61            ProviderError::execution(&self.catalog.provider.name, "", error)
62                .with_provider_kind("ai-sdk")
63                .with_source(source.clone())
64                .with_phase("input-serialization")
65        })?;
66
67        if input.len() > runtime.max_input_bytes {
68            return Err(ProviderError::validation(
69                &self.catalog.provider.name,
70                &call.action,
71                "ai_sdk_input_too_large",
72                format!("AI SDK input exceeds {} bytes", runtime.max_input_bytes),
73            )
74            .with_provider_kind("ai-sdk")
75            .with_source(source)
76            .with_phase("input-validation"));
77        }
78
79        let wrapper = SidecarWrapper::new(&self.path, &runtime.env).map_err(|error| {
80            ProviderError::execution(&self.catalog.provider.name, call.action.clone(), error)
81                .with_provider_kind("ai-sdk")
82                .with_source(source.clone())
83                .with_phase("runtime-load")
84        })?;
85        let started = Instant::now();
86        let sidecar = match run_bounded_sidecar(
87            &runtime.command,
88            &["--input-type=module", "--eval", wrapper.source()],
89            runtime.env,
90            &input,
91            runtime.timeout_ms,
92            runtime.max_output_bytes,
93        )
94        .await
95        {
96            Ok(sidecar) => sidecar,
97            Err(SidecarError::Timeout) => {
98                return Err(ProviderError::new(
99                    "ai_sdk_provider_timeout",
100                    &self.catalog.provider.name,
101                    Some(call.action.clone()),
102                    format!("AI SDK provider exceeded {}ms timeout", runtime.timeout_ms),
103                    "Increase tool.limits.timeout_ms or fix the provider handler.",
104                )
105                .with_provider_kind("ai-sdk")
106                .with_source(source)
107                .with_phase("execution"));
108            }
109            Err(error) => {
110                return Err(ProviderError::execution(
111                    &self.catalog.provider.name,
112                    call.action.clone(),
113                    error,
114                )
115                .with_provider_kind("ai-sdk")
116                .with_source(source)
117                .with_phase("execution"));
118            }
119        };
120        let output = sidecar.output;
121
122        tracing::debug!(
123            provider = %self.catalog.provider.name,
124            action = %call.action,
125            elapsed_ms = started.elapsed().as_millis(),
126            "AI SDK provider sidecar completed"
127        );
128
129        if sidecar.stdout_exceeded || sidecar.stderr_exceeded {
130            let stream = if sidecar.stdout_exceeded {
131                "stdout"
132            } else {
133                "stderr"
134            };
135            return Err(ProviderError::validation(
136                &self.catalog.provider.name,
137                &call.action,
138                "ai_sdk_output_too_large",
139                output_exceeded_message(stream, runtime.max_output_bytes),
140            )
141            .with_provider_kind("ai-sdk")
142            .with_source(source)
143            .with_phase("output-validation"));
144        }
145        if !output.status.success() {
146            let stderr = String::from_utf8_lossy(&output.stderr);
147            return Err(ProviderError::new(
148                "ai_sdk_provider_failed",
149                &self.catalog.provider.name,
150                Some(call.action),
151                format!("AI SDK provider failed: {}", redact_public(&stderr)),
152                "Fix the TypeScript provider handler and retry.",
153            )
154            .with_provider_kind("ai-sdk")
155            .with_source(source)
156            .with_phase("execution"));
157        }
158
159        let value = serde_json::from_slice(&output.stdout).map_err(|error| {
160            ProviderError::validation(
161                &self.catalog.provider.name,
162                &call.action,
163                "ai_sdk_invalid_json_output",
164                error.to_string(),
165            )
166            .with_provider_kind("ai-sdk")
167            .with_source(source)
168            .with_phase("output-validation")
169        })?;
170        Ok(ProviderOutput::json(value))
171    }
172}
173
174impl AiSdkProvider {
175    fn tool(&self, call: &ProviderCall) -> Result<&ProviderTool, ProviderError> {
176        self.catalog
177            .tools
178            .iter()
179            .find(|tool| tool.name == call.action)
180            .ok_or_else(|| {
181                ProviderError::validation(
182                    &self.catalog.provider.name,
183                    &call.action,
184                    "unknown_ai_sdk_action",
185                    format!("AI SDK provider has no action `{}`", call.action),
186                )
187            })
188    }
189}
190
191struct SidecarRuntime {
192    command: String,
193    env: Vec<(String, String)>,
194    timeout_ms: u64,
195    max_input_bytes: usize,
196    max_output_bytes: usize,
197}
198
199impl SidecarRuntime {
200    fn from_tool(
201        catalog: &ProviderCatalog,
202        tool: &ProviderTool,
203        call: &ProviderCall,
204        env_prefix: &str,
205    ) -> Result<Self, ProviderError> {
206        let meta = tool.meta.get("ai_sdk").or_else(|| tool.meta.get("sidecar"));
207        let command = meta
208            .and_then(|value| value.get("command"))
209            .and_then(Value::as_str)
210            .unwrap_or("node")
211            .to_owned();
212        let timeout_ms = tool
213            .limits
214            .as_ref()
215            .and_then(|limits| limits.timeout_ms)
216            .or_else(|| {
217                meta.and_then(|value| value.get("timeout_ms"))
218                    .and_then(Value::as_u64)
219            })
220            .unwrap_or(10_000);
221        let max_input_bytes = tool
222            .limits
223            .as_ref()
224            .and_then(|limits| limits.max_input_bytes)
225            .unwrap_or(64 * 1024);
226        let max_output_bytes = tool
227            .limits
228            .as_ref()
229            .and_then(|limits| limits.max_response_bytes)
230            .unwrap_or(256 * 1024);
231        Ok(Self {
232            command,
233            env: collect_provider_env(
234                &catalog.env,
235                &tool.env,
236                env_prefix,
237                &call.provider,
238                &call.action,
239            )?,
240            timeout_ms,
241            max_input_bytes,
242            max_output_bytes,
243        })
244    }
245}
246
247struct SidecarWrapper {
248    source: String,
249}
250
251impl SidecarWrapper {
252    fn new(provider_path: &std::path::Path, env: &[(String, String)]) -> std::io::Result<Self> {
253        let canonical = provider_path.canonicalize()?;
254        let module_path = canonical.display().to_string();
255        let env_keys: Vec<&str> = env.iter().map(|(key, _)| key.as_str()).collect();
256        let env_keys_json = serde_json::to_string(&env_keys).unwrap_or_else(|_| "[]".to_owned());
257        let source = format!(
258            r#"
259import {{ readFileSync }} from "node:fs";
260const chunks = [];
261for await (const chunk of process.stdin) chunks.push(chunk);
262const input = JSON.parse(Buffer.concat(chunks).toString("utf8") || "{{}}");
263const allowedEnv = new Set({env_keys_json});
264for (const key of Object.keys(process.env)) {{
265  if (!allowedEnv.has(key)) delete process.env[key];
266}}
267let providerSource = readFileSync({module_path:?}, "utf8");
268providerSource = removeDefaultManifest(providerSource);
269const module = await import("data:text/javascript;base64," + Buffer.from(providerSource).toString("base64"));
270const handler = module.call || module.default?.call;
271if (typeof handler !== "function") {{
272  throw new Error("TypeScript provider must export async function call(input)");
273}}
274const result = await handler(input);
275process.stdout.write(JSON.stringify(result ?? null));
276
277function removeDefaultManifest(source) {{
278  const marker = "export default";
279  const start = source.indexOf(marker);
280  if (start < 0) return source;
281  const open = source.indexOf("{{", start + marker.length);
282  if (open < 0) return source;
283  let depth = 0;
284  let inString = false;
285  let escaped = false;
286  for (let i = open; i < source.length; i++) {{
287    const ch = source[i];
288    if (inString) {{
289      if (escaped) escaped = false;
290      else if (ch === "\\\\") escaped = true;
291      else if (ch === "\"") inString = false;
292      continue;
293    }}
294    if (ch === "\"") inString = true;
295    else if (ch === "{{") depth++;
296    else if (ch === "}}") {{
297      depth--;
298      if (depth === 0) {{
299        let end = i + 1;
300        while (source[end] && /\\s/.test(source[end])) end++;
301        if (source[end] === ";") end++;
302        return source.slice(0, start) + source.slice(end);
303      }}
304    }}
305  }}
306  return source;
307}}
308"#
309        );
310        Ok(Self { source })
311    }
312
313    fn source(&self) -> &str {
314        &self.source
315    }
316}
317
318#[cfg(test)]
319#[path = "ai_sdk_tests.rs"]
320mod tests;