1use 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 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;