Python: Refactor ag-ui to clean up some patterns (#2363)

* Refactor ag-ui to clean up some patterns

* Mypy fixes

* Fix imports, typing, tests, logging.

* Fix test import error

* Fix imports again

* Fix thread handling

python/packages/ag-ui/tests/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Microsoft. All rights reserved.

python/packages/ag-ui/tests/test_ag_ui_client.py

@@ -1,10 +1,61 @@
+# Copyright (c) Microsoft. All rights reserved.
+
 """Tests for AGUIChatClient."""
 
 import json
+from collections.abc import AsyncGenerator, AsyncIterable, MutableSequence
+from typing import Any
 
-from agent_framework import ChatMessage, ChatOptions, FunctionCallContent, Role, ai_function
+from agent_framework import (
+    ChatMessage,
+    ChatOptions,
+    ChatResponseUpdate,
+    FunctionCallContent,
+    Role,
+    TextContent,
+    ai_function,
+)
+from agent_framework._types import ChatResponse
+from pytest import MonkeyPatch
 
 from agent_framework_ag_ui._client import AGUIChatClient, ServerFunctionCallContent
+from agent_framework_ag_ui._http_service import AGUIHttpService
+
+
+class TestableAGUIChatClient(AGUIChatClient):
+    """Testable wrapper exposing protected helpers."""
+
+    @property
+    def http_service(self) -> AGUIHttpService:
+        """Expose http service for monkeypatching."""
+        return self._http_service
+
+    def extract_state_from_messages(
+        self, messages: list[ChatMessage]
+    ) -> tuple[list[ChatMessage], dict[str, Any] | None]:
+        """Expose state extraction helper."""
+        return self._extract_state_from_messages(messages)
+
+    def convert_messages_to_agui_format(self, messages: list[ChatMessage]) -> list[dict[str, Any]]:
+        """Expose message conversion helper."""
+        return self._convert_messages_to_agui_format(messages)
+
+    def get_thread_id(self, chat_options: ChatOptions) -> str:
+        """Expose thread id helper."""
+        return self._get_thread_id(chat_options)
+
+    async def inner_get_streaming_response(
+        self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions
+    ) -> AsyncIterable[ChatResponseUpdate]:
+        """Proxy to protected streaming call."""
+        async for update in self._inner_get_streaming_response(messages=messages, chat_options=chat_options):
+            yield update
+
+    async def inner_get_response(
+        self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions
+    ) -> ChatResponse:
+        """Proxy to protected response call."""
+        return await self._inner_get_response(messages=messages, chat_options=chat_options)
 
 
 class TestAGUIChatClient:
@@ -12,25 +63,25 @@ class TestAGUIChatClient:
 
     async def test_client_initialization(self) -> None:
         """Test client initialization."""
-        client = AGUIChatClient(endpoint="http://localhost:8888/")
+        client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
 
-        assert client._http_service is not None
-        assert client._http_service.endpoint.startswith("http://localhost:8888")
+        assert client.http_service is not None
+        assert client.http_service.endpoint.startswith("http://localhost:8888")
 
     async def test_client_context_manager(self) -> None:
         """Test client as async context manager."""
-        async with AGUIChatClient(endpoint="http://localhost:8888/") as client:
+        async with TestableAGUIChatClient(endpoint="http://localhost:8888/") as client:
             assert client is not None
 
     async def test_extract_state_from_messages_no_state(self) -> None:
         """Test state extraction when no state is present."""
-        client = AGUIChatClient(endpoint="http://localhost:8888/")
+        client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
         messages = [
             ChatMessage(role="user", text="Hello"),
             ChatMessage(role="assistant", text="Hi there"),
         ]
 
-        result_messages, state = client._extract_state_from_messages(messages)
+        result_messages, state = client.extract_state_from_messages(messages)
 
         assert result_messages == messages
         assert state is None
@@ -39,7 +90,7 @@ class TestAGUIChatClient:
         """Test state extraction from last message."""
         import base64
 
-        client = AGUIChatClient(endpoint="http://localhost:8888/")
+        client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
 
         state_data = {"key": "value", "count": 42}
         state_json = json.dumps(state_data)
@@ -55,7 +106,7 @@ class TestAGUIChatClient:
             ),
         ]
 
-        result_messages, state = client._extract_state_from_messages(messages)
+        result_messages, state = client.extract_state_from_messages(messages)
 
         assert len(result_messages) == 1
         assert result_messages[0].text == "Hello"
@@ -65,7 +116,7 @@ class TestAGUIChatClient:
         """Test state extraction with invalid JSON."""
         import base64
 
-        client = AGUIChatClient(endpoint="http://localhost:8888/")
+        client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
 
         invalid_json = "not valid json"
         state_b64 = base64.b64encode(invalid_json.encode("utf-8")).decode("utf-8")
@@ -79,20 +130,20 @@ class TestAGUIChatClient:
             ),
         ]
 
-        result_messages, state = client._extract_state_from_messages(messages)
+        result_messages, state = client.extract_state_from_messages(messages)
 
         assert result_messages == messages
         assert state is None
 
     async def test_convert_messages_to_agui_format(self) -> None:
         """Test message conversion to AG-UI format."""
-        client = AGUIChatClient(endpoint="http://localhost:8888/")
+        client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
         messages = [
             ChatMessage(role=Role.USER, text="What is the weather?"),
             ChatMessage(role=Role.ASSISTANT, text="Let me check.", message_id="msg_123"),
         ]
 
-        agui_messages = client._convert_messages_to_agui_format(messages)
+        agui_messages = client.convert_messages_to_agui_format(messages)
 
         assert len(agui_messages) == 2
         assert agui_messages[0]["role"] == "user"
@@ -103,24 +154,24 @@ class TestAGUIChatClient:
 
     async def test_get_thread_id_from_metadata(self) -> None:
         """Test thread ID extraction from metadata."""
-        client = AGUIChatClient(endpoint="http://localhost:8888/")
+        client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
         chat_options = ChatOptions(metadata={"thread_id": "existing_thread_123"})
 
-        thread_id = client._get_thread_id(chat_options)
+        thread_id = client.get_thread_id(chat_options)
 
         assert thread_id == "existing_thread_123"
 
     async def test_get_thread_id_generation(self) -> None:
         """Test automatic thread ID generation."""
-        client = AGUIChatClient(endpoint="http://localhost:8888/")
+        client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
         chat_options = ChatOptions()
 
-        thread_id = client._get_thread_id(chat_options)
+        thread_id = client.get_thread_id(chat_options)
 
         assert thread_id.startswith("thread_")
         assert len(thread_id) > 7
 
