fix: cancellation turn-granularity, extension report gaps, and cache-aware cost pricing - #236
Open
arunSunnyKVS wants to merge 3 commits into
Open
Conversation
The CLI's AbortSignal was only checked between attacks/evaluators, so a long multi-turn attack (e.g. turns: 100) couldn't be interrupted until every turn finished. The extension had a parallel gap: its own stop signal wasn't recognized by runAllBrowser, so a cancelled run returned no report at all — dropping token/cost data and leaving a fake blank turn in the HTML report. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
WalkthroughThe change adds cancellation to attack execution, preserves partial stopped-run data, tracks cache-tier token usage, calculates tiered costs, and displays cost and final transcript details in the extension report. ChangesCancellation and cost reporting
Estimated code review effort: 3 (Moderate) | ~30 minutes Sequence Diagram(s)sequenceDiagram
participant Extension as Extension orchestrator
participant Browser as runAllBrowser
participant Evaluator as evaluatorLoop
participant Runner as runAttack
participant Tracker as TokenTracker
participant Popup as Extension popup
Extension->>Browser: request evaluator run
Browser->>Evaluator: run evaluators with AbortSignal
Evaluator->>Runner: start MCP or agent attack
Runner-->>Evaluator: stop before next turn and finalize
Evaluator->>Tracker: record token usage
Browser-->>Extension: run_stopped with partial results
Extension->>Popup: provide token usage and evaluator results
Popup->>Popup: calculate and display estimated cost
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 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: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
runners/extension/popup.js (1)
1423-1443: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winKeep
lastlimited to turns with an assistant response.
turnsForReport()outputs a turn when either the user or assistant entry exists. If an interrupted evaluator leaves an in-flight user message inraw.transcript,lastcan come from a turn whoseresponseis empty, andturnscan include that incomplete entry. Filter turns with a non-empty assistant response before selectinglastand before returningturns.🤖 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 `@runners/extension/popup.js` around lines 1423 - 1443, Update the turn handling in turnsForReport to retain only turns with a non-empty assistant response, then select last from that filtered collection and return it as turns. Preserve the existing fallback to empty prompt and response when no completed turns remain.
🧹 Nitpick comments (1)
runners/extension/popup.js (1)
14-15: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueImport the cost helpers directly from their source module.
estimateRunCostandformatUsdare defined incore/src/pricing/estimateCost.ts, so import them from that file instead of re-exporting them throughcore/src/browser.tsand then bundling from./dist/core.bundle.js.🤖 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 `@runners/extension/popup.js` around lines 14 - 15, Update the imports for estimateRunCost and formatUsd in the popup module to reference their defining source module directly, core/src/pricing/estimateCost.ts, instead of importing through core/src/browser.ts or ./dist/core.bundle.js. Keep both helper usages unchanged.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/runAllBrowser.ts`:
- Around line 217-222: Update the extension orchestrator’s event handling to add
a run_stopped branch alongside the existing attack_start and attack_done
branches. Forward the stopped event by calling broadcastProgress with its
reason/state so the popup displays the partial run reason emitted by
runAllBrowser.
---
Outside diff comments:
In `@runners/extension/popup.js`:
- Around line 1423-1443: Update the turn handling in turnsForReport to retain
only turns with a non-empty assistant response, then select last from that
filtered collection and return it as turns. Preserve the existing fallback to
empty prompt and response when no completed turns remain.
---
Nitpick comments:
In `@runners/extension/popup.js`:
- Around line 14-15: Update the imports for estimateRunCost and formatUsd in the
popup module to reference their defining source module directly,
core/src/pricing/estimateCost.ts, instead of importing through
core/src/browser.ts or ./dist/core.bundle.js. Keep both helper usages unchanged.
🪄 Autofix
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 Plus
Run ID: e5ddf7a0-b139-4175-a656-5bff67a48eb7
📒 Files selected for processing (9)
core/src/execute/attackRunner.tscore/src/execute/evaluatorLoop.tscore/src/execute/mcpAttackDriver.tscore/src/execute/runAgentLoop.tscore/src/execute/runAllBrowser.tscore/tests/attackRunner.test.tsrunners/extension/orchestrator.jsrunners/extension/popup.htmlrunners/extension/popup.js
Comment on lines
+217
to
+222
| if ((err as { code?: string })?.code === "OPFOR_STOP") { | ||
| stopReason = err instanceof Error ? err.message : "Run stopped by user."; | ||
| notify({ type: "run_stopped", reason: stopReason }); | ||
| pushPartialResult(stopReason); | ||
| break evaluatorLoop; | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ast-grep outline runners/extension/orchestrator.js --items all
rg -n -C 6 'run_stopped|onProgress|broadcastProgress|setRunStatus|finalizeUserInterruption' \
runners/extension/orchestrator.jsRepository: KeyValueSoftwareSystems/agent-opfor
Length of output: 12727
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "## runAllBrowser relevant section"
sed -n '180,235p' core/src/execute/runAllBrowser.ts
echo
echo "## extension run_stopped occurrences"
rg -n -C 8 'run_stopped|OPFOR_STOP|type === "attack_done"|emit\(|create.*Run|runAll\(' .
echo
echo "## executeAdaptiveRedTeamRun message send around stop"
sed -n '840,905p' runners/extension/orchestrator.jsRepository: KeyValueSoftwareSystems/agent-opfor
Length of output: 50394
Forward run_stopped from core to the extension popup.
runAllBrowser emits run_stopped for user-cancel/error paths, but runners/extension/orchestrator.js only calls broadcastProgress() for attack_start and attack_done. Add a matching branch and broadcast the stopped state so the extension UI can show the partial run reason.
🤖 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/runAllBrowser.ts` around lines 217 - 222, Update the
extension orchestrator’s event handling to add a run_stopped branch alongside
the existing attack_start and attack_done branches. Forward the stopped event by
calling broadcastProgress with its reason/state so the popup displays the
partial run reason emitted by runAllBrowser.
jithin23-kv
approved these changes
Aug 6, 2026
Cost estimation charged every input token at the full input rate, but `inputTokens` is the inclusive total — it already contains the cached tokens providers bill far more cheaply. The token counter was also discarding the cache split the AI SDK already reports. Input is now divided across its tiers and each priced at its own rate. Validated against LiteLLM billing for a 24-request run: previously reported $0.070749, actually billed $0.057939, now reports $0.057939 — every request agreeing to 12 decimal places. Reported fields are unchanged; only the cost figure moves. A run whose provider reports no cache split prices exactly as before, and a tier with no published rate falls back to the full input rate rather than to free. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
arunSunnyKVS
changed the title
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/tokenTracker.ts`:
- Around line 99-124: Update the token usage transform around the
inputTokenDetails schema and its noCache calculation to always derive noCache
from inputTokens minus cacheRead and cacheWrite, ignoring any provider-supplied
noCacheTokens override. Validate and reject inputs where cacheRead plus
cacheWrite exceeds inputTokens, and add regression tests covering both the
override case and the over-limit case.
🪄 Autofix
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 Plus
Run ID: b2540b77-5e97-4f72-bf8d-999e9b831573
📒 Files selected for processing (5)
core/src/execute/tokenTracker.tscore/src/pricing/estimateCost.tscore/src/pricing/types.tscore/tests/pricing.test.tscore/tests/tokenTracker.test.ts
The transform read `noCacheTokens` verbatim, so a provider reporting inputTokens 100 with noCache 100 and cacheRead 50 produced a 150-token split against a 100-token call — and estimateRunCost billed all 150. The invariant was asserted in the docs and tests but never enforced. Derive the fresh count from inputTokens instead. For a well-formed provider the two agree (the AI SDK builds inputTokens as the sum), so this costs nothing and makes the invariant hold by construction. A split claiming more cached tokens than there was input can't be divided at all, so it is dropped and the call prices at the full input rate. Addresses CodeRabbit review on #236. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Four bugs, surfaced while testing
opfor runwith Ctrl+C, the extension's Stop button, and comparing reported cost against actual billing.Cancellation (1) —
AbortSignalcancellation was only checked between attacks/evaluators, never between turns within one multi-turn attack. A long-running attack (e.g.turns: 100) was atomic — Ctrl+C wouldn't stop it until every turn finished, contradicting the "finishes in-flight attack" promise in the CLI's own interrupt message.Extension reporting (2, 3) — a cancelled evaluator with zero completed turns rendered a fake "1 turn" with a blank transcript, because
toResultViewModelhardcodeddetail: { prompt: "", response: "" }. And cancelled runs showed no token usage or cost at all:DomTarget.send()throws anErrortaggedcode: "OPFOR_STOP"on user cancel, whichrunAllBrowserdidn't recognise as a clean stop condition, so it fell through to an unhandledthrowand returned no report — silently dropping the run's token totals. Separately,costwas never plumbed through the extension even on successful runs.Cost accuracy (4) — cost estimation charged every input token at the full input rate. But
inputTokensis the inclusive total: it already contains the cached tokens that providers bill far more cheaply. The token counter was also discarding the cache split the AI SDK already reports. On a real 24-request run this overstated cost by 22%.Solution
AbortSignalone level deeper, intorunAttack()'s per-turn loop, so it's checked before every turn and breaks out the same way the existing early-stop-on-target-error path does — still callingfinalize()so a partial transcript is judged and reported.toResultViewModelto derivedetailfrom the last completed turn, matching the contract core's ownbuildReport.tsalready establishes for the CLI/SDK path.OPFOR_STOP-tagged-error branch torunAllBrowser's catch block (recognised structurally via.code, since core can't import the extension's error class) so a user-cancelled run returns a proper partial report. ThreadedtokenUsageByModelthroughorchestrator.js, aggregated it inpopup.jsvia the already-bundledestimateRunCost, and wired the result into both the downloaded report and a new Cost stat card.inputTokenDetails(so DeepSeek'sprompt_tokens_details.cached_tokensand Anthropic'scache_read_input_tokensboth work without provider-specific code).Validation against real billing
The cost change was verified against LiteLLM's own spend logs for a 24-request run:
Every one of the 24 requests agrees to 12 decimal places. The old formula reported $0.070749 for the same run — a 22.1% overstatement.
Changes
Cancellation
core/src/execute/attackRunner.ts—runAttack()accepts an optionalsignal, checked at the top of each turncore/src/execute/runAgentLoop.ts,core/src/execute/mcpAttackDriver.ts— forwardsignalthrough torunAttackcore/src/execute/evaluatorLoop.ts— pass the already-availablesignalinto both attack call sitesExtension reporting
core/src/execute/runAllBrowser.ts— recogniseOPFOR_STOP-tagged errors as a clean stop condition, same asTargetStopErrorrunners/extension/popup.js— fix thedetailfallback for zero-turn cancelled results; aggregatetokenUsageByModeland derivecost; populate the new Cost statrunners/extension/popup.html— add astatCostcard next to the existingstatTokenscardrunners/extension/orchestrator.js— threadtokenUsageByModelthrough all three result paths (success, error, stopped)Cost pricing
core/src/execute/tokenTracker.ts— capture the cache split instead of discarding it; carry it to the per-model bucketscore/src/pricing/estimateCost.ts— price each input tier at its own ratecore/src/pricing/types.ts— the "Not applied yet" doc comments, now appliedTests —
core/tests/attackRunner.test.ts,core/tests/pricing.test.ts,core/tests/tokenTracker.test.ts(18 new)Backward compatibility
cacheRead/cacheWritefall to 0,noCacheabsorbs the total, and the formula collapses to the old one).0is honoured (??not||) — DeepSeek genuinely publishescw: 0, and there's a test pinning this.Issue
N/A
How to test
npm run buildturnMode: "multi"and a highturnscount (e.g. 20+), Ctrl+C mid-attack, confirm it stops within one turn's latency instead of waiting for all turns.npm test— 294 tests pass (18 new).npm run typecheck/npm run lint— clean.Screenshots
N/A
Summary by CodeRabbit
New Features
Bug Fixes