From 662d8a11f9e8cf507c29fd40db361b490d18e2fc Mon Sep 17 00:00:00 2001 From: Joel Date: Sun, 27 Sep 2026 07:54:00 +0530 Subject: [PATCH 1/2] fix: preserve text after Prompt Shield documents --- pyrit/prompt_target/prompt_shield_target.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pyrit/prompt_target/prompt_shield_target.py b/pyrit/prompt_target/prompt_shield_target.py index 4da916a309..44d4c43d30 100644 --- a/pyrit/prompt_target/prompt_shield_target.py +++ b/pyrit/prompt_target/prompt_shield_target.py @@ -254,6 +254,7 @@ def _input_parser(self, input_str: str) -> dict[str, Any]: user_prompt += contents[0] else: documents.append(contents[0]) + user_prompt += contents[1] return {"userPrompt": user_prompt, "documents": documents if documents else []} From 594295ff25638b6cc69930b4e4273f38821b1eb6 Mon Sep 17 00:00:00 2001 From: Joel Date: Sun, 27 Sep 2026 07:54:04 +0530 Subject: [PATCH 2/2] test: cover Prompt Shield trailing document text --- .../target/test_prompt_shield_target.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/unit/prompt_target/target/test_prompt_shield_target.py b/tests/unit/prompt_target/target/test_prompt_shield_target.py index 96808fc3c0..f6c7a66367 100644 --- a/tests/unit/prompt_target/target/test_prompt_shield_target.py +++ b/tests/unit/prompt_target/target/test_prompt_shield_target.py @@ -84,6 +84,19 @@ async def test_prompt_shield_document_parsing( assert result == sample_delineated_prompt_as_dict +def test_prompt_shield_document_parsing_preserves_text_between_and_after_documents( + promptshield_target: PromptShieldTarget, +): + input_str = "q1 d1 q2 d2 q3" + + result = promptshield_target._input_parser(input_str) + + assert result == { + "userPrompt": "q1 q2 q3", + "documents": [" d1 ", " d2 "], + } + + async def test_prompt_shield_response_validation(promptshield_target: PromptShieldTarget): # This tests handling both an empty request and an empty response promptshield_target._validate_response(request_body={}, response_body={})