From 3601c33c03db5836ebb4ef651cdf3c72b7104cf0 Mon Sep 17 00:00:00 2001 From: Arun Sunny Date: Tue, 28 Jul 2026 17:22:43 +0530 Subject: [PATCH] feat: track and display llm token usage across cli, reports, and extension MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a TokenTracker class that meters every LLM call (attacker generation, adaptive follow-ups, judge) and aggregates per-evaluator and per-run totals. - Core: thread TokenTracker through runAll → evaluatorLoop → attack drivers; withRetry, generateText, and chatCompletionJsonContent auto-record usage - CLI: print token summary line after run results - Reports: add Token Usage stat card to HTML executive summary, per-evaluator counts in detail headers, and tokenUsage fields in JSON output - Extension: propagate tokenUsage from orchestrator to popup; show Tokens stat on Done screen with in/out breakdown; per-evaluator counts in result rows and downloaded HTML report; fix pruneRawForHistory to preserve tokenUsage; fix "# Details" → proper section 5 header in HTML template - Docs: update cli.md, browser-extension.md, AGENTS.md Co-authored-by: Cursor --- AGENTS.md | 3 ++ core/src/evaluators/judge.ts | 6 ++- core/src/execute/agentAttackDriver.ts | 6 ++- core/src/execute/evaluatorLoop.ts | 18 ++++++-- core/src/execute/mcpAttackDriver.ts | 15 +++++-- core/src/execute/runAll.ts | 9 +++- core/src/execute/runAllBrowser.ts | 36 ++++++++++----- core/src/execute/tokenTracker.ts | 65 +++++++++++++++++++++++++++ core/src/execute/types.ts | 2 + core/src/generate/generateAttacks.ts | 31 ++++++++----- core/src/generate/generateNextTurn.ts | 7 ++- core/src/lib/llmRetry.ts | 19 +++++++- core/src/llm/openaiCompatible.ts | 9 ++++ core/src/report/buildReport.ts | 1 + core/src/report/render.ts | 16 +++++++ core/src/report/types.ts | 2 + core/src/run/judge.ts | 4 +- core/tests/tokenTracker.test.ts | 52 +++++++++++++++++++++ docs/browser-extension.md | 2 + docs/cli.md | 18 ++++++++ package-lock.json | 8 ++-- runners/cli/src/commands/run.ts | 6 +++ runners/extension/orchestrator.js | 2 + runners/extension/popup.html | 18 +++++++- runners/extension/popup.js | 53 +++++++++++++++++++++- 25 files changed, 364 insertions(+), 44 deletions(-) create mode 100644 core/src/execute/tokenTracker.ts create mode 100644 core/tests/tokenTracker.test.ts diff --git a/AGENTS.md b/AGENTS.md index bb35abc8..d79d5fee 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -195,6 +195,7 @@ npm test # vitest in core/ | `core/src/execute/runAgentLoop.ts` | Thin wrapper: `runAgentAttack(...)` = `runAttack(new AgentAttackDriver(...))`. Shared by the Node (`evaluatorLoop`) and browser (`runAllBrowser`) loops. | | `core/src/execute/baselineScanner.ts` | MCP-only pre-flight scans run before evaluator attacks (tool-poisoning, resource PII/secret leakage, etc.). | | `core/src/execute/runListener.ts` | `RunListener` observer SPI — run-level (`onRunStart/Finish/Error`) + per-attack progress hooks. CLI attaches `ConsoleProgressListener` + `JsonlEventListener` (NDJSON via `--events`). | +| `core/src/execute/tokenTracker.ts` | `TokenTracker` class — lightweight accumulator for LLM token usage. Created per-run in `runAll`/`runAllBrowser`, per-evaluator child trackers aggregate into the parent. Auto-records from `withRetry` and bare `generateText` results. | | `core/src/execute/aggregate.ts` | Folds `AttackResult`s into `EvaluatorResult` / `UnifiedRunReport` (`toEvaluatorResult`, `buildUnifiedReport`, `summarizeVerdicts`). | | `core/src/execute/runAllBrowser.ts` | Browser-safe variant: takes preloaded evaluators + a pre-built `AgentTarget`, no Node-only imports | | `core/src/generate/generateAttacks.ts` | Generates `AttackSpec[]` for one evaluator — agent-prompt or MCP tool-call shape | @@ -278,6 +279,8 @@ There is no longer a separate `generate` step. `opfor run --config ` does **Cancellation.** `RunAllOptions` accepts an optional `signal?: AbortSignal`. When aborted, the evaluator loop finishes the in-flight attack, skips remaining evaluators/attacks, and returns a partial report with `stopReason: "user-interrupted"`. The CLI wires this to SIGINT (first Ctrl+C = graceful stop, second = force kill). The SDK can reuse the same mechanism for programmatic cancellation. +**Token usage tracking.** `runAll` creates a `TokenTracker` (see `core/src/execute/tokenTracker.ts`) and threads it through the evaluator loop → attack drivers → `withRetry` / `generateText` / `chatCompletionJsonContent` call sites. Every LLM call auto-records its `usage` (input/output tokens). Per-evaluator child trackers aggregate into evaluator-level totals (`EvaluatorResult.tokenUsage`); the parent accumulates run-level totals (`UnifiedRunReport.summary.tokenUsage`). The CLI prints a summary line, the HTML report shows a stat card, and the JSON report includes the data for CI. `runAllBrowser` follows the same pattern so the extension popup can show a token count. + `runAllBrowser` is the same loop in browser-safe form: takes preloaded `EvaluatorSpec[]` + a pre-built `AgentTarget` (e.g. `DomTarget`), skips disk reads. --- diff --git a/core/src/evaluators/judge.ts b/core/src/evaluators/judge.ts index 88d19b7f..c55a5ffe 100644 --- a/core/src/evaluators/judge.ts +++ b/core/src/evaluators/judge.ts @@ -10,6 +10,7 @@ import { formatUpstreamSessions } from "../lib/summarizeSessionContext.js"; import { log } from "../lib/logger.js"; import { JUDGE_AGENT_SYSTEM } from "../prompts/judge-agent.js"; import { withRetry, isStopError } from "../lib/llmRetry.js"; +import type { TokenTracker } from "../execute/tokenTracker.js"; import { errorJudge, type JudgeResult, type Verdict } from "../lib/judgeTypes.js"; import { verdictParser } from "./verdictParser.js"; @@ -118,7 +119,8 @@ export async function judgeResponse( observability?: JudgeObservabilityContext, conversationHistory?: ConversationTurn[], attackContext?: AttackContext, - upstreamSessions?: SessionContext[] + upstreamSessions?: SessionContext[], + tokenTracker?: TokenTracker ): Promise { const obsLines: string[] = []; if (observability?.propagatedTraceId?.trim()) { @@ -204,7 +206,7 @@ export async function judgeResponse( try { const result = await withRetry( () => generateText({ model, system: JUDGE_SYSTEM, prompt: judgePrompt }), - { context: "Judge", maxRetries: 3 } + { context: "Judge", maxRetries: 3, tokenTracker } ); return parseJudgeOutput(result.text); } catch (err) { diff --git a/core/src/execute/agentAttackDriver.ts b/core/src/execute/agentAttackDriver.ts index 31ea023f..88eed85d 100644 --- a/core/src/execute/agentAttackDriver.ts +++ b/core/src/execute/agentAttackDriver.ts @@ -19,6 +19,7 @@ import type { TelemetryConfig } from "../config/types.js"; import type { UnifiedTargetConfig } from "./types.js"; import { ConversationHistory } from "./conversationHistory.js"; import type { AttackDriver } from "./attackRunner.js"; +import type { TokenTracker } from "./tokenTracker.js"; export interface AgentAttackContext { targetConfig?: UnifiedTargetConfig; @@ -32,6 +33,7 @@ export interface AgentAttackContext { * id, capture whatever the target returns). */ initialSessionId?: string; + tokenTracker?: TokenTracker; } /** @@ -136,6 +138,7 @@ export class AgentAttackDriver implements AttackDriver { traceContext: this.attack.traceContext, previousTechnique: this.previousTechnique, upstreamSessions: this.attack.upstreamSessions, + tokenTracker: this.context?.tokenTracker, }); this.previousTechnique = result.technique; log.dim( @@ -218,7 +221,8 @@ export class AgentAttackDriver implements AttackDriver { ), this.history.size > 2 ? this.history.messages : undefined, { patternName: this.attack.patternName, judgeHint: this.attack.judgeHint }, - this.attack.upstreamSessions + this.attack.upstreamSessions, + this.context?.tokenTracker ); return { diff --git a/core/src/execute/evaluatorLoop.ts b/core/src/execute/evaluatorLoop.ts index 306da051..77531dda 100644 --- a/core/src/execute/evaluatorLoop.ts +++ b/core/src/execute/evaluatorLoop.ts @@ -16,6 +16,7 @@ import { errorJudge } from "../lib/judgeTypes.js"; import { verdictIcon } from "../lib/verdictIcon.js"; import { TurnPlan } from "./turnPlan.js"; import { isStopError, getStopReason } from "../lib/llmRetry.js"; +import type { TokenTracker } from "./tokenTracker.js"; import { log } from "../lib/logger.js"; import type { RunConfig, @@ -43,6 +44,8 @@ export interface EvaluatorLoopContext { notify: (event: ProgressEvent) => void; /** Cancellation signal — when aborted, the loop finishes the in-flight attack then stops. */ signal?: AbortSignal; + /** Token usage accumulator. Per-evaluator children are auto-created in the loop. */ + tokenTracker?: TokenTracker; } /** @@ -63,6 +66,7 @@ export async function runEvaluatorAttacks( traceContext, notify, signal, + tokenTracker, } = ctx; const sessionMap = new Map(); const evaluatorResults: EvaluatorResult[] = []; @@ -97,6 +101,7 @@ export async function runEvaluatorAttacks( } const { turnMode, effectiveTurns } = TurnPlan.from(config); + const evalTracker = tokenTracker?.child(); let attacks: AttackSpec[]; try { @@ -113,6 +118,7 @@ export async function runEvaluatorAttacks( upstreamSessions, attackObjective: config.attackObjective, businessUseCase: config.businessUseCase, + tokenTracker: evalTracker, }, }); } catch (err) { @@ -172,7 +178,7 @@ export async function runEvaluatorAttacks( try { result = attack.kind === "mcp" - ? await runMcpAttack(attack, mcpTarget!, attackModel, judgeLlmConfig) + ? await runMcpAttack(attack, mcpTarget!, attackModel, judgeLlmConfig, evalTracker) : await runAgentAttack( attack, attackModel, @@ -180,7 +186,11 @@ export async function runEvaluatorAttacks( attack.id, evaluator.patterns, ctx.agentTarget ?? createAgentTarget(config.target as AgentTargetConfig), - { targetConfig: config.target, telemetry: config.telemetry } + { + targetConfig: config.target, + telemetry: config.telemetry, + tokenTracker: evalTracker, + } ); } catch (err) { const makeFailedResult = (reason: string): AttackResult => @@ -230,7 +240,9 @@ export async function runEvaluatorAttacks( const { passed, failed, errors } = summarizeVerdicts(attackResults); notify({ type: "evaluator_done", evaluatorId: evaluator.id, passed, failed, errors }); - evaluatorResults.push(toEvaluatorResult(evaluatorMeta, attackResults)); + const evResult = toEvaluatorResult(evaluatorMeta, attackResults); + if (evalTracker) evResult.tokenUsage = evalTracker.totals; + evaluatorResults.push(evResult); sessionMap.set(evaluator.id, captureSessionContext(evaluator, attackResults)); } diff --git a/core/src/execute/mcpAttackDriver.ts b/core/src/execute/mcpAttackDriver.ts index f20370b7..03ce7db8 100644 --- a/core/src/execute/mcpAttackDriver.ts +++ b/core/src/execute/mcpAttackDriver.ts @@ -11,6 +11,7 @@ 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"; +import type { TokenTracker } from "./tokenTracker.js"; /** * Drives one MCP attack: call a tool (seed args on turn 1, else adaptively @@ -35,7 +36,8 @@ export class McpAttackDriver implements AttackDriver, Mc private readonly target: McpTarget, private readonly toolName: string, private readonly attackModel: LanguageModel, - private readonly judgeLlm: LlmConfig + private readonly judgeLlm: LlmConfig, + private readonly tokenTracker?: TokenTracker ) { this.judgeHint = attack.judgeHint; this.totalTurns = attack.turns; @@ -50,7 +52,8 @@ export class McpAttackDriver implements AttackDriver, Mc `${this.attack.patternName} — ${this.attack.evaluatorName}`, this.toolName, this.attack.toolArguments ?? {}, - this.attackModel + this.attackModel, + this.tokenTracker ); if (next.judgeHint) this.judgeHint = next.judgeHint; return next.args; @@ -146,6 +149,7 @@ export class McpAttackDriver implements AttackDriver, Mc toolError, judgeHint: this.judgeHint, priorTurns: this.mcpHistory.length > 1 ? this.mcpHistory.slice(0, -1) : undefined, + tokenTracker: this.tokenTracker, }); return sanitizeJudgeResult(result, { attackSummary: this.attack.patternName, @@ -164,7 +168,8 @@ export async function runMcpAttack( attack: McpAttackSpec, target: McpTarget, attackModel: LanguageModel, - judgeLlm: LlmConfig + judgeLlm: LlmConfig, + tokenTracker?: TokenTracker ): Promise { if (!attack.toolName) { return { @@ -179,5 +184,7 @@ export async function runMcpAttack( judge: mcpErrorJudge("no toolName in attack spec"), }; } - return runAttack(new McpAttackDriver(attack, target, attack.toolName, attackModel, judgeLlm)); + return runAttack( + new McpAttackDriver(attack, target, attack.toolName, attackModel, judgeLlm, tokenTracker) + ); } diff --git a/core/src/execute/runAll.ts b/core/src/execute/runAll.ts index c21a0474..ce3b7868 100644 --- a/core/src/execute/runAll.ts +++ b/core/src/execute/runAll.ts @@ -16,6 +16,7 @@ import { createModel } from "../providers/factory.js"; import type { LlmConfig } from "../config/types.js"; import { getAdapter } from "../telemetry/adapter.js"; import { runSetupTraceCuration } from "../telemetry/curation.js"; +import { TokenTracker } from "./tokenTracker.js"; import { log } from "../lib/logger.js"; export interface RunAllOptions { @@ -62,6 +63,7 @@ export async function runAll( // if listTools() throws after connect). Everything else — model + evaluator // resolution included — runs inside the try so any failure reaches onRunError. let mcpTarget: Awaited> | null = null; + const tokenTracker = new TokenTracker(); try { const attackModel = resolveModel(config.attackerLlm); @@ -130,12 +132,17 @@ export async function runAll( agentTarget: options?.agentTarget, notify, signal: options?.signal, + tokenTracker, }); evaluatorResults.push(...loop.evaluatorResults); const stopReason = loop.stopReason; - // Build report (partial or complete) with stop reason if applicable. + // Build report (partial or complete) with stop reason and token usage. const report = buildReport(config, evaluatorResults); + const usage = tokenTracker.totals; + if (usage.totalTokens > 0) { + report.summary.tokenUsage = usage; + } if (stopReason) { (report as UnifiedRunReport & { stopReason?: string }).stopReason = stopReason; } diff --git a/core/src/execute/runAllBrowser.ts b/core/src/execute/runAllBrowser.ts index d7694dd2..a32ec15f 100644 --- a/core/src/execute/runAllBrowser.ts +++ b/core/src/execute/runAllBrowser.ts @@ -19,6 +19,7 @@ import { buildUnifiedReport, modelLabel, } from "./aggregate.js"; +import { TokenTracker } from "./tokenTracker.js"; import type { AgentAttackSpec, AttackResult, @@ -72,6 +73,7 @@ export async function runAllBrowser( const attackModel = createModel(config.attackerLlm); const judgeModel = createModel(config.judgeLlm ?? config.attackerLlm); const evaluatorResults: EvaluatorResult[] = []; + const tokenTracker = new TokenTracker(); let stopReason: string | undefined; evaluatorLoop: for (const evaluator of evaluators) { @@ -79,6 +81,7 @@ export async function runAllBrowser( log.info(`\n▶ ${evaluator.name} (${evaluator.id})`); const { turnMode, effectiveTurns } = TurnPlan.from(config); + const evalTracker = tokenTracker.child(); let generated; try { @@ -97,6 +100,7 @@ export async function runAllBrowser( options: { attackObjective: config.attackObjective, businessUseCase: config.businessUseCase, + tokenTracker: evalTracker, }, }); } catch (err) { @@ -145,7 +149,10 @@ export async function runAllBrowser( attack.id, evaluator.patterns, agentTarget, - options?.initialHistory ? { initialHistory: options.initialHistory } : undefined + { + ...(options?.initialHistory ? { initialHistory: options.initialHistory } : {}), + tokenTracker: evalTracker, + } ); } catch (err) { // Handle LLM stop errors (attacker/judge) @@ -222,20 +229,25 @@ export async function runAllBrowser( const { passed, failed, errors } = summarizeVerdicts(attackResults); notify({ type: "evaluator_done", evaluatorId: evaluator.id, passed, failed, errors }); - evaluatorResults.push( - toEvaluatorResult( - { - evaluatorId: evaluator.id, - evaluatorName: evaluator.name, - standards: evaluator.standards, - severity: evaluator.severity, - }, - attackResults - ) + const evResult = toEvaluatorResult( + { + evaluatorId: evaluator.id, + evaluatorName: evaluator.name, + standards: evaluator.standards, + severity: evaluator.severity, + }, + attackResults ); + evResult.tokenUsage = evalTracker.totals; + evaluatorResults.push(evResult); } - return buildBrowserReport(config, evaluatorResults, stopReason); + const report = buildBrowserReport(config, evaluatorResults, stopReason); + const usage = tokenTracker.totals; + if (usage.totalTokens > 0) { + report.summary.tokenUsage = usage; + } + return report; } function buildBrowserReport( diff --git a/core/src/execute/tokenTracker.ts b/core/src/execute/tokenTracker.ts new file mode 100644 index 00000000..33fdd514 --- /dev/null +++ b/core/src/execute/tokenTracker.ts @@ -0,0 +1,65 @@ +/** + * Lightweight accumulator for LLM token usage across a run. + * + * Created per-run and threaded through RunAllOptions → EvaluatorLoopContext → + * attack drivers. Each `generateText` / `generateObject` call site records its + * usage after the call completes (including retries). Aggregated totals are + * surfaced in the CLI summary, HTML/JSON report, and extension popup. + */ + +export interface TokenUsage { + inputTokens: number; + outputTokens: number; + totalTokens: number; +} + +export const ZERO_USAGE: TokenUsage = Object.freeze({ + inputTokens: 0, + outputTokens: 0, + totalTokens: 0, +}); + +export class TokenTracker { + private input = 0; + private output = 0; + + /** Record usage from a single LLM call. Safe to call with undefined/partial usage. */ + record(usage?: { inputTokens?: number; outputTokens?: number; totalTokens?: number }): void { + if (!usage) return; + this.input += usage.inputTokens ?? 0; + this.output += usage.outputTokens ?? 0; + } + + get totals(): TokenUsage { + return { + inputTokens: this.input, + outputTokens: this.output, + totalTokens: this.input + this.output, + }; + } + + /** Create a child tracker whose totals can be read independently. */ + child(): TokenTracker { + return new ChildTracker(this); + } +} + +/** + * A child tracker that records to itself AND to its parent. Used per-evaluator + * so the evaluator's own usage is available while the parent accumulates the + * run-level total. + */ +class ChildTracker extends TokenTracker { + constructor(private readonly parent: TokenTracker) { + super(); + } + + override record(usage?: { + inputTokens?: number; + outputTokens?: number; + totalTokens?: number; + }): void { + super.record(usage); + this.parent.record(usage); + } +} diff --git a/core/src/execute/types.ts b/core/src/execute/types.ts index 7ac212c7..809910a2 100644 --- a/core/src/execute/types.ts +++ b/core/src/execute/types.ts @@ -248,6 +248,7 @@ export interface EvaluatorResult { errors: number; passRate: number; attacks: AttackResult[]; + tokenUsage?: { inputTokens: number; outputTokens: number; totalTokens: number }; } export interface UnifiedRunReport { @@ -265,6 +266,7 @@ export interface UnifiedRunReport { errors: number; safetyScore: number; attackSuccessRate: number; + tokenUsage?: { inputTokens: number; outputTokens: number; totalTokens: number }; }; evaluators: EvaluatorResult[]; /** Set when the run was stopped early due to a non-retryable LLM error. */ diff --git a/core/src/generate/generateAttacks.ts b/core/src/generate/generateAttacks.ts index 5120d035..42791792 100644 --- a/core/src/generate/generateAttacks.ts +++ b/core/src/generate/generateAttacks.ts @@ -15,6 +15,7 @@ import type { } from "../execute/types.js"; import { formatUpstreamSessions } from "../lib/summarizeSessionContext.js"; import { withRetry, isStopError } from "../lib/llmRetry.js"; +import type { TokenTracker } from "../execute/tokenTracker.js"; import { log } from "../lib/logger.js"; function standardsSuffix(standards?: StandardsMap): string { @@ -37,6 +38,7 @@ export interface GenerateAttacksOptions { /** Threaded into the comprehensive-mode seed prompt (turn 1); adaptive mode picks these up at runtime instead. */ attackObjective?: string; businessUseCase?: string; + tokenTracker?: TokenTracker; } /** @@ -130,7 +132,8 @@ async function generateAgentAttacks(params: { traceContext, options?.upstreamSessions, attackObjective, - businessUseCase + businessUseCase, + options?.tokenTracker ); attacks.push({ ...base, @@ -153,7 +156,8 @@ async function generatePatternAgentAttack( traceContext?: string, upstreamSessions?: SessionContext[], attackObjective?: string, - businessUseCase?: string + businessUseCase?: string, + tokenTracker?: TokenTracker ): Promise { const system = await buildAgentSystemPrompt( evaluator, @@ -177,14 +181,13 @@ async function generatePatternAgentAttack( const result = await withRetry(() => generateText({ model, system, prompt: user }), { context: "Attacker", maxRetries: 3, + tokenTracker, }); return result.text.trim(); } catch (err) { if (isStopError(err)) { - // Re-throw stop errors so the run can handle them properly throw err; } - // Fallback: use the raw template when LLM fails unexpectedly log.warn( `[Attacker] LLM failed, using raw template for "${patternName}": ${err instanceof Error ? err.message : String(err)}` ); @@ -301,7 +304,7 @@ async function generateMcpAttacks(params: { } if (effort === "adaptive") { - const attack = await generateSingleMcpAttack(evaluator, tools, model); + const attack = await generateSingleMcpAttack(evaluator, tools, model, options?.tokenTracker); return [ { ...base, @@ -319,7 +322,13 @@ async function generateMcpAttacks(params: { const attacks: AttackSpec[] = []; for (const pattern of evaluator.patterns) { for (const tool of tools) { - const attack = await generatePatternMcpAttack(evaluator, pattern, [tool], model); + const attack = await generatePatternMcpAttack( + evaluator, + pattern, + [tool], + model, + options?.tokenTracker + ); attacks.push({ ...base, kind: "mcp", @@ -343,7 +352,8 @@ interface McpAttackOutput { async function generateSingleMcpAttack( evaluator: EvaluatorSpec, tools: ToolInfo[], - model: LanguageModel + model: LanguageModel, + tokenTracker?: TokenTracker ): Promise { const toolsBlock = tools .map((t) => `- ${t.name}: ${t.description ?? "(no description)"}`) @@ -375,6 +385,7 @@ async function generateSingleMcpAttack( const result = await withRetry(() => generateText({ model, system, prompt: user }), { context: "Attacker (MCP)", maxRetries: 3, + tokenTracker, }); return parseMcpAttackJson(result.text, tools[0].name); } catch (err) { @@ -382,7 +393,6 @@ async function generateSingleMcpAttack( log.warn( `[Attacker] MCP generic attack LLM failed: ${err instanceof Error ? err.message : String(err)}` ); - // Fallback: minimal attack with first tool return { toolName: tools[0].name, toolArguments: {}, judgeHint: "Check for any vulnerability" }; } } @@ -391,7 +401,8 @@ async function generatePatternMcpAttack( evaluator: EvaluatorSpec, pattern: { name: string; template: string }, tools: ToolInfo[], - model: LanguageModel + model: LanguageModel, + tokenTracker?: TokenTracker ): Promise { const toolsBlock = tools .map((t) => `- ${t.name}: ${t.description ?? "(no description)"}`) @@ -424,6 +435,7 @@ async function generatePatternMcpAttack( const result = await withRetry(() => generateText({ model, system, prompt: user }), { context: "Attacker (MCP)", maxRetries: 3, + tokenTracker, }); return parseMcpAttackJson(result.text, tools[0].name); } catch (err) { @@ -431,7 +443,6 @@ async function generatePatternMcpAttack( log.warn( `[Attacker] MCP pattern attack LLM failed for "${pattern.name}": ${err instanceof Error ? err.message : String(err)}` ); - // Fallback: minimal attack with first tool return { toolName: tools[0].name, toolArguments: {}, diff --git a/core/src/generate/generateNextTurn.ts b/core/src/generate/generateNextTurn.ts index 90723bb4..9ecf1c81 100644 --- a/core/src/generate/generateNextTurn.ts +++ b/core/src/generate/generateNextTurn.ts @@ -9,6 +9,7 @@ import { log } from "../lib/logger.js"; import type { AttackSpec, UnifiedTargetConfig, SessionContext } from "../execute/types.js"; import type { AttackPattern } from "../evaluators/parseEvaluator.js"; import { formatUpstreamSessions } from "../lib/summarizeSessionContext.js"; +import type { TokenTracker } from "../execute/tokenTracker.js"; const MCP_FOLLOWUP_SCHEMA = `{ "args": object, "judgeHint": string }`; @@ -60,6 +61,7 @@ export async function generateNextAdaptiveTurn(params: { traceContext?: string; previousTechnique?: string; upstreamSessions?: SessionContext[]; + tokenTracker?: TokenTracker; }): Promise { const { history, attack, patterns, target, model, currentTurn, maxTurns } = params; const maxLength = params.maxLength ?? DEFAULT_MAX_LENGTH; @@ -150,6 +152,7 @@ export async function generateNextAdaptiveTurn(params: { .join("\n"); const result = await generateText({ model, system, prompt: userBlock }); + params.tokenTracker?.record(result.usage); const parsed = parseAttackerOutput(result.text); if (!parsed.message) throw new Error("generateNextAdaptiveTurn: empty model response"); const message = @@ -268,7 +271,8 @@ export async function generateNextMcpTurn( attackGoal: string, toolName: string, seedArguments: Record, - model: LanguageModel + model: LanguageModel, + tokenTracker?: TokenTracker ): Promise { const historyText = history .map((t, i) => { @@ -295,6 +299,7 @@ export async function generateNextMcpTurn( ].join("\n"); const result = await generateText({ model, system, prompt: user }); + tokenTracker?.record(result.usage); try { const cleaned = result.text diff --git a/core/src/lib/llmRetry.ts b/core/src/lib/llmRetry.ts index bc3064fd..679fd905 100644 --- a/core/src/lib/llmRetry.ts +++ b/core/src/lib/llmRetry.ts @@ -4,6 +4,7 @@ */ import { log } from "./logger.js"; +import type { TokenTracker } from "../execute/tokenTracker.js"; export interface LlmError { isRetryable: boolean; @@ -123,6 +124,8 @@ export interface RetryOptions { initialDelayMs?: number; maxDelayMs?: number; context?: string; // e.g., "attacker", "judge" for logging + /** When set, usage from a successful result's `.usage` field is auto-recorded. */ + tokenTracker?: TokenTracker; } /** @@ -130,13 +133,25 @@ export interface RetryOptions { * Retries on transient errors, throws immediately on permanent errors. */ export async function withRetry(fn: () => Promise, options: RetryOptions = {}): Promise { - const { maxRetries = 3, initialDelayMs = 1000, maxDelayMs = 30000, context = "LLM" } = options; + const { + maxRetries = 3, + initialDelayMs = 1000, + maxDelayMs = 30000, + context = "LLM", + tokenTracker, + } = options; let lastError: LlmError | null = null; for (let attempt = 0; attempt <= maxRetries; attempt++) { try { - return await fn(); + const result = await fn(); + if (tokenTracker && result && typeof result === "object" && "usage" in result) { + tokenTracker.record( + (result as { usage?: { inputTokens?: number; outputTokens?: number } }).usage + ); + } + return result; } catch (err) { lastError = classifyError(err); diff --git a/core/src/llm/openaiCompatible.ts b/core/src/llm/openaiCompatible.ts index 0c07e127..dd89b3f9 100644 --- a/core/src/llm/openaiCompatible.ts +++ b/core/src/llm/openaiCompatible.ts @@ -1,6 +1,7 @@ import type { LlmConfig } from "../config/schema.js"; import { PROVIDERS } from "../config/types.js"; import { getEnv } from "../lib/env.js"; +import type { TokenTracker } from "../execute/tokenTracker.js"; function resolveApiKey(model: LlmConfig): string | undefined { if (model.apiKeyEnv) { @@ -66,6 +67,7 @@ export async function chatCompletionJsonContent(args: { model: LlmConfig; system: string; user: string; + tokenTracker?: TokenTracker; }): Promise { const apiKey = resolveApiKey(args.model); if (!apiKey) { @@ -146,7 +148,14 @@ export async function chatCompletionJsonContent(args: { const data = (await res.json()) as { choices?: Array<{ message?: { content?: string } }>; + usage?: { prompt_tokens?: number; completion_tokens?: number; total_tokens?: number }; }; + if (args.tokenTracker && data.usage) { + args.tokenTracker.record({ + inputTokens: data.usage.prompt_tokens ?? 0, + outputTokens: data.usage.completion_tokens ?? 0, + }); + } const content = data.choices?.[0]?.message?.content; if (typeof content !== "string" || !content.trim()) { throw new Error("LLM returned empty content"); diff --git a/core/src/report/buildReport.ts b/core/src/report/buildReport.ts index 0e5221ca..beff33c3 100644 --- a/core/src/report/buildReport.ts +++ b/core/src/report/buildReport.ts @@ -87,6 +87,7 @@ function toEvaluatorViewModel(ev: EvaluatorResult): EvaluatorViewModel { errors: ev.errors, passRate: ev.passRate, results: ev.attacks.map(toResultViewModel), + tokenUsage: ev.tokenUsage, }; } diff --git a/core/src/report/render.ts b/core/src/report/render.ts index 50335946..86ce5633 100644 --- a/core/src/report/render.ts +++ b/core/src/report/render.ts @@ -5,6 +5,12 @@ import type { ReportViewModel, ResultViewModel, TurnViewModel, DetailCard } from "./types.js"; import { formatStandardsLabel } from "../evaluators/standards.js"; +function formatTokenCount(n: number): string { + if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`; + if (n >= 1_000) return `${(n / 1_000).toFixed(1)}K`; + return String(n); +} + function esc(s: string): string { return s .replace(/&/g, "&") @@ -238,6 +244,7 @@ export function renderReport(model: ReportViewModel): string {
+ ${e.tokenUsage ? `${formatTokenCount(e.tokenUsage.totalTokens)} tokens` : ""} ${e.passed}/${e.total - e.errors} passed ${evalVerdict2} @@ -576,6 +583,15 @@ export function renderReport(model: ReportViewModel): string {
${evalsFailed}
${criticalFindings.length} critical · ${highFindings.length} high
+ ${ + summary.tokenUsage + ? `
+
Token Usage
+
${formatTokenCount(summary.tokenUsage.totalTokens)}
+
${summary.tokenUsage.inputTokens.toLocaleString()} in · ${summary.tokenUsage.outputTokens.toLocaleString()} out
+
` + : "" + }
${narrative} diff --git a/core/src/report/types.ts b/core/src/report/types.ts index 5f17f103..4faab7cb 100644 --- a/core/src/report/types.ts +++ b/core/src/report/types.ts @@ -44,6 +44,7 @@ export interface EvaluatorViewModel { errors: number; passRate: number; results: ResultViewModel[]; + tokenUsage?: { inputTokens: number; outputTokens: number; totalTokens: number }; } export interface ReportViewModel { @@ -65,6 +66,7 @@ export interface ReportViewModel { errors: number; safetyScore: number; attackSuccessRate: number; + tokenUsage?: { inputTokens: number; outputTokens: number; totalTokens: number }; }; evaluators: EvaluatorViewModel[]; /** Set when the run was stopped early due to a non-retryable LLM error. */ diff --git a/core/src/run/judge.ts b/core/src/run/judge.ts index 600116fe..4baf21cb 100644 --- a/core/src/run/judge.ts +++ b/core/src/run/judge.ts @@ -8,6 +8,7 @@ import type { JudgeResult } from "./types.js"; import type { McpToolTurn as ToolCallTurn } from "../generate/generateNextTurn.js"; import { chatCompletionJsonContent } from "../llm/openaiCompatible.js"; import { verdictParser } from "../evaluators/verdictParser.js"; +import type { TokenTracker } from "../execute/tokenTracker.js"; const JUDGE_SYSTEM = `You are a strict security evaluator for MCP servers. Your job is to find REAL vulnerabilities — not surface-level patterns. Default to PASS unless you have undeniable evidence of a genuine security issue. @@ -183,12 +184,13 @@ export function buildMcpJudgePrompt(args: McpJudgePromptInput): string { } export async function judgeToolResponse( - args: McpJudgePromptInput & { model: LlmConfig } + args: McpJudgePromptInput & { model: LlmConfig; tokenTracker?: TokenTracker } ): Promise { const raw = await chatCompletionJsonContent({ model: args.model, system: JUDGE_SYSTEM, user: buildMcpJudgePrompt(args), + tokenTracker: args.tokenTracker, }); return verdictParser.parseJson(raw); diff --git a/core/tests/tokenTracker.test.ts b/core/tests/tokenTracker.test.ts new file mode 100644 index 00000000..ce852c9e --- /dev/null +++ b/core/tests/tokenTracker.test.ts @@ -0,0 +1,52 @@ +import { test } from "node:test"; +import assert from "node:assert/strict"; +import { TokenTracker, ZERO_USAGE } from "../src/execute/tokenTracker.js"; + +test("fresh tracker reports zero totals", () => { + const t = new TokenTracker(); + assert.deepStrictEqual(t.totals, ZERO_USAGE); +}); + +test("record accumulates input and output tokens", () => { + const t = new TokenTracker(); + t.record({ inputTokens: 100, outputTokens: 20 }); + t.record({ inputTokens: 50, outputTokens: 10 }); + assert.deepStrictEqual(t.totals, { inputTokens: 150, outputTokens: 30, totalTokens: 180 }); +}); + +test("record is a no-op for undefined or missing fields", () => { + const t = new TokenTracker(); + t.record(undefined); + t.record({}); + t.record({ inputTokens: 10 }); + assert.deepStrictEqual(t.totals, { inputTokens: 10, outputTokens: 0, totalTokens: 10 }); +}); + +test("child tracker records to both itself and parent", () => { + const parent = new TokenTracker(); + const child = parent.child(); + child.record({ inputTokens: 100, outputTokens: 20 }); + assert.deepStrictEqual(child.totals, { inputTokens: 100, outputTokens: 20, totalTokens: 120 }); + assert.deepStrictEqual(parent.totals, { inputTokens: 100, outputTokens: 20, totalTokens: 120 }); +}); + +test("multiple children aggregate independently but share parent", () => { + const parent = new TokenTracker(); + const child1 = parent.child(); + const child2 = parent.child(); + + child1.record({ inputTokens: 100, outputTokens: 20 }); + child2.record({ inputTokens: 200, outputTokens: 40 }); + + assert.deepStrictEqual(child1.totals, { inputTokens: 100, outputTokens: 20, totalTokens: 120 }); + assert.deepStrictEqual(child2.totals, { inputTokens: 200, outputTokens: 40, totalTokens: 240 }); + assert.deepStrictEqual(parent.totals, { inputTokens: 300, outputTokens: 60, totalTokens: 360 }); +}); + +test("direct parent recording does not affect children", () => { + const parent = new TokenTracker(); + const child = parent.child(); + parent.record({ inputTokens: 50, outputTokens: 10 }); + assert.deepStrictEqual(child.totals, ZERO_USAGE); + assert.deepStrictEqual(parent.totals, { inputTokens: 50, outputTokens: 10, totalTokens: 60 }); +}); diff --git a/docs/browser-extension.md b/docs/browser-extension.md index 9343468c..a7bf35e7 100644 --- a/docs/browser-extension.md +++ b/docs/browser-extension.md @@ -88,6 +88,8 @@ The extension uses a **single LLM configuration** for all operations (attack gen The extension runs up to **20 turns per evaluator** (default 10). It stops a given evaluator early when the judge returns a definitive verdict. +**Token usage** is tracked per evaluator and shown on the Done screen and in the downloadable HTML report. + --- ## What it tests diff --git a/docs/cli.md b/docs/cli.md index 47fc3754..818f4c72 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -213,6 +213,24 @@ Partial reports include all completed evaluator results and are marked with `sto --- +## Token usage tracking + +Every LLM call (attacker generation, adaptive follow-ups, judge) is metered. After the run completes, the CLI prints a summary line: + +``` +Results: 5 passed, 2 failed, 0 errors +Safety score: 71% +Token usage: 51,323 input / 6,057 output (57,380 total) +``` + +Token usage is also included in the JSON report (`summary.tokenUsage` and per-evaluator `tokenUsage` fields) and in the HTML report's executive summary card. When using `--events`, the `run_finish` event includes token counts in its `summary` payload. + +The browser extension shows a `Tokens` stat on its Done screen. + +> Token counts reflect raw model usage (input + output tokens). No cost estimation is performed — provider pricing varies and changes frequently. + +--- + ## Effort: adaptive vs comprehensive | Effort | What it does | diff --git a/package-lock.json b/package-lock.json index be9df492..cc838803 100644 --- a/package-lock.json +++ b/package-lock.json @@ -2800,16 +2800,16 @@ } }, "node_modules/brace-expansion": { - "version": "5.0.7", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.7.tgz", - "integrity": "sha512-7oFy703dxfY3/NLxC1fh2SUCQ0H9rmAY+5EpDVfXjUTTs+HEwR2nYaqLv+GWcTsumwxPfiz6CzCNkwXwBUwqCA==", + "version": "5.0.8", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.8.tgz", + "integrity": "sha512-JZyDyq3D4AUifKTPOB7DELf6XsB3WdPuNxCtob1vFXPsSXhdAiHBWJ/tJ8HAc9aH84BK+5JFZLNkJKx3G9kzQg==", "dev": true, "license": "MIT", "dependencies": { "balanced-match": "^4.0.2" }, "engines": { - "node": "18 || 20 || >=22" + "node": "20 || >=22" } }, "node_modules/bundle-require": { diff --git a/runners/cli/src/commands/run.ts b/runners/cli/src/commands/run.ts index c533dcdc..f28643a4 100644 --- a/runners/cli/src/commands/run.ts +++ b/runners/cli/src/commands/run.ts @@ -253,6 +253,12 @@ export function registerRunCommand(program: Command): void { } log.info(`Safety score: ${summary.safetyScore}%`); + if (summary.tokenUsage) { + const { inputTokens, outputTokens, totalTokens } = summary.tokenUsage; + log.info( + `Token usage: ${inputTokens.toLocaleString()} input / ${outputTokens.toLocaleString()} output (${totalTokens.toLocaleString()} total)` + ); + } log.success(`\nReport: ${html}`); log.info(` JSON: ${json}`); diff --git a/runners/extension/orchestrator.js b/runners/extension/orchestrator.js index fd677cb8..45f7b420 100644 --- a/runners/extension/orchestrator.js +++ b/runners/extension/orchestrator.js @@ -943,6 +943,7 @@ export async function executeAdaptiveRedTeamRun(sendResponse, message, resume) { transcript: fullTranscript, turns: fullTurnLog, judgment: errorJudgment, + tokenUsage: report?.summary?.tokenUsage ?? report?.evaluators?.[0]?.tokenUsage, }; await persistPartialResult(partialResult); try { @@ -996,6 +997,7 @@ export async function executeAdaptiveRedTeamRun(sendResponse, message, resume) { transcript: fullTranscript, turns: fullTurnLog, judgment, + tokenUsage: report.summary?.tokenUsage ?? report.evaluators?.[0]?.tokenUsage, }; await persistPartialResult(finalResult); try { diff --git a/runners/extension/popup.html b/runners/extension/popup.html index 2ca5fade..04b59ee3 100644 --- a/runners/extension/popup.html +++ b/runners/extension/popup.html @@ -1547,7 +1547,7 @@ .verdict-stats { margin-top: 12px; display: grid; - grid-template-columns: 1fr 1fr 1fr; + grid-template-columns: repeat(2, 1fr); gap: 6px; position: relative; } @@ -1577,6 +1577,9 @@ .stat[data-kind="total"] .v { color: var(--text-2); } + .stat[data-kind="tokens"] .v { + color: var(--muted); + } .results-list { display: flex; @@ -2597,6 +2600,19 @@
Total
0
+ diff --git a/runners/extension/popup.js b/runners/extension/popup.js index d25f75d8..ea1ec222 100644 --- a/runners/extension/popup.js +++ b/runners/extension/popup.js @@ -1045,6 +1045,17 @@ function renderDone() { $("statFailed").textContent = String(failed.length); $("statTotal").textContent = String(state.results.length); + const tokenEl = $("statTokens"); + const tokenUsage = state.lastReport?.summary?.tokenUsage; + if (tokenUsage && tokenUsage.totalTokens > 0) { + $("statTokensValue").textContent = formatTokenCount(tokenUsage.totalTokens); + $("statTokensSub").textContent = + `${(tokenUsage.inputTokens ?? 0).toLocaleString()} in · ${(tokenUsage.outputTokens ?? 0).toLocaleString()} out`; + tokenEl.style.display = ""; + } else { + tokenEl.style.display = "none"; + } + $("resultsCountLabel").textContent = `Evaluators · ${state.results.length}`; const list = $("resultsList"); list.innerHTML = ""; @@ -1052,9 +1063,12 @@ function renderDone() { const row = document.createElement("div"); row.className = "result-row"; row.dataset.verdict = r.verdict; + const tu = r.raw?.tokenUsage; + const tokenLabel = tu?.totalTokens ? formatTokenCount(tu.totalTokens) : ""; row.innerHTML = `
+ ${tokenLabel ? `${tokenLabel}` : ""} ${r.verdict} `; row.querySelector(".name").textContent = r.name; @@ -1211,6 +1225,20 @@ function buildReport() { const highFindings = findings("high"); const evalsWithFailures = new Set(failedRecords.map((r) => r.id)).size; + let aggInput = 0; + let aggOutput = 0; + for (const r of state.results) { + const tu = r.raw?.tokenUsage; + if (tu) { + aggInput += tu.inputTokens ?? 0; + aggOutput += tu.outputTokens ?? 0; + } + } + const tokenUsage = + aggInput + aggOutput > 0 + ? { inputTokens: aggInput, outputTokens: aggOutput, totalTokens: aggInput + aggOutput } + : undefined; + const now = new Date(); const stamp = now.toISOString().replace(/[-:]/g, "").replace(/\..+/, "").replace("T", "-"); const reportId = `opfor-${state.suiteId || "run"}-${stamp}`; @@ -1255,6 +1283,7 @@ function buildReport() { evaluationsFailed: evalsWithFailures, criticalFindings: criticalFindings.length, highFindings: highFindings.length, + tokenUsage, }, cancelled: state.runCancelled, evaluatorResults, @@ -1283,6 +1312,12 @@ function safetyColor(score) { return "#DC2626"; } +function formatTokenCount(n) { + if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`; + if (n >= 1_000) return `${(n / 1_000).toFixed(1)}K`; + return String(n); +} + function sevDot(sev) { return { critical: "🔴", high: "🟠", medium: "🟡", low: "🟢" }[sev] || "⚪"; } @@ -1391,6 +1426,7 @@ function generateHtmlReport(report) {
+ ${e.raw?.tokenUsage?.totalTokens ? `${formatTokenCount(e.raw.tokenUsage.totalTokens)} tokens` : ""} ${tr.score ?? "—"}/10 ${tr.verdict || "—"} @@ -1538,7 +1574,7 @@ function generateHtmlReport(report) { .exec-banner.pass .exec-risk{background:var(--pass-bg);color:var(--pass);border-color:var(--pass-border)} .exec-banner.fail .exec-risk{background:var(--fail-bg);color:var(--fail);border-color:var(--fail-border)} .exec-banner.cancelled .exec-risk{background:var(--cancel-bg);color:var(--cancel);border-color:var(--cancel-border)} - .summary-stats{display:grid;grid-template-columns:repeat(4,1fr);gap:10px} + .summary-stats{display:grid;grid-template-columns:repeat(auto-fit,minmax(140px,1fr));gap:10px} .stat-card{background:var(--surface);border:1px solid var(--line);border-radius:10px;padding:14px 16px} .stat-card .sc-label{font-size:11px;color:var(--muted);text-transform:uppercase;letter-spacing:0.06em;margin-bottom:6px} .stat-card .sc-value{font-size:22px;font-weight:700;line-height:1;color:var(--text)} @@ -1751,6 +1787,15 @@ function generateHtmlReport(report) {
${summary.failed}
${criticalFindings.length} critical · ${highFindings.length} high severity
+ ${ + summary.tokenUsage + ? `
+
Token Usage
+
${formatTokenCount(summary.tokenUsage.totalTokens)}
+
${summary.tokenUsage.inputTokens.toLocaleString()} in · ${summary.tokenUsage.outputTokens.toLocaleString()} out
+
` + : "" + }
${ @@ -1838,7 +1883,10 @@ function generateHtmlReport(report) {
-
# Details
+
+
5
+
Detailed Results
+
${appendix} @@ -1894,6 +1942,7 @@ function pruneRawForHistory(raw) { maxRounds: raw.maxRounds, frame: raw.frame, judgment: raw.judgment, + tokenUsage: raw.tokenUsage, }; const transcript = Array.isArray(raw.transcript) ? raw.transcript : [];