diff --git a/AGENTS.md b/AGENTS.md index 647d6cc0..22e21194 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -240,7 +240,7 @@ Shared helpers live in `detectors/utils.py`. Types live in `types/detector.py`. ### Docstrings -Use Google-style docstrings for public functions, classes, and modules. +Use Google-style docstrings for public functions, classes, and modules. Use single backticks (`` `foo` ``) for inline code references, not RST-style double backticks (`` ``foo`` ``). ### Logging Style diff --git a/src/strands_evals/experimental/redteam/strategies/__init__.py b/src/strands_evals/experimental/redteam/strategies/__init__.py index 5718b11a..b9ad8823 100644 --- a/src/strands_evals/experimental/redteam/strategies/__init__.py +++ b/src/strands_evals/experimental/redteam/strategies/__init__.py @@ -2,7 +2,13 @@ from .crescendo import CrescendoStrategy from .prompt_strategy import PromptStrategy from .prompt_strategy.gradual_escalation import get_template as _gradual_escalation_template -from .target_session import StrandsAgentSession, TargetCheckpoint, TargetSession, ToolUseEntry +from .target_session import ( + StrandsAgentSession, + StrandsMultiAgentSession, + TargetCheckpoint, + TargetSession, + ToolUseEntry, +) # Ready-made strategy instances users can pass to RedTeamExperiment(attack_strategies=[...]). # Strategy instances are shared across cases, so each must keep `__init__` for static @@ -19,6 +25,7 @@ "CrescendoStrategy", "PromptStrategy", "StrandsAgentSession", + "StrandsMultiAgentSession", "TargetCheckpoint", "TargetSession", "ToolUseEntry", diff --git a/src/strands_evals/experimental/redteam/strategies/target_session.py b/src/strands_evals/experimental/redteam/strategies/target_session.py index 1431c5e5..be29834c 100644 --- a/src/strands_evals/experimental/redteam/strategies/target_session.py +++ b/src/strands_evals/experimental/redteam/strategies/target_session.py @@ -1,25 +1,28 @@ """Target session protocol and implementations. A :class:`TargetSession` is the handle a strategy uses to talk to the system -under test. It replaces the older opaque ``call_target: Callable[[str], str]``: +under test. It replaces the older opaque `call_target: Callable[[str], str]`: besides sending a message, a session exposes snapshot/restore so a strategy can roll the target back to an earlier state (e.g. Crescendo backtracking past a refusal). The task runner builds a :class:`StrandsAgentSession` (which wraps a -``strands.Agent`` and is rewindable via the SDK snapshot API) from the agent the -experiment was given. A future multi-agent session would be a second -implementation; because the contract is a ``Protocol``, a customer can also -supply their own without subclassing anything here. +`strands.Agent` and is rewindable via the SDK snapshot API) from the agent +the experiment was given, or a :class:`StrandsMultiAgentSession` (which wraps a +`strands.multiagent.MultiAgentBase` such as a Graph or Swarm). Because the +contract is a `Protocol`, a customer can also supply their own without +subclassing anything here. """ from __future__ import annotations +import copy import logging from dataclasses import dataclass from typing import Any, Protocol, TypedDict from strands import Agent, Snapshot +from strands.multiagent.base import MultiAgentBase, Status from strands.types.content import Message logger = logging.getLogger(__name__) @@ -35,9 +38,9 @@ class ToolUseEntry(TypedDict): """A single tool invocation captured into a session's trace. The shape every session implementation produces and every consumer (the - report, the ``AttackSuccessEvaluator``) reads, declared once here. Named - ``ToolUseEntry`` rather than ``ToolUse`` to avoid colliding with the SDK's - ``strands.types.tools.ToolUse`` content block, which this is derived from but + report, the `AttackSuccessEvaluator`) reads, declared once here. Named + `ToolUseEntry` rather than `ToolUse` to avoid colliding with the SDK's + `strands.types.tools.ToolUse` content block, which this is derived from but is not the same shape. """ @@ -48,29 +51,29 @@ class ToolUseEntry(TypedDict): def _tool_uses_in(messages: list[Message]) -> list[ToolUseEntry]: """Extract tool-use entries from Strands messages, tolerating schema drift. - The single place that knows the Strands ``message -> content[] -> toolUse`` + The single place that knows the Strands `message -> content[] -> toolUse` schema, so the brittle dict-walking lives in exactly one tested function. It never raises on a malformed shape: a non-mapping message or content block is - skipped, and a ``toolUse`` block whose value isn't a mapping is recorded as a - ``MALFORMED_TOOL_NAME`` placeholder. Two reasons not to abort: + skipped, and a `toolUse` block whose value isn't a mapping is recorded as a + `MALFORMED_TOOL_NAME` placeholder. Two reasons not to abort: - - One malformed turn must not discard the whole case's trace. ``invoke`` runs - inside the experiment's per-case ``try/except``, which would turn a raise - into ``score=0`` -- silently mislabeling a possible breach as defended. + - One malformed turn must not discard the whole case's trace. `invoke` runs + inside the experiment's per-case `try/except`, which would turn a raise + into `score=0` -- silently mislabeling a possible breach as defended. - The trace length must stay honest. A backtracking strategy decides whether a turn drove a tool call by whether the trace grew, so a real-but-malformed tool block must still grow the trace (hence the placeholder) or the breach could be backtracked away. - A ``logger.warning`` surfaces the drift without making it case-fatal. + A `logger.warning` surfaces the drift without making it case-fatal. Args: messages: Strands messages to scan (typically the tail appended by one - ``Agent`` call). + `Agent` call). Returns: - One :class:`ToolUseEntry` per ``toolUse`` block found, in order; a block - with a non-mapping ``toolUse`` yields a ``MALFORMED_TOOL_NAME`` placeholder. + One :class:`ToolUseEntry` per `toolUse` block found, in order; a block + with a non-mapping `toolUse` yields a `MALFORMED_TOOL_NAME` placeholder. """ tool_uses: list[ToolUseEntry] = [] for message in messages: @@ -96,10 +99,10 @@ class TargetCheckpoint: Bundles the target's saved state with the session's trace length at capture time, so :meth:`TargetSession.restore` can roll back both the target's conversation and the tool trace together. Strategies treat it as opaque — - take one, pass it back to ``restore``. + take one, pass it back to `restore`. - ``agent_snapshot`` is deliberately typed ``Any``: each session decides what it - stores there (a single ``strands.Snapshot`` for :class:`StrandsAgentSession`, a + `agent_snapshot` is deliberately typed `Any`: each session decides what it + stores there (a single `strands.Snapshot` for :class:`StrandsAgentSession`, a composite of per-agent snapshots + orchestrator state for a future multi-agent session). The checkpoint is opaque to strategies either way. @@ -107,11 +110,11 @@ class TargetCheckpoint: pattern Crescendo uses (one live checkpoint at a time). Non-monotonic restores across multiple outstanding checkpoints (e.g. a future PAIR/TAP tree search) would need the checkpoint to carry the trace slice rather than just its length; - revisit ``trace_len`` then. + revisit `trace_len` then. Attributes: agent_snapshot: The session's saved target state, opaque to strategies. - trace_len: ``len(session.trace)`` when the checkpoint was taken. + trace_len: `len(session.trace)` when the checkpoint was taken. """ agent_snapshot: Any @@ -121,7 +124,7 @@ class TargetCheckpoint: class TargetSession(Protocol): """Contract a strategy uses to interact with the target under test. - A strategy receives a ``TargetSession`` in ``run_attack`` and drives the + A strategy receives a `TargetSession` in `run_attack` and drives the conversation through :meth:`invoke`. Strategies that backtrack (e.g. Crescendo) use :meth:`snapshot` / :meth:`restore` to roll the target back; strategies that only converse forward (e.g. GOAT, Bad Likert Judge) use @@ -132,7 +135,7 @@ class TargetSession(Protocol): """Tool-use entries captured across this session's :meth:`invoke` calls.""" def invoke(self, message: str) -> str: - """Send ``message`` to the target and return its text response. + """Send `message` to the target and return its text response. Appends any tool uses observed during the call to :attr:`trace`. @@ -157,7 +160,7 @@ def snapshot(self) -> TargetCheckpoint: ... def restore(self, checkpoint: TargetCheckpoint) -> None: - """Roll the target back to a previously captured ``checkpoint``. + """Roll the target back to a previously captured `checkpoint`. Also rolls :attr:`trace` back to its length at checkpoint time, so tool uses from rolled-back turns do not linger in the trajectory. The session @@ -171,10 +174,10 @@ def restore(self, checkpoint: TargetCheckpoint) -> None: class StrandsAgentSession: - """A :class:`TargetSession` backed by a ``strands.Agent``. + """A :class:`TargetSession` backed by a `strands.Agent`. Rewindable: :meth:`snapshot` / :meth:`restore` delegate to the SDK's - ``Agent.take_snapshot`` / ``Agent.load_snapshot``, which restore the agent's + `Agent.take_snapshot` / `Agent.load_snapshot`, which restore the agent's message history to its snapshot-time state; :meth:`restore` additionally truncates :attr:`trace` back to its snapshot-time length. """ @@ -187,7 +190,7 @@ def __init__(self, agent: Agent, *, baseline: Snapshot | None = None) -> None: baseline: A clean snapshot of the agent to roll back to in :meth:`reset`. The task runner captures this once, before the first case, so every case starts from the same as-constructed target state (system prompt - plus any seeded history). When ``None`` (e.g. a directly-constructed + plus any seeded history). When `None` (e.g. a directly-constructed session), :meth:`reset` falls back to clearing messages only. """ self._agent = agent @@ -200,7 +203,7 @@ def invoke(self, message: str) -> str: return str(result) def _send(self, message: str) -> tuple[Any, list[Message]]: - """Send ``message``; return the agent's result and the messages it appended.""" + """Send `message`; return the agent's result and the messages it appended.""" messages_before = len(self._agent.messages) result = self._agent(message) return result, self._agent.messages[messages_before:] @@ -228,4 +231,257 @@ def restore(self, checkpoint: TargetCheckpoint) -> None: del self.trace[checkpoint.trace_len :] -__all__ = ["MALFORMED_TOOL_NAME", "StrandsAgentSession", "TargetCheckpoint", "TargetSession", "ToolUseEntry"] +# Path of node ids from the root orchestrator to a leaf Agent. Tuple (not str) so +# the same node id repeated under different parents stays distinguishable. +_AgentPath = tuple[str, ...] + + +@dataclass +class _MultiAgentSnapshot: + """Composite snapshot for a multi-agent target. + + `agents` holds one `Snapshot` per leaf `Agent` keyed by its path in the + tree; `orchestrators` holds one `serialize_state` dict per + `MultiAgentBase` (root and every nested orchestrator). Stored opaquely + inside :class:`TargetCheckpoint.agent_snapshot`; strategies never inspect it. + """ + + agents: dict[_AgentPath, Snapshot] + orchestrators: dict[_AgentPath, dict[str, Any]] + + +class StrandsMultiAgentSession: + """A :class:`TargetSession` backed by a `strands.multiagent.MultiAgentBase`. + + Wraps a Graph, Swarm, or any other `MultiAgentBase` (including nested + orchestrators). At init we walk the tree once to build a path index of every + leaf `Agent`; :meth:`snapshot` then captures one `Agent.take_snapshot` + per leaf plus a `serialize_state` dict per orchestrator, and :meth:`restore` + pushes them back through `Agent.load_snapshot` and + `MultiAgentBase.deserialize_state` respectively. The composite is opaque to + strategies (it lives in :attr:`TargetCheckpoint.agent_snapshot`). + + Trace capture diffs each leaf agent's `messages` tail across an + :meth:`invoke` call -- the same approach :class:`StrandsAgentSession` uses, + extended to every leaf so tool uses anywhere in the tree are recorded. + """ + + def __init__(self, root: MultiAgentBase, *, baseline: _MultiAgentSnapshot | None = None) -> None: + """Initialize the session. + + Args: + root: The target orchestrator. Its tree is walked once and indexed; + topology must not change after construction (Strands graphs and + swarms are static after build, so this is the normal case). + baseline: A clean composite snapshot to roll back to in :meth:`reset`. + The task runner captures this once, before the first case, so + every case starts from the same as-constructed target state. + When `None`, :meth:`reset` falls back to clearing each leaf + agent's `messages` (matching :class:`StrandsAgentSession`). + """ + self._root = root + self._agent_index: dict[_AgentPath, Agent] = {} + self._orch_index: dict[_AgentPath, MultiAgentBase] = {} + self._index_tree(root, ()) + self._baseline = baseline + self.trace: list[ToolUseEntry] = [] + + def _index_tree(self, orch: MultiAgentBase, path: _AgentPath) -> None: + """Walk the orchestrator tree once and populate the path indexes. + + Recurses into nested `MultiAgentBase` executors so a Graph-of-Graphs or + a Graph containing a Swarm is fully covered. Non-Agent, non-MultiAgentBase + executors (custom `AgentBase` subclasses without `take_snapshot`) are + skipped silently with a warning -- their state cannot round-trip through + the SDK snapshot API and including them would mislead a backtracking + strategy into thinking it captured the whole tree. + """ + self._orch_index[path] = orch + for node_id, node in orch.nodes.items(): + child_path = path + (node_id,) + executor = node.executor + if isinstance(executor, MultiAgentBase): + self._index_tree(executor, child_path) + elif isinstance(executor, Agent): + self._agent_index[child_path] = executor + else: + logger.warning( + "path=<%s>, type=<%s> | executor is not Agent or MultiAgentBase, skipping snapshot coverage", + "/".join(child_path), + type(executor).__name__, + ) + + def invoke(self, message: str) -> str: + """Send `message` to the root and capture tool uses from every leaf. + + Diffs each indexed agent's `messages` tail before/after the call, in + the deterministic order leaves were registered during :meth:`_index_tree` + (insertion order of each `orch.nodes` dict, recursed depth-first). A + single agent invoked multiple times during one orchestrator call (e.g. a + revisited graph node) contributes every new message in order. The exact + across-leaves order is implementation-defined when leaves run + concurrently, but each leaf's own tool-use order is preserved. + + Out of scope: a single `Agent` instance reused as the executor for + multiple distinct node paths. The diff is keyed by path, so a shared + instance's tail would be scanned once per path and its tool uses + double-counted. `Graph` and `Swarm` give each node its own executor, so + this only affects hand-built `MultiAgentBase` subclasses that + deliberately share instances; rebuild distinct executors per node if + accurate trace counts matter there. + """ + before = {path: len(agent.messages) for path, agent in self._agent_index.items()} + result = self._root(message) + for path, agent in self._agent_index.items(): + self.trace.extend(_tool_uses_in(agent.messages[before[path] :])) + return _multi_agent_result_text(result) + + def reset(self) -> None: + """Roll the tree back to the baseline (or clear messages if no baseline).""" + if self._baseline is not None: + self._restore(self._baseline) + else: + # No baseline: best-effort clear of every leaf's messages, mirroring + # StrandsAgentSession's no-baseline fallback. Orchestrator bookkeeping + # (completed_nodes, results, ...) is left to the SDK to refresh on + # the next invoke; without a baseline we don't have a serialize_state + # payload to restore from. + for agent in self._agent_index.values(): + agent.messages.clear() + self.trace.clear() + + def snapshot(self) -> TargetCheckpoint: + """Capture a composite snapshot of every leaf agent and every orchestrator.""" + return TargetCheckpoint( + agent_snapshot=_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), + ) + + def restore(self, checkpoint: TargetCheckpoint) -> None: + """Roll the tree back to `checkpoint` and truncate the trace.""" + if not isinstance(checkpoint.agent_snapshot, _MultiAgentSnapshot): + raise TypeError( + f"StrandsMultiAgentSession.restore: expected _MultiAgentSnapshot, " + f"got {type(checkpoint.agent_snapshot).__name__}" + ) + self._restore(checkpoint.agent_snapshot) + del self.trace[checkpoint.trace_len :] + + def _restore(self, snapshot: _MultiAgentSnapshot) -> None: + """Push a composite snapshot back into the tree. + + Deep-copies the composite once at the boundary so a stored or baseline + snapshot can be replayed across cases without being mutated by the load. + Cloning the whole composite once shares the deepcopy cost across leaves + and orchestrators rather than copying piecewise. + + Orchestrators are restored FIRST, leaves LAST. `Graph.deserialize_state` + and `Swarm.deserialize_state` reset every node's executor state to + graph-build-time values when the payload has no ``next_nodes_to_execute`` + (the common case between attack turns: the orchestrator is PENDING or + COMPLETED). Running the orchestrator load AFTER the leaf loads would wipe + the leaves we just restored back to build-time, silently breaking + backtracking. Doing orchestrators first lets the per-leaf snapshots be + the final writers; for interrupted payloads the order is irrelevant + because `deserialize_state` does not touch leaves. + + After each `deserialize_state`, settled-status payloads get the + orchestrator's resume bookkeeping forced back to a fresh-invoke state. + Settled covers PENDING/COMPLETED/FAILED; in practice you'll see + COMPLETED between attack turns and FAILED on a target that errored + mid-attack, with PENDING included for completeness (a snapshot of a + never-invoked orchestrator). `Swarm.deserialize_state` always takes + the resume branch because `Swarm.serialize_state` always emits the + `next_nodes_to_execute` key (empty list for a settled swarm) and the + deserialize side checks key presence, not truthiness; without this + forcing, `_resume_from_session` stays True with `current_node=None` + and the next invoke crashes with `AttributeError`. The same forcing + is a no-op-or-better for `Graph` (which already takes the reset branch + for empty next-nodes) and any other `MultiAgentBase`. + """ + clone = copy.deepcopy(snapshot) + for path, orch_state in clone.orchestrators.items(): + orch = self._orch_index.get(path) + if orch is None: + logger.warning("path=<%s> | orchestrator path missing at restore, skipping", "/".join(path)) + continue + orch.deserialize_state(orch_state) + self._force_fresh_invoke_if_settled(orch, orch_state) + for path, agent_snap in clone.agents.items(): + agent = self._agent_index.get(path) + if agent is None: + # Topology shifted out from under us (caller mutated the tree + # post-construction). Drop the orphaned entry; the user should + # rebuild the session if they reshape the target. + logger.warning("path=<%s> | agent path missing at restore, skipping", "/".join(path)) + continue + agent.load_snapshot(agent_snap) + + @staticmethod + def _force_fresh_invoke_if_settled(orch: MultiAgentBase, orch_state: dict[str, Any]) -> None: + """Clear `_resume_from_session` for a settled-status payload. + + Settled = PENDING/COMPLETED/FAILED in the just-restored payload. For a + Swarm this is the load-bearing fix (see :meth:`_restore`); for a Graph + it's redundant with `deserialize_state`'s own reset path but harmless; + for any third-party `MultiAgentBase` it provides the same between-turn + guarantee. + + `_resume_from_session` is a private SDK attribute we deliberately reach + into. If it's missing on an orchestrator that returned a settled-status + payload, the SDK has likely renamed/restructured it -- log a warning + rather than silently no-op'ing, because the original Swarm bug + (silent score=0 / "defended" mislabels) is exactly what comes back when + this guard goes stale unnoticed. Custom `MultiAgentBase` subclasses + without the attribute trigger the same warning once per restore; tag + them with `_resume_from_session = False` to opt out. + """ + status = orch_state.get("status") + if status not in {Status.PENDING.value, Status.COMPLETED.value, Status.FAILED.value}: + return + if hasattr(orch, "_resume_from_session"): + orch._resume_from_session = False + return + logger.warning( + "orchestrator=<%s> | settled-status payload but no `_resume_from_session` attribute; " + "SDK may have renamed it -- update StrandsMultiAgentSession or strategy backtracks may " + "silently mislabel cases as defended", + type(orch).__name__, + ) + + +def _multi_agent_result_text(result: Any) -> str: + """Best-effort extract a textual response from a `MultiAgentResult`. + + The orchestrator returns a `MultiAgentResult` (a tree of `NodeResult`s), + not a single string. Strategies expect text, so we flatten the underlying + `AgentResult`s in node-result iteration order via `NodeResult.get_agent_results` + and concatenate their string forms. Falls back to `str(result)` when no + agent results surface (orchestrator failure, custom `MultiAgentBase` + returning a different shape) -- we hand the strategy a string rather than + raise inside the per-case `try/except` and silently mislabel a turn as + defended. + """ + results_dict = getattr(result, "results", None) + if isinstance(results_dict, dict): + parts: list[str] = [] + for node_result in results_dict.values(): + get_results = getattr(node_result, "get_agent_results", None) + if callable(get_results): + parts.extend(str(r) for r in get_results()) + if parts: + return "\n".join(parts) + return str(result) + + +__all__ = [ + "MALFORMED_TOOL_NAME", + "StrandsAgentSession", + "StrandsMultiAgentSession", + "TargetCheckpoint", + "TargetSession", + "ToolUseEntry", +] diff --git a/src/strands_evals/experimental/redteam/task.py b/src/strands_evals/experimental/redteam/task.py index aaf0982e..b53fe9fa 100644 --- a/src/strands_evals/experimental/redteam/task.py +++ b/src/strands_evals/experimental/redteam/task.py @@ -6,12 +6,13 @@ from collections.abc import Callable from typing import Any -from strands import Agent, Snapshot +from strands import Agent from strands.models.model import Model +from strands.multiagent.base import MultiAgentBase from .case import RedTeamCase from .strategies import AttackStrategy -from .strategies.target_session import StrandsAgentSession, TargetSession +from .strategies.target_session import StrandsAgentSession, StrandsMultiAgentSession, TargetSession logger = logging.getLogger(__name__) @@ -19,38 +20,52 @@ def _build_attacker_task( - agent: Agent | TargetSession, + agent: Agent | MultiAgentBase | TargetSession, by_label: dict[str, AttackStrategy], *, model: Model | str | None = None, run_meta: dict[str, dict[str, Any]] | None = None, ) -> Callable[[RedTeamCase], dict]: - """Build a red team task function for ``Experiment.run_evaluations``. - - Internal helper used by :class:`RedTeamExperiment`. Returns a ``task(case) - -> {"output": conversation, "trajectory": tool_uses}`` that looks up the - case's strategy (by ``metadata["strategy"]``) and delegates the multi-turn - loop to ``strategy.run_attack``, injecting a ``TargetSession`` that handles - target invocation, tool-trace capture, and per-case isolation. An ``Agent`` - is wrapped in a ``StrandsAgentSession``; a ``TargetSession`` is used as-is. - Either way the session is reset between cases. - - Each strategy owns its own turn budget; ``MAX_ALLOWED_TURNS`` is passed as a + """Build a red team task function for `Experiment.run_evaluations`. + + Internal helper used by :class:`RedTeamExperiment`. Returns a `task(case) + -> {"output": conversation, "trajectory": tool_uses}` that looks up the + case's strategy (by `metadata["strategy"]`) and delegates the multi-turn + loop to `strategy.run_attack`, injecting a `TargetSession` that handles + target invocation, tool-trace capture, and per-case isolation. An `Agent` + is wrapped in a `StrandsAgentSession`; a `MultiAgentBase` (Graph, Swarm, + nested orchestrator) is wrapped in a `StrandsMultiAgentSession`; a + `TargetSession` is used as-is. Either way the session is reset between + cases. + + Each strategy owns its own turn budget; `MAX_ALLOWED_TURNS` is passed as a hard ceiling so no strategy can run unbounded. The strategy's run metadata (turns_used, backtracks, ...) is recorded into - ``run_meta`` keyed by case name; the experiment owns that dict and joins it - onto the report (the base ``Experiment`` copies ``metadata`` into a fresh - ``EvaluationData``, so the strategy can't reach the report through it). + `run_meta` keyed by case name; the experiment owns that dict and joins it + onto the report (the base `Experiment` copies `metadata` into a fresh + `EvaluationData`, so the strategy can't reach the report through it). """ # Capture the target's clean state ONCE, before the first case, so every case # resets to the same as-constructed baseline. Must be here (build time), not in - # the session __init__: the agent is shared and reused, so by the time case N's - # session is built the agent already carries case N-1's conversation -- snapshotting - # then would bake a dirty baseline. Only an Agent is rewindable this way; a - # passed-in TargetSession owns its own reset. - initial_snapshot = agent.take_snapshot(preset="session") if isinstance(agent, Agent) else None + # the session __init__: the target is shared and reused, so by the time case N's + # session is built the target already carries case N-1's conversation -- + # snapshotting then would bake a dirty baseline. Only an Agent or + # MultiAgentBase is rewindable this way; a passed-in TargetSession owns its + # own reset. + initial_snapshot: Any + if isinstance(agent, Agent): + initial_snapshot = agent.take_snapshot(preset="session") + elif isinstance(agent, MultiAgentBase): + # The composite snapshot shape is internal to StrandsMultiAgentSession; + # capture it through the session's public snapshot() so this layer never + # has to import _MultiAgentSnapshot. The throwaway session shares the + # same indexes the per-case session will rebuild, so the baseline is + # apples-to-apples. + initial_snapshot = StrandsMultiAgentSession(agent).snapshot().agent_snapshot + else: + initial_snapshot = None def task_fn(case: RedTeamCase) -> dict: strategy = _resolve_case_strategy(case, by_label) @@ -81,22 +96,35 @@ def task_fn(case: RedTeamCase) -> dict: return task_fn -def _build_session(agent: Agent | TargetSession, *, baseline: Snapshot | None = None) -> TargetSession: - """Wrap an ``Agent`` in a ``StrandsAgentSession``, or pass a ``TargetSession`` through. +def _build_session( + agent: Agent | MultiAgentBase | TargetSession, + *, + baseline: Any = None, +) -> TargetSession: + """Wrap an `Agent` or `MultiAgentBase`, or pass a `TargetSession` through. Args: - agent: An ``Agent`` to wrap, or a ready ``TargetSession`` to use as-is. - baseline: A clean ``Snapshot`` the wrapped ``StrandsAgentSession`` resets to - between cases (ignored for a passed-in ``TargetSession``, which owns its - own reset). + agent: An `Agent` to wrap in :class:`StrandsAgentSession`, a + `MultiAgentBase` (Graph, Swarm, ...) to wrap in + :class:`StrandsMultiAgentSession`, or a ready `TargetSession` to + use as-is. + baseline: A clean snapshot the wrapped session resets to between cases. + For an `Agent` this is a `Snapshot`; for a `MultiAgentBase` + it's the opaque composite returned by + :meth:`StrandsMultiAgentSession.snapshot`. Ignored for a passed-in + `TargetSession`, which owns its own reset. Typed `Any` because + this is the boundary that fans into two session types whose + baseline shapes are distinct (and one is private). Raises: - TypeError: If ``agent`` is neither an ``Agent`` nor a ``TargetSession`` - (e.g. a bare callable -- a strategy needs snapshot/restore to manage - the target's state, which an opaque callable cannot provide). + TypeError: If `agent` is none of the supported types (e.g. a bare + callable -- a strategy needs snapshot/restore to manage the target's + state, which an opaque callable cannot provide). """ if isinstance(agent, Agent): return StrandsAgentSession(agent, baseline=baseline) + if isinstance(agent, MultiAgentBase): + return StrandsMultiAgentSession(agent, baseline=baseline) # Structural (not isinstance) check: TargetSession is a Protocol, and the method # set is checked by hand. The `trace` check is separate and load-bearing -- it's # the one member the task runner dereferences directly (it becomes the @@ -108,13 +136,14 @@ def _build_session(agent: Agent | TargetSession, *, baseline: Snapshot | None = if has_methods and isinstance(getattr(agent, "trace", None), list): return agent raise TypeError( - f"agent must be a strands.Agent or a TargetSession, got {type(agent).__name__!r}; " - "wrap a custom target in a TargetSession so the strategy can snapshot/restore its state." + f"agent must be a strands.Agent, strands.multiagent.MultiAgentBase, or a TargetSession, " + f"got {type(agent).__name__!r}; wrap a custom target in a TargetSession so the strategy " + "can snapshot/restore its state." ) def _resolve_case_strategy(case: RedTeamCase, by_label: dict[str, AttackStrategy]) -> AttackStrategy: - """Look up the strategy instance for a case from its ``metadata["strategy"]`` label.""" + """Look up the strategy instance for a case from its `metadata["strategy"]` label.""" metadata = case.metadata or {} label = metadata.get("strategy") if label is None: diff --git a/tests/strands_evals/experimental/redteam/test_multi_agent_session.py b/tests/strands_evals/experimental/redteam/test_multi_agent_session.py new file mode 100644 index 00000000..cd0a7859 --- /dev/null +++ b/tests/strands_evals/experimental/redteam/test_multi_agent_session.py @@ -0,0 +1,590 @@ +"""Tests for StrandsMultiAgentSession.""" + +import copy +from dataclasses import dataclass, field +from typing import Any + +import pytest +from strands import Agent +from strands.multiagent.base import MultiAgentBase + +from strands_evals.experimental.redteam.strategies.target_session import ( + StrandsMultiAgentSession, + TargetCheckpoint, + _MultiAgentSnapshot, +) + +# --------------------------------------------------------------------------- +# Fakes -- minimal MultiAgentBase the session can drive without a real LLM +# --------------------------------------------------------------------------- + + +@dataclass +class _FakeNode: + node_id: str + executor: Any + + +class _FakeOrchestrator(MultiAgentBase): + """A MultiAgentBase that exposes the surface StrandsMultiAgentSession uses. + + The real Graph/Swarm classes do far more, but the session only touches + `nodes`, `__call__`, `serialize_state`, and `deserialize_state` -- + so a fake with those is enough to exercise every code path without booting + a model. `__call__` walks the leaves and feeds each one the message, which + is all that's needed to drive the session's tail-diff trace capture. + """ + + id: str = "root" + + def __init__( + self, + nodes: dict[str, Any], + *, + result_text: str = "ok", + invoke_message: str | None = None, + ) -> None: + self.nodes = {nid: _FakeNode(nid, ex) for nid, ex in nodes.items()} + self._result_text = result_text + # Optional override for what the leaves are sent (defaults to forwarding). + self._invoke_message = invoke_message + # Seedable orchestrator-level state so serialize/deserialize can be + # exercised end-to-end. + self._orch_state: dict[str, Any] = {"version": 1, "scratch": None} + + async def invoke_async(self, task, invocation_state=None, **kwargs): # pragma: no cover - unused + raise NotImplementedError + + def __call__(self, task, invocation_state=None, **kwargs): + msg = self._invoke_message if self._invoke_message is not None else task + for node in self.nodes.values(): + ex = node.executor + if isinstance(ex, _FakeOrchestrator): + ex(msg) + elif isinstance(ex, Agent): + ex(msg) + return _FakeMultiAgentResult(self._result_text) + + def serialize_state(self) -> dict[str, Any]: + return {"type": "fake", "id": self.id, "orch_state": dict(self._orch_state)} + + def deserialize_state(self, payload: dict[str, Any]) -> None: + self._orch_state = dict(payload.get("orch_state", {})) + + +@dataclass +class _FakeMultiAgentResult: + """The minimum shape `_multi_agent_result_text` falls back to via `str()`.""" + + text: str + results: dict = field(default_factory=dict) + + def __str__(self) -> str: + return self.text + + +class _FakeNonAgentExecutor: + """An executor that is neither Agent nor MultiAgentBase -- must be skipped.""" + + +def _real_agent(seed_text: str | None = None) -> Agent: + """Build a real Agent (no model needed) so take_snapshot/load_snapshot + exercise the SDK path rather than mocks. Mirrors the helper in + test_target_session.py.""" + messages = [{"role": "user", "content": [{"text": seed_text}]}] if seed_text else [] + return Agent(model=None, messages=messages, callback_handler=None) + + +def _append_to(agent: Agent, text: str) -> None: + """Helper that pushes a turn the same way Agent.__call__ would, without + actually invoking a model.""" + agent.messages.append({"role": "user", "content": [{"text": text}]}) + + +# --------------------------------------------------------------------------- +# Tree indexing +# --------------------------------------------------------------------------- + + +class TestIndexTree: + def test_indexes_top_level_agents(self): + a = _real_agent() + b = _real_agent() + root = _FakeOrchestrator({"a": a, "b": b}) + s = StrandsMultiAgentSession(root) + assert set(s._agent_index.keys()) == {("a",), ("b",)} + assert s._agent_index[("a",)] is a + assert set(s._orch_index.keys()) == {()} + + def test_recurses_into_nested_orchestrators(self): + # Root contains a nested orchestrator with two leaves -- both must be indexed + # at their full path so a node id repeated under different parents stays + # distinguishable. + leaf_a = _real_agent() + leaf_b = _real_agent() + nested = _FakeOrchestrator({"x": leaf_a, "y": leaf_b}) + root = _FakeOrchestrator({"sub": nested}) + s = StrandsMultiAgentSession(root) + assert set(s._agent_index.keys()) == {("sub", "x"), ("sub", "y")} + assert set(s._orch_index.keys()) == {(), ("sub",)} + + def test_skips_unsupported_executors(self): + # A custom AgentBase subclass without take_snapshot can't round-trip; the + # session must not silently include it (would lie to a backtracking strategy + # about coverage) but must not raise either (other leaves are still valid). + a = _real_agent() + root = _FakeOrchestrator({"good": a, "bad": _FakeNonAgentExecutor()}) + s = StrandsMultiAgentSession(root) + assert set(s._agent_index.keys()) == {("good",)} + + +# --------------------------------------------------------------------------- +# invoke + trace capture +# --------------------------------------------------------------------------- + + +class TestInvoke: + def test_returns_string_response(self): + # No leaves needed — this test only checks that the orchestrator's result + # text is returned. Driving a real Agent here would call Bedrock. + root = _FakeOrchestrator({}, result_text="hello back") + assert StrandsMultiAgentSession(root).invoke("hi") == "hello back" + + def test_diffs_each_leaf_for_tool_uses(self): + # The fake __call__ doesn't run a real model, so we synthesize the message + # the agent would have appended (a tool-use block) before invoke captures + # the tail. Two leaves, each contributing one tool use; both must surface. + leaf_a = _real_agent() + leaf_b = _real_agent() + + class _Recording(_FakeOrchestrator): + def __call__(self, task, invocation_state=None, **kwargs): + leaf_a.messages.append({"content": [{"toolUse": {"name": "search_a", "input": {}}}]}) + leaf_b.messages.append({"content": [{"toolUse": {"name": "search_b", "input": {}}}]}) + return _FakeMultiAgentResult("ok") + + root = _Recording({"a": leaf_a, "b": leaf_b}) + s = StrandsMultiAgentSession(root) + s.invoke("hi") + names = sorted(t["name"] for t in s.trace) + assert names == ["search_a", "search_b"] + + def test_only_new_messages_are_scanned(self): + # A leaf that already has prior tool uses must NOT have them re-counted on + # the next invoke -- only the tail appended during this call. Use a no-op + # orchestrator so the prior (deliberately minimal) tool-use block doesn't + # have to satisfy the SDK's full message schema. + leaf = _real_agent() + leaf.messages.append({"content": [{"toolUse": {"name": "old", "input": {}}}]}) + + class _NoOp(_FakeOrchestrator): + def __call__(self, task, invocation_state=None, **kwargs): + return _FakeMultiAgentResult("ok") + + root = _NoOp({"a": leaf}) + s = StrandsMultiAgentSession(root) + s.invoke("hi") + assert s.trace == [] + + +# --------------------------------------------------------------------------- +# snapshot / restore +# --------------------------------------------------------------------------- + + +class TestSnapshotRestore: + def test_snapshot_captures_every_leaf_and_orchestrator(self): + leaf_a = _real_agent() + leaf_b = _real_agent() + nested = _FakeOrchestrator({"y": leaf_b}) + root = _FakeOrchestrator({"a": leaf_a, "sub": nested}) + s = StrandsMultiAgentSession(root) + + ck = s.snapshot() + assert isinstance(ck, TargetCheckpoint) + assert isinstance(ck.agent_snapshot, _MultiAgentSnapshot) + assert set(ck.agent_snapshot.agents.keys()) == {("a",), ("sub", "y")} + assert set(ck.agent_snapshot.orchestrators.keys()) == {(), ("sub",)} + + def test_restore_rolls_back_each_leaf_messages(self): + leaf_a = _real_agent("seed_a") + leaf_b = _real_agent("seed_b") + root = _FakeOrchestrator({"a": leaf_a, "b": leaf_b}) + s = StrandsMultiAgentSession(root) + + ck = s.snapshot() + # mutate every leaf past the snapshot + _append_to(leaf_a, "case turn a") + _append_to(leaf_b, "case turn b") + assert len(leaf_a.messages) == 2 + assert len(leaf_b.messages) == 2 + + s.restore(ck) + assert len(leaf_a.messages) == 1 + assert leaf_a.messages[0]["content"][0]["text"] == "seed_a" + assert len(leaf_b.messages) == 1 + assert leaf_b.messages[0]["content"][0]["text"] == "seed_b" + + def test_restore_rolls_back_orchestrator_state(self): + # serialize_state / deserialize_state are the orchestrator-bookkeeping + # channel; mutating it post-snapshot and then restoring must revert it. + root = _FakeOrchestrator({"a": _real_agent()}) + s = StrandsMultiAgentSession(root) + ck = s.snapshot() + root._orch_state["scratch"] = "dirty" + s.restore(ck) + assert root._orch_state["scratch"] is None + + def test_restore_truncates_trace(self): + root = _FakeOrchestrator({"a": _real_agent()}) + s = StrandsMultiAgentSession(root) + s.trace.append({"name": "before", "input": {}}) + ck = s.snapshot() + s.trace.append({"name": "after_snapshot", "input": {}}) + s.restore(ck) + assert s.trace == [{"name": "before", "input": {}}] + + def test_restore_rejects_wrong_checkpoint_payload(self): + # An opaque agent_snapshot that isn't a _MultiAgentSnapshot is a programmer + # error (someone mixed checkpoints across session types). Fail loud rather + # than silently passing through deepcopy and exploding deeper. + root = _FakeOrchestrator({"a": _real_agent()}) + s = StrandsMultiAgentSession(root) + bogus = TargetCheckpoint(agent_snapshot="not a multi-agent snapshot", trace_len=0) + with pytest.raises(TypeError, match="_MultiAgentSnapshot"): + s.restore(bogus) + + +# --------------------------------------------------------------------------- +# reset +# --------------------------------------------------------------------------- + + +class TestReset: + def test_reset_with_baseline_restores_as_constructed_state(self): + leaf = _real_agent("seed") + root = _FakeOrchestrator({"a": leaf}) + baseline = StrandsMultiAgentSession(root).snapshot().agent_snapshot + s = StrandsMultiAgentSession(root, baseline=baseline) + + # simulate a case mutating the target + _append_to(leaf, "case turn") + leaf.state.set("exfiltrated", True) + root._orch_state["scratch"] = "dirty" + s.trace.append({"name": "leftover", "input": {}}) + + s.reset() + + assert leaf.messages == [{"role": "user", "content": [{"text": "seed"}]}] + # full snapshot fields reset, not just messages -- agent state too + assert leaf.state.get("exfiltrated") is None + # orchestrator-level state rolled back through deserialize_state + assert root._orch_state["scratch"] is None + assert s.trace == [] + + def test_baseline_survives_repeated_resets(self): + # The baseline is captured once and replayed every case; load_snapshot must + # copy out of it so a case mutation can't poison subsequent resets. + leaf = _real_agent("seed") + root = _FakeOrchestrator({"a": leaf}) + baseline = StrandsMultiAgentSession(root).snapshot().agent_snapshot + s = StrandsMultiAgentSession(root, baseline=baseline) + for i in range(3): + _append_to(leaf, f"case{i}") + leaf.state.set("leak", i) + root._orch_state["scratch"] = f"dirty{i}" + s.reset() + assert leaf.messages == [{"role": "user", "content": [{"text": "seed"}]}] + assert leaf.state.get("leak") is None + assert root._orch_state["scratch"] is None + + def test_reset_without_baseline_clears_leaf_messages(self): + # No baseline: best-effort fallback clears each leaf's messages, mirroring + # StrandsAgentSession's no-baseline behavior. Orchestrator state is left as + # the SDK has it (next invoke refreshes it). + leaf_a = _real_agent("seed") + leaf_b = _real_agent() + _append_to(leaf_b, "added") + root = _FakeOrchestrator({"a": leaf_a, "b": leaf_b}) + s = StrandsMultiAgentSession(root) + s.trace.append({"name": "leftover", "input": {}}) + s.reset() + assert leaf_a.messages == [] + assert leaf_b.messages == [] + assert s.trace == [] + + +# --------------------------------------------------------------------------- +# restore ordering vs Graph/Swarm.deserialize_state +# --------------------------------------------------------------------------- + + +class _GraphLikeOrchestrator(_FakeOrchestrator): + """A `_FakeOrchestrator` whose `deserialize_state` mimics + `Graph`/`Swarm.deserialize_state` for completed payloads. + + The real SDK resets every leaf's `messages` / `state` / `_model_state` to the + values captured at `GraphBuilder.build()` time (`GraphNode.__post_init__`) + whenever the payload has no `next_nodes_to_execute` -- the normal between-turn + state for a `PENDING` or `COMPLETED` orchestrator. Without this fake, + `_FakeOrchestrator.deserialize_state` only round-trips an opaque dict and the + leaf-vs-orchestrator restore order is invisible. + + The fake snapshots each leaf's `messages` and `state` at the point the + orchestrator was wrapped (~ `GraphBuilder.build()` time) and replays them on + completed-payload restores. + """ + + def __init__(self, nodes: dict[str, Any]) -> None: + super().__init__(nodes) + # Capture per-leaf build-time messages + state the same way + # GraphNode.__post_init__ does. We deep-copy each leaf's `state` object + # whole and replay it on reset; this keeps us agnostic to whether the + # SDK exposes state as `AgentState` or `JSONSerializableDict`. + self._initial_messages: dict[str, Any] = {} + self._initial_state: dict[str, Any] = {} + for node_id, node in self.nodes.items(): + ex = node.executor + if hasattr(ex, "messages"): + self._initial_messages[node_id] = copy.deepcopy(ex.messages) + if hasattr(ex, "state"): + self._initial_state[node_id] = copy.deepcopy(ex.state) + + def serialize_state(self) -> dict[str, Any]: + # Match Graph: a settled orchestrator returns a payload with no + # `next_nodes_to_execute`, which is exactly the case that triggers leaf + # reset on the way back in. + return {"type": "graph-like", "next_nodes_to_execute": [], "orch_state": dict(self._orch_state)} + + def deserialize_state(self, payload: dict[str, Any]) -> None: + if not payload.get("next_nodes_to_execute"): + for node_id, node in self.nodes.items(): + ex = node.executor + if node_id in self._initial_messages and hasattr(ex, "messages"): + ex.messages = copy.deepcopy(self._initial_messages[node_id]) + if node_id in self._initial_state and hasattr(ex, "state"): + ex.state = copy.deepcopy(self._initial_state[node_id]) + self._orch_state = dict(payload.get("orch_state", {})) + + +class TestRestoreOrderVsGraphReset: + """Regression: leaves must outlive `deserialize_state`'s build-time reset. + + With the loops in the wrong order, every backtrack against a real + `Graph`/`Swarm` target restarts the conversation from build-time state, + silently breaking Crescendo's "escalate from accumulated context" guarantee. + """ + + def test_restore_preserves_post_build_messages_on_completed_payload(self): + # Build-time leaf has just the seed message; the case appended a turn + # AFTER build, then we snapshot. Restoring must end with snapshot-time + # state (seed + 1 turn), NOT build-time state (seed only) -- the latter + # is what the buggy ordering produces. + leaf = _real_agent("seed") + root = _GraphLikeOrchestrator({"a": leaf}) + _append_to(leaf, "case turn") # post-build mutation + assert len(leaf.messages) == 2 + + ck = StrandsMultiAgentSession(root).snapshot() + _append_to(leaf, "to be rolled back") # past the snapshot + assert len(leaf.messages) == 3 + + StrandsMultiAgentSession(root).restore(ck) + + # 2 messages = snapshot-time state; 1 would mean deserialize_state ran + # AFTER load_snapshot and reset the leaf to build-time. + assert len(leaf.messages) == 2 + assert leaf.messages[0]["content"][0]["text"] == "seed" + assert leaf.messages[1]["content"][0]["text"] == "case turn" + + def test_baseline_reset_preserves_seeded_post_build_state(self): + # The task runner captures the baseline AFTER the harness wraps the + # target -- so any leaf history seeded after `GraphBuilder.build()` lives + # in the baseline, and `reset()` must restore to it. With the loops in + # the wrong order, `deserialize_state` would wipe that seeded state and + # every case would start from build-time. + leaf = _real_agent("seed") + root = _GraphLikeOrchestrator({"a": leaf}) + _append_to(leaf, "post-build seed") # seeded after the orchestrator was built + assert len(leaf.messages) == 2 + + baseline = StrandsMultiAgentSession(root).snapshot().agent_snapshot + s = StrandsMultiAgentSession(root, baseline=baseline) + + _append_to(leaf, "case mutation") + leaf.state.set("exfiltrated", True) + s.reset() + + assert len(leaf.messages) == 2 + assert leaf.messages[1]["content"][0]["text"] == "post-build seed" + assert leaf.state.get("exfiltrated") is None + + +# --------------------------------------------------------------------------- +# Swarm-shaped resume-flag handling +# --------------------------------------------------------------------------- + + +class _SwarmLikeOrchestrator(_FakeOrchestrator): + """A `_FakeOrchestrator` whose `(de)serialize_state` mimics `Swarm`. + + Two SDK details matter for `_restore` correctness and aren't exercised by + the Graph-shaped fake: + + 1. `Swarm.serialize_state` ALWAYS emits `next_nodes_to_execute` (empty list + for a settled swarm with no current_node / no handoff). + 2. `Swarm.deserialize_state` checks key PRESENCE (`"key" in payload`), not + truthiness, so a settled-but-round-tripped swarm always takes the + resume branch -- setting `_resume_from_session=True` while leaving + `current_node=None`. The next invoke then dereferences `current_node.node_id` + and crashes; the experiment's per-case try/except records score=0 and + silently mislabels every Swarm case as defended. + + This fake reproduces both behaviors so the session's "force fresh invoke + when payload is settled" guard can be regression-tested without booting a + real swarm. + """ + + def __init__(self, nodes: dict[str, Any], *, status: str = "pending") -> None: + super().__init__(nodes) + self._status = status + self._resume_from_session = False + self.current_node: str | None = "a" # set by a fresh invoke + self.invoked_with: str | None = None + + def __call__(self, task, invocation_state=None, **kwargs): + # Mirror Swarm.invoke: when _resume_from_session is True, skip the + # current_node re-init, then dereference it -- crashes if None. We + # don't iterate leaves here (the base class would call them as real + # Agents with no model) -- the guard's effect is observable from the + # current_node behavior alone. + if not self._resume_from_session: + self.current_node = "a" + # The crash the real Swarm hits at swarm.py:419: AttributeError on a + # None current_node. The guard is what stops it. + _ = self.current_node.upper() # type: ignore[union-attr] + self.invoked_with = task + return _FakeMultiAgentResult(self._result_text) + + def serialize_state(self) -> dict[str, Any]: + # Always emit `next_nodes_to_execute` (empty for a settled swarm), and + # carry a `status` so the session's resume-flag guard can read it. + return { + "type": "swarm-like", + "status": self._status, + "next_nodes_to_execute": [], + "orch_state": dict(self._orch_state), + } + + def deserialize_state(self, payload: dict[str, Any]) -> None: + # Membership check, not truthiness -- the bug's root cause on the SDK side. + self._resume_from_session = "next_nodes_to_execute" in payload + if self._resume_from_session: + # _from_dict equivalent: empty next_node_ids leaves current_node alone, + # which is None on a freshly-restored fake. + self.current_node = None + self._orch_state = dict(payload.get("orch_state", {})) + + +class TestSwarmShapedResumeFlagGuard: + """Regression: a settled-payload restore must leave the orchestrator able + to start a fresh invoke, not stuck in a half-resumed state. + + Without the guard, `Swarm.deserialize_state` always takes the resume + branch (key presence, not truthiness), `current_node` is None, and the + next invoke crashes -- which the experiment's per-case try/except buries + as score=0 / "defended", silently mislabeling every Swarm case safe. + """ + + def test_invoke_after_restore_does_not_crash(self): + # Settled status (COMPLETED) + always-emit-next-nodes is exactly the + # Swarm-after-a-turn shape. Without _force_fresh_invoke_if_settled, + # this invoke raises AttributeError. + leaf = _real_agent("seed") + root = _SwarmLikeOrchestrator({"a": leaf}, status="completed") + s = StrandsMultiAgentSession(root) + + ck = s.snapshot() + s.restore(ck) + + # The next invoke must succeed; the guard must have cleared + # _resume_from_session so the orchestrator re-initializes current_node. + result = s.invoke("next attack turn") + assert result == "ok" + assert root.invoked_with == "next attack turn" + assert root._resume_from_session is False + + def test_reset_with_baseline_does_not_crash(self): + # The task runner calls session.reset() before every case -- so the + # bug surfaces on case 1 turn 1, not just on backtracks. Same guard, + # same fix; this test asserts the case-startup path specifically. + leaf = _real_agent("seed") + root = _SwarmLikeOrchestrator({"a": leaf}, status="completed") + baseline = StrandsMultiAgentSession(root).snapshot().agent_snapshot + s = StrandsMultiAgentSession(root, baseline=baseline) + + s.reset() + result = s.invoke("case 1 turn 1") + assert result == "ok" + assert root.invoked_with == "case 1 turn 1" + + def test_pending_and_failed_payloads_also_force_fresh_invoke(self): + # PENDING and FAILED are the other settled statuses; both must clear + # the resume flag too so a fresh invoke re-initializes cleanly. + for status in ("pending", "failed"): + leaf = _real_agent("seed") + root = _SwarmLikeOrchestrator({"a": leaf}, status=status) + s = StrandsMultiAgentSession(root) + s.restore(s.snapshot()) + assert root._resume_from_session is False, f"status={status} left resume flag set" + + def test_interrupted_payload_keeps_resume_flag(self): + # An INTERRUPTED payload genuinely needs to resume mid-execution; the + # guard must NOT clear the flag in that case (the orchestrator owns + # current_node for an interrupted resume). + leaf = _real_agent("seed") + root = _SwarmLikeOrchestrator({"a": leaf}, status="interrupted") + s = StrandsMultiAgentSession(root) + s.restore(s.snapshot()) + # _resume_from_session was set by deserialize_state and the guard left + # it alone because the payload status isn't settled. + assert root._resume_from_session is True + + def test_settled_payload_without_resume_attr_logs_warning(self, caplog): + # The guard reaches into a private SDK attribute. If a future SDK + # rename drops `_resume_from_session`, the guard would silently no-op + # and the original bug (score=0 cases looking "defended") would come + # back unnoticed. The warning is the tripwire that catches the rename. + class _NoResumeAttr(_FakeOrchestrator): + # serialize_state emits a settled `status` so the guard runs; + # deserialize_state never sets `_resume_from_session`, so the + # attribute is genuinely absent at the moment the guard checks. + def serialize_state(self) -> dict[str, Any]: + return {"type": "no-resume", "status": "completed"} + + def deserialize_state(self, payload: dict[str, Any]) -> None: + pass + + leaf = _real_agent("seed") + root = _NoResumeAttr({"a": leaf}) + s = StrandsMultiAgentSession(root) + with caplog.at_level( + "WARNING", logger="strands_evals.experimental.redteam.strategies.target_session" + ): + s.restore(s.snapshot()) + assert any("_resume_from_session" in rec.message for rec in caplog.records), ( + "missing-attribute path must log a rename warning" + ) + + def test_non_settled_payload_without_resume_attr_does_not_warn(self, caplog): + # A custom MultiAgentBase that doesn't carry `_resume_from_session` AND + # doesn't return a settled status (or returns no status at all) should + # not trip the warning -- the guard has no business firing there. + leaf = _real_agent("seed") + root = _FakeOrchestrator({"a": leaf}) # serialize_state emits no `status` + s = StrandsMultiAgentSession(root) + with caplog.at_level( + "WARNING", logger="strands_evals.experimental.redteam.strategies.target_session" + ): + s.restore(s.snapshot()) + assert not any("_resume_from_session" in rec.message for rec in caplog.records) diff --git a/tests/strands_evals/experimental/redteam/test_task.py b/tests/strands_evals/experimental/redteam/test_task.py index a8ca4bb9..566859a9 100644 --- a/tests/strands_evals/experimental/redteam/test_task.py +++ b/tests/strands_evals/experimental/redteam/test_task.py @@ -1,5 +1,6 @@ """Tests for _build_attacker_task.""" +from typing import Any from unittest.mock import MagicMock import pytest @@ -171,8 +172,8 @@ def test_task_fn_bare_callable_target_raises_type_error(): def test_task_fn_session_missing_trace_raises_type_error(): - """A session with all four methods but no ``trace`` is rejected up front, not - accepted and then crashed on ``session.trace`` (which the engine would swallow + """A session with all four methods but no `trace` is rejected up front, not + accepted and then crashed on `session.trace` (which the engine would swallow into a misleading 'defended' verdict).""" class _NoTrace: @@ -220,3 +221,52 @@ def test_task_fn_resets_agent_to_clean_baseline_per_case(): def test_max_allowed_turns_constant_present(): assert MAX_ALLOWED_TURNS >= 50 + + +def test_task_fn_routes_multi_agent_base_to_multi_agent_session(): + """A MultiAgentBase target is wrapped in StrandsMultiAgentSession (not + StrandsAgentSession) and the baseline is captured once at build time.""" + from strands.multiagent.base import MultiAgentBase + + from strands_evals.experimental.redteam.strategies.target_session import ( + StrandsMultiAgentSession, + _MultiAgentSnapshot, + ) + + class _FakeOrch(MultiAgentBase): + id = "root" + + def __init__(self): + self.nodes = {} # no leaves -- the strategy stub never invokes + self.serialize_calls = 0 + self.deserialize_calls = 0 + + async def invoke_async(self, task, invocation_state=None, **kwargs): + raise NotImplementedError + + def __call__(self, task, invocation_state=None, **kwargs): + return MagicMock(results={}, __str__=lambda self=None: "ok") + + def serialize_state(self): + self.serialize_calls += 1 + return {"type": "fake"} + + def deserialize_state(self, payload): + self.deserialize_calls += 1 + + root = _FakeOrch() + captured: dict[str, Any] = {} + + class _CapturingStrategy(_StubStrategy): + def run_attack(self, case, target_session, *, max_turns, model=None, **kwargs): + captured["session"] = target_session + return AttackRunResult(conversation=[], metadata={"turns_used": 0}) + + task = _build_attacker_task(root, _by_label(_CapturingStrategy())) + # baseline serialize_state captured once before any case runs + assert root.serialize_calls == 1 + task(_case("c0")) + # routed to the multi-agent session + assert isinstance(captured["session"], StrandsMultiAgentSession) + # baseline composite has the expected shape + assert isinstance(captured["session"]._baseline, _MultiAgentSnapshot)