Repository navigation
feat: Add bounded Scenario result hierarchy projections (#3145) - #3167
Nimit Jain (Nimit3418) wants to merge 1 commit into
Conversation
| 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), |
There was a problem hiding this comment.
🔴 Must fix: Read result metrics from their persisted fields.
executed_turns and outcome_reason are first-class AttackResultEntry columns, not values that ordinary producers put in attack_metadata or attribution_data. Neither delta query selects those columns. Persisting a real result with executed_turns=7 and an outcome reason, then reading this node, returns executed_turns=0 and outcome_reason=None. The other node/parent builders have the same problem. The synthetic tests hide it by inserting executed_turns into metadata.
Please include the canonical fields in the lightweight delta and project those values, with defaults only for genuinely absent legacy data. Cover the round trip with a real persisted result.
| .where( | ||
| and_( | ||
| AttackResultEntry.attribution_parent_id == scenario_uuid, | ||
| AttackResultEntry.id.in_(valid_uuids), |
There was a problem hiding this comment.
🔴 Must fix: Split child lookups below the backend parameter limit.
The node endpoint passes the entire stored child list into this single IN query. Compiling 2,100 child IDs for SQL Server produces 2,101 bind parameters after including the scenario ID, which exceeds SQL Server's 2,100-parameter limit. Duplicate valid references also consume parameters, even though they cannot add rows. The existing compilation test with two IDs misses this failure.
Please deduplicate IDs for fetching and split bounded batches using the backend's _MAX_BIND_VARS, reserving space for fixed parameters. Keep the original ordered references for assembly and make any detail clipping explicit.
| 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]: |
There was a problem hiding this comment.
🔴 Must fix: Apply the message limit in the database read.
_execute_get_conversation_messages() loads and groups every message piece in the conversation. Slicing its result only limits the response, not hydration. With 151 saved messages and limit_messages_per_conversation=5, this still hydrates all 151 before returning five. The related-conversation loop below repeats the same full read once per conversation, with no bound on the number of conversations.
Please use bounded, batched message/piece reads and an explicit conversation budget, while obtaining exact counts separately. A small node detail request should not have to load an arbitrarily large conversation history.
| 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) |
There was a problem hiding this comment.
🔴 Must fix: Keep cyclic components visible when another root exists.
The fallback runs only when the whole run has no roots. With three saved nodes, ordinary, A -> B, and B -> A, ordinary is a root, so neither A nor B is ever visited. The response contains only ordinary, with no unresolved cycle links or indication that two results vanished. The current cycle test passes because it contains only the cycle and takes the fallback.
Please also account for components not reached from the normal roots, using the same bounded cycle handling without inventing a parent relationship. Add a mixed ordinary-plus-cycle case so malformed lineage stays explicit rather than silently disappearing.
| 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, | ||
| ) |
There was a problem hiding this comment.
🔴 Must fix: Bound the hierarchy read before loading the full run.
get_snapshot_async() keeps fetching 500-row pages until every result is cached, and build_scenario_hierarchy() returns every root. With a cold SQLite cache and 1,001 saved results for one logical unit, this request reads all 1,001 results in three queries and returns 1,001 roots, about 923 KB of JSON. There is no result/page budget or truncation indicator. The depth guard does not bound breadth either: a shared DAG with 13 saved nodes expands to 127 payload nodes because descendants are rebuilt for each path.
Please make the hierarchy's node/edge budget drive bounded memory reads, with paging or explicit clipping, while keeping the run-wide totals exact. Reusing the full progress snapshot here does not satisfy #3145's bounded projection contract.
| 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, | ||
| ) |
There was a problem hiding this comment.
🔴 Must fix: Preserve the shared outcome statistics in group counts.
ScenarioAtomicGroupProgress already contains outcomes, but this copy leaves that field out, so every planned hierarchy group's counts.outcomes becomes null. For a group with one success and one undetermined result, the existing summary has a 100% decided-only rate and a 50% all-outcome rate. Both disappear from the hierarchy response, even though the overall response still has them.
Please carry the complete canonical counts through, including ag.outcomes, rather than copying only the older count fields. Add a parity case that checks both denominator policies at the group level.
| 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, | ||
| ) |
There was a problem hiding this comment.
🔴 Must fix: Reuse the shared counting policy for unmatched groups.
This fallback counts historical rows instead of selecting each logical unit's latest attempt. For an unmatched unit whose old attempt succeeded and latest attempt errored, it reports a completed success and 100% success. The shared summary correctly keeps those attempts unattributed. Also, ScenarioProducerCounts has target_facing, orchestration, and unknown fields, not attempts or errors, so Pydantic silently discards both supplied arguments and every producer count is zero. The type checker catches that mismatch too.
Please preserve unmatched attempts and their exact unattributed/raw totals, and use the existing logical-unit and producer calculators instead of defining another counting policy here. The hierarchy should agree with progress and SDK accounting.
| all_attack_retries.append( | ||
| AttackRetrySummary( | ||
| attack_result_id=result.attack_result_id, | ||
| atomic_attack_name=result.atomic_attack_name, | ||
| retries=result.retries, | ||
| ) |
There was a problem hiding this comment.
🔴 Must fix: Bound retry events inside each diagnostic item.
The outer attack_retries list is capped later, but each item here retains the entire result.retries list. A single result with 201 retry events gets 100 events and retries_truncated=True in its node, but all 201 events appear again in diagnostics, where attack_retries_truncated is false. One retry-heavy result can therefore bypass the diagnostic bound entirely.
Please cap diagnostic retry events too and expose truncation at the appropriate level, without changing total_retries or the exact diagnostic totals. A boundary/above-boundary test should cover a long retry list within one item, not only many one-event items.
| 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, |
There was a problem hiding this comment.
🔴 Must fix: Preserve the generic message evidence.
Message.get_value() defaults to piece zero, so a message with two text pieces loses the second piece here. A text-plus-image message likewise loses the image and its data type. msg.api_role also changes a stored simulated_assistant or simulated_tool into an ordinary assistant/tool role, with no simulated marker in this projection. A saved two-piece simulated reply currently comes back as just the first string with role="assistant".
Please reuse the generic message/piece projection and preserve the stored role and piece types rather than flattening each message to its first value. This is a read-only evidence view, not a target-facing API request. Apply the same fix to related conversations below.
Description
This PR implements bounded Scenario result hierarchy projections as requested in #3145. It provides a typed, bounded read-only projection resolving parent/child orchestration relationships, ordered direct child references, and generic conversation projections without requiring full-run result or message hydration.
Key Changes:
pyrit/models/scenario_hierarchy.py):ScenarioHierarchyProjection,ScenarioHierarchyNode,ScenarioOrchestrationParentSummary,ScenarioHierarchyGroup,ScenarioHierarchyDiagnostics, and generic read-onlyScenarioConversationProjection.child_index(0-indexed position), producerattempt_index(1-indexed), and explicit fallback toAttackResultRole.UNKNOWN.pyrit/memory/memory_interface.py):get_scenario_attack_result_deltas_by_ids_asyncexecuting a single batched query matchingid IN (...)andattribution_parent_id == scenario_uuidto ensure cross-run isolation.get_attack_result_conversation_projections_asyncto fetch conversation stats and bounded messages without hydrating unrelated run data.pyrit/backend/services/scenario_progress_read_model.py):get_run_hierarchy_node_asyncfetch only the targeted node and immediate children.active_pathand recursion depth guard (MAX_HIERARCHY_DEPTH = 10).unresolved_child_attack_result_idswithout failing the run.failed_attacks,attack_retries) are capped at 100 while preserving exact total counters.ag.technique_details.producer_countsandOutcomeStatisticsfrom FEAT: Add role-aware Scenario target-attempt accounting (#3043) #3062 and FEAT: Attack analytics SDK #3059 unchanged.pyrit/backend/routes/scenarios.py):GET /api/scenarios/runs/{scenario_result_id}/hierarchy.GET /api/scenarios/runs/{scenario_result_id}/hierarchy/nodes/{attack_result_id}.Closes #3145
Testing
tests/unit/models/test_scenario_hierarchy_models.py.tests/unit/memory/test_azure_sql_memory.py.tests/unit/backend/test_scenario_hierarchy_parity.pycovering ordinary results, nested sequential attacks, cyclic reference guards, retained FIX: preserve child-result links when execution fails in SequentialAttack #3048 failure links, missing child references, diagnostic boundary truncation, bounded single-node reads, and REST routes.make unit-test: all 26,162 tests passed.ruff checkandruff format: all checks passed cleanly.