keep web search calls replayable and out of the way of the answer

Author: u7k4rs6

pyrit/prompt_target/openai/openai_response_target.py

@@ -89,6 +89,10 @@ class MessagePieceType(str, Enum):
     MCP_APPROVAL_REQUEST = "mcp_approval_request"
 
 
+# Fields of a web_search_call input item (all required); the output item has the same ones.
+_WEB_SEARCH_CALL_INPUT_FIELDS: tuple[str, ...] = ("id", "type", "status", "action")
+
+
 class _ResponseToolCallContent(BaseModel):
     model_config = ConfigDict(extra="allow")
     type: str = Field(min_length=1, pattern=r"\S")
@@ -347,14 +351,14 @@ def _serialize_function_call(self, piece: MessagePiece) -> "ResponseFunctionTool
             "arguments": call.function.arguments,
         }
 
-    def _serialize_tool_call(self, piece: MessagePiece) -> dict[str, Any]:
+    def _serialize_tool_call(self, piece: MessagePiece) -> dict[str, Any] | None:
         stored = _ResponseToolCallContent.model_validate_json(piece.converted_value).model_dump(exclude_unset=True)
         if stored.get("type") == "web_search_call":
-            return {
-                "type": stored["type"],
-                "call_id": stored.get("call_id"),
-                "query": stored.get("query"),
-            }
+            # The input item takes exactly the output item's fields, and all of them are required. Older rows
+            # kept only type and id, which the API rejects, so those can't be replayed.
+            if any(stored.get(key) is None for key in _WEB_SEARCH_CALL_INPUT_FIELDS):
+                return None
+            return {key: stored[key] for key in _WEB_SEARCH_CALL_INPUT_FIELDS}
         filtered = {"type": stored["type"]}
         filtered.update({key: stored[key] for key in ("call_id", "query", "name", "arguments") if key in stored})
         return filtered
