Parcourir la source

fix: harden provider tool transcripts

Problem: duplicate call IDs, reserved payload overrides, and truncated EventAgent histories could produce invalid provider transcripts. Risk: public pre_messages now reject internal tool protocol state and reserved extra_body keys fail locally.
zhenyu.hu il y a 2 semaines
Parent
commit
94d093ff16

+ 2 - 0
src/agent_lab/application/contracts.py

@@ -46,4 +46,6 @@ class DebugRunRequest(BaseModel):
     def reject_tool_pre_messages(self) -> "DebugRunRequest":
         if any(message.role == "tool" for message in self.pre_messages):
             raise ValueError("pre_messages cannot include tool messages")
+        if any(message.tool_calls for message in self.pre_messages):
+            raise ValueError("pre_messages cannot include assistant tool calls")
         return self

+ 23 - 5
src/agent_lab/application/event_agent.py

@@ -200,11 +200,7 @@ class EventAgent:
         ]
         if system_prompt.strip():
             messages.append(ChatMessage(role="system", content=system_prompt))
-        messages.extend(
-            message
-            for message in history
-            if message.role in {"system", "user", "assistant"}
-        )
+        messages.extend(self._project_visible_history(history))
         messages.append(
             ChatMessage(
                 role="user",
@@ -213,6 +209,28 @@ class EventAgent:
         )
         return messages
 
+    def _project_visible_history(
+        self,
+        history: Sequence[ChatMessage],
+    ) -> list[ChatMessage]:
+        projected: list[ChatMessage] = []
+        for message in history:
+            if message.role == "tool":
+                continue
+            if message.role == "assistant":
+                if not message.content:
+                    continue
+                projected.append(
+                    ChatMessage(
+                        role="assistant",
+                        content=message.content,
+                        name=message.name,
+                    )
+                )
+                continue
+            projected.append(message)
+        return projected
+
     def summarize_replies(self, replies: Sequence[ChatMessage]) -> ChatMessage | None:
         if not replies:
             return None

+ 10 - 0
src/agent_lab/domain/messages.py

@@ -21,6 +21,16 @@ class ChatMessage(BaseModel):
             raise ValueError("tool_calls require assistant role")
         if self.role == "tool" and not self.tool_call_id:
             raise ValueError("tool messages require tool_call_id")
+        if self.role != "tool" and self.tool_call_id is not None:
+            raise ValueError("tool_call_id requires tool role")
+        if self.tool_calls:
+            seen_ids: set[str] = set()
+            for tool_call in self.tool_calls:
+                if tool_call.id in seen_ids:
+                    raise ValueError(
+                        f"duplicate assistant tool-call ID: {tool_call.id}"
+                    )
+                seen_ids.add(tool_call.id)
         return self
 
 

+ 53 - 0
src/agent_lab/infrastructure/chat_client.py

@@ -9,6 +9,20 @@ from agent_lab.domain.messages import ChatMessage, StreamItem
 from agent_lab.infrastructure.openai_compatible import ChatCompletionStreamParser
 
 
+RESERVED_EXTRA_BODY_FIELDS = frozenset(
+    {
+        "model",
+        "messages",
+        "tools",
+        "tool_choice",
+        "stream",
+        "stream_options",
+        "temperature",
+        "max_tokens",
+    }
+)
+
+
 class OpenAICompatibleChatClient:
     def __init__(
         self,
@@ -36,6 +50,8 @@ class OpenAICompatibleChatClient:
         tool_choice: dict[str, Any] | None = None,
     ) -> AsyncIterator[StreamItem]:
         self._validate_tool_transcript(messages)
+        self._validate_tool_choice(tools, tool_choice)
+        self._validate_extra_body(params.extra_body)
         payload: dict[str, Any] = {
             "model": self._resolve_model(params.model),
             "messages": [self._serialize_message(message) for message in messages],
@@ -97,6 +113,7 @@ class OpenAICompatibleChatClient:
     def _validate_tool_transcript(self, messages: list[ChatMessage]) -> None:
         pending: dict[str, None] = {}
         replied_ids: set[str] = set()
+        declared_ids: set[str] = set()
 
         for message in messages:
             if pending:
@@ -125,12 +142,48 @@ class OpenAICompatibleChatClient:
                 raise ValueError(f"orphan tool reply: {tool_call_id}")
 
             if message.tool_calls:
+                for tool_call in message.tool_calls:
+                    if tool_call.id in declared_ids:
+                        raise ValueError(
+                            "duplicate assistant tool-call ID across transcript: "
+                            f"{tool_call.id}"
+                        )
+                    declared_ids.add(tool_call.id)
                 pending = {tool_call.id: None for tool_call in message.tool_calls}
 
         if pending:
             unresolved = ", ".join(pending)
             raise ValueError(f"unresolved tool calls at end: {unresolved}")
 
+    def _validate_tool_choice(
+        self,
+        tools: list[dict[str, Any]],
+        tool_choice: dict[str, Any] | None,
+    ) -> None:
+        if tool_choice is None or tool_choice.get("type") != "function":
+            return
+        if not tools:
+            raise ValueError("forced tool choice requires tools")
+
+        function = tool_choice.get("function")
+        selected_name = function.get("name") if isinstance(function, dict) else None
+        available_names = {
+            tool["function"]["name"]
+            for tool in tools
+            if isinstance(tool.get("function"), dict)
+            and isinstance(tool["function"].get("name"), str)
+        }
+        if selected_name not in available_names:
+            raise ValueError(
+                f"forced tool choice references unknown tool: {selected_name}"
+            )
+
+    def _validate_extra_body(self, extra_body: dict[str, Any]) -> None:
+        conflicting_fields = RESERVED_EXTRA_BODY_FIELDS.intersection(extra_body)
+        if conflicting_fields:
+            field = sorted(conflicting_fields)[0]
+            raise ValueError(f"extra_body contains reserved field: {field}")
+
     def _serialize_message(self, message: ChatMessage) -> dict[str, Any]:
         if message.role == "tool":
             return {

+ 47 - 0
tests/test_event_agent.py

@@ -369,3 +369,50 @@ async def test_event_agent_llm_receives_history_and_agent_config_context():
         "message": "Use strict tool parameters.",
         "thinking": "disabled",
     }
+
+
+@pytest.mark.asyncio
+async def test_event_agent_projects_complete_tool_round_to_visible_history():
+    chat_client = ToolCallingChatClient({"message": "projected history"})
+    event = ToolCallEvent(
+        id="call_current",
+        name="handoff_note",
+        arguments={},
+        raw_arguments="{}",
+    )
+    prior_call = ToolCallEvent(
+        id="call_prior",
+        name="mock_search",
+        arguments={"query": "latency"},
+        raw_arguments='{"query":"latency"}',
+    )
+
+    await EventAgent(
+        enabled_tools=["handoff_note"],
+        chat_client=chat_client,
+    ).handle(
+        event,
+        history=[
+            ChatMessage(role="user", content="Find latency docs"),
+            ChatMessage(
+                role="assistant",
+                content="I checked the latency sources.",
+                tool_calls=[prior_call],
+            ),
+            ChatMessage(
+                role="tool",
+                content='{"results":["doc"]}',
+                tool_call_id="call_prior",
+            ),
+            ChatMessage(role="user", content="Prepare a handoff"),
+        ],
+    )
+
+    projected_messages = chat_client.calls[0]["messages"]
+    visible_assistant = next(
+        message
+        for message in projected_messages
+        if message.content == "I checked the latency sources."
+    )
+    assert visible_assistant.tool_calls == []
+    assert not any(message.role == "tool" for message in projected_messages)

+ 208 - 0
tests/test_websocket_api.py

@@ -413,6 +413,30 @@ def test_debug_run_request_rejects_tool_pre_messages():
         DebugRunRequest.model_validate(payload)
 
 
+def test_debug_run_request_rejects_assistant_tool_calls_in_pre_messages():
+    payload = _request_payload()
+    payload["pre_messages"] = [
+        {
+            "role": "assistant",
+            "content": "I will inspect this.",
+            "tool_calls": [
+                {
+                    "id": "call_1",
+                    "name": "mock_search",
+                    "arguments": {"query": "latency"},
+                    "raw_arguments": '{"query":"latency"}',
+                }
+            ],
+        }
+    ]
+
+    with pytest.raises(
+        ValidationError,
+        match="pre_messages cannot include assistant tool calls",
+    ):
+        DebugRunRequest.model_validate(payload)
+
+
 def test_chat_message_allows_internal_tool_replies():
     message = ChatMessage(
         role="tool",
@@ -482,6 +506,34 @@ def test_chat_message_rejects_tool_role_without_tool_call_id():
         ChatMessage(role="tool", content="orphan result")
 
 
+@pytest.mark.parametrize("role", ["system", "user", "assistant"])
+def test_chat_message_rejects_tool_call_id_on_non_tool_roles(role: str):
+    with pytest.raises(ValidationError, match="tool_call_id requires tool role"):
+        ChatMessage(role=role, content="invalid", tool_call_id="call_1")
+
+
+def test_chat_message_rejects_duplicate_assistant_tool_call_ids():
+    with pytest.raises(ValidationError, match="duplicate assistant tool-call ID: call_1"):
+        ChatMessage(
+            role="assistant",
+            content="checking twice",
+            tool_calls=[
+                ToolCallEvent(
+                    id="call_1",
+                    name="mock_search",
+                    arguments={"query": "first"},
+                    raw_arguments='{"query":"first"}',
+                ),
+                ToolCallEvent(
+                    id="call_1",
+                    name="handoff_note",
+                    arguments={"message": "second"},
+                    raw_arguments='{"message":"second"}',
+                ),
+            ],
+        )
+
+
 class HistoryCapturingChatClient:
     def __init__(self) -> None:
         self.calls = 0
@@ -492,6 +544,7 @@ class HistoryCapturingChatClient:
         messages: list[ChatMessage],
         tools: list[dict],
         params: AgentParams,
+        tool_choice: dict | None = None,
     ) -> AsyncIterator[StreamItem]:
         if tools:
             yield _event_tool_call_from_tools(tools, messages)
@@ -978,6 +1031,106 @@ async def test_openai_chat_client_rejects_unfinished_tool_calls_at_end():
     )
 
 
+@pytest.mark.asyncio
+async def test_openai_chat_client_rejects_reused_tool_call_id_across_transcript():
+    await _assert_tool_transcript_rejected(
+        [
+            ChatMessage(
+                role="assistant",
+                content="first round",
+                tool_calls=[
+                    ToolCallEvent(
+                        id="call_1",
+                        name="mock_search",
+                        arguments={},
+                        raw_arguments="{}",
+                    )
+                ],
+            ),
+            ChatMessage(role="tool", content="done", tool_call_id="call_1"),
+            ChatMessage(
+                role="assistant",
+                content="second round",
+                tool_calls=[
+                    ToolCallEvent(
+                        id="call_1",
+                        name="handoff_note",
+                        arguments={},
+                        raw_arguments="{}",
+                    )
+                ],
+            ),
+        ],
+        "duplicate assistant tool-call ID across transcript: call_1",
+    )
+
+
+async def _assert_tool_choice_rejected(
+    tools: list[dict],
+    tool_choice: dict,
+    error_match: str,
+) -> None:
+    network_called = False
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal network_called
+        network_called = True
+        return httpx.Response(200)
+
+    async with httpx.AsyncClient(
+        transport=httpx.MockTransport(handler),
+        base_url="https://llm.test/v1",
+    ) as http_client:
+        client = OpenAICompatibleChatClient(
+            api_key="",
+            base_url="https://llm.test/v1",
+            default_model="provider-default",
+            request_timeout_seconds=5,
+            http_client=http_client,
+        )
+        with pytest.raises(ValueError, match=error_match):
+            [
+                item
+                async for item in client.stream_chat(
+                    messages=[ChatMessage(role="user", content="use a tool")],
+                    tools=tools,
+                    params=AgentParams(model="provider-default"),
+                    tool_choice=tool_choice,
+                )
+            ]
+
+    assert network_called is False
+
+
+@pytest.mark.asyncio
+async def test_openai_chat_client_rejects_forced_choice_without_tools():
+    await _assert_tool_choice_rejected(
+        tools=[],
+        tool_choice={
+            "type": "function",
+            "function": {"name": "handoff_note"},
+        },
+        error_match="forced tool choice requires tools",
+    )
+
+
+@pytest.mark.asyncio
+async def test_openai_chat_client_rejects_forced_choice_for_unknown_tool():
+    await _assert_tool_choice_rejected(
+        tools=[
+            {
+                "type": "function",
+                "function": {"name": "mock_search", "parameters": {}},
+            }
+        ],
+        tool_choice={
+            "type": "function",
+            "function": {"name": "handoff_note"},
+        },
+        error_match="forced tool choice references unknown tool: handoff_note",
+    )
+
+
 def test_static_pre_message_role_selector_does_not_offer_tool():
     html = Path("src/agent_lab/presentation/static/index.html").read_text()
 
@@ -1387,3 +1540,58 @@ async def test_openai_chat_client_merges_agent_extra_body_into_payload():
     assert captured_payloads[0]["thinking"] == {"type": "disabled"}
     assert captured_payloads[0]["enable_search"] is False
     assert captured_payloads[0]["search_options"] == {"forced_search": False}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("reserved_key", "override_value"),
+    [
+        ("model", "overridden-model"),
+        ("messages", []),
+        ("tools", []),
+        (
+            "tool_choice",
+            {"type": "function", "function": {"name": "other_tool"}},
+        ),
+        ("stream", False),
+        ("stream_options", {"include_usage": False}),
+        ("temperature", 1.0),
+        ("max_tokens", 1),
+    ],
+)
+async def test_openai_chat_client_rejects_reserved_extra_body_fields_before_network(
+    reserved_key: str,
+    override_value: object,
+):
+    network_called = False
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal network_called
+        network_called = True
+        return httpx.Response(200)
+
+    async with httpx.AsyncClient(
+        transport=httpx.MockTransport(handler),
+        base_url="https://llm.test/v1",
+    ) as http_client:
+        client = OpenAICompatibleChatClient(
+            api_key="",
+            base_url="https://llm.test/v1",
+            default_model="provider-default",
+            request_timeout_seconds=5,
+            http_client=http_client,
+        )
+        with pytest.raises(
+            ValueError,
+            match=f"extra_body contains reserved field: {reserved_key}",
+        ):
+            [
+                item
+                async for item in client.stream_chat(
+                    messages=[ChatMessage(role="user", content="hi")],
+                    tools=[],
+                    params=AgentParams(extra_body={reserved_key: override_value}),
+                )
+            ]
+
+    assert network_called is False