Merge pull request #22 from SquidSec/fix/code-review-hardening

fix: harden diff path, LLM retries, and category filtering

CHANGELOG.md

@@ -4,6 +4,18 @@ All notable changes to **SquidGate** ([SquidSec](https://squidoffense.com/)) are
 
 Release tags: **`v{semver}-build.{N}`**, with floating **`v1`** / **`v1.0`**.
 
+## Unreleased
+
+### Fixed / Improved
+
+- **Diff acquisition:** prefer GitHub API (fork/shallow-safe); git fallback uses `merge-base` and multiple base candidates; honor `lines_before` / `lines_after` via unified context
+- **LLM resilience:** retries with exponential backoff on 429/5xx; Anthropic/Google re-prompt on JSON parse failure
+- **JSON extraction:** balanced-brace parser (nested objects / braces in strings); surface `parse_error` on the check summary
+- **Category filtering:** post-filter findings against `policy.categories` (plus CWE/text inference when `category` omitted)
+- **Prompts:** few-shot secret/injection examples; disabled-category instructions; `category` field in schema
+- **Checks:** single `shouldBlock` implementation; truncation warnings in check output
+- **`block_on: none`:** correctly never blocks merge
+
 ## [1.0.0] — 2026-07-28
 
 ### Added

__tests__/action-flow.test.ts

@@ -1,5 +1,6 @@
 import { loadConfig, filterFindings, shouldBlock } from '../src/config';
-import { callLlm } from '../src/llm'; // will be mocked
+import { callLlm } from '../src/llm';
+import type { Finding } from '../src/types';
 
 jest.mock('../src/llm');
 
@@ -9,12 +10,12 @@ describe('action claims via pieces', () => {
     const findings = [
       { confidence: 'high', severity: 'high' },
       { confidence: 'medium', severity: 'medium' },
-    ];
+    ] as Finding[];
 
-    const filtered = filterFindings(findings as any, cfg);
+    const filtered = filterFindings(findings, cfg);
     expect(filtered.length).toBeGreaterThanOrEqual(1);
 
-    const blocks = filtered.some((f: any) => shouldBlock(f.severity, cfg.policy.block_on));
+    const blocks = filtered.some((f) => shouldBlock(f.severity, cfg.policy.block_on));
     expect(blocks).toBe(true); // with default high, high finding blocks
   });
 
@@ -22,8 +23,28 @@ describe('action claims via pieces', () => {
     const mockCall = callLlm as jest.Mock;
     mockCall.mockResolvedValue({ findings: [], summary: '' });
 
-    // In real execution the callLlm receives system+user built from policy
-    // We just assert it is the function we expect to be used for all providers
     expect(typeof callLlm).toBe('function');
   });
+
+  it('category-disabled findings do not block even at high severity', () => {
+    const cfg = loadConfig('/nope');
+    cfg.policy.categories.injection = false;
+    const findings = [
+      {
+        file: 'a.py',
+        start_line: 1,
+        end_line: 1,
+        severity: 'critical',
+        title: 'SQLi',
+        description: 'SQL injection',
+        cwe: 'CWE-89',
+        owasp: null,
+        recommendation: 'fix',
+        confidence: 'high',
+        category: 'injection',
+      },
+    ] as Finding[];
+    const filtered = filterFindings(findings, cfg);
+    expect(filtered.length).toBe(0);
+  });
 });

__tests__/config.test.ts

@@ -1,6 +1,7 @@
-import { loadConfig, DEFAULT_CONFIG, shouldBlock, filterFindings } from '../src/config';
+import { loadConfig, DEFAULT_CONFIG, shouldBlock, filterFindings, inferCategory, deepMerge } from '../src/config';
 import * as fs from 'fs';
 import * as path from 'path';
+import type { Finding } from '../src/types';
 
 const tmpConfigPath = path.join(__dirname, 'temp-squidgate.yml');
 
@@ -37,11 +38,11 @@ policy:
   });
 
   it('applies action input overrides', () => {
-    const overrides: any = {
+    const overrides = {
       llm: { provider: 'anthropic', model: 'claude-3-5-sonnet-20241022' },
       policy: { block_on: 'critical' as const },
     };
-    const cfg = loadConfig('/no/file', overrides);
+    const cfg = loadConfig('/no/file', overrides as any);
     expect(cfg.llm.provider).toBe('anthropic');
     expect(cfg.llm.model).toBe('claude-3-5-sonnet-20241022');
     expect(cfg.policy.block_on).toBe('critical');
@@ -51,7 +52,7 @@ policy:
     expect(shouldBlock('critical', 'high')).toBe(true);
     expect(shouldBlock('high', 'high')).toBe(true);
     expect(shouldBlock('medium', 'high')).toBe(false);
-    expect(shouldBlock('high', 'none')).toBe(true);
+    expect(shouldBlock('high', 'none')).toBe(false); // none = never block
     expect(shouldBlock('info', 'low')).toBe(false);
   });
 
@@ -61,9 +62,76 @@ policy:
       { confidence: 'high' },
       { confidence: 'medium' },
       { confidence: 'low' },
-    ];
+    ] as Finding[];
     const filtered = filterFindings(findings, cfg);
     expect(filtered.length).toBe(1);
     expect(filtered[0].confidence).toBe('high');
   });
+
+  it('filterFindings drops disabled categories', () => {
+    const cfg = {
+      ...DEFAULT_CONFIG,
+      policy: {
+        ...DEFAULT_CONFIG.policy,
+        categories: { ...DEFAULT_CONFIG.policy.categories, xss: false, secrets: true },
+      },
+    };
+    const findings = [
+      {
+        file: 'a.ts',
+        start_line: 1,
+        end_line: 1,
+        severity: 'high',
+        title: 'Reflected XSS',
+        description: 'cross-site scripting via innerHTML',
+        cwe: 'CWE-79',
+        owasp: null,
+        recommendation: 'escape',
+        confidence: 'high',
+        category: 'xss',
+      },
+      {
+        file: 'b.ts',
+        start_line: 2,
+        end_line: 2,
+        severity: 'high',
+        title: 'API key',
+        description: 'hardcoded secret token',
+        cwe: 'CWE-798',
+        owasp: null,
+        recommendation: 'env',
+        confidence: 'high',
+        category: 'secrets',
+      },
+    ] as Finding[];
+    const filtered = filterFindings(findings, cfg);
+    expect(filtered.length).toBe(1);
+    expect(filtered[0].category).toBe('secrets');
+  });
+
+  it('inferCategory maps CWE/text when category omitted', () => {
+    const f = {
+      file: 'x.py',
+      start_line: 1,
+      end_line: 1,
+      severity: 'high',
+      title: 'Issue',
+      description: 'user input concatenated into SQL query',
+      cwe: 'CWE-89',
+      owasp: null,
+      recommendation: 'parameterize',
+      confidence: 'high',
+    } as Finding;
+    expect(inferCategory(f)).toBe('injection');
+  });
+
+  it('deepMerge merges nested objects recursively', () => {
+    const a = { policy: { categories: { a: true, b: true }, block_on: 'high' }, x: 1 };
+    const b = { policy: { categories: { b: false }, block_on: 'low' } };
+    const m = deepMerge(a, b);
+    expect(m.policy.block_on).toBe('low');
+    expect(m.policy.categories.a).toBe(true);
+    expect(m.policy.categories.b).toBe(false);
+    expect(m.x).toBe(1);
+  });
 });

__tests__/diff.test.ts

@@ -0,0 +1,12 @@
+import { extractBalancedObject } from '../src/llm';
+
+// Diff helpers that don't need git/network — truncation is internal.
+// We re-export behavior via a thin pure test of size messaging patterns used by diff.ts
+
+describe('diff helpers (pure)', () => {
+  it('balanced JSON extract used when model wraps objects', () => {
+    const noisy = 'Sure!\n{"findings":[],"summary":"ok"}\n';
+    const obj = extractBalancedObject(noisy);
+    expect(JSON.parse(obj!).summary).toBe('ok');
+  });
+});

