FIX isolate TAP configuration and per-run state

pyrit/executor/attack/multi_turn/tree_of_attacks.py

@@ -157,6 +157,35 @@ def threshold(self) -> float:
         return self.objective_scorer.threshold  # type: ignore[ty:unresolved-attribute]
 
 
+@dataclass(frozen=True, slots=True)
+class _TAPAttackConfiguration:
+    """Immutable configuration for the TAP search and node execution."""
+
+    tree_width: int
+    tree_depth: int
+    branching_factor: int
+    on_topic_checking_enabled: bool
+    desired_response_prefix: str
+    batch_size: int
+
+    def __post_init__(self) -> None:
+        """
+        Validate the TAP search limits.
+
+        Raises:
+            ValueError: If a search limit is less than one.
+        """
+        validations = (
+            (self.tree_depth, "The tree depth must be at least 1."),
+            (self.tree_width, "The tree width must be at least 1."),
+            (self.branching_factor, "The branching factor must be at least 1."),
+            (self.batch_size, "The batch size must be at least 1."),
+        )
+        for value, message in validations:
+            if value < 1:
+                raise ValueError(message)
+
+
 @dataclass
 class TAPAttackContext(MultiTurnAttackContext[Any]):
     """
@@ -179,6 +208,9 @@ class TAPAttackContext(MultiTurnAttackContext[Any]):
     best_objective_score: Score | None = None
     best_adversarial_conversation_id: str | None = None
 
+    # Visualization parent for first-level nodes in this execution
+    visualization_root_id: str = "root"
+
     @property
     def conversation_id(self) -> str | None:
         """The best objective-target conversation, or the first active branch."""
@@ -1357,32 +1389,20 @@ def __init__(
             ``score_blocked_content=True`` on the objective scorer (requires
             ``prompt_metadata["partial_content"]`` on the blocked piece).
         """
-        # Validate tree parameters
-        if tree_depth < 1:
-            raise ValueError("The tree depth must be at least 1.")
-        if tree_width < 1:
-            raise ValueError("The tree width must be at least 1.")
-        if branching_factor < 1:
-            raise ValueError("The branching factor must be at least 1.")
-        if batch_size < 1:
-            raise ValueError("The batch size must be at least 1.")
+        self._configuration = _TAPAttackConfiguration(
+            tree_width=tree_width,
+            tree_depth=tree_depth,
+            branching_factor=branching_factor,
+            on_topic_checking_enabled=on_topic_checking_enabled,
+            desired_response_prefix=desired_response_prefix,
+            batch_size=batch_size,
+        )
 
         # Initialize base class
         super().__init__(objective_target=objective_target, logger=logger, context_type=TAPAttackContext)
 
         self._memory = CentralMemory.get_memory_instance()
 
-        # Store tree configuration
-        self._tree_width = tree_width
-        self._tree_depth = tree_depth
-        self._branching_factor = branching_factor
-        self._vis_root_id = "root"
-
-        # Store execution configuration
-        self._on_topic_checking_enabled = on_topic_checking_enabled
-        self._desired_response_prefix = desired_response_prefix
-        self._batch_size = batch_size
-
         # Initialize adversarial configuration
         self._adversarial_chat = attack_adversarial_config.target
 
