FIX: Fixing json retry

doc/code/framework.md

@@ -231,7 +231,7 @@ If you are contributing to PyRIT, that work will most likely land in one of the
 - Targets should use message_normalizer along with TargetConfiguration to transform `Messages` into formats that target supports.
 - Because targets are so varied, it is reasonable to return multiple tool calls, or none at all.
 - One attack can have many targets (and in fact, converters and scorers can also use targets to convert/score the prompt).
-- **Does not own**: what to send or what to do with the response. A target sends a prepared `Message` and returns a response — it doesn't convert prompts (converters), score (scorers), manage the conversation or decide the next turn (attacks), or apply attack logic. Its retries stay at the target layer (e.g. `RateLimitException`).
+- **Does not own**: what to send or what to do with the response. A target sends a prepared `Message` and returns a response — it doesn't convert prompts (converters), score (scorers), manage the conversation or decide the next turn (attacks), apply attack logic, or persist prompts and responses to memory (the `prompt_normalizer` owns that). Its retries stay at the target layer (e.g. `RateLimitException`).
 
 **Framework Plans**:
 
@@ -310,7 +310,7 @@ The below talks about responsibilities of most modules in the PyRIT library
 
 **Responsibility**: Reshape prompts and conversations so components and targets can interoperate. There are two distinct modules:
 
-- **`prompt_normalizer`** applies converters and dispatches individual prompts to a `PromptTarget` (handling batching and memory persistence). `NormalizerRequest` and `ConverterConfiguration` describe what to send and which converters to apply.
+- **`prompt_normalizer`** applies converters and dispatches individual prompts to a `PromptTarget` (handling batching and memory persistence). It is the single component that writes each request and response to memory; targets never persist on their own. `NormalizerRequest` and `ConverterConfiguration` describe what to send and which converters to apply.
 - **`message_normalizer`** reshapes multi-message conversation payloads into the structure a given model expects — for example, handling system-message behavior (keep / squash / ignore), history squashing, and tokenizer chat templates.
 
 ## [Output](./output/0_output)

pyrit/executor/attack/component/adversarial_conversation_manager.py

@@ -16,7 +16,6 @@
     ComponentRole,
     InvalidJsonException,
     execution_context,
-    pyrit_json_retry,
     remove_markdown_json,
 )
 from pyrit.executor.attack.core.attack_config import (
@@ -34,7 +33,7 @@
     SeedPrompt,
     get_common_json_schema,
 )
-from pyrit.prompt_normalizer import PromptNormalizer
+from pyrit.prompt_normalizer import PromptNormalizer, send_json_with_retry_async
 
 if TYPE_CHECKING:
     from pathlib import Path