__tests__/llm.test.ts

@@ -1,4 +1,4 @@
-import { extractJson, normalizeResponse } from '../src/llm';
+import { extractJson, extractBalancedObject, normalizeResponse } from '../src/llm';
 
 describe('llm response parsing', () => {
   it('normalizes good response', () => {
@@ -35,6 +35,7 @@ describe('llm response parsing', () => {
     const out = extractJson('sorry, I cannot do that as an AI');
     expect(out.findings).toEqual([]);
     expect(out.summary).toMatch(/could not be parsed/);
+    expect(out.parse_error).toBeTruthy();
   });
 
   it('extracts from plain object even with surrounding text', () => {
@@ -42,4 +43,12 @@ describe('llm response parsing', () => {
     const out = extractJson(text);
     expect(out.summary).toBe('clean');
   });
+
+  it('extractBalancedObject handles nested braces in strings', () => {
+    const text = 'prefix {"findings":[{"description":"uses {braces} ok"}],"summary":"s"} trailing';
+    const obj = extractBalancedObject(text);
+    expect(obj).not.toBeNull();
+    const parsed = JSON.parse(obj!);
+    expect(parsed.findings[0].description).toContain('{braces}');
+  });
 });

__tests__/prompts.test.ts

@@ -8,7 +8,7 @@ describe('prompts', () => {
     expect(inferLanguage('main.py')).toBe('Python');
     expect(inferLanguage('app.go')).toBe('Go');
     expect(inferLanguage('unknown.xyz')).toBe('Unknown');
-    expect(inferLanguage('Dockerfile')).toBe('Unknown'); // no ext
+    expect(inferLanguage('Dockerfile')).toBe('Dockerfile');
   });
 
   it('buildSystemPrompt includes strict security standards and categories', () => {
@@ -19,6 +19,8 @@ describe('prompts', () => {
     expect(prompt).toContain('secrets, injection');
     expect(prompt).toContain('NEVER include any text before or after the JSON');
     expect(prompt).toContain('ONLY a single valid JSON object');
+    expect(prompt).toContain('EXAMPLE FINDING (secret)');
+    expect(prompt).toContain('"category"');
   });
 
   it('buildUserPrompt includes policy, files and diff', () => {
@@ -32,6 +34,11 @@ describe('prompts', () => {
     expect(prompt).toContain(diff);
   });
 
+  it('buildUserPrompt notes truncation', () => {
+    const prompt = buildUserPrompt('diff', [{ filename: 'a.ts' }], DEFAULT_CONFIG, true);
+    expect(prompt).toContain('truncated');
+  });
+
   it('buildSystemPrompt lists custom rules when present', () => {
     const cfg = {
       ...DEFAULT_CONFIG,
@@ -44,4 +51,17 @@ describe('prompts', () => {
     expect(p).toContain('CUSTOM RULES');
     expect(p).toContain('No eval ever');
   });
+
+  it('buildSystemPrompt lists disabled categories', () => {
+    const cfg = {
+      ...DEFAULT_CONFIG,
+      policy: {
+        ...DEFAULT_CONFIG.policy,
+        categories: { ...DEFAULT_CONFIG.policy.categories, xss: false },
+      },
+    };
+    const p = buildSystemPrompt(cfg);
+    expect(p).toContain('DISABLED CATEGORIES');
+    expect(p).toContain('xss');
+  });
 });

src/checks.ts

@@ -1,6 +1,7 @@
 import * as core from '@actions/core';
 import * as github from '@actions/github';
 import type { Finding } from './types';
+import { shouldBlock } from './config';
 
 export type { Finding };
 
@@ -12,39 +13,44 @@ export async function createCheckRun(
   findings: Finding[],
   summary: string,
   blockOn: string,
-  annotate: boolean
+  annotate: boolean,
+  extraSummaryNote?: string
 ): Promise<{ conclusion: string; blockingCount: number }> {
   const octokit = github.getOctokit(token);
 
-  const blockingFindings = findings.filter(f => {
-    const order: any = { critical: 4, high: 3, medium: 2, low: 1, info: 0 };
-    return (order[f.severity] || 0) >= (order[blockOn] || 0);
-  });
+  const blockingFindings = findings.filter((f) => shouldBlock(f.severity, blockOn));
   const blockingCount = blockingFindings.length;
   const conclusion = blockingCount > 0 ? 'failure' : 'success';
 
   const annotations = annotate
-    ? findings.slice(0, 50).map(f => ({
+    ? findings.slice(0, 50).map((f) => ({
         path: f.file,
-        start_line: Math.max(1, f.start_line),
-        end_line: Math.max(f.start_line, f.end_line),
-        annotation_level: (f.severity === 'critical' || f.severity === 'high' ? 'failure' : f.severity === 'medium' ? 'warning' : 'notice') as 'failure' | 'warning' | 'notice',
+        start_line: Math.max(1, f.start_line || 1),
+        end_line: Math.max(f.start_line || 1, f.end_line || f.start_line || 1),
+        annotation_level: (f.severity === 'critical' || f.severity === 'high'
+          ? 'failure'
+          : f.severity === 'medium'
+            ? 'warning'
+            : 'notice') as 'failure' | 'warning' | 'notice',
         title: f.title,
-        message: `${f.severity.toUpperCase()} [${f.confidence}] ${f.description}\n\nRecommendation: ${f.recommendation}${f.cwe ? `\nCWE: ${f.cwe}` : ''}${f.owasp ? `\nOWASP: ${f.owasp}` : ''}`,
+        message: `${f.severity.toUpperCase()} [${f.confidence}] ${f.description}\n\nRecommendation: ${f.recommendation}${f.cwe ? `\nCWE: ${f.cwe}` : ''}${f.owasp ? `\nOWASP: ${f.owasp}` : ''}${f.category ? `\nCategory: ${f.category}` : ''}`,
       }))
     : [];
 
-  const title = blockingCount > 0
-    ? `${blockingCount} blocking security finding(s)`
-    : findings.length > 0
-    ? `${findings.length} security finding(s) (below threshold)`
-    : 'No security issues found';
+  const title =
+    blockingCount > 0
+      ? `${blockingCount} blocking security finding(s)`
+      : findings.length > 0
+        ? `${findings.length} security finding(s) (below threshold)`
+        : 'No security issues found';
 
-  const checkOutput = {
-    title,
-    summary: summary + (blockingCount > 0 ? `\n\n**${blockingCount} finding(s) meet or exceed block threshold (${blockOn}).**` : ''),
-    annotations,
-  };
+  let fullSummary = summary || '';
+  if (blockingCount > 0) {
+    fullSummary += `\n\n**${blockingCount} finding(s) meet or exceed block threshold (${blockOn}).**`;
+  }
+  if (extraSummaryNote) {
+    fullSummary += `\n\n${extraSummaryNote}`;
+  }
 
   const check = await octokit.rest.checks.create({
     owner,
@@ -52,8 +58,12 @@ export async function createCheckRun(
     name: 'SquidGate',
     head_sha: headSha,
     status: 'completed',
-    conclusion: conclusion as any,
-    output: checkOutput,
+    conclusion: conclusion as 'success' | 'failure' | 'neutral' | 'cancelled' | 'skipped' | 'timed_out' | 'action_required',
+    output: {
+      title,
+      summary: fullSummary,
+      annotations,
+    },
   });
 
   core.info(`Check run created: ${check.data.html_url}`);
@@ -84,7 +94,9 @@ export async function postPrComment(
   for (const f of shown) {
     const sev = f.severity.toUpperCase();
     body += `### ${sev} — ${f.title}\n`;
-    body += `**File:** \`${f.file}:${f.start_line}\`  |  **Confidence:** ${f.confidence}\n\n`;
+    body += `**File:** \`${f.file}:${f.start_line}\`  |  **Confidence:** ${f.confidence}`;
+    if (f.category) body += `  |  **Category:** ${f.category}`;
+    body += `\n\n`;
     body += `${f.description}\n\n`;
     if (f.cwe || f.owasp) {
       body += `CWE: ${f.cwe || 'N/A'} | OWASP: ${f.owasp || 'N/A'}\n\n`;

src/config.ts

@@ -1,6 +1,6 @@
 import * as fs from 'fs';
 import * as yaml from 'js-yaml';
-import type { SecurityScanConfig } from './types';
+import type { Finding, FindingCategory, SecurityScanConfig } from './types';
 
 export type { SecurityScanConfig } from './types';
 
@@ -64,50 +64,71 @@ export const SEVERITY_ORDER: Record<string, number> = {
   none: -1,
 };
 
+export const CONFIDENCE_ORDER: Record<string, number> = {
+  high: 2,
+  medium: 1,
+  low: 0,
+};
+
 export function shouldBlock(severity: string, threshold: string): boolean {
+  if (threshold === 'none') return false;
   const sev = SEVERITY_ORDER[severity] ?? -1;
   const thr = SEVERITY_ORDER[threshold] ?? 999;
   return sev >= thr;
 }
 
-function deepMerge(target: any, source: any): any {
-  for (const key of Object.keys(source)) {
-    if (source[key] && typeof source[key] === 'object' && !Array.isArray(source[key])) {
-      target[key] = deepMerge(target[key] || {}, source[key]);
-    } else {
-      target[key] = source[key];
-    }
-  }
-  return target;
+function isPlainObject(v: unknown): v is Record<string, unknown> {
+  return typeof v === 'object' && v !== null && !Array.isArray(v);
 }
 
-export function loadConfig(configPath: string, overrides: Partial<SecurityScanConfig> = {}): SecurityScanConfig {
-  let config = JSON.parse(JSON.stringify(DEFAULT_CONFIG)); // deep clone
+/** Recursive deep merge (arrays replaced, not concatenated). */
+export function deepMerge<T extends Record<string, unknown>>(target: T, source: Record<string, unknown>): T {
+  const out: Record<string, unknown> = { ...target };
+  for (const key of Object.keys(source)) {
+    const sv = source[key];
+    const tv = out[key];
+    if (isPlainObject(sv) && isPlainObject(tv)) {
+      out[key] = deepMerge(tv, sv);
+    } else if (sv !== undefined) {
+      out[key] = sv;
+    }
+  }
+  return out as T;
+}
+
+export function loadConfig(
+  configPath: string,
+  overrides: Partial<SecurityScanConfig> = {}
+): SecurityScanConfig {
+  let config = JSON.parse(JSON.stringify(DEFAULT_CONFIG)) as SecurityScanConfig;
 
   if (fs.existsSync(configPath)) {
     try {
       const raw = fs.readFileSync(configPath, 'utf8');
-      const parsed = yaml.load(raw) as Partial<SecurityScanConfig>;
-      if (parsed) {
-        config = deepMerge(config, parsed);
+      const parsed = yaml.load(raw) as Partial<SecurityScanConfig> | null;
+      if (parsed && isPlainObject(parsed)) {
+        config = deepMerge(
+          config as unknown as Record<string, unknown>,
+          parsed as Record<string, unknown>
+        ) as unknown as SecurityScanConfig;
       }
-    } catch (e) {
-      // caller can warn
-      console.warn(`Failed to parse config at ${configPath}: ${e}`);
+    } catch (e: unknown) {
+      const msg = e instanceof Error ? e.message : String(e);
+      console.warn(`Failed to parse config at ${configPath}: ${msg}`);
     }
   }
 
-  // Apply overrides (inputs)
   if (overrides.llm) {
     config.llm = { ...config.llm, ...overrides.llm };
   }
   if (overrides.policy) {
-    config.policy = { ...config.policy, ...overrides.policy };
-    if (overrides.policy.categories) {
-      config.policy.categories = { ...config.policy.categories, ...overrides.policy.categories };
+    const { categories, custom_rules, ...rest } = overrides.policy;
+    config.policy = { ...config.policy, ...rest };
+    if (categories) {
+      config.policy.categories = { ...config.policy.categories, ...categories };
     }
-    if (overrides.policy.custom_rules) {
-      config.policy.custom_rules = overrides.policy.custom_rules;
+    if (custom_rules) {
+      config.policy.custom_rules = custom_rules;
     }
   }
   if (overrides.context) {
@@ -117,15 +138,73 @@ export function loadConfig(configPath: string, overrides: Partial<SecurityScanCo
     config.output = { ...config.output, ...overrides.output };
   }
 
-  return config as SecurityScanConfig;
+  return config;
 }
 
-export function filterFindings(findings: any[], config: SecurityScanConfig): any[] {
-  const minConf = config.policy.min_confidence;
-  const confOrder: Record<string, number> = { high: 2, medium: 1, low: 0 };
+/** Map free-text / CWE hints to a policy category when the model omits `category`. */
+export function inferCategory(finding: Finding): FindingCategory | null {
+  if (finding.category && typeof finding.category === 'string') {
+    return finding.category.toLowerCase().replace(/[\s-]+/g, '_');
+  }
 
-  return findings.filter((f: any) => {
-    if (!f.confidence || confOrder[f.confidence] < confOrder[minConf]) return false;
+  const blob = [
+    finding.title,
+    finding.description,
+    finding.cwe || '',
+    finding.owasp || '',
+    finding.recommendation,
+  ]
+    .join(' ')
+    .toLowerCase();
+
+  const rules: Array<[RegExp, FindingCategory]> = [
+    [/\b(secret|api[_ ]?key|private[_ ]?key|token|password|credential|cwe-798|cwe-259)\b/, 'secrets'],
+    [/\b(hardcoded).*(password|secret|key|token)\b/, 'hardcoded_credentials'],
+    [/\b(sql\s*inject|command\s*inject|os\s*command|cwe-89|cwe-78|cwe-77|sqli)\b/, 'injection'],
+    [/\b(xss|cross-site scripting|cwe-79)\b/, 'xss'],
+    [/\b(ssrf|server-side request|cwe-918)\b/, 'ssrf'],
+    [/\b(path\s*traversal|directory\s*traversal|cwe-22)\b/, 'path_traversal'],
+    [/\b(csrf|cross-site request forgery|cwe-352)\b/, 'csrf'],
+    [/\b(deserializ|pickle|yaml\.load|cwe-502)\b/, 'insecure_deserialization'],
+    [/\b(md5|sha1|ecb|weak\s*crypto|cwe-327|cwe-328)\b/, 'cryptography'],
+    [/\b(authn|authz|authorization|authentication|idor|broken access|cwe-287|cwe-862|cwe-863)\b/, 'authn_authz'],
+    [/\b(eval\(|exec\(|child_process|dangerous function)\b/, 'dangerous_functions'],
+    [/\b(supply.?chain|dependency|typosquat|malicious package)\b/, 'supply_chain'],
+    [/\b(misconfig|debug\s*=\s*true|permissive cors)\b/, 'misconfiguration'],
+  ];
+
+  for (const [re, cat] of rules) {
+    if (re.test(blob)) return cat;
+  }
+  return null;
+}
+
+export function filterFindings(findings: Finding[], config: SecurityScanConfig): Finding[] {
+  const minConf = config.policy.min_confidence;
+  const minScore = CONFIDENCE_ORDER[minConf] ?? 1;
+  const categories = config.policy.categories || {};
+
+  // If every known category is enabled (or categories empty), only confidence filters.
+  const disabled = new Set(
+    Object.entries(categories)
+      .filter(([, enabled]) => enabled === false)
+      .map(([k]) => k.toLowerCase())
+  );
+
+  return findings.filter((f) => {
+    const conf = CONFIDENCE_ORDER[f.confidence] ?? -1;
+    if (conf < minScore) return false;
+
+    if (disabled.size === 0) return true;
+
+    const cat = inferCategory(f);
+    if (!cat) {
+      // Unknown category: keep (do not drop on inference failure)
+      return true;
+    }
+    if (disabled.has(cat.toLowerCase())) return false;
+    // Explicit false for this key
+    if (categories[cat] === false) return false;
     return true;
   });
 }

src/diff.ts

@@ -0,0 +1,224 @@
+import * as core from '@actions/core';
+import * as github from '@actions/github';
+import * as fs from 'fs';
+import { exec } from '@actions/exec';
+
+export interface DiffResult {
+  diff: string;
+  truncated: boolean;
+  source: 'api' | 'git';
+  originalBytes: number;
+}
+
+export interface ChangedFile {
+  filename: string;
+  status: string;
+  additions: number;
+  deletions: number;
+  patch?: string;
+}
+
+async function captureExec(cmd: string, args: string[]): Promise<{ exit: number; stdout: string }> {
+  let stdout = '';
+  const exit = await exec(cmd, args, {
+    listeners: {
+      stdout: (data: Buffer) => {
+        stdout += data.toString();
+      },
+    },
+    silent: true,
+    ignoreReturnCode: true,
+  });
+  return { exit, stdout };
+}
+
+function truncateDiff(diff: string, maxBytes: number): { diff: string; truncated: boolean; originalBytes: number } {
+  const originalBytes = Buffer.byteLength(diff, 'utf8');
+  if (originalBytes <= maxBytes) {
+    return { diff, truncated: false, originalBytes };
+  }
+  // Prefer character cut near maxBytes (byte-aware enough for ASCII-heavy diffs)
+  const cut = diff.substring(0, maxBytes);
+  const msg =
+    `\n... [diff truncated: ${originalBytes} bytes → ${maxBytes} byte limit; ` +
+    `raise context.max_diff_bytes if needed]\n`;
+  return { diff: cut + msg, truncated: true, originalBytes };
+}
+
+/**
+ * Prefer GitHub API diff (reliable for forks / shallow clones).
+ * Fall back to local git with merge-base and multiple base candidates.
+ */
+export async function getPullRequestDiff(
+  token: string,
+  owner: string,
+  repo: string,
+  pullNumber: number,
+  maxBytes: number,
+  linesBefore: number = 30,
+  linesAfter: number = 30
+): Promise<DiffResult> {
+  // API-first: works for fork PRs and does not depend on local fetch depth.
+  try {
+    const result = await getDiffViaApi(token, owner, repo, pullNumber, maxBytes);
+    if (result.diff.trim()) {
+      if (result.truncated) {
+        core.warning(
+          `Diff truncated to ${maxBytes} bytes (original ~${result.originalBytes} bytes). ` +
+            `Security analysis may miss findings outside the kept window.`
+        );
+      }
+      return result;
+    }
+    core.warning('GitHub API returned empty diff; trying local git');
+  } catch (e: unknown) {
+    const msg = e instanceof Error ? e.message : String(e);
+    core.warning(`GitHub API diff failed (${msg}); falling back to local git`);
+  }
+
+  if (!fs.existsSync('.git')) {
+    throw new Error('No PR diff available (API failed and no local .git)');
+  }
+
+  const result = await getDiffViaGit(maxBytes, linesBefore, linesAfter);
+  if (result.truncated) {
+    core.warning(
+      `Diff truncated to ${maxBytes} bytes (original ~${result.originalBytes} bytes). ` +
+        `Security analysis may miss findings outside the kept window.`
+    );
+  }
+  return result;
+}
+
+async function getDiffViaApi(
+  token: string,
+  owner: string,
+  repo: string,
+  pullNumber: number,
+  maxBytes: number
+): Promise<DiffResult> {
+  const octokit = github.getOctokit(token);
+  const response = await octokit.rest.pulls.get({
+    owner,
+    repo,
+    pull_number: pullNumber,
+    mediaType: { format: 'diff' },
+  });
+
+  const raw = String(response.data);
+  const { diff, truncated, originalBytes } = truncateDiff(raw, maxBytes);
+  return { diff, truncated, source: 'api', originalBytes };
+}
+
+async function getDiffViaGit(
+  maxBytes: number,
+  linesBefore: number,
+  linesAfter: number
+): Promise<DiffResult> {
+  const baseRef = process.env.GITHUB_BASE_REF;
+  const headSha = process.env.GITHUB_SHA || 'HEAD';
+  // git --unified is symmetric; honor the larger of before/after.
+  const contextLines = Math.max(1, linesBefore || 0, linesAfter || 0);
+
+  if (baseRef) {
+    try {
+      await exec('git', ['fetch', 'origin', baseRef, '--depth=100'], { silent: true, ignoreReturnCode: true });
+    } catch {
+      /* best effort */
+    }
+    // Also try fetching the PR merge base refs GitHub Actions often provides
+    try {
+      await exec('git', ['fetch', 'origin', `pull/${process.env.GITHUB_REF_NAME || ''}`, '--depth=100'], {
+        silent: true,
+        ignoreReturnCode: true,
+      });
+    } catch {
+      /* best effort */
+    }
+  }
+
+  const candidates: string[] = [];
+  if (baseRef) {
+    candidates.push(`origin/${baseRef}`);
+    candidates.push(baseRef);
+  }
+  // merge-base against origin/base when available
+  if (baseRef) {
+    const mb = await captureExec('git', ['merge-base', `origin/${baseRef}`, headSha]);
+    if (mb.exit === 0 && mb.stdout.trim()) {
+      candidates.unshift(mb.stdout.trim());
+    }
+  }
+  candidates.push('HEAD^');
+
+  let lastError = 'no candidates produced diff';
+  for (const base of candidates) {
+    // Three-dot: changes on head since merge-base with base
+    for (const range of [`${base}...${headSha}`, `${base}..${headSha}`]) {
+      const { exit, stdout } = await captureExec('git', ['diff', `--unified=${contextLines}`, range]);
+      if ((exit === 0 || stdout.trim()) && stdout.trim()) {
+        const { diff, truncated, originalBytes } = truncateDiff(stdout, maxBytes);
+        return { diff, truncated, source: 'git', originalBytes };
+      }
+      lastError = `git diff ${range} exit=${exit}`;
+    }
+  }
+
+  // Last resort: single commit
+  const fallback = await captureExec('git', ['diff', `--unified=${contextLines}`, 'HEAD^..HEAD']);
+  if (fallback.stdout.trim()) {
+    const { diff, truncated, originalBytes } = truncateDiff(fallback.stdout, maxBytes);
+    return { diff, truncated, source: 'git', originalBytes };
+  }
+
+  throw new Error(`git diff produced no output (${lastError})`);
+}
+
+export async function getChangedFiles(
+  token: string,
+  owner: string,
+  repo: string,
+  pullNumber: number,
+  maxFiles: number
+): Promise<ChangedFile[]> {
+  try {
+    const octokit = github.getOctokit(token);
+    const { data: files } = await octokit.rest.pulls.listFiles({
+      owner,
+      repo,
+      pull_number: pullNumber,
+      per_page: Math.min(maxFiles, 100),
+    });
+    return files.slice(0, maxFiles).map((f) => ({
+      filename: f.filename,
+      status: f.status,
+      additions: f.additions,
+      deletions: f.deletions,
+      patch: f.patch,
+    }));
+  } catch (e: unknown) {
+    const msg = e instanceof Error ? e.message : String(e);
+    core.warning(`listFiles via API failed (${msg}), falling back to git`);
+    return getChangedFilesViaGit(maxFiles);
+  }
+}
+
+async function getChangedFilesViaGit(maxFiles: number): Promise<ChangedFile[]> {
+  const baseRef = process.env.GITHUB_BASE_REF;
+  const headSha = process.env.GITHUB_SHA || 'HEAD';
+  const ranges = baseRef
+    ? [`origin/${baseRef}...${headSha}`, 'HEAD^..HEAD']
+    : ['HEAD^..HEAD'];
+
+  for (const range of ranges) {
+    const { stdout } = await captureExec('git', ['diff', '--name-status', range]);
+    const lines = stdout.trim().split('\n').filter(Boolean).slice(0, maxFiles);
+    if (lines.length === 0) continue;
+    return lines.map((line) => {
+      const [status, ...rest] = line.split('\t');
+      const file = rest.join('\t');
+      return { filename: file, status: status || 'M', additions: 0, deletions: 0 };
+    });
+  }
+  return [];
+}

src/index.ts

@@ -1,7 +1,5 @@
 import * as core from '@actions/core';
 import * as github from '@actions/github';
-import * as fs from 'fs';
-import { exec } from '@actions/exec';
 
 import {
   DEFAULT_CONFIG,
@@ -10,112 +8,9 @@ import {
 } from './config';
 import { buildSystemPrompt, buildUserPrompt } from './prompts';
 import { callLlm } from './llm';
-import type { LlmResponse } from './types';
+import type { BlockOn, LlmResponse, SecurityScanConfig } from './types';
 import { createCheckRun, postPrComment } from './checks';
-
-async function getPullRequestDiff(
-  token: string,
-  owner: string,
-  repo: string,
-  pullNumber: number,
-  maxBytes: number,
-  contextLines: number = 30
-): Promise<string> {
-  if (fs.existsSync('.git')) {
-    try {
-      return await getDiffViaGit(maxBytes, contextLines);
-    } catch (e: any) {
-      core.warning(`Local git diff failed (${e.message}), falling back to GitHub API`);
-    }
-  }
-
-  const octokit = github.getOctokit(token);
-  try {
-    const response = await octokit.rest.pulls.get({
-      owner,
-      repo,
-      pull_number: pullNumber,
-      mediaType: { format: 'diff' },
-    });
-
-    let diff = String(response.data);
-    if (diff.length > maxBytes) {
-      diff = diff.substring(0, maxBytes) + '\n... [diff truncated]';
-      core.warning(`Diff truncated to ${maxBytes} bytes`);
-    }
-    return diff;
-  } catch (error: any) {
-    core.warning(`Failed to fetch diff via GitHub API: ${error.message}`);
-    throw error;
-  }
-}
-
-async function getDiffViaGit(maxBytes: number, contextLines: number = 30): Promise<string> {
-  const baseRef = process.env.GITHUB_BASE_REF;
-  const headSha = process.env.GITHUB_SHA || 'HEAD';
-
-  let base = 'HEAD^';
-  if (baseRef) {
-    base = `origin/${baseRef}`;
-    try {
-      await exec('git', ['fetch', 'origin', baseRef, '--depth=100'], { silent: true });
-    } catch {}
-  }
-
-  let diff = '';
-  const options: any = {
-    listeners: { stdout: (data: Buffer) => { diff += data.toString(); } },
-    silent: true,
-    ignoreReturnCode: true,
-  };
-
-  let exit = await exec('git', ['diff', `--unified=${contextLines}`, `${base}...${headSha}`], options);
-  if (exit !== 0 && !diff) {
-    await exec('git', ['diff', `--unified=${contextLines}`, 'HEAD^..HEAD'], options);
-  }
-
-  if (diff.length > maxBytes) {
-    diff = diff.substring(0, maxBytes) + '\n... [diff truncated]';
-  }
-  if (!diff.trim()) {
-    throw new Error('git diff produced no output');
-  }
-  return diff;
-}
-
-async function getChangedFiles(
-  token: string,
-  owner: string,
-  repo: string,
-  pullNumber: number,
-  maxFiles: number
-): Promise<Array<{ filename: string; status: string; additions: number; deletions: number; patch?: string }>> {
-  try {
-    const octokit = github.getOctokit(token);
-    const { data: files } = await octokit.rest.pulls.listFiles({
-      owner,
-      repo,
-      pull_number: pullNumber,
-      per_page: maxFiles,
-    });
-    return files.slice(0, maxFiles);
-  } catch (e) {
-    core.warning('listFiles via API failed, falling back to git');
-    return await getChangedFilesViaGit(maxFiles);
-  }
-}
-
-async function getChangedFilesViaGit(maxFiles: number) {
-  let out = '';
-  const opts: any = { listeners: { stdout: (d: Buffer) => { out += d.toString(); } }, silent: true, ignoreReturnCode: true };
-  await exec('git', ['diff', '--name-status', 'HEAD^..HEAD'], opts);
-  const lines = out.trim().split('\n').filter(Boolean).slice(0, maxFiles);
-  return lines.map(line => {
-    const [status, ...rest] = line.split('\t');
-    const file = rest.join('\t');
-    return { filename: file, status: status || 'M', additions: 0, deletions: 0 };
-  });
-}
+import { getChangedFiles, getPullRequestDiff } from './diff';
 
 async function run(): Promise<void> {
   try {
@@ -125,7 +20,7 @@ async function run(): Promise<void> {
     const configPath = core.getInput('config-path') || '.github/squidgate.yml';
     const overrideProvider = core.getInput('llm-provider');
     const overrideModel = core.getInput('llm-model');
-    const overrideBlockOn = core.getInput('block-on') as any;
+    const overrideBlockOn = core.getInput('block-on') as BlockOn | '';
     const baseUrl = core.getInput('llm-base-url') || undefined;
 
     const context = github.context;
@@ -141,7 +36,7 @@ async function run(): Promise<void> {
 
     core.info(`Analyzing PR #${pullNumber} @ ${headSha}`);
 
-    const overrides: any = {};
+    const overrides: Partial<SecurityScanConfig> = {};
     if (overrideProvider || overrideModel) {
       overrides.llm = {
         provider: overrideProvider || DEFAULT_CONFIG.llm.provider,
@@ -149,27 +44,45 @@ async function run(): Promise<void> {
       };
     }
     if (overrideBlockOn) {
-      overrides.policy = { block_on: overrideBlockOn };
+      overrides.policy = { block_on: overrideBlockOn } as SecurityScanConfig['policy'];
     }
 
     const config = loadConfig(configPath, overrides);
 
-    core.info(`Using provider=${config.llm.provider} model=${config.llm.model} block_on=${config.policy.block_on}`);
+    core.info(
+      `Using provider=${config.llm.provider} model=${config.llm.model} block_on=${config.policy.block_on}`
+    );
 
-    const diff = await getPullRequestDiff(
+    const diffResult = await getPullRequestDiff(
       token,
       owner,
       repo,
       pullNumber,
       config.context.max_diff_bytes,
-      config.context.lines_before
+      config.context.lines_before,
+      config.context.lines_after
+    );
+    const changedFiles = await getChangedFiles(
+      token,
+      owner,
+      repo,
+      pullNumber,
+      config.context.max_files
     );
-    const changedFiles = await getChangedFiles(token, owner, repo, pullNumber, config.context.max_files);
 
-    core.info(`Diff size: ${diff.length} bytes, ${changedFiles.length} files`);
+    core.info(
+      `Diff source=${diffResult.source} size=${diffResult.diff.length} bytes` +
+        (diffResult.truncated ? ` (truncated from ${diffResult.originalBytes})` : '') +
+        `, ${changedFiles.length} files`
+    );
 
     const systemPrompt = buildSystemPrompt(config);
-    const userPrompt = buildUserPrompt(diff, changedFiles, config);
+    const userPrompt = buildUserPrompt(
+      diffResult.diff,
+      changedFiles,
+      config,
+      diffResult.truncated
+    );
 
     core.info('Calling LLM for security analysis...');
     let llmResponse: LlmResponse;
@@ -182,21 +95,66 @@ async function run(): Promise<void> {
         userPrompt,
         baseUrl
       );
-    } catch (llmErr: any) {
-      core.error(`LLM call failed: ${llmErr.message}`);
+    } catch (llmErr: unknown) {
+      const msg = llmErr instanceof Error ? llmErr.message : String(llmErr);
+      core.error(`LLM call failed: ${msg}`);
       if (config.output.fail_on_error) {
         try {
-          await createCheckRun(token, owner, repo, headSha, [], `LLM call failed: ${llmErr.message}`, 'high', false);
-        } catch {}
-        core.setFailed(`LLM call failed: ${llmErr.message}`);
+          await createCheckRun(
+            token,
+            owner,
+            repo,
+            headSha,
+            [],
+            `LLM call failed: ${msg}`,
+            config.policy.block_on,
+            false,
+            'Check failed because the LLM provider call errored (fail_on_error: true).'
+          );
+        } catch {
+          /* ignore check create failure */
+        }
+        core.setFailed(`LLM call failed: ${msg}`);
       } else {
-        core.warning('fail_on_error is false, marking neutral.');
+        core.warning('fail_on_error is false, marking neutral / continuing without findings.');
+        try {
+          await createCheckRun(
+            token,
+            owner,
+            repo,
+            headSha,
+            [],
+            `LLM call failed (non-blocking): ${msg}`,
+            'none',
+            false,
+            'fail_on_error is false — check concluded success with no findings.'
+          );
+        } catch {
+          /* ignore */
+        }
       }
       return;
     }
 
+    if (llmResponse.parse_error) {
+      core.warning(`LLM JSON parse issue: ${llmResponse.parse_error}`);
+    }
+
     const filtered = filterFindings(llmResponse.findings, config);
-    core.info(`LLM returned ${llmResponse.findings.length} findings, ${filtered.length} after filtering`);
+    core.info(
+      `LLM returned ${llmResponse.findings.length} findings, ${filtered.length} after filtering`
+    );
+
+    const extraNotes: string[] = [];
+    if (diffResult.truncated) {
+      extraNotes.push(
+        `⚠️ Diff was truncated (${diffResult.originalBytes} → ${config.context.max_diff_bytes} bytes). ` +
+          `Findings may be incomplete.`
+      );
+    }
+    if (llmResponse.parse_error) {
+      extraNotes.push(`⚠️ LLM response parse issue: ${llmResponse.parse_error}`);
+    }
 
     const { conclusion, blockingCount } = await createCheckRun(
       token,
@@ -206,11 +164,20 @@ async function run(): Promise<void> {
       filtered,
       llmResponse.summary,
       config.policy.block_on,
-      config.output.annotate_lines
+      config.output.annotate_lines,
+      extraNotes.length ? extraNotes.join('\n') : undefined
     );
 
     if (config.output.comment_on_pr) {
-      await postPrComment(token, owner, repo, pullNumber, filtered, llmResponse.summary, blockingCount);
+      await postPrComment(
+        token,
+        owner,
+        repo,
+        pullNumber,
+        filtered,
+        llmResponse.summary,
+        blockingCount
+      );
     }
 
     core.setOutput('findings-count', filtered.length.toString());
@@ -218,19 +185,23 @@ async function run(): Promise<void> {
     core.setOutput('conclusion', conclusion);
 
     if (conclusion === 'failure') {
-      core.setFailed(`${blockingCount} security finding(s) reached or exceeded the block_on severity of '${config.policy.block_on}'`);
+      core.setFailed(
+        `${blockingCount} security finding(s) reached or exceeded the block_on severity of '${config.policy.block_on}'`
+      );
     } else {
       core.info('Security scan passed.');
     }
-  } catch (error: any) {
-    core.setFailed(error.message);
+  } catch (error: unknown) {
+    const msg = error instanceof Error ? error.message : String(error);
+    core.setFailed(msg);
   }
 }
 
 run();
 
-export { loadConfig, filterFindings, shouldBlock } from './config';
+export { loadConfig, filterFindings, shouldBlock, inferCategory, deepMerge } from './config';
 export { inferLanguage, buildSystemPrompt, buildUserPrompt } from './prompts';
-export { callLlm, extractJson, normalizeResponse } from './llm';
+export { callLlm, extractJson, extractBalancedObject, normalizeResponse } from './llm';
 export { createCheckRun, postPrComment } from './checks';
+export { getPullRequestDiff, getChangedFiles } from './diff';
 export type { Finding, LlmResponse, SecurityScanConfig } from './types';

src/llm.ts

@@ -3,35 +3,136 @@ import type { Finding, LlmResponse } from './types';
 
 export type { Finding, LlmResponse };
 
-export function normalizeResponse(resp: any): LlmResponse {
-  if (!resp) resp = {};
-  if (!Array.isArray(resp.findings)) resp.findings = [];
-  if (typeof resp.summary !== 'string') resp.summary = '';
-  return resp as LlmResponse;
+const RETRYABLE_STATUS = new Set([408, 425, 429, 500, 502, 503, 504]);
+const MAX_ATTEMPTS = 3;
+const BASE_DELAY_MS = 800;
+
+function sleep(ms: number): Promise<void> {
+  return new Promise((resolve) => setTimeout(resolve, ms));
 }
 
+export function normalizeResponse(resp: unknown): LlmResponse {
+  const r = (resp && typeof resp === 'object' ? resp : {}) as Record<string, unknown>;
+  const findings = Array.isArray(r.findings) ? (r.findings as Finding[]) : [];
+  const summary = typeof r.summary === 'string' ? r.summary : '';
+  return { findings, summary };
+}
+
+/**
+ * Extract the first balanced JSON object from text (handles nested braces).
+ * Falls back to code-fence contents, then empty findings with parse_error.
+ */
 export function extractJson(text: string): LlmResponse {
-  if (!text) return { findings: [], summary: 'LLM response could not be parsed.' };
-  // Try to find the largest JSON object
-  const jsonMatch = text.match(/\{[\s\S]*\}/);
-  if (jsonMatch) {
-    try {
-      const parsed = JSON.parse(jsonMatch[0]);
-      return normalizeResponse(parsed);
-    } catch (e) {
-      // fallthrough
-    }
+  if (!text || !text.trim()) {
+    return {
+      findings: [],
+      summary: 'LLM response could not be parsed.',
+      parse_error: 'empty response',
+      raw: text || '',
+    };
   }
-  // Try code fence
-  const fence = text.match(/```(?:json)?\s*([\s\S]*?)\s*```/);
+
+  // Prefer fenced JSON first (common for Anthropic/Google)
+  const fence = text.match(/```(?:json)?\s*([\s\S]*?)\s*```/i);
   if (fence) {
     try {
-      const parsed = JSON.parse(fence[1]);
-      return normalizeResponse(parsed);
-    } catch {}
+      return normalizeResponse(JSON.parse(fence[1]));
+    } catch {
+      /* try balanced object next */
+    }
   }
+
+  const sliced = extractBalancedObject(text);
+  if (sliced) {
+    try {
+      return normalizeResponse(JSON.parse(sliced));
+    } catch (e: unknown) {
+      const msg = e instanceof Error ? e.message : String(e);
+      core.warning(`JSON object found but failed to parse: ${msg}`);
+    }
+  }
+
+  // Last try: whole string
+  try {
+    return normalizeResponse(JSON.parse(text.trim()));
+  } catch {
+    /* fallthrough */
+  }
+
   core.warning('Could not parse LLM response as JSON. Returning empty findings.');
-  return { findings: [], summary: 'LLM response could not be parsed.' };
+  return {
+    findings: [],
+    summary: 'LLM response could not be parsed.',
+    parse_error: 'no valid JSON object in model output',
+    raw: text.length > 2000 ? text.slice(0, 2000) + '…' : text,
+  };
+}
+
+/** Find first top-level `{ ... }` with string/escape awareness. */
+export function extractBalancedObject(text: string): string | null {
+  const start = text.indexOf('{');
+  if (start < 0) return null;
+
+  let depth = 0;
+  let inString = false;
+  let escape = false;
+
+  for (let i = start; i < text.length; i++) {
+    const ch = text[i];
+    if (inString) {
+      if (escape) {
+        escape = false;
+      } else if (ch === '\\') {
+        escape = true;
+      } else if (ch === '"') {
+        inString = false;
+      }
+      continue;
+    }
+    if (ch === '"') {
+      inString = true;
+      continue;
+    }
+    if (ch === '{') depth++;
+    else if (ch === '}') {
+      depth--;
+      if (depth === 0) {
+        return text.slice(start, i + 1);
+      }
+    }
+  }
+  return null;
+}
+
+async function fetchWithRetry(
+  url: string,
+  init: RequestInit,
+  label: string
+): Promise<Response> {
+  let lastErr: Error | undefined;
+  for (let attempt = 1; attempt <= MAX_ATTEMPTS; attempt++) {
+    try {
+      const res = await fetch(url, init);
+      if (res.ok) return res;
+
+      const errText = await res.text();
+      const retryable = RETRYABLE_STATUS.has(res.status);
+      lastErr = new Error(`${label} error (${res.status}): ${errText}`);
+      if (!retryable || attempt === MAX_ATTEMPTS) throw lastErr;
+
+      const delay = BASE_DELAY_MS * Math.pow(2, attempt - 1) + Math.floor(Math.random() * 200);
+      core.warning(`${label} ${res.status}; retry ${attempt}/${MAX_ATTEMPTS} in ${delay}ms`);
+      await sleep(delay);
+    } catch (e: unknown) {
+      if (e instanceof Error && e.message.includes(' error (')) throw e;
+      lastErr = e instanceof Error ? e : new Error(String(e));
+      if (attempt === MAX_ATTEMPTS) throw lastErr;
+      const delay = BASE_DELAY_MS * Math.pow(2, attempt - 1);
+      core.warning(`${label} network error; retry ${attempt}/${MAX_ATTEMPTS} in ${delay}ms: ${lastErr.message}`);
+      await sleep(delay);
+    }
+  }
+  throw lastErr || new Error(`${label} failed`);
 }
 
 async function callOpenAICompatible(
@@ -46,8 +147,8 @@ async function callOpenAICompatible(
     provider === 'azure'
       ? `${baseUrl || process.env.AZURE_OPENAI_ENDPOINT}/openai/deployments/${model}/chat/completions?api-version=2024-02-15-preview`
       : provider === 'custom'
-      ? `${baseUrl || 'http://localhost:11434/v1'}/chat/completions`
-      : 'https://api.openai.com/v1/chat/completions';
+        ? `${(baseUrl || 'http://localhost:11434/v1').replace(/\/$/, '')}/chat/completions`
+        : 'https://api.openai.com/v1/chat/completions';
 
   const headers: Record<string, string> = {
     'Content-Type': 'application/json',
@@ -59,33 +160,29 @@ async function callOpenAICompatible(
     headers['Authorization'] = `Bearer ${apiKey}`;
   }
 
-  const body: any = {
+  const body = {
     model,
     messages: [
       { role: 'system', content: systemPrompt },
       { role: 'user', content: userPrompt },
     ],
     temperature: 0,
-    response_format: { type: 'json_object' },
+    response_format: { type: 'json_object' as const },
     max_tokens: 4000,
   };
 
-  const res = await fetch(url, {
-    method: 'POST',
-    headers,
-    body: JSON.stringify(body),
-  });
+  const res = await fetchWithRetry(
+    url,
+    { method: 'POST', headers, body: JSON.stringify(body) },
+    'LLM API'
+  );
 
-  if (!res.ok) {
-    const errText = await res.text();
-    throw new Error(`LLM API error (${res.status}): ${errText}`);
-  }
-
-  const data: any = await res.json();
+  const data = (await res.json()) as {
+    choices?: Array<{ message?: { content?: string } }>;
+  };
   const content = data.choices?.[0]?.message?.content || '{}';
   try {
-    const parsed = JSON.parse(content);
-    return normalizeResponse(parsed);
+    return normalizeResponse(JSON.parse(content));
   } catch {
     return extractJson(content);
   }
@@ -95,62 +192,84 @@ async function callAnthropic(
   model: string,
   apiKey: string,
   systemPrompt: string,
-  userPrompt: string
+  userPrompt: string,
+  retryJsonOnly = false
 ): Promise<LlmResponse> {
   const url = 'https://api.anthropic.com/v1/messages';
+  const userContent = retryJsonOnly
+    ? `${userPrompt}\n\nIMPORTANT: Your previous reply was not valid JSON. Reply with ONLY the JSON object, no markdown fences, no prose.`
+    : userPrompt;
 
-  const res = await fetch(url, {
-    method: 'POST',
-    headers: {
-      'Content-Type': 'application/json',
-      'x-api-key': apiKey,
-      'anthropic-version': '2023-06-01',
+  const res = await fetchWithRetry(
+    url,
+    {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+        'x-api-key': apiKey,
+        'anthropic-version': '2023-06-01',
+      },
+      body: JSON.stringify({
+        model,
+        max_tokens: 4000,
+        temperature: 0,
+        system: systemPrompt,
+        messages: [{ role: 'user', content: userContent }],
+      }),
     },
-    body: JSON.stringify({
-      model,
-      max_tokens: 4000,
-      temperature: 0,
-      system: systemPrompt,
-      messages: [{ role: 'user', content: userPrompt }],
-    }),
-  });
+    'Anthropic'
+  );
 
-  if (!res.ok) {
-    throw new Error(`Anthropic error: ${await res.text()}`);
+  const data = (await res.json()) as {
+    content?: Array<{ type?: string; text?: string }>;
+  };
+  const text = data.content?.find((c) => c.type === 'text')?.text || '{}';
+  const parsed = extractJson(text);
+  if (parsed.parse_error && !retryJsonOnly) {
+    core.warning('Anthropic JSON parse failed; retrying with stricter instruction');
+    return callAnthropic(model, apiKey, systemPrompt, userPrompt, true);
   }
-
-  const data: any = await res.json();
-  const text = data.content?.find((c: any) => c.type === 'text')?.text || '{}';
-  return extractJson(text);
+  return parsed;
 }
 
 async function callGoogle(
   model: string,
   apiKey: string,
   systemPrompt: string,
-  userPrompt: string
+  userPrompt: string,
+  retryJsonOnly = false
 ): Promise<LlmResponse> {
   const url = `https://generativelanguage.googleapis.com/v1beta/models/${model}:generateContent?key=${apiKey}`;
+  const textIn = retryJsonOnly
+    ? `${systemPrompt}\n\n${userPrompt}\n\nIMPORTANT: Previous reply was not valid JSON. Return ONLY the JSON object.`
+    : `${systemPrompt}\n\n${userPrompt}`;
 
-  const res = await fetch(url, {
-    method: 'POST',
-    headers: { 'Content-Type': 'application/json' },
-    body: JSON.stringify({
-      contents: [{ parts: [{ text: systemPrompt + '\n\n' + userPrompt }] }],
-      generationConfig: {
-        temperature: 0,
-        responseMimeType: 'application/json',
-      },
-    }),
-  });
+  const res = await fetchWithRetry(
+    url,
+    {
+      method: 'POST',
+      headers: { 'Content-Type': 'application/json' },
+      body: JSON.stringify({
+        contents: [{ parts: [{ text: textIn }] }],
+        generationConfig: {
+          temperature: 0,
+          responseMimeType: 'application/json',
+        },
+      }),
+    },
+    'Google'
+  );
 
-  if (!res.ok) {
-    throw new Error(`Google error: ${await res.text()}`);
-  }
-
-  const data: any = await res.json();
+  const data = (await res.json()) as {
+    candidates?: Array<{ content?: { parts?: Array<{ text?: string }> } }>;
+  };
   const text = data.candidates?.[0]?.content?.parts?.[0]?.text || '{}';
-  return extractJson(text);
+  const parsed = extractJson(text);
+  if (parsed.parse_error && !retryJsonOnly) {
+    core.warning('Google JSON parse failed; retrying with stricter instruction');
+    return callGoogle(model, apiKey, systemPrompt, userPrompt, true);
+  }
+  return parsed;
 }
 
 export async function callLlm(
@@ -162,11 +281,13 @@ export async function callLlm(
   baseUrl?: string
 ): Promise<LlmResponse> {
   if (provider === 'openai' || provider === 'azure' || provider === 'custom') {
-    return await callOpenAICompatible(provider, model, apiKey, systemPrompt, userPrompt, baseUrl);
-  } else if (provider === 'anthropic') {
-    return await callAnthropic(model, apiKey, systemPrompt, userPrompt);
-  } else if (provider === 'google') {
-    return await callGoogle(model, apiKey, systemPrompt, userPrompt);
+    return callOpenAICompatible(provider, model, apiKey, systemPrompt, userPrompt, baseUrl);
+  }
+  if (provider === 'anthropic') {
+    return callAnthropic(model, apiKey, systemPrompt, userPrompt);
+  }
+  if (provider === 'google') {
+    return callGoogle(model, apiKey, systemPrompt, userPrompt);
   }
   throw new Error(`Unsupported LLM provider: ${provider}`);
 }

src/prompts.ts

@@ -1,6 +1,10 @@
 import * as path from 'path';
+import type { SecurityScanConfig } from './types';
 
 export function inferLanguage(filename: string): string {
+  const base = path.basename(filename).toLowerCase();
+  if (base === 'dockerfile' || base.startsWith('dockerfile.')) return 'Dockerfile';
+
   const ext = path.extname(filename).toLowerCase();
   const map: Record<string, string> = {
     '.js': 'JavaScript',
@@ -32,15 +36,58 @@ export function inferLanguage(filename: string): string {
   return map[ext] || 'Unknown';
 }
 
-export function buildSystemPrompt(config: any): string {
+const FEW_SHOT = `
+EXAMPLE FINDING (secret):
+{
+  "file": "src/config.py",
+  "start_line": 12,
+  "end_line": 12,
+  "severity": "critical",
+  "title": "Hardcoded API key",
+  "description": "A live-looking API key is embedded in source on a changed line.",
+  "cwe": "CWE-798",
+  "owasp": "A07:2021",
+  "recommendation": "Load secrets from environment or a secret manager; rotate the exposed key.",
+  "confidence": "high",
+  "category": "secrets"
+}
+
+EXAMPLE FINDING (injection):
+{
+  "file": "api/users.py",
+  "start_line": 40,
+  "end_line": 42,
+  "severity": "high",
+  "title": "SQL injection via string concatenation",
+  "description": "User input is concatenated into a SQL query without parameterization.",
+  "cwe": "CWE-89",
+  "owasp": "A03:2021",
+  "recommendation": "Use parameterized queries or an ORM binder; never interpolate untrusted input into SQL.",
+  "confidence": "high",
+  "category": "injection"
+}
+`.trim();
+
+export function buildSystemPrompt(config: SecurityScanConfig): string {
   const enabledCategories = Object.entries(config.policy.categories)
     .filter(([, v]) => v)
     .map(([k]) => k)
     .join(', ');
 
-  const custom = config.policy.custom_rules.length > 0
-    ? 'CUSTOM RULES (must also enforce):\n' + config.policy.custom_rules.map((r: string) => '- ' + r).join('\n')
-    : '';
+  const disabledCategories = Object.entries(config.policy.categories)
+    .filter(([, v]) => !v)
+    .map(([k]) => k);
+
+  const custom =
+    config.policy.custom_rules.length > 0
+      ? 'CUSTOM RULES (must also enforce):\n' +
+        config.policy.custom_rules.map((r: string) => '- ' + r).join('\n')
+      : '';
+
+  const disabledNote =
+    disabledCategories.length > 0
+      ? `DISABLED CATEGORIES (do NOT report findings in these categories): ${disabledCategories.join(', ')}`
+      : 'All listed categories are enabled.';
 
   return `You are an expert application security code reviewer with deep knowledge of secure coding practices.
 
@@ -54,6 +101,7 @@ SECURITY STANDARDS TO ENFORCE (be conservative - prefer reporting over missing i
 - Modern cryptography standards (no MD5/SHA1 for security, no ECB, no hardcoded keys/IVs)
 
 DEFAULT ENABLED RULE CATEGORIES: ${enabledCategories}
+${disabledNote}
 
 ${custom}
 
@@ -65,6 +113,10 @@ STRICT RULES FOR ANALYSIS:
 5. Map to CWE and OWASP where applicable. Use null if none fit.
 6. Provide actionable recommendation with code suggestion when possible.
 7. Assign confidence: high (clear vulnerability), medium, low (possible issue).
+8. Set "category" to one of the enabled category keys (e.g. secrets, injection, xss).
+9. Do not invent findings outside the enabled categories.
+
+${FEW_SHOT}
 
 OUTPUT REQUIREMENTS:
 - Respond with ONLY a single valid JSON object. No markdown, no explanations outside the JSON.
@@ -81,7 +133,8 @@ OUTPUT REQUIREMENTS:
       "cwe": "CWE-XXX or null",
       "owasp": "A01:2021 or null",
       "recommendation": "how to fix it, preferably with example",
-      "confidence": "high|medium|low"
+      "confidence": "high|medium|low",
+      "category": "secrets|injection|authn_authz|cryptography|insecure_deserialization|path_traversal|ssrf|xss|csrf|supply_chain|hardcoded_credentials|dangerous_functions|misconfiguration"
     }
   ],
   "summary": "1-2 sentence overall assessment"
@@ -95,10 +148,11 @@ NEVER include any text before or after the JSON.`;
 export function buildUserPrompt(
   diff: string,
   files: Array<{ filename: string; status?: string }>,
-  config: any
+  config: SecurityScanConfig,
+  truncated = false
 ): string {
   const fileList = files
-    .map(f => {
+    .map((f) => {
       const lang = inferLanguage(f.filename);
       return `- ${f.filename} (${lang})`;
     })
@@ -108,6 +162,10 @@ export function buildUserPrompt(
 min_confidence: ${config.policy.min_confidence}
 categories: ${JSON.stringify(config.policy.categories)}`;
 
+  const truncNote = truncated
+    ? '\n\nNOTE: The unified diff was truncated due to size limits. Only analyze what is present; do not assume missing hunks are safe or unsafe.\n'
+    : '';
+
   return `## Pull Request Changes
 
 Files changed:
@@ -115,7 +173,7 @@ ${fileList}
 
 ## Policy
 ${policySummary}
-
+${truncNote}
 ## Unified Diff (with context)
 \`\`\`diff
 ${diff}

src/types.ts

@@ -1,27 +1,53 @@
+export type Severity = 'critical' | 'high' | 'medium' | 'low' | 'info';
+export type Confidence = 'high' | 'medium' | 'low';
+export type BlockOn = 'critical' | 'high' | 'medium' | 'low' | 'none';
+
+/** Policy category keys (matches DEFAULT_CONFIG.policy.categories). */
+export type FindingCategory =
+  | 'secrets'
+  | 'injection'
+  | 'authn_authz'
+  | 'cryptography'
+  | 'insecure_deserialization'
+  | 'path_traversal'
+  | 'ssrf'
+  | 'xss'
+  | 'csrf'
+  | 'supply_chain'
+  | 'hardcoded_credentials'
+  | 'dangerous_functions'
+  | 'misconfiguration'
+  | string;
+
 export interface Finding {
   file: string;
   start_line: number;
   end_line: number;
-  severity: 'critical' | 'high' | 'medium' | 'low' | 'info';
+  severity: Severity;
   title: string;
   description: string;
   cwe: string | null;
   owasp: string | null;
   recommendation: string;
-  confidence: 'high' | 'medium' | 'low';
+  confidence: Confidence;
+  /** Optional category used for policy.categories filtering. */
+  category?: FindingCategory | null;
 }
 
 export interface LlmResponse {
   findings: Finding[];
   summary: string;
+  /** Set when the model output could not be parsed as JSON. */
+  parse_error?: string;
+  raw?: string;
 }
 
 export interface SecurityScanConfig {
   version: number;
   llm: { provider: string; model: string };
   policy: {
-    block_on: 'critical' | 'high' | 'medium' | 'low' | 'none';
-    min_confidence: 'high' | 'medium' | 'low';
+    block_on: BlockOn;
+    min_confidence: Confidence;
     categories: Record<string, boolean>;
     custom_rules: string[];
   };