import asyncio import json from collections.abc import AsyncIterator import pytest from agent_lab.application.contracts import AgentParams, EventAgentParams from agent_lab.application.event_agent import EventAgent from agent_lab.application.tools import ToolDefinition, ToolExecutionContext, ToolRegistry from agent_lab.domain.events import ToolCallEvent from agent_lab.domain.messages import ChatMessage, StreamItem, TokenUsage class ToolCallingChatClient: def __init__(self, arguments: dict) -> None: self.arguments = arguments self.calls: list[dict] = [] async def stream_chat( self, messages: list[ChatMessage], tools: list[dict], params: AgentParams, tool_choice: dict | None = None, ) -> AsyncIterator[StreamItem]: self.calls.append( { "messages": list(messages), "tools": list(tools), "params": params, "tool_choice": tool_choice, } ) tool_name = tools[0]["function"]["name"] yield StreamItem.raw_response_chunk( { "choices": [ { "delta": { "tool_calls": [ { "index": 0, "function": {"name": tool_name}, } ] }, "finish_reason": None, } ] } ) yield StreamItem.provider_tool_call( ToolCallEvent( id="llm_call_1", name=tool_name, arguments=self.arguments, raw_arguments=json.dumps(self.arguments), ) ) class NoToolCallChatClient: def __init__(self) -> None: self.calls: list[dict] = [] async def stream_chat( self, messages: list[ChatMessage], tools: list[dict], params: AgentParams, tool_choice: dict | None = None, ) -> AsyncIterator[StreamItem]: self.calls.append( { "messages": list(messages), "tools": list(tools), "params": params, "tool_choice": tool_choice, } ) 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, tool_choice: dict | None = None, ) -> AsyncIterator[StreamItem]: yield self.item class WrongNameToolCallingChatClient(ToolCallingChatClient): async def stream_chat( self, messages: list[ChatMessage], tools: list[dict], params: AgentParams, tool_choice: dict | None = None, ) -> AsyncIterator[StreamItem]: self.calls.append({"messages": list(messages), "tools": list(tools)}) yield StreamItem.provider_tool_call( ToolCallEvent( id="wrong-call", name="wrong.tool", arguments=self.arguments, raw_arguments=json.dumps(self.arguments), ) ) class TrailingUsageChatClient: async def stream_chat( self, messages: list[ChatMessage], tools: list[dict], params: AgentParams, tool_choice: dict | None = None, ) -> AsyncIterator[StreamItem]: tool_name = tools[0]["function"]["name"] yield StreamItem.raw_response_chunk({"first": tool_name}) yield StreamItem.provider_tool_call( ToolCallEvent( id=f"provider-{tool_name}", name=tool_name, arguments={"query": tool_name}, raw_arguments=json.dumps({"query": tool_name}), ) ) yield StreamItem.usage_item( TokenUsage( prompt_tokens=3, completion_tokens=4, total_tokens=7, cached_tokens=1, ) ) class FailedResolutionUsageChatClient: async def stream_chat( self, messages: list[ChatMessage], tools: list[dict], params: AgentParams, tool_choice: dict | None = None, ) -> AsyncIterator[StreamItem]: yield StreamItem.message_delta("invalid fallback response") yield StreamItem.usage_item( TokenUsage(prompt_tokens=5, completion_tokens=2, total_tokens=7) ) @pytest.mark.asyncio async def test_event_agent_uses_deterministic_resolver_before_llm_fallback(): chat_client = ToolCallingChatClient({"message": "LLM generated handoff"}) agent = EventAgent( enabled_tools=["handoff_note"], chat_client=chat_client, params=AgentParams(model="event-model", temperature=0, max_tokens=80), ) event = ToolCallEvent( id="call_1", name="handoff_note", arguments={"message": "chat agent argument should be ignored"}, raw_arguments='{"message":"chat agent argument should be ignored"}', ) history = [ ChatMessage(role="user", content="debug this event flow"), ChatMessage(role="assistant", content="I need the event agent."), ] reply = await agent.handle(event, history=history) assert reply.role == "tool" assert reply.tool_call_id == "call_1" assert reply.name == "handoff_note" assert json.loads(reply.content) == { "tool": "handoff_note", "message": "I need the event agent.", } assert chat_client.calls == [] assert agent.raw_model_chunks([event]) == [ { "event_id": "call_1", "event_name": "handoff_note", "chunks": [], } ] @pytest.mark.asyncio async def test_event_agent_deterministic_mock_resolver_skips_llm(): chat_client = NoToolCallChatClient() agent = EventAgent( enabled_tools=["mock_search"], chat_client=chat_client, params=AgentParams(model="event-model", temperature=0, max_tokens=80), ) event = ToolCallEvent( id="call_1", name="mock_search", arguments={}, raw_arguments="{}", ) history = [ ChatMessage(role="user", content="Find latency docs"), ChatMessage(role="assistant", content="Need a search for latency docs"), ] reply = await agent.handle(event, history=history) payload = json.loads(reply.content) assert payload["tool"] == "mock_search" assert payload["query"] == "Need a search for latency docs" assert "event agent did not return arguments" not in reply.content assert chat_client.calls == [] @pytest.mark.asyncio async def test_event_agent_injects_clock_into_event_kernel_handler_metrics(): clock_values = iter([5.0, 5.030]) registry = ToolRegistry( [ ToolDefinition( name="timed.tool", description="Timed tool.", parameters={"type": "object"}, handler=lambda event: {"tool": event.name}, ) ] ) event = ToolCallEvent( id="timed-event", name="timed.tool", arguments={}, raw_arguments="{}", ) batch = await EventAgent( enabled_tools=[event.name], registry=registry, monotonic_clock=lambda: next(clock_values), ).handle_many([event], history=[]) assert batch.results[0].tool_latency_ms == 30 @pytest.mark.asyncio async def test_event_agent_calls_llm_once_for_incomplete_deterministic_arguments(): chat_client = ToolCallingChatClient({"message": "LLM generated handoff"}) registry = ToolRegistry( [ ToolDefinition( name="ambiguous_handoff", description="Resolve an ambiguous handoff.", parameters={ "type": "object", "properties": {"message": {"type": "string"}}, "required": ["message"], }, handler=lambda event: { "tool": event.name, "message": event.arguments["message"], }, argument_resolver=lambda event, context: {}, ) ] ) event = ToolCallEvent( id="call_1", name="ambiguous_handoff", arguments={}, raw_arguments="{}", ) reply = await EventAgent( enabled_tools=["ambiguous_handoff"], registry=registry, chat_client=chat_client, ).handle(event, history=[ChatMessage(role="user", content="ambiguous")]) assert json.loads(reply.content) == { "tool": "ambiguous_handoff", "message": "LLM generated handoff", } assert len(chat_client.calls) == 1 assert chat_client.calls[0]["tool_choice"] == { "type": "function", "function": {"name": "ambiguous_handoff"}, } @pytest.mark.asyncio async def test_event_agent_consumes_trailing_usage_and_records_provider_ttft(): clock_values = iter([10.0, 10.011, 10.050, 10.060, 10.080]) registry = ToolRegistry( [ ToolDefinition( name="fallback.search", description="Resolve with fallback.", parameters={ "type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"], }, handler=lambda event: {"query": event.arguments["query"]}, argument_resolver=lambda event, context: {}, ) ] ) event = ToolCallEvent( id="event-1", name="fallback.search", arguments={}, raw_arguments="{}", ) agent = EventAgent( enabled_tools=[event.name], registry=registry, chat_client=TrailingUsageChatClient(), monotonic_clock=lambda: next(clock_values), ) batch = await agent.handle_many([event], history=[]) assert batch.results[0].status.value == "success" assert agent.fallback_model_calls([event]) == [ { "event_id": "event-1", "event_name": "fallback.search", "usage": TokenUsage( prompt_tokens=3, completion_tokens=4, total_tokens=7, cached_tokens=1, ), "ttft_ms": 11, "elapsed_ms": 50, } ] @pytest.mark.asyncio async def test_event_agent_records_failed_fallback_resolution_attempt_usage(): clock_values = iter([2.0, 2.005, 2.020]) registry = ToolRegistry( [ ToolDefinition( name="fallback.search", description="Resolve with fallback.", parameters={ "type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"], }, handler=lambda event: {"query": event.arguments["query"]}, argument_resolver=lambda event, context: {}, ) ] ) event = ToolCallEvent( id="event-failed", name="fallback.search", arguments={}, raw_arguments="{}", ) agent = EventAgent( enabled_tools=[event.name], registry=registry, chat_client=FailedResolutionUsageChatClient(), monotonic_clock=lambda: next(clock_values), ) batch = await agent.handle_many([event], history=[]) assert batch.results[0].status.value == "invalid_arguments" assert agent.fallback_model_calls([event])[0] == { "event_id": "event-failed", "event_name": "fallback.search", "usage": TokenUsage(prompt_tokens=5, completion_tokens=2, total_tokens=7), "ttft_ms": 5, "elapsed_ms": 20, } @pytest.mark.asyncio async def test_event_agent_keeps_parallel_fallback_metrics_isolated_by_event(): registry = ToolRegistry( [ ToolDefinition( name=name, description="Resolve with fallback.", parameters={ "type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"], }, handler=lambda event: {"query": event.arguments["query"]}, argument_resolver=lambda event, context: {}, ) for name in ("fallback.first", "fallback.second") ] ) events = [ ToolCallEvent(id=f"event-{index}", name=name, arguments={}, raw_arguments="{}") for index, name in enumerate( ("fallback.first", "fallback.second"), start=1, ) ] agent = EventAgent( enabled_tools=[event.name for event in events], registry=registry, chat_client=TrailingUsageChatClient(), params=EventAgentParams(max_parallel_events=2), ) await agent.handle_many(events, history=[]) metrics = agent.fallback_model_calls(events) assert [(metric["event_id"], metric["event_name"]) for metric in metrics] == [ ("event-1", "fallback.first"), ("event-2", "fallback.second"), ] assert [metric["usage"].total_tokens for metric in metrics] == [7, 7] @pytest.mark.asyncio async def test_event_agent_rejects_fallback_provider_call_for_wrong_tool_name(): chat_client = WrongNameToolCallingChatClient({"message": "wrong"}) registry = ToolRegistry( [ ToolDefinition( name="expected.tool", description="Expected tool.", parameters={ "type": "object", "properties": {"message": {"type": "string"}}, "required": ["message"], }, handler=lambda event: {"tool": event.name}, argument_resolver=lambda event, context: {}, ) ] ) reply = await EventAgent( enabled_tools=["expected.tool"], registry=registry, chat_client=chat_client, ).handle( ToolCallEvent( id="call-1", name="expected.tool", arguments={}, raw_arguments="{}", ) ) assert json.loads(reply.content) == { "tool": "expected.tool", "error": "fallback returned tool wrong.tool for expected.tool", } @pytest.mark.asyncio async def test_event_agent_handle_many_starts_independent_handlers_concurrently(): started: list[str] = [] both_started = asyncio.Event() async def handler(event: ToolCallEvent) -> dict: started.append(event.name) if len(started) == 2: both_started.set() await asyncio.wait_for(both_started.wait(), timeout=0.2) return {"tool": event.name} registry = ToolRegistry( [ ToolDefinition( name=name, description=f"Handle {name}.", parameters={"type": "object"}, handler=handler, ) for name in ("event.first", "event.second") ] ) events = [ ToolCallEvent(id=f"call-{index}", name=name, arguments={}, raw_arguments="{}") for index, name in enumerate(("event.first", "event.second"), start=1) ] batch = await asyncio.wait_for( EventAgent( enabled_tools=[event.name for event in events], registry=registry, params=EventAgentParams(max_parallel_events=2), ).handle_many(events, history=[]), timeout=0.5, ) assert started == ["event.first", "event.second"] assert [result.event_name for result in batch.results] == [ "event.first", "event.second", ] assert [result.payload for result in batch.results] == [ {"tool": "event.first"}, {"tool": "event.second"}, ] @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="{}", ) registry = ToolRegistry( [ ToolDefinition( name="mock_search", description="Search mock external knowledge for the current turn.", parameters={ "type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"], }, handler=lambda resolved: { "tool": resolved.name, "query": resolved.arguments["query"], }, argument_resolver=lambda event, context: {}, ) ] ) agent = EventAgent( enabled_tools=["mock_search"], registry=registry, 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 == { "tool": "mock_search", "error": "missing required arguments: query", } @pytest.mark.asyncio async def test_event_agent_returns_registry_errors_for_disabled_and_unknown_tools(): registry = ToolRegistry( [ ToolDefinition( name="handoff_note", description="Send a note to the event agent.", parameters={"type": "object"}, handler=lambda event: {"tool": event.name, "message": "handled"}, ) ] ) disabled_reply = await EventAgent( enabled_tools=[], registry=registry, ).handle( ToolCallEvent( id="call_1", name="handoff_note", arguments={"message": "inspect this event"}, raw_arguments='{"message":"inspect this event"}', ) ) unknown_reply = await EventAgent( enabled_tools=["missing_tool"], registry=registry, ).handle( ToolCallEvent( id="call_2", name="missing_tool", arguments={}, raw_arguments="{}", ) ) assert json.loads(disabled_reply.content) == { "tool": "handoff_note", "error": "tool disabled", } assert json.loads(unknown_reply.content) == { "tool": "missing_tool", "error": "unknown tool", } @pytest.mark.asyncio async def test_event_agent_returns_structured_error_when_tool_handler_raises(): def fail_tool(event: ToolCallEvent) -> dict: raise RuntimeError("boom") registry = ToolRegistry( [ ToolDefinition( name="handoff_note", description="Send a note to the event agent.", parameters={"type": "object"}, handler=fail_tool, ) ] ) event = ToolCallEvent( id="call_1", name="handoff_note", arguments={"message": "inspect this event"}, raw_arguments='{"message":"inspect this event"}', ) reply = await EventAgent( enabled_tools=["handoff_note"], registry=registry, ).handle(event) assert reply.role == "tool" assert reply.tool_call_id == "call_1" assert json.loads(reply.content) == { "tool": "handoff_note", "error": "tool handler failed: boom", } @pytest.mark.asyncio async def test_event_agent_serializes_non_json_handler_payload_as_tool_error(): registry = ToolRegistry( [ ToolDefinition( name="bad_payload", description="Return an invalid payload.", parameters={"type": "object"}, handler=lambda event: {"invalid": object()}, ) ] ) reply = await EventAgent( enabled_tools=["bad_payload"], registry=registry, ).handle( ToolCallEvent( id="call-1", name="bad_payload", arguments={}, raw_arguments="{}", ) ) assert json.loads(reply.content) == { "tool": "bad_payload", "error": "event handler returned non-JSON payload", } @pytest.mark.asyncio async def test_event_agent_normalizes_handler_payload_snapshot_exceptions(): class ExplodingItemsDict(dict): def items(self): raise RuntimeError("payload items failed") registry = ToolRegistry( [ ToolDefinition( name="bad_payload", description="Return a payload that fails during snapshot.", parameters={"type": "object"}, handler=lambda event: ExplodingItemsDict(ok=True), ) ] ) reply = await EventAgent( enabled_tools=["bad_payload"], registry=registry, ).handle( ToolCallEvent( id="call-1", name="bad_payload", arguments={}, raw_arguments="{}", ) ) assert reply.role == "tool" assert reply.tool_call_id == "call-1" assert json.loads(reply.content) == { "tool": "bad_payload", "error": "event handler returned non-JSON payload", } @pytest.mark.asyncio async def test_event_agent_llm_receives_history_and_agent_config_context(): chat_client = ToolCallingChatClient( {"message": "Use strict tool parameters.", "thinking": "disabled"} ) registry = ToolRegistry( [ ToolDefinition( name="handoff_note", description="Send a note to the event agent.", parameters={ "type": "object", "properties": { "message": {"type": "string"}, "thinking": {"type": "string"}, }, "required": ["message", "thinking"], }, handler=lambda event: { "tool": event.name, "message": event.arguments["message"], "thinking": event.arguments["thinking"], }, ) ] ) event = ToolCallEvent( id="call_1", name="handoff_note", arguments={}, raw_arguments="{}", ) reply = await EventAgent( enabled_tools=["handoff_note"], registry=registry, chat_client=chat_client, params=AgentParams(model="event-model", temperature=0.4, max_tokens=120), ).handle( event, history=[ChatMessage(role="user", content="debug this")], system_prompt="Use strict tool parameters.", extra_body={"thinking": {"type": "disabled"}}, ) messages = chat_client.calls[0]["messages"] assert any(message.content == "debug this" for message in messages) assert any(message.content == "Use strict tool parameters." for message in messages) assert chat_client.calls[0]["params"].extra_body == {"thinking": {"type": "disabled"}} assert json.loads(reply.content) == { "tool": "handoff_note", "message": "Use strict tool parameters.", "thinking": "disabled", } @pytest.mark.asyncio async def test_event_agent_uses_configured_extra_body_when_override_is_omitted(): chat_client = ToolCallingChatClient({"message": "resolved"}) registry = ToolRegistry( [ ToolDefinition( name="handoff_note", description="Send a note to the event agent.", parameters={ "type": "object", "properties": {"message": {"type": "string"}}, "required": ["message"], }, handler=lambda event: {"tool": event.name}, argument_resolver=lambda event, context: {}, ) ] ) await EventAgent( enabled_tools=["handoff_note"], registry=registry, chat_client=chat_client, params=AgentParams(extra_body={"configured": True}), ).handle( ToolCallEvent( id="call-1", name="handoff_note", arguments={}, raw_arguments="{}", ) ) assert chat_client.calls[0]["params"].extra_body == {"configured": True} @pytest.mark.asyncio async def test_event_agent_preserves_explicit_empty_extra_body_override(): chat_client = ToolCallingChatClient({"message": "resolved"}) registry = ToolRegistry( [ ToolDefinition( name="handoff_note", description="Send a note to the event agent.", parameters={ "type": "object", "properties": {"message": {"type": "string"}}, "required": ["message"], }, handler=lambda event: {"tool": event.name}, argument_resolver=lambda event, context: {}, ) ] ) await EventAgent( enabled_tools=["handoff_note"], registry=registry, chat_client=chat_client, params=AgentParams(extra_body={"configured": True}), ).handle( ToolCallEvent( id="call-1", name="handoff_note", arguments={}, raw_arguments="{}", ), extra_body={}, ) assert chat_client.calls[0]["params"].extra_body == {} @pytest.mark.asyncio async def test_event_agent_projects_complete_tool_round_to_visible_history(): chat_client = ToolCallingChatClient({"message": "projected history"}) registry = ToolRegistry( [ ToolDefinition( name="handoff_note", description="Send a note to the event agent.", parameters={ "type": "object", "properties": {"message": {"type": "string"}}, "required": ["message"], }, handler=lambda resolved: { "tool": resolved.name, "message": resolved.arguments["message"], }, argument_resolver=lambda event, context: {}, ) ] ) 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"], registry=registry, 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)