Преглед изворни кода

fix: enforce typed event source boundaries

Problem: legacy event consumers allowed stale fixtures to mask incorrect typed event wiring. Risk: compatible providers may omit finish_reason; the client integration test verifies exactly-once flush at [DONE].
zhenyu.hu пре 2 недеља
родитељ
комит
b5535c9d4f

+ 1 - 4
src/agent_lab/application/event_agent.py

@@ -166,10 +166,7 @@ class EventAgent:
                 raw_chunks.append(item.raw_chunk)
                 continue
 
-            if (
-                item.kind in {"provider_tool_call", "event"}
-                and item.event is not None
-            ):
+            if item.kind == "provider_tool_call" and item.event is not None:
                 self._raw_chunks_by_event[event.id] = raw_chunks
                 return event.model_copy(
                     update={

+ 2 - 2
src/agent_lab/application/runtime.py

@@ -219,7 +219,7 @@ class DebugRuntime:
                     )
                     continue
 
-                if item.kind in {"text_event", "event"} and item.event is not None:
+                if item.kind == "text_event" and item.event is not None:
                     saw_event = True
                     event = self._event_name_only(item.event)
                     events.append(event)
@@ -491,7 +491,7 @@ class DebugRuntime:
                     )
                     continue
 
-                if item.kind in {"text_event", "event"} and item.event is not None:
+                if item.kind == "text_event" and item.event is not None:
                     saw_event = True
                     event = self._event_name_only(item.event)
                     events.append(event)

+ 61 - 0
tests/test_debug_runtime.py

@@ -168,6 +168,27 @@ class FakeChatClient:
         yield StreamItem.message_delta("final answer")
 
 
+class WrongSourceEventChatClient:
+    def __init__(self, item: StreamItem) -> None:
+        self.item = item
+        self.calls = 0
+
+    async def stream_chat(
+        self,
+        messages: list[ChatMessage],
+        tools: list[dict],
+        params: AgentParams,
+    ) -> AsyncIterator[StreamItem]:
+        if tools:
+            yield _event_tool_call_from_tools(tools, messages)
+            return
+        self.calls += 1
+        if self.calls == 1:
+            yield self.item
+            return
+        yield StreamItem.message_delta("unexpected continuation")
+
+
 class IncrementingClock:
     def __init__(self, current: float = 100.0, step: float = 0.01) -> None:
         self.current = current
