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

agentshield/adapter/
vercel_ai.rs
