From 0cdc24a9edfec3e1e9f35ba29c5bd3c3e56fe6aa Mon Sep 17 00:00:00 2001 From: Nimit Jain Date: Tue, 13 Oct 2026 10:36:11 +0530 Subject: [PATCH 01/22] 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/22] 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/22] 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/22] 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/22] 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/22] 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/22] 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/22] 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/22] 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/22] 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/22] 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/22] 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/22] 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/22] 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/22] 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/22] 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/22] 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/22] 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, ) From c639b998b2eb7ed4da41a41f017f01ae6f0047f0 Mon Sep 17 00:00:00 2001 From: Nimit Jain Date: Thu, 22 Oct 2026 11:04:49 +0530 Subject: [PATCH 19/22] fix: preserve child-error links in lightweight recount and align raw diagnostic retries --- .../backend/services/scenario_run_service.py | 2 +- pyrit/memory/memory_interface.py | 33 +- .../backend/test_scenario_hierarchy_parity.py | 402 ++++++++++++++++++ 3 files changed, 434 insertions(+), 3 deletions(-) diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index 49fdae90ec..af047f8bad 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -1773,7 +1773,7 @@ async def _recount_with_sdk_statistics_async( total_retries=summary.total_retries, result_role=res_metadata.result_role, attempt_index=res_metadata.attempt_index, - child_attack_result_ids=(), + child_attack_result_ids=summary.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 8dfd57723b..3f6b727176 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -444,6 +444,7 @@ class ScenarioAttemptSeedSummary: total_retries: int attribution_data: dict[str, Any] has_retries: bool = False + child_attack_result_ids: tuple[str, ...] = () @dataclass(frozen=True, slots=True, kw_only=True) @@ -7260,6 +7261,11 @@ def _build_scenario_history_ranked_units( + func.sum(pressure_attempt).over(partition_by=unit_partition) - 1 ) + raw_unit_retries = ( + func.sum(units.c.total_retries).over(partition_by=unit_partition) + + func.count().over(partition_by=unit_partition) + - 1 + ) ranked = select( units.c.scenario_result_id, units.c.unit_group_id, @@ -7269,6 +7275,7 @@ def _build_scenario_history_ranked_units( 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"), + raw_unit_retries.label("raw_unit_retries"), func.sum(case((is_error, pressure_attempt), else_=0)) .over(partition_by=unit_partition) .label("unit_errors"), @@ -7345,7 +7352,10 @@ def _build_scenario_history_aggregate_statement( "raw_total_errors" ), func.sum( - case((and_(ranked.c.unit_rank == 1, ranked.c.unit_retries > 0), ranked.c.unit_retries), else_=0) + case( + (and_(ranked.c.unit_rank == 1, ranked.c.raw_unit_retries > 0), ranked.c.raw_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"), @@ -7769,6 +7779,23 @@ def _row_to_scenario_attempt_seed_summary(row: Mapping[str, Any]) -> ScenarioAtt if row["atomic_attack_identifier"] else None ) + metadata = row.get("attack_metadata") + if isinstance(metadata, str): + try: + metadata = json.loads(metadata) + except Exception: + metadata = {} + metadata = metadata or {} + attr_data = dict(row.get("attribution_data") or {}) + for k in ("result_role", "attempt_index"): + if k not in attr_data and k in metadata: + attr_data[k] = metadata[k] + raw_child_ids = metadata.get("child_attack_result_ids") or attr_data.get("child_attack_result_ids") + if isinstance(raw_child_ids, (list, tuple)): + child_ids = tuple(str(cid) for cid in raw_child_ids if isinstance(cid, str) and cid) + else: + child_ids = () + return ScenarioAttemptSeedSummary( attack_result_id=str(row["id"]), objective=row["objective"], @@ -7777,8 +7804,9 @@ def _row_to_scenario_attempt_seed_summary(row: Mapping[str, Any]) -> ScenarioAtt outcome=AttackOutcome(row["outcome"]), timestamp=row["timestamp"], total_retries=row["total_retries"] or 0, - attribution_data=row["attribution_data"] or {}, + attribution_data=attr_data, has_retries=bool(row.get("has_retries", False)), + child_attack_result_ids=child_ids, ) def _execute_get_scenario_recount_summaries( @@ -7832,6 +7860,7 @@ def _execute_get_scenario_recount_summaries( AttackResultEntry.timestamp, AttackResultEntry.total_retries, AttackResultEntry.attribution_data, + AttackResultEntry.attack_metadata, case( ( or_( diff --git a/tests/unit/backend/test_scenario_hierarchy_parity.py b/tests/unit/backend/test_scenario_hierarchy_parity.py index ef05d9e1ae..8b2e7bb4b6 100644 --- a/tests/unit/backend/test_scenario_hierarchy_parity.py +++ b/tests/unit/backend/test_scenario_hierarchy_parity.py @@ -4760,3 +4760,405 @@ async def test_500_results_followed_by_retrying_attempt_diagnostic_count_parity( 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 + + +@pytest.mark.asyncio +async def test_legacy_identity_linked_child_error_recount_parity(sqlite_instance: SQLiteMemory) -> None: + """Test that a sequential unit with legacy identifier-only seed identity having two children + (failure followed by error) and an orchestration parent error referencing both children: + 1. SQL aggregate and SDK report 1 error and 1 retry. + 2. Lightweight recount preserves child-error links and reports 1 error and 1 retry. + 3. Cold one-node hierarchy with 500 preceding results reports 1 error and 1 retry, agreeing with warm hierarchy. + 4. Unlinked independent parent error control correctly reports 2 errors and 1 retry. + """ + base_time = datetime(2026, 1, 1, 0, 0, 0, tzinfo=UTC) + + # 1. Main test scenario with 500 preceding results followed by sequential unit + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="LegacyLinkedChildErrorScenario", + 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 legacy identifier-only seed identity + ordinary_results = [ + AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective=f"Objective {i}", + outcome=AttackOutcome.SUCCESS, + total_retries=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": "technique_legacy", + "result_role": AttackResultRole.TARGET_FACING.value, + }, + timestamp=base_time + timedelta(seconds=i), + ) + for i in range(500) + ] + + # Sequential unit with legacy identifier-only seed identity + seq_identifier = AtomicAttackIdentifier( + class_name="MockAttack", + seed_identifiers=[SeedIdentifier(class_name="SeedPrompt", value="sequential_seed_val")], + ) + child1_id = str(uuid.uuid4()) + child1 = AttackResult( + attack_result_id=child1_id, + conversation_id=str(uuid.uuid4()), + objective="Sequential Unit Objective", + outcome=AttackOutcome.FAILURE, + total_retries=0, + atomic_attack_identifier=seq_identifier, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "technique_legacy", + "result_role": AttackResultRole.TARGET_FACING.value, + "attempt_index": 1, + }, + timestamp=base_time + timedelta(seconds=500), + ) + + child2_id = str(uuid.uuid4()) + child2 = AttackResult( + attack_result_id=child2_id, + conversation_id=str(uuid.uuid4()), + objective="Sequential Unit Objective", + outcome=AttackOutcome.ERROR, + total_retries=0, + atomic_attack_identifier=seq_identifier, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "technique_legacy", + "result_role": AttackResultRole.TARGET_FACING.value, + "attempt_index": 2, + }, + timestamp=base_time + timedelta(seconds=501), + ) + + parent_id = str(uuid.uuid4()) + parent = AttackResult( + attack_result_id=parent_id, + conversation_id=str(uuid.uuid4()), + objective="Sequential Unit Objective", + outcome=AttackOutcome.ERROR, + total_retries=0, + atomic_attack_identifier=seq_identifier, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "technique_legacy", + "result_role": AttackResultRole.ORCHESTRATION.value, + }, + metadata={"child_attack_result_ids": [child1_id, child2_id]}, + timestamp=base_time + timedelta(seconds=502), + ) + + all_results = ordinary_results + [child1, child2, parent] + assert len(all_results) == 503 + await sqlite_instance.add_attack_results_to_memory_async(attack_results=all_results) + + # Verify SQL aggregate directly + aggregates = await sqlite_instance.get_scenario_history_aggregates_async(scenario_result_ids=[scenario_id]) + agg = aggregates[scenario_id] + assert agg.completed_units == 501 + assert agg.outcome_errors == 1 + assert agg.error_attempts == 1 + assert agg.total_retries == 1 + + service = ScenarioRunService(memory=sqlite_instance) + + # Verify lightweight recount via history list + run_list = await service.list_runs_async(limit=10) + matched_run = next(r for r in run_list.items if r.scenario_result_id == scenario_id) + assert matched_run.completed_attacks == 501 + + # Verify cold one-node hierarchy + 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.overall_counts.completed == 501 + assert cold_hier.overall_counts.errors == 1 + assert cold_hier.overall_counts.retries == 1 + + # Verify 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 + assert warm_hier.overall_counts.completed == 501 + assert warm_hier.overall_counts.errors == 1 + assert warm_hier.overall_counts.retries == 1 + + # Parity check between cold and warm + assert cold_hier.overall_counts.errors == warm_hier.overall_counts.errors == 1 + assert cold_hier.overall_counts.retries == warm_hier.overall_counts.retries == 1 + + # 2. Control: Unlinked independent parent error (where parent error does NOT wrap an already-failed child) + ctrl_uuid = uuid.uuid4() + ctrl_id = str(ctrl_uuid) + ctrl_scen = make_scenario_result( + id=ctrl_uuid, + scenario_name="UnlinkedParentErrorScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[ctrl_scen]) + + ctrl_child1_id = str(uuid.uuid4()) + ctrl_child1 = AttackResult( + attack_result_id=ctrl_child1_id, + conversation_id=str(uuid.uuid4()), + objective="Control Unit Objective", + outcome=AttackOutcome.FAILURE, + total_retries=0, + atomic_attack_identifier=seq_identifier, + attribution_parent_id=ctrl_id, + attribution_data={ + "parent_collection": "technique_legacy", + "result_role": AttackResultRole.TARGET_FACING.value, + "attempt_index": 1, + }, + timestamp=base_time, + ) + ctrl_child2_id = str(uuid.uuid4()) + ctrl_child2 = AttackResult( + attack_result_id=ctrl_child2_id, + conversation_id=str(uuid.uuid4()), + objective="Control Unit Objective", + outcome=AttackOutcome.FAILURE, + total_retries=0, + atomic_attack_identifier=seq_identifier, + attribution_parent_id=ctrl_id, + attribution_data={ + "parent_collection": "technique_legacy", + "result_role": AttackResultRole.TARGET_FACING.value, + "attempt_index": 2, + }, + timestamp=base_time + timedelta(seconds=1), + ) + # Parent error is independent because neither child is an ERROR + ctrl_parent = AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective="Control Unit Objective", + outcome=AttackOutcome.ERROR, + total_retries=0, + atomic_attack_identifier=seq_identifier, + attribution_parent_id=ctrl_id, + attribution_data={ + "parent_collection": "technique_legacy", + "result_role": AttackResultRole.ORCHESTRATION.value, + }, + metadata={"child_attack_result_ids": [ctrl_child1_id, ctrl_child2_id]}, + timestamp=base_time + timedelta(seconds=2), + ) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[ctrl_child1, ctrl_child2, ctrl_parent]) + + ctrl_aggregates = await sqlite_instance.get_scenario_history_aggregates_async(scenario_result_ids=[ctrl_id]) + ctrl_agg = ctrl_aggregates[ctrl_id] + assert ctrl_agg.outcome_errors == 1 + assert ctrl_agg.error_attempts == 1 + assert ctrl_agg.completed_units == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sequence_type", ["failed_sequence", "child_error_sequence"]) +async def test_clipped_cold_warm_hierarchy_diagnostic_retries_with_linked_children( + sqlite_instance: SQLiteMemory, + sequence_type: str, +) -> None: + """Test that 500 ordinary results followed by two target-facing children and their parent + (with explicit seed attribution and 0 persisted retries) reports identical diagnostics and progress + in both cold and warm hierarchy projections: + - overall_counts.retries == 1 (execution units exclude envelope) + - diagnostics.total_retries == 2 (raw diagnostic population includes all 3 saved results) + - 3-result control without preceding results also reports overall_counts.retries=1 and diagnostics.total_retries=2. + """ + base_time = datetime(2026, 1, 1, 0, 0, 0, tzinfo=UTC) + scenario_uuid = uuid.uuid4() + scenario_id = str(scenario_uuid) + + scen = make_scenario_result( + id=scenario_uuid, + scenario_name="ClippedDiagnosticRetriesScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) + + 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) + ] + + c1_outcome = AttackOutcome.FAILURE + c2_outcome = AttackOutcome.FAILURE if sequence_type == "failed_sequence" else AttackOutcome.ERROR + p_outcome = AttackOutcome.FAILURE if sequence_type == "failed_sequence" else AttackOutcome.ERROR + + child1_id = str(uuid.uuid4()) + child1 = AttackResult( + attack_result_id=child1_id, + conversation_id=str(uuid.uuid4()), + objective="Sequential Objective", + outcome=c1_outcome, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "attack_group", + "seed_group_id": "seed_sequential", + "result_role": AttackResultRole.TARGET_FACING.value, + "attempt_index": 1, + }, + timestamp=base_time + timedelta(seconds=500), + ) + + child2_id = str(uuid.uuid4()) + child2 = AttackResult( + attack_result_id=child2_id, + conversation_id=str(uuid.uuid4()), + objective="Sequential Objective", + outcome=c2_outcome, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "attack_group", + "seed_group_id": "seed_sequential", + "result_role": AttackResultRole.TARGET_FACING.value, + "attempt_index": 2, + }, + timestamp=base_time + timedelta(seconds=501), + ) + + parent = AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective="Sequential Objective", + outcome=p_outcome, + total_retries=0, + attribution_parent_id=scenario_id, + attribution_data={ + "parent_collection": "attack_group", + "seed_group_id": "seed_sequential", + "result_role": AttackResultRole.ORCHESTRATION.value, + }, + metadata={"child_attack_result_ids": [child1_id, child2_id]}, + timestamp=base_time + timedelta(seconds=502), + ) + + all_results = ordinary_results + [child1, child2, parent] + assert len(all_results) == 503 + await sqlite_instance.add_attack_results_to_memory_async(attack_results=all_results) + + service = ScenarioRunService(memory=sqlite_instance) + + # 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.overall_counts.retries == 1 + assert cold_hier.diagnostics.total_retries == 2 + + # Warm hierarchy 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 + assert warm_hier.overall_counts.retries == 1 + assert warm_hier.diagnostics.total_retries == 2 + + # Verify cold and warm agreement + assert cold_hier.overall_counts.retries == warm_hier.overall_counts.retries == 1 + assert cold_hier.diagnostics.total_retries == warm_hier.diagnostics.total_retries == 2 + + # 3-result control without preceding results + ctrl_uuid = uuid.uuid4() + ctrl_id = str(ctrl_uuid) + ctrl_scen = make_scenario_result( + id=ctrl_uuid, + scenario_name="ThreeResultControlScenario", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[ctrl_scen]) + + ctrl_c1_id = str(uuid.uuid4()) + ctrl_c1 = AttackResult( + attack_result_id=ctrl_c1_id, + conversation_id=str(uuid.uuid4()), + objective="Control Sequential Objective", + outcome=c1_outcome, + total_retries=0, + attribution_parent_id=ctrl_id, + attribution_data={ + "parent_collection": "attack_group", + "seed_group_id": "seed_ctrl_seq", + "result_role": AttackResultRole.TARGET_FACING.value, + "attempt_index": 1, + }, + timestamp=base_time, + ) + ctrl_c2_id = str(uuid.uuid4()) + ctrl_c2 = AttackResult( + attack_result_id=ctrl_c2_id, + conversation_id=str(uuid.uuid4()), + objective="Control Sequential Objective", + outcome=c2_outcome, + total_retries=0, + attribution_parent_id=ctrl_id, + attribution_data={ + "parent_collection": "attack_group", + "seed_group_id": "seed_ctrl_seq", + "result_role": AttackResultRole.TARGET_FACING.value, + "attempt_index": 2, + }, + timestamp=base_time + timedelta(seconds=1), + ) + ctrl_p = AttackResult( + attack_result_id=str(uuid.uuid4()), + conversation_id=str(uuid.uuid4()), + objective="Control Sequential Objective", + outcome=p_outcome, + total_retries=0, + attribution_parent_id=ctrl_id, + attribution_data={ + "parent_collection": "attack_group", + "seed_group_id": "seed_ctrl_seq", + "result_role": AttackResultRole.ORCHESTRATION.value, + }, + metadata={"child_attack_result_ids": [ctrl_c1_id, ctrl_c2_id]}, + timestamp=base_time + timedelta(seconds=2), + ) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[ctrl_c1, ctrl_c2, ctrl_p]) + + ctrl_cold = await service.get_run_hierarchy_async(scenario_result_id=ctrl_id, max_nodes=1) + assert ctrl_cold is not None + assert ctrl_cold.overall_counts.retries == 1 + assert ctrl_cold.diagnostics.total_retries == 2 + + await service.get_run_progress_async(scenario_result_id=ctrl_id, since=None, limit=10) + ctrl_warm = await service.get_run_hierarchy_async(scenario_result_id=ctrl_id, max_nodes=1) + assert ctrl_warm is not None + assert ctrl_warm.overall_counts.retries == 1 + assert ctrl_warm.diagnostics.total_retries == 2 From 6a5bd47ef39f627499f827bf602b7bef2b0b98b9 Mon Sep 17 00:00:00 2001 From: Roman Lutz Date: Sun, 25 Oct 2026 23:28:23 -0700 Subject: [PATCH 20/22] refactor: separate scenario hierarchy browsing from accounting Use native root pages and ordered child continuation without progress-cache hydration. Keep exact shared run statistics independent of clipping, and load selected saved conversations through canonical message-piece paging. Preserve existing unpaged attack APIs and explicit legacy previews. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 60a3baa0-0e4d-428b-8062-5101ad7866dd --- doc/gui/0_gui.md | 7 + pyrit/backend/README.md | 59 ++ pyrit/backend/mappers/attack_mappers.py | 18 +- pyrit/backend/models/attacks.py | 5 + pyrit/backend/routes/attacks.py | 28 +- pyrit/backend/routes/scenarios.py | 30 +- pyrit/backend/services/attack_service.py | 31 +- .../saved_conversation_read_service.py | 263 ++++++ .../services/scenario_progress_read_model.py | 409 +++----- .../backend/services/scenario_run_service.py | 411 ++++---- pyrit/common/pagination.py | 13 +- pyrit/memory/analytics_sql.py | 67 ++ pyrit/memory/memory_interface.py | 876 +++++++----------- pyrit/models/scenario_hierarchy.py | 8 + tests/unit/backend/test_attack_service.py | 4 +- .../test_saved_conversation_read_service.py | 189 ++++ .../test_scenario_hierarchy_pagination.py | 388 ++++++++ .../backend/test_scenario_hierarchy_parity.py | 155 ++-- .../test_scenario_progress_read_model.py | 8 +- tests/unit/common/test_pagination.py | 26 + tests/unit/memory/test_azure_sql_memory.py | 88 +- 21 files changed, 1950 insertions(+), 1133 deletions(-) create mode 100644 pyrit/backend/services/saved_conversation_read_service.py create mode 100644 tests/unit/backend/test_saved_conversation_read_service.py create mode 100644 tests/unit/backend/test_scenario_hierarchy_pagination.py diff --git a/doc/gui/0_gui.md b/doc/gui/0_gui.md index d7512b013a..25dcb6f1d7 100644 --- a/doc/gui/0_gui.md +++ b/doc/gui/0_gui.md @@ -646,6 +646,13 @@ In active runs and saved scenario results, **Atomic attack groups** defaults to Until you expand or collapse the section, its default follows the current group count as progress loads. Once you choose, the section keeps your choice during progress updates for the same run, even if the count crosses 20. Opening a different run resets to that run's count-based default. +Compound-attack hierarchy browsing is a separate GUI follow-up. The backend provides +paged roots and ordered child expansion, with exact run statistics independent of +the displayed page. Node reads return conversation metadata by default; clients +can page the selected saved transcript without loading every conversation. +See the [backend hierarchy API](../../pyrit/backend/README.md#scenario-result-hierarchy) +for continuation and clipping semantics. + ### Operations and Findings An operation groups related red-teaming work. Inside it you record findings: diff --git a/pyrit/backend/README.md b/pyrit/backend/README.md index d9fdeee66a..b1155bf945 100644 --- a/pyrit/backend/README.md +++ b/pyrit/backend/README.md @@ -48,6 +48,65 @@ Missing IDs, standalone conversations, and internal conversations return 404. Saving from the scanner creates a copy through `POST /api/attacks/save-conversation`; it does not change the source or send to a target. +For on-demand saved-result browsing, use +`GET /api/attacks/{attack_result_id}/conversations?limit=50` to page main, pruned, +adversarial, and preparation references in saved-reference order. This reads counts +and previews, not conversation bodies. The limit is 1-200; follow +`pagination.next_cursor` while `pagination.has_more` is true. + +Load only the selected conversation with +`GET /api/attacks/{attack_result_id}/messages?conversation_id={conversation_id}&limit=100`. +Here the limit is 1-500 **message pieces**, not turns or bytes. Results use the +canonical message, piece, and score views, ordered by sequence, timestamp, and ID. +A page can split a multimodal turn: `partial_message_sequences` identifies partial +turns. Append pieces by ID to their turn across pages, and follow +`pagination.next_cursor`. A page does not infer a final target-response status +from an incomplete transcript. + +Message cursors belong to one attack and conversation. Reference cursors also +validate the saved ordering and timestamp; restart reference paging if the saved +attack changes. Invalid or mismatched cursors return 400. Passing a cursor requires +a limit. Omitting both preserves the existing active-conversation listing and +full-transcript behavior. + +### Scenario result hierarchy + +These read-only APIs provide the backend foundation for compound-attack browsing. +The hierarchy GUI is separate work; the current scanner still shows flat executions. + +`GET /api/scenarios/runs/{scenario_result_id}/hierarchy?max_nodes=1000` returns a +bounded tree preview grouped by atomic attack group. Roots are saved results with +no incoming child reference from the same run, not results inferred from their +attempt indexes. The node budget is 1-5000, with at most 500 roots per page and +bounded descendant queries through depth 10. `has_more` and `next_cursor` refer +only to **root continuation**. `truncated` also covers deferred descendants, +depth limits, and other results not represented by the preview. + +Run/group completion, outcomes, producer counts, and raw error/retry diagnostics +use the shared exact accounting path independently of the selected page. Exact +statistics may scan lightweight summaries; a small tree request does not hydrate +the full progress snapshot, attack results, or conversation bodies. + +Expand one result with +`GET /api/scenarios/runs/{scenario_result_id}/hierarchy/nodes/{attack_result_id}?limit=200`. +The limit is 1-500 saved child-reference positions. Follow `next_children_cursor` +using `children_cursor`; `children_count` is the full reference count and +`child_offset` locates this page. `child_index` remains the absolute position under +the immediate parent. Missing, malformed, duplicate, cyclic, or cross-run references +remain explicit and consume positions, so even a wholly unresolved page can advance. +Child cursors validate the run, node, timestamp, and saved reference ordering; +restart child paging after those references change. + +Node expansion returns conversation metadata only by default. Use the generic +saved-conversation APIs above for continuation through any selected transcript. +`include_messages=true` explicitly opts into the older clipped conversation +previews, not a complete transcript. + +An entirely rootless saved graph, such as a closed cycle, returns a bounded +`diagnostics.rootless_result_ids` sample with its total count and clipping flag. +It does not invent lineage or promote a cycle member to a true root. The sampled +IDs remain available through the single-node and saved-conversation APIs. + ### Techniques | Endpoint | Result | diff --git a/pyrit/backend/mappers/attack_mappers.py b/pyrit/backend/mappers/attack_mappers.py index 4d35e78113..e2cbac4adb 100644 --- a/pyrit/backend/mappers/attack_mappers.py +++ b/pyrit/backend/mappers/attack_mappers.py @@ -11,6 +11,7 @@ from __future__ import annotations +import asyncio import logging import time import uuid @@ -53,6 +54,7 @@ _LEGACY_MANUAL_ATTACK_PLACEHOLDER = "Manual attack via GUI" if TYPE_CHECKING: + from pyrit.memory import MemoryInterface from pyrit.models.conversation_stats import ConversationStats # ============================================================================ @@ -326,7 +328,9 @@ async def _resolve_and_sign_media_async(*, value: str | None, data_type: str) -> Returns: The resolved (and signed, if a blob) URL, or ``None`` when *value* is empty. """ - resolved = _resolve_media_url(value=value, data_type=data_type) + if not value or data_type not in MEDIA_PATH_DATA_TYPES: + return None + resolved = await asyncio.to_thread(_resolve_media_url, value=value, data_type=data_type) if resolved and _is_azure_blob_url(resolved): return await _sign_blob_url_async(blob_url=resolved) return resolved @@ -368,6 +372,7 @@ def _score_lookup_key(*, piece: MessagePiece) -> str: async def _fetch_scores_by_piece_async( *, pyrit_messages: list[Message], + memory: MemoryInterface | None = None, ) -> dict[str, list[Score]]: """ Batch-fetch scores for every piece in ``pyrit_messages`` and group by piece id. @@ -382,8 +387,8 @@ async def _fetch_scores_by_piece_async( if not score_lookup_ids: return {} - memory = CentralMemory.get_memory_instance() - fetched = await memory.get_prompt_scores_async(prompt_ids=score_lookup_ids) + source = memory if memory is not None else CentralMemory.get_memory_instance() + fetched = await source.get_prompt_scores_async(prompt_ids=score_lookup_ids) grouped: dict[str, list[Score]] = {} for score in fetched: @@ -395,6 +400,7 @@ async def pyrit_messages_to_dto_async( pyrit_messages: list[Message], *, objective_score_id: uuid.UUID | str | None = None, + memory: MemoryInterface | None = None, ) -> list[MessageView]: """ Translate PyRIT messages to backend MessageView responses. @@ -406,15 +412,15 @@ async def pyrit_messages_to_dto_async( - Local files -> ``/api/media?path=...`` (served by the media endpoint) - Azure Blob Storage files -> signed URLs with SAS tokens - Scores are fetched from ``CentralMemory`` (``MessagePiece`` no longer carries - them) via a single batched ``get_prompt_scores`` call and attached to their + Scores are fetched from the supplied memory, or ``CentralMemory`` by default + (``MessagePiece`` no longer carries them), in one batched read and attached to their originating piece. When ``objective_score_id`` is provided, the matching score is marked as the attack's canonical objective score. Returns: List of MessageView responses for the API. """ - scores_by_piece = await _fetch_scores_by_piece_async(pyrit_messages=pyrit_messages) + scores_by_piece = await _fetch_scores_by_piece_async(pyrit_messages=pyrit_messages, memory=memory) messages: list[MessageView] = [] for msg in pyrit_messages: diff --git a/pyrit/backend/models/attacks.py b/pyrit/backend/models/attacks.py index a5fa4f90d5..3ed55acef4 100644 --- a/pyrit/backend/models/attacks.py +++ b/pyrit/backend/models/attacks.py @@ -342,6 +342,8 @@ class ConversationMessagesResponse(BaseModel): "None when the conversation does not end with a target response." ), ) + pagination: PaginationInfo | None = None + partial_message_sequences: list[int] = Field(default_factory=list) # ============================================================================ @@ -568,6 +570,7 @@ class ConversationSummary(BaseModel): message_count: int = Field(0, description="Number of messages in this conversation") last_message_preview: str | None = Field(None, description="Preview of the last message") created_at: datetime | None = Field(None, description="Timestamp of the first message") + conversation_type: str | None = None class AttackConversationsResponse(BaseModel): @@ -578,6 +581,8 @@ class AttackConversationsResponse(BaseModel): conversations: list[ConversationSummary] = Field( default_factory=list, description="All conversations including main" ) + pagination: PaginationInfo | None = None + total_count: int | None = None class CreateConversationRequest(BaseModel): diff --git a/pyrit/backend/routes/attacks.py b/pyrit/backend/routes/attacks.py index 4af5b210ae..9cce3b5d1c 100644 --- a/pyrit/backend/routes/attacks.py +++ b/pyrit/backend/routes/attacks.py @@ -387,11 +387,14 @@ async def get_conversation_messages( # pyrit-async-suffix-exempt conversation_id: str = Query( ..., max_length=MAX_IDENTIFIER_LENGTH, description="The conversation_id whose messages to return" ), + limit: int | None = Query(None, ge=1, le=500, description="Optional saved message-piece page size"), + cursor: CursorStr | None = Query(None, description="Opaque message-piece continuation cursor"), ) -> ConversationMessagesResponse: """ - Get all messages for a conversation belonging to an attack. + Get saved messages for a conversation belonging to an attack. - Returns prepended conversation and all messages in order. + Omitting limit preserves the existing full-conversation response. Paged reads can + split a multimodal turn; append pieces by their IDs and turn number across pages. Returns: ConversationMessagesResponse: All messages for the conversation. @@ -402,6 +405,8 @@ async def get_conversation_messages( # pyrit-async-suffix-exempt messages = await service.get_conversation_messages_async( attack_result_id=attack_result_id, conversation_id=conversation_id, + limit=limit, + cursor=cursor, ) except ValueError as e: raise HTTPException( @@ -422,24 +427,31 @@ async def get_conversation_messages( # pyrit-async-suffix-exempt "/{attack_result_id}/conversations", response_model=AttackConversationsResponse, responses={ + 400: {"model": ProblemDetail, "description": "Invalid conversation cursor"}, 404: {"model": ProblemDetail, "description": "Attack not found"}, }, ) -async def get_conversations( +async def get_conversations( # pyrit-async-suffix-exempt attack_result_id: IdentifierStr, -) -> AttackConversationsResponse: # pyrit-async-suffix-exempt + limit: int | None = Query(None, ge=1, le=200, description="Optional saved conversation-reference page size"), + cursor: CursorStr | None = Query(None, description="Opaque conversation-reference continuation cursor"), +) -> AttackConversationsResponse: """ - Get all conversations belonging to an attack. + Get saved conversations belonging to an attack. - Returns the main conversation and all related conversations with - message counts and preview text. + Without a limit, preserve the existing active-conversation list and timestamp + ordering. Paged reads include main, pruned, adversarial, and preparation contexts + in stored-reference order, with message counts and preview text. Returns: AttackConversationsResponse: All conversations for the attack. """ service = get_attack_service() - result = await service.get_conversations_async(attack_result_id=attack_result_id) + try: + result = await service.get_conversations_async(attack_result_id=attack_result_id, limit=limit, cursor=cursor) + except ValueError as exc: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from exc if not result: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, diff --git a/pyrit/backend/routes/scenarios.py b/pyrit/backend/routes/scenarios.py index a573bf0d43..cf3fea8479 100644 --- a/pyrit/backend/routes/scenarios.py +++ b/pyrit/backend/routes/scenarios.py @@ -386,6 +386,7 @@ async def get_scenario_run_progress( # pyrit-async-suffix-exempt "/runs/{scenario_result_id}/hierarchy", response_model=ScenarioHierarchyProjection, responses={ + 400: {"model": ProblemDetail, "description": "Invalid hierarchy cursor"}, 404: {"model": ProblemDetail, "description": "Run not found"}, }, ) @@ -393,6 +394,7 @@ 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"), + cursor: CursorStr | None = Query(None, description="Opaque root-page continuation cursor"), ) -> ScenarioHierarchyProjection: """ Get the bounded result hierarchy projection for a scenario run. @@ -402,10 +404,12 @@ 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, - max_nodes=max_nodes, - ) + try: + hierarchy = await service.get_run_hierarchy_async( + scenario_result_id=scenario_result_id, max_nodes=max_nodes, cursor=cursor + ) + except ValueError as exc: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from exc if hierarchy is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -418,6 +422,7 @@ async def get_scenario_run_hierarchy( # pyrit-async-suffix-exempt "/runs/{scenario_result_id}/hierarchy/nodes/{attack_result_id}", response_model=ScenarioHierarchyNode, responses={ + 400: {"model": ProblemDetail, "description": "Invalid child cursor"}, 404: {"model": ProblemDetail, "description": "Node or run not found"}, }, ) @@ -425,6 +430,9 @@ async def get_scenario_run_hierarchy_node( # pyrit-async-suffix-exempt *, scenario_result_id: IdentifierStr, attack_result_id: IdentifierStr, + limit: int = Query(200, ge=1, le=500, description="Maximum ordered child references in this page"), + children_cursor: CursorStr | None = Query(None, description="Opaque ordered-child continuation cursor"), + include_messages: bool = Query(False, description="Explicitly include legacy clipped conversation previews"), ) -> ScenarioHierarchyNode: """ Get single-node hierarchy details and generic conversation projections without full-run hydration. @@ -433,10 +441,16 @@ async def get_scenario_run_hierarchy_node( # pyrit-async-suffix-exempt 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, - ) + try: + node = await service.get_run_hierarchy_node_async( + scenario_result_id=scenario_result_id, + attack_result_id=attack_result_id, + max_immediate_children=limit, + children_cursor=children_cursor, + include_messages=include_messages, + ) + except ValueError as exc: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from exc if node is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, diff --git a/pyrit/backend/services/attack_service.py b/pyrit/backend/services/attack_service.py index 9509c23078..ecf70dc178 100644 --- a/pyrit/backend/services/attack_service.py +++ b/pyrit/backend/services/attack_service.py @@ -68,12 +68,14 @@ fingerprint_filters, normalize_label_filters, ) +from pyrit.backend.services.saved_conversation_read_service import SavedConversationReadService from pyrit.backend.services.target_service import get_target_service from pyrit.common.utils import to_sha256 from pyrit.executor.attack import attack_outcome_from_score from pyrit.memory import ( AttackResultKeysetCursor, CentralMemory, + MemoryInterface, data_serializer_factory, set_message_piece_sha256_async, ) @@ -132,9 +134,12 @@ class AttackService: Uses PyRIT memory (database) as the source of truth via AttackResult. """ - def __init__(self, *, message_send_service: MessageSendService | None = None) -> None: + def __init__( + self, *, message_send_service: MessageSendService | None = None, memory: MemoryInterface | None = None + ) -> None: """Initialize the attack service.""" - self._memory = CentralMemory.get_memory_instance() + self._memory = memory if memory is not None else CentralMemory.get_memory_instance() + self._saved_conversation_reader = SavedConversationReadService(memory=self._memory) self._message_send_service = ( message_send_service if message_send_service is not None else get_message_send_service() ) @@ -383,6 +388,8 @@ async def get_conversation_messages_async( *, attack_result_id: str, conversation_id: str, + limit: int | None = None, + cursor: str | None = None, ) -> ConversationMessagesResponse | None: """ Get all messages for a conversation belonging to an attack. @@ -390,6 +397,8 @@ async def get_conversation_messages_async( Args: attack_result_id: The AttackResult's primary key (used to verify existence). conversation_id: The conversation whose messages to return. + limit: Optional message-piece page size. None preserves the existing full-conversation read. + cursor: Continuation for an explicitly paged read. Returns: ConversationMessagesResponse if attack found, None otherwise. @@ -397,6 +406,13 @@ async def get_conversation_messages_async( Raises: ValueError: If the conversation does not belong to the attack. """ + if limit is not None: + return await self._saved_conversation_reader.get_messages_async( + attack_result_id=attack_result_id, conversation_id=conversation_id, limit=limit, cursor=cursor + ) + if cursor is not None: + raise ValueError("A message-page cursor requires a limit.") + # Check attack exists results = await self._memory.get_attack_results_async(attack_result_ids=[attack_result_id]) if not results: @@ -412,6 +428,7 @@ async def get_conversation_messages_async( backend_messages = await pyrit_messages_to_dto_async( list(pyrit_messages), objective_score_id=ar.last_score.id if ar.last_score else None, + memory=self._memory, ) return ConversationMessagesResponse( @@ -1076,7 +1093,9 @@ async def remove_human_score_async(self, *, attack_result_id: str) -> AttackSumm return await self.get_attack_async(attack_result_id=attack_result_id) - async def get_conversations_async(self, *, attack_result_id: str) -> AttackConversationsResponse | None: + async def get_conversations_async( + self, *, attack_result_id: str, limit: int | None = None, cursor: str | None = None + ) -> AttackConversationsResponse | None: """ Get all conversations belonging to an attack. @@ -1087,6 +1106,12 @@ async def get_conversations_async(self, *, attack_result_id: str) -> AttackConve Returns: AttackConversationsResponse if attack found, None otherwise. """ + if limit is not None: + return await self._saved_conversation_reader.list_conversations_async( + attack_result_id=attack_result_id, limit=limit, cursor=cursor + ) + if cursor is not None: + raise ValueError("A conversation-page cursor requires a limit.") results = await self._memory.get_attack_results_async(attack_result_ids=[attack_result_id]) if not results: return None diff --git a/pyrit/backend/services/saved_conversation_read_service.py b/pyrit/backend/services/saved_conversation_read_service.py new file mode 100644 index 0000000000..c32e6fa75f --- /dev/null +++ b/pyrit/backend/services/saved_conversation_read_service.py @@ -0,0 +1,263 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Bounded, read-only saved conversation views shared by attack and scenario APIs.""" + +import uuid +from typing import ClassVar + +from pyrit.backend.mappers._preview import format_last_message_preview +from pyrit.backend.mappers.attack_mappers import pyrit_messages_to_dto_async +from pyrit.backend.models.attacks import ( + AttackConversationsResponse, + ConversationMessagesResponse, + ConversationSummary, +) +from pyrit.backend.models.common import PaginationInfo +from pyrit.common.pagination import decode_keyset_cursor, encode_keyset_cursor, fingerprint_filters +from pyrit.memory import MemoryInterface +from pyrit.memory.memory_interface import MessagePieceKeysetCursor +from pyrit.models import ( + ConversationStats, + ConversationType, + ScenarioAttackResultConversations, + ScenarioConversationMessagePieceProjection, + ScenarioConversationMessageProjection, + ScenarioConversationProjection, +) + + +class SavedConversationReadService: + """Read canonical saved references/messages without constructing an executable attack.""" + + _LEGACY_DETAILS: ClassVar[dict[ConversationType | None, tuple[str, str]]] = { + None: ("main", "Main objective target conversation"), + ConversationType.PRUNED: ("pruned", "Pruned conversation"), + ConversationType.ADVERSARIAL: ("adversarial_chat", "Adversarial chat conversation"), + ConversationType.PREPARATION: ("preparation", "Preparation conversation"), + } + + def __init__(self, *, memory: MemoryInterface) -> None: + """Bind all saved reads to the caller's memory instance.""" + self._memory = memory + + async def list_conversations_async( + self, *, attack_result_id: str, limit: int = 50, cursor: str | None = None + ) -> AttackConversationsResponse | None: + """ + Page saved main/related conversations without loading their message bodies. + + Returns: + AttackConversationsResponse | None: A reference page, or None for an absent attack. + + Raises: + ValueError: If the cursor belongs to another attack or an obsolete reference ordering. + """ + page = await self._memory.get_attack_conversation_page_async(attack_result_id=attack_result_id, limit=limit) + if page is None: + return None + fingerprint = fingerprint_filters( + filters={ + "kind": "saved-conversations", + "attack_result_id": str(uuid.UUID(attack_result_id)), + "revision": page.reference_revision, + } + ) + anchor = decode_keyset_cursor(cursor=cursor, fingerprint=fingerprint) + if cursor is not None and ( + anchor is None + or anchor.position is None + or anchor.position < 0 + or anchor.position >= page.total_count + or anchor.identifier != str(uuid.UUID(attack_result_id)) + or anchor.timestamp != page.timestamp + ): + raise ValueError("Invalid or stale saved-conversation cursor.") + offset = anchor.position if anchor is not None and anchor.position is not None else 0 + if offset: + revision = page.reference_revision + page = await self._memory.get_attack_conversation_page_async( + attack_result_id=attack_result_id, offset=offset, limit=limit + ) + if page is None: + return None + if page.reference_revision != revision: + raise ValueError("Saved conversation references changed; restart pagination.") + stats = ( + await self._memory.get_conversation_stats_async( + conversation_ids=[reference.conversation_id for reference in page.conversations] + ) + if page.conversations + else {} + ) + conversations = [] + for reference in page.conversations: + stat = stats.get(reference.conversation_id, ConversationStats(message_count=0)) + conversations.append( + ConversationSummary( + conversation_id=reference.conversation_id, + conversation_type=reference.conversation_type.value if reference.conversation_type else "main", + message_count=stat.message_count, + last_message_preview=format_last_message_preview( + value=stat.last_message_preview, data_type=stat.last_message_data_type + ), + created_at=stat.created_at, + ) + ) + next_offset = offset + len(conversations) + has_more = next_offset < page.total_count + next_cursor = ( + encode_keyset_cursor( + timestamp=page.timestamp, + identifier=attack_result_id, + fingerprint=fingerprint, + position=next_offset, + ) + if has_more + else None + ) + return AttackConversationsResponse( + attack_result_id=attack_result_id, + main_conversation_id=page.main_conversation_id, + conversations=conversations, + total_count=page.total_count, + pagination=PaginationInfo(limit=limit, has_more=has_more, next_cursor=next_cursor), + ) + + async def get_messages_async( + self, *, attack_result_id: str, conversation_id: str, limit: int = 100, cursor: str | None = None + ) -> ConversationMessagesResponse | None: + """ + Page canonical messages for one explicitly selected saved conversation. + + Returns: + ConversationMessagesResponse | None: A piece-bounded page, or None for an absent attack. + + Raises: + ValueError: If the conversation is unrelated or the continuation cursor is invalid. + """ + references = await self._memory.get_attack_conversation_page_async( + attack_result_id=attack_result_id, conversation_id=conversation_id, limit=1 + ) + if references is None: + return None + if not references.conversations: + raise ValueError(f"Conversation '{conversation_id}' is not part of attack '{attack_result_id}'") + fingerprint = fingerprint_filters( + filters={ + "kind": "saved-message-pieces", + "attack_result_id": str(uuid.UUID(attack_result_id)), + "conversation_id": conversation_id, + } + ) + anchor = decode_keyset_cursor(cursor=cursor, fingerprint=fingerprint) + if cursor is not None and (anchor is None or anchor.position is None): + raise ValueError("Invalid saved-message cursor.") + piece_cursor = ( + MessagePieceKeysetCursor( + sequence=anchor.position, timestamp=anchor.timestamp, message_piece_id=anchor.identifier + ) + if anchor is not None and anchor.position is not None + else None + ) + page = await self._memory.get_conversation_message_page_async( + conversation_id=conversation_id, cursor=piece_cursor, limit=limit + ) + next_cursor = ( + encode_keyset_cursor( + timestamp=page.next_cursor.timestamp, + identifier=page.next_cursor.message_piece_id, + position=page.next_cursor.sequence, + fingerprint=fingerprint, + ) + if page.next_cursor is not None + else None + ) + return ConversationMessagesResponse( + conversation_id=conversation_id, + messages=await pyrit_messages_to_dto_async( + list(page.messages), memory=self._memory, objective_score_id=references.objective_score_id + ), + pagination=PaginationInfo(limit=limit, has_more=next_cursor is not None, next_cursor=next_cursor), + partial_message_sequences=list(page.partial_message_sequences), + ) + + async def get_legacy_projections_async( + self, + *, + attack_result_id: str, + include_messages: bool = False, + limit_messages_per_conversation: int = 100, + max_conversations: int = 50, + max_pieces_per_message: int = 100, + ) -> ScenarioAttackResultConversations: + """ + Adapt canonical saved records to the existing clipped hierarchy wire format. + + New clients load a selected conversation through ``get_messages_async`` instead. + + Returns: + ScenarioAttackResultConversations: Compatible bounded metadata and optional turn previews. + """ + page = await self._memory.get_attack_conversation_page_async( + attack_result_id=attack_result_id, limit=max_conversations + ) + if page is None: + return ScenarioAttackResultConversations(conversations=[], total_count=0, truncated=False) + stats = ( + await self._memory.get_conversation_stats_async( + conversation_ids=[reference.conversation_id for reference in page.conversations] + ) + if page.conversations + else {} + ) + conversations: list[ScenarioConversationProjection] = [] + for reference in page.conversations: + messages: list[ScenarioConversationMessageProjection] = [] + message_count = stats.get(reference.conversation_id, ConversationStats(message_count=0)).message_count + truncated = message_count > 0 + if include_messages: + ( + saved_messages, + truncated, + partial_sequences, + ) = await self._memory.get_bounded_conversation_messages_async( + conversation_id=reference.conversation_id, + limit=limit_messages_per_conversation, + max_pieces_per_message=max_pieces_per_message, + ) + for message in saved_messages: + first = message.message_pieces[0] + messages.append( + ScenarioConversationMessageProjection( + role=first.role, + content=message.get_value(), + sequence=first.sequence, + timestamp=first.timestamp, + pieces=[ + ScenarioConversationMessagePieceProjection( + data_type=piece.converted_value_data_type, + content=piece.original_value, + converted_value=piece.converted_value, + ) + for piece in message.message_pieces + ], + pieces_truncated=first.sequence in partial_sequences, + ) + ) + kind, description = self._LEGACY_DETAILS[reference.conversation_type] + conversations.append( + ScenarioConversationProjection( + conversation_id=reference.conversation_id, + conversation_type=kind, + description=description, + message_count=message_count, + messages=messages, + messages_truncated=truncated, + ) + ) + return ScenarioAttackResultConversations( + conversations=conversations, + total_count=page.total_count, + truncated=page.total_count > len(conversations), + ) diff --git a/pyrit/backend/services/scenario_progress_read_model.py b/pyrit/backend/services/scenario_progress_read_model.py index a3e027bc0b..742d071275 100644 --- a/pyrit/backend/services/scenario_progress_read_model.py +++ b/pyrit/backend/services/scenario_progress_read_model.py @@ -5,6 +5,7 @@ import asyncio import logging +import uuid from collections import OrderedDict from collections.abc import AsyncGenerator, Iterable, Mapping, Sequence from contextlib import asynccontextmanager @@ -864,9 +865,18 @@ def _read_child_attack_result_ids(*, attack_metadata: dict[str, Any]) -> list[st list[str]: The child IDs in stored order, or an empty list when none are recorded. """ child_ids = attack_metadata.get("child_attack_result_ids") - if isinstance(child_ids, list) and all(isinstance(child_id, str) for child_id in child_ids): - return list(child_ids) - return [] + if not isinstance(child_ids, list): + return [] + identifiers: list[str] = [] + for index, child_id in enumerate(child_ids): + if not isinstance(child_id, str): + identifiers.append(f"") + continue + try: + identifiers.append(str(uuid.UUID(child_id.lower()))) + except ValueError: + identifiers.append(child_id) + return identifiers @staticmethod def _synthesize_legacy_plan(*, deltas: list[ScenarioAttackResultDelta]) -> ScenarioRunPlan: @@ -924,8 +934,7 @@ def build_scenario_hierarchy( exact_total_retries: int | None = None, exact_retrying_attacks_count: int | None = None, exact_group_counts: Mapping[str, ScenarioProgressCounts] | None = None, - parent_deltas: Sequence[ScenarioAttackResultDelta] = (), - unverified_candidate_ids: set[str] | Sequence[str] | None = None, + root_result_ids: Sequence[str] | None = None, ) -> ScenarioHierarchyProjection: """ Assemble the scenario result hierarchy projection with cycle, depth, and node budget guards. @@ -943,10 +952,7 @@ def build_scenario_hierarchy( exact_total_retries: Optional exact run-wide retry count from database aggregates. exact_retrying_attacks_count: Optional exact run-wide count of retrying attacks from aggregates. exact_group_counts: Optional mapping of atomic group ID to exact counts from aggregates. - parent_deltas: Optional orchestration parent deltas to resolve roots and child references - beyond delta limits. - unverified_candidate_ids: Optional candidates whose parent status could not be verified - due to query clipping or budget exhaustion; excluded from root_results. + root_result_ids: Optional authoritative, run-scoped root page returned by storage. Returns: ScenarioHierarchyProjection: Bounded hierarchy projection with diagnostics. @@ -954,31 +960,16 @@ 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 results_by_id.values(): for child_id in result.child_attack_result_ids: referenced_child_ids.add(child_id) - # A true root is not referenced as a child by any parent and does not carry an attempt_index - # (which designates a child attempt dispatched by an orchestration parent). - unverified_set = set(unverified_candidate_ids) if unverified_candidate_ids else set() - root_results = [ - r - for r in results_by_id.values() - if r.attack_result_id not in referenced_child_ids - and r.attempt_index is None - and r.attack_result_id not in unverified_set - ] + root_results = [r for r in results_by_id.values() if r.attack_result_id not in referenced_child_ids] root_results.sort(key=lambda r: (r.timestamp, r.attack_result_id)) + if root_result_ids is not None: + root_results = [results_by_id[identifier] for identifier in root_result_ids if identifier in results_by_id] + reserved_root_ids = {result.attack_result_id for result in root_results} visited_ids: set[str] = set() total_nodes_count = 0 @@ -1005,56 +996,28 @@ def _build_node( total_nodes_count += 1 visited_ids.add(attack_result_id) - retries = result.retries - retries_truncated = False - if len(retries) > _MAX_RETRIES_PER_NODE: - retries = retries[:_MAX_RETRIES_PER_NODE] - retries_truncated = True - - executed_turns = _resolve_executed_turns(delta) - outcome_reason = ( - delta.outcome_reason - or delta.attribution_data.get("outcome_reason") - or delta.attack_metadata.get("outcome_reason") + node = ScenarioProgressReadModel._map_hierarchy_node( + delta=delta, result=result, parent_id=parent_id, child_index=child_index ) - - raw_child_ids = list(result.child_attack_result_ids) - child_refs_truncated = len(raw_child_ids) > _MAX_CHILD_REFERENCES - bounded_child_ids = raw_child_ids[:_MAX_CHILD_REFERENCES] if child_refs_truncated else raw_child_ids - - orchestration: ScenarioOrchestrationParentSummary | None = None - if result.result_role == AttackResultRole.ORCHESTRATION or result.child_attack_result_ids: - completion_policy = delta.attack_metadata.get("completion_policy") - orchestration = ScenarioOrchestrationParentSummary( - attack_result_id=result.attack_result_id, - result_role=result.result_role, - outcome=result.outcome, - outcome_reason=outcome_reason, - attempt_index=result.attempt_index, - execution_time_ms=result.execution_time_ms, - executed_turns=executed_turns, - completion_policy=str(completion_policy) if completion_policy is not None else None, - child_attack_result_ids=bounded_child_ids, - child_references_truncated=child_refs_truncated, - error_type=result.error_type, - error_message=result.error_message, - timestamp=result.timestamp, - total_retries=result.total_retries, - retries=retries, - retries_truncated=retries_truncated, - ) - current_active_path = active_path | {attack_result_id} children: list[ScenarioHierarchyNode] = [] unresolved_children: list[str] = [] + seen_child_ids: set[str] = set() + node.children_truncated = node.child_references_truncated + if node.child_references_truncated: + hierarchy_truncated = True - 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: + for idx, c_id in enumerate(node.child_attack_result_ids): + if c_id in current_active_path or c_id in seen_child_ids or c_id not in deltas_by_id: unresolved_children.append(c_id) - elif total_nodes_count >= max_nodes: + if c_id not in deltas_by_id and snapshot.has_more: + node.children_truncated = True + elif depth >= _MAX_HIERARCHY_DEPTH or total_nodes_count >= max_nodes - len(reserved_root_ids): + node.children_truncated = True hierarchy_truncated = True unresolved_children.append(c_id) else: + seen_child_ids.add(c_id) child_node = _build_node( attack_result_id=c_id, parent_id=attack_result_id, @@ -1110,47 +1073,19 @@ 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 "", - parent_attack_result_id=parent_id, - child_index=child_index, - result_role=result.result_role, - attempt_index=result.attempt_index, - atomic_attack_name=result.atomic_attack_name, - atomic_group_id=result.atomic_group_id, - seed_group_id=result.seed_group_id, - objective=delta.objective, - outcome=result.outcome, - outcome_reason=outcome_reason, - execution_time_ms=result.execution_time_ms, - executed_turns=executed_turns, - timestamp=result.timestamp, - score=result.score, - error_type=result.error_type, - error_message=result.error_message, - total_retries=result.total_retries, - retries=retries, - retries_truncated=retries_truncated, - child_attack_result_ids=bounded_child_ids, - child_references_truncated=child_refs_truncated, - children=children, - children_truncated=False, - unresolved_child_attack_result_ids=bounded_unresolved_children, - unresolved_child_references_truncated=unresolved_refs_truncated, - conversations=conversations, - conversations_count=node_conversations_count, - conversations_truncated=node_conversations_truncated, - orchestration=orchestration, + node.children = children + node.unresolved_child_attack_result_ids = unresolved_children[:_MAX_CHILD_REFERENCES] + node.unresolved_child_references_truncated = ( + len(unresolved_children) > _MAX_CHILD_REFERENCES or node.child_references_truncated ) + node.conversations = conversations + node.conversations_count = node_conversations_count + node.conversations_truncated = node_conversations_truncated + return node roots_by_group: dict[str, list[ScenarioHierarchyNode]] = {} for root_result in root_results: + reserved_root_ids.discard(root_result.attack_result_id) if total_nodes_count >= max_nodes: hierarchy_truncated = True break @@ -1165,24 +1100,8 @@ def _build_node( roots_by_group.setdefault(root_node.atomic_group_id, []).append(root_node) # Handle unreached/disconnected components (including detached cycles) in a single linear pass - if total_nodes_count < max_nodes: - # Nodes that are unverified or descended from unverified candidates can never be emitted as roots - unverified_descendants: set[str] = set(unverified_set) - unverified_queue = list(unverified_set) - while unverified_queue: - curr = unverified_queue.pop() - res = results_by_id.get(curr) - if res: - for cid in res.child_attack_result_ids: - if cid not in unverified_descendants: - unverified_descendants.add(cid) - unverified_queue.append(cid) - - unreached = [ - r - for r in results_by_id.values() - if r.attack_result_id not in visited_ids and r.attack_result_id not in unverified_descendants - ] + if root_result_ids is None and 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, @@ -1207,9 +1126,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 results_by_id.values()) or bool(unverified_set) - ): + 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(): @@ -1325,19 +1242,19 @@ def _build_node( ) failed_attacks = all_failed_attacks[:_MAX_DIAGNOSTIC_ITEMS] - attempts_by_unit: dict[tuple[str, str], list[ScenarioProgressResult]] = {} - for r in snapshot.results: - attempts_by_unit.setdefault((r.atomic_group_id, r.seed_group_id), []).append(r) - - diagnostic_retries = max( - snapshot.summary.overall.retries, - ScenarioProgressReadModel.total_retry_pressure( - attempts_per_unit=(len(attempts) for attempts in attempts_by_unit.values()), - persisted_retries=(r.total_retries for r in snapshot.results), - ), - ) if exact_total_retries is not None: - diagnostic_retries = max(diagnostic_retries, exact_total_retries) + diagnostic_retries = exact_total_retries + else: + raw_attempts_by_unit: dict[tuple[str, str], list[ScenarioProgressResult]] = {} + for r in snapshot.results: + raw_attempts_by_unit.setdefault((r.atomic_group_id, r.seed_group_id), []).append(r) + diagnostic_retries = max( + snapshot.summary.overall.retries, + ScenarioProgressReadModel.total_retry_pressure( + attempts_per_unit=(len(attempts) for attempts in raw_attempts_by_unit.values()), + persisted_retries=(r.total_retries for r in snapshot.results), + ), + ) total_persisted_retries = sum(r.total_retries for r in snapshot.results) attack_retries_count = ( @@ -1389,6 +1306,9 @@ def build_hierarchy_node( conversations_truncated: bool = False, children_truncated: bool = False, max_child_references: int = _MAX_CHILD_REFERENCES, + child_offset: int = 0, + child_limit: int | None = None, + next_children_cursor: str | None = None, ) -> ScenarioHierarchyNode: """ Assemble a single hierarchy node with its immediate children and generic conversations. @@ -1403,44 +1323,93 @@ def build_hierarchy_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. + child_offset: Position of the first reference in this page. + child_limit: Optional ordered-reference page size. + next_children_cursor: Continuation for the next reference page. Returns: ScenarioHierarchyNode: The assembled node. """ - conv_list: list[ScenarioConversationProjection] = [] + node = ScenarioProgressReadModel._map_hierarchy_node( + delta=delta, result=result, max_child_references=max_child_references + ) if isinstance(conversations, ScenarioAttackResultConversations): - conv_list = list(conversations.conversations) + node.conversations = list(conversations.conversations) if conversations_count is None: conversations_count = conversations.total_count conversations_truncated = conversations_truncated or conversations.truncated else: - conv_list = list(conversations) - - resolved_conversations_count: int | None = ( - conversations_count if conversations_count is not None else len(conv_list) - ) - + node.conversations = list(conversations) + node.conversations_count = conversations_count if conversations_count is not None else len(node.conversations) + node.conversations_truncated = conversations_truncated 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} + page_child_ids = ( + result.child_attack_result_ids[child_offset : child_offset + child_limit] + if child_limit is not None + else result.child_attack_result_ids[child_offset:] + ) + node.child_attack_result_ids = page_child_ids[:max_child_references] + node.child_references_truncated = len(result.child_attack_result_ids) > len(node.child_attack_result_ids) + node.child_offset = child_offset + node.next_children_cursor = next_children_cursor + node.children_truncated = children_truncated or len(page_child_ids) > max_child_references + if node.orchestration is not None: + node.orchestration.child_attack_result_ids = list(node.child_attack_result_ids) + node.orchestration.child_references_truncated = node.child_references_truncated + unresolved_children: list[str] = [] + seen_child_ids = set(result.child_attack_result_ids[:child_offset]) + for idx, c_id in enumerate(node.child_attack_result_ids, start=child_offset): + if ( + c_id == result.attack_result_id + or c_id in seen_child_ids + or c_id not in children_deltas_by_id + or c_id not in children_results_by_id + ): + unresolved_children.append(c_id) + continue + seen_child_ids.add(c_id) + node.children.append( + ScenarioProgressReadModel._map_hierarchy_node( + delta=children_deltas_by_id[c_id], + result=children_results_by_id[c_id], + parent_id=result.attack_result_id, + child_index=idx, + max_child_references=max_child_references, + ) + ) + node.unresolved_child_attack_result_ids = unresolved_children[:max_child_references] + node.unresolved_child_references_truncated = ( + len(unresolved_children) > max_child_references or len(page_child_ids) > max_child_references + ) + return node - retries = result.retries - retries_truncated = False - if len(retries) > _MAX_RETRIES_PER_NODE: - retries = retries[:_MAX_RETRIES_PER_NODE] - retries_truncated = True + @staticmethod + def _map_hierarchy_node( + *, + delta: ScenarioAttackResultDelta, + result: ScenarioProgressResult, + parent_id: str | None = None, + child_index: int | None = None, + max_child_references: int = _MAX_CHILD_REFERENCES, + ) -> ScenarioHierarchyNode: + """ + Map shared result fields once for tree previews and direct node expansion. + Returns: + ScenarioHierarchyNode: A bounded node before attaching resolved children or conversations. + """ + retries = result.retries[:_MAX_RETRIES_PER_NODE] + retries_truncated = len(result.retries) > _MAX_RETRIES_PER_NODE 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 + child_ids = result.child_attack_result_ids[:max_child_references] + child_refs_truncated = len(result.child_attack_result_ids) > max_child_references + orchestration = None if result.result_role == AttackResultRole.ORCHESTRATION or result.child_attack_result_ids: completion_policy = delta.attack_metadata.get("completion_policy") orchestration = ScenarioOrchestrationParentSummary( @@ -1452,7 +1421,7 @@ 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=bounded_child_ids, + child_attack_result_ids=child_ids, child_references_truncated=child_refs_truncated, error_type=result.error_type, error_message=result.error_message, @@ -1461,113 +1430,16 @@ def build_hierarchy_node( retries=retries, retries_truncated=retries_truncated, ) - - children: list[ScenarioHierarchyNode] = [] - unresolved_children: list[str] = [] - - seen_child_ids: set[str] = set() - for idx, c_id in enumerate(result.child_attack_result_ids): - if c_id == result.attack_result_id or c_id in seen_child_ids or c_id not in children_deltas_by_id: - unresolved_children.append(c_id) - else: - seen_child_ids.add(c_id) - c_delta = children_deltas_by_id[c_id] - c_res = children_results_by_id[c_id] - c_retries = c_res.retries - c_retries_truncated = False - if len(c_retries) > _MAX_RETRIES_PER_NODE: - c_retries = c_retries[:_MAX_RETRIES_PER_NODE] - c_retries_truncated = True - - c_executed_turns = _resolve_executed_turns(c_delta) - c_outcome_reason = ( - c_delta.outcome_reason - or c_delta.attribution_data.get("outcome_reason") - or c_delta.attack_metadata.get("outcome_reason") - ) - - c_child_ids = list(c_res.child_attack_result_ids) - c_child_refs_truncated = len(c_child_ids) > max_child_references - c_bounded_child_ids = c_child_ids[:max_child_references] if c_child_refs_truncated else c_child_ids - - c_orchestration: ScenarioOrchestrationParentSummary | None = None - if c_res.result_role == AttackResultRole.ORCHESTRATION or c_res.child_attack_result_ids: - c_policy = c_delta.attack_metadata.get("completion_policy") - c_orchestration = ScenarioOrchestrationParentSummary( - attack_result_id=c_res.attack_result_id, - result_role=c_res.result_role, - outcome=c_res.outcome, - outcome_reason=c_outcome_reason, - attempt_index=c_res.attempt_index, - execution_time_ms=c_res.execution_time_ms, - executed_turns=c_executed_turns, - completion_policy=str(c_policy) if c_policy is not None else None, - child_attack_result_ids=c_bounded_child_ids, - child_references_truncated=c_child_refs_truncated, - error_type=c_res.error_type, - error_message=c_res.error_message, - timestamp=c_res.timestamp, - total_retries=c_res.total_retries, - retries=c_retries, - retries_truncated=c_retries_truncated, - ) - - c_has_related = bool( - c_res.related_conversation_ids - or ScenarioProgressReadModel._read_related_conversation_ids(delta=c_delta) - ) - if c_res.conversation_id or c_has_related: - c_conv_count = None - c_conv_truncated = True - else: - c_conv_count = 0 - c_conv_truncated = False - - children.append( - ScenarioHierarchyNode( - attack_result_id=c_res.attack_result_id, - conversation_id=c_res.conversation_id or "", - parent_attack_result_id=result.attack_result_id, - child_index=idx, - result_role=c_res.result_role, - attempt_index=c_res.attempt_index, - atomic_attack_name=c_res.atomic_attack_name, - atomic_group_id=c_res.atomic_group_id, - seed_group_id=c_res.seed_group_id, - objective=c_delta.objective, - outcome=c_res.outcome, - outcome_reason=c_outcome_reason, - execution_time_ms=c_res.execution_time_ms, - executed_turns=c_executed_turns, - timestamp=c_res.timestamp, - score=c_res.score, - error_type=c_res.error_type, - error_message=c_res.error_message, - total_retries=c_res.total_retries, - retries=c_retries, - retries_truncated=c_retries_truncated, - child_attack_result_ids=c_bounded_child_ids, - child_references_truncated=c_child_refs_truncated, - children=[], - children_truncated=False, - unresolved_child_attack_result_ids=[], - conversations=[], - conversations_count=c_conv_count, - conversations_truncated=c_conv_truncated, - orchestration=c_orchestration, - ) - ) - - unresolved_refs_truncated = len(unresolved_children) > max_child_references - bounded_unresolved_children = ( - unresolved_children[:max_child_references] if unresolved_refs_truncated else unresolved_children + has_conversations = bool( + result.conversation_id + or result.related_conversation_ids + or ScenarioProgressReadModel._read_related_conversation_ids(delta=delta) ) - return ScenarioHierarchyNode( attack_result_id=result.attack_result_id, conversation_id=result.conversation_id or "", - parent_attack_result_id=None, - child_index=None, + 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, @@ -1585,14 +1457,11 @@ def build_hierarchy_node( total_retries=result.total_retries, retries=retries, retries_truncated=retries_truncated, - child_attack_result_ids=bounded_child_ids, + child_attack_result_ids=child_ids, child_references_truncated=child_refs_truncated, - children=children, - children_truncated=children_truncated, - unresolved_child_attack_result_ids=bounded_unresolved_children, - unresolved_child_references_truncated=unresolved_refs_truncated, - conversations=conv_list, - conversations_count=resolved_conversations_count, - conversations_truncated=conversations_truncated, + children_count=len(result.child_attack_result_ids), + children_truncated=bool(result.child_attack_result_ids), + conversations_count=None if has_conversations else 0, + conversations_truncated=has_conversations, orchestration=orchestration, ) diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index af047f8bad..3eee234e74 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -40,6 +40,7 @@ fingerprint_filters, normalize_label_filters, ) +from pyrit.backend.services.saved_conversation_read_service import SavedConversationReadService from pyrit.backend.services.scenario_configuration_resolver import ScenarioConfigurationResolver from pyrit.backend.services.scenario_progress_read_model import ( ResultUnitIdentity, @@ -230,6 +231,7 @@ def __init__( self._progress_read_model = ( progress_read_model if progress_read_model is not None else ScenarioProgressReadModel(memory=self._memory) ) + self._saved_conversation_reader = SavedConversationReadService(memory=self._memory) self._technique_metadata_cache: dict[str, dict[str, ScenarioTechniqueSummary]] = {} self._preparations: set[asyncio.Task[_PreparedRun]] = set() @@ -2492,11 +2494,116 @@ 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_hierarchy_aggregate_async( + self, *, scenario_result_id: str, plan: ScenarioRunPlan | None + ) -> ScenarioHistoryAggregate: + """ + Read exact shared run statistics independently of the selected hierarchy page. + + Returns: + ScenarioHistoryAggregate: The same SQL/SDK aggregate used by saved-run history. + """ + aggregates = await self._memory.get_scenario_history_aggregates_async( + scenario_result_ids=[scenario_result_id], + plan_scenario_ids=[scenario_result_id] if plan is not None else [], + ) + aggregates = await self._recount_with_sdk_statistics_async( + aggregates=aggregates, + plans={scenario_result_id: plan.atomic_groups if plan is not None else None}, + ) + return aggregates[scenario_result_id] + + @staticmethod + def _hierarchy_counts(*, aggregate: ScenarioHistoryAggregate, planned: int | None) -> ScenarioProgressCounts: + """ + Present the shared aggregate without choosing a second counting population. + + Returns: + ScenarioProgressCounts: Exact logical counts, outcomes, and producer diagnostics. + """ + return ScenarioProgressCounts( + planned=planned, + completed=aggregate.completed_units, + succeeded=aggregate.successful_units, + success_percentage=success_percentage( + succeeded=aggregate.successful_units, completed=aggregate.completed_units + ), + errors=aggregate.error_attempts, + retries=aggregate.total_retries, + producer_counts=ScenarioProducerCounts( + target_facing=ScenarioProducerCategoryCounts( + attempts=aggregate.target_facing_attempts, + errors=aggregate.target_facing_error_attempts, + retries=aggregate.target_facing_retries, + ), + orchestration=ScenarioProducerCategoryCounts( + attempts=aggregate.orchestration_attempts, + errors=aggregate.orchestration_error_attempts, + retries=aggregate.orchestration_retries, + ), + unknown=ScenarioProducerCategoryCounts( + attempts=aggregate.unknown_role_attempts, + errors=aggregate.unknown_role_error_attempts, + retries=aggregate.unknown_role_retries, + ), + ), + outcomes=compute_outcome_statistics( + { + AttackOutcome.SUCCESS: aggregate.outcome_successes, + AttackOutcome.FAILURE: aggregate.outcome_failures, + AttackOutcome.ERROR: aggregate.outcome_errors, + AttackOutcome.UNDETERMINED: aggregate.outcome_undetermined, + } + ), + ) + + async def _load_hierarchy_preview_async( + self, *, scenario_result_id: str, roots: Sequence[ScenarioAttackResultDelta], max_nodes: int + ) -> list[ScenarioAttackResultDelta]: + """ + Resolve only the bounded preview's descendants in breadth-first batches. + + Returns: + list[ScenarioAttackResultDelta]: Roots and saved descendants, never more than the node budget. + """ + deltas = {delta.attack_result_id: delta for delta in roots} + queried_ids = set(deltas) + frontier = list(roots) + for _ in range(10): + if not frontier or len(deltas) >= max_nodes: + break + pending = ( + child_id + for delta in frontier + for child_id in self._progress_read_model._read_child_attack_result_ids( + attack_metadata=delta.attack_metadata + )[:500] + ) + batch: list[str] = [] + batch_limit = min(500, max_nodes - len(deltas)) + for child_id in pending: + if child_id in queried_ids: + continue + queried_ids.add(child_id) + batch.append(child_id) + if len(batch) >= batch_limit: + break + if not batch: + break + frontier = list( + await self._memory.get_scenario_attack_result_deltas_by_ids_async( + scenario_result_id=scenario_result_id, attack_result_ids=batch + ) + ) + deltas.update((child.attack_result_id, child) for child in frontier) + return list(deltas.values()) + async def get_run_hierarchy_async( self, *, scenario_result_id: str, max_nodes: int = 1000, + cursor: str | None = None, ) -> ScenarioHierarchyProjection | None: """ Get the bounded result hierarchy projection for a scenario run. @@ -2504,6 +2611,7 @@ async def get_run_hierarchy_async( Args: scenario_result_id: The persisted scenario run ID. max_nodes: Maximum nodes to assemble in the hierarchy tree. + cursor: Optional continuation for the ascending page of true roots. Returns: ScenarioHierarchyProjection | None: The hierarchy projection, or None if the run does not exist. @@ -2520,181 +2628,51 @@ async def get_run_hierarchy_async( scenario_result_id, ) plan = None - plan_complete = plan is not None - terminal = header_result.scenario_run_state in ( - ScenarioRunState.COMPLETED, - ScenarioRunState.FAILED, - ScenarioRunState.CANCELLED, + if not 1 <= max_nodes <= 5000: + raise ValueError("Hierarchy node budget must be between 1 and 5000.") + fingerprint = fingerprint_filters(filters={"kind": "hierarchy-roots", "run": scenario_result_id}) + anchor = decode_keyset_cursor(cursor=cursor, fingerprint=fingerprint) + if cursor is not None and (anchor is None or anchor.position is not None): + raise ValueError("Invalid hierarchy root cursor.") + root_cursor = ( + AttackResultKeysetCursor(timestamp=anchor.timestamp, attack_result_id=anchor.identifier) + if anchor is not None + else None ) - snapshot = self.snapshot_active_run(scenario_result_id=scenario_result_id) - objective_scorer_identifier = header_result.objective_scorer_identifier - if not isinstance(objective_scorer_identifier, ComponentIdentifier): - objective_scorer_identifier = None - - progress_snapshot = await self._progress_read_model.get_snapshot_async( - scenario_result_id=scenario_result_id, - plan=plan, - plan_complete=plan_complete, - active_group_ids=snapshot.active_group_ids, - terminal=terminal, - objective_scorer_identifier=objective_scorer_identifier, - max_deltas=max(max_nodes, self._progress_read_model._STORAGE_PAGE_SIZE), + roots, has_more = await self._memory.get_scenario_root_deltas_async( + scenario_result_id=scenario_result_id, cursor=root_cursor, limit=min(500, max_nodes) + ) + deltas = await self._load_hierarchy_preview_async( + scenario_result_id=scenario_result_id, roots=roots, max_nodes=max_nodes + ) + aggregate = await self._get_hierarchy_aggregate_async(scenario_result_id=scenario_result_id, plan=plan) + summary_plan = plan or self._progress_read_model._synthesize_legacy_plan(deltas=deltas) + plan_lookup = ScenarioPlanLookup.from_plan(plan=summary_plan) + results = tuple( + self._progress_read_model._map_progress_delta(delta=delta, plan_lookup=plan_lookup) for delta in deltas + ) + summary = self._progress_read_model._build_progress_summary( + plan=summary_plan, + plan_complete=plan is not None, + results=results, + active_group_ids=(), + terminal=header_result.scenario_run_state + in (ScenarioRunState.COMPLETED, ScenarioRunState.FAILED, ScenarioRunState.CANCELLED), + objective_scorer_identifier=None, + technique_details_by_group=self._progress_read_model._build_technique_details_by_group( + deltas=deltas, results=results + ), + ) + raw_attempts = ( + aggregate.target_facing_attempts + aggregate.orchestration_attempts + aggregate.unknown_role_attempts + ) + progress_snapshot = ScenarioProgressSnapshot( + deltas=tuple(deltas), + results=results, + summary=summary, + plan=summary_plan, + has_more=has_more or raw_attempts > len(deltas), ) - - exact_overall_counts: ScenarioProgressCounts | None = None - exact_total_errors: int | None = None - exact_total_retries: int | None = None - exact_retrying_attacks_count: int | None = None - exact_group_counts: Mapping[str, ScenarioProgressCounts] | None = None - parent_deltas: list[ScenarioAttackResultDelta] = [] - unverified_candidate_ids: set[str] = set() - - if progress_snapshot.has_more: - plan_lookup = ScenarioPlanLookup.from_plan(plan=plan) - all_deltas_map: dict[str, ScenarioAttackResultDelta] = { - d.attack_result_id: d for d in progress_snapshot.deltas - } - all_results_map: dict[str, ScenarioProgressResult] = { - r.attack_result_id: r for r in progress_snapshot.results - } - - referenced_child_ids: set[str] = set() - for r in all_results_map.values(): - for cid in r.child_attack_result_ids: - referenced_child_ids.add(cid) - - unreferenced_roots = [ - r.attack_result_id for r in all_results_map.values() if r.attack_result_id not in referenced_child_ids - ] - reachable: set[str] = set() - queue = list(unreferenced_roots) - while queue: - curr = queue.pop() - if curr not in reachable: - reachable.add(curr) - res = all_results_map.get(curr) - if res: - queue.extend(res.child_attack_result_ids) - - unreached_cyclic = [cid for cid in all_results_map if cid not in reachable] - current_candidates = list(unreferenced_roots) if unreferenced_roots else list(unreached_cyclic) - - parent_budget = max(2500, self._progress_read_model._STORAGE_PAGE_SIZE * 5) - checked_child_ids: set[str] = set() - depth = 0 - layer_candidates = list(current_candidates) - while layer_candidates and depth < 10 and len(parent_deltas) < parent_budget: - depth += 1 - next_layer_candidates: list[str] = [] - layer_clipped = False - while len(parent_deltas) < parent_budget: - remaining_budget = max(0, parent_budget - len(parent_deltas)) - if remaining_budget <= 0: - break - page_limit = min(self._progress_read_model._STORAGE_PAGE_SIZE, remaining_budget) - new_parents, has_more = await self._memory.get_scenario_parent_deltas_for_child_ids_async( - scenario_result_id=scenario_result_id, - child_ids=layer_candidates, - limit=page_limit, - exclude_result_ids=list(all_deltas_map.keys()), - ) - for p in new_parents: - if p.attack_result_id not in all_deltas_map: - parent_deltas.append(p) - all_deltas_map[p.attack_result_id] = p - p_res = self._progress_read_model._map_progress_delta(delta=p, plan_lookup=plan_lookup) - all_results_map[p.attack_result_id] = p_res - next_layer_candidates.append(p.attack_result_id) - if len(parent_deltas) >= parent_budget: - break - - if len(parent_deltas) >= parent_budget: - layer_clipped = True - break - - if not has_more: - break - - if not new_parents: - layer_clipped = True - break - - if layer_clipped: - unverified_candidate_ids.update(layer_candidates) - unverified_candidate_ids.update(next_layer_candidates) - break - - checked_child_ids.update(layer_candidates) - layer_candidates = [cid for cid in next_layer_candidates if cid not in checked_child_ids] - - if not layer_candidates and len(parent_deltas) < parent_budget and depth < 10: - remaining_cyclic = [cid for cid in unreached_cyclic if cid not in checked_child_ids] - if remaining_cyclic: - layer_candidates = remaining_cyclic - - if layer_candidates: - unverified_candidate_ids.update(layer_candidates) - - # Any node in all_results_map that is not referenced as a child by another loaded node - # and whose parents were not queried to completion in the DB is unverified. - all_referenced_cids: set[str] = {cid for r in all_results_map.values() for cid in r.child_attack_result_ids} - for rid in all_results_map: - if rid not in all_referenced_cids and rid not in checked_child_ids: - unverified_candidate_ids.add(rid) - - plan_ids = [scenario_result_id] if plan is not None else [] - aggregates = await self._memory.get_scenario_history_aggregates_async( - scenario_result_ids=[scenario_result_id], - plan_scenario_ids=plan_ids, - ) - plans_map = {scenario_result_id: plan.atomic_groups if plan is not None else None} - aggregates = await self._recount_with_sdk_statistics_async( - aggregates=aggregates, - plans=plans_map, - ) - agg = aggregates.get(scenario_result_id) - if agg is not None: - producer_counts = ScenarioProducerCounts( - target_facing=ScenarioProducerCategoryCounts( - attempts=agg.target_facing_attempts, - errors=agg.target_facing_error_attempts, - retries=agg.target_facing_retries, - ), - orchestration=ScenarioProducerCategoryCounts( - attempts=agg.orchestration_attempts, - errors=agg.orchestration_error_attempts, - retries=agg.orchestration_retries, - ), - unknown=ScenarioProducerCategoryCounts( - attempts=agg.unknown_role_attempts, - errors=agg.unknown_role_error_attempts, - retries=agg.unknown_role_retries, - ), - ) - outcomes = compute_outcome_statistics( - { - AttackOutcome.SUCCESS: agg.outcome_successes, - AttackOutcome.FAILURE: agg.outcome_failures, - AttackOutcome.ERROR: agg.outcome_errors, - AttackOutcome.UNDETERMINED: agg.outcome_undetermined, - } - ) - exact_overall_counts = ScenarioProgressCounts( - planned=progress_snapshot.summary.overall.planned, - completed=agg.completed_units, - succeeded=agg.successful_units, - success_percentage=success_percentage( - succeeded=agg.successful_units, completed=agg.completed_units - ), - errors=agg.error_attempts, - retries=agg.total_retries, - producer_counts=producer_counts, - outcomes=outcomes, - ) - exact_total_errors = agg.raw_total_errors - exact_total_retries = agg.raw_total_retries - exact_retrying_attacks_count = agg.raw_retrying_attacks_count - exact_group_counts = agg.group_counts hierarchy = self._progress_read_model.build_scenario_hierarchy( scenario_result_id=scenario_result_id, @@ -2704,14 +2682,27 @@ 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, - exact_retrying_attacks_count=exact_retrying_attacks_count, - exact_group_counts=exact_group_counts, - parent_deltas=parent_deltas, - unverified_candidate_ids=unverified_candidate_ids, + overall_counts=self._hierarchy_counts(aggregate=aggregate, planned=summary.overall.planned), + exact_total_errors=aggregate.raw_total_errors, + exact_total_retries=aggregate.raw_total_retries, + exact_retrying_attacks_count=aggregate.raw_retrying_attacks_count, + exact_group_counts=aggregate.group_counts, + root_result_ids=[root.attack_result_id for root in roots], ) + hierarchy.has_more = has_more + if has_more and roots: + hierarchy.next_cursor = encode_keyset_cursor( + timestamp=roots[-1].timestamp, + identifier=roots[-1].attack_result_id, + fingerprint=fingerprint, + ) + if cursor is None and not roots and raw_attempts: + rootless, clipped = await self._memory.get_scenario_recount_summaries_async( + scenario_result_id=scenario_result_id, limit=min(100, max_nodes) + ) + hierarchy.diagnostics.rootless_result_ids = [result.attack_result_id for result in rootless] + hierarchy.diagnostics.rootless_results_count = raw_attempts + hierarchy.diagnostics.rootless_results_truncated = clipped # Obtain real conversation counts for the actual bounded preview nodes seen_conv_ids: set[str] = set() @@ -2773,6 +2764,8 @@ async def get_run_hierarchy_node_async( scenario_result_id: str, attack_result_id: str, max_immediate_children: int = 200, + children_cursor: str | None = None, + include_messages: bool = False, ) -> ScenarioHierarchyNode | None: """ Get bounded single-node hierarchy details and generic conversation projections without full-run hydration. @@ -2781,10 +2774,14 @@ async def get_run_hierarchy_node_async( 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. + children_cursor: Continuation bound to this run, node, and stored child ordering. + include_messages: Explicitly include compatible clipped conversation previews. Returns: ScenarioHierarchyNode | None: The node details, or None if the node is not found in this run. """ + if not 1 <= max_immediate_children <= 500: + raise ValueError("Hierarchy child page size must be between 1 and 500.") deltas = await self._memory.get_scenario_attack_result_deltas_by_ids_async( scenario_result_id=scenario_result_id, attack_result_ids=[attack_result_id], @@ -2800,6 +2797,10 @@ async def get_run_hierarchy_node_async( 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_lookup = ScenarioPlanLookup.from_plan(plan=plan) @@ -2809,29 +2810,52 @@ async def get_run_hierarchy_node_async( ) child_ids = target_result.child_attack_result_ids + fingerprint = fingerprint_filters( + filters={ + "kind": "hierarchy-children", + "run": scenario_result_id, + "node": attack_result_id, + "ordered_references": json.dumps(child_ids, separators=(",", ":")), + } + ) + anchor = decode_keyset_cursor(cursor=children_cursor, fingerprint=fingerprint) + if children_cursor is not None and ( + anchor is None + or anchor.position is None + or not 0 <= anchor.position < len(child_ids) + or anchor.identifier != target_delta.attack_result_id + or anchor.timestamp != target_delta.timestamp + ): + raise ValueError("Invalid or stale hierarchy child cursor.") + child_offset = anchor.position if anchor is not None and anchor.position is not None else 0 + page_child_ids = child_ids[child_offset : child_offset + max_immediate_children] + next_offset = child_offset + len(page_child_ids) + next_children_cursor = ( + encode_keyset_cursor( + timestamp=target_delta.timestamp, + identifier=target_delta.attack_result_id, + position=next_offset, + fingerprint=fingerprint, + ) + if next_offset < len(child_ids) + else None + ) 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 - + if page_child_ids: children_deltas = list( await self._memory.get_scenario_attack_result_deltas_by_ids_async( scenario_result_id=scenario_result_id, - attack_result_ids=fetch_child_ids, + attack_result_ids=page_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( + conversations = await self._saved_conversation_reader.get_legacy_projections_async( attack_result_id=attack_result_id, - include_messages=True, + include_messages=include_messages, ) return self._progress_read_model.build_hierarchy_node( @@ -2840,7 +2864,10 @@ async def get_run_hierarchy_node_async( children_deltas=children_deltas, children_results=children_results, conversations=conversations, - children_truncated=children_truncated, + children_truncated=child_offset > 0 or next_offset < len(child_ids), + child_offset=child_offset, + child_limit=max_immediate_children, + next_children_cursor=next_children_cursor, ) async def get_run_results_async(self, *, scenario_result_id: str) -> ScenarioResult | None: diff --git a/pyrit/common/pagination.py b/pyrit/common/pagination.py index 6b0649c089..e1c8e6b5ce 100644 --- a/pyrit/common/pagination.py +++ b/pyrit/common/pagination.py @@ -16,10 +16,11 @@ @dataclass(frozen=True, slots=True) class DecodedKeysetCursor: - """A UTC timestamp/UUID anchor for descending recency pagination.""" + """A UTC timestamp/UUID anchor with an optional ordered-reference or message-sequence position.""" timestamp: datetime identifier: str + position: int | None = None def normalize_label_filters( @@ -78,7 +79,7 @@ def fingerprint_filters(*, filters: Mapping[str, Any], length: int = 16) -> str: return hashlib.sha256(canonical.encode("utf-8")).hexdigest()[:length] -def encode_keyset_cursor(*, timestamp: datetime, identifier: str, fingerprint: str) -> str: +def encode_keyset_cursor(*, timestamp: datetime, identifier: str, fingerprint: str, position: int | None = None) -> str: """ Encode a filter-bound keyset anchor as an opaque cursor. @@ -86,6 +87,7 @@ def encode_keyset_cursor(*, timestamp: datetime, identifier: str, fingerprint: s timestamp (datetime): A timezone-aware recency key from the last returned row. identifier (str): That row's UUID, breaking ties at an identical timestamp. fingerprint (str): The effective filters and result-selection policy. + position (int | None): An optional ordered-reference or message-sequence position. Returns: str: A base64url-encoded cursor. @@ -96,6 +98,8 @@ def encode_keyset_cursor(*, timestamp: datetime, identifier: str, fingerprint: s "t": timestamp.isoformat(), "i": identifier, } + if position is not None: + payload["p"] = position raw = json.dumps(payload, separators=(",", ":")).encode("utf-8") return base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=") @@ -126,8 +130,11 @@ def decode_keyset_cursor(*, cursor: str | None, fingerprint: str) -> DecodedKeys return None timestamp_value = payload.get("t") identifier_value = payload.get("i") + position = payload.get("p") if not isinstance(timestamp_value, str) or not isinstance(identifier_value, str): return None + if position is not None and type(position) is not int: + return None try: timestamp = datetime.fromisoformat(timestamp_value) identifier = str(uuid.UUID(identifier_value)) @@ -139,7 +146,7 @@ def decode_keyset_cursor(*, cursor: str | None, fingerprint: str) -> DecodedKeys timestamp = timestamp.astimezone(UTC) except (OverflowError, OSError): return None - return DecodedKeysetCursor(timestamp=timestamp, identifier=identifier) + return DecodedKeysetCursor(timestamp=timestamp, identifier=identifier, position=position) def _canonicalize(value: Any) -> Any: diff --git a/pyrit/memory/analytics_sql.py b/pyrit/memory/analytics_sql.py index 6eacfee282..84a8d65519 100644 --- a/pyrit/memory/analytics_sql.py +++ b/pyrit/memory/analytics_sql.py @@ -191,6 +191,20 @@ class JsonArrayItems(FunctionElement[str]): inherit_cache = True +class JsonArrayValue(FunctionElement[str]): + """A saved array at a JSON path, or an empty array for a missing or non-array value.""" + + type = UnicodeText() + inherit_cache = True + + +class UuidHex(FunctionElement[str]): + """Normalize UUID string forms for comparison with a stored ID, without casting malformed references.""" + + type = UnicodeText() + inherit_cache = True + + class JsonArrayEmpty(FunctionElement[bool]): """ A scalar 0/1 check for a recorded empty array, independent of JSON whitespace. @@ -364,6 +378,59 @@ def _mssql_json_items(element: JsonArrayItems, compiler: Any, **kwargs: Any) -> return f"OPENJSON({', '.join(_arguments(element, compiler, **kwargs))})" +@compiles(JsonArrayValue, "sqlite") +def _sqlite_json_array_value(element: JsonArrayValue, compiler: Any, **kwargs: Any) -> str: + """ + Check the saved property's shape before expanding it. + + Returns: + str: A JSON array expression safe for JSON1 expansion. + """ + document, path = _arguments(element, compiler, **kwargs) + return f"CASE WHEN json_type({document}, {path}) = 'array' THEN json_extract({document}, {path}) ELSE '[]' END" + + +@compiles(JsonArrayValue, "mssql") +def _mssql_json_array_value(element: JsonArrayValue, compiler: Any, **kwargs: Any) -> str: + """ + Keep only saved JSON arrays without truncating their contents. + + Returns: + str: A JSON array expression safe for OPENJSON expansion. + """ + document, path = _arguments(element, compiler, **kwargs) + value = f"JSON_QUERY({document}, {path})" + normalized = _mssql_without_json_whitespace(value) + return f"CASE WHEN LEFT({normalized}, 1) = N'[' THEN {value} ELSE N'[]' END" + + +@compiles(UuidHex, "sqlite") +def _sqlite_uuid_hex(element: UuidHex, compiler: Any, **kwargs: Any) -> str: + """ + Match Python's case-insensitive UUID normalization without interpreting invalid text as an ID. + + Returns: + str: Lowercase UUID hex for canonical, undashed, braced, or URN string forms. + """ + value = _arguments(element, compiler, **kwargs)[0] + return f"replace(trim(replace(replace(lower({value}), 'urn:', ''), 'uuid:', ''), '{{}}'), '-', '')" + + +@compiles(UuidHex, "mssql") +def _mssql_uuid_hex(element: UuidHex, compiler: Any, **kwargs: Any) -> str: + """ + Normalize reference strings without SQL Server's truncating uniqueidentifier conversion. + + Exact UTF-16 width and binary collation also prevent trailing spaces or Unicode folding from matching an ID. + + Returns: + str: The same UUID-comparison expression as SQLite and the Python reader. + """ + value = _arguments(element, compiler, **kwargs)[0] + normalized = f"replace(trim(N'{{}}' FROM replace(replace(lower({value}), N'urn:', N''), N'uuid:', N'')), N'-', N'')" + return f"CASE WHEN DATALENGTH({normalized}) = 64 THEN {normalized} COLLATE Latin1_General_100_BIN2 END" + + @compiles(JsonArrayEmpty, "sqlite") def _sqlite_array_empty(element: JsonArrayEmpty, compiler: Any, **kwargs: Any) -> str: """ diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 3f6b727176..d86627f06b 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -57,7 +57,7 @@ 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 -from pyrit.memory.analytics_sql import JsonArrayItems, JsonArrayJoin, JsonContainer, JsonScalar +from pyrit.memory.analytics_sql import JsonArrayItems, JsonArrayJoin, JsonArrayValue, JsonContainer, JsonScalar, UuidHex if TYPE_CHECKING: from pyrit.memory.memory_embedding import MemoryEmbedding @@ -113,6 +113,7 @@ ConversationRetry, ConversationRetryReason, ConversationStats, + ConversationType, ConverterIdentifier, Finding, FindingCreate, @@ -127,11 +128,7 @@ Operation, PromptDataType, RetryEvent, - ScenarioAttackResultConversations, ScenarioAttackResultDelta, - ScenarioConversationMessagePieceProjection, - ScenarioConversationMessageProjection, - ScenarioConversationProjection, ScenarioIdentifier, ScenarioProducerCategoryCounts, ScenarioProducerCounts, @@ -447,6 +444,43 @@ class ScenarioAttemptSeedSummary: child_attack_result_ids: tuple[str, ...] = () +class MessagePieceKeysetCursor(NamedTuple): + """Ascending saved-message order, including pieces within the same turn.""" + + sequence: int + timestamp: datetime + message_piece_id: str + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ConversationMessagePage: + """A bounded page of canonical saved messages, possibly splitting a multimodal turn.""" + + messages: tuple[Message, ...] + next_cursor: MessagePieceKeysetCursor | None + partial_message_sequences: tuple[int, ...] = () + + +@dataclass(frozen=True, slots=True, kw_only=True) +class AttackConversationRecord: + """A saved main or related conversation reference, without presentation labels.""" + + conversation_id: str + conversation_type: ConversationType | None = None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class AttackConversationPage: + """Bounded conversation references and the revision of their stored ordering.""" + + main_conversation_id: str + conversations: tuple[AttackConversationRecord, ...] + total_count: int + timestamp: datetime + reference_revision: str + objective_score_id: str | None = None + + @dataclass(frozen=True, slots=True, kw_only=True) class _AttackResultQuery: """ @@ -7686,6 +7720,118 @@ def _row_to_scenario_attack_result_delta(row: Mapping[str, Any]) -> ScenarioAtta score=score, ) + @staticmethod + def _scenario_attack_result_delta_statement() -> Select: + """Return the shared lightweight result/score projection, without filters or pagination.""" + return select( + AttackResultEntry.id, + AttackResultEntry.conversation_id, + AttackResultEntry.objective, + AttackResultEntry.objective_sha256, + AttackResultEntry.atomic_attack_identifier, + AttackResultEntry.outcome, + AttackResultEntry.execution_time_ms, + AttackResultEntry.executed_turns, + AttackResultEntry.outcome_reason, + AttackResultEntry.timestamp, + AttackResultEntry.retry_events_json, + AttackResultEntry.total_retries, + AttackResultEntry.error_type, + AttackResultEntry.error_message, + AttackResultEntry.attribution_data, + AttackResultEntry.attack_metadata, + AttackResultEntry.pruned_conversation_ids, + AttackResultEntry.adversarial_chat_conversation_ids, + AttackResultEntry.preparation_conversation_ids, + ScoreEntry.id.label("score_id"), + ScoreEntry.score_value, + ScoreEntry.score_type, + ScoreEntry.status.label("score_status"), + ScoreEntry.score_rationale, + ScoreEntry.scorer_class_identifier, + ).outerjoin( + ScoreEntry, + func.coalesce(AttackResultEntry.human_score_id, AttackResultEntry.automated_score_id) == ScoreEntry.id, + ) + + def _build_scenario_root_deltas_statement( + self, + *, + scenario_result_id: str, + cursor: AttackResultKeysetCursor | None = None, + limit: int = 100, + ) -> Select: + """ + Page saved roots directly instead of hydrating a chronological prefix and searching for its ancestors. + + Keep reference expansion uncorrelated so SQLite reads the saved edges once, not once per candidate. + + Returns: + Select: A run-scoped, ascending keyset query with one sentinel row. + + Raises: + ValueError: If the run, cursor, or page limit is invalid. + """ + if not 1 <= limit <= 500: + raise ValueError("Hierarchy root limit must be between 1 and 500.") + scenario_uuid = uuid.UUID(scenario_result_id) + parents = AttackResultEntry.__table__.alias("hierarchy_parents") + references = JsonArrayItems( + JsonArrayValue(parents.c.attack_metadata, literal("$.child_attack_result_ids")) + ).table_valued("value", "type") + candidate_id = UuidHex(sql_cast(AttackResultEntry.id, String(36))) + referenced_id = UuidHex(references.c.value) + incoming_ids = ( + select(referenced_id) + .select_from(JsonArrayJoin(parents, references, literal(True))) + .where( + parents.c.attribution_parent_id == scenario_uuid, + sql_cast(references.c.type, String(8)).in_(("text", "1")), + referenced_id.is_not(None), + ) + ) + conditions = [ + AttackResultEntry.attribution_parent_id == scenario_uuid, + candidate_id.not_in(incoming_ids), + ] + attempt_id_order = self._get_scenario_attempt_id_order_expression(attempt_id=AttackResultEntry.id) + if cursor is not None: + cursor_id_order = self._get_scenario_attempt_id_order_expression( + attempt_id=literal(uuid.UUID(cursor.attack_result_id), type_=AttackResultEntry.id.type) + ) + conditions.append( + or_( + AttackResultEntry.timestamp > cursor.timestamp, + and_(AttackResultEntry.timestamp == cursor.timestamp, attempt_id_order > cursor_id_order), + ) + ) + return ( + self._scenario_attack_result_delta_statement() + .where(*conditions) + .order_by(AttackResultEntry.timestamp.asc(), attempt_id_order.asc()) + .limit(limit + 1) + ) + + def _execute_get_scenario_root_deltas( + self, + *, + scenario_result_id: str, + cursor: AttackResultKeysetCursor | None = None, + limit: int = 100, + ) -> tuple[list[ScenarioAttackResultDelta], bool]: + """ + Read one bounded page of true saved roots without loading their children or messages. + + Returns: + tuple[list[ScenarioAttackResultDelta], bool]: Root deltas and whether another root page exists. + """ + statement = self._build_scenario_root_deltas_statement( + scenario_result_id=scenario_result_id, cursor=cursor, limit=limit + ) + 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[:limit]], len(rows) > limit + def _execute_get_scenario_attack_result_deltas( self, *, @@ -7728,37 +7874,7 @@ def _execute_get_scenario_attack_result_deltas( ) 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, - ) + self._scenario_attack_result_delta_statement() .where(and_(*conditions)) .order_by(AttackResultEntry.timestamp.asc(), attempt_id_order.asc()) .limit(limit + 1) @@ -7920,44 +8036,9 @@ def _build_scenario_attack_result_deltas_by_ids_statement( max_batch_size = max(1, self._MAX_BIND_VARS - 10) chunk_uuids = valid_uuids[:max_batch_size] - return ( - select( - AttackResultEntry.id, - AttackResultEntry.conversation_id, - AttackResultEntry.objective, - AttackResultEntry.objective_sha256, - AttackResultEntry.atomic_attack_identifier, - AttackResultEntry.outcome, - AttackResultEntry.execution_time_ms, - AttackResultEntry.executed_turns, - AttackResultEntry.outcome_reason, - AttackResultEntry.timestamp, - AttackResultEntry.retry_events_json, - AttackResultEntry.total_retries, - AttackResultEntry.error_type, - AttackResultEntry.error_message, - AttackResultEntry.attribution_data, - AttackResultEntry.attack_metadata, - AttackResultEntry.pruned_conversation_ids, - AttackResultEntry.adversarial_chat_conversation_ids, - AttackResultEntry.preparation_conversation_ids, - ScoreEntry.id.label("score_id"), - ScoreEntry.score_value, - ScoreEntry.score_type, - ScoreEntry.status.label("score_status"), - ScoreEntry.score_rationale, - ScoreEntry.scorer_class_identifier, - ) - .outerjoin( - ScoreEntry, - func.coalesce(AttackResultEntry.human_score_id, AttackResultEntry.automated_score_id) == ScoreEntry.id, - ) - .where( - and_( - AttackResultEntry.attribution_parent_id == scenario_uuid, - AttackResultEntry.id.in_(chunk_uuids), - ) - ) + return self._scenario_attack_result_delta_statement().where( + AttackResultEntry.attribution_parent_id == scenario_uuid, + AttackResultEntry.id.in_(chunk_uuids), ) def _execute_get_scenario_attack_result_deltas_by_ids( @@ -8006,262 +8087,155 @@ def _execute_get_scenario_attack_result_deltas_by_ids( deltas.extend(self._row_to_scenario_attack_result_delta(row) for row in rows) return deltas - def _execute_get_scenario_orchestration_parent_deltas( + def _build_conversation_message_page_statement( self, *, - scenario_result_id: str, - ) -> list[ScenarioAttackResultDelta]: + conversation_id: str, + cursor: MessagePieceKeysetCursor | None = None, + limit: int = 100, + ) -> Select: """ - 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. + Build a canonical message-piece page, including continuation within a multimodal turn. Returns: - list[ScenarioAttackResultDelta]: The matched orchestration parent deltas. - """ - try: - scenario_uuid = uuid.UUID(scenario_result_id) - except (ValueError, TypeError): - return [] + Select: An ascending sequence/timestamp/UUID query with one sentinel. - role_match = self._get_scenario_role_match_condition( - role_expression=JsonScalar(AttackResultEntry.attribution_data, literal("$.result_role")), - role=AttackResultRole.ORCHESTRATION, - ) - metadata_text = sql_cast(AttackResultEntry.attack_metadata, String) - # Match only when child_attack_result_ids is not empty (contains at least one quoted ID) - child_ids_match = metadata_text.like('%child_attack_result_ids%["%') - - attempt_id_order = self._get_scenario_attempt_id_order_expression(attempt_id=AttackResultEntry.id) - - statement = ( - select( - AttackResultEntry.id, - AttackResultEntry.conversation_id, - AttackResultEntry.objective, - AttackResultEntry.objective_sha256, - AttackResultEntry.atomic_attack_identifier, - AttackResultEntry.outcome, - AttackResultEntry.execution_time_ms, - AttackResultEntry.executed_turns, - AttackResultEntry.outcome_reason, - AttackResultEntry.timestamp, - AttackResultEntry.retry_events_json, - AttackResultEntry.total_retries, - AttackResultEntry.error_type, - AttackResultEntry.error_message, - AttackResultEntry.attribution_data, - AttackResultEntry.attack_metadata, - AttackResultEntry.pruned_conversation_ids, - AttackResultEntry.adversarial_chat_conversation_ids, - AttackResultEntry.preparation_conversation_ids, - ScoreEntry.id.label("score_id"), - ScoreEntry.score_value, - ScoreEntry.score_type, - ScoreEntry.status.label("score_status"), - ScoreEntry.score_rationale, - ScoreEntry.scorer_class_identifier, - ) - .outerjoin( - ScoreEntry, - func.coalesce(AttackResultEntry.human_score_id, AttackResultEntry.automated_score_id) == ScoreEntry.id, + Raises: + ValueError: If the page limit or cursor UUID is invalid. + """ + if not 1 <= limit <= 500: + raise ValueError("Conversation message-piece limit must be between 1 and 500.") + id_order = self._get_scenario_attempt_id_order_expression(attempt_id=PromptMemoryEntry.id) + conditions = [PromptMemoryEntry.conversation_id == conversation_id] + if cursor is not None: + cursor_id = self._get_scenario_attempt_id_order_expression( + attempt_id=literal(uuid.UUID(cursor.message_piece_id), type_=PromptMemoryEntry.id.type) ) - .where( - and_( - AttackResultEntry.attribution_parent_id == scenario_uuid, - or_(role_match, child_ids_match), + conditions.append( + or_( + PromptMemoryEntry.sequence > cursor.sequence, + and_(PromptMemoryEntry.sequence == cursor.sequence, PromptMemoryEntry.timestamp > cursor.timestamp), + and_( + PromptMemoryEntry.sequence == cursor.sequence, + PromptMemoryEntry.timestamp == cursor.timestamp, + id_order > cursor_id, + ), ) ) - .order_by(AttackResultEntry.timestamp.asc(), attempt_id_order.asc()) - .limit(self._MAX_BIND_VARS) + return ( + select(PromptMemoryEntry) + .where(*conditions) + .order_by(PromptMemoryEntry.sequence.asc(), PromptMemoryEntry.timestamp.asc(), id_order.asc()) + .limit(limit + 1) ) - with closing(self._get_session()) as session: - rows = session.execute(statement).mappings().all() - - return [self._row_to_scenario_attack_result_delta(row) for row in rows] - - def _execute_get_scenario_parent_deltas_for_child_ids( + def _execute_get_conversation_message_page( self, *, - scenario_result_id: str, - child_ids: Sequence[str], - limit: int = 500, - exclude_result_ids: Sequence[str] | None = None, - ) -> tuple[list[ScenarioAttackResultDelta], bool]: + conversation_id: str, + cursor: MessagePieceKeysetCursor | None = None, + limit: int = 100, + ) -> ConversationMessagePage: """ - Return run-scoped parent deltas that reference any of the specified child result IDs. - - Paginates scanned database rows using keyset cursors across child chunks to guarantee - that a short returned list is a genuine proof of exhausted database rows, preventing - false-root generation when database exclusions are omitted or partial. - - Args: - scenario_result_id: Persisted scenario run ID. - child_ids: Child attack result IDs to find parents for. - limit: Maximum number of parent deltas to return. - exclude_result_ids: Optional result IDs already known to caller to exclude from query. + Read only the selected page of canonical message pieces. Returns: - tuple[list[ScenarioAttackResultDelta], bool]: The matched parent deltas and whether - more unscanned matching rows exist in the database. + ConversationMessagePage: Messages, continuation, and turns split by the page boundary. """ - try: - scenario_uuid = uuid.UUID(scenario_result_id) - except (ValueError, TypeError): - return [], False - valid_child_ids = [cid for cid in child_ids if cid] - if not valid_child_ids or limit <= 0: - return [], False - - valid_exclude_uuids: list[uuid.UUID] = [] - if exclude_result_ids: - for eid in exclude_result_ids: - try: - valid_exclude_uuids.append(uuid.UUID(str(eid))) - except (ValueError, TypeError): - continue - exclude_id_set = {str(u) for u in valid_exclude_uuids} - - seen_ids: set[str] = set() - unique_deltas: list[ScenarioAttackResultDelta] = [] - metadata_text = sql_cast(AttackResultEntry.attack_metadata, String) - attempt_id_order = self._get_scenario_attempt_id_order_expression(attempt_id=AttackResultEntry.id) - - _fixed_overhead = 10 - # To avoid exceeding SQL Server's 2,100 parameter limit (or SQLite's 500), - # budget the total parameters across fixed clauses, exclusions, and child ID matches. - max_sql_excludes = self._MAX_BIND_VARS // 2 - if len(valid_exclude_uuids) <= max_sql_excludes: - exclude_clauses = [ - AttackResultEntry.id.not_in(valid_exclude_uuids[j : j + 200]) - for j in range(0, len(valid_exclude_uuids), 200) - ] - 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), + statement = self._build_conversation_message_page_statement( + conversation_id=conversation_id, cursor=cursor, limit=limit ) - - has_more = False - max_scanned_rows = 10000 - total_scanned_rows = 0 - scan_page_size = min(500, self._MAX_BIND_VARS) - num_chunks = (len(valid_child_ids) + chunk_size - 1) // chunk_size - with closing(self._get_session()) as session: - for chunk_idx, i in enumerate(range(0, len(valid_child_ids), chunk_size)): - if len(unique_deltas) >= limit or total_scanned_rows >= max_scanned_rows: - has_more = True - break - - chunk = valid_child_ids[i : i + chunk_size] - child_matches = [metadata_text.like(f'%"{cid}"%') for cid in chunk] - base_where_conditions = [ - AttackResultEntry.attribution_parent_id == scenario_uuid, - or_(*child_matches), - *exclude_clauses, - ] + rows = list(session.scalars(statement)) + pieces = [row.get_message_piece() for row in rows[:limit]] + if not pieces: + return ConversationMessagePage(messages=(), next_cursor=None) + partial_sequences: set[int] = set() + if cursor is not None and pieces[0].sequence == cursor.sequence: + partial_sequences.add(cursor.sequence) + next_cursor = None + if len(rows) > limit: + last_piece = pieces[-1] + next_cursor = MessagePieceKeysetCursor( + sequence=last_piece.sequence, + timestamp=last_piece.timestamp, + message_piece_id=str(last_piece.id), + ) + if rows[limit].sequence == last_piece.sequence: + partial_sequences.add(last_piece.sequence) + return ConversationMessagePage( + messages=tuple(group_conversation_message_pieces_by_sequence(message_pieces=pieces)), + next_cursor=next_cursor, + partial_message_sequences=tuple(sorted(partial_sequences)), + ) - chunk_cursor: AttackResultKeysetCursor | None = None + def _execute_get_attack_conversation_page( + self, *, attack_result_id: str, offset: int = 0, limit: int = 50, conversation_id: str | None = None + ) -> AttackConversationPage | None: + """ + Read stored conversation references without loading attack relationships or message bodies. - 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, - ), - ) - ) + Returns: + AttackConversationPage | None: One bounded reference page, or None for an absent attack. - 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) + Raises: + ValueError: If the ID, offset, or page limit is invalid. + """ + if offset < 0 or not 1 <= limit <= 200: + raise ValueError("Conversation offset must be nonnegative and limit must be between 1 and 200.") + statement = select( + AttackResultEntry.conversation_id, + AttackResultEntry.pruned_conversation_ids, + AttackResultEntry.adversarial_chat_conversation_ids, + AttackResultEntry.preparation_conversation_ids, + AttackResultEntry.timestamp, + func.coalesce(AttackResultEntry.human_score_id, AttackResultEntry.automated_score_id).label( + "objective_score_id" + ), + ).where(AttackResultEntry.id == uuid.UUID(attack_result_id)) + with closing(self._get_session()) as session: + row = session.execute(statement).mappings().one_or_none() + if row is None: + return None + references: dict[str, AttackConversationRecord] = {} + main_id = row["conversation_id"] or "" + if main_id: + references[main_id] = AttackConversationRecord(conversation_id=main_id) + for column_name, kind in ( + ("pruned_conversation_ids", ConversationType.PRUNED), + ("adversarial_chat_conversation_ids", ConversationType.ADVERSARIAL), + ("preparation_conversation_ids", ConversationType.PREPARATION), + ): + values = row[column_name] + if not isinstance(values, list): + continue + for identifier in values: + if isinstance(identifier, str) and identifier: + references.setdefault( + identifier, AttackConversationRecord(conversation_id=identifier, conversation_type=kind) ) - 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"]), + ordered = tuple(references.values()) + revision = hashlib.sha256( + json.dumps( + [ + ( + reference.conversation_id, + reference.conversation_type.value if reference.conversation_type is not None else None, ) - - has_unprocessed_in_page = False - for row_idx, row in enumerate(rows[:scanned_in_query]): - delta = self._row_to_scenario_attack_result_delta(row) - if delta.attack_result_id not in seen_ids and delta.attack_result_id not in exclude_id_set: - seen_ids.add(delta.attack_result_id) - unique_deltas.append(delta) - if len(unique_deltas) >= limit: - if row_idx + 1 < len(rows): - has_unprocessed_in_page = True - break - - if len(unique_deltas) >= limit: - if has_more_in_chunk or has_unprocessed_in_page or chunk_idx + 1 < num_chunks: - has_more = True - break - - if not has_more_in_chunk: - break - - if total_scanned_rows >= max_scanned_rows: - has_more = True - break - - return unique_deltas[:limit], has_more + for reference in ordered + ] + ).encode("utf-8") + ).hexdigest() + if conversation_id is not None: + ordered = tuple(reference for reference in ordered if reference.conversation_id == conversation_id) + return AttackConversationPage( + main_conversation_id=main_id, + conversations=ordered[offset : offset + limit], + total_count=len(ordered), + timestamp=row["timestamp"], + reference_revision=revision, + objective_score_id=str(row["objective_score_id"]) if row["objective_score_id"] is not None else None, + ) def _fetch_bounded_conversation_messages( self, @@ -8269,17 +8243,17 @@ def _fetch_bounded_conversation_messages( conversation_id: str, limit: int, max_pieces_per_message: int = 100, - ) -> tuple[list[ScenarioConversationMessageProjection], bool]: + ) -> tuple[list[Message], bool, set[int]]: """ 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. + tuple[list[Message], bool, set[int]]: Canonical messages, more turns, and clipped turn sequences. """ if not conversation_id or limit < 1: - return [], False + return [], False, set() with closing(self._get_session()) as session: seq_stmt = ( @@ -8291,7 +8265,7 @@ def _fetch_bounded_conversation_messages( ) seqs = [row[0] for row in session.execute(seq_stmt).all()] if not seqs: - return [], False + return [], False, set() truncated = len(seqs) > limit target_seqs = seqs[:limit] @@ -8330,176 +8304,12 @@ def _fetch_bounded_conversation_messages( ) rows = session.execute(pieces_stmt).all() if not rows: - return [], False + return [], False, set() - truncated_seqs = {entry.sequence for entry, rn in rows if rn > piece_budget} + truncated_seqs: set[int] = {entry.sequence for entry, rn in rows if rn > piece_budget} bounded_pieces = [entry.get_message_piece() for entry, rn in rows if rn <= piece_budget] - grouped_messages = group_conversation_message_pieces_by_sequence(message_pieces=bounded_pieces) - - messages: list[ScenarioConversationMessageProjection] = [] - for msg in grouped_messages: - seq = msg.message_pieces[0].sequence if msg.message_pieces else 0 - pieces_truncated = seq in truncated_seqs - pieces_proj = [ - ScenarioConversationMessagePieceProjection( - data_type=piece.converted_value_data_type, - content=piece.original_value, - converted_value=piece.converted_value, - ) - for piece in msg.message_pieces - ] - first_piece = msg.message_pieces[0] if msg.message_pieces else None - raw_role = str(first_piece.role) if first_piece else "unknown" - ts = first_piece.timestamp if first_piece else None - content = msg.get_value() - - messages.append( - ScenarioConversationMessageProjection( - role=raw_role, - content=content, - sequence=seq, - timestamp=ts, - pieces=pieces_proj, - pieces_truncated=pieces_truncated, - ) - ) - - return messages, truncated - - def _execute_get_attack_result_conversation_projections( - self, - *, - attack_result_id: str, - include_messages: bool = False, - limit_messages_per_conversation: int = 100, - max_conversations: int = 50, - max_pieces_per_message: int = 100, - ) -> ScenarioAttackResultConversations: - """ - Read generic conversation projections for an attack result without full-run hydration. - - Applies bounded conversation budget and SQL LIMIT queries. - - Returns: - ScenarioAttackResultConversations: Container with projections, total count, and truncation flag. - """ - try: - ar_uuid = uuid.UUID(attack_result_id) - except (ValueError, TypeError): - return ScenarioAttackResultConversations(conversations=[], total_count=0, truncated=False) - - entries = self._query_entries( - AttackResultEntry, - conditions=AttackResultEntry.id == ar_uuid, - ) - if not entries: - return ScenarioAttackResultConversations(conversations=[], total_count=0, truncated=False) - - entry = entries[0] - conversations: list[ScenarioConversationProjection] = [] - all_conv_ids: list[str] = [] - - if entry.conversation_id: - all_conv_ids.append(entry.conversation_id) - - raw_pruned: Any = entry.pruned_conversation_ids - raw_adversarial: Any = entry.adversarial_chat_conversation_ids - raw_preparation: Any = entry.preparation_conversation_ids - - pruned_ids = [c for c in raw_pruned if isinstance(c, str)] if isinstance(raw_pruned, list) else [] - adversarial_ids = ( - [c for c in raw_adversarial if isinstance(c, str)] if isinstance(raw_adversarial, list) else [] - ) - preparation_ids = ( - [c for c in raw_preparation if isinstance(c, str)] if isinstance(raw_preparation, list) else [] - ) - - for c_id in pruned_ids + adversarial_ids + preparation_ids: - if c_id and c_id not in all_conv_ids: - all_conv_ids.append(c_id) - - total_conversations = len(all_conv_ids) - conversations_truncated = total_conversations > max_conversations - conv_ids_to_fetch = all_conv_ids[:max_conversations] - if not conv_ids_to_fetch: - return ScenarioAttackResultConversations(conversations=[], total_count=0, truncated=False) - - stats_map = self._execute_get_conversation_stats(conversation_ids=conv_ids_to_fetch) - - if entry.conversation_id: - stats = stats_map.get(entry.conversation_id) - msg_count = stats.message_count if stats else 0 - messages: list[ScenarioConversationMessageProjection] = [] - truncated = False - if include_messages and entry.conversation_id: - messages, truncated = self._fetch_bounded_conversation_messages( - conversation_id=entry.conversation_id, - limit=limit_messages_per_conversation, - max_pieces_per_message=max_pieces_per_message, - ) - conversations.append( - ScenarioConversationProjection( - conversation_id=entry.conversation_id, - conversation_type="main", - description="Main objective target conversation", - message_count=msg_count, - messages=messages, - messages_truncated=truncated, - ) - ) - - remaining_budget = max(0, max_conversations - (1 if entry.conversation_id else 0)) - related_convs = ( - [(c_id, "pruned", "Pruned conversation") for c_id in pruned_ids if c_id != entry.conversation_id] - + [ - (c_id, "adversarial_chat", "Adversarial chat conversation") - for c_id in adversarial_ids - if c_id != entry.conversation_id - ] - + [ - (c_id, "preparation", "Preparation conversation") - for c_id in preparation_ids - if c_id != entry.conversation_id - ] - ) - seen_related: set[str] = set() - deduped_related: list[tuple[str, str, str]] = [] - for item in related_convs: - if item[0] not in seen_related: - seen_related.add(item[0]) - deduped_related.append(item) - - bounded_related = deduped_related[:remaining_budget] - - for c_id, conv_type, desc in bounded_related: - if not c_id: - continue - stats = stats_map.get(c_id) - msg_count = stats.message_count if stats else 0 - messages = [] - truncated = False - if include_messages: - messages, truncated = self._fetch_bounded_conversation_messages( - conversation_id=c_id, - limit=limit_messages_per_conversation, - max_pieces_per_message=max_pieces_per_message, - ) - conversations.append( - ScenarioConversationProjection( - conversation_id=c_id, - conversation_type=conv_type, - description=desc, - message_count=msg_count, - messages=messages, - messages_truncated=truncated, - ) - ) - - return ScenarioAttackResultConversations( - conversations=conversations, - total_count=total_conversations, - truncated=conversations_truncated, - ) + messages = list(group_conversation_message_pieces_by_sequence(message_pieces=bounded_pieces)) + return messages, truncated, truncated_seqs def _execute_get_scenario_results( self, @@ -12002,6 +11812,26 @@ async def get_scenario_attack_result_deltas_async( limit=limit, ) + async def get_scenario_root_deltas_async( + self, + *, + scenario_result_id: str, + cursor: AttackResultKeysetCursor | None = None, + limit: int = 100, + ) -> tuple[list[ScenarioAttackResultDelta], bool]: + """ + Page true run-scoped roots without hydrating progress or discovering ancestors in Python. + + Returns: + tuple[list[ScenarioAttackResultDelta], bool]: Root deltas and whether another page exists. + """ + return await self._run_database_operation_async( + self._execute_get_scenario_root_deltas, + scenario_result_id=scenario_result_id, + cursor=cursor, + limit=limit, + ) + async def get_scenario_recount_summaries_async( self, *, @@ -12053,84 +11883,66 @@ async def get_scenario_attack_result_deltas_by_ids_async( attack_result_ids=attack_result_ids, ) - async def get_scenario_orchestration_parent_deltas_async( + async def get_attack_conversation_page_async( self, *, - scenario_result_id: str, - ) -> list[ScenarioAttackResultDelta]: + attack_result_id: str, + offset: int = 0, + limit: int = 50, + conversation_id: str | None = None, + ) -> AttackConversationPage | None: """ - 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. + Read one bounded page of stored main/related conversation references. Returns: - list[ScenarioAttackResultDelta]: The matched orchestration parent deltas. + AttackConversationPage | None: Saved references and ordering revision, or an absent attack. """ return await self._run_database_operation_async( - self._execute_get_scenario_orchestration_parent_deltas, - scenario_result_id=scenario_result_id, + self._execute_get_attack_conversation_page, + attack_result_id=attack_result_id, + offset=offset, + limit=limit, + conversation_id=conversation_id, ) - async def get_scenario_parent_deltas_for_child_ids_async( + async def get_conversation_message_page_async( self, *, - scenario_result_id: str, - child_ids: Sequence[str], - limit: int = 500, - exclude_result_ids: Sequence[str] | None = None, - ) -> tuple[list[ScenarioAttackResultDelta], bool]: + conversation_id: str, + cursor: MessagePieceKeysetCursor | None = None, + limit: int = 100, + ) -> ConversationMessagePage: """ - Return run-scoped parent deltas that reference any of the specified child result IDs. - - Args: - scenario_result_id: The ID of the scenario run. - child_ids: Child attack result IDs to find parents for. - limit: Maximum number of parent deltas to return. - exclude_result_ids: Optional result IDs already known to caller to exclude from query. + Read canonical saved messages with continuation through every message piece. Returns: - tuple[list[ScenarioAttackResultDelta], bool]: The matched parent deltas and whether - more unscanned matching rows exist in the database. + ConversationMessagePage: One bounded piece page grouped into canonical messages. """ 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, + self._execute_get_conversation_message_page, + conversation_id=conversation_id, + cursor=cursor, limit=limit, - exclude_result_ids=exclude_result_ids, ) - async def get_attack_result_conversation_projections_async( - self, - *, - attack_result_id: str, - include_messages: bool = False, - limit_messages_per_conversation: int = 100, - max_conversations: int = 50, - max_pieces_per_message: int = 100, - ) -> ScenarioAttackResultConversations: + async def get_bounded_conversation_messages_async( + self, *, conversation_id: str, limit: int = 100, max_pieces_per_message: int = 100 + ) -> tuple[list[Message], bool, set[int]]: """ - Read generic conversation projections for an attack result without full-run hydration. - - Args: - attack_result_id: The attack result ID. - include_messages: Whether to populate message pieces or only metadata/counts. - limit_messages_per_conversation: Bound on message pieces per conversation. - max_conversations: Maximum number of conversations to retrieve for this attack result. - max_pieces_per_message: Bound on message pieces within a single message. + Read a bounded turn preview for compatibility with existing clipped detail views. Returns: - ScenarioAttackResultConversations: Container with projections, total count, and truncation flag. + tuple[list[Message], bool, set[int]]: Canonical messages, more turns, and clipped turn sequences. + + Raises: + ValueError: If the turn or piece preview budget is invalid. """ + if not 1 <= limit <= 100 or not 1 <= max_pieces_per_message <= 500: + raise ValueError("Conversation previews allow 1 through 100 turns and 1 through 500 pieces per turn.") return await self._run_database_operation_async( - self._execute_get_attack_result_conversation_projections, - attack_result_id=attack_result_id, - include_messages=include_messages, - limit_messages_per_conversation=limit_messages_per_conversation, - max_conversations=max_conversations, + self._fetch_bounded_conversation_messages, + conversation_id=conversation_id, + limit=limit, max_pieces_per_message=max_pieces_per_message, ) diff --git a/pyrit/models/scenario_hierarchy.py b/pyrit/models/scenario_hierarchy.py index fc51d6cdc9..8be73713e9 100644 --- a/pyrit/models/scenario_hierarchy.py +++ b/pyrit/models/scenario_hierarchy.py @@ -141,6 +141,9 @@ class ScenarioHierarchyNode(BaseModel): # Resolved children in stored order (bounded traversal) children: list["ScenarioHierarchyNode"] = Field(default_factory=list) children_truncated: bool = False + children_count: int | None = Field(default=None, ge=0) + child_offset: int = Field(default=0, ge=0) + next_children_cursor: str | None = None # Missing, malformed, cyclic, or cross-run references that could not be safely resolved unresolved_child_attack_result_ids: list[str] = Field(default_factory=list) @@ -209,6 +212,9 @@ class ScenarioHierarchyDiagnostics(BaseModel): attack_retries_truncated: bool = False total_retries: int = 0 + rootless_result_ids: list[str] = Field(default_factory=list) + rootless_results_count: int = 0 + rootless_results_truncated: bool = False class ScenarioHierarchyProjection(BaseModel): @@ -224,6 +230,8 @@ class ScenarioHierarchyProjection(BaseModel): overall_counts: ScenarioProgressCounts diagnostics: ScenarioHierarchyDiagnostics = Field(default_factory=ScenarioHierarchyDiagnostics) truncated: bool = False + has_more: bool = False + next_cursor: str | None = None ScenarioHierarchyNode.model_rebuild() diff --git a/tests/unit/backend/test_attack_service.py b/tests/unit/backend/test_attack_service.py index 062b6d12ff..6dd95f5ffb 100644 --- a/tests/unit/backend/test_attack_service.py +++ b/tests/unit/backend/test_attack_service.py @@ -919,7 +919,7 @@ async def test_get_conversation_messages_marks_attack_objective_score(self, atta conversation_id="test-id", ) - mock_mapper.assert_awaited_once_with([], objective_score_id=objective_score_id) + mock_mapper.assert_awaited_once_with([], objective_score_id=objective_score_id, memory=mock_memory) async def test_get_conversation_messages_preserves_string_objective_score_id( self, attack_service, mock_memory @@ -940,7 +940,7 @@ async def test_get_conversation_messages_preserves_string_objective_score_id( conversation_id="test-id", ) - mock_mapper.assert_awaited_once_with([], objective_score_id=objective_score_id) + mock_mapper.assert_awaited_once_with([], objective_score_id=objective_score_id, memory=mock_memory) async def test_get_conversation_messages_raises_for_unrelated_conversation( self, attack_service, mock_memory diff --git a/tests/unit/backend/test_saved_conversation_read_service.py b/tests/unit/backend/test_saved_conversation_read_service.py new file mode 100644 index 0000000000..a863365f4b --- /dev/null +++ b/tests/unit/backend/test_saved_conversation_read_service.py @@ -0,0 +1,189 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Canonical saved-conversation pagination without attack reconstruction or writes.""" + +import uuid +from datetime import UTC, datetime +from unittest.mock import AsyncMock, patch + +import pytest + +from pyrit.backend.services.saved_conversation_read_service import SavedConversationReadService +from pyrit.memory import PromptMemoryEntry, SQLiteMemory, set_message_piece_sha256_async +from pyrit.models import AttackResult, ConversationReference, ConversationType, MessagePiece + + +@pytest.fixture +async def saved_attack(sqlite_instance: SQLiteMemory) -> AttackResult: + attack = AttackResult( + conversation_id=str(uuid.uuid4()), + objective="Saved objective", + related_conversations={ + ConversationReference(conversation_id=str(uuid.uuid4()), conversation_type=kind) + for kind in (ConversationType.PRUNED, ConversationType.ADVERSARIAL, ConversationType.PREPARATION) + }, + ) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[attack]) + return attack + + +@pytest.fixture +async def saved_pieces(sqlite_instance: SQLiteMemory, saved_attack: AttackResult) -> list[MessagePiece]: + pieces = [ + MessagePiece( + id=uuid.UUID(int=index + 1), + conversation_id=saved_attack.conversation_id, + original_value=f"original-{index}", + converted_value=f"converted-{index}", + role="user" if index < 7 else "simulated_assistant", + sequence=-1 if index < 7 else 0, + timestamp=datetime(2026, 1, 1, tzinfo=UTC), + ) + for index in range(10) + ] + for piece in pieces: + await set_message_piece_sha256_async(piece) + await sqlite_instance.add_message_pieces_to_memory_async(message_pieces=pieces) + return pieces + + +async def test_reference_pages_include_every_saved_context_without_message_hydration_async( + sqlite_instance: SQLiteMemory, saved_attack: AttackResult, saved_pieces: list[MessagePiece] +) -> None: + reader = SavedConversationReadService(memory=sqlite_instance) + with ( + patch.object(sqlite_instance, "get_attack_results_async", new_callable=AsyncMock) as full_attacks, + patch.object(sqlite_instance, "get_conversation_messages_async", new_callable=AsyncMock) as full_messages, + patch.object(PromptMemoryEntry, "get_message_piece", autospec=True) as hydrate, + ): + first = await reader.list_conversations_async(attack_result_id=saved_attack.attack_result_id, limit=2) + assert first is not None + assert first.total_count == 4 + assert first.pagination is not None + assert first.pagination.has_more + assert first.pagination.next_cursor is not None + second = await reader.list_conversations_async( + attack_result_id=saved_attack.attack_result_id, limit=2, cursor=first.pagination.next_cursor + ) + assert second is not None + assert second.pagination is not None + assert not second.pagination.has_more + assert second.pagination.next_cursor is None + hydrate.assert_not_called() + full_attacks.assert_not_awaited() + full_messages.assert_not_awaited() + references = first.conversations + second.conversations + assert [reference.conversation_id for reference in references] == [ + saved_attack.conversation_id, + *( + next( + reference.conversation_id + for reference in saved_attack.related_conversations + if reference.conversation_type == kind + ) + for kind in (ConversationType.PRUNED, ConversationType.ADVERSARIAL, ConversationType.PREPARATION) + ), + ] + assert [reference.conversation_type for reference in references] == ["main", "pruned", "adversarial", "preparation"] + assert references[0].message_count == 2 + assert references[1].message_count == 0 + + +async def test_piece_pages_continue_through_negative_sequences_and_multimodal_turns_async( + sqlite_instance: SQLiteMemory, saved_attack: AttackResult, saved_pieces: list[MessagePiece] +) -> None: + reader = SavedConversationReadService(memory=sqlite_instance) + cursor = None + returned: list[str] = [] + partial_sequences: list[list[int]] = [] + with ( + patch.object(sqlite_instance, "get_attack_results_async", new_callable=AsyncMock) as full_attacks, + patch.object(sqlite_instance, "get_conversation_messages_async", new_callable=AsyncMock) as full_messages, + ): + while True: + page = await reader.get_messages_async( + attack_result_id=saved_attack.attack_result_id, + conversation_id=saved_attack.conversation_id, + limit=3, + cursor=cursor, + ) + assert page is not None + pieces = [piece for message in page.messages for piece in message.message_pieces] + assert len(pieces) <= 3 + returned.extend(str(piece.id) for piece in pieces) + partial_sequences.append(page.partial_message_sequences) + assert page.pagination is not None + cursor = page.pagination.next_cursor + if not page.pagination.has_more: + break + full_attacks.assert_not_awaited() + full_messages.assert_not_awaited() + assert returned == [str(piece.id) for piece in saved_pieces] + assert len(set(returned)) == len(saved_pieces) + assert partial_sequences == [[-1], [-1], [-1, 0], [0]] + + +async def test_message_reader_maps_only_the_requested_pieces_not_the_sentinel_async( + sqlite_instance: SQLiteMemory, saved_pieces: list[MessagePiece] +) -> None: + with patch.object( + PromptMemoryEntry, "get_message_piece", autospec=True, side_effect=PromptMemoryEntry.get_message_piece + ) as hydrate: + page = await sqlite_instance.get_conversation_message_page_async( + conversation_id=saved_pieces[0].conversation_id, limit=3 + ) + assert hydrate.call_count == 3 + assert page.next_cursor is not None + assert page.next_cursor.sequence == -1 + assert page.partial_message_sequences == (-1,) + assert [piece.id for message in page.messages for piece in message.message_pieces] == [ + piece.id for piece in saved_pieces[:3] + ] + + +async def test_saved_cursors_reject_other_owners_and_changed_reference_order_async( + sqlite_instance: SQLiteMemory, saved_attack: AttackResult, saved_pieces: list[MessagePiece] +) -> None: + reader = SavedConversationReadService(memory=sqlite_instance) + references = await reader.list_conversations_async(attack_result_id=saved_attack.attack_result_id, limit=1) + assert references is not None + assert references.pagination is not None + reference_cursor = references.pagination.next_cursor + assert reference_cursor is not None + other_attack = AttackResult(conversation_id=str(uuid.uuid4()), objective="Other saved attack") + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[other_attack]) + with pytest.raises(ValueError, match="cursor"): + await reader.list_conversations_async(attack_result_id=other_attack.attack_result_id, cursor=reference_cursor) + message_page = await reader.get_messages_async( + attack_result_id=saved_attack.attack_result_id, conversation_id=saved_attack.conversation_id, limit=1 + ) + assert message_page is not None + assert message_page.pagination is not None + with pytest.raises(ValueError, match="cursor"): + await reader.get_messages_async( + attack_result_id=saved_attack.attack_result_id, + conversation_id=next(iter(saved_attack.related_conversations)).conversation_id, + limit=1, + cursor=message_page.pagination.next_cursor, + ) + with pytest.raises(ValueError, match="not part of attack"): + await reader.get_messages_async( + attack_result_id=saved_attack.attack_result_id, conversation_id=other_attack.conversation_id + ) + await sqlite_instance.update_attack_result_by_id_async( + attack_result_id=saved_attack.attack_result_id, + update_fields={ + "pruned_conversation_ids": [ + str(uuid.uuid4()), + next( + reference.conversation_id + for reference in saved_attack.related_conversations + if reference.conversation_type == ConversationType.PRUNED + ), + ], + "timestamp": saved_attack.timestamp, + }, + ) + with pytest.raises(ValueError, match="cursor"): + await reader.list_conversations_async(attack_result_id=saved_attack.attack_result_id, cursor=reference_cursor) diff --git a/tests/unit/backend/test_scenario_hierarchy_pagination.py b/tests/unit/backend/test_scenario_hierarchy_pagination.py new file mode 100644 index 0000000000..dda40cc68e --- /dev/null +++ b/tests/unit/backend/test_scenario_hierarchy_pagination.py @@ -0,0 +1,388 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Paged hierarchy navigation, independent counts, and on-demand canonical conversations.""" + +import uuid +from datetime import UTC, datetime +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi.testclient import TestClient +from sqlalchemy import text +from sqlalchemy.dialects import sqlite + +from pyrit.backend.main import app +from pyrit.backend.services.attack_service import AttackService +from pyrit.backend.services.message_send_service import MessageSendService +from pyrit.backend.services.scenario_run_service import ScenarioRunService +from pyrit.memory import SQLiteMemory, set_message_piece_sha256_async +from pyrit.models import AttackOutcome, AttackResult, AttackResultRole, MessagePiece, ScenarioHierarchyNode +from unit.mocks import get_mock_target_identifier, make_scenario_result + + +async def _store_run_async(memory: SQLiteMemory) -> str: + identifier = uuid.uuid4() + await memory.add_scenario_results_to_memory_async( + scenario_results=[ + make_scenario_result( + id=identifier, + scenario_name="Hierarchy pagination", + scenario_version=1, + objective_target_identifier=get_mock_target_identifier(), + attack_results={}, + ) + ] + ) + return str(identifier) + + +def _result(*, run_id: str, index: int, child_ids: list[Any] | None = None) -> AttackResult: + return AttackResult( + attack_result_id=str(uuid.UUID(int=index + 100)), + conversation_id=str(uuid.UUID(int=index + 10000)), + objective=f"Objective {index}", + outcome=AttackOutcome.SUCCESS, + timestamp=datetime(2026, 1, 1, tzinfo=UTC), + attribution_parent_id=run_id, + attribution_data={ + "parent_collection": "atomic", + "seed_group_id": f"seed-{index}", + "result_role": ( + AttackResultRole.ORCHESTRATION.value if child_ids is not None else AttackResultRole.TARGET_FACING.value + ), + }, + metadata={"child_attack_result_ids": child_ids} if child_ids is not None else {}, + ) + + +def _flatten(nodes: list[ScenarioHierarchyNode]) -> list[ScenarioHierarchyNode]: + return [descendant for node in nodes for descendant in [node, *_flatten(node.children)]] + + +async def test_root_continuation_reserves_every_returned_root_and_counts_emitted_dag_nodes_async( + sqlite_instance: SQLiteMemory, +) -> None: + run_id = await _store_run_async(sqlite_instance) + children = [_result(run_id=run_id, index=index) for index in range(6, 9)] + child_ids = [child.attack_result_id for child in children] + roots = [_result(run_id=run_id, index=index) for index in range(4)] + roots += [ + _result(run_id=run_id, index=4, child_ids=child_ids), + _result(run_id=run_id, index=5, child_ids=child_ids[:2]), + ] + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[*roots, *children]) + service = ScenarioRunService(memory=sqlite_instance) + with patch.object(service._progress_read_model, "get_snapshot_async", new_callable=AsyncMock) as progress: + first = await service.get_run_hierarchy_async(scenario_result_id=run_id, max_nodes=4) + assert first is not None + assert first.has_more + assert first.next_cursor is not None + second = await service.get_run_hierarchy_async(scenario_result_id=run_id, max_nodes=4, cursor=first.next_cursor) + assert second is not None + progress.assert_not_awaited() + first_roots = [node for group in first.groups for node in group.roots] + second_roots = [node for group in second.groups for node in group.roots] + assert [node.attack_result_id for node in first_roots + second_roots] == [ + result.attack_result_id for result in roots + ] + assert len(_flatten(second_roots)) == 4 + assert len(second_roots[0].children) == 2 + assert second_roots[1].attack_result_id == roots[-1].attack_result_id + assert second_roots[1].children == [] + assert not second.has_more + assert second.next_cursor is None + assert second.overall_counts == first.overall_counts + assert second.truncated + + +@pytest.mark.parametrize("form", ["uppercase", "hex", "braced", "urn"]) +async def test_native_roots_and_node_mapping_agree_on_saved_uuid_forms_async( + sqlite_instance: SQLiteMemory, form: str +) -> None: + run_id = await _store_run_async(sqlite_instance) + child = _result(run_id=run_id, index=1001) + identifier = child.attack_result_id + reference = { + "uppercase": identifier.upper(), + "hex": uuid.UUID(identifier).hex.upper(), + "braced": "{" + identifier + "}", + "urn": "URN:UUID:" + identifier, + }[form] + parent = _result(run_id=run_id, index=1000, child_ids=[reference]) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[parent, child]) + hierarchy = await ScenarioRunService(memory=sqlite_instance).get_run_hierarchy_async( + scenario_result_id=run_id, max_nodes=5 + ) + assert hierarchy is not None + roots = [node for group in hierarchy.groups for node in group.roots] + assert [node.attack_result_id for node in roots] == [parent.attack_result_id] + assert roots[0].children[0].attack_result_id == child.attack_result_id + assert roots[0].children[0].child_index == 0 + + +async def test_native_roots_ignore_non_array_metadata_and_cross_run_parents_async( + sqlite_instance: SQLiteMemory, +) -> None: + run_id = await _store_run_async(sqlite_instance) + other_run = await _store_run_async(sqlite_instance) + results = [_result(run_id=run_id, index=index) for index in range(5)] + results[3].metadata = {"child_attack_result_ids": results[0].attack_result_id} + results[4].metadata = {"child_attack_result_ids": {"child": results[1].attack_result_id}} + foreign_parent = _result(run_id=other_run, index=20, child_ids=[results[2].attack_result_id]) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[*results, foreign_parent]) + roots, has_more = await sqlite_instance.get_scenario_root_deltas_async(scenario_result_id=run_id, limit=10) + assert {root.attack_result_id for root in roots} == {result.attack_result_id for result in results} + assert not has_more + + +@pytest.mark.parametrize("reference", [None, True, 7, {}, [], "not-a-uuid", "00000000-0000-0000-0000-000000000064 "]) +async def test_malformed_references_cannot_hide_saved_roots_async( + sqlite_instance: SQLiteMemory, reference: Any +) -> None: + run_id = await _store_run_async(sqlite_instance) + child = _result(run_id=run_id, index=0) + parent = _result(run_id=run_id, index=1, child_ids=[reference]) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[child, parent]) + roots, has_more = await sqlite_instance.get_scenario_root_deltas_async(scenario_result_id=run_id, limit=2) + assert {root.attack_result_id for root in roots} == {child.attack_result_id, parent.attack_result_id} + assert not has_more + + +def test_native_root_query_expands_saved_references_once_not_once_per_candidate(sqlite_instance: SQLiteMemory) -> None: + statement = sqlite_instance._build_scenario_root_deltas_statement(scenario_result_id=str(uuid.uuid4()), limit=2) + compiled = statement.compile(dialect=sqlite.dialect(), compile_kwargs={"literal_binds": True}) + with sqlite_instance._get_session() as session: + plan = session.execute(text(f"EXPLAIN QUERY PLAN {compiled}")).all() + details = " ".join(str(row[-1]).upper() for row in plan) + assert "LIST SUBQUERY" in details + assert "CORRELATED" not in details + + +async def test_empty_run_has_an_empty_page_and_zero_exact_counts_async(sqlite_instance: SQLiteMemory) -> None: + run_id = await _store_run_async(sqlite_instance) + page = await ScenarioRunService(memory=sqlite_instance).get_run_hierarchy_async(scenario_result_id=run_id) + assert page is not None + assert page.groups == [] + assert page.overall_counts.completed == 0 + assert page.diagnostics.total_retries == 0 + assert not page.has_more + assert page.next_cursor is None + assert not page.truncated + + +async def test_child_pages_preserve_malformed_positions_and_continue_past_empty_pages_async( + sqlite_instance: SQLiteMemory, +) -> None: + run_id = await _store_run_async(sqlite_instance) + child = _result(run_id=run_id, index=1) + missing_id = str(uuid.uuid4()) + parent = _result( + run_id=run_id, + index=0, + child_ids=[child.attack_result_id.upper(), 7, None, {}, missing_id, child.attack_result_id], + ) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[parent, child]) + service = ScenarioRunService(memory=sqlite_instance) + first = await service.get_run_hierarchy_node_async( + scenario_result_id=run_id, attack_result_id=parent.attack_result_id, max_immediate_children=2 + ) + assert first is not None + assert first.children_count == 6 + assert first.children[0].child_index == 0 + assert first.unresolved_child_attack_result_ids == [""] + assert first.next_children_cursor is not None + second = await service.get_run_hierarchy_node_async( + scenario_result_id=run_id, + attack_result_id=parent.attack_result_id, + max_immediate_children=2, + children_cursor=first.next_children_cursor, + ) + assert second is not None + assert second.child_offset == 2 + assert second.children == [] + assert len(second.unresolved_child_attack_result_ids) == 2 + assert second.next_children_cursor is not None + third = await service.get_run_hierarchy_node_async( + scenario_result_id=run_id, + attack_result_id=parent.attack_result_id, + max_immediate_children=2, + children_cursor=second.next_children_cursor, + ) + assert third is not None + assert third.child_offset == 4 + assert third.children == [] + assert third.unresolved_child_attack_result_ids == [missing_id, child.attack_result_id] + assert third.next_children_cursor is None + + +async def test_children_beyond_the_old_500_reference_cap_are_accessible_async(sqlite_instance: SQLiteMemory) -> None: + run_id = await _store_run_async(sqlite_instance) + child_ids = [str(uuid.UUID(int=index + 100)) for index in range(503)] + parent = _result(run_id=run_id, index=1000, child_ids=child_ids) + await sqlite_instance.add_attack_results_to_memory_async( + attack_results=[parent, _result(run_id=run_id, index=0), _result(run_id=run_id, index=502)] + ) + service = ScenarioRunService(memory=sqlite_instance) + first = await service.get_run_hierarchy_node_async( + scenario_result_id=run_id, attack_result_id=parent.attack_result_id, max_immediate_children=500 + ) + assert first is not None + assert len(first.children) == 1 + assert len(first.child_attack_result_ids) == 500 + assert first.next_children_cursor is not None + last = await service.get_run_hierarchy_node_async( + scenario_result_id=run_id, + attack_result_id=parent.attack_result_id, + max_immediate_children=500, + children_cursor=first.next_children_cursor, + ) + assert last is not None + assert last.child_offset == 500 + assert last.children[0].child_index == 502 + assert last.children[0].attack_result_id == child_ids[502] + assert last.next_children_cursor is None + + +async def test_rootless_cycles_are_diagnostic_samples_not_invented_roots_async(sqlite_instance: SQLiteMemory) -> None: + run_id = await _store_run_async(sqlite_instance) + first = _result(run_id=run_id, index=0, child_ids=[str(uuid.UUID(int=101))]) + second = _result(run_id=run_id, index=1, child_ids=[first.attack_result_id]) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[first, second]) + service = ScenarioRunService(memory=sqlite_instance) + hierarchy = await service.get_run_hierarchy_async(scenario_result_id=run_id, max_nodes=1) + assert hierarchy is not None + assert not any(group.roots for group in hierarchy.groups) + assert hierarchy.truncated + assert not hierarchy.has_more + assert hierarchy.diagnostics.rootless_results_count == 2 + assert hierarchy.diagnostics.rootless_result_ids == [first.attack_result_id] + assert hierarchy.diagnostics.rootless_results_truncated + node = await service.get_run_hierarchy_node_async( + scenario_result_id=run_id, attack_result_id=hierarchy.diagnostics.rootless_result_ids[0] + ) + assert node is not None + assert node.children[0].attack_result_id == second.attack_result_id + + +@pytest.mark.parametrize("has_resolvable_tail", [False, True]) +async def test_preview_does_not_probe_every_missing_reference_to_fill_the_node_budget_async( + sqlite_instance: SQLiteMemory, has_resolvable_tail: bool +) -> None: + run_id = await _store_run_async(sqlite_instance) + tail = _result(run_id=run_id, index=1) + parent = _result( + run_id=run_id, index=0, child_ids=[*[str(uuid.uuid4()) for _ in range(499)], tail.attack_result_id] + ) + await sqlite_instance.add_attack_results_to_memory_async( + attack_results=[parent, tail] if has_resolvable_tail else [parent] + ) + service = ScenarioRunService(memory=sqlite_instance) + with patch.object( + sqlite_instance, + "get_scenario_attack_result_deltas_by_ids_async", + wraps=sqlite_instance.get_scenario_attack_result_deltas_by_ids_async, + ) as descendants: + hierarchy = await service.get_run_hierarchy_async(scenario_result_id=run_id, max_nodes=2) + assert hierarchy is not None + descendants.assert_awaited_once() + assert len(descendants.call_args.kwargs["attack_result_ids"]) == 1 + root = hierarchy.groups[0].roots[0] + assert len(root.unresolved_child_attack_result_ids) == 500 + assert root.children_truncated is has_resolvable_tail + if has_resolvable_tail: + node = await service.get_run_hierarchy_node_async( + scenario_result_id=run_id, attack_result_id=parent.attack_result_id, max_immediate_children=500 + ) + assert node is not None + assert node.children[0].attack_result_id == tail.attack_result_id + assert node.children[0].child_index == 499 + + +async def test_child_cursors_reject_other_nodes_and_reordered_references_async(sqlite_instance: SQLiteMemory) -> None: + run_id = await _store_run_async(sqlite_instance) + child_ids = [str(uuid.uuid4()), str(uuid.uuid4())] + parent = _result(run_id=run_id, index=0, child_ids=child_ids) + other = _result(run_id=run_id, index=1, child_ids=child_ids) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[parent, other]) + service = ScenarioRunService(memory=sqlite_instance) + first = await service.get_run_hierarchy_node_async( + scenario_result_id=run_id, attack_result_id=parent.attack_result_id, max_immediate_children=1 + ) + assert first is not None + assert first.next_children_cursor is not None + with pytest.raises(ValueError, match="cursor"): + await service.get_run_hierarchy_node_async( + scenario_result_id=run_id, + attack_result_id=other.attack_result_id, + children_cursor=first.next_children_cursor, + ) + await sqlite_instance.update_attack_result_by_id_async( + attack_result_id=parent.attack_result_id, + update_fields={"attack_metadata": {"child_attack_result_ids": child_ids[::-1]}, "timestamp": parent.timestamp}, + ) + with pytest.raises(ValueError, match="cursor"): + await service.get_run_hierarchy_node_async( + scenario_result_id=run_id, + attack_result_id=parent.attack_result_id, + children_cursor=first.next_children_cursor, + ) + + +async def test_api_journey_loads_only_the_selected_saved_conversation_async( + sqlite_instance: SQLiteMemory, compatibility_headers: dict[str, str] +) -> None: + run_id = await _store_run_async(sqlite_instance) + child = _result(run_id=run_id, index=1) + parent = _result(run_id=run_id, index=0, child_ids=[child.attack_result_id]) + await sqlite_instance.add_attack_results_to_memory_async(attack_results=[parent, child]) + pieces = [ + MessagePiece( + conversation_id=child.conversation_id, + original_value=f"Saved piece {index}", + role="user", + sequence=0, + ) + for index in range(3) + ] + for piece in pieces: + await set_message_piece_sha256_async(piece) + await sqlite_instance.add_message_pieces_to_memory_async(message_pieces=pieces) + scenario_service = ScenarioRunService(memory=sqlite_instance) + sends = MagicMock(spec=MessageSendService) + attack_service = AttackService(memory=sqlite_instance, message_send_service=sends) + client = TestClient(app, headers=compatibility_headers) + with ( + patch("pyrit.backend.routes.scenarios.get_scenario_run_service", return_value=scenario_service), + patch("pyrit.backend.routes.attacks.get_attack_service", return_value=attack_service), + ): + roots = client.get(f"/api/scenarios/runs/{run_id}/hierarchy", params={"max_nodes": 1}) + assert roots.status_code == 200 + assert roots.json()["groups"][0]["roots"][0]["attack_result_id"] == parent.attack_result_id + node = client.get(f"/api/scenarios/runs/{run_id}/hierarchy/nodes/{child.attack_result_id}") + assert node.status_code == 200 + assert node.json()["conversations"][0]["messages"] == [] + references = client.get(f"/api/attacks/{child.attack_result_id}/conversations", params={"limit": 1}) + assert references.status_code == 200 + assert references.json()["conversations"][0]["message_count"] == 1 + messages_url = f"/api/attacks/{child.attack_result_id}/messages" + first = client.get(messages_url, params={"conversation_id": child.conversation_id, "limit": 2}) + assert first.status_code == 200 + first_body = first.json() + assert first_body["partial_message_sequences"] == [0] + second = client.get( + messages_url, + params={ + "conversation_id": child.conversation_id, + "limit": 2, + "cursor": first_body["pagination"]["next_cursor"], + }, + ) + assert second.status_code == 200 + all_pieces = first_body["messages"][0]["message_pieces"] + second.json()["messages"][0]["message_pieces"] + assert len(all_pieces) == 3 + assert {piece["id"] for piece in all_pieces} == {str(piece.id) for piece in pieces} + bad_cursor = client.get(f"/api/scenarios/runs/{run_id}/hierarchy", params={"cursor": "malformed"}) + assert bad_cursor.status_code == 400 + assert sends.mock_calls == [] diff --git a/tests/unit/backend/test_scenario_hierarchy_parity.py b/tests/unit/backend/test_scenario_hierarchy_parity.py index 8b2e7bb4b6..c355bdd766 100644 --- a/tests/unit/backend/test_scenario_hierarchy_parity.py +++ b/tests/unit/backend/test_scenario_hierarchy_parity.py @@ -16,6 +16,7 @@ from pyrit.analytics import compute_outcome_statistics from pyrit.backend.main import app +from pyrit.backend.services.saved_conversation_read_service import SavedConversationReadService from pyrit.backend.services.scenario_progress_read_model import ( ScenarioProgressReadModel, ScenarioProgressSnapshot, @@ -216,6 +217,22 @@ def _make_test_snapshot( ) +def test_exact_retry_diagnostics_override_partial_preview_arithmetic() -> None: + snapshot = _make_test_snapshot( + deltas=[_make_test_delta(attack_result_id="first"), _make_test_delta(attack_result_id="second")] + ) + assert snapshot.summary.overall.retries == 1 + projection = ScenarioProgressReadModel.build_scenario_hierarchy( + scenario_result_id="scenario-run-1", + scenario_name="TestScenario", + scenario_version=1, + status=ScenarioRunState.COMPLETED, + snapshot=snapshot, + exact_total_retries=0, + ) + assert projection.diagnostics.total_retries == 0 + + 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) @@ -477,7 +494,8 @@ 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( + conversation_reader = MagicMock(spec=SavedConversationReadService) + conversation_reader.get_legacy_projections_async = AsyncMock( return_value=ScenarioAttackResultConversations( conversations=[ ScenarioConversationProjection( @@ -500,6 +518,7 @@ async def mock_get_deltas( read_model = ScenarioProgressReadModel(memory=memory) service = ScenarioRunService(memory=memory, progress_read_model=read_model) + service._saved_conversation_reader = conversation_reader node = await service.get_run_hierarchy_node_async( scenario_result_id="scenario-run-1", @@ -511,8 +530,9 @@ async def mock_get_deltas( 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" + conversation_reader.get_legacy_projections_async.assert_awaited_once_with( + attack_result_id="parent-1", include_messages=False + ) # Querying non-existent node returns None missing = await service.get_run_hierarchy_node_async( @@ -687,7 +707,7 @@ async def test_sqlite_memory_batched_reads_and_projections_parity(sqlite_instanc assert cross_run_id not in delta_ids # 2. Test conversation projections - conv_projections = await sqlite_instance.get_attack_result_conversation_projections_async( + conv_projections = await SavedConversationReadService(memory=sqlite_instance).get_legacy_projections_async( attack_result_id=child_id, include_messages=True, ) @@ -1112,7 +1132,7 @@ async def test_sqlite_bounded_multimodal_conversation_projection(sqlite_instance 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( + projs = await SavedConversationReadService(memory=sqlite_instance).get_legacy_projections_async( attack_result_id=ar_id, include_messages=True, limit_messages_per_conversation=2, @@ -1251,7 +1271,7 @@ async def test_sqlite_conversation_list_clipping_explicit(sqlite_instance: SQLit 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( + conv_result = await SavedConversationReadService(memory=sqlite_instance).get_legacy_projections_async( attack_result_id=ar_id, max_conversations=50, ) @@ -1410,7 +1430,7 @@ async def test_sqlite_message_pieces_bounding(sqlite_instance: SQLiteMemory) -> ) # 1. Fetch with default max_pieces_per_message=100 - projs = await sqlite_instance.get_attack_result_conversation_projections_async( + projs = await SavedConversationReadService(memory=sqlite_instance).get_legacy_projections_async( attack_result_id=ar_id, include_messages=True, max_pieces_per_message=100, @@ -1422,7 +1442,7 @@ async def test_sqlite_message_pieces_bounding(sqlite_instance: SQLiteMemory) -> 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( + projs_untruncated = await SavedConversationReadService(memory=sqlite_instance).get_legacy_projections_async( attack_result_id=ar_id, include_messages=True, max_pieces_per_message=200, @@ -1791,7 +1811,6 @@ async def test_persisted_conversation_statistics_parity_between_hierarchy_and_no 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, @@ -1799,8 +1818,14 @@ async def test_persisted_conversation_statistics_parity_between_hierarchy_and_no 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 + assert node.conversations[0].messages == [] + assert node.conversations[0].messages_truncated is True + explicit_preview = await service.get_run_hierarchy_node_async( + scenario_result_id=scenario_id, attack_result_id=ar_id, include_messages=True + ) + assert explicit_preview is not None + assert len(explicit_preview.conversations[0].messages) == 3 + assert explicit_preview.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) @@ -1914,8 +1939,7 @@ async def test_bounded_conversation_messages_piece_budget_applied_at_sql_layer( ) 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( + projs = await SavedConversationReadService(memory=sqlite_instance).get_legacy_projections_async( attack_result_id=ar_id, include_messages=True, max_pieces_per_message=100, @@ -1936,18 +1960,14 @@ async def test_bounded_conversation_messages_piece_budget_applied_at_sql_layer( 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( + messages_capped_10, _, partial_sequences = await sqlite_instance.get_bounded_conversation_messages_async( 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 + assert len(messages_capped_10[0].message_pieces) == 10 + assert len(messages_capped_10[1].message_pieces) == 5 + assert partial_sequences == {0} async def test_diagnostics_total_retries_count_event_mismatch_parity( @@ -2162,7 +2182,6 @@ async def test_parent_with_2500_children_bounds_unresolved_references(sqlite_ins 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, @@ -2170,8 +2189,20 @@ async def test_parent_with_2500_children_bounds_unresolved_references(sqlite_ins ) 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 + assert node.unresolved_child_attack_result_ids == [] + assert node.children_count == 2500 + assert node.child_attack_result_ids == child_ids[:200] + assert node.next_children_cursor is not None + missing_page = await service.get_run_hierarchy_node_async( + scenario_result_id=scenario_id, + attack_result_id=parent_id, + max_immediate_children=200, + children_cursor=node.next_children_cursor, + ) + assert missing_page is not None + assert missing_page.children == [] + assert missing_page.unresolved_child_attack_result_ids == child_ids[200:400] + assert missing_page.next_children_cursor is not None # 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) @@ -2209,25 +2240,29 @@ async def test_cold_hierarchy_hydration_bounded_by_max_nodes(sqlite_instance: SQ ] 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) + with ( + patch.object( + sqlite_instance, "get_scenario_root_deltas_async", wraps=sqlite_instance.get_scenario_root_deltas_async + ) as roots, + patch.object( + sqlite_instance, + "_row_to_scenario_attack_result_delta", + wraps=sqlite_instance._row_to_scenario_attack_result_delta, + ) as hydrate, + patch.object(sqlite_instance, "get_scenario_attack_result_deltas_async", new_callable=AsyncMock) as full_deltas, + patch.object(service._progress_read_model, "get_snapshot_async", new_callable=AsyncMock) as progress, + ): + hierarchy = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + roots.assert_awaited_once() + assert roots.call_args.kwargs["limit"] == 1 + assert hydrate.call_count == 1 + full_deltas.assert_not_awaited() + progress.assert_not_awaited() 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 + assert hierarchy.has_more is True + assert hierarchy.next_cursor is not None def test_build_scenario_hierarchy_unqueried_conversation_not_false_zero() -> None: @@ -3222,12 +3257,10 @@ async def test_500_cyclic_results_with_outer_parent_at_501_parity(sqlite_instanc 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. - """ +async def test_shared_child_with_1001_parents_has_bounded_native_root_pages_async( + sqlite_instance: SQLiteMemory, +) -> None: + """Page the actual parent roots without reverse-discovering or hydrating their complete fan-in.""" scenario_uuid = uuid.uuid4() scenario_id = str(scenario_uuid) base_time = datetime(2026, 1, 1, 0, 0, 0, tzinfo=UTC) @@ -3274,21 +3307,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, has_more = await sqlite_instance.get_scenario_parent_deltas_for_child_ids_async( - scenario_result_id=scenario_id, - child_ids=[child_id], - limit=500, - ) - assert len(deltas) == 500 - assert has_more is True - - deltas_10, has_more_10 = await sqlite_instance.get_scenario_parent_deltas_for_child_ids_async( - scenario_result_id=scenario_id, - child_ids=[child_id], - limit=10, - ) - assert len(deltas_10) == 10 - assert has_more_10 is True + service = ScenarioRunService(memory=sqlite_instance) + with patch.object( + sqlite_instance, + "_row_to_scenario_attack_result_delta", + wraps=sqlite_instance._row_to_scenario_attack_result_delta, + ) as mapped: + page = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=10) + assert page is not None + roots = [node for group in page.groups for node in group.roots] + assert len(roots) == 10 + assert {root.attack_result_id for root in roots} == {parent.attack_result_id for parent in parents[:10]} + assert all(not root.children for root in roots) + assert page.has_more + assert page.next_cursor is not None + assert mapped.call_count == 10 @pytest.mark.asyncio diff --git a/tests/unit/backend/test_scenario_progress_read_model.py b/tests/unit/backend/test_scenario_progress_read_model.py index 16ec6ee78b..57efdfa111 100644 --- a/tests/unit/backend/test_scenario_progress_read_model.py +++ b/tests/unit/backend/test_scenario_progress_read_model.py @@ -1097,7 +1097,13 @@ def _strip(obj): ), ({"parent_collection": "attack", "attempt_index": True}, {}, "unknown", [], None), ({"parent_collection": "attack", "attempt_index": "2"}, {"child_attack_result_ids": "c1"}, "unknown", [], None), - ({"parent_collection": "attack"}, {"child_attack_result_ids": ["c1", 7]}, "unknown", [], None), + ( + {"parent_collection": "attack"}, + {"child_attack_result_ids": ["c1", 7]}, + "unknown", + ["c1", ""], + None, + ), ], ) def test_map_progress_delta_reads_legacy_and_malformed_rows_conservatively( diff --git a/tests/unit/common/test_pagination.py b/tests/unit/common/test_pagination.py index afcc6eb041..e2cf34bb8f 100644 --- a/tests/unit/common/test_pagination.py +++ b/tests/unit/common/test_pagination.py @@ -71,6 +71,32 @@ def test_encode_keyset_cursor_preserves_history_wire_format() -> None: } +@pytest.mark.parametrize("position", [-1, 0, 500]) +def test_keyset_cursor_preserves_ordered_positions(position: int) -> None: + cursor = encode_keyset_cursor( + timestamp=datetime(2026, 9, 21, tzinfo=UTC), + identifier="00000000-0000-0000-0000-000000000001", + fingerprint="cohort", + position=position, + ) + decoded = decode_keyset_cursor(cursor=cursor, fingerprint="cohort") + assert decoded is not None + assert decoded.position == position + + +@pytest.mark.parametrize("position", [True, False, "1", 1.0, [], {}]) +def test_keyset_cursor_rejects_non_integer_positions(position: object) -> None: + payload = { + "v": 1, + "f": "cohort", + "t": "2026-09-21T00:00:00+00:00", + "i": "00000000-0000-0000-0000-000000000001", + "p": position, + } + cursor = base64.urlsafe_b64encode(json.dumps(payload).encode()).decode() + assert decode_keyset_cursor(cursor=cursor, fingerprint="cohort") is None + + def test_fingerprint_filters_preserves_history_default_digest() -> None: assert fingerprint_filters(filters={}) == "44136fa355b3678a" assert fingerprint_filters(filters={}, length=64) == ( diff --git a/tests/unit/memory/test_azure_sql_memory.py b/tests/unit/memory/test_azure_sql_memory.py index f07b056809..048bf51f14 100644 --- a/tests/unit/memory/test_azure_sql_memory.py +++ b/tests/unit/memory/test_azure_sql_memory.py @@ -16,7 +16,7 @@ from pyrit.common.singleton import Singleton from pyrit.converter.base64_converter import Base64Converter from pyrit.memory import AzureSQLMemory, EmbeddingDataEntry, MemoryInterface, PromptMemoryEntry -from pyrit.memory.memory_interface import AttackResultKeysetCursor, _AttackResultQuery +from pyrit.memory.memory_interface import AttackResultKeysetCursor, MessagePieceKeysetCursor, _AttackResultQuery from pyrit.memory.memory_models import ScenarioResultEntry from pyrit.memory.storage.serializers import set_message_piece_sha256_async from pyrit.models import Conversation, MessagePiece @@ -1274,53 +1274,47 @@ def test_scenario_attack_result_deltas_by_ids_statement_type_hints() -> None: assert hints.get("return") is Select -def test_scenario_parent_deltas_for_child_ids_bounds_parameters_with_large_exclusions_mssql( - memory_interface: AzureSQLMemory, +def test_hierarchy_roots_compile_native_array_matching_and_bounded_keysets_mssql( + uninitialized_memory_interface: AzureSQLMemory, ) -> None: - """Verify that _execute_get_scenario_parent_deltas_for_child_ids never exceeds SQL Server bind limits.""" - scenario_id = str(uuid.uuid4()) - # Case 1: 1,100 children and 1,100 exclusions - previously combined to ~2,900 params exceeding 2,100 MSSQL limit - child_ids = [str(uuid.uuid4()) for _ in range(1100)] - exclude_ids = [str(uuid.uuid4()) for _ in range(1100)] - - session = MagicMock() - session.execute.return_value.mappings.return_value.all.return_value = [] - with patch.object(memory_interface, "_get_session", return_value=session): - deltas, has_more = memory_interface._execute_get_scenario_parent_deltas_for_child_ids( - scenario_result_id=scenario_id, - child_ids=child_ids, - exclude_result_ids=exclude_ids, - ) - - assert deltas == [] - assert has_more is False - assert session.execute.call_count >= 1 - for call in session.execute.call_args_list: - stmt = call.args[0] - compiled = stmt.compile(dialect=mssql.dialect()) - assert len(compiled.params) <= memory_interface._MAX_BIND_VARS - assert len(compiled.params) <= 2000 - - # Case 2: 1,000 children and 2,500 exclusions (budget boundary) - child_ids_2 = [str(uuid.uuid4()) for _ in range(1000)] - exclude_ids_2 = [str(uuid.uuid4()) for _ in range(2500)] - - session.reset_mock() - with patch.object(memory_interface, "_get_session", return_value=session): - deltas, has_more = memory_interface._execute_get_scenario_parent_deltas_for_child_ids( - scenario_result_id=scenario_id, - child_ids=child_ids_2, - exclude_result_ids=exclude_ids_2, - ) - - assert deltas == [] - assert has_more is False - assert session.execute.call_count >= 1 - for call in session.execute.call_args_list: - stmt = call.args[0] - compiled = stmt.compile(dialect=mssql.dialect()) - assert len(compiled.params) <= memory_interface._MAX_BIND_VARS - assert len(compiled.params) <= 2000 + statement = uninitialized_memory_interface._build_scenario_root_deltas_statement( + scenario_result_id=str(uuid.uuid4()), + cursor=AttackResultKeysetCursor(timestamp=datetime(2026, 1, 1, tzinfo=UTC), attack_result_id=str(uuid.uuid4())), + limit=5, + ) + compiled = statement.compile(dialect=mssql.dialect()) + sql = str(compiled) + assert "OPENJSON" in sql + assert "OUTER APPLY" in sql + assert "JSON_QUERY" in sql + assert "NOT IN" in sql + assert "IS NOT NULL" in sql + assert "DATALENGTH" in sql + assert "Latin1_General_100_BIN2" in sql + assert "TOP" in sql + assert "LOWER(CAST" in sql.upper() + assert len(compiled.params) < 30 + + +def test_message_piece_page_compiles_sequence_timestamp_uuid_order_mssql( + uninitialized_memory_interface: AzureSQLMemory, +) -> None: + statement = uninitialized_memory_interface._build_conversation_message_page_statement( + conversation_id=str(uuid.uuid4()), + cursor=MessagePieceKeysetCursor( + sequence=-1, timestamp=datetime(2026, 1, 1, tzinfo=UTC), message_piece_id=str(uuid.uuid4()) + ), + limit=3, + ) + compiled = statement.compile(dialect=mssql.dialect()) + sql = str(compiled) + assert "TOP" in sql + assert "ORDER BY" in sql + assert "sequence ASC" in sql + assert "timestamp ASC" in sql + assert "LOWER(CAST" in sql.upper() + assert -1 in compiled.params.values() + assert len(compiled.params) < 15 def test_get_metadata_conditions_binds_special_character_keys( From 56dc1a79c63c1ef72a2130edae62be7c7a8159ff Mon Sep 17 00:00:00 2001 From: Roman Lutz Date: Mon, 26 Oct 2026 00:03:02 -0700 Subject: [PATCH 21/22] refactor: remove duplicate hierarchy conversation previews Reuse canonical conversation metadata for selected nodes and load all transcript bodies through the generic saved-conversation API. Remove unreleased projection models, preview adapters, and eager tree conversation statistics. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 60a3baa0-0e4d-428b-8062-5101ad7866dd --- doc/gui/0_gui.md | 5 +- pyrit/backend/README.md | 11 +- pyrit/backend/models/attacks.py | 11 +- pyrit/backend/routes/scenarios.py | 2 - .../saved_conversation_read_service.py | 98 +--- .../services/scenario_progress_read_model.py | 65 +-- .../backend/services/scenario_run_service.py | 68 +-- pyrit/memory/memory_interface.py | 95 ---- pyrit/models/__init__.py | 10 +- pyrit/models/conversation_summary.py | 16 + pyrit/models/scenario_hierarchy.py | 62 +-- .../test_scenario_hierarchy_pagination.py | 3 +- .../backend/test_scenario_hierarchy_parity.py | 511 +++++++----------- .../models/test_scenario_hierarchy_models.py | 105 +--- 14 files changed, 247 insertions(+), 815 deletions(-) create mode 100644 pyrit/models/conversation_summary.py diff --git a/doc/gui/0_gui.md b/doc/gui/0_gui.md index 25dcb6f1d7..6eea876ca7 100644 --- a/doc/gui/0_gui.md +++ b/doc/gui/0_gui.md @@ -648,8 +648,9 @@ Until you expand or collapse the section, its default follows the current group Compound-attack hierarchy browsing is a separate GUI follow-up. The backend provides paged roots and ordered child expansion, with exact run statistics independent of -the displayed page. Node reads return conversation metadata by default; clients -can page the selected saved transcript without loading every conversation. +the displayed page. Tree previews defer conversation reads. Node expansion uses the +same conversation metadata as the generic attack API; clients load transcript bodies +only through the selected saved conversation's paginated message endpoint. See the [backend hierarchy API](../../pyrit/backend/README.md#scenario-result-hierarchy) for continuation and clipping semantics. diff --git a/pyrit/backend/README.md b/pyrit/backend/README.md index b1155bf945..83ef374dd2 100644 --- a/pyrit/backend/README.md +++ b/pyrit/backend/README.md @@ -97,10 +97,13 @@ remain explicit and consume positions, so even a wholly unresolved page can adva Child cursors validate the run, node, timestamp, and saved reference ordering; restart child paging after those references change. -Node expansion returns conversation metadata only by default. Use the generic -saved-conversation APIs above for continuation through any selected transcript. -`include_messages=true` explicitly opts into the older clipped conversation -previews, not a complete transcript. +Tree previews leave conversation metadata deferred: `conversations` is empty, +`conversations_count` is unknown, and `conversations_truncated` is true when saved +references exist. The node's `conversation_id` still links to its main conversation. +Node expansion returns up to 50 canonical conversation summaries, the exact reference +count, and a clipping flag. These use the same metadata model as the generic attack API. +Use the saved-conversation APIs above for reference continuation and all transcript +bodies. There is no hierarchy-specific transcript format or conversation-preview mode. An entirely rootless saved graph, such as a closed cycle, returns a bounded `diagnostics.rootless_result_ids` sample with its total count and clipping flag. diff --git a/pyrit/backend/models/attacks.py b/pyrit/backend/models/attacks.py index 3ed55acef4..7c16fa51b7 100644 --- a/pyrit/backend/models/attacks.py +++ b/pyrit/backend/models/attacks.py @@ -27,6 +27,7 @@ AttackResult, ChatMessageRole, ConversationReference, + ConversationSummary, Message, MessagePiece, PromptDataType, @@ -563,16 +564,6 @@ def _validate_update(self) -> "UpdateAttackRequest": # ============================================================================ -class ConversationSummary(BaseModel): - """Summary of a conversation (message count, preview, timestamp).""" - - conversation_id: str = Field(..., description="Unique conversation identifier") - message_count: int = Field(0, description="Number of messages in this conversation") - last_message_preview: str | None = Field(None, description="Preview of the last message") - created_at: datetime | None = Field(None, description="Timestamp of the first message") - conversation_type: str | None = None - - class AttackConversationsResponse(BaseModel): """Response listing all conversations belonging to an attack.""" diff --git a/pyrit/backend/routes/scenarios.py b/pyrit/backend/routes/scenarios.py index cf3fea8479..ad99eda09b 100644 --- a/pyrit/backend/routes/scenarios.py +++ b/pyrit/backend/routes/scenarios.py @@ -432,7 +432,6 @@ async def get_scenario_run_hierarchy_node( # pyrit-async-suffix-exempt attack_result_id: IdentifierStr, limit: int = Query(200, ge=1, le=500, description="Maximum ordered child references in this page"), children_cursor: CursorStr | None = Query(None, description="Opaque ordered-child continuation cursor"), - include_messages: bool = Query(False, description="Explicitly include legacy clipped conversation previews"), ) -> ScenarioHierarchyNode: """ Get single-node hierarchy details and generic conversation projections without full-run hydration. @@ -447,7 +446,6 @@ async def get_scenario_run_hierarchy_node( # pyrit-async-suffix-exempt attack_result_id=attack_result_id, max_immediate_children=limit, children_cursor=children_cursor, - include_messages=include_messages, ) except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from exc diff --git a/pyrit/backend/services/saved_conversation_read_service.py b/pyrit/backend/services/saved_conversation_read_service.py index c32e6fa75f..784ae3406f 100644 --- a/pyrit/backend/services/saved_conversation_read_service.py +++ b/pyrit/backend/services/saved_conversation_read_service.py @@ -4,39 +4,23 @@ """Bounded, read-only saved conversation views shared by attack and scenario APIs.""" import uuid -from typing import ClassVar from pyrit.backend.mappers._preview import format_last_message_preview from pyrit.backend.mappers.attack_mappers import pyrit_messages_to_dto_async from pyrit.backend.models.attacks import ( AttackConversationsResponse, ConversationMessagesResponse, - ConversationSummary, ) from pyrit.backend.models.common import PaginationInfo from pyrit.common.pagination import decode_keyset_cursor, encode_keyset_cursor, fingerprint_filters from pyrit.memory import MemoryInterface from pyrit.memory.memory_interface import MessagePieceKeysetCursor -from pyrit.models import ( - ConversationStats, - ConversationType, - ScenarioAttackResultConversations, - ScenarioConversationMessagePieceProjection, - ScenarioConversationMessageProjection, - ScenarioConversationProjection, -) +from pyrit.models import ConversationStats, ConversationSummary class SavedConversationReadService: """Read canonical saved references/messages without constructing an executable attack.""" - _LEGACY_DETAILS: ClassVar[dict[ConversationType | None, tuple[str, str]]] = { - None: ("main", "Main objective target conversation"), - ConversationType.PRUNED: ("pruned", "Pruned conversation"), - ConversationType.ADVERSARIAL: ("adversarial_chat", "Adversarial chat conversation"), - ConversationType.PREPARATION: ("preparation", "Preparation conversation"), - } - def __init__(self, *, memory: MemoryInterface) -> None: """Bind all saved reads to the caller's memory instance.""" self._memory = memory @@ -181,83 +165,3 @@ async def get_messages_async( pagination=PaginationInfo(limit=limit, has_more=next_cursor is not None, next_cursor=next_cursor), partial_message_sequences=list(page.partial_message_sequences), ) - - async def get_legacy_projections_async( - self, - *, - attack_result_id: str, - include_messages: bool = False, - limit_messages_per_conversation: int = 100, - max_conversations: int = 50, - max_pieces_per_message: int = 100, - ) -> ScenarioAttackResultConversations: - """ - Adapt canonical saved records to the existing clipped hierarchy wire format. - - New clients load a selected conversation through ``get_messages_async`` instead. - - Returns: - ScenarioAttackResultConversations: Compatible bounded metadata and optional turn previews. - """ - page = await self._memory.get_attack_conversation_page_async( - attack_result_id=attack_result_id, limit=max_conversations - ) - if page is None: - return ScenarioAttackResultConversations(conversations=[], total_count=0, truncated=False) - stats = ( - await self._memory.get_conversation_stats_async( - conversation_ids=[reference.conversation_id for reference in page.conversations] - ) - if page.conversations - else {} - ) - conversations: list[ScenarioConversationProjection] = [] - for reference in page.conversations: - messages: list[ScenarioConversationMessageProjection] = [] - message_count = stats.get(reference.conversation_id, ConversationStats(message_count=0)).message_count - truncated = message_count > 0 - if include_messages: - ( - saved_messages, - truncated, - partial_sequences, - ) = await self._memory.get_bounded_conversation_messages_async( - conversation_id=reference.conversation_id, - limit=limit_messages_per_conversation, - max_pieces_per_message=max_pieces_per_message, - ) - for message in saved_messages: - first = message.message_pieces[0] - messages.append( - ScenarioConversationMessageProjection( - role=first.role, - content=message.get_value(), - sequence=first.sequence, - timestamp=first.timestamp, - pieces=[ - ScenarioConversationMessagePieceProjection( - data_type=piece.converted_value_data_type, - content=piece.original_value, - converted_value=piece.converted_value, - ) - for piece in message.message_pieces - ], - pieces_truncated=first.sequence in partial_sequences, - ) - ) - kind, description = self._LEGACY_DETAILS[reference.conversation_type] - conversations.append( - ScenarioConversationProjection( - conversation_id=reference.conversation_id, - conversation_type=kind, - description=description, - message_count=message_count, - messages=messages, - messages_truncated=truncated, - ) - ) - return ScenarioAttackResultConversations( - conversations=conversations, - total_count=page.total_count, - truncated=page.total_count > len(conversations), - ) diff --git a/pyrit/backend/services/scenario_progress_read_model.py b/pyrit/backend/services/scenario_progress_read_model.py index 742d071275..8c072e2c07 100644 --- a/pyrit/backend/services/scenario_progress_read_model.py +++ b/pyrit/backend/services/scenario_progress_read_model.py @@ -30,12 +30,11 @@ AttackResultRole, AttackTechniqueIdentifier, ComponentIdentifier, + ConversationSummary, ScenarioAtomicGroupProgress, - ScenarioAttackResultConversations, ScenarioAttackResultDelta, ScenarioAttackTechniqueDetails, ScenarioComponentIdentity, - ScenarioConversationProjection, ScenarioDisplayGroupProgress, ScenarioExecutionUnit, ScenarioHierarchyDiagnostics, @@ -928,7 +927,6 @@ def build_scenario_hierarchy( status: ScenarioRunState, snapshot: ScenarioProgressSnapshot, max_nodes: int = _MAX_HIERARCHY_NODES, - conversation_stats: Mapping[str, Any] | None = None, overall_counts: ScenarioProgressCounts | None = None, exact_total_errors: int | None = None, exact_total_retries: int | None = None, @@ -946,7 +944,6 @@ 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. 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. @@ -1030,57 +1027,11 @@ def _build_node( else: unresolved_children.append(c_id) - conversations: list[ScenarioConversationProjection] = [] - if result.conversation_id: - msg_count: int | None = None - messages_truncated: bool = True - if conversation_stats is not None and result.conversation_id in conversation_stats: - stat = conversation_stats[result.conversation_id] - c = getattr(stat, "message_count", None) - if c is None and isinstance(stat, int): - c = stat - if isinstance(c, int): - msg_count = c - messages_truncated = c > 0 - - conversations.append( - ScenarioConversationProjection( - conversation_id=result.conversation_id, - conversation_type=str(delta.attack_metadata.get("conversation_type") or "main"), - description=str(delta.attack_metadata.get("conversation_description") or "") or None, - message_count=msg_count, - messages=[], - messages_truncated=messages_truncated, - ) - ) - - candidate_ids = ( - result.related_conversation_ids - or getattr(result, "related_conversations", None) - or ScenarioProgressReadModel._read_related_conversation_ids(delta=delta) - ) - has_other_conversations = False - if candidate_ids: - for c in candidate_ids: - cid = getattr(c, "conversation_id", c) - if isinstance(cid, str) and cid and cid != result.conversation_id: - has_other_conversations = True - break - if has_other_conversations: - node_conversations_count = None - node_conversations_truncated = True - else: - node_conversations_count = len(conversations) - node_conversations_truncated = False - node.children = children node.unresolved_child_attack_result_ids = unresolved_children[:_MAX_CHILD_REFERENCES] node.unresolved_child_references_truncated = ( len(unresolved_children) > _MAX_CHILD_REFERENCES or node.child_references_truncated ) - node.conversations = conversations - node.conversations_count = node_conversations_count - node.conversations_truncated = node_conversations_truncated return node roots_by_group: dict[str, list[ScenarioHierarchyNode]] = {} @@ -1301,7 +1252,7 @@ def build_hierarchy_node( result: ScenarioProgressResult, children_deltas: Sequence[ScenarioAttackResultDelta] = (), children_results: Sequence[ScenarioProgressResult] = (), - conversations: Sequence[ScenarioConversationProjection] | ScenarioAttackResultConversations = (), + conversations: Sequence[ConversationSummary] = (), conversations_count: int | None = None, conversations_truncated: bool = False, children_truncated: bool = False, @@ -1311,14 +1262,14 @@ def build_hierarchy_node( next_children_cursor: str | None = None, ) -> ScenarioHierarchyNode: """ - Assemble a single hierarchy node with its immediate children and generic conversations. + Assemble a single hierarchy node with its immediate children and shared conversation metadata. Args: delta: The attack result delta for the target node. result: The mapped progress result for the target node. children_deltas: Deltas for immediate children resolved in a single batch. children_results: Mapped results for immediate children. - conversations: Generic read-only conversation projections. + conversations: Canonical read-only conversation summaries. 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. @@ -1333,13 +1284,7 @@ def build_hierarchy_node( node = ScenarioProgressReadModel._map_hierarchy_node( delta=delta, result=result, max_child_references=max_child_references ) - if isinstance(conversations, ScenarioAttackResultConversations): - node.conversations = list(conversations.conversations) - if conversations_count is None: - conversations_count = conversations.total_count - conversations_truncated = conversations_truncated or conversations.truncated - else: - node.conversations = list(conversations) + node.conversations = list(conversations) node.conversations_count = conversations_count if conversations_count is not None else len(node.conversations) node.conversations_truncated = conversations_truncated children_deltas_by_id = {d.attack_result_id: d for d in children_deltas} diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index 3eee234e74..1f0b742453 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -2681,7 +2681,6 @@ async def get_run_hierarchy_async( status=header_result.scenario_run_state, snapshot=progress_snapshot, max_nodes=max_nodes, - conversation_stats=None, overall_counts=self._hierarchy_counts(aggregate=aggregate, planned=summary.overall.planned), exact_total_errors=aggregate.raw_total_errors, exact_total_retries=aggregate.raw_total_retries, @@ -2704,58 +2703,6 @@ async def get_run_hierarchy_async( hierarchy.diagnostics.rootless_results_count = raw_attempts hierarchy.diagnostics.rootless_results_truncated = clipped - # 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( @@ -2765,7 +2712,6 @@ async def get_run_hierarchy_node_async( attack_result_id: str, max_immediate_children: int = 200, children_cursor: str | None = None, - include_messages: bool = False, ) -> ScenarioHierarchyNode | None: """ Get bounded single-node hierarchy details and generic conversation projections without full-run hydration. @@ -2775,7 +2721,6 @@ async def get_run_hierarchy_node_async( attack_result_id: The target attack result ID. max_immediate_children: Maximum number of immediate child nodes to resolve. children_cursor: Continuation bound to this run, node, and stored child ordering. - include_messages: Explicitly include compatible clipped conversation previews. Returns: ScenarioHierarchyNode | None: The node details, or None if the node is not found in this run. @@ -2853,17 +2798,22 @@ async def get_run_hierarchy_node_async( self._progress_read_model._map_progress_delta(delta=d, plan_lookup=plan_lookup) for d in children_deltas ] - conversations = await self._saved_conversation_reader.get_legacy_projections_async( - attack_result_id=attack_result_id, - include_messages=include_messages, + conversations = await self._saved_conversation_reader.list_conversations_async( + attack_result_id=attack_result_id ) + if conversations is None: + return None return self._progress_read_model.build_hierarchy_node( delta=target_delta, result=target_result, children_deltas=children_deltas, children_results=children_results, - conversations=conversations, + conversations=conversations.conversations, + conversations_count=conversations.total_count, + conversations_truncated=( + conversations.pagination.has_more if conversations.pagination is not None else False + ), children_truncated=child_offset > 0 or next_offset < len(child_ids), child_offset=child_offset, child_limit=max_immediate_children, diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index d86627f06b..7f17b12784 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -8237,80 +8237,6 @@ def _execute_get_attack_conversation_page( objective_score_id=str(row["objective_score_id"]) if row["objective_score_id"] is not None else None, ) - def _fetch_bounded_conversation_messages( - self, - *, - conversation_id: str, - limit: int, - max_pieces_per_message: int = 100, - ) -> tuple[list[Message], bool, set[int]]: - """ - Fetch conversation messages applying SQL LIMIT directly at the database layer. - - Preserves all multimodal message pieces and raw stored roles. - - Returns: - tuple[list[Message], bool, set[int]]: Canonical messages, more turns, and clipped turn sequences. - """ - if not conversation_id or limit < 1: - return [], False, set() - - with closing(self._get_session()) as session: - seq_stmt = ( - select(PromptMemoryEntry.sequence) - .where(PromptMemoryEntry.conversation_id == str(conversation_id)) - .distinct() - .order_by(PromptMemoryEntry.sequence.asc()) - .limit(limit + 1) - ) - seqs = [row[0] for row in session.execute(seq_stmt).all()] - if not seqs: - return [], False, set() - - truncated = len(seqs) > limit - target_seqs = seqs[:limit] - - piece_budget = max(1, max_pieces_per_message) - subq = ( - select( - PromptMemoryEntry.id.label("entry_id"), - func.row_number() - .over( - partition_by=PromptMemoryEntry.sequence, - order_by=( - PromptMemoryEntry.timestamp.asc(), - PromptMemoryEntry.id.asc(), - ), - ) - .label("rn"), - ) - .where( - and_( - PromptMemoryEntry.conversation_id == str(conversation_id), - PromptMemoryEntry.sequence.in_(target_seqs), - ) - ) - .subquery() - ) - 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(), - ) - ) - rows = session.execute(pieces_stmt).all() - if not rows: - return [], False, set() - - truncated_seqs: set[int] = {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] - messages = list(group_conversation_message_pieces_by_sequence(message_pieces=bounded_pieces)) - return messages, truncated, truncated_seqs - def _execute_get_scenario_results( self, *, @@ -11925,27 +11851,6 @@ async def get_conversation_message_page_async( limit=limit, ) - async def get_bounded_conversation_messages_async( - self, *, conversation_id: str, limit: int = 100, max_pieces_per_message: int = 100 - ) -> tuple[list[Message], bool, set[int]]: - """ - Read a bounded turn preview for compatibility with existing clipped detail views. - - Returns: - tuple[list[Message], bool, set[int]]: Canonical messages, more turns, and clipped turn sequences. - - Raises: - ValueError: If the turn or piece preview budget is invalid. - """ - if not 1 <= limit <= 100 or not 1 <= max_pieces_per_message <= 500: - raise ValueError("Conversation previews allow 1 through 100 turns and 1 through 500 pieces per turn.") - return await self._run_database_operation_async( - self._fetch_bounded_conversation_messages, - conversation_id=conversation_id, - limit=limit, - max_pieces_per_message=max_pieces_per_message, - ) - def get_scenario_results( self, *, diff --git a/pyrit/models/__init__.py b/pyrit/models/__init__.py index b65d45c0c9..f18db5ee5b 100644 --- a/pyrit/models/__init__.py +++ b/pyrit/models/__init__.py @@ -66,6 +66,7 @@ StoredPreset, ) from pyrit.models.conversation_stats import ConversationStats + from pyrit.models.conversation_summary import ConversationSummary from pyrit.models.embeddings import EmbeddingData, EmbeddingResponse, EmbeddingSupport, EmbeddingUsageInformation from pyrit.models.finding import Finding, FindingCreate, FindingEvidence, FindingSeverity from pyrit.models.harm_definition import HarmDefinition, ScaleDescription, get_all_harm_definitions @@ -163,10 +164,6 @@ scenario_dataset_size_from_limit, ) from pyrit.models.scenario_hierarchy import ( - ScenarioAttackResultConversations, - ScenarioConversationMessagePieceProjection, - ScenarioConversationMessageProjection, - ScenarioConversationProjection, ScenarioHierarchyDiagnostics, ScenarioHierarchyGroup, ScenarioHierarchyNode, @@ -367,6 +364,7 @@ "ConversationRetry": "pyrit.models.messages.conversation_retry", "ConversationRetryReason": "pyrit.models.messages.conversation_retry", "ConversationStats": "pyrit.models.conversation_stats", + "ConversationSummary": "pyrit.models.conversation_summary", "Finding": "pyrit.models.finding", "FindingCreate": "pyrit.models.finding", "FindingEvidence": "pyrit.models.finding", @@ -481,10 +479,6 @@ "ScenarioDisplayGroupProgress": "pyrit.models.scenario_progress", "ScenarioExecutionStatistics": "pyrit.models.scenario_progress", "ScenarioExecutionUnit": "pyrit.models.scenario_progress", - "ScenarioConversationMessagePieceProjection": "pyrit.models.scenario_hierarchy", - "ScenarioAttackResultConversations": "pyrit.models.scenario_hierarchy", - "ScenarioConversationMessageProjection": "pyrit.models.scenario_hierarchy", - "ScenarioConversationProjection": "pyrit.models.scenario_hierarchy", "ScenarioHierarchyDiagnostics": "pyrit.models.scenario_hierarchy", "ScenarioHierarchyGroup": "pyrit.models.scenario_hierarchy", "ScenarioHierarchyNode": "pyrit.models.scenario_hierarchy", diff --git a/pyrit/models/conversation_summary.py b/pyrit/models/conversation_summary.py new file mode 100644 index 0000000000..7019265be0 --- /dev/null +++ b/pyrit/models/conversation_summary.py @@ -0,0 +1,16 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from datetime import datetime + +from pydantic import BaseModel, Field + + +class ConversationSummary(BaseModel): + """Read-only conversation metadata shared by attack and hierarchy views.""" + + conversation_id: str = Field(..., description="Unique conversation identifier") + message_count: int = Field(0, description="Number of messages in this conversation") + last_message_preview: str | None = Field(None, description="Preview of the last message") + created_at: datetime | None = Field(None, description="Timestamp of the first message") + conversation_type: str | None = None diff --git a/pyrit/models/scenario_hierarchy.py b/pyrit/models/scenario_hierarchy.py index 8be73713e9..bfde364933 100644 --- a/pyrit/models/scenario_hierarchy.py +++ b/pyrit/models/scenario_hierarchy.py @@ -1,11 +1,12 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -from typing import ClassVar, NamedTuple, Self +from typing import ClassVar, Self from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, model_validator from pyrit.models.catalog.scenario import AttackErrorSummary, AttackRetrySummary +from pyrit.models.conversation_summary import ConversationSummary from pyrit.models.results.attack_result import AttackOutcome, AttackResultRole from pyrit.models.results.scenario_result import ScenarioRunState from pyrit.models.retry_event import RetryEvent @@ -53,53 +54,6 @@ def _bound_child_references(self) -> Self: return self -class ScenarioConversationMessagePieceProjection(BaseModel): - """ - Projection of an individual message piece preserving data type and multimodal content. - """ - - model_config = ConfigDict(extra="ignore") - - data_type: str = "text" - content: str = "" - converted_value: str = "" - - -class ScenarioConversationMessageProjection(BaseModel): - """ - Generic, read-only message projection for conversation views. - - Contains display text, conversation role, sequence order, multimodal pieces, and linked score indicators. - """ - - model_config = ConfigDict(extra="ignore") - - role: str - content: str - sequence: int - timestamp: AwareDatetime | None = None - scores: list[ScenarioProgressScore] = Field(default_factory=list) - pieces: list[ScenarioConversationMessagePieceProjection] = Field(default_factory=list) - pieces_truncated: bool = False - - -class ScenarioConversationProjection(BaseModel): - """ - Generic read-only conversation projection belonging to an attack attempt. - - Carries conversation metadata and optional message lists for detailed views. - """ - - model_config = ConfigDict(extra="ignore") - - conversation_id: str - conversation_type: str = "main" - description: str | None = None - message_count: int | None = None - messages: list[ScenarioConversationMessageProjection] = Field(default_factory=list) - messages_truncated: bool = False - - class ScenarioHierarchyNode(BaseModel): """ One node in the scenario result hierarchy. @@ -149,8 +103,8 @@ class ScenarioHierarchyNode(BaseModel): 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) + # Canonical conversation metadata, loaded on node expansion + conversations: list[ConversationSummary] = Field(default_factory=list) conversations_count: int | None = None conversations_truncated: bool = False @@ -170,14 +124,6 @@ def _bound_child_references(self) -> Self: return self -class ScenarioAttackResultConversations(NamedTuple): - """Container for bounded conversation projections, total count, and list truncation flag.""" - - conversations: list[ScenarioConversationProjection] - total_count: int - truncated: bool - - class ScenarioHierarchyGroup(BaseModel): """One atomic attack group's hierarchy projection.""" diff --git a/tests/unit/backend/test_scenario_hierarchy_pagination.py b/tests/unit/backend/test_scenario_hierarchy_pagination.py index dda40cc68e..43062a8f64 100644 --- a/tests/unit/backend/test_scenario_hierarchy_pagination.py +++ b/tests/unit/backend/test_scenario_hierarchy_pagination.py @@ -362,10 +362,11 @@ async def test_api_journey_loads_only_the_selected_saved_conversation_async( assert roots.json()["groups"][0]["roots"][0]["attack_result_id"] == parent.attack_result_id node = client.get(f"/api/scenarios/runs/{run_id}/hierarchy/nodes/{child.attack_result_id}") assert node.status_code == 200 - assert node.json()["conversations"][0]["messages"] == [] + assert "messages" not in node.json()["conversations"][0] references = client.get(f"/api/attacks/{child.attack_result_id}/conversations", params={"limit": 1}) assert references.status_code == 200 assert references.json()["conversations"][0]["message_count"] == 1 + assert node.json()["conversations"] == references.json()["conversations"] messages_url = f"/api/attacks/{child.attack_result_id}/messages" first = client.get(messages_url, params={"conversation_id": child.conversation_id, "limit": 2}) assert first.status_code == 200 diff --git a/tests/unit/backend/test_scenario_hierarchy_parity.py b/tests/unit/backend/test_scenario_hierarchy_parity.py index c355bdd766..1f56fa9ca4 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 asyncio import hashlib import uuid from datetime import UTC, datetime, timedelta @@ -16,6 +17,8 @@ from pyrit.analytics import compute_outcome_statistics from pyrit.backend.main import app +from pyrit.backend.models.attacks import AttackConversationsResponse +from pyrit.backend.models.common import PaginationInfo from pyrit.backend.services.saved_conversation_read_service import SavedConversationReadService from pyrit.backend.services.scenario_progress_read_model import ( ScenarioProgressReadModel, @@ -33,15 +36,13 @@ AttackSeedGroup, ComponentIdentifier, ConversationReference, + ConversationSummary, ConversationType, MessagePiece, RetryEvent, ScenarioAtomicGroupProgress, - ScenarioAttackResultConversations, ScenarioAttackResultDelta, ScenarioAttackTechniqueDetails, - ScenarioConversationMessageProjection, - ScenarioConversationProjection, ScenarioHierarchyGroup, ScenarioHierarchyNode, ScenarioHierarchyProjection, @@ -495,24 +496,19 @@ 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) conversation_reader = MagicMock(spec=SavedConversationReadService) - conversation_reader.get_legacy_projections_async = AsyncMock( - return_value=ScenarioAttackResultConversations( + conversation_reader.list_conversations_async = AsyncMock( + return_value=AttackConversationsResponse( + attack_result_id="parent-1", + main_conversation_id="conv_1", conversations=[ - ScenarioConversationProjection( + ConversationSummary( conversation_id="conv_1", conversation_type="main", message_count=1, - messages=[ - ScenarioConversationMessageProjection( - role="user", - content="Test prompt", - sequence=1, - ) - ], ) ], total_count=1, - truncated=False, + pagination=PaginationInfo(limit=50, has_more=False), ) ) @@ -530,9 +526,7 @@ async def mock_get_deltas( assert len(node.children) == 1 assert node.children[0].attack_result_id == "child-1" assert len(node.conversations) == 1 - conversation_reader.get_legacy_projections_async.assert_awaited_once_with( - attack_result_id="parent-1", include_messages=False - ) + conversation_reader.list_conversations_async.assert_awaited_once_with(attack_result_id="parent-1") # Querying non-existent node returns None missing = await service.get_run_hierarchy_node_async( @@ -540,6 +534,11 @@ async def mock_get_deltas( attack_result_id="non-existent", ) assert missing is None + conversation_reader.list_conversations_async.return_value = None + assert ( + await service.get_run_hierarchy_node_async(scenario_result_id="scenario-run-1", attack_result_id="parent-1") + is None + ) def test_rest_routes_hierarchy_and_node(client: TestClient) -> None: @@ -706,22 +705,20 @@ async def test_sqlite_memory_batched_reads_and_projections_parity(sqlite_instanc # 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 SavedConversationReadService(memory=sqlite_instance).get_legacy_projections_async( - attack_result_id=child_id, - include_messages=True, - ) - assert conv_projections.total_count == 1 - assert conv_projections.truncated is False - assert len(conv_projections.conversations) == 1 - proj = conv_projections.conversations[0] - assert proj.conversation_id == child_conv_id - assert proj.message_count == 2 - assert len(proj.messages) == 2 - assert proj.messages[0].role == "user" - assert proj.messages[0].content == "Prompt to target" - assert proj.messages[1].role == "assistant" - assert proj.messages[1].content == "Response from target" + reader = SavedConversationReadService(memory=sqlite_instance) + references = await reader.list_conversations_async(attack_result_id=child_id) + assert references is not None + assert references.total_count == 1 + assert len(references.conversations) == 1 + assert references.conversations[0].conversation_id == child_conv_id + assert references.conversations[0].message_count == 2 + transcript = await reader.get_messages_async(attack_result_id=child_id, conversation_id=child_conv_id) + assert transcript is not None + assert len(transcript.messages) == 2 + assert transcript.messages[0].role == "user" + assert transcript.messages[0].message_pieces[0].original_value == "Prompt to target" + assert transcript.messages[1].role == "assistant" + assert transcript.messages[1].message_pieces[0].original_value == "Response from target" # 3. Test service single-node read service = ScenarioRunService(memory=sqlite_instance) @@ -1084,7 +1081,7 @@ async def test_sqlite_bounded_multimodal_conversation_projection(sqlite_instance # Create real image on disk img_file = tmp_path / "diagram.png" - img_file.write_bytes(b"dummy_png_bytes") + await asyncio.to_thread(img_file.write_bytes, b"dummy_png_bytes") img_path = str(img_file) # Sequence 0: 2 pieces (text + image) @@ -1131,36 +1128,33 @@ async def test_sqlite_bounded_multimodal_conversation_projection(sqlite_instance 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 SavedConversationReadService(memory=sqlite_instance).get_legacy_projections_async( - attack_result_id=ar_id, - include_messages=True, - limit_messages_per_conversation=2, - ) - - assert projs.total_count == 1 - assert projs.truncated is False - assert len(projs.conversations) == 1 - proj = projs.conversations[0] - assert proj.conversation_id == conv_id - assert proj.message_count == 3 # Exact count from COUNT query - assert proj.messages_truncated is True - assert len(proj.messages) == 2 # Bounded by SQL LIMIT - - # Verify multimodal pieces in message 0 - msg0 = proj.messages[0] + reader = SavedConversationReadService(memory=sqlite_instance) + references = await reader.list_conversations_async(attack_result_id=ar_id) + assert references is not None + assert references.total_count == 1 + assert references.conversations[0].message_count == 3 + transcript = await reader.get_messages_async(attack_result_id=ar_id, conversation_id=conv_id, limit=3) + assert transcript is not None + assert transcript.pagination is not None + assert transcript.pagination.has_more + assert transcript.partial_message_sequences == [] + assert len(transcript.messages) == 2 + msg0 = transcript.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 len(msg0.message_pieces) == 2 + assert msg0.message_pieces[0].original_value_data_type == "text" + assert msg0.message_pieces[0].original_value == "Analyze blueprint image" + assert msg0.message_pieces[1].original_value_data_type == "image_path" + assert msg0.message_pieces[1].original_value == img_path + assert msg0.message_pieces[1].converted_value == img_path + msg1 = transcript.messages[1] assert msg1.role == "simulated_assistant" - assert msg1.content == "I detect architectural blueprints." + assert msg1.message_pieces[0].original_value == "I detect architectural blueprints." + remaining = await reader.get_messages_async( + attack_result_id=ar_id, conversation_id=conv_id, limit=3, cursor=transcript.pagination.next_cursor + ) + assert remaining is not None + assert [message.sequence for message in remaining.messages] == [2] def test_hierarchy_repeated_attempt_diagnostic_retries_consistency() -> None: @@ -1270,14 +1264,24 @@ async def test_sqlite_conversation_list_clipping_explicit(sqlite_instance: SQLit ) 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 SavedConversationReadService(memory=sqlite_instance).get_legacy_projections_async( - attack_result_id=ar_id, - max_conversations=50, - ) + reader = SavedConversationReadService(memory=sqlite_instance) + conv_result = await reader.list_conversations_async(attack_result_id=ar_id, limit=50) + assert conv_result is not None assert conv_result.total_count == 55 - assert conv_result.truncated is True + assert conv_result.pagination is not None + assert conv_result.pagination.has_more assert len(conv_result.conversations) == 50 + remaining = await reader.list_conversations_async( + attack_result_id=ar_id, limit=50, cursor=conv_result.pagination.next_cursor + ) + assert remaining is not None + assert remaining.pagination is not None + assert not remaining.pagination.has_more + assert len(remaining.conversations) == 5 + assert {conversation.conversation_id for conversation in conv_result.conversations + remaining.conversations} == { + main_conv_id, + *prep_conv_ids, + } # 2. Service single-node read projects exact count and truncation flag on node service = ScenarioRunService(memory=sqlite_instance) @@ -1332,8 +1336,8 @@ async def test_executed_turns_zero_vs_invalid_metadata(sqlite_instance: SQLiteMe 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.""" +@pytest.mark.parametrize("count", [2, 3, 4]) +def test_node_preserves_canonical_conversation_summary_sequences(count: int) -> None: delta = _make_test_delta( attack_result_id="test-node", outcome=AttackOutcome.SUCCESS, @@ -1345,47 +1349,19 @@ def test_eliminate_tuple_duck_typing_raw_tuples() -> None: 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( + conversations = tuple(ConversationSummary(conversation_id=f"conv-{index}") for index in range(count)) + node = ScenarioProgressReadModel.build_hierarchy_node( delta=delta, result=result, - conversations=(conv1, conv2, conv3, conv4), + conversations=conversations, ) - assert len(node4.conversations) == 4 - assert node4.conversations == [conv1, conv2, conv3, conv4] - assert node4.conversations_count == 4 - assert node4.conversations_truncated is False + assert node.conversations == list(conversations) + assert node.conversations_count == count + assert node.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.""" + """Continue through a large turn without permanently clipping any saved pieces.""" scenario_uuid = uuid.uuid4() scenario_id = str(scenario_uuid) scen = make_scenario_result( @@ -1429,28 +1405,26 @@ async def test_sqlite_message_pieces_bounding(sqlite_instance: SQLiteMemory) -> sqlite_instance._insert_entry, PromptMemoryEntry(entry=piece) ) - # 1. Fetch with default max_pieces_per_message=100 - projs = await SavedConversationReadService(memory=sqlite_instance).get_legacy_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 SavedConversationReadService(memory=sqlite_instance).get_legacy_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 + reader = SavedConversationReadService(memory=sqlite_instance) + first = await reader.get_messages_async(attack_result_id=ar_id, conversation_id=conv_id, limit=100) + assert first is not None + assert len(first.messages) == 1 + assert len(first.messages[0].message_pieces) == 100 + assert first.partial_message_sequences == [0] + assert first.pagination is not None + second = await reader.get_messages_async( + attack_result_id=ar_id, conversation_id=conv_id, limit=100, cursor=first.pagination.next_cursor + ) + assert second is not None + assert len(second.messages[0].message_pieces) == 50 + assert second.pagination is not None + assert not second.pagination.has_more + complete = await reader.get_messages_async(attack_result_id=ar_id, conversation_id=conv_id, limit=200) + assert complete is not None + assert complete.partial_message_sequences == [] + assert [piece.id for page in (first, second) for piece in page.messages[0].message_pieces] == [ + piece.id for piece in complete.messages[0].message_pieces + ] async def test_unloaded_conversations_in_full_hierarchy(sqlite_instance: SQLiteMemory) -> None: @@ -1493,7 +1467,8 @@ async def test_unloaded_conversations_in_full_hierarchy(sqlite_instance: SQLiteM 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 + assert root_node.conversations == [] + assert root_node.conversation_id == main_conv_id def test_diagnostic_total_retries_scope_unmatched_attempt() -> None: @@ -1692,8 +1667,9 @@ async def test_persisted_result_with_malformed_related_conversations_metadata( 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 + assert root_node.conversations == [] + assert root_node.conversations_count is None + assert root_node.conversations_truncated is True def test_read_related_conversation_ids_shapes() -> None: @@ -1764,10 +1740,13 @@ def test_build_hierarchy_node_deduplicates_duplicate_child_ids() -> None: assert node.unresolved_child_attack_result_ids == ["child-1"] -async def test_persisted_conversation_statistics_parity_between_hierarchy_and_node( +@pytest.mark.parametrize("saved_message_count", [0, 3]) +@pytest.mark.parametrize("metadata_count", [None, 0, 99, "unknown"]) +async def test_selected_node_uses_canonical_saved_stats_and_tree_defers_conversations_async( sqlite_instance: SQLiteMemory, + saved_message_count: int, + metadata_count: int | str | None, ) -> 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( @@ -1789,12 +1768,11 @@ async def test_persisted_conversation_statistics_parity_between_hierarchy_and_no outcome=AttackOutcome.SUCCESS, attribution_parent_id=scenario_id, attribution_data={"parent_collection": "group_1", "result_role": AttackResultRole.TARGET_FACING.value}, - metadata={}, # No count metadata! + metadata={} if metadata_count is None else {"conversation_message_count": metadata_count}, ) await sqlite_instance.add_attack_results_to_memory_async(attack_results=[ar]) - # Insert 3 saved messages into SQLite - for seq in range(3): + for seq in range(saved_message_count): piece = MessagePiece( id=uuid.uuid4(), conversation_id=conv_id, @@ -1811,88 +1789,30 @@ async def test_persisted_conversation_statistics_parity_between_hierarchy_and_no service = ScenarioRunService(memory=sqlite_instance) - node = await service.get_run_hierarchy_node_async( - scenario_result_id=scenario_id, - attack_result_id=ar_id, - ) + with patch.object( + sqlite_instance, "get_conversation_stats_async", wraps=sqlite_instance.get_conversation_stats_async + ) as stats: + hierarchy = await service.get_run_hierarchy_async(scenario_result_id=scenario_id) + stats.assert_not_awaited() + node = await service.get_run_hierarchy_node_async(scenario_result_id=scenario_id, attack_result_id=ar_id) + stats.assert_awaited_once_with(conversation_ids=[conv_id]) assert node is not None assert len(node.conversations) == 1 - assert node.conversations[0].message_count == 3 - assert node.conversations[0].messages == [] - assert node.conversations[0].messages_truncated is True - explicit_preview = await service.get_run_hierarchy_node_async( - scenario_result_id=scenario_id, attack_result_id=ar_id, include_messages=True - ) - assert explicit_preview is not None - assert len(explicit_preview.conversations[0].messages) == 3 - assert explicit_preview.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 node.conversations[0].message_count == saved_message_count + assert "messages" not in node.conversations[0].model_dump() 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 + assert root_node.conversations == [] + assert root_node.conversation_id == conv_id + assert root_node.conversations_count is None + assert root_node.conversations_truncated + reader = SavedConversationReadService(memory=sqlite_instance) + references = await reader.list_conversations_async(attack_result_id=ar_id) + assert references is not None + assert references.conversations == node.conversations + transcript = await reader.get_messages_async(attack_result_id=ar_id, conversation_id=conv_id, limit=2) + assert transcript is not None + assert len(transcript.messages) == min(saved_message_count, 2) async def test_bounded_conversation_messages_piece_budget_applied_at_sql_layer( @@ -1939,35 +1859,30 @@ async def test_bounded_conversation_messages_piece_budget_applied_at_sql_layer( ) await sqlite_instance.add_attack_results_to_memory_async(attack_results=[ar]) - projs = await SavedConversationReadService(memory=sqlite_instance).get_legacy_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 - - messages_capped_10, _, partial_sequences = await sqlite_instance.get_bounded_conversation_messages_async( - conversation_id=conv_id, - limit=10, - max_pieces_per_message=10, - ) - assert len(messages_capped_10[0].message_pieces) == 10 - assert len(messages_capped_10[1].message_pieces) == 5 - assert partial_sequences == {0} + reader = SavedConversationReadService(memory=sqlite_instance) + with patch.object( + PromptMemoryEntry, "get_message_piece", autospec=True, side_effect=PromptMemoryEntry.get_message_piece + ) as hydrate: + page = await reader.get_messages_async(attack_result_id=ar_id, conversation_id=conv_id, limit=100) + # Canonical score lookup reads the same saved pieces once more. + assert hydrate.call_count <= 200 + assert {str(call.args[0].id) for call in hydrate.call_args_list} == {str(piece.id) for piece in pieces_seq0[:100]} + assert page is not None + assert page.partial_message_sequences == [0] + returned: list[MessagePiece] = [] + while True: + returned.extend(piece for message in page.messages for piece in message.message_pieces) + assert sum(len(message.message_pieces) for message in page.messages) <= 100 + assert page.pagination is not None + if not page.pagination.has_more: + break + page = await reader.get_messages_async( + attack_result_id=ar_id, conversation_id=conv_id, limit=100, cursor=page.pagination.next_cursor + ) + assert page is not None + assert [piece.id for piece in returned] == [piece.id for piece in pieces_seq0 + pieces_seq1] + assert [message.sequence for message in page.messages] == [0, 1] + assert len(page.messages[-1].message_pieces) == 5 async def test_diagnostics_total_retries_count_event_mismatch_parity( @@ -2030,8 +1945,7 @@ async def test_diagnostics_total_retries_count_event_mismatch_parity( @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.""" +async def test_clipped_tree_defers_stats_until_the_selected_node_is_opened_async(sqlite_instance: SQLiteMemory) -> None: scenario_uuid = uuid.uuid4() scenario_id = str(scenario_uuid) scen = make_scenario_result( @@ -2088,22 +2002,30 @@ async def test_clipped_tree_queries_actual_preview_node_conversation_stats(sqlit 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) + with patch.object( + sqlite_instance, "get_conversation_stats_async", wraps=sqlite_instance.get_conversation_stats_async + ) as stats: + hierarchy = await service.get_run_hierarchy_async(scenario_result_id=scenario_id, max_nodes=1) + stats.assert_not_awaited() + node = await service.get_run_hierarchy_node_async(scenario_result_id=scenario_id, attack_result_id=parent_id) + stats.assert_awaited_once_with(conversation_ids=[parent_conv_id]) 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 + assert root.conversation_id == parent_conv_id + assert root.conversations == [] + assert root.conversations_count is None + assert root.conversations_truncated + assert node is not None + assert node.conversations[0].message_count == 3 -@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.""" +@pytest.mark.parametrize("metadata_count", [None, 0]) +async def test_conversation_stats_failure_is_deferred_and_never_becomes_false_zero_async( + sqlite_instance: SQLiteMemory, metadata_count: int | None +) -> None: scenario_uuid = uuid.uuid4() scenario_id = str(scenario_uuid) scen = make_scenario_result( @@ -2124,6 +2046,7 @@ async def test_conversation_stats_query_failure_does_not_become_false_zero(sqlit outcome=AttackOutcome.SUCCESS, attribution_parent_id=scenario_id, attribution_data={"parent_collection": "group_1", "result_role": AttackResultRole.TARGET_FACING.value}, + metadata={} if metadata_count is None else {"conversation_message_count": metadata_count}, ) await sqlite_instance.add_attack_results_to_memory_async(attack_results=[ar]) @@ -2134,10 +2057,12 @@ async def test_conversation_stats_query_failure_does_not_become_false_zero(sqlit 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 + assert root.conversations == [] + assert root.conversations_count is None + assert root.conversations_truncated + sqlite_instance.get_conversation_stats_async.assert_not_awaited() + with pytest.raises(RuntimeError, match="Database error"): + await service.get_run_hierarchy_node_async(scenario_result_id=scenario_id, attack_result_id=ar_id) @pytest.mark.asyncio @@ -2266,7 +2191,7 @@ async def test_cold_hierarchy_hydration_bounded_by_max_nodes(sqlite_instance: SQ 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.""" + """Unqueried conversation metadata stays unknown rather than reporting a fabricated zero.""" delta = _make_test_delta( attack_result_id="root-1", role=AttackResultRole.TARGET_FACING, @@ -2275,19 +2200,18 @@ def test_build_scenario_hierarchy_unqueried_conversation_not_false_zero() -> Non 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 + root = projection.groups[0].roots[0] + assert root.conversation_id == "conv-unqueried" + assert root.conversations == [] + assert root.conversations_count is None + assert root.conversations_truncated @pytest.mark.asyncio @@ -2454,65 +2378,6 @@ async def test_repeated_attempts_latest_error_beyond_cap_reported_correctly(sqli assert hierarchy.overall_counts.completed == 500 -@pytest.mark.asyncio -async def test_failed_conversation_stats_read_never_falls_back_to_metadata(sqlite_instance: SQLiteMemory) -> None: - """Test that a failed stats query leaves counts unknown/incomplete even if count metadata is present.""" - scenario_uuid = uuid.uuid4() - scenario_id = str(scenario_uuid) - scen = make_scenario_result( - id=scenario_uuid, - scenario_name="FailedStatsMetadataScenario", - scenario_version=1, - objective_target_identifier=get_mock_target_identifier(), - attack_results={}, - ) - await sqlite_instance.add_scenario_results_to_memory_async(scenario_results=[scen]) - - conv_id = str(uuid.uuid4()) - ar_id = str(uuid.uuid4()) - # 3 saved messages in SQLite - now = datetime.now(UTC) - for i in range(3): - piece = MessagePiece( - id=uuid.uuid4(), - conversation_id=conv_id, - role="user", - original_value_data_type="text", - original_value=f"msg {i}", - converted_value=f"msg {i}", - sequence=i, - timestamp=now + timedelta(seconds=i), - ) - await set_message_piece_sha256_async(piece) - await sqlite_instance._run_database_operation_async( - sqlite_instance._insert_entry, PromptMemoryEntry(entry=piece) - ) - - # Attack result with metadata={"conversation_message_count": 0} - ar = AttackResult( - attack_result_id=ar_id, - conversation_id=conv_id, - objective="Failed stats metadata fallback test", - outcome=AttackOutcome.SUCCESS, - attribution_parent_id=scenario_id, - attribution_data={"parent_collection": "group_1", "result_role": AttackResultRole.TARGET_FACING.value}, - metadata={"conversation_message_count": 0}, - ) - await sqlite_instance.add_attack_results_to_memory_async(attack_results=[ar]) - - service = ScenarioRunService(memory=sqlite_instance) - # Simulate DB query failure - sqlite_instance.get_conversation_stats_async = AsyncMock(side_effect=RuntimeError("DB query failure")) - - hierarchy = await service.get_run_hierarchy_async(scenario_result_id=scenario_id) - assert hierarchy is not None - root = hierarchy.groups[0].roots[0] - assert len(root.conversations) == 1 - # Must NOT return 0 or False from metadata! Must be None and truncated=True! - assert root.conversations[0].message_count is None - assert root.conversations[0].messages_truncated is True - - @pytest.mark.asyncio async def test_cold_hierarchy_projection_preserves_parent_beyond_delta_page_limit(sqlite_instance): """ diff --git a/tests/unit/models/test_scenario_hierarchy_models.py b/tests/unit/models/test_scenario_hierarchy_models.py index 7b3646a884..d09206608f 100644 --- a/tests/unit/models/test_scenario_hierarchy_models.py +++ b/tests/unit/models/test_scenario_hierarchy_models.py @@ -9,19 +9,15 @@ from pyrit.models import ( AttackOutcome, AttackResultRole, + ConversationSummary, RetryEvent, - ScenarioAttackResultConversations, ScenarioAttackTechniqueDetails, - ScenarioConversationMessagePieceProjection, - ScenarioConversationMessageProjection, - ScenarioConversationProjection, ScenarioHierarchyDiagnostics, ScenarioHierarchyGroup, ScenarioHierarchyNode, ScenarioHierarchyProjection, ScenarioOrchestrationParentSummary, ScenarioProgressCounts, - ScenarioProgressScore, ScenarioRunState, ) from pyrit.models.catalog.scenario import AttackErrorSummary, AttackRetrySummary @@ -76,37 +72,20 @@ def test_scenario_orchestration_parent_summary_validation() -> None: assert summary.retries_truncated is True -def test_scenario_conversation_projections() -> None: +def test_shared_conversation_summary_is_metadata_only() -> 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( + conv = ConversationSummary( conversation_id="conv-1", conversation_type="main", - description="Main objective target conversation", message_count=1, - messages=[msg], - messages_truncated=False, + last_message_preview="Response text", + created_at=now, ) assert conv.conversation_id == "conv-1" assert conv.conversation_type == "main" - assert len(conv.messages) == 1 + assert conv.message_count == 1 + assert conv.created_at == now + assert "messages" not in conv.model_dump() def test_scenario_hierarchy_node_recursive_tree() -> None: @@ -245,33 +224,6 @@ def test_scenario_hierarchy_projection_envelope() -> None: 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", @@ -312,7 +264,7 @@ def test_scenario_hierarchy_node_conversations_count_and_truncated() -> None: outcome=AttackOutcome.SUCCESS, timestamp=now, conversations=[ - ScenarioConversationProjection(conversation_id="conv-1"), + ConversationSummary(conversation_id="conv-1"), ], conversations_count=55, conversations_truncated=True, @@ -322,23 +274,6 @@ def test_scenario_hierarchy_node_conversations_count_and_truncated() -> None: assert len(node.conversations) == 1 -def test_scenario_attack_result_conversations_namedtuple() -> None: - conv = ScenarioConversationProjection(conversation_id="conv-1") - result = ScenarioAttackResultConversations( - conversations=[conv], - total_count=55, - truncated=True, - ) - assert result.conversations == [conv] - assert result.total_count == 55 - assert result.truncated is True - # Verify tuple unpacking works - c_list, count, trunc = result - assert c_list == [conv] - assert count == 55 - assert trunc is True - - def test_scenario_hierarchy_node_child_references_bounding() -> None: now = datetime.now(UTC) child_ids = [f"child-{i}" for i in range(600)] @@ -375,28 +310,6 @@ def test_scenario_orchestration_parent_summary_child_references_bounding() -> No assert summary.child_attack_result_ids == child_ids[:500] -def test_scenario_conversation_message_pieces_truncated_flag() -> None: - now = datetime.now(UTC) - pieces = [ - ScenarioConversationMessagePieceProjection( - data_type="text", - content=f"piece-{i}", - converted_value=f"piece-{i}", - ) - for i in range(100) - ] - msg = ScenarioConversationMessageProjection( - role="user", - content="test prompt", - sequence=0, - timestamp=now, - pieces=pieces, - pieces_truncated=True, - ) - assert len(msg.pieces) == 100 - assert msg.pieces_truncated is True - - def test_scenario_hierarchy_node_unresolved_child_references_bounding() -> None: now = datetime.now(UTC) unresolved_ids = [f"unresolved-{i}" for i in range(600)] From d4b33b6799c7961e1b23e031fdd3d131d65e73da Mon Sep 17 00:00:00 2001 From: Roman Lutz Date: Mon, 26 Oct 2026 21:20:38 -0700 Subject: [PATCH 22/22] fix: keep GUI conversation reads scoped to main and pruned Use the same active conversation population for hierarchy metadata, paginated lists, and saved transcripts. Exclude internal contexts before counting or paging and preserve existing unpaged behavior. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 60a3baa0-0e4d-428b-8062-5101ad7866dd --- doc/gui/0_gui.md | 2 + pyrit/backend/README.md | 9 ++- pyrit/backend/routes/attacks.py | 8 +-- pyrit/backend/services/attack_service.py | 7 +-- .../saved_conversation_read_service.py | 4 +- pyrit/memory/memory_interface.py | 19 ++---- .../test_saved_conversation_read_service.py | 27 +++++---- .../test_scenario_hierarchy_pagination.py | 59 +++++++++++++++++-- .../backend/test_scenario_hierarchy_parity.py | 19 +++--- 9 files changed, 101 insertions(+), 53 deletions(-) diff --git a/doc/gui/0_gui.md b/doc/gui/0_gui.md index 6eea876ca7..007057ed1b 100644 --- a/doc/gui/0_gui.md +++ b/doc/gui/0_gui.md @@ -651,6 +651,8 @@ paged roots and ordered child expansion, with exact run statistics independent o the displayed page. Tree previews defer conversation reads. Node expansion uses the same conversation metadata as the generic attack API; clients load transcript bodies only through the selected saved conversation's paginated message endpoint. +Only main and pruned conversations are shown; adversarial and preparation contexts +remain internal. See the [backend hierarchy API](../../pyrit/backend/README.md#scenario-result-hierarchy) for continuation and clipping semantics. diff --git a/pyrit/backend/README.md b/pyrit/backend/README.md index 83ef374dd2..3f9b0e265b 100644 --- a/pyrit/backend/README.md +++ b/pyrit/backend/README.md @@ -49,8 +49,8 @@ Saving from the scanner creates a copy through `POST /api/attacks/save-conversat it does not change the source or send to a target. For on-demand saved-result browsing, use -`GET /api/attacks/{attack_result_id}/conversations?limit=50` to page main, pruned, -adversarial, and preparation references in saved-reference order. This reads counts +`GET /api/attacks/{attack_result_id}/conversations?limit=50` to page main and pruned +references in saved-reference order. This reads counts and previews, not conversation bodies. The limit is 1-200; follow `pagination.next_cursor` while `pagination.has_more` is true. @@ -68,6 +68,9 @@ validate the saved ordering and timestamp; restart reference paging if the saved attack changes. Invalid or mismatched cursors return 400. Passing a cursor requires a limit. Omitting both preserves the existing active-conversation listing and full-transcript behavior. +Both paged and unpaged reads expose only main and pruned conversations. +Adversarial and preparation contexts are excluded from lists, reference totals, +and transcript reads, including hierarchy node expansion. ### Scenario result hierarchy @@ -100,7 +103,7 @@ restart child paging after those references change. Tree previews leave conversation metadata deferred: `conversations` is empty, `conversations_count` is unknown, and `conversations_truncated` is true when saved references exist. The node's `conversation_id` still links to its main conversation. -Node expansion returns up to 50 canonical conversation summaries, the exact reference +Node expansion returns up to 50 main/pruned conversation summaries, the exact reference count, and a clipping flag. These use the same metadata model as the generic attack API. Use the saved-conversation APIs above for reference continuation and all transcript bodies. There is no hierarchy-specific transcript format or conversation-preview mode. diff --git a/pyrit/backend/routes/attacks.py b/pyrit/backend/routes/attacks.py index 9cce3b5d1c..5d6befe252 100644 --- a/pyrit/backend/routes/attacks.py +++ b/pyrit/backend/routes/attacks.py @@ -391,7 +391,7 @@ async def get_conversation_messages( # pyrit-async-suffix-exempt cursor: CursorStr | None = Query(None, description="Opaque message-piece continuation cursor"), ) -> ConversationMessagesResponse: """ - Get saved messages for a conversation belonging to an attack. + Get saved messages for a main or pruned conversation belonging to an attack. Omitting limit preserves the existing full-conversation response. Paged reads can split a multimodal turn; append pieces by their IDs and turn number across pages. @@ -437,11 +437,11 @@ async def get_conversations( # pyrit-async-suffix-exempt cursor: CursorStr | None = Query(None, description="Opaque conversation-reference continuation cursor"), ) -> AttackConversationsResponse: """ - Get saved conversations belonging to an attack. + Get main and pruned conversations belonging to an attack. Without a limit, preserve the existing active-conversation list and timestamp - ordering. Paged reads include main, pruned, adversarial, and preparation contexts - in stored-reference order, with message counts and preview text. + ordering. Paged reads use stored-reference order, with message counts and + preview text. Both paths exclude adversarial and preparation contexts. Returns: AttackConversationsResponse: All conversations for the attack. diff --git a/pyrit/backend/services/attack_service.py b/pyrit/backend/services/attack_service.py index ecf70dc178..12aea030e2 100644 --- a/pyrit/backend/services/attack_service.py +++ b/pyrit/backend/services/attack_service.py @@ -392,7 +392,7 @@ async def get_conversation_messages_async( cursor: str | None = None, ) -> ConversationMessagesResponse | None: """ - Get all messages for a conversation belonging to an attack. + Get messages for a main or pruned conversation belonging to an attack. Args: attack_result_id: The AttackResult's primary key (used to verify existence). @@ -1097,10 +1097,9 @@ async def get_conversations_async( self, *, attack_result_id: str, limit: int | None = None, cursor: str | None = None ) -> AttackConversationsResponse | None: """ - Get all conversations belonging to an attack. + Get main and pruned conversations belonging to an attack. - Includes the main conversation and all related conversations from the - AttackResult. Each entry is enriched with message count, a preview, + Each entry is enriched with message count, a preview, and the earliest message timestamp using a single batched query. Returns: diff --git a/pyrit/backend/services/saved_conversation_read_service.py b/pyrit/backend/services/saved_conversation_read_service.py index 784ae3406f..a78bff901d 100644 --- a/pyrit/backend/services/saved_conversation_read_service.py +++ b/pyrit/backend/services/saved_conversation_read_service.py @@ -29,7 +29,7 @@ async def list_conversations_async( self, *, attack_result_id: str, limit: int = 50, cursor: str | None = None ) -> AttackConversationsResponse | None: """ - Page saved main/related conversations without loading their message bodies. + Page saved main and pruned conversations without loading their message bodies. Returns: AttackConversationsResponse | None: A reference page, or None for an absent attack. @@ -112,7 +112,7 @@ async def get_messages_async( self, *, attack_result_id: str, conversation_id: str, limit: int = 100, cursor: str | None = None ) -> ConversationMessagesResponse | None: """ - Page canonical messages for one explicitly selected saved conversation. + Page canonical messages for one selected main or pruned conversation. Returns: ConversationMessagesResponse | None: A piece-bounded page, or None for an absent attack. diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 7f17b12784..b6d7fa9b6f 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -8173,7 +8173,7 @@ def _execute_get_attack_conversation_page( self, *, attack_result_id: str, offset: int = 0, limit: int = 50, conversation_id: str | None = None ) -> AttackConversationPage | None: """ - Read stored conversation references without loading attack relationships or message bodies. + Read main and pruned references without loading attack relationships or message bodies. Returns: AttackConversationPage | None: One bounded reference page, or None for an absent attack. @@ -8186,8 +8186,6 @@ def _execute_get_attack_conversation_page( statement = select( AttackResultEntry.conversation_id, AttackResultEntry.pruned_conversation_ids, - AttackResultEntry.adversarial_chat_conversation_ids, - AttackResultEntry.preparation_conversation_ids, AttackResultEntry.timestamp, func.coalesce(AttackResultEntry.human_score_id, AttackResultEntry.automated_score_id).label( "objective_score_id" @@ -8201,18 +8199,13 @@ def _execute_get_attack_conversation_page( main_id = row["conversation_id"] or "" if main_id: references[main_id] = AttackConversationRecord(conversation_id=main_id) - for column_name, kind in ( - ("pruned_conversation_ids", ConversationType.PRUNED), - ("adversarial_chat_conversation_ids", ConversationType.ADVERSARIAL), - ("preparation_conversation_ids", ConversationType.PREPARATION), - ): - values = row[column_name] - if not isinstance(values, list): - continue + values = row["pruned_conversation_ids"] + if isinstance(values, list): for identifier in values: if isinstance(identifier, str) and identifier: references.setdefault( - identifier, AttackConversationRecord(conversation_id=identifier, conversation_type=kind) + identifier, + AttackConversationRecord(conversation_id=identifier, conversation_type=ConversationType.PRUNED), ) ordered = tuple(references.values()) revision = hashlib.sha256( @@ -11818,7 +11811,7 @@ async def get_attack_conversation_page_async( conversation_id: str | None = None, ) -> AttackConversationPage | None: """ - Read one bounded page of stored main/related conversation references. + Read one bounded page of stored main and pruned conversation references. Returns: AttackConversationPage | None: Saved references and ordering revision, or an absent attack. diff --git a/tests/unit/backend/test_saved_conversation_read_service.py b/tests/unit/backend/test_saved_conversation_read_service.py index a863365f4b..4ab5be658b 100644 --- a/tests/unit/backend/test_saved_conversation_read_service.py +++ b/tests/unit/backend/test_saved_conversation_read_service.py @@ -48,7 +48,7 @@ async def saved_pieces(sqlite_instance: SQLiteMemory, saved_attack: AttackResult return pieces -async def test_reference_pages_include_every_saved_context_without_message_hydration_async( +async def test_reference_pages_include_only_main_and_pruned_without_message_hydration_async( sqlite_instance: SQLiteMemory, saved_attack: AttackResult, saved_pieces: list[MessagePiece] ) -> None: reader = SavedConversationReadService(memory=sqlite_instance) @@ -57,14 +57,14 @@ async def test_reference_pages_include_every_saved_context_without_message_hydra patch.object(sqlite_instance, "get_conversation_messages_async", new_callable=AsyncMock) as full_messages, patch.object(PromptMemoryEntry, "get_message_piece", autospec=True) as hydrate, ): - first = await reader.list_conversations_async(attack_result_id=saved_attack.attack_result_id, limit=2) + first = await reader.list_conversations_async(attack_result_id=saved_attack.attack_result_id, limit=1) assert first is not None - assert first.total_count == 4 + assert first.total_count == 2 assert first.pagination is not None assert first.pagination.has_more assert first.pagination.next_cursor is not None second = await reader.list_conversations_async( - attack_result_id=saved_attack.attack_result_id, limit=2, cursor=first.pagination.next_cursor + attack_result_id=saved_attack.attack_result_id, limit=1, cursor=first.pagination.next_cursor ) assert second is not None assert second.pagination is not None @@ -76,16 +76,13 @@ async def test_reference_pages_include_every_saved_context_without_message_hydra references = first.conversations + second.conversations assert [reference.conversation_id for reference in references] == [ saved_attack.conversation_id, - *( - next( - reference.conversation_id - for reference in saved_attack.related_conversations - if reference.conversation_type == kind - ) - for kind in (ConversationType.PRUNED, ConversationType.ADVERSARIAL, ConversationType.PREPARATION) + next( + reference.conversation_id + for reference in saved_attack.related_conversations + if reference.conversation_type == ConversationType.PRUNED ), ] - assert [reference.conversation_type for reference in references] == ["main", "pruned", "adversarial", "preparation"] + assert [reference.conversation_type for reference in references] == ["main", "pruned"] assert references[0].message_count == 2 assert references[1].message_count == 0 @@ -163,7 +160,11 @@ async def test_saved_cursors_reject_other_owners_and_changed_reference_order_asy with pytest.raises(ValueError, match="cursor"): await reader.get_messages_async( attack_result_id=saved_attack.attack_result_id, - conversation_id=next(iter(saved_attack.related_conversations)).conversation_id, + conversation_id=next( + reference.conversation_id + for reference in saved_attack.related_conversations + if reference.conversation_type == ConversationType.PRUNED + ), limit=1, cursor=message_page.pagination.next_cursor, ) diff --git a/tests/unit/backend/test_scenario_hierarchy_pagination.py b/tests/unit/backend/test_scenario_hierarchy_pagination.py index 43062a8f64..19203cec09 100644 --- a/tests/unit/backend/test_scenario_hierarchy_pagination.py +++ b/tests/unit/backend/test_scenario_hierarchy_pagination.py @@ -18,7 +18,15 @@ from pyrit.backend.services.message_send_service import MessageSendService from pyrit.backend.services.scenario_run_service import ScenarioRunService from pyrit.memory import SQLiteMemory, set_message_piece_sha256_async -from pyrit.models import AttackOutcome, AttackResult, AttackResultRole, MessagePiece, ScenarioHierarchyNode +from pyrit.models import ( + AttackOutcome, + AttackResult, + AttackResultRole, + ConversationReference, + ConversationType, + MessagePiece, + ScenarioHierarchyNode, +) from unit.mocks import get_mock_target_identifier, make_scenario_result @@ -335,6 +343,14 @@ async def test_api_journey_loads_only_the_selected_saved_conversation_async( ) -> None: run_id = await _store_run_async(sqlite_instance) child = _result(run_id=run_id, index=1) + related = { + kind: str(uuid.uuid4()) + for kind in (ConversationType.PRUNED, ConversationType.ADVERSARIAL, ConversationType.PREPARATION) + } + child.related_conversations = { + ConversationReference(conversation_id=identifier, conversation_type=kind) + for kind, identifier in related.items() + } parent = _result(run_id=run_id, index=0, child_ids=[child.attack_result_id]) await sqlite_instance.add_attack_results_to_memory_async(attack_results=[parent, child]) pieces = [ @@ -346,9 +362,13 @@ async def test_api_journey_loads_only_the_selected_saved_conversation_async( ) for index in range(3) ] - for piece in pieces: + context_pieces = [ + MessagePiece(conversation_id=identifier, original_value=f"Saved {kind.value}", role="user", sequence=0) + for kind, identifier in related.items() + ] + for piece in pieces + context_pieces: await set_message_piece_sha256_async(piece) - await sqlite_instance.add_message_pieces_to_memory_async(message_pieces=pieces) + await sqlite_instance.add_message_pieces_to_memory_async(message_pieces=pieces + context_pieces) scenario_service = ScenarioRunService(memory=sqlite_instance) sends = MagicMock(spec=MessageSendService) attack_service = AttackService(memory=sqlite_instance, message_send_service=sends) @@ -363,11 +383,38 @@ async def test_api_journey_loads_only_the_selected_saved_conversation_async( node = client.get(f"/api/scenarios/runs/{run_id}/hierarchy/nodes/{child.attack_result_id}") assert node.status_code == 200 assert "messages" not in node.json()["conversations"][0] - references = client.get(f"/api/attacks/{child.attack_result_id}/conversations", params={"limit": 1}) + assert node.json()["conversations_count"] == 2 + assert node.json()["conversations_truncated"] is False + conversations_url = f"/api/attacks/{child.attack_result_id}/conversations" + unpaged = client.get(conversations_url) + assert unpaged.status_code == 200 + assert {conversation["conversation_id"] for conversation in unpaged.json()["conversations"]} == { + child.conversation_id, + related[ConversationType.PRUNED], + } + references = client.get(conversations_url, params={"limit": 1}) assert references.status_code == 200 - assert references.json()["conversations"][0]["message_count"] == 1 - assert node.json()["conversations"] == references.json()["conversations"] + reference_body = references.json() + assert reference_body["conversations"][0]["message_count"] == 1 + assert reference_body["total_count"] == 2 + remaining = client.get( + conversations_url, params={"limit": 1, "cursor": reference_body["pagination"]["next_cursor"]} + ) + assert remaining.status_code == 200 + assert remaining.json()["pagination"]["has_more"] is False + assert node.json()["conversations"] == reference_body["conversations"] + remaining.json()["conversations"] messages_url = f"/api/attacks/{child.attack_result_id}/messages" + for limit in (None, 2): + for kind, identifier in related.items(): + params: dict[str, str | int] = {"conversation_id": identifier} + if limit is not None: + params["limit"] = limit + response = client.get(messages_url, params=params) + if kind == ConversationType.PRUNED: + assert response.status_code == 200 + assert response.json()["messages"][0]["message_pieces"][0]["original_value"] == "Saved pruned" + else: + assert response.status_code == 400 first = client.get(messages_url, params={"conversation_id": child.conversation_id, "limit": 2}) assert first.status_code == 200 first_body = first.json() diff --git a/tests/unit/backend/test_scenario_hierarchy_parity.py b/tests/unit/backend/test_scenario_hierarchy_parity.py index 1f56fa9ca4..eef2033667 100644 --- a/tests/unit/backend/test_scenario_hierarchy_parity.py +++ b/tests/unit/backend/test_scenario_hierarchy_parity.py @@ -1247,15 +1247,18 @@ async def test_sqlite_conversation_list_clipping_explicit(sqlite_instance: SQLit 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)] + pruned_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 + ConversationReference(conversation_id=cid, conversation_type=ConversationType.PRUNED) + for cid in pruned_conv_ids + } + | { + ConversationReference(conversation_id=str(uuid.uuid4()), conversation_type=kind) + for kind in (ConversationType.ADVERSARIAL, ConversationType.PREPARATION) }, objective="Conversation clipping test", outcome=AttackOutcome.SUCCESS, @@ -1280,7 +1283,7 @@ async def test_sqlite_conversation_list_clipping_explicit(sqlite_instance: SQLit assert len(remaining.conversations) == 5 assert {conversation.conversation_id for conversation in conv_result.conversations + remaining.conversations} == { main_conv_id, - *prep_conv_ids, + *pruned_conv_ids, } # 2. Service single-node read projects exact count and truncation flag on node @@ -1443,14 +1446,14 @@ async def test_unloaded_conversations_in_full_hierarchy(sqlite_instance: SQLiteM 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)] + pruned_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 + ConversationReference(conversation_id=cid, conversation_type=ConversationType.PRUNED) + for cid in pruned_conv_ids }, objective="Unloaded conversations test", outcome=AttackOutcome.SUCCESS,