diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py index 2bc10c172e..324933cbeb 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py @@ -649,6 +649,34 @@ def _latest_assistant_contents(messages: list[Message]) -> list[Content] | None: return None +def _unemitted_exposed_function_results(response: AgentResponse, flow: FlowState) -> list[Content]: + """Return finalized function results that complete calls exposed during this run.""" + exposed_call_ids = set(flow.tool_calls_by_id) + emitted_call_ids = { + str(result["toolCallId"]) for result in flow.tool_results if isinstance(result.get("toolCallId"), str) + } + latest_assistant_result_ids = { + str(content.call_id) + for content in _latest_assistant_contents(list(response.messages or [])) or [] + if content.type == "function_result" and content.call_id + } + results: list[Content] = [] + for message in response.messages or []: + for content in message.contents or []: + call_id = content.call_id + if ( + content.type != "function_result" + or not call_id + or call_id not in exposed_call_ids + or call_id in emitted_call_ids + or call_id in latest_assistant_result_ids + ): + continue + results.append(content) + emitted_call_ids.add(call_id) + return results + + def _text_from_contents(contents: list[Content]) -> str | None: """Return normalized assistant text from a content list when present.""" text_parts: list[str] = [] @@ -795,7 +823,12 @@ async def run_workflow_stream( if pending_interrupt_ids else _resume_to_workflow_responses(resume_payload) ) - responses = _merge_workflow_response_sources(resume_responses, _extract_responses_from_messages(messages)) + message_responses = { + request_id: value + for request_id, value in _extract_responses_from_messages(messages).items() + if request_id in pending_interrupt_ids + } + responses = _merge_workflow_response_sources(resume_responses, message_responses) responses, response_error = _coerce_responses_for_pending_requests_strict(responses, pending_before_run) if response_error is not None: yield RunStartedEvent(run_id=run_id, thread_id=thread_id) @@ -973,13 +1006,16 @@ def _drain_open_message() -> list[TextMessageEndEvent]: for item in output_payload: yield item continue + if isinstance(output_payload, AgentResponse): + for result in _unemitted_exposed_function_results(output_payload, flow): + for out_event in _emit_content(result, flow, predictive_handler=None, skip_text=False): + yield out_event contents = _workflow_payload_to_contents(output_payload) if contents: output_text = _text_from_contents(contents) - if output_text and output_text == last_assistant_text: - continue + skip_text = bool(output_text and output_text == last_assistant_text) for content in contents: - for out_event in _emit_content(content, flow, predictive_handler=None, skip_text=False): + for out_event in _emit_content(content, flow, predictive_handler=None, skip_text=skip_text): yield out_event if flow.message_id and flow.accumulated_text: last_assistant_text = flow.accumulated_text.strip() or last_assistant_text diff --git a/python/packages/ag-ui/tests/ag_ui/test_handoff_replay.py b/python/packages/ag-ui/tests/ag_ui/test_handoff_replay.py new file mode 100644 index 0000000000..bd3b473ffd --- /dev/null +++ b/python/packages/ag-ui/tests/ag_ui/test_handoff_replay.py @@ -0,0 +1,173 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Acceptance tests for replayable AG-UI workflow handoffs.""" + +from __future__ import annotations + +import json +from collections.abc import AsyncIterator, Sequence +from typing import Any, cast + +import httpx +from agent_framework import Agent, ChatOptions, ChatResponseUpdate, Content, Message, Workflow +from agent_framework.orchestrations import HandoffBuilder +from conftest import StreamingChatClientStub # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports] +from fastapi import FastAPI + +from agent_framework_ag_ui import AgentFrameworkWorkflow, AGUIChatClient, add_agent_framework_fastapi_endpoint +from agent_framework_ag_ui._workflow_run import run_workflow_stream + + +def _as_handoff_agent(agent: Any) -> Agent: + return cast(Agent, agent) + + +def _unmatched_function_call_ids(messages: Sequence[Message]) -> set[str]: + pending: set[str] = set() + for message in messages: + for content in message.contents: + if content.type == "function_call" and content.call_id: + pending.add(content.call_id) + elif content.type == "function_result" and content.call_id: + pending.discard(content.call_id) + return pending + + +def _build_handoff_workflow( + triage_requests: list[list[Message]] | None = None, +) -> Workflow: + triage_invocation = 0 + specialist_invocation = 0 + + async def triage_stream( + messages: Sequence[Message], options: Any, **kwargs: Any + ) -> AsyncIterator[ChatResponseUpdate]: + nonlocal triage_invocation + del options, kwargs + unmatched_call_ids = _unmatched_function_call_ids(messages) + if unmatched_call_ids: + raise ValueError(f"No tool output found for function call {sorted(unmatched_call_ids)[0]}") + if triage_requests is not None: + triage_requests.append(list(messages)) + + if triage_invocation == 0: + yield ChatResponseUpdate( + contents=[ + Content.from_function_call( + call_id="handoff-call", + name="handoff_to_specialist", + arguments={}, + ) + ], + role=None, + ) + else: + yield ChatResponseUpdate(contents=[Content.from_text("Triage follow-up response")], role="assistant") + triage_invocation += 1 + + async def specialist_stream( + messages: Sequence[Message], options: Any, **kwargs: Any + ) -> AsyncIterator[ChatResponseUpdate]: + nonlocal specialist_invocation + del options, kwargs + unmatched_call_ids = _unmatched_function_call_ids(messages) + if unmatched_call_ids: + raise ValueError(f"No tool output found for function call {sorted(unmatched_call_ids)[0]}") + specialist_invocation += 1 + yield ChatResponseUpdate( + contents=[Content.from_text(f"Specialist response {specialist_invocation}")], + role="assistant", + ) + + triage = Agent( + id="triage", + name="triage", + client=StreamingChatClientStub(triage_stream), + require_per_service_call_history_persistence=True, + ) + specialist = Agent( + id="specialist", + name="specialist", + client=StreamingChatClientStub(specialist_stream), + require_per_service_call_history_persistence=True, + ) + return ( + HandoffBuilder( + participants=[_as_handoff_agent(triage), _as_handoff_agent(specialist)], + termination_condition=lambda conversation: bool( + conversation and conversation[-1].role == "assistant" and conversation[-1].text + ), + ) + .with_start_agent(_as_handoff_agent(triage)) + .build() + ) + + +async def test_real_handoff_runner_emits_replayable_tool_result() -> None: + """A real HandoffBuilder run should expose its synthetic result under the original call ID.""" + events = [ + event + async for event in run_workflow_stream( + {"messages": [{"role": "user", "content": "Route this request"}]}, + _build_handoff_workflow(), + ) + ] + + handoff_events = [ + event + for event in events + if event.type in {"TOOL_CALL_START", "TOOL_CALL_END", "TOOL_CALL_RESULT"} + and getattr(event, "tool_call_id", None) == "handoff-call" + ] + assert [event.type for event in handoff_events] == [ + "TOOL_CALL_START", + "TOOL_CALL_END", + "TOOL_CALL_RESULT", + ] + assert json.loads(handoff_events[-1].content) == {"handoff_to": "specialist"} # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + assert events.index(handoff_events[-1]) < next(i for i, event in enumerate(events) if event.type == "RUN_FINISHED") + + +async def test_outer_agent_replays_balanced_handoff_on_second_turn() -> None: + """An outer Agent should replay a balanced handoff transcript to the workflow on turn two.""" + triage_requests: list[list[Message]] = [] + workflow_runner = AgentFrameworkWorkflow( + workflow_factory=lambda _thread_id: _build_handoff_workflow(triage_requests) + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint(app, workflow_runner, path="/workflow", keepalive_seconds=None) + + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as http_client: + agui_client = AGUIChatClient( + endpoint="http://testserver/workflow", + http_client=http_client, + ) + outer_options: ChatOptions = {"metadata": {"thread_id": "handoff-thread"}} + outer_agent = Agent( + name="outer", + client=cast(Any, agui_client), + default_options=outer_options, + ) + session = outer_agent.create_session() + + first_response = await outer_agent.run("Route this request", session=session) + second_response = await outer_agent.run("Continue on the same thread", session=session) + + first_contents = [content for message in first_response.messages for content in message.contents] + assert [content.call_id for content in first_contents if content.type == "function_call"] == ["handoff-call"] + assert [content.call_id for content in first_contents if content.type == "function_result"] == ["handoff-call"] + first_call_index = next( + index + for index, message in enumerate(first_response.messages) + if any(content.type == "function_call" and content.call_id == "handoff-call" for content in message.contents) + ) + first_result_index = next( + index + for index, message in enumerate(first_response.messages) + if any(content.type == "function_result" and content.call_id == "handoff-call" for content in message.contents) + ) + assert first_call_index < first_result_index + assert second_response.text == "Triage follow-up response" + assert len(triage_requests) == 2 + assert _unmatched_function_call_ids(triage_requests[1]) == set() diff --git a/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py b/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py index 5a0426258c..8af9a33be2 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py @@ -460,8 +460,8 @@ async def scripted_stream(messages: Any, options: Any, **kwargs: Any) -> AsyncIt ] -async def test_workflow_stream_does_not_repeat_tool_call_from_final_response() -> None: - """A final response containing streamed history should not repeat its tool call.""" +async def test_workflow_stream_reconciles_result_from_final_response() -> None: + """A final response should complete, but not repeat, its exposed tool call.""" function_call = Content.from_function_call( call_id="weather-call", name="get_weather", @@ -496,6 +496,158 @@ async def participant(message: Any, ctx: WorkflowContext[Any, AgentResponse | Ag tool_call_starts = [event for event in events if event.type == "TOOL_CALL_START"] assert [event.tool_call_id for event in tool_call_starts] == ["weather-call"] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + tool_results = [event for event in events if event.type == "TOOL_CALL_RESULT"] + assert [event.tool_call_id for event in tool_results] == ["weather-call"] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + assert tool_results[0].content == "Sunny in Seattle" # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + + +async def test_workflow_stream_does_not_repeat_finalized_result_from_latest_assistant_message() -> None: + """A finalized assistant result should complete its exposed call exactly once.""" + function_call = Content.from_function_call( + call_id="weather-call", + name="get_weather", + arguments={"city": "Seattle"}, + ) + function_result = Content.from_function_result(call_id="weather-call", result="Sunny in Seattle") + + @executor(id="participant") + async def participant(message: Any, ctx: WorkflowContext[Any, AgentResponse | AgentResponseUpdate]) -> None: + del message + await ctx.yield_output(AgentResponseUpdate(contents=[function_call], role=None)) + await ctx.yield_output( + AgentResponse( + messages=[ + Message(role="assistant", contents=[function_call]), + Message( + role="assistant", + contents=[function_result, Content.from_text("The weather is sunny.")], + ), + ] + ) + ) + + workflow = WorkflowBuilder(start_executor=participant, output_from="all").build() + events = [ + event + async for event in run_workflow_stream( + {"messages": [{"role": "user", "content": "What is the weather in Seattle?"}]}, + workflow, + ) + ] + + tool_ends = [event for event in events if event.type == "TOOL_CALL_END"] + assert [event.tool_call_id for event in tool_ends] == ["weather-call"] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + tool_results = [event for event in events if event.type == "TOOL_CALL_RESULT"] + assert [event.tool_call_id for event in tool_results] == ["weather-call"] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + assert tool_results[0].content == "Sunny in Seattle" # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + assert [event.delta for event in events if event.type == "TEXT_MESSAGE_CONTENT"] == [ # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + "The weather is sunny." + ] + + +async def test_workflow_stream_preserves_finalized_result_when_text_was_already_streamed() -> None: + """Duplicate finalized text must not suppress its exposed tool result.""" + function_call = Content.from_function_call( + call_id="weather-call", + name="get_weather", + arguments={"city": "Seattle"}, + ) + function_result = Content.from_function_result(call_id="weather-call", result="Sunny in Seattle") + response_text = "The weather is sunny." + + @executor(id="participant") + async def participant(message: Any, ctx: WorkflowContext[Any, AgentResponse | AgentResponseUpdate]) -> None: + del message + await ctx.yield_output(AgentResponseUpdate(contents=[function_call], role=None)) + await ctx.yield_output(AgentResponseUpdate(contents=[Content.from_text(response_text)], role="assistant")) + await ctx.yield_output( + AgentResponse( + messages=[ + Message(role="assistant", contents=[function_call]), + Message(role="assistant", contents=[function_result, Content.from_text(response_text)]), + ] + ) + ) + + workflow = WorkflowBuilder(start_executor=participant, output_from="all").build() + events = [ + event + async for event in run_workflow_stream( + {"messages": [{"role": "user", "content": "What is the weather in Seattle?"}]}, + workflow, + ) + ] + + tool_ends = [event for event in events if event.type == "TOOL_CALL_END"] + assert [event.tool_call_id for event in tool_ends] == ["weather-call"] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + tool_results = [event for event in events if event.type == "TOOL_CALL_RESULT"] + assert [event.tool_call_id for event in tool_results] == ["weather-call"] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + assert tool_results[0].content == "Sunny in Seattle" # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + assert [event.delta for event in events if event.type == "TEXT_MESSAGE_CONTENT"] == [response_text] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + + +async def test_workflow_stream_keeps_unexposed_finalized_results_private() -> None: + """Finalized results must not publish calls that were never exposed during the run.""" + private_call = Content.from_function_call( + call_id="private-call", + name="internal_tool", + arguments={}, + ) + + @executor(id="participant") + async def participant(message: Any, ctx: WorkflowContext[Any, AgentResponse]) -> None: + del message + await ctx.yield_output( + AgentResponse( + messages=[ + Message(role="assistant", contents=[private_call]), + Message( + role="tool", + contents=[Content.from_function_result(call_id="private-call", result="private result")], + ), + Message(role="assistant", contents=[Content.from_text("Public response")]), + ] + ) + ) + + workflow = WorkflowBuilder(start_executor=participant, output_from="all").build() + events = [event async for event in run_workflow_stream({"messages": [{"role": "user", "content": "go"}]}, workflow)] + + assert not [event for event in events if event.type.startswith("TOOL_CALL")] + assert [event.delta for event in events if event.type == "TEXT_MESSAGE_CONTENT"] == ["Public response"] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + + +async def test_workflow_stream_matches_multiple_finalized_results_by_call_id() -> None: + """Each exposed call should receive only the finalized result carrying its call ID.""" + first_call = Content.from_function_call(call_id="call-a", name="first_tool", arguments={"value": "a"}) + second_call = Content.from_function_call(call_id="call-b", name="second_tool", arguments={"value": "b"}) + + @executor(id="participant") + async def participant(message: Any, ctx: WorkflowContext[Any, AgentResponse | AgentResponseUpdate]) -> None: + del message + await ctx.yield_output(AgentResponseUpdate(contents=[first_call], role=None)) + await ctx.yield_output(AgentResponseUpdate(contents=[second_call], role=None)) + await ctx.yield_output( + AgentResponse( + messages=[ + Message(role="assistant", contents=[first_call, second_call]), + Message(role="tool", contents=[Content.from_function_result(call_id="call-b", result="result-b")]), + Message(role="tool", contents=[Content.from_function_result(call_id="call-a", result="result-a")]), + ] + ) + ) + + workflow = WorkflowBuilder(start_executor=participant, output_from="all").build() + events = [event async for event in run_workflow_stream({"messages": [{"role": "user", "content": "go"}]}, workflow)] + + tool_results = [event for event in events if event.type == "TOOL_CALL_RESULT"] + assert { + event.tool_call_id: event.content # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + for event in tool_results + } == { + "call-a": "result-a", + "call-b": "result-b", + } async def test_workflow_run_passthroughs_ag_ui_base_events():