-    async def test_get_streaming_response(self, monkeypatch) -> None:
+    async def test_get_streaming_response(self, monkeypatch: MonkeyPatch) -> None:
         """Test streaming response method."""
         mock_events = [
             {"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
@@ -129,26 +180,32 @@ class TestAGUIChatClient:
             {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
         ]
 
-        async def mock_post_run(*args, **kwargs):
+        async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
             for event in mock_events:
                 yield event
 
-        client = AGUIChatClient(endpoint="http://localhost:8888/")
-        monkeypatch.setattr(client._http_service, "post_run", mock_post_run)
+        client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
+        monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
 
         messages = [ChatMessage(role="user", text="Test message")]
         chat_options = ChatOptions()
 
-        updates = []
-        async for update in client._inner_get_streaming_response(messages=messages, chat_options=chat_options):
+        updates: list[ChatResponseUpdate] = []
+        async for update in client.inner_get_streaming_response(messages=messages, chat_options=chat_options):
             updates.append(update)
 
         assert len(updates) == 4
+        assert updates[0].additional_properties is not None
         assert updates[0].additional_properties["thread_id"] == "thread_1"
-        assert updates[1].contents[0].text == "Hello"
-        assert updates[2].contents[0].text == " world"
 
-    async def test_get_response_non_streaming(self, monkeypatch) -> None:
+        first_content = updates[1].contents[0]
+        second_content = updates[2].contents[0]
+        assert isinstance(first_content, TextContent)
+        assert isinstance(second_content, TextContent)
+        assert first_content.text == "Hello"
+        assert second_content.text == " world"
+
+    async def test_get_response_non_streaming(self, monkeypatch: MonkeyPatch) -> None:
         """Test non-streaming response method."""
         mock_events = [
             {"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
@@ -156,23 +213,23 @@ class TestAGUIChatClient:
             {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
         ]
 
-        async def mock_post_run(*args, **kwargs):
+        async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
             for event in mock_events:
                 yield event
 
-        client = AGUIChatClient(endpoint="http://localhost:8888/")
-        monkeypatch.setattr(client._http_service, "post_run", mock_post_run)
+        client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
+        monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
 
         messages = [ChatMessage(role="user", text="Test message")]
         chat_options = ChatOptions()
 
-        response = await client._inner_get_response(messages=messages, chat_options=chat_options)
+        response = await client.inner_get_response(messages=messages, chat_options=chat_options)
 
         assert response is not None
         assert len(response.messages) > 0
         assert "Complete response" in response.text
 
-    async def test_tool_handling(self, monkeypatch) -> None:
+    async def test_tool_handling(self, monkeypatch: MonkeyPatch) -> None:
         """Test that client tool metadata is sent to server.
 
         Client tool metadata (name, description, schema) is sent to server for planning.
@@ -191,28 +248,29 @@ class TestAGUIChatClient:
             {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
         ]
 
-        async def mock_post_run(*args, **kwargs):
+        async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
             # Client tool metadata should be sent to server
-            tools = kwargs.get("tools")
+            tools: list[dict[str, Any]] | None = kwargs.get("tools")
             assert tools is not None
             assert len(tools) == 1
-            assert tools[0]["name"] == "test_tool"
-            assert tools[0]["description"] == "Test tool."
-            assert "parameters" in tools[0]
+            tool_entry = tools[0]
+            assert tool_entry["name"] == "test_tool"
+            assert tool_entry["description"] == "Test tool."
+            assert "parameters" in tool_entry
             for event in mock_events:
                 yield event
 
-        client = AGUIChatClient(endpoint="http://localhost:8888/")
-        monkeypatch.setattr(client._http_service, "post_run", mock_post_run)
+        client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
+        monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
 
         messages = [ChatMessage(role="user", text="Test with tools")]
         chat_options = ChatOptions(tools=[test_tool])
 
-        response = await client._inner_get_response(messages=messages, chat_options=chat_options)
+        response = await client.inner_get_response(messages=messages, chat_options=chat_options)
 
         assert response is not None
 
-    async def test_server_tool_calls_unwrapped_after_invocation(self, monkeypatch) -> None:
+    async def test_server_tool_calls_unwrapped_after_invocation(self, monkeypatch: MonkeyPatch) -> None:
         """Ensure server-side tool calls are exposed as FunctionCallContent after processing."""
 
         mock_events = [
@@ -222,17 +280,17 @@ class TestAGUIChatClient:
             {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
         ]
 
-        async def mock_post_run(*args, **kwargs):
+        async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
             for event in mock_events:
                 yield event
 
-        client = AGUIChatClient(endpoint="http://localhost:8888/")
-        monkeypatch.setattr(client._http_service, "post_run", mock_post_run)
+        client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
+        monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
 
         messages = [ChatMessage(role="user", text="Test server tool execution")]
         chat_options = ChatOptions()
 
-        updates = []
+        updates: list[ChatResponseUpdate] = []
         async for update in client.get_streaming_response(messages, chat_options=chat_options):
             updates.append(update)
 
@@ -245,7 +303,7 @@ class TestAGUIChatClient:
             isinstance(content, ServerFunctionCallContent) for update in updates for content in update.contents
         )
 
-    async def test_server_tool_calls_not_executed_locally(self, monkeypatch) -> None:
+    async def test_server_tool_calls_not_executed_locally(self, monkeypatch: MonkeyPatch) -> None:
         """Server tools should not trigger local function invocation even when client tools exist."""
 
         @ai_function
@@ -260,18 +318,18 @@ class TestAGUIChatClient:
             {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
         ]
 
-        async def mock_post_run(*args, **kwargs):
+        async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
             for event in mock_events:
                 yield event
 
-        async def fake_auto_invoke(*args, **kwargs):
+        async def fake_auto_invoke(*args: object, **kwargs: Any) -> None:
             function_call = kwargs.get("function_call_content") or args[0]
             raise AssertionError(f"Unexpected local execution of server tool: {getattr(function_call, 'name', '?')}")
 
         monkeypatch.setattr("agent_framework._tools._auto_invoke_function", fake_auto_invoke)
 
-        client = AGUIChatClient(endpoint="http://localhost:8888/")
-        monkeypatch.setattr(client._http_service, "post_run", mock_post_run)
+        client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
+        monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
 
         messages = [ChatMessage(role="user", text="Test server tool execution")]
         chat_options = ChatOptions(tool_choice="auto", tools=[client_tool])
@@ -279,7 +337,7 @@ class TestAGUIChatClient:
         async for _ in client.get_streaming_response(messages, chat_options=chat_options):
             pass
 
-    async def test_state_transmission(self, monkeypatch) -> None:
+    async def test_state_transmission(self, monkeypatch: MonkeyPatch) -> None:
         """Test state is properly transmitted to server."""
         import base64
 
@@ -302,16 +360,16 @@ class TestAGUIChatClient:
             {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
         ]
 
-        async def mock_post_run(*args, **kwargs):
+        async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
             assert kwargs.get("state") == state_data
             for event in mock_events:
                 yield event
 
-        client = AGUIChatClient(endpoint="http://localhost:8888/")
-        monkeypatch.setattr(client._http_service, "post_run", mock_post_run)
+        client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
+        monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
 
         chat_options = ChatOptions()
 
-        response = await client._inner_get_response(messages=messages, chat_options=chat_options)
+        response = await client.inner_get_response(messages=messages, chat_options=chat_options)
 
         assert response is not None

python/packages/ag-ui/tests/test_agent_wrapper_comprehensive.py

@@ -3,21 +3,30 @@
 """Comprehensive tests for AgentFrameworkAgent (_agent.py)."""
 
 import json
+import sys
+from collections.abc import AsyncIterator, MutableSequence
+from pathlib import Path
+from typing import Any
 
 import pytest
-from agent_framework import ChatAgent, TextContent
+from agent_framework import ChatAgent, ChatMessage, ChatOptions, TextContent
 from agent_framework._types import ChatResponseUpdate
+from pydantic import BaseModel
+
+sys.path.insert(0, str(Path(__file__).parent))
+from test_helpers_ag_ui import StreamingChatClientStub
 
 
 async def test_agent_initialization_basic():
     """Test basic agent initialization without state schema."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     wrapper = AgentFrameworkAgent(agent=agent)
 
     assert wrapper.name == "test_agent"
@@ -30,12 +39,13 @@ async def test_agent_initialization_with_state_schema():
     """Test agent initialization with state_schema."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
-    state_schema = {"document": {"type": "string"}}
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
+    state_schema: dict[str, dict[str, Any]] = {"document": {"type": "string"}}
     wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
 
     assert wrapper.config.state_schema == state_schema
@@ -45,31 +55,56 @@ async def test_agent_initialization_with_predict_state_config():
     """Test agent initialization with predict_state_config."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     predict_config = {"document": {"tool": "write_doc", "tool_argument": "content"}}
     wrapper = AgentFrameworkAgent(agent=agent, predict_state_config=predict_config)
 
     assert wrapper.config.predict_state_config == predict_config
 
 
+async def test_agent_initialization_with_pydantic_state_schema():
+    """Test agent initialization when state_schema is provided as Pydantic model/class."""
+    from agent_framework.ag_ui import AgentFrameworkAgent
+
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
+
+    class MyState(BaseModel):
+        document: str
+        tags: list[str] = []
+
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
+
+    wrapper_class_schema = AgentFrameworkAgent(agent=agent, state_schema=MyState)
+    wrapper_instance_schema = AgentFrameworkAgent(agent=agent, state_schema=MyState(document="hi"))
+
+    expected_properties = MyState.model_json_schema().get("properties", {})
+    assert wrapper_class_schema.config.state_schema == expected_properties
+    assert wrapper_instance_schema.config.state_schema == expected_properties
+
+
 async def test_run_started_event_emission():
     """Test RunStartedEvent is emitted at start of run."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     wrapper = AgentFrameworkAgent(agent=agent)
 
     input_data = {"messages": [{"role": "user", "content": "Hi"}]}
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
@@ -83,11 +118,12 @@ async def test_predict_state_custom_event_emission():
     """Test PredictState CustomEvent is emitted when predict_state_config is present."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     predict_config = {
         "document": {"tool": "write_doc", "tool_argument": "content"},
         "summary": {"tool": "summarize", "tool_argument": "text"},
@@ -96,7 +132,7 @@ async def test_predict_state_custom_event_emission():
 
     input_data = {"messages": [{"role": "user", "content": "Hi"}]}
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
@@ -114,11 +150,12 @@ async def test_initial_state_snapshot_with_schema():
     """Test initial StateSnapshotEvent emission when state_schema present."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     state_schema = {"document": {"type": "string"}}
     wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
 
@@ -127,7 +164,7 @@ async def test_initial_state_snapshot_with_schema():
         "state": {"document": "Initial content"},
     }
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
@@ -143,17 +180,18 @@ async def test_state_initialization_object_type():
     """Test state initialization with object type in schema."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
-    state_schema = {"recipe": {"type": "object", "properties": {}}}
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
+    state_schema: dict[str, dict[str, Any]] = {"recipe": {"type": "object", "properties": {}}}
     wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
 
     input_data = {"messages": [{"role": "user", "content": "Hi"}]}
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
@@ -169,17 +207,18 @@ async def test_state_initialization_array_type():
     """Test state initialization with array type in schema."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
-    state_schema = {"steps": {"type": "array", "items": {}}}
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
+    state_schema: dict[str, dict[str, Any]] = {"steps": {"type": "array", "items": {}}}
     wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
 
     input_data = {"messages": [{"role": "user", "content": "Hi"}]}
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
@@ -195,16 +234,17 @@ async def test_run_finished_event_emission():
     """Test RunFinishedEvent is emitted at end of run."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     wrapper = AgentFrameworkAgent(agent=agent)
 
     input_data = {"messages": [{"role": "user", "content": "Hi"}]}
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
@@ -216,11 +256,12 @@ async def test_tool_result_confirm_changes_accepted():
     """Test confirm_changes tool result handling when accepted."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            yield ChatResponseUpdate(contents=[TextContent(text="Document updated")])
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        yield ChatResponseUpdate(contents=[TextContent(text="Document updated")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     wrapper = AgentFrameworkAgent(
         agent=agent,
         state_schema={"document": {"type": "string"}},
@@ -228,8 +269,8 @@ async def test_tool_result_confirm_changes_accepted():
     )
 
     # Simulate tool result message with acceptance
-    tool_result = {"accepted": True, "steps": []}
-    input_data = {
+    tool_result: dict[str, Any] = {"accepted": True, "steps": []}
+    input_data: dict[str, Any] = {
         "messages": [
             {
                 "role": "tool",  # Tool result from UI
@@ -240,7 +281,7 @@ async def test_tool_result_confirm_changes_accepted():
         "state": {"document": "Updated content"},
     }
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
@@ -262,16 +303,17 @@ async def test_tool_result_confirm_changes_rejected():
     """Test confirm_changes tool result handling when rejected."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            yield ChatResponseUpdate(contents=[TextContent(text="OK")])
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        yield ChatResponseUpdate(contents=[TextContent(text="OK")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     wrapper = AgentFrameworkAgent(agent=agent)
 
     # Simulate tool result message with rejection
-    tool_result = {"accepted": False, "steps": []}
-    input_data = {
+    tool_result: dict[str, Any] = {"accepted": False, "steps": []}
+    input_data: dict[str, Any] = {
         "messages": [
             {
                 "role": "tool",
@@ -281,7 +323,7 @@ async def test_tool_result_confirm_changes_rejected():
         ],
     }
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
@@ -295,22 +337,23 @@ async def test_tool_result_function_approval_accepted():
     """Test function approval tool result when steps are accepted."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            yield ChatResponseUpdate(contents=[TextContent(text="OK")])
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        yield ChatResponseUpdate(contents=[TextContent(text="OK")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     wrapper = AgentFrameworkAgent(agent=agent)
 
     # Simulate tool result with multiple steps
-    tool_result = {
+    tool_result: dict[str, Any] = {
         "accepted": True,
         "steps": [
             {"id": "step1", "description": "Send email", "status": "enabled"},
             {"id": "step2", "description": "Create calendar event", "status": "enabled"},
         ],
     }
-    input_data = {
+    input_data: dict[str, Any] = {
         "messages": [
             {
                 "role": "tool",
@@ -320,7 +363,7 @@ async def test_tool_result_function_approval_accepted():
         ],
     }
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
@@ -340,19 +383,20 @@ async def test_tool_result_function_approval_rejected():
     """Test function approval tool result when rejected."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            yield ChatResponseUpdate(contents=[TextContent(text="OK")])
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        yield ChatResponseUpdate(contents=[TextContent(text="OK")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     wrapper = AgentFrameworkAgent(agent=agent)
 
     # Simulate tool result rejection with steps
-    tool_result = {
+    tool_result: dict[str, Any] = {
         "accepted": False,
         "steps": [{"id": "step1", "description": "Send email", "status": "disabled"}],
     }
-    input_data = {
+    input_data: dict[str, Any] = {
         "messages": [
             {
                 "role": "tool",
@@ -362,7 +406,7 @@ async def test_tool_result_function_approval_rejected():
         ],
     }
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
@@ -376,17 +420,16 @@ async def test_thread_metadata_tracking():
     """Test that thread metadata includes ag_ui_thread_id and ag_ui_run_id."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    thread_metadata = {}
+    thread_metadata: dict[str, Any] = {}
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            # Capture thread metadata from kwargs
-            nonlocal thread_metadata
-            if "thread" in kwargs:
-                thread_metadata = kwargs["thread"].metadata
-            yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        if chat_options.metadata:
+            thread_metadata.update(chat_options.metadata)
+        yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     wrapper = AgentFrameworkAgent(agent=agent)
 
     input_data = {
@@ -395,28 +438,28 @@ async def test_thread_metadata_tracking():
         "run_id": "test_run_456",
     }
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
-    # Check thread metadata was set
-    # Note: This test may need adjustment based on actual thread passing mechanism
+    assert thread_metadata.get("ag_ui_thread_id") == "test_thread_123"
+    assert thread_metadata.get("ag_ui_run_id") == "test_run_456"
 
 
 async def test_state_context_injection():
     """Test that current state is injected into thread metadata."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    thread_metadata = {}
+    thread_metadata: dict[str, Any] = {}
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            # Track if state context message was added
-            nonlocal thread_metadata
-            # In actual implementation, thread is passed and state is in metadata
-            yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        if chat_options.metadata:
+            thread_metadata.update(chat_options.metadata)
+        yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     wrapper = AgentFrameworkAgent(
         agent=agent,
         state_schema={"document": {"type": "string"}},
@@ -427,27 +470,31 @@ async def test_state_context_injection():
         "state": {"document": "Test content"},
     }
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
-    # State should be injected - this is validated by agent execution flow
+    current_state = thread_metadata.get("current_state")
+    if isinstance(current_state, str):
+        current_state = json.loads(current_state)
+    assert current_state == {"document": "Test content"}
 
 
 async def test_no_messages_provided():
     """Test handling when no messages are provided."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     wrapper = AgentFrameworkAgent(agent=agent)
 
-    input_data = {"messages": []}
+    input_data: dict[str, Any] = {"messages": []}
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
@@ -461,16 +508,17 @@ async def test_message_end_event_emission():
     """Test TextMessageEndEvent is emitted for assistant messages."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            yield ChatResponseUpdate(contents=[TextContent(text="Hello world")])
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        yield ChatResponseUpdate(contents=[TextContent(text="Hello world")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     wrapper = AgentFrameworkAgent(agent=agent)
 
-    input_data = {"messages": [{"role": "user", "content": "Hi"}]}
+    input_data: dict[str, Any] = {"messages": [{"role": "user", "content": "Hi"}]}
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
@@ -488,19 +536,20 @@ async def test_error_handling_with_exception():
     """Test that exceptions during agent execution are re-raised."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class FailingChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            if False:
-                yield
-            raise RuntimeError("Simulated failure")
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        if False:
+            yield ChatResponseUpdate(contents=[])
+        raise RuntimeError("Simulated failure")
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=FailingChatClient())
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     wrapper = AgentFrameworkAgent(agent=agent)
 
-    input_data = {"messages": [{"role": "user", "content": "Hi"}]}
+    input_data: dict[str, Any] = {"messages": [{"role": "user", "content": "Hi"}]}
 
     with pytest.raises(RuntimeError, match="Simulated failure"):
-        async for event in wrapper.run_agent(input_data):
+        async for _ in wrapper.run_agent(input_data):
             pass
 
 
@@ -508,18 +557,18 @@ async def test_json_decode_error_in_tool_result():
     """Test handling of orphaned tool result - should be sanitized out."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            # Should not be called since orphaned tool result is dropped
-            if False:
-                yield
-            raise AssertionError("ChatClient should not be called with orphaned tool result")
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        if False:
+            yield ChatResponseUpdate(contents=[])
+        raise AssertionError("ChatClient should not be called with orphaned tool result")
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     wrapper = AgentFrameworkAgent(agent=agent)
 
     # Send invalid JSON as tool result without preceding tool call
-    input_data = {
+    input_data: dict[str, Any] = {
         "messages": [
             {
                 "role": "tool",
@@ -529,7 +578,7 @@ async def test_json_decode_error_in_tool_result():
         ],
     }
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
@@ -545,11 +594,12 @@ 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
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            yield ChatResponseUpdate(contents=[TextContent(text="Response")])
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        yield ChatResponseUpdate(contents=[TextContent(text="Response")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     wrapper = AgentFrameworkAgent(
         agent=agent,
         state_schema={"document": {"type": "string"}},
@@ -558,8 +608,8 @@ async def test_suppressed_summary_with_document_state():
     )
 
     # Simulate confirmation with document state
-    tool_result = {"accepted": True, "steps": []}
-    input_data = {
+    tool_result: dict[str, Any] = {"accepted": True, "steps": []}
+    input_data: dict[str, Any] = {
         "messages": [
             {
                 "role": "tool",
@@ -570,7 +620,7 @@ async def test_suppressed_summary_with_document_state():
         "state": {"document": "This is the beginning of a document. It contains important information."},
     }
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 

python/packages/ag-ui/tests/test_backend_tool_rendering.py

@@ -2,6 +2,8 @@
 
 """Tests for backend tool rendering."""
 
+from typing import cast
+
 from ag_ui.core import (
     TextMessageContentEvent,
     TextMessageStartEvent,
@@ -119,6 +121,9 @@ async def test_multiple_tool_results():
         assert isinstance(events[end_idx], ToolCallEndEvent)
         assert isinstance(events[result_idx], ToolCallResultEvent)
 
-        assert events[end_idx].tool_call_id == f"tool-{i + 1}"
-        assert events[result_idx].tool_call_id == f"tool-{i + 1}"
-        assert f"Result {i + 1}" in events[result_idx].content
+        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

@@ -14,7 +14,7 @@ from agent_framework_ag_ui._confirmation_strategies import (
 
 
 @pytest.fixture
-def sample_steps():
+def sample_steps() -> list[dict[str, str]]:
     """Sample steps for testing approval messages."""
     return [
         {"description": "Step 1: Do something", "status": "enabled"},
@@ -24,7 +24,7 @@ def sample_steps():
 
 
 @pytest.fixture
-def all_enabled_steps():
+def all_enabled_steps() -> list[dict[str, str]]:
     """All steps enabled."""
     return [
         {"description": "Task A", "status": "enabled"},
@@ -34,7 +34,7 @@ def all_enabled_steps():
 
 
 @pytest.fixture
-def empty_steps():
+def empty_steps() -> list[dict[str, str]]:
     """Empty steps list."""
     return []
 
@@ -42,7 +42,7 @@ def empty_steps():
 class TestDefaultConfirmationStrategy:
     """Tests for DefaultConfirmationStrategy."""
 
-    def test_on_approval_accepted_with_enabled_steps(self, sample_steps):
+    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)
 
@@ -52,7 +52,7 @@ class TestDefaultConfirmationStrategy:
         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):
+    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)
 
@@ -61,28 +61,28 @@ class TestDefaultConfirmationStrategy:
         assert "Task B" in message
         assert "Task C" in message
 
-    def test_on_approval_accepted_with_empty_steps(self, empty_steps):
+    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):
+    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):
+    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):
+    def test_on_state_rejected(self) -> None:
         strategy = DefaultConfirmationStrategy()
         message = strategy.on_state_rejected()
 
@@ -93,7 +93,7 @@ class TestDefaultConfirmationStrategy:
 class TestTaskPlannerConfirmationStrategy:
     """Tests for TaskPlannerConfirmationStrategy."""
 
-    def test_on_approval_accepted_with_enabled_steps(self, sample_steps):
+    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)
 
@@ -103,7 +103,7 @@ class TestTaskPlannerConfirmationStrategy:
         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):
+    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)
 
@@ -112,28 +112,28 @@ class TestTaskPlannerConfirmationStrategy:
         assert "2. Task B" in message
         assert "3. Task C" in message
 
-    def test_on_approval_accepted_with_empty_steps(self, empty_steps):
+    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):
+    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):
+    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):
+    def test_on_state_rejected(self) -> None:
         strategy = TaskPlannerConfirmationStrategy()
         message = strategy.on_state_rejected()
 
@@ -144,7 +144,7 @@ class TestTaskPlannerConfirmationStrategy:
 class TestRecipeConfirmationStrategy:
     """Tests for RecipeConfirmationStrategy."""
 
-    def test_on_approval_accepted_with_enabled_steps(self, sample_steps):
+    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)
 
@@ -154,7 +154,7 @@ class TestRecipeConfirmationStrategy:
         assert "Step 3" not in message
         assert "Recipe updated successfully!" in message
 
-    def test_on_approval_accepted_with_all_enabled(self, all_enabled_steps):
+    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)
 
@@ -163,28 +163,28 @@ class TestRecipeConfirmationStrategy:
         assert "2. Task B" in message
         assert "3. Task C" in message
 
-    def test_on_approval_accepted_with_empty_steps(self, empty_steps):
+    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):
+    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):
+    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):
+    def test_on_state_rejected(self) -> None:
         strategy = RecipeConfirmationStrategy()
         message = strategy.on_state_rejected()
 
@@ -195,7 +195,7 @@ class TestRecipeConfirmationStrategy:
 class TestDocumentWriterConfirmationStrategy:
     """Tests for DocumentWriterConfirmationStrategy."""
 
-    def test_on_approval_accepted_with_enabled_steps(self, sample_steps):
+    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)
 
@@ -205,7 +205,7 @@ class TestDocumentWriterConfirmationStrategy:
         assert "Step 3" not in message
         assert "Document updated successfully!" in message
 
-    def test_on_approval_accepted_with_all_enabled(self, all_enabled_steps):
+    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)
 
@@ -214,27 +214,27 @@ class TestDocumentWriterConfirmationStrategy:
         assert "2. Task B" in message
         assert "3. Task C" in message
 
-    def test_on_approval_accepted_with_empty_steps(self, empty_steps):
+    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):
+    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):
+    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):
+    def test_on_state_rejected(self) -> None:
         strategy = DocumentWriterConfirmationStrategy()
         message = strategy.on_state_rejected()
 

python/packages/ag-ui/tests/test_document_writer_flow.py

@@ -2,7 +2,7 @@
 
 """Tests for document writer predictive state flow with confirm_changes."""
 
-from ag_ui.core import EventType
+from ag_ui.core import EventType, StateDeltaEvent, ToolCallArgsEvent, ToolCallEndEvent, ToolCallStartEvent
 from agent_framework import FunctionCallContent, FunctionResultContent, TextContent
 from agent_framework._types import AgentRunResponseUpdate
 
@@ -35,16 +35,12 @@ async def test_streaming_document_with_state_deltas():
     assert any(e.type == EventType.TOOL_CALL_ARGS for e in events1)
 
     # Second chunk - incomplete JSON, should try partial extraction
-    tool_call_chunk2 = FunctionCallContent(
-        call_id="call_123",
-        name=None,  # Name only in first chunk
-        arguments=" upon a time",
-    )
+    tool_call_chunk2 = FunctionCallContent(call_id="call_123", name="write_document_local", arguments=" upon a time")
     update2 = AgentRunResponseUpdate(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 e.type == EventType.STATE_DELTA]
+    state_deltas = [e for e in events2 if isinstance(e, StateDeltaEvent)]
     assert len(state_deltas) >= 1
 
     # Check JSON Patch format
@@ -62,7 +58,7 @@ async def test_confirm_changes_emission():
         "document": {"tool": "write_document_local", "tool_argument": "document"},
     }
 
-    current_state = {}
+    current_state: dict[str, str] = {}
 
     bridge = AgentFrameworkEventBridge(
         run_id="test_run",
@@ -90,15 +86,13 @@ async def test_confirm_changes_emission():
     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 e.type == EventType.TOOL_CALL_START and e.tool_call_name == "confirm_changes"
-    ]
+    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 e.type == EventType.TOOL_CALL_ARGS and e.delta == "{}"]
+    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 e.type == EventType.TOOL_CALL_END]
+    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
 
@@ -141,7 +135,7 @@ async def test_no_confirm_for_non_predictive_tools():
         "document": {"tool": "write_document_local", "tool_argument": "document"},
     }
 
