Python: (ag-ui): Add Workflow Support, Harden Streaming Semantics, and add Dynamic Handoff Demo (#3911)

* fix Workflow.as_agent() streaming regression in ag-ui

* Address PR feedback

* workflows wip

* wip

* wip

* Workflow AG-UI demo

* Fixes for handoff workflow demo

* Fixes to workflows support in AG-UI

* Fixes

* Add headers to some demo files

* Fix comment

* Fixes for store

* Make _input_schema lazy-loaded

* fix mypy

* revert session change to handoff only for now

---------

Co-authored-by: Eduard van Valkenburg <eavanvalkenburg@users.noreply.github.com>
This commit is contained in:
Evan Mattson
2026-02-23 20:59:56 +09:00
committed by GitHub
Unverified
parent b1c7c7c844
commit d8b9409e96
60 changed files with 8349 additions and 512 deletions
@@ -1,5 +1,6 @@
# Copyright (c) Microsoft. All rights reserved.
import re
from collections.abc import AsyncIterable, Awaitable, Mapping, Sequence
from typing import Any, cast
from unittest.mock import AsyncMock, MagicMock
@@ -15,18 +16,21 @@ from agent_framework import (
ResponseStream,
WorkflowEvent,
resolve_agent_id,
tool,
)
from agent_framework._clients import BaseChatClient
from agent_framework._middleware import ChatMiddlewareLayer, FunctionInvocationContext, MiddlewareTermination
from agent_framework._tools import FunctionInvocationLayer, FunctionTool, tool
from agent_framework._tools import FunctionInvocationLayer, FunctionTool
from agent_framework.orchestrations import HandoffAgentUserRequest, HandoffBuilder
from agent_framework_orchestrations._handoff import (
HANDOFF_FUNCTION_RESULT_KEY,
HandoffAgentExecutor,
HandoffConfiguration,
_AutoHandoffMiddleware, # pyright: ignore[reportPrivateUsage]
get_handoff_tool_name,
)
from agent_framework_orchestrations._orchestrator_helpers import clean_conversation_for_handoff
class MockChatClient(ChatMiddlewareLayer[Any], FunctionInvocationLayer[Any], BaseChatClient[Any]):
@@ -130,6 +134,67 @@ class MockHandoffAgent(Agent):
super().__init__(client=MockChatClient(name=name, handoff_to=handoff_to), name=name, id=name)
class ContextAwareRefundClient(ChatMiddlewareLayer[Any], FunctionInvocationLayer[Any], BaseChatClient[Any]):
"""Mock client that expects prior user context to remain available on resume."""
def __init__(self) -> None:
ChatMiddlewareLayer.__init__(self)
FunctionInvocationLayer.__init__(self)
BaseChatClient.__init__(self)
self._call_index = 0
def _inner_get_response(
self,
*,
messages: Sequence[Message],
stream: bool,
options: Mapping[str, Any],
**kwargs: Any,
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
del kwargs
del options
contents = self._next_contents(messages)
if stream:
return self._build_streaming_response(contents)
async def _get() -> ChatResponse:
return ChatResponse(messages=[Message(role="assistant", contents=contents)], response_id="context-aware")
return _get()
def _build_streaming_response(self, contents: list[Content]) -> ResponseStream[ChatResponseUpdate, ChatResponse]:
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=contents, role="assistant", finish_reason="stop")
def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse:
return ChatResponse.from_updates(updates)
return ResponseStream(_stream(), finalizer=_finalize)
def _next_contents(self, messages: Sequence[Message]) -> list[Content]:
user_text = " ".join(message.text or "" for message in messages if message.role == "user")
order_match = re.search(r"\b(\d{4,12})\b", user_text)
order_id = order_match.group(1) if order_match else None
asks_refund = any(token in user_text.lower() for token in ("broken", "damaged", "refund", "cracked"))
if self._call_index == 0:
reply = "Refund Agent: Please share your order number."
elif self._call_index == 1:
if order_id:
reply = f"Refund Agent: Thanks, I found order {order_id}. Why do you need the refund?"
else:
reply = "Refund Agent: I still need your order number."
else:
if order_id and asks_refund:
reply = f"Refund Agent: Got it for order {order_id}. I can proceed with your refund."
else:
reply = "Refund Agent: I still need your order number."
self._call_index += 1
return [Content.from_text(text=reply)]
async def _drain(stream: AsyncIterable[WorkflowEvent]) -> list[WorkflowEvent]:
return [event async for event in stream]
@@ -168,6 +233,567 @@ async def test_handoff():
assert request.source_executor_id == escalation.name
def _latest_request_info_event(events: list[WorkflowEvent]) -> WorkflowEvent[Any]:
request_events = [event for event in events if event.type == "request_info"]
assert request_events
request_event = request_events[-1]
assert isinstance(request_event.data, HandoffAgentUserRequest)
return request_event
def _request_text(event: WorkflowEvent[Any]) -> str:
request_payload = cast(HandoffAgentUserRequest, event.data)
messages = request_payload.agent_response.messages
assert messages
return messages[-1].text or ""
async def test_resume_keeps_prior_user_context_for_same_agent() -> None:
"""Ensure same-agent request_info resumes retain prior turn context."""
refund_agent = Agent(
id="refund_agent",
name="refund_agent",
client=ContextAwareRefundClient(),
)
workflow = (
HandoffBuilder(participants=[refund_agent], termination_condition=lambda _: False)
.with_start_agent(refund_agent)
.build()
)
first_events = await _drain(workflow.run("My order arrived damaged.", stream=True))
first_request = _latest_request_info_event(first_events)
assert "order number" in _request_text(first_request).lower()
second_events = await _drain(
workflow.run(
stream=True,
responses={first_request.request_id: [Message(role="user", text="Order 2939393")]},
)
)
second_request = _latest_request_info_event(second_events)
second_text = _request_text(second_request).lower()
assert "order 2939393" in second_text
assert "order number" not in second_text
third_events = await _drain(
workflow.run(
stream=True,
responses={second_request.request_id: [Message(role="user", text="It arrived broken and unusable.")]},
)
)
third_request = _latest_request_info_event(third_events)
third_text = _request_text(third_request).lower()
assert "order 2939393" in third_text
assert "order number" not in third_text
async def test_tool_approval_responses_are_not_replayed_from_history() -> None:
"""Ensure persisted history does not re-execute previously approved tool calls."""
execution_count = 0
@tool(name="submit_refund_counted", approval_mode="always_require")
def submit_refund_counted() -> str:
nonlocal execution_count
execution_count += 1
return "ok"
class ApprovalReplayClient(ChatMiddlewareLayer[Any], FunctionInvocationLayer[Any], BaseChatClient[Any]):
def __init__(self) -> None:
ChatMiddlewareLayer.__init__(self)
FunctionInvocationLayer.__init__(self)
BaseChatClient.__init__(self)
self._call_index = 0
def _inner_get_response(
self,
*,
messages: Sequence[Message],
stream: bool,
options: Mapping[str, Any],
**kwargs: Any,
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
del messages
del options
del kwargs
if self._call_index == 0:
contents = [
Content.from_function_call(
call_id="refund-call-1",
name="submit_refund_counted",
arguments={},
)
]
elif self._call_index == 1:
contents = [Content.from_text(text="Refund approved and recorded.")]
else:
contents = [Content.from_text(text="No additional tool work needed.")]
self._call_index += 1
if stream:
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=contents, role="assistant", finish_reason="stop")
return ResponseStream(_stream(), finalizer=lambda updates: ChatResponse.from_updates(updates))
async def _get() -> ChatResponse:
return ChatResponse(
messages=[Message(role="assistant", contents=contents)],
response_id="approval-replay",
)
return _get()
agent = Agent(
id="refund_agent",
name="refund_agent",
client=ApprovalReplayClient(),
tools=[submit_refund_counted],
)
workflow = (
HandoffBuilder(participants=[agent], termination_condition=lambda _: False).with_start_agent(agent).build()
)
first_events = await _drain(workflow.run("start", stream=True))
first_requests = [event for event in first_events if event.type == "request_info"]
assert first_requests
first_request = first_requests[-1]
assert isinstance(first_request.data, Content)
approval_response = first_request.data.to_function_approval_response(approved=True)
second_events = await _drain(workflow.run(stream=True, responses={first_request.request_id: approval_response}))
second_request = _latest_request_info_event(second_events)
await _drain(
workflow.run(
stream=True,
responses={second_request.request_id: [Message(role="user", text="Thanks, what's next?")]},
)
)
assert execution_count == 1
async def test_handoff_resume_preserves_approval_function_call_for_stateless_runs() -> None:
"""Approval resume turns must replay matching function calls when store=False."""
@tool(name="submit_refund", approval_mode="always_require")
def submit_refund() -> str:
return "ok"
class StrictStatelessApprovalClient(ChatMiddlewareLayer[Any], FunctionInvocationLayer[Any], BaseChatClient[Any]):
def __init__(self) -> None:
ChatMiddlewareLayer.__init__(self)
FunctionInvocationLayer.__init__(self)
BaseChatClient.__init__(self)
self._call_index = 0
self.resume_validated = False
def _inner_get_response(
self,
*,
messages: Sequence[Message],
stream: bool,
options: Mapping[str, Any],
**kwargs: Any,
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
del options
del kwargs
if self._call_index == 0:
contents = [
Content.from_function_call(
call_id="refund-call-1",
name="submit_refund",
arguments={},
)
]
else:
function_call_ids = {
content.call_id
for message in messages
for content in message.contents
if content.type == "function_call" and content.call_id
}
function_result_ids = {
content.call_id
for message in messages
for content in message.contents
if content.type == "function_result" and content.call_id
}
missing_call_ids = sorted(function_result_ids - function_call_ids)
if missing_call_ids:
raise AssertionError(
f"No tool call found for function call output with call_id {missing_call_ids[0]}."
)
self.resume_validated = True
contents = [Content.from_text(text="Refund submitted.")]
self._call_index += 1
if stream:
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=contents, role="assistant", finish_reason="stop")
return ResponseStream(_stream(), finalizer=lambda updates: ChatResponse.from_updates(updates))
async def _get() -> ChatResponse:
return ChatResponse(
messages=[Message(role="assistant", contents=contents)],
response_id="strict-stateless",
)
return _get()
client = StrictStatelessApprovalClient()
agent = Agent(
id="refund_agent",
name="refund_agent",
client=client,
tools=[submit_refund],
)
workflow = (
HandoffBuilder(participants=[agent], termination_condition=lambda _: False).with_start_agent(agent).build()
)
first_events = await _drain(workflow.run("start", stream=True))
approval_requests = [
event for event in first_events if event.type == "request_info" and isinstance(event.data, Content)
]
assert approval_requests
first_request = approval_requests[0]
approval_response = first_request.data.to_function_approval_response(True)
await _drain(workflow.run(stream=True, responses={first_request.request_id: approval_response}))
assert client.resume_validated is True
async def test_handoff_replay_serializes_handoff_function_results() -> None:
"""Returning to the same agent must not replay dict tool outputs."""
class ReplaySafeHandoffClient(ChatMiddlewareLayer[Any], FunctionInvocationLayer[Any], BaseChatClient[Any]):
def __init__(self, name: str, handoff_sequence: list[str | None]) -> None:
ChatMiddlewareLayer.__init__(self)
FunctionInvocationLayer.__init__(self)
BaseChatClient.__init__(self)
self._name = name
self._handoff_sequence = handoff_sequence
self._call_index = 0
def _inner_get_response(
self,
*,
messages: Sequence[Message],
stream: bool,
options: Mapping[str, Any],
**kwargs: Any,
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
del options
del kwargs
for message in messages:
for content in message.contents:
if content.type == "function_result" and isinstance(content.result, dict):
raise AssertionError("Expected replayed function_result payloads to be JSON strings.")
handoff_to = (
self._handoff_sequence[self._call_index] if self._call_index < len(self._handoff_sequence) else None
)
call_id = f"{self._name}-handoff-{self._call_index}" if handoff_to else None
contents = _build_reply_contents(self._name, handoff_to, call_id)
self._call_index += 1
if stream:
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=contents, role="assistant", finish_reason="stop")
return ResponseStream(_stream(), finalizer=lambda updates: ChatResponse.from_updates(updates))
async def _get() -> ChatResponse:
return ChatResponse(messages=[Message(role="assistant", contents=contents)], response_id="replay-safe")
return _get()
triage = Agent(
id="triage",
name="triage",
client=ReplaySafeHandoffClient(name="triage", handoff_sequence=["specialist", None]),
)
specialist = Agent(
id="specialist",
name="specialist",
client=ReplaySafeHandoffClient(name="specialist", handoff_sequence=["triage"]),
)
workflow = (
HandoffBuilder(participants=[triage, specialist], termination_condition=lambda _: False)
.with_start_agent(triage)
.build()
)
events = await _drain(workflow.run("start", stream=True))
requests = [event for event in events if event.type == "request_info"]
assert requests
assert requests[-1].source_executor_id == triage.name
async def test_handoff_resume_preserves_approved_tool_output_for_stateless_runs() -> None:
"""Approved calls must keep function_call/function_result pairs for later replays."""
submit_call_id = "call_submit_refund_approved"
@tool(name="submit_refund", approval_mode="always_require")
def submit_refund() -> str:
return "submitted"
class RefundReplayClient(ChatMiddlewareLayer[Any], FunctionInvocationLayer[Any], BaseChatClient[Any]):
def __init__(self) -> None:
ChatMiddlewareLayer.__init__(self)
FunctionInvocationLayer.__init__(self)
BaseChatClient.__init__(self)
self._call_index = 0
self.resume_validated = False
def _inner_get_response(
self,
*,
messages: Sequence[Message],
stream: bool,
options: Mapping[str, Any],
**kwargs: Any,
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
del options
del kwargs
if self._call_index == 0:
contents = [Content.from_function_call(call_id=submit_call_id, name="submit_refund", arguments={})]
elif self._call_index == 1:
contents = _build_reply_contents("refund_agent", "order_agent", "refund-order-handoff-1")
else:
function_call_ids = {
content.call_id
for message in messages
for content in message.contents
if content.type == "function_call" and content.call_id
}
function_result_ids = {
content.call_id
for message in messages
for content in message.contents
if content.type == "function_result" and content.call_id
}
if submit_call_id in function_call_ids and submit_call_id not in function_result_ids:
raise AssertionError(f"No tool output found for function call {submit_call_id}.")
self.resume_validated = True
contents = [Content.from_text(text="Refund agent resumed.")]
self._call_index += 1
if stream:
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=contents, role="assistant", finish_reason="stop")
return ResponseStream(_stream(), finalizer=lambda updates: ChatResponse.from_updates(updates))
async def _get() -> ChatResponse:
return ChatResponse(
messages=[Message(role="assistant", contents=contents)],
response_id="refund-replay",
)
return _get()
class OrderReplayClient(ChatMiddlewareLayer[Any], FunctionInvocationLayer[Any], BaseChatClient[Any]):
def __init__(self) -> None:
ChatMiddlewareLayer.__init__(self)
FunctionInvocationLayer.__init__(self)
BaseChatClient.__init__(self)
self._call_index = 0
def _inner_get_response(
self,
*,
messages: Sequence[Message],
stream: bool,
options: Mapping[str, Any],
**kwargs: Any,
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
del messages
del options
del kwargs
if self._call_index == 0:
contents = [Content.from_text(text="Would you like a replacement or a refund?")]
else:
contents = _build_reply_contents("order_agent", "refund_agent", "order-refund-handoff-1")
self._call_index += 1
if stream:
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=contents, role="assistant", finish_reason="stop")
return ResponseStream(_stream(), finalizer=lambda updates: ChatResponse.from_updates(updates))
async def _get() -> ChatResponse:
return ChatResponse(messages=[Message(role="assistant", contents=contents)], response_id="order-replay")
return _get()
refund_client = RefundReplayClient()
refund_agent = Agent(
id="refund_agent",
name="refund_agent",
client=refund_client,
tools=[submit_refund],
)
order_agent = Agent(
id="order_agent",
name="order_agent",
client=OrderReplayClient(),
)
workflow = (
HandoffBuilder(participants=[refund_agent, order_agent], termination_condition=lambda _: False)
.with_start_agent(refund_agent)
.build()
)
first_events = await _drain(workflow.run("start", stream=True))
approval_requests = [
event for event in first_events if event.type == "request_info" and isinstance(event.data, Content)
]
assert approval_requests
approval_request = approval_requests[-1]
approval_response = approval_request.data.to_function_approval_response(True)
second_events = await _drain(workflow.run(stream=True, responses={approval_request.request_id: approval_response}))
order_request = _latest_request_info_event(second_events)
assert order_request.source_executor_id == order_agent.name
await _drain(
workflow.run(
stream=True,
responses={order_request.request_id: [Message(role="user", text="Please continue with refund.")]},
)
)
assert refund_client.resume_validated is True
def test_handoff_clone_disables_provider_side_storage() -> None:
"""Handoff executors should force store=False to avoid stale provider call state."""
triage = MockHandoffAgent(name="triage")
workflow = HandoffBuilder(participants=[triage]).with_start_agent(triage).build()
executor = workflow.executors[resolve_agent_id(triage)]
assert isinstance(executor, HandoffAgentExecutor)
assert executor._agent.default_options.get("store") is False
async def test_handoff_clears_stale_service_session_id_before_run() -> None:
"""Stale service session IDs must be dropped before each handoff agent turn."""
triage = MockHandoffAgent(name="triage", handoff_to="specialist")
specialist = MockHandoffAgent(name="specialist")
workflow = HandoffBuilder(participants=[triage, specialist]).with_start_agent(triage).build()
triage_executor = workflow.executors[resolve_agent_id(triage)]
assert isinstance(triage_executor, HandoffAgentExecutor)
triage_executor._session.service_session_id = "resp_stale_value"
await _drain(workflow.run("My order is damaged", stream=True))
assert triage_executor._session.service_session_id is None
def test_clean_conversation_for_handoff_keeps_text_only_history() -> None:
"""Tool-control messages must be excluded from persisted handoff history."""
function_call = Content.from_function_call(
call_id="handoff-call-1",
name="handoff_to_refund_agent",
arguments={"context": "route to refund"},
)
approval_response = Content.from_function_approval_response(
approved=True,
id="approval-1",
function_call=function_call,
)
conversation = [
Message(role="user", text="My order arrived damaged."),
Message(
role="assistant",
contents=[
function_call,
Content.from_text(text="Triage Agent: Routing you to Refund."),
],
),
Message(role="tool", contents=[Content.from_function_result(call_id="handoff-call-1", result="ok")]),
Message(role="user", contents=[approval_response]),
Message(
role="assistant",
contents=[Content.from_function_call(call_id="handoff-call-2", name="handoff_to_order_agent")],
),
]
cleaned = clean_conversation_for_handoff(conversation)
assert [message.role for message in cleaned] == ["user", "assistant"]
assert [message.text for message in cleaned] == [
"My order arrived damaged.",
"Triage Agent: Routing you to Refund.",
]
def test_persist_missing_approved_function_results_handles_runtime_and_fallback_outputs() -> None:
"""Persisted history should retain approved call outputs across runtime shapes."""
agent = MockHandoffAgent(name="triage")
executor = HandoffAgentExecutor(agent, handoffs=[])
call_with_runtime_result = "call-runtime-result"
call_with_approval_only = "call-approval-only"
executor._full_conversation = [
Message(
role="assistant",
contents=[
Content.from_function_call(call_id=call_with_runtime_result, name="submit_refund", arguments={}),
Content.from_function_call(call_id=call_with_approval_only, name="submit_refund", arguments={}),
],
)
]
approval_response = Content.from_function_approval_response(
approved=True,
id=call_with_approval_only,
function_call=Content.from_function_call(call_id=call_with_approval_only, name="submit_refund", arguments={}),
)
runtime_messages = [
Message(
role="tool",
contents=[Content.from_function_result(call_id=call_with_runtime_result, result='{"submitted":true}')],
),
Message(role="user", contents=[approval_response]),
]
executor._persist_missing_approved_function_results(runtime_tool_messages=runtime_messages, response_messages=[])
persisted_tool_messages = [message for message in executor._full_conversation if message.role == "tool"]
assert persisted_tool_messages
persisted_results = [
content
for message in persisted_tool_messages
for content in message.contents
if content.type == "function_result" and content.call_id
]
result_by_call_id = {content.call_id: content.result for content in persisted_results}
assert result_by_call_id[call_with_runtime_result] == '{"submitted":true}'
assert result_by_call_id[call_with_approval_only] == '{"status":"approved"}'
async def test_autonomous_mode_yields_output_without_user_request():
"""Ensure autonomous interaction mode yields output without requesting user input."""
triage = MockHandoffAgent(name="triage", handoff_to="specialist")
@@ -278,6 +904,61 @@ async def test_handoff_async_termination_condition() -> None:
assert termination_call_count > 0
async def test_handoff_terminates_without_request_info_when_latest_response_meets_condition() -> None:
"""Termination triggered by the latest assistant response should not emit request_info."""
class FinalizingClient(ChatMiddlewareLayer[Any], FunctionInvocationLayer[Any], BaseChatClient[Any]):
def __init__(self) -> None:
ChatMiddlewareLayer.__init__(self)
FunctionInvocationLayer.__init__(self)
BaseChatClient.__init__(self)
def _inner_get_response(
self,
*,
messages: Sequence[Message],
stream: bool,
options: Mapping[str, Any],
**kwargs: Any,
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
del messages, options, kwargs
contents = [Content.from_text(text="Replacement request submitted. Case complete.")]
if stream:
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=contents, role="assistant", finish_reason="stop")
return ResponseStream(_stream(), finalizer=lambda updates: ChatResponse.from_updates(updates))
async def _get() -> ChatResponse:
return ChatResponse(messages=[Message(role="assistant", contents=contents)], response_id="finalizing")
return _get()
agent = Agent(id="order_agent", name="order_agent", client=FinalizingClient())
workflow = (
HandoffBuilder(
participants=[agent],
termination_condition=lambda conv: any(
message.role == "assistant" and "case complete." in (message.text or "").lower() for message in conv
),
)
.with_start_agent(agent)
.build()
)
events = await _drain(workflow.run("ship replacement", stream=True))
requests = [event for event in events if event.type == "request_info"]
assert not requests
outputs = [event for event in events if event.type == "output"]
assert outputs
conversation_outputs = [event for event in outputs if isinstance(event.data, list)]
assert len(conversation_outputs) == 1
async def test_tool_choice_preserved_from_agent_config():
"""Verify that agent-level tool_choice configuration is preserved and not overridden."""
# Create a mock chat client that records the tool_choice used