diff --git a/pyrit/analytics/scenario_statistics.py b/pyrit/analytics/scenario_statistics.py index 1af0195953..98b66e1cdd 100644 --- a/pyrit/analytics/scenario_statistics.py +++ b/pyrit/analytics/scenario_statistics.py @@ -160,8 +160,8 @@ class ScenarioAttempt: attempt_id: str total_retries: int result_role: AttackResultRole | None - attempt_index: int | None - child_attack_result_ids: tuple[str, ...] + attempt_index: int | None = None + child_attack_result_ids: tuple[str, ...] = () @property def attack_result_id(self) -> str: diff --git a/pyrit/backend/routes/scenarios.py b/pyrit/backend/routes/scenarios.py index 6ef2bdee5c..a573bf0d43 100644 --- a/pyrit/backend/routes/scenarios.py +++ b/pyrit/backend/routes/scenarios.py @@ -29,7 +29,14 @@ get_scenario_run_service, ) from pyrit.backend.services.scenario_service import get_scenario_service -from pyrit.models import ScenarioQueueSnapshot, ScenarioResult, ScenarioRunProgress, ScenarioRunState +from pyrit.models import ( + ScenarioHierarchyNode, + ScenarioHierarchyProjection, + ScenarioQueueSnapshot, + ScenarioResult, + ScenarioRunProgress, + ScenarioRunState, +) from pyrit.models.catalog import ( RegisteredScenario, RunScenarioRequest, @@ -375,6 +382,69 @@ async def get_scenario_run_progress( # pyrit-async-suffix-exempt return progress +@router.get( + "/runs/{scenario_result_id}/hierarchy", + response_model=ScenarioHierarchyProjection, + responses={ + 404: {"model": ProblemDetail, "description": "Run not found"}, + }, +) +async def get_scenario_run_hierarchy( # pyrit-async-suffix-exempt + *, + scenario_result_id: IdentifierStr, + max_nodes: int = Query(1000, ge=1, le=5000, description="Maximum nodes to assemble in the hierarchy tree"), +) -> ScenarioHierarchyProjection: + """ + Get the bounded result hierarchy projection for a scenario run. + + Returns: + ScenarioHierarchyProjection: Bounded hierarchy projection connecting orchestration parents, + target-facing children, and generic read-only conversations. + """ + service = get_scenario_run_service() + hierarchy = await service.get_run_hierarchy_async( + scenario_result_id=scenario_result_id, + max_nodes=max_nodes, + ) + if hierarchy is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Scenario run '{scenario_result_id}' not found", + ) + return hierarchy + + +@router.get( + "/runs/{scenario_result_id}/hierarchy/nodes/{attack_result_id}", + response_model=ScenarioHierarchyNode, + responses={ + 404: {"model": ProblemDetail, "description": "Node or run not found"}, + }, +) +async def get_scenario_run_hierarchy_node( # pyrit-async-suffix-exempt + *, + scenario_result_id: IdentifierStr, + attack_result_id: IdentifierStr, +) -> ScenarioHierarchyNode: + """ + Get single-node hierarchy details and generic conversation projections without full-run hydration. + + Returns: + ScenarioHierarchyNode: The target node with resolved immediate children and conversations. + """ + service = get_scenario_run_service() + node = await service.get_run_hierarchy_node_async( + scenario_result_id=scenario_result_id, + attack_result_id=attack_result_id, + ) + if node is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Hierarchy node '{attack_result_id}' not found in scenario run '{scenario_result_id}'", + ) + return node + + @router.post( "/runs/{scenario_result_id}/cancel", response_model=ScenarioRunSummary, diff --git a/pyrit/backend/services/scenario_progress_read_model.py b/pyrit/backend/services/scenario_progress_read_model.py index faf2528c2d..a3e027bc0b 100644 --- a/pyrit/backend/services/scenario_progress_read_model.py +++ b/pyrit/backend/services/scenario_progress_read_model.py @@ -6,7 +6,7 @@ import asyncio import logging from collections import OrderedDict -from collections.abc import AsyncGenerator, Iterable, Sequence +from collections.abc import AsyncGenerator, Iterable, Mapping, Sequence from contextlib import asynccontextmanager from dataclasses import dataclass, field from typing import Any, Literal @@ -23,18 +23,27 @@ from pyrit.memory import AttackResultKeysetCursor from pyrit.memory.memory_interface import MemoryInterface from pyrit.models import ( + AttackOutcome, AttackResult, AttackResultMetadata, + AttackResultRole, AttackTechniqueIdentifier, ComponentIdentifier, ScenarioAtomicGroupProgress, + ScenarioAttackResultConversations, ScenarioAttackResultDelta, ScenarioAttackTechniqueDetails, ScenarioComponentIdentity, + ScenarioConversationProjection, ScenarioDisplayGroupProgress, ScenarioExecutionUnit, + ScenarioHierarchyDiagnostics, + ScenarioHierarchyGroup, + ScenarioHierarchyNode, + ScenarioHierarchyProjection, ScenarioObjectiveScorer, ScenarioObjectiveScorerMetrics, + ScenarioOrchestrationParentSummary, ScenarioProgressCounts, ScenarioProgressResult, ScenarioProgressSummary, @@ -43,6 +52,7 @@ ScenarioRunPlanAtomicGroup, ScenarioRunPlanGroupKind, ScenarioRunPlanSeedGroup, + ScenarioRunState, ScenarioScorerIdentity, ScenarioSeedGroupProgress, ScenarioTechniqueProgress, @@ -50,6 +60,7 @@ ScorerIdentifier, project_behavioral_identity, ) +from pyrit.models.catalog.scenario import AttackErrorSummary, AttackRetrySummary from pyrit.score.scorer_evaluation.scorer_metrics_io import find_objective_metrics_by_eval_hash logger = logging.getLogger(__name__) @@ -64,6 +75,29 @@ # applied by ``project_behavioral_identity``. _TECHNIQUE_SEEDS_CHILD = "technique_seeds" _TECHNIQUE_SEED_DISPLAY_PARAMS = ("value", "data_type") +_MAX_HIERARCHY_DEPTH = 10 +_MAX_HIERARCHY_NODES = 1000 +_MAX_DIAGNOSTIC_ITEMS = 100 +_MAX_RETRIES_PER_NODE = 100 +_MAX_CHILD_REFERENCES = 500 + + +def _resolve_executed_turns(delta: ScenarioAttackResultDelta) -> int: + """ + Safely resolve executed turns without falling through on 0 or crashing on invalid metadata. + + Returns: + int: The resolved turn count. + """ + if delta.executed_turns is not None: + return delta.executed_turns + raw_val = delta.attack_metadata.get("executed_turns") + if raw_val is not None: + try: + return int(raw_val) + except (ValueError, TypeError): + pass + return 0 @dataclass(frozen=True, slots=True) @@ -74,6 +108,7 @@ class ScenarioProgressSnapshot: results: tuple[ScenarioProgressResult, ...] summary: ScenarioProgressSummary plan: ScenarioRunPlan + has_more: bool = False @dataclass(frozen=True, slots=True) @@ -103,6 +138,7 @@ class _ProgressCacheEntry: cursor: AttackResultKeysetCursor | None = None summary: ScenarioProgressSummary | None = None summary_state: _ProgressSummaryState | None = None + has_more: bool = True class ScenarioProgressReadModel: @@ -126,6 +162,7 @@ async def get_snapshot_async( active_group_ids: Sequence[str], terminal: bool, objective_scorer_identifier: ComponentIdentifier | None, + max_deltas: int | None = None, ) -> ScenarioProgressSnapshot: """ Refresh and return the mapped progress state for one run. @@ -137,6 +174,7 @@ async def get_snapshot_async( active_group_ids (Sequence[str]): Groups that are running. terminal (bool): Whether the run has stopped. objective_scorer_identifier (ComponentIdentifier | None): Objective scorer identity. + max_deltas (int | None): Maximum deltas to hydrate, bounding cold queries. Returns: ScenarioProgressSnapshot: Deltas, mapped results, summary, and effective plan. @@ -144,7 +182,12 @@ async def get_snapshot_async( async with self._refresh_run_async(scenario_result_id=scenario_result_id): entry = self._get_cache_entry(scenario_result_id=scenario_result_id, plan=plan, terminal=terminal) first_new_index = len(entry.results) - deltas_updated = await self._hydrate_new_deltas_async(scenario_result_id=scenario_result_id, entry=entry) + deltas_updated = await self._hydrate_new_deltas_async( + scenario_result_id=scenario_result_id, + entry=entry, + terminal=terminal, + max_deltas=max_deltas, + ) return self._build_snapshot( entry=entry, first_new_index=first_new_index, @@ -255,19 +298,34 @@ def _build_snapshot( results=tuple(entry.results), summary=entry.summary, plan=summary_plan, + has_more=entry.has_more, ) - async def _hydrate_new_deltas_async(self, *, scenario_result_id: str, entry: _ProgressCacheEntry) -> bool: + async def _hydrate_new_deltas_async( + self, + *, + scenario_result_id: str, + entry: _ProgressCacheEntry, + terminal: bool = False, + max_deltas: int | None = None, + ) -> bool: """ Hydrate persisted deltas while preserving their timestamp/ID keyset order. Returns: bool: Whether any delta was appended or replaced, so the summary rebuilds. """ + if max_deltas is not None and len(entry.deltas) >= max_deltas: + return False changed = False while True: + fetch_limit = ( + min(self._STORAGE_PAGE_SIZE, max_deltas - len(entry.deltas)) + if max_deltas is not None + else self._STORAGE_PAGE_SIZE + ) page, has_more = await self._memory.get_scenario_attack_result_deltas_async( - scenario_result_id=scenario_result_id, cursor=entry.cursor, limit=self._STORAGE_PAGE_SIZE + scenario_result_id=scenario_result_id, cursor=entry.cursor, limit=fetch_limit ) # The keyset is ordered by (timestamp, id), so an edited row whose timestamp # moved forward is emitted again after the cursor. Such a re-emission is an @@ -297,7 +355,11 @@ async def _hydrate_new_deltas_async(self, *, scenario_result_id: str, entry: _Pr timestamp=last.timestamp, attack_result_id=last.attack_result_id, ) + if max_deltas is not None and len(entry.deltas) >= max_deltas: + entry.has_more = has_more + return changed if not has_more: + entry.has_more = False return changed if not page: raise RuntimeError("Scenario progress storage returned an empty page with has_more=True.") @@ -771,9 +833,28 @@ def _map_progress_delta( child_attack_result_ids=ScenarioProgressReadModel._read_child_attack_result_ids( attack_metadata=delta.attack_metadata ), + related_conversation_ids=ScenarioProgressReadModel._read_related_conversation_ids(delta=delta), attempt_index=result_metadata.attempt_index, ) + @staticmethod + def _read_related_conversation_ids(*, delta: ScenarioAttackResultDelta) -> list[str]: + """ + Read related conversation IDs from typed delta or validate legacy metadata fallback. + + Returns: + list[str]: Validated list of string conversation IDs. + """ + raw_ids: Any = delta.related_conversation_ids + if isinstance(raw_ids, (list, tuple, set)): + valid: list[str] = [str(cid) for cid in raw_ids if isinstance(cid, str) and cid] + if valid: + return valid + fallback: Any = delta.attack_metadata.get("related_conversation_ids") + if isinstance(fallback, (list, tuple, set)): + return [str(cid) for cid in fallback if isinstance(cid, str) and cid] + return [] + @staticmethod def _read_child_attack_result_ids(*, attack_metadata: dict[str, Any]) -> list[str]: """ @@ -827,3 +908,691 @@ def _synthesize_legacy_plan(*, deltas: list[ScenarioAttackResultDelta]) -> Scena seen_seed_ids.add(mapped.seed_group_id) group.seed_group_ids.append(mapped.seed_group_id) return ScenarioRunPlan(atomic_groups=list(groups.values()), seed_groups=list(seeds.values())) + + @staticmethod + def build_scenario_hierarchy( + *, + scenario_result_id: str, + scenario_name: str, + scenario_version: int, + status: ScenarioRunState, + snapshot: ScenarioProgressSnapshot, + max_nodes: int = _MAX_HIERARCHY_NODES, + conversation_stats: Mapping[str, Any] | None = None, + overall_counts: ScenarioProgressCounts | None = None, + exact_total_errors: int | None = None, + exact_total_retries: int | None = None, + exact_retrying_attacks_count: int | None = None, + exact_group_counts: Mapping[str, ScenarioProgressCounts] | None = None, + parent_deltas: Sequence[ScenarioAttackResultDelta] = (), + unverified_candidate_ids: set[str] | Sequence[str] | None = None, + ) -> ScenarioHierarchyProjection: + """ + Assemble the scenario result hierarchy projection with cycle, depth, and node budget guards. + + Args: + scenario_result_id: Persisted scenario run ID. + scenario_name: Scenario name. + scenario_version: Scenario version. + status: Terminal or running state of the scenario. + snapshot: Hydrated progress snapshot containing deltas, mapped results, and summary. + max_nodes: Maximum nodes to assemble in the hierarchy tree. + conversation_stats: Optional mapping of conversation IDs to stats for real message count previews. + overall_counts: Optional exact run-wide progress counts overriding truncated snapshot totals. + exact_total_errors: Optional exact run-wide error count from database aggregates. + exact_total_retries: Optional exact run-wide retry count from database aggregates. + exact_retrying_attacks_count: Optional exact run-wide count of retrying attacks from aggregates. + exact_group_counts: Optional mapping of atomic group ID to exact counts from aggregates. + parent_deltas: Optional orchestration parent deltas to resolve roots and child references + beyond delta limits. + unverified_candidate_ids: Optional candidates whose parent status could not be verified + due to query clipping or budget exhaustion; excluded from root_results. + + Returns: + ScenarioHierarchyProjection: Bounded hierarchy projection with diagnostics. + """ + deltas_by_id: dict[str, ScenarioAttackResultDelta] = {d.attack_result_id: d for d in snapshot.deltas} + results_by_id: dict[str, ScenarioProgressResult] = {r.attack_result_id: r for r in snapshot.results} + + plan_lookup = ScenarioPlanLookup.from_plan(plan=snapshot.plan) + for p_delta in parent_deltas: + if p_delta.attack_result_id not in deltas_by_id: + deltas_by_id[p_delta.attack_result_id] = p_delta + if p_delta.attack_result_id not in results_by_id: + results_by_id[p_delta.attack_result_id] = ScenarioProgressReadModel._map_progress_delta( + delta=p_delta, plan_lookup=plan_lookup + ) + + referenced_child_ids: set[str] = set() + for result in results_by_id.values(): + for child_id in result.child_attack_result_ids: + referenced_child_ids.add(child_id) + + # A true root is not referenced as a child by any parent and does not carry an attempt_index + # (which designates a child attempt dispatched by an orchestration parent). + unverified_set = set(unverified_candidate_ids) if unverified_candidate_ids else set() + root_results = [ + r + for r in results_by_id.values() + if r.attack_result_id not in referenced_child_ids + and r.attempt_index is None + and r.attack_result_id not in unverified_set + ] + root_results.sort(key=lambda r: (r.timestamp, r.attack_result_id)) + + visited_ids: set[str] = set() + total_nodes_count = 0 + hierarchy_truncated = bool(snapshot.has_more) + + def _build_node( + *, + attack_result_id: str, + parent_id: str | None, + child_index: int | None, + active_path: set[str], + depth: int, + ) -> ScenarioHierarchyNode | None: + nonlocal total_nodes_count, hierarchy_truncated + delta = deltas_by_id.get(attack_result_id) + result = results_by_id.get(attack_result_id) + if delta is None or result is None: + return None + + if total_nodes_count >= max_nodes: + hierarchy_truncated = True + return None + + total_nodes_count += 1 + visited_ids.add(attack_result_id) + + retries = result.retries + retries_truncated = False + if len(retries) > _MAX_RETRIES_PER_NODE: + retries = retries[:_MAX_RETRIES_PER_NODE] + retries_truncated = True + + executed_turns = _resolve_executed_turns(delta) + outcome_reason = ( + delta.outcome_reason + or delta.attribution_data.get("outcome_reason") + or delta.attack_metadata.get("outcome_reason") + ) + + raw_child_ids = list(result.child_attack_result_ids) + child_refs_truncated = len(raw_child_ids) > _MAX_CHILD_REFERENCES + bounded_child_ids = raw_child_ids[:_MAX_CHILD_REFERENCES] if child_refs_truncated else raw_child_ids + + orchestration: ScenarioOrchestrationParentSummary | None = None + if result.result_role == AttackResultRole.ORCHESTRATION or result.child_attack_result_ids: + completion_policy = delta.attack_metadata.get("completion_policy") + orchestration = ScenarioOrchestrationParentSummary( + attack_result_id=result.attack_result_id, + result_role=result.result_role, + outcome=result.outcome, + outcome_reason=outcome_reason, + attempt_index=result.attempt_index, + execution_time_ms=result.execution_time_ms, + executed_turns=executed_turns, + completion_policy=str(completion_policy) if completion_policy is not None else None, + child_attack_result_ids=bounded_child_ids, + child_references_truncated=child_refs_truncated, + error_type=result.error_type, + error_message=result.error_message, + timestamp=result.timestamp, + total_retries=result.total_retries, + retries=retries, + retries_truncated=retries_truncated, + ) + + current_active_path = active_path | {attack_result_id} + children: list[ScenarioHierarchyNode] = [] + unresolved_children: list[str] = [] + + for idx, c_id in enumerate(result.child_attack_result_ids): + if c_id in current_active_path or depth >= _MAX_HIERARCHY_DEPTH or c_id not in deltas_by_id: + unresolved_children.append(c_id) + elif total_nodes_count >= max_nodes: + hierarchy_truncated = True + unresolved_children.append(c_id) + else: + child_node = _build_node( + attack_result_id=c_id, + parent_id=attack_result_id, + child_index=idx, + active_path=current_active_path, + depth=depth + 1, + ) + if child_node is not None: + children.append(child_node) + else: + unresolved_children.append(c_id) + + conversations: list[ScenarioConversationProjection] = [] + if result.conversation_id: + msg_count: int | None = None + messages_truncated: bool = True + if conversation_stats is not None and result.conversation_id in conversation_stats: + stat = conversation_stats[result.conversation_id] + c = getattr(stat, "message_count", None) + if c is None and isinstance(stat, int): + c = stat + if isinstance(c, int): + msg_count = c + messages_truncated = c > 0 + + conversations.append( + ScenarioConversationProjection( + conversation_id=result.conversation_id, + conversation_type=str(delta.attack_metadata.get("conversation_type") or "main"), + description=str(delta.attack_metadata.get("conversation_description") or "") or None, + message_count=msg_count, + messages=[], + messages_truncated=messages_truncated, + ) + ) + + candidate_ids = ( + result.related_conversation_ids + or getattr(result, "related_conversations", None) + or ScenarioProgressReadModel._read_related_conversation_ids(delta=delta) + ) + has_other_conversations = False + if candidate_ids: + for c in candidate_ids: + cid = getattr(c, "conversation_id", c) + if isinstance(cid, str) and cid and cid != result.conversation_id: + has_other_conversations = True + break + if has_other_conversations: + node_conversations_count = None + node_conversations_truncated = True + else: + node_conversations_count = len(conversations) + node_conversations_truncated = False + + unresolved_refs_truncated = len(unresolved_children) > _MAX_CHILD_REFERENCES + bounded_unresolved_children = ( + unresolved_children[:_MAX_CHILD_REFERENCES] if unresolved_refs_truncated else unresolved_children + ) + + return ScenarioHierarchyNode( + attack_result_id=result.attack_result_id, + conversation_id=result.conversation_id or "", + parent_attack_result_id=parent_id, + child_index=child_index, + result_role=result.result_role, + attempt_index=result.attempt_index, + atomic_attack_name=result.atomic_attack_name, + atomic_group_id=result.atomic_group_id, + seed_group_id=result.seed_group_id, + objective=delta.objective, + outcome=result.outcome, + outcome_reason=outcome_reason, + execution_time_ms=result.execution_time_ms, + executed_turns=executed_turns, + timestamp=result.timestamp, + score=result.score, + error_type=result.error_type, + error_message=result.error_message, + total_retries=result.total_retries, + retries=retries, + retries_truncated=retries_truncated, + child_attack_result_ids=bounded_child_ids, + child_references_truncated=child_refs_truncated, + children=children, + children_truncated=False, + unresolved_child_attack_result_ids=bounded_unresolved_children, + unresolved_child_references_truncated=unresolved_refs_truncated, + conversations=conversations, + conversations_count=node_conversations_count, + conversations_truncated=node_conversations_truncated, + orchestration=orchestration, + ) + + roots_by_group: dict[str, list[ScenarioHierarchyNode]] = {} + for root_result in root_results: + if total_nodes_count >= max_nodes: + hierarchy_truncated = True + break + root_node = _build_node( + attack_result_id=root_result.attack_result_id, + parent_id=None, + child_index=None, + active_path=set(), + depth=0, + ) + if root_node is not None: + roots_by_group.setdefault(root_node.atomic_group_id, []).append(root_node) + + # Handle unreached/disconnected components (including detached cycles) in a single linear pass + if total_nodes_count < max_nodes: + # Nodes that are unverified or descended from unverified candidates can never be emitted as roots + unverified_descendants: set[str] = set(unverified_set) + unverified_queue = list(unverified_set) + while unverified_queue: + curr = unverified_queue.pop() + res = results_by_id.get(curr) + if res: + for cid in res.child_attack_result_ids: + if cid not in unverified_descendants: + unverified_descendants.add(cid) + unverified_queue.append(cid) + + unreached = [ + r + for r in results_by_id.values() + if r.attack_result_id not in visited_ids and r.attack_result_id not in unverified_descendants + ] + unreached.sort( + key=lambda r: ( + r.timestamp, + r.attack_result_id, + ) + ) + for next_root in unreached: + if total_nodes_count >= max_nodes: + hierarchy_truncated = True + break + if next_root.attack_result_id in visited_ids: + continue + root_node = _build_node( + attack_result_id=next_root.attack_result_id, + parent_id=None, + child_index=None, + active_path=set(), + depth=0, + ) + if root_node is not None: + roots_by_group.setdefault(root_node.atomic_group_id, []).append(root_node) + else: + visited_ids.add(next_root.attack_result_id) + + if not hierarchy_truncated and ( + any(r.attack_result_id not in visited_ids for r in results_by_id.values()) or bool(unverified_set) + ): + hierarchy_truncated = True + + for g_roots in roots_by_group.values(): + g_roots.sort(key=lambda n: n.timestamp) + + groups: list[ScenarioHierarchyGroup] = [] + seen_group_ids: set[str] = set() + + for ag in snapshot.summary.atomic_groups: + seen_group_ids.add(ag.id) + group_roots = roots_by_group.get(ag.id, []) + if exact_group_counts and ag.id in exact_group_counts: + exact_g = exact_group_counts[ag.id] + counts = ScenarioProgressCounts( + completed=exact_g.completed, + planned=ag.planned, + succeeded=exact_g.succeeded, + success_percentage=exact_g.success_percentage, + errors=exact_g.errors, + retries=exact_g.retries, + producer_counts=exact_g.producer_counts, + outcomes=exact_g.outcomes, + ) + else: + counts = ScenarioProgressCounts( + completed=ag.completed, + planned=ag.planned, + succeeded=ag.succeeded, + success_percentage=ag.success_percentage, + errors=ag.errors, + retries=ag.retries, + producer_counts=ag.producer_counts, + outcomes=ag.outcomes, + ) + groups.append( + ScenarioHierarchyGroup( + atomic_group_id=ag.id, + atomic_attack_name=ag.atomic_attack_name, + display_group=ag.display_group, + technique=ag.technique_details, + roots=group_roots, + counts=counts, + ) + ) + + for gid, group_roots in roots_by_group.items(): + if gid not in seen_group_ids: + first_root = group_roots[0] + group_results = [r for r in snapshot.results if r.atomic_group_id == gid] + attempts_by_unit: dict[ResultUnitIdentity, list[ScenarioProgressResult]] = {} + for r in group_results: + unit = ResultUnitIdentity(atomic_group_id=r.atomic_group_id, seed_group_id=r.seed_group_id) + attempts_by_unit.setdefault(unit, []).append(r) + + if exact_group_counts and gid in exact_group_counts: + counts = exact_group_counts[gid] + else: + counts = count_execution_units( + units=list(attempts_by_unit.keys()), + attempts_by_unit=attempts_by_unit, + planned=None, + producer_attempts=group_results, + ) + groups.append( + ScenarioHierarchyGroup( + atomic_group_id=gid, + atomic_attack_name=first_root.atomic_attack_name, + display_group=first_root.atomic_attack_name, + technique=None, + roots=group_roots, + counts=counts, + ) + ) + + all_failed_attacks: list[AttackErrorSummary] = [] + all_attack_retries: list[AttackRetrySummary] = [] + + for delta, result in zip(snapshot.deltas, snapshot.results, strict=True): + if result.outcome == AttackOutcome.ERROR: + all_failed_attacks.append( + AttackErrorSummary( + atomic_attack_name=result.atomic_attack_name, + objective=delta.objective, + error_type=result.error_type, + error_message=result.error_message, + total_retries=result.total_retries, + ) + ) + if result.retries or (result.total_retries and result.total_retries > 0): + retries_list = list(result.retries) if result.retries else [] + retries_truncated = False + if len(retries_list) > _MAX_RETRIES_PER_NODE: + retries_list = retries_list[:_MAX_RETRIES_PER_NODE] + retries_truncated = True + all_attack_retries.append( + AttackRetrySummary( + attack_result_id=result.attack_result_id, + atomic_attack_name=result.atomic_attack_name, + total_retries=result.total_retries, + retries=retries_list, + retries_truncated=retries_truncated, + ) + ) + + failed_attacks_count = ( + max(len(all_failed_attacks), exact_total_errors) + if exact_total_errors is not None + else len(all_failed_attacks) + ) + failed_attacks_truncated = bool( + (snapshot.has_more and failed_attacks_count > len(all_failed_attacks)) + or failed_attacks_count > _MAX_DIAGNOSTIC_ITEMS + ) + failed_attacks = all_failed_attacks[:_MAX_DIAGNOSTIC_ITEMS] + + attempts_by_unit: dict[tuple[str, str], list[ScenarioProgressResult]] = {} + for r in snapshot.results: + attempts_by_unit.setdefault((r.atomic_group_id, r.seed_group_id), []).append(r) + + diagnostic_retries = max( + snapshot.summary.overall.retries, + ScenarioProgressReadModel.total_retry_pressure( + attempts_per_unit=(len(attempts) for attempts in attempts_by_unit.values()), + persisted_retries=(r.total_retries for r in snapshot.results), + ), + ) + if exact_total_retries is not None: + diagnostic_retries = max(diagnostic_retries, exact_total_retries) + + total_persisted_retries = sum(r.total_retries for r in snapshot.results) + attack_retries_count = ( + max(len(all_attack_retries), exact_retrying_attacks_count) + if exact_retrying_attacks_count is not None + else len(all_attack_retries) + ) + attack_retries_truncated = bool( + ( + snapshot.has_more + and (diagnostic_retries > total_persisted_retries or attack_retries_count > len(all_attack_retries)) + ) + or attack_retries_count > _MAX_DIAGNOSTIC_ITEMS + ) + attack_retries = all_attack_retries[:_MAX_DIAGNOSTIC_ITEMS] + + diagnostics = ScenarioHierarchyDiagnostics( + failed_attacks=failed_attacks, + failed_attacks_count=failed_attacks_count, + failed_attacks_truncated=failed_attacks_truncated, + attack_retries=attack_retries, + attack_retries_count=attack_retries_count, + attack_retries_truncated=attack_retries_truncated, + total_retries=diagnostic_retries, + ) + + effective_overall_counts = overall_counts if overall_counts is not None else snapshot.summary.overall + + return ScenarioHierarchyProjection( + scenario_result_id=scenario_result_id, + scenario_name=scenario_name, + scenario_version=scenario_version, + status=status, + groups=groups, + overall_counts=effective_overall_counts, + diagnostics=diagnostics, + truncated=hierarchy_truncated, + ) + + @staticmethod + def build_hierarchy_node( + *, + delta: ScenarioAttackResultDelta, + result: ScenarioProgressResult, + children_deltas: Sequence[ScenarioAttackResultDelta] = (), + children_results: Sequence[ScenarioProgressResult] = (), + conversations: Sequence[ScenarioConversationProjection] | ScenarioAttackResultConversations = (), + conversations_count: int | None = None, + conversations_truncated: bool = False, + children_truncated: bool = False, + max_child_references: int = _MAX_CHILD_REFERENCES, + ) -> ScenarioHierarchyNode: + """ + Assemble a single hierarchy node with its immediate children and generic conversations. + + Args: + delta: The attack result delta for the target node. + result: The mapped progress result for the target node. + children_deltas: Deltas for immediate children resolved in a single batch. + children_results: Mapped results for immediate children. + conversations: Generic read-only conversation projections. + conversations_count: Total conversations count for this node. + conversations_truncated: Flag indicating if conversations list was truncated. + children_truncated: Flag indicating if immediate children list was truncated. + max_child_references: Maximum number of child references to retain before truncating. + + Returns: + ScenarioHierarchyNode: The assembled node. + """ + conv_list: list[ScenarioConversationProjection] = [] + if isinstance(conversations, ScenarioAttackResultConversations): + conv_list = list(conversations.conversations) + if conversations_count is None: + conversations_count = conversations.total_count + conversations_truncated = conversations_truncated or conversations.truncated + else: + conv_list = list(conversations) + + resolved_conversations_count: int | None = ( + conversations_count if conversations_count is not None else len(conv_list) + ) + + children_deltas_by_id = {d.attack_result_id: d for d in children_deltas} + children_results_by_id = {r.attack_result_id: r for r in children_results} + + retries = result.retries + retries_truncated = False + if len(retries) > _MAX_RETRIES_PER_NODE: + retries = retries[:_MAX_RETRIES_PER_NODE] + retries_truncated = True + + executed_turns = _resolve_executed_turns(delta) + outcome_reason = ( + delta.outcome_reason + or delta.attribution_data.get("outcome_reason") + or delta.attack_metadata.get("outcome_reason") + ) + + raw_child_ids = list(result.child_attack_result_ids) + child_refs_truncated = len(raw_child_ids) > max_child_references + bounded_child_ids = raw_child_ids[:max_child_references] if child_refs_truncated else raw_child_ids + + orchestration: ScenarioOrchestrationParentSummary | None = None + if result.result_role == AttackResultRole.ORCHESTRATION or result.child_attack_result_ids: + completion_policy = delta.attack_metadata.get("completion_policy") + orchestration = ScenarioOrchestrationParentSummary( + attack_result_id=result.attack_result_id, + result_role=result.result_role, + outcome=result.outcome, + outcome_reason=outcome_reason, + attempt_index=result.attempt_index, + execution_time_ms=result.execution_time_ms, + executed_turns=executed_turns, + completion_policy=str(completion_policy) if completion_policy is not None else None, + child_attack_result_ids=bounded_child_ids, + child_references_truncated=child_refs_truncated, + error_type=result.error_type, + error_message=result.error_message, + timestamp=result.timestamp, + total_retries=result.total_retries, + retries=retries, + retries_truncated=retries_truncated, + ) + + children: list[ScenarioHierarchyNode] = [] + unresolved_children: list[str] = [] + + seen_child_ids: set[str] = set() + for idx, c_id in enumerate(result.child_attack_result_ids): + if c_id == result.attack_result_id or c_id in seen_child_ids or c_id not in children_deltas_by_id: + unresolved_children.append(c_id) + else: + seen_child_ids.add(c_id) + c_delta = children_deltas_by_id[c_id] + c_res = children_results_by_id[c_id] + c_retries = c_res.retries + c_retries_truncated = False + if len(c_retries) > _MAX_RETRIES_PER_NODE: + c_retries = c_retries[:_MAX_RETRIES_PER_NODE] + c_retries_truncated = True + + c_executed_turns = _resolve_executed_turns(c_delta) + c_outcome_reason = ( + c_delta.outcome_reason + or c_delta.attribution_data.get("outcome_reason") + or c_delta.attack_metadata.get("outcome_reason") + ) + + c_child_ids = list(c_res.child_attack_result_ids) + c_child_refs_truncated = len(c_child_ids) > max_child_references + c_bounded_child_ids = c_child_ids[:max_child_references] if c_child_refs_truncated else c_child_ids + + c_orchestration: ScenarioOrchestrationParentSummary | None = None + if c_res.result_role == AttackResultRole.ORCHESTRATION or c_res.child_attack_result_ids: + c_policy = c_delta.attack_metadata.get("completion_policy") + c_orchestration = ScenarioOrchestrationParentSummary( + attack_result_id=c_res.attack_result_id, + result_role=c_res.result_role, + outcome=c_res.outcome, + outcome_reason=c_outcome_reason, + attempt_index=c_res.attempt_index, + execution_time_ms=c_res.execution_time_ms, + executed_turns=c_executed_turns, + completion_policy=str(c_policy) if c_policy is not None else None, + child_attack_result_ids=c_bounded_child_ids, + child_references_truncated=c_child_refs_truncated, + error_type=c_res.error_type, + error_message=c_res.error_message, + timestamp=c_res.timestamp, + total_retries=c_res.total_retries, + retries=c_retries, + retries_truncated=c_retries_truncated, + ) + + c_has_related = bool( + c_res.related_conversation_ids + or ScenarioProgressReadModel._read_related_conversation_ids(delta=c_delta) + ) + if c_res.conversation_id or c_has_related: + c_conv_count = None + c_conv_truncated = True + else: + c_conv_count = 0 + c_conv_truncated = False + + children.append( + ScenarioHierarchyNode( + attack_result_id=c_res.attack_result_id, + conversation_id=c_res.conversation_id or "", + parent_attack_result_id=result.attack_result_id, + child_index=idx, + result_role=c_res.result_role, + attempt_index=c_res.attempt_index, + atomic_attack_name=c_res.atomic_attack_name, + atomic_group_id=c_res.atomic_group_id, + seed_group_id=c_res.seed_group_id, + objective=c_delta.objective, + outcome=c_res.outcome, + outcome_reason=c_outcome_reason, + execution_time_ms=c_res.execution_time_ms, + executed_turns=c_executed_turns, + timestamp=c_res.timestamp, + score=c_res.score, + error_type=c_res.error_type, + error_message=c_res.error_message, + total_retries=c_res.total_retries, + retries=c_retries, + retries_truncated=c_retries_truncated, + child_attack_result_ids=c_bounded_child_ids, + child_references_truncated=c_child_refs_truncated, + children=[], + children_truncated=False, + unresolved_child_attack_result_ids=[], + conversations=[], + conversations_count=c_conv_count, + conversations_truncated=c_conv_truncated, + orchestration=c_orchestration, + ) + ) + + unresolved_refs_truncated = len(unresolved_children) > max_child_references + bounded_unresolved_children = ( + unresolved_children[:max_child_references] if unresolved_refs_truncated else unresolved_children + ) + + return ScenarioHierarchyNode( + attack_result_id=result.attack_result_id, + conversation_id=result.conversation_id or "", + parent_attack_result_id=None, + child_index=None, + result_role=result.result_role, + attempt_index=result.attempt_index, + atomic_attack_name=result.atomic_attack_name, + atomic_group_id=result.atomic_group_id, + seed_group_id=result.seed_group_id, + objective=delta.objective, + outcome=result.outcome, + outcome_reason=outcome_reason, + execution_time_ms=result.execution_time_ms, + executed_turns=executed_turns, + timestamp=result.timestamp, + score=result.score, + error_type=result.error_type, + error_message=result.error_message, + total_retries=result.total_retries, + retries=retries, + retries_truncated=retries_truncated, + child_attack_result_ids=bounded_child_ids, + child_references_truncated=child_refs_truncated, + children=children, + children_truncated=children_truncated, + unresolved_child_attack_result_ids=bounded_unresolved_children, + unresolved_child_references_truncated=unresolved_refs_truncated, + conversations=conv_list, + conversations_count=resolved_conversations_count, + conversations_truncated=conversations_truncated, + orchestration=orchestration, + ) diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index 286605d9f9..49fdae90ec 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -23,11 +23,13 @@ from pydantic import TypeAdapter, ValidationError -from pyrit.analytics.outcome_statistics import success_percentage +from pyrit.analytics.outcome_statistics import compute_outcome_statistics, success_percentage from pyrit.analytics.scenario_statistics import ( ScenarioAttempt, - compute_scenario_statistics, + ScenarioExecutionUnit, + count_execution_units, resolve_attack_result_attempt, + resolve_execution_unit, select_execution_attempts, ) from pyrit.backend.models.common import PaginationInfo, filter_sensitive_fields @@ -41,12 +43,15 @@ from pyrit.backend.services.scenario_configuration_resolver import ScenarioConfigurationResolver from pyrit.backend.services.scenario_progress_read_model import ( ResultUnitIdentity, + ScenarioPlanLookup, ScenarioProgressReadModel, ScenarioProgressSnapshot, ) from pyrit.common.async_compatibility import run_legacy_sync_async from pyrit.memory import AttackResultKeysetCursor, CentralMemory, SQLiteMemory from pyrit.memory.memory_interface import ( + MemoryInterface, + ScenarioAttemptSeedSummary, ScenarioHistoryAggregate, ScenarioHistoryKeysetCursor, ScenarioHistoryRunRecord, @@ -54,12 +59,20 @@ from pyrit.models import ( SCENARIO_RUN_PLAN_METADATA_KEY, SCENARIO_RUN_STARTED_AT_METADATA_KEY, + AtomicAttackIdentifier, AttackOutcome, AttackResult, + AttackResultMetadata, ComponentIdentifier, ScenarioAttackResultDelta, + ScenarioHierarchyNode, + ScenarioHierarchyProjection, ScenarioIdentifier, + ScenarioProducerCategoryCounts, + ScenarioProducerCounts, + ScenarioProgressCounts, ScenarioProgressHeader, + ScenarioProgressResult, ScenarioQueueEntry, ScenarioQueueSnapshot, ScenarioResult, @@ -196,7 +209,13 @@ class ScenarioRunService: multiple replicas require a shared database-backed scheduler or lease. """ - def __init__(self, *, max_concurrent_runs: int = _DEFAULT_MAX_CONCURRENT_RUNS) -> None: + def __init__( + self, + *, + max_concurrent_runs: int = _DEFAULT_MAX_CONCURRENT_RUNS, + memory: MemoryInterface | None = None, + progress_read_model: ScenarioProgressReadModel | None = None, + ) -> None: """ Initialize the scenario run service. @@ -205,10 +224,12 @@ def __init__(self, *, max_concurrent_runs: int = _DEFAULT_MAX_CONCURRENT_RUNS) - """ if max_concurrent_runs < 1: raise ValueError("max_concurrent_runs must be at least 1.") - self._memory = CentralMemory.get_memory_instance() + self._memory = memory if memory is not None else CentralMemory.get_memory_instance() self._active_tasks: dict[str, _ActiveTask] = {} self._configuration_resolver = ScenarioConfigurationResolver() - self._progress_read_model = ScenarioProgressReadModel(memory=self._memory) + self._progress_read_model = ( + progress_read_model if progress_read_model is not None else ScenarioProgressReadModel(memory=self._memory) + ) self._technique_metadata_cache: dict[str, dict[str, ScenarioTechniqueSummary]] = {} self._preparations: set[asyncio.Task[_PreparedRun]] = set() @@ -1556,11 +1577,16 @@ def _build_run_summary( retry_events = getattr(attack_result, "retry_events", None) if isinstance(retry_events, list) and retry_events: overload_events.extend(retry_events) + truncated_retries = retry_events[:100] + retries_truncated = len(retry_events) > 100 + total_r = max(0, retries) if isinstance(retries, int) else len(retry_events) attack_retries.append( AttackRetrySummary( attack_result_id=str(attack_result.attack_result_id), atomic_attack_name=atomic_attack_name, - retries=retry_events, + total_retries=total_r, + retries=truncated_retries, + retries_truncated=retries_truncated, ) ) @@ -1671,6 +1697,7 @@ async def _recount_with_sdk_statistics_async( *, aggregates: dict[str, ScenarioHistoryAggregate], plans: Mapping[str, list[ScenarioRunPlanAtomicGroup] | None], + max_deltas: int | None = None, ) -> dict[str, ScenarioHistoryAggregate]: """ Recount runs the SQL aggregate can't resolve with the shared statistics. @@ -1686,10 +1713,157 @@ async def _recount_with_sdk_statistics_async( if not run_ids: return aggregates recounted = dict(aggregates) - for scenario_result in await self._memory.get_scenario_results_async(scenario_result_ids=run_ids): - run_id = str(scenario_result.id) - # A plan this service rejected counts as a legacy run, as in the SQL path. - overall = compute_scenario_statistics(scenario_result, use_saved_plan=plans.get(run_id) is not None).overall + for run_id in run_ids: + cursor: AttackResultKeysetCursor | None = None + all_summaries: list[ScenarioAttemptSeedSummary] = [] + while True: + remaining_limit = min(500, max_deltas - len(all_summaries)) if max_deltas is not None else 500 + if remaining_limit <= 0: + break + page, has_more = await self._memory.get_scenario_recount_summaries_async( + scenario_result_id=run_id, + cursor=cursor, + limit=remaining_limit, + ) + all_summaries.extend(page) + if not has_more or not page or (max_deltas is not None and len(all_summaries) >= max_deltas): + break + cursor = AttackResultKeysetCursor( + timestamp=page[-1].timestamp, attack_result_id=page[-1].attack_result_id + ) + + saved_plan: ScenarioRunPlan | None = None + header = await self._memory.get_scenario_result_header_async(scenario_result_id=run_id) + if header is not None: + try: + saved_plan = self._load_run_plan(scenario_result=header) + except (ValidationError, ValueError): + logger.warning( + "Scenario run %s has an invalid saved run plan; recounting as legacy run.", + run_id, + ) + plan_lookup = ScenarioPlanLookup.from_plan(plan=saved_plan) + attempts: list[ScenarioAttempt] = [] + for summary in all_summaries: + atomic_name = str(summary.attribution_data.get("parent_collection") or "") + eval_hash = summary.attribution_data.get("parent_eval_hash") + attr_seed = summary.attribution_data.get("seed_group_id") + typed_identifier = ( + AtomicAttackIdentifier.from_component_identifier(summary.atomic_attack_identifier) + if isinstance(summary.atomic_attack_identifier, ComponentIdentifier) + else None + ) + unit = resolve_execution_unit( + atomic_attack_name=atomic_name, + technique_eval_hash=str(eval_hash) if eval_hash is not None else None, + attributed_seed_group_id=str(attr_seed) if attr_seed else None, + atomic_attack_identifier=typed_identifier, + objective=summary.objective, + objective_sha256=summary.objective_sha256, + plan_lookup=plan_lookup, + ) + res_metadata = AttackResultMetadata.from_metadata(metadata=summary.attribution_data) + attempts.append( + ScenarioAttempt( + unit=unit, + atomic_attack_name=atomic_name, + outcome=summary.outcome, + timestamp=summary.timestamp, + attempt_id=summary.attack_result_id, + total_retries=summary.total_retries, + result_role=res_metadata.result_role, + attempt_index=res_metadata.attempt_index, + child_attack_result_ids=(), + ) + ) + attempts.sort(key=lambda a: (a.timestamp, a.attempt_id)) + attempts_by_unit: dict[ScenarioExecutionUnit, list[ScenarioAttempt]] = {} + for a in attempts: + attempts_by_unit.setdefault(a.unit, []).append(a) + + if saved_plan is not None: + planned_units = list( + dict.fromkeys( + ScenarioExecutionUnit(atomic_group_id=group.id, seed_group_id=seed_id) + for group in saved_plan.atomic_groups + for seed_id in group.seed_group_ids + ) + ) + counted_units = planned_units + planned = len(planned_units) + else: + counted_units = list(attempts_by_unit.keys()) + planned = None + + overall = count_execution_units( + units=counted_units, + attempts_by_unit=attempts_by_unit, + planned=planned, + producer_attempts=attempts, + ) + + recounted_group_counts: dict[str, ScenarioProgressCounts] = {} + if saved_plan is not None: + for group in saved_plan.atomic_groups: + g_units = [ + ScenarioExecutionUnit(atomic_group_id=group.id, seed_group_id=seed_id) + for seed_id in group.seed_group_ids + ] + g_producer_attempts = [a for a in attempts if a.unit.atomic_group_id == group.id] + g_counts = count_execution_units( + units=g_units, + attempts_by_unit=attempts_by_unit, + planned=len(g_units), + producer_attempts=g_producer_attempts, + ) + recounted_group_counts[group.id] = g_counts + + unplanned_group_ids = {a.unit.atomic_group_id for a in attempts} - { + g.id for g in saved_plan.atomic_groups + } + for gid in unplanned_group_ids: + g_units = [u for u in attempts_by_unit if u.atomic_group_id == gid] + g_producer_attempts = [a for a in attempts if a.unit.atomic_group_id == gid] + g_counts = count_execution_units( + units=g_units, + attempts_by_unit=attempts_by_unit, + planned=None, + producer_attempts=g_producer_attempts, + ) + recounted_group_counts[gid] = g_counts + else: + for gid in {a.unit.atomic_group_id for a in attempts}: + g_units = [u for u in attempts_by_unit if u.atomic_group_id == gid] + g_producer_attempts = [a for a in attempts if a.unit.atomic_group_id == gid] + g_counts = count_execution_units( + units=g_units, + attempts_by_unit=attempts_by_unit, + planned=None, + producer_attempts=g_producer_attempts, + ) + recounted_group_counts[gid] = g_counts + + is_truncated = max_deltas is not None and len(all_summaries) >= max_deltas + raw_total_errors = ( + aggregates[run_id].raw_total_errors + if is_truncated + else sum(1 for d in all_summaries if d.outcome == AttackOutcome.ERROR) + ) + raw_total_retries = ( + aggregates[run_id].raw_total_retries + if is_truncated + else ScenarioProgressReadModel.total_retry_pressure( + attempts_per_unit=(len(a) for a in attempts_by_unit.values()), + persisted_retries=(d.total_retries for d in all_summaries), + ) + ) + raw_retrying_attacks_count = ( + aggregates[run_id].raw_retrying_attacks_count + if is_truncated + else sum(1 for d in all_summaries if d.has_retries or d.total_retries > 0) + ) + merged_group_counts = {**aggregates[run_id].group_counts, **recounted_group_counts} + recounted[run_id] = replace( aggregates[run_id], unit_count=overall.completed, @@ -1697,6 +1871,14 @@ async def _recount_with_sdk_statistics_async( successful_units=overall.succeeded, error_attempts=overall.errors, total_retries=overall.retries, + outcome_successes=overall.outcomes.successes if overall.outcomes else 0, + outcome_failures=overall.outcomes.failures if overall.outcomes else 0, + outcome_errors=overall.outcomes.errors if overall.outcomes else 0, + outcome_undetermined=overall.outcomes.undetermined if overall.outcomes else 0, + raw_total_errors=raw_total_errors, + raw_total_retries=raw_total_retries, + raw_retrying_attacks_count=raw_retrying_attacks_count, + group_counts=merged_group_counts, target_facing_attempts=overall.producer_counts.target_facing.attempts, target_facing_error_attempts=overall.producer_counts.target_facing.errors, target_facing_retries=overall.producer_counts.target_facing.retries, @@ -2310,6 +2492,357 @@ def _decode_progress_cursor( raise ValueError("Cursor timestamp must include a timezone.") return AttackResultKeysetCursor(timestamp=timestamp, attack_result_id=attack_result_id) + async def get_run_hierarchy_async( + self, + *, + scenario_result_id: str, + max_nodes: int = 1000, + ) -> ScenarioHierarchyProjection | None: + """ + Get the bounded result hierarchy projection for a scenario run. + + Args: + scenario_result_id: The persisted scenario run ID. + max_nodes: Maximum nodes to assemble in the hierarchy tree. + + Returns: + ScenarioHierarchyProjection | None: The hierarchy projection, or None if the run does not exist. + """ + header_result = await self._memory.get_scenario_result_header_async(scenario_result_id=scenario_result_id) + if header_result is None: + return None + + try: + plan = self._load_run_plan(scenario_result=header_result) + except (ValidationError, ValueError): + logger.warning( + "Scenario run %s has invalid persisted plan metadata; treating the plan as unavailable.", + scenario_result_id, + ) + plan = None + plan_complete = plan is not None + terminal = header_result.scenario_run_state in ( + ScenarioRunState.COMPLETED, + ScenarioRunState.FAILED, + ScenarioRunState.CANCELLED, + ) + snapshot = self.snapshot_active_run(scenario_result_id=scenario_result_id) + objective_scorer_identifier = header_result.objective_scorer_identifier + if not isinstance(objective_scorer_identifier, ComponentIdentifier): + objective_scorer_identifier = None + + progress_snapshot = await self._progress_read_model.get_snapshot_async( + scenario_result_id=scenario_result_id, + plan=plan, + plan_complete=plan_complete, + active_group_ids=snapshot.active_group_ids, + terminal=terminal, + objective_scorer_identifier=objective_scorer_identifier, + max_deltas=max(max_nodes, self._progress_read_model._STORAGE_PAGE_SIZE), + ) + + exact_overall_counts: ScenarioProgressCounts | None = None + exact_total_errors: int | None = None + exact_total_retries: int | None = None + exact_retrying_attacks_count: int | None = None + exact_group_counts: Mapping[str, ScenarioProgressCounts] | None = None + parent_deltas: list[ScenarioAttackResultDelta] = [] + unverified_candidate_ids: set[str] = set() + + if progress_snapshot.has_more: + plan_lookup = ScenarioPlanLookup.from_plan(plan=plan) + all_deltas_map: dict[str, ScenarioAttackResultDelta] = { + d.attack_result_id: d for d in progress_snapshot.deltas + } + all_results_map: dict[str, ScenarioProgressResult] = { + r.attack_result_id: r for r in progress_snapshot.results + } + + referenced_child_ids: set[str] = set() + for r in all_results_map.values(): + for cid in r.child_attack_result_ids: + referenced_child_ids.add(cid) + + unreferenced_roots = [ + r.attack_result_id for r in all_results_map.values() if r.attack_result_id not in referenced_child_ids + ] + reachable: set[str] = set() + queue = list(unreferenced_roots) + while queue: + curr = queue.pop() + if curr not in reachable: + reachable.add(curr) + res = all_results_map.get(curr) + if res: + queue.extend(res.child_attack_result_ids) + + unreached_cyclic = [cid for cid in all_results_map if cid not in reachable] + current_candidates = list(unreferenced_roots) if unreferenced_roots else list(unreached_cyclic) + + parent_budget = max(2500, self._progress_read_model._STORAGE_PAGE_SIZE * 5) + checked_child_ids: set[str] = set() + depth = 0 + layer_candidates = list(current_candidates) + while layer_candidates and depth < 10 and len(parent_deltas) < parent_budget: + depth += 1 + next_layer_candidates: list[str] = [] + layer_clipped = False + while len(parent_deltas) < parent_budget: + remaining_budget = max(0, parent_budget - len(parent_deltas)) + if remaining_budget <= 0: + break + page_limit = min(self._progress_read_model._STORAGE_PAGE_SIZE, remaining_budget) + new_parents, has_more = await self._memory.get_scenario_parent_deltas_for_child_ids_async( + scenario_result_id=scenario_result_id, + child_ids=layer_candidates, + limit=page_limit, + exclude_result_ids=list(all_deltas_map.keys()), + ) + for p in new_parents: + if p.attack_result_id not in all_deltas_map: + parent_deltas.append(p) + all_deltas_map[p.attack_result_id] = p + p_res = self._progress_read_model._map_progress_delta(delta=p, plan_lookup=plan_lookup) + all_results_map[p.attack_result_id] = p_res + next_layer_candidates.append(p.attack_result_id) + if len(parent_deltas) >= parent_budget: + break + + if len(parent_deltas) >= parent_budget: + layer_clipped = True + break + + if not has_more: + break + + if not new_parents: + layer_clipped = True + break + + if layer_clipped: + unverified_candidate_ids.update(layer_candidates) + unverified_candidate_ids.update(next_layer_candidates) + break + + checked_child_ids.update(layer_candidates) + layer_candidates = [cid for cid in next_layer_candidates if cid not in checked_child_ids] + + if not layer_candidates and len(parent_deltas) < parent_budget and depth < 10: + remaining_cyclic = [cid for cid in unreached_cyclic if cid not in checked_child_ids] + if remaining_cyclic: + layer_candidates = remaining_cyclic + + if layer_candidates: + unverified_candidate_ids.update(layer_candidates) + + # Any node in all_results_map that is not referenced as a child by another loaded node + # and whose parents were not queried to completion in the DB is unverified. + all_referenced_cids: set[str] = {cid for r in all_results_map.values() for cid in r.child_attack_result_ids} + for rid in all_results_map: + if rid not in all_referenced_cids and rid not in checked_child_ids: + unverified_candidate_ids.add(rid) + + plan_ids = [scenario_result_id] if plan is not None else [] + aggregates = await self._memory.get_scenario_history_aggregates_async( + scenario_result_ids=[scenario_result_id], + plan_scenario_ids=plan_ids, + ) + plans_map = {scenario_result_id: plan.atomic_groups if plan is not None else None} + aggregates = await self._recount_with_sdk_statistics_async( + aggregates=aggregates, + plans=plans_map, + ) + agg = aggregates.get(scenario_result_id) + if agg is not None: + producer_counts = ScenarioProducerCounts( + target_facing=ScenarioProducerCategoryCounts( + attempts=agg.target_facing_attempts, + errors=agg.target_facing_error_attempts, + retries=agg.target_facing_retries, + ), + orchestration=ScenarioProducerCategoryCounts( + attempts=agg.orchestration_attempts, + errors=agg.orchestration_error_attempts, + retries=agg.orchestration_retries, + ), + unknown=ScenarioProducerCategoryCounts( + attempts=agg.unknown_role_attempts, + errors=agg.unknown_role_error_attempts, + retries=agg.unknown_role_retries, + ), + ) + outcomes = compute_outcome_statistics( + { + AttackOutcome.SUCCESS: agg.outcome_successes, + AttackOutcome.FAILURE: agg.outcome_failures, + AttackOutcome.ERROR: agg.outcome_errors, + AttackOutcome.UNDETERMINED: agg.outcome_undetermined, + } + ) + exact_overall_counts = ScenarioProgressCounts( + planned=progress_snapshot.summary.overall.planned, + completed=agg.completed_units, + succeeded=agg.successful_units, + success_percentage=success_percentage( + succeeded=agg.successful_units, completed=agg.completed_units + ), + errors=agg.error_attempts, + retries=agg.total_retries, + producer_counts=producer_counts, + outcomes=outcomes, + ) + exact_total_errors = agg.raw_total_errors + exact_total_retries = agg.raw_total_retries + exact_retrying_attacks_count = agg.raw_retrying_attacks_count + exact_group_counts = agg.group_counts + + hierarchy = self._progress_read_model.build_scenario_hierarchy( + scenario_result_id=scenario_result_id, + scenario_name=header_result.scenario_name, + scenario_version=header_result.scenario_version, + status=header_result.scenario_run_state, + snapshot=progress_snapshot, + max_nodes=max_nodes, + conversation_stats=None, + overall_counts=exact_overall_counts, + exact_total_errors=exact_total_errors, + exact_total_retries=exact_total_retries, + exact_retrying_attacks_count=exact_retrying_attacks_count, + exact_group_counts=exact_group_counts, + parent_deltas=parent_deltas, + unverified_candidate_ids=unverified_candidate_ids, + ) + + # Obtain real conversation counts for the actual bounded preview nodes + seen_conv_ids: set[str] = set() + actual_conv_ids: list[str] = [] + + def _collect_node_conv_ids(node: ScenarioHierarchyNode) -> None: + if node.conversation_id and node.conversation_id not in seen_conv_ids: + seen_conv_ids.add(node.conversation_id) + actual_conv_ids.append(node.conversation_id) + for c in node.conversations: + if c.conversation_id and c.conversation_id not in seen_conv_ids: + seen_conv_ids.add(c.conversation_id) + actual_conv_ids.append(c.conversation_id) + for child in node.children: + _collect_node_conv_ids(child) + + for group in hierarchy.groups: + for root in group.roots: + _collect_node_conv_ids(root) + + if actual_conv_ids: + conversation_stats: dict[str, Any] | None = None + try: + conversation_stats = await self._memory.get_conversation_stats_async(conversation_ids=actual_conv_ids) + except Exception: + logger.warning("Could not fetch conversation statistics for hierarchy preview.", exc_info=True) + conversation_stats = None + + if conversation_stats is not None: + queried_set = set(actual_conv_ids) + + def _apply_stats(node: ScenarioHierarchyNode) -> None: + for conv in node.conversations: + if conv.conversation_id in queried_set: + if conv.conversation_id in conversation_stats: + stat = conversation_stats[conv.conversation_id] + c = getattr(stat, "message_count", None) + if c is None and isinstance(stat, int): + c = stat + if isinstance(c, int): + conv.message_count = c + conv.messages_truncated = c > 0 + else: + # Queried successfully, but no messages exist in memory + conv.message_count = 0 + conv.messages_truncated = False + for child in node.children: + _apply_stats(child) + + for group in hierarchy.groups: + for root in group.roots: + _apply_stats(root) + + return hierarchy + + async def get_run_hierarchy_node_async( + self, + *, + scenario_result_id: str, + attack_result_id: str, + max_immediate_children: int = 200, + ) -> ScenarioHierarchyNode | None: + """ + Get bounded single-node hierarchy details and generic conversation projections without full-run hydration. + + Args: + scenario_result_id: The persisted scenario run ID. + attack_result_id: The target attack result ID. + max_immediate_children: Maximum number of immediate child nodes to resolve. + + Returns: + ScenarioHierarchyNode | None: The node details, or None if the node is not found in this run. + """ + deltas = await self._memory.get_scenario_attack_result_deltas_by_ids_async( + scenario_result_id=scenario_result_id, + attack_result_ids=[attack_result_id], + ) + if not deltas: + return None + + target_delta = deltas[0] + + header_result = await self._memory.get_scenario_result_header_async(scenario_result_id=scenario_result_id) + plan = None + if header_result is not None: + try: + plan = self._load_run_plan(scenario_result=header_result) + except (ValidationError, ValueError): + plan = None + + plan_lookup = ScenarioPlanLookup.from_plan(plan=plan) + target_result = self._progress_read_model._map_progress_delta( + delta=target_delta, + plan_lookup=plan_lookup, + ) + + child_ids = target_result.child_attack_result_ids + children_deltas: list[ScenarioAttackResultDelta] = [] + children_results: list[ScenarioProgressResult] = [] + children_truncated = False + if child_ids: + if len(child_ids) > max_immediate_children: + children_truncated = True + fetch_child_ids = child_ids[:max_immediate_children] + else: + fetch_child_ids = child_ids + + children_deltas = list( + await self._memory.get_scenario_attack_result_deltas_by_ids_async( + scenario_result_id=scenario_result_id, + attack_result_ids=fetch_child_ids, + ) + ) + children_results = [ + self._progress_read_model._map_progress_delta(delta=d, plan_lookup=plan_lookup) for d in children_deltas + ] + + conversations = await self._memory.get_attack_result_conversation_projections_async( + attack_result_id=attack_result_id, + include_messages=True, + ) + + return self._progress_read_model.build_hierarchy_node( + delta=target_delta, + result=target_result, + children_deltas=children_deltas, + children_results=children_results, + conversations=conversations, + children_truncated=children_truncated, + ) + async def get_run_results_async(self, *, scenario_result_id: str) -> ScenarioResult | None: """ Get the ScenarioResult for a completed scenario run. diff --git a/pyrit/memory/__init__.py b/pyrit/memory/__init__.py index 42aae06309..50e8af9f68 100644 --- a/pyrit/memory/__init__.py +++ b/pyrit/memory/__init__.py @@ -19,6 +19,7 @@ from pyrit.memory.memory_interface import ( AttackResultKeysetCursor, MemoryInterface, + ScenarioAttemptSeedSummary, ScenarioHistoryAggregate, ScenarioHistoryKeysetCursor, ScenarioHistoryRunRecord, @@ -62,6 +63,7 @@ "ImagePathDataTypeSerializer": "pyrit.memory.storage", "MemoryInterface": "pyrit.memory.memory_interface", "MemoryEmbedding": "pyrit.memory.memory_embedding", + "ScenarioAttemptSeedSummary": "pyrit.memory.memory_interface", "ScenarioHistoryKeysetCursor": "pyrit.memory.memory_interface", "ScenarioHistoryRunRecord": "pyrit.memory.memory_interface", "ScenarioHistoryAggregate": "pyrit.memory.memory_interface", diff --git a/pyrit/memory/azure_sql_memory.py b/pyrit/memory/azure_sql_memory.py index abf3fcb9c4..76353222ef 100644 --- a/pyrit/memory/azure_sql_memory.py +++ b/pyrit/memory/azure_sql_memory.py @@ -693,57 +693,61 @@ def _execute_get_conversation_stats(self, *, conversation_ids: Sequence[str]) -> if not conversation_ids: return {} - placeholders = ", ".join(f":cid{i}" for i in range(len(conversation_ids))) - params = {f"cid{i}": cid for i, cid in enumerate(conversation_ids)} - - sql = text( - f""" - SELECT - aggregate_rows.conversation_id, - aggregate_rows.msg_count, - latest.last_preview, - latest.last_data_type, - aggregate_rows.created_at - FROM ( - SELECT - pme.conversation_id, - COUNT(DISTINCT pme.sequence) AS msg_count, - MIN(pme.timestamp) AS created_at - FROM "PromptMemoryEntries" pme - WHERE pme.conversation_id IN ({placeholders}) - GROUP BY pme.conversation_id - ) AS aggregate_rows - OUTER APPLY ( - SELECT TOP 1 - LEFT(p2.converted_value, {ConversationStats.PREVIEW_FETCH_MAX_LEN}) AS last_preview, - p2.converted_value_data_type AS last_data_type - FROM "PromptMemoryEntries" p2 - WHERE p2.conversation_id = aggregate_rows.conversation_id - ORDER BY p2.sequence DESC, p2.id DESC - ) AS latest - """ - ) + batch_size = max(1, self._MAX_BIND_VARS) + result: dict[str, ConversationStats] = {} with closing(self._get_session()) as session: - rows = session.execute(sql, params).fetchall() + for start in range(0, len(conversation_ids), batch_size): + chunk = conversation_ids[start : start + batch_size] + placeholders = ", ".join(f":cid{i}" for i in range(len(chunk))) + params = {f"cid{i}": cid for i, cid in enumerate(chunk)} + + sql = text( + f""" + SELECT + aggregate_rows.conversation_id, + aggregate_rows.msg_count, + latest.last_preview, + latest.last_data_type, + aggregate_rows.created_at + FROM ( + SELECT + pme.conversation_id, + COUNT(DISTINCT pme.sequence) AS msg_count, + MIN(pme.timestamp) AS created_at + FROM "PromptMemoryEntries" pme + WHERE pme.conversation_id IN ({placeholders}) + GROUP BY pme.conversation_id + ) AS aggregate_rows + OUTER APPLY ( + SELECT TOP 1 + LEFT(p2.converted_value, {ConversationStats.PREVIEW_FETCH_MAX_LEN}) AS last_preview, + p2.converted_value_data_type AS last_data_type + FROM "PromptMemoryEntries" p2 + WHERE p2.conversation_id = aggregate_rows.conversation_id + ORDER BY p2.sequence DESC, p2.id DESC + ) AS latest + """ + ) - result: dict[str, ConversationStats] = {} - for row in rows: - conv_id, msg_count, last_preview, last_data_type, raw_created_at = row + rows = session.execute(sql, params).fetchall() - created_at = None - if raw_created_at is not None: - if isinstance(raw_created_at, str): - created_at = datetime.fromisoformat(raw_created_at) - else: - created_at = raw_created_at + for row in rows: + conv_id, msg_count, last_preview, last_data_type, raw_created_at = row - result[conv_id] = ConversationStats( - message_count=msg_count, - last_message_preview=last_preview, - last_message_data_type=last_data_type, - created_at=created_at, - ) + created_at = None + if raw_created_at is not None: + if isinstance(raw_created_at, str): + created_at = datetime.fromisoformat(raw_created_at) + else: + created_at = raw_created_at + + result[conv_id] = ConversationStats( + message_count=msg_count, + last_message_preview=last_preview, + last_message_data_type=last_data_type, + created_at=created_at, + ) return result diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index b522f82562..8dfd57723b 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -15,7 +15,7 @@ from collections.abc import Callable, Collection, Iterator, Mapping, MutableSequence, Sequence from contextlib import closing from contextvars import ContextVar -from dataclasses import dataclass +from dataclasses import dataclass, field from datetime import UTC, datetime, timedelta from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Literal, NamedTuple, ParamSpec, TypeVar, cast @@ -24,6 +24,7 @@ from sqlalchemy import ( Column, MetaData, + Select, String, Unicode, and_, @@ -52,6 +53,7 @@ from sqlalchemy.sql import visitors from sqlalchemy.sql.elements import ClauseElement +from pyrit.analytics.outcome_statistics import compute_outcome_statistics, success_percentage from pyrit.common.async_compatibility import legacy_sync_override, run_legacy_sync_async from pyrit.common.deprecation import print_deprecation_message from pyrit.common.pagination import DecodedKeysetCursor @@ -125,8 +127,15 @@ Operation, PromptDataType, RetryEvent, + ScenarioAttackResultConversations, ScenarioAttackResultDelta, + ScenarioConversationMessagePieceProjection, + ScenarioConversationMessageProjection, + ScenarioConversationProjection, ScenarioIdentifier, + ScenarioProducerCategoryCounts, + ScenarioProducerCounts, + ScenarioProgressCounts, ScenarioProgressScore, ScenarioResult, ScenarioRunState, @@ -144,6 +153,7 @@ SeedRecord, SeedType, TargetIdentifier, + config_hash, group_conversation_message_pieces_by_sequence, sort_message_pieces, ) @@ -359,6 +369,17 @@ class ScenarioHistoryAggregate: unknown_role_error_attempts: int unknown_role_retries: int + outcome_successes: int = 0 + outcome_failures: int = 0 + outcome_errors: int = 0 + outcome_undetermined: int = 0 + + raw_total_errors: int = 0 + raw_total_retries: int = 0 + raw_retrying_attacks_count: int = 0 + + group_counts: dict[str, ScenarioProgressCounts] = field(default_factory=dict) + @classmethod def empty(cls, *, scenario_result_id: str) -> "ScenarioHistoryAggregate": """ @@ -385,6 +406,14 @@ def empty(cls, *, scenario_result_id: str) -> "ScenarioHistoryAggregate": unknown_role_attempts=0, unknown_role_error_attempts=0, unknown_role_retries=0, + outcome_successes=0, + outcome_failures=0, + outcome_errors=0, + outcome_undetermined=0, + raw_total_errors=0, + raw_total_retries=0, + raw_retrying_attacks_count=0, + group_counts={}, ) @@ -396,6 +425,27 @@ class ScenarioRunStateRecord: state: ScenarioRunState +@dataclass(frozen=True, slots=True, kw_only=True) +class ScenarioAttemptSeedSummary: + """ + Lightweight projection of attempt seed and outcome fields for scenario recount. + + Contains only the scalar fields required to resolve logical execution units and + aggregate attempt counts during legacy recount. Never loads prompt entries, score + records, retry events, or large conversation payloads. + """ + + attack_result_id: str + objective: str + objective_sha256: str | None + atomic_attack_identifier: AtomicAttackIdentifier | None + outcome: AttackOutcome + timestamp: datetime + total_retries: int + attribution_data: dict[str, Any] + has_retries: bool = False + + @dataclass(frozen=True, slots=True, kw_only=True) class _AttackResultQuery: """ @@ -6981,6 +7031,15 @@ def _execute_get_scenario_history_aggregates( .mappings() .all() ) + group_rows = ( + session.execute( + self._build_scenario_history_group_aggregate_statement( + entry_ids=entry_ids, plan_entry_ids=plan_entry_ids + ) + ) + .mappings() + .all() + ) name_rows = session.execute( select(AttackResultEntry.attribution_parent_id, self._get_scenario_attempt_unit_expressions()[0]) .where(AttackResultEntry.attribution_parent_id.in_(entry_ids)) @@ -7000,6 +7059,58 @@ def _execute_get_scenario_history_aggregates( ).all() } + groups_by_run: dict[str, dict[str, ScenarioProgressCounts]] = {} + for grow in group_rows: + if grow["scenario_result_id"] is None or not grow["unit_group_id"]: + continue + run_id = str(grow["scenario_result_id"]) + raw_gid = str(grow["unit_group_id"]) + canonical_gid = raw_gid + if "\x1f" in raw_gid: + name, _, eval_hash = raw_gid.partition("\x1f") + canonical_gid = config_hash({"atomic_attack_name": name, "technique_eval_hash": eval_hash}) + g_outcomes = compute_outcome_statistics( + { + AttackOutcome.SUCCESS: grow["outcome_successes"] or 0, + AttackOutcome.FAILURE: grow["outcome_failures"] or 0, + AttackOutcome.ERROR: grow["outcome_errors"] or 0, + AttackOutcome.UNDETERMINED: grow["outcome_undetermined"] or 0, + } + ) + g_producer_counts = ScenarioProducerCounts( + target_facing=ScenarioProducerCategoryCounts( + attempts=grow["target_facing_attempts"] or 0, + errors=grow["target_facing_error_attempts"] or 0, + retries=grow["target_facing_retries"] or 0, + ), + orchestration=ScenarioProducerCategoryCounts( + attempts=grow["orchestration_attempts"] or 0, + errors=grow["orchestration_error_attempts"] or 0, + retries=grow["orchestration_retries"] or 0, + ), + unknown=ScenarioProducerCategoryCounts( + attempts=grow["unknown_role_attempts"] or 0, + errors=grow["unknown_role_error_attempts"] or 0, + retries=grow["unknown_role_retries"] or 0, + ), + ) + g_completed = grow["completed_units"] or 0 + g_succeeded = grow["successful_units"] or 0 + counts = ScenarioProgressCounts( + completed=g_completed, + planned=None, + succeeded=g_succeeded, + success_percentage=success_percentage(succeeded=g_succeeded, completed=g_completed), + errors=grow["error_attempts"] or 0, + retries=grow["total_retries"] or 0, + producer_counts=g_producer_counts, + outcomes=g_outcomes, + ) + run_groups = groups_by_run.setdefault(run_id, {}) + run_groups[canonical_gid] = counts + if canonical_gid != raw_gid: + run_groups[raw_gid] = counts + names_by_run: dict[str, list[str]] = {} for scenario_result_id, atomic_attack_name in name_rows: if scenario_result_id is None or not atomic_attack_name: @@ -7028,22 +7139,32 @@ def _execute_get_scenario_history_aggregates( unknown_role_attempts=row["unknown_role_attempts"] or 0, unknown_role_error_attempts=row["unknown_role_error_attempts"] or 0, unknown_role_retries=row["unknown_role_retries"] or 0, + outcome_successes=row["outcome_successes"] or 0, + outcome_failures=row["outcome_failures"] or 0, + outcome_errors=row["outcome_errors"] or 0, + outcome_undetermined=row["outcome_undetermined"] or 0, + raw_total_errors=row["raw_total_errors"] or 0, + raw_total_retries=row["raw_total_retries"] or 0, + raw_retrying_attacks_count=row["raw_retrying_attacks_count"] or 0, + group_counts=groups_by_run.get(run_id, {}), ) return aggregates - def _build_scenario_history_aggregate_statement( + def _build_scenario_history_ranked_units( self, *, entry_ids: Sequence[uuid.UUID], plan_entry_ids: Sequence[uuid.UUID], - ) -> Any: + ) -> tuple[Any, Any, Any, Any, Any, Any]: """ - Build the statement that reduces linked attempts to one metrics row per run. + Build the intermediate ranked-units subquery and role match conditions for scenario history. Returns: - Any: A statement selecting one aggregate row per scenario run with attempts. + tuple: (ranked, counted, completed, is_target_facing, is_orchestration, is_unknown) """ - atomic_name, technique_hash, attributed_seed_group_id, _ = self._get_scenario_attempt_unit_expressions() + atomic_name, technique_hash, attributed_seed_group_id, identifier_seed_key = ( + self._get_scenario_attempt_unit_expressions() + ) parents = AttackResultEntry.__table__.alias("history_parents") children = AttackResultEntry.__table__.alias("history_children") links = JsonArrayItems(parents.c.attack_metadata, literal("$.child_attack_result_ids")).table_valued( @@ -7060,9 +7181,6 @@ def _build_scenario_history_aggregate_statement( .distinct() .subquery("history_duplicate_errors") ) - # Explicit seed attribution is authoritative, then the objective hash. Runs with attempts identified only by - # their atomic identifier's seeds are flagged instead (see _get_scenario_history_runs_needing_sdk_statistics), - # because SQL can't derive the logical seed group from those seeds. attempts = ( select( AttackResultEntry.id.label("attempt_id"), @@ -7070,7 +7188,9 @@ def _build_scenario_history_aggregate_statement( atomic_name.label("atomic_attack_name"), technique_hash.label("technique_eval_hash"), attributed_seed_group_id.label("attributed_seed_group_id"), - func.coalesce(attributed_seed_group_id, AttackResultEntry.objective_sha256, "").label("seed_group_id"), + func.coalesce( + attributed_seed_group_id, identifier_seed_key, AttackResultEntry.objective_sha256, "" + ).label("seed_group_id"), AttackResultEntry.objective_sha256.label("objective_sha256"), AttackResultEntry.outcome.label("outcome"), AttackResultEntry.timestamp.label("timestamp"), @@ -7081,6 +7201,16 @@ def _build_scenario_history_aggregate_statement( ), else_=0, ).label("total_retries"), + case( + ( + or_( + func.coalesce(AttackResultEntry.total_retries, 0) > 0, + AttackResultEntry.retry_events_json.is_not(None), + ), + 1, + ), + else_=0, + ).label("has_retries"), JsonScalar(AttackResultEntry.attribution_data, literal("$.result_role")).label("result_role"), JsonScalar(AttackResultEntry.attribution_data, literal("$.attempt_index")).label("attempt_index"), case((duplicate_errors.c.parent_id.is_not(None), 1), else_=0).label("duplicates_child_error"), @@ -7132,8 +7262,11 @@ def _build_scenario_history_aggregate_statement( ) ranked = select( units.c.scenario_result_id, + units.c.unit_group_id, units.c.timestamp, units.c.outcome.label("latest_outcome"), + units.c.outcome.label("attempt_outcome"), + units.c.has_retries.label("has_retries"), func.max(units.c.is_planned).over(partition_by=unit_partition).label("is_planned"), unit_retries.label("unit_retries"), func.sum(case((is_error, pressure_attempt), else_=0)) @@ -7142,6 +7275,7 @@ def _build_scenario_history_aggregate_statement( is_completion.label("is_completion"), units.c.result_role.label("result_role"), units.c.total_retries.label("attempt_retries"), + func.max(units.c.has_retries).over(partition_by=unit_partition).label("unit_has_retries"), func.row_number() .over( partition_by=unit_partition, @@ -7165,6 +7299,23 @@ def _build_scenario_history_aggregate_statement( ranked.c.result_role.is_(None), and_(not_(is_target_facing), not_(is_orchestration)), ) + return ranked, counted, completed, is_target_facing, is_orchestration, is_unknown + + def _build_scenario_history_aggregate_statement( + self, + *, + entry_ids: Sequence[uuid.UUID], + plan_entry_ids: Sequence[uuid.UUID], + ) -> Any: + """ + Build the statement that reduces linked attempts to one metrics row per run. + + Returns: + Any: A statement selecting one aggregate row per scenario run with attempts. + """ + ranked, counted, completed, is_target_facing, is_orchestration, is_unknown = ( + self._build_scenario_history_ranked_units(entry_ids=entry_ids, plan_entry_ids=plan_entry_ids) + ) return ( select( ranked.c.scenario_result_id, @@ -7178,6 +7329,25 @@ def _build_scenario_history_aggregate_statement( func.sum(case((and_(counted, ranked.c.unit_retries > 0), ranked.c.unit_retries), else_=0)).label( "total_retries" ), + func.sum( + case((and_(completed, ranked.c.latest_outcome == AttackOutcome.SUCCESS.value), 1), else_=0) + ).label("outcome_successes"), + func.sum( + case((and_(completed, ranked.c.latest_outcome == AttackOutcome.FAILURE.value), 1), else_=0) + ).label("outcome_failures"), + func.sum( + case((and_(completed, ranked.c.latest_outcome == AttackOutcome.ERROR.value), 1), else_=0) + ).label("outcome_errors"), + func.sum( + case((and_(completed, ranked.c.latest_outcome == AttackOutcome.UNDETERMINED.value), 1), else_=0) + ).label("outcome_undetermined"), + func.sum(case((ranked.c.attempt_outcome == AttackOutcome.ERROR.value, 1), else_=0)).label( + "raw_total_errors" + ), + func.sum( + case((and_(ranked.c.unit_rank == 1, ranked.c.unit_retries > 0), ranked.c.unit_retries), else_=0) + ).label("raw_total_retries"), + func.sum(case((ranked.c.has_retries == 1, 1), else_=0)).label("raw_retrying_attacks_count"), func.sum(case((is_target_facing, 1), else_=0)).label("target_facing_attempts"), func.sum( case( @@ -7207,6 +7377,74 @@ def _build_scenario_history_aggregate_statement( .order_by(ranked.c.scenario_result_id) ) + def _build_scenario_history_group_aggregate_statement( + self, + *, + entry_ids: Sequence[uuid.UUID], + plan_entry_ids: Sequence[uuid.UUID], + ) -> Any: + """ + Build the statement that reduces linked attempts to metrics rows per group. + + Returns: + Any: A statement selecting aggregate rows grouped by scenario run and atomic group. + """ + ranked, counted, completed, is_target_facing, is_orchestration, is_unknown = ( + self._build_scenario_history_ranked_units(entry_ids=entry_ids, plan_entry_ids=plan_entry_ids) + ) + return ( + select( + ranked.c.scenario_result_id, + ranked.c.unit_group_id, + func.sum(case((completed, 1), else_=0)).label("completed_units"), + func.sum( + case((and_(completed, ranked.c.latest_outcome == AttackOutcome.SUCCESS.value), 1), else_=0) + ).label("successful_units"), + func.sum(case((counted, ranked.c.unit_errors), else_=0)).label("error_attempts"), + func.sum(case((and_(counted, ranked.c.unit_retries > 0), ranked.c.unit_retries), else_=0)).label( + "total_retries" + ), + func.sum( + case((and_(completed, ranked.c.latest_outcome == AttackOutcome.SUCCESS.value), 1), else_=0) + ).label("outcome_successes"), + func.sum( + case((and_(completed, ranked.c.latest_outcome == AttackOutcome.FAILURE.value), 1), else_=0) + ).label("outcome_failures"), + func.sum( + case((and_(completed, ranked.c.latest_outcome == AttackOutcome.ERROR.value), 1), else_=0) + ).label("outcome_errors"), + func.sum( + case((and_(completed, ranked.c.latest_outcome == AttackOutcome.UNDETERMINED.value), 1), else_=0) + ).label("outcome_undetermined"), + func.sum(case((is_target_facing, 1), else_=0)).label("target_facing_attempts"), + func.sum( + case( + (and_(is_target_facing, ranked.c.latest_outcome == AttackOutcome.ERROR.value), 1), + else_=0, + ) + ).label("target_facing_error_attempts"), + func.sum(case((is_target_facing, ranked.c.attempt_retries), else_=0)).label("target_facing_retries"), + func.sum(case((is_orchestration, 1), else_=0)).label("orchestration_attempts"), + func.sum( + case( + (and_(is_orchestration, ranked.c.latest_outcome == AttackOutcome.ERROR.value), 1), + else_=0, + ) + ).label("orchestration_error_attempts"), + func.sum(case((is_orchestration, ranked.c.attempt_retries), else_=0)).label("orchestration_retries"), + func.sum(case((is_unknown, 1), else_=0)).label("unknown_role_attempts"), + func.sum( + case( + (and_(is_unknown, ranked.c.latest_outcome == AttackOutcome.ERROR.value), 1), + else_=0, + ) + ).label("unknown_role_error_attempts"), + func.sum(case((is_unknown, ranked.c.attempt_retries), else_=0)).label("unknown_role_retries"), + ) + .group_by(ranked.c.scenario_result_id, ranked.c.unit_group_id) + .order_by(ranked.c.scenario_result_id, ranked.c.unit_group_id) + ) + def _build_scenario_history_unit_statement(self, *, attempts: Any, plan_entry_ids: Sequence[uuid.UUID]) -> Any: """ Resolve every persisted attempt to the logical work unit whose counters it feeds. @@ -7228,6 +7466,7 @@ def _build_scenario_history_unit_statement(self, *, attempts: Any, plan_entry_id attempts.c.outcome, attempts.c.timestamp, attempts.c.total_retries, + attempts.c.has_retries, attempts.c.result_role, attempts.c.attempt_index, attempts.c.duplicates_child_error, @@ -7310,6 +7549,7 @@ def _build_scenario_history_unit_statement(self, *, attempts: Any, plan_entry_id attempts.c.outcome, attempts.c.timestamp, attempts.c.total_retries, + attempts.c.has_retries, attempts.c.result_role, attempts.c.attempt_index, attempts.c.duplicates_child_error, @@ -7338,6 +7578,7 @@ def _build_scenario_history_unit_statement(self, *, attempts: Any, plan_entry_id matched.c.outcome, matched.c.timestamp, matched.c.total_retries, + matched.c.has_retries, matched.c.result_role, matched.c.attempt_index, matched.c.duplicates_child_error, @@ -7381,6 +7622,60 @@ def _execute_get_unique_scenario_labels(self) -> dict[str, list[str]]: label_values.setdefault(key, set()).add(value) return {key: sorted(values) for key, values in sorted(label_values.items())} + @staticmethod + def _row_to_scenario_attack_result_delta(row: Mapping[str, Any]) -> ScenarioAttackResultDelta: + retry_events = [ + RetryEvent.model_validate(event) + for event in (json.loads(row["retry_events_json"]) if row["retry_events_json"] else []) + ] + atomic_identifier = ( + AtomicAttackIdentifier.model_validate(row["atomic_attack_identifier"]) + if row["atomic_attack_identifier"] + else None + ) + score = None + if row["score_id"] is not None: + scorer_identifier = ( + ComponentIdentifier.model_validate(row["scorer_class_identifier"]) + if row["scorer_class_identifier"] + else None + ) + score = ScenarioProgressScore( + scorer_name=scorer_identifier.class_name if scorer_identifier else "Unknown", + score_type=row["score_type"], + status=ScoreStatus(row["score_status"]), + score_value=row["score_value"], + score_rationale=row["score_rationale"], + ) + related_cids: list[str] = [] + for col in ( + row["pruned_conversation_ids"], + row["adversarial_chat_conversation_ids"], + row["preparation_conversation_ids"], + ): + if isinstance(col, list): + related_cids.extend(str(c) for c in col if c) + return ScenarioAttackResultDelta( + attack_result_id=str(row["id"]), + conversation_id=row["conversation_id"], + objective=row["objective"], + objective_sha256=row["objective_sha256"], + atomic_attack_identifier=atomic_identifier, + outcome=AttackOutcome(row["outcome"]), + execution_time_ms=row["execution_time_ms"], + executed_turns=row["executed_turns"], + outcome_reason=row["outcome_reason"], + timestamp=row["timestamp"], + retry_events=retry_events, + total_retries=row["total_retries"] or 0, + error_type=row["error_type"], + error_message=row["error_message"], + attribution_data=row["attribution_data"] or {}, + attack_metadata=row["attack_metadata"] or {}, + related_conversation_ids=related_cids, + score=score, + ) + def _execute_get_scenario_attack_result_deltas( self, *, @@ -7431,6 +7726,8 @@ def _execute_get_scenario_attack_result_deltas( AttackResultEntry.atomic_attack_identifier, AttackResultEntry.outcome, AttackResultEntry.execution_time_ms, + AttackResultEntry.executed_turns, + AttackResultEntry.outcome_reason, AttackResultEntry.timestamp, AttackResultEntry.retry_events_json, AttackResultEntry.total_retries, @@ -7438,6 +7735,9 @@ def _execute_get_scenario_attack_result_deltas( AttackResultEntry.error_message, AttackResultEntry.attribution_data, AttackResultEntry.attack_metadata, + AttackResultEntry.pruned_conversation_ids, + AttackResultEntry.adversarial_chat_conversation_ids, + AttackResultEntry.preparation_conversation_ids, ScoreEntry.id.label("score_id"), ScoreEntry.score_value, ScoreEntry.score_type, @@ -7457,51 +7757,720 @@ def _execute_get_scenario_attack_result_deltas( rows = session.execute(statement).mappings().all() has_more = len(rows) > limit + deltas: list[ScenarioAttackResultDelta] = [ + self._row_to_scenario_attack_result_delta(row) for row in rows[:limit] + ] + return deltas, has_more + + @staticmethod + def _row_to_scenario_attempt_seed_summary(row: Mapping[str, Any]) -> ScenarioAttemptSeedSummary: + atomic_identifier = ( + AtomicAttackIdentifier.model_validate(row["atomic_attack_identifier"]) + if row["atomic_attack_identifier"] + else None + ) + return ScenarioAttemptSeedSummary( + attack_result_id=str(row["id"]), + objective=row["objective"], + objective_sha256=row["objective_sha256"], + atomic_attack_identifier=atomic_identifier, + outcome=AttackOutcome(row["outcome"]), + timestamp=row["timestamp"], + total_retries=row["total_retries"] or 0, + attribution_data=row["attribution_data"] or {}, + has_retries=bool(row.get("has_retries", False)), + ) + + def _execute_get_scenario_recount_summaries( + self, + *, + scenario_result_id: str, + cursor: AttackResultKeysetCursor | None = None, + limit: int = 500, + ) -> tuple[list[ScenarioAttemptSeedSummary], bool]: + """ + Return bounded scenario attempt seed summaries in ascending keyset order for recount. + + This projection selects only minimal scalar fields required for execution unit + resolution and recount metrics without loading score entries, prompt rows, or + retry events. + + Returns: + tuple[list[ScenarioAttemptSeedSummary], bool]: The page and whether more rows exist. + + Raises: + ValueError: If the limit or cursor identifiers are invalid. + """ + if limit < 1 or limit > 500: + raise ValueError("Scenario recount limit must be between 1 and 500.") + + scenario_uuid = uuid.UUID(scenario_result_id) + conditions: list[Any] = [AttackResultEntry.attribution_parent_id == scenario_uuid] + attempt_id_order = self._get_scenario_attempt_id_order_expression(attempt_id=AttackResultEntry.id) + if cursor is not None: + cursor_uuid = uuid.UUID(cursor.attack_result_id) + cursor_id_order = self._get_scenario_attempt_id_order_expression( + attempt_id=literal(cursor_uuid, type_=AttackResultEntry.id.type) + ) + conditions.append( + or_( + AttackResultEntry.timestamp > cursor.timestamp, + and_( + AttackResultEntry.timestamp == cursor.timestamp, + attempt_id_order > cursor_id_order, + ), + ) + ) + + statement = ( + select( + AttackResultEntry.id, + AttackResultEntry.objective, + AttackResultEntry.objective_sha256, + AttackResultEntry.atomic_attack_identifier, + AttackResultEntry.outcome, + AttackResultEntry.timestamp, + AttackResultEntry.total_retries, + AttackResultEntry.attribution_data, + case( + ( + or_( + func.coalesce(AttackResultEntry.total_retries, 0) > 0, + AttackResultEntry.retry_events_json.is_not(None), + ), + 1, + ), + else_=0, + ).label("has_retries"), + ) + .where(and_(*conditions)) + .order_by(AttackResultEntry.timestamp.asc(), attempt_id_order.asc()) + .limit(limit + 1) + ) + with closing(self._get_session()) as session: + rows = session.execute(statement).mappings().all() + + has_more = len(rows) > limit + summaries: list[ScenarioAttemptSeedSummary] = [ + self._row_to_scenario_attempt_seed_summary(row) for row in rows[:limit] + ] + return summaries, has_more + + def _build_scenario_attack_result_deltas_by_ids_statement( + self, + *, + scenario_result_id: str, + attack_result_ids: Sequence[str], + ) -> Select: + """ + Build the batched SELECT statement for scenario result deltas by ID. + + Ensures cross-run references are rejected by constraining + attribution_parent_id to the scenario UUID. + Bounds parameters to avoid exceeding SQL Server's parameter limit. + + Returns: + Select: The constructed select statement. + """ + try: + scenario_uuid = uuid.UUID(scenario_result_id) + except (ValueError, TypeError): + scenario_uuid = uuid.UUID(int=0) + valid_uuids: list[uuid.UUID] = [] + seen: set[uuid.UUID] = set() + for aid in attack_result_ids: + try: + u = uuid.UUID(str(aid)) + if u not in seen: + seen.add(u) + valid_uuids.append(u) + except (ValueError, TypeError): + continue + + # Strictly reserve parameter slots for fixed query parameters like attribution_parent_id + max_batch_size = max(1, self._MAX_BIND_VARS - 10) + chunk_uuids = valid_uuids[:max_batch_size] + + return ( + select( + AttackResultEntry.id, + AttackResultEntry.conversation_id, + AttackResultEntry.objective, + AttackResultEntry.objective_sha256, + AttackResultEntry.atomic_attack_identifier, + AttackResultEntry.outcome, + AttackResultEntry.execution_time_ms, + AttackResultEntry.executed_turns, + AttackResultEntry.outcome_reason, + AttackResultEntry.timestamp, + AttackResultEntry.retry_events_json, + AttackResultEntry.total_retries, + AttackResultEntry.error_type, + AttackResultEntry.error_message, + AttackResultEntry.attribution_data, + AttackResultEntry.attack_metadata, + AttackResultEntry.pruned_conversation_ids, + AttackResultEntry.adversarial_chat_conversation_ids, + AttackResultEntry.preparation_conversation_ids, + ScoreEntry.id.label("score_id"), + ScoreEntry.score_value, + ScoreEntry.score_type, + ScoreEntry.status.label("score_status"), + ScoreEntry.score_rationale, + ScoreEntry.scorer_class_identifier, + ) + .outerjoin( + ScoreEntry, + func.coalesce(AttackResultEntry.human_score_id, AttackResultEntry.automated_score_id) == ScoreEntry.id, + ) + .where( + and_( + AttackResultEntry.attribution_parent_id == scenario_uuid, + AttackResultEntry.id.in_(chunk_uuids), + ) + ) + ) + + def _execute_get_scenario_attack_result_deltas_by_ids( + self, + *, + scenario_result_id: str, + attack_result_ids: Sequence[str], + ) -> list[ScenarioAttackResultDelta]: + """ + Return lightweight scenario-linked result deltas for specific attack results in chunked queries. + + Deduplicates IDs and chunks queries to respect SQL parameter limits. + + Returns: + list[ScenarioAttackResultDelta]: The matched deltas. + """ + try: + uuid.UUID(scenario_result_id) + except (ValueError, TypeError): + return [] + + valid_uuids: list[uuid.UUID] = [] + seen: set[uuid.UUID] = set() + for aid in attack_result_ids: + try: + u = uuid.UUID(str(aid)) + if u not in seen: + seen.add(u) + valid_uuids.append(u) + except (ValueError, TypeError): + continue + + if not valid_uuids: + return [] + + max_batch_size = max(1, self._MAX_BIND_VARS - 10) deltas: list[ScenarioAttackResultDelta] = [] - for row in rows[:limit]: - retry_events = [ - RetryEvent.model_validate(event) - for event in (json.loads(row["retry_events_json"]) if row["retry_events_json"] else []) + with closing(self._get_session()) as session: + for start in range(0, len(valid_uuids), max_batch_size): + chunk = valid_uuids[start : start + max_batch_size] + statement = self._build_scenario_attack_result_deltas_by_ids_statement( + scenario_result_id=scenario_result_id, + attack_result_ids=[str(u) for u in chunk], + ) + rows = session.execute(statement).mappings().all() + deltas.extend(self._row_to_scenario_attack_result_delta(row) for row in rows) + return deltas + + def _execute_get_scenario_orchestration_parent_deltas( + self, + *, + scenario_result_id: str, + ) -> list[ScenarioAttackResultDelta]: + """ + Return bounded scenario-linked orchestration parent result deltas for a run. + + Finds results that have result_role == ORCHESTRATION or have child_attack_result_ids in metadata. + + Returns: + list[ScenarioAttackResultDelta]: The matched orchestration parent deltas. + """ + try: + scenario_uuid = uuid.UUID(scenario_result_id) + except (ValueError, TypeError): + return [] + + role_match = self._get_scenario_role_match_condition( + role_expression=JsonScalar(AttackResultEntry.attribution_data, literal("$.result_role")), + role=AttackResultRole.ORCHESTRATION, + ) + metadata_text = sql_cast(AttackResultEntry.attack_metadata, String) + # Match only when child_attack_result_ids is not empty (contains at least one quoted ID) + child_ids_match = metadata_text.like('%child_attack_result_ids%["%') + + attempt_id_order = self._get_scenario_attempt_id_order_expression(attempt_id=AttackResultEntry.id) + + statement = ( + select( + AttackResultEntry.id, + AttackResultEntry.conversation_id, + AttackResultEntry.objective, + AttackResultEntry.objective_sha256, + AttackResultEntry.atomic_attack_identifier, + AttackResultEntry.outcome, + AttackResultEntry.execution_time_ms, + AttackResultEntry.executed_turns, + AttackResultEntry.outcome_reason, + AttackResultEntry.timestamp, + AttackResultEntry.retry_events_json, + AttackResultEntry.total_retries, + AttackResultEntry.error_type, + AttackResultEntry.error_message, + AttackResultEntry.attribution_data, + AttackResultEntry.attack_metadata, + AttackResultEntry.pruned_conversation_ids, + AttackResultEntry.adversarial_chat_conversation_ids, + AttackResultEntry.preparation_conversation_ids, + ScoreEntry.id.label("score_id"), + ScoreEntry.score_value, + ScoreEntry.score_type, + ScoreEntry.status.label("score_status"), + ScoreEntry.score_rationale, + ScoreEntry.scorer_class_identifier, + ) + .outerjoin( + ScoreEntry, + func.coalesce(AttackResultEntry.human_score_id, AttackResultEntry.automated_score_id) == ScoreEntry.id, + ) + .where( + and_( + AttackResultEntry.attribution_parent_id == scenario_uuid, + or_(role_match, child_ids_match), + ) + ) + .order_by(AttackResultEntry.timestamp.asc(), attempt_id_order.asc()) + .limit(self._MAX_BIND_VARS) + ) + + with closing(self._get_session()) as session: + rows = session.execute(statement).mappings().all() + + return [self._row_to_scenario_attack_result_delta(row) for row in rows] + + def _execute_get_scenario_parent_deltas_for_child_ids( + self, + *, + scenario_result_id: str, + child_ids: Sequence[str], + limit: int = 500, + exclude_result_ids: Sequence[str] | None = None, + ) -> tuple[list[ScenarioAttackResultDelta], bool]: + """ + Return run-scoped parent deltas that reference any of the specified child result IDs. + + Paginates scanned database rows using keyset cursors across child chunks to guarantee + that a short returned list is a genuine proof of exhausted database rows, preventing + false-root generation when database exclusions are omitted or partial. + + Args: + scenario_result_id: Persisted scenario run ID. + child_ids: Child attack result IDs to find parents for. + limit: Maximum number of parent deltas to return. + exclude_result_ids: Optional result IDs already known to caller to exclude from query. + + Returns: + tuple[list[ScenarioAttackResultDelta], bool]: The matched parent deltas and whether + more unscanned matching rows exist in the database. + """ + try: + scenario_uuid = uuid.UUID(scenario_result_id) + except (ValueError, TypeError): + return [], False + valid_child_ids = [cid for cid in child_ids if cid] + if not valid_child_ids or limit <= 0: + return [], False + + valid_exclude_uuids: list[uuid.UUID] = [] + if exclude_result_ids: + for eid in exclude_result_ids: + try: + valid_exclude_uuids.append(uuid.UUID(str(eid))) + except (ValueError, TypeError): + continue + exclude_id_set = {str(u) for u in valid_exclude_uuids} + + seen_ids: set[str] = set() + unique_deltas: list[ScenarioAttackResultDelta] = [] + metadata_text = sql_cast(AttackResultEntry.attack_metadata, String) + attempt_id_order = self._get_scenario_attempt_id_order_expression(attempt_id=AttackResultEntry.id) + + _fixed_overhead = 10 + # To avoid exceeding SQL Server's 2,100 parameter limit (or SQLite's 500), + # budget the total parameters across fixed clauses, exclusions, and child ID matches. + max_sql_excludes = self._MAX_BIND_VARS // 2 + if len(valid_exclude_uuids) <= max_sql_excludes: + exclude_clauses = [ + AttackResultEntry.id.not_in(valid_exclude_uuids[j : j + 200]) + for j in range(0, len(valid_exclude_uuids), 200) ] - atomic_identifier = ( - AtomicAttackIdentifier.model_validate(row["atomic_attack_identifier"]) - if row["atomic_attack_identifier"] - else None + exclude_param_count = len(valid_exclude_uuids) + else: + # When exclusions are too large, don't bind them in SQL; Python filters them via exclude_id_set. + exclude_clauses = [] + exclude_param_count = 0 + + chunk_size = min( + len(valid_child_ids), + max(1, self._MAX_BIND_VARS - _fixed_overhead - exclude_param_count), + ) + + has_more = False + max_scanned_rows = 10000 + total_scanned_rows = 0 + scan_page_size = min(500, self._MAX_BIND_VARS) + num_chunks = (len(valid_child_ids) + chunk_size - 1) // chunk_size + + with closing(self._get_session()) as session: + for chunk_idx, i in enumerate(range(0, len(valid_child_ids), chunk_size)): + if len(unique_deltas) >= limit or total_scanned_rows >= max_scanned_rows: + has_more = True + break + + chunk = valid_child_ids[i : i + chunk_size] + child_matches = [metadata_text.like(f'%"{cid}"%') for cid in chunk] + base_where_conditions = [ + AttackResultEntry.attribution_parent_id == scenario_uuid, + or_(*child_matches), + *exclude_clauses, + ] + + chunk_cursor: AttackResultKeysetCursor | None = None + + while len(unique_deltas) < limit and total_scanned_rows < max_scanned_rows: + where_conditions = list(base_where_conditions) + if chunk_cursor is not None: + cursor_uuid = uuid.UUID(chunk_cursor.attack_result_id) + cursor_id_order = self._get_scenario_attempt_id_order_expression( + attempt_id=literal(cursor_uuid, type_=AttackResultEntry.id.type) + ) + where_conditions.append( + or_( + AttackResultEntry.timestamp > chunk_cursor.timestamp, + and_( + AttackResultEntry.timestamp == chunk_cursor.timestamp, + attempt_id_order > cursor_id_order, + ), + ) + ) + + statement = ( + select( + AttackResultEntry.id, + AttackResultEntry.conversation_id, + AttackResultEntry.objective, + AttackResultEntry.objective_sha256, + AttackResultEntry.atomic_attack_identifier, + AttackResultEntry.outcome, + AttackResultEntry.execution_time_ms, + AttackResultEntry.executed_turns, + AttackResultEntry.outcome_reason, + AttackResultEntry.timestamp, + AttackResultEntry.retry_events_json, + AttackResultEntry.total_retries, + AttackResultEntry.error_type, + AttackResultEntry.error_message, + AttackResultEntry.attribution_data, + AttackResultEntry.attack_metadata, + AttackResultEntry.pruned_conversation_ids, + AttackResultEntry.adversarial_chat_conversation_ids, + AttackResultEntry.preparation_conversation_ids, + ScoreEntry.id.label("score_id"), + ScoreEntry.score_value, + ScoreEntry.score_type, + ScoreEntry.status.label("score_status"), + ScoreEntry.score_rationale, + ScoreEntry.scorer_class_identifier, + ) + .outerjoin( + ScoreEntry, + func.coalesce(AttackResultEntry.human_score_id, AttackResultEntry.automated_score_id) + == ScoreEntry.id, + ) + .where(and_(*where_conditions)) + .order_by(AttackResultEntry.timestamp.asc(), attempt_id_order.asc()) + .limit(scan_page_size + 1) + ) + rows = session.execute(statement).mappings().all() + if not rows: + break + + scanned_in_query = min(len(rows), scan_page_size) + total_scanned_rows += scanned_in_query + has_more_in_chunk = len(rows) > scan_page_size + + last_scanned_row = rows[scanned_in_query - 1] + chunk_cursor = AttackResultKeysetCursor( + timestamp=last_scanned_row["timestamp"], + attack_result_id=str(last_scanned_row["id"]), + ) + + has_unprocessed_in_page = False + for row_idx, row in enumerate(rows[:scanned_in_query]): + delta = self._row_to_scenario_attack_result_delta(row) + if delta.attack_result_id not in seen_ids and delta.attack_result_id not in exclude_id_set: + seen_ids.add(delta.attack_result_id) + unique_deltas.append(delta) + if len(unique_deltas) >= limit: + if row_idx + 1 < len(rows): + has_unprocessed_in_page = True + break + + if len(unique_deltas) >= limit: + if has_more_in_chunk or has_unprocessed_in_page or chunk_idx + 1 < num_chunks: + has_more = True + break + + if not has_more_in_chunk: + break + + if total_scanned_rows >= max_scanned_rows: + has_more = True + break + + return unique_deltas[:limit], has_more + + def _fetch_bounded_conversation_messages( + self, + *, + conversation_id: str, + limit: int, + max_pieces_per_message: int = 100, + ) -> tuple[list[ScenarioConversationMessageProjection], bool]: + """ + Fetch conversation messages applying SQL LIMIT directly at the database layer. + + Preserves all multimodal message pieces and raw stored roles. + + Returns: + tuple[list[ScenarioConversationMessageProjection], bool]: Messages and whether more messages exist. + """ + if not conversation_id or limit < 1: + return [], False + + with closing(self._get_session()) as session: + seq_stmt = ( + select(PromptMemoryEntry.sequence) + .where(PromptMemoryEntry.conversation_id == str(conversation_id)) + .distinct() + .order_by(PromptMemoryEntry.sequence.asc()) + .limit(limit + 1) + ) + seqs = [row[0] for row in session.execute(seq_stmt).all()] + if not seqs: + return [], False + + truncated = len(seqs) > limit + target_seqs = seqs[:limit] + + piece_budget = max(1, max_pieces_per_message) + subq = ( + select( + PromptMemoryEntry.id.label("entry_id"), + func.row_number() + .over( + partition_by=PromptMemoryEntry.sequence, + order_by=( + PromptMemoryEntry.timestamp.asc(), + PromptMemoryEntry.id.asc(), + ), + ) + .label("rn"), + ) + .where( + and_( + PromptMemoryEntry.conversation_id == str(conversation_id), + PromptMemoryEntry.sequence.in_(target_seqs), + ) + ) + .subquery() ) - score = None - if row["score_id"] is not None: - scorer_identifier = ( - ComponentIdentifier.model_validate(row["scorer_class_identifier"]) - if row["scorer_class_identifier"] - else None + pieces_stmt = ( + select(PromptMemoryEntry, subq.c.rn) + .join(subq, PromptMemoryEntry.id == subq.c.entry_id) + .where(subq.c.rn <= piece_budget + 1) + .order_by( + PromptMemoryEntry.sequence.asc(), + PromptMemoryEntry.timestamp.asc(), + PromptMemoryEntry.id.asc(), ) - score = ScenarioProgressScore( - scorer_name=scorer_identifier.class_name if scorer_identifier else "Unknown", - score_type=row["score_type"], - status=ScoreStatus(row["score_status"]), - score_value=row["score_value"], - score_rationale=row["score_rationale"], + ) + rows = session.execute(pieces_stmt).all() + if not rows: + return [], False + + truncated_seqs = {entry.sequence for entry, rn in rows if rn > piece_budget} + bounded_pieces = [entry.get_message_piece() for entry, rn in rows if rn <= piece_budget] + grouped_messages = group_conversation_message_pieces_by_sequence(message_pieces=bounded_pieces) + + messages: list[ScenarioConversationMessageProjection] = [] + for msg in grouped_messages: + seq = msg.message_pieces[0].sequence if msg.message_pieces else 0 + pieces_truncated = seq in truncated_seqs + pieces_proj = [ + ScenarioConversationMessagePieceProjection( + data_type=piece.converted_value_data_type, + content=piece.original_value, + converted_value=piece.converted_value, + ) + for piece in msg.message_pieces + ] + first_piece = msg.message_pieces[0] if msg.message_pieces else None + raw_role = str(first_piece.role) if first_piece else "unknown" + ts = first_piece.timestamp if first_piece else None + content = msg.get_value() + + messages.append( + ScenarioConversationMessageProjection( + role=raw_role, + content=content, + sequence=seq, + timestamp=ts, + pieces=pieces_proj, + pieces_truncated=pieces_truncated, + ) ) - deltas.append( - ScenarioAttackResultDelta( - attack_result_id=str(row["id"]), - conversation_id=row["conversation_id"], - objective=row["objective"], - objective_sha256=row["objective_sha256"], - atomic_attack_identifier=atomic_identifier, - outcome=AttackOutcome(row["outcome"]), - execution_time_ms=row["execution_time_ms"], - timestamp=row["timestamp"], - retry_events=retry_events, - total_retries=row["total_retries"] or 0, - error_type=row["error_type"], - error_message=row["error_message"], - attribution_data=row["attribution_data"] or {}, - attack_metadata=row["attack_metadata"] or {}, - score=score, + + return messages, truncated + + def _execute_get_attack_result_conversation_projections( + self, + *, + attack_result_id: str, + include_messages: bool = False, + limit_messages_per_conversation: int = 100, + max_conversations: int = 50, + max_pieces_per_message: int = 100, + ) -> ScenarioAttackResultConversations: + """ + Read generic conversation projections for an attack result without full-run hydration. + + Applies bounded conversation budget and SQL LIMIT queries. + + Returns: + ScenarioAttackResultConversations: Container with projections, total count, and truncation flag. + """ + try: + ar_uuid = uuid.UUID(attack_result_id) + except (ValueError, TypeError): + return ScenarioAttackResultConversations(conversations=[], total_count=0, truncated=False) + + entries = self._query_entries( + AttackResultEntry, + conditions=AttackResultEntry.id == ar_uuid, + ) + if not entries: + return ScenarioAttackResultConversations(conversations=[], total_count=0, truncated=False) + + entry = entries[0] + conversations: list[ScenarioConversationProjection] = [] + all_conv_ids: list[str] = [] + + if entry.conversation_id: + all_conv_ids.append(entry.conversation_id) + + raw_pruned: Any = entry.pruned_conversation_ids + raw_adversarial: Any = entry.adversarial_chat_conversation_ids + raw_preparation: Any = entry.preparation_conversation_ids + + pruned_ids = [c for c in raw_pruned if isinstance(c, str)] if isinstance(raw_pruned, list) else [] + adversarial_ids = ( + [c for c in raw_adversarial if isinstance(c, str)] if isinstance(raw_adversarial, list) else [] + ) + preparation_ids = ( + [c for c in raw_preparation if isinstance(c, str)] if isinstance(raw_preparation, list) else [] + ) + + for c_id in pruned_ids + adversarial_ids + preparation_ids: + if c_id and c_id not in all_conv_ids: + all_conv_ids.append(c_id) + + total_conversations = len(all_conv_ids) + conversations_truncated = total_conversations > max_conversations + conv_ids_to_fetch = all_conv_ids[:max_conversations] + if not conv_ids_to_fetch: + return ScenarioAttackResultConversations(conversations=[], total_count=0, truncated=False) + + stats_map = self._execute_get_conversation_stats(conversation_ids=conv_ids_to_fetch) + + if entry.conversation_id: + stats = stats_map.get(entry.conversation_id) + msg_count = stats.message_count if stats else 0 + messages: list[ScenarioConversationMessageProjection] = [] + truncated = False + if include_messages and entry.conversation_id: + messages, truncated = self._fetch_bounded_conversation_messages( + conversation_id=entry.conversation_id, + limit=limit_messages_per_conversation, + max_pieces_per_message=max_pieces_per_message, + ) + conversations.append( + ScenarioConversationProjection( + conversation_id=entry.conversation_id, + conversation_type="main", + description="Main objective target conversation", + message_count=msg_count, + messages=messages, + messages_truncated=truncated, + ) + ) + + remaining_budget = max(0, max_conversations - (1 if entry.conversation_id else 0)) + related_convs = ( + [(c_id, "pruned", "Pruned conversation") for c_id in pruned_ids if c_id != entry.conversation_id] + + [ + (c_id, "adversarial_chat", "Adversarial chat conversation") + for c_id in adversarial_ids + if c_id != entry.conversation_id + ] + + [ + (c_id, "preparation", "Preparation conversation") + for c_id in preparation_ids + if c_id != entry.conversation_id + ] + ) + seen_related: set[str] = set() + deduped_related: list[tuple[str, str, str]] = [] + for item in related_convs: + if item[0] not in seen_related: + seen_related.add(item[0]) + deduped_related.append(item) + + bounded_related = deduped_related[:remaining_budget] + + for c_id, conv_type, desc in bounded_related: + if not c_id: + continue + stats = stats_map.get(c_id) + msg_count = stats.message_count if stats else 0 + messages = [] + truncated = False + if include_messages: + messages, truncated = self._fetch_bounded_conversation_messages( + conversation_id=c_id, + limit=limit_messages_per_conversation, + max_pieces_per_message=max_pieces_per_message, + ) + conversations.append( + ScenarioConversationProjection( + conversation_id=c_id, + conversation_type=conv_type, + description=desc, + message_count=msg_count, + messages=messages, + messages_truncated=truncated, ) ) - return deltas, has_more + + return ScenarioAttackResultConversations( + conversations=conversations, + total_count=total_conversations, + truncated=conversations_truncated, + ) def _execute_get_scenario_results( self, @@ -11004,6 +11973,138 @@ async def get_scenario_attack_result_deltas_async( limit=limit, ) + async def get_scenario_recount_summaries_async( + self, + *, + scenario_result_id: str, + cursor: AttackResultKeysetCursor | None = None, + limit: int = 500, + ) -> tuple[list[ScenarioAttemptSeedSummary], bool]: + """ + Return bounded scenario attempt seed summaries in ascending keyset order for recount. + + This projection selects only the minimal scalar fields needed for execution-unit + matching and recount metrics. Never loads scores, prompt entries, retry events, or + large conversation payloads. + + Returns: + tuple[list[ScenarioAttemptSeedSummary], bool]: The page and whether more rows exist. + + Raises: + ValueError: If the limit is invalid. + """ + return await self._run_database_operation_async( + self._execute_get_scenario_recount_summaries, + scenario_result_id=scenario_result_id, + cursor=cursor, + limit=limit, + ) + + async def get_scenario_attack_result_deltas_by_ids_async( + self, + *, + scenario_result_id: str, + attack_result_ids: Sequence[str], + ) -> list[ScenarioAttackResultDelta]: + """ + Return lightweight scenario-linked result deltas for specific attack results in a single query. + + Ensures child results belonging to other scenarios or non-existent IDs are isolated safely. + + Args: + scenario_result_id: The ID of the scenario run. + attack_result_ids: Sequence of attack result IDs to retrieve. + + Returns: + list[ScenarioAttackResultDelta]: The matched deltas. + """ + return await self._run_database_operation_async( + self._execute_get_scenario_attack_result_deltas_by_ids, + scenario_result_id=scenario_result_id, + attack_result_ids=attack_result_ids, + ) + + async def get_scenario_orchestration_parent_deltas_async( + self, + *, + scenario_result_id: str, + ) -> list[ScenarioAttackResultDelta]: + """ + Return lightweight scenario-linked orchestration parent result deltas for a run. + + Finds results that have result_role == ORCHESTRATION or have child_attack_result_ids in metadata. + + Args: + scenario_result_id: The ID of the scenario run. + + Returns: + list[ScenarioAttackResultDelta]: The matched orchestration parent deltas. + """ + return await self._run_database_operation_async( + self._execute_get_scenario_orchestration_parent_deltas, + scenario_result_id=scenario_result_id, + ) + + async def get_scenario_parent_deltas_for_child_ids_async( + self, + *, + scenario_result_id: str, + child_ids: Sequence[str], + limit: int = 500, + exclude_result_ids: Sequence[str] | None = None, + ) -> tuple[list[ScenarioAttackResultDelta], bool]: + """ + Return run-scoped parent deltas that reference any of the specified child result IDs. + + Args: + scenario_result_id: The ID of the scenario run. + child_ids: Child attack result IDs to find parents for. + limit: Maximum number of parent deltas to return. + exclude_result_ids: Optional result IDs already known to caller to exclude from query. + + Returns: + tuple[list[ScenarioAttackResultDelta], bool]: The matched parent deltas and whether + more unscanned matching rows exist in the database. + """ + return await self._run_database_operation_async( + self._execute_get_scenario_parent_deltas_for_child_ids, + scenario_result_id=scenario_result_id, + child_ids=child_ids, + limit=limit, + exclude_result_ids=exclude_result_ids, + ) + + async def get_attack_result_conversation_projections_async( + self, + *, + attack_result_id: str, + include_messages: bool = False, + limit_messages_per_conversation: int = 100, + max_conversations: int = 50, + max_pieces_per_message: int = 100, + ) -> ScenarioAttackResultConversations: + """ + Read generic conversation projections for an attack result without full-run hydration. + + Args: + attack_result_id: The attack result ID. + include_messages: Whether to populate message pieces or only metadata/counts. + limit_messages_per_conversation: Bound on message pieces per conversation. + max_conversations: Maximum number of conversations to retrieve for this attack result. + max_pieces_per_message: Bound on message pieces within a single message. + + Returns: + ScenarioAttackResultConversations: Container with projections, total count, and truncation flag. + """ + return await self._run_database_operation_async( + self._execute_get_attack_result_conversation_projections, + attack_result_id=attack_result_id, + include_messages=include_messages, + limit_messages_per_conversation=limit_messages_per_conversation, + max_conversations=max_conversations, + max_pieces_per_message=max_pieces_per_message, + ) + def get_scenario_results( self, *, diff --git a/pyrit/memory/sqlite_memory.py b/pyrit/memory/sqlite_memory.py index 047c577f84..15387228f8 100644 --- a/pyrit/memory/sqlite_memory.py +++ b/pyrit/memory/sqlite_memory.py @@ -729,58 +729,62 @@ def _execute_get_conversation_stats(self, *, conversation_ids: Sequence[str]) -> if not conversation_ids: return {} - placeholders = ", ".join(f":cid{i}" for i in range(len(conversation_ids))) - params = {f"cid{i}": cid for i, cid in enumerate(conversation_ids)} - - sql = text( - f""" - WITH aggregate_rows AS ( - SELECT - conversation_id, - COUNT(DISTINCT sequence) AS msg_count, - MIN(timestamp) AS created_at - FROM "PromptMemoryEntries" - WHERE conversation_id IN ({placeholders}) - GROUP BY conversation_id - ) - SELECT - aggregate_rows.conversation_id, - aggregate_rows.msg_count, - SUBSTR(latest.converted_value, 1, {ConversationStats.PREVIEW_FETCH_MAX_LEN}) AS last_preview, - latest.converted_value_data_type AS last_data_type, - aggregate_rows.created_at - FROM aggregate_rows - LEFT JOIN "PromptMemoryEntries" latest - ON latest.id = ( - SELECT p2.id - FROM "PromptMemoryEntries" p2 - WHERE p2.conversation_id = aggregate_rows.conversation_id - ORDER BY p2.sequence DESC, p2.id DESC - LIMIT 1 - ) - """ - ) + batch_size = max(1, self._MAX_BIND_VARS) + result: dict[str, ConversationStats] = {} with closing(self._get_session()) as session: - rows = session.execute(sql, params).fetchall() + for start in range(0, len(conversation_ids), batch_size): + chunk = conversation_ids[start : start + batch_size] + placeholders = ", ".join(f":cid{i}" for i in range(len(chunk))) + params = {f"cid{i}": cid for i, cid in enumerate(chunk)} + + sql = text( + f""" + WITH aggregate_rows AS ( + SELECT + conversation_id, + COUNT(DISTINCT sequence) AS msg_count, + MIN(timestamp) AS created_at + FROM "PromptMemoryEntries" + WHERE conversation_id IN ({placeholders}) + GROUP BY conversation_id + ) + SELECT + aggregate_rows.conversation_id, + aggregate_rows.msg_count, + SUBSTR(latest.converted_value, 1, {ConversationStats.PREVIEW_FETCH_MAX_LEN}) AS last_preview, + latest.converted_value_data_type AS last_data_type, + aggregate_rows.created_at + FROM aggregate_rows + LEFT JOIN "PromptMemoryEntries" latest + ON latest.id = ( + SELECT p2.id + FROM "PromptMemoryEntries" p2 + WHERE p2.conversation_id = aggregate_rows.conversation_id + ORDER BY p2.sequence DESC, p2.id DESC + LIMIT 1 + ) + """ + ) - result: dict[str, ConversationStats] = {} - for row in rows: - conv_id, msg_count, last_preview, last_data_type, raw_created_at = row - - created_at = None - if raw_created_at is not None: - if isinstance(raw_created_at, str): - created_at = datetime.fromisoformat(raw_created_at) - else: - created_at = raw_created_at - - result[conv_id] = ConversationStats( - message_count=msg_count, - last_message_preview=last_preview, - last_message_data_type=last_data_type, - created_at=created_at, - ) + rows = session.execute(sql, params).fetchall() + + for row in rows: + conv_id, msg_count, last_preview, last_data_type, raw_created_at = row + + created_at = None + if raw_created_at is not None: + if isinstance(raw_created_at, str): + created_at = datetime.fromisoformat(raw_created_at) + else: + created_at = raw_created_at + + result[conv_id] = ConversationStats( + message_count=msg_count, + last_message_preview=last_preview, + last_message_data_type=last_data_type, + created_at=created_at, + ) return result diff --git a/pyrit/models/__init__.py b/pyrit/models/__init__.py index 1fcd1483bd..b65d45c0c9 100644 --- a/pyrit/models/__init__.py +++ b/pyrit/models/__init__.py @@ -162,6 +162,17 @@ ScenarioDatasetSizeEstimateKind, scenario_dataset_size_from_limit, ) + from pyrit.models.scenario_hierarchy import ( + ScenarioAttackResultConversations, + ScenarioConversationMessagePieceProjection, + ScenarioConversationMessageProjection, + ScenarioConversationProjection, + ScenarioHierarchyDiagnostics, + ScenarioHierarchyGroup, + ScenarioHierarchyNode, + ScenarioHierarchyProjection, + ScenarioOrchestrationParentSummary, + ) from pyrit.models.scenario_progress import ( SCENARIO_RUN_PLAN_METADATA_KEY, SCENARIO_RUN_PLAN_VERSION, @@ -470,6 +481,15 @@ "ScenarioDisplayGroupProgress": "pyrit.models.scenario_progress", "ScenarioExecutionStatistics": "pyrit.models.scenario_progress", "ScenarioExecutionUnit": "pyrit.models.scenario_progress", + "ScenarioConversationMessagePieceProjection": "pyrit.models.scenario_hierarchy", + "ScenarioAttackResultConversations": "pyrit.models.scenario_hierarchy", + "ScenarioConversationMessageProjection": "pyrit.models.scenario_hierarchy", + "ScenarioConversationProjection": "pyrit.models.scenario_hierarchy", + "ScenarioHierarchyDiagnostics": "pyrit.models.scenario_hierarchy", + "ScenarioHierarchyGroup": "pyrit.models.scenario_hierarchy", + "ScenarioHierarchyNode": "pyrit.models.scenario_hierarchy", + "ScenarioHierarchyProjection": "pyrit.models.scenario_hierarchy", + "ScenarioOrchestrationParentSummary": "pyrit.models.scenario_hierarchy", "ScenarioObjectiveScorer": "pyrit.models.scenario_progress", "ScenarioObjectiveScorerMetrics": "pyrit.models.scenario_progress", "ScenarioProgressCounts": "pyrit.models.scenario_progress", diff --git a/pyrit/models/catalog/scenario.py b/pyrit/models/catalog/scenario.py index ee87a3d657..748098cfb9 100644 --- a/pyrit/models/catalog/scenario.py +++ b/pyrit/models/catalog/scenario.py @@ -528,9 +528,11 @@ class AttackRetrySummary(BaseModel): attack_result_id: str = Field(..., description="Stable ID of the attack result (used to de-duplicate)") atomic_attack_name: str = Field(..., description="Atomic-attack cell that retried") + total_retries: int = Field(0, ge=0, description="Total retry attempts recorded for this attack") retries: list[RetryEvent] = Field( default_factory=list, description="Retry attempts, each with component role/name, endpoint, and exception" ) + retries_truncated: bool = Field(False, description="Whether the retries list was truncated to a maximum budget") class ScenarioOverloadSummary(BaseModel): diff --git a/pyrit/models/scenario_hierarchy.py b/pyrit/models/scenario_hierarchy.py new file mode 100644 index 0000000000..fc51d6cdc9 --- /dev/null +++ b/pyrit/models/scenario_hierarchy.py @@ -0,0 +1,231 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from typing import ClassVar, NamedTuple, Self + +from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, model_validator + +from pyrit.models.catalog.scenario import AttackErrorSummary, AttackRetrySummary +from pyrit.models.results.attack_result import AttackOutcome, AttackResultRole +from pyrit.models.results.scenario_result import ScenarioRunState +from pyrit.models.retry_event import RetryEvent +from pyrit.models.scenario_progress import ( + ScenarioAttackTechniqueDetails, + ScenarioProgressCounts, + ScenarioProgressScore, +) + + +class ScenarioOrchestrationParentSummary(BaseModel): + """ + Typed summary of an orchestration parent (e.g. SequentialAttack, Adaptive). + + Surfaces execution metrics, completion policies, ordered child result IDs, + and retry pressure without equating orchestration results with direct target calls. + """ + + model_config = ConfigDict(extra="ignore") + + max_child_references: ClassVar[int] = 500 + + attack_result_id: str + result_role: AttackResultRole = AttackResultRole.UNKNOWN + outcome: AttackOutcome + outcome_reason: str | None = None + attempt_index: int | None = Field(default=None, ge=1) + execution_time_ms: int = 0 + executed_turns: int = 0 + completion_policy: str | None = None + child_attack_result_ids: list[str] = Field(default_factory=list) + child_references_truncated: bool = False + error_type: str | None = None + error_message: str | None = None + timestamp: AwareDatetime + total_retries: int = 0 + retries: list[RetryEvent] = Field(default_factory=list) + retries_truncated: bool = False + + @model_validator(mode="after") + def _bound_child_references(self) -> Self: + if len(self.child_attack_result_ids) > self.max_child_references: + self.child_attack_result_ids = self.child_attack_result_ids[: self.max_child_references] + self.child_references_truncated = True + return self + + +class ScenarioConversationMessagePieceProjection(BaseModel): + """ + Projection of an individual message piece preserving data type and multimodal content. + """ + + model_config = ConfigDict(extra="ignore") + + data_type: str = "text" + content: str = "" + converted_value: str = "" + + +class ScenarioConversationMessageProjection(BaseModel): + """ + Generic, read-only message projection for conversation views. + + Contains display text, conversation role, sequence order, multimodal pieces, and linked score indicators. + """ + + model_config = ConfigDict(extra="ignore") + + role: str + content: str + sequence: int + timestamp: AwareDatetime | None = None + scores: list[ScenarioProgressScore] = Field(default_factory=list) + pieces: list[ScenarioConversationMessagePieceProjection] = Field(default_factory=list) + pieces_truncated: bool = False + + +class ScenarioConversationProjection(BaseModel): + """ + Generic read-only conversation projection belonging to an attack attempt. + + Carries conversation metadata and optional message lists for detailed views. + """ + + model_config = ConfigDict(extra="ignore") + + conversation_id: str + conversation_type: str = "main" + description: str | None = None + message_count: int | None = None + messages: list[ScenarioConversationMessageProjection] = Field(default_factory=list) + messages_truncated: bool = False + + +class ScenarioHierarchyNode(BaseModel): + """ + One node in the scenario result hierarchy. + + Can represent an orchestration parent (with children) or a leaf target-facing attack. + Maintains ordered child links, parent-relative index, and unresolved child references. + """ + + model_config = ConfigDict(extra="ignore") + + max_child_references: ClassVar[int] = 500 + + attack_result_id: str + conversation_id: str = "" + parent_attack_result_id: str | None = None + child_index: int | None = None + result_role: AttackResultRole = AttackResultRole.UNKNOWN + attempt_index: int | None = Field(default=None, ge=1) + atomic_attack_name: str + atomic_group_id: str + seed_group_id: str + objective: str + outcome: AttackOutcome + outcome_reason: str | None = None + execution_time_ms: int = 0 + executed_turns: int = 0 + timestamp: AwareDatetime + score: ScenarioProgressScore | None = None + error_type: str | None = None + error_message: str | None = None + total_retries: int = 0 + retries: list[RetryEvent] = Field(default_factory=list) + retries_truncated: bool = False + + # Stored references from producer + child_attack_result_ids: list[str] = Field(default_factory=list) + child_references_truncated: bool = False + + # Resolved children in stored order (bounded traversal) + children: list["ScenarioHierarchyNode"] = Field(default_factory=list) + children_truncated: bool = False + + # Missing, malformed, cyclic, or cross-run references that could not be safely resolved + unresolved_child_attack_result_ids: list[str] = Field(default_factory=list) + unresolved_child_references_truncated: bool = False + + # Generic read-only conversations for this attack + conversations: list[ScenarioConversationProjection] = Field(default_factory=list) + conversations_count: int | None = None + conversations_truncated: bool = False + + # Orchestration details if role is ORCHESTRATION or has children + orchestration: ScenarioOrchestrationParentSummary | None = None + + @model_validator(mode="after") + def _bound_child_references(self) -> Self: + if len(self.child_attack_result_ids) > self.max_child_references: + self.child_attack_result_ids = self.child_attack_result_ids[: self.max_child_references] + self.child_references_truncated = True + if len(self.unresolved_child_attack_result_ids) > self.max_child_references: + self.unresolved_child_attack_result_ids = self.unresolved_child_attack_result_ids[ + : self.max_child_references + ] + self.unresolved_child_references_truncated = True + return self + + +class ScenarioAttackResultConversations(NamedTuple): + """Container for bounded conversation projections, total count, and list truncation flag.""" + + conversations: list[ScenarioConversationProjection] + total_count: int + truncated: bool + + +class ScenarioHierarchyGroup(BaseModel): + """One atomic attack group's hierarchy projection.""" + + model_config = ConfigDict(extra="ignore") + + atomic_group_id: str + atomic_attack_name: str + display_group: str + # Behavioral technique identity projected once per atomic group + technique: ScenarioAttackTechniqueDetails | None = None + roots: list[ScenarioHierarchyNode] = Field(default_factory=list) + # Reuses existing #3062 / #3043 progress counters unchanged + counts: ScenarioProgressCounts + + +class ScenarioHierarchyDiagnostics(BaseModel): + """ + Bounded diagnostic detail with exact aggregate totals. + + Diagnostics lists are capped to bound payload sizes, but total counts + remain exact and truncation is signaled explicitly. + """ + + model_config = ConfigDict(extra="ignore") + + failed_attacks: list[AttackErrorSummary] = Field(default_factory=list) + failed_attacks_count: int = 0 + failed_attacks_truncated: bool = False + + attack_retries: list[AttackRetrySummary] = Field(default_factory=list) + attack_retries_count: int = 0 + attack_retries_truncated: bool = False + + total_retries: int = 0 + + +class ScenarioHierarchyProjection(BaseModel): + """Complete consumed, typed, bounded Scenario result hierarchy projection.""" + + model_config = ConfigDict(extra="ignore") + + scenario_result_id: str + scenario_name: str + scenario_version: int + status: ScenarioRunState + groups: list[ScenarioHierarchyGroup] = Field(default_factory=list) + overall_counts: ScenarioProgressCounts + diagnostics: ScenarioHierarchyDiagnostics = Field(default_factory=ScenarioHierarchyDiagnostics) + truncated: bool = False + + +ScenarioHierarchyNode.model_rebuild() +ScenarioHierarchyGroup.model_rebuild() +ScenarioHierarchyProjection.model_rebuild() diff --git a/pyrit/models/scenario_progress.py b/pyrit/models/scenario_progress.py index 0d535615aa..0e00c9b9fc 100644 --- a/pyrit/models/scenario_progress.py +++ b/pyrit/models/scenario_progress.py @@ -184,6 +184,8 @@ class ScenarioProgressResult(BaseModel): result_role: AttackResultRole = AttackResultRole.UNKNOWN #: Ordered results this orchestration parent ran, as persisted by ``SequentialAttack``. child_attack_result_ids: list[str] = Field(default_factory=list) + #: Related conversation IDs for this attack attempt (e.g. pruned, adversarial, preparation). + related_conversation_ids: list[str] = Field(default_factory=list) #: 1-based position of this result among its orchestration parent's children. attempt_index: int | None = Field(default=None, ge=1) @@ -401,8 +403,11 @@ class ScenarioAttackResultDelta(BaseModel): total_retries: int = 0 error_type: str | None = None error_message: str | None = None + executed_turns: int | None = None + outcome_reason: str | None = None attribution_data: dict[str, Any] = Field(default_factory=dict) attack_metadata: dict[str, Any] = Field(default_factory=dict) + related_conversation_ids: list[str] = Field(default_factory=list) score: ScenarioProgressScore | None = None diff --git a/tests/integration/targets/test_a2a_target_integration.py b/tests/integration/targets/test_a2a_target_integration.py index 03145afa3b..e61f16f316 100644 --- a/tests/integration/targets/test_a2a_target_integration.py +++ b/tests/integration/targets/test_a2a_target_integration.py @@ -265,7 +265,7 @@ async def test_a2a_http_input_required_continuation( request = _make_msg(text="ask confirmation", conversation_id=cid) first = await target.send_prompt_async(message=request) assert first[0].get_value() == "Are you sure?" - _persist_turn(memory=sqlite_instance, target=target, request=request, response=first[0]) + await _persist_turn_async(memory=sqlite_instance, target=target, request=request, response=first[0]) second = await target.send_prompt_async(message=_make_msg(text="yes", conversation_id=cid)) assert second[0].get_value() == "Confirmed." assert executor.requests[1].task_id == executor.requests[0].task_id diff --git a/tests/unit/backend/test_scenario_hierarchy_parity.py b/tests/unit/backend/test_scenario_hierarchy_parity.py new file mode 100644 index 0000000000..ef05d9e1ae --- /dev/null +++ b/tests/unit/backend/test_scenario_hierarchy_parity.py @@ -0,0 +1,4762 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Unit and parity tests for bounded Scenario result hierarchy projections.""" + +import hashlib +import uuid +from datetime import UTC, datetime, timedelta +from pathlib import Path +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import status +from fastapi.testclient import TestClient + +from pyrit.analytics import compute_outcome_statistics +from pyrit.backend.main import app +from pyrit.backend.services.scenario_progress_read_model import ( + ScenarioProgressReadModel, + ScenarioProgressSnapshot, +) +from pyrit.backend.services.scenario_run_service import ScenarioRunService +from pyrit.memory import MemoryInterface, PromptMemoryEntry, SQLiteMemory +from pyrit.memory.storage.serializers import set_message_piece_sha256_async +from pyrit.models import ( + SCENARIO_RUN_PLAN_METADATA_KEY, + AtomicAttackIdentifier, + AttackOutcome, + AttackResult, + AttackResultRole, + AttackSeedGroup, + ComponentIdentifier, + ConversationReference, + ConversationType, + MessagePiece, + RetryEvent, + ScenarioAtomicGroupProgress, + ScenarioAttackResultConversations, + ScenarioAttackResultDelta, + ScenarioAttackTechniqueDetails, + ScenarioConversationMessageProjection, + ScenarioConversationProjection, + ScenarioHierarchyGroup, + ScenarioHierarchyNode, + ScenarioHierarchyProjection, + ScenarioProgressCounts, + ScenarioProgressResult, + ScenarioProgressSummary, + ScenarioRunPlan, + ScenarioRunPlanAtomicGroup, + ScenarioRunPlanSeedGroup, + ScenarioRunState, + SeedIdentifier, + SeedObjective, + SeedPrompt, +) +from unit.mocks import get_mock_target_identifier, make_scenario_result + + +@pytest.fixture +def client(compatibility_headers: dict[str, str]) -> TestClient: + """Create a test client for the FastAPI app.""" + return TestClient(app, headers=compatibility_headers) + + +def _make_test_delta( + *, + attack_result_id: str, + outcome: AttackOutcome = AttackOutcome.SUCCESS, + role: AttackResultRole = AttackResultRole.TARGET_FACING, + child_ids: list[str] | None = None, + parent_collection: str = "attack_group", + seed_group_id: str = "seed_1", + objective: str = "Test objective", + retries: list[RetryEvent] | None = None, + error_type: str | None = None, + error_message: str | None = None, + attempt_index: int | None = None, + conversation_id: str = "conv_1", + seconds_offset: int = 0, +) -> ScenarioAttackResultDelta: + return ScenarioAttackResultDelta( + attack_result_id=attack_result_id, + conversation_id=conversation_id, + objective=objective, + objective_sha256="hash123", + outcome=outcome, + execution_time_ms=100, + timestamp=datetime(2025, 1, 1, tzinfo=UTC) + timedelta(seconds=seconds_offset), + retry_events=retries or [], + total_retries=len(retries or []), + error_type=error_type, + error_message=error_message, + attack_metadata={ + "child_attack_result_ids": child_ids or [], + "completion_policy": "ALL_CHILDREN", + "executed_turns": 1, + }, + attribution_data={ + "parent_collection": parent_collection, + "seed_group_id": seed_group_id, + "result_role": role.value, + "attempt_index": attempt_index, + }, + ) + + +def _make_test_result( + *, + attack_result_id: str, + atomic_group_id: str = "group-1", + atomic_attack_name: str = "attack", + seed_group_id: str = "seed-1", + outcome: AttackOutcome = AttackOutcome.SUCCESS, + conversation_id: str = "conv_1", + total_retries: int = 0, + retries: list[RetryEvent] | None = None, + error_type: str | None = None, + error_message: str | None = None, + child_attack_result_ids: list[str] | None = None, + related_conversation_ids: list[str] | None = None, + role: AttackResultRole = AttackResultRole.TARGET_FACING, +) -> ScenarioProgressResult: + return ScenarioProgressResult( + attack_result_id=attack_result_id, + conversation_id=conversation_id, + atomic_group_id=atomic_group_id, + atomic_attack_name=atomic_attack_name, + seed_group_id=seed_group_id, + outcome=outcome, + execution_time_ms=100, + timestamp=datetime(2025, 1, 1, tzinfo=UTC), + total_retries=total_retries, + retries=retries or [], + error_type=error_type, + error_message=error_message, + result_role=role, + child_attack_result_ids=child_attack_result_ids or [], + related_conversation_ids=related_conversation_ids or [], + ) + + +def _make_test_snapshot( + *, + deltas: list[ScenarioAttackResultDelta], + plan: ScenarioRunPlan | None = None, + atomic_groups: list[ScenarioAtomicGroupProgress] | None = None, + summary: ScenarioProgressSummary | None = None, +) -> ScenarioProgressSnapshot: + results: list[ScenarioProgressResult] = [] + for d in deltas: + role_str = d.attribution_data.get("result_role", AttackResultRole.UNKNOWN.value) + role = ( + AttackResultRole(role_str) if role_str in AttackResultRole._value2member_map_ else AttackResultRole.UNKNOWN + ) + results.append( + ScenarioProgressResult( + attack_result_id=d.attack_result_id, + conversation_id=d.conversation_id, + atomic_group_id=str(d.attribution_data.get("parent_collection") or ""), + atomic_attack_name=str(d.attribution_data.get("parent_collection") or ""), + seed_group_id=str(d.attribution_data.get("seed_group_id") or ""), + outcome=d.outcome, + execution_time_ms=d.execution_time_ms, + timestamp=d.timestamp, + total_retries=d.total_retries, + retries=d.retry_events, + error_type=d.error_type, + error_message=d.error_message, + score=d.score, + result_role=role, + child_attack_result_ids=list(d.attack_metadata.get("child_attack_result_ids") or []), + attempt_index=d.attribution_data.get("attempt_index"), + ) + ) + effective_plan = plan or ScenarioRunPlan( + atomic_groups=[ + ScenarioRunPlanAtomicGroup( + id="attack_group", + atomic_attack_name="attack_group", + display_group="attack_group", + technique_eval_hash="", + seed_group_ids=["seed_1"], + ) + ], + seed_groups=[ScenarioRunPlanSeedGroup(id="seed_1", objective_sha256="hash123", objective="Test objective")], + ) + if summary is None: + attempts_by_unit: dict[tuple[str, str], list[ScenarioProgressResult]] = {} + for r in results: + unit_key = (r.atomic_group_id, r.seed_group_id) + attempts_by_unit.setdefault(unit_key, []).append(r) + + calculated_retries = 0 + for attempts in attempts_by_unit.values(): + calculated_retries += max(len(attempts) - 1, 0) + sum(a.total_retries for a in attempts) + + effective_summary = ScenarioProgressSummary( + overall=ScenarioProgressCounts( + completed=len([r for r in results if r.outcome != AttackOutcome.ERROR]), + planned=1, + succeeded=len([r for r in results if r.outcome == AttackOutcome.SUCCESS]), + errors=len([r for r in results if r.outcome == AttackOutcome.ERROR]), + retries=calculated_retries, + ), + atomic_groups=atomic_groups or [], + ) + else: + effective_summary = summary + return ScenarioProgressSnapshot( + deltas=tuple(deltas), + results=tuple(results), + summary=effective_summary, + plan=effective_plan, + ) + + +def test_hierarchy_ordinary_flat_results() -> None: + delta1 = _make_test_delta(attack_result_id="attack-1", seconds_offset=1) + delta2 = _make_test_delta(attack_result_id="attack-2", seconds_offset=2) + snapshot = _make_test_snapshot(deltas=[delta1, delta2]) + + projection = ScenarioProgressReadModel.build_scenario_hierarchy( + scenario_result_id="scenario-run-1", + scenario_name="TestScenario", + scenario_version=1, + status=ScenarioRunState.COMPLETED, + snapshot=snapshot, + ) + + assert projection.scenario_result_id == "scenario-run-1" + assert len(projection.groups) == 1 + group = projection.groups[0] + assert len(group.roots) == 2 + assert [r.attack_result_id for r in group.roots] == ["attack-1", "attack-2"] + for root in group.roots: + assert root.children == [] + assert root.unresolved_child_attack_result_ids == [] + assert root.parent_attack_result_id is None + assert root.child_index is None + assert root.orchestration is None + + +def test_hierarchy_nested_sequential_attack_with_relative_indexes() -> None: + child1 = _make_test_delta( + attack_result_id="child-1", + role=AttackResultRole.TARGET_FACING, + attempt_index=1, + seconds_offset=1, + ) + child2 = _make_test_delta( + attack_result_id="child-2", + role=AttackResultRole.TARGET_FACING, + attempt_index=2, + seconds_offset=2, + ) + parent = _make_test_delta( + attack_result_id="parent-1", + role=AttackResultRole.ORCHESTRATION, + child_ids=["child-1", "child-2"], + seconds_offset=3, + ) + + snapshot = _make_test_snapshot(deltas=[child1, child2, parent]) + projection = ScenarioProgressReadModel.build_scenario_hierarchy( + scenario_result_id="scenario-run-1", + scenario_name="TestScenario", + scenario_version=1, + status=ScenarioRunState.COMPLETED, + snapshot=snapshot, + ) + + assert len(projection.groups) == 1 + group = projection.groups[0] + # Only parent is root; children are not roots + assert len(group.roots) == 1 + root = group.roots[0] + assert root.attack_result_id == "parent-1" + assert root.result_role == AttackResultRole.ORCHESTRATION + assert root.orchestration is not None + assert root.orchestration.child_attack_result_ids == ["child-1", "child-2"] + assert len(root.children) == 2 + assert root.children[0].attack_result_id == "child-1" + assert root.children[0].child_index == 0 + assert root.children[0].parent_attack_result_id == "parent-1" + assert root.children[1].attack_result_id == "child-2" + assert root.children[1].child_index == 1 + assert root.children[1].parent_attack_result_id == "parent-1" + assert root.unresolved_child_attack_result_ids == [] + + +def test_hierarchy_cyclic_references_guarded() -> None: + # A -> B and B -> A cycle + delta_a = _make_test_delta( + attack_result_id="node-a", + role=AttackResultRole.ORCHESTRATION, + child_ids=["node-b"], + seconds_offset=1, + ) + delta_b = _make_test_delta( + attack_result_id="node-b", + role=AttackResultRole.ORCHESTRATION, + child_ids=["node-a"], + seconds_offset=2, + ) + + snapshot = _make_test_snapshot(deltas=[delta_a, delta_b]) + projection = ScenarioProgressReadModel.build_scenario_hierarchy( + scenario_result_id="scenario-run-1", + scenario_name="TestScenario", + scenario_version=1, + status=ScenarioRunState.COMPLETED, + snapshot=snapshot, + ) + + # Traversal must terminate safely without recursion error + assert projection is not None + assert len(projection.groups) == 1 + # Check that cycle link is placed into unresolved_child_attack_result_ids + all_unresolved = [ + uid + for group in projection.groups + for root in group.roots + for uid in ( + root.unresolved_child_attack_result_ids + + [u for c in root.children for u in c.unresolved_child_attack_result_ids] + ) + ] + assert any("node-a" in u or "node-b" in u for u in all_unresolved) + + +def test_hierarchy_retained_failure_link_issue_3048() -> None: + # Issue #3048: An orchestration parent encounters an error, but retains completed + failed children + child1 = _make_test_delta( + attack_result_id="child-1", + outcome=AttackOutcome.SUCCESS, + role=AttackResultRole.TARGET_FACING, + attempt_index=1, + seconds_offset=1, + ) + child2 = _make_test_delta( + attack_result_id="child-2", + outcome=AttackOutcome.ERROR, + role=AttackResultRole.TARGET_FACING, + error_type="TargetError", + error_message="Endpoint failure", + attempt_index=2, + seconds_offset=2, + ) + parent = _make_test_delta( + attack_result_id="parent-error", + outcome=AttackOutcome.ERROR, + role=AttackResultRole.ORCHESTRATION, + child_ids=["child-1", "child-2"], + error_type="ScenarioPartialFailureException", + error_message="Aborted after child 2 failed", + seconds_offset=3, + ) + + snapshot = _make_test_snapshot(deltas=[child1, child2, parent]) + projection = ScenarioProgressReadModel.build_scenario_hierarchy( + scenario_result_id="scenario-run-1", + scenario_name="TestScenario", + scenario_version=1, + status=ScenarioRunState.FAILED, + snapshot=snapshot, + ) + + group = projection.groups[0] + assert len(group.roots) == 1 + root = group.roots[0] + assert root.attack_result_id == "parent-error" + assert root.outcome == AttackOutcome.ERROR + assert root.error_type == "ScenarioPartialFailureException" + assert len(root.children) == 2 + assert root.children[0].attack_result_id == "child-1" + assert root.children[0].outcome == AttackOutcome.SUCCESS + assert root.children[1].attack_result_id == "child-2" + assert root.children[1].outcome == AttackOutcome.ERROR + assert root.unresolved_child_attack_result_ids == [] + + +def test_hierarchy_missing_and_cross_run_child_references() -> None: + # Parent references a child ID that does not exist in deltas (or was from another run) + parent = _make_test_delta( + attack_result_id="parent-1", + role=AttackResultRole.ORCHESTRATION, + child_ids=["missing-child-99", "cross-run-child-42"], + ) + snapshot = _make_test_snapshot(deltas=[parent]) + + projection = ScenarioProgressReadModel.build_scenario_hierarchy( + scenario_result_id="scenario-run-1", + scenario_name="TestScenario", + scenario_version=1, + status=ScenarioRunState.COMPLETED, + snapshot=snapshot, + ) + + root = projection.groups[0].roots[0] + assert root.children == [] + assert root.unresolved_child_attack_result_ids == ["missing-child-99", "cross-run-child-42"] + + +def test_hierarchy_diagnostics_bounding_at_100() -> None: + # Create 105 errors and 120 retry events + deltas: list[ScenarioAttackResultDelta] = [] + for i in range(105): + retry = RetryEvent(exception_type="RateLimit", attempt_number=1) if i < 120 else None + retries = [retry] if retry else [] + deltas.append( + _make_test_delta( + attack_result_id=f"err-{i}", + outcome=AttackOutcome.ERROR, + error_type="TestError", + error_message=f"Error {i}", + retries=retries, + seed_group_id=f"seed_{i}", + seconds_offset=i, + ) + ) + deltas.extend( + _make_test_delta( + attack_result_id=f"retry-{i}", + outcome=AttackOutcome.SUCCESS, + retries=[RetryEvent(exception_type="RateLimit", attempt_number=1)], + seed_group_id=f"seed_{i}", + seconds_offset=i, + ) + for i in range(105, 120) + ) + + snapshot = _make_test_snapshot(deltas=deltas) + projection = ScenarioProgressReadModel.build_scenario_hierarchy( + scenario_result_id="scenario-run-1", + scenario_name="TestScenario", + scenario_version=1, + status=ScenarioRunState.COMPLETED, + snapshot=snapshot, + ) + + diag = projection.diagnostics + assert len(diag.failed_attacks) == 100 + assert diag.failed_attacks_count == 105 + assert diag.failed_attacks_truncated is True + + assert len(diag.attack_retries) == 100 + assert diag.attack_retries_count == 120 + assert diag.attack_retries_truncated is True + assert diag.total_retries == 120 + + +@pytest.mark.asyncio +async def test_service_get_run_hierarchy_node_async_bounded_read() -> None: + memory = MagicMock(spec=MemoryInterface) + parent_delta = _make_test_delta( + attack_result_id="parent-1", + role=AttackResultRole.ORCHESTRATION, + child_ids=["child-1"], + ) + child_delta = _make_test_delta( + attack_result_id="child-1", + role=AttackResultRole.TARGET_FACING, + ) + + async def mock_get_deltas( + *, scenario_result_id: str, attack_result_ids: list[str] + ) -> list[ScenarioAttackResultDelta]: + res = [] + for aid in attack_result_ids: + if aid == "parent-1": + res.append(parent_delta) + elif aid == "child-1": + res.append(child_delta) + return res + + memory.get_scenario_attack_result_deltas_by_ids_async = AsyncMock(side_effect=mock_get_deltas) + memory.get_scenario_result_header_async = AsyncMock(return_value=None) + memory.get_attack_result_conversation_projections_async = AsyncMock( + return_value=ScenarioAttackResultConversations( + conversations=[ + ScenarioConversationProjection( + conversation_id="conv_1", + conversation_type="main", + message_count=1, + messages=[ + ScenarioConversationMessageProjection( + role="user", + content="Test prompt", + sequence=1, + ) + ], + ) + ], + total_count=1, + truncated=False, + ) + ) + + read_model = ScenarioProgressReadModel(memory=memory) + service = ScenarioRunService(memory=memory, progress_read_model=read_model) + + node = await service.get_run_hierarchy_node_async( + scenario_result_id="scenario-run-1", + attack_result_id="parent-1", + ) + + assert node is not None + assert node.attack_result_id == "parent-1" + assert len(node.children) == 1 + assert node.children[0].attack_result_id == "child-1" + assert len(node.conversations) == 1 + assert len(node.conversations[0].messages) == 1 + assert node.conversations[0].messages[0].content == "Test prompt" + + # Querying non-existent node returns None + missing = await service.get_run_hierarchy_node_async( + scenario_result_id="scenario-run-1", + attack_result_id="non-existent", + ) + assert missing is None + + +def test_rest_routes_hierarchy_and_node(client: TestClient) -> None: + now = datetime.now(UTC) + node = ScenarioHierarchyNode( + attack_result_id="node-1", + result_role=AttackResultRole.TARGET_FACING, + atomic_attack_name="attack", + atomic_group_id="group-1", + seed_group_id="seed-1", + objective="obj", + outcome=AttackOutcome.SUCCESS, + timestamp=now, + ) + group = ScenarioHierarchyGroup( + atomic_group_id="group-1", + atomic_attack_name="attack", + display_group="display", + technique=ScenarioAttackTechniqueDetails(component_name="Technique"), + roots=[node], + counts=ScenarioProgressCounts(completed=1, planned=1, succeeded=1, success_percentage=100, errors=0, retries=0), + ) + projection = ScenarioHierarchyProjection( + scenario_result_id="run-123", + scenario_name="TestScenario", + scenario_version=1, + status=ScenarioRunState.COMPLETED, + groups=[group], + overall_counts=ScenarioProgressCounts( + completed=1, planned=1, succeeded=1, success_percentage=100, errors=0, retries=0 + ), + ) + + with patch("pyrit.backend.routes.scenarios.get_scenario_run_service") as mock_get_svc: + mock_svc = MagicMock() + mock_svc.get_run_hierarchy_async = AsyncMock(return_value=projection) + mock_svc.get_run_hierarchy_node_async = AsyncMock(return_value=node) + mock_get_svc.return_value = mock_svc + + # 1. Hierarchy projection endpoint + resp = client.get("/api/scenarios/runs/run-123/hierarchy") + assert resp.status_code == status.HTTP_200_OK + data = resp.json() + assert data["scenario_result_id"] == "run-123" + assert len(data["groups"]) == 1 + assert data["groups"][0]["roots"][0]["attack_result_id"] == "node-1" + + # 2. Node detail endpoint + resp_node = client.get("/api/scenarios/runs/run-123/hierarchy/nodes/node-1") + assert resp_node.status_code == status.HTTP_200_OK + node_data = resp_node.json() + assert node_data["attack_result_id"] == "node-1" + + # 3. 404 on unknown run + mock_svc.get_run_hierarchy_async.return_value = None + resp_404 = client.get("/api/scenarios/runs/unknown-run/hierarchy") + assert resp_404.status_code == status.HTTP_404_NOT_FOUND + + # 4. 404 on unknown node + mock_svc.get_run_hierarchy_node_async.return_value = None + resp_node_404 = client.get("/api/scenarios/runs/run-123/hierarchy/nodes/unknown-node") + assert resp_node_404.status_code == status.HTTP_404_NOT_FOUND + + +@pytest.mark.asyncio +async def test_sqlite_memory_batched_reads_and_projections_parity(sqlite_instance: SQLiteMemory) -> None: + """Test batched memory delta reads, conversation projections, and cross-run isolation on SQLite.""" + scenario_uuid_1 = uuid.uuid4() + scenario_uuid_2 = uuid.uuid4() + scenario_id_1 = str(scenario_uuid_1) + scenario_id_2 = str(scenario_uuid_2) + + scen_1 = make_scenario_result( + id=scenario_uuid_1, + scenario_name="Scenario1", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + scen_2 = make_scenario_result( + id=scenario_uuid_2, + scenario_name="Scenario2", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen_1, scen_2]) + + parent_id = str(uuid.uuid4()) + child_id = str(uuid.uuid4()) + cross_run_id = str(uuid.uuid4()) + child_conv_id = str(uuid.uuid4()) + + parent_result = AttackResult( + attack_result_id=parent_id, + conversation_id="parent_conv", + objective="Orchestration objective", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id_1, + attribution_data={"parent_collection": "group_1", "result_role": AttackResultRole.ORCHESTRATION.value}, + metadata={"child_attack_result_ids": [child_id, cross_run_id]}, + ) + child_result = AttackResult( + attack_result_id=child_id, + conversation_id=child_conv_id, + objective="Child target-facing objective", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id_1, + attribution_data={ + "parent_collection": "group_1", + "result_role": AttackResultRole.TARGET_FACING.value, + "attempt_index": 1, + }, + ) + cross_run_result = AttackResult( + attack_result_id=cross_run_id, + conversation_id="cross_conv", + objective="Cross-run child objective", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id_2, + attribution_data={"parent_collection": "group_2", "result_role": AttackResultRole.TARGET_FACING.value}, + ) + + await sqlite_instance.add_attack_results_to_memory_async( + attack_results=[parent_result, child_result, cross_run_result] + ) + + # Insert messages into child conversation + piece_user = MessagePiece( + id=uuid.uuid4(), + conversation_id=child_conv_id, + role="user", + original_value_data_type="text", + original_value="Prompt to target", + converted_value="Prompt to target", + sequence=0, + ) + piece_assistant = MessagePiece( + id=uuid.uuid4(), + conversation_id=child_conv_id, + role="assistant", + original_value_data_type="text", + original_value="Response from target", + converted_value="Response from target", + sequence=1, + ) + await set_message_piece_sha256_async(piece_user) + await set_message_piece_sha256_async(piece_assistant) + await sqlite_instance._run_database_operation_async( + sqlite_instance._insert_entry, PromptMemoryEntry(entry=piece_user) + ) + await sqlite_instance._run_database_operation_async( + sqlite_instance._insert_entry, PromptMemoryEntry(entry=piece_assistant) + ) + + # 1. Test batched delta reads + deltas = await sqlite_instance.get_scenario_attack_result_deltas_by_ids_async( + scenario_result_id=scenario_id_1, + attack_result_ids=[parent_id, child_id, cross_run_id], + ) + delta_ids = [d.attack_result_id for d in deltas] + assert parent_id in delta_ids + assert child_id in delta_ids + # Cross-run ID belongs to scenario_id_2, so it MUST NOT be returned! + assert cross_run_id not in delta_ids + + # 2. Test conversation projections + conv_projections = await sqlite_instance.get_attack_result_conversation_projections_async( + attack_result_id=child_id, + include_messages=True, + ) + assert conv_projections.total_count == 1 + assert conv_projections.truncated is False + assert len(conv_projections.conversations) == 1 + proj = conv_projections.conversations[0] + assert proj.conversation_id == child_conv_id + assert proj.message_count == 2 + assert len(proj.messages) == 2 + assert proj.messages[0].role == "user" + assert proj.messages[0].content == "Prompt to target" + assert proj.messages[1].role == "assistant" + assert proj.messages[1].content == "Response from target" + + # 3. Test service single-node read + service = ScenarioRunService(memory=sqlite_instance) + node = await service.get_run_hierarchy_node_async( + scenario_result_id=scenario_id_1, + attack_result_id=parent_id, + ) + assert node is not None + assert node.attack_result_id == parent_id + assert len(node.children) == 1 + assert node.children[0].attack_result_id == child_id + # Cross-run child is safely recognized as unresolved without failing the run + assert node.unresolved_child_attack_result_ids == [cross_run_id] + + +@pytest.mark.asyncio +async def test_sqlite_canonical_columns_round_trip(sqlite_instance: SQLiteMemory) -> None: + """Test real SQLite round-trip of canonical executed_turns and outcome_reason columns.""" + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="CanonicalScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + ar_uuid = uuid.uuid4() + ar_id = str(ar_uuid) + ar = AttackResult( + attack_result_id=ar_id, + conversation_id="conv_canonical", + objective="Probe objective", + outcome=AttackOutcome.FAILURE, + executed_turns=7, + outcome_reason="Target actively refused harmful request", + attribution_parent_id=scenario_id, + attribution_data={"parent_collection": "group_1", "result_role": AttackResultRole.TARGET_FACING.value}, + metadata={}, # Explicitly empty metadata to ensure canonical DB columns are read + ) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[ar]) + + # 1. Delta query directly from SQLite selects canonical columns + deltas = await sqlite_instance.get_scenario_attack_result_deltas_by_ids_async( + scenario_result_id=scenario_id, + attack_result_ids=[ar_id], + ) + assert len(deltas) == 1 + delta = deltas[0] + assert delta.executed_turns == 7 + assert delta.outcome_reason == "Target actively refused harmful request" + assert "executed_turns" not in delta.attack_metadata + assert "outcome_reason" not in delta.attack_metadata + + # 2. Hierarchy node query through service projects canonical values + service = ScenarioRunService(memory=sqlite_instance) + node = await service.get_run_hierarchy_node_async( + scenario_result_id=scenario_id, + attack_result_id=ar_id, + ) + assert node is not None + assert node.executed_turns == 7 + assert node.outcome_reason == "Target actively refused harmful request" + + +def test_hierarchy_disconnected_cyclic_component_mixed() -> None: + """Test that disconnected cyclic components (e.g. A <-> B) are not dropped when standard roots exist.""" + # 1. Ordinary tree: root-1 -> child-1 + root1 = _make_test_delta( + attack_result_id="root-1", + role=AttackResultRole.ORCHESTRATION, + child_ids=["child-1"], + seconds_offset=1, + ) + child1 = _make_test_delta( + attack_result_id="child-1", + role=AttackResultRole.TARGET_FACING, + seconds_offset=2, + ) + + # 2. Detached cyclic component: cycle-a -> cycle-b and cycle-b -> cycle-a + # Neither cycle-a nor cycle-b is reachable from root-1, and both are in each other's child_ids + cycle_a = _make_test_delta( + attack_result_id="cycle-a", + role=AttackResultRole.ORCHESTRATION, + child_ids=["cycle-b"], + seconds_offset=3, + ) + cycle_b = _make_test_delta( + attack_result_id="cycle-b", + role=AttackResultRole.ORCHESTRATION, + child_ids=["cycle-a"], + seconds_offset=4, + ) + + snapshot = _make_test_snapshot(deltas=[root1, child1, cycle_a, cycle_b]) + projection = ScenarioProgressReadModel.build_scenario_hierarchy( + scenario_result_id="scenario-run-1", + scenario_name="TestScenario", + scenario_version=1, + status=ScenarioRunState.COMPLETED, + snapshot=snapshot, + ) + + # All 4 nodes must be accounted for in the hierarchy! + all_node_ids = set() + + def _collect_ids(nodes: list[ScenarioHierarchyNode]) -> None: + for n in nodes: + all_node_ids.add(n.attack_result_id) + _collect_ids(n.children) + + for g in projection.groups: + _collect_ids(g.roots) + + assert "root-1" in all_node_ids + assert "child-1" in all_node_ids + assert "cycle-a" in all_node_ids + assert "cycle-b" in all_node_ids + assert len(all_node_ids) == 4 + + # Verify root-1 has child-1 + group = projection.groups[0] + root_node = next(r for r in group.roots if r.attack_result_id == "root-1") + assert len(root_node.children) == 1 + assert root_node.children[0].attack_result_id == "child-1" + + # Verify cycle was guarded and placed in unresolved_child_attack_result_ids without infinite loop + cycle_root = next(r for r in group.roots if r.attack_result_id in ("cycle-a", "cycle-b")) + assert len(cycle_root.children) == 1 + cycle_child = cycle_root.children[0] + assert cycle_root.attack_result_id in cycle_child.unresolved_child_attack_result_ids + + +def test_hierarchy_bounded_node_budget() -> None: + """Test that max_nodes bounds tree traversal and sets truncated=True while preserving exact overall counts.""" + # Create 6 deltas in a chain: node-0 -> node-1 -> node-2 -> node-3 -> node-4 -> node-5 + deltas = [ + _make_test_delta( + attack_result_id=f"node-{i}", + role=AttackResultRole.ORCHESTRATION if i < 5 else AttackResultRole.TARGET_FACING, + child_ids=[f"node-{i + 1}"] if i < 5 else [], + seconds_offset=i, + ) + for i in range(6) + ] + + snapshot = _make_test_snapshot(deltas=deltas) + projection = ScenarioProgressReadModel.build_scenario_hierarchy( + scenario_result_id="scenario-run-1", + scenario_name="TestScenario", + scenario_version=1, + status=ScenarioRunState.COMPLETED, + snapshot=snapshot, + max_nodes=3, + ) + + assert projection.truncated is True + + # Count nodes projected in hierarchy + node_count = 0 + + def _count_nodes(nodes: list[ScenarioHierarchyNode]) -> None: + nonlocal node_count + for n in nodes: + node_count += 1 + _count_nodes(n.children) + + for g in projection.groups: + _count_nodes(g.roots) + + assert node_count == 3 + # Top-level aggregate totals must remain exact (not clipped) + assert projection.overall_counts.completed == 6 + + +def test_hierarchy_group_progress_outcomes_parity() -> None: + """Test that group.counts.outcomes is preserved from atomic group progress.""" + outcomes = compute_outcome_statistics({AttackOutcome.SUCCESS: 1, AttackOutcome.UNDETERMINED: 1}) + ag_progress = ScenarioAtomicGroupProgress( + id="attack_group", + atomic_attack_name="attack_group", + display_group="attack_group", + status="COMPLETED", + completed=2, + planned=2, + succeeded=1, + success_percentage=50, + errors=0, + retries=0, + outcomes=outcomes, + ) + + delta1 = _make_test_delta(attack_result_id="a-1", outcome=AttackOutcome.SUCCESS, seconds_offset=1) + delta2 = _make_test_delta(attack_result_id="a-2", outcome=AttackOutcome.UNDETERMINED, seconds_offset=2) + + plan = ScenarioRunPlan( + atomic_groups=[ + ScenarioRunPlanAtomicGroup( + id="attack_group", + atomic_attack_name="attack_group", + display_group="attack_group", + technique_eval_hash="", + seed_group_ids=["seed_1"], + ) + ], + seed_groups=[ScenarioRunPlanSeedGroup(id="seed_1", objective_sha256="hash123", objective="Test objective")], + ) + + summary = ScenarioProgressSummary( + overall=ScenarioProgressCounts( + completed=2, + planned=2, + succeeded=1, + success_percentage=50, + errors=0, + retries=0, + outcomes=outcomes, + ), + atomic_groups=[ag_progress], + ) + + snapshot = ScenarioProgressSnapshot( + deltas=(delta1, delta2), + results=tuple( + ScenarioProgressResult( + attack_result_id=d.attack_result_id, + conversation_id=d.conversation_id, + atomic_group_id="attack_group", + atomic_attack_name="attack_group", + seed_group_id="seed_1", + outcome=d.outcome, + execution_time_ms=d.execution_time_ms, + timestamp=d.timestamp, + total_retries=d.total_retries, + result_role=AttackResultRole.TARGET_FACING, + ) + for d in [delta1, delta2] + ), + summary=summary, + plan=plan, + ) + + projection = ScenarioProgressReadModel.build_scenario_hierarchy( + scenario_result_id="scenario-run-1", + scenario_name="TestScenario", + scenario_version=1, + status=ScenarioRunState.COMPLETED, + snapshot=snapshot, + ) + + assert len(projection.groups) == 1 + group = projection.groups[0] + assert group.counts.outcomes is not None + assert group.counts.outcomes.success_rate_decided == 1.0 + assert group.counts.outcomes.success_rate_all == 0.5 + + +def test_hierarchy_unplanned_groups_shared_counting() -> None: + """Test that unplanned groups use shared counting policy (latest attempt resolution + valid producer counts).""" + # 2 attempts on the same unit in an unplanned group: + # Attempt 1 failed, Attempt 2 succeeded + delta1 = _make_test_delta( + attack_result_id="unplanned-1", + outcome=AttackOutcome.ERROR, + role=AttackResultRole.TARGET_FACING, + parent_collection="unplanned_group", + seed_group_id="seed_u", + seconds_offset=1, + ) + delta2 = _make_test_delta( + attack_result_id="unplanned-2", + outcome=AttackOutcome.SUCCESS, + role=AttackResultRole.TARGET_FACING, + parent_collection="unplanned_group", + seed_group_id="seed_u", + seconds_offset=2, + ) + + # Empty plan (no planned atomic groups) + plan = ScenarioRunPlan(atomic_groups=[], seed_groups=[]) + snapshot = _make_test_snapshot(deltas=[delta1, delta2], plan=plan) + + projection = ScenarioProgressReadModel.build_scenario_hierarchy( + scenario_result_id="scenario-run-1", + scenario_name="TestScenario", + scenario_version=1, + status=ScenarioRunState.COMPLETED, + snapshot=snapshot, + ) + + assert len(projection.groups) == 1 + group = projection.groups[0] + assert group.atomic_group_id == "unplanned_group" + # Latest attempt succeeded, so 1 completed unit, 1 success, 1 historical error + assert group.counts.completed == 1 + assert group.counts.succeeded == 1 + assert group.counts.errors == 1 + # Producer counts include both historical attempts + assert group.counts.producer_counts.target_facing.attempts == 2 + assert group.counts.producer_counts.target_facing.errors == 1 + assert group.counts.producer_counts.target_facing.retries == 0 + + +def test_hierarchy_diagnostics_cap_retries_per_item() -> None: + """Test that retries list inside individual AttackRetrySummary is bounded to 100 with retries_truncated=True.""" + retries = [RetryEvent(exception_type="RateLimitError", attempt_number=i) for i in range(250)] + delta = _make_test_delta( + attack_result_id="heavy-retry", + outcome=AttackOutcome.SUCCESS, + retries=retries, + ) + + snapshot = _make_test_snapshot(deltas=[delta]) + projection = ScenarioProgressReadModel.build_scenario_hierarchy( + scenario_result_id="scenario-run-1", + scenario_name="TestScenario", + scenario_version=1, + status=ScenarioRunState.COMPLETED, + snapshot=snapshot, + ) + + diag = projection.diagnostics + assert len(diag.attack_retries) == 1 + item = diag.attack_retries[0] + assert item.attack_result_id == "heavy-retry" + assert item.total_retries == 250 + assert len(item.retries) == 100 + assert item.retries_truncated is True + + +@pytest.mark.asyncio +async def test_sqlite_bounded_multimodal_conversation_projection(sqlite_instance: SQLiteMemory, tmp_path: Path) -> None: + """Test multimodal piece projection, raw role preservation, and database-level SQL LIMIT bounding.""" + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="MultimodalScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + ar_uuid = uuid.uuid4() + ar_id = str(ar_uuid) + conv_id = str(uuid.uuid4()) + ar = AttackResult( + attack_result_id=ar_id, + conversation_id=conv_id, + objective="Multimodal test", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={"parent_collection": "group_1", "result_role": AttackResultRole.TARGET_FACING.value}, + ) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[ar]) + + # Create real image on disk + img_file = tmp_path / "diagram.png" + img_file.write_bytes(b"dummy_png_bytes") + img_path = str(img_file) + + # Sequence 0: 2 pieces (text + image) + p0_text = MessagePiece( + id=uuid.uuid4(), + conversation_id=conv_id, + role="user", + original_value_data_type="text", + original_value="Analyze blueprint image", + converted_value="Analyze blueprint image", + sequence=0, + ) + p0_img = MessagePiece( + id=uuid.uuid4(), + conversation_id=conv_id, + role="user", + original_value_data_type="image_path", + original_value=img_path, + converted_value=img_path, + sequence=0, + ) + # Sequence 1: simulated_assistant role + p1 = MessagePiece( + id=uuid.uuid4(), + conversation_id=conv_id, + role="simulated_assistant", + original_value_data_type="text", + original_value="I detect architectural blueprints.", + converted_value="I detect architectural blueprints.", + sequence=1, + ) + # Sequence 2: 3rd message + p2 = MessagePiece( + id=uuid.uuid4(), + conversation_id=conv_id, + role="user", + original_value_data_type="text", + original_value="Proceed with analysis.", + converted_value="Proceed with analysis.", + sequence=2, + ) + + for p in [p0_text, p0_img, p1, p2]: + await set_message_piece_sha256_async(p) + await sqlite_instance._run_database_operation_async(sqlite_instance._insert_entry, PromptMemoryEntry(entry=p)) + + # Read with limit=2 + projs = await sqlite_instance.get_attack_result_conversation_projections_async( + attack_result_id=ar_id, + include_messages=True, + limit_messages_per_conversation=2, + ) + + assert projs.total_count == 1 + assert projs.truncated is False + assert len(projs.conversations) == 1 + proj = projs.conversations[0] + assert proj.conversation_id == conv_id + assert proj.message_count == 3 # Exact count from COUNT query + assert proj.messages_truncated is True + assert len(proj.messages) == 2 # Bounded by SQL LIMIT + + # Verify multimodal pieces in message 0 + msg0 = proj.messages[0] + assert msg0.role == "user" + assert len(msg0.pieces) == 2 + assert msg0.pieces[0].data_type == "text" + assert msg0.pieces[0].content == "Analyze blueprint image" + assert msg0.pieces[1].data_type == "image_path" + assert msg0.pieces[1].content == img_path + assert msg0.pieces[1].converted_value == img_path + + # Verify raw custom role in message 1 + msg1 = proj.messages[1] + assert msg1.role == "simulated_assistant" + assert msg1.content == "I detect architectural blueprints." + + +def test_hierarchy_repeated_attempt_diagnostic_retries_consistency() -> None: + """Test that units with repeated attempts without row retries match overall_counts and diagnostics retries.""" + # 2 attempts on the same unit: attempt 1 error, attempt 2 success, both with total_retries=0 + delta1 = _make_test_delta( + attack_result_id="attempt-1", + outcome=AttackOutcome.ERROR, + role=AttackResultRole.TARGET_FACING, + parent_collection="group_1", + seed_group_id="seed_1", + seconds_offset=1, + ) + delta2 = _make_test_delta( + attack_result_id="attempt-2", + outcome=AttackOutcome.SUCCESS, + role=AttackResultRole.TARGET_FACING, + parent_collection="group_1", + seed_group_id="seed_1", + seconds_offset=2, + ) + + snapshot = _make_test_snapshot(deltas=[delta1, delta2]) + projection = ScenarioProgressReadModel.build_scenario_hierarchy( + scenario_result_id="scenario-run-1", + scenario_name="TestScenario", + scenario_version=1, + status=ScenarioRunState.COMPLETED, + snapshot=snapshot, + ) + + # Re-attempt of same unit causes retry pressure of 1 + assert projection.overall_counts.retries == 1 + # Diagnostics total_retries must be consistent with overall_counts.retries + assert projection.diagnostics.total_retries == 1 + assert projection.diagnostics.total_retries == projection.overall_counts.retries + + +def test_hierarchy_traversal_halts_immediately_on_budget_exhaustion() -> None: + """Test that with 1,000 root results and max_nodes=1, traversal stops immediately without quadratic loop.""" + deltas = [ + _make_test_delta( + attack_result_id=f"root-{i}", + role=AttackResultRole.TARGET_FACING, + parent_collection="group_1", + seed_group_id=f"seed_{i}", + seconds_offset=i, + ) + for i in range(1000) + ] + + snapshot = _make_test_snapshot(deltas=deltas) + + import time + + start_time = time.perf_counter() + projection = ScenarioProgressReadModel.build_scenario_hierarchy( + scenario_result_id="scenario-run-1", + scenario_name="TestScenario", + scenario_version=1, + status=ScenarioRunState.COMPLETED, + snapshot=snapshot, + max_nodes=1, + ) + elapsed = time.perf_counter() - start_time + + assert elapsed < 0.05, f"Expected <0.05s, got {elapsed:.4f}s" + assert projection.truncated is True + # Count total nodes in projection + node_count = sum(len(g.roots) for g in projection.groups) + assert node_count == 1 + # Aggregate total remains exact + assert projection.overall_counts.completed == 1000 + + +@pytest.mark.asyncio +async def test_sqlite_conversation_list_clipping_explicit(sqlite_instance: SQLiteMemory) -> None: + """Test explicit conversation-list clipping flags and exact count with 55 conversations.""" + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="ConversationClippingScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + ar_uuid = uuid.uuid4() + ar_id = str(ar_uuid) + main_conv_id = str(uuid.uuid4()) + # 54 preparation conversations + 1 main conversation = 55 total conversations + prep_conv_ids = [str(uuid.uuid4()) for _ in range(54)] + + ar = AttackResult( + attack_result_id=ar_id, + conversation_id=main_conv_id, + related_conversations={ + ConversationReference(conversation_id=cid, conversation_type=ConversationType.PREPARATION) + for cid in prep_conv_ids + }, + objective="Conversation clipping test", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={"parent_collection": "group_1", "result_role": AttackResultRole.TARGET_FACING.value}, + ) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[ar]) + + # 1. Direct memory read returns ScenarioAttackResultConversations with total_count=55 and truncated=True + conv_result = await sqlite_instance.get_attack_result_conversation_projections_async( + attack_result_id=ar_id, + max_conversations=50, + ) + assert conv_result.total_count == 55 + assert conv_result.truncated is True + assert len(conv_result.conversations) == 50 + + # 2. Service single-node read projects exact count and truncation flag on node + service = ScenarioRunService(memory=sqlite_instance) + node = await service.get_run_hierarchy_node_async( + scenario_result_id=scenario_id, + attack_result_id=ar_id, + ) + assert node is not None + assert node.conversations_count == 55 + assert node.conversations_truncated is True + assert len(node.conversations) == 50 + + +async def test_executed_turns_zero_vs_invalid_metadata(sqlite_instance: SQLiteMemory) -> None: + """Test that executed_turns=0 with invalid string in metadata safely projects 0 without ValueError.""" + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="ExecutedTurnsZeroScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + ar_uuid = uuid.uuid4() + ar_id = str(ar_uuid) + ar = AttackResult( + attack_result_id=ar_id, + conversation_id=str(uuid.uuid4()), + objective="Executed turns zero test", + outcome=AttackOutcome.SUCCESS, + executed_turns=0, + metadata={"executed_turns": "invalid"}, + attribution_parent_id=scenario_id, + attribution_data={"parent_collection": "group_1", "result_role": AttackResultRole.TARGET_FACING.value}, + ) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[ar]) + + service = ScenarioRunService(memory=sqlite_instance) + node = await service.get_run_hierarchy_node_async( + scenario_result_id=scenario_id, + attack_result_id=ar_id, + ) + assert node is not None + assert node.executed_turns == 0 + + hierarchy = await service.get_run_hierarchy_async(scenario_result_id=scenario_id) + assert hierarchy is not None + assert len(hierarchy.groups) == 1 + assert hierarchy.groups[0].roots[0].executed_turns == 0 + + +def test_eliminate_tuple_duck_typing_raw_tuples() -> None: + """Test passing 2, 3, and 4 conversation projections as a raw Python tuple preserves all items.""" + delta = _make_test_delta( + attack_result_id="test-node", + outcome=AttackOutcome.SUCCESS, + ) + result = _make_test_result( + attack_result_id="test-node", + atomic_group_id="group-1", + atomic_attack_name="attack", + seed_group_id="seed-1", + outcome=AttackOutcome.SUCCESS, + ) + conv1 = ScenarioConversationProjection(conversation_id="conv-1") + conv2 = ScenarioConversationProjection(conversation_id="conv-2") + conv3 = ScenarioConversationProjection(conversation_id="conv-3") + conv4 = ScenarioConversationProjection(conversation_id="conv-4") + + # 2-tuple + node2 = ScenarioProgressReadModel.build_hierarchy_node( + delta=delta, + result=result, + conversations=(conv1, conv2), + ) + assert len(node2.conversations) == 2 + assert node2.conversations == [conv1, conv2] + assert node2.conversations_count == 2 + assert node2.conversations_truncated is False + + # 3-tuple (specifically targeted in Roman's review comment 14) + node3 = ScenarioProgressReadModel.build_hierarchy_node( + delta=delta, + result=result, + conversations=(conv1, conv2, conv3), + ) + assert len(node3.conversations) == 3 + assert node3.conversations == [conv1, conv2, conv3] + assert node3.conversations_count == 3 + assert node3.conversations_truncated is False + + # 4-tuple + node4 = ScenarioProgressReadModel.build_hierarchy_node( + delta=delta, + result=result, + conversations=(conv1, conv2, conv3, conv4), + ) + assert len(node4.conversations) == 4 + assert node4.conversations == [conv1, conv2, conv3, conv4] + assert node4.conversations_count == 4 + assert node4.conversations_truncated is False + + +async def test_sqlite_message_pieces_bounding(sqlite_instance: SQLiteMemory) -> None: + """Test persisting 1 message sequence with 150 pieces and verifying piece truncation at max_pieces_per_message.""" + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="MessagePiecesBoundingScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + ar_uuid = uuid.uuid4() + ar_id = str(ar_uuid) + conv_id = str(uuid.uuid4()) + + ar = AttackResult( + attack_result_id=ar_id, + conversation_id=conv_id, + objective="Message pieces bounding test", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={"parent_collection": "group_1", "result_role": AttackResultRole.TARGET_FACING.value}, + ) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[ar]) + + # Insert 150 pieces for sequence 0 + now = datetime.now(UTC) + for i in range(150): + piece = MessagePiece( + id=uuid.uuid4(), + conversation_id=conv_id, + role="user", + original_value_data_type="text", + original_value=f"Piece content {i}", + converted_value=f"Piece content {i}", + sequence=0, + timestamp=now + timedelta(milliseconds=i), + ) + await set_message_piece_sha256_async(piece) + await sqlite_instance._run_database_operation_async( + sqlite_instance._insert_entry, PromptMemoryEntry(entry=piece) + ) + + # 1. Fetch with default max_pieces_per_message=100 + projs = await sqlite_instance.get_attack_result_conversation_projections_async( + attack_result_id=ar_id, + include_messages=True, + max_pieces_per_message=100, + ) + assert len(projs.conversations) == 1 + assert len(projs.conversations[0].messages) == 1 + msg = projs.conversations[0].messages[0] + assert len(msg.pieces) == 100 + assert msg.pieces_truncated is True + + # 2. Fetch with higher limit max_pieces_per_message=200 + projs_untruncated = await sqlite_instance.get_attack_result_conversation_projections_async( + attack_result_id=ar_id, + include_messages=True, + max_pieces_per_message=200, + ) + assert len(projs_untruncated.conversations[0].messages) == 1 + msg_untruncated = projs_untruncated.conversations[0].messages[0] + assert len(msg_untruncated.pieces) == 150 + assert msg_untruncated.pieces_truncated is False + + +async def test_unloaded_conversations_in_full_hierarchy(sqlite_instance: SQLiteMemory) -> None: + """Test that full hierarchy marks nodes with unhydrated related conversations as truncated with count=None.""" + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="UnloadedConversationsScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + ar_uuid = uuid.uuid4() + ar_id = str(ar_uuid) + main_conv_id = str(uuid.uuid4()) + prep_conv_ids = [str(uuid.uuid4()) for _ in range(55)] + + ar = AttackResult( + attack_result_id=ar_id, + conversation_id=main_conv_id, + related_conversations={ + ConversationReference(conversation_id=cid, conversation_type=ConversationType.PREPARATION) + for cid in prep_conv_ids + }, + objective="Unloaded conversations test", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={"parent_collection": "group_1", "result_role": AttackResultRole.TARGET_FACING.value}, + ) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[ar]) + + service = ScenarioRunService(memory=sqlite_instance) + hierarchy = await service.get_run_hierarchy_async(scenario_result_id=scenario_id) + assert hierarchy is not None + assert len(hierarchy.groups) == 1 + root_node = hierarchy.groups[0].roots[0] + assert root_node.attack_result_id == ar_id + assert root_node.conversations_count is None + assert root_node.conversations_truncated is True + assert len(root_node.conversations) == 1 + + +def test_diagnostic_total_retries_scope_unmatched_attempt() -> None: + """Test that diagnostics.total_retries reconciles with unmatched attempts having retry events.""" + plan = ScenarioRunPlan( + atomic_groups=[ + ScenarioRunPlanAtomicGroup( + id="group-1", + atomic_attack_name="planned_attack", + display_group="Group 1", + technique_eval_hash="", + seed_group_ids=["seed-1"], + ) + ], + seed_groups=[ScenarioRunPlanSeedGroup(id="seed-1", objective_sha256="hash123", objective="Planned objective")], + ) + + planned_delta = _make_test_delta( + attack_result_id="att-planned", + outcome=AttackOutcome.SUCCESS, + parent_collection="planned_attack", + seed_group_id="seed-1", + ) + planned_result = _make_test_result( + attack_result_id="att-planned", + atomic_group_id="group-1", + atomic_attack_name="planned_attack", + seed_group_id="seed-1", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + retries=[], + ) + + # Unmatched error attempt with 3 retry events + retry_events = [ + RetryEvent(exception_type="RateLimitError", attempt_number=1), + RetryEvent(exception_type="RateLimitError", attempt_number=2), + RetryEvent(exception_type="TimeoutError", attempt_number=3), + ] + unmatched_delta = _make_test_delta( + attack_result_id="att-unmatched", + outcome=AttackOutcome.ERROR, + retries=retry_events, + error_type="TimeoutError", + error_message="Timed out", + parent_collection="unmatched_attack", + seed_group_id="unmatched-seed", + ) + unmatched_result = _make_test_result( + attack_result_id="att-unmatched", + atomic_group_id="unmatched-group", + atomic_attack_name="unmatched_attack", + seed_group_id="unmatched-seed", + outcome=AttackOutcome.ERROR, + total_retries=3, + retries=retry_events, + ) + + summary = ScenarioProgressSummary( + overall=ScenarioProgressCounts( + completed=1, + planned=1, + succeeded=1, + success_percentage=100, + errors=0, + retries=0, # Canonical planned summary had 0 retries + ), + atomic_groups=[], + ) + + snapshot = ScenarioProgressSnapshot( + deltas=(planned_delta, unmatched_delta), + results=(planned_result, unmatched_result), + summary=summary, + plan=plan, + ) + + hierarchy = ScenarioProgressReadModel.build_scenario_hierarchy( + scenario_result_id="scen-1", + scenario_name="RetryScopeScenario", + scenario_version=1, + status=ScenarioRunState.COMPLETED, + snapshot=snapshot, + ) + + assert hierarchy.diagnostics.total_retries == 3 + assert hierarchy.diagnostics.attack_retries_count == 1 + assert len(hierarchy.diagnostics.attack_retries) == 1 + assert hierarchy.diagnostics.attack_retries[0].attack_result_id == "att-unmatched" + assert hierarchy.diagnostics.attack_retries[0].total_retries == 3 + + +async def test_immediate_children_bounding_in_run_service(sqlite_instance: SQLiteMemory) -> None: + """ + Test that get_run_hierarchy_node_async bounds immediate children to + max_immediate_children with children_truncated=True. + """ + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="ImmediateChildrenBoundingScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + # 250 child attack results + child_ids = [str(uuid.uuid4()) for _ in range(250)] + parent_id = str(uuid.uuid4()) + + parent_ar = AttackResult( + attack_result_id=parent_id, + conversation_id=str(uuid.uuid4()), + objective="Orchestration parent with 250 children", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={"parent_collection": "group_1", "result_role": AttackResultRole.ORCHESTRATION.value}, + metadata={"child_attack_result_ids": child_ids}, + ) + child_ars = [ + AttackResult( + attack_result_id=cid, + conversation_id=str(uuid.uuid4()), + objective=f"Child {i}", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={"parent_collection": "group_1", "result_role": AttackResultRole.TARGET_FACING.value}, + ) + for i, cid in enumerate(child_ids) + ] + + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[parent_ar] + child_ars) + + service = ScenarioRunService(memory=sqlite_instance) + node = await service.get_run_hierarchy_node_async( + scenario_result_id=scenario_id, + attack_result_id=parent_id, + max_immediate_children=200, + ) + + assert node is not None + assert node.children_truncated is True + assert len(node.children) == 200 + + +async def test_persisted_result_with_malformed_related_conversations_metadata( + sqlite_instance: SQLiteMemory, +) -> None: + """Test that malformed related_conversation_ids metadata does not cause TypeError in progress or node reader.""" + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="MalformedRelatedConvScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + ar_uuid = uuid.uuid4() + ar_id = str(ar_uuid) + ar = AttackResult( + attack_result_id=ar_id, + conversation_id=str(uuid.uuid4()), + objective="Malformed related_conversation_ids test", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={"parent_collection": "group_1", "result_role": AttackResultRole.TARGET_FACING.value}, + metadata={"related_conversation_ids": 7}, + ) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[ar]) + + service = ScenarioRunService(memory=sqlite_instance) + + # 1. get_run_progress_async should succeed without raising TypeError + progress = await service.get_run_progress_async(scenario_result_id=scenario_id, since=None, limit=25) + assert progress is not None + assert len(progress.results) == 1 + assert progress.results[0].related_conversation_ids == [] + + # 2. get_run_hierarchy_node_async should succeed without raising TypeError + node = await service.get_run_hierarchy_node_async( + scenario_result_id=scenario_id, + attack_result_id=ar_id, + ) + assert node is not None + assert node.conversations_count == 1 + assert node.conversations_truncated is False + assert len(node.conversations) == 1 + + # 3. get_run_hierarchy_async should also succeed without raising TypeError + hierarchy = await service.get_run_hierarchy_async(scenario_result_id=scenario_id) + assert hierarchy is not None + assert len(hierarchy.groups) == 1 + root_node = hierarchy.groups[0].roots[0] + assert root_node.conversations_count == 1 + assert root_node.conversations_truncated is False + + +def test_read_related_conversation_ids_shapes() -> None: + """Test _read_related_conversation_ids with typed list, legacy valid list, and malformed fallbacks.""" + # Typed list + delta1 = _make_test_delta(attack_result_id="d1") + delta1.related_conversation_ids = ["conv-1", "conv-2"] + assert ScenarioProgressReadModel._read_related_conversation_ids(delta=delta1) == ["conv-1", "conv-2"] + + # Legacy valid metadata list + delta2 = _make_test_delta(attack_result_id="d2") + delta2.related_conversation_ids = [] + delta2.attack_metadata = {"related_conversation_ids": ["conv-3", "conv-4"]} + assert ScenarioProgressReadModel._read_related_conversation_ids(delta=delta2) == ["conv-3", "conv-4"] + + # Legacy metadata with non-string elements filtered out + delta3 = _make_test_delta(attack_result_id="d3") + delta3.related_conversation_ids = [] + delta3.attack_metadata = {"related_conversation_ids": ["conv-5", 42, None, "conv-6"]} + assert ScenarioProgressReadModel._read_related_conversation_ids(delta=delta3) == ["conv-5", "conv-6"] + + # Typed list with non-string elements filtered + delta4 = _make_test_delta(attack_result_id="d4") + delta4.related_conversation_ids = ["conv-7", 100, None] # type: ignore[list-item] + assert ScenarioProgressReadModel._read_related_conversation_ids(delta=delta4) == ["conv-7"] + + # Malformed typed non-list fallbacks safely + delta5 = _make_test_delta(attack_result_id="d5") + delta5.related_conversation_ids = 123 # type: ignore[assignment] + delta5.attack_metadata = {"related_conversation_ids": ["conv-fallback"]} + assert ScenarioProgressReadModel._read_related_conversation_ids(delta=delta5) == ["conv-fallback"] + + # Malformed metadata shapes (int, str, dict) return empty list + for malformed in [7, "conv-single", {"id": "conv-1"}, None, True]: + delta_bad = _make_test_delta(attack_result_id="dbad") + delta_bad.related_conversation_ids = [] + delta_bad.attack_metadata = {"related_conversation_ids": malformed} + assert ScenarioProgressReadModel._read_related_conversation_ids(delta=delta_bad) == [] + + +async def test_invalid_scenario_id_returns_empty_deltas(sqlite_instance: SQLiteMemory) -> None: + """Test that non-UUID scenario_result_id returns empty list without raising ValueError.""" + deltas = await sqlite_instance.get_scenario_attack_result_deltas_by_ids_async( + scenario_result_id="not-a-valid-uuid", + attack_result_ids=["bad-id"], + ) + assert deltas == [] + + +def test_build_hierarchy_node_deduplicates_duplicate_child_ids() -> None: + """Test that duplicate child IDs in stored order are deduplicated and tracked as unresolved.""" + delta = _make_test_delta(attack_result_id="parent") + result = _make_test_result( + attack_result_id="parent", + child_attack_result_ids=["child-1", "child-1"], + ) + c1_delta = _make_test_delta(attack_result_id="child-1") + c1_result = _make_test_result(attack_result_id="child-1") + + node = ScenarioProgressReadModel.build_hierarchy_node( + delta=delta, + result=result, + children_deltas=[c1_delta], + children_results=[c1_result], + ) + assert len(node.children) == 1 + assert node.children[0].attack_result_id == "child-1" + assert node.unresolved_child_attack_result_ids == ["child-1"] + + +async def test_persisted_conversation_statistics_parity_between_hierarchy_and_node( + sqlite_instance: SQLiteMemory, +) -> None: + """Test that real conversation message counts are reported in both full hierarchy and node detail.""" + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="RealStatsScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + conv_id = str(uuid.uuid4()) + ar_uuid = uuid.uuid4() + ar_id = str(ar_uuid) + ar = AttackResult( + attack_result_id=ar_id, + conversation_id=conv_id, + objective="Real conversation stats test", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={"parent_collection": "group_1", "result_role": AttackResultRole.TARGET_FACING.value}, + metadata={}, # No count metadata! + ) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[ar]) + + # Insert 3 saved messages into SQLite + for seq in range(3): + piece = MessagePiece( + id=uuid.uuid4(), + conversation_id=conv_id, + role="user" if seq % 2 == 0 else "assistant", + original_value_data_type="text", + original_value=f"Message {seq}", + converted_value=f"Message {seq}", + sequence=seq, + ) + await set_message_piece_sha256_async(piece) + await sqlite_instance._run_database_operation_async( + sqlite_instance._insert_entry, PromptMemoryEntry(entry=piece) + ) + + service = ScenarioRunService(memory=sqlite_instance) + + # 1. Node detail endpoint: full messages loaded, exact count=3, messages_truncated=False + node = await service.get_run_hierarchy_node_async( + scenario_result_id=scenario_id, + attack_result_id=ar_id, + ) + assert node is not None + assert len(node.conversations) == 1 + assert node.conversations[0].message_count == 3 + assert len(node.conversations[0].messages) == 3 + assert node.conversations[0].messages_truncated is False + + # 2. Full hierarchy endpoint: real count=3 via bounded stats read, empty messages list, messages_truncated=True + hierarchy = await service.get_run_hierarchy_async(scenario_result_id=scenario_id) + assert hierarchy is not None + root_node = hierarchy.groups[0].roots[0] + assert len(root_node.conversations) == 1 + assert root_node.conversations[0].message_count == 3 + assert len(root_node.conversations[0].messages) == 0 + assert root_node.conversations[0].messages_truncated is True + + # 3. Setting metadata={"conversation_message_count": 99} does not override real database stats + ar2_uuid = uuid.uuid4() + ar2_id = str(ar2_uuid) + ar2 = AttackResult( + attack_result_id=ar2_id, + conversation_id=conv_id, + objective="Metadata precedence test", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={"parent_collection": "group_1", "result_role": AttackResultRole.TARGET_FACING.value}, + metadata={"conversation_message_count": 99}, + ) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[ar2]) + hierarchy2 = await service.get_run_hierarchy_async(scenario_result_id=scenario_id) + assert hierarchy2 is not None + root_node2 = [r for r in hierarchy2.groups[0].roots if r.attack_result_id == ar2_id][0] + assert root_node2.conversations[0].message_count == 3 # Still real count 3, not 99! + + +async def test_persisted_result_with_non_numeric_conversation_message_count_metadata( + sqlite_instance: SQLiteMemory, +) -> None: + """Test that metadata={'conversation_message_count': 'unknown'} does not raise ValueError in hierarchy.""" + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="NonNumericMetadataScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + ar_uuid = uuid.uuid4() + ar_id = str(ar_uuid) + ar = AttackResult( + attack_result_id=ar_id, + conversation_id=str(uuid.uuid4()), + objective="Non-numeric metadata test", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={"parent_collection": "group_1", "result_role": AttackResultRole.TARGET_FACING.value}, + metadata={"conversation_message_count": "unknown"}, + ) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[ar]) + + service = ScenarioRunService(memory=sqlite_instance) + + # Must succeed without raising ValueError + hierarchy = await service.get_run_hierarchy_async(scenario_result_id=scenario_id) + assert hierarchy is not None + root_node = hierarchy.groups[0].roots[0] + assert len(root_node.conversations) == 1 + # Real SQLite has 0 messages for this new conversation, so count is 0 and not truncated + assert root_node.conversations[0].message_count == 0 + assert root_node.conversations[0].messages_truncated is False + + +async def test_bounded_conversation_messages_piece_budget_applied_at_sql_layer( + sqlite_instance: SQLiteMemory, +) -> None: + """Test that piece budget is applied at SQL layer, preventing hydration of unbounded pieces (e.g. 1,001 rows).""" + ar_id = str(uuid.uuid4()) + conv_id = str(uuid.uuid4()) + now = datetime.now(UTC) + + # Insert 1,001 pieces for sequence 0 + pieces_seq0 = [ + MessagePiece( + role="user", + original_value=f"piece-{i}", + converted_value=f"piece-{i}", + conversation_id=conv_id, + sequence=0, + timestamp=now + timedelta(milliseconds=i), + ) + for i in range(1001) + ] + # Insert 5 pieces for sequence 1 + pieces_seq1 = [ + MessagePiece( + role="assistant", + original_value=f"response-{i}", + converted_value=f"response-{i}", + conversation_id=conv_id, + sequence=1, + timestamp=now + timedelta(seconds=1, milliseconds=i), + ) + for i in range(5) + ] + for p in pieces_seq0 + pieces_seq1: + await set_message_piece_sha256_async(p) + await sqlite_instance.add_message_pieces_to_memory_async(message_pieces=pieces_seq0 + pieces_seq1) + + ar = AttackResult( + attack_result_id=ar_id, + conversation_id=conv_id, + objective="Piece bounding test", + outcome=AttackOutcome.SUCCESS, + ) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[ar]) + + # 1. Fetch via get_attack_result_conversation_projections_async with max_pieces_per_message=100 + projs = await sqlite_instance.get_attack_result_conversation_projections_async( + attack_result_id=ar_id, + include_messages=True, + max_pieces_per_message=100, + ) + assert len(projs.conversations) == 1 + messages = projs.conversations[0].messages + assert len(messages) == 2 + + # Sequence 0 (1,001 pieces) is capped at 100 with pieces_truncated=True + msg0 = messages[0] + assert msg0.sequence == 0 + assert len(msg0.pieces) == 100 + assert msg0.pieces_truncated is True + + # Sequence 1 (5 pieces) is fully returned with pieces_truncated=False + msg1 = messages[1] + assert msg1.sequence == 1 + assert len(msg1.pieces) == 5 + assert msg1.pieces_truncated is False + + # 2. Verify SQL-level row bounding: _fetch_bounded_conversation_messages only fetches + # at most (max_pieces_per_message + 1) rows per sequence directly from the DB. + # Testing directly with max_pieces_per_message=10: + messages_capped_10, _ = sqlite_instance._fetch_bounded_conversation_messages( + conversation_id=conv_id, + limit=10, + max_pieces_per_message=10, + ) + assert len(messages_capped_10[0].pieces) == 10 + assert messages_capped_10[0].pieces_truncated is True + assert len(messages_capped_10[1].pieces) == 5 + assert messages_capped_10[1].pieces_truncated is False + + +async def test_diagnostics_total_retries_count_event_mismatch_parity( + sqlite_instance: SQLiteMemory, +) -> None: + """ + Test that when total_retries=0 but retry events exist, run detail, + producer counts, diagnostic item, and diagnostics.total_retries all agree on 0. + """ + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="MismatchRetryScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + retry_events = [ + RetryEvent(exception_type="RateLimitError", attempt_number=1), + RetryEvent(exception_type="RateLimitError", attempt_number=2), + RetryEvent(exception_type="TimeoutError", attempt_number=3), + ] + ar_id = str(uuid.uuid4()) + ar = AttackResult( + attack_result_id=ar_id, + conversation_id=str(uuid.uuid4()), + objective="Retry mismatch test", + outcome=AttackOutcome.ERROR, + total_retries=0, # Persisted count is 0 + retry_events=retry_events, # 3 retry events recorded + error_type="TimeoutError", + error_message="Timed out", + attribution_parent_id=scenario_id, + attribution_data={"parent_collection": "group_1", "result_role": AttackResultRole.TARGET_FACING.value}, + ) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[ar]) + + service = ScenarioRunService(memory=sqlite_instance) + + # 1. Run summary / detail reports total_retries = 0 + run_summary = await service.get_run_async(scenario_result_id=scenario_id) + assert run_summary is not None + assert run_summary.total_retries == 0 + + # 2. Hierarchy diagnostics and overall counts report total_retries = 0 + hierarchy = await service.get_run_hierarchy_async(scenario_result_id=scenario_id) + assert hierarchy is not None + assert hierarchy.overall_counts.retries == 0 + assert hierarchy.diagnostics.total_retries == 0 + assert hierarchy.diagnostics.attack_retries_count == 1 + assert len(hierarchy.diagnostics.attack_retries) == 1 + + # The individual diagnostic item reports total_retries = 0 while preserving the 3 events + diag_item = hierarchy.diagnostics.attack_retries[0] + assert diag_item.total_retries == 0 + assert len(diag_item.retries) == 3 + + +@pytest.mark.asyncio +async def test_clipped_tree_queries_actual_preview_node_conversation_stats(sqlite_instance: SQLiteMemory) -> None: + """Test that when child is saved before parent and max_nodes=1, the actual returned parent gets real stats.""" + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="ClippedTreeStatsScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + # 1. Parent with its own 3-message conversation + parent_conv_id = str(uuid.uuid4()) + for seq in range(3): + piece = MessagePiece( + id=uuid.uuid4(), + conversation_id=parent_conv_id, + role="user" if seq % 2 == 0 else "assistant", + original_value_data_type="text", + original_value=f"Parent message {seq}", + converted_value=f"Parent message {seq}", + sequence=seq, + ) + await set_message_piece_sha256_async(piece) + await sqlite_instance._run_database_operation_async( + sqlite_instance._insert_entry, PromptMemoryEntry(entry=piece) + ) + + child_conv_id = str(uuid.uuid4()) + child_id = str(uuid.uuid4()) + parent_id = str(uuid.uuid4()) + + # Child is saved FIRST in memory + child_ar = AttackResult( + attack_result_id=child_id, + conversation_id=child_conv_id, + objective="Child attack", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={"parent_collection": "group_1", "result_role": AttackResultRole.TARGET_FACING.value}, + ) + # Parent references child + parent_ar = AttackResult( + attack_result_id=parent_id, + conversation_id=parent_conv_id, + objective="Parent attack", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={"parent_collection": "group_1", "result_role": AttackResultRole.ORCHESTRATION.value}, + metadata={"child_attack_result_ids": [child_id]}, + ) + # Add child first, then parent + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[child_ar, parent_ar]) + + service = ScenarioRunService(memory=sqlite_instance) + + # With max_nodes=1, parent root is returned. It must report real count=3, not 0! + hierarchy = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert hierarchy is not None + assert hierarchy.truncated is True + assert len(hierarchy.groups[0].roots) == 1 + root = hierarchy.groups[0].roots[0] + assert root.attack_result_id == parent_id + assert len(root.conversations) == 1 + assert root.conversations[0].conversation_id == parent_conv_id + assert root.conversations[0].message_count == 3 + assert root.conversations[0].messages_truncated is True + + +@pytest.mark.asyncio +async def test_conversation_stats_query_failure_does_not_become_false_zero(sqlite_instance: SQLiteMemory) -> None: + """Test that a failed stats query does not turn uncounted conversations into message_count=0.""" + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="FailedStatsScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + conv_id = str(uuid.uuid4()) + ar_id = str(uuid.uuid4()) + ar = AttackResult( + attack_result_id=ar_id, + conversation_id=conv_id, + objective="Failed stats test", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={"parent_collection": "group_1", "result_role": AttackResultRole.TARGET_FACING.value}, + ) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[ar]) + + service = ScenarioRunService(memory=sqlite_instance) + # Mock get_conversation_stats_async to fail with an exception + sqlite_instance.get_conversation_stats_async = AsyncMock(side_effect=RuntimeError("Database error")) + + hierarchy = await service.get_run_hierarchy_async(scenario_result_id=scenario_id) + assert hierarchy is not None + root = hierarchy.groups[0].roots[0] + assert len(root.conversations) == 1 + # Must NOT become 0! Must be None and marked truncated/incomplete + assert root.conversations[0].message_count is None + assert root.conversations[0].messages_truncated is True + + +@pytest.mark.asyncio +async def test_parent_with_2500_children_bounds_unresolved_references(sqlite_instance: SQLiteMemory) -> None: + """Test that a parent with 2,500 children bounds unresolved references to 500 in both detail and full hierarchy.""" + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="MassiveChildrenScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + parent_id = str(uuid.uuid4()) + child_ids = [str(uuid.uuid4()) for _ in range(2500)] + + # Add parent and the first 200 children into memory + children = [ + AttackResult( + attack_result_id=cid, + conversation_id=str(uuid.uuid4()), + objective=f"Child {i}", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={"parent_collection": "group_1", "result_role": AttackResultRole.TARGET_FACING.value}, + ) + for i, cid in enumerate(child_ids[:200]) + ] + parent = AttackResult( + attack_result_id=parent_id, + conversation_id=str(uuid.uuid4()), + objective="Parent with 2500 children", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={"parent_collection": "group_1", "result_role": AttackResultRole.ORCHESTRATION.value}, + metadata={"child_attack_result_ids": child_ids}, + ) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[parent, *children]) + + service = ScenarioRunService(memory=sqlite_instance) + + # 1. Single node detail: with 200 immediate children loaded, unresolved references are capped at 500 (not 2,300) + node = await service.get_run_hierarchy_node_async( + scenario_result_id=scenario_id, + attack_result_id=parent_id, + max_immediate_children=200, + ) + assert node is not None + assert len(node.children) == 200 + assert len(node.unresolved_child_attack_result_ids) == 500 + assert node.unresolved_child_references_truncated is True + + # 2. Full hierarchy with max_nodes=1: parent has 2,500 unresolved children capped at 500 (not 2,500) + hierarchy = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert hierarchy is not None + root = hierarchy.groups[0].roots[0] + assert len(root.unresolved_child_attack_result_ids) == 500 + assert root.unresolved_child_references_truncated is True + + +@pytest.mark.asyncio +async def test_cold_hierarchy_hydration_bounded_by_max_nodes(sqlite_instance: SQLiteMemory) -> None: + """Test that a cold max_nodes=1 request only fetches 1 delta and does not hydrate 1,001 results across 3 pages.""" + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="ColdHydrationScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + # Insert 1,001 attack results + results = [ + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective=f"Attack {i}", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={"parent_collection": "group_1", "result_role": AttackResultRole.TARGET_FACING.value}, + ) + for i in range(1001) + ] + await sqlite_instance.add_attack_results_to_memory_async(attack_results=results) + + # Track delta queries + original_get_deltas = sqlite_instance.get_scenario_attack_result_deltas_async + query_limits: list[int] = [] + + async def tracking_get_deltas(*args: Any, **kwargs: Any) -> Any: + query_limits.append(kwargs.get("limit", 0)) + return await original_get_deltas(*args, **kwargs) + + sqlite_instance.get_scenario_attack_result_deltas_async = tracking_get_deltas # type: ignore[assignment] + + service = ScenarioRunService(memory=sqlite_instance) + + # Cold max_nodes=1 hierarchy call + hierarchy = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert hierarchy is not None + assert hierarchy.truncated is True + # Must only have executed 1 page query (capped at storage page size), NOT 3 page queries (500, 500, 1)! + assert len(query_limits) == 1 + assert query_limits[0] == 500 + + +def test_build_scenario_hierarchy_unqueried_conversation_not_false_zero() -> None: + """Test that missing-key branch in build_scenario_hierarchy produces None and truncated=True, never false zero.""" + delta = _make_test_delta( + attack_result_id="root-1", + role=AttackResultRole.TARGET_FACING, + parent_collection="group_1", + seed_group_id="seed_1", + conversation_id="conv-unqueried", + ) + snapshot = _make_test_snapshot(deltas=[delta]) + # Empty stats passed (e.g. stats read failure or unqueried) + projection = ScenarioProgressReadModel.build_scenario_hierarchy( + scenario_result_id="scenario-run-1", + scenario_name="TestScenario", + scenario_version=1, + status=ScenarioRunState.COMPLETED, + snapshot=snapshot, + conversation_stats={}, + ) + conv = projection.groups[0].roots[0].conversations[0] + # Must NOT report 0 and False! + assert conv.message_count is None + assert conv.messages_truncated is True + + +@pytest.mark.asyncio +async def test_cold_hierarchy_preserves_run_wide_totals_for_tail_error(sqlite_instance: SQLiteMemory) -> None: + """Test that max_nodes=1 on cold run with 501 results accurately reports 501 units, 1 error, and 7 retries.""" + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="TailErrorScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + # 500 successes, 501st has an error and 7 retries + now = datetime.now(UTC) + results = [ + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective=f"Attack {i}", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={"parent_collection": f"group_{i}", "result_role": AttackResultRole.TARGET_FACING.value}, + timestamp=now + timedelta(seconds=i), + ) + for i in range(500) + ] + retry_events = [ + RetryEvent( + attempt_number=r, + timestamp=now + timedelta(seconds=500 + r), + exception_message="Network glitch", + exception_type="NetworkError", + ) + for r in range(1, 8) + ] + tail_error = AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective="Attack 500 (failed)", + outcome=AttackOutcome.ERROR, + error_type="FatalError", + error_message="Failed after retries", + attribution_parent_id=scenario_id, + attribution_data={"parent_collection": "group_500", "result_role": AttackResultRole.TARGET_FACING.value}, + timestamp=now + timedelta(seconds=510), + retry_events=retry_events, + total_retries=7, + ) + results.append(tail_error) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=results) + + service = ScenarioRunService(memory=sqlite_instance) + + # 1. Cold hierarchy with max_nodes=1 + hierarchy_cold = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert hierarchy_cold is not None + assert hierarchy_cold.truncated is True + assert len(hierarchy_cold.groups[0].roots) == 1 + + # Exact run-wide totals must be retained! + assert hierarchy_cold.overall_counts.completed == 501 + assert hierarchy_cold.overall_counts.errors == 1 + assert hierarchy_cold.overall_counts.retries == 7 + assert hierarchy_cold.overall_counts.producer_counts.target_facing.attempts == 501 + assert hierarchy_cold.diagnostics.failed_attacks_count == 1 + assert hierarchy_cold.diagnostics.failed_attacks_truncated is True + assert hierarchy_cold.diagnostics.total_retries == 7 + assert hierarchy_cold.diagnostics.attack_retries_truncated is True + + # 2. Read progress (which hydrates all 501 deltas into cache) + progress = await service.get_run_progress_async(scenario_result_id=scenario_id, since=None, limit=1000) + assert progress is not None + + # 3. Repeat hierarchy with max_nodes=1 on warm cache + hierarchy_warm = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert hierarchy_warm is not None + + # Cold and warm hierarchy calls must have identical totals without any database writes! + assert hierarchy_warm.overall_counts == hierarchy_cold.overall_counts + assert hierarchy_warm.diagnostics.total_retries == hierarchy_cold.diagnostics.total_retries == 7 + assert hierarchy_warm.diagnostics.failed_attacks_count == hierarchy_cold.diagnostics.failed_attacks_count == 1 + + +@pytest.mark.asyncio +async def test_repeated_attempts_latest_error_beyond_cap_reported_correctly(sqlite_instance: SQLiteMemory) -> None: + """Test that a latest error beyond the 500 cap is reported as an error, not earlier success.""" + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="RepeatedAttemptsScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + now = datetime.now(UTC) + results: list[AttackResult] = [] + # Unit 0 Attempt 1: SUCCESS at time 0 + results.append( + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective="Unit 0", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "unit_0", + "seed_group_id": "seed_0", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=now, + ) + ) + # Units 1..499: SUCCESS at times 1..499 + results.extend( + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective=f"Unit {i}", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": f"unit_{i}", + "seed_group_id": f"seed_{i}", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=now + timedelta(seconds=i), + ) + for i in range(1, 500) + ) + # Unit 0 Attempt 2: ERROR at time 501 (beyond the 500-delta cap) + results.append( + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective="Unit 0", + outcome=AttackOutcome.ERROR, + error_type="Attempt2Error", + error_message="Attempt 2 failed", + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "unit_0", + "seed_group_id": "seed_0", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=now + timedelta(seconds=501), + ) + ) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=results) + + service = ScenarioRunService(memory=sqlite_instance) + + hierarchy = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert hierarchy is not None + # Latest attempt for unit_0 was an error, so overall errors must be 1 (not 0)! + assert hierarchy.overall_counts.errors == 1 + assert hierarchy.overall_counts.succeeded == 499 + assert hierarchy.overall_counts.completed == 500 + + +@pytest.mark.asyncio +async def test_failed_conversation_stats_read_never_falls_back_to_metadata(sqlite_instance: SQLiteMemory) -> None: + """Test that a failed stats query leaves counts unknown/incomplete even if count metadata is present.""" + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="FailedStatsMetadataScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + conv_id = str(uuid.uuid4()) + ar_id = str(uuid.uuid4()) + # 3 saved messages in SQLite + now = datetime.now(UTC) + for i in range(3): + piece = MessagePiece( + id=uuid.uuid4(), + conversation_id=conv_id, + role="user", + original_value_data_type="text", + original_value=f"msg {i}", + converted_value=f"msg {i}", + sequence=i, + timestamp=now + timedelta(seconds=i), + ) + await set_message_piece_sha256_async(piece) + await sqlite_instance._run_database_operation_async( + sqlite_instance._insert_entry, PromptMemoryEntry(entry=piece) + ) + + # Attack result with metadata={"conversation_message_count": 0} + ar = AttackResult( + attack_result_id=ar_id, + conversation_id=conv_id, + objective="Failed stats metadata fallback test", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={"parent_collection": "group_1", "result_role": AttackResultRole.TARGET_FACING.value}, + metadata={"conversation_message_count": 0}, + ) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[ar]) + + service = ScenarioRunService(memory=sqlite_instance) + # Simulate DB query failure + sqlite_instance.get_conversation_stats_async = AsyncMock(side_effect=RuntimeError("DB query failure")) + + hierarchy = await service.get_run_hierarchy_async(scenario_result_id=scenario_id) + assert hierarchy is not None + root = hierarchy.groups[0].roots[0] + assert len(root.conversations) == 1 + # Must NOT return 0 or False from metadata! Must be None and truncated=True! + assert root.conversations[0].message_count is None + assert root.conversations[0].messages_truncated is True + + +@pytest.mark.asyncio +async def test_cold_hierarchy_projection_preserves_parent_beyond_delta_page_limit(sqlite_instance): + """ + Test that a sequential orchestration parent persisted after 500 children is correctly + identified as the root on a cold request (where delta pagination only loads 500 deltas), + matching warm cache behavior and never mistaking child 0 for an independent root. + """ + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + base_time = datetime.now(UTC) + + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="Cold Page Limit Test", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + # Persist 500 children first + child_ids = [str(uuid.uuid4()) for _ in range(500)] + child_results = [ + AttackResult( + attack_result_id=cid, + conversation_id=str(uuid.uuid4()), + objective=f"Child attack {i}", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "seq_group", + "result_role": AttackResultRole.TARGET_FACING.value, + "attempt_index": i + 1, + }, + timestamp=base_time + timedelta(seconds=i), + ) + for i, cid in enumerate(child_ids) + ] + await sqlite_instance.add_attack_results_to_memory_async(attack_results=child_results) + + # Persist 1 Sequential orchestration parent after the 500 children + parent_id = str(uuid.uuid4()) + parent_result = AttackResult( + attack_result_id=parent_id, + conversation_id=str(uuid.uuid4()), + objective="Sequential parent attack", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "seq_group", + "result_role": AttackResultRole.ORCHESTRATION.value, + }, + metadata={"child_attack_result_ids": child_ids}, + timestamp=base_time + timedelta(seconds=501), + ) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[parent_result]) + + service = ScenarioRunService(memory=sqlite_instance) + + # 1. Cold hierarchy request with max_nodes=1 + hierarchy_cold = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert hierarchy_cold is not None + assert len(hierarchy_cold.groups) == 1 + cold_roots = hierarchy_cold.groups[0].roots + assert len(cold_roots) == 1 + + cold_root = cold_roots[0] + # The root MUST be the parent, NOT child 0! + assert cold_root.attack_result_id == parent_id + assert cold_root.parent_attack_result_id is None + assert cold_root.child_index is None + assert cold_root.result_role == AttackResultRole.ORCHESTRATION + assert cold_root.attempt_index is None + + # 2. Warm cache: read progress to hydrate all deltas into cache + await service.get_run_progress_async(scenario_result_id=scenario_id, since=None, limit=500) + + # 3. Warm hierarchy request with max_nodes=1 must match cold exactly + hierarchy_warm = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert hierarchy_warm is not None + warm_roots = hierarchy_warm.groups[0].roots + assert len(warm_roots) == 1 + warm_root = warm_roots[0] + assert warm_root.attack_result_id == parent_id + assert warm_root.parent_attack_result_id is None + assert warm_root.child_index is None + assert warm_root.result_role == AttackResultRole.ORCHESTRATION + + # 4. Request with max_nodes=5 must nest children properly under the parent + hierarchy_expanded = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=5) + assert hierarchy_expanded is not None + expanded_root = hierarchy_expanded.groups[0].roots[0] + assert expanded_root.attack_result_id == parent_id + # 1 root + 4 children = 5 nodes + assert len(expanded_root.children) == 4 + for idx, child in enumerate(expanded_root.children): + assert child.attack_result_id == child_ids[idx] + assert child.parent_attack_result_id == parent_id + assert child.child_index == idx + assert child.attempt_index == idx + 1 + assert child.result_role == AttackResultRole.TARGET_FACING + assert len(expanded_root.unresolved_child_attack_result_ids) == 496 + assert hierarchy_expanded.truncated is True + + +@pytest.mark.asyncio +async def test_empty_child_array_metadata_parent_resolution_parity( + sqlite_instance: SQLiteMemory, +) -> None: + """ + Test that when 500 children have empty metadata {"child_attack_result_ids": []}, + they are not mistaken for orchestration parents and parent at 501 is resolved correctly. + """ + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + base_time = datetime(2026, 1, 1, 0, 0, 0, tzinfo=UTC) + + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="EmptyChildArrayScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + child_ids = [str(uuid.uuid4()) for _ in range(500)] + child_results = [ + AttackResult( + attack_result_id=cid, + conversation_id=str(uuid.uuid4()), + objective=f"Child attack {i}", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "seq_group", + "result_role": AttackResultRole.TARGET_FACING.value, + "attempt_index": i + 1, + }, + metadata={"child_attack_result_ids": []}, + timestamp=base_time + timedelta(seconds=i), + ) + for i, cid in enumerate(child_ids) + ] + await sqlite_instance.add_attack_results_to_memory_async(attack_results=child_results) + + parent_id = str(uuid.uuid4()) + parent_result = AttackResult( + attack_result_id=parent_id, + conversation_id=str(uuid.uuid4()), + objective="Sequential parent attack", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "seq_group", + "result_role": AttackResultRole.ORCHESTRATION.value, + }, + metadata={"child_attack_result_ids": child_ids}, + timestamp=base_time + timedelta(seconds=501), + ) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[parent_result]) + + service = ScenarioRunService(memory=sqlite_instance) + + # 1. Cold hierarchy request with max_nodes=1 + hierarchy_cold = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert hierarchy_cold is not None + assert len(hierarchy_cold.groups) == 1 + cold_roots = hierarchy_cold.groups[0].roots + assert len(cold_roots) == 1 + cold_root = cold_roots[0] + assert cold_root.attack_result_id == parent_id + assert cold_root.parent_attack_result_id is None + assert cold_root.child_index is None + assert cold_root.result_role == AttackResultRole.ORCHESTRATION + + # 2. Warm hierarchy request with max_nodes=1 + await service.get_run_progress_async(scenario_result_id=scenario_id, since=None, limit=500) + hierarchy_warm = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert hierarchy_warm is not None + warm_root = hierarchy_warm.groups[0].roots[0] + assert warm_root.attack_result_id == parent_id + assert warm_root.parent_attack_result_id is None + assert warm_root.child_index is None + + +@pytest.mark.asyncio +async def test_501_orchestration_parents_with_outer_parent_targeted_resolution( + sqlite_instance: SQLiteMemory, +) -> None: + """ + Test targeted parent resolution when there are 501 inner orchestration parents + and 1 outer parent referencing all of them. + """ + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + base_time = datetime(2026, 1, 1, 0, 0, 0, tzinfo=UTC) + + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="ManyParentsScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + # 501 inner parents, each having 1 leaf child + inner_parent_ids = [str(uuid.uuid4()) for _ in range(501)] + all_results: list[AttackResult] = [] + for i, pid in enumerate(inner_parent_ids): + leaf_id = str(uuid.uuid4()) + leaf = AttackResult( + attack_result_id=leaf_id, + conversation_id=str(uuid.uuid4()), + objective=f"Leaf {i}", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "group_orch", + "result_role": AttackResultRole.TARGET_FACING.value, + "attempt_index": 1, + }, + timestamp=base_time + timedelta(seconds=i * 2), + ) + parent = AttackResult( + attack_result_id=pid, + conversation_id=str(uuid.uuid4()), + objective=f"Inner parent {i}", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "group_orch", + "result_role": AttackResultRole.ORCHESTRATION.value, + "attempt_index": i + 1, + }, + metadata={"child_attack_result_ids": [leaf_id]}, + timestamp=base_time + timedelta(seconds=i * 2 + 1), + ) + all_results.extend([leaf, parent]) + + # 1 Outer parent referencing all 501 inner parents + outer_id = str(uuid.uuid4()) + outer_parent = AttackResult( + attack_result_id=outer_id, + conversation_id=str(uuid.uuid4()), + objective="Outer orchestration parent", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "group_orch", + "result_role": AttackResultRole.ORCHESTRATION.value, + }, + metadata={"child_attack_result_ids": inner_parent_ids}, + timestamp=base_time + timedelta(seconds=1100), + ) + all_results.append(outer_parent) + + await sqlite_instance.add_attack_results_to_memory_async(attack_results=all_results) + + service = ScenarioRunService(memory=sqlite_instance) + + # Cold max_nodes=1 hierarchy call: root MUST resolve to outer_parent! + hierarchy_cold = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert hierarchy_cold is not None + assert len(hierarchy_cold.groups) == 1 + roots = hierarchy_cold.groups[0].roots + assert len(roots) == 1 + assert roots[0].attack_result_id == outer_id + assert roots[0].parent_attack_result_id is None + assert roots[0].child_index is None + + +@pytest.mark.asyncio +async def test_unreached_node_zero_inference_attempt_index_parity( + sqlite_instance: SQLiteMemory, +) -> None: + """ + Test zero-inference on unreached nodes: when parent_id is None, child_index MUST be None, + and attempt_index must be preserved on both hierarchy and detail projections. + """ + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="UnreachedZeroInferenceScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + ar_id = str(uuid.uuid4()) + ar = AttackResult( + attack_result_id=ar_id, + conversation_id=str(uuid.uuid4()), + objective="Disconnected node", + outcome=AttackOutcome.SUCCESS, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "group_1", + "result_role": AttackResultRole.TARGET_FACING.value, + "attempt_index": 7, + }, + ) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[ar]) + + service = ScenarioRunService(memory=sqlite_instance) + + # 1. Hierarchy projection + hierarchy = await service.get_run_hierarchy_async(scenario_result_id=scenario_id) + assert hierarchy is not None + assert len(hierarchy.groups) == 1 + node = hierarchy.groups[0].roots[0] + assert node.attack_result_id == ar_id + assert node.parent_attack_result_id is None + assert node.child_index is None + assert node.attempt_index == 7 + + # 2. Detail projection + detail = await service.get_run_hierarchy_node_async( + scenario_result_id=scenario_id, + attack_result_id=ar_id, + ) + assert detail is not None + assert detail.attack_result_id == ar_id + assert detail.parent_attack_result_id is None + assert detail.child_index is None + assert detail.attempt_index == 7 + + +@pytest.mark.asyncio +async def test_501_units_with_tail_error_and_unmatched_tail_error_diagnostics_parity( + sqlite_instance: SQLiteMemory, +) -> None: + """ + Test 501 planned units with tail error + unmatched tail error. + Verify cold max_nodes=1 overall counts, outcome breakdown, group counts, + and diagnostics match warm breakdown. + """ + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + base_time = datetime(2026, 1, 1, 0, 0, 0, tzinfo=UTC) + + # Plan with 501 units + seed_groups = [ + ScenarioRunPlanSeedGroup( + id=f"seed_{i}", + objective=f"Objective {i}", + objective_sha256=f"hash_{i}", + ) + for i in range(501) + ] + atomic_group = ScenarioRunPlanAtomicGroup( + id="group_1", + atomic_attack_name="attack_type_1", + display_group="Attack Type 1", + technique_eval_hash="hash_tech_1", + seed_group_ids=[s.id for s in seed_groups], + ) + plan = ScenarioRunPlan( + atomic_groups=[atomic_group], + seed_groups=seed_groups, + ) + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="TailErrorParityScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + metadata={"run_plan": plan.model_dump(mode="json")}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + # 500 successful units + all_results: list[AttackResult] = [ + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective=f"Objective {i}", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "attack_type_1", + "parent_eval_hash": "hash_tech_1", + "seed_group_id": f"seed_{i}", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=i), + ) + for i in range(500) + ] + + # Unit 500 (the 501st unit): attempt fails with 10 retries + tail_planned_id = str(uuid.uuid4()) + tail_planned_result = AttackResult( + attack_result_id=tail_planned_id, + conversation_id=str(uuid.uuid4()), + objective="Objective 500", + outcome=AttackOutcome.ERROR, + total_retries=10, + error_type="PlannedTailError", + error_message="Planned tail failure", + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "attack_type_1", + "parent_eval_hash": "hash_tech_1", + "seed_group_id": "seed_500", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=500), + ) + all_results.append(tail_planned_result) + + # Unmatched / unplanned tail attempt with 6 retries + unmatched_id = str(uuid.uuid4()) + unmatched_result = AttackResult( + attack_result_id=unmatched_id, + conversation_id=str(uuid.uuid4()), + objective="Unmatched objective", + outcome=AttackOutcome.ERROR, + total_retries=6, + error_type="UnmatchedTailError", + error_message="Unmatched tail failure", + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "unmatched_collection", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=501), + ) + all_results.append(unmatched_result) + + await sqlite_instance.add_attack_results_to_memory_async(attack_results=all_results) + + service = ScenarioRunService(memory=sqlite_instance) + + # 1. Cold hierarchy request with max_nodes=1 + cold_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert cold_hier is not None + assert cold_hier.truncated is True + + # 2. Warm hierarchy request with max_nodes=1 + await service.get_run_progress_async(scenario_result_id=scenario_id, since=None, limit=600) + warm_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert warm_hier is not None + + # Parity checks: + # A. Overall logical outcomes: 500 successes, 1 error (unmatched attempt does not inflate logical unit error count) + assert cold_hier.overall_counts.completed == 501 + assert cold_hier.overall_counts.succeeded == 500 + assert cold_hier.overall_counts.errors == 1 + assert cold_hier.overall_counts.outcomes is not None + assert warm_hier.overall_counts.outcomes is not None + assert cold_hier.overall_counts.outcomes.successes == 500 + assert cold_hier.overall_counts.outcomes.errors == 1 + assert cold_hier.overall_counts.outcomes.successes == warm_hier.overall_counts.outcomes.successes + assert cold_hier.overall_counts.outcomes.errors == warm_hier.overall_counts.outcomes.errors + assert ( + cold_hier.overall_counts.outcomes.success_rate_decided == warm_hier.overall_counts.outcomes.success_rate_decided + ) + + # B. Diagnostics: raw diagnostics report ALL failed attacks (both planned and unmatched) + assert cold_hier.diagnostics.failed_attacks_count == 2 + assert cold_hier.diagnostics.total_retries >= 16 + + # C. Per-group counts parity + assert cold_hier.groups[0].counts.completed == 501 + assert cold_hier.groups[0].counts.succeeded == 500 + assert cold_hier.groups[0].counts.errors == 1 + assert cold_hier.groups[0].counts.completed == warm_hier.groups[0].counts.completed + assert cold_hier.groups[0].counts.succeeded == warm_hier.groups[0].counts.succeeded + assert cold_hier.groups[0].counts.errors == warm_hier.groups[0].counts.errors + + +@pytest.mark.asyncio +async def test_legacy_recount_paged_deltas_without_hydrating_scenario_results( + sqlite_instance: SQLiteMemory, +) -> None: + """ + Test that legacy recount for runs needing SDK statistics pages deltas (batches of 500) + and does NOT call get_scenario_results_async (avoiding memory bloat). + """ + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + base_time = datetime(2026, 1, 1, 0, 0, 0, tzinfo=UTC) + + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="LegacyRecountScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + # 501 results with seed identifiers (triggers needs_sdk_statistics) + results = [ + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective=f"Atomic seed attack {i}", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + atomic_attack_identifier=AtomicAttackIdentifier( + class_name="MockAttack", + seed_identifiers=[SeedIdentifier(class_name="SeedPrompt", value=f"seed_value_{i}")], + ), + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "group_legacy", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=i), + ) + for i in range(501) + ] + await sqlite_instance.add_attack_results_to_memory_async(attack_results=results) + + # Spy on memory methods + original_get_scenario_results = sqlite_instance.get_scenario_results_async + original_get_deltas = sqlite_instance.get_scenario_attack_result_deltas_async + original_get_summaries = sqlite_instance.get_scenario_recount_summaries_async + + scenario_results_called = False + deltas_called = False + summary_page_sizes: list[int] = [] + + async def tracking_get_scenario_results(*args: Any, **kwargs: Any) -> Any: + nonlocal scenario_results_called + scenario_results_called = True + return await original_get_scenario_results(*args, **kwargs) + + async def tracking_get_deltas(*args: Any, **kwargs: Any) -> Any: + nonlocal deltas_called + deltas_called = True + return await original_get_deltas(*args, **kwargs) + + async def tracking_get_summaries(*args: Any, **kwargs: Any) -> Any: + summary_page_sizes.append(kwargs.get("limit", 0)) + return await original_get_summaries(*args, **kwargs) + + sqlite_instance.get_scenario_results_async = tracking_get_scenario_results # type: ignore[assignment] + sqlite_instance.get_scenario_attack_result_deltas_async = tracking_get_deltas # type: ignore[assignment] + sqlite_instance.get_scenario_recount_summaries_async = tracking_get_summaries # type: ignore[assignment] + + service = ScenarioRunService(memory=sqlite_instance) + + # Fetch run summary list which triggers recount for runs needing SDK statistics + run_list = await service.list_runs_async(limit=10) + assert run_list is not None + matched_run = next((r for r in run_list.items if r.scenario_result_id == scenario_id), None) + assert matched_run is not None + assert matched_run.completed_attacks == 501 + + # Verified: full ScenarioResult was NOT loaded, deltas were NOT hydrated, summaries were paged! + assert scenario_results_called is False + assert deltas_called is False + assert len(summary_page_sizes) >= 2 # Paged in batches of 500 + + +@pytest.mark.asyncio +async def test_501_unplanned_group_attempts_parity_between_cold_and_warm_cache( + sqlite_instance: SQLiteMemory, +) -> None: + """ + Test that a scenario with 501 unplanned group attempts maintains exact parity + in group counts, overall counts, and outcomes between cold (max_nodes=1) and warm cache. + """ + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + base_time = datetime(2026, 1, 1, 0, 0, 0, tzinfo=UTC) + + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="UnplannedParityScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + # 500 successful unplanned attempts + all_results: list[AttackResult] = [ + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective=f"Unplanned objective {i}", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "unplanned_attack_group", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=i), + ) + for i in range(500) + ] + + # 501st unplanned attempt is an error with retries + all_results.append( + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective="Unplanned objective 500", + outcome=AttackOutcome.ERROR, + total_retries=5, + error_type="UnplannedExecutionError", + error_message="Execution failed", + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "unplanned_attack_group", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=500), + ) + ) + + await sqlite_instance.add_attack_results_to_memory_async(attack_results=all_results) + service = ScenarioRunService(memory=sqlite_instance) + + # 1. Cold hierarchy with max_nodes=1 + cold_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert cold_hier is not None + assert cold_hier.truncated is True + + # 2. Warm hierarchy + await service.get_run_progress_async(scenario_result_id=scenario_id, since=None, limit=600) + warm_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert warm_hier is not None + + # Parity assertions: + assert cold_hier.overall_counts.completed == 501 + assert cold_hier.overall_counts.succeeded == 500 + assert cold_hier.overall_counts.errors == 1 + assert cold_hier.overall_counts.completed == warm_hier.overall_counts.completed + assert cold_hier.overall_counts.succeeded == warm_hier.overall_counts.succeeded + assert cold_hier.overall_counts.errors == warm_hier.overall_counts.errors + + # Outcomes parity + assert cold_hier.overall_counts.outcomes is not None + assert warm_hier.overall_counts.outcomes is not None + assert cold_hier.overall_counts.outcomes.successes == warm_hier.overall_counts.outcomes.successes + assert cold_hier.overall_counts.outcomes.errors == warm_hier.overall_counts.outcomes.errors + + # Group counts parity for the unplanned group + assert len(cold_hier.groups) == len(warm_hier.groups) == 1 + cold_group = cold_hier.groups[0] + warm_group = warm_hier.groups[0] + assert cold_group.counts.completed == warm_group.counts.completed == 501 + assert cold_group.counts.succeeded == warm_group.counts.succeeded == 500 + assert cold_group.counts.errors == warm_group.counts.errors == 1 + assert cold_group.counts.retries == warm_group.counts.retries == 5 + + +@pytest.mark.asyncio +async def test_500_cyclic_results_with_outer_parent_at_501_parity(sqlite_instance: SQLiteMemory) -> None: + """ + Test that a 500-result cyclic component (A -> B -> C -> ... -> A) with an outer parent + at position 501 pointing to A discovers the outer parent on cold max_nodes=1, maintaining + exact parity with warm cache. + """ + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + base_time = datetime(2026, 1, 1, 0, 0, 0, tzinfo=UTC) + + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="CyclicScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + cyclic_ids = [str(uuid.uuid4()) for _ in range(500)] + outer_parent_id = str(uuid.uuid4()) + + results: list[AttackResult] = [ + AttackResult( + attack_result_id=cyclic_ids[i], + conversation_id=str(uuid.uuid4()), + objective=f"Cyclic objective {i}", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "attack_group", + "seed_group_id": f"seed_{i}", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + metadata={ + "child_attack_result_ids": [cyclic_ids[(i + 1) % 500]], + "completion_policy": "ALL_CHILDREN", + "executed_turns": 1, + }, + timestamp=base_time + timedelta(seconds=i), + ) + for i in range(500) + ] + + results.append( + AttackResult( + attack_result_id=outer_parent_id, + conversation_id=str(uuid.uuid4()), + objective="Outer parent objective", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "attack_group", + "seed_group_id": "outer_seed", + "result_role": AttackResultRole.ORCHESTRATION.value, + }, + metadata={ + "child_attack_result_ids": [cyclic_ids[-1]], + "completion_policy": "ALL_CHILDREN", + "executed_turns": 1, + }, + timestamp=base_time + timedelta(seconds=500), + ) + ) + + await sqlite_instance.add_attack_results_to_memory_async(attack_results=results) + service = ScenarioRunService(memory=sqlite_instance) + + cold_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert cold_hier is not None + assert cold_hier.truncated is True + + await service.get_run_progress_async(scenario_result_id=scenario_id, since=None, limit=600) + warm_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert warm_hier is not None + + cold_root_ids = [r.attack_result_id for g in cold_hier.groups for r in g.roots] + warm_root_ids = [r.attack_result_id for g in warm_hier.groups for r in g.roots] + assert outer_parent_id in cold_root_ids + assert outer_parent_id in warm_root_ids + assert cold_root_ids == warm_root_ids + + +@pytest.mark.asyncio +async def test_shared_child_with_1001_parents_bounded_parent_query(sqlite_instance: SQLiteMemory) -> None: + """ + Test that a child with 1,001 parents queries parents in bounded chunks with SQL LIMIT, + never hydrating all 1,001 rows in memory. + """ + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + base_time = datetime(2026, 1, 1, 0, 0, 0, tzinfo=UTC) + + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="ParentFaninScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + child_id = str(uuid.uuid4()) + child_result = AttackResult( + attack_result_id=child_id, + conversation_id=str(uuid.uuid4()), + objective="Shared child objective", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={"parent_collection": "attack_group", "seed_group_id": "seed_0"}, + timestamp=base_time, + ) + + parents = [ + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective=f"Parent objective {i}", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={"parent_collection": "attack_group", "seed_group_id": f"seed_{i + 1}"}, + metadata={ + "child_attack_result_ids": [child_id], + "completion_policy": "ALL_CHILDREN", + "executed_turns": 1, + }, + timestamp=base_time + timedelta(seconds=i + 1), + ) + for i in range(1001) + ] + + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[child_result, *parents]) + + deltas, has_more = await sqlite_instance.get_scenario_parent_deltas_for_child_ids_async( + scenario_result_id=scenario_id, + child_ids=[child_id], + limit=500, + ) + assert len(deltas) == 500 + assert has_more is True + + deltas_10, has_more_10 = await sqlite_instance.get_scenario_parent_deltas_for_child_ids_async( + scenario_result_id=scenario_id, + child_ids=[child_id], + limit=10, + ) + assert len(deltas_10) == 10 + assert has_more_10 is True + + +@pytest.mark.asyncio +async def test_raw_retrying_attacks_count_counts_all_attempts_not_units_parity( + sqlite_instance: SQLiteMemory, +) -> None: + """ + Test that raw_retrying_attacks_count counts all retrying attempts (even for the same logical unit), + maintaining parity between cold SQL aggregates and warm cache. + """ + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + base_time = datetime(2026, 1, 1, 0, 0, 0, tzinfo=UTC) + + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="RetryingAttemptsScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + attempt_1 = AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective="Objective 1", + outcome=AttackOutcome.ERROR, + total_retries=2, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "attack_group", + "seed_group_id": "seed_1", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time, + ) + attempt_2 = AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective="Objective 1", + outcome=AttackOutcome.SUCCESS, + total_retries=3, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "attack_group", + "seed_group_id": "seed_1", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=1), + ) + + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[attempt_1, attempt_2]) + service = ScenarioRunService(memory=sqlite_instance) + + cold_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=10) + assert cold_hier is not None + + await service.get_run_progress_async(scenario_result_id=scenario_id, since=None, limit=10) + warm_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=10) + assert warm_hier is not None + + assert cold_hier.diagnostics.attack_retries_count == 2 + assert warm_hier.diagnostics.attack_retries_count == 2 + assert cold_hier.diagnostics.total_retries == 6 + assert warm_hier.diagnostics.total_retries == 6 + assert cold_hier.diagnostics.failed_attacks_count == 1 + assert warm_hier.diagnostics.failed_attacks_count == 1 + + +@pytest.mark.asyncio +async def test_recount_with_saved_plan_preserves_counting_policy_and_unmatched_identifier_only_error( + sqlite_instance: SQLiteMemory, +) -> None: + """ + Test that legacy recount loads the saved plan from the scenario result header and preserves + the planned counting policy: planned=1, completed=1, succeeded=1, 100% success percentage, + with the unmatched identifier-only error reflected in producer attempts and raw errors. + """ + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + base_time = datetime(2026, 1, 1, 0, 0, 0, tzinfo=UTC) + + planned_obj = "Planned Objective 1" + planned_sha = hashlib.sha256(planned_obj.encode("utf-8")).hexdigest() + + saved_plan = ScenarioRunPlan( + scenario_registry_name="PlanRecountScenario", + atomic_groups=[ + ScenarioRunPlanAtomicGroup( + id="group_1", + display_group="display_1", + atomic_attack_name="attack_1", + technique_eval_hash="hash_tech_1", + seed_group_ids=["seed_1"], + ) + ], + seed_groups=[ + ScenarioRunPlanSeedGroup( + id="seed_1", + objective=planned_obj, + objective_sha256=planned_sha, + ) + ], + ) + + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="PlanRecountScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + metadata={SCENARIO_RUN_PLAN_METADATA_KEY: saved_plan.model_dump()}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + success_attempt = AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective=planned_obj, + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "attack_1", + "parent_eval_hash": "hash_tech_1", + "seed_group_id": "seed_1", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time, + ) + + error_identifier = AtomicAttackIdentifier( + class_name="MockAttack", + seed_identifiers=[SeedIdentifier(class_name="SeedPrompt", value="unmatched_seed_val")], + ) + error_attempt = AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective="Unmatched Error Objective", + atomic_attack_identifier=error_identifier, + outcome=AttackOutcome.ERROR, + total_retries=1, + error_type="UnmatchedError", + error_message="Failure", + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "attack_1", + "parent_eval_hash": "hash_tech_1", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=1), + ) + + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[success_attempt, error_attempt]) + service = ScenarioRunService(memory=sqlite_instance) + + run_list = await service.list_runs_async(limit=10) + assert run_list is not None + matched_run = next((r for r in run_list.items if r.scenario_result_id == scenario_id), None) + assert matched_run is not None + + assert matched_run.completed_attacks == 1 + assert matched_run.total_attacks == 1 + assert matched_run.objective_achieved_rate == 100 + + cold_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=10) + assert cold_hier is not None + + await service.get_run_progress_async(scenario_result_id=scenario_id, since=None, limit=10) + warm_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=10) + assert warm_hier is not None + + assert cold_hier.overall_counts.completed == 1 + assert cold_hier.overall_counts.planned == 1 + assert cold_hier.overall_counts.succeeded == 1 + assert cold_hier.overall_counts.errors == 0 + assert cold_hier.overall_counts.success_percentage == 100.0 + + assert warm_hier.overall_counts.completed == 1 + assert warm_hier.overall_counts.planned == 1 + assert warm_hier.overall_counts.succeeded == 1 + assert warm_hier.overall_counts.errors == 0 + assert warm_hier.overall_counts.success_percentage == 100.0 + + assert cold_hier.overall_counts.producer_counts.target_facing.attempts == 2 + assert cold_hier.overall_counts.producer_counts.target_facing.errors == 1 + assert warm_hier.overall_counts.producer_counts.target_facing.attempts == 2 + assert warm_hier.overall_counts.producer_counts.target_facing.errors == 1 + + assert cold_hier.diagnostics.failed_attacks_count == 1 + assert warm_hier.diagnostics.failed_attacks_count == 1 + + +@pytest.mark.asyncio +async def test_recount_updates_group_counts_and_raw_aggregates_parity(sqlite_instance: SQLiteMemory) -> None: + """ + Test that recount updates group_counts and raw aggregates so that cold hierarchy + matches warm hierarchy when SDK recount is triggered. + """ + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + base_time = datetime(2026, 1, 1, 0, 0, 0, tzinfo=UTC) + + alpha_obj = "Alpha Objective" + alpha_sha = hashlib.sha256(alpha_obj.encode("utf-8")).hexdigest() + + saved_plan = ScenarioRunPlan( + scenario_registry_name="MultiGroupRecountScenario", + atomic_groups=[ + ScenarioRunPlanAtomicGroup( + id="group_alpha", + display_group="display_alpha", + atomic_attack_name="attack_alpha", + technique_eval_hash="eval_alpha", + seed_group_ids=["seed_alpha_1"], + ) + ], + seed_groups=[ + ScenarioRunPlanSeedGroup( + id="seed_alpha_1", + objective=alpha_obj, + objective_sha256=alpha_sha, + ) + ], + ) + + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="MultiGroupRecountScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + metadata={SCENARIO_RUN_PLAN_METADATA_KEY: saved_plan.model_dump()}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + res_1 = AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective=alpha_obj, + outcome=AttackOutcome.SUCCESS, + total_retries=1, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "attack_alpha", + "parent_eval_hash": "eval_alpha", + "seed_group_id": "seed_alpha_1", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time, + ) + + ident = AtomicAttackIdentifier( + class_name="MockAttack", + seed_identifiers=[SeedIdentifier(class_name="SeedPrompt", value="alpha_seed_val")], + ) + res_2 = AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective=alpha_obj, + atomic_attack_identifier=ident, + outcome=AttackOutcome.SUCCESS, + total_retries=2, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "attack_alpha", + "parent_eval_hash": "eval_alpha", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=1), + ) + + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[res_1, res_2]) + service = ScenarioRunService(memory=sqlite_instance) + + cold_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=10) + assert cold_hier is not None + + await service.get_run_progress_async(scenario_result_id=scenario_id, since=None, limit=10) + warm_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=10) + assert warm_hier is not None + + assert len(cold_hier.groups) == len(warm_hier.groups) == 1 + cold_grp = cold_hier.groups[0] + warm_grp = warm_hier.groups[0] + assert cold_grp.counts.completed == warm_grp.counts.completed == 1 + assert cold_grp.counts.succeeded == warm_grp.counts.succeeded == 1 + assert cold_grp.counts.retries == warm_grp.counts.retries == 4 + assert cold_hier.diagnostics.total_retries == warm_hier.diagnostics.total_retries == 4 + assert cold_hier.diagnostics.attack_retries_count == warm_hier.diagnostics.attack_retries_count == 2 + + +@pytest.mark.asyncio +async def test_500_leaves_500_inner_parents_one_outer_root_parity(sqlite_instance: SQLiteMemory) -> None: + """Test 3-layer hierarchy (500 leaves -> 500 inner parents -> 1 outer root) parity.""" + scenario_id = str(uuid.uuid4()) + base_time = datetime(2025, 1, 1, 0, 0, 0, tzinfo=UTC) + + scen = make_scenario_result( + id=uuid.UUID(scenario_id), + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + leaf_ids = [str(uuid.uuid4()) for _ in range(500)] + inner_ids = [str(uuid.uuid4()) for _ in range(500)] + outer_root_id = str(uuid.uuid4()) + + results: list[AttackResult] = [ + AttackResult( + attack_result_id=leaf_ids[i], + conversation_id=str(uuid.uuid4()), + objective=f"Leaf objective {i}", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "attack_group", + "seed_group_id": f"seed_{i}", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=i), + ) + for i in range(500) + ] + + results.extend( + AttackResult( + attack_result_id=inner_ids[i], + conversation_id=str(uuid.uuid4()), + objective=f"Inner parent objective {i}", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "attack_group", + "seed_group_id": f"seed_{i}", + "result_role": AttackResultRole.ORCHESTRATION.value, + }, + metadata={ + "child_attack_result_ids": [leaf_ids[i]], + "completion_policy": "ALL_CHILDREN", + "executed_turns": 1, + }, + timestamp=base_time + timedelta(seconds=500 + i), + ) + for i in range(500) + ) + + results.append( + AttackResult( + attack_result_id=outer_root_id, + conversation_id=str(uuid.uuid4()), + objective="Outer root objective", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "attack_group", + "seed_group_id": "root_seed", + "result_role": AttackResultRole.ORCHESTRATION.value, + }, + metadata={ + "child_attack_result_ids": inner_ids, + "completion_policy": "ALL_CHILDREN", + "executed_turns": 1, + }, + timestamp=base_time + timedelta(seconds=1000), + ) + ) + + await sqlite_instance.add_attack_results_to_memory_async(attack_results=results) + service = ScenarioRunService(memory=sqlite_instance) + + cold_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert cold_hier is not None + + await service.get_run_progress_async(scenario_result_id=scenario_id, since=None, limit=1100) + warm_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert warm_hier is not None + + cold_root_ids = [r.attack_result_id for g in cold_hier.groups for r in g.roots] + warm_root_ids = [r.attack_result_id for g in warm_hier.groups for r in g.roots] + assert cold_root_ids == warm_root_ids == [outer_root_id] + + +@pytest.mark.asyncio +async def test_501_distinct_units_prefix_recount_preserves_exact_statistics_parity( + sqlite_instance: SQLiteMemory, +) -> None: + """Test that 501 distinct units preserve exact run-wide statistics under cold max_nodes=1.""" + scenario_id = str(uuid.uuid4()) + base_time = datetime(2025, 1, 1, 0, 0, 0, tzinfo=UTC) + + scen = make_scenario_result( + id=uuid.UUID(scenario_id), + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + results: list[AttackResult] = [ + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective=f"Objective {i}", + outcome=AttackOutcome.SUCCESS if i < 500 else AttackOutcome.ERROR, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "attack_group", + "seed_group_id": f"seed_{i}", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=i), + ) + for i in range(501) + ] + + await sqlite_instance.add_attack_results_to_memory_async(attack_results=results) + service = ScenarioRunService(memory=sqlite_instance) + + cold_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert cold_hier is not None + + await service.get_run_progress_async(scenario_result_id=scenario_id, since=None, limit=600) + warm_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert warm_hier is not None + + assert cold_hier.overall_counts.completed == warm_hier.overall_counts.completed == 501 + assert ( + cold_hier.overall_counts.producer_counts.target_facing.attempts + == warm_hier.overall_counts.producer_counts.target_facing.attempts + == 501 + ) + assert cold_hier.overall_counts.outcomes is not None + assert warm_hier.overall_counts.outcomes is not None + assert cold_hier.overall_counts.outcomes.errors == warm_hier.overall_counts.outcomes.errors == 1 + + +@pytest.mark.asyncio +async def test_501_attempts_of_one_unit_with_7_retries_diagnostic_pressure_parity( + sqlite_instance: SQLiteMemory, +) -> None: + """Test that 501 attempts of 1 unit with 7 retries yields total retry pressure of 507 cold and warm.""" + scenario_id = str(uuid.uuid4()) + base_time = datetime(2025, 1, 1, 0, 0, 0, tzinfo=UTC) + + scen = make_scenario_result( + id=uuid.UUID(scenario_id), + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + results: list[AttackResult] = [ + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective="Single shared objective", + outcome=AttackOutcome.SUCCESS, + total_retries=7 if i == 500 else 0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "attack_group", + "seed_group_id": "seed_0", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=i), + ) + for i in range(501) + ] + + await sqlite_instance.add_attack_results_to_memory_async(attack_results=results) + service = ScenarioRunService(memory=sqlite_instance) + + cold_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert cold_hier is not None + + await service.get_run_progress_async(scenario_result_id=scenario_id, since=None, limit=600) + warm_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert warm_hier is not None + + assert cold_hier.diagnostics.total_retries == 507 + assert warm_hier.diagnostics.total_retries == 507 + + +@pytest.mark.asyncio +async def test_two_omitted_retrying_tail_attempts_diagnostic_pressure_parity( + sqlite_instance: SQLiteMemory, +) -> None: + """Test that two omitted tail attempts recording 7 and 9 retries yield total retry pressure of 17.""" + scenario_id = str(uuid.uuid4()) + base_time = datetime(2025, 1, 1, 0, 0, 0, tzinfo=UTC) + + scen = make_scenario_result( + id=uuid.UUID(scenario_id), + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + results: list[AttackResult] = [ + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective=f"Prefix objective {i}", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "attack_group", + "seed_group_id": f"seed_{i}", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=i), + ) + for i in range(500) + ] + + results.append( + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective="Tail unit objective", + outcome=AttackOutcome.ERROR, + total_retries=7, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "tail_group", + "seed_group_id": "seed_tail", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=500), + ) + ) + results.append( + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective="Tail unit objective", + outcome=AttackOutcome.SUCCESS, + total_retries=9, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "tail_group", + "seed_group_id": "seed_tail", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=501), + ) + ) + + await sqlite_instance.add_attack_results_to_memory_async(attack_results=results) + service = ScenarioRunService(memory=sqlite_instance) + + cold_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert cold_hier is not None + + await service.get_run_progress_async(scenario_result_id=scenario_id, since=None, limit=600) + warm_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert warm_hier is not None + + assert cold_hier.diagnostics.total_retries == 17 + assert warm_hier.diagnostics.total_retries == 17 + + +@pytest.mark.asyncio +async def test_repeated_legacy_unit_500_errors_followed_by_recovery_parity( + sqlite_instance: SQLiteMemory, +) -> None: + """Test that 500 errors followed by 1 recovery for one legacy unit reports 1 success cold and warm.""" + scenario_id = str(uuid.uuid4()) + base_time = datetime(2025, 1, 1, 0, 0, 0, tzinfo=UTC) + + scen = make_scenario_result( + id=uuid.UUID(scenario_id), + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + results: list[AttackResult] = [ + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective="Repeated unit objective", + outcome=AttackOutcome.ERROR if i < 500 else AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "attack_group", + "seed_group_id": "seed_0", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=i), + ) + for i in range(501) + ] + + await sqlite_instance.add_attack_results_to_memory_async(attack_results=results) + service = ScenarioRunService(memory=sqlite_instance) + + cold_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert cold_hier is not None + + await service.get_run_progress_async(scenario_result_id=scenario_id, since=None, limit=600) + warm_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert warm_hier is not None + + assert cold_hier.overall_counts.completed == 1 + assert cold_hier.overall_counts.succeeded == 1 + assert warm_hier.overall_counts.completed == 1 + assert warm_hier.overall_counts.succeeded == 1 + + +@pytest.mark.asyncio +async def test_500_leaves_and_501_parents_leaf_0_not_emitted_as_root_parity( + sqlite_instance: SQLiteMemory, +) -> None: + """Test DAG where 500 leaves, 500 parents referencing leaf 1, and parent 501 referencing leaf 0 + does not emit leaf 0 as a root. + """ + scenario_id = str(uuid.uuid4()) + base_time = datetime(2025, 1, 1, 0, 0, 0, tzinfo=UTC) + + scen = make_scenario_result( + id=uuid.UUID(scenario_id), + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + # 500 leaf results (leaf 0 to leaf 499) + leaf_ids = [str(uuid.uuid4()) for _ in range(500)] + leaf_results = [ + AttackResult( + attack_result_id=leaf_ids[i], + conversation_id=str(uuid.uuid4()), + objective=f"Leaf objective {i}", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "attack_group", + "seed_group_id": f"seed_{i}", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=i), + ) + for i in range(500) + ] + + # 500 parents (parents 0 to 499): each references leaf 1, and between them cover leaves 1..499 + parent_results: list[AttackResult] = [] + for i in range(500): + children = [leaf_ids[1], leaf_ids[(i % 499) + 1]] + parent_results.append( + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective=f"Parent objective {i}", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "orch_group", + "result_role": AttackResultRole.ORCHESTRATION.value, + }, + metadata={"child_attack_result_ids": children}, + timestamp=base_time + timedelta(seconds=500 + i), + ) + ) + + # 1 final parent (parent 500) saved at index 1000: references leaf 0 + final_parent_id = str(uuid.uuid4()) + parent_results.append( + AttackResult( + attack_result_id=final_parent_id, + conversation_id=str(uuid.uuid4()), + objective="Parent of leaf 0", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "orch_group", + "result_role": AttackResultRole.ORCHESTRATION.value, + }, + metadata={"child_attack_result_ids": [leaf_ids[0]]}, + timestamp=base_time + timedelta(seconds=1000), + ) + ) + + await sqlite_instance.add_attack_results_to_memory_async(attack_results=leaf_results + parent_results) + service = ScenarioRunService(memory=sqlite_instance) + + cold_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert cold_hier is not None + + await service.get_run_progress_async(scenario_result_id=scenario_id, since=None, limit=1200) + warm_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert warm_hier is not None + + # Leaf 0 must NEVER be emitted as a root in cold or warm hierarchy + all_cold_roots = [r for g in cold_hier.groups for r in g.roots] + all_warm_roots = [r for g in warm_hier.groups for r in g.roots] + assert len(all_cold_roots) >= 1 + assert len(all_warm_roots) >= 1 + assert all(r.attack_result_id != leaf_ids[0] for r in all_cold_roots) + assert all(r.attack_result_id != leaf_ids[0] for r in all_warm_roots) + # The roots must be actual orchestration parents + assert all(r.result_role == AttackResultRole.ORCHESTRATION for r in all_cold_roots) + assert all(r.result_role == AttackResultRole.ORCHESTRATION for r in all_warm_roots) + + +@pytest.mark.asyncio +async def test_legacy_identity_500_errors_and_1_success_recounts_single_unit_parity( + sqlite_instance: SQLiteMemory, +) -> None: + """Test that 500 identifier-only errors followed by 1 attributed success for the same seed unifies into + 1 unit cold and warm. + """ + scenario_id = str(uuid.uuid4()) + base_time = datetime(2025, 1, 1, 0, 0, 0, tzinfo=UTC) + + scen = make_scenario_result( + id=uuid.UUID(scenario_id), + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + seed_group = AttackSeedGroup(seeds=[SeedObjective(value="Shared seed objective"), SeedPrompt(value="context")]) + + # 500 identifier-only ERROR attempts (no seed_group_id in attribution_data, triggers needs_sdk_statistics) + results: list[AttackResult] = [ + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective="Shared seed objective", + outcome=AttackOutcome.ERROR, + total_retries=0, + atomic_attack_identifier=AtomicAttackIdentifier.build( + attack_identifier=ComponentIdentifier(class_name="MockAttack", class_module="tests"), + seed_group=seed_group, + ), + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "legacy_group", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=i), + ) + for i in range(500) + ] + + # 1 explicitly attributed SUCCESS attempt for the same seed + results.append( + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective="Shared seed objective", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "legacy_group", + "seed_group_id": seed_group.logical_id, + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=500), + ) + ) + + await sqlite_instance.add_attack_results_to_memory_async(attack_results=results) + service = ScenarioRunService(memory=sqlite_instance) + + cold_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert cold_hier is not None + + await service.get_run_progress_async(scenario_result_id=scenario_id, since=None, limit=600) + warm_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert warm_hier is not None + + # Unified into 1 completed unit, 1 successful unit (100% success), and 500 retries / retry pressure + assert cold_hier.overall_counts.completed == 1 + assert warm_hier.overall_counts.completed == 1 + assert cold_hier.overall_counts.succeeded == 1 + assert warm_hier.overall_counts.succeeded == 1 + assert cold_hier.overall_counts.errors == 500 + assert warm_hier.overall_counts.errors == 500 + assert cold_hier.overall_counts.retries == 500 + assert warm_hier.overall_counts.retries == 500 + assert cold_hier.overall_counts.success_percentage == 100.0 + assert warm_hier.overall_counts.success_percentage == 100.0 + assert cold_hier.diagnostics.total_retries == 500 + assert warm_hier.diagnostics.total_retries == 500 + + +@pytest.mark.asyncio +async def test_saved_plan_two_seeds_shared_objective_different_context_500_attempts_parity( + sqlite_instance: SQLiteMemory, +) -> None: + """Test saved plan containing two seeds sharing an objective but different contexts, with 500 attempts.""" + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + base_time = datetime(2025, 1, 1, 0, 0, 0, tzinfo=UTC) + + obj_text = "Shared objective between seeds" + obj_hash = hashlib.sha256(obj_text.encode("utf-8")).hexdigest() + + sg_0 = AttackSeedGroup(seeds=[SeedObjective(value=obj_text), SeedPrompt(value="context_a")]) + sg_1 = AttackSeedGroup(seeds=[SeedObjective(value=obj_text), SeedPrompt(value="context_b")]) + + seed_groups = [ + ScenarioRunPlanSeedGroup( + id=sg_0.logical_id, + objective=obj_text, + objective_sha256=obj_hash, + ), + ScenarioRunPlanSeedGroup( + id=sg_1.logical_id, + objective=obj_text, + objective_sha256=obj_hash, + ), + ] + atomic_group = ScenarioRunPlanAtomicGroup( + id="group_1", + atomic_attack_name="attack_type_1", + display_group="Attack Type 1", + technique_eval_hash="hash_tech_1", + seed_group_ids=[sg_0.logical_id, sg_1.logical_id], + ) + plan = ScenarioRunPlan( + atomic_groups=[atomic_group], + seed_groups=seed_groups, + ) + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="SharedObjectiveScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + metadata={"run_plan": plan.model_dump(mode="json")}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + # 500 attempts across the two seeds (250 each), with distinct seed contexts in identifier + results: list[AttackResult] = [] + for i in range(500): + sg = sg_0 if i % 2 == 0 else sg_1 + results.append( + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective=obj_text, + outcome=AttackOutcome.ERROR, + total_retries=0, + atomic_attack_identifier=AtomicAttackIdentifier.build( + attack_identifier=ComponentIdentifier(class_name="attack_type_1", class_module="tests"), + seed_group=sg, + ), + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "attack_type_1", + "parent_eval_hash": "hash_tech_1", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=i), + ) + ) + + await sqlite_instance.add_attack_results_to_memory_async(attack_results=results) + service = ScenarioRunService(memory=sqlite_instance) + + cold_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert cold_hier is not None + + await service.get_run_progress_async(scenario_result_id=scenario_id, since=None, limit=600) + warm_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert warm_hier is not None + + # Both seeds resolved; 500 errors total; retry pressure: (250-1) + (250-1) = 498 + assert cold_hier.overall_counts.errors == 500 + assert warm_hier.overall_counts.errors == 500 + assert cold_hier.overall_counts.retries == 498 + assert warm_hier.overall_counts.retries == 498 + assert cold_hier.diagnostics.total_retries == 498 + assert warm_hier.diagnostics.total_retries == 498 + assert cold_hier.overall_counts.completed == 2 + assert warm_hier.overall_counts.completed == 2 + + +@pytest.mark.asyncio +async def test_11_parent_chain_boundary_guards_depth_cutoff_mssql_and_sqlite( + sqlite_instance: SQLiteMemory, +) -> None: + """ + Test that an 11-parent deep chain (exceeding depth boundary 10) above 500 leaves with max_nodes=1 + never promotes unverified candidates (at depth 10) to roots in cold response, and marks hierarchy truncated. + """ + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + base_time = datetime(2026, 1, 1, 0, 0, 0, tzinfo=UTC) + + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="ElevenParentChainScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + # 500 leaves + leaf_ids = [str(uuid.uuid4()) for _ in range(500)] + leaf_results = [ + AttackResult( + attack_result_id=lid, + conversation_id=str(uuid.uuid4()), + objective=f"Leaf {i}", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "leaves_group", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=i), + ) + for i, lid in enumerate(leaf_ids) + ] + + # Chain of 11 parents: + # Parent 10 references all 500 leaves + # Parent 9 references Parent 10 + # ... + # Parent 0 (True Root) references Parent 1 + parent_ids = [str(uuid.uuid4()) for _ in range(11)] + parents: list[AttackResult] = [] + for depth in range(11): + pid = parent_ids[depth] + children = leaf_ids if depth == 10 else [parent_ids[depth + 1]] + parents.append( + AttackResult( + attack_result_id=pid, + conversation_id=str(uuid.uuid4()), + objective=f"Parent depth {depth}", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "orch_group", + "result_role": AttackResultRole.ORCHESTRATION.value, + }, + metadata={"child_attack_result_ids": children}, + timestamp=base_time + timedelta(seconds=500 + depth), + ) + ) + + await sqlite_instance.add_attack_results_to_memory_async(attack_results=leaf_results + parents) + service = ScenarioRunService(memory=sqlite_instance) + + # Cold hierarchy with max_nodes=1 loads 1 leaf initially, then seeks parents up to depth 10. + # Depth 10 reaches parent_ids[1]. parent_ids[0] (True Root) is not loaded due to depth <= 10. + # Parent 1 is unverified. None of Parent 1, Parent 2, ..., Parent 10 should be promoted to root! + cold_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert cold_hier is not None + assert cold_hier.truncated is True + + cold_roots = [r for g in cold_hier.groups for r in g.roots] + cold_root_ids = {r.attack_result_id for r in cold_roots} + # Intermediate parents must NOT be emitted as fake roots + for pid in parent_ids[1:]: + assert pid not in cold_root_ids + + +@pytest.mark.asyncio +async def test_parent_budget_boundary_guards_unverified_parents( + sqlite_instance: SQLiteMemory, +) -> None: + """ + Test that at the 2,500-parent budget boundary, parents whose own parents were not checked + are marked unverified and never emitted as apparent roots. + """ + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + base_time = datetime(2026, 1, 1, 0, 0, 0, tzinfo=UTC) + + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="ParentBudgetBoundaryScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + # 1 leaf + leaf_id = str(uuid.uuid4()) + leaf_result = AttackResult( + attack_result_id=leaf_id, + conversation_id=str(uuid.uuid4()), + objective="Leaf", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "leaf_group", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time, + ) + + # 2,500 parents all referencing leaf_id + parent_budget = 2500 + parent_ids = [str(uuid.uuid4()) for _ in range(parent_budget)] + parents = [ + AttackResult( + attack_result_id=pid, + conversation_id=str(uuid.uuid4()), + objective=f"Parent {i}", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "orch_group", + "result_role": AttackResultRole.ORCHESTRATION.value, + }, + metadata={"child_attack_result_ids": [leaf_id]}, + timestamp=base_time + timedelta(seconds=1 + i), + ) + for i, pid in enumerate(parent_ids) + ] + + # 1 top-level outer root referencing the first 10 parents + outer_root_id = str(uuid.uuid4()) + outer_root = AttackResult( + attack_result_id=outer_root_id, + conversation_id=str(uuid.uuid4()), + objective="Outer Root", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "root_group", + "result_role": AttackResultRole.ORCHESTRATION.value, + }, + metadata={"child_attack_result_ids": parent_ids[:10]}, + timestamp=base_time + timedelta(seconds=3000), + ) + + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[leaf_result] + parents + [outer_root]) + service = ScenarioRunService(memory=sqlite_instance) + + cold_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert cold_hier is not None + assert cold_hier.truncated is True + + cold_roots = [r for g in cold_hier.groups for r in g.roots] + cold_root_ids = {r.attack_result_id for r in cold_roots} + # None of the 2,500 parents whose parents were never verified may be an apparent root + for pid in parent_ids[:10]: + assert pid not in cold_root_ids + + +@pytest.mark.asyncio +async def test_2001_legacy_attempts_recount_uses_lightweight_summaries_without_full_delta_hydration( + sqlite_instance: SQLiteMemory, +) -> None: + """ + Test that a large legacy run with 2,001 attempts recounts using get_scenario_recount_summaries_async + without hydrating full ScenarioAttackResultDelta objects or calling get_scenario_attack_result_deltas_async. + """ + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + base_time = datetime(2026, 1, 1, 0, 0, 0, tzinfo=UTC) + + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="LegacyRecount2001Scenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + # 2,001 legacy attempts without seed group IDs in attribution (needing SDK statistics) + count = 2001 + results = [ + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective=f"Objective {i}", + outcome=AttackOutcome.SUCCESS if i % 2 == 0 else AttackOutcome.FAILURE, + total_retries=1 if i % 10 == 0 else 0, + atomic_attack_identifier=AtomicAttackIdentifier( + class_name="MockAttack", + seed_identifiers=[SeedIdentifier(class_name="SeedPrompt", value=f"seed_val_{i}")], + ), + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": f"technique_{(i % 5)}", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=i), + ) + for i in range(count) + ] + await sqlite_instance.add_attack_results_to_memory_async(attack_results=results) + + deltas_call_count = 0 + summaries_page_limits: list[int] = [] + + original_get_deltas = sqlite_instance.get_scenario_attack_result_deltas_async + original_get_summaries = sqlite_instance.get_scenario_recount_summaries_async + + async def tracking_get_deltas(*args: Any, **kwargs: Any) -> Any: + nonlocal deltas_call_count + deltas_call_count += 1 + return await original_get_deltas(*args, **kwargs) + + async def tracking_get_summaries(*args: Any, **kwargs: Any) -> Any: + summaries_page_limits.append(kwargs.get("limit", 0)) + return await original_get_summaries(*args, **kwargs) + + sqlite_instance.get_scenario_attack_result_deltas_async = tracking_get_deltas # type: ignore[assignment] + sqlite_instance.get_scenario_recount_summaries_async = tracking_get_summaries # type: ignore[assignment] + + service = ScenarioRunService(memory=sqlite_instance) + run_list = await service.list_runs_async(limit=10) + assert run_list is not None + matched_run = next((r for r in run_list.items if r.scenario_result_id == scenario_id), None) + assert matched_run is not None + assert matched_run.completed_attacks == count + + # Recount MUST use lightweight summaries and NEVER call full delta hydration + assert deltas_call_count == 0 + # 2,001 items paged with limit 500 => pages: 500, 500, 500, 500, 500 + assert len(summaries_page_limits) == 5 + assert summaries_page_limits == [500, 500, 500, 500, 500] + + +@pytest.mark.asyncio +async def test_1002_results_500_leaves_501_parents_leaf_0_not_emitted_as_root_parity( + sqlite_instance: SQLiteMemory, +) -> None: + """Test exact 1,002-result SQLite case where 500 leaves, 501 parents referencing leaf 1 covering + leaves 1..499, and 1 final parent referencing leaf 0 does not emit leaf 0 as a root in cold hierarchy, + and agrees with warm hierarchy. + """ + scenario_id = str(uuid.uuid4()) + base_time = datetime(2025, 1, 1, 0, 0, 0, tzinfo=UTC) + + scen = make_scenario_result( + id=uuid.UUID(scenario_id), + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + # 500 leaf results (leaf 0 to leaf 499) + leaf_ids = [str(uuid.uuid4()) for _ in range(500)] + leaf_results = [ + AttackResult( + attack_result_id=leaf_ids[i], + conversation_id=str(uuid.uuid4()), + objective=f"Leaf objective {i}", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "attack_group", + "seed_group_id": f"seed_{i}", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=i), + ) + for i in range(500) + ] + + # 501 parents (parents 0 to 500): all reference leaf 1, and between them cover leaves 1..499 + parent_results: list[AttackResult] = [] + for i in range(501): + children = [leaf_ids[1], leaf_ids[(i % 499) + 1]] + parent_results.append( + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective=f"Parent objective {i}", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "orch_group", + "result_role": AttackResultRole.ORCHESTRATION.value, + }, + metadata={"child_attack_result_ids": children}, + timestamp=base_time + timedelta(seconds=500 + i), + ) + ) + + # 1 final parent (parent 501) saved at index 1001: references leaf 0 + final_parent_id = str(uuid.uuid4()) + parent_results.append( + AttackResult( + attack_result_id=final_parent_id, + conversation_id=str(uuid.uuid4()), + objective="Parent of leaf 0", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "orch_group", + "result_role": AttackResultRole.ORCHESTRATION.value, + }, + metadata={"child_attack_result_ids": [leaf_ids[0]]}, + timestamp=base_time + timedelta(seconds=1001), + ) + ) + + assert len(leaf_results) + len(parent_results) == 1002 + await sqlite_instance.add_attack_results_to_memory_async(attack_results=leaf_results + parent_results) + service = ScenarioRunService(memory=sqlite_instance) + + cold_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert cold_hier is not None + + await service.get_run_progress_async(scenario_result_id=scenario_id, since=None, limit=1200) + warm_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert warm_hier is not None + + # Leaf 0 must NEVER be emitted as a root in cold or warm hierarchy + all_cold_roots = [r for g in cold_hier.groups for r in g.roots] + all_warm_roots = [r for g in warm_hier.groups for r in g.roots] + assert len(all_cold_roots) >= 1 + assert len(all_warm_roots) >= 1 + assert all(r.attack_result_id != leaf_ids[0] for r in all_cold_roots) + assert all(r.attack_result_id != leaf_ids[0] for r in all_warm_roots) + # The roots must be actual orchestration parents + assert all(r.result_role == AttackResultRole.ORCHESTRATION for r in all_cold_roots) + assert all(r.result_role == AttackResultRole.ORCHESTRATION for r in all_warm_roots) + # Cold and warm root IDs must agree + cold_root_ids = {r.attack_result_id for r in all_cold_roots} + warm_root_ids = {r.attack_result_id for r in all_warm_roots} + assert cold_root_ids == warm_root_ids + + +@pytest.mark.asyncio +@pytest.mark.parametrize("leaf_490_parents_count", [1, 2, 3, 11]) +async def test_chunk_boundary_unprocessed_rows_500_leaves_parents_leaf_499_parity( + sqlite_instance: SQLiteMemory, + leaf_490_parents_count: int, +) -> None: + """Test 1,000, 1,001, 1,002, and 1,010-result SQLite chunk-boundary cases. + 500 leaves split into chunks of 490 and 10 in SQLite. + 490 parents reference leaves 0..489 (1 each). + leaf_490_parents_count parents reference leaf 490. + 9 parents reference leaves 491..499 (1 each). + Total parents: 490 + leaf_490_parents_count + 9. + When leaf_490_parents_count == 2 (total 501 parents): chunk 2 returns 11 parents. + Parent loop consumes 10 parents before reaching limit 500, leaving the 11th parent (leaf 499) unprocessed. + The database has_more flag must be True so leaf 499 is not emitted as a false root. + """ + scenario_id = str(uuid.uuid4()) + base_time = datetime(2025, 1, 1, 0, 0, 0, tzinfo=UTC) + + scen = make_scenario_result( + id=uuid.UUID(scenario_id), + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + # 500 leaf results (leaves 0 to 499) + leaf_ids = [str(uuid.uuid4()) for _ in range(500)] + leaf_results = [ + AttackResult( + attack_result_id=leaf_ids[i], + conversation_id=str(uuid.uuid4()), + objective=f"Leaf objective {i}", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "attack_group", + "seed_group_id": f"seed_{i}", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=i), + ) + for i in range(500) + ] + + parent_results: list[AttackResult] = [] + # 490 parents referencing leaves 0..489 (1 each) + for i in range(490): + parent_results.append( + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective=f"Parent {i}", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "orch_group", + "result_role": AttackResultRole.ORCHESTRATION.value, + }, + metadata={"child_attack_result_ids": [leaf_ids[i]]}, + timestamp=base_time + timedelta(seconds=500 + len(parent_results)), + ) + ) + + # leaf_490_parents_count parents referencing leaf 490 + for i in range(leaf_490_parents_count): + parent_results.append( + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective=f"Parent 490_{i}", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "orch_group", + "result_role": AttackResultRole.ORCHESTRATION.value, + }, + metadata={"child_attack_result_ids": [leaf_ids[490]]}, + timestamp=base_time + timedelta(seconds=500 + len(parent_results)), + ) + ) + + # 9 parents referencing leaves 491..499 (1 each) + for i in range(491, 500): + parent_results.append( + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective=f"Parent {i}", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "orch_group", + "result_role": AttackResultRole.ORCHESTRATION.value, + }, + metadata={"child_attack_result_ids": [leaf_ids[i]]}, + timestamp=base_time + timedelta(seconds=500 + len(parent_results)), + ) + ) + + await sqlite_instance.add_attack_results_to_memory_async(attack_results=leaf_results + parent_results) + service = ScenarioRunService(memory=sqlite_instance) + + cold_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert cold_hier is not None + + await service.get_run_progress_async(scenario_result_id=scenario_id, since=None, limit=1200) + warm_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert warm_hier is not None + + # Leaf 499 must NEVER be emitted as a root in cold or warm hierarchy + all_cold_roots = [r for g in cold_hier.groups for r in g.roots] + all_warm_roots = [r for g in warm_hier.groups for r in g.roots] + assert len(all_cold_roots) >= 1 + assert len(all_warm_roots) >= 1 + assert all(r.attack_result_id != leaf_ids[499] for r in all_cold_roots) + assert all(r.attack_result_id != leaf_ids[499] for r in all_warm_roots) + # The roots must be actual orchestration parents + assert all(r.result_role == AttackResultRole.ORCHESTRATION for r in all_cold_roots) + assert all(r.result_role == AttackResultRole.ORCHESTRATION for r in all_warm_roots) + # Cold and warm root IDs must agree + cold_root_ids = {r.attack_result_id for r in all_cold_roots} + warm_root_ids = {r.attack_result_id for r in all_warm_roots} + assert cold_root_ids == warm_root_ids + + +@pytest.mark.asyncio +@pytest.mark.parametrize("control_retries", [0, 1, 3]) +async def test_500_results_followed_by_retrying_attempt_diagnostic_count_parity( + sqlite_instance: SQLiteMemory, + control_retries: int, +) -> None: + """Test that 500 ordinary results followed by 1 result with retry events (or total_retries) + correctly preserves raw_retrying_attacks_count=1 and attack_retries_count=1 in cold hierarchy, + agreeing with warm hierarchy and SQL aggregate across controls (total_retries=0, 1, 3). + """ + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + base_time = datetime(2026, 1, 1, 0, 0, 0, tzinfo=UTC) + + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="RetryEventOnlyScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + # 500 ordinary results with total_retries=0 and no retry events + ordinary_results = [ + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective=f"Objective {i}", + outcome=AttackOutcome.SUCCESS, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "attack_group", + "seed_group_id": f"seed_{i}", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=i), + ) + for i in range(500) + ] + + # 1 legacy result at index 500 (attempt 501) with saved RetryEvent + retry_events = [ + RetryEvent( + attempt_number=1, + exception_type="RateLimitError", + timestamp=base_time + timedelta(seconds=500), + ) + ] + if control_retries > 1: + retry_events.append( + RetryEvent( + attempt_number=2, + exception_type="RateLimitError", + timestamp=base_time + timedelta(seconds=501), + ) + ) + + retrying_result = AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective="Objective 500", + outcome=AttackOutcome.SUCCESS, + total_retries=control_retries, + retry_events=retry_events, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "attack_group", + "seed_group_id": "seed_500", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=500), + ) + + all_results = ordinary_results + [retrying_result] + assert len(all_results) == 501 + await sqlite_instance.add_attack_results_to_memory_async(attack_results=all_results) + service = ScenarioRunService(memory=sqlite_instance) + + cold_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert cold_hier is not None + + await service.get_run_progress_async(scenario_result_id=scenario_id, since=None, limit=600) + warm_hier = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + assert warm_hier is not None + + # Parity assertions: SQL aggregate and recount must preserve raw_retrying_attacks_count=1 + assert cold_hier.diagnostics.attack_retries_count == 1 + assert warm_hier.diagnostics.attack_retries_count == 1 + assert cold_hier.diagnostics.attack_retries_count == warm_hier.diagnostics.attack_retries_count diff --git a/tests/unit/memory/test_azure_sql_memory.py b/tests/unit/memory/test_azure_sql_memory.py index e8a88bff5a..f07b056809 100644 --- a/tests/unit/memory/test_azure_sql_memory.py +++ b/tests/unit/memory/test_azure_sql_memory.py @@ -15,7 +15,7 @@ from pyrit.common.singleton import Singleton from pyrit.converter.base64_converter import Base64Converter -from pyrit.memory import AzureSQLMemory, EmbeddingDataEntry, PromptMemoryEntry +from pyrit.memory import AzureSQLMemory, EmbeddingDataEntry, MemoryInterface, PromptMemoryEntry from pyrit.memory.memory_interface import AttackResultKeysetCursor, _AttackResultQuery from pyrit.memory.memory_models import ScenarioResultEntry from pyrit.memory.storage.serializers import set_message_piece_sha256_async @@ -612,6 +612,46 @@ def test_get_conversation_stats_uses_one_latest_row_apply( assert "p2.converted_value_data_type AS last_data_type" in sql +def test_get_conversation_stats_chunks_within_bind_budget( + uninitialized_memory_interface: AzureSQLMemory, +) -> None: + """Conversation statistics query chunks parameters within _MAX_BIND_VARS and combines results.""" + session = MagicMock() + session.execute.side_effect = lambda sql, params: MagicMock( + fetchall=lambda: [(cid, 3, "preview", "text", "2026-10-16T12:00:00+00:00") for cid in params.values()] + ) + + max_vars = uninitialized_memory_interface._MAX_BIND_VARS + assert max_vars == 2000 + + # 1. At the limit: exactly 2,000 IDs -> 1 query, 2,000 parameters + at_limit_ids = [f"conv_{i}" for i in range(max_vars)] + with patch.object(uninitialized_memory_interface, "_get_session", return_value=session): + stats_at_limit = uninitialized_memory_interface._execute_get_conversation_stats(conversation_ids=at_limit_ids) + assert len(stats_at_limit) == max_vars + assert session.execute.call_count == 1 + params_at_limit = session.execute.call_args[0][1] + assert len(params_at_limit) == max_vars + assert len(params_at_limit) <= 2100 + + session.reset_mock() + + # 2. Above the limit: 2,101 IDs -> exactly 2 queries (2,000 and 101 parameters) + above_limit_ids = [f"conv_{i}" for i in range(2101)] + with patch.object(uninitialized_memory_interface, "_get_session", return_value=session): + stats_above_limit = uninitialized_memory_interface._execute_get_conversation_stats( + conversation_ids=above_limit_ids + ) + assert len(stats_above_limit) == 2101 + assert session.execute.call_count == 2 + first_call_params = session.execute.call_args_list[0][0][1] + second_call_params = session.execute.call_args_list[1][0][1] + assert len(first_call_params) == 2000 + assert len(second_call_params) == 101 + assert len(first_call_params) <= max_vars + assert len(second_call_params) <= max_vars + + def test_scenario_history_conditions_bind_or_within_label_and_registry_values( memory_interface: AzureSQLMemory, ) -> None: @@ -1169,6 +1209,120 @@ def test_scenario_role_match_condition_uses_binary_collation_and_datalength_guar assert "datalength(result_role) = datalength(N'target_facing')" in compiled_mssql +def test_scenario_attack_result_deltas_by_ids_statement_compiles_mssql( + memory_interface: AzureSQLMemory, +) -> None: + scenario_id = str(uuid.uuid4()) + child_ids = [str(uuid.uuid4()), str(uuid.uuid4())] + + stmt = memory_interface._build_scenario_attack_result_deltas_by_ids_statement( + scenario_result_id=scenario_id, + attack_result_ids=child_ids, + ) + compiled = stmt.compile(dialect=mssql.dialect()) + compiled_str = str(compiled) + + assert "[AttackResultEntries].attribution_parent_id =" in compiled_str + assert "[AttackResultEntries].id IN (" in compiled_str + assert "LEFT OUTER JOIN [ScoreEntries]" in compiled_str + assert "[AttackResultEntries].executed_turns" in compiled_str + assert "[AttackResultEntries].outcome_reason" in compiled_str + + +def test_scenario_attack_result_deltas_by_ids_chunks_over_2100_ids_mssql( + memory_interface: AzureSQLMemory, +) -> None: + scenario_id = str(uuid.uuid4()) + # Generate 2,500 distinct IDs - exceeds SQL Server 2,100 bind parameter limit + large_id_count = 2500 + child_ids = [str(uuid.uuid4()) for _ in range(large_id_count)] + + # 1. Statement generator bounds parameters strictly to _MAX_BIND_VARS - 10 + stmt = memory_interface._build_scenario_attack_result_deltas_by_ids_statement( + scenario_result_id=scenario_id, + attack_result_ids=child_ids, + ) + compiled = stmt.compile(dialect=mssql.dialect()) + # Total bind variables (child IDs + attribution_parent_id) must not exceed 2000 + assert len(compiled.params) <= 2000 + + # 2. Execution chunks IDs in batches and executes each statement within limits + session = MagicMock() + session.execute.return_value.mappings.return_value.all.return_value = [] + with patch.object(memory_interface, "_get_session", return_value=session): + deltas = memory_interface._execute_get_scenario_attack_result_deltas_by_ids( + scenario_result_id=scenario_id, + attack_result_ids=child_ids, + ) + + assert deltas == [] + # 2,500 IDs with chunk size 1990 requires exactly 2 batches + assert session.execute.call_count == 2 + for call in session.execute.call_args_list: + chunk_stmt = call.args[0] + chunk_compiled = chunk_stmt.compile(dialect=mssql.dialect()) + assert len(chunk_compiled.params) <= 2000 + + +def test_scenario_attack_result_deltas_by_ids_statement_type_hints() -> None: + """Test that type annotations on _build_scenario_attack_result_deltas_by_ids_statement evaluate cleanly.""" + import typing + + from sqlalchemy.sql import Select + + hints = typing.get_type_hints(MemoryInterface._build_scenario_attack_result_deltas_by_ids_statement) + assert hints.get("return") is Select + + +def test_scenario_parent_deltas_for_child_ids_bounds_parameters_with_large_exclusions_mssql( + memory_interface: AzureSQLMemory, +) -> None: + """Verify that _execute_get_scenario_parent_deltas_for_child_ids never exceeds SQL Server bind limits.""" + scenario_id = str(uuid.uuid4()) + # Case 1: 1,100 children and 1,100 exclusions - previously combined to ~2,900 params exceeding 2,100 MSSQL limit + child_ids = [str(uuid.uuid4()) for _ in range(1100)] + exclude_ids = [str(uuid.uuid4()) for _ in range(1100)] + + session = MagicMock() + session.execute.return_value.mappings.return_value.all.return_value = [] + with patch.object(memory_interface, "_get_session", return_value=session): + deltas, has_more = memory_interface._execute_get_scenario_parent_deltas_for_child_ids( + scenario_result_id=scenario_id, + child_ids=child_ids, + exclude_result_ids=exclude_ids, + ) + + assert deltas == [] + assert has_more is False + assert session.execute.call_count >= 1 + for call in session.execute.call_args_list: + stmt = call.args[0] + compiled = stmt.compile(dialect=mssql.dialect()) + assert len(compiled.params) <= memory_interface._MAX_BIND_VARS + assert len(compiled.params) <= 2000 + + # Case 2: 1,000 children and 2,500 exclusions (budget boundary) + child_ids_2 = [str(uuid.uuid4()) for _ in range(1000)] + exclude_ids_2 = [str(uuid.uuid4()) for _ in range(2500)] + + session.reset_mock() + with patch.object(memory_interface, "_get_session", return_value=session): + deltas, has_more = memory_interface._execute_get_scenario_parent_deltas_for_child_ids( + scenario_result_id=scenario_id, + child_ids=child_ids_2, + exclude_result_ids=exclude_ids_2, + ) + + assert deltas == [] + assert has_more is False + assert session.execute.call_count >= 1 + for call in session.execute.call_args_list: + stmt = call.args[0] + compiled = stmt.compile(dialect=mssql.dialect()) + assert len(compiled.params) <= memory_interface._MAX_BIND_VARS + assert len(compiled.params) <= 2000 + + def test_get_metadata_conditions_binds_special_character_keys( uninitialized_memory_interface: AzureSQLMemory, ) -> None: diff --git a/tests/unit/models/test_scenario_hierarchy_models.py b/tests/unit/models/test_scenario_hierarchy_models.py new file mode 100644 index 0000000000..7b3646a884 --- /dev/null +++ b/tests/unit/models/test_scenario_hierarchy_models.py @@ -0,0 +1,415 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from datetime import UTC, datetime + +import pytest +from pydantic import ValidationError + +from pyrit.models import ( + AttackOutcome, + AttackResultRole, + RetryEvent, + ScenarioAttackResultConversations, + ScenarioAttackTechniqueDetails, + ScenarioConversationMessagePieceProjection, + ScenarioConversationMessageProjection, + ScenarioConversationProjection, + ScenarioHierarchyDiagnostics, + ScenarioHierarchyGroup, + ScenarioHierarchyNode, + ScenarioHierarchyProjection, + ScenarioOrchestrationParentSummary, + ScenarioProgressCounts, + ScenarioProgressScore, + ScenarioRunState, +) +from pyrit.models.catalog.scenario import AttackErrorSummary, AttackRetrySummary + + +def test_scenario_orchestration_parent_summary_defaults() -> None: + now = datetime.now(UTC) + summary = ScenarioOrchestrationParentSummary( + attack_result_id="attack-123", + outcome=AttackOutcome.SUCCESS, + timestamp=now, + ) + + assert summary.attack_result_id == "attack-123" + assert summary.result_role == AttackResultRole.UNKNOWN + assert summary.outcome == AttackOutcome.SUCCESS + assert summary.attempt_index is None + assert summary.child_attack_result_ids == [] + assert summary.completion_policy is None + assert summary.retries == [] + assert summary.retries_truncated is False + + +def test_scenario_orchestration_parent_summary_validation() -> None: + now = datetime.now(UTC) + # attempt_index must be >= 1 when provided + with pytest.raises(ValidationError): + ScenarioOrchestrationParentSummary( + attack_result_id="attack-123", + outcome=AttackOutcome.SUCCESS, + timestamp=now, + attempt_index=0, + ) + + summary = ScenarioOrchestrationParentSummary( + attack_result_id="attack-123", + result_role=AttackResultRole.ORCHESTRATION, + outcome=AttackOutcome.ERROR, + attempt_index=2, + completion_policy="FIRST_SUCCESS", + child_attack_result_ids=["child-1", "child-2"], + error_type="TargetError", + error_message="Connection timed out", + timestamp=now, + total_retries=3, + retries=[RetryEvent(exception_type="RateLimit", attempt_number=1)], + retries_truncated=True, + ) + assert summary.attempt_index == 2 + assert summary.completion_policy == "FIRST_SUCCESS" + assert summary.child_attack_result_ids == ["child-1", "child-2"] + assert summary.retries_truncated is True + + +def test_scenario_conversation_projections() -> None: + now = datetime.now(UTC) + msg = ScenarioConversationMessageProjection( + role="assistant", + content="Response text", + sequence=1, + timestamp=now, + scores=[ + ScenarioProgressScore( + scorer_name="ObjectiveScorer", + score_type="float_scale", + status="complete", + score_value="0.9", + ) + ], + ) + assert msg.role == "assistant" + assert msg.content == "Response text" + assert len(msg.scores) == 1 + + conv = ScenarioConversationProjection( + conversation_id="conv-1", + conversation_type="main", + description="Main objective target conversation", + message_count=1, + messages=[msg], + messages_truncated=False, + ) + assert conv.conversation_id == "conv-1" + assert conv.conversation_type == "main" + assert len(conv.messages) == 1 + + +def test_scenario_hierarchy_node_recursive_tree() -> None: + now = datetime.now(UTC) + child1 = ScenarioHierarchyNode( + attack_result_id="child-1", + result_role=AttackResultRole.TARGET_FACING, + attempt_index=1, + atomic_attack_name="Crescendo", + atomic_group_id="group-1", + seed_group_id="seed-1", + objective="test objective", + outcome=AttackOutcome.FAILURE, + timestamp=now, + ) + child2 = ScenarioHierarchyNode( + attack_result_id="child-2", + result_role=AttackResultRole.TARGET_FACING, + attempt_index=2, + atomic_attack_name="PromptSending", + atomic_group_id="group-1", + seed_group_id="seed-1", + objective="test objective", + outcome=AttackOutcome.SUCCESS, + timestamp=now, + ) + parent = ScenarioHierarchyNode( + attack_result_id="parent-1", + result_role=AttackResultRole.ORCHESTRATION, + atomic_attack_name="SequentialAttack", + atomic_group_id="group-1", + seed_group_id="seed-1", + objective="test objective", + outcome=AttackOutcome.SUCCESS, + timestamp=now, + child_attack_result_ids=["child-1", "child-2", "unresolved-3"], + children=[child1, child2], + unresolved_child_attack_result_ids=["unresolved-3"], + orchestration=ScenarioOrchestrationParentSummary( + attack_result_id="parent-1", + result_role=AttackResultRole.ORCHESTRATION, + outcome=AttackOutcome.SUCCESS, + child_attack_result_ids=["child-1", "child-2", "unresolved-3"], + completion_policy="FIRST_SUCCESS", + timestamp=now, + ), + ) + + assert len(parent.children) == 2 + assert parent.children[0].attack_result_id == "child-1" + assert parent.children[0].attempt_index == 1 + assert parent.children[1].attack_result_id == "child-2" + assert parent.children[1].attempt_index == 2 + assert parent.unresolved_child_attack_result_ids == ["unresolved-3"] + assert parent.orchestration is not None + assert parent.orchestration.completion_policy == "FIRST_SUCCESS" + + # Round trip serialization + dumped = parent.model_dump(mode="json") + reloaded = ScenarioHierarchyNode.model_validate(dumped) + assert reloaded.attack_result_id == "parent-1" + assert len(reloaded.children) == 2 + assert reloaded.unresolved_child_attack_result_ids == ["unresolved-3"] + + +def test_scenario_hierarchy_diagnostics_bounding() -> None: + diagnostics = ScenarioHierarchyDiagnostics( + failed_attacks=[ + AttackErrorSummary( + atomic_attack_name="attack", + objective="obj", + error_type="Error", + error_message="msg", + total_retries=1, + ) + ], + failed_attacks_count=105, + failed_attacks_truncated=True, + attack_retries=[ + AttackRetrySummary( + attack_result_id="att-1", + atomic_attack_name="attack", + retries=[], + ) + ], + attack_retries_count=120, + attack_retries_truncated=True, + total_retries=250, + ) + + assert len(diagnostics.failed_attacks) == 1 + assert diagnostics.failed_attacks_count == 105 + assert diagnostics.failed_attacks_truncated is True + assert len(diagnostics.attack_retries) == 1 + assert diagnostics.attack_retries_count == 120 + assert diagnostics.attack_retries_truncated is True + assert diagnostics.total_retries == 250 + + +def test_scenario_hierarchy_projection_envelope() -> None: + now = datetime.now(UTC) + node = ScenarioHierarchyNode( + attack_result_id="node-1", + result_role=AttackResultRole.TARGET_FACING, + atomic_attack_name="attack", + atomic_group_id="group-1", + seed_group_id="seed-1", + objective="obj", + outcome=AttackOutcome.SUCCESS, + timestamp=now, + ) + group = ScenarioHierarchyGroup( + atomic_group_id="group-1", + atomic_attack_name="attack", + display_group="display", + technique=ScenarioAttackTechniqueDetails(component_name="Technique"), + roots=[node], + counts=ScenarioProgressCounts(completed=1, planned=1, succeeded=1, success_percentage=100, errors=0, retries=0), + ) + projection = ScenarioHierarchyProjection( + scenario_result_id="scen-123", + scenario_name="TestScenario", + scenario_version=1, + status=ScenarioRunState.COMPLETED, + groups=[group], + overall_counts=ScenarioProgressCounts( + completed=1, planned=1, succeeded=1, success_percentage=100, errors=0, retries=0 + ), + diagnostics=ScenarioHierarchyDiagnostics(), + ) + + assert projection.scenario_result_id == "scen-123" + assert len(projection.groups) == 1 + assert projection.groups[0].roots[0].attack_result_id == "node-1" + assert projection.overall_counts.completed == 1 + assert projection.truncated is False + + +def test_scenario_conversation_multimodal_pieces() -> None: + now = datetime.now(UTC) + pieces = [ + ScenarioConversationMessagePieceProjection( + data_type="text", + content="Look at this image:", + converted_value="Look at this image:", + ), + ScenarioConversationMessagePieceProjection( + data_type="image_path", + content="/path/to/img.png", + converted_value="/path/to/img.png", + ), + ] + msg = ScenarioConversationMessageProjection( + role="simulated_assistant", + content="Look at this image:", + sequence=1, + timestamp=now, + pieces=pieces, + ) + assert msg.role == "simulated_assistant" + assert len(msg.pieces) == 2 + assert msg.pieces[0].data_type == "text" + assert msg.pieces[1].data_type == "image_path" + + +def test_attack_retry_summary_bounded_fields() -> None: + summary = AttackRetrySummary( + attack_result_id="att-1", + atomic_attack_name="attack", + total_retries=250, + retries=[RetryEvent(exception_type="RateLimit", attempt_number=i) for i in range(100)], + retries_truncated=True, + ) + assert summary.total_retries == 250 + assert len(summary.retries) == 100 + assert summary.retries_truncated is True + + +def test_scenario_hierarchy_projection_truncated_flag() -> None: + projection = ScenarioHierarchyProjection( + scenario_result_id="scen-123", + scenario_name="TestScenario", + scenario_version=1, + status=ScenarioRunState.COMPLETED, + groups=[], + overall_counts=ScenarioProgressCounts( + completed=0, planned=0, succeeded=0, success_percentage=0, errors=0, retries=0 + ), + diagnostics=ScenarioHierarchyDiagnostics(), + truncated=True, + ) + assert projection.truncated is True + + +def test_scenario_hierarchy_node_conversations_count_and_truncated() -> None: + now = datetime.now(UTC) + node = ScenarioHierarchyNode( + attack_result_id="att-1", + atomic_attack_name="attack", + atomic_group_id="group-1", + seed_group_id="seed-1", + objective="obj", + outcome=AttackOutcome.SUCCESS, + timestamp=now, + conversations=[ + ScenarioConversationProjection(conversation_id="conv-1"), + ], + conversations_count=55, + conversations_truncated=True, + ) + assert node.conversations_count == 55 + assert node.conversations_truncated is True + assert len(node.conversations) == 1 + + +def test_scenario_attack_result_conversations_namedtuple() -> None: + conv = ScenarioConversationProjection(conversation_id="conv-1") + result = ScenarioAttackResultConversations( + conversations=[conv], + total_count=55, + truncated=True, + ) + assert result.conversations == [conv] + assert result.total_count == 55 + assert result.truncated is True + # Verify tuple unpacking works + c_list, count, trunc = result + assert c_list == [conv] + assert count == 55 + assert trunc is True + + +def test_scenario_hierarchy_node_child_references_bounding() -> None: + now = datetime.now(UTC) + child_ids = [f"child-{i}" for i in range(600)] + node = ScenarioHierarchyNode( + attack_result_id="att-1", + atomic_attack_name="attack", + atomic_group_id="group-1", + seed_group_id="seed-1", + objective="obj", + outcome=AttackOutcome.SUCCESS, + timestamp=now, + child_attack_result_ids=child_ids, + ) + assert ScenarioHierarchyNode.max_child_references == 500 + assert node.max_child_references == 500 + assert len(node.child_attack_result_ids) == 500 + assert node.child_references_truncated is True + assert node.child_attack_result_ids == child_ids[:500] + + +def test_scenario_orchestration_parent_summary_child_references_bounding() -> None: + now = datetime.now(UTC) + child_ids = [f"child-{i}" for i in range(600)] + summary = ScenarioOrchestrationParentSummary( + attack_result_id="parent-1", + outcome=AttackOutcome.SUCCESS, + timestamp=now, + child_attack_result_ids=child_ids, + ) + assert ScenarioOrchestrationParentSummary.max_child_references == 500 + assert summary.max_child_references == 500 + assert len(summary.child_attack_result_ids) == 500 + assert summary.child_references_truncated is True + assert summary.child_attack_result_ids == child_ids[:500] + + +def test_scenario_conversation_message_pieces_truncated_flag() -> None: + now = datetime.now(UTC) + pieces = [ + ScenarioConversationMessagePieceProjection( + data_type="text", + content=f"piece-{i}", + converted_value=f"piece-{i}", + ) + for i in range(100) + ] + msg = ScenarioConversationMessageProjection( + role="user", + content="test prompt", + sequence=0, + timestamp=now, + pieces=pieces, + pieces_truncated=True, + ) + assert len(msg.pieces) == 100 + assert msg.pieces_truncated is True + + +def test_scenario_hierarchy_node_unresolved_child_references_bounding() -> None: + now = datetime.now(UTC) + unresolved_ids = [f"unresolved-{i}" for i in range(600)] + node = ScenarioHierarchyNode( + attack_result_id="node-1", + atomic_attack_name="attack", + atomic_group_id="group-1", + seed_group_id="seed-1", + objective="obj", + outcome=AttackOutcome.SUCCESS, + timestamp=now, + unresolved_child_attack_result_ids=unresolved_ids, + ) + assert len(node.unresolved_child_attack_result_ids) == 500 + assert node.unresolved_child_references_truncated is True + assert node.unresolved_child_attack_result_ids == unresolved_ids[:500]