agent/agent-framework
2
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: [BREAKING] simplify ag-ui run logic, fix mcp bugs, fix anthropic client issues in ag-ui (#3322)

Browse Source 
* Refactor ag-ui to simplify flow

* Refactoring

* Fix backend tool

* Update tests

* Improvements

* Fix mypy

* Fixes

* Fix json serialize errors
This commit is contained in:
Evan Mattson
2026-01-23 14:10:46 +09:00
committed by GitHub
Unverified
parent 9f893a32a6
commit 5436354a83
42 changed files with 2789 additions and 5395 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/ag-ui/tests/test_agent_wrapper_comprehensive.py
+39 -52
View File
@@ -420,17 +420,25 @@ async def test_tool_result_function_approval_rejected():
async def test_thread_metadata_tracking():
"""Test that thread metadata includes ag_ui_thread_id and ag_ui_run_id."""
"""Test that thread metadata includes ag_ui_thread_id and ag_ui_run_id.
AG-UI internal metadata is stored in thread.metadata for orchestration,
but filtered out before passing to the chat client's options.metadata.
"""
from agent_framework.ag_ui import AgentFrameworkAgent
thread_metadata: dict[str, Any] = {}
captured_thread: dict[str, Any] = {}
captured_options: dict[str, Any] = {}
async def stream_fn(
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
metadata = options.get("metadata")
if metadata:
thread_metadata.update(metadata)
# Capture the thread object from kwargs
thread = kwargs.get("thread")
if thread and hasattr(thread, "metadata"):
captured_thread["metadata"] = thread.metadata
# Capture options to verify internal keys are NOT passed to chat client
captured_options.update(options)
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
@@ -446,22 +454,37 @@ async def test_thread_metadata_tracking():
async for event in wrapper.run_agent(input_data):
events.append(event)
# AG-UI internal metadata should be stored in thread.metadata
thread_metadata = captured_thread.get("metadata", {})
assert thread_metadata.get("ag_ui_thread_id") == "test_thread_123"
assert thread_metadata.get("ag_ui_run_id") == "test_run_456"
# Internal metadata should NOT be passed to chat client options
options_metadata = captured_options.get("metadata", {})
assert "ag_ui_thread_id" not in options_metadata
assert "ag_ui_run_id" not in options_metadata
async def test_state_context_injection():
"""Test that current state is injected into thread metadata."""
"""Test that current state is injected into thread metadata.
AG-UI internal metadata (including current_state) is stored in thread.metadata
for orchestration, but filtered out before passing to the chat client's options.metadata.
"""
from agent_framework_ag_ui import AgentFrameworkAgent
thread_metadata: dict[str, Any] = {}
captured_thread: dict[str, Any] = {}
captured_options: dict[str, Any] = {}
async def stream_fn(
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
metadata = options.get("metadata")
if metadata:
thread_metadata.update(metadata)
# Capture the thread object from kwargs
thread = kwargs.get("thread")
if thread and hasattr(thread, "metadata"):
captured_thread["metadata"] = thread.metadata
# Capture options to verify internal keys are NOT passed to chat client
captured_options.update(options)
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
@@ -479,11 +502,17 @@ async def test_state_context_injection():
async for event in wrapper.run_agent(input_data):
events.append(event)
# Current state should be stored in thread.metadata
thread_metadata = captured_thread.get("metadata", {})
current_state = thread_metadata.get("current_state")
if isinstance(current_state, str):
current_state = json.loads(current_state)
assert current_state == {"document": "Test content"}
# Internal metadata should NOT be passed to chat client options
options_metadata = captured_options.get("metadata", {})
assert "current_state" not in options_metadata
async def test_no_messages_provided():
"""Test handling when no messages are provided."""
@@ -595,48 +624,6 @@ async def test_json_decode_error_in_tool_result():
assert len(tool_events) == 0
async def test_suppressed_summary_with_document_state():
"""Test suppressed summary uses document state for confirmation message."""
from agent_framework.ag_ui import AgentFrameworkAgent, DocumentWriterConfirmationStrategy
async def stream_fn(
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="Response")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
wrapper = AgentFrameworkAgent(
agent=agent,
state_schema={"document": {"type": "string"}},
predict_state_config={"document": {"tool": "write_doc", "tool_argument": "content"}},
confirmation_strategy=DocumentWriterConfirmationStrategy(),
)
# Simulate confirmation with document state
tool_result: dict[str, Any] = {"accepted": True, "steps": []}
input_data: dict[str, Any] = {
"messages": [
{
"role": "tool",
"content": json.dumps(tool_result),
"toolCallId": "confirm_123",
}
],
"state": {"document": "This is the beginning of a document. It contains important information."},
}
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
events.append(event)
# Should generate fallback summary from document state
text_events = [e for e in events if e.type == "TEXT_MESSAGE_CONTENT"]
assert len(text_events) > 0
# Should contain some reference to the document
full_text = "".join(e.delta for e in text_events)
assert "written" in full_text.lower() or "document" in full_text.lower()
async def test_agent_with_use_service_thread_is_false():
"""Test that when use_service_thread is False, the AgentThread used to run the agent is NOT set to the service thread ID."""
from agent_framework.ag_ui import AgentFrameworkAgent
python/packages/ag-ui/tests/test_backend_tool_rendering.py
-129
View File
@@ -1,129 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for backend tool rendering."""
from typing import cast
from ag_ui.core import (
TextMessageContentEvent,
TextMessageStartEvent,
ToolCallArgsEvent,
ToolCallEndEvent,
ToolCallResultEvent,
ToolCallStartEvent,
)
from agent_framework import AgentResponseUpdate, Content
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
async def test_tool_call_flow():
"""Test complete tool call flow: call -> args -> end -> result."""
bridge = AgentFrameworkEventBridge(run_id="test-run", thread_id="test-thread")
# Step 1: Tool call starts
tool_call = Content.from_function_call(
call_id="weather-123",
name="get_weather",
arguments={"location": "Seattle"},
)
update1 = AgentResponseUpdate(contents=[tool_call])
events1 = await bridge.from_agent_run_update(update1)
# Should have: ToolCallStartEvent, ToolCallArgsEvent
assert len(events1) == 2
assert isinstance(events1[0], ToolCallStartEvent)
assert isinstance(events1[1], ToolCallArgsEvent)
start_event = events1[0]
assert start_event.tool_call_id == "weather-123"
assert start_event.tool_call_name == "get_weather"
args_event = events1[1]
assert "Seattle" in args_event.delta
# Step 2: Tool result comes back
tool_result = Content.from_function_result(
call_id="weather-123",
result="Weather in Seattle: Rainy, 52°F",
)
update2 = AgentResponseUpdate(contents=[tool_result])
events2 = await bridge.from_agent_run_update(update2)
# Should have: ToolCallEndEvent, ToolCallResultEvent
assert len(events2) == 2
assert isinstance(events2[0], ToolCallEndEvent)
assert isinstance(events2[1], ToolCallResultEvent)
end_event = events2[0]
assert end_event.tool_call_id == "weather-123"
result_event = events2[1]
assert result_event.tool_call_id == "weather-123"
assert "Seattle" in result_event.content
assert "Rainy" in result_event.content
async def test_text_with_tool_call():
"""Test agent response with both text and tool calls."""
bridge = AgentFrameworkEventBridge(run_id="test-run", thread_id="test-thread")
# Agent says something then calls a tool
text_content = Content.from_text(text="Let me check the weather for you.")
tool_call = Content.from_function_call(
call_id="weather-456",
name="get_forecast",
arguments={"location": "San Francisco", "days": 3},
)
update = AgentResponseUpdate(contents=[text_content, tool_call])
events = await bridge.from_agent_run_update(update)
# Should have: TextMessageStart, TextMessageContent, ToolCallStart, ToolCallArgs
assert len(events) == 4
assert isinstance(events[0], TextMessageStartEvent)
assert isinstance(events[1], TextMessageContentEvent)
assert isinstance(events[2], ToolCallStartEvent)
assert isinstance(events[3], ToolCallArgsEvent)
text_event = events[1]
assert "check the weather" in text_event.delta
tool_start = events[2]
assert tool_start.tool_call_name == "get_forecast"
async def test_multiple_tool_results():
"""Test handling multiple tool results in sequence."""
bridge = AgentFrameworkEventBridge(run_id="test-run", thread_id="test-thread")
# Multiple tool results
results = [
Content.from_function_result(call_id="tool-1", result="Result 1"),
Content.from_function_result(call_id="tool-2", result="Result 2"),
Content.from_function_result(call_id="tool-3", result="Result 3"),
]
update = AgentResponseUpdate(contents=results)
events = await bridge.from_agent_run_update(update)
# Should have 3 pairs of ToolCallEndEvent + ToolCallResultEvent = 6 events
assert len(events) == 6
# Verify the pattern: End, Result, End, Result, End, Result
for i in range(3):
end_idx = i * 2
result_idx = i * 2 + 1
assert isinstance(events[end_idx], ToolCallEndEvent)
assert isinstance(events[result_idx], ToolCallResultEvent)
end_event = cast(ToolCallEndEvent, events[end_idx])
result_event = cast(ToolCallResultEvent, events[result_idx])
assert end_event.tool_call_id == f"tool-{i + 1}"
assert result_event.tool_call_id == f"tool-{i + 1}"
assert f"Result {i + 1}" in result_event.content
python/packages/ag-ui/tests/test_confirmation_strategies_comprehensive.py
-275
View File
@@ -1,275 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""Comprehensive tests for all confirmation strategies."""
import pytest
from agent_framework_ag_ui._confirmation_strategies import (
ConfirmationStrategy,
DefaultConfirmationStrategy,
DocumentWriterConfirmationStrategy,
RecipeConfirmationStrategy,
TaskPlannerConfirmationStrategy,
)
@pytest.fixture
def sample_steps() -> list[dict[str, str]]:
"""Sample steps for testing approval messages."""
return [
{"description": "Step 1: Do something", "status": "enabled"},
{"description": "Step 2: Do another thing", "status": "enabled"},
{"description": "Step 3: Disabled step", "status": "disabled"},
]
@pytest.fixture
def all_enabled_steps() -> list[dict[str, str]]:
"""All steps enabled."""
return [
{"description": "Task A", "status": "enabled"},
{"description": "Task B", "status": "enabled"},
{"description": "Task C", "status": "enabled"},
]
@pytest.fixture
def empty_steps() -> list[dict[str, str]]:
"""Empty steps list."""
return []
class TestDefaultConfirmationStrategy:
"""Tests for DefaultConfirmationStrategy."""
def test_on_approval_accepted_with_enabled_steps(self, sample_steps: list[dict[str, str]]) -> None:
strategy = DefaultConfirmationStrategy()
message = strategy.on_approval_accepted(sample_steps)
assert "Executing 2 approved steps" in message
assert "Step 1: Do something" in message
assert "Step 2: Do another thing" in message
assert "Step 3" not in message # Disabled step shouldn't appear
assert "All steps completed successfully!" in message
def test_on_approval_accepted_with_all_enabled(self, all_enabled_steps: list[dict[str, str]]) -> None:
strategy = DefaultConfirmationStrategy()
message = strategy.on_approval_accepted(all_enabled_steps)
assert "Executing 3 approved steps" in message
assert "Task A" in message
assert "Task B" in message
assert "Task C" in message
def test_on_approval_accepted_with_empty_steps(self, empty_steps: list[dict[str, str]]) -> None:
strategy = DefaultConfirmationStrategy()
message = strategy.on_approval_accepted(empty_steps)
assert "Executing 0 approved steps" in message
assert "All steps completed successfully!" in message
def test_on_approval_rejected(self, sample_steps: list[dict[str, str]]) -> None:
strategy = DefaultConfirmationStrategy()
message = strategy.on_approval_rejected(sample_steps)
assert "No problem!" in message
assert "What would you like me to change" in message
def test_on_state_confirmed(self) -> None:
strategy = DefaultConfirmationStrategy()
message = strategy.on_state_confirmed()
assert "Changes confirmed" in message
assert "successfully" in message
def test_on_state_rejected(self) -> None:
strategy = DefaultConfirmationStrategy()
message = strategy.on_state_rejected()
assert "No problem!" in message
assert "What would you like me to change" in message
class TestTaskPlannerConfirmationStrategy:
"""Tests for TaskPlannerConfirmationStrategy."""
def test_on_approval_accepted_with_enabled_steps(self, sample_steps: list[dict[str, str]]) -> None:
strategy = TaskPlannerConfirmationStrategy()
message = strategy.on_approval_accepted(sample_steps)
assert "Executing your requested tasks" in message
assert "1. Step 1: Do something" in message
assert "2. Step 2: Do another thing" in message
assert "Step 3" not in message
assert "All tasks completed successfully!" in message
def test_on_approval_accepted_with_all_enabled(self, all_enabled_steps: list[dict[str, str]]) -> None:
strategy = TaskPlannerConfirmationStrategy()
message = strategy.on_approval_accepted(all_enabled_steps)
assert "Executing your requested tasks" in message
assert "1. Task A" in message
assert "2. Task B" in message
assert "3. Task C" in message
def test_on_approval_accepted_with_empty_steps(self, empty_steps: list[dict[str, str]]) -> None:
strategy = TaskPlannerConfirmationStrategy()
message = strategy.on_approval_accepted(empty_steps)
assert "Executing your requested tasks" in message
assert "All tasks completed successfully!" in message
def test_on_approval_rejected(self, sample_steps: list[dict[str, str]]) -> None:
strategy = TaskPlannerConfirmationStrategy()
message = strategy.on_approval_rejected(sample_steps)
assert "No problem!" in message
assert "revise the plan" in message
def test_on_state_confirmed(self) -> None:
strategy = TaskPlannerConfirmationStrategy()
message = strategy.on_state_confirmed()
assert "Tasks confirmed" in message
assert "ready to execute" in message
def test_on_state_rejected(self) -> None:
strategy = TaskPlannerConfirmationStrategy()
message = strategy.on_state_rejected()
assert "No problem!" in message
assert "adjust the task list" in message
class TestRecipeConfirmationStrategy:
"""Tests for RecipeConfirmationStrategy."""
def test_on_approval_accepted_with_enabled_steps(self, sample_steps: list[dict[str, str]]) -> None:
strategy = RecipeConfirmationStrategy()
message = strategy.on_approval_accepted(sample_steps)
assert "Updating your recipe" in message
assert "1. Step 1: Do something" in message
assert "2. Step 2: Do another thing" in message
assert "Step 3" not in message
assert "Recipe updated successfully!" in message
def test_on_approval_accepted_with_all_enabled(self, all_enabled_steps: list[dict[str, str]]) -> None:
strategy = RecipeConfirmationStrategy()
message = strategy.on_approval_accepted(all_enabled_steps)
assert "Updating your recipe" in message
assert "1. Task A" in message
assert "2. Task B" in message
assert "3. Task C" in message
def test_on_approval_accepted_with_empty_steps(self, empty_steps: list[dict[str, str]]) -> None:
strategy = RecipeConfirmationStrategy()
message = strategy.on_approval_accepted(empty_steps)
assert "Updating your recipe" in message
assert "Recipe updated successfully!" in message
def test_on_approval_rejected(self, sample_steps: list[dict[str, str]]) -> None:
strategy = RecipeConfirmationStrategy()
message = strategy.on_approval_rejected(sample_steps)
assert "No problem!" in message
assert "ingredients or steps" in message
def test_on_state_confirmed(self) -> None:
strategy = RecipeConfirmationStrategy()
message = strategy.on_state_confirmed()
assert "Recipe changes applied" in message
assert "successfully" in message
def test_on_state_rejected(self) -> None:
strategy = RecipeConfirmationStrategy()
message = strategy.on_state_rejected()
assert "No problem!" in message
assert "adjust in the recipe" in message
class TestDocumentWriterConfirmationStrategy:
"""Tests for DocumentWriterConfirmationStrategy."""
def test_on_approval_accepted_with_enabled_steps(self, sample_steps: list[dict[str, str]]) -> None:
strategy = DocumentWriterConfirmationStrategy()
message = strategy.on_approval_accepted(sample_steps)
assert "Applying your edits" in message
assert "1. Step 1: Do something" in message
assert "2. Step 2: Do another thing" in message
assert "Step 3" not in message
assert "Document updated successfully!" in message
def test_on_approval_accepted_with_all_enabled(self, all_enabled_steps: list[dict[str, str]]) -> None:
strategy = DocumentWriterConfirmationStrategy()
message = strategy.on_approval_accepted(all_enabled_steps)
assert "Applying your edits" in message
assert "1. Task A" in message
assert "2. Task B" in message
assert "3. Task C" in message
def test_on_approval_accepted_with_empty_steps(self, empty_steps: list[dict[str, str]]) -> None:
strategy = DocumentWriterConfirmationStrategy()
message = strategy.on_approval_accepted(empty_steps)
assert "Applying your edits" in message
assert "Document updated successfully!" in message
def test_on_approval_rejected(self, sample_steps: list[dict[str, str]]) -> None:
strategy = DocumentWriterConfirmationStrategy()
message = strategy.on_approval_rejected(sample_steps)
assert "No problem!" in message
assert "keep or modify" in message
def test_on_state_confirmed(self) -> None:
strategy = DocumentWriterConfirmationStrategy()
message = strategy.on_state_confirmed()
assert "Document edits applied!" in message
def test_on_state_rejected(self) -> None:
strategy = DocumentWriterConfirmationStrategy()
message = strategy.on_state_rejected()
assert "No problem!" in message
assert "change about the document" in message
class TestConfirmationStrategyInterface:
"""Tests for ConfirmationStrategy abstract base class."""
def test_cannot_instantiate_abstract_class(self):
"""Verify ConfirmationStrategy is abstract and cannot be instantiated."""
with pytest.raises(TypeError):
ConfirmationStrategy() # type: ignore
def test_all_strategies_implement_interface(self):
"""Verify all concrete strategies implement the full interface."""
strategies = [
DefaultConfirmationStrategy(),
TaskPlannerConfirmationStrategy(),
RecipeConfirmationStrategy(),
DocumentWriterConfirmationStrategy(),
]
sample_steps = [{"description": "Test", "status": "enabled"}]
for strategy in strategies:
# All should have these methods
assert callable(strategy.on_approval_accepted)
assert callable(strategy.on_approval_rejected)
assert callable(strategy.on_state_confirmed)
assert callable(strategy.on_state_rejected)
# All should return strings
assert isinstance(strategy.on_approval_accepted(sample_steps), str)
assert isinstance(strategy.on_approval_rejected(sample_steps), str)
assert isinstance(strategy.on_state_confirmed(), str)
assert isinstance(strategy.on_state_rejected(), str)
python/packages/ag-ui/tests/test_document_writer_flow.py
-236
View File
@@ -1,236 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for document writer predictive state flow with confirm_changes."""
from ag_ui.core import EventType, StateDeltaEvent, ToolCallArgsEvent, ToolCallEndEvent, ToolCallStartEvent
from agent_framework import AgentResponseUpdate, Content
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
async def test_streaming_document_with_state_deltas():
"""Test that streaming tool arguments emit progressive StateDeltaEvents."""
predict_config = {
"document": {"tool": "write_document_local", "tool_argument": "document"},
}
bridge = AgentFrameworkEventBridge(
run_id="test_run",
thread_id="test_thread",
predict_state_config=predict_config,
)
# Simulate streaming tool call - first chunk with name
tool_call_start = Content.from_function_call(
call_id="call_123",
name="write_document_local",
arguments='{"document":"Once',
)
update1 = AgentResponseUpdate(contents=[tool_call_start])
events1 = await bridge.from_agent_run_update(update1)
# Should have ToolCallStartEvent and ToolCallArgsEvent
assert any(e.type == EventType.TOOL_CALL_START for e in events1)
assert any(e.type == EventType.TOOL_CALL_ARGS for e in events1)
# Second chunk - incomplete JSON, should try partial extraction
tool_call_chunk2 = Content.from_function_call(
call_id="call_123", name="write_document_local", arguments=" upon a time"
)
update2 = AgentResponseUpdate(contents=[tool_call_chunk2])
events2 = await bridge.from_agent_run_update(update2)
# Should emit StateDeltaEvent with partial document
state_deltas = [e for e in events2 if isinstance(e, StateDeltaEvent)]
assert len(state_deltas) >= 1
# Check JSON Patch format
delta = state_deltas[0]
assert isinstance(delta.delta, list)
assert len(delta.delta) > 0
assert delta.delta[0]["op"] == "replace"
assert delta.delta[0]["path"] == "/document"
assert "Once upon a time" in delta.delta[0]["value"]
async def test_confirm_changes_emission():
"""Test that confirm_changes tool call is emitted after predictive tool completion."""
predict_config = {
"document": {"tool": "write_document_local", "tool_argument": "document"},
}
current_state: dict[str, str] = {}
bridge = AgentFrameworkEventBridge(
run_id="test_run",
thread_id="test_thread",
predict_state_config=predict_config,
current_state=current_state,
)
# Set current tool name (simulating earlier tool call start)
bridge.current_tool_call_name = "write_document_local"
bridge.pending_state_updates = {"document": "A short story"}
# Tool result
tool_result = Content.from_function_result(
call_id="call_123",
result="Document written.",
)
update = AgentResponseUpdate(contents=[tool_result])
events = await bridge.from_agent_run_update(update)
# Should have: ToolCallEndEvent, ToolCallResultEvent, StateSnapshotEvent, confirm_changes sequence
assert any(e.type == EventType.TOOL_CALL_END for e in events)
assert any(e.type == EventType.TOOL_CALL_RESULT for e in events)
assert any(e.type == EventType.STATE_SNAPSHOT for e in events)
# Check for confirm_changes tool call
confirm_starts = [e for e in events if isinstance(e, ToolCallStartEvent) and e.tool_call_name == "confirm_changes"]
assert len(confirm_starts) == 1
confirm_args = [e for e in events if isinstance(e, ToolCallArgsEvent) and e.delta == "{}"]
assert len(confirm_args) >= 1
confirm_ends = [e for e in events if isinstance(e, ToolCallEndEvent)]
# At least 2: one for write_document_local, one for confirm_changes
assert len(confirm_ends) >= 2
# Check that stop flag is set
assert bridge.should_stop_after_confirm is True
async def test_text_suppression_before_confirm():
"""Test that text messages are suppressed when confirm_changes is pending."""
predict_config = {
"document": {"tool": "write_document_local", "tool_argument": "document"},
}
bridge = AgentFrameworkEventBridge(
run_id="test_run",
thread_id="test_thread",
predict_state_config=predict_config,
)
# Set flag indicating we're waiting for confirmation
bridge.should_stop_after_confirm = True
# Text content that should be suppressed
text = Content.from_text(text="I have written a story about pirates.")
update = AgentResponseUpdate(contents=[text])
events = await bridge.from_agent_run_update(update)
# Should NOT emit TextMessageContentEvent
text_events = [e for e in events if e.type == EventType.TEXT_MESSAGE_CONTENT]
assert len(text_events) == 0
# But should save the text
assert bridge.suppressed_summary == "I have written a story about pirates."
async def test_no_confirm_for_non_predictive_tools():
"""Test that confirm_changes is NOT emitted for regular tool calls."""
predict_config = {
"document": {"tool": "write_document_local", "tool_argument": "document"},
}
current_state: dict[str, str] = {}
bridge = AgentFrameworkEventBridge(
run_id="test_run",
thread_id="test_thread",
predict_state_config=predict_config,
current_state=current_state,
)
# Different tool (not in predict_state_config)
bridge.current_tool_call_name = "get_weather"
tool_result = Content.from_function_result(
call_id="call_456",
result="Sunny, 72°F",
)
update = AgentResponseUpdate(contents=[tool_result])
events = await bridge.from_agent_run_update(update)
# Should NOT have confirm_changes
confirm_starts = [e for e in events if isinstance(e, ToolCallStartEvent) and e.tool_call_name == "confirm_changes"]
assert len(confirm_starts) == 0
# Stop flag should NOT be set
assert bridge.should_stop_after_confirm is False
async def test_state_delta_deduplication():
"""Test that duplicate state values don't emit multiple StateDeltaEvents."""
predict_config = {
"document": {"tool": "write_document_local", "tool_argument": "document"},
}
bridge = AgentFrameworkEventBridge(
run_id="test_run",
thread_id="test_thread",
predict_state_config=predict_config,
)
# First tool call with document
tool_call1 = Content.from_function_call(
call_id="call_1",
name="write_document_local",
arguments='{"document":"Same text"}',
)
update1 = AgentResponseUpdate(contents=[tool_call1])
events1 = await bridge.from_agent_run_update(update1)
# Count state deltas
state_deltas_1 = [e for e in events1 if isinstance(e, StateDeltaEvent)]
assert len(state_deltas_1) >= 1
# Second tool call with SAME document (shouldn't emit new delta)
bridge.current_tool_call_name = "write_document_local"
tool_call2 = Content.from_function_call(
call_id="call_2",
name="write_document_local",
arguments='{"document":"Same text"}', # Identical content
)
update2 = AgentResponseUpdate(contents=[tool_call2])
events2 = await bridge.from_agent_run_update(update2)
# Should NOT emit state delta (same value)
state_deltas_2 = [e for e in events2 if e.type == EventType.STATE_DELTA]
assert len(state_deltas_2) == 0
async def test_predict_state_config_multiple_fields():
"""Test predictive state with multiple state fields."""
predict_config = {
"title": {"tool": "create_post", "tool_argument": "title"},
"content": {"tool": "create_post", "tool_argument": "body"},
}
bridge = AgentFrameworkEventBridge(
run_id="test_run",
thread_id="test_thread",
predict_state_config=predict_config,
)
# Tool call with both fields
tool_call = Content.from_function_call(
call_id="call_999",
name="create_post",
arguments='{"title":"My Post","body":"Post content"}',
)
update = AgentResponseUpdate(contents=[tool_call])
events = await bridge.from_agent_run_update(update)
# Should emit StateDeltaEvent for both fields
state_deltas = [e for e in events if isinstance(e, StateDeltaEvent)]
assert len(state_deltas) >= 2
# Check both fields are present
paths = [delta.delta[0]["path"] for delta in state_deltas]
assert "/title" in paths
assert "/content" in paths
python/packages/ag-ui/tests/test_events_comprehensive.py
-917
View File
@@ -1,917 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""Comprehensive tests for AgentFrameworkEventBridge (_events.py)."""
import json
from agent_framework import (
AgentResponseUpdate,
Content,
)
async def test_basic_text_message_conversion():
"""Test basic TextContent to AG-UI events."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
update = AgentResponseUpdate(contents=[Content.from_text(text="Hello")])
events = await bridge.from_agent_run_update(update)
assert len(events) == 2
assert events[0].type == "TEXT_MESSAGE_START"
assert events[0].role == "assistant"
assert events[1].type == "TEXT_MESSAGE_CONTENT"
assert events[1].delta == "Hello"
async def test_text_message_streaming():
"""Test streaming TextContent with multiple chunks."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
update1 = AgentResponseUpdate(contents=[Content.from_text(text="Hello ")])
update2 = AgentResponseUpdate(contents=[Content.from_text(text="world")])
events1 = await bridge.from_agent_run_update(update1)
events2 = await bridge.from_agent_run_update(update2)
# First update: START + CONTENT
assert len(events1) == 2
assert events1[0].type == "TEXT_MESSAGE_START"
assert events1[1].delta == "Hello "
# Second update: just CONTENT (same message)
assert len(events2) == 1
assert events2[0].type == "TEXT_MESSAGE_CONTENT"
assert events2[0].delta == "world"
# Both content events should have same message_id
assert events1[1].message_id == events2[0].message_id
async def test_skip_text_content_for_structured_outputs():
"""Test that text content is skipped when skip_text_content=True."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread", skip_text_content=True)
update = AgentResponseUpdate(contents=[Content.from_text(text='{"result": "data"}')])
events = await bridge.from_agent_run_update(update)
# No events should be emitted
assert len(events) == 0
async def test_skip_text_content_for_empty_text():
"""Test streaming TextContent with empty chunks."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
update1 = AgentResponseUpdate(contents=[Content.from_text(text="Hello ")])
update2 = AgentResponseUpdate(contents=[Content.from_text(text="")]) # Empty chunk
update3 = AgentResponseUpdate(contents=[Content.from_text(text="world")])
events1 = await bridge.from_agent_run_update(update1)
events2 = await bridge.from_agent_run_update(update2)
events3 = await bridge.from_agent_run_update(update3)
# First update: START + CONTENT
assert len(events1) == 2
assert events1[0].type == "TEXT_MESSAGE_START"
assert events1[1].delta == "Hello "
# Second update: should skip empty chunk, no events
assert len(events2) == 0
# Third update: just CONTENT (same message)
assert len(events3) == 1
assert events3[0].type == "TEXT_MESSAGE_CONTENT"
assert events3[0].delta == "world"
# Both content events should have same message_id
assert events1[1].message_id == events3[0].message_id
async def test_tool_call_with_name():
"""Test FunctionCallContent with name emits ToolCallStartEvent."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
update = AgentResponseUpdate(contents=[Content.from_function_call(name="search_web", call_id="call_123")])
events = await bridge.from_agent_run_update(update)
assert len(events) == 1
assert events[0].type == "TOOL_CALL_START"
assert events[0].tool_call_name == "search_web"
assert events[0].tool_call_id == "call_123"
async def test_tool_call_streaming_args():
"""Test streaming tool call arguments."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
# First chunk: name only
update1 = AgentResponseUpdate(contents=[Content.from_function_call(name="search_web", call_id="call_123")])
events1 = await bridge.from_agent_run_update(update1)
# Second chunk: arguments chunk 1 (name can be empty string for continuation)
update2 = AgentResponseUpdate(
contents=[Content.from_function_call(name="", call_id="call_123", arguments='{"query": "')]
)
events2 = await bridge.from_agent_run_update(update2)
# Third chunk: arguments chunk 2
update3 = AgentResponseUpdate(contents=[Content.from_function_call(name="", call_id="call_123", arguments='AI"}')])
events3 = await bridge.from_agent_run_update(update3)
# First update: ToolCallStartEvent
assert len(events1) == 1
assert events1[0].type == "TOOL_CALL_START"
# Second update: ToolCallArgsEvent
assert len(events2) == 1
assert events2[0].type == "TOOL_CALL_ARGS"
assert events2[0].delta == '{"query": "'
# Third update: ToolCallArgsEvent
assert len(events3) == 1
assert events3[0].type == "TOOL_CALL_ARGS"
assert events3[0].delta == 'AI"}'
# All should have same tool_call_id
assert events1[0].tool_call_id == events2[0].tool_call_id == events3[0].tool_call_id
async def test_streaming_tool_call_no_duplicate_start_events():
"""Test that streaming tool calls emit exactly one ToolCallStartEvent.
This is a regression test for the Anthropic streaming fix where input_json_delta
events were incorrectly passing the tool name, causing duplicate ToolCallStartEvents.
The correct behavior is:
- Initial FunctionCallContent with name -> emits ToolCallStartEvent
- Subsequent FunctionCallContent with name="" -> emits only ToolCallArgsEvent
See: https://github.com/microsoft/agent-framework/pull/3051
"""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
# Simulate streaming tool call: first chunk has name, subsequent chunks have name=""
update1 = AgentResponseUpdate(contents=[Content.from_function_call(name="get_weather", call_id="call_789")])
update2 = AgentResponseUpdate(
contents=[Content.from_function_call(name="", call_id="call_789", arguments='{"loc":')]
)
update3 = AgentResponseUpdate(contents=[Content.from_function_call(name="", call_id="call_789", arguments='"SF"}')])
events1 = await bridge.from_agent_run_update(update1)
events2 = await bridge.from_agent_run_update(update2)
events3 = await bridge.from_agent_run_update(update3)
# Count all ToolCallStartEvents - should be exactly 1
all_events = events1 + events2 + events3
tool_call_start_count = sum(1 for e in all_events if e.type == "TOOL_CALL_START")
assert tool_call_start_count == 1, f"Expected 1 ToolCallStartEvent, got {tool_call_start_count}"
# Verify event types
assert events1[0].type == "TOOL_CALL_START"
assert events2[0].type == "TOOL_CALL_ARGS"
assert events3[0].type == "TOOL_CALL_ARGS"
async def test_tool_result_with_dict():
"""Test FunctionResultContent with dict result."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
result_data = {"status": "success", "count": 42}
update = AgentResponseUpdate(contents=[Content.from_function_result(call_id="call_123", result=result_data)])
events = await bridge.from_agent_run_update(update)
# Should emit ToolCallEndEvent + ToolCallResultEvent
assert len(events) == 2
assert events[0].type == "TOOL_CALL_END"
assert events[0].tool_call_id == "call_123"
assert events[1].type == "TOOL_CALL_RESULT"
assert events[1].tool_call_id == "call_123"
assert events[1].role == "tool"
# Result should be JSON-serialized
assert json.loads(events[1].content) == result_data
async def test_tool_result_with_string():
"""Test FunctionResultContent with string result."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
update = AgentResponseUpdate(contents=[Content.from_function_result(call_id="call_123", result="Search complete")])
events = await bridge.from_agent_run_update(update)
assert len(events) == 2
assert events[0].type == "TOOL_CALL_END"
assert events[1].type == "TOOL_CALL_RESULT"
assert events[1].content == "Search complete"
async def test_tool_result_with_none():
"""Test FunctionResultContent with None result."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
update = AgentResponseUpdate(contents=[Content.from_function_result(call_id="call_123", result=None)])
events = await bridge.from_agent_run_update(update)
assert len(events) == 2
assert events[0].type == "TOOL_CALL_END"
assert events[1].type == "TOOL_CALL_RESULT"
# prepare_function_call_results serializes None as JSON "null"
assert events[1].content == "null"
async def test_multiple_tool_results_in_sequence():
"""Test multiple tool results processed sequentially."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
update = AgentResponseUpdate(
contents=[
Content.from_function_result(call_id="call_1", result="Result 1"),
Content.from_function_result(call_id="call_2", result="Result 2"),
]
)
events = await bridge.from_agent_run_update(update)
# Each result emits: ToolCallEndEvent + ToolCallResultEvent = 4 events total
assert len(events) == 4
assert events[0].tool_call_id == "call_1"
assert events[1].tool_call_id == "call_1"
assert events[2].tool_call_id == "call_2"
assert events[3].tool_call_id == "call_2"
async def test_function_approval_request_basic():
"""Test FunctionApprovalRequestContent conversion."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
# Set require_confirmation=False to test just the function_approval_request event
bridge = AgentFrameworkEventBridge(
run_id="test_run",
thread_id="test_thread",
require_confirmation=False,
)
func_call = Content.from_function_call(
call_id="call_123",
name="send_email",
arguments={"to": "user@example.com", "subject": "Test"},
)
approval = Content.from_function_approval_request(
id="approval_001",
function_call=func_call,
)
update = AgentResponseUpdate(contents=[approval])
events = await bridge.from_agent_run_update(update)
# Should emit: ToolCallEndEvent + CustomEvent
assert len(events) == 2
# First: ToolCallEndEvent to close the tool call
assert events[0].type == "TOOL_CALL_END"
assert events[0].tool_call_id == "call_123"
# Second: CustomEvent with approval details
assert events[1].type == "CUSTOM"
assert events[1].name == "function_approval_request"
assert events[1].value["id"] == "approval_001"
assert events[1].value["function_call"]["name"] == "send_email"
async def test_empty_predict_state_config():
"""Test behavior with no predictive state configuration."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(
run_id="test_run",
thread_id="test_thread",
predict_state_config={}, # Empty config
)
# Tool call with arguments
update = AgentResponseUpdate(
contents=[
Content.from_function_call(name="write_doc", call_id="call_1", arguments='{"content": "test"}'),
Content.from_function_result(call_id="call_1", result="Done"),
]
)
events = await bridge.from_agent_run_update(update)
# Should NOT emit StateDeltaEvent or confirm_changes
event_types = [e.type for e in events]
assert "STATE_DELTA" not in event_types
assert "STATE_SNAPSHOT" not in event_types
# Should have: ToolCallStart, ToolCallArgs, ToolCallEnd, ToolCallResult
assert event_types == [
"TOOL_CALL_START",
"TOOL_CALL_ARGS",
"TOOL_CALL_END",
"TOOL_CALL_RESULT",
]
async def test_tool_not_in_predict_state_config():
"""Test tool that doesn't match any predict_state_config entry."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(
run_id="test_run",
thread_id="test_thread",
predict_state_config={
"document": {"tool": "write_document", "tool_argument": "content"},
},
)
# Different tool name
update = AgentResponseUpdate(
contents=[
Content.from_function_call(name="search_web", call_id="call_1", arguments='{"query": "AI"}'),
Content.from_function_result(call_id="call_1", result="Results"),
]
)
events = await bridge.from_agent_run_update(update)
# Should NOT emit StateDeltaEvent or confirm_changes
event_types = [e.type for e in events]
assert "STATE_DELTA" not in event_types
assert "STATE_SNAPSHOT" not in event_types
async def test_state_management_tracking():
"""Test current_state and pending_state_updates tracking."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
initial_state = {"document": ""}
bridge = AgentFrameworkEventBridge(
run_id="test_run",
thread_id="test_thread",
predict_state_config={
"document": {"tool": "write_doc", "tool_argument": "content"},
},
current_state=initial_state,
)
# Streaming tool call
update1 = AgentResponseUpdate(
contents=[
Content.from_function_call(name="write_doc", call_id="call_1"),
Content.from_function_call(name="", call_id="call_1", arguments='{"content": "Hello"}'),
]
)
await bridge.from_agent_run_update(update1)
# Check pending_state_updates was populated
assert "document" in bridge.pending_state_updates
assert bridge.pending_state_updates["document"] == "Hello"
# Tool result should update current_state
update2 = AgentResponseUpdate(contents=[Content.from_function_result(call_id="call_1", result="Done")])
await bridge.from_agent_run_update(update2)
# current_state should be updated
assert bridge.current_state["document"] == "Hello"
# pending_state_updates should be cleared
assert len(bridge.pending_state_updates) == 0
async def test_wildcard_tool_argument():
"""Test tool_argument='*' uses all arguments as state value."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(
run_id="test_run",
thread_id="test_thread",
predict_state_config={
"recipe": {"tool": "create_recipe", "tool_argument": "*"},
},
current_state={},
)
# Complete tool call with dict arguments
update = AgentResponseUpdate(
contents=[
Content.from_function_call(
name="create_recipe",
call_id="call_1",
arguments={"title": "Pasta", "ingredients": ["pasta", "sauce"]},
),
Content.from_function_result(call_id="call_1", result="Created"),
]
)
events = await bridge.from_agent_run_update(update)
# Find StateDeltaEvent
delta_events = [e for e in events if e.type == "STATE_DELTA"]
assert len(delta_events) > 0
# Value should be the entire arguments dict
delta = delta_events[0].delta[0]
assert delta["path"] == "/recipe"
assert delta["value"] == {"title": "Pasta", "ingredients": ["pasta", "sauce"]}
async def test_run_lifecycle_events():
"""Test RunStartedEvent and RunFinishedEvent creation."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
started = bridge.create_run_started_event()
assert started.type == "RUN_STARTED"
assert started.run_id == "test_run"
assert started.thread_id == "test_thread"
finished = bridge.create_run_finished_event(result={"status": "complete"})
assert finished.type == "RUN_FINISHED"
assert finished.run_id == "test_run"
assert finished.thread_id == "test_thread"
assert finished.result == {"status": "complete"}
async def test_message_lifecycle_events():
"""Test TextMessageStartEvent and TextMessageEndEvent creation."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
start = bridge.create_message_start_event("msg_123", role="assistant")
assert start.type == "TEXT_MESSAGE_START"
assert start.message_id == "msg_123"
assert start.role == "assistant"
end = bridge.create_message_end_event("msg_123")
assert end.type == "TEXT_MESSAGE_END"
assert end.message_id == "msg_123"
async def test_state_event_creation():
"""Test StateSnapshotEvent and StateDeltaEvent creation helpers."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
# StateSnapshotEvent
snapshot = bridge.create_state_snapshot_event({"document": "content"})
assert snapshot.type == "STATE_SNAPSHOT"
assert snapshot.snapshot == {"document": "content"}
# StateDeltaEvent with JSON Patch
delta = bridge.create_state_delta_event([{"op": "replace", "path": "/document", "value": "new content"}])
assert delta.type == "STATE_DELTA"
assert len(delta.delta) == 1
assert delta.delta[0]["op"] == "replace"
assert delta.delta[0]["path"] == "/document"
assert delta.delta[0]["value"] == "new content"
async def test_state_snapshot_after_tool_result():
"""Test StateSnapshotEvent emission after tool result with pending updates."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(
run_id="test_run",
thread_id="test_thread",
predict_state_config={
"document": {"tool": "write_doc", "tool_argument": "content"},
},
current_state={"document": ""},
)
# Tool call with streaming args
update1 = AgentResponseUpdate(
contents=[
Content.from_function_call(name="write_doc", call_id="call_1"),
Content.from_function_call(name="", call_id="call_1", arguments='{"content": "Test"}'),
]
)
await bridge.from_agent_run_update(update1)
# Tool result should trigger StateSnapshotEvent
update2 = AgentResponseUpdate(contents=[Content.from_function_result(call_id="call_1", result="Done")])
events = await bridge.from_agent_run_update(update2)
# Should have: ToolCallEnd, ToolCallResult, StateSnapshot, ToolCallStart (confirm_changes), ToolCallArgs, ToolCallEnd
snapshot_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
assert len(snapshot_events) == 1
assert snapshot_events[0].snapshot["document"] == "Test"
async def test_message_id_persistence_across_chunks():
"""Test that message_id persists across multiple text chunks."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
# First chunk
update1 = AgentResponseUpdate(contents=[Content.from_text(text="Hello ")])
events1 = await bridge.from_agent_run_update(update1)
message_id = events1[0].message_id
# Second chunk
update2 = AgentResponseUpdate(contents=[Content.from_text(text="world")])
events2 = await bridge.from_agent_run_update(update2)
# Should use same message_id
assert events2[0].message_id == message_id
assert bridge.current_message_id == message_id
async def test_tool_call_id_tracking():
"""Test tool_call_id tracking across streaming chunks."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
# First chunk with name
update1 = AgentResponseUpdate(contents=[Content.from_function_call(name="search", call_id="call_1")])
await bridge.from_agent_run_update(update1)
assert bridge.current_tool_call_id == "call_1"
assert bridge.current_tool_call_name == "search"
# Second chunk with args but no name
update2 = AgentResponseUpdate(
contents=[Content.from_function_call(name="", call_id="call_1", arguments='{"q":"AI"}')]
)
events2 = await bridge.from_agent_run_update(update2)
# Should still track same tool call
assert bridge.current_tool_call_id == "call_1"
assert events2[0].tool_call_id == "call_1"
async def test_tool_name_reset_after_result():
"""Test current_tool_call_name is reset after tool result."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(
run_id="test_run",
thread_id="test_thread",
predict_state_config={
"document": {"tool": "write_doc", "tool_argument": "content"},
},
)
# Tool call
update1 = AgentResponseUpdate(
contents=[
Content.from_function_call(name="write_doc", call_id="call_1"),
Content.from_function_call(name="", call_id="call_1", arguments='{"content": "Test"}'),
]
)
await bridge.from_agent_run_update(update1)
assert bridge.current_tool_call_name == "write_doc"
# Tool result with predictive state (should trigger confirm_changes and reset)
update2 = AgentResponseUpdate(contents=[Content.from_function_result(call_id="call_1", result="Done")])
await bridge.from_agent_run_update(update2)
# Tool name should be reset
assert bridge.current_tool_call_name is None
async def test_function_approval_with_wildcard_argument():
"""Test function approval with wildcard * argument."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(
run_id="test_run",
thread_id="test_thread",
predict_state_config={
"payload": {"tool": "submit", "tool_argument": "*"},
},
)
approval_content = Content.from_function_approval_request(
id="approval_1",
function_call=Content.from_function_call(
name="submit", call_id="call_1", arguments='{"key1": "value1", "key2": "value2"}'
),
)
update = AgentResponseUpdate(contents=[approval_content])
events = await bridge.from_agent_run_update(update)
# Should emit StateSnapshotEvent with entire parsed args as value
snapshot_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
assert len(snapshot_events) == 1
assert snapshot_events[0].snapshot["payload"] == {"key1": "value1", "key2": "value2"}
async def test_function_approval_missing_argument():
"""Test function approval when specified argument is not in parsed args."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(
run_id="test_run",
thread_id="test_thread",
predict_state_config={
"data": {"tool": "process", "tool_argument": "missing_field"},
},
)
approval_content = Content.from_function_approval_request(
id="approval_1",
function_call=Content.from_function_call(
name="process", call_id="call_1", arguments='{"other_field": "value"}'
),
)
update = AgentResponseUpdate(contents=[approval_content])
events = await bridge.from_agent_run_update(update)
# Should not emit StateSnapshotEvent since argument not found
snapshot_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
assert len(snapshot_events) == 0
async def test_empty_predict_state_config_no_deltas():
"""Test with empty predict_state_config (no predictive updates)."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread", predict_state_config={})
# Tool call with arguments
update = AgentResponseUpdate(
contents=[
Content.from_function_call(name="search", call_id="call_1"),
Content.from_function_call(name="", call_id="call_1", arguments='{"query": "test"}'),
]
)
events = await bridge.from_agent_run_update(update)
# Should not emit any StateDeltaEvents
delta_events = [e for e in events if e.type == "STATE_DELTA"]
assert len(delta_events) == 0
async def test_tool_with_no_matching_config():
"""Test tool call for tool not in predict_state_config."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(
run_id="test_run",
thread_id="test_thread",
predict_state_config={"document": {"tool": "write_doc", "tool_argument": "content"}},
)
# Tool call for different tool
update = AgentResponseUpdate(
contents=[
Content.from_function_call(name="search_web", call_id="call_1"),
Content.from_function_call(name="", call_id="call_1", arguments='{"query": "test"}'),
]
)
events = await bridge.from_agent_run_update(update)
# Should not emit StateDeltaEvents
delta_events = [e for e in events if e.type == "STATE_DELTA"]
assert len(delta_events) == 0
async def test_tool_call_without_name_or_id():
"""Test handling FunctionCallContent with no name and no call_id."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
# This should not crash but log an error
update = AgentResponseUpdate(contents=[Content.from_function_call(name="", call_id="", arguments='{"arg": "val"}')])
events = await bridge.from_agent_run_update(update)
# Should emit ToolCallArgsEvent with generated ID
assert len(events) >= 1
async def test_state_delta_count_logging():
"""Test that state delta count increments and logs at intervals."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(
run_id="test_run",
thread_id="test_thread",
predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}},
)
# Emit multiple state deltas with different content each time
for i in range(15):
update = AgentResponseUpdate(
contents=[
Content.from_function_call(name="", call_id="call_1", arguments=f'{{"text": "Content variation {i}"}}'),
]
)
# Set the tool name to match config
bridge.current_tool_call_name = "write"
await bridge.from_agent_run_update(update)
# State delta count should have incremented (one per unique state update)
assert bridge.state_delta_count >= 1
# Tests for list type tool results (MCP tool serialization)
async def test_tool_result_with_empty_list():
"""Test FunctionResultContent with empty list result."""
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
update = AgentResponseUpdate(contents=[Content.from_function_result(call_id="call_123", result=[])])
events = await bridge.from_agent_run_update(update)
assert len(events) == 2
assert events[0].type == "TOOL_CALL_END"
assert events[1].type == "TOOL_CALL_RESULT"
# Empty list serializes as JSON empty array
assert events[1].content == "[]"
async def test_tool_result_with_single_text_content():
"""Test FunctionResultContent with single TextContent-like item (MCP tool result)."""
from dataclasses import dataclass
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
@dataclass
class MockTextContent:
text: str
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
update = AgentResponseUpdate(
contents=[Content.from_function_result(call_id="call_123", result=[MockTextContent("Hello from MCP tool!")])]
)
events = await bridge.from_agent_run_update(update)
assert len(events) == 2
assert events[0].type == "TOOL_CALL_END"
assert events[1].type == "TOOL_CALL_RESULT"
# TextContent text is extracted and serialized as JSON array
assert events[1].content == '["Hello from MCP tool!"]'
async def test_tool_result_with_multiple_text_contents():
"""Test FunctionResultContent with multiple TextContent-like items (MCP tool result)."""
from dataclasses import dataclass
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
@dataclass
class MockTextContent:
text: str
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
update = AgentResponseUpdate(
contents=[
Content.from_function_result(
call_id="call_123",
result=[MockTextContent("First result"), MockTextContent("Second result")],
)
]
)
events = await bridge.from_agent_run_update(update)
assert len(events) == 2
assert events[0].type == "TOOL_CALL_END"
assert events[1].type == "TOOL_CALL_RESULT"
# Multiple TextContent items should return JSON array
assert events[1].content == '["First result", "Second result"]'
async def test_tool_result_with_model_dump_objects():
"""Test FunctionResultContent with Pydantic BaseModel objects."""
from pydantic import BaseModel
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
class MockModel(BaseModel):
value: int
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
update = AgentResponseUpdate(
contents=[Content.from_function_result(call_id="call_123", result=[MockModel(value=1), MockModel(value=2)])]
)
events = await bridge.from_agent_run_update(update)
assert len(events) == 2
assert events[1].type == "TOOL_CALL_RESULT"
# Should be properly serialized JSON array without double escaping
assert events[1].content == '[{"value": 1}, {"value": 2}]'
async def test_function_call_with_dataclass_arguments():
"""Test FunctionCallContent with dataclass arguments is serialized correctly.
This test verifies the fix for the AG-UI JSON serialization error when
HandoffAgentUserRequest (a dataclass) is passed as FunctionCallContent.arguments.
"""
from dataclasses import dataclass
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
@dataclass
class TestRequest:
field1: str
field2: int
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
# FunctionCallContent with a dataclass as arguments (not a string)
update = AgentResponseUpdate(
contents=[
Content.from_function_call(
name="request_info",
call_id="call_dataclass",
arguments=TestRequest(field1="value", field2=42),
)
]
)
events = await bridge.from_agent_run_update(update)
# Should have ToolCallStartEvent and ToolCallArgsEvent
tool_args_events = [e for e in events if e.type == "TOOL_CALL_ARGS"]
assert len(tool_args_events) == 1
# Verify the delta is valid JSON
delta = tool_args_events[0].delta
parsed = json.loads(delta)
assert parsed == {"field1": "value", "field2": 42}
async def test_function_call_with_nested_dataclass_arguments():
"""Test FunctionCallContent with nested dataclass arguments is serialized correctly.
This test covers the scenario where HandoffAgentUserRequest contains an AgentResponse
with nested content objects.
"""
from dataclasses import dataclass
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
@dataclass
class InnerContent:
text: str
@dataclass
class AgentResponseMock:
contents: list[InnerContent]
@dataclass
class HandoffRequest:
agent_response: AgentResponseMock
bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
# Simulate a HandoffAgentUserRequest-like structure
update = AgentResponseUpdate(
contents=[
Content.from_function_call(
name="request_info",
call_id="call_nested",
arguments=HandoffRequest(
agent_response=AgentResponseMock(contents=[InnerContent(text="Hello from agent")])
),
)
]
)
events = await bridge.from_agent_run_update(update)
# Should have ToolCallStartEvent and ToolCallArgsEvent
tool_args_events = [e for e in events if e.type == "TOOL_CALL_ARGS"]
assert len(tool_args_events) == 1
# Verify the delta is valid JSON and contains nested structure
delta = tool_args_events[0].delta
parsed = json.loads(delta)
assert "agent_response" in parsed
assert parsed["agent_response"]["contents"] == [{"text": "Hello from agent"}]
python/packages/ag-ui/tests/test_helpers.py
+502
View File
@@ -0,0 +1,502 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for orchestration helper functions."""
from agent_framework import ChatMessage, Content
from agent_framework_ag_ui._orchestration._helpers import (
approval_steps,
build_safe_metadata,
ensure_tool_call_entry,
is_state_context_message,
is_step_based_approval,
latest_approval_response,
pending_tool_call_ids,
schema_has_steps,
select_approval_tool_name,
tool_name_for_call_id,
)
class TestPendingToolCallIds:
"""Tests for pending_tool_call_ids function."""
def test_empty_messages(self):
"""Returns empty set for empty messages list."""
result = pending_tool_call_ids([])
assert result == set()
def test_no_tool_calls(self):
"""Returns empty set when no tool calls in messages."""
messages = [
ChatMessage(role="user", contents=[Content.from_text("Hello")]),
ChatMessage(role="assistant", contents=[Content.from_text("Hi there")]),
]
result = pending_tool_call_ids(messages)
assert result == set()
def test_pending_tool_call(self):
"""Returns pending tool call ID when no result exists."""
messages = [
ChatMessage(
role="assistant",
contents=[Content.from_function_call(call_id="call_123", name="get_weather", arguments="{}")],
),
]
result = pending_tool_call_ids(messages)
assert result == {"call_123"}
def test_resolved_tool_call(self):
"""Returns empty set when tool call has result."""
messages = [
ChatMessage(
role="assistant",
contents=[Content.from_function_call(call_id="call_123", name="get_weather", arguments="{}")],
),
ChatMessage(
role="tool",
contents=[Content.from_function_result(call_id="call_123", result="sunny")],
),
]
result = pending_tool_call_ids(messages)
assert result == set()
def test_multiple_tool_calls_some_resolved(self):
"""Returns only unresolved tool call IDs."""
messages = [
ChatMessage(
role="assistant",
contents=[
Content.from_function_call(call_id="call_1", name="tool_a", arguments="{}"),
Content.from_function_call(call_id="call_2", name="tool_b", arguments="{}"),
Content.from_function_call(call_id="call_3", name="tool_c", arguments="{}"),
],
),
ChatMessage(
role="tool",
contents=[Content.from_function_result(call_id="call_1", result="result_a")],
),
ChatMessage(
role="tool",
contents=[Content.from_function_result(call_id="call_3", result="result_c")],
),
]
result = pending_tool_call_ids(messages)
assert result == {"call_2"}
class TestIsStateContextMessage:
"""Tests for is_state_context_message function."""
def test_state_context_message(self):
"""Returns True for state context message."""
message = ChatMessage(
role="system",
contents=[Content.from_text("Current state of the application: {}")],
)
assert is_state_context_message(message) is True
def test_non_system_message(self):
"""Returns False for non-system message."""
message = ChatMessage(
role="user",
contents=[Content.from_text("Current state of the application: {}")],
)
assert is_state_context_message(message) is False
def test_system_message_without_state_prefix(self):
"""Returns False for system message without state prefix."""
message = ChatMessage(
role="system",
contents=[Content.from_text("You are a helpful assistant.")],
)
assert is_state_context_message(message) is False
def test_empty_contents(self):
"""Returns False for message with empty contents."""
message = ChatMessage(role="system", contents=[])
assert is_state_context_message(message) is False
class TestEnsureToolCallEntry:
"""Tests for ensure_tool_call_entry function."""
def test_creates_new_entry(self):
"""Creates new entry when ID not found."""
tool_calls_by_id: dict = {}
pending_tool_calls: list = []
entry = ensure_tool_call_entry("call_123", tool_calls_by_id, pending_tool_calls)
assert entry["id"] == "call_123"
assert entry["type"] == "function"
assert entry["function"]["name"] == ""
assert entry["function"]["arguments"] == ""
assert "call_123" in tool_calls_by_id
assert len(pending_tool_calls) == 1
def test_returns_existing_entry(self):
"""Returns existing entry when ID found."""
existing_entry = {
"id": "call_123",
"type": "function",
"function": {"name": "get_weather", "arguments": '{"city": "NYC"}'},
}
tool_calls_by_id = {"call_123": existing_entry}
pending_tool_calls: list = []
entry = ensure_tool_call_entry("call_123", tool_calls_by_id, pending_tool_calls)
assert entry is existing_entry
assert entry["function"]["name"] == "get_weather"
assert len(pending_tool_calls) == 0 # Not added again
class TestToolNameForCallId:
"""Tests for tool_name_for_call_id function."""
def test_returns_tool_name(self):
"""Returns tool name for valid entry."""
tool_calls_by_id = {
"call_123": {
"id": "call_123",
"function": {"name": "get_weather", "arguments": "{}"},
}
}
result = tool_name_for_call_id(tool_calls_by_id, "call_123")
assert result == "get_weather"
def test_returns_none_for_missing_id(self):
"""Returns None when ID not found."""
tool_calls_by_id: dict = {}
result = tool_name_for_call_id(tool_calls_by_id, "call_123")
assert result is None
def test_returns_none_for_missing_function(self):
"""Returns None when function key missing."""
tool_calls_by_id = {"call_123": {"id": "call_123"}}
result = tool_name_for_call_id(tool_calls_by_id, "call_123")
assert result is None
def test_returns_none_for_non_dict_function(self):
"""Returns None when function is not a dict."""
tool_calls_by_id = {"call_123": {"id": "call_123", "function": "not_a_dict"}}
result = tool_name_for_call_id(tool_calls_by_id, "call_123")
assert result is None
def test_returns_none_for_empty_name(self):
"""Returns None when name is empty."""
tool_calls_by_id = {"call_123": {"id": "call_123", "function": {"name": "", "arguments": "{}"}}}
result = tool_name_for_call_id(tool_calls_by_id, "call_123")
assert result is None
class TestSchemaHasSteps:
"""Tests for schema_has_steps function."""
def test_schema_with_steps_array(self):
"""Returns True when schema has steps array property."""
schema = {"properties": {"steps": {"type": "array"}}}
assert schema_has_steps(schema) is True
def test_schema_without_steps(self):
"""Returns False when schema doesn't have steps."""
schema = {"properties": {"name": {"type": "string"}}}
assert schema_has_steps(schema) is False
def test_schema_with_non_array_steps(self):
"""Returns False when steps is not array type."""
schema = {"properties": {"steps": {"type": "string"}}}
assert schema_has_steps(schema) is False
def test_non_dict_schema(self):
"""Returns False for non-dict schema."""
assert schema_has_steps(None) is False
assert schema_has_steps("not a dict") is False
assert schema_has_steps([]) is False
def test_missing_properties(self):
"""Returns False when properties key is missing."""
schema = {"type": "object"}
assert schema_has_steps(schema) is False
def test_non_dict_properties(self):
"""Returns False when properties is not a dict."""
schema = {"properties": "not a dict"}
assert schema_has_steps(schema) is False
def test_non_dict_steps(self):
"""Returns False when steps is not a dict."""
schema = {"properties": {"steps": "not a dict"}}
assert schema_has_steps(schema) is False
class TestSelectApprovalToolName:
"""Tests for select_approval_tool_name function."""
def test_none_client_tools(self):
"""Returns None when client_tools is None."""
result = select_approval_tool_name(None)
assert result is None
def test_empty_client_tools(self):
"""Returns None when client_tools is empty."""
result = select_approval_tool_name([])
assert result is None
def test_finds_approval_tool(self):
"""Returns tool name when tool has steps schema."""
class MockTool:
name = "generate_task_steps"
def parameters(self):
return {"properties": {"steps": {"type": "array"}}}
result = select_approval_tool_name([MockTool()])
assert result == "generate_task_steps"
def test_skips_tool_without_name(self):
"""Skips tools without name attribute."""
class MockToolNoName:
def parameters(self):
return {"properties": {"steps": {"type": "array"}}}
result = select_approval_tool_name([MockToolNoName()])
assert result is None
def test_skips_tool_without_parameters_method(self):
"""Skips tools without callable parameters method."""
class MockToolNoParams:
name = "some_tool"
parameters = "not callable"
result = select_approval_tool_name([MockToolNoParams()])
assert result is None
def test_skips_tool_without_steps_schema(self):
"""Skips tools that don't have steps in schema."""
class MockToolNoSteps:
name = "other_tool"
def parameters(self):
return {"properties": {"data": {"type": "string"}}}
result = select_approval_tool_name([MockToolNoSteps()])
assert result is None
class TestBuildSafeMetadata:
"""Tests for build_safe_metadata function."""
def test_none_metadata(self):
"""Returns empty dict for None metadata."""
result = build_safe_metadata(None)
assert result == {}
def test_empty_metadata(self):
"""Returns empty dict for empty metadata."""
result = build_safe_metadata({})
assert result == {}
def test_string_values_under_limit(self):
"""Preserves string values under 512 chars."""
metadata = {"key1": "short value", "key2": "another value"}
result = build_safe_metadata(metadata)
assert result == metadata
def test_truncates_long_string_values(self):
"""Truncates string values over 512 chars."""
long_value = "x" * 1000
metadata = {"key": long_value}
result = build_safe_metadata(metadata)
assert len(result["key"]) == 512
assert result["key"] == "x" * 512
def test_non_string_values_serialized(self):
"""Serializes non-string values to JSON."""
metadata = {"count": 42, "items": ["a", "b"]}
result = build_safe_metadata(metadata)
assert result["count"] == "42"
assert result["items"] == '["a", "b"]'
def test_truncates_serialized_values(self):
"""Truncates serialized JSON values over 512 chars."""
long_list = list(range(200)) # Will serialize to >512 chars
metadata = {"data": long_list}
result = build_safe_metadata(metadata)
assert len(result["data"]) == 512
class TestLatestApprovalResponse:
"""Tests for latest_approval_response function."""
def test_empty_messages(self):
"""Returns None for empty messages."""
result = latest_approval_response([])
assert result is None
def test_no_approval_response(self):
"""Returns None when no approval response in last message."""
messages = [
ChatMessage(role="assistant", contents=[Content.from_text("Hello")]),
]
result = latest_approval_response(messages)
assert result is None
def test_finds_approval_response(self):
"""Returns approval response from last message."""
# Create a function call content first
fc = Content.from_function_call(call_id="call_123", name="test_tool", arguments="{}")
approval_content = Content.from_function_approval_response(
approved=True,
id="approval_123",
function_call=fc,
)
messages = [
ChatMessage(role="user", contents=[approval_content]),
]
result = latest_approval_response(messages)
assert result is approval_content
class TestApprovalSteps:
"""Tests for approval_steps function."""
def test_steps_from_ag_ui_state_args(self):
"""Extracts steps from ag_ui_state_args."""
fc = Content.from_function_call(call_id="call_123", name="test_tool", arguments="{}")
approval = Content.from_function_approval_response(
approved=True,
id="approval_123",
function_call=fc,
additional_properties={"ag_ui_state_args": {"steps": [{"id": 1}, {"id": 2}]}},
)
result = approval_steps(approval)
assert result == [{"id": 1}, {"id": 2}]
def test_steps_from_function_call(self):
"""Extracts steps from function call arguments."""
fc = Content.from_function_call(
call_id="call_123",
name="test",
arguments='{"steps": [{"step": 1}]}',
)
approval = Content.from_function_approval_response(
approved=True,
id="approval_123",
function_call=fc,
)
result = approval_steps(approval)
assert result == [{"step": 1}]
def test_empty_steps_when_no_state_args(self):
"""Returns empty list when no ag_ui_state_args."""
fc = Content.from_function_call(call_id="call_123", name="test_tool", arguments="{}")
approval = Content.from_function_approval_response(
approved=True,
id="approval_123",
function_call=fc,
)
result = approval_steps(approval)
assert result == []
def test_empty_steps_when_state_args_not_dict(self):
"""Returns empty list when ag_ui_state_args is not a dict."""
fc = Content.from_function_call(call_id="call_123", name="test_tool", arguments="{}")
approval = Content.from_function_approval_response(
approved=True,
id="approval_123",
function_call=fc,
additional_properties={"ag_ui_state_args": "not a dict"},
)
result = approval_steps(approval)
assert result == []
def test_empty_steps_when_steps_not_list(self):
"""Returns empty list when steps is not a list."""
fc = Content.from_function_call(call_id="call_123", name="test_tool", arguments="{}")
approval = Content.from_function_approval_response(
approved=True,
id="approval_123",
function_call=fc,
additional_properties={"ag_ui_state_args": {"steps": "not a list"}},
)
result = approval_steps(approval)
assert result == []
class TestIsStepBasedApproval:
"""Tests for is_step_based_approval function."""
def test_returns_true_when_has_steps(self):
"""Returns True when approval has steps."""
fc = Content.from_function_call(call_id="call_123", name="test_tool", arguments="{}")
approval = Content.from_function_approval_response(
approved=True,
id="approval_123",
function_call=fc,
additional_properties={"ag_ui_state_args": {"steps": [{"id": 1}]}},
)
result = is_step_based_approval(approval, None)
assert result is True
def test_returns_false_no_steps_no_function_call(self):
"""Returns False when no steps and no function call."""
# Create content directly to have no function_call
approval = Content(
type="function_approval_response",
function_call=None,
)
result = is_step_based_approval(approval, None)
assert result is False
def test_returns_false_no_predict_config(self):
"""Returns False when no predict_state_config."""
fc = Content.from_function_call(call_id="call_123", name="some_tool", arguments="{}")
approval = Content.from_function_approval_response(
approved=True,
id="approval_123",
function_call=fc,
)
result = is_step_based_approval(approval, None)
assert result is False
def test_returns_true_when_tool_matches_config(self):
"""Returns True when tool matches predict_state_config with steps."""
fc = Content.from_function_call(call_id="call_123", name="generate_steps", arguments="{}")
approval = Content.from_function_approval_response(
approved=True,
id="approval_123",
function_call=fc,
)
config = {"steps": {"tool": "generate_steps", "tool_argument": "steps"}}
result = is_step_based_approval(approval, config)
assert result is True
def test_returns_false_when_tool_not_in_config(self):
"""Returns False when tool not in predict_state_config."""
fc = Content.from_function_call(call_id="call_123", name="other_tool", arguments="{}")
approval = Content.from_function_approval_response(
approved=True,
id="approval_123",
function_call=fc,
)
config = {"steps": {"tool": "generate_steps", "tool_argument": "steps"}}
result = is_step_based_approval(approval, config)
assert result is False
def test_returns_false_when_tool_arg_not_steps(self):
"""Returns False when tool_argument is not 'steps'."""
fc = Content.from_function_call(call_id="call_123", name="generate_steps", arguments="{}")
approval = Content.from_function_approval_response(
approved=True,
id="approval_123",
function_call=fc,
)
config = {"document": {"tool": "generate_steps", "tool_argument": "content"}}
result = is_step_based_approval(approval, config)
assert result is False
python/packages/ag-ui/tests/test_human_in_the_loop.py
-180
View File
@@ -1,180 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for human in the loop (function approval requests)."""
from agent_framework import AgentResponseUpdate, Content
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
async def test_function_approval_request_emission():
"""Test that CustomEvent is emitted for FunctionApprovalRequestContent."""
# Set require_confirmation=False to test just the function_approval_request event
bridge = AgentFrameworkEventBridge(
run_id="test_run",
thread_id="test_thread",
require_confirmation=False,
)
# Create approval request
func_call = Content.from_function_call(
call_id="call_123",
name="send_email",
arguments={"to": "user@example.com", "subject": "Test"},
)
approval_request = Content.from_function_approval_request(
id="approval_001",
function_call=func_call,
)
update = AgentResponseUpdate(contents=[approval_request])
events = await bridge.from_agent_run_update(update)
# Should emit ToolCallEndEvent + CustomEvent for approval request
assert len(events) == 2
# First event: ToolCallEndEvent to close the tool call
assert events[0].type == "TOOL_CALL_END"
assert events[0].tool_call_id == "call_123"
# Second event: CustomEvent with approval details
event = events[1]
assert event.type == "CUSTOM"
assert event.name == "function_approval_request"
assert event.value["id"] == "approval_001"
assert event.value["function_call"]["call_id"] == "call_123"
assert event.value["function_call"]["name"] == "send_email"
assert event.value["function_call"]["arguments"]["to"] == "user@example.com"
assert event.value["function_call"]["arguments"]["subject"] == "Test"
async def test_function_approval_request_with_confirm_changes():
"""Test that confirm_changes is also emitted when require_confirmation=True."""
bridge = AgentFrameworkEventBridge(
run_id="test_run",
thread_id="test_thread",
require_confirmation=True,
)
func_call = Content.from_function_call(
call_id="call_456",
name="delete_file",
arguments={"path": "/tmp/test.txt"},
)
approval_request = Content.from_function_approval_request(
id="approval_002",
function_call=func_call,
)
update = AgentResponseUpdate(contents=[approval_request])
events = await bridge.from_agent_run_update(update)
# Should emit: ToolCallEndEvent, CustomEvent, and confirm_changes (Start, Args, End) = 5 events
assert len(events) == 5
# Check ToolCallEndEvent
assert events[0].type == "TOOL_CALL_END"
assert events[0].tool_call_id == "call_456"
# Check function_approval_request CustomEvent
assert events[1].type == "CUSTOM"
assert events[1].name == "function_approval_request"
# Check confirm_changes tool call events
assert events[2].type == "TOOL_CALL_START"
assert events[2].tool_call_name == "confirm_changes"
assert events[3].type == "TOOL_CALL_ARGS"
# Verify confirm_changes includes function info for Dojo UI
import json
args = json.loads(events[3].delta)
assert args["function_name"] == "delete_file"
assert args["function_call_id"] == "call_456"
assert args["function_arguments"] == {"path": "/tmp/test.txt"}
assert args["steps"] == [
{
"description": "Execute delete_file",
"status": "enabled",
}
]
assert events[4].type == "TOOL_CALL_END"
async def test_multiple_approval_requests():
"""Test handling multiple approval requests in one update."""
# Set require_confirmation=False to simplify the test
bridge = AgentFrameworkEventBridge(
run_id="test_run",
thread_id="test_thread",
require_confirmation=False,
)
func_call_1 = Content.from_function_call(
call_id="call_1",
name="create_event",
arguments={"title": "Meeting"},
)
approval_1 = Content.from_function_approval_request(
id="approval_1",
function_call=func_call_1,
)
func_call_2 = Content.from_function_call(
call_id="call_2",
name="book_room",
arguments={"room": "Conference A"},
)
approval_2 = Content.from_function_approval_request(
id="approval_2",
function_call=func_call_2,
)
update = AgentResponseUpdate(contents=[approval_1, approval_2])
events = await bridge.from_agent_run_update(update)
# Should emit ToolCallEndEvent + CustomEvent for each approval (4 events total)
assert len(events) == 4
# Events should alternate: End, Custom, End, Custom
assert events[0].type == "TOOL_CALL_END"
assert events[0].tool_call_id == "call_1"
assert events[1].type == "CUSTOM"
assert events[1].name == "function_approval_request"
assert events[1].value["id"] == "approval_1"
assert events[2].type == "TOOL_CALL_END"
assert events[2].tool_call_id == "call_2"
assert events[3].type == "CUSTOM"
assert events[3].name == "function_approval_request"
assert events[3].value["id"] == "approval_2"
async def test_function_approval_request_sets_stop_flag():
"""Test that function approval request sets should_stop_after_confirm flag.
This ensures the orchestrator stops the run after emitting the approval request,
allowing the UI to send back an approval response.
"""
bridge = AgentFrameworkEventBridge(
run_id="test_run",
thread_id="test_thread",
)
assert bridge.should_stop_after_confirm is False
func_call = Content.from_function_call(
call_id="call_stop_test",
name="get_datetime",
arguments={},
)
approval_request = Content.from_function_approval_request(
id="approval_stop_test",
function_call=func_call,
)
update = AgentResponseUpdate(contents=[approval_request])
await bridge.from_agent_run_update(update)
assert bridge.should_stop_after_confirm is True
python/packages/ag-ui/tests/test_message_adapters.py
+96 -38
View File
@@ -644,49 +644,107 @@ def test_agent_framework_to_agui_function_result_multiple_text_contents():
assert agui_msg["content"] == '["First result", "Second result"]'
def test_agui_tool_approval_with_dataclass_modified_args():
"""Test that agui_messages_to_agent_framework handles dataclass in modified args.
# Additional tests for better coverage
This tests the fix for json.dumps() serialization errors at line 274
when modified_args contains non-serializable objects via make_json_safe.
"""
from dataclasses import dataclass
@dataclass
class ModifiedData:
field1: str
field2: int
def test_extract_text_from_contents_empty():
"""Test extracting text from empty contents."""
result = extract_text_from_contents([])
assert result == ""
# Create AG-UI format messages that simulate tool approval flow
# where modified args could contain a dataclass after parsing
# First, an assistant message with a tool call (string arguments)
assistant_msg = {
"id": "msg-1",
"role": "assistant",
"content": "",
"tool_calls": [
{
"id": "call-test",
"type": "function",
"function": {
"name": "update_state",
"arguments": '{"data": "original"}', # String args
},
}
],
}
def test_extract_text_from_contents_multiple():
"""Test extracting text from multiple text contents."""
contents = [
Content.from_text("Hello "),
Content.from_text("World"),
]
result = extract_text_from_contents(contents)
assert result == "Hello World"
# Then a user approval message (the approval path will merge modified args)
approval_msg = {
"id": "msg-2",
"role": "user",
"content": '{"approved": true}',
"toolCallId": "call-test",
}
# This should NOT raise TypeError
result = agui_messages_to_agent_framework([assistant_msg, approval_msg])
def test_extract_text_from_contents_non_text():
"""Test extracting text ignores non-text contents."""
contents = [
Content.from_text("Hello"),
Content.from_function_call(call_id="call_1", name="tool", arguments="{}"),
]
result = extract_text_from_contents(contents)
assert result == "Hello"
def test_agui_to_agent_framework_with_tool_calls():
"""Test converting AG-UI message with tool_calls."""
messages = [
{
"role": "assistant",
"content": "",
"tool_calls": [
{
"id": "call_123",
"type": "function",
"function": {"name": "get_weather", "arguments": '{"city": "NYC"}'},
}
],
}
]
result = agui_messages_to_agent_framework(messages)
assert len(result) == 1
assert len(result[0].contents) == 1
assert result[0].contents[0].type == "function_call"
assert result[0].contents[0].name == "get_weather"
def test_agui_to_agent_framework_tool_result():
"""Test converting AG-UI tool result message."""
messages = [
{
"role": "assistant",
"content": "",
"tool_calls": [
{
"id": "call_123",
"type": "function",
"function": {"name": "get_weather", "arguments": "{}"},
}
],
},
{
"role": "tool",
"content": "Sunny",
"toolCallId": "call_123",
},
]
result = agui_messages_to_agent_framework(messages)
# Should have processed both messages without error
assert len(result) == 2
# Second message should be tool result
tool_msg = result[1]
assert tool_msg.role == Role.TOOL
assert tool_msg.contents[0].type == "function_result"
assert tool_msg.contents[0].result == "Sunny"
def test_agui_messages_to_snapshot_format_empty():
"""Test converting empty messages to snapshot format."""
result = agui_messages_to_snapshot_format([])
assert result == []
def test_agui_messages_to_snapshot_format_basic():
"""Test converting messages to snapshot format."""
messages = [
{"role": "user", "content": "Hello", "id": "msg_1"},
{"role": "assistant", "content": "Hi there", "id": "msg_2"},
]
result = agui_messages_to_snapshot_format(messages)
assert len(result) == 2
assert result[0]["role"] == "user"
assert result[0]["content"] == "Hello"
assert result[1]["role"] == "assistant"
assert result[1]["content"] == "Hi there"
python/packages/ag-ui/tests/test_orchestrators.py
-307
View File
@@ -1,307 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for AG-UI orchestrators."""
from collections.abc import AsyncGenerator
from typing import Any
from unittest.mock import MagicMock
from ag_ui.core import BaseEvent, RunFinishedEvent
from agent_framework import (
AgentResponseUpdate,
AgentThread,
BaseChatClient,
ChatAgent,
ChatResponseUpdate,
Content,
FunctionInvocationConfiguration,
ai_function,
)
from agent_framework_ag_ui._agent import AgentConfig
from agent_framework_ag_ui._orchestrators import DefaultOrchestrator, ExecutionContext
@ai_function
def server_tool() -> str:
"""Server-executable tool."""
return "server"
def _create_mock_chat_agent(
tools: list[Any] | None = None,
response_format: Any = None,
capture_tools: list[Any] | None = None,
capture_messages: list[Any] | None = None,
) -> ChatAgent:
"""Create a ChatAgent with mocked chat client for testing.
Args:
tools: Tools to configure on the agent.
response_format: Response format to configure.
capture_tools: If provided, tools passed to run_stream will be appended here.
capture_messages: If provided, messages passed to run_stream will be appended here.
"""
mock_chat_client = MagicMock(spec=BaseChatClient)
mock_chat_client.function_invocation_configuration = FunctionInvocationConfiguration()
agent = ChatAgent(
chat_client=mock_chat_client,
tools=tools or [server_tool],
response_format=response_format,
)
# Create a mock run_stream that captures parameters and yields a simple response
async def mock_run_stream(
messages: list[Any],
*,
# thread: AgentThread,
# tools: list[Any] | None = None,
# **kwargs: Any,
# ) -> AsyncGenerator[AgentRunResponseUpdate, None]:
# self.seen_tools = tools
# yield AgentRunResponseUpdate(
# contents=[TextContent(text="ok")],
# role="assistant",
# response_id=thread.metadata.get("ag_ui_run_id"), # type: ignore[attr-defined] (metadata always created in orchestrator)
# raw_representation=ChatResponseUpdate(
# contents=[TextContent(text="ok")],
# conversation_id=thread.metadata.get("ag_ui_thread_id"), # type: ignore[attr-defined] (metadata always created in orchestrator)
# response_id=thread.metadata.get("ag_ui_run_id"), # type: ignore[attr-defined] (metadata always created in orchestrator)
# ),
# )
thread: AgentThread,
tools: list[Any] | None = None,
**kwargs: Any,
) -> AsyncGenerator[AgentResponseUpdate, None]:
if capture_tools is not None and tools is not None:
capture_tools.extend(tools)
if capture_messages is not None:
capture_messages.extend(messages)
yield AgentResponseUpdate(
contents=[Content.from_text(text="ok")],
role="assistant",
response_id=thread.metadata.get("ag_ui_run_id"), # type: ignore[attr-defined] (metadata always created in orchestrator)
raw_representation=ChatResponseUpdate(
contents=[Content.from_text(text="ok")],
conversation_id=thread.metadata.get("ag_ui_thread_id"), # type: ignore[attr-defined] (metadata always created in orchestrator)
response_id=thread.metadata.get("ag_ui_run_id"), # type: ignore[attr-defined] (metadata always created in orchestrator)
),
)
# Patch the run_stream method
agent.run_stream = mock_run_stream # type: ignore[method-assign]
return agent
async def test_default_orchestrator_merges_client_tools() -> None:
"""Client tool declarations are merged with server tools before running agent."""
captured_tools: list[Any] = []
agent = _create_mock_chat_agent(tools=[server_tool], capture_tools=captured_tools)
orchestrator = DefaultOrchestrator()
input_data = {
"messages": [
{
"role": "user",
"content": [{"type": "input_text", "text": "Hello"}],
}
],
"tools": [
{
"name": "get_weather",
"description": "Client weather lookup.",
"parameters": {
"type": "object",
"properties": {"location": {"type": "string"}},
"required": ["location"],
},
}
],
}
context = ExecutionContext(
input_data=input_data,
agent=agent,
config=AgentConfig(),
)
events = []
async for event in orchestrator.run(context):
events.append(event)
assert len(captured_tools) > 0
tool_names = [getattr(tool, "name", "?") for tool in captured_tools]
assert "server_tool" in tool_names
assert "get_weather" in tool_names
assert agent.chat_client.function_invocation_configuration.additional_tools
async def test_default_orchestrator_with_camel_case_ids() -> None:
"""Client tool is able to extract camelCase IDs."""
agent = _create_mock_chat_agent()
orchestrator = DefaultOrchestrator()
input_data = {
"runId": "test-camelcase-runid",
"threadId": "test-camelcase-threadid",
"messages": [
{
"role": "user",
"content": [{"type": "input_text", "text": "Hello"}],
}
],
"tools": [],
}
context = ExecutionContext(
input_data=input_data,
agent=agent,
config=AgentConfig(),
)
events = []
async for event in orchestrator.run(context):
events.append(event)
# assert the last event has the expected run_id and thread_id
assert isinstance(events[-1], RunFinishedEvent)
last_event = events[-1]
assert last_event.run_id == "test-camelcase-runid"
assert last_event.thread_id == "test-camelcase-threadid"
async def test_default_orchestrator_with_snake_case_ids() -> None:
"""Client tool is able to extract snake_case IDs."""
agent = _create_mock_chat_agent()
orchestrator = DefaultOrchestrator()
input_data = {
"run_id": "test-snakecase-runid",
"thread_id": "test-snakecase-threadid",
"messages": [
{
"role": "user",
"content": [{"type": "input_text", "text": "Hello"}],
}
],
"tools": [],
}
context = ExecutionContext(
input_data=input_data,
agent=agent,
config=AgentConfig(),
)
events: list[BaseEvent] = []
async for event in orchestrator.run(context):
events.append(event)
# assert the last event has the expected run_id and thread_id
assert isinstance(events[-1], RunFinishedEvent)
last_event = events[-1]
assert last_event.run_id == "test-snakecase-runid"
assert last_event.thread_id == "test-snakecase-threadid"
async def test_state_context_injected_when_tool_call_state_mismatch() -> None:
"""State context should be injected when current state differs from tool call args."""
captured_messages: list[Any] = []
agent = _create_mock_chat_agent(tools=[], capture_messages=captured_messages)
orchestrator = DefaultOrchestrator()
tool_recipe = {"title": "Salad", "special_preferences": []}
current_recipe = {"title": "Salad", "special_preferences": ["Vegetarian"]}
input_data = {
"state": {"recipe": current_recipe},
"messages": [
{"role": "system", "content": "Instructions"},
{
"role": "assistant",
"tool_calls": [
{
"id": "call_1",
"type": "function",
"function": {"name": "update_recipe", "arguments": {"recipe": tool_recipe}},
}
],
},
{"role": "user", "content": "What are the dietary preferences?"},
],
}
context = ExecutionContext(
input_data=input_data,
agent=agent,
config=AgentConfig(
state_schema={"recipe": {"type": "object"}},
predict_state_config={"recipe": {"tool": "update_recipe", "tool_argument": "recipe"}},
require_confirmation=False,
),
)
async for _event in orchestrator.run(context):
pass
assert len(captured_messages) > 0
state_messages = []
for msg in captured_messages:
role_value = msg.role.value if hasattr(msg.role, "value") else str(msg.role)
if role_value != "system":
continue
for content in msg.contents or []:
if content.type == "text" and content.text.startswith("Current state of the application:"):
state_messages.append(content.text)
assert state_messages
assert "Vegetarian" in state_messages[0]
async def test_state_context_not_injected_when_tool_call_matches_state() -> None:
"""State context should be skipped when tool call args match current state."""
captured_messages: list[Any] = []
agent = _create_mock_chat_agent(tools=[], capture_messages=captured_messages)
orchestrator = DefaultOrchestrator()
input_data = {
"messages": [
{"role": "system", "content": "Instructions"},
{
"role": "assistant",
"tool_calls": [
{
"id": "call_1",
"type": "function",
"function": {"name": "update_recipe", "arguments": {"recipe": {}}},
}
],
},
{"role": "user", "content": "What are the dietary preferences?"},
],
}
context = ExecutionContext(
input_data=input_data,
agent=agent,
config=AgentConfig(
state_schema={"recipe": {"type": "object"}},
predict_state_config={"recipe": {"tool": "update_recipe", "tool_argument": "recipe"}},
require_confirmation=False,
),
)
async for _event in orchestrator.run(context):
pass
assert len(captured_messages) > 0
state_messages = []
for msg in captured_messages:
role_value = msg.role.value if hasattr(msg.role, "value") else str(msg.role)
if role_value != "system":
continue
for content in msg.contents or []:
if content.type == "text" and content.text.startswith("Current state of the application:"):
state_messages.append(content.text)
assert not state_messages
python/packages/ag-ui/tests/test_orchestrators_coverage.py
-929
View File
@@ -1,929 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""Comprehensive tests for orchestrator coverage."""
import sys
from collections.abc import AsyncGenerator
from pathlib import Path
from types import SimpleNamespace
from typing import Any
from agent_framework import AgentResponseUpdate, ChatMessage, Content, ai_function
from pydantic import BaseModel
from agent_framework_ag_ui._agent import AgentConfig
from agent_framework_ag_ui._orchestrators import DefaultOrchestrator, HumanInTheLoopOrchestrator
sys.path.insert(0, str(Path(__file__).parent))
from utils_test_ag_ui import StubAgent, TestExecutionContext
@ai_function(approval_mode="always_require")
def approval_tool(param: str) -> str:
"""Tool requiring approval."""
return f"executed: {param}"
DEFAULT_OPTIONS: dict[str, Any] = {"tools": [approval_tool], "response_format": None}
async def test_human_in_the_loop_json_decode_error() -> None:
"""Test HumanInTheLoopOrchestrator handles invalid JSON in tool result."""
orchestrator = HumanInTheLoopOrchestrator()
input_data: dict[str, Any] = {
"messages": [
{
"role": "tool",
"content": [{"type": "text", "text": "not valid json {"}],
}
],
}
messages = [
ChatMessage(
role="tool",
contents=[Content.from_text(text="not valid json {")],
additional_properties={"is_tool_result": True},
)
]
agent = StubAgent(
default_options={"tools": [approval_tool], "response_format": None},
updates=[AgentResponseUpdate(contents=[Content.from_text(text="response")], role="assistant")],
)
context = TestExecutionContext(
input_data=input_data,
agent=agent,
config=AgentConfig(),
)
context.set_messages(messages, normalize=False)
assert orchestrator.can_handle(context)
events: list[Any] = []
async for event in orchestrator.run(context):
events.append(event)
# Should emit RunErrorEvent for invalid JSON
error_events: list[Any] = [e for e in events if e.type == "RUN_ERROR"]
assert len(error_events) == 1
assert "Invalid tool result format" in error_events[0].message
async def test_sanitize_tool_history_confirm_changes() -> None:
"""Test sanitize_tool_history logic for confirm_changes synthetic result."""
from agent_framework import ChatMessage
# Create messages that will trigger confirm_changes synthetic result injection
messages = [
ChatMessage(
role="assistant",
contents=[
Content.from_function_call(
name="confirm_changes",
call_id="call_confirm_123",
arguments='{"changes": "test"}',
)
],
),
ChatMessage(
role="user",
contents=[Content.from_text(text='{"accepted": true}')],
),
]
# The sanitize_tool_history function is internal to DefaultOrchestrator.run
# We'll test it indirectly by checking the orchestrator processes it correctly
orchestrator = DefaultOrchestrator()
# Use pre-constructed ChatMessage objects to bypass message adapter
input_data: dict[str, Any] = {"messages": []}
agent = StubAgent(
default_options=DEFAULT_OPTIONS,
)
context = TestExecutionContext(
input_data=input_data,
agent=agent,
config=AgentConfig(),
)
# Override the messages property to use our pre-constructed messages
context.set_messages(messages)
events: list[Any] = []
async for event in orchestrator.run(context):
events.append(event)
# Agent should receive synthetic tool result
assert len(agent.messages_received) > 0
tool_messages = [
msg
for msg in agent.messages_received
if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "tool"
]
assert len(tool_messages) == 1
assert str(tool_messages[0].contents[0].call_id) == "call_confirm_123"
assert tool_messages[0].contents[0].result == "Confirmed"
async def test_sanitize_tool_history_orphaned_tool_result() -> None:
"""Test sanitize_tool_history removes orphaned tool results."""
from agent_framework import ChatMessage
# Tool result without preceding assistant tool call
messages = [
ChatMessage(
role="tool",
contents=[Content.from_function_result(call_id="orphan_123", result="orphaned data")],
),
ChatMessage(
role="user",
contents=[Content.from_text(text="Hello")],
),
]
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {"messages": []}
agent = StubAgent(
default_options=DEFAULT_OPTIONS,
)
context = TestExecutionContext(
input_data=input_data,
agent=agent,
config=AgentConfig(),
)
context.set_messages(messages)
events: list[Any] = []
async for event in orchestrator.run(context):
events.append(event)
# Orphaned tool result should be filtered out
tool_messages = [
msg
for msg in agent.messages_received
if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "tool"
]
assert len(tool_messages) == 0
async def test_orphaned_tool_result_sanitization() -> None:
"""Test that orphaned tool results are filtered out."""
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {
"messages": [
{
"role": "tool",
"content": [{"type": "tool_result", "tool_call_id": "orphan_123", "content": "result"}],
},
{
"role": "user",
"content": [{"type": "text", "text": "Hello"}],
},
],
}
agent = StubAgent(
default_options=DEFAULT_OPTIONS,
)
context = TestExecutionContext(
input_data=input_data,
agent=agent,
config=AgentConfig(),
)
events: list[Any] = []
async for event in orchestrator.run(context):
events.append(event)
# Orphaned tool result should be filtered, only user message remains
tool_messages = [
msg
for msg in agent.messages_received
if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "tool"
]
assert len(tool_messages) == 0
async def test_deduplicate_messages_empty_tool_results() -> None:
"""Test deduplicate_messages prefers non-empty tool results."""
from agent_framework import ChatMessage
messages = [
ChatMessage(
role="assistant",
contents=[Content.from_function_call(name="test_tool", call_id="call_789", arguments="{}")],
),
ChatMessage(
role="tool",
contents=[Content.from_function_result(call_id="call_789", result="")],
),
ChatMessage(
role="tool",
contents=[Content.from_function_result(call_id="call_789", result="real data")],
),
]
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {"messages": []}
agent = StubAgent(
default_options=DEFAULT_OPTIONS,
)
context = TestExecutionContext(
input_data=input_data,
agent=agent,
config=AgentConfig(),
)
context.set_messages(messages)
events: list[Any] = []
async for event in orchestrator.run(context):
events.append(event)
# Should have only one tool result with actual data
tool_messages = [
msg
for msg in agent.messages_received
if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "tool"
]
assert len(tool_messages) == 1
assert tool_messages[0].contents[0].result == "real data"
async def test_deduplicate_messages_duplicate_assistant_tool_calls() -> None:
"""Test deduplicate_messages removes duplicate assistant tool call messages."""
from agent_framework import ChatMessage
messages = [
ChatMessage(
role="assistant",
contents=[Content.from_function_call(name="test_tool", call_id="call_abc", arguments="{}")],
),
ChatMessage(
role="assistant",
contents=[Content.from_function_call(name="test_tool", call_id="call_abc", arguments="{}")],
),
ChatMessage(
role="tool",
contents=[Content.from_function_result(call_id="call_abc", result="result")],
),
]
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {"messages": []}
agent = StubAgent(
default_options=DEFAULT_OPTIONS,
)
context = TestExecutionContext(
input_data=input_data,
agent=agent,
config=AgentConfig(),
)
context.set_messages(messages)
events: list[Any] = []
async for event in orchestrator.run(context):
events.append(event)
# Should have only one assistant message
assistant_messages = [
msg
for msg in agent.messages_received
if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "assistant"
]
assert len(assistant_messages) == 1
async def test_deduplicate_messages_duplicate_system_messages() -> None:
"""Test that deduplication logic is invoked for system messages."""
from agent_framework import ChatMessage
messages = [
ChatMessage(
role="system",
contents=[Content.from_text(text="You are a helpful assistant.")],
),
ChatMessage(
role="system",
contents=[Content.from_text(text="You are a helpful assistant.")],
),
ChatMessage(
role="user",
contents=[Content.from_text(text="Hello")],
),
]
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {"messages": []}
agent = StubAgent(
default_options=DEFAULT_OPTIONS,
)
context = TestExecutionContext(
input_data=input_data,
agent=agent,
config=AgentConfig(),
)
context.set_messages(messages)
events: list[Any] = []
async for event in orchestrator.run(context):
events.append(event)
# Deduplication uses hash() which may not deduplicate identical content
# This test verifies deduplication logic runs without errors
system_messages = [
msg
for msg in agent.messages_received
if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "system"
]
# At least one system message should be present
assert len(system_messages) >= 1
async def test_state_context_injection() -> None:
"""Test state context message injection for first request."""
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {
"messages": [
{
"role": "user",
"content": [{"type": "text", "text": "Hello"}],
}
],
"state": {"items": ["apple", "banana"]},
}
agent = StubAgent(
default_options=DEFAULT_OPTIONS,
)
context = TestExecutionContext(
input_data=input_data,
agent=agent,
config=AgentConfig(state_schema={"items": {"type": "array"}}),
)
events: list[Any] = []
async for event in orchestrator.run(context):
events.append(event)
# Should inject system message with current state
system_messages = [
msg
for msg in agent.messages_received
if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "system"
]
assert len(system_messages) == 1
assert "apple" in system_messages[0].contents[0].text
assert "banana" in system_messages[0].contents[0].text
async def test_state_context_injection_with_tool_calls_and_input_state() -> None:
"""Test state context is injected when state is provided, even with tool calls."""
from agent_framework import ChatMessage
messages = [
ChatMessage(
role="assistant",
contents=[Content.from_function_call(name="get_weather", call_id="call_xyz", arguments="{}")],
),
ChatMessage(
role="tool",
contents=[Content.from_function_result(call_id="call_xyz", result="sunny")],
),
ChatMessage(
role="user",
contents=[Content.from_text(text="Thanks")],
),
]
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {"messages": [], "state": {"weather": "sunny"}}
agent = StubAgent(
default_options=DEFAULT_OPTIONS,
)
context = TestExecutionContext(
input_data=input_data,
agent=agent,
config=AgentConfig(state_schema={"weather": {"type": "string"}}),
)
context.set_messages(messages)
events: list[Any] = []
async for event in orchestrator.run(context):
events.append(event)
# Should inject state context system message because input state is provided
system_messages = [
msg
for msg in agent.messages_received
if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "system"
]
assert len(system_messages) == 1
async def test_structured_output_processing() -> None:
"""Test structured output extraction and state update."""
class RecipeState(BaseModel):
ingredients: list[str]
message: str
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {
"messages": [
{
"role": "user",
"content": [{"type": "text", "text": "Add tomato"}],
}
],
}
# Agent with structured output
agent = StubAgent(
default_options=DEFAULT_OPTIONS,
updates=[
AgentResponseUpdate(
contents=[Content.from_text(text='{"ingredients": ["tomato"], "message": "Added tomato"}')],
role="assistant",
)
],
)
agent.default_options["response_format"] = RecipeState
context = TestExecutionContext(
input_data=input_data,
agent=agent,
config=AgentConfig(state_schema={"ingredients": {"type": "array"}}),
)
events: list[Any] = []
async for event in orchestrator.run(context):
events.append(event)
# Should emit StateSnapshotEvent with ingredients
state_events: list[Any] = [e for e in events if e.type == "STATE_SNAPSHOT"]
assert len(state_events) >= 1
# Should emit TextMessage with message field
text_content_events: list[Any] = [e for e in events if e.type == "TEXT_MESSAGE_CONTENT"]
assert len(text_content_events) >= 1
assert any("Added tomato" in e.delta for e in text_content_events)
async def test_duplicate_client_tools_filtered() -> None:
"""Test that client tools duplicating server tools are filtered out."""
@ai_function
def get_weather(location: str) -> str:
"""Get weather for location."""
return f"Weather in {location}"
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {
"messages": [
{
"role": "user",
"content": [{"type": "text", "text": "Hello"}],
}
],
"tools": [
{
"name": "get_weather",
"description": "Client weather tool.",
"parameters": {
"type": "object",
"properties": {"location": {"type": "string"}},
"required": ["location"],
},
}
],
}
agent = StubAgent(
default_options=DEFAULT_OPTIONS,
)
agent.default_options["tools"] = [get_weather]
context = TestExecutionContext(
input_data=input_data,
agent=agent,
config=AgentConfig(),
)
events: list[Any] = []
async for event in orchestrator.run(context):
events.append(event)
# tools parameter should not be passed since client tool duplicates server tool
assert agent.tools_received is None
async def test_unique_client_tools_merged() -> None:
"""Test that unique client tools are merged with server tools."""
@ai_function
def server_tool() -> str:
"""Server tool."""
return "server"
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {
"messages": [
{
"role": "user",
"content": [{"type": "text", "text": "Hello"}],
}
],
"tools": [
{
"name": "client_tool",
"description": "Unique client tool.",
"parameters": {
"type": "object",
"properties": {"param": {"type": "string"}},
"required": ["param"],
},
}
],
}
agent = StubAgent(
default_options=DEFAULT_OPTIONS,
)
agent.default_options["tools"] = [server_tool]
context = TestExecutionContext(
input_data=input_data,
agent=agent,
config=AgentConfig(),
)
events: list[Any] = []
async for event in orchestrator.run(context):
events.append(event)
# tools parameter should be passed with both server and client tools
assert agent.tools_received is not None
tool_names = [getattr(tool, "name", None) for tool in agent.tools_received]
assert "server_tool" in tool_names
assert "client_tool" in tool_names
async def test_empty_messages_handling() -> None:
"""Test orchestrator handles empty message list gracefully."""
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {"messages": []}
agent = StubAgent(
default_options=DEFAULT_OPTIONS,
)
context = TestExecutionContext(
input_data=input_data,
agent=agent,
config=AgentConfig(),
)
events: list[Any] = []
async for event in orchestrator.run(context):
events.append(event)
# Should emit run lifecycle events but not call agent
assert len(agent.messages_received) == 0
run_started = [e for e in events if e.type == "RUN_STARTED"]
run_finished = [e for e in events if e.type == "RUN_FINISHED"]
assert len(run_started) == 1
assert len(run_finished) == 1
async def test_all_messages_filtered_handling() -> None:
"""Test orchestrator handles case where all messages are filtered out."""
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {
"messages": [
{
"role": "tool",
"content": [{"type": "tool_result", "tool_call_id": "orphan", "content": "data"}],
}
]
}
agent = StubAgent(
default_options=DEFAULT_OPTIONS,
)
context = TestExecutionContext(
input_data=input_data,
agent=agent,
config=AgentConfig(),
)
events: list[Any] = []
async for event in orchestrator.run(context):
events.append(event)
# Should finish without calling agent
assert len(agent.messages_received) == 0
run_finished = [e for e in events if e.type == "RUN_FINISHED"]
assert len(run_finished) == 1
async def test_confirm_changes_with_invalid_json_fallback() -> None:
"""Test confirm_changes with invalid JSON falls back to normal processing."""
from agent_framework import ChatMessage
messages = [
ChatMessage(
role="assistant",
contents=[
Content.from_function_call(
name="confirm_changes",
call_id="call_confirm_invalid",
arguments='{"changes": "test"}',
)
],
),
ChatMessage(
role="user",
contents=[Content.from_text(text="invalid json {")],
),
]
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {"messages": []}
agent = StubAgent(
default_options=DEFAULT_OPTIONS,
)
context = TestExecutionContext(
input_data=input_data,
agent=agent,
config=AgentConfig(),
)
context.set_messages(messages)
events: list[Any] = []
async for event in orchestrator.run(context):
events.append(event)
# Invalid JSON should fall back - user message should be included
user_messages = [
msg
for msg in agent.messages_received
if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "user"
]
assert len(user_messages) == 1
async def test_confirm_changes_closes_active_message_before_finish() -> None:
"""Confirm-changes flow closes any active text message before run finishes."""
from ag_ui.core import TextMessageEndEvent, TextMessageStartEvent
updates = [
AgentResponseUpdate(
contents=[
Content.from_function_call(
name="write_document_local",
call_id="call_1",
arguments='{"document": "Draft"}',
)
]
),
AgentResponseUpdate(contents=[Content.from_function_result(call_id="call_1", result="Done")]),
]
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {"messages": [{"role": "user", "content": "Start"}]}
agent = StubAgent(
default_options=DEFAULT_OPTIONS,
updates=updates,
)
context = TestExecutionContext(
input_data=input_data,
agent=agent,
config=AgentConfig(
predict_state_config={"document": {"tool": "write_document_local", "tool_argument": "document"}},
require_confirmation=True,
),
)
events: list[Any] = []
async for event in orchestrator.run(context):
events.append(event)
start_events = [e for e in events if isinstance(e, TextMessageStartEvent)]
end_events = [e for e in events if isinstance(e, TextMessageEndEvent)]
assert len(start_events) == 1
assert len(end_events) == 1
assert end_events[0].message_id == start_events[0].message_id
end_index = events.index(end_events[0])
finished_index = events.index([e for e in events if e.type == "RUN_FINISHED"][0])
assert end_index < finished_index
async def test_tool_result_kept_when_call_id_matches() -> None:
"""Test tool result is kept when call_id matches pending tool calls."""
from agent_framework import ChatMessage
messages = [
ChatMessage(
role="assistant",
contents=[Content.from_function_call(name="get_data", call_id="call_match", arguments="{}")],
),
ChatMessage(
role="tool",
contents=[Content.from_function_result(call_id="call_match", result="data")],
),
]
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {"messages": []}
agent = StubAgent(
default_options=DEFAULT_OPTIONS,
)
context = TestExecutionContext(
input_data=input_data,
agent=agent,
config=AgentConfig(),
)
context.set_messages(messages)
events: list[Any] = []
async for event in orchestrator.run(context):
events.append(event)
# Tool result should be kept
tool_messages = [
msg
for msg in agent.messages_received
if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "tool"
]
assert len(tool_messages) == 1
assert tool_messages[0].contents[0].result == "data"
async def test_agent_protocol_fallback_paths() -> None:
"""Test fallback paths for non-ChatAgent implementations."""
class CustomAgent:
"""Custom agent without ChatAgent type."""
def __init__(self) -> None:
self.default_options: dict[str, Any] = {"tools": [], "response_format": None}
self.chat_client = SimpleNamespace(function_invocation_configuration=SimpleNamespace())
self.messages_received: list[Any] = []
async def run_stream(
self,
messages: list[Any],
*,
thread: Any = None,
tools: list[Any] | None = None,
**kwargs: Any,
) -> AsyncGenerator[AgentResponseUpdate, None]:
self.messages_received = messages
yield AgentResponseUpdate(contents=[Content.from_text(text="response")], role="assistant")
from agent_framework import ChatMessage
messages = [ChatMessage(role="user", contents=[Content.from_text(text="Hello")])]
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {"messages": []}
agent = CustomAgent()
context = TestExecutionContext(
input_data=input_data,
agent=agent, # type: ignore
config=AgentConfig(),
)
context.set_messages(messages)
events: list[Any] = []
async for event in orchestrator.run(context):
events.append(event)
# Should work with custom agent implementation
assert len(agent.messages_received) > 0
async def test_initial_state_snapshot_with_array_schema() -> None:
"""Test state initialization with array type schema."""
from agent_framework import ChatMessage
messages = [ChatMessage(role="user", contents=[Content.from_text(text="Hello")])]
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {"messages": [], "state": {}}
agent = StubAgent(
default_options=DEFAULT_OPTIONS,
)
context = TestExecutionContext(
input_data=input_data,
agent=agent,
config=AgentConfig(state_schema={"items": {"type": "array"}}),
)
context.set_messages(messages)
events: list[Any] = []
async for event in orchestrator.run(context):
events.append(event)
# Should emit state snapshot with empty array for items
state_events: list[Any] = [e for e in events if e.type == "STATE_SNAPSHOT"]
assert len(state_events) >= 1
async def test_response_format_skip_text_content() -> None:
"""Test that response_format causes skip_text_content to be set."""
class OutputModel(BaseModel):
result: str
from agent_framework import ChatMessage
messages = [ChatMessage(role="user", contents=[Content.from_text(text="Hello")])]
orchestrator = DefaultOrchestrator()
input_data: dict[str, Any] = {"messages": []}
agent = StubAgent(
default_options=DEFAULT_OPTIONS,
)
agent.default_options["response_format"] = OutputModel
context = TestExecutionContext(
input_data=input_data,
agent=agent,
config=AgentConfig(),
)
context.set_messages(messages)
events: list[Any] = []
async for event in orchestrator.run(context):
events.append(event)
# Test passes if no errors occur - verifies response_format code path
assert len(events) > 0
async def test_human_in_the_loop_handles_none_additional_properties() -> None:
"""Test that HumanInTheLoopOrchestrator handles None additional_properties gracefully.
This test ensures the null safety fix for msg.additional_properties.get() works.
"""
orchestrator = HumanInTheLoopOrchestrator()
# Create a message with None additional_properties
msg = ChatMessage(
role="user",
contents=[Content.from_text(text="Hello")],
)
# Explicitly set additional_properties to None
msg.additional_properties = None # type: ignore[assignment]
agent = StubAgent() # Use default StubAgent
config = AgentConfig()
context = TestExecutionContext(
input_data={"messages": [{"role": "user", "content": "Hello"}]},
agent=agent,
config=config,
)
context.set_messages([msg])
# can_handle should return False (not crash) when additional_properties is None
result = orchestrator.can_handle(context)
assert result is False
async def test_default_orchestrator_handles_none_default_options() -> None:
"""Test that DefaultOrchestrator handles None default_options gracefully.
This test ensures the null safety fix for context.agent.default_options.get() works.
"""
orchestrator = DefaultOrchestrator()
# Use StubAgent with default_options set to None
agent = StubAgent(default_options=None)
config = AgentConfig()
context = TestExecutionContext(
input_data={"messages": [{"role": "user", "content": "Hello"}]},
agent=agent,
config=config,
)
# This should NOT crash when accessing default_options.get()
events: list[Any] = []
async for event in orchestrator.run(context):
events.append(event)
# Just check a few events to verify it's working
if len(events) > 3:
break
# Test passes if no AttributeError occurred
assert True
python/packages/ag-ui/tests/test_predictive_state.py
+320
View File
@@ -0,0 +1,320 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for predictive state handling."""
from ag_ui.core import StateDeltaEvent
from agent_framework_ag_ui._orchestration._predictive_state import PredictiveStateHandler
class TestPredictiveStateHandlerInit:
"""Tests for PredictiveStateHandler initialization."""
def test_default_init(self):
"""Initializes with default values."""
handler = PredictiveStateHandler()
assert handler.predict_state_config == {}
assert handler.current_state == {}
assert handler.streaming_tool_args == ""
assert handler.last_emitted_state == {}
assert handler.state_delta_count == 0
assert handler.pending_state_updates == {}
def test_init_with_config(self):
"""Initializes with provided config."""
config = {"document": {"tool": "write_doc", "tool_argument": "content"}}
state = {"document": "initial"}
handler = PredictiveStateHandler(predict_state_config=config, current_state=state)
assert handler.predict_state_config == config
assert handler.current_state == state
class TestResetStreaming:
"""Tests for reset_streaming method."""
def test_resets_streaming_state(self):
"""Resets streaming-related state."""
handler = PredictiveStateHandler()
handler.streaming_tool_args = "some accumulated args"
handler.state_delta_count = 5
handler.reset_streaming()
assert handler.streaming_tool_args == ""
assert handler.state_delta_count == 0
class TestExtractStateValue:
"""Tests for extract_state_value method."""
def test_no_config(self):
"""Returns None when no config."""
handler = PredictiveStateHandler()
result = handler.extract_state_value("some_tool", {"arg": "value"})
assert result is None
def test_no_args(self):
"""Returns None when args is None."""
handler = PredictiveStateHandler(predict_state_config={"key": {"tool": "tool", "tool_argument": "arg"}})
result = handler.extract_state_value("tool", None)
assert result is None
def test_empty_args(self):
"""Returns None when args is empty string."""
handler = PredictiveStateHandler(predict_state_config={"key": {"tool": "tool", "tool_argument": "arg"}})
result = handler.extract_state_value("tool", "")
assert result is None
def test_tool_not_in_config(self):
"""Returns None when tool not in config."""
handler = PredictiveStateHandler(predict_state_config={"key": {"tool": "other_tool", "tool_argument": "arg"}})
result = handler.extract_state_value("some_tool", {"arg": "value"})
assert result is None
def test_extracts_specific_argument(self):
"""Extracts value from specific tool argument."""
handler = PredictiveStateHandler(
predict_state_config={"document": {"tool": "write_doc", "tool_argument": "content"}}
)
result = handler.extract_state_value("write_doc", {"content": "Hello world"})
assert result == ("document", "Hello world")
def test_extracts_with_wildcard(self):
"""Extracts entire args with * wildcard."""
handler = PredictiveStateHandler(predict_state_config={"data": {"tool": "update_data", "tool_argument": "*"}})
args = {"key1": "value1", "key2": "value2"}
result = handler.extract_state_value("update_data", args)
assert result == ("data", args)
def test_extracts_from_json_string(self):
"""Extracts value from JSON string args."""
handler = PredictiveStateHandler(
predict_state_config={"document": {"tool": "write_doc", "tool_argument": "content"}}
)
result = handler.extract_state_value("write_doc", '{"content": "Hello world"}')
assert result == ("document", "Hello world")
def test_argument_not_in_args(self):
"""Returns None when tool_argument not in args."""
handler = PredictiveStateHandler(
predict_state_config={"document": {"tool": "write_doc", "tool_argument": "content"}}
)
result = handler.extract_state_value("write_doc", {"other": "value"})
assert result is None
class TestIsPredictiveTool:
"""Tests for is_predictive_tool method."""
def test_none_tool_name(self):
"""Returns False for None tool name."""
handler = PredictiveStateHandler(predict_state_config={"key": {"tool": "some_tool", "tool_argument": "arg"}})
assert handler.is_predictive_tool(None) is False
def test_no_config(self):
"""Returns False when no config."""
handler = PredictiveStateHandler()
assert handler.is_predictive_tool("some_tool") is False
def test_tool_in_config(self):
"""Returns True when tool is in config."""
handler = PredictiveStateHandler(predict_state_config={"key": {"tool": "some_tool", "tool_argument": "arg"}})
assert handler.is_predictive_tool("some_tool") is True
def test_tool_not_in_config(self):
"""Returns False when tool not in config."""
handler = PredictiveStateHandler(predict_state_config={"key": {"tool": "other_tool", "tool_argument": "arg"}})
assert handler.is_predictive_tool("some_tool") is False
class TestEmitStreamingDeltas:
"""Tests for emit_streaming_deltas method."""
def test_no_tool_name(self):
"""Returns empty list for None tool name."""
handler = PredictiveStateHandler(predict_state_config={"key": {"tool": "tool", "tool_argument": "arg"}})
result = handler.emit_streaming_deltas(None, '{"arg": "value"}')
assert result == []
def test_no_config(self):
"""Returns empty list when no config."""
handler = PredictiveStateHandler()
result = handler.emit_streaming_deltas("some_tool", '{"arg": "value"}')
assert result == []
def test_accumulates_args(self):
"""Accumulates argument chunks."""
handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}})
handler.emit_streaming_deltas("write", '{"text')
handler.emit_streaming_deltas("write", '": "hello')
assert handler.streaming_tool_args == '{"text": "hello'
def test_emits_delta_on_complete_json(self):
"""Emits delta when JSON is complete."""
handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}})
events = handler.emit_streaming_deltas("write", '{"text": "hello"}')
assert len(events) == 1
assert isinstance(events[0], StateDeltaEvent)
assert events[0].delta[0]["path"] == "/doc"
assert events[0].delta[0]["value"] == "hello"
assert events[0].delta[0]["op"] == "replace"
def test_emits_delta_on_partial_json(self):
"""Emits delta from partial JSON using regex."""
handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}})
# First chunk - partial
events = handler.emit_streaming_deltas("write", '{"text": "hel')
assert len(events) == 1
assert events[0].delta[0]["value"] == "hel"
def test_does_not_emit_duplicate_deltas(self):
"""Does not emit delta when value unchanged."""
handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}})
# First emission
events1 = handler.emit_streaming_deltas("write", '{"text": "hello"}')
assert len(events1) == 1
# Reset and emit same value again
handler.streaming_tool_args = ""
events2 = handler.emit_streaming_deltas("write", '{"text": "hello"}')
assert len(events2) == 0 # No duplicate
def test_emits_delta_on_value_change(self):
"""Emits delta when value changes."""
handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}})
# First value
events1 = handler.emit_streaming_deltas("write", '{"text": "hello"}')
assert len(events1) == 1
# Reset and new value
handler.streaming_tool_args = ""
events2 = handler.emit_streaming_deltas("write", '{"text": "world"}')
assert len(events2) == 1
assert events2[0].delta[0]["value"] == "world"
def test_tracks_pending_updates(self):
"""Tracks pending state updates."""
handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}})
handler.emit_streaming_deltas("write", '{"text": "hello"}')
assert handler.pending_state_updates == {"doc": "hello"}
class TestEmitPartialDeltas:
"""Tests for _emit_partial_deltas method."""
def test_unescapes_newlines(self):
"""Unescapes \\n in partial values."""
handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}})
handler.streaming_tool_args = '{"text": "line1\\nline2'
events = handler._emit_partial_deltas("write")
assert len(events) == 1
assert events[0].delta[0]["value"] == "line1\nline2"
def test_handles_escaped_quotes_partially(self):
"""Handles escaped quotes - regex stops at quote character."""
handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}})
# The regex pattern [^"]* stops at ANY quote, including escaped ones.
# This is expected behavior for partial streaming - the full JSON
# will be parsed correctly when complete.
handler.streaming_tool_args = '{"text": "say \\"hi'
events = handler._emit_partial_deltas("write")
assert len(events) == 1
# Captures "say \" then the backslash gets converted to empty string
# by the replace("\\\\", "\\") first, then replace('\\"', '"')
# but since there's no closing quote, we get "say \"
# After .replace("\\\\", "\\") -> "say \"
# After .replace('\\"', '"') -> "say " (but actually still "say \" due to order)
# The actual result: backslash is preserved since it's not a valid escape sequence
assert events[0].delta[0]["value"] == "say \\"
def test_unescapes_backslashes(self):
"""Unescapes \\\\ in partial values."""
handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}})
handler.streaming_tool_args = '{"text": "path\\\\to\\\\file'
events = handler._emit_partial_deltas("write")
assert len(events) == 1
assert events[0].delta[0]["value"] == "path\\to\\file"
class TestEmitCompleteDeltas:
"""Tests for _emit_complete_deltas method."""
def test_emits_for_matching_tool(self):
"""Emits delta for tool matching config."""
handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}})
events = handler._emit_complete_deltas("write", {"text": "content"})
assert len(events) == 1
assert events[0].delta[0]["value"] == "content"
def test_skips_non_matching_tool(self):
"""Skips tools not matching config."""
handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}})
events = handler._emit_complete_deltas("other_tool", {"text": "content"})
assert len(events) == 0
def test_handles_wildcard_argument(self):
"""Handles * wildcard for entire args."""
handler = PredictiveStateHandler(predict_state_config={"data": {"tool": "update", "tool_argument": "*"}})
args = {"key1": "val1", "key2": "val2"}
events = handler._emit_complete_deltas("update", args)
assert len(events) == 1
assert events[0].delta[0]["value"] == args
def test_skips_missing_argument(self):
"""Skips when tool_argument not in args."""
handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}})
events = handler._emit_complete_deltas("write", {"other": "value"})
assert len(events) == 0
class TestCreateDeltaEvent:
"""Tests for _create_delta_event method."""
def test_creates_event(self):
"""Creates StateDeltaEvent with correct structure."""
handler = PredictiveStateHandler()
event = handler._create_delta_event("key", "value")
assert isinstance(event, StateDeltaEvent)
assert event.delta[0]["op"] == "replace"
assert event.delta[0]["path"] == "/key"
assert event.delta[0]["value"] == "value"
def test_increments_count(self):
"""Increments state_delta_count."""
handler = PredictiveStateHandler()
handler._create_delta_event("key", "value")
assert handler.state_delta_count == 1
handler._create_delta_event("key", "value2")
assert handler.state_delta_count == 2
class TestApplyPendingUpdates:
"""Tests for apply_pending_updates method."""
def test_applies_pending_to_current(self):
"""Applies pending updates to current state."""
handler = PredictiveStateHandler(current_state={"existing": "value"})
handler.pending_state_updates = {"doc": "new content", "count": 5}
handler.apply_pending_updates()
assert handler.current_state == {"existing": "value", "doc": "new content", "count": 5}
def test_clears_pending_updates(self):
"""Clears pending updates after applying."""
handler = PredictiveStateHandler()
handler.pending_state_updates = {"doc": "content"}
handler.apply_pending_updates()
assert handler.pending_state_updates == {}
def test_overwrites_existing_keys(self):
"""Overwrites existing keys in current state."""
handler = PredictiveStateHandler(current_state={"doc": "old"})
handler.pending_state_updates = {"doc": "new"}
handler.apply_pending_updates()
assert handler.current_state["doc"] == "new"
python/packages/ag-ui/tests/test_run.py
+373
View File
@@ -0,0 +1,373 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for _run.py helper functions and FlowState."""
from agent_framework import ChatMessage, Content
from agent_framework_ag_ui._run import (
FlowState,
_build_safe_metadata,
_create_state_context_message,
_has_only_tool_calls,
_inject_state_context,
_should_suppress_intermediate_snapshot,
)
class TestBuildSafeMetadata:
"""Tests for _build_safe_metadata function."""
def test_none_metadata(self):
"""Returns empty dict for None."""
result = _build_safe_metadata(None)
assert result == {}
def test_empty_metadata(self):
"""Returns empty dict for empty dict."""
result = _build_safe_metadata({})
assert result == {}
def test_short_string_values(self):
"""Preserves short string values."""
metadata = {"key1": "short", "key2": "value"}
result = _build_safe_metadata(metadata)
assert result == metadata
def test_truncates_long_strings(self):
"""Truncates strings over 512 chars."""
long_value = "x" * 1000
metadata = {"key": long_value}
result = _build_safe_metadata(metadata)
assert len(result["key"]) == 512
def test_serializes_non_strings(self):
"""Serializes non-string values to JSON."""
metadata = {"count": 42, "items": [1, 2, 3]}
result = _build_safe_metadata(metadata)
assert result["count"] == "42"
assert result["items"] == "[1, 2, 3]"
def test_truncates_serialized_values(self):
"""Truncates serialized values over 512 chars."""
long_list = list(range(200))
metadata = {"data": long_list}
result = _build_safe_metadata(metadata)
assert len(result["data"]) == 512
class TestHasOnlyToolCalls:
"""Tests for _has_only_tool_calls function."""
def test_only_tool_calls(self):
"""Returns True when only function_call content."""
contents = [
Content.from_function_call(call_id="call_1", name="tool1", arguments="{}"),
]
assert _has_only_tool_calls(contents) is True
def test_tool_call_with_text(self):
"""Returns False when both tool call and text."""
contents = [
Content.from_text("Some text"),
Content.from_function_call(call_id="call_1", name="tool1", arguments="{}"),
]
assert _has_only_tool_calls(contents) is False
def test_only_text(self):
"""Returns False when only text."""
contents = [Content.from_text("Just text")]
assert _has_only_tool_calls(contents) is False
def test_empty_contents(self):
"""Returns False for empty contents."""
assert _has_only_tool_calls([]) is False
def test_tool_call_with_empty_text(self):
"""Returns True when text content has empty text."""
contents = [
Content.from_text(""),
Content.from_function_call(call_id="call_1", name="tool1", arguments="{}"),
]
assert _has_only_tool_calls(contents) is True
class TestShouldSuppressIntermediateSnapshot:
"""Tests for _should_suppress_intermediate_snapshot function."""
def test_no_tool_name(self):
"""Returns False when no tool name."""
result = _should_suppress_intermediate_snapshot(
None, {"key": {"tool": "write_doc", "tool_argument": "content"}}, False
)
assert result is False
def test_no_config(self):
"""Returns False when no config."""
result = _should_suppress_intermediate_snapshot("write_doc", None, False)
assert result is False
def test_confirmation_required(self):
"""Returns False when confirmation is required."""
config = {"key": {"tool": "write_doc", "tool_argument": "content"}}
result = _should_suppress_intermediate_snapshot("write_doc", config, True)
assert result is False
def test_tool_not_in_config(self):
"""Returns False when tool not in config."""
config = {"key": {"tool": "other_tool", "tool_argument": "content"}}
result = _should_suppress_intermediate_snapshot("write_doc", config, False)
assert result is False
def test_suppresses_predictive_tool(self):
"""Returns True for predictive tool without confirmation."""
config = {"document": {"tool": "write_doc", "tool_argument": "content"}}
result = _should_suppress_intermediate_snapshot("write_doc", config, False)
assert result is True
class TestFlowState:
"""Tests for FlowState dataclass."""
def test_default_values(self):
"""Tests default initialization."""
flow = FlowState()
assert flow.message_id is None
assert flow.tool_call_id is None
assert flow.tool_call_name is None
assert flow.waiting_for_approval is False
assert flow.current_state == {}
assert flow.accumulated_text == ""
assert flow.pending_tool_calls == []
assert flow.tool_calls_by_id == {}
assert flow.tool_results == []
assert flow.tool_calls_ended == set()
def test_get_tool_name(self):
"""Tests get_tool_name method."""
flow = FlowState()
flow.tool_calls_by_id = {"call_123": {"function": {"name": "get_weather", "arguments": "{}"}}}
assert flow.get_tool_name("call_123") == "get_weather"
assert flow.get_tool_name("nonexistent") is None
assert flow.get_tool_name(None) is None
def test_get_tool_name_empty_name(self):
"""Tests get_tool_name with empty name."""
flow = FlowState()
flow.tool_calls_by_id = {"call_123": {"function": {"name": "", "arguments": "{}"}}}
assert flow.get_tool_name("call_123") is None
def test_get_pending_without_end(self):
"""Tests get_pending_without_end method."""
flow = FlowState()
flow.pending_tool_calls = [
{"id": "call_1", "function": {"name": "tool1"}},
{"id": "call_2", "function": {"name": "tool2"}},
{"id": "call_3", "function": {"name": "tool3"}},
]
flow.tool_calls_ended = {"call_1", "call_3"}
result = flow.get_pending_without_end()
assert len(result) == 1
assert result[0]["id"] == "call_2"
class TestCreateStateContextMessage:
"""Tests for _create_state_context_message function."""
def test_no_state(self):
"""Returns None when no state."""
result = _create_state_context_message({}, {"properties": {}})
assert result is None
def test_no_schema(self):
"""Returns None when no schema."""
result = _create_state_context_message({"key": "value"}, {})
assert result is None
def test_creates_message(self):
"""Creates state context message."""
from agent_framework import Role
state = {"document": "Hello world"}
schema = {"properties": {"document": {"type": "string"}}}
result = _create_state_context_message(state, schema)
assert result is not None
assert result.role == Role.SYSTEM
assert len(result.contents) == 1
assert "Hello world" in result.contents[0].text
assert "Current state" in result.contents[0].text
class TestInjectStateContext:
"""Tests for _inject_state_context function."""
def test_no_state_message(self):
"""Returns original messages when no state context needed."""
messages = [ChatMessage(role="user", contents=[Content.from_text("Hello")])]
result = _inject_state_context(messages, {}, {})
assert result == messages
def test_empty_messages(self):
"""Returns empty list for empty messages."""
result = _inject_state_context([], {"key": "value"}, {"properties": {}})
assert result == []
def test_last_message_not_user(self):
"""Returns original messages when last message is not from user."""
messages = [
ChatMessage(role="user", contents=[Content.from_text("Hello")]),
ChatMessage(role="assistant", contents=[Content.from_text("Hi")]),
]
state = {"key": "value"}
schema = {"properties": {"key": {"type": "string"}}}
result = _inject_state_context(messages, state, schema)
assert result == messages
def test_injects_before_last_user_message(self):
"""Injects state context before last user message."""
from agent_framework import Role
messages = [
ChatMessage(role="system", contents=[Content.from_text("You are helpful")]),
ChatMessage(role="user", contents=[Content.from_text("Hello")]),
]
state = {"document": "content"}
schema = {"properties": {"document": {"type": "string"}}}
result = _inject_state_context(messages, state, schema)
assert len(result) == 3
# System message first
assert result[0].role == Role.SYSTEM
assert "helpful" in result[0].contents[0].text
# State context second
assert result[1].role == Role.SYSTEM
assert "Current state" in result[1].contents[0].text
# User message last
assert result[2].role == Role.USER
assert "Hello" in result[2].contents[0].text
# Additional tests for _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")
events = _emit_text(content, flow)
assert len(events) == 2 # TextMessageStartEvent + TextMessageContentEvent
assert flow.message_id is not None
assert flow.accumulated_text == "Hello world"
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("")
events = _emit_text(content, flow)
assert len(events) == 0
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")
events = _emit_text(content, flow)
assert len(events) == 1 # Only TextMessageContentEvent, no new start
assert flow.message_id == "existing-id"
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")
events = _emit_text(content, flow)
assert len(events) == 0
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")
events = _emit_text(content, flow, skip_text=True)
assert len(events) == 0
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",
name="get_weather",
arguments='{"city": "NYC"}',
)
events = _emit_tool_call(content, flow)
assert len(events) >= 1 # At least ToolCallStartEvent
assert flow.tool_call_id == "call_123"
assert flow.tool_call_name == "get_weather"
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="{}")
events = _emit_tool_call(content, flow)
assert len(events) >= 1
assert flow.tool_call_id is not None # ID should be generated
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
messages = [ChatMessage(role="user", contents=[Content.from_text("Hello")])]
result = _extract_approved_state_updates(messages, None)
assert result == {}
def test_extract_approved_state_updates_no_approval():
"""Test _extract_approved_state_updates returns empty when no approval content."""
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 = [ChatMessage(role="user", contents=[Content.from_text("Hello")])]
result = _extract_approved_state_updates(messages, handler)
assert result == {}
python/packages/ag-ui/tests/test_shared_state.py
-108
View File
@@ -1,108 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for shared state management."""
import sys
from pathlib import Path
from typing import Any
import pytest
from ag_ui.core import StateSnapshotEvent
from agent_framework import ChatAgent, ChatResponseUpdate, Content
from agent_framework_ag_ui._agent import AgentFrameworkAgent
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
sys.path.insert(0, str(Path(__file__).parent))
from utils_test_ag_ui import StreamingChatClientStub, stream_from_updates
@pytest.fixture
def mock_agent() -> ChatAgent:
"""Create a mock agent for testing."""
updates = [ChatResponseUpdate(contents=[Content.from_text(text="Hello!")])]
chat_client = StreamingChatClientStub(stream_from_updates(updates))
return ChatAgent(name="test_agent", instructions="Test agent", chat_client=chat_client)
def test_state_snapshot_event():
"""Test creating state snapshot events."""
bridge = AgentFrameworkEventBridge(run_id="test-run", thread_id="test-thread")
state = {
"recipe": {
"name": "Chocolate Chip Cookies",
"ingredients": ["flour", "sugar", "chocolate chips"],
"instructions": ["Mix ingredients", "Bake at 350°F"],
"servings": 24,
}
}
event = bridge.create_state_snapshot_event(state)
assert isinstance(event, StateSnapshotEvent)
assert event.snapshot == state
assert event.snapshot["recipe"]["name"] == "Chocolate Chip Cookies"
assert len(event.snapshot["recipe"]["ingredients"]) == 3
def test_state_delta_event():
"""Test creating state delta events using JSON Patch format."""
bridge = AgentFrameworkEventBridge(run_id="test-run", thread_id="test-thread")
# JSON Patch operations (RFC 6902)
delta = [
{"op": "add", "path": "/recipe/ingredients/-", "value": "vanilla extract"},
{"op": "replace", "path": "/recipe/servings", "value": 30},
]
event = bridge.create_state_delta_event(delta)
assert event.delta == delta
assert len(event.delta) == 2
assert event.delta[0]["op"] == "add"
assert event.delta[1]["op"] == "replace"
async def test_agent_with_initial_state(mock_agent: ChatAgent) -> None:
"""Test agent emits state snapshot when initial state provided."""
state_schema: dict[str, Any] = {"recipe": {"type": "object", "properties": {"name": {"type": "string"}}}}
agent = AgentFrameworkAgent(
agent=mock_agent,
state_schema=state_schema,
)
initial_state = {"recipe": {"name": "Test Recipe"}}
input_data: dict[str, Any] = {
"messages": [{"role": "user", "content": "Hello"}],
"state": initial_state,
}
events: list[Any] = []
async for event in agent.run_agent(input_data):
events.append(event)
# Should have RunStartedEvent, StateSnapshotEvent, RunFinishedEvent at minimum
snapshot_events = [e for e in events if isinstance(e, StateSnapshotEvent)]
assert len(snapshot_events) == 1
assert snapshot_events[0].snapshot == initial_state
async def test_agent_without_state_schema(mock_agent: ChatAgent) -> None:
"""Test agent doesn't emit state events without state schema."""
agent = AgentFrameworkAgent(agent=mock_agent)
input_data: dict[str, Any] = {
"messages": [{"role": "user", "content": "Hello"}],
"state": {"some": "state"},
}
events: list[Any] = []
async for event in agent.run_agent(input_data):
events.append(event)
# Should NOT have any StateSnapshotEvent
snapshot_events = [e for e in events if isinstance(e, StateSnapshotEvent)]
assert len(snapshot_events) == 0
python/packages/ag-ui/tests/test_state_manager.py
-105
View File
@@ -1,105 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from ag_ui.core import CustomEvent, EventType
from agent_framework import ChatMessage
from agent_framework_ag_ui._events import AgentFrameworkEventBridge
from agent_framework_ag_ui._orchestration._state_manager import StateManager
def test_state_manager_initializes_defaults_and_snapshot() -> None:
state_manager = StateManager(
state_schema={"items": {"type": "array"}, "metadata": {"type": "object"}},
predict_state_config=None,
require_confirmation=True,
)
current_state = state_manager.initialize({"metadata": {"a": 1}})
bridge = AgentFrameworkEventBridge(run_id="run", thread_id="thread", current_state=current_state)
snapshot_event = state_manager.initial_snapshot_event(bridge)
assert snapshot_event is not None
assert snapshot_event.snapshot["items"] == []
assert snapshot_event.snapshot["metadata"] == {"a": 1}
def test_state_manager_predict_state_event_shape() -> None:
state_manager = StateManager(
state_schema=None,
predict_state_config={"doc": {"tool": "write_document_local", "tool_argument": "document"}},
require_confirmation=True,
)
predict_event = state_manager.predict_state_event()
assert isinstance(predict_event, CustomEvent)
assert predict_event.type == EventType.CUSTOM
assert predict_event.name == "PredictState"
assert predict_event.value[0]["state_key"] == "doc"
def test_state_context_only_when_new_user_turn() -> None:
state_manager = StateManager(
state_schema={"items": {"type": "array"}},
predict_state_config=None,
require_confirmation=True,
)
state_manager.initialize({"items": [1]})
assert state_manager.state_context_message(is_new_user_turn=False, conversation_has_tool_calls=False) is None
message = state_manager.state_context_message(is_new_user_turn=True, conversation_has_tool_calls=False)
assert isinstance(message, ChatMessage)
assert message.contents[0].type == "text"
assert "Current state of the application" in message.contents[0].text
def test_state_manager_with_dataclass_in_state() -> None:
"""Test that state containing dataclasses can be serialized without crashing.
This test ensures the fix for JSON serialization errors when state
contains dataclass or other non-JSON-serializable objects.
"""
from dataclasses import dataclass
@dataclass
class UserData:
name: str
age: int
state_manager = StateManager(
state_schema={"user": {"type": "object"}},
predict_state_config=None,
require_confirmation=True,
)
# Initialize with a dataclass object in the state
state_manager.initialize({"user": UserData(name="Alice", age=30)})
# This should NOT raise TypeError when generating the context message
message = state_manager.state_context_message(is_new_user_turn=True, conversation_has_tool_calls=False)
assert message is not None
assert isinstance(message, ChatMessage)
# The dataclass should be serialized to JSON in the message
assert "Alice" in message.contents[0].text
assert "30" in message.contents[0].text
def test_state_manager_with_pydantic_in_state() -> None:
"""Test that state containing Pydantic models can be serialized without crashing."""
from pydantic import BaseModel
class UserModel(BaseModel):
email: str
active: bool
state_manager = StateManager(
state_schema={"user": {"type": "object"}},
predict_state_config=None,
require_confirmation=True,
)
# Initialize with a Pydantic model in the state
state_manager.initialize({"user": UserModel(email="test@example.com", active=True)})
# This should NOT raise TypeError
message = state_manager.state_context_message(is_new_user_turn=True, conversation_has_tool_calls=False)
assert message is not None
assert "test@example.com" in message.contents[0].text
python/packages/ag-ui/tests/test_tooling.py
+92
View File
@@ -129,3 +129,95 @@ def test_collect_server_tools_with_mcp_tools_via_public_property() -> None:
assert "regular_tool" in names
assert "mcp_function" in names
assert len(tools) == 2
# Additional tests for tooling coverage
def test_collect_server_tools_no_default_options() -> None:
"""collect_server_tools returns empty list when agent has no default_options."""
class MockAgent:
pass
agent = MockAgent()
tools = collect_server_tools(agent)
assert tools == []
def test_register_additional_client_tools_no_tools() -> None:
"""register_additional_client_tools does nothing with None tools."""
mock_chat_client = MagicMock()
agent = ChatAgent(chat_client=mock_chat_client)
# Should not raise
register_additional_client_tools(agent, None)
def test_register_additional_client_tools_no_chat_client() -> None:
"""register_additional_client_tools does nothing when agent has no chat_client."""
from agent_framework_ag_ui._orchestration._tooling import register_additional_client_tools
class MockAgent:
pass
agent = MockAgent()
tools = [DummyTool("x")]
# Should not raise
register_additional_client_tools(agent, tools)
def test_merge_tools_no_client_tools() -> None:
"""merge_tools returns None when no client tools."""
server = [DummyTool("a")]
result = merge_tools(server, None)
assert result is None
def test_merge_tools_all_duplicates() -> None:
"""merge_tools returns None when all client tools duplicate server tools."""
server = [DummyTool("a"), DummyTool("b")]
client = [DummyTool("a"), DummyTool("b")]
result = merge_tools(server, client)
assert result is None
def test_merge_tools_empty_server() -> None:
"""merge_tools works with empty server tools."""
server: list = []
client = [DummyTool("a"), DummyTool("b")]
result = merge_tools(server, client)
assert result is not None
assert len(result) == 2
def test_merge_tools_with_approval_tools_no_client() -> None:
"""merge_tools returns server tools when they have approval mode even without client tools."""
class ApprovalTool:
def __init__(self, name: str):
self.name = name
self.approval_mode = "always_require"
server = [ApprovalTool("write_doc")]
result = merge_tools(server, None)
assert result is not None
assert len(result) == 1
assert result[0].name == "write_doc"
def test_merge_tools_with_approval_tools_all_duplicates() -> None:
"""merge_tools returns server tools with approval mode even when client duplicates."""
class ApprovalTool:
def __init__(self, name: str):
self.name = name
self.approval_mode = "always_require"
server = [ApprovalTool("write_doc")]
client = [DummyTool("write_doc")] # Same name as server
result = merge_tools(server, client)
assert result is not None
assert len(result) == 1
assert result[0].approval_mode == "always_require"
python/packages/ag-ui/tests/test_types.py
+81 -1
View File
@@ -2,7 +2,7 @@
"""Tests for type definitions in _types.py."""
from agent_framework_ag_ui._types import AgentState, PredictStateConfig, RunMetadata
from agent_framework_ag_ui._types import AgentState, AGUIRequest, PredictStateConfig, RunMetadata
class TestPredictStateConfig:
@@ -143,3 +143,83 @@ class TestAgentState:
assert len(state["messages"]) == 2
assert "metadata" in state["messages"][0]
assert "tool_calls" in state["messages"][1]
class TestAGUIRequest:
"""Test AGUIRequest Pydantic model."""
def test_agui_request_minimal(self) -> None:
"""Test creating AGUIRequest with only required fields."""
request = AGUIRequest(messages=[{"role": "user", "content": "Hello"}])
assert len(request.messages) == 1
assert request.messages[0]["content"] == "Hello"
assert request.run_id is None
assert request.thread_id is None
assert request.state is None
assert request.tools is None
assert request.context is None
assert request.forwarded_props is None
assert request.parent_run_id is None
def test_agui_request_all_fields(self) -> None:
"""Test creating AGUIRequest with all fields populated."""
request = AGUIRequest(
messages=[{"role": "user", "content": "Hello"}],
run_id="run-123",
thread_id="thread-456",
state={"counter": 0},
tools=[{"name": "search", "description": "Search tool"}],
context=[{"type": "document", "content": "Some context"}],
forwarded_props={"custom_key": "custom_value"},
parent_run_id="parent-run-789",
)
assert request.run_id == "run-123"
assert request.thread_id == "thread-456"
assert request.state == {"counter": 0}
assert request.tools == [{"name": "search", "description": "Search tool"}]
assert request.context == [{"type": "document", "content": "Some context"}]
assert request.forwarded_props == {"custom_key": "custom_value"}
assert request.parent_run_id == "parent-run-789"
def test_agui_request_model_dump_excludes_none(self) -> None:
"""Test that model_dump(exclude_none=True) excludes None fields."""
request = AGUIRequest(
messages=[{"role": "user", "content": "test"}],
tools=[{"name": "my_tool"}],
context=[{"id": "ctx1"}],
)
dumped = request.model_dump(exclude_none=True)
assert "messages" in dumped
assert "tools" in dumped
assert "context" in dumped
assert "run_id" not in dumped
assert "thread_id" not in dumped
assert "state" not in dumped
assert "forwarded_props" not in dumped
assert "parent_run_id" not in dumped
def test_agui_request_model_dump_includes_all_set_fields(self) -> None:
"""Test that model_dump preserves all explicitly set fields.
This is critical for the fix - ensuring tools, context, forwarded_props,
and parent_run_id are not stripped during request validation.
"""
request = AGUIRequest(
messages=[{"role": "user", "content": "test"}],
tools=[{"name": "client_tool", "parameters": {"type": "object"}}],
context=[{"type": "snippet", "content": "code here"}],
forwarded_props={"auth_token": "secret", "user_id": "user-1"},
parent_run_id="parent-456",
)
dumped = request.model_dump(exclude_none=True)
# Verify all fields are preserved (the main bug fix)
assert dumped["tools"] == [{"name": "client_tool", "parameters": {"type": "object"}}]
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"
python/packages/ag-ui/tests/test_utils.py
+170
View File
@@ -356,3 +356,173 @@ def test_convert_tools_to_agui_format_with_multiple_tools():
assert len(result) == 2
assert result[0]["name"] == "tool1"
assert result[1]["name"] == "tool2"
# Additional tests for utils coverage
def test_safe_json_parse_with_dict():
"""Test safe_json_parse with dict input."""
from agent_framework_ag_ui._utils import safe_json_parse
input_dict = {"key": "value"}
result = safe_json_parse(input_dict)
assert result == input_dict
def test_safe_json_parse_with_json_string():
"""Test safe_json_parse with JSON string."""
from agent_framework_ag_ui._utils import safe_json_parse
result = safe_json_parse('{"key": "value"}')
assert result == {"key": "value"}
def test_safe_json_parse_with_invalid_json():
"""Test safe_json_parse with invalid JSON."""
from agent_framework_ag_ui._utils import safe_json_parse
result = safe_json_parse("not json")
assert result is None
def test_safe_json_parse_with_non_dict_json():
"""Test safe_json_parse with JSON that parses to non-dict."""
from agent_framework_ag_ui._utils import safe_json_parse
result = safe_json_parse("[1, 2, 3]")
assert result is None
def test_safe_json_parse_with_none():
"""Test safe_json_parse with None input."""
from agent_framework_ag_ui._utils import safe_json_parse
result = safe_json_parse(None)
assert result is None
def test_get_role_value_with_enum():
"""Test get_role_value with enum role."""
from agent_framework import ChatMessage, Content, Role
from agent_framework_ag_ui._utils import get_role_value
message = ChatMessage(role=Role.USER, contents=[Content.from_text("test")])
result = get_role_value(message)
assert result == "user"
def test_get_role_value_with_string():
"""Test get_role_value with string role."""
from agent_framework_ag_ui._utils import get_role_value
class MockMessage:
role = "assistant"
result = get_role_value(MockMessage())
assert result == "assistant"
def test_get_role_value_with_none():
"""Test get_role_value with no role."""
from agent_framework_ag_ui._utils import get_role_value
class MockMessage:
pass
result = get_role_value(MockMessage())
assert result == ""
def test_normalize_agui_role_developer():
"""Test normalize_agui_role maps developer to system."""
from agent_framework_ag_ui._utils import normalize_agui_role
assert normalize_agui_role("developer") == "system"
def test_normalize_agui_role_valid():
"""Test normalize_agui_role with valid roles."""
from agent_framework_ag_ui._utils import normalize_agui_role
assert normalize_agui_role("user") == "user"
assert normalize_agui_role("assistant") == "assistant"
assert normalize_agui_role("system") == "system"
assert normalize_agui_role("tool") == "tool"
def test_normalize_agui_role_invalid():
"""Test normalize_agui_role with invalid role defaults to user."""
from agent_framework_ag_ui._utils import normalize_agui_role
assert normalize_agui_role("invalid") == "user"
assert normalize_agui_role(123) == "user"
def test_extract_state_from_tool_args():
"""Test extract_state_from_tool_args."""
from agent_framework_ag_ui._utils import extract_state_from_tool_args
# Specific key
assert extract_state_from_tool_args({"key": "value"}, "key") == "value"
# Wildcard
args = {"a": 1, "b": 2}
assert extract_state_from_tool_args(args, "*") == args
# Missing key
assert extract_state_from_tool_args({"other": "value"}, "key") is None
# None args
assert extract_state_from_tool_args(None, "key") is None
def test_convert_agui_tools_to_agent_framework():
"""Test convert_agui_tools_to_agent_framework."""
from agent_framework_ag_ui._utils import convert_agui_tools_to_agent_framework
agui_tools = [
{
"name": "test_tool",
"description": "A test tool",
"parameters": {"type": "object", "properties": {"arg": {"type": "string"}}},
}
]
result = convert_agui_tools_to_agent_framework(agui_tools)
assert result is not None
assert len(result) == 1
assert result[0].name == "test_tool"
assert result[0].description == "A test tool"
assert result[0].declaration_only is True
def test_convert_agui_tools_to_agent_framework_none():
"""Test convert_agui_tools_to_agent_framework with None."""
from agent_framework_ag_ui._utils import convert_agui_tools_to_agent_framework
result = convert_agui_tools_to_agent_framework(None)
assert result is None
def test_convert_agui_tools_to_agent_framework_empty():
"""Test convert_agui_tools_to_agent_framework with empty list."""
from agent_framework_ag_ui._utils import convert_agui_tools_to_agent_framework
result = convert_agui_tools_to_agent_framework([])
assert result is None
def test_make_json_safe_unconvertible():
"""Test make_json_safe with object that has no standard conversion."""
class NoConversion:
__slots__ = () # No __dict__
from agent_framework_ag_ui._utils import make_json_safe
result = make_json_safe(NoConversion())
# Falls back to str()
assert isinstance(result, str)
python/packages/ag-ui/tests/utils_test_ag_ui.py
-14
View File
@@ -20,9 +20,6 @@ from agent_framework import (
)
from agent_framework._clients import TOptions_co
from agent_framework_ag_ui._message_adapters import _deduplicate_messages, _sanitize_tool_history
from agent_framework_ag_ui._orchestrators import ExecutionContext
if sys.version_info >= (3, 12):
from typing import override # type: ignore # pragma: no cover
else:
@@ -125,14 +122,3 @@ class StubAgent(AgentProtocol):
def get_new_thread(self, **kwargs: Any) -> AgentThread:
return AgentThread()
class TestExecutionContext(ExecutionContext):
"""ExecutionContext helper that allows setting messages for tests."""
def set_messages(self, messages: list[ChatMessage], *, normalize: bool = True) -> None:
if normalize:
self._messages = _deduplicate_messages(_sanitize_tool_history(messages))
else:
self._messages = messages
self._snapshot_messages = None