@complior/engine 0.9.0

package/src/domain/scanner/fix-diff-builder.ts

@@ -0,0 +1,477 @@
+import type { CodeContext, CodeContextLine, FixDiff } from '../../types/common.types.js';
+
+const CONTEXT_LINES = 5;
+
+/** Extract surrounding source lines for a finding with file+line. */
+export const buildCodeContext = (
+  fileContent: string,
+  line: number,
+): CodeContext => {
+  const allLines = fileContent.split('\n');
+  const startLine = Math.max(1, line - CONTEXT_LINES);
+  const endLine = Math.min(allLines.length, line + CONTEXT_LINES);
+  const lines: CodeContextLine[] = [];
+  for (let i = startLine; i <= endLine; i++) {
+    lines.push({ num: i, content: allLines[i - 1] ?? '' });
+  }
+  return { lines, startLine, highlightLine: line };
+};
+
+/**
+ * Find the end of a statement starting at `startIdx` (0-based).
+ * Tracks paren/brace depth for multi-line calls like:
+ * ```
+ try {
+   * const r = await client.messages.create({
+   *   model: "claude-3",
+   *   messages: [...],
+   * });
+ } catch (err) {
+   console.error('LLM call failed:', err);
+   throw err;
+ }
+ * ```
+ */
+const findStatementEnd = (allLines: readonly string[], startIdx: number): number => {
+  let depth = 0;
+  for (let i = startIdx; i < Math.min(allLines.length, startIdx + 30); i++) {
+    const line = allLines[i];
+    if (line === undefined) continue;
+    for (const ch of line) {
+      if (ch === '(' || ch === '{' || ch === '[') depth++;
+      if (ch === ')' || ch === '}' || ch === ']') depth--;
+    }
+    if (depth <= 0) return i;
+  }
+  return startIdx;
+};
+
+/** Generate a production-quality before/after fix diff for any fixable finding. */
+export const buildFixDiff = (
+  fileContent: string,
+  line: number,
+  filePath: string,
+  checkId: string,
+): FixDiff | undefined => {
+  const allLines = fileContent.split('\n');
+  const targetLine = allLines[line - 1];
+  if (targetLine === undefined && !checkId.startsWith('l3-banned-')) return undefined;
+
+  // Bare LLM calls are informational — no fix diff needed
+  if (checkId.includes('bare')) return undefined;
+  if (checkId.startsWith('l4-nhi-') && checkId !== 'l4-nhi-clean')
+    return buildNhiDiff(allLines, line, filePath, targetLine!);
+  if (checkId === 'l4-security-risk') {
+    const secDiff = buildSecurityRiskDiff(allLines, line, filePath, targetLine!);
+    if (secDiff) return secDiff;
+    // Fallback: security-risk "Hardcoded API key" → treat like NHI
+    return buildNhiDiff(allLines, line, filePath, targetLine!);
+  }
+  if (checkId === 'l4-ast-missing-error-handling')
+    return buildErrorHandlingDiff(allLines, line, filePath, fileContent);
+  if (checkId.startsWith('l3-banned-'))
+    return buildBannedDepDiff(allLines, filePath, checkId);
+
+  // External tool findings — route to same fix builders
+  if (checkId.startsWith('ext-semgrep-complior-bare-call'))
+    return undefined;
+  if (checkId.startsWith('ext-semgrep-complior-unsafe-deser'))
+    return buildSecurityRiskDiff(allLines, line, filePath, targetLine!) ?? buildNhiDiff(allLines, line, filePath, targetLine!);
+  if (checkId.startsWith('ext-semgrep-complior-injection'))
+    return buildSecurityRiskDiff(allLines, line, filePath, targetLine!);
+  if (checkId.startsWith('ext-semgrep-complior-missing-error-handling'))
+    return buildErrorHandlingDiff(allLines, line, filePath, fileContent);
+  if (checkId.startsWith('ext-detect-secrets-'))
+    return buildNhiDiff(allLines, line, filePath, targetLine!);
+  if (checkId.startsWith('ext-bandit-')) {
+    // B301 (pickle), B603 (shell=True), etc. → security pattern replacement
+    const secDiff = buildSecurityRiskDiff(allLines, line, filePath, targetLine!);
+    if (secDiff) return secDiff;
+    // B105 (hardcoded password) and others → NHI env var replacement
+    return buildNhiDiff(allLines, line, filePath, targetLine!);
+  }
+
+  return undefined;
+};
+
+// ---------------------------------------------------------------------------
+// NHI diff — replace hardcoded secrets with env vars
+// ---------------------------------------------------------------------------
+
+const PY_ASSIGN = /^(\s*)(\w+)\s*=\s*(['"`])(.+?)\3/;
+const TS_ASSIGN = /^(\s*(?:export\s+)?(?:const|let|var)\s+)(\w+)(?:\s*:\s*\w+)?\s*=\s*(['"`])(.+?)\3/;
+const OBJ_PROP = /^(\s*)(\w+)\s*:\s*(['"`])(.+?)\3/;
+const CONN_STRING = /(mongodb|postgres(?:ql)?|mysql|redis):\/\/[^'"`\s]+/;
+
+/** Derive env var name from a key name: aws_secret_key → AWS_SECRET_KEY */
+const toEnvKey = (name: string): string => name.replace(/[^a-zA-Z0-9]/g, '_').toUpperCase();
+
+/** Handle multi-line private key assignments (e.g. PEM blocks in triple-quotes). */
+const buildPrivateKeyDiff = (
+  allLines: readonly string[],
+  line: number,
+  filePath: string,
+  targetLine: string,
+  isPython: boolean,
+): FixDiff | undefined => {
+  const startIdx = line - 1;
+
+  // Detect variable name from the assignment
+  const assignMatch = isPython
+    ? targetLine.match(/^(\s*)(\w+)\s*=/)
+    : targetLine.match(/^(\s*(?:export\s+)?(?:const|let|var)\s+)(\w+)/);
+  if (!assignMatch) return undefined;
+
+  const [, prefix, varName] = assignMatch;
+  const envKey = (varName ?? 'PRIVATE_KEY').toUpperCase();
+
+  // Find the closing triple-quote or backtick
+  let endIdx = startIdx;
+  const tripleQuote = targetLine.includes('"""') ? '"""' : targetLine.includes("'''") ? "'''" : '`';
+  // Count opening quotes on start line — if there's only one occurrence, it's unclosed
+  const quoteCount = (targetLine.split(tripleQuote).length - 1);
+  if (quoteCount >= 2) {
+    // Single-line: both open and close on same line
+    endIdx = startIdx;
+  } else {
+    // Multi-line: find closing quote
+    for (let i = startIdx + 1; i < Math.min(allLines.length, startIdx + 50); i++) {
+      if ((allLines[i] ?? '').includes(tripleQuote)) {
+        endIdx = i;
+        break;
+      }
+    }
+  }
+
+  const before: string[] = [];
+  for (let i = startIdx; i <= endIdx; i++) {
+    before.push(allLines[i] ?? '');
+  }
+
+  const after: string[] = isPython
+    ? [`${prefix}${varName} = os.environ.get('${envKey}', '')`]
+    : [`${prefix}${varName} = process.env.${envKey} ?? '';`];
+
+  return {
+    before,
+    after,
+    startLine: line,
+    filePath,
+    importLine: isPython ? 'import os' : undefined,
+  };
+};
+
+const buildNhiDiff = (
+  allLines: readonly string[],
+  line: number,
+  filePath: string,
+  targetLine: string,
+): FixDiff | undefined => {
+  const isPython = filePath.endsWith('.py');
+
+  // Multi-line private key: replace entire block with env var
+  if (/-----BEGIN\s+(\w+\s+)?PRIVATE\s+KEY-----/.test(targetLine)) {
+    return buildPrivateKeyDiff(allLines, line, filePath, targetLine, isPython);
+  }
+
+  // Skip lines already using env vars
+  if (/process\.env\.|os\.environ|os\.getenv/.test(targetLine)) return undefined;
+
+  // Connection string
+  if (CONN_STRING.test(targetLine)) {
+    const envVar = 'DATABASE_URL';
+    if (isPython) {
+      const pyMatch = targetLine.match(PY_ASSIGN);
+      if (pyMatch) {
+        const [, indent, varName] = pyMatch;
+        return {
+          before: [targetLine],
+          after: [`${indent}${varName} = os.environ.get('${envVar}', '')`],
+          startLine: line,
+          filePath,
+          importLine: 'import os',
+        };
+      }
+    }
+    const tsMatch = targetLine.match(TS_ASSIGN);
+    if (tsMatch) {
+      const [, prefix, varName] = tsMatch;
+      return {
+        before: [targetLine],
+        after: [`${prefix}${varName} = process.env.${envVar} ?? ''`],
+        startLine: line,
+        filePath,
+      };
+    }
+    // Object property with connection string
+    const objMatch = targetLine.match(OBJ_PROP);
+    if (objMatch) {
+      const [, indent, propName] = objMatch;
+      const comma = targetLine.trimEnd().endsWith(',') ? ',' : '';
+      return {
+        before: [targetLine],
+        after: [`${indent}${propName}: process.env.${toEnvKey(propName)} ?? ''${comma}`],
+        startLine: line,
+        filePath,
+      };
+    }
+  }
+
+  // Python: multi-line string that's not a private key (e.g. triple-quoted secret)
+  if (isPython && (targetLine.includes('"""') || targetLine.includes("'''"))) {
+    const tripleQ = targetLine.includes('"""') ? '"""' : "'''";
+    const qCount = targetLine.split(tripleQ).length - 1;
+    if (qCount === 1) {
+      // Opening quote only → multi-line string, find closing
+      return buildPrivateKeyDiff(allLines, line, filePath, targetLine, true);
+    }
+  }
+
+  // Python assignment
+  if (isPython) {
+    const match = targetLine.match(PY_ASSIGN);
+    if (match) {
+      const [, indent, varName] = match;
+      const envKey = varName.toUpperCase();
+      return {
+        before: [targetLine],
+        after: [`${indent}${varName} = os.environ.get('${envKey}', '')`],
+        startLine: line,
+        filePath,
+        importLine: 'import os',
+      };
+    }
+    return undefined;
+  }
+
+  // TypeScript/JavaScript variable assignment: const X = '...'
+  const tsMatch = targetLine.match(TS_ASSIGN);
+  if (tsMatch) {
+    const [, prefix, varName] = tsMatch;
+    const envKey = varName.toUpperCase();
+    return {
+      before: [targetLine],
+      after: [`${prefix}${varName} = process.env.${envKey} ?? ''`],
+      startLine: line,
+      filePath,
+    };
+  }
+
+  // Object property assignment: key: 'secret-value'
+  const objMatch = targetLine.match(OBJ_PROP);
+  if (objMatch) {
+    const [, indent, propName, , value] = objMatch;
+    // Only fix if value looks like a real secret (long, not a URL scheme, not a placeholder)
+    if (value.length >= 8 && !/^https?:\/\//.test(value)) {
+      const comma = targetLine.trimEnd().endsWith(',') ? ',' : '';
+      return {
+        before: [targetLine],
+        after: [`${indent}${propName}: process.env.${toEnvKey(propName)} ?? ''${comma}`],
+        startLine: line,
+        filePath,
+      };
+    }
+  }
+
+  return undefined;
+};
+
+// ---------------------------------------------------------------------------
+// 1c. Security risk diff — pattern-specific replacements
+// ---------------------------------------------------------------------------
+
+interface SecurityPattern {
+  readonly test: RegExp;
+  readonly replace: (line: string) => string;
+  readonly importLine?: string;
+}
+
+const SECURITY_PATTERNS: readonly SecurityPattern[] = [
+  // eval() → disabled
+  { test: /\beval\s*\(/, replace: (l) => l.replace(/\beval\s*\([^)]*\)/, '/* COMPLIOR: eval() disabled — Art. 15 */ undefined') },
+  // new Function(...) → disabled
+  { test: /\bnew\s+Function\s*\(/, replace: (l) => l.replace(/\bnew\s+Function\s*\([^)]*\)/, '/* COMPLIOR: new Function() disabled */ undefined') },
+  // vm.runInNewContext / vm.runInThisContext → disabled
+  { test: /\bvm\.(runInNewContext|runInThisContext)\s*\(/, replace: (l) => l.replace(/\bvm\.(runInNewContext|runInThisContext)\s*\([^)]*\)/, '/* COMPLIOR: vm execution disabled */ undefined') },
+  // pickle.load(f) → json.load(f)
+  { test: /\bpickle\.load\s*\(/, replace: (l) => l.replace(/\bpickle\.load\b/, 'json.load'), importLine: 'import json' },
+  // pickle.loads(d) → json.loads(d)
+  { test: /\bpickle\.loads\s*\(/, replace: (l) => l.replace(/\bpickle\.loads\b/, 'json.loads'), importLine: 'import json' },
+  // hashlib.md5() → hashlib.sha256()
+  { test: /\bhashlib\.md5\s*\(/, replace: (l) => l.replace(/\bhashlib\.md5\b/, 'hashlib.sha256') },
+  // hashlib.sha1() → hashlib.sha256()
+  { test: /\bhashlib\.sha1\s*\(/, replace: (l) => l.replace(/\bhashlib\.sha1\b/, 'hashlib.sha256') },
+  // verify=False → verify=True
+  { test: /\bverify\s*=\s*False\b/, replace: (l) => l.replace(/\bverify\s*=\s*False\b/, 'verify=True') },
+  // rejectUnauthorized: false → true
+  { test: /rejectUnauthorized\s*:\s*false/, replace: (l) => l.replace(/rejectUnauthorized\s*:\s*false/, 'rejectUnauthorized: true') },
+  // shell=True → shell=False
+  { test: /\bshell\s*=\s*True\b/, replace: (l) => l.replace(/\bshell\s*=\s*True\b/, 'shell=False') },
+  // os.system(cmd) → subprocess.run(cmd.split(), check=True)
+  { test: /\bos\.system\s*\(/, replace: (l) => l.replace(/\bos\.system\s*\(([^)]+)\)/, 'subprocess.run($1.split(), check=True)'), importLine: 'import subprocess' },
+  // torch.load(x) → torch.load(x, weights_only=True)
+  { test: /\btorch\.load\s*\(/, replace: (l) => {
+    // Skip if already has weights_only
+    if (l.includes('weights_only')) return l;
+    return l.replace(/\btorch\.load\s*\(([^)]+)\)/, 'torch.load($1, weights_only=True)');
+  }},
+];
+
+const buildSecurityRiskDiff = (
+  _allLines: readonly string[],
+  line: number,
+  filePath: string,
+  targetLine: string,
+): FixDiff | undefined => {
+  for (const pattern of SECURITY_PATTERNS) {
+    if (pattern.test.test(targetLine)) {
+      const replaced = pattern.replace(targetLine);
+      if (replaced === targetLine) continue;
+      return {
+        before: [targetLine],
+        after: [replaced],
+        startLine: line,
+        filePath,
+        importLine: pattern.importLine,
+      };
+    }
+  }
+  return undefined;
+};
+
+// ---------------------------------------------------------------------------
+// 1d. Error handling diff — wrap LLM calls in try/catch
+// ---------------------------------------------------------------------------
+
+const LLM_CALL = /\.(messages\.create|chat\.completions\.create|chat\.complete|chat|generateContent|invoke|images\.generate|embeddings\.create|send)\s*\(|\b(generateText|streamText|generateObject)\s*\(\s*\{/;
+
+const buildErrorHandlingDiff = (
+  allLines: readonly string[],
+  line: number,
+  filePath: string,
+  fileContent: string,
+): FixDiff | undefined => {
+  const startIdx = line - 1;
+  const targetLine = allLines[startIdx];
+  if (targetLine === undefined) return undefined;
+
+  // Only wrap if line contains an LLM call pattern
+  if (!LLM_CALL.test(fileContent.split('\n')[startIdx] ?? '')) return undefined;
+
+  const endIdx = findStatementEnd(allLines, startIdx);
+  const isPython = filePath.endsWith('.py');
+
+  const before: string[] = [];
+  for (let i = startIdx; i <= endIdx; i++) {
+    before.push(allLines[i] ?? '');
+  }
+
+  // Detect indentation from the first line
+  const indentMatch = (allLines[startIdx] ?? '').match(/^(\s*)/);
+  const indent = indentMatch ? indentMatch[1] : '';
+  const innerIndent = indent + (isPython ? '    ' : '  ');
+
+  const after: string[] = [];
+  if (isPython) {
+    after.push(`${indent}try:`);
+    for (const bLine of before) {
+      after.push(`${innerIndent}${bLine.trimStart()}`);
+    }
+    after.push(`${indent}except Exception as e:`);
+    after.push(`${innerIndent}print(f"LLM call failed: {e}")`);
+    after.push(`${innerIndent}raise`);
+  } else {
+    after.push(`${indent}try {`);
+    for (const bLine of before) {
+      after.push(`${innerIndent}${bLine.trimStart()}`);
+    }
+    after.push(`${indent}} catch (err) {`);
+    after.push(`${innerIndent}console.error('LLM call failed:', err);`);
+    after.push(`${innerIndent}throw err;`);
+    after.push(`${indent}}`);
+  }
+
+  return {
+    before,
+    after,
+    startLine: line,
+    filePath,
+  };
+};
+
+// ---------------------------------------------------------------------------
+// 1e. Banned dep diff — remove prohibited dependencies from manifest files
+// ---------------------------------------------------------------------------
+
+const buildBannedDepDiff = (
+  allLines: readonly string[],
+  filePath: string,
+  checkId: string,
+): FixDiff | undefined => {
+  // Extract package name from checkId: "l3-banned-emotion-recognition" → "emotion-recognition"
+  const packageName = checkId.replace('l3-banned-', '');
+  if (!packageName) return undefined;
+
+  const isPip = filePath.endsWith('requirements.txt');
+
+  // Find the line containing the package
+  let lineIdx = -1;
+  for (let i = 0; i < allLines.length; i++) {
+    const l = allLines[i] ?? '';
+    if (isPip) {
+      // requirements.txt: line starts with package name
+      if (l.trim().startsWith(packageName)) { lineIdx = i; break; }
+    } else {
+      // package.json: line contains "packageName"
+      if (l.includes(`"${packageName}"`)) { lineIdx = i; break; }
+    }
+  }
+
+  if (lineIdx === -1) return undefined;
+
+  const targetLine = allLines[lineIdx]!;
+
+  if (isPip) {
+    // Simply remove the line
+    return {
+      before: [targetLine],
+      after: [],
+      startLine: lineIdx + 1,
+      filePath,
+    };
+  }
+
+  // package.json: remove line, handle trailing comma
+  const hasTrailingComma = targetLine.trimEnd().endsWith(',');
+  if (hasTrailingComma) {
+    // Just remove the line
+    return {
+      before: [targetLine],
+      after: [],
+      startLine: lineIdx + 1,
+      filePath,
+    };
+  }
+
+  // No trailing comma — this is the last entry. Remove trailing comma from previous line too
+  if (lineIdx > 0) {
+    const prevLine = allLines[lineIdx - 1]!;
+    if (prevLine.trimEnd().endsWith(',')) {
+      const cleanPrev = prevLine.replace(/,(\s*)$/, '$1');
+      return {
+        before: [prevLine, targetLine],
+        after: [cleanPrev],
+        startLine: lineIdx, // 1-based (prev line)
+        filePath,
+      };
+    }
+  }
+
+  // Fallback: just remove the line
+  return {
+    before: [targetLine],
+    after: [],
+    startLine: lineIdx + 1,
+    filePath,
+  };
+};

package/src/domain/scanner/import-graph.test.ts

@@ -0,0 +1,162 @@
+import { describe, it, expect } from 'vitest';
+import { extractImports, isAiPackageImport, buildImportGraph } from './import-graph.js';
+import type { FileInfo } from '../../ports/scanner.port.js';
+
+const makeFile = (path: string, content: string): FileInfo => ({
+  path: `/project/${path}`,
+  relativePath: path,
+  content,
+  extension: '.' + (path.split('.').pop() ?? 'ts'),
+});
+
+describe('extractImports', () => {
+  it('extracts ES import statements', () => {
+    const content = `import { OpenAI } from 'openai';\nimport Anthropic from '@anthropic-ai/sdk';`;
+    const imports = extractImports(content, '.ts');
+    expect(imports).toContain('openai');
+    expect(imports).toContain('@anthropic-ai/sdk');
+  });
+
+  it('extracts require statements', () => {
+    const content = `const openai = require('openai');\nconst fs = require('fs');`;
+    const imports = extractImports(content, '.js');
+    expect(imports).toContain('openai');
+    expect(imports).toContain('fs');
+  });
+
+  it('extracts relative imports', () => {
+    const content = `import { client } from './llm-client.js';\nimport { config } from '../config.js';`;
+    const imports = extractImports(content, '.ts');
+    expect(imports).toContain('./llm-client.js');
+    expect(imports).toContain('../config.js');
+  });
+
+  it('extracts Python imports', () => {
+    const content = `import openai\nfrom anthropic import Anthropic\nimport os`;
+    const imports = extractImports(content, '.py');
+    expect(imports).toContain('openai');
+    expect(imports).toContain('anthropic');
+    expect(imports).toContain('os');
+  });
+
+  it('handles side-effect imports', () => {
+    const content = `import 'reflect-metadata';`;
+    const imports = extractImports(content, '.ts');
+    expect(imports).toContain('reflect-metadata');
+  });
+});
+
+describe('isAiPackageImport', () => {
+  it('detects direct AI SDK packages', () => {
+    expect(isAiPackageImport('openai')).toBe(true);
+    expect(isAiPackageImport('@anthropic-ai/sdk')).toBe(true);
+    expect(isAiPackageImport('@ai-sdk/openai')).toBe(true);
+    expect(isAiPackageImport('langchain')).toBe(true);
+  });
+
+  it('detects scoped subpaths', () => {
+    expect(isAiPackageImport('@ai-sdk/openai/chat')).toBe(true);
+    expect(isAiPackageImport('@langchain/core/messages')).toBe(true);
+  });
+
+  it('detects Python AI packages', () => {
+    expect(isAiPackageImport('openai')).toBe(true);
+    expect(isAiPackageImport('anthropic')).toBe(true);
+    expect(isAiPackageImport('transformers')).toBe(true);
+    expect(isAiPackageImport('torch')).toBe(true);
+  });
+
+  it('rejects non-AI packages', () => {
+    expect(isAiPackageImport('fs')).toBe(false);
+    expect(isAiPackageImport('express')).toBe(false);
+    expect(isAiPackageImport('react')).toBe(false);
+    expect(isAiPackageImport('lodash')).toBe(false);
+  });
+});
+
+describe('buildImportGraph', () => {
+  it('detects direct AI SDK imports', () => {
+    const files = [
+      makeFile('src/llm.ts', `import { OpenAI } from 'openai';\nconst client = new OpenAI();`),
+      makeFile('src/utils.ts', `export const helper = () => {};`),
+    ];
+
+    const graph = buildImportGraph(files);
+    expect(graph.directAiFiles.has('src/llm.ts')).toBe(true);
+    expect(graph.directAiFiles.has('src/utils.ts')).toBe(false);
+    expect(graph.aiRelevantFiles.has('src/llm.ts')).toBe(true);
+  });
+
+  it('propagates AI relevance transitively', () => {
+    const files = [
+      makeFile('src/openai-client.ts', `import { OpenAI } from 'openai';\nexport const client = new OpenAI();`),
+      makeFile('src/chat-service.ts', `import { client } from './openai-client.js';\nexport const chat = () => client.chat.completions.create({});`),
+      makeFile('src/api.ts', `import { chat } from './chat-service.js';\nexport const handler = () => chat();`),
+      makeFile('src/unrelated.ts', `import { readFile } from 'fs';\nexport const x = 1;`),
+    ];
+
+    const graph = buildImportGraph(files);
+
+    // Direct
+    expect(graph.directAiFiles.has('src/openai-client.ts')).toBe(true);
+
+    // Transitive through imports
+    expect(graph.aiRelevantFiles.has('src/chat-service.ts')).toBe(true);
+    expect(graph.aiRelevantFiles.has('src/api.ts')).toBe(true);
+
+    // Unrelated file should not be AI-relevant
+    expect(graph.aiRelevantFiles.has('src/unrelated.ts')).toBe(false);
+  });
+
+  it('handles circular imports without infinite loop', () => {
+    const files = [
+      makeFile('src/a.ts', `import { b } from './b.js';\nimport { OpenAI } from 'openai';`),
+      makeFile('src/b.ts', `import { a } from './a.js';\nexport const b = 1;`),
+    ];
+
+    const graph = buildImportGraph(files);
+    expect(graph.aiRelevantFiles.has('src/a.ts')).toBe(true);
+    expect(graph.aiRelevantFiles.has('src/b.ts')).toBe(true);
+  });
+
+  it('tracks AI relevance depth', () => {
+    const files = [
+      makeFile('src/llm.ts', `import { OpenAI } from 'openai';`),
+      makeFile('src/service.ts', `import { x } from './llm.js';`),
+      makeFile('src/handler.ts', `import { y } from './service.js';`),
+    ];
+
+    const graph = buildImportGraph(files);
+    expect(graph.nodes.get('src/llm.ts')?.aiRelevanceDepth).toBe(0);
+    expect(graph.nodes.get('src/service.ts')?.aiRelevanceDepth).toBe(1);
+    expect(graph.nodes.get('src/handler.ts')?.aiRelevanceDepth).toBe(2);
+  });
+
+  it('returns non-AI-relevant files with depth -1', () => {
+    const files = [
+      makeFile('src/utils.ts', `export const add = (a: number, b: number) => a + b;`),
+    ];
+
+    const graph = buildImportGraph(files);
+    const node = graph.nodes.get('src/utils.ts');
+    expect(node?.isAiRelevant).toBe(false);
+    expect(node?.aiRelevanceDepth).toBe(-1);
+  });
+
+  it('handles Python imports', () => {
+    const files = [
+      makeFile('src/llm.py', `import openai\nclient = openai.OpenAI()`),
+      makeFile('src/utils.py', `def helper():\n    pass`),
+    ];
+
+    const graph = buildImportGraph(files);
+    expect(graph.directAiFiles.has('src/llm.py')).toBe(true);
+    expect(graph.directAiFiles.has('src/utils.py')).toBe(false);
+  });
+
+  it('handles empty file list', () => {
+    const graph = buildImportGraph([]);
+    expect(graph.nodes.size).toBe(0);
+    expect(graph.aiRelevantFiles.size).toBe(0);
+  });
+});