@@ -1482,7 +1502,7 @@ def __init__(
         # Use the adversarial chat target for scoring, as in CrescendoAttack
         self._scoring_target = self._adversarial_chat
 
-        if self._on_topic_checking_enabled and not self._scoring_target:
+        if self._configuration.on_topic_checking_enabled and not self._scoring_target:
             raise ValueError("On-topic checking is enabled but no scoring target is available.")
 
         self._prompt_normalizer = prompt_normalizer or PromptNormalizer()
@@ -1573,6 +1593,7 @@ async def _setup_async(self, *, context: TAPAttackContext) -> None:
 
         context.tree_visualization = Tree()
         context.tree_visualization.create_node("Root", "root")
+        context.visualization_root_id = "root"
 
         context.nodes = []
         context.best_conversation_id = None
@@ -1586,10 +1607,10 @@ async def _setup_async(self, *, context: TAPAttackContext) -> None:
         context.executed_turns = get_prepended_turn_count(context.prepended_conversation)
 
         # Validate that prepended conversation doesn't exceed tree_depth
-        if context.executed_turns >= self._tree_depth:
+        if context.executed_turns >= self._configuration.tree_depth:
             raise ValueError(
                 f"Prepended conversation has {context.executed_turns} turns, "
-                f"which equals or exceeds tree_depth={self._tree_depth}. "
+                f"which equals or exceeds tree_depth={self._configuration.tree_depth}. "
                 f"Reduce prepended turns or increase tree_depth."
             )
 
@@ -1601,7 +1622,7 @@ async def _setup_async(self, *, context: TAPAttackContext) -> None:
             node_id = f"prepended_{turn}"
             context.tree_visualization.create_node(f"{turn}: (prepended)", node_id, parent=vis_parent)
             vis_parent = node_id
-        self._vis_root_id = vis_parent
+        context.visualization_root_id = vis_parent
 
     async def _perform_async(self, *, context: TAPAttackContext) -> TAPAttackResult:
         """
@@ -1629,11 +1650,13 @@ async def _perform_async(self, *, context: TAPAttackContext) -> TAPAttackResult:
         """
         self._logger.info(f"Starting TAP attack with objective: {context.objective}")
         self._logger.info(
-            f"Tree dimensions - Width: {self._tree_width}, Depth: {self._tree_depth}, "
-            f"Branching factor: {self._branching_factor}"
+            f"Tree dimensions - Width: {self._configuration.tree_width}, "
+            f"Depth: {self._configuration.tree_depth}, "
+            f"Branching factor: {self._configuration.branching_factor}"
         )
         self._logger.info(
-            f"Execution settings - Batch size: {self._batch_size}, On-topic checking: {self._on_topic_checking_enabled}"
+            f"Execution settings - Batch size: {self._configuration.batch_size}, "
+            f"On-topic checking: {self._configuration.on_topic_checking_enabled}"
         )
 
         # TAP Attack Execution Algorithm:
@@ -1656,9 +1679,9 @@ async def _perform_async(self, *, context: TAPAttackContext) -> TAPAttackResult:
         # Execute tree exploration iterations
         # Note: executed_turns is initialized in _setup_async with prepended conversation count
         # Start from executed_turns + 1 so prepended turns count toward tree_depth
-        for turn in range(context.executed_turns + 1, self._tree_depth + 1):
+        for turn in range(context.executed_turns + 1, self._configuration.tree_depth + 1):
             context.executed_turns = turn
-            self._logger.info(f"Starting TAP turn {turn}/{self._tree_depth}")
+            self._logger.info(f"Starting TAP turn {turn}/{self._configuration.tree_depth}")
 
             # Prepare nodes for current iteration
             await self._prepare_nodes_for_iteration_async(context)
@@ -1787,7 +1810,7 @@ async def _initialize_first_level_nodes_async(self, context: TAPAttackContext) -
             context.next_message is not None and self._modality_router.objective_target_requires_media_on_first_turn
         )
 
-        for i in range(self._tree_width):
+        for i in range(self._configuration.tree_width):
             # Historically only node 0 consumed next_message so sibling roots could
             # explore alternative starts. For edit-only objectives (no {text} path),
             # every root must receive seed media to build a valid first-turn request.
@@ -1804,7 +1827,7 @@ async def _initialize_first_level_nodes_async(self, context: TAPAttackContext) -
                 )
 
             context.nodes.append(node)
-            node._vis_node_id = self._vis_root_id
+            node._vis_node_id = context.visualization_root_id
 
         # Clear next_message after initialization (it's been used by the first node)
         context.next_message = None
@@ -1824,7 +1847,7 @@ def _branch_existing_nodes(self, context: TAPAttackContext) -> None:
         cloned_nodes = []
 
         for node in context.nodes:
