diff --git a/doc/code/executor/3_attack_configuration.ipynb b/doc/code/executor/3_attack_configuration.ipynb index b5d1df1477..0602296fdf 100644 --- a/doc/code/executor/3_attack_configuration.ipynb +++ b/doc/code/executor/3_attack_configuration.ipynb @@ -462,7 +462,7 @@ "Beyond the call arguments, attacks are tuned at construction time with three configuration objects:\n", "\n", "- **`AttackConverterConfig`** — request/response [converters](../converters/0_converters.ipynb)\n", - " applied to every prompt and response.\n", + " applied to live attack prompts and responses, plus selected roles in prepended history.\n", "- **`AttackScoringConfig`** — the objective scorer plus any auxiliary\n", " [scorers](../scoring/0_scoring.ipynb).\n", "- **`AttackAdversarialConfig`** — the adversarial target (a model PyRIT controls) that multi-turn\n", @@ -470,7 +470,27 @@ "\n", "Converter and scoring configs apply to single- and multi-turn attacks alike; the adversarial config\n", "only applies to attacks that drive a conversation. Below builds a converter config — it's just a\n", - "plain object you hand to the attack constructor." + "plain object you hand to the attack constructor.\n", + "\n", + "Request converters apply to prepended `user` messages by default. Prepended `assistant` messages\n", + "represent simulated target output, so PyRIT leaves them unchanged unless the attack explicitly\n", + "opts in. For example:\n", + "\n", + "```python\n", + "from pyrit.executor.attack import PrependedConversationConfig\n", + "\n", + "attack = PromptSendingAttack(\n", + " objective_target=target,\n", + " attack_converter_config=converter_config,\n", + " prepended_conversation_config=PrependedConversationConfig(\n", + " apply_converters_to_roles=[\"user\", \"assistant\"],\n", + " ),\n", + ")\n", + "```\n", + "\n", + "PyRIT applies these role-specific conversions while the prepended messages are still structured.\n", + "If the target cannot accept editable history, target normalization then formats the converted and\n", + "unconverted history with the first live request without broadening the selected converter scope." ] }, { diff --git a/doc/code/executor/3_attack_configuration.py b/doc/code/executor/3_attack_configuration.py index d7597a2e27..74e2cd6bdd 100644 --- a/doc/code/executor/3_attack_configuration.py +++ b/doc/code/executor/3_attack_configuration.py @@ -166,7 +166,7 @@ # Beyond the call arguments, attacks are tuned at construction time with three configuration objects: # # - **`AttackConverterConfig`** — request/response [converters](../converters/0_converters.ipynb) -# applied to every prompt and response. +# applied to live attack prompts and responses, plus selected roles in prepended history. # - **`AttackScoringConfig`** — the objective scorer plus any auxiliary # [scorers](../scoring/0_scoring.ipynb). # - **`AttackAdversarialConfig`** — the adversarial target (a model PyRIT controls) that multi-turn @@ -175,6 +175,26 @@ # Converter and scoring configs apply to single- and multi-turn attacks alike; the adversarial config # only applies to attacks that drive a conversation. Below builds a converter config — it's just a # plain object you hand to the attack constructor. +# +# Request converters apply to prepended `user` messages by default. Prepended `assistant` messages +# represent simulated target output, so PyRIT leaves them unchanged unless the attack explicitly +# opts in. For example: +# +# ```python +# from pyrit.executor.attack import PrependedConversationConfig +# +# attack = PromptSendingAttack( +# objective_target=target, +# attack_converter_config=converter_config, +# prepended_conversation_config=PrependedConversationConfig( +# apply_converters_to_roles=["user", "assistant"], +# ), +# ) +# ``` +# +# PyRIT applies these role-specific conversions while the prepended messages are still structured. +# If the target cannot accept editable history, target normalization then formats the converted and +# unconverted history with the first live request without broadening the selected converter scope. # %% from pyrit.converter import Base64Converter diff --git a/doc/code/framework.md b/doc/code/framework.md index 278bf9c466..b75ecf231b 100644 --- a/doc/code/framework.md +++ b/doc/code/framework.md @@ -233,6 +233,9 @@ If you are contributing to PyRIT, that work will most likely land in one of the - This is often an LLM, but it doesn't have to be. For Cross-Domain Prompt Injection Attacks, the target might be a storage account that a later target has a reference to. Message and conversation should be generic enough to handle this extra data. - Target capabilities should be used to see if a target is compatible with the capabilities that the other components want to use. - Targets should use message_normalizer along with TargetConfiguration to transform `Messages` into formats that target supports. +- A target may observe an internal, caller-owned send context at the provider-invocation boundary, + after target-side waits and immediately before irreversible provider I/O, but the caller owns any + bootstrap-history identity, replay, or branching state. - 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), 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`). @@ -314,10 +317,14 @@ The below talks about responsibilities of most modules in the PyRIT library ## [Normalizers](./targets/11_message_normalizer) -**Responsibility**: Reshape prompts and conversations so components and targets can interoperate. There are two distinct modules: +**Responsibility**: Reshape prompts and conversations so components and targets can interoperate. -- **`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. +- **`prompt_normalizer`** applies converters, persists requests and responses, and dispatches prompts to a `PromptTarget`. Targets do not persist messages. +- **`message_normalizer`** reshapes conversations into target-compatible payloads. It owns target-facing representation, not attack policy or conversation state. +- Prepended-history identity and delivery state belong to the attack execution context. Targets and normalizers consume only the narrow send-time view needed for provider adaptation. +- **Does not own**: the conversation of record. Memory is canonical; a normalized payload is an ephemeral target-facing view that is never written back. + +See [message normalizers](./targets/11_message_normalizer) for capability behavior, processing order, and prepended-history lifecycle details. ## [Output](./output/0_output) diff --git a/doc/code/scenarios/0_scenarios.ipynb b/doc/code/scenarios/0_scenarios.ipynb index c96b9b1abd..9ed5f78114 100644 --- a/doc/code/scenarios/0_scenarios.ipynb +++ b/doc/code/scenarios/0_scenarios.ipynb @@ -83,7 +83,7 @@ " - `max_retries`: Number of retry attempts on failure (default: 0)\n", " - `memory_labels`: Optional labels for tracking (optional)\n", " - `include_baseline`: Whether to prepend a baseline attack (defaults to the scenario type's\n", - " `BASELINE_ATTACK_POLICY`; most scenarios default it on, `Jailbreak` defaults it off)\n", + " `BASELINE_ATTACK_POLICY`; most scenarios, including `Jailbreak`, default it on)\n", "\n", "### Example Structure\n", "\n", @@ -241,17 +241,42 @@ " airt.jailbreak\u001b[0m\n", " Class: Jailbreak\n", " Description:\n", - " Jailbreak scenario implementation for PyRIT. This scenario tests how\n", - " vulnerable models are to jailbreak attacks by applying various\n", - " single-turn jailbreak templates to a set of test prompts. The responses\n", - " are scored to determine if the jailbreak was successful.\n", + " Jailbreak scenario implementation for PyRIT. Tests how vulnerable a\n", + " model is to jailbreak templates. A run is the cross-product of three\n", + " selectors: - **dataset** — the harmful objectives (HarmBench). -\n", + " **techniques** — two delivery methods for each jailbreak:\n", + " ``prompt_sending`` (the template rendered inline into the user message)\n", + " and ``jailbreak_system_prompt`` (the template set as the system prompt\n", + " with the objective sent as the user turn). - **jailbreaks** — which\n", + " jailbreak templates to run (a random ``num_jailbreaks`` sample or an\n", + " explicit ``jailbreak_names`` set). ``prompt_sending`` applies each\n", + " template as a ``TextJailbreakConverter`` on the outgoing request, so the\n", + " objective is rendered inline into the template's ``{{prompt}}`` slot.\n", + " ``jailbreak_system_prompt`` instead sets the template as a native system\n", + " prompt and sends the objective as its own user turn, so it is only built\n", + " for targets that natively support editable history and system prompts\n", + " (it is skipped for incapable targets, or raises if it is the only\n", + " selected technique). Responses are scored to determine whether the\n", + " jailbreak succeeded (non-refusal).\n", " Aggregate Techniques:\n", - " - all, simple, complex\n", - " Available Techniques (4):\n", - " prompt_sending, many_shot, skeleton, role_play\n", - " Default Technique: simple\n", - " Default Datasets (1, max 4 per dataset):\n", - " airt_harms\n", + " - all, default, single_turn\n", + " Available Techniques (2):\n", + " prompt_sending, jailbreak_system_prompt\n", + " Default Technique: default\n", + " Default Datasets (1):\n", + " harmbench\n", + " Supported Parameters:\n", + " - objective_target (any): Target system under attack: a registered target name or a PromptTarget instance.\n", + " - scenario_techniques (any): Techniques to execute; defaults to the scenario's default aggregate when omitted.\n", + " - technique_converters (any): Mapping of concrete technique name to extra request converters to append.\n", + " - dataset_config (any): Dataset source configuration; defaults to the scenario's default when omitted.\n", + " - memory_labels (any): Additional labels applied to every attack run in the scenario.\n", + " - max_concurrency (int) [default: 4]: Maximum number of concurrent units of work for the scenario.\n", + " - max_retries (int) [default: 0]: Maximum number of automatic retries if the scenario raises an exception.\n", + " - include_baseline (bool): Whether to prepend a baseline atomic attack; None defers to BASELINE_ATTACK_POLICY.\n", + " - num_jailbreaks (int): Draw this many random jailbreak templates for the run. Mutually exclusive with jailbreak_names.\n", + " - num_jailbreak_attempts (int) [default: 1]: Number of times to try each (technique x jailbreak template x objective).\n", + " - jailbreak_names (list[str]): Explicit jailbreak template file names to run (e.g. aim.yaml dan_11.yaml). When omitted, a random sample is drawn. Mutually exclusive with num_jailbreaks.\n", "\u001b[1m\u001b[36m\n", " airt.leakage\u001b[0m\n", " Class: Leakage\n", @@ -439,8 +464,8 @@ "each objective directly to the target without any converters or multi-turn techniques. This is\n", "controlled by the `include_baseline` scenario parameter, supplied through the CLI, config, or\n", "`set_params_from_args` before `initialize_async`; when omitted, each scenario falls back to its\n", - "own `BASELINE_ATTACK_POLICY` class attribute (most scenarios default it on; `Jailbreak` defaults\n", - "it off). See\n", + "own `BASELINE_ATTACK_POLICY` class attribute (most scenarios, including `Jailbreak`, default it\n", + "on). See\n", "[Common Scenario Parameters](./1_common_scenario_parameters.ipynb) for a worked example.\n", "\n", "Custom scenarios should choose their `BASELINE_ATTACK_POLICY` based on whether an unmodified\n", @@ -449,8 +474,7 @@ "- **`Enabled`** — the baseline is prepended by default and the caller can opt out. Use when an\n", " unmodified-prompt run is a meaningful comparison point (most scenarios).\n", "- **`Disabled`** — the baseline is supported but omitted by default; the caller must opt in. Use\n", - " when the scenario is already dominated by a large set of templates/techniques that already\n", - " exercise the unmodified surface (e.g., `Jailbreak`).\n", + " when an unmodified-prompt comparison is valid but not useful enough to run by default.\n", "- **`Forbidden`** — the baseline is unavailable and passing `include_baseline=True` raises. Use\n", " when the scenario's semantics make a single-shot unmodified prompt meaningless as a comparator\n", " (e.g., benchmarks comparing across adversarial models, or multi-turn-only scenarios)." diff --git a/doc/code/scenarios/0_scenarios.py b/doc/code/scenarios/0_scenarios.py index 459fb6754f..aa01af4930 100644 --- a/doc/code/scenarios/0_scenarios.py +++ b/doc/code/scenarios/0_scenarios.py @@ -85,7 +85,7 @@ # - `max_retries`: Number of retry attempts on failure (default: 0) # - `memory_labels`: Optional labels for tracking (optional) # - `include_baseline`: Whether to prepend a baseline attack (defaults to the scenario type's -# `BASELINE_ATTACK_POLICY`; most scenarios default it on, `Jailbreak` defaults it off) +# `BASELINE_ATTACK_POLICY`; most scenarios, including `Jailbreak`, default it on) # # ### Example Structure # @@ -180,8 +180,8 @@ async def _build_atomic_attacks_async(self, *, context): # each objective directly to the target without any converters or multi-turn techniques. This is # controlled by the `include_baseline` scenario parameter, supplied through the CLI, config, or # `set_params_from_args` before `initialize_async`; when omitted, each scenario falls back to its -# own `BASELINE_ATTACK_POLICY` class attribute (most scenarios default it on; `Jailbreak` defaults -# it off). See +# own `BASELINE_ATTACK_POLICY` class attribute (most scenarios, including `Jailbreak`, default it +# on). See # [Common Scenario Parameters](./1_common_scenario_parameters.ipynb) for a worked example. # # Custom scenarios should choose their `BASELINE_ATTACK_POLICY` based on whether an unmodified @@ -190,8 +190,7 @@ async def _build_atomic_attacks_async(self, *, context): # - **`Enabled`** — the baseline is prepended by default and the caller can opt out. Use when an # unmodified-prompt run is a meaningful comparison point (most scenarios). # - **`Disabled`** — the baseline is supported but omitted by default; the caller must opt in. Use -# when the scenario is already dominated by a large set of templates/techniques that already -# exercise the unmodified surface (e.g., `Jailbreak`). +# when an unmodified-prompt comparison is valid but not useful enough to run by default. # - **`Forbidden`** — the baseline is unavailable and passing `include_baseline=True` raises. Use # when the scenario's semantics make a single-shot unmodified prompt meaningless as a comparator # (e.g., benchmarks comparing across adversarial models, or multi-turn-only scenarios). diff --git a/doc/code/targets/0_prompt_targets.md b/doc/code/targets/0_prompt_targets.md index 3a7943e45d..1da66a6bfd 100644 --- a/doc/code/targets/0_prompt_targets.md +++ b/doc/code/targets/0_prompt_targets.md @@ -9,15 +9,49 @@ Prompt Targets are endpoints for where to send prompts. For example, a target co Prompt targets are found [here](https://github.com/microsoft/PyRIT/tree/main/pyrit/prompt_target/) in code. -## Send_Prompt_Async +## `send_prompt_async` -The main entry method follow the following signature: +The main entry method has the following signature: -``` -async def send_prompt_async(self, *, message: Message) -> Message: +```python +async def send_prompt_async( + self, + *, + message: Message, + send_context: TargetSendContext | None = None, +) -> list[Message]: ``` -A `Message` object is a normalized object with all the information a target will need to send a prompt, including a way to get a history for that prompt (in the cases that also needs to be sent). This is discussed in more depth [here](../memory/3_memory_data_types.md). +A `Message` object contains the current request and the identifiers needed to load its conversation +history. This is discussed in more depth [here](../memory/3_memory_data_types.md). + +`send_context` is an internal protocol that lets caller-owned execution state select persisted history +and observe the provider-attempt boundary. Attacks with prepended history own the concrete +`PrependedHistorySendContext`; targets do not construct it, clone it, or decide whether its seed should +be replayed. Before a provider send, `PromptTarget` loads memory history, asks the protocol for the +caller-approved target view, and then runs the target's capability-normalization pipeline. Request +converters have already run by this point, so role-specific converter choices remain intact even when +the target must receive one flattened request. The context is ephemeral and does not replace the +structured messages stored in memory. + +For a stateful target without editable history, the initial bootstrap is flattened once. Later sends +retain replayable memory history in the normalized view so a target such as `WebsocketTarget` can +restore a replaced provider session, while the existing provider session still receives only the +current request. A stateful TAP clone bootstraps its new provider session with the complete replayable +duplicated branch, even when the attack started without an explicit prepended seed. Stateless targets +continue to receive only the explicit prepended seed plus each current request, never prior live branch +turns. A stateful TAP target without editable history can flatten only text converter output when +branching; non-text converter output requires an editable-history target, a stateless target, or +`branching_factor=1` so copied media is never replayed under a different role. + +The provider-attempt signal is emitted after shared target-side rate limiting. Targets that need more +precise setup, such as WebSocket, Playwright, or conversation-keyed HTTP targets, emit it immediately +before the irreversible provider operation. Cancellation before that point leaves a one-time +bootstrap available for retry. + +`send_prompt_async` is the final public orchestration method. Custom target subclasses implement +`_send_prompt_to_target_async(*, normalized_conversation: list[Message]) -> list[Message]` instead of +overriding `send_prompt_async`. ## Chat-style targets vs general targets diff --git a/doc/code/targets/11_message_normalizer.ipynb b/doc/code/targets/11_message_normalizer.ipynb index 1fd52230fa..43cb028149 100644 --- a/doc/code/targets/11_message_normalizer.ipynb +++ b/doc/code/targets/11_message_normalizer.ipynb @@ -16,6 +16,35 @@ "\n", "The `MessageNormalizer` classes handle these conversions, making it easy to work with any target regardless of its expected input format.\n", "\n", + "## Memory is canonical, the normalized payload is not\n", + "\n", + "A normalizer builds a **target-facing view at send time**. It is never written back to memory.\n", + "\n", + "This matters most for a target that cannot accept editable history. If you prepend eight\n", + "structured turns to a conversation, memory keeps eight structured turns, and the UI, scorers,\n", + "resume, and exports all see eight turns. `HistorySquashNormalizer` flattens those turns into\n", + "a single prompt only for the wire, then discards the flattened copy. The two representations\n", + "are expected to differ.\n", + "\n", + "Flattening the conversation in memory instead would break scoring, resume, and evaluation for\n", + "the sake of one target's wire format. If a prepended piece is not text (an image, for example),\n", + "the flattened view holds a text placeholder and the normalizer logs a warning, because the\n", + "target receives a description instead of the media.\n", + "\n", + "For prepended history on a target without editable history, PyRIT:\n", + "\n", + "1. Converts and persists the structured prepended messages.\n", + "2. Converts the live request.\n", + "3. Applies the per-send `EDITABLE_HISTORY` normalizer selected by\n", + " `PrependedConversationConfig`.\n", + "4. Applies the target's remaining capability normalizers.\n", + "5. Serializes the normalized view and invokes the provider.\n", + "\n", + "The attack-owned `PrependedHistorySendContext` records the persisted prepended-message boundary.\n", + "Stateful targets consume it after the first provider attempt; stateless targets reuse it for each\n", + "current request. Targets interact with that state only through an internal `TargetSendContext`\n", + "protocol at the send boundary.\n", + "\n", "## Base Classes\n", "\n", "There are two base normalizer types:\n", @@ -25,10 +54,40 @@ "Some normalizers implement both interfaces." ] }, + { + "cell_type": "markdown", + "id": "1", + "metadata": {}, + "source": [ + "## Two History-Squashing Scopes\n", + "\n", + "PyRIT uses `HistorySquashNormalizer` in two places that adapt different target capabilities:\n", + "\n", + "- **Multi-turn support** answers whether the target can continue a conversation across sends.\n", + "- **Editable-history support** answers whether PyRIT can supply or rewrite earlier turns before\n", + " sending the current message.\n", + "\n", + "| Target capabilities | Prepended-history behavior |\n", + "|---|---|\n", + "| Multi-turn with editable history | Send the structured history directly; no history squashing is needed. |\n", + "| Multi-turn without editable history | A context-scoped `HistorySquashNormalizer` encodes the prepended history and first live message into the initial request. Later turns use the target's own conversation state. |\n", + "| Single-turn without editable history | The context-scoped `HistorySquashNormalizer` first produces one message. The target's ordinary use of the same normalizer then sees one message and does nothing. |\n", + "\n", + "The first-turn distinction comes from the attack-owned prepended-history send context, not from a\n", + "separate normalizer implementation. The context applies its configured `HistorySquashNormalizer`\n", + "once to bootstrap prepended history for a target that cannot accept caller-supplied prior turns.\n", + "Its key use case is a multi-turn, server-managed target without editable history.\n", + "\n", + "The target's ordinary capability pipeline independently uses `HistorySquashNormalizer` for a target\n", + "without multi-turn support. Both scopes preserve original-versus-converted text views and keep\n", + "non-text pieces from the current request separate. The context-scoped use can supply a custom\n", + "formatter; the ordinary use defaults to `[Conversation History]` and `[Current Message]` sections." + ] + }, { "cell_type": "code", "execution_count": null, - "id": "1", + "id": "2", "metadata": {}, "outputs": [ { @@ -61,7 +120,7 @@ }, { "cell_type": "markdown", - "id": "2", + "id": "3", "metadata": {}, "source": [ "## ChatMessageNormalizer\n", @@ -77,7 +136,7 @@ { "cell_type": "code", "execution_count": null, - "id": "3", + "id": "4", "metadata": {}, "outputs": [ { @@ -107,7 +166,7 @@ { "cell_type": "code", "execution_count": null, - "id": "4", + "id": "5", "metadata": {}, "outputs": [ { @@ -135,7 +194,7 @@ { "cell_type": "code", "execution_count": null, - "id": "5", + "id": "6", "metadata": {}, "outputs": [ { @@ -173,7 +232,7 @@ }, { "cell_type": "markdown", - "id": "6", + "id": "7", "metadata": {}, "source": [ "## GenericSystemSquashNormalizer\n", @@ -201,7 +260,7 @@ { "cell_type": "code", "execution_count": null, - "id": "7", + "id": "8", "metadata": {}, "outputs": [ { @@ -236,7 +295,7 @@ }, { "cell_type": "markdown", - "id": "8", + "id": "9", "metadata": {}, "source": [ "## ConversationContextNormalizer\n", @@ -261,7 +320,7 @@ { "cell_type": "code", "execution_count": null, - "id": "9", + "id": "10", "metadata": {}, "outputs": [ { @@ -289,7 +348,7 @@ }, { "cell_type": "markdown", - "id": "10", + "id": "11", "metadata": {}, "source": [ "## TokenizerTemplateNormalizer\n", @@ -316,7 +375,7 @@ { "cell_type": "code", "execution_count": null, - "id": "11", + "id": "12", "metadata": {}, "outputs": [ { @@ -357,7 +416,7 @@ }, { "cell_type": "markdown", - "id": "12", + "id": "13", "metadata": {}, "source": [ "### System Message Behavior\n", @@ -373,7 +432,7 @@ { "cell_type": "code", "execution_count": null, - "id": "13", + "id": "14", "metadata": {}, "outputs": [ { @@ -416,7 +475,7 @@ }, { "cell_type": "markdown", - "id": "14", + "id": "15", "metadata": {}, "source": [ "### Using Custom Models\n", @@ -427,7 +486,7 @@ { "cell_type": "code", "execution_count": null, - "id": "15", + "id": "16", "metadata": {}, "outputs": [ { @@ -467,7 +526,7 @@ }, { "cell_type": "markdown", - "id": "16", + "id": "17", "metadata": {}, "source": [ "## Creating Custom Normalizers\n", @@ -478,7 +537,7 @@ { "cell_type": "code", "execution_count": null, - "id": "17", + "id": "18", "metadata": {}, "outputs": [ { diff --git a/doc/code/targets/11_message_normalizer.py b/doc/code/targets/11_message_normalizer.py index 87e8cb4441..e6fbf13177 100644 --- a/doc/code/targets/11_message_normalizer.py +++ b/doc/code/targets/11_message_normalizer.py @@ -5,7 +5,7 @@ # extension: .py # format_name: percent # format_version: '1.3' -# jupytext_version: 1.19.4 +# jupytext_version: 1.19.5 # --- # %% [markdown] @@ -20,6 +20,35 @@ # # The `MessageNormalizer` classes handle these conversions, making it easy to work with any target regardless of its expected input format. # +# ## Memory is canonical, the normalized payload is not +# +# A normalizer builds a **target-facing view at send time**. It is never written back to memory. +# +# This matters most for a target that cannot accept editable history. If you prepend eight +# structured turns to a conversation, memory keeps eight structured turns, and the UI, scorers, +# resume, and exports all see eight turns. `HistorySquashNormalizer` flattens those turns into +# a single prompt only for the wire, then discards the flattened copy. The two representations +# are expected to differ. +# +# Flattening the conversation in memory instead would break scoring, resume, and evaluation for +# the sake of one target's wire format. If a prepended piece is not text (an image, for example), +# the flattened view holds a text placeholder and the normalizer logs a warning, because the +# target receives a description instead of the media. +# +# For prepended history on a target without editable history, PyRIT: +# +# 1. Converts and persists the structured prepended messages. +# 2. Converts the live request. +# 3. Applies the per-send `EDITABLE_HISTORY` normalizer selected by +# `PrependedConversationConfig`. +# 4. Applies the target's remaining capability normalizers. +# 5. Serializes the normalized view and invokes the provider. +# +# The attack-owned `PrependedHistorySendContext` records the persisted prepended-message boundary. +# Stateful targets consume it after the first provider attempt; stateless targets reuse it for each +# current request. Targets interact with that state only through an internal `TargetSendContext` +# protocol at the send boundary. +# # ## Base Classes # # There are two base normalizer types: @@ -28,6 +57,31 @@ # # Some normalizers implement both interfaces. +# %% [markdown] +# ## Two History-Squashing Scopes +# +# PyRIT uses `HistorySquashNormalizer` in two places that adapt different target capabilities: +# +# - **Multi-turn support** answers whether the target can continue a conversation across sends. +# - **Editable-history support** answers whether PyRIT can supply or rewrite earlier turns before +# sending the current message. +# +# | Target capabilities | Prepended-history behavior | +# |---|---| +# | Multi-turn with editable history | Send the structured history directly; no history squashing is needed. | +# | Multi-turn without editable history | A context-scoped `HistorySquashNormalizer` encodes the prepended history and first live message into the initial request. Later turns use the target's own conversation state. | +# | Single-turn without editable history | The context-scoped `HistorySquashNormalizer` first produces one message. The target's ordinary use of the same normalizer then sees one message and does nothing. | +# +# The first-turn distinction comes from the attack-owned prepended-history send context, not from a +# separate normalizer implementation. The context applies its configured `HistorySquashNormalizer` +# once to bootstrap prepended history for a target that cannot accept caller-supplied prior turns. +# Its key use case is a multi-turn, server-managed target without editable history. +# +# The target's ordinary capability pipeline independently uses `HistorySquashNormalizer` for a target +# without multi-turn support. Both scopes preserve original-versus-converted text views and keep +# non-text pieces from the current request separate. The context-scoped use can supply a custom +# formatter; the ordinary use defaults to `[Conversation History]` and `[Current Message]` sections. + # %% from pyrit.models import Message diff --git a/doc/scanner/1_pyrit_scan.ipynb b/doc/scanner/1_pyrit_scan.ipynb index 5edd3c5c51..711975eb20 100644 --- a/doc/scanner/1_pyrit_scan.ipynb +++ b/doc/scanner/1_pyrit_scan.ipynb @@ -220,33 +220,24 @@ " Jailbreak scenario implementation for PyRIT. Tests how vulnerable a\n", " model is to jailbreak templates. A run is the cross-product of three\n", " selectors: - **dataset** — the harmful objectives (HarmBench). -\n", - " **techniques** — the *attack techniques* each jailbreak is delivered\n", - " through. Two deliveries are on by default: ``prompt_sending`` (the\n", - " template rendered inline into the user message) and\n", - " ``jailbreak_system_prompt`` (the template set as the system prompt with\n", - " the objective sent as the user turn). The registry techniques\n", - " (``role_play_*``, ``many_shot``, ``tap``, …) are opt-in. -\n", - " **jailbreaks** — which jailbreak templates to run (a random\n", - " ``num_jailbreaks`` sample or an explicit ``jailbreak_names`` set).\n", - " ``prompt_sending`` applies each template as a ``TextJailbreakConverter``\n", - " on the outgoing request, so the objective is rendered inline into the\n", - " template's ``{{prompt}}`` slot; this keeps that delivery target-agnostic\n", - " and lets it compose with every technique. ``jailbreak_system_prompt``\n", - " instead sets the template as a native system prompt and sends the\n", - " objective as its own user turn, so it is only built for targets that\n", - " natively support editable history and system prompts (it is skipped for\n", - " incapable targets, or raises if it is the only selected technique).\n", - " Responses are scored to determine whether the jailbreak succeeded\n", - " (non-refusal).\n", + " **techniques** — two delivery methods for each jailbreak:\n", + " ``prompt_sending`` (the template rendered inline into the user message)\n", + " and ``jailbreak_system_prompt`` (the template set as the system prompt\n", + " with the objective sent as the user turn). - **jailbreaks** — which\n", + " jailbreak templates to run (a random ``num_jailbreaks`` sample or an\n", + " explicit ``jailbreak_names`` set). ``prompt_sending`` applies each\n", + " template as a ``TextJailbreakConverter`` on the outgoing request, so the\n", + " objective is rendered inline into the template's ``{{prompt}}`` slot.\n", + " ``jailbreak_system_prompt`` instead sets the template as a native system\n", + " prompt and sends the objective as its own user turn, so it is only built\n", + " for targets that natively support editable history and system prompts\n", + " (it is skipped for incapable targets, or raises if it is the only\n", + " selected technique). Responses are scored to determine whether the\n", + " jailbreak succeeded (non-refusal).\n", " Aggregate Techniques:\n", - " - all, default, core, light, multi_turn, single_turn\n", - " Available Techniques (16):\n", - " context_compliance, crescendo_history_lecture,\n", - " crescendo_journalist_interview, crescendo_movie_director,\n", - " crescendo_simulated, flip, many_shot, red_teaming,\n", - " role_play_movie_script, role_play_persuasion,\n", - " role_play_persuasion_written, role_play_trivia_game,\n", - " role_play_video_game, tap, prompt_sending, jailbreak_system_prompt\n", + " - all, default, single_turn\n", + " Available Techniques (2):\n", + " prompt_sending, jailbreak_system_prompt\n", " Default Technique: default\n", " Default Datasets (1):\n", " harmbench\n", @@ -256,11 +247,11 @@ " - technique_converters (any): Mapping of concrete technique name to extra request converters to append.\n", " - dataset_config (any): Dataset source configuration; defaults to the scenario's default when omitted.\n", " - memory_labels (any): Additional labels applied to every attack run in the scenario.\n", - " - max_concurrency (int) [default: '4']: Maximum number of concurrent units of work for the scenario.\n", - " - max_retries (int) [default: '0']: Maximum number of automatic retries if the scenario raises an exception.\n", + " - max_concurrency (int) [default: 4]: Maximum number of concurrent units of work for the scenario.\n", + " - max_retries (int) [default: 0]: Maximum number of automatic retries if the scenario raises an exception.\n", " - include_baseline (bool): Whether to prepend a baseline atomic attack; None defers to BASELINE_ATTACK_POLICY.\n", " - num_jailbreaks (int): Draw this many random jailbreak templates for the run. Mutually exclusive with jailbreak_names.\n", - " - num_jailbreak_attempts (int) [default: '1']: Number of times to try each (technique x jailbreak template x objective).\n", + " - num_jailbreak_attempts (int) [default: 1]: Number of times to try each (technique x jailbreak template x objective).\n", " - jailbreak_names (list[str]): Explicit jailbreak template file names to run (e.g. aim.yaml dan_11.yaml). When omitted, a random sample is drawn. Mutually exclusive with num_jailbreaks.\n", "\u001b[1m\u001b[36m\n", " airt.leakage\u001b[0m\n", diff --git a/doc/scanner/airt.ipynb b/doc/scanner/airt.ipynb index 6c5a75a604..74c871973a 100644 --- a/doc/scanner/airt.ipynb +++ b/doc/scanner/airt.ipynb @@ -486,9 +486,11 @@ "objective inline into the template as a request converter (target-agnostic), and\n", "`jailbreak_system_prompt` sets the template as a native system prompt with the objective sent as\n", "the user turn (only for targets that natively support editable history + system prompts — it is\n", - "skipped for incapable targets). Registry techniques like `role_play_*`, `many_shot`, and `tap` are\n", - "opt-in. Results are grouped by jailbreak template, and a baseline (the un-jailbroken objective) is\n", - "included by default so complying with the bare objective is itself visible.\n", + "skipped for incapable targets). These are the only delivery techniques exposed by Jailbreak.\n", + "Generic simulated, multi-turn, or non-composable registry techniques are intentionally excluded\n", + "because they cannot preserve Jailbreak's per-template delivery semantics. Results are grouped by\n", + "jailbreak template, and a baseline (the un-jailbroken objective) is included by default so\n", + "complying with the bare objective is itself visible.\n", "\n", "```bash\n", "pyrit_scan run airt.jailbreak \\\n", @@ -498,10 +500,10 @@ " --max-dataset-size 1\n", "```\n", "\n", - "**Available techniques:** ALL, DEFAULT (`prompt_sending` + `jailbreak_system_prompt`), plus registry\n", - "techniques (`role_play_*`, `many_shot`, `tap`, …). By default a small random sample of jailbreak\n", - "templates runs; pass `num_jailbreaks` (random count) or `jailbreak_names` (explicit) to widen or\n", - "pin the selection." + "**Available technique selectors:** ALL, DEFAULT, and SINGLE_TURN currently select both\n", + "`prompt_sending` and `jailbreak_system_prompt`; either concrete technique can also be selected\n", + "directly. By default a small random sample of jailbreak templates runs; pass `num_jailbreaks`\n", + "(random count) or `jailbreak_names` (explicit) to widen or pin the selection." ] }, { diff --git a/doc/scanner/airt.py b/doc/scanner/airt.py index cc583894d9..1e9366d0c4 100644 --- a/doc/scanner/airt.py +++ b/doc/scanner/airt.py @@ -177,9 +177,11 @@ # objective inline into the template as a request converter (target-agnostic), and # `jailbreak_system_prompt` sets the template as a native system prompt with the objective sent as # the user turn (only for targets that natively support editable history + system prompts — it is -# skipped for incapable targets). Registry techniques like `role_play_*`, `many_shot`, and `tap` are -# opt-in. Results are grouped by jailbreak template, and a baseline (the un-jailbroken objective) is -# included by default so complying with the bare objective is itself visible. +# skipped for incapable targets). These are the only delivery techniques exposed by Jailbreak. +# Generic simulated, multi-turn, or non-composable registry techniques are intentionally excluded +# because they cannot preserve Jailbreak's per-template delivery semantics. Results are grouped by +# jailbreak template, and a baseline (the un-jailbroken objective) is included by default so +# complying with the bare objective is itself visible. # # ```bash # pyrit_scan run airt.jailbreak \ @@ -189,10 +191,10 @@ # --max-dataset-size 1 # ``` # -# **Available techniques:** ALL, DEFAULT (`prompt_sending` + `jailbreak_system_prompt`), plus registry -# techniques (`role_play_*`, `many_shot`, `tap`, …). By default a small random sample of jailbreak -# templates runs; pass `num_jailbreaks` (random count) or `jailbreak_names` (explicit) to widen or -# pin the selection. +# **Available technique selectors:** ALL, DEFAULT, and SINGLE_TURN currently select both +# `prompt_sending` and `jailbreak_system_prompt`; either concrete technique can also be selected +# directly. By default a small random sample of jailbreak templates runs; pass `num_jailbreaks` +# (random count) or `jailbreak_names` (explicit) to widen or pin the selection. # %% from pyrit.scenario.airt import Jailbreak, JailbreakTechnique diff --git a/pyrit/executor/attack/component/conversation_manager.py b/pyrit/executor/attack/component/conversation_manager.py index d92f99bbde..893da6c752 100644 --- a/pyrit/executor/attack/component/conversation_manager.py +++ b/pyrit/executor/attack/component/conversation_manager.py @@ -12,8 +12,11 @@ from pyrit.executor.attack.component.prepended_conversation_config import ( PrependedConversationConfig, ) +from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, +) from pyrit.memory import CentralMemory -from pyrit.message_normalizer import ConversationContextNormalizer, GenericSystemSquashNormalizer +from pyrit.message_normalizer import ConversationContextNormalizer from pyrit.models import ( ChatMessageRole, ComponentIdentifier, @@ -24,6 +27,7 @@ ) from pyrit.prompt_normalizer.prompt_normalizer import PromptNormalizer from pyrit.prompt_target import CapabilityName, PromptTarget +from pyrit.prompt_target.common.target_history import filter_non_replayable_messages if TYPE_CHECKING: from collections.abc import Sequence @@ -275,17 +279,11 @@ async def initialize_context_async( This is the primary method for setting up an attack context. It: 1. Merges memory_labels from attack strategy with context labels - 2. Processes prepended_conversation based on target type and config + 2. Persists prepended_conversation structurally with role-scoped converters 3. Updates context.executed_turns for multi-turn attacks - 4. Sets context.next_message if there's an unanswered user message - - For chat-capable PromptTarget: - - Adds prepended messages to memory with simulated_assistant role - - All messages get new UUIDs - For non-chat PromptTarget: - - Normalizes the prepended conversation to a string and prepends it to - ``context.next_message`` (using ``config.message_normalizer`` when provided). + For all PromptTarget types, prepended messages are added to memory with + simulated_assistant roles and new UUIDs. Args: context: The attack context to initialize. @@ -300,15 +298,13 @@ async def initialize_context_async( ConversationState with turn_count and last_assistant_message_scores. Raises: - ValueError: If conversation_id is empty, or if prepended_conversation - requires a chat-capable PromptTarget but target is not one. + ValueError: If conversation_id is empty. """ if not conversation_id: raise ValueError("conversation_id cannot be empty") # Merge memory labels: attack strategy labels + context labels context.memory_labels = combine_dict(existing_dict=memory_labels, new_dict=context.memory_labels) - state = ConversationState() prepended_conversation = context.prepended_conversation @@ -316,19 +312,7 @@ async def initialize_context_async( logger.debug(f"No prepended conversation for context initialization: {conversation_id}") return state - # Targets that don't natively support editable history cannot consume a - # prepended multi-message conversation as-is — route them to the - # single-string fallback path via capability-based routing. - is_chat_target = target.configuration.includes(capability=CapabilityName.EDITABLE_HISTORY) - if not is_chat_target: - return await self._handle_non_chat_target_async( - context=context, - prepended_conversation=prepended_conversation, - config=prepended_conversation_config, - ) - - # Process prepended conversation for objective target - return await self._process_prepended_for_chat_target_async( + return await self._process_prepended_conversation_async( context=context, prepended_conversation=prepended_conversation, conversation_id=conversation_id, @@ -336,76 +320,9 @@ async def initialize_context_async( prepended_conversation_config=prepended_conversation_config, max_turns=max_turns, target_identifier=target.get_identifier(), + target=target, ) - async def _handle_non_chat_target_async( - self, - *, - context: AttackContext[Any], - prepended_conversation: list[Message], - config: PrependedConversationConfig | None, - ) -> ConversationState: - """ - Handle prepended conversation for non-chat targets. - - Args: - context: The attack context. - prepended_conversation: Messages to prepend. - config: Configuration for non-chat target behavior. - - Returns: - Empty ConversationState (non-chat targets don't track turns). - """ - if config is None: - config = PrependedConversationConfig() - - normalizer = config.get_message_normalizer() - messages_to_normalize = prepended_conversation - if isinstance(normalizer, ConversationContextNormalizer): - messages_to_normalize = await GenericSystemSquashNormalizer().normalize_async(prepended_conversation) - - normalized_context = await normalizer.normalize_string_async(messages_to_normalize) - - next_message = context.next_message - if next_message is None: - next_message = Message.from_prompt(prompt=context.objective, role="user") - context.next_message = next_message - - if normalized_context: - # Find an existing text piece to prepend to - text_piece = None - for piece in next_message.message_pieces: - if piece.original_value_data_type == "text": - text_piece = piece - break - - if text_piece: - # Prepend context to the existing text piece - context_prefix = f"{normalized_context}\n\n" - if text_piece.original_value != normalized_context and not text_piece.original_value.startswith( - context_prefix - ): - text_piece.original_value = f"{context_prefix}{text_piece.original_value}" - if text_piece.converted_value != normalized_context and not text_piece.converted_value.startswith( - context_prefix - ): - text_piece.converted_value = f"{context_prefix}{text_piece.converted_value}" - else: - # No text piece found (multimodal message), add a new text piece at the beginning - context_piece = MessagePiece( - id=uuid.uuid4(), - role="user", - original_value=normalized_context, - converted_value=normalized_context, - original_value_data_type="text", - converted_value_data_type="text", - ) - # Create a new message with the context piece prepended - context.next_message = Message(message_pieces=[context_piece] + list(next_message.message_pieces)) - - logger.debug(f"Normalized prepended conversation for non-chat target: {len(normalized_context)} characters") - return ConversationState() - async def add_prepended_conversation_to_memory_async( self, *, @@ -415,9 +332,10 @@ async def add_prepended_conversation_to_memory_async( prepended_conversation_config: PrependedConversationConfig | None = None, max_turns: int | None = None, target_identifier: ComponentIdentifier | None = None, + target: PromptTarget | None = None, ) -> int: """ - Add prepended conversation messages to memory for a chat target. + Add prepended conversation messages to memory for a target. This is a lower-level method that handles adding messages to memory without modifying any attack context state. It can be called directly by attacks @@ -437,6 +355,7 @@ async def add_prepended_conversation_to_memory_async( max_turns: If provided, validates that turn count doesn't exceed this limit. target_identifier (ComponentIdentifier | None): The target the conversation is held with, if known. Recorded once per conversation. + target (PromptTarget | None): Target that will receive the first live request. Returns: The number of turns (assistant messages) added. @@ -444,23 +363,25 @@ async def add_prepended_conversation_to_memory_async( Raises: ValueError: If max_turns is exceeded by the prepended conversation. """ - # Filter valid messages - valid_messages = [msg for msg in prepended_conversation if msg and msg.message_pieces] + valid_messages = self.get_persistable_prepended_messages(prepended_conversation=prepended_conversation) if not valid_messages: return 0 - self._memory.add_conversation_to_memory( - conversation=Conversation(conversation_id=conversation_id, target_identifier=target_identifier) - ) + if target and target_identifier is None: + target_identifier = target.get_identifier() - # Get roles that should have converters applied - apply_to_roles = ( - prepended_conversation_config.apply_converters_to_roles if prepended_conversation_config else None + # Assistant history represents simulated target output, so the absent-config + # path must use the same safe role default as an explicit default config. + config = prepended_conversation_config or PrependedConversationConfig() + apply_to_roles = config.apply_converters_to_roles + requires_prepended_adaptation = bool( + target and not target.configuration.includes(capability=CapabilityName.EDITABLE_HISTORY) ) turn_count = 0 + prepared_messages: list[Message] = [] - for i, message in enumerate(valid_messages): + for message in valid_messages: message_copy = message.duplicate() message_copy.set_simulated_role() @@ -485,14 +406,68 @@ async def add_prepended_conversation_to_memory_async( request_converters=request_converters, apply_to_roles=apply_to_roles, ) + if requires_prepended_adaptation: + self._validate_flattenable_converter_output( + source_message=message, + converted_message=message_copy, + ) - # Add to memory - self._memory.add_message_to_memory(request=message_copy) - logger.debug(f"Added prepended message {i + 1}/{len(valid_messages)} to memory") + prepared_messages.append(message_copy) + + self._memory.add_conversation_to_memory( + conversation=Conversation(conversation_id=conversation_id, target_identifier=target_identifier) + ) + for i, message in enumerate(prepared_messages): + self._memory.add_message_to_memory(request=message) + logger.debug(f"Added prepended message {i + 1}/{len(prepared_messages)} to memory") return turn_count - async def _process_prepended_for_chat_target_async( + @staticmethod + def get_persistable_prepended_messages( + *, + prepended_conversation: list[Message], + ) -> list[Message]: + """ + Return prepended messages that can be recovered from memory at send time. + + Args: + prepended_conversation: Candidate prepended messages. + + Returns: + list[Message]: Non-empty messages containing at least one persistable piece. + """ + persistable_messages = [ + message + for message in prepended_conversation + if message and message.message_pieces and any(not piece.not_in_memory for piece in message.message_pieces) + ] + return filter_non_replayable_messages(messages=persistable_messages) + + @staticmethod + def create_prepended_history_send_context( + *, + target: PromptTarget, + conversation_id: str, + prepended_messages: list[Message], + ) -> PrependedHistorySendContext | None: + """ + Build persisted-prefix state for a target without editable history. + + Returns: + PrependedHistorySendContext | None: Per-send state, or ``None`` for + editable-history targets or empty prepended history. + """ + if not prepended_messages or target.configuration.includes(capability=CapabilityName.EDITABLE_HISTORY): + return None + + return PrependedHistorySendContext( + conversation_id=conversation_id, + seed_message_ids=tuple(message.get_piece().id for message in prepended_messages), + replay_seed_each_send=not target.configuration.includes(capability=CapabilityName.MULTI_TURN), + ) + + async def _process_prepended_conversation_async( self, *, context: AttackContext[Any], @@ -502,9 +477,10 @@ async def _process_prepended_for_chat_target_async( prepended_conversation_config: PrependedConversationConfig | None, max_turns: int | None, target_identifier: ComponentIdentifier | None = None, + target: PromptTarget, ) -> ConversationState: """ - Process prepended conversation for a chat target. + Process prepended conversation for a target. Adds messages to memory with: - New UUIDs for all pieces @@ -520,6 +496,7 @@ async def _process_prepended_for_chat_target_async( max_turns: Maximum turns for validation. target_identifier (ComponentIdentifier | None): The objective target the conversation is held with, if known. + target: The objective target that will receive the conversation. Returns: ConversationState with turn_count and scores. @@ -527,11 +504,14 @@ async def _process_prepended_for_chat_target_async( state = ConversationState() is_multi_turn = max_turns is not None - # Filter valid messages - valid_messages = [msg for msg in prepended_conversation if msg and msg.message_pieces] + valid_messages = self.get_persistable_prepended_messages(prepended_conversation=prepended_conversation) if not valid_messages: return state + existing_message_ids = { + message.get_piece().id for message in self.get_conversation(conversation_id=conversation_id) + } + # Use the lower-level method to add messages to memory state.turn_count = await self.add_prepended_conversation_to_memory_async( prepended_conversation=prepended_conversation, @@ -540,8 +520,18 @@ async def _process_prepended_for_chat_target_async( prepended_conversation_config=prepended_conversation_config, max_turns=max_turns, target_identifier=target_identifier, + target=target, + ) + persisted_messages = [ + message + for message in self.get_conversation(conversation_id) + if message.get_piece().id not in existing_message_ids + ] + context.prepended_history_send_context = self.create_prepended_history_send_context( + target=target, + conversation_id=conversation_id, + prepended_messages=persisted_messages, ) - # Update context for multi-turn attacks to reflect prepended_conversation final_prepended_message = valid_messages[-1] @@ -557,7 +547,11 @@ async def _process_prepended_for_chat_target_async( # true_false scores with score_value=False so attacks can use the rationale for # feedback without re-scoring. memory_pieces = self._memory.get_message_pieces(conversation_id=conversation_id) - assistant_piece_ids = [str(piece.id) for piece in memory_pieces if piece.api_role == "assistant"] + assistant_pieces = [piece for piece in memory_pieces if piece.api_role == "assistant"] + last_assistant_sequence = max((piece.sequence for piece in assistant_pieces), default=None) + assistant_piece_ids = [ + str(piece.id) for piece in assistant_pieces if piece.sequence == last_assistant_sequence + ] existing_scores = ( self._memory.get_prompt_scores(prompt_ids=assistant_piece_ids) if assistant_piece_ids else [] ) @@ -570,12 +564,42 @@ async def _process_prepended_for_chat_target_async( return state + @staticmethod + def _validate_flattenable_converter_output( + *, + source_message: Message, + converted_message: Message, + ) -> None: + """ + Reject non-text output produced by this prepended conversion pass. + + Raises: + ValueError: If an applied converter produced non-text prepended history. + """ + output_types = { + converted_piece.converted_value_data_type + for source_piece, converted_piece in zip( + source_message.message_pieces, + converted_message.message_pieces, + strict=True, + ) + if len(converted_piece.converter_identifiers) > len(source_piece.converter_identifiers) + and not converted_piece.not_in_memory + and converted_piece.converted_value_data_type != "text" + } + if output_types: + raise ValueError( + "Cannot flatten prepended conversation for a target without editable history after " + f"request converters produced non-text output types {sorted(output_types)}. Prepended " + "conversion must produce text." + ) + async def _apply_converters_async( self, *, message: Message, request_converters: list[ConverterConfiguration], - apply_to_roles: list[ChatMessageRole] | None, + apply_to_roles: list[ChatMessageRole], ) -> None: """ Apply converters to message pieces. @@ -583,16 +607,15 @@ async def _apply_converters_async( Args: message: The message containing pieces to convert. request_converters: Converter configurations to apply. - apply_to_roles: If provided, only apply to pieces with these roles. - If None, apply to all roles. + apply_to_roles: Only apply to pieces with these roles. """ - for piece in message.message_pieces: - # Filter by role if specified - if apply_to_roles is not None and piece.api_role not in apply_to_roles: - continue - - temp_message = Message(message_pieces=[piece]) - await self._prompt_normalizer.convert_values_async( - message=temp_message, - converter_configurations=request_converters, - ) + if message.api_role not in apply_to_roles: + return + + # Apply to the complete message so ConverterConfiguration.indexes_to_apply remains relative + # to the original piece list. Converting one temporary piece at a time would reset every + # selected piece to index zero. + await self._prompt_normalizer.convert_values_async( + message=message, + converter_configurations=request_converters, + ) diff --git a/pyrit/executor/attack/component/prepended_conversation_config.py b/pyrit/executor/attack/component/prepended_conversation_config.py index a0daedfd6e..3ec3e35b61 100644 --- a/pyrit/executor/attack/component/prepended_conversation_config.py +++ b/pyrit/executor/attack/component/prepended_conversation_config.py @@ -4,13 +4,22 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import get_args +from typing import TYPE_CHECKING from pyrit.message_normalizer import ( ConversationContextNormalizer, + HistorySquashNormalizer, + MessageListNormalizer, MessageStringNormalizer, ) -from pyrit.models import ChatMessageRole +from pyrit.prompt_target.common.target_capabilities import CapabilityName + +if TYPE_CHECKING: + from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, + ) + from pyrit.models import ChatMessageRole, Message + from pyrit.prompt_target.common.prompt_target import PromptTarget @dataclass @@ -21,21 +30,23 @@ class PrependedConversationConfig: This class provides control over: - Which message roles should have request converters applied - - How to normalize conversation history for non-chat objective targets + - How targets without editable history format prepended messages with live requests - Non-chat objective targets always normalize the prepended conversation into the - first turn (via ``message_normalizer``; default: ConversationContextNormalizer). + Prepended messages remain role-structured in memory. Request converters are applied to + configured roles before a target without editable history renders that history and the + applicable live request together (via ``message_normalizer``; default: ConversationContextNormalizer). + Those converters must produce text because string normalization cannot preserve converted + image, audio, or other non-text output. """ - # Roles for which request converters should be applied to prepended messages. - # By default, converters are applied to all roles. - # Example: ["user"] to apply converters only to user messages. - apply_converters_to_roles: list[ChatMessageRole] = field(default_factory=lambda: list(get_args(ChatMessageRole))) + # Request converters default to prepended user messages only. Assistant history is + # simulated target output and must be explicitly opted in with ["assistant"]. + apply_converters_to_roles: list[ChatMessageRole] = field(default_factory=lambda: ["user"]) - # Optional normalizer to format conversation history into a single text block. + # Optional normalizer to format prepended history and a live request as one text block. # Must implement MessageStringNormalizer (e.g., TokenizerTemplateNormalizer or ConversationContextNormalizer). - # When None and normalization is needed (e.g., for non-chat targets), a default - # ConversationContextNormalizer is used that produces "Turn N: User/Assistant" format. + # When None and adaptation is needed, a default ConversationContextNormalizer is used + # that produces "Turn N: User/Assistant" format. message_normalizer: MessageStringNormalizer | None = None def get_message_normalizer(self) -> MessageStringNormalizer: @@ -47,3 +58,34 @@ def get_message_normalizer(self) -> MessageStringNormalizer: ConversationContextNormalizer if none was configured. """ return self.message_normalizer or ConversationContextNormalizer() + + def get_normalizer_overrides( + self, + *, + target: PromptTarget, + prepended_history_send_context: PrependedHistorySendContext | None, + ) -> dict[CapabilityName, MessageListNormalizer[Message]]: + """ + Build per-send target normalizer overrides for prepended history. + + Args: + target: Target that receives the live request. + prepended_history_send_context: Explicit persisted seed boundary for + this attack execution. + + Returns: + Overrides keyed by the capability they adapt. + """ + if ( + target.configuration.includes(capability=CapabilityName.EDITABLE_HISTORY) + or prepended_history_send_context is None + or not prepended_history_send_context.should_include_seed + ): + return {} + + return { + CapabilityName.EDITABLE_HISTORY: HistorySquashNormalizer( + expected_history_message_count=prepended_history_send_context.bootstrap_message_count, + message_normalizer=self.get_message_normalizer(), + ) + } diff --git a/pyrit/executor/attack/component/prepended_history_send_context.py b/pyrit/executor/attack/component/prepended_history_send_context.py new file mode 100644 index 0000000000..a11df417fc --- /dev/null +++ b/pyrit/executor/attack/component/prepended_history_send_context.py @@ -0,0 +1,246 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import uuid + from typing import Any + + from pyrit.models import Message + + +@dataclass(slots=True) +class PrependedHistorySendContext: + """ + Per-execution state for delivering persisted bootstrap history. + + An explicit seed is identified by persisted message-piece IDs instead of + inferred from response roles or error codes. A stateful cloned conversation + may instead have an empty seed plus copied branch messages to bootstrap its + new provider session. A context may be used by only one send at a time. + Stateful targets consume the bootstrap when provider invocation begins; + stateless targets replay their explicit seed for every send. + """ + + conversation_id: str + seed_message_ids: tuple[uuid.UUID, ...] + replay_seed_each_send: bool + bootstrap_message_ids: tuple[uuid.UUID, ...] | None = None + _seed_consumed: bool = field(default=False, init=False, repr=False) + _send_in_progress: bool = field(default=False, init=False, repr=False) + _provider_attempt_marked: bool = field(default=False, init=False, repr=False) + _provider_attempted_task: asyncio.Task[Any] | None = field(default=None, init=False, repr=False) + _provider_attempt_count: int = field(default=0, init=False, repr=False) + + def __post_init__(self) -> None: + """ + Validate the persisted seed identity. + + Raises: + ValueError: If the conversation ID or seed IDs are invalid. + """ + if not self.conversation_id: + raise ValueError("conversation_id must not be empty") + if not self.seed_message_ids and not self.bootstrap_message_ids: + raise ValueError("seed_message_ids must not be empty unless bootstrap_message_ids are provided") + if len(set(self.seed_message_ids)) != len(self.seed_message_ids): + raise ValueError("seed_message_ids must be unique") + if self.bootstrap_message_ids is not None: + if not self.bootstrap_message_ids: + raise ValueError("bootstrap_message_ids must not be empty") + if len(set(self.bootstrap_message_ids)) != len(self.bootstrap_message_ids): + raise ValueError("bootstrap_message_ids must be unique") + if not set(self.seed_message_ids).issubset(self.bootstrap_message_ids): + raise ValueError("bootstrap_message_ids must include every seed message") + + @property + def seed_message_count(self) -> int: + """Number of messages in the explicit seed prefix.""" + return len(self.seed_message_ids) + + @property + def bootstrap_message_count(self) -> int: + """Number of historical messages needed to bootstrap the next provider session.""" + return len(self.bootstrap_message_ids or self.seed_message_ids) + + @property + def should_include_seed(self) -> bool: + """Whether the seed prefix should be included in the next send.""" + return self.replay_seed_each_send or not self._seed_consumed + + @property + def is_seed_consumed(self) -> bool: + """Whether a stateful target has consumed the seed prefix.""" + return self._seed_consumed + + @property + def provider_attempt_count(self) -> int: + """Number of sends that reached provider invocation.""" + return self._provider_attempt_count + + @property + def provider_attempted_by_current_task(self) -> bool: + """Whether the current send task reached provider invocation.""" + try: + current_task = asyncio.current_task() + except RuntimeError: + return False + return current_task is not None and current_task is self._provider_attempted_task + + def begin_send(self) -> None: + """ + Acquire this context for one complete send. + + Raises: + RuntimeError: If another send using this context is still active. + """ + if self._send_in_progress: + raise RuntimeError( + "Concurrent sends for the same prepended history send context are not supported. " + "Wait for the active send to finish before sending another request." + ) + self._send_in_progress = True + self._provider_attempt_marked = False + self._provider_attempted_task = None + + def mark_provider_attempted(self) -> None: + """ + Record that provider invocation has begun for the active send. + + Stateful targets consume their seed history at this boundary even if the + provider later fails or the task is cancelled. + + Raises: + RuntimeError: If no send currently owns the context. + """ + if not self._send_in_progress: + raise RuntimeError("Cannot mark a provider attempt without an active send.") + if self._provider_attempt_marked: + return + try: + current_task = asyncio.current_task() + except RuntimeError: + current_task = None + self._provider_attempt_marked = True + self._provider_attempted_task = current_task + self._provider_attempt_count += 1 + if not self.replay_seed_each_send: + self._seed_consumed = True + + def finish_send(self) -> None: + """Release this context after the active send completes or is cancelled.""" + if not self._send_in_progress: + return + self._send_in_progress = False + + def select_history(self, *, messages: list[Message]) -> list[Message]: + """ + Select history needed for delivery or provider-session restoration. + + Args: + messages: Persisted conversation messages after failed exchanges + have been removed. + + Returns: + The pending bootstrap messages before a stateful target consumes + them, the explicit seed for a stateless target, or all replayable + persisted history after a stateful target has been bootstrapped. + + Raises: + ValueError: If an expected persisted seed message is missing. + """ + if not self.should_include_seed: + return list(messages) + + messages_by_id = {message.get_piece().id: message for message in messages} + selected_message_ids = self.bootstrap_message_ids or self.seed_message_ids + missing_ids = [message_id for message_id in selected_message_ids if message_id not in messages_by_id] + if missing_ids: + raise ValueError( + "The persisted prepended history no longer matches the prepended history send context. " + f"Missing {len(missing_ids)} message(s)." + ) + return [messages_by_id[message_id] for message_id in selected_message_ids] + + def remap_for_duplicate_conversation( + self, + *, + conversation_id: str, + source_messages: list[Message], + duplicated_messages: list[Message], + ) -> PrependedHistorySendContext: + """ + Remap this explicit seed boundary to a duplicated conversation. + + The explicit prepended seed identity is always remapped. A stateless clone + continues to replay only that seed. A stateful clone opens a new provider + session, so every replayable duplicated branch message becomes its one-time + bootstrap boundary. + + Args: + conversation_id: Conversation ID assigned to the duplicated messages. + source_messages: Messages from the source conversation in persisted order. + duplicated_messages: Their duplicates in the same persisted order. + + Returns: + A new context with seed IDs from the duplicated conversation. + A new logical conversation starts with an unconsumed seed boundary. + + Raises: + ValueError: If the duplicated messages do not match the source structure + or an explicit seed piece cannot be remapped. + """ + if len(source_messages) != len(duplicated_messages): + raise ValueError("Duplicated conversation does not match the source message count.") + + duplicated_ids_by_source_id: dict[uuid.UUID, uuid.UUID] = {} + for source_message, duplicated_message in zip(source_messages, duplicated_messages, strict=True): + if ( + source_message.api_role != duplicated_message.api_role + or source_message.sequence != duplicated_message.sequence + or len(source_message.message_pieces) != len(duplicated_message.message_pieces) + ): + raise ValueError("Duplicated conversation does not preserve the source message structure.") + + for source_piece, duplicated_piece in zip( + source_message.message_pieces, + duplicated_message.message_pieces, + strict=True, + ): + if ( + source_piece.api_role != duplicated_piece.api_role + or source_piece.sequence != duplicated_piece.sequence + or duplicated_piece.conversation_id != conversation_id + ): + raise ValueError("Duplicated conversation does not preserve the source piece structure.") + duplicated_ids_by_source_id[source_piece.id] = duplicated_piece.id + + current_bootstrap_ids = self.bootstrap_message_ids or self.seed_message_ids + ids_to_remap = set(self.seed_message_ids) | set(current_bootstrap_ids) + missing_ids = [message_id for message_id in ids_to_remap if message_id not in duplicated_ids_by_source_id] + if missing_ids: + raise ValueError(f"Could not remap {len(missing_ids)} bootstrap message(s).") + + remapped_seed_ids = tuple(duplicated_ids_by_source_id[message_id] for message_id in self.seed_message_ids) + if conversation_id != self.conversation_id and not self.replay_seed_each_send: + remapped_bootstrap_ids = tuple(message.get_piece().id for message in duplicated_messages) + else: + remapped_bootstrap_ids = tuple( + duplicated_ids_by_source_id[message_id] for message_id in current_bootstrap_ids + ) + + duplicated_context = PrependedHistorySendContext( + conversation_id=conversation_id, + seed_message_ids=remapped_seed_ids, + replay_seed_each_send=self.replay_seed_each_send, + bootstrap_message_ids=remapped_bootstrap_ids, + ) + # Provider bootstrap consumption belongs to the logical conversation, not copied memory. + if conversation_id == self.conversation_id: + duplicated_context._seed_consumed = self._seed_consumed + return duplicated_context diff --git a/pyrit/executor/attack/core/attack_strategy.py b/pyrit/executor/attack/core/attack_strategy.py index 11c58b60ad..a6fee0f286 100644 --- a/pyrit/executor/attack/core/attack_strategy.py +++ b/pyrit/executor/attack/core/attack_strategy.py @@ -42,12 +42,20 @@ from pyrit.prompt_target.common.target_requirements import TargetRequirements if TYPE_CHECKING: + from pyrit.executor.attack.component.prepended_conversation_config import ( + PrependedConversationConfig, + ) + from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, + ) from pyrit.executor.attack.core.attack_config import ( AttackAdversarialConfig, AttackScoringConfig, ) from pyrit.executor.attack.core.attack_result_attribution import AttackResultAttribution + from pyrit.message_normalizer import MessageListNormalizer from pyrit.prompt_target import PromptTarget + from pyrit.prompt_target.common.target_capabilities import CapabilityName AttackStrategyContextT = TypeVar("AttackStrategyContextT", bound="AttackContext[Any]") AttackStrategyResultT = TypeVar("AttackStrategyResultT", bound="AttackResult") @@ -88,6 +96,13 @@ class AttackContext(StrategyContext, ABC, Generic[AttackParamsT]): _prepended_conversation_override: list[Message] | None = None _memory_labels_override: dict[str, str] | None = None + # Per-execution prepended-history boundary and send lifecycle. Never persisted. + prepended_history_send_context: PrependedHistorySendContext | None = field( + default=None, + repr=False, + compare=False, + ) + # Optional attribution from an upstream orchestrator (e.g. Scenario). When # set, the persistence path stamps attribution_parent_id + attribution_data # onto the resulting AttackResult so it can be located later for hydration @@ -422,6 +437,7 @@ def __init__( objective_target: PromptTarget, context_type: type[AttackStrategyContextT], params_type: type[AttackParamsT] = AttackParameters, # type: ignore[ty:invalid-parameter-default] + prepended_conversation_config: PrependedConversationConfig | None = None, logger: logging.Logger = logger, ) -> None: """ @@ -433,6 +449,9 @@ def __init__( params_type (type[AttackParamsT]): The type of parameters this strategy accepts. Defaults to AttackParameters. Use AttackParameters.excluding() to create a params type that rejects certain fields. + prepended_conversation_config (PrependedConversationConfig | None): Policy for + prepended conversations. Controls converter role scope and target-facing + history formatting. logger (logging.Logger): Logger instance for logging events. """ super().__init__( @@ -442,15 +461,40 @@ def __init__( ), logger=logger, ) + # Local import avoids the component package's import cycle through attack config. + from pyrit.executor.attack.component.prepended_conversation_config import ( + PrependedConversationConfig, + ) + type(self).TARGET_REQUIREMENTS.validate(target=objective_target) self._objective_target = objective_target self._params_type = params_type + self._prepended_conversation_config = prepended_conversation_config or PrependedConversationConfig() # Guard so subclasses that set converters before calling super() aren't clobbered if not hasattr(self, "_request_converters"): self._request_converters: list[Any] = [] if not hasattr(self, "_response_converters"): self._response_converters: list[Any] = [] + def _get_prepended_normalizer_overrides( + self, + *, + prepended_history_send_context: PrependedHistorySendContext | None, + ) -> dict[CapabilityName, MessageListNormalizer[Message]]: + """ + Resolve prepended-history overrides for one target send. + + Args: + prepended_history_send_context: Persisted seed boundary for this execution. + + Returns: + Overrides keyed by the capability they adapt. + """ + return self._prepended_conversation_config.get_normalizer_overrides( + target=self._objective_target, + prepended_history_send_context=prepended_history_send_context, + ) + def _create_identifier( self, *, @@ -478,6 +522,10 @@ def _create_identifier( objective_target = TargetIdentifier.from_component_identifier(self.get_objective_target().get_identifier()) + prepended_config = self._prepended_conversation_config + merged_params["prepended_conversation_converter_roles"] = list(prepended_config.apply_converters_to_roles) + all_children["prepended_conversation_formatter"] = prepended_config.get_message_normalizer().get_identifier() + # Add scorer if present objective_scorer: ScorerIdentifier | None = None scoring_config = self.get_attack_scoring_config() diff --git a/pyrit/executor/attack/multi_turn/chunked_request.py b/pyrit/executor/attack/multi_turn/chunked_request.py index 6db7537722..b90eca4d09 100644 --- a/pyrit/executor/attack/multi_turn/chunked_request.py +++ b/pyrit/executor/attack/multi_turn/chunked_request.py @@ -9,7 +9,7 @@ from pyrit.common.apply_defaults import REQUIRED_VALUE, apply_defaults from pyrit.exceptions import ComponentRole, execution_context -from pyrit.executor.attack.component import ConversationManager +from pyrit.executor.attack.component import ConversationManager, PrependedConversationConfig from pyrit.executor.attack.core.attack_config import ( AttackConverterConfig, AttackScoringConfig, @@ -105,6 +105,7 @@ def __init__( attack_converter_config: AttackConverterConfig | None = None, attack_scoring_config: AttackScoringConfig | None = None, prompt_normalizer: PromptNormalizer | None = None, + prepended_conversation_config: PrependedConversationConfig | None = None, ) -> None: """ Initialize the chunked request attack strategy. @@ -119,6 +120,8 @@ def __init__( attack_converter_config (AttackConverterConfig | None): Configuration for converters. attack_scoring_config (AttackScoringConfig | None): Configuration for scoring components. prompt_normalizer (PromptNormalizer | None): Normalizer for handling prompts. + prepended_conversation_config: Configuration for prepended-conversation + conversion and target-facing formatting. Raises: ValueError: If chunk_size or total_length are invalid. @@ -150,6 +153,7 @@ def __init__( logger=logger, context_type=ChunkedRequestAttackContext, params_type=ChunkedRequestAttackParameters, + prepended_conversation_config=prepended_conversation_config, ) # Store chunk configuration @@ -246,6 +250,7 @@ async def _setup_async(self, *, context: ChunkedRequestAttackContext) -> None: target=self._objective_target, conversation_id=context.session.conversation_id, request_converters=self._request_converters, + prepended_conversation_config=self._prepended_conversation_config, memory_labels=self._memory_labels, ) @@ -288,6 +293,10 @@ async def _perform_async(self, *, context: ChunkedRequestAttackContext) -> Attac conversation_id=context.session.conversation_id, request_converter_configurations=self._request_converters, response_converter_configurations=self._response_converters, + normalizer_overrides=self._get_prepended_normalizer_overrides( + prepended_history_send_context=context.prepended_history_send_context, + ), + send_context=context.prepended_history_send_context, ) # Store the response diff --git a/pyrit/executor/attack/multi_turn/crescendo.py b/pyrit/executor/attack/multi_turn/crescendo.py index 68ca358db5..b9fb98a677 100644 --- a/pyrit/executor/attack/multi_turn/crescendo.py +++ b/pyrit/executor/attack/multi_turn/crescendo.py @@ -130,14 +130,10 @@ class CrescendoAttack(MultiTurnAttackStrategy[CrescendoAttackContext, CrescendoA You can learn more about the Crescendo attack [@russinovich2024crescendo]. """ - # Crescendo fundamentally relies on multi-turn conversation history to - # gradually escalate prompts; history-squash adaptation would collapse the - # conversation into a single prompt and silently break the attack's - # semantics. Declare MULTI_TURN as native_required so adaptation is - # rejected at construction time. + # Crescendo relies on native live history for both gradual escalation and + # backtracking. Squashed or non-editable history changes those semantics. TARGET_REQUIREMENTS = TargetRequirements( - required=frozenset({CapabilityName.EDITABLE_HISTORY, CapabilityName.MULTI_TURN}), - native_required=frozenset({CapabilityName.MULTI_TURN}), + native_required=frozenset({CapabilityName.EDITABLE_HISTORY, CapabilityName.MULTI_TURN}), ) # Default system prompt template path for Crescendo attack @@ -174,13 +170,19 @@ def __init__( max_turns (int): Maximum number of turns allowed. prepended_conversation_config (PrependedConversationConfiguration | None): Configuration for how to process prepended conversations. Controls converter - application by role, message normalization, and non-chat target behavior. + application by role and request formatting for targets without editable history. Raises: - ValueError: If objective_target does not natively support editable history. + ValueError: If the objective target does not natively support multi-turn conversations + and editable history. """ # Initialize base class - super().__init__(objective_target=objective_target, logger=logger, context_type=CrescendoAttackContext) + super().__init__( + objective_target=objective_target, + logger=logger, + context_type=CrescendoAttackContext, + prepended_conversation_config=prepended_conversation_config, + ) self._memory = CentralMemory.get_memory_instance() @@ -264,9 +266,6 @@ def __init__( self._max_backtracks = max_backtracks self._max_turns = max_turns - # Store the prepended conversation configuration - self._prepended_conversation_config = prepended_conversation_config - def get_attack_scoring_config(self) -> AttackScoringConfig | None: """ Get the attack scoring configuration used by this strategy. @@ -642,6 +641,10 @@ async def _send_prompt_to_objective_target_async( conversation_id=context.session.conversation_id, request_converter_configurations=self._request_converters, response_converter_configurations=self._response_converters, + normalizer_overrides=self._get_prepended_normalizer_overrides( + prepended_history_send_context=context.prepended_history_send_context, + ), + send_context=context.prepended_history_send_context, ) if not response: diff --git a/pyrit/executor/attack/multi_turn/multi_prompt_sending.py b/pyrit/executor/attack/multi_turn/multi_prompt_sending.py index d431ec20db..eb065fb6d0 100644 --- a/pyrit/executor/attack/multi_turn/multi_prompt_sending.py +++ b/pyrit/executor/attack/multi_turn/multi_prompt_sending.py @@ -10,7 +10,7 @@ from pyrit.common.apply_defaults import REQUIRED_VALUE, apply_defaults from pyrit.common.utils import get_kwarg_param from pyrit.exceptions import ComponentRole, execution_context -from pyrit.executor.attack.component import ConversationManager +from pyrit.executor.attack.component import ConversationManager, PrependedConversationConfig from pyrit.executor.attack.core.attack_config import ( AttackConverterConfig, AttackScoringConfig, @@ -142,6 +142,7 @@ def __init__( attack_converter_config: AttackConverterConfig | None = None, attack_scoring_config: AttackScoringConfig | None = None, prompt_normalizer: PromptNormalizer | None = None, + prepended_conversation_config: PrependedConversationConfig | None = None, ) -> None: """ Initialize the multi-prompt sending attack strategy. @@ -151,6 +152,8 @@ def __init__( attack_converter_config (AttackConverterConfig | None): Configuration for converters. attack_scoring_config (AttackScoringConfig | None): Configuration for scoring components. prompt_normalizer (PromptNormalizer | None): Normalizer for handling prompts. + prepended_conversation_config: Configuration for prepended-conversation + conversion and target-facing formatting. Raises: ValueError: If the objective scorer is not a true/false scorer. @@ -161,6 +164,7 @@ def __init__( logger=logger, context_type=MultiTurnAttackContext, params_type=MultiPromptSendingAttackParameters, + prepended_conversation_config=prepended_conversation_config, ) # Initialize the converter configuration @@ -224,6 +228,7 @@ async def _setup_async(self, *, context: MultiTurnAttackContext[Any]) -> None: target=self._objective_target, conversation_id=context.session.conversation_id, request_converters=self._request_converters, + prepended_conversation_config=self._prepended_conversation_config, memory_labels=self._memory_labels, ) @@ -366,6 +371,10 @@ async def _send_prompt_to_objective_target_async( conversation_id=context.session.conversation_id, request_converter_configurations=self._request_converters, response_converter_configurations=self._response_converters, + normalizer_overrides=self._get_prepended_normalizer_overrides( + prepended_history_send_context=context.prepended_history_send_context, + ), + send_context=context.prepended_history_send_context, ) async def _evaluate_response_async(self, *, response: Message, objective: str) -> Score | None: diff --git a/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py b/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py index db7a44755b..0aa621b7f5 100644 --- a/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py +++ b/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Any, TypeVar from pyrit.common.logger import logger +from pyrit.executor.attack.component.conversation_manager import ConversationManager from pyrit.executor.attack.core.attack_parameters import AttackParameters, AttackParamsT from pyrit.executor.attack.core.attack_strategy import ( AttackContext, @@ -21,6 +22,9 @@ from pyrit.prompt_target import CapabilityName if TYPE_CHECKING: + from pyrit.executor.attack.component.prepended_conversation_config import ( + PrependedConversationConfig, + ) from pyrit.models import ( Message, Score, @@ -77,6 +81,7 @@ def __init__( objective_target: PromptTarget, context_type: type[MultiTurnAttackStrategyContextT], params_type: type[AttackParamsT] = AttackParameters, # type: ignore[ty:invalid-parameter-default] + prepended_conversation_config: PrependedConversationConfig | None = None, logger: logging.Logger = logger, ) -> None: """ @@ -86,12 +91,15 @@ def __init__( objective_target (PromptTarget): The target system to attack. context_type (type[MultiTurnAttackContext]): The type of context this strategy will use. params_type (type[AttackParamsT]): The type of parameters this strategy accepts. + prepended_conversation_config (PrependedConversationConfig | None): Policy for + prepended conversations. See ``AttackStrategy``. logger (logging.Logger): Logger instance for logging events and messages. """ super().__init__( objective_target=objective_target, context_type=context_type, params_type=params_type, + prepended_conversation_config=prepended_conversation_config, logger=logger, ) @@ -101,26 +109,16 @@ def _rotate_conversation_for_single_turn_target( context: MultiTurnAttackContext[Any], ) -> None: """ - Create a fresh conversation_id for the objective target if it is a single-turn target. - - For single-turn targets, each turn must use a separate conversation_id because the target - rejects conversations with prior messages. The prior turn's conversation_id is recorded - as a PRUNED related conversation on the attack context. - - System messages (e.g., from prepended conversation) are duplicated into the new - conversation so that the target retains its system prompt context. - - For multi-turn targets this method is a no-op. - - This should be called before each turn (except the first) when sending prompts to the - objective target. + Rotate an unseeded single-turn target conversation before later sends. - Args: - context: The current attack context. + An explicit target normalization context already selects only the + persisted seed plus the current request, so rotating in that case would + lose the seed and change the target-facing payload. """ if self._objective_target.configuration.includes(capability=CapabilityName.MULTI_TURN): return - + if context.prepended_history_send_context: + return if context.executed_turns == 0: return @@ -133,25 +131,31 @@ def _rotate_conversation_for_single_turn_target( ) ) - # Duplicate system messages (e.g., system prompt from prepended conversation) - # into the new conversation so the target retains its configuration. memory = CentralMemory.get_memory_instance() messages = memory.get_conversation_messages(conversation_id=old_conversation_id) - system_messages = [m for m in messages if m.api_role == "system"] + system_messages = [message for message in messages if message.api_role == "system"] if system_messages: new_conversation_id, pieces = memory.duplicate_messages(messages=system_messages) memory.add_conversation_to_memory( conversation=Conversation( - conversation_id=new_conversation_id, target_identifier=self._objective_target.get_identifier() + conversation_id=new_conversation_id, + target_identifier=self._objective_target.get_identifier(), ) ) memory.add_message_pieces_to_memory(message_pieces=pieces) context.session.conversation_id = new_conversation_id + persisted_messages = list(memory.get_conversation_messages(conversation_id=new_conversation_id)) + context.prepended_history_send_context = ConversationManager.create_prepended_history_send_context( + target=self._objective_target, + conversation_id=new_conversation_id, + prepended_messages=persisted_messages, + ) else: context.session.conversation_id = str(uuid.uuid4()) self._logger.debug( - f"Rotated conversation_id for single-turn target: " - f"{old_conversation_id} -> {context.session.conversation_id}" + "Rotated conversation_id for single-turn target: %s -> %s", + old_conversation_id, + context.session.conversation_id, ) diff --git a/pyrit/executor/attack/multi_turn/red_teaming.py b/pyrit/executor/attack/multi_turn/red_teaming.py index c9cb196f36..44893bb755 100644 --- a/pyrit/executor/attack/multi_turn/red_teaming.py +++ b/pyrit/executor/attack/multi_turn/red_teaming.py @@ -14,6 +14,7 @@ from pyrit.exceptions import ComponentRole, execution_context from pyrit.executor.attack.component import ( ConversationManager, + PrependedConversationConfig, _AdversarialConversationManager, get_adversarial_chat_messages, ) @@ -93,6 +94,7 @@ def __init__( attack_converter_config: AttackConverterConfig | None = None, attack_scoring_config: AttackScoringConfig | None = None, prompt_normalizer: PromptNormalizer | None = None, + prepended_conversation_config: PrependedConversationConfig | None = None, max_turns: int = 10, score_last_turn_only: bool = False, ) -> None: @@ -105,6 +107,8 @@ def __init__( attack_converter_config: Configuration for attack converters. Defaults to None. attack_scoring_config: Configuration for attack scoring. Defaults to None. prompt_normalizer: The prompt normalizer to use for sending prompts. Defaults to None. + prepended_conversation_config: Configuration for prepended-conversation + converter roles and target-facing formatting. Defaults to None. max_turns (int): Maximum number of turns for the attack. Defaults to 10. score_last_turn_only (bool): If True, only score the final turn instead of every turn. This reduces LLM calls when intermediate scores are not needed (e.g., for @@ -115,7 +119,12 @@ def __init__( ValueError: If objective_scorer is not provided in attack_scoring_config. """ # Initialize base class - super().__init__(objective_target=objective_target, logger=logger, context_type=MultiTurnAttackContext) + super().__init__( + objective_target=objective_target, + logger=logger, + context_type=MultiTurnAttackContext, + prepended_conversation_config=prepended_conversation_config, + ) self._memory = CentralMemory.get_memory_instance() # Initialize converter configuration @@ -169,7 +178,7 @@ def __init__( # Initialize utilities self._prompt_normalizer = prompt_normalizer or PromptNormalizer() - self._conversation_manager = ConversationManager() + self._conversation_manager = ConversationManager(prompt_normalizer=self._prompt_normalizer) # set the maximum number of turns for the attack if max_turns <= 0: @@ -268,6 +277,7 @@ async def _setup_async(self, *, context: MultiTurnAttackContext[Any]) -> None: target=self._objective_target, conversation_id=context.session.conversation_id, request_converters=self._request_converters, + prepended_conversation_config=self._prepended_conversation_config, max_turns=self._max_turns, memory_labels=self._memory_labels, ) @@ -468,7 +478,6 @@ async def _send_prompt_to_objective_target_async( """ logger.info(f"Sending prompt to target: {message.get_value()[:50]}...") - # For single-turn targets, rotate conversation_id so each turn starts fresh self._rotate_conversation_for_single_turn_target(context=context) with execution_context( @@ -485,6 +494,10 @@ async def _send_prompt_to_objective_target_async( request_converter_configurations=self._request_converters, response_converter_configurations=self._response_converters, target=self._objective_target, + normalizer_overrides=self._get_prepended_normalizer_overrides( + prepended_history_send_context=context.prepended_history_send_context, + ), + send_context=context.prepended_history_send_context, ) if response is None: diff --git a/pyrit/executor/attack/multi_turn/tree_of_attacks.py b/pyrit/executor/attack/multi_turn/tree_of_attacks.py index a7eaf72dbd..8ad0bed536 100644 --- a/pyrit/executor/attack/multi_turn/tree_of_attacks.py +++ b/pyrit/executor/attack/multi_turn/tree_of_attacks.py @@ -4,6 +4,7 @@ from __future__ import annotations import asyncio +import copy import enum import logging import uuid @@ -31,6 +32,9 @@ build_conversation_context_string_async, ) from pyrit.executor.attack.component.modality_router import _ModalityFeedbackRouter +from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, +) from pyrit.executor.attack.core.attack_config import ( AttackAdversarialConfig, AttackConverterConfig, @@ -44,7 +48,6 @@ AttackOutcome, AttackResult, ComponentIdentifier, - Conversation, ConversationReference, ConversationType, Message, @@ -54,6 +57,7 @@ ) from pyrit.prompt_normalizer import ConverterConfiguration, PromptNormalizer from pyrit.prompt_target import CapabilityName, PromptTarget +from pyrit.prompt_target.common.target_history import filter_non_replayable_messages from pyrit.prompt_target.common.target_requirements import TargetRequirements from pyrit.score import ( FloatScaleThresholdScorer, @@ -92,6 +96,42 @@ class TAPSystemPromptPaths(enum.Enum): ) +def _validate_stateful_clone_history_compatibility( + *, + objective_target: PromptTarget, + messages: list[Message], +) -> None: + """ + Reject converted history that cannot be replayed into a cloned provider session. + + Raises: + ValueError: If a stateful target cannot preserve converted media while cloning. + """ + if not objective_target.configuration.includes( + capability=CapabilityName.MULTI_TURN + ) or objective_target.configuration.includes(capability=CapabilityName.EDITABLE_HISTORY): + return + + non_text_output_types = { + piece.converted_value_data_type + for message in messages + for piece in message.message_pieces + if piece.converted_value_data_type != "text" + and ( + piece.original_value != piece.converted_value + or piece.original_value_data_type != piece.converted_value_data_type + ) + } + if non_text_output_types: + raise ValueError( + "Tree of Attacks cannot clone a stateful objective-target conversation without editable history " + "when persisted request or response history contains converted non-text output " + f"{sorted(non_text_output_types)}. Copied media cannot be flattened without changing converter " + "role scoping. Use an editable-history target, text-output converters, a stateless target, " + "or branching_factor=1." + ) + + class TAPAttackScoringConfig(AttackScoringConfig): """ Scoring configuration specifically for Tree of Attacks with Pruning (TAP). @@ -350,6 +390,7 @@ def __init__( parent_id: str | None = None, prompt_normalizer: PromptNormalizer | None = None, initial_prompt: Message | None = None, + prepended_conversation_config: PrependedConversationConfig | None = None, ) -> None: """ Initialize a tree node. @@ -379,6 +420,9 @@ def __init__( prompt_normalizer (PromptNormalizer | None): Normalizer for handling prompts and responses. initial_prompt (Message | None): Initial message to send for the first turn, bypassing adversarial chat generation. Supports multimodal messages. + prepended_conversation_config (PrependedConversationConfig | None): + Configuration for prepended-conversation converter roles and + target-facing formatting. """ # Store configuration self._objective_target = objective_target @@ -396,6 +440,7 @@ def __init__( self._attack_strategy_name = attack_strategy_name self._memory_labels = memory_labels or {} self._modality_router = modality_router + self._prepended_conversation_config = prepended_conversation_config or PrependedConversationConfig() self._use_score_as_feedback = use_score_as_feedback # Initialize utilities @@ -421,7 +466,7 @@ def __init__( self.last_prompt_sent: str | None = None self.last_response: Message | None = None self.error_message: str | None = None - + self._prepended_history_send_context: PrependedHistorySendContext | None = None # Context from prepended conversation (for adversarial chat system prompt) self._conversation_context: str | None = None @@ -464,6 +509,8 @@ async def initialize_with_prepended_conversation_async( """ if not prepended_conversation: return + if prepended_conversation_config: + self._prepended_conversation_config = prepended_conversation_config # Use ConversationManager to add messages to memory conversation_manager = ConversationManager( @@ -476,8 +523,16 @@ async def initialize_with_prepended_conversation_async( request_converters=self._request_converters, prepended_conversation_config=prepended_conversation_config, target_identifier=self._objective_target.get_identifier(), + target=self._objective_target, + ) + persisted_messages = list( + self._memory.get_conversation_messages(conversation_id=self.objective_target_conversation_id) + ) + self._prepended_history_send_context = conversation_manager.create_prepended_history_send_context( + target=self._objective_target, + conversation_id=self.objective_target_conversation_id, + prepended_messages=persisted_messages, ) - # Build context string for adversarial chat system prompt (like Crescendo) # The adversarial chat uses this in its system prompt rather than in conversation history self._conversation_context = await build_conversation_context_string_async(prepended_conversation) @@ -609,10 +664,7 @@ async def _send_prompt_to_target_async(self, prompt: str) -> Message: Side Effects: - Sets self.last_response to the target's response text """ - # For single-turn targets, generate a fresh conversation ID before each send - # to ensure the target always receives a clean conversation without prior history. - if not self._objective_target.configuration.includes(capability=CapabilityName.MULTI_TURN): - self.objective_target_conversation_id = str(uuid.uuid4()) + self._rotate_unseeded_single_turn_conversation() # Build the request message via the modality router so prior media (if any) # is included when the objective target accepts it. @@ -636,6 +688,11 @@ async def _send_prompt_to_target_async(self, prompt: str) -> Message: response_converter_configurations=self._response_converters, conversation_id=self.objective_target_conversation_id, target=self._objective_target, + normalizer_overrides=self._prepended_conversation_config.get_normalizer_overrides( + target=self._objective_target, + prepended_history_send_context=self._prepended_history_send_context, + ), + send_context=self._prepended_history_send_context, ) # Store the full response so subsequent turns can forward media when supported. @@ -671,9 +728,7 @@ async def _send_initial_prompt_to_target_async(self) -> Message: if self._initial_prompt is None: raise ValueError("_initial_prompt must be set before calling this method") - # For single-turn targets, generate a fresh conversation ID - if not self._objective_target.configuration.includes(capability=CapabilityName.MULTI_TURN): - self.objective_target_conversation_id = str(uuid.uuid4()) + self._rotate_unseeded_single_turn_conversation() assert self._objective is not None initial_prompt = self._initial_prompt @@ -712,6 +767,11 @@ async def _send_initial_prompt_to_target_async(self) -> Message: response_converter_configurations=self._response_converters, conversation_id=self.objective_target_conversation_id, target=self._objective_target, + normalizer_overrides=self._prepended_conversation_config.get_normalizer_overrides( + target=self._objective_target, + prepended_history_send_context=self._prepended_history_send_context, + ), + send_context=self._prepended_history_send_context, ) # Store the full response so subsequent turns can forward media when supported. @@ -720,6 +780,14 @@ async def _send_initial_prompt_to_target_async(self) -> Message: return response + def _rotate_unseeded_single_turn_conversation(self) -> None: + """Isolate unseeded single-turn sends without discarding an explicit branch boundary.""" + if ( + not self._objective_target.configuration.includes(capability=CapabilityName.MULTI_TURN) + and self._prepended_history_send_context is None + ): + self.objective_target_conversation_id = str(uuid.uuid4()) + async def _score_response_async(self, *, response: Message, objective: str) -> None: """ Score the response from the objective target using the configured scorers. @@ -873,6 +941,13 @@ def duplicate(self) -> _TreeOfAttacksNode: of multiple attack variations from promising nodes. The tree expands by duplicating successful nodes and pruning unsuccessful ones. """ + source_messages = filter_non_replayable_messages( + messages=list(self._memory.get_conversation_messages(conversation_id=self.objective_target_conversation_id)) + ) + _validate_stateful_clone_history_compatibility( + objective_target=self._objective_target, + messages=source_messages, + ) duplicate_node = _TreeOfAttacksNode( objective_target=self._objective_target, adversarial_chat=self._adversarial_chat, @@ -892,30 +967,36 @@ def duplicate(self) -> _TreeOfAttacksNode: desired_response_prefix=self._desired_response_prefix, parent_id=self.node_id, prompt_normalizer=self._prompt_normalizer, + prepended_conversation_config=self._prepended_conversation_config, ) - # Duplicate the conversations to preserve history - # For single-turn targets, duplicate only the system messages (e.g., system prompt - # from prepended conversation) so the target retains its configuration without - # carrying over attack turn history that would cause validation errors. - if self._objective_target.configuration.includes(capability=CapabilityName.MULTI_TURN): - duplicate_node.objective_target_conversation_id = self._memory.duplicate_conversation( - conversation_id=self.objective_target_conversation_id + duplicate_node.objective_target_conversation_id = self._memory.duplicate_conversation( + conversation_id=self.objective_target_conversation_id + ) + duplicated_messages = filter_non_replayable_messages( + messages=list( + self._memory.get_conversation_messages(conversation_id=duplicate_node.objective_target_conversation_id) ) - else: - messages = self._memory.get_conversation_messages(conversation_id=self.objective_target_conversation_id) - system_messages = [m for m in messages if m.api_role == "system"] - if system_messages: - new_id, pieces = self._memory.duplicate_messages(messages=system_messages) - self._memory.add_conversation_to_memory( - conversation=Conversation( - conversation_id=new_id, target_identifier=self._objective_target.get_identifier() - ) + ) + if self._prepended_history_send_context: + duplicate_node._prepended_history_send_context = ( + self._prepended_history_send_context.remap_for_duplicate_conversation( + conversation_id=duplicate_node.objective_target_conversation_id, + source_messages=source_messages, + duplicated_messages=duplicated_messages, ) - self._memory.add_message_pieces_to_memory(message_pieces=pieces) - duplicate_node.objective_target_conversation_id = new_id - else: - duplicate_node.objective_target_conversation_id = str(uuid.uuid4()) + ) + elif ( + duplicated_messages + and self._objective_target.configuration.includes(capability=CapabilityName.MULTI_TURN) + and not self._objective_target.configuration.includes(capability=CapabilityName.EDITABLE_HISTORY) + ): + duplicate_node._prepended_history_send_context = PrependedHistorySendContext( + conversation_id=duplicate_node.objective_target_conversation_id, + seed_message_ids=(), + replay_seed_each_send=False, + bootstrap_message_ids=tuple(message.get_piece().id for message in duplicated_messages), + ) duplicate_node.adversarial_chat_conversation_id = self._memory.duplicate_conversation( conversation_id=self.adversarial_chat_conversation_id @@ -923,6 +1004,7 @@ def duplicate(self) -> _TreeOfAttacksNode: # Copy conversation context for adversarial chat system prompt duplicate_node._conversation_context = self._conversation_context + duplicate_node.last_response = copy.deepcopy(self.last_response) # Copy visualization position so the clone starts from the same tree position duplicate_node._vis_node_id = self._vis_node_id @@ -1454,7 +1536,7 @@ def __init__( batch_size (int): Number of nodes to process in parallel per batch. Defaults to 10. prepended_conversation_config (PrependedConversationConfig | None): Configuration for how to process prepended conversations. Controls converter - application by role, message normalization, and non-chat target behavior. + application by role and request formatting for targets without editable history. Raises: ValueError: If attack_scoring_config uses a non-FloatScaleThresholdScorer objective scorer, @@ -1480,7 +1562,12 @@ def __init__( ) # Initialize base class - super().__init__(objective_target=objective_target, logger=logger, context_type=TAPAttackContext) + super().__init__( + objective_target=objective_target, + logger=logger, + context_type=TAPAttackContext, + prepended_conversation_config=prepended_conversation_config, + ) self._memory = CentralMemory.get_memory_instance() self._node_executor = _TreeOfAttacksNodeExecutor( @@ -1592,9 +1679,6 @@ def __init__( self._prompt_normalizer = prompt_normalizer or PromptNormalizer() - # Store the prepended conversation configuration - self._prepended_conversation_config = prepended_conversation_config - def _load_adversarial_prompts(self) -> None: """Load the adversarial chat prompt template and seed prompt from the default paths.""" # Load prompt template @@ -2110,6 +2194,7 @@ def _create_attack_node( parent_id=parent_id, prompt_normalizer=self._prompt_normalizer, initial_prompt=initial_prompt, + prepended_conversation_config=self._prepended_conversation_config, ) # Add the adversarial chat conversation ID to the context's tracking (ensuring uniqueness) diff --git a/pyrit/executor/attack/single_turn/prompt_sending.py b/pyrit/executor/attack/single_turn/prompt_sending.py index 508b72d924..c37889d7e9 100644 --- a/pyrit/executor/attack/single_turn/prompt_sending.py +++ b/pyrit/executor/attack/single_turn/prompt_sending.py @@ -76,7 +76,9 @@ def __init__( a params type that rejects certain fields. prepended_conversation_config (PrependedConversationConfiguration | None): Configuration for how to process prepended conversations. Controls converter - application by role, message normalization, and non-chat target behavior. + application by role and request formatting for targets without editable history. + Request converters apply to prepended user messages by default; include + ``"assistant"`` explicitly to transform simulated assistant history. Raises: ValueError: If the objective scorer is not a true/false scorer. @@ -87,6 +89,7 @@ def __init__( logger=logger, context_type=SingleTurnAttackContext, params_type=params_type, + prepended_conversation_config=prepended_conversation_config, ) # Initialize the converter configuration @@ -115,9 +118,6 @@ def __init__( self._max_attempts_on_failure = max_attempts_on_failure - # Store the prepended conversation configuration - self._prepended_conversation_config = prepended_conversation_config - def get_attack_scoring_config(self) -> AttackScoringConfig | None: """ Get the attack scoring configuration used by this strategy. @@ -323,6 +323,10 @@ async def _send_prompt_to_objective_target_async( conversation_id=context.conversation_id, request_converter_configurations=self._request_converters, response_converter_configurations=self._response_converters, + normalizer_overrides=self._get_prepended_normalizer_overrides( + prepended_history_send_context=context.prepended_history_send_context, + ), + send_context=context.prepended_history_send_context, ) async def _evaluate_response_async( diff --git a/pyrit/executor/attack/single_turn/single_turn_attack_strategy.py b/pyrit/executor/attack/single_turn/single_turn_attack_strategy.py index 1f699e3654..7dd7c32d17 100644 --- a/pyrit/executor/attack/single_turn/single_turn_attack_strategy.py +++ b/pyrit/executor/attack/single_turn/single_turn_attack_strategy.py @@ -15,6 +15,9 @@ from pyrit.models import AttackResult if TYPE_CHECKING: + from pyrit.executor.attack.component.prepended_conversation_config import ( + PrependedConversationConfig, + ) from pyrit.prompt_target import PromptTarget @@ -48,6 +51,7 @@ def __init__( objective_target: PromptTarget, context_type: type[SingleTurnAttackContext[Any]] = SingleTurnAttackContext, params_type: type[AttackParamsT] = AttackParameters, # type: ignore[ty:invalid-parameter-default] + prepended_conversation_config: PrependedConversationConfig | None = None, logger: logging.Logger = logger, ) -> None: """ @@ -57,11 +61,14 @@ def __init__( objective_target (PromptTarget): The target system to attack. context_type (type[SingleTurnAttackContext]): The type of context this strategy will use. params_type (type[AttackParamsT]): The type of parameters this strategy accepts. + prepended_conversation_config (PrependedConversationConfig | None): Policy for + prepended conversations. See ``AttackStrategy``. logger (logging.Logger): Logger instance for logging events and messages. """ super().__init__( objective_target=objective_target, context_type=context_type, params_type=params_type, + prepended_conversation_config=prepended_conversation_config, logger=logger, ) diff --git a/pyrit/executor/attack/single_turn/skeleton_key.py b/pyrit/executor/attack/single_turn/skeleton_key.py index 5e9d33fbf3..e283153587 100644 --- a/pyrit/executor/attack/single_turn/skeleton_key.py +++ b/pyrit/executor/attack/single_turn/skeleton_key.py @@ -7,6 +7,7 @@ from pyrit.common.apply_defaults import REQUIRED_VALUE, apply_defaults from pyrit.common.path import EXECUTOR_SEED_PROMPT_PATH +from pyrit.executor.attack.component import PrependedConversationConfig from pyrit.executor.attack.core.attack_config import AttackConverterConfig, AttackScoringConfig from pyrit.executor.attack.core.attack_parameters import AttackParameters from pyrit.executor.attack.single_turn.prompt_sending import PromptSendingAttack @@ -59,6 +60,7 @@ def __init__( skeleton_key_prompt: str | None = None, skeleton_key_acceptance: str | None = None, max_attempts_on_failure: int = 0, + prepended_conversation_config: PrependedConversationConfig | None = None, ) -> None: """ Initialize the skeleton key attack strategy. @@ -73,6 +75,10 @@ def __init__( skeleton_key_acceptance (str | None): The simulated assistant acceptance response to prepend. If not provided, uses the default acceptance response. max_attempts_on_failure (int): Maximum number of attempts to retry on failure. + prepended_conversation_config (PrependedConversationConfig | None): Policy for the + skeleton key exchange this attack prepends. Controls which roles receive request + converters and how a target without editable history renders that exchange with + the live request. """ super().__init__( objective_target=objective_target, @@ -81,6 +87,7 @@ def __init__( prompt_normalizer=prompt_normalizer, max_attempts_on_failure=max_attempts_on_failure, params_type=SkeletonKeyAttackParameters, + prepended_conversation_config=prepended_conversation_config, ) self._skeleton_key_prompt = ( diff --git a/pyrit/executor/attack/streaming/barge_in.py b/pyrit/executor/attack/streaming/barge_in.py index 717e269d71..4e892bea41 100644 --- a/pyrit/executor/attack/streaming/barge_in.py +++ b/pyrit/executor/attack/streaming/barge_in.py @@ -28,6 +28,7 @@ if TYPE_CHECKING: from collections.abc import AsyncIterator + from pyrit.executor.attack.component import PrependedConversationConfig from pyrit.prompt_target import PromptTarget, RealtimeTarget logger = logging.getLogger(__name__) @@ -71,6 +72,7 @@ def __init__( attack_converter_config: AttackConverterConfig | None = None, prompt_normalizer: PromptNormalizer | None = None, params_type: type[AttackParamsT] = AttackParameters, # type: ignore[ty:invalid-parameter-default] + prepended_conversation_config: PrependedConversationConfig | None = None, ) -> None: """ Initialize the streaming barge-in attack. @@ -82,6 +84,9 @@ def __init__( prompt_normalizer: Normalizer used to apply converters and persist messages. Defaults to a fresh ``PromptNormalizer``. params_type: Attack parameter dataclass type. + prepended_conversation_config: Configuration for prepended-conversation + converter-role selection. Its message formatter is not used by the + direct realtime streaming path. Raises: ValueError: If ``objective_target`` does not declare the ``STREAMING_AUDIO`` @@ -91,6 +96,7 @@ def __init__( objective_target=objective_target, context_type=BargeInAttackContext, params_type=params_type, + prepended_conversation_config=prepended_conversation_config, logger=logger, ) self._realtime_target = cast("RealtimeTarget", objective_target) @@ -126,16 +132,28 @@ async def _setup_async(self, *, context: BargeInAttackContext[Any]) -> None: Prepended messages are recorded in memory but are NOT pushed into the live realtime session beyond the system prompt — the model only conditions on the system message - and live audio chunks. + and live audio chunks. The direct streaming path does not call + ``PromptTarget.send_prompt_async``, so it cannot consume a target-normalization + context or a prepended-history formatter. """ if not context.conversation_id: context.conversation_id = str(uuid.uuid4()) + existing_message_ids = { + message.get_piece().id for message in self._conversation_manager.get_conversation(context.conversation_id) + } + context.prepended_history_send_context = None await self._conversation_manager.initialize_context_async( context=context, target=self._objective_target, conversation_id=context.conversation_id, request_converters=self._request_converters, + prepended_conversation_config=self._prepended_conversation_config, ) + persisted_messages = self._conversation_manager.get_conversation(context.conversation_id) + context.prepended_conversation = [ + message for message in persisted_messages if message.get_piece().id not in existing_message_ids + ] + context.prepended_history_send_context = None async def _teardown_async(self, *, context: BargeInAttackContext[Any]) -> None: """No-op teardown — connection / dispatcher are closed inside the session's ``run_async``.""" diff --git a/pyrit/message_normalizer/_helpers.py b/pyrit/message_normalizer/_helpers.py index 081ff5b2fe..cbc73e16fd 100644 --- a/pyrit/message_normalizer/_helpers.py +++ b/pyrit/message_normalizer/_helpers.py @@ -2,18 +2,43 @@ # Licensed under the MIT license. """ -Internal helpers shared by squash-style message normalizers. - -Squash normalizers (e.g. ``HistorySquashNormalizer``, -``GenericSystemSquashNormalizer``) collapse several input messages into one -fresh user-role ``Message`` built via ``Message.from_prompt``. Because that -factory creates a brand-new piece with empty ``prompt_metadata``, callers must -explicitly carry request-level metadata (such as the JSON schema key) forward -so downstream normalizers in the pipeline still see it. ``build_squashed_user_message`` -centralizes that propagation rule. +Internal helpers for squash-style message normalizers. + +Some squash normalizers, such as ``GenericSystemSquashNormalizer``, collapse +several input messages into one fresh user-role ``Message`` built via +``Message.from_prompt``. Because that factory creates a brand-new piece with +empty ``prompt_metadata``, callers must explicitly carry request-level metadata +forward so downstream normalizers still see it. +``build_squashed_user_message`` centralizes that propagation rule. """ -from pyrit.models import Message +from pyrit.models import Message, MessagePiece + + +def format_message_piece_for_context(*, piece: MessagePiece) -> str: + """ + Format one message piece for inclusion in textual conversation history. + + Non-text pieces use their context description when available and otherwise + use a modality placeholder, so local asset paths are not exposed as prompt + text. + + Args: + piece (MessagePiece): The piece to represent as context. + + Returns: + str: The textual representation of the piece. + """ + data_type = piece.converted_value_data_type or piece.original_value_data_type + if data_type != "text": + description = piece.prompt_metadata.get("context_description") + if description: + return f"[{data_type.capitalize()} - {description}]" + return f"[{data_type.capitalize()}]" + + if piece.original_value != piece.converted_value: + return f"{piece.converted_value} (original: {piece.original_value})" + return piece.converted_value def build_squashed_user_message(*, new_message_content: str, source_messages: list[Message]) -> Message: diff --git a/pyrit/message_normalizer/conversation_context_normalizer.py b/pyrit/message_normalizer/conversation_context_normalizer.py index a7b7d60a17..b3a7bb78e7 100644 --- a/pyrit/message_normalizer/conversation_context_normalizer.py +++ b/pyrit/message_normalizer/conversation_context_normalizer.py @@ -1,7 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. - +from pyrit.message_normalizer._helpers import format_message_piece_for_context from pyrit.message_normalizer.message_normalizer import MessageStringNormalizer from pyrit.models import Message, MessagePiece @@ -72,19 +72,4 @@ def _format_piece_content(self, piece: MessagePiece) -> str: Returns: The formatted content string. """ - data_type = piece.converted_value_data_type or piece.original_value_data_type - - # For non-text pieces, use metadata description or placeholder - if data_type != "text": - if piece.prompt_metadata and "context_description" in piece.prompt_metadata: - description = piece.prompt_metadata["context_description"] - return f"[{data_type.capitalize()} - {description}]" - return f"[{data_type.capitalize()}]" - - # For text pieces, include both original and converted if different - original = piece.original_value - converted = piece.converted_value - - if original != converted: - return f"{converted} (original: {original})" - return converted + return format_message_piece_for_context(piece=piece) diff --git a/pyrit/message_normalizer/history_squash_normalizer.py b/pyrit/message_normalizer/history_squash_normalizer.py index f0369af682..8fd428c2f8 100644 --- a/pyrit/message_normalizer/history_squash_normalizer.py +++ b/pyrit/message_normalizer/history_squash_normalizer.py @@ -1,64 +1,280 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -from pyrit.message_normalizer._helpers import build_squashed_user_message -from pyrit.message_normalizer.message_normalizer import MessageListNormalizer -from pyrit.models import Message +import copy +import logging +import uuid + +from pyrit.message_normalizer._helpers import format_message_piece_for_context +from pyrit.message_normalizer.conversation_context_normalizer import ConversationContextNormalizer +from pyrit.message_normalizer.generic_system_squash import GenericSystemSquashNormalizer +from pyrit.message_normalizer.message_normalizer import MessageListNormalizer, MessageStringNormalizer +from pyrit.models import Message, MessagePiece + +logger = logging.getLogger(__name__) class HistorySquashNormalizer(MessageListNormalizer[Message]): """ - Squashes a multi-turn conversation into a single user message. + Combine conversation history and the current request into one message. + + The same implementation serves two normalization scopes. Prepended-conversation + flows create per-send overrides with a string formatter and the explicit history + count from the caller-owned prepended-history send context. The ordinary target capability + pipeline uses the default formatter whenever a target does not support multiple + turns. - Previous turns are formatted as labeled context and prepended to the - latest message. Used by the normalization pipeline to adapt prompts - for targets that do not support multi-turn conversations. + The surrounding pipeline controls when the normalizer runs; this class does not + track turn state. In both scopes, historical content becomes text while non-text + pieces from the current request remain separate target-facing pieces. """ + def __init__( + self, + *, + message_normalizer: MessageStringNormalizer | None = None, + expected_history_message_count: int | None = None, + ) -> None: + """ + Initialize the normalizer. + + Args: + message_normalizer (MessageStringNormalizer | None): Optional formatter + for the combined text. When + omitted, use the labeled conversation-history format. + expected_history_message_count (int | None): Optional exact number of + messages expected before the current request. The one-shot + prepended-history scope sets this value; the ordinary capability + pipeline does not. + + Raises: + ValueError: If expected_history_message_count is less than one. + """ + if expected_history_message_count is not None and expected_history_message_count < 1: + raise ValueError("expected_history_message_count must be at least 1") + self._message_normalizer = message_normalizer + self._expected_history_message_count = expected_history_message_count + async def normalize_async(self, messages: list[Message]) -> list[Message]: """ - Combine all messages into a single user message. + Combine history and the current request into one target-facing message. - When there is only one message it is returned unchanged. Otherwise - all prior turns are formatted as ``Role: content`` lines under a - ``[Conversation History]`` header and the last message's content - appears under a ``[Current Message]`` header. + When there is only one message it is returned unchanged. Otherwise, the + configured formatter receives both history and current text. The combined + text replaces the current request's text pieces, while current non-text + pieces retain their original positions. Args: - messages: The conversation messages to squash. + messages (list[Message]): The conversation messages to squash. Returns: - list[Message]: A single-element list containing the squashed message. + list[Message]: A single-element list containing the target-facing message. Raises: - ValueError: If the messages list is empty. + ValueError: If messages is empty, the expected history count does not + match, or converted history cannot be represented as text. """ if not messages: raise ValueError("Messages list cannot be empty") + self._validate_expected_message_count(messages=messages) if len(messages) == 1: return list(messages) - history_lines = self._format_history(messages=messages[:-1]) - current_parts = [piece.converted_value for piece in messages[-1].message_pieces] + history = messages[:-1] + live_request = messages[-1] + self._validate_flattenable_converter_output(messages=history) + self._warn_on_non_text_history(messages=history) - combined = ( - "[Conversation History]\n" + "\n".join(history_lines) + "\n\n[Current Message]\n" + "\n".join(current_parts) - ) + original_view = self._build_original_view(messages=messages) + converted_view = self._build_converted_view(messages=messages) + original_text = await self._normalize_context_async(messages=original_view) + converted_text = original_text + if self._contains_converted_values(messages=messages): + converted_text = await self._normalize_context_async(messages=converted_view) + + return [ + self._build_target_request( + live_request=live_request, + original_text=original_text, + converted_text=converted_text, + ) + ] + + def _validate_expected_message_count(self, *, messages: list[Message]) -> None: + """ + Validate the optional history-boundary contract. + + Args: + messages (list[Message]): History followed by the current request. - return [build_squashed_user_message(new_message_content=combined, source_messages=messages)] + Raises: + ValueError: If the configured history count does not match the input. + """ + if self._expected_history_message_count is None: + return - def _format_history(self, *, messages: list[Message]) -> list[str]: + expected_count = self._expected_history_message_count + 1 + if len(messages) != expected_count: + raise ValueError( + "History squash expected " + f"{self._expected_history_message_count} history messages and one current request, " + f"but received {len(messages)} messages." + ) + + async def _normalize_context_async(self, *, messages: list[Message]) -> str: """ - Format prior messages as ``Role: content`` lines. + Format the text portion of history and the current request. Args: - messages: The history messages to format. + messages (list[Message]): Original-view or converted-view messages to + format. Returns: - list[str]: One line per message piece. + str: The combined target-facing text. """ - lines: list[str] = [] - for msg in messages: - lines.extend(f"{piece.api_role.capitalize()}: {piece.converted_value}" for piece in msg.message_pieces) - return lines + if self._message_normalizer is None: + return self._format_default_context(messages=messages) + + messages_to_normalize = self._filter_live_non_text_pieces(messages=messages) + if isinstance(self._message_normalizer, ConversationContextNormalizer): + messages_to_normalize = await GenericSystemSquashNormalizer().normalize_async(messages_to_normalize) + return await self._message_normalizer.normalize_string_async(messages_to_normalize) + + @staticmethod + def _format_default_context(*, messages: list[Message]) -> str: + """ + Format history with labeled roles and current text in a separate section. + + Args: + messages (list[Message]): History followed by the current request. + + Returns: + str: The default labeled history representation. + """ + history_lines = [ + f"{piece.api_role.capitalize()}: {format_message_piece_for_context(piece=piece)}" + for message in messages[:-1] + for piece in message.message_pieces + ] + current_parts = [ + piece.converted_value for piece in messages[-1].message_pieces if piece.converted_value_data_type == "text" + ] + + sections = ["[Conversation History]\n" + "\n".join(history_lines)] + if current_parts: + sections.append("[Current Message]\n" + "\n".join(current_parts)) + return "\n\n".join(sections) + + @staticmethod + def _filter_live_non_text_pieces(*, messages: list[Message]) -> list[Message]: + filtered = copy.deepcopy(messages) + live_message = filtered[-1] + text_pieces = [piece for piece in live_message.message_pieces if piece.converted_value_data_type == "text"] + if not text_pieces: + filtered.pop() + else: + live_message.message_pieces = text_pieces + return filtered + + @staticmethod + def _build_original_view(*, messages: list[Message]) -> list[Message]: + original_messages = copy.deepcopy(messages) + for message in original_messages: + for piece in message.message_pieces: + piece.converted_value = piece.original_value + piece.converted_value_data_type = piece.original_value_data_type + return original_messages + + @staticmethod + def _build_converted_view(*, messages: list[Message]) -> list[Message]: + converted_messages = copy.deepcopy(messages) + for message in converted_messages: + for piece in message.message_pieces: + piece.original_value = piece.converted_value + piece.original_value_data_type = piece.converted_value_data_type + return converted_messages + + @staticmethod + def _contains_converted_values(*, messages: list[Message]) -> bool: + return any( + piece.converter_identifiers + or piece.original_value != piece.converted_value + or piece.original_value_data_type != piece.converted_value_data_type + for message in messages + for piece in message.message_pieces + ) + + @staticmethod + def _validate_flattenable_converter_output(*, messages: list[Message]) -> None: + output_types = { + piece.converted_value_data_type + for message in messages + for piece in message.message_pieces + if piece.converted_value_data_type != "text" + and ( + piece.original_value != piece.converted_value + or piece.original_value_data_type != piece.converted_value_data_type + ) + } + if output_types: + raise ValueError( + "Cannot flatten conversation history after request converters produced " + f"non-text output types {sorted(output_types)}. Historical conversion must produce text." + ) + + @staticmethod + def _warn_on_non_text_history(*, messages: list[Message]) -> None: + """Warn when native non-text history becomes target-facing text.""" + flattened_types = sorted( + { + piece.converted_value_data_type + for message in messages + for piece in message.message_pieces + if piece.converted_value_data_type != "text" + } + ) + if flattened_types: + logger.warning( + "Conversation history contains non-text pieces %s. History squashing " + "represents them as text placeholders for the target; memory keeps the " + "original pieces.", + flattened_types, + ) + + @staticmethod + def _build_target_request( + *, + live_request: Message, + original_text: str, + converted_text: str, + ) -> Message: + request = copy.deepcopy(live_request) + template_piece = next( + (piece for piece in request.message_pieces if piece.converted_value_data_type == "text"), + request.get_piece(), + ) + text_piece = MessagePiece( + id=uuid.uuid4(), + role=template_piece.role, + original_value=original_text, + converted_value=converted_text, + original_value_data_type="text", + converted_value_data_type="text", + conversation_id=template_piece.conversation_id, + sequence=template_piece.sequence, + prompt_metadata=dict(template_piece.prompt_metadata), + ) + target_pieces: list[MessagePiece] = [] + text_inserted = False + for piece in request.message_pieces: + if piece.converted_value_data_type == "text": + if not text_inserted: + target_pieces.append(text_piece) + text_inserted = True + continue + target_pieces.append(piece) + if not text_inserted: + target_pieces.insert(0, text_piece) + request.message_pieces = target_pieces + return request diff --git a/pyrit/message_normalizer/message_normalizer.py b/pyrit/message_normalizer/message_normalizer.py index 4e15ae4697..dfd941dd6f 100644 --- a/pyrit/message_normalizer/message_normalizer.py +++ b/pyrit/message_normalizer/message_normalizer.py @@ -6,7 +6,7 @@ from pydantic import BaseModel -from pyrit.models import Message +from pyrit.models import ComponentIdentifier, Identifiable, Message # Type alias for system message handling strategies SystemMessageBehavior = Literal["keep", "squash", "ignore"] @@ -38,6 +38,12 @@ async def normalize_async(self, messages: list[Message]) -> list[T]: Returns: A list of normalized items of type T. + + Note: + Output metadata is authoritative. A normalizer that creates replacement + ``Message`` or ``MessagePiece`` objects must preserve any + ``prompt_metadata`` required by downstream consumers. The target stamps + the active conversation ID but does not merge removed metadata back in. """ async def normalize_to_dicts_async(self, messages: list[Message]) -> list[dict[str, Any]]: @@ -57,7 +63,7 @@ async def normalize_to_dicts_async(self, messages: list[Message]) -> list[dict[s return [item.model_dump(exclude_none=True) for item in normalized] -class MessageStringNormalizer(abc.ABC): +class MessageStringNormalizer(Identifiable, abc.ABC): """ Abstract base class for normalizers that return a string representation. @@ -76,6 +82,18 @@ async def normalize_string_async(self, messages: list[Message]) -> str: A string representation of the messages. """ + def _build_identifier(self) -> ComponentIdentifier: + """ + Build the formatter's behavioral identifier. + + Stateless custom formatters receive class identity by default. Formatters + with behavior-changing configuration should override this method. + + Returns: + The formatter's behavioral identifier. + """ + return ComponentIdentifier.of(self) + async def apply_system_message_behavior_async( messages: list[Message], behavior: SystemMessageBehavior diff --git a/pyrit/message_normalizer/tokenizer_template_normalizer.py b/pyrit/message_normalizer/tokenizer_template_normalizer.py index 319af35970..a3bc85d2e1 100644 --- a/pyrit/message_normalizer/tokenizer_template_normalizer.py +++ b/pyrit/message_normalizer/tokenizer_template_normalizer.py @@ -1,5 +1,6 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import json import logging from dataclasses import dataclass from typing import TYPE_CHECKING, ClassVar, Literal, cast @@ -11,7 +12,7 @@ SystemMessageBehavior, apply_system_message_behavior_async, ) -from pyrit.models import Message +from pyrit.models import ComponentIdentifier, Message if TYPE_CHECKING: from transformers import PreTrainedTokenizerBase @@ -228,3 +229,30 @@ async def normalize_string_async(self, messages: list[Message]) -> str: add_generation_prompt=True, ) ) + + def _build_identifier(self) -> ComponentIdentifier: + """ + Build an identifier for the tokenizer template and role behavior. + + Returns: + The tokenizer formatter's behavioral identifier. + """ + tokenizer_type = f"{type(self.tokenizer).__module__}.{type(self.tokenizer).__qualname__}" + return ComponentIdentifier.of( + self, + params={ + "system_message_behavior": self.system_message_behavior, + "tokenizer_type": tokenizer_type, + "tokenizer_name_or_path": str(getattr(self.tokenizer, "name_or_path", "")), + "chat_template": json.dumps( + getattr(self.tokenizer, "chat_template", None), + sort_keys=True, + default=str, + ), + "special_tokens_map": json.dumps( + getattr(self.tokenizer, "special_tokens_map", {}), + sort_keys=True, + default=str, + ), + }, + ) diff --git a/pyrit/prompt_normalizer/normalizer_request.py b/pyrit/prompt_normalizer/normalizer_request.py index cf62e8ebe0..ccd2bf22d3 100644 --- a/pyrit/prompt_normalizer/normalizer_request.py +++ b/pyrit/prompt_normalizer/normalizer_request.py @@ -1,12 +1,16 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +from collections.abc import Mapping from dataclasses import dataclass +from pyrit.message_normalizer import MessageListNormalizer from pyrit.models import Message from pyrit.prompt_normalizer.converter_configuration import ( ConverterConfiguration, ) +from pyrit.prompt_target.common.target_capabilities import CapabilityName +from pyrit.prompt_target.common.target_send_context import TargetSendContext @dataclass @@ -19,6 +23,8 @@ class NormalizerRequest: request_converter_configurations: list[ConverterConfiguration] response_converter_configurations: list[ConverterConfiguration] conversation_id: str | None + normalizer_overrides: dict[CapabilityName, MessageListNormalizer[Message]] + send_context: TargetSendContext | None def __init__( self, @@ -27,6 +33,8 @@ def __init__( request_converter_configurations: list[ConverterConfiguration] | None = None, response_converter_configurations: list[ConverterConfiguration] | None = None, conversation_id: str | None = None, + normalizer_overrides: Mapping[CapabilityName, MessageListNormalizer[Message]] | None = None, + send_context: TargetSendContext | None = None, ) -> None: """ Initialize a normalizer request. @@ -38,6 +46,8 @@ def __init__( response_converter_configurations (list[ConverterConfiguration]): Configurations for converting the response. Defaults to an empty list. conversation_id (str | None): The ID of the conversation. Defaults to None. + normalizer_overrides: Optional per-send target normalizer overrides. + send_context: Optional internal target-send coordination contract. """ if response_converter_configurations is None: response_converter_configurations = [] @@ -47,3 +57,5 @@ def __init__( self.request_converter_configurations = request_converter_configurations self.response_converter_configurations = response_converter_configurations self.conversation_id = conversation_id + self.normalizer_overrides = dict(normalizer_overrides or {}) + self.send_context = send_context diff --git a/pyrit/prompt_normalizer/prompt_normalizer.py b/pyrit/prompt_normalizer/prompt_normalizer.py index 412c71fc71..8c9e6cad54 100644 --- a/pyrit/prompt_normalizer/prompt_normalizer.py +++ b/pyrit/prompt_normalizer/prompt_normalizer.py @@ -8,6 +8,7 @@ import tempfile import traceback import wave +from collections.abc import Mapping from pathlib import Path from typing import Any from uuid import uuid4 @@ -19,6 +20,7 @@ get_execution_context, ) from pyrit.memory import CentralMemory, MemoryInterface, set_message_piece_sha256_async +from pyrit.message_normalizer import MessageListNormalizer from pyrit.models import ( ComponentIdentifier, Conversation, @@ -27,8 +29,9 @@ construct_response_from_request, ) from pyrit.prompt_normalizer import ConverterConfiguration, NormalizerRequest -from pyrit.prompt_target import PromptTarget +from pyrit.prompt_target import CapabilityName, PromptTarget from pyrit.prompt_target.batch_helper import batch_task_async +from pyrit.prompt_target.common.target_send_context import TargetSendContext logger = logging.getLogger(__name__) @@ -71,6 +74,8 @@ async def send_prompt_async( conversation_id: str | None = None, request_converter_configurations: list[ConverterConfiguration] | None = None, response_converter_configurations: list[ConverterConfiguration] | None = None, + normalizer_overrides: Mapping[CapabilityName, MessageListNormalizer[Message]] | None = None, + send_context: TargetSendContext | None = None, ) -> Message: """ Send a single request to a target. @@ -83,6 +88,9 @@ async def send_prompt_async( converting the request. Defaults to an empty list. response_converter_configurations (list[ConverterConfiguration], optional): Configurations for converting the response. Defaults to an empty list. + normalizer_overrides: Optional per-send target normalizer overrides. + send_context: Optional internal coordination contract for caller-owned + history selection and send lifecycle state. Returns: Message: The response received from the target. @@ -108,17 +116,26 @@ async def send_prompt_async( for piece in request.message_pieces: piece.conversation_id = conversation_id - # Apply request converters - await self.convert_values_async(converter_configurations=request_converter_configurations, message=request) + await self.convert_values_async( + converter_configurations=request_converter_configurations, + message=request, + ) await self._calc_hash_async(request=request) responses = None - try: - responses = await target.send_prompt_async(message=request) + responses = await target.send_prompt_async( + message=request, + normalizer_overrides=normalizer_overrides, + send_context=send_context, + ) self.memory.add_message_to_memory(request=request) - except EmptyResponseException: + except EmptyResponseException as ex: + if send_context and not send_context.provider_attempted_by_current_task: + cid = request.message_pieces[0].conversation_id if request.message_pieces else None + raise Exception(f"Error normalizing prompt with conversation ID: {cid}") from ex + # Empty responses are retried, but we don't want them to stop execution self.memory.add_message_to_memory(request=request) @@ -132,6 +149,10 @@ async def send_prompt_async( ] except Exception as ex: + if send_context and not send_context.provider_attempted_by_current_task: + cid = request.message_pieces[0].conversation_id if request.message_pieces else None + raise Exception(f"Error normalizing prompt with conversation ID: {cid}") from ex + # Ensure request to memory before processing exception self.memory.add_message_to_memory(request=request) @@ -207,6 +228,8 @@ async def send_prompt_batch_to_target_async( [request.request_converter_configurations for request in requests], [request.response_converter_configurations for request in requests], [request.conversation_id for request in requests], + [request.normalizer_overrides for request in requests], + [request.send_context for request in requests], ] batch_item_keys = [ @@ -214,9 +237,11 @@ async def send_prompt_batch_to_target_async( "request_converter_configurations", "response_converter_configurations", "conversation_id", + "normalizer_overrides", + "send_context", ] - results: list[Message] = await batch_task_async( + responses: list[Message] = await batch_task_async( prompt_target=target, batch_size=batch_size, items_to_batch=batch_items, @@ -224,7 +249,7 @@ async def send_prompt_batch_to_target_async( task_arguments=batch_item_keys, target=target, ) - return results + return responses async def convert_values_async( self, diff --git a/pyrit/prompt_target/common/conversation_normalization_pipeline.py b/pyrit/prompt_target/common/conversation_normalization_pipeline.py index 37d8b024bb..1b81993ea7 100644 --- a/pyrit/prompt_target/common/conversation_normalization_pipeline.py +++ b/pyrit/prompt_target/common/conversation_normalization_pipeline.py @@ -2,7 +2,7 @@ # Licensed under the MIT license. import logging -from collections.abc import Mapping +from collections.abc import Callable, Mapping from typing import Any from pyrit.message_normalizer import ( @@ -26,10 +26,17 @@ # Single registry: add new normalizable capabilities here and nowhere else. # Order in the list determines pipeline execution order. # --------------------------------------------------------------------------- -_NORMALIZER_REGISTRY: list[tuple[CapabilityName, MessageListNormalizer[Message]]] = [ - (CapabilityName.SYSTEM_PROMPT, GenericSystemSquashNormalizer()), - (CapabilityName.MULTI_TURN, HistorySquashNormalizer()), - (CapabilityName.JSON_SCHEMA, JsonSchemaNormalizer()), +NormalizerFactory = Callable[[TargetCapabilities], MessageListNormalizer[Message]] + +_NORMALIZER_REGISTRY: list[tuple[CapabilityName, NormalizerFactory | None]] = [ + # Editable-history adaptation is intentionally per-send only. Prepended + # conversation flows provide an explicit HistorySquashNormalizer override + # whose boundary comes from persisted message IDs. It must run before + # cardinality-changing target normalizers such as system-message squashing. + (CapabilityName.EDITABLE_HISTORY, None), + (CapabilityName.SYSTEM_PROMPT, lambda _: GenericSystemSquashNormalizer()), + (CapabilityName.MULTI_TURN, lambda _: HistorySquashNormalizer()), + (CapabilityName.JSON_SCHEMA, lambda _: JsonSchemaNormalizer()), ] # Derived constant — no manual maintenance required. @@ -50,7 +57,12 @@ class ConversationNormalizationPipeline: pipeline ordering, and default normalizers are all derived from it. """ - def __init__(self, normalizers: tuple[MessageListNormalizer[Message], ...] = ()) -> None: + def __init__( + self, + normalizers: tuple[MessageListNormalizer[Message], ...] = (), + *, + adapted_capabilities: frozenset[CapabilityName] = frozenset(), + ) -> None: """ Initialize the normalization pipeline with an ordered sequence of normalizers. @@ -58,8 +70,10 @@ def __init__(self, normalizers: tuple[MessageListNormalizer[Message], ...] = ()) normalizers (tuple[MessageListNormalizer[Message], ...]): Ordered normalizers to apply during ``normalize_async``. Defaults to an empty tuple (pass-through). + adapted_capabilities: Capabilities handled by the normalizer sequence. """ self._normalizers = normalizers + self._adapted_capabilities = adapted_capabilities @classmethod def from_capabilities( @@ -75,8 +89,10 @@ def from_capabilities( For each capability in ``_NORMALIZER_REGISTRY`` (in order): * If the target already supports the capability, no normalizer is added. - * If the capability is missing and the policy is ``ADAPT``, the - corresponding normalizer (from overrides or defaults) is added. + * If the capability is missing and an explicit override exists, that + override is added regardless of the target's sparse policy mapping. + * Otherwise, if the policy is ``ADAPT``, the default normalizer is added + when one exists. * If the capability is missing and the policy is ``RAISE``, no normalizer is added (validation is deferred to ``TargetConfiguration.ensure_can_handle()``). @@ -94,11 +110,18 @@ def from_capabilities( """ overrides = normalizer_overrides or {} normalizers: list[MessageListNormalizer[Message]] = [] + adapted_capabilities: set[CapabilityName] = set() - for capability, default_normalizer in _NORMALIZER_REGISTRY: + for capability, default_normalizer_factory in _NORMALIZER_REGISTRY: if capabilities.includes(capability=capability): continue + override = overrides.get(capability) + if override is not None: + normalizers.append(override) + adapted_capabilities.add(capability) + continue + # ``behaviors`` is treated as a sparse mapping: a missing entry means # RAISE (no adaptation; validation deferred to # ``TargetConfiguration.ensure_can_handle``). This keeps the pipeline @@ -111,11 +134,14 @@ def from_capabilities( # Validation is deferred to TargetConfiguration.ensure_can_handle(), # which should be called in the request flow once the full end-to-end # workflow is implemented. - if behavior == UnsupportedCapabilityBehavior.ADAPT: - normalizer = overrides.get(capability, default_normalizer) - normalizers.append(normalizer) + if behavior == UnsupportedCapabilityBehavior.ADAPT and default_normalizer_factory is not None: + normalizers.append(default_normalizer_factory(capabilities)) + adapted_capabilities.add(capability) - return cls(normalizers=tuple(normalizers)) + return cls( + normalizers=tuple(normalizers), + adapted_capabilities=frozenset(adapted_capabilities), + ) async def normalize_async(self, *, messages: list[Message]) -> list[Message]: """ @@ -141,3 +167,7 @@ def normalizers(self) -> tuple[MessageListNormalizer[Message], ...]: tuple[MessageListNormalizer[Message], ...]: The normalizer sequence. """ return self._normalizers + + def has_normalizer_for(self, *, capability: CapabilityName) -> bool: + """Return whether this pipeline adapts the specified capability.""" + return capability in self._adapted_capabilities diff --git a/pyrit/prompt_target/common/prompt_target.py b/pyrit/prompt_target/common/prompt_target.py index aaf918f4cf..f3abfd211b 100644 --- a/pyrit/prompt_target/common/prompt_target.py +++ b/pyrit/prompt_target/common/prompt_target.py @@ -3,9 +3,11 @@ import abc import logging +from collections.abc import Mapping from typing import Any, ClassVar, Literal, final from pyrit.memory import CentralMemory, MemoryInterface +from pyrit.message_normalizer import MessageListNormalizer from pyrit.models import ( ComponentIdentifier, Conversation, @@ -21,6 +23,14 @@ get_known_capabilities, ) from pyrit.prompt_target.common.target_configuration import TargetConfiguration +from pyrit.prompt_target.common.target_history import filter_non_replayable_messages +from pyrit.prompt_target.common.target_send_context import ( + TargetSendContext, + _activate_target_send_context, + _mark_active_provider_attempted, + _reset_target_send_context, +) +from pyrit.prompt_target.common.utils import _marks_provider_attempt logger = logging.getLogger(__name__) @@ -58,6 +68,7 @@ class PromptTarget(Identifiable): # Per-instance overrides are also possible via the ``custom_configuration`` # constructor parameter, which takes precedence over the class-level value. _DEFAULT_CONFIGURATION: TargetConfiguration = TargetConfiguration(capabilities=TargetCapabilities()) + _MANAGES_PROVIDER_ATTEMPT_BOUNDARY: ClassVar[bool] = False # Declarative auth facts consumed by the create-target service and catalog. # Kept off ``TargetCapabilities`` (auth is a construction/credential axis, not @@ -84,6 +95,8 @@ def __init_subclass__(cls, **kwargs: object) -> None: after ``self``. """ super().__init_subclass__(**kwargs) + if "_send_prompt_to_target_async" in cls.__dict__ and "_MANAGES_PROVIDER_ATTEMPT_BOUNDARY" not in cls.__dict__: + cls._MANAGES_PROVIDER_ATTEMPT_BOUNDARY = False # Local import to avoid a circular dependency at package init time. from pyrit.common.brick_contract import enforce_keyword_only_init @@ -132,23 +145,31 @@ def __init__( logging.basicConfig(level=logging.INFO) @final - async def send_prompt_async(self, *, message: Message) -> list[Message]: + async def send_prompt_async( + self, + *, + message: Message, + normalizer_overrides: Mapping[CapabilityName, MessageListNormalizer[Message]] | None = None, + send_context: TargetSendContext | None = None, + ) -> list[Message]: """ Validate, normalize, and send a prompt to the target. This is the public entry point called by the prompt normalizer. It: - 1. Validates the message, fetches the conversation from memory, appends ``message``, and runs - the normalization pipeline (system‑squash, history‑squash, etc.). - 2. Validates the normalized conversation against the target's capabilities. - 3. Delegates to ``_send_prompt_to_target_async`` with the normalized - conversation. + 1. Validates the message. + 2. Loads memory history and runs the target's normalization pipeline. + 3. Validates the normalized conversation against the target's capabilities. + 4. Delegates to ``_send_prompt_to_target_async`` with the normalized conversation. Subclasses MUST NOT override this method. Override ``_send_prompt_to_target_async`` instead. Args: message (Message): The message to send. + normalizer_overrides: Optional per-send target normalizer overrides. + send_context: Optional internal coordination contract for caller-owned + history selection and send lifecycle state. Returns: list[Message]: Response messages from the target. @@ -157,11 +178,42 @@ async def send_prompt_async(self, *, message: Message) -> list[Message]: ValueError: If the message or normalized conversation are empty. """ message.validate() - normalized_conversation = await self._get_normalized_conversation_async(message=message) - if not normalized_conversation: - raise ValueError("Normalization pipeline returned an empty conversation. Cannot send an empty request.") - self._validate_request(normalized_conversation=normalized_conversation) - return await self._send_prompt_to_target_async(normalized_conversation=normalized_conversation) + conversation_id = message.get_piece().conversation_id or "" + if send_context and send_context.conversation_id != conversation_id: + raise ValueError("Target send context conversation_id does not match the current request conversation_id.") + if send_context: + send_context.begin_send() + + try: + normalized_conversation = await self._get_normalized_conversation_async( + message=message, + normalizer_overrides=normalizer_overrides, + send_context=send_context, + ) + if not normalized_conversation: + raise ValueError("Normalization pipeline returned an empty conversation. Cannot send an empty request.") + self._validate_request(normalized_conversation=normalized_conversation) + active_context_token = ( + _activate_target_send_context(send_context=send_context) if send_context is not None else None + ) + try: + target_send = self._send_prompt_to_target_async + target_marks_provider_attempt = self._MANAGES_PROVIDER_ATTEMPT_BOUNDARY or _marks_provider_attempt( + target_send + ) + if send_context and not target_marks_provider_attempt: + send_context.mark_provider_attempted() + return await target_send(normalized_conversation=normalized_conversation) + finally: + if active_context_token is not None: + _reset_target_send_context(token=active_context_token) + finally: + if send_context: + send_context.finish_send() + + def _mark_provider_attempted(self) -> None: + """Notify caller-owned state immediately before irreversible provider I/O.""" + _mark_active_provider_attempted() @abc.abstractmethod async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: @@ -220,85 +272,53 @@ def _validate_request(self, *, normalized_conversation: list[Message]) -> None: if not self.configuration.includes(capability=CapabilityName.MULTI_TURN) and len(normalized_conversation) > 1: raise ValueError(f"This target only supports a single turn conversation. {custom_configuration_message}") - async def _get_normalized_conversation_async(self, *, message: Message) -> list[Message]: + async def _get_normalized_conversation_async( + self, + *, + message: Message, + normalizer_overrides: Mapping[CapabilityName, MessageListNormalizer[Message]] | None = None, + send_context: TargetSendContext | None = None, + ) -> list[Message]: """ - Fetch the conversation from memory, append the current message, and run the - normalization pipeline. + Build the target-facing conversation and run the normalization pipeline. + + Memory history is loaded and the current message is appended before the + target normalization pipeline runs. The original conversation in memory is never mutated. The returned list is an ephemeral copy intended only for building the API request body. - After normalization, the metadata from the original ``message`` is copied - onto the last normalized message so that downstream code (e.g. - ``construct_response_from_request``) propagates the correct - ``conversation_id`` and request lineage to the response. + After normalization, every output piece is stamped with the current + conversation ID. Normalizers own all other output metadata; removed + ``prompt_metadata`` keys are not restored. Args: message (Message): The current message to append. + normalizer_overrides: Optional per-send target normalizer overrides. + send_context: Optional internal coordination contract for caller-approved + persisted history. Returns: list[Message]: The normalized conversation (possibly with system prompt squashed, history squashed, etc.). """ conversation_id = message.message_pieces[0].conversation_id - conversation = ( + persisted_messages = ( list(self._memory.get_conversation_messages(conversation_id=conversation_id)) if conversation_id else [] ) + persisted_messages = filter_non_replayable_messages(messages=persisted_messages) + conversation = send_context.select_history(messages=persisted_messages) if send_context else persisted_messages conversation.append(message) - normalized = await self.configuration.normalize_async(messages=conversation) + normalized = await self.configuration.normalize_async( + messages=conversation, + normalizer_overrides=normalizer_overrides, + ) if normalized: - # Normalizers may create new Message objects (via Message.from_prompt) with - # random conversation_ids. Stamp the correct conversation_id on every - # message (idempotent for originals, fixes new ones). Full lineage is only - # propagated to the last message — it's the one targets use to build the - # response, and earlier messages carry their own legitimate metadata. for msg in normalized: for piece in msg.message_pieces: piece.conversation_id = conversation_id - self._propagate_lineage(source=message, target_message=normalized[-1]) - if len(normalized) > len(conversation): - logger.warning( - "Normalization produced more messages than the input conversation " - "(%d → %d). Only the last normalized message has full lineage " - "metadata. Additional new messages have conversation_id set but " - "require manual lineage updates if needed.", - len(conversation), - len(normalized), - ) return normalized - @staticmethod - def _propagate_lineage(*, source: Message, target_message: Message) -> None: - """ - Copy request-lineage metadata from ``source`` onto every piece in ``target_message``. - - Normalizers may create brand-new ``Message`` objects (e.g. ``HistorySquashNormalizer`` - uses ``Message.from_prompt``) that carry fresh random ``conversation_id`` values and - lack request lineage. This method restores the original metadata so that the response - built from the normalized message stays part of the correct conversation and retains - traceability. - - ``prompt_metadata`` is handled by provenance so that metadata-editing normalizers - are honored. A piece that shares the source piece's ``id`` is the same logical piece - (possibly a copy whose metadata a normalizer intentionally edited or stripped, e.g. - ``JsonSchemaNormalizer``) — its metadata is kept as-is. A piece with a different - ``id`` is brand-new (e.g. a squashed message), so the source's request metadata is - restored, with any keys the normalizer set on the new piece taking precedence. - - Args: - source: The original (pre-normalization) message whose metadata is authoritative. - target_message: The normalized message whose pieces will be updated in place. - """ - source_piece = source.message_pieces[0] - for piece in target_message.message_pieces: - normalized_metadata = dict(piece.prompt_metadata) - is_new_piece = piece.id != source_piece.id - piece.copy_lineage_from(source=source_piece) - if is_new_piece: - piece.prompt_metadata = {**dict(source_piece.prompt_metadata), **normalized_metadata} - else: - piece.prompt_metadata = normalized_metadata - def set_model_name(self, *, model_name: str) -> None: """ Set the model name for this target. @@ -453,6 +473,7 @@ def apply_capabilities(self, *, capabilities: TargetCapabilities) -> None: self._configuration = TargetConfiguration( capabilities=capabilities, policy=self._configuration.policy, + normalizer_overrides=self._configuration.normalizer_overrides, ) @classmethod diff --git a/pyrit/prompt_target/common/target_capabilities.py b/pyrit/prompt_target/common/target_capabilities.py index 382feed1c3..c087e20c7b 100644 --- a/pyrit/prompt_target/common/target_capabilities.py +++ b/pyrit/prompt_target/common/target_capabilities.py @@ -39,7 +39,7 @@ class CapabilityHandlingPolicy: Design invariants ----------------- * The policy is never consulted if the capability is already supported. - * Non-adaptable capabilities (e.g. ``supports_editable_history``) are not + * Non-adaptable capabilities (e.g. ``supports_multi_message_pieces``) are not represented here; requesting them on a target that lacks them always raises immediately. """ @@ -64,7 +64,7 @@ def get_behavior(self, *, capability: CapabilityName) -> UnsupportedCapabilityBe Raises: KeyError: If no behavior exists for the capability. This occurs for - non-adaptable capabilities (e.g., supports_editable_history). + non-adaptable capabilities (e.g., supports_multi_message_pieces). """ try: return self.behaviors[capability] diff --git a/pyrit/prompt_target/common/target_configuration.py b/pyrit/prompt_target/common/target_configuration.py index 6613b9dbb0..c5b19fa259 100644 --- a/pyrit/prompt_target/common/target_configuration.py +++ b/pyrit/prompt_target/common/target_configuration.py @@ -3,6 +3,7 @@ import logging from collections.abc import Mapping +from types import MappingProxyType from typing import Any from pyrit.message_normalizer import MessageListNormalizer @@ -18,7 +19,7 @@ logger = logging.getLogger(__name__) -# Default policy: RAISE on all adaptable capabilities. +# Default policy: preserve each capability's ordinary global behavior. _DEFAULT_POLICY = CapabilityHandlingPolicy() @@ -53,16 +54,17 @@ def __init__( Args: capabilities (TargetCapabilities): The target's declared capabilities. policy (CapabilityHandlingPolicy | None): How to handle each missing - capability. Defaults to RAISE for all adaptable capabilities. + capability. Defaults to each capability's ordinary global behavior. normalizer_overrides (Mapping[CapabilityName, MessageListNormalizer[Any]] | None): Optional overrides for specific capability normalizers. """ self._capabilities = capabilities self._policy = policy or _DEFAULT_POLICY + self._normalizer_overrides = dict(normalizer_overrides or {}) self._pipeline = ConversationNormalizationPipeline.from_capabilities( capabilities=self._capabilities, policy=self._policy, - normalizer_overrides=normalizer_overrides, + normalizer_overrides=self._normalizer_overrides, ) @property @@ -80,6 +82,11 @@ def pipeline(self) -> ConversationNormalizationPipeline: """The resolved normalization pipeline.""" return self._pipeline + @property + def normalizer_overrides(self) -> Mapping[CapabilityName, MessageListNormalizer[Any]]: + """Read-only view of construction-time normalizer overrides.""" + return MappingProxyType(self._normalizer_overrides) + def includes(self, *, capability: CapabilityName) -> bool: """ Check whether the target includes support for the given capability. @@ -110,6 +117,9 @@ def ensure_can_handle(self, *, capability: CapabilityName) -> None: if self._capabilities.includes(capability=capability): return + if self._pipeline.has_normalizer_for(capability=capability): + return + try: behavior = self._policy.get_behavior(capability=capability) except KeyError: @@ -118,18 +128,35 @@ def ensure_can_handle(self, *, capability: CapabilityName) -> None: ) from None if behavior == UnsupportedCapabilityBehavior.RAISE: raise ValueError(f"Target does not support '{capability.value}' and the handling policy is RAISE.") + raise ValueError( + f"Target does not support '{capability.value}', but no default or configured normalizer can adapt it." + ) - async def normalize_async(self, *, messages: list[Message]) -> list[Message]: + async def normalize_async( + self, + *, + messages: list[Message], + normalizer_overrides: Mapping[CapabilityName, MessageListNormalizer[Any]] | None = None, + ) -> list[Message]: """ Run the normalization pipeline over the given messages. Args: messages (list[Message]): The full conversation to normalize. + normalizer_overrides: Per-send replacements for capability normalizers. Returns: list[Message]: The (possibly adapted) message list. """ - return await self._pipeline.normalize_async(messages=messages) + pipeline = self._pipeline + if normalizer_overrides: + merged_overrides = {**self._normalizer_overrides, **normalizer_overrides} + pipeline = ConversationNormalizationPipeline.from_capabilities( + capabilities=self._capabilities, + policy=self._policy, + normalizer_overrides=merged_overrides, + ) + return await pipeline.normalize_async(messages=messages) def as_identifier_params(self) -> dict[str, Any]: """ @@ -161,9 +188,11 @@ def as_identifier_params(self) -> dict[str, Any]: for capability, behavior in self._policy.behaviors.items() if not caps.includes(capability=capability) }, - # Stable, ordered representation of the resolved normalization - # pipeline. Captures the effect of ``normalizer_overrides`` since - # the pipeline is built from defaults + overrides. + # Stable, ordered representation of the pipeline this configuration was + # built with. Per-send ``normalizer_overrides`` are NOT reflected here: + # ``normalize_async`` builds a throwaway pipeline and never mutates + # ``self._pipeline``. Attack-owned overrides are represented by the + # attack identifier instead. "normalization_pipeline": [ f"{type(normalizer).__module__}.{type(normalizer).__qualname__}" for normalizer in self._pipeline.normalizers diff --git a/pyrit/prompt_target/common/target_history.py b/pyrit/prompt_target/common/target_history.py new file mode 100644 index 0000000000..17b1faa0da --- /dev/null +++ b/pyrit/prompt_target/common/target_history.py @@ -0,0 +1,46 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from pyrit.models import Message + + +def filter_non_replayable_messages(*, messages: list[Message]) -> list[Message]: + """ + Remove failed request/error-response pairs from target-facing history. + + ``processing`` and ``unknown`` responses represent failed exchanges, not + provider-authored conversation turns. ``blocked`` and ``empty`` responses + are retained because they are real provider round trips. + + Args: + messages: Persisted messages that may contain failed exchanges. + + Returns: + Messages that are safe to include in a later target-facing payload. + """ + non_replayable_errors = {"processing", "unknown"} + excluded_indexes: set[int] = set() + + for index, error_response in enumerate(messages): + if not any(piece.response_error in non_replayable_errors for piece in error_response.message_pieces): + continue + + excluded_indexes.add(index) + if index == 0: + continue + + request = messages[index - 1] + request_piece = request.get_piece() + error_piece = error_response.get_piece() + is_adjacent_request = ( + request.api_role == "user" + and error_response.api_role == "assistant" + and bool(request_piece.conversation_id) + and request_piece.conversation_id == error_piece.conversation_id + and request_piece.sequence >= 0 + and error_piece.sequence == request_piece.sequence + 1 + ) + if is_adjacent_request: + excluded_indexes.add(index - 1) + + return [message for index, message in enumerate(messages) if index not in excluded_indexes] diff --git a/pyrit/prompt_target/common/target_send_context.py b/pyrit/prompt_target/common/target_send_context.py new file mode 100644 index 0000000000..de589c3374 --- /dev/null +++ b/pyrit/prompt_target/common/target_send_context.py @@ -0,0 +1,65 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from __future__ import annotations + +from contextvars import ContextVar, Token +from typing import TYPE_CHECKING, Protocol + +if TYPE_CHECKING: + from pyrit.models import Message + + +class TargetSendContext(Protocol): + """Internal contract coordinating one target send with caller-owned state.""" + + conversation_id: str + + @property + def provider_attempted_by_current_task(self) -> bool: + """Whether the current task reached provider invocation.""" + ... + + def begin_send(self) -> None: + """Acquire caller-owned state for one complete send.""" + ... + + def select_history(self, *, messages: list[Message]) -> list[Message]: + """Select the caller-approved persisted history for this send.""" + ... + + def mark_provider_attempted(self) -> None: + """Record that provider invocation has begun.""" + ... + + def finish_send(self) -> None: + """Release caller-owned state after the send.""" + ... + + +_ACTIVE_TARGET_SEND_CONTEXT: ContextVar[TargetSendContext | None] = ContextVar( + "_ACTIVE_TARGET_SEND_CONTEXT", + default=None, +) + + +def _activate_target_send_context(*, send_context: TargetSendContext) -> Token[TargetSendContext | None]: + """ + Expose one caller-owned context to target-side provider-boundary helpers. + + Returns: + A token that restores the previous task-local context. + """ + return _ACTIVE_TARGET_SEND_CONTEXT.set(send_context) + + +def _reset_target_send_context(*, token: Token[TargetSendContext | None]) -> None: + """Restore the task-local target-send context after target invocation.""" + _ACTIVE_TARGET_SEND_CONTEXT.reset(token) + + +def _mark_active_provider_attempted() -> None: + """Mark provider invocation for the active target send, if one exists.""" + send_context = _ACTIVE_TARGET_SEND_CONTEXT.get() + if send_context: + send_context.mark_provider_attempted() diff --git a/pyrit/prompt_target/common/utils.py b/pyrit/prompt_target/common/utils.py index a28b164252..25bf593903 100644 --- a/pyrit/prompt_target/common/utils.py +++ b/pyrit/prompt_target/common/utils.py @@ -4,6 +4,7 @@ import asyncio import logging from collections.abc import Callable +from functools import wraps from typing import Any from pyrit.exceptions import PyritException @@ -14,9 +15,17 @@ TokenUsage, construct_response_from_request, ) +from pyrit.prompt_target.common.target_send_context import _mark_active_provider_attempted logger = logging.getLogger(__name__) +_PROVIDER_ATTEMPT_MARKING_WRAPPERS: set[Callable[..., Any]] = set() + + +def _marks_provider_attempt(func: Callable[..., Any]) -> bool: + unbound_func = getattr(func, "__func__", func) + return unbound_func in _PROVIDER_ATTEMPT_MARKING_WRAPPERS + def validate_temperature(temperature: float | None) -> None: """ @@ -58,14 +67,18 @@ def limit_requests_per_minute(func: Callable[..., Any]) -> Callable[..., Any]: Callable: The decorated function with a sleep introduced. """ + @wraps(func) async def set_max_rpm_async(*args: Any, **kwargs: Any) -> Any: self = args[0] rpm = getattr(self, "_max_requests_per_minute", None) if rpm and rpm > 0: await asyncio.sleep(60 / rpm) + if not getattr(self, "_MANAGES_PROVIDER_ATTEMPT_BOUNDARY", False): + _mark_active_provider_attempted() return await func(*args, **kwargs) + _PROVIDER_ATTEMPT_MARKING_WRAPPERS.add(set_max_rpm_async) return set_max_rpm_async diff --git a/pyrit/prompt_target/http_target/httpx_api_target.py b/pyrit/prompt_target/http_target/httpx_api_target.py index 465ae0fa1f..02ba17a0f0 100644 --- a/pyrit/prompt_target/http_target/httpx_api_target.py +++ b/pyrit/prompt_target/http_target/httpx_api_target.py @@ -50,6 +50,7 @@ class HTTPXAPITarget(HTTPTarget): ), ) ) + _MANAGES_PROVIDER_ATTEMPT_BOUNDARY = True def __init__( self, @@ -166,6 +167,7 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me logger.info(f"HTTPXApiTarget: uploading file={filename} via {self.method} to {self.http_url}") + self._mark_provider_attempted() response = await client.request( method=self.method, url=self.http_url, @@ -177,6 +179,7 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me else: # No file upload, handle based on HTTP method logger.info(f"HTTPXApiTarget: sending {self.method} to {self.http_url} with possible JSON/form.") + self._mark_provider_attempted() response = await client.request( method=self.method, url=self.http_url, diff --git a/pyrit/prompt_target/playwright_copilot_target.py b/pyrit/prompt_target/playwright_copilot_target.py index a662e7f1a1..ce937868ed 100644 --- a/pyrit/prompt_target/playwright_copilot_target.py +++ b/pyrit/prompt_target/playwright_copilot_target.py @@ -94,6 +94,7 @@ class PlaywrightCopilotTarget(PromptTarget): ), ) ) + _MANAGES_PROVIDER_ATTEMPT_BOUNDARY = True # Placeholder text constants PLACEHOLDER_GENERATING_RESPONSE: str = "generating response" @@ -258,6 +259,8 @@ async def _interact_with_copilot_async(self, message: Message) -> str | list[tup either as a single text string or a list of (data, data_type) tuples. """ selectors = self._get_selectors() + if any(piece.converted_value_data_type == "text" for piece in message.message_pieces): + await self._clear_text_input_async(input_selector=selectors.input_selector) # Handle multimodal input - process all pieces in the request for piece in message.message_pieces: @@ -290,6 +293,7 @@ async def _wait_for_response_async(self, selectors: CopilotSelectors) -> str | l initial_group_count = len(initial_ai_message_groups) logger.debug(f"Initial message group count before sending: {initial_group_count}") + self._mark_provider_attempted() await self._page.click(selectors.send_button_selector) # Wait for the next AI message to appear @@ -797,6 +801,13 @@ async def _send_text_async(self, *, text: str, input_selector: str) -> None: await self._page.locator(input_selector).click() # Focus first await self._page.locator(input_selector).type(text) + async def _clear_text_input_async(self, *, input_selector: str) -> None: + """Clear locally staged text so a cancelled send can be retried safely.""" + input_locator = self._page.locator(input_selector) + await input_locator.click() + await input_locator.press("ControlOrMeta+A") + await input_locator.press("Backspace") + async def _upload_image_async(self, image_path: str) -> None: """ Handle image upload through Copilot's dropdown interface. @@ -817,6 +828,7 @@ async def _upload_image_async(self, image_path: str) -> None: async with self._page.expect_file_chooser() as fc_info: await add_files_button.click() file_chooser = await fc_info.value + self._mark_provider_attempted() await file_chooser.set_files(image_path) # Check for login requirement in Consumer Copilot diff --git a/pyrit/prompt_target/playwright_target.py b/pyrit/prompt_target/playwright_target.py index 6dcd5d378e..9d29316843 100644 --- a/pyrit/prompt_target/playwright_target.py +++ b/pyrit/prompt_target/playwright_target.py @@ -65,6 +65,7 @@ class PlaywrightTarget(PromptTarget): ), ) ) + _MANAGES_PROVIDER_ATTEMPT_BOUNDARY = True def __init__( self, @@ -118,6 +119,7 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me ) try: + self._mark_provider_attempted() text = await self._interaction_func(self._page, message) except Exception as e: raise RuntimeError(f"An error occurred during interaction: {str(e)}") from e diff --git a/pyrit/prompt_target/websocket_copilot_target.py b/pyrit/prompt_target/websocket_copilot_target.py index 1652e176d7..26ccb78da8 100644 --- a/pyrit/prompt_target/websocket_copilot_target.py +++ b/pyrit/prompt_target/websocket_copilot_target.py @@ -91,6 +91,7 @@ class WebSocketCopilotTarget(PromptTarget): ), ) ) + _MANAGES_PROVIDER_ATTEMPT_BOUNDARY = True def __init__( self, @@ -526,9 +527,10 @@ async def _connect_and_send_async( ) as websocket: for input_msg in inputs: payload = self._dict_to_websocket(input_msg) - await websocket.send(payload) - is_user_input = input_msg.get("type") == CopilotMessageType.USER_PROMPT + if is_user_input: + self._mark_provider_attempted() + await websocket.send(payload) max_message_iterations = 1000 iteration_count = 0 diff --git a/pyrit/prompt_target/websocket_target.py b/pyrit/prompt_target/websocket_target.py index b2ab2be695..ba4cc6cd97 100644 --- a/pyrit/prompt_target/websocket_target.py +++ b/pyrit/prompt_target/websocket_target.py @@ -38,6 +38,7 @@ class WebsocketTarget(PromptTarget): _DEFAULT_CONFIGURATION: TargetConfiguration = TargetConfiguration( capabilities=TargetCapabilities(supports_multi_turn=True) ) + _MANAGES_PROVIDER_ATTEMPT_BOUNDARY = True def __init__( self, @@ -247,6 +248,7 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me conversation_id=conversation_id, conversation_history=normalized_conversation[:-1], ) + self._mark_provider_attempted() result = await self._send_text_async( text=request.converted_value, conversation_id=conversation_id, diff --git a/pyrit/scenario/scenarios/airt/jailbreak.py b/pyrit/scenario/scenarios/airt/jailbreak.py index 710c7eba7c..9bfcc63f19 100644 --- a/pyrit/scenario/scenarios/airt/jailbreak.py +++ b/pyrit/scenario/scenarios/airt/jailbreak.py @@ -107,23 +107,14 @@ def _extra_default_factories() -> dict[str, AttackTechniqueFactory]: @cache def _build_jailbreak_technique() -> type[ScenarioTechnique]: """ - Build the Jailbreak technique class dynamically from every registered factory plus the - scenario-local defaults. - - The technique axis is the set of *attack techniques* a jailbreak is delivered through: the two - default deliveries (``prompt_sending`` and ``jailbreak_system_prompt``) plus whatever techniques - are registered (``role_play_*``, ``many_shot``, ``tap``, …). Jailbreak templates are a separate - selector (``num_jailbreaks`` / ``jailbreak_names``), so only the two deliveries are on by default - — crossing every template with every registered technique explodes quickly. + Build the Jailbreak technique class from its two scenario-owned delivery methods. Returns: type[ScenarioTechnique]: The dynamically generated technique enum class. """ - registry = AttackTechniqueRegistry.get_registry_singleton() - factories = list(registry.get_factories_or_raise().values()) + list(_extra_default_factories().values()) return AttackTechniqueRegistry.build_technique_class_from_factories( # type: ignore[return-value, ty:invalid-return-type] class_name="JailbreakTechnique", - factories=factories, + factories=list(_extra_default_factories().values()), default_names=set(_DEFAULT_TECHNIQUES), ) @@ -136,24 +127,23 @@ class Jailbreak(Scenario): selectors: - **dataset** — the harmful objectives (HarmBench). - - **techniques** — the *attack techniques* each jailbreak is delivered through. Two deliveries - are on by default: ``prompt_sending`` (the template rendered inline into the user message) and + - **techniques** — two delivery methods for each jailbreak: ``prompt_sending`` (the template + rendered inline into the user message) and ``jailbreak_system_prompt`` (the template set as the system prompt with the objective sent as - the user turn). The registry techniques (``role_play_*``, ``many_shot``, ``tap``, …) are - opt-in. + the user turn). - **jailbreaks** — which jailbreak templates to run (a random ``num_jailbreaks`` sample or an explicit ``jailbreak_names`` set). ``prompt_sending`` applies each template as a ``TextJailbreakConverter`` on the outgoing request, - so the objective is rendered inline into the template's ``{{prompt}}`` slot; this keeps that - delivery target-agnostic and lets it compose with every technique. ``jailbreak_system_prompt`` - instead sets the template as a native system prompt and sends the objective as its own user turn, - so it is only built for targets that natively support editable history and system prompts (it is - skipped for incapable targets, or raises if it is the only selected technique). Responses are - scored to determine whether the jailbreak succeeded (non-refusal). + so the objective is rendered inline into the template's ``{{prompt}}`` slot. + ``jailbreak_system_prompt`` instead sets the template as a native system prompt and sends the + objective as its own user turn, so it is only built for targets that natively support editable + history and system prompts (it is skipped for incapable targets, or raises if it is the only + selected technique). Responses are scored to determine whether the jailbreak succeeded + (non-refusal). """ - VERSION: int = 3 + VERSION: int = 4 #: Baseline (an un-jailbroken prompt-send over the objectives) is included by default: a model #: that complies with the bare objective is itself interesting signal. Callers opt out per run @@ -232,6 +222,30 @@ def __init__( scenario_result_id=scenario_result_id, ) + def _resolve_scenario_techniques(self, *, scenario_techniques: Any) -> list[ScenarioTechnique]: + """ + Resolve techniques while rejecting stale or incompatible enum members. + + Args: + scenario_techniques (Any): Requested Jailbreak technique members. + + Returns: + list[ScenarioTechnique]: Compatible concrete techniques. + + Raises: + ValueError: If a caller supplies members from an older or different + technique enum. + """ + if scenario_techniques: + incompatible = [item for item in scenario_techniques if not isinstance(item, self._technique_class)] + if incompatible: + values = [getattr(item, "value", repr(item)) for item in incompatible] + raise ValueError( + "Jailbreak received stale or incompatible techniques " + f"{values}. Select 'prompt_sending' or 'jailbreak_system_prompt'." + ) + return super()._resolve_scenario_techniques(scenario_techniques=scenario_techniques) + def _resolve_templates(self) -> list[str]: """ Resolve the jailbreak templates for this run, replaying the persisted set on resume. @@ -286,13 +300,12 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list """ Build one atomic attack per (technique x jailbreak template x dataset x attempt). - ``prompt_sending`` (and any opt-in registry techniques) deliver each jailbreak template as a - ``TextJailbreakConverter`` appended to that technique's request converters, so the objective - is rendered inline into the template's ``{{prompt}}`` slot on the wire — target-agnostic and - composable with every technique. ``jailbreak_system_prompt`` instead delivers the template as - a native system prompt (no converter) with the objective sent as its own user turn, so it is - only built when the objective target natively supports editable history and system prompts. - Results group by jailbreak template so per-template ASR rolls up naturally. + ``prompt_sending`` delivers each jailbreak template as a ``TextJailbreakConverter`` so the + objective is rendered inline into the template's ``{{prompt}}`` slot on the wire. + ``jailbreak_system_prompt`` instead delivers the template as a native system prompt (no + converter) with the objective sent as its own user turn, so it is only built when the + objective target natively supports editable history and system prompts. Results group by + jailbreak template so per-template ASR rolls up naturally. Args: context (ScenarioContext): The resolved runtime inputs for this run. @@ -314,17 +327,20 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list num_attempts = self.params.get("num_jailbreak_attempts", 1) technique_factories = resolve_technique_factories(context=context, extra_factories=_extra_default_factories()) + selected_names = {technique.value for technique in context.scenario_techniques} + missing = selected_names - set(technique_factories) + if missing: + raise ValueError( + "Jailbreak selected techniques that are no longer available: " + f"{sorted(missing)}. Refresh the plan and select a supported delivery method." + ) - # ``jailbreak_system_prompt`` is delivered separately (native system prompt, no converter); - # every other technique goes through the inline converter path. + prompt_sending_factory = technique_factories.get(_PROMPT_SENDING) system_selected = _JAILBREAK_SYSTEM_PROMPT in technique_factories - converter_factories = { - name: factory for name, factory in technique_factories.items() if name != _JAILBREAK_SYSTEM_PROMPT - } build_system_delivery = system_selected and self._target_supports_system_delivery(self._objective_target) if system_selected and not build_system_delivery: - if not converter_factories: + if prompt_sending_factory is None: raise ValueError( "The 'jailbreak_system_prompt' technique needs a target that natively supports " "editable history and system prompts. Choose a capable target or a different technique." @@ -353,22 +369,21 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list for template_file_name in self._resolved_jailbreaks: template_stem = Path(template_file_name).stem - if converter_factories: + if prompt_sending_factory is not None: jailbreak_converter = TextJailbreakConverter( jailbreak_template=TextJailBreak(template_file_name=template_file_name) ) - # Within the extra-converter stack, apply the jailbreak first (wrap the raw objective - # in the template), then any per-technique converters the caller layered on via - # ``--techniques :converter.*``. (A technique's own built-in converters, if any, - # still run ahead of this extra stack inside the factory.) + # Apply the jailbreak before any caller-supplied prompt_sending converters. technique_converters = { - technique_name: [jailbreak_converter, *self._technique_converters.get(technique_name, [])] - for technique_name in converter_factories + _PROMPT_SENDING: [ + jailbreak_converter, + *self._technique_converters.get(_PROMPT_SENDING, []), + ] } atomic_attacks.extend( self._build_delivery_attacks( builder=builder, - technique_factories=converter_factories, + technique_factories={_PROMPT_SENDING: prompt_sending_factory}, technique_converters=technique_converters, dataset_groups=context.seed_groups_by_dataset, template_stem=template_stem, diff --git a/tests/unit/executor/attack/component/test_conversation_manager.py b/tests/unit/executor/attack/component/test_conversation_manager.py index 550e4b631d..72fe89e431 100644 --- a/tests/unit/executor/attack/component/test_conversation_manager.py +++ b/tests/unit/executor/attack/component/test_conversation_manager.py @@ -17,12 +17,14 @@ - get_prepended_turn_count: Counts assistant messages in a conversation """ +import base64 import uuid from unittest.mock import AsyncMock, MagicMock import pytest from unit.mocks import get_mock_scorer_identifier +from pyrit.converter import Base64Converter, Converter, ConverterResult from pyrit.executor.attack import ConversationManager, ConversationState from pyrit.executor.attack.component import PrependedConversationConfig from pyrit.executor.attack.component.conversation_manager import ( @@ -33,10 +35,10 @@ ) from pyrit.executor.attack.core import AttackContext from pyrit.executor.attack.core.attack_parameters import AttackParameters -from pyrit.message_normalizer import ConversationContextNormalizer -from pyrit.models import ComponentIdentifier, Message, MessagePiece, Score +from pyrit.message_normalizer import ConversationContextNormalizer, HistorySquashNormalizer +from pyrit.models import ComponentIdentifier, Message, MessagePiece, PromptDataType, Score from pyrit.prompt_normalizer import ConverterConfiguration, PromptNormalizer -from pyrit.prompt_target import PromptTarget +from pyrit.prompt_target import CapabilityName, PromptTarget def _mock_target_id(name: str = "MockTarget") -> ComponentIdentifier: @@ -59,6 +61,26 @@ class _TestAttackContext(AttackContext): last_score: Score | None = None +class _ImageOutputConverter(Converter): + """A deterministic text-to-image converter for prepended-history adaptation tests.""" + + SUPPORTED_INPUT_TYPES: tuple[PromptDataType, ...] = ("text",) + SUPPORTED_OUTPUT_TYPES: tuple[PromptDataType, ...] = ("image_path",) + + async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text") -> ConverterResult: + return ConverterResult(output_text="converted.png", output_type="image_path") + + +class _ImageToImageConverter(Converter): + """A deterministic image-to-image converter for lossy adaptation tests.""" + + SUPPORTED_INPUT_TYPES: tuple[PromptDataType, ...] = ("image_path",) + SUPPORTED_OUTPUT_TYPES: tuple[PromptDataType, ...] = ("image_path",) + + async def convert_async(self, *, prompt: str, input_type: PromptDataType = "image_path") -> ConverterResult: + return ConverterResult(output_text="converted.png", output_type="image_path") + + # ============================================================================= # Fixtures # ============================================================================= @@ -710,105 +732,41 @@ async def test_converts_assistant_to_simulated_assistant( assert stored[0].get_piece().role == "simulated_assistant" assert stored[0].get_piece().api_role == "assistant" - async def test_normalizes_for_non_chat_target_by_default( - self, - attack_identifier: ComponentIdentifier, - mock_prompt_target: MagicMock, - sample_conversation: list[Message], - ) -> None: - """Test that prepended conversation is normalized for non-chat targets by default.""" - manager = ConversationManager() - conversation_id = str(uuid.uuid4()) - context = _TestAttackContext(params=AttackParameters(objective="Test objective")) - context.prepended_conversation = sample_conversation - context.next_message = None - - # By default, should normalize (not raise) - matching PrependedConversationConfig field default - await manager.initialize_context_async( - context=context, - target=mock_prompt_target, - conversation_id=conversation_id, - ) - - # next_message should now contain the normalized prepended context - assert context.next_message is not None - text_value = context.next_message.get_piece().original_value - assert len(text_value) > 0 - - async def test_normalizes_for_non_chat_target_when_configured( + async def test_stores_prepended_conversation_for_non_editable_target( self, attack_identifier: ComponentIdentifier, mock_prompt_target: MagicMock, sample_conversation: list[Message], ) -> None: - """Test that non-chat target normalizes prepended conversation when configured.""" manager = ConversationManager() conversation_id = str(uuid.uuid4()) context = _TestAttackContext(params=AttackParameters(objective="Test objective")) context.prepended_conversation = sample_conversation - context.next_message = Message.from_prompt(prompt="Next message", role="user") - - config = PrependedConversationConfig() - await manager.initialize_context_async( + state = await manager.initialize_context_async( context=context, target=mock_prompt_target, conversation_id=conversation_id, - prepended_conversation_config=config, ) - # next_message should now contain the prepended context - assert context.next_message is not None - text_value = context.next_message.get_piece().original_value - assert "Next message" in text_value - assert "Hello" in text_value or "doing well" in text_value - - @pytest.mark.parametrize( - "prepended_conversation_config", - [ - None, - PrependedConversationConfig(message_normalizer=ConversationContextNormalizer()), - ], - ) - async def test_system_prompt_for_non_chat_target_preserves_instruction_and_objective( - self, - attack_identifier: ComponentIdentifier, - mock_prompt_target: MagicMock, - prepended_conversation_config: PrependedConversationConfig | None, - ) -> None: - manager = ConversationManager() - context = _TestAttackContext(params=AttackParameters(objective="Explain saponification")) - context.prepended_conversation = [Message.from_system_prompt("You are a chemistry tutor")] - - await manager.initialize_context_async( - context=context, - target=mock_prompt_target, - conversation_id=str(uuid.uuid4()), - prepended_conversation_config=prepended_conversation_config, - ) - - assert context.next_message is not None - assert context.next_message.get_value() == "Turn 1:\nuser: You are a chemistry tutor\n\nExplain saponification" - - await manager.initialize_context_async( - context=context, - target=mock_prompt_target, - conversation_id=str(uuid.uuid4()), - prepended_conversation_config=prepended_conversation_config, - ) - - assert context.next_message.get_value() == "Turn 1:\nuser: You are a chemistry tutor\n\nExplain saponification" + stored = manager.get_conversation(conversation_id) + assert len(stored) == 2 + assert [message.api_role for message in stored] == ["user", "assistant"] + assert stored[1].get_piece().role == "simulated_assistant" + assert state.turn_count == 1 + assert context.next_message is None - async def test_system_prompt_for_non_chat_target_preserves_supplied_next_message( + async def test_non_editable_target_does_not_rewrite_supplied_next_message( self, attack_identifier: ComponentIdentifier, mock_prompt_target: MagicMock, ) -> None: manager = ConversationManager() + next_message = Message.from_prompt(prompt="Caller-supplied question", role="user") context = _TestAttackContext( params=AttackParameters( objective="Unused objective", - next_message=Message.from_prompt(prompt="Caller-supplied question", role="user"), + next_message=next_message, ) ) context.prepended_conversation = [Message.from_system_prompt("Follow the policy")] @@ -819,48 +777,45 @@ async def test_system_prompt_for_non_chat_target_preserves_supplied_next_message conversation_id=str(uuid.uuid4()), ) - assert context.next_message is not None - assert context.next_message.get_value() == "Turn 1:\nuser: Follow the policy\n\nCaller-supplied question" + assert context.next_message is next_message + assert context.next_message.get_value() == "Caller-supplied question" - async def test_system_prompt_for_non_chat_target_preserves_multimodal_next_message( + async def test_non_editable_target_persists_history_without_using_formatter( self, attack_identifier: ComponentIdentifier, + mock_prompt_normalizer: MagicMock, mock_prompt_target: MagicMock, + sample_conversation: list[Message], ) -> None: - manager = ConversationManager() - image_piece = MessagePiece( - role="user", - original_value="diagram.png", - original_value_data_type="image_path", - ) - context = _TestAttackContext( - params=AttackParameters( - objective="Unused objective", - next_message=Message(message_pieces=[image_piece]), - ) - ) - context.prepended_conversation = [Message.from_system_prompt("Describe images precisely")] + manager = ConversationManager(prompt_normalizer=mock_prompt_normalizer) + conversation_id = str(uuid.uuid4()) + context = _TestAttackContext(params=AttackParameters(objective="Test objective")) + context.prepended_conversation = sample_conversation + message_normalizer = MagicMock(spec=ConversationContextNormalizer) + config = PrependedConversationConfig(message_normalizer=message_normalizer) await manager.initialize_context_async( context=context, target=mock_prompt_target, - conversation_id=str(uuid.uuid4()), + conversation_id=conversation_id, + prepended_conversation_config=config, ) - assert context.next_message is not None - assert len(context.next_message.message_pieces) == 2 - assert context.next_message.message_pieces[0].converted_value == "Turn 1:\nuser: Describe images precisely" - assert context.next_message.message_pieces[1].converted_value == "diagram.png" - assert context.next_message.message_pieces[1].original_value_data_type == "image_path" - - await manager.initialize_context_async( - context=context, - target=mock_prompt_target, - conversation_id=str(uuid.uuid4()), + assert context.prepended_history_send_context is not None + assert context.prepended_history_send_context.conversation_id == conversation_id + stored = manager.get_conversation(conversation_id) + assert len(stored) == len(sample_conversation) + assert context.prepended_history_send_context.seed_message_ids == tuple( + message.get_piece().id for message in stored ) - - assert len(context.next_message.message_pieces) == 2 - assert context.next_message.message_pieces[0].converted_value == "Turn 1:\nuser: Describe images precisely" + normalizer = config.get_normalizer_overrides( + target=mock_prompt_target, + prepended_history_send_context=context.prepended_history_send_context, + )[CapabilityName.EDITABLE_HISTORY] + assert isinstance(normalizer, HistorySquashNormalizer) + assert normalizer._message_normalizer is message_normalizer + assert normalizer._expected_history_message_count == len(sample_conversation) + message_normalizer.normalize_string_async.assert_not_called() async def test_returns_turn_count_for_multi_turn_attacks( self, @@ -959,6 +914,62 @@ async def test_multipart_message_extracts_scores_from_all_pieces( assert score1.id in returned_ids assert score2.id in returned_ids + async def test_scores_come_only_from_the_last_assistant_turn( + self, + attack_identifier: ComponentIdentifier, + mock_chat_target: MagicMock, + ) -> None: + """Only the final assistant turn's scores are surfaced, not every assistant turn.""" + manager = ConversationManager() + conversation_id = str(uuid.uuid4()) + context = _TestAttackContext(params=AttackParameters(objective="Test objective")) + + early_piece = MessagePiece( + role="assistant", + original_value="early reply", + conversation_id=str(uuid.uuid4()), + ) + final_piece = MessagePiece( + role="assistant", + original_value="final reply", + conversation_id=str(uuid.uuid4()), + ) + manager._memory.add_message_pieces_to_memory(message_pieces=[early_piece, final_piece]) + + def _false_score(piece: MessagePiece, rationale: str) -> Score: + return Score( + score_type="true_false", + score_value="false", + score_category=["test"], + score_value_description=rationale, + score_rationale=rationale, + score_metadata={}, + message_piece_id=str(piece.id), + scorer_class_identifier=get_mock_scorer_identifier(), + ) + + early_score = _false_score(early_piece, "early") + final_score = _false_score(final_piece, "final") + manager._memory.add_scores_to_memory(scores=[early_score, final_score]) + + context.prepended_conversation = [ + Message.from_prompt(prompt="first ask", role="user"), + Message(message_pieces=[early_piece]), + Message.from_prompt(prompt="second ask", role="user"), + Message(message_pieces=[final_piece]), + ] + + state = await manager.initialize_context_async( + context=context, + target=mock_chat_target, + conversation_id=conversation_id, + max_turns=10, + ) + + assert [score.id for score in state.last_assistant_message_scores] == [final_score.id] + assert context.last_score is not None + assert context.last_score.id == final_score.id + async def test_prepended_conversation_ignores_true_scores( self, attack_identifier: ComponentIdentifier, @@ -1064,127 +1075,261 @@ async def test_prepended_conversation_ignores_true_scores( class TestPrependedConversationConfigSettings: """Tests for PrependedConversationConfig settings in initialize_context_async.""" - # ------------------------------------------------------------------------- - # non_chat_target_behavior Tests - # ------------------------------------------------------------------------- - - async def test_non_chat_target_behavior_normalize_is_default( + async def test_non_editable_target_converts_selected_roles_before_storage( self, attack_identifier: ComponentIdentifier, mock_prompt_target: MagicMock, sample_conversation: list[Message], ) -> None: - """Test that non-chat targets normalize by default (no config), matching dataclass field default.""" manager = ConversationManager() - conversation_id = str(uuid.uuid4()) context = _TestAttackContext(params=AttackParameters(objective="Test objective")) context.prepended_conversation = sample_conversation - context.next_message = None + context.next_message = Message.from_prompt(prompt="live request", role="user") + converter_config = ConverterConfiguration.from_converters(converters=[Base64Converter()]) - # Should normalize by default (matching PrependedConversationConfig field default) + conversation_id = str(uuid.uuid4()) await manager.initialize_context_async( context=context, target=mock_prompt_target, conversation_id=conversation_id, + request_converters=converter_config, ) - # next_message should contain normalized context - assert context.next_message is not None - text_value = context.next_message.get_piece().original_value - assert len(text_value) > 0 + stored = manager.get_conversation(conversation_id) + encoded_user = base64.b64encode(b"Hello, how are you?").decode() + assert stored[0].get_piece().converted_value == encoded_user + assert stored[1].get_piece().converted_value == "I'm doing well, thank you!" + assert context.next_message.get_piece().converted_value == "live request" - async def test_non_chat_target_behavior_normalize_first_turn_creates_next_message( + async def test_non_editable_target_converts_assistant_history_only_when_opted_in( self, attack_identifier: ComponentIdentifier, mock_prompt_target: MagicMock, sample_conversation: list[Message], ) -> None: - """Test that normalize_first_turn creates next_message when none exists.""" manager = ConversationManager() - conversation_id = str(uuid.uuid4()) context = _TestAttackContext(params=AttackParameters(objective="Test objective")) context.prepended_conversation = sample_conversation - context.next_message = None - - config = PrependedConversationConfig() + context.next_message = Message.from_prompt(prompt="live request", role="user") + converter_config = ConverterConfiguration.from_converters(converters=[Base64Converter()]) + config = PrependedConversationConfig(apply_converters_to_roles=["assistant"]) + conversation_id = str(uuid.uuid4()) await manager.initialize_context_async( context=context, target=mock_prompt_target, conversation_id=conversation_id, + request_converters=converter_config, prepended_conversation_config=config, ) - # Should have created a next_message with the normalized context - assert context.next_message is not None - text_value = context.next_message.get_piece().original_value - assert len(text_value) > 0 + stored = manager.get_conversation(conversation_id) + encoded_assistant = base64.b64encode(b"I'm doing well, thank you!").decode() + assert stored[0].get_piece().converted_value == "Hello, how are you?" + assert stored[1].get_piece().converted_value == encoded_assistant - async def test_non_chat_target_behavior_normalize_first_turn_prepends_to_existing_message( + async def test_non_editable_target_rejects_non_text_output_from_current_converter( self, attack_identifier: ComponentIdentifier, mock_prompt_target: MagicMock, sample_conversation: list[Message], ) -> None: - """Test that normalize_first_turn prepends context to existing next_message.""" manager = ConversationManager() - conversation_id = str(uuid.uuid4()) context = _TestAttackContext(params=AttackParameters(objective="Test objective")) context.prepended_conversation = sample_conversation - context.next_message = Message.from_prompt(prompt="My question", role="user") + converter_config = ConverterConfiguration.from_converters(converters=[_ImageOutputConverter()]) - config = PrependedConversationConfig() + with pytest.raises(ValueError, match="non-text output types.*image_path"): + await manager.initialize_context_async( + context=context, + target=mock_prompt_target, + conversation_id=str(uuid.uuid4()), + request_converters=converter_config, + ) + + assert sample_conversation[0].get_piece().converted_value_data_type == "text" + + async def test_non_editable_target_rejects_same_modality_non_text_output( + self, + attack_identifier: ComponentIdentifier, + mock_prompt_target: MagicMock, + ) -> None: + manager = ConversationManager() + context = _TestAttackContext(params=AttackParameters(objective="Test objective")) + context.prepended_conversation = [ + Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="original.png", + converted_value="original.png", + original_value_data_type="image_path", + converted_value_data_type="image_path", + conversation_id="seed", + ) + ] + ) + ] + converter_config = ConverterConfiguration.from_converters(converters=[_ImageToImageConverter()]) + + with pytest.raises(ValueError, match="non-text output types.*image_path"): + await manager.initialize_context_async( + context=context, + target=mock_prompt_target, + conversation_id=str(uuid.uuid4()), + request_converters=converter_config, + ) + + async def test_prepended_conversion_failure_does_not_partially_write_history( + self, + attack_identifier: ComponentIdentifier, + mock_prompt_normalizer: MagicMock, + mock_prompt_target: MagicMock, + sample_conversation: list[Message], + ) -> None: + manager = ConversationManager(prompt_normalizer=mock_prompt_normalizer) + context = _TestAttackContext(params=AttackParameters(objective="Test objective")) + context.prepended_conversation = sample_conversation + mock_prompt_normalizer.convert_values_async.side_effect = [None, ValueError("second conversion failed")] + config = PrependedConversationConfig(apply_converters_to_roles=["user", "assistant"]) + conversation_id = str(uuid.uuid4()) + + with pytest.raises(ValueError, match="second conversion failed"): + await manager.initialize_context_async( + context=context, + target=mock_prompt_target, + conversation_id=conversation_id, + request_converters=[ConverterConfiguration(converters=[])], + prepended_conversation_config=config, + ) + + assert manager.get_conversation(conversation_id) == [] + + async def test_non_persisted_prepended_message_is_not_counted_in_context( + self, + attack_identifier: ComponentIdentifier, + mock_prompt_target: MagicMock, + ) -> None: + manager = ConversationManager() + context = _TestAttackContext(params=AttackParameters(objective="Test objective")) + piece = MessagePiece( + role="user", + original_value="ephemeral", + conversation_id="seed", + ) + piece.not_in_memory = True + context.prepended_conversation = [Message(message_pieces=[piece])] + conversation_id = str(uuid.uuid4()) await manager.initialize_context_async( context=context, target=mock_prompt_target, conversation_id=conversation_id, - prepended_conversation_config=config, ) - # Should have prepended context to existing message - text_value = context.next_message.get_piece().original_value - assert "My question" in text_value - # Context should come before the original question - question_index = text_value.find("My question") - assert question_index > 0 # Context should be prepended + assert manager.get_conversation(conversation_id) == [] - async def test_non_chat_target_behavior_normalize_returns_empty_state( + async def test_non_persisted_piece_does_not_constrain_flattening( self, attack_identifier: ComponentIdentifier, mock_prompt_target: MagicMock, - sample_conversation: list[Message], ) -> None: - """Test that normalize_first_turn returns empty ConversationState (no turn tracking).""" manager = ConversationManager() - conversation_id = str(uuid.uuid4()) context = _TestAttackContext(params=AttackParameters(objective="Test objective")) - context.prepended_conversation = sample_conversation + ephemeral_piece = MessagePiece( + role="user", + original_value="ephemeral", + conversation_id="seed", + sequence=0, + ) + ephemeral_piece.not_in_memory = True + context.prepended_conversation = [ + Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="persisted", + conversation_id="seed", + sequence=0, + ), + ephemeral_piece, + ] + ) + ] + conversation_id = str(uuid.uuid4()) - config = PrependedConversationConfig() + await manager.initialize_context_async( + context=context, + target=mock_prompt_target, + conversation_id=conversation_id, + request_converters=[ + ConverterConfiguration( + converters=[_ImageOutputConverter()], + indexes_to_apply=[1], + ) + ], + ) - state = await manager.initialize_context_async( + stored = manager.get_conversation(conversation_id) + assert len(stored) == 1 + assert [piece.converted_value for piece in stored[0].message_pieces] == ["persisted"] + + async def test_non_editable_target_preserves_converter_piece_indexes( + self, + attack_identifier: ComponentIdentifier, + mock_prompt_target: MagicMock, + ) -> None: + manager = ConversationManager() + context = _TestAttackContext(params=AttackParameters(objective="Test objective")) + context.prepended_conversation = [ + Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="first piece", + conversation_id="seed", + sequence=0, + ), + MessagePiece( + role="user", + original_value="second piece", + conversation_id="seed", + sequence=0, + ), + ] + ) + ] + context.next_message = Message.from_prompt(prompt="live request", role="user") + converter_config = [ + ConverterConfiguration( + converters=[Base64Converter()], + indexes_to_apply=[0], + ) + ] + + conversation_id = str(uuid.uuid4()) + await manager.initialize_context_async( context=context, target=mock_prompt_target, conversation_id=conversation_id, - prepended_conversation_config=config, + request_converters=converter_config, ) - # Non-chat targets don't track turns - assert state.turn_count == 0 - assert state.last_assistant_message_scores == [] + stored_pieces = manager.get_conversation(conversation_id)[0].message_pieces + assert stored_pieces[0].converted_value == base64.b64encode(b"first piece").decode() + assert stored_pieces[1].converted_value == "second piece" # ------------------------------------------------------------------------- # apply_converters_to_roles Tests # ------------------------------------------------------------------------- - async def test_apply_converters_to_roles_default_applies_to_all( + async def test_apply_converters_to_roles_default_applies_to_user_only( self, attack_identifier: ComponentIdentifier, mock_chat_target: MagicMock, sample_conversation: list[Message], ) -> None: - """Test that converters are applied to all roles by default.""" + """Test that converters are applied only to user history by default.""" mock_normalizer = MagicMock(spec=PromptNormalizer) mock_normalizer.convert_values_async = AsyncMock() manager = ConversationManager(prompt_normalizer=mock_normalizer) @@ -1201,8 +1346,7 @@ async def test_apply_converters_to_roles_default_applies_to_all( request_converters=converter_config, ) - # convert_values_async should be called for each message (both user and assistant) - assert mock_normalizer.convert_values_async.call_count == 2 + mock_normalizer.convert_values_async.assert_awaited_once() async def test_apply_converters_to_roles_user_only( self, @@ -1295,66 +1439,42 @@ async def test_apply_converters_to_roles_empty_list_skips_all( async def test_message_normalizer_default_uses_conversation_context_normalizer( self, attack_identifier: ComponentIdentifier, + mock_prompt_normalizer: MagicMock, mock_prompt_target: MagicMock, sample_conversation: list[Message], ) -> None: - """Test that default normalizer produces Turn N format.""" - manager = ConversationManager() + """Test that default formatting is supplied by an explicit per-send override.""" + manager = ConversationManager(prompt_normalizer=mock_prompt_normalizer) conversation_id = str(uuid.uuid4()) context = _TestAttackContext(params=AttackParameters(objective="Test objective")) context.prepended_conversation = sample_conversation - context.next_message = None - - config = PrependedConversationConfig() await manager.initialize_context_async( context=context, target=mock_prompt_target, conversation_id=conversation_id, - prepended_conversation_config=config, ) - # Default ConversationContextNormalizer produces "Turn N:" format - assert context.next_message is not None - text_value = context.next_message.get_piece().original_value - assert "Turn 1" in text_value or "turn 1" in text_value.lower() + assert context.prepended_history_send_context is not None + normalizer = PrependedConversationConfig().get_normalizer_overrides( + target=mock_prompt_target, + prepended_history_send_context=context.prepended_history_send_context, + )[CapabilityName.EDITABLE_HISTORY] + assert isinstance(normalizer, HistorySquashNormalizer) + assert isinstance(normalizer._message_normalizer, ConversationContextNormalizer) - async def test_message_normalizer_custom_normalizer_is_used( + def test_message_normalizer_is_not_overridden_for_editable_target( self, - attack_identifier: ComponentIdentifier, - mock_prompt_target: MagicMock, - sample_conversation: list[Message], + mock_chat_target: MagicMock, ) -> None: - """Test that custom message_normalizer is used when provided.""" - from pyrit.message_normalizer import MessageStringNormalizer + mock_chat_target.configuration.includes.return_value = True - # Create a mock normalizer that returns a specific format - mock_normalizer = MagicMock(spec=MessageStringNormalizer) - mock_normalizer.normalize_string_async = AsyncMock(return_value="CUSTOM_FORMAT: test content") - - manager = ConversationManager() - conversation_id = str(uuid.uuid4()) - context = _TestAttackContext(params=AttackParameters(objective="Test objective")) - context.prepended_conversation = sample_conversation - context.next_message = None - - config = PrependedConversationConfig( - message_normalizer=mock_normalizer, - ) - - await manager.initialize_context_async( - context=context, - target=mock_prompt_target, - conversation_id=conversation_id, - prepended_conversation_config=config, + overrides = PrependedConversationConfig().get_normalizer_overrides( + target=mock_chat_target, + prepended_history_send_context=None, ) - # Verify custom normalizer was called - mock_normalizer.normalize_string_async.assert_called_once() - # Verify the custom format is in the message - assert context.next_message is not None - text_value = context.next_message.get_piece().original_value - assert "CUSTOM_FORMAT: test content" in text_value + assert overrides == {} # ------------------------------------------------------------------------- # Chat Target Behavior (Config has no effect) diff --git a/tests/unit/executor/attack/component/test_prepended_conversation_config.py b/tests/unit/executor/attack/component/test_prepended_conversation_config.py index b1c34a6770..546a7f1dbf 100644 --- a/tests/unit/executor/attack/component/test_prepended_conversation_config.py +++ b/tests/unit/executor/attack/component/test_prepended_conversation_config.py @@ -1,17 +1,15 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -from typing import get_args from unittest.mock import MagicMock from pyrit.executor.attack.component.prepended_conversation_config import PrependedConversationConfig from pyrit.message_normalizer import ConversationContextNormalizer -from pyrit.models import ChatMessageRole -def test_default_init_apply_converters_to_all_roles(): +def test_default_init_apply_converters_to_user_role(): config = PrependedConversationConfig() - assert config.apply_converters_to_roles == list(get_args(ChatMessageRole)) + assert config.apply_converters_to_roles == ["user"] def test_default_init_message_normalizer_is_none(): diff --git a/tests/unit/executor/attack/component/test_prepended_history_send_context.py b/tests/unit/executor/attack/component/test_prepended_history_send_context.py new file mode 100644 index 0000000000..43561ce4e5 --- /dev/null +++ b/tests/unit/executor/attack/component/test_prepended_history_send_context.py @@ -0,0 +1,222 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import uuid + +import pytest + +from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, +) +from pyrit.models import ChatMessageRole, Message + + +def _message( + *, + role: ChatMessageRole, + value: str, + sequence: int, + conversation_id: str = "conversation", +) -> Message: + message = Message.from_prompt(prompt=value, role=role) + piece = message.get_piece() + piece.sequence = sequence + piece.conversation_id = conversation_id + return message + + +def test_context_requires_unique_persisted_seed_ids() -> None: + message_id = uuid.uuid4() + + with pytest.raises(ValueError, match="conversation_id"): + PrependedHistorySendContext( + conversation_id="", + seed_message_ids=(message_id,), + replay_seed_each_send=False, + ) + with pytest.raises(ValueError, match="seed_message_ids"): + PrependedHistorySendContext( + conversation_id="conversation", + seed_message_ids=(), + replay_seed_each_send=False, + ) + with pytest.raises(ValueError, match="unique"): + PrependedHistorySendContext( + conversation_id="conversation", + seed_message_ids=(message_id, message_id), + replay_seed_each_send=False, + ) + + +def test_stateful_context_consumes_explicit_boundary_at_provider_attempt() -> None: + first = _message(role="system", value="system", sequence=0) + second = _message(role="user", value="seed", sequence=1) + unrelated = _message(role="assistant", value="later response", sequence=2) + context = PrependedHistorySendContext( + conversation_id="conversation", + seed_message_ids=(first.get_piece().id, second.get_piece().id), + replay_seed_each_send=False, + ) + + context.begin_send() + selected = context.select_history(messages=[unrelated, second, first]) + context.mark_provider_attempted() + context.finish_send() + + assert selected == [first, second] + assert context.is_seed_consumed + assert context.provider_attempt_count == 1 + + context.begin_send() + assert context.select_history(messages=[first, second, unrelated]) == [first, second, unrelated] + context.finish_send() + + +def test_stateless_context_replays_explicit_boundary_after_provider_attempt() -> None: + seed = _message(role="user", value="seed", sequence=0) + context = PrependedHistorySendContext( + conversation_id="conversation", + seed_message_ids=(seed.get_piece().id,), + replay_seed_each_send=True, + ) + + for _ in range(2): + context.begin_send() + assert context.select_history(messages=[seed]) == [seed] + context.mark_provider_attempted() + context.finish_send() + + assert not context.is_seed_consumed + assert context.provider_attempt_count == 2 + + +def test_context_rejects_concurrent_send_until_active_send_finishes() -> None: + context = PrependedHistorySendContext( + conversation_id="conversation", + seed_message_ids=(uuid.uuid4(),), + replay_seed_each_send=False, + ) + + context.begin_send() + with pytest.raises(RuntimeError, match="Concurrent sends"): + context.begin_send() + + context.finish_send() + context.begin_send() + context.finish_send() + + +def test_context_rejects_provider_attempt_without_active_send() -> None: + context = PrependedHistorySendContext( + conversation_id="conversation", + seed_message_ids=(uuid.uuid4(),), + replay_seed_each_send=False, + ) + + with pytest.raises(RuntimeError, match="without an active send"): + context.mark_provider_attempted() + + +def test_context_counts_one_provider_attempt_per_send_task() -> None: + context = PrependedHistorySendContext( + conversation_id="conversation", + seed_message_ids=(uuid.uuid4(),), + replay_seed_each_send=False, + ) + + context.begin_send() + context.mark_provider_attempted() + context.mark_provider_attempted() + context.finish_send() + + assert context.provider_attempt_count == 1 + + +def test_context_rejects_missing_persisted_boundary_message() -> None: + context = PrependedHistorySendContext( + conversation_id="conversation", + seed_message_ids=(uuid.uuid4(),), + replay_seed_each_send=False, + ) + + with pytest.raises(ValueError, match="Missing 1 message"): + context.select_history(messages=[]) + + +def test_context_remaps_only_explicit_seed_for_duplicate_conversation() -> None: + seed = _message(role="system", value="seed", sequence=0) + live_request = _message(role="user", value="live", sequence=1) + source_messages = [seed, live_request] + duplicated_messages = [message.duplicate() for message in source_messages] + for message in duplicated_messages: + for piece in message.message_pieces: + piece.conversation_id = "duplicate" + + context = PrependedHistorySendContext( + conversation_id="conversation", + seed_message_ids=(seed.get_piece().id,), + replay_seed_each_send=True, + ) + + duplicate = context.remap_for_duplicate_conversation( + conversation_id="duplicate", + source_messages=source_messages, + duplicated_messages=duplicated_messages, + ) + + assert duplicate.seed_message_ids == (duplicated_messages[0].get_piece().id,) + assert duplicate.bootstrap_message_ids == (duplicated_messages[0].get_piece().id,) + assert duplicated_messages[1].get_piece().id not in duplicate.seed_message_ids + assert duplicate.select_history(messages=duplicated_messages) == [duplicated_messages[0]] + + +def test_stateful_context_bootstraps_full_duplicated_branch() -> None: + seed = _message(role="system", value="seed", sequence=0) + live_request = _message(role="user", value="live", sequence=1) + response = _message(role="assistant", value="response", sequence=2) + source_messages = [seed, live_request, response] + duplicated_messages = [ + _message(role=message.api_role, value=message.get_value(), sequence=message.sequence) + for message in source_messages + ] + for message in duplicated_messages: + message.get_piece().conversation_id = "duplicate" + + context = PrependedHistorySendContext( + conversation_id="source", + seed_message_ids=(seed.get_piece().id,), + replay_seed_each_send=False, + ) + + duplicate = context.remap_for_duplicate_conversation( + conversation_id="duplicate", + source_messages=source_messages, + duplicated_messages=duplicated_messages, + ) + + assert duplicate.seed_message_ids == (duplicated_messages[0].get_piece().id,) + assert duplicate.bootstrap_message_ids == tuple(message.get_piece().id for message in duplicated_messages) + assert duplicate.select_history(messages=duplicated_messages) == duplicated_messages + + +def test_context_remap_resets_consumed_state_for_new_conversation() -> None: + seed = _message(role="user", value="seed", sequence=0) + duplicated_seed = seed.duplicate() + duplicated_seed.get_piece().conversation_id = "duplicate" + context = PrependedHistorySendContext( + conversation_id="conversation", + seed_message_ids=(seed.get_piece().id,), + replay_seed_each_send=False, + ) + context.begin_send() + context.mark_provider_attempted() + context.finish_send() + + duplicate = context.remap_for_duplicate_conversation( + conversation_id="duplicate", + source_messages=[seed], + duplicated_messages=[duplicated_seed], + ) + + assert not duplicate.is_seed_consumed + assert duplicate.select_history(messages=[duplicated_seed]) == [duplicated_seed] diff --git a/tests/unit/executor/attack/core/test_attack_strategy_prepended_policy.py b/tests/unit/executor/attack/core/test_attack_strategy_prepended_policy.py new file mode 100644 index 0000000000..e812d4cf25 --- /dev/null +++ b/tests/unit/executor/attack/core/test_attack_strategy_prepended_policy.py @@ -0,0 +1,109 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Tests for prepended-conversation policy owned by ``AttackStrategy``.""" + +import inspect +import uuid +from typing import Any + +import pytest + +from pyrit.executor.attack.component import PrependedConversationConfig +from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, +) +from pyrit.executor.attack.core import AttackStrategy +from pyrit.executor.attack.core.attack_config import AttackScoringConfig +from pyrit.executor.attack.multi_turn.chunked_request import ChunkedRequestAttack +from pyrit.executor.attack.multi_turn.crescendo import CrescendoAttack +from pyrit.executor.attack.multi_turn.multi_prompt_sending import MultiPromptSendingAttack +from pyrit.executor.attack.multi_turn.pair import PAIRAttack +from pyrit.executor.attack.multi_turn.red_teaming import RedTeamingAttack +from pyrit.executor.attack.multi_turn.tree_of_attacks import TreeOfAttacksWithPruningAttack +from pyrit.executor.attack.single_turn.prompt_sending import PromptSendingAttack +from pyrit.executor.attack.single_turn.skeleton_key import SkeletonKeyAttack +from pyrit.message_normalizer import HistorySquashNormalizer +from pyrit.models import Message, MessagePiece +from pyrit.prompt_target import CapabilityName, PromptTarget, TargetCapabilities, TargetConfiguration +from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory + + +class _NonEditableHistoryTarget(PromptTarget): + _DEFAULT_CONFIGURATION = TargetConfiguration(capabilities=TargetCapabilities(supports_editable_history=False)) + + async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: + request = normalized_conversation[-1] + return [ + MessagePiece( + role="assistant", + original_value="response", + conversation_id=request.get_piece().conversation_id, + ).to_message() + ] + + +@pytest.mark.usefixtures("patch_central_database") +def test_attack_strategy_owns_prepended_policy_and_identifier() -> None: + config = PrependedConversationConfig(apply_converters_to_roles=["user", "assistant"]) + attack = PromptSendingAttack( + objective_target=_NonEditableHistoryTarget(), + prepended_conversation_config=config, + ) + + identifier = attack.get_identifier() + assert attack._prepended_conversation_config is config + assert identifier.params["prepended_conversation_converter_roles"] == ["user", "assistant"] + assert "prepended_conversation_formatter" in identifier.children + + +@pytest.mark.usefixtures("patch_central_database") +def test_attack_strategy_resolves_per_send_history_override() -> None: + attack = PromptSendingAttack(objective_target=_NonEditableHistoryTarget()) + context = PrependedHistorySendContext( + conversation_id="conversation", + seed_message_ids=(uuid.uuid4(),), + replay_seed_each_send=False, + ) + + overrides = attack._get_prepended_normalizer_overrides(prepended_history_send_context=context) + + assert isinstance(overrides[CapabilityName.EDITABLE_HISTORY], HistorySquashNormalizer) + + +@pytest.mark.parametrize( + "attack_class", + [ + ChunkedRequestAttack, + CrescendoAttack, + MultiPromptSendingAttack, + PAIRAttack, + PromptSendingAttack, + RedTeamingAttack, + SkeletonKeyAttack, + TreeOfAttacksWithPruningAttack, + ], +) +def test_techniques_can_specify_prepended_policy(attack_class: type[AttackStrategy[Any, Any]]) -> None: + """Each attack that creates or accepts prepended history exposes the policy.""" + parameter = inspect.signature(attack_class.__init__).parameters.get("prepended_conversation_config") + + assert parameter is not None + assert parameter.kind is inspect.Parameter.KEYWORD_ONLY + + +@pytest.mark.usefixtures("patch_central_database") +def test_technique_factory_forwards_prepended_policy() -> None: + config = PrependedConversationConfig(apply_converters_to_roles=["user", "assistant"]) + factory = AttackTechniqueFactory( + name="policy_test", + attack_class=PromptSendingAttack, + attack_kwargs={"prepended_conversation_config": config}, + ) + + attack = factory.create( + objective_target=_NonEditableHistoryTarget(), + attack_scoring_config=AttackScoringConfig(), + ).attack + + assert attack._prepended_conversation_config is config diff --git a/tests/unit/executor/attack/multi_turn/test_chunked_request.py b/tests/unit/executor/attack/multi_turn/test_chunked_request.py index 7ac1d00af2..106499dbbe 100644 --- a/tests/unit/executor/attack/multi_turn/test_chunked_request.py +++ b/tests/unit/executor/attack/multi_turn/test_chunked_request.py @@ -5,18 +5,29 @@ Tests for ChunkedRequestAttack. """ +import uuid from unittest.mock import AsyncMock, MagicMock import pytest +from pyrit.executor.attack.component import PrependedConversationConfig +from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, +) from pyrit.executor.attack.core.attack_parameters import AttackParameters from pyrit.executor.attack.multi_turn import ( ChunkedRequestAttack, ChunkedRequestAttackContext, ) +from pyrit.message_normalizer import HistorySquashNormalizer, MessageStringNormalizer from pyrit.models import ComponentIdentifier, Message, MessagePiece from pyrit.prompt_normalizer import PromptNormalizer -from pyrit.prompt_target import PromptTarget +from pyrit.prompt_target import ( + CapabilityName, + PromptTarget, + TargetCapabilities, + TargetConfiguration, +) def _mock_target_id(name: str = "MockTarget") -> ComponentIdentifier: @@ -257,6 +268,38 @@ def test_generate_chunk_prompts_with_objective(self): class TestChunkedRequestAttackExecution: """Tests for the main attack execution logic.""" + async def test_perform_async_forwards_prepended_formatter_override(self): + mock_target = _make_mock_target() + mock_target.configuration = TargetConfiguration(capabilities=TargetCapabilities(supports_multi_turn=True)) + mock_normalizer = MagicMock(spec=PromptNormalizer) + mock_normalizer.send_prompt_async = AsyncMock( + return_value=Message.from_prompt(prompt="chunk response", role="assistant") + ) + formatter = MagicMock(spec=MessageStringNormalizer) + config = PrependedConversationConfig(message_normalizer=formatter) + attack = ChunkedRequestAttack( + objective_target=mock_target, + prompt_normalizer=mock_normalizer, + prepended_conversation_config=config, + chunk_size=100, + total_length=100, + ) + context = ChunkedRequestAttackContext(params=AttackParameters(objective="Extract the secret")) + target_context = PrependedHistorySendContext( + conversation_id=context.session.conversation_id, + seed_message_ids=(uuid.uuid4(),), + replay_seed_each_send=False, + ) + context.prepended_history_send_context = target_context + + await attack._perform_async(context=context) + + send_kwargs = mock_normalizer.send_prompt_async.await_args.kwargs + override = send_kwargs["normalizer_overrides"][CapabilityName.EDITABLE_HISTORY] + assert isinstance(override, HistorySquashNormalizer) + assert override._message_normalizer is formatter + assert send_kwargs["send_context"] is target_context + async def test_perform_async_sets_atomic_attack_identifier(self): """Test that _perform_async sets atomic_attack_identifier in the correct AtomicAttack format.""" mock_target = _make_mock_target() diff --git a/tests/unit/executor/attack/multi_turn/test_crescendo.py b/tests/unit/executor/attack/multi_turn/test_crescendo.py index 83720e485b..ea9b718a70 100644 --- a/tests/unit/executor/attack/multi_turn/test_crescendo.py +++ b/tests/unit/executor/attack/multi_turn/test_crescendo.py @@ -458,12 +458,14 @@ def test_init_rejects_adversarial_chat_missing_native_capability( """Adversarial chat must natively support MULTI_TURN and SYSTEM_PROMPT.""" from pyrit.prompt_target.common.target_capabilities import CapabilityName - mock_adversarial_chat.configuration.includes.side_effect = lambda *, capability: ( - capability != CapabilityName(missing_capability) - ) + missing = { + "multi_turn": CapabilityName.MULTI_TURN, + "system_prompt": CapabilityName.SYSTEM_PROMPT, + }[missing_capability] + mock_adversarial_chat.configuration.includes.side_effect = lambda *, capability: capability != missing adversarial_config = AttackAdversarialConfig(target=mock_adversarial_chat) - with pytest.raises(ValueError, match=f"CrescendoAttack .*{missing_capability}"): + with pytest.raises(ValueError, match=f"supports_{missing_capability}"): CrescendoAttack( objective_target=mock_objective_target, attack_adversarial_config=adversarial_config, diff --git a/tests/unit/executor/attack/multi_turn/test_multi_prompt_sending.py b/tests/unit/executor/attack/multi_turn/test_multi_prompt_sending.py index e6a4467089..40dee2a434 100644 --- a/tests/unit/executor/attack/multi_turn/test_multi_prompt_sending.py +++ b/tests/unit/executor/attack/multi_turn/test_multi_prompt_sending.py @@ -16,6 +16,11 @@ MultiPromptSendingAttackParameters, MultiTurnAttackContext, ) +from pyrit.executor.attack.component import PrependedConversationConfig +from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, +) +from pyrit.message_normalizer import HistorySquashNormalizer, MessageStringNormalizer from pyrit.models import ( AttackOutcome, AttackResult, @@ -25,7 +30,12 @@ Score, ) from pyrit.prompt_normalizer import ConverterConfiguration, PromptNormalizer -from pyrit.prompt_target import PromptTarget +from pyrit.prompt_target import ( + CapabilityName, + PromptTarget, + TargetCapabilities, + TargetConfiguration, +) from pyrit.score import Scorer, TrueFalseScorer @@ -311,6 +321,35 @@ async def test_setup_updates_conversation_state_with_converters(self, mock_targe class TestPromptSending: """Tests for sending prompts to target""" + async def test_send_prompt_forwards_prepended_formatter_override( + self, mock_target, mock_prompt_normalizer, basic_context, sample_response + ): + mock_target.configuration = TargetConfiguration(capabilities=TargetCapabilities(supports_multi_turn=True)) + formatter = MagicMock(spec=MessageStringNormalizer) + attack = MultiPromptSendingAttack( + objective_target=mock_target, + prompt_normalizer=mock_prompt_normalizer, + prepended_conversation_config=PrependedConversationConfig(message_normalizer=formatter), + ) + target_context = PrependedHistorySendContext( + conversation_id=basic_context.session.conversation_id, + seed_message_ids=(uuid.uuid4(),), + replay_seed_each_send=False, + ) + basic_context.prepended_history_send_context = target_context + mock_prompt_normalizer.send_prompt_async.return_value = sample_response + + await attack._send_prompt_to_objective_target_async( + current_message=Message.from_prompt(prompt="test prompt", role="user"), + context=basic_context, + ) + + send_kwargs = mock_prompt_normalizer.send_prompt_async.await_args.kwargs + override = send_kwargs["normalizer_overrides"][CapabilityName.EDITABLE_HISTORY] + assert isinstance(override, HistorySquashNormalizer) + assert override._message_normalizer is formatter + assert send_kwargs["send_context"] is target_context + async def test_send_prompt_to_target_with_all_configurations( self, mock_target, mock_prompt_normalizer, basic_context, sample_response ): diff --git a/tests/unit/executor/attack/multi_turn/test_prepended_history_normalization.py b/tests/unit/executor/attack/multi_turn/test_prepended_history_normalization.py new file mode 100644 index 0000000000..067d3b4b6b --- /dev/null +++ b/tests/unit/executor/attack/multi_turn/test_prepended_history_normalization.py @@ -0,0 +1,731 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Focused regression tests for prepended-history target normalization.""" + +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from pyrit.converter import Converter, ConverterResult +from pyrit.executor.attack.component import ConversationManager, PrependedConversationConfig +from pyrit.executor.attack.component.modality_router import _ModalityFeedbackRouter +from pyrit.executor.attack.core import AttackAdversarialConfig, AttackScoringConfig +from pyrit.executor.attack.core.attack_parameters import AttackParameters +from pyrit.executor.attack.multi_turn.multi_turn_attack_strategy import ( + MultiTurnAttackContext, + MultiTurnAttackStrategy, +) +from pyrit.executor.attack.multi_turn.tree_of_attacks import ( + TreeOfAttacksWithPruningAttack, + _TreeOfAttacksNode, +) +from pyrit.memory import CentralMemory +from pyrit.message_normalizer import MessageStringNormalizer +from pyrit.models import ( + ComponentIdentifier, + Conversation, + ConversationReference, + ConversationType, + Message, + MessagePiece, + PromptDataType, +) +from pyrit.prompt_normalizer import ConverterConfiguration, PromptNormalizer +from pyrit.prompt_target import ( + CapabilityName, + PromptTarget, + TargetCapabilities, + TargetConfiguration, +) +from pyrit.score import TrueFalseScorer + + +class _RecordingTarget(PromptTarget): + def __init__(self, *, supports_multi_turn: bool = False, supports_editable_history: bool = False) -> None: + super().__init__( + custom_configuration=TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_turn=supports_multi_turn, + supports_editable_history=supports_editable_history, + ) + ) + ) + self.prompt_sent: list[str] = [] + self.normalized_requests: list[Message] = [] + + async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: + request = normalized_conversation[-1] + self.prompt_sent.append(request.get_value()) + self.normalized_requests.append(request) + return [ + MessagePiece( + role="assistant", + original_value="response", + conversation_id=request.get_piece().conversation_id, + ).to_message() + ] + + +class _ConversationKeyedRecordingTarget(_RecordingTarget): + def __init__(self) -> None: + super().__init__(supports_multi_turn=True) + self.prompts_by_conversation: dict[str, list[str]] = {} + + async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: + request = normalized_conversation[-1] + conversation_id = request.get_piece().conversation_id + assert conversation_id + self.prompts_by_conversation.setdefault(conversation_id, []).append(request.get_value()) + return await super()._send_prompt_to_target_async(normalized_conversation=normalized_conversation) + + +class _ImageOutputConverter(Converter): + SUPPORTED_INPUT_TYPES: tuple[PromptDataType, ...] = ("text",) + SUPPORTED_OUTPUT_TYPES: tuple[PromptDataType, ...] = ("image_path",) + + def __init__(self, *, output_path: str) -> None: + super().__init__() + self._output_path = output_path + + async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text") -> ConverterResult: + return ConverterResult(output_text=self._output_path, output_type="image_path") + + +class _TextOutputConverter(Converter): + SUPPORTED_INPUT_TYPES: tuple[PromptDataType, ...] = ("image_path",) + SUPPORTED_OUTPUT_TYPES: tuple[PromptDataType, ...] = ("text",) + + async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text") -> ConverterResult: + return ConverterResult(output_text="converted text", output_type="text") + + +def _make_context() -> MultiTurnAttackContext[AttackParameters]: + return MultiTurnAttackContext(params=AttackParameters(objective="test objective")) + + +def _rotate( + *, + target: PromptTarget, + context: MultiTurnAttackContext[AttackParameters], +) -> MagicMock: + strategy = MagicMock() + strategy._objective_target = target + strategy._logger = MagicMock() + MultiTurnAttackStrategy._rotate_conversation_for_single_turn_target(strategy, context=context) + return strategy + + +def _seed_conversation( + *, + conversation_id: str, + target: PromptTarget, + messages: list[Message], +) -> None: + memory = CentralMemory.get_memory_instance() + memory.add_conversation_to_memory( + conversation=Conversation( + conversation_id=conversation_id, + target_identifier=target.get_identifier(), + ) + ) + for message in messages: + for piece in message.message_pieces: + piece.conversation_id = conversation_id + memory.add_message_to_memory(request=message) + + +@pytest.mark.usefixtures("patch_central_database") +async def test_single_turn_target_replays_seed_without_rotation(): + target = _RecordingTarget() + prompt_normalizer = PromptNormalizer() + manager = ConversationManager(prompt_normalizer=prompt_normalizer) + config = PrependedConversationConfig() + conversation_id = "conversation" + prepended = Message.from_system_prompt("system") + + await manager.add_prepended_conversation_to_memory_async( + prepended_conversation=[prepended], + conversation_id=conversation_id, + prepended_conversation_config=config, + target=target, + ) + persisted = manager.get_conversation(conversation_id) + target_context = manager.create_prepended_history_send_context( + target=target, + conversation_id=conversation_id, + prepended_messages=persisted, + ) + assert target_context is not None + + for prompt in ["first", "second"]: + await prompt_normalizer.send_prompt_async( + message=Message.from_prompt(prompt=prompt, role="user"), + target=target, + conversation_id=conversation_id, + normalizer_overrides=config.get_normalizer_overrides( + target=target, + prepended_history_send_context=target_context, + ), + send_context=target_context, + ) + + assert target.prompt_sent == [ + "Turn 1:\nuser: ### Instructions ###\n\nsystem\n\n######\n\nfirst", + "Turn 1:\nuser: ### Instructions ###\n\nsystem\n\n######\n\nsecond", + ] + + +@pytest.mark.usefixtures("patch_central_database") +def test_rotation_is_noop_for_multi_turn_target(): + target = _RecordingTarget(supports_multi_turn=True) + context = _make_context() + context.executed_turns = 1 + original_id = context.session.conversation_id + + _rotate(target=target, context=context) + + assert context.session.conversation_id == original_id + assert not context.related_conversations + + +@pytest.mark.usefixtures("patch_central_database") +def test_rotation_is_noop_for_seeded_single_turn_target(): + target = _RecordingTarget() + context = _make_context() + context.executed_turns = 2 + original_id = context.session.conversation_id + seed = Message.from_prompt(prompt="seed", role="user") + context.prepended_history_send_context = ConversationManager.create_prepended_history_send_context( + target=target, + conversation_id=original_id, + prepended_messages=[seed], + ) + + _rotate(target=target, context=context) + + assert context.session.conversation_id == original_id + assert not context.related_conversations + + +@pytest.mark.usefixtures("patch_central_database") +def test_rotation_moves_unseeded_single_turn_target_to_fresh_conversation(): + target = _RecordingTarget() + context = _make_context() + context.executed_turns = 1 + original_id = context.session.conversation_id + + _rotate(target=target, context=context) + + assert context.session.conversation_id != original_id + assert ( + ConversationReference( + conversation_id=original_id, + conversation_type=ConversationType.PRUNED, + description="single-turn target prior turn 1", + ) + in context.related_conversations + ) + + +@pytest.mark.usefixtures("patch_central_database") +async def test_rotation_preserves_system_payload_for_later_single_turn_send(): + target = _RecordingTarget() + context = _make_context() + context.executed_turns = 1 + old_id = context.session.conversation_id + _seed_conversation( + conversation_id=old_id, + target=target, + messages=[ + Message.from_system_prompt("system"), + Message.from_prompt(prompt="old request", role="user"), + ], + ) + + _rotate(target=target, context=context) + + assert context.prepended_history_send_context is not None + config = PrependedConversationConfig() + await PromptNormalizer().send_prompt_async( + message=Message.from_prompt(prompt="current request", role="user"), + target=target, + conversation_id=context.session.conversation_id, + normalizer_overrides=config.get_normalizer_overrides( + target=target, + prepended_history_send_context=context.prepended_history_send_context, + ), + send_context=context.prepended_history_send_context, + ) + assert target.prompt_sent == ["Turn 1:\nuser: ### Instructions ###\n\nsystem\n\n######\n\ncurrent request"] + + +def _make_tap_node(*, target: PromptTarget) -> _TreeOfAttacksNode: + adversarial_chat = MagicMock(spec=PromptTarget) + adversarial_chat.configuration = TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_turn=True, + supports_system_prompt=True, + supports_editable_history=True, + ) + ) + adversarial_chat.get_identifier.return_value = ComponentIdentifier( + class_name="AdversarialTarget", + class_module="tests", + ) + scorer = MagicMock() + scorer.get_identifier.return_value = ComponentIdentifier(class_name="Scorer", class_module="tests") + seed = MagicMock() + seed.render_template_value.return_value = "template" + return _TreeOfAttacksNode( + objective_target=target, + adversarial_chat=adversarial_chat, + adversarial_chat_seed_prompt=seed, + adversarial_chat_prompt_template=seed, + adversarial_chat_system_seed_prompt=seed, + desired_response_prefix="Sure,", + objective_scorer=scorer, + on_topic_scorer=None, + request_converters=[], + response_converters=[], + auxiliary_scorers=None, + attack_id=ComponentIdentifier(class_name="TAP", class_module="tests"), + attack_strategy_name="TAP", + modality_router=_ModalityFeedbackRouter( + adversarial_chat=adversarial_chat, + objective_target=target, + ), + ) + + +def _set_tap_seed_boundary( + *, + node: _TreeOfAttacksNode, + target: PromptTarget, + seed_messages: list[Message], +) -> None: + _seed_conversation( + conversation_id=node.objective_target_conversation_id, + target=target, + messages=seed_messages, + ) + node._prepended_history_send_context = ConversationManager.create_prepended_history_send_context( + target=target, + conversation_id=node.objective_target_conversation_id, + prepended_messages=seed_messages, + ) + assert node._prepended_history_send_context is not None + + +def _branch_tap_node( + *, + node: _TreeOfAttacksNode, + branching_factor: int, +) -> list[_TreeOfAttacksNode]: + context = MagicMock() + context.nodes = [node] + context.related_conversations = set() + attack = MagicMock() + attack._configuration.branching_factor = branching_factor + TreeOfAttacksWithPruningAttack._branch_existing_nodes(attack, context) + return context.nodes + + +@pytest.mark.usefixtures("patch_central_database") +async def test_tap_seeded_stateless_retained_and_cloned_branches_replay_only_original_seed(): + target = _RecordingTarget() + formatter = MagicMock(spec=MessageStringNormalizer) + + async def format_messages(messages: list[Message]) -> str: + return "|".join(message.get_value() for message in messages) + + formatter.normalize_string_async = AsyncMock(side_effect=format_messages) + node = _make_tap_node(target=target) + node._prepended_conversation_config = PrependedConversationConfig(message_normalizer=formatter) + seed = Message.from_prompt(prompt="original seed", role="user") + _set_tap_seed_boundary(node=node, target=target, seed_messages=[seed]) + node._objective = "objective" + await node._send_prompt_to_target_async("depth one") + original_context = node._prepended_history_send_context + assert original_context is not None + + retained, cloned = _branch_tap_node(node=node, branching_factor=2) + assert retained is node + assert retained._prepended_history_send_context is original_context + assert cloned._prepended_history_send_context is not None + assert cloned._prepended_history_send_context.seed_message_count == 1 + cloned_messages = CentralMemory.get_memory_instance().get_conversation_messages( + conversation_id=cloned.objective_target_conversation_id + ) + assert cloned._prepended_history_send_context.seed_message_ids == (cloned_messages[0].get_piece().id,) + assert cloned._prepended_history_send_context.seed_message_ids != original_context.seed_message_ids + + for branch, prompt in [(retained, "depth two retained"), (cloned, "depth two cloned")]: + branch._objective = "objective" + await branch._send_prompt_to_target_async(prompt) + + deep_clone = cloned.duplicate() + deep_clone._objective = "objective" + await deep_clone._send_prompt_to_target_async("depth three cloned") + + assert target.prompt_sent == [ + "original seed|depth one", + "original seed|depth two retained", + "original seed|depth two cloned", + "original seed|depth three cloned", + ] + formatted_values = [ + [message.get_value() for message in call.args[0]] for call in formatter.normalize_string_async.await_args_list + ] + assert formatted_values == [ + ["original seed", "depth one"], + ["original seed", "depth two retained"], + ["original seed", "depth two cloned"], + ["original seed", "depth three cloned"], + ] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_tap_stateful_clone_bootstraps_duplicated_branch_once(): + target = _ConversationKeyedRecordingTarget() + formatter = MagicMock(spec=MessageStringNormalizer) + + async def format_messages(messages: list[Message]) -> str: + return "|".join(message.get_value() for message in messages) + + formatter.normalize_string_async = AsyncMock(side_effect=format_messages) + node = _make_tap_node(target=target) + node._prepended_conversation_config = PrependedConversationConfig(message_normalizer=formatter) + _set_tap_seed_boundary( + node=node, + target=target, + seed_messages=[Message.from_prompt(prompt="original seed", role="user")], + ) + node._objective = "objective" + parent_conversation_id = node.objective_target_conversation_id + + await node._send_prompt_to_target_async("parent first") + assert node._prepended_history_send_context + assert node._prepended_history_send_context.is_seed_consumed + + cloned = node.duplicate() + cloned._objective = "objective" + cloned_conversation_id = cloned.objective_target_conversation_id + assert cloned_conversation_id != parent_conversation_id + assert cloned._prepended_history_send_context + assert not cloned._prepended_history_send_context.is_seed_consumed + + await node._send_prompt_to_target_async("parent second") + await cloned._send_prompt_to_target_async("clone first") + await cloned._send_prompt_to_target_async("clone second") + + assert target.prompts_by_conversation == { + parent_conversation_id: ["original seed|parent first", "parent second"], + cloned_conversation_id: ["original seed|parent first|response|clone first", "clone second"], + } + + +@pytest.mark.usefixtures("patch_central_database") +async def test_tap_unseeded_stateful_clone_bootstraps_duplicated_branch_once(): + target = _ConversationKeyedRecordingTarget() + formatter = MagicMock(spec=MessageStringNormalizer) + + async def format_messages(messages: list[Message]) -> str: + return "|".join(message.get_value() for message in messages) + + formatter.normalize_string_async = AsyncMock(side_effect=format_messages) + node = _make_tap_node(target=target) + node._prepended_conversation_config = PrependedConversationConfig(message_normalizer=formatter) + node._objective = "objective" + parent_conversation_id = node.objective_target_conversation_id + + await node._send_prompt_to_target_async("parent first") + assert node._prepended_history_send_context is None + + cloned = node.duplicate() + cloned._objective = "objective" + cloned_conversation_id = cloned.objective_target_conversation_id + assert cloned._prepended_history_send_context + assert cloned._prepended_history_send_context.seed_message_count == 0 + assert cloned._prepended_history_send_context.bootstrap_message_count == 2 + + await node._send_prompt_to_target_async("parent second") + await cloned._send_prompt_to_target_async("clone first") + await cloned._send_prompt_to_target_async("clone second") + + assert target.prompts_by_conversation == { + parent_conversation_id: ["parent first", "parent second"], + cloned_conversation_id: ["parent first|response|clone first", "clone second"], + } + + +@pytest.mark.usefixtures("patch_central_database") +def test_tap_branch_preserves_multimodal_last_response(): + target = _RecordingTarget() + node = _make_tap_node(target=target) + node.last_response = Message( + message_pieces=[ + MessagePiece(role="assistant", original_value="text response"), + MessagePiece( + role="assistant", + original_value="response.png", + original_value_data_type="image_path", + ), + ] + ) + + duplicate = node.duplicate() + + assert duplicate.last_response == node.last_response + assert duplicate.last_response is not node.last_response + assert [piece.original_value_data_type for piece in duplicate.last_response.message_pieces] == [ + "text", + "image_path", + ] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_tap_unseeded_stateless_send_retains_current_only_payload(): + target = _RecordingTarget() + node = _make_tap_node(target=target) + node._objective = "objective" + previous_conversation_id = node.objective_target_conversation_id + _seed_conversation( + conversation_id=previous_conversation_id, + target=target, + messages=[Message.from_prompt(prompt="prior request", role="user")], + ) + + await node._send_prompt_to_target_async("current request") + + assert node.objective_target_conversation_id != previous_conversation_id + assert target.prompt_sent == ["current request"] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_tap_unseeded_stateless_retained_and_cloned_branches_send_current_only(): + target = _RecordingTarget() + node = _make_tap_node(target=target) + node._objective = "objective" + await node._send_prompt_to_target_async("depth one") + + retained, cloned = _branch_tap_node(node=node, branching_factor=2) + assert retained._prepended_history_send_context is None + assert cloned._prepended_history_send_context is None + for branch, prompt in [(retained, "depth two retained"), (cloned, "depth two cloned")]: + branch._objective = "objective" + await branch._send_prompt_to_target_async(prompt) + + deep_clone = cloned.duplicate() + deep_clone._objective = "objective" + await deep_clone._send_prompt_to_target_async("depth three cloned") + + assert target.prompt_sent == [ + "depth one", + "depth two retained", + "depth two cloned", + "depth three cloned", + ] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_tap_branching_factor_one_preserves_retained_seed_boundary(): + target = _RecordingTarget() + formatter = MagicMock(spec=MessageStringNormalizer) + + async def format_messages(messages: list[Message]) -> str: + return "|".join(message.get_value() for message in messages) + + formatter.normalize_string_async = AsyncMock(side_effect=format_messages) + node = _make_tap_node(target=target) + node._prepended_conversation_config = PrependedConversationConfig(message_normalizer=formatter) + _set_tap_seed_boundary( + node=node, + target=target, + seed_messages=[Message.from_prompt(prompt="original seed", role="user")], + ) + node._objective = "objective" + await node._send_prompt_to_target_async("depth one") + original_context = node._prepended_history_send_context + original_boundary = original_context.seed_message_ids if original_context else () + + branches = _branch_tap_node(node=node, branching_factor=1) + + assert branches == [node] + assert node._prepended_history_send_context is original_context + assert node._prepended_history_send_context is not None + assert node._prepended_history_send_context.seed_message_ids == original_boundary + await node._send_prompt_to_target_async("depth two retained") + assert target.prompt_sent == [ + "original seed|depth one", + "original seed|depth two retained", + ] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_tap_clone_does_not_replay_non_text_live_converter_output(tmp_path: Path): + image_path = tmp_path / "converted.png" + image_path.write_bytes(b"test image") + target = _RecordingTarget() + target._configuration = TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_message_pieces=True, + input_modalities=frozenset( + { + frozenset({"text"}), + frozenset({"image_path"}), + frozenset({"text", "image_path"}), + } + ), + ) + ) + formatter = MagicMock(spec=MessageStringNormalizer) + + async def format_messages(messages: list[Message]) -> str: + return "|".join(message.get_value() for message in messages) + + formatter.normalize_string_async = AsyncMock(side_effect=format_messages) + node = _make_tap_node(target=target) + node._request_converters = ConverterConfiguration.from_converters( + converters=[_ImageOutputConverter(output_path=str(image_path))] + ) + node._prepended_conversation_config = PrependedConversationConfig(message_normalizer=formatter) + _set_tap_seed_boundary( + node=node, + target=target, + seed_messages=[Message.from_prompt(prompt="original seed", role="user")], + ) + node._objective = "objective" + + await node._send_prompt_to_target_async("depth one") + cloned = node.duplicate() + cloned._objective = "objective" + await cloned._send_prompt_to_target_async("depth two") + + assert [piece.converted_value_data_type for piece in target.normalized_requests[-1].message_pieces] == [ + "text", + "image_path", + ] + assert target.normalized_requests[-1].get_values() == ["original seed", str(image_path)] + formatted_values = [ + [message.get_value() for message in call.args[0]] for call in formatter.normalize_string_async.await_args_list + ] + assert formatted_values == [ + ["original seed", "depth one"], + ["original seed"], + ["original seed", "depth two"], + ["original seed"], + ] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_tap_stateful_clone_rejects_non_text_converter_history(tmp_path: Path): + image_path = tmp_path / "converted.png" + image_path.write_bytes(b"test image") + target = _RecordingTarget() + target._configuration = TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_turn=True, + supports_multi_message_pieces=True, + input_modalities=frozenset( + { + frozenset({"text"}), + frozenset({"image_path"}), + frozenset({"text", "image_path"}), + } + ), + ) + ) + node = _make_tap_node(target=target) + node._request_converters = ConverterConfiguration.from_converters( + converters=[_ImageOutputConverter(output_path=str(image_path))] + ) + node._objective = "objective" + + await node._send_prompt_to_target_async("depth one") + + with pytest.raises(ValueError, match="cannot clone.*non-text output.*image_path"): + node.duplicate() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_tap_stateful_clone_accepts_converter_pipeline_with_final_text_output(tmp_path: Path): + image_path = tmp_path / "converted.png" + image_path.write_bytes(b"test image") + target = _ConversationKeyedRecordingTarget() + formatter = MagicMock(spec=MessageStringNormalizer) + + async def format_messages(messages: list[Message]) -> str: + return "|".join(message.get_value() for message in messages) + + formatter.normalize_string_async = AsyncMock(side_effect=format_messages) + node = _make_tap_node(target=target) + node._request_converters = ConverterConfiguration.from_converters( + converters=[ + _ImageOutputConverter(output_path=str(image_path)), + _TextOutputConverter(), + ] + ) + node._prepended_conversation_config = PrependedConversationConfig(message_normalizer=formatter) + node._objective = "objective" + + await node._send_prompt_to_target_async("depth one") + cloned = node.duplicate() + cloned._objective = "objective" + await cloned._send_prompt_to_target_async("depth two") + + assert cloned._prepended_history_send_context + assert cloned._prepended_history_send_context.is_seed_consumed + assert target.normalized_requests[-1].get_piece().converted_value == "converted text|response|converted text" + + +@pytest.fixture +def adversarial_config() -> AttackAdversarialConfig: + target = MagicMock(spec=PromptTarget) + target.configuration = TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_turn=True, + supports_system_prompt=True, + supports_editable_history=True, + ) + ) + return AttackAdversarialConfig(target=target) + + +@pytest.fixture +def scoring_config() -> AttackScoringConfig: + scorer = MagicMock(spec=TrueFalseScorer) + return AttackScoringConfig(objective_scorer=scorer) + + +@pytest.mark.usefixtures("patch_central_database") +def test_crescendo_requires_native_editable_history( + adversarial_config: AttackAdversarialConfig, + scoring_config: AttackScoringConfig, +): + from pyrit.executor.attack.multi_turn.crescendo import CrescendoAttack + + target = _RecordingTarget(supports_multi_turn=True, supports_editable_history=False) + with pytest.raises(ValueError, match=CapabilityName.EDITABLE_HISTORY.value): + CrescendoAttack( + objective_target=target, + attack_adversarial_config=adversarial_config, + attack_scoring_config=scoring_config, + ) + + +@pytest.mark.usefixtures("patch_central_database") +def test_multi_prompt_sending_requires_native_multi_turn(): + from pyrit.executor.attack.multi_turn.multi_prompt_sending import MultiPromptSendingAttack + + with pytest.raises(ValueError, match=CapabilityName.MULTI_TURN.value): + MultiPromptSendingAttack(objective_target=_RecordingTarget()) + + +@pytest.mark.usefixtures("patch_central_database") +def test_chunked_request_requires_native_multi_turn(): + from pyrit.executor.attack.multi_turn.chunked_request import ChunkedRequestAttack + + with pytest.raises(ValueError, match=CapabilityName.MULTI_TURN.value): + ChunkedRequestAttack(objective_target=_RecordingTarget()) diff --git a/tests/unit/executor/attack/multi_turn/test_red_teaming.py b/tests/unit/executor/attack/multi_turn/test_red_teaming.py index 74930594f4..8ba062cc05 100644 --- a/tests/unit/executor/attack/multi_turn/test_red_teaming.py +++ b/tests/unit/executor/attack/multi_turn/test_red_teaming.py @@ -19,7 +19,10 @@ RedTeamingAttack, RTASystemPromptPaths, ) +from pyrit.executor.attack.component import ConversationManager, PrependedConversationConfig from pyrit.executor.attack.core.attack_config import DEFAULT_ADVERSARIAL_FIRST_MESSAGE +from pyrit.memory import CentralMemory +from pyrit.message_normalizer import MessageStringNormalizer from pyrit.models import ( AttackOutcome, AttackResult, @@ -33,7 +36,10 @@ ) from pyrit.prompt_normalizer import PromptNormalizer from pyrit.prompt_target import PromptTarget +from pyrit.prompt_target.common.target_capabilities import TargetCapabilities +from pyrit.prompt_target.common.target_configuration import TargetConfiguration from pyrit.score import Scorer, TrueFalseScorer +from tests.unit.mocks import MockPromptTarget def _adversarial_reply_message(next_message: str = "Adversarial next message") -> Message: @@ -298,6 +304,7 @@ def test_init_with_all_custom_configurations( assert attack._request_converters == converter_config.request_converters assert attack._response_converters == converter_config.response_converters assert attack._prompt_normalizer == mock_prompt_normalizer + assert attack._conversation_manager._prompt_normalizer is mock_prompt_normalizer assert attack._max_turns == 20 def test_init_without_objective_scorer_raises_error( @@ -644,6 +651,32 @@ async def test_setup_initializes_conversation_session( assert basic_context.session is not None assert isinstance(basic_context.session, ConversationSession) + async def test_setup_forwards_prepended_conversation_config( + self, + mock_objective_target: MagicMock, + mock_objective_scorer: MagicMock, + mock_adversarial_chat: MagicMock, + basic_context: MultiTurnAttackContext, + ): + """Setup must use the configured prepended-conversation formatter.""" + prepended_conversation_config = PrependedConversationConfig() + attack = RedTeamingAttack( + objective_target=mock_objective_target, + attack_adversarial_config=AttackAdversarialConfig(target=mock_adversarial_chat), + attack_scoring_config=AttackScoringConfig(objective_scorer=mock_objective_scorer), + prepended_conversation_config=prepended_conversation_config, + ) + + with patch.object( + attack._conversation_manager, + "initialize_context_async", + new_callable=AsyncMock, + return_value=ConversationState(turn_count=0), + ) as mock_initialize: + await attack._setup_async(context=basic_context) + + assert mock_initialize.call_args.kwargs["prepended_conversation_config"] is prepended_conversation_config + async def test_setup_updates_turn_count_from_prepended_conversation( self, mock_objective_target: MagicMock, @@ -916,6 +949,64 @@ async def test_generate_next_prompt_raises_on_none_response( await attack._generate_next_prompt_async(context=basic_context) +@pytest.mark.usefixtures("patch_central_database") +class TestObjectiveTargetSending: + """Tests for sending prompts to the objective target.""" + + async def test_second_turn_uses_configured_message_normalizer_without_rotation( + self, + mock_objective_scorer: MagicMock, + mock_adversarial_chat: MagicMock, + basic_context: MultiTurnAttackContext, + ) -> None: + """A stateless target must reuse formatting without attack-specific rotation.""" + objective_target = MockPromptTarget() + objective_target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + message_normalizer = MagicMock(spec=MessageStringNormalizer) + message_normalizer.normalize_string_async = AsyncMock(return_value="custom formatted request") + attack = RedTeamingAttack( + objective_target=objective_target, + attack_adversarial_config=AttackAdversarialConfig(target=mock_adversarial_chat), + attack_scoring_config=AttackScoringConfig(objective_scorer=mock_objective_scorer), + prepended_conversation_config=PrependedConversationConfig(message_normalizer=message_normalizer), + ) + basic_context.session = ConversationSession() + old_conversation_id = basic_context.session.conversation_id + memory = CentralMemory.get_memory_instance() + system_piece = MessagePiece( + original_value="You are a helpful assistant.", + role="system", + conversation_id=old_conversation_id, + sequence=0, + ) + memory.add_message_pieces_to_memory( + message_pieces=[ + system_piece, + MessagePiece( + original_value="First request", + role="user", + conversation_id=old_conversation_id, + sequence=1, + ), + ] + ) + basic_context.prepended_history_send_context = ConversationManager.create_prepended_history_send_context( + target=objective_target, + conversation_id=old_conversation_id, + prepended_messages=[system_piece.to_message()], + ) + basic_context.executed_turns = 1 + + await attack._send_prompt_to_objective_target_async( + context=basic_context, + message=Message.from_prompt(prompt="Second request", role="user"), + ) + + assert basic_context.session.conversation_id == old_conversation_id + assert objective_target.prompt_sent == ["custom formatted request"] + message_normalizer.normalize_string_async.assert_awaited_once() + + @pytest.mark.usefixtures("patch_central_database") class TestResponseScoring: """Tests for response scoring logic.""" diff --git a/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py b/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py index ec81a8ff6e..7cc2ce0625 100644 --- a/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py +++ b/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py @@ -1,17 +1,44 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import uuid from unittest.mock import AsyncMock, MagicMock, patch import pytest +from pyrit.executor.attack.component import PrependedConversationConfig +from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, +) from pyrit.executor.attack.core.attack_parameters import AttackParameters from pyrit.executor.attack.multi_turn.multi_turn_attack_strategy import ( ConversationSession, MultiTurnAttackContext, ) from pyrit.memory import CentralMemory -from pyrit.models import ConversationType, MessagePiece +from pyrit.models import ConversationType, Message, MessagePiece +from pyrit.prompt_target import PromptTarget +from pyrit.prompt_target.common.target_capabilities import TargetCapabilities +from pyrit.prompt_target.common.target_configuration import TargetConfiguration + + +class _SingleTurnPromptTarget(PromptTarget): + _DEFAULT_CONFIGURATION = TargetConfiguration(capabilities=TargetCapabilities()) + + def __init__(self) -> None: + super().__init__() + self.normalized_conversations: list[list[Message]] = [] + + async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: + self.normalized_conversations.append(normalized_conversation) + request_piece = normalized_conversation[-1].get_piece() + return [ + MessagePiece( + role="assistant", + original_value="response", + conversation_id=request_piece.conversation_id, + ).to_message() + ] def _make_context() -> MultiTurnAttackContext: @@ -88,6 +115,22 @@ def test_noop_on_first_turn(self): assert context.session.conversation_id == original_id assert len(context.related_conversations) == 0 + def test_pending_first_turn_context_suppresses_rotation_after_prepended_turns(self): + strategy = _make_strategy(supports_multi_turn=False) + context = _make_context() + context.executed_turns = 2 + original_id = context.session.conversation_id + context.prepended_history_send_context = PrependedHistorySendContext( + conversation_id=original_id, + seed_message_ids=(uuid.uuid4(),), + replay_seed_each_send=True, + ) + + strategy._rotate_conversation_for_single_turn_target(context=context) + + assert context.session.conversation_id == original_id + assert len(context.related_conversations) == 0 + def test_rotates_on_second_turn_for_single_turn_target(self): strategy = _make_strategy(supports_multi_turn=False) context = _make_context() @@ -176,6 +219,42 @@ def test_system_prompt_preserved_across_multiple_rotations(self): ) memory.add_message_pieces_to_memory(message_pieces=[user_piece]) + async def test_rotated_system_prompt_is_normalized_with_next_request(self): + target = _SingleTurnPromptTarget() + strategy = _make_strategy(supports_multi_turn=False) + strategy._objective_target = target + context = _make_context() + old_id = context.session.conversation_id + _seed_conversation( + conversation_id=old_id, + system_prompt="You are a helpful assistant.", + user_text="First request", + ) + context.executed_turns = 1 + + strategy._rotate_conversation_for_single_turn_target(context=context) + + next_request = Message.from_prompt(prompt="Second request", role="user") + next_request.get_piece().conversation_id = context.session.conversation_id + config = PrependedConversationConfig() + await target.send_prompt_async( + message=next_request, + normalizer_overrides=config.get_normalizer_overrides( + target=target, + prepended_history_send_context=context.prepended_history_send_context, + ), + send_context=context.prepended_history_send_context, + ) + + assert context.prepended_history_send_context is not None + assert not context.prepended_history_send_context.is_seed_consumed + assert len(target.normalized_conversations) == 1 + normalized_conversation = target.normalized_conversations[0] + assert len(normalized_conversation) == 1 + assert normalized_conversation[0].get_value() == ( + "Turn 1:\nuser: ### Instructions ###\n\nYou are a helpful assistant.\n\n######\n\nSecond request" + ) + def test_no_system_prompt_yields_fresh_conversation_id(self): """When there is no system prompt, rotation still generates a new conversation_id.""" strategy = _make_strategy(supports_multi_turn=False) @@ -414,8 +493,8 @@ def _make_tap_node(self, *, supports_multi_turn: bool): ), ) - def test_single_turn_target_duplicates_only_system_messages(self): - """For single-turn targets, only system messages are copied to the duplicate node.""" + def test_single_turn_target_duplicates_logical_history_without_seed_boundary(self): + """Single-turn branches retain memory without replaying live turns as seed history.""" node = self._make_tap_node(supports_multi_turn=False) memory = CentralMemory.get_memory_instance() @@ -445,11 +524,11 @@ def test_single_turn_target_duplicates_only_system_messages(self): # The duplicate should have a different conversation_id assert duplicate.objective_target_conversation_id != node.objective_target_conversation_id - # The duplicate's conversation should contain only the system message dup_messages = memory.get_conversation_messages(conversation_id=duplicate.objective_target_conversation_id) - assert len(dup_messages) == 1 - assert dup_messages[0].api_role == "system" - assert dup_messages[0].get_value() == "TAP system prompt" + assert [message.api_role for message in dup_messages] == ["system", "user", "assistant"] + # No explicit prepended seed was initialized, so copied live turns do not + # become a new replay boundary. + assert duplicate._prepended_history_send_context is None def test_multi_turn_target_duplicates_full_conversation(self): """For multi-turn targets, the full conversation is duplicated.""" @@ -485,8 +564,8 @@ def test_multi_turn_target_duplicates_full_conversation(self): roles = [m.api_role for m in dup_messages] assert roles == ["system", "user", "assistant"] - def test_single_turn_no_system_messages_yields_fresh_id(self): - """For single-turn targets with no system messages, a fresh empty conversation is created.""" + def test_single_turn_no_system_messages_retains_full_history(self): + """Single-turn branches do not discard non-system logical history.""" node = self._make_tap_node(supports_multi_turn=False) memory = CentralMemory.get_memory_instance() @@ -503,7 +582,7 @@ def test_single_turn_no_system_messages_yields_fresh_id(self): assert duplicate.objective_target_conversation_id != node.objective_target_conversation_id dup_messages = memory.get_conversation_messages(conversation_id=duplicate.objective_target_conversation_id) - assert len(dup_messages) == 0 + assert [message.get_value() for message in dup_messages] == ["Attack prompt"] def test_adversarial_chat_always_fully_duplicated(self): """The adversarial chat conversation should always be fully duplicated regardless of target type.""" @@ -540,8 +619,8 @@ def test_adversarial_chat_always_fully_duplicated(self): roles = [m.api_role for m in dup_adv_messages] assert roles == ["system", "user"] - def test_single_turn_multiple_system_messages_all_duplicated(self): - """For single-turn targets with multiple system messages, all are duplicated.""" + def test_single_turn_multiple_system_messages_and_user_are_duplicated(self): + """All logical branch messages are duplicated in their original order.""" node = self._make_tap_node(supports_multi_turn=False) memory = CentralMemory.get_memory_instance() @@ -568,9 +647,12 @@ def test_single_turn_multiple_system_messages_all_duplicated(self): duplicate = node.duplicate() dup_messages = memory.get_conversation_messages(conversation_id=duplicate.objective_target_conversation_id) - assert all(m.api_role == "system" for m in dup_messages) - dup_values = sorted(m.get_value() for m in dup_messages) - assert dup_values == ["System prompt A", "System prompt B"] + assert [message.api_role for message in dup_messages] == ["system", "user", "system"] + assert [message.get_value() for message in dup_messages] == [ + "System prompt A", + "Attack prompt", + "System prompt B", + ] def test_single_turn_empty_conversation_yields_fresh_id(self): """For single-turn targets with empty conversation, a fresh ID is produced.""" @@ -625,7 +707,7 @@ def test_system_message_content_preserved_exactly(self): duplicate = node.duplicate() dup_messages = memory.get_conversation_messages(conversation_id=duplicate.objective_target_conversation_id) - assert len(dup_messages) == 1 + assert len(dup_messages) == 2 assert dup_messages[0].get_value() == long_prompt def test_original_conversation_untouched_after_duplicate(self): @@ -683,7 +765,7 @@ def test_single_turn_multipiece_system_message_duplicated(self): duplicate = node.duplicate() dup_messages = memory.get_conversation_messages(conversation_id=duplicate.objective_target_conversation_id) - assert len(dup_messages) == 1 + assert len(dup_messages) == 2 assert dup_messages[0].api_role == "system" assert len(dup_messages[0].message_pieces) == 2 dup_values = {p.converted_value for p in dup_messages[0].message_pieces} @@ -801,10 +883,10 @@ def test_branching_single_turn_target_preserves_system_across_depths(self): """Simulate TAP branching across 2 depths and verify system prompts survive. Depth 1: Create a node, seed system + user + assistant messages. - Depth 2: duplicate() the node (simulating branching). For single-turn targets, - only the system message should be in the duplicate's conversation. + Depth 2: duplicate() the node (simulating branching). The full logical + history should be in the duplicate's conversation. Then simulate another turn on the duplicate (add user + assistant). - Depth 3: duplicate() again. System message should still be there. + Depth 3: duplicate() again. The complete branch history should still be there. """ memory = CentralMemory.get_memory_instance() node = self._make_tap_node(supports_multi_turn=False) @@ -830,39 +912,36 @@ def test_branching_single_turn_target_preserves_system_across_depths(self): ) memory.add_message_pieces_to_memory(message_pieces=[sys_piece, user_piece, asst_piece]) - # Depth 2: branch (duplicate) — single-turn means only system msg is copied + # Depth 2: branch with the full logical history. branch1 = node.duplicate() branch1_msgs = memory.get_conversation_messages(conversation_id=branch1.objective_target_conversation_id) - assert len(branch1_msgs) == 1 - assert branch1_msgs[0].api_role == "system" - assert branch1_msgs[0].get_value() == "You are a red team assistant." + assert [message.api_role for message in branch1_msgs] == ["system", "user", "assistant"] # Simulate depth-2 turn on branch1: add user + assistant on branch1's conversation user2 = MessagePiece( original_value="Now tell me about Y", role="user", conversation_id=branch1.objective_target_conversation_id, - sequence=1, + sequence=3, ) asst2 = MessagePiece( original_value="Here is info about Y", role="assistant", conversation_id=branch1.objective_target_conversation_id, - sequence=2, + sequence=4, ) memory.add_message_pieces_to_memory(message_pieces=[user2, asst2]) - # Verify branch1 now has system + user + assistant + # Verify branch1 now has the inherited and new turns. branch1_full = memory.get_conversation_messages(conversation_id=branch1.objective_target_conversation_id) - assert [m.api_role for m in branch1_full] == ["system", "user", "assistant"] + assert [m.api_role for m in branch1_full] == ["system", "user", "assistant", "user", "assistant"] # Depth 3: branch again from branch1 branch2 = branch1.duplicate() branch2_msgs = memory.get_conversation_messages(conversation_id=branch2.objective_target_conversation_id) - assert len(branch2_msgs) == 1 - assert branch2_msgs[0].api_role == "system" + assert [m.api_role for m in branch2_msgs] == ["system", "user", "assistant", "user", "assistant"] assert branch2_msgs[0].get_value() == "You are a red team assistant." def test_branching_multi_turn_target_preserves_full_history(self): diff --git a/tests/unit/executor/attack/multi_turn/test_tree_of_attacks.py b/tests/unit/executor/attack/multi_turn/test_tree_of_attacks.py index 3c36e96fb4..bd2b34edb9 100644 --- a/tests/unit/executor/attack/multi_turn/test_tree_of_attacks.py +++ b/tests/unit/executor/attack/multi_turn/test_tree_of_attacks.py @@ -1921,8 +1921,8 @@ def mock_score_response(*args, **kwargs): assert node.auxiliary_scores["AuxScorer1"].get_value() == 0.8 assert node.auxiliary_scores["AuxScorer2"].get_value() == 0.6 - async def test_node_single_turn_target_generates_new_conv_id(self, node_components): - """Test that single-turn targets get a fresh conversation_id before each send.""" + async def test_node_unseeded_single_turn_target_rotates_conversation_id(self, node_components): + """An unseeded single-turn node keeps prior live history out of the next payload.""" node_components["objective_target"].capabilities.supports_multi_turn = False node_components["objective_target"].configuration.includes.side_effect = lambda capability: False node = _TreeOfAttacksNode(**node_components) @@ -1945,8 +1945,10 @@ async def test_node_single_turn_target_generates_new_conv_id(self, node_componen with patch.object(node, "_score_response_async", new_callable=AsyncMock): await node._send_prompt_to_target_async("test prompt") - # Conversation ID should have changed for single-turn target assert node.objective_target_conversation_id != original_conv_id + send_kwargs = node._prompt_normalizer.send_prompt_async.await_args.kwargs + assert send_kwargs["conversation_id"] == node.objective_target_conversation_id + assert send_kwargs["send_context"] is None async def test_node_multi_turn_target_keeps_conv_id(self, node_components): """Test that multi-turn targets keep the same conversation_id.""" diff --git a/tests/unit/executor/attack/single_turn/test_prompt_sending.py b/tests/unit/executor/attack/single_turn/test_prompt_sending.py index 5d4fda2209..288a2dc9d5 100644 --- a/tests/unit/executor/attack/single_turn/test_prompt_sending.py +++ b/tests/unit/executor/attack/single_turn/test_prompt_sending.py @@ -1,20 +1,24 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import base64 import uuid from unittest.mock import AsyncMock, MagicMock, patch import pytest -from unit.mocks import get_mock_scorer_identifier, get_mock_target_identifier +from unit.mocks import MockPromptTarget, get_mock_scorer_identifier, get_mock_target_identifier from pyrit.converter import Base64Converter, StringJoinConverter from pyrit.executor.attack import ( AttackConverterConfig, AttackParameters, AttackScoringConfig, + PrependedConversationConfig, PromptSendingAttack, SingleTurnAttackContext, ) +from pyrit.memory import CentralMemory +from pyrit.message_normalizer import TokenizerTemplateNormalizer from pyrit.models import ( AttackOutcome, AttackResult, @@ -27,6 +31,8 @@ ) from pyrit.prompt_normalizer import ConverterConfiguration, PromptNormalizer from pyrit.prompt_target import PromptTarget +from pyrit.prompt_target.common.target_capabilities import TargetCapabilities +from pyrit.prompt_target.common.target_configuration import TargetConfiguration from pyrit.score import Scorer, TrueFalseScorer @@ -276,10 +282,105 @@ async def test_setup_updates_conversation_state_with_converters(self, mock_targe target=mock_target, conversation_id=basic_context.conversation_id, request_converters=converter_config, - prepended_conversation_config=None, + prepended_conversation_config=PrependedConversationConfig(), memory_labels={}, ) + async def test_default_converter_scoping_preserves_simulated_assistant_history(self): + target = MockPromptTarget() + converter_config = AttackConverterConfig( + request_converters=ConverterConfiguration.from_converters(converters=[Base64Converter()]) + ) + attack = PromptSendingAttack(objective_target=target, attack_converter_config=converter_config) + prepended_user = "prepended user request" + simulated_response = "simulated assistant response" + final_request = "live final request" + + result = await attack.execute_async( + objective="Test objective", + prepended_conversation=[ + Message.from_prompt(prompt=prepended_user, role="user"), + Message.from_prompt(prompt=simulated_response, role="assistant"), + ], + next_message=Message.from_prompt(prompt=final_request, role="user"), + ) + + pieces = CentralMemory.get_memory_instance().get_message_pieces(conversation_id=result.conversation_id) + assistant_piece = next(piece for piece in pieces if piece.original_value == simulated_response) + final_piece = next(piece for piece in pieces if piece.original_value == final_request) + + assert assistant_piece.role == "simulated_assistant" + assert assistant_piece.converted_value == assistant_piece.original_value + assert assistant_piece.converter_identifiers == [] + assert final_piece.converted_value != final_piece.original_value + assert [identifier.class_name for identifier in final_piece.converter_identifiers] == ["Base64Converter"] + + async def test_explicit_assistant_role_opt_in_converts_simulated_history(self): + target = MockPromptTarget() + converter_config = AttackConverterConfig( + request_converters=ConverterConfiguration.from_converters(converters=[Base64Converter()]) + ) + attack = PromptSendingAttack( + objective_target=target, + attack_converter_config=converter_config, + prepended_conversation_config=PrependedConversationConfig(apply_converters_to_roles=["assistant"]), + ) + simulated_response = "assistant history explicitly converted" + + result = await attack.execute_async( + objective="Test objective", + prepended_conversation=[Message.from_prompt(prompt=simulated_response, role="assistant")], + next_message=Message.from_prompt(prompt="live request", role="user"), + ) + + pieces = CentralMemory.get_memory_instance().get_message_pieces(conversation_id=result.conversation_id) + assistant_piece = next(piece for piece in pieces if piece.original_value == simulated_response) + + assert assistant_piece.role == "simulated_assistant" + assert assistant_piece.converted_value != assistant_piece.original_value + assert [identifier.class_name for identifier in assistant_piece.converter_identifiers] == ["Base64Converter"] + + async def test_non_chat_target_converts_history_by_role_before_flattening(self): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + converter_config = AttackConverterConfig( + request_converters=ConverterConfiguration.from_converters(converters=[Base64Converter()]) + ) + attack = PromptSendingAttack(objective_target=target, attack_converter_config=converter_config) + prepended_user = "prepended user request" + simulated_response = "simulated assistant response" + final_request = "live final request" + + await attack.execute_async( + objective="Test objective", + prepended_conversation=[ + Message.from_prompt(prompt=prepended_user, role="user"), + Message.from_prompt(prompt=simulated_response, role="assistant"), + ], + next_message=Message.from_prompt(prompt=final_request, role="user"), + ) + + encoded_user = base64.b64encode(prepended_user.encode()).decode() + encoded_final_request = base64.b64encode(final_request.encode()).decode() + assert target.prompt_sent == [ + (f"Turn 1:\nuser: {encoded_user}\nassistant: {simulated_response}\nTurn 2:\nuser: {encoded_final_request}") + ] + + async def test_retry_setup_creates_fresh_conversation(self, basic_context): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + attack = PromptSendingAttack(objective_target=target) + basic_context.prepended_conversation = [ + Message.from_prompt(prompt="prepended", role="user"), + ] + + await attack._setup_async(context=basic_context) + first_conversation_id = basic_context.conversation_id + + await attack._setup_async(context=basic_context) + + assert basic_context.conversation_id != first_conversation_id + @pytest.mark.usefixtures("patch_central_database") class TestPromptPreparation: @@ -1151,6 +1252,34 @@ def test_attack_has_unique_identifier(self, mock_target): assert id1.hash == id2.hash assert id1.class_name == id2.class_name == "PromptSendingAttack" + def test_attack_identifier_owns_prepended_formatter_provenance(self, mock_target): + tokenizer = MagicMock() + tokenizer.name_or_path = "example/tokenizer" + tokenizer.chat_template = "{{ messages }}" + tokenizer.special_tokens_map = {"bos_token": ""} + keep_formatter = TokenizerTemplateNormalizer( + tokenizer=tokenizer, + system_message_behavior="keep", + ) + ignore_formatter = TokenizerTemplateNormalizer( + tokenizer=tokenizer, + system_message_behavior="ignore", + ) + + keep_attack = PromptSendingAttack( + objective_target=mock_target, + prepended_conversation_config=PrependedConversationConfig(message_normalizer=keep_formatter), + ) + ignore_attack = PromptSendingAttack( + objective_target=mock_target, + prepended_conversation_config=PrependedConversationConfig(message_normalizer=ignore_formatter), + ) + + keep_identifier = keep_attack.get_identifier() + ignore_identifier = ignore_attack.get_identifier() + assert "prepended_conversation_formatter" in keep_identifier.children + assert keep_identifier.hash != ignore_identifier.hash + async def test_retry_stores_unsuccessful_conversation_and_updates_id( self, mock_target, mock_true_false_scorer, basic_context, sample_response, failure_score ): diff --git a/tests/unit/executor/attack/streaming/test_barge_in.py b/tests/unit/executor/attack/streaming/test_barge_in.py index 7a4bc289fe..60f3a49e95 100644 --- a/tests/unit/executor/attack/streaming/test_barge_in.py +++ b/tests/unit/executor/attack/streaming/test_barge_in.py @@ -10,9 +10,15 @@ import pytest +from pyrit.converter import Base64Converter from pyrit.executor.attack import BargeInAttack, BargeInAttackContext -from pyrit.executor.attack.core import AttackParameters +from pyrit.executor.attack.component import PrependedConversationConfig +from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, +) +from pyrit.executor.attack.core import AttackConverterConfig, AttackParameters from pyrit.models import AttackOutcome, Message, MessagePiece +from pyrit.prompt_normalizer import ConverterConfiguration from pyrit.prompt_target import RealtimeTarget if TYPE_CHECKING: @@ -156,15 +162,74 @@ async def test_setup_async_persists_prepended_conversation_to_memory(vad_target) # All three messages share the context's conversation_id post-setup. for m in add_calls: assert m.message_pieces[0].conversation_id == ctx.conversation_id + assert ctx.prepended_history_send_context is None -async def test_setup_async_no_op_when_prepended_conversation_empty(vad_target): - """Empty prepended_conversation: no memory writes, no crash.""" +async def test_converted_system_prompt_is_passed_to_streaming_session(vad_target): + attack = BargeInAttack( + objective_target=vad_target, + attack_converter_config=AttackConverterConfig( + request_converters=ConverterConfiguration.from_converters(converters=[Base64Converter()]) + ), + prepended_conversation_config=PrependedConversationConfig(apply_converters_to_roles=["system"]), + ) + ctx = BargeInAttackContext( + params=AttackParameters( + objective="o", + prepended_conversation=[Message.from_prompt(prompt="You are strict.", role="system")], + ), + audio_chunks=_aiter([b"\x00" * 96]), + ) + await attack._setup_async(context=ctx) + fake_session = _fake_session() + + with patch.object(RealtimeTarget, "open_streaming_session", return_value=fake_session) as factory: + await attack._perform_async(context=ctx) + + persisted = attack._conversation_manager.get_conversation(ctx.conversation_id) + passed_to_session = factory.call_args.kwargs["prepended_conversation"] + assert passed_to_session == persisted + assert passed_to_session[0].get_piece().converted_value == "WW91IGFyZSBzdHJpY3Qu" + + +async def test_setup_reusing_conversation_passes_only_new_prepended_messages(vad_target): + attack = BargeInAttack(objective_target=vad_target) + conversation_id = "existing-conversation" + await attack._conversation_manager.add_prepended_conversation_to_memory_async( + prepended_conversation=[Message.from_prompt(prompt="old system", role="system")], + conversation_id=conversation_id, + target=vad_target, + ) + ctx = BargeInAttackContext( + params=AttackParameters( + objective="o", + prepended_conversation=[Message.from_prompt(prompt="new system", role="system")], + ), + audio_chunks=_aiter([b"\x00" * 96]), + conversation_id=conversation_id, + ) + + await attack._setup_async(context=ctx) + + assert [message.get_value() for message in ctx.prepended_conversation] == ["new system"] + assert [message.get_value() for message in attack._conversation_manager.get_conversation(conversation_id)] == [ + "old system", + "new system", + ] + + +async def test_setup_async_clears_unused_normalization_context_when_prepended_empty(vad_target): + """The direct streaming path does not retain target normalization state.""" attack = BargeInAttack(objective_target=vad_target) ctx = BargeInAttackContext( params=AttackParameters(objective="o"), # no prepended_conversation audio_chunks=_aiter([b"\x00" * 96]), ) + ctx.prepended_history_send_context = PrependedHistorySendContext( + conversation_id=ctx.conversation_id, + seed_message_ids=(Message.from_prompt(prompt="unused", role="user").get_piece().id,), + replay_seed_each_send=True, + ) add_calls: list[Any] = [] with patch.object(attack._conversation_manager._memory, "add_message_to_memory") as mock_add: @@ -172,6 +237,7 @@ async def test_setup_async_no_op_when_prepended_conversation_empty(vad_target): await attack._setup_async(context=ctx) assert add_calls == [] + assert ctx.prepended_history_send_context is None # ---- _perform_async: session factory passthrough ---------------------------------------------- diff --git a/tests/unit/message_normalizer/test_chat_normalizer_tokenizer.py b/tests/unit/message_normalizer/test_chat_normalizer_tokenizer.py index a0e439e4e7..eb17c6f658 100644 --- a/tests/unit/message_normalizer/test_chat_normalizer_tokenizer.py +++ b/tests/unit/message_normalizer/test_chat_normalizer_tokenizer.py @@ -43,6 +43,18 @@ def test_model_aliases_contains_expected_aliases(self): assert "qwen" in aliases assert "llama3" in aliases + def test_identifier_includes_behavior_changing_configuration(self): + tokenizer = MagicMock() + tokenizer.name_or_path = "example/tokenizer" + tokenizer.chat_template = "{{ messages }}" + tokenizer.special_tokens_map = {"bos_token": ""} + + keep = TokenizerTemplateNormalizer(tokenizer=tokenizer, system_message_behavior="keep") + ignore = TokenizerTemplateNormalizer(tokenizer=tokenizer, system_message_behavior="ignore") + + assert keep.get_identifier().hash != ignore.get_identifier().hash + assert keep.get_identifier().params["chat_template"] == '"{{ messages }}"' + class TestFromModel: """Tests for the from_model factory method.""" diff --git a/tests/unit/message_normalizer/test_history_squash_normalizer.py b/tests/unit/message_normalizer/test_history_squash_normalizer.py index 7ac3648271..c9dad2c2fc 100644 --- a/tests/unit/message_normalizer/test_history_squash_normalizer.py +++ b/tests/unit/message_normalizer/test_history_squash_normalizer.py @@ -1,10 +1,12 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +from unittest.mock import AsyncMock, MagicMock + import pytest -from pyrit.message_normalizer import HistorySquashNormalizer -from pyrit.models import Message, MessagePiece +from pyrit.message_normalizer import HistorySquashNormalizer, MessageStringNormalizer +from pyrit.models import JSON_SCHEMA_METADATA_KEY, Message, MessagePiece from pyrit.models.literals import ChatMessageRole @@ -25,6 +27,21 @@ async def test_history_squash_single_message_returns_unchanged(): assert result[0].api_role == "user" +def test_history_squash_rejects_invalid_expected_history_count(): + with pytest.raises(ValueError, match="expected_history_message_count must be at least 1"): + HistorySquashNormalizer(expected_history_message_count=0) + + +async def test_history_squash_rejects_unexpected_history_count(): + messages = [ + _make_message("user", "history"), + _make_message("user", "current"), + ] + + with pytest.raises(ValueError, match="expected 2 history messages.*received 2 messages"): + await HistorySquashNormalizer(expected_history_message_count=2).normalize_async(messages) + + async def test_history_squash_two_turns(): messages = [ _make_message("user", "hello"), @@ -44,6 +61,23 @@ async def test_history_squash_two_turns(): assert "how are you?" in text +async def test_history_squash_uses_configured_formatter(): + formatter = MagicMock(spec=MessageStringNormalizer) + formatter.normalize_string_async = AsyncMock(return_value="custom format") + messages = [ + _make_message("user", "history"), + _make_message("user", "current"), + ] + + result = await HistorySquashNormalizer( + message_normalizer=formatter, + expected_history_message_count=1, + ).normalize_async(messages) + + assert result[0].get_value() == "custom format" + formatter.normalize_string_async.assert_awaited_once() + + async def test_history_squash_includes_system_in_history(): messages = [ _make_message("system", "You are helpful"), @@ -80,6 +114,124 @@ async def test_history_squash_multi_piece_message(): assert "part2" in text +async def test_history_squash_preserves_original_and_converted_views(): + history = _make_message("user", "original history") + history.get_piece().converted_value = "converted history" + current = _make_message("user", "original current") + current.get_piece().converted_value = "converted current" + + result = await HistorySquashNormalizer().normalize_async([history, current]) + + piece = result[0].get_piece() + assert "User: original history" in piece.original_value + assert "original current" in piece.original_value + assert "User: converted history" in piece.converted_value + assert "converted current" in piece.converted_value + + +async def test_history_squash_preserves_live_multimodal_piece_order(): + conversation_id = "test-conv-id" + current = Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="diagram.png", + original_value_data_type="image_path", + conversation_id=conversation_id, + ), + MessagePiece( + role="user", + original_value="What does this show?", + conversation_id=conversation_id, + ), + ] + ) + + result = await HistorySquashNormalizer().normalize_async([_make_message("assistant", "Earlier response"), current]) + + pieces = result[0].message_pieces + assert [piece.converted_value_data_type for piece in pieces] == ["image_path", "text"] + assert pieces[0].converted_value == "diagram.png" + assert "What does this show?" in pieces[1].converted_value + assert "diagram.png" not in pieces[1].converted_value + + +async def test_history_squash_preserves_entirely_non_text_live_request(): + conversation_id = "test-conv-id" + current = Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="diagram.png", + original_value_data_type="image_path", + conversation_id=conversation_id, + ) + ] + ) + + result = await HistorySquashNormalizer().normalize_async([_make_message("assistant", "Earlier response"), current]) + + pieces = result[0].message_pieces + assert [piece.converted_value_data_type for piece in pieces] == ["text", "image_path"] + assert pieces[0].converted_value == "[Conversation History]\nAssistant: Earlier response" + assert pieces[1].converted_value == "diagram.png" + + +async def test_history_squash_describes_non_text_history_without_exposing_path(): + history = Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="C:\\private\\diagram.png", + original_value_data_type="image_path", + prompt_metadata={"context_description": "architecture diagram"}, + ) + ] + ) + + result = await HistorySquashNormalizer().normalize_async([history, _make_message("user", "What does it show?")]) + + text = result[0].get_value() + assert "User: [Image_path - architecture diagram]" in text + assert "C:\\private\\diagram.png" not in text + + +async def test_history_squash_uses_live_text_metadata_for_combined_piece(): + schema = {"type": "object"} + conversation_id = "test-conv-id" + current = Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="diagram.png", + original_value_data_type="image_path", + conversation_id=conversation_id, + ), + MessagePiece( + role="user", + original_value="Describe this image", + conversation_id=conversation_id, + prompt_metadata={JSON_SCHEMA_METADATA_KEY: schema}, + ), + ] + ) + + result = await HistorySquashNormalizer().normalize_async([_make_message("assistant", "Earlier response"), current]) + + text_piece = result[0].message_pieces[1] + assert text_piece.converted_value_data_type == "text" + assert text_piece.prompt_metadata == {JSON_SCHEMA_METADATA_KEY: schema} + + +async def test_history_squash_rejects_converted_non_text_history(): + history = _make_message("user", "original") + history.get_piece().converted_value = "converted.png" + history.get_piece().converted_value_data_type = "image_path" + + with pytest.raises(ValueError, match="non-text output types.*image_path"): + await HistorySquashNormalizer().normalize_async([history, _make_message("user", "current")]) + + async def test_history_squash_preserves_original_list(): """Normalize should not mutate the input list.""" messages = [ diff --git a/tests/unit/prompt_normalizer/test_prompt_normalizer.py b/tests/unit/prompt_normalizer/test_prompt_normalizer.py index 025d6d6c77..bdc0f00ae8 100644 --- a/tests/unit/prompt_normalizer/test_prompt_normalizer.py +++ b/tests/unit/prompt_normalizer/test_prompt_normalizer.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import asyncio import os import tempfile import wave @@ -24,7 +25,11 @@ execution_context, get_execution_context, ) +from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, +) from pyrit.memory import CentralMemory +from pyrit.message_normalizer import MessageListNormalizer from pyrit.models import ( Message, MessagePiece, @@ -36,7 +41,9 @@ from pyrit.prompt_normalizer.converter_configuration import ( ConverterConfiguration, ) -from pyrit.prompt_target import PromptTarget +from pyrit.prompt_target import CapabilityName, PromptTarget +from pyrit.prompt_target.common.target_capabilities import TargetCapabilities +from pyrit.prompt_target.common.target_configuration import TargetConfiguration @pytest.fixture @@ -93,6 +100,14 @@ def output_supported(self, output_type: PromptDataType) -> bool: return output_type == "text" +class ContextFailingConverter(Converter): + SUPPORTED_INPUT_TYPES: tuple[PromptDataType, ...] = ("text",) + SUPPORTED_OUTPUT_TYPES: tuple[PromptDataType, ...] = ("text",) + + async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text") -> ConverterResult: + raise ValueError("conversion failed") + + def assert_message_piece_hashes_set(request: Message): assert request assert request.message_pieces @@ -115,6 +130,148 @@ async def test_send_prompt_async_multiple_converters(mock_memory_instance, seed_ assert prompt_target.prompt_sent == ["S_G_V_s_b_G_8_="] +async def test_send_prompt_async_forwards_normalizer_overrides_and_context(mock_memory_instance): + prompt_target = MagicMock(spec=PromptTarget) + prompt_target.get_identifier.return_value = get_mock_target_identifier("MockTarget") + prompt_target.send_prompt_async = AsyncMock( + return_value=[MessagePiece(role="assistant", original_value="first").to_message()] + ) + normalizer = PromptNormalizer() + conversation_id = "prepended-conversation" + message_normalizer = MagicMock(spec=MessageListNormalizer) + normalizer_overrides = {CapabilityName.EDITABLE_HISTORY: message_normalizer} + target_context = PrependedHistorySendContext( + conversation_id=conversation_id, + seed_message_ids=(uuid4(),), + replay_seed_each_send=False, + ) + + await normalizer.send_prompt_async( + message=Message.from_prompt(prompt="first request", role="user"), + target=prompt_target, + conversation_id=conversation_id, + normalizer_overrides=normalizer_overrides, + send_context=target_context, + ) + + call = prompt_target.send_prompt_async.await_args + assert call.kwargs["normalizer_overrides"] == normalizer_overrides + assert call.kwargs["send_context"] is target_context + + +async def test_send_prompt_async_conversion_failure_does_not_call_target(mock_memory_instance): + prompt_target = MagicMock(spec=PromptTarget) + prompt_target.get_identifier.return_value = get_mock_target_identifier("MockTarget") + prompt_target.send_prompt_async = AsyncMock() + conversation_id = "prepended-conversation" + target_context = PrependedHistorySendContext( + conversation_id=conversation_id, + seed_message_ids=(uuid4(),), + replay_seed_each_send=False, + ) + converter_config = ConverterConfiguration.from_converters(converters=[ContextFailingConverter()]) + + with pytest.raises(ValueError, match="conversion failed"): + await PromptNormalizer().send_prompt_async( + message=Message.from_prompt(prompt="request", role="user"), + target=prompt_target, + conversation_id=conversation_id, + request_converter_configurations=converter_config, + send_context=target_context, + ) + + prompt_target.send_prompt_async.assert_not_awaited() + mock_memory_instance.add_message_to_memory.assert_not_called() + + +async def test_send_prompt_async_target_failure_is_persisted(mock_memory_instance): + prompt_target = MagicMock(spec=PromptTarget) + prompt_target.get_identifier.return_value = get_mock_target_identifier("MockTarget") + prompt_target.send_prompt_async = AsyncMock(side_effect=ValueError("normalization failed")) + conversation_id = "prepended-conversation" + target_context = PrependedHistorySendContext( + conversation_id=conversation_id, + seed_message_ids=(uuid4(),), + replay_seed_each_send=False, + ) + + with pytest.raises(Exception, match="Error normalizing prompt with conversation ID"): + await PromptNormalizer().send_prompt_async( + message=Message.from_prompt(prompt="request", role="user"), + target=prompt_target, + conversation_id=conversation_id, + send_context=target_context, + ) + + assert target_context.provider_attempt_count == 0 + mock_memory_instance.add_message_to_memory.assert_not_called() + + +async def test_concurrent_rejection_is_not_misclassified_as_provider_attempt(mock_memory_instance): + conversation_id = "prepended-conversation" + seed = Message.from_prompt(prompt="seed", role="user") + seed.get_piece().conversation_id = conversation_id + mock_memory_instance.get_conversation_messages.return_value = [seed] + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities(supports_multi_turn=True)) + provider_started = asyncio.Event() + provider_release = asyncio.Event() + + async def wait_in_provider(*, normalized_conversation: list[Message]) -> list[Message]: + provider_started.set() + await provider_release.wait() + raise RuntimeError("provider failed") + + target._send_prompt_to_target_async = AsyncMock(side_effect=wait_in_provider) # type: ignore[method-assign] + target_context = PrependedHistorySendContext( + conversation_id=conversation_id, + seed_message_ids=(seed.get_piece().id,), + replay_seed_each_send=False, + ) + normalizer = PromptNormalizer() + first_send = asyncio.create_task( + normalizer.send_prompt_async( + message=Message.from_prompt(prompt="first request", role="user"), + target=target, + conversation_id=conversation_id, + send_context=target_context, + ) + ) + await provider_started.wait() + + with pytest.raises(Exception, match="Error normalizing prompt"): + await normalizer.send_prompt_async( + message=Message.from_prompt(prompt="concurrent request", role="user"), + target=target, + conversation_id=conversation_id, + send_context=target_context, + ) + + mock_memory_instance.add_message_to_memory.assert_not_called() + provider_release.set() + with pytest.raises(Exception, match="Error sending prompt"): + await first_send + persisted_values = [ + call.kwargs["request"].get_value() for call in mock_memory_instance.add_message_to_memory.call_args_list + ] + assert "concurrent request" not in persisted_values + + +async def test_send_prompt_async_empty_response_exception_is_persisted(mock_memory_instance): + prompt_target = MagicMock(spec=PromptTarget) + prompt_target.get_identifier.return_value = get_mock_target_identifier("MockTarget") + prompt_target.send_prompt_async = AsyncMock(side_effect=EmptyResponseException(message="normalization failed")) + conversation_id = "prepended-conversation" + response = await PromptNormalizer().send_prompt_async( + message=Message.from_prompt(prompt="request", role="user"), + target=prompt_target, + conversation_id=conversation_id, + ) + + assert response.get_piece().response_error == "empty" + assert mock_memory_instance.add_message_to_memory.call_count == 2 + + async def test_send_prompt_async_no_response_adds_memory(mock_memory_instance, seed_group): prompt_target = MagicMock() prompt_target.send_prompt_async = AsyncMock(return_value=None) diff --git a/tests/unit/prompt_target/target/test_conversation_normalization_pipeline.py b/tests/unit/prompt_target/target/test_conversation_normalization_pipeline.py index 91c33c31b3..5841d7f13d 100644 --- a/tests/unit/prompt_target/target/test_conversation_normalization_pipeline.py +++ b/tests/unit/prompt_target/target/test_conversation_normalization_pipeline.py @@ -152,6 +152,28 @@ def test_from_capabilities_uses_override_normalizer(): assert pipeline.normalizers[0] is mock_normalizer +def test_explicit_editable_history_override_works_with_sparse_raise_policy(): + mock_normalizer = MagicMock(spec=MessageListNormalizer) + caps = TargetCapabilities( + supports_multi_turn=True, + supports_editable_history=False, + supports_system_prompt=True, + supports_json_schema=True, + ) + sparse_policy = CapabilityHandlingPolicy( + behaviors={CapabilityName.SYSTEM_PROMPT: UnsupportedCapabilityBehavior.RAISE} + ) + + pipeline = ConversationNormalizationPipeline.from_capabilities( + capabilities=caps, + policy=sparse_policy, + normalizer_overrides={CapabilityName.EDITABLE_HISTORY: mock_normalizer}, + ) + + assert pipeline.normalizers == (mock_normalizer,) + assert pipeline.has_normalizer_for(capability=CapabilityName.EDITABLE_HISTORY) + + # --------------------------------------------------------------------------- # normalize_async — pass-through # --------------------------------------------------------------------------- diff --git a/tests/unit/prompt_target/target/test_normalize_async_integration.py b/tests/unit/prompt_target/target/test_normalize_async_integration.py index 407f1222a2..f6a61e1882 100644 --- a/tests/unit/prompt_target/target/test_normalize_async_integration.py +++ b/tests/unit/prompt_target/target/test_normalize_async_integration.py @@ -3,7 +3,9 @@ from __future__ import annotations +import asyncio import json +import logging from typing import TYPE_CHECKING from unittest.mock import AsyncMock, MagicMock, patch @@ -13,9 +15,22 @@ import pytest from openai.types.chat import ChatCompletion from openai.types.responses import ResponseOutputMessage, ResponseOutputText +from unit.mocks import MockPromptTarget +from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, +) +from pyrit.memory import CentralMemory from pyrit.memory.memory_interface import MemoryInterface -from pyrit.models import Message, MessagePiece +from pyrit.message_normalizer import ( + ConversationContextNormalizer, + HistorySquashNormalizer, + MessageListNormalizer, + MessageStringNormalizer, + TokenizerTemplateNormalizer, +) +from pyrit.models import ComponentIdentifier, Message, MessagePiece, PromptResponseError +from pyrit.prompt_normalizer import PromptNormalizer from pyrit.prompt_target import AzureMLChatTarget, OpenAIChatTarget from pyrit.prompt_target.common.target_capabilities import ( CapabilityHandlingPolicy, @@ -42,6 +57,34 @@ def _make_message(*, role: str, content: str, conversation_id: str = "conv1") -> return Message(message_pieces=[_make_message_piece(role=role, content=content, conversation_id=conversation_id)]) +def _make_normalizer_overrides( + *, + send_context: PrependedHistorySendContext, + formatter: MessageStringNormalizer | None = None, +) -> dict[CapabilityName, HistorySquashNormalizer]: + if not send_context.should_include_seed: + return {} + return { + CapabilityName.EDITABLE_HISTORY: HistorySquashNormalizer( + message_normalizer=formatter or ConversationContextNormalizer(), + expected_history_message_count=send_context.seed_message_count, + ) + } + + +def _make_prepended_history_send_context( + *, + prepended_messages: list[Message], + target_supports_multi_turn: bool = False, + conversation_id: str = "conv1", +) -> PrependedHistorySendContext: + return PrependedHistorySendContext( + conversation_id=conversation_id, + seed_message_ids=tuple(message.get_piece().id for message in prepended_messages), + replay_seed_each_send=not target_supports_multi_turn, + ) + + def _create_mock_chat_completion(content: str = "hi") -> MagicMock: mock = MagicMock(spec=ChatCompletion) mock.choices = [MagicMock()] @@ -489,3 +532,908 @@ async def test_get_normalized_conversation_passthrough_when_no_adaptation_needed assert result[0].get_value() == "be nice" assert result[1].get_piece().api_role == "user" assert result[1].get_value() == "hello" + + +# --------------------------------------------------------------------------- +# Prepended history adaptation +# --------------------------------------------------------------------------- + + +@pytest.mark.usefixtures("patch_central_database") +async def test_non_editable_target_adapts_prepended_history_without_mutating_memory(): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + + prepended_user = _make_message(role="user", content="original history") + prepended_user.get_piece().converted_value = "converted history" + prepended_user.get_piece().converter_identifiers = [ + ComponentIdentifier(class_name="TestConverter", class_module="tests") + ] + prepended_assistant = _make_message(role="simulated_assistant", content="assistant history") + live_request = _make_message(role="user", content="original live") + live_request.get_piece().converted_value = "converted live" + memory_messages: MutableSequence[Message] = [prepended_user, prepended_assistant] + + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.return_value = memory_messages + target._memory = mock_memory + target_context = _make_prepended_history_send_context(prepended_messages=list(memory_messages)) + normalizer_overrides = _make_normalizer_overrides(send_context=target_context) + + result = await target._get_normalized_conversation_async( + message=live_request, + normalizer_overrides=normalizer_overrides, + send_context=target_context, + ) + + assert len(result) == 1 + assert result[0].get_piece().original_value == ( + "Turn 1:\nuser: original history\nassistant: assistant history\nTurn 2:\nuser: original live" + ) + assert result[0].get_piece().converted_value == ( + "Turn 1:\nuser: converted history\nassistant: assistant history\nTurn 2:\nuser: converted live" + ) + assert len(memory_messages) == 2 + assert memory_messages[0].get_piece().original_value == "original history" + + +@pytest.mark.usefixtures("patch_central_database") +async def test_non_editable_target_preserves_system_history_and_multimodal_live_request(): + target = MockPromptTarget() + target._configuration = TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_message_pieces=True, + input_modalities=frozenset({frozenset({"text", "image_path"})}), + ) + ) + system_message = _make_message(role="system", content="Describe images precisely") + live_request = Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="diagram.png", + converted_value="diagram.png", + original_value_data_type="image_path", + converted_value_data_type="image_path", + conversation_id="conv1", + ) + ] + ) + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.return_value = [system_message] + target._memory = mock_memory + target_context = _make_prepended_history_send_context(prepended_messages=[system_message]) + normalizer_overrides = _make_normalizer_overrides(send_context=target_context) + + result = await target._get_normalized_conversation_async( + message=live_request, + normalizer_overrides=normalizer_overrides, + send_context=target_context, + ) + + assert len(result) == 1 + assert len(result[0].message_pieces) == 2 + assert result[0].message_pieces[0].converted_value == "Turn 1:\nuser: Describe images precisely" + assert result[0].message_pieces[1].converted_value == "diagram.png" + assert result[0].message_pieces[1].converted_value_data_type == "image_path" + + +@pytest.mark.usefixtures("patch_central_database") +async def test_editable_history_override_runs_before_system_prompt_adaptation(): + target = MockPromptTarget() + target._configuration = TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_turn=True, + supports_system_prompt=False, + supports_editable_history=False, + ), + policy=CapabilityHandlingPolicy( + behaviors={ + CapabilityName.MULTI_TURN: UnsupportedCapabilityBehavior.RAISE, + CapabilityName.SYSTEM_PROMPT: UnsupportedCapabilityBehavior.ADAPT, + } + ), + ) + prepended = _make_message(role="system", content="system") + live_request = _make_message(role="user", content="live") + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.return_value = [prepended] + target._memory = mock_memory + target_context = _make_prepended_history_send_context( + prepended_messages=[prepended], + target_supports_multi_turn=True, + ) + + result = await target._get_normalized_conversation_async( + message=live_request, + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, + ) + + assert [message.get_value() for message in result] == [ + "Turn 1:\nuser: ### Instructions ###\n\nsystem\n\n######\n\nlive" + ] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_first_turn_normalization_preserves_live_multimodal_piece_order(): + target = MockPromptTarget() + target._configuration = TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_message_pieces=True, + input_modalities=frozenset({frozenset({"text", "image_path"})}), + ) + ) + live_request = Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="diagram.png", + converted_value="diagram.png", + original_value_data_type="image_path", + converted_value_data_type="image_path", + conversation_id="conv1", + sequence=0, + prompt_metadata={"piece": "image", "image-only": "preserved-on-image"}, + ), + MessagePiece( + role="user", + original_value="What does this show?", + converted_value="What does this show?", + original_value_data_type="text", + converted_value_data_type="text", + conversation_id="conv1", + sequence=0, + prompt_metadata={"piece": "live-text", "trace": "preserved"}, + ), + ] + ) + mock_memory = MagicMock(spec=MemoryInterface) + prepended = _make_message(role="user", content="prepended") + mock_memory.get_conversation_messages.return_value = [prepended] + target._memory = mock_memory + target_context = _make_prepended_history_send_context(prepended_messages=[prepended]) + normalizer_overrides = _make_normalizer_overrides(send_context=target_context) + + result = await target._get_normalized_conversation_async( + message=live_request, + normalizer_overrides=normalizer_overrides, + send_context=target_context, + ) + + assert [piece.converted_value_data_type for piece in result[0].message_pieces] == ["image_path", "text"] + assert result[0].message_pieces[0].converted_value == "diagram.png" + assert "What does this show?" in result[0].message_pieces[1].converted_value + assert result[0].message_pieces[0].prompt_metadata == { + "piece": "image", + "image-only": "preserved-on-image", + } + assert result[0].message_pieces[1].prompt_metadata == {"piece": "live-text", "trace": "preserved"} + + +@pytest.mark.usefixtures("patch_central_database") +async def test_history_squash_does_not_restore_adapted_json_schema_metadata(): + target = MockPromptTarget() + target._configuration = TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_message_pieces=True, + supports_system_prompt=True, + input_modalities=frozenset({frozenset({"text", "image_path"})}), + ) + ) + prepended = _make_message(role="user", content="prepended") + live_request = Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="diagram.png", + original_value_data_type="image_path", + conversation_id="conv1", + prompt_metadata={"piece": "image"}, + ), + MessagePiece( + role="user", + original_value="Describe this as JSON", + conversation_id="conv1", + prompt_metadata={ + "response_format": "json", + "json_schema": {"type": "object", "properties": {"description": {"type": "string"}}}, + }, + ), + ] + ) + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.return_value = [prepended] + target._memory = mock_memory + target_context = _make_prepended_history_send_context(prepended_messages=[prepended]) + + result = await target._get_normalized_conversation_async( + message=live_request, + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, + ) + + image_piece, text_piece = result[0].message_pieces + assert image_piece.prompt_metadata == {"piece": "image"} + assert text_piece.prompt_metadata == {"response_format": "json"} + assert '"description"' in text_piece.converted_value + + +@pytest.mark.usefixtures("patch_central_database") +async def test_custom_normalizer_output_metadata_is_authoritative(): + target = MockPromptTarget() + target._configuration = TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_turn=True, + supports_multi_message_pieces=True, + supports_system_prompt=True, + ) + ) + source = _make_message(role="user", content="source") + source.get_piece().prompt_metadata = {"source-only": "must not be restored"} + replacement = Message.from_prompt(prompt="replacement", role="user") + normalizer = MagicMock(spec=MessageListNormalizer) + normalizer.normalize_async = AsyncMock(return_value=[replacement]) + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.return_value = [] + target._memory = mock_memory + + result = await target._get_normalized_conversation_async( + message=source, + normalizer_overrides={CapabilityName.EDITABLE_HISTORY: normalizer}, + ) + + assert result[0].get_piece().conversation_id == "conv1" + assert result[0].get_piece().prompt_metadata == {} + + +@pytest.mark.usefixtures("patch_central_database") +async def test_prepended_history_adapter_is_used_only_when_explicitly_passed(): + target = MockPromptTarget() + target._configuration = TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_turn=True, + supports_system_prompt=True, + ) + ) + prepended = _make_message(role="user", content="prepended") + prior_live = _make_message(role="user", content="first live") + prior_response = _make_message(role="assistant", content="first response") + second_live = _make_message(role="user", content="second live") + + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.side_effect = [ + [prepended], + [prepended, prior_live, prior_response], + ] + target._memory = mock_memory + target_context = _make_prepended_history_send_context( + prepended_messages=[prepended], + target_supports_multi_turn=True, + ) + + await target.send_prompt_async( + message=prior_live, + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, + ) + await target.send_prompt_async( + message=second_live, + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, + ) + + assert target.prompt_sent == ["Turn 1:\nuser: prepended\nTurn 2:\nuser: first live", "second live"] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_non_editable_multi_turn_target_retains_history_after_response(): + target = MockPromptTarget() + target._configuration = TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_turn=True, + supports_system_prompt=True, + ) + ) + prepended = _make_message(role="user", content="prepended") + first_live = _make_message(role="user", content="first live") + second_live = _make_message(role="user", content="second live") + prior_response = _make_message(role="assistant", content="first response") + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.side_effect = [ + [prepended], + [prepended, first_live, prior_response], + ] + target._memory = mock_memory + target._send_prompt_to_target_async = AsyncMock( # type: ignore[method-assign] + return_value=[_make_message(role="assistant", content="response")] + ) + target_context = _make_prepended_history_send_context( + prepended_messages=[prepended], + target_supports_multi_turn=True, + ) + + await target.send_prompt_async( + message=first_live, + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, + ) + await target.send_prompt_async( + message=second_live, + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, + ) + + first_payload, second_payload = target._send_prompt_to_target_async.await_args_list + assert len(first_payload.kwargs["normalized_conversation"]) == 1 + assert [message.get_value() for message in second_payload.kwargs["normalized_conversation"]] == [ + "prepended", + "first live", + "first response", + "second live", + ] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_stateless_target_replays_only_seed_and_current_request(): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + prepended = _make_message(role="user", content="prepended") + first_live = _make_message(role="user", content="first live") + first_response = _make_message(role="assistant", content="first response") + second_live = _make_message(role="user", content="second live") + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.side_effect = [ + [prepended], + [prepended, first_live, first_response], + ] + target._memory = mock_memory + target_context = _make_prepended_history_send_context(prepended_messages=[prepended]) + + await target.send_prompt_async( + message=first_live, + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, + ) + await target.send_prompt_async( + message=second_live, + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, + ) + + assert target.prompt_sent == [ + "Turn 1:\nuser: prepended\nTurn 2:\nuser: first live", + "Turn 1:\nuser: prepended\nTurn 2:\nuser: second live", + ] + + +@pytest.mark.parametrize("response_error", ["blocked", "empty", "processing"]) +@pytest.mark.usefixtures("patch_central_database") +async def test_stateful_target_consumes_seed_after_provider_outcome(response_error: PromptResponseError): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities(supports_multi_turn=True)) + prepended = _make_message(role="user", content="prepended") + first_live = _make_message(role="user", content="first live") + provider_response = _make_message(role="assistant", content=f"{response_error} response") + provider_response.get_piece().response_error = response_error + second_live = _make_message(role="user", content="second live") + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.side_effect = [ + [prepended], + [prepended, first_live, provider_response], + ] + target._memory = mock_memory + target._send_prompt_to_target_async = AsyncMock( # type: ignore[method-assign] + side_effect=[[provider_response], [_make_message(role="assistant", content="second response")]] + ) + target_context = _make_prepended_history_send_context( + prepended_messages=[prepended], + target_supports_multi_turn=True, + ) + + await target.send_prompt_async( + message=first_live, + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, + ) + await target.send_prompt_async( + message=second_live, + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, + ) + + first_payload, second_payload = target._send_prompt_to_target_async.await_args_list + assert "prepended" in first_payload.kwargs["normalized_conversation"][0].get_value() + assert second_payload.kwargs["normalized_conversation"][-1].get_value() == "second live" + + +@pytest.mark.usefixtures("patch_central_database") +async def test_non_editable_target_uses_custom_prepended_formatter(): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + mock_memory = MagicMock(spec=MemoryInterface) + prepended = _make_message(role="user", content="prepended") + mock_memory.get_conversation_messages.return_value = [prepended] + target._memory = mock_memory + formatter = MagicMock(spec=MessageStringNormalizer) + formatter.normalize_string_async = AsyncMock(return_value="CUSTOM HISTORY") + target_context = _make_prepended_history_send_context(prepended_messages=[prepended]) + normalizer_overrides = _make_normalizer_overrides( + send_context=target_context, + formatter=formatter, + ) + + result = await target._get_normalized_conversation_async( + message=_make_message(role="user", content="live"), + normalizer_overrides=normalizer_overrides, + send_context=target_context, + ) + + assert result[0].get_value() == "CUSTOM HISTORY" + formatter.normalize_string_async.assert_awaited_once() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_non_editable_target_rejects_non_text_converted_prepended_history(): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + prepended = _make_message(role="user", content="original") + prepended.get_piece().converted_value = "converted.png" + prepended.get_piece().converted_value_data_type = "image_path" + prepended.get_piece().converter_identifiers = [ + ComponentIdentifier(class_name="ImageConverter", class_module="tests") + ] + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.return_value = [prepended] + target._memory = mock_memory + target_context = _make_prepended_history_send_context(prepended_messages=[prepended]) + normalizer_overrides = _make_normalizer_overrides(send_context=target_context) + + with pytest.raises(ValueError, match="non-text output types.*image_path"): + await target.send_prompt_async( + message=_make_message(role="user", content="live"), + normalizer_overrides=normalizer_overrides, + send_context=target_context, + ) + + +@pytest.mark.usefixtures("patch_central_database") +async def test_non_editable_target_rejects_same_modality_non_text_conversion(): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + prepended = Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="original.png", + converted_value="converted.png", + original_value_data_type="image_path", + converted_value_data_type="image_path", + conversation_id="conv1", + ) + ] + ) + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.return_value = [prepended] + target._memory = mock_memory + target_context = _make_prepended_history_send_context(prepended_messages=[prepended]) + normalizer_overrides = _make_normalizer_overrides(send_context=target_context) + + with pytest.raises(ValueError, match="non-text output types.*image_path"): + await target.send_prompt_async( + message=_make_message(role="user", content="live"), + normalizer_overrides=normalizer_overrides, + send_context=target_context, + ) + + +@pytest.mark.usefixtures("patch_central_database") +async def test_non_editable_target_allows_preexisting_non_text_history_with_converter_provenance(): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + prepended = Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="existing.png", + converted_value="existing.png", + original_value_data_type="image_path", + converted_value_data_type="image_path", + conversation_id="conv1", + converter_identifiers=[ComponentIdentifier(class_name="PriorConverter", class_module="tests")], + ) + ] + ) + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.return_value = [prepended] + target._memory = mock_memory + target_context = _make_prepended_history_send_context(prepended_messages=[prepended]) + normalizer_overrides = _make_normalizer_overrides(send_context=target_context) + + result = await target._get_normalized_conversation_async( + message=_make_message(role="user", content="live"), + normalizer_overrides=normalizer_overrides, + send_context=target_context, + ) + + assert result[0].get_value() == "Turn 1:\nuser: [Image_path]\nTurn 2:\nuser: live" + + +@pytest.mark.usefixtures("patch_central_database") +async def test_non_editable_target_warns_when_non_text_history_becomes_a_placeholder( + caplog: pytest.LogCaptureFixture, +) -> None: + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + prepended = Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="existing.png", + converted_value="existing.png", + original_value_data_type="image_path", + converted_value_data_type="image_path", + conversation_id="conv1", + ) + ] + ) + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.return_value = [prepended] + target._memory = mock_memory + target_context = _make_prepended_history_send_context(prepended_messages=[prepended]) + + with caplog.at_level(logging.WARNING, logger="pyrit.message_normalizer.history_squash_normalizer"): + await target._get_normalized_conversation_async( + message=_make_message(role="user", content="live"), + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, + ) + + assert "image_path" in caplog.text + assert "text placeholders" in caplog.text + + +@pytest.mark.usefixtures("patch_central_database") +async def test_target_normalization_failure_can_be_retried(): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + mock_memory = MagicMock(spec=MemoryInterface) + prepended = _make_message(role="user", content="prepended") + mock_memory.get_conversation_messages.return_value = [prepended] + target._memory = mock_memory + formatter = MagicMock(spec=MessageStringNormalizer) + formatter.normalize_string_async = AsyncMock(side_effect=[ValueError("format failed"), "formatted request"]) + target_context = _make_prepended_history_send_context(prepended_messages=[prepended]) + normalizer_overrides = _make_normalizer_overrides( + send_context=target_context, + formatter=formatter, + ) + live_request = _make_message(role="user", content="live") + + with pytest.raises(ValueError, match="format failed"): + await target.send_prompt_async( + message=live_request, + normalizer_overrides=normalizer_overrides, + send_context=target_context, + ) + + await target.send_prompt_async( + message=live_request, + normalizer_overrides=normalizer_overrides, + send_context=target_context, + ) + assert target.prompt_sent == ["formatted request"] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_prompt_normalizer_retry_excludes_persisted_processing_exchange(): + target = MockPromptTarget() + prompt_normalizer = PromptNormalizer() + conversation_id = "processing-retry" + successful_response = _make_message(role="assistant", content="successful response") + retry_response = _make_message(role="assistant", content="retry response") + + with patch.object(target, "_send_prompt_to_target_async", new_callable=AsyncMock) as send: + send.side_effect = [[successful_response], RuntimeError("private provider failure"), [retry_response]] + await prompt_normalizer.send_prompt_async( + message=_make_message(role="user", content="successful request"), + target=target, + conversation_id=conversation_id, + ) + with pytest.raises(Exception, match="Error sending prompt"): + await prompt_normalizer.send_prompt_async( + message=_make_message(role="user", content="failed request"), + target=target, + conversation_id=conversation_id, + ) + await prompt_normalizer.send_prompt_async( + message=_make_message(role="user", content="retry"), + target=target, + conversation_id=conversation_id, + ) + + retry_payload = send.await_args_list[2].kwargs["normalized_conversation"] + assert [message.get_value() for message in retry_payload] == [ + "successful request", + "successful response", + "retry", + ] + persisted = list(CentralMemory.get_memory_instance().get_conversation_messages(conversation_id=conversation_id)) + processing_index = next( + index for index, message in enumerate(persisted) if message.get_piece().response_error == "processing" + ) + failed_request = persisted[processing_index - 1].get_piece() + processing_error = persisted[processing_index].get_piece() + assert failed_request.api_role == "user" + assert processing_error.original_prompt_id != failed_request.id + assert processing_error.sequence == failed_request.sequence + 1 + + +@pytest.mark.usefixtures("patch_central_database") +async def test_prompt_normalizer_retry_excludes_persisted_unknown_exchange(): + target = MockPromptTarget() + prompt_normalizer = PromptNormalizer() + unknown_response = _make_message(role="assistant", content="unknown provider failure") + unknown_response.get_piece().response_error = "unknown" + retry_response = _make_message(role="assistant", content="retry response") + + with patch.object(target, "_send_prompt_to_target_async", new_callable=AsyncMock) as send: + send.side_effect = [[unknown_response], [retry_response]] + await prompt_normalizer.send_prompt_async( + message=_make_message(role="user", content="failed request"), + target=target, + conversation_id="unknown-retry", + ) + await prompt_normalizer.send_prompt_async( + message=_make_message(role="user", content="retry"), + target=target, + conversation_id="unknown-retry", + ) + + retry_payload = send.await_args_list[1].kwargs["normalized_conversation"] + assert [message.get_value() for message in retry_payload] == ["retry"] + + +@pytest.mark.parametrize("response_error", ["blocked", "empty"]) +@pytest.mark.usefixtures("patch_central_database") +async def test_prompt_normalizer_retains_provider_round_trip(response_error: PromptResponseError): + target = MockPromptTarget() + prompt_normalizer = PromptNormalizer() + provider_response = _make_message(role="assistant", content=f"{response_error} response") + provider_response.get_piece().response_error = response_error + next_response = _make_message(role="assistant", content="next response") + + with patch.object(target, "_send_prompt_to_target_async", new_callable=AsyncMock) as send: + send.side_effect = [[provider_response], [next_response]] + await prompt_normalizer.send_prompt_async( + message=_make_message(role="user", content="first request"), + target=target, + conversation_id=f"{response_error}-round-trip", + ) + await prompt_normalizer.send_prompt_async( + message=_make_message(role="user", content="second request"), + target=target, + conversation_id=f"{response_error}-round-trip", + ) + + second_payload = send.await_args_list[1].kwargs["normalized_conversation"] + assert [message.get_value() for message in second_payload] == [ + "first request", + f"{response_error} response", + "second request", + ] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_target_normalization_cancellation_propagates(): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + mock_memory = MagicMock(spec=MemoryInterface) + prepended = _make_message(role="user", content="prepended") + mock_memory.get_conversation_messages.return_value = [prepended] + target._memory = mock_memory + formatter = MagicMock(spec=MessageStringNormalizer) + formatter.normalize_string_async = AsyncMock(side_effect=asyncio.CancelledError()) + target_context = _make_prepended_history_send_context(prepended_messages=[prepended]) + normalizer_overrides = _make_normalizer_overrides( + send_context=target_context, + formatter=formatter, + ) + + with pytest.raises(asyncio.CancelledError): + await target.send_prompt_async( + message=_make_message(role="user", content="live"), + normalizer_overrides=normalizer_overrides, + send_context=target_context, + ) + assert target_context.provider_attempt_count == 0 + assert not target_context.is_seed_consumed + + +@pytest.mark.usefixtures("patch_central_database") +async def test_rate_limit_cancellation_does_not_consume_stateful_seed(): + target = MockPromptTarget(rpm=1) + target._configuration = TargetConfiguration(capabilities=TargetCapabilities(supports_multi_turn=True)) + mock_memory = MagicMock(spec=MemoryInterface) + prepended = _make_message(role="user", content="prepended") + mock_memory.get_conversation_messages.return_value = [prepended] + target._memory = mock_memory + target_context = _make_prepended_history_send_context( + prepended_messages=[prepended], + target_supports_multi_turn=True, + ) + sleep_started = asyncio.Event() + sleep_release = asyncio.Event() + + async def wait_for_rate_limit(delay: float) -> None: + sleep_started.set() + await sleep_release.wait() + + with patch("pyrit.prompt_target.common.utils.asyncio.sleep", side_effect=wait_for_rate_limit): + send_task = asyncio.create_task( + target.send_prompt_async( + message=_make_message(role="user", content="live"), + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, + ) + ) + await sleep_started.wait() + send_task.cancel() + with pytest.raises(asyncio.CancelledError): + await send_task + + assert target_context.provider_attempt_count == 0 + assert not target_context.is_seed_consumed + + +@pytest.mark.usefixtures("patch_central_database") +async def test_provider_failure_propagates_after_normalization(): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities(supports_multi_turn=True)) + mock_memory = MagicMock(spec=MemoryInterface) + prepended = _make_message(role="user", content="prepended") + mock_memory.get_conversation_messages.return_value = [prepended] + target._memory = mock_memory + target._send_prompt_to_target_async = AsyncMock(side_effect=RuntimeError("provider failed")) # type: ignore[method-assign] + target_context = _make_prepended_history_send_context( + prepended_messages=[prepended], + target_supports_multi_turn=True, + ) + normalizer_overrides = _make_normalizer_overrides(send_context=target_context) + + with pytest.raises(RuntimeError, match="provider failed"): + await target.send_prompt_async( + message=_make_message(role="user", content="live"), + normalizer_overrides=normalizer_overrides, + send_context=target_context, + ) + assert target_context.provider_attempt_count == 1 + assert target_context.is_seed_consumed + + +@pytest.mark.usefixtures("patch_central_database") +async def test_provider_cancellation_consumes_stateful_seed_and_releases_context(): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities(supports_multi_turn=True)) + prepended = _make_message(role="user", content="prepended") + first_live = _make_message(role="user", content="first live") + second_live = _make_message(role="user", content="second live") + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.side_effect = [ + [prepended], + [prepended, first_live], + ] + target._memory = mock_memory + provider_started = asyncio.Event() + provider_release = asyncio.Event() + + async def wait_in_provider(*, normalized_conversation: list[Message]) -> list[Message]: + provider_started.set() + await provider_release.wait() + return [_make_message(role="assistant", content="response")] + + target._send_prompt_to_target_async = AsyncMock(side_effect=wait_in_provider) # type: ignore[method-assign] + target_context = _make_prepended_history_send_context( + prepended_messages=[prepended], + target_supports_multi_turn=True, + ) + first_send = asyncio.create_task( + target.send_prompt_async( + message=first_live, + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, + ) + ) + await provider_started.wait() + first_send.cancel() + with pytest.raises(asyncio.CancelledError): + await first_send + + assert target_context.is_seed_consumed + assert target_context.provider_attempt_count == 1 + + target._send_prompt_to_target_async = AsyncMock( # type: ignore[method-assign] + return_value=[_make_message(role="assistant", content="second response")] + ) + await target.send_prompt_async( + message=second_live, + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, + ) + payload = target._send_prompt_to_target_async.await_args.kwargs["normalized_conversation"] + assert [message.get_value() for message in payload] == ["prepended", "first live", "second live"] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_concurrent_sends_with_one_context_are_rejected(): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + mock_memory = MagicMock(spec=MemoryInterface) + prepended = _make_message(role="user", content="prepended") + mock_memory.get_conversation_messages.return_value = [prepended] + target._memory = mock_memory + started = asyncio.Event() + release = asyncio.Event() + + async def wait_to_format(messages: list[Message]) -> str: + started.set() + await release.wait() + return "formatted request" + + formatter = MagicMock(spec=MessageStringNormalizer) + formatter.normalize_string_async = AsyncMock(side_effect=wait_to_format) + target_context = _make_prepended_history_send_context(prepended_messages=[prepended]) + normalizer_overrides = _make_normalizer_overrides( + send_context=target_context, + formatter=formatter, + ) + first_send = asyncio.create_task( + target.send_prompt_async( + message=_make_message(role="user", content="first"), + normalizer_overrides=normalizer_overrides, + send_context=target_context, + ) + ) + await started.wait() + try: + with pytest.raises(RuntimeError, match="Concurrent sends"): + await target.send_prompt_async( + message=_make_message(role="user", content="second"), + normalizer_overrides=normalizer_overrides, + send_context=target_context, + ) + finally: + release.set() + + await first_send + assert target.prompt_sent == ["formatted request"] + formatter.normalize_string_async.assert_awaited_once() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_tokenizer_formatter_receives_live_request_before_generation_prompt(): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + mock_memory = MagicMock(spec=MemoryInterface) + prepended = _make_message(role="user", content="prepended") + mock_memory.get_conversation_messages.return_value = [prepended] + target._memory = mock_memory + tokenizer = MagicMock() + tokenizer.apply_chat_template.return_value = "TOKENIZED REQUEST" + formatter = TokenizerTemplateNormalizer(tokenizer=tokenizer) + target_context = _make_prepended_history_send_context(prepended_messages=[prepended]) + normalizer_overrides = _make_normalizer_overrides( + send_context=target_context, + formatter=formatter, + ) + + result = await target._get_normalized_conversation_async( + message=_make_message(role="user", content="live"), + normalizer_overrides=normalizer_overrides, + send_context=target_context, + ) + + tokenizer_messages = tokenizer.apply_chat_template.call_args.args[0] + assert tokenizer_messages[-1] == {"role": "user", "content": "live"} + assert tokenizer.apply_chat_template.call_args.kwargs["add_generation_prompt"] is True + assert result[0].get_value() == "TOKENIZED REQUEST" diff --git a/tests/unit/prompt_target/target/test_playwright_copilot_target.py b/tests/unit/prompt_target/target/test_playwright_copilot_target.py index f933f4cd13..5e22f125ad 100644 --- a/tests/unit/prompt_target/target/test_playwright_copilot_target.py +++ b/tests/unit/prompt_target/target/test_playwright_copilot_target.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import asyncio from unittest.mock import AsyncMock, MagicMock, call, patch import pytest @@ -296,11 +297,13 @@ async def __aexit__(self, *args): mock_page.expect_file_chooser = MagicMock(return_value=MockFileChooserContextManager(mock_file_chooser)) - await target._upload_image_async("/path/to/image.jpg") + with patch.object(target, "_mark_provider_attempted") as mark_provider_attempted: + await target._upload_image_async("/path/to/image.jpg") dropdown_locator.click.assert_awaited_once() file_picker_locator.wait_for.assert_awaited_once_with(state="visible", timeout=5000) file_picker_locator.click.assert_awaited_once() + mark_provider_attempted.assert_called_once_with() mock_file_chooser.set_files.assert_awaited_once_with("/path/to/image.jpg") async def test_wait_for_response_async_success(self, mock_page): @@ -314,10 +317,14 @@ async def test_wait_for_response_async_success(self, mock_page): mock_page.query_selector_all.return_value = [AsyncMock()] # Mock response extraction - with patch.object(target, "_extract_multimodal_content_async", return_value="Response text") as mock_extract: + with ( + patch.object(target, "_extract_multimodal_content_async", return_value="Response text") as mock_extract, + patch.object(target, "_mark_provider_attempted") as mark_provider_attempted, + ): result = await target._wait_for_response_async(selectors) assert result == "Response text" + mark_provider_attempted.assert_called_once_with() mock_page.click.assert_awaited_once_with(selectors.send_button_selector) mock_extract.assert_awaited_once() @@ -372,6 +379,7 @@ async def test_interact_with_copilot_async_multimodal(self, mock_page, multimoda # Mock the helper methods with ( + patch.object(target, "_clear_text_input_async") as mock_clear_text, patch.object(target, "_send_text_async") as mock_send_text, patch.object(target, "_upload_image_async") as mock_upload_image, patch.object(target, "_wait_for_response_async", return_value="AI response") as mock_wait, @@ -379,11 +387,36 @@ async def test_interact_with_copilot_async_multimodal(self, mock_page, multimoda result = await target._interact_with_copilot_async(multimodal_request) # Verify text and image handling + mock_clear_text.assert_awaited_once() mock_send_text.assert_awaited_once() mock_upload_image.assert_awaited_once_with("/path/to/image.jpg") mock_wait.assert_awaited_once() assert result == "AI response" + async def test_interact_cancellation_while_staging_text_does_not_mark_provider_attempt( + self, mock_page, text_request_piece + ): + target = PlaywrightCopilotTarget(page=mock_page) + request = Message(message_pieces=[text_request_piece]) + staging_started = asyncio.Event() + + async def stage_text(*, text: str, input_selector: str) -> None: + staging_started.set() + await asyncio.Event().wait() + + with ( + patch.object(target, "_clear_text_input_async"), + patch.object(target, "_send_text_async", side_effect=stage_text), + patch.object(target, "_mark_provider_attempted") as mark_provider_attempted, + ): + task = asyncio.create_task(target._interact_with_copilot_async(request)) + await staging_started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + mark_provider_attempted.assert_not_called() + def test_constants(self, mock_page): """Test that class constants are defined correctly.""" target = PlaywrightCopilotTarget(page=mock_page) diff --git a/tests/unit/prompt_target/target/test_prompt_target.py b/tests/unit/prompt_target/target/test_prompt_target.py index 3de86263e7..453870e382 100644 --- a/tests/unit/prompt_target/target/test_prompt_target.py +++ b/tests/unit/prompt_target/target/test_prompt_target.py @@ -153,7 +153,7 @@ async def test_send_prompt_async_with_delay( # --------------------------------------------------------------------------- -# _propagate_lineage — metadata preservation after normalization +# Normalizer metadata and conversation ownership # --------------------------------------------------------------------------- _LINEAGE_CONVERSATION_ID = "original-conv-id-12345" @@ -192,8 +192,8 @@ def _make_mock_chat_completion(content: str = "response") -> MagicMock: @pytest.mark.usefixtures("patch_central_database") async def test_history_squash_preserves_metadata_on_normalized_message(): """ - After history squash, _propagate_lineage should restore the original request's - conversation ID and prompt metadata onto the squashed message. + History squash preserves the current request's metadata, and the target stamps + the active conversation ID on its output. """ target = OpenAIChatTarget( model_name="gpt-4o", @@ -236,8 +236,7 @@ async def test_history_squash_preserves_metadata_on_normalized_message(): async def test_response_preserves_metadata_after_history_squash(): """ End-to-end: after history squash the response must carry the original - request's conversation ID and prompt metadata, not the random values - created by the normalizer. + request's conversation ID and prompt metadata. """ target = OpenAIChatTarget( model_name="gpt-4o", @@ -282,8 +281,8 @@ async def test_response_preserves_metadata_after_history_squash(): @pytest.mark.usefixtures("patch_central_database") async def test_system_squash_preserves_metadata(): """ - GenericSystemSquashNormalizer also creates messages via Message.from_prompt. - _propagate_lineage should restore the original metadata after system squash too. + GenericSystemSquashNormalizer preserves the current request's metadata when + it builds the replacement user message. """ target = OpenAIChatTarget( model_name="gpt-4o", @@ -324,10 +323,10 @@ async def test_system_squash_preserves_metadata(): @pytest.mark.usefixtures("patch_central_database") -async def test_history_squash_propagates_lineage_to_all_pieces(): +async def test_history_squash_preserves_metadata_on_all_output_pieces(): """ - When the squashed message contains multiple pieces, _propagate_lineage - must stamp every piece — not just the first one. + Every piece produced by history squash keeps the current request's metadata + and receives the active conversation ID. """ target = OpenAIChatTarget( model_name="gpt-4o", @@ -372,12 +371,10 @@ async def test_history_squash_propagates_lineage_to_all_pieces(): @pytest.mark.usefixtures("patch_central_database") -async def test_conversation_id_stamped_on_all_but_full_lineage_only_on_last(): +async def test_conversation_id_stamped_without_merging_normalizer_output_metadata(): """ - conversation_id is stamped on every normalized message (including new ones - created by the normalizer). Full lineage is only propagated to the last - message. Earlier messages keep their own metadata. A warning is logged when - the normalizer increases message count. + The target stamps conversation_id on every normalized output while leaving + each normalizer-produced message's metadata authoritative. """ target = OpenAIChatTarget( model_name="gpt-4o", @@ -405,14 +402,19 @@ async def test_conversation_id_stamped_on_all_but_full_lineage_only_on_last(): converted_value_data_type="text", ) new_msg = Message(message_pieces=[new_piece]) + replacement_piece = MessagePiece( + role="user", + conversation_id="another-normalizer-uuid", + original_value="replacement", + converted_value="replacement", + original_value_data_type="text", + converted_value_data_type="text", + ) + replacement_msg = Message(message_pieces=[replacement_piece]) with patch.object(target.configuration, "normalize_async", new_callable=AsyncMock) as mock_normalize: - mock_normalize.return_value = [history_msg, new_msg, user_msg] - - import logging - - with patch.object(logging.getLogger("pyrit.prompt_target.common.prompt_target"), "warning") as mock_warn: - normalized = await target._get_normalized_conversation_async(message=user_msg) + mock_normalize.return_value = [history_msg, new_msg, replacement_msg] + normalized = await target._get_normalized_conversation_async(message=user_msg) # All messages should carry the correct conversation_id. for msg in normalized: @@ -422,24 +424,17 @@ async def test_conversation_id_stamped_on_all_but_full_lineage_only_on_last(): # History message's other metadata should be untouched. assert normalized[0].message_pieces[0].prompt_metadata == {"original": "history_meta"} - # New middle message should NOT have full lineage overwritten. + # New messages keep exactly the metadata produced by the normalizer. assert normalized[1].message_pieces[0].prompt_metadata == {} - - # Last message should carry full lineage. - last_piece = normalized[-1].message_pieces[0] - assert last_piece.prompt_metadata == _LINEAGE_PROMPT_METADATA - - # Warning should fire because message count increased (2 → 3). - mock_warn.assert_called_once() + assert normalized[-1].message_pieces[0].prompt_metadata == {} @pytest.mark.usefixtures("patch_central_database") -async def test_json_schema_stripped_for_non_schema_target_survives_lineage(): +async def test_json_schema_stripped_for_non_schema_target_remains_authoritative(): """ Regression: for a non-schema target (default ADAPT) the embedded json_schema is - removed by JsonSchemaNormalizer and must NOT be re-introduced by - _propagate_lineage copying the original (unstripped) request metadata back onto - the normalized message. + removed by JsonSchemaNormalizer and must not be reintroduced from the source + request metadata. """ target = OpenAIChatTarget( model_name="gpt-4o", @@ -471,11 +466,10 @@ async def test_json_schema_stripped_for_non_schema_target_survives_lineage(): @pytest.mark.usefixtures("patch_central_database") -async def test_json_schema_only_metadata_fully_stripped_survives_lineage(): +async def test_json_schema_only_metadata_fully_stripped_remains_authoritative(): """ Regression: even when json_schema is the ONLY metadata key, the strip leaves empty - metadata and _propagate_lineage must not restore the original json_schema (the piece - is the same logical piece, identified by id, so its stripped metadata is authoritative). + metadata and the target must not restore the original json_schema. """ target = OpenAIChatTarget( model_name="gpt-4o", diff --git a/tests/unit/prompt_target/target/test_target_capabilities.py b/tests/unit/prompt_target/target/test_target_capabilities.py index 30c47ce1ca..5ef35c81b8 100644 --- a/tests/unit/prompt_target/target/test_target_capabilities.py +++ b/tests/unit/prompt_target/target/test_target_capabilities.py @@ -71,11 +71,11 @@ def test_capability_handling_policy_get_behavior_for_all_default_keys(self): def test_capability_handling_policy_rejects_capability_without_policy(self): policy = CapabilityHandlingPolicy() - with pytest.raises(KeyError, match="No policy for capability 'supports_editable_history'"): - policy.get_behavior(capability=CapabilityName.EDITABLE_HISTORY) + with pytest.raises(KeyError, match="No policy for capability 'supports_multi_message_pieces'"): + policy.get_behavior(capability=CapabilityName.MULTI_MESSAGE_PIECES) - with pytest.raises(AttributeError, match="supports_editable_history"): - _ = policy.supports_editable_history + with pytest.raises(AttributeError, match="supports_multi_message_pieces"): + _ = policy.supports_multi_message_pieces def test_capability_handling_policy_rejects_unknown_attribute(self): policy = CapabilityHandlingPolicy() @@ -88,6 +88,7 @@ def test_normalizable_capabilities(self): frozenset( { CapabilityName.MULTI_TURN, + CapabilityName.EDITABLE_HISTORY, CapabilityName.SYSTEM_PROMPT, CapabilityName.JSON_SCHEMA, } diff --git a/tests/unit/prompt_target/target/test_target_configuration.py b/tests/unit/prompt_target/target/test_target_configuration.py index 36f0c62fc5..9ed1d6cafe 100644 --- a/tests/unit/prompt_target/target/test_target_configuration.py +++ b/tests/unit/prompt_target/target/test_target_configuration.py @@ -1,12 +1,15 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +from unittest.mock import AsyncMock, MagicMock + import pytest from pyrit.message_normalizer import ( GenericSystemSquashNormalizer, HistorySquashNormalizer, JsonSchemaNormalizer, + MessageListNormalizer, ) from pyrit.models import Message, MessagePiece from pyrit.models.literals import ChatMessageRole @@ -47,14 +50,22 @@ def _make(role: ChatMessageRole, content: str) -> Message: def test_init_with_defaults_uses_raise_policy(): - caps = TargetCapabilities(supports_multi_turn=True, supports_system_prompt=True) + caps = TargetCapabilities( + supports_multi_turn=True, + supports_editable_history=True, + supports_system_prompt=True, + ) config = TargetConfiguration(capabilities=caps) # Default policy is RAISE for all adaptable capabilities assert config.policy.get_behavior(capability=CapabilityName.MULTI_TURN) == UnsupportedCapabilityBehavior.RAISE def test_init_with_explicit_policy(adapt_all_policy): - caps = TargetCapabilities(supports_multi_turn=True, supports_system_prompt=True) + caps = TargetCapabilities( + supports_multi_turn=True, + supports_editable_history=True, + supports_system_prompt=True, + ) config = TargetConfiguration(capabilities=caps, policy=adapt_all_policy) assert config.policy is adapt_all_policy @@ -92,7 +103,11 @@ def test_init_missing_capability_raise_policy_skips_normalizer(): def test_init_missing_json_schema_default_policy_adds_normalizer(): # Default policy adapts JSON_SCHEMA; a target lacking native support gets the JSON-schema normalizer. - caps = TargetCapabilities(supports_multi_turn=True, supports_system_prompt=True) + caps = TargetCapabilities( + supports_multi_turn=True, + supports_editable_history=True, + supports_system_prompt=True, + ) config = TargetConfiguration(capabilities=caps) assert len(config.pipeline.normalizers) == 1 assert isinstance(config.pipeline.normalizers[0], JsonSchemaNormalizer) @@ -101,6 +116,7 @@ def test_init_missing_json_schema_default_policy_adds_normalizer(): def test_init_supports_json_schema_no_normalizer(): caps = TargetCapabilities( supports_multi_turn=True, + supports_editable_history=True, supports_system_prompt=True, supports_json_schema=True, ) @@ -157,6 +173,47 @@ def test_init_sparse_policy_missing_json_schema_no_normalizer(): assert config.pipeline.normalizers == () +def test_default_configuration_does_not_adapt_editable_history(): + caps = TargetCapabilities( + supports_multi_turn=True, + supports_editable_history=False, + supports_system_prompt=True, + supports_json_schema=True, + ) + config = TargetConfiguration(capabilities=caps) + + assert not config.pipeline.has_normalizer_for(capability=CapabilityName.EDITABLE_HISTORY) + with pytest.raises(ValueError, match="no handling policy"): + config.ensure_can_handle(capability=CapabilityName.EDITABLE_HISTORY) + + +async def test_per_send_override_preserves_construction_overrides_with_sparse_policy(make_message): + construction_normalizer = MagicMock(spec=MessageListNormalizer) + construction_normalizer.normalize_async = AsyncMock(side_effect=lambda messages: messages) + per_send_normalizer = MagicMock(spec=MessageListNormalizer) + per_send_normalizer.normalize_async = AsyncMock(side_effect=lambda messages: messages) + caps = TargetCapabilities( + supports_multi_turn=True, + supports_editable_history=False, + supports_system_prompt=False, + supports_json_schema=True, + ) + sparse_policy = CapabilityHandlingPolicy(behaviors={}) + config = TargetConfiguration( + capabilities=caps, + policy=sparse_policy, + normalizer_overrides={CapabilityName.SYSTEM_PROMPT: construction_normalizer}, + ) + + await config.normalize_async( + messages=[make_message("user", "hello")], + normalizer_overrides={CapabilityName.EDITABLE_HISTORY: per_send_normalizer}, + ) + + construction_normalizer.normalize_async.assert_awaited_once() + per_send_normalizer.normalize_async.assert_awaited_once() + + # --------------------------------------------------------------------------- # Properties # --------------------------------------------------------------------------- @@ -232,11 +289,43 @@ def test_ensure_can_handle_raises_when_capability_missing_from_policy(): config.ensure_can_handle(capability=CapabilityName.JSON_SCHEMA) +@pytest.mark.parametrize( + "behaviors", + [ + {}, + {CapabilityName.EDITABLE_HISTORY: UnsupportedCapabilityBehavior.RAISE}, + ], +) +def test_ensure_can_handle_explicit_override_precedes_sparse_or_raise_policy( + behaviors: dict[CapabilityName, UnsupportedCapabilityBehavior], +): + normalizer = MagicMock(spec=MessageListNormalizer) + config = TargetConfiguration( + capabilities=TargetCapabilities(supports_editable_history=False), + policy=CapabilityHandlingPolicy(behaviors=behaviors), + normalizer_overrides={CapabilityName.EDITABLE_HISTORY: normalizer}, + ) + + config.ensure_can_handle(capability=CapabilityName.EDITABLE_HISTORY) + + +def test_ensure_can_handle_adapt_requires_available_normalizer(): + config = TargetConfiguration( + capabilities=TargetCapabilities(supports_editable_history=False), + policy=CapabilityHandlingPolicy( + behaviors={CapabilityName.EDITABLE_HISTORY: UnsupportedCapabilityBehavior.ADAPT} + ), + ) + + with pytest.raises(ValueError, match="no default or configured normalizer"): + config.ensure_can_handle(capability=CapabilityName.EDITABLE_HISTORY) + + def test_ensure_can_handle_raises_valueerror_for_non_normalizable_capability(): - caps = TargetCapabilities(supports_multi_turn=True, supports_system_prompt=True, supports_editable_history=False) + caps = TargetCapabilities(supports_multi_turn=True, supports_system_prompt=True) config = TargetConfiguration(capabilities=caps) with pytest.raises(ValueError, match="no handling policy"): - config.ensure_can_handle(capability=CapabilityName.EDITABLE_HISTORY) + config.ensure_can_handle(capability=CapabilityName.MULTI_MESSAGE_PIECES) # --------------------------------------------------------------------------- diff --git a/tests/unit/prompt_target/target/test_target_history.py b/tests/unit/prompt_target/target/test_target_history.py new file mode 100644 index 0000000000..dda5313525 --- /dev/null +++ b/tests/unit/prompt_target/target/test_target_history.py @@ -0,0 +1,71 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import pytest + +from pyrit.models import ChatMessageRole, Message, PromptResponseError +from pyrit.prompt_target.common.target_history import filter_non_replayable_messages + + +def _message( + *, + role: ChatMessageRole, + value: str, + sequence: int, + conversation_id: str = "conversation", +) -> Message: + message = Message.from_prompt(prompt=value, role=role) + piece = message.get_piece() + piece.sequence = sequence + piece.conversation_id = conversation_id + return message + + +@pytest.mark.parametrize("response_error", ["processing", "unknown"]) +def test_filter_removes_adjacent_failed_exchange(response_error: PromptResponseError) -> None: + successful = _message(role="user", value="successful", sequence=0) + failed_request = _message(role="user", value="failed request", sequence=1) + error_response = _message(role="assistant", value="private stack trace", sequence=2) + error_response.get_piece().response_error = response_error + + filtered = filter_non_replayable_messages(messages=[successful, failed_request, error_response]) + + assert filtered == [successful] + + +@pytest.mark.parametrize( + ("preceding_role", "preceding_conversation_id", "preceding_sequence"), + [ + ("assistant", "conversation", 4), + ("user", "other-conversation", 4), + ("user", "conversation", 3), + ], +) +def test_filter_does_not_remove_unrelated_preceding_message( + preceding_role: ChatMessageRole, + preceding_conversation_id: str, + preceding_sequence: int, +) -> None: + preceding = _message( + role=preceding_role, + value="unrelated", + sequence=preceding_sequence, + conversation_id=preceding_conversation_id, + ) + error_response = _message(role="assistant", value="private stack trace", sequence=5) + error_response.get_piece().response_error = "processing" + + filtered = filter_non_replayable_messages(messages=[preceding, error_response]) + + assert filtered == [preceding] + + +@pytest.mark.parametrize("response_error", ["blocked", "empty"]) +def test_filter_retains_provider_round_trip(response_error: PromptResponseError) -> None: + request = _message(role="user", value="request", sequence=0) + response = _message(role="assistant", value="provider response", sequence=1) + response.get_piece().response_error = response_error + + filtered = filter_non_replayable_messages(messages=[request, response]) + + assert filtered == [request, response] diff --git a/tests/unit/prompt_target/target/test_target_requirements.py b/tests/unit/prompt_target/target/test_target_requirements.py index 9c34ce326e..809b1a6f93 100644 --- a/tests/unit/prompt_target/target/test_target_requirements.py +++ b/tests/unit/prompt_target/target/test_target_requirements.py @@ -40,6 +40,7 @@ def test_chat_target_requirements_shape(): CapabilityName.EDITABLE_HISTORY, CapabilityName.MULTI_TURN, } + assert CHAT_TARGET_REQUIREMENTS.native_required == set() def test_requirements_are_frozen(): diff --git a/tests/unit/prompt_target/target/test_websocket_target.py b/tests/unit/prompt_target/target/test_websocket_target.py index 281dc5d5cd..5eac8e842a 100644 --- a/tests/unit/prompt_target/target/test_websocket_target.py +++ b/tests/unit/prompt_target/target/test_websocket_target.py @@ -13,11 +13,26 @@ from websockets.protocol import State from pyrit.exceptions import EmptyResponseException +from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, +) from pyrit.memory import SQLiteMemory from pyrit.models import Message, MessagePiece from pyrit.prompt_target import WebsocketTarget +class _OverriddenWebsocketTarget(WebsocketTarget): + async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: + request = normalized_conversation[-1].get_piece() + return [ + MessagePiece( + role="assistant", + original_value="response", + conversation_id=request.conversation_id, + ).to_message() + ] + + @pytest.fixture def response_parser() -> Callable[[str | bytes], str | None]: def parse_response(message: str | bytes) -> str | None: @@ -78,6 +93,10 @@ def test_init_invalid_endpoint_raises( ) +def test_overridden_managed_target_uses_compatibility_boundary() -> None: + assert not _OverriddenWebsocketTarget._MANAGES_PROVIDER_ATTEMPT_BOUNDARY + + def test_init_empty_protocol_identifier_raises( response_parser: Callable[[str | bytes], str | None], message_builder: Callable[[str], str | bytes], @@ -249,6 +268,45 @@ async def test_send_prompt_async_failure_discards_connection(websocket_target: W assert "conversation" not in websocket_target._existing_conversation +async def test_cancellation_before_websocket_send_does_not_consume_seed( + websocket_target: WebsocketTarget, + sqlite_instance: SQLiteMemory, +) -> None: + seed = create_message(value="Seed") + sqlite_instance.add_message_to_memory(request=seed) + send_context = PrependedHistorySendContext( + conversation_id="conversation", + seed_message_ids=(seed.get_piece().id,), + replay_seed_each_send=False, + ) + connection_started = asyncio.Event() + connection_release = asyncio.Event() + + async def wait_for_connection( + *, + conversation_id: str, + conversation_history: list[Message], + ) -> ClientConnection: + connection_started.set() + await connection_release.wait() + return AsyncMock(spec=ClientConnection) + + with patch.object(websocket_target, "_get_or_create_connection_async", side_effect=wait_for_connection): + send_task = asyncio.create_task( + websocket_target.send_prompt_async( + message=create_message(value="Current"), + send_context=send_context, + ) + ) + await connection_started.wait() + send_task.cancel() + with pytest.raises(asyncio.CancelledError): + await send_task + + assert send_context.provider_attempt_count == 0 + assert not send_context.is_seed_consumed + + async def test_get_or_create_connection_async_restores_history_on_reconnect( response_parser: Callable[[str | bytes], str | None], message_builder: Callable[[str], str | bytes], @@ -289,6 +347,75 @@ async def test_get_or_create_connection_async_restores_history_on_reconnect( assert target._existing_conversation == {"conversation": replacement_connection} +async def test_consumed_context_retains_history_for_reconnect( + response_parser: Callable[[str | bytes], str | None], + message_builder: Callable[[str], str | bytes], + sqlite_instance: SQLiteMemory, +) -> None: + stale_connection = AsyncMock(spec=ClientConnection) + stale_connection.state = State.CLOSED + replacement_connection = AsyncMock(spec=ClientConnection) + replacement_connection.state = State.OPEN + restore_callback = AsyncMock() + target = WebsocketTarget( + endpoint="wss://example.com", + protocol_identifier="test-protocol", + initialization_strings=[], + response_parser=response_parser, + message_builder=message_builder, + conversation_restore_callback=restore_callback, + discard_initial_messages=0, + existing_convo={"conversation": stale_connection}, + ) + seed = create_message(value="Seed") + first_request = create_message(value="First request") + first_response = MessagePiece( + role="assistant", + original_value="First response", + converted_value="First response", + conversation_id="conversation", + ).to_message() + for sequence, message in enumerate([seed, first_request, first_response]): + for piece in message.message_pieces: + piece.sequence = sequence + sqlite_instance.add_message_to_memory(request=message) + + send_context = PrependedHistorySendContext( + conversation_id="conversation", + seed_message_ids=(seed.get_piece().id,), + replay_seed_each_send=False, + ) + send_context.begin_send() + send_context.mark_provider_attempted() + send_context.finish_send() + + with ( + patch.object( + target, + "_connect_async", + new_callable=AsyncMock, + return_value=replacement_connection, + ), + patch.object( + target, + "_send_text_async", + new_callable=AsyncMock, + return_value="Second response", + ), + ): + await target.send_prompt_async( + message=create_message(value="Second request"), + send_context=send_context, + ) + + restored_history = restore_callback.await_args.args[1] + assert [message.get_value() for message in restored_history] == [ + "Seed", + "First request", + "First response", + ] + + async def test_get_or_create_connection_async_fails_when_history_cannot_be_restored( websocket_target: WebsocketTarget, ) -> None: diff --git a/tests/unit/scenario/airt/test_jailbreak.py b/tests/unit/scenario/airt/test_jailbreak.py index 5cce529015..b055e087aa 100644 --- a/tests/unit/scenario/airt/test_jailbreak.py +++ b/tests/unit/scenario/airt/test_jailbreak.py @@ -16,6 +16,7 @@ from pyrit.prompt_target import PromptTarget from pyrit.registry import TargetRegistry from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry +from pyrit.registry.components.scenario_registry import ScenarioRegistry from pyrit.scenario.core import BaselineAttackPolicy from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory from pyrit.scenario.scenarios.airt.jailbreak import ( @@ -41,13 +42,7 @@ def _technique_class(): @pytest.fixture(autouse=True) def reset_technique_registry(): - """Populate the attack-technique registry so the dynamic technique class can be built. - - Mirrors the RapidResponse test setup: reset the registries, register a mock adversarial - target (so factory construction does not fall back to a real target), and register the core - technique factories. The build cache is cleared around each test so the class reflects the - freshly-registered factories. - """ + """Populate the attack-technique registry used by the shared matrix factory resolver.""" AttackTechniqueRegistry.reset_registry_singleton() TargetRegistry.reset_registry_singleton() _build_jailbreak_technique.cache_clear() @@ -309,9 +304,8 @@ async def test_jailbreak_delivered_as_request_converter( """The crux: the jailbreak template reaches the target as a ``TextJailbreakConverter`` on the technique's outgoing requests (not as prepended framing on the seed group). - Delivery via ``factory.create(extra_request_converters=...)`` is what keeps the scenario - target-agnostic and composable with every technique. Also assert the seed groups carry no - prepended jailbreak framing. + Delivery via ``factory.create(extra_request_converters=...)`` keeps the prompt-sending path + target-agnostic. Also assert the seed groups carry no prepended jailbreak framing. """ captured: list[Any] = [] original_create = AttackTechniqueFactory.create @@ -328,9 +322,7 @@ def _spy_create(self, **kwargs): assert captured, "Expected factory.create to be called" converters = [c for extra in captured if extra for cc in extra for c in cc.converters] - assert any(isinstance(c, TextJailbreakConverter) for c in converters), ( - "Expected a TextJailbreakConverter to be threaded to factory.create" - ) + assert sum(isinstance(c, TextJailbreakConverter) for c in converters) == 1 # The objective seed groups themselves carry no jailbreak framing (converter delivery only). for attack in scenario._atomic_attacks: @@ -376,49 +368,35 @@ def _spy_create(self, **kwargs): "Jailbreak converter must be applied before caller-supplied converters" ) - async def test_simulated_conversation_techniques_produce_attacks_with_jailbreak( + async def test_stale_incompatible_technique_is_rejected( self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups ): - """Regression: simulated-conversation techniques (``role_play_*``, ``crescendo_*``) must still - produce atomic attacks when crossed with a jailbreak template, and each must receive the - jailbreak converter. - - Converter delivery leaves the objective seed group unframed, so it stays compatible with the - simulated-conversation seed technique. (Delivering the jailbreak as a system-role framing seed - instead collided with that technique's seed range and silently produced zero attacks.) - """ - technique_class = _build_jailbreak_technique() - techniques = [ - technique_class("role_play_movie_script"), - technique_class("crescendo_simulated"), - ] - captured: list[Any] = [] - original_create = AttackTechniqueFactory.create - - def _spy_create(self, **kwargs): - captured.append(kwargs.get("extra_request_converters")) - return original_create(self, **kwargs) - + registry_factories = list(AttackTechniqueRegistry.get_registry_singleton().get_factories_or_raise().values()) + legacy_class = AttackTechniqueRegistry.build_technique_class_from_factories( + class_name="LegacyJailbreakTechnique", + factories=registry_factories, + ) with _patch_seed_groups(mock_memory_seed_groups): - with patch.object(AttackTechniqueFactory, "create", _spy_create): - scenario = Jailbreak(objective_scorer=mock_objective_scorer) - scenario.set_params_from_args( - args=_default_args( - mock_objective_target, scenario_techniques=techniques, jailbreak_names=["aim.yaml"] - ) + scenario = Jailbreak(objective_scorer=mock_objective_scorer) + scenario.set_params_from_args( + args=_default_args( + mock_objective_target, + scenario_techniques=[legacy_class("tap")], + jailbreak_names=["aim.yaml"], ) + ) + with pytest.raises(ValueError, match="stale or incompatible"): await scenario.initialize_async() - names = {a.atomic_attack_name for a in scenario._atomic_attacks} - assert "role_play_movie_script_aim_harmbench" in names - assert "crescendo_simulated_aim_harmbench" in names - # Every build of a simulated-conversation technique must still carry the jailbreak - # converter. Assert both techniques captured a non-empty converter stack (so the check - # can't pass vacuously on a dropped/None stack) and each contains the jailbreak converter. - populated = [extra for extra in captured if extra] - assert len(populated) == 2, "Expected both simulated-conversation techniques to receive converters" - assert all( - any(isinstance(c, TextJailbreakConverter) for cc in extra for c in cc.converters) for extra in populated - ), "Each simulated-conversation technique must receive the jailbreak converter" + + async def test_missing_runtime_factory_is_rejected( + self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups + ): + with _patch_seed_groups(mock_memory_seed_groups): + with patch("pyrit.scenario.scenarios.airt.jailbreak.resolve_technique_factories", return_value={}): + scenario = Jailbreak(objective_scorer=mock_objective_scorer) + scenario.set_params_from_args(args=_default_args(mock_objective_target, jailbreak_names=["aim.yaml"])) + with pytest.raises(ValueError, match="no longer available.*prompt_sending"): + await scenario.initialize_async() async def test_all_templates_produce_attacks( self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups @@ -708,15 +686,17 @@ def test_default_techniques_are_the_two_deliveries(self): assert default_values == set(_DEFAULT_TECHNIQUES) assert default_values == {_PROMPT_SENDING, _JAILBREAK_SYSTEM_PROMPT} - def test_registry_techniques_are_available(self): + def test_only_scenario_delivery_techniques_are_available(self): technique_class = _technique_class() available = {t.value for t in technique_class.get_all_techniques()} - assert {_PROMPT_SENDING, _JAILBREAK_SYSTEM_PROMPT}.issubset(available) - # The "normal ones available like from rapid response" are exposed as opt-in techniques. - assert {"role_play_movie_script", "many_shot", "tap"}.issubset(available) + assert available == {_PROMPT_SENDING, _JAILBREAK_SYSTEM_PROMPT} + + def test_registry_metadata_lists_only_scenario_deliveries(self): + metadata = ScenarioRegistry()._build_metadata("airt.jailbreak", Jailbreak) + assert set(metadata.all_techniques) == {_PROMPT_SENDING, _JAILBREAK_SYSTEM_PROMPT} - def test_scenario_version_is_three(self): - assert Jailbreak.VERSION == 3 + def test_scenario_version_is_four(self): + assert Jailbreak.VERSION == 4 def test_default_dataset_is_harmbench(self): assert Jailbreak.required_datasets() == ["harmbench"]