diff --git a/core/src/execute/agentAttackDriver.ts b/core/src/execute/agentAttackDriver.ts new file mode 100644 index 0000000..b77c300 --- /dev/null +++ b/core/src/execute/agentAttackDriver.ts @@ -0,0 +1,225 @@ +// Browser-safe agent AttackDriver — no Node-only imports. +// Fills the AttackDriver holes for agent/chatbot attacks; runAttack owns the loop. + +import type { LanguageModel } from "ai"; +import { generateNextAdaptiveTurn } from "../generate/generateNextTurn.js"; +import type { AttackPattern } from "../evaluators/parseEvaluator.js"; +import { judgeResponse } from "../evaluators/judge.js"; +import { errorJudge } from "../lib/judgeTypes.js"; +import type { JudgeObservabilityContext } from "../evaluators/judge.js"; +import { isTargetError } from "../targets/agentTarget.js"; +import type { AgentTarget } from "../targets/agentTarget.js"; +import { newOtelTraceId } from "../lib/tracePropagation.js"; +import { randomUUID } from "../lib/random.js"; +import { getAdapter } from "../telemetry/adapter.js"; +import { log } from "../lib/logger.js"; +import type { AgentAttackSpec, AttackResult, AgentTurnRecord } from "./types.js"; +import type { TelemetryConfig } from "../config/types.js"; +import type { UnifiedTargetConfig } from "./types.js"; +import { ConversationHistory } from "./conversationHistory.js"; +import type { AttackDriver } from "./attackRunner.js"; + +export interface AgentAttackContext { + targetConfig?: UnifiedTargetConfig; + telemetry?: TelemetryConfig; + /** Prior conversation turns to seed the history (used for resume). */ + initialHistory?: { role: "user" | "assistant"; content: string }[]; +} + +/** + * Drives one agent attack: build a prompt (seed on turn 1, else an adaptive + * follow-up), send it, record the turn, stop on a target error, and judge the + * whole transcript once at the end. + */ +export class AgentAttackDriver implements AttackDriver { + readonly startTurn: number; + readonly totalTurns: number; + + private readonly turns: AgentTurnRecord[] = []; + private readonly history: ConversationHistory; + // Previous turn's technique, threaded into the next turn's user-block as + // PREVIOUS_TECHNIQUE. Undefined on turn 1 and on the first turn of a resumed + // run (no prior technique yet), so the refusal-pivot rule (STEP 5) skips there. + private previousTechnique: string | undefined; + private finalPrompt: string; + private finalResponse = ""; + private readonly propagation: TelemetryConfig["propagation"]; + private readonly attackTraceId: string | undefined; + // One sessionId per attack — every turn reuses it so a stateful agent under + // test can thread its own chat history across turns. + private readonly attackSessionId = randomUUID(); + + constructor( + private readonly attack: AgentAttackSpec, + private readonly attackModel: LanguageModel, + private readonly judgeModel: LanguageModel, + private readonly attackIndex: string, + private readonly patterns: AttackPattern[], + private readonly target: AgentTarget, + private readonly context?: AgentAttackContext + ) { + this.history = new ConversationHistory(context?.initialHistory); + this.finalPrompt = attack.prompt ?? ""; + // For resume: if we already have a completed turn, seed finalPrompt/Response from it. + if (this.history.turnCount >= 1) { + this.finalPrompt = this.history.lastUser(); + this.finalResponse = this.history.lastAssistant(); + } + // Resume: seed turns already in the transcript so the result reports the + // full conversation, not just turns run after resume. + const seeded = this.history.messages; + for (let i = 0; i + 1 < seeded.length; i += 2) { + if (seeded[i].role === "user" && seeded[i + 1].role === "assistant") { + this.turns.push({ + kind: "agent", + turnIndex: this.turns.length + 1, + prompt: seeded[i].content, + response: seeded[i + 1].content, + }); + } + } + + this.propagation = context?.telemetry?.propagation; + const hasPropagation = + Boolean(this.propagation?.headers && Object.keys(this.propagation.headers).length > 0) || + Boolean(this.propagation?.traceIdBodyField?.trim()); + this.attackTraceId = + hasPropagation && (this.propagation?.traceIdStrategy ?? "per-attack") === "per-attack" + ? newOtelTraceId() + : undefined; + + this.startTurn = this.history.turnCount + 1; + this.totalTurns = attack.turns; + } + + async buildTurn(turnNo: number): Promise { + // Comprehensive mode seeds attack.prompt → turn 1 uses it verbatim. Adaptive + // mode sets attack.prompt = "" → generateNextAdaptiveTurn picks the opening. + if (turnNo === 1 && this.attack.prompt) { + return this.finalPrompt; + } + + const result = await generateNextAdaptiveTurn({ + history: this.history.messages, + attack: this.attack, + patterns: this.patterns, + target: this.context?.targetConfig ?? { + kind: "agent", + name: "", + description: "", + type: "http-endpoint", + }, + model: this.attackModel, + currentTurn: turnNo, + maxTurns: this.attack.turns, + attackObjective: this.attack.attackObjective, + businessUseCase: this.attack.businessUseCase, + siteSnapshot: this.attack.siteSnapshot, + maxLength: this.attack.maxMessageLength, + traceContext: this.attack.traceContext, + previousTechnique: this.previousTechnique, + upstreamSessions: this.attack.upstreamSessions, + }); + this.previousTechnique = result.technique; + log.dim( + `[attacker] turn=${turnNo}/${this.attack.turns} technique=${result.technique ?? "unknown"} hook=${ + result.lastReplyHook?.slice(0, 60) ?? "-" + }` + ); + return result.message; + } + + async execute(prompt: string): Promise { + return this.target.send(prompt, { + sessionId: this.attackSessionId, + history: this.history.messages, + propagation: this.propagation, + attackTraceId: this.attackTraceId, + attackIndex: Number.isFinite(Number(this.attackIndex)) ? Number(this.attackIndex) : undefined, + }); + } + + record(turnNo: number, prompt: string, response: string): void { + this.history.push(prompt, response); + this.finalPrompt = prompt; + this.finalResponse = response; + this.turns.push({ kind: "agent", turnIndex: turnNo, prompt, response }); + } + + async shouldEarlyStop(_turnNo: number, _prompt: string, response: string): Promise { + return isTargetError(response); + } + + async finalize(): Promise { + await this.target.close(); + + const finalJudge = + this.turns.length === 0 + ? errorJudge("no turns completed") + : isTargetError(this.finalResponse) + ? errorJudge(this.finalResponse) + : await judgeResponse( + { + id: this.attack.evaluatorId, + name: this.attack.evaluatorName, + severity: this.attack.severity, + standards: this.attack.standards, + description: this.attack.description ?? "", + passCriteria: this.attack.passCriteria, + failCriteria: this.attack.failCriteria, + patterns: [], + }, + this.finalPrompt, + this.finalResponse, + this.judgeModel, + await buildJudgeObservability( + this.context?.telemetry, + this.attackTraceId, + this.finalResponse + ), + this.history.size > 2 ? this.history.messages : undefined, + { patternName: this.attack.patternName, judgeHint: this.attack.judgeHint }, + this.attack.upstreamSessions + ); + + return { + kind: "agent", + attackId: this.attack.id, + evaluatorId: this.attack.evaluatorId, + patternName: this.attack.patternName, + prompt: this.finalPrompt, + response: this.finalResponse, + judge: finalJudge, + turns: this.turns.length > 1 ? this.turns : undefined, + }; + } +} + +async function buildJudgeObservability( + telemetry: TelemetryConfig | undefined, + attackTraceId: string | undefined, + finalResponse: string +): Promise { + if (!telemetry || !attackTraceId) return undefined; + const obs: JudgeObservabilityContext = { propagatedTraceId: attackTraceId }; + const adapter = getAdapter(telemetry.provider); + if (adapter && telemetry.enrichJudgeFromTrace && !isTargetError(finalResponse)) { + log.info(` → fetching ${telemetry.provider} trace for judge...`); + const traceJson = + (await adapter.fetchTraceForJudge(telemetry, attackTraceId, { + // Budget defaults sized for the completeness poll (wait for the final turn + // to ingest): ~1s + 7×1.5s ≈ 11.5s cap before returning best-effort. + initialDelayMs: telemetry.traceFetchInitialDelayMs ?? 1000, + maxAttempts: telemetry.traceFetchMaxAttempts ?? 8, + retryDelayMs: telemetry.traceFetchRetryDelayMs ?? 1500, + maxChars: telemetry.enrichJudgeTraceJsonMaxChars ?? 40_000, + // Completeness signal: the trace is "done" once this final turn's + // response has been ingested, not the instant any span appears. + expectedResponse: finalResponse, + })) ?? undefined; + if (traceJson) obs.traceJson = traceJson; + const ok = traceJson && !traceJson.startsWith("["); + log.info(` → trace ${ok ? "fetched ✓" : "not found ✗"}`); + } + return obs; +} diff --git a/core/src/execute/attackRunner.ts b/core/src/execute/attackRunner.ts new file mode 100644 index 0000000..4d48386 --- /dev/null +++ b/core/src/execute/attackRunner.ts @@ -0,0 +1,45 @@ +import type { AttackResult } from "./types.js"; + +/** + * Per-kind strategy for one attack. Agent and MCP attacks share a control flow — + * loop over turns, build an input, execute it against the target, record the + * result, optionally stop early, then judge once and assemble a result — but + * differ in every step. A driver fills those holes; {@link runAttack} owns the + * order so the two kinds can't drift (the duplication the review flagged). + * + * `TInput` is what a turn sends (an agent prompt, an MCP tool call); `TOutput` is + * what the target returns (a response string, a tool response + error). + */ +export interface AttackDriver { + /** First turn to run (1, or later when resuming a partial transcript). */ + readonly startTurn: number; + /** Last turn to run, inclusive. */ + readonly totalTurns: number; + + /** Build the input for one turn (seed, or an adaptively-generated follow-up). */ + buildTurn(turnNo: number): Promise; + /** Send the input to the target and return its output. */ + execute(input: TInput): Promise; + /** Record the completed turn (append to the transcript / turn list). */ + record(turnNo: number, input: TInput, output: TOutput): void; + /** Whether to stop after this turn — may run a mid-attack judge check. */ + shouldEarlyStop(turnNo: number, input: TInput, output: TOutput): Promise; + /** Judge the completed attack and assemble its result. */ + finalize(): Promise; +} + +/** + * Template Method for running one attack: the invariant skeleton every attack + * kind shares. The `driver` supplies the kind-specific behavior. + */ +export async function runAttack( + driver: AttackDriver +): Promise { + for (let turnNo = driver.startTurn; turnNo <= driver.totalTurns; turnNo++) { + const input = await driver.buildTurn(turnNo); + const output = await driver.execute(input); + driver.record(turnNo, input, output); + if (await driver.shouldEarlyStop(turnNo, input, output)) break; + } + return driver.finalize(); +} diff --git a/core/src/execute/mcpAttackDriver.ts b/core/src/execute/mcpAttackDriver.ts new file mode 100644 index 0000000..f20370b --- /dev/null +++ b/core/src/execute/mcpAttackDriver.ts @@ -0,0 +1,183 @@ +// MCP AttackDriver — Node-only (MCP transport). Fills the AttackDriver holes for +// MCP tool-call attacks; runAttack owns the loop. Extracted from runAll.ts. + +import { generateNextMcpTurn } from "../generate/generateNextTurn.js"; +import type { McpToolTurn } from "../generate/generateNextTurn.js"; +import { judgeToolResponse, sanitizeJudgeResult } from "../run/judge.js"; +import { errorJudge as mcpErrorJudge, type JudgeResult } from "../lib/judgeTypes.js"; +import { log } from "../lib/logger.js"; +import type { LanguageModel } from "ai"; +import type { LlmConfig } from "../config/types.js"; +import type { McpTarget, McpToolCallResult } from "../targets/mcpTarget.js"; +import type { McpAttackSpec, McpTurnRecord, AttackResult } from "./types.js"; +import { runAttack, type AttackDriver } from "./attackRunner.js"; + +/** + * Drives one MCP attack: call a tool (seed args on turn 1, else adaptively + * generated), record the tool turn, and run a per-turn judge that early-stops on + * the first FAIL. Judges the last turn once at the end otherwise. + * + * Dependencies are injected narrowly (the attacker model built once per run, the + * judge's LLM config, the resolved non-empty tool name) rather than the whole + * RunConfig — same construction contract style as AgentAttackDriver. + */ +export class McpAttackDriver implements AttackDriver, McpToolCallResult> { + readonly startTurn = 1; + readonly totalTurns: number; + + private readonly turns: McpTurnRecord[] = []; + private readonly mcpHistory: McpToolTurn[] = []; + private judgeHint: string | undefined; + private earlyStopJudge: JudgeResult | null = null; + + constructor( + private readonly attack: McpAttackSpec, + private readonly target: McpTarget, + private readonly toolName: string, + private readonly attackModel: LanguageModel, + private readonly judgeLlm: LlmConfig + ) { + this.judgeHint = attack.judgeHint; + this.totalTurns = attack.turns; + } + + async buildTurn(turnNo: number): Promise> { + if (turnNo === 1) { + return this.attack.toolArguments ?? {}; + } + const next = await generateNextMcpTurn( + this.mcpHistory, + `${this.attack.patternName} — ${this.attack.evaluatorName}`, + this.toolName, + this.attack.toolArguments ?? {}, + this.attackModel + ); + if (next.judgeHint) this.judgeHint = next.judgeHint; + return next.args; + } + + async execute(toolArguments: Record): Promise { + return this.target.callTool(this.toolName, toolArguments); + } + + record(turnNo: number, toolArguments: Record, output: McpToolCallResult): void { + const { response, toolError } = output; + this.mcpHistory.push({ toolName: this.toolName, toolArguments, response, toolError }); + this.turns.push({ + kind: "mcp", + turnIndex: turnNo, + toolName: this.toolName, + toolArguments, + response, + toolError, + }); + } + + async shouldEarlyStop( + turnNo: number, + toolArguments: Record, + output: McpToolCallResult + ): Promise { + // Multi-turn: run a per-turn judge check and early-stop on FAIL. + if (this.attack.turns > 1 && turnNo < this.attack.turns) { + const midJudge = await this.judge(toolArguments, output); + if (midJudge.verdict === "FAIL") { + log.info(` ⚡ Early stop at turn ${turnNo}/${this.attack.turns} — vulnerability found`); + this.earlyStopJudge = midJudge; + return true; + } + } + return false; + } + + async finalize(): Promise { + const lastTurn = this.turns[this.turns.length - 1]; + + let finalJudge: JudgeResult; + if (this.earlyStopJudge) { + finalJudge = this.earlyStopJudge; + } else if (!lastTurn) { + finalJudge = mcpErrorJudge("no turns completed"); + } else { + finalJudge = await this.judge(lastTurn.toolArguments, { + response: lastTurn.response, + toolError: lastTurn.toolError, + }); + } + + return { + kind: "mcp", + attackId: this.attack.id, + evaluatorId: this.attack.evaluatorId, + patternName: this.attack.patternName, + toolName: this.toolName, + toolArguments: lastTurn?.toolArguments ?? this.attack.toolArguments, + toolResponse: lastTurn?.response, + toolError: lastTurn?.toolError, + judge: finalJudge, + turns: this.turns.length > 1 ? this.turns : undefined, + }; + } + + /** + * Judge one tool turn and sanitize hallucinated evidence. Shared by the mid-turn + * early-stop check and the final judgement (previously duplicated verbatim). + * `priorTurns` excludes the turn under judgement (the last recorded turn). + */ + private async judge( + toolArguments: Record, + output: McpToolCallResult + ): Promise { + const { response, toolError } = output; + const result = await judgeToolResponse({ + model: this.judgeLlm, + evaluator: { + id: this.attack.evaluatorId, + name: this.attack.evaluatorName, + standards: this.attack.standards, + severity: this.attack.severity, + passCriteria: this.attack.passCriteria, + failCriteria: this.attack.failCriteria, + }, + attackSummary: this.attack.patternName, + toolName: this.toolName, + toolArguments, + toolResponse: response, + toolError, + judgeHint: this.judgeHint, + priorTurns: this.mcpHistory.length > 1 ? this.mcpHistory.slice(0, -1) : undefined, + }); + return sanitizeJudgeResult(result, { + attackSummary: this.attack.patternName, + toolArguments, + toolResponse: response, + toolError, + }); + } +} + +/** + * Run one MCP attack. Short-circuits to an ERROR result when the spec carries no + * tool name (nothing to call); otherwise drives the shared attack loop. + */ +export async function runMcpAttack( + attack: McpAttackSpec, + target: McpTarget, + attackModel: LanguageModel, + judgeLlm: LlmConfig +): Promise { + if (!attack.toolName) { + return { + kind: "mcp", + attackId: attack.id, + evaluatorId: attack.evaluatorId, + patternName: attack.patternName, + toolName: "", + toolArguments: {}, + toolResponse: "", + toolError: "no toolName in attack spec", + judge: mcpErrorJudge("no toolName in attack spec"), + }; + } + return runAttack(new McpAttackDriver(attack, target, attack.toolName, attackModel, judgeLlm)); +} diff --git a/core/src/execute/runAgentLoop.ts b/core/src/execute/runAgentLoop.ts index 7b3b53f..2f552b0 100644 --- a/core/src/execute/runAgentLoop.ts +++ b/core/src/execute/runAgentLoop.ts @@ -1,23 +1,16 @@ // Browser-safe agent attack runner — no Node-only imports. // Used by both runAll.ts (Node) and runAllBrowser.ts (browser/extension). // The caller is responsible for creating and passing the AgentTarget. +// +// Thin wrapper: the loop is the shared runAttack Template Method, the agent +// behavior is AgentAttackDriver. import type { LanguageModel } from "ai"; -import { generateNextAdaptiveTurn } from "../generate/generateNextTurn.js"; import type { AttackPattern } from "../evaluators/parseEvaluator.js"; -import { judgeResponse } from "../evaluators/judge.js"; -import { errorJudge } from "../lib/judgeTypes.js"; -import type { JudgeObservabilityContext } from "../evaluators/judge.js"; -import { isTargetError } from "../targets/agentTarget.js"; import type { AgentTarget } from "../targets/agentTarget.js"; -import { newOtelTraceId } from "../lib/tracePropagation.js"; -import { randomUUID } from "../lib/random.js"; -import { getAdapter } from "../telemetry/adapter.js"; -import { log } from "../lib/logger.js"; -import type { AgentAttackSpec, AttackResult, AgentTurnRecord } from "./types.js"; -import type { TelemetryConfig } from "../config/types.js"; -import type { UnifiedTargetConfig } from "./types.js"; -import { ConversationHistory } from "./conversationHistory.js"; +import type { AgentAttackSpec, AttackResult } from "./types.js"; +import { runAttack } from "./attackRunner.js"; +import { AgentAttackDriver, type AgentAttackContext } from "./agentAttackDriver.js"; export async function runAgentAttack( attack: AgentAttackSpec, @@ -26,171 +19,9 @@ export async function runAgentAttack( attackIndex: string, patterns: AttackPattern[], target: AgentTarget, - context?: { - targetConfig?: UnifiedTargetConfig; - telemetry?: TelemetryConfig; - /** Prior conversation turns to seed the history (used for resume). */ - initialHistory?: { role: "user" | "assistant"; content: string }[]; - } + context?: AgentAttackContext ): Promise { - const turns: AgentTurnRecord[] = []; - const history = new ConversationHistory(context?.initialHistory); - // Parallel meta channel for attacker tag output (not sent to target). - // Used to thread PREVIOUS_TECHNIQUE into the next turn's user-block. - // Resume limitation: on resumed runs (initialHistory non-empty), this - // array starts empty even when startTurn > 1, so the first resumed turn - // has no PREVIOUS_TECHNIQUE and the refusal-pivot rule (STEP 5) skips. - // Subsequent turns within the resumed session populate and work normally. - // To eliminate: persist meta alongside history in the resume context - // (Phase 2 alongside AgentTurnRecord.technique field). - const attackerMeta: Array<{ technique?: string; lastReplyHook?: string }> = []; - let finalPrompt = attack.prompt ?? ""; - let finalResponse = ""; - - // For resume: if we already have a completed turn, seed finalPrompt/Response from it - if (history.turnCount >= 1) { - finalPrompt = history.lastUser(); - finalResponse = history.lastAssistant(); - } - - const propagation = context?.telemetry?.propagation; - const hasPropagation = - Boolean(propagation?.headers && Object.keys(propagation.headers).length > 0) || - Boolean(propagation?.traceIdBodyField?.trim()); - const attackTraceId = - hasPropagation && (propagation?.traceIdStrategy ?? "per-attack") === "per-attack" - ? newOtelTraceId() - : undefined; - - // One sessionId per attack — every turn of this attack reuses it so a - // stateful agent under test can thread its own chat history across turns. - // Agents that don't read sessionIdField just ignore it. - const attackSessionId = randomUUID(); - - const startTurn = history.turnCount + 1; - for (let t = startTurn; t <= attack.turns; t++) { - let prompt: string; - // Mode discriminator. Comprehensive mode seeds `attack.prompt` from the - // pattern template via generateAttacks → turn 1 uses that seed verbatim. - // Adaptive mode sets `attack.prompt = ""` (sentinel) → falls through to - // generateNextAdaptiveTurn for turn 1, letting the OPENING-variant prompt - // pick the opening angle dynamically. - if (t === 1 && attack.prompt) { - prompt = finalPrompt; - // Push empty meta so attackerMeta indices stay aligned with turn numbers. - attackerMeta.push({}); - } else { - const result = await generateNextAdaptiveTurn({ - history: history.messages, - attack, - patterns, - target: context?.targetConfig ?? { - kind: "agent", - name: "", - description: "", - type: "http-endpoint", - }, - model: attackModel, - currentTurn: t, - maxTurns: attack.turns, - attackObjective: attack.attackObjective, - businessUseCase: attack.businessUseCase, - siteSnapshot: attack.siteSnapshot, - maxLength: attack.maxMessageLength, - traceContext: attack.traceContext, - previousTechnique: attackerMeta[attackerMeta.length - 1]?.technique, - upstreamSessions: attack.upstreamSessions, - }); - prompt = result.message; - attackerMeta.push({ technique: result.technique, lastReplyHook: result.lastReplyHook }); - log.dim( - `[attacker] turn=${t}/${attack.turns} technique=${result.technique ?? "unknown"} hook=${ - result.lastReplyHook?.slice(0, 60) ?? "-" - }` - ); - } - - const response = await target.send(prompt, { - sessionId: attackSessionId, - history: history.messages, - propagation, - attackTraceId, - attackIndex: Number.isFinite(Number(attackIndex)) ? Number(attackIndex) : undefined, - }); - - history.push(prompt, response); - - finalPrompt = prompt; - finalResponse = response; - turns.push({ kind: "agent", turnIndex: t, prompt, response }); - - if (isTargetError(response)) break; - } - - await target.close(); - - const finalJudge = - turns.length === 0 - ? errorJudge("no turns completed") - : isTargetError(finalResponse) - ? errorJudge(finalResponse) - : await judgeResponse( - { - id: attack.evaluatorId, - name: attack.evaluatorName, - severity: attack.severity, - standards: attack.standards, - description: attack.description ?? "", - passCriteria: attack.passCriteria, - failCriteria: attack.failCriteria, - patterns: [], - }, - finalPrompt, - finalResponse, - judgeModel, - await buildJudgeObservability(context?.telemetry, attackTraceId, finalResponse), - history.size > 2 ? history.messages : undefined, - { patternName: attack.patternName, judgeHint: attack.judgeHint }, - attack.upstreamSessions - ); - - return { - kind: "agent", - attackId: attack.id, - evaluatorId: attack.evaluatorId, - patternName: attack.patternName, - prompt: finalPrompt, - response: finalResponse, - judge: finalJudge, - turns: turns.length > 1 ? turns : undefined, - }; -} - -async function buildJudgeObservability( - telemetry: TelemetryConfig | undefined, - attackTraceId: string | undefined, - finalResponse: string -): Promise { - if (!telemetry || !attackTraceId) return undefined; - const obs: JudgeObservabilityContext = { propagatedTraceId: attackTraceId }; - const adapter = getAdapter(telemetry.provider); - if (adapter && telemetry.enrichJudgeFromTrace && !isTargetError(finalResponse)) { - log.info(` → fetching ${telemetry.provider} trace for judge...`); - const traceJson = - (await adapter.fetchTraceForJudge(telemetry, attackTraceId, { - // Budget defaults sized for the completeness poll (wait for the final turn - // to ingest): ~1s + 7×1.5s ≈ 11.5s cap before returning best-effort. - initialDelayMs: telemetry.traceFetchInitialDelayMs ?? 1000, - maxAttempts: telemetry.traceFetchMaxAttempts ?? 8, - retryDelayMs: telemetry.traceFetchRetryDelayMs ?? 1500, - maxChars: telemetry.enrichJudgeTraceJsonMaxChars ?? 40_000, - // Completeness signal: the trace is "done" once this final turn's - // response has been ingested, not the instant any span appears. - expectedResponse: finalResponse, - })) ?? undefined; - if (traceJson) obs.traceJson = traceJson; - const ok = traceJson && !traceJson.startsWith("["); - log.info(` → trace ${ok ? "fetched ✓" : "not found ✗"}`); - } - return obs; + return runAttack( + new AgentAttackDriver(attack, attackModel, judgeModel, attackIndex, patterns, target, context) + ); } diff --git a/core/src/execute/runAll.ts b/core/src/execute/runAll.ts index 562438f..7a901ce 100644 --- a/core/src/execute/runAll.ts +++ b/core/src/execute/runAll.ts @@ -6,15 +6,12 @@ import type { LanguageModel } from "ai"; import type { RunConfig, AttackSpec, - McpAttackSpec, AttackResult, EvaluatorResult, UnifiedRunReport, - McpTurnRecord, SessionContext, } from "./types.js"; import { generateAttacks, type ToolInfo } from "../generate/generateAttacks.js"; -import { generateNextMcpTurn } from "../generate/generateNextTurn.js"; import { createAgentTarget, TargetStopError } from "../targets/agentTarget.js"; import type { AgentTarget } from "../targets/agentTarget.js"; import { createMcpTarget } from "../targets/mcpTarget.js"; @@ -24,13 +21,14 @@ import type { EvaluatorSpec } from "../evaluators/parseEvaluator.js"; import { loadSkillCatalog, resolveSuiteEvaluatorIds } from "../config/loadSkillCatalog.js"; import { loadCatalog } from "../catalog/loadCatalog.js"; import { runAgentAttack } from "./runAgentLoop.js"; +import { runMcpAttack } from "./mcpAttackDriver.js"; import { summarizeVerdicts, toEvaluatorResult, buildUnifiedReport, modelLabel, } from "./aggregate.js"; -import { judgeToolResponse, sanitizeJudgeResult } from "../run/judge.js"; +import { judgeToolResponse } from "../run/judge.js"; import { errorJudge as mcpErrorJudge } from "../lib/judgeTypes.js"; // scanResources is available for direct client usage; scan-mode evaluators // use target.listResources()/readResource() through the McpTarget interface. @@ -68,7 +66,11 @@ export async function runAll( const notify = options?.onProgress ?? (() => {}); const attackModel = resolveModel(config.attackerLlm); - const judgeModel = resolveModel(config.judgeLlm ?? config.attackerLlm); + // Single source of truth for the judge LLM: explicit judge model, else the + // attacker model. Reused by the baseline scans and the MCP dispatch below so + // the fallback can't drift between paths. + const judgeLlmConfig = config.judgeLlm ?? config.attackerLlm; + const judgeModel = resolveModel(judgeLlmConfig); const isMcp = config.target.kind === "mcp"; const evaluators = await resolveEvaluators( @@ -99,7 +101,7 @@ export async function runAll( try { // ── MCP pre-flight scans (always run for MCP targets) ────────────── if (isMcp && mcpTarget) { - const judgeModelConfig = config.judgeLlm ?? config.attackerLlm; + const judgeModelConfig = judgeLlmConfig; log.info(`\n── MCP baseline scans ──`); const resourceResults = await scanResources({ target: mcpTarget, judgeModelConfig, notify }); @@ -210,7 +212,7 @@ export async function runAll( try { result = attack.kind === "mcp" - ? await runMcpAttack(attack, mcpTarget!, judgeModel, config) + ? await runMcpAttack(attack, mcpTarget!, attackModel, judgeLlmConfig) : await runAgentAttack( attack, attackModel, @@ -335,145 +337,6 @@ export async function runAll( return report; } -// --------------------------------------------------------------------------- -// MCP attack runner -// --------------------------------------------------------------------------- - -async function runMcpAttack( - attack: McpAttackSpec, - target: Awaited>, - judgeModel: LanguageModel, - config: RunConfig -): Promise { - const attackModel = resolveModel(config.attackerLlm); - const toolName = attack.toolName ?? ""; - const turns: McpTurnRecord[] = []; - const mcpHistory: import("../generate/generateNextTurn.js").McpToolTurn[] = []; - - if (!toolName) { - return { - kind: "mcp", - attackId: attack.id, - evaluatorId: attack.evaluatorId, - patternName: attack.patternName, - toolName: "", - toolArguments: {}, - toolResponse: "", - toolError: "no toolName in attack spec", - judge: mcpErrorJudge("no toolName in attack spec"), - }; - } - - let judgeHint = attack.judgeHint; - let earlyStopJudge: import("../run/types.js").JudgeResult | null = null; - - for (let t = 1; t <= attack.turns; t++) { - let toolArguments: Record; - - if (t === 1) { - toolArguments = attack.toolArguments ?? {}; - } else { - const next = await generateNextMcpTurn( - mcpHistory, - `${attack.patternName} — ${attack.evaluatorName}`, - toolName, - attack.toolArguments ?? {}, - attackModel - ); - toolArguments = next.args; - if (next.judgeHint) judgeHint = next.judgeHint; - } - - const { response, toolError } = await target.callTool(toolName, toolArguments); - - mcpHistory.push({ toolName, toolArguments, response, toolError }); - turns.push({ kind: "mcp", turnIndex: t, toolName, toolArguments, response, toolError }); - - // Multi-turn: run a per-turn judge check and early-stop on FAIL - if (attack.turns > 1 && t < attack.turns) { - const midJudge = await judgeToolResponse({ - model: config.judgeLlm ?? config.attackerLlm, - evaluator: { - id: attack.evaluatorId, - name: attack.evaluatorName, - standards: attack.standards, - severity: attack.severity, - passCriteria: attack.passCriteria, - failCriteria: attack.failCriteria, - }, - attackSummary: attack.patternName, - toolName, - toolArguments, - toolResponse: response, - toolError, - judgeHint, - priorTurns: mcpHistory.length > 1 ? mcpHistory.slice(0, -1) : undefined, - }).then((r) => - sanitizeJudgeResult(r, { - attackSummary: attack.patternName, - toolArguments, - toolResponse: response, - toolError, - }) - ); - - if (midJudge.verdict === "FAIL") { - log.info(` ⚡ Early stop at turn ${t}/${attack.turns} — vulnerability found`); - earlyStopJudge = midJudge; - break; - } - } - } - - const lastTurn = turns[turns.length - 1]; - - let finalJudge: import("../run/types.js").JudgeResult; - if (earlyStopJudge) { - finalJudge = earlyStopJudge; - } else if (!lastTurn) { - finalJudge = mcpErrorJudge("no turns completed"); - } else { - finalJudge = await judgeToolResponse({ - model: config.judgeLlm ?? config.attackerLlm, - evaluator: { - id: attack.evaluatorId, - name: attack.evaluatorName, - standards: attack.standards, - severity: attack.severity, - passCriteria: attack.passCriteria, - failCriteria: attack.failCriteria, - }, - attackSummary: attack.patternName, - toolName: lastTurn.toolName, - toolArguments: lastTurn.toolArguments, - toolResponse: lastTurn.response, - toolError: lastTurn.toolError, - judgeHint, - priorTurns: mcpHistory.length > 1 ? mcpHistory.slice(0, -1) : undefined, - }).then((r) => - sanitizeJudgeResult(r, { - attackSummary: attack.patternName, - toolArguments: lastTurn.toolArguments, - toolResponse: lastTurn.response, - toolError: lastTurn.toolError, - }) - ); - } - - return { - kind: "mcp", - attackId: attack.id, - evaluatorId: attack.evaluatorId, - patternName: attack.patternName, - toolName, - toolArguments: lastTurn?.toolArguments ?? attack.toolArguments, - toolResponse: lastTurn?.response, - toolError: lastTurn?.toolError, - judge: finalJudge, - turns: turns.length > 1 ? turns : undefined, - }; -} - // --------------------------------------------------------------------------- // Built-in MCP scans (always run for MCP targets) // --------------------------------------------------------------------------- diff --git a/core/tests/agentAttackDriver.test.ts b/core/tests/agentAttackDriver.test.ts new file mode 100644 index 0000000..a266d23 --- /dev/null +++ b/core/tests/agentAttackDriver.test.ts @@ -0,0 +1,135 @@ +/** + * AgentAttackDriver resume characterization. + * + * The extension threads a prior transcript back in via `initialHistory` to + * resume a paused run. These pin the seeding behavior: turns already in the + * resumed transcript are reported in the final result, and a run resumed at its + * last turn judges that transcript instead of reporting "no turns completed". + * + * Both cases run zero *new* turns (startTurn > totalTurns), so finalize() is the + * only work: a fake target + a local judge server exercise it without an + * attacker LLM or a real agent target. + */ +import { test, before, after } from "node:test"; +import assert from "node:assert/strict"; +import { createServer, type Server } from "node:http"; +import { setEnvProvider } from "../src/lib/env.js"; + +setEnvProvider(() => "fake-test-api-key"); + +interface ServerState { + server: Server; + port: number; +} + +let srv: ServerState; + +// Agent judge output is line-based (Verdict/Score/Confidence/Evidence/Reasoning). +const JUDGE_PASS = "Verdict: PASS\nScore: 10\nConfidence: 90\nEvidence: N/A\nReasoning: safe."; + +function chatCompletion(content: string): string { + return JSON.stringify({ + id: "t", + object: "chat.completion", + created: 0, + model: "m", + choices: [{ index: 0, message: { role: "assistant", content }, finish_reason: "stop" }], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }); +} + +before(async () => { + srv = await new Promise((resolve) => { + const server = createServer((req, res) => { + let body = ""; + req.on("data", (c: Buffer) => (body += c.toString())); + req.on("end", () => { + if ((req.url ?? "").startsWith("/v1/chat/completions")) { + res.writeHead(200, { "Content-Type": "application/json" }); + res.end(chatCompletion(JUDGE_PASS)); + return; + } + res.writeHead(404); + res.end("no"); + }); + }); + server.listen(0, "127.0.0.1", () => { + const { port } = server.address() as { port: number }; + resolve({ server, port }); + }); + }); +}); + +after(async () => { + await new Promise((resolve, reject) => + srv.server.close((e) => (e ? reject(e) : resolve())) + ); +}); + +const { runAgentAttack } = await import("../src/execute/runAgentLoop.js"); +const { createModel } = await import("../src/providers/factory.js"); +type FakeTarget = Parameters[5]; + +/** Target whose send() must never be called (these resumes run zero new turns). */ +function fakeTarget(): FakeTarget { + return { + async send() { + throw new Error("send() should not run when resuming at/after the last turn"); + }, + async close() {}, + }; +} + +function judgeModel() { + return createModel({ + provider: "openai-compatible" as const, + model: "m", + apiKeyEnv: "K", + baseURL: `http://127.0.0.1:${srv.port}/v1`, + }); +} + +function agentAttack(turns: number) { + return { + kind: "agent" as const, + id: "att-1", + evaluatorId: "e1", + evaluatorName: "Eval", + severity: "high", + patternName: "pattern", + passCriteria: "defended", + failCriteria: "leaked", + prompt: "seed prompt", + turns, + }; +} + +test("resume at the last turn judges the transcript instead of reporting no turns", async () => { + const model = judgeModel(); + const result = await runAgentAttack(agentAttack(1), model, model, "0", [], fakeTarget(), { + initialHistory: [ + { role: "user", content: "hi" }, + { role: "assistant", content: "I can't help with that." }, + ], + }); + // Before the fix this errored with "no turns completed" (turns[] was empty on + // resume); now the seeded turn is judged. + assert.strictEqual(result.judge.verdict, "PASS"); +}); + +test("resume reports the full pre-resume transcript in turns[]", async () => { + const model = judgeModel(); + const result = await runAgentAttack(agentAttack(2), model, model, "0", [], fakeTarget(), { + initialHistory: [ + { role: "user", content: "turn 1 q" }, + { role: "assistant", content: "turn 1 a" }, + { role: "user", content: "turn 2 q" }, + { role: "assistant", content: "turn 2 a" }, + ], + }); + assert.strictEqual(result.kind, "agent"); + // Both completed turns are seeded and reported, not dropped. + assert.strictEqual(result.turns?.length, 2); + assert.strictEqual(result.turns?.[0]?.turnIndex, 1); + assert.strictEqual(result.turns?.[1]?.turnIndex, 2); +}); diff --git a/core/tests/attackRunner.test.ts b/core/tests/attackRunner.test.ts new file mode 100644 index 0000000..1c37e08 --- /dev/null +++ b/core/tests/attackRunner.test.ts @@ -0,0 +1,103 @@ +/** + * PR8 — AttackRunner (Template Method) contract. + * + * Pins the invariant skeleton every attack kind shares: per turn build → execute + * → record → shouldEarlyStop, in that order; early-stop breaks the loop; finalize + * runs once at the end and its result is returned; turns run from startTurn to + * totalTurns inclusive. + */ +import { test } from "node:test"; +import assert from "node:assert/strict"; +import { runAttack, type AttackDriver } from "../src/execute/attackRunner.js"; +import type { AttackResult } from "../src/execute/types.js"; + +const RESULT = { kind: "agent", attackId: "a", evaluatorId: "e", patternName: "p" } as AttackResult; + +/** Records the exact call order and turn numbers for assertions. */ +function trackingDriver(opts: { + startTurn: number; + totalTurns: number; + stopAfter?: number; +}): AttackDriver & { calls: string[] } { + const calls: string[] = []; + return { + calls, + startTurn: opts.startTurn, + totalTurns: opts.totalTurns, + async buildTurn(t) { + calls.push(`build:${t}`); + return `in:${t}`; + }, + async execute(input) { + calls.push(`execute:${input}`); + return `out:${input}`; + }, + record(t, input, output) { + calls.push(`record:${t}:${input}:${output}`); + }, + async shouldEarlyStop(t) { + calls.push(`stop?:${t}`); + return opts.stopAfter === t; + }, + async finalize() { + calls.push("finalize"); + return RESULT; + }, + }; +} + +test("runs each turn build→execute→record→shouldEarlyStop, then finalize once", async () => { + const driver = trackingDriver({ startTurn: 1, totalTurns: 2 }); + const result = await runAttack(driver); + assert.strictEqual(result, RESULT); + assert.deepStrictEqual(driver.calls, [ + "build:1", + "execute:in:1", + "record:1:in:1:out:in:1", + "stop?:1", + "build:2", + "execute:in:2", + "record:2:in:2:out:in:2", + "stop?:2", + "finalize", + ]); +}); + +test("early-stop breaks the loop but still finalizes", async () => { + const driver = trackingDriver({ startTurn: 1, totalTurns: 5, stopAfter: 2 }); + await runAttack(driver); + // Turn 3+ never build; finalize still runs. + assert.deepStrictEqual(driver.calls, [ + "build:1", + "execute:in:1", + "record:1:in:1:out:in:1", + "stop?:1", + "build:2", + "execute:in:2", + "record:2:in:2:out:in:2", + "stop?:2", + "finalize", + ]); +}); + +test("honors startTurn when resuming a partial transcript", async () => { + const driver = trackingDriver({ startTurn: 3, totalTurns: 4 }); + await runAttack(driver); + assert.deepStrictEqual(driver.calls, [ + "build:3", + "execute:in:3", + "record:3:in:3:out:in:3", + "stop?:3", + "build:4", + "execute:in:4", + "record:4:in:4:out:in:4", + "stop?:4", + "finalize", + ]); +}); + +test("runs zero turns when startTurn exceeds totalTurns, still finalizes", async () => { + const driver = trackingDriver({ startTurn: 3, totalTurns: 2 }); + await runAttack(driver); + assert.deepStrictEqual(driver.calls, ["finalize"]); +}); diff --git a/core/tests/mcpAttackDriver.test.ts b/core/tests/mcpAttackDriver.test.ts new file mode 100644 index 0000000..d6e1d81 --- /dev/null +++ b/core/tests/mcpAttackDriver.test.ts @@ -0,0 +1,239 @@ +/** + * PR8 — McpAttackDriver characterization. + * + * The MCP attack loop had no test coverage before it was extracted from runAll + * into an AttackDriver. These pin the behaviors that must be preserved: the + * empty-tool short-circuit, a single-turn call → judge → result, the multi-turn + * per-turn judge that early-stops on the first FAIL, and the continuation path + * where a PASS turn runs the turn generator and the loop judges the last turn. + * + * A fake target plus a local server exercise the loop without a real MCP server + * or attacker LLM: the server backs both the judge and the attacker turn + * generator on one endpoint, routing by a marker unique to the generator prompt. + */ +import { test, before, after } from "node:test"; +import assert from "node:assert/strict"; +import { createServer, type Server } from "node:http"; +import { setEnvProvider } from "../src/lib/env.js"; + +setEnvProvider(() => "fake-test-api-key"); + +interface ServerState { + server: Server; + port: number; + llmCalls: number; + /** Requests routed to the MCP turn generator (buildTurn's turn>1 branch). */ + turnGenCalls: number; + nextVerdict: "PASS" | "FAIL"; + reset(verdict: "PASS" | "FAIL"): void; +} + +let srv: ServerState; + +/** Wrap a model's completion text in an OpenAI chat.completion envelope. */ +function chatCompletion(content: string): string { + return JSON.stringify({ + id: "t", + object: "chat.completion", + created: 0, + model: "m", + choices: [{ index: 0, message: { role: "assistant", content }, finish_reason: "stop" }], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }); +} + +function judgeJson(verdict: "PASS" | "FAIL"): string { + const body = + verdict === "FAIL" + ? { verdict: "FAIL", score: 2, confidence: 90, evidence: "AKIA", reasoning: "leaked" } + : { verdict: "PASS", score: 10, confidence: 90, evidence: "N/A", reasoning: "safe" }; + return chatCompletion(JSON.stringify(body)); +} + +/** The attacker's next-turn generator expects a JSON object with args + judgeHint. */ +function turnGenJson(): string { + return chatCompletion( + JSON.stringify({ args: { q: "escalated" }, judgeHint: "probe the boundary" }) + ); +} + +// Unique marker in generateNextMcpTurn's prompt — lets the fake server tell a +// turn-generation request apart from a judge request on the shared endpoint. +const TURN_GEN_MARKER = "escalate the attack"; + +before(async () => { + srv = await new Promise((resolve) => { + let llmCalls = 0; + let turnGenCalls = 0; + let nextVerdict: "PASS" | "FAIL" = "PASS"; + const server = createServer((req, res) => { + let body = ""; + req.on("data", (c: Buffer) => (body += c.toString())); + req.on("end", () => { + if ((req.url ?? "").startsWith("/v1/chat/completions")) { + llmCalls++; + res.writeHead(200, { "Content-Type": "application/json" }); + // One endpoint backs both the attacker turn generator and the judge; + // route by the attacker prompt's unique marker. + if (body.includes(TURN_GEN_MARKER)) { + turnGenCalls++; + res.end(turnGenJson()); + } else { + res.end(judgeJson(nextVerdict)); + } + return; + } + res.writeHead(404); + res.end("no"); + }); + }); + server.listen(0, "127.0.0.1", () => { + const { port } = server.address() as { port: number }; + resolve({ + server, + port, + get llmCalls() { + return llmCalls; + }, + get turnGenCalls() { + return turnGenCalls; + }, + get nextVerdict() { + return nextVerdict; + }, + set nextVerdict(v) { + nextVerdict = v; + }, + reset(v) { + llmCalls = 0; + turnGenCalls = 0; + nextVerdict = v; + }, + }); + }); + }); +}); + +after(async () => { + await new Promise((resolve, reject) => + srv.server.close((e) => (e ? reject(e) : resolve())) + ); +}); + +const { runMcpAttack } = await import("../src/execute/mcpAttackDriver.js"); +const { createModel } = await import("../src/providers/factory.js"); +type FakeTarget = Parameters[1]; + +function fakeTarget(toolResponse: string): FakeTarget { + return { + async callTool() { + return { response: toolResponse }; + }, + async listTools() { + return []; + }, + async listResources() { + return []; + }, + async readResource() { + return ""; + }, + async close() {}, + }; +} + +/** Narrow driver deps: [attackModel, judgeLlm] both pointed at the local server. */ +function deps() { + const llm = { + provider: "openai-compatible" as const, + model: "m", + apiKeyEnv: "K", + baseURL: `http://127.0.0.1:${srv.port}/v1`, + }; + return [createModel(llm), llm] as const; +} + +function mcpAttack(overrides: Record = {}) { + return { + kind: "mcp" as const, + id: "att-1", + evaluatorId: "e1", + evaluatorName: "Eval", + severity: "high", + patternName: "pattern", + passCriteria: "defended", + failCriteria: "leaked", + turns: 1, + toolName: "lookup", + toolArguments: { q: "x" }, + ...overrides, + }; +} + +test("empty toolName short-circuits to an ERROR result without any LLM call", async () => { + srv.reset("PASS"); + const result = await runMcpAttack(mcpAttack({ toolName: "" }), fakeTarget("unused"), ...deps()); + assert.strictEqual(result.kind, "mcp"); + assert.strictEqual(result.judge.verdict, "ERROR"); + assert.strictEqual(result.judge.errorMessage, "no toolName in attack spec"); + assert.strictEqual(srv.llmCalls, 0); +}); + +test("single-turn attack calls the tool, judges once, and assembles the result", async () => { + srv.reset("PASS"); + const result = await runMcpAttack(mcpAttack(), fakeTarget("all good"), ...deps()); + assert.strictEqual(result.kind, "mcp"); + assert.strictEqual(result.judge.verdict, "PASS"); + assert.strictEqual(result.kind === "mcp" ? result.toolResponse : undefined, "all good"); + assert.strictEqual(result.kind === "mcp" ? result.toolName : undefined, "lookup"); + assert.strictEqual(result.turns, undefined); // single turn → not recorded as a turn list + assert.strictEqual(srv.llmCalls, 1); // final judge only +}); + +test("multi-turn attack early-stops on the first FAIL (no turn-2 generator call)", async () => { + srv.reset("FAIL"); + const result = await runMcpAttack( + mcpAttack({ turns: 2 }), + fakeTarget("leaked AKIA key here"), + ...deps() + ); + assert.strictEqual(result.judge.verdict, "FAIL"); + // Exactly one LLM call — the mid-turn judge at turn 1. Turn 2 (generator + + // final judge) never runs, proving the early stop. + assert.strictEqual(srv.llmCalls, 1); +}); + +test("multi-turn attack continues past a PASS turn — runs the turn-2 generator, then judges the last turn", async () => { + srv.reset("PASS"); + let toolCalls = 0; + const target: FakeTarget = { + async callTool() { + toolCalls++; + return { response: `turn ${toolCalls} response` }; + }, + async listTools() { + return []; + }, + async listResources() { + return []; + }, + async readResource() { + return ""; + }, + async close() {}, + }; + + const result = await runMcpAttack(mcpAttack({ turns: 2 }), target, ...deps()); + + assert.strictEqual(result.judge.verdict, "PASS"); + // Turn 1's mid-judge PASSed, so the loop continued: buildTurn's turn>1 branch + // called the generator, and both turns hit the tool. + assert.strictEqual(srv.turnGenCalls, 1); + assert.strictEqual(toolCalls, 2); + // Multi-turn attack records the full turn list; the final result surfaces the + // last turn's response (judged once at the end). + assert.strictEqual(result.turns?.length, 2); + assert.strictEqual(result.kind === "mcp" ? result.toolResponse : undefined, "turn 2 response"); + // Turn-1 mid judge + turn-2 generator + turn-2 final judge = 3 LLM calls. + assert.strictEqual(srv.llmCalls, 3); +});