@@ -377,7 +381,14 @@ async def _serialize_piece_async(self, *, piece: MessagePiece, message_index: in
         if data_type == "function_call":
             return _SerializedPiece(item=dict(self._serialize_function_call(piece)), placement="top_level")
         if data_type == "tool_call":
-            return _SerializedPiece(item=self._serialize_tool_call(piece), placement="top_level")
+            tool_call = self._serialize_tool_call(piece)
+            if tool_call is None:
+                logger.warning(
+                    "Skipping web_search_call in message index %d: it lacks the id, status or action the API needs.",
+                    message_index,
+                )
+                return None
+            return _SerializedPiece(item=tool_call, placement="top_level")
         if data_type == "function_call_output":
             return _SerializedPiece(item=dict(self._serialize_function_call_output(piece)), placement="top_level")
         raise ValueError(f"Unsupported data type '{data_type}' in message index {message_index}")
@@ -610,7 +621,14 @@ async def _construct_message_from_response_async(self, response: Response, reque
         # reasoning commonly precedes the actual message in provider output, so
         # retain it for memory/debugging after the actionable response pieces.
         # This must stay ahead of the metadata writes below, which target the first piece.
-        extracted_response_pieces.sort(key=lambda piece: piece.converted_value_data_type == "reasoning")
+        # Hosted tool records (e.g. web_search_call) also precede the answer; keep them after it too. Replay
+        # order is unaffected because tool records are always emitted before the message item.
+        extracted_response_pieces.sort(
+            key=lambda piece: (
+                piece.converted_value_data_type == "reasoning",
+                piece.converted_value_data_type == "tool_call",
+            )
+        )
 
         # Capture token usage and the stop reason in the first piece's metadata. This also runs on
         # the truncated path: usage is populated on token-limit responses and is most valuable
@@ -834,20 +852,9 @@ def _parse_response_output_section(
             piece_type = "function_call"
 
         elif section_type == MessagePieceType.WEB_SEARCH_CALL:
-            # Forward web_search_call with only API-expected fields
-            # Note: web search may have different field structure than function calls
-            web_search_data = {
-                "type": "web_search_call",
-            }
-            # Add optional fields if they exist
-            if hasattr(section, "call_id") and section.call_id:
-                web_search_data["call_id"] = section.call_id
-            if hasattr(section, "query") and section.query:
-                web_search_data["query"] = section.query
-            if hasattr(section, "id") and section.id:
-                web_search_data["id"] = section.id
-
-            piece_value = json.dumps(web_search_data, separators=(",", ":"))
+            # Keep the whole item (id, status and the search action). The API needs all of it to accept the
+            # call back as input on the next turn.
+            piece_value = json.dumps(section.model_dump(mode="json", exclude_none=True), separators=(",", ":"))
             piece_type = "tool_call"
 
         elif section_type == "custom_tool_call":

tests/unit/backend/test_conversation_editor.py

@@ -476,7 +476,10 @@ async def test_provider_tool_extensions_round_trip_without_execution_async(
     ) -> None:
         request = draft()
         request.target_registry_name = "responses"
-        payload = '{"type":"web_search_call","call_id":"web-1","query":"query","extension":{"keep":true}}'
+        payload = (
+            '{"type":"web_search_call","id":"ws-1","status":"completed",'
+            '"action":{"type":"search","query":"query"},"extension":{"keep":true}}'
+        )
         request.messages = [
             ConversationMessageRequest(
                 role="simulated_assistant",

tests/unit/prompt_target/target/test_openai_response_adapters.py

@@ -48,12 +48,20 @@ def test_tool_preflight_and_serialization_reject_invalid_provider_fields(payload
 @pytest.mark.usefixtures("patch_central_database")
 def test_tool_preflight_accepts_empty_history_and_preserves_provider_extensions() -> None:
     target = object.__new__(OpenAIResponseTarget)
-    payload = '{"type":"web_search_call","call_id":"web-1","query":"query","extension":{"key":"value"}}'
+    payload = (
+        '{"type":"web_search_call","id":"ws-1","status":"completed",'
+        '"action":{"type":"search","query":"query"},"extension":{"key":"value"}}'
+    )
     piece = MessagePiece(role="simulated_assistant", original_value_data_type="tool_call", original_value=payload)
     target.validate_tool_history([])
     target.validate_tool_history([Message(message_pieces=[piece])])
     assert piece.converted_value == payload
-    assert target._serialize_tool_call(piece) == {"type": "web_search_call", "call_id": "web-1", "query": "query"}
+    assert target._serialize_tool_call(piece) == {
+        "id": "ws-1",
+        "type": "web_search_call",
+        "status": "completed",
+        "action": {"type": "search", "query": "query"},
+    }
 
 
 @pytest.mark.usefixtures("patch_central_database")

tests/unit/prompt_target/target/test_openai_response_target.py

@@ -7,12 +7,18 @@
 from collections.abc import MutableSequence
 from tempfile import NamedTemporaryFile
 from types import SimpleNamespace
-from typing import Any
+from typing import Any, get_type_hints
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from openai import BadRequestError, RateLimitError
-from openai.types.responses import ResponseOutputMessage, ResponseOutputRefusal, ResponseOutputText
+from openai.types.responses import (
+    Response,
+    ResponseFunctionWebSearchParam,
+    ResponseOutputMessage,
+    ResponseOutputRefusal,
+    ResponseOutputText,
+)
 from unit.mocks import (
     get_audio_message_piece,
     get_image_message_piece,
@@ -53,8 +59,6 @@ def create_mock_response(response_dict: dict = None) -> MagicMock:
     Returns:
         A mock object that simulates the OpenAI SDK response with Pydantic-style attribute access.
     """
-    from openai.types.responses import Response
-
     if response_dict is None:
         response_dict = openai_response_json_dict()
 
@@ -970,7 +974,15 @@ async def test_build_input_for_multi_modal_async_preserves_mixed_payload_contrac
                 ),
                 MessagePiece(
                     role="assistant",
-                    original_value=json.dumps({"type": "web_search_call", "call_id": "web-1", "id": "drop-id"}),
+                    original_value=json.dumps(
+                        {
+                            "type": "web_search_call",
+                            "id": "ws-1",
+                            "status": "completed",
+                            "action": {"type": "search", "query": "query"},
+                            "extra": "drop",
+                        }
+                    ),
                     original_value_data_type="tool_call",
                 ),
                 MessagePiece(
@@ -1036,7 +1048,12 @@ async def test_build_input_for_multi_modal_async_preserves_mixed_payload_contrac
             "name": "lookup",
             "arguments": '{"value":1}',
         },
-        {"type": "web_search_call", "call_id": "web-1", "query": None},
+        {
+            "id": "ws-1",
+            "type": "web_search_call",
+            "status": "completed",
+            "action": {"type": "search", "query": "query"},
+        },
         {
             "type": "provider_tool_call",
             "call_id": "tool-1",
@@ -1055,6 +1072,86 @@ async def test_build_input_for_multi_modal_async_preserves_mixed_payload_contrac
     ]
 
 
+def _web_search_response() -> Response:
+    return Response.model_validate(
+        {
+            "id": "resp-1",
+            "object": "response",
+            "created_at": 1,
+            "model": "gpt-4.1",
+            "parallel_tool_calls": True,
+            "tool_choice": "auto",
+            "tools": [{"type": "web_search_preview"}],
+            "status": "completed",
+            "output": [
+                {
+                    "type": "web_search_call",
+                    "id": "ws-1",
+                    "status": "completed",
+                    "action": {"type": "search", "query": "pyrit release"},
+                },
+                {
+                    "type": "message",
+                    "id": "msg-1",
+                    "role": "assistant",
+                    "status": "completed",
+                    "content": [{"type": "output_text", "text": "PyRIT 0.13 is out.", "annotations": []}],
+                },
+            ],
+        }
+    )
+
+
+async def test_web_search_answer_is_the_response_value_async(
+    target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece
+):
+    reply = await target._construct_message_from_response_async(_web_search_response(), dummy_text_message_piece)
+
+    assert reply.get_value() == "PyRIT 0.13 is out."
+    assert [piece.converted_value_data_type for piece in reply.message_pieces] == ["text", "tool_call"]
+
+
+async def test_web_search_call_replays_as_a_valid_input_item_async(
+    target: OpenAIResponseTarget, dummy_text_message_piece: MessagePiece
+):
+    reply = await target._construct_message_from_response_async(_web_search_response(), dummy_text_message_piece)
+    follow_up = Message.from_prompt(prompt="Anything else?", role="user")
+
+    items = await target._build_input_for_multi_modal_async([dummy_text_message_piece.to_message(), reply, follow_up])
+
+    web_search = {
+        "id": "ws-1",
+        "type": "web_search_call",
+        "status": "completed",
+        "action": {"type": "search", "query": "pyrit release"},
+    }
+    assert items[1] == web_search
+    assert set(web_search) == set(get_type_hints(ResponseFunctionWebSearchParam))
+    assert items[2] == {"role": "assistant", "content": [{"type": "output_text", "text": "PyRIT 0.13 is out."}]}
+
+
+async def test_web_search_call_stored_without_its_action_is_skipped_on_replay_async(
+    target: OpenAIResponseTarget, caplog: pytest.LogCaptureFixture
+):
+    reply = Message(
+        message_pieces=[
+            MessagePiece(role="assistant", original_value="answer", conversation_id="c-1"),
+            MessagePiece(
+                role="assistant",
+                original_value='{"type":"web_search_call","id":"ws-1"}',
+                original_value_data_type="tool_call",
+                conversation_id="c-1",
+            ),
+        ]
+    )
+
+    with caplog.at_level(logging.WARNING):
+        items = await target._build_input_for_multi_modal_async([reply])
+
+    assert items == [{"role": "assistant", "content": [{"type": "output_text", "text": "answer"}]}]
+    assert "Skipping web_search_call" in caplog.text
+
+
 @pytest.mark.parametrize("data_type", ["function_call", "tool_call", "function_call_output"])
 async def test_build_input_for_multi_modal_async_validates_malformed_artifact_async(
     target: OpenAIResponseTarget, data_type: PromptDataType

tests/unit/prompt_target/target/test_tool_call_history.py

@@ -345,7 +345,10 @@ def test_responses_preflight_and_serializer_reject_invalid_tool_fields(payload:
 def test_responses_preflight_retains_provider_extensions_and_uses_converted_values() -> None:
     target = _target(OpenAIResponseTarget)
     assert isinstance(target, OpenAIResponseTarget)
-    payload = '{"type":"web_search_call","call_id":"web-1","query":"edited","extension":{"key":"value"}}'
+    payload = (
+        '{"type":"web_search_call","id":"ws-1","status":"completed",'
+        '"action":{"type":"search","query":"edited"},"extension":{"key":"value"}}'
+    )
     piece = MessagePiece(
         role="simulated_assistant",
         original_value_data_type="tool_call",
@@ -354,7 +357,12 @@ def test_responses_preflight_retains_provider_extensions_and_uses_converted_valu
     )
     target.validate_tool_history([piece.to_message()])
     assert piece.converted_value == payload
-    assert target._serialize_tool_call(piece) == {"type": "web_search_call", "call_id": "web-1", "query": "edited"}
+    assert target._serialize_tool_call(piece) == {
+        "id": "ws-1",
+        "type": "web_search_call",
+        "status": "completed",
+        "action": {"type": "search", "query": "edited"},
+    }
 
 
 @pytest.mark.parametrize("target_type", [OpenAIChatTarget, LiteLLMChatTarget])