-    current_state = {}
+    current_state: dict[str, str] = {}
 
     bridge = AgentFrameworkEventBridge(
         run_id="test_run",
@@ -162,9 +156,7 @@ async def test_no_confirm_for_non_predictive_tools():
     events = await bridge.from_agent_run_update(update)
 
     # Should NOT have confirm_changes
-    confirm_starts = [
-        e for e in events if e.type == EventType.TOOL_CALL_START and e.tool_call_name == "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
@@ -193,14 +185,14 @@ async def test_state_delta_deduplication():
     events1 = await bridge.from_agent_run_update(update1)
 
     # Count state deltas
-    state_deltas_1 = [e for e in events1 if e.type == EventType.STATE_DELTA]
+    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 = FunctionCallContent(
         call_id="call_2",
-        name=None,
+        name="write_document_local",
         arguments='{"document":"Same text"}',  # Identical content
     )
     update2 = AgentRunResponseUpdate(contents=[tool_call2])
@@ -234,7 +226,7 @@ async def test_predict_state_config_multiple_fields():
     events = await bridge.from_agent_run_update(update)
 
     # Should emit StateDeltaEvent for both fields
-    state_deltas = [e for e in events if e.type == EventType.STATE_DELTA]
+    state_deltas = [e for e in events if isinstance(e, StateDeltaEvent)]
     assert len(state_deltas) >= 2
 
     # Check both fields are present

python/packages/ag-ui/tests/test_endpoint.py

@@ -3,7 +3,8 @@
 """Tests for FastAPI endpoint creation (_endpoint.py)."""
 
 import json
-from typing import Any
+import sys
+from pathlib import Path
 
 from agent_framework import ChatAgent, TextContent
 from agent_framework._types import ChatResponseUpdate
@@ -13,22 +14,20 @@ from fastapi.testclient import TestClient
 from agent_framework_ag_ui._agent import AgentFrameworkAgent
 from agent_framework_ag_ui._endpoint import add_agent_framework_fastapi_endpoint
 
+sys.path.insert(0, str(Path(__file__).parent))
+from test_helpers_ag_ui import StreamingChatClientStub, stream_from_updates
 
-class MockChatClient:
-    """Mock chat client for testing."""
 
-    def __init__(self, response_text: str = "Test response"):
-        self.response_text = response_text
-
-    async def get_streaming_response(self, messages: list[Any], chat_options: Any, **kwargs: Any):
-        """Mock streaming response."""
-        yield ChatResponseUpdate(contents=[TextContent(text=self.response_text)])
+def build_chat_client(response_text: str = "Test response") -> StreamingChatClientStub:
+    """Create a typed chat client stub for endpoint tests."""
+    updates = [ChatResponseUpdate(contents=[TextContent(text=response_text)])]
+    return StreamingChatClientStub(stream_from_updates(updates))
 
 
 async def test_add_endpoint_with_agent_protocol():
     """Test adding endpoint with raw AgentProtocol."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/test-agent")
 
@@ -42,7 +41,7 @@ async def test_add_endpoint_with_agent_protocol():
 async def test_add_endpoint_with_wrapped_agent():
     """Test adding endpoint with pre-wrapped AgentFrameworkAgent."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
     wrapped_agent = AgentFrameworkAgent(agent=agent, name="wrapped")
 
     add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/wrapped-agent")
@@ -57,7 +56,7 @@ async def test_add_endpoint_with_wrapped_agent():
 async def test_endpoint_with_state_schema():
     """Test endpoint with state_schema parameter."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
     state_schema = {"document": {"type": "string"}}
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/stateful", state_schema=state_schema)
@@ -70,10 +69,37 @@ async def test_endpoint_with_state_schema():
     assert response.status_code == 200
 
 
+async def test_endpoint_with_default_state_seed():
+    """Test endpoint seeds default state when client omits it."""
+    app = FastAPI()
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
+    state_schema = {"proverbs": {"type": "array"}}
+    default_state = {"proverbs": ["Keep the original."]}
+
+    add_agent_framework_fastapi_endpoint(
+        app,
+        agent,
+        path="/default-state",
+        state_schema=state_schema,
+        default_state=default_state,
+    )
+
+    client = TestClient(app)
+    response = client.post("/default-state", json={"messages": [{"role": "user", "content": "Hello"}]})
+
+    assert response.status_code == 200
+
+    content = response.content.decode("utf-8")
+    lines = [line for line in content.split("\n") if line.startswith("data: ")]
+    snapshots = [json.loads(line[6:]) for line in lines if json.loads(line[6:]).get("type") == "STATE_SNAPSHOT"]
+    assert snapshots, "Expected a STATE_SNAPSHOT event"
+    assert snapshots[0]["snapshot"]["proverbs"] == default_state["proverbs"]
+
+
 async def test_endpoint_with_predict_state_config():
     """Test endpoint with predict_state_config parameter."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
     predict_config = {"document": {"tool": "write_doc", "tool_argument": "content"}}
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/predictive", predict_state_config=predict_config)
@@ -87,7 +113,7 @@ async def test_endpoint_with_predict_state_config():
 async def test_endpoint_request_logging():
     """Test that endpoint logs request details."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/logged")
 
@@ -107,7 +133,7 @@ async def test_endpoint_request_logging():
 async def test_endpoint_event_streaming():
     """Test that endpoint streams events correctly."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient("Streamed response"))
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client("Streamed response"))
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/stream")
 
@@ -141,14 +167,14 @@ async def test_endpoint_event_streaming():
 async def test_endpoint_error_handling():
     """Test endpoint error handling during request parsing."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/failing")
 
     client = TestClient(app)
 
     # Send invalid JSON to trigger parsing error before streaming
-    response = client.post("/failing", data="invalid json", headers={"content-type": "application/json"})
+    response = client.post("/failing", data=b"invalid json", headers={"content-type": "application/json"})  # type: ignore
 
     # The exception handler catches it and returns JSON error
     assert response.status_code == 200
@@ -160,8 +186,8 @@ async def test_endpoint_error_handling():
 async def test_endpoint_multiple_paths():
     """Test adding multiple endpoints with different paths."""
     app = FastAPI()
-    agent1 = ChatAgent(name="agent1", instructions="First agent", chat_client=MockChatClient("Response 1"))
-    agent2 = ChatAgent(name="agent2", instructions="Second agent", chat_client=MockChatClient("Response 2"))
+    agent1 = ChatAgent(name="agent1", instructions="First agent", chat_client=build_chat_client("Response 1"))
+    agent2 = ChatAgent(name="agent2", instructions="Second agent", chat_client=build_chat_client("Response 2"))
 
     add_agent_framework_fastapi_endpoint(app, agent1, path="/agent1")
     add_agent_framework_fastapi_endpoint(app, agent2, path="/agent2")
@@ -178,7 +204,7 @@ async def test_endpoint_multiple_paths():
 async def test_endpoint_default_path():
     """Test endpoint with default path."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
 
     add_agent_framework_fastapi_endpoint(app, agent)
 
@@ -191,7 +217,7 @@ async def test_endpoint_default_path():
 async def test_endpoint_response_headers():
     """Test that endpoint sets correct response headers."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/headers")
 
@@ -207,7 +233,7 @@ async def test_endpoint_response_headers():
 async def test_endpoint_empty_messages():
     """Test endpoint with empty messages list."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/empty")
 
