agent/agent-framework
1
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
* 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
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/ag-ui/AGENTS.md
+11 -1
View File
@@ -5,17 +5,27 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard.
## Main Classes
- **`AgentFrameworkAgent`** - Wraps agents for AG-UI compatibility
- **`AgentFrameworkWorkflow`** - Wraps native `Workflow` objects, or accepts `workflow_factory(thread_id)` for thread-scoped workflow instances without subclassing
- **`AGUIChatClient`** - Chat client that speaks AG-UI protocol
- **`AGUIHttpService`** - HTTP service for AG-UI endpoints
- **`AGUIEventConverter`** - Converts between Agent Framework and AG-UI events
- **`add_agent_framework_fastapi_endpoint()`** - Add AG-UI endpoint to FastAPI app
- **`add_agent_framework_fastapi_endpoint()`** - Add AG-UI endpoint to FastAPI app (`SupportsAgentRun` or `Workflow`)
## Types
- **`AGUIRequest`** / **`AGUIChatOptions`** - Request types
- **`availableInterrupts` / `resume`** - Optional interrupt configuration and continuation payloads
- **`AgentState`** / **`RunMetadata`** - State management types
- **`PredictStateConfig`** - Configuration for state prediction
## Protocol Notes
- Outbound custom events are emitted as AG-UI `CUSTOM`.
- Usage metadata from `Content(type="usage")` is surfaced as `CUSTOM` events with `name="usage"`.
- Inbound custom event aliases are accepted: `CUSTOM`, `CUSTOM_EVENT`, and `custom_event`.
- Multimodal user inputs support both legacy (`text`, `binary`) and draft-style (`image`, `audio`, `video`, `document`) shapes.
- `RUN_FINISHED.interrupt` can be emitted for pause/request-info flows, and interruption metadata is preserved in converters.
## Usage
```python
python/packages/ag-ui/README.md
+46
View File
@@ -36,6 +36,44 @@ add_agent_framework_fastapi_endpoint(app, agent, "/")
# Run with: uvicorn main:app --reload
```
### Server (Host a Workflow)
```python
from fastapi import FastAPI
from agent_framework import WorkflowBuilder, WorkflowContext, executor
from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
@executor(id="start")
async def start(message: str, ctx: WorkflowContext) -> None:
await ctx.yield_output(f"Workflow received: {message}")
workflow = WorkflowBuilder(start_executor=start).build()
app = FastAPI()
add_agent_framework_fastapi_endpoint(app, workflow, "/")
```
### Server (Thread-Scoped WorkflowBuilder)
Use `workflow_factory` when your workflow keeps runtime state (for example pending `request_info` interrupts) and must be isolated per AG-UI thread:
```python
from fastapi import FastAPI
from agent_framework import Workflow, WorkflowBuilder
from agent_framework.ag_ui import AgentFrameworkWorkflow, add_agent_framework_fastapi_endpoint
def build_workflow_for_thread(thread_id: str) -> Workflow:
# Build a fresh workflow instance for each thread id.
return WorkflowBuilder(start_executor=...).build()
app = FastAPI()
thread_scoped_workflow = AgentFrameworkWorkflow(
workflow_factory=build_workflow_for_thread,
name="my_workflow",
)
add_agent_framework_fastapi_endpoint(app, thread_scoped_workflow, "/")
```
### Client (Connect to an AG-UI Server)
```python
@@ -59,6 +97,7 @@ The `AGUIChatClient` supports:
- Hybrid tool execution (client-side + server-side tools)
- Automatic thread management for conversation continuity
- Integration with `Agent` for client-side history management
- Interrupt metadata passthrough (`availableInterrupts` and `resume`)
## Documentation
@@ -81,6 +120,13 @@ This integration supports all 7 AG-UI features:
6. **Shared State**: Bidirectional state sync between client and server
7. **Predictive State Updates**: Stream tool arguments as optimistic state updates during execution
Additional compatibility and draft support:
- Native `Workflow` endpoint registration via `add_agent_framework_fastapi_endpoint(...)`
- Workflow-to-AG-UI event mapping (run/step/activity/tool/custom events)
- Custom event compatibility for inbound `CUSTOM`, `CUSTOM_EVENT`, and `custom_event`
- Pragmatic multimodal input parsing for both legacy (`binary`) and draft media-part shapes
- Pragmatic interrupt/resume handling (`availableInterrupts`, `resume`, and `RUN_FINISHED.interrupt`)
## Security: Authentication & Authorization
The AG-UI endpoint does not enforce authentication by default. **For production deployments, you should add authentication** using FastAPI's dependency injection system via the `dependencies` parameter.
python/packages/ag-ui/agent_framework_ag_ui/__init__.py
+3
View File
@@ -10,6 +10,7 @@ from ._endpoint import add_agent_framework_fastapi_endpoint
from ._event_converters import AGUIEventConverter
from ._http_service import AGUIHttpService
from ._types import AgentState, AGUIChatOptions, AGUIRequest, PredictStateConfig, RunMetadata
from ._workflow import AgentFrameworkWorkflow, WorkflowFactory
try:
__version__ = importlib.metadata.version(__name__)
@@ -21,6 +22,8 @@ DEFAULT_TAGS = ["AG-UI"]
__all__ = [
"AgentFrameworkAgent",
"AgentFrameworkWorkflow",
"WorkflowFactory",
"add_agent_framework_fastapi_endpoint",
"AGUIChatClient",
"AGUIChatOptions",
python/packages/ag-ui/agent_framework_ag_ui/_agent.py
+3 -3
View File
@@ -8,7 +8,7 @@ from typing import Any, cast
from ag_ui.core import BaseEvent
from agent_framework import SupportsAgentRun
from ._run import run_agent_stream
from ._agent_run import run_agent_stream
class AgentConfig:
@@ -101,11 +101,11 @@ class AgentFrameworkAgent:
require_confirmation=require_confirmation,
)
async def run_agent(
async def run(
self,
input_data: dict[str, Any],
) -> AsyncGenerator[BaseEvent, None]:
"""Run the agent and yield AG-UI events.
"""Run the wrapped agent and yield AG-UI events.
Args:
input_data: The AG-UI run input containing messages, state, etc.
python/packages/ag-ui/agent_framework_ag_ui/_run.py → python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py
+51 -283
View File
@@ -2,20 +2,18 @@
"""Simplified AG-UI orchestration - single linear flow."""
from __future__ import annotations
from __future__ import annotations # noqa: I001
import json
import logging
import uuid
from collections.abc import AsyncIterable, Awaitable
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, cast
from ag_ui.core import (
BaseEvent,
CustomEvent,
MessagesSnapshotEvent,
RunFinishedEvent,
RunStartedEvent,
StateSnapshotEvent,
TextMessageContentEvent,
@@ -23,7 +21,6 @@ from ag_ui.core import (
TextMessageStartEvent,
ToolCallArgsEvent,
ToolCallEndEvent,
ToolCallResultEvent,
ToolCallStartEvent,
)
from agent_framework import (
@@ -45,6 +42,14 @@ from agent_framework.exceptions import AgentInvalidResponseException
from ._message_adapters import normalize_agui_input_messages
from ._orchestration._predictive_state import PredictiveStateHandler
from ._orchestration._tooling import collect_server_tools, merge_tools, register_additional_client_tools
from ._run_common import (
FlowState,
_build_run_finished_event, # type: ignore
_emit_content, # type: ignore
_extract_resume_payload, # type: ignore
_has_only_tool_calls, # type: ignore
_normalize_resume_interrupts, # type: ignore
)
from ._utils import (
convert_agui_tools_to_agent_framework,
generate_event_id,
@@ -86,20 +91,6 @@ def _build_safe_metadata(thread_metadata: dict[str, Any] | None) -> dict[str, An
return safe_metadata
def _has_only_tool_calls(contents: list[Any]) -> bool:
"""Check if contents have only tool calls (no text).
Args:
contents: List of content items
Returns:
True if there are tool calls but no text content
"""
has_tool_call = any(getattr(c, "type", None) == "function_call" for c in contents)
has_text = any(getattr(c, "type", None) == "text" and getattr(c, "text", None) for c in contents)
return has_tool_call and not has_text
def _should_suppress_intermediate_snapshot(
tool_name: str | None,
predict_state_config: dict[str, dict[str, str]] | None,
@@ -164,31 +155,24 @@ def _extract_approved_state_updates(
return updates
@dataclass
class FlowState:
"""Minimal explicit state for a single AG-UI run."""
message_id: str | None = None # Current text message being streamed
tool_call_id: str | None = None # Current tool call being streamed
tool_call_name: str | None = None # Name of current tool call
waiting_for_approval: bool = False # Stop after approval request
current_state: dict[str, Any] = field(default_factory=dict) # pyright: ignore[reportUnknownVariableType]
accumulated_text: str = "" # For MessagesSnapshotEvent
pending_tool_calls: list[dict[str, Any]] = field(default_factory=list) # pyright: ignore[reportUnknownVariableType]
tool_calls_by_id: dict[str, dict[str, Any]] = field(default_factory=dict) # pyright: ignore[reportUnknownVariableType]
tool_results: list[dict[str, Any]] = field(default_factory=list) # pyright: ignore[reportUnknownVariableType]
tool_calls_ended: set[str] = field(default_factory=set) # pyright: ignore[reportUnknownVariableType]
def get_tool_name(self, call_id: str | None) -> str | None:
"""Get tool name by call ID."""
if not call_id or call_id not in self.tool_calls_by_id:
return None
name = self.tool_calls_by_id[call_id]["function"].get("name")
return str(name) if name else None
def get_pending_without_end(self) -> list[dict[str, Any]]:
"""Get tool calls that started but never received an end event (declaration-only)."""
return [tc for tc in self.pending_tool_calls if tc.get("id") not in self.tool_calls_ended]
def _resume_to_tool_messages(resume_payload: Any) -> list[dict[str, Any]]:
"""Convert a resume payload into AG-UI tool messages for approval continuation."""
result: list[dict[str, Any]] = []
for interrupt in _normalize_resume_interrupts(resume_payload):
value = interrupt.get("value")
content: str
if isinstance(value, str):
content = value
else:
content = json.dumps(make_json_safe(value))
result.append(
{
"role": "tool",
"toolCallId": interrupt["id"],
"content": content,
}
)
return result
async def _normalize_response_stream(response_stream: Any) -> AsyncIterable[Any]:
@@ -303,242 +287,6 @@ def _inject_state_context(
return result
def _emit_text(content: Content, flow: FlowState, skip_text: bool = False) -> list[BaseEvent]:
"""Emit TextMessage events for TextContent."""
if not content.text:
return []
# Skip if we're in structured output mode or waiting for approval
if skip_text or flow.waiting_for_approval:
return []
events: list[BaseEvent] = []
if not flow.message_id:
flow.message_id = generate_event_id()
events.append(TextMessageStartEvent(message_id=flow.message_id, role="assistant"))
events.append(TextMessageContentEvent(message_id=flow.message_id, delta=content.text))
flow.accumulated_text += content.text
return events
def _emit_tool_call(
content: Content,
flow: FlowState,
predictive_handler: PredictiveStateHandler | None = None,
) -> list[BaseEvent]:
"""Emit ToolCall events for FunctionCallContent."""
events: list[BaseEvent] = []
tool_call_id = content.call_id or flow.tool_call_id or generate_event_id()
# Emit start event when we have a new tool call
if content.name and tool_call_id != flow.tool_call_id:
flow.tool_call_id = tool_call_id
flow.tool_call_name = content.name
if predictive_handler:
predictive_handler.reset_streaming()
events.append(
ToolCallStartEvent(
tool_call_id=tool_call_id,
tool_call_name=content.name,
parent_message_id=flow.message_id,
)
)
# Track for MessagesSnapshotEvent
tool_entry = {
"id": tool_call_id,
"type": "function",
"function": {"name": content.name, "arguments": ""},
}
flow.pending_tool_calls.append(tool_entry)
flow.tool_calls_by_id[tool_call_id] = tool_entry
elif tool_call_id:
flow.tool_call_id = tool_call_id
# Emit args if present
if content.arguments:
delta = (
content.arguments if isinstance(content.arguments, str) else json.dumps(make_json_safe(content.arguments))
)
events.append(ToolCallArgsEvent(tool_call_id=tool_call_id, delta=delta))
# Track args for MessagesSnapshotEvent
if tool_call_id in flow.tool_calls_by_id:
flow.tool_calls_by_id[tool_call_id]["function"]["arguments"] += delta
# Emit predictive state deltas
if predictive_handler and flow.tool_call_name:
delta_events = predictive_handler.emit_streaming_deltas(flow.tool_call_name, delta)
events.extend(delta_events)
return events
def _emit_tool_result(
content: Content,
flow: FlowState,
predictive_handler: PredictiveStateHandler | None = None,
) -> list[BaseEvent]:
"""Emit ToolCallResult events for function_result content."""
events: list[BaseEvent] = []
# Cannot emit tool result without a call_id to associate it with
if not content.call_id:
return events
events.append(ToolCallEndEvent(tool_call_id=content.call_id))
flow.tool_calls_ended.add(content.call_id) # Track ended tool calls
result_content = content.result if content.result is not None else ""
message_id = generate_event_id()
events.append(
ToolCallResultEvent(
message_id=message_id,
tool_call_id=content.call_id,
content=result_content,
role="tool",
)
)
# Track for MessagesSnapshotEvent
flow.tool_results.append(
{
"id": message_id,
"role": "tool",
"toolCallId": content.call_id,
"content": result_content,
}
)
# Apply predictive state updates and emit snapshot
if predictive_handler:
predictive_handler.apply_pending_updates()
if flow.current_state:
events.append(StateSnapshotEvent(snapshot=flow.current_state))
# Reset tool tracking and message context
# After tool result, any subsequent text should start a new message
flow.tool_call_id = None
flow.tool_call_name = None
# Close any open text message before resetting message_id (issue #3568)
# This handles the case where a TextMessageStartEvent was emitted for tool-only
# messages (Feature #4) but needs to be closed before starting a new message
if flow.message_id:
logger.debug("Closing text message (issue #3568 fix): message_id=%s", flow.message_id)
events.append(TextMessageEndEvent(message_id=flow.message_id))
flow.message_id = None # Reset so next text content starts a new message
return events
def _emit_approval_request(
content: Content,
flow: FlowState,
predictive_handler: PredictiveStateHandler | None = None,
require_confirmation: bool = True,
) -> list[BaseEvent]:
"""Emit events for function approval request."""
events: list[BaseEvent] = []
# function_call is required for approval requests - skip if missing
func_call = content.function_call
if not func_call:
logger.warning("Approval request content missing function_call, skipping")
return events
func_name = func_call.name or ""
func_call_id = func_call.call_id
# Extract state from function arguments if predictive
if predictive_handler and func_name:
parsed_args = func_call.parse_arguments()
result = predictive_handler.extract_state_value(func_name, parsed_args)
if result:
state_key, state_value = result
flow.current_state[state_key] = state_value
events.append(StateSnapshotEvent(snapshot=flow.current_state))
# End the original tool call
if func_call_id:
events.append(ToolCallEndEvent(tool_call_id=func_call_id))
flow.tool_calls_ended.add(func_call_id) # Track ended tool calls
# Emit custom event for UI
events.append(
CustomEvent(
name="function_approval_request",
value={
"id": content.id,
"function_call": {
"call_id": func_call_id,
"name": func_name,
"arguments": make_json_safe(func_call.parse_arguments()),
},
},
)
)
# Emit confirm_changes tool call for UI compatibility
# The complete sequence (Start -> Args -> End) signals the UI to show the confirmation dialog
if require_confirmation:
confirm_id = generate_event_id()
events.append(
ToolCallStartEvent(
tool_call_id=confirm_id,
tool_call_name="confirm_changes",
parent_message_id=flow.message_id,
)
)
args: dict[str, Any] = {
"function_name": func_name,
"function_call_id": func_call_id,
"function_arguments": make_json_safe(func_call.parse_arguments()) or {},
"steps": [{"description": f"Execute {func_name}", "status": "enabled"}],
}
args_json = json.dumps(args)
events.append(ToolCallArgsEvent(tool_call_id=confirm_id, delta=args_json))
events.append(ToolCallEndEvent(tool_call_id=confirm_id))
# Track confirm_changes in pending_tool_calls for MessagesSnapshotEvent
# The frontend needs to see this in the snapshot to render the confirmation dialog
confirm_entry = {
"id": confirm_id,
"type": "function",
"function": {"name": "confirm_changes", "arguments": args_json},
}
flow.pending_tool_calls.append(confirm_entry)
flow.tool_calls_by_id[confirm_id] = confirm_entry
flow.tool_calls_ended.add(confirm_id) # Mark as ended since we emit End event
flow.waiting_for_approval = True
return events
def _emit_content(
content: Any,
flow: FlowState,
predictive_handler: PredictiveStateHandler | None = None,
skip_text: bool = False,
require_confirmation: bool = True,
) -> list[BaseEvent]:
"""Emit appropriate events for any content type."""
content_type = getattr(content, "type", None)
if content_type == "text":
return _emit_text(content, flow, skip_text)
elif content_type == "function_call":
return _emit_tool_call(content, flow, predictive_handler)
elif content_type == "function_result":
return _emit_tool_result(content, flow, predictive_handler)
elif content_type == "function_approval_request":
return _emit_approval_request(content, flow, predictive_handler, require_confirmation)
return []
def _is_confirm_changes_response(messages: list[Any]) -> bool:
"""Check if the last message is a confirm_changes tool result (state confirmation flow).
@@ -831,7 +579,14 @@ async def run_agent_stream(
)
# Normalize messages
raw_messages = input_data.get("messages", [])
available_interrupts = input_data.get("available_interrupts") or input_data.get("availableInterrupts")
raw_messages = list(cast(list[dict[str, Any]], input_data.get("messages", []) or []))
resume_messages = _resume_to_tool_messages(_extract_resume_payload(input_data))
if available_interrupts:
logger.debug("Received available interrupts metadata: %s", available_interrupts)
if resume_messages:
logger.info(f"Appending {len(resume_messages)} synthesized resume message(s) to AG-UI input.")
raw_messages.extend(resume_messages)
messages, snapshot_messages = normalize_agui_input_messages(raw_messages)
# Check for structured output mode (skip text content)
@@ -847,7 +602,7 @@ async def run_agent_stream(
if not messages:
logger.warning("No messages provided in AG-UI input")
yield RunStartedEvent(run_id=run_id, thread_id=thread_id)
yield RunFinishedEvent(run_id=run_id, thread_id=thread_id)
yield _build_run_finished_event(run_id=run_id, thread_id=thread_id)
return
# Prepare tools
@@ -906,7 +661,7 @@ async def run_agent_stream(
yield StateSnapshotEvent(snapshot=flow.current_state)
for event in _handle_step_based_approval(messages):
yield event
yield RunFinishedEvent(run_id=run_id, thread_id=thread_id)
yield _build_run_finished_event(run_id=run_id, thread_id=thread_id)
return
# Inject state context message so the model knows current application state
@@ -1099,6 +854,19 @@ async def run_agent_stream(
flow.tool_calls_by_id[confirm_id] = confirm_entry
flow.tool_calls_ended.add(confirm_id) # Mark as ended since we emit End event
flow.waiting_for_approval = True
flow.interrupts = [
{
"id": str(confirm_id),
"value": {
"type": "function_approval_request",
"function_call": {
"call_id": tool_call_id,
"name": tool_name,
"arguments": function_arguments,
},
},
}
]
# Close any open message
if flow.message_id:
@@ -1122,4 +890,4 @@ async def run_agent_stream(
# Always emit RunFinished - confirm_changes tool call is complete (Start -> Args -> End)
# The UI will show confirmation dialog and send a new request when user responds
yield RunFinishedEvent(run_id=run_id, thread_id=thread_id)
yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=flow.interrupts)
python/packages/ag-ui/agent_framework_ag_ui/_client.py
+5
View File
@@ -439,6 +439,11 @@ class AGUIChatClient(
messages=agui_messages,
state=state,
tools=agui_tools,
available_interrupts=cast(
list[dict[str, Any]] | None,
options.get("available_interrupts") or options.get("availableInterrupts"),
),
resume=cast(dict[str, Any] | None, options.get("resume")),
):
logger.debug(f"[AGUIChatClient] Raw AG-UI event: {event}")
update = converter.convert_event(event)
python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py
+59 -25
View File
@@ -9,21 +9,23 @@ import logging
from collections.abc import AsyncGenerator, Sequence
from typing import Any
from ag_ui.core import RunErrorEvent
from ag_ui.encoder import EventEncoder
from agent_framework import SupportsAgentRun
from fastapi import FastAPI
from agent_framework import SupportsAgentRun, Workflow
from fastapi import FastAPI, HTTPException
from fastapi.params import Depends
from fastapi.responses import StreamingResponse
from ._agent import AgentFrameworkAgent
from ._types import AGUIRequest
from ._workflow import AgentFrameworkWorkflow
logger = logging.getLogger(__name__)
def add_agent_framework_fastapi_endpoint(
app: FastAPI,
agent: SupportsAgentRun | AgentFrameworkAgent,
agent: SupportsAgentRun | AgentFrameworkAgent | Workflow | AgentFrameworkWorkflow,
path: str = "/",
state_schema: Any | None = None,
predict_state_config: dict[str, dict[str, str]] | None = None,
@@ -49,17 +51,24 @@ def add_agent_framework_fastapi_endpoint(
authentication checks, rate limiting, or other middleware-like behavior.
Example: `dependencies=[Depends(verify_api_key)]`
"""
if isinstance(agent, SupportsAgentRun):
wrapped_agent = AgentFrameworkAgent(
protocol_runner: AgentFrameworkAgent | AgentFrameworkWorkflow
if isinstance(agent, AgentFrameworkWorkflow):
protocol_runner = agent
elif isinstance(agent, AgentFrameworkAgent):
protocol_runner = agent
elif isinstance(agent, Workflow):
protocol_runner = AgentFrameworkWorkflow(workflow=agent)
elif isinstance(agent, SupportsAgentRun):
protocol_runner = AgentFrameworkAgent(
agent=agent,
state_schema=state_schema,
predict_state_config=predict_state_config,
)
else:
wrapped_agent = agent
raise TypeError("agent must be SupportsAgentRun, Workflow, AgentFrameworkAgent, or AgentFrameworkWorkflow.")
@app.post(path, tags=tags or ["AG-UI"], dependencies=dependencies, response_model=None) # type: ignore[arg-type]
async def agent_endpoint(request_body: AGUIRequest) -> StreamingResponse | dict[str, str]:
async def agent_endpoint(request_body: AGUIRequest) -> StreamingResponse:
"""Handle AG-UI agent requests.
Note: Function is accessed via FastAPI's decorator registration,
@@ -82,25 +91,50 @@ def add_agent_framework_fastapi_endpoint(
async def event_generator() -> AsyncGenerator[str]:
encoder = EventEncoder()
event_count = 0
async for event in wrapped_agent.run_agent(input_data):
event_count += 1
event_type_name = getattr(event, "type", type(event).__name__)
# Log important events at INFO level
if "TOOL_CALL" in str(event_type_name) or "RUN" in str(event_type_name):
if hasattr(event, "model_dump"):
event_data = event.model_dump(exclude_none=True)
logger.info(f"[{path}] Event {event_count}: {event_type_name} - {event_data}")
else:
logger.info(f"[{path}] Event {event_count}: {event_type_name}")
try:
async for event in protocol_runner.run(input_data):
event_count += 1
event_type_name = getattr(event, "type", type(event).__name__)
# Log important events at INFO level
if "TOOL_CALL" in str(event_type_name) or "RUN" in str(event_type_name):
if hasattr(event, "model_dump"):
event_data = event.model_dump(exclude_none=True)
logger.info(f"[{path}] Event {event_count}: {event_type_name} - {event_data}")
else:
logger.info(f"[{path}] Event {event_count}: {event_type_name}")
encoded = encoder.encode(event)
logger.debug(
f"[{path}] Encoded as: {encoded[:200]}..."
if len(encoded) > 200
else f"[{path}] Encoded as: {encoded}"
try:
encoded = encoder.encode(event)
except Exception as encode_error:
logger.exception("[%s] Failed to encode event %s", path, event_type_name)
run_error = RunErrorEvent(
message="An internal error has occurred while streaming events.",
code=type(encode_error).__name__,
)
try:
yield encoder.encode(run_error)
except Exception:
logger.exception("[%s] Failed to encode RUN_ERROR event", path)
return
logger.debug(
f"[{path}] Encoded as: {encoded[:200]}..."
if len(encoded) > 200
else f"[{path}] Encoded as: {encoded}"
)
yield encoded
logger.info(f"[{path}] Completed streaming {event_count} events")
except Exception as stream_error:
logger.exception("[%s] Streaming failed", path)
run_error = RunErrorEvent(
message="An internal error has occurred while streaming events.",
code=type(stream_error).__name__,
)
yield encoded
logger.info(f"[{path}] Completed streaming {event_count} events")
try:
yield encoder.encode(run_error)
except Exception:
logger.exception("[%s] Failed to encode RUN_ERROR event", path)
return StreamingResponse(
event_generator(),
@@ -113,4 +147,4 @@ def add_agent_framework_fastapi_endpoint(
)
except Exception as e:
logger.error(f"Error in agent endpoint: {e}", exc_info=True)
return {"error": "An internal error has occurred."}
raise HTTPException(status_code=500, detail="An internal error has occurred.") from e
python/packages/ag-ui/agent_framework_ag_ui/_event_converters.py
+33 -5
View File
@@ -55,7 +55,8 @@ class AGUIEventConverter:
update = converter.convert_event(event)
assert update.contents[0].text == "Hello"
"""
event_type = event.get("type", "")
raw_event_type = str(event.get("type", ""))
event_type = raw_event_type.upper()
if event_type == "RUN_STARTED":
return self._handle_run_started(event)
@@ -77,6 +78,8 @@ class AGUIEventConverter:
return self._handle_run_finished(event)
elif event_type == "RUN_ERROR":
return self._handle_run_error(event)
elif event_type in {"CUSTOM", "CUSTOM_EVENT"}:
return self._handle_custom_event(event, raw_event_type)
return None
@@ -176,14 +179,20 @@ class AGUIEventConverter:
def _handle_run_finished(self, event: dict[str, Any]) -> ChatResponseUpdate:
"""Handle RUN_FINISHED event."""
additional_properties: dict[str, Any] = {
"thread_id": self.thread_id,
"run_id": self.run_id,
}
if "interrupt" in event:
additional_properties["interrupt"] = event.get("interrupt")
if "result" in event:
additional_properties["result"] = event.get("result")
return ChatResponseUpdate(
role="assistant",
finish_reason="stop",
contents=[],
additional_properties={
"thread_id": self.thread_id,
"run_id": self.run_id,
},
additional_properties=additional_properties,
)
def _handle_run_error(self, event: dict[str, Any]) -> ChatResponseUpdate:
@@ -204,3 +213,22 @@ class AGUIEventConverter:
"run_id": self.run_id,
},
)
def _handle_custom_event(self, event: dict[str, Any], raw_event_type: str) -> ChatResponseUpdate:
"""Handle CUSTOM/CUSTOM_EVENT events.
Custom events are surfaced as metadata so callers can inspect protocol-specific payloads.
"""
return ChatResponseUpdate(
role="assistant",
contents=[],
additional_properties={
"thread_id": self.thread_id,
"run_id": self.run_id,
"ag_ui_custom_event": {
"name": event.get("name"),
"value": event.get("value"),
"raw_type": raw_event_type,
},
},
)
python/packages/ag-ui/agent_framework_ag_ui/_http_service.py
+12 -1
View File
@@ -66,6 +66,8 @@ class AGUIHttpService:
messages: list[dict[str, Any]],
state: dict[str, Any] | None = None,
tools: list[dict[str, Any]] | None = None,
available_interrupts: list[dict[str, Any]] | None = None,
resume: dict[str, Any] | None = None,
) -> AsyncIterable[dict[str, Any]]:
"""Post a run request and stream AG-UI events.
@@ -75,6 +77,8 @@ class AGUIHttpService:
messages: List of messages in AG-UI format
state: Optional state object to send to server
tools: Optional list of tools available to the agent
available_interrupts: Optional list of interrupt descriptors available for resumption
resume: Optional resume payload to continue a paused run
Yields:
AG-UI event dictionaries parsed from SSE stream
@@ -109,9 +113,16 @@ class AGUIHttpService:
if tools is not None:
request_data["tools"] = tools
if available_interrupts is not None:
request_data["availableInterrupts"] = available_interrupts
if resume is not None:
request_data["resume"] = resume
logger.debug(
f"Posting run to {self.endpoint}: thread_id={thread_id}, run_id={run_id}, "
f"messages={len(messages)}, has_state={state is not None}, has_tools={tools is not None}"
f"messages={len(messages)}, has_state={state is not None}, has_tools={tools is not None}, "
f"has_available_interrupts={available_interrupts is not None}, has_resume={resume is not None}"
)
# Stream the response using SSE
python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py
+237 -28
View File
@@ -4,6 +4,8 @@
from __future__ import annotations
import base64
import binascii
import json
import logging
from typing import Any, cast
@@ -253,12 +255,235 @@ def _deduplicate_messages(messages: list[Message]) -> list[Message]:
return unique_messages
def _parse_multimodal_media_part(part: dict[str, Any]) -> Content | None:
"""Convert a multimodal media part into Agent Framework content."""
part_type = str(part.get("type", "")).lower()
source = part.get("source")
mime_type = cast(
str | None,
part.get("mimeType")
or part.get("mime_type")
or {
"image": "image/*",
"audio": "audio/*",
"video": "video/*",
"document": "application/octet-stream",
"binary": "application/octet-stream",
}.get(part_type, "application/octet-stream"),
)
url = cast(str | None, part.get("url") or part.get("uri"))
data = cast(str | None, part.get("data"))
binary_id = cast(str | None, part.get("id"))
if isinstance(source, dict):
source_dict = cast(dict[str, Any], source)
source_type = str(source_dict.get("type", "")).lower()
source_mime = source_dict.get("mimeType") or source_dict.get("mime_type")
if isinstance(source_mime, str) and source_mime:
mime_type = source_mime
if source_type in {"url", "uri"}:
url = cast(str | None, source_dict.get("url") or source_dict.get("uri"))
elif source_type in {"base64", "data", "binary"}:
data = cast(str | None, source_dict.get("data"))
elif source_type in {"id", "file"}:
binary_id = cast(str | None, source_dict.get("id"))
else:
url = cast(str | None, source_dict.get("url") or source_dict.get("uri") or url)
data = cast(str | None, source_dict.get("data") or data)
binary_id = cast(str | None, source_dict.get("id") or binary_id)
if isinstance(url, str) and url:
return Content.from_uri(uri=url, media_type=mime_type)
if isinstance(data, str) and data:
if data.startswith("data:"):
return Content.from_uri(uri=data, media_type=mime_type)
try:
decoded = base64.b64decode(data, validate=True)
return Content.from_data(data=decoded, media_type=mime_type or "application/octet-stream")
except (binascii.Error, ValueError):
logger.debug("Strict base64 decode failed for AG-UI media payload (mime_type=%s).", mime_type)
try:
decoded = base64.b64decode(data)
return Content.from_data(data=decoded, media_type=mime_type or "application/octet-stream")
except (binascii.Error, ValueError):
logger.warning(
"Failed to decode AG-UI media payload as base64; falling back to data URI (mime_type=%s).",
mime_type,
exc_info=True,
)
# Best effort fallback for malformed payloads.
return Content.from_uri(
uri=f"data:{mime_type or 'application/octet-stream'};base64,{data}",
media_type=mime_type,
)
if isinstance(binary_id, str) and binary_id:
return Content.from_uri(uri=f"ag-ui://binary/{binary_id}", media_type=mime_type)
return None
def _convert_agui_content_to_framework(content: Any) -> list[Content]:
"""Convert AG-UI content payloads to Agent Framework Content entries."""
if isinstance(content, str):
return [Content.from_text(text=content)]
if isinstance(content, list):
converted: list[Content] = []
for item in content:
if isinstance(item, str):
converted.append(Content.from_text(text=item))
continue
if not isinstance(item, dict):
converted.append(Content.from_text(text=str(item)))
continue
part = cast(dict[str, Any], item)
part_type = str(part.get("type", "")).lower()
if part_type in {"text", "input_text"}:
converted.append(Content.from_text(text=str(part.get("text", ""))))
continue
if part_type in {"binary", "image", "audio", "video", "document"}:
media_content = _parse_multimodal_media_part(part)
if media_content is not None:
converted.append(media_content)
continue
text_value = part.get("text")
if isinstance(text_value, str):
converted.append(Content.from_text(text=text_value))
else:
converted.append(Content.from_text(text=str(part)))
return converted
if content is None:
return []
return [Content.from_text(text=str(content))]
def _normalize_snapshot_content(content: Any) -> Any:
"""Normalize AG-UI message content for snapshot payloads.
Preserve multimodal fidelity whenever non-text parts are present.
"""
if isinstance(content, list):
has_non_text_parts = False
normalized_parts: list[dict[str, Any]] = []
text_parts: list[str] = []
def _legacy_binary_part(part: dict[str, Any]) -> dict[str, Any]:
"""Convert draft/legacy multimodal parts to AG-UI snapshot binary shape."""
normalized: dict[str, Any] = {"type": "binary"}
mime_type = cast(str | None, part.get("mimeType") or part.get("mime_type"))
url = cast(str | None, part.get("url") or part.get("uri"))
data = cast(str | None, part.get("data"))
binary_id = cast(str | None, part.get("id"))
source = part.get("source")
if isinstance(source, dict):
source_part = cast(dict[str, Any], source)
source_mime = source_part.get("mimeType") or source_part.get("mime_type")
if isinstance(source_mime, str) and source_mime:
mime_type = source_mime
source_type = str(source_part.get("type", "")).lower()
if source_type in {"url", "uri"}:
url = cast(str | None, source_part.get("url") or source_part.get("uri"))
elif source_type in {"base64", "data", "binary"}:
data = cast(str | None, source_part.get("data"))
elif source_type in {"id", "file"}:
binary_id = cast(str | None, source_part.get("id"))
else:
url = cast(str | None, source_part.get("url") or source_part.get("uri") or url)
data = cast(str | None, source_part.get("data") or data)
binary_id = cast(str | None, source_part.get("id") or binary_id)
if isinstance(mime_type, str) and mime_type:
normalized["mimeType"] = mime_type
if isinstance(url, str) and url:
normalized["url"] = url
if isinstance(data, str) and data:
normalized["data"] = data
if isinstance(binary_id, str) and binary_id:
normalized["id"] = binary_id
return normalized
for item in content:
if isinstance(item, str):
text_parts.append(item)
normalized_parts.append({"type": "text", "text": item})
continue
if not isinstance(item, dict):
item_text = str(item)
text_parts.append(item_text)
normalized_parts.append({"type": "text", "text": item_text})
continue
part = cast(dict[str, Any], item).copy()
part_type = str(part.get("type", "")).lower()
if part_type == "input_text":
part["type"] = "text"
part_type = "text"
elif part_type == "input_image":
part["type"] = "binary"
part_type = "binary"
if part_type == "text":
text_parts.append(str(part.get("text", "")))
else:
has_non_text_parts = True
if part_type in {"binary", "image", "audio", "video", "document"}:
normalized_parts.append(_legacy_binary_part(part))
continue
if "mime_type" in part and "mimeType" not in part:
part["mimeType"] = part.get("mime_type")
source = part.get("source")
if isinstance(source, dict):
source_part = cast(dict[str, Any], source)
if "mime_type" in source_part and "mimeType" not in source_part:
source_part["mimeType"] = source_part.get("mime_type")
normalized_parts.append(part)
if has_non_text_parts:
return normalized_parts
return "".join(text_parts)
if content is None:
return ""
return content
def normalize_agui_input_messages(
messages: list[dict[str, Any]],
*,
sanitize_tool_history: bool = True,
) -> tuple[list[Message], list[dict[str, Any]]]:
"""Normalize raw AG-UI messages into provider and snapshot formats."""
"""Normalize raw AG-UI messages into provider and snapshot formats.
Args:
messages: Raw AG-UI messages.
sanitize_tool_history: Apply agent-run specific tool history repair logic.
Keep enabled for standard agent runs; disable for native workflow runs
where pending-request responses must come explicitly from interrupt resume.
"""
provider_messages = agui_messages_to_agent_framework(messages)
provider_messages = _sanitize_tool_history(provider_messages)
if sanitize_tool_history:
provider_messages = _sanitize_tool_history(provider_messages)
provider_messages = _deduplicate_messages(provider_messages)
snapshot_messages = agui_messages_to_snapshot_format(messages)
return provider_messages, snapshot_messages
@@ -562,10 +787,10 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Mes
tool_calls = msg.get("tool_calls") or msg.get("toolCalls")
if tool_calls:
contents: list[Any] = []
# Include any assistant text content if present
content_text = msg.get("content")
if isinstance(content_text, str) and content_text:
contents.append(Content.from_text(text=content_text))
# Include any assistant content if present
content_value = msg.get("content")
if content_value not in (None, ""):
contents.extend(_convert_agui_content_to_framework(content_value))
# Convert each tool call entry
for tc in tool_calls:
if not isinstance(tc, dict):
@@ -620,12 +845,12 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Mes
chat_msg = Message(role=role, contents=approval_contents) # type: ignore[call-overload]
else:
# Regular text message
# Regular message content (text or multimodal)
content = msg.get("content", "")
if isinstance(content, str):
chat_msg = Message(role=role, contents=[Content.from_text(text=content)]) # type: ignore[call-overload]
else:
chat_msg = Message(role=role, contents=[Content.from_text(text=str(content))]) # type: ignore[call-overload]
converted_contents = _convert_agui_content_to_framework(content)
if not converted_contents:
converted_contents = [Content.from_text(text="")]
chat_msg = Message(role=role, contents=converted_contents) # type: ignore[call-overload]
if "id" in msg:
chat_msg.message_id = msg["id"]
@@ -760,23 +985,7 @@ def agui_messages_to_snapshot_format(messages: list[dict[str, Any]]) -> list[dic
normalized_msg["id"] = generate_event_id()
# Normalize content field
content = normalized_msg.get("content")
if isinstance(content, list):
# Convert content array format to simple string
text_parts: list[str] = []
for item in content:
if isinstance(item, dict):
# Convert 'input_text' to 'text' type
if item.get("type") == "input_text":
text_parts.append(str(item.get("text", "")))
elif item.get("type") == "text":
text_parts.append(str(item.get("text", "")))
else:
# Other types - just extract text field if present
text_parts.append(str(item.get("text", "")))
normalized_msg["content"] = "".join(text_parts)
elif content is None:
normalized_msg["content"] = ""
normalized_msg["content"] = _normalize_snapshot_content(normalized_msg.get("content"))
tool_calls = normalized_msg.get("tool_calls") or normalized_msg.get("toolCalls")
if isinstance(tool_calls, list):
python/packages/ag-ui/agent_framework_ag_ui/_orchestration/_helpers.py
-1
View File
@@ -2,7 +2,6 @@
"""Helper functions for orchestration logic.
Most orchestration helpers have been moved inline to _run.py.
This module retains utilities that may be useful for testing or extensions.
"""
python/packages/ag-ui/agent_framework_ag_ui/_run_common.py
+378
View File
@@ -0,0 +1,378 @@
# Copyright (c) Microsoft. All rights reserved.
"""Shared AG-UI run helpers used by agent and workflow runners."""
from __future__ import annotations
import json
import logging
from dataclasses import dataclass, field
from typing import Any, cast
from ag_ui.core import (
BaseEvent,
CustomEvent,
RunFinishedEvent,
StateSnapshotEvent,
TextMessageContentEvent,
TextMessageEndEvent,
TextMessageStartEvent,
ToolCallArgsEvent,
ToolCallEndEvent,
ToolCallResultEvent,
ToolCallStartEvent,
)
from agent_framework import Content
from ._orchestration._predictive_state import PredictiveStateHandler
from ._utils import generate_event_id, make_json_safe
logger = logging.getLogger(__name__)
def _has_only_tool_calls(contents: list[Any]) -> bool:
"""Check if contents have only tool calls (no text)."""
has_tool_call = any(getattr(c, "type", None) == "function_call" for c in contents)
has_text = any(getattr(c, "type", None) == "text" and getattr(c, "text", None) for c in contents)
return has_tool_call and not has_text
def _normalize_resume_interrupts(resume_payload: Any) -> list[dict[str, Any]]:
"""Normalize resume payload to a list of interrupt responses."""
if resume_payload is None:
return []
if isinstance(resume_payload, list):
candidates = resume_payload
elif isinstance(resume_payload, dict):
resume_dict = cast(dict[str, Any], resume_payload)
if isinstance(resume_dict.get("interrupts"), list):
candidates = cast(list[Any], resume_dict["interrupts"])
elif isinstance(resume_dict.get("interrupt"), list):
candidates = cast(list[Any], resume_dict["interrupt"])
else:
candidates = [resume_dict]
else:
return []
normalized: list[dict[str, Any]] = []
for item in candidates:
if not isinstance(item, dict):
continue
item_dict = cast(dict[str, Any], item)
interrupt_id = item_dict.get("id") or item_dict.get("interruptId") or item_dict.get("toolCallId")
if not interrupt_id:
continue
if "value" in item_dict:
value = item_dict.get("value")
elif "response" in item_dict:
value = item_dict.get("response")
else:
value = {k: v for k, v in item_dict.items() if k not in {"id", "interruptId", "toolCallId", "type"}}
normalized.append({"id": str(interrupt_id), "value": value})
return normalized
def _extract_resume_payload(input_data: dict[str, Any]) -> Any:
"""Extract resume payload from standard and forwarded-props request locations."""
resume_payload = input_data.get("resume")
if resume_payload is not None:
return resume_payload
forwarded_props = input_data.get("forwarded_props") or input_data.get("forwardedProps")
if not isinstance(forwarded_props, dict):
return None
forwarded_props_dict = cast(dict[str, Any], forwarded_props)
command = forwarded_props_dict.get("command")
if isinstance(command, dict):
command_dict = cast(dict[str, Any], command)
if command_dict.get("resume") is not None:
return command_dict.get("resume")
return forwarded_props_dict.get("resume")
def _build_run_finished_event(
run_id: str, thread_id: str, interrupts: list[dict[str, Any]] | None = None
) -> RunFinishedEvent:
"""Create a RUN_FINISHED event, optionally carrying interrupt metadata."""
if interrupts:
return RunFinishedEvent(run_id=run_id, thread_id=thread_id, interrupt=interrupts) # type: ignore[call-arg]
return RunFinishedEvent(run_id=run_id, thread_id=thread_id)
@dataclass
class FlowState:
"""Minimal explicit state for a single AG-UI run."""
message_id: str | None = None
tool_call_id: str | None = None
tool_call_name: str | None = None
waiting_for_approval: bool = False
current_state: dict[str, Any] = field(default_factory=dict) # pyright: ignore[reportUnknownVariableType]
accumulated_text: str = ""
pending_tool_calls: list[dict[str, Any]] = field(default_factory=list) # pyright: ignore[reportUnknownVariableType]
tool_calls_by_id: dict[str, dict[str, Any]] = field(default_factory=dict) # pyright: ignore[reportUnknownVariableType]
tool_results: list[dict[str, Any]] = field(default_factory=list) # pyright: ignore[reportUnknownVariableType]
tool_calls_ended: set[str] = field(default_factory=set) # pyright: ignore[reportUnknownVariableType]
interrupts: list[dict[str, Any]] = field(default_factory=list) # pyright: ignore[reportUnknownVariableType]
def get_tool_name(self, call_id: str | None) -> str | None:
"""Get tool name by call ID."""
if not call_id or call_id not in self.tool_calls_by_id:
return None
name = self.tool_calls_by_id[call_id]["function"].get("name")
return str(name) if name else None
def get_pending_without_end(self) -> list[dict[str, Any]]:
"""Get tool calls that started but never received an end event (declaration-only)."""
return [tc for tc in self.pending_tool_calls if tc.get("id") not in self.tool_calls_ended]
def _emit_text(content: Content, flow: FlowState, skip_text: bool = False) -> list[BaseEvent]:
"""Emit TextMessage events for TextContent."""
if not content.text:
return []
if skip_text or flow.waiting_for_approval:
return []
events: list[BaseEvent] = []
if not flow.message_id:
flow.message_id = generate_event_id()
flow.accumulated_text = ""
events.append(TextMessageStartEvent(message_id=flow.message_id, role="assistant"))
elif flow.accumulated_text and content.text == flow.accumulated_text:
# Guard against full-message replay chunks that can appear after streaming deltas.
logger.debug("Skipping duplicate full-text delta for message_id=%s", flow.message_id)
return []
events.append(TextMessageContentEvent(message_id=flow.message_id, delta=content.text))
flow.accumulated_text += content.text
return events
def _emit_tool_call(
content: Content,
flow: FlowState,
predictive_handler: PredictiveStateHandler | None = None,
) -> list[BaseEvent]:
"""Emit ToolCall events for FunctionCallContent."""
events: list[BaseEvent] = []
tool_call_id = content.call_id or flow.tool_call_id or generate_event_id()
if content.name and tool_call_id != flow.tool_call_id:
flow.tool_call_id = tool_call_id
flow.tool_call_name = content.name
if predictive_handler:
predictive_handler.reset_streaming()
events.append(
ToolCallStartEvent(
tool_call_id=tool_call_id,
tool_call_name=content.name,
parent_message_id=flow.message_id,
)
)
tool_entry = {
"id": tool_call_id,
"type": "function",
"function": {"name": content.name, "arguments": ""},
}
flow.pending_tool_calls.append(tool_entry)
flow.tool_calls_by_id[tool_call_id] = tool_entry
elif tool_call_id:
flow.tool_call_id = tool_call_id
if content.arguments:
delta = (
content.arguments if isinstance(content.arguments, str) else json.dumps(make_json_safe(content.arguments))
)
events.append(ToolCallArgsEvent(tool_call_id=tool_call_id, delta=delta))
if tool_call_id in flow.tool_calls_by_id:
flow.tool_calls_by_id[tool_call_id]["function"]["arguments"] += delta
if predictive_handler and flow.tool_call_name:
delta_events = predictive_handler.emit_streaming_deltas(flow.tool_call_name, delta)
events.extend(delta_events)
return events
def _emit_tool_result(
content: Content,
flow: FlowState,
predictive_handler: PredictiveStateHandler | None = None,
) -> list[BaseEvent]:
"""Emit ToolCallResult events for function_result content."""
events: list[BaseEvent] = []
if not content.call_id:
return events
events.append(ToolCallEndEvent(tool_call_id=content.call_id))
flow.tool_calls_ended.add(content.call_id)
raw_result = content.result if content.result is not None else ""
result_content = raw_result if isinstance(raw_result, str) else json.dumps(make_json_safe(raw_result))
message_id = generate_event_id()
events.append(
ToolCallResultEvent(
message_id=message_id,
tool_call_id=content.call_id,
content=result_content,
role="tool",
)
)
flow.tool_results.append(
{
"id": message_id,
"role": "tool",
"toolCallId": content.call_id,
"content": result_content,
}
)
if predictive_handler:
predictive_handler.apply_pending_updates()
if flow.current_state:
events.append(StateSnapshotEvent(snapshot=flow.current_state))
flow.tool_call_id = None
flow.tool_call_name = None
if flow.message_id:
logger.debug("Closing text message (issue #3568 fix): message_id=%s", flow.message_id)
events.append(TextMessageEndEvent(message_id=flow.message_id))
flow.message_id = None
flow.accumulated_text = ""
return events
def _emit_approval_request(
content: Content,
flow: FlowState,
predictive_handler: PredictiveStateHandler | None = None,
require_confirmation: bool = True,
) -> list[BaseEvent]:
"""Emit events for function approval request."""
events: list[BaseEvent] = []
func_call = content.function_call
if not func_call:
logger.warning("Approval request content missing function_call, skipping")
return events
func_name = func_call.name or ""
func_call_id = func_call.call_id
if predictive_handler and func_name:
parsed_args = func_call.parse_arguments()
result = predictive_handler.extract_state_value(func_name, parsed_args)
if result:
state_key, state_value = result
flow.current_state[state_key] = state_value
events.append(StateSnapshotEvent(snapshot=flow.current_state))
if func_call_id:
events.append(ToolCallEndEvent(tool_call_id=func_call_id))
flow.tool_calls_ended.add(func_call_id)
events.append(
CustomEvent(
name="function_approval_request",
value={
"id": content.id,
"function_call": {
"call_id": func_call_id,
"name": func_name,
"arguments": make_json_safe(func_call.parse_arguments()),
},
},
)
)
interrupt_id = func_call_id or content.id
if interrupt_id:
flow.interrupts = [
{
"id": str(interrupt_id),
"value": {
"type": "function_approval_request",
"function_call": {
"call_id": func_call_id,
"name": func_name,
"arguments": make_json_safe(func_call.parse_arguments()),
},
},
}
]
if require_confirmation:
confirm_id = generate_event_id()
events.append(
ToolCallStartEvent(
tool_call_id=confirm_id,
tool_call_name="confirm_changes",
parent_message_id=flow.message_id,
)
)
args: dict[str, Any] = {
"function_name": func_name,
"function_call_id": func_call_id,
"function_arguments": make_json_safe(func_call.parse_arguments()) or {},
"steps": [{"description": f"Execute {func_name}", "status": "enabled"}],
}
args_json = json.dumps(args)
events.append(ToolCallArgsEvent(tool_call_id=confirm_id, delta=args_json))
events.append(ToolCallEndEvent(tool_call_id=confirm_id))
confirm_entry = {
"id": confirm_id,
"type": "function",
"function": {"name": "confirm_changes", "arguments": args_json},
}
flow.pending_tool_calls.append(confirm_entry)
flow.tool_calls_by_id[confirm_id] = confirm_entry
flow.tool_calls_ended.add(confirm_id)
flow.waiting_for_approval = True
return events
def _emit_usage(content: Content) -> list[BaseEvent]:
"""Emit usage details as a protocol-level custom event."""
usage_details = make_json_safe(content.usage_details or {})
return [CustomEvent(name="usage", value=usage_details)]
def _emit_content(
content: Any,
flow: FlowState,
predictive_handler: PredictiveStateHandler | None = None,
skip_text: bool = False,
require_confirmation: bool = True,
) -> list[BaseEvent]:
"""Emit appropriate events for any content type."""
content_type = getattr(content, "type", None)
if content_type == "text":
return _emit_text(content, flow, skip_text)
if content_type == "function_call":
return _emit_tool_call(content, flow, predictive_handler)
if content_type == "function_result":
return _emit_tool_result(content, flow, predictive_handler)
if content_type == "function_approval_request":
return _emit_approval_request(content, flow, predictive_handler, require_confirmation)
if content_type == "usage":
return _emit_usage(content)
logger.debug("Skipping unsupported content type in AG-UI emitter: %s", content_type)
return []
python/packages/ag-ui/agent_framework_ag_ui/_types.py
+20 -1
View File
@@ -6,7 +6,7 @@ import sys
from typing import Any, Generic
from agent_framework import ChatOptions
from pydantic import BaseModel, Field
from pydantic import AliasChoices, BaseModel, Field
if sys.version_info >= (3, 13):
from typing import TypeVar # type: ignore # pragma: no cover
@@ -53,10 +53,12 @@ class AGUIRequest(BaseModel):
)
run_id: str | None = Field(
None,
validation_alias=AliasChoices("run_id", "runId"),
description="Optional run identifier for tracking",
)
thread_id: str | None = Field(
None,
validation_alias=AliasChoices("thread_id", "threadId"),
description="Optional thread identifier for conversation context",
)
state: dict[str, Any] | None = Field(
@@ -73,12 +75,23 @@ class AGUIRequest(BaseModel):
)
forwarded_props: dict[str, Any] | None = Field(
None,
validation_alias=AliasChoices("forwarded_props", "forwardedProps"),
description="Additional properties forwarded to the agent",
)
parent_run_id: str | None = Field(
None,
validation_alias=AliasChoices("parent_run_id", "parentRunId"),
description="ID of the run that spawned this run",
)
available_interrupts: list[dict[str, Any]] | None = Field(
None,
validation_alias=AliasChoices("availableInterrupts", "available_interrupts"),
description="List of interrupts that can be resumed by the server",
)
resume: dict[str, Any] | None = Field(
None,
description="Resume payload containing interrupt responses",
)
# region AG-UI Chat Options TypedDict
@@ -140,6 +153,12 @@ class AGUIChatOptions(ChatOptions[ResponseModelT], Generic[ResponseModelT], tota
context: dict[str, Any]
"""Shared context/state to send to the server."""
available_interrupts: list[dict[str, Any]]
"""Interrupt descriptors available for resumption."""
resume: dict[str, Any]
"""Interrupt resume payload to continue a paused run."""
# ChatOptions fields not applicable for AG-UI
store: None # type: ignore[misc]
"""Not applicable for AG-UI protocol."""
python/packages/ag-ui/agent_framework_ag_ui/_workflow.py
+82
View File
@@ -0,0 +1,82 @@
# Copyright (c) Microsoft. All rights reserved.
"""Workflow wrapper for AG-UI protocol compatibility."""
from __future__ import annotations
import uuid
from collections.abc import AsyncGenerator, Callable
from typing import Any
from ag_ui.core import BaseEvent
from agent_framework import Workflow
from ._workflow_run import run_workflow_stream
WorkflowFactory = Callable[[str], Workflow]
class AgentFrameworkWorkflow:
"""Base AG-UI workflow wrapper.
Can wrap a native ``Workflow`` or be subclassed for custom ``run`` behavior.
"""
def __init__(
self,
workflow: Workflow | None = None,
*,
workflow_factory: WorkflowFactory | None = None,
name: str | None = None,
description: str | None = None,
) -> None:
if workflow is not None and workflow_factory is not None:
raise ValueError("Pass either workflow= or workflow_factory=, not both.")
self.workflow = workflow
self._workflow_factory = workflow_factory
self._workflow_by_thread: dict[str, Workflow] = {}
self.name = name if name is not None else getattr(workflow, "name", "workflow")
self.description = description if description is not None else getattr(workflow, "description", "")
@staticmethod
def _thread_id_from_input(input_data: dict[str, Any]) -> str:
"""Resolve a stable thread id from AG-UI input payload."""
thread_id = input_data.get("thread_id") or input_data.get("threadId")
if thread_id is not None:
return str(thread_id)
return str(uuid.uuid4())
def _resolve_workflow(self, thread_id: str) -> Workflow:
"""Get the workflow instance for the current run."""
if self.workflow is not None:
return self.workflow
if self._workflow_factory is None:
raise NotImplementedError("No workflow is attached. Override run or pass workflow=/workflow_factory=.")
workflow = self._workflow_by_thread.get(thread_id)
if workflow is None:
workflow = self._workflow_factory(thread_id)
if not isinstance(workflow, Workflow):
raise TypeError("workflow_factory must return a Workflow instance.")
self._workflow_by_thread[thread_id] = workflow
return workflow
def clear_thread_workflow(self, thread_id: str) -> None:
"""Drop a single cached thread workflow instance."""
self._workflow_by_thread.pop(thread_id, None)
def clear_workflow_cache(self) -> None:
"""Drop all cached thread workflow instances."""
self._workflow_by_thread.clear()
async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]:
"""Run the wrapped workflow and yield AG-UI events.
Subclasses may override this to provide custom AG-UI streams.
"""
thread_id = self._thread_id_from_input(input_data)
workflow = self._resolve_workflow(thread_id)
async for event in run_workflow_stream(input_data, workflow):
yield event
python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py
+727
View File
@@ -0,0 +1,727 @@
# Copyright (c) Microsoft. All rights reserved.
"""Native AG-UI orchestration for MAF Workflow streams."""
from __future__ import annotations
import json
import logging
import uuid
from collections.abc import AsyncGenerator
from typing import Any, cast, get_args, get_origin
from ag_ui.core import (
ActivitySnapshotEvent,
BaseEvent,
CustomEvent,
RunErrorEvent,
RunStartedEvent,
StepFinishedEvent,
StepStartedEvent,
TextMessageEndEvent,
ToolCallArgsEvent,
ToolCallEndEvent,
ToolCallStartEvent,
)
from agent_framework import AgentResponse, AgentResponseUpdate, Content, Message, Workflow, WorkflowRunState
from ._message_adapters import normalize_agui_input_messages
from ._run_common import (
FlowState,
_build_run_finished_event,
_emit_content,
_extract_resume_payload,
_normalize_resume_interrupts,
)
from ._utils import generate_event_id, make_json_safe
logger = logging.getLogger(__name__)
_TERMINAL_STATES: set[str] = {
WorkflowRunState.IDLE.value,
WorkflowRunState.IDLE_WITH_PENDING_REQUESTS.value,
WorkflowRunState.CANCELLED.value,
}
_WORKFLOW_EVENT_BASE_FIELDS: set[str] = {
"type",
"data",
"origin",
"state",
"details",
"executor_id",
"_request_id",
"_source_executor_id",
"_request_type",
"_response_type",
"iteration",
}
_INTERRUPT_CARD_EVENT_NAME = "WorkflowInterruptEvent"
async def _pending_request_events(workflow: Workflow) -> dict[str, Any]:
"""Best-effort retrieval of pending request_info events from workflow context."""
runner_context = getattr(workflow, "_runner_context", None)
if runner_context is None:
return {}
get_pending = getattr(runner_context, "get_pending_request_info_events", None)
if get_pending is None:
return {}
try:
pending = await get_pending()
except Exception: # pragma: no cover - defensive for internal API drift
logger.warning("Could not read pending workflow requests", exc_info=True)
return {}
if isinstance(pending, dict):
return cast(dict[str, Any], pending)
return {}
def _interrupt_entry_for_request_event(request_event: Any) -> dict[str, Any] | None:
"""Build AG-UI interrupt payload from a workflow request_info event."""
request_id = getattr(request_event, "request_id", None)
if request_id is None:
return None
request_data = make_json_safe(getattr(request_event, "data", None))
if isinstance(request_data, dict):
value: Any = request_data
else:
value = {"data": request_data}
return {"id": str(request_id), "value": value}
def _interrupts_from_pending_requests(pending_events: dict[str, Any]) -> list[dict[str, Any]]:
"""Convert pending workflow request events into AG-UI interrupt descriptors."""
interrupts: list[dict[str, Any]] = []
for request_event in pending_events.values():
entry = _interrupt_entry_for_request_event(request_event)
if entry is not None:
interrupts.append(entry)
return interrupts
def _request_payload_from_request_event(request_event: Any) -> dict[str, Any] | None:
"""Build the normalized request_info payload from a workflow request event."""
request_id = getattr(request_event, "request_id", None)
if not request_id:
return None
request_type = getattr(request_event, "request_type", None)
response_type = getattr(request_event, "response_type", None)
request_data = make_json_safe(getattr(request_event, "data", None))
return {
"request_id": request_id,
"source_executor_id": getattr(request_event, "source_executor_id", None),
"request_type": getattr(request_type, "__name__", str(request_type) if request_type else None),
"response_type": getattr(response_type, "__name__", str(response_type) if response_type else None),
"data": request_data,
}
def _extract_responses_from_messages(messages: list[Message]) -> dict[str, Any]:
"""Extract request-info responses from incoming tool/function-result messages."""
responses: dict[str, Any] = {}
for message in messages:
for content in message.contents:
if content.type != "function_result" or not content.call_id:
continue
value = _coerce_json_value(content.result)
responses[str(content.call_id)] = value
return responses
def _resume_to_workflow_responses(resume_payload: Any) -> dict[str, Any]:
"""Convert AG-UI resume payloads into workflow responses."""
responses: dict[str, Any] = {}
for interrupt in _normalize_resume_interrupts(resume_payload):
value = _coerce_json_value(interrupt.get("value"))
responses[str(interrupt["id"])] = value
return responses
def _coerce_json_value(value: Any) -> Any:
"""Parse JSON strings when possible; otherwise return the original value."""
if not isinstance(value, str):
return value
stripped = value.strip()
if not stripped:
return value
try:
return json.loads(stripped)
except json.JSONDecodeError:
return value
def _response_type_name(request_event: Any) -> str:
"""Return a stable string name for a request's expected response type."""
response_type = getattr(request_event, "response_type", None)
if response_type is None:
return "unknown"
return getattr(response_type, "__name__", str(response_type))
def _coerce_content(value: Any) -> Content | None:
"""Best-effort conversion of JSON-like payloads into Content."""
if isinstance(value, Content):
return value
candidate = _coerce_json_value(value)
if not isinstance(candidate, dict):
return None
content_payload = dict(candidate)
if "type" not in content_payload and {"approved", "id", "function_call"}.issubset(content_payload):
content_payload["type"] = "function_approval_response"
try:
return Content.from_dict(content_payload)
except Exception:
return None
def _coerce_message_content(content_payload: Any) -> Content | None:
"""Best-effort conversion of AG-UI message content items into Content."""
if isinstance(content_payload, Content):
return content_payload
if isinstance(content_payload, str):
return Content.from_text(text=content_payload)
if isinstance(content_payload, dict):
content_dict = dict(content_payload)
if content_dict.get("type") == "text":
if isinstance(content_dict.get("text"), str):
return Content.from_text(text=cast(str, content_dict["text"]))
if isinstance(content_dict.get("content"), str):
return Content.from_text(text=cast(str, content_dict["content"]))
try:
return Content.from_dict(content_dict)
except Exception:
return None
return None
def _coerce_message(value: Any) -> Message | None:
"""Best-effort conversion of JSON-like payloads into Message."""
if isinstance(value, Message):
return value
candidate = _coerce_json_value(value)
if isinstance(candidate, str):
return Message(role="user", contents=[Content.from_text(text=candidate)])
if not isinstance(candidate, dict):
return None
role = str(candidate.get("role") or "user")
author_name = candidate.get("author_name") or candidate.get("authorName")
message_id = candidate.get("message_id") or candidate.get("messageId")
contents_payload = candidate.get("contents")
if contents_payload is None and "content" in candidate:
contents_payload = candidate.get("content")
normalized_contents: list[Content] = []
if isinstance(contents_payload, list):
for item in contents_payload:
parsed_content = _coerce_message_content(item)
if parsed_content is None:
return None
normalized_contents.append(parsed_content)
elif contents_payload is not None:
parsed_content = _coerce_message_content(contents_payload)
if parsed_content is None:
return None
normalized_contents.append(parsed_content)
else:
normalized_contents.append(Content.from_text(text=""))
return Message(
role=role,
contents=normalized_contents,
author_name=str(author_name) if isinstance(author_name, str) else None,
message_id=str(message_id) if isinstance(message_id, str) else None,
)
def _coerce_response_for_request(request_event: Any, value: Any) -> Any | None:
"""Coerce a candidate value into the request's expected response type."""
response_type = getattr(request_event, "response_type", None)
candidate = _coerce_json_value(value)
if response_type is None:
return candidate
target_type = get_origin(response_type) or response_type
if target_type is Any:
return candidate
if target_type is dict:
return candidate if isinstance(candidate, dict) else None
if target_type is list:
if not isinstance(candidate, list):
return None
item_types = get_args(response_type)
if not item_types:
return candidate
item_type = get_origin(item_types[0]) or item_types[0]
if item_type is Message:
converted_messages: list[Message] = []
for item in candidate:
message = _coerce_message(item)
if message is None:
return None
converted_messages.append(message)
return converted_messages
if item_type is Content:
converted_contents: list[Content] = []
for item in candidate:
content = _coerce_content(item)
if content is None:
return None
converted_contents.append(content)
return converted_contents
return candidate
if target_type is str:
if isinstance(value, str):
return value
if isinstance(candidate, str):
return candidate
return json.dumps(make_json_safe(candidate))
if target_type is Message:
return _coerce_message(candidate)
if target_type is Content:
return _coerce_content(candidate)
if target_type is bool:
return candidate if isinstance(candidate, bool) else None
if target_type is int:
return candidate if isinstance(candidate, int) and not isinstance(candidate, bool) else None
if target_type is float:
return candidate if isinstance(candidate, (int, float)) and not isinstance(candidate, bool) else None
if isinstance(target_type, type):
return candidate if isinstance(candidate, target_type) else None
# Unknown typing metadata: preserve value as-is.
return candidate
def _single_pending_response_from_value(pending_events: dict[str, Any], value: Any) -> dict[str, Any]:
"""Map a scalar resume payload to the single pending request (if unambiguous)."""
if value is None or len(pending_events) != 1:
return {}
request_event = next(iter(pending_events.values()))
request_id = getattr(request_event, "request_id", None)
if not request_id:
return {}
coerced_value = _coerce_response_for_request(request_event, value)
if coerced_value is None:
logger.info(
"Ignoring pending request response for request_id=%s: expected %s",
request_id,
_response_type_name(request_event),
)
return {}
return {str(request_id): coerced_value}
def _coerce_responses_for_pending_requests(
responses: dict[str, Any],
pending_events: dict[str, Any],
) -> dict[str, Any]:
"""Coerce resume responses to the expected types for known pending requests."""
if not responses or not pending_events:
return responses
normalized: dict[str, Any] = {}
pending_by_id = {str(request_id): event for request_id, event in pending_events.items()}
for request_id, value in responses.items():
request_key = str(request_id)
request_event = pending_by_id.get(request_key)
if request_event is None:
normalized[request_key] = value
continue
coerced_value = _coerce_response_for_request(request_event, value)
if coerced_value is None:
logger.info(
"Ignoring resume response for request_id=%s: expected %s",
request_key,
_response_type_name(request_event),
)
continue
normalized[request_key] = coerced_value
return normalized
def _latest_user_text(messages: list[Message]) -> str | None:
"""Get the most recent user text message, if present."""
for message in reversed(messages):
role_field = message.role
if isinstance(role_field, str):
role = role_field
else:
role = str(getattr(role_field, "value", role_field))
if role != "user":
continue
for content in reversed(message.contents):
if content.type != "text":
continue
text_value = getattr(content, "text", None)
if isinstance(text_value, str) and text_value.strip():
return text_value
return None
def _workflow_interrupt_event_value(request_payload: dict[str, Any]) -> str | None:
"""Build a string payload for interrupt-card custom events."""
request_data = request_payload.get("data")
if request_data is None:
return None
if isinstance(request_data, str):
return request_data
return json.dumps(make_json_safe(request_data))
def _message_role_value(message: Message) -> str:
"""Normalize Message.role to its string value."""
role = message.role
if isinstance(role, str):
return role
return str(getattr(role, "value", role))
def _latest_assistant_contents(messages: list[Message]) -> list[Content] | None:
"""Return contents from the most recent assistant message."""
for message in reversed(messages):
if _message_role_value(message) != "assistant":
continue
contents = list(message.contents or [])
if contents:
return contents
return None
def _text_from_contents(contents: list[Content]) -> str | None:
"""Return normalized assistant text from a content list when present."""
text_parts: list[str] = []
for content in contents:
if content.type != "text":
continue
text_value = getattr(content, "text", None)
if not isinstance(text_value, str):
continue
if not text_value:
continue
text_parts.append(text_value)
if not text_parts:
return None
return "".join(text_parts).strip() or None
def _workflow_payload_to_contents(payload: Any) -> list[Content] | None:
"""Best-effort conversion from workflow payloads to chat content fragments."""
if payload is None:
return None
if isinstance(payload, Content):
return [payload]
if isinstance(payload, str):
return [Content.from_text(text=payload)]
if isinstance(payload, Message):
if _message_role_value(payload) != "assistant":
return None
return list(payload.contents or [])
if isinstance(payload, AgentResponseUpdate):
role_field = payload.role
if role_field is None:
return None
if isinstance(role_field, str):
role = role_field
else:
role = str(getattr(role_field, "value", role_field))
if role != "assistant":
return None
return list(payload.contents or [])
if isinstance(payload, AgentResponse):
return _latest_assistant_contents(list(payload.messages or []))
if isinstance(payload, list):
if payload and all(isinstance(item, Message) for item in payload):
return _latest_assistant_contents(cast(list[Message], payload))
contents: list[Content] = []
for item in payload:
item_contents = _workflow_payload_to_contents(item)
if item_contents is None:
return None
contents.extend(item_contents)
return contents if contents else None
return None
def _event_name(event: Any) -> str:
event_type = getattr(event, "type", None)
if isinstance(event_type, str) and event_type:
return event_type
return type(event).__name__
def _custom_event_value(event: Any) -> Any:
if getattr(event, "data", None) is not None:
return make_json_safe(getattr(event, "data"))
event_dict = cast(dict[str, Any], getattr(event, "__dict__", {}) or {})
custom_fields = {
key: make_json_safe(value)
for key, value in event_dict.items()
if key not in _WORKFLOW_EVENT_BASE_FIELDS and not key.startswith("_")
}
return custom_fields if custom_fields else None
def _details_message(details: Any) -> str:
if details is None:
return "Workflow execution failed."
if hasattr(details, "message"):
message = getattr(details, "message")
if isinstance(message, str) and message:
return message
return str(details)
def _details_code(details: Any) -> str | None:
if details is None:
return None
if hasattr(details, "error_type"):
error_type = getattr(details, "error_type")
if isinstance(error_type, str) and error_type:
return error_type
return None
async def run_workflow_stream(
input_data: dict[str, Any],
workflow: Workflow,
) -> AsyncGenerator[BaseEvent]:
"""Run a Workflow and emit AG-UI protocol events."""
thread_id = input_data.get("thread_id") or input_data.get("threadId") or str(uuid.uuid4())
run_id = input_data.get("run_id") or input_data.get("runId") or str(uuid.uuid4())
available_interrupts = input_data.get("available_interrupts") or input_data.get("availableInterrupts")
if available_interrupts:
logger.debug("Received available interrupts metadata: %s", available_interrupts)
raw_messages = list(cast(list[dict[str, Any]], input_data.get("messages", []) or []))
messages, _ = normalize_agui_input_messages(raw_messages, sanitize_tool_history=False)
flow = FlowState()
interrupts: list[dict[str, Any]] = []
run_started_emitted = False
terminal_emitted = False
run_error_emitted = False
last_assistant_text: str | None = None
resume_payload = _extract_resume_payload(input_data)
responses = _resume_to_workflow_responses(resume_payload)
responses.update(_extract_responses_from_messages(messages))
pending_before_run = await _pending_request_events(workflow)
responses = _coerce_responses_for_pending_requests(responses, pending_before_run)
pending_interrupts = _interrupts_from_pending_requests(pending_before_run)
if not responses and pending_before_run:
responses.update(_single_pending_response_from_value(pending_before_run, resume_payload))
if not responses and pending_before_run:
responses.update(_single_pending_response_from_value(pending_before_run, _latest_user_text(messages)))
if not responses and pending_before_run:
yield RunStartedEvent(run_id=run_id, thread_id=thread_id)
for request_event in pending_before_run.values():
request_payload = _request_payload_from_request_event(request_event)
if request_payload is None:
continue
request_id = str(request_payload["request_id"])
yield ToolCallStartEvent(tool_call_id=request_id, tool_call_name="request_info")
yield ToolCallArgsEvent(tool_call_id=request_id, delta=json.dumps(request_payload))
yield ToolCallEndEvent(tool_call_id=request_id)
yield CustomEvent(name="request_info", value=request_payload)
interrupt_event_value = _workflow_interrupt_event_value(request_payload)
if interrupt_event_value is not None:
yield CustomEvent(name=_INTERRUPT_CARD_EVENT_NAME, value=interrupt_event_value)
yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=pending_interrupts)
return
if not responses and not messages:
yield RunStartedEvent(run_id=run_id, thread_id=thread_id)
yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=pending_interrupts)
return
def _drain_open_message() -> list[TextMessageEndEvent]:
"""Close any open assistant text message and clear flow state."""
if not flow.message_id:
return []
current_message_id = flow.message_id
flow.message_id = None
flow.accumulated_text = ""
return [TextMessageEndEvent(message_id=current_message_id)]
try:
if responses:
event_stream = workflow.run(responses=responses, stream=True)
else:
event_stream = workflow.run(message=messages, stream=True)
async for event in event_stream:
event_type = getattr(event, "type", None)
if event_type == "started":
if not run_started_emitted:
yield RunStartedEvent(run_id=run_id, thread_id=thread_id)
run_started_emitted = True
continue
if not run_started_emitted:
yield RunStartedEvent(run_id=run_id, thread_id=thread_id)
run_started_emitted = True
if event_type == "failed":
details = getattr(event, "details", None)
yield RunErrorEvent(message=_details_message(details), code=_details_code(details))
run_error_emitted = True
terminal_emitted = True
continue
if event_type == "status":
state = getattr(event, "state", None)
if isinstance(state, str):
state_value = state
else:
state_value = str(getattr(state, "value", state))
if state_value in _TERMINAL_STATES and not terminal_emitted:
if not interrupts:
interrupts.extend(_interrupts_from_pending_requests(await _pending_request_events(workflow)))
yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=interrupts)
terminal_emitted = True
elif state_value not in _TERMINAL_STATES:
yield CustomEvent(name="status", value={"state": state_value})
continue
if event_type == "superstep_started":
for end_event in _drain_open_message():
yield end_event
iteration = getattr(event, "iteration", None)
yield StepStartedEvent(step_name=f"superstep:{iteration}")
continue
if event_type == "superstep_completed":
iteration = getattr(event, "iteration", None)
yield StepFinishedEvent(step_name=f"superstep:{iteration}")
continue
if event_type in {"executor_invoked", "executor_completed", "executor_failed"}:
executor_id = getattr(event, "executor_id", None)
status = {
"executor_invoked": "in_progress",
"executor_completed": "completed",
"executor_failed": "failed",
}[event_type]
if isinstance(executor_id, str) and executor_id:
if event_type == "executor_invoked":
for end_event in _drain_open_message():
yield end_event
yield StepStartedEvent(step_name=executor_id)
else:
yield StepFinishedEvent(step_name=executor_id)
executor_payload: dict[str, Any] = {
"executor_id": executor_id,
"status": status,
}
if event_type == "executor_failed":
executor_payload["details"] = make_json_safe(getattr(event, "details", None))
else:
executor_payload["data"] = make_json_safe(getattr(event, "data", None))
yield ActivitySnapshotEvent(
message_id=f"executor:{executor_id}" if executor_id else generate_event_id(),
activity_type="executor",
content=executor_payload,
)
continue
if event_type == "request_info":
for end_event in _drain_open_message():
yield end_event
request_payload = _request_payload_from_request_event(event)
if request_payload is None:
continue
request_id = request_payload["request_id"]
request_data = request_payload.get("data")
if isinstance(request_data, dict):
interrupt_value: Any = request_data
else:
interrupt_value = {"data": request_data}
interrupts.append({"id": str(request_id), "value": interrupt_value})
args_delta = json.dumps(request_payload)
yield ToolCallStartEvent(tool_call_id=str(request_id), tool_call_name="request_info")
yield ToolCallArgsEvent(tool_call_id=str(request_id), delta=args_delta)
yield ToolCallEndEvent(tool_call_id=str(request_id))
yield CustomEvent(name="request_info", value=request_payload)
interrupt_event_value = _workflow_interrupt_event_value(request_payload)
if interrupt_event_value is not None:
yield CustomEvent(name=_INTERRUPT_CARD_EVENT_NAME, value=interrupt_event_value)
continue
if event_type in {"output", "data"}:
output_payload = getattr(event, "data", None)
if isinstance(output_payload, BaseEvent):
yield output_payload
continue
if (
isinstance(output_payload, list)
and output_payload
and all(isinstance(item, BaseEvent) for item in output_payload)
):
for item in output_payload:
yield item
continue
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
for content in contents:
for out_event in _emit_content(content, flow, predictive_handler=None, skip_text=False):
yield out_event
if flow.message_id and flow.accumulated_text:
last_assistant_text = flow.accumulated_text.strip() or last_assistant_text
elif output_text:
last_assistant_text = output_text
else:
yield CustomEvent(name="workflow_output", value=make_json_safe(output_payload))
continue
# Fall back to custom events for diagnostics, orchestration events, and custom workflow events.
yield CustomEvent(name=_event_name(event), value=_custom_event_value(event))
except Exception as exc:
logger.exception("Workflow AG-UI stream failed: %s", exc)
if not run_started_emitted:
yield RunStartedEvent(run_id=run_id, thread_id=thread_id)
run_started_emitted = True
if not run_error_emitted:
yield RunErrorEvent(message=str(exc), code=type(exc).__name__)
run_error_emitted = True
terminal_emitted = True
for end_event in _drain_open_message():
yield end_event
if not run_started_emitted:
yield RunStartedEvent(run_id=run_id, thread_id=thread_id)
if not terminal_emitted and not run_error_emitted:
if not interrupts:
interrupts.extend(_interrupts_from_pending_requests(await _pending_request_events(workflow)))
yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=interrupts)
python/packages/ag-ui/agent_framework_ag_ui_examples/README.md
+4
View File
@@ -85,6 +85,7 @@ Complete examples for all AG-UI features are available:
- `document_writer_agent(client)` - Predictive state updates (Feature 7)
- `research_assistant_agent(client)` - Research with progress events
- `task_planner_agent(client)` - Task planning with approvals
- `subgraphs_agent()` - Deterministic travel-planning subgraphs flow (Dojo `subgraphs` feature)
### Using Example Agents
@@ -130,6 +131,7 @@ The server exposes endpoints at:
- `/tool_based_generative_ui` - Custom UI components with `ui_generator_agent`
- `/shared_state` - Recipe management with `recipe_agent`
- `/predictive_state_updates` - Document writing with `document_writer_agent`
- `/subgraphs` - Travel planner with interrupt-driven flight/hotel choices via `subgraphs_agent`
### Complete FastAPI Example
@@ -145,6 +147,7 @@ from agent_framework_ag_ui_examples.agents import (
ui_generator_agent,
recipe_agent,
document_writer_agent,
subgraphs_agent,
)
app = FastAPI(title="AG-UI Examples")
@@ -160,6 +163,7 @@ add_agent_framework_fastapi_endpoint(app, task_steps_agent_wrapped(client), "/ag
add_agent_framework_fastapi_endpoint(app, ui_generator_agent(client), "/tool_based_generative_ui")
add_agent_framework_fastapi_endpoint(app, recipe_agent(client), "/shared_state")
add_agent_framework_fastapi_endpoint(app, document_writer_agent(client), "/predictive_state_updates")
add_agent_framework_fastapi_endpoint(app, subgraphs_agent(), "/subgraphs")
```
## Architecture
python/packages/ag-ui/agent_framework_ag_ui_examples/agents/__init__.py
+2
View File
@@ -7,6 +7,7 @@ from .human_in_the_loop_agent import human_in_the_loop_agent
from .recipe_agent import recipe_agent
from .research_assistant_agent import research_assistant_agent
from .simple_agent import simple_agent
from .subgraphs_agent import subgraphs_agent
from .task_planner_agent import task_planner_agent
from .task_steps_agent import task_steps_agent_wrapped
from .ui_generator_agent import ui_generator_agent
@@ -18,6 +19,7 @@ __all__ = [
"recipe_agent",
"research_assistant_agent",
"simple_agent",
"subgraphs_agent",
"task_planner_agent",
"task_steps_agent_wrapped",
"ui_generator_agent",
python/packages/ag-ui/agent_framework_ag_ui_examples/agents/subgraphs_agent.py
+405
View File
@@ -0,0 +1,405 @@
# Copyright (c) Microsoft. All rights reserved.
"""Subgraphs travel planner built with MAF workflow primitives."""
import json
import uuid
from copy import deepcopy
from dataclasses import dataclass
from typing import Any
from ag_ui.core import (
BaseEvent,
StateSnapshotEvent,
TextMessageContentEvent,
TextMessageEndEvent,
TextMessageStartEvent,
)
from agent_framework import (
Executor,
Message,
Workflow,
WorkflowBuilder,
WorkflowContext,
handler,
response_handler,
)
from agent_framework_ag_ui import AgentFrameworkWorkflow
STATIC_FLIGHTS: list[dict[str, str]] = [
{
"airline": "KLM",
"departure": "Amsterdam (AMS)",
"arrival": "San Francisco (SFO)",
"price": "$650",
"duration": "11h 30m",
},
{
"airline": "United",
"departure": "Amsterdam (AMS)",
"arrival": "San Francisco (SFO)",
"price": "$720",
"duration": "12h 15m",
},
]
STATIC_HOTELS: list[dict[str, str]] = [
{
"name": "Hotel Zephyr",
"location": "Fisherman's Wharf",
"price_per_night": "$280/night",
"rating": "4.2 stars",
},
{
"name": "The Ritz-Carlton",
"location": "Nob Hill",
"price_per_night": "$550/night",
"rating": "4.8 stars",
},
{
"name": "Hotel Zoe",
"location": "Union Square",
"price_per_night": "$320/night",
"rating": "4.4 stars",
},
]
STATIC_EXPERIENCES: list[dict[str, str]] = [
{
"name": "Pier 39",
"type": "activity",
"description": "Iconic waterfront destination with shops and sea lions",
"location": "Fisherman's Wharf",
},
{
"name": "Golden Gate Bridge",
"type": "activity",
"description": "World-famous suspension bridge with stunning views",
"location": "Golden Gate",
},
{
"name": "Swan Oyster Depot",
"type": "restaurant",
"description": "Historic seafood counter serving fresh oysters",
"location": "Polk Street",
},
{
"name": "Tartine Bakery",
"type": "restaurant",
"description": "Artisanal bakery famous for bread and pastries",
"location": "Mission District",
},
]
_STATE_KEY = "subgraphs_state"
@dataclass
class _PresentFlights:
pass
@dataclass
class _PresentHotels:
pass
@dataclass
class _PlanExperiences:
pass
@dataclass
class _FinalizeTrip:
pass
def _initial_state() -> dict[str, Any]:
return {
"itinerary": {},
"experiences": [],
"flights": [],
"hotels": [],
"planning_step": "start",
"active_agent": "supervisor",
}
def _emit_text_events(text: str) -> list[BaseEvent]:
message_id = str(uuid.uuid4())
return [
TextMessageStartEvent(message_id=message_id, role="assistant"),
TextMessageContentEvent(message_id=message_id, delta=text),
TextMessageEndEvent(message_id=message_id),
]
async def _emit_text(ctx: WorkflowContext[Any, BaseEvent], text: str) -> None:
for event in _emit_text_events(text):
await ctx.yield_output(event)
async def _emit_state_snapshot(ctx: WorkflowContext[Any, BaseEvent], state: dict[str, Any]) -> None:
await ctx.yield_output(StateSnapshotEvent(snapshot=deepcopy(state)))
def _flight_interrupt_value() -> dict[str, Any]:
return {
"message": "Choose the flight you want. I recommend KLM because it is cheaper and usually on time.",
"options": deepcopy(STATIC_FLIGHTS),
"recommendation": deepcopy(STATIC_FLIGHTS[0]),
"agent": "flights",
}
def _hotel_interrupt_value() -> dict[str, Any]:
return {
"message": "Choose your hotel. I recommend Hotel Zoe for the best value in a central location.",
"options": deepcopy(STATIC_HOTELS),
"recommendation": deepcopy(STATIC_HOTELS[2]),
"agent": "hotels",
}
def _normalize_flight(value: Any) -> dict[str, str] | None:
if isinstance(value, str):
try:
value = json.loads(value)
except json.JSONDecodeError:
return None
if isinstance(value, dict) and value.get("airline"):
return {
"airline": str(value.get("airline", "")),
"departure": str(value.get("departure", "")),
"arrival": str(value.get("arrival", "")),
"price": str(value.get("price", "")),
"duration": str(value.get("duration", "")),
}
return None
def _normalize_hotel(value: Any) -> dict[str, str] | None:
if isinstance(value, str):
try:
value = json.loads(value)
except json.JSONDecodeError:
return None
if isinstance(value, dict) and value.get("name"):
return {
"name": str(value.get("name", "")),
"location": str(value.get("location", "")),
"price_per_night": str(value.get("price_per_night", "")),
"rating": str(value.get("rating", "")),
}
return None
def _load_state(ctx: WorkflowContext[Any, BaseEvent]) -> dict[str, Any]:
state = ctx.get_state(_STATE_KEY)
if isinstance(state, dict):
return state
new_state = _initial_state()
ctx.set_state(_STATE_KEY, new_state)
return new_state
class _SupervisorExecutor(Executor):
def __init__(self) -> None:
super().__init__(id="supervisor_agent")
@handler
async def start(self, message: list[Message], ctx: WorkflowContext[_PresentFlights, BaseEvent]) -> None:
del message
state = _initial_state()
ctx.set_state(_STATE_KEY, state)
await _emit_state_snapshot(ctx, state)
await _emit_text(
ctx,
"Supervisor: I will coordinate our specialist agents to plan your San Francisco trip end to end.",
)
state["active_agent"] = "flights"
state["planning_step"] = "collecting_flights"
state["flights"] = deepcopy(STATIC_FLIGHTS)
ctx.set_state(_STATE_KEY, state)
await _emit_state_snapshot(ctx, state)
await ctx.send_message(_PresentFlights(), target_id="flights_agent")
@handler
async def finalize(self, message: _FinalizeTrip, ctx: WorkflowContext[Any, BaseEvent]) -> None:
del message
state = _load_state(ctx)
state["active_agent"] = "supervisor"
state["planning_step"] = "complete"
ctx.set_state(_STATE_KEY, state)
await _emit_state_snapshot(ctx, state)
await _emit_text(ctx, "Supervisor: Your travel planning is complete and your itinerary is ready.")
class _FlightsExecutor(Executor):
def __init__(self) -> None:
super().__init__(id="flights_agent")
@handler
async def present_options(self, message: _PresentFlights, ctx: WorkflowContext[_PresentHotels, BaseEvent]) -> None:
del message
await _emit_text(
ctx,
"Flights Agent: I found two flight options from Amsterdam to San Francisco. "
"KLM is recommended for the best value and schedule.",
)
await ctx.request_info(_flight_interrupt_value(), dict, request_id="flights-choice")
@response_handler
async def handle_selection(
self,
original_request: dict,
response: dict,
ctx: WorkflowContext[_PresentHotels, BaseEvent],
) -> None:
del original_request
state = _load_state(ctx)
selected_flight = _normalize_flight(response)
if selected_flight is None:
state["active_agent"] = "flights"
state["planning_step"] = "collecting_flights"
state["flights"] = deepcopy(STATIC_FLIGHTS)
ctx.set_state(_STATE_KEY, state)
await _emit_state_snapshot(ctx, state)
await _emit_text(ctx, "Flights Agent: Please choose a flight option from the selection card to continue.")
await ctx.request_info(_flight_interrupt_value(), dict, request_id="flights-choice")
return
itinerary = state.setdefault("itinerary", {})
itinerary["flight"] = selected_flight
state["active_agent"] = "flights"
state["planning_step"] = "booking_flight"
ctx.set_state(_STATE_KEY, state)
await _emit_state_snapshot(ctx, state)
await _emit_text(
ctx,
f"Flights Agent: Great choice. I will book the {selected_flight['airline']} flight. "
"Now I am routing you to Hotels Agent for accommodation.",
)
state["active_agent"] = "hotels"
state["planning_step"] = "collecting_hotels"
state["hotels"] = deepcopy(STATIC_HOTELS)
ctx.set_state(_STATE_KEY, state)
await _emit_state_snapshot(ctx, state)
await ctx.send_message(_PresentHotels(), target_id="hotels_agent")
class _HotelsExecutor(Executor):
def __init__(self) -> None:
super().__init__(id="hotels_agent")
@handler
async def present_options(self, message: _PresentHotels, ctx: WorkflowContext[_PlanExperiences, BaseEvent]) -> None:
del message
await _emit_text(
ctx,
"Hotels Agent: I found three accommodation options in San Francisco. "
"Hotel Zoe is recommended for the best balance of location, quality, and price.",
)
await ctx.request_info(_hotel_interrupt_value(), dict, request_id="hotels-choice")
@response_handler
async def handle_selection(
self,
original_request: dict,
response: dict,
ctx: WorkflowContext[_PlanExperiences, BaseEvent],
) -> None:
del original_request
state = _load_state(ctx)
selected_hotel = _normalize_hotel(response)
if selected_hotel is None:
state["active_agent"] = "hotels"
state["planning_step"] = "collecting_hotels"
state["hotels"] = deepcopy(STATIC_HOTELS)
ctx.set_state(_STATE_KEY, state)
await _emit_state_snapshot(ctx, state)
await _emit_text(ctx, "Hotels Agent: Please choose a hotel option from the selection card to continue.")
await ctx.request_info(_hotel_interrupt_value(), dict, request_id="hotels-choice")
return
itinerary = state.setdefault("itinerary", {})
itinerary["hotel"] = selected_hotel
state["active_agent"] = "hotels"
state["planning_step"] = "booking_hotel"
ctx.set_state(_STATE_KEY, state)
await _emit_state_snapshot(ctx, state)
await _emit_text(
ctx,
f"Hotels Agent: Excellent, {selected_hotel['name']} is booked. "
"I am routing you to Experiences Agent for activities and restaurants.",
)
state["active_agent"] = "experiences"
state["planning_step"] = "curating_experiences"
state["experiences"] = deepcopy(STATIC_EXPERIENCES)
ctx.set_state(_STATE_KEY, state)
await _emit_state_snapshot(ctx, state)
await ctx.send_message(_PlanExperiences(), target_id="experiences_agent")
class _ExperiencesExecutor(Executor):
def __init__(self) -> None:
super().__init__(id="experiences_agent")
@handler
async def plan(self, message: _PlanExperiences, ctx: WorkflowContext[_FinalizeTrip, BaseEvent]) -> None:
del message
await _emit_text(
ctx,
"Experiences Agent: I planned activities and restaurants including "
"Pier 39, Golden Gate Bridge, Swan Oyster Depot, and Tartine Bakery.",
)
await ctx.send_message(_FinalizeTrip(), target_id="supervisor_agent")
def _build_subgraphs_workflow() -> Workflow:
supervisor = _SupervisorExecutor()
flights = _FlightsExecutor()
hotels = _HotelsExecutor()
experiences = _ExperiencesExecutor()
return (
WorkflowBuilder(
name="subgraphs",
description="Travel planning supervisor with flights/hotels/experiences subgraphs.",
start_executor=supervisor,
)
.add_edge(supervisor, flights)
.add_edge(flights, hotels)
.add_edge(hotels, experiences)
.add_edge(experiences, supervisor)
.build()
)
def _build_subgraphs_workflow_for_thread(thread_id: str) -> Workflow:
"""Create a workflow instance scoped to a single AG-UI thread."""
del thread_id
return _build_subgraphs_workflow()
def subgraphs_agent() -> AgentFrameworkWorkflow:
"""Create the subgraphs travel planner agent."""
return AgentFrameworkWorkflow(
workflow_factory=_build_subgraphs_workflow_for_thread,
name="subgraphs",
description="Travel planning workflow with interrupt-driven selections.",
)
python/packages/ag-ui/agent_framework_ag_ui_examples/agents/task_steps_agent.py
+8 -15
View File
@@ -24,6 +24,8 @@ from agent_framework import Agent, Content, Message, SupportsChatGetResponse, to
from agent_framework.ag_ui import AgentFrameworkAgent
from pydantic import BaseModel, Field
from agent_framework_ag_ui import AgentFrameworkWorkflow
class StepStatus(str, Enum):
"""Status of a task step."""
@@ -105,38 +107,29 @@ def _create_task_steps_agent(client: SupportsChatGetResponse[Any]) -> AgentFrame
# Wrap the agent's run method to add step execution simulation
class TaskStepsAgentWithExecution:
class TaskStepsAgentWithExecution(AgentFrameworkWorkflow):
"""Wrapper that adds step execution simulation after plan generation.
This wrapper delegates to AgentFrameworkAgent but is recognized as compatible
by add_agent_framework_fastapi_endpoint since it implements run_agent().
by add_agent_framework_fastapi_endpoint since it implements run().
"""
def __init__(self, base_agent: AgentFrameworkAgent):
"""Initialize wrapper with base agent."""
super().__init__(name=base_agent.name, description=base_agent.description)
self._base_agent = base_agent
@property
def name(self) -> str:
"""Delegate to base agent."""
return self._base_agent.name
@property
def description(self) -> str:
"""Delegate to base agent."""
return self._base_agent.description
def __getattr__(self, name: str) -> Any:
"""Delegate all other attribute access to base agent."""
return getattr(self._base_agent, name)
async def run_agent(self, input_data: dict[str, Any]) -> AsyncGenerator[Any]:
async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[Any]:
"""Run the agent and then simulate step execution."""
import logging
import uuid
logger = logging.getLogger(__name__)
logger.info("TaskStepsAgentWithExecution.run_agent() called - wrapper is active")
logger.info("TaskStepsAgentWithExecution.run() called - wrapper is active")
# First, run the base agent to generate the plan - buffer text messages
final_state: dict[str, Any] = {}
@@ -144,7 +137,7 @@ class TaskStepsAgentWithExecution:
tool_call_id: str | None = None
buffered_text_events: list[Any] = [] # Buffer text from first LLM call
async for event in self._base_agent.run_agent(input_data):
async for event in self._base_agent.run(input_data):
event_type_str = str(event.type) if hasattr(event, "type") else type(event).__name__
logger.info(f"Processing event: {event_type_str}")
python/packages/ag-ui/agent_framework_ag_ui_examples/server/main.py
+8
View File
@@ -21,6 +21,7 @@ from ..agents.document_writer_agent import document_writer_agent
from ..agents.human_in_the_loop_agent import human_in_the_loop_agent
from ..agents.recipe_agent import recipe_agent
from ..agents.simple_agent import simple_agent
from ..agents.subgraphs_agent import subgraphs_agent
from ..agents.task_steps_agent import task_steps_agent_wrapped
from ..agents.ui_generator_agent import ui_generator_agent
from ..agents.weather_agent import weather_agent
@@ -123,6 +124,13 @@ add_agent_framework_fastapi_endpoint(
path="/tool_based_generative_ui",
)
# Subgraphs - deterministic travel planner with interrupt-driven selections
add_agent_framework_fastapi_endpoint(
app=app,
agent=subgraphs_agent(),
path="/subgraphs",
)
def main():
"""Run the server."""
python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py
+28
View File
@@ -356,3 +356,31 @@ class TestAGUIChatClient:
response = await client.inner_get_response(messages=messages, options=chat_options)
assert response is not None
async def test_interrupt_options_transmission(self, monkeypatch: MonkeyPatch) -> None:
"""Interrupt option fields are forwarded to the HTTP service."""
available_interrupts = [{"id": "req_1", "type": "request_info"}]
resume_payload = {"interrupts": [{"id": "req_1", "value": "approved"}]}
mock_events = [
{"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
{"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
]
async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
assert kwargs.get("available_interrupts") == available_interrupts
assert kwargs.get("resume") == resume_payload
for event in mock_events:
yield event
client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
messages = [Message(role="user", text="continue")]
options = {
"available_interrupts": available_interrupts,
"resume": resume_payload,
}
response = await client.inner_get_response(messages=messages, options=options)
assert response is not None
python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py
+97 -20
View File
@@ -103,7 +103,7 @@ async def test_run_started_event_emission(streaming_chat_client_stub):
input_data = {"messages": [{"role": "user", "content": "Hi"}]}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
# First event should be RunStartedEvent
@@ -131,7 +131,7 @@ async def test_predict_state_custom_event_emission(streaming_chat_client_stub):
input_data = {"messages": [{"role": "user", "content": "Hi"}]}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
# Find PredictState event
@@ -144,6 +144,83 @@ async def test_predict_state_custom_event_emission(streaming_chat_client_stub):
assert {"state_key": "summary", "tool": "summarize", "tool_argument": "text"} in predict_value
async def test_usage_content_emits_custom_usage_event(streaming_chat_client_stub):
"""Usage content from the wrapped agent should be surfaced as a custom usage event."""
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
del messages, options, kwargs
yield ChatResponseUpdate(
contents=[
Content.from_usage(
{
"input_token_count": 10,
"output_token_count": 4,
"total_token_count": 14,
}
)
]
)
agent = Agent(name="usage_agent", instructions="Usage test", client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
events: list[Any] = []
async for event in wrapper.run({"messages": [{"role": "user", "content": "Hi"}]}):
events.append(event)
usage_events = [event for event in events if event.type == "CUSTOM" and event.name == "usage"]
assert len(usage_events) == 1
assert usage_events[0].value["input_token_count"] == 10
assert usage_events[0].value["output_token_count"] == 4
assert usage_events[0].value["total_token_count"] == 14
async def test_multimodal_input_is_forwarded_to_agent_run(streaming_chat_client_stub):
"""Multimodal AG-UI input should be converted and passed through to agent.run."""
from agent_framework.ag_ui import AgentFrameworkAgent
captured_messages: list[Message] = []
async def stream_fn(
messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
del options, kwargs
captured_messages[:] = list(messages)
yield ChatResponseUpdate(contents=[Content.from_text(text="Processed multimodal input")])
agent = Agent(name="multimodal_agent", instructions="Multimodal test", client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
input_data = {
"messages": [
{
"role": "user",
"content": [
{"type": "text", "text": "What is in this image?"},
{
"type": "image",
"source": {"type": "url", "url": "https://example.com/cat.png", "mimeType": "image/png"},
},
],
}
]
}
_ = [event async for event in wrapper.run(input_data)]
assert len(captured_messages) == 1
message = captured_messages[0]
assert message.role == "user"
assert len(message.contents) == 2
assert message.contents[0].type == "text"
assert message.contents[0].text == "What is in this image?"
assert message.contents[1].type == "uri"
assert message.contents[1].uri == "https://example.com/cat.png"
async def test_initial_state_snapshot_with_schema(streaming_chat_client_stub):
"""Test initial StateSnapshotEvent emission when state_schema present."""
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -163,7 +240,7 @@ async def test_initial_state_snapshot_with_schema(streaming_chat_client_stub):
}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
# Find StateSnapshotEvent
@@ -190,7 +267,7 @@ async def test_state_initialization_object_type(streaming_chat_client_stub):
input_data = {"messages": [{"role": "user", "content": "Hi"}]}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
# Find StateSnapshotEvent
@@ -217,7 +294,7 @@ async def test_state_initialization_array_type(streaming_chat_client_stub):
input_data = {"messages": [{"role": "user", "content": "Hi"}]}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
# Find StateSnapshotEvent
@@ -243,7 +320,7 @@ async def test_run_finished_event_emission(streaming_chat_client_stub):
input_data = {"messages": [{"role": "user", "content": "Hi"}]}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
# Last event should be RunFinishedEvent
@@ -280,7 +357,7 @@ async def test_tool_result_confirm_changes_accepted(streaming_chat_client_stub):
}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
# Should emit text message confirming acceptance
@@ -322,7 +399,7 @@ async def test_tool_result_confirm_changes_rejected(streaming_chat_client_stub):
}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
# Should emit text message asking what to change
@@ -362,7 +439,7 @@ async def test_tool_result_function_approval_accepted(streaming_chat_client_stub
}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
# Should list enabled steps
@@ -405,7 +482,7 @@ async def test_tool_result_function_approval_rejected(streaming_chat_client_stub
}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
# Should ask what to change about the plan
@@ -441,7 +518,7 @@ async def test_thread_metadata_tracking(streaming_chat_client_stub):
}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
# AG-UI internal metadata should NOT be passed to chat client options
@@ -479,7 +556,7 @@ async def test_state_context_injection(streaming_chat_client_stub):
}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
# Current state should NOT be passed to chat client options
@@ -502,7 +579,7 @@ async def test_no_messages_provided(streaming_chat_client_stub):
input_data: dict[str, Any] = {"messages": []}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
# Should emit RunStartedEvent and RunFinishedEvent only
@@ -526,7 +603,7 @@ async def test_message_end_event_emission(streaming_chat_client_stub):
input_data: dict[str, Any] = {"messages": [{"role": "user", "content": "Hi"}]}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
# Should have TextMessageEndEvent before RunFinishedEvent
@@ -556,7 +633,7 @@ async def test_error_handling_with_exception(streaming_chat_client_stub):
input_data: dict[str, Any] = {"messages": [{"role": "user", "content": "Hi"}]}
with pytest.raises(RuntimeError, match="Simulated failure"):
async for _ in wrapper.run_agent(input_data):
async for _ in wrapper.run(input_data):
pass
@@ -586,7 +663,7 @@ async def test_json_decode_error_in_tool_result(streaming_chat_client_stub):
}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
# Orphaned tool result should be sanitized out
@@ -616,7 +693,7 @@ async def test_agent_with_use_service_session_is_false(streaming_chat_client_stu
input_data = {"messages": [{"role": "user", "content": "Hi"}], "thread_id": "conv_123456"}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
assert request_service_session_id is None # type: ignore[attr-defined] (service_session_id should be set)
@@ -643,7 +720,7 @@ async def test_agent_with_use_service_session_is_true(streaming_chat_client_stub
input_data = {"messages": [{"role": "user", "content": "Hi"}], "thread_id": "conv_123456"}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
request_service_session_id = agent.client.last_service_session_id
assert request_service_session_id == "conv_123456" # type: ignore[attr-defined] (service_session_id should be set)
@@ -714,7 +791,7 @@ async def test_function_approval_mode_executes_tool(streaming_chat_client_stub):
}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
# Verify the run completed successfully
@@ -802,7 +879,7 @@ async def test_function_approval_mode_rejection(streaming_chat_client_stub):
}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
# Verify the run completed
python/packages/ag-ui/tests/ag_ui/test_endpoint.py
+62 -2
View File
@@ -3,9 +3,18 @@
"""Tests for FastAPI endpoint creation (_endpoint.py)."""
import json
from typing import Any
import pytest
from agent_framework import Agent, ChatResponseUpdate, Content
from ag_ui.core import RunStartedEvent
from agent_framework import (
Agent,
ChatResponseUpdate,
Content,
WorkflowBuilder,
WorkflowContext,
executor,
)
from agent_framework.orchestrations import SequentialBuilder
from fastapi import FastAPI, Header, HTTPException
from fastapi.params import Depends
@@ -13,6 +22,7 @@ from fastapi.testclient import TestClient
from agent_framework_ag_ui import add_agent_framework_fastapi_endpoint
from agent_framework_ag_ui._agent import AgentFrameworkAgent
from agent_framework_ag_ui._workflow import AgentFrameworkWorkflow
@pytest.fixture
@@ -55,6 +65,32 @@ async def test_add_endpoint_with_wrapped_agent(build_chat_client):
assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
async def test_add_endpoint_with_workflow_protocol():
"""Test adding endpoint with native Workflow support."""
@executor(id="start")
async def start(message: Any, ctx: WorkflowContext) -> None:
await ctx.yield_output("Workflow response")
app = FastAPI()
workflow = WorkflowBuilder(start_executor=start).build()
add_agent_framework_fastapi_endpoint(app, workflow, path="/workflow")
client = TestClient(app)
response = client.post("/workflow", json={"messages": [{"role": "user", "content": "Hello"}]})
assert response.status_code == 200
assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
content = response.content.decode("utf-8")
lines = [line for line in content.split("\n") if line.startswith("data: ")]
event_types = [json.loads(line[6:]).get("type") for line in lines]
assert "RUN_STARTED" in event_types
assert "TEXT_MESSAGE_CONTENT" in event_types
assert "RUN_FINISHED" in event_types
async def test_endpoint_with_state_schema(build_chat_client):
"""Test endpoint with state_schema parameter."""
app = FastAPI()
@@ -403,8 +439,32 @@ async def test_endpoint_internal_error_handling(build_chat_client):
mock_deepcopy.side_effect = Exception("Simulated internal error")
response = client.post("/error-test", json={"messages": [{"role": "user", "content": "Hello"}]})
assert response.status_code == 500
assert response.json() == {"detail": "An internal error has occurred."}
async def test_endpoint_streaming_error_emits_run_error_event():
"""Streaming exceptions should emit RUN_ERROR instead of terminating silently."""
class FailingStreamWorkflow(AgentFrameworkWorkflow):
async def run(self, input_data: dict[str, Any]):
del input_data
yield RunStartedEvent(run_id="run-1", thread_id="thread-1")
raise RuntimeError("stream exploded")
app = FastAPI()
add_agent_framework_fastapi_endpoint(app, FailingStreamWorkflow(), path="/stream-error")
client = TestClient(app)
response = client.post("/stream-error", json={"messages": [{"role": "user", "content": "Hello"}]})
assert response.status_code == 200
assert response.json() == {"error": "An internal error has occurred."}
content = response.content.decode("utf-8")
lines = [line for line in content.split("\n") if line.startswith("data: ")]
event_types = [json.loads(line[6:]).get("type") for line in lines]
assert "RUN_STARTED" in event_types
assert "RUN_ERROR" in event_types
async def test_endpoint_with_dependencies_blocks_unauthorized(build_chat_client):
python/packages/ag-ui/tests/ag_ui/test_event_converters.py
+51
View File
@@ -207,6 +207,26 @@ class TestAGUIEventConverter:
assert update.additional_properties["thread_id"] == "thread_123"
assert update.additional_properties["run_id"] == "run_456"
def test_run_finished_event_with_interrupt(self) -> None:
"""RUN_FINISHED interrupt metadata is preserved in additional_properties."""
converter = AGUIEventConverter()
converter.thread_id = "thread_123"
converter.run_id = "run_456"
event = {
"type": "RUN_FINISHED",
"threadId": "thread_123",
"runId": "run_456",
"interrupt": [{"id": "req_1", "value": {"question": "Continue?"}}],
"result": {"status": "paused"},
}
update = converter.convert_event(event)
assert update is not None
assert update.additional_properties["interrupt"] == [{"id": "req_1", "value": {"question": "Continue?"}}]
assert update.additional_properties["result"] == {"status": "paused"}
def test_run_error_event(self) -> None:
"""Test conversion of RUN_ERROR event."""
converter = AGUIEventConverter()
@@ -239,6 +259,37 @@ class TestAGUIEventConverter:
assert update is None
def test_custom_event_conversion(self) -> None:
"""CUSTOM events are converted to update metadata."""
converter = AGUIEventConverter()
event = {
"type": "CUSTOM",
"name": "progress",
"value": {"percent": 10},
}
update = converter.convert_event(event)
assert update is not None
assert update.additional_properties["ag_ui_custom_event"]["name"] == "progress"
assert update.additional_properties["ag_ui_custom_event"]["value"] == {"percent": 10}
assert update.additional_properties["ag_ui_custom_event"]["raw_type"] == "CUSTOM"
def test_custom_event_alias_conversion(self) -> None:
"""CUSTOM_EVENT/custom_event aliases map to CUSTOM behavior."""
converter = AGUIEventConverter()
events = [
{"type": "CUSTOM_EVENT", "name": "alias_upper", "value": {"v": 1}},
{"type": "custom_event", "name": "alias_lower", "value": {"v": 2}},
]
updates = [converter.convert_event(event) for event in events]
assert updates[0] is not None
assert updates[1] is not None
assert updates[0].additional_properties["ag_ui_custom_event"]["raw_type"] == "CUSTOM_EVENT"
assert updates[1].additional_properties["ag_ui_custom_event"]["raw_type"] == "custom_event"
def test_full_conversation_flow(self) -> None:
"""Test complete conversation flow with multiple event types."""
converter = AGUIEventConverter()
python/packages/ag-ui/tests/ag_ui/test_http_service.py
+15 -3
View File
@@ -107,8 +107,8 @@ async def test_post_run_successful_streaming(mock_http_client, sample_events):
assert call_args.kwargs["headers"] == {"Accept": "text/event-stream"}
async def test_post_run_with_state_and_tools(mock_http_client):
"""Test posting run with state and tools."""
async def test_post_run_with_state_tools_and_interrupts(mock_http_client):
"""Test posting run with state, tools, and interrupt metadata."""
async def mock_aiter_lines():
return
@@ -127,8 +127,18 @@ async def test_post_run_with_state_and_tools(mock_http_client):
state = {"user_context": {"name": "Alice"}}
tools = [{"type": "function", "function": {"name": "test_tool"}}]
available_interrupts = [{"id": "req_1", "type": "request_info"}]
resume = {"interrupts": [{"id": "req_1", "value": "approved"}]}
async for _ in service.post_run(thread_id="thread_123", run_id="run_456", messages=[], state=state, tools=tools):
async for _ in service.post_run(
thread_id="thread_123",
run_id="run_456",
messages=[],
state=state,
tools=tools,
available_interrupts=available_interrupts,
resume=resume,
):
pass
# Verify state and tools were included in request
@@ -136,6 +146,8 @@ async def test_post_run_with_state_and_tools(mock_http_client):
request_data = call_args.kwargs["json"]
assert request_data["state"] == state
assert request_data["tools"] == tools
assert request_data["availableInterrupts"] == available_interrupts
assert request_data["resume"] == resume
async def test_post_run_http_error(mock_http_client):
python/packages/ag-ui/tests/ag_ui/test_message_adapters.py
+122
View File
@@ -2,7 +2,9 @@
"""Tests for message adapters."""
import base64
import json
import logging
import pytest
from agent_framework import Content, Message
@@ -406,6 +408,101 @@ def test_agui_non_string_content():
assert "nested" in messages[0].contents[0].text
def test_agui_multimodal_legacy_binary_to_agent_framework():
"""Legacy text/binary multimodal content converts to text + media Content."""
messages = agui_messages_to_agent_framework(
[
{
"role": "user",
"content": [
{"type": "text", "text": "See this image"},
{"type": "binary", "mimeType": "image/png", "url": "https://example.com/image.png"},
],
}
]
)
assert len(messages) == 1
assert len(messages[0].contents) == 2
assert messages[0].contents[0].type == "text"
assert messages[0].contents[0].text == "See this image"
assert messages[0].contents[1].type == "uri"
assert messages[0].contents[1].uri == "https://example.com/image.png"
assert messages[0].contents[1].media_type == "image/png"
def test_agui_multimodal_draft_source_base64_to_agent_framework():
"""Draft-style media source payload converts into data Content."""
payload = base64.b64encode(b"abc").decode("utf-8")
messages = agui_messages_to_agent_framework(
[
{
"role": "user",
"content": [
{
"type": "audio",
"source": {"type": "base64", "data": payload, "mimeType": "audio/wav"},
}
],
}
]
)
assert len(messages) == 1
assert len(messages[0].contents) == 1
assert messages[0].contents[0].type == "data"
assert messages[0].contents[0].media_type == "audio/wav"
assert isinstance(messages[0].contents[0].uri, str)
assert messages[0].contents[0].uri.startswith("data:audio/wav;base64,")
def test_agui_multimodal_invalid_base64_logs_warning(caplog):
"""Malformed base64 payloads should log and fall back to data URI."""
with caplog.at_level(logging.WARNING):
messages = agui_messages_to_agent_framework(
[
{
"role": "user",
"content": [
{
"type": "image",
"source": {"type": "base64", "data": "abc", "mimeType": "image/png"},
}
],
}
]
)
assert len(messages) == 1
assert len(messages[0].contents) == 1
assert messages[0].contents[0].type in {"data", "uri"}
assert messages[0].contents[0].uri == "data:image/png;base64,abc"
assert any("Failed to decode AG-UI media payload as base64" in record.message for record in caplog.records)
def test_agui_multimodal_mixed_order_preserved():
"""Mixed text/media multimodal input keeps content ordering."""
messages = agui_messages_to_agent_framework(
[
{
"role": "user",
"content": [
{"type": "text", "text": "First"},
{"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}},
{"type": "text", "text": "Last"},
],
}
]
)
assert len(messages[0].contents) == 3
assert messages[0].contents[0].type == "text"
assert messages[0].contents[0].text == "First"
assert messages[0].contents[1].type == "uri"
assert messages[0].contents[2].type == "text"
assert messages[0].contents[2].text == "Last"
def test_agui_message_without_id():
"""Test message without ID field."""
messages = agui_messages_to_agent_framework([{"role": "user", "content": "No ID"}])
@@ -414,6 +511,31 @@ def test_agui_message_without_id():
assert messages[0].message_id is None
def test_agui_snapshot_format_preserves_multimodal_content():
"""Snapshot normalization emits legacy binary parts for multimodal content."""
normalized = agui_messages_to_snapshot_format(
[
{
"role": "user",
"content": [
{"type": "input_text", "text": "Caption"},
{
"type": "image",
"source": {"type": "url", "url": "https://example.com/image.png", "mime_type": "image/png"},
},
],
}
]
)
assert isinstance(normalized[0]["content"], list)
content_parts = normalized[0]["content"]
assert content_parts[0]["type"] == "text"
assert content_parts[1]["type"] == "binary"
assert content_parts[1]["mimeType"] == "image/png"
assert content_parts[1]["url"] == "https://example.com/image.png"
def test_agui_with_tool_calls_to_agent_framework():
"""Assistant message with tool_calls is converted to FunctionCallContent."""
agui_msg = {
python/packages/ag-ui/tests/ag_ui/test_message_hygiene.py
+2 -2
View File
@@ -66,7 +66,7 @@ def test_convert_approval_results_to_tool_messages() -> None:
results ended up in user messages instead of tool messages, causing OpenAI to
reject the request with 'tool_call_ids did not have response messages'.
"""
from agent_framework_ag_ui._run import _convert_approval_results_to_tool_messages
from agent_framework_ag_ui._agent_run import _convert_approval_results_to_tool_messages
# Simulate what happens after _resolve_approval_responses:
# A user message contains function_result content (the executed tool result)
@@ -106,7 +106,7 @@ def test_convert_approval_results_preserves_other_user_content() -> None:
the function_result content should be extracted to a tool message while the
remaining content stays in the user message.
"""
from agent_framework_ag_ui._run import _convert_approval_results_to_tool_messages
from agent_framework_ag_ui._agent_run import _convert_approval_results_to_tool_messages
messages = [
Message(
python/packages/ag-ui/tests/ag_ui/test_public_exports.py
+24
View File
@@ -0,0 +1,24 @@
# Copyright (c) Microsoft. All rights reserved.
"""Public export coverage for AG-UI package surfaces."""
def test_agent_framework_ag_ui_exports_workflow() -> None:
"""Runtime package should export AgentFrameworkWorkflow."""
from agent_framework_ag_ui import AgentFrameworkWorkflow
assert AgentFrameworkWorkflow.__name__ == "AgentFrameworkWorkflow"
def test_core_ag_ui_lazy_exports_include_only_stable_api() -> None:
"""Core facade should expose only the stable high-level AG-UI API."""
from agent_framework import ag_ui
assert hasattr(ag_ui, "AgentFrameworkWorkflow")
assert hasattr(ag_ui, "AgentFrameworkAgent")
assert hasattr(ag_ui, "AGUIChatClient")
assert hasattr(ag_ui, "add_agent_framework_fastapi_endpoint")
assert not hasattr(ag_ui, "WorkflowFactory")
assert not hasattr(ag_ui, "AGUIRequest")
assert not hasattr(ag_ui, "RunMetadata")
python/packages/ag-ui/tests/ag_ui/test_run.py
+126 -27
View File
@@ -1,6 +1,6 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for _run.py helper functions and FlowState."""
"""Tests for _agent_run.py helper functions and FlowState."""
import pytest
from ag_ui.core import (
@@ -10,17 +10,25 @@ from ag_ui.core import (
from agent_framework import AgentResponseUpdate, Content, Message, ResponseStream
from agent_framework.exceptions import AgentInvalidResponseException
from agent_framework_ag_ui._run import (
FlowState,
from agent_framework_ag_ui._agent_run import (
_build_safe_metadata,
_create_state_context_message,
_emit_content,
_emit_tool_result,
_has_only_tool_calls,
_inject_state_context,
_normalize_response_stream,
_resume_to_tool_messages,
_should_suppress_intermediate_snapshot,
)
from agent_framework_ag_ui._run_common import (
FlowState,
_build_run_finished_event,
_emit_approval_request,
_emit_content,
_emit_text,
_emit_tool_call,
_emit_tool_result,
_extract_resume_payload,
_has_only_tool_calls,
)
class TestBuildSafeMetadata:
@@ -150,6 +158,7 @@ class TestFlowState:
assert flow.tool_calls_by_id == {}
assert flow.tool_results == []
assert flow.tool_calls_ended == set()
assert flow.interrupts == []
def test_get_tool_name(self):
"""Tests get_tool_name method."""
@@ -308,13 +317,11 @@ class TestInjectStateContext:
assert "Hello" in result[2].contents[0].text
# Additional tests for _run.py functions
# Additional tests for _agent_run.py functions
def test_emit_text_basic():
"""Test _emit_text emits correct events."""
from agent_framework_ag_ui._run import _emit_text
flow = FlowState()
content = Content.from_text("Hello world")
@@ -327,8 +334,6 @@ def test_emit_text_basic():
def test_emit_text_skip_empty():
"""Test _emit_text skips empty text."""
from agent_framework_ag_ui._run import _emit_text
flow = FlowState()
content = Content.from_text("")
@@ -339,8 +344,6 @@ def test_emit_text_skip_empty():
def test_emit_text_continues_existing_message():
"""Test _emit_text continues existing message."""
from agent_framework_ag_ui._run import _emit_text
flow = FlowState()
flow.message_id = "existing-id"
content = Content.from_text("more text")
@@ -351,10 +354,21 @@ def test_emit_text_continues_existing_message():
assert flow.message_id == "existing-id"
def test_emit_text_skips_duplicate_full_message_delta():
"""Test _emit_text skips replayed full-message chunks on an open message."""
flow = FlowState()
flow.message_id = "existing-id"
flow.accumulated_text = "Case complete."
content = Content.from_text("Case complete.")
events = _emit_text(content, flow)
assert events == []
assert flow.accumulated_text == "Case complete."
def test_emit_text_skips_when_waiting_for_approval():
"""Test _emit_text skips when waiting for approval."""
from agent_framework_ag_ui._run import _emit_text
flow = FlowState()
flow.waiting_for_approval = True
content = Content.from_text("should skip")
@@ -366,8 +380,6 @@ def test_emit_text_skips_when_waiting_for_approval():
def test_emit_text_skips_when_skip_text_flag():
"""Test _emit_text skips with skip_text flag."""
from agent_framework_ag_ui._run import _emit_text
flow = FlowState()
content = Content.from_text("should skip")
@@ -378,8 +390,6 @@ def test_emit_text_skips_when_skip_text_flag():
def test_emit_tool_call_basic():
"""Test _emit_tool_call emits correct events."""
from agent_framework_ag_ui._run import _emit_tool_call
flow = FlowState()
content = Content.from_function_call(
call_id="call_123",
@@ -396,8 +406,6 @@ def test_emit_tool_call_basic():
def test_emit_tool_call_generates_id():
"""Test _emit_tool_call generates ID when not provided."""
from agent_framework_ag_ui._run import _emit_tool_call
flow = FlowState()
# Create content without call_id
content = Content(type="function_call", name="test_tool", arguments="{}")
@@ -452,9 +460,100 @@ def test_emit_tool_result_no_open_message():
assert len(text_end_events) == 0
def test_emit_tool_result_serializes_non_string_result():
"""Non-string tool results should be serialized before emitting TOOL_CALL_RESULT."""
flow = FlowState()
content = Content.from_function_result(call_id="call_789", result={"ok": True, "items": [1, 2]})
events = _emit_tool_result(content, flow, predictive_handler=None)
result_event = next(event for event in events if getattr(event, "type", None) == "TOOL_CALL_RESULT")
assert isinstance(result_event.content, str)
assert '"ok": true' in result_event.content
assert flow.tool_results[0]["content"] == result_event.content
def test_emit_content_usage_emits_custom_usage_event():
"""Usage content should be emitted as a custom usage event."""
flow = FlowState()
content = Content.from_usage({"input_token_count": 3, "output_token_count": 2, "total_token_count": 5})
events = _emit_content(content, flow)
assert len(events) == 1
assert events[0].type == "CUSTOM"
assert events[0].name == "usage"
assert events[0].value["total_token_count"] == 5
def test_emit_approval_request_populates_interrupt_metadata():
"""Approval requests should populate FlowState interrupts for RUN_FINISHED metadata."""
flow = FlowState(message_id="msg-1")
function_call = Content.from_function_call(call_id="call_123", name="write_doc", arguments={"content": "x"})
approval_content = Content.from_function_approval_request(id="approval_1", function_call=function_call)
_emit_approval_request(approval_content, flow)
assert flow.waiting_for_approval is True
assert len(flow.interrupts) == 1
assert flow.interrupts[0]["id"] == "call_123"
assert flow.interrupts[0]["value"]["type"] == "function_approval_request"
def test_resume_to_tool_messages_from_interrupts_payload():
"""Resume payload interrupt responses map to tool messages."""
resume = {
"interrupts": [
{"id": "req_1", "value": {"accepted": True, "steps": []}},
{"id": "req_2", "value": "plain value"},
]
}
messages = _resume_to_tool_messages(resume)
assert len(messages) == 2
assert messages[0]["role"] == "tool"
assert messages[0]["toolCallId"] == "req_1"
assert '"accepted": true' in messages[0]["content"]
assert messages[1]["content"] == "plain value"
def test_extract_resume_payload_prefers_top_level_resume():
"""Top-level resume should take precedence over forwarded props."""
payload = {
"resume": {"interrupts": [{"id": "req_1", "value": "approved"}]},
"forwarded_props": {"command": {"resume": "ignored"}},
}
result = _extract_resume_payload(payload)
assert result == {"interrupts": [{"id": "req_1", "value": "approved"}]}
def test_extract_resume_payload_reads_forwarded_command_resume():
"""Forwarded command.resume should be treated as a resume payload."""
payload = {
"forwarded_props": {
"command": {"resume": '{"airline":"KLM","departure":"Amsterdam (AMS)","arrival":"San Francisco (SFO)"}'}
}
}
result = _extract_resume_payload(payload)
assert isinstance(result, str)
assert "KLM" in result
def test_build_run_finished_event_with_interrupt():
"""RUN_FINISHED helper should preserve interrupt payloads."""
event = _build_run_finished_event("run-1", "thread-1", interrupts=[{"id": "req_1", "value": {"x": 1}}])
dumped = event.model_dump()
assert dumped["run_id"] == "run-1"
assert dumped["thread_id"] == "thread-1"
assert dumped["interrupt"] == [{"id": "req_1", "value": {"x": 1}}]
def test_extract_approved_state_updates_no_handler():
"""Test _extract_approved_state_updates returns empty with no handler."""
from agent_framework_ag_ui._run import _extract_approved_state_updates
from agent_framework_ag_ui._agent_run import _extract_approved_state_updates
messages = [Message(role="user", contents=[Content.from_text("Hello")])]
result = _extract_approved_state_updates(messages, None)
@@ -463,8 +562,8 @@ def test_extract_approved_state_updates_no_handler():
def test_extract_approved_state_updates_no_approval():
"""Test _extract_approved_state_updates returns empty when no approval content."""
from agent_framework_ag_ui._agent_run import _extract_approved_state_updates
from agent_framework_ag_ui._orchestration._predictive_state import PredictiveStateHandler
from agent_framework_ag_ui._run import _extract_approved_state_updates
handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "content"}})
messages = [Message(role="user", contents=[Content.from_text("Hello")])]
@@ -481,7 +580,7 @@ class TestBuildMessagesSnapshot:
This is a regression test for issue #3619 where tool calls and content
were incorrectly merged into a single assistant message.
"""
from agent_framework_ag_ui._run import FlowState, _build_messages_snapshot
from agent_framework_ag_ui._agent_run import FlowState, _build_messages_snapshot
flow = FlowState()
flow.message_id = "msg-123"
@@ -518,7 +617,7 @@ class TestBuildMessagesSnapshot:
def test_only_tool_calls_no_text(self):
"""Test snapshot with only tool calls and no accumulated text."""
from agent_framework_ag_ui._run import FlowState, _build_messages_snapshot
from agent_framework_ag_ui._agent_run import FlowState, _build_messages_snapshot
flow = FlowState()
flow.message_id = "msg-123"
@@ -538,7 +637,7 @@ class TestBuildMessagesSnapshot:
def test_only_text_no_tool_calls(self):
"""Test snapshot with only text and no tool calls."""
from agent_framework_ag_ui._run import FlowState, _build_messages_snapshot
from agent_framework_ag_ui._agent_run import FlowState, _build_messages_snapshot
flow = FlowState()
flow.message_id = "msg-123"
@@ -558,7 +657,7 @@ class TestBuildMessagesSnapshot:
def test_preserves_snapshot_messages(self):
"""Test that existing snapshot messages are preserved."""
from agent_framework_ag_ui._run import FlowState, _build_messages_snapshot
from agent_framework_ag_ui._agent_run import FlowState, _build_messages_snapshot
flow = FlowState()
flow.pending_tool_calls = []
python/packages/ag-ui/tests/ag_ui/test_service_thread_id.py
+3 -3
View File
@@ -32,7 +32,7 @@ async def test_service_thread_id_when_there_are_updates(stub_agent):
}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
assert isinstance(events[0], RunStartedEvent)
@@ -54,7 +54,7 @@ async def test_service_thread_id_when_no_user_message(stub_agent):
}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
assert len(events) == 2
@@ -74,7 +74,7 @@ async def test_service_thread_id_when_user_supplied_thread_id(stub_agent):
input_data: dict[str, Any] = {"messages": [{"role": "user", "content": "Hi"}], "threadId": "conv_12345"}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
assert isinstance(events[0], RunStartedEvent)
python/packages/ag-ui/tests/ag_ui/test_structured_output.py
+7 -7
View File
@@ -52,7 +52,7 @@ async def test_structured_output_with_recipe(streaming_chat_client_stub, stream_
input_data = {"messages": [{"role": "user", "content": "Make pasta"}]}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
# Should emit StateSnapshotEvent with recipe
@@ -94,7 +94,7 @@ async def test_structured_output_with_steps(streaming_chat_client_stub, stream_f
input_data = {"messages": [{"role": "user", "content": "Do steps"}]}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
# Should emit StateSnapshotEvent with steps
@@ -129,7 +129,7 @@ async def test_structured_output_with_no_schema_match(streaming_chat_client_stub
input_data = {"messages": [{"role": "user", "content": "Generate data"}]}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
# Should emit StateSnapshotEvent but with no state updates since no schema fields match
@@ -164,7 +164,7 @@ async def test_structured_output_without_schema(streaming_chat_client_stub, stre
input_data = {"messages": [{"role": "user", "content": "Generate data"}]}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
# Should emit StateSnapshotEvent with both data and info fields
@@ -194,7 +194,7 @@ async def test_no_structured_output_when_no_response_format(streaming_chat_clien
input_data = {"messages": [{"role": "user", "content": "Hi"}]}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
# Should emit text content normally
@@ -224,7 +224,7 @@ async def test_structured_output_with_message_field(streaming_chat_client_stub,
input_data = {"messages": [{"role": "user", "content": "Make salad"}]}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
# Should emit the message as text
@@ -256,7 +256,7 @@ async def test_empty_updates_no_structured_processing(streaming_chat_client_stub
input_data = {"messages": [{"role": "user", "content": "Test"}]}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
async for event in wrapper.run(input_data):
events.append(event)
# Should only have start and end events
python/packages/ag-ui/tests/ag_ui/test_subgraphs_example_agent.py
+238
View File
@@ -0,0 +1,238 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for the subgraphs example agent used by Dojo."""
from __future__ import annotations
import json
from typing import Any
from agent_framework_ag_ui_examples.agents.subgraphs_agent import subgraphs_agent
async def _run(agent: Any, payload: dict[str, Any]) -> list[Any]:
return [event async for event in agent.run(payload)]
async def test_subgraphs_example_initial_run_emits_flight_interrupt() -> None:
"""Initial run should publish flight options and pause with an interrupt."""
agent = subgraphs_agent()
events = await _run(
agent,
{
"thread_id": "thread-subgraphs-initial",
"run_id": "run-initial",
"messages": [{"role": "user", "content": "Help me plan a trip to San Francisco"}],
},
)
event_types = [event.type for event in events]
assert event_types[0] == "RUN_STARTED"
assert "STATE_SNAPSHOT" in event_types
assert "STEP_STARTED" in event_types
assert "STEP_FINISHED" in event_types
assert "TEXT_MESSAGE_CONTENT" in event_types
assert "RUN_FINISHED" in event_types
started_steps = [event.step_name for event in events if event.type == "STEP_STARTED"]
finished_steps = [event.step_name for event in events if event.type == "STEP_FINISHED"]
assert "supervisor_agent" in started_steps
assert "flights_agent" in started_steps
assert "supervisor_agent" in finished_steps
assert "flights_agent" in finished_steps
finished = [event for event in events if event.type == "RUN_FINISHED"][0]
interrupt_payload = finished.model_dump().get("interrupt")
assert isinstance(interrupt_payload, list)
assert interrupt_payload
assert interrupt_payload[0]["value"]["agent"] == "flights"
assert len(interrupt_payload[0]["value"]["options"]) == 2
assert interrupt_payload[0]["value"]["options"][0]["airline"] == "KLM"
custom_event_names = [event.name for event in events if event.type == "CUSTOM"]
assert "WorkflowInterruptEvent" in custom_event_names
async def test_subgraphs_example_resume_flow_reaches_completion() -> None:
"""Flight + hotel resume payloads should complete the itinerary state."""
agent = subgraphs_agent()
thread_id = "thread-subgraphs-complete"
first_events = await _run(
agent,
{
"thread_id": thread_id,
"run_id": "run-1",
"messages": [{"role": "user", "content": "I want to visit San Francisco from Amsterdam"}],
},
)
first_interrupt = [event for event in first_events if event.type == "RUN_FINISHED"][0].model_dump()["interrupt"][0]
second_events = await _run(
agent,
{
"thread_id": thread_id,
"run_id": "run-2",
"resume": {
"interrupts": [
{
"id": first_interrupt["id"],
"value": json.dumps(
{
"airline": "United",
"departure": "Amsterdam (AMS)",
"arrival": "San Francisco (SFO)",
"price": "$720",
"duration": "12h 15m",
}
),
}
]
},
},
)
second_finished = [event for event in second_events if event.type == "RUN_FINISHED"][0].model_dump()
second_interrupt = second_finished.get("interrupt")
assert isinstance(second_interrupt, list)
assert second_interrupt[0]["value"]["agent"] == "hotels"
third_events = await _run(
agent,
{
"thread_id": thread_id,
"run_id": "run-3",
"resume": {
"interrupts": [
{
"id": second_interrupt[0]["id"],
"value": json.dumps(
{
"name": "The Ritz-Carlton",
"location": "Nob Hill",
"price_per_night": "$550/night",
"rating": "4.8 stars",
}
),
}
]
},
},
)
third_finished = [event for event in third_events if event.type == "RUN_FINISHED"][0].model_dump()
assert "interrupt" not in third_finished
snapshots = [event.snapshot for event in third_events if event.type == "STATE_SNAPSHOT"]
assert snapshots
final_snapshot = snapshots[-1]
assert final_snapshot["planning_step"] == "complete"
assert final_snapshot["active_agent"] == "supervisor"
assert final_snapshot["itinerary"]["flight"]["airline"] == "United"
assert final_snapshot["itinerary"]["hotel"]["name"] == "The Ritz-Carlton"
assert len(final_snapshot["experiences"]) == 4
async def test_subgraphs_example_requires_structured_resume_for_selection() -> None:
"""Agent should re-issue interrupts when user sends plain text instead of resume payload."""
agent = subgraphs_agent()
thread_id = "thread-subgraphs-text"
first_events = await _run(
agent,
{
"thread_id": thread_id,
"run_id": "run-a",
"messages": [{"role": "user", "content": "Plan a trip for me"}],
},
)
first_finished = [event for event in first_events if event.type == "RUN_FINISHED"][0].model_dump()
assert isinstance(first_finished.get("interrupt"), list)
assert first_finished["interrupt"][0]["value"]["agent"] == "flights"
second_events = await _run(
agent,
{
"thread_id": thread_id,
"run_id": "run-b",
"messages": [{"role": "user", "content": "Let's do the United flight"}],
},
)
second_finished = [event for event in second_events if event.type == "RUN_FINISHED"][0].model_dump()
assert isinstance(second_finished.get("interrupt"), list)
assert second_finished["interrupt"][0]["value"]["agent"] == "flights"
assert "TOOL_CALL_START" in [event.type for event in second_events]
assert "TEXT_MESSAGE_CONTENT" not in [event.type for event in second_events]
third_events = await _run(
agent,
{
"thread_id": thread_id,
"run_id": "run-c",
"resume": {
"interrupts": [
{
"id": second_finished["interrupt"][0]["id"],
"value": json.dumps(
{
"airline": "United",
"departure": "Amsterdam (AMS)",
"arrival": "San Francisco (SFO)",
"price": "$720",
"duration": "12h 15m",
}
),
}
]
},
},
)
third_finished = [event for event in third_events if event.type == "RUN_FINISHED"][0].model_dump()
assert isinstance(third_finished.get("interrupt"), list)
assert third_finished["interrupt"][0]["value"]["agent"] == "hotels"
third_snapshots = [event.snapshot for event in third_events if event.type == "STATE_SNAPSHOT"]
assert third_snapshots[-1]["itinerary"]["flight"]["airline"] == "United"
async def test_subgraphs_example_forwarded_command_resume_reaches_hotels_interrupt() -> None:
"""CopilotKit-style forwarded command.resume should continue workflow interrupts."""
agent = subgraphs_agent()
thread_id = "thread-subgraphs-forwarded-resume"
first_events = await _run(
agent,
{
"thread_id": thread_id,
"run_id": "run-forwarded-1",
"messages": [{"role": "user", "content": "Plan my trip"}],
},
)
first_interrupt = [event for event in first_events if event.type == "RUN_FINISHED"][0].model_dump()["interrupt"][0]
second_events = await _run(
agent,
{
"thread_id": thread_id,
"run_id": "run-forwarded-2",
"messages": [],
"forwarded_props": {
"command": {
"resume": json.dumps(
{
"airline": "KLM",
"departure": "Amsterdam (AMS)",
"arrival": "San Francisco (SFO)",
"price": "$650",
"duration": "11h 30m",
}
)
}
},
},
)
second_finished = [event for event in second_events if event.type == "RUN_FINISHED"][0].model_dump()
second_interrupt = second_finished.get("interrupt")
assert isinstance(second_interrupt, list)
assert second_interrupt[0]["value"]["agent"] == "hotels"
assert second_interrupt[0]["id"] != first_interrupt["id"]
python/packages/ag-ui/tests/ag_ui/test_types.py
+27
View File
@@ -183,6 +183,21 @@ class TestAGUIRequest:
assert request.forwarded_props == {"custom_key": "custom_value"}
assert request.parent_run_id == "parent-run-789"
def test_agui_request_camel_case_aliases(self) -> None:
"""Test AGUIRequest accepts camelCase aliases from AG-UI HTTP clients."""
request = AGUIRequest(
messages=[{"role": "user", "content": "Hello"}],
runId="run-camel-1",
threadId="thread-camel-1",
forwardedProps={"k": "v"},
parentRunId="parent-camel-1",
)
assert request.run_id == "run-camel-1"
assert request.thread_id == "thread-camel-1"
assert request.forwarded_props == {"k": "v"}
assert request.parent_run_id == "parent-camel-1"
def test_agui_request_model_dump_excludes_none(self) -> None:
"""Test that model_dump(exclude_none=True) excludes None fields."""
request = AGUIRequest(
@@ -223,3 +238,15 @@ class TestAGUIRequest:
assert dumped["context"] == [{"type": "snippet", "content": "code here"}]
assert dumped["forwarded_props"] == {"auth_token": "secret", "user_id": "user-1"}
assert dumped["parent_run_id"] == "parent-456"
def test_agui_request_available_interrupts_alias_round_trip(self) -> None:
"""availableInterrupts should deserialize, while dumps remain snake_case."""
request = AGUIRequest(
messages=[{"role": "user", "content": "Hello"}],
availableInterrupts=[{"id": "req_1", "value": {"choice": "A"}}],
)
assert request.available_interrupts == [{"id": "req_1", "value": {"choice": "A"}}]
dumped = request.model_dump(exclude_none=True)
assert dumped["available_interrupts"] == [{"id": "req_1", "value": {"choice": "A"}}]
assert "availableInterrupts" not in dumped
python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py
+112
View File
@@ -0,0 +1,112 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for AgentFrameworkWorkflow wrapper behavior."""
from __future__ import annotations
from typing import Any, cast
import pytest
from agent_framework import Workflow, WorkflowBuilder, WorkflowContext, executor
from agent_framework_ag_ui import AgentFrameworkWorkflow
async def _run(agent: AgentFrameworkWorkflow, payload: dict[str, Any]) -> list[Any]:
return [event async for event in agent.run(payload)]
async def test_workflow_wrapper_rejects_workflow_and_factory_at_once() -> None:
"""Workflow wrapper should reject ambiguous workflow source configuration."""
@executor(id="start")
async def start(message: Any, ctx: WorkflowContext) -> None:
del message
await ctx.yield_output("ok")
workflow = WorkflowBuilder(start_executor=start).build()
with pytest.raises(ValueError, match="workflow_factory"):
AgentFrameworkWorkflow(workflow=workflow, workflow_factory=lambda _thread_id: workflow)
async def test_workflow_wrapper_factory_is_thread_scoped() -> None:
"""Thread-scoped workflow factories should isolate workflow instances by thread id."""
@executor(id="requester")
async def requester(message: Any, ctx: WorkflowContext) -> None:
del message
await ctx.request_info({"message": "Choose an option", "options": ["a", "b"]}, dict, request_id="choice")
factory_calls: dict[str, int] = {}
def workflow_factory(thread_id: str) -> Workflow:
factory_calls[thread_id] = factory_calls.get(thread_id, 0) + 1
return WorkflowBuilder(start_executor=requester).build()
agent = AgentFrameworkWorkflow(workflow_factory=workflow_factory)
first_events = await _run(
agent,
{
"thread_id": "thread-a",
"messages": [{"role": "user", "content": "start"}],
},
)
first_finished = [event for event in first_events if event.type == "RUN_FINISHED"][0].model_dump()
first_interrupt = first_finished.get("interrupt")
assert isinstance(first_interrupt, list)
assert first_interrupt[0]["id"] == "choice"
assert factory_calls["thread-a"] == 1
second_events = await _run(
agent,
{
"thread_id": "thread-a",
"messages": [],
"resume": {"interrupts": [{"id": "choice", "value": {"selection": "a"}}]},
},
)
second_types = [event.type for event in second_events]
assert "RUN_ERROR" not in second_types
second_finished = [event for event in second_events if event.type == "RUN_FINISHED"][0].model_dump()
assert "interrupt" not in second_finished
assert factory_calls["thread-a"] == 1
third_events = await _run(
agent,
{
"thread_id": "thread-b",
"messages": [{"role": "user", "content": "start"}],
},
)
third_finished = [event for event in third_events if event.type == "RUN_FINISHED"][0].model_dump()
third_interrupt = third_finished.get("interrupt")
assert isinstance(third_interrupt, list)
assert third_interrupt[0]["id"] == "choice"
assert factory_calls["thread-b"] == 1
agent.clear_thread_workflow("thread-a")
await _run(
agent,
{
"thread_id": "thread-a",
"messages": [{"role": "user", "content": "restart"}],
},
)
assert factory_calls["thread-a"] == 2
async def test_workflow_wrapper_without_workflow_raises_not_implemented() -> None:
"""Without workflow/workflow_factory, run should raise NotImplementedError."""
agent = AgentFrameworkWorkflow()
with pytest.raises(NotImplementedError, match="No workflow is attached"):
_ = [event async for event in agent.run({"messages": [{"role": "user", "content": "start"}]})]
async def test_workflow_wrapper_factory_return_type_is_validated() -> None:
"""Factory outputs must be Workflow instances."""
agent = AgentFrameworkWorkflow(workflow_factory=lambda _thread_id: cast(Any, object()))
with pytest.raises(TypeError, match="workflow_factory must return a Workflow instance"):
_ = [event async for event in agent.run({"thread_id": "thread-a", "messages": []})]
python/packages/ag-ui/tests/ag_ui/test_workflow_run.py
+679
View File
@@ -0,0 +1,679 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for native workflow AG-UI runner."""
import json
from types import SimpleNamespace
from typing import Any, cast
from ag_ui.core import EventType, StateSnapshotEvent
from agent_framework import (
AgentResponse,
Content,
Executor,
Message,
WorkflowBuilder,
WorkflowContext,
WorkflowEvent,
executor,
handler,
response_handler,
)
from typing_extensions import Never
from agent_framework_ag_ui._workflow_run import (
_coerce_message,
_coerce_response_for_request,
run_workflow_stream,
)
class ProgressEvent(WorkflowEvent):
"""Custom workflow event used to validate CUSTOM mapping."""
def __init__(self, progress: int) -> None:
super().__init__("custom_progress", data={"progress": progress})
async def test_workflow_run_maps_custom_and_text_events():
"""Custom workflow events and yielded text are mapped to AG-UI events."""
@executor(id="start")
async def start(message: Any, ctx: WorkflowContext[Never, str]) -> None:
await ctx.add_event(ProgressEvent(10))
await ctx.yield_output("Hello workflow")
workflow = WorkflowBuilder(start_executor=start).build()
input_data = {"messages": [{"role": "user", "content": "go"}]}
events = [event async for event in run_workflow_stream(input_data, workflow)]
event_types = [event.type for event in events]
assert "RUN_STARTED" in event_types
assert "CUSTOM" in event_types
assert "TEXT_MESSAGE_CONTENT" in event_types
assert "STEP_STARTED" in event_types
assert "STEP_FINISHED" in event_types
assert "RUN_FINISHED" in event_types
custom_events = [event for event in events if event.type == "CUSTOM" and event.name == "custom_progress"]
assert len(custom_events) == 1
assert custom_events[0].value == {"progress": 10}
async def test_workflow_run_request_info_emits_interrupt_and_resume_works():
"""request_info should emit interrupt metadata and resume should continue run."""
@executor(id="requester")
async def requester(message: Any, ctx: WorkflowContext) -> None:
await ctx.request_info("Need approval", str)
workflow = WorkflowBuilder(start_executor=requester).build()
first_run_events = [
event async for event in run_workflow_stream({"messages": [{"role": "user", "content": "go"}]}, workflow)
]
run_finished_events = [event for event in first_run_events if event.type == "RUN_FINISHED"]
assert len(run_finished_events) == 1
interrupt_payload = run_finished_events[0].model_dump().get("interrupt")
assert isinstance(interrupt_payload, list)
assert len(interrupt_payload) == 1
request_id = str(interrupt_payload[0]["id"])
assert request_id
resumed_events = [
event
async for event in run_workflow_stream(
{"messages": [], "resume": {"interrupts": [{"id": request_id, "value": "approved"}]}},
workflow,
)
]
resumed_types = [event.type for event in resumed_events]
assert "RUN_STARTED" in resumed_types
assert "RUN_FINISHED" in resumed_types
assert "RUN_ERROR" not in resumed_types
async def test_workflow_run_request_info_closes_open_text_message() -> None:
"""Text output should end before request_info interrupt events begin."""
@executor(id="requester")
async def requester(message: Any, ctx: WorkflowContext) -> None:
del message
await ctx.yield_output("Please confirm this action.")
await ctx.request_info("Need approval", str, request_id="approval-1")
workflow = WorkflowBuilder(start_executor=requester).build()
events = [event async for event in run_workflow_stream({"messages": [{"role": "user", "content": "go"}]}, workflow)]
content_index = next(i for i, event in enumerate(events) if event.type == "TEXT_MESSAGE_CONTENT")
end_index = next(i for i, event in enumerate(events) if event.type == "TEXT_MESSAGE_END")
request_start_index = next(
i
for i, event in enumerate(events)
if event.type == "TOOL_CALL_START" and getattr(event, "tool_call_id", None) == "approval-1"
)
assert content_index < end_index < request_start_index
async def test_workflow_run_request_info_interrupt_uses_raw_dict_value():
"""Dict request payloads should be surfaced directly in RUN_FINISHED.interrupt.value."""
@executor(id="requester")
async def requester(message: Any, ctx: WorkflowContext) -> None:
await ctx.request_info(
{
"message": "Choose a flight",
"options": [{"airline": "KLM"}],
"recommendation": {"airline": "KLM"},
"agent": "flights",
},
dict,
request_id="flights-choice",
)
workflow = WorkflowBuilder(start_executor=requester).build()
events = [event async for event in run_workflow_stream({"messages": [{"role": "user", "content": "go"}]}, workflow)]
run_finished = [event for event in events if event.type == "RUN_FINISHED"][0].model_dump()
interrupt_payload = run_finished.get("interrupt")
assert isinstance(interrupt_payload, list)
assert interrupt_payload[0]["id"] == "flights-choice"
assert interrupt_payload[0]["value"]["agent"] == "flights"
assert interrupt_payload[0]["value"]["message"] == "Choose a flight"
async def test_workflow_run_resume_from_forwarded_command_payload() -> None:
"""forwarded_props.command.resume should resume a pending dict request."""
@executor(id="requester")
async def requester(message: Any, ctx: WorkflowContext) -> None:
del message
await ctx.request_info({"options": [{"airline": "KLM"}]}, dict, request_id="flights-choice")
workflow = WorkflowBuilder(start_executor=requester).build()
_ = [event async for event in run_workflow_stream({"messages": [{"role": "user", "content": "go"}]}, workflow)]
resumed_events = [
event
async for event in run_workflow_stream(
{
"messages": [],
"forwarded_props": {
"command": {"resume": json.dumps({"airline": "KLM", "departure": "AMS", "arrival": "SFO"})}
},
},
workflow,
)
]
resumed_types = [event.type for event in resumed_events]
assert "RUN_ERROR" not in resumed_types
finished = [event for event in resumed_events if event.type == "RUN_FINISHED"][0].model_dump()
assert "interrupt" not in finished
async def test_workflow_run_structured_user_json_resumes_single_pending_request() -> None:
"""A JSON user reply should resume a single pending dict request without heuristics."""
@executor(id="requester")
async def requester(message: Any, ctx: WorkflowContext) -> None:
del message
await ctx.request_info({"options": [{"name": "Hotel Zoe"}]}, dict, request_id="hotel-choice")
workflow = WorkflowBuilder(start_executor=requester).build()
_ = [event async for event in run_workflow_stream({"messages": [{"role": "user", "content": "go"}]}, workflow)]
resumed_events = [
event
async for event in run_workflow_stream(
{
"messages": [{"role": "user", "content": json.dumps({"name": "Hotel Zoe"})}],
},
workflow,
)
]
resumed_types = [event.type for event in resumed_events]
assert "RUN_ERROR" not in resumed_types
finished = [event for event in resumed_events if event.type == "RUN_FINISHED"][0].model_dump()
assert "interrupt" not in finished
async def test_workflow_run_resume_content_response_from_json_payload() -> None:
"""JSON resume payloads should coerce into Content responses for approval requests."""
class ApprovalExecutor(Executor):
def __init__(self) -> None:
super().__init__(id="approval_executor")
@handler
async def start(self, message: Any, ctx: WorkflowContext) -> None:
del message
function_call = Content.from_function_call(
call_id="refund-call",
name="submit_refund",
arguments={"order_id": "12345", "amount": "$89.99"},
)
approval_request = Content.from_function_approval_request(id="approval-1", function_call=function_call)
await ctx.request_info(approval_request, Content, request_id="approval-1")
@response_handler
async def handle_approval(self, original_request: Content, response: Content, ctx: WorkflowContext) -> None:
del original_request
status = "approved" if bool(response.approved) else "rejected"
await ctx.yield_output(f"Refund tool call {status}.")
workflow = WorkflowBuilder(start_executor=ApprovalExecutor()).build()
first_events = [
event async for event in run_workflow_stream({"messages": [{"role": "user", "content": "go"}]}, workflow)
]
first_finished = [event for event in first_events if event.type == "RUN_FINISHED"][0].model_dump()
interrupt_payload = cast(list[dict[str, Any]], first_finished.get("interrupt"))
interrupt_value = cast(dict[str, Any], interrupt_payload[0]["value"])
resumed_events = [
event
async for event in run_workflow_stream(
{
"messages": [],
"resume": {
"interrupts": [
{
"id": "approval-1",
"value": {
"type": "function_approval_response",
"approved": True,
"id": interrupt_value.get("id", "approval-1"),
"function_call": interrupt_value.get("function_call"),
},
}
]
},
},
workflow,
)
]
resumed_types = [event.type for event in resumed_events]
assert "RUN_ERROR" not in resumed_types
assert "TEXT_MESSAGE_CONTENT" in resumed_types
resumed_finished = [event for event in resumed_events if event.type == "RUN_FINISHED"][0].model_dump()
assert "interrupt" not in resumed_finished
text_deltas = [event.delta for event in resumed_events if event.type == "TEXT_MESSAGE_CONTENT"]
assert any("approved" in delta for delta in text_deltas)
async def test_workflow_run_resume_message_list_from_json_payload() -> None:
"""Resume payloads should coerce AG-UI message dictionaries into list[Message] responses."""
class MessageRequestExecutor(Executor):
def __init__(self) -> None:
super().__init__(id="message_request_executor")
@handler
async def start(self, message: Any, ctx: WorkflowContext) -> None:
del message
await ctx.request_info({"prompt": "Need user follow-up"}, list[Message], request_id="handoff-user-input")
@response_handler
async def handle_user_input(
self, original_request: dict, response: list[Message], ctx: WorkflowContext
) -> None:
del original_request
user_text = response[0].text if response else ""
await ctx.yield_output(f"Captured response: {user_text}")
workflow = WorkflowBuilder(start_executor=MessageRequestExecutor()).build()
_ = [event async for event in run_workflow_stream({"messages": [{"role": "user", "content": "start"}]}, workflow)]
resumed_events = [
event
async for event in run_workflow_stream(
{
"messages": [],
"resume": {
"interrupts": [
{
"id": "handoff-user-input",
"value": [
{
"role": "user",
"contents": [{"type": "text", "text": "Please ship a replacement instead."}],
}
],
}
]
},
},
workflow,
)
]
resumed_types = [event.type for event in resumed_events]
assert "RUN_ERROR" not in resumed_types
assert "TEXT_MESSAGE_CONTENT" in resumed_types
resumed_finished = [event for event in resumed_events if event.type == "RUN_FINISHED"][0].model_dump()
assert "interrupt" not in resumed_finished
text_deltas = [event.delta for event in resumed_events if event.type == "TEXT_MESSAGE_CONTENT"]
assert any("replacement" in delta for delta in text_deltas)
async def test_workflow_run_non_chat_output_maps_to_custom_output_event():
"""Non-chat workflow outputs are emitted as CUSTOM workflow_output events."""
@executor(id="structured")
async def structured(message: Any, ctx: WorkflowContext[Never, dict[str, int]]) -> None:
await ctx.yield_output({"count": 3})
workflow = WorkflowBuilder(start_executor=structured).build()
events = [event async for event in run_workflow_stream({"messages": [{"role": "user", "content": "go"}]}, workflow)]
output_custom = [event for event in events if event.type == "CUSTOM" and event.name == "workflow_output"]
assert len(output_custom) == 1
assert output_custom[0].value == {"count": 3}
async def test_workflow_run_passthroughs_ag_ui_base_events():
"""Workflow outputs that are AG-UI BaseEvent instances should be emitted directly."""
@executor(id="stateful")
async def stateful(message: Any, ctx: WorkflowContext[Never, StateSnapshotEvent]) -> None:
await ctx.yield_output(StateSnapshotEvent(type=EventType.STATE_SNAPSHOT, snapshot={"active_agent": "flights"}))
workflow = WorkflowBuilder(start_executor=stateful).build()
events = [event async for event in run_workflow_stream({"messages": [{"role": "user", "content": "go"}]}, workflow)]
snapshots = [event for event in events if event.type == "STATE_SNAPSHOT"]
assert len(snapshots) == 1
assert snapshots[0].snapshot["active_agent"] == "flights"
async def test_workflow_run_plain_text_follow_up_does_not_infer_interrupt_response():
"""User follow-up text should not be coerced into request_info responses for workflows."""
@executor(id="requester")
async def requester(message: Any, ctx: WorkflowContext) -> None:
del message
await ctx.request_info(
{
"message": "Choose a flight",
"options": [{"airline": "KLM"}, {"airline": "United"}],
"agent": "flights",
},
dict,
request_id="flights-choice",
)
workflow = WorkflowBuilder(start_executor=requester).build()
_ = [event async for event in run_workflow_stream({"messages": [{"role": "user", "content": "go"}]}, workflow)]
follow_up_events = [
event
async for event in run_workflow_stream(
{
"messages": [
{
"role": "assistant",
"content": "",
"tool_calls": [
{
"id": "flights-choice",
"type": "function",
"function": {"name": "request_info", "arguments": "{}"},
}
],
},
{"role": "user", "content": "I prefer KLM please"},
]
},
workflow,
)
]
follow_up_types = [event.type for event in follow_up_events]
assert "RUN_ERROR" not in follow_up_types
assert "TOOL_CALL_START" in follow_up_types
run_finished = [event for event in follow_up_events if event.type == "RUN_FINISHED"][0].model_dump()
interrupt_payload = run_finished.get("interrupt")
assert isinstance(interrupt_payload, list)
assert interrupt_payload[0]["id"] == "flights-choice"
assert interrupt_payload[0]["value"]["agent"] == "flights"
async def test_workflow_run_empty_turn_with_pending_request_preserves_interrupts():
"""An empty turn should still return pending workflow interrupts without errors."""
@executor(id="requester")
async def requester(message: Any, ctx: WorkflowContext) -> None:
del message
await ctx.request_info({"prompt": "choose"}, dict, request_id="pick-one")
workflow = WorkflowBuilder(start_executor=requester).build()
_ = [event async for event in run_workflow_stream({"messages": [{"role": "user", "content": "go"}]}, workflow)]
events = [event async for event in run_workflow_stream({"messages": []}, workflow)]
types = [event.type for event in events]
assert types[0] == "RUN_STARTED"
assert "RUN_FINISHED" in types
assert "RUN_ERROR" not in types
finished = [event for event in events if event.type == "RUN_FINISHED"][0].model_dump()
interrupts = finished.get("interrupt")
assert isinstance(interrupts, list)
assert interrupts[0]["id"] == "pick-one"
async def test_workflow_run_agent_response_output_uses_latest_assistant_message_only() -> None:
"""Conversation payload outputs should not flatten full history into one assistant message."""
@executor(id="responder")
async def responder(message: Any, ctx: WorkflowContext[Never, AgentResponse]) -> None:
del message
response = AgentResponse(
messages=[
Message(role="user", contents=[Content.from_text("My order arrived damaged")]),
Message(
role="assistant",
contents=[Content.from_text("Order Agent: Got it. I submitted the replacement request.")],
),
]
)
await ctx.yield_output(response)
workflow = WorkflowBuilder(start_executor=responder).build()
events = [event async for event in run_workflow_stream({"messages": [{"role": "user", "content": "go"}]}, workflow)]
text_deltas = [event.delta for event in events if event.type == "TEXT_MESSAGE_CONTENT"]
assert text_deltas == ["Order Agent: Got it. I submitted the replacement request."]
async def test_workflow_run_skips_duplicate_text_from_conversation_snapshot() -> None:
"""Do not emit duplicate assistant text when a snapshot repeats the latest output."""
@executor(id="responder")
async def responder(message: Any, ctx: WorkflowContext[Never, Any]) -> None:
del message
duplicate_text = "Order Agent: Got it. I submitted the replacement request."
await ctx.yield_output(duplicate_text)
await ctx.yield_output(
AgentResponse(
messages=[
Message(role="user", contents=[Content.from_text("standard")]),
Message(role="assistant", contents=[Content.from_text(duplicate_text)]),
]
)
)
workflow = WorkflowBuilder(start_executor=responder).build()
events = [event async for event in run_workflow_stream({"messages": [{"role": "user", "content": "go"}]}, workflow)]
text_deltas = [event.delta for event in events if event.type == "TEXT_MESSAGE_CONTENT"]
assert text_deltas == ["Order Agent: Got it. I submitted the replacement request."]
async def test_workflow_run_skips_consecutive_duplicate_text_outputs() -> None:
"""Do not emit duplicate assistant text when consecutive outputs are identical."""
@executor(id="responder")
async def responder(message: Any, ctx: WorkflowContext[Never, Any]) -> None:
del message
duplicate_text = "Order Agent: Replacement processed. Case complete."
await ctx.yield_output(duplicate_text)
await ctx.yield_output(duplicate_text)
workflow = WorkflowBuilder(start_executor=responder).build()
events = [event async for event in run_workflow_stream({"messages": [{"role": "user", "content": "go"}]}, workflow)]
text_deltas = [event.delta for event in events if event.type == "TEXT_MESSAGE_CONTENT"]
assert text_deltas == ["Order Agent: Replacement processed. Case complete."]
async def test_workflow_run_skips_final_snapshot_when_streamed_chunks_already_match() -> None:
"""Do not append full snapshot text when prior chunk outputs already formed the same message."""
@executor(id="responder")
async def responder(message: Any, ctx: WorkflowContext[Never, Any]) -> None:
del message
full_text = (
"Your replacement request for order 28939393 has been submitted with expedited shipping, "
"as you requested.\n\nCase complete."
)
await ctx.yield_output(
"Your replacement request for order 28939393 has been submitted with expedited shipping, "
)
await ctx.yield_output("as you requested.\n\nCase complete.")
await ctx.yield_output(
AgentResponse(
messages=[
Message(role="user", contents=[Content.from_text("My order is 28939393.")]),
Message(role="assistant", contents=[Content.from_text(full_text)]),
]
)
)
workflow = WorkflowBuilder(start_executor=responder).build()
events = [event async for event in run_workflow_stream({"messages": [{"role": "user", "content": "go"}]}, workflow)]
text_deltas = [event.delta for event in events if event.type == "TEXT_MESSAGE_CONTENT"]
assert text_deltas == [
"Your replacement request for order 28939393 has been submitted with expedited shipping, ",
"as you requested.\n\nCase complete.",
]
async def test_workflow_run_usage_content_emits_custom_usage_event() -> None:
"""Usage output from workflows should be surfaced as a custom usage event."""
@executor(id="usage")
async def usage(message: Any, ctx: WorkflowContext[Never, Content]) -> None:
del message
await ctx.yield_output(
Content.from_usage(
{
"input_token_count": 12,
"output_token_count": 6,
"total_token_count": 18,
}
)
)
workflow = WorkflowBuilder(start_executor=usage).build()
events = [event async for event in run_workflow_stream({"messages": [{"role": "user", "content": "go"}]}, workflow)]
usage_events = [event for event in events if event.type == "CUSTOM" and event.name == "usage"]
assert len(usage_events) == 1
assert usage_events[0].value["input_token_count"] == 12
assert usage_events[0].value["output_token_count"] == 6
assert usage_events[0].value["total_token_count"] == 18
async def test_workflow_run_accepts_multimodal_input_messages() -> None:
"""Workflow runner should normalize multimodal input into workflow Message content."""
class CapturingWorkflow:
def __init__(self) -> None:
self.captured_message: list[Message] | None = None
def run(self, **kwargs: Any):
self.captured_message = cast(list[Message] | None, kwargs.get("message"))
async def _stream():
yield SimpleNamespace(type="started")
return _stream()
workflow = CapturingWorkflow()
events = [
event
async for event in run_workflow_stream(
{
"messages": [
{
"role": "user",
"content": [
{"type": "text", "text": "Please analyze this image"},
{
"type": "image",
"source": {
"type": "url",
"url": "https://example.com/diagram.png",
"mimeType": "image/png",
},
},
],
}
]
},
cast(Any, workflow),
)
]
event_types = [event.type for event in events]
assert "RUN_STARTED" in event_types
assert "RUN_FINISHED" in event_types
assert "RUN_ERROR" not in event_types
assert workflow.captured_message is not None
assert len(workflow.captured_message) == 1
user_message = workflow.captured_message[0]
assert user_message.role == "user"
assert len(user_message.contents) == 2
assert user_message.contents[0].type == "text"
assert user_message.contents[0].text == "Please analyze this image"
assert user_message.contents[1].type == "uri"
assert user_message.contents[1].uri == "https://example.com/diagram.png"
def test_coerce_message_accepts_string_payload() -> None:
"""String values should coerce into a user Message with one text content."""
message = _coerce_message("Please continue")
assert message is not None
assert message.role == "user"
assert len(message.contents) == 1
assert message.contents[0].type == "text"
assert message.contents[0].text == "Please continue"
def test_coerce_message_accepts_content_key_variant() -> None:
"""The 'content' key variant should map into Message.contents."""
message = _coerce_message({"role": "assistant", "content": {"type": "text", "content": "Done"}})
assert message is not None
assert message.role == "assistant"
assert len(message.contents) == 1
assert message.contents[0].type == "text"
assert message.contents[0].text == "Done"
def test_coerce_response_for_request_bool_int_float_and_mismatch() -> None:
"""Scalar coercion should enforce bool/int/float rules and return None on mismatches."""
bool_request = SimpleNamespace(response_type=bool)
assert _coerce_response_for_request(bool_request, True) is True
assert _coerce_response_for_request(bool_request, "true") is True
assert _coerce_response_for_request(bool_request, 1) is None
int_request = SimpleNamespace(response_type=int)
assert _coerce_response_for_request(int_request, 7) == 7
assert _coerce_response_for_request(int_request, "7") == 7
assert _coerce_response_for_request(int_request, True) is None
float_request = SimpleNamespace(response_type=float)
assert _coerce_response_for_request(float_request, 2) == 2
assert _coerce_response_for_request(float_request, "2.5") == 2.5
assert _coerce_response_for_request(float_request, True) is None
dict_request = SimpleNamespace(response_type=dict)
assert _coerce_response_for_request(dict_request, "[1,2,3]") is None
async def test_workflow_run_emits_run_error_when_stream_raises() -> None:
"""Unexpected stream exceptions should be converted into RUN_ERROR events."""
class FailingWorkflow:
def run(self, **kwargs: Any):
del kwargs
async def _stream():
raise RuntimeError("workflow stream exploded")
yield # pragma: no cover
return _stream()
events = [
event
async for event in run_workflow_stream(
{"messages": [{"role": "user", "content": "go"}]},
cast(Any, FailingWorkflow()),
)
]
event_types = [event.type for event in events]
assert event_types[0] == "RUN_STARTED"
assert "RUN_ERROR" in event_types
run_error = next(event for event in events if event.type == "RUN_ERROR")
assert "workflow stream exploded" in run_error.message
python/packages/core/agent_framework/_tools.py
+19 -2
View File
@@ -14,6 +14,7 @@ from collections.abc import (
Mapping,
Sequence,
)
from contextlib import suppress
from functools import partial, wraps
from time import perf_counter, time_ns
from typing import (
@@ -288,15 +289,18 @@ class FunctionTool(SerializationMixin):
self.func = func
self._instance = None # Store the instance for bound methods
# Initialize schema cache (will be lazily populated)
self._input_schema_cached: dict[str, Any] | None = None
# Track if schema was supplied as JSON dict (for optimization)
if isinstance(input_model, Mapping):
self._schema_supplied = True
self._input_schema: dict[str, Any] = dict(input_model)
self._input_schema_cached = dict(input_model)
self.input_model: type[BaseModel] | None = None
else:
self._schema_supplied = False
self.input_model = self._resolve_input_model(input_model)
self._input_schema = self.input_model.model_json_schema()
# Defer schema generation to avoid issues with forward references
self._cached_parameters: dict[str, Any] | None = None
self.approval_mode = approval_mode or "never_require"
if max_invocations is not None and max_invocations < 1:
@@ -546,6 +550,19 @@ class FunctionTool(SerializationMixin):
self._invocation_duration_histogram.record(duration, attributes=attributes)
logger.info("Function duration: %fs", duration)
@property
def _input_schema(self) -> dict[str, Any]:
"""Get the input schema, generating it lazily if needed."""
if self._input_schema_cached is None:
if self.input_model is not None:
# Try to rebuild the model in case it has forward references
with suppress(Exception):
self.input_model.model_rebuild(force=True, raise_errors=False)
self._input_schema_cached = self.input_model.model_json_schema()
else:
self._input_schema_cached = {}
return self._input_schema_cached
def parameters(self) -> dict[str, Any]:
"""Create the JSON schema of the parameters.
python/packages/core/agent_framework/ag_ui/__init__.py
+1 -2
View File
@@ -20,10 +20,9 @@ IMPORT_PATH = "agent_framework_ag_ui"
PACKAGE_NAME = "agent-framework-ag-ui"
_IMPORTS = [
"AgentFrameworkAgent",
"AgentFrameworkWorkflow",
"add_agent_framework_fastapi_endpoint",
"AGUIChatClient",
"AGUIEventConverter",
"AGUIHttpService",
]
python/packages/core/agent_framework/ag_ui/__init__.pyi
+4
View File
@@ -2,9 +2,11 @@
from agent_framework_ag_ui import (
AgentFrameworkAgent,
AgentFrameworkWorkflow,
AGUIChatClient,
AGUIEventConverter,
AGUIHttpService,
__version__,
add_agent_framework_fastapi_endpoint,
)
@@ -13,5 +15,7 @@ __all__ = [
"AGUIEventConverter",
"AGUIHttpService",
"AgentFrameworkAgent",
"AgentFrameworkWorkflow",
"__version__",
"add_agent_framework_fastapi_endpoint",
]
python/packages/core/tests/workflow/test_full_conversation.py
+3 -11
View File
@@ -362,9 +362,7 @@ async def test_run_request_with_full_history_clears_service_session_id() -> None
"""Replaying a full conversation (including function calls) via AgentExecutorRequest must
clear service_session_id so the API does not receive both previous_response_id and the
same function-call items in input — which would cause a 'Duplicate item' API error."""
tool_agent = _ToolHistoryAgent(
id="tool_agent", name="ToolAgent", summary_text="Done."
)
tool_agent = _ToolHistoryAgent(id="tool_agent", name="ToolAgent", summary_text="Done.")
tool_exec = AgentExecutor(tool_agent, id="tool_agent")
spy_agent = _SessionIdCapturingAgent(id="spy_agent", name="SpyAgent")
@@ -393,9 +391,7 @@ async def test_from_response_preserves_service_session_id() -> None:
"""from_response hands off a prior agent's full conversation to the next executor.
The receiving executor's service_session_id is preserved so the API can continue
the conversation using previous_response_id."""
tool_agent = _ToolHistoryAgent(
id="tool_agent2", name="ToolAgent", summary_text="Done."
)
tool_agent = _ToolHistoryAgent(id="tool_agent2", name="ToolAgent", summary_text="Done.")
tool_exec = AgentExecutor(tool_agent, id="tool_agent2")
spy_agent = _SessionIdCapturingAgent(id="spy_agent2", name="SpyAgent")
@@ -403,11 +399,7 @@ async def test_from_response_preserves_service_session_id() -> None:
# Simulate a prior run on the spy executor.
spy_exec._session.service_session_id = "resp_PREVIOUS_RUN" # pyright: ignore[reportPrivateUsage]
wf = (
WorkflowBuilder(start_executor=tool_exec, output_executors=[spy_exec])
.add_edge(tool_exec, spy_exec)
.build()
)
wf = WorkflowBuilder(start_executor=tool_exec, output_executors=[spy_exec]).add_edge(tool_exec, spy_exec).build()
result = await wf.run("start")
assert result.get_outputs() is not None
python/packages/orchestrations/agent_framework_orchestrations/_handoff.py
+152 -24
View File
@@ -41,7 +41,7 @@ from agent_framework import Agent, SupportsAgentRun
from agent_framework._middleware import FunctionInvocationContext, FunctionMiddleware
from agent_framework._sessions import AgentSession
from agent_framework._tools import FunctionTool, tool
from agent_framework._types import AgentResponse, AgentResponseUpdate, Message
from agent_framework._types import AgentResponse, AgentResponseUpdate, Content, Message
from agent_framework._workflows._agent_executor import AgentExecutor, AgentExecutorRequest, AgentExecutorResponse
from agent_framework._workflows._agent_utils import resolve_agent_id
from agent_framework._workflows._checkpoint import CheckpointStorage
@@ -267,6 +267,88 @@ class HandoffAgentExecutor(AgentExecutor):
return cloned_agent
def _persist_pending_approval_function_calls(self) -> None:
"""Persist pending approval function calls for stateless provider resumes.
Handoff workflows force ``store=False`` and replay conversation state from ``_full_conversation``.
When a run pauses on function approval, ``AgentExecutor`` returns ``None`` and the assistant
function-call message is not returned as an ``AgentResponse``. Without persisting that call, the
next turn may submit only a function result, which responses-style APIs reject.
"""
pending_calls: list[Content] = []
for request in self._pending_agent_requests.values():
if request.type != "function_approval_request":
continue
function_call = getattr(request, "function_call", None)
if isinstance(function_call, Content) and function_call.type == "function_call":
pending_calls.append(function_call)
if not pending_calls:
return
self._full_conversation.append(
Message(
role="assistant",
contents=pending_calls,
author_name=self._agent.name,
)
)
def _persist_missing_approved_function_results(
self,
*,
runtime_tool_messages: list[Message],
response_messages: list[Message],
) -> None:
"""Persist fallback function_result entries for approved calls when missing.
In approval resumes, function invocation can execute approved tools without
always surfacing those tool outputs in the returned ``AgentResponse.messages``.
For stateless handoff replays, we must keep call/output pairs balanced.
"""
candidate_results: dict[str, Content] = {}
for message in runtime_tool_messages:
for content in message.contents:
if content.type == "function_result":
call_id = getattr(content, "call_id", None)
if isinstance(call_id, str) and call_id:
candidate_results[call_id] = content
continue
if content.type != "function_approval_response" or not content.approved:
continue
function_call = getattr(content, "function_call", None)
call_id = getattr(function_call, "call_id", None) or getattr(content, "id", None)
if isinstance(call_id, str) and call_id and call_id not in candidate_results:
# Fallback content for approved calls when runtime messages do not include
# a concrete function_result payload.
candidate_results[call_id] = Content.from_function_result(
call_id=call_id,
result='{"status":"approved"}',
)
if not candidate_results:
return
observed_result_call_ids: set[str] = set()
for message in [*self._full_conversation, *response_messages]:
for content in message.contents:
if content.type == "function_result" and isinstance(content.call_id, str) and content.call_id:
observed_result_call_ids.add(content.call_id)
missing_call_ids = sorted(set(candidate_results.keys()) - observed_result_call_ids)
if not missing_call_ids:
return
self._full_conversation.append(
Message(
role="tool",
contents=[candidate_results[call_id] for call_id in missing_call_ids],
author_name=self._agent.name,
)
)
def _clone_chat_agent(self, agent: Agent) -> Agent:
"""Produce a deep copy of the Agent while preserving runtime configuration."""
options = agent.default_options
@@ -287,6 +369,10 @@ class HandoffAgentExecutor(AgentExecutor):
# Disable parallel tool calls to prevent the agent from invoking multiple handoff tools at once.
cloned_options: dict[str, Any] = {
"allow_multiple_tool_calls": False,
# Handoff workflows already manage full conversation context explicitly
# across executors. Keep provider-side conversation storage disabled to
# avoid stale tool-call state (Responses API previous_response chains).
"store": False,
"frequency_penalty": options.get("frequency_penalty"),
"instructions": options.get("instructions"),
"logit_bias": dict(logit_bias) if logit_bias else None,
@@ -297,7 +383,6 @@ class HandoffAgentExecutor(AgentExecutor):
"response_format": options.get("response_format"),
"seed": options.get("seed"),
"stop": options.get("stop"),
"store": options.get("store"),
"temperature": options.get("temperature"),
"tool_choice": options.get("tool_choice"),
"tools": all_tools if all_tools else None,
@@ -366,14 +451,44 @@ class HandoffAgentExecutor(AgentExecutor):
self, ctx: WorkflowContext[AgentExecutorResponse, AgentResponse | AgentResponseUpdate]
) -> None:
"""Override to support handoff."""
incoming_messages = list(self._cache)
cleaned_incoming_messages = clean_conversation_for_handoff(incoming_messages)
runtime_tool_messages = [
message
for message in incoming_messages
if any(
content.type
in {
"function_result",
"function_approval_response",
}
for content in message.contents
)
or message.role == "tool"
]
# When the full conversation is empty, it means this is the first run.
# Broadcast the initial cache to all other agents. Subsequent runs won't
# need this since responses are broadcast after each agent run and user input.
if self._is_start_agent and not self._full_conversation:
await self._broadcast_messages(self._cache.copy(), cast(WorkflowContext[AgentExecutorRequest], ctx))
await self._broadcast_messages(cleaned_incoming_messages, cast(WorkflowContext[AgentExecutorRequest], ctx))
# Append the cache to the full conversation history
self._full_conversation.extend(self._cache)
# Persist only cleaned chat history between turns to avoid replaying stale tool calls.
self._full_conversation.extend(cleaned_incoming_messages)
# Always run with full conversation context for request_info resumes.
# Keep runtime tool-control messages for this run only (e.g., approval responses).
self._cache = list(self._full_conversation)
self._cache.extend(runtime_tool_messages)
# Handoff workflows are orchestrator-stateful and provider-stateless by design.
# If an existing session still has a service conversation id, clear it to avoid
# replaying stale unresolved tool calls across resumed turns.
if (
cast(Agent, self._agent).default_options.get("store") is False
and self._session.service_session_id is not None
):
self._session.service_session_id = None
# Check termination condition before running the agent
if await self._check_terminate_and_yield(cast(WorkflowContext[Never, list[Message]], ctx)):
@@ -392,17 +507,26 @@ class HandoffAgentExecutor(AgentExecutor):
# A function approval request is issued by the base AgentExecutor
if response is None:
if cast(Agent, self._agent).default_options.get("store") is False:
self._persist_pending_approval_function_calls()
# Agent did not complete (e.g., waiting for user input); do not emit response
logger.debug("AgentExecutor %s: Agent did not complete, awaiting user input", self.id)
return
# Remove function call related content from the agent response for full conversation history
# Remove function call related content from the agent response for broadcast.
# This prevents replaying stale tool artifacts to other agents.
cleaned_response = clean_conversation_for_handoff(response.messages)
# Append the agent response to the full conversation history. This list removes
# function call related content such that the result stays consistent regardless
# of which agent yields the final output.
self._full_conversation.extend(cleaned_response)
# Broadcast the cleaned response to all other agents
# For internal tracking, preserve the full response (including function_calls)
# in _full_conversation so that Azure OpenAI can match function_calls with
# function_results when the workflow resumes after user approvals.
self._full_conversation.extend(response.messages)
self._persist_missing_approved_function_results(
runtime_tool_messages=runtime_tool_messages,
response_messages=response.messages,
)
# Broadcast only the cleaned response to other agents (without function_calls/results)
await self._broadcast_messages(cleaned_response, cast(WorkflowContext[AgentExecutorRequest], ctx))
# Check if a handoff was requested
@@ -422,6 +546,12 @@ class HandoffAgentExecutor(AgentExecutor):
self._autonomous_mode_turns = 0 # Reset autonomous mode turn counter on handoff
return
# Re-evaluate termination after appending and broadcasting this response.
# Without this check, workflows that become terminal due to the latest assistant
# message would still emit request_info and require an unnecessary extra resume.
if await self._check_terminate_and_yield(cast(WorkflowContext[Never, list[Message]], ctx)):
return
# Handle case where no handoff was requested
if self._autonomous_mode and self._autonomous_mode_turns < self._autonomous_mode_turn_limit:
# In autonomous mode, continue running the agent until a handoff is requested
@@ -497,22 +627,20 @@ class HandoffAgentExecutor(AgentExecutor):
last_message = response.messages[-1]
for content in last_message.contents:
if content.type == "function_result":
if not content.result:
continue
parsed_result: dict[str, Any] | None = None
if isinstance(content.result, dict):
parsed_result = content.result
elif isinstance(content.result, str):
payload = content.result
parsed_payload: dict[str, Any] | None = None
if isinstance(payload, dict):
parsed_payload = payload
elif isinstance(payload, str):
try:
loaded_result = json.loads(content.result)
maybe_payload = json.loads(payload)
except json.JSONDecodeError:
continue
if isinstance(loaded_result, dict):
parsed_result = loaded_result
maybe_payload = None
if isinstance(maybe_payload, dict):
parsed_payload = maybe_payload
if parsed_result is not None:
handoff_target = parsed_result.get(HANDOFF_FUNCTION_RESULT_KEY)
if parsed_payload:
handoff_target = parsed_payload.get(HANDOFF_FUNCTION_RESULT_KEY)
if isinstance(handoff_target, str):
return handoff_target
else:
python/packages/orchestrations/agent_framework_orchestrations/_orchestrator_helpers.py
+20 -44
View File
@@ -14,57 +14,33 @@ logger = logging.getLogger(__name__)
def clean_conversation_for_handoff(conversation: list[Message]) -> list[Message]:
"""Remove tool-related content from conversation for clean handoffs.
"""Keep only plain text chat history for handoff routing.
During handoffs, tool calls can cause API errors because:
1. Assistant messages with tool_calls must be followed by tool responses
2. Tool response messages must follow an assistant message with tool_calls
Handoff executors must not replay prior tool-control artifacts (function calls,
tool outputs, approval payloads) into future model turns, or providers may reject
the next request due to unmatched tool-call state.
This creates a cleaned copy removing ALL tool-related content.
Removes:
- function_approval_request and function_call from assistant messages
- Tool response messages (role="tool")
- Messages with only tool calls and no text
Preserves:
- User messages
- Assistant messages with text content
Args:
conversation: Original conversation with potential tool content
Returns:
Cleaned conversation safe for handoff routing
This helper builds a text-only copy of the conversation:
- Drops all non-text content from every message.
- Drops messages with no remaining text content.
- Preserves original roles and author names for retained text messages.
"""
cleaned: list[Message] = []
for msg in conversation:
# Skip tool response messages entirely
if msg.role == "tool":
# Keep only plain text history for handoff routing. Tool-control content
# (function_call/function_result/approval payloads) is runtime-only and
# must not be replayed in future model turns.
text_parts = [content.text for content in msg.contents if content.type == "text" and content.text]
if not text_parts:
continue
# Check for tool-related content
has_tool_content = False
if msg.contents:
has_tool_content = any(
content.type in ("function_approval_request", "function_call") for content in msg.contents
)
# If no tool content, keep original
if not has_tool_content:
cleaned.append(msg)
continue
# Has tool content - only keep if it also has text
if msg.text and msg.text.strip():
# Create fresh text-only message while preserving additional_properties
msg_copy = Message(
role=msg.role,
text=msg.text,
author_name=msg.author_name,
additional_properties=dict(msg.additional_properties) if msg.additional_properties else None,
)
cleaned.append(msg_copy)
msg_copy = Message(
role=msg.role,
text=" ".join(text_parts),
author_name=msg.author_name,
additional_properties=dict(msg.additional_properties) if msg.additional_properties else None,
)
cleaned.append(msg_copy)
return cleaned
python/packages/orchestrations/tests/test_handoff.py
+682 -1
View File
@@ -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