diff --git a/lib/audit/dns-engine.ts b/lib/audit/dns-engine.ts index a064a3c2..2b360876 100644 --- a/lib/audit/dns-engine.ts +++ b/lib/audit/dns-engine.ts @@ -1,9 +1,9 @@ import Anthropic from "@anthropic-ai/sdk"; -import { zodOutputFormat } from "@anthropic-ai/sdk/helpers/zod"; -// zod/v4 — the SDK's zodOutputFormat calls z.toJSONSchema(), a v4-only API. +// zod/v4: backendAi and the SDK both call z.toJSONSchema(), a v4-only API. // See the note in claude-engine.ts. import { z } from "zod/v4"; import { env } from "../env"; +import { generateStructuredOutput, hasBackendAiTextProvider } from "../lx/backendAi"; import { collectDnsRecords, dnsBaselineFindings, type DnsRecords } from "./dns"; import { scoreFindings } from "./score"; import type { ClaudeAuditResult } from "./claude-engine"; @@ -97,26 +97,19 @@ export async function dnsAudit(targetUrl: string): Promise { // AI step — best effort. Fall back to the deterministic baseline on any // failure so this free scan never hard-fails the worker. - if (env.anthropicApiKey) { + if (hasBackendAiTextProvider({ anthropicApiKey: env.anthropicApiKey, openaiApiKey: env.openaiApiKey })) { try { - const client = new Anthropic({ apiKey: env.anthropicApiKey }); - const stream = client.messages.stream({ - model: "claude-sonnet-4-6", - max_tokens: 16000, - thinking: { type: "disabled" }, - output_config: { - effort: "medium", - // eslint-disable-next-line @typescript-eslint/no-explicit-any - format: zodOutputFormat(ResultSchema as any), - }, - system: [ - { type: "text", text: SYSTEM_PROMPT, cache_control: { type: "ephemeral" } }, - ], - messages: [{ role: "user", content: buildUserPrompt(records) }], - }); - const response = await stream.finalMessage(); - const parsed = response.parsed_output; - if (!parsed) throw new Error(`no parsed output (stop_reason=${response.stop_reason})`); + const { output: parsed } = await generateStructuredOutput({ + name: "dns_audit", + // eslint-disable-next-line @typescript-eslint/no-explicit-any + schema: ResultSchema as any, + system: SYSTEM_PROMPT, + user: buildUserPrompt(records), + maxTokens: 16000, + anthropicModel: "claude-sonnet-4-6", + anthropic: env.anthropicApiKey ? new Anthropic({ apiKey: env.anthropicApiKey }) : null, + anthropicCacheSystemPrompt: true, + }) as { output: z.infer }; const aiFindings: Finding[] = parsed.findings.map( (f: z.infer) => ({ diff --git a/tests/dns-engine.test.ts b/tests/dns-engine.test.ts new file mode 100644 index 00000000..5b9e7cb1 --- /dev/null +++ b/tests/dns-engine.test.ts @@ -0,0 +1,50 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const calls = vi.hoisted(() => ({ args: [] as any[], answer: null as null | (() => unknown) })); + +vi.mock("@/lib/lx/backendAi", () => ({ + hasBackendAiTextProvider: () => true, + generateStructuredOutput: async (args: unknown) => { + calls.args.push(args); + return calls.answer!(); + }, +})); +vi.mock("@/lib/env", () => ({ env: { anthropicApiKey: "", openaiApiKey: "" } })); +vi.mock("@/lib/audit/dns", () => ({ + collectDnsRecords: async () => ({ domain: "example.com", a: [], aaaa: [], mx: [], ns: [], txt: [], dkim: [] }), + dnsBaselineFindings: () => [ + { section: "DNS", check_key: "dns.inventory", status: "pass", title: "Inventory", priority: 5 }, + ], +})); + +import { dnsAudit } from "@/lib/audit/dns-engine"; + +describe("dnsAudit AI step", () => { + beforeEach(() => { + calls.args = []; + }); + + it("runs with no paid key and takes c0mpute's answer", async () => { + calls.answer = () => ({ + provider: "c0mpute", + output: { + score: 80, + findings: [{ check_key: "spf", status: "warn", title: "SPF is soft", priority: 2 }], + markdown: "# DNS", + }, + }); + const result = await dnsAudit("https://example.com"); + expect(calls.args).toHaveLength(1); + expect(calls.args[0].anthropic).toBeNull(); + expect(result.score).toBe(80); + expect(result.findings.map((f) => f.check_key)).toEqual(["dns.inventory", "dns.spf"]); + }); + + it("keeps the deterministic baseline when every provider fails", async () => { + calls.answer = () => { + throw new Error("Backend AI generation failed"); + }; + const result = await dnsAudit("https://example.com"); + expect(result.findings.map((f) => f.check_key)).toEqual(["dns.inventory"]); + }); +});