agentshield/adapter/
vercel_ai.rs

1use std::path::Path;
2
3use once_cell::sync::Lazy;
4use regex::Regex;
5
6use crate::config::ScanPathFilter;
7use crate::error::Result;
8use crate::ir::capability::{project_declared_description, project_declared_permissions};
9use crate::ir::taint_builder::build_data_surface;
10use crate::ir::*;
11
12static VERCEL_TOOL_RE: Lazy<Regex> = Lazy::new(|| {
13    Regex::new(r#"(?s)(?:const|let|var|export\s+const|export\s+let)\s+(\w+)\s*=\s*tool\s*\(\s*\{(?:\s*description\s*:\s*["'`]([^"'`]+)["'`])?"#)
14        .expect("static regex pattern is valid")
15});
16
17static VERCEL_PROPERTY_DESC_RE: Lazy<Regex> = Lazy::new(|| {
18    Regex::new(r#"(\w+)\s*:\s*z\.\w+\([^)]*\)\.describe\(\s*["'`]([^"'`]+)["'`]\s*\)"#)
19        .expect("static regex pattern is valid")
20});
21
22/// Vercel AI SDK adapter (`ai` and `@ai-sdk/*`).
23///
24/// Detects Vercel AI SDK projects by looking for:
25/// - `package.json` with `ai` or `@ai-sdk/*` dependencies
26/// - `ai.config.ts` or `ai.config.js`
27/// - TypeScript / JavaScript files importing `from "ai"` or `from "@ai-sdk/"`
28pub struct VercelAiAdapter;
29
30impl super::Adapter for VercelAiAdapter {
31    fn framework(&self) -> Framework {
32        Framework::VercelAi
33    }
34
35    fn detect(&self, root: &Path) -> bool {
36        let package_json = root.join("package.json");
37        if package_json.exists() {
38            if let Some(content) = super::read_file_capped(&package_json) {
39                if content.contains("\"ai\"")
40                    || content.contains("@ai-sdk/")
41                    || content.contains("\"@ai-sdk/openai\"")
42                    || content.contains("\"@ai-sdk/anthropic\"")
43                {
44                    return true;
45                }
46            }
47        }
48
49        if root.join("ai.config.ts").exists() || root.join("ai.config.js").exists() {
50            return true;
51        }
52
53        let walker = ignore::WalkBuilder::new(root)
54            .hidden(true)
55            .git_ignore(true)
56            .max_depth(Some(4))
57            .build();
58
59        for entry in walker.flatten() {
60            let path = entry.path();
61            if !path.is_file() {
62                continue;
63            }
64
65            let ext = path.extension().and_then(|e| e.to_str()).unwrap_or_default();
66            if !matches!(ext, "ts" | "tsx" | "js" | "jsx" | "mjs") {
67                continue;
68            }
69
70            if let Some(content) = super::read_file_capped(path) {
71                if content.contains("from 'ai'")
72                    || content.contains("from \"ai\"")
73                    || content.contains("from '@ai-sdk/")
74                    || content.contains("from \"@ai-sdk/")
75                {
76                    return true;
77                }
78            }
79        }
80
81        false
82    }
83
84    fn load(&self, root: &Path, ignore_tests: bool) -> Result<Vec<ScanTarget>> {
85        let filter = ScanPathFilter::for_ignore_tests(ignore_tests);
86        self.load_with_filter(root, &filter)
87    }
88
89    fn load_with_filter(&self, root: &Path, filter: &ScanPathFilter) -> Result<Vec<ScanTarget>> {
90        let name = root
91            .file_name()
92            .map(|n| n.to_string_lossy().to_string())
93            .unwrap_or_else(|| "vercel-ai-project".into());
94
95        let mut source_files = Vec::new();
96        super::mcp::collect_source_files_with_filter(root, filter, &mut source_files)?;
97
98        source_files.retain(|sf| {
99            matches!(
100                sf.language,
101                Language::TypeScript | Language::JavaScript
102            )
103        });
104
105        let mut tools = Vec::new();
106        for sf in &source_files {
107            tools.extend(extract_vercel_ai_tools_from_source(&sf.path, &sf.content));
108        }
109
110        for tool in &mut tools {
111            project_declared_permissions(tool);
112            project_declared_description(tool);
113        }
114
115        let execution = super::pipeline::build_execution_surface(&source_files);
116        let data = build_data_surface(&tools, &execution);
117        let dependencies = super::mcp::parse_dependencies(root, filter);
118        let provenance = super::mcp::parse_provenance(root, filter);
119
120        Ok(vec![ScanTarget {
121            name,
122            framework: Framework::VercelAi,
123            root_path: root.to_path_buf(),
124            tools,
125            execution,
126            data,
127            dependencies,
128            provenance,
129            source_files,
130        }])
131    }
132}
133
134pub(crate) fn extract_vercel_ai_tools_from_source(path: &Path, content: &str) -> Vec<ToolSurface> {
135    let mut tools = Vec::new();
136
137    for cap in VERCEL_TOOL_RE.captures_iter(content) {
138        let tool_name = cap[1].to_string();
139        let description = cap.get(2).map(|m| m.as_str().to_string());
140        let match_start = cap.get(0).map(|m| m.start()).unwrap_or(0);
141        let line = content[..match_start].lines().count() + 1;
142
143        let max_end = (match_start + 1500).min(content.len());
144        let mut end = max_end;
145        while !content.is_char_boundary(end) {
146            end -= 1;
147        }
148        let tool_snippet = &content[match_start..end];
149        let mut parameter_descriptions = Vec::new();
150
151        for prop_cap in VERCEL_PROPERTY_DESC_RE.captures_iter(tool_snippet) {
152            let prop_name = prop_cap[1].to_string();
153            let prop_desc = prop_cap[2].to_string();
154            parameter_descriptions.push(format!("{prop_name}: {prop_desc}"));
155        }
156
157        let combined_desc = if !parameter_descriptions.is_empty() {
158            let params_text = parameter_descriptions.join("\n");
159            match description {
160                Some(desc) => Some(format!("{desc}\nParameters:\n{params_text}")),
161                None => Some(format!("Parameters:\n{params_text}")),
162            }
163        } else {
164            description
165        };
166
167        tools.push(ToolSurface {
168            name: tool_name,
169            description: combined_desc,
170            input_schema: None,
171            output_schema: None,
172            declared_permissions: Vec::new(),
173            declared_capabilities: std::collections::BTreeSet::new(),
174            capability_declarations: Vec::new(),
175            observed_capabilities: std::collections::BTreeSet::new(),
176            capability_evidence: Vec::new(),
177            capability_observation_complete: false,
178            defined_at: Some(SourceLocation {
179                file: path.to_path_buf(),
180                line,
181                column: 0,
182                end_line: Some(line),
183                end_column: None,
184            }),
185        });
186    }
187
188    tools
189}
190
191#[cfg(test)]
192mod tests {
193    use super::*;
194    use crate::adapter::Adapter;
195
196    #[test]
197    fn detects_vercel_ai_package_json() {
198        let dir = tempfile::tempdir().unwrap();
199        let pkg = dir.path().join("package.json");
200        std::fs::write(&pkg, r#"{"dependencies": {"ai": "^3.1.0"}}"#).unwrap();
201
202        let adapter = VercelAiAdapter;
203        assert!(adapter.detect(dir.path()));
204    }
205
206    #[test]
207    fn detects_vercel_ai_import() {
208        let dir = tempfile::tempdir().unwrap();
209        let src = dir.path().join("tool.ts");
210        std::fs::write(&src, "import { tool } from 'ai';\n").unwrap();
211
212        let adapter = VercelAiAdapter;
213        assert!(adapter.detect(dir.path()));
214    }
215
216    #[test]
217    fn extracts_vercel_ai_tool_declarations() {
218        let content = r#"
219import { tool } from 'ai';
220import { z } from 'zod';
221
222export const getWeather = tool({
223  description: 'Get the current weather for a city',
224  parameters: z.object({
225    city: z.string().describe('The target city name'),
226  }),
227  execute: async ({ city }) => {
228    return fetch(`https://api.weather.com/${city}`);
229  },
230});
231"#;
232        let tools = extract_vercel_ai_tools_from_source(Path::new("tools.ts"), content);
233        assert_eq!(tools.len(), 1);
234        assert_eq!(tools[0].name, "getWeather");
235        assert!(tools[0].description.as_deref().unwrap().contains("weather for a city"));
236        assert!(tools[0].description.as_deref().unwrap().contains("The target city name"));
237        assert_eq!(tools[0].defined_at.as_ref().unwrap().line, 5);
238    }
239
240    #[test]
241    fn extracts_vercel_ai_tool_with_unicode_characters() {
242        // Multi-byte UTF-8 string (emojis & accented text) spanning across the 1500-byte boundary
243        let unicode_padding = "🚀 Informação confidencial e análise de segurança ".repeat(40);
244        let content = format!(
245            r#"
246import {{ tool }} from 'ai';
247import {{ z }} from 'zod';
248
249export const analyzeData = tool({{
250  description: 'Análise de dados avançada com IA',
251  parameters: z.object({{
252    query: z.string().describe('Consulta SQL para execução'),
253  }}),
254  // {unicode_padding}
255  execute: async ({{ query }}) => {{
256    return db.query(query);
257  }},
258}});
259"#
260        );
261
262        let tools = extract_vercel_ai_tools_from_source(Path::new("tools.ts"), &content);
263        assert_eq!(tools.len(), 1);
264        assert_eq!(tools[0].name, "analyzeData");
265        assert!(tools[0].description.as_deref().unwrap().contains("Consulta SQL"));
266    }
267}
agentshield/adapter/vercel_ai.rs

agentshield/adapter/
vercel_ai.rs
