chore: import upstream snapshot with attribution

python/packages/a2a/AGENTS.md

@@ -0,0 +1,58 @@
+# A2A Package (agent-framework-a2a)
+
+Agent-to-Agent (A2A) protocol support for inter-agent communication.
+
+## Main Classes
+
+- **`A2AAgent`** - Client to connect to remote A2A-compliant agents.
+- **`A2AExecutor`** - Bridge to expose Agent Framework agents via the A2A protocol.
+- **`A2AServiceSessionId`** - Typed durable A2A continuation state shape stored in `AgentSession.service_session_id`.
+- **`A2AAgentSession`** - Deprecated compatibility session wrapper; prefer `AgentSession` + `A2AServiceSessionId`.
+
+## Usage
+
+### A2AAgent (Client)
+
+```python
+from agent_framework.a2a import A2AAgent
+
+# Connect to a remote A2A agent
+a2a_agent = A2AAgent(url="http://remote-agent/a2a")
+response = await a2a_agent.run("Hello!")
+```
+
+### A2AExecutor (Server/Bridge)
+
+```python
+from agent_framework.a2a import A2AExecutor
+from a2a.server.request_handlers import DefaultRequestHandler
+from a2a.server.routes import create_agent_card_routes, create_jsonrpc_routes
+from a2a.server.tasks import InMemoryTaskStore
+from starlette.applications import Starlette
+
+# Create an A2A executor for your agent
+executor = A2AExecutor(agent=my_agent)
+
+# Set up the request handler (agent_card is required)
+request_handler = DefaultRequestHandler(
+    agent_executor=executor,
+    task_store=InMemoryTaskStore(),
+    agent_card=my_agent_card,
+)
+
+# Build a Starlette app with A2A routes
+app = Starlette(
+    routes=[
+        *create_agent_card_routes(my_agent_card),
+        *create_jsonrpc_routes(request_handler, "/"),
+    ]
+)
+```
+
+## Import Path
+
+```python
+from agent_framework.a2a import A2AAgent, A2AExecutor, A2AServiceSessionId
+# or directly:
+from agent_framework_a2a import A2AAgent, A2AExecutor, A2AServiceSessionId
+```

python/packages/a2a/LICENSE

@@ -0,0 +1,21 @@
+    MIT License
+
+    Copyright (c) Microsoft Corporation.
+
+    Permission is hereby granted, free of charge, to any person obtaining a copy
+    of this software and associated documentation files (the "Software"), to deal
+    in the Software without restriction, including without limitation the rights
+    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+    copies of the Software, and to permit persons to whom the Software is
+    furnished to do so, subject to the following conditions:
+
+    The above copyright notice and this permission notice shall be included in all
+    copies or substantial portions of the Software.
+
+    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+    SOFTWARE

python/packages/a2a/README.md

@@ -0,0 +1,64 @@
+# Get Started with Microsoft Agent Framework A2A
+
+Please install this package via pip:
+
+```bash
+pip install agent-framework-a2a --pre
+```
+
+## A2A Agent Integration
+
+The A2A agent integration enables communication with remote A2A-compliant agents using the standardized A2A protocol. This allows your Agent Framework applications to connect to agents running on different platforms, languages, or services.
+
+### A2AAgent (Client)
+
+The `A2AAgent` class is a client that wraps an A2A Client to connect the Agent Framework with external A2A-compliant agents.
+
+```python
+from agent_framework.a2a import A2AAgent
+
+# Connect to a remote A2A agent
+a2a_agent = A2AAgent(url="http://remote-agent/a2a")
+response = await a2a_agent.run("Hello!")
+```
+
+### A2AExecutor (Hosting)
+
+The `A2AExecutor` class bridges local AI agents built with the `agent_framework` library to the A2A protocol, allowing them to be hosted and accessed by other A2A-compliant clients.
+
+```python
+from agent_framework.a2a import A2AExecutor
+from a2a.server.apps import A2AStarletteApplication
+from a2a.server.request_handlers import DefaultRequestHandler
+from a2a.server.tasks import InMemoryTaskStore
+
+# Create an A2A executor for your agent
+executor = A2AExecutor(agent=my_agent)
+
+# Set up the request handler and server application
+request_handler = DefaultRequestHandler(
+    agent_executor=executor,
+    task_store=InMemoryTaskStore(),
+)
+
+app = A2AStarletteApplication(
+    agent_card=my_agent_card,
+    http_handler=request_handler,
+).build()
+```
+
+### Basic Usage Example
+
+See the [A2A agent examples](../../samples/04-hosting/a2a/) which demonstrate:
+
+- Connecting to remote A2A agents
+- Hosting local agents via A2A protocol
+- Sending messages and receiving responses
+- Handling different content types (text, files, data)
+- Streaming responses and real-time interaction
+
+## Security considerations
+
+The hosting example above focuses on protocol wiring and does not add authentication or authorization by itself. Production A2A hosts should protect their HTTP or JSON-RPC entry points with the deployment's normal auth layer and verify that each caller is allowed to access the requested agent, task, or session.
+
+Task, thread, context, and session identifiers used by an A2A host are routing handles, not bearer credentials. Do not rely on client-supplied identifiers alone to select or mutate persisted state; bind them to authenticated user, tenant, or workspace context first.

python/packages/a2a/agent_framework_a2a/__init__.py

@@ -0,0 +1,20 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import importlib.metadata
+
+from ._a2a_executor import A2AExecutor
+from ._agent import A2AAgent, A2AAgentSession, A2AContinuationToken, A2AServiceSessionId
+
+try:
+    __version__ = importlib.metadata.version(__name__)
+except importlib.metadata.PackageNotFoundError:
+    __version__ = "0.0.0"  # Fallback for development mode
+
+__all__ = [
+    "A2AAgent",
+    "A2AAgentSession",
+    "A2AContinuationToken",
+    "A2AExecutor",
+    "A2AServiceSessionId",
+    "__version__",
+]

python/packages/a2a/agent_framework_a2a/_a2a_executor.py

@@ -0,0 +1,300 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import base64
+import logging
+import uuid
+from asyncio import CancelledError
+from collections.abc import Mapping
+from functools import partial
+from typing import Any
+
+from a2a.helpers import new_task_from_user_message
+from a2a.server.agent_execution import AgentExecutor, RequestContext
+from a2a.server.events import EventQueue
+from a2a.server.tasks import TaskUpdater
+from a2a.types import Part, TaskState
+from agent_framework import (
+    AgentResponseUpdate,
+    AgentSession,
+    Message,
+    SupportsAgentRun,
+)
+from typing_extensions import override
+
+from ._utils import get_uri_data
+
+logger = logging.getLogger("agent_framework.a2a")
+
+
+class A2AExecutor(AgentExecutor):
+    """Execute AI agents using the A2A (Agent-to-Agent) protocol.
+
+    The A2AExecutor bridges AI agents built with the agent_framework library and the A2A protocol,
+    enabling structured agent execution with event-driven communication. It handles execution
+    contexts, delegates history management to the agent's session, and converts agent
+    responses into A2A protocol events.
+
+    The executor supports executing an Agent or WorkflowAgent. It provides comprehensive
+    error handling with task status updates and supports various content types including text,
+    binary data, and URI-based content.
+
+    Example:
+        .. code-block:: python
+
+            from a2a.server.request_handlers import DefaultRequestHandler
+            from a2a.server.routes import create_jsonrpc_routes, create_agent_card_routes
+            from a2a.server.tasks import InMemoryTaskStore
+            from a2a.types import AgentCapabilities, AgentCard, AgentInterface
+            from agent_framework.a2a import A2AExecutor
+            from agent_framework.openai import OpenAIResponsesClient
+            from starlette.applications import Starlette
+
+            public_agent_card = AgentCard(
+                name="Food Agent",
+                description="A simple agent that provides food-related information.",
+                version="1.0.0",
+                default_input_modes=["text"],
+                default_output_modes=["text"],
+                capabilities=AgentCapabilities(streaming=True),
+                supported_interfaces=[
+                    AgentInterface(url="http://localhost:9999/", protocol_binding="JSONRPC"),
+                ],
+                skills=[],
+            )
+
+            # Create an agent
+            agent = OpenAIResponsesClient().as_agent(
+                name="Food Agent",
+                instructions="A simple agent that provides food-related information.",
+            )
+
+            # Set up the A2A server with the A2AExecutor enabled for streaming
+            # and passing custom keyword arguments to the agent's run method.
+            request_handler = DefaultRequestHandler(
+                agent_executor=A2AExecutor(agent, stream=True, run_kwargs={"client_kwargs": {"max_tokens": 500}}),
+                task_store=InMemoryTaskStore(),
+                agent_card=public_agent_card,
+            )
+
+            app = Starlette(
+                routes=[
+                    *create_agent_card_routes(public_agent_card),
+                    *create_jsonrpc_routes(request_handler, "/"),
+                ],
+            )
+
+    Args:
+        agent: The AI agent to execute.
+        stream: Whether to stream the agent response. Defaults to False.
+        run_kwargs: Additional keyword arguments to pass to the agent's run method.
+    """
+
+    def __init__(self, agent: SupportsAgentRun, stream: bool = False, run_kwargs: Mapping[str, Any] | None = None):
+        """Initialize the A2AExecutor with the specified agent.
+
+        Args:
+            agent: The AI agent or workflow to execute.
+            stream: Whether to stream the agent response. Defaults to False.
+            run_kwargs: Additional keyword arguments to pass to the agent's run method.
+                Cannot contain 'session' or 'stream' as these are managed by the executor.
+
+        Raises:
+            ValueError: If run_kwargs contains 'session' or 'stream'.
+        """
+        super().__init__()
+        self._agent: SupportsAgentRun = agent
+        self._stream: bool = stream
+        if run_kwargs:
+            if "session" in run_kwargs:
+                raise ValueError("run_kwargs cannot contain 'session' as it is managed by the executor.")
+            if "stream" in run_kwargs:
+                raise ValueError("run_kwargs cannot contain 'stream' as it is managed by the executor.")
+        self._run_kwargs: Mapping[str, Any] = run_kwargs or {}
+
+    @override
+    async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
+        """Cancel agent execution for the given request context.
+
+        Uses a TaskUpdater to send a cancellation event through the provided event queue.
+
+        Args:
+            context: The request context identifying the task to cancel.
+            event_queue: The event queue to publish the cancellation event to.
+
+        Raises:
+            ValueError: If context_id is not provided in the RequestContext.
+        """
+        if context.context_id is None:
+            raise ValueError("Context ID must be provided in the RequestContext")
+
+        updater = TaskUpdater(
+            event_queue=event_queue,
+            task_id=context.task_id or "",
+            context_id=context.context_id,
+        )
+
+        await updater.cancel()
+
+    @override
+    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
+        """Execute the agent with the given context and event queue.
+
+        Orchestrates the agent execution process: sets up the agent session,
+        executes the agent, processes response messages, and handles errors with appropriate task status updates.
+        """
+        if context.context_id is None:
+            raise ValueError("Context ID must be provided in the RequestContext")
+        if context.message is None:
+            raise ValueError("Message must be provided in the RequestContext")
+
+        query = context.get_user_input()
+        task = context.current_task
+
+        if not task:
+            task = new_task_from_user_message(context.message)
+            await event_queue.enqueue_event(task)
+
+        updater = TaskUpdater(event_queue, task.id, context.context_id)
+        await updater.submit()
+
+        try:
+            await updater.start_work()
+
+            session = self._agent.create_session(session_id=task.context_id)
+
+            if self._stream:
+                await self._run_stream(query, session, updater)
+            else:
+                await self._run(query, session, updater)
+
+            # Mark as complete
+            await updater.complete()
+        except CancelledError:
+            await updater.update_status(state=TaskState.TASK_STATE_CANCELED)
+        except Exception as e:
+            logger.exception("A2AExecutor encountered an error during execution.", exc_info=e)
+            await updater.update_status(
+                state=TaskState.TASK_STATE_FAILED,
+                message=updater.new_agent_message([Part(text=str(e))]),
+            )
+
+    async def _run_stream(self, query: Any, session: AgentSession, updater: TaskUpdater) -> None:
+        """Run the agent in streaming mode and publish updates to the task updater."""
+        response_stream = self._agent.run(query, session=session, stream=True, **self._run_kwargs)
+        streamed_artifact_ids: set[str] = set()
+        # Generate a stable artifact ID for the entire stream so all chunks share the same ID.
+        # This ensures clients can coalesce streaming tokens into a single artifact/message
+        # per the A2A spec (TaskArtifactUpdateEvent with append=True on same artifactId).
+        default_artifact_id = str(uuid.uuid4())
+        await (
+            response_stream.with_transform_hook(
+                partial(
+                    self.handle_events,
+                    updater=updater,
+                    streamed_artifact_ids=streamed_artifact_ids,
+                    default_artifact_id=default_artifact_id,
+                )
+            )
+        ).get_final_response()
+
+    async def _run(self, query: Any, session: AgentSession, updater: TaskUpdater) -> None:
+        """Run the agent in non-streaming mode and publish messages to the task updater."""
+        response = await self._agent.run(query, session=session, stream=False, **self._run_kwargs)
+        response_messages = response.messages
+
+        if not isinstance(response_messages, list):
+            response_messages = [response_messages]
+
+        for message in response_messages:
+            await self.handle_events(message, updater)
+
+    async def handle_events(
+        self,
+        item: Message | AgentResponseUpdate,
+        updater: TaskUpdater,
+        streamed_artifact_ids: set[str] | None = None,
+        default_artifact_id: str | None = None,
+    ) -> None:
+        """Convert agent response items (Messages or Updates) to A2A protocol events.
+
+        Processes Message or AgentResponseUpdate objects and converts them into A2A protocol format.
+        Handles text, data, and URI content. USER role messages are skipped.
+
+        Users can override this method in a subclass to implement custom transformations
+        from their agent's output format to A2A protocol events.
+
+        Args:
+            item: The agent response item (Message or AgentResponseUpdate) to process.
+            updater: The task updater to publish events to.
+            streamed_artifact_ids: A set of artifact IDs that have already been streamed.
+                Used to track which artifacts need append=True on subsequent chunks.
+            default_artifact_id: A stable artifact ID to use when the item does not provide one.
+                This ensures all streaming chunks for a single response share the same artifact ID,
+                allowing clients to coalesce them into a single message.
+
+        Example:
+            .. code-block:: python
+
+                class CustomA2AExecutor(A2AExecutor):
+                    async def handle_events(
+                        self,
+                        item: Message | AgentResponseUpdate,
+                        updater: TaskUpdater,
+                        streamed_artifact_ids: set[str] | None = None,
+                        default_artifact_id: str | None = None,
+                    ) -> None:
+                        # Custom logic to transform item contents
+                        if item.role == "assistant" and item.contents:
+                            parts = [Part(text=f"Custom: {item.contents[0].text}")]
+                            await updater.update_status(
+                                state=TaskState.TASK_STATE_WORKING,
+                                message=updater.new_agent_message(parts=parts),
+                            )
+                        else:
+                            await super().handle_events(item, updater)
+        """
+        role = getattr(item, "role", None)
+        if role == "user":
+            # This is a user message, we can ignore it in the context of task updates
+            return
+
+        parts: list[Part] = []
+        metadata = getattr(item, "additional_properties", None)
+
+        # AgentResponseUpdate uses 'contents', Message uses 'contents'
+        contents = getattr(item, "contents", [])
+
+        for content in contents:
+            if content.type == "text" and content.text:
+                parts.append(Part(text=content.text))
+            elif content.type == "data" and content.uri:
+                base64_str = get_uri_data(content.uri)
+                parts.append(Part(raw=base64.b64decode(base64_str), media_type=content.media_type or ""))
+            elif content.type == "uri" and content.uri:
+                parts.append(Part(url=content.uri, media_type=content.media_type or ""))
+            else:
+                # Silently skip unsupported content types
+                logger.warning("A2AExecutor does not yet support content type: %s. Omitted.", content.type)
+
+        if parts:
+            if isinstance(item, AgentResponseUpdate):
+                # Resolve artifact ID: use item's message_id if available, otherwise fall back
+                # to the stable default_artifact_id so all streaming chunks share the same ID.
+                artifact_id = item.message_id or default_artifact_id
+                # For streaming updates, we send TaskArtifactUpdateEvent via add_artifact
+                await updater.add_artifact(
+                    parts=parts,
+                    artifact_id=artifact_id,
+                    metadata=metadata,
+                    append=(
+                        True if streamed_artifact_ids is not None and artifact_id in streamed_artifact_ids else None
+                    ),
+                )
+                if artifact_id and streamed_artifact_ids is not None:
+                    streamed_artifact_ids.add(artifact_id)
+            else:
+                # For final messages, we send TaskStatusUpdateEvent with 'working' state
+                await updater.update_status(
+                    state=TaskState.TASK_STATE_WORKING,
+                    message=updater.new_agent_message(parts=parts, metadata=metadata),
+                )

python/packages/a2a/agent_framework_a2a/_utils.py

@@ -0,0 +1,24 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import re
+
+URI_PATTERN = re.compile(r"^data:(?P<media_type>[^;,]+(?:;[^;,=]+=[^;,]+)*);base64,(?P<base64_data>[A-Za-z0-9+/=]+)\Z")
+
+
+def get_uri_data(uri: str) -> str:
+    """Extracts the base64-encoded data from a data URI.
+
+    Args:
+        uri: The data URI to parse.
+
+    Returns:
+        The base64-encoded data part of the URI.
+
+    Raises:
+        ValueError: If the URI format is invalid.
+    """
+    match = URI_PATTERN.match(uri)
+    if not match:
+        raise ValueError(f"Invalid data URI format: {uri}")
+
+    return match.group("base64_data")

python/packages/a2a/pyproject.toml

@@ -0,0 +1,98 @@
+[project]
+name = "agent-framework-a2a"
+description = "A2A integration for Microsoft Agent Framework."
+authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
+readme = "README.md"
+requires-python = ">=3.10"
+version = "1.0.0b260709"
+license-files = ["LICENSE"]
+urls.homepage = "https://aka.ms/agent-framework"
+urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
+urls.release_notes = "https://github.com/microsoft/agent-framework/releases?q=tag%3Apython-1&expanded=true"
+urls.issues = "https://github.com/microsoft/agent-framework/issues"
+classifiers = [
+  "License :: OSI Approved :: MIT License",
+  "Development Status :: 4 - Beta",
+  "Intended Audience :: Developers",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+  "Programming Language :: Python :: 3.13",
+  "Programming Language :: Python :: 3.14",
+  "Typing :: Typed",
+]
+dependencies = [
+    "agent-framework-core>=1.11.0,<2",
+    "a2a-sdk>=1.0.0,<2",
+]
+
+[tool.uv]
+prerelease = "if-necessary-or-explicit"
+environments = [
+    "sys_platform == 'darwin'",
+    "sys_platform == 'linux'",
+    "sys_platform == 'win32'"
+]
+
+[tool.uv-dynamic-versioning]
+fallback-version = "0.0.0"
+
+[tool.pytest.ini_options]
+testpaths = 'tests'
+addopts = "-ra -q -r fEX"
+asyncio_mode = "auto"
+asyncio_default_fixture_loop_scope = "function"
+filterwarnings = [
+    "ignore:Support for class-based `config` is deprecated:DeprecationWarning:pydantic.*"
+]
+timeout = 120
+markers = [
+    "integration: marks tests as integration tests that require external services",
+]
+
+[tool.ruff]
+extend = "../../pyproject.toml"
+
+[tool.coverage.run]
+omit = [
+    "**/__init__.py"
+]
+
+[tool.pyright]
+extends = "../../pyproject.toml"
+include = ["agent_framework_a2a"]
+
+[tool.mypy]
+plugins = ['pydantic.mypy']
+strict = true
+python_version = "3.10"
+ignore_missing_imports = true
+disallow_untyped_defs = true
+no_implicit_optional = true
+check_untyped_defs = true
+warn_return_any = true
+show_error_codes = true
+warn_unused_ignores = false
+disallow_incomplete_defs = true
+disallow_untyped_decorators = true
+
+[tool.bandit]
+targets = ["agent_framework_a2a"]
+exclude_dirs = ["tests"]
+
+[tool.poe]
+executor.type = "uv"
+include = "../../shared_tasks.toml"
+
+[tool.poe.tasks.mypy]
+help = "Run MyPy for this package."
+cmd = "mypy --config-file $POE_ROOT/pyproject.toml agent_framework_a2a"
+
+[tool.poe.tasks.test]
+help = "Run the default unit test suite for this package."
+cmd = 'pytest -m "not integration" --cov=agent_framework_a2a --cov-report=term-missing:skip-covered tests'
+
+[build-system]
+requires = ["flit-core >= 3.11,<4.0"]
+build-backend = "flit_core.buildapi"

python/packages/a2a/tests/test_a2a_executor.py

@@ -0,0 +1,907 @@
+# Copyright (c) Microsoft. All rights reserved.
+from asyncio import CancelledError
+from typing import Any, cast
+from unittest.mock import AsyncMock, MagicMock, patch
+from uuid import uuid4
+
+from a2a.types import Part, Task, TaskState
+from agent_framework import (
+    AgentResponseUpdate,
+    Content,
+    Message,
+    SupportsAgentRun,
+)
+from agent_framework._types import AgentResponse
+from agent_framework.a2a import A2AExecutor
+from pytest import fixture, raises
+
+
+@fixture
+def mock_agent() -> MagicMock:
+    """Fixture that provides a mock SupportsAgentRun."""
+    agent = MagicMock(spec=SupportsAgentRun)
+    agent.run = AsyncMock()
+    return agent
+
+
+@fixture
+def mock_request_context() -> MagicMock:
+    """Fixture that provides a mock RequestContext."""
+    request_context = MagicMock()
+    request_context.context_id = str(uuid4())
+    request_context.get_user_input = MagicMock(return_value="Test query")
+    request_context.current_task = None
+    request_context.message = None
+    return request_context
+
+
+@fixture
+def mock_event_queue() -> MagicMock:
+    """Fixture that provides a mock EventQueue."""
+    queue = AsyncMock()
+    queue.enqueue_event = AsyncMock()
+    return queue
+
+
+@fixture
+def mock_task() -> Task:
+    """Fixture that provides a mock Task."""
+    task = MagicMock(spec=Task)
+    task.id = str(uuid4())
+    task.context_id = str(uuid4())
+    task.state = TaskState.TASK_STATE_COMPLETED
+    return task
+
+
+@fixture
+def mock_task_updater() -> MagicMock:
+    """Fixture that provides a mock TaskUpdater."""
+    updater = MagicMock()
+    updater.submit = AsyncMock()
+    updater.start_work = AsyncMock()
+    updater.complete = AsyncMock()
+    updater.update_status = AsyncMock()
+    updater.new_agent_message = MagicMock()
+    return updater
+
+
+@fixture
+def executor(mock_agent: MagicMock) -> A2AExecutor:
+    """Fixture that provides an A2AExecutor."""
+    return A2AExecutor(agent=mock_agent)
+
+
+class TestA2AExecutorInitialization:
+    """Tests for A2AExecutor initialization."""
+
+    def test_initialization_with_agent_only(self, mock_agent: MagicMock) -> None:
+        """Arrange: Create mock agent
+        Act: Initialize A2AExecutor with only agent
+        Assert: Executor is created with default values
+        """
+        # Act
+        executor = A2AExecutor(agent=mock_agent)
+
+        # Assert
+        assert executor._agent is mock_agent
+        assert executor._stream is False
+        assert executor._run_kwargs == {}
+
+    def test_initialization_with_stream_and_kwargs(self, mock_agent: MagicMock) -> None:
+        """Arrange: Create mock agent
+        Act: Initialize A2AExecutor with stream and run_kwargs
+        Assert: Executor is created with specified values
+        """
+        # Arrange
+        run_kwargs = {"temperature": 0.5}
+
+        # Act
+        executor = A2AExecutor(agent=mock_agent, stream=True, run_kwargs=run_kwargs)
+
+        # Assert
+        assert executor._agent is mock_agent
+        assert executor._stream is True
+        assert executor._run_kwargs == run_kwargs
+
+    def test_initialization_with_invalid_run_kwargs(self, mock_agent: MagicMock) -> None:
+        """Arrange: Create mock agent
+        Act: Initialize A2AExecutor with reserved keys in run_kwargs
+        Assert: ValueError is raised
+        """
+        # Act & Assert
+        with raises(ValueError, match="run_kwargs cannot contain 'session'"):
+            A2AExecutor(agent=mock_agent, run_kwargs={"session": "something"})
+
+        with raises(ValueError, match="run_kwargs cannot contain 'stream'"):
+            A2AExecutor(agent=mock_agent, run_kwargs={"stream": True})
+
+
+class TestA2AExecutorCancel:
+    """Tests for the cancel method."""
+
+    async def test_cancel_method_completes(
+        self,
+        executor: A2AExecutor,
+        mock_request_context: MagicMock,
+        mock_event_queue: MagicMock,
+    ) -> None:
+        """Arrange: Create executor with dependencies
+        Act: Call cancel method
+        Assert: Method completes without raising error
+        """
+        # Arrange
+        mock_request_context.task_id = "task-123"
+
+        # Act & Assert (should not raise)
+        await executor.cancel(mock_request_context, mock_event_queue)  # type: ignore
+
+    async def test_cancel_handles_different_contexts(
+        self,
+        executor: A2AExecutor,
+        mock_event_queue: MagicMock,
+    ) -> None:
+        """Arrange: Create executor with multiple request contexts
+        Act: Call cancel with different contexts
+        Assert: Each cancel completes successfully
+        """
+        # Arrange
+        context1 = MagicMock()
+        context1.context_id = "ctx-1"
+        context1.task_id = "task-1"
+        context2 = MagicMock()
+        context2.context_id = "ctx-2"
+        context2.task_id = "task-2"
+
+        # Act & Assert
+        await executor.cancel(context1, mock_event_queue)  # type: ignore
+        await executor.cancel(context2, mock_event_queue)  # type: ignore
+
+    async def test_cancel_raises_error_when_context_id_missing(
+        self,
+        executor: A2AExecutor,
+        mock_event_queue: MagicMock,
+    ) -> None:
+        """Arrange: Create context without context_id
+        Act: Call cancel method
+        Assert: ValueError is raised
+        """
+        # Arrange
+        mock_context = MagicMock()
+        mock_context.context_id = None
+
+        # Act & Assert
+        with raises(ValueError) as excinfo:
+            await executor.cancel(mock_context, mock_event_queue)  # type: ignore
+
+        # Assert
+        assert "Context ID" in str(excinfo.value)
+
+
+class TestA2AExecutorExecute:
+    """Tests for the execute method."""
+
+    async def test_execute_with_existing_task_succeeds(
+        self,
+        executor: A2AExecutor,
+        mock_request_context: MagicMock,
+        mock_event_queue: MagicMock,
+        mock_task: Task,
+    ) -> None:
+        """Arrange: Create executor with mocked dependencies and existing task
+        Act: Call execute method
+        Assert: Execution completes successfully
+        """
+        # Arrange
+        mock_request_context.get_user_input = MagicMock(return_value="Hello")
+        mock_request_context.current_task = mock_task
+        mock_request_context.context_id = "ctx-123"
+        mock_request_context.message = MagicMock()
+
+        response_message = Message(role="assistant", contents=[Content.from_text(text="Hello back")])
+        response = MagicMock(spec=AgentResponse)
+        response.messages = [response_message]
+        cast(Any, executor._agent).run = AsyncMock(return_value=response)
+        cast(Any, executor._agent).create_session = MagicMock()
+
+        with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
+            mock_updater = MagicMock()
+            mock_updater.submit = AsyncMock()
+            mock_updater.start_work = AsyncMock()
+            mock_updater.complete = AsyncMock()
+            mock_updater.update_status = AsyncMock()
+            mock_updater.new_agent_message = MagicMock(return_value="message_obj")
+            mock_updater_class.return_value = mock_updater
+
+            # Act
+            await executor.execute(mock_request_context, mock_event_queue)
+
+            # Assert
+            mock_updater.submit.assert_called_once()
+            mock_updater.start_work.assert_called_once()
+            mock_updater.complete.assert_called_once()
+            cast(Any, executor._agent.create_session).assert_called_once()
+            cast(Any, executor._agent.run).assert_called_once()
+
+    async def test_execute_creates_task_when_not_exists(
+        self,
+        executor: A2AExecutor,
+        mock_request_context: MagicMock,
+        mock_event_queue: MagicMock,
+    ) -> None:
+        """Arrange: Create executor with request context without task
+        Act: Call execute method
+        Assert: New task is created and enqueued
+        """
+        # Arrange
+        mock_message = MagicMock()
+        mock_request_context.get_user_input = MagicMock(return_value="Hello")
+        mock_request_context.current_task = None
+        mock_request_context.message = mock_message
+        mock_request_context.context_id = "ctx-123"
+
+        response_message = Message(role="assistant", contents=[Content.from_text(text="Response")])
+        response = MagicMock(spec=AgentResponse)
+        response.messages = [response_message]
+        cast(Any, executor._agent).run = AsyncMock(return_value=response)
+        cast(Any, executor._agent).create_session = MagicMock()
+
+        with patch("agent_framework_a2a._a2a_executor.new_task_from_user_message") as mock_new_task:
+            mock_task = MagicMock(spec=Task)
+            mock_task.id = "task-new"
+            mock_task.context_id = "ctx-123"
+            mock_new_task.return_value = mock_task
+
+            with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
+                mock_updater = MagicMock()
+                mock_updater.submit = AsyncMock()
+                mock_updater.start_work = AsyncMock()
+                mock_updater.complete = AsyncMock()
+                mock_updater.update_status = AsyncMock()
+                mock_updater.new_agent_message = MagicMock(return_value="message_obj")
+                mock_updater_class.return_value = mock_updater
+
+                # Act
+                await executor.execute(mock_request_context, mock_event_queue)
+
+                # Assert
+                mock_new_task.assert_called_once()
+                mock_event_queue.enqueue_event.assert_called_once()
+
+    async def test_execute_raises_error_when_context_id_missing(
+        self,
+        executor: A2AExecutor,
+        mock_request_context: MagicMock,
+        mock_event_queue: MagicMock,
+    ) -> None:
+        """Arrange: Create context without context_id
+        Act: Call execute method
+        Assert: ValueError is raised
+        """
+        # Arrange
+        mock_request_context.context_id = None
+        mock_request_context.message = MagicMock()
+
+        # Act & Assert
+        with raises(ValueError) as excinfo:
+            await executor.execute(mock_request_context, mock_event_queue)
+
+        # Assert
+        assert "Context ID" in str(excinfo.value)
+
+    async def test_execute_raises_error_when_message_missing(
+        self,
+        executor: A2AExecutor,
+        mock_request_context: MagicMock,
+        mock_event_queue: MagicMock,
+    ) -> None:
+        """Arrange: Create context without message
+        Act: Call execute method
+        Assert: ValueError is raised
+        """
+        # Arrange
+        mock_request_context.context_id = "ctx-123"
+        mock_request_context.message = None
+
+        # Act & Assert
+        with raises(ValueError) as excinfo:
+            await executor.execute(mock_request_context, mock_event_queue)
+
+        # Assert
+        assert "Message" in str(excinfo.value)
+
+    async def test_execute_handles_cancelled_error(
+        self,
+        executor: A2AExecutor,
+        mock_request_context: MagicMock,
+        mock_event_queue: MagicMock,
+        mock_task: Task,
+    ) -> None:
+        """Arrange: Create executor that raises CancelledError
+        Act: Call execute method
+        Assert: Error is caught and task is marked as canceled
+        """
+        # Arrange
+        mock_request_context.get_user_input = MagicMock(return_value="Hello")
+        mock_request_context.current_task = mock_task
+        mock_request_context.context_id = "ctx-123"
+        mock_request_context.message = MagicMock()
+
+        cast(Any, executor._agent).run = AsyncMock(side_effect=CancelledError())
+        cast(Any, executor._agent).create_session = MagicMock()
+
+        with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
+            mock_updater = MagicMock()
+            mock_updater.submit = AsyncMock()
+            mock_updater.start_work = AsyncMock()
+            mock_updater.update_status = AsyncMock()
+            mock_updater_class.return_value = mock_updater
+
+            # Act
+            await executor.execute(mock_request_context, mock_event_queue)  # type: ignore
+
+            # Assert
+            mock_updater.update_status.assert_called()
+            call_args_list = mock_updater.update_status.call_args_list
+            assert any(call[1].get("state") == TaskState.TASK_STATE_CANCELED for call in call_args_list)
+
+    async def test_execute_handles_generic_exception(
+        self,
+        executor: A2AExecutor,
+        mock_request_context: MagicMock,
+        mock_event_queue: MagicMock,
+        mock_task: Task,
+    ) -> None:
+        """Arrange: Create executor that raises generic exception
+        Act: Call execute method
+        Assert: Error is caught and task is marked as failed
+        """
+        # Arrange
+        mock_request_context.get_user_input = MagicMock(return_value="Hello")
+        mock_request_context.current_task = mock_task
+        mock_request_context.context_id = "ctx-123"
+        mock_request_context.message = MagicMock()
+
+        error_message = "Test error"
+        cast(Any, executor._agent).run = AsyncMock(side_effect=ValueError(error_message))
+        cast(Any, executor._agent).create_session = MagicMock()
+
+        with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
+            mock_updater = MagicMock()
+            mock_updater.submit = AsyncMock()
+            mock_updater.start_work = AsyncMock()
+            mock_updater.update_status = AsyncMock()
+            mock_updater.new_agent_message = MagicMock(return_value="error_message_obj")
+            mock_updater_class.return_value = mock_updater
+
+            # Act
+            await executor.execute(mock_request_context, mock_event_queue)
+
+            # Assert
+            mock_updater.new_agent_message.assert_called_once()
+            args, _ = mock_updater.new_agent_message.call_args
+            parts = args[0]
+            assert len(parts) == 1
+            assert isinstance(parts[0], Part)
+            assert parts[0].text == error_message
+
+            call_args_list = mock_updater.update_status.call_args_list
+            assert any(
+                call[1].get("state") == TaskState.TASK_STATE_FAILED and call[1].get("message") == "error_message_obj"
+                for call in call_args_list
+            )
+
+    async def test_execute_processes_multiple_response_messages(
+        self,
+        executor: A2AExecutor,
+        mock_request_context: MagicMock,
+        mock_event_queue: MagicMock,
+        mock_task: Task,
+    ) -> None:
+        """Arrange: Create executor that returns multiple response messages
+        Act: Call execute method
+        Assert: All messages are processed through handle_events
+        """
+        # Arrange
+        mock_request_context.get_user_input = MagicMock(return_value="Hello")
+        mock_request_context.current_task = mock_task
+        mock_request_context.context_id = "ctx-123"
+        mock_request_context.message = MagicMock()
+
+        response_message1 = Message(role="assistant", contents=[Content.from_text(text="First")])
+        response_message2 = Message(role="assistant", contents=[Content.from_text(text="Second")])
+        response = MagicMock(spec=AgentResponse)
+        response.messages = [response_message1, response_message2]
+        cast(Any, executor._agent).run = AsyncMock(return_value=response)
+        cast(Any, executor._agent).create_session = MagicMock()
+
+        # Mock handle_events
+        cast(Any, executor).handle_events = AsyncMock()
+
+        with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
+            mock_updater = MagicMock()
+            mock_updater.submit = AsyncMock()
+            mock_updater.start_work = AsyncMock()
+            mock_updater.complete = AsyncMock()
+            mock_updater.update_status = AsyncMock()
+            mock_updater_class.return_value = mock_updater
+
+            # Act
+            await executor.execute(mock_request_context, mock_event_queue)
+
+            # Assert
+            assert cast(Any, executor.handle_events).call_count == 2
+
+    async def test_execute_passes_query_to_run(
+        self,
+        executor: A2AExecutor,
+        mock_request_context: MagicMock,
+        mock_event_queue: MagicMock,
+        mock_task: Task,
+    ) -> None:
+        """Arrange: Create executor with request
+        Act: Call execute method
+        Assert: Query text is passed to run method with default stream and kwargs
+        """
+        # Arrange
+        query_text = "Hello agent"
+        mock_request_context.get_user_input = MagicMock(return_value=query_text)
+        mock_request_context.current_task = mock_task
+        mock_request_context.context_id = "ctx-123"
+        mock_request_context.message = MagicMock()
+
+        response_message = Message(role="assistant", contents=[Content.from_text(text="Response")])
+        response = MagicMock(spec=AgentResponse)
+        response.messages = [response_message]
+        cast(Any, executor._agent).run = AsyncMock(return_value=response)
+        cast(Any, executor._agent).create_session = MagicMock()
+
+        with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
+            mock_updater = MagicMock()
+            mock_updater.submit = AsyncMock()
+            mock_updater.start_work = AsyncMock()
+            mock_updater.complete = AsyncMock()
+            mock_updater.update_status = AsyncMock()
+            mock_updater.new_agent_message = MagicMock(return_value="message_obj")
+            mock_updater_class.return_value = mock_updater
+
+            # Act
+            await executor.execute(mock_request_context, mock_event_queue)
+
+            # Assert
+            cast(Any, executor._agent.run).assert_called_once_with(
+                query_text, session=executor._agent.create_session(), stream=False
+            )
+
+    async def test_execute_with_stream_enabled(
+        self,
+        mock_agent: MagicMock,
+        mock_request_context: MagicMock,
+        mock_event_queue: MagicMock,
+        mock_task: Task,
+    ) -> None:
+        """Arrange: Create executor with stream=True
+        Act: Call execute method
+        Assert: _run_stream is called and passes stream=True to run
+        """
+        # Arrange
+        executor = A2AExecutor(agent=mock_agent, stream=True)
+        query_text = "Hello agent"
+        mock_request_context.get_user_input = MagicMock(return_value=query_text)
+        mock_request_context.current_task = mock_task
+        mock_request_context.context_id = "ctx-123"
+        mock_request_context.message = MagicMock()
+
+        mock_response_stream = MagicMock()
+        mock_response_stream.with_transform_hook = MagicMock(return_value=mock_response_stream)
+        mock_response_stream.get_final_response = AsyncMock()
+        mock_agent.run = MagicMock(return_value=mock_response_stream)
+        mock_agent.create_session = MagicMock()
+
+        with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
+            mock_updater = MagicMock()
+            mock_updater.submit = AsyncMock()
+            mock_updater.start_work = AsyncMock()
+            mock_updater.complete = AsyncMock()
+            mock_updater.update_status = AsyncMock()
+            mock_updater_class.return_value = mock_updater
+
+            # Act
+            await executor.execute(mock_request_context, mock_event_queue)
+
+            # Assert
+            mock_agent.run.assert_called_once_with(query_text, session=mock_agent.create_session(), stream=True)
+            mock_response_stream.with_transform_hook.assert_called_once()
+            mock_response_stream.get_final_response.assert_called_once()
+
+    async def test_execute_with_run_kwargs(
+        self,
+        mock_agent: MagicMock,
+        mock_request_context: MagicMock,
+        mock_event_queue: MagicMock,
+        mock_task: Task,
+    ) -> None:
+        """Arrange: Create executor with run_kwargs
+        Act: Call execute method
+        Assert: run_kwargs are passed to run method
+        """
+        # Arrange
+        run_kwargs = {"temperature": 0.5, "max_tokens": 100}
+        executor = A2AExecutor(agent=mock_agent, run_kwargs=run_kwargs)
+        query_text = "Hello agent"
+        mock_request_context.get_user_input = MagicMock(return_value=query_text)
+        mock_request_context.current_task = mock_task
+        mock_request_context.context_id = "ctx-123"
+        mock_request_context.message = MagicMock()
+
+        response_message = Message(role="assistant", contents=[Content.from_text(text="Response")])
+        response = MagicMock(spec=AgentResponse)
+        response.messages = [response_message]
+        mock_agent.run = AsyncMock(return_value=response)
+        mock_agent.create_session = MagicMock()
+
+        with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
+            mock_updater = MagicMock()
+            mock_updater.submit = AsyncMock()
+            mock_updater.start_work = AsyncMock()
+            mock_updater.complete = AsyncMock()
+            mock_updater.update_status = AsyncMock()
+            mock_updater_class.return_value = mock_updater
+
+            # Act
+            await executor.execute(mock_request_context, mock_event_queue)
+
+            # Assert
+            mock_agent.run.assert_called_once_with(
+                query_text, session=mock_agent.create_session(), stream=False, **run_kwargs
+            )
+
+
+class TestA2AExecutorHandleEvents:
+    """Tests for A2AExecutor.handle_events method."""
+
+    async def test_run_method_with_single_message(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
+        """Test the private _run method with a single message (not a list)."""
+        # Arrange
+        query = "test query"
+        session = MagicMock()
+        response_message = Message(role="assistant", contents=[Content.from_text(text="Response")])
+        response = MagicMock(spec=AgentResponse)
+        response.messages = response_message  # Not a list
+        cast(Any, executor._agent).run = AsyncMock(return_value=response)
+        cast(Any, executor).handle_events = AsyncMock()
+
+        # Act
+        await executor._run(query, session, mock_updater)
+
+        # Assert
+        cast(Any, executor.handle_events).assert_called_once_with(response_message, mock_updater)
+
+    @fixture
+    def mock_updater(self) -> MagicMock:
+        """Create a mock execution context."""
+        updater = MagicMock()
+        updater.update_status = AsyncMock()
+        updater.new_agent_message = MagicMock(return_value="mock_message")
+        return updater
+
+    async def test_ignore_user_messages(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
+        """Test that messages from USER role are ignored."""
+        # Arrange
+        message = Message(
+            contents=[Content.from_text(text="User input")],
+            role="user",
+        )
+
+        # Act
+        await executor.handle_events(message, mock_updater)
+
+        # Assert
+        mock_updater.update_status.assert_not_called()
+
+    async def test_ignore_messages_with_no_contents(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
+        """Test that messages with no contents are ignored."""
+        # Arrange
+        message = Message(
+            contents=[],
+            role="assistant",
+        )
+
+        # Act
+        await executor.handle_events(message, mock_updater)
+
+        # Assert
+        mock_updater.update_status.assert_not_called()
+
+    async def test_handle_text_content(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
+        """Test handling messages with text content."""
+        # Arrange
+        text = "Hello, this is a test message"
+        message = Message(
+            contents=[Content.from_text(text=text)],
+            role="assistant",
+        )
+
+        # Act
+        await executor.handle_events(message, mock_updater)
+
+        # Assert
+        mock_updater.update_status.assert_called_once()
+        call_args = mock_updater.update_status.call_args
+        assert call_args.kwargs["state"] == TaskState.TASK_STATE_WORKING
+        assert mock_updater.new_agent_message.called
+
+    async def test_handle_multiple_text_contents(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
+        """Test handling messages with multiple text contents."""
+        # Arrange
+        message = Message(
+            contents=[
+                Content.from_text(text="First message"),
+                Content.from_text(text="Second message"),
+            ],
+            role="assistant",
+        )
+
+        # Act
+        await executor.handle_events(message, mock_updater)
+
+        # Assert
+        mock_updater.update_status.assert_called_once()
+        assert mock_updater.new_agent_message.called
+
+    async def test_handle_data_content(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
+        """Test handling messages with data content."""
+        # Arrange
+        data = b"test file data"
+        message = Message(
+            contents=[Content.from_data(data=data, media_type="application/octet-stream")],
+            role="assistant",
+        )
+
+        # Act
+        await executor.handle_events(message, mock_updater)
+
+        # Assert
+        mock_updater.update_status.assert_called_once()
+        call_args = mock_updater.update_status.call_args
+        assert call_args.kwargs["state"] == TaskState.TASK_STATE_WORKING
+
+    async def test_handle_uri_content(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
+        """Test handling messages with URI content."""
+        # Arrange
+        uri = "https://example.com/file.pdf"
+        message = Message(
+            contents=[Content.from_uri(uri=uri, media_type="application/pdf")],
+            role="assistant",
+        )
+
+        # Act
+        await executor.handle_events(message, mock_updater)
+
+        # Assert
+        mock_updater.update_status.assert_called_once()
+        call_args = mock_updater.update_status.call_args
+        assert call_args.kwargs["state"] == TaskState.TASK_STATE_WORKING
+
+    async def test_handle_mixed_content_types(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
+        """Test handling messages with mixed content types."""
+        # Arrange
+        data = b"file data"
+
+        message = Message(
+            contents=[
+                Content.from_text(text="Processing file..."),
+                Content.from_data(data=data, media_type="application/octet-stream"),
+                Content.from_uri(uri="https://example.com/reference.pdf", media_type="application/pdf"),
+            ],
+            role="assistant",
+        )
+
+        # Act
+        await executor.handle_events(message, mock_updater)
+
+        # Assert
+        mock_updater.update_status.assert_called_once()
+        call_args = mock_updater.update_status.call_args
+        assert call_args.kwargs["state"] == TaskState.TASK_STATE_WORKING
+
+    async def test_handle_with_additional_properties(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
+        """Test handling messages with additional properties metadata."""
+        # Arrange
+        additional_props = {"custom_field": "custom_value", "priority": "high"}
+        message = Message(
+            contents=[Content.from_text(text="Test message")],
+            role="assistant",
+            additional_properties=additional_props,
+        )
+
+        # Act
+        await executor.handle_events(message, mock_updater)
+
+        # Assert
+        mock_updater.update_status.assert_called_once()
+        mock_updater.new_agent_message.assert_called_once()
+        call_args = mock_updater.new_agent_message.call_args
+        assert call_args.kwargs["metadata"] == additional_props
+
+    async def test_handle_with_no_additional_properties(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
+        """Test handling messages without additional properties."""
+        # Arrange
+        message = Message(
+            contents=[Content.from_text(text="Test message")],
+            role="assistant",
+            additional_properties=None,
+        )
+
+        # Act
+        await executor.handle_events(message, mock_updater)
+
+        # Assert
+        mock_updater.update_status.assert_called_once()
+        mock_updater.new_agent_message.assert_called_once()
+        call_args = mock_updater.new_agent_message.call_args
+        assert call_args.kwargs["metadata"] == {}
+
+    async def test_parts_list_passed_to_new_agent_message(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
+        """Test that parts list is correctly passed to new_agent_message."""
+        # Arrange
+        message = Message(
+            contents=[
+                Content.from_text(text="Message 1"),
+                Content.from_text(text="Message 2"),
+            ],
+            role="assistant",
+        )
+
+        # Act
+        await executor.handle_events(message, mock_updater)
+
+        # Assert
+        mock_updater.new_agent_message.assert_called_once()
+        call_kwargs = mock_updater.new_agent_message.call_args.kwargs
+        assert "parts" in call_kwargs
+        parts_list = call_kwargs["parts"]
+        assert len(parts_list) == 2
+
+    async def test_task_state_always_working(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
+        """Test that task state is always set to working."""
+        # Arrange
+        message = Message(
+            contents=[Content.from_text(text="Any message")],
+            role="assistant",
+        )
+
+        # Act
+        await executor.handle_events(message, mock_updater)
+
+        # Assert
+        call_kwargs = mock_updater.update_status.call_args.kwargs
+        assert call_kwargs["state"] == TaskState.TASK_STATE_WORKING
+
+    async def test_handle_agent_response_update_no_streamed_set(
+        self, executor: A2AExecutor, mock_updater: MagicMock
+    ) -> None:
+        """Test handling AgentResponseUpdate (streaming) without a tracking set."""
+        # Arrange
+        update = AgentResponseUpdate(
+            contents=[Content.from_text(text="Streaming chunk")],
+            role="assistant",
+            message_id="msg-1",
+        )
+        mock_updater.add_artifact = AsyncMock()
+
+        # Act
+        await executor.handle_events(update, mock_updater)
+
+        # Assert
+        mock_updater.add_artifact.assert_called_once()
+        call_kwargs = mock_updater.add_artifact.call_args.kwargs
+        assert call_kwargs["artifact_id"] == "msg-1"
+        assert call_kwargs["append"] is None
+
+    async def test_handle_agent_response_update_first_time(
+        self, executor: A2AExecutor, mock_updater: MagicMock
+    ) -> None:
+        """Test handling AgentResponseUpdate (streaming) for the first time with a tracking set."""
+        # Arrange
+        update = AgentResponseUpdate(
+            contents=[Content.from_text(text="Streaming chunk")],
+            role="assistant",
+            message_id="msg-1",
+        )
+        mock_updater.add_artifact = AsyncMock()
+        streamed_artifact_ids: set[str] = set()
+
+        # Act
+        await executor.handle_events(update, mock_updater, streamed_artifact_ids=streamed_artifact_ids)
+
+        # Assert
+        mock_updater.add_artifact.assert_called_once()
+        call_kwargs = mock_updater.add_artifact.call_args.kwargs
+        assert call_kwargs["append"] is None
+        assert "msg-1" in streamed_artifact_ids
+
+    async def test_handle_agent_response_update_subsequent_time(
+        self, executor: A2AExecutor, mock_updater: MagicMock
+    ) -> None:
+        """Test handling AgentResponseUpdate (streaming) for subsequent times with a tracking set."""
+        # Arrange
+        update = AgentResponseUpdate(
+            contents=[Content.from_text(text="Next chunk")],
+            role="assistant",
+            message_id="msg-1",
+        )
+        mock_updater.add_artifact = AsyncMock()
+        streamed_artifact_ids = {"msg-1"}
+
+        # Act
+        await executor.handle_events(update, mock_updater, streamed_artifact_ids=streamed_artifact_ids)
+
+        # Assert
+        mock_updater.add_artifact.assert_called_once()
+        call_kwargs = mock_updater.add_artifact.call_args.kwargs
+        assert call_kwargs["append"] is True
+
+    async def test_handle_unsupported_content_type(self, executor: A2AExecutor, mock_updater: MagicMock) -> None:
+        """Test handling messages with unsupported content types."""
+        # Arrange
+        message = Message(
+            contents=[Content(type=cast(Any, "unknown"), text="Some text")],  # type: ignore[arg-type]
+            role="assistant",
+        )
+
+        # Act
+        with patch("agent_framework_a2a._a2a_executor.logger") as mock_logger:
+            await executor.handle_events(message, mock_updater)
+
+        # Assert
+        mock_logger.warning.assert_called_once()
+        mock_updater.update_status.assert_not_called()
+
+
+class TestA2AExecutorIntegration:
+    """Integration tests for A2AExecutor."""
+
+    async def test_full_execution_flow_with_responses(
+        self,
+        executor: A2AExecutor,
+        mock_request_context: MagicMock,
+        mock_event_queue: MagicMock,
+        mock_task: Task,
+    ) -> None:
+        """Arrange: Create executor with all mocked dependencies
+        Act: Execute full flow from request to completion
+        Assert: All components interact correctly
+        """
+        # Arrange
+        mock_request_context.get_user_input = MagicMock(return_value="Hello agent")
+        mock_request_context.current_task = mock_task
+        mock_request_context.context_id = "ctx-123"
+        mock_request_context.message = MagicMock()
+
+        response = MagicMock(spec=AgentResponse)
+        response_message = MagicMock(spec=Message)
+        response.messages = [response_message]
+        response_message.contents = [Content.from_text(text="Hello user")]
+        response_message.role = "assistant"
+        response_message.additional_properties = None
+
+        cast(Any, executor._agent).run = AsyncMock(return_value=response)
+        cast(Any, executor._agent).create_session = MagicMock()
+        cast(Any, executor).handle_events = AsyncMock()
+
+        with patch("agent_framework_a2a._a2a_executor.TaskUpdater") as mock_updater_class:
+            mock_updater = MagicMock()
+            mock_updater.submit = AsyncMock()
+            mock_updater.start_work = AsyncMock()
+            mock_updater.complete = AsyncMock()
+            mock_updater.update_status = AsyncMock()
+            mock_updater_class.return_value = mock_updater
+
+            # Act
+            await executor.execute(mock_request_context, mock_event_queue)
+
+            # Assert
+            mock_updater.submit.assert_called_once()
+            mock_updater.start_work.assert_called_once()
+            cast(Any, executor.handle_events).assert_called_once()
+            mock_updater.complete.assert_called_once()

python/packages/a2a/tests/test_utils.py

@@ -0,0 +1,51 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import pytest
+
+from agent_framework_a2a._utils import get_uri_data
+
+
+def test_get_uri_data_valid() -> None:
+    """Test get_uri_data with valid data URIs."""
+    # Simple text/plain
+    uri = "data:text/plain;base64,SGVsbG8sIFdvcmxkIQ=="
+    assert get_uri_data(uri) == "SGVsbG8sIFdvcmxkIQ=="
+
+    # Image png
+    uri = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
+    assert get_uri_data(uri) == "iVBORw0KGgoAAAANSUhEUgfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
+
+    # Application octet-stream
+    uri = "data:application/octet-stream;base64,AQIDBA=="
+    assert get_uri_data(uri) == "AQIDBA=="
+
+    # Media type with parameters
+    uri = "data:text/plain;charset=utf-8;base64,SGVsbG8sIFdvcmxkIQ=="
+    assert get_uri_data(uri) == "SGVsbG8sIFdvcmxkIQ=="
+
+    # Media type with multiple parameters
+    uri = "data:text/plain;charset=utf-8;name=hello.txt;base64,SGVsbG8sIFdvcmxkIQ=="
+    assert get_uri_data(uri) == "SGVsbG8sIFdvcmxkIQ=="
+
+
+def test_get_uri_data_invalid_format() -> None:
+    """Test get_uri_data with invalid URI formats."""
+    invalid_uris = [
+        "not-a-uri",
+        "http://example.com",
+        "data:text/plain;SGVsbG8sIFdvcmxkIQ==",  # Missing base64 marker
+        "data:base64,SGVsbG8sIFdvcmxkIQ==",  # Missing media type
+        "data:text/plain;foo;base64,SGVsbG8sIFdvcmxkIQ==",  # Parameter without value
+        "data:text/plain;base64;base64,SGVsbG8sIFdvcmxkIQ==",  # base64 used as a parameter name
+        "data:text/plain;base64,SGVsbG8sIFdvcmxkIQ== extra",
+        "data:text/plain;base64,SGVsbG8sIFdvcmxkIQ==\n",
+    ]
+    for uri in invalid_uris:
+        with pytest.raises(ValueError, match="Invalid data URI format"):
+            get_uri_data(uri)
+
+
+def test_get_uri_data_empty() -> None:
+    """Test get_uri_data with empty string."""
+    with pytest.raises(ValueError, match="Invalid data URI format"):
+        get_uri_data("")

python/packages/ag-ui/AGENTS.md

@@ -0,0 +1,51 @@
+# AG-UI Package (agent-framework-ag-ui)
+
+AG-UI protocol integration for building agent UIs with the AG-UI standard.
+
+## Main Classes
+
+- **`AgentFrameworkAgent`** - Wraps agents for AG-UI compatibility
+- **`AgentFrameworkWorkflow`** - Wraps native `Workflow` objects, or accepts `workflow_factory(thread_id)` for thread-scoped workflow instances without subclassing
+- **`AGUIChatClient`** - Chat client that speaks AG-UI protocol
+- **`AGUIHttpService`** - HTTP service for AG-UI endpoints
+- **`AGUIEventConverter`** - Converts between Agent Framework and AG-UI events
+- **`add_agent_framework_fastapi_endpoint()`** - Add AG-UI endpoint to FastAPI app (`SupportsAgentRun` or `Workflow`)
+- **`InMemoryAGUIThreadSnapshotStore`** - Memory-only latest AG-UI Thread Snapshot store for local development, demos, and tests
+
+## Types
+
+- **`AGUIRequest`** / **`AGUIChatOptions`** - Request types
+- **`AGUIThreadSnapshot`** / **`AGUIThreadSnapshotStore`** - Replayable thread snapshot model and scoped async store protocol
+- **`availableInterrupts` / `resume`** - Optional canonical AG-UI `Interrupt` and `ResumeEntry` protocol data
+- **`AgentState`** / **`RunMetadata`** - State management types
+- **`PredictStateConfig`** - Configuration for state prediction
+
+## Protocol Notes
+
+- Outbound custom events are emitted as AG-UI `CUSTOM`.
+- Usage metadata from `Content(type="usage")` is surfaced as `CUSTOM` events with `name="usage"`.
+- Inbound custom event aliases are accepted: `CUSTOM`, `CUSTOM_EVENT`, and `custom_event`.
+- Multimodal user inputs support both legacy (`text`, `binary`) and draft-style (`image`, `audio`, `video`, `document`) shapes.
+- Interrupted runs complete with `RUN_FINISHED.outcome.type == "interrupt"` and canonical `outcome.interrupts`; do not document or add new flows that depend on the legacy top-level `RUN_FINISHED.interrupt` field.
+- `Interrupt` and `ResumeEntry` come from the `ag-ui-protocol` package (`ag_ui.core`), not from an Agent Framework-specific interrupt model.
+- SSE keepalive is endpoint-owned transport behavior configured through
+  `add_agent_framework_fastapi_endpoint(keepalive_seconds=...)`. It emits SSE comments only; do not add `PING`,
+  `HEARTBEAT`, or `KEEPALIVE` AG-UI events, and do not add runner-level keepalive settings.
+
+## Usage
+
+```python
+from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
+from fastapi import FastAPI
+
+app = FastAPI()
+add_agent_framework_fastapi_endpoint(app, agent)
+```
+
+## Import Path
+
+```python
+from agent_framework.ag_ui import AGUIChatClient, add_agent_framework_fastapi_endpoint
+# or directly:
+from agent_framework_ag_ui import AGUIChatClient
+```

python/packages/ag-ui/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) Microsoft Corporation.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

python/packages/ag-ui/README.md

@@ -0,0 +1,378 @@
+# Agent Framework AG-UI Integration
+
+AG-UI protocol integration for Agent Framework, enabling seamless integration with AG-UI's web interface and streaming protocol.
+
+## Installation
+
+```bash
+pip install agent-framework-ag-ui
+```
+
+## Quick Start
+
+### Server (Host an AI Agent)
+
+```python
+from fastapi import FastAPI
+from agent_framework import Agent
+from agent_framework.openai import OpenAIChatCompletionClient
+from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
+
+# Create your agent
+agent = Agent(
+    name="my_agent",
+    instructions="You are a helpful assistant.",
+    client=OpenAIChatCompletionClient(
+        azure_endpoint="https://your-resource.openai.azure.com/",
+        model="gpt-4o-mini",
+        api_key="your-api-key",
+    ),
+)
+
+# Create FastAPI app and add AG-UI endpoint
+app = FastAPI()
+add_agent_framework_fastapi_endpoint(app, agent, "/")
+
+# Run with: uvicorn main:app --reload
+```
+
+### Server (Host a Workflow)
+
+```python
+from fastapi import FastAPI
+from agent_framework import WorkflowBuilder, WorkflowContext, executor
+from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
+
+@executor(id="start")
+async def start(message: str, ctx: WorkflowContext) -> None:
+    await ctx.yield_output(f"Workflow received: {message}")
+
+workflow = WorkflowBuilder(start_executor=start).build()
+
+app = FastAPI()
+add_agent_framework_fastapi_endpoint(app, workflow, "/")
+```
+
+### Server (Thread-Scoped WorkflowBuilder)
+
+Use `workflow_factory` when your workflow keeps runtime state (for example pending `request_info` interrupts) and must be isolated per AG-UI thread:
+
+```python
+from fastapi import FastAPI
+from agent_framework import Workflow, WorkflowBuilder
+from agent_framework.ag_ui import AgentFrameworkWorkflow, add_agent_framework_fastapi_endpoint
+
+def build_workflow_for_thread(thread_id: str) -> Workflow:
+    # Build a fresh workflow instance for each thread id.
+    return WorkflowBuilder(start_executor=...).build()
+
+app = FastAPI()
+thread_scoped_workflow = AgentFrameworkWorkflow(
+    workflow_factory=build_workflow_for_thread,
+    name="my_workflow",
+)
+add_agent_framework_fastapi_endpoint(app, thread_scoped_workflow, "/")
+```
+
+### Client (Connect to an AG-UI Server)
+
+```python
+import asyncio
+from agent_framework.ag_ui import AGUIChatClient
+
+async def main():
+    async with AGUIChatClient(endpoint="http://localhost:8000/") as client:
+        # Stream responses
+        async for update in client.get_response("Hello!", stream=True):
+            for content in update.contents:
+                if content.type == "text" and content.text:
+                    print(content.text, end="", flush=True)
+        print()
+
+asyncio.run(main())
+```
+
+The `AGUIChatClient` supports:
+- Streaming and non-streaming responses
+- Hybrid tool execution (client-side + server-side tools)
+- Automatic thread management for conversation continuity
+- Integration with `Agent` for client-side history management
+- Canonical interrupt/resume passthrough (`availableInterrupts` and `resume`)
+
+## Tool Return Helpers
+
+Use `state_update` when a backend tool needs to send different payloads to the model, the UI, and shared state. The `text` value remains the LLM-bound tool result, `tool_result` becomes the AG-UI `ToolCallResultEvent.content` for frontend rendering, and `state` is merged into durable shared state.
+
+```python
+from agent_framework import Content, tool
+from agent_framework.ag_ui import state_update
+
+@tool
+async def get_weather(city: str) -> Content:
+    data = await fetch_weather(city)
+    return state_update(
+        text=f"{city}: {data['temp']}°C and {data['conditions']}",
+        tool_result={
+            "component": "weather-card",
+            "city": city,
+            "temperature": data["temp"],
+            "conditions": data["conditions"],
+            "humidity": data["humidity"],
+        },
+        state={"weather": {"city": city, **data}},
+    )
+```
+
+## Documentation
+
+- **[Getting Started Tutorial](getting_started/)** - Step-by-step guide to building AG-UI servers and clients
+  - Server setup with FastAPI
+  - Client examples using `AGUIChatClient`
+  - Hybrid tool execution (client-side + server-side)
+  - Thread management and conversation continuity
+- **[Examples](agent_framework_ag_ui_examples/)** - Complete examples for AG-UI features
+
+## Interrupts and Resume
+
+Agent Framework AG-UI uses the canonical AG-UI interrupt protocol. Paused agent approval and workflow
+`request_info` runs finish with `RUN_FINISHED.outcome.type == "interrupt"` and a non-empty
+`RUN_FINISHED.outcome.interrupts` array. Agent Framework does not define a separate interrupt model; use
+`ag_ui.core.Interrupt` and `ag_ui.core.ResumeEntry` when constructing typed request data in Python.
+
+Tool approval interrupts use `reason: "tool_call"` and include `toolCallId` when the pause is bound to a tool call.
+Workflow `request_info` interrupts use `reason: "input_required"`. Framework-specific details needed for resume
+validation live in each interrupt's `metadata`, while generic clients can render the human-readable `message` and
+`responseSchema`.
+
+Interrupted terminal event shape:
+
+```json
+{
+  "type": "RUN_FINISHED",
+  "outcome": {
+    "type": "interrupt",
+    "interrupts": [
+      {
+        "id": "approval_1",
+        "reason": "tool_call",
+        "message": "Approve tool call get_weather?",
+        "toolCallId": "tool_call_1",
+        "responseSchema": {
+          "type": "object",
+          "properties": {
+            "accepted": { "type": "boolean" },
+            "arguments": { "type": "object" }
+          },
+          "required": ["accepted"]
+        },
+        "metadata": {
+          "agent_framework": {
+            "type": "function_approval_request",
+            "function_call": {
+              "call_id": "tool_call_1",
+              "name": "get_weather",
+              "arguments": {
+                "city": "Seattle"
+              }
+            }
+          }
+        }
+      }
+    ]
+  }
+}
+```
+
+Resume the paused thread with a canonical `resume` array. Each entry addresses exactly one open interrupt by
+`interruptId`; `status` is `resolved` or `cancelled`; resolved entries carry the approval or workflow response payload.
+
+```json
+{
+  "threadId": "thread-1",
+  "messages": [],
+  "resume": [
+    {
+      "interruptId": "approval_1",
+      "status": "resolved",
+      "payload": {
+        "approved": true
+      }
+    }
+  ]
+}
+```
+
+This is a clean release-candidate breaking change before `1.0.0`: new interrupted runs use
+`RUN_FINISHED.outcome.interrupts` and do not emit a stable top-level `RUN_FINISHED.interrupt` field. Normal
+non-interrupted runs continue to finish with valid `RUN_FINISHED` terminal events.
+
+## Public API Review Notes
+
+The Python package is currently in release candidate stage and is targeting the released `1.0.0` API surface. The preferred application import path is `agent_framework.ag_ui`; direct package imports from `agent_framework_ag_ui` are also supported.
+
+Review focus: whether these names are the right stable contract for Python users, and whether the protocol interrupt fields below match AG-UI's expected pause/resume shape.
+
+| Surface | Public exports |
+| --- | --- |
+| `agent_framework.ag_ui` facade | `AgentFrameworkAgent`, `AgentFrameworkWorkflow`, `AGUIChatClient`, `AGUIEventConverter`, `AGUIHttpService`, `AGUIThreadSnapshot`, `AGUIThreadSnapshotStore`, `InMemoryAGUIThreadSnapshotStore`, `SnapshotScopeResolver`, `add_agent_framework_fastapi_endpoint`, `state_update`, `__version__` |
+| Direct `agent_framework_ag_ui` package | Facade exports plus `AGUIChatOptions`, `AGUIRequest`, `AGUIThreadID`, `AgentState`, `DEFAULT_MAX_THREAD_SNAPSHOTS`, `DEFAULT_TAGS`, `PredictStateConfig`, `RunMetadata`, `SnapshotScope`, `WorkflowFactory` |
+| AG-UI protocol package (`ag_ui.core`) | `Interrupt`, `ResumeEntry`, `RunFinishedInterruptOutcome`, and related run outcome models |
+
+Interrupt support is protocol data rather than a separate Agent Framework Python class. Requests accept canonical `availableInterrupts`/`available_interrupts` and `resume` values; `AGUIChatClient` and `AGUIHttpService.post_run(...)` forward those fields with AG-UI wire aliases; agent approval and workflow `request_info` pauses emit `RUN_FINISHED.outcome.interrupts`; `AGUIEventConverter` preserves canonical interrupt outcome metadata on the final `ChatResponseUpdate`; and thread snapshot hydration replays the canonical interrupt outcome when a scoped snapshot stores an unresolved pause.
+
+## Features
+
+This integration supports all 7 AG-UI features:
+
+1. **Agentic Chat**: Basic streaming chat with tool calling support
+2. **Backend Tool Rendering**: Tools executed on backend with results streamed to client
+3. **Human in the Loop**: Function approval requests for user confirmation before tool execution
+4. **Agentic Generative UI**: Async tools for long-running operations with progress updates
+5. **Tool-based Generative UI**: Custom UI components rendered on frontend based on tool calls
+6. **Shared State**: Bidirectional state sync between client and server
+7. **Predictive State Updates**: Stream tool arguments as optimistic state updates during execution
+
+Additional compatibility and draft support:
+- Native `Workflow` endpoint registration via `add_agent_framework_fastapi_endpoint(...)`
+- Workflow-to-AG-UI event mapping (run/step/activity/tool/custom events)
+- Custom event compatibility for inbound `CUSTOM`, `CUSTOM_EVENT`, and `custom_event`
+- Pragmatic multimodal input parsing for both legacy (`binary`) and draft media-part shapes
+- Canonical interrupt/resume handling (`availableInterrupts`, `resume`, and `RUN_FINISHED.outcome.interrupts`)
+
+## Security: Authentication & Authorization
+
+The AG-UI endpoint does not enforce authentication by default. **For production deployments, you should add authentication** using FastAPI's dependency injection system via the `dependencies` parameter.
+
+### API Key Authentication Example
+
+```python
+import os
+from fastapi import Depends, FastAPI, HTTPException, Security
+from fastapi.security import APIKeyHeader
+from agent_framework import Agent
+from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
+
+# Configure API key authentication
+API_KEY_HEADER = APIKeyHeader(name="X-API-Key", auto_error=False)
+EXPECTED_API_KEY = os.environ.get("AG_UI_API_KEY")
+
+async def verify_api_key(api_key: str | None = Security(API_KEY_HEADER)) -> None:
+    """Verify the API key provided in the request header."""
+    if not api_key or api_key != EXPECTED_API_KEY:
+        raise HTTPException(status_code=401, detail="Invalid or missing API key")
+
+# Create agent and app
+agent = Agent(name="my_agent", instructions="...", client=...)
+app = FastAPI()
+
+# Register endpoint WITH authentication
+add_agent_framework_fastapi_endpoint(
+    app,
+    agent,
+    "/",
+    dependencies=[Depends(verify_api_key)],  # Authentication enforced here
+)
+```
+
+### Other Authentication Options
+
+The `dependencies` parameter accepts any FastAPI dependency, enabling integration with:
+
+- **OAuth 2.0 / OpenID Connect** - Use `fastapi.security.OAuth2PasswordBearer`
+- **JWT Tokens** - Validate tokens with libraries like `python-jose`
+- **Azure AD / Entra ID** - Use `azure-identity` for Microsoft identity platform
+- **Rate Limiting** - Add request throttling dependencies
+- **Custom Authentication** - Implement your organization's auth requirements
+
+For a complete authentication example, see [getting_started/server.py](getting_started/server.py).
+
+## AG-UI Thread Snapshots
+
+AG-UI Thread Snapshot persistence is opt-in and disabled by default. Existing endpoints keep their current behavior
+unless you provide a `snapshot_store`.
+
+Thread snapshots let an AG-UI frontend recover replayable UI state after a refresh. When snapshot persistence is
+enabled, the endpoint stores the latest replayable snapshot for an AG-UI Thread within an application-defined
+Snapshot Scope. A Hydrate Request is an AG-UI request with a known `threadId`, `messages: []`, and no `resume`
+payload. Hydration replays the stored Shared State, message snapshot, and canonical interrupt outcome when available,
+then finishes without invoking the wrapped agent or workflow.
+
+Use the built-in in-memory store for local development, demos, and tests:
+
+```python
+from fastapi import FastAPI
+
+from agent_framework.ag_ui import InMemoryAGUIThreadSnapshotStore, add_agent_framework_fastapi_endpoint
+
+app = FastAPI()
+agent = ...
+snapshot_store = InMemoryAGUIThreadSnapshotStore(max_snapshots=500)
+
+
+def resolve_snapshot_scope(request):
+    # Local demo scope. Production apps should derive the scope from authenticated user or tenant context.
+    del request
+    return "local-demo"
+
+
+add_agent_framework_fastapi_endpoint(
+    app,
+    agent,
+    "/",
+    snapshot_store=snapshot_store,
+    snapshot_scope_resolver=resolve_snapshot_scope,
+)
+```
+
+A frontend can then hydrate the latest stored snapshot for the scoped thread:
+
+```json
+{
+  "threadId": "thread-1",
+  "messages": []
+}
+```
+
+Endpoint configuration requires `snapshot_scope_resolver` whenever a snapshot store is configured, including when
+the store is already set on a pre-wrapped `AgentFrameworkAgent` or `AgentFrameworkWorkflow`. The resolver returns
+the application-defined Snapshot Scope used with the AG-UI Thread id as the storage key.
+
+AG-UI Thread ids identify AG-UI Threads; they do not authorize snapshot access. Do not treat a thread id as a bearer
+credential or tenant boundary. Production applications must authenticate and authorize every AG-UI endpoint request
+and choose a Snapshot Scope that represents the app's real access boundary, such as an authenticated user, tenant,
+or workspace. Do not rely on untrusted client-provided fields by themselves to choose that boundary.
+
+Tool approval resumes are validated against server-owned Approval State. The default Approval State store is
+process-local and bounded, and stores only approval-specific state needed to validate and continue pending approvals.
+It is not an authentication, tenant authorization, or distributed durability mechanism; production applications remain
+responsible for endpoint authentication, tenant authorization, and deployment/storage architecture that matches their
+availability and worker topology requirements.
+
+Stored snapshots are untrusted application data with confidentiality impact. They may contain sensitive user text,
+model output, tool results, function arguments, UI payloads, Shared State, and interrupt data. The built-in
+`InMemoryAGUIThreadSnapshotStore` is in-memory only, process-local, bounded, latest-only, and not durable production
+storage. It is cleared on process restart and is not shared across workers.
+
+No file-backed AG-UI snapshot store is provided by the package. Applications that need durable persistence should
+provide an app-owned implementation of the `AGUIThreadSnapshotStore` protocol and own storage hardening, including
+encryption, access control, retention, audit, data residency, and deletion behavior.
+
+## Architecture
+
+The package uses a clean, orchestrator-based architecture:
+
+- **AgentFrameworkAgent**: Lightweight wrapper that delegates to orchestrators
+- **Orchestrators**: Handle different execution flows (default, human-in-the-loop, etc.)
+- **Confirmation Strategies**: Domain-specific confirmation messages (extensible)
+- **AgentFrameworkEventBridge**: Converts Agent Framework events to AG-UI events
+- **Message Adapters**: Bidirectional conversion between AG-UI and Agent Framework message formats
+- **FastAPI Endpoint**: Streaming HTTP endpoint with Server-Sent Events (SSE)
+
+## Next Steps
+
+1. **New to AG-UI?** Start with the [Getting Started Tutorial](getting_started/)
+2. **Want to see examples?** Check out the [Examples](agent_framework_ag_ui_examples/) for AG-UI features
+
+## License
+
+MIT

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

@@ -0,0 +1,56 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""AG-UI protocol integration for Agent Framework."""
+
+import importlib.metadata
+
+from ._agent import AgentFrameworkAgent
+from ._client import AGUIChatClient
+from ._endpoint import add_agent_framework_fastapi_endpoint
+from ._event_converters import AGUIEventConverter
+from ._http_service import AGUIHttpService
+from ._snapshots import (
+    DEFAULT_MAX_THREAD_SNAPSHOTS,
+    AGUIThreadID,
+    AGUIThreadSnapshot,
+    AGUIThreadSnapshotStore,
+    InMemoryAGUIThreadSnapshotStore,
+    SnapshotScope,
+    SnapshotScopeResolver,
+)
+from ._state import state_update
+from ._types import AgentState, AGUIChatOptions, AGUIRequest, PredictStateConfig, RunMetadata
+from ._workflow import AgentFrameworkWorkflow, WorkflowFactory
+
+try:
+    __version__ = importlib.metadata.version(__name__)
+except importlib.metadata.PackageNotFoundError:
+    __version__ = "0.0.0"
+
+# Default OpenAPI tags for AG-UI endpoints
+DEFAULT_TAGS = ["AG-UI"]
+
+__all__ = [
+    "AgentFrameworkAgent",
+    "AgentFrameworkWorkflow",
+    "WorkflowFactory",
+    "add_agent_framework_fastapi_endpoint",
+    "AGUIChatClient",
+    "AGUIChatOptions",
+    "AGUIEventConverter",
+    "AGUIHttpService",
+    "AGUIRequest",
+    "AGUIThreadID",
+    "AGUIThreadSnapshot",
+    "AGUIThreadSnapshotStore",
+    "AgentState",
+    "InMemoryAGUIThreadSnapshotStore",
+    "PredictStateConfig",
+    "RunMetadata",
+    "SnapshotScope",
+    "SnapshotScopeResolver",
+    "DEFAULT_MAX_THREAD_SNAPSHOTS",
+    "DEFAULT_TAGS",
+    "state_update",
+    "__version__",
+]

python/packages/ag-ui/agent_framework_ag_ui/_agent.py

@@ -0,0 +1,146 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""AgentFrameworkAgent wrapper for AG-UI protocol."""
+
+from collections.abc import AsyncGenerator
+from typing import Any, cast
+
+from ag_ui.core import BaseEvent
+from agent_framework import SupportsAgentRun
+
+from ._agent_run import PendingApprovalEntry, PendingApprovalKey, run_agent_stream
+from ._approval_state import InMemoryAGUIApprovalStateStore
+from ._snapshots import AGUIThreadSnapshotStore
+
+
+class AgentConfig:
+    """Configuration for agent wrapper."""
+
+    def __init__(
+        self,
+        state_schema: Any | None = None,
+        predict_state_config: dict[str, dict[str, str]] | None = None,
+        use_service_session: bool = False,
+        require_confirmation: bool = True,
+        snapshot_store: AGUIThreadSnapshotStore | None = None,
+    ):
+        """Initialize agent configuration.
+
+        Args:
+            state_schema: Optional state schema for state management; accepts dict or Pydantic model/class
+            predict_state_config: Configuration for predictive state updates
+            use_service_session: Whether the agent session is service-managed
+            require_confirmation: Whether predictive updates require user confirmation before applying
+            snapshot_store: Optional AG-UI Thread Snapshot store. Snapshot persistence remains inactive unless
+                endpoint setup also provides an explicit Snapshot Scope resolver.
+        """
+        self.state_schema = self._normalize_state_schema(state_schema)
+        self.predict_state_config = predict_state_config or {}
+        self.use_service_session = use_service_session
+        self.require_confirmation = require_confirmation
+        self.snapshot_store = snapshot_store
+
+    @staticmethod
+    def _normalize_state_schema(state_schema: Any | None) -> dict[str, Any]:
+        """Accept dict or Pydantic model/class and return a properties dict."""
+        if state_schema is None:
+            return {}
+
+        if isinstance(state_schema, dict):
+            return cast(dict[str, Any], state_schema)
+
+        base_model_type: type[Any] | None
+        try:
+            from pydantic import BaseModel as ImportedBaseModel
+
+            base_model_type = ImportedBaseModel
+        except Exception:  # pragma: no cover
+            base_model_type = None
+
+        if base_model_type is not None and isinstance(state_schema, base_model_type):
+            schema_dict = state_schema.__class__.model_json_schema()
+            return schema_dict.get("properties", {}) or {}
+
+        if base_model_type is not None and isinstance(state_schema, type) and issubclass(state_schema, base_model_type):
+            schema_dict = state_schema.model_json_schema()  # type: ignore[union-attr]
+            return schema_dict.get("properties", {}) or {}  # type: ignore
+
+        return {}
+
+
+class AgentFrameworkAgent:
+    """Wraps Agent Framework agents for AG-UI protocol compatibility.
+
+    Translates between Agent Framework's SupportsAgentRun and AG-UI's event-based
+    protocol. Follows a simple linear flow: RunStarted -> content events -> RunFinished.
+    """
+
+    def __init__(
+        self,
+        agent: SupportsAgentRun,
+        name: str | None = None,
+        description: str | None = None,
+        state_schema: Any | None = None,
+        predict_state_config: dict[str, dict[str, str]] | None = None,
+        require_confirmation: bool = True,
+        use_service_session: bool = False,
+        snapshot_store: AGUIThreadSnapshotStore | None = None,
+    ):
+        """Initialize the AG-UI compatible agent wrapper.
+
+        Args:
+            agent: The Agent Framework agent to wrap
+            name: Optional name for the agent
+            description: Optional description
+            state_schema: Optional state schema for state management; accepts dict or Pydantic model/class
+            predict_state_config: Configuration for predictive state updates
+            require_confirmation: Whether predictive updates require user confirmation before applying
+            use_service_session: Whether the agent session is service-managed
+            snapshot_store: Optional AG-UI Thread Snapshot store. Snapshot persistence remains inactive unless
+                endpoint setup also provides an explicit Snapshot Scope resolver.
+        """
+        self.agent = agent
+        self.name = name or getattr(agent, "name", "agent")
+        self.description = description or getattr(agent, "description", "")
+
+        self.config = AgentConfig(
+            state_schema=state_schema,
+            predict_state_config=predict_state_config,
+            use_service_session=use_service_session,
+            require_confirmation=require_confirmation,
+            snapshot_store=snapshot_store,
+        )
+
+        # Server-side Approval State. Populated when approval requests are emitted
+        # and consumed when resume decisions arrive.
+        self._approval_state_store = InMemoryAGUIApprovalStateStore()
+        self._pending_approvals = cast(
+            dict[PendingApprovalKey, PendingApprovalEntry],
+            self._approval_state_store.pending_approvals,
+        )
+
+    @property
+    def snapshot_store(self) -> AGUIThreadSnapshotStore | None:
+        """Configured AG-UI Thread Snapshot store, if any."""
+        return self.config.snapshot_store
+
+    async def run(
+        self,
+        input_data: dict[str, Any],
+    ) -> AsyncGenerator[BaseEvent, None]:
+        """Run the wrapped agent and yield AG-UI events.
+
+        Args:
+            input_data: The AG-UI run input containing messages, state, etc.
+
+        Yields:
+            AG-UI events
+        """
+        async for event in run_agent_stream(
+            input_data,
+            self.agent,
+            self.config,
+            pending_approvals=self._pending_approvals,
+            approval_state_store=self._approval_state_store,
+        ):
+            yield event

python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py

@@ -0,0 +1,59 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Server-side AG-UI approval state storage."""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+from typing import Any
+
+ApprovalScope = str
+"""Application-defined scope for server-side AG-UI Approval State."""
+
+DEFAULT_MAX_APPROVAL_STATES = 10_000
+_APPROVAL_SCOPE_INPUT_KEY = "__ag_ui_approval_scope"
+_APPROVAL_THREAD_SEPARATOR = "\x1f"
+
+
+def approval_state_thread_id(*, scope: object | None, thread_id: str) -> str:
+    """Return the storage thread key for Approval State.
+
+    ``None`` is the only unscoped value. A provided scope must be a non-empty
+    string so accidental empty or malformed scopes cannot collapse into the
+    unscoped namespace.
+    """
+    if scope is None:
+        return thread_id
+    if not isinstance(scope, str) or not scope:
+        raise ValueError("scope must be a non-empty string when provided.")
+    return f"{scope}{_APPROVAL_THREAD_SEPARATOR}{thread_id}"
+
+
+class InMemoryAGUIApprovalStateStore:
+    """Bounded process-local server-side store for AG-UI Approval State.
+
+    The default store keeps only pending approval entries. It does not store
+    general ``AgentSession.state`` or AG-UI Thread Snapshots.
+    """
+
+    def __init__(self, *, max_entries: int = DEFAULT_MAX_APPROVAL_STATES) -> None:
+        """Initialize the process-local Approval State store.
+
+        Keyword Args:
+            max_entries: Maximum pending approval entries to retain.
+
+        Raises:
+            ValueError: If ``max_entries`` is less than 1.
+        """
+        if max_entries < 1:
+            raise ValueError("max_entries must be greater than 0.")
+        self.max_entries = max_entries
+        self.pending_approvals: OrderedDict[tuple[str, str], Any] = OrderedDict()
+        self.tool_approval_states: OrderedDict[str, dict[str, Any]] = OrderedDict()
+
+    def evict_oldest(self) -> None:
+        """Evict oldest pending approval entries until the store is within bounds."""
+        while len(self.pending_approvals) > self.max_entries:
+            self.pending_approvals.popitem(last=False)
+        while len(self.tool_approval_states) > self.max_entries:
+            self.tool_approval_states.popitem(last=False)

python/packages/ag-ui/agent_framework_ag_ui/_client.py

@@ -0,0 +1,468 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""AG-UI Chat Client implementation."""
+
+from __future__ import annotations
+
+import json
+import logging
+import sys
+import uuid
+from collections.abc import AsyncIterable, Awaitable, Mapping, MutableSequence, Sequence
+from functools import wraps
+from typing import TYPE_CHECKING, Any, Generic, TypedDict, cast
+
+import httpx
+from agent_framework import (
+    BaseChatClient,
+    ChatResponse,
+    ChatResponseUpdate,
+    Content,
+    FunctionTool,
+    Message,
+    ResponseStream,
+)
+from agent_framework._middleware import ChatMiddlewareLayer
+from agent_framework._tools import FunctionInvocationConfiguration, FunctionInvocationLayer
+from agent_framework.observability import ChatTelemetryLayer
+
+from ._event_converters import AGUIEventConverter
+from ._http_service import AGUIHttpService, _serialize_available_interrupts, _serialize_resume
+from ._message_adapters import agent_framework_messages_to_agui
+from ._utils import convert_tools_to_agui_format
+
+if sys.version_info >= (3, 13):
+    from typing import TypeVar  # pragma: no cover
+else:
+    from typing_extensions import TypeVar  # pragma: no cover
+if sys.version_info >= (3, 12):
+    from typing import override  # pragma: no cover
+else:
+    from typing_extensions import override  # pragma: no cover
+if sys.version_info >= (3, 11):
+    from typing import Self, TypedDict  # pragma: no cover
+else:
+    from typing_extensions import Self, TypedDict  # pragma: no cover
+
+if TYPE_CHECKING:
+    from agent_framework._middleware import ChatAndFunctionMiddlewareTypes
+
+    from ._types import AGUIChatOptions
+
+logger: logging.Logger = logging.getLogger("agent_framework.ag_ui")
+
+
+def _unwrap_server_function_call_contents(contents: MutableSequence[Content | dict[str, Any]]) -> None:
+    """Replace server_function_call instances with their underlying call content."""
+    for idx, content in enumerate(contents):
+        if content.type == "server_function_call":  # type: ignore[union-attr]
+            contents[idx] = content.function_call  # type: ignore[assignment, union-attr]
+
+
+BaseChatClientT = TypeVar("BaseChatClientT", bound=type[BaseChatClient[Any]])
+
+AGUIChatOptionsT = TypeVar(
+    "AGUIChatOptionsT",
+    bound=TypedDict,  # type: ignore[valid-type]
+    default="AGUIChatOptions",
+    covariant=True,
+)
+
+
+def _apply_server_function_call_unwrap(client: BaseChatClientT) -> BaseChatClientT:
+    """Class decorator that unwraps server-side function calls after tool handling."""
+
+    original_get_response = client.get_response
+
+    @wraps(original_get_response)
+    def response_wrapper(
+        self, *args: Any, stream: bool = False, **kwargs: Any
+    ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
+        if stream:
+            stream_response = original_get_response(self, *args, stream=True, **kwargs)
+            if isinstance(stream_response, ResponseStream):
+                return stream_response.with_transform_hook(_map_update)
+            return ResponseStream(_stream_wrapper_impl(stream_response))
+        return _response_wrapper_impl(self, original_get_response, *args, **kwargs)
+
+    async def _response_wrapper_impl(self, original_func: Any, *args: Any, **kwargs: Any) -> ChatResponse:
+        """Non-streaming wrapper implementation."""
+        response = await original_func(self, *args, stream=False, **kwargs)
+        if response.messages:
+            for message in response.messages:
+                _unwrap_server_function_call_contents(cast(MutableSequence[Content | dict[str, Any]], message.contents))
+        return response
+
+    async def _stream_wrapper_impl(stream: Any) -> AsyncIterable[ChatResponseUpdate]:
+        """Streaming wrapper implementation."""
+        if isinstance(stream, Awaitable):
+            stream = await stream
+        async for update in stream:
+            _unwrap_server_function_call_contents(cast(MutableSequence[Content | dict[str, Any]], update.contents))
+            yield update
+
+    def _map_update(update: ChatResponseUpdate) -> ChatResponseUpdate:
+        _unwrap_server_function_call_contents(cast(MutableSequence[Content | dict[str, Any]], update.contents))
+        return update
+
+    client.get_response = response_wrapper  # type: ignore[assignment]
+    return client
+
+
+@_apply_server_function_call_unwrap
+class AGUIChatClient(
+    FunctionInvocationLayer[AGUIChatOptionsT],
+    ChatMiddlewareLayer[AGUIChatOptionsT],
+    ChatTelemetryLayer[AGUIChatOptionsT],
+    BaseChatClient[AGUIChatOptionsT],
+    Generic[AGUIChatOptionsT],
+):
+    """Chat client for communicating with AG-UI compliant servers.
+
+    This client implements the BaseChatClient interface and automatically handles:
+    - Thread ID management for conversation continuity
+    - State synchronization between client and server
+    - Server-Sent Events (SSE) streaming
+    - Event conversion to Agent Framework types
+    - MiddlewareTypes, telemetry, and function invocation support
+
+    Important: Message History Management
+        This client sends exactly the messages it receives to the server. It does NOT
+        automatically maintain conversation history. The server must handle history via thread_id.
+
+        For stateless servers: Use Agent wrapper which will send full message history on each
+        request. However, even with Agent, the server must echo back all context for the
+        agent to maintain history across turns.
+
+    Important: Tool Handling (Hybrid Execution - matches .NET)
+        1. Client tool metadata sent to server - LLM knows about both client and server tools
+        2. Server has its own tools that execute server-side
+        3. When LLM calls a client tool, function invocation executes it locally
+        4. Both client and server tools work together (hybrid pattern)
+
+        The wrapping Agent's function invocation handles client tool execution
+        automatically when the server's LLM decides to call them.
+
+    Examples:
+        Direct usage (server manages thread history):
+
+        .. code-block:: python
+
+            from agent_framework.ag_ui import AGUIChatClient
+
+            client = AGUIChatClient(endpoint="http://localhost:8888/")
+
+            # First message - thread ID auto-generated
+            response = await client.get_response("Hello!")
+            thread_id = response.additional_properties.get("thread_id")
+
+            # Second message - server retrieves history using thread_id
+            response2 = await client.get_response(
+                "How are you?",
+                metadata={"thread_id": thread_id}
+            )
+
+        Recommended usage with Agent (client manages history):
+
+        .. code-block:: python
+
+            from agent_framework import Agent
+            from agent_framework.ag_ui import AGUIChatClient
+
+            client = AGUIChatClient(endpoint="http://localhost:8888/")
+            agent = Agent(name="assistant", client=client)
+            session = agent.create_session()
+
+            # Agent automatically maintains history and sends full context
+            response = await agent.run("Hello!", session=session)
+            response2 = await agent.run("How are you?", session=session)
+
+        Streaming usage:
+
+        .. code-block:: python
+
+            async for update in client.get_response("Tell me a story", stream=True):
+                if update.contents:
+                    for content in update.contents:
+                        if hasattr(content, "text"):
+                            print(content.text, end="", flush=True)
+
+        Context manager:
+
+        .. code-block:: python
+
+            async with AGUIChatClient(endpoint="http://localhost:8888/") as client:
+                response = await client.get_response("Hello!")
+                print(response.messages[0].text)
+
+        Using custom ChatOptions with type safety:
+
+        .. code-block:: python
+
+            from typing import TypedDict
+            from agent_framework_ag_ui import AGUIChatClient, AGUIChatOptions
+
+            class MyOptions(AGUIChatOptions, total=False):
+                my_custom_option: str
+
+            client: AGUIChatClient[MyOptions] = AGUIChatClient(endpoint="http://localhost:8888/")
+            response = await client.get_response("Hello", options={"my_custom_option": "value"})
+    """
+
+    OTEL_PROVIDER_NAME = "agui"
+
+    def __init__(
+        self,
+        *,
+        endpoint: str,
+        http_client: httpx.AsyncClient | None = None,
+        timeout: float = 60.0,
+        additional_properties: dict[str, Any] | None = None,
+        middleware: Sequence[ChatAndFunctionMiddlewareTypes] | None = None,
+        function_invocation_configuration: FunctionInvocationConfiguration | None = None,
+    ) -> None:
+        """Initialize the AG-UI chat client.
+
+        Args:
+            endpoint: The AG-UI server endpoint URL (e.g., "http://localhost:8888/")
+            http_client: Optional httpx.AsyncClient instance. If None, one will be created.
+            timeout: Request timeout in seconds (default: 60.0)
+            additional_properties: Additional properties to store
+            middleware: Optional middleware to apply to the client.
+            function_invocation_configuration: Optional function invocation configuration override.
+        """
+        super().__init__(
+            additional_properties=additional_properties,
+            middleware=middleware,
+            function_invocation_configuration=function_invocation_configuration,
+        )
+        self._http_service = AGUIHttpService(
+            endpoint=endpoint,
+            http_client=http_client,
+            timeout=timeout,
+        )
+
+    async def close(self) -> None:
+        """Close the HTTP client."""
+        await self._http_service.close()
+
+    async def __aenter__(self) -> Self:
+        """Enter async context manager."""
+        return self
+
+    async def __aexit__(self, *args: Any) -> None:
+        """Exit async context manager."""
+        await self.close()
+
+    def _register_server_tool_placeholder(self, tool_name: str) -> None:
+        """Register a declaration-only placeholder so function invocation skips execution."""
+
+        config = getattr(self, "function_invocation_configuration", None)
+        if not isinstance(config, dict):
+            return
+        additional_tools = list(config.get("additional_tools", []))
+        if any(getattr(tool, "name", None) == tool_name for tool in additional_tools):
+            return
+
+        placeholder: FunctionTool = FunctionTool(
+            name=tool_name,
+            description="Server-managed tool placeholder (AG-UI)",
+            func=None,
+        )
+        additional_tools.append(placeholder)
+        config["additional_tools"] = additional_tools
+        registered: set[str] = getattr(self, "_registered_server_tools", set())
+        registered.add(tool_name)
+        self._registered_server_tools = registered
+        logger.debug(f"[AGUIChatClient] Registered server placeholder: {tool_name}")
+
+    def _extract_state_from_messages(self, messages: Sequence[Message]) -> tuple[list[Message], dict[str, Any] | None]:
+        """Extract state from last message if present.
+
+        Args:
+            messages: List of chat messages
+
+        Returns:
+            Tuple of (messages_without_state, state_dict)
+        """
+        if not messages:
+            return list(messages), None
+
+        last_message = messages[-1]
+
+        for content in last_message.contents:
+            if isinstance(content, Content) and content.type == "data" and content.media_type == "application/json":
+                try:
+                    uri = content.uri
+                    if uri.startswith("data:application/json;base64,"):  # type: ignore[union-attr]
+                        import base64
+
+                        encoded_data = uri.split(",", 1)[1]  # type: ignore[union-attr]
+                        decoded_bytes = base64.b64decode(encoded_data)
+                        state = json.loads(decoded_bytes.decode("utf-8"))
+
+                        messages_without_state = list(messages[:-1]) if len(messages) > 1 else []
+                        return messages_without_state, state
+                except (json.JSONDecodeError, ValueError, KeyError) as e:
+                    logger.warning(f"Failed to extract state from message: {e}")
+
+        return list(messages), None
+
+    def _convert_messages_to_agui_format(self, messages: list[Message]) -> list[dict[str, Any]]:
+        """Convert Agent Framework messages to AG-UI format.
+
+        Args:
+            messages: List of Message objects
+
+        Returns:
+            List of AG-UI formatted message dictionaries
+        """
+        return agent_framework_messages_to_agui(messages)
+
+    def _get_thread_id(self, options: Mapping[str, Any]) -> str:
+        """Get or generate thread ID from chat options.
+
+        Args:
+            options: Chat options containing metadata
+
+        Returns:
+            Thread ID string
+        """
+        thread_id = None
+        metadata = options.get("metadata")
+        if metadata:
+            thread_id = metadata.get("thread_id")
+
+        if not thread_id:
+            thread_id = f"thread_{uuid.uuid4().hex}"
+
+        return thread_id
+
+    @override
+    def _inner_get_response(
+        self,
+        *,
+        messages: Sequence[Message],
+        stream: bool = False,
+        options: Mapping[str, Any],
+        **kwargs: Any,
+    ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
+        """Internal method to get non-streaming response.
+
+        Keyword Args:
+            messages: List of chat messages
+            stream: Whether to stream the response.
+            options: Chat options for the request
+            **kwargs: Additional keyword arguments
+
+        Returns:
+            ChatResponse object
+        """
+        if stream:
+            return ResponseStream(
+                self._streaming_impl(
+                    messages=messages,
+                    options=options,
+                    **kwargs,
+                ),
+                finalizer=ChatResponse.from_updates,
+            )
+
+        async def _get_response() -> ChatResponse:
+            return await ChatResponse.from_update_generator(
+                self._streaming_impl(
+                    messages=messages,
+                    options=options,
+                    **kwargs,
+                )
+            )
+
+        return _get_response()
+
+    async def _streaming_impl(
+        self,
+        *,
+        messages: Sequence[Message],
+        options: Mapping[str, Any],
+        **kwargs: Any,
+    ) -> AsyncIterable[ChatResponseUpdate]:
+        """Internal method to get streaming response.
+
+        Keyword Args:
+            messages: Sequence of chat messages
+            options: Chat options for the request
+            **kwargs: Additional keyword arguments
+
+        Yields:
+            ChatResponseUpdate objects
+        """
+        messages_to_send, state = self._extract_state_from_messages(messages)
+
+        thread_id = self._get_thread_id(options)
+        run_id = f"run_{uuid.uuid4().hex}"
+
+        agui_messages = self._convert_messages_to_agui_format(messages_to_send)
+
+        # Send client tools to server so LLM knows about them
+        # Client tools execute via Agent's function invocation wrapper
+        agui_tools = convert_tools_to_agui_format(options.get("tools"))
+
+        # Build set of client tool names (matches .NET clientToolSet)
+        # Used to distinguish client vs server tools in response stream
+        client_tool_set: set[str] = set()
+        tools = options.get("tools")
+        if tools:
+            for tool in tools:
+                if hasattr(tool, "name"):
+                    client_tool_set.add(tool.name)
+        self._last_client_tool_set = client_tool_set
+
+        logger.debug(
+            "[AGUIChatClient] Preparing request",
+            extra={
+                "thread_id": thread_id,
+                "run_id": run_id,
+                "client_tools": list(client_tool_set),
+                "messages": [msg.text for msg in messages_to_send if msg.text],
+            },
+        )
+        logger.debug(f"[AGUIChatClient] Client tool set: {client_tool_set}")
+
+        converter = AGUIEventConverter()
+
+        available_interrupts = options.get("available_interrupts", options.get("availableInterrupts"))
+
+        async for event in self._http_service.post_run(
+            thread_id=thread_id,
+            run_id=run_id,
+            messages=agui_messages,
+            state=state,
+            tools=agui_tools,
+            available_interrupts=_serialize_available_interrupts(cast(Sequence[Any] | None, available_interrupts)),
+            resume=_serialize_resume(options.get("resume")),
+        ):
+            logger.debug(f"[AGUIChatClient] Raw AG-UI event: {event}")
+            update = converter.convert_event(event)
+            if update is not None:
+                logger.debug(
+                    "[AGUIChatClient] Converted update",
+                    extra={"role": update.role, "contents": [type(c).__name__ for c in update.contents]},
+                )
+                # Distinguish client vs server tools
+                for i, content in enumerate(update.contents):
+                    if content.type == "function_call":
+                        logger.debug(
+                            f"[AGUIChatClient] Function call: {content.name}, in client_tool_set: {content.name in client_tool_set}"
+                        )
+                        if content.name in client_tool_set:
+                            # Client tool - let function invocation execute it
+                            if not content.additional_properties:
+                                content.additional_properties = {}
+                            content.additional_properties["agui_thread_id"] = thread_id
+                        else:
+                            # Server tool - wrap so function invocation ignores it
+                            logger.debug(f"[AGUIChatClient] Wrapping server tool: {content.name}")
+                            self._register_server_tool_placeholder(content.name)  # type: ignore[arg-type]
+                            update.contents[i] = Content(type="server_function_call", function_call=content)  # type: ignore
+
+                yield update

python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py

@@ -0,0 +1,257 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""FastAPI endpoint creation for AG-UI agents."""
+
+from __future__ import annotations
+
+import copy
+import logging
+from collections.abc import AsyncGenerator, Sequence
+from inspect import isawaitable
+from typing import Any, cast
+
+from ag_ui.core import RunErrorEvent
+from ag_ui.encoder import EventEncoder
+from agent_framework import SupportsAgentRun, Workflow
+from fastapi import FastAPI, HTTPException
+from fastapi.params import Depends
+from fastapi.responses import Response, StreamingResponse
+
+from ._agent import AgentFrameworkAgent
+from ._approval_state import _APPROVAL_SCOPE_INPUT_KEY
+from ._snapshots import (
+    _DEFAULT_STATE_INPUT_KEY,
+    _SNAPSHOT_SCOPE_INPUT_KEY,
+    AGUIThreadSnapshotStore,
+    SnapshotScopeResolver,
+)
+from ._types import AGUIRequest
+from ._workflow import AgentFrameworkWorkflow
+
+logger = logging.getLogger(__name__)
+
+_KEEPALIVE_COMMENT = "keepalive"
+
+
+def _get_snapshot_store(
+    protocol_runner: AgentFrameworkAgent | AgentFrameworkWorkflow,
+) -> AGUIThreadSnapshotStore | None:
+    if isinstance(protocol_runner, AgentFrameworkAgent):
+        return protocol_runner.config.snapshot_store
+    return protocol_runner.snapshot_store
+
+
+def _set_snapshot_store(
+    protocol_runner: AgentFrameworkAgent | AgentFrameworkWorkflow,
+    snapshot_store: AGUIThreadSnapshotStore,
+) -> None:
+    if isinstance(protocol_runner, AgentFrameworkAgent):
+        protocol_runner.config.snapshot_store = snapshot_store
+        return
+    protocol_runner.snapshot_store = snapshot_store
+
+
+def _configure_snapshot_persistence(
+    protocol_runner: AgentFrameworkAgent | AgentFrameworkWorkflow,
+    *,
+    snapshot_store: AGUIThreadSnapshotStore | None,
+    snapshot_scope_resolver: SnapshotScopeResolver | None,
+) -> None:
+    existing_snapshot_store = _get_snapshot_store(protocol_runner)
+    if snapshot_store is not None:
+        if existing_snapshot_store is not None and existing_snapshot_store is not snapshot_store:
+            raise ValueError("snapshot_store is already configured on the AG-UI runner.")
+        if existing_snapshot_store is None:
+            _set_snapshot_store(protocol_runner, snapshot_store)
+        existing_snapshot_store = snapshot_store
+
+    if existing_snapshot_store is not None and snapshot_scope_resolver is None:
+        raise ValueError(
+            "snapshot_scope_resolver is required when snapshot_store is configured. "
+            "AG-UI Thread ids identify threads but do not authorize snapshot access; "
+            "provide a resolver that returns an explicit Snapshot Scope."
+        )
+
+
+def _validate_keepalive_seconds(keepalive_seconds: float | None) -> None:
+    if keepalive_seconds is not None and not keepalive_seconds > 0:
+        raise ValueError("keepalive_seconds must be positive or None.")
+
+
+def add_agent_framework_fastapi_endpoint(
+    app: FastAPI,
+    agent: SupportsAgentRun | AgentFrameworkAgent | Workflow | AgentFrameworkWorkflow,
+    path: str = "/",
+    state_schema: Any | None = None,
+    predict_state_config: dict[str, dict[str, str]] | None = None,
+    allow_origins: list[str] | None = None,
+    default_state: dict[str, Any] | None = None,
+    tags: list[str] | None = None,
+    dependencies: Sequence[Depends] | None = None,
+    snapshot_store: AGUIThreadSnapshotStore | None = None,
+    snapshot_scope_resolver: SnapshotScopeResolver | None = None,
+    keepalive_seconds: float | None = 15,
+) -> None:
+    """Add an AG-UI endpoint to a FastAPI app.
+
+    Args:
+        app: The FastAPI application
+        agent: The agent to expose (can be raw SupportsAgentRun or wrapped)
+        path: The endpoint path
+        state_schema: Optional state schema for shared state management; accepts dict or Pydantic model/class
+        predict_state_config: Optional predictive state update configuration.
+            Format: {"state_key": {"tool": "tool_name", "tool_argument": "arg_name"}}
+        allow_origins: CORS origins (not yet implemented)
+        default_state: Optional initial state to seed when the client does not provide state keys
+        tags: OpenAPI tags for endpoint categorization (defaults to ["AG-UI"])
+        dependencies: Optional FastAPI dependencies for authentication/authorization.
+            These dependencies run before the endpoint handler. Use this to add
+            authentication checks, rate limiting, or other middleware-like behavior.
+            Example: `dependencies=[Depends(verify_api_key)]`
+        snapshot_store: Optional AG-UI Thread Snapshot store. Snapshot persistence is opt-in and requires an
+            explicit Snapshot Scope resolver.
+        snapshot_scope_resolver: Optional resolver for the application-defined Snapshot Scope. Required whenever
+            a snapshot store is configured because an AG-UI Thread id is not an authorization boundary.
+        keepalive_seconds: Endpoint SSE keepalive interval in seconds. Defaults to 15. Positive values emit fixed
+            SSE comments while the stream is open. None disables keepalive and preserves the non-keepalive response
+            path. Keepalive comments are transport traffic and do not change AG-UI events.
+    """
+    _validate_keepalive_seconds(keepalive_seconds)
+
+    protocol_runner: AgentFrameworkAgent | AgentFrameworkWorkflow
+    if isinstance(agent, AgentFrameworkWorkflow):
+        protocol_runner = agent
+    elif isinstance(agent, AgentFrameworkAgent):
+        protocol_runner = agent
+    elif isinstance(agent, Workflow):
+        protocol_runner = AgentFrameworkWorkflow(workflow=agent)
+    elif isinstance(agent, SupportsAgentRun):
+        protocol_runner = AgentFrameworkAgent(
+            agent=agent,
+            state_schema=state_schema,
+            predict_state_config=predict_state_config,
+            snapshot_store=snapshot_store,
+        )
+    else:
+        raise TypeError("agent must be SupportsAgentRun, Workflow, AgentFrameworkAgent, or AgentFrameworkWorkflow.")
+
+    _configure_snapshot_persistence(
+        protocol_runner,
+        snapshot_store=snapshot_store,
+        snapshot_scope_resolver=snapshot_scope_resolver,
+    )
+
+    @app.post(path, tags=tags or ["AG-UI"], dependencies=dependencies, response_model=None)  # type: ignore[arg-type]
+    async def agent_endpoint(request_body: AGUIRequest) -> Response:
+        """Handle AG-UI agent requests.
+
+        Note: Function is accessed via FastAPI's decorator registration,
+        despite appearing unused to static analysis.
+        """
+        try:
+            input_data = request_body.model_dump(exclude_none=True)
+            snapshot_persistence_active = False
+            if snapshot_scope_resolver is not None:
+                snapshot_scope = snapshot_scope_resolver(request_body)
+                if isawaitable(snapshot_scope):
+                    snapshot_scope = await snapshot_scope
+                input_data[_APPROVAL_SCOPE_INPUT_KEY] = snapshot_scope
+                if _get_snapshot_store(protocol_runner) is not None:
+                    input_data[_SNAPSHOT_SCOPE_INPUT_KEY] = snapshot_scope
+                    snapshot_persistence_active = True
+            if default_state:
+                if snapshot_persistence_active:
+                    # Defer default application to the runner so defaults only fill keys
+                    # missing from both the stored snapshot state and the request state.
+                    input_data[_DEFAULT_STATE_INPUT_KEY] = copy.deepcopy(default_state)
+                else:
+                    state = input_data.setdefault("state", {})
+                    for key, value in default_state.items():
+                        if key not in state:
+                            state[key] = copy.deepcopy(value)
+            logger.debug(
+                f"[{path}] Received request - Run ID: {input_data.get('run_id', 'no-run-id')}, "
+                f"Thread ID: {input_data.get('thread_id', 'no-thread-id')}, "
+                f"Messages: {len(input_data.get('messages', []))}"
+            )
+            logger.info(f"Received request at {path}: {input_data.get('run_id', 'no-run-id')}")
+
+            keepalive_enabled = keepalive_seconds is not None
+
+            def prepare_frame(encoded: str) -> str | bytes:
+                if keepalive_enabled:
+                    return encoded.encode("utf-8")
+                return encoded
+
+            async def event_generator() -> AsyncGenerator[str | bytes]:
+                encoder = EventEncoder()
+                event_count = 0
+                try:
+                    async for event in protocol_runner.run(input_data):
+                        event_count += 1
+                        event_type_name = getattr(event, "type", type(event).__name__)
+                        # Log important events at INFO level
+                        if "TOOL_CALL" in str(event_type_name) or "RUN" in str(event_type_name):
+                            if hasattr(event, "model_dump"):
+                                event_data = event.model_dump(exclude_none=True)
+                                logger.info(f"[{path}] Event {event_count}: {event_type_name} - {event_data}")
+                            else:
+                                logger.info(f"[{path}] Event {event_count}: {event_type_name}")
+
+                        try:
+                            encoded = encoder.encode(event)
+                        except Exception as encode_error:
+                            logger.exception("[%s] Failed to encode event %s", path, event_type_name)
+                            run_error = RunErrorEvent(
+                                message="An internal error has occurred while streaming events.",
+                                code=type(encode_error).__name__,
+                            )
+                            try:
+                                yield prepare_frame(encoder.encode(run_error))
+                            except Exception:
+                                logger.exception("[%s] Failed to encode RUN_ERROR event", path)
+                            return
+
+                        logger.debug(
+                            f"[{path}] Encoded as: {encoded[:200]}..."
+                            if len(encoded) > 200
+                            else f"[{path}] Encoded as: {encoded}"
+                        )
+                        yield prepare_frame(encoded)
+
+                    logger.info(f"[{path}] Completed streaming {event_count} events")
+                except Exception as stream_error:
+                    logger.exception("[%s] Streaming failed", path)
+                    run_error = RunErrorEvent(
+                        message="An internal error has occurred while streaming events.",
+                        code=type(stream_error).__name__,
+                    )
+                    try:
+                        yield prepare_frame(encoder.encode(run_error))
+                    except Exception:
+                        logger.exception("[%s] Failed to encode RUN_ERROR event", path)
+
+            headers = {
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+                "X-Accel-Buffering": "no",
+            }
+            if keepalive_seconds is not None:
+                from sse_starlette.event import ServerSentEvent
+                from sse_starlette.sse import EventSourceResponse
+
+                return EventSourceResponse(
+                    event_generator(),
+                    ping=cast(int, keepalive_seconds),
+                    ping_message_factory=lambda: ServerSentEvent(comment=_KEEPALIVE_COMMENT),
+                    headers=headers,
+                    media_type="text/event-stream",
+                )
+            return StreamingResponse(
+                event_generator(),
+                media_type="text/event-stream",
+                headers=headers,
+            )
+        except Exception as e:
+            logger.error(f"Error in agent endpoint: {e}", exc_info=True)
+            raise HTTPException(status_code=500, detail="An internal error has occurred.") from e

python/packages/ag-ui/agent_framework_ag_ui/_event_converters.py

@@ -0,0 +1,279 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Event converter for AG-UI protocol events to Agent Framework types."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from agent_framework import (
+    ChatResponseUpdate,
+    Content,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class AGUIEventConverter:
+    """Converter for AG-UI events to Agent Framework types.
+
+    Handles conversion of AG-UI protocol events to ChatResponseUpdate objects
+    while maintaining state, aggregating content, and tracking metadata.
+    """
+
+    def __init__(self) -> None:
+        """Initialize the converter with fresh state."""
+        self.current_message_id: str | None = None
+        self.current_tool_call_id: str | None = None
+        self.current_tool_name: str | None = None
+        self.accumulated_tool_args: str = ""
+        self.thread_id: str | None = None
+        self.run_id: str | None = None
+
+    @staticmethod
+    def _get_tool_call_id(event: dict[str, Any]) -> str | None:
+        """Return the tool call ID from either AG-UI field spelling."""
+        tool_call_id = event.get("toolCallId")
+        if tool_call_id is None:
+            tool_call_id = event.get("tool_call_id")
+        if tool_call_id is None:
+            return None
+        return str(tool_call_id)
+
+    def convert_event(self, event: dict[str, Any]) -> ChatResponseUpdate | None:
+        """Convert a single AG-UI event to ChatResponseUpdate.
+
+        Args:
+            event: AG-UI event dictionary
+
+        Returns:
+            ChatResponseUpdate if event produces content, None otherwise
+
+        Examples:
+            RUN_STARTED event:
+
+            .. code-block:: python
+
+                converter = AGUIEventConverter()
+                event = {"type": "RUN_STARTED", "threadId": "t1", "runId": "r1"}
+                update = converter.convert_event(event)
+                assert update.additional_properties["thread_id"] == "t1"
+
+            TEXT_MESSAGE_CONTENT event:
+
+            .. code-block:: python
+
+                event = {"type": "TEXT_MESSAGE_CONTENT", "messageId": "m1", "delta": "Hello"}
+                update = converter.convert_event(event)
+                assert update.contents[0].text == "Hello"
+        """
+        raw_event_type = str(event.get("type", ""))
+        event_type = raw_event_type.upper()
+
+        if event_type == "RUN_STARTED":
+            return self._handle_run_started(event)
+        elif event_type == "TEXT_MESSAGE_START":
+            return self._handle_text_message_start(event)
+        elif event_type == "TEXT_MESSAGE_CONTENT":
+            return self._handle_text_message_content(event)
+        elif event_type == "TEXT_MESSAGE_END":
+            return self._handle_text_message_end(event)
+        elif event_type == "TOOL_CALL_START":
+            return self._handle_tool_call_start(event)
+        elif event_type == "TOOL_CALL_ARGS":
+            return self._handle_tool_call_args(event)
+        elif event_type == "TOOL_CALL_END":
+            return self._handle_tool_call_end(event)
+        elif event_type == "TOOL_CALL_RESULT":
+            return self._handle_tool_call_result(event)
+        elif event_type == "RUN_FINISHED":
+            return self._handle_run_finished(event)
+        elif event_type == "RUN_ERROR":
+            return self._handle_run_error(event)
+        elif event_type in {"CUSTOM", "CUSTOM_EVENT"}:
+            return self._handle_custom_event(event, raw_event_type)
+
+        return None
+
+    def _handle_run_started(self, event: dict[str, Any]) -> ChatResponseUpdate:
+        """Handle RUN_STARTED event."""
+        self.thread_id = event.get("threadId")
+        self.run_id = event.get("runId")
+
+        return ChatResponseUpdate(
+            role="assistant",
+            contents=[],
+            additional_properties={
+                "thread_id": self.thread_id,
+                "run_id": self.run_id,
+            },
+        )
+
+    def _handle_text_message_start(self, event: dict[str, Any]) -> ChatResponseUpdate | None:
+        """Handle TEXT_MESSAGE_START event."""
+        self.current_message_id = event.get("messageId")
+        return ChatResponseUpdate(
+            role="assistant",
+            message_id=self.current_message_id,
+            contents=[],
+        )
+
+    def _handle_text_message_content(self, event: dict[str, Any]) -> ChatResponseUpdate:
+        """Handle TEXT_MESSAGE_CONTENT event."""
+        message_id = event.get("messageId")
+        delta = event.get("delta", "")
+
+        if message_id != self.current_message_id:
+            self.current_message_id = message_id
+
+        return ChatResponseUpdate(
+            role="assistant",
+            message_id=self.current_message_id,
+            contents=[Content.from_text(text=delta)],
+        )
+
+    def _handle_text_message_end(self, event: dict[str, Any]) -> ChatResponseUpdate | None:
+        """Handle TEXT_MESSAGE_END event."""
+        return None
+
+    def _handle_tool_call_start(self, event: dict[str, Any]) -> ChatResponseUpdate:
+        """Handle TOOL_CALL_START event."""
+        self.current_tool_call_id = self._get_tool_call_id(event)
+        self.current_tool_name = event.get("toolName") or event.get("toolCallName") or event.get("tool_call_name")
+        self.accumulated_tool_args = ""
+
+        return ChatResponseUpdate(
+            role="assistant",
+            contents=[
+                Content.from_function_call(
+                    call_id=self.current_tool_call_id or "",
+                    name=self.current_tool_name or "",
+                    arguments="",
+                )
+            ],
+        )
+
+    def _handle_tool_call_args(self, event: dict[str, Any]) -> ChatResponseUpdate | None:
+        """Handle TOOL_CALL_ARGS event."""
+        event_tool_call_id = self._get_tool_call_id(event)
+        if event_tool_call_id is not None:
+            if self.current_tool_call_id and event_tool_call_id != self.current_tool_call_id:
+                logger.warning(
+                    "Ignoring TOOL_CALL_ARGS for toolCallId=%s while current toolCallId=%s",
+                    event_tool_call_id,
+                    self.current_tool_call_id,
+                )
+                return None
+            if not self.current_tool_call_id:
+                self.current_tool_call_id = event_tool_call_id
+
+        delta = event.get("delta", "")
+        self.accumulated_tool_args += delta
+
+        return ChatResponseUpdate(
+            role="assistant",
+            contents=[
+                Content.from_function_call(
+                    call_id=self.current_tool_call_id or "",
+                    name=self.current_tool_name or "",
+                    arguments=delta,
+                )
+            ],
+        )
+
+    def _handle_tool_call_end(self, event: dict[str, Any]) -> ChatResponseUpdate | None:
+        """Handle TOOL_CALL_END event."""
+        event_tool_call_id = self._get_tool_call_id(event)
+        if (
+            self.current_tool_call_id is None
+            or event_tool_call_id is None
+            or event_tool_call_id == self.current_tool_call_id
+        ):
+            self.current_tool_call_id = None
+            self.current_tool_name = None
+            self.accumulated_tool_args = ""
+        return None
+
+    def _handle_tool_call_result(self, event: dict[str, Any]) -> ChatResponseUpdate:
+        """Handle TOOL_CALL_RESULT event."""
+        tool_call_id = event.get("toolCallId", "")
+        result = event.get("result") if event.get("result") is not None else event.get("content")
+
+        return ChatResponseUpdate(
+            role="tool",
+            contents=[
+                Content.from_function_result(
+                    call_id=tool_call_id,
+                    result=result,
+                )
+            ],
+        )
+
+    def _handle_run_finished(self, event: dict[str, Any]) -> ChatResponseUpdate:
+        """Handle RUN_FINISHED event."""
+        additional_properties: dict[str, Any] = {
+            "thread_id": self.thread_id,
+            "run_id": self.run_id,
+        }
+        if "interrupt" in event:
+            additional_properties["interrupt"] = event.get("interrupt")
+        if "outcome" in event:
+            outcome = event.get("outcome")
+            additional_properties["outcome"] = outcome
+            if not isinstance(outcome, dict):
+                logger.warning(
+                    "RUN_FINISHED outcome should be an object; got %s. Preserving raw outcome.",
+                    type(outcome).__name__,
+                )
+            elif outcome.get("type") == "interrupt":
+                interrupts = outcome.get("interrupts")
+                if isinstance(interrupts, list):
+                    additional_properties["interrupts"] = interrupts
+        if "result" in event:
+            additional_properties["result"] = event.get("result")
+
+        return ChatResponseUpdate(
+            role="assistant",
+            finish_reason="stop",
+            contents=[],
+            additional_properties=additional_properties,
+        )
+
+    def _handle_run_error(self, event: dict[str, Any]) -> ChatResponseUpdate:
+        """Handle RUN_ERROR event."""
+        error_message = event.get("message", "Unknown error")
+
+        return ChatResponseUpdate(
+            role="assistant",
+            finish_reason="content_filter",
+            contents=[
+                Content.from_error(
+                    message=error_message,
+                    error_code="RUN_ERROR",
+                )
+            ],
+            additional_properties={
+                "thread_id": self.thread_id,
+                "run_id": self.run_id,
+            },
+        )
+
+    def _handle_custom_event(self, event: dict[str, Any], raw_event_type: str) -> ChatResponseUpdate:
+        """Handle CUSTOM/CUSTOM_EVENT events.
+
+        Custom events are surfaced as metadata so callers can inspect protocol-specific payloads.
+        """
+        return ChatResponseUpdate(
+            role="assistant",
+            contents=[],
+            additional_properties={
+                "thread_id": self.thread_id,
+                "run_id": self.run_id,
+                "ag_ui_custom_event": {
+                    "name": event.get("name"),
+                    "value": event.get("value"),
+                    "raw_type": raw_event_type,
+                },
+            },
+        )

python/packages/ag-ui/agent_framework_ag_ui/_http_service.py

@@ -0,0 +1,265 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""HTTP service for AG-UI protocol communication."""
+
+from __future__ import annotations
+
+import json
+import logging
+from collections.abc import AsyncIterable, Mapping, Sequence
+from typing import Any, cast
+
+import httpx
+from ag_ui.core import Interrupt, ResumeEntry
+
+logger = logging.getLogger(__name__)
+
+
+def _json_safe_protocol_value(value: Any) -> Any:
+    """Convert protocol values to JSON-compatible data using AG-UI aliases."""
+    model_dump = getattr(value, "model_dump", None)
+    if callable(model_dump):
+        return _json_safe_protocol_value(model_dump(by_alias=True, exclude_none=True))
+    if isinstance(value, Mapping):
+        return {key: _json_safe_protocol_value(item) for key, item in value.items()}
+    if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
+        return [_json_safe_protocol_value(item) for item in value]
+    return value
+
+
+def _serialize_available_interrupts(available_interrupts: Sequence[Any] | None) -> list[dict[str, Any]] | None:
+    """Serialize typed or compatible interrupt inputs to canonical AG-UI JSON."""
+    if available_interrupts is None:
+        return None
+    serialized: list[dict[str, Any]] = []
+    for interrupt in available_interrupts:
+        if isinstance(interrupt, Mapping) and "reason" not in interrupt:
+            interrupt = dict(interrupt)
+            interrupt_type = interrupt.pop("type", None)
+            if interrupt_type == "request_info" or interrupt_type is None:
+                interrupt["reason"] = "input_required"
+            elif isinstance(interrupt_type, str):
+                interrupt["reason"] = interrupt_type
+        serialized.append(
+            cast(dict[str, Any], Interrupt.model_validate(interrupt).model_dump(by_alias=True, exclude_none=True))
+        )
+    return serialized
+
+
+def _serialize_resume_entry(entry: Any) -> dict[str, Any]:
+    """Serialize one typed or legacy resume entry to canonical AG-UI JSON."""
+    model_dump = getattr(entry, "model_dump", None)
+    if callable(model_dump):
+        entry = model_dump(by_alias=True, exclude_none=True)
+
+    if not isinstance(entry, Mapping):
+        raise ValueError("Each resume entry must be an object.")
+
+    entry_dict = cast(Mapping[str, Any], entry)
+    interrupt_id = (
+        entry_dict.get("interruptId")
+        or entry_dict.get("interrupt_id")
+        or entry_dict.get("id")
+        or entry_dict.get("toolCallId")
+    )
+    if not interrupt_id:
+        raise ValueError("Each resume entry must include interruptId.")
+
+    status = entry_dict.get("status") or "resolved"
+    payload = (
+        entry_dict.get("payload")
+        if "payload" in entry_dict
+        else entry_dict.get("value")
+        if "value" in entry_dict
+        else entry_dict.get("response")
+        if "response" in entry_dict
+        else {
+            key: value
+            for key, value in entry_dict.items()
+            if key not in {"id", "interruptId", "interrupt_id", "toolCallId", "type", "status"}
+        }
+    )
+
+    serialized: dict[str, Any] = {"interruptId": str(interrupt_id), "status": str(status)}
+    if status != "cancelled" or payload:
+        serialized["payload"] = _json_safe_protocol_value(payload)
+    return cast(dict[str, Any], ResumeEntry.model_validate(serialized).model_dump(by_alias=True, exclude_none=True))
+
+
+def _serialize_resume(resume: Any) -> Any:  # noqa: ANN401
+    """Serialize typed or compatible resume inputs to canonical AG-UI JSON."""
+    if resume is None:
+        return None
+    if isinstance(resume, Sequence) and not isinstance(resume, (str, bytes, bytearray)):
+        return [_serialize_resume_entry(entry) for entry in resume]
+    if isinstance(resume, Mapping):
+        resume_dict = cast(Mapping[str, Any], resume)
+        if isinstance(resume_dict.get("interrupts"), Sequence) and not isinstance(
+            resume_dict.get("interrupts"), (str, bytes, bytearray)
+        ):
+            return [_serialize_resume_entry(entry) for entry in cast(Sequence[Any], resume_dict["interrupts"])]
+        if isinstance(resume_dict.get("interrupt"), Sequence) and not isinstance(
+            resume_dict.get("interrupt"), (str, bytes, bytearray)
+        ):
+            return [_serialize_resume_entry(entry) for entry in cast(Sequence[Any], resume_dict["interrupt"])]
+        if any(key in resume_dict for key in ("interruptId", "interrupt_id", "id", "toolCallId")):
+            return [_serialize_resume_entry(resume_dict)]
+    return _json_safe_protocol_value(resume)
+
+
+class AGUIHttpService:
+    """HTTP service for AG-UI protocol communication.
+
+    Handles HTTP POST requests and Server-Sent Events (SSE) stream parsing
+    for the AG-UI protocol.
+
+    Examples:
+        Basic usage:
+
+        .. code-block:: python
+
+            service = AGUIHttpService("http://localhost:8888/")
+            async for event in service.post_run(
+                thread_id="thread_123",
+                run_id="run_456",
+                messages=[{"role": "user", "content": "Hello"}]
+            ):
+                print(event["type"])
+
+        With context manager:
+
+        .. code-block:: python
+
+            async with AGUIHttpService("http://localhost:8888/") as service:
+                async for event in service.post_run(...):
+                    print(event)
+    """
+
+    def __init__(
+        self,
+        endpoint: str,
+        http_client: httpx.AsyncClient | None = None,
+        timeout: float = 60.0,
+    ) -> None:
+        """Initialize the HTTP service.
+
+        Args:
+            endpoint: AG-UI server endpoint URL (e.g., "http://localhost:8888/")
+            http_client: Optional httpx AsyncClient. If None, creates a new one.
+            timeout: Request timeout in seconds (default: 60.0)
+        """
+        self.endpoint = endpoint.rstrip("/")
+        self._owns_client = http_client is None
+        self.http_client = http_client or httpx.AsyncClient(timeout=timeout)
+
+    async def post_run(
+        self,
+        thread_id: str,
+        run_id: str,
+        messages: list[dict[str, Any]],
+        state: dict[str, Any] | None = None,
+        tools: list[dict[str, Any]] | None = None,
+        available_interrupts: Sequence[Any] | None = None,
+        resume: Any = None,
+    ) -> AsyncIterable[dict[str, Any]]:
+        """Post a run request and stream AG-UI events.
+
+        Args:
+            thread_id: Thread identifier for conversation continuity
+            run_id: Unique run identifier
+            messages: List of messages in AG-UI format
+            state: Optional state object to send to server
+            tools: Optional list of tools available to the agent
+            available_interrupts: Optional list of interrupt descriptors available for resumption
+            resume: Optional resume payload to continue a paused run
+
+        Yields:
+            AG-UI event dictionaries parsed from SSE stream
+
+        Raises:
+            httpx.HTTPStatusError: If the HTTP request fails
+            ValueError: If SSE parsing encounters invalid data
+
+        Examples:
+            .. code-block:: python
+
+                service = AGUIHttpService("http://localhost:8888/")
+                async for event in service.post_run(
+                    thread_id="thread_abc",
+                    run_id="run_123",
+                    messages=[{"role": "user", "content": "Hello"}],
+                    state={"user_context": {"name": "Alice"}}
+                ):
+                    if event["type"] == "TEXT_MESSAGE_CONTENT":
+                        print(event["delta"])
+        """
+        # Build request payload
+        request_data: dict[str, Any] = {
+            "thread_id": thread_id,
+            "run_id": run_id,
+            "messages": messages,
+        }
+
+        if state is not None:
+            request_data["state"] = state
+
+        if tools is not None:
+            request_data["tools"] = tools
+
+        serialized_available_interrupts = _serialize_available_interrupts(available_interrupts)
+        if serialized_available_interrupts is not None:
+            request_data["availableInterrupts"] = serialized_available_interrupts
+
+        serialized_resume = _serialize_resume(resume)
+        if serialized_resume is not None:
+            request_data["resume"] = serialized_resume
+
+        logger.debug(
+            f"Posting run to {self.endpoint}: thread_id={thread_id}, run_id={run_id}, "
+            f"messages={len(messages)}, has_state={state is not None}, has_tools={tools is not None}, "
+            f"has_available_interrupts={available_interrupts is not None}, has_resume={resume is not None}"
+        )
+
+        # Stream the response using SSE
+        async with self.http_client.stream(
+            "POST",
+            self.endpoint,
+            json=request_data,
+            headers={"Accept": "text/event-stream"},
+        ) as response:
+            try:
+                response.raise_for_status()
+            except httpx.HTTPStatusError as e:
+                logger.error(f"HTTP request failed: {e.response.status_code} - {e.response.text}")
+                raise
+
+            async for line in response.aiter_lines():
+                # Parse Server-Sent Events format
+                if line.startswith("data: "):
+                    data = line[6:]  # Remove "data: " prefix
+                    try:
+                        event = json.loads(data)
+                        logger.debug(f"Received event: {event.get('type', 'UNKNOWN')}")
+                        yield event
+                    except json.JSONDecodeError as e:
+                        logger.warning(f"Failed to parse SSE data: {data}. Error: {e}")
+                        # Continue processing other events instead of failing
+                        continue
+
+    async def close(self) -> None:
+        """Close the HTTP client if owned by this service.
+
+        Only closes the client if it was created by this service instance.
+        If an external client was provided, it remains the caller's
+        responsibility to close it.
+        """
+        if self._owns_client and self.http_client:
+            await self.http_client.aclose()
+
+    async def __aenter__(self) -> AGUIHttpService:
+        """Enter async context manager."""
+        return self
+
+    async def __aexit__(self, *args: Any) -> None:
+        """Exit async context manager and clean up resources."""
+        await self.close()

python/packages/ag-ui/agent_framework_ag_ui/_orchestration/__init__.py

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

python/packages/ag-ui/agent_framework_ag_ui/_orchestration/_helpers.py

@@ -0,0 +1,247 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Helper functions for orchestration logic.
+
+This module retains utilities that may be useful for testing or extensions.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from typing import Any
+
+from agent_framework import (
+    Content,
+    Message,
+)
+
+from .._utils import get_role_value
+
+logger = logging.getLogger(__name__)
+
+
+def pending_tool_call_ids(messages: list[Message]) -> set[str]:
+    """Get IDs of tool calls without corresponding results.
+
+    Args:
+        messages: List of messages to scan
+
+    Returns:
+        Set of pending tool call IDs
+    """
+    pending_ids: set[str] = set()
+    resolved_ids: set[str] = set()
+    for msg in messages:
+        for content in msg.contents:
+            if content.type == "function_call" and content.call_id:
+                pending_ids.add(str(content.call_id))
+            elif content.type == "function_result" and content.call_id:
+                resolved_ids.add(str(content.call_id))
+    return pending_ids - resolved_ids
+
+
+def is_state_context_message(message: Message) -> bool:
+    """Check if a message is a state context system message.
+
+    Args:
+        message: Message to check
+
+    Returns:
+        True if this is a state context message
+    """
+    if get_role_value(message) != "system":
+        return False
+    for content in message.contents:
+        if content.type == "text" and content.text.startswith("Current state of the application:"):  # type: ignore[union-attr]
+            return True
+    return False
+
+
+def ensure_tool_call_entry(
+    tool_call_id: str,
+    tool_calls_by_id: dict[str, dict[str, Any]],
+    pending_tool_calls: list[dict[str, Any]],
+) -> dict[str, Any]:
+    """Get or create a tool call entry in the tracking dicts.
+
+    Args:
+        tool_call_id: The tool call ID
+        tool_calls_by_id: Dict mapping IDs to tool call entries
+        pending_tool_calls: List of pending tool calls
+
+    Returns:
+        The tool call entry dict
+    """
+    entry = tool_calls_by_id.get(tool_call_id)
+    if entry is None:
+        entry = {
+            "id": tool_call_id,
+            "type": "function",
+            "function": {
+                "name": "",
+                "arguments": "",
+            },
+        }
+        tool_calls_by_id[tool_call_id] = entry
+        pending_tool_calls.append(entry)
+    return entry
+
+
+def tool_name_for_call_id(
+    tool_calls_by_id: dict[str, dict[str, Any]],
+    tool_call_id: str,
+) -> str | None:
+    """Get the tool name for a given call ID.
+
+    Args:
+        tool_calls_by_id: Dict mapping IDs to tool call entries
+        tool_call_id: The tool call ID to look up
+
+    Returns:
+        Tool name or None if not found
+    """
+    entry = tool_calls_by_id.get(tool_call_id)
+    if not entry:
+        return None
+    function = entry.get("function")
+    if not isinstance(function, dict):
+        return None
+    name = function.get("name")
+    return str(name) if name else None
+
+
+def schema_has_steps(schema: Any) -> bool:
+    """Check if a schema has a steps array property.
+
+    Args:
+        schema: JSON schema to check
+
+    Returns:
+        True if schema has steps array
+    """
+    if not isinstance(schema, dict):
+        return False
+    properties = schema.get("properties")
+    if not isinstance(properties, dict):
+        return False
+    steps_schema = properties.get("steps")
+    if not isinstance(steps_schema, dict):
+        return False
+    return steps_schema.get("type") == "array"
+
+
+def select_approval_tool_name(client_tools: list[Any] | None) -> str | None:
+    """Select appropriate approval tool from client tools.
+
+    Args:
+        client_tools: List of client tool definitions
+
+    Returns:
+        Name of approval tool, or None if not found
+    """
+    if not client_tools:
+        return None
+    for tool in client_tools:
+        tool_name = getattr(tool, "name", None)
+        if not tool_name:
+            continue
+        params_fn = getattr(tool, "parameters", None)
+        if not callable(params_fn):
+            continue
+        schema = params_fn()
+        if schema_has_steps(schema):
+            return str(tool_name)
+    return None
+
+
+def build_safe_metadata(thread_metadata: dict[str, Any] | None) -> dict[str, Any]:
+    """Build metadata dict with truncated string values for Azure compatibility.
+
+    Azure has a 512 character limit per metadata value.
+
+    Args:
+        thread_metadata: Raw metadata dict
+
+    Returns:
+        Metadata with string values truncated to 512 chars
+    """
+    if not thread_metadata:
+        return {}
+    safe_metadata: dict[str, Any] = {}
+    for key, value in thread_metadata.items():
+        value_str = value if isinstance(value, str) else json.dumps(value)
+        if len(value_str) > 512:
+            value_str = value_str[:512]
+        safe_metadata[key] = value_str
+    return safe_metadata
+
+
+def latest_approval_response(messages: list[Message]) -> Content | None:
+    """Get the latest approval response from messages.
+
+    Args:
+        messages: Messages to search
+
+    Returns:
+        Latest approval response or None
+    """
+    if not messages:
+        return None
+    last_message = messages[-1]
+    for content in last_message.contents:
+        if content.type == "function_approval_response":
+            return content
+    return None
+
+
+def approval_steps(approval: Content) -> list[Any]:
+    """Extract steps from an approval response.
+
+    Args:
+        approval: Approval response content
+
+    Returns:
+        List of steps, or empty list if none
+    """
+    state_args = approval.additional_properties.get("ag_ui_state_args", None)
+    if isinstance(state_args, dict):
+        steps = state_args.get("steps")
+        if isinstance(steps, list):
+            return steps
+
+    if approval.function_call:
+        parsed_args = approval.function_call.parse_arguments()
+        if isinstance(parsed_args, dict):
+            steps = parsed_args.get("steps")
+            if isinstance(steps, list):
+                return steps
+
+    return []
+
+
+def is_step_based_approval(
+    approval: Content,
+    predict_state_config: dict[str, dict[str, str]] | None,
+) -> bool:
+    """Check if an approval is step-based.
+
+    Args:
+        approval: Approval response to check
+        predict_state_config: Predictive state configuration
+
+    Returns:
+        True if this is a step-based approval
+    """
+    steps = approval_steps(approval)
+    if steps:
+        return True
+    if not approval.function_call:
+        return False
+    if not predict_state_config:
+        return False
+    tool_name = approval.function_call.name
+    for config in predict_state_config.values():
+        if config.get("tool") == tool_name and config.get("tool_argument") == "steps":
+            return True
+    return False

python/packages/ag-ui/agent_framework_ag_ui/_orchestration/_predictive_state.py

@@ -0,0 +1,232 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Predictive state handling utilities."""
+
+from __future__ import annotations
+
+import json
+import logging
+import re
+from typing import Any
+
+from ag_ui.core import StateDeltaEvent
+
+from .._utils import safe_json_parse
+
+logger = logging.getLogger(__name__)
+
+
+class PredictiveStateHandler:
+    """Handles predictive state updates from streaming tool calls."""
+
+    def __init__(
+        self,
+        predict_state_config: dict[str, dict[str, str]] | None = None,
+        current_state: dict[str, Any] | None = None,
+    ) -> None:
+        """Initialize the handler.
+
+        Args:
+            predict_state_config: Configuration mapping state keys to tool/argument pairs
+            current_state: Reference to current state dict
+        """
+        self.predict_state_config = predict_state_config or {}
+        self.current_state = current_state or {}
+        self.streaming_tool_args: str = ""
+        self.last_emitted_state: dict[str, Any] = {}
+        self.state_delta_count: int = 0
+        self.pending_state_updates: dict[str, Any] = {}
+
+    def reset_streaming(self) -> None:
+        """Reset streaming state for a new tool call."""
+        self.streaming_tool_args = ""
+        self.state_delta_count = 0
+
+    def extract_state_value(
+        self,
+        tool_name: str,
+        args: dict[str, Any] | str | None,
+    ) -> tuple[str, Any] | None:
+        """Extract state value from tool arguments based on config.
+
+        Args:
+            tool_name: Name of the tool being called
+            args: Tool arguments (dict or JSON string)
+
+        Returns:
+            Tuple of (state_key, state_value) or None if no match
+        """
+        if not self.predict_state_config:
+            return None
+
+        parsed_args = safe_json_parse(args) if isinstance(args, str) else args
+        if not parsed_args:
+            return None
+
+        for state_key, config in self.predict_state_config.items():
+            if config["tool"] != tool_name:
+                continue
+            tool_arg_name = config["tool_argument"]
+            if tool_arg_name == "*":
+                return (state_key, parsed_args)
+            if tool_arg_name in parsed_args:
+                return (state_key, parsed_args[tool_arg_name])
+
+        return None
+
+    def is_predictive_tool(self, tool_name: str | None) -> bool:
+        """Check if a tool is configured for predictive state.
+
+        Args:
+            tool_name: Name of the tool to check
+
+        Returns:
+            True if tool is in predictive state config
+        """
+        if not tool_name or not self.predict_state_config:
+            return False
+        for config in self.predict_state_config.values():
+            if config["tool"] == tool_name:
+                return True
+        return False
+
+    def emit_streaming_deltas(
+        self,
+        tool_name: str | None,
+        argument_chunk: str,
+    ) -> list[StateDeltaEvent]:
+        """Process streaming argument chunk and emit state deltas.
+
+        Args:
+            tool_name: Name of the current tool
+            argument_chunk: New chunk of JSON arguments
+
+        Returns:
+            List of state delta events to emit
+        """
+        events: list[StateDeltaEvent] = []
+        if not tool_name or not self.predict_state_config:
+            return events
+
+        self.streaming_tool_args += argument_chunk
+        logger.debug(
+            "Predictive state: accumulated %s chars for tool '%s'",
+            len(self.streaming_tool_args),
+            tool_name,
+        )
+
+        # Try to parse complete JSON first
+        parsed_args = None
+        try:
+            parsed_args = json.loads(self.streaming_tool_args)
+        except json.JSONDecodeError:
+            # Fall back to regex matching for partial JSON
+            events.extend(self._emit_partial_deltas(tool_name))
+
+        if parsed_args:
+            events.extend(self._emit_complete_deltas(tool_name, parsed_args))
+
+        return events
+
+    def _emit_partial_deltas(self, tool_name: str) -> list[StateDeltaEvent]:
+        """Emit deltas from partial JSON using regex matching.
+
+        Args:
+            tool_name: Name of the current tool
+
+        Returns:
+            List of state delta events
+        """
+        events: list[StateDeltaEvent] = []
+
+        for state_key, config in self.predict_state_config.items():
+            if config["tool"] != tool_name:
+                continue
+            tool_arg_name = config["tool_argument"]
+            pattern = rf'"{re.escape(tool_arg_name)}":\s*"([^"]*)'
+            match = re.search(pattern, self.streaming_tool_args)
+
+            if match:
+                partial_value = match.group(1).replace("\\n", "\n").replace('\\"', '"').replace("\\\\", "\\")
+
+                if state_key not in self.last_emitted_state or self.last_emitted_state[state_key] != partial_value:
+                    event = self._create_delta_event(state_key, partial_value)
+                    events.append(event)
+                    self.last_emitted_state[state_key] = partial_value
+                    self.pending_state_updates[state_key] = partial_value
+
+        return events
+
+    def _emit_complete_deltas(
+        self,
+        tool_name: str,
+        parsed_args: dict[str, Any],
+    ) -> list[StateDeltaEvent]:
+        """Emit deltas from complete parsed JSON.
+
+        Args:
+            tool_name: Name of the current tool
+            parsed_args: Fully parsed arguments dict
+
+        Returns:
+            List of state delta events
+        """
+        events: list[StateDeltaEvent] = []
+
+        for state_key, config in self.predict_state_config.items():
+            if config["tool"] != tool_name:
+                continue
+            tool_arg_name = config["tool_argument"]
+
+            if tool_arg_name == "*":
+                state_value = parsed_args
+            elif tool_arg_name in parsed_args:
+                state_value = parsed_args[tool_arg_name]
+            else:
+                continue
+
+            if state_key not in self.last_emitted_state or self.last_emitted_state[state_key] != state_value:
+                event = self._create_delta_event(state_key, state_value)
+                events.append(event)
+                self.last_emitted_state[state_key] = state_value
+                self.pending_state_updates[state_key] = state_value
+
+        return events
+
+    def _create_delta_event(self, state_key: str, value: Any) -> StateDeltaEvent:
+        """Create a state delta event with logging.
+
+        Args:
+            state_key: The state key being updated
+            value: The new value
+
+        Returns:
+            StateDeltaEvent instance
+        """
+        self.state_delta_count += 1
+        if self.state_delta_count % 10 == 1:
+            logger.info(
+                "StateDeltaEvent #%s for '%s': op=replace, path=/%s, value_length=%s",
+                self.state_delta_count,
+                state_key,
+                state_key,
+                len(str(value)),
+            )
+        elif self.state_delta_count % 100 == 0:
+            logger.info(f"StateDeltaEvent #{self.state_delta_count} emitted")
+
+        return StateDeltaEvent(
+            delta=[
+                {
+                    "op": "replace",
+                    "path": f"/{state_key}",
+                    "value": value,
+                }
+            ],
+        )
+
+    def apply_pending_updates(self) -> None:
+        """Apply pending updates to current state and clear them."""
+        for key, value in self.pending_state_updates.items():
+            self.current_state[key] = value
+        self.pending_state_updates.clear()

python/packages/ag-ui/agent_framework_ag_ui/_orchestration/_tooling.py

@@ -0,0 +1,126 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tool handling helpers."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+from agent_framework import BaseChatClient
+from agent_framework._tools import _append_unique_tools  # pyright: ignore[reportPrivateUsage]
+
+if TYPE_CHECKING:
+    from agent_framework import SupportsAgentRun
+
+logger = logging.getLogger(__name__)
+
+
+def _collect_mcp_tool_functions(mcp_tools: list[Any]) -> list[Any]:
+    """Extract functions from connected MCP tools.
+
+    Args:
+        mcp_tools: List of MCP tool instances.
+
+    Returns:
+        Functions from connected MCP tools.
+    """
+    functions: list[Any] = []
+    for mcp_tool in mcp_tools:
+        if getattr(mcp_tool, "is_connected", False) and hasattr(mcp_tool, "functions"):
+            functions.extend(mcp_tool.functions)
+    return functions
+
+
+def collect_server_tools(agent: SupportsAgentRun) -> list[Any]:
+    """Collect server tools from an agent.
+
+    This includes both regular tools from default_options and MCP tools.
+    MCP tools are stored separately for lifecycle management but their
+    functions need to be included for tool execution during approval flows.
+
+    Args:
+        agent: Agent instance to collect tools from. Works with Agent
+            or any agent with default_options and optional mcp_tools attributes.
+
+    Returns:
+        List of tools including both regular tools and connected MCP tool functions.
+    """
+    # Get tools from default_options
+    default_options = getattr(agent, "default_options", None)
+    if default_options is None:
+        return []
+
+    tools_from_agent = default_options.get("tools") if isinstance(default_options, dict) else None
+    server_tools = list(tools_from_agent) if tools_from_agent else []
+
+    # Include functions from connected MCP tools (only available on Agent)
+    mcp_tools = getattr(agent, "mcp_tools", None)
+    if mcp_tools:
+        _append_unique_tools(
+            server_tools,
+            _collect_mcp_tool_functions(mcp_tools),
+            duplicate_error_message="Tool names must be unique. Consider setting `tool_name_prefix` on the MCPTool.",
+        )
+
+    logger.info(f"[TOOLS] Agent has {len(server_tools)} configured tools")
+    for tool in server_tools:
+        tool_name = getattr(tool, "name", "unknown")
+        approval_mode = getattr(tool, "approval_mode", None)
+        logger.info(f"[TOOLS]   - {tool_name}: approval_mode={approval_mode}")
+    return server_tools
+
+
+def register_additional_client_tools(agent: SupportsAgentRun, client_tools: list[Any] | None) -> None:
+    """Register client tools as additional declaration-only tools to avoid server execution.
+
+    Args:
+        agent: Agent instance to register tools on. Works with Agent
+            or any agent with a client attribute.
+        client_tools: List of client tools to register.
+    """
+    if not client_tools:
+        return
+
+    client = getattr(agent, "client", None)
+    if client is None:
+        return
+
+    if isinstance(client, BaseChatClient) and client.function_invocation_configuration is not None:  # type: ignore[attr-defined]
+        client.function_invocation_configuration["additional_tools"] = client_tools  # type: ignore[attr-defined]
+        logger.debug(f"[TOOLS] Registered {len(client_tools)} client tools as additional_tools (declaration-only)")
+
+
+def _has_approval_tools(tools: list[Any]) -> bool:
+    """Check if any tools require approval."""
+    return any(getattr(tool, "approval_mode", None) == "always_require" for tool in tools)
+
+
+def merge_tools(server_tools: list[Any], client_tools: list[Any] | None) -> list[Any] | None:
+    """Combine server and client tools without overriding server metadata.
+
+    IMPORTANT: When server tools have approval_mode="always_require", we MUST return
+    them so they get passed to the streaming response handler. Otherwise, the approval
+    check in _try_execute_function_calls won't find the tool and won't trigger approval.
+    """
+    if not client_tools:
+        # Even without client tools, we must pass server tools if any require approval
+        if server_tools and _has_approval_tools(server_tools):
+            logger.info(
+                f"[TOOLS] No client tools but server has approval tools - "
+                f"passing {len(server_tools)} server tools for approval mode"
+            )
+            return server_tools
+        logger.info("[TOOLS] No client tools - not passing tools= parameter (using agent's configured tools)")
+        return None
+
+    combined_tools = _append_unique_tools(
+        list(server_tools),
+        client_tools,
+        duplicate_error_message="Tool names must be unique.",
+    )
+    logger.info(
+        f"[TOOLS] Passing tools= parameter with {len(combined_tools)} tools "
+        f"({len(server_tools)} server + {len(client_tools)} client)"
+    )
+    return combined_tools

python/packages/ag-ui/agent_framework_ag_ui/_snapshots.py

@@ -0,0 +1,234 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""AG-UI Thread Snapshot storage primitives."""
+
+from __future__ import annotations
+
+import copy
+import logging
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Protocol, TypeAlias, runtime_checkable
+
+if TYPE_CHECKING:
+    from ._types import AGUIRequest
+
+SnapshotScope: TypeAlias = str
+"""Application-defined scope for authorizing access to AG-UI Thread Snapshots."""
+
+AGUIThreadID: TypeAlias = str
+"""AG-UI Thread identifier within a Snapshot Scope."""
+
+SnapshotScopeResolver: TypeAlias = Callable[["AGUIRequest"], str | Awaitable[str]]
+"""Callable that resolves the Snapshot Scope for an AG-UI endpoint request."""
+
+_SnapshotKey: TypeAlias = tuple[SnapshotScope, AGUIThreadID]
+
+DEFAULT_MAX_THREAD_SNAPSHOTS = 1_000
+_SNAPSHOT_SCOPE_INPUT_KEY = "__ag_ui_snapshot_scope"
+_DEFAULT_STATE_INPUT_KEY = "__ag_ui_default_state"
+logger = logging.getLogger(__name__)
+
+
+@dataclass(slots=True)
+class AGUIThreadSnapshot:
+    """Replayable AG-UI Thread state.
+
+    AG-UI Thread Snapshots intentionally contain only data that can be replayed
+    to a UI: message snapshots, optional Shared State, and optional interruption
+    state. They do not include raw events, request metadata, auth claims,
+    diagnostics, traces, or provider responses.
+
+    Attributes:
+        messages: Replayable AG-UI message snapshots.
+        state: Optional AG-UI Shared State snapshot.
+        interrupt: Optional interruption state from ``RUN_FINISHED.outcome.interrupts``.
+    """
+
+    messages: list[dict[str, Any]] = field(default_factory=list)
+    state: dict[str, Any] | None = None
+    interrupt: list[dict[str, Any]] | None = None
+
+
+@runtime_checkable
+class AGUIThreadSnapshotStore(Protocol):
+    """Async store for latest AG-UI Thread Snapshots keyed by scope and thread id."""
+
+    async def save(
+        self,
+        *,
+        scope: SnapshotScope,
+        thread_id: AGUIThreadID,
+        snapshot: AGUIThreadSnapshot,
+    ) -> None:
+        """Save the latest snapshot for an AG-UI Thread within a Snapshot Scope.
+
+        Args:
+            scope: Application-defined Snapshot Scope. This is part of the
+                storage key and must represent the app's authorization boundary.
+            thread_id: AG-UI Thread id within the scope.
+            snapshot: Snapshot to save.
+        """
+        ...
+
+    async def get(
+        self,
+        *,
+        scope: SnapshotScope,
+        thread_id: AGUIThreadID,
+    ) -> AGUIThreadSnapshot | None:
+        """Get the latest snapshot for an AG-UI Thread within a Snapshot Scope.
+
+        Args:
+            scope: Application-defined Snapshot Scope.
+            thread_id: AG-UI Thread id within the scope.
+
+        Returns:
+            The latest snapshot, or ``None`` when no snapshot exists for the key.
+        """
+        ...
+
+    async def delete(
+        self,
+        *,
+        scope: SnapshotScope,
+        thread_id: AGUIThreadID,
+    ) -> bool:
+        """Delete the latest snapshot for an AG-UI Thread within a Snapshot Scope.
+
+        Args:
+            scope: Application-defined Snapshot Scope.
+            thread_id: AG-UI Thread id within the scope.
+
+        Returns:
+            ``True`` when a snapshot was deleted, otherwise ``False``.
+        """
+        ...
+
+    async def clear(self, *, scope: SnapshotScope | None = None) -> None:
+        """Clear saved snapshots.
+
+        Args:
+            scope: Optional Snapshot Scope to clear. When omitted, all in-memory
+                snapshots are cleared.
+        """
+        ...
+
+
+class InMemoryAGUIThreadSnapshotStore:
+    """Bounded memory-only latest snapshot store for local development, demos, and tests.
+
+    This store keeps at most one snapshot per ``(scope, thread_id)`` key. It is
+    process-local and not durable production storage.
+    """
+
+    def __init__(self, *, max_snapshots: int = DEFAULT_MAX_THREAD_SNAPSHOTS) -> None:
+        """Initialize the in-memory snapshot store.
+
+        Keyword Args:
+            max_snapshots: Maximum number of scoped thread snapshots to retain.
+
+        Raises:
+            ValueError: If ``max_snapshots`` is less than 1.
+        """
+        if max_snapshots < 1:
+            raise ValueError("max_snapshots must be greater than 0.")
+        self._max_snapshots = max_snapshots
+        self._snapshots: dict[_SnapshotKey, AGUIThreadSnapshot] = {}
+
+    async def save(
+        self,
+        *,
+        scope: SnapshotScope,
+        thread_id: AGUIThreadID,
+        snapshot: AGUIThreadSnapshot,
+    ) -> None:
+        """Save the latest snapshot for an AG-UI Thread within a Snapshot Scope."""
+        key = self._key(scope=scope, thread_id=thread_id)
+        if key in self._snapshots:
+            del self._snapshots[key]
+        self._snapshots[key] = copy.deepcopy(snapshot)
+        self._evict_oldest()
+
+    async def get(
+        self,
+        *,
+        scope: SnapshotScope,
+        thread_id: AGUIThreadID,
+    ) -> AGUIThreadSnapshot | None:
+        """Get the latest snapshot for an AG-UI Thread within a Snapshot Scope."""
+        snapshot = self._snapshots.get(self._key(scope=scope, thread_id=thread_id))
+        return copy.deepcopy(snapshot) if snapshot is not None else None
+
+    async def delete(
+        self,
+        *,
+        scope: SnapshotScope,
+        thread_id: AGUIThreadID,
+    ) -> bool:
+        """Delete the latest snapshot for an AG-UI Thread within a Snapshot Scope."""
+        key = self._key(scope=scope, thread_id=thread_id)
+        if key not in self._snapshots:
+            return False
+        del self._snapshots[key]
+        return True
+
+    async def clear(self, *, scope: SnapshotScope | None = None) -> None:
+        """Clear saved snapshots, optionally limited to one Snapshot Scope."""
+        if scope is None:
+            self._snapshots.clear()
+            return
+
+        normalized_scope = self._normalize_key_part(scope, "scope")
+        for key in list(self._snapshots):
+            if key[0] == normalized_scope:
+                del self._snapshots[key]
+
+    @classmethod
+    def _key(cls, *, scope: SnapshotScope, thread_id: AGUIThreadID) -> _SnapshotKey:
+        return (
+            cls._normalize_key_part(scope, "scope"),
+            cls._normalize_key_part(thread_id, "thread_id"),
+        )
+
+    @staticmethod
+    def _normalize_key_part(value: str, name: str) -> str:
+        if not isinstance(value, str):
+            raise TypeError(f"{name} must be a string.")
+        if not value:
+            raise ValueError(f"{name} must be a non-empty string.")
+        return value
+
+    def _evict_oldest(self) -> None:
+        while len(self._snapshots) > self._max_snapshots:
+            del self._snapshots[next(iter(self._snapshots))]
+
+
+async def _clear_thread_snapshot_interrupt(
+    *,
+    snapshot_store: AGUIThreadSnapshotStore,
+    scope: SnapshotScope,
+    thread_id: AGUIThreadID,
+    interrupt_ids: set[str] | None = None,
+) -> None:
+    """Clear completed interruption state from the latest replayable thread snapshot."""
+    try:
+        snapshot = await snapshot_store.get(scope=scope, thread_id=thread_id)
+        if snapshot is None or snapshot.interrupt is None:
+            return
+        if interrupt_ids is None:
+            snapshot.interrupt = None
+        else:
+            remaining_interrupts = [
+                interrupt
+                for interrupt in snapshot.interrupt
+                if str(interrupt.get("id") or interrupt.get("interruptId")) not in interrupt_ids
+            ]
+            snapshot.interrupt = remaining_interrupts or None
+        await snapshot_store.save(scope=scope, thread_id=thread_id, snapshot=snapshot)
+    except Exception:
+        logger.exception(
+            "Failed to clear AG-UI Thread Snapshot interrupt for scope=%s thread_id=%s; keeping previous snapshot.",
+            scope,
+            thread_id,
+        )

python/packages/ag-ui/agent_framework_ag_ui/_state.py

@@ -0,0 +1,137 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Deterministic tool-driven AG-UI state updates and display payloads.
+
+Tools wired into the :mod:`agent_framework_ag_ui` endpoint can push a
+deterministic state update or a per-call tool result display payload by
+returning :func:`state_update`. Unlike ``predict_state_config`` — which emits
+``StateDeltaEvent``s optimistically from LLM-predicted tool call arguments —
+``state_update`` runs *after* the tool executes, so AG-UI state and display
+content always reflect the tool's actual return value.
+
+See issue https://github.com/microsoft/agent-framework/issues/3167 for the
+motivating discussion.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Mapping
+from typing import Any
+
+from agent_framework import Content
+
+from ._utils import make_json_safe
+
+__all__ = ["TOOL_RESULT_DISPLAY_KEY", "TOOL_RESULT_STATE_KEY", "state_update"]
+
+
+TOOL_RESULT_STATE_KEY = "__ag_ui_tool_result_state__"
+"""Reserved ``Content.additional_properties`` key used to carry a tool-driven
+state snapshot from a tool return value through to the AG-UI emitter."""
+
+TOOL_RESULT_DISPLAY_KEY = "__ag_ui_tool_result_display__"
+"""Reserved ``Content.additional_properties`` key used to carry UI-only tool result display content from a tool return value through to the AG-UI emitter."""
+
+_UNSET = object()
+
+
+def _serialize_tool_result(value: Any) -> str:  # noqa: ANN401
+    return value if isinstance(value, str) else json.dumps(make_json_safe(value))
+
+
+def state_update(
+    text: str = "",
+    *,
+    state: Mapping[str, Any] | None = None,
+    tool_result: Any = _UNSET,  # noqa: ANN401
+) -> Content:
+    """Build a tool return value that updates AG-UI shared state or display content.
+
+    Return the result of this helper from an agent tool to push a state update
+    or UI-only display payload to AG-UI clients using the actual tool output,
+    rather than LLM-predicted tool arguments.
+
+    When the AG-UI endpoint emits the tool result, it will:
+
+    * Forward ``text`` to the LLM as the normal ``function_result`` content.
+    * Use ``tool_result`` as the ``ToolCallResultEvent.content`` payload shown
+      to AG-UI clients, falling back to ``text`` when no display payload is set.
+    * Merge ``state`` into ``FlowState.current_state``.
+    * Emit a deterministic ``StateSnapshotEvent`` after the ``ToolCallResult``
+      event so frontends observe the updated state deterministically. If
+      predictive state is enabled, a predictive snapshot may be emitted first.
+
+    Example:
+        .. code-block:: python
+
+            from agent_framework import Content, tool
+            from agent_framework_ag_ui import state_update
+
+
+            @tool
+            async def get_weather(city: str) -> Content:
+                data = await _fetch_weather(city)
+                return state_update(
+                    text=f"Weather in {city}: {data['temp']}°C {data['conditions']}",
+                    state={"weather": {"city": city, **data}},
+                )
+
+    Example:
+        .. code-block:: python
+
+            from agent_framework import Content, tool
+            from agent_framework_ag_ui import state_update
+
+
+            @tool
+            async def get_weather(city: str) -> Content:
+                data = await _fetch_weather(city)
+                return state_update(
+                    text=f"{city}: {data['temp']}°C and {data['conditions']}",
+                    tool_result={
+                        "component": "weather-card",
+                        "city": city,
+                        "temperature": data["temp"],
+                        "conditions": data["conditions"],
+                        "humidity": data["humidity"],
+                    },
+                    state={"weather": {"city": city, **data}},
+                )
+
+    Args:
+        text: Text passed back to the LLM as the ``function_result`` content.
+            Defaults to an empty string for tools whose only output is a state
+            update.
+        state: A mapping merged into the AG-UI shared state via JSON-compatible
+            ``dict.update`` semantics. Nested dicts are replaced, not deep-merged.
+        tool_result: JSON-safe payload emitted to AG-UI clients as
+            ``ToolCallResultEvent.content`` for frontend rendering. The LLM
+            still receives ``text``. If ``text`` is empty, the serialized
+            display payload is also used as the LLM-bound text fallback.
+
+    Returns:
+        A ``Content`` object with ``type="text"``. The state payload rides in
+        ``additional_properties`` under :data:`TOOL_RESULT_STATE_KEY`
+        (``"__ag_ui_tool_result_state__"``), and the display payload rides
+        under :data:`TOOL_RESULT_DISPLAY_KEY`
+        (``"__ag_ui_tool_result_display__"``). Both reserved keys are extracted
+        by the AG-UI emitter.
+
+    Raises:
+        TypeError: If ``state`` is not a ``Mapping``.
+    """
+    if state is not None and not isinstance(state, Mapping):
+        raise TypeError(f"state_update() 'state' must be a Mapping, got {type(state).__name__}")
+    additional_properties: dict[str, Any] = {}
+    if state is not None:
+        additional_properties[TOOL_RESULT_STATE_KEY] = dict(state)
+    if tool_result is not _UNSET:
+        display_content = _serialize_tool_result(tool_result)
+        additional_properties[TOOL_RESULT_DISPLAY_KEY] = display_content
+        if not text:
+            text = display_content
+    return Content.from_text(
+        text,
+        additional_properties=additional_properties,
+    )

python/packages/ag-ui/agent_framework_ag_ui/_types.py

@@ -0,0 +1,218 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Type definitions for AG-UI integration."""
+
+from __future__ import annotations
+
+import sys
+from typing import Annotated, Any, Generic
+
+from ag_ui.core import Interrupt, ResumeEntry
+from agent_framework import ChatOptions
+from pydantic import AliasChoices, BaseModel, BeforeValidator, Field
+
+if sys.version_info >= (3, 13):
+    from typing import TypeVar  # pragma: no cover
+else:
+    from typing_extensions import TypeVar  # pragma: no cover
+if sys.version_info >= (3, 11):
+    from typing import TypedDict  # pragma: no cover
+else:
+    from typing_extensions import TypedDict  # pragma: no cover
+
+
+AGUIChatOptionsT = TypeVar("AGUIChatOptionsT", bound=TypedDict, default="AGUIChatOptions", covariant=True)  # type: ignore[valid-type]
+ResponseModelT = TypeVar("ResponseModelT", bound=BaseModel | None, default=None)
+
+
+def _coerce_legacy_resume_entry(value: Any) -> Any:  # noqa: ANN401
+    if not isinstance(value, dict):
+        return value
+
+    interrupt_id = value.get("interruptId") or value.get("interrupt_id") or value.get("id") or value.get("toolCallId")
+    if not interrupt_id:
+        return value
+
+    if "payload" in value:
+        payload = value.get("payload")
+    elif "value" in value:
+        payload = value.get("value")
+    elif "response" in value:
+        payload = value.get("response")
+    else:
+        payload = {
+            key: item
+            for key, item in value.items()
+            if key not in {"id", "interruptId", "interrupt_id", "toolCallId", "type", "status"}
+        }
+
+    entry: dict[str, Any] = {"interruptId": str(interrupt_id), "status": value.get("status", "resolved")}
+    if payload is not None:
+        entry["payload"] = payload
+    return entry
+
+
+def _coerce_legacy_resume(value: Any) -> Any:  # noqa: ANN401
+    if value is None:
+        return value
+    if isinstance(value, dict):
+        if "interrupts" in value:
+            value = value["interrupts"]
+        elif "interrupt" in value:
+            value = value["interrupt"]
+        elif any(key in value for key in ("interruptId", "interrupt_id", "id", "toolCallId")):
+            value = [value]
+        else:
+            return value
+    if not isinstance(value, list):
+        return value
+    return [_coerce_legacy_resume_entry(entry) for entry in value]
+
+
+class PredictStateConfig(TypedDict):
+    """Configuration for predictive state updates."""
+
+    state_key: str
+    tool: str
+    tool_argument: str | None
+
+
+class RunMetadata(TypedDict):
+    """Metadata for agent run."""
+
+    run_id: str
+    thread_id: str
+    predict_state: list[PredictStateConfig] | None
+
+
+class AgentState(TypedDict):
+    """Base state for AG-UI agents."""
+
+    messages: list[Any] | None
+
+
+class AGUIRequest(BaseModel):
+    """Request model for AG-UI endpoints."""
+
+    messages: list[dict[str, Any]] = Field(
+        ...,
+        description="AG-UI format messages array",
+    )
+    run_id: str | None = Field(
+        None,
+        validation_alias=AliasChoices("run_id", "runId"),
+        description="Optional run identifier for tracking",
+    )
+    thread_id: str | None = Field(
+        None,
+        validation_alias=AliasChoices("thread_id", "threadId"),
+        description="Optional thread identifier for conversation context",
+    )
+    state: dict[str, Any] | None = Field(
+        None,
+        description="Optional shared state for agentic generative UI",
+    )
+    tools: list[dict[str, Any]] | None = Field(
+        None,
+        description="Client-side tools to advertise to the LLM",
+    )
+    context: list[dict[str, Any]] | None = Field(
+        None,
+        description="List of context objects provided to the agent",
+    )
+    forwarded_props: dict[str, Any] | None = Field(
+        None,
+        validation_alias=AliasChoices("forwarded_props", "forwardedProps"),
+        description="Additional properties forwarded to the agent",
+    )
+    parent_run_id: str | None = Field(
+        None,
+        validation_alias=AliasChoices("parent_run_id", "parentRunId"),
+        description="ID of the run that spawned this run",
+    )
+    available_interrupts: list[Interrupt] | None = Field(
+        None,
+        validation_alias=AliasChoices("availableInterrupts", "available_interrupts"),
+        description="Canonical AG-UI interrupts that can be resumed by the server",
+    )
+    resume: Annotated[list[ResumeEntry], BeforeValidator(_coerce_legacy_resume)] | None = Field(
+        None,
+        description="Resume payload for continuing interrupted runs",
+    )
+
+
+# region AG-UI Chat Options TypedDict
+
+
+class AGUIChatOptions(ChatOptions[ResponseModelT], Generic[ResponseModelT], total=False):
+    """AG-UI protocol-specific chat options dict.
+
+    Extends base ChatOptions for the AG-UI (Agent-UI) protocol.
+    AG-UI is a streaming protocol for connecting AI agents to user interfaces.
+    Options are forwarded to the remote AG-UI server.
+
+    See: https://github.com/ag-ui/ag-ui-protocol
+
+    Keys:
+        # Inherited from ChatOptions (forwarded to remote server):
+        model: The model identifier (forwarded as-is to server).
+        temperature: Sampling temperature.
+        top_p: Nucleus sampling parameter.
+        max_tokens: Maximum tokens to generate.
+        stop: Stop sequences.
+        tools: List of tools - sent to server so LLM knows about client tools.
+            Server executes its own tools; client tools execute locally via
+            function invocation middleware.
+        tool_choice: How the model should use tools.
+        metadata: Metadata dict containing thread_id for conversation continuity.
+
+        # Options with limited support (depends on remote server):
+        frequency_penalty: Forwarded if remote server supports it.
+        presence_penalty: Forwarded if remote server supports it.
+        seed: Forwarded if remote server supports it.
+        response_format: Forwarded if remote server supports it.
+        logit_bias: Forwarded if remote server supports it.
+        user: Forwarded if remote server supports it.
+
+        # Options not typically used in AG-UI:
+        store: Not applicable for AG-UI protocol.
+        allow_multiple_tool_calls: Handled by underlying server.
+
+        # AG-UI-specific options:
+        forward_props: Additional properties to forward to the AG-UI server.
+            Useful for passing custom parameters to specific server implementations.
+        context: Shared context/state to send to the server.
+
+    Note:
+        AG-UI is a protocol bridge - actual option support depends on the
+        remote server implementation. The client sends all options to the
+        server, which decides how to handle them.
+
+        Thread ID management:
+        - Pass ``thread_id`` in ``metadata`` to maintain conversation continuity
+        - If not provided, a new thread ID is auto-generated
+    """
+
+    # AG-UI-specific options
+    forward_props: dict[str, Any]
+    """Additional properties to forward to the AG-UI server."""
+
+    context: dict[str, Any]
+    """Shared context/state to send to the server."""
+
+    available_interrupts: list[Interrupt]
+    """Canonical AG-UI interrupt descriptors available for resumption."""
+
+    resume: list[ResumeEntry]
+    """Canonical AG-UI resume entries to continue a paused run."""
+
+    # ChatOptions fields not applicable for AG-UI
+    store: None  # type: ignore[misc]
+    """Not applicable for AG-UI protocol."""
+
+
+AGUI_OPTION_TRANSLATIONS: dict[str, str] = {}
+"""Maps ChatOptions keys to AG-UI parameter names (protocol uses standard names)."""
+
+
+# endregion

python/packages/ag-ui/agent_framework_ag_ui/_utils.py

@@ -0,0 +1,292 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Utility functions for AG-UI integration."""
+
+from __future__ import annotations
+
+import copy
+import json
+import uuid
+from collections.abc import Callable, MutableMapping, Sequence
+from dataclasses import asdict, is_dataclass
+from datetime import date, datetime
+from typing import Any
+
+from agent_framework import AgentResponseUpdate, ChatResponseUpdate, FunctionTool
+
+# Role mapping constants
+AGUI_TO_FRAMEWORK_ROLE: dict[str, str] = {
+    "user": "user",
+    "assistant": "assistant",
+    "system": "system",
+}
+
+FRAMEWORK_TO_AGUI_ROLE: dict[str, str] = {
+    "user": "user",
+    "assistant": "assistant",
+    "system": "system",
+}
+
+ALLOWED_AGUI_ROLES: set[str] = {"user", "assistant", "system", "tool", "reasoning"}
+
+
+def generate_event_id() -> str:
+    """Generate a unique event ID."""
+    return str(uuid.uuid4())
+
+
+def safe_json_parse(value: Any) -> dict[str, Any] | None:
+    """Safely parse a value as JSON dict.
+
+    Args:
+        value: String or dict to parse
+
+    Returns:
+        Parsed dict or None if parsing fails
+    """
+    if isinstance(value, dict):
+        return value
+    if isinstance(value, str):
+        try:
+            parsed = json.loads(value)
+            if isinstance(parsed, dict):
+                return parsed
+        except json.JSONDecodeError:
+            pass
+    return None
+
+
+def canonical_function_arguments(function_call: Any) -> str | None:
+    """Return a stable representation of function-call arguments."""
+    if function_call is None:
+        return None
+
+    try:
+        parsed_arguments = function_call.parse_arguments()
+    except Exception:
+        parsed_arguments = getattr(function_call, "arguments", None)
+
+    if parsed_arguments is None:
+        parsed_arguments = {}
+
+    return json.dumps(make_json_safe(parsed_arguments), sort_keys=True, separators=(",", ":"))
+
+
+def get_role_value(message: Any) -> str:
+    """Extract role string from a message object.
+
+    Handles both enum roles (with .value) and string roles.
+
+    Args:
+        message: Message object with role attribute
+
+    Returns:
+        Role as lowercase string, or empty string if not found
+    """
+    role = getattr(message, "role", None)
+    if role is None:
+        return ""
+    if hasattr(role, "value"):
+        return str(role.value)
+    return str(role)
+
+
+def normalize_agui_role(raw_role: Any) -> str:
+    """Normalize an AG-UI role to a standard role string.
+
+    Args:
+        raw_role: Raw role value from AG-UI message
+
+    Returns:
+        Normalized role string (user, assistant, system, tool, or reasoning)
+    """
+    if not isinstance(raw_role, str):
+        return "user"
+    role = raw_role.lower()
+    if role == "developer":
+        return "system"
+    if role in ALLOWED_AGUI_ROLES:
+        return role
+    return "user"
+
+
+def extract_state_from_tool_args(
+    args: dict[str, Any] | None,
+    tool_arg_name: str,
+) -> Any:
+    """Extract state value from tool arguments based on config.
+
+    Args:
+        args: Parsed tool arguments dict
+        tool_arg_name: Name of the argument to extract, or "*" for entire args
+
+    Returns:
+        Extracted state value, or None if not found
+    """
+    if not args:
+        return None
+    if tool_arg_name == "*":
+        return args
+    return args.get(tool_arg_name)
+
+
+def merge_state(current: dict[str, Any], update: dict[str, Any]) -> dict[str, Any]:
+    """Merge state updates.
+
+    Args:
+        current: Current state dictionary
+        update: Update to apply
+
+    Returns:
+        Merged state
+    """
+    result = copy.deepcopy(current)
+    result.update(update)
+    return result
+
+
+def make_json_safe(obj: Any) -> Any:  # noqa: ANN401
+    """Make an object JSON serializable.
+
+    Args:
+        obj: Object to make JSON safe
+
+    Returns:
+        JSON-serializable version of the object
+    """
+    if obj is None or isinstance(obj, (str, int, float, bool)):
+        return obj
+    if isinstance(obj, (datetime, date)):
+        return obj.isoformat()
+    if is_dataclass(obj):
+        # asdict may return nested non-dataclass objects, so recursively make them safe
+        return make_json_safe(asdict(obj))  # type: ignore[arg-type]
+    if hasattr(obj, "model_dump"):
+        return make_json_safe(obj.model_dump())
+    if hasattr(obj, "to_dict"):
+        return make_json_safe(obj.to_dict())
+    if hasattr(obj, "dict"):
+        return make_json_safe(obj.dict())
+    if hasattr(obj, "__dict__"):
+        return {key: make_json_safe(value) for key, value in vars(obj).items()}  # type: ignore[misc]
+    if isinstance(obj, (list, tuple)):
+        return [make_json_safe(item) for item in obj]  # type: ignore[misc]
+    if isinstance(obj, dict):
+        return {key: make_json_safe(value) for key, value in obj.items()}  # type: ignore[misc]
+    return str(obj)
+
+
+def convert_agui_tools_to_agent_framework(
+    agui_tools: list[dict[str, Any]] | None,
+) -> list[FunctionTool] | None:
+    """Convert AG-UI tool definitions to Agent Framework FunctionTool declarations.
+
+    Creates declaration-only FunctionTool instances (no executable implementation).
+    These are used to tell the LLM about available tools. The actual execution
+    happens on the client side via function invocation mixin.
+
+    CRITICAL: These tools MUST have func=None so that declaration_only returns True.
+    This prevents the server from trying to execute client-side tools.
+
+    Args:
+        agui_tools: List of AG-UI tool definitions with name, description, parameters
+
+    Returns:
+        List of FunctionTool declarations, or None if no tools provided
+    """
+    if not agui_tools:
+        return None
+
+    result: list[FunctionTool] = []
+    for tool_def in agui_tools:
+        # Create declaration-only FunctionTool (func=None means no implementation)
+        # When func=None, the declaration_only property returns True,
+        # which tells the function invocation mixin to return the function call
+        # without executing it (so it can be sent back to the client)
+        func: FunctionTool = FunctionTool(
+            name=tool_def.get("name", ""),
+            description=tool_def.get("description", ""),
+            func=None,  # CRITICAL: Makes declaration_only=True
+            input_model=tool_def.get("parameters", {}),
+        )
+        result.append(func)
+
+    return result
+
+
+def convert_tools_to_agui_format(
+    tools: (
+        FunctionTool
+        | Callable[..., Any]
+        | MutableMapping[str, Any]
+        | Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]]
+        | None
+    ),
+) -> list[dict[str, Any]] | None:
+    """Convert tools to AG-UI format.
+
+    This sends only the metadata (name, description, JSON schema) to the server.
+    The actual executable implementation stays on the client side.
+    The function invocation mixin handles client-side execution when
+    the server requests a function.
+
+    Args:
+        tools: Tools to convert (single tool or sequence of tools)
+
+    Returns:
+        List of tool specifications in AG-UI format, or None if no tools provided
+    """
+    if not tools:
+        return None
+
+    # Normalize to list
+    if not isinstance(tools, list):
+        tool_list: list[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]] = [tools]  # type: ignore[list-item]
+    else:
+        tool_list = tools  # type: ignore[assignment]
+
+    results: list[dict[str, Any]] = []
+
+    for tool_item in tool_list:
+        if isinstance(tool_item, dict):
+            # Already in dict format, pass through
+            results.append(tool_item)  # type: ignore[arg-type]
+        elif isinstance(tool_item, FunctionTool):
+            # Convert FunctionTool to AG-UI tool format
+            results.append(
+                {
+                    "name": tool_item.name,
+                    "description": tool_item.description,
+                    "parameters": tool_item.parameters(),
+                }
+            )
+        elif callable(tool_item):
+            # Convert callable to FunctionTool first, then to AG-UI format
+            from agent_framework import tool
+
+            ai_func = tool(tool_item)
+            results.append(
+                {
+                    "name": ai_func.name,
+                    "description": ai_func.description,
+                    "parameters": ai_func.parameters(),
+                }
+            )
+        # Note: dict-based hosted tools (CodeInterpreter, WebSearch, etc.) are passed through
+        # as-is in the first branch. Non-FunctionTool, non-dict items are skipped.
+
+    return results if results else None
+
+
+def get_conversation_id_from_update(update: AgentResponseUpdate) -> str | None:
+    """Extract conversation ID from AgentResponseUpdate metadata.
+
+    Args:
+        update: AgentRunResponseUpdate instance
+    Returns:
+        Conversation ID if present, else None
+
+    """
+    if isinstance(update.raw_representation, ChatResponseUpdate):
+        return update.raw_representation.conversation_id
+    return None

python/packages/ag-ui/agent_framework_ag_ui/_workflow.py

@@ -0,0 +1,404 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Workflow wrapper for AG-UI protocol compatibility."""
+
+from __future__ import annotations
+
+import copy
+import logging
+import uuid
+from collections.abc import AsyncGenerator, Callable
+from typing import Any, cast
+
+from ag_ui.core import (
+    BaseEvent,
+    MessagesSnapshotEvent,
+    RunErrorEvent,
+    RunFinishedEvent,
+    RunStartedEvent,
+    StateSnapshotEvent,
+    TextMessageContentEvent,
+    TextMessageEndEvent,
+    TextMessageStartEvent,
+    ToolCallArgsEvent,
+    ToolCallResultEvent,
+    ToolCallStartEvent,
+)
+from agent_framework import Workflow
+
+from ._message_adapters import agui_messages_to_snapshot_format
+from ._run_common import (
+    _build_run_finished_event,
+    _extract_resume_payload,
+    _normalize_resume_interrupts,
+    _reconstruct_messages_from_thread_snapshot,
+)
+from ._snapshots import (
+    _DEFAULT_STATE_INPUT_KEY,
+    _SNAPSHOT_SCOPE_INPUT_KEY,
+    AGUIThreadSnapshot,
+    AGUIThreadSnapshotStore,
+    _clear_thread_snapshot_interrupt,
+)
+from ._utils import generate_event_id, make_json_safe
+from ._workflow_run import run_workflow_stream
+
+logger = logging.getLogger(__name__)
+
+WorkflowFactory = Callable[[str], Workflow]
+
+
+def _cancelled_resume_interrupt_ids(resume_payload: Any) -> set[str]:
+    """Return cancelled interrupt ids from a resume payload."""
+    return {
+        str(interrupt["id"])
+        for interrupt in _normalize_resume_interrupts(resume_payload)
+        if interrupt.get("status") == "cancelled"
+    }
+
+
+def _event_messages_to_snapshot_dicts(messages: list[Any]) -> list[dict[str, Any]]:
+    """Convert AG-UI message event models to plain snapshot dictionaries."""
+    safe_messages = make_json_safe(messages)
+    if not isinstance(safe_messages, list):
+        return []
+    return [cast(dict[str, Any], message) for message in safe_messages if isinstance(message, dict)]
+
+
+class _WorkflowSnapshotBuilder:
+    """Capture replayable workflow protocol output without retaining raw events."""
+
+    def __init__(self, raw_messages: list[dict[str, Any]]) -> None:
+        self._synthesized_messages = agui_messages_to_snapshot_format(raw_messages)
+        self._emitted_messages: list[dict[str, Any]] | None = None
+        self._open_text_message: dict[str, Any] | None = None
+        self._tool_call_message: dict[str, Any] | None = None
+        self._tool_calls_by_id: dict[str, dict[str, Any]] = {}
+        self.state: dict[str, Any] | None = None
+        self.interrupt: list[dict[str, Any]] | None = None
+
+    def observe(self, event: BaseEvent) -> None:
+        """Fold one replayable AG-UI event into the latest snapshot state."""
+        if isinstance(event, StateSnapshotEvent):
+            state = make_json_safe(event.snapshot)
+            if isinstance(state, dict):
+                self.state = cast(dict[str, Any], state)
+            return
+
+        if isinstance(event, MessagesSnapshotEvent):
+            self._emitted_messages = _event_messages_to_snapshot_dicts(list(event.messages))
+            return
+
+        if isinstance(event, RunFinishedEvent):
+            outcome = getattr(event, "outcome", None)
+            interrupt = (
+                make_json_safe(getattr(outcome, "interrupts", None))
+                if getattr(outcome, "type", None) == "interrupt"
+                else None
+            )
+            if isinstance(interrupt, list):
+                self.interrupt = [cast(dict[str, Any], item) for item in interrupt if isinstance(item, dict)]
+            return
+
+        if self._emitted_messages is not None:
+            return
+
+        if isinstance(event, TextMessageStartEvent):
+            self._observe_text_start(event)
+        elif isinstance(event, TextMessageContentEvent):
+            self._observe_text_content(event)
+        elif isinstance(event, TextMessageEndEvent):
+            self._observe_text_end(event)
+        elif isinstance(event, ToolCallStartEvent):
+            self._observe_tool_call_start(event)
+        elif isinstance(event, ToolCallArgsEvent):
+            self._observe_tool_call_args(event)
+        elif isinstance(event, ToolCallResultEvent):
+            self._observe_tool_call_result(event)
+
+    def build(self) -> AGUIThreadSnapshot:
+        """Return the replayable thread snapshot."""
+        self._flush_open_text_message()
+        messages = self._emitted_messages if self._emitted_messages is not None else self._synthesized_messages
+        return AGUIThreadSnapshot(messages=messages, state=self.state, interrupt=self.interrupt)
+
+    def _observe_text_start(self, event: TextMessageStartEvent) -> None:
+        if self._open_text_message is not None and self._open_text_message.get("id") != event.message_id:
+            self._flush_open_text_message()
+        self._open_text_message = {"id": event.message_id, "role": event.role, "content": ""}
+
+    def _observe_text_content(self, event: TextMessageContentEvent) -> None:
+        if self._open_text_message is None or self._open_text_message.get("id") != event.message_id:
+            self._open_text_message = {"id": event.message_id, "role": "assistant", "content": ""}
+        self._open_text_message["content"] = f"{self._open_text_message.get('content', '')}{event.delta}"
+
+    def _observe_text_end(self, event: TextMessageEndEvent) -> None:
+        if self._open_text_message is None or self._open_text_message.get("id") != event.message_id:
+            return
+        self._flush_open_text_message()
+
+    def _observe_tool_call_start(self, event: ToolCallStartEvent) -> None:
+        parent_message_id = event.parent_message_id
+        if (
+            self._open_text_message is not None
+            and parent_message_id is not None
+            and self._open_text_message.get("id") == parent_message_id
+            and self._open_text_message.get("content")
+        ):
+            self._open_text_message["id"] = generate_event_id()
+        self._flush_open_text_message()
+        if self._tool_call_message is None or (
+            parent_message_id is not None and self._tool_call_message.get("id") != parent_message_id
+        ):
+            self._tool_call_message = {
+                "id": parent_message_id or generate_event_id(),
+                "role": "assistant",
+                "tool_calls": [],
+            }
+            self._synthesized_messages.append(self._tool_call_message)
+
+        tool_call = {
+            "id": event.tool_call_id,
+            "type": "function",
+            "function": {"name": event.tool_call_name, "arguments": ""},
+        }
+        cast(list[dict[str, Any]], self._tool_call_message["tool_calls"]).append(tool_call)
+        self._tool_calls_by_id[event.tool_call_id] = tool_call
+
+    def _observe_tool_call_args(self, event: ToolCallArgsEvent) -> None:
+        tool_call = self._tool_calls_by_id.get(event.tool_call_id)
+        if tool_call is None:
+            return
+        function_payload = cast(dict[str, Any], tool_call["function"])
+        function_payload["arguments"] = f"{function_payload.get('arguments', '')}{event.delta}"
+
+    def _observe_tool_call_result(self, event: ToolCallResultEvent) -> None:
+        self._synthesized_messages.append(
+            {
+                "id": event.message_id,
+                "role": "tool",
+                "toolCallId": event.tool_call_id,
+                "content": event.content,
+            }
+        )
+        # A result closes the current tool-call group; later tool calls start a new
+        # assistant message so replayed transcripts keep results adjacent to their
+        # tool_calls message, which provider APIs require.
+        self._tool_call_message = None
+
+    def _flush_open_text_message(self) -> None:
+        if self._open_text_message is None:
+            return
+        if self._open_text_message.get("content"):
+            self._synthesized_messages.append(self._open_text_message)
+            # Text between tool calls closes the current tool-call group as well.
+            self._tool_call_message = None
+        self._open_text_message = None
+
+
+async def _hydrate_workflow_thread_snapshot(
+    *,
+    snapshot_store: AGUIThreadSnapshotStore,
+    scope: str,
+    thread_id: str,
+    run_id: str,
+) -> AsyncGenerator[BaseEvent]:
+    """Replay the latest stored workflow AG-UI Thread Snapshot without invoking the workflow."""
+    yield RunStartedEvent(run_id=run_id, thread_id=thread_id)
+    snapshot = await snapshot_store.get(scope=scope, thread_id=thread_id)
+    if snapshot is None:
+        yield _build_run_finished_event(run_id=run_id, thread_id=thread_id)
+        return
+
+    if snapshot.state is not None:
+        yield StateSnapshotEvent(snapshot=snapshot.state)
+    if snapshot.messages:
+        yield MessagesSnapshotEvent(messages=snapshot.messages)  # type: ignore[arg-type]
+    yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=snapshot.interrupt)
+
+
+class AgentFrameworkWorkflow:
+    """Base AG-UI workflow wrapper.
+
+    Can wrap a native ``Workflow`` or be subclassed for custom ``run`` behavior.
+    """
+
+    def __init__(
+        self,
+        workflow: Workflow | None = None,
+        *,
+        workflow_factory: WorkflowFactory | None = None,
+        name: str | None = None,
+        description: str | None = None,
+        snapshot_store: AGUIThreadSnapshotStore | None = None,
+    ) -> None:
+        """Initialize the AG-UI workflow wrapper.
+
+        Args:
+            workflow: Optional workflow instance to expose.
+            workflow_factory: Optional factory for thread-scoped workflow instances.
+            name: Optional workflow name.
+            description: Optional workflow description.
+            snapshot_store: Optional AG-UI Thread Snapshot store. Snapshot persistence remains inactive unless
+                endpoint setup also provides an explicit Snapshot Scope resolver.
+        """
+        if workflow is not None and workflow_factory is not None:
+            raise ValueError("Pass either workflow= or workflow_factory=, not both.")
+
+        self.workflow = workflow
+        self._workflow_factory = workflow_factory
+        # Cache keyed by (snapshot_scope, thread_id): the Snapshot Scope is the
+        # authorization boundary, so the same thread id under different scopes
+        # must never share an in-memory workflow instance.
+        self._workflow_by_thread: dict[tuple[str | None, str], Workflow] = {}
+        self.name = name if name is not None else getattr(workflow, "name", "workflow")
+        self.description = description if description is not None else getattr(workflow, "description", "")
+        self.snapshot_store = snapshot_store
+
+    @staticmethod
+    def _thread_id_from_input(input_data: dict[str, Any]) -> str:
+        """Resolve a stable thread id from AG-UI input payload."""
+        thread_id = input_data.get("thread_id") or input_data.get("threadId")
+        if thread_id is not None:
+            return str(thread_id)
+        return str(uuid.uuid4())
+
+    def _resolve_workflow(self, thread_id: str, snapshot_scope: str | None = None) -> Workflow:
+        """Get the workflow instance for the current run."""
+        if self.workflow is not None:
+            return self.workflow
+
+        if self._workflow_factory is None:
+            raise NotImplementedError("No workflow is attached. Override run or pass workflow=/workflow_factory=.")
+
+        cache_key = (snapshot_scope, thread_id)
+        workflow = self._workflow_by_thread.get(cache_key)
+        if workflow is None:
+            workflow = self._workflow_factory(thread_id)
+            if not isinstance(workflow, Workflow):
+                raise TypeError("workflow_factory must return a Workflow instance.")
+            self._workflow_by_thread[cache_key] = workflow
+        return workflow
+
+    def clear_thread_workflow(self, thread_id: str, snapshot_scope: str | None = None) -> None:
+        """Drop cached workflow instances for a thread, optionally limited to one Snapshot Scope."""
+        if snapshot_scope is not None:
+            self._workflow_by_thread.pop((snapshot_scope, thread_id), None)
+            return
+        for key in [key for key in self._workflow_by_thread if key[1] == thread_id]:
+            del self._workflow_by_thread[key]
+
+    def clear_workflow_cache(self) -> None:
+        """Drop all cached thread workflow instances."""
+        self._workflow_by_thread.clear()
+
+    async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]:
+        """Run the wrapped workflow and yield AG-UI events.
+
+        Subclasses may override this to provide custom AG-UI streams.
+        """
+        thread_id = self._thread_id_from_input(input_data)
+        run_id = str(input_data.get("run_id") or input_data.get("runId") or uuid.uuid4())
+        snapshot_scope = cast(str | None, input_data.get(_SNAPSHOT_SCOPE_INPUT_KEY))
+        raw_messages = list(cast(list[dict[str, Any]], input_data.get("messages", []) or []))
+        resume_payload = _extract_resume_payload(input_data)
+        snapshot_store = self.snapshot_store
+
+        if snapshot_store is not None and snapshot_scope is not None and not raw_messages and resume_payload is None:
+            async for event in _hydrate_workflow_thread_snapshot(
+                snapshot_store=snapshot_store,
+                scope=snapshot_scope,
+                thread_id=thread_id,
+                run_id=run_id,
+            ):
+                yield event
+            return
+
+        # Load the stored snapshot for follow-up turns so the workflow runs with the
+        # full persisted thread history instead of just the latest request messages.
+        stored_snapshot: AGUIThreadSnapshot | None = None
+        if snapshot_store is not None and snapshot_scope is not None:
+            stored_snapshot = await snapshot_store.get(scope=snapshot_scope, thread_id=thread_id)
+            if stored_snapshot is not None and resume_payload is None:
+                raw_messages = _reconstruct_messages_from_thread_snapshot(
+                    stored_messages=stored_snapshot.messages,
+                    incoming_messages=raw_messages,
+                    stored_interrupt=stored_snapshot.interrupt,
+                )
+                input_data["messages"] = raw_messages
+
+        # Merge stored state with request overrides, then fill endpoint-deferred
+        # defaults only for keys missing from both.
+        request_state = input_data.get("state")
+        deferred_default_state = cast(dict[str, Any] | None, input_data.get(_DEFAULT_STATE_INPUT_KEY))
+        effective_state: dict[str, Any] = {}
+        if stored_snapshot is not None and stored_snapshot.state is not None:
+            effective_state.update(stored_snapshot.state)
+        if isinstance(request_state, dict):
+            effective_state.update(cast(dict[str, Any], request_state))
+        if deferred_default_state:
+            for key, value in deferred_default_state.items():
+                if key not in effective_state:
+                    effective_state[key] = copy.deepcopy(value)
+        if effective_state:
+            input_data["state"] = effective_state
+
+        workflow = self._resolve_workflow(thread_id, snapshot_scope)
+        builder_seed_messages = raw_messages
+        if resume_payload is not None and stored_snapshot is not None:
+            # Resume requests carry only the synthesized interrupt response, so seed
+            # the builder with stored history to avoid persisting a truncated thread.
+            builder_seed_messages = [
+                copy.deepcopy(message) for message in stored_snapshot.messages
+            ] + builder_seed_messages
+        snapshot_builder = (
+            _WorkflowSnapshotBuilder(builder_seed_messages)
+            if snapshot_store is not None and snapshot_scope is not None
+            else None
+        )
+        if snapshot_builder is not None and effective_state:
+            # Seed builder state so a run that emits no StateSnapshotEvent still
+            # persists the latest known Shared State instead of dropping it.
+            state_snapshot = make_json_safe(effective_state)
+            if isinstance(state_snapshot, dict):
+                snapshot_builder.state = cast(dict[str, Any], state_snapshot)
+        run_error_emitted = False
+        async for event in run_workflow_stream(input_data, workflow):
+            if snapshot_builder is not None:
+                snapshot_builder.observe(event)
+            if isinstance(event, RunErrorEvent):
+                run_error_emitted = True
+                if (
+                    getattr(event, "code", None) == "WORKFLOW_RESUME_CANCELLED"
+                    and snapshot_store is not None
+                    and snapshot_scope is not None
+                ):
+                    await _clear_thread_snapshot_interrupt(
+                        snapshot_store=snapshot_store,
+                        scope=snapshot_scope,
+                        thread_id=thread_id,
+                        interrupt_ids=_cancelled_resume_interrupt_ids(resume_payload),
+                    )
+            yield event
+
+        if (
+            snapshot_builder is not None
+            and not run_error_emitted
+            and snapshot_store is not None
+            and snapshot_scope is not None
+        ):
+            try:
+                await snapshot_store.save(
+                    scope=snapshot_scope,
+                    thread_id=thread_id,
+                    snapshot=snapshot_builder.build(),
+                )
+            except Exception:
+                # RUN_FINISHED has already been yielded; a store failure must not
+                # surface as a second terminal RUN_ERROR event. The previous
+                # snapshot stays available for hydration.
+                logger.exception(
+                    "Failed to save AG-UI Thread Snapshot for scope=%s thread_id=%s; keeping previous snapshot.",
+                    snapshot_scope,
+                    thread_id,
+                )

python/packages/ag-ui/agent_framework_ag_ui/py.typed

@@ -0,0 +1 @@
+# Marker file for PEP 561

python/packages/ag-ui/agent_framework_ag_ui_examples/.env.example

@@ -0,0 +1,3 @@
+AZURE_OPENAI_ENDPOINT=https://your-resource.openai.azure.com/
+AZURE_OPENAI_API_KEY=your-api-key-here
+PORT=8000

python/packages/ag-ui/agent_framework_ag_ui_examples/.vscode/settings.json

@@ -0,0 +1,5 @@
+{
+    "python.analysis.extraPaths": [
+        "${workspaceFolder}/packages/ag-ui/examples"
+    ]
+}

python/packages/ag-ui/agent_framework_ag_ui_examples/README.md

@@ -0,0 +1,360 @@
+# Agent Framework AG-UI Integration
+
+AG-UI protocol integration for Agent Framework, enabling seamless integration with AG-UI's web interface and streaming protocol.
+
+## Installation
+
+```bash
+pip install agent-framework-ag-ui
+```
+
+## Quick Start
+
+### Using Example Agents with Any Chat Client
+
+All example agents are factory functions that accept any `SupportsChatGetResponse`-compatible chat client:
+
+```python
+from fastapi import FastAPI
+from agent_framework.openai import OpenAIChatCompletionClient
+from agent_framework.openai import OpenAIChatClient
+from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
+from agent_framework_ag_ui_examples.agents import simple_agent, weather_agent
+
+app = FastAPI()
+
+# Option 1: Use Azure OpenAI
+azure_client = OpenAIChatCompletionClient(model="gpt-4")
+add_agent_framework_fastapi_endpoint(app, simple_agent(azure_client), "/chat")
+
+# Option 2: Use OpenAI
+openai_client = OpenAIChatClient(model="gpt-4o")
+add_agent_framework_fastapi_endpoint(app, weather_agent(openai_client), "/weather")
+
+# Run with: uvicorn main:app --reload
+```
+
+### Creating Your Own Agent
+
+```python
+from fastapi import FastAPI
+from agent_framework import Agent
+from agent_framework.openai import OpenAIChatCompletionClient
+from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
+
+# Create your agent
+agent = Agent(
+    name="my_agent",
+    instructions="You are a helpful assistant.",
+    client=OpenAIChatCompletionClient(model="gpt-4o"),
+)
+
+# Create FastAPI app and add AG-UI endpoint
+app = FastAPI()
+add_agent_framework_fastapi_endpoint(app, agent, "/agent")
+
+# Run with: uvicorn main:app --reload
+```
+
+## Features
+
+This integration supports all 7 AG-UI features:
+
+1. **Agentic Chat**: Basic streaming chat with tool calling support
+2. **Backend Tool Rendering**: Tools executed on backend with results streamed via ToolCallResultEvent
+3. **Human in the Loop**: Function approval requests for user confirmation before tool execution
+4. **Agentic Generative UI**: Async tools for long-running operations with progress updates
+5. **Tool-based Generative UI**: Custom UI components rendered on frontend based on tool calls
+6. **Shared State**: Bidirectional state sync using StateSnapshotEvent and StateDeltaEvent
+7. **Predictive State Updates**: Stream tool arguments as optimistic state updates during execution
+
+## Examples
+
+All example agents are implemented as **factory functions** that accept any chat client implementing `SupportsChatGetResponse`. This provides maximum flexibility to use Azure OpenAI, OpenAI, Anthropic, or any custom chat client implementation.
+
+### Available Example Agents
+
+Complete examples for all AG-UI features are available:
+
+- `simple_agent(client)` - Basic agentic chat (Feature 1)
+- `weather_agent(client)` - Backend tool rendering (Feature 2)
+- `human_in_the_loop_agent(client)` - Human-in-the-loop with step customization (Feature 3)
+- `task_steps_agent_wrapped(client)` - Agentic generative UI with step execution (Feature 4)
+- `ui_generator_agent(client)` - Tool-based generative UI (Feature 5)
+- `recipe_agent(client)` - Shared state management (Feature 6)
+- `document_writer_agent(client)` - Predictive state updates (Feature 7)
+- `research_assistant_agent(client)` - Research with progress events
+- `task_planner_agent(client)` - Task planning with approvals
+- `subgraphs_agent()` - Deterministic travel-planning subgraphs flow (Dojo `subgraphs` feature)
+
+### Using Example Agents
+
+```python
+from agent_framework.openai import OpenAIChatCompletionClient
+from agent_framework.openai import OpenAIChatClient
+from agent_framework_ag_ui_examples.agents import (
+    simple_agent,
+    weather_agent,
+    recipe_agent,
+)
+
+# Create a chat client (use any SupportsChatGetResponse implementation)
+azure_client = OpenAIChatCompletionClient(model="gpt-4")
+openai_client = OpenAIChatClient(model="gpt-4o")
+
+# Create agent instances by calling the factory functions
+agent1 = simple_agent(azure_client)
+agent2 = weather_agent(openai_client)
+agent3 = recipe_agent(azure_client)
+```
+
+### Running the Example Server
+
+The example server demonstrates all 7 AG-UI features:
+
+```bash
+# Install the package
+pip install agent-framework-ag-ui
+
+# Run the example server
+python -m agent_framework_ag_ui_examples
+
+# Or with debug logging
+ENABLE_DEBUG_LOGGING=1 python -m agent_framework_ag_ui_examples
+```
+
+The server exposes endpoints at:
+- `/agentic_chat` - Simple chat with `simple_agent`
+- `/backend_tool_rendering` - Weather tools with `weather_agent`
+- `/human_in_the_loop` - Step approval with `human_in_the_loop_agent`
+- `/agentic_generative_ui` - Task steps with `task_steps_agent_wrapped`
+- `/tool_based_generative_ui` - Custom UI components with `ui_generator_agent`
+- `/shared_state` - Recipe management with `recipe_agent`
+- `/predictive_state_updates` - Document writing with `document_writer_agent`
+- `/subgraphs` - Travel planner with interrupt-driven flight/hotel choices via `subgraphs_agent`
+
+### Interrupt and Resume Shape
+
+Human-in-the-loop and workflow examples use the canonical AG-UI protocol shape. A paused run finishes with
+`RUN_FINISHED.outcome.type == "interrupt"` and renders prompts from `RUN_FINISHED.outcome.interrupts`; it does not
+depend on a stable top-level `RUN_FINISHED.interrupt` field.
+
+Resume interrupted example threads with a canonical `resume` array:
+
+```json
+{
+  "threadId": "thread-1",
+  "messages": [],
+  "resume": [
+    {
+      "interruptId": "interrupt_1",
+      "status": "resolved",
+      "payload": {
+        "approved": true
+      }
+    }
+  ]
+}
+```
+
+### Complete FastAPI Example
+
+```python
+from fastapi import FastAPI
+from agent_framework.openai import OpenAIChatCompletionClient
+from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
+from agent_framework_ag_ui_examples.agents import (
+    simple_agent,
+    weather_agent,
+    human_in_the_loop_agent,
+    task_steps_agent_wrapped,
+    ui_generator_agent,
+    recipe_agent,
+    document_writer_agent,
+    subgraphs_agent,
+)
+
+app = FastAPI(title="AG-UI Examples")
+
+# Create a chat client (shared across all agents, or create individual ones)
+client = OpenAIChatCompletionClient(model="gpt-4")
+
+# Add all example endpoints
+add_agent_framework_fastapi_endpoint(app, simple_agent(client), "/agentic_chat")
+add_agent_framework_fastapi_endpoint(app, weather_agent(client), "/backend_tool_rendering")
+add_agent_framework_fastapi_endpoint(app, human_in_the_loop_agent(client), "/human_in_the_loop")
+add_agent_framework_fastapi_endpoint(app, task_steps_agent_wrapped(client), "/agentic_generative_ui")  # type: ignore[arg-type]
+add_agent_framework_fastapi_endpoint(app, ui_generator_agent(client), "/tool_based_generative_ui")
+add_agent_framework_fastapi_endpoint(app, recipe_agent(client), "/shared_state")
+add_agent_framework_fastapi_endpoint(app, document_writer_agent(client), "/predictive_state_updates")
+add_agent_framework_fastapi_endpoint(app, subgraphs_agent(), "/subgraphs")
+```
+
+## Architecture
+
+The package uses a clean, orchestrator-based architecture:
+
+- **AgentFrameworkAgent**: Lightweight wrapper that delegates to orchestrators
+- **Orchestrators**: Handle different execution flows (default, human-in-the-loop, etc.)
+- **Confirmation Strategies**: Domain-specific confirmation messages (extensible)
+- **AgentFrameworkEventBridge**: Converts AgentResponseUpdate to AG-UI events
+- **Message Adapters**: Bidirectional conversion between AG-UI and Agent Framework message formats
+- **FastAPI Endpoint**: Streaming HTTP endpoint with Server-Sent Events (SSE)
+
+### Key Design Patterns
+
+- **Orchestrator Pattern**: Separates flow control from protocol translation
+- **Strategy Pattern**: Pluggable confirmation message strategies
+- **Context Object**: Lazy-loaded execution context passed to orchestrators
+- **Event Bridge**: Stateless translation of Agent Framework events to AG-UI events
+
+## Advanced Usage
+
+### Creating Custom Agent Factories
+
+You can create your own agent factories following the same pattern as the examples:
+
+```python
+from agent_framework import Agent, tool
+from agent_framework import SupportsChatGetResponse
+from agent_framework.ag_ui import AgentFrameworkAgent
+
+@tool
+def my_tool(param: str) -> str:
+    """My custom tool."""
+    return f"Result: {param}"
+
+def my_custom_agent(client: SupportsChatGetResponse) -> AgentFrameworkAgent:
+    """Create a custom agent with the specified chat client.
+
+    Args:
+        client: The chat client to use for the agent
+
+    Returns:
+        A configured AgentFrameworkAgent instance
+    """
+    agent = Agent(
+        name="my_custom_agent",
+        instructions="Custom instructions here",
+        client=client,
+        tools=[my_tool],
+    )
+
+    return AgentFrameworkAgent(
+        agent=agent,
+        name="MyCustomAgent",
+        description="My custom agent description",
+    )
+
+# Use it
+from agent_framework.openai import OpenAIChatCompletionClient
+client = OpenAIChatCompletionClient()
+agent = my_custom_agent(client)
+```
+
+### Shared State
+
+State is injected as system messages and updated via predictive state updates:
+
+```python
+from agent_framework import Agent
+from agent_framework.openai import OpenAIChatCompletionClient
+from agent_framework.ag_ui import AgentFrameworkAgent
+
+# Create your agent
+agent = Agent(
+    name="recipe_agent",
+    client=OpenAIChatCompletionClient(model="gpt-4o"),
+)
+
+state_schema = {
+    "recipe": {
+        "type": "object",
+        "properties": {
+            "name": {"type": "string"},
+            "ingredients": {"type": "array"}
+        }
+    }
+}
+
+# Configure which tool updates which state fields
+predict_state_config = {
+    "recipe": {"tool": "update_recipe", "tool_argument": "recipe_data"}
+}
+
+wrapped_agent = AgentFrameworkAgent(
+    agent=agent,
+    state_schema=state_schema,
+    predict_state_config=predict_state_config,
+)
+```
+
+### Predictive State Updates
+
+Predictive state updates automatically stream tool arguments as optimistic state updates:
+
+```python
+from agent_framework import Agent
+from agent_framework.openai import OpenAIChatCompletionClient
+from agent_framework.ag_ui import AgentFrameworkAgent
+
+# Create your agent
+agent = Agent(
+    name="document_writer",
+    client=OpenAIChatCompletionClient(model="gpt-4o"),
+)
+
+predict_state_config = {
+    "current_title": {"tool": "write_document", "tool_argument": "title"},
+    "current_content": {"tool": "write_document", "tool_argument": "content"},
+}
+
+wrapped_agent = AgentFrameworkAgent(
+    agent=agent,
+    state_schema={"current_title": {"type": "string"}, "current_content": {"type": "string"}},
+    predict_state_config=predict_state_config,
+    require_confirmation=True,  # User can approve/reject changes
+)
+```
+
+### Human in the Loop
+
+Human-in-the-loop is automatically handled when tools are marked for approval:
+
+```python
+from agent_framework import tool
+
+@tool(approval_mode="always_require")
+def sensitive_action(param: str) -> str:
+    """This action requires user approval."""
+    return f"Executed with {param}"
+
+# The orchestrator automatically detects approval responses and handles them
+```
+
+### Custom Orchestrators
+
+Add custom execution flows by implementing the Orchestrator pattern:
+
+```python
+from agent_framework.ag_ui._orchestrators import Orchestrator, ExecutionContext
+
+class MyCustomOrchestrator(Orchestrator):
+    def can_handle(self, context: ExecutionContext) -> bool:
+        # Return True if this orchestrator should handle the request
+        return context.input_data.get("custom_mode") == True
+
+    async def run(self, context: ExecutionContext):
+        # Custom execution logic
+        yield RunStartedEvent(...)
+        # ... your custom flow
+        yield RunFinishedEvent(...)
+
+wrapped_agent = AgentFrameworkAgent(
+    agent=your_agent,
+    orchestrators=[MyCustomOrchestrator(), DefaultOrchestrator()],
+)
+
+## License
+
+MIT

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

@@ -0,0 +1,7 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Example agents for AG-UI demonstration."""
+
+from . import agents
+
+__all__ = ["agents"]

python/packages/ag-ui/agent_framework_ag_ui_examples/__main__.py

@@ -0,0 +1,8 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Entry point for running the AG-UI examples server as a module."""
+
+from .server.main import main
+
+if __name__ == "__main__":
+    main()

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/__init__.py

@@ -0,0 +1,27 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Example agents for AG-UI demonstration."""
+
+from .document_writer_agent import document_writer_agent
+from .human_in_the_loop_agent import human_in_the_loop_agent
+from .recipe_agent import recipe_agent
+from .research_assistant_agent import research_assistant_agent
+from .simple_agent import simple_agent
+from .subgraphs_agent import subgraphs_agent
+from .task_planner_agent import task_planner_agent
+from .task_steps_agent import task_steps_agent_wrapped
+from .ui_generator_agent import ui_generator_agent
+from .weather_agent import weather_agent
+
+__all__ = [
+    "document_writer_agent",
+    "human_in_the_loop_agent",
+    "recipe_agent",
+    "research_assistant_agent",
+    "simple_agent",
+    "subgraphs_agent",
+    "task_planner_agent",
+    "task_steps_agent_wrapped",
+    "ui_generator_agent",
+    "weather_agent",
+]

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/document_writer_agent.py

@@ -0,0 +1,69 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Example agent demonstrating predictive state updates with document writing."""
+
+from __future__ import annotations
+
+from agent_framework import Agent, SupportsChatGetResponse, tool
+from agent_framework.ag_ui import AgentFrameworkAgent
+
+
+@tool(approval_mode="always_require")
+def write_document(document: str) -> str:
+    """Write a document. Use markdown formatting to format the document.
+
+    It's good to format the document extensively so it's easy to read.
+    You can use all kinds of markdown.
+    However, do not use italic or strike-through formatting, it's reserved for another purpose.
+    You MUST write the full document, even when changing only a few words.
+    When making edits to the document, try to make them minimal - do not change every word.
+    Keep stories SHORT!
+
+    Args:
+        document: The complete document content in markdown format
+
+    Returns:
+        Confirmation that the document was written
+    """
+    return "Document written."
+
+
+_DOCUMENT_WRITER_INSTRUCTIONS = (
+    "You are a helpful assistant for writing documents. "
+    "To write the document, you MUST use the write_document tool. "
+    "You MUST write the full document, even when changing only a few words. "
+    "When you wrote the document, DO NOT repeat it as a message. "
+    "Just briefly summarize the changes you made. 2 sentences max. "
+    "\n\n"
+    "The current state of the document will be provided to you. "
+    "When editing, make minimal changes - do not change every word unless requested."
+)
+
+
+def document_writer_agent(client: SupportsChatGetResponse) -> AgentFrameworkAgent:
+    """Create a document writer agent with predictive state updates.
+
+    Args:
+        client: The chat client to use for the agent
+
+    Returns:
+        A configured AgentFrameworkAgent instance with document writing capabilities
+    """
+    agent = Agent(
+        name="document_writer",
+        instructions=_DOCUMENT_WRITER_INSTRUCTIONS,
+        client=client,
+        tools=[write_document],
+    )
+
+    return AgentFrameworkAgent(
+        agent=agent,
+        name="DocumentWriter",
+        description="Writes and edits documents with predictive state updates",
+        state_schema={
+            "document": {"type": "string", "description": "The current document content"},
+        },
+        predict_state_config={
+            "document": {"tool": "write_document", "tool_argument": "document"},
+        },
+    )

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/human_in_the_loop_agent.py

@@ -0,0 +1,86 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Human-in-the-loop agent demonstrating step customization (Feature 5)."""
+
+from enum import Enum
+from typing import Any
+
+from agent_framework import Agent, SupportsChatGetResponse, tool
+from pydantic import BaseModel, Field
+
+
+class StepStatus(str, Enum):
+    """Status of a task step."""
+
+    ENABLED = "enabled"
+    DISABLED = "disabled"
+
+
+class TaskStep(BaseModel):
+    """A single step in a task execution plan."""
+
+    description: str = Field(..., description="The text of the step in imperative form (e.g., 'Dig hole', 'Open door')")
+    status: StepStatus = Field(default=StepStatus.ENABLED, description="Whether the step is enabled or disabled")
+
+
+@tool(
+    name="generate_task_steps",
+    description="Generate execution steps for a task",
+    approval_mode="always_require",
+)
+def generate_task_steps(steps: list[TaskStep]) -> str:
+    """Make up 10 steps (only a couple of words per step) that are required for a task.
+
+    The step should be in imperative form (i.e. Dig hole, Open door, ...).
+    Each step will have status='enabled' by default.
+
+    Args:
+        steps: An array of 10 step objects, each containing description and status
+
+    Returns:
+        Confirmation message
+    """
+    return f"Generated {len(steps)} execution steps for the task."
+
+
+def human_in_the_loop_agent(client: SupportsChatGetResponse[Any]) -> Agent[Any]:
+    """Create a human-in-the-loop agent using tool-based approach for predictive state.
+
+    Args:
+        client: The chat client to use for the agent
+
+    Returns:
+        A configured Agent instance with human-in-the-loop capabilities
+    """
+    return Agent(
+        name="human_in_the_loop_agent",
+        instructions="""You are a helpful assistant that can perform any task by breaking it down into steps.
+
+    When asked to perform a task, you MUST call the `generate_task_steps` function with the proper
+    number of steps per the request.
+
+    Rules for steps:
+    - Each step description should be in imperative form (e.g., "Dig hole", "Open door", "Prepare ingredients")
+    - Each step should be brief (only a couple of words)
+    - All steps must have status='enabled' initially
+
+    Example steps for "Build a robot":
+    1. "Design blueprint"
+    2. "Gather components"
+    3. "Assemble frame"
+    4. "Install motors"
+    5. "Wire electronics"
+    6. "Program controller"
+    7. "Test movements"
+    8. "Add sensors"
+    9. "Calibrate systems"
+    10. "Final testing"
+
+    IMPORTANT: When you call generate_task_steps, the user will be shown the steps and asked to approve.
+    Do NOT output any text along with the function call - just call the function.
+    After the user approves and the function executes, THEN provide a brief acknowledgment like:
+    "The plan has been created with X steps selected."
+    """,
+        client=client,
+        tools=[generate_task_steps],
+    )

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/recipe_agent.py

@@ -0,0 +1,134 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Recipe agent example demonstrating shared state management (Feature 3)."""
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Any
+
+from agent_framework import Agent, SupportsChatGetResponse, tool
+from agent_framework.ag_ui import AgentFrameworkAgent
+from pydantic import BaseModel, Field
+
+
+class SkillLevel(str, Enum):
+    """The skill level required for the recipe."""
+
+    BEGINNER = "Beginner"
+    INTERMEDIATE = "Intermediate"
+    ADVANCED = "Advanced"
+
+
+class CookingTime(str, Enum):
+    """The cooking time of the recipe."""
+
+    FIVE_MIN = "5 min"
+    FIFTEEN_MIN = "15 min"
+    THIRTY_MIN = "30 min"
+    FORTY_FIVE_MIN = "45 min"
+    SIXTY_PLUS_MIN = "60+ min"
+
+
+class Ingredient(BaseModel):
+    """An ingredient with its details."""
+
+    icon: str = Field(..., description="Emoji icon representing the ingredient (e.g., 🥕)")
+    name: str = Field(..., description="Name of the ingredient")
+    amount: str = Field(..., description="Amount or quantity of the ingredient")
+
+
+class Recipe(BaseModel):
+    """A complete recipe."""
+
+    title: str = Field(..., description="The title of the recipe")
+    skill_level: SkillLevel = Field(..., description="The skill level required")
+    special_preferences: list[str] = Field(
+        default_factory=list, description="Dietary preferences (e.g., Vegetarian, Gluten-free)"
+    )
+    cooking_time: CookingTime = Field(..., description="The estimated cooking time")
+    ingredients: list[Ingredient] = Field(..., description="Complete list of ingredients")
+    instructions: list[str] = Field(..., description="Step-by-step cooking instructions")
+
+
+@tool
+def update_recipe(recipe: Recipe) -> str:
+    """Update the recipe with new or modified content.
+
+    You MUST write the complete recipe with ALL fields, even when changing only a few items.
+    When modifying an existing recipe, include ALL existing ingredients and instructions plus your changes.
+    NEVER delete existing data - only add or modify.
+
+    Args:
+        recipe: The complete recipe object with all details
+
+    Returns:
+        Confirmation that the recipe was updated
+    """
+    return "Recipe updated."
+
+
+_RECIPE_INSTRUCTIONS = """You are a helpful recipe assistant that creates and modifies recipes.
+
+    CRITICAL RULES:
+    1. You will receive the current recipe state in the system context
+    2. To update the recipe, you MUST use the update_recipe tool
+    3. When modifying a recipe, ALWAYS include ALL existing data plus your changes in the tool call
+    4. NEVER delete existing ingredients or instructions - only add or modify
+    5. After calling the tool, provide a brief conversational message (1-2 sentences)
+
+    When creating a NEW recipe:
+    - Provide all required fields: title, skill_level, cooking_time, ingredients, instructions
+    - Use actual emojis for ingredient icons (🥕 🧄 🧅 🍅 🌿 🍗 🥩 🧀)
+    - Leave special_preferences empty unless specified
+    - Message: "Here's your recipe!" or similar
+
+    When MODIFYING or IMPROVING an existing recipe:
+    - Include ALL existing ingredients + any new ones
+    - Include ALL existing instructions + any new/modified ones
+    - Update other fields as needed
+    - Message: Explain what you improved (e.g., "I upgraded the ingredients to premium quality")
+    - When asked to "improve", enhance with:
+      * Better ingredients (upgrade quality, add complementary flavors)
+      * More detailed instructions
+      * Professional techniques
+      * Adjust skill_level if complexity changes
+      * Add relevant special_preferences
+
+    Example improvements:
+    - Upgrade "chicken" → "organic free-range chicken breast"
+    - Add herbs: basil, oregano, thyme
+    - Add aromatics: garlic, shallots
+    - Add finishing touches: lemon zest, fresh parsley
+    - Make instructions more detailed and professional
+    """
+
+
+def recipe_agent(client: SupportsChatGetResponse[Any]) -> AgentFrameworkAgent:
+    """Create a recipe agent with streaming state updates.
+
+    Args:
+        client: The chat client to use for the agent
+
+    Returns:
+        A configured AgentFrameworkAgent instance with recipe management
+    """
+    agent = Agent(
+        name="recipe_agent",
+        instructions=_RECIPE_INSTRUCTIONS,
+        client=client,
+        tools=[update_recipe],
+    )
+
+    return AgentFrameworkAgent(
+        agent=agent,
+        name="RecipeAgent",
+        description="Creates and modifies recipes with streaming state updates",
+        state_schema={
+            "recipe": {"type": "object", "description": "The current recipe"},
+        },
+        predict_state_config={
+            "recipe": {"tool": "update_recipe", "tool_argument": "recipe"},
+        },
+        require_confirmation=False,
+    )

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/research_assistant_agent.py

@@ -0,0 +1,111 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Example agent demonstrating agentic generative UI with custom events during execution."""
+
+import asyncio
+from typing import Any
+
+from agent_framework import Agent, SupportsChatGetResponse, tool
+from agent_framework.ag_ui import AgentFrameworkAgent
+
+
+@tool
+async def research_topic(topic: str) -> str:
+    """Research a topic and generate a comprehensive report.
+
+    Args:
+        topic: The topic to research
+
+    Returns:
+        Research report
+    """
+    # Simulate multi-step research process
+    steps = [
+        ("Searching databases", 1.0),
+        ("Analyzing sources", 1.5),
+        ("Synthesizing information", 1.0),
+        ("Generating report", 0.5),
+    ]
+
+    results: list[str] = []
+    for step_name, duration in steps:
+        await asyncio.sleep(duration)
+        results.append(f"- {step_name}: completed")
+
+    return f"Research report on '{topic}':\n" + "\n".join(results)
+
+
+@tool
+async def create_presentation(title: str, num_slides: int) -> str:
+    """Create a presentation with multiple slides.
+
+    Args:
+        title: Presentation title
+        num_slides: Number of slides to create
+
+    Returns:
+        Presentation summary
+    """
+    # Simulate slide generation
+    slides: list[str] = []
+    for i in range(num_slides):
+        await asyncio.sleep(0.5)
+        slides.append(f"Slide {i + 1}: Content for {title}")
+
+    return f"Created presentation '{title}' with {num_slides} slides:\n" + "\n".join(slides)
+
+
+@tool
+async def analyze_data(dataset: str) -> str:
+    """Analyze a dataset and produce insights.
+
+    Args:
+        dataset: The dataset name to analyze
+
+    Returns:
+        Analysis results
+    """
+    # Simulate data analysis phases
+    phases = [
+        ("Loading data", 0.8),
+        ("Cleaning data", 1.0),
+        ("Running statistical analysis", 1.2),
+        ("Generating visualizations", 0.7),
+    ]
+
+    insights: list[str] = []
+    for phase_name, duration in phases:
+        await asyncio.sleep(duration)
+        insights.append(f"- {phase_name}: done")
+
+    return f"Analysis of '{dataset}':\n" + "\n".join(insights)
+
+
+_RESEARCH_ASSISTANT_INSTRUCTIONS = (
+    "You are a research and analysis assistant. "
+    "You can research topics, create presentations, and analyze data. "
+    "Use the available tools to help users with their research needs."
+)
+
+
+def research_assistant_agent(client: SupportsChatGetResponse[Any]) -> AgentFrameworkAgent:
+    """Create a research assistant agent.
+
+    Args:
+        client: The chat client to use for the agent
+
+    Returns:
+        A configured AgentFrameworkAgent instance with research capabilities
+    """
+    agent = Agent(
+        name="research_assistant",
+        instructions=_RESEARCH_ASSISTANT_INSTRUCTIONS,
+        client=client,
+        tools=[research_topic, create_presentation, analyze_data],
+    )
+
+    return AgentFrameworkAgent(
+        agent=agent,
+        name="ResearchAssistant",
+        description="Research assistant that emits progress events during task execution",
+    )

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/simple_agent.py

@@ -0,0 +1,23 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Simple agentic chat example (Feature 1: Agentic Chat)."""
+
+from typing import Any
+
+from agent_framework import Agent, SupportsChatGetResponse
+
+
+def simple_agent(client: SupportsChatGetResponse[Any]) -> Agent[Any]:
+    """Create a simple chat agent.
+
+    Args:
+        client: The chat client to use for the agent
+
+    Returns:
+        A configured Agent instance
+    """
+    return Agent[Any](
+        name="simple_chat_agent",
+        instructions="You are a helpful assistant. Be concise and friendly.",
+        client=client,
+    )

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/subgraphs_agent.py

@@ -0,0 +1,405 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Subgraphs travel planner built with MAF workflow primitives."""
+
+import json
+import uuid
+from copy import deepcopy
+from dataclasses import dataclass
+from typing import Any
+
+from ag_ui.core import (
+    BaseEvent,
+    StateSnapshotEvent,
+    TextMessageContentEvent,
+    TextMessageEndEvent,
+    TextMessageStartEvent,
+)
+from agent_framework import (
+    Executor,
+    Message,
+    Workflow,
+    WorkflowBuilder,
+    WorkflowContext,
+    handler,
+    response_handler,
+)
+
+from agent_framework_ag_ui import AgentFrameworkWorkflow
+
+STATIC_FLIGHTS: list[dict[str, str]] = [
+    {
+        "airline": "KLM",
+        "departure": "Amsterdam (AMS)",
+        "arrival": "San Francisco (SFO)",
+        "price": "$650",
+        "duration": "11h 30m",
+    },
+    {
+        "airline": "United",
+        "departure": "Amsterdam (AMS)",
+        "arrival": "San Francisco (SFO)",
+        "price": "$720",
+        "duration": "12h 15m",
+    },
+]
+
+STATIC_HOTELS: list[dict[str, str]] = [
+    {
+        "name": "Hotel Zephyr",
+        "location": "Fisherman's Wharf",
+        "price_per_night": "$280/night",
+        "rating": "4.2 stars",
+    },
+    {
+        "name": "The Ritz-Carlton",
+        "location": "Nob Hill",
+        "price_per_night": "$550/night",
+        "rating": "4.8 stars",
+    },
+    {
+        "name": "Hotel Zoe",
+        "location": "Union Square",
+        "price_per_night": "$320/night",
+        "rating": "4.4 stars",
+    },
+]
+
+STATIC_EXPERIENCES: list[dict[str, str]] = [
+    {
+        "name": "Pier 39",
+        "type": "activity",
+        "description": "Iconic waterfront destination with shops and sea lions",
+        "location": "Fisherman's Wharf",
+    },
+    {
+        "name": "Golden Gate Bridge",
+        "type": "activity",
+        "description": "World-famous suspension bridge with stunning views",
+        "location": "Golden Gate",
+    },
+    {
+        "name": "Swan Oyster Depot",
+        "type": "restaurant",
+        "description": "Historic seafood counter serving fresh oysters",
+        "location": "Polk Street",
+    },
+    {
+        "name": "Tartine Bakery",
+        "type": "restaurant",
+        "description": "Artisanal bakery famous for bread and pastries",
+        "location": "Mission District",
+    },
+]
+
+_STATE_KEY = "subgraphs_state"
+
+
+@dataclass
+class _PresentFlights:
+    pass
+
+
+@dataclass
+class _PresentHotels:
+    pass
+
+
+@dataclass
+class _PlanExperiences:
+    pass
+
+
+@dataclass
+class _FinalizeTrip:
+    pass
+
+
+def _initial_state() -> dict[str, Any]:
+    return {
+        "itinerary": {},
+        "experiences": [],
+        "flights": [],
+        "hotels": [],
+        "planning_step": "start",
+        "active_agent": "supervisor",
+    }
+
+
+def _emit_text_events(text: str) -> list[BaseEvent]:
+    message_id = str(uuid.uuid4())
+    return [
+        TextMessageStartEvent(message_id=message_id, role="assistant"),
+        TextMessageContentEvent(message_id=message_id, delta=text),
+        TextMessageEndEvent(message_id=message_id),
+    ]
+
+
+async def _emit_text(ctx: WorkflowContext[Any, BaseEvent], text: str) -> None:
+    for event in _emit_text_events(text):
+        await ctx.yield_output(event)
+
+
+async def _emit_state_snapshot(ctx: WorkflowContext[Any, BaseEvent], state: dict[str, Any]) -> None:
+    await ctx.yield_output(StateSnapshotEvent(snapshot=deepcopy(state)))
+
+
+def _flight_interrupt_value() -> dict[str, Any]:
+    return {
+        "message": "Choose the flight you want. I recommend KLM because it is cheaper and usually on time.",
+        "options": deepcopy(STATIC_FLIGHTS),
+        "recommendation": deepcopy(STATIC_FLIGHTS[0]),
+        "agent": "flights",
+    }
+
+
+def _hotel_interrupt_value() -> dict[str, Any]:
+    return {
+        "message": "Choose your hotel. I recommend Hotel Zoe for the best value in a central location.",
+        "options": deepcopy(STATIC_HOTELS),
+        "recommendation": deepcopy(STATIC_HOTELS[2]),
+        "agent": "hotels",
+    }
+
+
+def _normalize_flight(value: Any) -> dict[str, str] | None:
+    if isinstance(value, str):
+        try:
+            value = json.loads(value)
+        except json.JSONDecodeError:
+            return None
+    if isinstance(value, dict) and value.get("airline"):
+        return {
+            "airline": str(value.get("airline", "")),
+            "departure": str(value.get("departure", "")),
+            "arrival": str(value.get("arrival", "")),
+            "price": str(value.get("price", "")),
+            "duration": str(value.get("duration", "")),
+        }
+    return None
+
+
+def _normalize_hotel(value: Any) -> dict[str, str] | None:
+    if isinstance(value, str):
+        try:
+            value = json.loads(value)
+        except json.JSONDecodeError:
+            return None
+    if isinstance(value, dict) and value.get("name"):
+        return {
+            "name": str(value.get("name", "")),
+            "location": str(value.get("location", "")),
+            "price_per_night": str(value.get("price_per_night", "")),
+            "rating": str(value.get("rating", "")),
+        }
+    return None
+
+
+def _load_state(ctx: WorkflowContext[Any, BaseEvent]) -> dict[str, Any]:
+    state = ctx.get_state(_STATE_KEY)
+    if isinstance(state, dict):
+        return state
+    new_state = _initial_state()
+    ctx.set_state(_STATE_KEY, new_state)
+    return new_state
+
+
+class _SupervisorExecutor(Executor):
+    def __init__(self) -> None:
+        super().__init__(id="supervisor_agent")
+
+    @handler
+    async def start(self, message: list[Message], ctx: WorkflowContext[_PresentFlights, BaseEvent]) -> None:
+        del message
+        state = _initial_state()
+        ctx.set_state(_STATE_KEY, state)
+        await _emit_state_snapshot(ctx, state)
+
+        await _emit_text(
+            ctx,
+            "Supervisor: I will coordinate our specialist agents to plan your San Francisco trip end to end.",
+        )
+
+        state["active_agent"] = "flights"
+        state["planning_step"] = "collecting_flights"
+        state["flights"] = deepcopy(STATIC_FLIGHTS)
+        ctx.set_state(_STATE_KEY, state)
+        await _emit_state_snapshot(ctx, state)
+
+        await ctx.send_message(_PresentFlights(), target_id="flights_agent")
+
+    @handler
+    async def finalize(self, message: _FinalizeTrip, ctx: WorkflowContext[Any, BaseEvent]) -> None:
+        del message
+        state = _load_state(ctx)
+        state["active_agent"] = "supervisor"
+        state["planning_step"] = "complete"
+        ctx.set_state(_STATE_KEY, state)
+        await _emit_state_snapshot(ctx, state)
+        await _emit_text(ctx, "Supervisor: Your travel planning is complete and your itinerary is ready.")
+
+
+class _FlightsExecutor(Executor):
+    def __init__(self) -> None:
+        super().__init__(id="flights_agent")
+
+    @handler
+    async def present_options(self, message: _PresentFlights, ctx: WorkflowContext[_PresentHotels, BaseEvent]) -> None:
+        del message
+        await _emit_text(
+            ctx,
+            "Flights Agent: I found two flight options from Amsterdam to San Francisco. "
+            "KLM is recommended for the best value and schedule.",
+        )
+        await ctx.request_info(_flight_interrupt_value(), dict, request_id="flights-choice")
+
+    @response_handler
+    async def handle_selection(
+        self,
+        original_request: dict,
+        response: dict,
+        ctx: WorkflowContext[_PresentHotels, BaseEvent],
+    ) -> None:
+        del original_request
+        state = _load_state(ctx)
+        selected_flight = _normalize_flight(response)
+
+        if selected_flight is None:
+            state["active_agent"] = "flights"
+            state["planning_step"] = "collecting_flights"
+            state["flights"] = deepcopy(STATIC_FLIGHTS)
+            ctx.set_state(_STATE_KEY, state)
+            await _emit_state_snapshot(ctx, state)
+            await _emit_text(ctx, "Flights Agent: Please choose a flight option from the selection card to continue.")
+            await ctx.request_info(_flight_interrupt_value(), dict, request_id="flights-choice")
+            return
+
+        itinerary = state.setdefault("itinerary", {})
+        itinerary["flight"] = selected_flight
+
+        state["active_agent"] = "flights"
+        state["planning_step"] = "booking_flight"
+        ctx.set_state(_STATE_KEY, state)
+        await _emit_state_snapshot(ctx, state)
+
+        await _emit_text(
+            ctx,
+            f"Flights Agent: Great choice. I will book the {selected_flight['airline']} flight. "
+            "Now I am routing you to Hotels Agent for accommodation.",
+        )
+
+        state["active_agent"] = "hotels"
+        state["planning_step"] = "collecting_hotels"
+        state["hotels"] = deepcopy(STATIC_HOTELS)
+        ctx.set_state(_STATE_KEY, state)
+        await _emit_state_snapshot(ctx, state)
+
+        await ctx.send_message(_PresentHotels(), target_id="hotels_agent")
+
+
+class _HotelsExecutor(Executor):
+    def __init__(self) -> None:
+        super().__init__(id="hotels_agent")
+
+    @handler
+    async def present_options(self, message: _PresentHotels, ctx: WorkflowContext[_PlanExperiences, BaseEvent]) -> None:
+        del message
+        await _emit_text(
+            ctx,
+            "Hotels Agent: I found three accommodation options in San Francisco. "
+            "Hotel Zoe is recommended for the best balance of location, quality, and price.",
+        )
+        await ctx.request_info(_hotel_interrupt_value(), dict, request_id="hotels-choice")
+
+    @response_handler
+    async def handle_selection(
+        self,
+        original_request: dict,
+        response: dict,
+        ctx: WorkflowContext[_PlanExperiences, BaseEvent],
+    ) -> None:
+        del original_request
+        state = _load_state(ctx)
+        selected_hotel = _normalize_hotel(response)
+
+        if selected_hotel is None:
+            state["active_agent"] = "hotels"
+            state["planning_step"] = "collecting_hotels"
+            state["hotels"] = deepcopy(STATIC_HOTELS)
+            ctx.set_state(_STATE_KEY, state)
+            await _emit_state_snapshot(ctx, state)
+            await _emit_text(ctx, "Hotels Agent: Please choose a hotel option from the selection card to continue.")
+            await ctx.request_info(_hotel_interrupt_value(), dict, request_id="hotels-choice")
+            return
+
+        itinerary = state.setdefault("itinerary", {})
+        itinerary["hotel"] = selected_hotel
+
+        state["active_agent"] = "hotels"
+        state["planning_step"] = "booking_hotel"
+        ctx.set_state(_STATE_KEY, state)
+        await _emit_state_snapshot(ctx, state)
+
+        await _emit_text(
+            ctx,
+            f"Hotels Agent: Excellent, {selected_hotel['name']} is booked. "
+            "I am routing you to Experiences Agent for activities and restaurants.",
+        )
+
+        state["active_agent"] = "experiences"
+        state["planning_step"] = "curating_experiences"
+        state["experiences"] = deepcopy(STATIC_EXPERIENCES)
+        ctx.set_state(_STATE_KEY, state)
+        await _emit_state_snapshot(ctx, state)
+
+        await ctx.send_message(_PlanExperiences(), target_id="experiences_agent")
+
+
+class _ExperiencesExecutor(Executor):
+    def __init__(self) -> None:
+        super().__init__(id="experiences_agent")
+
+    @handler
+    async def plan(self, message: _PlanExperiences, ctx: WorkflowContext[_FinalizeTrip, BaseEvent]) -> None:
+        del message
+        await _emit_text(
+            ctx,
+            "Experiences Agent: I planned activities and restaurants including "
+            "Pier 39, Golden Gate Bridge, Swan Oyster Depot, and Tartine Bakery.",
+        )
+        await ctx.send_message(_FinalizeTrip(), target_id="supervisor_agent")
+
+
+def _build_subgraphs_workflow() -> Workflow:
+    supervisor = _SupervisorExecutor()
+    flights = _FlightsExecutor()
+    hotels = _HotelsExecutor()
+    experiences = _ExperiencesExecutor()
+
+    return (
+        WorkflowBuilder(
+            name="subgraphs",
+            description="Travel planning supervisor with flights/hotels/experiences subgraphs.",
+            start_executor=supervisor,
+        )
+        .add_edge(supervisor, flights)
+        .add_edge(flights, hotels)
+        .add_edge(hotels, experiences)
+        .add_edge(experiences, supervisor)
+        .build()
+    )
+
+
+def _build_subgraphs_workflow_for_thread(thread_id: str) -> Workflow:
+    """Create a workflow instance scoped to a single AG-UI thread."""
+    del thread_id
+    return _build_subgraphs_workflow()
+
+
+def subgraphs_agent() -> AgentFrameworkWorkflow:
+    """Create the subgraphs travel planner agent."""
+    return AgentFrameworkWorkflow(
+        workflow_factory=_build_subgraphs_workflow_for_thread,
+        name="subgraphs",
+        description="Travel planning workflow with interrupt-driven selections.",
+    )

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/task_planner_agent.py

@@ -0,0 +1,84 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Example agent demonstrating human-in-the-loop with function approvals."""
+
+from typing import Any
+
+from agent_framework import Agent, SupportsChatGetResponse, tool
+from agent_framework.ag_ui import AgentFrameworkAgent
+
+
+@tool(approval_mode="always_require")
+def create_calendar_event(title: str, date: str, time: str) -> str:
+    """Create a calendar event.
+
+    Args:
+        title: The event title
+        date: The event date (YYYY-MM-DD)
+        time: The event time (HH:MM)
+
+    Returns:
+        Confirmation message
+    """
+    return f"Calendar event '{title}' created for {date} at {time}"
+
+
+@tool(approval_mode="always_require")
+def send_email(to: str, subject: str, body: str) -> str:
+    """Send an email.
+
+    Args:
+        to: Recipient email address
+        subject: Email subject
+        body: Email body text
+
+    Returns:
+        Confirmation message
+    """
+    return f"Email sent to {to} with subject '{subject}'"
+
+
+@tool(approval_mode="always_require")
+def book_meeting_room(room_name: str, date: str, start_time: str, end_time: str) -> str:
+    """Book a meeting room.
+
+    Args:
+        room_name: The meeting room name
+        date: The booking date (YYYY-MM-DD)
+        start_time: Start time (HH:MM)
+        end_time: End time (HH:MM)
+
+    Returns:
+        Confirmation message
+    """
+    return f"Meeting room '{room_name}' booked for {date} from {start_time} to {end_time}"
+
+
+_TASK_PLANNER_INSTRUCTIONS = (
+    "You are a helpful assistant that plans and executes tasks. "
+    "You have access to calendar, email, and meeting room booking functions. "
+    "All of these actions require user approval before execution."
+)
+
+
+def task_planner_agent(client: SupportsChatGetResponse[Any]) -> AgentFrameworkAgent:
+    """Create a task planner agent with user approval for actions.
+
+    Args:
+        client: The chat client to use for the agent
+
+    Returns:
+        A configured AgentFrameworkAgent instance with task planning capabilities
+    """
+    agent = Agent(
+        name="task_planner",
+        instructions=_TASK_PLANNER_INSTRUCTIONS,
+        client=client,
+        tools=[create_calendar_event, send_email, book_meeting_room],
+    )
+
+    return AgentFrameworkAgent(
+        agent=agent,
+        name="TaskPlanner",
+        description="Plans and executes tasks with user approval",
+    )

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/task_steps_agent.py

@@ -0,0 +1,338 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Task steps agent demonstrating agentic generative UI (Feature 6)."""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncGenerator
+from enum import Enum
+from typing import Any
+
+from ag_ui.core import (
+    EventType,
+    MessagesSnapshotEvent,
+    RunFinishedEvent,
+    StateDeltaEvent,
+    StateSnapshotEvent,
+    TextMessageContentEvent,
+    TextMessageEndEvent,
+    TextMessageStartEvent,
+    ToolCallStartEvent,
+)
+from agent_framework import Agent, Content, Message, SupportsChatGetResponse, tool
+from agent_framework.ag_ui import AgentFrameworkAgent
+from pydantic import BaseModel, Field
+
+from agent_framework_ag_ui import AgentFrameworkWorkflow
+
+
+class StepStatus(str, Enum):
+    """Status of a task step."""
+
+    PENDING = "pending"
+    COMPLETED = "completed"
+
+
+class TaskStep(BaseModel):
+    """A single step in a task."""
+
+    description: str = Field(
+        ..., description="The text of the step in gerund form (e.g., 'Digging hole', 'Opening door')"
+    )
+    status: StepStatus = Field(default=StepStatus.PENDING, description="The status of the step")
+
+
+@tool
+def generate_task_steps(steps: list[TaskStep]) -> str:
+    """Generate a list of task steps for completing a task.
+
+    Args:
+        steps: Complete list of task steps with descriptions and status
+
+    Returns:
+        Confirmation that steps were generated
+    """
+    return "Steps generated."
+
+
+def _create_task_steps_agent(client: SupportsChatGetResponse[Any]) -> AgentFrameworkAgent:
+    """Create the task steps agent using tool-based approach for streaming.
+
+    Args:
+        client: The chat client to use for the agent
+
+    Returns:
+        A configured AgentFrameworkAgent instance
+    """
+    agent = Agent[Any](
+        name="task_steps_agent",
+        instructions="""You are a helpful assistant that breaks down tasks into actionable steps.
+
+    When asked to perform a task, you MUST:
+    1. Use the generate_task_steps tool to create the steps
+    2. Pay attention to how many steps the user requests (if specified)
+    3. If no specific number is mentioned, use a reasonable number of steps (typically 5-10)
+    4. Each step description should be in gerund form (e.g., "Designing spacecraft", "Training astronauts")
+    5. Each step should be brief (only 2-4 words)
+    6. All steps must have status='pending'
+    7. After calling the tool, provide a brief conversational message (one sentence) saying you created the plan
+
+    Example steps for "Build a treehouse in 5 steps":
+    - "Selecting location"
+    - "Gathering materials"
+    - "Assembling frame"
+    - "Installing platform"
+    - "Adding finishing touches"
+    """,
+        client=client,
+        tools=[generate_task_steps],
+    )
+
+    return AgentFrameworkAgent(
+        agent=agent,
+        name="TaskStepsAgent",
+        description="Generates task steps with streaming state updates",
+        state_schema={
+            "steps": {"type": "array", "description": "The list of task steps"},
+        },
+        predict_state_config={
+            "steps": {
+                "tool": "generate_task_steps",
+                "tool_argument": "steps",
+            }
+        },
+        require_confirmation=False,  # Agentic generative UI updates automatically without confirmation
+    )
+
+
+# Wrap the agent's run method to add step execution simulation
+class TaskStepsAgentWithExecution(AgentFrameworkWorkflow):
+    """Wrapper that adds step execution simulation after plan generation.
+
+    This wrapper delegates to AgentFrameworkAgent but is recognized as compatible
+    by add_agent_framework_fastapi_endpoint since it implements run().
+    """
+
+    def __init__(self, base_agent: AgentFrameworkAgent):
+        """Initialize wrapper with base agent."""
+        super().__init__(name=base_agent.name, description=base_agent.description)
+        self._base_agent = base_agent
+
+    def __getattr__(self, name: str) -> Any:
+        """Delegate all other attribute access to base agent."""
+        return getattr(self._base_agent, name)
+
+    async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[Any]:
+        """Run the agent and then simulate step execution."""
+        import logging
+        import uuid
+
+        logger = logging.getLogger(__name__)
+        logger.info("TaskStepsAgentWithExecution.run() called - wrapper is active")
+
+        # First, run the base agent to generate the plan - buffer text messages
+        final_state: dict[str, Any] = {}
+        run_finished_event: Any = None
+        tool_call_id: str | None = None
+        buffered_text_events: list[Any] = []  # Buffer text from first LLM call
+
+        async for event in self._base_agent.run(input_data):
+            event_type_str = str(event.type) if hasattr(event, "type") else type(event).__name__
+            logger.info(f"Processing event: {event_type_str}")
+
+            match event:
+                case StateSnapshotEvent(snapshot=snapshot):
+                    final_state = snapshot.copy() if snapshot else {}
+                    logger.info(f"Captured STATE_SNAPSHOT event with state: {final_state}")
+                    yield event
+                case StateDeltaEvent(delta=delta):
+                    # Apply state delta to final_state
+                    if delta:
+                        for patch in delta:
+                            if patch.get("op") == "replace" and patch.get("path") == "/steps":
+                                final_state["steps"] = patch.get("value", [])
+                                logger.info(
+                                    f"Applied STATE_DELTA: updated steps to {len(final_state.get('steps', []))} items"
+                                )
+                    logger.info(f"Yielding event immediately: {event_type_str}")
+                    yield event
+                case RunFinishedEvent():
+                    run_finished_event = event
+                    logger.info("Captured RUN_FINISHED event - will send after step execution and summary")
+                case ToolCallStartEvent(tool_call_id=call_id):
+                    tool_call_id = call_id
+                    logger.info(f"Captured tool_call_id: {tool_call_id}")
+                    yield event
+                case TextMessageStartEvent() | TextMessageContentEvent() | TextMessageEndEvent():
+                    buffered_text_events.append(event)
+                    logger.info(f"Buffered {event_type_str} from first LLM call")
+                case _:
+                    logger.info(f"Yielding event immediately: {event_type_str}")
+                    yield event
+
+        logger.info(f"Base agent completed. Final state: {final_state}")
+
+        # Now simulate executing the steps
+        if final_state and "steps" in final_state:
+            steps = final_state["steps"]
+            logger.info(f"Starting step execution simulation for {len(steps)} steps")
+
+            for i in range(len(steps)):
+                logger.info(f"Simulating execution of step {i + 1}/{len(steps)}: {steps[i].get('description')}")
+                await asyncio.sleep(1.0)  # Simulate work
+
+                # Update step to completed
+                steps[i]["status"] = "completed"
+                logger.info(f"Step {i + 1} marked as completed")
+
+                # Send delta event with manual JSON patch format
+                delta_event = StateDeltaEvent(
+                    type=EventType.STATE_DELTA,
+                    delta=[
+                        {
+                            "op": "replace",
+                            "path": f"/steps/{i}/status",
+                            "value": "completed",
+                        }
+                    ],
+                )
+                logger.info(f"Yielding StateDeltaEvent for step {i + 1}")
+                yield delta_event
+
+            # Send final snapshot
+            final_snapshot = StateSnapshotEvent(
+                type=EventType.STATE_SNAPSHOT,
+                snapshot={"steps": steps},
+            )
+            logger.info("Yielding final StateSnapshotEvent with all steps completed")
+            yield final_snapshot
+
+            # SECOND LLM call: Stream summary from chat client directly
+            logger.info("Making SECOND LLM call to generate summary after step execution")
+
+            # Get the underlying chat agent and client
+            chat_agent = self._base_agent.agent
+            client = chat_agent.client  # type: ignore
+
+            # Build messages for summary call
+
+            original_messages = input_data.get("messages", [])
+
+            # Convert to Message objects if needed
+            messages: list[Message] = []
+            for msg in original_messages:
+                if isinstance(msg, dict):
+                    content_str = msg.get("content", "")
+                    if isinstance(content_str, str):
+                        messages.append(
+                            Message(
+                                role=msg.get("role", "user"),
+                                contents=[Content.from_text(text=content_str)],
+                            )
+                        )
+                elif isinstance(msg, Message):
+                    messages.append(msg)
+
+            # Add completion message
+            messages.append(
+                Message(
+                    role="user",
+                    contents=[
+                        Content.from_text(
+                            text="The steps have been successfully executed. Provide a brief one-sentence summary."
+                        )
+                    ],
+                )
+            )
+
+            # Stream the LLM response and manually emit text events
+            logger.info("Calling chat client for summary")
+
+            message_id = str(uuid.uuid4())
+
+            try:
+                # Emit TEXT_MESSAGE_START
+                yield TextMessageStartEvent(
+                    type=EventType.TEXT_MESSAGE_START,
+                    message_id=message_id,
+                    role="assistant",
+                )
+                # Small delay to ensure START event is processed before CONTENT events
+                await asyncio.sleep(0.01)
+
+                # Stream completion
+                accumulated_text = ""
+                async for chunk in client.get_response(messages=messages, stream=True):
+                    # chunk is ChatResponseUpdate
+                    if hasattr(chunk, "text") and chunk.text:
+                        accumulated_text += chunk.text
+                        # Emit TEXT_MESSAGE_CONTENT
+                        yield TextMessageContentEvent(
+                            type=EventType.TEXT_MESSAGE_CONTENT,
+                            message_id=message_id,
+                            delta=chunk.text,
+                        )
+
+                # Emit TEXT_MESSAGE_END
+                yield TextMessageEndEvent(
+                    type=EventType.TEXT_MESSAGE_END,
+                    message_id=message_id,
+                )
+                logger.info(f"Summary complete: {accumulated_text}")
+
+                # Build complete message for persistence
+                summary_message = {
+                    "role": "assistant",
+                    "content": accumulated_text,
+                    "id": message_id,
+                }
+                final_messages = list(original_messages)
+                final_messages.append(summary_message)
+
+                # Emit MessagesSnapshotEvent to persist in history
+                yield MessagesSnapshotEvent(
+                    type=EventType.MESSAGES_SNAPSHOT,
+                    messages=final_messages,
+                )
+            except Exception as e:
+                logger.error(f"Error generating summary: {e}")
+                # Generate a new message ID for the error
+                error_message_id = str(uuid.uuid4())
+                # Yield TEXT_MESSAGE_START for error
+                yield TextMessageStartEvent(
+                    type=EventType.TEXT_MESSAGE_START,
+                    message_id=error_message_id,
+                    role="assistant",
+                )
+                # Yield error message content
+                yield TextMessageContentEvent(
+                    type=EventType.TEXT_MESSAGE_CONTENT,
+                    message_id=error_message_id,
+                    delta=f"[Summary generation error: {e!s}]",
+                )
+                # Yield TEXT_MESSAGE_END for error
+                yield TextMessageEndEvent(
+                    type=EventType.TEXT_MESSAGE_END,
+                    message_id=error_message_id,
+                )
+        else:
+            logger.warning(f"No steps found in final_state to execute. final_state={final_state}")
+
+        # Finally send the original RUN_FINISHED event
+        if run_finished_event:
+            logger.info("Yielding original RUN_FINISHED event")
+            yield run_finished_event
+
+
+def task_steps_agent_wrapped(client: SupportsChatGetResponse[Any]) -> TaskStepsAgentWithExecution:
+    """Create a task steps agent with execution simulation.
+
+    Args:
+        client: The chat client to use for the agent
+
+    Returns:
+        A wrapped agent instance with step execution simulation
+    """
+    base_agent = _create_task_steps_agent(client)
+    return TaskStepsAgentWithExecution(base_agent)

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/ui_generator_agent.py

@@ -0,0 +1,193 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Example agent demonstrating Tool-based Generative UI (Feature 5)."""
+
+from __future__ import annotations
+
+import sys
+from typing import TYPE_CHECKING, TypedDict
+
+from agent_framework import Agent, FunctionTool, SupportsChatGetResponse
+from agent_framework.ag_ui import AgentFrameworkAgent
+
+if sys.version_info >= (3, 13):
+    from typing import TypeVar  # pragma: no cover
+else:
+    from typing_extensions import TypeVar  # pragma: no cover
+if sys.version_info >= (3, 11):
+    from typing import TypedDict  # pragma: no cover
+else:
+    from typing_extensions import TypedDict  # pragma: no cover
+
+if TYPE_CHECKING:
+    from agent_framework import ChatOptions
+
+# Declaration-only tools (func=None) - actual rendering happens on the client side
+generate_haiku = FunctionTool(
+    name="generate_haiku",
+    description="""Generate a haiku with image and gradient background (FRONTEND_RENDER).
+
+    This tool generates UI for displaying a haiku with an image and gradient background.
+    The frontend should render this as a custom haiku component.""",
+    func=None,  # Makes declaration_only=True so client renders the UI
+    input_model={
+        "type": "object",
+        "properties": {
+            "english": {
+                "type": "array",
+                "items": {"type": "string"},
+                "description": "English haiku lines (exactly 3 lines)",
+                "minItems": 3,
+                "maxItems": 3,
+            },
+            "japanese": {
+                "type": "array",
+                "items": {"type": "string"},
+                "description": "Japanese haiku lines (exactly 3 lines)",
+                "minItems": 3,
+                "maxItems": 3,
+            },
+            "image_name": {
+                "type": "string",
+                "description": """Image filename for visual accompaniment. Must be one of:
+            - "Osaka_Castle_Turret_Stone_Wall_Pine_Trees_Daytime.jpg"
+            - "Tokyo_Skyline_Night_Tokyo_Tower_Mount_Fuji_View.jpg"
+            - "Itsukushima_Shrine_Miyajima_Floating_Torii_Gate_Sunset_Long_Exposure.jpg"
+            - "Takachiho_Gorge_Waterfall_River_Lush_Greenery_Japan.jpg"
+            - "Bonsai_Tree_Potted_Japanese_Art_Green_Foliage.jpeg"
+            - "Shirakawa-go_Gassho-zukuri_Thatched_Roof_Village_Aerial_View.jpg"
+            - "Ginkaku-ji_Silver_Pavilion_Kyoto_Japanese_Garden_Pond_Reflection.jpg"
+            - "Senso-ji_Temple_Asakusa_Cherry_Blossoms_Kimono_Umbrella.jpg"
+            - "Cherry_Blossoms_Sakura_Night_View_City_Lights_Japan.jpg"
+            - "Mount_Fuji_Lake_Reflection_Cherry_Blossoms_Sakura_Spring.jpg"
+            """,
+            },
+            "gradient": {
+                "type": "string",
+                "description": 'CSS gradient string for background (e.g., "linear-gradient(135deg, #667eea 0%, #764ba2 100%)")',
+            },
+        },
+        "required": ["english", "japanese", "image_name", "gradient"],
+    },
+)
+
+create_chart = FunctionTool(
+    name="create_chart",
+    description="""Create an interactive chart (FRONTEND_RENDER).
+
+    This tool creates chart specifications for frontend rendering.
+    The frontend should render this as an interactive chart component.""",
+    func=None,  # Makes declaration_only=True so client renders the UI
+    input_model={
+        "type": "object",
+        "properties": {
+            "chart_type": {
+                "type": "string",
+                "description": "Type of chart (bar, line, pie, scatter)",
+            },
+            "data_points": {
+                "type": "array",
+                "items": {"type": "object"},
+                "description": "Data points for the chart",
+            },
+            "title": {
+                "type": "string",
+                "description": "Chart title",
+            },
+        },
+        "required": ["chart_type", "data_points", "title"],
+    },
+)
+
+display_timeline = FunctionTool(
+    name="display_timeline",
+    description="""Display an interactive timeline (FRONTEND_RENDER).
+
+    This tool creates timeline specifications for frontend rendering.
+    The frontend should render this as an interactive timeline component.""",
+    func=None,  # Makes declaration_only=True so client renders the UI
+    input_model={
+        "type": "object",
+        "properties": {
+            "events": {
+                "type": "array",
+                "items": {"type": "object"},
+                "description": "Events to display on the timeline",
+            },
+            "start_date": {
+                "type": "string",
+                "description": "Timeline start date",
+            },
+            "end_date": {
+                "type": "string",
+                "description": "Timeline end date",
+            },
+        },
+        "required": ["events", "start_date", "end_date"],
+    },
+)
+
+show_comparison_table = FunctionTool(
+    name="show_comparison_table",
+    description="""Show a comparison table (FRONTEND_RENDER).
+
+    This tool creates table specifications for frontend rendering.
+    The frontend should render this as an interactive comparison table.""",
+    func=None,  # Makes declaration_only=True so client renders the UI
+    input_model={
+        "type": "object",
+        "properties": {
+            "items": {
+                "type": "array",
+                "items": {"type": "object"},
+                "description": "Items to compare",
+            },
+            "columns": {
+                "type": "array",
+                "items": {"type": "string"},
+                "description": "Column names",
+            },
+        },
+        "required": ["items", "columns"],
+    },
+)
+
+
+_UI_GENERATOR_INSTRUCTIONS = """You MUST use the provided tools to generate content. Never respond with plain text descriptions.
+
+    For haiku requests:
+    - Call generate_haiku tool with all 4 required parameters
+    - English: 3 lines
+    - Japanese: 3 lines
+    - image_name: Choose from available images
+    - gradient: CSS gradient string
+
+    For other requests, use the appropriate tool (create_chart, display_timeline, show_comparison_table).
+    """
+
+OptionsT = TypeVar("OptionsT", bound=TypedDict, default="ChatOptions")  # type: ignore[valid-type]
+
+
+def ui_generator_agent(client: SupportsChatGetResponse[OptionsT]) -> AgentFrameworkAgent:
+    """Create a UI generator agent with custom React component rendering.
+
+    Args:
+        client: The chat client to use for the agent
+
+    Returns:
+        A configured AgentFrameworkAgent instance with UI generation capabilities
+    """
+    agent = Agent(
+        name="ui_generator",
+        instructions=_UI_GENERATOR_INSTRUCTIONS,
+        client=client,
+        tools=[generate_haiku, create_chart, display_timeline, show_comparison_table],
+        # Force tool usage - the LLM MUST call a tool, cannot respond with plain text
+        default_options={"tool_choice": "required"},  # type: ignore
+    )
+
+    return AgentFrameworkAgent(
+        agent=agent,
+        name="UIGenerator",
+        description="Generates custom UI components through tool calls",
+    )

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/weather_agent.py

@@ -0,0 +1,81 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Weather agent example demonstrating backend tool rendering."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from agent_framework import Agent, SupportsChatGetResponse, tool
+
+
+@tool
+def get_weather(location: str) -> dict[str, Any]:
+    """Get the current weather for a location.
+
+    Args:
+        location: The city or location to get weather for.
+
+    Returns:
+        Weather information as a dictionary with temperatures in Celsius.
+    """
+    # Simulated weather data with structured format (temperatures in Celsius for dojo UI)
+    weather_data = {
+        "seattle": {"temperature": 11, "conditions": "rainy", "humidity": 75, "wind_speed": 12, "feels_like": 10},
+        "san francisco": {"temperature": 14, "conditions": "foggy", "humidity": 85, "wind_speed": 8, "feels_like": 13},
+        "new york city": {"temperature": 18, "conditions": "sunny", "humidity": 60, "wind_speed": 10, "feels_like": 17},
+        "miami": {"temperature": 29, "conditions": "hot and humid", "humidity": 90, "wind_speed": 5, "feels_like": 32},
+        "chicago": {"temperature": 9, "conditions": "windy", "humidity": 65, "wind_speed": 20, "feels_like": 6},
+    }
+
+    location_lower = location.lower()
+    if location_lower in weather_data:
+        return weather_data[location_lower]
+
+    return {
+        "temperature": 21,
+        "conditions": "partly cloudy",
+        "humidity": 50,
+        "wind_speed": 10,
+        "feels_like": 20,
+    }
+
+
+@tool
+def get_forecast(location: str, days: int = 3) -> str:
+    """Get the weather forecast for a location.
+
+    Args:
+        location: The city or location to get forecast for.
+        days: Number of days to forecast (default: 3).
+
+    Returns:
+        Forecast information string.
+    """
+    forecast: list[str] = []
+    for day in range(1, min(days, 7) + 1):
+        forecast.append(f"Day {day}: Partly cloudy, {60 + day * 2}°F")
+
+    return f"{days}-day forecast for {location}:\n" + "\n".join(forecast)
+
+
+def weather_agent(client: SupportsChatGetResponse[Any]) -> Agent[Any]:
+    """Create a weather agent with get_weather and get_forecast tools.
+
+    Args:
+        client: The chat client to use for the agent
+
+    Returns:
+        A configured Agent instance with weather tools
+    """
+    return Agent[Any](
+        name="weather_agent",
+        instructions=(
+            "You are a helpful weather assistant. "
+            "Use the get_weather and get_forecast functions to help users with weather information. "
+            "Always provide friendly and informative responses. "
+            "First return the weather result, and then return details about the forecast."
+        ),
+        client=client,
+        tools=[get_weather, get_forecast],
+    )

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/weather_state_agent.py

@@ -0,0 +1,92 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Deterministic tool-driven AG-UI state example.
+
+This sample demonstrates how a tool can push a *deterministic* state update
+to the AG-UI frontend based on its actual return value — in contrast to
+``predict_state_config`` which fires optimistically from LLM-predicted tool
+call arguments. See issue https://github.com/microsoft/agent-framework/issues/3167.
+
+The :func:`agent_framework_ag_ui.state_update` helper wraps a text result
+together with a state snapshot. When a tool returns one of these, the AG-UI
+endpoint merges the snapshot into the shared state and emits a
+``StateSnapshotEvent`` after the tool result.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from agent_framework import Agent, Content, SupportsChatGetResponse, tool
+from agent_framework.ag_ui import AgentFrameworkAgent
+
+from agent_framework_ag_ui import state_update
+
+# Simulated weather database — in the issue's motivating example the tool
+# would instead call a real weather API.
+_WEATHER_DB: dict[str, dict[str, Any]] = {
+    "seattle": {"temperature": 11, "conditions": "rainy", "humidity": 75},
+    "san francisco": {"temperature": 14, "conditions": "foggy", "humidity": 85},
+    "new york city": {"temperature": 18, "conditions": "sunny", "humidity": 60},
+    "miami": {"temperature": 29, "conditions": "hot and humid", "humidity": 90},
+    "chicago": {"temperature": 9, "conditions": "windy", "humidity": 65},
+}
+
+
+@tool
+async def get_weather(location: str) -> Content:
+    """Fetch current weather for a location and push it into AG-UI shared state.
+
+    Unlike ``predict_state_config`` — which derives state optimistically from
+    LLM-predicted tool call arguments — this tool uses ``state_update`` to
+    forward the *actual* fetched weather to the frontend. The ``text`` goes
+    back to the LLM as the normal tool result, and the ``state`` dict is merged
+    into the AG-UI shared state.
+
+    Args:
+        location: City name to look up.
+
+    Returns:
+        A :class:`Content` carrying both the LLM-visible text result and a
+        deterministic state snapshot.
+    """
+    key = location.lower()
+    data = _WEATHER_DB.get(
+        key,
+        {"temperature": 21, "conditions": "partly cloudy", "humidity": 50},
+    )
+    weather_record = {"location": location, **data}
+    return state_update(
+        text=(
+            f"The weather in {location} is {data['conditions']} at "
+            f"{data['temperature']}°C with {data['humidity']}% humidity."
+        ),
+        state={"weather": weather_record},
+    )
+
+
+def weather_state_agent(client: SupportsChatGetResponse[Any]) -> AgentFrameworkAgent:
+    """Create an AG-UI agent with a deterministic tool-driven state tool."""
+    agent = Agent[Any](
+        name="weather_state_agent",
+        instructions=(
+            "You are a weather assistant. When a user asks about the weather "
+            "in a city, call the get_weather tool and use its output to give a "
+            "friendly, concise reply. The tool also updates the shared UI state "
+            "so the frontend can render a weather card from the `weather` key."
+        ),
+        client=client,
+        tools=[get_weather],
+    )
+
+    return AgentFrameworkAgent(
+        agent=agent,
+        name="WeatherStateAgent",
+        description="Weather agent that deterministically updates shared state from tool results.",
+        state_schema={
+            "weather": {
+                "type": "object",
+                "description": "Last fetched weather record",
+            },
+        },
+    )

python/packages/ag-ui/agent_framework_ag_ui_examples/server/__init__.py

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

python/packages/ag-ui/agent_framework_ag_ui_examples/server/api/backend_tool_rendering.py

@@ -0,0 +1,28 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Backend tool rendering endpoint."""
+
+from typing import Any, cast
+
+from agent_framework._clients import SupportsChatGetResponse
+from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
+from agent_framework.openai import OpenAIChatCompletionClient
+from fastapi import FastAPI
+
+from ...agents.weather_agent import weather_agent
+
+
+def register_backend_tool_rendering(app: FastAPI) -> None:
+    """Register the backend tool rendering endpoint.
+
+    Args:
+        app: The FastAPI application.
+    """
+    # Create a chat client and call the factory function
+    client = cast(SupportsChatGetResponse[Any], OpenAIChatCompletionClient())
+
+    add_agent_framework_fastapi_endpoint(
+        app,
+        weather_agent(client),
+        "/backend_tool_rendering",
+    )

python/packages/ag-ui/agent_framework_ag_ui_examples/server/main.py

@@ -0,0 +1,173 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Example FastAPI server with AG-UI endpoints."""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any, cast
+
+import uvicorn
+from agent_framework import ChatOptions
+from agent_framework._clients import SupportsChatGetResponse
+from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
+from agent_framework.openai import OpenAIChatCompletionClient
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+
+from ..agents.document_writer_agent import document_writer_agent
+from ..agents.human_in_the_loop_agent import human_in_the_loop_agent
+from ..agents.recipe_agent import recipe_agent
+from ..agents.simple_agent import simple_agent
+from ..agents.subgraphs_agent import subgraphs_agent
+from ..agents.task_steps_agent import task_steps_agent_wrapped
+from ..agents.ui_generator_agent import ui_generator_agent
+from ..agents.weather_agent import weather_agent
+from ..agents.weather_state_agent import weather_state_agent
+
+AnthropicClient: type[Any] | None
+try:
+    import agent_framework.anthropic as _anthropic_namespace
+except ImportError:
+    # If the Anthropic client isn't installed, we can still run the server with Azure OpenAI as the default chat client
+    AnthropicClient = None
+else:
+    AnthropicClient = cast(type[Any] | None, getattr(_anthropic_namespace, "AnthropicClient", None))
+
+# Configure logging to file and console (disabled by default - set ENABLE_DEBUG_LOGGING=1 to enable)
+if os.getenv("ENABLE_DEBUG_LOGGING"):
+    log_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "ag_ui_server.log")
+
+    # Remove any existing handlers
+    root_logger = logging.getLogger()
+    for handler in root_logger.handlers[:]:
+        root_logger.removeHandler(handler)
+
+    # Configure new handlers
+    file_handler = logging.FileHandler(log_file, mode="w")
+    file_handler.setLevel(logging.INFO)
+    file_handler.setFormatter(logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s"))
+
+    console_handler = logging.StreamHandler()
+    console_handler.setLevel(logging.INFO)
+    console_handler.setFormatter(logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s"))
+
+    root_logger.addHandler(file_handler)
+    root_logger.addHandler(console_handler)
+    root_logger.setLevel(logging.INFO)
+
+    # Explicitly set log levels for our modules
+    logging.getLogger("agent_framework_ag_ui").setLevel(logging.INFO)
+    logging.getLogger("agent_framework").setLevel(logging.INFO)
+
+    logger = logging.getLogger(__name__)
+    logger.info(f"AG-UI Examples Server starting... Logs writing to: {log_file}")
+
+app = FastAPI(title="Agent Framework AG-UI Example Server")
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Create a shared chat client for all agents
+# You can use different chat clients for different agents if needed
+# Set CHAT_CLIENT=anthropic to use Anthropic, defaults to Azure OpenAI
+client: SupportsChatGetResponse[ChatOptions] = cast(
+    SupportsChatGetResponse[ChatOptions],
+    AnthropicClient()
+    if AnthropicClient is not None and os.getenv("CHAT_CLIENT", "").lower() == "anthropic"
+    else OpenAIChatCompletionClient(),
+)
+
+# Agentic Chat - basic chat agent
+add_agent_framework_fastapi_endpoint(
+    app=app,
+    agent=simple_agent(client),
+    path="/agentic_chat",
+)
+
+# Backend Tool Rendering - agent with tools
+add_agent_framework_fastapi_endpoint(
+    app=app,
+    agent=weather_agent(client),
+    path="/backend_tool_rendering",
+)
+
+# Shared State - recipe agent with structured output
+add_agent_framework_fastapi_endpoint(
+    app=app,
+    agent=recipe_agent(client),
+    path="/shared_state",
+)
+
+# Predictive State Updates - document writer with predictive state
+add_agent_framework_fastapi_endpoint(
+    app=app,
+    agent=document_writer_agent(client),
+    path="/predictive_state_updates",
+)
+
+# Human in the Loop - human-in-the-loop agent with step customization
+add_agent_framework_fastapi_endpoint(
+    app=app,
+    agent=human_in_the_loop_agent(client),
+    path="/human_in_the_loop",
+    state_schema={"steps": {"type": "array"}},
+    predict_state_config={"steps": {"tool": "generate_task_steps", "tool_argument": "steps"}},
+)
+
+# Agentic Generative UI - task steps agent with streaming state updates
+add_agent_framework_fastapi_endpoint(
+    app=app,
+    agent=task_steps_agent_wrapped(client),
+    path="/agentic_generative_ui",
+)
+
+# Tool-based Generative UI - UI generator with frontend-rendered tools
+add_agent_framework_fastapi_endpoint(
+    app=app,
+    agent=ui_generator_agent(client),
+    path="/tool_based_generative_ui",
+)
+
+# Subgraphs - deterministic travel planner with interrupt-driven selections
+add_agent_framework_fastapi_endpoint(
+    app=app,
+    agent=subgraphs_agent(),
+    path="/subgraphs",
+)
+
+# Deterministic Tool-Driven State - tool returns state_update() to push snapshot
+# from actual tool output (see issue #3167).
+add_agent_framework_fastapi_endpoint(
+    app=app,
+    agent=weather_state_agent(client),
+    path="/deterministic_state",
+)
+
+
+def main():
+    """Run the server."""
+    port = int(os.getenv("PORT", "8887"))
+    host = os.getenv("HOST", "127.0.0.1")
+
+    print(f"\nAG-UI Examples Server starting on http://{host}:{port}")
+    print("Set ENABLE_DEBUG_LOGGING=1 for detailed request logging\n")
+
+    # Use log_config=None to prevent uvicorn from reconfiguring logging
+    # This preserves our file + console logging setup
+    uvicorn.run(
+        app,
+        host=host,
+        port=port,
+        log_config=None,
+    )
+
+
+if __name__ == "__main__":
+    main()

python/packages/ag-ui/getting_started/README.md

@@ -0,0 +1,486 @@
+# Getting Started with AG-UI (Python)
+
+The AG-UI (Agent UI) protocol provides a standardized way for client applications to interact with AI agents over HTTP. This tutorial demonstrates how to build both server and client applications using the AG-UI protocol with Python.
+
+## Quick Start - Client Examples
+
+If you want to quickly try out the AG-UI client, we provide three ready-to-use examples:
+
+### Basic Interactive Client (`client.py`)
+
+A simple command-line chat client that demonstrates:
+- Streaming responses in real-time
+- Automatic thread management for conversation continuity
+- Direct `AGUIChatClient` usage (caller manages message history)
+
+**Run:**
+```bash
+python client.py
+```
+
+**Note:** This example sends only the current message to the server. The server is responsible for maintaining conversation history using the thread_id.
+
+### Advanced Features Client (`client_advanced.py`)
+
+Demonstrates advanced capabilities:
+- Tool/function calling
+- Both streaming and non-streaming responses
+- Multi-turn conversations
+- Error handling patterns
+
+**Run:**
+```bash
+python client_advanced.py
+```
+
+**Note:** This example shows direct `AGUIChatClient` usage. Tool execution and conversation continuity depend on server-side configuration and capabilities.
+
+### Agent Integration (`client_with_agent.py`)
+
+Best practice example using `Agent` wrapper with **AgentThread**
+- **AgentThread** maintains conversation state
+- Client-side conversation history management via `thread.message_store`
+- **Hybrid tool execution**: client-side + server-side tools simultaneously
+- Full conversation history sent on each request
+- Tool calling with conversation context
+
+**To demonstrate hybrid tools:**
+
+1. **Start server with server-side tool** (Terminal 1):
+   ```bash
+   # Server has get_time_zone tool
+   python server.py
+   ```
+
+2. **Run client with client-side tool** (Terminal 2):
+   ```bash
+   # Client has get_weather tool
+   python client_with_agent.py
+   ```
+
+All examples require a running AG-UI server (see Step 1 below for setup).
+
+## Understanding AG-UI Architecture
+
+### Thread Management
+
+The AG-UI protocol supports two approaches to conversation history:
+
+1. **Server-Managed Threads** (client.py, client_advanced.py)
+   - Client sends only the current message + thread_id
+   - Server maintains full conversation history
+   - Requires server to support stateful thread storage
+   - Lighter network payload
+
+2. **Client-Managed History** (client_with_agent.py)
+   - Client maintains full conversation history locally
+   - Full message history sent with each request
+   - Works with any AG-UI server (stateful or stateless)
+
+The `Agent` wrapper (used in client_with_agent.py) collects messages from local storage and sends the full history to `AGUIChatClient`, which then forwards everything to the server.
+
+### Tool/Function Calling
+
+The AG-UI protocol supports **hybrid tool execution** - both client-side AND server-side tools can coexist in the same conversation.
+
+**The Hybrid Pattern** (client_with_agent.py):
+```
+Client defines:           Server defines:
+- get_weather()          - get_current_time()
+- read_sensors()         - get_server_forecast()
+
+User: "What's the weather in SF and what time is it?"
+    ↓
+Agent sends: full history + tool definitions for get_weather, read_sensors
+    ↓
+Server LLM decides: "I need get_weather('SF') and get_current_time()"
+    ↓
+Server executes get_current_time() → "2025-11-11 14:30:00 UTC"
+Server sends function call request → get_weather('SF')
+    ↓
+Agent intercepts get_weather call → executes locally
+    ↓
+Client sends result → "Sunny, 72°F"
+    ↓
+Server combines both results → "It's sunny and 72°F in SF, and the current time is 2:30 PM UTC"
+    ↓
+Client receives final response
+```
+
+**How it works:**
+
+1. **Client-Side Tools** (`client_with_agent.py`):
+   - Tools defined in Agent's `tools` parameter execute locally
+   - Tool metadata (name, description, schema) sent to server for planning
+   - When server requests client tool → client intercepts → executes locally → sends result
+
+2. **Server-Side Tools**:
+   - Defined in server agent's configuration
+   - Server executes directly without client involvement
+   - Results included in server's response
+
+3. **Hybrid Pattern (Both Together)**:
+   - Server LLM sees ALL tool definitions (client + server)
+   - Decides which to use based on task
+   - Server tools execute server-side
+   - Client tools execute client-side
+
+**Direct AGUIChatClient Usage** (client_advanced.py):
+Even without Agent wrapper, client-side tools work:
+- Tools passed in ChatOptions execute locally
+- Server can also have its own tools
+- Hybrid execution works automatically
+
+### Interrupts and Resume Entries
+
+Human-in-the-loop approvals and workflow input requests pause by emitting a terminal `RUN_FINISHED` event whose
+`outcome.type` is `"interrupt"`. Generic AG-UI clients should read prompts from `RUN_FINISHED.outcome.interrupts`
+and resume the same `threadId` with a canonical `resume` array of `ResumeEntry` values.
+
+```json
+{
+  "threadId": "thread-1",
+  "messages": [],
+  "resume": [
+    {
+      "interruptId": "approval_1",
+      "status": "resolved",
+      "payload": {
+        "approved": true
+      }
+    }
+  ]
+}
+```
+
+`Interrupt` and `ResumeEntry` are AG-UI protocol models from `ag_ui.core`; Agent Framework does not define a
+separate interrupt model. New interrupted runs use `RUN_FINISHED.outcome.interrupts`, not a stable top-level
+`RUN_FINISHED.interrupt` field.
+
+## What is AG-UI?
+
+AG-UI is a protocol that enables:
+- **Remote agent hosting**: Host AI agents as web services that can be accessed by multiple clients
+- **Streaming responses**: Real-time streaming of agent responses using Server-Sent Events (SSE)
+- **Standardized communication**: Consistent message format for agent interactions
+- **Thread management**: Maintain conversation context across multiple requests
+- **Advanced features**: Human-in-the-loop, state management, tool rendering
+
+## Prerequisites
+
+Before you begin, ensure you have the following:
+
+- Python 3.10 or later
+- Azure OpenAI service endpoint and deployment configured
+- Azure CLI installed and authenticated (for DefaultAzureCredential)
+- User has the `Cognitive Services OpenAI Contributor` role for the Azure OpenAI resource
+
+**Note**: These samples use Azure OpenAI models. For more information, see [how to deploy Azure OpenAI models with Azure AI Foundry](https://learn.microsoft.com/en-us/azure/ai-foundry/how-to/deploy-models-openai).
+
+**Note**: These samples use `DefaultAzureCredential` for authentication. Make sure you're authenticated with Azure (e.g., via `az login`, or environment variables). For more information, see the [Azure Identity documentation](https://learn.microsoft.com/python/api/azure-identity/azure.identity.defaultazurecredential).
+
+> **Warning**
+> The AG-UI protocol is still under development and subject to change.
+> We will keep these samples updated as the protocol evolves.
+
+## Step 1: Creating an AG-UI Server
+
+The AG-UI server hosts your AI agent and exposes it via HTTP endpoints using FastAPI.
+
+### Install Required Packages
+
+```bash
+pip install agent-framework-ag-ui
+```
+
+Or using uv:
+
+```bash
+uv pip install agent-framework-ag-ui
+```
+
+### Server Code
+
+Create a file named `server.py`:
+
+```python
+# Copyright (c) Microsoft. All rights reserved.
+
+"""AG-UI server example."""
+
+import os
+
+from agent_framework import Agent
+from agent_framework.openai import OpenAIChatCompletionClient
+from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
+from fastapi import FastAPI
+
+# Read required configuration
+endpoint = os.environ.get("AZURE_OPENAI_ENDPOINT")
+model = os.environ.get("AZURE_OPENAI_MODEL")
+api_key = os.environ.get("AZURE_OPENAI_API_KEY")
+
+if not endpoint:
+    raise ValueError("AZURE_OPENAI_ENDPOINT environment variable is required")
+if not model:
+    raise ValueError("AZURE_OPENAI_MODEL environment variable is required")
+if not api_key:
+    raise ValueError("AZURE_OPENAI_API_KEY environment variable is required")
+
+# Create the AI agent
+agent = Agent(
+    name="AGUIAssistant",
+    instructions="You are a helpful assistant.",
+    client=OpenAIChatCompletionClient(
+        azure_endpoint=endpoint,
+        model=model,
+        api_key=api_key,
+    ),
+)
+
+# Create FastAPI app
+app = FastAPI(title="AG-UI Server")
+
+# Register the AG-UI endpoint
+add_agent_framework_fastapi_endpoint(app, agent, "/")
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run(app, host="127.0.0.1", port=5100)
+```
+
+### Key Concepts
+
+- **`add_agent_framework_fastapi_endpoint`**: Registers the AG-UI endpoint with automatic request/response handling and SSE streaming
+- **`Agent`**: The agent that will handle incoming requests
+- **FastAPI Integration**: Uses FastAPI's native async support for streaming responses
+- **Instructions**: The agent is created with default instructions, which can be overridden by client messages
+- **Configuration**: `OpenAIChatCompletionClient` can read from environment variables (`AZURE_OPENAI_ENDPOINT`, `AZURE_OPENAI_MODEL`, `AZURE_OPENAI_API_KEY`) or accept parameters directly
+
+**Alternative (simpler)**: Use environment variables only:
+
+```python
+# No need to read environment variables manually
+agent = Agent(
+    name="AGUIAssistant",
+    instructions="You are a helpful assistant.",
+    client=OpenAIChatCompletionClient(),  # Reads from environment automatically
+)
+```
+
+### Configure and Run the Server
+
+Set the required environment variables:
+
+```bash
+export AZURE_OPENAI_ENDPOINT="https://your-resource.openai.azure.com/"
+export AZURE_OPENAI_MODEL="gpt-4o-mini"
+# Optional: Set API key if not using DefaultAzureCredential
+# export AZURE_OPENAI_API_KEY="your-api-key"
+```
+
+Run the server:
+
+```bash
+python server.py
+```
+
+Or using uvicorn directly:
+
+```bash
+uvicorn server:app --host 127.0.0.1 --port 5100
+```
+
+The server will start listening on `http://127.0.0.1:5100`.
+
+## Step 2: Creating an AG-UI Client
+
+The AG-UI client connects to the remote server and displays streaming responses. The `AGUIChatClient` is a built-in implementation that integrates with the Agent Framework's standard chat interface.
+
+### Install Required Packages
+
+The `AGUIChatClient` is included in the `agent-framework-ag-ui` package (already installed if you installed the server packages).
+
+```bash
+pip install agent-framework-ag-ui
+```
+
+### Client Code
+
+Create a file named `client.py`:
+
+```python
+# Copyright (c) Microsoft. All rights reserved.
+
+"""AG-UI client example using AGUIChatClient."""
+
+import asyncio
+import os
+
+from agent_framework.ag_ui import AGUIChatClient
+
+
+async def main():
+    """Main client loop demonstrating AGUIChatClient usage."""
+    # Get server URL from environment or use default
+    server_url = os.environ.get("AGUI_SERVER_URL", "http://127.0.0.1:5100/")
+    print(f"Connecting to AG-UI server at: {server_url}\n")
+
+    # Create client with context manager for automatic cleanup
+    async with AGUIChatClient(endpoint=server_url) as client:
+        thread_id: str | None = None
+
+        try:
+            while True:
+                # Get user input
+                message = input("\nUser (:q or quit to exit): ")
+                if not message.strip():
+                    print("Request cannot be empty.")
+                    continue
+
+                if message.lower() in (":q", "quit"):
+                    break
+
+                # Send message and stream the response
+                print("\nAssistant: ", end="", flush=True)
+
+                # Use metadata to maintain conversation continuity
+                metadata = {"thread_id": thread_id} if thread_id else None
+
+                async for update in client.get_response(message, metadata=metadata, stream=True):
+                    # Extract thread ID from first update
+                    if not thread_id and update.additional_properties:
+                        thread_id = update.additional_properties.get("thread_id")
+                        if thread_id:
+                            print(f"\n[Thread: {thread_id}]")
+                            print("Assistant: ", end="", flush=True)
+
+                    # Stream text content as it arrives
+                    for content in update.contents:
+                        if content.type == "text" and content.text:
+                            print(content.text, end="", flush=True)
+
+                print()  # New line after response
+
+        except KeyboardInterrupt:
+            print("\n\nExiting...")
+        except Exception as e:
+            print(f"\nAn error occurred: {e}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())
+```
+
+### Key Concepts
+
+- **`AGUIChatClient`**: Built-in client that implements the Agent Framework's `BaseChatClient` interface
+- **Automatic Event Handling**: The client automatically converts AG-UI events to Agent Framework types
+- **Thread Management**: Pass `thread_id` in metadata to maintain conversation context across requests
+- **Streaming Responses**: Use `get_response(..., stream=True)` for real-time streaming or `get_response(..., stream=False)` for non-streaming
+- **Context Manager**: Use `async with` for automatic cleanup of HTTP connections
+- **Standard Interface**: Works with all Agent Framework patterns (Agent, tools, etc.)
+- **Hybrid Tool Execution**: Supports both client-side and server-side tools executing together in the same conversation
+
+### Configure and Run the Client
+
+Optionally set a custom server URL:
+
+```bash
+export AGUI_SERVER_URL="http://127.0.0.1:5100/"
+```
+
+Run the client (in a separate terminal):
+
+```bash
+python client.py
+```
+
+## Step 3: Testing the Complete System
+
+### Expected Output
+
+```
+$ python client.py
+Connecting to AG-UI server at: http://127.0.0.1:5100/
+
+User (:q or quit to exit): What is the capital of France?
+
+[Thread: abc123]
+Assistant: The capital of France is Paris. It is known for its rich history, culture,
+and iconic landmarks such as the Eiffel Tower and the Louvre Museum.
+
+User (:q or quit to exit): Tell me a fun fact about space
+```
+
+## Troubleshooting
+
+### Connection Refused
+
+Ensure the server is running before starting the client:
+
+```bash
+# Terminal 1
+python server.py
+
+# Terminal 2 (after server starts)
+python client.py
+```
+
+### Authentication Errors
+
+Make sure you're authenticated with Azure:
+
+```bash
+az login
+```
+
+Verify you have the correct role assignment on the Azure OpenAI resource.
+
+### Streaming Not Working
+
+Check that your client timeout is sufficient:
+
+```python
+httpx.AsyncClient(timeout=60.0)  # 60 seconds should be enough
+```
+
+For long-running agents, increase the timeout accordingly.
+
+### No Events Received
+
+Ensure you're using the correct `Accept` header:
+
+```python
+headers={"Accept": "text/event-stream"}
+```
+
+And parsing SSE format correctly (lines starting with `data: `).
+
+### Thread Context Lost
+
+The client automatically manages thread continuity. If context is lost:
+
+1. Check that `threadId` is being captured from `RUN_STARTED` events
+2. Ensure the same client instance is used across messages
+3. Verify the server is receiving the `thread_id` in subsequent requests
+
+### Event Type Mismatches
+
+Remember that event types are UPPERCASE with underscores (`RUN_STARTED`, not `run_started`) and field names are camelCase (`threadId`, not `thread_id`).
+
+### Import Errors
+
+Make sure all packages are installed:
+
+```bash
+pip install agent-framework-ag-ui agent-framework-core fastapi uvicorn httpx
+```
+
+Or check your virtual environment is activated:
+
+```bash
+source venv/bin/activate  # Linux/macOS
+venv\Scripts\activate     # Windows
+```

python/packages/ag-ui/getting_started/client.py

@@ -0,0 +1,78 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""AG-UI client example using AGUIChatClient.
+
+This example demonstrates how to use the AGUIChatClient to connect to
+a remote AG-UI server and interact with it using the Agent Framework's
+standard chat interface.
+"""
+
+import asyncio
+import os
+from typing import cast
+
+from agent_framework import ChatResponse, ChatResponseUpdate, Message, ResponseStream
+from agent_framework.ag_ui import AGUIChatClient
+
+
+async def main():
+    """Main client loop demonstrating AGUIChatClient usage."""
+    # Get server URL from environment or use default
+    server_url = os.environ.get("AGUI_SERVER_URL", "http://127.0.0.1:5100/")
+    print(f"Connecting to AG-UI server at: {server_url}\n")
+    print("Using AGUIChatClient with automatic thread management and Agent Framework integration.\n")
+
+    # Create client with context manager for automatic cleanup
+    async with AGUIChatClient(endpoint=server_url) as client:
+        thread_id: str | None = None
+
+        try:
+            while True:
+                # Get user input
+                message = input("\nUser (:q or quit to exit): ")
+                if not message.strip():
+                    print("Request cannot be empty.")
+                    continue
+
+                if message.lower() in (":q", "quit"):
+                    break
+
+                # Send message and stream the response
+                print("\nAssistant: ", end="", flush=True)
+
+                # Use metadata to maintain conversation continuity
+                metadata = {"thread_id": thread_id} if thread_id else None
+
+                stream = client.get_response(
+                    [Message(role="user", contents=[message])],
+                    stream=True,
+                    options={"metadata": metadata} if metadata else None,
+                )
+                stream = cast(ResponseStream[ChatResponseUpdate, ChatResponse], stream)
+                async for update in stream:
+                    # Extract and display thread ID from first update
+                    if not thread_id and update.additional_properties:
+                        thread_id = update.additional_properties.get("thread_id")
+                        if thread_id:
+                            print(f"\n\033[93m[Thread: {thread_id}]\033[0m", end="", flush=True)
+                            print("\nAssistant: ", end="", flush=True)
+
+                    # Display text content as it streams
+                    for content in update.contents:
+                        if content.type == "text" and content.text:
+                            print(f"\033[96m{content.text}\033[0m", end="", flush=True)
+
+                    # Display finish reason if present
+                    if update.finish_reason:
+                        print(f"\n\033[92m[Finished: {update.finish_reason}]\033[0m", end="", flush=True)
+
+                print()  # New line after response
+
+        except KeyboardInterrupt:
+            print("\n\nExiting...")
+        except Exception as e:
+            print(f"\n\033[91mAn error occurred: {e}\033[0m")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

python/packages/ag-ui/getting_started/client_advanced.py

@@ -0,0 +1,245 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Advanced AG-UI client example with tools and features.
+
+This example demonstrates advanced AGUIChatClient features including:
+- Tool/function calling
+- Non-streaming responses
+- Multiple conversation turns
+- Error handling
+"""
+
+from __future__ import annotations
+
+import asyncio
+import os
+from typing import cast
+
+from agent_framework import ChatResponse, ChatResponseUpdate, Message, ResponseStream, tool
+from agent_framework.ag_ui import AGUIChatClient
+
+
+@tool
+def get_weather(location: str) -> str:
+    """Get the current weather for a location.
+
+    Args:
+        location: The city or location name
+    """
+    # Simulate weather lookup
+    weather_data = {
+        "seattle": "Rainy, 55°F",
+        "san francisco": "Foggy, 62°F",
+        "new york": "Sunny, 68°F",
+        "london": "Cloudy, 52°F",
+    }
+    return weather_data.get(location.lower(), f"Weather data not available for {location}")
+
+
+@tool
+def calculate(a: float, b: float, operation: str) -> str:
+    """Perform basic arithmetic operations.
+
+    Args:
+        a: First number
+        b: Second number
+        operation: Operation to perform (add, subtract, multiply, divide)
+    """
+    try:
+        if operation == "add":
+            result = a + b
+        elif operation == "subtract":
+            result = a - b
+        elif operation == "multiply":
+            result = a * b
+        elif operation == "divide":
+            result = a / b
+        else:
+            return f"Unsupported operation: {operation}"
+        return f"The result is: {result}"
+    except Exception as e:
+        return f"Error calculating: {e}"
+
+
+async def streaming_example(client: AGUIChatClient, thread_id: str | None = None):
+    """Demonstrate streaming responses."""
+    print("\n" + "=" * 60)
+    print("STREAMING EXAMPLE")
+    print("=" * 60)
+
+    metadata = {"thread_id": thread_id} if thread_id else None
+
+    print("\nUser: Tell me a short joke\n")
+    print("Assistant: ", end="", flush=True)
+
+    stream = client.get_response(
+        [Message(role="user", contents=["Tell me a short joke"])],
+        stream=True,
+        options={"metadata": metadata} if metadata else None,
+    )
+    stream = cast(ResponseStream[ChatResponseUpdate, ChatResponse], stream)
+    async for update in stream:
+        if not thread_id and update.additional_properties:
+            thread_id = update.additional_properties.get("thread_id")
+
+        for content in update.contents:
+            if content.type == "text" and content.text:  # type: ignore[attr-defined]
+                print(content.text, end="", flush=True)  # type: ignore[attr-defined]
+
+    print("\n")
+    return thread_id
+
+
+async def non_streaming_example(client: AGUIChatClient, thread_id: str | None = None):
+    """Demonstrate non-streaming responses."""
+    print("\n" + "=" * 60)
+    print("NON-STREAMING EXAMPLE")
+    print("=" * 60)
+
+    metadata = {"thread_id": thread_id} if thread_id else None
+
+    print("\nUser: What is 2 + 2?\n")
+
+    response = await client.get_response([Message(role="user", contents=["What is 2 + 2?"])], metadata=metadata)
+
+    print(f"Assistant: {response.text}")
+
+    if response.additional_properties:
+        thread_id = response.additional_properties.get("thread_id")
+        print(f"\n[Thread: {thread_id}]")
+
+    return thread_id
+
+
+async def tool_example(client: AGUIChatClient, thread_id: str | None = None):
+    """Demonstrate sending tool definitions to the server.
+
+    IMPORTANT: When using AGUIChatClient directly (without Agent wrapper):
+    - Tools are sent as DEFINITIONS only
+    - No automatic client-side execution (no function invocation middleware)
+    - Server must have matching tool implementations to execute them
+
+    For CLIENT-SIDE tool execution (like .NET AGUIClient sample):
+    - Use Agent wrapper with tools
+    - See client_with_agent.py for the hybrid pattern
+    - Agent middleware intercepts and executes client tools locally
+    - Server can have its own tools that execute server-side
+    - Both client and server tools work together in same conversation
+
+    This example sends tool definitions and assumes server-side execution.
+    """
+    print("\n" + "=" * 60)
+    print("TOOL DEFINITION EXAMPLE")
+    print("=" * 60)
+
+    metadata = {"thread_id": thread_id} if thread_id else None
+
+    print("\nUser: What's the weather in Seattle?\n")
+    print("Sending tool definitions to server...")
+    print("(Server must be configured with matching tools to execute them)\n")
+
+    response = await client.get_response(
+        [Message(role="user", contents=["What's the weather in Seattle?"])],
+        tools=[get_weather, calculate],
+        metadata=metadata,
+    )
+
+    print(f"Assistant: {response.text}")
+
+    # Show tool calls if any
+    tool_called = False
+    for message in response.messages:
+        for content in message.contents:
+            if content.type == "function_call":  # type: ignore[attr-defined]
+                print(f"\n[Tool Called: {content.name}]")  # type: ignore[attr-defined]
+                tool_called = True
+
+    if not tool_called:
+        print("\n[Note: No tools were called - server may not be configured for tool execution]")
+
+    if response.additional_properties:
+        thread_id = response.additional_properties.get("thread_id")
+
+    return thread_id
+
+
+async def conversation_example(client: AGUIChatClient):
+    """Demonstrate multi-turn conversation.
+
+    Note: Conversation continuity depends on the server maintaining thread state.
+    Some servers may require explicit message history to be sent with each request.
+    """
+    print("\n" + "=" * 60)
+    print("MULTI-TURN CONVERSATION EXAMPLE")
+    print("=" * 60)
+    print("\nNote: This example uses thread_id for context. Server must support thread-based state.\n")
+
+    # First turn
+    print("User: My name is Alice\n")
+    response1 = await client.get_response([Message(role="user", contents=["My name is Alice"])])
+    print(f"Assistant: {response1.text}")
+    thread_id = response1.additional_properties.get("thread_id")
+    print(f"\n[Thread: {thread_id}]")
+
+    # Second turn - using same thread
+    print("\nUser: What's my name?\n")
+    response2 = await client.get_response(
+        [Message(role="user", contents=["What's my name?"])], options={"metadata": {"thread_id": thread_id}}
+    )
+    print(f"Assistant: {response2.text}")
+
+    # Check if context was maintained
+    if "alice" not in response2.text.lower():
+        print("\n[Note: Server may not maintain thread context - consider using Agent for history management]")
+
+    # Third turn
+    print("\nUser: Can you also tell me what 10 * 5 is?\n")
+    response3 = await client.get_response(
+        [Message(role="user", contents=["Can you also tell me what 10 * 5 is?"])],
+        options={"metadata": {"thread_id": thread_id}},
+        tools=[calculate],
+    )
+    print(f"Assistant: {response3.text}")
+
+
+async def main():
+    """Run all examples."""
+    # Get server URL from environment or use default
+    server_url = os.environ.get("AGUI_SERVER_URL", "http://127.0.0.1:5100/")
+
+    print("=" * 60)
+    print("AG-UI Chat Client Advanced Examples")
+    print("=" * 60)
+    print(f"\nServer: {server_url}")
+    print("\nThese examples demonstrate various AGUIChatClient features:")
+    print("  1. Streaming responses")
+    print("  2. Non-streaming responses")
+    print("  3. Tool/function calling")
+    print("  4. Multi-turn conversations")
+
+    try:
+        async with AGUIChatClient(endpoint=server_url) as client:
+            # Run examples in sequence
+            thread_id = await streaming_example(client)
+            thread_id = await non_streaming_example(client, thread_id)
+            await tool_example(client, thread_id)
+
+            # Separate conversation with new thread
+            await conversation_example(client)
+
+            print("\n" + "=" * 60)
+            print("All examples completed successfully!")
+            print("=" * 60)
+
+    except ConnectionError as e:
+        print(f"\n\033[91mConnection Error: {e}\033[0m")
+        print("\nMake sure an AG-UI server is running at the specified endpoint.")
+    except Exception as e:
+        print(f"\n\033[91mError: {e}\033[0m")
+        import traceback
+
+        traceback.print_exc()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

python/packages/ag-ui/getting_started/client_with_agent.py

@@ -0,0 +1,150 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Example showing Agent with AGUIChatClient for hybrid tool execution.
+
+This demonstrates the HYBRID pattern matching .NET AGUIClient implementation:
+
+1. AgentSession Pattern (like .NET):
+   - Create session with agent.create_session()
+   - Pass session to agent.run(stream=True) on each turn
+   - Session maintains conversation context via context providers
+
+2. Hybrid Tool Execution:
+   - AGUIChatClient uses function invocation mixin
+   - Client-side tools (get_weather) can execute locally when server requests them
+   - Server may also have its own tools that execute server-side
+   - Both work together: server LLM decides which tool to call, decorator handles client execution
+
+This matches .NET pattern: session maintains state, tools execute on appropriate side.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+
+from agent_framework import Agent, tool
+from agent_framework.ag_ui import AGUIChatClient
+
+# Enable debug logging
+logging.basicConfig(
+    level=logging.DEBUG,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger(__name__)
+
+
+@tool(description="Get the current weather for a location.")
+def get_weather(location: str) -> str:
+    """Get the current weather for a location.
+
+    Args:
+        location: The city or location name
+    """
+    print(f"[CLIENT] get_weather tool called with location: {location}")
+    weather_data = {
+        "seattle": "Rainy, 55°F",
+        "san francisco": "Foggy, 62°F",
+        "new york": "Sunny, 68°F",
+        "london": "Cloudy, 52°F",
+    }
+    result = weather_data.get(location.lower(), f"Weather data not available for {location}")
+    print(f"[CLIENT] get_weather returning: {result}")
+    return result
+
+
+async def main():
+    """Demonstrate Agent + AGUIChatClient hybrid tool execution.
+
+    This matches the .NET pattern from Program.cs where:
+    - AIAgent agent = chatClient.CreateAIAgent(tools: [...])
+    - AgentSession session = agent.CreateSession()
+    - RunStreamingAsync(messages, session)
+
+    Python equivalent:
+    - agent = Agent(client=AGUIChatClient(...), tools=[...])
+    - session = agent.create_session()  # Creates session
+    - agent.run(message, stream=True, session=session)  # Session tracks context
+    """
+    server_url = os.environ.get("AGUI_SERVER_URL", "http://127.0.0.1:5100/")
+
+    print("=" * 70)
+    print("Agent + AGUIChatClient: Hybrid Tool Execution")
+    print("=" * 70)
+    print(f"\nServer: {server_url}")
+    print("\nThis example demonstrates:")
+    print("  1. AgentSession maintains conversation state (like .NET)")
+    print("  2. Client-side tools execute locally via function invocation mixin")
+    print("  3. Server may have additional tools that execute server-side")
+    print("  4. HYBRID: Client and server tools work together simultaneously\n")
+
+    try:
+        # Create remote client in async context manager
+        async with AGUIChatClient(endpoint=server_url) as remote_client:
+            # Wrap in Agent for conversation history management
+            agent = Agent(
+                name="remote_assistant",
+                instructions="You are a helpful assistant. Remember user information across the conversation.",
+                client=remote_client,
+                tools=[get_weather],
+            )
+
+            # Create a session to maintain conversation state (like .NET AgentSession)
+            session = agent.create_session()
+
+            print("=" * 70)
+            print("CONVERSATION WITH HISTORY")
+            print("=" * 70)
+
+            # Turn 1: Introduce
+            print("\nUser: My name is Alice and I live in Seattle\n")
+            async for chunk in agent.run("My name is Alice and I live in Seattle", stream=True, session=session):
+                if chunk.text:
+                    print(chunk.text, end="", flush=True)
+            print("\n")
+
+            # Turn 2: Ask about name (tests history)
+            print("User: What's my name?\n")
+            async for chunk in agent.run("What's my name?", stream=True, session=session):
+                if chunk.text:
+                    print(chunk.text, end="", flush=True)
+            print("\n")
+
+            # Turn 3: Ask about location (tests history)
+            print("User: Where do I live?\n")
+            async for chunk in agent.run("Where do I live?", stream=True, session=session):
+                if chunk.text:
+                    print(chunk.text, end="", flush=True)
+            print("\n")
+
+            # Turn 4: Test client-side tool (get_weather is client-side)
+            print("User: What's the weather forecast for today in Seattle?\n")
+            async for chunk in agent.run(
+                "What's the weather forecast for today in Seattle?",
+                stream=True,
+                session=session,
+            ):
+                if chunk.text:
+                    print(chunk.text, end="", flush=True)
+            print("\n")
+
+            # Turn 5: Test server-side tool (get_time_zone is server-side only)
+            print("User: What time zone is Seattle in?\n")
+            async for chunk in agent.run("What time zone is Seattle in?", stream=True, session=session):
+                if chunk.text:
+                    print(chunk.text, end="", flush=True)
+            print("\n")
+
+    except ConnectionError as e:
+        print(f"\n\033[91mConnection Error: {e}\033[0m")
+        print("\nMake sure an AG-UI server is running at the specified endpoint.")
+    except Exception as e:
+        print(f"\n\033[91mError: {e}\033[0m")
+        import traceback
+
+        traceback.print_exc()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

python/packages/ag-ui/getting_started/server.py

@@ -0,0 +1,144 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""AG-UI server example with server-side tools."""
+
+from __future__ import annotations
+
+import logging
+import os
+
+from agent_framework import Agent, tool
+from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
+from agent_framework.openai import OpenAIChatCompletionClient
+from dotenv import load_dotenv
+from fastapi import Depends, FastAPI, HTTPException, Security
+from fastapi.security import APIKeyHeader
+
+load_dotenv()
+
+# Enable debug logging
+logging.basicConfig(
+    level=logging.DEBUG,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger(__name__)
+
+
+# Read required configuration
+endpoint = os.environ.get("AZURE_OPENAI_ENDPOINT")
+model = os.environ.get("AZURE_OPENAI_MODEL")
+
+if not endpoint:
+    raise ValueError("AZURE_OPENAI_ENDPOINT environment variable is required")
+if not model:
+    raise ValueError("AZURE_OPENAI_MODEL environment variable is required")
+
+
+# ============================================================================
+# AUTHENTICATION EXAMPLE
+# ============================================================================
+# This demonstrates how to secure the AG-UI endpoint with API key authentication.
+# In production, you should use a more robust authentication mechanism such as:
+# - OAuth 2.0 / OpenID Connect
+# - JWT tokens with proper validation
+# - Azure AD / Entra ID integration
+# - Your organization's identity provider
+#
+# The API key should be stored securely (e.g., Azure Key Vault, environment variables)
+# and rotated regularly.
+# ============================================================================
+
+# API key header configuration
+API_KEY_HEADER = APIKeyHeader(name="X-API-Key", auto_error=False)
+
+# Get the expected API key from environment variable
+# In production, use a secrets manager like Azure Key Vault
+EXPECTED_API_KEY = os.environ.get("AG_UI_API_KEY")
+
+
+async def verify_api_key(api_key: str | None = Security(API_KEY_HEADER)) -> None:
+    """Verify the API key provided in the request header.
+
+    Args:
+        api_key: The API key from the X-API-Key header
+
+    Raises:
+        HTTPException: If the API key is missing or invalid
+    """
+    if not EXPECTED_API_KEY:
+        # If no API key is configured, log a warning but allow the request
+        # This maintains backward compatibility but warns about the security risk
+        logger.warning(
+            "AG_UI_API_KEY environment variable not set. "
+            "The endpoint is accessible without authentication. "
+            "Set AG_UI_API_KEY to enable API key authentication."
+        )
+        return
+
+    if not api_key:
+        raise HTTPException(
+            status_code=401,
+            detail="Missing API key. Provide X-API-Key header.",
+        )
+
+    if api_key != EXPECTED_API_KEY:
+        raise HTTPException(
+            status_code=403,
+            detail="Invalid API key.",
+        )
+
+
+# Server-side tool (executes on server)
+@tool(description="Get the time zone for a location.")
+def get_time_zone(location: str) -> str:
+    """Get the time zone for a location.
+
+    Args:
+        location: The city or location name
+    """
+    print(f"[SERVER] get_time_zone tool called with location: {location}")
+    timezone_data = {
+        "seattle": "Pacific Time (UTC-8)",
+        "san francisco": "Pacific Time (UTC-8)",
+        "new york": "Eastern Time (UTC-5)",
+        "london": "Greenwich Mean Time (UTC+0)",
+    }
+    result = timezone_data.get(location.lower(), f"Time zone data not available for {location}")
+    print(f"[SERVER] get_time_zone returning: {result}")
+    return result
+
+
+# Create the AI agent with ONLY server-side tools
+# IMPORTANT: Do NOT include tools that the client provides!
+# In this example:
+# - get_time_zone: SERVER-ONLY tool (only server has this)
+# - get_weather: CLIENT-ONLY tool (client provides this, server should NOT include it)
+# The client will send get_weather tool metadata so the LLM knows about it,
+# and the function invocation mixin on AGUIChatClient will execute it client-side.
+# This matches the .NET AG-UI hybrid execution pattern.
+agent = Agent(
+    name="AGUIAssistant",
+    instructions="You are a helpful assistant. Use get_weather for weather and get_time_zone for time zones.",
+    client=OpenAIChatCompletionClient(
+        azure_endpoint=endpoint,
+        model=model,
+    ),
+    tools=[get_time_zone],  # ONLY server-side tools
+)
+
+# Create FastAPI app
+app = FastAPI(title="AG-UI Server")
+
+# Register the AG-UI endpoint with authentication
+# The dependencies parameter accepts FastAPI Depends() objects that run before the handler
+add_agent_framework_fastapi_endpoint(
+    app,
+    agent,
+    "/",
+    dependencies=[Depends(verify_api_key)],
+)
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run(app, host="127.0.0.1", port=5100, log_level="debug", access_log=True)

python/packages/ag-ui/pyproject.toml

@@ -0,0 +1,82 @@
+[project]
+name = "agent-framework-ag-ui"
+version = "1.0.0rc8"
+description = "AG-UI protocol integration for Agent Framework"
+readme = "README.md"
+license-files = ["LICENSE"]
+authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
+requires-python = ">=3.10"
+urls.homepage = "https://aka.ms/agent-framework"
+urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
+urls.release_notes = "https://github.com/microsoft/agent-framework/releases?q=tag%3Apython-1&expanded=true"
+urls.issues = "https://github.com/microsoft/agent-framework/issues"
+classifiers = [
+  "License :: OSI Approved :: MIT License",
+  "Development Status :: 4 - Beta",
+  "Intended Audience :: Developers",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+  "Programming Language :: Python :: 3.13",
+  "Typing :: Typed",
+]
+dependencies = [
+    "agent-framework-core>=1.11.0,<2",
+    "ag-ui-protocol>=0.1.19,<0.2",
+    "fastapi>=0.121.0,<0.138.1",
+    "sse-starlette>=3.4.5,<4",
+    "uvicorn[standard]>=0.30.0,<1"
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest==9.1.1",
+    "httpx==0.28.1",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["agent_framework_ag_ui", "agent_framework_ag_ui_examples"]
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+testpaths = ["tests/ag_ui"]
+pythonpath = [".", "tests/ag_ui"]
+markers = [
+    "integration: marks tests as integration tests that require external services",
+]
+
+[tool.ruff]
+line-length = 120
+target-version = "py311"
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "N", "W"]
+ignore = ["E501"]
+
+[tool.mypy]
+python_version = "3.11"
+warn_return_any = true
+warn_unused_configs = true
+disallow_untyped_defs = false
+
+[tool.pyright]
+include = ["agent_framework_ag_ui"]
+exclude = ["tests", "tests/ag_ui", "examples"]
+typeCheckingMode = "basic"
+
+[tool.poe]
+executor.type = "uv"
+include = "../../shared_tasks.toml"
+
+[tool.poe.tasks.mypy]
+help = "Run MyPy for this package."
+cmd = "mypy --config-file $POE_ROOT/pyproject.toml agent_framework_ag_ui"
+
+[tool.poe.tasks.test]
+help = "Run the default unit test suite for this package."
+cmd = 'pytest -m "not integration" --cov=agent_framework_ag_ui --cov-report=term-missing:skip-covered -n auto --dist worksteal tests/ag_ui'

python/packages/ag-ui/tests/ag_ui/conftest.py

@@ -0,0 +1,350 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Shared test fixtures and stubs for AG-UI tests."""
+
+import sys
+from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Mapping, MutableSequence, Sequence
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any, Generic, Literal, TypedDict, cast, overload  # noqa: F401
+
+import pytest
+from agent_framework import (
+    AgentResponse,
+    AgentResponseUpdate,
+    AgentSession,
+    BaseChatClient,
+    ChatOptions,
+    ChatResponse,
+    ChatResponseUpdate,
+    Content,
+    Message,
+    ServiceSessionId,
+    SupportsAgentRun,
+    SupportsChatGetResponse,
+)
+from agent_framework._clients import OptionsCoT
+from agent_framework._middleware import ChatMiddlewareLayer
+from agent_framework._tools import FunctionInvocationLayer
+from agent_framework._types import ResponseStream
+from agent_framework.observability import ChatTelemetryLayer
+
+if sys.version_info >= (3, 12):
+    from typing import override  # type: ignore # pragma: no cover
+else:
+    from typing_extensions import override  # type: ignore[import] # pragma: no cover
+
+StreamFn = Callable[..., AsyncIterable[ChatResponseUpdate]]
+ResponseFn = Callable[..., Awaitable[ChatResponse]]
+
+
+def pytest_configure() -> None:
+    """Ensure this test directory is on sys.path so helper modules can be imported by name."""
+    test_dir = str(Path(__file__).resolve().parent)
+    if test_dir not in sys.path:
+        sys.path.insert(0, test_dir)
+
+
+class StreamingChatClientStub(
+    FunctionInvocationLayer[OptionsCoT],
+    ChatMiddlewareLayer[OptionsCoT],
+    ChatTelemetryLayer[OptionsCoT],
+    BaseChatClient[OptionsCoT],
+    Generic[OptionsCoT],
+):
+    """Typed streaming stub that satisfies SupportsChatGetResponse."""
+
+    def __init__(self, stream_fn: StreamFn, response_fn: ResponseFn | None = None) -> None:
+        super().__init__(middleware=[])
+        self._stream_fn = stream_fn
+        self._response_fn = response_fn
+        self.last_session: AgentSession | None = None
+        self.last_service_session_id: str | ServiceSessionId | None = None
+
+    @overload
+    def get_response(
+        self,
+        messages: Sequence[Message],
+        *,
+        stream: Literal[False] = ...,
+        options: ChatOptions[Any],
+        **kwargs: Any,
+    ) -> Awaitable[ChatResponse[Any]]: ...
+
+    @overload
+    def get_response(
+        self,
+        messages: Sequence[Message],
+        *,
+        stream: Literal[False] = ...,
+        options: OptionsCoT | ChatOptions[None] | None = ...,
+        **kwargs: Any,
+    ) -> Awaitable[ChatResponse[Any]]: ...
+
+    @overload
+    def get_response(
+        self,
+        messages: Sequence[Message],
+        *,
+        stream: Literal[True],
+        options: OptionsCoT | ChatOptions[Any] | None = ...,
+        **kwargs: Any,
+    ) -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: ...
+
+    def get_response(
+        self,
+        messages: Sequence[Message],
+        *,
+        stream: bool = False,
+        options: OptionsCoT | ChatOptions[Any] | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]]:
+        client_kwargs = kwargs.get("client_kwargs")
+        if isinstance(client_kwargs, Mapping):
+            self.last_session = cast(AgentSession | None, client_kwargs.get("session"))
+        else:
+            self.last_session = None
+        self.last_service_session_id = self.last_session.service_session_id if self.last_session else None
+        if stream:
+            return super().get_response(
+                messages=messages,
+                stream=True,
+                options=options,
+                **kwargs,
+            )
+        return super().get_response(
+            messages=messages,
+            stream=False,
+            options=options,
+            **kwargs,
+        )
+
+    @override
+    def _inner_get_response(
+        self,
+        *,
+        messages: Sequence[Message],
+        stream: bool = False,
+        options: Mapping[str, Any],
+        **kwargs: Any,
+    ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
+        if stream:
+
+            def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse:
+                return ChatResponse.from_updates(updates)
+
+            return ResponseStream(self._stream_fn(messages, options, **kwargs), finalizer=_finalize)
+
+        return self._get_response_impl(messages, options, **kwargs)
+
+    async def _get_response_impl(
+        self, messages: Sequence[Message], options: Mapping[str, Any], **kwargs: Any
+    ) -> ChatResponse:
+        """Non-streaming implementation."""
+        if self._response_fn is not None:
+            return await self._response_fn(messages, options, **kwargs)
+
+        contents: list[Any] = []
+        async for update in self._stream_fn(list(messages), dict(options), **kwargs):
+            contents.extend(update.contents)
+
+        return ChatResponse(
+            messages=[Message(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[Message], options: dict[str, Any], **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        for update in updates:
+            yield update
+
+    return _stream
+
+
+class StubAgent(SupportsAgentRun):
+    """Minimal SupportsAgentRun stub for orchestrator tests."""
+
+    def __init__(
+        self,
+        updates: list[AgentResponseUpdate] | None = None,
+        *,
+        agent_id: str = "stub-agent",
+        agent_name: str | None = "stub-agent",
+        default_options: Any | None = None,
+        client: Any | None = None,
+    ) -> None:
+        self.id = agent_id
+        self.name = agent_name
+        self.description = "stub agent"
+        self.updates = updates or [AgentResponseUpdate(contents=[Content.from_text(text="response")], role="assistant")]
+        self.default_options: dict[str, Any] = (
+            default_options if isinstance(default_options, dict) else {"tools": None, "response_format": None}
+        )
+        self.client = client or SimpleNamespace(function_invocation_configuration=None)
+        self.messages_received: list[Any] = []
+        self.tools_received: list[Any] | None = None
+        self.last_session: AgentSession | None = None
+
+    @overload
+    def run(
+        self,
+        messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
+
+    @overload
+    def run(
+        self,
+        messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = None,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
+    def run(
+        self,
+        messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+        *,
+        stream: bool = False,
+        session: AgentSession | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
+        if stream:
+
+            async def _stream() -> AsyncIterator[AgentResponseUpdate]:
+                if messages is None:
+                    self.messages_received = []
+                elif isinstance(messages, (str, Content, Message)):
+                    self.messages_received = [messages]
+                else:
+                    self.messages_received = list(messages)
+                self.last_session = session
+                self.tools_received = kwargs.get("tools")
+                for update in self.updates:
+                    yield update
+
+            def _finalize(updates: Sequence[AgentResponseUpdate]) -> AgentResponse:
+                return AgentResponse.from_updates(updates)
+
+            return ResponseStream(_stream(), finalizer=_finalize)
+
+        async def _get_response() -> AgentResponse[Any]:
+            return AgentResponse(messages=[], response_id="stub-response")
+
+        return _get_response()
+
+    def create_session(self, **kwargs: Any) -> AgentSession:
+        return AgentSession(session_id=kwargs.get("session_id"))
+
+    def get_session(self, service_session_id: str | ServiceSessionId, *, session_id: str | None = None) -> AgentSession:
+        return AgentSession(session_id=session_id, service_session_id=service_session_id)
+
+
+# Fixtures
+
+
+@pytest.fixture
+def streaming_chat_client_stub() -> type[SupportsChatGetResponse]:
+    """Return the StreamingChatClientStub class for creating test instances."""
+    return StreamingChatClientStub  # type: ignore[return-value]
+
+
+@pytest.fixture
+def stream_from_updates_fixture() -> Callable[[list[ChatResponseUpdate]], StreamFn]:
+    """Return the stream_from_updates helper function."""
+    return stream_from_updates
+
+
+@pytest.fixture
+def stub_agent() -> type[SupportsAgentRun]:
+    """Return the StubAgent class for creating test instances."""
+    return StubAgent  # type: ignore[return-value]
+
+
+# ── Fixtures for golden / integration tests ──
+
+
+@pytest.fixture
+def collect_events() -> Callable[..., Any]:
+    """Return an async helper that collects all events from an async generator."""
+
+    async def _collect(async_gen: AsyncIterable[Any]) -> list[Any]:
+        return [event async for event in async_gen]
+
+    return _collect
+
+
+@pytest.fixture
+def make_agent_wrapper() -> Callable[..., Any]:
+    """Factory that builds an AgentFrameworkAgent from a stream function.
+
+    Usage::
+
+        agent = make_agent_wrapper(
+            stream_fn=stream_from_updates(updates),
+            state_schema=...,
+        )
+        events = [e async for e in agent.run(payload)]
+    """
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    def _factory(
+        stream_fn: StreamFn,
+        *,
+        state_schema: Any | None = None,
+        predict_state_config: dict[str, dict[str, str]] | None = None,
+        require_confirmation: bool = True,
+    ) -> Any:
+        client = StreamingChatClientStub(stream_fn)
+        stub = StubAgent(client=client)
+        return AgentFrameworkAgent(
+            agent=stub,
+            state_schema=state_schema,
+            predict_state_config=predict_state_config,
+            require_confirmation=require_confirmation,
+        )
+
+    return _factory
+
+
+@pytest.fixture
+def make_app() -> Callable[..., Any]:
+    """Factory that builds a FastAPI app with an AG-UI endpoint.
+
+    Usage::
+
+        app = make_app(agent_or_wrapper, path="/test")
+    """
+    from fastapi import FastAPI
+
+    from agent_framework_ag_ui import add_agent_framework_fastapi_endpoint
+
+    def _factory(
+        agent: Any,
+        *,
+        path: str = "/",
+        state_schema: Any | None = None,
+        predict_state_config: dict[str, dict[str, str]] | None = None,
+        default_state: dict[str, Any] | None = None,
+    ) -> FastAPI:
+        app = FastAPI()
+        add_agent_framework_fastapi_endpoint(
+            app,
+            agent,
+            path=path,
+            state_schema=state_schema,
+            predict_state_config=predict_state_config,
+            default_state=default_state,
+        )
+        return app
+
+    return _factory

python/packages/ag-ui/tests/ag_ui/event_stream.py

@@ -0,0 +1,210 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""EventStream assertion helper for AG-UI regression tests."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+class EventStream:
+    """Wraps a list of AG-UI events with structured assertion methods.
+
+    Usage:
+        events = [event async for event in agent.run(payload)]
+        stream = EventStream(events)
+        stream.assert_bookends()
+        stream.assert_text_messages_balanced()
+    """
+
+    def __init__(self, events: list[Any]) -> None:
+        self.events = events
+
+    def __len__(self) -> int:
+        return len(self.events)
+
+    def __iter__(self):
+        return iter(self.events)
+
+    def types(self) -> list[str]:
+        """Return ordered list of event type strings."""
+        return [self._type_str(e) for e in self.events]
+
+    def get(self, event_type: str) -> list[Any]:
+        """Filter events matching the given type string."""
+        return [e for e in self.events if self._type_str(e) == event_type]
+
+    def first(self, event_type: str) -> Any:
+        """Return the first event matching the given type, or raise."""
+        matches = self.get(event_type)
+        if not matches:
+            raise ValueError(f"No event of type {event_type!r} found. Available: {self.types()}")
+        return matches[0]
+
+    def last(self, event_type: str) -> Any:
+        """Return the last event matching the given type, or raise."""
+        matches = self.get(event_type)
+        if not matches:
+            raise ValueError(f"No event of type {event_type!r} found. Available: {self.types()}")
+        return matches[-1]
+
+    def snapshot(self) -> dict[str, Any]:
+        """Return the latest StateSnapshotEvent snapshot dict."""
+        return self.last("STATE_SNAPSHOT").snapshot
+
+    def messages_snapshot(self) -> list[Any]:
+        """Return the latest MessagesSnapshotEvent messages list."""
+        return self.last("MESSAGES_SNAPSHOT").messages
+
+    def run_finished_interrupts(self, event: Any | None = None) -> list[dict[str, Any]]:
+        """Return canonical interrupts from a RUN_FINISHED event."""
+        target = event or self.last("RUN_FINISHED")
+        dumped = self._event_dump(target)
+        assert "interrupt" not in dumped
+        outcome = dumped.get("outcome")
+        assert isinstance(outcome, dict), f"Expected RUN_FINISHED.outcome, got {dumped}"
+        assert outcome.get("type") == "interrupt"
+        interrupts = outcome.get("interrupts")
+        assert isinstance(interrupts, list), f"Expected outcome.interrupts, got {outcome}"
+        return interrupts
+
+    @staticmethod
+    def interrupt_metadata_value(interrupt: dict[str, Any]) -> dict[str, Any]:
+        """Return Agent Framework interruption details from canonical interrupt metadata."""
+        metadata = interrupt.get("metadata")
+        assert isinstance(metadata, dict)
+        agent_framework_metadata = metadata.get("agent_framework")
+        assert isinstance(agent_framework_metadata, dict)
+        value = agent_framework_metadata.get("value")
+        assert isinstance(value, dict)
+        return value
+
+    # ── Structural assertions ──
+
+    def assert_bookends(self) -> None:
+        """Assert first event is RUN_STARTED and last is RUN_FINISHED."""
+        types = self.types()
+        assert types, "Event stream is empty"
+        assert types[0] == "RUN_STARTED", f"Expected RUN_STARTED first, got {types[0]}"
+        assert types[-1] == "RUN_FINISHED", f"Expected RUN_FINISHED last, got {types[-1]}"
+
+    def assert_has_run_lifecycle(self) -> None:
+        """Assert RUN_STARTED is first and RUN_FINISHED exists (may not be last).
+
+        Use this instead of assert_bookends() for workflow resume streams where
+        _drain_open_message() can emit TEXT_MESSAGE_END after RUN_FINISHED.
+        """
+        types = self.types()
+        assert types, "Event stream is empty"
+        assert types[0] == "RUN_STARTED", f"Expected RUN_STARTED first, got {types[0]}"
+        assert "RUN_FINISHED" in types, f"Expected RUN_FINISHED in stream. Types: {types}"
+
+    def assert_strict_types(self, expected: list[str]) -> None:
+        """Assert exact type sequence match."""
+        actual = self.types()
+        assert actual == expected, f"Event type mismatch.\nExpected: {expected}\nActual:   {actual}"
+
+    def assert_ordered_types(self, expected: list[str]) -> None:
+        """Assert expected types appear as a subsequence (in order, not necessarily contiguous)."""
+        actual = self.types()
+        actual_idx = 0
+        for expected_type in expected:
+            found = False
+            while actual_idx < len(actual):
+                if actual[actual_idx] == expected_type:
+                    actual_idx += 1
+                    found = True
+                    break
+                actual_idx += 1
+            if not found:
+                raise AssertionError(
+                    f"Expected subsequence type {expected_type!r} not found after index {actual_idx}.\n"
+                    f"Expected subsequence: {expected}\n"
+                    f"Actual types: {actual}"
+                )
+
+    def assert_text_messages_balanced(self) -> None:
+        """Assert every TEXT_MESSAGE_START has a matching TEXT_MESSAGE_END with the same message_id."""
+        starts: dict[str, int] = {}
+        ends: set[str] = set()
+        for i, event in enumerate(self.events):
+            t = self._type_str(event)
+            if t == "TEXT_MESSAGE_START":
+                mid = event.message_id
+                assert mid not in starts, f"Duplicate TEXT_MESSAGE_START for message_id={mid}"
+                starts[mid] = i
+            elif t == "TEXT_MESSAGE_END":
+                mid = event.message_id
+                assert mid in starts, f"TEXT_MESSAGE_END for unknown message_id={mid}"
+                assert mid not in ends, f"Duplicate TEXT_MESSAGE_END for message_id={mid}"
+                ends.add(mid)
+
+        unclosed = set(starts.keys()) - ends
+        assert not unclosed, f"Unclosed text messages: {unclosed}"
+
+    def assert_tool_calls_balanced(self) -> None:
+        """Assert every TOOL_CALL_START has a matching TOOL_CALL_END with the same tool_call_id."""
+        starts: dict[str, int] = {}
+        ends: set[str] = set()
+        for i, event in enumerate(self.events):
+            t = self._type_str(event)
+            if t == "TOOL_CALL_START":
+                tid = event.tool_call_id
+                assert tid not in starts, f"Duplicate TOOL_CALL_START for tool_call_id={tid}"
+                starts[tid] = i
+            elif t == "TOOL_CALL_END":
+                tid = event.tool_call_id
+                assert tid in starts, f"TOOL_CALL_END for unknown tool_call_id={tid}"
+                assert tid not in ends, f"Duplicate TOOL_CALL_END for tool_call_id={tid}"
+                ends.add(tid)
+
+        unclosed = set(starts.keys()) - ends
+        assert not unclosed, f"Unclosed tool calls: {unclosed}"
+
+    def assert_no_run_error(self) -> None:
+        """Assert no RUN_ERROR events exist."""
+        errors = self.get("RUN_ERROR")
+        if errors:
+            messages = [getattr(e, "message", str(e)) for e in errors]
+            raise AssertionError(f"Found {len(errors)} RUN_ERROR event(s): {messages}")
+
+    def assert_has_type(self, event_type: str) -> None:
+        """Assert at least one event of the given type exists."""
+        assert event_type in self.types(), f"Expected {event_type!r} in stream. Available: {self.types()}"
+
+    def assert_message_ids_consistent(self) -> None:
+        """Assert TEXT_MESSAGE_CONTENT events reference valid, open message_ids."""
+        open_messages: set[str] = set()
+        for event in self.events:
+            t = self._type_str(event)
+            if t == "TEXT_MESSAGE_START":
+                open_messages.add(event.message_id)
+            elif t == "TEXT_MESSAGE_END":
+                open_messages.discard(event.message_id)
+            elif t == "TEXT_MESSAGE_CONTENT":
+                mid = event.message_id
+                assert mid in open_messages, f"TEXT_MESSAGE_CONTENT references message_id={mid} which is not open"
+
+    # ── Internal ──
+
+    @staticmethod
+    def _type_str(event: Any) -> str:
+        """Extract event type as a plain string."""
+        t = getattr(event, "type", None)
+        if t is None:
+            return type(event).__name__
+        if isinstance(t, str):
+            return t
+        return getattr(t, "value", str(t))
+
+    @staticmethod
+    def _event_dump(event: Any) -> dict[str, Any]:
+        """Serialize an event object or return a raw event dict."""
+        if isinstance(event, dict):
+            return event
+        if hasattr(event, "model_dump"):
+            return event.model_dump(by_alias=True, exclude_none=True)
+        raw = getattr(event, "raw", None)
+        if isinstance(raw, dict):
+            return raw
+        raise TypeError(f"Unsupported event type: {type(event).__name__}")

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

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

python/packages/ag-ui/tests/ag_ui/golden/conftest.py

@@ -0,0 +1,13 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Conftest for golden tests — ensures parent test dir is importable."""
+
+import sys
+from pathlib import Path
+
+
+def pytest_configure() -> None:
+    """Ensure parent test directory is on sys.path for helper module imports."""
+    parent_test_dir = str(Path(__file__).resolve().parent.parent)
+    if parent_test_dir not in sys.path:
+        sys.path.insert(0, parent_test_dir)

python/packages/ag-ui/tests/ag_ui/golden/test_scenario_agentic_chat.py

@@ -0,0 +1,140 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Golden event-stream tests for the basic agentic chat scenario."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from agent_framework import AgentResponseUpdate, Content
+from conftest import StubAgent  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+from event_stream import EventStream  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+
+from agent_framework_ag_ui import AgentFrameworkAgent
+
+
+def _build_agent(updates: list[AgentResponseUpdate], **kwargs: Any) -> AgentFrameworkAgent:
+    stub = StubAgent(updates=updates)
+    return AgentFrameworkAgent(agent=stub, **kwargs)
+
+
+async def _run(agent: AgentFrameworkAgent, payload: dict[str, Any]) -> EventStream:
+    return EventStream([event async for event in agent.run(payload)])
+
+
+BASIC_PAYLOAD: dict[str, Any] = {
+    "thread_id": "thread-chat",
+    "run_id": "run-chat",
+    "messages": [{"role": "user", "content": "Hello"}],
+}
+
+
+def _text_update(text: str) -> AgentResponseUpdate:
+    return AgentResponseUpdate(contents=[Content.from_text(text=text)], role="assistant")
+
+
+def _snapshot_role(msg: Any) -> str:
+    """Extract role string from a snapshot message (Pydantic model or dict)."""
+    role = getattr(msg, "role", None) or (msg.get("role") if isinstance(msg, dict) else None)
+    if role is None:
+        return ""
+    return str(getattr(role, "value", role))
+
+
+def _snapshot_content(msg: Any) -> str:
+    """Extract content string from a snapshot message."""
+    content = getattr(msg, "content", None) or (msg.get("content") if isinstance(msg, dict) else "")
+    return str(content) if content else ""
+
+
+# ── Golden stream tests ──
+
+
+async def test_basic_chat_golden_event_sequence() -> None:
+    """Assert the exact event type sequence for a single text response."""
+    agent = _build_agent([_text_update("Hi there!")])
+    stream = await _run(agent, BASIC_PAYLOAD)
+
+    stream.assert_strict_types(
+        [
+            "RUN_STARTED",
+            "TEXT_MESSAGE_START",
+            "TEXT_MESSAGE_CONTENT",
+            "TEXT_MESSAGE_END",
+            "MESSAGES_SNAPSHOT",
+            "RUN_FINISHED",
+        ]
+    )
+
+
+async def test_basic_chat_bookends() -> None:
+    """RUN_STARTED is first, RUN_FINISHED is last."""
+    agent = _build_agent([_text_update("reply")])
+    stream = await _run(agent, BASIC_PAYLOAD)
+    stream.assert_bookends()
+
+
+async def test_basic_chat_text_messages_balanced() -> None:
+    """Every TEXT_MESSAGE_START has a matching TEXT_MESSAGE_END."""
+    agent = _build_agent([_text_update("reply")])
+    stream = await _run(agent, BASIC_PAYLOAD)
+    stream.assert_text_messages_balanced()
+
+
+async def test_basic_chat_no_errors() -> None:
+    """No RUN_ERROR events in a normal flow."""
+    agent = _build_agent([_text_update("reply")])
+    stream = await _run(agent, BASIC_PAYLOAD)
+    stream.assert_no_run_error()
+
+
+async def test_basic_chat_message_id_consistency() -> None:
+    """All text events reference the same message_id."""
+    agent = _build_agent([_text_update("reply")])
+    stream = await _run(agent, BASIC_PAYLOAD)
+
+    start = stream.first("TEXT_MESSAGE_START")
+    content = stream.first("TEXT_MESSAGE_CONTENT")
+    end = stream.first("TEXT_MESSAGE_END")
+    assert start.message_id == content.message_id == end.message_id
+
+
+async def test_multi_chunk_text_golden_sequence() -> None:
+    """Streaming multiple chunks produces START + multiple CONTENT + END."""
+    agent = _build_agent([_text_update("Hello "), _text_update("world!")])
+    stream = await _run(agent, BASIC_PAYLOAD)
+
+    stream.assert_strict_types(
+        [
+            "RUN_STARTED",
+            "TEXT_MESSAGE_START",
+            "TEXT_MESSAGE_CONTENT",
+            "TEXT_MESSAGE_CONTENT",
+            "TEXT_MESSAGE_END",
+            "MESSAGES_SNAPSHOT",
+            "RUN_FINISHED",
+        ]
+    )
+    stream.assert_text_messages_balanced()
+    stream.assert_message_ids_consistent()
+
+
+async def test_messages_snapshot_contains_assistant_reply() -> None:
+    """MessagesSnapshotEvent includes the assistant's accumulated text."""
+    agent = _build_agent([_text_update("Hello there")])
+    stream = await _run(agent, BASIC_PAYLOAD)
+
+    snapshot = stream.messages_snapshot()
+    assistant_msgs = [m for m in snapshot if _snapshot_role(m) == "assistant"]
+    assert assistant_msgs, "No assistant message in snapshot"
+    assert any("Hello there" in _snapshot_content(m) for m in assistant_msgs)
+
+
+async def test_empty_messages_produces_start_and_finish() -> None:
+    """Empty message list still produces RUN_STARTED and RUN_FINISHED."""
+    agent = _build_agent([_text_update("reply")])
+    payload = {"thread_id": "t1", "run_id": "r1", "messages": []}
+    stream = await _run(agent, payload)
+
+    stream.assert_bookends()
+    assert "TEXT_MESSAGE_START" not in stream.types()

python/packages/ag-ui/tests/ag_ui/golden/test_scenario_backend_tools.py

@@ -0,0 +1,236 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Golden event-stream tests for the backend (server-side) tools scenario."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from agent_framework import AgentResponseUpdate, Content
+from conftest import StubAgent  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+from event_stream import EventStream  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+
+from agent_framework_ag_ui import AgentFrameworkAgent
+
+
+def _build_agent(updates: list[AgentResponseUpdate], **kwargs: Any) -> AgentFrameworkAgent:
+    stub = StubAgent(updates=updates)
+    return AgentFrameworkAgent(agent=stub, **kwargs)
+
+
+async def _run(agent: AgentFrameworkAgent, payload: dict[str, Any]) -> EventStream:
+    return EventStream([event async for event in agent.run(payload)])
+
+
+PAYLOAD: dict[str, Any] = {
+    "thread_id": "thread-tools",
+    "run_id": "run-tools",
+    "messages": [{"role": "user", "content": "What's the weather?"}],
+}
+
+
+# ── Golden stream tests ──
+
+
+async def test_tool_call_lifecycle_golden_sequence() -> None:
+    """Assert the full event sequence for a tool call → result → text response."""
+    updates = [
+        # LLM calls the tool
+        AgentResponseUpdate(
+            contents=[Content.from_function_call(name="get_weather", call_id="call-1", arguments='{"city": "SF"}')],
+            role="assistant",
+        ),
+        # Tool result comes back
+        AgentResponseUpdate(
+            contents=[Content.from_function_result(call_id="call-1", result="72°F and sunny")],
+            role="assistant",
+        ),
+        # LLM responds with text
+        AgentResponseUpdate(
+            contents=[Content.from_text(text="It's 72°F and sunny in SF!")],
+            role="assistant",
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    stream.assert_ordered_types(
+        [
+            "RUN_STARTED",
+            "TEXT_MESSAGE_START",  # Synthetic start for tool-only message
+            "TOOL_CALL_START",
+            "TOOL_CALL_ARGS",
+            "TOOL_CALL_END",
+            "TOOL_CALL_RESULT",
+            "TEXT_MESSAGE_END",  # End of synthetic message
+            "TEXT_MESSAGE_START",  # New message for text response
+            "TEXT_MESSAGE_CONTENT",
+            "TEXT_MESSAGE_END",
+            "MESSAGES_SNAPSHOT",
+            "RUN_FINISHED",
+        ]
+    )
+
+
+async def test_tool_calls_balanced() -> None:
+    """Every TOOL_CALL_START has a matching TOOL_CALL_END."""
+    updates = [
+        AgentResponseUpdate(
+            contents=[Content.from_function_call(name="get_weather", call_id="call-1", arguments='{"city": "SF"}')],
+            role="assistant",
+        ),
+        AgentResponseUpdate(
+            contents=[Content.from_function_result(call_id="call-1", result="72°F")],
+            role="assistant",
+        ),
+        AgentResponseUpdate(
+            contents=[Content.from_text(text="It's 72°F!")],
+            role="assistant",
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    stream.assert_tool_calls_balanced()
+
+
+async def test_text_messages_balanced_with_tools() -> None:
+    """Text messages are properly balanced even around tool calls."""
+    updates = [
+        AgentResponseUpdate(
+            contents=[Content.from_function_call(name="get_weather", call_id="call-1", arguments='{"city": "SF"}')],
+            role="assistant",
+        ),
+        AgentResponseUpdate(
+            contents=[Content.from_function_result(call_id="call-1", result="72°F")],
+            role="assistant",
+        ),
+        AgentResponseUpdate(
+            contents=[Content.from_text(text="It's 72°F!")],
+            role="assistant",
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    stream.assert_text_messages_balanced()
+
+
+async def test_tool_call_id_matches_result() -> None:
+    """TOOL_CALL_START and TOOL_CALL_RESULT reference the same tool_call_id."""
+    updates = [
+        AgentResponseUpdate(
+            contents=[Content.from_function_call(name="get_weather", call_id="call-1", arguments="{}")],
+            role="assistant",
+        ),
+        AgentResponseUpdate(
+            contents=[Content.from_function_result(call_id="call-1", result="72°F")],
+            role="assistant",
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    start = stream.first("TOOL_CALL_START")
+    result = stream.first("TOOL_CALL_RESULT")
+    assert start.tool_call_id == result.tool_call_id == "call-1"
+
+
+async def test_tool_result_content_preserved() -> None:
+    """TOOL_CALL_RESULT event carries the tool's result content."""
+    updates = [
+        AgentResponseUpdate(
+            contents=[Content.from_function_call(name="get_weather", call_id="call-1", arguments="{}")],
+            role="assistant",
+        ),
+        AgentResponseUpdate(
+            contents=[Content.from_function_result(call_id="call-1", result="72°F and sunny")],
+            role="assistant",
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    result = stream.first("TOOL_CALL_RESULT")
+    assert result.content == "72°F and sunny"
+
+
+async def test_no_run_error_on_tool_flow() -> None:
+    """Tool call flow doesn't produce RUN_ERROR."""
+    updates = [
+        AgentResponseUpdate(
+            contents=[Content.from_function_call(name="get_weather", call_id="call-1", arguments="{}")],
+            role="assistant",
+        ),
+        AgentResponseUpdate(
+            contents=[Content.from_function_result(call_id="call-1", result="72°F")],
+            role="assistant",
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    stream.assert_no_run_error()
+    stream.assert_bookends()
+
+
+async def test_multiple_sequential_tool_calls() -> None:
+    """Multiple sequential tool calls each produce balanced START/END pairs."""
+    updates = [
+        AgentResponseUpdate(
+            contents=[Content.from_function_call(name="tool_a", call_id="call-a", arguments="{}")],
+            role="assistant",
+        ),
+        AgentResponseUpdate(
+            contents=[Content.from_function_result(call_id="call-a", result="result-a")],
+            role="assistant",
+        ),
+        AgentResponseUpdate(
+            contents=[Content.from_function_call(name="tool_b", call_id="call-b", arguments="{}")],
+            role="assistant",
+        ),
+        AgentResponseUpdate(
+            contents=[Content.from_function_result(call_id="call-b", result="result-b")],
+            role="assistant",
+        ),
+        AgentResponseUpdate(
+            contents=[Content.from_text(text="Done!")],
+            role="assistant",
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    stream.assert_tool_calls_balanced()
+    stream.assert_text_messages_balanced()
+    stream.assert_bookends()
+
+    # Both tool calls should appear
+    starts = stream.get("TOOL_CALL_START")
+    assert len(starts) == 2
+    assert {s.tool_call_name for s in starts} == {"tool_a", "tool_b"}
+
+
+async def test_messages_snapshot_includes_tool_calls() -> None:
+    """MessagesSnapshotEvent includes tool call and result messages."""
+    updates = [
+        AgentResponseUpdate(
+            contents=[Content.from_function_call(name="get_weather", call_id="call-1", arguments='{"city":"SF"}')],
+            role="assistant",
+        ),
+        AgentResponseUpdate(
+            contents=[Content.from_function_result(call_id="call-1", result="72°F")],
+            role="assistant",
+        ),
+        AgentResponseUpdate(
+            contents=[Content.from_text(text="It's warm!")],
+            role="assistant",
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    stream.assert_has_type("MESSAGES_SNAPSHOT")
+    snapshot = stream.messages_snapshot()
+    # Should have: user message, assistant with tool_calls, tool result, assistant text
+    assert len(snapshot) >= 3

python/packages/ag-ui/tests/ag_ui/golden/test_scenario_deterministic_state.py

@@ -0,0 +1,364 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Golden event-stream tests for the deterministic tool-driven state scenario.
+
+Covers issue https://github.com/microsoft/agent-framework/issues/3167 — a tool
+returning :func:`agent_framework_ag_ui.state_update` must push a deterministic
+``StateSnapshotEvent`` derived from its actual return value, orthogonal to the
+optimistic ``predict_state_config`` path. These golden tests pin the user-visible
+event stream so additive changes cannot silently regress it.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from agent_framework import AgentResponseUpdate, Content
+from conftest import StubAgent  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+from event_stream import EventStream  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+
+from agent_framework_ag_ui import AgentFrameworkAgent, state_update
+
+STATE_SCHEMA = {
+    "weather": {"type": "object", "description": "Last fetched weather"},
+}
+
+
+def _build_agent(updates: list[AgentResponseUpdate], **kwargs: Any) -> AgentFrameworkAgent:
+    stub = StubAgent(updates=updates)
+    kwargs.setdefault("state_schema", STATE_SCHEMA)
+    return AgentFrameworkAgent(agent=stub, **kwargs)
+
+
+async def _run(agent: AgentFrameworkAgent, payload: dict[str, Any]) -> EventStream:
+    return EventStream([event async for event in agent.run(payload)])
+
+
+PAYLOAD: dict[str, Any] = {
+    "thread_id": "thread-det-state",
+    "run_id": "run-det-state",
+    "messages": [{"role": "user", "content": "What's the weather in SF?"}],
+    "state": {"weather": {}},
+}
+
+
+def _tool_call(call_id: str, name: str, arguments: str) -> AgentResponseUpdate:
+    return AgentResponseUpdate(
+        contents=[Content.from_function_call(name=name, call_id=call_id, arguments=arguments)],
+        role="assistant",
+    )
+
+
+def _tool_result_with_state(call_id: str, text: str, state: dict[str, Any]) -> AgentResponseUpdate:
+    """Build a function_result update whose inner item carries a state marker.
+
+    This mirrors what the core framework produces when a real ``@tool`` returns
+    :func:`state_update`: ``parse_result`` keeps the ``Content`` as-is, and
+    ``Content.from_function_result`` preserves its ``additional_properties``
+    inside ``items``.
+    """
+    return AgentResponseUpdate(
+        contents=[
+            Content.from_function_result(
+                call_id=call_id,
+                result=[state_update(text=text, state=state)],
+            )
+        ],
+        role="assistant",
+    )
+
+
+def _tool_result_with_display(call_id: str, text: str, tool_result: Any, **kwargs: Any) -> AgentResponseUpdate:
+    """Build a function_result update carrying an optional UI display marker."""
+    return AgentResponseUpdate(
+        contents=[
+            Content.from_function_result(
+                call_id=call_id,
+                result=[state_update(text=text, tool_result=tool_result, **kwargs)],
+            )
+        ],
+        role="assistant",
+    )
+
+
+# ── Golden stream tests ──
+
+
+async def test_deterministic_state_emits_snapshot_after_tool_result() -> None:
+    """The happy path: STATE_SNAPSHOT follows TOOL_CALL_RESULT in order."""
+    updates = [
+        _tool_call("call-1", "get_weather", '{"city": "SF"}'),
+        _tool_result_with_state(
+            "call-1",
+            text="Weather in SF: 14°C foggy",
+            state={"weather": {"city": "SF", "temp": 14, "conditions": "foggy"}},
+        ),
+        AgentResponseUpdate(
+            contents=[Content.from_text(text="It's 14°C and foggy in SF.")],
+            role="assistant",
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    stream.assert_bookends()
+    stream.assert_no_run_error()
+    stream.assert_tool_calls_balanced()
+    stream.assert_text_messages_balanced()
+
+    # Ordered subsequence: the deterministic STATE_SNAPSHOT must follow the
+    # TOOL_CALL_RESULT. This is the central contract for #3167.
+    stream.assert_ordered_types(
+        [
+            "RUN_STARTED",
+            "TOOL_CALL_START",
+            "TOOL_CALL_ARGS",
+            "TOOL_CALL_END",
+            "TOOL_CALL_RESULT",
+            "STATE_SNAPSHOT",
+            "RUN_FINISHED",
+        ]
+    )
+
+    # The final STATE_SNAPSHOT must carry the tool-driven state.
+    snapshot = stream.snapshot()
+    assert snapshot["weather"] == {"city": "SF", "temp": 14, "conditions": "foggy"}
+
+
+async def test_deterministic_state_does_not_fire_for_plain_tool_result() -> None:
+    """Regression guard: tools returning plain strings must NOT emit a new STATE_SNAPSHOT.
+
+    The initial STATE_SNAPSHOT fires once from the schema + initial payload
+    state. A plain (non-state_update) tool result must not add another one.
+    """
+    updates = [
+        _tool_call("call-1", "get_weather", '{"city": "SF"}'),
+        AgentResponseUpdate(
+            contents=[Content.from_function_result(call_id="call-1", result="14°C foggy")],
+            role="assistant",
+        ),
+        AgentResponseUpdate(
+            contents=[Content.from_text(text="It's 14°C and foggy.")],
+            role="assistant",
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    stream.assert_bookends()
+    stream.assert_no_run_error()
+
+    snapshots = stream.get("STATE_SNAPSHOT")
+    # Only the initial snapshot (from state_schema + payload state) should exist.
+    # No deterministic snapshot should have been added by the plain tool result.
+    assert len(snapshots) == 1, (
+        f"Expected exactly 1 STATE_SNAPSHOT (initial only) for plain tool result; "
+        f"got {len(snapshots)}. Snapshots: {[s.snapshot for s in snapshots]}"
+    )
+
+
+async def test_deterministic_state_merges_into_initial_state() -> None:
+    """The tool-driven snapshot must merge into, not replace, pre-existing state keys."""
+    payload = dict(PAYLOAD)
+    payload["state"] = {"weather": {}, "user_preferences": {"unit": "C"}}
+
+    updates = [
+        _tool_call("call-1", "get_weather", '{"city": "SF"}'),
+        _tool_result_with_state(
+            "call-1",
+            text="Weather: 14°C",
+            state={"weather": {"city": "SF", "temp": 14}},
+        ),
+    ]
+    agent = _build_agent(updates, state_schema={**STATE_SCHEMA, "user_preferences": {"type": "object"}})
+    stream = await _run(agent, payload)
+
+    stream.assert_bookends()
+    stream.assert_no_run_error()
+
+    final_snapshot = stream.snapshot()
+    assert final_snapshot["weather"] == {"city": "SF", "temp": 14}
+    assert final_snapshot["user_preferences"] == {"unit": "C"}, (
+        "Pre-existing state keys must survive the deterministic merge"
+    )
+
+
+async def test_deterministic_state_llm_visible_text_is_clean() -> None:
+    """The LLM-visible TOOL_CALL_RESULT content must not leak the state marker key."""
+    updates = [
+        _tool_call("call-1", "get_weather", '{"city": "SF"}'),
+        _tool_result_with_state(
+            "call-1",
+            text="Weather in SF: 14°C foggy",
+            state={"weather": {"city": "SF", "temp": 14}},
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    result = stream.first("TOOL_CALL_RESULT")
+    assert result.content == "Weather in SF: 14°C foggy"
+    # The marker key must never appear in the content sent back to the LLM.
+    assert "__ag_ui_tool_result_state__" not in result.content
+    assert "weather" not in result.content  # not as a raw state dump
+
+
+async def test_deterministic_state_multiple_tools_merge_in_order() -> None:
+    """Two state-updating tools in one run merge in order; later wins on key collisions."""
+    updates = [
+        _tool_call("call-a", "get_weather", '{"city": "SF"}'),
+        _tool_result_with_state(
+            "call-a",
+            text="First result",
+            state={"weather": {"city": "SF", "temp": 14}, "source": "primary"},
+        ),
+        _tool_call("call-b", "get_weather_refined", '{"city": "SF"}'),
+        _tool_result_with_state(
+            "call-b",
+            text="Refined result",
+            state={"source": "refined"},
+        ),
+        AgentResponseUpdate(
+            contents=[Content.from_text(text="Here you go.")],
+            role="assistant",
+        ),
+    ]
+    agent = _build_agent(
+        updates,
+        state_schema={**STATE_SCHEMA, "source": {"type": "string"}},
+    )
+    stream = await _run(agent, PAYLOAD)
+
+    stream.assert_bookends()
+    stream.assert_tool_calls_balanced()
+    stream.assert_no_run_error()
+
+    # Two tool-driven snapshots emitted (one per tool) plus the initial snapshot.
+    snapshots = stream.get("STATE_SNAPSHOT")
+    assert len(snapshots) >= 2, f"Expected at least 2 STATE_SNAPSHOTs; got {len(snapshots)}"
+
+    final = stream.snapshot()
+    assert final["weather"] == {"city": "SF", "temp": 14}
+    # Later tool must override earlier tool on the shared key.
+    assert final["source"] == "refined"
+
+
+async def test_deterministic_state_coexists_with_predict_state_config() -> None:
+    """Predictive state and deterministic state must coexist without clobbering each other."""
+    predict_config = {
+        "draft": {
+            "tool": "write_draft",
+            "tool_argument": "body",
+        }
+    }
+    updates = [
+        # Predictive tool: its argument "body" populates state.draft optimistically.
+        _tool_call("call-1", "write_draft", '{"body": "Hello world"}'),
+        # Then a deterministic tool result landing a different key.
+        _tool_result_with_state(
+            "call-1",
+            text="Draft saved",
+            state={"weather": {"city": "SF", "temp": 14}},
+        ),
+    ]
+    agent = _build_agent(
+        updates,
+        state_schema={**STATE_SCHEMA, "draft": {"type": "string"}},
+        predict_state_config=predict_config,
+        require_confirmation=False,
+    )
+    payload = dict(PAYLOAD)
+    payload["state"] = {"weather": {}, "draft": ""}
+    stream = await _run(agent, payload)
+
+    stream.assert_bookends()
+    stream.assert_no_run_error()
+    stream.assert_tool_calls_balanced()
+
+    # The final observed state must contain both the deterministic and predictive contributions.
+    final = stream.snapshot()
+    assert final["weather"] == {"city": "SF", "temp": 14}, f"Deterministic state missing from final snapshot: {final}"
+
+
+async def test_tool_result_display_payload_reaches_ui_event_only() -> None:
+    """Rich display payload overrides TOOL_CALL_RESULT without leaking marker keys."""
+    updates = [
+        _tool_call("call-1", "get_weather", '{"city": "SF"}'),
+        _tool_result_with_display(
+            "call-1",
+            text="Weather in SF: 14°C foggy",
+            tool_result={"city": "SF", "temp": 14, "conditions": "foggy"},
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    stream.assert_bookends()
+    stream.assert_no_run_error()
+    stream.assert_tool_calls_balanced()
+
+    result = stream.first("TOOL_CALL_RESULT")
+    assert result.content == '{"city": "SF", "temp": 14, "conditions": "foggy"}'
+    assert "__ag_ui_tool_result_display__" not in result.content
+    assert "__ag_ui_tool_result_state__" not in result.content
+
+
+async def test_tool_result_display_falls_back_to_text_when_unset() -> None:
+    """Without a display marker, the UI event keeps the existing text content."""
+    updates = [
+        _tool_call("call-1", "get_weather", '{"city": "SF"}'),
+        _tool_result_with_state(
+            "call-1",
+            text="Weather in SF: 14°C foggy",
+            state={"weather": {"city": "SF", "temp": 14}},
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    stream.assert_bookends()
+    stream.assert_no_run_error()
+    stream.assert_tool_calls_balanced()
+
+    result = stream.first("TOOL_CALL_RESULT")
+    assert result.content == "Weather in SF: 14°C foggy"
+    assert "__ag_ui_tool_result_display__" not in result.content
+    assert "__ag_ui_tool_result_state__" not in result.content
+
+
+async def test_tool_result_display_coexists_with_state_snapshot() -> None:
+    """Display and durable state markers produce one deterministic state snapshot."""
+    updates = [
+        _tool_call("call-1", "get_weather", '{"city": "SF"}'),
+        _tool_result_with_display(
+            "call-1",
+            text="Weather in SF: 14°C foggy",
+            tool_result={"city": "SF", "temp": 14, "conditions": "foggy"},
+            state={"weather": {"city": "SF", "temp": 14, "conditions": "foggy"}},
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    stream.assert_bookends()
+    stream.assert_no_run_error()
+    stream.assert_tool_calls_balanced()
+    stream.assert_ordered_types(["TOOL_CALL_RESULT", "STATE_SNAPSHOT", "RUN_FINISHED"])
+
+    result = stream.first("TOOL_CALL_RESULT")
+    assert result.content == '{"city": "SF", "temp": 14, "conditions": "foggy"}'
+
+    result_idx = stream.events.index(result)
+    deterministic_snapshots = [
+        event
+        for event in stream.events[result_idx + 1 :]
+        if getattr(getattr(event, "type", None), "value", getattr(event, "type", None)) == "STATE_SNAPSHOT"
+    ]
+    assert len(deterministic_snapshots) == 1
+    assert deterministic_snapshots[0].snapshot["weather"] == {
+        "city": "SF",
+        "temp": 14,
+        "conditions": "foggy",
+    }
+    assert "__ag_ui_tool_result_display__" not in str(deterministic_snapshots[0].snapshot)
+    assert "__ag_ui_tool_result_state__" not in str(deterministic_snapshots[0].snapshot)

python/packages/ag-ui/tests/ag_ui/golden/test_scenario_generative_ui_agent.py

@@ -0,0 +1,90 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Golden event-stream tests for the generative UI (workflow-as-agent) scenario."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from agent_framework import WorkflowBuilder, WorkflowContext, executor
+from event_stream import EventStream  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+
+from agent_framework_ag_ui import AgentFrameworkWorkflow
+
+
+async def _run(wrapper: AgentFrameworkWorkflow, payload: dict[str, Any]) -> EventStream:
+    return EventStream([event async for event in wrapper.run(payload)])
+
+
+PAYLOAD: dict[str, Any] = {
+    "thread_id": "thread-gen-ui-agent",
+    "run_id": "run-gen-ui-agent",
+    "messages": [{"role": "user", "content": "Generate a UI"}],
+}
+
+
+# ── Golden stream tests ──
+
+
+async def test_workflow_agent_golden_sequence() -> None:
+    """Workflow-as-agent: emits step events and text content."""
+
+    @executor(id="generator")
+    async def generator(message: Any, ctx: WorkflowContext[Any, str]) -> None:
+        await ctx.yield_output("Here is your generated UI content!")
+
+    workflow = WorkflowBuilder(start_executor=generator).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+    stream = await _run(wrapper, PAYLOAD)
+
+    stream.assert_bookends()
+    stream.assert_no_run_error()
+    stream.assert_text_messages_balanced()
+
+    # Should have step events for the executor
+    stream.assert_has_type("STEP_STARTED")
+    stream.assert_has_type("STEP_FINISHED")
+
+    # Should have text message content
+    stream.assert_has_type("TEXT_MESSAGE_CONTENT")
+
+
+async def test_workflow_agent_step_names_match() -> None:
+    """Step started/finished events reference the executor name."""
+
+    @executor(id="my_executor")
+    async def my_executor(message: Any, ctx: WorkflowContext[Any, str]) -> None:
+        await ctx.yield_output("Done!")
+
+    workflow = WorkflowBuilder(start_executor=my_executor).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+    stream = await _run(wrapper, PAYLOAD)
+
+    started = [e for e in stream.get("STEP_STARTED") if getattr(e, "step_name", "") == "my_executor"]
+    finished = [e for e in stream.get("STEP_FINISHED") if getattr(e, "step_name", "") == "my_executor"]
+    assert started, "Expected STEP_STARTED for 'my_executor'"
+    assert finished, "Expected STEP_FINISHED for 'my_executor'"
+
+
+async def test_workflow_agent_ordered_events() -> None:
+    """Workflow events follow expected ordering: RUN_STARTED → STEP_STARTED → content → STEP_FINISHED → RUN_FINISHED."""
+
+    @executor(id="my_step")
+    async def my_step(message: Any, ctx: WorkflowContext[Any, str]) -> None:
+        await ctx.yield_output("Generated content")
+
+    workflow = WorkflowBuilder(start_executor=my_step).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+    stream = await _run(wrapper, PAYLOAD)
+
+    stream.assert_ordered_types(
+        [
+            "RUN_STARTED",
+            "STEP_STARTED",
+            "TEXT_MESSAGE_START",
+            "TEXT_MESSAGE_CONTENT",
+            "STEP_FINISHED",
+            "TEXT_MESSAGE_END",
+            "RUN_FINISHED",
+        ]
+    )

python/packages/ag-ui/tests/ag_ui/golden/test_scenario_generative_ui_tool.py

@@ -0,0 +1,135 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Golden event-stream tests for the client-side (declaration-only) tools scenario."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from agent_framework import AgentResponseUpdate, Content
+from conftest import StubAgent  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+from event_stream import EventStream  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+
+from agent_framework_ag_ui import AgentFrameworkAgent
+
+
+def _build_agent(updates: list[AgentResponseUpdate], **kwargs: Any) -> AgentFrameworkAgent:
+    stub = StubAgent(updates=updates)
+    return AgentFrameworkAgent(agent=stub, **kwargs)
+
+
+async def _run(agent: AgentFrameworkAgent, payload: dict[str, Any]) -> EventStream:
+    return EventStream([event async for event in agent.run(payload)])
+
+
+PAYLOAD: dict[str, Any] = {
+    "thread_id": "thread-gen-ui-tool",
+    "run_id": "run-gen-ui-tool",
+    "messages": [{"role": "user", "content": "Show me a chart"}],
+    "tools": [
+        {
+            "type": "function",
+            "function": {
+                "name": "render_chart",
+                "description": "Render a chart in the UI",
+                "parameters": {
+                    "type": "object",
+                    "properties": {"data": {"type": "array"}},
+                },
+            },
+        }
+    ],
+}
+
+
+# ── Golden stream tests ──
+
+
+async def test_declaration_only_tool_golden_sequence() -> None:
+    """Declaration-only tool: TOOL_CALL_START/ARGS emitted, TOOL_CALL_END at stream end."""
+    # The LLM calls a client-side tool (no server-side execution)
+    updates = [
+        AgentResponseUpdate(
+            contents=[
+                Content.from_function_call(
+                    name="render_chart",
+                    call_id="call-chart",
+                    arguments='{"data": [1, 2, 3]}',
+                )
+            ],
+            role="assistant",
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    stream.assert_bookends()
+    stream.assert_no_run_error()
+
+    # Tool call start and args should be present
+    stream.assert_has_type("TOOL_CALL_START")
+    stream.assert_has_type("TOOL_CALL_ARGS")
+
+    # TOOL_CALL_END should be emitted (via get_pending_without_end)
+    stream.assert_has_type("TOOL_CALL_END")
+    stream.assert_tool_calls_balanced()
+
+
+async def test_declaration_only_tool_no_tool_call_result() -> None:
+    """Declaration-only tools should NOT produce TOOL_CALL_RESULT events."""
+    updates = [
+        AgentResponseUpdate(
+            contents=[
+                Content.from_function_call(
+                    name="render_chart",
+                    call_id="call-chart",
+                    arguments='{"data": [1, 2, 3]}',
+                )
+            ],
+            role="assistant",
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    assert "TOOL_CALL_RESULT" not in stream.types(), "Declaration-only tools should not have TOOL_CALL_RESULT"
+
+
+async def test_declaration_only_tool_text_messages_balanced() -> None:
+    """Text messages remain balanced even with declaration-only tools."""
+    updates = [
+        AgentResponseUpdate(
+            contents=[
+                Content.from_function_call(
+                    name="render_chart",
+                    call_id="call-chart",
+                    arguments='{"data": [1, 2, 3]}',
+                )
+            ],
+            role="assistant",
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    stream.assert_text_messages_balanced()
+
+
+async def test_declaration_only_tool_messages_snapshot() -> None:
+    """MessagesSnapshotEvent includes the tool call for declaration-only tools."""
+    updates = [
+        AgentResponseUpdate(
+            contents=[
+                Content.from_function_call(
+                    name="render_chart",
+                    call_id="call-chart",
+                    arguments='{"data": [1, 2, 3]}',
+                )
+            ],
+            role="assistant",
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    stream.assert_has_type("MESSAGES_SNAPSHOT")

python/packages/ag-ui/tests/ag_ui/golden/test_scenario_hitl.py

@@ -0,0 +1,194 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Golden event-stream tests for the HITL (human-in-the-loop) approval scenario."""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from agent_framework import AgentResponseUpdate, Content
+from conftest import StubAgent  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+from event_stream import EventStream  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+
+from agent_framework_ag_ui import AgentFrameworkAgent
+
+PREDICT_CONFIG = {
+    "tasks": {
+        "tool": "generate_task_steps",
+        "tool_argument": "steps",
+    }
+}
+
+STATE_SCHEMA = {
+    "tasks": {"type": "array", "items": {"type": "object"}},
+}
+
+
+def _build_agent(updates: list[AgentResponseUpdate], **kwargs: Any) -> AgentFrameworkAgent:
+    stub = StubAgent(updates=updates)
+    return AgentFrameworkAgent(
+        agent=stub,
+        state_schema=STATE_SCHEMA,
+        predict_state_config=PREDICT_CONFIG,
+        require_confirmation=True,
+        **kwargs,
+    )
+
+
+async def _run(agent: AgentFrameworkAgent, payload: dict[str, Any]) -> EventStream:
+    return EventStream([event async for event in agent.run(payload)])
+
+
+STEPS = [
+    {"description": "Step 1: Plan", "status": "enabled"},
+    {"description": "Step 2: Execute", "status": "enabled"},
+]
+
+
+PAYLOAD: dict[str, Any] = {
+    "thread_id": "thread-hitl",
+    "run_id": "run-hitl",
+    "messages": [{"role": "user", "content": "Plan my tasks"}],
+    "state": {"tasks": []},
+}
+
+
+# ── Turn 1: Tool call → confirm_changes → interrupt ──
+
+
+async def test_hitl_turn1_golden_sequence() -> None:
+    """Turn 1 emits tool call, confirm_changes, and finishes with interrupt."""
+    updates = [
+        AgentResponseUpdate(
+            contents=[
+                Content.from_function_call(
+                    name="generate_task_steps",
+                    call_id="call-steps",
+                    arguments=json.dumps({"steps": STEPS}),
+                )
+            ],
+            role="assistant",
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    # Should have: tool call start/args/end for the primary tool,
+    # then TOOL_CALL_END, STATE_SNAPSHOT, confirm_changes cycle
+    stream.assert_bookends()
+    stream.assert_no_run_error()
+
+    # confirm_changes tool call should be present
+    tool_starts = stream.get("TOOL_CALL_START")
+    tool_names = [getattr(s, "tool_call_name", None) for s in tool_starts]
+    assert "generate_task_steps" in tool_names
+    assert "confirm_changes" in tool_names
+
+    # RUN_FINISHED should have interrupt metadata
+    interrupt = stream.run_finished_interrupts()
+    assert len(interrupt) > 0
+
+
+async def test_hitl_turn1_tool_calls_balanced() -> None:
+    """All tool calls in turn 1 (primary + confirm_changes) are balanced."""
+    updates = [
+        AgentResponseUpdate(
+            contents=[
+                Content.from_function_call(
+                    name="generate_task_steps",
+                    call_id="call-steps",
+                    arguments=json.dumps({"steps": STEPS}),
+                )
+            ],
+            role="assistant",
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    stream.assert_tool_calls_balanced()
+
+
+async def test_hitl_turn1_text_messages_balanced() -> None:
+    """Text messages are balanced even in the approval flow."""
+    updates = [
+        AgentResponseUpdate(
+            contents=[
+                Content.from_function_call(
+                    name="generate_task_steps",
+                    call_id="call-steps",
+                    arguments=json.dumps({"steps": STEPS}),
+                )
+            ],
+            role="assistant",
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    stream.assert_text_messages_balanced()
+
+
+# ── Turn 2: Resume with approval → confirmation message → no interrupt ──
+
+
+async def test_hitl_turn2_resume_with_approval() -> None:
+    """Resuming with confirm_changes result emits confirmation text and finishes cleanly."""
+    # Turn 2: user sends confirm_changes result as resume
+    # The agent wrapper sees a confirm_changes response and emits a confirmation message
+    confirm_result = json.dumps(
+        {
+            "accepted": True,
+            "steps": STEPS,
+        }
+    )
+
+    # Build payload with resume containing the approval
+    # For confirm_changes, the messages should include the tool result
+    payload: dict[str, Any] = {
+        "thread_id": "thread-hitl",
+        "run_id": "run-hitl-2",
+        "messages": [
+            {"role": "user", "content": "Plan my tasks"},
+            {
+                "role": "assistant",
+                "tool_calls": [
+                    {
+                        "id": "confirm-id-1",
+                        "type": "function",
+                        "function": {"name": "confirm_changes", "arguments": json.dumps({"steps": STEPS})},
+                    }
+                ],
+            },
+            {
+                "role": "tool",
+                "toolCallId": "confirm-id-1",
+                "content": confirm_result,
+            },
+        ],
+        "state": {"tasks": []},
+    }
+
+    # In turn 2, the agent sees the confirm_changes result and emits a confirmation text
+    updates = [
+        AgentResponseUpdate(
+            contents=[Content.from_text(text="Tasks confirmed!")],
+            role="assistant",
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, payload)
+
+    stream.assert_bookends()
+    stream.assert_text_messages_balanced()
+    stream.assert_no_run_error()
+
+    # Should have text message content (the confirmation message)
+    text_events = stream.get("TEXT_MESSAGE_CONTENT")
+    assert text_events, "Expected confirmation text message"
+
+    # RUN_FINISHED should NOT have interrupt (approval completed)
+    finished = stream.last("RUN_FINISHED")
+    dumped = finished.model_dump(by_alias=True, exclude_none=True)
+    assert "outcome" not in dumped, f"Expected no interrupt after approval, got {dumped.get('outcome')}"

python/packages/ag-ui/tests/ag_ui/golden/test_scenario_predictive_state.py

@@ -0,0 +1,130 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Golden event-stream tests for the predictive state scenario."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from agent_framework import AgentResponseUpdate, Content
+from conftest import StubAgent  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+from event_stream import EventStream  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+
+from agent_framework_ag_ui import AgentFrameworkAgent
+
+PREDICT_CONFIG = {
+    "document": {
+        "tool": "update_document",
+        "tool_argument": "content",
+    }
+}
+
+STATE_SCHEMA = {
+    "document": {"type": "string"},
+}
+
+
+def _build_agent(updates: list[AgentResponseUpdate], **kwargs: Any) -> AgentFrameworkAgent:
+    stub = StubAgent(updates=updates)
+    return AgentFrameworkAgent(
+        agent=stub,
+        state_schema=STATE_SCHEMA,
+        predict_state_config=PREDICT_CONFIG,
+        require_confirmation=False,
+        **kwargs,
+    )
+
+
+async def _run(agent: AgentFrameworkAgent, payload: dict[str, Any]) -> EventStream:
+    return EventStream([event async for event in agent.run(payload)])
+
+
+PAYLOAD: dict[str, Any] = {
+    "thread_id": "thread-predict",
+    "run_id": "run-predict",
+    "messages": [{"role": "user", "content": "Write a document"}],
+    "state": {"document": ""},
+}
+
+
+# ── Golden stream tests ──
+
+
+async def test_predictive_state_emits_deltas_during_tool_args() -> None:
+    """STATE_DELTA events are emitted as tool arguments stream in."""
+    updates = [
+        AgentResponseUpdate(
+            contents=[Content.from_function_call(name="update_document", call_id="call-1", arguments="")],
+            role="assistant",
+        ),
+        AgentResponseUpdate(
+            contents=[
+                Content.from_function_call(name="update_document", call_id="call-1", arguments='{"content": "Hello')
+            ],
+            role="assistant",
+        ),
+        AgentResponseUpdate(
+            contents=[Content.from_function_call(name="update_document", call_id="call-1", arguments=' world"}')],
+            role="assistant",
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    stream.assert_bookends()
+    stream.assert_no_run_error()
+
+    # PredictState custom event should be present
+    custom_events = stream.get("CUSTOM")
+    predict_events = [e for e in custom_events if getattr(e, "name", None) == "PredictState"]
+    assert predict_events, "Expected PredictState custom event"
+
+    # STATE_DELTA events should be emitted during tool arg streaming
+    assert "STATE_DELTA" in stream.types(), "Expected STATE_DELTA events during predictive streaming"
+
+
+async def test_predictive_state_snapshot_after_tool_end() -> None:
+    """STATE_SNAPSHOT is emitted when a predictive tool completes (no confirmation)."""
+    updates = [
+        AgentResponseUpdate(
+            contents=[
+                Content.from_function_call(
+                    name="update_document", call_id="call-1", arguments='{"content": "Final text"}'
+                )
+            ],
+            role="assistant",
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    stream.assert_bookends()
+
+    # Should have initial state snapshot + updated snapshot after tool completion
+    snapshots = stream.get("STATE_SNAPSHOT")
+    assert len(snapshots) >= 1, "Expected at least one STATE_SNAPSHOT"
+
+
+async def test_predictive_state_ordered_events() -> None:
+    """Event ordering: RUN_STARTED → PredictState → STATE_SNAPSHOT → TOOL_CALL_* → STATE_SNAPSHOT → RUN_FINISHED."""
+    updates = [
+        AgentResponseUpdate(
+            contents=[
+                Content.from_function_call(name="update_document", call_id="call-1", arguments='{"content": "doc"}')
+            ],
+            role="assistant",
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    stream.assert_ordered_types(
+        [
+            "RUN_STARTED",
+            "CUSTOM",  # PredictState
+            "STATE_SNAPSHOT",  # Initial state
+            "TOOL_CALL_START",
+            "TOOL_CALL_ARGS",
+            "RUN_FINISHED",
+        ]
+    )

python/packages/ag-ui/tests/ag_ui/golden/test_scenario_shared_state.py

@@ -0,0 +1,110 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Golden event-stream tests for the shared state (structured output) scenario."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from agent_framework import AgentResponseUpdate, Content
+from conftest import StubAgent  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+from event_stream import EventStream  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+from pydantic import BaseModel
+
+from agent_framework_ag_ui import AgentFrameworkAgent
+
+
+class RecipeState(BaseModel):
+    recipe_title: str = ""
+    ingredients: list[str] = []
+    message: str = ""
+
+
+def _build_agent(updates: list[AgentResponseUpdate], **kwargs: Any) -> AgentFrameworkAgent:
+    stub = StubAgent(
+        updates=updates,
+        default_options={"tools": None, "response_format": RecipeState},
+    )
+    return AgentFrameworkAgent(
+        agent=stub,
+        state_schema={
+            "recipe_title": {"type": "string"},
+            "ingredients": {"type": "array", "items": {"type": "string"}},
+        },
+        **kwargs,
+    )
+
+
+async def _run(agent: AgentFrameworkAgent, payload: dict[str, Any]) -> EventStream:
+    return EventStream([event async for event in agent.run(payload)])
+
+
+PAYLOAD: dict[str, Any] = {
+    "thread_id": "thread-state",
+    "run_id": "run-state",
+    "messages": [{"role": "user", "content": "Give me a pasta recipe"}],
+    "state": {"recipe_title": "", "ingredients": []},
+}
+
+
+# ── Golden stream tests ──
+
+
+async def test_shared_state_emits_state_snapshot() -> None:
+    """Structured output agent emits STATE_SNAPSHOT with parsed model fields."""
+    # The structured output agent gets a response that the framework parses as RecipeState
+    updates = [
+        AgentResponseUpdate(
+            contents=[
+                Content.from_text(
+                    text='{"recipe_title": "Pasta Carbonara", "ingredients": ["pasta", "eggs", "cheese"], "message": "Here is your recipe!"}'
+                )
+            ],
+            role="assistant",
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    stream.assert_bookends()
+    stream.assert_no_run_error()
+
+    # Should have STATE_SNAPSHOT with the initial state at minimum
+    stream.assert_has_type("STATE_SNAPSHOT")
+
+
+async def test_shared_state_initial_snapshot_on_first_update() -> None:
+    """When state_schema and state are provided, initial STATE_SNAPSHOT is emitted after RUN_STARTED."""
+    updates = [
+        AgentResponseUpdate(
+            contents=[Content.from_text(text='{"recipe_title": "Test", "ingredients": [], "message": "hi"}')],
+            role="assistant",
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    # RUN_STARTED should be followed by STATE_SNAPSHOT (initial state)
+    stream.assert_ordered_types(["RUN_STARTED", "STATE_SNAPSHOT"])
+
+
+async def test_shared_state_text_emitted_from_message_field() -> None:
+    """Structured output's 'message' field is emitted as text message events."""
+    updates = [
+        AgentResponseUpdate(
+            contents=[
+                Content.from_text(
+                    text='{"recipe_title": "Pasta", "ingredients": ["pasta"], "message": "Enjoy your pasta!"}'
+                )
+            ],
+            role="assistant",
+        ),
+    ]
+    agent = _build_agent(updates)
+    stream = await _run(agent, PAYLOAD)
+
+    # Text should be emitted from the message field
+    text_contents = stream.get("TEXT_MESSAGE_CONTENT")
+    if text_contents:
+        combined = "".join(getattr(e, "delta", "") for e in text_contents)
+        assert "Enjoy your pasta!" in combined

python/packages/ag-ui/tests/ag_ui/golden/test_scenario_subgraphs.py

@@ -0,0 +1,207 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Golden event-stream tests for the workflow HITL (subgraphs) scenario.
+
+Extends the existing test_subgraphs_example_agent.py with EventStream assertions
+on full event ordering, balancing, and interrupt structure.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from event_stream import EventStream  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+
+from agent_framework_ag_ui_examples.agents.subgraphs_agent import subgraphs_agent
+
+
+async def _run(agent: Any, payload: dict[str, Any]) -> EventStream:
+    return EventStream([event async for event in agent.run(payload)])
+
+
+# ── Turn 1: Initial request → flight interrupt ──
+
+
+async def test_subgraphs_turn1_golden_bookends() -> None:
+    """Turn 1 starts with RUN_STARTED and ends with RUN_FINISHED."""
+    agent = subgraphs_agent()
+    stream = await _run(
+        agent,
+        {
+            "thread_id": "thread-sub-golden-1",
+            "run_id": "run-1",
+            "messages": [{"role": "user", "content": "Plan a trip to San Francisco"}],
+        },
+    )
+    stream.assert_bookends()
+
+
+async def test_subgraphs_turn1_no_errors() -> None:
+    """Turn 1 completes without errors."""
+    agent = subgraphs_agent()
+    stream = await _run(
+        agent,
+        {
+            "thread_id": "thread-sub-golden-2",
+            "run_id": "run-1",
+            "messages": [{"role": "user", "content": "Plan a trip"}],
+        },
+    )
+    stream.assert_no_run_error()
+
+
+async def test_subgraphs_turn1_has_step_events() -> None:
+    """Turn 1 emits STEP_STARTED and STEP_FINISHED for workflow executors."""
+    agent = subgraphs_agent()
+    stream = await _run(
+        agent,
+        {
+            "thread_id": "thread-sub-golden-3",
+            "run_id": "run-1",
+            "messages": [{"role": "user", "content": "Plan a trip"}],
+        },
+    )
+    stream.assert_has_type("STEP_STARTED")
+    stream.assert_has_type("STEP_FINISHED")
+
+
+async def test_subgraphs_turn1_interrupt_structure() -> None:
+    """Turn 1 RUN_FINISHED carries flight interrupt with correct structure."""
+    agent = subgraphs_agent()
+    stream = await _run(
+        agent,
+        {
+            "thread_id": "thread-sub-golden-4",
+            "run_id": "run-1",
+            "messages": [{"role": "user", "content": "Plan a trip to SF"}],
+        },
+    )
+
+    interrupt = stream.run_finished_interrupts()
+    assert len(interrupt) > 0
+    interrupt_value = stream.interrupt_metadata_value(interrupt[0])
+    assert interrupt_value["agent"] == "flights"
+    assert len(interrupt_value["options"]) == 2
+
+
+async def test_subgraphs_turn1_text_messages_balanced() -> None:
+    """All text messages in turn 1 are properly balanced."""
+    agent = subgraphs_agent()
+    stream = await _run(
+        agent,
+        {
+            "thread_id": "thread-sub-golden-5",
+            "run_id": "run-1",
+            "messages": [{"role": "user", "content": "Plan a trip"}],
+        },
+    )
+    stream.assert_text_messages_balanced()
+
+
+async def test_subgraphs_turn1_ordered_flow() -> None:
+    """Turn 1 event ordering: RUN_STARTED → STATE_SNAPSHOT → STEP_* → TOOL_CALL_* → RUN_FINISHED."""
+    agent = subgraphs_agent()
+    stream = await _run(
+        agent,
+        {
+            "thread_id": "thread-sub-golden-6",
+            "run_id": "run-1",
+            "messages": [{"role": "user", "content": "Plan a trip"}],
+        },
+    )
+    stream.assert_ordered_types(
+        [
+            "RUN_STARTED",
+            "STATE_SNAPSHOT",
+            "STEP_STARTED",
+            "RUN_FINISHED",
+        ]
+    )
+
+
+# ── Multi-turn: Flight selection → hotel interrupt → completion ──
+
+
+async def test_subgraphs_full_flow_event_ordering() -> None:
+    """Complete 3-turn flow maintains proper event ordering throughout."""
+    agent = subgraphs_agent()
+    thread_id = "thread-sub-golden-full"
+
+    # Turn 1
+    stream1 = await _run(
+        agent,
+        {
+            "thread_id": thread_id,
+            "run_id": "run-1",
+            "messages": [{"role": "user", "content": "Plan a trip to SF from Amsterdam"}],
+        },
+    )
+    stream1.assert_bookends()
+    stream1.assert_no_run_error()
+
+    # Extract flight interrupt
+    interrupt1 = stream1.run_finished_interrupts()[0]
+
+    # Turn 2: Select flight
+    stream2 = await _run(
+        agent,
+        {
+            "thread_id": thread_id,
+            "run_id": "run-2",
+            "resume": {
+                "interrupts": [
+                    {
+                        "id": interrupt1["id"],
+                        "value": json.dumps(
+                            {
+                                "airline": "United",
+                                "departure": "Amsterdam (AMS)",
+                                "arrival": "San Francisco (SFO)",
+                                "price": "$720",
+                                "duration": "12h 15m",
+                            }
+                        ),
+                    }
+                ]
+            },
+        },
+    )
+    stream2.assert_bookends()
+    stream2.assert_no_run_error()
+
+    # Should now have hotel interrupt
+    interrupt2 = stream2.run_finished_interrupts()
+    assert stream2.interrupt_metadata_value(interrupt2[0])["agent"] == "hotels"
+
+    # Turn 3: Select hotel
+    stream3 = await _run(
+        agent,
+        {
+            "thread_id": thread_id,
+            "run_id": "run-3",
+            "resume": {
+                "interrupts": [
+                    {
+                        "id": interrupt2[0]["id"],
+                        "value": json.dumps(
+                            {
+                                "name": "The Ritz-Carlton",
+                                "location": "Nob Hill",
+                                "price_per_night": "$550/night",
+                                "rating": "4.8 stars",
+                            }
+                        ),
+                    }
+                ]
+            },
+        },
+    )
+    stream3.assert_bookends()
+    stream3.assert_no_run_error()
+    stream3.assert_text_messages_balanced()
+
+    # Final turn should not have interrupt
+    finished3 = stream3.last("RUN_FINISHED")
+    final_dump = finished3.model_dump(by_alias=True, exclude_none=True)
+    assert "outcome" not in final_dump, f"Expected no interrupt after completion, got {final_dump.get('outcome')}"

python/packages/ag-ui/tests/ag_ui/golden/test_scenario_workflow.py

@@ -0,0 +1,945 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Comprehensive golden event-stream tests for AgentFrameworkWorkflow.
+
+Covers the full matrix of workflow-specific AG-UI patterns:
+- request_info → TOOL_CALL lifecycle and balancing
+- Executor step events and activity snapshots
+- Text output, dict output, BaseEvent passthrough, AgentResponse output
+- Text deduplication across workflow outputs
+- Workflow error handling → RUN_ERROR
+- Multi-turn interrupt/resume round-trips
+- Empty turns with pending requests
+- Custom workflow events
+- Text message draining on request_info and executor boundaries
+"""
+
+import json
+from typing import Any, cast
+
+from ag_ui.core import EventType, StateSnapshotEvent
+from agent_framework import (
+    AgentResponse,
+    Content,
+    Executor,
+    Message,
+    WorkflowBuilder,
+    WorkflowContext,
+    WorkflowEvent,
+    executor,
+    handler,
+    response_handler,
+)
+from event_stream import EventStream  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+
+from agent_framework_ag_ui import AgentFrameworkWorkflow
+
+
+async def _run(wrapper: AgentFrameworkWorkflow, payload: dict[str, Any]) -> EventStream:
+    return EventStream([event async for event in wrapper.run(payload)])
+
+
+def _payload(
+    msg: str = "go",
+    *,
+    thread_id: str = "thread-wf",
+    run_id: str = "run-wf",
+    **extra: Any,
+) -> dict[str, Any]:
+    return {"thread_id": thread_id, "run_id": run_id, "messages": [{"role": "user", "content": msg}], **extra}
+
+
+# ──────────────────────────────────────────────────────────────────────
+# 1. Basic workflow text output
+# ──────────────────────────────────────────────────────────────────────
+
+
+async def test_workflow_text_output_golden_sequence() -> None:
+    """Simple text output: RUN_STARTED → STEP_STARTED → TEXT_* → STEP_FINISHED → TEXT_MESSAGE_END → RUN_FINISHED."""
+
+    @executor(id="greeter")
+    async def greeter(message: Any, ctx: WorkflowContext[Any, str]) -> None:
+        await ctx.yield_output("Hello from workflow!")
+
+    workflow = WorkflowBuilder(start_executor=greeter).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+    stream = await _run(wrapper, _payload())
+
+    stream.assert_bookends()
+    stream.assert_no_run_error()
+    stream.assert_text_messages_balanced()
+    stream.assert_has_type("TEXT_MESSAGE_START")
+    stream.assert_has_type("TEXT_MESSAGE_CONTENT")
+    stream.assert_has_type("TEXT_MESSAGE_END")
+
+    # Verify actual content
+    deltas = [e.delta for e in stream.get("TEXT_MESSAGE_CONTENT")]
+    assert "Hello from workflow!" in deltas
+
+
+async def test_workflow_text_output_message_id_consistency() -> None:
+    """All text events for a single output share the same message_id."""
+
+    @executor(id="echo")
+    async def echo(message: Any, ctx: WorkflowContext[Any, str]) -> None:
+        await ctx.yield_output("echo reply")
+
+    workflow = WorkflowBuilder(start_executor=echo).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+    stream = await _run(wrapper, _payload())
+
+    stream.assert_message_ids_consistent()
+
+
+# ──────────────────────────────────────────────────────────────────────
+# 2. Executor step events and activity snapshots
+# ──────────────────────────────────────────────────────────────────────
+
+
+async def test_workflow_executor_lifecycle_events() -> None:
+    """Executor invocation produces STEP_STARTED, ACTIVITY_SNAPSHOT, STEP_FINISHED."""
+
+    @executor(id="worker")
+    async def worker(message: Any, ctx: WorkflowContext[Any, str]) -> None:
+        await ctx.yield_output("done")
+
+    workflow = WorkflowBuilder(start_executor=worker).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+    stream = await _run(wrapper, _payload())
+
+    # Step events with executor ID
+    started = [e for e in stream.get("STEP_STARTED") if getattr(e, "step_name", "") == "worker"]
+    finished = [e for e in stream.get("STEP_FINISHED") if getattr(e, "step_name", "") == "worker"]
+    assert started, "Expected STEP_STARTED for 'worker'"
+    assert finished, "Expected STEP_FINISHED for 'worker'"
+
+    # Activity snapshots
+    activities = stream.get("ACTIVITY_SNAPSHOT")
+    assert activities, "Expected ACTIVITY_SNAPSHOT events"
+    # Check one of them has executor payload
+    executor_activities = [a for a in activities if getattr(a, "activity_type", None) == "executor"]
+    assert executor_activities, "Expected executor-type activity snapshots"
+
+
+async def test_workflow_executor_step_ordering() -> None:
+    """STEP_STARTED comes before content, STEP_FINISHED comes after."""
+
+    @executor(id="orderer")
+    async def orderer(message: Any, ctx: WorkflowContext[Any, str]) -> None:
+        await ctx.yield_output("ordered output")
+
+    workflow = WorkflowBuilder(start_executor=orderer).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+    stream = await _run(wrapper, _payload())
+
+    stream.assert_ordered_types(
+        [
+            "RUN_STARTED",
+            "STEP_STARTED",
+            "TEXT_MESSAGE_START",
+            "TEXT_MESSAGE_CONTENT",
+            "STEP_FINISHED",
+            "RUN_FINISHED",
+        ]
+    )
+
+
+# ──────────────────────────────────────────────────────────────────────
+# 3. Dict output → CUSTOM workflow_output
+# ──────────────────────────────────────────────────────────────────────
+
+
+async def test_workflow_dict_output_maps_to_custom_event() -> None:
+    """Non-chat dict output is emitted as CUSTOM workflow_output event."""
+
+    @executor(id="structured")
+    async def structured(message: Any, ctx: WorkflowContext[Any, dict[str, int]]) -> None:
+        await ctx.yield_output({"count": 42, "status": 1})
+
+    workflow = WorkflowBuilder(start_executor=structured).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+    stream = await _run(wrapper, _payload())
+
+    stream.assert_bookends()
+    stream.assert_no_run_error()
+
+    customs = [e for e in stream.get("CUSTOM") if getattr(e, "name", None) == "workflow_output"]
+    assert len(customs) == 1
+    assert customs[0].value == {"count": 42, "status": 1}
+
+    # Should NOT have TEXT_MESSAGE events for dict output
+    assert "TEXT_MESSAGE_CONTENT" not in stream.types()
+
+
+# ──────────────────────────────────────────────────────────────────────
+# 4. BaseEvent passthrough
+# ──────────────────────────────────────────────────────────────────────
+
+
+async def test_workflow_base_event_passthrough() -> None:
+    """AG-UI BaseEvent outputs are yielded directly, not wrapped."""
+
+    @executor(id="stateful")
+    async def stateful(message: Any, ctx: WorkflowContext[Any, StateSnapshotEvent]) -> None:
+        await ctx.yield_output(StateSnapshotEvent(type=EventType.STATE_SNAPSHOT, snapshot={"active_agent": "flights"}))
+
+    workflow = WorkflowBuilder(start_executor=stateful).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+    stream = await _run(wrapper, _payload())
+
+    stream.assert_bookends()
+    snapshots = stream.get("STATE_SNAPSHOT")
+    assert len(snapshots) == 1
+    assert snapshots[0].snapshot["active_agent"] == "flights"
+
+
+# ──────────────────────────────────────────────────────────────────────
+# 5. AgentResponse output (conversation payload)
+# ──────────────────────────────────────────────────────────────────────
+
+
+async def test_workflow_agent_response_output_extracts_latest_assistant() -> None:
+    """AgentResponse output uses only the latest assistant message, not full history."""
+
+    @executor(id="responder")
+    async def responder(message: Any, ctx: WorkflowContext[Any, AgentResponse]) -> None:
+        response = AgentResponse(
+            messages=[
+                Message(role="user", contents=[Content.from_text("My order is damaged")]),
+                Message(role="assistant", contents=[Content.from_text("I'll process your replacement.")]),
+            ]
+        )
+        await ctx.yield_output(response)
+
+    workflow = WorkflowBuilder(start_executor=responder).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+    stream = await _run(wrapper, _payload())
+
+    stream.assert_bookends()
+    stream.assert_text_messages_balanced()
+
+    deltas = [e.delta for e in stream.get("TEXT_MESSAGE_CONTENT")]
+    assert deltas == ["I'll process your replacement."]
+
+
+# ──────────────────────────────────────────────────────────────────────
+# 6. Custom workflow events
+# ──────────────────────────────────────────────────────────────────────
+
+
+class ProgressEvent(WorkflowEvent):
+    """Custom workflow event for testing CUSTOM event mapping."""
+
+    def __init__(self, progress: int) -> None:
+        super().__init__(cast(Any, "custom_progress"), data={"progress": progress})
+
+
+async def test_workflow_custom_events() -> None:
+    """Custom workflow events are mapped to CUSTOM AG-UI events."""
+
+    @executor(id="progress_tracker")
+    async def progress_tracker(message: Any, ctx: WorkflowContext[Any, str]) -> None:
+        await ctx.add_event(ProgressEvent(25))
+        await ctx.yield_output("In progress...")
+        await ctx.add_event(ProgressEvent(100))
+
+    workflow = WorkflowBuilder(start_executor=progress_tracker).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+    stream = await _run(wrapper, _payload())
+
+    stream.assert_bookends()
+    stream.assert_no_run_error()
+
+    progress_events = [e for e in stream.get("CUSTOM") if getattr(e, "name", None) == "custom_progress"]
+    assert len(progress_events) == 2
+    assert progress_events[0].value == {"progress": 25}
+    assert progress_events[1].value == {"progress": 100}
+
+
+# ──────────────────────────────────────────────────────────────────────
+# 7. request_info → TOOL_CALL lifecycle
+# ──────────────────────────────────────────────────────────────────────
+
+
+async def test_workflow_request_info_tool_call_lifecycle() -> None:
+    """request_info emits TOOL_CALL_START/ARGS/END cycle plus CUSTOM request_info."""
+
+    @executor(id="requester")
+    async def requester(message: Any, ctx: WorkflowContext) -> None:
+        await ctx.request_info("Need approval", str, request_id="req-1")
+
+    workflow = WorkflowBuilder(start_executor=requester).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+    stream = await _run(wrapper, _payload())
+
+    stream.assert_bookends()
+    stream.assert_no_run_error()
+
+    # Tool call lifecycle
+    stream.assert_ordered_types(
+        [
+            "RUN_STARTED",
+            "TOOL_CALL_START",
+            "TOOL_CALL_ARGS",
+            "TOOL_CALL_END",
+            "CUSTOM",  # request_info
+            "RUN_FINISHED",
+        ]
+    )
+
+    # Verify tool call details
+    start = stream.first("TOOL_CALL_START")
+    assert start.tool_call_id == "req-1"
+    assert start.tool_call_name == "request_info"
+
+    # TOOL_CALL_ARGS should contain the request payload
+    args = stream.first("TOOL_CALL_ARGS")
+    assert args.tool_call_id == "req-1"
+    parsed_args = json.loads(args.delta)
+    assert parsed_args["request_id"] == "req-1"
+
+    # Tool calls should be balanced
+    stream.assert_tool_calls_balanced()
+
+
+async def test_workflow_request_info_interrupt_in_run_finished() -> None:
+    """request_info populates RUN_FINISHED.outcome.interrupts with the request metadata."""
+
+    @executor(id="requester")
+    async def requester(message: Any, ctx: WorkflowContext) -> None:
+        await ctx.request_info(
+            {"message": "Choose a flight", "options": [{"airline": "KLM"}], "agent": "flights"},
+            dict,
+            request_id="flights-choice",
+        )
+
+    workflow = WorkflowBuilder(start_executor=requester).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+    stream = await _run(wrapper, _payload())
+
+    interrupt = stream.run_finished_interrupts()
+    assert len(interrupt) == 1
+    assert interrupt[0]["id"] == "flights-choice"
+    assert stream.interrupt_metadata_value(interrupt[0])["agent"] == "flights"
+
+
+async def test_workflow_request_info_emits_interrupt_card_event() -> None:
+    """request_info with dict data emits a WorkflowInterruptEvent custom event."""
+
+    @executor(id="requester")
+    async def requester(message: Any, ctx: WorkflowContext) -> None:
+        await ctx.request_info(
+            {"message": "Pick one", "options": ["A", "B"]},
+            dict,
+            request_id="pick-1",
+        )
+
+    workflow = WorkflowBuilder(start_executor=requester).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+    stream = await _run(wrapper, _payload())
+
+    interrupt_cards = [e for e in stream.get("CUSTOM") if getattr(e, "name", None) == "WorkflowInterruptEvent"]
+    assert interrupt_cards, "Expected WorkflowInterruptEvent custom event"
+
+
+# ──────────────────────────────────────────────────────────────────────
+# 8. Text message draining on request_info boundary
+# ──────────────────────────────────────────────────────────────────────
+
+
+async def test_workflow_text_drained_before_request_info() -> None:
+    """Open text message is closed (TEXT_MESSAGE_END) before request_info tool calls begin."""
+
+    @executor(id="text_then_request")
+    async def text_then_request(message: Any, ctx: WorkflowContext) -> None:
+        await ctx.yield_output("Please confirm this action.")  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+        await ctx.request_info("Need approval", str, request_id="approval-1")
+
+    workflow = WorkflowBuilder(start_executor=text_then_request).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+    stream = await _run(wrapper, _payload())
+
+    stream.assert_text_messages_balanced()
+    stream.assert_tool_calls_balanced()
+
+    # TEXT_MESSAGE_END must appear before TOOL_CALL_START
+    types = stream.types()
+    text_end_idx = types.index("TEXT_MESSAGE_END")
+    tool_start_idx = types.index("TOOL_CALL_START")
+    assert text_end_idx < tool_start_idx, (
+        f"TEXT_MESSAGE_END (idx={text_end_idx}) must come before TOOL_CALL_START (idx={tool_start_idx})"
+    )
+
+
+# ──────────────────────────────────────────────────────────────────────
+# 9. Text deduplication
+# ──────────────────────────────────────────────────────────────────────
+
+
+async def test_workflow_skips_duplicate_text_from_snapshot() -> None:
+    """Duplicate text from AgentResponse snapshot is not re-emitted."""
+
+    @executor(id="deduper")
+    async def deduper(message: Any, ctx: WorkflowContext[Any, Any]) -> None:
+        text = "Order processed successfully."
+        await ctx.yield_output(text)
+        # Snapshot repeats the same text
+        await ctx.yield_output(
+            AgentResponse(
+                messages=[
+                    Message(role="user", contents=[Content.from_text("process order")]),
+                    Message(role="assistant", contents=[Content.from_text(text)]),
+                ]
+            )
+        )
+
+    workflow = WorkflowBuilder(start_executor=deduper).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+    stream = await _run(wrapper, _payload())
+
+    stream.assert_text_messages_balanced()
+    deltas = [e.delta for e in stream.get("TEXT_MESSAGE_CONTENT")]
+    # Text should appear only once
+    assert deltas == ["Order processed successfully."]
+
+
+async def test_workflow_skips_consecutive_duplicate_outputs() -> None:
+    """Consecutive identical text outputs are deduplicated."""
+
+    @executor(id="repeater")
+    async def repeater(message: Any, ctx: WorkflowContext[Any, Any]) -> None:
+        text = "Done!"
+        await ctx.yield_output(text)
+        await ctx.yield_output(text)
+
+    workflow = WorkflowBuilder(start_executor=repeater).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+    stream = await _run(wrapper, _payload())
+
+    stream.assert_text_messages_balanced()
+    deltas = [e.delta for e in stream.get("TEXT_MESSAGE_CONTENT")]
+    assert deltas == ["Done!"]
+
+
+async def test_workflow_emits_distinct_consecutive_outputs() -> None:
+    """Distinct text outputs are all emitted, not incorrectly deduplicated."""
+
+    @executor(id="multisayer")
+    async def multisayer(message: Any, ctx: WorkflowContext[Any, str]) -> None:
+        await ctx.yield_output("First part. ")
+        await ctx.yield_output("Second part.")
+
+    workflow = WorkflowBuilder(start_executor=multisayer).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+    stream = await _run(wrapper, _payload())
+
+    stream.assert_text_messages_balanced()
+    deltas = [e.delta for e in stream.get("TEXT_MESSAGE_CONTENT")]
+    assert deltas == ["First part. ", "Second part."]
+
+
+# ──────────────────────────────────────────────────────────────────────
+# 10. Workflow error handling → RUN_ERROR
+# ──────────────────────────────────────────────────────────────────────
+
+
+async def test_workflow_error_emits_run_error_event() -> None:
+    """Exceptions during workflow streaming produce RUN_ERROR events."""
+
+    class FailingWorkflow:
+        def run(self, **kwargs: Any):
+            async def _stream():
+                raise RuntimeError("workflow exploded")
+                yield  # pragma: no cover
+
+            return _stream()
+
+    wrapper = AgentFrameworkWorkflow(workflow=cast(Any, FailingWorkflow()))
+    stream = await _run(wrapper, _payload())
+
+    # Should still have RUN_STARTED
+    stream.assert_has_type("RUN_STARTED")
+    # Should have RUN_ERROR
+    stream.assert_has_type("RUN_ERROR")
+    error = stream.first("RUN_ERROR")
+    assert "workflow exploded" in error.message
+
+
+async def test_workflow_error_preserves_bookend_structure() -> None:
+    """Even on error, RUN_STARTED is the first event."""
+
+    class FailingWorkflow:
+        def run(self, **kwargs: Any):
+            async def _stream():
+                raise ValueError("bad input")
+                yield  # pragma: no cover
+
+            return _stream()
+
+    wrapper = AgentFrameworkWorkflow(workflow=cast(Any, FailingWorkflow()))
+    stream = await _run(wrapper, _payload())
+
+    types = stream.types()
+    assert types[0] == "RUN_STARTED"
+    assert "RUN_ERROR" in types
+
+
+# ──────────────────────────────────────────────────────────────────────
+# 11. Multi-turn request_info interrupt/resume
+# ──────────────────────────────────────────────────────────────────────
+
+
+async def test_workflow_interrupt_resume_round_trip() -> None:
+    """Turn 1: request_info → interrupt. Turn 2: resume → completion."""
+
+    class RequesterExecutor(Executor):
+        def __init__(self) -> None:
+            super().__init__(id="requester")
+
+        @handler
+        async def start(self, message: Any, ctx: WorkflowContext) -> None:
+            await ctx.request_info("Choose an option", str, request_id="choice-1")
+
+        @response_handler
+        async def handle_choice(self, original: str, response: str, ctx: WorkflowContext) -> None:
+            await ctx.yield_output(f"You chose: {response}")  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+
+    workflow = WorkflowBuilder(start_executor=RequesterExecutor()).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+
+    # Turn 1
+    stream1 = await _run(wrapper, _payload(thread_id="thread-resume", run_id="run-1"))
+    stream1.assert_bookends()
+    stream1.assert_no_run_error()
+    stream1.assert_tool_calls_balanced()
+
+    interrupt1 = stream1.run_finished_interrupts()
+    assert interrupt1, "Expected interrupt"
+    assert interrupt1[0]["id"] == "choice-1"
+
+    # Turn 2: resume
+    stream2 = await _run(
+        wrapper,
+        {
+            "thread_id": "thread-resume",
+            "run_id": "run-2",
+            "messages": [],
+            "resume": {"interrupts": [{"id": "choice-1", "value": "Option A"}]},
+        },
+    )
+    stream2.assert_has_run_lifecycle()
+    stream2.assert_no_run_error()
+    stream2.assert_text_messages_balanced()
+
+    # Should have the response text
+    deltas = [e.delta for e in stream2.get("TEXT_MESSAGE_CONTENT")]
+    assert any("Option A" in d for d in deltas), f"Expected 'Option A' in deltas: {deltas}"
+
+    # No interrupt after resume
+    finished2 = stream2.last("RUN_FINISHED")
+    dumped2 = finished2.model_dump(by_alias=True, exclude_none=True)
+    assert "outcome" not in dumped2
+
+
+async def test_workflow_forwarded_props_resume() -> None:
+    """forwarded_props.command.resume should resume with canonical resume entries."""
+
+    @executor(id="requester")
+    async def requester(message: Any, ctx: WorkflowContext) -> None:
+        await ctx.request_info({"options": [{"name": "A"}]}, dict, request_id="pick")
+
+    workflow = WorkflowBuilder(start_executor=requester).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+
+    # Turn 1
+    await _run(wrapper, _payload(thread_id="thread-fwd", run_id="run-1"))
+
+    # Turn 2 via forwarded_props
+    stream2 = await _run(
+        wrapper,
+        {
+            "thread_id": "thread-fwd",
+            "run_id": "run-2",
+            "messages": [],
+            "forwarded_props": {
+                "command": {"resume": [{"interruptId": "pick", "status": "resolved", "payload": {"name": "A"}}]}
+            },
+        },
+    )
+    stream2.assert_bookends()
+    stream2.assert_no_run_error()
+
+    finished = stream2.last("RUN_FINISHED")
+    assert "outcome" not in finished.model_dump(by_alias=True, exclude_none=True)
+
+
+# ──────────────────────────────────────────────────────────────────────
+# 12. Empty turns with pending requests
+# ──────────────────────────────────────────────────────────────────────
+
+
+async def test_workflow_empty_turn_with_pending_request_emits_run_error() -> None:
+    """An empty turn with a pending request must provide canonical resume entries."""
+
+    @executor(id="requester")
+    async def requester(message: Any, ctx: WorkflowContext) -> None:
+        await ctx.request_info({"prompt": "choose"}, dict, request_id="pick-one")
+
+    workflow = WorkflowBuilder(start_executor=requester).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+
+    # Turn 1: trigger the request
+    await _run(wrapper, _payload(thread_id="thread-empty", run_id="run-1"))
+
+    # Turn 2: empty messages, no resume
+    stream2 = await _run(
+        wrapper,
+        {
+            "thread_id": "thread-empty",
+            "run_id": "run-2",
+            "messages": [],
+        },
+    )
+    stream2.assert_has_type("RUN_STARTED")
+    run_error = stream2.last("RUN_ERROR")
+    assert run_error.code == "WORKFLOW_RESUME_REQUIRED"
+
+
+async def test_workflow_empty_turn_no_pending_requests() -> None:
+    """Empty turn with no pending requests produces clean bookends."""
+
+    @executor(id="noop")
+    async def noop(message: Any, ctx: WorkflowContext[Any, str]) -> None:
+        await ctx.yield_output("done")
+
+    workflow = WorkflowBuilder(start_executor=noop).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+
+    # Run once to completion
+    await _run(wrapper, _payload(thread_id="thread-empty-clean", run_id="run-1"))
+
+    # Empty turn
+    stream2 = await _run(
+        wrapper,
+        {
+            "thread_id": "thread-empty-clean",
+            "run_id": "run-2",
+            "messages": [],
+        },
+    )
+    stream2.assert_bookends()
+    stream2.assert_no_run_error()
+
+
+# ──────────────────────────────────────────────────────────────────────
+# 13. Usage content as CUSTOM event
+# ──────────────────────────────────────────────────────────────────────
+
+
+async def test_workflow_usage_output_maps_to_custom_event() -> None:
+    """Usage Content outputs are surfaced as custom usage events."""
+
+    @executor(id="usage_reporter")
+    async def usage_reporter(message: Any, ctx: WorkflowContext[Any, Content]) -> None:
+        await ctx.yield_output(
+            Content.from_usage({"input_token_count": 100, "output_token_count": 50, "total_token_count": 150})
+        )
+
+    workflow = WorkflowBuilder(start_executor=usage_reporter).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+    stream = await _run(wrapper, _payload())
+
+    stream.assert_bookends()
+    stream.assert_no_run_error()
+
+    usage_events = [e for e in stream.get("CUSTOM") if getattr(e, "name", None) == "usage"]
+    assert len(usage_events) == 1
+    assert usage_events[0].value["input_token_count"] == 100
+    assert usage_events[0].value["total_token_count"] == 150
+
+
+# ──────────────────────────────────────────────────────────────────────
+# 14. Approval flow (Content-based request_info)
+# ──────────────────────────────────────────────────────────────────────
+
+
+async def test_workflow_approval_flow_round_trip() -> None:
+    """function_approval_request via request_info, then resume with approval response."""
+
+    class ApprovalExecutor(Executor):
+        def __init__(self) -> None:
+            super().__init__(id="approval_exec")
+
+        @handler
+        async def start(self, message: Any, ctx: WorkflowContext) -> None:
+            function_call = Content.from_function_call(
+                call_id="refund-call",
+                name="submit_refund",
+                arguments={"order_id": "12345", "amount": "$89.99"},
+            )
+            approval_request = Content.from_function_approval_request(id="approval-1", function_call=function_call)
+            await ctx.request_info(approval_request, Content, request_id="approval-1")
+
+        @response_handler
+        async def handle_approval(self, original_request: Content, response: Content, ctx: WorkflowContext) -> None:
+            status = "approved" if bool(response.approved) else "rejected"
+            await ctx.yield_output(f"Refund {status}.")  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+
+    workflow = WorkflowBuilder(start_executor=ApprovalExecutor()).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+
+    # Turn 1: request approval
+    stream1 = await _run(wrapper, _payload(thread_id="thread-approval", run_id="run-1"))
+    stream1.assert_bookends()
+    stream1.assert_no_run_error()
+
+    interrupt1 = stream1.run_finished_interrupts()
+    assert interrupt1, "Expected approval interrupt"
+    interrupt_value = stream1.interrupt_metadata_value(interrupt1[0])
+
+    # Turn 2: approve
+    stream2 = await _run(
+        wrapper,
+        {
+            "thread_id": "thread-approval",
+            "run_id": "run-2",
+            "messages": [],
+            "resume": {
+                "interrupts": [
+                    {
+                        "id": "approval-1",
+                        "value": {
+                            "type": "function_approval_response",
+                            "approved": True,
+                            "id": interrupt_value.get("id", "approval-1"),
+                            "function_call": interrupt_value.get("function_call"),
+                        },
+                    }
+                ]
+            },
+        },
+    )
+    stream2.assert_has_run_lifecycle()
+    stream2.assert_no_run_error()
+    stream2.assert_text_messages_balanced()
+
+    deltas = [e.delta for e in stream2.get("TEXT_MESSAGE_CONTENT")]
+    assert any("approved" in d for d in deltas)
+
+    # No more interrupt
+    finished2 = stream2.last("RUN_FINISHED")
+    assert "outcome" not in finished2.model_dump(by_alias=True, exclude_none=True)
+
+
+# ──────────────────────────────────────────────────────────────────────
+# 15. Message list request/response coercion
+# ──────────────────────────────────────────────────────────────────────
+
+
+async def test_workflow_message_list_resume() -> None:
+    """Resume with list[Message] payload coerces correctly into workflow response."""
+
+    class MessageRequestExecutor(Executor):
+        def __init__(self) -> None:
+            super().__init__(id="msg_request")
+
+        @handler
+        async def start(self, message: Any, ctx: WorkflowContext) -> None:
+            await ctx.request_info({"prompt": "Need follow-up"}, list[Message], request_id="handoff")
+
+        @response_handler
+        async def handle_input(self, original: dict, response: list[Message], ctx: WorkflowContext) -> None:
+            user_text = response[0].text if response else ""
+            await ctx.yield_output(f"Got: {user_text}")  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+
+    workflow = WorkflowBuilder(start_executor=MessageRequestExecutor()).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+
+    # Turn 1
+    await _run(wrapper, _payload(thread_id="thread-msg", run_id="run-1"))
+
+    # Turn 2: resume with message list
+    stream2 = await _run(
+        wrapper,
+        {
+            "thread_id": "thread-msg",
+            "run_id": "run-2",
+            "messages": [],
+            "resume": {
+                "interrupts": [
+                    {
+                        "id": "handoff",
+                        "value": [
+                            {"role": "user", "contents": [{"type": "text", "text": "Ship a replacement"}]},
+                        ],
+                    }
+                ]
+            },
+        },
+    )
+    stream2.assert_has_run_lifecycle()
+    stream2.assert_no_run_error()
+    stream2.assert_text_messages_balanced()
+
+    deltas = [e.delta for e in stream2.get("TEXT_MESSAGE_CONTENT")]
+    assert any("replacement" in d for d in deltas)
+
+
+# ──────────────────────────────────────────────────────────────────────
+# 16. Plain text follow-up does NOT infer interrupt response
+# ──────────────────────────────────────────────────────────────────────
+
+
+async def test_workflow_plain_text_does_not_resume_pending_dict_request() -> None:
+    """Plain text user follow-up should fail instead of being coerced into a dict response."""
+
+    @executor(id="requester")
+    async def requester(message: Any, ctx: WorkflowContext) -> None:
+        await ctx.request_info(
+            {"message": "Choose a flight", "options": [{"airline": "KLM"}], "agent": "flights"},
+            dict,
+            request_id="flights-choice",
+        )
+
+    workflow = WorkflowBuilder(start_executor=requester).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+
+    # Turn 1
+    await _run(wrapper, _payload(thread_id="thread-nocoerce", run_id="run-1"))
+
+    # Turn 2: plain text follow-up with request_info tool call in history
+    stream2 = await _run(
+        wrapper,
+        {
+            "thread_id": "thread-nocoerce",
+            "run_id": "run-2",
+            "messages": [
+                {
+                    "role": "assistant",
+                    "content": "",
+                    "tool_calls": [
+                        {
+                            "id": "flights-choice",
+                            "type": "function",
+                            "function": {"name": "request_info", "arguments": "{}"},
+                        }
+                    ],
+                },
+                {"role": "user", "content": "I prefer KLM please"},
+            ],
+        },
+    )
+    stream2.assert_has_type("RUN_STARTED")
+    run_error = stream2.last("RUN_ERROR")
+    assert run_error.code == "WORKFLOW_RESUME_REQUIRED"
+
+
+# ──────────────────────────────────────────────────────────────────────
+# 17. Workflow factory (thread-scoped workflows)
+# ──────────────────────────────────────────────────────────────────────
+
+
+async def test_workflow_factory_thread_scoping() -> None:
+    """workflow_factory creates separate workflow instances per thread_id."""
+
+    def make_workflow(thread_id: str):
+        @executor(id="echo")
+        async def echo(message: Any, ctx: WorkflowContext[Any, str]) -> None:
+            await ctx.yield_output(f"Thread: {thread_id}")
+
+        return WorkflowBuilder(start_executor=echo).build()
+
+    wrapper = AgentFrameworkWorkflow(workflow_factory=make_workflow)
+
+    stream_a = await _run(wrapper, _payload(thread_id="thread-a", run_id="run-a"))
+    stream_b = await _run(wrapper, _payload(thread_id="thread-b", run_id="run-b"))
+
+    stream_a.assert_bookends()
+    stream_b.assert_bookends()
+
+    deltas_a = [e.delta for e in stream_a.get("TEXT_MESSAGE_CONTENT")]
+    deltas_b = [e.delta for e in stream_b.get("TEXT_MESSAGE_CONTENT")]
+    assert any("thread-a" in d for d in deltas_a)
+    assert any("thread-b" in d for d in deltas_b)
+
+
+# ──────────────────────────────────────────────────────────────────────
+# 18. Multiple request_info calls in sequence
+# ──────────────────────────────────────────────────────────────────────
+
+
+async def test_workflow_sequential_request_info_interrupts() -> None:
+    """Two chained executors each requesting info: first triggers interrupt, resume, then second triggers interrupt.
+
+    This mirrors the subgraphs_agent pattern where separate executors handle sequential interactions.
+    """
+
+    class NameRequester(Executor):
+        def __init__(self) -> None:
+            super().__init__(id="name_requester")
+
+        @handler
+        async def start(self, message: Any, ctx: WorkflowContext[str]) -> None:
+            await ctx.request_info("What's your name?", str, request_id="name-req")
+
+        @response_handler
+        async def handle_name(self, original: str, response: str, ctx: WorkflowContext[str]) -> None:
+            await ctx.send_message(response)
+
+    class DestRequester(Executor):
+        def __init__(self) -> None:
+            super().__init__(id="dest_requester")
+
+        @handler
+        async def start(self, message: str, ctx: WorkflowContext[str]) -> None:
+            self._name = message
+            await ctx.request_info("Where to?", str, request_id="dest-req")
+
+        @response_handler
+        async def handle_dest(self, original: str, response: str, ctx: WorkflowContext[str]) -> None:
+            await ctx.yield_output(f"Booking for {self._name} to {response}")  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+
+    name_requester = NameRequester()
+    dest_requester = DestRequester()
+    workflow = WorkflowBuilder(start_executor=name_requester).add_chain([name_requester, dest_requester]).build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+
+    # Turn 1
+    stream1 = await _run(wrapper, _payload(thread_id="thread-seq", run_id="run-1"))
+    stream1.assert_bookends()
+    stream1.assert_tool_calls_balanced()
+    interrupt1 = stream1.run_finished_interrupts()
+    assert interrupt1[0]["id"] == "name-req"
+
+    # Turn 2: answer name → triggers second executor's request_info
+    stream2 = await _run(
+        wrapper,
+        {
+            "thread_id": "thread-seq",
+            "run_id": "run-2",
+            "messages": [],
+            "resume": {"interrupts": [{"id": "name-req", "value": "Alice"}]},
+        },
+    )
+    stream2.assert_has_run_lifecycle()
+    stream2.assert_tool_calls_balanced()
+    interrupt2 = stream2.run_finished_interrupts()
+    assert interrupt2[0]["id"] == "dest-req"
+
+    # Turn 3: answer destination → completion
+    stream3 = await _run(
+        wrapper,
+        {
+            "thread_id": "thread-seq",
+            "run_id": "run-3",
+            "messages": [],
+            "resume": {"interrupts": [{"id": "dest-req", "value": "Paris"}]},
+        },
+    )
+    stream3.assert_has_run_lifecycle()
+    stream3.assert_no_run_error()
+    stream3.assert_text_messages_balanced()
+
+    deltas = [e.delta for e in stream3.get("TEXT_MESSAGE_CONTENT")]
+    assert any("Alice" in d and "Paris" in d for d in deltas)
+    assert "outcome" not in stream3.last("RUN_FINISHED").model_dump(by_alias=True, exclude_none=True)

python/packages/ag-ui/tests/ag_ui/sse_helpers.py

@@ -0,0 +1,72 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""SSE parsing helpers for AG-UI HTTP round-trip tests."""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from event_stream import EventStream  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+
+
+def parse_sse_response(response_content: bytes) -> list[dict[str, Any]]:
+    """Parse raw SSE bytes from TestClient into a list of event dicts.
+
+    Each SSE event is a ``data: {...}`` line followed by a blank line.
+    """
+    text = response_content.decode("utf-8")
+    events: list[dict[str, Any]] = []
+    decode_errors: list[str] = []
+    for line in text.splitlines():
+        if line.startswith("data: "):
+            payload = line[6:]
+            try:
+                events.append(json.loads(payload))
+            except json.JSONDecodeError as exc:
+                decode_errors.append(f"payload={payload!r}, error={exc}")
+                continue
+    if decode_errors:
+        joined = "; ".join(decode_errors)
+        raise AssertionError(f"Failed to decode one or more SSE data lines: {joined}")
+    return events
+
+
+def parse_sse_to_event_stream(response_content: bytes) -> EventStream:
+    """Parse SSE bytes and wrap in EventStream for structured assertions.
+
+    Returns an EventStream over lightweight SimpleNamespace objects that
+    mirror AG-UI event attributes (type, message_id, tool_call_id, etc.)
+    so that EventStream assertion methods work.
+    """
+    from types import SimpleNamespace
+
+    raw_events = parse_sse_response(response_content)
+    events: list[Any] = []
+    for raw in raw_events:
+        # Normalize camelCase keys to snake_case attributes that EventStream expects
+        ns = SimpleNamespace()
+        ns.type = raw.get("type", "")
+        ns.raw = raw
+        # Map common camelCase fields
+        for camel, snake in _FIELD_MAP.items():
+            if camel in raw:
+                setattr(ns, snake, raw[camel])
+        # Also keep camelCase as attributes for direct access
+        for key, value in raw.items():
+            if not hasattr(ns, key):
+                setattr(ns, key, value)
+        events.append(ns)
+    return EventStream(events)
+
+
+_FIELD_MAP: dict[str, str] = {
+    "messageId": "message_id",
+    "runId": "run_id",
+    "threadId": "thread_id",
+    "toolCallId": "tool_call_id",
+    "toolCallName": "tool_call_name",
+    "toolName": "tool_call_name",
+    "parentMessageId": "parent_message_id",
+    "stepName": "step_name",
+}

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

@@ -0,0 +1,564 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for AGUIChatClient."""
+
+import json
+from collections.abc import AsyncGenerator, Awaitable, MutableSequence
+from typing import Any
+
+from ag_ui.core import Interrupt, ResumeEntry
+from agent_framework import (
+    ChatOptions,
+    ChatResponse,
+    ChatResponseUpdate,
+    Content,
+    Message,
+    ResponseStream,
+    tool,
+)
+from pytest import MonkeyPatch
+
+from agent_framework_ag_ui._client import AGUIChatClient
+from agent_framework_ag_ui._http_service import AGUIHttpService
+
+
+class StubAGUIChatClient(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[Message]) -> tuple[list[Message], dict[str, Any] | None]:
+        """Expose state extraction helper."""
+        return self._extract_state_from_messages(messages)
+
+    def convert_messages_to_agui_format(self, messages: list[Message]) -> list[dict[str, Any]]:
+        """Expose message conversion helper."""
+        return self._convert_messages_to_agui_format(messages)
+
+    def get_thread_id(self, options: ChatOptions[Any] | dict[str, Any] | None) -> str:
+        """Expose thread id helper."""
+        return self._get_thread_id(options)  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+
+    def inner_get_response(
+        self,
+        *,
+        messages: MutableSequence[Message],
+        options: ChatOptions[Any] | dict[str, Any] | None,
+        stream: bool = False,
+    ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
+        """Proxy to protected response call."""
+        return self._inner_get_response(messages=messages, options=options, stream=stream)  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+
+
+class TestAGUIChatClient:
+    """Test suite for AGUIChatClient."""
+
+    async def test_client_initialization(self) -> None:
+        """Test client initialization."""
+        client = StubAGUIChatClient(endpoint="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 StubAGUIChatClient(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 = StubAGUIChatClient(endpoint="http://localhost:8888/")
+        messages = [
+            Message(role="user", contents=["Hello"]),
+            Message(role="assistant", contents=["Hi there"]),
+        ]
+
+        result_messages, state = client.extract_state_from_messages(messages)
+
+        assert result_messages == messages
+        assert state is None
+
+    async def test_extract_state_from_messages_with_state(self) -> None:
+        """Test state extraction from last message."""
+        import base64
+
+        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
+
+        state_data = {"key": "value", "count": 42}
+        state_json = json.dumps(state_data)
+        state_b64 = base64.b64encode(state_json.encode("utf-8")).decode("utf-8")
+
+        messages = [
+            Message(role="user", contents=["Hello"]),
+            Message(
+                role="user",
+                contents=[Content.from_uri(uri=f"data:application/json;base64,{state_b64}")],
+            ),
+        ]
+
+        result_messages, state = client.extract_state_from_messages(messages)
+
+        assert len(result_messages) == 1
+        assert result_messages[0].text == "Hello"
+        assert state == state_data
+
+    async def test_extract_state_invalid_json(self) -> None:
+        """Test state extraction with invalid JSON."""
+        import base64
+
+        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
+
+        invalid_json = "not valid json"
+        state_b64 = base64.b64encode(invalid_json.encode("utf-8")).decode("utf-8")
+
+        messages = [
+            Message(
+                role="user",
+                contents=[Content.from_uri(uri=f"data:application/json;base64,{state_b64}")],
+            ),
+        ]
+
+        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 = StubAGUIChatClient(endpoint="http://localhost:8888/")
+        messages = [
+            Message(role="user", contents=["What is the weather?"]),
+            Message(role="assistant", contents=["Let me check."], message_id="msg_123"),
+        ]
+
+        agui_messages = client.convert_messages_to_agui_format(messages)
+
+        assert len(agui_messages) == 2
+        assert agui_messages[0]["role"] == "user"
+        assert agui_messages[0]["content"] == "What is the weather?"
+        assert agui_messages[1]["role"] == "assistant"
+        assert agui_messages[1]["content"] == "Let me check."
+        assert agui_messages[1]["id"] == "msg_123"
+
+    async def test_get_thread_id_from_metadata(self) -> None:
+        """Test thread ID extraction from metadata."""
+        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
+        chat_options = ChatOptions(metadata={"thread_id": "existing_thread_123"})
+
+        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 = StubAGUIChatClient(endpoint="http://localhost:8888/")
+        chat_options = ChatOptions()
+
+        thread_id = client.get_thread_id(chat_options)
+
+        assert thread_id.startswith("thread_")
+        assert len(thread_id) > 7
+
+    async def test_get_response_streaming(self, monkeypatch: MonkeyPatch) -> None:
+        """Test streaming response method."""
+        mock_events = [
+            {"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
+            {"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": "Hello"},
+            {"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": " world"},
+            {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
+        ]
+
+        async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
+            for event in mock_events:
+                yield event
+
+        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
+        monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
+
+        messages = [Message(role="user", contents=["Test message"])]
+        chat_options = ChatOptions()
+
+        updates: list[ChatResponseUpdate] = []
+        stream = client.inner_get_response(messages=messages, stream=True, options=chat_options)
+        assert isinstance(stream, ResponseStream)
+        async for update in stream:
+            updates.append(update)
+
+        assert len(updates) == 4
+        assert updates[0].additional_properties is not None
+        assert updates[0].additional_properties["thread_id"] == "thread_1"
+
+        first_content = updates[1].contents[0]
+        second_content = updates[2].contents[0]
+        assert first_content.type == "text"
+        assert second_content.type == "text"
+        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"},
+            {"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": "Complete response"},
+            {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
+        ]
+
+        async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
+            for event in mock_events:
+                yield event
+
+        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
+        monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
+
+        messages = [Message(role="user", contents=["Test message"])]
+        chat_options: dict[str, Any] = {}
+
+        response = await client.inner_get_response(messages=messages, 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: MonkeyPatch) -> None:
+        """Test that client tool metadata is sent to server.
+
+        Client tool metadata (name, description, schema) is sent to server for planning.
+        When server requests a client function, function invocation mixin
+        intercepts and executes it locally. This matches .NET AG-UI implementation.
+        """
+        from agent_framework import tool
+
+        @tool
+        def test_tool(param: str) -> str:
+            """Test tool."""
+            return "result"
+
+        mock_events = [
+            {"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
+            {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
+        ]
+
+        async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
+            # Client tool metadata should be sent to server
+            tools: list[dict[str, Any]] | None = kwargs.get("tools")
+            assert tools is not None
+            assert len(tools) == 1
+            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 = StubAGUIChatClient(endpoint="http://localhost:8888/")
+        monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
+
+        messages = [Message(role="user", contents=["Test with tools"])]
+        chat_options = ChatOptions(tools=[test_tool])
+
+        response = await client.inner_get_response(messages=messages, options=chat_options)
+
+        assert response is not 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 = [
+            {"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
+            {"type": "TOOL_CALL_START", "toolCallId": "call_1", "toolName": "get_time_zone"},
+            {"type": "TOOL_CALL_ARGS", "toolCallId": "call_1", "delta": '{"location": "Seattle"}'},
+            {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
+        ]
+
+        async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
+            for event in mock_events:
+                yield event
+
+        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
+        monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
+
+        messages = [Message(role="user", contents=["Test server tool execution"])]
+
+        updates: list[ChatResponseUpdate] = []
+        async for update in client.get_response(messages, stream=True):
+            updates.append(update)
+
+        function_calls = [
+            content for update in updates for content in update.contents if content.type == "function_call"
+        ]
+        assert function_calls
+        assert function_calls[0].name == "get_time_zone"
+
+        assert not any(content.type == "server_function_call" for update in updates for content in update.contents)
+
+    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."""
+
+        @tool
+        def client_tool() -> str:
+            """Client tool stub."""
+            return "client"
+
+        mock_events = [
+            {"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
+            {"type": "TOOL_CALL_START", "toolCallId": "call_1", "toolName": "get_time_zone"},
+            {"type": "TOOL_CALL_ARGS", "toolCallId": "call_1", "delta": '{"location": "Seattle"}'},
+            {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
+        ]
+
+        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: 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 = StubAGUIChatClient(endpoint="http://localhost:8888/")
+        monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
+
+        messages = [Message(role="user", contents=["Test server tool execution"])]
+
+        async for _ in client.get_response(
+            messages, stream=True, options={"tool_choice": "auto", "tools": [client_tool]}
+        ):
+            pass
+
+    async def test_state_transmission(self, monkeypatch: MonkeyPatch) -> None:
+        """Test state is properly transmitted to server."""
+        import base64
+
+        state_data = {"user_id": "123", "session": "abc"}
+        state_json = json.dumps(state_data)
+        state_b64 = base64.b64encode(state_json.encode("utf-8")).decode("utf-8")
+
+        messages = [
+            Message(role="user", contents=["Hello"]),
+            Message(
+                role="user",
+                contents=[Content.from_uri(uri=f"data:application/json;base64,{state_b64}")],
+            ),
+        ]
+
+        mock_events = [
+            {"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
+            {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
+        ]
+
+        async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
+            assert kwargs.get("state") == state_data
+            for event in mock_events:
+                yield event
+
+        client = StubAGUIChatClient(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, options=chat_options)
+
+        assert response is not None
+
+    async def test_extract_state_from_empty_messages(self) -> None:
+        """Empty messages list returns empty list and None state."""
+        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
+        result_messages, state = client.extract_state_from_messages([])
+        assert result_messages == []
+        assert state is None
+
+    async def test_register_server_tool_non_dict_config(self) -> None:
+        """Non-dict function_invocation_configuration is a no-op."""
+        client = StubAGUIChatClient(
+            endpoint="http://localhost:8888/",
+            function_invocation_configuration=None,  # type: ignore[arg-type]
+        )
+        # Should not raise
+        client._register_server_tool_placeholder("some_tool")
+
+    async def test_non_streaming_response(self, monkeypatch: MonkeyPatch) -> None:
+        """Non-streaming path collects updates into ChatResponse."""
+        mock_events = [
+            {"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
+            {"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": "Hello"},
+            {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
+        ]
+
+        async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
+            for event in mock_events:
+                yield event
+
+        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
+        monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
+
+        messages = [Message(role="user", contents=["Test"])]
+        response = await client.inner_get_response(messages=messages, options={}, stream=False)
+
+        assert response is not None
+        assert len(response.messages) > 0
+
+    async def test_client_tool_sets_additional_properties(self, monkeypatch: MonkeyPatch) -> None:
+        """Client tool content gets agui_thread_id additional property."""
+
+        @tool
+        def my_tool(param: str) -> str:
+            """My tool."""
+            return "result"
+
+        mock_events = [
+            {"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
+            {"type": "TOOL_CALL_START", "toolCallId": "call_1", "toolName": "my_tool"},
+            {"type": "TOOL_CALL_ARGS", "toolCallId": "call_1", "delta": '{"param": "test"}'},
+            {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
+        ]
+
+        async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
+            for event in mock_events:
+                yield event
+
+        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
+        monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
+
+        messages = [Message(role="user", contents=["Test"])]
+        updates: list[ChatResponseUpdate] = []
+        stream = client.inner_get_response(messages=messages, stream=True, options={"tools": [my_tool]})
+        assert isinstance(stream, ResponseStream)
+        async for update in stream:
+            updates.append(update)
+
+        # Find the function_call content - it should have agui_thread_id
+        found = False
+        for update in updates:
+            for content in update.contents:
+                if content.type == "function_call" and content.name == "my_tool":
+                    assert content.additional_properties is not None
+                    assert "agui_thread_id" in content.additional_properties
+                    found = True
+                    break
+        assert found, "Expected to find function_call content for my_tool"
+
+    async def test_tool_call_args_id_mismatch_does_not_execute_current_client_tool(
+        self, monkeypatch: MonkeyPatch
+    ) -> None:
+        """Mismatched TOOL_CALL_ARGS must not be rebound to the latest client tool."""
+        executed: list[int] = []
+
+        @tool
+        def danger_tool(amount: int) -> str:
+            """Record an invocation for the regression assertion."""
+            executed.append(amount)
+            return f"danger={amount}"
+
+        call_count = 0
+
+        async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
+            nonlocal call_count
+            call_count += 1
+            if call_count == 1:
+                mock_events = [
+                    {"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
+                    {"type": "TOOL_CALL_START", "toolCallId": "safe", "toolName": "safe_tool"},
+                    {"type": "TOOL_CALL_START", "toolCallId": "danger", "toolName": "danger_tool"},
+                    {"type": "TOOL_CALL_ARGS", "toolCallId": "safe", "delta": '{"amount": 100}'},
+                    {"type": "TOOL_CALL_END", "toolCallId": "safe"},
+                    {"type": "TOOL_CALL_END", "toolCallId": "danger"},
+                    {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
+                ]
+            else:
+                mock_events = [
+                    {"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_2"},
+                    {"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": "done"},
+                    {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_2"},
+                ]
+
+            for event in mock_events:
+                yield event
+
+        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
+        monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
+
+        response = await client.get_response(
+            [Message(role="user", contents=["Test"])],
+            options={"tools": [danger_tool]},
+        )
+
+        assert response.text == "done"
+        assert executed == []
+
+    async def test_interrupt_options_transmission(self, monkeypatch: MonkeyPatch) -> None:
+        """Interrupt option fields are forwarded to the HTTP service."""
+        available_interrupts = [{"id": "req_1", "type": "request_info"}]
+        expected_available_interrupts = [{"id": "req_1", "reason": "input_required"}]
+        resume_payload = {"interrupts": [{"id": "req_1", "value": "approved"}]}
+        expected_resume_payload = [{"interruptId": "req_1", "status": "resolved", "payload": "approved"}]
+
+        mock_events = [
+            {"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
+            {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
+        ]
+
+        async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
+            assert kwargs.get("available_interrupts") == expected_available_interrupts
+            assert kwargs.get("resume") == expected_resume_payload
+            for event in mock_events:
+                yield event
+
+        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
+        monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
+
+        messages = [Message(role="user", contents=["continue"])]
+        options = {
+            "available_interrupts": available_interrupts,
+            "resume": resume_payload,
+        }
+
+        response = await client.inner_get_response(messages=messages, options=options)
+        assert response is not None
+
+    async def test_typed_interrupt_options_forward_canonical_protocol_shape(self, monkeypatch: MonkeyPatch) -> None:
+        """Typed interrupt options are forwarded as canonical protocol JSON."""
+        mock_events = [
+            {"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
+            {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
+        ]
+
+        async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
+            assert kwargs.get("available_interrupts") == [
+                {
+                    "id": "approval_1",
+                    "reason": "tool_call",
+                    "toolCallId": "call_1",
+                    "responseSchema": {"type": "object"},
+                }
+            ]
+            assert kwargs.get("resume") == [
+                {"interruptId": "approval_1", "status": "resolved", "payload": {"approved": True}}
+            ]
+            for event in mock_events:
+                yield event
+
+        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
+        monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
+
+        options: dict[str, Any] = {
+            "available_interrupts": [
+                Interrupt(
+                    id="approval_1",
+                    reason="tool_call",
+                    tool_call_id="call_1",
+                    response_schema={"type": "object"},
+                )
+            ],
+            "resume": [ResumeEntry(interrupt_id="approval_1", status="resolved", payload={"approved": True})],
+        }
+
+        response = await client.inner_get_response(
+            messages=[Message(role="user", contents=["continue"])],
+            options=options,
+        )
+
+        assert response is not None

python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py

@@ -0,0 +1,558 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for TOOL_CALL_RESULT event emission on approval resume flows."""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from agent_framework import AgentResponseUpdate, Content, FunctionTool
+from conftest import StubAgent  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+
+from agent_framework_ag_ui._agent import AgentConfig
+from agent_framework_ag_ui._agent_run import run_agent_stream
+
+
+def _make_weather_tool() -> FunctionTool:
+    """Create a real executable weather tool with approval_mode='always_require'."""
+
+    def get_weather(city: str) -> str:
+        return f"Sunny in {city}"
+
+    return FunctionTool(
+        name="get_weather",
+        description="Get the weather for a city",
+        func=get_weather,
+        approval_mode="always_require",
+    )
+
+
+async def test_approval_resume_emits_tool_call_result() -> None:
+    """After approving a tool call, the resume stream should contain a TOOL_CALL_RESULT event.
+
+    The message format follows the AG-UI approval pattern:
+    - assistant message with tool_calls
+    - tool message with {"accepted": true} content and toolCallId
+    """
+    tool_name = "get_weather"
+    call_id = "call_abc123"
+    weather_tool = _make_weather_tool()
+
+    agent = StubAgent(
+        updates=[AgentResponseUpdate(contents=[Content.from_text(text="The weather is sunny.")], role="assistant")],
+        default_options={"tools": [weather_tool]},
+    )
+    config = AgentConfig()
+
+    # Build resume messages: user query, assistant tool call, approval response
+    resume_messages: list[dict[str, Any]] = [
+        {"role": "user", "content": "What's the weather in Seattle?"},
+        {
+            "role": "assistant",
+            "content": "",
+            "tool_calls": [
+                {
+                    "id": call_id,
+                    "type": "function",
+                    "function": {
+                        "name": tool_name,
+                        "arguments": json.dumps({"city": "Seattle"}),
+                    },
+                }
+            ],
+        },
+        {
+            "role": "tool",
+            "content": json.dumps({"accepted": True}),
+            "toolCallId": call_id,
+        },
+    ]
+
+    input_data: dict[str, Any] = {
+        "thread_id": "thread-approval-result",
+        "run_id": "run-resume",
+        "messages": resume_messages,
+    }
+
+    events: list[Any] = []
+    async for event in run_agent_stream(input_data, agent, config):
+        events.append(event)
+
+    event_types = [getattr(e, "type", None) for e in events]
+
+    assert "RUN_STARTED" in event_types, f"Expected RUN_STARTED, got types: {event_types}"
+    assert "RUN_FINISHED" in event_types, f"Expected RUN_FINISHED, got types: {event_types}"
+
+    # TOOL_CALL_RESULT must be present for the approved tool
+    tool_result_events = [e for e in events if getattr(e, "type", None) == "TOOL_CALL_RESULT"]
+
+    assert len(tool_result_events) > 0, (
+        f"Expected at least one TOOL_CALL_RESULT event for the approved tool, "
+        f"but found none. Event types in stream: {event_types}"
+    )
+
+    result_event = tool_result_events[0]
+    assert result_event.tool_call_id == call_id, (
+        f"Expected TOOL_CALL_RESULT with tool_call_id={call_id}, got tool_call_id={result_event.tool_call_id}"
+    )
+    # Verify the result contains the actual tool execution output
+    assert result_event.content == "Sunny in Seattle"
+
+
+async def test_approval_resume_result_has_content() -> None:
+    """TOOL_CALL_RESULT event from an approved tool should contain the execution result."""
+    tool_name = "get_weather"
+    call_id = "call_content_check"
+    weather_tool = _make_weather_tool()
+
+    agent = StubAgent(
+        updates=[AgentResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")],
+        default_options={"tools": [weather_tool]},
+    )
+    config = AgentConfig()
+
+    resume_messages: list[dict[str, Any]] = [
+        {"role": "user", "content": "Check the weather"},
+        {
+            "role": "assistant",
+            "content": "",
+            "tool_calls": [
+                {
+                    "id": call_id,
+                    "type": "function",
+                    "function": {
+                        "name": tool_name,
+                        "arguments": json.dumps({"city": "Portland"}),
+                    },
+                }
+            ],
+        },
+        {
+            "role": "tool",
+            "content": json.dumps({"accepted": True}),
+            "toolCallId": call_id,
+        },
+    ]
+
+    input_data: dict[str, Any] = {
+        "thread_id": "thread-result-content",
+        "run_id": "run-resume-2",
+        "messages": resume_messages,
+    }
+
+    events: list[Any] = []
+    async for event in run_agent_stream(input_data, agent, config):
+        events.append(event)
+
+    tool_result_events = [e for e in events if getattr(e, "type", None) == "TOOL_CALL_RESULT"]
+    assert len(tool_result_events) == 1
+
+    result_event = tool_result_events[0]
+    assert result_event.tool_call_id == call_id
+    assert result_event.role == "tool"
+    # Verify the result contains the actual tool execution output (string returned directly)
+    assert result_event.content == "Sunny in Portland"
+
+
+async def test_approval_resume_snapshot_replaces_approval_payload_with_tool_result() -> None:
+    """Approved HITL tools persist their executed result in MESSAGES_SNAPSHOT for replay."""
+    from agent_framework_ag_ui._message_adapters import normalize_agui_input_messages
+
+    call_id = "call_snapshot_replay"
+    weather_tool = _make_weather_tool()
+    agent = StubAgent(
+        updates=[AgentResponseUpdate(contents=[Content.from_text(text="The weather is sunny.")], role="assistant")],
+        default_options={"tools": [weather_tool]},
+    )
+    config = AgentConfig()
+    resume_messages: list[dict[str, Any]] = [
+        {"role": "user", "content": "What's the weather in Seattle?"},
+        {
+            "role": "assistant",
+            "content": "",
+            "tool_calls": [
+                {
+                    "id": call_id,
+                    "type": "function",
+                    "function": {
+                        "name": "get_weather",
+                        "arguments": json.dumps({"city": "Seattle"}),
+                    },
+                }
+            ],
+        },
+        {
+            "role": "tool",
+            "content": json.dumps({"accepted": True}),
+            "toolCallId": call_id,
+        },
+    ]
+
+    events: list[Any] = []
+    async for event in run_agent_stream(
+        {
+            "thread_id": "thread-snapshot-replay",
+            "run_id": "run-snapshot-replay",
+            "messages": resume_messages,
+        },
+        agent,
+        config,
+    ):
+        events.append(event)
+
+    snapshots = [event.messages for event in events if getattr(event, "type", None) == "MESSAGES_SNAPSHOT"]
+    assert snapshots
+    snapshot_messages = [
+        message.model_dump(by_alias=True, exclude_none=True) if hasattr(message, "model_dump") else message
+        for message in snapshots[-1]
+    ]
+    tool_messages = [message for message in snapshot_messages if message.get("role") == "tool"]
+    assert any(
+        message.get("toolCallId") == call_id and message.get("content") == "Sunny in Seattle"
+        for message in tool_messages
+    )
+    assert not any(message.get("content") == json.dumps({"accepted": True}) for message in tool_messages)
+
+    replay_messages = snapshot_messages + [{"role": "user", "content": "What is the weather now?"}]
+    provider_messages, _ = normalize_agui_input_messages(replay_messages)
+
+    assert not any(
+        content.type == "function_approval_response"
+        for message in provider_messages
+        for content in message.contents or []
+    )
+    assert any(
+        content.type == "function_result" and content.call_id == call_id and content.result == "Sunny in Seattle"
+        for message in provider_messages
+        for content in message.contents or []
+    )
+
+
+async def test_no_approval_no_extra_tool_result() -> None:
+    """When no approval response is present, no extra TOOL_CALL_RESULT events should be emitted."""
+    agent = StubAgent(updates=[AgentResponseUpdate(contents=[Content.from_text(text="Hello.")], role="assistant")])
+    config = AgentConfig()
+
+    input_data: dict[str, Any] = {
+        "thread_id": "thread-no-approval",
+        "run_id": "run-normal",
+        "messages": [{"role": "user", "content": "Hi"}],
+    }
+
+    events: list[Any] = []
+    async for event in run_agent_stream(input_data, agent, config):
+        events.append(event)
+
+    tool_result_events = [e for e in events if getattr(e, "type", None) == "TOOL_CALL_RESULT"]
+    assert len(tool_result_events) == 0, f"Unexpected TOOL_CALL_RESULT events: {tool_result_events}"
+
+
+async def test_rejection_does_not_emit_tool_call_result() -> None:
+    """Rejected tool calls should not produce TOOL_CALL_RESULT events."""
+    tool_name = "get_weather"
+    call_id = "call_rejected"
+    weather_tool = _make_weather_tool()
+
+    agent = StubAgent(
+        updates=[AgentResponseUpdate(contents=[Content.from_text(text="OK, I won't check.")], role="assistant")],
+        default_options={"tools": [weather_tool]},
+    )
+    config = AgentConfig()
+
+    resume_messages: list[dict[str, Any]] = [
+        {"role": "user", "content": "What's the weather?"},
+        {
+            "role": "assistant",
+            "content": "",
+            "tool_calls": [
+                {
+                    "id": call_id,
+                    "type": "function",
+                    "function": {
+                        "name": tool_name,
+                        "arguments": json.dumps({"city": "Denver"}),
+                    },
+                }
+            ],
+        },
+        {
+            "role": "tool",
+            "content": json.dumps({"accepted": False}),
+            "toolCallId": call_id,
+        },
+    ]
+
+    input_data: dict[str, Any] = {
+        "thread_id": "thread-rejection",
+        "run_id": "run-rejected",
+        "messages": resume_messages,
+    }
+
+    events: list[Any] = []
+    async for event in run_agent_stream(input_data, agent, config):
+        events.append(event)
+
+    tool_result_events = [e for e in events if getattr(e, "type", None) == "TOOL_CALL_RESULT"]
+    assert len(tool_result_events) == 0, (
+        f"Expected no TOOL_CALL_RESULT for rejected tool, got {len(tool_result_events)}"
+    )
+
+
+def _make_temperature_tool() -> FunctionTool:
+    """Create a real executable temperature tool with approval_mode='always_require'."""
+
+    def get_temperature(city: str) -> str:
+        return f"72F in {city}"
+
+    return FunctionTool(
+        name="get_temperature",
+        description="Get the temperature for a city",
+        func=get_temperature,
+        approval_mode="always_require",
+    )
+
+
+async def test_mixed_approve_reject_emits_only_approved_tool_result() -> None:
+    """When one tool call is approved and another rejected, only the approved one produces a TOOL_CALL_RESULT event."""
+    weather_tool = _make_weather_tool()
+    temperature_tool = _make_temperature_tool()
+    approved_call_id = "call_approved"
+    rejected_call_id = "call_rejected"
+
+    agent = StubAgent(
+        updates=[AgentResponseUpdate(contents=[Content.from_text(text="Here are the results.")], role="assistant")],
+        default_options={"tools": [weather_tool, temperature_tool]},
+    )
+    config = AgentConfig()
+
+    resume_messages: list[dict[str, Any]] = [
+        {"role": "user", "content": "Weather and temperature in Seattle?"},
+        {
+            "role": "assistant",
+            "content": "",
+            "tool_calls": [
+                {
+                    "id": approved_call_id,
+                    "type": "function",
+                    "function": {
+                        "name": "get_weather",
+                        "arguments": json.dumps({"city": "Seattle"}),
+                    },
+                },
+                {
+                    "id": rejected_call_id,
+                    "type": "function",
+                    "function": {
+                        "name": "get_temperature",
+                        "arguments": json.dumps({"city": "Seattle"}),
+                    },
+                },
+            ],
+        },
+        {
+            "role": "tool",
+            "content": json.dumps({"accepted": True}),
+            "toolCallId": approved_call_id,
+        },
+        {
+            "role": "tool",
+            "content": json.dumps({"accepted": False}),
+            "toolCallId": rejected_call_id,
+        },
+    ]
+
+    input_data: dict[str, Any] = {
+        "thread_id": "thread-mixed",
+        "run_id": "run-mixed",
+        "messages": resume_messages,
+    }
+
+    events: list[Any] = []
+    async for event in run_agent_stream(input_data, agent, config):
+        events.append(event)
+
+    tool_result_events = [e for e in events if getattr(e, "type", None) == "TOOL_CALL_RESULT"]
+
+    # Only the approved tool call should produce a TOOL_CALL_RESULT event
+    assert len(tool_result_events) == 1, (
+        f"Expected exactly 1 TOOL_CALL_RESULT (approved only), got {len(tool_result_events)}"
+    )
+    assert tool_result_events[0].tool_call_id == approved_call_id
+    assert tool_result_events[0].content == "Sunny in Seattle"
+
+
+async def test_approval_resume_zero_updates_emits_tool_result() -> None:
+    """When the agent produces zero updates, TOOL_CALL_RESULT events should still be emitted via the fallback path."""
+    tool_name = "get_weather"
+    call_id = "call_zero_updates"
+    weather_tool = _make_weather_tool()
+
+    agent = StubAgent(
+        updates=[],
+        default_options={"tools": [weather_tool]},
+    )
+    config = AgentConfig()
+
+    resume_messages: list[dict[str, Any]] = [
+        {"role": "user", "content": "What's the weather?"},
+        {
+            "role": "assistant",
+            "content": "",
+            "tool_calls": [
+                {
+                    "id": call_id,
+                    "type": "function",
+                    "function": {
+                        "name": tool_name,
+                        "arguments": json.dumps({"city": "Boston"}),
+                    },
+                }
+            ],
+        },
+        {
+            "role": "tool",
+            "content": json.dumps({"accepted": True}),
+            "toolCallId": call_id,
+        },
+    ]
+
+    input_data: dict[str, Any] = {
+        "thread_id": "thread-zero-updates",
+        "run_id": "run-zero-updates",
+        "messages": resume_messages,
+    }
+
+    events: list[Any] = []
+    async for event in run_agent_stream(input_data, agent, config):
+        events.append(event)
+
+    event_types = [getattr(e, "type", None) for e in events]
+    assert "RUN_STARTED" in event_types
+
+    tool_result_events = [e for e in events if getattr(e, "type", None) == "TOOL_CALL_RESULT"]
+    assert len(tool_result_events) == 1, (
+        f"Expected 1 TOOL_CALL_RESULT in zero-updates fallback path, got {len(tool_result_events)}"
+    )
+    assert tool_result_events[0].tool_call_id == call_id
+    assert tool_result_events[0].content == "Sunny in Boston"
+
+
+async def test_resolve_approval_responses_returns_only_approved() -> None:
+    """_resolve_approval_responses should return only approved results; rejection results go into messages only."""
+    from agent_framework import Message
+
+    from agent_framework_ag_ui._agent_run import _resolve_approval_responses
+
+    weather_tool = _make_weather_tool()
+    temperature_tool = _make_temperature_tool()
+    approved_call_id = "call_a"
+    rejected_call_id = "call_r"
+
+    messages: list[Any] = [
+        Message(role="user", contents=[Content.from_text(text="Hi")]),
+        Message(
+            role="assistant",
+            contents=[
+                Content(
+                    type="function_approval_request",
+                    id=approved_call_id,
+                    function_call=Content(
+                        type="function_call",
+                        name="get_weather",
+                        call_id=approved_call_id,
+                        arguments='{"city": "NYC"}',
+                    ),
+                ),
+                Content(
+                    type="function_approval_request",
+                    id=rejected_call_id,
+                    function_call=Content(
+                        type="function_call",
+                        name="get_temperature",
+                        call_id=rejected_call_id,
+                        arguments='{"city": "NYC"}',
+                    ),
+                ),
+            ],
+        ),
+        Message(
+            role="user",
+            contents=[
+                Content(
+                    type="function_approval_response",
+                    id=approved_call_id,
+                    approved=True,
+                    function_call=Content(
+                        type="function_call",
+                        name="get_weather",
+                        call_id=approved_call_id,
+                        arguments='{"city": "NYC"}',
+                    ),
+                ),
+                Content(
+                    type="function_approval_response",
+                    id=rejected_call_id,
+                    approved=False,
+                    function_call=Content(
+                        type="function_call",
+                        name="get_temperature",
+                        call_id=rejected_call_id,
+                        arguments='{"city": "NYC"}',
+                    ),
+                ),
+            ],
+        ),
+    ]
+
+    agent = StubAgent(
+        updates=[],
+        default_options={"tools": [weather_tool, temperature_tool]},
+    )
+
+    results = await _resolve_approval_responses(messages, [weather_tool, temperature_tool], agent, {})
+
+    # Return value should only contain approved results
+    assert len(results) == 1
+    assert results[0].call_id == approved_call_id
+    assert results[0].type == "function_result"
+
+    # Rejection result should be written into messages (by _replace_approval_contents_with_results)
+    all_contents = [c for msg in messages for c in msg.contents]
+    rejection_results = [c for c in all_contents if c.type == "function_result" and c.call_id == rejected_call_id]
+    assert len(rejection_results) == 1
+    assert "rejected" in str(rejection_results[0].result).lower()
+
+
+class TestApprovalToolResultDisplayChannel:
+    """Approved tools using ``state_update(..., tool_result=...)`` must route the
+    display payload to the UI event while ``flow.tool_results`` still receives
+    the LLM-bound text. The HITL approval emitter is separate from the standard
+    streaming emitter, so it gets its own coverage.
+    """
+
+    def test_approval_emits_display_payload_when_marker_present(self) -> None:
+        from agent_framework_ag_ui import state_update
+        from agent_framework_ag_ui._agent_run import _make_approval_tool_result_events
+
+        display_payload = {"city": "Seattle", "temp": 14, "conditions": "foggy"}
+        inner = state_update(text="14°C, foggy", tool_result=display_payload)
+        resolved = Content.from_function_result(call_id="call_disp", result=[inner])
+
+        events = _make_approval_tool_result_events([resolved])
+
+        assert len(events) == 1
+        # UI event must carry the serialized display payload, NOT the LLM text.
+        assert json.loads(events[0].content) == display_payload
+        assert events[0].content != "14°C, foggy"
+
+    def test_approval_falls_back_to_text_when_no_marker(self) -> None:
+        """Backward compat: without a display marker, behaviour is unchanged."""
+        from agent_framework_ag_ui._agent_run import _make_approval_tool_result_events
+
+        resolved = Content.from_function_result(call_id="call_plain", result="Sunny in Seattle")
+
+        events = _make_approval_tool_result_events([resolved])
+
+        assert len(events) == 1
+        assert events[0].content == "Sunny in Seattle"

python/packages/ag-ui/tests/ag_ui/test_approval_state.py

@@ -0,0 +1,43 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for server-side AG-UI approval state storage."""
+
+import pytest
+
+from agent_framework_ag_ui._approval_state import InMemoryAGUIApprovalStateStore, approval_state_thread_id
+
+
+def test_approval_state_thread_id_allows_unscoped_thread() -> None:
+    assert approval_state_thread_id(scope=None, thread_id="thread-1") == "thread-1"
+
+
+def test_approval_state_thread_id_scopes_thread() -> None:
+    scoped_thread_id = approval_state_thread_id(scope="tenant-a", thread_id="thread-1")
+
+    assert scoped_thread_id != "thread-1"
+    assert "tenant-a" in scoped_thread_id
+    assert "thread-1" in scoped_thread_id
+
+
+@pytest.mark.parametrize("scope", ["", object()])
+def test_approval_state_thread_id_rejects_invalid_scope(scope: object) -> None:
+    with pytest.raises(ValueError, match="scope must be a non-empty string"):
+        approval_state_thread_id(scope=scope, thread_id="thread-1")
+
+
+def test_approval_state_store_rejects_invalid_max_entries() -> None:
+    with pytest.raises(ValueError, match="max_entries must be greater than 0"):
+        InMemoryAGUIApprovalStateStore(max_entries=0)
+
+
+def test_approval_state_store_evicts_oldest_entries() -> None:
+    store = InMemoryAGUIApprovalStateStore(max_entries=1)
+    store.pending_approvals[("thread-1", "call-1")] = "first"
+    store.pending_approvals[("thread-2", "call-2")] = "second"
+    store.tool_approval_states["thread-1"] = {"call_id": "call-1"}
+    store.tool_approval_states["thread-2"] = {"call_id": "call-2"}
+
+    store.evict_oldest()
+
+    assert list(store.pending_approvals.items()) == [(("thread-2", "call-2"), "second")]
+    assert list(store.tool_approval_states.items()) == [("thread-2", {"call_id": "call-2"})]

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

@@ -0,0 +1,469 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for AG-UI event converter."""
+
+import logging
+from typing import Any, cast
+
+import pytest
+from agent_framework import ChatResponse
+
+from agent_framework_ag_ui._event_converters import AGUIEventConverter
+
+
+class TestAGUIEventConverter:
+    """Test suite for AGUIEventConverter."""
+
+    def test_run_started_event(self) -> None:
+        """Test conversion of RUN_STARTED event."""
+        converter = AGUIEventConverter()
+        event = {
+            "type": "RUN_STARTED",
+            "threadId": "thread_123",
+            "runId": "run_456",
+        }
+
+        update = converter.convert_event(event)
+
+        assert update is not None
+        assert update.additional_properties is not None
+        assert update.additional_properties is not None
+        assert update.role == "assistant"
+        assert update.additional_properties["thread_id"] == "thread_123"
+        assert update.additional_properties["run_id"] == "run_456"
+        assert converter.thread_id == "thread_123"
+        assert converter.run_id == "run_456"
+
+    def test_text_message_start_event(self) -> None:
+        """Test conversion of TEXT_MESSAGE_START event."""
+        converter = AGUIEventConverter()
+        event = {
+            "type": "TEXT_MESSAGE_START",
+            "messageId": "msg_789",
+        }
+
+        update = converter.convert_event(event)
+
+        assert update is not None
+        assert update.role == "assistant"
+        assert update.message_id == "msg_789"
+        assert converter.current_message_id == "msg_789"
+
+    def test_text_message_content_event(self) -> None:
+        """Test conversion of TEXT_MESSAGE_CONTENT event."""
+        converter = AGUIEventConverter()
+        event = {
+            "type": "TEXT_MESSAGE_CONTENT",
+            "messageId": "msg_1",
+            "delta": "Hello",
+        }
+
+        update = converter.convert_event(event)
+
+        assert update is not None
+        assert update.role == "assistant"
+        assert update.message_id == "msg_1"
+        assert len(update.contents) == 1
+        assert update.contents[0].text == "Hello"
+
+    def test_text_message_streaming(self) -> None:
+        """Test streaming text across multiple TEXT_MESSAGE_CONTENT events."""
+        converter = AGUIEventConverter()
+        events = [
+            {"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": "Hello"},
+            {"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": " world"},
+            {"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": "!"},
+        ]
+
+        updates = cast(list[Any], [converter.convert_event(event) for event in events])
+
+        assert all(update is not None for update in updates)
+        assert all(update.message_id == "msg_1" for update in updates)
+        assert updates[0].contents[0].text == "Hello"
+        assert updates[1].contents[0].text == " world"
+        assert updates[2].contents[0].text == "!"
+
+    def test_text_message_end_event(self) -> None:
+        """Test conversion of TEXT_MESSAGE_END event."""
+        converter = AGUIEventConverter()
+        event = {
+            "type": "TEXT_MESSAGE_END",
+            "messageId": "msg_1",
+        }
+
+        update = converter.convert_event(event)
+
+        assert update is None
+
+    def test_tool_call_start_event(self) -> None:
+        """Test conversion of TOOL_CALL_START event."""
+        converter = AGUIEventConverter()
+        event = {
+            "type": "TOOL_CALL_START",
+            "toolCallId": "call_123",
+            "toolName": "get_weather",
+        }
+
+        update = converter.convert_event(event)
+
+        assert update is not None
+        assert update.role == "assistant"
+        assert len(update.contents) == 1
+        assert update.contents[0].call_id == "call_123"
+        assert update.contents[0].name == "get_weather"
+        assert update.contents[0].arguments == ""
+        assert converter.current_tool_call_id == "call_123"
+        assert converter.current_tool_name == "get_weather"
+
+    def test_tool_call_start_with_tool_call_name(self) -> None:
+        """Ensure TOOL_CALL_START with toolCallName still sets the tool name."""
+        converter = AGUIEventConverter()
+        event = {
+            "type": "TOOL_CALL_START",
+            "toolCallId": "call_abc",
+            "toolCallName": "get_weather",
+        }
+
+        update = converter.convert_event(event)
+
+        assert update is not None
+        assert update.contents[0].name == "get_weather"
+        assert converter.current_tool_name == "get_weather"
+
+    def test_tool_call_start_with_tool_call_name_snake_case(self) -> None:
+        """Support tool_call_name snake_case field for backwards compatibility."""
+        converter = AGUIEventConverter()
+        event = {
+            "type": "TOOL_CALL_START",
+            "toolCallId": "call_snake",
+            "tool_call_name": "get_weather",
+        }
+
+        update = converter.convert_event(event)
+
+        assert update is not None
+        assert update.contents[0].name == "get_weather"
+        assert converter.current_tool_name == "get_weather"
+
+    def test_tool_call_args_streaming(self) -> None:
+        """Test streaming tool arguments across multiple TOOL_CALL_ARGS events."""
+        converter = AGUIEventConverter()
+        converter.current_tool_call_id = "call_123"
+        converter.current_tool_name = "search"
+
+        events = [
+            {"type": "TOOL_CALL_ARGS", "delta": '{"query": "'},
+            {"type": "TOOL_CALL_ARGS", "delta": 'latest news"}'},
+        ]
+
+        updates = cast(list[Any], [converter.convert_event(event) for event in events])
+
+        assert all(update is not None for update in updates)
+        assert updates[0].contents[0].arguments == '{"query": "'
+        assert updates[1].contents[0].arguments == 'latest news"}'
+        assert converter.accumulated_tool_args == '{"query": "latest news"}'
+
+    def test_tool_call_end_event(self) -> None:
+        """Test conversion of TOOL_CALL_END event."""
+        converter = AGUIEventConverter()
+        converter.accumulated_tool_args = '{"location": "Seattle"}'
+
+        event = {
+            "type": "TOOL_CALL_END",
+            "toolCallId": "call_123",
+        }
+
+        update = converter.convert_event(event)
+
+        assert update is None
+        assert converter.accumulated_tool_args == ""
+
+    def test_tool_call_result_event(self) -> None:
+        """Test conversion of TOOL_CALL_RESULT event."""
+        converter = AGUIEventConverter()
+        event = {
+            "type": "TOOL_CALL_RESULT",
+            "toolCallId": "call_123",
+            "result": {"temperature": 22, "condition": "sunny"},
+        }
+
+        update = converter.convert_event(event)
+
+        assert update is not None
+        assert update.role == "tool"
+        assert len(update.contents) == 1
+        assert update.contents[0].call_id == "call_123"
+        assert update.contents[0].result == '{"temperature": 22, "condition": "sunny"}'
+
+    def test_run_finished_event(self) -> None:
+        """Test conversion of RUN_FINISHED event."""
+        converter = AGUIEventConverter()
+        converter.thread_id = "thread_123"
+        converter.run_id = "run_456"
+
+        event = {
+            "type": "RUN_FINISHED",
+            "threadId": "thread_123",
+            "runId": "run_456",
+        }
+
+        update = converter.convert_event(event)
+
+        assert update is not None
+        assert update.role == "assistant"
+        assert update.finish_reason == "stop"
+        assert update.additional_properties["thread_id"] == "thread_123"  # type: ignore[index]  # pyrefly: ignore[unsupported-operation]  # ty: ignore[not-subscriptable]
+        assert update.additional_properties["run_id"] == "run_456"  # type: ignore[index]  # pyrefly: ignore[unsupported-operation]  # ty: ignore[not-subscriptable]
+
+    def test_run_finished_event_with_interrupt(self) -> None:
+        """RUN_FINISHED interrupt metadata is preserved in additional_properties."""
+        converter = AGUIEventConverter()
+        converter.thread_id = "thread_123"
+        converter.run_id = "run_456"
+
+        event = {
+            "type": "RUN_FINISHED",
+            "threadId": "thread_123",
+            "runId": "run_456",
+            "interrupt": [{"id": "req_1", "value": {"question": "Continue?"}}],
+            "result": {"status": "paused"},
+        }
+
+        update = converter.convert_event(event)
+
+        assert update is not None
+        assert update.additional_properties is not None
+        assert update.additional_properties["interrupt"] == [{"id": "req_1", "value": {"question": "Continue?"}}]
+        assert update.additional_properties["result"] == {"status": "paused"}
+
+    def test_run_finished_event_with_canonical_interrupt_outcome(self) -> None:
+        """RUN_FINISHED outcome.interrupts metadata is preserved in additional_properties."""
+        converter = AGUIEventConverter()
+        converter.thread_id = "thread_123"
+        converter.run_id = "run_456"
+
+        outcome = {
+            "type": "interrupt",
+            "interrupts": [
+                {
+                    "id": "req_1",
+                    "reason": "input_required",
+                    "message": "Choose a value",
+                    "responseSchema": {"type": "string"},
+                    "metadata": {"agent_framework": {"request_type": "str"}},
+                }
+            ],
+        }
+        event = {
+            "type": "RUN_FINISHED",
+            "threadId": "thread_123",
+            "runId": "run_456",
+            "outcome": outcome,
+        }
+
+        update = converter.convert_event(event)
+
+        assert update is not None
+        assert update.additional_properties is not None
+        assert update.additional_properties["outcome"] == outcome
+        assert update.additional_properties["interrupts"] == outcome["interrupts"]
+
+    def test_run_finished_event_with_success_outcome_preserves_normal_completion(self) -> None:
+        """Non-interrupt RUN_FINISHED outcome metadata stays a normal stop update."""
+        converter = AGUIEventConverter()
+        converter.thread_id = "thread_123"
+        converter.run_id = "run_456"
+
+        event = {
+            "type": "RUN_FINISHED",
+            "threadId": "thread_123",
+            "runId": "run_456",
+            "outcome": {"type": "success"},
+        }
+
+        update = converter.convert_event(event)
+
+        assert update is not None
+        assert update.finish_reason == "stop"
+        assert update.additional_properties is not None
+        assert update.additional_properties["outcome"] == {"type": "success"}
+        assert "interrupts" not in update.additional_properties
+
+    def test_run_finished_event_with_non_dict_outcome_preserves_and_warns(
+        self, caplog: pytest.LogCaptureFixture
+    ) -> None:
+        """Malformed non-object outcome metadata is preserved but logged."""
+        converter = AGUIEventConverter()
+        converter.thread_id = "thread_123"
+        converter.run_id = "run_456"
+
+        event = {
+            "type": "RUN_FINISHED",
+            "threadId": "thread_123",
+            "runId": "run_456",
+            "outcome": "malformed",
+        }
+
+        with caplog.at_level(logging.WARNING, logger="agent_framework_ag_ui._event_converters"):
+            update = converter.convert_event(event)
+
+        assert update is not None
+        assert update.additional_properties is not None
+        assert update.additional_properties["outcome"] == "malformed"
+        assert "interrupts" not in update.additional_properties
+        assert "RUN_FINISHED outcome should be an object" in caplog.text
+
+    def test_run_error_event(self) -> None:
+        """Test conversion of RUN_ERROR event."""
+        converter = AGUIEventConverter()
+        converter.thread_id = "thread_123"
+        converter.run_id = "run_456"
+
+        event = {
+            "type": "RUN_ERROR",
+            "message": "Connection timeout",
+        }
+
+        update = converter.convert_event(event)
+
+        assert update is not None
+        assert update.role == "assistant"
+        assert update.finish_reason == "content_filter"
+        assert len(update.contents) == 1
+        assert update.contents[0].message == "Connection timeout"
+        assert update.contents[0].error_code == "RUN_ERROR"
+
+    def test_unknown_event_type(self) -> None:
+        """Test handling of unknown event types."""
+        converter = AGUIEventConverter()
+        event = {
+            "type": "UNKNOWN_EVENT",
+            "data": "some data",
+        }
+
+        update = converter.convert_event(event)
+
+        assert update is None
+
+    def test_custom_event_conversion(self) -> None:
+        """CUSTOM events are converted to update metadata."""
+        converter = AGUIEventConverter()
+        event = {
+            "type": "CUSTOM",
+            "name": "progress",
+            "value": {"percent": 10},
+        }
+
+        update = converter.convert_event(event)
+
+        assert update is not None
+        assert update.additional_properties is not None
+        assert update.additional_properties["ag_ui_custom_event"]["name"] == "progress"
+        assert update.additional_properties["ag_ui_custom_event"]["value"] == {"percent": 10}
+        assert update.additional_properties["ag_ui_custom_event"]["raw_type"] == "CUSTOM"
+
+    def test_custom_event_alias_conversion(self) -> None:
+        """CUSTOM_EVENT/custom_event aliases map to CUSTOM behavior."""
+        converter = AGUIEventConverter()
+        events = [
+            {"type": "CUSTOM_EVENT", "name": "alias_upper", "value": {"v": 1}},
+            {"type": "custom_event", "name": "alias_lower", "value": {"v": 2}},
+        ]
+
+        updates = cast(list[Any], [converter.convert_event(event) for event in events])
+
+        assert updates[0] is not None
+        assert updates[1] is not None
+        assert updates[0].additional_properties["ag_ui_custom_event"]["raw_type"] == "CUSTOM_EVENT"
+        assert updates[1].additional_properties["ag_ui_custom_event"]["raw_type"] == "custom_event"
+
+    def test_full_conversation_flow(self) -> None:
+        """Test complete conversation flow with multiple event types."""
+        converter = AGUIEventConverter()
+
+        events = [
+            {"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
+            {"type": "TEXT_MESSAGE_START", "messageId": "msg_1"},
+            {"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": "I'll check"},
+            {"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": " the weather."},
+            {"type": "TEXT_MESSAGE_END", "messageId": "msg_1"},
+            {"type": "TOOL_CALL_START", "toolCallId": "call_1", "toolName": "get_weather"},
+            {"type": "TOOL_CALL_ARGS", "delta": '{"location": "Seattle"}'},
+            {"type": "TOOL_CALL_END", "toolCallId": "call_1"},
+            {"type": "TOOL_CALL_RESULT", "toolCallId": "call_1", "result": "Sunny, 72°F"},
+            {"type": "TEXT_MESSAGE_START", "messageId": "msg_2"},
+            {"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_2", "delta": "It's sunny!"},
+            {"type": "TEXT_MESSAGE_END", "messageId": "msg_2"},
+            {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
+        ]
+
+        updates = cast(list[Any], [converter.convert_event(event) for event in events])
+        non_none_updates = [u for u in updates if u is not None]
+
+        assert len(non_none_updates) == 10
+        assert converter.thread_id == "thread_1"
+        assert converter.run_id == "run_1"
+
+    def test_multiple_tool_calls(self) -> None:
+        """Test handling multiple tool calls in sequence."""
+        converter = AGUIEventConverter()
+
+        events = [
+            {"type": "TOOL_CALL_START", "toolCallId": "call_1", "toolName": "search"},
+            {"type": "TOOL_CALL_ARGS", "delta": '{"query": "weather"}'},
+            {"type": "TOOL_CALL_END", "toolCallId": "call_1"},
+            {"type": "TOOL_CALL_START", "toolCallId": "call_2", "toolName": "fetch"},
+            {"type": "TOOL_CALL_ARGS", "delta": '{"url": "http://api.weather.com"}'},
+            {"type": "TOOL_CALL_END", "toolCallId": "call_2"},
+        ]
+
+        updates = cast(list[Any], [converter.convert_event(event) for event in events])
+        non_none_updates = [u for u in updates if u is not None]
+
+        assert len(non_none_updates) == 4
+        assert non_none_updates[0].contents[0].name == "search"
+        assert non_none_updates[2].contents[0].name == "fetch"
+
+    def test_tool_call_args_must_match_current_tool_call_id(self) -> None:
+        """TOOL_CALL_ARGS for another call must not be rebound to the current tool."""
+        converter = AGUIEventConverter()
+
+        events = [
+            {"type": "TOOL_CALL_START", "toolCallId": "safe", "toolName": "safe_tool"},
+            {"type": "TOOL_CALL_START", "toolCallId": "danger", "toolName": "danger_tool"},
+            {"type": "TOOL_CALL_ARGS", "toolCallId": "safe", "delta": '{"amount": 100}'},
+            {"type": "TOOL_CALL_END", "toolCallId": "safe"},
+            {"type": "TOOL_CALL_END", "toolCallId": "danger"},
+        ]
+
+        updates = [update for event in events if (update := converter.convert_event(event)) is not None]
+        response = ChatResponse.from_updates(updates)
+        function_calls = [
+            (content.call_id, content.name, content.arguments)
+            for message in response.messages
+            for content in message.contents
+            if content.type == "function_call"
+        ]
+
+        assert ("danger", "danger_tool", "") in function_calls
+        assert ("danger", "danger_tool", '{"amount": 100}') not in function_calls
+        assert all(call[2] != '{"amount": 100}' for call in function_calls)
+
+    def test_tool_call_end_must_match_current_tool_call_id(self) -> None:
+        """TOOL_CALL_END for another call must not clear the current call state."""
+        converter = AGUIEventConverter()
+
+        converter.convert_event({"type": "TOOL_CALL_START", "toolCallId": "danger", "toolName": "danger_tool"})
+        converter.convert_event({"type": "TOOL_CALL_ARGS", "toolCallId": "danger", "delta": '{"amount":'})
+        update = converter.convert_event({"type": "TOOL_CALL_END", "toolCallId": "safe"})
+
+        assert update is None
+        assert converter.current_tool_call_id == "danger"
+        assert converter.current_tool_name == "danger_tool"
+        assert converter.accumulated_tool_args == '{"amount":'
+
+        converter.convert_event({"type": "TOOL_CALL_END", "toolCallId": "danger"})
+
+        assert converter.current_tool_call_id is None
+        assert converter.current_tool_name is None
+        assert converter.accumulated_tool_args == ""

python/packages/ag-ui/tests/ag_ui/test_forwarded_props_in_metadata.py

@@ -0,0 +1,83 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for forwarded_props inclusion in AG-UI session metadata."""
+
+import json
+from typing import Any
+
+from agent_framework_ag_ui._agent_run import AG_UI_INTERNAL_METADATA_KEYS, _build_safe_metadata
+
+
+class TestForwardedPropsInSessionMetadata:
+    """Verify that forwarded_props is surfaced in session metadata and filtered from LLM metadata."""
+
+    def test_forwarded_props_in_internal_metadata_keys(self):
+        """forwarded_props is listed in AG_UI_INTERNAL_METADATA_KEYS to prevent LLM leakage."""
+        assert "forwarded_props" in AG_UI_INTERNAL_METADATA_KEYS
+
+    def test_forwarded_props_filtered_from_client_metadata(self):
+        """forwarded_props is filtered out when building LLM-bound client metadata."""
+        session_metadata: dict[str, Any] = {
+            "ag_ui_thread_id": "t1",
+            "ag_ui_run_id": "r1",
+            "forwarded_props": '{"custom_flag": true}',
+        }
+
+        client_metadata = {k: v for k, v in session_metadata.items() if k not in AG_UI_INTERNAL_METADATA_KEYS}
+
+        assert "forwarded_props" not in client_metadata
+        assert "ag_ui_thread_id" not in client_metadata
+
+
+class TestBuildSafeMetadata:
+    """Verify _build_safe_metadata handles various value types correctly."""
+
+    def test_string_value_unchanged(self):
+        result = _build_safe_metadata({"key": "hello"})
+        assert result == {"key": "hello"}
+
+    def test_dict_value_serialized_to_json(self):
+        result = _build_safe_metadata({"fp": {"flag": True, "source": "frontend"}})
+        assert "fp" in result
+        assert isinstance(result["fp"], str)
+        # Must be valid, decodable JSON
+        decoded = json.loads(result["fp"])
+        assert decoded == {"flag": True, "source": "frontend"}
+
+    def test_empty_dict_serialized_to_json(self):
+        result = _build_safe_metadata({"fp": {}})
+        assert result["fp"] == "{}"
+        assert json.loads(result["fp"]) == {}
+
+    def test_value_within_limit_kept(self):
+        value = "x" * 512
+        result = _build_safe_metadata({"key": value})
+        assert result["key"] == value
+
+    def test_value_exceeding_limit_dropped(self):
+        """Values exceeding 512 chars are dropped entirely (not truncated)."""
+        value = "x" * 513
+        result = _build_safe_metadata({"key": value})
+        assert "key" not in result
+
+    def test_json_value_exceeding_limit_dropped(self):
+        """JSON-serialized dict exceeding 512 chars is dropped, not truncated into invalid JSON."""
+        big_dict = {f"key_{i}": "v" * 100 for i in range(50)}
+        result = _build_safe_metadata({"forwarded_props": big_dict})
+        assert "forwarded_props" not in result
+
+    def test_other_keys_preserved_when_one_dropped(self):
+        """Dropping one oversized key does not affect other keys."""
+        result = _build_safe_metadata(
+            {
+                "small": "ok",
+                "big": "x" * 600,
+            }
+        )
+        assert result == {"small": "ok"}
+
+    def test_none_input_returns_empty(self):
+        assert _build_safe_metadata(None) == {}
+
+    def test_empty_input_returns_empty(self):
+        assert _build_safe_metadata({}) == {}

python/packages/ag-ui/tests/ag_ui/test_helpers.py

@@ -0,0 +1,504 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for orchestration helper functions."""
+
+from typing import Any
+
+from agent_framework import Content, Message
+
+from agent_framework_ag_ui._orchestration._helpers import (
+    approval_steps,
+    build_safe_metadata,
+    ensure_tool_call_entry,
+    is_state_context_message,
+    is_step_based_approval,
+    latest_approval_response,
+    pending_tool_call_ids,
+    schema_has_steps,
+    select_approval_tool_name,
+    tool_name_for_call_id,
+)
+
+
+class TestPendingToolCallIds:
+    """Tests for pending_tool_call_ids function."""
+
+    def test_empty_messages(self):
+        """Returns empty set for empty messages list."""
+        result = pending_tool_call_ids([])
+        assert result == set()
+
+    def test_no_tool_calls(self):
+        """Returns empty set when no tool calls in messages."""
+        messages = [
+            Message(role="user", contents=[Content.from_text("Hello")]),
+            Message(role="assistant", contents=[Content.from_text("Hi there")]),
+        ]
+        result = pending_tool_call_ids(messages)
+        assert result == set()
+
+    def test_pending_tool_call(self):
+        """Returns pending tool call ID when no result exists."""
+        messages = [
+            Message(
+                role="assistant",
+                contents=[Content.from_function_call(call_id="call_123", name="get_weather", arguments="{}")],
+            ),
+        ]
+        result = pending_tool_call_ids(messages)
+        assert result == {"call_123"}
+
+    def test_resolved_tool_call(self):
+        """Returns empty set when tool call has result."""
+        messages = [
+            Message(
+                role="assistant",
+                contents=[Content.from_function_call(call_id="call_123", name="get_weather", arguments="{}")],
+            ),
+            Message(
+                role="tool",
+                contents=[Content.from_function_result(call_id="call_123", result="sunny")],
+            ),
+        ]
+        result = pending_tool_call_ids(messages)
+        assert result == set()
+
+    def test_multiple_tool_calls_some_resolved(self):
+        """Returns only unresolved tool call IDs."""
+        messages = [
+            Message(
+                role="assistant",
+                contents=[
+                    Content.from_function_call(call_id="call_1", name="tool_a", arguments="{}"),
+                    Content.from_function_call(call_id="call_2", name="tool_b", arguments="{}"),
+                    Content.from_function_call(call_id="call_3", name="tool_c", arguments="{}"),
+                ],
+            ),
+            Message(
+                role="tool",
+                contents=[Content.from_function_result(call_id="call_1", result="result_a")],
+            ),
+            Message(
+                role="tool",
+                contents=[Content.from_function_result(call_id="call_3", result="result_c")],
+            ),
+        ]
+        result = pending_tool_call_ids(messages)
+        assert result == {"call_2"}
+
+
+class TestIsStateContextMessage:
+    """Tests for is_state_context_message function."""
+
+    def test_state_context_message(self):
+        """Returns True for state context message."""
+        message = Message(
+            role="system",
+            contents=[Content.from_text("Current state of the application: {}")],
+        )
+        assert is_state_context_message(message) is True
+
+    def test_non_system_message(self):
+        """Returns False for non-system message."""
+        message = Message(
+            role="user",
+            contents=[Content.from_text("Current state of the application: {}")],
+        )
+        assert is_state_context_message(message) is False
+
+    def test_system_message_without_state_prefix(self):
+        """Returns False for system message without state prefix."""
+        message = Message(
+            role="system",
+            contents=[Content.from_text("You are a helpful assistant.")],
+        )
+        assert is_state_context_message(message) is False
+
+    def test_empty_contents(self):
+        """Returns False for message with empty contents."""
+        message = Message(role="system", contents=[])
+        assert is_state_context_message(message) is False
+
+
+class TestEnsureToolCallEntry:
+    """Tests for ensure_tool_call_entry function."""
+
+    def test_creates_new_entry(self):
+        """Creates new entry when ID not found."""
+        tool_calls_by_id: dict = {}
+        pending_tool_calls: list = []
+
+        entry = ensure_tool_call_entry("call_123", tool_calls_by_id, pending_tool_calls)
+
+        assert entry["id"] == "call_123"
+        assert entry["type"] == "function"
+        assert entry["function"]["name"] == ""
+        assert entry["function"]["arguments"] == ""
+        assert "call_123" in tool_calls_by_id
+        assert len(pending_tool_calls) == 1
+
+    def test_returns_existing_entry(self):
+        """Returns existing entry when ID found."""
+        existing_entry: dict[str, Any] = {
+            "id": "call_123",
+            "type": "function",
+            "function": {"name": "get_weather", "arguments": '{"city": "NYC"}'},
+        }
+        tool_calls_by_id: dict[str, dict[str, Any]] = {"call_123": existing_entry}
+        pending_tool_calls: list[dict[str, Any]] = []
+
+        entry = ensure_tool_call_entry("call_123", tool_calls_by_id, pending_tool_calls)
+
+        assert entry is existing_entry
+        assert entry["function"]["name"] == "get_weather"
+        assert len(pending_tool_calls) == 0  # Not added again
+
+
+class TestToolNameForCallId:
+    """Tests for tool_name_for_call_id function."""
+
+    def test_returns_tool_name(self):
+        """Returns tool name for valid entry."""
+        tool_calls_by_id = {
+            "call_123": {
+                "id": "call_123",
+                "function": {"name": "get_weather", "arguments": "{}"},
+            }
+        }
+        result = tool_name_for_call_id(tool_calls_by_id, "call_123")
+        assert result == "get_weather"
+
+    def test_returns_none_for_missing_id(self):
+        """Returns None when ID not found."""
+        tool_calls_by_id: dict = {}
+        result = tool_name_for_call_id(tool_calls_by_id, "call_123")
+        assert result is None
+
+    def test_returns_none_for_missing_function(self):
+        """Returns None when function key missing."""
+        tool_calls_by_id = {"call_123": {"id": "call_123"}}
+        result = tool_name_for_call_id(tool_calls_by_id, "call_123")
+        assert result is None
+
+    def test_returns_none_for_non_dict_function(self):
+        """Returns None when function is not a dict."""
+        tool_calls_by_id = {"call_123": {"id": "call_123", "function": "not_a_dict"}}
+        result = tool_name_for_call_id(tool_calls_by_id, "call_123")
+        assert result is None
+
+    def test_returns_none_for_empty_name(self):
+        """Returns None when name is empty."""
+        tool_calls_by_id = {"call_123": {"id": "call_123", "function": {"name": "", "arguments": "{}"}}}
+        result = tool_name_for_call_id(tool_calls_by_id, "call_123")
+        assert result is None
+
+
+class TestSchemaHasSteps:
+    """Tests for schema_has_steps function."""
+
+    def test_schema_with_steps_array(self):
+        """Returns True when schema has steps array property."""
+        schema = {"properties": {"steps": {"type": "array"}}}
+        assert schema_has_steps(schema) is True
+
+    def test_schema_without_steps(self):
+        """Returns False when schema doesn't have steps."""
+        schema = {"properties": {"name": {"type": "string"}}}
+        assert schema_has_steps(schema) is False
+
+    def test_schema_with_non_array_steps(self):
+        """Returns False when steps is not array type."""
+        schema = {"properties": {"steps": {"type": "string"}}}
+        assert schema_has_steps(schema) is False
+
+    def test_non_dict_schema(self):
+        """Returns False for non-dict schema."""
+        assert schema_has_steps(None) is False
+        assert schema_has_steps("not a dict") is False
+        assert schema_has_steps([]) is False
+
+    def test_missing_properties(self):
+        """Returns False when properties key is missing."""
+        schema = {"type": "object"}
+        assert schema_has_steps(schema) is False
+
+    def test_non_dict_properties(self):
+        """Returns False when properties is not a dict."""
+        schema = {"properties": "not a dict"}
+        assert schema_has_steps(schema) is False
+
+    def test_non_dict_steps(self):
+        """Returns False when steps is not a dict."""
+        schema = {"properties": {"steps": "not a dict"}}
+        assert schema_has_steps(schema) is False
+
+
+class TestSelectApprovalToolName:
+    """Tests for select_approval_tool_name function."""
+
+    def test_none_client_tools(self):
+        """Returns None when client_tools is None."""
+        result = select_approval_tool_name(None)
+        assert result is None
+
+    def test_empty_client_tools(self):
+        """Returns None when client_tools is empty."""
+        result = select_approval_tool_name([])
+        assert result is None
+
+    def test_finds_approval_tool(self):
+        """Returns tool name when tool has steps schema."""
+
+        class MockTool:
+            name = "generate_task_steps"
+
+            def parameters(self):
+                return {"properties": {"steps": {"type": "array"}}}
+
+        result = select_approval_tool_name([MockTool()])
+        assert result == "generate_task_steps"
+
+    def test_skips_tool_without_name(self):
+        """Skips tools without name attribute."""
+
+        class MockToolNoName:
+            def parameters(self):
+                return {"properties": {"steps": {"type": "array"}}}
+
+        result = select_approval_tool_name([MockToolNoName()])
+        assert result is None
+
+    def test_skips_tool_without_parameters_method(self):
+        """Skips tools without callable parameters method."""
+
+        class MockToolNoParams:
+            name = "some_tool"
+            parameters = "not callable"
+
+        result = select_approval_tool_name([MockToolNoParams()])
+        assert result is None
+
+    def test_skips_tool_without_steps_schema(self):
+        """Skips tools that don't have steps in schema."""
+
+        class MockToolNoSteps:
+            name = "other_tool"
+
+            def parameters(self):
+                return {"properties": {"data": {"type": "string"}}}
+
+        result = select_approval_tool_name([MockToolNoSteps()])
+        assert result is None
+
+
+class TestBuildSafeMetadata:
+    """Tests for build_safe_metadata function."""
+
+    def test_none_metadata(self):
+        """Returns empty dict for None metadata."""
+        result = build_safe_metadata(None)
+        assert result == {}
+
+    def test_empty_metadata(self):
+        """Returns empty dict for empty metadata."""
+        result = build_safe_metadata({})
+        assert result == {}
+
+    def test_string_values_under_limit(self):
+        """Preserves string values under 512 chars."""
+        metadata = {"key1": "short value", "key2": "another value"}
+        result = build_safe_metadata(metadata)
+        assert result == metadata
+
+    def test_truncates_long_string_values(self):
+        """Truncates string values over 512 chars."""
+        long_value = "x" * 1000
+        metadata = {"key": long_value}
+        result = build_safe_metadata(metadata)
+        assert len(result["key"]) == 512
+        assert result["key"] == "x" * 512
+
+    def test_non_string_values_serialized(self):
+        """Serializes non-string values to JSON."""
+        metadata = {"count": 42, "items": ["a", "b"]}
+        result = build_safe_metadata(metadata)
+        assert result["count"] == "42"
+        assert result["items"] == '["a", "b"]'
+
+    def test_truncates_serialized_values(self):
+        """Truncates serialized JSON values over 512 chars."""
+        long_list = list(range(200))  # Will serialize to >512 chars
+        metadata = {"data": long_list}
+        result = build_safe_metadata(metadata)
+        assert len(result["data"]) == 512
+
+
+class TestLatestApprovalResponse:
+    """Tests for latest_approval_response function."""
+
+    def test_empty_messages(self):
+        """Returns None for empty messages."""
+        result = latest_approval_response([])
+        assert result is None
+
+    def test_no_approval_response(self):
+        """Returns None when no approval response in last message."""
+        messages = [
+            Message(role="assistant", contents=[Content.from_text("Hello")]),
+        ]
+        result = latest_approval_response(messages)
+        assert result is None
+
+    def test_finds_approval_response(self):
+        """Returns approval response from last message."""
+        # Create a function call content first
+        fc = Content.from_function_call(call_id="call_123", name="test_tool", arguments="{}")
+        approval_content = Content.from_function_approval_response(
+            approved=True,
+            id="approval_123",
+            function_call=fc,
+        )
+        messages = [
+            Message(role="user", contents=[approval_content]),
+        ]
+        result = latest_approval_response(messages)
+        assert result is approval_content
+
+
+class TestApprovalSteps:
+    """Tests for approval_steps function."""
+
+    def test_steps_from_ag_ui_state_args(self):
+        """Extracts steps from ag_ui_state_args."""
+        fc = Content.from_function_call(call_id="call_123", name="test_tool", arguments="{}")
+        approval = Content.from_function_approval_response(
+            approved=True,
+            id="approval_123",
+            function_call=fc,
+            additional_properties={"ag_ui_state_args": {"steps": [{"id": 1}, {"id": 2}]}},
+        )
+        result = approval_steps(approval)
+        assert result == [{"id": 1}, {"id": 2}]
+
+    def test_steps_from_function_call(self):
+        """Extracts steps from function call arguments."""
+        fc = Content.from_function_call(
+            call_id="call_123",
+            name="test",
+            arguments='{"steps": [{"step": 1}]}',
+        )
+        approval = Content.from_function_approval_response(
+            approved=True,
+            id="approval_123",
+            function_call=fc,
+        )
+        result = approval_steps(approval)
+        assert result == [{"step": 1}]
+
+    def test_empty_steps_when_no_state_args(self):
+        """Returns empty list when no ag_ui_state_args."""
+        fc = Content.from_function_call(call_id="call_123", name="test_tool", arguments="{}")
+        approval = Content.from_function_approval_response(
+            approved=True,
+            id="approval_123",
+            function_call=fc,
+        )
+        result = approval_steps(approval)
+        assert result == []
+
+    def test_empty_steps_when_state_args_not_dict(self):
+        """Returns empty list when ag_ui_state_args is not a dict."""
+        fc = Content.from_function_call(call_id="call_123", name="test_tool", arguments="{}")
+        approval = Content.from_function_approval_response(
+            approved=True,
+            id="approval_123",
+            function_call=fc,
+            additional_properties={"ag_ui_state_args": "not a dict"},
+        )
+        result = approval_steps(approval)
+        assert result == []
+
+    def test_empty_steps_when_steps_not_list(self):
+        """Returns empty list when steps is not a list."""
+        fc = Content.from_function_call(call_id="call_123", name="test_tool", arguments="{}")
+        approval = Content.from_function_approval_response(
+            approved=True,
+            id="approval_123",
+            function_call=fc,
+            additional_properties={"ag_ui_state_args": {"steps": "not a list"}},
+        )
+        result = approval_steps(approval)
+        assert result == []
+
+
+class TestIsStepBasedApproval:
+    """Tests for is_step_based_approval function."""
+
+    def test_returns_true_when_has_steps(self):
+        """Returns True when approval has steps."""
+        fc = Content.from_function_call(call_id="call_123", name="test_tool", arguments="{}")
+        approval = Content.from_function_approval_response(
+            approved=True,
+            id="approval_123",
+            function_call=fc,
+            additional_properties={"ag_ui_state_args": {"steps": [{"id": 1}]}},
+        )
+        result = is_step_based_approval(approval, None)
+        assert result is True
+
+    def test_returns_false_no_steps_no_function_call(self):
+        """Returns False when no steps and no function call."""
+        # Create content directly to have no function_call
+        approval = Content(
+            type="function_approval_response",
+            function_call=None,
+        )
+        result = is_step_based_approval(approval, None)
+        assert result is False
+
+    def test_returns_false_no_predict_config(self):
+        """Returns False when no predict_state_config."""
+        fc = Content.from_function_call(call_id="call_123", name="some_tool", arguments="{}")
+        approval = Content.from_function_approval_response(
+            approved=True,
+            id="approval_123",
+            function_call=fc,
+        )
+        result = is_step_based_approval(approval, None)
+        assert result is False
+
+    def test_returns_true_when_tool_matches_config(self):
+        """Returns True when tool matches predict_state_config with steps."""
+        fc = Content.from_function_call(call_id="call_123", name="generate_steps", arguments="{}")
+        approval = Content.from_function_approval_response(
+            approved=True,
+            id="approval_123",
+            function_call=fc,
+        )
+        config = {"steps": {"tool": "generate_steps", "tool_argument": "steps"}}
+        result = is_step_based_approval(approval, config)
+        assert result is True
+
+    def test_returns_false_when_tool_not_in_config(self):
+        """Returns False when tool not in predict_state_config."""
+        fc = Content.from_function_call(call_id="call_123", name="other_tool", arguments="{}")
+        approval = Content.from_function_approval_response(
+            approved=True,
+            id="approval_123",
+            function_call=fc,
+        )
+        config = {"steps": {"tool": "generate_steps", "tool_argument": "steps"}}
+        result = is_step_based_approval(approval, config)
+        assert result is False
+
+    def test_returns_false_when_tool_arg_not_steps(self):
+        """Returns False when tool_argument is not 'steps'."""
+        fc = Content.from_function_call(call_id="call_123", name="generate_steps", arguments="{}")
+        approval = Content.from_function_approval_response(
+            approved=True,
+            id="approval_123",
+            function_call=fc,
+        )
+        config = {"document": {"tool": "generate_steps", "tool_argument": "content"}}
+        result = is_step_based_approval(approval, config)
+        assert result is False

python/packages/ag-ui/tests/ag_ui/test_http_round_trip.py

@@ -0,0 +1,348 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""HTTP round-trip tests: POST → SSE bytes → parse → validate event sequence.
+
+These tests exercise the full HTTP pipeline using FastAPI TestClient,
+parsing the raw SSE byte stream and validating through EventStream assertions.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from agent_framework import AgentResponseUpdate, Content, WorkflowBuilder, WorkflowContext, executor
+from conftest import StubAgent  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+from sse_helpers import (  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+    parse_sse_response,
+    parse_sse_to_event_stream,
+)
+
+from agent_framework_ag_ui import AgentFrameworkAgent, AgentFrameworkWorkflow, add_agent_framework_fastapi_endpoint
+
+
+def _build_app_with_agent(updates: list[AgentResponseUpdate], **kwargs: Any) -> FastAPI:
+    stub = StubAgent(updates=updates)
+    agent = AgentFrameworkAgent(agent=stub, **kwargs)
+    app = FastAPI()
+    add_agent_framework_fastapi_endpoint(app, agent)
+    return app
+
+
+def _build_app_with_workflow(workflow_builder: WorkflowBuilder) -> FastAPI:
+    workflow = workflow_builder.build()
+    wrapper = AgentFrameworkWorkflow(workflow=workflow)
+    app = FastAPI()
+    add_agent_framework_fastapi_endpoint(app, wrapper)
+    return app
+
+
+USER_PAYLOAD: dict[str, Any] = {
+    "messages": [{"role": "user", "content": "Hello"}],
+    "threadId": "thread-http",
+    "runId": "run-http",
+}
+
+
+# ── Agentic chat SSE round-trip ──
+
+
+def test_agentic_chat_sse_round_trip() -> None:
+    """Full HTTP round-trip: POST → SSE bytes → parse → validate event sequence."""
+    app = _build_app_with_agent(
+        [
+            AgentResponseUpdate(contents=[Content.from_text(text="Hi there!")], role="assistant"),
+        ]
+    )
+    client = TestClient(app)
+    response = client.post("/", json=USER_PAYLOAD)
+
+    assert response.status_code == 200
+    assert "text/event-stream" in response.headers["content-type"]
+
+    stream = parse_sse_to_event_stream(response.content)
+    stream.assert_bookends()
+    stream.assert_text_messages_balanced()
+    stream.assert_no_run_error()
+    stream.assert_ordered_types(
+        [
+            "RUN_STARTED",
+            "TEXT_MESSAGE_START",
+            "TEXT_MESSAGE_CONTENT",
+            "TEXT_MESSAGE_END",
+            "MESSAGES_SNAPSHOT",
+            "RUN_FINISHED",
+        ]
+    )
+
+
+# ── Tool call SSE round-trip ──
+
+
+def test_tool_call_sse_round_trip() -> None:
+    """Tool call events survive SSE encoding/parsing round-trip."""
+    app = _build_app_with_agent(
+        [
+            AgentResponseUpdate(
+                contents=[Content.from_function_call(name="get_weather", call_id="call-1", arguments='{"city": "SF"}')],
+                role="assistant",
+            ),
+            AgentResponseUpdate(
+                contents=[Content.from_function_result(call_id="call-1", result="72°F")],
+                role="assistant",
+            ),
+            AgentResponseUpdate(
+                contents=[Content.from_text(text="It's warm!")],
+                role="assistant",
+            ),
+        ]
+    )
+    client = TestClient(app)
+    response = client.post("/", json=USER_PAYLOAD)
+
+    stream = parse_sse_to_event_stream(response.content)
+    stream.assert_bookends()
+    stream.assert_tool_calls_balanced()
+    stream.assert_text_messages_balanced()
+
+    # Verify tool call details survive SSE encoding
+    start = stream.first("TOOL_CALL_START")
+    assert start.tool_call_name == "get_weather"
+    assert start.tool_call_id == "call-1"
+
+
+# ── SSE encoding fidelity ──
+
+
+def test_sse_event_encoding_fidelity() -> None:
+    """Every event from agent.run() produces a valid SSE data: line that round-trips."""
+    app = _build_app_with_agent(
+        [
+            AgentResponseUpdate(contents=[Content.from_text(text="Hello world")], role="assistant"),
+        ]
+    )
+    client = TestClient(app)
+    response = client.post("/", json=USER_PAYLOAD)
+
+    raw_events = parse_sse_response(response.content)
+    assert len(raw_events) > 0, "No SSE events parsed"
+
+    # Every event should have a 'type' field
+    for event in raw_events:
+        assert "type" in event, f"Event missing 'type': {event}"
+
+    # Event types should include the expected ones
+    event_types = [e["type"] for e in raw_events]
+    assert "RUN_STARTED" in event_types
+    assert "RUN_FINISHED" in event_types
+
+
+# ── camelCase request field acceptance ──
+
+
+def test_camel_case_request_fields_accepted() -> None:
+    """Request with camelCase fields (runId, threadId) is correctly parsed."""
+    app = _build_app_with_agent(
+        [
+            AgentResponseUpdate(contents=[Content.from_text(text="ok")], role="assistant"),
+        ]
+    )
+    client = TestClient(app)
+    response = client.post(
+        "/",
+        json={
+            "messages": [{"role": "user", "content": "hi"}],
+            "runId": "camel-run",
+            "threadId": "camel-thread",
+        },
+    )
+    assert response.status_code == 200
+
+    stream = parse_sse_to_event_stream(response.content)
+    stream.assert_bookends()
+
+
+# ── Workflow SSE round-trip ──
+
+
+def test_workflow_sse_round_trip() -> None:
+    """Workflow events survive SSE encoding/parsing."""
+
+    @executor(id="greeter")
+    async def greeter(message: Any, ctx: WorkflowContext[Any, str]) -> None:
+        await ctx.yield_output("Hello from workflow!")
+
+    app = _build_app_with_workflow(WorkflowBuilder(start_executor=greeter))
+    client = TestClient(app)
+    response = client.post("/", json=USER_PAYLOAD)
+
+    assert response.status_code == 200
+    stream = parse_sse_to_event_stream(response.content)
+    stream.assert_bookends()
+    stream.assert_no_run_error()
+    stream.assert_text_messages_balanced()
+    stream.assert_has_type("STEP_STARTED")
+
+
+# ── Error handling ──
+
+
+def test_empty_messages_returns_valid_sse() -> None:
+    """Empty messages list still returns a valid SSE stream with bookends."""
+    app = _build_app_with_agent(
+        [
+            AgentResponseUpdate(contents=[Content.from_text(text="ok")], role="assistant"),
+        ]
+    )
+    client = TestClient(app)
+    response = client.post("/", json={"messages": []})
+
+    assert response.status_code == 200
+    stream = parse_sse_to_event_stream(response.content)
+    stream.assert_bookends()
+
+
+def test_sse_response_headers() -> None:
+    """SSE response has correct headers for event streaming."""
+    app = _build_app_with_agent(
+        [
+            AgentResponseUpdate(contents=[Content.from_text(text="ok")], role="assistant"),
+        ]
+    )
+    client = TestClient(app)
+    response = client.post("/", json=USER_PAYLOAD)
+
+    assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
+    assert response.headers.get("cache-control") == "no-cache"
+
+
+# ── MCP tool call SSE round-trip ──
+
+
+def test_mcp_tool_call_sse_round_trip() -> None:
+    """MCP tool call + result events survive SSE encoding/parsing round-trip."""
+    app = _build_app_with_agent(
+        [
+            AgentResponseUpdate(
+                contents=[
+                    Content.from_mcp_server_tool_call(
+                        call_id="mcp-1",
+                        tool_name="search",
+                        server_name="brave",
+                        arguments={"query": "weather"},
+                    )
+                ],
+                role="assistant",
+            ),
+            AgentResponseUpdate(
+                contents=[
+                    Content.from_mcp_server_tool_result(
+                        call_id="mcp-1",
+                        output={"results": ["sunny"]},
+                    )
+                ],
+                role="assistant",
+            ),
+            AgentResponseUpdate(
+                contents=[Content.from_text(text="It's sunny!")],
+                role="assistant",
+            ),
+        ]
+    )
+    client = TestClient(app)
+    response = client.post("/", json=USER_PAYLOAD)
+
+    assert response.status_code == 200
+    stream = parse_sse_to_event_stream(response.content)
+    stream.assert_bookends()
+    stream.assert_tool_calls_balanced()
+    stream.assert_text_messages_balanced()
+    stream.assert_no_run_error()
+
+    # Verify MCP tool call details survive SSE encoding
+    start = stream.first("TOOL_CALL_START")
+    assert start.tool_call_name == "search"
+    assert start.tool_call_id == "mcp-1"
+
+    # Verify the result came through
+    result = stream.first("TOOL_CALL_RESULT")
+    assert "sunny" in result.content
+
+
+# ── Text reasoning SSE round-trip ──
+
+
+def test_text_reasoning_sse_round_trip() -> None:
+    """Text reasoning events survive SSE encoding/parsing round-trip."""
+    app = _build_app_with_agent(
+        [
+            AgentResponseUpdate(
+                contents=[
+                    Content.from_text_reasoning(
+                        id="reason-1",
+                        text="The user wants weather info, I should use a tool.",
+                    )
+                ],
+                role="assistant",
+            ),
+            AgentResponseUpdate(
+                contents=[Content.from_text(text="Let me check the weather.")],
+                role="assistant",
+            ),
+        ]
+    )
+    client = TestClient(app)
+    response = client.post("/", json=USER_PAYLOAD)
+
+    assert response.status_code == 200
+    stream = parse_sse_to_event_stream(response.content)
+    stream.assert_bookends()
+    stream.assert_text_messages_balanced()
+    stream.assert_no_run_error()
+    stream.assert_has_type("REASONING_START")
+    stream.assert_has_type("REASONING_MESSAGE_CONTENT")
+    stream.assert_has_type("REASONING_END")
+
+    # Verify reasoning content survives SSE encoding
+    raw_events = parse_sse_response(response.content)
+    reasoning_content = [e for e in raw_events if e["type"] == "REASONING_MESSAGE_CONTENT"]
+    assert len(reasoning_content) == 1
+    assert "weather" in reasoning_content[0]["delta"]
+
+
+def test_text_reasoning_with_encrypted_value_sse_round_trip() -> None:
+    """Reasoning with protected_data emits ReasoningEncryptedValue through SSE."""
+    app = _build_app_with_agent(
+        [
+            AgentResponseUpdate(
+                contents=[
+                    Content.from_text_reasoning(
+                        id="reason-enc",
+                        text="visible reasoning",
+                        protected_data="encrypted-payload-abc123",
+                    )
+                ],
+                role="assistant",
+            ),
+            AgentResponseUpdate(
+                contents=[Content.from_text(text="Done.")],
+                role="assistant",
+            ),
+        ]
+    )
+    client = TestClient(app)
+    response = client.post("/", json=USER_PAYLOAD)
+
+    assert response.status_code == 200
+    stream = parse_sse_to_event_stream(response.content)
+    stream.assert_bookends()
+    stream.assert_no_run_error()
+    stream.assert_has_type("REASONING_ENCRYPTED_VALUE")
+
+    raw_events = parse_sse_response(response.content)
+    encrypted = [e for e in raw_events if e["type"] == "REASONING_ENCRYPTED_VALUE"]
+    assert len(encrypted) == 1
+    assert encrypted[0]["encryptedValue"] == "encrypted-payload-abc123"
+    assert encrypted[0]["entityId"] == "reason-enc"
+    assert encrypted[0]["subtype"] == "message"

python/packages/ag-ui/tests/ag_ui/test_http_service.py

@@ -0,0 +1,364 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for AGUIHttpService."""
+
+import json
+from unittest.mock import AsyncMock, Mock
+
+import httpx
+import pytest
+from ag_ui.core import Interrupt, ResumeEntry
+
+from agent_framework_ag_ui._http_service import AGUIHttpService
+
+
+@pytest.fixture
+def mock_http_client():
+    """Create a mock httpx.AsyncClient."""
+    client = AsyncMock(spec=httpx.AsyncClient)
+    return client
+
+
+@pytest.fixture
+def sample_events():
+    """Sample AG-UI events for testing."""
+    return [
+        {"type": "RUN_STARTED", "threadId": "thread_123", "runId": "run_456"},
+        {"type": "TEXT_MESSAGE_START", "messageId": "msg_1", "role": "assistant"},
+        {"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": "Hello"},
+        {"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": " world"},
+        {"type": "TEXT_MESSAGE_END", "messageId": "msg_1"},
+        {"type": "RUN_FINISHED", "threadId": "thread_123", "runId": "run_456"},
+    ]
+
+
+def create_sse_response(events: list[dict]) -> str:
+    """Create SSE formatted response from events."""
+    lines = []
+    for event in events:
+        lines.append(f"data: {json.dumps(event)}\n")
+    return "\n".join(lines)
+
+
+async def test_http_service_initialization():
+    """Test AGUIHttpService initialization."""
+    # Test with default client
+    service = AGUIHttpService("http://localhost:8888/")
+    assert service.endpoint == "http://localhost:8888"
+    assert service._owns_client is True
+    assert isinstance(service.http_client, httpx.AsyncClient)
+    await service.close()
+
+    # Test with custom client
+    custom_client = httpx.AsyncClient()
+    service = AGUIHttpService("http://localhost:8888/", http_client=custom_client)
+    assert service._owns_client is False
+    assert service.http_client is custom_client
+    # Shouldn't close the custom client
+    await service.close()
+    await custom_client.aclose()
+
+
+async def test_http_service_strips_trailing_slash():
+    """Test that endpoint trailing slash is stripped."""
+    service = AGUIHttpService("http://localhost:8888/")
+    assert service.endpoint == "http://localhost:8888"
+    await service.close()
+
+
+async def test_post_run_successful_streaming(mock_http_client, sample_events):
+    """Test successful streaming of events."""
+
+    # Create async generator for lines
+    async def mock_aiter_lines():
+        sse_data = create_sse_response(sample_events)
+        for line in sse_data.split("\n"):
+            if line:
+                yield line
+
+    # Create mock response
+    mock_response = AsyncMock()
+    mock_response.status_code = 200
+    # aiter_lines is called as a method, so it should return a new generator each time
+    mock_response.aiter_lines = mock_aiter_lines
+
+    # Setup mock streaming context manager
+    mock_stream_context = AsyncMock()
+    mock_stream_context.__aenter__.return_value = mock_response
+    mock_stream_context.__aexit__.return_value = None
+    mock_http_client.stream.return_value = mock_stream_context
+
+    service = AGUIHttpService("http://localhost:8888/", http_client=mock_http_client)
+
+    events = []
+    async for event in service.post_run(
+        thread_id="thread_123", run_id="run_456", messages=[{"role": "user", "content": "Hello"}]
+    ):
+        events.append(event)
+
+    assert len(events) == len(sample_events)
+    assert events[0]["type"] == "RUN_STARTED"
+    assert events[-1]["type"] == "RUN_FINISHED"
+
+    # Verify request was made correctly
+    mock_http_client.stream.assert_called_once()
+    call_args = mock_http_client.stream.call_args
+    assert call_args.args[0] == "POST"
+    assert call_args.args[1] == "http://localhost:8888"
+    assert call_args.kwargs["headers"] == {"Accept": "text/event-stream"}
+
+
+async def test_post_run_with_state_tools_and_interrupts(mock_http_client):
+    """Test posting run with state, tools, and interrupt metadata."""
+
+    async def mock_aiter_lines():
+        return
+        yield  # Make it an async generator
+
+    mock_response = AsyncMock()
+    mock_response.status_code = 200
+    mock_response.aiter_lines = mock_aiter_lines
+
+    mock_stream_context = AsyncMock()
+    mock_stream_context.__aenter__.return_value = mock_response
+    mock_stream_context.__aexit__.return_value = None
+    mock_http_client.stream.return_value = mock_stream_context
+
+    service = AGUIHttpService("http://localhost:8888/", http_client=mock_http_client)
+
+    state = {"user_context": {"name": "Alice"}}
+    tools = [{"type": "function", "function": {"name": "test_tool"}}]
+    available_interrupts = [{"id": "req_1", "type": "request_info"}]
+    expected_available_interrupts = [{"id": "req_1", "reason": "input_required"}]
+    resume = {"interrupts": [{"id": "req_1", "value": "approved"}]}
+    expected_resume = [{"interruptId": "req_1", "status": "resolved", "payload": "approved"}]
+
+    async for _ in service.post_run(
+        thread_id="thread_123",
+        run_id="run_456",
+        messages=[],
+        state=state,
+        tools=tools,
+        available_interrupts=available_interrupts,
+        resume=resume,
+    ):
+        pass
+
+    # Verify state and tools were included in request
+    call_args = mock_http_client.stream.call_args
+    request_data = call_args.kwargs["json"]
+    assert request_data["state"] == state
+    assert request_data["tools"] == tools
+    assert request_data["availableInterrupts"] == expected_available_interrupts
+    assert request_data["resume"] == expected_resume
+
+
+async def test_post_run_serializes_typed_interrupts_and_resume_with_protocol_aliases(mock_http_client):
+    """Typed protocol interrupt and resume models are serialized to canonical wire fields."""
+
+    async def mock_aiter_lines():
+        return
+        yield  # Make it an async generator
+
+    mock_response = AsyncMock()
+    mock_response.status_code = 200
+    mock_response.aiter_lines = mock_aiter_lines
+
+    mock_stream_context = AsyncMock()
+    mock_stream_context.__aenter__.return_value = mock_response
+    mock_stream_context.__aexit__.return_value = None
+    mock_http_client.stream.return_value = mock_stream_context
+
+    service = AGUIHttpService("http://localhost:8888/", http_client=mock_http_client)
+
+    async for _ in service.post_run(
+        thread_id="thread_123",
+        run_id="run_456",
+        messages=[],
+        available_interrupts=[
+            Interrupt(
+                id="approval_1",
+                reason="tool_call",
+                tool_call_id="call_1",
+                response_schema={"type": "object"},
+            )
+        ],
+        resume=[ResumeEntry(interrupt_id="approval_1", status="resolved", payload={"approved": True})],
+    ):
+        pass
+
+    request_data = mock_http_client.stream.call_args.kwargs["json"]
+    assert request_data["availableInterrupts"] == [
+        {
+            "id": "approval_1",
+            "reason": "tool_call",
+            "toolCallId": "call_1",
+            "responseSchema": {"type": "object"},
+        }
+    ]
+    assert request_data["resume"] == [
+        {"interruptId": "approval_1", "status": "resolved", "payload": {"approved": True}}
+    ]
+
+
+async def test_post_run_serializes_legacy_single_resume_mapping_as_canonical_list(mock_http_client):
+    """Legacy single-entry resume mappings are sent as canonical ResumeEntry arrays."""
+
+    async def mock_aiter_lines():
+        return
+        yield  # Make it an async generator
+
+    mock_response = AsyncMock()
+    mock_response.status_code = 200
+    mock_response.aiter_lines = mock_aiter_lines
+
+    mock_stream_context = AsyncMock()
+    mock_stream_context.__aenter__.return_value = mock_response
+    mock_stream_context.__aexit__.return_value = None
+    mock_http_client.stream.return_value = mock_stream_context
+
+    service = AGUIHttpService("http://localhost:8888/", http_client=mock_http_client)
+
+    async for _ in service.post_run(
+        thread_id="thread_123",
+        run_id="run_456",
+        messages=[],
+        resume={"id": "approval_1", "value": {"approved": True}},
+    ):
+        pass
+
+    request_data = mock_http_client.stream.call_args.kwargs["json"]
+    assert request_data["resume"] == [
+        {"interruptId": "approval_1", "status": "resolved", "payload": {"approved": True}}
+    ]
+
+
+async def test_post_run_serializes_legacy_interrupt_without_type(mock_http_client):
+    """Legacy interrupt metadata without reason/type does not crash client serialization."""
+
+    async def mock_aiter_lines():
+        return
+        yield  # Make it an async generator
+
+    mock_response = AsyncMock()
+    mock_response.status_code = 200
+    mock_response.aiter_lines = mock_aiter_lines
+
+    mock_stream_context = AsyncMock()
+    mock_stream_context.__aenter__.return_value = mock_response
+    mock_stream_context.__aexit__.return_value = None
+    mock_http_client.stream.return_value = mock_stream_context
+
+    service = AGUIHttpService("http://localhost:8888/", http_client=mock_http_client)
+
+    async for _ in service.post_run(
+        thread_id="thread_123",
+        run_id="run_456",
+        messages=[],
+        available_interrupts=[{"id": "req_1", "value": {"prompt": "Choose"}}],
+    ):
+        pass
+
+    request_data = mock_http_client.stream.call_args.kwargs["json"]
+    assert request_data["availableInterrupts"][0]["id"] == "req_1"
+    assert request_data["availableInterrupts"][0]["reason"] == "input_required"
+
+
+async def test_post_run_http_error(mock_http_client):
+    """Test handling of HTTP errors."""
+    mock_response = Mock()
+    mock_response.status_code = 500
+    mock_response.text = "Internal Server Error"
+
+    def raise_http_error():
+        raise httpx.HTTPStatusError("Server error", request=Mock(), response=mock_response)
+
+    mock_response_async = AsyncMock()
+    mock_response_async.raise_for_status = raise_http_error
+
+    mock_stream_context = AsyncMock()
+    mock_stream_context.__aenter__.return_value = mock_response_async
+    mock_stream_context.__aexit__.return_value = None
+    mock_http_client.stream.return_value = mock_stream_context
+
+    service = AGUIHttpService("http://localhost:8888/", http_client=mock_http_client)
+
+    with pytest.raises(httpx.HTTPStatusError):
+        async for _ in service.post_run(thread_id="thread_123", run_id="run_456", messages=[]):
+            pass
+
+
+async def test_post_run_invalid_json(mock_http_client):
+    """Test handling of invalid JSON in SSE stream."""
+    invalid_sse = "data: {invalid json}\n\ndata: " + json.dumps({"type": "RUN_FINISHED"}) + "\n"
+
+    async def mock_aiter_lines():
+        for line in invalid_sse.split("\n"):
+            if line:
+                yield line
+
+    mock_response = AsyncMock()
+    mock_response.status_code = 200
+    mock_response.aiter_lines = mock_aiter_lines
+
+    mock_stream_context = AsyncMock()
+    mock_stream_context.__aenter__.return_value = mock_response
+    mock_stream_context.__aexit__.return_value = None
+    mock_http_client.stream.return_value = mock_stream_context
+
+    service = AGUIHttpService("http://localhost:8888/", http_client=mock_http_client)
+
+    events = []
+    async for event in service.post_run(thread_id="thread_123", run_id="run_456", messages=[]):
+        events.append(event)
+
+    # Should skip invalid JSON and continue with valid events
+    assert len(events) == 1
+    assert events[0]["type"] == "RUN_FINISHED"
+
+
+async def test_context_manager():
+    """Test context manager functionality."""
+    async with AGUIHttpService("http://localhost:8888/") as service:
+        assert service.http_client is not None
+        assert service._owns_client is True
+
+    # Client should be closed after exiting context
+
+
+async def test_context_manager_with_external_client():
+    """Test context manager doesn't close external client."""
+    external_client = httpx.AsyncClient()
+
+    async with AGUIHttpService("http://localhost:8888/", http_client=external_client) as service:
+        assert service.http_client is external_client
+        assert service._owns_client is False
+
+    # External client should still be open
+    # (caller's responsibility to close)
+    await external_client.aclose()
+
+
+async def test_post_run_empty_response(mock_http_client):
+    """Test handling of empty response stream."""
+
+    async def mock_aiter_lines():
+        return
+        yield  # Make it an async generator
+
+    mock_response = AsyncMock()
+    mock_response.status_code = 200
+    mock_response.aiter_lines = mock_aiter_lines
+
+    mock_stream_context = AsyncMock()
+    mock_stream_context.__aenter__.return_value = mock_response
+    mock_stream_context.__aexit__.return_value = None
+    mock_http_client.stream.return_value = mock_stream_context
+
+    service = AGUIHttpService("http://localhost:8888/", http_client=mock_http_client)
+
+    events = []
+    async for event in service.post_run(thread_id="thread_123", run_id="run_456", messages=[]):
+        events.append(event)
+
+    assert len(events) == 0

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

@@ -0,0 +1,404 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from typing import Any
+
+from agent_framework import Content, Message
+
+from agent_framework_ag_ui._message_adapters import _deduplicate_messages, _sanitize_tool_history
+
+
+def test_sanitize_tool_history_filters_out_confirm_changes_only_message() -> None:
+    """Test that assistant messages with ONLY confirm_changes are filtered out entirely.
+
+    When an assistant message contains only a confirm_changes tool call (no other tools),
+    the entire message should be filtered out because confirm_changes is a synthetic
+    tool for the approval UI flow that shouldn't be sent to the LLM.
+    """
+    messages = [
+        Message(
+            role="assistant",
+            contents=[
+                Content.from_function_call(
+                    name="confirm_changes",
+                    call_id="call_confirm_123",
+                    arguments='{"changes": "test"}',
+                )
+            ],
+        ),
+        Message(
+            role="user",
+            contents=[Content.from_text(text='{"accepted": true}')],
+        ),
+    ]
+
+    sanitized = _sanitize_tool_history(messages)
+
+    # Assistant message with only confirm_changes should be filtered out
+    assistant_messages = [
+        msg for msg in sanitized if (msg.role if hasattr(msg.role, "value") else str(msg.role)) == "assistant"
+    ]
+    assert len(assistant_messages) == 0
+
+    # No synthetic tool result should be injected since confirm_changes was filtered out
+    tool_messages = [msg for msg in sanitized if (msg.role if hasattr(msg.role, "value") else str(msg.role)) == "tool"]
+    assert len(tool_messages) == 0
+
+
+def test_deduplicate_messages_prefers_non_empty_tool_results() -> None:
+    messages = [
+        Message(
+            role="tool",
+            contents=[Content.from_function_result(call_id="call1", result="")],
+        ),
+        Message(
+            role="tool",
+            contents=[Content.from_function_result(call_id="call1", result="result data")],
+        ),
+    ]
+
+    deduped = _deduplicate_messages(messages)
+    assert len(deduped) == 1
+    assert deduped[0].contents[0].result == "result data"
+
+
+def test_convert_approval_results_to_tool_messages() -> None:
+    """Test that function_result content in user messages gets converted to tool messages.
+
+    This is a regression test for the MCP tool double-call bug where approved tool
+    results ended up in user messages instead of tool messages, causing OpenAI to
+    reject the request with 'tool_call_ids did not have response messages'.
+    """
+    from agent_framework_ag_ui._agent_run import _convert_approval_results_to_tool_messages
+
+    # Simulate what happens after _resolve_approval_responses:
+    # A user message contains function_result content (the executed tool result)
+    messages = [
+        Message(
+            role="assistant",
+            contents=[
+                Content.from_function_call(call_id="call_123", name="my_mcp_tool", arguments="{}"),
+            ],
+        ),
+        Message(
+            role="user",
+            contents=[
+                Content.from_function_result(call_id="call_123", result="tool execution result"),
+            ],
+        ),
+    ]
+
+    _convert_approval_results_to_tool_messages(messages)
+
+    # After conversion, the function result should be in a tool message, not user message
+    assert len(messages) == 2
+
+    # First message unchanged
+    assert messages[0].role == "assistant"
+
+    # Second message should now be role="tool"
+    assert messages[1].role == "tool"
+    assert messages[1].contents[0].type == "function_result"
+    assert messages[1].contents[0].call_id == "call_123"
+
+
+def test_convert_approval_results_preserves_other_user_content() -> None:
+    """Test that user messages with mixed content are handled correctly.
+
+    If a user message has both function_result content and other content (like text),
+    the function_result content should be extracted to a tool message while the
+    remaining content stays in the user message.
+    """
+    from agent_framework_ag_ui._agent_run import _convert_approval_results_to_tool_messages
+
+    messages = [
+        Message(
+            role="assistant",
+            contents=[
+                Content.from_function_call(call_id="call_123", name="my_tool", arguments="{}"),
+            ],
+        ),
+        Message(
+            role="user",
+            contents=[
+                Content.from_text(text="User also said something"),
+                Content.from_function_result(call_id="call_123", result="tool result"),
+            ],
+        ),
+    ]
+
+    _convert_approval_results_to_tool_messages(messages)
+
+    # Should have 3 messages now: assistant, tool (with result), user (with text)
+    # OpenAI requires tool messages immediately after the assistant message with the tool call
+    assert len(messages) == 3
+
+    # First message unchanged
+    assert messages[0].role == "assistant"
+
+    # Second message should be tool with result (must come right after assistant per OpenAI requirements)
+    assert messages[1].role == "tool"
+    assert messages[1].contents[0].type == "function_result"
+
+    # Third message should be user with just text
+    assert messages[2].role == "user"
+    assert len(messages[2].contents) == 1
+    assert messages[2].contents[0].type == "text"
+
+
+def test_sanitize_tool_history_filters_confirm_changes_keeps_other_tools() -> None:
+    """Test that confirm_changes is filtered but other tools are preserved.
+
+    When an assistant message contains both a real tool call and confirm_changes,
+    confirm_changes should be filtered out while the real tool call is kept.
+    No synthetic result is injected for confirm_changes since it's filtered.
+    """
+    messages = [
+        # User asks something
+        Message(
+            role="user",
+            contents=[Content.from_text(text="What time is it?")],
+        ),
+        # Assistant calls MCP tool + confirm_changes
+        Message(
+            role="assistant",
+            contents=[
+                Content.from_function_call(call_id="call_1", name="get_datetime", arguments="{}"),
+                Content.from_function_call(call_id="call_c1", name="confirm_changes", arguments="{}"),
+            ],
+        ),
+        # Tool result for the actual MCP tool
+        Message(
+            role="tool",
+            contents=[Content.from_function_result(call_id="call_1", result="2024-01-01 12:00:00")],
+        ),
+        # User asks something else
+        Message(
+            role="user",
+            contents=[Content.from_text(text="What's the date?")],
+        ),
+    ]
+
+    sanitized = _sanitize_tool_history(messages)
+
+    # Find the assistant message
+    assistant_messages = [
+        msg for msg in sanitized if (msg.role if hasattr(msg.role, "value") else str(msg.role)) == "assistant"
+    ]
+    assert len(assistant_messages) == 1
+
+    # Assistant message should only have get_datetime, not confirm_changes
+    function_call_names = [c.name for c in assistant_messages[0].contents if c.type == "function_call"]
+    assert "get_datetime" in function_call_names
+    assert "confirm_changes" not in function_call_names
+
+    # Only one tool message (for call_1), no synthetic for confirm_changes
+    tool_messages = [msg for msg in sanitized if (msg.role 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_1"
+
+
+def test_sanitize_tool_history_filters_confirm_changes_from_assistant_messages() -> None:
+    """Test that confirm_changes is removed from assistant messages sent to LLM.
+
+    This is a regression test for the human-in-the-loop bug where the LLM would see
+    confirm_changes with function_arguments containing the original steps (e.g., 5 steps)
+    even when the user only approved a subset (e.g., 2 steps), causing the LLM to
+    respond with "Here's your 5-step plan" instead of "Here's your 2-step plan".
+    """
+    messages = [
+        Message(
+            role="user",
+            contents=[Content.from_text(text="Build a robot")],
+        ),
+        # Assistant message with both generate_task_steps and confirm_changes
+        Message(
+            role="assistant",
+            contents=[
+                Content.from_function_call(
+                    call_id="call_1",
+                    name="generate_task_steps",
+                    arguments='{"steps": [{"description": "Step 1"}, {"description": "Step 2"}]}',
+                ),
+                Content.from_function_call(
+                    call_id="call_c1",
+                    name="confirm_changes",
+                    arguments='{"function_arguments": {"steps": [{"description": "Step 1"}, {"description": "Step 2"}]}}',
+                ),
+            ],
+        ),
+        # Approval response
+        Message(
+            role="user",
+            contents=[
+                Content.from_function_approval_response(
+                    approved=True,
+                    id="call_1",
+                    function_call=Content.from_function_call(
+                        call_id="call_1",
+                        name="generate_task_steps",
+                        arguments='{"steps": [{"description": "Step 1"}]}',  # Only 1 step approved
+                    ),
+                ),
+            ],
+        ),
+    ]
+
+    sanitized = _sanitize_tool_history(messages)
+
+    # Find the assistant message in sanitized output
+    assistant_messages = [
+        msg for msg in sanitized if (msg.role if hasattr(msg.role, "value") else str(msg.role)) == "assistant"
+    ]
+
+    assert len(assistant_messages) == 1
+
+    # The assistant message should NOT contain confirm_changes
+    assistant_contents = assistant_messages[0].contents or []
+    function_call_names = [c.name for c in assistant_contents if c.type == "function_call"]
+    assert "generate_task_steps" in function_call_names
+    assert "confirm_changes" not in function_call_names
+
+    # No synthetic tool result for confirm_changes (it was filtered from the message)
+    tool_messages = [msg for msg in sanitized if (msg.role if hasattr(msg.role, "value") else str(msg.role)) == "tool"]
+    # No tool results expected since there are no completed tool calls
+    # (the approval response is handled separately by the framework)
+    tool_call_ids = {str(msg.contents[0].call_id) for msg in tool_messages}
+    assert "call_c1" not in tool_call_ids  # No synthetic result for confirm_changes
+
+
+# ---------------------------------------------------------------------------
+# Tests for _clean_resolved_approvals_from_snapshot
+# ---------------------------------------------------------------------------
+
+
+def test_clean_resolved_approvals_from_snapshot() -> None:
+    """Approval payload in snapshot should be replaced with the actual tool result."""
+    import json
+
+    from agent_framework_ag_ui._agent_run import _clean_resolved_approvals_from_snapshot
+
+    # Snapshot still has the approval payload
+    snapshot_messages: list[dict[str, Any]] = [  # type: ignore[name-defined]
+        {"role": "user", "content": "What time is it?", "id": "msg_1"},
+        {
+            "role": "assistant",
+            "content": "",
+            "tool_calls": [
+                {"id": "call_123", "type": "function", "function": {"name": "get_datetime", "arguments": "{}"}}
+            ],
+            "id": "msg_2",
+        },
+        {
+            "role": "tool",
+            "content": json.dumps({"accepted": True}),
+            "toolCallId": "call_123",
+            "id": "msg_3",
+        },
+    ]
+
+    # Resolved provider messages have the actual tool result
+    resolved_messages = [
+        Message(role="user", contents=[Content.from_text(text="What time is it?")]),
+        Message(
+            role="assistant",
+            contents=[Content.from_function_call(call_id="call_123", name="get_datetime", arguments="{}")],
+        ),
+        Message(
+            role="tool",
+            contents=[Content.from_function_result(call_id="call_123", result="2024-01-01 12:00:00")],
+        ),
+    ]
+
+    _clean_resolved_approvals_from_snapshot(snapshot_messages, resolved_messages)
+
+    # The approval payload should now be replaced with the tool result
+    tool_snap = snapshot_messages[2]
+    assert tool_snap["content"] == "2024-01-01 12:00:00"
+
+
+def test_clean_resolved_approvals_from_snapshot_no_approvals() -> None:
+    """When there are no approval payloads, snapshot should be unchanged."""
+    from agent_framework_ag_ui._agent_run import _clean_resolved_approvals_from_snapshot  # type: ignore
+
+    snapshot_messages: list[dict[str, Any]] = [  # type: ignore[name-defined]
+        {"role": "user", "content": "Hello", "id": "msg_1"},
+        {"role": "assistant", "content": "Hi there", "id": "msg_2"},
+    ]
+    original = [dict(m) for m in snapshot_messages]
+
+    resolved_messages = [
+        Message(role="user", contents=[Content.from_text(text="Hello")]),
+        Message(role="assistant", contents=[Content.from_text(text="Hi there")]),
+    ]
+
+    _clean_resolved_approvals_from_snapshot(snapshot_messages, resolved_messages)
+
+    # Nothing should have changed
+    assert snapshot_messages == original
+
+
+def test_cleaned_snapshot_prevents_approval_reprocessing() -> None:
+    """After snapshot cleaning, approval payload is replaced so it won't re-trigger on next turn.
+
+    Simulates what happens on Turn 2: the approval is processed, the tool executes,
+    and _clean_resolved_approvals_from_snapshot replaces the approval payload with the
+    real tool result. On Turn 3, CopilotKit re-sends the cleaned snapshot, which no
+    longer contains an approval payload — so no function_approval_response is produced.
+    """
+    import json
+
+    from agent_framework_ag_ui._agent_run import _clean_resolved_approvals_from_snapshot
+    from agent_framework_ag_ui._message_adapters import normalize_agui_input_messages
+
+    # Turn 2 snapshot: still has the raw approval payload
+    snapshot_messages: list[dict[str, Any]] = [  # type: ignore[name-defined]
+        {"role": "user", "content": "What time is it?", "id": "msg_1"},
+        {
+            "role": "assistant",
+            "content": "",
+            "tool_calls": [
+                {"id": "call_789", "type": "function", "function": {"name": "get_datetime", "arguments": "{}"}}
+            ],
+            "id": "msg_2",
+        },
+        {
+            "role": "tool",
+            "content": json.dumps({"accepted": True}),
+            "toolCallId": "call_789",
+            "id": "msg_3",
+        },
+    ]
+
+    # Resolved provider messages after tool execution
+    resolved_messages = [
+        Message(role="user", contents=[Content.from_text(text="What time is it?")]),
+        Message(
+            role="assistant",
+            contents=[Content.from_function_call(call_id="call_789", name="get_datetime", arguments="{}")],
+        ),
+        Message(
+            role="tool",
+            contents=[Content.from_function_result(call_id="call_789", result="2024-01-01 12:00:00")],
+        ),
+    ]
+
+    # Fix B: clean the snapshot
+    _clean_resolved_approvals_from_snapshot(snapshot_messages, resolved_messages)
+
+    # Snapshot should now have the real tool result
+    assert snapshot_messages[2]["content"] == "2024-01-01 12:00:00"
+
+    # Simulate Turn 3: CopilotKit re-sends the cleaned snapshot + new messages
+    turn3_messages: list[dict[str, Any]] = list(snapshot_messages) + [  # type: ignore[name-defined]
+        {"role": "assistant", "content": "It is 12:00 PM.", "id": "msg_4"},
+        {"role": "user", "content": "Thanks!", "id": "msg_5"},
+    ]
+
+    provider_messages, _ = normalize_agui_input_messages(turn3_messages)
+
+    # No function_approval_response should exist — the approval payload is gone
+    for msg in provider_messages:
+        for content in msg.contents or []:
+            assert content.type != "function_approval_response", (
+                f"Stale approval was re-processed on subsequent turn: {content}"
+            )

python/packages/ag-ui/tests/ag_ui/test_multi_turn.py

@@ -0,0 +1,332 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Multi-turn conversation tests: POST → collect events → extract snapshot → POST again.
+
+These tests catch round-trip fidelity bugs: if MessagesSnapshotEvent produces a
+malformed message list, the second turn will fail during normalize_agui_input_messages()
+or produce incorrect behavior.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from agent_framework import AgentResponseUpdate, Content
+from conftest import StubAgent  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+from sse_helpers import (  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+    parse_sse_response,
+    parse_sse_to_event_stream,
+)
+
+from agent_framework_ag_ui import AgentFrameworkAgent, add_agent_framework_fastapi_endpoint
+
+
+def _build_app_with_agent(updates: list[AgentResponseUpdate], **kwargs: Any) -> FastAPI:
+    stub = StubAgent(updates=updates)
+    agent = AgentFrameworkAgent(agent=stub, **kwargs)
+    app = FastAPI()
+    add_agent_framework_fastapi_endpoint(app, agent)
+    return app
+
+
+def _extract_snapshot_messages(response_content: bytes) -> list[dict[str, Any]]:
+    """Extract the latest MessagesSnapshotEvent.messages from SSE response bytes."""
+    raw_events = parse_sse_response(response_content)
+    snapshot_msgs: list[dict[str, Any]] | None = None
+    for event in raw_events:
+        if event.get("type") == "MESSAGES_SNAPSHOT":
+            snapshot_msgs = event.get("messages", [])
+    assert snapshot_msgs is not None, "No MESSAGES_SNAPSHOT event found"
+    return snapshot_msgs
+
+
+# ── Basic multi-turn chat ──
+
+
+def test_basic_multi_turn_chat() -> None:
+    """Turn 1: user→assistant. Turn 2: user→assistant with prior history from snapshot."""
+    app = _build_app_with_agent(
+        [
+            AgentResponseUpdate(contents=[Content.from_text(text="Hello! How can I help?")], role="assistant"),
+        ]
+    )
+    client = TestClient(app)
+
+    # Turn 1
+    resp1 = client.post(
+        "/",
+        json={
+            "messages": [{"role": "user", "content": "Hi there"}],
+            "threadId": "thread-multi",
+            "runId": "run-1",
+        },
+    )
+    assert resp1.status_code == 200
+    stream1 = parse_sse_to_event_stream(resp1.content)
+    stream1.assert_bookends()
+    stream1.assert_text_messages_balanced()
+
+    # Extract snapshot messages from turn 1
+    snapshot_messages = _extract_snapshot_messages(resp1.content)
+
+    # Turn 2: send snapshot messages + new user message
+    turn2_messages = list(snapshot_messages) + [{"role": "user", "content": "Tell me more"}]
+    resp2 = client.post(
+        "/",
+        json={
+            "messages": turn2_messages,
+            "threadId": "thread-multi",
+            "runId": "run-2",
+        },
+    )
+    assert resp2.status_code == 200
+    stream2 = parse_sse_to_event_stream(resp2.content)
+    stream2.assert_bookends()
+    stream2.assert_text_messages_balanced()
+    stream2.assert_no_run_error()
+
+
+# ── Tool call history round-trip ──
+
+
+def test_tool_call_history_round_trips() -> None:
+    """Turn 1: tool call + result. Turn 2: snapshot messages correctly reconstruct tool history."""
+    app = _build_app_with_agent(
+        [
+            AgentResponseUpdate(
+                contents=[Content.from_function_call(name="get_weather", call_id="call-1", arguments='{"city": "SF"}')],
+                role="assistant",
+            ),
+            AgentResponseUpdate(
+                contents=[Content.from_function_result(call_id="call-1", result="72°F")],
+                role="assistant",
+            ),
+            AgentResponseUpdate(
+                contents=[Content.from_text(text="It's warm!")],
+                role="assistant",
+            ),
+        ]
+    )
+    client = TestClient(app)
+
+    # Turn 1
+    resp1 = client.post(
+        "/",
+        json={
+            "messages": [{"role": "user", "content": "What's the weather?"}],
+            "threadId": "thread-tool-multi",
+            "runId": "run-1",
+        },
+    )
+    assert resp1.status_code == 200
+    stream1 = parse_sse_to_event_stream(resp1.content)
+    stream1.assert_tool_calls_balanced()
+
+    # Extract snapshot and verify it has tool history
+    snapshot_messages = _extract_snapshot_messages(resp1.content)
+    roles = [m.get("role") for m in snapshot_messages]
+    assert "tool" in roles or "assistant" in roles, f"Expected tool/assistant messages in snapshot, got: {roles}"
+
+    # Turn 2: send snapshot + new question
+    turn2_messages = list(snapshot_messages) + [{"role": "user", "content": "What about tomorrow?"}]
+    resp2 = client.post(
+        "/",
+        json={
+            "messages": turn2_messages,
+            "threadId": "thread-tool-multi",
+            "runId": "run-2",
+        },
+    )
+    assert resp2.status_code == 200
+    stream2 = parse_sse_to_event_stream(resp2.content)
+    stream2.assert_bookends()
+    stream2.assert_no_run_error()
+
+
+# ── Approval interrupt/resume round-trip ──
+
+
+async def test_approval_interrupt_resume_round_trip() -> None:
+    """Turn 1: approval request → interrupt with confirm_changes. Turn 2: confirm_changes result → confirmation text.
+
+    The confirm_changes flow uses a specific message format that bypasses the agent
+    and directly emits a confirmation text message.
+    """
+    from event_stream import EventStream  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+
+    steps = [{"description": "Execute task", "status": "enabled"}]
+
+    # Build agent with predictive state and confirmation
+    stub = StubAgent(
+        updates=[
+            AgentResponseUpdate(
+                contents=[
+                    Content.from_function_call(
+                        name="generate_task_steps",
+                        call_id="call-steps",
+                        arguments=json.dumps({"steps": steps}),
+                    )
+                ],
+                role="assistant",
+            ),
+        ]
+    )
+    agent = AgentFrameworkAgent(
+        agent=stub,
+        state_schema={"tasks": {"type": "array"}},
+        predict_state_config={"tasks": {"tool": "generate_task_steps", "tool_argument": "steps"}},
+        require_confirmation=True,
+    )
+
+    # Turn 1
+    events1 = [
+        e
+        async for e in agent.run(
+            {
+                "thread_id": "thread-approval-multi",
+                "run_id": "run-1",
+                "messages": [{"role": "user", "content": "Plan my tasks"}],
+                "state": {"tasks": []},
+            }
+        )
+    ]
+    stream1 = EventStream(events1)
+    stream1.assert_bookends()
+    stream1.assert_tool_calls_balanced()
+
+    # Should have interrupt with function_approval_request
+    interrupt1 = stream1.run_finished_interrupts()
+    assert interrupt1, "Expected interrupt in RUN_FINISHED"
+
+    # Verify confirm_changes tool call was emitted
+    tool_starts = stream1.get("TOOL_CALL_START")
+    tool_names = [getattr(s, "tool_call_name", None) for s in tool_starts]
+    assert "confirm_changes" in tool_names, f"Expected confirm_changes in tool calls, got {tool_names}"
+
+    # Turn 2: Direct confirm_changes response (the way CopilotKit sends it)
+    # Construct the messages as CopilotKit would - with the confirm_changes tool call
+    # and a tool result
+    confirm_tool = [s for s in tool_starts if getattr(s, "tool_call_name", None) == "confirm_changes"][0]
+    confirm_id = confirm_tool.tool_call_id
+    confirm_args = None
+    for e in stream1.get("TOOL_CALL_ARGS"):
+        if e.tool_call_id == confirm_id:
+            confirm_args = e.delta
+            break
+
+    turn2_messages = [
+        {"role": "user", "content": "Plan my tasks"},
+        {
+            "role": "assistant",
+            "tool_calls": [
+                {
+                    "id": confirm_id,
+                    "type": "function",
+                    "function": {"name": "confirm_changes", "arguments": confirm_args or "{}"},
+                },
+            ],
+        },
+        {
+            "role": "tool",
+            "toolCallId": confirm_id,
+            "content": json.dumps({"accepted": True, "steps": steps}),
+        },
+    ]
+
+    events2 = [
+        e
+        async for e in agent.run(
+            {
+                "thread_id": "thread-approval-multi",
+                "run_id": "run-2",
+                "messages": turn2_messages,
+                "state": {"tasks": []},
+            }
+        )
+    ]
+    stream2 = EventStream(events2)
+    stream2.assert_bookends()
+    stream2.assert_text_messages_balanced()
+    stream2.assert_no_run_error()
+
+    # Turn 2 should have confirmation text (the approval handler generates it)
+    text_events = stream2.get("TEXT_MESSAGE_CONTENT")
+    assert text_events, "Expected confirmation text message in turn 2"
+
+    # Turn 2 should NOT have interrupt (approval completed)
+    finished2 = stream2.last("RUN_FINISHED")
+    dumped2 = finished2.model_dump(by_alias=True, exclude_none=True)
+    assert "outcome" not in dumped2, f"Expected no interrupt after approval, got {dumped2.get('outcome')}"
+
+
+# ── Workflow interrupt/resume round-trip ──
+# Note: Workflow tests use async agent.run() directly instead of HTTP TestClient
+# because the sync TestClient runs in a different event loop, which conflicts
+# with the workflow's asyncio Queue.
+
+
+async def test_workflow_interrupt_resume_round_trip() -> None:
+    """Turn 1: workflow request_info → interrupt. Turn 2: resume → completion."""
+    from event_stream import EventStream  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+
+    from agent_framework_ag_ui_examples.agents.subgraphs_agent import subgraphs_agent
+
+    agent = subgraphs_agent()
+
+    # Turn 1: initial request → flight interrupt
+    events1 = [
+        event
+        async for event in agent.run(
+            {
+                "messages": [{"role": "user", "content": "Plan a trip to SF"}],
+                "thread_id": "thread-wf-multi",
+                "run_id": "run-1",
+            }
+        )
+    ]
+    stream1 = EventStream(events1)
+    stream1.assert_bookends()
+    stream1.assert_no_run_error()
+
+    interrupt1 = stream1.run_finished_interrupts()
+    assert interrupt1, "Expected flight interrupt"
+    assert stream1.interrupt_metadata_value(interrupt1[0])["agent"] == "flights"
+
+    # Turn 2: resume with flight selection
+    events2 = [
+        event
+        async for event in agent.run(
+            {
+                "messages": [],
+                "thread_id": "thread-wf-multi",
+                "run_id": "run-2",
+                "resume": {
+                    "interrupts": [
+                        {
+                            "id": interrupt1[0]["id"],
+                            "value": json.dumps(
+                                {
+                                    "airline": "United",
+                                    "departure": "Amsterdam (AMS)",
+                                    "arrival": "San Francisco (SFO)",
+                                    "price": "$720",
+                                    "duration": "12h 15m",
+                                }
+                            ),
+                        }
+                    ],
+                },
+            }
+        )
+    ]
+    stream2 = EventStream(events2)
+    stream2.assert_bookends()
+    stream2.assert_no_run_error()
+
+    # Should now have hotel interrupt
+    interrupt2 = stream2.run_finished_interrupts()
+    assert interrupt2, "Expected hotel interrupt"
+    assert stream2.interrupt_metadata_value(interrupt2[0])["agent"] == "hotels"

python/packages/ag-ui/tests/ag_ui/test_predictive_state.py

@@ -0,0 +1,320 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for predictive state handling."""
+
+from ag_ui.core import StateDeltaEvent
+
+from agent_framework_ag_ui._orchestration._predictive_state import PredictiveStateHandler
+
+
+class TestPredictiveStateHandlerInit:
+    """Tests for PredictiveStateHandler initialization."""
+
+    def test_default_init(self):
+        """Initializes with default values."""
+        handler = PredictiveStateHandler()
+        assert handler.predict_state_config == {}
+        assert handler.current_state == {}
+        assert handler.streaming_tool_args == ""
+        assert handler.last_emitted_state == {}
+        assert handler.state_delta_count == 0
+        assert handler.pending_state_updates == {}
+
+    def test_init_with_config(self):
+        """Initializes with provided config."""
+        config = {"document": {"tool": "write_doc", "tool_argument": "content"}}
+        state = {"document": "initial"}
+        handler = PredictiveStateHandler(predict_state_config=config, current_state=state)
+        assert handler.predict_state_config == config
+        assert handler.current_state == state
+
+
+class TestResetStreaming:
+    """Tests for reset_streaming method."""
+
+    def test_resets_streaming_state(self):
+        """Resets streaming-related state."""
+        handler = PredictiveStateHandler()
+        handler.streaming_tool_args = "some accumulated args"
+        handler.state_delta_count = 5
+
+        handler.reset_streaming()
+
+        assert handler.streaming_tool_args == ""
+        assert handler.state_delta_count == 0
+
+
+class TestExtractStateValue:
+    """Tests for extract_state_value method."""
+
+    def test_no_config(self):
+        """Returns None when no config."""
+        handler = PredictiveStateHandler()
+        result = handler.extract_state_value("some_tool", {"arg": "value"})
+        assert result is None
+
+    def test_no_args(self):
+        """Returns None when args is None."""
+        handler = PredictiveStateHandler(predict_state_config={"key": {"tool": "tool", "tool_argument": "arg"}})
+        result = handler.extract_state_value("tool", None)
+        assert result is None
+
+    def test_empty_args(self):
+        """Returns None when args is empty string."""
+        handler = PredictiveStateHandler(predict_state_config={"key": {"tool": "tool", "tool_argument": "arg"}})
+        result = handler.extract_state_value("tool", "")
+        assert result is None
+
+    def test_tool_not_in_config(self):
+        """Returns None when tool not in config."""
+        handler = PredictiveStateHandler(predict_state_config={"key": {"tool": "other_tool", "tool_argument": "arg"}})
+        result = handler.extract_state_value("some_tool", {"arg": "value"})
+        assert result is None
+
+    def test_extracts_specific_argument(self):
+        """Extracts value from specific tool argument."""
+        handler = PredictiveStateHandler(
+            predict_state_config={"document": {"tool": "write_doc", "tool_argument": "content"}}
+        )
+        result = handler.extract_state_value("write_doc", {"content": "Hello world"})
+        assert result == ("document", "Hello world")
+
+    def test_extracts_with_wildcard(self):
+        """Extracts entire args with * wildcard."""
+        handler = PredictiveStateHandler(predict_state_config={"data": {"tool": "update_data", "tool_argument": "*"}})
+        args = {"key1": "value1", "key2": "value2"}
+        result = handler.extract_state_value("update_data", args)
+        assert result == ("data", args)
+
+    def test_extracts_from_json_string(self):
+        """Extracts value from JSON string args."""
+        handler = PredictiveStateHandler(
+            predict_state_config={"document": {"tool": "write_doc", "tool_argument": "content"}}
+        )
+        result = handler.extract_state_value("write_doc", '{"content": "Hello world"}')
+        assert result == ("document", "Hello world")
+
+    def test_argument_not_in_args(self):
+        """Returns None when tool_argument not in args."""
+        handler = PredictiveStateHandler(
+            predict_state_config={"document": {"tool": "write_doc", "tool_argument": "content"}}
+        )
+        result = handler.extract_state_value("write_doc", {"other": "value"})
+        assert result is None
+
+
+class TestIsPredictiveTool:
+    """Tests for is_predictive_tool method."""
+
+    def test_none_tool_name(self):
+        """Returns False for None tool name."""
+        handler = PredictiveStateHandler(predict_state_config={"key": {"tool": "some_tool", "tool_argument": "arg"}})
+        assert handler.is_predictive_tool(None) is False
+
+    def test_no_config(self):
+        """Returns False when no config."""
+        handler = PredictiveStateHandler()
+        assert handler.is_predictive_tool("some_tool") is False
+
+    def test_tool_in_config(self):
+        """Returns True when tool is in config."""
+        handler = PredictiveStateHandler(predict_state_config={"key": {"tool": "some_tool", "tool_argument": "arg"}})
+        assert handler.is_predictive_tool("some_tool") is True
+
+    def test_tool_not_in_config(self):
+        """Returns False when tool not in config."""
+        handler = PredictiveStateHandler(predict_state_config={"key": {"tool": "other_tool", "tool_argument": "arg"}})
+        assert handler.is_predictive_tool("some_tool") is False
+
+
+class TestEmitStreamingDeltas:
+    """Tests for emit_streaming_deltas method."""
+
+    def test_no_tool_name(self):
+        """Returns empty list for None tool name."""
+        handler = PredictiveStateHandler(predict_state_config={"key": {"tool": "tool", "tool_argument": "arg"}})
+        result = handler.emit_streaming_deltas(None, '{"arg": "value"}')
+        assert result == []
+
+    def test_no_config(self):
+        """Returns empty list when no config."""
+        handler = PredictiveStateHandler()
+        result = handler.emit_streaming_deltas("some_tool", '{"arg": "value"}')
+        assert result == []
+
+    def test_accumulates_args(self):
+        """Accumulates argument chunks."""
+        handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}})
+        handler.emit_streaming_deltas("write", '{"text')
+        handler.emit_streaming_deltas("write", '": "hello')
+        assert handler.streaming_tool_args == '{"text": "hello'
+
+    def test_emits_delta_on_complete_json(self):
+        """Emits delta when JSON is complete."""
+        handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}})
+        events = handler.emit_streaming_deltas("write", '{"text": "hello"}')
+        assert len(events) == 1
+        assert isinstance(events[0], StateDeltaEvent)
+        assert events[0].delta[0]["path"] == "/doc"
+        assert events[0].delta[0]["value"] == "hello"
+        assert events[0].delta[0]["op"] == "replace"
+
+    def test_emits_delta_on_partial_json(self):
+        """Emits delta from partial JSON using regex."""
+        handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}})
+        # First chunk - partial
+        events = handler.emit_streaming_deltas("write", '{"text": "hel')
+        assert len(events) == 1
+        assert events[0].delta[0]["value"] == "hel"
+
+    def test_does_not_emit_duplicate_deltas(self):
+        """Does not emit delta when value unchanged."""
+        handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}})
+        # First emission
+        events1 = handler.emit_streaming_deltas("write", '{"text": "hello"}')
+        assert len(events1) == 1
+
+        # Reset and emit same value again
+        handler.streaming_tool_args = ""
+        events2 = handler.emit_streaming_deltas("write", '{"text": "hello"}')
+        assert len(events2) == 0  # No duplicate
+
+    def test_emits_delta_on_value_change(self):
+        """Emits delta when value changes."""
+        handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}})
+        # First value
+        events1 = handler.emit_streaming_deltas("write", '{"text": "hello"}')
+        assert len(events1) == 1
+
+        # Reset and new value
+        handler.streaming_tool_args = ""
+        events2 = handler.emit_streaming_deltas("write", '{"text": "world"}')
+        assert len(events2) == 1
+        assert events2[0].delta[0]["value"] == "world"
+
+    def test_tracks_pending_updates(self):
+        """Tracks pending state updates."""
+        handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}})
+        handler.emit_streaming_deltas("write", '{"text": "hello"}')
+        assert handler.pending_state_updates == {"doc": "hello"}
+
+
+class TestEmitPartialDeltas:
+    """Tests for _emit_partial_deltas method."""
+
+    def test_unescapes_newlines(self):
+        """Unescapes \\n in partial values."""
+        handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}})
+        handler.streaming_tool_args = '{"text": "line1\\nline2'
+        events = handler._emit_partial_deltas("write")
+        assert len(events) == 1
+        assert events[0].delta[0]["value"] == "line1\nline2"
+
+    def test_handles_escaped_quotes_partially(self):
+        """Handles escaped quotes - regex stops at quote character."""
+        handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}})
+        # The regex pattern [^"]* stops at ANY quote, including escaped ones.
+        # This is expected behavior for partial streaming - the full JSON
+        # will be parsed correctly when complete.
+        handler.streaming_tool_args = '{"text": "say \\"hi'
+        events = handler._emit_partial_deltas("write")
+        assert len(events) == 1
+        # Captures "say \" then the backslash gets converted to empty string
+        # by the replace("\\\\", "\\") first, then replace('\\"', '"')
+        # but since there's no closing quote, we get "say \"
+        # After .replace("\\\\", "\\") -> "say \"
+        # After .replace('\\"', '"') -> "say "  (but actually still "say \" due to order)
+        # The actual result: backslash is preserved since it's not a valid escape sequence
+        assert events[0].delta[0]["value"] == "say \\"
+
+    def test_unescapes_backslashes(self):
+        """Unescapes \\\\ in partial values."""
+        handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}})
+        handler.streaming_tool_args = '{"text": "path\\\\to\\\\file'
+        events = handler._emit_partial_deltas("write")
+        assert len(events) == 1
+        assert events[0].delta[0]["value"] == "path\\to\\file"
+
+
+class TestEmitCompleteDeltas:
+    """Tests for _emit_complete_deltas method."""
+
+    def test_emits_for_matching_tool(self):
+        """Emits delta for tool matching config."""
+        handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}})
+        events = handler._emit_complete_deltas("write", {"text": "content"})
+        assert len(events) == 1
+        assert events[0].delta[0]["value"] == "content"
+
+    def test_skips_non_matching_tool(self):
+        """Skips tools not matching config."""
+        handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}})
+        events = handler._emit_complete_deltas("other_tool", {"text": "content"})
+        assert len(events) == 0
+
+    def test_handles_wildcard_argument(self):
+        """Handles * wildcard for entire args."""
+        handler = PredictiveStateHandler(predict_state_config={"data": {"tool": "update", "tool_argument": "*"}})
+        args = {"key1": "val1", "key2": "val2"}
+        events = handler._emit_complete_deltas("update", args)
+        assert len(events) == 1
+        assert events[0].delta[0]["value"] == args
+
+    def test_skips_missing_argument(self):
+        """Skips when tool_argument not in args."""
+        handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}})
+        events = handler._emit_complete_deltas("write", {"other": "value"})
+        assert len(events) == 0
+
+
+class TestCreateDeltaEvent:
+    """Tests for _create_delta_event method."""
+
+    def test_creates_event(self):
+        """Creates StateDeltaEvent with correct structure."""
+        handler = PredictiveStateHandler()
+        event = handler._create_delta_event("key", "value")
+
+        assert isinstance(event, StateDeltaEvent)
+        assert event.delta[0]["op"] == "replace"
+        assert event.delta[0]["path"] == "/key"
+        assert event.delta[0]["value"] == "value"
+
+    def test_increments_count(self):
+        """Increments state_delta_count."""
+        handler = PredictiveStateHandler()
+        handler._create_delta_event("key", "value")
+        assert handler.state_delta_count == 1
+        handler._create_delta_event("key", "value2")
+        assert handler.state_delta_count == 2
+
+
+class TestApplyPendingUpdates:
+    """Tests for apply_pending_updates method."""
+
+    def test_applies_pending_to_current(self):
+        """Applies pending updates to current state."""
+        handler = PredictiveStateHandler(current_state={"existing": "value"})
+        handler.pending_state_updates = {"doc": "new content", "count": 5}
+
+        handler.apply_pending_updates()
+
+        assert handler.current_state == {"existing": "value", "doc": "new content", "count": 5}
+
+    def test_clears_pending_updates(self):
+        """Clears pending updates after applying."""
+        handler = PredictiveStateHandler()
+        handler.pending_state_updates = {"doc": "content"}
+
+        handler.apply_pending_updates()
+
+        assert handler.pending_state_updates == {}
+
+    def test_overwrites_existing_keys(self):
+        """Overwrites existing keys in current state."""
+        handler = PredictiveStateHandler(current_state={"doc": "old"})
+        handler.pending_state_updates = {"doc": "new"}
+
+        handler.apply_pending_updates()
+
+        assert handler.current_state["doc"] == "new"

python/packages/ag-ui/tests/ag_ui/test_public_exports.py

@@ -0,0 +1,66 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Public export coverage for AG-UI package surfaces."""
+
+
+def test_agent_framework_ag_ui_exports_workflow() -> None:
+    """Runtime package should export AgentFrameworkWorkflow."""
+    from agent_framework_ag_ui import AgentFrameworkWorkflow
+
+    assert AgentFrameworkWorkflow.__name__ == "AgentFrameworkWorkflow"
+
+
+def test_core_ag_ui_lazy_exports_include_only_stable_api() -> None:
+    """Core facade should expose only the stable high-level AG-UI API."""
+    from agent_framework import ag_ui
+
+    assert hasattr(ag_ui, "AgentFrameworkWorkflow")
+    assert hasattr(ag_ui, "AgentFrameworkAgent")
+    assert hasattr(ag_ui, "AGUIChatClient")
+    assert hasattr(ag_ui, "add_agent_framework_fastapi_endpoint")
+    assert hasattr(ag_ui, "state_update")
+
+    assert not hasattr(ag_ui, "WorkflowFactory")
+    assert not hasattr(ag_ui, "AGUIRequest")
+    assert not hasattr(ag_ui, "RunMetadata")
+
+
+def test_agent_framework_ag_ui_exports_state_update() -> None:
+    """Runtime package should export the ``state_update`` helper."""
+    from agent_framework_ag_ui import state_update
+
+    assert callable(state_update)
+
+
+def test_agent_framework_ag_ui_exports_snapshot_primitives() -> None:
+    """Runtime package should export AG-UI Thread Snapshot primitives."""
+    from agent_framework_ag_ui import (
+        DEFAULT_MAX_THREAD_SNAPSHOTS,
+        AGUIThreadSnapshot,
+        AGUIThreadSnapshotStore,
+        InMemoryAGUIThreadSnapshotStore,
+    )
+
+    assert AGUIThreadSnapshot.__name__ == "AGUIThreadSnapshot"
+    assert AGUIThreadSnapshotStore.__name__ == "AGUIThreadSnapshotStore"
+    assert InMemoryAGUIThreadSnapshotStore.__name__ == "InMemoryAGUIThreadSnapshotStore"
+    assert DEFAULT_MAX_THREAD_SNAPSHOTS >= 1
+
+
+def test_core_ag_ui_lazy_exports_include_event_converter_and_http_service() -> None:
+    """Core facade must expose AGUIEventConverter, AGUIHttpService, and __version__."""
+    from agent_framework import ag_ui
+
+    assert hasattr(ag_ui, "AGUIEventConverter")
+    assert hasattr(ag_ui, "AGUIHttpService")
+    assert hasattr(ag_ui, "__version__")
+
+
+def test_core_ag_ui_lazy_exports_include_snapshot_primitives() -> None:
+    """Core facade must expose snapshot primitives needed for endpoint configuration."""
+    from agent_framework import ag_ui
+
+    assert hasattr(ag_ui, "AGUIThreadSnapshot")
+    assert hasattr(ag_ui, "AGUIThreadSnapshotStore")
+    assert hasattr(ag_ui, "InMemoryAGUIThreadSnapshotStore")
+    assert hasattr(ag_ui, "SnapshotScopeResolver")

python/packages/ag-ui/tests/ag_ui/test_run_common.py

@@ -0,0 +1,550 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for _run_common.py edge cases."""
+
+import logging
+
+import pytest
+from ag_ui.core import EventType
+from agent_framework import Content
+
+from agent_framework_ag_ui import state_update
+from agent_framework_ag_ui._orchestration._predictive_state import PredictiveStateHandler
+from agent_framework_ag_ui._run_common import (
+    FlowState,
+    _build_run_finished_event,
+    _emit_mcp_tool_result,
+    _emit_tool_result,
+    _extract_resume_payload,
+    _extract_tool_result_state,
+    _normalize_resume_interrupts,
+    _reconstruct_messages_from_thread_snapshot,
+    _strict_resume_entries,
+)
+from agent_framework_ag_ui._state import TOOL_RESULT_DISPLAY_KEY, TOOL_RESULT_STATE_KEY
+
+
+class TestNormalizeResumeInterrupts:
+    """Tests for _normalize_resume_interrupts edge cases."""
+
+    def test_plain_list_of_dicts(self):
+        """Resume payload as a plain list of interrupt dicts."""
+        result = _normalize_resume_interrupts([{"id": "x", "value": "y"}])
+        assert result == [{"id": "x", "value": "y"}]
+
+    def test_dict_with_singular_interrupt_key(self):
+        """Resume dict using 'interrupt' (singular) instead of 'interrupts'."""
+        result = _normalize_resume_interrupts({"interrupt": [{"id": "x", "value": "y"}]})
+        assert result == [{"id": "x", "value": "y"}]
+
+    def test_dict_without_interrupts_key_wraps_as_candidate(self):
+        """Resume dict without interrupts/interrupt key wraps the dict itself."""
+        result = _normalize_resume_interrupts({"id": "x", "value": "y"})
+        assert result == [{"id": "x", "value": "y"}]
+
+    def test_non_dict_items_in_list_are_skipped(self):
+        """Non-dict items in candidate list are silently skipped."""
+        result = _normalize_resume_interrupts([None, "string", {"id": "x", "value": "y"}])
+        assert result == [{"id": "x", "value": "y"}]
+
+    def test_items_missing_id_are_skipped(self):
+        """Dict items without any id field are skipped."""
+        result = _normalize_resume_interrupts([{"name": "test"}])
+        assert result == []
+
+    def test_response_key_used_as_value(self):
+        """'response' key is used as value when 'value' is absent."""
+        result = _normalize_resume_interrupts([{"id": "x", "response": "approved"}])
+        assert result == [{"id": "x", "value": "approved"}]
+
+    def test_neither_value_nor_response_uses_remaining_fields(self):
+        """When neither 'value' nor 'response' key exists, remaining fields become value."""
+        result = _normalize_resume_interrupts([{"id": "x", "extra": "data", "more": 42}])
+        assert result == [{"id": "x", "value": {"extra": "data", "more": 42}}]
+
+    def test_none_payload_returns_empty(self):
+        """None resume payload returns empty list."""
+        assert _normalize_resume_interrupts(None) == []
+
+    def test_non_dict_non_list_returns_empty(self):
+        """Non-dict, non-list payload returns empty list."""
+        assert _normalize_resume_interrupts(42) == []
+
+    def test_interrupt_id_key_used_as_id(self):
+        """interruptId key is accepted as identifier."""
+        result = _normalize_resume_interrupts([{"interruptId": "abc", "value": "yes"}])
+        assert result == [{"id": "abc", "value": "yes"}]
+
+    def test_tool_call_id_key_used_as_id(self):
+        """toolCallId key is accepted as identifier."""
+        result = _normalize_resume_interrupts([{"toolCallId": "tc1", "value": "done"}])
+        assert result == [{"id": "tc1", "value": "done"}]
+
+    def test_canonical_resume_entry_uses_interrupt_id_and_payload(self):
+        """Canonical ResumeEntry dictionaries preserve status and map payload to legacy runner values."""
+        result = _normalize_resume_interrupts(
+            [{"interrupt_id": "req_1", "status": "resolved", "payload": {"approved": True}}]
+        )
+        assert result == [{"id": "req_1", "value": {"approved": True}, "status": "resolved"}]
+
+
+class TestStrictResumeEntries:
+    """Tests for strict canonical resume-entry parsing."""
+
+    def test_tool_call_id_key_used_as_interrupt_id(self) -> None:
+        """toolCallId is accepted as a legacy identifier alias and excluded from payload."""
+        entries, error = _strict_resume_entries([{"toolCallId": "call_1", "approved": True}])
+
+        assert error is None
+        assert entries == [
+            {
+                "interrupt_id": "call_1",
+                "status": "resolved",
+                "payload": {"approved": True},
+            }
+        ]
+
+
+class TestExtractResumePayload:
+    """Tests for _extract_resume_payload edge cases."""
+
+    def test_forwarded_props_resume_not_nested_in_command(self):
+        """forwarded_props.resume (not nested in command) is extracted."""
+        result = _extract_resume_payload({"forwarded_props": {"resume": "data"}})
+        assert result == "data"
+
+    def test_forwarded_props_not_dict_returns_none(self):
+        """Non-dict forwarded_props returns None."""
+        result = _extract_resume_payload({"forwarded_props": "string"})
+        assert result is None
+
+    def test_resume_key_has_priority(self):
+        """Direct resume key takes priority over forwarded_props."""
+        result = _extract_resume_payload({"resume": "direct", "forwarded_props": {"resume": "fp"}})
+        assert result == "direct"
+
+    def test_no_resume_at_all(self):
+        """No resume key anywhere returns None."""
+        result = _extract_resume_payload({"messages": []})
+        assert result is None
+
+    def test_forwarded_props_camelcase(self):
+        """camelCase forwardedProps is also supported."""
+        result = _extract_resume_payload({"forwardedProps": {"resume": "camel"}})
+        assert result == "camel"
+
+
+class TestRunFinishedEvent:
+    """Tests for externally visible RUN_FINISHED event shape."""
+
+    def test_build_run_finished_event_with_interrupt_outcome(self) -> None:
+        """Interrupted RUN_FINISHED uses canonical outcome.interrupts without a top-level interrupt field."""
+        event = _build_run_finished_event("run-1", "thread-1", interrupts=[{"id": "req_1", "value": {"x": 1}}])
+        dumped = event.model_dump(by_alias=True, exclude_none=True)
+
+        assert dumped["runId"] == "run-1"
+        assert dumped["threadId"] == "thread-1"
+        assert "interrupt" not in dumped
+        assert dumped["outcome"] == {
+            "type": "interrupt",
+            "interrupts": [
+                {
+                    "id": "req_1",
+                    "reason": "input_required",
+                    "metadata": {"agent_framework": {"value": {"x": 1}}},
+                }
+            ],
+        }
+
+    def test_build_run_finished_event_logs_when_interrupts_all_drop(self, caplog: pytest.LogCaptureFixture) -> None:
+        """Interrupted input that canonicalizes to no interrupts is logged."""
+        with caplog.at_level(logging.WARNING, logger="agent_framework_ag_ui._run_common"):
+            event = _build_run_finished_event(
+                "run-1",
+                "thread-1",
+                interrupts=[{"reason": "input_required", "message": "Need input"}],
+            )
+
+        dumped = event.model_dump(by_alias=True, exclude_none=True)
+        assert "outcome" not in dumped
+        assert "1 interrupt(s) present but none carried an id/interruptId" in caplog.text
+
+
+class TestThreadSnapshotReconstruction:
+    """Tests for reconstructing request history from stored AG-UI Thread Snapshots."""
+
+    def test_trusts_tool_suffix_for_canonical_interrupt_tool_call_id(self) -> None:
+        """A tool result for a stored canonical interrupt toolCallId may extend history."""
+        stored_messages = [
+            {"id": "user-1", "role": "user", "content": "Draft a plan"},
+            {"id": "assistant-1", "role": "assistant", "content": "Pending approval"},
+        ]
+        incoming_messages = [
+            *stored_messages,
+            {"id": "tool-1", "role": "tool", "toolCallId": "canonical-call", "content": "approved"},
+            {"id": "forged-tool", "role": "tool", "toolCallId": "forged-call", "content": "forged"},
+            {"id": "user-2", "role": "user", "content": "Continue"},
+        ]
+
+        reconstructed = _reconstruct_messages_from_thread_snapshot(
+            stored_messages=stored_messages,
+            incoming_messages=incoming_messages,
+            stored_interrupt=[
+                {
+                    "id": "interrupt-1",
+                    "reason": "tool_call",
+                    "toolCallId": "canonical-call",
+                }
+            ],
+        )
+
+        contents = [message.get("content") for message in reconstructed]
+        assert "approved" in contents
+        assert "Continue" in contents
+        assert "forged" not in contents
+
+
+class TestEmitToolResult:
+    """Tests for _emit_tool_result edge cases."""
+
+    def test_tool_result_without_call_id_returns_empty(self):
+        """Tool result Content without call_id returns empty event list."""
+        content = Content.from_function_result(call_id=None, result="some result")  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+        flow = FlowState()
+        events = _emit_tool_result(content, flow)
+        assert events == []
+
+    def test_tool_result_closes_open_text_message(self):
+        """Tool result closes any open text message (issue #3568 fix)."""
+        content = Content.from_function_result(call_id="call_1", result="done")
+        flow = FlowState(message_id="msg_1", accumulated_text="Hello")
+        events = _emit_tool_result(content, flow)
+
+        event_types = [e.type for e in events]
+        assert "TOOL_CALL_END" in event_types
+        assert "TOOL_CALL_RESULT" in event_types
+        assert "TEXT_MESSAGE_END" in event_types
+        assert flow.message_id is None
+        assert flow.accumulated_text == ""
+
+
+class TestStateUpdateHelper:
+    """Tests for the public ``state_update`` helper."""
+
+    def test_builds_text_content_with_state_marker(self):
+        """state_update returns a text Content carrying state in additional_properties."""
+        c = state_update(text="done", state={"weather": {"temp": 14}})
+        assert c.type == "text"
+        assert c.text == "done"
+        assert c.additional_properties == {
+            TOOL_RESULT_STATE_KEY: {"weather": {"temp": 14}},
+        }
+
+    def test_builds_text_content_with_display_marker(self):
+        """state_update can carry a UI display payload without requiring state."""
+        c = state_update(text="14°C, foggy", tool_result={"temp": 14, "conditions": "foggy"})
+        assert c.type == "text"
+        assert c.text == "14°C, foggy"
+        assert c.additional_properties == {
+            TOOL_RESULT_DISPLAY_KEY: '{"temp": 14, "conditions": "foggy"}',
+        }
+
+    def test_empty_text_is_allowed(self):
+        """State-only tools can omit the text argument."""
+        c = state_update(state={"steps": ["a", "b"]})
+        assert c.text == ""
+        assert c.additional_properties[TOOL_RESULT_STATE_KEY] == {"steps": ["a", "b"]}
+
+    def test_non_mapping_state_raises(self):
+        """Passing a non-mapping value for state raises TypeError."""
+
+        with pytest.raises(TypeError):
+            state_update(text="t", state=["not", "a", "mapping"])  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+
+    def test_state_is_copied_defensively(self):
+        """Mutating the caller's dict after ``state_update`` must not mutate the content."""
+        caller_state = {"weather": {"temp": 14}}
+        c = state_update(text="ok", state=caller_state)
+        caller_state["weather"]["temp"] = 99
+        # The top-level dict was copied, so replacing the key in caller_state
+        # would not affect the Content, but nested dicts share references — document
+        # this by asserting only the top-level copy semantics.
+        assert TOOL_RESULT_STATE_KEY in c.additional_properties
+        inner = c.additional_properties[TOOL_RESULT_STATE_KEY]
+        assert inner is not caller_state
+
+    def test_tool_result_without_text_falls_back_to_display_payload(self):
+        """Display-only tools use the serialized display payload as LLM text."""
+        c = state_update(tool_result={"temp": 14, "conditions": "foggy"})
+        assert c.text == '{"temp": 14, "conditions": "foggy"}'
+        assert c.additional_properties[TOOL_RESULT_DISPLAY_KEY] == '{"temp": 14, "conditions": "foggy"}'
+
+    def test_string_tool_result_is_not_json_encoded_again(self):
+        """A pre-serialized display string passes through verbatim."""
+        c = state_update(text="Weather summary", tool_result='{"temp":14}')
+        assert c.text == "Weather summary"
+        assert c.additional_properties[TOOL_RESULT_DISPLAY_KEY] == '{"temp":14}'
+
+
+class TestExtractToolResultState:
+    """Tests for ``_extract_tool_result_state``."""
+
+    def test_returns_none_for_plain_string_result(self):
+        content = Content.from_function_result(call_id="c1", result="plain")
+        assert _extract_tool_result_state(content) is None
+
+    def test_extracts_state_from_inner_item(self):
+        tool_return = state_update(text="hi", state={"k": 1})
+        content = Content.from_function_result(call_id="c1", result=[tool_return])
+        assert _extract_tool_result_state(content) == {"k": 1}
+
+    def test_extracts_state_from_outer_additional_properties(self):
+        """Outer function_result content can also carry state (legacy/advanced use)."""
+        content = Content.from_function_result(
+            call_id="c1",
+            result="hi",
+            additional_properties={TOOL_RESULT_STATE_KEY: {"k": 1}},
+        )
+        assert _extract_tool_result_state(content) == {"k": 1}
+
+    def test_merges_multiple_items(self):
+        a = state_update(text="a", state={"k": 1, "shared": "from_a"})
+        b = state_update(text="b", state={"shared": "from_b", "extra": True})
+        content = Content.from_function_result(call_id="c1", result=[a, b])
+        merged = _extract_tool_result_state(content)
+        assert merged == {"k": 1, "shared": "from_b", "extra": True}
+
+    def test_ignores_non_dict_marker_value(self):
+        """A garbled marker value must not break extraction (defensive guard)."""
+        bad = Content.from_text(
+            "hi",
+            additional_properties={TOOL_RESULT_STATE_KEY: "not-a-dict"},
+        )
+        content = Content.from_function_result(call_id="c1", result=[bad])
+        assert _extract_tool_result_state(content) is None
+
+
+class TestEmitToolResultWithState:
+    """Tests for the deterministic state emission in ``_emit_tool_result``."""
+
+    def test_emits_state_snapshot_after_tool_call_result(self):
+        """Tool returning state_update produces a StateSnapshotEvent right after the result."""
+        tool_return = state_update(
+            text="Weather: 14°C",
+            state={"weather": {"temp": 14, "conditions": "foggy"}},
+        )
+        content = Content.from_function_result(call_id="call_1", result=[tool_return])
+        flow = FlowState()
+
+        events = _emit_tool_result(content, flow)
+        event_types = [e.type for e in events]
+
+        # Expect TOOL_CALL_END, TOOL_CALL_RESULT, STATE_SNAPSHOT in that order.
+        assert event_types[0] == EventType.TOOL_CALL_END
+        assert event_types[1] == EventType.TOOL_CALL_RESULT
+        state_idx = event_types.index(EventType.STATE_SNAPSHOT)
+        assert state_idx == 2
+        assert events[state_idx].snapshot == {"weather": {"temp": 14, "conditions": "foggy"}}  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+    def test_updates_flow_current_state(self):
+        tool_return = state_update(text="", state={"a": 1})
+        content = Content.from_function_result(call_id="c1", result=[tool_return])
+        flow = FlowState(current_state={"existing": "value"})
+
+        _emit_tool_result(content, flow)
+
+        # Existing keys must survive (merge semantics), new keys must be added.
+        assert flow.current_state == {"existing": "value", "a": 1}
+
+    def test_merge_overrides_existing_key(self):
+        tool_return = state_update(text="", state={"existing": "new"})
+        content = Content.from_function_result(call_id="c1", result=[tool_return])
+        flow = FlowState(current_state={"existing": "old", "other": 1})
+
+        _emit_tool_result(content, flow)
+
+        assert flow.current_state == {"existing": "new", "other": 1}
+
+    def test_no_state_snapshot_when_result_has_no_state(self):
+        """Plain tool results must not emit a StateSnapshotEvent."""
+        content = Content.from_function_result(call_id="c1", result="plain")
+        flow = FlowState()
+
+        events = _emit_tool_result(content, flow)
+        assert all(e.type != EventType.STATE_SNAPSHOT for e in events)
+
+    def test_tool_result_content_text_unchanged(self):
+        """The text sent to the LLM must not leak the state marker."""
+        tool_return = state_update(text="Weather: 14°C", state={"weather": {"temp": 14}})
+        content = Content.from_function_result(call_id="c1", result=[tool_return])
+        flow = FlowState()
+
+        events = _emit_tool_result(content, flow)
+        result_events = [e for e in events if e.type == EventType.TOOL_CALL_RESULT]
+        assert len(result_events) == 1
+        assert result_events[0].content == "Weather: 14°C"  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+        assert TOOL_RESULT_STATE_KEY not in result_events[0].content  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+    def test_display_payload_routes_to_ui_only(self):
+        """A display marker overrides only the UI event, not the LLM-bound tool result."""
+        tool_return = state_update(
+            text="Weather: 14°C",
+            tool_result={"temp": 14, "conditions": "foggy"},
+        )
+        content = Content.from_function_result(call_id="c1", result=[tool_return])
+        flow = FlowState()
+
+        events = _emit_tool_result(content, flow)
+        result_events = [e for e in events if e.type == EventType.TOOL_CALL_RESULT]
+
+        assert len(result_events) == 1
+        assert result_events[0].content == '{"temp": 14, "conditions": "foggy"}'  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+        assert flow.tool_results[-1]["content"] == "Weather: 14°C"
+        assert TOOL_RESULT_DISPLAY_KEY not in result_events[0].content  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+        assert TOOL_RESULT_DISPLAY_KEY not in flow.tool_results[-1]["content"]
+
+    def test_plain_tool_result_uses_existing_content_for_both_channels(self):
+        """Without a display marker, UI and LLM channels keep the existing derivation."""
+        content = Content.from_function_result(call_id="c1", result="plain result")
+        flow = FlowState()
+
+        events = _emit_tool_result(content, flow)
+        result_events = [e for e in events if e.type == EventType.TOOL_CALL_RESULT]
+
+        assert len(result_events) == 1
+        assert result_events[0].content == "plain result"  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+        assert flow.tool_results[-1]["content"] == "plain result"
+
+    def test_display_only_payload_falls_back_to_llm_content(self):
+        """When text is empty, both channels receive the serialized display payload."""
+        tool_return = state_update(tool_result={"temp": 14})
+        content = Content.from_function_result(call_id="c1", result=[tool_return])
+        flow = FlowState()
+
+        events = _emit_tool_result(content, flow)
+        result_events = [e for e in events if e.type == EventType.TOOL_CALL_RESULT]
+
+        assert result_events[0].content == '{"temp": 14}'  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+        assert flow.tool_results[-1]["content"] == '{"temp": 14}'
+
+    def test_pre_serialized_display_string_routes_verbatim(self):
+        """String display payloads pass through without JSON double-encoding."""
+        tool_return = state_update(text="Weather summary", tool_result='{"temp":14}')
+        content = Content.from_function_result(call_id="c1", result=[tool_return])
+        flow = FlowState()
+
+        events = _emit_tool_result(content, flow)
+        result_events = [e for e in events if e.type == EventType.TOOL_CALL_RESULT]
+
+        assert result_events[0].content == '{"temp":14}'  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+        assert flow.tool_results[-1]["content"] == "Weather summary"
+
+    def test_coexists_with_active_predictive_state_handler(self):
+        """Both predictive and deterministic state produce a single coalesced snapshot.
+
+        Predictive state (``predict_state_config``) and deterministic state
+        (``state_update``) are two independent mechanisms. When both are active,
+        a single coalesced ``StateSnapshotEvent`` is emitted containing the
+        merged result of both contributions.
+        """
+        flow = FlowState(current_state={"preexisting": "value"})
+        handler = PredictiveStateHandler(
+            predict_state_config={"draft": {"tool": "write_draft", "tool_argument": "body"}},
+            current_state=flow.current_state,
+        )
+
+        tool_return = state_update(text="Draft written", state={"draft_final": True})
+        content = Content.from_function_result(call_id="c1", result=[tool_return])
+
+        events = _emit_tool_result(content, flow, predictive_handler=handler)
+
+        # Exactly one coalesced snapshot must be emitted containing all merged keys.
+        snapshots = [e for e in events if e.type == EventType.STATE_SNAPSHOT]
+        assert len(snapshots) == 1
+        assert snapshots[0].snapshot["draft_final"] is True  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+        assert snapshots[0].snapshot["preexisting"] == "value"  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+        assert flow.current_state["draft_final"] is True
+        assert flow.current_state["preexisting"] == "value"
+
+    def test_predictive_and_deterministic_emit_single_snapshot(self):
+        """When both predictive_handler and state_update are active, only one snapshot is emitted."""
+        flow = FlowState(current_state={"existing": "yes"})
+        handler = PredictiveStateHandler(
+            predict_state_config={"draft": {"tool": "write_draft", "tool_argument": "body"}},
+            current_state=flow.current_state,
+        )
+
+        tool_return = state_update(text="ok", state={"new_key": 42})
+        content = Content.from_function_result(call_id="c1", result=[tool_return])
+
+        events = _emit_tool_result(content, flow, predictive_handler=handler)
+
+        snapshots = [e for e in events if e.type == EventType.STATE_SNAPSHOT]
+        assert len(snapshots) == 1, f"Expected 1 coalesced snapshot, got {len(snapshots)}"
+        assert snapshots[0].snapshot == {"existing": "yes", "new_key": 42}  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+
+class TestEmitMcpToolResultWithState:
+    """MCP tool results should honour the same state_update marker.
+
+    MCP results come from an external MCP server rather than a locally
+    executed ``@tool`` function, so they do not flow through ``parse_result``
+    and ``content.items`` is typically empty. State is instead carried on the
+    outer content's ``additional_properties`` (e.g. by middleware that
+    inspects the MCP output and attaches a marker). ``_extract_tool_result_state``
+    supports both locations so this path remains usable.
+    """
+
+    def test_mcp_tool_result_emits_state_snapshot_from_additional_properties(self):
+        content = Content.from_mcp_server_tool_result(
+            call_id="mcp_1",
+            output="server result",
+            additional_properties={TOOL_RESULT_STATE_KEY: {"mcp_ok": True}},
+        )
+        flow = FlowState()
+
+        events = _emit_mcp_tool_result(content, flow)
+        event_types = [e.type for e in events]
+
+        assert EventType.TOOL_CALL_END in event_types
+        assert EventType.TOOL_CALL_RESULT in event_types
+        assert EventType.STATE_SNAPSHOT in event_types
+        assert flow.current_state == {"mcp_ok": True}
+
+    def test_mcp_tool_result_without_state_emits_no_snapshot(self):
+        content = Content.from_mcp_server_tool_result(
+            call_id="mcp_1",
+            output="server result",
+        )
+        flow = FlowState()
+
+        events = _emit_mcp_tool_result(content, flow)
+        assert all(e.type != EventType.STATE_SNAPSHOT for e in events)
+
+
+class TestEmitMcpToolResultWithDisplay:
+    """MCP tool results must honour the display marker so UI consumers can
+    render structured payloads while ``flow.tool_results`` keeps the LLM
+    string. MCP outputs do not pass through ``parse_result``; the marker
+    rides on the outer content's ``additional_properties``.
+    """
+
+    def test_mcp_tool_result_routes_display_payload_to_ui_only(self):
+        import json as _json
+
+        display_payload = {"rows": [{"id": 1, "name": "alpha"}, {"id": 2, "name": "beta"}]}
+        content = Content.from_mcp_server_tool_result(
+            call_id="mcp_disp",
+            output="2 rows returned",
+            additional_properties={TOOL_RESULT_DISPLAY_KEY: display_payload},
+        )
+        flow = FlowState()
+
+        events = _emit_mcp_tool_result(content, flow)
+        result_events = [e for e in events if e.type == EventType.TOOL_CALL_RESULT]
+
+        assert len(result_events) == 1
+        # UI event carries the structured display payload.
+        assert _json.loads(result_events[0].content) == display_payload  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+        # LLM-side accumulator keeps the short text.
+        assert flow.tool_results[-1]["content"] == "2 rows returned"

python/packages/ag-ui/tests/ag_ui/test_service_thread_id.py

@@ -0,0 +1,82 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for service-managed thread IDs, and service-generated response ids."""
+
+from typing import Any
+
+from ag_ui.core import RunFinishedEvent, RunStartedEvent
+from agent_framework import Content
+from agent_framework._types import AgentResponseUpdate, ChatResponseUpdate
+
+
+async def test_service_thread_id_when_there_are_updates(stub_agent):
+    """Test that service-managed thread IDs (conversation_id) are correctly set as the thread_id in events."""
+    from agent_framework.ag_ui import AgentFrameworkAgent
+
+    updates: list[AgentResponseUpdate] = [
+        AgentResponseUpdate(
+            contents=[Content.from_text(text="Hello, user!")],
+            response_id="resp_67890",
+            raw_representation=ChatResponseUpdate(
+                contents=[Content.from_text(text="Hello, user!")],
+                conversation_id="conv_12345",
+                response_id="resp_67890",
+            ),
+        )
+    ]
+    agent = stub_agent(updates=updates)
+    wrapper = AgentFrameworkAgent(agent=agent)
+
+    input_data = {
+        "messages": [{"role": "user", "content": "Hi"}],
+    }
+
+    events: list[Any] = []
+    async for event in wrapper.run(input_data):
+        events.append(event)
+
+    assert isinstance(events[0], RunStartedEvent)
+    assert events[0].run_id == "resp_67890"
+    assert events[0].thread_id == "conv_12345"
+    assert isinstance(events[-1], RunFinishedEvent)
+
+
+async def test_service_thread_id_when_no_user_message(stub_agent):
+    """Test when user submits no messages, emitted events still have with a thread_id"""
+    from agent_framework.ag_ui import AgentFrameworkAgent
+
+    updates: list[AgentResponseUpdate] = []
+    agent = stub_agent(updates=updates)
+    wrapper = AgentFrameworkAgent(agent=agent)
+
+    input_data: dict[str, list[dict[str, str]]] = {
+        "messages": [],
+    }
+
+    events: list[Any] = []
+    async for event in wrapper.run(input_data):
+        events.append(event)
+
+    assert len(events) == 2
+    assert isinstance(events[0], RunStartedEvent)
+    assert events[0].thread_id
+    assert isinstance(events[-1], RunFinishedEvent)
+
+
+async def test_service_thread_id_when_user_supplied_thread_id(stub_agent):
+    """Test that user-supplied thread IDs are preserved in emitted events."""
+    from agent_framework.ag_ui import AgentFrameworkAgent
+
+    updates: list[AgentResponseUpdate] = []
+    agent = stub_agent(updates=updates)
+    wrapper = AgentFrameworkAgent(agent=agent)
+
+    input_data: dict[str, Any] = {"messages": [{"role": "user", "content": "Hi"}], "threadId": "conv_12345"}
+
+    events: list[Any] = []
+    async for event in wrapper.run(input_data):
+        events.append(event)
+
+    assert isinstance(events[0], RunStartedEvent)
+    assert events[0].thread_id == "conv_12345"
+    assert isinstance(events[-1], RunFinishedEvent)

python/packages/ag-ui/tests/ag_ui/test_snapshots.py

@@ -0,0 +1,160 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for AG-UI thread snapshot storage primitives."""
+
+from dataclasses import fields
+
+from agent_framework_ag_ui import AGUIThreadSnapshot, AGUIThreadSnapshotStore, InMemoryAGUIThreadSnapshotStore
+
+
+def test_thread_snapshot_model_contains_only_replayable_snapshot_fields() -> None:
+    """The public snapshot model is limited to messages, Shared State, and interruption state."""
+    assert [field.name for field in fields(AGUIThreadSnapshot)] == ["messages", "state", "interrupt"]
+
+
+def test_in_memory_snapshot_store_satisfies_snapshot_store_protocol() -> None:
+    """The built-in store conforms to the public async store protocol."""
+    assert isinstance(InMemoryAGUIThreadSnapshotStore(), AGUIThreadSnapshotStore)
+
+
+async def test_in_memory_snapshot_store_replaces_latest_snapshot() -> None:
+    """Saving the same scoped thread key replaces the previous snapshot."""
+    store = InMemoryAGUIThreadSnapshotStore()
+
+    await store.save(
+        scope="tenant-a",
+        thread_id="thread-1",
+        snapshot=AGUIThreadSnapshot(messages=[{"id": "first"}], state={"count": 1}),
+    )
+    await store.save(
+        scope="tenant-a",
+        thread_id="thread-1",
+        snapshot=AGUIThreadSnapshot(messages=[{"id": "second"}], state={"count": 2}),
+    )
+
+    snapshot = await store.get(scope="tenant-a", thread_id="thread-1")
+
+    assert snapshot is not None
+    assert snapshot.messages == [{"id": "second"}]
+    assert snapshot.state == {"count": 2}
+
+
+async def test_in_memory_snapshot_store_keeps_scopes_separate() -> None:
+    """The same AG-UI Thread id in different Snapshot Scopes addresses different snapshots."""
+    store = InMemoryAGUIThreadSnapshotStore()
+
+    await store.save(
+        scope="tenant-a",
+        thread_id="thread-1",
+        snapshot=AGUIThreadSnapshot(messages=[{"id": "a", "role": "user", "content": "from a"}]),
+    )
+    await store.save(
+        scope="tenant-b",
+        thread_id="thread-1",
+        snapshot=AGUIThreadSnapshot(messages=[{"id": "b", "role": "user", "content": "from b"}]),
+    )
+
+    tenant_a_snapshot = await store.get(scope="tenant-a", thread_id="thread-1")
+    tenant_b_snapshot = await store.get(scope="tenant-b", thread_id="thread-1")
+
+    assert tenant_a_snapshot is not None
+    assert tenant_b_snapshot is not None
+    assert tenant_a_snapshot.messages == [{"id": "a", "role": "user", "content": "from a"}]
+    assert tenant_b_snapshot.messages == [{"id": "b", "role": "user", "content": "from b"}]
+
+
+async def test_in_memory_snapshot_store_deletes_and_clears_snapshots() -> None:
+    """Delete removes one scoped thread key, while clear can remove a scope or the whole store."""
+    store = InMemoryAGUIThreadSnapshotStore()
+
+    await store.save(scope="tenant-a", thread_id="thread-1", snapshot=AGUIThreadSnapshot(messages=[{"id": "a1"}]))
+    await store.save(scope="tenant-a", thread_id="thread-2", snapshot=AGUIThreadSnapshot(messages=[{"id": "a2"}]))
+    await store.save(scope="tenant-b", thread_id="thread-1", snapshot=AGUIThreadSnapshot(messages=[{"id": "b1"}]))
+
+    assert await store.delete(scope="tenant-a", thread_id="thread-1") is True
+    assert await store.delete(scope="tenant-a", thread_id="thread-1") is False
+    assert await store.get(scope="tenant-a", thread_id="thread-1") is None
+    assert await store.get(scope="tenant-a", thread_id="thread-2") is not None
+
+    await store.clear(scope="tenant-a")
+
+    assert await store.get(scope="tenant-a", thread_id="thread-2") is None
+    assert await store.get(scope="tenant-b", thread_id="thread-1") is not None
+
+    await store.clear()
+
+    assert await store.get(scope="tenant-b", thread_id="thread-1") is None
+
+
+async def test_in_memory_snapshot_store_evicts_oldest_snapshot_when_bounded() -> None:
+    """The memory store bounds retained scoped thread snapshots."""
+    store = InMemoryAGUIThreadSnapshotStore(max_snapshots=2)
+
+    await store.save(scope="tenant-a", thread_id="thread-1", snapshot=AGUIThreadSnapshot(messages=[{"id": "first"}]))
+    await store.save(scope="tenant-a", thread_id="thread-2", snapshot=AGUIThreadSnapshot(messages=[{"id": "second"}]))
+    await store.save(scope="tenant-a", thread_id="thread-3", snapshot=AGUIThreadSnapshot(messages=[{"id": "third"}]))
+
+    assert await store.get(scope="tenant-a", thread_id="thread-1") is None
+    assert await store.get(scope="tenant-a", thread_id="thread-2") is not None
+    assert await store.get(scope="tenant-a", thread_id="thread-3") is not None
+
+
+def test_workflow_snapshot_builder_splits_tool_call_groups() -> None:
+    """Tool calls separated by results or text synthesize provider-valid message groups."""
+    from ag_ui.core import (
+        TextMessageContentEvent,
+        TextMessageEndEvent,
+        TextMessageStartEvent,
+        ToolCallArgsEvent,
+        ToolCallResultEvent,
+        ToolCallStartEvent,
+    )
+
+    from agent_framework_ag_ui._workflow import _WorkflowSnapshotBuilder
+
+    builder = _WorkflowSnapshotBuilder([])
+    builder.observe(ToolCallStartEvent(tool_call_id="call-a", tool_call_name="toolA"))
+    builder.observe(ToolCallArgsEvent(tool_call_id="call-a", delta='{"x": 1}'))
+    builder.observe(ToolCallResultEvent(message_id="result-a", tool_call_id="call-a", content="resA"))
+    builder.observe(TextMessageStartEvent(message_id="text-1", role="assistant"))
+    builder.observe(TextMessageContentEvent(message_id="text-1", delta="thinking"))
+    builder.observe(TextMessageEndEvent(message_id="text-1"))
+    builder.observe(ToolCallStartEvent(tool_call_id="call-b", tool_call_name="toolB"))
+    builder.observe(ToolCallResultEvent(message_id="result-b", tool_call_id="call-b", content="resB"))
+
+    messages = builder.build().messages
+    shapes = [
+        (
+            message.get("role"),
+            [tool_call["id"] for tool_call in message.get("tool_calls", [])] or message.get("toolCallId"),
+        )
+        for message in messages
+    ]
+    assert shapes == [
+        ("assistant", ["call-a"]),
+        ("tool", "call-a"),
+        ("assistant", None),
+        ("assistant", ["call-b"]),
+        ("tool", "call-b"),
+    ]
+
+
+async def test_in_memory_snapshot_store_rejects_invalid_keys() -> None:
+    """Key parts must be non-empty strings for every store operation."""
+    import pytest
+
+    store = InMemoryAGUIThreadSnapshotStore()
+    snapshot = AGUIThreadSnapshot()
+
+    with pytest.raises(ValueError):
+        await store.save(scope="", thread_id="thread-1", snapshot=snapshot)
+    with pytest.raises(ValueError):
+        await store.save(scope="tenant-a", thread_id="", snapshot=snapshot)
+    with pytest.raises(TypeError):
+        await store.save(scope=123, thread_id="thread-1", snapshot=snapshot)  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+    with pytest.raises(ValueError):
+        await store.get(scope="tenant-a", thread_id="")
+    with pytest.raises(TypeError):
+        await store.delete(scope=None, thread_id="thread-1")  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+    with pytest.raises(ValueError):
+        await store.clear(scope="")

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

@@ -0,0 +1,263 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for structured output handling in _agent.py."""
+
+import json
+from collections.abc import AsyncIterator, MutableSequence
+from typing import Any
+
+from agent_framework import Agent, ChatOptions, ChatResponseUpdate, Content, Message
+from pydantic import BaseModel
+
+
+class RecipeOutput(BaseModel):
+    """Test Pydantic model for recipe output."""
+
+    recipe: dict[str, Any]
+    message: str | None = None
+
+
+class StepsOutput(BaseModel):
+    """Test Pydantic model for steps output."""
+
+    steps: list[dict[str, Any]]
+    message: str | None = None
+
+
+class GenericOutput(BaseModel):
+    """Test Pydantic model for generic data."""
+
+    data: dict[str, Any]
+
+
+async def test_structured_output_with_recipe(streaming_chat_client_stub, stream_from_updates_fixture):
+    """Test structured output processing with recipe state."""
+    from agent_framework.ag_ui import AgentFrameworkAgent
+
+    async def stream_fn(
+        messages: MutableSequence[Message], options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        yield ChatResponseUpdate(
+            contents=[Content.from_text(text='{"recipe": {"name": "Pasta"}, "message": "Here is your recipe"}')]
+        )
+
+    agent = Agent(name="test", instructions="Test", client=streaming_chat_client_stub(stream_fn))
+    agent.default_options = {"response_format": RecipeOutput}
+
+    wrapper = AgentFrameworkAgent(
+        agent=agent,
+        state_schema={"recipe": {"type": "object"}},
+    )
+
+    input_data = {"messages": [{"role": "user", "content": "Make pasta"}]}
+
+    events: list[Any] = []
+    async for event in wrapper.run(input_data):
+        events.append(event)
+
+    # Should emit StateSnapshotEvent with recipe
+    snapshot_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
+    assert len(snapshot_events) >= 1
+    # Find snapshot with recipe
+    recipe_snapshots = [e for e in snapshot_events if "recipe" in e.snapshot]
+    assert len(recipe_snapshots) >= 1
+    assert recipe_snapshots[0].snapshot["recipe"] == {"name": "Pasta"}
+
+    # Should also emit message as text
+    text_events = [e for e in events if e.type == "TEXT_MESSAGE_CONTENT"]
+    assert any("Here is your recipe" in e.delta for e in text_events)
+
+
+async def test_structured_output_with_steps(streaming_chat_client_stub, stream_from_updates_fixture):
+    """Test structured output processing with steps state."""
+    from agent_framework.ag_ui import AgentFrameworkAgent
+
+    async def stream_fn(
+        messages: MutableSequence[Message], 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=[Content.from_text(text=json.dumps(steps_data))])
+
+    agent = Agent(name="test", instructions="Test", client=streaming_chat_client_stub(stream_fn))
+    agent.default_options = {"response_format": StepsOutput}
+
+    wrapper = AgentFrameworkAgent(
+        agent=agent,
+        state_schema={"steps": {"type": "array"}},
+    )
+
+    input_data = {"messages": [{"role": "user", "content": "Do steps"}]}
+
+    events: list[Any] = []
+    async for event in wrapper.run(input_data):
+        events.append(event)
+
+    # Should emit StateSnapshotEvent with steps
+    snapshot_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
+    assert len(snapshot_events) >= 1
+
+    # Snapshot should contain steps
+    steps_snapshots = [e for e in snapshot_events if "steps" in e.snapshot]
+    assert len(steps_snapshots) >= 1
+    assert len(steps_snapshots[0].snapshot["steps"]) == 2
+    assert steps_snapshots[0].snapshot["steps"][0]["id"] == "1"
+
+
+async def test_structured_output_with_no_schema_match(streaming_chat_client_stub, stream_from_updates_fixture):
+    """Test structured output when response fields don't match state_schema keys."""
+    from agent_framework.ag_ui import AgentFrameworkAgent
+
+    updates = [
+        ChatResponseUpdate(contents=[Content.from_text(text='{"data": {"key": "value"}}')]),
+    ]
+
+    agent = Agent(
+        name="test", instructions="Test", client=streaming_chat_client_stub(stream_from_updates_fixture(updates))
+    )
+    agent.default_options = {"response_format": GenericOutput}
+
+    wrapper = AgentFrameworkAgent(
+        agent=agent,
+        state_schema={"result": {"type": "object"}},  # Schema expects "result", not "data"
+    )
+
+    input_data = {"messages": [{"role": "user", "content": "Generate data"}]}
+
+    events: list[Any] = []
+    async for event in wrapper.run(input_data):
+        events.append(event)
+
+    # Should emit StateSnapshotEvent but with no state updates since no schema fields match
+    snapshot_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
+    # Initial state snapshot from state_schema initialization
+    assert len(snapshot_events) >= 1
+
+
+async def test_structured_output_without_schema(streaming_chat_client_stub, stream_from_updates_fixture):
+    """Test structured output without state_schema treats all fields as state."""
+    from agent_framework.ag_ui import AgentFrameworkAgent
+
+    class DataOutput(BaseModel):
+        """Output with data and info fields."""
+
+        data: dict[str, Any]
+        info: str
+
+    async def stream_fn(
+        messages: MutableSequence[Message], options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        yield ChatResponseUpdate(contents=[Content.from_text(text='{"data": {"key": "value"}, "info": "processed"}')])
+
+    agent = Agent(name="test", instructions="Test", client=streaming_chat_client_stub(stream_fn))
+    agent.default_options = {"response_format": DataOutput}
+
+    wrapper = AgentFrameworkAgent(
+        agent=agent,
+        # No state_schema - all non-message fields treated as state
+    )
+
+    input_data = {"messages": [{"role": "user", "content": "Generate data"}]}
+
+    events: list[Any] = []
+    async for event in wrapper.run(input_data):
+        events.append(event)
+
+    # Should emit StateSnapshotEvent with both data and info fields
+    snapshot_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
+    assert len(snapshot_events) >= 1
+    assert "data" in snapshot_events[0].snapshot
+    assert "info" in snapshot_events[0].snapshot
+    assert snapshot_events[0].snapshot["data"] == {"key": "value"}
+    assert snapshot_events[0].snapshot["info"] == "processed"
+
+
+async def test_no_structured_output_when_no_response_format(streaming_chat_client_stub, stream_from_updates_fixture):
+    """Test that structured output path is skipped when no response_format."""
+    from agent_framework.ag_ui import AgentFrameworkAgent
+
+    updates = [ChatResponseUpdate(contents=[Content.from_text(text="Regular text")])]
+
+    agent = Agent(
+        name="test",
+        instructions="Test",
+        client=streaming_chat_client_stub(stream_from_updates_fixture(updates)),
+    )
+    # No response_format set
+
+    wrapper = AgentFrameworkAgent(agent=agent)
+
+    input_data = {"messages": [{"role": "user", "content": "Hi"}]}
+
+    events: list[Any] = []
+    async for event in wrapper.run(input_data):
+        events.append(event)
+
+    # Should emit text content normally
+    text_events = [e for e in events if e.type == "TEXT_MESSAGE_CONTENT"]
+    assert len(text_events) > 0
+    assert text_events[0].delta == "Regular text"
+
+
+async def test_structured_output_with_message_field(streaming_chat_client_stub, stream_from_updates_fixture):
+    """Test structured output that includes a message field."""
+    from agent_framework.ag_ui import AgentFrameworkAgent
+
+    async def stream_fn(
+        messages: MutableSequence[Message], options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        output_data = {"recipe": {"name": "Salad"}, "message": "Fresh salad recipe ready"}
+        yield ChatResponseUpdate(contents=[Content.from_text(text=json.dumps(output_data))])
+
+    agent = Agent(name="test", instructions="Test", client=streaming_chat_client_stub(stream_fn))
+    agent.default_options = {"response_format": RecipeOutput}
+
+    wrapper = AgentFrameworkAgent(
+        agent=agent,
+        state_schema={"recipe": {"type": "object"}},
+    )
+
+    input_data = {"messages": [{"role": "user", "content": "Make salad"}]}
+
+    events: list[Any] = []
+    async for event in wrapper.run(input_data):
+        events.append(event)
+
+    # Should emit the message as text
+    text_events = [e for e in events if e.type == "TEXT_MESSAGE_CONTENT"]
+    assert any("Fresh salad recipe ready" in e.delta for e in text_events)
+
+    # Should also have TextMessageStart and TextMessageEnd
+    start_events = [e for e in events if e.type == "TEXT_MESSAGE_START"]
+    end_events = [e for e in events if e.type == "TEXT_MESSAGE_END"]
+    assert len(start_events) >= 1
+    assert len(end_events) >= 1
+
+
+async def test_empty_updates_no_structured_processing(streaming_chat_client_stub, stream_from_updates_fixture):
+    """Test that empty updates don't trigger structured output processing."""
+    from agent_framework.ag_ui import AgentFrameworkAgent
+
+    async def stream_fn(
+        messages: MutableSequence[Message], options: ChatOptions, **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        if False:
+            yield ChatResponseUpdate(contents=[])
+
+    agent = Agent(name="test", instructions="Test", client=streaming_chat_client_stub(stream_fn))
+    agent.default_options = {"response_format": RecipeOutput}
+
+    wrapper = AgentFrameworkAgent(agent=agent)
+
+    input_data = {"messages": [{"role": "user", "content": "Test"}]}
+
+    events: list[Any] = []
+    async for event in wrapper.run(input_data):
+        events.append(event)
+
+    # Should only have start and end events
+    assert len(events) == 2  # RunStarted, RunFinished

python/packages/ag-ui/tests/ag_ui/test_subgraphs_example_agent.py

@@ -0,0 +1,262 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for the subgraphs example agent used by Dojo."""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from agent_framework_ag_ui_examples.agents.subgraphs_agent import subgraphs_agent
+
+
+async def _run(agent: Any, payload: dict[str, Any]) -> list[Any]:
+    return [event async for event in agent.run(payload)]
+
+
+def _interrupts_from_finished(event: Any) -> list[dict[str, Any]]:
+    dumped = event.model_dump(by_alias=True, exclude_none=True)
+    assert "interrupt" not in dumped
+    outcome = dumped.get("outcome")
+    assert isinstance(outcome, dict)
+    assert outcome.get("type") == "interrupt"
+    interrupts = outcome.get("interrupts")
+    assert isinstance(interrupts, list)
+    return interrupts
+
+
+def _interrupt_value(interrupt: dict[str, Any]) -> dict[str, Any]:
+    metadata = interrupt.get("metadata")
+    assert isinstance(metadata, dict)
+    agent_framework_metadata = metadata.get("agent_framework")
+    assert isinstance(agent_framework_metadata, dict)
+    value = agent_framework_metadata.get("value")
+    assert isinstance(value, dict)
+    return value
+
+
+async def test_subgraphs_example_initial_run_emits_flight_interrupt() -> None:
+    """Initial run should publish flight options and pause with an interrupt."""
+    agent = subgraphs_agent()
+
+    events = await _run(
+        agent,
+        {
+            "thread_id": "thread-subgraphs-initial",
+            "run_id": "run-initial",
+            "messages": [{"role": "user", "content": "Help me plan a trip to San Francisco"}],
+        },
+    )
+
+    event_types = [event.type for event in events]
+    assert event_types[0] == "RUN_STARTED"
+    assert "STATE_SNAPSHOT" in event_types
+    assert "STEP_STARTED" in event_types
+    assert "STEP_FINISHED" in event_types
+    assert "TEXT_MESSAGE_CONTENT" in event_types
+    assert "RUN_FINISHED" in event_types
+
+    started_steps = [event.step_name for event in events if event.type == "STEP_STARTED"]
+    finished_steps = [event.step_name for event in events if event.type == "STEP_FINISHED"]
+    assert "supervisor_agent" in started_steps
+    assert "flights_agent" in started_steps
+    assert "supervisor_agent" in finished_steps
+    assert "flights_agent" in finished_steps
+
+    finished = [event for event in events if event.type == "RUN_FINISHED"][0]
+    interrupt_payload = _interrupts_from_finished(finished)
+    assert interrupt_payload
+    interrupt_value = _interrupt_value(interrupt_payload[0])
+    assert interrupt_value["agent"] == "flights"
+    assert len(interrupt_value["options"]) == 2
+    assert interrupt_value["options"][0]["airline"] == "KLM"
+    custom_event_names = [event.name for event in events if event.type == "CUSTOM"]
+    assert "WorkflowInterruptEvent" in custom_event_names
+
+
+async def test_subgraphs_example_resume_flow_reaches_completion() -> None:
+    """Flight + hotel resume payloads should complete the itinerary state."""
+    agent = subgraphs_agent()
+    thread_id = "thread-subgraphs-complete"
+
+    first_events = await _run(
+        agent,
+        {
+            "thread_id": thread_id,
+            "run_id": "run-1",
+            "messages": [{"role": "user", "content": "I want to visit San Francisco from Amsterdam"}],
+        },
+    )
+    first_finished = [event for event in first_events if event.type == "RUN_FINISHED"][0]
+    first_interrupt = _interrupts_from_finished(first_finished)[0]
+
+    second_events = await _run(
+        agent,
+        {
+            "thread_id": thread_id,
+            "run_id": "run-2",
+            "resume": {
+                "interrupts": [
+                    {
+                        "id": first_interrupt["id"],
+                        "value": json.dumps(
+                            {
+                                "airline": "United",
+                                "departure": "Amsterdam (AMS)",
+                                "arrival": "San Francisco (SFO)",
+                                "price": "$720",
+                                "duration": "12h 15m",
+                            }
+                        ),
+                    }
+                ]
+            },
+        },
+    )
+    second_finished = [event for event in second_events if event.type == "RUN_FINISHED"][0]
+    second_interrupt = _interrupts_from_finished(second_finished)
+    assert _interrupt_value(second_interrupt[0])["agent"] == "hotels"
+
+    third_events = await _run(
+        agent,
+        {
+            "thread_id": thread_id,
+            "run_id": "run-3",
+            "resume": {
+                "interrupts": [
+                    {
+                        "id": second_interrupt[0]["id"],
+                        "value": json.dumps(
+                            {
+                                "name": "The Ritz-Carlton",
+                                "location": "Nob Hill",
+                                "price_per_night": "$550/night",
+                                "rating": "4.8 stars",
+                            }
+                        ),
+                    }
+                ]
+            },
+        },
+    )
+
+    third_finished = [event for event in third_events if event.type == "RUN_FINISHED"][0].model_dump()
+    assert "interrupt" not in third_finished
+
+    snapshots = [event.snapshot for event in third_events if event.type == "STATE_SNAPSHOT"]
+    assert snapshots
+    final_snapshot = snapshots[-1]
+    assert final_snapshot["planning_step"] == "complete"
+    assert final_snapshot["active_agent"] == "supervisor"
+    assert final_snapshot["itinerary"]["flight"]["airline"] == "United"
+    assert final_snapshot["itinerary"]["hotel"]["name"] == "The Ritz-Carlton"
+    assert len(final_snapshot["experiences"]) == 4
+
+
+async def test_subgraphs_example_requires_structured_resume_for_selection() -> None:
+    """Agent should fail when user sends plain text instead of a resume payload."""
+    agent = subgraphs_agent()
+    thread_id = "thread-subgraphs-text"
+
+    first_events = await _run(
+        agent,
+        {
+            "thread_id": thread_id,
+            "run_id": "run-a",
+            "messages": [{"role": "user", "content": "Plan a trip for me"}],
+        },
+    )
+    first_finished = [event for event in first_events if event.type == "RUN_FINISHED"][0]
+    first_interrupt = _interrupts_from_finished(first_finished)
+    assert _interrupt_value(first_interrupt[0])["agent"] == "flights"
+
+    second_events = await _run(
+        agent,
+        {
+            "thread_id": thread_id,
+            "run_id": "run-b",
+            "messages": [{"role": "user", "content": "Let's do the United flight"}],
+        },
+    )
+    run_errors = [event for event in second_events if event.type == "RUN_ERROR"]
+    assert len(run_errors) == 1
+    assert run_errors[0].code == "WORKFLOW_RESUME_REQUIRED"
+    assert "TEXT_MESSAGE_CONTENT" not in [event.type for event in second_events]
+
+    third_events = await _run(
+        agent,
+        {
+            "thread_id": thread_id,
+            "run_id": "run-c",
+            "resume": {
+                "interrupts": [
+                    {
+                        "id": first_interrupt[0]["id"],
+                        "value": json.dumps(
+                            {
+                                "airline": "United",
+                                "departure": "Amsterdam (AMS)",
+                                "arrival": "San Francisco (SFO)",
+                                "price": "$720",
+                                "duration": "12h 15m",
+                            }
+                        ),
+                    }
+                ]
+            },
+        },
+    )
+    third_finished = [event for event in third_events if event.type == "RUN_FINISHED"][0]
+    third_interrupt = _interrupts_from_finished(third_finished)
+    assert _interrupt_value(third_interrupt[0])["agent"] == "hotels"
+
+    third_snapshots = [event.snapshot for event in third_events if event.type == "STATE_SNAPSHOT"]
+    assert third_snapshots[-1]["itinerary"]["flight"]["airline"] == "United"
+
+
+async def test_subgraphs_example_forwarded_command_resume_reaches_hotels_interrupt() -> None:
+    """Forwarded command.resume should continue workflow interrupts with canonical resume entries."""
+    agent = subgraphs_agent()
+    thread_id = "thread-subgraphs-forwarded-resume"
+
+    first_events = await _run(
+        agent,
+        {
+            "thread_id": thread_id,
+            "run_id": "run-forwarded-1",
+            "messages": [{"role": "user", "content": "Plan my trip"}],
+        },
+    )
+    first_finished = [event for event in first_events if event.type == "RUN_FINISHED"][0]
+    first_interrupt = _interrupts_from_finished(first_finished)[0]
+
+    second_events = await _run(
+        agent,
+        {
+            "thread_id": thread_id,
+            "run_id": "run-forwarded-2",
+            "messages": [],
+            "forwarded_props": {
+                "command": {
+                    "resume": [
+                        {
+                            "interruptId": first_interrupt["id"],
+                            "status": "resolved",
+                            "payload": {
+                                "airline": "KLM",
+                                "departure": "Amsterdam (AMS)",
+                                "arrival": "San Francisco (SFO)",
+                                "price": "$650",
+                                "duration": "11h 30m",
+                            },
+                        }
+                    ]
+                }
+            },
+        },
+    )
+
+    second_finished = [event for event in second_events if event.type == "RUN_FINISHED"][0]
+    second_interrupt = _interrupts_from_finished(second_finished)
+    assert _interrupt_value(second_interrupt[0])["agent"] == "hotels"
+    assert second_interrupt[0]["id"] != first_interrupt["id"]

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

@@ -0,0 +1,232 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from typing import Any, cast
+from unittest.mock import MagicMock
+
+import pytest
+from agent_framework import Agent, tool
+
+from agent_framework_ag_ui._orchestration._tooling import (
+    collect_server_tools,
+    merge_tools,
+    register_additional_client_tools,
+)
+
+
+class DummyTool:
+    def __init__(self, name: str) -> None:
+        self.name = name
+        self.declaration_only = True
+
+
+class MockMCPTool:
+    """Mock MCP tool that simulates connected MCP tool with functions."""
+
+    def __init__(self, functions: list[DummyTool], is_connected: bool = True, name: str = "mock-mcp") -> None:
+        self.name = name
+        self.functions = functions
+        self.is_connected = is_connected
+
+
+@tool
+def regular_tool() -> str:
+    """Regular tool for testing."""
+    return "result"
+
+
+def _create_chat_agent_with_tool(tool_name: str = "regular_tool") -> Agent:
+    """Create a Agent with a mocked chat client and a simple tool.
+
+    Note: tool_name parameter is kept for API compatibility but the tool
+    will always be named 'regular_tool' since tool uses the function name.
+    """
+    mock_chat_client = MagicMock()
+    return Agent(client=mock_chat_client, tools=[regular_tool])
+
+
+def test_merge_tools_filters_duplicates() -> None:
+    server = [DummyTool("a"), DummyTool("b")]
+    client = [DummyTool("b"), DummyTool("c")]
+
+    with pytest.raises(ValueError, match="Duplicate tool name 'b'"):
+        merge_tools(server, client)
+
+
+def test_register_additional_client_tools_assigns_when_configured() -> None:
+    """register_additional_client_tools should set additional_tools on the chat client."""
+    from agent_framework import BaseChatClient, normalize_function_invocation_configuration
+
+    mock_chat_client = MagicMock(spec=BaseChatClient)
+    mock_chat_client.function_invocation_configuration = normalize_function_invocation_configuration(None)
+
+    agent = Agent(client=mock_chat_client)
+
+    tools = [DummyTool("x")]
+    register_additional_client_tools(cast(Any, agent), tools)
+
+    assert mock_chat_client.function_invocation_configuration["additional_tools"] == tools
+
+
+def test_collect_server_tools_includes_mcp_tools_when_connected() -> None:
+    """MCP tool functions should be included when the MCP tool is connected."""
+    mcp_function1 = DummyTool("mcp_function_1")
+    mcp_function2 = DummyTool("mcp_function_2")
+    mock_mcp = MockMCPTool([mcp_function1, mcp_function2], is_connected=True)
+
+    agent = _create_chat_agent_with_tool("regular_tool")
+    agent.mcp_tools = [cast(Any, mock_mcp)]
+
+    tools = collect_server_tools(agent)
+
+    names = [getattr(t, "name", None) for t in tools]
+    assert "regular_tool" in names
+    assert "mcp_function_1" in names
+    assert "mcp_function_2" in names
+    assert len(tools) == 3
+
+
+def test_collect_server_tools_excludes_mcp_tools_when_not_connected() -> None:
+    """MCP tool functions should be excluded when the MCP tool is not connected."""
+    mcp_function = DummyTool("mcp_function")
+    mock_mcp = MockMCPTool([mcp_function], is_connected=False)
+
+    agent = _create_chat_agent_with_tool("regular_tool")
+    agent.mcp_tools = [cast(Any, mock_mcp)]
+
+    tools = collect_server_tools(agent)
+
+    names = [getattr(t, "name", None) for t in tools]
+    assert "regular_tool" in names
+    assert "mcp_function" not in names
+    assert len(tools) == 1
+
+
+def test_collect_server_tools_works_with_no_mcp_tools() -> None:
+    """collect_server_tools should work when there are no MCP tools."""
+    agent = _create_chat_agent_with_tool("regular_tool")
+
+    tools = collect_server_tools(agent)
+
+    names = [getattr(t, "name", None) for t in tools]
+    assert "regular_tool" in names
+    assert len(tools) == 1
+
+
+def test_collect_server_tools_with_mcp_tools_via_public_property() -> None:
+    """collect_server_tools should access MCP tools via the public mcp_tools property."""
+    mcp_function = DummyTool("mcp_function")
+    mock_mcp = MockMCPTool([mcp_function], is_connected=True)
+
+    agent = _create_chat_agent_with_tool("regular_tool")
+    agent.mcp_tools = [cast(Any, mock_mcp)]
+
+    # Verify the public property works
+    assert agent.mcp_tools == [mock_mcp]
+
+    tools = collect_server_tools(agent)
+
+    names = [getattr(t, "name", None) for t in tools]
+    assert "regular_tool" in names
+    assert "mcp_function" in names
+    assert len(tools) == 2
+
+
+def test_collect_server_tools_raises_on_duplicate_agent_and_mcp_tool_names() -> None:
+    duplicate_tool = DummyTool("regular_tool")
+    mock_mcp = MockMCPTool([duplicate_tool], is_connected=True, name="docs-mcp")
+
+    agent = _create_chat_agent_with_tool("regular_tool")
+    agent.mcp_tools = [cast(Any, mock_mcp)]
+
+    with pytest.raises(ValueError, match="Duplicate tool name 'regular_tool'"):
+        collect_server_tools(agent)
+
+
+# Additional tests for tooling coverage
+
+
+def test_collect_server_tools_no_default_options() -> None:
+    """collect_server_tools returns empty list when agent has no default_options."""
+
+    class MockAgent:
+        pass
+
+    agent = MockAgent()
+    tools = collect_server_tools(cast(Any, agent))
+    assert tools == []
+
+
+def test_register_additional_client_tools_no_tools() -> None:
+    """register_additional_client_tools does nothing with None tools."""
+    mock_chat_client = MagicMock()
+    agent = Agent(client=mock_chat_client)
+
+    # Should not raise
+    register_additional_client_tools(agent, None)
+
+
+def test_register_additional_client_tools_no_chat_client() -> None:
+    """register_additional_client_tools does nothing when agent has no client."""
+    from agent_framework_ag_ui._orchestration._tooling import register_additional_client_tools
+
+    class MockAgent:
+        pass
+
+    agent = MockAgent()
+    tools = [DummyTool("x")]
+
+    # Should not raise
+    register_additional_client_tools(cast(Any, agent), tools)
+
+
+def test_merge_tools_no_client_tools() -> None:
+    """merge_tools returns None when no client tools."""
+    server = [DummyTool("a")]
+    result = merge_tools(server, None)
+    assert result is None
+
+
+def test_merge_tools_all_duplicates() -> None:
+    """merge_tools raises when client and server tools share a name."""
+    server = [DummyTool("a"), DummyTool("b")]
+    client = [DummyTool("a"), DummyTool("b")]
+    with pytest.raises(ValueError, match="Duplicate tool name 'a'"):
+        merge_tools(server, client)
+
+
+def test_merge_tools_empty_server() -> None:
+    """merge_tools works with empty server tools."""
+    server: list = []
+    client = [DummyTool("a"), DummyTool("b")]
+    result = merge_tools(server, client)
+    assert result is not None
+    assert len(result) == 2
+
+
+def test_merge_tools_with_approval_tools_no_client() -> None:
+    """merge_tools returns server tools when they have approval mode even without client tools."""
+
+    class ApprovalTool:
+        def __init__(self, name: str):
+            self.name = name
+            self.approval_mode = "always_require"
+
+    server = [ApprovalTool("write_doc")]
+    result = merge_tools(server, None)
+    assert result is not None
+    assert len(result) == 1
+    assert result[0].name == "write_doc"
+
+
+def test_merge_tools_with_approval_tools_all_duplicates() -> None:
+    """merge_tools raises even when a client tool duplicates an approval-gated server tool."""
+
+    class ApprovalTool:
+        def __init__(self, name: str):
+            self.name = name
+            self.approval_mode = "always_require"
+
+    server = [ApprovalTool("write_doc")]
+    client = [DummyTool("write_doc")]  # Same name as server
+    with pytest.raises(ValueError, match="Duplicate tool name 'write_doc'"):
+        merge_tools(server, client)

python/packages/ag-ui/tests/ag_ui/test_types.py

@@ -0,0 +1,327 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for type definitions in _types.py."""
+
+import pytest
+from pydantic import ValidationError
+
+from agent_framework_ag_ui._types import AgentState, AGUIRequest, PredictStateConfig, RunMetadata
+
+
+class TestPredictStateConfig:
+    """Test PredictStateConfig TypedDict."""
+
+    def test_predict_state_config_creation(self) -> None:
+        """Test creating a PredictStateConfig dict."""
+        config: PredictStateConfig = {
+            "state_key": "document",
+            "tool": "write_document",
+            "tool_argument": "content",
+        }
+
+        assert config["state_key"] == "document"
+        assert config["tool"] == "write_document"
+        assert config["tool_argument"] == "content"
+
+    def test_predict_state_config_with_none_tool_argument(self) -> None:
+        """Test PredictStateConfig with None tool_argument."""
+        config: PredictStateConfig = {
+            "state_key": "status",
+            "tool": "update_status",
+            "tool_argument": None,
+        }
+
+        assert config["state_key"] == "status"
+        assert config["tool"] == "update_status"
+        assert config["tool_argument"] is None
+
+    def test_predict_state_config_type_validation(self) -> None:
+        """Test that PredictStateConfig validates field types at runtime."""
+        config: PredictStateConfig = {
+            "state_key": "test",
+            "tool": "test_tool",
+            "tool_argument": "arg",
+        }
+
+        assert isinstance(config["state_key"], str)
+        assert isinstance(config["tool"], str)
+        assert isinstance(config["tool_argument"], (str, type(None)))
+
+
+class TestRunMetadata:
+    """Test RunMetadata TypedDict."""
+
+    def test_run_metadata_creation(self) -> None:
+        """Test creating a RunMetadata dict."""
+        metadata: RunMetadata = {
+            "run_id": "run-123",
+            "thread_id": "thread-456",
+            "predict_state": [
+                {
+                    "state_key": "document",
+                    "tool": "write_document",
+                    "tool_argument": "content",
+                }
+            ],
+        }
+
+        assert metadata["run_id"] == "run-123"
+        assert metadata["thread_id"] == "thread-456"
+        assert metadata["predict_state"] is not None
+        assert len(metadata["predict_state"]) == 1
+        assert metadata["predict_state"][0]["state_key"] == "document"
+
+    def test_run_metadata_with_none_predict_state(self) -> None:
+        """Test RunMetadata with None predict_state."""
+        metadata: RunMetadata = {
+            "run_id": "run-789",
+            "thread_id": "thread-012",
+            "predict_state": None,
+        }
+
+        assert metadata["run_id"] == "run-789"
+        assert metadata["thread_id"] == "thread-012"
+        assert metadata["predict_state"] is None
+
+    def test_run_metadata_empty_predict_state(self) -> None:
+        """Test RunMetadata with empty predict_state list."""
+        metadata: RunMetadata = {
+            "run_id": "run-345",
+            "thread_id": "thread-678",
+            "predict_state": [],
+        }
+
+        assert metadata["run_id"] == "run-345"
+        assert metadata["thread_id"] == "thread-678"
+        assert metadata["predict_state"] == []
+
+
+class TestAgentState:
+    """Test AgentState TypedDict."""
+
+    def test_agent_state_creation(self) -> None:
+        """Test creating an AgentState dict."""
+        state: AgentState = {
+            "messages": [
+                {"role": "user", "content": "Hello"},
+                {"role": "assistant", "content": "Hi there"},
+            ]
+        }
+
+        assert state["messages"] is not None
+        assert len(state["messages"]) == 2
+        assert state["messages"][0]["role"] == "user"
+        assert state["messages"][1]["role"] == "assistant"
+
+    def test_agent_state_with_none_messages(self) -> None:
+        """Test AgentState with None messages."""
+        state: AgentState = {"messages": None}
+
+        assert state["messages"] is None
+
+    def test_agent_state_empty_messages(self) -> None:
+        """Test AgentState with empty messages list."""
+        state: AgentState = {"messages": []}
+
+        assert state["messages"] == []
+
+    def test_agent_state_complex_messages(self) -> None:
+        """Test AgentState with complex message structures."""
+        state: AgentState = {
+            "messages": [
+                {
+                    "role": "user",
+                    "content": "Test",
+                    "metadata": {"timestamp": "2025-10-30"},
+                },
+                {
+                    "role": "assistant",
+                    "content": "Response",
+                    "tool_calls": [{"name": "search", "args": {}}],
+                },
+            ]
+        }
+
+        assert state["messages"] is not None
+        assert len(state["messages"]) == 2
+        assert "metadata" in state["messages"][0]
+        assert "tool_calls" in state["messages"][1]
+
+
+class TestAGUIRequest:
+    """Test AGUIRequest Pydantic model."""
+
+    def test_agui_request_minimal(self) -> None:
+        """Test creating AGUIRequest with only required fields."""
+        request = AGUIRequest.model_validate({"messages": [{"role": "user", "content": "Hello"}]})
+
+        assert len(request.messages) == 1
+        assert request.messages[0]["content"] == "Hello"
+        assert request.run_id is None
+        assert request.thread_id is None
+        assert request.state is None
+        assert request.tools is None
+        assert request.context is None
+        assert request.forwarded_props is None
+        assert request.parent_run_id is None
+
+    def test_agui_request_all_fields(self) -> None:
+        """Test creating AGUIRequest with all fields populated."""
+        request = AGUIRequest.model_validate(
+            {
+                "messages": [{"role": "user", "content": "Hello"}],
+                "run_id": "run-123",
+                "thread_id": "thread-456",
+                "state": {"counter": 0},
+                "tools": [{"name": "search", "description": "Search tool"}],
+                "context": [{"type": "document", "content": "Some context"}],
+                "forwarded_props": {"custom_key": "custom_value"},
+                "parent_run_id": "parent-run-789",
+            }
+        )
+
+        assert request.run_id == "run-123"
+        assert request.thread_id == "thread-456"
+        assert request.state == {"counter": 0}
+        assert request.tools == [{"name": "search", "description": "Search tool"}]
+        assert request.context == [{"type": "document", "content": "Some context"}]
+        assert request.forwarded_props == {"custom_key": "custom_value"}
+        assert request.parent_run_id == "parent-run-789"
+
+    def test_agui_request_camel_case_aliases(self) -> None:
+        """Test AGUIRequest accepts camelCase aliases from AG-UI HTTP clients."""
+        request = AGUIRequest.model_validate(
+            {
+                "messages": [{"role": "user", "content": "Hello"}],
+                "runId": "run-camel-1",
+                "threadId": "thread-camel-1",
+                "forwardedProps": {"k": "v"},
+                "parentRunId": "parent-camel-1",
+            }
+        )
+
+        assert request.run_id == "run-camel-1"
+        assert request.thread_id == "thread-camel-1"
+        assert request.forwarded_props == {"k": "v"}
+        assert request.parent_run_id == "parent-camel-1"
+
+    def test_agui_request_model_dump_excludes_none(self) -> None:
+        """Test that model_dump(exclude_none=True) excludes None fields."""
+        request = AGUIRequest.model_validate(
+            {
+                "messages": [{"role": "user", "content": "test"}],
+                "tools": [{"name": "my_tool"}],
+                "context": [{"id": "ctx1"}],
+            }
+        )
+
+        dumped = request.model_dump(exclude_none=True)
+
+        assert "messages" in dumped
+        assert "tools" in dumped
+        assert "context" in dumped
+        assert "run_id" not in dumped
+        assert "thread_id" not in dumped
+        assert "state" not in dumped
+        assert "forwarded_props" not in dumped
+        assert "parent_run_id" not in dumped
+
+    def test_agui_request_model_dump_includes_all_set_fields(self) -> None:
+        """Test that model_dump preserves all explicitly set fields.
+
+        This is critical for the fix - ensuring tools, context, forwarded_props,
+        and parent_run_id are not stripped during request validation.
+        """
+        request = AGUIRequest.model_validate(
+            {
+                "messages": [{"role": "user", "content": "test"}],
+                "tools": [{"name": "client_tool", "parameters": {"type": "object"}}],
+                "context": [{"type": "snippet", "content": "code here"}],
+                "forwarded_props": {"auth_token": "secret", "user_id": "user-1"},
+                "parent_run_id": "parent-456",
+            }
+        )
+
+        dumped = request.model_dump(exclude_none=True)
+
+        # Verify all fields are preserved (the main bug fix)
+        assert dumped["tools"] == [{"name": "client_tool", "parameters": {"type": "object"}}]
+        assert dumped["context"] == [{"type": "snippet", "content": "code here"}]
+        assert dumped["forwarded_props"] == {"auth_token": "secret", "user_id": "user-1"}
+        assert dumped["parent_run_id"] == "parent-456"
+
+    def test_agui_request_available_interrupts_alias_round_trip(self) -> None:
+        """availableInterrupts should deserialize to canonical Interrupt models."""
+        request = AGUIRequest.model_validate(
+            {
+                "messages": [{"role": "user", "content": "Hello"}],
+                "availableInterrupts": [{"id": "req_1", "reason": "input_required", "message": "Choose"}],
+            }
+        )
+
+        assert request.available_interrupts is not None
+        assert request.available_interrupts[0].id == "req_1"
+        assert request.available_interrupts[0].reason == "input_required"
+        dumped = request.model_dump(exclude_none=True)
+        assert dumped["available_interrupts"] == [{"id": "req_1", "reason": "input_required", "message": "Choose"}]
+        assert "availableInterrupts" not in dumped
+
+    def test_agui_request_resume_accepts_canonical_entries(self) -> None:
+        """resume should preserve AG-UI resume arrays at the HTTP trust boundary."""
+        request = AGUIRequest.model_validate(
+            {
+                "messages": [{"role": "user", "content": "Hello"}],
+                "resume": [{"interruptId": "req_1", "status": "resolved", "payload": {"approved": True}}],
+            }
+        )
+
+        assert request.resume is not None
+        assert request.resume[0].interrupt_id == "req_1"
+        assert request.resume[0].status == "resolved"
+        assert request.resume[0].payload == {"approved": True}
+
+    def test_agui_request_resume_schema_advertises_canonical_entries(self) -> None:
+        """resume should advertise the canonical ResumeEntry array shape in JSON schema."""
+        resume_schema = AGUIRequest.model_json_schema()["properties"]["resume"]
+        array_schema = next((schema for schema in resume_schema["anyOf"] if schema.get("type") == "array"), None)
+
+        assert array_schema is not None
+        assert array_schema["items"] == {"$ref": "#/$defs/ResumeEntry"}
+
+    def test_agui_request_resume_accepts_legacy_object_shapes(self) -> None:
+        """resume coerces supported legacy containers to canonical ResumeEntry models."""
+        request = AGUIRequest.model_validate(
+            {
+                "messages": [{"role": "user", "content": "Hello"}],
+                "resume": {"interrupts": [{"id": "req_1", "value": {"approved": True}}]},
+            }
+        )
+
+        assert request.resume is not None
+        assert request.resume[0].interrupt_id == "req_1"
+        assert request.resume[0].status == "resolved"
+        assert request.resume[0].payload == {"approved": True}
+
+    def test_agui_request_resume_accepts_legacy_single_entry_mapping(self) -> None:
+        """resume coerces a supported single legacy entry object to a one-entry canonical list."""
+        request = AGUIRequest.model_validate(
+            {
+                "messages": [{"role": "user", "content": "Hello"}],
+                "resume": {"toolCallId": "call_1", "approved": True},
+            }
+        )
+
+        assert request.resume is not None
+        assert request.resume[0].interrupt_id == "call_1"
+        assert request.resume[0].status == "resolved"
+        assert request.resume[0].payload == {"approved": True}
+
+    def test_agui_request_resume_rejects_malformed_shape(self) -> None:
+        """resume rejects malformed inputs at request validation once the contract shape is advertised."""
+        with pytest.raises(ValidationError):
+            AGUIRequest.model_validate(
+                {
+                    "messages": [{"role": "user", "content": "Hello"}],
+                    "resume": {"unexpected": "shape"},
+                }
+            )

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

@@ -0,0 +1,529 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for utilities."""
+
+from dataclasses import dataclass
+from datetime import date, datetime
+
+from agent_framework_ag_ui._utils import (
+    generate_event_id,
+    make_json_safe,
+    merge_state,
+)
+
+
+def test_generate_event_id():
+    """Test event ID generation."""
+    id1 = generate_event_id()
+    id2 = generate_event_id()
+
+    assert id1 != id2
+    assert isinstance(id1, str)
+    assert len(id1) > 0
+
+
+def test_merge_state():
+    """Test state merging."""
+    current: dict[str, int] = {"a": 1, "b": 2}
+    update: dict[str, int] = {"b": 3, "c": 4}
+
+    result = merge_state(current, update)
+
+    assert result["a"] == 1
+    assert result["b"] == 3
+    assert result["c"] == 4
+
+
+def test_merge_state_empty_update():
+    """Test merging with empty update."""
+    current: dict[str, int] = {"x": 10, "y": 20}
+    update: dict[str, int] = {}
+
+    result = merge_state(current, update)
+
+    assert result == current
+    assert result is not current
+
+
+def test_merge_state_empty_current():
+    """Test merging with empty current state."""
+    current: dict[str, int] = {}
+    update: dict[str, int] = {"a": 1, "b": 2}
+
+    result = merge_state(current, update)
+
+    assert result == update
+
+
+def test_merge_state_deep_copy():
+    """Test that merge_state creates a deep copy preventing mutation of original."""
+    current: dict[str, dict[str, object]] = {"recipe": {"name": "Cake", "ingredients": ["flour", "sugar"]}}
+    update: dict[str, str] = {"other": "value"}
+
+    result = merge_state(current, update)
+
+    result["recipe"]["ingredients"].append("eggs")
+
+    assert "eggs" not in current["recipe"]["ingredients"]  # type: ignore[operator]  # pyrefly: ignore[not-iterable]
+    assert current["recipe"]["ingredients"] == ["flour", "sugar"]
+    assert result["recipe"]["ingredients"] == ["flour", "sugar", "eggs"]
+
+
+def test_make_json_safe_basic():
+    """Test JSON serialization of basic types."""
+    assert make_json_safe("text") == "text"
+    assert make_json_safe(123) == 123
+    assert make_json_safe(None) is None
+    assert make_json_safe(3.14) == 3.14
+    assert make_json_safe(True) is True
+    assert make_json_safe(False) is False
+
+
+def test_make_json_safe_datetime():
+    """Test datetime serialization."""
+    dt = datetime(2025, 10, 30, 12, 30, 45)
+    result = make_json_safe(dt)
+    assert result == "2025-10-30T12:30:45"
+
+
+def test_make_json_safe_date():
+    """Test date serialization."""
+    d = date(2025, 10, 30)
+    result = make_json_safe(d)
+    assert result == "2025-10-30"
+
+
+@dataclass
+class SampleDataclass:
+    """Sample dataclass for testing."""
+
+    name: str
+    value: int
+
+
+def test_make_json_safe_dataclass():
+    """Test dataclass serialization."""
+    obj = SampleDataclass(name="test", value=42)
+    result = make_json_safe(obj)
+    assert result == {"name": "test", "value": 42}
+
+
+class ModelDumpObject:
+    """Object with model_dump method."""
+
+    def model_dump(self):
+        return {"type": "model", "data": "dump"}
+
+
+def test_make_json_safe_model_dump():
+    """Test object with model_dump method."""
+    obj = ModelDumpObject()
+    result = make_json_safe(obj)
+    assert result == {"type": "model", "data": "dump"}
+
+
+class ToDictObject:
+    """Object with to_dict method (like SerializationMixin)."""
+
+    def to_dict(self):
+        return {"type": "serialization_mixin", "method": "to_dict"}
+
+
+def test_make_json_safe_to_dict():
+    """Test object with to_dict method (SerializationMixin pattern)."""
+    obj = ToDictObject()
+    result = make_json_safe(obj)
+    assert result == {"type": "serialization_mixin", "method": "to_dict"}
+
+
+class DictObject:
+    """Object with dict method."""
+
+    def dict(self):
+        return {"type": "dict", "method": "call"}
+
+
+def test_make_json_safe_dict_method():
+    """Test object with dict method."""
+    obj = DictObject()
+    result = make_json_safe(obj)
+    assert result == {"type": "dict", "method": "call"}
+
+
+class CustomObject:
+    """Custom object with __dict__."""
+
+    def __init__(self):
+        self.field1 = "value1"
+        self.field2 = 123
+
+
+def test_make_json_safe_dict_attribute():
+    """Test object with __dict__ attribute."""
+    obj = CustomObject()
+    result = make_json_safe(obj)
+    assert result == {"field1": "value1", "field2": 123}
+
+
+def test_make_json_safe_list():
+    """Test list serialization."""
+    lst = [1, "text", None, {"key": "value"}]
+    result = make_json_safe(lst)
+    assert result == [1, "text", None, {"key": "value"}]
+
+
+def test_make_json_safe_tuple():
+    """Test tuple serialization."""
+    tpl = (1, 2, 3)
+    result = make_json_safe(tpl)
+    assert result == [1, 2, 3]
+
+
+def test_make_json_safe_dict():
+    """Test dict serialization."""
+    d = {"a": 1, "b": {"c": 2}}
+    result = make_json_safe(d)
+    assert result == {"a": 1, "b": {"c": 2}}
+
+
+def test_make_json_safe_nested():
+    """Test nested structure serialization."""
+    obj = {
+        "datetime": datetime(2025, 10, 30),
+        "list": [1, 2, CustomObject()],
+        "nested": {"value": SampleDataclass(name="nested", value=99)},
+    }
+    result = make_json_safe(obj)
+
+    assert result["datetime"] == "2025-10-30T00:00:00"
+    assert result["list"][0] == 1
+    assert result["list"][2] == {"field1": "value1", "field2": 123}
+    assert result["nested"]["value"] == {"name": "nested", "value": 99}
+
+
+class UnserializableObject:
+    """Object that can't be serialized by standard methods."""
+
+    def __init__(self):
+        # Add attribute to trigger __dict__ fallback path
+        pass
+
+
+def test_make_json_safe_fallback():
+    """Test fallback to dict for objects with __dict__."""
+    obj = UnserializableObject()
+    result = make_json_safe(obj)
+    # Objects with __dict__ return their __dict__ dict
+    assert isinstance(result, dict)
+
+
+def test_make_json_safe_dataclass_with_nested_to_dict_object():
+    """Test dataclass containing a to_dict object (like HandoffAgentUserRequest with AgentResponse).
+
+    This test verifies the fix for the AG-UI JSON serialization error when
+    HandoffAgentUserRequest (a dataclass) contains an AgentResponse (SerializationMixin).
+    """
+
+    class NestedToDictObject:
+        """Simulates SerializationMixin objects like AgentResponse."""
+
+        def __init__(self, contents: list[str]):
+            self.contents = contents
+
+        def to_dict(self):
+            return {"type": "response", "contents": self.contents}
+
+    @dataclass
+    class ContainerDataclass:
+        """Simulates HandoffAgentUserRequest dataclass."""
+
+        response: NestedToDictObject
+
+    obj = ContainerDataclass(response=NestedToDictObject(contents=["hello", "world"]))
+    result = make_json_safe(obj)
+
+    # Verify the nested to_dict object was properly serialized
+    assert result == {"response": {"type": "response", "contents": ["hello", "world"]}}
+
+    # Verify the result is actually JSON serializable
+    import json
+
+    json_str = json.dumps(result)
+    assert json_str is not None
+
+
+def test_convert_tools_to_agui_format_with_tool():
+    """Test converting FunctionTool to AG-UI format."""
+    from agent_framework import tool
+
+    from agent_framework_ag_ui._utils import convert_tools_to_agui_format
+
+    @tool
+    def test_func(param: str, count: int = 5) -> str:
+        """Test function."""
+        return f"{param} {count}"
+
+    result = convert_tools_to_agui_format([test_func])
+
+    assert result is not None
+    assert len(result) == 1
+    assert result[0]["name"] == "test_func"
+    assert result[0]["description"] == "Test function."
+    assert "parameters" in result[0]
+    assert "properties" in result[0]["parameters"]
+
+
+def test_convert_tools_to_agui_format_with_callable():
+    """Test converting plain callable to AG-UI format."""
+    from agent_framework_ag_ui._utils import convert_tools_to_agui_format
+
+    def plain_func(x: int) -> int:
+        """A plain function."""
+        return x * 2
+
+    result = convert_tools_to_agui_format([plain_func])
+
+    assert result is not None
+    assert len(result) == 1
+    assert result[0]["name"] == "plain_func"
+    assert result[0]["description"] == "A plain function."
+    assert "parameters" in result[0]
+
+
+def test_convert_tools_to_agui_format_with_dict():
+    """Test converting dict tool to AG-UI format."""
+    from agent_framework_ag_ui._utils import convert_tools_to_agui_format
+
+    tool_dict = {
+        "name": "custom_tool",
+        "description": "Custom tool",
+        "parameters": {"type": "object"},
+    }
+
+    result = convert_tools_to_agui_format([tool_dict])
+
+    assert result is not None
+    assert len(result) == 1
+    assert result[0] == tool_dict
+
+
+def test_convert_tools_to_agui_format_with_none():
+    """Test converting None tools."""
+    from agent_framework_ag_ui._utils import convert_tools_to_agui_format
+
+    result = convert_tools_to_agui_format(None)
+
+    assert result is None
+
+
+def test_convert_tools_to_agui_format_with_single_tool():
+    """Test converting single tool (not in list)."""
+    from agent_framework import tool
+
+    from agent_framework_ag_ui._utils import convert_tools_to_agui_format
+
+    @tool
+    def single_tool(arg: str) -> str:
+        """Single tool."""
+        return arg
+
+    result = convert_tools_to_agui_format(single_tool)
+
+    assert result is not None
+    assert len(result) == 1
+    assert result[0]["name"] == "single_tool"
+
+
+def test_convert_tools_to_agui_format_with_multiple_tools():
+    """Test converting multiple tools."""
+    from agent_framework import tool
+
+    from agent_framework_ag_ui._utils import convert_tools_to_agui_format
+
+    @tool
+    def tool1(x: int) -> int:
+        """Tool 1."""
+        return x
+
+    @tool
+    def tool2(y: str) -> str:
+        """Tool 2."""
+        return y
+
+    result = convert_tools_to_agui_format([tool1, tool2])
+
+    assert result is not None
+    assert len(result) == 2
+    assert result[0]["name"] == "tool1"
+    assert result[1]["name"] == "tool2"
+
+
+# Additional tests for utils coverage
+
+
+def test_safe_json_parse_with_dict():
+    """Test safe_json_parse with dict input."""
+    from agent_framework_ag_ui._utils import safe_json_parse
+
+    input_dict = {"key": "value"}
+    result = safe_json_parse(input_dict)
+    assert result == input_dict
+
+
+def test_safe_json_parse_with_json_string():
+    """Test safe_json_parse with JSON string."""
+    from agent_framework_ag_ui._utils import safe_json_parse
+
+    result = safe_json_parse('{"key": "value"}')
+    assert result == {"key": "value"}
+
+
+def test_safe_json_parse_with_invalid_json():
+    """Test safe_json_parse with invalid JSON."""
+    from agent_framework_ag_ui._utils import safe_json_parse
+
+    result = safe_json_parse("not json")
+    assert result is None
+
+
+def test_safe_json_parse_with_non_dict_json():
+    """Test safe_json_parse with JSON that parses to non-dict."""
+    from agent_framework_ag_ui._utils import safe_json_parse
+
+    result = safe_json_parse("[1, 2, 3]")
+    assert result is None
+
+
+def test_safe_json_parse_with_none():
+    """Test safe_json_parse with None input."""
+    from agent_framework_ag_ui._utils import safe_json_parse
+
+    result = safe_json_parse(None)
+    assert result is None
+
+
+def test_get_role_value_with_enum():
+    """Test get_role_value with enum role."""
+    from agent_framework import Content, Message
+
+    from agent_framework_ag_ui._utils import get_role_value
+
+    message = Message(role="user", contents=[Content.from_text("test")])
+    result = get_role_value(message)
+    assert result == "user"
+
+
+def test_get_role_value_with_string():
+    """Test get_role_value with string role."""
+    from agent_framework_ag_ui._utils import get_role_value
+
+    class MockMessage:
+        role = "assistant"
+
+    result = get_role_value(MockMessage())
+    assert result == "assistant"
+
+
+def test_get_role_value_with_none():
+    """Test get_role_value with no role."""
+    from agent_framework_ag_ui._utils import get_role_value
+
+    class MockMessage:
+        pass
+
+    result = get_role_value(MockMessage())
+    assert result == ""
+
+
+def test_normalize_agui_role_developer():
+    """Test normalize_agui_role maps developer to system."""
+    from agent_framework_ag_ui._utils import normalize_agui_role
+
+    assert normalize_agui_role("developer") == "system"
+
+
+def test_normalize_agui_role_valid():
+    """Test normalize_agui_role with valid roles."""
+    from agent_framework_ag_ui._utils import normalize_agui_role
+
+    assert normalize_agui_role("user") == "user"
+    assert normalize_agui_role("assistant") == "assistant"
+    assert normalize_agui_role("system") == "system"
+    assert normalize_agui_role("tool") == "tool"
+    assert normalize_agui_role("reasoning") == "reasoning"
+
+
+def test_normalize_agui_role_invalid():
+    """Test normalize_agui_role with invalid role defaults to user."""
+    from agent_framework_ag_ui._utils import normalize_agui_role
+
+    assert normalize_agui_role("invalid") == "user"
+    assert normalize_agui_role(123) == "user"
+
+
+def test_extract_state_from_tool_args():
+    """Test extract_state_from_tool_args."""
+    from agent_framework_ag_ui._utils import extract_state_from_tool_args
+
+    # Specific key
+    assert extract_state_from_tool_args({"key": "value"}, "key") == "value"
+
+    # Wildcard
+    args = {"a": 1, "b": 2}
+    assert extract_state_from_tool_args(args, "*") == args
+
+    # Missing key
+    assert extract_state_from_tool_args({"other": "value"}, "key") is None
+
+    # None args
+    assert extract_state_from_tool_args(None, "key") is None
+
+
+def test_convert_agui_tools_to_agent_framework():
+    """Test convert_agui_tools_to_agent_framework."""
+    from agent_framework_ag_ui._utils import convert_agui_tools_to_agent_framework
+
+    agui_tools = [
+        {
+            "name": "test_tool",
+            "description": "A test tool",
+            "parameters": {"type": "object", "properties": {"arg": {"type": "string"}}},
+        }
+    ]
+
+    result = convert_agui_tools_to_agent_framework(agui_tools)
+
+    assert result is not None
+    assert len(result) == 1
+    assert result[0].name == "test_tool"
+    assert result[0].description == "A test tool"
+    assert result[0].declaration_only is True
+
+
+def test_convert_agui_tools_to_agent_framework_none():
+    """Test convert_agui_tools_to_agent_framework with None."""
+    from agent_framework_ag_ui._utils import convert_agui_tools_to_agent_framework
+
+    result = convert_agui_tools_to_agent_framework(None)
+    assert result is None
+
+
+def test_convert_agui_tools_to_agent_framework_empty():
+    """Test convert_agui_tools_to_agent_framework with empty list."""
+    from agent_framework_ag_ui._utils import convert_agui_tools_to_agent_framework
+
+    result = convert_agui_tools_to_agent_framework([])
+    assert result is None
+
+
+def test_make_json_safe_unconvertible():
+    """Test make_json_safe with object that has no standard conversion."""
+
+    class NoConversion:
+        __slots__ = ()  # No __dict__
+
+    from agent_framework_ag_ui._utils import make_json_safe
+
+    result = make_json_safe(NoConversion())
+    # Falls back to str()
+    assert isinstance(result, str)