@@ -550,6 +571,46 @@ async def test_runtime_routes_chat_events_through_event_agent_then_continues_cha
     assert business_outputs[4]["content"] == "final answer"
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "item",
+    [
+        StreamItem.provider_tool_call(
+            ToolCallEvent(
+                id="provider_call_1",
+                name="handoff_note",
+                arguments={"message": "wrong source"},
+                raw_arguments='{"message":"wrong source"}',
+            )
+        ),
+        StreamItem.event(
+            ToolCallEvent(
+                id="legacy_event_1",
+                name="handoff_note",
+                arguments={},
+                raw_arguments="{}",
+            )
+        ),
+    ],
+    ids=["provider_tool_call", "legacy_event"],
+)
+async def test_runtime_ignores_non_text_event_sources(item: StreamItem):
+    request = DebugRunRequest(
+        user_message="debug this",
+        system_prompts=[],
+        pre_messages=[],
+        chat_agent=AgentParams(model="fake-model"),
+        event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=2),
+    )
+    client = WrongSourceEventChatClient(item)
+
+    outputs = [message async for message in DebugRuntime(client).run(request)]
+
+    assert client.calls == 1
+    assert "event" not in _message_types(outputs)
+    assert "tool_result" not in _message_types(outputs)
+
+
 @pytest.mark.asyncio
 async def test_runtime_emits_audit_events_and_backend_logs(caplog):
     request = DebugRunRequest(

+ 58 - 0
tests/test_event_agent.py

@@ -76,6 +76,19 @@ class NoToolCallChatClient:
         yield StreamItem.message_delta("I should have called the tool.")
 
 
+class WrongSourceToolCallChatClient:
+    def __init__(self, item: StreamItem) -> None:
+        self.item = item
+
+    async def stream_chat(
+        self,
+        messages: list[ChatMessage],
+        tools: list[dict],
+        params: AgentParams,
+    ) -> AsyncIterator[StreamItem]:
+        yield self.item
+
+
 @pytest.mark.asyncio
 async def test_event_agent_resolves_tool_arguments_with_llm_tool_call():
     chat_client = ToolCallingChatClient({"message": "LLM generated handoff"})
@@ -174,6 +187,51 @@ async def test_event_agent_falls_back_to_context_arguments_when_llm_returns_no_t
     assert chat_client.calls[0]["tools"][0]["function"]["name"] == "mock_search"
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "item",
+    [
+        StreamItem.text_event(
+            ToolCallEvent(
+                id="text_event_1",
+                name="mock_search",
+                arguments={"query": "wrong source"},
+                raw_arguments='{"query":"wrong source"}',
+            )
+        ),
+        StreamItem.event(
+            ToolCallEvent(
+                id="legacy_event_1",
+                name="mock_search",
+                arguments={"query": "legacy arguments"},
+                raw_arguments='{"query":"legacy arguments"}',
+            )
+        ),
+    ],
+    ids=["text_event", "legacy_event"],
+)
+async def test_event_agent_ignores_non_provider_tool_call_sources(item: StreamItem):
+    event = ToolCallEvent(
+        id="call_1",
+        name="mock_search",
+        arguments={},
+        raw_arguments="{}",
+    )
+    agent = EventAgent(
+        enabled_tools=["mock_search"],
+        chat_client=WrongSourceToolCallChatClient(item),
+    )
+
+    reply = await agent.handle(
+        event,
+        history=[ChatMessage(role="assistant", content="fallback query")],
+    )
+
+    payload = json.loads(reply.content)
+    assert payload["tool"] == "mock_search"
+    assert payload["query"] == "fallback query"
+
+
 @pytest.mark.asyncio
 async def test_event_agent_returns_registry_errors_for_disabled_and_unknown_tools():
     registry = ToolRegistry(

+ 56 - 2
tests/test_websocket_api.py

@@ -139,7 +139,7 @@ def _event_tool_call_from_tools(
         "query": content,
         "title": content,
     }
-    return StreamItem.event(
+    return StreamItem.provider_tool_call(
         ToolCallEvent(
             id="event_agent_call_1",
             name=tool_name,
@@ -442,7 +442,7 @@ class HistoryCapturingChatClient:
         self.calls += 1
         if self.calls == 1:
             yield StreamItem.message_delta("Need event help.")
-            yield StreamItem.event(
+            yield StreamItem.text_event(
                 ToolCallEvent(
                     id="call_1",
                     name="handoff_note",
@@ -544,6 +544,60 @@ async def test_openai_chat_client_streams_sse_chunks_through_parser():
     ]
 
 
+@pytest.mark.asyncio
+async def test_openai_chat_client_flushes_one_provider_tool_call_at_done():
+    def handler(request: httpx.Request) -> httpx.Response:
+        return httpx.Response(
+            200,
+            content=(
+                b'data: {"choices":[{"delta":{"tool_calls":[{"index":0,'
+                b'"id":"call_1","function":{"name":"mock_search",'
+                b'"arguments":"{\\"query\\":\\""}}]},"finish_reason":null}]}\n\n'
+                b'data: {"choices":[{"delta":{"tool_calls":[{"index":0,'
+                b'"function":{"arguments":"latency docs\\"}"}}]},'
+                b'"finish_reason":null}]}\n\n'
+                b"data: [DONE]\n\n"
+            ),
+        )
+
+    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,
+        )
+        items = [
+            item
+            async for item in client.stream_chat(
+                messages=[ChatMessage(role="user", content="find docs")],
+                tools=[
+                    {
+                        "type": "function",
+                        "function": {"name": "mock_search", "parameters": {}},
+                    }
+                ],
+                params=AgentParams(model="provider-default"),
+            )
+        ]
+
+    provider_calls = [
+        item.event for item in items if item.kind == "provider_tool_call"
+    ]
+    assert provider_calls == [
+        ToolCallEvent(
+            id="call_1",
+            name="mock_search",
+            arguments={"query": "latency docs"},
+            raw_arguments='{"query":"latency docs"}',
+        )
+    ]
+
+
 @pytest.mark.asyncio
 async def test_openai_chat_client_uses_default_model_when_request_model_is_blank():
     captured_payloads: list[dict] = []