refactor: unify agent and mcp attack loops behind one attack runner#152
Conversation
The turn loop existed twice — runAgentAttack and a 134-line runMcpAttack embedded in runAll — sharing a control flow but drifting in detail. Introduce a Template Method: runAttack owns the invariant loop (build -> execute -> record -> shouldEarlyStop, then finalize once), and an AttackDriver fills the kind-specific holes. - attackRunner.ts: AttackDriver<TInput,TOutput> interface + runAttack (a stateless function; no needless class). - agentAttackDriver.ts: AgentAttackDriver — agent logic moved verbatim. - mcpAttackDriver.ts: McpAttackDriver + runMcpAttack — extracted out of runAll (runAll: 985 -> ~840 lines). The MCP mid-turn and final judge blocks, previously duplicated verbatim, are consolidated into one private judge() helper. - runAgentLoop.ts: runAgentAttack is now a thin wrapper with an unchanged signature, so runAll/runAllBrowser/SDK/extension callers are untouched. Behavior-preserving: agent target-error early-stop + judge-once + target.close(); MCP mid-turn judge only when turns>1 && t<turns with priorTurns=history.slice(0,-1), early-stop-on-FAIL, sanitizeJudgeResult, empty-toolName short-circuit, and result shapes are all reproduced. Guarded by the runAll.smoke + orchestrator.equivalence drift tests plus new attackRunner + mcpAttackDriver tests (the MCP loop had no coverage before). Drivers take narrowly-injected dependencies (the run-level attacker model, the judge config) rather than the whole RunConfig; the agent driver tracks the previous technique as a scalar instead of a full per-turn meta array. The Judge SPI is intentionally deferred — each driver owns its judging, and the agent/MCP judges take different model types, which the SPI will unify in the hunt work. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
Warning Review limit reached
Next review available in: 45 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
WalkthroughThis PR extracts a shared attack runner, moves agent and MCP attack orchestration into dedicated drivers, rewires the existing entry points to use them, and adds runner and MCP driver tests. ChangesAttack Runner Refactor
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant runAttack
participant AttackDriver
participant Target
participant Judge
Caller->>runAttack: runAttack(driver)
loop each turn
runAttack->>AttackDriver: buildTurn(turn)
runAttack->>AttackDriver: execute(input)
AttackDriver->>Target: send/callTool
Target-->>AttackDriver: response/result
runAttack->>AttackDriver: record(turn, input, output)
runAttack->>AttackDriver: shouldEarlyStop(...)
alt early stop
runAttack->>runAttack: break loop
end
end
runAttack->>AttackDriver: finalize()
AttackDriver->>Judge: judgeResponse/judgeToolResponse
Judge-->>AttackDriver: verdict
AttackDriver-->>Caller: AttackResult
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (3)
core/tests/mcpAttackDriver.test.ts (1)
149-180: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winGap: no test exercises turn continuation (mid-turn PASS → turn 2).
All three tests either single-turn or early-stop at turn 1, so
buildTurn's turn>1 branch (generateNextMcpTurncall,judgeHintupdate) and the populatedpriorTurnsbranch injudge()(mcpAttackDriver.ts Line 148) are never exercised. A 2-turn scenario where turn 1 returns PASS and turn 2 is judged would cover this.Given the existing local server already backs both
attackModelandjudgeLlm(deps()), extending its handler to also serve a turn-generation response looks feasible without a real attacker LLM.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@core/tests/mcpAttackDriver.test.ts` around lines 149 - 180, The MCP attack driver tests are missing coverage for the continued multi-turn path, so add a 2-turn case that keeps turn 1 as PASS and then exercises turn 2. Update the existing test setup in mcpAttackDriver.test.ts using the current deps()/local server so the attack model can return a generated next turn, then assert that buildTurn() takes the turn > 1 branch, generateNextMcpTurn() is used, judgeHint is updated, and judge() receives populated priorTurns. This should complement the existing single-turn and early-stop tests without changing their assertions.core/src/execute/runAll.ts (1)
100-100: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate
config.judgeLlm ?? config.attackerLlmcomputation.Same raw fallback expression is computed independently at Line 100 (baseline scans) and Lines 215-216 (MCP dispatch), in different block scopes. Hoisting it once near Line 69 avoids drift if the fallback logic changes.
♻️ Proposed refactor
const attackModel = resolveModel(config.attackerLlm); const judgeModel = resolveModel(config.judgeLlm ?? config.attackerLlm); + const judgeLlmConfig = config.judgeLlm ?? config.attackerLlm;- const judgeModelConfig = config.judgeLlm ?? config.attackerLlm; + const judgeModelConfig = judgeLlmConfig;? await runMcpAttack( attack, mcpTarget!, attackModel, - config.judgeLlm ?? config.attackerLlm + judgeLlmConfig )Also applies to: 211-216
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@core/src/execute/runAll.ts` at line 100, The fallback selection for the judge model is duplicated as config.judgeLlm ?? config.attackerLlm in both the baseline scans path and the MCP dispatch path. Hoist this logic into a single shared variable in runAll near the top of the function, then reuse that variable wherever judgeModelConfig is needed so the fallback behavior stays consistent and only needs to be changed in one place.core/src/execute/mcpAttackDriver.ts (1)
93-106: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winNon-actionable error messages.
"no toolName in attack spec"and"no turns completed"describe the failure but not the fix. Per coding guidelines, error messages should be actionable.✏️ Proposed fix
- toolError: "no toolName in attack spec", - judge: mcpErrorJudge("no toolName in attack spec"), + toolError: "MCP attack spec is missing a toolName — set `toolName` to one of the target's discovered tools.", + judge: mcpErrorJudge( + "MCP attack spec is missing a toolName — set `toolName` to one of the target's discovered tools." + ),- finalJudge = mcpErrorJudge("no turns completed"); + finalJudge = mcpErrorJudge( + "MCP attack completed no turns — verify `attack.turns` is a positive integer." + );Note:
core/tests/mcpAttackDriver.test.ts(Line 154) asserts the exact string"no toolName in attack spec", so the test would need updating alongside this change.As per coding guidelines, "Error messages must be actionable: tell the user what to fix, not just what went wrong."
Also applies to: 163-181
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@core/src/execute/mcpAttackDriver.ts` around lines 93 - 106, The error strings used by mcpAttackDriver are not actionable, so update the messages in finalize() and the related attack-spec validation path to tell the user what to fix instead of only stating the failure. Adjust the mcpErrorJudge callers and any validation that currently emits "no toolName in attack spec" or "no turns completed" to include the missing requirement or next step. Update core/tests/mcpAttackDriver.test.ts to match the new actionable messages, especially the assertion around the toolName validation.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@core/src/execute/agentAttackDriver.ts`:
- Around line 38-39: The resumed attack flow in AgentAttackDriver is not
preserving prior turn records, so final results can report missing turns or fail
with no turns completed. Update the AgentAttackDriver state and finalization
logic to seed this.turns from the resumed ConversationHistory/initialHistory, or
derive the returned turns from the full history before calling errorJudge and
building the result. Make sure the logic in the constructor/state setup and the
completion path that returns turns uses the combined transcript, not only turns
executed in the current invocation.
- Line 145: The no-turns branch in agentAttackDriver is too vague, so update the
error passed through errorJudge("no turns completed") to tell the user what to
correct or retry. Use the no-turns path in execute/agentAttackDriver and make
the message actionable by naming the likely fix in the current flow, so the
caller gets guidance instead of only a failure state.
In `@core/src/execute/attackRunner.ts`:
- Around line 13-29: The attack runner currently skips cleanup when buildTurn,
execute, record, or shouldEarlyStop throws, so AgentAttackDriver.finalize() is
never reached and target.close() is not guaranteed. Add a dedicated cleanup
method on AttackDriver/AgentAttackDriver for resource shutdown, move
target.close() out of finalize(), and call the new cleanup from a finally block
in the runner so it always executes even on failed turns.
---
Nitpick comments:
In `@core/src/execute/mcpAttackDriver.ts`:
- Around line 93-106: The error strings used by mcpAttackDriver are not
actionable, so update the messages in finalize() and the related attack-spec
validation path to tell the user what to fix instead of only stating the
failure. Adjust the mcpErrorJudge callers and any validation that currently
emits "no toolName in attack spec" or "no turns completed" to include the
missing requirement or next step. Update core/tests/mcpAttackDriver.test.ts to
match the new actionable messages, especially the assertion around the toolName
validation.
In `@core/src/execute/runAll.ts`:
- Line 100: The fallback selection for the judge model is duplicated as
config.judgeLlm ?? config.attackerLlm in both the baseline scans path and the
MCP dispatch path. Hoist this logic into a single shared variable in runAll near
the top of the function, then reuse that variable wherever judgeModelConfig is
needed so the fallback behavior stays consistent and only needs to be changed in
one place.
In `@core/tests/mcpAttackDriver.test.ts`:
- Around line 149-180: The MCP attack driver tests are missing coverage for the
continued multi-turn path, so add a 2-turn case that keeps turn 1 as PASS and
then exercises turn 2. Update the existing test setup in mcpAttackDriver.test.ts
using the current deps()/local server so the attack model can return a generated
next turn, then assert that buildTurn() takes the turn > 1 branch,
generateNextMcpTurn() is used, judgeHint is updated, and judge() receives
populated priorTurns. This should complement the existing single-turn and
early-stop tests without changing their assertions.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 0d51ee90-b2e2-4260-a6d0-051269969425
📒 Files selected for processing (7)
core/src/execute/agentAttackDriver.tscore/src/execute/attackRunner.tscore/src/execute/mcpAttackDriver.tscore/src/execute/runAgentLoop.tscore/src/execute/runAll.tscore/tests/attackRunner.test.tscore/tests/mcpAttackDriver.test.ts
|
Changes made
|
There was a problem hiding this comment.
🧹 Nitpick comments (1)
core/tests/mcpAttackDriver.test.ts (1)
60-83: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse a shared marker for the mock routing.
Matching on the literal"escalate the attack"ties this test to the current prompt text ingenerateNextMcpTurn, so a wording tweak can misroute generator calls to the judge path. Export a shared marker constant or use a dedicated header instead.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@core/tests/mcpAttackDriver.test.ts` around lines 60 - 83, Use a shared routing marker for the mock endpoint instead of matching the literal prompt text in mcpAttackDriver.test.ts. The current body check in the test server depends on the exact wording from generateNextMcpTurn, so update the test and the turn-generation path to share a dedicated constant or request header that clearly identifies attacker turn-generation calls. Keep the routing logic in the server handler and the generateNextMcpTurn prompt aligned through that shared symbol.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@core/tests/mcpAttackDriver.test.ts`:
- Around line 60-83: Use a shared routing marker for the mock endpoint instead
of matching the literal prompt text in mcpAttackDriver.test.ts. The current body
check in the test server depends on the exact wording from generateNextMcpTurn,
so update the test and the turn-generation path to share a dedicated constant or
request header that clearly identifies attacker turn-generation calls. Keep the
routing logic in the server handler and the generateNextMcpTurn prompt aligned
through that shared symbol.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: deb32197-fc49-4c42-8d27-cacb5fa1706b
📒 Files selected for processing (2)
core/src/execute/runAll.tscore/tests/mcpAttackDriver.test.ts
🚧 Files skipped from review as they are similar to previous changes (1)
- core/src/execute/runAll.ts
|
Additional change (resume drops prior turns): live bug via the extension's pause/resume — on resume the driver dropped all pre-pause turns from the reported transcript (and a resume-at-last-turn returned "no turns completed"). |
3 participants
What & why
The attack turn loop existed twice —
runAgentAttackand a 134-linerunMcpAttackembedded inrunAll— sharing a control flow but free to drift.This introduces a Template Method:
runAttackowns the invariant loop(build → execute → record → shouldEarlyStop, then finalize once), and an
AttackDriverfills the kind-specific holes.attackRunner.ts—AttackDriver<TInput,TOutput>+runAttack(a statelessfunction; no needless class).
agentAttackDriver.ts—AgentAttackDriver, agent logic moved verbatim.mcpAttackDriver.ts—McpAttackDriver+runMcpAttack, extracted out ofrunAll(985 → ~840 lines). The MCP mid-turn and final judge blocks(previously duplicated verbatim) are consolidated into one
judge()helper.runAgentLoop.ts—runAgentAttackis now a thin wrapper with anunchanged signature →
runAll/runAllBrowser/SDK/extension untouched.Behavior-preserving
Agent target-error early-stop + judge-once +
target.close(); MCP mid-turn judgeonly when
turns>1 && t<turnswithpriorTurns = history.slice(0,-1),early-stop-on-FAIL,
sanitizeJudgeResult, empty-tool short-circuit, and resultshapes are all reproduced. Guarded by the
runAll.smoke+orchestrator.equivalencedrift tests, plus new
attackRunner(Template-Method contract) andmcpAttackDrivertests — the MCP loop had zero coverage before.Clean-code (hardened after review)
rather than the whole
RunConfig; no per-attack model rebuild.array (dead
lastReplyHook+ alignment ceremony removed).the agent/MCP judges take different model types (
LanguageModelvsLlmConfig);unifying that is hunt-work (PR13).
Deferred (logged in the plan)
Pre-existing latent
target.close()lifecycle (shared-target / no try-finally) andfake-LLM test-harness consolidation.
Verification
full suite 117 tests, 0 fail (3 skips) · typecheck 0 · lint 0 · full build (incl.
browser bundle) passes.
🤖 Generated with Claude Code
Summary by CodeRabbit