diff --git a/src/app/endpoints/a2a.py b/src/app/endpoints/a2a.py index 99c53b8d6..c87da2dcd 100644 --- a/src/app/endpoints/a2a.py +++ b/src/app/endpoints/a2a.py @@ -27,14 +27,26 @@ TaskState, TaskStatus, TaskStatusUpdateEvent, - TextPart, +) +from a2a.types import ( + TextPart as A2ATextPart, ) from a2a.utils import new_agent_text_message, new_task from fastapi import APIRouter, Depends, HTTPException, Request, status -from llama_stack_api.openai_responses import ( - OpenAIResponseObjectStream, +from llama_stack_client import APIConnectionError, APIStatusError +from pydantic_ai import AgentRunResultEvent +from pydantic_ai.exceptions import AgentRunError +from pydantic_ai.messages import ( + AgentStreamEvent, + FunctionToolCallEvent, + NativeToolCallPart, + PartDeltaEvent, + PartEndEvent, + PartStartEvent, + TextPart, + TextPartDelta, ) -from llama_stack_client import APIConnectionError, AsyncLlamaStackClient +from pydantic_ai.run import AgentRunResult from starlette.responses import Response, StreamingResponse from a2a_storage import A2AContextStore, A2AStorageFactory @@ -47,18 +59,13 @@ from constants import MEDIA_TYPE_EVENT_STREAM from log import get_logger from models.api.requests import QueryRequest -from models.common.responses.types import ResponseInput from models.config import Action -from utils.conversation_compaction import ( - apply_compaction_blocking, - store_compacted_turn, -) +from utils.agents.query import map_agent_inference_error +from utils.conversation_compaction import apply_compaction_blocking from utils.mcp_headers import McpHeaders, mcp_headers_dependency -from utils.responses import ( - extract_text_from_response_item, - prepare_responses_params, -) -from utils.suid import normalize_conversation_id, to_llama_stack_conversation_id +from utils.pydantic_ai import build_agent +from utils.responses import prepare_responses_params +from utils.suid import normalize_conversation_id from version import __version__ logger = get_logger(__name__) @@ -103,23 +110,30 @@ async def _get_context_store() -> A2AContextStore: return _CONTEXT_STORE -def _convert_responses_content_to_a2a_parts(output: list[Any]) -> list[Part]: - """Convert Responses API output to A2A Parts. +def _build_a2a_parts_from_agent_result( + run_result: Optional[AgentRunResult[str]], + accumulated_text: list[str], +) -> list[Part]: + """Convert a pydantic-ai agent run result to A2A Parts. - Args: - output: List of Responses API output items + Prefers the authoritative text from the agent result, falling back to + the text accumulated from stream deltas — same pattern as + ``utils/agents/streaming.py:429``. + + Parameters: + run_result: Completed agent run result, or None if no result event arrived. + accumulated_text: Text parts accumulated during streaming. Returns: - List of A2A Part objects + List of A2A Part objects for the final artifact. """ - parts: list[Part] = [] - - for output_item in output: - text = extract_text_from_response_item(output_item) - if text: - parts.append(Part(root=TextPart(text=text))) - - return parts + if run_result is not None: + final_text = run_result.response.text or "".join(accumulated_text) + else: + final_text = "".join(accumulated_text) + if not final_text: + return [] + return [Part(root=A2ATextPart(text=final_text))] class TaskResultAggregator: @@ -208,6 +222,7 @@ def __init__( self.auth_token: str = auth_token self.mcp_headers: McpHeaders = mcp_headers or {} self.request_headers: Optional[Mapping[str, str]] = request_headers + self._run_result: Optional[AgentRunResult[str]] = None async def execute( self, @@ -356,22 +371,15 @@ async def _process_task_streaming( # pylint: disable=too-many-locals configuration.compaction, ) responses_params = compaction.params - # Stream response from LLM using the Responses API - stream = await client.responses.create(**responses_params.model_dump()) - except APIConnectionError as e: - error_message = ( - f"Unable to connect to Llama Stack backend service: {e!s}. " - "The service may be temporarily unavailable. Please try again later." - ) - logger.error( - "APIConnectionError in A2A request: %s", - str(e), - exc_info=True, - ) + + agent = build_agent(client, responses_params, configuration.skills) + except (AgentRunError, APIStatusError, APIConnectionError, RuntimeError) as e: + error_response = map_agent_inference_error(e, query_request.model or "") + logger.error("Error preparing A2A agent: %s", str(e), exc_info=True) await task_updater.update_status( TaskState.failed, message=new_agent_text_message( - error_message, + error_response.detail.response, context_id=context_id, task_id=task_id, ), @@ -412,19 +420,31 @@ async def _process_task_streaming( # pylint: disable=too-many-locals ) ) - # Process stream using generator and aggregator pattern. In compacted - # mode the conversation parameter is not sent, so the turn is stored - # explicitly once the response completes (see _convert_stream_to_events). - async for a2a_event in self._convert_stream_to_events( - stream, - task_id, - context_id, - conversation_id, - client, - compaction.original_input if compaction.compacted else None, - ): - aggregator.process_event(a2a_event) - await event_queue.enqueue_event(a2a_event) + # Run the pydantic-ai agent and convert stream events to A2A events. + prompt = user_input + try: + async for a2a_event in self._convert_stream_to_events( + agent, + prompt, + task_id, + context_id, + conversation_id=conversation_id, + ): + aggregator.process_event(a2a_event) + await event_queue.enqueue_event(a2a_event) + except (AgentRunError, APIStatusError, APIConnectionError, RuntimeError) as e: + error_response = map_agent_inference_error(e, responses_params.model) + logger.error("Error during A2A agent run: %s", str(e), exc_info=True) + await task_updater.update_status( + TaskState.failed, + message=new_agent_text_message( + error_response.detail.response, + context_id=context_id, + task_id=task_id, + ), + final=True, + ) + return # Publish the final task result event if aggregator.task_state == TaskState.working: @@ -441,22 +461,22 @@ async def _process_task_streaming( # pylint: disable=too-many-locals final=True, ) - async def _convert_stream_to_events( # pylint: disable=too-many-branches,too-many-locals,too-many-arguments,too-many-positional-arguments + async def _convert_stream_to_events( self, - stream: AsyncIterator[OpenAIResponseObjectStream], + agent: Any, + prompt: str, task_id: str, context_id: str, - conversation_id: Optional[str], - client: Optional[AsyncLlamaStackClient] = None, - compacted_original_input: Optional[ResponseInput] = None, + conversation_id: Optional[str] = None, ) -> AsyncIterator[Any]: - """Convert Responses API stream chunks to A2A events. + """Run an agent and convert stream events to A2A events. - Args: - stream: The Responses API response stream - task_id: The task ID for this execution - context_id: The context ID for this execution - conversation_id: The conversation ID for this A2A context + Parameters: + agent: Pydantic-ai Agent to execute. + prompt: User input text. + task_id: The task ID for this execution. + context_id: The context ID for this execution. + conversation_id: The conversation ID for this A2A context. Yields: A2A events (TaskStatusUpdateEvent or TaskArtifactUpdateEvent) @@ -466,61 +486,91 @@ async def _convert_stream_to_events( # pylint: disable=too-many-branches,too-ma artifact_id = str(uuid.uuid4()) text_parts: list[str] = [] + run_result: Optional[AgentRunResult[str]] = None + + async with agent.run_stream_events(prompt) as stream: + async for event in stream: + if isinstance(event, AgentRunResultEvent): + run_result = event.result + self._run_result = run_result + continue + a2a_event = self._dispatch_agent_event( + event, task_id, context_id, text_parts, artifact_id + ) + if a2a_event is not None: + yield a2a_event + + a2a_parts = _build_a2a_parts_from_agent_result(run_result, text_parts) + + yield TaskArtifactUpdateEvent( + task_id=task_id, + last_chunk=True, + context_id=context_id, + artifact=Artifact( + artifact_id=artifact_id, + parts=a2a_parts, + metadata={"conversation_id": str(conversation_id or "")}, + ), + ) - async for chunk in stream: - event_type = getattr(chunk, "type", None) - - # Skip response.created - conversation is already created - if event_type == "response.created": - continue + def _dispatch_agent_event( # pylint: disable=too-many-arguments,too-many-positional-arguments + self, + event: AgentStreamEvent | AgentRunResultEvent, + task_id: str, + context_id: str, + text_parts: list[str], + artifact_id: str, + ) -> Optional[TaskStatusUpdateEvent]: + """Map a single pydantic-ai stream event to an A2A status update. + + Parameters: + event: Pydantic-ai stream event. + task_id: The task ID for this execution. + context_id: The context ID for this execution. + text_parts: Mutable list accumulating text deltas. + artifact_id: Artifact ID (unused here, reserved for future use). + + Returns: + A2A TaskStatusUpdateEvent, or None if the event is not mapped. + """ + _ = artifact_id + + if isinstance(event, PartStartEvent): + if isinstance(event.part, TextPart): + text_parts.append(event.part.content) + return self._text_status_event(event.part.content, task_id, context_id) + + elif isinstance(event, PartDeltaEvent): + if isinstance(event.delta, TextPartDelta): + text_parts.append(event.delta.content_delta) + return self._text_status_event( + event.delta.content_delta, task_id, context_id + ) - # Text streaming - emit as working status with text delta - if event_type == "response.output_text.delta": - delta = getattr(chunk, "delta", "") - if delta: - text_parts.append(delta) - yield TaskStatusUpdateEvent( - task_id=task_id, - status=TaskStatus( - state=TaskState.working, - message=new_agent_text_message( - delta, - context_id=context_id, - task_id=task_id, - ), - timestamp=datetime.now(UTC).isoformat(), - ), + elif isinstance(event, FunctionToolCallEvent): + return TaskStatusUpdateEvent( + task_id=task_id, + status=TaskStatus( + state=TaskState.working, + message=new_agent_text_message( + f"Tool call: {event.part.tool_call_id} ({event.part.tool_name})", context_id=context_id, - final=False, - ) - - # Tool call events - elif event_type == "response.function_call_arguments.done": - item_id = getattr(chunk, "item_id", "") - yield TaskStatusUpdateEvent( - task_id=task_id, - status=TaskStatus( - state=TaskState.working, - message=new_agent_text_message( - f"Tool call: {item_id}", - context_id=context_id, - task_id=task_id, - ), - timestamp=datetime.now(UTC).isoformat(), + task_id=task_id, ), - context_id=context_id, - final=False, - ) + timestamp=datetime.now(UTC).isoformat(), + ), + context_id=context_id, + final=False, + ) - # MCP call completion - elif event_type == "response.mcp_call.arguments.done": - item_id = getattr(chunk, "item_id", "") - yield TaskStatusUpdateEvent( + elif isinstance(event, PartEndEvent): + if isinstance(event.part, NativeToolCallPart): + return TaskStatusUpdateEvent( task_id=task_id, status=TaskStatus( state=TaskState.working, message=new_agent_text_message( - f"MCP call: {item_id}", + f"Tool call: {event.part.tool_call_id} ({event.part.tool_name})", context_id=context_id, task_id=task_id, ), @@ -530,45 +580,39 @@ async def _convert_stream_to_events( # pylint: disable=too-many-branches,too-ma final=False, ) - # Response completed - emit final artifact - elif event_type == "response.completed": - response_obj = getattr(chunk, "response", None) - final_text = "".join(text_parts) - - if response_obj: - output = getattr(response_obj, "output", []) - # In compacted mode the conversation parameter was not sent, - # so persist this turn ourselves to keep the recent-turn - # buffer and audit history intact for the next request. - if ( - compacted_original_input is not None - and client is not None - and conversation_id - ): - await store_compacted_turn( - client, - to_llama_stack_conversation_id(conversation_id), - compacted_original_input, - output, - ) - a2a_parts = _convert_responses_content_to_a2a_parts(output) - if not a2a_parts and final_text: - a2a_parts = [Part(root=TextPart(text=final_text))] - else: - a2a_parts = ( - [Part(root=TextPart(text=final_text))] if final_text else [] - ) - - yield TaskArtifactUpdateEvent( - task_id=task_id, - last_chunk=True, + # AgentRunResultEvent and other events are not mapped to status updates. + return None + + def _text_status_event( + self, + text: str, + task_id: str, + context_id: str, + ) -> TaskStatusUpdateEvent: + """Build a working-status A2A event carrying a text delta. + + Parameters: + text: The text delta content. + task_id: The task ID. + context_id: The context ID. + + Returns: + TaskStatusUpdateEvent with the text delta. + """ + return TaskStatusUpdateEvent( + task_id=task_id, + status=TaskStatus( + state=TaskState.working, + message=new_agent_text_message( + text, context_id=context_id, - artifact=Artifact( - artifact_id=artifact_id, - parts=a2a_parts, - metadata={"conversation_id": str(conversation_id or "")}, - ), - ) + task_id=task_id, + ), + timestamp=datetime.now(UTC).isoformat(), + ), + context_id=context_id, + final=False, + ) async def cancel( self, diff --git a/tests/unit/app/endpoints/test_a2a.py b/tests/unit/app/endpoints/test_a2a.py index e35b82b8e..c317e0bba 100644 --- a/tests/unit/app/endpoints/test_a2a.py +++ b/tests/unit/app/endpoints/test_a2a.py @@ -2,6 +2,7 @@ # pylint: disable=redefined-outer-name # pylint: disable=protected-access +# pylint: disable=too-many-lines from typing import Any @@ -22,12 +23,24 @@ from a2a.utils import new_agent_text_message from fastapi import HTTPException, Request from llama_stack_client import APIConnectionError +from pydantic_ai import AgentRunResultEvent +from pydantic_ai.exceptions import AgentRunError +from pydantic_ai.messages import ( + FunctionToolCallEvent, + NativeToolCallPart, + PartDeltaEvent, + PartEndEvent, + PartStartEvent, + TextPartDelta, + ToolCallPart, +) +from pydantic_ai.messages import TextPart as PydanticTextPart from pytest_mock import MockerFixture from app.endpoints.a2a import ( A2AAgentExecutor, TaskResultAggregator, - _convert_responses_content_to_a2a_parts, + _build_a2a_parts_from_agent_result, _get_context_store, _get_task_store, a2a_health_check, @@ -139,69 +152,69 @@ def setup_minimal_configuration_fixture(mocker: MockerFixture) -> AppConfig: # ----------------------------- -# Tests for _convert_responses_content_to_a2a_parts +# Tests for _build_a2a_parts_from_agent_result # ----------------------------- -class TestConvertResponsesContentToA2AParts: - """Tests for the content conversion function.""" - - def test_convert_empty_output(self, mocker: MockerFixture) -> None: - """Test converting empty output returns empty list.""" - mocker.patch( - "app.endpoints.a2a.extract_text_from_response_item", - return_value=None, - ) - result = _convert_responses_content_to_a2a_parts([]) - assert not result +class TestBuildA2APartsFromAgentResult: + """Tests for the agent result to A2A parts conversion function.""" + + def test_result_with_response_text(self, mocker: MockerFixture) -> None: + """Test conversion when result.response.text is available.""" + mock_result = mocker.MagicMock() + mock_result.response.text = "Hello, world!" + result = _build_a2a_parts_from_agent_result(mock_result, []) + assert len(result) == 1 + text = result[0].root.text # pyright: ignore[reportAttributeAccessIssue] + assert text == "Hello, world!" - def test_convert_single_output_item(self, mocker: MockerFixture) -> None: - """Test converting single output item with text.""" - mocker.patch( - "app.endpoints.a2a.extract_text_from_response_item", - return_value="Hello, world!", - ) - mock_output_item = mocker.MagicMock() - result = _convert_responses_content_to_a2a_parts([mock_output_item]) + def test_result_falls_back_to_accumulated_text(self, mocker: MockerFixture) -> None: + """Test conversion uses accumulated text when response.text is empty.""" + mock_result = mocker.MagicMock() + mock_result.response.text = "" + result = _build_a2a_parts_from_agent_result(mock_result, ["Hello, ", "world!"]) assert len(result) == 1 - assert result[0].root is not None text = result[0].root.text # pyright: ignore[reportAttributeAccessIssue] assert text == "Hello, world!" - def test_convert_multiple_output_items(self, mocker: MockerFixture) -> None: - """Test converting multiple output items.""" - extract_mock = mocker.patch( - "app.endpoints.a2a.extract_text_from_response_item", + def test_result_prefers_response_text_over_accumulated( + self, mocker: MockerFixture + ) -> None: + """Test that response.text takes priority over accumulated text.""" + mock_result = mocker.MagicMock() + mock_result.response.text = "Final response" + result = _build_a2a_parts_from_agent_result( + mock_result, ["accumulated", " text"] ) - extract_mock.side_effect = ["First", "Second"] - - mock_item1 = mocker.MagicMock() - mock_item2 = mocker.MagicMock() + assert len(result) == 1 + text = result[0].root.text # pyright: ignore[reportAttributeAccessIssue] + assert text == "Final response" - result = _convert_responses_content_to_a2a_parts([mock_item1, mock_item2]) - assert len(result) == 2 - assert result[0].root is not None - assert result[1].root is not None - text1 = result[0].root.text # pyright: ignore[reportAttributeAccessIssue] - text2 = result[1].root.text # pyright: ignore[reportAttributeAccessIssue] - assert text1 == "First" - assert text2 == "Second" + def test_empty_result_returns_empty_list(self, mocker: MockerFixture) -> None: + """Test that empty result returns empty list.""" + mock_result = mocker.MagicMock() + mock_result.response.text = "" + result = _build_a2a_parts_from_agent_result(mock_result, []) + assert not result - def test_convert_output_items_with_none_text(self, mocker: MockerFixture) -> None: - """Test that output items with no text are filtered out.""" - extract_mock = mocker.patch( - "app.endpoints.a2a.extract_text_from_response_item", - ) - extract_mock.side_effect = ["Valid text", None, "Another valid"] + def test_none_response_text_uses_accumulated(self, mocker: MockerFixture) -> None: + """Test conversion when response.text is None.""" + mock_result = mocker.MagicMock() + mock_result.response.text = None + result = _build_a2a_parts_from_agent_result(mock_result, ["fallback text"]) + assert len(result) == 1 + text = result[0].root.text # pyright: ignore[reportAttributeAccessIssue] + assert text == "fallback text" - mock_items = [mocker.MagicMock(), mocker.MagicMock(), mocker.MagicMock()] + def test_none_run_result_uses_accumulated(self) -> None: + """Test conversion when run_result is None.""" + result = _build_a2a_parts_from_agent_result(None, ["accumulated", " text"]) + assert len(result) == 1 + text = result[0].root.text # pyright: ignore[reportAttributeAccessIssue] + assert text == "accumulated text" - result = _convert_responses_content_to_a2a_parts(mock_items) - assert len(result) == 2 - assert result[0].root is not None - assert result[1].root is not None - text1 = result[0].root.text # pyright: ignore[reportAttributeAccessIssue] - text2 = result[1].root.text # pyright: ignore[reportAttributeAccessIssue] - assert text1 == "Valid text" - assert text2 == "Another valid" + def test_none_run_result_empty_text_returns_empty(self) -> None: + """Test that None run_result with no accumulated text returns empty.""" + result = _build_a2a_parts_from_agent_result(None, []) + assert not result # ----------------------------- @@ -725,12 +738,12 @@ async def test_process_task_streaming_handles_api_connection_error_on_models_lis assert "Unable to connect to Llama Stack" in str(exc_info.value.detail) @pytest.mark.asyncio - async def test_process_task_streaming_handles_api_connection_error_on_retrieve_response( + async def test_process_task_streaming_handles_api_connection_error( # pylint: disable=too-many-locals self, mocker: MockerFixture, setup_configuration: AppConfig, # pylint: disable=unused-argument ) -> None: - """Test _process_task_streaming handles APIConnectionError from retrieve_response().""" + """Test _process_task_streaming handles APIConnectionError during agent run.""" executor = A2AAgentExecutor(auth_token="test-token") # Mock the context with valid input @@ -760,49 +773,138 @@ async def test_process_task_streaming_handles_api_connection_error_on_retrieve_r "app.endpoints.a2a._get_context_store", return_value=mock_context_store ) - # Mock the client to succeed on models.list() + # Mock the client mock_client = mocker.AsyncMock() - mock_models = [mocker.MagicMock()] # Return a list of models + mock_models = [mocker.MagicMock()] mock_client.models.list = mocker.AsyncMock(return_value=mock_models) + mocker.patch( + "app.endpoints.a2a.AsyncLlamaStackClientHolder" + ).return_value.get_client.return_value = mock_client - # Mock responses.create to raise APIConnectionError + # Mock prepare_responses_params + mock_responses_params = mocker.Mock() + mock_responses_params.model = "test-model" + mock_responses_params.conversation = "conv_x" + mocker.patch( + "app.endpoints.a2a.prepare_responses_params", + new=mocker.AsyncMock(return_value=mock_responses_params), + ) + + # Mock compaction + compaction_result = mocker.Mock() + compaction_result.params = mock_responses_params + compaction_result.summarized = False + compaction_result.compacted = False + compaction_result.original_input = None + mocker.patch( + "app.endpoints.a2a.apply_compaction_blocking", + new=mocker.AsyncMock(return_value=compaction_result), + ) + + # Mock build_agent to return an agent whose run_stream_events raises mock_request = httpx.Request("POST", "http://test-llama-stack/responses") - mock_client.responses.create = mocker.AsyncMock( + mock_agent = mocker.MagicMock() + mock_stream_ctx = mocker.AsyncMock() + mock_stream_ctx.__aenter__ = mocker.AsyncMock( side_effect=APIConnectionError( - message="Connection timeout during streaming", request=mock_request + message="Connection timeout during streaming", + request=mock_request, ) ) + mock_agent.run_stream_events.return_value = mock_stream_ctx + mocker.patch( + "app.endpoints.a2a.build_agent", + return_value=mock_agent, + ) + + await executor._process_task_streaming( + context, task_updater, context.task_id, context.context_id + ) + + # Verify failure status was sent + update_calls = task_updater.update_status.call_args_list + failure_calls = [c for c in update_calls if c[0][0] == TaskState.failed] + assert len(failure_calls) == 1 + assert failure_calls[0][1]["final"] is True + + @pytest.mark.asyncio + async def test_process_task_streaming_handles_agent_run_error( # pylint: disable=too-many-locals + self, + mocker: MockerFixture, + setup_configuration: AppConfig, # pylint: disable=unused-argument + ) -> None: + """Test _process_task_streaming handles AgentRunError during agent run.""" + executor = A2AAgentExecutor(auth_token="test-token") + + mock_message = mocker.MagicMock() + mock_message.role = "user" + mock_message.parts = [Part(root=TextPart(text="Hello"))] + mock_message.metadata = {} + + context = mocker.MagicMock(spec=RequestContext) + context.task_id = "task-123" + context.context_id = "ctx-456" + context.message = mock_message + context.get_user_input.return_value = "Hello" + event_queue = mocker.AsyncMock(spec=EventQueue) + + task_updater = mocker.MagicMock() + task_updater.update_status = mocker.AsyncMock() + task_updater.event_queue = event_queue + + mock_context_store = mocker.AsyncMock() + mock_context_store.get.return_value = None + mocker.patch( + "app.endpoints.a2a._get_context_store", return_value=mock_context_store + ) + + mock_client = mocker.AsyncMock() + mock_client.models.list = mocker.AsyncMock(return_value=[mocker.MagicMock()]) mocker.patch( "app.endpoints.a2a.AsyncLlamaStackClientHolder" ).return_value.get_client.return_value = mock_client - # Mock prepare_responses_params to return valid params mock_responses_params = mocker.Mock() - mock_responses_params.model_dump.return_value = { - "input": "Hello", - "model": "test-model", - } + mock_responses_params.model = "test-model" + mock_responses_params.conversation = "conv_x" mocker.patch( "app.endpoints.a2a.prepare_responses_params", new=mocker.AsyncMock(return_value=mock_responses_params), ) + compaction_result = mocker.Mock() + compaction_result.params = mock_responses_params + compaction_result.summarized = False + compaction_result.compacted = False + compaction_result.original_input = None + mocker.patch( + "app.endpoints.a2a.apply_compaction_blocking", + new=mocker.AsyncMock(return_value=compaction_result), + ) + + mock_agent = mocker.MagicMock() + mock_stream_ctx = mocker.AsyncMock() + mock_stream_ctx.__aenter__ = mocker.AsyncMock( + side_effect=AgentRunError("Agent execution failed") + ) + mock_agent.run_stream_events.return_value = mock_stream_ctx + mocker.patch( + "app.endpoints.a2a.build_agent", + return_value=mock_agent, + ) + await executor._process_task_streaming( context, task_updater, context.task_id, context.context_id ) - # Verify failure status was sent - task_updater.update_status.assert_called_once() - call_args = task_updater.update_status.call_args - assert call_args[0][0] == TaskState.failed - assert call_args[1]["final"] is True - # Verify error message contains helpful info - error_message = call_args[1]["message"] - assert "Unable to connect to Llama Stack backend service" in str(error_message) + update_calls = task_updater.update_status.call_args_list + failure_calls = [c for c in update_calls if c[0][0] == TaskState.failed] + assert len(failure_calls) == 1 + assert failure_calls[0][1]["final"] is True @pytest.mark.asyncio - async def test_process_task_streaming_applies_compaction( + async def test_process_task_streaming_applies_compaction( # pylint: disable=too-many-locals self, mocker: MockerFixture, setup_configuration: AppConfig, # pylint: disable=unused-argument @@ -832,14 +934,8 @@ async def test_process_task_streaming_applies_compaction( "app.endpoints.a2a._get_context_store", return_value=mock_context_store ) - async def _empty_stream() -> Any: - """Yield no stream events.""" - return - yield # pragma: no cover (makes this an async generator) - mock_client = mocker.AsyncMock() mock_client.models.list = mocker.AsyncMock(return_value=[mocker.MagicMock()]) - mock_client.responses.create = mocker.AsyncMock(return_value=_empty_stream()) mocker.patch( "app.endpoints.a2a.AsyncLlamaStackClientHolder" ).return_value.get_client.return_value = mock_client @@ -847,7 +943,7 @@ async def _empty_stream() -> Any: mock_params = mocker.Mock() mock_params.model = "test-model" mock_params.conversation = "conv_x" - mock_params.model_dump.return_value = {"input": "Hello", "model": "test-model"} + mock_params.skills = None mocker.patch( "app.endpoints.a2a.prepare_responses_params", new=mocker.AsyncMock(return_value=mock_params), @@ -856,18 +952,73 @@ async def _empty_stream() -> Any: compaction_result = mocker.Mock() compaction_result.params = mock_params compaction_result.summarized = False + compaction_result.compacted = False compaction_result.original_input = None apply = mocker.patch( "app.endpoints.a2a.apply_compaction_blocking", new=mocker.AsyncMock(return_value=compaction_result), ) + # Mock build_agent to return a mock agent that yields an AgentRunResultEvent + mock_run_result = mocker.MagicMock() + mock_run_result.response.text = "Response" + result_event = mocker.MagicMock(spec=AgentRunResultEvent) + result_event.result = mock_run_result + + async def _event_stream() -> Any: + yield result_event + + mock_stream_ctx = mocker.AsyncMock() + mock_stream_ctx.__aenter__ = mocker.AsyncMock(return_value=_event_stream()) + mock_stream_ctx.__aexit__ = mocker.AsyncMock(return_value=False) + mock_agent = mocker.MagicMock() + mock_agent.run_stream_events.return_value = mock_stream_ctx + mocker.patch( + "app.endpoints.a2a.build_agent", + return_value=mock_agent, + ) + await executor._process_task_streaming( # pylint: disable=protected-access context, task_updater, context.task_id, context.context_id ) apply.assert_awaited_once() + @pytest.mark.asyncio + async def test_convert_stream_artifact_carries_conversation_id( + self, mocker: MockerFixture + ) -> None: + """Test that the final artifact metadata includes the conversation_id.""" + executor = A2AAgentExecutor(auth_token="test-token") + + mock_run_result = mocker.MagicMock() + mock_run_result.response.text = "answer" + result_event = mocker.MagicMock(spec=AgentRunResultEvent) + result_event.result = mock_run_result + + async def _event_stream() -> Any: + yield result_event + + mock_stream_ctx = mocker.AsyncMock() + mock_stream_ctx.__aenter__ = mocker.AsyncMock(return_value=_event_stream()) + mock_stream_ctx.__aexit__ = mocker.AsyncMock(return_value=False) + mock_agent = mocker.MagicMock() + mock_agent.run_stream_events.return_value = mock_stream_ctx + + conversation_id = "abc123" + events = [ + e + async for e in executor._convert_stream_to_events( + mock_agent, "hi", "task-1", "ctx-1", conversation_id=conversation_id + ) + ] + + artifact_events = [e for e in events if isinstance(e, TaskArtifactUpdateEvent)] + assert len(artifact_events) == 1 + metadata = artifact_events[0].artifact.metadata + assert metadata is not None + assert metadata["conversation_id"] == conversation_id + @pytest.mark.asyncio async def test_cancel_raises_not_implemented(self, mocker: MockerFixture) -> None: """Test that cancel raises NotImplementedError.""" @@ -924,6 +1075,112 @@ async def test_get_task_store_returns_store( assert store is not None +# ----------------------------- +# Tests for _dispatch_agent_event +# ----------------------------- +class TestDispatchAgentEvent: + """Tests for mapping pydantic-ai stream events to A2A events.""" + + def _make_executor(self) -> A2AAgentExecutor: + """Create an executor for testing.""" + return A2AAgentExecutor(auth_token="test-token") + + def test_text_part_start_emits_status_update(self) -> None: + """Test that PartStartEvent with TextPart emits a working status.""" + executor = self._make_executor() + text_parts: list[str] = [] + event = PartStartEvent( + index=0, + part=PydanticTextPart(content="Hello"), + event_kind="part_start", + ) + result = executor._dispatch_agent_event( + event, "task-1", "ctx-1", text_parts, "art-1" + ) + assert result is not None + assert isinstance(result, TaskStatusUpdateEvent) + assert result.status.state == TaskState.working + assert text_parts == ["Hello"] + + def test_text_part_delta_emits_status_update(self) -> None: + """Test that PartDeltaEvent with TextPartDelta emits a working status.""" + executor = self._make_executor() + text_parts: list[str] = ["Hello"] + event = PartDeltaEvent( + index=0, + delta=TextPartDelta(content_delta=", world!"), + event_kind="part_delta", + ) + result = executor._dispatch_agent_event( + event, "task-1", "ctx-1", text_parts, "art-1" + ) + assert result is not None + assert isinstance(result, TaskStatusUpdateEvent) + assert result.status.state == TaskState.working + assert text_parts == ["Hello", ", world!"] + + def test_function_tool_call_emits_status_update( + self, mocker: MockerFixture + ) -> None: + """Test that FunctionToolCallEvent emits a tool call status.""" + executor = self._make_executor() + text_parts: list[str] = [] + tool_call_part = mocker.MagicMock(spec=ToolCallPart) + tool_call_part.tool_name = "search_docs" + tool_call_part.tool_call_id = "call_xyz789" + event = FunctionToolCallEvent( + part=tool_call_part, + event_kind="function_tool_call", + ) + result = executor._dispatch_agent_event( + event, "task-1", "ctx-1", text_parts, "art-1" + ) + assert result is not None + assert isinstance(result, TaskStatusUpdateEvent) + assert "Tool call: call_xyz789 (search_docs)" in str(result.status.message) + + def test_native_tool_call_end_emits_status_update( + self, mocker: MockerFixture + ) -> None: + """Test that PartEndEvent with NativeToolCallPart emits an MCP call status.""" + executor = self._make_executor() + text_parts: list[str] = [] + native_part = mocker.MagicMock(spec=NativeToolCallPart) + native_part.tool_name = "file_search" + native_part.tool_call_id = "call_abc123" + event = PartEndEvent( + index=0, + part=native_part, + event_kind="part_end", + ) + result = executor._dispatch_agent_event( + event, "task-1", "ctx-1", text_parts, "art-1" + ) + assert result is not None + assert isinstance(result, TaskStatusUpdateEvent) + assert "Tool call: call_abc123 (file_search)" in str(result.status.message) + + def test_agent_run_result_returns_none(self, mocker: MockerFixture) -> None: + """Test that AgentRunResultEvent is not mapped to a status update.""" + executor = self._make_executor() + text_parts: list[str] = [] + event = mocker.MagicMock(spec=AgentRunResultEvent) + result = executor._dispatch_agent_event( + event, "task-1", "ctx-1", text_parts, "art-1" + ) + assert result is None + + def test_unknown_event_returns_none(self, mocker: MockerFixture) -> None: + """Test that unknown events return None.""" + executor = self._make_executor() + text_parts: list[str] = [] + event = mocker.MagicMock() + result = executor._dispatch_agent_event( + event, "task-1", "ctx-1", text_parts, "art-1" + ) + assert result is None + + # ----------------------------- # Integration-style tests for endpoint handlers # -----------------------------