From ca86c839ccf6a9b82544c75f4e85be5b24d74ade Mon Sep 17 00:00:00 2001 From: ezl-keygraph Date: Wed, 6 May 2026 04:07:38 +0530 Subject: [PATCH] feat(ai): steer notes field for analysis-only mode (#329) --- apps/worker/src/ai/queue-schemas.ts | 106 +++++++++++++++----- apps/worker/src/services/agent-execution.ts | 2 +- 2 files changed, 81 insertions(+), 27 deletions(-) diff --git a/apps/worker/src/ai/queue-schemas.ts b/apps/worker/src/ai/queue-schemas.ts index 5daf422..311cfdd 100644 --- a/apps/worker/src/ai/queue-schemas.ts +++ b/apps/worker/src/ai/queue-schemas.ts @@ -17,15 +17,27 @@ import type { AgentName } from '../types/agents.js'; // === Common Fields === -const baseVulnerability = z.object({ - ID: z.string(), - vulnerability_type: z.string(), - externally_exploitable: z.boolean(), - confidence: z.string(), - notes: z.string().optional(), -}); +const ANALYSIS_NOTES_DESCRIPTION = + 'Plain context for defenders (caveats, scope, what is at risk). Not attack steps.'; -// === Per-Vuln-Type Schemas === +function notesField(exploit: boolean) { + const f = z.string().optional(); + return exploit ? f : f.describe(ANALYSIS_NOTES_DESCRIPTION); +} + +function makeBase(exploit: boolean) { + return z.object({ + ID: z.string(), + vulnerability_type: z.string(), + externally_exploitable: z.boolean(), + confidence: z.string(), + notes: notesField(exploit), + }); +} + +// === Per-Vuln-Type Schemas (used for type inference; notes description is mode-agnostic for types) === + +const baseVulnerability = makeBase(true); const InjectionVulnerability = baseVulnerability.extend({ source: z.string().optional(), @@ -87,14 +99,6 @@ export type AuthFinding = z.infer; export type SsrfFinding = z.infer; export type AuthzFinding = z.infer; -// === Queue Wrapper Schemas === - -const InjectionQueueSchema = z.object({ vulnerabilities: z.array(InjectionVulnerability) }); -const XssQueueSchema = z.object({ vulnerabilities: z.array(XssVulnerability) }); -const AuthQueueSchema = z.object({ vulnerabilities: z.array(AuthVulnerability) }); -const SsrfQueueSchema = z.object({ vulnerabilities: z.array(SsrfVulnerability) }); -const AuthzQueueSchema = z.object({ vulnerabilities: z.array(AuthzVulnerability) }); - // === Convert to JSON Schema for SDK === // NOTE: The SDK's AJV validator expects draft-07. Zod defaults to draft-2020-12 which @@ -103,15 +107,65 @@ function toOutputFormat(zodSchema: z.ZodType): JsonSchemaOutputFormat { return { type: 'json_schema', schema: z.toJSONSchema(zodSchema, { target: 'draft-07' }) as Record }; } -// === Lookup Maps === +// === Per-Mode Output Format Builders === +// Two maps cached at module load; the only per-mode difference is the +// description on the `notes` field, which steers the LLM's writing. -const VULN_AGENT_OUTPUT_FORMAT: Partial> = { - 'injection-vuln': toOutputFormat(InjectionQueueSchema), - 'xss-vuln': toOutputFormat(XssQueueSchema), - 'auth-vuln': toOutputFormat(AuthQueueSchema), - 'ssrf-vuln': toOutputFormat(SsrfQueueSchema), - 'authz-vuln': toOutputFormat(AuthzQueueSchema), -}; +function buildOutputFormats(exploit: boolean): Partial> { + const base = makeBase(exploit); + return { + 'injection-vuln': toOutputFormat(z.object({ vulnerabilities: z.array(base.extend({ + source: z.string().optional(), + combined_sources: z.string().optional(), + path: z.string().optional(), + sink_call: z.string().optional(), + slot_type: z.string().optional(), + sanitization_observed: z.string().optional(), + concat_occurrences: z.string().optional(), + verdict: z.string().optional(), + mismatch_reason: z.string().optional(), + witness_payload: z.string().optional(), + })) })), + 'xss-vuln': toOutputFormat(z.object({ vulnerabilities: z.array(base.extend({ + source: z.string().optional(), + source_detail: z.string().optional(), + path: z.string().optional(), + sink_function: z.string().optional(), + render_context: z.string().optional(), + encoding_observed: z.string().optional(), + verdict: z.string().optional(), + mismatch_reason: z.string().optional(), + witness_payload: z.string().optional(), + })) })), + 'auth-vuln': toOutputFormat(z.object({ vulnerabilities: z.array(base.extend({ + source_endpoint: z.string().optional(), + vulnerable_code_location: z.string().optional(), + missing_defense: z.string().optional(), + exploitation_hypothesis: z.string().optional(), + suggested_exploit_technique: z.string().optional(), + })) })), + 'ssrf-vuln': toOutputFormat(z.object({ vulnerabilities: z.array(base.extend({ + source_endpoint: z.string().optional(), + vulnerable_parameter: z.string().optional(), + vulnerable_code_location: z.string().optional(), + missing_defense: z.string().optional(), + exploitation_hypothesis: z.string().optional(), + suggested_exploit_technique: z.string().optional(), + })) })), + 'authz-vuln': toOutputFormat(z.object({ vulnerabilities: z.array(base.extend({ + endpoint: z.string().optional(), + vulnerable_code_location: z.string().optional(), + role_context: z.string().optional(), + guard_evidence: z.string().optional(), + side_effect: z.string().optional(), + reason: z.string().optional(), + minimal_witness: z.string().optional(), + })) })), + }; +} + +const OUTPUT_FORMATS_EXPLOIT = buildOutputFormats(true); +const OUTPUT_FORMATS_ANALYSIS = buildOutputFormats(false); const VULN_AGENT_QUEUE_FILENAMES: Partial> = { 'injection-vuln': 'injection_exploitation_queue.json', @@ -122,8 +176,8 @@ const VULN_AGENT_QUEUE_FILENAMES: Partial> = { }; /** Returns the structured output format for a vuln agent, or undefined for non-vuln agents. */ -export function getOutputFormat(agentName: AgentName): JsonSchemaOutputFormat | undefined { - return VULN_AGENT_OUTPUT_FORMAT[agentName]; +export function getOutputFormat(agentName: AgentName, exploit = true): JsonSchemaOutputFormat | undefined { + return (exploit ? OUTPUT_FORMATS_EXPLOIT : OUTPUT_FORMATS_ANALYSIS)[agentName]; } /** Returns the queue filename for a vuln agent, or undefined for non-vuln agents. */ diff --git a/apps/worker/src/services/agent-execution.ts b/apps/worker/src/services/agent-execution.ts index e45d22f..cac686c 100644 --- a/apps/worker/src/services/agent-execution.ts +++ b/apps/worker/src/services/agent-execution.ts @@ -142,7 +142,7 @@ export class AgentExecutionService { await auditSession.startAgent(agentName, prompt, attemptNumber); // 5. Execute agent - const outputFormat = getOutputFormat(agentName); + const outputFormat = getOutputFormat(agentName, distributedConfig?.exploit ?? true); const result: ClaudePromptResult = await runClaudePrompt( prompt, repoPath,