@@ -413,7 +412,6 @@ def __init__(
         adversarial_first_user_message: SeedPrompt | None = None,
         adversarial_next_user_message: SeedPrompt | None = None,
         max_turns: int = 1,
-        raise_on_invalid_json: bool = True,
         prompt_normalizer: PromptNormalizer | None = None,
         conversation_id: str | None = None,
         objective: str | None = None,
@@ -439,9 +437,6 @@ def __init__(
                 text directly.
             max_turns: Maximum number of turns; rendered into the adversarial system prompt as
                 ``max_turns``. Defaults to 1.
-            raise_on_invalid_json: When True (default), a reply that fails to match the resolved
-                schema raises ``InvalidJsonException`` (retried via ``pyrit_json_retry``). When False,
-                the raw reply text is returned as ``next_message`` instead of raising.
             prompt_normalizer: The prompt normalizer to send through. Defaults to a new one.
             conversation_id: The adversarial-chat conversation id this manager drives. A fresh
                 id is generated when None.
@@ -468,7 +463,6 @@ def __init__(
         self._adversarial_first_user_message = adversarial_first_user_message
         self._adversarial_next_user_message = adversarial_next_user_message
         self._max_turns = max_turns
-        self._raise_on_invalid_json = raise_on_invalid_json
         self._prompt_normalizer = prompt_normalizer or PromptNormalizer()
         self._conversation_id = conversation_id or str(uuid4())
         self._objective = objective
@@ -757,7 +751,7 @@ async def generate_adversarial_reply_async(
 
         Raises:
             ValueError: If no response is received from the adversarial chat.
-            InvalidJsonException: If ``raise_on_invalid_json`` is True and the reply is invalid.
+            InvalidJsonException: If the reply is not valid JSON after the retry budget is exhausted.
         """
         return await self._send_and_parse_async(
             prompt_text=prompt_text,
@@ -803,7 +797,6 @@ def _build_objective_message(
             )
         return Message.from_prompt(prompt=reply.next_message, role="user")
 
-    @pyrit_json_retry
     async def _send_and_parse_async(
         self,
         *,
@@ -814,10 +807,10 @@ async def _send_and_parse_async(
         """
         Send one user turn to the adversarial chat and parse its reply.
 
-        This is the single place adversarial-chat JSON retry lives: when the reply fails to match the
-        resolved schema, ``InvalidJsonException`` propagates and ``pyrit_json_retry`` re-sends the turn
-        until it parses or the attempt budget is exhausted. When ``raise_on_invalid_json`` is False, an
-        unparseable reply is returned as raw text instead.
+        This is the single place adversarial-chat JSON retry lives. It delegates to
+        ``send_json_with_retry_async`` so each retry sends on a clean conversation history:
+        the failed turn is rolled back out of memory before the turn is resent, instead of
+        replaying the model's own malformed reply.
 
         When a ``modality_router`` is configured, the outgoing message is built via
         ``build_adversarial_input_message`` so first-turn seed media (``seed_message``) and prior
@@ -834,7 +827,7 @@ async def _send_and_parse_async(
 
         Raises:
             ValueError: If no response is received from the adversarial chat.
-            InvalidJsonException: If ``raise_on_invalid_json`` is True and the reply is invalid.
+            InvalidJsonException: If the reply is not valid JSON after the retry budget is exhausted.
         """
         prompt_metadata = _build_adversarial_prompt_metadata(response_json_schema=self._response_json_schema)
 
@@ -852,32 +845,26 @@ async def _send_and_parse_async(
                 prompt_metadata=prompt_metadata or None,
             )
 
+        if self._memory_labels:
+            for piece in message.message_pieces:
+                piece.labels = self._memory_labels
+
+        schema = self._response_json_schema
+
+        def _parse(response: Message) -> AdversarialReply:
+            return _parse_adversarial_reply(response.get_value(), schema=schema)
+
         with execution_context(
             component_role=ComponentRole.ADVERSARIAL_CHAT,
             attack_strategy_name=self._attack_strategy_name,
             component_identifier=self._adversarial_target.get_identifier(),
             objective_target_conversation_id=self._objective_target_conversation_id,
             objective=self._objective,
         ):
-            if self._memory_labels:
-                for piece in message.message_pieces:
-                    piece.labels = self._memory_labels
-            response = await self._prompt_normalizer.send_prompt_async(
+            return await send_json_with_retry_async(
+                normalizer=self._prompt_normalizer,
+                target=self._adversarial_target,
                 message=message,
                 conversation_id=self._conversation_id,
-                target=self._adversarial_target,
+                parse=_parse,
             )
-
-        if not response:
-            raise ValueError("No response received from adversarial chat")
-
-        raw = response.get_value()
-
-        schema = self._response_json_schema
-        if not self._raise_on_invalid_json:
-            try:
-                return _parse_adversarial_reply(raw, schema=schema)
-            except InvalidJsonException:
-                return AdversarialReply(next_message=raw, raw=raw)
-
-        return _parse_adversarial_reply(raw, schema=schema)

pyrit/memory/alembic/versions/a1c3e5d7f9b0_add_retries_to_conversations.py

@@ -0,0 +1,35 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+"""
+add retries to Conversations.
+
+Revision ID: a1c3e5d7f9b0
+Revises: e5f7a9c1b3d2
+Create Date: 2026-01-01 00:00:00.000000
+"""
+
+from collections.abc import Sequence
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "a1c3e5d7f9b0"
+down_revision: str | None = "e5f7a9c1b3d2"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def upgrade() -> None:
+    """Apply this schema upgrade."""
+    # ### commands auto generated by Alembic and reviewed by author ###
+    op.add_column("Conversations", sa.Column("retries", sa.JSON(), nullable=True))
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Revert this schema upgrade."""
+    # ### commands auto generated by Alembic and reviewed by author ###
+    op.drop_column("Conversations", "retries")
+    # ### end Alembic commands ###

pyrit/memory/memory_interface.py

@@ -15,7 +15,7 @@
 from sqlalchemy import MetaData, and_, not_, or_, select
 from sqlalchemy.engine.base import Engine
 from sqlalchemy.exc import IntegrityError, SQLAlchemyError
-from sqlalchemy.orm.attributes import InstrumentedAttribute
+from sqlalchemy.orm.attributes import InstrumentedAttribute, flag_modified
 
 if TYPE_CHECKING:
     from pyrit.memory.memory_embedding import MemoryEmbedding
@@ -53,6 +53,8 @@
     AttackTechniqueIdentifier,
     ComponentIdentifier,
     Conversation,
+    ConversationRetry,
+    ConversationRetryReason,
     ConversationStats,
     ConverterIdentifier,
     IdentifierFilter,
@@ -522,6 +524,82 @@ def _insert_conversation(self, *, conversation: Conversation) -> None:
                 logger.exception(f"Error registering conversation {conversation.conversation_id}: {e}")
                 raise
 
+    def add_conversation_retry(self, *, conversation_id: str, sequence: int, reason: ConversationRetryReason) -> None:
+        """
+        Append a retry record to the conversation-scoped metadata for ``conversation_id``.
+
+        Records that a turn had to be retried (e.g. because its response failed JSON
+        validation and was rolled back out of memory). The conversation's ``Conversations``
+        row is updated in place; if no row exists yet it is created. This is distinct from
+        the insert-only ``_insert_conversation``.
+
+        Args:
+            conversation_id (str): The conversation whose turn was retried.
+            sequence (int): The sequence the retried turn's request occupies.
+            reason (ConversationRetryReason): Why the turn was retried.
+
+        Raises:
+            SQLAlchemyError: If the database update fails; the transaction is rolled back first.
+        """
+        record = ConversationRetry(sequence=sequence, reason=reason).model_dump(mode="json")
+        with closing(self.get_session()) as session:
+            try:
+                entry = session.get(ConversationEntry, str(conversation_id))
+                if entry is None:
+                    entry = ConversationEntry(conversation=Conversation(conversation_id=str(conversation_id)))
+                    session.add(entry)
+                entry.retries = [*(entry.retries or []), record]
+                flag_modified(entry, "retries")
+                session.commit()
+            except SQLAlchemyError as e:
+                session.rollback()
+                logger.exception(f"Error recording retry for conversation {conversation_id}: {e}")
+                raise
+
+    def delete_conversation_pieces_after_sequence(self, *, conversation_id: str, sequence: int) -> int:
+        """
+        Delete all message pieces in a conversation whose sequence is greater than ``sequence``.
+
+        Rolls a conversation back to a baseline so a failed turn can be resent on a clean
+        history. Dependent ``EmbeddingData`` rows for the deleted pieces are removed first to
+        avoid orphaned foreign keys. Pieces at or below ``sequence`` (e.g. the system prompt
+        and any prior good turns) are left intact.
+
+        Args:
+            conversation_id (str): The conversation to roll back.
+            sequence (int): The baseline sequence; pieces with a greater sequence are deleted.
+
+        Returns:
+            int: The number of ``PromptMemoryEntries`` deleted.
+
+        Raises:
+            SQLAlchemyError: If the deletion fails; the transaction is rolled back first.
+        """
+        with closing(self.get_session()) as session:
+            try:
+                pieces = (
+                    session.query(PromptMemoryEntry)
+                    .filter(
+                        PromptMemoryEntry.conversation_id == str(conversation_id),
+                        PromptMemoryEntry.sequence > sequence,
+                    )
+                    .all()
+                )
+                if not pieces:
+                    return 0
+                piece_ids = [piece.id for piece in pieces]
+                session.query(EmbeddingDataEntry).filter(EmbeddingDataEntry.id.in_(piece_ids)).delete(
+                    synchronize_session=False
+                )
+                for piece in pieces:
+                    session.delete(piece)
+                session.commit()
+                return len(pieces)
+            except SQLAlchemyError as e:
+                session.rollback()
+                logger.exception(f"Error deleting conversation pieces for {conversation_id}: {e}")
+                raise
+
     @classmethod
     def _persist_target_identifier(cls, *, session: Any, target_identifier: TargetIdentifier) -> None:
         """

pyrit/memory/memory_models.py

@@ -48,6 +48,7 @@
     ComponentIdentifier,
     Conversation,
     ConversationReference,
+    ConversationRetry,
     ConversationType,
     ConverterIdentifier,
     EvaluationIdentifier,
@@ -1017,6 +1018,10 @@ class ConversationEntry(Base):
         String(64), ForeignKey(f"{TargetIdentifierEntry.__tablename__}.hash"), nullable=True
     )
 
+    # JSON-serialized list of ConversationRetry records (turns that were retried in
+    # this conversation). Nullable for backwards compatibility with existing databases.
+    retries: Mapped[list[dict[str, Any]] | None] = mapped_column(JSON, nullable=True)
+
     # Version of PyRIT used when this entry was created. Nullable for backwards
     # compatibility with existing databases.
     pyrit_version = mapped_column(String, nullable=True)
@@ -1031,6 +1036,7 @@ def __init__(self, *, conversation: Conversation) -> None:
         self.conversation_id = conversation.conversation_id
         self.target_identifier = conversation.target_identifier.model_dump() if conversation.target_identifier else None
         self.target_identifier_hash = conversation.target_identifier.hash if conversation.target_identifier else None
+        self.retries = [retry.model_dump(mode="json") for retry in conversation.retries] or None
         self.pyrit_version = pyrit.__version__
 
     def get_conversation(self) -> Conversation:
@@ -1042,7 +1048,12 @@ def get_conversation(self) -> Conversation:
         """
         stored_version = self.pyrit_version or LEGACY_PYRIT_VERSION
         target_id = _load_identifier(self.target_identifier, pyrit_version=stored_version)
-        return Conversation(conversation_id=self.conversation_id, target_identifier=target_id)
+        retries = [ConversationRetry.model_validate(retry) for retry in self.retries or []]
+        return Conversation(
+            conversation_id=self.conversation_id,
+            target_identifier=target_id,
+            retries=retries,
+        )
 
 
 class EmbeddingDataEntry(Base):

pyrit/models/__init__.py

@@ -75,6 +75,7 @@
     ToolCall,
 )
 from pyrit.models.messages.conversation_reference import ConversationReference, ConversationType
+from pyrit.models.messages.conversation_retry import ConversationRetry, ConversationRetryReason
 from pyrit.models.parameter import (
     ComponentType,
     Parameter,
@@ -141,6 +142,8 @@
     "ConverterIdentifier",
     "Conversation",
     "ConversationReference",
+    "ConversationRetry",
+    "ConversationRetryReason",
     "ConversationStats",
     "ConversationType",
     "construct_response_from_request",

pyrit/models/messages/conversation_retry.py

@@ -0,0 +1,34 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+from __future__ import annotations
+
+from enum import Enum
+
+from pydantic import BaseModel, ConfigDict
+
+
+class ConversationRetryReason(str, Enum):
+    """Why a turn in a conversation had to be retried."""
+
+    JSON_PARSING = "json_parsing"
+
+
+class ConversationRetry(BaseModel):
+    """
+    Record of a single retried turn within a conversation.
+
+    A retry happens when a turn's response failed validation (e.g. malformed JSON)
+    and the failed turn was rolled back out of memory so the turn could be resent on
+    a clean history. The record is conversation-scoped metadata: it captures which
+    turn was retried and why, without keeping the discarded turn's pieces around.
+    """
+
+    model_config = ConfigDict(frozen=True)
+
+    # The sequence the retried turn's request occupies. Stable across attempts
+    # because the rollback resets the sequence so the eventual successful request
+    # lands at the same index.
+    sequence: int
+
+    reason: ConversationRetryReason