@@ -220,7 +246,7 @@ async def test_endpoint_empty_messages():
 async def test_endpoint_complex_input():
     """Test endpoint with complex input data."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/complex")
 

python/packages/ag-ui/tests/test_event_converters.py

@@ -1,3 +1,5 @@
+# Copyright (c) Microsoft. All rights reserved.
+
 """Tests for AG-UI event converter."""
 
 from agent_framework import FinishReason, Role

python/packages/ag-ui/tests/test_helpers_ag_ui.py

@@ -0,0 +1,138 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Shared test stubs for AG-UI tests."""
+
+from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, MutableSequence
+from types import SimpleNamespace
+from typing import Any
+
+from agent_framework import (
+    AgentProtocol,
+    AgentRunResponse,
+    AgentRunResponseUpdate,
+    AgentThread,
+    ChatMessage,
+    ChatOptions,
+    TextContent,
+)
+from agent_framework._clients import BaseChatClient
+from agent_framework._types import ChatResponse, ChatResponseUpdate
+
+from agent_framework_ag_ui._orchestrators import ExecutionContext
+
+StreamFn = Callable[..., AsyncIterator[ChatResponseUpdate]]
+ResponseFn = Callable[..., Awaitable[ChatResponse]]
+
+
+class StreamingChatClientStub(BaseChatClient):
+    """Typed streaming stub that satisfies ChatClientProtocol."""
+
+    def __init__(self, stream_fn: StreamFn, response_fn: ResponseFn | None = None) -> None:
+        super().__init__()
+        self._stream_fn = stream_fn
+        self._response_fn = response_fn
+
+    async def _inner_get_streaming_response(
+        self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        async for update in self._stream_fn(messages, chat_options, **kwargs):
+            yield update
+
+    async def _inner_get_response(
+        self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> ChatResponse:
+        if self._response_fn is not None:
+            return await self._response_fn(messages, chat_options, **kwargs)
+
+        contents: list[Any] = []
+        async for update in self._stream_fn(messages, chat_options, **kwargs):
+            contents.extend(update.contents)
+
+        return ChatResponse(
+            messages=[ChatMessage(role="assistant", contents=contents)],
+            response_id="stub-response",
+        )
+
+
+def stream_from_updates(updates: list[ChatResponseUpdate]) -> StreamFn:
+    """Create a stream function that yields from a static list of updates."""
+
+    async def _stream(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        for update in updates:
+            yield update
+
+    return _stream
+
+
+class StubAgent(AgentProtocol):
+    """Minimal AgentProtocol stub for orchestrator tests."""
+
+    def __init__(
+        self,
+        updates: list[AgentRunResponseUpdate] | None = None,
+        *,
+        agent_id: str = "stub-agent",
+        agent_name: str | None = "stub-agent",
+        chat_options: Any | None = None,
+        chat_client: Any | None = None,
+    ) -> None:
+        self._id = agent_id
+        self._name = agent_name
+        self._description = "stub agent"
+        self.updates = updates or [AgentRunResponseUpdate(contents=[TextContent(text="response")], role="assistant")]
+        self.chat_options = chat_options or SimpleNamespace(tools=None, response_format=None)
+        self.chat_client = chat_client or SimpleNamespace(function_invocation_configuration=None)
+        self.messages_received: list[Any] = []
+        self.tools_received: list[Any] | None = None
+
+    @property
+    def id(self) -> str:
+        return self._id
+
+    @property
+    def name(self) -> str | None:
+        return self._name
+
+    @property
+    def display_name(self) -> str:
+        return self._name or self._id
+
+    @property
+    def description(self) -> str | None:
+        return self._description
+
+    async def run(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
+        *,
+        thread: AgentThread | None = None,
+        **kwargs: Any,
+    ) -> AgentRunResponse:
+        return AgentRunResponse(messages=[], response_id="stub-response")
+
+    def run_stream(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
+        *,
+        thread: AgentThread | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterable[AgentRunResponseUpdate]:
+        async def _stream() -> AsyncIterator[AgentRunResponseUpdate]:
+            self.messages_received = [] if messages is None else list(messages)  # type: ignore[arg-type]
+            self.tools_received = kwargs.get("tools")
+            for update in self.updates:
+                yield update
+
+        return _stream()
+
+    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]) -> None:
+        self._messages = messages

python/packages/ag-ui/tests/test_message_hygiene.py

@@ -0,0 +1,53 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from agent_framework import ChatMessage, FunctionCallContent, FunctionResultContent, TextContent
+
+from agent_framework_ag_ui._orchestration._message_hygiene import (
+    deduplicate_messages,
+    sanitize_tool_history,
+)
+
+
+def test_sanitize_tool_history_injects_confirm_changes_result() -> None:
+    messages = [
+        ChatMessage(
+            role="assistant",
+            contents=[
+                FunctionCallContent(
+                    name="confirm_changes",
+                    call_id="call_confirm_123",
+                    arguments='{"changes": "test"}',
+                )
+            ],
+        ),
+        ChatMessage(
+            role="user",
+            contents=[TextContent(text='{"accepted": true}')],
+        ),
+    ]
+
+    sanitized = sanitize_tool_history(messages)
+
+    tool_messages = [
+        msg for msg in sanitized 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"
+
+
+def test_deduplicate_messages_prefers_non_empty_tool_results() -> None:
+    messages = [
+        ChatMessage(
+            role="tool",
+            contents=[FunctionResultContent(call_id="call1", result="")],
+        ),
+        ChatMessage(
+            role="tool",
+            contents=[FunctionResultContent(call_id="call1", result="result data")],
+        ),
+    ]
+
+    deduped = deduplicate_messages(messages)
+    assert len(deduped) == 1
+    assert deduped[0].contents[0].result == "result data"

python/packages/ag-ui/tests/test_orchestrators.py

@@ -1,3 +1,5 @@
+# Copyright (c) Microsoft. All rights reserved.
+
 """Tests for AG-UI orchestrators."""
 
 from collections.abc import AsyncGenerator
@@ -34,6 +36,7 @@ class DummyAgent:
         *,
         thread: Any,
         tools: list[Any] | None = None,
+        **kwargs: Any,
     ) -> AsyncGenerator[AgentRunResponseUpdate, None]:
         self.seen_tools = tools
         yield AgentRunResponseUpdate(contents=[TextContent(text="ok")], role="assistant")

python/packages/ag-ui/tests/test_orchestrators_coverage.py

@@ -2,7 +2,9 @@
 
 """Comprehensive tests for orchestrator coverage."""
 
+import sys
 from collections.abc import AsyncGenerator
+from pathlib import Path
 from types import SimpleNamespace
 from typing import Any
 
@@ -15,11 +17,10 @@ from agent_framework import (
 from pydantic import BaseModel
 
 from agent_framework_ag_ui._agent import AgentConfig
-from agent_framework_ag_ui._orchestrators import (
-    DefaultOrchestrator,
-    ExecutionContext,
-    HumanInTheLoopOrchestrator,
-)
+from agent_framework_ag_ui._orchestrators import DefaultOrchestrator, HumanInTheLoopOrchestrator
+
+sys.path.insert(0, str(Path(__file__).parent))
+from test_helpers_ag_ui import StubAgent, TestExecutionContext
 
 
 @ai_function(approval_mode="always_require")
@@ -28,34 +29,14 @@ def approval_tool(param: str) -> str:
     return f"executed: {param}"
 
 
-class MockAgent:
-    """Mock agent for testing."""
-
-    def __init__(self, updates: list[AgentRunResponseUpdate] | None = None) -> None:
-        self.updates = updates or [AgentRunResponseUpdate(contents=[TextContent(text="response")], role="assistant")]
-        self.chat_options = SimpleNamespace(tools=[approval_tool], response_format=None)
-        self.chat_client = SimpleNamespace(function_invocation_configuration=None)
-        self.messages_received: list[Any] = []
-        self.tools_received: list[Any] | None = None
-
-    async def run_stream(
-        self,
-        messages: list[Any],
-        *,
-        thread: Any = None,
-        tools: list[Any] | None = None,
-    ) -> AsyncGenerator[AgentRunResponseUpdate, None]:
-        self.messages_received = messages
-        self.tools_received = tools
-        for update in self.updates:
-            yield update
+DEFAULT_CHAT_OPTIONS = SimpleNamespace(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 = {
+    input_data: dict[str, Any] = {
         "messages": [
             {
                 "role": "tool",
@@ -72,21 +53,25 @@ async def test_human_in_the_loop_json_decode_error() -> None:
         )
     ]
 
-    context = ExecutionContext(
+    agent = StubAgent(
+        chat_options=SimpleNamespace(tools=[approval_tool], response_format=None),
+        updates=[AgentRunResponseUpdate(contents=[TextContent(text="response")], role="assistant")],
+    )
+    context = TestExecutionContext(
         input_data=input_data,
-        agent=MockAgent(),
+        agent=agent,
         config=AgentConfig(),
     )
-    context._messages = messages
+    context.set_messages(messages)
 
     assert orchestrator.can_handle(context)
 
-    events = []
+    events: list[Any] = []
     async for event in orchestrator.run(context):
         events.append(event)
 
     # Should emit RunErrorEvent for invalid JSON
-    error_events = [e for e in events if e.type == "RUN_ERROR"]
+    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
 
@@ -118,18 +103,20 @@ async def test_sanitize_tool_history_confirm_changes() -> None:
     orchestrator = DefaultOrchestrator()
 
     # Use pre-constructed ChatMessage objects to bypass message adapter
-    input_data = {"messages": []}
+    input_data: dict[str, Any] = {"messages": []}
 
-    agent = MockAgent()
-    context = ExecutionContext(
+    agent = StubAgent(
+        chat_options=DEFAULT_CHAT_OPTIONS,
+    )
+    context = TestExecutionContext(
         input_data=input_data,
         agent=agent,
         config=AgentConfig(),
     )
     # Override the messages property to use our pre-constructed messages
-    context._messages = messages
+    context.set_messages(messages)
 
-    events = []
+    events: list[Any] = []
     async for event in orchestrator.run(context):
         events.append(event)
 
@@ -162,16 +149,18 @@ async def test_sanitize_tool_history_orphaned_tool_result() -> None:
     ]
 
     orchestrator = DefaultOrchestrator()
-    input_data = {"messages": []}
-    agent = MockAgent()
-    context = ExecutionContext(
+    input_data: dict[str, Any] = {"messages": []}
+    agent = StubAgent(
+        chat_options=DEFAULT_CHAT_OPTIONS,
+    )
+    context = TestExecutionContext(
         input_data=input_data,
         agent=agent,
         config=AgentConfig(),
     )
-    context._messages = messages
+    context.set_messages(messages)
 
-    events = []
+    events: list[Any] = []
     async for event in orchestrator.run(context):
         events.append(event)
 
@@ -188,7 +177,7 @@ async def test_orphaned_tool_result_sanitization() -> None:
     """Test that orphaned tool results are filtered out."""
     orchestrator = DefaultOrchestrator()
 
-    input_data = {
+    input_data: dict[str, Any] = {
         "messages": [
             {
                 "role": "tool",
@@ -201,14 +190,16 @@ async def test_orphaned_tool_result_sanitization() -> None:
         ],
     }
 
-    agent = MockAgent()
-    context = ExecutionContext(
+    agent = StubAgent(
+        chat_options=DEFAULT_CHAT_OPTIONS,
+    )
+    context = TestExecutionContext(
         input_data=input_data,
         agent=agent,
         config=AgentConfig(),
     )
 
-    events = []
+    events: list[Any] = []
     async for event in orchestrator.run(context):
         events.append(event)
 
@@ -241,16 +232,18 @@ async def test_deduplicate_messages_empty_tool_results() -> None:
     ]
 
     orchestrator = DefaultOrchestrator()
-    input_data = {"messages": []}
-    agent = MockAgent()
-    context = ExecutionContext(
+    input_data: dict[str, Any] = {"messages": []}
+    agent = StubAgent(
+        chat_options=DEFAULT_CHAT_OPTIONS,
+    )
+    context = TestExecutionContext(
         input_data=input_data,
         agent=agent,
         config=AgentConfig(),
     )
-    context._messages = messages
+    context.set_messages(messages)
 
-    events = []
+    events: list[Any] = []
     async for event in orchestrator.run(context):
         events.append(event)
 
@@ -284,16 +277,18 @@ async def test_deduplicate_messages_duplicate_assistant_tool_calls() -> None:
     ]
 
     orchestrator = DefaultOrchestrator()
-    input_data = {"messages": []}
-    agent = MockAgent()
-    context = ExecutionContext(
+    input_data: dict[str, Any] = {"messages": []}
+    agent = StubAgent(
+        chat_options=DEFAULT_CHAT_OPTIONS,
+    )
+    context = TestExecutionContext(
         input_data=input_data,
         agent=agent,
         config=AgentConfig(),
     )
-    context._messages = messages
+    context.set_messages(messages)
 
-    events = []
+    events: list[Any] = []
     async for event in orchestrator.run(context):
         events.append(event)
 
@@ -326,16 +321,18 @@ async def test_deduplicate_messages_duplicate_system_messages() -> None:
     ]
 
     orchestrator = DefaultOrchestrator()
-    input_data = {"messages": []}
-    agent = MockAgent()
-    context = ExecutionContext(
+    input_data: dict[str, Any] = {"messages": []}
+    agent = StubAgent(
+        chat_options=DEFAULT_CHAT_OPTIONS,
+    )
+    context = TestExecutionContext(
         input_data=input_data,
         agent=agent,
         config=AgentConfig(),
     )
-    context._messages = messages
+    context.set_messages(messages)
 
-    events = []
+    events: list[Any] = []
     async for event in orchestrator.run(context):
         events.append(event)
 
@@ -354,7 +351,7 @@ async def test_state_context_injection() -> None:
     """Test state context message injection for first request."""
     orchestrator = DefaultOrchestrator()
 
-    input_data = {
+    input_data: dict[str, Any] = {
         "messages": [
             {
                 "role": "user",
@@ -364,14 +361,16 @@ async def test_state_context_injection() -> None:
         "state": {"items": ["apple", "banana"]},
     }
 
-    agent = MockAgent()
-    context = ExecutionContext(
+    agent = StubAgent(
+        chat_options=DEFAULT_CHAT_OPTIONS,
+    )
+    context = TestExecutionContext(
         input_data=input_data,
         agent=agent,
         config=AgentConfig(state_schema={"items": {"type": "array"}}),
     )
 
-    events = []
+    events: list[Any] = []
     async for event in orchestrator.run(context):
         events.append(event)
 
@@ -406,16 +405,18 @@ async def test_no_state_context_injection_with_tool_calls() -> None:
     ]
 
     orchestrator = DefaultOrchestrator()
-    input_data = {"messages": [], "state": {"weather": "sunny"}}
-    agent = MockAgent()
-    context = ExecutionContext(
+    input_data: dict[str, Any] = {"messages": [], "state": {"weather": "sunny"}}
+    agent = StubAgent(
+        chat_options=DEFAULT_CHAT_OPTIONS,
+    )
+    context = TestExecutionContext(
         input_data=input_data,
         agent=agent,
         config=AgentConfig(state_schema={"weather": {"type": "string"}}),
     )
-    context._messages = messages
+    context.set_messages(messages)
 
-    events = []
+    events: list[Any] = []
     async for event in orchestrator.run(context):
         events.append(event)
 
@@ -437,7 +438,7 @@ async def test_structured_output_processing() -> None:
 
     orchestrator = DefaultOrchestrator()
 
-    input_data = {
+    input_data: dict[str, Any] = {
         "messages": [
             {
                 "role": "user",
@@ -447,32 +448,33 @@ async def test_structured_output_processing() -> None:
     }
 
     # Agent with structured output
-    agent = MockAgent(
+    agent = StubAgent(
+        chat_options=DEFAULT_CHAT_OPTIONS,
         updates=[
             AgentRunResponseUpdate(
                 contents=[TextContent(text='{"ingredients": ["tomato"], "message": "Added tomato"}')],
                 role="assistant",
             )
-        ]
+        ],
     )
     agent.chat_options.response_format = RecipeState
 
-    context = ExecutionContext(
+    context = TestExecutionContext(
         input_data=input_data,
         agent=agent,
         config=AgentConfig(state_schema={"ingredients": {"type": "array"}}),
     )
 
-    events = []
+    events: list[Any] = []
     async for event in orchestrator.run(context):
         events.append(event)
 
     # Should emit StateSnapshotEvent with ingredients
-    state_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
+    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 = [e for e in events if e.type == "TEXT_MESSAGE_CONTENT"]
+    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)
 
@@ -487,7 +489,7 @@ async def test_duplicate_client_tools_filtered() -> None:
 
     orchestrator = DefaultOrchestrator()
 
-    input_data = {
+    input_data: dict[str, Any] = {
         "messages": [
             {
                 "role": "user",
@@ -507,16 +509,18 @@ async def test_duplicate_client_tools_filtered() -> None:
         ],
     }
 
-    agent = MockAgent()
+    agent = StubAgent(
+        chat_options=DEFAULT_CHAT_OPTIONS,
+    )
     agent.chat_options.tools = [get_weather]
 
-    context = ExecutionContext(
+    context = TestExecutionContext(
         input_data=input_data,
         agent=agent,
         config=AgentConfig(),
     )
 
-    events = []
+    events: list[Any] = []
     async for event in orchestrator.run(context):
         events.append(event)
 
@@ -534,7 +538,7 @@ async def test_unique_client_tools_merged() -> None:
 
     orchestrator = DefaultOrchestrator()
 
-    input_data = {
+    input_data: dict[str, Any] = {
         "messages": [
             {
                 "role": "user",
@@ -554,16 +558,18 @@ async def test_unique_client_tools_merged() -> None:
         ],
     }
 
-    agent = MockAgent()
+    agent = StubAgent(
+        chat_options=DEFAULT_CHAT_OPTIONS,
+    )
     agent.chat_options.tools = [server_tool]
 
-    context = ExecutionContext(
+    context = TestExecutionContext(
         input_data=input_data,
         agent=agent,
         config=AgentConfig(),
     )
 
-    events = []
+    events: list[Any] = []
     async for event in orchestrator.run(context):
         events.append(event)
 
@@ -578,16 +584,18 @@ async def test_empty_messages_handling() -> None:
     """Test orchestrator handles empty message list gracefully."""
     orchestrator = DefaultOrchestrator()
 
-    input_data = {"messages": []}
+    input_data: dict[str, Any] = {"messages": []}
 
-    agent = MockAgent()
-    context = ExecutionContext(
+    agent = StubAgent(
+        chat_options=DEFAULT_CHAT_OPTIONS,
+    )
+    context = TestExecutionContext(
         input_data=input_data,
         agent=agent,
         config=AgentConfig(),
     )
 
-    events = []
+    events: list[Any] = []
     async for event in orchestrator.run(context):
         events.append(event)
 
@@ -603,7 +611,7 @@ async def test_all_messages_filtered_handling() -> None:
     """Test orchestrator handles case where all messages are filtered out."""
     orchestrator = DefaultOrchestrator()
 
-    input_data = {
+    input_data: dict[str, Any] = {
         "messages": [
             {
                 "role": "tool",
@@ -612,14 +620,16 @@ async def test_all_messages_filtered_handling() -> None:
         ]
     }
 
-    agent = MockAgent()
-    context = ExecutionContext(
+    agent = StubAgent(
+        chat_options=DEFAULT_CHAT_OPTIONS,
+    )
+    context = TestExecutionContext(
         input_data=input_data,
         agent=agent,
         config=AgentConfig(),
     )
 
-    events = []
+    events: list[Any] = []
     async for event in orchestrator.run(context):
         events.append(event)
 
@@ -651,16 +661,18 @@ async def test_confirm_changes_with_invalid_json_fallback() -> None:
     ]
 
     orchestrator = DefaultOrchestrator()
-    input_data = {"messages": []}
-    agent = MockAgent()
-    context = ExecutionContext(
+    input_data: dict[str, Any] = {"messages": []}
+    agent = StubAgent(
+        chat_options=DEFAULT_CHAT_OPTIONS,
+    )
+    context = TestExecutionContext(
         input_data=input_data,
         agent=agent,
         config=AgentConfig(),
     )
-    context._messages = messages
+    context.set_messages(messages)
 
-    events = []
+    events: list[Any] = []
     async for event in orchestrator.run(context):
         events.append(event)
 
@@ -689,16 +701,18 @@ async def test_tool_result_kept_when_call_id_matches() -> None:
     ]
 
     orchestrator = DefaultOrchestrator()
-    input_data = {"messages": []}
-    agent = MockAgent()
-    context = ExecutionContext(
+    input_data: dict[str, Any] = {"messages": []}
+    agent = StubAgent(
+        chat_options=DEFAULT_CHAT_OPTIONS,
+    )
+    context = TestExecutionContext(
         input_data=input_data,
         agent=agent,
         config=AgentConfig(),
     )
-    context._messages = messages
+    context.set_messages(messages)
 
-    events = []
+    events: list[Any] = []
     async for event in orchestrator.run(context):
         events.append(event)
 
@@ -729,6 +743,7 @@ async def test_agent_protocol_fallback_paths() -> None:
             *,
             thread: Any = None,
             tools: list[Any] | None = None,
+            **kwargs: Any,
         ) -> AsyncGenerator[AgentRunResponseUpdate, None]:
             self.messages_received = messages
             yield AgentRunResponseUpdate(contents=[TextContent(text="response")], role="assistant")
@@ -738,16 +753,16 @@ async def test_agent_protocol_fallback_paths() -> None:
     messages = [ChatMessage(role="user", contents=[TextContent(text="Hello")])]
 
     orchestrator = DefaultOrchestrator()
-    input_data = {"messages": []}
+    input_data: dict[str, Any] = {"messages": []}
     agent = CustomAgent()
-    context = ExecutionContext(
+    context = TestExecutionContext(
         input_data=input_data,
         agent=agent,  # type: ignore
         config=AgentConfig(),
     )
-    context._messages = messages
+    context.set_messages(messages)
 
-    events = []
+    events: list[Any] = []
     async for event in orchestrator.run(context):
         events.append(event)
 
@@ -762,21 +777,23 @@ async def test_initial_state_snapshot_with_array_schema() -> None:
     messages = [ChatMessage(role="user", contents=[TextContent(text="Hello")])]
 
     orchestrator = DefaultOrchestrator()
-    input_data = {"messages": [], "state": {}}
-    agent = MockAgent()
-    context = ExecutionContext(
+    input_data: dict[str, Any] = {"messages": [], "state": {}}
+    agent = StubAgent(
+        chat_options=DEFAULT_CHAT_OPTIONS,
+    )
+    context = TestExecutionContext(
         input_data=input_data,
         agent=agent,
         config=AgentConfig(state_schema={"items": {"type": "array"}}),
     )
-    context._messages = messages
+    context.set_messages(messages)
 
-    events = []
+    events: list[Any] = []
     async for event in orchestrator.run(context):
         events.append(event)
 
     # Should emit state snapshot with empty array for items
-    state_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
+    state_events: list[Any] = [e for e in events if e.type == "STATE_SNAPSHOT"]
     assert len(state_events) >= 1
 
 
@@ -791,19 +808,21 @@ async def test_response_format_skip_text_content() -> None:
     messages = [ChatMessage(role="user", contents=[TextContent(text="Hello")])]
 
     orchestrator = DefaultOrchestrator()
-    input_data = {"messages": []}
+    input_data: dict[str, Any] = {"messages": []}
 
-    agent = MockAgent()
+    agent = StubAgent(
+        chat_options=DEFAULT_CHAT_OPTIONS,
+    )
     agent.chat_options.response_format = OutputModel
 
-    context = ExecutionContext(
+    context = TestExecutionContext(
         input_data=input_data,
         agent=agent,
         config=AgentConfig(),
     )
-    context._messages = messages
+    context.set_messages(messages)
 
-    events = []
+    events: list[Any] = []
     async for event in orchestrator.run(context):
         events.append(event)
 

python/packages/ag-ui/tests/test_shared_state.py

@@ -2,6 +2,10 @@
 
 """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, TextContent
@@ -10,20 +14,16 @@ from agent_framework._types import ChatResponseUpdate
 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 test_helpers_ag_ui import StreamingChatClientStub, stream_from_updates
+
 
 @pytest.fixture
-def mock_agent():
+def mock_agent() -> ChatAgent:
     """Create a mock agent for testing."""
-
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            yield ChatResponseUpdate(contents=[TextContent(text="Hello!")])
-
-    return ChatAgent(
-        name="test_agent",
-        instructions="Test agent",
-        chat_client=MockChatClient(),
-    )
+    updates = [ChatResponseUpdate(contents=[TextContent(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():
@@ -65,9 +65,9 @@ def test_state_delta_event():
     assert event.delta[1]["op"] == "replace"
 
 
-async def test_agent_with_initial_state(mock_agent):
+async def test_agent_with_initial_state(mock_agent: ChatAgent) -> None:
     """Test agent emits state snapshot when initial state provided."""
-    state_schema = {"recipe": {"type": "object", "properties": {"name": {"type": "string"}}}}
+    state_schema: dict[str, Any] = {"recipe": {"type": "object", "properties": {"name": {"type": "string"}}}}
 
     agent = AgentFrameworkAgent(
         agent=mock_agent,
@@ -76,12 +76,12 @@ async def test_agent_with_initial_state(mock_agent):
 
     initial_state = {"recipe": {"name": "Test Recipe"}}
 
-    input_data = {
+    input_data: dict[str, Any] = {
         "messages": [{"role": "user", "content": "Hello"}],
         "state": initial_state,
     }
 
-    events = []
+    events: list[Any] = []
     async for event in agent.run_agent(input_data):
         events.append(event)
 
@@ -91,16 +91,16 @@ async def test_agent_with_initial_state(mock_agent):
     assert snapshot_events[0].snapshot == initial_state
 
 
-async def test_agent_without_state_schema(mock_agent):
+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 = {
+    input_data: dict[str, Any] = {
         "messages": [{"role": "user", "content": "Hello"}],
         "state": {"some": "state"},
     }
 
-    events = []
+    events: list[Any] = []
     async for event in agent.run_agent(input_data):
         events.append(event)
 

python/packages/ag-ui/tests/test_state_manager.py

@@ -0,0 +1,51 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from ag_ui.core import CustomEvent, EventType
+from agent_framework import ChatMessage, TextContent
+
+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 isinstance(message.contents[0], TextContent)
+    assert "Current state of the application" in message.contents[0].text

python/packages/ag-ui/tests/test_structured_output.py

@@ -3,12 +3,18 @@
 """Tests for structured output handling in _agent.py."""
 
 import json
+import sys
+from collections.abc import AsyncIterator, MutableSequence
+from pathlib import Path
 from typing import Any
 
-from agent_framework import ChatAgent, ChatOptions, TextContent
+from agent_framework import ChatAgent, ChatMessage, ChatOptions, TextContent
 from agent_framework._types import ChatResponseUpdate
 from pydantic import BaseModel
 
+sys.path.insert(0, str(Path(__file__).parent))
+from test_helpers_ag_ui import StreamingChatClientStub, stream_from_updates
+
 
 class RecipeOutput(BaseModel):
     """Test Pydantic model for recipe output."""
@@ -34,14 +40,14 @@ async def test_structured_output_with_recipe():
     """Test structured output processing with recipe state."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            # Simulate structured output
-            yield ChatResponseUpdate(
-                contents=[TextContent(text='{"recipe": {"name": "Pasta"}, "message": "Here is your recipe"}')]
-            )
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        yield ChatResponseUpdate(
+            contents=[TextContent(text='{"recipe": {"name": "Pasta"}, "message": "Here is your recipe"}')]
+        )
 
-    agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     agent.chat_options = ChatOptions(response_format=RecipeOutput)
 
     wrapper = AgentFrameworkAgent(
@@ -51,7 +57,7 @@ async def test_structured_output_with_recipe():
 
     input_data = {"messages": [{"role": "user", "content": "Make pasta"}]}
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
@@ -72,17 +78,18 @@ async def test_structured_output_with_steps():
     """Test structured output processing with steps state."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            steps_data = {
-                "steps": [
-                    {"id": "1", "description": "Step 1", "status": "pending"},
-                    {"id": "2", "description": "Step 2", "status": "pending"},
-                ]
-            }
-            yield ChatResponseUpdate(contents=[TextContent(text=json.dumps(steps_data))])
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        steps_data = {
+            "steps": [
+                {"id": "1", "description": "Step 1", "status": "pending"},
+                {"id": "2", "description": "Step 2", "status": "pending"},
+            ]
+        }
+        yield ChatResponseUpdate(contents=[TextContent(text=json.dumps(steps_data))])
 
-    agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     agent.chat_options = ChatOptions(response_format=StepsOutput)
 
     wrapper = AgentFrameworkAgent(
@@ -92,7 +99,7 @@ async def test_structured_output_with_steps():
 
     input_data = {"messages": [{"role": "user", "content": "Do steps"}]}
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
@@ -111,12 +118,13 @@ async def test_structured_output_with_no_schema_match():
     """Test structured output when response fields don't match state_schema keys."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            # Response has "data" field but schema expects "result" field
-            yield ChatResponseUpdate(contents=[TextContent(text='{"data": {"key": "value"}}')])
+    updates = [
+        ChatResponseUpdate(contents=[TextContent(text='{"data": {"key": "value"}}')]),
+    ]
 
-    agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(
+        name="test", instructions="Test", chat_client=StreamingChatClientStub(stream_from_updates(updates))
+    )
     agent.chat_options = ChatOptions(response_format=GenericOutput)
 
     wrapper = AgentFrameworkAgent(
@@ -126,7 +134,7 @@ async def test_structured_output_with_no_schema_match():
 
     input_data = {"messages": [{"role": "user", "content": "Generate data"}]}
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
@@ -146,11 +154,12 @@ async def test_structured_output_without_schema():
         data: dict[str, Any]
         info: str
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            yield ChatResponseUpdate(contents=[TextContent(text='{"data": {"key": "value"}, "info": "processed"}')])
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        yield ChatResponseUpdate(contents=[TextContent(text='{"data": {"key": "value"}, "info": "processed"}')])
 
-    agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     agent.chat_options = ChatOptions(response_format=DataOutput)
 
     wrapper = AgentFrameworkAgent(
@@ -160,7 +169,7 @@ async def test_structured_output_without_schema():
 
     input_data = {"messages": [{"role": "user", "content": "Generate data"}]}
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
@@ -177,18 +186,20 @@ async def test_no_structured_output_when_no_response_format():
     """Test that structured output path is skipped when no response_format."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            yield ChatResponseUpdate(contents=[TextContent(text="Regular text")])
+    updates = [ChatResponseUpdate(contents=[TextContent(text="Regular text")])]
 
-    agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(
+        name="test",
+        instructions="Test",
+        chat_client=StreamingChatClientStub(stream_from_updates(updates)),
+    )
     # No response_format set
 
     wrapper = AgentFrameworkAgent(agent=agent)
 
     input_data = {"messages": [{"role": "user", "content": "Hi"}]}
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
@@ -202,12 +213,13 @@ async def test_structured_output_with_message_field():
     """Test structured output that includes a message field."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            output_data = {"recipe": {"name": "Salad"}, "message": "Fresh salad recipe ready"}
-            yield ChatResponseUpdate(contents=[TextContent(text=json.dumps(output_data))])
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        output_data = {"recipe": {"name": "Salad"}, "message": "Fresh salad recipe ready"}
+        yield ChatResponseUpdate(contents=[TextContent(text=json.dumps(output_data))])
 
-    agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     agent.chat_options = ChatOptions(response_format=RecipeOutput)
 
     wrapper = AgentFrameworkAgent(
@@ -217,7 +229,7 @@ async def test_structured_output_with_message_field():
 
     input_data = {"messages": [{"role": "user", "content": "Make salad"}]}
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
@@ -236,20 +248,20 @@ async def test_empty_updates_no_structured_processing():
     """Test that empty updates don't trigger structured output processing."""
     from agent_framework.ag_ui import AgentFrameworkAgent
 
-    class MockChatClient:
-        async def get_streaming_response(self, messages, chat_options, **kwargs):
-            # Return nothing
-            if False:
-                yield
+    async def stream_fn(
+        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        if False:
+            yield ChatResponseUpdate(contents=[])
 
-    agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
+    agent = ChatAgent(name="test", instructions="Test", chat_client=StreamingChatClientStub(stream_fn))
     agent.chat_options = ChatOptions(response_format=RecipeOutput)
 
     wrapper = AgentFrameworkAgent(agent=agent)
 
     input_data = {"messages": [{"role": "user", "content": "Test"}]}
 
-    events = []
+    events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 

python/packages/ag-ui/tests/test_tooling.py

@@ -0,0 +1,36 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from types import SimpleNamespace
+
+from agent_framework_ag_ui._orchestration._tooling import merge_tools, register_additional_client_tools
+
+
+class DummyTool:
+    def __init__(self, name: str) -> None:
+        self.name = name
+        self.declaration_only = True
+
+
+def test_merge_tools_filters_duplicates() -> None:
+    server = [DummyTool("a"), DummyTool("b")]
+    client = [DummyTool("b"), DummyTool("c")]
+
+    merged = merge_tools(server, client)
+
+    assert merged is not None
+    names = [getattr(t, "name", None) for t in merged]
+    assert names == ["a", "b", "c"]
+
+
+def test_register_additional_client_tools_assigns_when_configured() -> None:
+    class Fic:
+        def __init__(self) -> None:
+            self.additional_tools = None
+
+    holder = SimpleNamespace(function_invocation_configuration=Fic())
+    agent = SimpleNamespace(chat_client=holder)
+
+    tools = [DummyTool("x")]
+    register_additional_client_tools(agent, tools)
+
+    assert holder.function_invocation_configuration.additional_tools == tools

python/packages/ag-ui/tests/test_utils.py

@@ -20,8 +20,8 @@ def test_generate_event_id():
 
 def test_merge_state():
     """Test state merging."""
-    current = {"a": 1, "b": 2}
-    update = {"b": 3, "c": 4}
+    current: dict[str, int] = {"a": 1, "b": 2}
+    update: dict[str, int] = {"b": 3, "c": 4}
 
     result = merge_state(current, update)
 
@@ -32,8 +32,8 @@ def test_merge_state():
 
 def test_merge_state_empty_update():
     """Test merging with empty update."""
-    current = {"x": 10, "y": 20}
-    update = {}
+    current: dict[str, int] = {"x": 10, "y": 20}
+    update: dict[str, int] = {}
 
     result = merge_state(current, update)
 
@@ -43,8 +43,8 @@ def test_merge_state_empty_update():
 
 def test_merge_state_empty_current():
     """Test merging with empty current state."""
-    current = {}
-    update = {"a": 1, "b": 2}
+    current: dict[str, int] = {}
+    update: dict[str, int] = {"a": 1, "b": 2}
 
     result = merge_state(current, update)
 
@@ -53,8 +53,8 @@ def test_merge_state_empty_current():
 
 def test_merge_state_deep_copy():
     """Test that merge_state creates a deep copy preventing mutation of original."""
-    current = {"recipe": {"name": "Cake", "ingredients": ["flour", "sugar"]}}
-    update = {"other": "value"}
+    current: dict[str, dict[str, object]] = {"recipe": {"name": "Cake", "ingredients": ["flour", "sugar"]}}
+    update: dict[str, str] = {"other": "value"}
 
     result = merge_state(current, update)