From 0cdc24a9edfec3e1e9f35ba29c5bd3c3e56fe6aa Mon Sep 17 00:00:00 2001 From: Nimit Jain Date: Tue, 13 Oct 2026 10:36:11 +0530 Subject: [PATCH 01/18] feat: Add bounded Scenario result hierarchy projections (#3145) --- pyrit/backend/routes/scenarios.py | 68 +- .../services/scenario_progress_read_model.py | 427 ++++++++++++ .../backend/services/scenario_run_service.py | 139 +++- pyrit/memory/memory_interface.py | 302 ++++++++ pyrit/models/__init__.py | 16 + pyrit/models/scenario_hierarchy.py | 177 +++++ .../backend/test_scenario_hierarchy_parity.py | 645 ++++++++++++++++++ tests/unit/memory/test_azure_sql_memory.py | 22 + .../models/test_scenario_hierarchy_models.py | 242 +++++++ 9 files changed, 2034 insertions(+), 4 deletions(-) create mode 100644 pyrit/models/scenario_hierarchy.py create mode 100644 tests/unit/backend/test_scenario_hierarchy_parity.py create mode 100644 tests/unit/models/test_scenario_hierarchy_models.py diff --git a/pyrit/backend/routes/scenarios.py b/pyrit/backend/routes/scenarios.py index 70966e8a40..1a8cf3a888 100644 --- a/pyrit/backend/routes/scenarios.py +++ b/pyrit/backend/routes/scenarios.py @@ -26,7 +26,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, @@ -335,6 +342,65 @@ 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, +) -> 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) + 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 0035e68511..0b9c4c1942 100644 --- a/pyrit/backend/services/scenario_progress_read_model.py +++ b/pyrit/backend/services/scenario_progress_read_model.py @@ -22,18 +22,27 @@ from pyrit.memory import AttackResultKeysetCursor from pyrit.memory.memory_interface import MemoryInterface from pyrit.models import ( + AttackOutcome, AttackResult, AttackResultMetadata, + AttackResultRole, AttackTechniqueIdentifier, ComponentIdentifier, ScenarioAtomicGroupProgress, ScenarioAttackResultDelta, ScenarioAttackTechniqueDetails, ScenarioComponentIdentity, + ScenarioConversationProjection, ScenarioDisplayGroupProgress, ScenarioExecutionUnit, + ScenarioHierarchyDiagnostics, + ScenarioHierarchyGroup, + ScenarioHierarchyNode, + ScenarioHierarchyProjection, ScenarioObjectiveScorer, ScenarioObjectiveScorerMetrics, + ScenarioOrchestrationParentSummary, + ScenarioProducerCounts, ScenarioProgressCounts, ScenarioProgressResult, ScenarioProgressSummary, @@ -42,6 +51,7 @@ ScenarioRunPlanAtomicGroup, ScenarioRunPlanGroupKind, ScenarioRunPlanSeedGroup, + ScenarioRunState, ScenarioScorerIdentity, ScenarioSeedGroupProgress, ScenarioTechniqueProgress, @@ -49,6 +59,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__) @@ -63,6 +74,9 @@ # applied by ``project_behavioral_identity``. _TECHNIQUE_SEEDS_CHILD = "technique_seeds" _TECHNIQUE_SEED_DISPLAY_PARAMS = ("value", "data_type") +_MAX_HIERARCHY_DEPTH = 10 +_MAX_DIAGNOSTIC_ITEMS = 100 +_MAX_RETRIES_PER_NODE = 100 @dataclass(frozen=True, slots=True) @@ -754,3 +768,416 @@ 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, + ) -> ScenarioHierarchyProjection: + """ + Assemble the scenario result hierarchy projection with cycle and depth 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. + + 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} + + referenced_child_ids: set[str] = set() + for result in snapshot.results: + for child_id in result.child_attack_result_ids: + referenced_child_ids.add(child_id) + + root_results = [r for r in snapshot.results if r.attack_result_id not in referenced_child_ids] + if not root_results and snapshot.results: + root_results = list(snapshot.results) + + def _build_node( + *, + attack_result_id: str, + parent_id: str | None, + child_index: int | None, + active_path: set[str], + depth: int, + ) -> ScenarioHierarchyNode | None: + 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 + + retries = result.retries + retries_truncated = False + if len(retries) > _MAX_RETRIES_PER_NODE: + retries = retries[:_MAX_RETRIES_PER_NODE] + retries_truncated = True + + 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=delta.attribution_data.get("outcome_reason") + or delta.attack_metadata.get("outcome_reason"), + attempt_index=result.attempt_index, + execution_time_ms=result.execution_time_ms, + executed_turns=int(delta.attack_metadata.get("executed_turns") or 0), + completion_policy=str(completion_policy) if completion_policy is not None else None, + child_attack_result_ids=list(result.child_attack_result_ids), + 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) + 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: + 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=int(delta.attack_metadata.get("conversation_message_count") or 0), + ) + ) + + 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=delta.attribution_data.get("outcome_reason") + or delta.attack_metadata.get("outcome_reason"), + execution_time_ms=result.execution_time_ms, + executed_turns=int(delta.attack_metadata.get("executed_turns") or 0), + 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=list(result.child_attack_result_ids), + children=children, + unresolved_child_attack_result_ids=unresolved_children, + conversations=conversations, + orchestration=orchestration, + ) + + roots_by_group: dict[str, list[ScenarioHierarchyNode]] = {} + for root_result in root_results: + 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) + + 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, []) + 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, + ) + 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] + completed = sum(1 for r in group_results if r.outcome != AttackOutcome.ERROR) + succeeded = sum(1 for r in group_results if r.outcome == AttackOutcome.SUCCESS) + errors = sum(1 for r in group_results if r.outcome == AttackOutcome.ERROR) + retries = sum(r.total_retries for r in group_results) + producer_counts = ScenarioProducerCounts( + attempts=len(group_results), + errors=errors, + ) + counts = ScenarioProgressCounts( + completed=completed, + planned=None, + succeeded=succeeded, + success_percentage=int((succeeded / completed) * 100) if completed > 0 else 0, + errors=errors, + retries=retries, + producer_counts=producer_counts, + ) + 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] = [] + total_retries_count = 0 + + for delta, result in zip(snapshot.deltas, snapshot.results, strict=True): + total_retries_count += result.total_retries + 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: + all_attack_retries.append( + AttackRetrySummary( + attack_result_id=result.attack_result_id, + atomic_attack_name=result.atomic_attack_name, + retries=result.retries, + ) + ) + + failed_attacks_count = len(all_failed_attacks) + failed_attacks_truncated = failed_attacks_count > _MAX_DIAGNOSTIC_ITEMS + failed_attacks = all_failed_attacks[:_MAX_DIAGNOSTIC_ITEMS] + + attack_retries_count = len(all_attack_retries) + attack_retries_truncated = 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=total_retries_count, + ) + + return ScenarioHierarchyProjection( + scenario_result_id=scenario_result_id, + scenario_name=scenario_name, + scenario_version=scenario_version, + status=status, + groups=groups, + overall_counts=snapshot.summary.overall, + diagnostics=diagnostics, + ) + + @staticmethod + def build_hierarchy_node( + *, + delta: ScenarioAttackResultDelta, + result: ScenarioProgressResult, + children_deltas: Sequence[ScenarioAttackResultDelta] = (), + children_results: Sequence[ScenarioProgressResult] = (), + conversations: Sequence[ScenarioConversationProjection] = (), + ) -> 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. + + Returns: + ScenarioHierarchyNode: The assembled node. + """ + 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 + + 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=delta.attribution_data.get("outcome_reason") + or delta.attack_metadata.get("outcome_reason"), + attempt_index=result.attempt_index, + execution_time_ms=result.execution_time_ms, + executed_turns=int(delta.attack_metadata.get("executed_turns") or 0), + completion_policy=str(completion_policy) if completion_policy is not None else None, + child_attack_result_ids=list(result.child_attack_result_ids), + 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] = [] + + for idx, c_id in enumerate(result.child_attack_result_ids): + if c_id == result.attack_result_id or c_id not in children_deltas_by_id: + unresolved_children.append(c_id) + else: + 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_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_delta.attribution_data.get("outcome_reason") + or c_delta.attack_metadata.get("outcome_reason"), + attempt_index=c_res.attempt_index, + execution_time_ms=c_res.execution_time_ms, + executed_turns=int(c_delta.attack_metadata.get("executed_turns") or 0), + completion_policy=str(c_policy) if c_policy is not None else None, + child_attack_result_ids=list(c_res.child_attack_result_ids), + 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, + ) + + 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_delta.attribution_data.get("outcome_reason") + or c_delta.attack_metadata.get("outcome_reason"), + execution_time_ms=c_res.execution_time_ms, + executed_turns=int(c_delta.attack_metadata.get("executed_turns") or 0), + 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=list(c_res.child_attack_result_ids), + children=[], + unresolved_child_attack_result_ids=[], + conversations=[], + orchestration=c_orchestration, + ) + ) + + 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=delta.attribution_data.get("outcome_reason") or delta.attack_metadata.get("outcome_reason"), + execution_time_ms=result.execution_time_ms, + executed_turns=int(delta.attack_metadata.get("executed_turns") or 0), + 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=list(result.child_attack_result_ids), + children=children, + unresolved_child_attack_result_ids=unresolved_children, + conversations=list(conversations), + orchestration=orchestration, + ) diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index da4b48311c..14321360ec 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -36,12 +36,14 @@ 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, ScenarioHistoryAggregate, ScenarioHistoryKeysetCursor, ScenarioHistoryRunRecord, @@ -53,8 +55,11 @@ AttackResult, ComponentIdentifier, ScenarioAttackResultDelta, + ScenarioHierarchyNode, + ScenarioHierarchyProjection, ScenarioIdentifier, ScenarioProgressHeader, + ScenarioProgressResult, ScenarioQueueEntry, ScenarioQueueSnapshot, ScenarioResult, @@ -187,7 +192,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. @@ -196,10 +207,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() @@ -2217,6 +2230,126 @@ 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, + ) -> ScenarioHierarchyProjection | None: + """ + Get the bounded result hierarchy projection for a scenario run. + + Args: + scenario_result_id: The persisted scenario run ID. + + 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, + ) + + return 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, + ) + + async def get_run_hierarchy_node_async( + self, + *, + scenario_result_id: str, + attack_result_id: str, + ) -> 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. + + 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] = [] + if 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=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, + ) + 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/memory_interface.py b/pyrit/memory/memory_interface.py index 738ef53e1f..6fc1086507 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -110,6 +110,8 @@ PromptDataType, RetryEvent, ScenarioAttackResultDelta, + ScenarioConversationMessageProjection, + ScenarioConversationProjection, ScenarioIdentifier, ScenarioProgressScore, ScenarioResult, @@ -6583,6 +6585,257 @@ def _execute_get_scenario_attack_result_deltas( ) return deltas, 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. + + Returns: + Select: The constructed select statement. + """ + scenario_uuid = uuid.UUID(scenario_result_id) + valid_uuids: list[uuid.UUID] = [] + for aid in attack_result_ids: + try: + valid_uuids.append(uuid.UUID(str(aid))) + except (ValueError, TypeError): + continue + + return ( + select( + AttackResultEntry.id, + AttackResultEntry.conversation_id, + AttackResultEntry.objective, + AttackResultEntry.objective_sha256, + AttackResultEntry.atomic_attack_identifier, + AttackResultEntry.outcome, + AttackResultEntry.execution_time_ms, + AttackResultEntry.timestamp, + AttackResultEntry.retry_events_json, + AttackResultEntry.total_retries, + AttackResultEntry.error_type, + AttackResultEntry.error_message, + AttackResultEntry.attribution_data, + AttackResultEntry.attack_metadata, + 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_(valid_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 a single query. + + Returns: + list[ScenarioAttackResultDelta]: The matched deltas. + """ + valid_uuids: list[uuid.UUID] = [] + for aid in attack_result_ids: + try: + valid_uuids.append(uuid.UUID(str(aid))) + except (ValueError, TypeError): + continue + + if not valid_uuids: + return [] + + statement = self._build_scenario_attack_result_deltas_by_ids_statement( + scenario_result_id=scenario_result_id, + attack_result_ids=attack_result_ids, + ) + with closing(self._get_session()) as session: + rows = session.execute(statement).mappings().all() + + deltas: list[ScenarioAttackResultDelta] = [] + for row in rows: + 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"], + ) + 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 deltas + + def _execute_get_attack_result_conversation_projections( + self, + *, + attack_result_id: str, + include_messages: bool = False, + limit_messages_per_conversation: int = 100, + ) -> list[ScenarioConversationProjection]: + """ + Read generic conversation projections for an attack result without full-run hydration. + + Returns: + list[ScenarioConversationProjection]: Generic read-only conversation projections. + """ + try: + ar_uuid = uuid.UUID(attack_result_id) + except (ValueError, TypeError): + return [] + + entries = self._query_entries( + AttackResultEntry, + conditions=AttackResultEntry.id == ar_uuid, + ) + if not entries: + return [] + + entry = entries[0] + conversations: list[ScenarioConversationProjection] = [] + conv_ids_to_fetch: list[str] = [] + + if entry.conversation_id: + conv_ids_to_fetch.append(entry.conversation_id) + + pruned_ids = entry.pruned_conversation_ids or [] + adversarial_ids = entry.adversarial_chat_conversation_ids or [] + preparation_ids = entry.preparation_conversation_ids or [] + + for c_id in pruned_ids + adversarial_ids + preparation_ids: + if c_id and c_id not in conv_ids_to_fetch: + conv_ids_to_fetch.append(c_id) + + if not conv_ids_to_fetch: + return [] + + 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: + pyrit_messages = self._execute_get_conversation_messages(conversation_id=entry.conversation_id) + truncated = len(pyrit_messages) > limit_messages_per_conversation + for msg in pyrit_messages[:limit_messages_per_conversation]: + content = msg.get_value() + piece = msg.message_pieces[0] if msg.message_pieces else None + seq = piece.sequence if piece else 0 + ts = piece.timestamp if piece else None + messages.append( + ScenarioConversationMessageProjection( + role=msg.api_role, + content=content, + sequence=seq, + timestamp=ts, + ) + ) + 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, + ) + ) + + related_convs = ( + [(c_id, "pruned", "Pruned conversation") for c_id in pruned_ids] + + [(c_id, "adversarial_chat", "Adversarial chat conversation") for c_id in adversarial_ids] + + [(c_id, "preparation", "Preparation conversation") for c_id in preparation_ids] + ) + for c_id, conv_type, desc in related_convs: + 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: + pyrit_messages = self._execute_get_conversation_messages(conversation_id=c_id) + truncated = len(pyrit_messages) > limit_messages_per_conversation + for msg in pyrit_messages[:limit_messages_per_conversation]: + content = msg.get_value() + piece = msg.message_pieces[0] if msg.message_pieces else None + seq = piece.sequence if piece else 0 + ts = piece.timestamp if piece else None + messages.append( + ScenarioConversationMessageProjection( + role=msg.api_role, + content=content, + sequence=seq, + timestamp=ts, + ) + ) + conversations.append( + ScenarioConversationProjection( + conversation_id=c_id, + conversation_type=conv_type, + description=desc, + message_count=msg_count, + messages=messages, + messages_truncated=truncated, + ) + ) + + return conversations + def _execute_get_scenario_results( self, *, @@ -10055,6 +10308,55 @@ async def get_scenario_attack_result_deltas_async( 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_attack_result_conversation_projections_async( + self, + *, + attack_result_id: str, + include_messages: bool = False, + limit_messages_per_conversation: int = 100, + ) -> list[ScenarioConversationProjection]: + """ + 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. + + Returns: + list[ScenarioConversationProjection]: Generic read-only conversation projections. + """ + 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, + ) + def get_scenario_results( self, *, diff --git a/pyrit/models/__init__.py b/pyrit/models/__init__.py index 0eaa0ae466..d5d5834fc8 100644 --- a/pyrit/models/__init__.py +++ b/pyrit/models/__init__.py @@ -147,6 +147,15 @@ ScenarioDatasetSizeEstimateKind, scenario_dataset_size_from_limit, ) + from pyrit.models.scenario_hierarchy import ( + ScenarioConversationMessageProjection, + ScenarioConversationProjection, + ScenarioHierarchyDiagnostics, + ScenarioHierarchyGroup, + ScenarioHierarchyNode, + ScenarioHierarchyProjection, + ScenarioOrchestrationParentSummary, + ) from pyrit.models.scenario_progress import ( SCENARIO_RUN_PLAN_METADATA_KEY, SCENARIO_RUN_PLAN_VERSION, @@ -436,6 +445,13 @@ "ScenarioDisplayGroupProgress": "pyrit.models.scenario_progress", "ScenarioExecutionStatistics": "pyrit.models.scenario_progress", "ScenarioExecutionUnit": "pyrit.models.scenario_progress", + "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/scenario_hierarchy.py b/pyrit/models/scenario_hierarchy.py new file mode 100644 index 0000000000..779a16f23d --- /dev/null +++ b/pyrit/models/scenario_hierarchy.py @@ -0,0 +1,177 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from pydantic import AwareDatetime, BaseModel, ConfigDict, Field + +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") + + 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) + 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 + + +class ScenarioConversationMessageProjection(BaseModel): + """ + Generic, read-only message projection for conversation views. + + Contains display text, conversation role, sequence order, 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) + + +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 = 0 + 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") + + 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) + + # Resolved children in stored order (bounded traversal) + children: list["ScenarioHierarchyNode"] = Field(default_factory=list) + + # Missing, malformed, cyclic, or cross-run references that could not be safely resolved + unresolved_child_attack_result_ids: list[str] = Field(default_factory=list) + + # Generic read-only conversations for this attack + conversations: list[ScenarioConversationProjection] = Field(default_factory=list) + + # Orchestration details if role is ORCHESTRATION or has children + orchestration: ScenarioOrchestrationParentSummary | None = None + + +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) + + +ScenarioHierarchyNode.model_rebuild() +ScenarioHierarchyGroup.model_rebuild() +ScenarioHierarchyProjection.model_rebuild() 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..b48ed9d4e9 --- /dev/null +++ b/tests/unit/backend/test_scenario_hierarchy_parity.py @@ -0,0 +1,645 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Unit and parity tests for bounded Scenario result hierarchy projections.""" + +import uuid +from datetime import UTC, datetime, timedelta +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import status +from fastapi.testclient import TestClient + +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 ( + AttackOutcome, + AttackResult, + AttackResultRole, + MessagePiece, + RetryEvent, + ScenarioAttackResultDelta, + ScenarioAttackTechniqueDetails, + ScenarioConversationMessageProjection, + ScenarioConversationProjection, + ScenarioHierarchyGroup, + ScenarioHierarchyNode, + ScenarioHierarchyProjection, + ScenarioProgressCounts, + ScenarioProgressResult, + ScenarioProgressSummary, + ScenarioRunPlan, + ScenarioRunPlanAtomicGroup, + ScenarioRunPlanSeedGroup, + ScenarioRunState, +) +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_snapshot( + *, + deltas: list[ScenarioAttackResultDelta], + plan: ScenarioRunPlan | 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")], + ) + 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=sum(r.total_retries for r in results), + ), + atomic_groups=[], + ) + return ScenarioProgressSnapshot( + deltas=tuple(deltas), + results=tuple(results), + summary=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, + 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)], + 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=[ + ScenarioConversationProjection( + conversation_id="conv_1", + conversation_type="main", + message_count=1, + messages=[ + ScenarioConversationMessageProjection( + role="user", + content="Test prompt", + sequence=1, + ) + ], + ) + ] + ) + + 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 len(conv_projections) == 1 + proj = conv_projections[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] diff --git a/tests/unit/memory/test_azure_sql_memory.py b/tests/unit/memory/test_azure_sql_memory.py index 78af9f916f..d1b917e536 100644 --- a/tests/unit/memory/test_azure_sql_memory.py +++ b/tests/unit/memory/test_azure_sql_memory.py @@ -982,3 +982,25 @@ def test_scenario_role_match_condition_uses_binary_collation_and_datalength_guar compiled_mssql = str(condition.compile(dialect=mssql.dialect(), compile_kwargs={"literal_binds": True})) assert "(result_role COLLATE Latin1_General_100_BIN2) = 'target_facing'" in compiled_mssql 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: + import uuid + + from sqlalchemy.dialects import mssql + + 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 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..5752220783 --- /dev/null +++ b/tests/unit/models/test_scenario_hierarchy_models.py @@ -0,0 +1,242 @@ +# 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, + ScenarioAttackTechniqueDetails, + 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 From a7367c06c3d4149d13c7cbec50dce6731589f2c4 Mon Sep 17 00:00:00 2001 From: Nimit Jain Date: Wed, 14 Oct 2026 10:54:42 +0530 Subject: [PATCH 02/18] fix: address all review feedback on bounded scenario hierarchy (#3167) --- pyrit/backend/routes/scenarios.py | 6 +- .../services/scenario_progress_read_model.py | 126 +++-- .../backend/services/scenario_run_service.py | 10 +- pyrit/memory/memory_interface.py | 256 ++++++---- pyrit/models/__init__.py | 2 + pyrit/models/catalog/scenario.py | 2 + pyrit/models/scenario_hierarchy.py | 16 +- pyrit/models/scenario_progress.py | 2 + .../backend/test_scenario_hierarchy_parity.py | 437 +++++++++++++++++- tests/unit/memory/test_azure_sql_memory.py | 41 +- .../models/test_scenario_hierarchy_models.py | 58 +++ 11 files changed, 828 insertions(+), 128 deletions(-) diff --git a/pyrit/backend/routes/scenarios.py b/pyrit/backend/routes/scenarios.py index 1a8cf3a888..b50b377981 100644 --- a/pyrit/backend/routes/scenarios.py +++ b/pyrit/backend/routes/scenarios.py @@ -352,6 +352,7 @@ async def get_scenario_run_progress( # pyrit-async-suffix-exempt 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. @@ -361,7 +362,10 @@ async def get_scenario_run_hierarchy( # pyrit-async-suffix-exempt 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) + 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, diff --git a/pyrit/backend/services/scenario_progress_read_model.py b/pyrit/backend/services/scenario_progress_read_model.py index 0b9c4c1942..4eb99cbd47 100644 --- a/pyrit/backend/services/scenario_progress_read_model.py +++ b/pyrit/backend/services/scenario_progress_read_model.py @@ -42,7 +42,6 @@ ScenarioObjectiveScorer, ScenarioObjectiveScorerMetrics, ScenarioOrchestrationParentSummary, - ScenarioProducerCounts, ScenarioProgressCounts, ScenarioProgressResult, ScenarioProgressSummary, @@ -75,6 +74,7 @@ _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 @@ -777,9 +777,10 @@ def build_scenario_hierarchy( scenario_version: int, status: ScenarioRunState, snapshot: ScenarioProgressSnapshot, + max_nodes: int = _MAX_HIERARCHY_NODES, ) -> ScenarioHierarchyProjection: """ - Assemble the scenario result hierarchy projection with cycle and depth guards. + Assemble the scenario result hierarchy projection with cycle, depth, and node budget guards. Args: scenario_result_id: Persisted scenario run ID. @@ -787,6 +788,7 @@ def build_scenario_hierarchy( 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. Returns: ScenarioHierarchyProjection: Bounded hierarchy projection with diagnostics. @@ -800,8 +802,10 @@ def build_scenario_hierarchy( referenced_child_ids.add(child_id) root_results = [r for r in snapshot.results if r.attack_result_id not in referenced_child_ids] - if not root_results and snapshot.results: - root_results = list(snapshot.results) + + visited_ids: set[str] = set() + total_nodes_count = 0 + hierarchy_truncated = False def _build_node( *, @@ -811,17 +815,32 @@ def _build_node( 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 = delta.executed_turns or int(delta.attack_metadata.get("executed_turns") or 0) + outcome_reason = ( + delta.outcome_reason + or delta.attribution_data.get("outcome_reason") + or delta.attack_metadata.get("outcome_reason") + ) + orchestration: ScenarioOrchestrationParentSummary | None = None if result.result_role == AttackResultRole.ORCHESTRATION or result.child_attack_result_ids: completion_policy = delta.attack_metadata.get("completion_policy") @@ -829,11 +848,10 @@ def _build_node( attack_result_id=result.attack_result_id, result_role=result.result_role, outcome=result.outcome, - outcome_reason=delta.attribution_data.get("outcome_reason") - or delta.attack_metadata.get("outcome_reason"), + outcome_reason=outcome_reason, attempt_index=result.attempt_index, execution_time_ms=result.execution_time_ms, - executed_turns=int(delta.attack_metadata.get("executed_turns") or 0), + executed_turns=executed_turns, completion_policy=str(completion_policy) if completion_policy is not None else None, child_attack_result_ids=list(result.child_attack_result_ids), error_type=result.error_type, @@ -851,6 +869,9 @@ def _build_node( 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, @@ -887,10 +908,9 @@ def _build_node( seed_group_id=result.seed_group_id, objective=delta.objective, outcome=result.outcome, - outcome_reason=delta.attribution_data.get("outcome_reason") - or delta.attack_metadata.get("outcome_reason"), + outcome_reason=outcome_reason, execution_time_ms=result.execution_time_ms, - executed_turns=int(delta.attack_metadata.get("executed_turns") or 0), + executed_turns=executed_turns, timestamp=result.timestamp, score=result.score, error_type=result.error_type, @@ -917,6 +937,24 @@ def _build_node( 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) + unreached = [r for r in snapshot.results if r.attack_result_id not in visited_ids] + while unreached: + unreached.sort(key=lambda r: r.timestamp) + next_root = unreached[0] + 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) + unreached = [r for r in snapshot.results if r.attack_result_id not in visited_ids] + for g_roots in roots_by_group.values(): g_roots.sort(key=lambda n: n.timestamp) @@ -934,6 +972,7 @@ def _build_node( errors=ag.errors, retries=ag.retries, producer_counts=ag.producer_counts, + outcomes=ag.outcomes, ) groups.append( ScenarioHierarchyGroup( @@ -950,22 +989,16 @@ def _build_node( 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] - completed = sum(1 for r in group_results if r.outcome != AttackOutcome.ERROR) - succeeded = sum(1 for r in group_results if r.outcome == AttackOutcome.SUCCESS) - errors = sum(1 for r in group_results if r.outcome == AttackOutcome.ERROR) - retries = sum(r.total_retries for r in group_results) - producer_counts = ScenarioProducerCounts( - attempts=len(group_results), - errors=errors, - ) - counts = ScenarioProgressCounts( - completed=completed, + 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) + + counts = count_execution_units( + units=list(attempts_by_unit.keys()), + attempts_by_unit=attempts_by_unit, planned=None, - succeeded=succeeded, - success_percentage=int((succeeded / completed) * 100) if completed > 0 else 0, - errors=errors, - retries=retries, - producer_counts=producer_counts, + producer_attempts=group_results, ) groups.append( ScenarioHierarchyGroup( @@ -995,11 +1028,18 @@ def _build_node( ) ) if result.retries: + retries_list = result.retries + 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, - retries=result.retries, + total_retries=result.total_retries, + retries=retries_list, + retries_truncated=retries_truncated, ) ) @@ -1029,6 +1069,7 @@ def _build_node( groups=groups, overall_counts=snapshot.summary.overall, diagnostics=diagnostics, + truncated=hierarchy_truncated, ) @staticmethod @@ -1062,6 +1103,13 @@ def build_hierarchy_node( retries = retries[:_MAX_RETRIES_PER_NODE] retries_truncated = True + executed_turns = delta.executed_turns or int(delta.attack_metadata.get("executed_turns") or 0) + outcome_reason = ( + delta.outcome_reason + or delta.attribution_data.get("outcome_reason") + or delta.attack_metadata.get("outcome_reason") + ) + orchestration: ScenarioOrchestrationParentSummary | None = None if result.result_role == AttackResultRole.ORCHESTRATION or result.child_attack_result_ids: completion_policy = delta.attack_metadata.get("completion_policy") @@ -1069,11 +1117,10 @@ def build_hierarchy_node( attack_result_id=result.attack_result_id, result_role=result.result_role, outcome=result.outcome, - outcome_reason=delta.attribution_data.get("outcome_reason") - or delta.attack_metadata.get("outcome_reason"), + outcome_reason=outcome_reason, attempt_index=result.attempt_index, execution_time_ms=result.execution_time_ms, - executed_turns=int(delta.attack_metadata.get("executed_turns") or 0), + executed_turns=executed_turns, completion_policy=str(completion_policy) if completion_policy is not None else None, child_attack_result_ids=list(result.child_attack_result_ids), error_type=result.error_type, @@ -1099,6 +1146,13 @@ def build_hierarchy_node( c_retries = c_retries[:_MAX_RETRIES_PER_NODE] c_retries_truncated = True + c_executed_turns = c_delta.executed_turns or int(c_delta.attack_metadata.get("executed_turns") or 0) + c_outcome_reason = ( + c_delta.outcome_reason + or c_delta.attribution_data.get("outcome_reason") + or c_delta.attack_metadata.get("outcome_reason") + ) + 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") @@ -1106,11 +1160,10 @@ def build_hierarchy_node( attack_result_id=c_res.attack_result_id, result_role=c_res.result_role, outcome=c_res.outcome, - outcome_reason=c_delta.attribution_data.get("outcome_reason") - or c_delta.attack_metadata.get("outcome_reason"), + outcome_reason=c_outcome_reason, attempt_index=c_res.attempt_index, execution_time_ms=c_res.execution_time_ms, - executed_turns=int(c_delta.attack_metadata.get("executed_turns") or 0), + executed_turns=c_executed_turns, completion_policy=str(c_policy) if c_policy is not None else None, child_attack_result_ids=list(c_res.child_attack_result_ids), error_type=c_res.error_type, @@ -1134,10 +1187,9 @@ def build_hierarchy_node( seed_group_id=c_res.seed_group_id, objective=c_delta.objective, outcome=c_res.outcome, - outcome_reason=c_delta.attribution_data.get("outcome_reason") - or c_delta.attack_metadata.get("outcome_reason"), + outcome_reason=c_outcome_reason, execution_time_ms=c_res.execution_time_ms, - executed_turns=int(c_delta.attack_metadata.get("executed_turns") or 0), + executed_turns=c_executed_turns, timestamp=c_res.timestamp, score=c_res.score, error_type=c_res.error_type, @@ -1165,9 +1217,9 @@ def build_hierarchy_node( seed_group_id=result.seed_group_id, objective=delta.objective, outcome=result.outcome, - outcome_reason=delta.attribution_data.get("outcome_reason") or delta.attack_metadata.get("outcome_reason"), + outcome_reason=outcome_reason, execution_time_ms=result.execution_time_ms, - executed_turns=int(delta.attack_metadata.get("executed_turns") or 0), + executed_turns=executed_turns, timestamp=result.timestamp, score=result.score, error_type=result.error_type, diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index 14321360ec..840bc76d26 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -1476,11 +1476,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, ) ) @@ -2234,12 +2239,14 @@ 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. @@ -2282,6 +2289,7 @@ async def get_run_hierarchy_async( scenario_version=header_result.scenario_version, status=header_result.scenario_run_state, snapshot=progress_snapshot, + max_nodes=max_nodes, ) async def get_run_hierarchy_node_async( diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 6fc1086507..3d057fd5f4 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -110,6 +110,7 @@ PromptDataType, RetryEvent, ScenarioAttackResultDelta, + ScenarioConversationMessagePieceProjection, ScenarioConversationMessageProjection, ScenarioConversationProjection, ScenarioIdentifier, @@ -6513,6 +6514,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, @@ -6573,6 +6576,8 @@ def _execute_get_scenario_attack_result_deltas( atomic_attack_identifier=atomic_identifier, outcome=AttackOutcome(row["outcome"]), execution_time_ms=row["execution_time_ms"], + executed_turns=row["executed_turns"] or 0, + outcome_reason=row["outcome_reason"], timestamp=row["timestamp"], retry_events=retry_events, total_retries=row["total_retries"] or 0, @@ -6596,18 +6601,27 @@ def _build_scenario_attack_result_deltas_by_ids_statement( 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. """ scenario_uuid = uuid.UUID(scenario_result_id) valid_uuids: list[uuid.UUID] = [] + seen: set[uuid.UUID] = set() for aid in attack_result_ids: try: - valid_uuids.append(uuid.UUID(str(aid))) + 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, @@ -6617,6 +6631,8 @@ def _build_scenario_attack_result_deltas_by_ids_statement( 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, @@ -6638,7 +6654,7 @@ def _build_scenario_attack_result_deltas_by_ids_statement( .where( and_( AttackResultEntry.attribution_parent_id == scenario_uuid, - AttackResultEntry.id.in_(valid_uuids), + AttackResultEntry.id.in_(chunk_uuids), ) ) ) @@ -6650,73 +6666,164 @@ def _execute_get_scenario_attack_result_deltas_by_ids( attack_result_ids: Sequence[str], ) -> list[ScenarioAttackResultDelta]: """ - Return lightweight scenario-linked result deltas for specific attack results in a single query. + 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. """ valid_uuids: list[uuid.UUID] = [] + seen: set[uuid.UUID] = set() for aid in attack_result_ids: try: - valid_uuids.append(uuid.UUID(str(aid))) + 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 [] - statement = self._build_scenario_attack_result_deltas_by_ids_statement( - scenario_result_id=scenario_result_id, - attack_result_ids=attack_result_ids, - ) + max_batch_size = max(1, self._MAX_BIND_VARS - 10) + deltas: list[ScenarioAttackResultDelta] = [] with closing(self._get_session()) as session: - rows = session.execute(statement).mappings().all() + 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() + for row in rows: + 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"], + ) + 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"], + executed_turns=row["executed_turns"] or 0, + 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 {}, + score=score, + ) + ) + return deltas - deltas: list[ScenarioAttackResultDelta] = [] - for row in rows: - 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 + def _fetch_bounded_conversation_messages( + self, + *, + conversation_id: str, + limit: int, + ) -> 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) ) - 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"], + seqs = [row[0] for row in session.execute(seq_stmt).all()] + if not seqs: + return [], False + + truncated = len(seqs) > limit + target_seqs = seqs[:limit] + + pieces_stmt = ( + select(PromptMemoryEntry) + .where( + and_( + PromptMemoryEntry.conversation_id == str(conversation_id), + PromptMemoryEntry.sequence.in_(target_seqs), + ) ) - 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, + .order_by( + PromptMemoryEntry.sequence.asc(), + PromptMemoryEntry.timestamp.asc(), + PromptMemoryEntry.id.asc(), ) ) - return deltas + entries = session.execute(pieces_stmt).scalars().all() + if not entries: + return [], False + + message_pieces = [e.get_message_piece() for e in entries] + grouped_messages = group_conversation_message_pieces_by_sequence(message_pieces=message_pieces) + + messages: list[ScenarioConversationMessageProjection] = [] + for msg in grouped_messages: + 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" + seq = first_piece.sequence if first_piece else 0 + 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, + ) + ) + + return messages, truncated def _execute_get_attack_result_conversation_projections( self, @@ -6724,10 +6831,13 @@ def _execute_get_attack_result_conversation_projections( attack_result_id: str, include_messages: bool = False, limit_messages_per_conversation: int = 100, + max_conversations: int = 50, ) -> list[ScenarioConversationProjection]: """ Read generic conversation projections for an attack result without full-run hydration. + Applies bounded conversation budget and SQL LIMIT queries. + Returns: list[ScenarioConversationProjection]: Generic read-only conversation projections. """ @@ -6758,6 +6868,7 @@ def _execute_get_attack_result_conversation_projections( if c_id and c_id not in conv_ids_to_fetch: conv_ids_to_fetch.append(c_id) + conv_ids_to_fetch = conv_ids_to_fetch[:max_conversations] if not conv_ids_to_fetch: return [] @@ -6769,21 +6880,10 @@ def _execute_get_attack_result_conversation_projections( messages: list[ScenarioConversationMessageProjection] = [] truncated = False if include_messages and entry.conversation_id: - pyrit_messages = self._execute_get_conversation_messages(conversation_id=entry.conversation_id) - truncated = len(pyrit_messages) > limit_messages_per_conversation - for msg in pyrit_messages[:limit_messages_per_conversation]: - content = msg.get_value() - piece = msg.message_pieces[0] if msg.message_pieces else None - seq = piece.sequence if piece else 0 - ts = piece.timestamp if piece else None - messages.append( - ScenarioConversationMessageProjection( - role=msg.api_role, - content=content, - sequence=seq, - timestamp=ts, - ) - ) + messages, truncated = self._fetch_bounded_conversation_messages( + conversation_id=entry.conversation_id, + limit=limit_messages_per_conversation, + ) conversations.append( ScenarioConversationProjection( conversation_id=entry.conversation_id, @@ -6795,11 +6895,13 @@ def _execute_get_attack_result_conversation_projections( ) ) + 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] + [(c_id, "adversarial_chat", "Adversarial chat conversation") for c_id in adversarial_ids] + [(c_id, "preparation", "Preparation conversation") for c_id in preparation_ids] - ) + )[:remaining_budget] + for c_id, conv_type, desc in related_convs: if not c_id: continue @@ -6808,21 +6910,10 @@ def _execute_get_attack_result_conversation_projections( messages = [] truncated = False if include_messages: - pyrit_messages = self._execute_get_conversation_messages(conversation_id=c_id) - truncated = len(pyrit_messages) > limit_messages_per_conversation - for msg in pyrit_messages[:limit_messages_per_conversation]: - content = msg.get_value() - piece = msg.message_pieces[0] if msg.message_pieces else None - seq = piece.sequence if piece else 0 - ts = piece.timestamp if piece else None - messages.append( - ScenarioConversationMessageProjection( - role=msg.api_role, - content=content, - sequence=seq, - timestamp=ts, - ) - ) + messages, truncated = self._fetch_bounded_conversation_messages( + conversation_id=c_id, + limit=limit_messages_per_conversation, + ) conversations.append( ScenarioConversationProjection( conversation_id=c_id, @@ -10338,6 +10429,7 @@ async def get_attack_result_conversation_projections_async( attack_result_id: str, include_messages: bool = False, limit_messages_per_conversation: int = 100, + max_conversations: int = 50, ) -> list[ScenarioConversationProjection]: """ Read generic conversation projections for an attack result without full-run hydration. @@ -10346,6 +10438,7 @@ async def get_attack_result_conversation_projections_async( 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. Returns: list[ScenarioConversationProjection]: Generic read-only conversation projections. @@ -10355,6 +10448,7 @@ async def get_attack_result_conversation_projections_async( attack_result_id=attack_result_id, include_messages=include_messages, limit_messages_per_conversation=limit_messages_per_conversation, + max_conversations=max_conversations, ) def get_scenario_results( diff --git a/pyrit/models/__init__.py b/pyrit/models/__init__.py index d5d5834fc8..d4161cbfae 100644 --- a/pyrit/models/__init__.py +++ b/pyrit/models/__init__.py @@ -148,6 +148,7 @@ scenario_dataset_size_from_limit, ) from pyrit.models.scenario_hierarchy import ( + ScenarioConversationMessagePieceProjection, ScenarioConversationMessageProjection, ScenarioConversationProjection, ScenarioHierarchyDiagnostics, @@ -445,6 +446,7 @@ "ScenarioDisplayGroupProgress": "pyrit.models.scenario_progress", "ScenarioExecutionStatistics": "pyrit.models.scenario_progress", "ScenarioExecutionUnit": "pyrit.models.scenario_progress", + "ScenarioConversationMessagePieceProjection": "pyrit.models.scenario_hierarchy", "ScenarioConversationMessageProjection": "pyrit.models.scenario_hierarchy", "ScenarioConversationProjection": "pyrit.models.scenario_hierarchy", "ScenarioHierarchyDiagnostics": "pyrit.models.scenario_hierarchy", diff --git a/pyrit/models/catalog/scenario.py b/pyrit/models/catalog/scenario.py index 27dc39082c..6018db7a3f 100644 --- a/pyrit/models/catalog/scenario.py +++ b/pyrit/models/catalog/scenario.py @@ -535,9 +535,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 index 779a16f23d..9d4fbda5b4 100644 --- a/pyrit/models/scenario_hierarchy.py +++ b/pyrit/models/scenario_hierarchy.py @@ -41,11 +41,23 @@ class ScenarioOrchestrationParentSummary(BaseModel): retries_truncated: bool = False +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, and linked score indicators. + Contains display text, conversation role, sequence order, multimodal pieces, and linked score indicators. """ model_config = ConfigDict(extra="ignore") @@ -55,6 +67,7 @@ class ScenarioConversationMessageProjection(BaseModel): sequence: int timestamp: AwareDatetime | None = None scores: list[ScenarioProgressScore] = Field(default_factory=list) + pieces: list[ScenarioConversationMessagePieceProjection] = Field(default_factory=list) class ScenarioConversationProjection(BaseModel): @@ -170,6 +183,7 @@ class ScenarioHierarchyProjection(BaseModel): groups: list[ScenarioHierarchyGroup] = Field(default_factory=list) overall_counts: ScenarioProgressCounts diagnostics: ScenarioHierarchyDiagnostics = Field(default_factory=ScenarioHierarchyDiagnostics) + truncated: bool = False ScenarioHierarchyNode.model_rebuild() diff --git a/pyrit/models/scenario_progress.py b/pyrit/models/scenario_progress.py index 76fa287231..33f35ab174 100644 --- a/pyrit/models/scenario_progress.py +++ b/pyrit/models/scenario_progress.py @@ -398,6 +398,8 @@ class ScenarioAttackResultDelta(BaseModel): total_retries: int = 0 error_type: str | None = None error_message: str | None = None + executed_turns: int = 0 + outcome_reason: str | None = None attribution_data: dict[str, Any] = Field(default_factory=dict) attack_metadata: dict[str, Any] = Field(default_factory=dict) score: ScenarioProgressScore | None = None diff --git a/tests/unit/backend/test_scenario_hierarchy_parity.py b/tests/unit/backend/test_scenario_hierarchy_parity.py index b48ed9d4e9..053ecd13be 100644 --- a/tests/unit/backend/test_scenario_hierarchy_parity.py +++ b/tests/unit/backend/test_scenario_hierarchy_parity.py @@ -5,12 +5,14 @@ import uuid from datetime import UTC, datetime, timedelta +from pathlib import Path 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, @@ -25,6 +27,7 @@ AttackResultRole, MessagePiece, RetryEvent, + ScenarioAtomicGroupProgress, ScenarioAttackResultDelta, ScenarioAttackTechniqueDetails, ScenarioConversationMessageProjection, @@ -95,6 +98,8 @@ 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: @@ -134,7 +139,7 @@ def _make_test_snapshot( ], seed_groups=[ScenarioRunPlanSeedGroup(id="seed_1", objective_sha256="hash123", objective="Test objective")], ) - summary = ScenarioProgressSummary( + effective_summary = summary or ScenarioProgressSummary( overall=ScenarioProgressCounts( completed=len([r for r in results if r.outcome != AttackOutcome.ERROR]), planned=1, @@ -142,12 +147,12 @@ def _make_test_snapshot( errors=len([r for r in results if r.outcome == AttackOutcome.ERROR]), retries=sum(r.total_retries for r in results), ), - atomic_groups=[], + atomic_groups=atomic_groups or [], ) return ScenarioProgressSnapshot( deltas=tuple(deltas), results=tuple(results), - summary=summary, + summary=effective_summary, plan=effective_plan, ) @@ -643,3 +648,429 @@ async def test_sqlite_memory_batched_reads_and_projections_parity(sqlite_instanc 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", + attempt_index=1, + 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", + attempt_index=2, + 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 len(projs) == 1 + proj = projs[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." diff --git a/tests/unit/memory/test_azure_sql_memory.py b/tests/unit/memory/test_azure_sql_memory.py index d1b917e536..471568439e 100644 --- a/tests/unit/memory/test_azure_sql_memory.py +++ b/tests/unit/memory/test_azure_sql_memory.py @@ -987,10 +987,6 @@ def test_scenario_role_match_condition_uses_binary_collation_and_datalength_guar def test_scenario_attack_result_deltas_by_ids_statement_compiles_mssql( memory_interface: AzureSQLMemory, ) -> None: - import uuid - - from sqlalchemy.dialects import mssql - scenario_id = str(uuid.uuid4()) child_ids = [str(uuid.uuid4()), str(uuid.uuid4())] @@ -1004,3 +1000,40 @@ def test_scenario_attack_result_deltas_by_ids_statement_compiles_mssql( 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 diff --git a/tests/unit/models/test_scenario_hierarchy_models.py b/tests/unit/models/test_scenario_hierarchy_models.py index 5752220783..04a123daea 100644 --- a/tests/unit/models/test_scenario_hierarchy_models.py +++ b/tests/unit/models/test_scenario_hierarchy_models.py @@ -11,6 +11,7 @@ AttackResultRole, RetryEvent, ScenarioAttackTechniqueDetails, + ScenarioConversationMessagePieceProjection, ScenarioConversationMessageProjection, ScenarioConversationProjection, ScenarioHierarchyDiagnostics, @@ -240,3 +241,60 @@ def test_scenario_hierarchy_projection_envelope() -> None: 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 From 5fde2ca7ab6cb3159f944d727f1abea5cd668967 Mon Sep 17 00:00:00 2001 From: Nimit Jain Date: Wed, 14 Oct 2026 17:22:42 +0530 Subject: [PATCH 03/18] fix: resolve import NameError, O(N^2) loop, retry consistency, and conversation clipping (#3167) --- .../services/scenario_progress_read_model.py | 82 +++++-- pyrit/memory/memory_interface.py | 60 +++-- pyrit/models/__init__.py | 2 + pyrit/models/scenario_hierarchy.py | 12 + .../targets/test_a2a_target_integration.py | 2 +- .../backend/test_scenario_hierarchy_parity.py | 211 +++++++++++++++--- tests/unit/memory/test_azure_sql_memory.py | 12 +- .../models/test_scenario_hierarchy_models.py | 39 ++++ 8 files changed, 351 insertions(+), 69 deletions(-) diff --git a/pyrit/backend/services/scenario_progress_read_model.py b/pyrit/backend/services/scenario_progress_read_model.py index 4eb99cbd47..9681893b4a 100644 --- a/pyrit/backend/services/scenario_progress_read_model.py +++ b/pyrit/backend/services/scenario_progress_read_model.py @@ -29,6 +29,7 @@ AttackTechniqueIdentifier, ComponentIdentifier, ScenarioAtomicGroupProgress, + ScenarioAttackResultConversations, ScenarioAttackResultDelta, ScenarioAttackTechniqueDetails, ScenarioComponentIdentity, @@ -922,11 +923,16 @@ def _build_node( children=children, unresolved_child_attack_result_ids=unresolved_children, conversations=conversations, + conversations_count=len(conversations), + conversations_truncated=False, 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, @@ -937,23 +943,30 @@ def _build_node( 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) - unreached = [r for r in snapshot.results if r.attack_result_id not in visited_ids] - while unreached: - unreached.sort(key=lambda r: r.timestamp) - next_root = unreached[0] - 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) + # Handle unreached/disconnected components (including detached cycles) in a single linear pass + if total_nodes_count < max_nodes: unreached = [r for r in snapshot.results if r.attack_result_id not in visited_ids] + unreached.sort(key=lambda r: r.timestamp) + 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 snapshot.results): + hierarchy_truncated = True for g_roots in roots_by_group.values(): g_roots.sort(key=lambda n: n.timestamp) @@ -1013,10 +1026,8 @@ def _build_node( all_failed_attacks: list[AttackErrorSummary] = [] all_attack_retries: list[AttackRetrySummary] = [] - total_retries_count = 0 for delta, result in zip(snapshot.deltas, snapshot.results, strict=True): - total_retries_count += result.total_retries if result.outcome == AttackOutcome.ERROR: all_failed_attacks.append( AttackErrorSummary( @@ -1058,7 +1069,7 @@ def _build_node( attack_retries=attack_retries, attack_retries_count=attack_retries_count, attack_retries_truncated=attack_retries_truncated, - total_retries=total_retries_count, + total_retries=snapshot.summary.overall.retries, ) return ScenarioHierarchyProjection( @@ -1079,7 +1090,11 @@ def build_hierarchy_node( result: ScenarioProgressResult, children_deltas: Sequence[ScenarioAttackResultDelta] = (), children_results: Sequence[ScenarioProgressResult] = (), - conversations: Sequence[ScenarioConversationProjection] = (), + conversations: Sequence[ScenarioConversationProjection] + | ScenarioAttackResultConversations + | tuple[Sequence[ScenarioConversationProjection], int, bool] = (), + conversations_count: int | None = None, + conversations_truncated: bool = False, ) -> ScenarioHierarchyNode: """ Assemble a single hierarchy node with its immediate children and generic conversations. @@ -1090,10 +1105,31 @@ def build_hierarchy_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. 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 + elif isinstance(conversations, tuple) and len(conversations) == 3: + first, second, third = conversations + if isinstance(first, Sequence): + conv_list = [c for c in first if isinstance(c, ScenarioConversationProjection)] + if conversations_count is None and isinstance(second, int): + conversations_count = second + if isinstance(third, bool): + conversations_truncated = conversations_truncated or third + else: + conv_list = list(conversations) + + resolved_conversations_count: int = 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} @@ -1201,6 +1237,8 @@ def build_hierarchy_node( children=[], unresolved_child_attack_result_ids=[], conversations=[], + conversations_count=0, + conversations_truncated=False, orchestration=c_orchestration, ) ) @@ -1230,6 +1268,8 @@ def build_hierarchy_node( child_attack_result_ids=list(result.child_attack_result_ids), children=children, unresolved_child_attack_result_ids=unresolved_children, - conversations=list(conversations), + conversations=conv_list, + conversations_count=resolved_conversations_count, + conversations_truncated=conversations_truncated, orchestration=orchestration, ) diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 3d057fd5f4..d1173bdfb4 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -22,6 +22,7 @@ from sqlalchemy import ( MetaData, + Select, String, Unicode, and_, @@ -109,6 +110,7 @@ Observation, PromptDataType, RetryEvent, + ScenarioAttackResultConversations, ScenarioAttackResultDelta, ScenarioConversationMessagePieceProjection, ScenarioConversationMessageProjection, @@ -6832,45 +6834,47 @@ def _execute_get_attack_result_conversation_projections( include_messages: bool = False, limit_messages_per_conversation: int = 100, max_conversations: int = 50, - ) -> list[ScenarioConversationProjection]: + ) -> ScenarioAttackResultConversations: """ Read generic conversation projections for an attack result without full-run hydration. Applies bounded conversation budget and SQL LIMIT queries. Returns: - list[ScenarioConversationProjection]: Generic read-only conversation projections. + ScenarioAttackResultConversations: Container with projections, total count, and truncation flag. """ try: ar_uuid = uuid.UUID(attack_result_id) except (ValueError, TypeError): - return [] + return ScenarioAttackResultConversations(conversations=[], total_count=0, truncated=False) entries = self._query_entries( AttackResultEntry, conditions=AttackResultEntry.id == ar_uuid, ) if not entries: - return [] + return ScenarioAttackResultConversations(conversations=[], total_count=0, truncated=False) entry = entries[0] conversations: list[ScenarioConversationProjection] = [] - conv_ids_to_fetch: list[str] = [] + all_conv_ids: list[str] = [] if entry.conversation_id: - conv_ids_to_fetch.append(entry.conversation_id) + all_conv_ids.append(entry.conversation_id) pruned_ids = entry.pruned_conversation_ids or [] adversarial_ids = entry.adversarial_chat_conversation_ids or [] preparation_ids = entry.preparation_conversation_ids or [] for c_id in pruned_ids + adversarial_ids + preparation_ids: - if c_id and c_id not in conv_ids_to_fetch: - conv_ids_to_fetch.append(c_id) + if c_id and c_id not in all_conv_ids: + all_conv_ids.append(c_id) - conv_ids_to_fetch = conv_ids_to_fetch[:max_conversations] + 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 [] + return ScenarioAttackResultConversations(conversations=[], total_count=0, truncated=False) stats_map = self._execute_get_conversation_stats(conversation_ids=conv_ids_to_fetch) @@ -6897,12 +6901,28 @@ def _execute_get_attack_result_conversation_projections( 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] - + [(c_id, "adversarial_chat", "Adversarial chat conversation") for c_id in adversarial_ids] - + [(c_id, "preparation", "Preparation conversation") for c_id in preparation_ids] - )[:remaining_budget] + [(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 related_convs: + for c_id, conv_type, desc in bounded_related: if not c_id: continue stats = stats_map.get(c_id) @@ -6925,7 +6945,11 @@ def _execute_get_attack_result_conversation_projections( ) ) - return conversations + return ScenarioAttackResultConversations( + conversations=conversations, + total_count=total_conversations, + truncated=conversations_truncated, + ) def _execute_get_scenario_results( self, @@ -10430,7 +10454,7 @@ async def get_attack_result_conversation_projections_async( include_messages: bool = False, limit_messages_per_conversation: int = 100, max_conversations: int = 50, - ) -> list[ScenarioConversationProjection]: + ) -> ScenarioAttackResultConversations: """ Read generic conversation projections for an attack result without full-run hydration. @@ -10441,7 +10465,7 @@ async def get_attack_result_conversation_projections_async( max_conversations: Maximum number of conversations to retrieve for this attack result. Returns: - list[ScenarioConversationProjection]: Generic read-only conversation projections. + ScenarioAttackResultConversations: Container with projections, total count, and truncation flag. """ return await self._run_database_operation_async( self._execute_get_attack_result_conversation_projections, diff --git a/pyrit/models/__init__.py b/pyrit/models/__init__.py index d4161cbfae..2a7fed5d9a 100644 --- a/pyrit/models/__init__.py +++ b/pyrit/models/__init__.py @@ -148,6 +148,7 @@ scenario_dataset_size_from_limit, ) from pyrit.models.scenario_hierarchy import ( + ScenarioAttackResultConversations, ScenarioConversationMessagePieceProjection, ScenarioConversationMessageProjection, ScenarioConversationProjection, @@ -447,6 +448,7 @@ "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", diff --git a/pyrit/models/scenario_hierarchy.py b/pyrit/models/scenario_hierarchy.py index 9d4fbda5b4..47291c7fbd 100644 --- a/pyrit/models/scenario_hierarchy.py +++ b/pyrit/models/scenario_hierarchy.py @@ -1,6 +1,8 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +from typing import NamedTuple + from pydantic import AwareDatetime, BaseModel, ConfigDict, Field from pyrit.models.catalog.scenario import AttackErrorSummary, AttackRetrySummary @@ -130,11 +132,21 @@ class ScenarioHierarchyNode(BaseModel): # Generic read-only conversations for this attack conversations: list[ScenarioConversationProjection] = Field(default_factory=list) + conversations_count: int = 0 + conversations_truncated: bool = False # Orchestration details if role is ORCHESTRATION or has children orchestration: ScenarioOrchestrationParentSummary | None = None +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.""" 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 index 053ecd13be..4051837455 100644 --- a/tests/unit/backend/test_scenario_hierarchy_parity.py +++ b/tests/unit/backend/test_scenario_hierarchy_parity.py @@ -25,9 +25,12 @@ AttackOutcome, AttackResult, AttackResultRole, + ConversationReference, + ConversationType, MessagePiece, RetryEvent, ScenarioAtomicGroupProgress, + ScenarioAttackResultConversations, ScenarioAttackResultDelta, ScenarioAttackTechniqueDetails, ScenarioConversationMessageProjection, @@ -139,16 +142,28 @@ def _make_test_snapshot( ], seed_groups=[ScenarioRunPlanSeedGroup(id="seed_1", objective_sha256="hash123", objective="Test objective")], ) - effective_summary = summary or 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=sum(r.total_retries for r in results), - ), - atomic_groups=atomic_groups or [], - ) + 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), @@ -357,6 +372,7 @@ def test_hierarchy_diagnostics_bounding_at_100() -> None: error_type="TestError", error_message=f"Error {i}", retries=retries, + seed_group_id=f"seed_{i}", seconds_offset=i, ) ) @@ -365,6 +381,7 @@ def test_hierarchy_diagnostics_bounding_at_100() -> None: 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) @@ -417,20 +434,24 @@ async def mock_get_deltas( 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=[ - ScenarioConversationProjection( - conversation_id="conv_1", - conversation_type="main", - message_count=1, - messages=[ - ScenarioConversationMessageProjection( - role="user", - content="Test prompt", - sequence=1, - ) - ], - ) - ] + 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) @@ -626,8 +647,10 @@ async def test_sqlite_memory_batched_reads_and_projections_parity(sqlite_instanc attack_result_id=child_id, include_messages=True, ) - assert len(conv_projections) == 1 - proj = conv_projections[0] + 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 @@ -1053,8 +1076,10 @@ async def test_sqlite_bounded_multimodal_conversation_projection(sqlite_instance limit_messages_per_conversation=2, ) - assert len(projs) == 1 - proj = projs[0] + 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 @@ -1074,3 +1099,133 @@ async def test_sqlite_bounded_multimodal_conversation_projection(sqlite_instance 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", + attempt_index=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", + attempt_index=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, + ) + + # 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 diff --git a/tests/unit/memory/test_azure_sql_memory.py b/tests/unit/memory/test_azure_sql_memory.py index 471568439e..0f40187c48 100644 --- a/tests/unit/memory/test_azure_sql_memory.py +++ b/tests/unit/memory/test_azure_sql_memory.py @@ -14,7 +14,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 from pyrit.memory.memory_models import ScenarioResultEntry from pyrit.memory.storage.serializers import set_message_piece_sha256_async @@ -1037,3 +1037,13 @@ def test_scenario_attack_result_deltas_by_ids_chunks_over_2100_ids_mssql( 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 diff --git a/tests/unit/models/test_scenario_hierarchy_models.py b/tests/unit/models/test_scenario_hierarchy_models.py index 04a123daea..21e0d8fe62 100644 --- a/tests/unit/models/test_scenario_hierarchy_models.py +++ b/tests/unit/models/test_scenario_hierarchy_models.py @@ -10,6 +10,7 @@ AttackOutcome, AttackResultRole, RetryEvent, + ScenarioAttackResultConversations, ScenarioAttackTechniqueDetails, ScenarioConversationMessagePieceProjection, ScenarioConversationMessageProjection, @@ -298,3 +299,41 @@ def test_scenario_hierarchy_projection_truncated_flag() -> None: 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 From 689eef904a23b7e498efa69be2723b55dd3fb6d8 Mon Sep 17 00:00:00 2001 From: Nimit Jain Date: Thu, 15 Oct 2026 00:15:01 +0530 Subject: [PATCH 04/18] fix: resolve remaining edge cases on executed_turns, piece limits, and reference bounds (#3167) --- .../services/scenario_progress_read_model.py | 118 ++++-- .../backend/services/scenario_run_service.py | 12 +- pyrit/memory/memory_interface.py | 43 +- pyrit/models/scenario_hierarchy.py | 28 +- pyrit/models/scenario_progress.py | 5 +- .../backend/test_scenario_hierarchy_parity.py | 385 ++++++++++++++++++ .../models/test_scenario_hierarchy_models.py | 58 +++ 7 files changed, 615 insertions(+), 34 deletions(-) diff --git a/pyrit/backend/services/scenario_progress_read_model.py b/pyrit/backend/services/scenario_progress_read_model.py index 9681893b4a..6d1fb5cb48 100644 --- a/pyrit/backend/services/scenario_progress_read_model.py +++ b/pyrit/backend/services/scenario_progress_read_model.py @@ -78,6 +78,25 @@ _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) @@ -713,6 +732,9 @@ def _map_progress_delta( child_attack_result_ids=ScenarioProgressReadModel._read_child_attack_result_ids( attack_metadata=delta.attack_metadata ), + related_conversation_ids=list( + delta.related_conversation_ids or delta.attack_metadata.get("related_conversation_ids") or [] + ), attempt_index=result_metadata.attempt_index, ) @@ -835,13 +857,17 @@ def _build_node( retries = retries[:_MAX_RETRIES_PER_NODE] retries_truncated = True - executed_turns = delta.executed_turns or int(delta.attack_metadata.get("executed_turns") or 0) + 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") @@ -854,7 +880,8 @@ def _build_node( 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=list(result.child_attack_result_ids), + 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, @@ -897,6 +924,20 @@ def _build_node( ) ) + related_conv_ids = ( + delta.related_conversation_ids + or getattr(result, "related_conversations", None) + or getattr(result, "related_conversation_ids", None) + or delta.attack_metadata.get("related_conversation_ids") + or [] + ) + if related_conv_ids: + node_conversations_count = None + node_conversations_truncated = True + else: + node_conversations_count = len(conversations) + node_conversations_truncated = False + return ScenarioHierarchyNode( attack_result_id=result.attack_result_id, conversation_id=result.conversation_id or "", @@ -919,12 +960,14 @@ def _build_node( total_retries=result.total_retries, retries=retries, retries_truncated=retries_truncated, - child_attack_result_ids=list(result.child_attack_result_ids), + child_attack_result_ids=bounded_child_ids, + child_references_truncated=child_refs_truncated, children=children, + children_truncated=False, unresolved_child_attack_result_ids=unresolved_children, conversations=conversations, - conversations_count=len(conversations), - conversations_truncated=False, + conversations_count=node_conversations_count, + conversations_truncated=node_conversations_truncated, orchestration=orchestration, ) @@ -1062,6 +1105,19 @@ def _build_node( attack_retries_truncated = attack_retries_count > _MAX_DIAGNOSTIC_ITEMS attack_retries = all_attack_retries[:_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) + + total_retry_pressure = 0 + for attempts in attempts_by_unit.values(): + total_retry_pressure += max(0, len(attempts) - 1) + sum( + max(a.total_retries, len(a.retries)) for a in attempts + ) + + diagnostic_items_retries = sum(max(r.total_retries, len(r.retries)) for r in all_attack_retries) + diagnostic_retries = max(snapshot.summary.overall.retries, total_retry_pressure, diagnostic_items_retries) + diagnostics = ScenarioHierarchyDiagnostics( failed_attacks=failed_attacks, failed_attacks_count=failed_attacks_count, @@ -1069,7 +1125,7 @@ def _build_node( attack_retries=attack_retries, attack_retries_count=attack_retries_count, attack_retries_truncated=attack_retries_truncated, - total_retries=snapshot.summary.overall.retries, + total_retries=diagnostic_retries, ) return ScenarioHierarchyProjection( @@ -1090,11 +1146,11 @@ def build_hierarchy_node( result: ScenarioProgressResult, children_deltas: Sequence[ScenarioAttackResultDelta] = (), children_results: Sequence[ScenarioProgressResult] = (), - conversations: Sequence[ScenarioConversationProjection] - | ScenarioAttackResultConversations - | tuple[Sequence[ScenarioConversationProjection], int, bool] = (), + 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. @@ -1107,6 +1163,8 @@ def build_hierarchy_node( 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. @@ -1117,18 +1175,12 @@ def build_hierarchy_node( if conversations_count is None: conversations_count = conversations.total_count conversations_truncated = conversations_truncated or conversations.truncated - elif isinstance(conversations, tuple) and len(conversations) == 3: - first, second, third = conversations - if isinstance(first, Sequence): - conv_list = [c for c in first if isinstance(c, ScenarioConversationProjection)] - if conversations_count is None and isinstance(second, int): - conversations_count = second - if isinstance(third, bool): - conversations_truncated = conversations_truncated or third else: conv_list = list(conversations) - resolved_conversations_count: int = conversations_count if conversations_count is not None else len(conv_list) + 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} @@ -1139,13 +1191,17 @@ def build_hierarchy_node( retries = retries[:_MAX_RETRIES_PER_NODE] retries_truncated = True - executed_turns = delta.executed_turns or int(delta.attack_metadata.get("executed_turns") or 0) + 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") @@ -1158,7 +1214,8 @@ def build_hierarchy_node( 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=list(result.child_attack_result_ids), + 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, @@ -1182,13 +1239,17 @@ def build_hierarchy_node( c_retries = c_retries[:_MAX_RETRIES_PER_NODE] c_retries_truncated = True - c_executed_turns = c_delta.executed_turns or int(c_delta.attack_metadata.get("executed_turns") or 0) + 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") @@ -1201,7 +1262,8 @@ def build_hierarchy_node( 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=list(c_res.child_attack_result_ids), + 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, @@ -1233,12 +1295,14 @@ def build_hierarchy_node( total_retries=c_res.total_retries, retries=c_retries, retries_truncated=c_retries_truncated, - child_attack_result_ids=list(c_res.child_attack_result_ids), + 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=0, - conversations_truncated=False, + conversations_count=None, + conversations_truncated=True, orchestration=c_orchestration, ) ) @@ -1265,8 +1329,10 @@ def build_hierarchy_node( total_retries=result.total_retries, retries=retries, retries_truncated=retries_truncated, - child_attack_result_ids=list(result.child_attack_result_ids), + child_attack_result_ids=bounded_child_ids, + child_references_truncated=child_refs_truncated, children=children, + children_truncated=children_truncated, unresolved_child_attack_result_ids=unresolved_children, conversations=conv_list, conversations_count=resolved_conversations_count, diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index 840bc76d26..7dc8cd4935 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -2297,6 +2297,7 @@ async def get_run_hierarchy_node_async( *, 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. @@ -2304,6 +2305,7 @@ async def get_run_hierarchy_node_async( 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. @@ -2334,11 +2336,18 @@ async def get_run_hierarchy_node_async( 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=child_ids, + attack_result_ids=fetch_child_ids, ) ) children_results = [ @@ -2356,6 +2365,7 @@ async def get_run_hierarchy_node_async( 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: diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index d1173bdfb4..9a48c8a61c 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -6525,6 +6525,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, @@ -6569,6 +6572,15 @@ def _execute_get_scenario_attack_result_deltas( score_value=row["score_value"], score_rationale=row["score_rationale"], ) + related_cids = [ + str(c) + for c in ( + list(row["pruned_conversation_ids"] or []) + + list(row["adversarial_chat_conversation_ids"] or []) + + list(row["preparation_conversation_ids"] or []) + ) + if c + ] deltas.append( ScenarioAttackResultDelta( attack_result_id=str(row["id"]), @@ -6578,7 +6590,7 @@ def _execute_get_scenario_attack_result_deltas( atomic_attack_identifier=atomic_identifier, outcome=AttackOutcome(row["outcome"]), execution_time_ms=row["execution_time_ms"], - executed_turns=row["executed_turns"] or 0, + executed_turns=row["executed_turns"], outcome_reason=row["outcome_reason"], timestamp=row["timestamp"], retry_events=retry_events, @@ -6587,6 +6599,7 @@ def _execute_get_scenario_attack_result_deltas( error_message=row["error_message"], attribution_data=row["attribution_data"] or {}, attack_metadata=row["attack_metadata"] or {}, + related_conversation_ids=related_cids, score=score, ) ) @@ -6642,6 +6655,9 @@ def _build_scenario_attack_result_deltas_by_ids_statement( 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, @@ -6723,6 +6739,15 @@ def _execute_get_scenario_attack_result_deltas_by_ids( score_value=row["score_value"], score_rationale=row["score_rationale"], ) + related_cids = [ + str(c) + for c in ( + list(row["pruned_conversation_ids"] or []) + + list(row["adversarial_chat_conversation_ids"] or []) + + list(row["preparation_conversation_ids"] or []) + ) + if c + ] deltas.append( ScenarioAttackResultDelta( attack_result_id=str(row["id"]), @@ -6732,7 +6757,7 @@ def _execute_get_scenario_attack_result_deltas_by_ids( atomic_attack_identifier=atomic_identifier, outcome=AttackOutcome(row["outcome"]), execution_time_ms=row["execution_time_ms"], - executed_turns=row["executed_turns"] or 0, + executed_turns=row["executed_turns"], outcome_reason=row["outcome_reason"], timestamp=row["timestamp"], retry_events=retry_events, @@ -6741,6 +6766,7 @@ def _execute_get_scenario_attack_result_deltas_by_ids( error_message=row["error_message"], attribution_data=row["attribution_data"] or {}, attack_metadata=row["attack_metadata"] or {}, + related_conversation_ids=related_cids, score=score, ) ) @@ -6751,6 +6777,7 @@ def _fetch_bounded_conversation_messages( *, 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. @@ -6801,13 +6828,16 @@ def _fetch_bounded_conversation_messages( messages: list[ScenarioConversationMessageProjection] = [] for msg in grouped_messages: + pieces_count = len(msg.message_pieces) + pieces_truncated = pieces_count > max_pieces_per_message + bounded_pieces = msg.message_pieces[:max_pieces_per_message] if pieces_truncated else msg.message_pieces 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 + for piece in bounded_pieces ] first_piece = msg.message_pieces[0] if msg.message_pieces else None raw_role = str(first_piece.role) if first_piece else "unknown" @@ -6822,6 +6852,7 @@ def _fetch_bounded_conversation_messages( sequence=seq, timestamp=ts, pieces=pieces_proj, + pieces_truncated=pieces_truncated, ) ) @@ -6834,6 +6865,7 @@ def _execute_get_attack_result_conversation_projections( 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. @@ -6887,6 +6919,7 @@ def _execute_get_attack_result_conversation_projections( 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( @@ -6933,6 +6966,7 @@ def _execute_get_attack_result_conversation_projections( 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( @@ -10454,6 +10488,7 @@ async def get_attack_result_conversation_projections_async( 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. @@ -10463,6 +10498,7 @@ async def get_attack_result_conversation_projections_async( 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. @@ -10473,6 +10509,7 @@ async def get_attack_result_conversation_projections_async( 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( diff --git a/pyrit/models/scenario_hierarchy.py b/pyrit/models/scenario_hierarchy.py index 47291c7fbd..4dfb83c71a 100644 --- a/pyrit/models/scenario_hierarchy.py +++ b/pyrit/models/scenario_hierarchy.py @@ -1,9 +1,9 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -from typing import NamedTuple +from typing import ClassVar, NamedTuple, Self -from pydantic import AwareDatetime, BaseModel, ConfigDict, Field +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 @@ -26,6 +26,8 @@ class ScenarioOrchestrationParentSummary(BaseModel): model_config = ConfigDict(extra="ignore") + max_child_references: ClassVar[int] = 500 + attack_result_id: str result_role: AttackResultRole = AttackResultRole.UNKNOWN outcome: AttackOutcome @@ -35,6 +37,7 @@ class ScenarioOrchestrationParentSummary(BaseModel): 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 @@ -42,6 +45,13 @@ class ScenarioOrchestrationParentSummary(BaseModel): 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): """ @@ -70,6 +80,7 @@ class ScenarioConversationMessageProjection(BaseModel): 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): @@ -99,6 +110,8 @@ class ScenarioHierarchyNode(BaseModel): 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 @@ -123,21 +136,30 @@ class ScenarioHierarchyNode(BaseModel): # 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) # Generic read-only conversations for this attack conversations: list[ScenarioConversationProjection] = Field(default_factory=list) - conversations_count: int = 0 + 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 + return self + class ScenarioAttackResultConversations(NamedTuple): """Container for bounded conversation projections, total count, and list truncation flag.""" diff --git a/pyrit/models/scenario_progress.py b/pyrit/models/scenario_progress.py index 33f35ab174..2b5933379c 100644 --- a/pyrit/models/scenario_progress.py +++ b/pyrit/models/scenario_progress.py @@ -183,6 +183,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) @@ -398,10 +400,11 @@ class ScenarioAttackResultDelta(BaseModel): total_retries: int = 0 error_type: str | None = None error_message: str | None = None - executed_turns: int = 0 + 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/unit/backend/test_scenario_hierarchy_parity.py b/tests/unit/backend/test_scenario_hierarchy_parity.py index 4051837455..3ad5ca9c5c 100644 --- a/tests/unit/backend/test_scenario_hierarchy_parity.py +++ b/tests/unit/backend/test_scenario_hierarchy_parity.py @@ -97,6 +97,41 @@ def _make_test_delta( ) +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], @@ -1229,3 +1264,353 @@ async def test_sqlite_conversation_list_clipping_explicit(sqlite_instance: SQLit 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 diff --git a/tests/unit/models/test_scenario_hierarchy_models.py b/tests/unit/models/test_scenario_hierarchy_models.py index 21e0d8fe62..3aa149b48f 100644 --- a/tests/unit/models/test_scenario_hierarchy_models.py +++ b/tests/unit/models/test_scenario_hierarchy_models.py @@ -337,3 +337,61 @@ def test_scenario_attack_result_conversations_namedtuple() -> None: 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 From f33a5b9b49ea6bb2f1729547d5bd869cca376fa1 Mon Sep 17 00:00:00 2001 From: Nimit Jain Date: Thu, 15 Oct 2026 12:13:20 +0530 Subject: [PATCH 05/18] fix: harden related conversation handling against malformed metadata and database shapes (#3167) --- .../services/scenario_progress_read_model.py | 58 +++++++-- pyrit/memory/memory_interface.py | 58 +++++---- .../backend/test_scenario_hierarchy_parity.py | 123 ++++++++++++++++++ 3 files changed, 205 insertions(+), 34 deletions(-) diff --git a/pyrit/backend/services/scenario_progress_read_model.py b/pyrit/backend/services/scenario_progress_read_model.py index 6d1fb5cb48..3c21ebe020 100644 --- a/pyrit/backend/services/scenario_progress_read_model.py +++ b/pyrit/backend/services/scenario_progress_read_model.py @@ -732,12 +732,28 @@ def _map_progress_delta( child_attack_result_ids=ScenarioProgressReadModel._read_child_attack_result_ids( attack_metadata=delta.attack_metadata ), - related_conversation_ids=list( - delta.related_conversation_ids or delta.attack_metadata.get("related_conversation_ids") or [] - ), + 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]: """ @@ -924,14 +940,19 @@ def _build_node( ) ) - related_conv_ids = ( - delta.related_conversation_ids + candidate_ids = ( + result.related_conversation_ids or getattr(result, "related_conversations", None) - or getattr(result, "related_conversation_ids", None) - or delta.attack_metadata.get("related_conversation_ids") - or [] + or ScenarioProgressReadModel._read_related_conversation_ids(delta=delta) ) - if related_conv_ids: + 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: @@ -1227,10 +1248,12 @@ def build_hierarchy_node( 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 not in children_deltas_by_id: + 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 @@ -1272,6 +1295,17 @@ def build_hierarchy_node( 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, @@ -1301,8 +1335,8 @@ def build_hierarchy_node( children_truncated=False, unresolved_child_attack_result_ids=[], conversations=[], - conversations_count=None, - conversations_truncated=True, + conversations_count=c_conv_count, + conversations_truncated=c_conv_truncated, orchestration=c_orchestration, ) ) diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 9a48c8a61c..b32b61cbac 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -6572,15 +6572,14 @@ def _execute_get_scenario_attack_result_deltas( score_value=row["score_value"], score_rationale=row["score_rationale"], ) - related_cids = [ - str(c) - for c in ( - list(row["pruned_conversation_ids"] or []) - + list(row["adversarial_chat_conversation_ids"] or []) - + list(row["preparation_conversation_ids"] or []) - ) - if c - ] + 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) deltas.append( ScenarioAttackResultDelta( attack_result_id=str(row["id"]), @@ -6621,7 +6620,10 @@ def _build_scenario_attack_result_deltas_by_ids_statement( Returns: Select: The constructed select statement. """ - scenario_uuid = uuid.UUID(scenario_result_id) + 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: @@ -6691,6 +6693,11 @@ def _execute_get_scenario_attack_result_deltas_by_ids( 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: @@ -6739,15 +6746,14 @@ def _execute_get_scenario_attack_result_deltas_by_ids( score_value=row["score_value"], score_rationale=row["score_rationale"], ) - related_cids = [ - str(c) - for c in ( - list(row["pruned_conversation_ids"] or []) - + list(row["adversarial_chat_conversation_ids"] or []) - + list(row["preparation_conversation_ids"] or []) - ) - if c - ] + 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) deltas.append( ScenarioAttackResultDelta( attack_result_id=str(row["id"]), @@ -6894,9 +6900,17 @@ def _execute_get_attack_result_conversation_projections( if entry.conversation_id: all_conv_ids.append(entry.conversation_id) - pruned_ids = entry.pruned_conversation_ids or [] - adversarial_ids = entry.adversarial_chat_conversation_ids or [] - preparation_ids = entry.preparation_conversation_ids or [] + 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: diff --git a/tests/unit/backend/test_scenario_hierarchy_parity.py b/tests/unit/backend/test_scenario_hierarchy_parity.py index 3ad5ca9c5c..6f2735973b 100644 --- a/tests/unit/backend/test_scenario_hierarchy_parity.py +++ b/tests/unit/backend/test_scenario_hierarchy_parity.py @@ -1614,3 +1614,126 @@ async def test_immediate_children_bounding_in_run_service(sqlite_instance: SQLit 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"] From 5d04b4f279c328d8d0b0facb921929db042e1f47 Mon Sep 17 00:00:00 2001 From: Nimit Jain Date: Thu, 15 Oct 2026 13:08:18 +0530 Subject: [PATCH 06/18] fix: use real conversation statistics and safely handle unavailable message counts (#3167) --- .../services/scenario_progress_read_model.py | 36 ++++- .../backend/services/scenario_run_service.py | 19 +++ pyrit/models/scenario_hierarchy.py | 2 +- .../backend/test_scenario_hierarchy_parity.py | 126 ++++++++++++++++++ 4 files changed, 180 insertions(+), 3 deletions(-) diff --git a/pyrit/backend/services/scenario_progress_read_model.py b/pyrit/backend/services/scenario_progress_read_model.py index 3c21ebe020..22a93152de 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 Iterable, Sequence +from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass, field from typing import Any, Literal @@ -817,6 +817,7 @@ def build_scenario_hierarchy( status: ScenarioRunState, snapshot: ScenarioProgressSnapshot, max_nodes: int = _MAX_HIERARCHY_NODES, + conversation_stats: Mapping[str, Any] | None = None, ) -> ScenarioHierarchyProjection: """ Assemble the scenario result hierarchy projection with cycle, depth, and node budget guards. @@ -828,6 +829,7 @@ def build_scenario_hierarchy( 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. Returns: ScenarioHierarchyProjection: Bounded hierarchy projection with diagnostics. @@ -931,12 +933,42 @@ def _build_node( conversations: list[ScenarioConversationProjection] = [] if result.conversation_id: + msg_count: int | None = None + messages_truncated: bool = True + if conversation_stats is not None: + if 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 + else: + msg_count = 0 + messages_truncated = False + else: + raw_count = delta.attack_metadata.get("conversation_message_count") + if raw_count is not None: + try: + parsed_count = int(raw_count) + msg_count = parsed_count + messages_truncated = parsed_count > 0 + except (ValueError, TypeError): + msg_count = None + messages_truncated = True + else: + msg_count = None + messages_truncated = True + 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=int(delta.attack_metadata.get("conversation_message_count") or 0), + message_count=msg_count, + messages=[], + messages_truncated=messages_truncated, ) ) diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index 7dc8cd4935..1d32597876 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -2283,6 +2283,24 @@ async def get_run_hierarchy_async( objective_scorer_identifier=objective_scorer_identifier, ) + # Obtain real conversation counts through a bounded batch of conversation statistics + seen_conv_ids: set[str] = set() + conv_ids_to_fetch: list[str] = [] + for r in progress_snapshot.results: + if r.conversation_id and r.conversation_id not in seen_conv_ids: + seen_conv_ids.add(r.conversation_id) + conv_ids_to_fetch.append(r.conversation_id) + if len(conv_ids_to_fetch) >= max_nodes: + break + + conversation_stats: dict[str, Any] = {} + if conv_ids_to_fetch: + try: + conversation_stats = await self._memory.get_conversation_stats_async(conversation_ids=conv_ids_to_fetch) + except Exception: + logger.warning("Could not fetch conversation statistics for hierarchy preview.", exc_info=True) + conversation_stats = {} + return self._progress_read_model.build_scenario_hierarchy( scenario_result_id=scenario_result_id, scenario_name=header_result.scenario_name, @@ -2290,6 +2308,7 @@ async def get_run_hierarchy_async( status=header_result.scenario_run_state, snapshot=progress_snapshot, max_nodes=max_nodes, + conversation_stats=conversation_stats, ) async def get_run_hierarchy_node_async( diff --git a/pyrit/models/scenario_hierarchy.py b/pyrit/models/scenario_hierarchy.py index 4dfb83c71a..783c9634bc 100644 --- a/pyrit/models/scenario_hierarchy.py +++ b/pyrit/models/scenario_hierarchy.py @@ -95,7 +95,7 @@ class ScenarioConversationProjection(BaseModel): conversation_id: str conversation_type: str = "main" description: str | None = None - message_count: int = 0 + message_count: int | None = None messages: list[ScenarioConversationMessageProjection] = Field(default_factory=list) messages_truncated: bool = False diff --git a/tests/unit/backend/test_scenario_hierarchy_parity.py b/tests/unit/backend/test_scenario_hierarchy_parity.py index 6f2735973b..0179cfdcca 100644 --- a/tests/unit/backend/test_scenario_hierarchy_parity.py +++ b/tests/unit/backend/test_scenario_hierarchy_parity.py @@ -1737,3 +1737,129 @@ def test_build_hierarchy_node_deduplicates_duplicate_child_ids() -> None: 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 From 0e7abdadaaba20098cfd2626d18efb94710590cb Mon Sep 17 00:00:00 2001 From: Nimit Jain Date: Thu, 15 Oct 2026 14:06:21 +0530 Subject: [PATCH 07/18] fix: apply SQL piece budgeting and reconcile diagnostic retries with shared retry policy (#3167) --- .../services/scenario_progress_read_model.py | 15 +- pyrit/memory/memory_interface.py | 39 +++-- .../backend/test_scenario_hierarchy_parity.py | 139 ++++++++++++++++++ 3 files changed, 174 insertions(+), 19 deletions(-) diff --git a/pyrit/backend/services/scenario_progress_read_model.py b/pyrit/backend/services/scenario_progress_read_model.py index 22a93152de..8bb68b6a4c 100644 --- a/pyrit/backend/services/scenario_progress_read_model.py +++ b/pyrit/backend/services/scenario_progress_read_model.py @@ -1162,14 +1162,13 @@ def _build_node( for r in snapshot.results: attempts_by_unit.setdefault((r.atomic_group_id, r.seed_group_id), []).append(r) - total_retry_pressure = 0 - for attempts in attempts_by_unit.values(): - total_retry_pressure += max(0, len(attempts) - 1) + sum( - max(a.total_retries, len(a.retries)) for a in attempts - ) - - diagnostic_items_retries = sum(max(r.total_retries, len(r.retries)) for r in all_attack_retries) - diagnostic_retries = max(snapshot.summary.overall.retries, total_retry_pressure, diagnostic_items_retries) + 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), + ), + ) diagnostics = ScenarioHierarchyDiagnostics( failed_attacks=failed_attacks, diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index b32b61cbac..f4e587721d 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -6811,43 +6811,60 @@ def _fetch_bounded_conversation_messages( truncated = len(seqs) > limit target_seqs = seqs[:limit] - pieces_stmt = ( - select(PromptMemoryEntry) + 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() + ) + 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(), ) ) - entries = session.execute(pieces_stmt).scalars().all() - if not entries: + rows = session.execute(pieces_stmt).all() + if not rows: return [], False - message_pieces = [e.get_message_piece() for e in entries] - grouped_messages = group_conversation_message_pieces_by_sequence(message_pieces=message_pieces) + 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: - pieces_count = len(msg.message_pieces) - pieces_truncated = pieces_count > max_pieces_per_message - bounded_pieces = msg.message_pieces[:max_pieces_per_message] if pieces_truncated else msg.message_pieces + 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 bounded_pieces + 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" - seq = first_piece.sequence if first_piece else 0 ts = first_piece.timestamp if first_piece else None content = msg.get_value() diff --git a/tests/unit/backend/test_scenario_hierarchy_parity.py b/tests/unit/backend/test_scenario_hierarchy_parity.py index 0179cfdcca..cacd9f0106 100644 --- a/tests/unit/backend/test_scenario_hierarchy_parity.py +++ b/tests/unit/backend/test_scenario_hierarchy_parity.py @@ -1863,3 +1863,142 @@ async def test_persisted_result_with_non_numeric_conversation_message_count_meta # 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 From e4d0ca06778a079c38b29c5665e2364ea8674801 Mon Sep 17 00:00:00 2001 From: Nimit Jain Date: Fri, 16 Oct 2026 06:24:47 +0530 Subject: [PATCH 08/18] fix: bound cold delta hydration, truncate unresolved child IDs, and bind preview stats to visible nodes (#3167) --- .../services/scenario_progress_read_model.py | 76 ++++-- .../backend/services/scenario_run_service.py | 77 ++++-- pyrit/models/scenario_hierarchy.py | 6 + .../backend/test_scenario_hierarchy_parity.py | 247 ++++++++++++++++++ .../models/test_scenario_hierarchy_models.py | 18 ++ 5 files changed, 385 insertions(+), 39 deletions(-) diff --git a/pyrit/backend/services/scenario_progress_read_model.py b/pyrit/backend/services/scenario_progress_read_model.py index 8bb68b6a4c..f5439e5b01 100644 --- a/pyrit/backend/services/scenario_progress_read_model.py +++ b/pyrit/backend/services/scenario_progress_read_model.py @@ -107,6 +107,7 @@ class ScenarioProgressSnapshot: results: tuple[ScenarioProgressResult, ...] summary: ScenarioProgressSummary plan: ScenarioRunPlan + has_more: bool = False @dataclass(frozen=True, slots=True) @@ -128,6 +129,7 @@ class _ProgressCacheEntry: cursor: AttackResultKeysetCursor | None = None summary: ScenarioProgressSummary | None = None summary_state: _ProgressSummaryState | None = None + has_more: bool = True class ScenarioProgressReadModel: @@ -151,6 +153,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. @@ -162,6 +165,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. @@ -169,7 +173,12 @@ async def get_snapshot_async( async with self._cache_lock: entry = self._get_cache_entry(scenario_result_id=scenario_result_id, plan=plan, terminal=terminal) first_new_index = len(entry.results) - await self._hydrate_new_deltas_async(scenario_result_id=scenario_result_id, entry=entry) + 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, @@ -247,13 +256,32 @@ 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) -> None: - """Hydrate every persisted delta after the cached keyset cursor.""" + async def _hydrate_new_deltas_async( + self, + *, + scenario_result_id: str, + entry: _ProgressCacheEntry, + terminal: bool = False, + max_deltas: int | None = None, + ) -> None: + """Hydrate persisted deltas after the cached keyset cursor up to max_deltas if specified.""" + if not entry.has_more and max_deltas is not None and len(entry.deltas) >= max_deltas: + return + if not entry.has_more and max_deltas is None and terminal: + return while True: + if max_deltas is not None and len(entry.deltas) >= max_deltas: + return + 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 ) entry.deltas.extend(page) if page: @@ -263,6 +291,8 @@ async def _hydrate_new_deltas_async(self, *, scenario_result_id: str, entry: _Pr attack_result_id=last.attack_result_id, ) if not has_more: + if terminal: + entry.has_more = False return if not page: raise RuntimeError("Scenario progress storage returned an empty page with has_more=True.") @@ -846,7 +876,7 @@ def build_scenario_hierarchy( visited_ids: set[str] = set() total_nodes_count = 0 - hierarchy_truncated = False + hierarchy_truncated = bool(snapshot.has_more) def _build_node( *, @@ -935,18 +965,14 @@ def _build_node( if result.conversation_id: msg_count: int | None = None messages_truncated: bool = True - if conversation_stats is not None: - if 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 - else: - msg_count = 0 - messages_truncated = False + 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 else: raw_count = delta.attack_metadata.get("conversation_message_count") if raw_count is not None: @@ -991,6 +1017,11 @@ def _build_node( 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 "", @@ -1017,7 +1048,8 @@ def _build_node( child_references_truncated=child_refs_truncated, children=children, children_truncated=False, - unresolved_child_attack_result_ids=unresolved_children, + 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, @@ -1372,6 +1404,11 @@ def build_hierarchy_node( ) ) + 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 "", @@ -1398,7 +1435,8 @@ def build_hierarchy_node( child_references_truncated=child_refs_truncated, children=children, children_truncated=children_truncated, - unresolved_child_attack_result_ids=unresolved_children, + 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, diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index 1d32597876..9145a4f1c8 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -2281,36 +2281,73 @@ async def get_run_hierarchy_async( 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), ) - # Obtain real conversation counts through a bounded batch of conversation statistics - seen_conv_ids: set[str] = set() - conv_ids_to_fetch: list[str] = [] - for r in progress_snapshot.results: - if r.conversation_id and r.conversation_id not in seen_conv_ids: - seen_conv_ids.add(r.conversation_id) - conv_ids_to_fetch.append(r.conversation_id) - if len(conv_ids_to_fetch) >= max_nodes: - break - - conversation_stats: dict[str, Any] = {} - if conv_ids_to_fetch: - try: - conversation_stats = await self._memory.get_conversation_stats_async(conversation_ids=conv_ids_to_fetch) - except Exception: - logger.warning("Could not fetch conversation statistics for hierarchy preview.", exc_info=True) - conversation_stats = {} - - return self._progress_read_model.build_scenario_hierarchy( + 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=conversation_stats, + conversation_stats=None, ) + # 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, *, diff --git a/pyrit/models/scenario_hierarchy.py b/pyrit/models/scenario_hierarchy.py index 783c9634bc..fc51d6cdc9 100644 --- a/pyrit/models/scenario_hierarchy.py +++ b/pyrit/models/scenario_hierarchy.py @@ -144,6 +144,7 @@ class ScenarioHierarchyNode(BaseModel): # 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) @@ -158,6 +159,11 @@ 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 diff --git a/tests/unit/backend/test_scenario_hierarchy_parity.py b/tests/unit/backend/test_scenario_hierarchy_parity.py index cacd9f0106..f44404c063 100644 --- a/tests/unit/backend/test_scenario_hierarchy_parity.py +++ b/tests/unit/backend/test_scenario_hierarchy_parity.py @@ -6,6 +6,7 @@ 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 @@ -2002,3 +2003,249 @@ async def test_diagnostics_total_retries_count_event_mismatch_parity( 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 diff --git a/tests/unit/models/test_scenario_hierarchy_models.py b/tests/unit/models/test_scenario_hierarchy_models.py index 3aa149b48f..7b3646a884 100644 --- a/tests/unit/models/test_scenario_hierarchy_models.py +++ b/tests/unit/models/test_scenario_hierarchy_models.py @@ -395,3 +395,21 @@ def test_scenario_conversation_message_pieces_truncated_flag() -> None: ) 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] From 6af0e8d3915ae7026a98290f3df51cba57953084 Mon Sep 17 00:00:00 2001 From: Nimit Jain Date: Fri, 16 Oct 2026 12:28:22 +0530 Subject: [PATCH 09/18] fix: decouple run-wide aggregates from bounded hierarchy and enforce zero-trust conversation stats (#3167) --- .../services/scenario_progress_read_model.py | 44 ++-- .../backend/services/scenario_run_service.py | 54 +++++ .../backend/test_scenario_hierarchy_parity.py | 223 ++++++++++++++++++ 3 files changed, 301 insertions(+), 20 deletions(-) diff --git a/pyrit/backend/services/scenario_progress_read_model.py b/pyrit/backend/services/scenario_progress_read_model.py index f5439e5b01..2dd5c819de 100644 --- a/pyrit/backend/services/scenario_progress_read_model.py +++ b/pyrit/backend/services/scenario_progress_read_model.py @@ -848,6 +848,9 @@ def build_scenario_hierarchy( 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, ) -> ScenarioHierarchyProjection: """ Assemble the scenario result hierarchy projection with cycle, depth, and node budget guards. @@ -860,6 +863,9 @@ def build_scenario_hierarchy( 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. Returns: ScenarioHierarchyProjection: Bounded hierarchy projection with diagnostics. @@ -973,19 +979,6 @@ def _build_node( if isinstance(c, int): msg_count = c messages_truncated = c > 0 - else: - raw_count = delta.attack_metadata.get("conversation_message_count") - if raw_count is not None: - try: - parsed_count = int(raw_count) - msg_count = parsed_count - messages_truncated = parsed_count > 0 - except (ValueError, TypeError): - msg_count = None - messages_truncated = True - else: - msg_count = None - messages_truncated = True conversations.append( ScenarioConversationProjection( @@ -1182,14 +1175,13 @@ def _build_node( ) ) - failed_attacks_count = len(all_failed_attacks) - failed_attacks_truncated = failed_attacks_count > _MAX_DIAGNOSTIC_ITEMS + failed_attacks_count = 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] - attack_retries_count = len(all_attack_retries) - attack_retries_truncated = attack_retries_count > _MAX_DIAGNOSTIC_ITEMS - attack_retries = all_attack_retries[:_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) @@ -1201,6 +1193,16 @@ def _build_node( 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 = len(all_attack_retries) + attack_retries_truncated = bool( + (snapshot.has_more and diagnostic_retries > total_persisted_retries) + or attack_retries_count > _MAX_DIAGNOSTIC_ITEMS + ) + attack_retries = all_attack_retries[:_MAX_DIAGNOSTIC_ITEMS] diagnostics = ScenarioHierarchyDiagnostics( failed_attacks=failed_attacks, @@ -1212,13 +1214,15 @@ def _build_node( 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=snapshot.summary.overall, + overall_counts=effective_overall_counts, diagnostics=diagnostics, truncated=hierarchy_truncated, ) diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index 9145a4f1c8..e68842e381 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -58,6 +58,9 @@ ScenarioHierarchyNode, ScenarioHierarchyProjection, ScenarioIdentifier, + ScenarioProducerCategoryCounts, + ScenarioProducerCounts, + ScenarioProgressCounts, ScenarioProgressHeader, ScenarioProgressResult, ScenarioQueueEntry, @@ -2284,6 +2287,54 @@ async def get_run_hierarchy_async( 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 + + if progress_snapshot.has_more: + 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, + ), + ) + 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, + ) + exact_total_errors = agg.error_attempts + exact_total_retries = agg.total_retries + hierarchy = self._progress_read_model.build_scenario_hierarchy( scenario_result_id=scenario_result_id, scenario_name=header_result.scenario_name, @@ -2292,6 +2343,9 @@ async def get_run_hierarchy_async( 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, ) # Obtain real conversation counts for the actual bounded preview nodes diff --git a/tests/unit/backend/test_scenario_hierarchy_parity.py b/tests/unit/backend/test_scenario_hierarchy_parity.py index f44404c063..14254c9cf5 100644 --- a/tests/unit/backend/test_scenario_hierarchy_parity.py +++ b/tests/unit/backend/test_scenario_hierarchy_parity.py @@ -2249,3 +2249,226 @@ def test_build_scenario_hierarchy_unqueried_conversation_not_false_zero() -> Non # 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 From 7c489c59791856c0201691ac09d9daf0f0bc015f Mon Sep 17 00:00:00 2001 From: Nimit Jain Date: Fri, 16 Oct 2026 15:30:03 +0530 Subject: [PATCH 10/18] fix: resolve run-scoped roots for truncated deltas and chunk conversation stats (#3167) --- .../services/scenario_progress_read_model.py | 34 ++- .../backend/services/scenario_run_service.py | 5 + pyrit/memory/azure_sql_memory.py | 96 ++++--- pyrit/memory/memory_interface.py | 262 ++++++++++-------- pyrit/memory/sqlite_memory.py | 102 +++---- .../backend/test_scenario_hierarchy_parity.py | 105 +++++++ tests/unit/memory/test_azure_sql_memory.py | 45 +++ 7 files changed, 439 insertions(+), 210 deletions(-) diff --git a/pyrit/backend/services/scenario_progress_read_model.py b/pyrit/backend/services/scenario_progress_read_model.py index 2dd5c819de..33ffa98c30 100644 --- a/pyrit/backend/services/scenario_progress_read_model.py +++ b/pyrit/backend/services/scenario_progress_read_model.py @@ -851,6 +851,7 @@ def build_scenario_hierarchy( overall_counts: ScenarioProgressCounts | None = None, exact_total_errors: int | None = None, exact_total_retries: int | None = None, + parent_deltas: Sequence[ScenarioAttackResultDelta] = (), ) -> ScenarioHierarchyProjection: """ Assemble the scenario result hierarchy projection with cycle, depth, and node budget guards. @@ -866,6 +867,8 @@ def build_scenario_hierarchy( 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. + parent_deltas: Optional orchestration parent deltas to resolve roots and child references + beyond delta limits. Returns: ScenarioHierarchyProjection: Bounded hierarchy projection with diagnostics. @@ -873,12 +876,28 @@ def build_scenario_hierarchy( 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 snapshot.results: + for result in results_by_id.values(): for child_id in result.child_attack_result_ids: referenced_child_ids.add(child_id) - root_results = [r for r in snapshot.results if r.attack_result_id not in referenced_child_ids] + # 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). + 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 + ] + root_results.sort(key=lambda r: (r.timestamp, r.attack_result_id)) visited_ids: set[str] = set() total_nodes_count = 0 @@ -1066,18 +1085,21 @@ def _build_node( # Handle unreached/disconnected components (including detached cycles) in a single linear pass if total_nodes_count < max_nodes: - unreached = [r for r in snapshot.results if r.attack_result_id not in visited_ids] - unreached.sort(key=lambda r: r.timestamp) + unreached = [r for r in results_by_id.values() if r.attack_result_id not in visited_ids] + 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 + fallback_child_index = ( + next_root.attempt_index - 1 if next_root.attempt_index is not None else None + ) root_node = _build_node( attack_result_id=next_root.attack_result_id, parent_id=None, - child_index=None, + child_index=fallback_child_index, active_path=set(), depth=0, ) @@ -1086,7 +1108,7 @@ def _build_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 snapshot.results): + if not hierarchy_truncated and any(r.attack_result_id not in visited_ids for r in results_by_id.values()): hierarchy_truncated = True for g_roots in roots_by_group.values(): diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index e68842e381..ffab9d2f4c 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -2290,8 +2290,12 @@ async def get_run_hierarchy_async( exact_overall_counts: ScenarioProgressCounts | None = None exact_total_errors: int | None = None exact_total_retries: int | None = None + parent_deltas: list[ScenarioAttackResultDelta] = [] if progress_snapshot.has_more: + parent_deltas = await self._memory.get_scenario_orchestration_parent_deltas_async( + scenario_result_id=scenario_result_id + ) 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], @@ -2346,6 +2350,7 @@ async def get_run_hierarchy_async( overall_counts=exact_overall_counts, exact_total_errors=exact_total_errors, exact_total_retries=exact_total_retries, + parent_deltas=parent_deltas, ) # Obtain real conversation counts for the actual bounded preview nodes diff --git a/pyrit/memory/azure_sql_memory.py b/pyrit/memory/azure_sql_memory.py index 104f04da0e..a65d2ad6c5 100644 --- a/pyrit/memory/azure_sql_memory.py +++ b/pyrit/memory/azure_sql_memory.py @@ -633,57 +633,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 f4e587721d..cbff5908af 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -6466,6 +6466,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, *, @@ -6547,61 +6601,9 @@ def _execute_get_scenario_attack_result_deltas( rows = session.execute(statement).mappings().all() has_more = len(rows) > limit - 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 []) - ] - 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) - 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"], - 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, - ) - ) + deltas: list[ScenarioAttackResultDelta] = [ + self._row_to_scenario_attack_result_delta(row) for row in rows[:limit] + ] return deltas, has_more def _build_scenario_attack_result_deltas_by_ids_statement( @@ -6722,62 +6724,83 @@ def _execute_get_scenario_attack_result_deltas_by_ids( attack_result_ids=[str(u) for u in chunk], ) rows = session.execute(statement).mappings().all() - for row in rows: - 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) - 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"], - 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, - ) - ) + 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) + 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 _fetch_bounded_conversation_messages( self, *, @@ -10512,6 +10535,27 @@ async def get_scenario_attack_result_deltas_by_ids_async( 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_attack_result_conversation_projections_async( self, *, diff --git a/pyrit/memory/sqlite_memory.py b/pyrit/memory/sqlite_memory.py index 01db1fe859..c1ebc80dcf 100644 --- a/pyrit/memory/sqlite_memory.py +++ b/pyrit/memory/sqlite_memory.py @@ -695,58 +695,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/tests/unit/backend/test_scenario_hierarchy_parity.py b/tests/unit/backend/test_scenario_hierarchy_parity.py index 14254c9cf5..b479adbf26 100644 --- a/tests/unit/backend/test_scenario_hierarchy_parity.py +++ b/tests/unit/backend/test_scenario_hierarchy_parity.py @@ -2472,3 +2472,108 @@ async def test_failed_conversation_stats_read_never_falls_back_to_metadata(sqlit # 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 diff --git a/tests/unit/memory/test_azure_sql_memory.py b/tests/unit/memory/test_azure_sql_memory.py index 0f40187c48..0e07ec948b 100644 --- a/tests/unit/memory/test_azure_sql_memory.py +++ b/tests/unit/memory/test_azure_sql_memory.py @@ -492,6 +492,51 @@ 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: From 04d713ea86b51645ab4b2f28a11e4895ec190aaf Mon Sep 17 00:00:00 2001 From: Nimit Jain Date: Sat, 17 Oct 2026 12:02:46 +0530 Subject: [PATCH 11/18] fix: enforce exact cold-warm parity, targeted parent resolution, and bounded recount (#3167) --- .../services/scenario_progress_read_model.py | 74 ++- .../backend/services/scenario_run_service.py | 147 ++++- pyrit/memory/memory_interface.py | 356 +++++++++++- .../backend/test_scenario_hierarchy_parity.py | 546 ++++++++++++++++++ 4 files changed, 1078 insertions(+), 45 deletions(-) diff --git a/pyrit/backend/services/scenario_progress_read_model.py b/pyrit/backend/services/scenario_progress_read_model.py index 33ffa98c30..2a91620010 100644 --- a/pyrit/backend/services/scenario_progress_read_model.py +++ b/pyrit/backend/services/scenario_progress_read_model.py @@ -851,6 +851,8 @@ def build_scenario_hierarchy( 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] = (), ) -> ScenarioHierarchyProjection: """ @@ -867,6 +869,8 @@ def build_scenario_hierarchy( 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. @@ -1093,13 +1097,10 @@ def _build_node( break if next_root.attack_result_id in visited_ids: continue - fallback_child_index = ( - next_root.attempt_index - 1 if next_root.attempt_index is not None else None - ) root_node = _build_node( attack_result_id=next_root.attack_result_id, parent_id=None, - child_index=fallback_child_index, + child_index=None, active_path=set(), depth=0, ) @@ -1120,16 +1121,29 @@ def _build_node( for ag in snapshot.summary.atomic_groups: seen_group_ids.add(ag.id) group_roots = roots_by_group.get(ag.id, []) - 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, - ) + 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, @@ -1150,12 +1164,15 @@ def _build_node( unit = ResultUnitIdentity(atomic_group_id=r.atomic_group_id, seed_group_id=r.seed_group_id) attempts_by_unit.setdefault(unit, []).append(r) - counts = count_execution_units( - units=list(attempts_by_unit.keys()), - attempts_by_unit=attempts_by_unit, - planned=None, - producer_attempts=group_results, - ) + 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, @@ -1197,7 +1214,11 @@ def _build_node( ) ) - failed_attacks_count = exact_total_errors if exact_total_errors is not None else len(all_failed_attacks) + 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 @@ -1219,9 +1240,16 @@ def _build_node( diagnostic_retries = max(diagnostic_retries, exact_total_retries) total_persisted_retries = sum(r.total_retries for r in snapshot.results) - attack_retries_count = len(all_attack_retries) + 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) + ( + 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] diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index ffab9d2f4c..71d6ab36eb 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -23,8 +23,13 @@ from pydantic import TypeAdapter, ValidationError -from pyrit.analytics.outcome_statistics import success_percentage -from pyrit.analytics.scenario_statistics import compute_scenario_statistics +from pyrit.analytics.outcome_statistics import compute_outcome_statistics, success_percentage +from pyrit.analytics.scenario_statistics import ( + ScenarioAttempt, + ScenarioExecutionUnit, + count_execution_units, + resolve_execution_unit, +) from pyrit.backend.models.common import PaginationInfo, filter_sensitive_fields from pyrit.backend.models.scenarios import ScenarioRunListResponse from pyrit.backend.services.pagination import ( @@ -51,8 +56,10 @@ from pyrit.models import ( SCENARIO_RUN_PLAN_METADATA_KEY, SCENARIO_RUN_STARTED_AT_METADATA_KEY, + AtomicAttackIdentifier, AttackOutcome, AttackResult, + AttackResultMetadata, ComponentIdentifier, ScenarioAttackResultDelta, ScenarioHierarchyNode, @@ -1617,10 +1624,76 @@ 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_deltas: list[ScenarioAttackResultDelta] = [] + while True: + page, has_more = await self._memory.get_scenario_attack_result_deltas_async( + scenario_result_id=run_id, + cursor=cursor, + limit=500, + ) + all_deltas.extend(page) + if not has_more or not page: + break + cursor = AttackResultKeysetCursor( + timestamp=page[-1].timestamp, attack_result_id=page[-1].attack_result_id + ) + + use_plan = plans.get(run_id) is not None + atomic_groups = plans.get(run_id) or [] if use_plan else [] + groups_by_identity: dict[tuple[str, str], ScenarioRunPlanAtomicGroup] = {} + grouped_by_name: dict[str, list[ScenarioRunPlanAtomicGroup]] = {} + for group in atomic_groups: + groups_by_identity[(group.atomic_attack_name, group.technique_eval_hash)] = group + grouped_by_name.setdefault(group.atomic_attack_name, []).append(group) + plan_lookup = ScenarioPlanLookup( + groups_by_identity=groups_by_identity, + groups_by_name={name: tuple(g) for name, g in grouped_by_name.items()}, + seed_ids_by_group_and_objective={}, + ) + attempts: list[ScenarioAttempt] = [] + for delta in all_deltas: + atomic_name = str(delta.attribution_data.get("parent_collection") or "") + eval_hash = delta.attribution_data.get("parent_eval_hash") + attr_seed = delta.attribution_data.get("seed_group_id") + typed_identifier = ( + AtomicAttackIdentifier.from_component_identifier(delta.atomic_attack_identifier) + if isinstance(delta.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=delta.objective, + objective_sha256=delta.objective_sha256, + plan_lookup=plan_lookup, + ) + res_metadata = AttackResultMetadata.from_metadata(metadata=delta.attribution_data) + attempts.append( + ScenarioAttempt( + unit=unit, + atomic_attack_name=atomic_name, + outcome=delta.outcome, + timestamp=delta.timestamp, + attempt_id=delta.attack_result_id, + total_retries=delta.total_retries, + result_role=res_metadata.result_role, + ) + ) + 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) + + overall = count_execution_units( + units=list(attempts_by_unit.keys()), + attempts_by_unit=attempts_by_unit, + planned=None, + producer_attempts=attempts, + ) recounted[run_id] = replace( aggregates[run_id], unit_count=overall.completed, @@ -1628,6 +1701,10 @@ 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, 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, @@ -2290,12 +2367,47 @@ async def get_run_hierarchy_async( 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] = [] if progress_snapshot.has_more: - parent_deltas = await self._memory.get_scenario_orchestration_parent_deltas_async( - scenario_result_id=scenario_result_id - ) + 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) + + current_candidates = [ + r.attack_result_id for r in all_results_map.values() if r.attack_result_id not in referenced_child_ids + ] + + depth = 0 + while current_candidates and depth < 10: + depth += 1 + new_parents = await self._memory.get_scenario_parent_deltas_for_child_ids_async( + scenario_result_id=scenario_result_id, + child_ids=current_candidates, + ) + if not new_parents: + break + next_candidates: list[str] = [] + 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_candidates.append(p.attack_result_id) + current_candidates = next_candidates + 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], @@ -2325,6 +2437,14 @@ async def get_run_hierarchy_async( 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, @@ -2335,9 +2455,12 @@ async def get_run_hierarchy_async( errors=agg.error_attempts, retries=agg.total_retries, producer_counts=producer_counts, + outcomes=outcomes, ) - exact_total_errors = agg.error_attempts - exact_total_retries = agg.total_retries + 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, @@ -2350,6 +2473,8 @@ async def get_run_hierarchy_async( 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, ) diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index cbff5908af..a52435429e 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -14,7 +14,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 @@ -45,6 +45,7 @@ from sqlalchemy.orm.attributes import InstrumentedAttribute, flag_modified from sqlalchemy.orm.session import Session +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 @@ -116,6 +117,9 @@ ScenarioConversationMessageProjection, ScenarioConversationProjection, ScenarioIdentifier, + ScenarioProducerCategoryCounts, + ScenarioProducerCounts, + ScenarioProgressCounts, ScenarioProgressScore, ScenarioResult, ScenarioRunState, @@ -133,6 +137,7 @@ SeedRecord, SeedType, TargetIdentifier, + config_hash, group_conversation_message_pieces_by_sequence, sort_message_pieces, ) @@ -282,6 +287,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": """ @@ -308,6 +324,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={}, ) @@ -6126,6 +6150,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)) @@ -6145,6 +6178,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: @@ -6173,25 +6258,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]: """ - 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, is_target_facing, is_orchestration, is_unknown) """ - atomic_name, technique_hash, attributed_seed_group_id, _ = self._get_scenario_attempt_unit_expressions() - # 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. + atomic_name, technique_hash, attributed_seed_group_id, identifier_seed_key = ( + self._get_scenario_attempt_unit_expressions() + ) attempts = ( select( AttackResultEntry.id.label("attempt_id"), @@ -6199,7 +6291,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"), @@ -6210,6 +6304,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"), ) .where(AttackResultEntry.attribution_parent_id.in_(entry_ids)) @@ -6228,6 +6332,7 @@ 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"), func.max(units.c.is_planned).over(partition_by=unit_partition).label("is_planned"), @@ -6235,6 +6340,7 @@ def _build_scenario_history_aggregate_statement( func.sum(case((is_error, 1), else_=0)).over(partition_by=unit_partition).label("unit_errors"), 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, @@ -6256,6 +6362,23 @@ def _build_scenario_history_aggregate_statement( ranked.c.result_role.is_(None), and_(not_(is_target_facing), not_(is_orchestration)), ) + return ranked, counted, 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, 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, @@ -6269,6 +6392,38 @@ 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_(counted, ranked.c.latest_outcome == AttackOutcome.SUCCESS.value), 1), else_=0) + ).label("outcome_successes"), + func.sum( + case((and_(counted, ranked.c.latest_outcome == AttackOutcome.FAILURE.value), 1), else_=0) + ).label("outcome_failures"), + func.sum(case((and_(counted, ranked.c.latest_outcome == AttackOutcome.ERROR.value), 1), else_=0)).label( + "outcome_errors" + ), + func.sum( + case((and_(counted, ranked.c.latest_outcome == AttackOutcome.UNDETERMINED.value), 1), else_=0) + ).label("outcome_undetermined"), + func.sum(case((ranked.c.unit_rank == 1, ranked.c.unit_errors), 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( + ( + and_( + ranked.c.unit_rank == 1, + or_( + ranked.c.unit_retries > 0, + ranked.c.attempt_retries > 0, + ranked.c.unit_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( @@ -6298,6 +6453,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, 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((counted, 1), else_=0)).label("completed_units"), + func.sum( + case((and_(counted, 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_(counted, ranked.c.latest_outcome == AttackOutcome.SUCCESS.value), 1), else_=0) + ).label("outcome_successes"), + func.sum( + case((and_(counted, ranked.c.latest_outcome == AttackOutcome.FAILURE.value), 1), else_=0) + ).label("outcome_failures"), + func.sum(case((and_(counted, ranked.c.latest_outcome == AttackOutcome.ERROR.value), 1), else_=0)).label( + "outcome_errors" + ), + func.sum( + case((and_(counted, 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. @@ -6319,6 +6542,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, unplanned_group_id.label("unit_group_id"), attempts.c.seed_group_id.label("unit_seed_id"), @@ -6399,6 +6623,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.atomic_attack_name, unplanned_group_id.label("unplanned_group_id"), @@ -6425,6 +6650,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, func.coalesce(matched.c.atomic_group_id, matched.c.unplanned_group_id).label("unit_group_id"), func.coalesce(matched.c.planned_seed_group_id, matched.c.seed_group_id).label("unit_seed_id"), @@ -6750,7 +6976,8 @@ def _execute_get_scenario_orchestration_parent_deltas( role=AttackResultRole.ORCHESTRATION, ) metadata_text = sql_cast(AttackResultEntry.attack_metadata, String) - child_ids_match = metadata_text.like("%child_attack_result_ids%") + # 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) @@ -6801,6 +7028,91 @@ def _execute_get_scenario_orchestration_parent_deltas( 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], + ) -> list[ScenarioAttackResultDelta]: + """ + Return run-scoped parent deltas that reference any of the specified child result IDs. + + Args: + scenario_result_id: Persisted scenario run ID. + child_ids: Child attack result IDs to find parents for. + + Returns: + list[ScenarioAttackResultDelta]: The matched parent deltas. + """ + try: + scenario_uuid = uuid.UUID(scenario_result_id) + except (ValueError, TypeError): + return [] + valid_child_ids = [cid for cid in child_ids if cid] + if not valid_child_ids: + return [] + + 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) + + chunk_size = min(len(valid_child_ids), max(1, self._MAX_BIND_VARS - 10)) + with closing(self._get_session()) as session: + for i in range(0, len(valid_child_ids), chunk_size): + chunk = valid_child_ids[i : i + chunk_size] + child_matches = [metadata_text.like(f'%"{cid}"%') for cid in chunk] + 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_(*child_matches), + ) + ) + .order_by(AttackResultEntry.timestamp.asc(), attempt_id_order.asc()) + ) + rows = session.execute(statement).mappings().all() + deltas.extend(self._row_to_scenario_attack_result_delta(row) for row in rows) + + seen_ids: set[str] = set() + unique_deltas: list[ScenarioAttackResultDelta] = [] + for d in deltas: + if d.attack_result_id not in seen_ids: + seen_ids.add(d.attack_result_id) + unique_deltas.append(d) + return unique_deltas + def _fetch_bounded_conversation_messages( self, *, @@ -10556,6 +10868,28 @@ async def get_scenario_orchestration_parent_deltas_async( 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], + ) -> list[ScenarioAttackResultDelta]: + """ + 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. + + Returns: + list[ScenarioAttackResultDelta]: The matched parent deltas. + """ + 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, + ) + async def get_attack_result_conversation_projections_async( self, *, diff --git a/tests/unit/backend/test_scenario_hierarchy_parity.py b/tests/unit/backend/test_scenario_hierarchy_parity.py index b479adbf26..0223ea612b 100644 --- a/tests/unit/backend/test_scenario_hierarchy_parity.py +++ b/tests/unit/backend/test_scenario_hierarchy_parity.py @@ -23,6 +23,7 @@ from pyrit.memory import MemoryInterface, PromptMemoryEntry, SQLiteMemory from pyrit.memory.storage.serializers import set_message_piece_sha256_async from pyrit.models import ( + AtomicAttackIdentifier, AttackOutcome, AttackResult, AttackResultRole, @@ -46,6 +47,7 @@ ScenarioRunPlanAtomicGroup, ScenarioRunPlanSeedGroup, ScenarioRunState, + SeedIdentifier, ) from unit.mocks import get_mock_target_identifier, make_scenario_result @@ -2577,3 +2579,547 @@ async def test_cold_hierarchy_projection_preserves_parent_beyond_delta_page_limi 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 + + scenario_results_called = False + delta_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: + delta_page_sizes.append(kwargs.get("limit", 0)) + return await original_get_deltas(*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] + + 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 paged! + assert scenario_results_called is False + assert len(delta_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 From fbf33012d71d10a4a2c565a8bae16a9f3dcfc843 Mon Sep 17 00:00:00 2001 From: Nimit Jain Date: Sat, 17 Oct 2026 14:43:20 +0530 Subject: [PATCH 12/18] fix: bound parent queries, handle cyclic discovery, and preserve planned recount (#3167) --- .../backend/services/scenario_run_service.py | 146 +++++- pyrit/memory/memory_interface.py | 61 ++- .../backend/test_scenario_hierarchy_parity.py | 446 ++++++++++++++++++ 3 files changed, 601 insertions(+), 52 deletions(-) diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index 71d6ab36eb..1ae4b9349b 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -1609,6 +1609,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. @@ -1628,30 +1629,32 @@ async def _recount_with_sdk_statistics_async( cursor: AttackResultKeysetCursor | None = None all_deltas: list[ScenarioAttackResultDelta] = [] while True: + remaining_limit = min(500, max_deltas - len(all_deltas)) if max_deltas is not None else 500 + if remaining_limit <= 0: + break page, has_more = await self._memory.get_scenario_attack_result_deltas_async( scenario_result_id=run_id, cursor=cursor, - limit=500, + limit=remaining_limit, ) all_deltas.extend(page) - if not has_more or not page: + if not has_more or not page or (max_deltas is not None and len(all_deltas) >= max_deltas): break cursor = AttackResultKeysetCursor( timestamp=page[-1].timestamp, attack_result_id=page[-1].attack_result_id ) - use_plan = plans.get(run_id) is not None - atomic_groups = plans.get(run_id) or [] if use_plan else [] - groups_by_identity: dict[tuple[str, str], ScenarioRunPlanAtomicGroup] = {} - grouped_by_name: dict[str, list[ScenarioRunPlanAtomicGroup]] = {} - for group in atomic_groups: - groups_by_identity[(group.atomic_attack_name, group.technique_eval_hash)] = group - grouped_by_name.setdefault(group.atomic_attack_name, []).append(group) - plan_lookup = ScenarioPlanLookup( - groups_by_identity=groups_by_identity, - groups_by_name={name: tuple(g) for name, g in grouped_by_name.items()}, - seed_ids_by_group_and_objective={}, - ) + 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 delta in all_deltas: atomic_name = str(delta.attribution_data.get("parent_collection") or "") @@ -1688,12 +1691,85 @@ async def _recount_with_sdk_statistics_async( 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=list(attempts_by_unit.keys()), + units=counted_units, attempts_by_unit=attempts_by_unit, - planned=None, + 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_deltas) >= max_deltas + raw_total_errors = ( + aggregates[run_id].raw_total_errors + if is_truncated + else sum(1 for d in all_deltas if d.outcome == AttackOutcome.ERROR) + ) + raw_total_retries = ( + aggregates[run_id].raw_total_retries if is_truncated else sum(d.total_retries for d in all_deltas) + ) + raw_retrying_attacks_count = ( + aggregates[run_id].raw_retrying_attacks_count + if is_truncated + else sum(1 for d in all_deltas if d.total_retries > 0 or bool(d.retry_events)) + ) + + merged_group_counts = {**aggregates[run_id].group_counts, **recounted_group_counts} + recounted[run_id] = replace( aggregates[run_id], unit_count=overall.completed, @@ -1705,6 +1781,10 @@ async def _recount_with_sdk_statistics_async( 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, @@ -2385,19 +2465,36 @@ async def get_run_hierarchy_async( for cid in r.child_attack_result_ids: referenced_child_ids.add(cid) - current_candidates = [ + 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 = self._progress_read_model._STORAGE_PAGE_SIZE + checked_child_ids: set[str] = set() depth = 0 - while current_candidates and depth < 10: + while current_candidates and depth < 10 and len(parent_deltas) < parent_budget: depth += 1 + remaining_budget = max(0, parent_budget - len(parent_deltas)) + if remaining_budget <= 0: + break new_parents = await self._memory.get_scenario_parent_deltas_for_child_ids_async( scenario_result_id=scenario_result_id, child_ids=current_candidates, + limit=remaining_budget, ) - if not new_parents: - break + checked_child_ids.update(current_candidates) next_candidates: list[str] = [] for p in new_parents: if p.attack_result_id not in all_deltas_map: @@ -2406,7 +2503,13 @@ async def get_run_hierarchy_async( 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_candidates.append(p.attack_result_id) + if len(parent_deltas) >= parent_budget: + break current_candidates = next_candidates + if not current_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: + current_candidates = remaining_cyclic plan_ids = [scenario_result_id] if plan is not None else [] aggregates = await self._memory.get_scenario_history_aggregates_async( @@ -2417,6 +2520,7 @@ async def get_run_hierarchy_async( aggregates = await self._recount_with_sdk_statistics_async( aggregates=aggregates, plans=plans_map, + max_deltas=max(max_nodes, self._progress_read_model._STORAGE_PAGE_SIZE), ) agg = aggregates.get(scenario_result_id) if agg is not None: diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index a52435429e..4cfcb65713 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -6335,6 +6335,8 @@ def _build_scenario_history_ranked_units( 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, 1), else_=0)).over(partition_by=unit_partition).label("unit_errors"), @@ -6404,26 +6406,11 @@ def _build_scenario_history_aggregate_statement( func.sum( case((and_(counted, ranked.c.latest_outcome == AttackOutcome.UNDETERMINED.value), 1), else_=0) ).label("outcome_undetermined"), - func.sum(case((ranked.c.unit_rank == 1, ranked.c.unit_errors), 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( - ( - and_( - ranked.c.unit_rank == 1, - or_( - ranked.c.unit_retries > 0, - ranked.c.attempt_retries > 0, - ranked.c.unit_has_retries == 1, - ), - ), - 1, - ), - else_=0, - ) - ).label("raw_retrying_attacks_count"), + func.sum(case((ranked.c.attempt_outcome == AttackOutcome.ERROR.value, 1), else_=0)).label( + "raw_total_errors" + ), + func.sum(ranked.c.attempt_retries).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( @@ -7033,6 +7020,7 @@ def _execute_get_scenario_parent_deltas_for_child_ids( *, scenario_result_id: str, child_ids: Sequence[str], + limit: int = 500, ) -> list[ScenarioAttackResultDelta]: """ Return run-scoped parent deltas that reference any of the specified child result IDs. @@ -7040,6 +7028,7 @@ def _execute_get_scenario_parent_deltas_for_child_ids( 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. Returns: list[ScenarioAttackResultDelta]: The matched parent deltas. @@ -7049,16 +7038,20 @@ def _execute_get_scenario_parent_deltas_for_child_ids( except (ValueError, TypeError): return [] valid_child_ids = [cid for cid in child_ids if cid] - if not valid_child_ids: + if not valid_child_ids or limit <= 0: return [] - deltas: list[ScenarioAttackResultDelta] = [] + 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) chunk_size = min(len(valid_child_ids), max(1, self._MAX_BIND_VARS - 10)) with closing(self._get_session()) as session: for i in range(0, len(valid_child_ids), chunk_size): + remaining_limit = max(0, limit - len(unique_deltas)) + if remaining_limit <= 0: + break chunk = valid_child_ids[i : i + chunk_size] child_matches = [metadata_text.like(f'%"{cid}"%') for cid in chunk] statement = ( @@ -7101,17 +7094,20 @@ def _execute_get_scenario_parent_deltas_for_child_ids( ) ) .order_by(AttackResultEntry.timestamp.asc(), attempt_id_order.asc()) + .limit(remaining_limit + 1) ) rows = session.execute(statement).mappings().all() - deltas.extend(self._row_to_scenario_attack_result_delta(row) for row in rows) - - seen_ids: set[str] = set() - unique_deltas: list[ScenarioAttackResultDelta] = [] - for d in deltas: - if d.attack_result_id not in seen_ids: - seen_ids.add(d.attack_result_id) - unique_deltas.append(d) - return unique_deltas + for row in rows: + delta = self._row_to_scenario_attack_result_delta(row) + if delta.attack_result_id not in seen_ids: + seen_ids.add(delta.attack_result_id) + unique_deltas.append(delta) + if len(unique_deltas) >= limit: + break + if len(unique_deltas) >= limit: + break + + return unique_deltas[:limit] def _fetch_bounded_conversation_messages( self, @@ -10873,6 +10869,7 @@ async def get_scenario_parent_deltas_for_child_ids_async( *, scenario_result_id: str, child_ids: Sequence[str], + limit: int = 500, ) -> list[ScenarioAttackResultDelta]: """ Return run-scoped parent deltas that reference any of the specified child result IDs. @@ -10880,6 +10877,7 @@ async def get_scenario_parent_deltas_for_child_ids_async( 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. Returns: list[ScenarioAttackResultDelta]: The matched parent deltas. @@ -10888,6 +10886,7 @@ async def get_scenario_parent_deltas_for_child_ids_async( self._execute_get_scenario_parent_deltas_for_child_ids, scenario_result_id=scenario_result_id, child_ids=child_ids, + limit=limit, ) async def get_attack_result_conversation_projections_async( diff --git a/tests/unit/backend/test_scenario_hierarchy_parity.py b/tests/unit/backend/test_scenario_hierarchy_parity.py index 0223ea612b..28a1f8dd39 100644 --- a/tests/unit/backend/test_scenario_hierarchy_parity.py +++ b/tests/unit/backend/test_scenario_hierarchy_parity.py @@ -3,6 +3,7 @@ """Unit and parity tests for bounded Scenario result hierarchy projections.""" +import hashlib import uuid from datetime import UTC, datetime, timedelta from pathlib import Path @@ -23,6 +24,7 @@ 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, @@ -3123,3 +3125,447 @@ async def test_501_unplanned_group_attempts_parity_between_cold_and_warm_cache( 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[0]], + "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 = 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 + + deltas_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 + + +@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 From 544ccbf600e53bd9b8cba2e02879069850d7c365 Mon Sep 17 00:00:00 2001 From: Nimit Jain Date: Sun, 18 Oct 2026 13:06:00 +0530 Subject: [PATCH 13/18] fix: exclude known rows from parent queries, handle multi-layer roots, and preserve exact run statistics (#3167) --- .../backend/services/scenario_run_service.py | 84 +++-- pyrit/memory/memory_interface.py | 40 ++- .../backend/test_scenario_hierarchy_parity.py | 327 +++++++++++++++++- 3 files changed, 398 insertions(+), 53 deletions(-) diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index 1ae4b9349b..142b23b098 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -1754,47 +1754,44 @@ async def _recount_with_sdk_statistics_async( recounted_group_counts[gid] = g_counts is_truncated = max_deltas is not None and len(all_deltas) >= max_deltas - raw_total_errors = ( - aggregates[run_id].raw_total_errors - if is_truncated - else sum(1 for d in all_deltas if d.outcome == AttackOutcome.ERROR) - ) - raw_total_retries = ( - aggregates[run_id].raw_total_retries if is_truncated else sum(d.total_retries for d in all_deltas) - ) - raw_retrying_attacks_count = ( - aggregates[run_id].raw_retrying_attacks_count - if is_truncated - else sum(1 for d in all_deltas if d.total_retries > 0 or bool(d.retry_events)) - ) - - merged_group_counts = {**aggregates[run_id].group_counts, **recounted_group_counts} - - recounted[run_id] = replace( - aggregates[run_id], - unit_count=overall.completed, - completed_units=overall.completed, - 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, - orchestration_attempts=overall.producer_counts.orchestration.attempts, - orchestration_error_attempts=overall.producer_counts.orchestration.errors, - orchestration_retries=overall.producer_counts.orchestration.retries, - unknown_role_attempts=overall.producer_counts.unknown.attempts, - unknown_role_error_attempts=overall.producer_counts.unknown.errors, - unknown_role_retries=overall.producer_counts.unknown.retries, - ) + if is_truncated: + # Bounded recount read only a prefix of deltas. Preserve exact run-wide + # statistics from the database aggregates rather than replacing them with prefix counts. + recounted[run_id] = aggregates[run_id] + else: + raw_total_errors = sum(1 for d in all_deltas if d.outcome == AttackOutcome.ERROR) + raw_total_retries = 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_deltas), + ) + raw_retrying_attacks_count = sum(1 for d in all_deltas if d.total_retries > 0 or bool(d.retry_events)) + merged_group_counts = {**aggregates[run_id].group_counts, **recounted_group_counts} + + recounted[run_id] = replace( + aggregates[run_id], + unit_count=overall.completed, + completed_units=overall.completed, + 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, + orchestration_attempts=overall.producer_counts.orchestration.attempts, + orchestration_error_attempts=overall.producer_counts.orchestration.errors, + orchestration_retries=overall.producer_counts.orchestration.retries, + unknown_role_attempts=overall.producer_counts.unknown.attempts, + unknown_role_error_attempts=overall.producer_counts.unknown.errors, + unknown_role_retries=overall.producer_counts.unknown.retries, + ) return recounted def _build_history_summary( @@ -2481,7 +2478,7 @@ async def get_run_hierarchy_async( 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 = self._progress_read_model._STORAGE_PAGE_SIZE + parent_budget = max(2500, self._progress_read_model._STORAGE_PAGE_SIZE * 5) checked_child_ids: set[str] = set() depth = 0 while current_candidates and depth < 10 and len(parent_deltas) < parent_budget: @@ -2492,7 +2489,8 @@ async def get_run_hierarchy_async( new_parents = await self._memory.get_scenario_parent_deltas_for_child_ids_async( scenario_result_id=scenario_result_id, child_ids=current_candidates, - limit=remaining_budget, + limit=min(self._progress_read_model._STORAGE_PAGE_SIZE, remaining_budget), + exclude_result_ids=list(all_deltas_map.keys()), ) checked_child_ids.update(current_candidates) next_candidates: list[str] = [] diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 4cfcb65713..45023a2c69 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -6409,7 +6409,9 @@ def _build_scenario_history_aggregate_statement( func.sum(case((ranked.c.attempt_outcome == AttackOutcome.ERROR.value, 1), else_=0)).label( "raw_total_errors" ), - func.sum(ranked.c.attempt_retries).label("raw_total_retries"), + 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( @@ -7021,6 +7023,7 @@ def _execute_get_scenario_parent_deltas_for_child_ids( scenario_result_id: str, child_ids: Sequence[str], limit: int = 500, + exclude_result_ids: Sequence[str] | None = None, ) -> list[ScenarioAttackResultDelta]: """ Return run-scoped parent deltas that reference any of the specified child result IDs. @@ -7029,6 +7032,7 @@ def _execute_get_scenario_parent_deltas_for_child_ids( 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: list[ScenarioAttackResultDelta]: The matched parent deltas. @@ -7041,12 +7045,27 @@ def _execute_get_scenario_parent_deltas_for_child_ids( if not valid_child_ids or limit <= 0: return [] + 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) - chunk_size = min(len(valid_child_ids), max(1, self._MAX_BIND_VARS - 10)) + exclude_clauses = [ + AttackResultEntry.id.not_in(valid_exclude_uuids[j : j + 200]) + for j in range(0, len(valid_exclude_uuids), 200) + ] + max_exclude_batch = min(200, len(valid_exclude_uuids)) + chunk_size = min(len(valid_child_ids), max(1, self._MAX_BIND_VARS - 10 - max_exclude_batch)) + with closing(self._get_session()) as session: for i in range(0, len(valid_child_ids), chunk_size): remaining_limit = max(0, limit - len(unique_deltas)) @@ -7054,6 +7073,11 @@ def _execute_get_scenario_parent_deltas_for_child_ids( break chunk = valid_child_ids[i : i + chunk_size] child_matches = [metadata_text.like(f'%"{cid}"%') for cid in chunk] + where_conditions = [ + AttackResultEntry.attribution_parent_id == scenario_uuid, + or_(*child_matches), + *exclude_clauses, + ] statement = ( select( AttackResultEntry.id, @@ -7087,19 +7111,14 @@ def _execute_get_scenario_parent_deltas_for_child_ids( func.coalesce(AttackResultEntry.human_score_id, AttackResultEntry.automated_score_id) == ScoreEntry.id, ) - .where( - and_( - AttackResultEntry.attribution_parent_id == scenario_uuid, - or_(*child_matches), - ) - ) + .where(and_(*where_conditions)) .order_by(AttackResultEntry.timestamp.asc(), attempt_id_order.asc()) .limit(remaining_limit + 1) ) rows = session.execute(statement).mappings().all() for row in rows: delta = self._row_to_scenario_attack_result_delta(row) - if delta.attack_result_id not in seen_ids: + 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: @@ -10870,6 +10889,7 @@ async def get_scenario_parent_deltas_for_child_ids_async( scenario_result_id: str, child_ids: Sequence[str], limit: int = 500, + exclude_result_ids: Sequence[str] | None = None, ) -> list[ScenarioAttackResultDelta]: """ Return run-scoped parent deltas that reference any of the specified child result IDs. @@ -10878,6 +10898,7 @@ async def get_scenario_parent_deltas_for_child_ids_async( 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: list[ScenarioAttackResultDelta]: The matched parent deltas. @@ -10887,6 +10908,7 @@ async def get_scenario_parent_deltas_for_child_ids_async( 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( diff --git a/tests/unit/backend/test_scenario_hierarchy_parity.py b/tests/unit/backend/test_scenario_hierarchy_parity.py index 28a1f8dd39..d1c1724d0f 100644 --- a/tests/unit/backend/test_scenario_hierarchy_parity.py +++ b/tests/unit/backend/test_scenario_hierarchy_parity.py @@ -3187,7 +3187,7 @@ async def test_500_cyclic_results_with_outer_parent_at_501_parity(sqlite_instanc "result_role": AttackResultRole.ORCHESTRATION.value, }, metadata={ - "child_attack_result_ids": [cyclic_ids[0]], + "child_attack_result_ids": [cyclic_ids[-1]], "completion_policy": "ALL_CHILDREN", "executed_turns": 1, }, @@ -3569,3 +3569,328 @@ async def test_recount_updates_group_counts_and_raw_aggregates_parity(sqlite_ins 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 From a1b11089c375f482692e8b9fc8d8581d74b27072 Mon Sep 17 00:00:00 2001 From: Nimit Jain Date: Sun, 18 Oct 2026 23:12:38 +0530 Subject: [PATCH 14/18] fix: page parent queries, guard unverified roots, and preserve legacy-identity recount (#3167) --- .../services/scenario_progress_read_model.py | 22 +- .../backend/services/scenario_run_service.py | 146 +++++---- .../backend/test_scenario_hierarchy_parity.py | 281 ++++++++++++++++++ tests/unit/memory/test_azure_sql_memory.py | 9 +- 4 files changed, 390 insertions(+), 68 deletions(-) diff --git a/pyrit/backend/services/scenario_progress_read_model.py b/pyrit/backend/services/scenario_progress_read_model.py index 2a91620010..8c32940ddb 100644 --- a/pyrit/backend/services/scenario_progress_read_model.py +++ b/pyrit/backend/services/scenario_progress_read_model.py @@ -854,6 +854,7 @@ def build_scenario_hierarchy( 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. @@ -873,6 +874,8 @@ def build_scenario_hierarchy( 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. @@ -896,11 +899,20 @@ def build_scenario_hierarchy( # 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 + 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 ] + if not root_results: + 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 + ] root_results.sort(key=lambda r: (r.timestamp, r.attack_result_id)) visited_ids: set[str] = set() @@ -1090,7 +1102,13 @@ def _build_node( # Handle unreached/disconnected components (including detached cycles) in a single linear pass if total_nodes_count < max_nodes: unreached = [r for r in results_by_id.values() if r.attack_result_id not in visited_ids] - unreached.sort(key=lambda r: (r.timestamp, r.attack_result_id)) + unreached.sort( + key=lambda r: ( + r.attack_result_id in unverified_set, + r.timestamp, + r.attack_result_id, + ) + ) for next_root in unreached: if total_nodes_count >= max_nodes: hierarchy_truncated = True diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index 142b23b098..d21cdd2b2f 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -1754,44 +1754,51 @@ async def _recount_with_sdk_statistics_async( recounted_group_counts[gid] = g_counts is_truncated = max_deltas is not None and len(all_deltas) >= max_deltas - if is_truncated: - # Bounded recount read only a prefix of deltas. Preserve exact run-wide - # statistics from the database aggregates rather than replacing them with prefix counts. - recounted[run_id] = aggregates[run_id] - else: - raw_total_errors = sum(1 for d in all_deltas if d.outcome == AttackOutcome.ERROR) - raw_total_retries = ScenarioProgressReadModel.total_retry_pressure( + raw_total_errors = ( + aggregates[run_id].raw_total_errors + if is_truncated + else sum(1 for d in all_deltas 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_deltas), ) - raw_retrying_attacks_count = sum(1 for d in all_deltas if d.total_retries > 0 or bool(d.retry_events)) - merged_group_counts = {**aggregates[run_id].group_counts, **recounted_group_counts} - - recounted[run_id] = replace( - aggregates[run_id], - unit_count=overall.completed, - completed_units=overall.completed, - 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, - orchestration_attempts=overall.producer_counts.orchestration.attempts, - orchestration_error_attempts=overall.producer_counts.orchestration.errors, - orchestration_retries=overall.producer_counts.orchestration.retries, - unknown_role_attempts=overall.producer_counts.unknown.attempts, - unknown_role_error_attempts=overall.producer_counts.unknown.errors, - unknown_role_retries=overall.producer_counts.unknown.retries, - ) + ) + raw_retrying_attacks_count = ( + aggregates[run_id].raw_retrying_attacks_count + if is_truncated + else sum(1 for d in all_deltas if d.total_retries > 0 or bool(d.retry_events)) + ) + merged_group_counts = {**aggregates[run_id].group_counts, **recounted_group_counts} + + recounted[run_id] = replace( + aggregates[run_id], + unit_count=overall.completed, + completed_units=overall.completed, + 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, + orchestration_attempts=overall.producer_counts.orchestration.attempts, + orchestration_error_attempts=overall.producer_counts.orchestration.errors, + orchestration_retries=overall.producer_counts.orchestration.retries, + unknown_role_attempts=overall.producer_counts.unknown.attempts, + unknown_role_error_attempts=overall.producer_counts.unknown.errors, + unknown_role_retries=overall.producer_counts.unknown.retries, + ) return recounted def _build_history_summary( @@ -2447,6 +2454,7 @@ async def get_run_hierarchy_async( 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) @@ -2481,33 +2489,53 @@ async def get_run_hierarchy_async( parent_budget = max(2500, self._progress_read_model._STORAGE_PAGE_SIZE * 5) checked_child_ids: set[str] = set() depth = 0 - while current_candidates and depth < 10 and len(parent_deltas) < parent_budget: + layer_candidates = list(current_candidates) + while layer_candidates and depth < 10 and len(parent_deltas) < parent_budget: depth += 1 - remaining_budget = max(0, parent_budget - len(parent_deltas)) - if remaining_budget <= 0: + 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 = 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(new_parents) < page_limit: + break + + if len(parent_deltas) >= parent_budget: + layer_clipped = True + break + + if layer_clipped: + unverified_candidate_ids.update(layer_candidates) break - new_parents = await self._memory.get_scenario_parent_deltas_for_child_ids_async( - scenario_result_id=scenario_result_id, - child_ids=current_candidates, - limit=min(self._progress_read_model._STORAGE_PAGE_SIZE, remaining_budget), - exclude_result_ids=list(all_deltas_map.keys()), - ) - checked_child_ids.update(current_candidates) - next_candidates: list[str] = [] - 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_candidates.append(p.attack_result_id) - if len(parent_deltas) >= parent_budget: - break - current_candidates = next_candidates - if not current_candidates and len(parent_deltas) < parent_budget and depth < 10: + + 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: - current_candidates = remaining_cyclic + layer_candidates = remaining_cyclic + + if layer_candidates: + unverified_candidate_ids.update(layer_candidates) plan_ids = [scenario_result_id] if plan is not None else [] aggregates = await self._memory.get_scenario_history_aggregates_async( @@ -2518,7 +2546,6 @@ async def get_run_hierarchy_async( aggregates = await self._recount_with_sdk_statistics_async( aggregates=aggregates, plans=plans_map, - max_deltas=max(max_nodes, self._progress_read_model._STORAGE_PAGE_SIZE), ) agg = aggregates.get(scenario_result_id) if agg is not None: @@ -2578,6 +2605,7 @@ async def get_run_hierarchy_async( 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 diff --git a/tests/unit/backend/test_scenario_hierarchy_parity.py b/tests/unit/backend/test_scenario_hierarchy_parity.py index d1c1724d0f..6864c8b27d 100644 --- a/tests/unit/backend/test_scenario_hierarchy_parity.py +++ b/tests/unit/backend/test_scenario_hierarchy_parity.py @@ -29,6 +29,8 @@ AttackOutcome, AttackResult, AttackResultRole, + AttackSeedGroup, + ComponentIdentifier, ConversationReference, ConversationType, MessagePiece, @@ -50,6 +52,8 @@ ScenarioRunPlanSeedGroup, ScenarioRunState, SeedIdentifier, + SeedObjective, + SeedPrompt, ) from unit.mocks import get_mock_target_identifier, make_scenario_result @@ -3894,3 +3898,280 @@ async def test_repeated_legacy_unit_500_errors_followed_by_recovery_parity( 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 diff --git a/tests/unit/memory/test_azure_sql_memory.py b/tests/unit/memory/test_azure_sql_memory.py index 0e07ec948b..7e6c7e7196 100644 --- a/tests/unit/memory/test_azure_sql_memory.py +++ b/tests/unit/memory/test_azure_sql_memory.py @@ -498,10 +498,7 @@ def test_get_conversation_stats_chunks_within_bind_budget( """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() - ] + 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 @@ -510,9 +507,7 @@ def test_get_conversation_stats_chunks_within_bind_budget( # 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 - ) + 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] From dfb3ce3779be1cbd2822d78d42d3114fb7cee523 Mon Sep 17 00:00:00 2001 From: Nimit Jain Date: Tue, 20 Oct 2026 02:06:08 +0530 Subject: [PATCH 15/18] fix: lightweight recount projection, guard boundary roots, and bound exclusion SQL params (#3167) - Add ScenarioAttemptSeedSummary dataclass and memory methods to recount legacy runs with lightweight scalar projections instead of hydrating full ScenarioAttackResultDelta objects with outer joins and JSON parsing. - Remove root promotion fallback for unverified candidates and exclude unverified candidates and their descendants from being emitted as roots during depth-10 and 2,500-parent budget cutoffs. - Dynamically budget bind variables in _execute_get_scenario_parent_deltas_for_child_ids to ensure statement parameters strictly stay within _MAX_BIND_VARS, preventing SQL Server 2,100 parameter overflow when large exclusion lists are present. - Add regression tests verifying depth-10 chain truncation, 2,500-parent boundary guards, lightweight recount without delta hydration, and MSSQL parameter compilation limits with 1,100+ and 2,500+ exclusions. --- .../services/scenario_progress_read_model.py | 29 +- .../backend/services/scenario_run_service.py | 59 ++-- pyrit/memory/__init__.py | 2 + pyrit/memory/memory_interface.py | 155 +++++++++- .../backend/test_scenario_hierarchy_parity.py | 267 +++++++++++++++++- tests/unit/memory/test_azure_sql_memory.py | 47 +++ 6 files changed, 515 insertions(+), 44 deletions(-) diff --git a/pyrit/backend/services/scenario_progress_read_model.py b/pyrit/backend/services/scenario_progress_read_model.py index 8c32940ddb..b462ef011e 100644 --- a/pyrit/backend/services/scenario_progress_read_model.py +++ b/pyrit/backend/services/scenario_progress_read_model.py @@ -907,12 +907,6 @@ def build_scenario_hierarchy( and r.attempt_index is None and r.attack_result_id not in unverified_set ] - if not root_results: - 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 - ] root_results.sort(key=lambda r: (r.timestamp, r.attack_result_id)) visited_ids: set[str] = set() @@ -1101,10 +1095,25 @@ def _build_node( # Handle unreached/disconnected components (including detached cycles) in a single linear pass if total_nodes_count < max_nodes: - unreached = [r for r in results_by_id.values() if r.attack_result_id not in visited_ids] + # 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.attack_result_id in unverified_set, r.timestamp, r.attack_result_id, ) @@ -1127,7 +1136,9 @@ def _build_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()): + 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(): diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index d21cdd2b2f..d8fdf44694 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -49,6 +49,7 @@ from pyrit.memory import AttackResultKeysetCursor, CentralMemory, SQLiteMemory from pyrit.memory.memory_interface import ( MemoryInterface, + ScenarioAttemptSeedSummary, ScenarioHistoryAggregate, ScenarioHistoryKeysetCursor, ScenarioHistoryRunRecord, @@ -1627,18 +1628,18 @@ async def _recount_with_sdk_statistics_async( recounted = dict(aggregates) for run_id in run_ids: cursor: AttackResultKeysetCursor | None = None - all_deltas: list[ScenarioAttackResultDelta] = [] + all_summaries: list[ScenarioAttemptSeedSummary] = [] while True: - remaining_limit = min(500, max_deltas - len(all_deltas)) if max_deltas is not None else 500 + 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_attack_result_deltas_async( + page, has_more = await self._memory.get_scenario_recount_summaries_async( scenario_result_id=run_id, cursor=cursor, limit=remaining_limit, ) - all_deltas.extend(page) - if not has_more or not page or (max_deltas is not None and len(all_deltas) >= max_deltas): + 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 @@ -1656,13 +1657,13 @@ async def _recount_with_sdk_statistics_async( ) plan_lookup = ScenarioPlanLookup.from_plan(plan=saved_plan) attempts: list[ScenarioAttempt] = [] - for delta in all_deltas: - atomic_name = str(delta.attribution_data.get("parent_collection") or "") - eval_hash = delta.attribution_data.get("parent_eval_hash") - attr_seed = delta.attribution_data.get("seed_group_id") + 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(delta.atomic_attack_identifier) - if isinstance(delta.atomic_attack_identifier, ComponentIdentifier) + AtomicAttackIdentifier.from_component_identifier(summary.atomic_attack_identifier) + if isinstance(summary.atomic_attack_identifier, ComponentIdentifier) else None ) unit = resolve_execution_unit( @@ -1670,19 +1671,19 @@ async def _recount_with_sdk_statistics_async( 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=delta.objective, - objective_sha256=delta.objective_sha256, + objective=summary.objective, + objective_sha256=summary.objective_sha256, plan_lookup=plan_lookup, ) - res_metadata = AttackResultMetadata.from_metadata(metadata=delta.attribution_data) + res_metadata = AttackResultMetadata.from_metadata(metadata=summary.attribution_data) attempts.append( ScenarioAttempt( unit=unit, atomic_attack_name=atomic_name, - outcome=delta.outcome, - timestamp=delta.timestamp, - attempt_id=delta.attack_result_id, - total_retries=delta.total_retries, + outcome=summary.outcome, + timestamp=summary.timestamp, + attempt_id=summary.attack_result_id, + total_retries=summary.total_retries, result_role=res_metadata.result_role, ) ) @@ -1753,24 +1754,24 @@ async def _recount_with_sdk_statistics_async( ) recounted_group_counts[gid] = g_counts - is_truncated = max_deltas is not None and len(all_deltas) >= max_deltas + 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_deltas if d.outcome == AttackOutcome.ERROR) + 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_deltas), + 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_deltas if d.total_retries > 0 or bool(d.retry_events)) + else sum(1 for d in all_summaries if d.total_retries > 0) ) merged_group_counts = {**aggregates[run_id].group_counts, **recounted_group_counts} @@ -2515,15 +2516,16 @@ async def get_run_hierarchy_async( if len(parent_deltas) >= parent_budget: break - if len(new_parents) < page_limit: - break - if len(parent_deltas) >= parent_budget: layer_clipped = True break + if len(new_parents) < page_limit: + break + if layer_clipped: unverified_candidate_ids.update(layer_candidates) + unverified_candidate_ids.update(next_layer_candidates) break checked_child_ids.update(layer_candidates) @@ -2537,6 +2539,13 @@ async def get_run_hierarchy_async( 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], 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/memory_interface.py b/pyrit/memory/memory_interface.py index 45023a2c69..5eaf2d771d 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -343,6 +343,26 @@ 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] + + @dataclass(frozen=True, slots=True, kw_only=True) class _AttackResultQuery: """ @@ -6821,6 +6841,89 @@ def _execute_get_scenario_attack_result_deltas( ] 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 {}, + ) + + 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, + ) + .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, *, @@ -7059,12 +7162,25 @@ def _execute_get_scenario_parent_deltas_for_child_ids( metadata_text = sql_cast(AttackResultEntry.attack_metadata, String) attempt_id_order = self._get_scenario_attempt_id_order_expression(attempt_id=AttackResultEntry.id) - exclude_clauses = [ - AttackResultEntry.id.not_in(valid_exclude_uuids[j : j + 200]) - for j in range(0, len(valid_exclude_uuids), 200) - ] - max_exclude_batch = min(200, len(valid_exclude_uuids)) - chunk_size = min(len(valid_child_ids), max(1, self._MAX_BIND_VARS - 10 - max_exclude_batch)) + _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) + ] + 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), + ) with closing(self._get_session()) as session: for i in range(0, len(valid_child_ids), chunk_size): @@ -10838,6 +10954,33 @@ 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, *, diff --git a/tests/unit/backend/test_scenario_hierarchy_parity.py b/tests/unit/backend/test_scenario_hierarchy_parity.py index 6864c8b27d..e5080ed4ff 100644 --- a/tests/unit/backend/test_scenario_hierarchy_parity.py +++ b/tests/unit/backend/test_scenario_hierarchy_parity.py @@ -3006,9 +3006,11 @@ async def test_legacy_recount_paged_deltas_without_hydrating_scenario_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 - delta_page_sizes: list[int] = [] + deltas_called = False + summary_page_sizes: list[int] = [] async def tracking_get_scenario_results(*args: Any, **kwargs: Any) -> Any: nonlocal scenario_results_called @@ -3016,11 +3018,17 @@ async def tracking_get_scenario_results(*args: Any, **kwargs: Any) -> Any: return await original_get_scenario_results(*args, **kwargs) async def tracking_get_deltas(*args: Any, **kwargs: Any) -> Any: - delta_page_sizes.append(kwargs.get("limit", 0)) + 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) @@ -3031,9 +3039,10 @@ async def tracking_get_deltas(*args: Any, **kwargs: Any) -> Any: assert matched_run is not None assert matched_run.completed_attacks == 501 - # Verified: full ScenarioResult was NOT loaded, deltas were paged! + # Verified: full ScenarioResult was NOT loaded, deltas were NOT hydrated, summaries were paged! assert scenario_results_called is False - assert len(delta_page_sizes) >= 2 # Paged in batches of 500 + assert deltas_called is False + assert len(summary_page_sizes) >= 2 # Paged in batches of 500 @pytest.mark.asyncio @@ -4175,3 +4184,253 @@ async def test_saved_plan_two_seeds_shared_objective_different_context_500_attem 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] diff --git a/tests/unit/memory/test_azure_sql_memory.py b/tests/unit/memory/test_azure_sql_memory.py index 7e6c7e7196..fcd13e13c7 100644 --- a/tests/unit/memory/test_azure_sql_memory.py +++ b/tests/unit/memory/test_azure_sql_memory.py @@ -1087,3 +1087,50 @@ def test_scenario_attack_result_deltas_by_ids_statement_type_hints() -> None: 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 = 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 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 = 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 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 From 1d14bd23763b891a173bee9400725d00cde57c48 Mon Sep 17 00:00:00 2001 From: Nimit Jain Date: Tue, 20 Oct 2026 10:24:15 +0530 Subject: [PATCH 16/18] fix: keyset-page parent discovery across exclusion fallbacks and preserve retry-event diagnostics (#3167) --- .../backend/services/scenario_run_service.py | 10 +- pyrit/memory/memory_interface.py | 178 ++++++++++----- .../backend/test_scenario_hierarchy_parity.py | 204 +++++++++++++++++- tests/unit/memory/test_azure_sql_memory.py | 6 +- 4 files changed, 336 insertions(+), 62 deletions(-) diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index d8fdf44694..633fe1e159 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -1771,7 +1771,7 @@ async def _recount_with_sdk_statistics_async( raw_retrying_attacks_count = ( aggregates[run_id].raw_retrying_attacks_count if is_truncated - else sum(1 for d in all_summaries if d.total_retries > 0) + 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} @@ -2500,7 +2500,7 @@ async def get_run_hierarchy_async( if remaining_budget <= 0: break page_limit = min(self._progress_read_model._STORAGE_PAGE_SIZE, remaining_budget) - new_parents = await self._memory.get_scenario_parent_deltas_for_child_ids_async( + 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, @@ -2520,7 +2520,11 @@ async def get_run_hierarchy_async( layer_clipped = True break - if len(new_parents) < page_limit: + if not has_more: + break + + if not new_parents: + layer_clipped = True break if layer_clipped: diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 5eaf2d771d..61cf32bed3 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -361,6 +361,7 @@ class ScenarioAttemptSeedSummary: timestamp: datetime total_retries: int attribution_data: dict[str, Any] + has_retries: bool = False @dataclass(frozen=True, slots=True, kw_only=True) @@ -6857,6 +6858,7 @@ def _row_to_scenario_attempt_seed_summary(row: Mapping[str, Any]) -> ScenarioAtt 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( @@ -6910,6 +6912,16 @@ def _execute_get_scenario_recount_summaries( 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()) @@ -7127,10 +7139,14 @@ def _execute_get_scenario_parent_deltas_for_child_ids( child_ids: Sequence[str], limit: int = 500, exclude_result_ids: Sequence[str] | None = None, - ) -> list[ScenarioAttackResultDelta]: + ) -> 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. @@ -7138,15 +7154,16 @@ def _execute_get_scenario_parent_deltas_for_child_ids( exclude_result_ids: Optional result IDs already known to caller to exclude from query. Returns: - list[ScenarioAttackResultDelta]: The matched parent deltas. + 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 [] + return [], False valid_child_ids = [cid for cid in child_ids if cid] if not valid_child_ids or limit <= 0: - return [] + return [], False valid_exclude_uuids: list[uuid.UUID] = [] if exclude_result_ids: @@ -7182,67 +7199,117 @@ def _execute_get_scenario_parent_deltas_for_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 i in range(0, len(valid_child_ids), chunk_size): - remaining_limit = max(0, limit - len(unique_deltas)) - if remaining_limit <= 0: + 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] - where_conditions = [ + base_where_conditions = [ AttackResultEntry.attribution_parent_id == scenario_uuid, or_(*child_matches), *exclude_clauses, ] - 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, + + 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) ) - .outerjoin( - ScoreEntry, - func.coalesce(AttackResultEntry.human_score_id, AttackResultEntry.automated_score_id) - == ScoreEntry.id, + 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"]), ) - .where(and_(*where_conditions)) - .order_by(AttackResultEntry.timestamp.asc(), attempt_id_order.asc()) - .limit(remaining_limit + 1) - ) - rows = session.execute(statement).mappings().all() - for row in rows: - 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: - break - if len(unique_deltas) >= limit: + + for row in 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: + break + + if len(unique_deltas) >= limit: + if has_more_in_chunk 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] + return unique_deltas[:limit], has_more def _fetch_bounded_conversation_messages( self, @@ -11033,7 +11100,7 @@ async def get_scenario_parent_deltas_for_child_ids_async( child_ids: Sequence[str], limit: int = 500, exclude_result_ids: Sequence[str] | None = None, - ) -> list[ScenarioAttackResultDelta]: + ) -> tuple[list[ScenarioAttackResultDelta], bool]: """ Return run-scoped parent deltas that reference any of the specified child result IDs. @@ -11044,7 +11111,8 @@ async def get_scenario_parent_deltas_for_child_ids_async( exclude_result_ids: Optional result IDs already known to caller to exclude from query. Returns: - list[ScenarioAttackResultDelta]: The matched parent deltas. + 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, diff --git a/tests/unit/backend/test_scenario_hierarchy_parity.py b/tests/unit/backend/test_scenario_hierarchy_parity.py index e5080ed4ff..33d490453b 100644 --- a/tests/unit/backend/test_scenario_hierarchy_parity.py +++ b/tests/unit/backend/test_scenario_hierarchy_parity.py @@ -3278,19 +3278,21 @@ async def test_shared_child_with_1001_parents_bounded_parent_query(sqlite_instan await sqlite_instance.add_attack_results_to_memory_async(attack_results=[child_result, *parents]) - deltas = await sqlite_instance.get_scenario_parent_deltas_for_child_ids_async( + 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 = await sqlite_instance.get_scenario_parent_deltas_for_child_ids_async( + 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 @@ -4434,3 +4436,201 @@ async def tracking_get_summaries(*args: Any, **kwargs: Any) -> Any: # 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("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 fcd13e13c7..ea1d019701 100644 --- a/tests/unit/memory/test_azure_sql_memory.py +++ b/tests/unit/memory/test_azure_sql_memory.py @@ -1101,13 +1101,14 @@ def test_scenario_parent_deltas_for_child_ids_bounds_parameters_with_large_exclu 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_parent_deltas_for_child_ids( + 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] @@ -1121,13 +1122,14 @@ def test_scenario_parent_deltas_for_child_ids_bounds_parameters_with_large_exclu session.reset_mock() with patch.object(memory_interface, "_get_session", return_value=session): - deltas = memory_interface._execute_get_scenario_parent_deltas_for_child_ids( + 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] From 294f3de8bd943801553011cc83c626e5fe515a7f Mon Sep 17 00:00:00 2001 From: Nimit Jain Date: Tue, 20 Oct 2026 12:35:27 +0530 Subject: [PATCH 17/18] fix: account for unprocessed page rows in parent query has_more contract (#3167) --- pyrit/memory/memory_interface.py | 7 +- .../backend/test_scenario_hierarchy_parity.py | 130 ++++++++++++++++++ 2 files changed, 135 insertions(+), 2 deletions(-) diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 61cf32bed3..8a96b41606 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -7289,16 +7289,19 @@ def _execute_get_scenario_parent_deltas_for_child_ids( attack_result_id=str(last_scanned_row["id"]), ) - for row in rows[:scanned_in_query]: + 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 chunk_idx + 1 < num_chunks: + if has_more_in_chunk or has_unprocessed_in_page or chunk_idx + 1 < num_chunks: has_more = True break diff --git a/tests/unit/backend/test_scenario_hierarchy_parity.py b/tests/unit/backend/test_scenario_hierarchy_parity.py index 33d490453b..dea7f1b537 100644 --- a/tests/unit/backend/test_scenario_hierarchy_parity.py +++ b/tests/unit/backend/test_scenario_hierarchy_parity.py @@ -4543,6 +4543,136 @@ async def test_1002_results_500_leaves_501_parents_leaf_0_not_emitted_as_root_pa 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( From 625b470c027b5587bd9fa77a249044046b2ae4ca Mon Sep 17 00:00:00 2001 From: Nimit Jain Date: Tue, 20 Oct 2026 13:41:42 +0530 Subject: [PATCH 18/18] fix: preserve keyset delta polling, completed units aggregation, and retry diagnostics --- pyrit/analytics/scenario_statistics.py | 4 +- .../services/scenario_progress_read_model.py | 16 ++++--- .../backend/services/scenario_run_service.py | 2 + pyrit/memory/memory_interface.py | 42 +++++++++---------- .../backend/test_scenario_hierarchy_parity.py | 4 -- 5 files changed, 32 insertions(+), 36 deletions(-) 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/services/scenario_progress_read_model.py b/pyrit/backend/services/scenario_progress_read_model.py index 1142dd927b..a3e027bc0b 100644 --- a/pyrit/backend/services/scenario_progress_read_model.py +++ b/pyrit/backend/services/scenario_progress_read_model.py @@ -315,14 +315,10 @@ async def _hydrate_new_deltas_async( Returns: bool: Whether any delta was appended or replaced, so the summary rebuilds. """ - if not entry.has_more and max_deltas is not None and len(entry.deltas) >= max_deltas: - return False - if not entry.has_more and max_deltas is None and terminal: + if max_deltas is not None and len(entry.deltas) >= max_deltas: return False changed = False while True: - if max_deltas is not None and len(entry.deltas) >= max_deltas: - return changed fetch_limit = ( min(self._STORAGE_PAGE_SIZE, max_deltas - len(entry.deltas)) if max_deltas is not None @@ -359,9 +355,11 @@ async def _hydrate_new_deltas_async( 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: - if terminal: - entry.has_more = False + entry.has_more = False return changed if not page: raise RuntimeError("Scenario progress storage returned an empty page with has_more=True.") @@ -1300,8 +1298,8 @@ def _build_node( total_retries=result.total_retries, ) ) - if result.retries: - retries_list = result.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] diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index 31c9324678..49fdae90ec 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -1772,6 +1772,8 @@ async def _recount_with_sdk_statistics_async( 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)) diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 5f80d20e87..8dfd57723b 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -7155,12 +7155,12 @@ def _build_scenario_history_ranked_units( *, entry_ids: Sequence[uuid.UUID], plan_entry_ids: Sequence[uuid.UUID], - ) -> tuple[Any, Any, Any, Any, Any]: + ) -> tuple[Any, Any, Any, Any, Any, Any]: """ Build the intermediate ranked-units subquery and role match conditions for scenario history. Returns: - tuple: (ranked, counted, is_target_facing, is_orchestration, is_unknown) + tuple: (ranked, counted, completed, is_target_facing, is_orchestration, is_unknown) """ atomic_name, technique_hash, attributed_seed_group_id, identifier_seed_key = ( self._get_scenario_attempt_unit_expressions() @@ -7299,7 +7299,7 @@ def _build_scenario_history_ranked_units( ranked.c.result_role.is_(None), and_(not_(is_target_facing), not_(is_orchestration)), ) - return ranked, counted, is_target_facing, is_orchestration, is_unknown + return ranked, counted, completed, is_target_facing, is_orchestration, is_unknown def _build_scenario_history_aggregate_statement( self, @@ -7313,8 +7313,8 @@ def _build_scenario_history_aggregate_statement( Returns: Any: A statement selecting one aggregate row per scenario run with attempts. """ - ranked, counted, is_target_facing, is_orchestration, is_unknown = self._build_scenario_history_ranked_units( - entry_ids=entry_ids, plan_entry_ids=plan_entry_ids + 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( @@ -7330,16 +7330,16 @@ def _build_scenario_history_aggregate_statement( "total_retries" ), func.sum( - case((and_(counted, ranked.c.latest_outcome == AttackOutcome.SUCCESS.value), 1), else_=0) + case((and_(completed, ranked.c.latest_outcome == AttackOutcome.SUCCESS.value), 1), else_=0) ).label("outcome_successes"), func.sum( - case((and_(counted, ranked.c.latest_outcome == AttackOutcome.FAILURE.value), 1), else_=0) + case((and_(completed, ranked.c.latest_outcome == AttackOutcome.FAILURE.value), 1), else_=0) ).label("outcome_failures"), - func.sum(case((and_(counted, ranked.c.latest_outcome == AttackOutcome.ERROR.value), 1), else_=0)).label( - "outcome_errors" - ), func.sum( - case((and_(counted, ranked.c.latest_outcome == AttackOutcome.UNDETERMINED.value), 1), else_=0) + 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" @@ -7389,32 +7389,32 @@ def _build_scenario_history_group_aggregate_statement( Returns: Any: A statement selecting aggregate rows grouped by scenario run and atomic group. """ - ranked, counted, is_target_facing, is_orchestration, is_unknown = self._build_scenario_history_ranked_units( - entry_ids=entry_ids, plan_entry_ids=plan_entry_ids + 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((counted, 1), else_=0)).label("completed_units"), + func.sum(case((completed, 1), else_=0)).label("completed_units"), func.sum( - case((and_(counted, ranked.c.latest_outcome == AttackOutcome.SUCCESS.value), 1), else_=0) + 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_(counted, ranked.c.latest_outcome == AttackOutcome.SUCCESS.value), 1), else_=0) + case((and_(completed, ranked.c.latest_outcome == AttackOutcome.SUCCESS.value), 1), else_=0) ).label("outcome_successes"), func.sum( - case((and_(counted, ranked.c.latest_outcome == AttackOutcome.FAILURE.value), 1), else_=0) + case((and_(completed, ranked.c.latest_outcome == AttackOutcome.FAILURE.value), 1), else_=0) ).label("outcome_failures"), - func.sum(case((and_(counted, ranked.c.latest_outcome == AttackOutcome.ERROR.value), 1), else_=0)).label( - "outcome_errors" - ), func.sum( - case((and_(counted, ranked.c.latest_outcome == AttackOutcome.UNDETERMINED.value), 1), else_=0) + 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( diff --git a/tests/unit/backend/test_scenario_hierarchy_parity.py b/tests/unit/backend/test_scenario_hierarchy_parity.py index dea7f1b537..ef05d9e1ae 100644 --- a/tests/unit/backend/test_scenario_hierarchy_parity.py +++ b/tests/unit/backend/test_scenario_hierarchy_parity.py @@ -972,7 +972,6 @@ def test_hierarchy_unplanned_groups_shared_counting() -> None: role=AttackResultRole.TARGET_FACING, parent_collection="unplanned_group", seed_group_id="seed_u", - attempt_index=1, seconds_offset=1, ) delta2 = _make_test_delta( @@ -981,7 +980,6 @@ def test_hierarchy_unplanned_groups_shared_counting() -> None: role=AttackResultRole.TARGET_FACING, parent_collection="unplanned_group", seed_group_id="seed_u", - attempt_index=2, seconds_offset=2, ) @@ -1154,7 +1152,6 @@ def test_hierarchy_repeated_attempt_diagnostic_retries_consistency() -> None: role=AttackResultRole.TARGET_FACING, parent_collection="group_1", seed_group_id="seed_1", - attempt_index=1, seconds_offset=1, ) delta2 = _make_test_delta( @@ -1163,7 +1160,6 @@ def test_hierarchy_repeated_attempt_diagnostic_retries_consistency() -> None: role=AttackResultRole.TARGET_FACING, parent_collection="group_1", seed_group_id="seed_1", - attempt_index=2, seconds_offset=2, )