From f96964dfd318e83533df1ca3e2731ef851436e8b Mon Sep 17 00:00:00 2001 From: Utkarsh Bahuguna Date: Sat, 10 Oct 2026 23:46:09 +0530 Subject: [PATCH 1/5] don't answer a Copilot conversation from a session that never saw its earlier turns --- pyrit/prompt_target/github_copilot_target.py | 14 ++++++ .../target/test_github_copilot_target.py | 50 +++++++++++++++++++ 2 files changed, 64 insertions(+) diff --git a/pyrit/prompt_target/github_copilot_target.py b/pyrit/prompt_target/github_copilot_target.py index c066c7e999..8c622d0709 100644 --- a/pyrit/prompt_target/github_copilot_target.py +++ b/pyrit/prompt_target/github_copilot_target.py @@ -37,6 +37,10 @@ class GitHubCopilotTarget(PromptTarget): """ Send text requests through the GitHub Copilot SDK. + Conversation history lives in the Copilot session, so a conversation can only be continued by the target + instance that started it. Continuing one that has earlier turns but no live session raises instead of + sending the last turn to a fresh session. + Capture INFO logs for session mapping and SDK/runtime version diagnostics. """ @@ -196,6 +200,7 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me session = await self._get_or_create_session_async( conversation_id=conversation_id, initial_system_prompt=initial_system_prompt, + has_earlier_turns=any(message.api_role != "system" for message in normalized_conversation[:-1]), ) try: reply_text = await self._send_text_async( @@ -345,6 +350,7 @@ async def _get_or_create_session_async( *, conversation_id: str, initial_system_prompt: str | None, + has_earlier_turns: bool = False, ) -> "CopilotSession": async with self._lifecycle_condition: conversation = self._conversations[conversation_id] @@ -361,6 +367,14 @@ async def _get_or_create_session_async( existing_session = conversation.session if existing_session is not None: return existing_session + if has_earlier_turns: + # The history lives in the Copilot session, not in what we send. A new session would answer the + # last turn as if it were the first, e.g. for a copied conversation or after a restart. + raise ValueError( + f"GitHubCopilotTarget has no Copilot session for conversation '{conversation_id}', which already " + "has earlier turns. Copilot keeps the history inside its session, so those turns can't be " + "restored; continue in a new conversation." + ) client = await self._get_or_start_client_async() session_id = str(uuid4()) diff --git a/tests/unit/prompt_target/target/test_github_copilot_target.py b/tests/unit/prompt_target/target/test_github_copilot_target.py index 7bcbc4f10b..733b63d24a 100644 --- a/tests/unit/prompt_target/target/test_github_copilot_target.py +++ b/tests/unit/prompt_target/target/test_github_copilot_target.py @@ -1239,6 +1239,56 @@ async def send_and_wait_async(*_args: Any, **_kwargs: Any) -> Any: client.stop.assert_awaited_once() +@pytest.mark.usefixtures("patch_central_database") +async def test_continuing_a_conversation_without_its_session_raises_async( + *, + sdk: Any, + client: NonCallableMagicMock, + sqlite_instance: MemoryInterface, +) -> None: + """A copied or restarted conversation must not be answered by a fresh session that never saw its turns.""" + conversation_id = "conversation-started-elsewhere" + session = client.create_session.return_value + session.send_and_wait.side_effect = [_assistant_reply("FIRST")] + first_target = GitHubCopilotTarget(model_name="gpt-5-mini") + await _send_normalized_async( + target=first_target, original_value="My secret word is avocado.", conversation_id=conversation_id + ) + + # A new target instance (e.g. after a backend restart) has the memory but not the Copilot session. + second_target = GitHubCopilotTarget(model_name="gpt-5-mini") + with pytest.raises(Exception, match="Error sending prompt") as raised: + await _send_normalized_async( + target=second_target, original_value="What is my secret word?", conversation_id=conversation_id + ) + assert isinstance(raised.value.__cause__, ValueError) + assert "no Copilot session for conversation" in str(raised.value.__cause__) + + client.create_session.assert_awaited_once() + assert session.send_and_wait.await_count == 1 + await first_target.cleanup_target_async() + await second_target.cleanup_target_async() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_system_prompt_alone_still_starts_a_session_async( + *, + sdk: Any, + client: NonCallableMagicMock, + sqlite_instance: MemoryInterface, +) -> None: + conversation_id = "system-only-conversation" + client.create_session.return_value.send_and_wait.side_effect = [_assistant_reply("OK")] + target = GitHubCopilotTarget(model_name="gpt-5-mini") + await target.set_system_prompt_async(system_prompt="Be brief.", conversation_id=conversation_id) + + response = await _send_normalized_async(target=target, original_value="hi", conversation_id=conversation_id) + + assert response.get_piece().converted_value == "OK" + client.create_session.assert_awaited_once() + await target.cleanup_target_async() + + @pytest.mark.usefixtures("patch_central_database", "sdk") def test_target_advertises_native_text_only_capabilities() -> None: capabilities = GitHubCopilotTarget(model_name="gpt-4o").capabilities From c73fcf85b2cc0a573ed8c6c743fe1d6e1070129d Mon Sep 17 00:00:00 2001 From: Utkarsh Bahuguna Date: Sat, 10 Oct 2026 23:46:09 +0530 Subject: [PATCH 2/5] send our own A2A contextId so turn 2 works when the agent doesn't return one --- pyrit/prompt_target/a2a_target.py | 6 ++- .../prompt_target/target/test_a2a_target.py | 51 +++++++++++++++++++ 2 files changed, 55 insertions(+), 2 deletions(-) diff --git a/pyrit/prompt_target/a2a_target.py b/pyrit/prompt_target/a2a_target.py index b607f701a3..ecc701a3f6 100644 --- a/pyrit/prompt_target/a2a_target.py +++ b/pyrit/prompt_target/a2a_target.py @@ -55,7 +55,7 @@ async def async_auth_flow( # pyrit-async-suffix-exempt @dataclass class _A2AConversationState: - """Server-issued identifiers that continue one PyRIT conversation on the agent.""" + """Identifiers that continue one PyRIT conversation on the agent (our context ID until the agent sends its own).""" context_id: str | None = None open_task_id: str | None = None @@ -426,7 +426,9 @@ async def _send_message_async(self, *, normalized_conversation: list[Message], c "The target requires server-side context continuity for multi-turn conversations, " "and earlier turns cannot be restored." ) - state = self._conversations.setdefault(conversation_id, _A2AConversationState()) + # Name the context ourselves so turn 2 can continue it even when the agent's reply carries no + # contextId (a plain Message may omit it) or the first turn came back as an error. + state = self._conversations.setdefault(conversation_id, _A2AConversationState(context_id=str(uuid.uuid4()))) prompt_text = message_piece.converted_value req = a2a_pb2.SendMessageRequest( diff --git a/tests/unit/prompt_target/target/test_a2a_target.py b/tests/unit/prompt_target/target/test_a2a_target.py index 2369a0a7b2..ff0402e690 100644 --- a/tests/unit/prompt_target/target/test_a2a_target.py +++ b/tests/unit/prompt_target/target/test_a2a_target.py @@ -12,6 +12,7 @@ from pyrit.exceptions import EmptyResponseException from pyrit.models import Message, MessagePiece +from pyrit.prompt_normalizer import PromptNormalizer from pyrit.prompt_target import A2ATarget from pyrit.prompt_target.a2a_target import _A2AConversationState from pyrit.prompt_target.common.target_capabilities import TargetCapabilities @@ -190,6 +191,56 @@ async def test_a2a_send_prompt_message_result(mock_send): assert payload["params"]["message"]["parts"][0]["text"] == "Hello A2A Agent" +def _message_payload_without_context(text: str, *, message_id: str) -> dict[str, Any]: + payload = _message_payload(text, message_id=message_id) + del payload["result"]["contextId"] + return payload + + +@pytest.mark.usefixtures("patch_central_database") +@patch("httpx.AsyncClient.send") +async def test_a2a_continues_its_own_context_when_the_agent_sends_none(mock_send): + """A plain Message reply may omit contextId; turn 2 must still continue the same context.""" + mock_send.side_effect = [ + _rpc_response(_message_payload_without_context("First reply", message_id="m-1")), + _rpc_response(_message_payload_without_context("Second reply", message_id="m-2")), + ] + target = A2ATarget(endpoint=ENDPOINT) + normalizer = PromptNormalizer() + + first = await normalizer.send_prompt_async( + message=_user_message("turn one"), conversation_id="conv-1", target=target + ) + second = await normalizer.send_prompt_async( + message=_user_message("turn two"), conversation_id="conv-1", target=target + ) + + assert first.get_piece().converted_value == "First reply" + assert second.get_piece().converted_value == "Second reply" + sent_contexts = [ + json.loads(call.args[0].content)["params"]["message"]["contextId"] for call in mock_send.call_args_list + ] + assert sent_contexts[0] + assert sent_contexts == [sent_contexts[0], sent_contexts[0]] + + +@pytest.mark.usefixtures("patch_central_database") +@patch("httpx.AsyncClient.send") +async def test_a2a_adopts_the_context_the_agent_returns(mock_send): + mock_send.side_effect = [ + _rpc_response(_message_payload("First reply", context_id="ctx-from-agent")), + _rpc_response(_message_payload("Second reply", context_id="ctx-from-agent")), + ] + target = A2ATarget(endpoint=ENDPOINT) + normalizer = PromptNormalizer() + + await normalizer.send_prompt_async(message=_user_message("turn one"), conversation_id="conv-1", target=target) + await normalizer.send_prompt_async(message=_user_message("turn two"), conversation_id="conv-1", target=target) + + second_request = json.loads(mock_send.call_args_list[1].args[0].content) + assert second_request["params"]["message"]["contextId"] == "ctx-from-agent" + + @pytest.mark.usefixtures("patch_central_database") @patch("httpx.AsyncClient.send") async def test_a2a_send_prompt_task_completed(mock_send): From 4044c5390c0cdb6b234ae5a84eb591f7d3da1c29 Mon Sep 17 00:00:00 2001 From: Utkarsh Bahuguna Date: Sat, 10 Oct 2026 23:46:09 +0530 Subject: [PATCH 3/5] reject repeat-send into copies of a conversation for targets that keep their own history --- .../backend/services/message_send_service.py | 25 ++++++++++- .../unit/backend/test_message_send_service.py | 45 +++++++++++++++++++ 2 files changed, 69 insertions(+), 1 deletion(-) diff --git a/pyrit/backend/services/message_send_service.py b/pyrit/backend/services/message_send_service.py index 57e2064233..0b42c9e036 100644 --- a/pyrit/backend/services/message_send_service.py +++ b/pyrit/backend/services/message_send_service.py @@ -50,7 +50,7 @@ MessagePiece, ) from pyrit.prompt_normalizer import ConverterConfiguration, PromptNormalizer -from pyrit.prompt_target import PromptTarget +from pyrit.prompt_target import CapabilityName, PromptTarget from pyrit.prompt_target.common.target_send_context import TargetSendContext logger = logging.getLogger(__name__) @@ -186,6 +186,8 @@ async def submit_async(self, *, attack_result_id: str, request: MessageSendReque conversations[conversation_id] = owner owned_request = request.model_copy(deep=True) validated = await self._validate_message_async(attack_result_id=attack_result_id, request=owned_request) + if owned_request.count > 1: + await self._validate_repeat_history_async(request=owned_request, target=validated.target) if self._closing: raise ManualSendQueueFullError("Manual message operations are shutting down") operation = _Send( @@ -311,6 +313,27 @@ async def _validate_message_async(self, *, attack_result_id: str, request: AddMe applied_identifiers=resolve_applied_converter_identifiers(request.pieces), ) + async def _validate_repeat_history_async(self, *, request: MessageSendRequest, target: PromptTarget | None) -> None: + """ + Reject repeating a send into copies of a conversation the target can't be handed. + + Repeat-send copies the existing turns into new conversations. A target without editable history + keeps its own copy of the conversation (for example a Copilot session or an A2A context), so it + would answer each copy without the earlier turns while memory shows them. + + Raises: + ValueError: If the target lacks editable history and the conversation has turns besides a + system prompt. + """ + if target is None or target.configuration.includes(capability=CapabilityName.EDITABLE_HISTORY): + return + history = await self._memory.get_conversation_messages_async(conversation_id=request.target_conversation_id) + if any(message.api_role != "system" for message in history): + raise ValueError( + f"{type(target).__name__} keeps its own conversation history, so a send can't be repeated into " + "copies of a conversation that already has turns. Repeat the send from a new conversation instead." + ) + async def _execute_validated_message_async( self, *, diff --git a/tests/unit/backend/test_message_send_service.py b/tests/unit/backend/test_message_send_service.py index 1ac6b74247..a6cd662715 100644 --- a/tests/unit/backend/test_message_send_service.py +++ b/tests/unit/backend/test_message_send_service.py @@ -59,6 +59,8 @@ PromptDataType, ) from pyrit.prompt_normalizer import ConverterConfiguration, PromptNormalizer +from pyrit.prompt_target.common.target_capabilities import TargetCapabilities +from pyrit.prompt_target.common.target_configuration import TargetConfiguration from unit.async_utils import wait_for_completion_async from unit.backend.mocks import ( _make_matching_target_mock, @@ -2984,6 +2986,49 @@ async def hold_async(*, normalized_conversation: list[Message]) -> list[Message] await service.get_status_async(attack_result_id=ar.attack_result_id, send_id=status.send_id) ).state == MessageSendState.COMPLETED + async def test_repeat_send_rejects_copying_history_into_a_stateful_target_async( + self, + *, + real_send_context: tuple[MessageSendService, AttackResult, MockPromptTarget, Base64Converter], + sqlite_instance: SQLiteMemory, + ) -> None: + """Copies would be answered by a target that never saw the earlier turns, so refuse up front.""" + service, ar, target, _ = real_send_context + first = await service.submit_async( + attack_result_id=ar.attack_result_id, request=_submission(conversation_id=ar.conversation_id) + ) + assert (await _settle_send_async(service=service, status=first)).state == MessageSendState.COMPLETED + target._configuration = TargetConfiguration( + capabilities=TargetCapabilities(supports_multi_turn=True, supports_multi_message_pieces=True) + ) + + with pytest.raises(ValueError, match="keeps its own conversation history"): + await service.submit_async( + attack_result_id=ar.attack_result_id, + request=_submission(conversation_id=ar.conversation_id, submission_id="repeat", count=3), + ) + + pieces = await sqlite_instance.get_message_pieces_async(conversation_id=ar.conversation_id) + assert [piece.role for piece in pieces] == ["user", "assistant"] + assert not service._scheduler._conversations + + async def test_repeat_send_into_a_stateful_target_is_fine_on_a_new_conversation_async( + self, + *, + real_send_context: tuple[MessageSendService, AttackResult, MockPromptTarget, Base64Converter], + ) -> None: + service, ar, target, _ = real_send_context + target._configuration = TargetConfiguration( + capabilities=TargetCapabilities(supports_multi_turn=True, supports_multi_message_pieces=True) + ) + + status = await service.submit_async( + attack_result_id=ar.attack_result_id, + request=_submission(conversation_id=ar.conversation_id, count=3), + ) + + assert (await _settle_send_async(service=service, status=status)).state == MessageSendState.COMPLETED + @pytest.mark.parametrize("invalid", ["attack", "conversation", "target", "converter"]) async def test_validation_does_not_accept_or_write_async( self, From 00bc5da97432bf713e40f72dae35ec1be3a476e9 Mon Sep 17 00:00:00 2001 From: Utkarsh Bahuguna Date: Sun, 11 Oct 2026 10:43:09 +0530 Subject: [PATCH 4/5] tighten tests: async names, typed args, public paths --- tests/unit/prompt_target/target/test_a2a_target.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/prompt_target/target/test_a2a_target.py b/tests/unit/prompt_target/target/test_a2a_target.py index ff0402e690..7e5ac7c6d6 100644 --- a/tests/unit/prompt_target/target/test_a2a_target.py +++ b/tests/unit/prompt_target/target/test_a2a_target.py @@ -199,7 +199,7 @@ def _message_payload_without_context(text: str, *, message_id: str) -> dict[str, @pytest.mark.usefixtures("patch_central_database") @patch("httpx.AsyncClient.send") -async def test_a2a_continues_its_own_context_when_the_agent_sends_none(mock_send): +async def test_a2a_continues_its_own_context_when_the_agent_sends_none_async(mock_send: AsyncMock) -> None: """A plain Message reply may omit contextId; turn 2 must still continue the same context.""" mock_send.side_effect = [ _rpc_response(_message_payload_without_context("First reply", message_id="m-1")), @@ -226,7 +226,7 @@ async def test_a2a_continues_its_own_context_when_the_agent_sends_none(mock_send @pytest.mark.usefixtures("patch_central_database") @patch("httpx.AsyncClient.send") -async def test_a2a_adopts_the_context_the_agent_returns(mock_send): +async def test_a2a_adopts_the_context_the_agent_returns_async(mock_send: AsyncMock) -> None: mock_send.side_effect = [ _rpc_response(_message_payload("First reply", context_id="ctx-from-agent")), _rpc_response(_message_payload("Second reply", context_id="ctx-from-agent")), From bc77f3fbae55a8a03ede561be7ce9cdab6cd440e Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Tue, 13 Oct 2026 23:27:52 -0700 Subject: [PATCH 5/5] FIX Validate configured history replay before repeat sends Clarify editable-history semantics and check native support, configured adapters, or WebSocket restoration before copying conversation history. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 23cf5e2b-2223-44e3-ab65-252dcc5ce4ee --- .github/instructions/targets.instructions.md | 5 ++ doc/code/framework.md | 4 ++ .../backend/services/message_send_service.py | 26 +++---- pyrit/models/target/target_capabilities.py | 6 +- pyrit/prompt_target/common/prompt_target.py | 17 +++++ pyrit/prompt_target/websocket_target.py | 5 ++ .../unit/backend/test_message_send_service.py | 53 +++++++++++++-- .../target/test_prompt_target.py | 68 ++++++++++++++++++- .../target/test_websocket_target.py | 53 +++++++++++++++ 9 files changed, 215 insertions(+), 22 deletions(-) diff --git a/.github/instructions/targets.instructions.md b/.github/instructions/targets.instructions.md index f25d11c60b..a5e5d17e9b 100644 --- a/.github/instructions/targets.instructions.md +++ b/.github/instructions/targets.instructions.md @@ -82,6 +82,11 @@ def __init__(self, endpoint: str, api_key: str) -> None: ... # missing * ``__init__`` and forward it to ``super().__init__`` so callers can override capabilities per-instance (this is required for HTTP / Playwright targets whose capabilities depend on deployment configuration). +- ``supports_editable_history`` describes native support for caller-supplied + history changes, not whether the target stores state. Callers use + ``validate_history_replay()`` before branching stored history. Targets with + an explicit restoration path can override this check without claiming + editable history. Live session and context checks remain at send time. ## Identifiable Pattern diff --git a/doc/code/framework.md b/doc/code/framework.md index 7b788af485..9a22bc2489 100644 --- a/doc/code/framework.md +++ b/doc/code/framework.md @@ -244,6 +244,10 @@ 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. +- Editable history means native support for caller-supplied history changes. Its absence does not + imply server-side state. Callers that branch stored history use `validate_history_replay()` to check + native support, configured adaptation, or provider-specific restoration before creating copies. + Targets still check live session and context continuity when sending. - 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 diff --git a/pyrit/backend/services/message_send_service.py b/pyrit/backend/services/message_send_service.py index 0b42c9e036..5351370600 100644 --- a/pyrit/backend/services/message_send_service.py +++ b/pyrit/backend/services/message_send_service.py @@ -50,7 +50,7 @@ MessagePiece, ) from pyrit.prompt_normalizer import ConverterConfiguration, PromptNormalizer -from pyrit.prompt_target import CapabilityName, PromptTarget +from pyrit.prompt_target import PromptTarget from pyrit.prompt_target.common.target_send_context import TargetSendContext logger = logging.getLogger(__name__) @@ -315,24 +315,24 @@ async def _validate_message_async(self, *, attack_result_id: str, request: AddMe async def _validate_repeat_history_async(self, *, request: MessageSendRequest, target: PromptTarget | None) -> None: """ - Reject repeating a send into copies of a conversation the target can't be handed. - - Repeat-send copies the existing turns into new conversations. A target without editable history - keeps its own copy of the conversation (for example a Copilot session or an A2A context), so it - would answer each copy without the earlier turns while memory shows them. + Check that the target can initialize the repeated conversations from stored history. Raises: - ValueError: If the target lacks editable history and the conversation has turns besides a - system prompt. + ValueError: If non-system history cannot be replayed through native support, + a configured adapter, or a provider-specific restoration path. """ - if target is None or target.configuration.includes(capability=CapabilityName.EDITABLE_HISTORY): + if target is None: return history = await self._memory.get_conversation_messages_async(conversation_id=request.target_conversation_id) - if any(message.api_role != "system" for message in history): + if not any(message.api_role != "system" for message in history): + return + try: + target.validate_history_replay() + except ValueError as exc: raise ValueError( - f"{type(target).__name__} keeps its own conversation history, so a send can't be repeated into " - "copies of a conversation that already has turns. Repeat the send from a new conversation instead." - ) + f"{type(target).__name__} cannot initialize copies of this conversation from stored history. " + "Repeat the send from a new conversation or configure a supported history adapter or restoration path." + ) from exc async def _execute_validated_message_async( self, diff --git a/pyrit/models/target/target_capabilities.py b/pyrit/models/target/target_capabilities.py index 930dc8fa0b..554f916bb1 100644 --- a/pyrit/models/target/target_capabilities.py +++ b/pyrit/models/target/target_capabilities.py @@ -91,9 +91,9 @@ class attribute. Users can override individual capabilities per instance #: format), which ensures the output is valid JSON. supports_json_output: bool = False - #: Whether the target allows the attack history to be modified. Implies that the - #: target supports multi-turn interactions and that the attack history is not - #: immutable once set. + #: Whether the target natively honors caller-supplied changes to conversation + #: history. Implies multi-turn support. False does not imply server-side state + #: or rule out a configured adapter or restoration path for a new conversation. supports_editable_history: bool = False #: Whether the target natively supports system prompts. diff --git a/pyrit/prompt_target/common/prompt_target.py b/pyrit/prompt_target/common/prompt_target.py index cf04d20e52..933e1e71e5 100644 --- a/pyrit/prompt_target/common/prompt_target.py +++ b/pyrit/prompt_target/common/prompt_target.py @@ -216,6 +216,23 @@ def validate_tool_history(self, messages: Sequence[Message]) -> None: """ validate_tool_conversation(messages) + def validate_history_replay(self) -> None: + """ + Check configuration support for initializing a new conversation from stored history. + + Native editable history or an explicit history adapter is sufficient. + Targets with a provider-specific restoration path can override this check. + This does not validate message content or guarantee that a live session exists. + + Raises: + ValueError: If neither native support nor a configured history adapter is available. + """ + if not self.configuration.includes( + capability=CapabilityName.MULTI_TURN + ) and self.configuration.pipeline.has_normalizer_for(capability=CapabilityName.MULTI_TURN): + return + self.configuration.ensure_can_handle(capability=CapabilityName.EDITABLE_HISTORY) + @final async def send_prompt_async( self, diff --git a/pyrit/prompt_target/websocket_target.py b/pyrit/prompt_target/websocket_target.py index 613aef456b..715711a84a 100644 --- a/pyrit/prompt_target/websocket_target.py +++ b/pyrit/prompt_target/websocket_target.py @@ -119,6 +119,11 @@ def _build_identifier(self) -> ComponentIdentifier: """ return self._create_identifier(params={"protocol_identifier": self._protocol_identifier}) + def validate_history_replay(self) -> None: + """Check for a restoration callback or configured history support.""" + if self._conversation_restore_callback is None: + super().validate_history_replay() + async def _connect_async(self) -> ClientConnection: """ Open a connection to the configured WebSocket endpoint. diff --git a/tests/unit/backend/test_message_send_service.py b/tests/unit/backend/test_message_send_service.py index a57dd55b6d..b874a6f15b 100644 --- a/tests/unit/backend/test_message_send_service.py +++ b/tests/unit/backend/test_message_send_service.py @@ -46,6 +46,7 @@ from pyrit.common.utils import to_sha256 from pyrit.converter import Base64Converter, Converter, ConverterResult from pyrit.memory import CentralMemory, SQLiteMemory +from pyrit.message_normalizer import HistorySquashNormalizer from pyrit.models import ( AtomicAttackIdentifier, AttackIdentifier, @@ -59,7 +60,7 @@ PromptDataType, ) from pyrit.prompt_normalizer import ConverterConfiguration, PromptNormalizer -from pyrit.prompt_target.common.target_capabilities import TargetCapabilities +from pyrit.prompt_target.common.target_capabilities import CapabilityName, TargetCapabilities from pyrit.prompt_target.common.target_configuration import TargetConfiguration from unit.async_utils import wait_for_completion_async from unit.backend.mocks import ( @@ -2986,7 +2987,7 @@ async def hold_async(*, normalized_conversation: list[Message]) -> list[Message] await service.get_status_async(attack_result_id=ar.attack_result_id, send_id=status.send_id) ).state == MessageSendState.COMPLETED - async def test_repeat_send_rejects_copying_history_into_a_stateful_target_async( + async def test_repeat_send_rejects_copying_history_without_replay_support_async( self, *, real_send_context: tuple[MessageSendService, AttackResult, MockPromptTarget, Base64Converter], @@ -3002,7 +3003,7 @@ async def test_repeat_send_rejects_copying_history_into_a_stateful_target_async( capabilities=TargetCapabilities(supports_multi_turn=True, supports_multi_message_pieces=True) ) - with pytest.raises(ValueError, match="keeps its own conversation history"): + with pytest.raises(ValueError, match="cannot initialize copies"): await service.submit_async( attack_result_id=ar.attack_result_id, request=_submission(conversation_id=ar.conversation_id, submission_id="repeat", count=3), @@ -3012,15 +3013,57 @@ async def test_repeat_send_rejects_copying_history_into_a_stateful_target_async( assert [piece.role for piece in pieces] == ["user", "assistant"] assert not service._scheduler._conversations - async def test_repeat_send_into_a_stateful_target_is_fine_on_a_new_conversation_async( + @pytest.mark.parametrize("capability", [CapabilityName.EDITABLE_HISTORY, CapabilityName.MULTI_TURN]) + async def test_repeat_send_replays_history_through_a_configured_adapter_async( self, *, real_send_context: tuple[MessageSendService, AttackResult, MockPromptTarget, Base64Converter], + capability: CapabilityName, ) -> None: service, ar, target, _ = real_send_context + first = await service.submit_async( + attack_result_id=ar.attack_result_id, request=_submission(conversation_id=ar.conversation_id) + ) + assert (await _settle_send_async(service=service, status=first)).state == MessageSendState.COMPLETED target._configuration = TargetConfiguration( - capabilities=TargetCapabilities(supports_multi_turn=True, supports_multi_message_pieces=True) + capabilities=TargetCapabilities(supports_multi_turn=capability == CapabilityName.EDITABLE_HISTORY), + normalizer_overrides={capability: HistorySquashNormalizer()}, + ) + + with patch.object( + target, "_send_prompt_to_target_async", wraps=target._send_prompt_to_target_async + ) as provider: + status = await service.submit_async( + attack_result_id=ar.attack_result_id, + request=_submission(conversation_id=ar.conversation_id, submission_id="repeat", count=3), + ) + assert (await _settle_send_async(service=service, status=status)).state == MessageSendState.COMPLETED + + assert provider.await_count == 3 + for call in provider.call_args_list: + conversation = call.kwargs["normalized_conversation"] + assert len(conversation) == 1 + prompt = conversation[0].get_piece().converted_value + assert "[Conversation History]" in prompt + assert "User: Hello" in prompt + assert "Assistant: default" in prompt + assert "[Current Message]" in prompt + + @pytest.mark.parametrize("system_history", [False, True]) + async def test_repeat_send_without_replay_support_is_fine_on_a_new_conversation_async( + self, + *, + real_send_context: tuple[MessageSendService, AttackResult, MockPromptTarget, Base64Converter], + system_history: bool, + ) -> None: + service, ar, target, _ = real_send_context + target._configuration = TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_turn=True, supports_multi_message_pieces=True, supports_system_prompt=True + ) ) + if system_history: + await target.set_system_prompt_async(system_prompt="Be brief.", conversation_id=ar.conversation_id) status = await service.submit_async( attack_result_id=ar.attack_result_id, diff --git a/tests/unit/prompt_target/target/test_prompt_target.py b/tests/unit/prompt_target/target/test_prompt_target.py index feed09a5b4..3ea93d2c7e 100644 --- a/tests/unit/prompt_target/target/test_prompt_target.py +++ b/tests/unit/prompt_target/target/test_prompt_target.py @@ -8,10 +8,11 @@ import pytest from openai.types.chat import ChatCompletion -from unit.mocks import get_sample_conversations, openai_chat_response_json_dict +from unit.mocks import MockPromptTarget, get_sample_conversations, openai_chat_response_json_dict from pyrit.executor.attack.core.attack_strategy import AttackStrategy from pyrit.memory.memory_interface import MemoryInterface +from pyrit.message_normalizer import HistorySquashNormalizer from pyrit.models import ( ChatMessageRole, ComponentIdentifier, @@ -120,6 +121,71 @@ def test_validate_history_preserves_provider_validation(azure_openai_target: Ope azure_openai_target.validate_history(history) +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize( + ("configuration", "allowed"), + [ + pytest.param( + TargetConfiguration( + capabilities=TargetCapabilities(supports_multi_turn=True, supports_editable_history=True) + ), + True, + id="native-editable-history", + ), + pytest.param( + TargetConfiguration(capabilities=TargetCapabilities(supports_multi_turn=True)), + False, + id="no-replay-path", + ), + pytest.param( + TargetConfiguration( + capabilities=TargetCapabilities(supports_multi_turn=True), + policy=CapabilityHandlingPolicy( + behaviors={CapabilityName.EDITABLE_HISTORY: UnsupportedCapabilityBehavior.ADAPT} + ), + ), + False, + id="policy-without-adapter", + ), + pytest.param( + TargetConfiguration( + capabilities=TargetCapabilities(supports_multi_turn=True), + normalizer_overrides={CapabilityName.EDITABLE_HISTORY: HistorySquashNormalizer()}, + ), + True, + id="explicit-history-adapter", + ), + pytest.param( + TargetConfiguration( + capabilities=TargetCapabilities(), + policy=CapabilityHandlingPolicy( + behaviors={CapabilityName.MULTI_TURN: UnsupportedCapabilityBehavior.ADAPT} + ), + ), + True, + id="single-turn-history-adaptation", + ), + pytest.param( + TargetConfiguration( + capabilities=TargetCapabilities(supports_multi_turn=True), + normalizer_overrides={CapabilityName.MULTI_TURN: HistorySquashNormalizer()}, + ), + False, + id="unused-multi-turn-adapter", + ), + ], +) +def test_validate_history_replay_checks_configured_path(*, configuration: TargetConfiguration, allowed: bool) -> None: + target = MockPromptTarget(custom_configuration=configuration) + if allowed: + target.validate_history_replay() + else: + with pytest.raises(ValueError, match="supports_editable_history"): + target.validate_history_replay() + assert target.configuration is configuration + assert not target.prompt_sent + + async def test_set_system_prompt(azure_openai_target: OpenAIChatTarget, mock_attack_strategy: AttackStrategy): ( await azure_openai_target.set_system_prompt_async( diff --git a/tests/unit/prompt_target/target/test_websocket_target.py b/tests/unit/prompt_target/target/test_websocket_target.py index 2b6358b3c0..2dd3a23438 100644 --- a/tests/unit/prompt_target/target/test_websocket_target.py +++ b/tests/unit/prompt_target/target/test_websocket_target.py @@ -18,6 +18,7 @@ ) from pyrit.memory import SQLiteMemory from pyrit.models import Message, MessagePiece +from pyrit.prompt_normalizer import PromptNormalizer from pyrit.prompt_target import WebsocketTarget @@ -68,6 +69,58 @@ def create_message(*, value: str = "Hello", conversation_id: str = "conversation ).to_message() +@pytest.mark.usefixtures("patch_central_database") +def test_validate_history_replay_requires_a_restore_path(websocket_target: WebsocketTarget) -> None: + with pytest.raises(ValueError, match="supports_editable_history"): + websocket_target.validate_history_replay() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_validate_history_replay_allows_restoration_on_a_new_conversation_async( + *, + response_parser: Callable[[str | bytes], str | None], + message_builder: Callable[[str], str | bytes], + sqlite_instance: SQLiteMemory, +) -> None: + 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, + ) + history = [ + create_message(value="Earlier request"), + MessagePiece(role="assistant", original_value="Earlier reply", conversation_id="conversation").to_message(), + ] + for message in history: + await sqlite_instance.add_message_to_memory_async(request=message) + connection = AsyncMock(spec=ClientConnection) + connection.state = State.OPEN + connection.__aiter__.return_value = ['{"message": "Restored reply"}'] + + target.validate_history_replay() + assert not target.capabilities.supports_editable_history + restore_callback.assert_not_awaited() + with patch.object(target, "_connect_async", return_value=connection): + try: + response = await PromptNormalizer().send_prompt_async( + message=create_message(value="Follow-up"), conversation_id="conversation", target=target + ) + finally: + await target.cleanup_target_async() + + assert response.get_piece().converted_value == "Restored reply" + restore_callback.assert_awaited_once() + restored_connection, restored_history = restore_callback.call_args.args + assert restored_connection is connection + assert [message.get_piece().converted_value for message in restored_history] == ["Earlier request", "Earlier reply"] + connection.send.assert_awaited_once_with('{"message": "Follow-up"}') + + def test_init_invalid_endpoint_raises( response_parser: Callable[[str | bytes], str | None], message_builder: Callable[[str], str | bytes],