-            for _ in range(self._branching_factor - 1):
+            for _ in range(self._configuration.branching_factor - 1):
                 cloned_node = node.duplicate()
                 # Add the adversarial chat conversation ID of the duplicated node to the context's tracking
                 context.related_conversations.add(
@@ -1863,12 +1886,12 @@ async def _send_prompts_to_all_nodes_async(self, context: TAPAttackContext) -> N
             node._vis_node_id = vis_id
 
         # Process nodes in batches
-        for batch_start in range(0, len(context.nodes), self._batch_size):
-            batch_end = min(batch_start + self._batch_size, len(context.nodes))
+        for batch_start in range(0, len(context.nodes), self._configuration.batch_size):
+            batch_end = min(batch_start + self._configuration.batch_size, len(context.nodes))
             batch_nodes = context.nodes[batch_start:batch_end]
 
             self._logger.debug(
-                f"Processing batch {batch_start // self._batch_size + 1} "
+                f"Processing batch {batch_start // self._configuration.batch_size + 1} "
                 f"(nodes {batch_start + 1}-{batch_end} of {len(context.nodes)})"
             )
 
@@ -1918,8 +1941,8 @@ def _prune_nodes_to_maintain_width(self, context: TAPAttackContext) -> None:
         completed_nodes = self._get_completed_nodes_sorted_by_score(context.nodes)
 
         # Keep nodes up to width limit
-        nodes_to_keep = completed_nodes[: self._tree_width]
-        nodes_to_prune = completed_nodes[self._tree_width :]
+        nodes_to_keep = completed_nodes[: self._configuration.tree_width]
+        nodes_to_prune = completed_nodes[self._configuration.tree_width :]
 
         # Mark pruned nodes in visualization and track their conversation IDs
         for node in nodes_to_prune:
@@ -2012,7 +2035,7 @@ def _create_attack_node(
             attack_strategy_name=self.__class__.__name__,
             modality_router=self._modality_router,
             memory_labels=context.memory_labels,
-            desired_response_prefix=self._desired_response_prefix,
+            desired_response_prefix=self._configuration.desired_response_prefix,
             parent_id=parent_id,
             prompt_normalizer=self._prompt_normalizer,
             initial_prompt=initial_prompt,
@@ -2109,7 +2132,7 @@ def _create_on_topic_scorer(self, objective: str) -> Scorer | None:
                 - `None` if `on_topic_checking_enabled` is `False` or no scoring_target
                 is available
         """
-        if not self._on_topic_checking_enabled:
+        if not self._configuration.on_topic_checking_enabled:
             return None
 
         return TrueFalseInverterScorer(

tests/unit/executor/attack/multi_turn/test_pair.py

@@ -101,10 +101,10 @@ def test_init_applies_pair_structural_defaults(self, objective_target, adversari
             attack_adversarial_config=adversarial_config,
         )
 
-        assert attack._tree_width == 3
-        assert attack._tree_depth == 5
-        assert attack._branching_factor == 1
-        assert attack._on_topic_checking_enabled is False
+        assert attack._configuration.tree_width == 3
+        assert attack._configuration.tree_depth == 5
+        assert attack._configuration.branching_factor == 1
+        assert attack._configuration.on_topic_checking_enabled is False
 
     def test_branching_factor_is_not_exposed_in_signature(self):
         """branching_factor is definitional for PAIR (always 1) and must not be a public init kwarg."""
@@ -126,17 +126,17 @@ def test_tree_width_override(self, objective_target, adversarial_config):
             attack_adversarial_config=adversarial_config,
             tree_width=7,
         )
-        assert attack._tree_width == 7
-        assert attack._branching_factor == 1
+        assert attack._configuration.tree_width == 7
+        assert attack._configuration.branching_factor == 1
 
     def test_tree_depth_override(self, objective_target, adversarial_config):
         attack = PAIRAttack(
             objective_target=objective_target,
             attack_adversarial_config=adversarial_config,
             tree_depth=12,
         )
-        assert attack._tree_depth == 12
-        assert attack._on_topic_checking_enabled is False
+        assert attack._configuration.tree_depth == 12
+        assert attack._configuration.on_topic_checking_enabled is False
 
     def test_is_subclass_of_tap(self):
         assert issubclass(PAIRAttack, TreeOfAttacksWithPruningAttack)