From 01790cf998d9431b5b4626db885dd22cbeb65a60 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Wed, 29 Jul 2026 16:39:18 -0700 Subject: [PATCH] MAINT: Order AttackResults by timestamp for perf Make the indexable AttackResultEntry.timestamp column the single source of truth for History recency, replacing the non-indexable JSON-extract sort expression. Adds serving indexes + an Alembic migration that backfills timestamp from the legacy attack_metadata.updated_at/created_at keys. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 5ae8d0da-5198-4ac0-acd9-65e6f7b5d283 --- pyrit/backend/mappers/attack_mappers.py | 10 +- pyrit/backend/services/attack_service.py | 50 +++--- ...7e9f1a3b5c6_index_attack_result_recency.py | 153 ++++++++++++++++++ pyrit/memory/azure_sql_memory.py | 31 +--- pyrit/memory/memory_interface.py | 88 +++------- pyrit/memory/memory_models.py | 8 +- pyrit/memory/sqlite_memory.py | 31 +--- tests/unit/backend/test_attack_service.py | 28 ++-- tests/unit/backend/test_mappers.py | 16 ++ .../test_interface_attack_results.py | 70 ++------ tests/unit/memory/test_azure_sql_memory.py | 20 --- tests/unit/memory/test_migration.py | 138 ++++++++++++++++ 12 files changed, 389 insertions(+), 254 deletions(-) create mode 100644 pyrit/memory/alembic/versions/d7e9f1a3b5c6_index_attack_result_recency.py diff --git a/pyrit/backend/mappers/attack_mappers.py b/pyrit/backend/mappers/attack_mappers.py index fd856791a9..5276aad277 100644 --- a/pyrit/backend/mappers/attack_mappers.py +++ b/pyrit/backend/mappers/attack_mappers.py @@ -231,22 +231,22 @@ def _resolve_summary_timestamps(ar: AttackResult) -> tuple[datetime, datetime]: """ Resolve ``created_at`` / ``updated_at`` for a summary. - Resolution order for ``created_at``: explicit metadata override, then the - persisted ``AttackResult.timestamp``, and finally ``datetime.now`` as a - last-resort fallback for never-persisted results. + ``updated_at`` is the persisted ``AttackResult.timestamp`` — the single indexed recency + key that manual edits (add message, branch, promote conversation) bump. ``created_at`` + comes from the display-only ``metadata.created_at`` override, falling back to + ``AttackResult.timestamp`` and finally ``datetime.now`` for never-persisted results. Returns: A ``(created_at, updated_at)`` tuple. """ created_str = ar.metadata.get("created_at") - updated_str = ar.metadata.get("updated_at") if created_str: created_at = datetime.fromisoformat(created_str) elif ar.timestamp is not None: created_at = ar.timestamp else: created_at = datetime.now(timezone.utc) - updated_at = datetime.fromisoformat(updated_str) if updated_str else created_at + updated_at = ar.timestamp if ar.timestamp is not None else created_at return created_at, updated_at diff --git a/pyrit/backend/services/attack_service.py b/pyrit/backend/services/attack_service.py index ca6a00f36c..cab5916755 100644 --- a/pyrit/backend/services/attack_service.py +++ b/pyrit/backend/services/attack_service.py @@ -358,9 +358,9 @@ async def create_attack_async(self, *, request: CreateAttackRequest) -> CreateAt ), ), outcome=AttackOutcome.UNDETERMINED, + timestamp=now, metadata={ "created_at": now.isoformat(), - "updated_at": now.isoformat(), }, labels=labels, ) @@ -414,15 +414,11 @@ async def update_attack_async(self, *, attack_result_id: str, request: UpdateAtt } new_outcome = outcome_map.get(request.outcome, AttackOutcome.UNDETERMINED) - ar = results[0] - updated_metadata = dict(ar.metadata) if ar.metadata else {} - updated_metadata["updated_at"] = datetime.now(timezone.utc).isoformat() - self._memory.update_attack_result_by_id( attack_result_id=attack_result_id, update_fields={ "outcome": new_outcome.value, - "attack_metadata": updated_metadata, + "timestamp": datetime.now(timezone.utc), }, ) @@ -527,14 +523,11 @@ async def create_related_conversation_async( # Add to pruned_conversation_ids so user-created branches are visible in the GUI history panel. existing_pruned = ar.get_pruned_conversation_ids() - updated_metadata = dict(ar.metadata or {}) - updated_metadata["updated_at"] = now.isoformat() - self._memory.update_attack_result_by_id( attack_result_id=attack_result_id, update_fields={ "pruned_conversation_ids": existing_pruned + [new_conversation_id], - "attack_metadata": updated_metadata, + "timestamp": now, }, ) @@ -590,8 +583,6 @@ async def update_main_conversation_async( updated_pruned.append(ar.conversation_id) now = datetime.now(timezone.utc) - updated_metadata = dict(ar.metadata or {}) - updated_metadata["updated_at"] = now.isoformat() self._memory.update_attack_result_by_id( attack_result_id=attack_result_id, @@ -599,7 +590,7 @@ async def update_main_conversation_async( "conversation_id": target_conv_id, "pruned_conversation_ids": updated_pruned if updated_pruned else None, "adversarial_chat_conversation_ids": updated_adversarial if updated_adversarial else None, - "attack_metadata": updated_metadata, + "timestamp": now, }, ) @@ -752,12 +743,12 @@ async def _update_attack_after_message_async( self, *, attack_result_id: str, ar: AttackResult, request: AddMessageRequest ) -> None: """ - Update attack metadata and converter tracking after a message is added. - """ - updated_metadata = dict(ar.metadata or {}) - updated_metadata["updated_at"] = datetime.now(timezone.utc).isoformat() + Update attack recency and converter tracking after a message is added. - update_fields: dict[str, Any] = {"attack_metadata": updated_metadata} + Bumps the attack's ``timestamp`` column (the single indexed recency key) so the edited + conversation re-floats to the top of the History view. + """ + update_fields: dict[str, Any] = {"timestamp": datetime.now(timezone.utc)} if request.converter_ids: converter_objs = get_converter_service().get_converter_objects_for_ids(converter_ids=request.converter_ids) @@ -924,18 +915,16 @@ def _encode_attack_cursor(*, cursor: AttackResultsKeysetCursor, fingerprint: str """ Encode a keyset anchor and its filter fingerprint into an opaque pagination cursor. - The anchor's recency string and timestamp can contain ``.``/``:``/``-`` (ISO - timestamps), so the payload is JSON-serialized and base64url-encoded rather than - joined with a delimiter, keeping the cursor an unambiguous opaque token for - ``_decode_attack_cursor``. + The anchor's timestamp can contain ``.``/``:``/``-`` (ISO timestamps), so the payload + is JSON-serialized and base64url-encoded rather than joined with a delimiter, keeping + the cursor an unambiguous opaque token for ``_decode_attack_cursor``. Returns: - An opaque base64url cursor string encoding ``{fingerprint, recency, timestamp, + An opaque base64url cursor string encoding ``{fingerprint, timestamp, attack_result_id}``. """ payload = { "f": fingerprint, - "r": cursor.recency, "t": cursor.timestamp.isoformat(), "i": cursor.attack_result_id, } @@ -947,12 +936,12 @@ def _decode_attack_cursor(*, cursor: str | None, fingerprint: str) -> AttackResu """ Decode the opaque list-attacks cursor into a keyset (seek) anchor. - The cursor encodes the previous page's last-row recency anchor together with a + The cursor encodes the previous page's last-row timestamp anchor together with a fingerprint of the filter set it was generated for (see ``_attack_filter_fingerprint``). A cursor is honored only when its fingerprint matches the current request's filters; - malformed, legacy (offset/attack-result-id), or filter-mismatched cursors fall back to - the first page (``None``) so a stale cursor degrades gracefully instead of raising or - seeking within the wrong result set. + malformed, legacy (offset/attack-result-id/recency-string), or filter-mismatched cursors + fall back to the first page (``None``) so a stale cursor degrades gracefully instead of + raising or seeking within the wrong result set. Returns: The decoded ``AttackResultsKeysetCursor``, or ``None`` to start at the first page. @@ -966,10 +955,9 @@ def _decode_attack_cursor(*, cursor: str | None, fingerprint: str) -> AttackResu return None if not isinstance(payload, dict) or payload.get("f") != fingerprint: return None - recency = payload.get("r") raw_timestamp = payload.get("t") attack_result_id = payload.get("i") - if not isinstance(recency, str) or not isinstance(raw_timestamp, str) or not isinstance(attack_result_id, str): + if not isinstance(raw_timestamp, str) or not isinstance(attack_result_id, str): return None try: timestamp = datetime.fromisoformat(raw_timestamp) @@ -989,7 +977,7 @@ def _decode_attack_cursor(*, cursor: str | None, fingerprint: str) -> AttackResu # A crafted cursor near datetime's min/max with a large UTC offset overflows the # representable range when shifted to UTC; treat it as malformed and restart at page one. return None - return AttackResultsKeysetCursor(recency=recency, timestamp=timestamp, attack_result_id=attack_result_id) + return AttackResultsKeysetCursor(timestamp=timestamp, attack_result_id=attack_result_id) # ======================================================================== # Private Helper Methods - Duplicate / Branch diff --git a/pyrit/memory/alembic/versions/d7e9f1a3b5c6_index_attack_result_recency.py b/pyrit/memory/alembic/versions/d7e9f1a3b5c6_index_attack_result_recency.py new file mode 100644 index 0000000000..99d30b23f4 --- /dev/null +++ b/pyrit/memory/alembic/versions/d7e9f1a3b5c6_index_attack_result_recency.py @@ -0,0 +1,153 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Index AttackResultEntries for the History recency query and make the ``timestamp`` +column the single source of truth for "last updated". + +Adds an index on ``conversation_id`` (serves the per-conversation dedup window) and a +composite ``(timestamp, id)`` index (serves the recency ORDER BY + keyset seek). Backfills +``timestamp`` from the legacy ``attack_metadata.updated_at``/``created_at`` JSON keys so +manually-edited conversations keep their current History order, and drops the now-redundant +``updated_at`` key from the JSON metadata. + +Revision ID: d7e9f1a3b5c6 +Revises: 3f6e8a0c2d4b +Create Date: 2026-07-20 12:00:00.000000 +""" + +from __future__ import annotations + +import logging +from collections.abc import Sequence # noqa: TC003 +from datetime import datetime + +import sqlalchemy as sa +from alembic import op + +from pyrit.memory.memory_models import CustomUUID, UTCDateTime + +# revision identifiers, used by Alembic. +revision: str = "d7e9f1a3b5c6" +down_revision: str | Sequence[str] | None = "3f6e8a0c2d4b" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +logger = logging.getLogger(__name__) + + +def _attack_results_table() -> sa.Table: + """ + Build a lightweight typed table clause for reading/writing the columns this migration + touches, so ``UTCDateTime``/``JSON`` bind and result processors run on both backends. + + Returns: + sa.Table: A minimal ``AttackResultEntries`` table with the id, attack_metadata, and + timestamp columns. + """ + metadata_obj = sa.MetaData() + return sa.Table( + "AttackResultEntries", + metadata_obj, + sa.Column("id", CustomUUID(), primary_key=True), + sa.Column("attack_metadata", sa.JSON()), + sa.Column("timestamp", UTCDateTime()), + extend_existing=True, + ) + + +def _parse_iso(value: object) -> datetime | None: + """ + Parse an ISO-8601 string (tolerating a trailing ``Z``) into a datetime, or ``None``. + + Returns: + datetime | None: The parsed datetime, or ``None`` when ``value`` is not a parseable + ISO string. + """ + if not isinstance(value, str) or not value: + return None + try: + return datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError: + return None + + +def upgrade() -> None: + """Apply this schema upgrade.""" + op.create_index( + "ix_AttackResultEntries_conversation_id", + "AttackResultEntries", + ["conversation_id"], + ) + op.create_index( + "ix_AttackResultEntries_timestamp_id", + "AttackResultEntries", + ["timestamp", "id"], + ) + _backfill_timestamp_from_metadata() + + +def downgrade() -> None: + """Revert this schema upgrade.""" + # Restore metadata.updated_at from the timestamp column so the legacy JSON recency sort + # reflects each row's last-updated time again. The original created/updated split cannot be + # perfectly recovered (it was collapsed into the single timestamp), so updated_at is set to + # the current timestamp value. + _restore_updated_at_to_metadata() + op.drop_index("ix_AttackResultEntries_timestamp_id", table_name="AttackResultEntries") + op.drop_index("ix_AttackResultEntries_conversation_id", table_name="AttackResultEntries") + + +def _backfill_timestamp_from_metadata() -> None: + """ + Set ``timestamp = COALESCE(parse(metadata.updated_at), parse(metadata.created_at), + timestamp)`` for existing rows and drop the now-redundant ``updated_at`` JSON key. + + Only rows that carried a JSON recency key change; programmatic rows already have + ``timestamp`` equal to their creation time. Idempotent: rerunning is a no-op once the + ``updated_at`` key is gone and ``timestamp`` already reflects it. + """ + bind = op.get_bind() + table = _attack_results_table() + + rows = bind.execute(sa.select(table.c.id, table.c.attack_metadata, table.c.timestamp)).fetchall() + + updated_rows = 0 + for row in rows: + metadata = dict(row.attack_metadata) if isinstance(row.attack_metadata, dict) else {} + new_timestamp = _parse_iso(metadata.get("updated_at")) or _parse_iso(metadata.get("created_at")) + + had_updated_at = "updated_at" in metadata + new_metadata = {k: v for k, v in metadata.items() if k != "updated_at"} + + timestamp_changes = new_timestamp is not None and new_timestamp != row.timestamp + if not timestamp_changes and not had_updated_at: + continue + + values: dict[str, object] = {"attack_metadata": new_metadata or None} + if timestamp_changes: + values["timestamp"] = new_timestamp + bind.execute(sa.update(table).where(table.c.id == row.id).values(**values)) + updated_rows += 1 + + if updated_rows: + logger.info(f"Attack recency backfill: updated {updated_rows} AttackResultEntries row(s).") + + +def _restore_updated_at_to_metadata() -> None: + """ + Downgrade inverse: write ``metadata.updated_at = timestamp.isoformat()`` for every row so + the legacy JSON recency expression orders rows by their last-updated time again. + """ + bind = op.get_bind() + table = _attack_results_table() + + rows = bind.execute(sa.select(table.c.id, table.c.attack_metadata, table.c.timestamp)).fetchall() + + for row in rows: + if row.timestamp is None: + continue + metadata = dict(row.attack_metadata) if isinstance(row.attack_metadata, dict) else {} + metadata["updated_at"] = row.timestamp.isoformat() + bind.execute(sa.update(table).where(table.c.id == row.id).values(attack_metadata=metadata)) diff --git a/pyrit/memory/azure_sql_memory.py b/pyrit/memory/azure_sql_memory.py index 6f3092c192..102e8d3c2b 100644 --- a/pyrit/memory/azure_sql_memory.py +++ b/pyrit/memory/azure_sql_memory.py @@ -8,7 +8,7 @@ from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Any, Literal, cast -from sqlalchemy import and_, column, create_engine, event, exists, func, select, text +from sqlalchemy import and_, create_engine, event, exists, text from sqlalchemy.engine.base import Engine from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.orm import InstrumentedAttribute, sessionmaker @@ -441,35 +441,6 @@ def _get_condition_json_array_match( combined = joiner.join(conditions) return text(f"""ISJSON("{table_name}".{column_name}) = 1 AND ({combined})""").bindparams(**bindparams_dict) - def _attack_results_recency_expr(self) -> Any: - """ - Return the Azure SQL recency sort expression reproducing the History-view sort. - - Coalesces the ``attack_metadata`` ``updated_at`` key, falling back to ``created_at`` - then an empty string. Only JSON *string* values are honored: each accessor is an - ``OPENJSON`` subquery filtered to ``type = 1`` (JSON string), so a non-string value - (e.g. a JSON bool or number) yields no row and coalesces to ``""`` instead. A raw - ``JSON_VALUE`` would instead stringify such a scalar (``false`` -> ``"false"``, ``2`` -> - ``"2"``), diverging from ``_attack_result_recency_key``'s Python anchor -- which collapses - non-strings to ``""`` -- so the keyset seek would anchor on a value the database never - sorted by and skip rows at a page boundary. Mirroring SQLite's ``json_type`` guard keeps - recency a pure text sort key that agrees with the Python key on both backends. Uses ORM - ``func``/``select`` expressions (not raw ``text``) so they adapt to the aliased subquery - SQLAlchemy emits when a collection ``joinedload`` is combined with ``limit``. Backs both - the recency ORDER BY and the keyset seek predicate so they stay in lockstep. - - Returns: - Any: A SQLAlchemy expression yielding the coalesced recency string. - """ - metadata = AttackResultEntry.attack_metadata - - def _string_valued(json_key: str) -> Any: - table = func.openjson(metadata).table_valued(column("key"), column("value"), column("type")) - rows = table.alias(f"oj_{json_key}") - return select(rows.c["value"]).where(rows.c["key"] == json_key, rows.c["type"] == 1).scalar_subquery() - - return func.coalesce(_string_valued("updated_at"), _string_valued("created_at"), "") - def _get_attack_result_label_condition(self, *, labels: dict[str, str | Sequence[str]]) -> Any: """ Azure SQL implementation for filtering AttackResults by labels. diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index ae886f1f82..3d523fecc4 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -87,45 +87,18 @@ IdentifierModel = TypeVar("IdentifierModel", bound=ComponentIdentifier) -def _attack_result_recency_key(metadata: dict[str, Any]) -> str: - """ - Compute the recency sort key for an attack result exactly as the SQL layer does. - - Mirrors the ``COALESCE(, , '')`` expression built by - ``MemoryInterface._attack_results_recency_expr``. Recency is always an ISO-8601 timestamp - *string*; only string values are honored (via ``isinstance``) so a non-string - ``updated_at``/``created_at`` (e.g. a JSON bool or number) falls through exactly like the - SQLite expression, which keeps only JSON ``text`` values. This matters because SQLite's - ``json_extract`` returns a typed scalar for non-string JSON (``0`` for a bool) and sorts - numbers before text: a stringified anchor would then never advance past such a row and - repeat it. Keeping this in lockstep with the SQL expression is what makes the keyset anchor - land on the same row the database ordered by, so pagination never skips or repeats a result. - - Returns: - str: The coalesced recency key (``""`` when neither timestamp is a string). - """ - value = metadata.get("updated_at") - if not isinstance(value, str): - value = metadata.get("created_at") - if not isinstance(value, str): - value = "" - return value - - class AttackResultsKeysetCursor(NamedTuple): """ Keyset (seek) anchor identifying the last attack result on a page. - Captures the full recency ordering tuple — the coalesced ``updated_at``/``created_at`` - recency string, the row ``timestamp``, and the unique ``attack_result_id`` — so the next - page can seek to rows strictly after this one under the - ``recency DESC, timestamp DESC, id DESC`` ordering. Unlike a numeric offset — which shifts - every later row when a row is inserted or deleted between page loads — a keyset anchor - stays pinned to its row, so inserts and deletions elsewhere no longer skip or duplicate - results at the page boundary. + Captures the recency ordering tuple — the row ``timestamp`` (the attack's last-updated + time; see ``AttackResultEntry.timestamp``) and the unique ``attack_result_id`` — so the + next page can seek to rows strictly after this one under the ``timestamp DESC, id DESC`` + ordering. Unlike a numeric offset — which shifts every later row when a row is inserted or + deleted between page loads — a keyset anchor stays pinned to its row, so inserts and + deletions elsewhere no longer skip or duplicate results at the page boundary. """ - recency: str timestamp: datetime attack_result_id: str @@ -138,7 +111,6 @@ def from_attack_result(cls, result: AttackResult) -> "AttackResultsKeysetCursor" AttackResultsKeysetCursor: Anchor capturing the result's recency sort key. """ return cls( - recency=_attack_result_recency_key(result.metadata), timestamp=result.timestamp, attack_result_id=result.attack_result_id, ) @@ -370,62 +342,42 @@ def _get_condition_json_array_match( Any: A database-specific SQLAlchemy condition. """ - @abc.abstractmethod - def _attack_results_recency_expr(self) -> Any: - """ - Return the SQL expression for an attack result's recency sort key. - - Concrete subclasses translate this into their SQL dialect's JSON accessor - (SQLite ``json_extract`` / Azure SQL ``JSON_VALUE``) so the sort key matches the - per-backend JSON abstraction the attack-result filters already use. It reproduces the - previous service-side Python sort: ``COALESCE(, - , '')``. This single expression backs both the recency - ORDER BY and the keyset seek predicate, guaranteeing they stay in lockstep (a keyset - seek that used a different expression than the ORDER BY would skip or repeat rows). - - Returns: - Any: A SQLAlchemy expression yielding the coalesced recency string. - """ - def _attack_results_recency_order_by(self) -> list[Any]: """ Return the ORDER BY clauses that reproduce the History-view recency sort. - Orders by the recency expression (``updated_at`` falling back to ``created_at`` then - an empty string), all descending, with the real ``timestamp`` column and ``id`` as - deterministic descending tie-breaks (required for stable keyset pagination). + Orders by the indexed ``timestamp`` column (an attack result's last-updated time — + bumped whenever the conversation is edited) descending, with ``id`` as a deterministic + descending tie-break (required for stable keyset pagination). Both columns are covered + by the composite ``(timestamp, id)`` index, so this ordering is index-served rather + than requiring a full sort. Returns: list[Any]: SQLAlchemy ORDER BY clauses (all descending) for the recency sort, suitable for splatting into ``_query_entries(order_by=...)``. """ - recency = self._attack_results_recency_expr() - return [recency.desc(), AttackResultEntry.timestamp.desc(), AttackResultEntry.id.desc()] + return [AttackResultEntry.timestamp.desc(), AttackResultEntry.id.desc()] def _attack_results_keyset_seek_condition(self, *, after: AttackResultsKeysetCursor) -> Any: """ Build the keyset seek predicate selecting rows strictly after ``after``. - Under the ``recency DESC, timestamp DESC, id DESC`` ordering, a row comes after the - anchor when its ordering tuple is lexicographically smaller. The predicate is - OR-expanded (rather than a row-value ``(a, b, c) < (x, y, z)`` comparison) because - Azure SQL does not support row-value tuple comparisons. It uses the same recency - expression as the ORDER BY so the seek lands exactly on the ordering boundary. The - ``id`` bound is a ``uuid.UUID`` so it is compared through the column's own type - (``CHAR(36)`` on SQLite, native ``uniqueidentifier`` on Azure), matching how the - ORDER BY sorts it on each backend. + Under the ``timestamp DESC, id DESC`` ordering, a row comes after the anchor when its + ordering tuple is lexicographically smaller. The predicate is OR-expanded (rather than + a row-value ``(a, b) < (x, y)`` comparison) because Azure SQL does not support + row-value tuple comparisons. The ``id`` bound is a ``uuid.UUID`` so it is compared + through the column's own type (``CHAR(36)`` on SQLite, native ``uniqueidentifier`` on + Azure), matching how the ORDER BY sorts it on each backend. Returns: Any: A SQLAlchemy boolean condition for the rows following the anchor. """ - recency = self._attack_results_recency_expr() timestamp = AttackResultEntry.timestamp entry_id = AttackResultEntry.id anchor_id = uuid.UUID(after.attack_result_id) return or_( - recency < after.recency, - and_(recency == after.recency, timestamp < after.timestamp), - and_(recency == after.recency, timestamp == after.timestamp, entry_id < anchor_id), + timestamp < after.timestamp, + and_(timestamp == after.timestamp, entry_id < anchor_id), ) def get_all_embeddings(self) -> Sequence[EmbeddingDataEntry]: diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index 24e8fb875c..04e8d3da75 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -1504,7 +1504,13 @@ class AttackResultEntry(Base): """ __tablename__ = "AttackResultEntries" - __table_args__ = {"extend_existing": True} + __table_args__ = ( + # Serves the PARTITION BY conversation_id dedup window in _query_paginated_attack_results. + Index("ix_AttackResultEntries_conversation_id", "conversation_id"), + # Serves the History recency ORDER BY timestamp DESC, id DESC and its keyset seek. + Index("ix_AttackResultEntries_timestamp_id", "timestamp", "id"), + {"extend_existing": True}, + ) id = mapped_column(CustomUUID, nullable=False, primary_key=True) conversation_id = mapped_column(String, nullable=False) objective = mapped_column(Unicode, nullable=False) diff --git a/pyrit/memory/sqlite_memory.py b/pyrit/memory/sqlite_memory.py index ed57c7e683..210c3f1da6 100644 --- a/pyrit/memory/sqlite_memory.py +++ b/pyrit/memory/sqlite_memory.py @@ -8,7 +8,7 @@ from pathlib import Path from typing import Any, Literal -from sqlalchemy import and_, case, create_engine, exists, func, or_, text +from sqlalchemy import and_, create_engine, exists, func, or_, text from sqlalchemy.engine.base import Engine from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.orm import InstrumentedAttribute, sessionmaker @@ -271,35 +271,6 @@ def _get_condition_json_array_match( combined = joiner.join(conditions) return text(f"({combined})").bindparams(**bindparams_dict) - def _attack_results_recency_expr(self) -> Any: - """ - Return the SQLite recency sort expression reproducing the History-view sort. - - Coalesces the ``attack_metadata`` ``updated_at`` key, falling back to ``created_at`` - then an empty string. Only JSON *string* values are honored: ``json_type`` gates each - accessor so a non-string value (e.g. a JSON bool, which ``json_extract`` would return - as the integer ``0``) collapses to ``NULL`` -> ``""`` instead. This keeps recency a - pure text sort key -- SQLite orders numbers before text, so a typed scalar would break - keyset seek monotonicity (the anchor could never advance past such a row). Backs both - the recency ORDER BY and the keyset seek predicate so they stay in lockstep, matching - ``_attack_result_recency_key``'s string-only Python key. - - Returns: - Any: A SQLAlchemy expression yielding the coalesced recency string. - """ - metadata = AttackResultEntry.attack_metadata - return func.coalesce( - case( - (func.json_type(metadata, "$.updated_at") == "text", func.json_extract(metadata, "$.updated_at")), - else_=None, - ), - case( - (func.json_type(metadata, "$.created_at") == "text", func.json_extract(metadata, "$.created_at")), - else_=None, - ), - "", - ) - def get_all_table_models(self) -> list[type[Base]]: """ Return a list of all table models used in the database by inspecting the Base registry. diff --git a/tests/unit/backend/test_attack_service.py b/tests/unit/backend/test_attack_service.py index a4e67966d9..a5b80b2a8b 100644 --- a/tests/unit/backend/test_attack_service.py +++ b/tests/unit/backend/test_attack_service.py @@ -104,6 +104,7 @@ def make_attack_result( ), outcome=outcome, attack_result_id=effective_ar_id, + timestamp=updated, metadata={ "created_at": created.isoformat(), "updated_at": updated.isoformat(), @@ -1043,8 +1044,8 @@ async def test_update_attack_updates_outcome_error(self, attack_service, mock_me call_kwargs = mock_memory.update_attack_result_by_id.call_args[1] assert call_kwargs["update_fields"]["outcome"] == "error" - async def test_update_attack_refreshes_updated_at(self, attack_service, mock_memory) -> None: - """Test that update_attack refreshes the updated_at metadata.""" + async def test_update_attack_bumps_timestamp(self, attack_service, mock_memory) -> None: + """Test that update_attack bumps the timestamp recency column and does not write metadata.""" old_time = datetime(2024, 1, 1, tzinfo=timezone.utc) ar = make_attack_result(conversation_id="test-id", updated_at=old_time) mock_memory.get_attack_results.return_value = [ar] @@ -1054,8 +1055,10 @@ async def test_update_attack_refreshes_updated_at(self, attack_service, mock_mem attack_result_id="test-id", request=UpdateAttackRequest(outcome="success") ) - call_kwargs = mock_memory.update_attack_result_by_id.call_args[1] - assert call_kwargs["update_fields"]["attack_metadata"]["updated_at"] != old_time.isoformat() + update_fields = mock_memory.update_attack_result_by_id.call_args[1]["update_fields"] + assert isinstance(update_fields["timestamp"], datetime) + assert update_fields["timestamp"] > old_time + assert "attack_metadata" not in update_fields # ============================================================================ @@ -1339,8 +1342,8 @@ async def test_add_message_raises_when_messages_not_found_after_update(self, att with pytest.raises(ValueError, match="messages not found after update"): await attack_service.add_message_async(attack_result_id="test-id", request=request) - async def test_add_message_persists_updated_at_timestamp(self, attack_service, mock_memory) -> None: - """Should persist updated_at in attack_metadata via update_attack_result.""" + async def test_add_message_bumps_timestamp(self, attack_service, mock_memory) -> None: + """Should bump the timestamp recency column via update_attack_result (no metadata write).""" ar = make_attack_result(conversation_id="test-id") ar.metadata = {"created_at": "2026-01-01T00:00:00+00:00"} mock_memory.get_attack_results.return_value = [ar] @@ -1359,9 +1362,9 @@ async def test_add_message_persists_updated_at_timestamp(self, attack_service, m mock_memory.update_attack_result_by_id.assert_called_once() call_kwargs = mock_memory.update_attack_result_by_id.call_args[1] assert call_kwargs["attack_result_id"] == "test-id" - persisted_metadata = call_kwargs["update_fields"]["attack_metadata"] - assert "updated_at" in persisted_metadata - assert persisted_metadata["created_at"] == "2026-01-01T00:00:00+00:00" + update_fields = call_kwargs["update_fields"] + assert isinstance(update_fields["timestamp"], datetime) + assert "attack_metadata" not in update_fields async def test_converter_ids_propagate_even_when_preconverted(self, attack_service, mock_memory) -> None: """Test that converter identifiers propagate to attack_identifier even when pieces are preconverted.""" @@ -1474,13 +1477,12 @@ def decode(cursor, fp=fingerprint): decoded = decode(valid) assert decoded is not None assert decoded.attack_result_id == anchor.attack_result_id - assert decoded.recency == anchor.metadata.get("updated_at") + assert decoded.timestamp == anchor.timestamp # A crafted cursor carrying a naive (tz-less) timestamp is rejected: service-minted anchors # are always timezone-aware, and a naive anchor would bind inconsistently in the seek. naive_payload = { "f": fingerprint, - "r": anchor.metadata.get("updated_at"), "t": "2026-01-01T00:00:00", "i": str(uuid.uuid4()), } @@ -1491,7 +1493,6 @@ def decode(cursor, fp=fingerprint): # tie-break matches the UTC-normalized timestamp column (service cursors are already UTC). offset_payload = { "f": fingerprint, - "r": anchor.metadata.get("updated_at"), "t": "2026-01-01T00:00:00+05:00", "i": str(uuid.uuid4()), } @@ -1505,7 +1506,6 @@ def decode(cursor, fp=fingerprint): # representable range when normalized to UTC decodes to None instead of raising. overflow_payload = { "f": fingerprint, - "r": anchor.metadata.get("updated_at"), "t": "0001-01-01T00:00:00+23:59", "i": str(uuid.uuid4()), } @@ -2099,7 +2099,7 @@ async def test_creates_conversation_and_adds_to_related(self, attack_service, mo call_kwargs = mock_memory.update_attack_result_by_id.call_args[1] assert call_kwargs["attack_result_id"] == "attack-1" assert result.conversation_id in call_kwargs["update_fields"]["pruned_conversation_ids"] - assert "updated_at" in call_kwargs["update_fields"]["attack_metadata"] + assert isinstance(call_kwargs["update_fields"]["timestamp"], datetime) async def test_rejects_source_conversation_from_different_attack(self, attack_service, mock_memory): """Should raise ValueError when source_conversation_id doesn't belong to the attack.""" diff --git a/tests/unit/backend/test_mappers.py b/tests/unit/backend/test_mappers.py index 667242c613..bbc091a1b0 100644 --- a/tests/unit/backend/test_mappers.py +++ b/tests/unit/backend/test_mappers.py @@ -416,6 +416,22 @@ async def test_created_at_metadata_still_wins_over_ar_timestamp(self) -> None: assert summary.created_at == metadata_ts + async def test_updated_at_uses_ar_timestamp_ignoring_metadata_updated_at(self) -> None: + """``updated_at`` is the persisted ``ar.timestamp``; a stale ``metadata['updated_at']`` is ignored.""" + ar_ts = datetime(2026, 4, 17, 12, 0, 0, tzinfo=timezone.utc) + stale = datetime(2020, 1, 1, 0, 0, 0, tzinfo=timezone.utc) + ar = AttackResult( + conversation_id="attack-1", + objective="test", + outcome=AttackOutcome.SUCCESS, + timestamp=ar_ts, + metadata={"created_at": stale.isoformat(), "updated_at": stale.isoformat()}, + ) + summary = await attack_result_to_summary_async(ar, stats=ConversationStats(message_count=0)) + + assert summary.updated_at == ar_ts + assert summary.created_at == stale + async def test_created_at_falls_back_to_now_when_both_absent(self) -> None: """When neither metadata nor ar.timestamp is set, fall back to datetime.now().""" ar = AttackResult( diff --git a/tests/unit/memory/memory_interface/test_interface_attack_results.py b/tests/unit/memory/memory_interface/test_interface_attack_results.py index 902b30f8ab..83b2e0c6a9 100644 --- a/tests/unit/memory/memory_interface/test_interface_attack_results.py +++ b/tests/unit/memory/memory_interface/test_interface_attack_results.py @@ -1713,16 +1713,16 @@ def test_get_attack_results_pagination_duplicate_and_tie_heavy_pages_disjoint_an assert len(set(paged)) == num_convs -def test_get_attack_results_pagination_order_parity_with_python_sort(sqlite_instance: MemoryInterface): - """Paginated SQL ordering matches the previous Python metadata sort for distinct metadata.""" - attack_results = [_make_attack_result(f"conv-{i}", ts_offset=i, updated_at_offset=(i * 7) % 50) for i in range(15)] +def test_get_attack_results_pagination_order_parity_with_timestamp_sort(sqlite_instance: MemoryInterface): + """Paginated SQL ordering matches a Python sort on the timestamp recency column.""" + attack_results = [_make_attack_result(f"conv-{i}", ts_offset=(i * 7) % 50) for i in range(15)] sqlite_instance.add_attack_results_to_memory(attack_results=attack_results) new_order = [r.conversation_id for r in sqlite_instance.get_attack_results(limit=100)] all_unpaged = list(sqlite_instance.get_attack_results()) expected = sorted( all_unpaged, - key=lambda ar: ar.metadata.get("updated_at", ar.metadata.get("created_at", "")), + key=lambda ar: (ar.timestamp, ar.attack_result_id), reverse=True, ) assert new_order == [r.conversation_id for r in expected] @@ -1799,7 +1799,7 @@ def test_get_attack_results_paginated_empty_metadata_orders_newest_first(sqlite_ def test_get_attack_results_pagination_with_ids_raises(sqlite_instance: MemoryInterface): """limit/keyset pagination cannot be combined with id-batched lookups.""" - anchor = AttackResultsKeysetCursor(recency="", timestamp=_BASE_TS, attack_result_id=str(uuid.uuid4())) + anchor = AttackResultsKeysetCursor(timestamp=_BASE_TS, attack_result_id=str(uuid.uuid4())) with pytest.raises(ValueError, match="pagination cannot be combined"): sqlite_instance.get_attack_results(attack_result_ids=[str(uuid.uuid4())], limit=10) with pytest.raises(ValueError, match="pagination cannot be combined"): @@ -1942,63 +1942,23 @@ def test_get_attack_results_keyset_paginates_recency_and_timestamp_ties(sqlite_i assert len(set(paged)) == 8 -def test_attack_result_recency_key_matches_sql_order(sqlite_instance: MemoryInterface): - """The Python recency anchor key reproduces the DB recency ordering (updated_at > created_at > '').""" +def test_attack_result_keyset_order_matches_sql_order(sqlite_instance: MemoryInterface): + """The keyset anchor ordering (timestamp, id) reproduces the DB recency ordering.""" seeded = [ - _make_attack_result("conv-upd", ts_offset=1, updated_at_offset=50, created_at_offset=1), - _make_attack_result("conv-created-only", ts_offset=2, created_at_offset=40), - _make_attack_result("conv-none", ts_offset=3), - _make_attack_result("conv-upd-high", ts_offset=4, updated_at_offset=90), + _make_attack_result("conv-a", ts_offset=1, created_at_offset=1), + _make_attack_result("conv-b", ts_offset=2, created_at_offset=40), + _make_attack_result("conv-c", ts_offset=3), + _make_attack_result("conv-d", ts_offset=4), ] sqlite_instance.add_attack_results_to_memory(attack_results=seeded) sql_order = list(sqlite_instance.get_attack_results(limit=100)) python_order = sorted( sqlite_instance.get_attack_results(), - key=lambda ar: (AttackResultsKeysetCursor.from_attack_result(ar).recency, ar.timestamp, ar.attack_result_id), + key=lambda ar: ( + AttackResultsKeysetCursor.from_attack_result(ar).timestamp, + ar.attack_result_id, + ), reverse=True, ) assert [r.conversation_id for r in sql_order] == [r.conversation_id for r in python_order] - - -@pytest.mark.parametrize("bad_recency", [False, 0, 1700000000]) -def test_get_attack_results_keyset_terminates_with_non_string_recency( - sqlite_instance: MemoryInterface, bad_recency: object -): - """Non-string recency scalars must not break keyset seek monotonicity. - - ``attack_metadata`` is ``dict[str, Any]``, so a caller can persist a non-string - ``updated_at`` (a JSON bool or number). SQLite's ``json_extract`` returns such a value as a - typed scalar and orders numbers before text, so the seek predicate ``recency < anchor`` - could never advance past the row — pagination would loop forever. The ``json_type`` string - guard collapses non-strings to ``""`` so the row sorts last and the keyset drains cleanly. - """ - string_rows = [_make_attack_result(f"conv-{i}", ts_offset=i, updated_at_offset=i) for i in range(4)] - bad_row = AttackResult( - conversation_id="conv-bad", - objective="Objective conv-bad", - outcome=AttackOutcome.UNDETERMINED, - metadata={"updated_at": bad_recency}, - timestamp=_BASE_TS, - ) - sqlite_instance.add_attack_results_to_memory(attack_results=[*string_rows, bad_row]) - - full = list(sqlite_instance.get_attack_results(limit=100)) - assert len(full) == 5 - - # Bounded manual drain: a regression re-introduces an infinite loop, so cap the page count - # and fail on non-termination rather than hanging the suite. - drained: list[str] = [] - after: AttackResultsKeysetCursor | None = None - for _ in range(len(full) + 2): - page = list(sqlite_instance.get_attack_results(limit=1, after=after)) - if not page: - break - drained.append(page[0].conversation_id) - after = _after(page) - else: - pytest.fail("keyset pagination did not terminate for non-string recency") - - assert drained == [r.conversation_id for r in full] - assert len(set(drained)) == len(drained) - assert drained[-1] == "conv-bad" # non-string recency collapses to "" and sorts last diff --git a/tests/unit/memory/test_azure_sql_memory.py b/tests/unit/memory/test_azure_sql_memory.py index 4d57d7f873..c2a1cfafee 100644 --- a/tests/unit/memory/test_azure_sql_memory.py +++ b/tests/unit/memory/test_azure_sql_memory.py @@ -290,26 +290,6 @@ def test_get_message_pieces_memory_label_conditions_uses_attack_result_labels( assert 'JSON_VALUE("PromptMemoryEntries".labels' not in compiled -def test_attack_results_recency_expr_guards_non_string_values(uninitialized_memory_interface: AzureSQLMemory): - """Azure recency sort honors only JSON *string* values, mirroring SQLite's json_type guard. - - A raw ``JSON_VALUE`` would stringify a non-string scalar (a JSON bool ``false`` -> ``"false"``), - diverging from ``_attack_result_recency_key``'s Python anchor (which collapses non-strings to - ``""``) and skipping rows at a page boundary. Each accessor is instead an ``OPENJSON`` subquery - filtered to ``type = 1`` (JSON string), so non-strings coalesce to ``""`` and both agree. - """ - from sqlalchemy.dialects import mssql - - expr = uninitialized_memory_interface._attack_results_recency_expr() - compiled = expr.compile(dialect=mssql.dialect(), compile_kwargs={"literal_binds": False}) - sql = str(compiled).lower() - # One OPENJSON accessor each for updated_at and created_at; the raw JSON_VALUE path is gone. - assert sql.count("openjson(") == 2 - assert "json_value" not in sql - # Both accessors guard on JSON string type (OPENJSON type == 1). - assert sorted(v for k, v in compiled.params.items() if k.startswith("type")) == [1, 1] - - def test_get_message_pieces_memory_label_conditions_bind_params(uninitialized_memory_interface: AzureSQLMemory): """Test that bind parameters are created for AttackResultEntry labels.""" conditions = uninitialized_memory_interface._get_message_pieces_memory_label_conditions( diff --git a/tests/unit/memory/test_migration.py b/tests/unit/memory/test_migration.py index 1fcfdb602f..d9844823ad 100644 --- a/tests/unit/memory/test_migration.py +++ b/tests/unit/memory/test_migration.py @@ -2241,6 +2241,144 @@ def test_attack_identifier_migration_backfills_graph_and_result_link(): engine.dispose() +# ============================================================================= +# Attack recency index + timestamp backfill migration (d7e9f1a3b5c6) +# ============================================================================= + + +_ATTACK_RECENCY_REV = "d7e9f1a3b5c6" +_ATTACK_RECENCY_PREV_REV = "3f6e8a0c2d4b" + + +def _seed_attack_result_with_metadata(connection, *, attack_id, conversation_id, timestamp, attack_metadata): + """Insert an AttackResultEntry row carrying legacy JSON recency metadata.""" + connection.execute( + text( + 'INSERT INTO "AttackResultEntries" ' + "(id, conversation_id, objective, executed_turns, execution_time_ms, outcome, " + "timestamp, attack_metadata) " + "VALUES (:id, :conv, 'obj', 1, 0, 'success', :timestamp, :metadata)" + ), + { + "id": attack_id, + "conv": conversation_id, + "timestamp": timestamp, + "metadata": json.dumps(attack_metadata) if attack_metadata is not None else None, + }, + ) + + +def test_attack_recency_migration_script_metadata(): + """The attack-recency migration declares the expected revision chain.""" + from pyrit.memory.alembic.versions import d7e9f1a3b5c6_index_attack_result_recency as mig + + assert mig.revision == _ATTACK_RECENCY_REV + assert mig.down_revision == _ATTACK_RECENCY_PREV_REV + assert mig.branch_labels is None + assert mig.depends_on is None + + +def test_attack_recency_upgrade_creates_indexes_and_backfills_timestamp(): + """Upgrading adds both indexes, backfills timestamp from the legacy JSON keys, and + strips the redundant ``updated_at`` key while preserving ``created_at``.""" + edited_id = str(uuid.uuid4()) + prog_id = str(uuid.uuid4()) + + with tempfile.TemporaryDirectory() as temp_dir: + engine = create_engine(f"sqlite:///{os.path.join(temp_dir, 'attack-recency.db')}") + try: + with engine.begin() as connection: + config = _config_for(connection) + command.upgrade(config, _ATTACK_RECENCY_PREV_REV) + + # A manually-edited GUI conversation: timestamp column is stale (creation time), + # recency lives in metadata.updated_at. + _seed_attack_result_with_metadata( + connection, + attack_id=edited_id, + conversation_id="conv-edited", + timestamp="2020-01-01 00:00:00", + attack_metadata={ + "created_at": "2020-01-01T00:00:00+00:00", + "updated_at": "2026-06-01T12:00:00+00:00", + }, + ) + # A programmatic attack: no JSON recency keys, timestamp column is authoritative. + _seed_attack_result_with_metadata( + connection, + attack_id=prog_id, + conversation_id="conv-prog", + timestamp="2023-03-03 08:00:00", + attack_metadata=None, + ) + + command.upgrade(config, _ATTACK_RECENCY_REV) + + index_names = {ix["name"] for ix in inspect(connection).get_indexes("AttackResultEntries")} + rows = dict( + connection.execute(text('SELECT id, attack_metadata FROM "AttackResultEntries"')).fetchall() + ) + timestamps = dict( + connection.execute(text('SELECT id, timestamp FROM "AttackResultEntries"')).fetchall() + ) + + assert "ix_AttackResultEntries_conversation_id" in index_names + assert "ix_AttackResultEntries_timestamp_id" in index_names + + edited_metadata = json.loads(rows[edited_id]) + assert "updated_at" not in edited_metadata + assert edited_metadata["created_at"] == "2020-01-01T00:00:00+00:00" + # Timestamp advanced to the former updated_at so the edited row keeps its recency. + assert str(timestamps[edited_id]).startswith("2026-06-01 12:00:00") + + # Programmatic row is untouched: no metadata, original timestamp preserved. + assert rows[prog_id] is None + assert str(timestamps[prog_id]).startswith("2023-03-03 08:00:00") + finally: + engine.dispose() + + +def test_attack_recency_downgrade_restores_updated_at_and_drops_indexes(): + """Downgrading drops both indexes and rewrites ``metadata.updated_at`` from the + timestamp column so the legacy JSON recency sort works again.""" + edited_id = str(uuid.uuid4()) + + with tempfile.TemporaryDirectory() as temp_dir: + engine = create_engine(f"sqlite:///{os.path.join(temp_dir, 'attack-recency-down.db')}") + try: + with engine.begin() as connection: + config = _config_for(connection) + command.upgrade(config, _ATTACK_RECENCY_REV) + + _seed_attack_result_with_metadata( + connection, + attack_id=edited_id, + conversation_id="conv-edited", + timestamp="2026-06-01 12:00:00", + attack_metadata={"created_at": "2026-06-01T00:00:00+00:00"}, + ) + + index_names_up = {ix["name"] for ix in inspect(connection).get_indexes("AttackResultEntries")} + assert "ix_AttackResultEntries_timestamp_id" in index_names_up + + command.downgrade(config, _ATTACK_RECENCY_PREV_REV) + + index_names_down = {ix["name"] for ix in inspect(connection).get_indexes("AttackResultEntries")} + restored = json.loads( + connection.execute( + text('SELECT attack_metadata FROM "AttackResultEntries" WHERE id = :id'), + {"id": edited_id}, + ).scalar_one() + ) + + assert "ix_AttackResultEntries_conversation_id" not in index_names_down + assert "ix_AttackResultEntries_timestamp_id" not in index_names_down + assert restored["created_at"] == "2026-06-01T00:00:00+00:00" + assert restored["updated_at"].startswith("2026-06-01T12:00:00") + finally: + engine.dispose() + + _STRING_TYPES_REQUIRING_LENGTH = {"String", "VARCHAR", "NVARCHAR", "Unicode"}