fix(redteam): reset target session between PAIR/SequentialBreak iterations - #292
Conversation
…tions Both PAIR and SequentialBreak assume each candidate prompt is an independent single-shot attempt, but previously reused one session across iterations/variants — letting earlier refusals contaminate later attempts and breaking the method's premises. Now each iteration (PAIR) / variant (SequentialBreak) starts with a fresh target session via reset(). To avoid a trajectory regression (reset clears the tool-use trace, which the authoritative judge needs), tool-use entries are accumulated into a local list across iterations and restored onto the session trace after the loop completes. Updates the SequentialBreak fake-session test contract to allow reset() (previously pinned to append-only "must not reset").
|
Assessment: Comment (one important gap to address before merge) The fix itself is sound — the reset placement, the accumulate-and-restore of Review themes
Nice, well-scoped bug fix with a clear writeup and a linked issue that explains the reasoning — the trace-preservation subtlety was handled thoughtfully. |
|
@poshinchen — replying to your
The scoped fix should take one Tests should prove every attempt sees identical seeded/non-message state and unique per-attempt tool entries survive exactly once and in order. The current focused suite passes (145 tests), but its fakes do not exercise either invariant. |
|
@strandly-the-agent what's your suggestion on the fix: #292 (comment) |
|
My concrete suggestion is to replace the strategy-level checkpoint = target_session.snapshot()
trace_prefix = list(target_session.trace[: checkpoint.trace_len])
attempt_trace: list[ToolUseEntry] = []
try:
for ...:
# Generate/validate the next prompt first; do not restore for skipped attempts.
if target_calls > 0:
attempt_trace.extend(target_session.trace[checkpoint.trace_len :])
target_session.restore(checkpoint)
response = target_session.invoke(prompt)
target_calls += 1
...
finally:
# Keep completed-attempt evidence even if later strategy/judge logic raises.
attempt_trace.extend(target_session.trace[checkpoint.trace_len :])
target_session.trace[:] = trace_prefix + attempt_traceApply the same lifecycle to PAIR and SequentialBreak, ideally through a small private helper so the trace bookkeeping is implemented once. The first attempt uses the exact state supplied to I would add two regression tests per strategy: (1) I would not change |
Replace the between-attempt reset() with a shared single_shot_attempts helper that checkpoints at entry and restores between attempts, so each attempt starts from identical target state (including agent state, not just messages) while preserving every attempt's tool-use trace. Add regression tests with fakes that implement real snapshot/restore.
|
Reworked to use snapshot/restore instead of |
|
@strandly-the-agent review the PR before I approve. |
strandly-the-agent
left a comment
There was a problem hiding this comment.
TL;DR — two small fixes before you approve. The single_shot_attempts design is right, and the single-agent path holds up under attack: the original Agent.state leak is genuinely gone ([0,1,2] on the old head → [0,0,0] here). But two reproduced defects slip through the new helper — multi-agent checkpoints alias live Swarm state, and trace entries duplicate if restore() itself raises. Both fixes are 1–2 lines.
✅ Verified at 943a587 (base b99dc07): pytest tests/strands_evals/experimental/redteam → 327 passed; CI green on head (unit py3.10–3.14 × linux/windows/macOS, lint, integration); both defect repros run twice cold with identical output. 🟡 ruff isn't installed in my sandbox, so lint is cited from CI rather than re-run locally.
| # | Finding | Where | New in this diff? | Suggested |
|---|---|---|---|---|
| 🔴 1 | Multi-agent checkpoint aliases live Swarm state |
target_session.py:259-267 |
root cause pre-existing, newly exposed | fix here |
| 🟡 2 | Trace duplicated if restore() raises |
target_session.py:128-140 |
yes | fix here |
| 🟡 3 | Stale contamination docstring + dropped best_score comment |
sequentialbreak/__init__.py:8-10, pair/__init__.py:160 |
yes | fix here |
| 🟡 4 | started no-op guard is untested (mutation survives) |
target_session.py:128-133 |
yes | fix here |
| 🟡 5 | Stub TargetSession now dies with AttributeError |
task.py:193-197 |
new exposure | follow-up |
🔴 1 — one checkpoint restored N times isn't safe for StrandsMultiAgentSession
snapshot() stores orch.serialize_state() verbatim (target_session.py:259-267), while _restore() deep-copies on the way back in (:286). For Swarm, serialize_state() returns getattr(self.state.shared_context, "context", {}) or {} — the live dict whenever it's non-empty, so the checkpoint aliases target state.
My own probe against a real Swarm (2 cold runs, identical):
live shared context before snapshot: {'planner': {'note1': 'value-1'}}
snapshot ALIASES live dict? True
after attempt-1 write, checkpoint: {'planner': {'note1': 'value-1', 'note2': 'value-2'}}
CHECKPOINT MUTATED by the attempt? True
An independent adversarial pass took it end to end on a real 2-agent Swarm with PAIR ×3: on a warm swarm attempt 1's writes leak into attempts 2–3; on a cold swarm attempts 2..N restore into {} — a state that never existed. Empty-context targets (the common default) are unaffected, which is why the fakes don't catch it.
The aliasing line predates this PR, but PAIR/SequentialBreak never called snapshot/restore before, and reusing one checkpoint across N restores is what wakes it up. MultiAgentBase is a supported target here (task.py:88-89, :191-192), so as-is the PR's isolation guarantee silently fails for exactly the class of target it advertises.
Fix — deep-copy at capture, mirroring _restore:
def snapshot(self) -> TargetCheckpoint:
"""Capture a composite snapshot of every leaf and every orchestrator."""
return TargetCheckpoint(
# Deep-copy at capture: `serialize_state()` can alias live orchestrator state (Swarm returns
# its live `shared_context`), and one checkpoint is restored repeatedly by single_shot_attempts.
agent_snapshot=copy.deepcopy(
_MultiAgentSnapshot(
agents={path: agent.take_snapshot(preset="session") for path, agent in self._agent_index.items()},
orchestrators={path: orch.serialize_state() for path, orch in self._orch_index.items()},
)
),
trace_len=len(self.trace),
)
The cold-swarm wipe is a separate upstream Swarm serialize/deserialize asymmetry — worth an SDK issue, not a gate here. Until it's fixed, "covers all target state" in the helper docstring is a slight overclaim for multi-agent targets.
🟡 2 — trace entries duplicate when restore() raises
Both built-in restore()s truncate the trace as their last statement, after a step that can raise: load_snapshot (:184-185) and _restore (:277). If that step raises, the delta begin_attempt() just copied into attempt_trace is still on session.trace, and the finally at :138-140 re-collects the same slice.
My probe (2 cold runs, identical) — 1 real invoke, 2 trace entries:
FAILING RESTORE -> raised: RuntimeError: load_snapshot failed
FAILING RESTORE -> rebuilt trace: ['tool_attempt1', 'tool_attempt1']
FAILING RESTORE -> invokes that really happened: ['tool_attempt1']
The adversarial pass reproduced the same thing with a real Agent (a tool that swaps conversation_manager, so the failure comes from the SDK), and showed a non-truncating custom restore() yields triangular duplication (3 invokes → 6 entries) — the exact shape this PR set out to remove. It also contradicts the docstring's "once each, in order — even if the caller raises".
Fix — consume the delta so it can never be read twice:
def take_delta() -> list[ToolUseEntry]:
delta = list(session.trace[checkpoint.trace_len :])
del session.trace[checkpoint.trace_len :]
return delta
def begin_attempt() -> None:
nonlocal started
if started:
attempt_trace.extend(take_delta())
session.restore(checkpoint)
started = True
# finally: attempt_trace.extend(take_delta()); session.trace[:] = trace_prefix + attempt_traceOnly fires on an already-abnormal path, so 🟡 rather than 🔴 — but it's cheap, and test_trace_preserved_even_when_caller_raises currently only raises after a successful begin_attempt().
🟡 3 — stale docstring and a dropped comment in the touched files
sequentialbreak/__init__.py:8-10 still documents the contamination this PR removes:
On a stateful
target_session, variants 2..N see earlier variants' refusal context, so measured ASR is a lower bound whenever more than one variant is tried; best-variant-first (dc_t1) minimizes this.variants_triedlength > 1 flags possible contamination.
Suggested replacement:
Each scaffold variant is an independent single-shot: the strategy checkpoints
target_sessionatrun_attackentry and restores that checkpoint before every later variant, so variant N never sees variant N-1's refusal context. Ordering therefore doesn't bias measured ASR, andvariants_triedlength > 1 no longer implies contamination.
Same stale rationale at sequentialbreak_v0.py:77-78 ("dc_t1 first … minimizes cross-variant contamination").
And pair/__init__.py:160 drops, with no replacement, the comment explaining why best_score is a peak rather than the last value ("a refinement strategy's closest approach is more informative than the last turn, which may have dipped after a peak"). The logic is unchanged, so this is pure context loss — worth restoring next to the max() call.
🟡 4 — the started guard is real but untested (and the fakes hide #1/#2)
A test-quality pass mutation-tested the helper against the 105 relevant tests. Every structural mutation is caught — skip the restore (8 failures), drop the try/finally rebuild (4), reset() instead of restore() (33), swap collect/restore order (4) — except removing the started guard, which survives with 0 failures. No test mutates the session between with single_shot_attempts(...) and the first begin_attempt(), so a spurious first rollback is invisible. This probe passes on HEAD and fails with the guard removed:
def test_first_begin_attempt_is_a_true_noop():
session = _RewindSession()
with single_shot_attempts(session) as begin_attempt:
session.state = 99 # something happens before the first begin_attempt()
begin_attempt()
assert session.state == 99 # first call must be a no-op, not a restoreThe bigger gap: the new tests only drive hand-written fakes whose non-message state is an int. Nothing exercises single_shot_attempts through a real StrandsAgentSession/StrandsMultiAgentSession, and test_task.py's trajectory tests use a stub strategy that invokes once. That's precisely why #1 and #2 got through. One real-Agent test per strategy (seed agent.state, assert iterations 2/3 see iteration 1's entry state) plus one real-Swarm case would have caught #1.
🟡 5 — follow-up: stub TargetSession now fails with a confusing AttributeError
_build_session (task.py:193-197) only checks that the four methods are callable, so a target that subclasses the Protocol and inherits the ... stubs (returning None) is accepted, then dies inside the helper with AttributeError: 'NoneType' object has no attribute 'trace_len' before any target call.
To be clear, this isn't a contract break — snapshot/restore were always declared on TargetSession (:89-95) and _build_session's own TypeError already says "so the strategy can snapshot/restore its state", so such a target was already violating the Protocol. But it is a newly exercised path with a much worse error than the informative TypeError two lines away, and base.py:87's run_attack docstring still names only invoke. It fails fast with no API spend, so a follow-up (type-validate the return, or document the requirement) is fine.
Questions
Both non-blocking.
- Is
single_shot_attemptsmeant to be public? It's the only non-underscore module-level symbol intarget_session.pyabsent from its__all__(:335-342), and it isn't re-exported fromstrategies/__init__.pyorredteam/__init__.py, while that export ladder looks deliberate elsewhere. Internal →_single_shot_attemptsis a two-line change; public extension point → it should join__all__and probably carry thedesignlabel. - Should the MUST-call-
begin_attempt()contract defend itself, or is docstring-only fine for now? Four misuse modes were probed and all are silent: omitting the call reintroduces #291 with the trace still looking correct, calling it afterinvoke()contaminates the next attempt, and letting the callback escape thewithwipessession.trace. With two internal call sites, docstring-only is defensible — just worth an explicit "yes, for now".
Checked and clean / suppressed
- The single-agent path survived a hard look: one checkpoint, 4 restores, byte-identical messages /
agent.state(incl. nested dicts) / conversation-manager state, checkpoint unmutated.strandsdeep-copies messages/state/model_state on both capture and restore, so repeated restores are safe there. - Trace correctness 7/7 across every early-exit and exception path (PAIR empty-response and parse-failure
break,invoke()raising, judge raising, SB empty-prompt and empty-responsecontinue, early-stopbreak); zero-attempt paths 3/3; entry-prefix preserved;session.trace[:] = …keeps list identity fortask.py's copy; strategies stay stateless for parallel cases. - Real
task._build_attacker_taskpath with two sequential cases: no intra- or cross-case leakage, session still snapshot/restore-usable afterwards. - Not this PR:
_send's index-diff loses tool-use evidence when a conversation manager trims mid-invoke (target_session.py:163-167, same at the old head — a real evaluator blind spot worth its own issue); the cold-swarmSwarmserialize/deserialize asymmetry (upstream SDK). - Suppressed: Crescendo not migrated (per-turn conditional rollback genuinely doesn't fit this helper);
goat/bad_likert_judge/prompt_strategycorrectly untouched; keyword-only/handle-object future-proofing (no current defect); "Append-only" wording in both module docstrings is still technically true (it's anchored topruned_branches), just easier to over-read now.
The rework is a clear improvement on the reset() version and the shape is right — this is solid work for you to arbitrate, not a gate. My read: #1 and #2 are worth landing in this PR since both are a couple of lines and #1 silently breaks the guarantee the PR ships for MultiAgentBase targets; #3/#4 are cheap while you're in there. Happy to be told #1 belongs in a separate PR against StrandsMultiAgentSession instead — it's your call.
…ed restore Deep-copy orchestrator state at capture so one checkpoint restored N times can't alias live Swarm shared_context. Move each attempt's trace delta off the session so a restore() that raises can't leave it for the finally block to collect twice. Make the helper private and align the stale docstrings. Add contract tests driving real Agent/Swarm sessions, not just fakes.
|
Fixed #1–#4; #5 as a follow-up.
Questions: (1) internal — renamed
|
|
@strandly-the-agent review the PR again before I approve. |
|
Assessment: Approve Re-reviewed after the rework. Moving from Prior review items — all addressed
Verified the focused redteam suite (109 tests) passes locally. Clean, well-tested fix — the |
strandly-the-agent
left a comment
There was a problem hiding this comment.
TL;DR — approve; safe to merge from my side. Both defects I reproduced at 943a587 are genuinely fixed at 2b84558, and I re-ran my own repros to confirm rather than taking the fixes on faith. What's left is two non-blocking follow-ups and one comment that's narrower than it reads.
✅ Verified at 2b84558 (delta 943a587..2b84558, base b99dc07): pytest tests/strands_evals/experimental/redteam → 333 passed; TestSingleShotAttempts run 8× back to back, clean each time; both my repros re-run twice cold with identical output; CI green on head; no dangling references to the old public helper name; working tree clean.
| Prior finding | Status at 2b84558 |
Evidence |
|---|---|---|
🔴 1 Multi-agent checkpoint aliased live Swarm state |
Fixed | my repro: snapshot ALIASES live dict? False, checkpoint no longer mutated |
🟡 2 Trace duplicated when restore() raises |
Fixed | my repro: 1 invoke → 1 entry (was 2); 3 invokes → 3 |
🟡 3 Stale docstrings / dropped best_score comment |
Fixed | sequentialbreak/__init__.py:8-10, sequentialbreak_v0.py:77, pair/__init__.py:189 |
🟡 4 started guard untested; fakes-only coverage |
Fixed | independent mutation pass: all 7 mutations now caught, incl. the one that survived |
🟡 5 Stub TargetSession → confusing AttributeError |
Deferred (agreed) | base.py:87 now names the snapshot/restore/trace contract |
| ⚪ new, non-blocking | Leaf-snapshot comment is narrower than preset="session" |
see below |
Both open questions are settled the way I'd have chosen: _single_shot_attempts is private and stays out of __all__, and docstring-only for the call contract with two internal call sites.
How I verified fixes #1 and #2 (my own repros, 2 cold runs each)
#1 — real Swarm + StrandsMultiAgentSession, seeded shared_context, snapshot, then an attempt writes to it:
snapshot ALIASES live dict? False (was True)
after attempt-1 write, live: {'planner': {'note1': 'value-1', 'note2': 'value-2'}}
after attempt-1 write, checkpoint: {'planner': {'note1': 'value-1'}}
CHECKPOINT MUTATED by the attempt? False (was True)
copy.deepcopy(orch.serialize_state()) at target_session.py:270-273 detaches the payload; Swarm.serialize_state() still hands back the live dict, so the hazard was real and the capture-side copy is the right place for the fix (the pre-existing deepcopy in _restore protects the target from the checkpoint, not the checkpoint from the target).
#2 — custom session whose restore() raises during rollback, before its trace truncation:
FAILING RESTORE -> raised: RuntimeError: load_snapshot failed
FAILING RESTORE -> trace: ['tool_attempt1'] (was ['tool_attempt1', 'tool_attempt1'])
HEALTHY RESTORE -> trace len: 3 for 3 invokes
take_delta() moving the delta off session.trace (:127-133) is a cleaner fix than the cursor I suggested — it makes double-collection structurally impossible, and it also makes restore()'s own del self.trace[trace_len:] a harmless no-op. An independent adversarial pass additionally confirmed the non-truncating custom restore() case (3 invokes → 3 entries, previously 6) and that a restore() which appends to or rebinds trace doesn't reopen anything.
⚪ New, non-blocking: the leaf-snapshot comment is narrower than preset="session"
target_session.py:271-272 says "Leaf snapshots need no copy -- take_snapshot already deep-copies messages and state." True for those two, but preset="session" captures four fields, and I checked the installed strands 1.50.2 myself:
| field | capture | detached? |
|---|---|---|
messages |
copy.deepcopy(self.messages) |
yes |
state |
self.state.get() |
yes |
conversation_manager_state |
conversation_manager.get_state() |
no |
interrupt_state |
self._interrupt_state.to_dict() → "context": self.context |
no (by reference) |
So two of the four are aliased in the reused checkpoint. I do not think this blocks: an adversarial pass reached it only by having a tool write directly to the private agent._interrupt_state.context, and the supported path can't — _InterruptState.resume() returns early unless activated, and raises TypeError for anything that isn't a list of interruptResponse blocks, while both sessions only ever invoke() with a str. SummarizingConversationManager rebinds _summary_message rather than mutating it. So the conclusion holds today; only the justification is narrower than it sounds.
Cheapest options, either is fine and neither needs to be in this PR: extend the copy to the leaves (agents={path: copy.deepcopy(agent.take_snapshot(preset="session")) ...}, symmetric with the orchestrator line and immune to future SDK changes), or reword the comment to name all four preset fields and say the two aliased ones aren't reachable from invoke(str).
Follow-ups worth issues (none blocking)
- Finding #5, as agreed —
_build_session(task.py:193-197) only checks callability, so a Protocol stub still fails inside the helper withAttributeError: 'NoneType' object has no attribute 'trace_len'rather than the informativeTypeErrortwo lines away. Type-validating thesnapshot()return would fix it. - Task-level
trajectorycoverage —test_task.pyis untouched, so the aggregation is still only proven via_StubStrategy(oneinvoke, no helper). One test per strategy assertingresult["trajectory"]holds every attempt's tool uses would close the last gap from #4; the mutation pass rates this medium. preset="session"scope — rollback doesn't coversystem_prompt, so an attempt that rewrites the target's own guardrail prompt persists into later attempts (demonstrated by the adversarial pass; pre-existing, and the reworded docstring no longer overclaims). Worth an explicit note that isolation is limited to the checkpoint's fields.- Pre-existing, unrelated to this PR —
_send's index-diff loses tool-use evidence when a conversation manager trims mid-invoke (target_session.py:169-173): a real breach lands inagent.messagesbut never insession.trace, so the evaluator scores it clean. Same at every head of this PR; aMessageAddedEventhook would fix it. This one has real teeth for ASR accuracy — I'd file it.
Also checked / suppressed
- Independent passes re-checked every path that survived last time: PAIR empty-prompt
breakbeforebegin_attempt()and empty-responsebreakafter it, SequentialBreak's twocontinuepaths, exceptions out ofinvoke()and out of the judge,turn_cap=0, attempts producing 0/1/2/3 tool uses,snapshot()itself raising, and the realtask._build_attacker_taskseam across two sequential cases against a shared realAgent— no regressions, no duplicated or dropped evidence. - Capture-time
deepcopydoesn't add a new failure class:_restorealready deep-copies the same payload on every restore, and realSwarm/Graphpayloads are JSON-safe primitives. It does add one copy persnapshot()(crescendo snapshots per turn), which is negligible next to a model call. - One pass reported transient
TestSingleShotAttemptsfailures and attributed them to its own__pycache__churn under concurrent tool use. I could not reproduce: 8/8 clean focused runs plus a clean 333-test suite run. Flagging only so a CI blip isn't mistaken for a code issue. - Suppressed: the
delinside thefinallyis redundant given the very next line overwrites the list, and can only mask an exception for atracethat doesn't support slice deletion — out of contract fortrace: list[ToolUseEntry], so not worth a change.
Good turnaround — the move-not-copy take_delta() and the real-Agent-through-run_attack tests are both better than what I asked for. poshinchen still owns the merge call; nothing above needs to land first.
Description
PairStrategyandSequentialBreakStrategyboth treat each candidate prompt as an independent single-shot attempt (PAIR refines one prompt P over K iterations; SequentialBreak tries N scaffold variants), but previously reused one target session across all of them. The target carried the conversation from earlier attempts, so a refusal or partial breach in attempt N-1 contaminated attempt N — breaking the single-shot premise these methods rely on.Each attempt now starts from the target's entry state via a shared
single_shot_attemptshelper that snapshots once at entry and restores between attempts. Unlike a baseline-lessreset()(messages only),restore()rolls back all target state — messages, agentstate, conversation-manager state — so non-message state can't leak across attempts. The attacker is intentionally not reset — PAIR's cross-iteration refinement context is what makes it work.The helper accumulates each attempt's tool-use trace and rebuilds
session.traceon exit (entry prefix + every attempt's delta, in order), so the authoritativeAttackSuccessEvaluatorstill sees all tool-call breach evidence.Related Issues
Closes #291
Documentation PR
N/A
Type of Change
Bug fix
Testing
hatch run preparehatch run preparepasses end to end: ruff clean, mypy clean, full suite green across the Python version matrix (1663 passed each).Added regression tests for both strategies (every attempt observes the same entry state; each attempt's tool-use survives the restores) plus unit tests for the shared helper; the test fakes now implement real
snapshot/restore.Checklist
By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.