Merge branch 'main' into feature-azure-functions

python/packages/a2a/agent_framework_a2a/_agent.py

@@ -388,6 +388,17 @@ class A2AAgent(BaseAgent):
         if task.artifacts is not None:
             for artifact in task.artifacts:
                 messages.append(self._artifact_to_chat_message(artifact))
+        elif task.history is not None and len(task.history) > 0:
+            # Include the last history item as the agent response
+            history_item = task.history[-1]
+            contents = self._a2a_parts_to_contents(history_item.parts)
+            messages.append(
+                ChatMessage(
+                    role=Role.ASSISTANT if history_item.role == A2ARole.agent else Role.USER,
+                    contents=contents,
+                    raw_representation=history_item,
+                )
+            )
 
         return messages
 

python/packages/a2a/pyproject.toml

@@ -4,7 +4,7 @@ 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.0b251104"
+version = "1.0.0b251111"
 license-files = ["LICENSE"]
 urls.homepage = "https://aka.ms/agent-framework"
 urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
@@ -19,6 +19,7 @@ classifiers = [
   "Programming Language :: Python :: 3.11",
   "Programming Language :: Python :: 3.12",
   "Programming Language :: Python :: 3.13",
+  "Programming Language :: Python :: 3.14",
   "Typing :: Typed",
 ]
 dependencies = [

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,103 @@
+# 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 ChatAgent
+from agent_framework.azure import AzureOpenAIChatClient
+from agent_framework_ag_ui import add_agent_framework_fastapi_endpoint
+
+# Create your agent
+agent = ChatAgent(
+    name="my_agent",
+    instructions="You are a helpful assistant.",
+    chat_client=AzureOpenAIChatClient(
+        endpoint="https://your-resource.openai.azure.com/",
+        deployment_name="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
+```
+
+### Client (Connect to an AG-UI Server)
+
+```python
+import asyncio
+from agent_framework import TextContent
+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_streaming_response("Hello!"):
+            for content in update.contents:
+                if isinstance(content, TextContent):
+                    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 `ChatAgent` for client-side history management
+
+## 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
+
+## 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
+
+## 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,37 @@
+# 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 ._confirmation_strategies import (
+    ConfirmationStrategy,
+    DefaultConfirmationStrategy,
+    DocumentWriterConfirmationStrategy,
+    RecipeConfirmationStrategy,
+    TaskPlannerConfirmationStrategy,
+)
+from ._endpoint import add_agent_framework_fastapi_endpoint
+from ._event_converters import AGUIEventConverter
+from ._http_service import AGUIHttpService
+
+try:
+    __version__ = importlib.metadata.version(__name__)
+except importlib.metadata.PackageNotFoundError:
+    __version__ = "0.0.0"
+
+__all__ = [
+    "AgentFrameworkAgent",
+    "add_agent_framework_fastapi_endpoint",
+    "AGUIChatClient",
+    "AGUIEventConverter",
+    "AGUIHttpService",
+    "ConfirmationStrategy",
+    "DefaultConfirmationStrategy",
+    "TaskPlannerConfirmationStrategy",
+    "RecipeConfirmationStrategy",
+    "DocumentWriterConfirmationStrategy",
+    "__version__",
+]

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

@@ -0,0 +1,160 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""AgentFrameworkAgent wrapper for AG-UI protocol - Clean Architecture."""
+
+from collections.abc import AsyncGenerator
+from typing import Any
+
+from ag_ui.core import BaseEvent
+from agent_framework import AgentProtocol
+
+from ._confirmation_strategies import ConfirmationStrategy, DefaultConfirmationStrategy
+from ._orchestrators import (
+    DefaultOrchestrator,
+    ExecutionContext,
+    HumanInTheLoopOrchestrator,
+    Orchestrator,
+)
+
+
+class AgentConfig:
+    """Configuration for agent wrapper."""
+
+    def __init__(
+        self,
+        state_schema: dict[str, Any] | None = None,
+        predict_state_config: dict[str, dict[str, str]] | None = None,
+        require_confirmation: bool = True,
+    ):
+        """Initialize agent configuration.
+
+        Args:
+            state_schema: Optional state schema for state management
+            predict_state_config: Configuration for predictive state updates
+            require_confirmation: Whether predictive updates require confirmation
+        """
+        self.state_schema = state_schema or {}
+        self.predict_state_config = predict_state_config or {}
+        self.require_confirmation = require_confirmation
+
+
+class AgentFrameworkAgent:
+    """Wraps Agent Framework agents for AG-UI protocol compatibility.
+
+    Translates between Agent Framework's AgentProtocol and AG-UI's event-based
+    protocol. Uses orchestrators to handle different execution flows (standard
+    execution, human-in-the-loop, etc.). Orchestrators are checked in order;
+    the first matching orchestrator handles the request.
+
+    Supports predictive state updates for agentic generative UI, with optional
+    confirmation requirements configurable per use case.
+    """
+
+    def __init__(
+        self,
+        agent: AgentProtocol,
+        name: str | None = None,
+        description: str | None = None,
+        state_schema: dict[str, Any] | None = None,
+        predict_state_config: dict[str, dict[str, str]] | None = None,
+        require_confirmation: bool = True,
+        orchestrators: list[Orchestrator] | None = None,
+        confirmation_strategy: ConfirmationStrategy | 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
+            predict_state_config: Configuration for predictive state updates.
+                Format: {"state_key": {"tool": "tool_name", "tool_argument": "arg_name"}}
+            require_confirmation: Whether predictive updates require confirmation.
+                Set to False for agentic generative UI that updates automatically.
+            orchestrators: Custom orchestrators (auto-configured if None).
+                Orchestrators are checked in order; first match handles the request.
+            confirmation_strategy: Strategy for generating confirmation messages.
+                Defaults to DefaultConfirmationStrategy if None.
+        """
+        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,
+            require_confirmation=require_confirmation,
+        )
+
+        # Configure orchestrators
+        if orchestrators is None:
+            self.orchestrators = self._default_orchestrators()
+        else:
+            self.orchestrators = orchestrators
+
+        # Configure confirmation strategy
+        if confirmation_strategy is None:
+            self.confirmation_strategy: ConfirmationStrategy = DefaultConfirmationStrategy()
+        else:
+            self.confirmation_strategy = confirmation_strategy
+
+    def _default_orchestrators(self) -> list[Orchestrator]:
+        """Create default orchestrator chain.
+
+        Returns:
+            List of orchestrators in priority order. First matching orchestrator
+            handles the request, so order matters.
+        """
+        return [
+            HumanInTheLoopOrchestrator(),  # Handle tool approval responses
+            # Add more specialized orchestrators here as needed
+            DefaultOrchestrator(),  # Fallback: standard agent execution
+        ]
+
+    async def run_agent(
+        self,
+        input_data: dict[str, Any],
+    ) -> AsyncGenerator[BaseEvent, None]:
+        """Run the agent and yield AG-UI events.
+
+        This is the ONLY public method - much simpler than the original 376-line
+        implementation. All orchestration logic has been extracted into dedicated
+        Orchestrator classes.
+
+        The method creates an ExecutionContext with all needed data, then finds
+        the first orchestrator that can handle the request and delegates to it.
+
+        Args:
+            input_data: The AG-UI run input containing messages, state, etc.
+
+        Yields:
+            AG-UI events
+
+        Raises:
+            RuntimeError: If no orchestrator matches (should never happen if
+                DefaultOrchestrator is last in the chain)
+        """
+        # Create execution context with all needed data
+        context = ExecutionContext(
+            input_data=input_data,
+            agent=self.agent,
+            config=self.config,
+            confirmation_strategy=self.confirmation_strategy,
+        )
+
+        # Find matching orchestrator and execute
+        for orchestrator in self.orchestrators:
+            if orchestrator.can_handle(context):
+                async for event in orchestrator.run(context):
+                    yield event
+                return
+
+        # Should never reach here if DefaultOrchestrator is last
+        raise RuntimeError("No orchestrator matched - check configuration")
+
+
+__all__ = [
+    "AgentFrameworkAgent",
+    "AgentConfig",
+]

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

@@ -0,0 +1,407 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""AG-UI Chat Client implementation."""
+
+import json
+import logging
+import uuid
+from collections.abc import AsyncIterable, MutableSequence
+from functools import wraps
+from typing import Any, TypeVar, cast
+
+import httpx
+from agent_framework import (
+    AIFunction,
+    BaseChatClient,
+    ChatMessage,
+    ChatOptions,
+    ChatResponse,
+    ChatResponseUpdate,
+    DataContent,
+    FunctionCallContent,
+)
+from agent_framework._middleware import use_chat_middleware
+from agent_framework._tools import use_function_invocation
+from agent_framework._types import BaseContent, Contents
+from agent_framework.observability import use_observability
+
+from ._event_converters import AGUIEventConverter
+from ._http_service import AGUIHttpService
+from ._message_adapters import agent_framework_messages_to_agui
+from ._utils import convert_tools_to_agui_format
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+class ServerFunctionCallContent(BaseContent):
+    """Wrapper for server function calls to prevent client re-execution.
+
+    All function calls from the remote server are server-side executions.
+    This wrapper prevents @use_function_invocation from trying to execute them again.
+    """
+
+    function_call_content: FunctionCallContent
+
+    def __init__(self, function_call_content: FunctionCallContent) -> None:
+        """Initialize with the function call content."""
+        super().__init__(type="server_function_call")
+        self.function_call_content = function_call_content
+
+
+def _unwrap_server_function_call_contents(contents: MutableSequence[Contents | dict[str, Any]]) -> None:
+    """Replace ServerFunctionCallContent instances with their underlying call content."""
+    for idx, content in enumerate(contents):
+        if isinstance(content, ServerFunctionCallContent):
+            contents[idx] = content.function_call_content  # type: ignore[assignment]
+
+
+TBaseChatClient = TypeVar("TBaseChatClient", bound=type[BaseChatClient])
+
+
+def _apply_server_function_call_unwrap(chat_client: TBaseChatClient) -> TBaseChatClient:
+    """Class decorator that unwraps server-side function calls after tool handling."""
+
+    original_get_streaming_response = chat_client.get_streaming_response
+
+    @wraps(original_get_streaming_response)
+    async def streaming_wrapper(self, *args: Any, **kwargs: Any) -> AsyncIterable[ChatResponseUpdate]:
+        async for update in original_get_streaming_response(self, *args, **kwargs):
+            _unwrap_server_function_call_contents(cast(MutableSequence[Contents | dict[str, Any]], update.contents))
+            yield update
+
+    chat_client.get_streaming_response = streaming_wrapper  # type: ignore[assignment]
+
+    original_get_response = chat_client.get_response
+
+    @wraps(original_get_response)
+    async def response_wrapper(self, *args: Any, **kwargs: Any) -> ChatResponse:
+        response = await original_get_response(self, *args, **kwargs)
+        if response.messages:
+            for message in response.messages:
+                _unwrap_server_function_call_contents(
+                    cast(MutableSequence[Contents | dict[str, Any]], message.contents)
+                )
+        return response
+
+    chat_client.get_response = response_wrapper  # type: ignore[assignment]
+    return chat_client
+
+
+@_apply_server_function_call_unwrap
+@use_function_invocation
+@use_observability
+@use_chat_middleware
+class AGUIChatClient(BaseChatClient):
+    """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
+
+    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 ChatAgent wrapper which will send full message history on each
+        request. However, even with ChatAgent, 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, @use_function_invocation executes it locally
+        4. Both client and server tools work together (hybrid pattern)
+
+        The wrapping ChatAgent's @use_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 ChatAgent (client manages history):
+
+        .. code-block:: python
+
+            from agent_framework import ChatAgent
+            from agent_framework.ag_ui import AGUIChatClient
+
+            client = AGUIChatClient(endpoint="http://localhost:8888/")
+            agent = ChatAgent(name="assistant", client=client)
+            thread = await agent.get_new_thread()
+
+            # ChatAgent automatically maintains history and sends full context
+            response = await agent.run("Hello!", thread=thread)
+            response2 = await agent.run("How are you?", thread=thread)
+
+        Streaming usage:
+
+        .. code-block:: python
+
+            async for update in client.get_streaming_response("Tell me a story"):
+                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)
+    """
+
+    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,
+        **kwargs: Any,
+    ) -> 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
+            **kwargs: Additional arguments passed to BaseChatClient
+        """
+        super().__init__(additional_properties=additional_properties, **kwargs)
+        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) -> "AGUIChatClient":
+        """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 config:
+            return
+        if any(getattr(tool, "name", None) == tool_name for tool in config.additional_tools):
+            return
+
+        placeholder: AIFunction[Any, Any] = AIFunction(
+            name=tool_name,
+            description="Server-managed tool placeholder (AG-UI)",
+            func=None,
+        )
+        config.additional_tools = list(config.additional_tools) + [placeholder]
+        registered: set[str] = getattr(self, "_registered_server_tools", set())
+        registered.add(tool_name)
+        self._registered_server_tools = registered  # type: ignore[attr-defined]
+        from agent_framework._logging import get_logger
+
+        logger = get_logger()
+        logger.debug(f"[AGUIChatClient] Registered server placeholder: {tool_name}")
+
+    def _extract_state_from_messages(
+        self, messages: MutableSequence[ChatMessage]
+    ) -> tuple[list[ChatMessage], 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, DataContent) and content.media_type == "application/json":
+                try:
+                    uri = content.uri
+                    if uri.startswith("data:application/json;base64,"):
+                        import base64
+
+                        encoded_data = uri.split(",", 1)[1]
+                        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:
+                    from agent_framework._logging import get_logger
+
+                    logger = get_logger()
+                    logger.warning(f"Failed to extract state from message: {e}")
+
+        return list(messages), None
+
+    def _convert_messages_to_agui_format(self, messages: list[ChatMessage]) -> list[dict[str, Any]]:
+        """Convert Agent Framework messages to AG-UI format.
+
+        Args:
+            messages: List of ChatMessage objects
+
+        Returns:
+            List of AG-UI formatted message dictionaries
+        """
+        return agent_framework_messages_to_agui(messages)
+
+    def _get_thread_id(self, chat_options: ChatOptions) -> str:
+        """Get or generate thread ID from chat options.
+
+        Args:
+            chat_options: Chat options containing metadata
+
+        Returns:
+            Thread ID string
+        """
+        thread_id = None
+        if chat_options.metadata:
+            thread_id = chat_options.metadata.get("thread_id")
+
+        if not thread_id:
+            thread_id = f"thread_{uuid.uuid4().hex}"
+
+        return thread_id
+
+    async def _inner_get_response(
+        self,
+        *,
+        messages: MutableSequence[ChatMessage],
+        chat_options: ChatOptions,
+        **kwargs: Any,
+    ) -> ChatResponse:
+        """Internal method to get non-streaming response.
+
+        Keyword Args:
+            messages: List of chat messages
+            chat_options: Chat options for the request
+            **kwargs: Additional keyword arguments
+
+        Returns:
+            ChatResponse object
+        """
+        return await ChatResponse.from_chat_response_generator(
+            self._inner_get_streaming_response(
+                messages=messages,
+                chat_options=chat_options,
+                **kwargs,
+            )
+        )
+
+    async def _inner_get_streaming_response(
+        self,
+        *,
+        messages: MutableSequence[ChatMessage],
+        chat_options: ChatOptions,
+        **kwargs: Any,
+    ) -> AsyncIterable[ChatResponseUpdate]:
+        """Internal method to get streaming response.
+
+        Keyword Args:
+            messages: List of chat messages
+            chat_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(chat_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 ChatAgent's @use_function_invocation wrapper
+        agui_tools = convert_tools_to_agui_format(chat_options.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()
+        if chat_options.tools:
+            for tool in chat_options.tools:
+                if hasattr(tool, "name"):
+                    client_tool_set.add(tool.name)  # type: ignore[arg-type]
+        self._last_client_tool_set = client_tool_set  # type: ignore[attr-defined]
+
+        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()
+
+        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,
+        ):
+            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 isinstance(content, FunctionCallContent):
+                        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 @use_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 @use_function_invocation ignores it
+                            logger.debug(f"[AGUIChatClient] Wrapping server tool: {content.name}")
+                            self._register_server_tool_placeholder(content.name)
+                            update.contents[i] = ServerFunctionCallContent(content)  # type: ignore
+
+                yield update

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

@@ -0,0 +1,175 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Confirmation strategies for human-in-the-loop approval flows.
+
+Each agent can provide a custom confirmation strategy to generate domain-specific
+messages when users approve or reject changes/actions.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class ConfirmationStrategy(ABC):
+    """Strategy for generating confirmation messages during human-in-the-loop flows."""
+
+    @abstractmethod
+    def on_approval_accepted(self, steps: list[dict[str, Any]]) -> str:
+        """Generate message when user approves function execution.
+
+        Args:
+            steps: List of approved steps with 'description', 'status', etc.
+
+        Returns:
+            Message to display to user
+        """
+        ...
+
+    @abstractmethod
+    def on_approval_rejected(self, steps: list[dict[str, Any]]) -> str:
+        """Generate message when user rejects function execution.
+
+        Args:
+            steps: List of rejected steps
+
+        Returns:
+            Message to display to user
+        """
+        ...
+
+    @abstractmethod
+    def on_state_confirmed(self) -> str:
+        """Generate message when user confirms predictive state changes.
+
+        Returns:
+            Message to display to user
+        """
+        ...
+
+    @abstractmethod
+    def on_state_rejected(self) -> str:
+        """Generate message when user rejects predictive state changes.
+
+        Returns:
+            Message to display to user
+        """
+        ...
+
+
+class DefaultConfirmationStrategy(ConfirmationStrategy):
+    """Generic confirmation messages suitable for most agents.
+
+    This preserves the original behavior from v1.
+    """
+
+    def on_approval_accepted(self, steps: list[dict[str, Any]]) -> str:
+        """Generate generic approval message with step list."""
+        enabled_steps = [s for s in steps if s.get("status") == "enabled"]
+
+        message_parts = [f"Executing {len(enabled_steps)} approved steps:\n\n"]
+
+        for i, step in enumerate(enabled_steps, 1):
+            message_parts.append(f"{i}. {step['description']}\n")
+
+        message_parts.append("\nAll steps completed successfully!")
+
+        return "".join(message_parts)
+
+    def on_approval_rejected(self, steps: list[dict[str, Any]]) -> str:
+        """Generate generic rejection message."""
+        return "No problem! What would you like me to change about the plan?"
+
+    def on_state_confirmed(self) -> str:
+        """Generate generic state confirmation message."""
+        return "Changes confirmed and applied successfully!"
+
+    def on_state_rejected(self) -> str:
+        """Generate generic state rejection message."""
+        return "No problem! What would you like me to change?"
+
+
+class TaskPlannerConfirmationStrategy(ConfirmationStrategy):
+    """Domain-specific confirmation messages for task planning agents."""
+
+    def on_approval_accepted(self, steps: list[dict[str, Any]]) -> str:
+        """Generate task-specific approval message."""
+        enabled_steps = [s for s in steps if s.get("status") == "enabled"]
+
+        message_parts = ["Executing your requested tasks:\n\n"]
+
+        for i, step in enumerate(enabled_steps, 1):
+            message_parts.append(f"{i}. {step['description']}\n")
+
+        message_parts.append("\nAll tasks completed successfully!")
+
+        return "".join(message_parts)
+
+    def on_approval_rejected(self, steps: list[dict[str, Any]]) -> str:
+        """Generate task-specific rejection message."""
+        return "No problem! Let me revise the plan. What would you like me to change?"
+
+    def on_state_confirmed(self) -> str:
+        """Task planners typically don't use state confirmation."""
+        return "Tasks confirmed and ready to execute!"
+
+    def on_state_rejected(self) -> str:
+        """Task planners typically don't use state confirmation."""
+        return "No problem! How should I adjust the task list?"
+
+
+class RecipeConfirmationStrategy(ConfirmationStrategy):
+    """Domain-specific confirmation messages for recipe agents."""
+
+    def on_approval_accepted(self, steps: list[dict[str, Any]]) -> str:
+        """Generate recipe-specific approval message."""
+        enabled_steps = [s for s in steps if s.get("status") == "enabled"]
+
+        message_parts = ["Updating your recipe:\n\n"]
+
+        for i, step in enumerate(enabled_steps, 1):
+            message_parts.append(f"{i}. {step['description']}\n")
+
+        message_parts.append("\nRecipe updated successfully!")
+
+        return "".join(message_parts)
+
+    def on_approval_rejected(self, steps: list[dict[str, Any]]) -> str:
+        """Generate recipe-specific rejection message."""
+        return "No problem! What ingredients or steps should I change?"
+
+    def on_state_confirmed(self) -> str:
+        """Generate recipe-specific state confirmation message."""
+        return "Recipe changes applied successfully!"
+
+    def on_state_rejected(self) -> str:
+        """Generate recipe-specific state rejection message."""
+        return "No problem! What would you like me to adjust in the recipe?"
+
+
+class DocumentWriterConfirmationStrategy(ConfirmationStrategy):
+    """Domain-specific confirmation messages for document writing agents."""
+
+    def on_approval_accepted(self, steps: list[dict[str, Any]]) -> str:
+        """Generate document-specific approval message."""
+        enabled_steps = [s for s in steps if s.get("status") == "enabled"]
+
+        message_parts = ["Applying your edits:\n\n"]
+
+        for i, step in enumerate(enabled_steps, 1):
+            message_parts.append(f"{i}. {step['description']}\n")
+
+        message_parts.append("\nDocument updated successfully!")
+
+        return "".join(message_parts)
+
+    def on_approval_rejected(self, steps: list[dict[str, Any]]) -> str:
+        """Generate document-specific rejection message."""
+        return "No problem! Which changes should I keep or modify?"
+
+    def on_state_confirmed(self) -> str:
+        """Generate document-specific state confirmation message."""
+        return "Document edits applied!"
+
+    def on_state_rejected(self) -> str:
+        """Generate document-specific state rejection message."""
+        return "No problem! What should I change about the document?"

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

@@ -0,0 +1,94 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""FastAPI endpoint creation for AG-UI agents."""
+
+import logging
+from typing import Any
+
+from ag_ui.encoder import EventEncoder
+from agent_framework import AgentProtocol
+from fastapi import FastAPI, Request
+from fastapi.responses import StreamingResponse
+
+from ._agent import AgentFrameworkAgent
+
+logger = logging.getLogger(__name__)
+
+
+def add_agent_framework_fastapi_endpoint(
+    app: FastAPI,
+    agent: AgentProtocol | AgentFrameworkAgent,
+    path: str = "/",
+    state_schema: dict[str, Any] | None = None,
+    predict_state_config: dict[str, dict[str, str]] | None = None,
+    allow_origins: list[str] | None = None,
+) -> None:
+    """Add an AG-UI endpoint to a FastAPI app.
+
+    Args:
+        app: The FastAPI application
+        agent: The agent to expose (can be raw AgentProtocol or wrapped)
+        path: The endpoint path
+        state_schema: Optional state schema for shared state management
+        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)
+    """
+    if isinstance(agent, AgentProtocol):
+        wrapped_agent = AgentFrameworkAgent(
+            agent=agent,
+            state_schema=state_schema,
+            predict_state_config=predict_state_config,
+        )
+    else:
+        wrapped_agent = agent
+
+    @app.post(path)
+    async def agent_endpoint(request: Request):  # type: ignore[misc]
+        """Handle AG-UI agent requests.
+
+        Note: Function is accessed via FastAPI's decorator registration,
+        despite appearing unused to static analysis.
+        """
+        try:
+            input_data = await request.json()
+            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')}")
+
+            async def event_generator():
+                encoder = EventEncoder()
+                event_count = 0
+                async for event in wrapped_agent.run_agent(input_data):
+                    event_count += 1
+                    logger.debug(f"[{path}] Event {event_count}: {type(event).__name__}")
+
+                    # Log event payload for debugging
+                    if hasattr(event, "model_dump"):
+                        event_data = event.model_dump(exclude_none=True)
+                        logger.debug(f"[{path}] Event payload: {event_data}")
+
+                    encoded = encoder.encode(event)
+                    logger.debug(
+                        f"[{path}] Encoded as: {encoded[:200]}..."
+                        if len(encoded) > 200
+                        else f"[{path}] Encoded as: {encoded}"
+                    )
+                    yield encoded
+                logger.info(f"[{path}] Completed streaming {event_count} events")
+
+            return StreamingResponse(
+                event_generator(),
+                media_type="text/event-stream",
+                headers={
+                    "Cache-Control": "no-cache",
+                    "Connection": "keep-alive",
+                    "X-Accel-Buffering": "no",
+                },
+            )
+        except Exception as e:
+            logger.error(f"Error in agent endpoint: {e}", exc_info=True)
+            return {"error": str(e)}

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

@@ -0,0 +1,209 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Event converter for AG-UI protocol events to Agent Framework types."""
+
+from typing import Any
+
+from agent_framework import (
+    ChatResponseUpdate,
+    ErrorContent,
+    FinishReason,
+    FunctionCallContent,
+    FunctionResultContent,
+    Role,
+    TextContent,
+)
+
+
+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
+
+    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"
+        """
+        event_type = event.get("type", "")
+
+        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)
+
+        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=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=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=Role.ASSISTANT,
+            message_id=self.current_message_id,
+            contents=[TextContent(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 = event.get("toolCallId")
+        self.current_tool_name = event.get("toolName") or event.get("toolCallName") or event.get("tool_call_name")
+        self.accumulated_tool_args = ""
+
+        return ChatResponseUpdate(
+            role=Role.ASSISTANT,
+            contents=[
+                FunctionCallContent(
+                    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:
+        """Handle TOOL_CALL_ARGS event."""
+        delta = event.get("delta", "")
+        self.accumulated_tool_args += delta
+
+        return ChatResponseUpdate(
+            role=Role.ASSISTANT,
+            contents=[
+                FunctionCallContent(
+                    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."""
+        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=Role.TOOL,
+            contents=[
+                FunctionResultContent(
+                    call_id=tool_call_id,
+                    result=result,
+                )
+            ],
+        )
+
+    def _handle_run_finished(self, event: dict[str, Any]) -> ChatResponseUpdate:
+        """Handle RUN_FINISHED event."""
+        return ChatResponseUpdate(
+            role=Role.ASSISTANT,
+            finish_reason=FinishReason.STOP,
+            contents=[],
+            additional_properties={
+                "thread_id": self.thread_id,
+                "run_id": self.run_id,
+            },
+        )
+
+    def _handle_run_error(self, event: dict[str, Any]) -> ChatResponseUpdate:
+        """Handle RUN_ERROR event."""
+        error_message = event.get("message", "Unknown error")
+
+        return ChatResponseUpdate(
+            role=Role.ASSISTANT,
+            finish_reason=FinishReason.CONTENT_FILTER,
+            contents=[
+                ErrorContent(
+                    message=error_message,
+                    error_code="RUN_ERROR",
+                )
+            ],
+            additional_properties={
+                "thread_id": self.thread_id,
+                "run_id": self.run_id,
+            },
+        )

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

@@ -0,0 +1,693 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Event bridge for converting Agent Framework events to AG-UI protocol."""
+
+import json
+import logging
+import re
+from typing import Any
+
+from ag_ui.core import (
+    BaseEvent,
+    CustomEvent,
+    EventType,
+    MessagesSnapshotEvent,
+    RunFinishedEvent,
+    RunStartedEvent,
+    StateDeltaEvent,
+    StateSnapshotEvent,
+    TextMessageContentEvent,
+    TextMessageEndEvent,
+    TextMessageStartEvent,
+    ToolCallArgsEvent,
+    ToolCallEndEvent,
+    ToolCallResultEvent,
+    ToolCallStartEvent,
+)
+from agent_framework import (
+    AgentRunResponseUpdate,
+    FunctionApprovalRequestContent,
+    FunctionCallContent,
+    FunctionResultContent,
+    TextContent,
+)
+
+from ._utils import generate_event_id
+
+logger = logging.getLogger(__name__)
+
+
+class AgentFrameworkEventBridge:
+    """Converts Agent Framework responses to AG-UI events."""
+
+    def __init__(
+        self,
+        run_id: str,
+        thread_id: str,
+        predict_state_config: dict[str, dict[str, str]] | None = None,
+        current_state: dict[str, Any] | None = None,
+        skip_text_content: bool = False,
+        input_messages: list[Any] | None = None,
+        require_confirmation: bool = True,
+    ) -> None:
+        """
+        Initialize the event bridge.
+
+        Args:
+            run_id: The run identifier.
+            thread_id: The thread identifier.
+            predict_state_config: Configuration for predictive state updates.
+                Format: {"state_key": {"tool": "tool_name", "tool_argument": "arg_name"}}
+            current_state: Reference to the current state dict for tracking updates.
+            skip_text_content: If True, skip emitting TextMessageContentEvents (for structured outputs).
+            input_messages: The input messages from the conversation history.
+            require_confirmation: Whether predictive state updates require user confirmation.
+        """
+        self.run_id = run_id
+        self.thread_id = thread_id
+        self.current_message_id: str | None = None
+        self.current_tool_call_id: str | None = None
+        self.current_tool_call_name: str | None = None  # Track the tool name across streaming chunks
+        self.predict_state_config = predict_state_config or {}
+        self.current_state = current_state or {}
+        self.pending_state_updates: dict[str, Any] = {}  # Track updates from tool calls
+        self.skip_text_content = skip_text_content
+        self.require_confirmation = require_confirmation
+
+        # For predictive state updates: accumulate streaming arguments
+        self.streaming_tool_args: str = ""  # Accumulated JSON string
+        self.last_emitted_state: dict[str, Any] = {}  # Track last emitted state to avoid duplicates
+        self.state_delta_count: int = 0  # Counter for sampling log output
+        self.should_stop_after_confirm: bool = False  # Flag to stop run after confirm_changes
+        self.suppressed_summary: str = ""  # Store LLM summary to show after confirmation
+
+        # For MessagesSnapshotEvent: track tool calls and results
+        self.input_messages = input_messages or []
+        self.pending_tool_calls: list[dict[str, Any]] = []  # Track tool calls for assistant message
+        self.tool_results: list[dict[str, Any]] = []  # Track tool results
+
+    async def from_agent_run_update(self, update: AgentRunResponseUpdate) -> list[BaseEvent]:
+        """
+        Convert an AgentRunResponseUpdate to AG-UI events.
+
+        Args:
+            update: The agent run update to convert.
+
+        Returns:
+            List of AG-UI events.
+        """
+        events: list[BaseEvent] = []
+
+        for content in update.contents:
+            if isinstance(content, TextContent):
+                # Skip text content if using structured outputs (it's just the JSON)
+                if self.skip_text_content:
+                    continue
+
+                # Skip text content if we're about to emit confirm_changes
+                # The summary should only appear after user confirms
+                if self.should_stop_after_confirm:
+                    logger.debug("Skipping text content - waiting for confirm_changes response")
+                    # Save the summary text to show after confirmation
+                    self.suppressed_summary += content.text
+                    continue
+
+                if not self.current_message_id:
+                    self.current_message_id = generate_event_id()
+                    start_event = TextMessageStartEvent(
+                        message_id=self.current_message_id,
+                        role="assistant",
+                    )
+                    events.append(start_event)
+
+                event = TextMessageContentEvent(
+                    message_id=self.current_message_id,
+                    delta=content.text,
+                )
+                events.append(event)
+
+            elif isinstance(content, FunctionCallContent):
+                # Log tool calls for debugging
+                if content.name:
+                    logger.debug(f"Tool call: {content.name} (call_id: {content.call_id})")
+
+                if not content.name and not content.call_id and not self.current_tool_call_name:
+                    args_preview = str(content.arguments)[:50] if content.arguments else "None"
+                    logger.warning(f"FunctionCallContent missing name and call_id. Args: {args_preview}")
+
+                # Get or use existing tool call ID - all chunks of same tool call share the same call_id
+                # Important: the first chunk might have name but no call_id yet
+                if content.call_id:
+                    tool_call_id = content.call_id
+                elif self.current_tool_call_id:
+                    tool_call_id = self.current_tool_call_id
+                else:
+                    # Generate a new ID for this tool call
+                    tool_call_id = (
+                        generate_event_id()
+                    )  # Handle streaming tool calls - name comes in first chunk, arguments in subsequent chunks
+                if content.name:
+                    # This is a new tool call or the first chunk with the name
+                    self.current_tool_call_id = tool_call_id
+                    self.current_tool_call_name = content.name
+
+                    tool_start_event = ToolCallStartEvent(
+                        tool_call_id=tool_call_id,
+                        tool_call_name=content.name,
+                        parent_message_id=self.current_message_id,
+                    )
+                    logger.info(f"Emitting ToolCallStartEvent with name='{content.name}', id='{tool_call_id}'")
+                    events.append(tool_start_event)
+
+                    # Track tool call for MessagesSnapshotEvent
+                    # Initialize a new tool call entry
+                    self.pending_tool_calls.append(
+                        {
+                            "id": tool_call_id,
+                            "type": "function",
+                            "function": {
+                                "name": content.name,
+                                "arguments": "",  # Will accumulate as we get argument chunks
+                            },
+                        }
+                    )
+                else:
+                    # Subsequent chunk without name - update our tracked ID if needed
+                    if tool_call_id:
+                        self.current_tool_call_id = tool_call_id
+
+                # Emit arguments if present
+                if content.arguments:
+                    # content.arguments is already a JSON string from the LLM for streaming calls
+                    # For non-streaming it could be a dict, so we need to handle both
+                    if isinstance(content.arguments, str):
+                        delta_str = content.arguments
+                    else:
+                        # If it's a dict, convert to JSON
+                        delta_str = json.dumps(content.arguments)
+
+                    logger.info(f"Emitting ToolCallArgsEvent with delta: {delta_str!r}..., id='{tool_call_id}'")
+                    args_event = ToolCallArgsEvent(
+                        tool_call_id=tool_call_id,
+                        delta=delta_str,
+                    )
+                    events.append(args_event)
+
+                    # Accumulate arguments for MessagesSnapshotEvent
+                    if self.pending_tool_calls:
+                        # Find the matching tool call and append the delta
+                        for tool_call in self.pending_tool_calls:
+                            if tool_call["id"] == tool_call_id:
+                                tool_call["function"]["arguments"] += delta_str
+                                break
+
+                    # Predictive state updates - accumulate streaming arguments and emit deltas
+                    # Use current_tool_call_name since content.name is only present on first chunk
+                    if self.current_tool_call_name and self.predict_state_config:
+                        # Accumulate the argument string
+                        if isinstance(content.arguments, str):
+                            self.streaming_tool_args += content.arguments
+                        else:
+                            self.streaming_tool_args += json.dumps(content.arguments)
+
+                        logger.debug(
+                            f"Predictive state: accumulated {len(self.streaming_tool_args)} chars for tool '{self.current_tool_call_name}'"
+                        )
+
+                        # Try to parse accumulated arguments (may be incomplete JSON)
+                        # We use a lenient approach: try standard parsing first, then try to extract partial values
+                        parsed_args = None
+                        try:
+                            parsed_args = json.loads(self.streaming_tool_args)
+                        except json.JSONDecodeError:
+                            # JSON is incomplete - try to extract partial string values
+                            # For streaming "document" field, we can extract: {"document": "text...
+                            # Look for pattern: {"field": "value (incomplete)
+                            for state_key, config in self.predict_state_config.items():
+                                if config["tool"] == self.current_tool_call_name:
+                                    tool_arg_name = config["tool_argument"]
+
+                                    # Try to extract partial string value for this argument
+                                    # Pattern: "argument_name": "partial text
+                                    pattern = rf'"{re.escape(tool_arg_name)}":\s*"([^"]*)'
+                                    match = re.search(pattern, self.streaming_tool_args)
+
+                                    if match:
+                                        partial_value = match.group(1)
+                                        # Unescape common sequences
+                                        partial_value = (
+                                            partial_value.replace("\\n", "\n").replace('\\"', '"').replace("\\\\", "\\")
+                                        )
+
+                                        # Emit delta if we have new content
+                                        if (
+                                            state_key not in self.last_emitted_state
+                                            or self.last_emitted_state[state_key] != partial_value
+                                        ):
+                                            state_delta_event = StateDeltaEvent(
+                                                delta=[
+                                                    {
+                                                        "op": "replace",
+                                                        "path": f"/{state_key}",
+                                                        "value": partial_value,
+                                                    }
+                                                ],
+                                            )
+
+                                            self.state_delta_count += 1
+                                            if self.state_delta_count % 10 == 1:
+                                                value_preview = (
+                                                    str(partial_value)[:100] + "..."
+                                                    if len(str(partial_value)) > 100
+                                                    else str(partial_value)
+                                                )
+                                                logger.info(
+                                                    f"StateDeltaEvent #{self.state_delta_count} for '{state_key}': "
+                                                    f"op=replace, path=/{state_key}, value={value_preview}"
+                                                )
+                                            elif self.state_delta_count % 100 == 0:
+                                                logger.info(f"StateDeltaEvent #{self.state_delta_count} emitted")
+
+                                            events.append(state_delta_event)
+                                            self.last_emitted_state[state_key] = partial_value
+                                            self.pending_state_updates[state_key] = partial_value
+
+                        # If we successfully parsed complete JSON, process it
+                        if parsed_args:
+                            # Check if this tool matches any predictive state config
+                            for state_key, config in self.predict_state_config.items():
+                                if config["tool"] == self.current_tool_call_name:
+                                    tool_arg_name = config["tool_argument"]
+
+                                    # Extract the state value
+                                    if tool_arg_name == "*":
+                                        state_value = parsed_args
+                                    elif tool_arg_name in parsed_args:
+                                        state_value = parsed_args[tool_arg_name]
+                                    else:
+                                        continue
+
+                                    # Only emit if state has changed from last emission
+                                    if (
+                                        state_key not in self.last_emitted_state
+                                        or self.last_emitted_state[state_key] != state_value
+                                    ):
+                                        # Emit StateDeltaEvent for real-time UI updates (JSON Patch format)
+                                        state_delta_event = StateDeltaEvent(
+                                            delta=[
+                                                {
+                                                    "op": "replace",  # Use replace since field exists in schema
+                                                    "path": f"/{state_key}",  # JSON Pointer path with leading slash
+                                                    "value": state_value,
+                                                }
+                                            ],
+                                        )
+
+                                        # Increment counter and log every 10th emission with sample data
+                                        self.state_delta_count += 1
+                                        if self.state_delta_count % 10 == 1:  # Log 1st, 11th, 21st, etc.
+                                            value_preview = (
+                                                str(state_value)[:100] + "..."
+                                                if len(str(state_value)) > 100
+                                                else str(state_value)
+                                            )
+                                            logger.info(
+                                                f"StateDeltaEvent #{self.state_delta_count} for '{state_key}': "
+                                                f"op=replace, path=/{state_key}, value={value_preview}"
+                                            )
+                                        elif self.state_delta_count % 100 == 0:  # Also log every 100th
+                                            logger.info(f"StateDeltaEvent #{self.state_delta_count} emitted")
+
+                                        events.append(state_delta_event)
+
+                                        # Track what we emitted
+                                        self.last_emitted_state[state_key] = state_value
+                                        self.pending_state_updates[state_key] = state_value
+
+                    # Legacy predictive state check (for when arguments are complete)
+                    if content.name and content.arguments:
+                        parsed_args = content.parse_arguments()
+
+                        if parsed_args:
+                            logger.info(f"Checking predict_state_config: {self.predict_state_config}")
+                            for state_key, config in self.predict_state_config.items():
+                                logger.info(f"Checking state_key='{state_key}', config={config}")
+                                if config["tool"] == content.name:
+                                    tool_arg_name = config["tool_argument"]
+                                    logger.info(
+                                        f"MATCHED tool '{content.name}' for state key '{state_key}', arg='{tool_arg_name}'"
+                                    )
+
+                                    # If tool_argument is "*", use all arguments as the state value
+                                    if tool_arg_name == "*":
+                                        state_value = parsed_args
+                                        logger.info(f"Using all args as state value, keys: {list(state_value.keys())}")
+                                    elif tool_arg_name in parsed_args:
+                                        state_value = parsed_args[tool_arg_name]
+                                        logger.info(f"Using specific arg '{tool_arg_name}' as state value")
+                                    else:
+                                        logger.warning(f"Tool argument '{tool_arg_name}' not found in parsed args")
+                                        continue
+
+                                    # Emit predictive delta (JSON Patch format)
+                                    state_delta_event = StateDeltaEvent(
+                                        delta=[
+                                            {
+                                                "op": "replace",  # Use replace since field exists in schema
+                                                "path": f"/{state_key}",  # JSON Pointer path with leading slash
+                                                "value": state_value,
+                                            }
+                                        ],
+                                    )
+                                    logger.info(
+                                        f"Emitting StateDeltaEvent for key '{state_key}', value type: {type(state_value)}"
+                                    )
+                                    events.append(state_delta_event)
+
+                                    # Track pending update for later snapshot
+                                    self.pending_state_updates[state_key] = state_value
+
+                # Note: ToolCallEndEvent is emitted when we receive FunctionResultContent,
+                # not here during streaming, since we don't know when the stream is complete
+
+            elif isinstance(content, FunctionResultContent):
+                # First emit ToolCallEndEvent to close the tool call
+                if content.call_id:
+                    end_event = ToolCallEndEvent(
+                        tool_call_id=content.call_id,
+                    )
+                    logger.info(f"Emitting ToolCallEndEvent for completed tool call '{content.call_id}'")
+                    events.append(end_event)
+
+                    # Log total StateDeltaEvent count for this tool call
+                    if self.state_delta_count > 0:
+                        logger.info(
+                            f"Tool call '{content.call_id}' complete: emitted {self.state_delta_count} StateDeltaEvents total"
+                        )
+
+                    # Reset streaming accumulator and counter for next tool call
+                    self.streaming_tool_args = ""
+                    self.state_delta_count = 0
+
+                # Tool result - emit ToolCallResultEvent
+                result_message_id = generate_event_id()
+
+                # Preserve structured data for backend tool rendering
+                # Serialize dicts to JSON string, otherwise convert to string
+                if isinstance(content.result, dict):
+                    result_content = json.dumps(content.result)  # type: ignore[arg-type]
+                elif content.result is not None:
+                    result_content = str(content.result)
+                else:
+                    result_content = ""
+
+                result_event = ToolCallResultEvent(
+                    message_id=result_message_id,
+                    tool_call_id=content.call_id,
+                    content=result_content,
+                    role="tool",
+                )
+                events.append(result_event)
+
+                # Track tool result for MessagesSnapshotEvent
+                # AG-UI protocol expects: { role: "tool", toolCallId: ..., content: ... }
+                # Use camelCase for Pydantic's alias_generator=to_camel
+                self.tool_results.append(
+                    {
+                        "id": result_message_id,
+                        "role": "tool",
+                        "toolCallId": content.call_id,
+                        "content": result_content,
+                    }
+                )
+
+                # Emit MessagesSnapshotEvent with the complete conversation including tool calls and results
+                # This is required for CopilotKit's useCopilotAction to detect tool result
+                if self.pending_tool_calls and self.tool_results:
+                    # Import message adapter
+                    from ._message_adapters import agent_framework_messages_to_agui
+
+                    # Build assistant message with tool_calls
+                    assistant_message = {
+                        "id": generate_event_id(),
+                        "role": "assistant",
+                        "tool_calls": self.pending_tool_calls.copy(),  # Copy the accumulated tool calls
+                    }
+
+                    # Convert Agent Framework messages to AG-UI format (adds required 'id' field)
+                    converted_input_messages = agent_framework_messages_to_agui(self.input_messages)
+
+                    # Build complete messages array: input messages + assistant message + tool results
+                    all_messages = converted_input_messages + [assistant_message] + self.tool_results.copy()
+
+                    # Emit MessagesSnapshotEvent using the proper event type
+                    # Note: messages are dict[str, Any] but Pydantic will validate them as Message types
+                    messages_snapshot_event = MessagesSnapshotEvent(
+                        type=EventType.MESSAGES_SNAPSHOT,
+                        messages=all_messages,  # type: ignore[arg-type]
+                    )
+                    logger.info(f"Emitting MessagesSnapshotEvent with {len(all_messages)} messages")
+                    events.append(messages_snapshot_event)
+
+                # After tool execution, emit StateSnapshotEvent if we have pending state updates
+                if self.pending_state_updates:
+                    # Update the current state with pending updates
+                    for key, value in self.pending_state_updates.items():
+                        self.current_state[key] = value
+
+                    # Log the state structure for debugging
+                    logger.info(f"Emitting StateSnapshotEvent with keys: {list(self.current_state.keys())}")
+                    if "recipe" in self.current_state:
+                        recipe = self.current_state["recipe"]
+                        logger.info(
+                            f"Recipe fields: title={recipe.get('title')}, "
+                            f"skill_level={recipe.get('skill_level')}, "
+                            f"ingredients_count={len(recipe.get('ingredients', []))}, "
+                            f"instructions_count={len(recipe.get('instructions', []))}"
+                        )
+
+                    # Emit complete state snapshot
+                    state_snapshot_event = StateSnapshotEvent(
+                        snapshot=self.current_state,
+                    )
+                    events.append(state_snapshot_event)
+
+                    # Check if this was a predictive state update tool (e.g., write_document_local)
+                    # If so, emit a confirm_changes tool call for the UI modal
+                    tool_was_predictive = False
+                    logger.debug(
+                        f"Checking predictive state: current_tool='{self.current_tool_call_name}', "
+                        f"predict_config={list(self.predict_state_config.keys()) if self.predict_state_config else 'None'}"
+                    )
+                    for state_key, config in self.predict_state_config.items():
+                        # Check if this tool call matches a predictive config
+                        # We need to match against self.current_tool_call_name
+                        if self.current_tool_call_name and config["tool"] == self.current_tool_call_name:
+                            logger.info(
+                                f"Tool '{self.current_tool_call_name}' matches predictive config for state key '{state_key}'"
+                            )
+                            tool_was_predictive = True
+                            break
+
+                    if tool_was_predictive and self.require_confirmation:
+                        # Emit confirm_changes tool call sequence
+                        confirm_call_id = generate_event_id()
+
+                        logger.info("Emitting confirm_changes tool call for predictive update")
+
+                        # Track confirm_changes tool call for MessagesSnapshotEvent (so it persists after RUN_FINISHED)
+                        self.pending_tool_calls.append(
+                            {
+                                "id": confirm_call_id,
+                                "type": "function",
+                                "function": {
+                                    "name": "confirm_changes",
+                                    "arguments": "{}",
+                                },
+                            }
+                        )
+
+                        # Start the confirm_changes tool call
+                        confirm_start = ToolCallStartEvent(
+                            tool_call_id=confirm_call_id,
+                            tool_call_name="confirm_changes",
+                        )
+                        events.append(confirm_start)
+
+                        # Empty args for confirm_changes
+                        confirm_args = ToolCallArgsEvent(
+                            tool_call_id=confirm_call_id,
+                            delta="{}",
+                        )
+                        events.append(confirm_args)
+
+                        # End the confirm_changes tool call
+                        confirm_end = ToolCallEndEvent(
+                            tool_call_id=confirm_call_id,
+                        )
+                        events.append(confirm_end)
+
+                        # Emit MessagesSnapshotEvent so confirm_changes persists after RUN_FINISHED
+                        # Import message adapter
+                        from ._message_adapters import agent_framework_messages_to_agui
+
+                        # Build assistant message with pending confirm_changes tool call
+                        assistant_message = {
+                            "id": generate_event_id(),
+                            "role": "assistant",
+                            "tool_calls": self.pending_tool_calls.copy(),  # Includes confirm_changes
+                        }
+
+                        # Convert Agent Framework messages to AG-UI format (adds required 'id' field)
+                        converted_input_messages = agent_framework_messages_to_agui(self.input_messages)
+
+                        # Build complete messages array: input messages + assistant message + any tool results
+                        all_messages = converted_input_messages + [assistant_message] + self.tool_results.copy()
+
+                        # Emit MessagesSnapshotEvent
+                        # Note: messages are dict[str, Any] but Pydantic will validate them as Message types
+                        messages_snapshot_event = MessagesSnapshotEvent(
+                            type=EventType.MESSAGES_SNAPSHOT,
+                            messages=all_messages,  # type: ignore[arg-type]
+                        )
+                        logger.info(
+                            f"Emitting MessagesSnapshotEvent for confirm_changes with {len(all_messages)} messages"
+                        )
+                        events.append(messages_snapshot_event)
+
+                        # Set flag to stop the run after this - we're waiting for user response
+                        self.should_stop_after_confirm = True
+                        logger.info("Set flag to stop run after confirm_changes")
+                    elif tool_was_predictive:
+                        logger.info("Skipping confirm_changes - require_confirmation is False")
+
+                    # Clear pending updates and reset tool name tracker
+                    self.pending_state_updates.clear()
+                    self.last_emitted_state.clear()
+                    self.current_tool_call_name = None  # Reset for next tool call
+
+            elif isinstance(content, FunctionApprovalRequestContent):
+                # Human in the loop - function approval request
+                logger.info("=== FUNCTION APPROVAL REQUEST ===")
+                logger.info(f"  Function: {content.function_call.name}")
+                logger.info(f"  Call ID: {content.function_call.call_id}")
+
+                # Parse the arguments to extract state for predictive UI updates
+                parsed_args = content.function_call.parse_arguments()
+                logger.info(f"  Parsed args keys: {list(parsed_args.keys()) if parsed_args else 'None'}")
+
+                # Check if this matches our predict_state_config and emit state
+                if parsed_args and self.predict_state_config:
+                    logger.info(f"  Checking predict_state_config: {self.predict_state_config}")
+                    for state_key, config in self.predict_state_config.items():
+                        if config["tool"] == content.function_call.name:
+                            tool_arg_name = config["tool_argument"]
+                            logger.info(
+                                f"  MATCHED tool '{content.function_call.name}' for state key '{state_key}', arg='{tool_arg_name}'"
+                            )
+
+                            # Extract the state value
+                            if tool_arg_name == "*":
+                                state_value = parsed_args
+                            elif tool_arg_name in parsed_args:
+                                state_value = parsed_args[tool_arg_name]
+                            else:
+                                logger.warning(f"  Tool argument '{tool_arg_name}' not found in parsed args")
+                                continue
+
+                            # Update current state
+                            self.current_state[state_key] = state_value
+                            logger.info(
+                                f"Emitting StateSnapshotEvent for key '{state_key}', value type: {type(state_value)}"
+                            )
+
+                            # Emit state snapshot
+                            state_snapshot = StateSnapshotEvent(
+                                snapshot=self.current_state,
+                            )
+                            events.append(state_snapshot)
+
+                # The tool call has been streamed already (Start/Args events)
+                # Now we need to close it with an End event before the agent waits for approval
+                if content.function_call.call_id:
+                    end_event = ToolCallEndEvent(
+                        tool_call_id=content.function_call.call_id,
+                    )
+                    logger.info(
+                        f"Emitting ToolCallEndEvent for approval-required tool '{content.function_call.call_id}'"
+                    )
+                    events.append(end_event)
+
+                # Emit custom event for approval request
+                # Note: In AG-UI protocol, the frontend handles interrupts automatically
+                # when it sees a tool call with the configured name (via predict_state_config)
+                # This custom event is for additional metadata if needed
+                approval_event = CustomEvent(
+                    name="function_approval_request",
+                    value={
+                        "id": content.id,
+                        "function_call": {
+                            "call_id": content.function_call.call_id,
+                            "name": content.function_call.name,
+                            "arguments": content.function_call.parse_arguments(),
+                        },
+                    },
+                )
+                logger.info(f"Emitting function_approval_request custom event for '{content.function_call.name}'")
+                events.append(approval_event)
+
+        return events
+
+    def create_run_started_event(self) -> RunStartedEvent:
+        """Create a run started event."""
+        return RunStartedEvent(
+            run_id=self.run_id,
+            thread_id=self.thread_id,
+        )
+
+    def create_run_finished_event(self, result: Any = None) -> RunFinishedEvent:
+        """Create a run finished event."""
+        return RunFinishedEvent(
+            run_id=self.run_id,
+            thread_id=self.thread_id,
+            result=result,
+        )
+
+    def create_message_start_event(self, message_id: str, role: str = "assistant") -> TextMessageStartEvent:
+        """Create a message start event."""
+        return TextMessageStartEvent(
+            message_id=message_id,
+            role=role,  # type: ignore
+        )
+
+    def create_message_end_event(self, message_id: str) -> TextMessageEndEvent:
+        """Create a message end event."""
+        return TextMessageEndEvent(
+            message_id=message_id,
+        )
+
+    def create_state_snapshot_event(self, state: dict[str, Any]) -> StateSnapshotEvent:
+        """Create a state snapshot event.
+
+        Args:
+            state: The complete state snapshot.
+
+        Returns:
+            StateSnapshotEvent.
+        """
+        return StateSnapshotEvent(
+            snapshot=state,
+        )
+
+    def create_state_delta_event(self, delta: list[dict[str, Any]]) -> StateDeltaEvent:
+        """Create a state delta event using JSON Patch format (RFC 6902).
+
+        Args:
+            delta: List of JSON Patch operations.
+
+        Returns:
+            StateDeltaEvent.
+        """
+        return StateDeltaEvent(
+            delta=delta,
+        )

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

@@ -0,0 +1,157 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""HTTP service for AG-UI protocol communication."""
+
+import json
+import logging
+from collections.abc import AsyncIterable
+from typing import Any
+
+import httpx
+
+logger = logging.getLogger(__name__)
+
+
+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,
+    ) -> 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
+
+        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
+
+        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}"
+        )
+
+        # 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/_message_adapters.py

@@ -0,0 +1,291 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Message format conversion between AG-UI and Agent Framework."""
+
+from typing import Any, cast
+
+from agent_framework import (
+    ChatMessage,
+    FunctionApprovalResponseContent,
+    FunctionCallContent,
+    FunctionResultContent,
+    Role,
+    TextContent,
+)
+
+# Role mapping constants
+_AGUI_TO_FRAMEWORK_ROLE = {
+    "user": Role.USER,
+    "assistant": Role.ASSISTANT,
+    "system": Role.SYSTEM,
+}
+
+_FRAMEWORK_TO_AGUI_ROLE = {
+    Role.USER: "user",
+    Role.ASSISTANT: "assistant",
+    Role.SYSTEM: "system",
+}
+
+
+def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[ChatMessage]:
+    """Convert AG-UI messages to Agent Framework format.
+
+    Args:
+        messages: List of AG-UI messages
+
+    Returns:
+        List of Agent Framework ChatMessage objects
+    """
+    result: list[ChatMessage] = []
+    for msg in messages:
+        # Check for backend tool rendering results FIRST (may not have role field)
+        if "actionExecutionId" in msg or "actionName" in msg:
+            # Backend tool rendering - convert to FunctionResultContent
+            from agent_framework import FunctionResultContent
+
+            tool_call_id = msg.get("actionExecutionId", "")
+            result_content = msg.get("result", msg.get("content", ""))
+
+            chat_msg = ChatMessage(
+                role=Role.TOOL,  # Tool results must be tool role
+                contents=[FunctionResultContent(call_id=tool_call_id, result=result_content)],
+            )
+
+            if "id" in msg:
+                chat_msg.message_id = msg["id"]
+
+            result.append(chat_msg)
+            continue
+
+        # If assistant message includes tool calls, convert to FunctionCallContent(s)
+        tool_calls = msg.get("tool_calls") or msg.get("toolCalls")
+        if tool_calls:
+            contents: list[Any] = []
+            # Include any assistant text content if present
+            content_text = msg.get("content")
+            if isinstance(content_text, str) and content_text:
+                contents.append(TextContent(text=content_text))
+            # Convert each tool call entry
+            for tc in tool_calls:
+                if not isinstance(tc, dict):
+                    continue
+                # Cast to typed dict for proper type inference
+                tc_dict = cast(dict[str, Any], tc)
+                tc_type = tc_dict.get("type")
+                if tc_type == "function":
+                    func_data = tc_dict.get("function", {})
+                    func_dict = cast(dict[str, Any], func_data) if isinstance(func_data, dict) else {}
+
+                    call_id = str(tc_dict.get("id", ""))
+                    name = str(func_dict.get("name", ""))
+                    arguments = func_dict.get("arguments")
+
+                    contents.append(
+                        FunctionCallContent(
+                            call_id=call_id,
+                            name=name,
+                            arguments=arguments,
+                        )
+                    )
+            chat_msg = ChatMessage(role=Role.ASSISTANT, contents=contents)
+            if "id" in msg:
+                chat_msg.message_id = msg["id"]
+            result.append(chat_msg)
+            continue
+
+        role_str = msg.get("role", "user")
+
+        # Handle tool result messages (with role="tool")
+        if role_str == "tool":
+            # Check if this is a standard tool result (has tool_call_id or toolCallId)
+            tool_call_id = msg.get("tool_call_id") or msg.get("toolCallId")
+            result_content = msg.get("content", "")
+
+            # Distinguish between backend tool results and approval responses
+            # Approval responses have {"accepted": ...} structure
+            is_approval = False
+            if result_content:
+                import json
+
+                try:
+                    parsed_content = json.loads(result_content)
+                    is_approval = "accepted" in parsed_content
+                except (json.JSONDecodeError, TypeError):
+                    is_approval = False
+
+            # Backend tool results have non-empty content WITHOUT "accepted" field
+            if tool_call_id and result_content and not is_approval:
+                # Tool execution result - convert to FunctionResultContent with correct role
+                from agent_framework import FunctionResultContent
+
+                chat_msg = ChatMessage(
+                    role=Role.TOOL,
+                    contents=[FunctionResultContent(call_id=tool_call_id, result=result_content)],
+                )
+
+                if "id" in msg:
+                    chat_msg.message_id = msg["id"]
+
+                result.append(chat_msg)
+                continue
+            else:
+                # Human-in-the-loop approval response - mark for special handling
+                content = msg.get("content", "")
+                chat_msg = ChatMessage(
+                    role=Role.USER,  # Approval responses are user messages
+                    contents=[TextContent(text=content)],
+                    additional_properties={"is_tool_result": True, "tool_call_id": msg.get("toolCallId", "")},
+                )
+
+                if "id" in msg:
+                    chat_msg.message_id = msg["id"]
+
+                result.append(chat_msg)
+                continue
+
+        role = _AGUI_TO_FRAMEWORK_ROLE.get(role_str, Role.USER)
+
+        # Check if this message contains function approvals
+        if "function_approvals" in msg and msg["function_approvals"]:
+            # Convert function approvals to FunctionApprovalResponseContent
+            approval_contents: list[Any] = []
+            for approval in msg["function_approvals"]:
+                # Create FunctionCallContent with the modified arguments
+                func_call = FunctionCallContent(
+                    call_id=approval.get("call_id", ""),
+                    name=approval.get("name", ""),
+                    arguments=approval.get("arguments", {}),
+                )
+
+                # Create the approval response
+                approval_response = FunctionApprovalResponseContent(
+                    approved=approval.get("approved", True),
+                    id=approval.get("id", ""),
+                    function_call=func_call,
+                )
+                approval_contents.append(approval_response)
+
+            chat_msg = ChatMessage(role=role, contents=approval_contents)  # type: ignore[arg-type]
+        else:
+            # Regular text message
+            content = msg.get("content", "")
+            if isinstance(content, str):
+                chat_msg = ChatMessage(role=role, contents=[TextContent(text=content)])
+            else:
+                chat_msg = ChatMessage(role=role, contents=[TextContent(text=str(content))])
+
+        if "id" in msg:
+            chat_msg.message_id = msg["id"]
+
+        result.append(chat_msg)
+
+    return result
+
+
+def agent_framework_messages_to_agui(messages: list[ChatMessage] | list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """Convert Agent Framework messages to AG-UI format.
+
+    Args:
+        messages: List of Agent Framework ChatMessage objects or AG-UI dicts (already converted)
+
+    Returns:
+        List of AG-UI message dictionaries
+    """
+    from ._utils import generate_event_id
+
+    result: list[dict[str, Any]] = []
+    for msg in messages:
+        # If already a dict (AG-UI format), ensure it has an ID and normalize keys for Pydantic
+        if isinstance(msg, dict):
+            # Always work on a copy to avoid mutating input
+            normalized_msg = msg.copy()
+            # Ensure ID exists
+            if "id" not in normalized_msg:
+                normalized_msg["id"] = generate_event_id()
+            # Normalize tool_call_id to toolCallId for Pydantic's alias_generator=to_camel
+            if normalized_msg.get("role") == "tool":
+                if "tool_call_id" in normalized_msg:
+                    normalized_msg["toolCallId"] = normalized_msg["tool_call_id"]
+                    del normalized_msg["tool_call_id"]
+                elif "toolCallId" not in normalized_msg:
+                    # Tool message missing toolCallId - add empty string to satisfy schema
+                    normalized_msg["toolCallId"] = ""
+            # Always append the normalized copy, not the original
+            result.append(normalized_msg)
+            continue
+
+        # Convert ChatMessage to AG-UI format
+        role = _FRAMEWORK_TO_AGUI_ROLE.get(msg.role, "user")
+
+        content_text = ""
+        tool_calls: list[dict[str, Any]] = []
+        tool_result_call_id: str | None = None
+
+        for content in msg.contents:
+            if isinstance(content, TextContent):
+                content_text += content.text
+            elif isinstance(content, FunctionCallContent):
+                tool_calls.append(
+                    {
+                        "id": content.call_id,
+                        "type": "function",
+                        "function": {
+                            "name": content.name,
+                            "arguments": content.arguments,
+                        },
+                    }
+                )
+            elif isinstance(content, FunctionResultContent):
+                # Tool result content - extract call_id and result
+                tool_result_call_id = content.call_id
+                # Serialize result to string
+                if isinstance(content.result, dict):
+                    import json
+
+                    content_text = json.dumps(content.result)  # type: ignore
+                elif content.result is not None:
+                    content_text = str(content.result)
+
+        agui_msg: dict[str, Any] = {
+            "id": msg.message_id if msg.message_id else generate_event_id(),  # Always include id
+            "role": role,
+            "content": content_text,
+        }
+
+        if tool_calls:
+            agui_msg["tool_calls"] = tool_calls
+
+        # If this is a tool result message, add toolCallId (using camelCase for Pydantic)
+        if tool_result_call_id:
+            agui_msg["toolCallId"] = tool_result_call_id
+            # Tool result messages should have role="tool"
+            agui_msg["role"] = "tool"
+
+        result.append(agui_msg)
+
+    return result
+
+
+def extract_text_from_contents(contents: list[Any]) -> str:
+    """Extract text from Agent Framework contents.
+
+    Args:
+        contents: List of content objects
+
+    Returns:
+        Concatenated text
+    """
+    text_parts: list[str] = []
+    for content in contents:
+        if isinstance(content, TextContent):
+            text_parts.append(content.text)
+        elif hasattr(content, "text"):
+            text_parts.append(content.text)
+    return "".join(text_parts)
+
+
+__all__ = [
+    "agui_messages_to_agent_framework",
+    "agent_framework_messages_to_agui",
+    "extract_text_from_contents",
+]

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

@@ -0,0 +1,493 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Orchestrators for multi-turn agent flows."""
+
+import json
+import logging
+import uuid
+from abc import ABC, abstractmethod
+from collections.abc import AsyncGenerator
+from typing import TYPE_CHECKING, Any
+
+from ag_ui.core import (
+    BaseEvent,
+    RunErrorEvent,
+    TextMessageContentEvent,
+    TextMessageEndEvent,
+    TextMessageStartEvent,
+)
+from agent_framework import AgentProtocol, AgentThread, ChatAgent, TextContent
+
+from ._utils import convert_agui_tools_to_agent_framework, generate_event_id
+
+if TYPE_CHECKING:
+    from ._agent import AgentConfig
+    from ._confirmation_strategies import ConfirmationStrategy
+
+
+logger = logging.getLogger(__name__)
+
+
+class ExecutionContext:
+    """Shared context for orchestrators."""
+
+    def __init__(
+        self,
+        input_data: dict[str, Any],
+        agent: AgentProtocol,
+        config: "AgentConfig",  # noqa: F821
+        confirmation_strategy: "ConfirmationStrategy | None" = None,  # noqa: F821
+    ):
+        """Initialize execution context.
+
+        Args:
+            input_data: AG-UI run input containing messages, state, etc.
+            agent: The Agent Framework agent to execute
+            config: Agent configuration
+            confirmation_strategy: Strategy for generating confirmation messages
+        """
+        self.input_data = input_data
+        self.agent = agent
+        self.config = config
+        self.confirmation_strategy = confirmation_strategy
+
+        # Lazy-loaded properties
+        self._messages = None
+        self._last_message = None
+        self._run_id: str | None = None
+        self._thread_id: str | None = None
+
+    @property
+    def messages(self):
+        """Get converted Agent Framework messages (lazy loaded)."""
+        if self._messages is None:
+            from ._message_adapters import agui_messages_to_agent_framework
+
+            raw = self.input_data.get("messages", [])
+            self._messages = agui_messages_to_agent_framework(raw)
+        return self._messages
+
+    @property
+    def last_message(self):
+        """Get the last message in the conversation (lazy loaded)."""
+        if self._last_message is None and self.messages:
+            self._last_message = self.messages[-1]
+        return self._last_message
+
+    @property
+    def run_id(self) -> str:
+        """Get or generate run ID."""
+        if self._run_id is None:
+            self._run_id = self.input_data.get("run_id") or str(uuid.uuid4())
+        # This should never be None after the if block above, but satisfy type checkers
+        if self._run_id is None:  # pragma: no cover
+            raise RuntimeError("Failed to initialize run_id")
+        return self._run_id
+
+    @property
+    def thread_id(self) -> str:
+        """Get or generate thread ID."""
+        if self._thread_id is None:
+            self._thread_id = self.input_data.get("thread_id") or str(uuid.uuid4())
+        # This should never be None after the if block above, but satisfy type checkers
+        if self._thread_id is None:  # pragma: no cover
+            raise RuntimeError("Failed to initialize thread_id")
+        return self._thread_id
+
+
+class Orchestrator(ABC):
+    """Base orchestrator for agent execution flows."""
+
+    @abstractmethod
+    def can_handle(self, context: ExecutionContext) -> bool:
+        """Determine if this orchestrator handles the current request.
+
+        Args:
+            context: Execution context with input data and agent
+
+        Returns:
+            True if this orchestrator should handle the request
+        """
+        ...
+
+    @abstractmethod
+    async def run(
+        self,
+        context: ExecutionContext,
+    ) -> AsyncGenerator[BaseEvent, None]:
+        """Execute the orchestration and yield events.
+
+        Args:
+            context: Execution context
+
+        Yields:
+            AG-UI events
+        """
+        # This is never executed - just satisfies mypy's requirement for async generators
+        if False:  # pragma: no cover
+            yield
+        raise NotImplementedError
+
+
+class HumanInTheLoopOrchestrator(Orchestrator):
+    """Handles tool approval responses from user."""
+
+    def can_handle(self, context: ExecutionContext) -> bool:
+        """Check if last message is a tool approval response.
+
+        Args:
+            context: Execution context
+
+        Returns:
+            True if last message is a tool result
+        """
+        msg = context.last_message
+        if not msg:
+            return False
+
+        return bool(msg.additional_properties.get("is_tool_result", False))
+
+    async def run(
+        self,
+        context: ExecutionContext,
+    ) -> AsyncGenerator[BaseEvent, None]:
+        """Process approval response and generate confirmation events.
+
+        This implementation is extracted from the legacy _agent.py lines 144-244.
+
+        Args:
+            context: Execution context
+
+        Yields:
+            AG-UI events (TextMessage, RunFinished)
+        """
+        from ._confirmation_strategies import DefaultConfirmationStrategy
+        from ._events import AgentFrameworkEventBridge
+
+        logger.info("=== TOOL RESULT DETECTED (HumanInTheLoopOrchestrator) ===")
+
+        # Create event bridge for run events
+        event_bridge = AgentFrameworkEventBridge(
+            run_id=context.run_id,
+            thread_id=context.thread_id,
+        )
+
+        # CRITICAL: Every AG-UI run must start with RunStartedEvent
+        yield event_bridge.create_run_started_event()
+
+        # Get confirmation strategy (use default if none provided)
+        strategy = context.confirmation_strategy
+        if strategy is None:
+            strategy = DefaultConfirmationStrategy()
+
+        # Parse the tool result content
+        tool_content_text = ""
+        last_message = context.last_message
+        if last_message:
+            for content in last_message.contents:
+                if isinstance(content, TextContent):
+                    tool_content_text = content.text
+                    break
+
+        try:
+            tool_result = json.loads(tool_content_text)
+            accepted = tool_result.get("accepted", False)
+            steps = tool_result.get("steps", [])
+
+            logger.info(f"  Accepted: {accepted}")
+            logger.info(f"  Steps count: {len(steps)}")
+
+            # Emit a text message confirming execution
+            message_id = generate_event_id()
+
+            yield TextMessageStartEvent(message_id=message_id, role="assistant")
+
+            # Check if this is confirm_changes (no steps) or function approval (has steps)
+            if not steps:
+                # This is confirm_changes for predictive state updates
+                if accepted:
+                    confirmation_message = strategy.on_state_confirmed()
+                else:
+                    confirmation_message = strategy.on_state_rejected()
+            elif accepted:
+                # User approved - execute the enabled steps (function approval flow)
+                confirmation_message = strategy.on_approval_accepted(steps)
+            else:
+                # User rejected
+                confirmation_message = strategy.on_approval_rejected(steps)
+
+            yield TextMessageContentEvent(
+                message_id=message_id,
+                delta=confirmation_message,
+            )
+
+            yield TextMessageEndEvent(message_id=message_id)
+
+            # Emit run finished
+            yield event_bridge.create_run_finished_event()
+
+        except json.JSONDecodeError:
+            logger.error(f"Failed to parse tool result: {tool_content_text}")
+            yield RunErrorEvent(message=f"Invalid tool result format: {tool_content_text[:100]}")
+            yield event_bridge.create_run_finished_event()
+
+
+class DefaultOrchestrator(Orchestrator):
+    """Standard agent execution (no special handling)."""
+
+    def can_handle(self, context: ExecutionContext) -> bool:
+        """Always returns True as this is the fallback orchestrator.
+
+        Args:
+            context: Execution context
+
+        Returns:
+            Always True
+        """
+        return True
+
+    async def run(
+        self,
+        context: ExecutionContext,
+    ) -> AsyncGenerator[BaseEvent, None]:
+        """Standard agent run with event translation.
+
+        This implements the default agent execution flow using the event bridge
+        to translate Agent Framework events to AG-UI events.
+
+        Args:
+            context: Execution context
+
+        Yields:
+            AG-UI events
+        """
+        from ._events import AgentFrameworkEventBridge
+
+        logger.info(f"Starting default agent run for thread_id={context.thread_id}, run_id={context.run_id}")
+
+        # Initialize state tracking
+        initial_state = context.input_data.get("state", {})
+        current_state: dict[str, Any] = initial_state.copy() if initial_state else {}
+
+        # Check if agent uses structured outputs (response_format)
+        # Use isinstance to narrow type for proper attribute access
+        response_format = None
+        if isinstance(context.agent, ChatAgent):
+            response_format = context.agent.chat_options.response_format
+        skip_text_content = response_format is not None
+
+        # Create event bridge
+        event_bridge = AgentFrameworkEventBridge(
+            run_id=context.run_id,
+            thread_id=context.thread_id,
+            predict_state_config=context.config.predict_state_config,
+            current_state=current_state,
+            skip_text_content=skip_text_content,
+            input_messages=context.input_data.get("messages", []),
+            require_confirmation=context.config.require_confirmation,
+        )
+
+        yield event_bridge.create_run_started_event()
+
+        # Emit PredictState custom event if we have predictive state config
+        if context.config.predict_state_config:
+            from ag_ui.core import CustomEvent, EventType
+
+            predict_state_value = [
+                {
+                    "state_key": state_key,
+                    "tool": config["tool"],
+                    "tool_argument": config["tool_argument"],
+                }
+                for state_key, config in context.config.predict_state_config.items()
+            ]
+
+            yield CustomEvent(
+                type=EventType.CUSTOM,
+                name="PredictState",
+                value=predict_state_value,
+            )
+
+        # If we have a state schema, ensure we emit initial state snapshot
+        if context.config.state_schema:
+            # Initialize missing state fields with appropriate empty values based on schema type
+            for key, schema in context.config.state_schema.items():
+                if key not in current_state:
+                    # Default to empty object; use empty array if schema specifies "array" type
+                    current_state[key] = [] if isinstance(schema, dict) and schema.get("type") == "array" else {}  # type: ignore
+            yield event_bridge.create_state_snapshot_event(current_state)
+
+        # Create thread for context tracking
+        thread = AgentThread()
+        thread.metadata = {  # type: ignore[attr-defined]
+            "ag_ui_thread_id": context.thread_id,
+            "ag_ui_run_id": context.run_id,
+        }
+
+        # Inject current state into thread metadata so agent can access it
+        if current_state:
+            thread.metadata["current_state"] = current_state  # type: ignore[attr-defined]
+
+        # Add incoming AG-UI messages to the thread history
+        if context.messages:
+            await thread.on_new_messages(context.messages)
+
+        # Use the full incoming message batch to preserve tool-call adjacency
+        if not context.messages:
+            logger.warning("No messages provided in AG-UI input")
+            yield event_bridge.create_run_finished_event()
+            return
+
+        # Inject current state as system message context if we have state
+        messages_to_run: list[Any] = []
+        if current_state and context.config.state_schema:
+            state_json = json.dumps(current_state, indent=2)
+            from agent_framework import ChatMessage
+
+            state_context_msg = ChatMessage(
+                role="system",
+                contents=[
+                    TextContent(
+                        text=f"""Current state of the application:
+{state_json}
+
+When modifying state, you MUST include ALL existing data plus your changes.
+For example, if adding a new ingredient, include all existing ingredients PLUS the new one.
+Never replace existing data - always append or merge."""
+                    )
+                ],
+            )
+            messages_to_run.append(state_context_msg)
+
+        # Preserve order from client to satisfy provider constraints (assistant tool_calls must
+        # immediately precede tool result messages). Using the full batch avoids reordering.
+        messages_to_run.extend(context.messages)
+
+        # Handle client tools for hybrid execution
+        # Client sends tool metadata, server merges with its own tools.
+        # Client tools have func=None (declaration-only), so @use_function_invocation
+        # will return the function call without executing (passes back to client).
+        from agent_framework import BaseChatClient
+
+        client_tools = convert_agui_tools_to_agent_framework(context.input_data.get("tools"))
+
+        # Extract server tools - use type narrowing when possible
+        server_tools: list[Any] = []
+        if isinstance(context.agent, ChatAgent):
+            server_tools = context.agent.chat_options.tools or []
+        else:
+            # AgentProtocol allows duck-typed implementations - fallback to attribute access
+            # This supports test mocks and custom agent implementations
+            try:
+                chat_options_attr = getattr(context.agent, "chat_options", None)
+                if chat_options_attr is not None:
+                    server_tools = getattr(chat_options_attr, "tools", None) or []
+            except AttributeError:
+                pass
+
+        # Register client tools as additional (declaration-only) so they are not executed on server
+        if client_tools:
+            if isinstance(context.agent, ChatAgent):
+                # Type-safe path for ChatAgent
+                chat_client = context.agent.chat_client
+                if (
+                    isinstance(chat_client, BaseChatClient)
+                    and chat_client.function_invocation_configuration is not None
+                ):
+                    chat_client.function_invocation_configuration.additional_tools = client_tools
+                    logger.debug(
+                        f"[TOOLS] Registered {len(client_tools)} client tools as additional_tools (declaration-only)"
+                    )
+            else:
+                # Fallback for AgentProtocol implementations (test mocks, custom agents)
+                try:
+                    chat_client_attr = getattr(context.agent, "chat_client", None)
+                    if chat_client_attr is not None:
+                        fic = getattr(chat_client_attr, "function_invocation_configuration", None)
+                        if fic is not None:
+                            fic.additional_tools = client_tools  # type: ignore[attr-defined]
+                            logger.debug(
+                                f"[TOOLS] Registered {len(client_tools)} client tools as additional_tools (declaration-only)"
+                            )
+                except AttributeError:
+                    pass
+
+        combined_tools: list[Any] = []
+        if server_tools:
+            combined_tools.extend(server_tools)
+        if client_tools:
+            combined_tools.extend(client_tools)
+
+        # Collect all updates to get the final structured output
+        all_updates: list[Any] = []
+        async for update in context.agent.run_stream(messages_to_run, thread=thread, tools=combined_tools or None):
+            all_updates.append(update)
+            events = await event_bridge.from_agent_run_update(update)
+            for event in events:
+                yield event
+
+        # After agent completes, check if we should stop (waiting for user to confirm changes)
+        if event_bridge.should_stop_after_confirm:
+            logger.info("Stopping run after confirm_changes - waiting for user response")
+            yield event_bridge.create_run_finished_event()
+            return
+
+        # After streaming completes, check if agent has response_format and extract structured output
+        if all_updates and response_format:
+            from agent_framework import AgentRunResponse
+            from pydantic import BaseModel
+
+            logger.info(f"Processing structured output, update count: {len(all_updates)}")
+
+            # Convert streaming updates to final response to get the structured output
+            final_response = AgentRunResponse.from_agent_run_response_updates(
+                all_updates, output_format_type=response_format
+            )
+
+            if final_response.value and isinstance(final_response.value, BaseModel):
+                # Convert Pydantic model to dict
+                response_dict = final_response.value.model_dump(mode="json", exclude_none=True)
+                logger.info(f"Received structured output: {list(response_dict.keys())}")
+
+                # Extract state fields based on state_schema
+                state_updates: dict[str, Any] = {}
+
+                if context.config.state_schema:
+                    # Use state_schema to determine which fields are state
+                    for state_key in context.config.state_schema.keys():
+                        if state_key in response_dict:
+                            state_updates[state_key] = response_dict[state_key]
+                else:
+                    # No schema: treat all non-message fields as state
+                    state_updates = {k: v for k, v in response_dict.items() if k != "message"}
+
+                # Apply state updates if any found
+                if state_updates:
+                    current_state.update(state_updates)
+
+                    # Emit StateSnapshotEvent with the updated state
+                    state_snapshot = event_bridge.create_state_snapshot_event(current_state)
+                    yield state_snapshot
+                    logger.info(f"Emitted StateSnapshotEvent with updates: {list(state_updates.keys())}")
+
+                # If there's a message field, emit it as chat text
+                if "message" in response_dict and response_dict["message"]:
+                    message_id = generate_event_id()
+                    yield TextMessageStartEvent(message_id=message_id, role="assistant")
+                    yield TextMessageContentEvent(message_id=message_id, delta=response_dict["message"])
+                    yield TextMessageEndEvent(message_id=message_id)
+                    logger.info(f"Emitted conversational message: {response_dict['message'][:100]}...")
+
+        if event_bridge.current_message_id:
+            yield event_bridge.create_message_end_event(event_bridge.current_message_id)
+
+        yield event_bridge.create_run_finished_event()
+        logger.info(f"Completed agent run for thread_id={context.thread_id}, run_id={context.run_id}")
+
+
+__all__ = [
+    "Orchestrator",
+    "ExecutionContext",
+    "HumanInTheLoopOrchestrator",
+    "DefaultOrchestrator",
+]

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

@@ -0,0 +1,27 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Type definitions for AG-UI integration."""
+
+from typing import Any, TypedDict
+
+
+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

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

@@ -0,0 +1,166 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Utility functions for AG-UI integration."""
+
+import copy
+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 AIFunction, ToolProtocol
+
+
+def generate_event_id() -> str:
+    """Generate a unique event ID."""
+    return str(uuid.uuid4())
+
+
+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):
+        return asdict(obj)  # type: ignore[arg-type]
+    if hasattr(obj, "model_dump"):
+        return obj.model_dump()  # type: ignore[no-any-return]
+    if hasattr(obj, "dict"):
+        return obj.dict()  # type: ignore[no-any-return]
+    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[AIFunction[Any, Any]] | None:
+    """Convert AG-UI tool definitions to Agent Framework AIFunction declarations.
+
+    Creates declaration-only AIFunction instances (no executable implementation).
+    These are used to tell the LLM about available tools. The actual execution
+    happens on the client side via @use_function_invocation.
+
+    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 AIFunction declarations, or None if no tools provided
+    """
+    if not agui_tools:
+        return None
+
+    result: list[AIFunction[Any, Any]] = []
+    for tool_def in agui_tools:
+        # Create declaration-only AIFunction (func=None means no implementation)
+        # When func=None, the declaration_only property returns True,
+        # which tells @use_function_invocation to return the function call
+        # without executing it (so it can be sent back to the client)
+        func: AIFunction[Any, Any] = AIFunction(
+            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: (
+        ToolProtocol
+        | Callable[..., Any]
+        | MutableMapping[str, Any]
+        | Sequence[ToolProtocol | 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 @use_function_invocation decorator 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[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]] = [tools]  # type: ignore[list-item]
+    else:
+        tool_list = tools  # type: ignore[assignment]
+
+    results: list[dict[str, Any]] = []
+
+    for tool in tool_list:
+        if isinstance(tool, dict):
+            # Already in dict format, pass through
+            results.append(tool)  # type: ignore[arg-type]
+        elif isinstance(tool, AIFunction):
+            # Convert AIFunction to AG-UI tool format
+            results.append(
+                {
+                    "name": tool.name,
+                    "description": tool.description,
+                    "parameters": tool.parameters(),
+                }
+            )
+        elif callable(tool):
+            # Convert callable to AIFunction first, then to AG-UI format
+            from agent_framework import ai_function
+
+            ai_func = ai_function(tool)
+            results.append(
+                {
+                    "name": ai_func.name,
+                    "description": ai_func.description,
+                    "parameters": ai_func.parameters(),
+                }
+            )
+        elif isinstance(tool, ToolProtocol):
+            # Handle other ToolProtocol implementations
+            # For now, we'll skip non-AIFunction tools as they may not have
+            # the parameters() method. This matches .NET behavior which only
+            # converts AIFunctionDeclaration instances.
+            continue
+
+    return results if results else None

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,243 @@
+# 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
+
+```python
+from fastapi import FastAPI
+from agent_framework import ChatAgent
+from agent_framework.azure import AzureOpenAIChatClient
+from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
+
+# Create your agent
+agent = ChatAgent(
+    name="my_agent",
+    instructions="You are a helpful assistant.",
+    chat_client=AzureOpenAIChatClient(model_id="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
+
+Complete examples for all features are in the `examples/` directory:
+
+- `examples/agents/simple_agent.py` - Basic agentic chat
+- `examples/agents/weather_agent.py` - Backend tool rendering
+- `examples/agents/task_planner_agent.py` - Human in the loop with approvals
+- `examples/agents/research_assistant_agent.py` - Agentic generative UI
+- `examples/agents/ui_generator_agent.py` - Tool-based generative UI
+- `examples/agents/recipe_agent.py` - Shared state management
+- `examples/agents/document_writer_agent.py` - Predictive state updates
+- `examples/server/main.py` - FastAPI server with all endpoints
+
+Run the example server:
+
+```bash
+cd examples/server
+uvicorn main:app --reload
+```
+
+To enable debug logging:
+
+```bash
+ENABLE_DEBUG_LOGGING=1 uvicorn main:app --reload
+```
+
+The server exposes endpoints at:
+- `/agentic_chat`
+- `/backend_tool_rendering`
+- `/human_in_the_loop`
+- `/agentic_generative_ui`
+- `/tool_based_generative_ui`
+- `/shared_state`
+- `/predictive_state_updates`
+
+## 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 AgentRunResponseUpdate 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
+
+### Shared State
+
+State is injected as system messages and updated via predictive state updates:
+
+```python
+from agent_framework import ChatAgent
+from agent_framework.azure import AzureOpenAIChatClient
+from agent_framework.ag_ui import AgentFrameworkAgent
+
+# Create your agent
+agent = ChatAgent(
+    name="recipe_agent",
+    chat_client=AzureOpenAIChatClient(model_id="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 ChatAgent
+from agent_framework.azure import AzureOpenAIChatClient
+from agent_framework.ag_ui import AgentFrameworkAgent
+
+# Create your agent
+agent = ChatAgent(
+    name="document_writer",
+    chat_client=AzureOpenAIChatClient(model_id="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
+)
+```
+
+### Custom Confirmation Strategies
+
+Provide domain-specific confirmation messages:
+
+```python
+from typing import Any
+from agent_framework import ChatAgent
+from agent_framework.azure import AzureOpenAIChatClient
+from agent_framework.ag_ui import AgentFrameworkAgent, ConfirmationStrategy
+
+class CustomConfirmationStrategy(ConfirmationStrategy):
+    def on_approval_accepted(self, steps: list[dict[str, Any]]) -> str:
+        return "Your custom approval message!"
+    
+    def on_approval_rejected(self, steps: list[dict[str, Any]]) -> str:
+        return "Your custom rejection message!"
+    
+    def on_state_confirmed(self) -> str:
+        return "State changes confirmed!"
+    
+    def on_state_rejected(self) -> str:
+        return "State changes rejected!"
+
+agent = ChatAgent(
+    name="custom_agent",
+    chat_client=AzureOpenAIChatClient(model_id="gpt-4o"),
+)
+
+wrapped_agent = AgentFrameworkAgent(
+    agent=agent,
+    confirmation_strategy=CustomConfirmationStrategy(),
+)
+```
+
+### Human in the Loop
+
+Human-in-the-loop is automatically handled when tools are marked for approval:
+
+```python
+from agent_framework import ai_function
+
+@ai_function(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()],
+)
+
+## Documentation
+
+For detailed documentation, see [DESIGN.md](DESIGN.md).
+
+## 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,25 @@
+# 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 agent as simple_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",
+    "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,58 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Example agent demonstrating predictive state updates with document writing."""
+
+from agent_framework import ChatAgent, ai_function
+from agent_framework.azure import AzureOpenAIChatClient
+
+from agent_framework_ag_ui import AgentFrameworkAgent, DocumentWriterConfirmationStrategy
+
+
+@ai_function
+def write_document_local(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."
+
+
+agent = ChatAgent(
+    name="document_writer",
+    instructions=(
+        "You are a helpful assistant for writing documents. "
+        "To write the document, you MUST use the write_document_local 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."
+    ),
+    chat_client=AzureOpenAIChatClient(),
+    tools=[write_document_local],
+)
+
+document_writer_agent = 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_local", "tool_argument": "document"},
+    },
+    confirmation_strategy=DocumentWriterConfirmationStrategy(),
+)

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

@@ -0,0 +1,76 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Human-in-the-loop agent demonstrating step customization (Feature 5)."""
+
+from enum import Enum
+
+from agent_framework import ChatAgent, ai_function
+from agent_framework.azure import AzureOpenAIChatClient
+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")
+
+
+@ai_function(
+    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."
+
+
+# Create the human-in-the-loop agent using tool-based approach for predictive state
+human_in_the_loop_agent = ChatAgent(
+    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"
+
+    After calling the function, provide a brief acknowledgment like:
+    "I've created a plan with 10 steps. You can customize which steps to enable before I proceed."
+    """,
+    chat_client=AzureOpenAIChatClient(),
+    tools=[generate_task_steps],
+)

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

@@ -0,0 +1,122 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Recipe agent example demonstrating shared state management (Feature 3)."""
+
+from enum import Enum
+
+from agent_framework import ChatAgent, ai_function
+from agent_framework.azure import AzureOpenAIChatClient
+from pydantic import BaseModel, Field
+
+from agent_framework_ag_ui import AgentFrameworkAgent, RecipeConfirmationStrategy
+
+
+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")
+
+
+@ai_function
+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."
+
+
+# Create the recipe agent using tool-based approach for streaming
+agent = ChatAgent(
+    name="recipe_agent",
+    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
+    """,
+    chat_client=AzureOpenAIChatClient(),
+    tools=[update_recipe],
+)
+
+recipe_agent = 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"},
+    },
+    confirmation_strategy=RecipeConfirmationStrategy(),
+)

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

@@ -0,0 +1,100 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Example agent demonstrating agentic generative UI with custom events during execution."""
+
+import asyncio
+
+from agent_framework import ChatAgent, ai_function
+from agent_framework.azure import AzureOpenAIChatClient
+
+from agent_framework_ag_ui import AgentFrameworkAgent
+
+
+@ai_function
+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)
+
+
+@ai_function
+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)
+
+
+@ai_function
+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)
+
+
+agent = ChatAgent(
+    name="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."
+    ),
+    chat_client=AzureOpenAIChatClient(),
+    tools=[research_topic, create_presentation, analyze_data],
+)
+
+research_assistant_agent = 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,13 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Simple agentic chat example (Feature 1: Agentic Chat)."""
+
+from agent_framework import ChatAgent
+from agent_framework.azure import AzureOpenAIChatClient
+
+# Create a simple chat agent
+agent = ChatAgent(
+    name="simple_chat_agent",
+    instructions="You are a helpful assistant. Be concise and friendly.",
+    chat_client=AzureOpenAIChatClient(),
+)

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

@@ -0,0 +1,73 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Example agent demonstrating human-in-the-loop with function approvals."""
+
+from agent_framework import ChatAgent, ai_function
+from agent_framework.azure import AzureOpenAIChatClient
+
+from agent_framework_ag_ui import AgentFrameworkAgent, TaskPlannerConfirmationStrategy
+
+
+@ai_function(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}"
+
+
+@ai_function(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}'"
+
+
+@ai_function(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}"
+
+
+agent = ChatAgent(
+    name="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."
+    ),
+    chat_client=AzureOpenAIChatClient(),
+    tools=[create_calendar_event, send_email, book_meeting_room],
+)
+
+task_planner_agent = AgentFrameworkAgent(
+    agent=agent,
+    name="TaskPlanner",
+    description="Plans and executes tasks with user approval",
+    confirmation_strategy=TaskPlannerConfirmationStrategy(),
+)

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

@@ -0,0 +1,318 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Task steps agent demonstrating agentic generative UI (Feature 6)."""
+
+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 ChatAgent, ai_function
+from agent_framework.azure import AzureOpenAIChatClient
+from pydantic import BaseModel, Field
+
+from agent_framework_ag_ui import AgentFrameworkAgent
+
+
+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")
+
+
+@ai_function
+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."
+
+
+# Create the task steps agent using tool-based approach for streaming
+agent = ChatAgent(
+    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"
+    """,
+    chat_client=AzureOpenAIChatClient(),
+    tools=[generate_task_steps],
+)
+
+task_steps_agent = 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:
+    """Wrapper that adds step execution simulation after plan generation.
+
+    This wrapper delegates to AgentFrameworkAgent but is recognized as compatible
+    by add_agent_framework_fastapi_endpoint since it implements run_agent().
+    """
+
+    def __init__(self, base_agent: AgentFrameworkAgent):
+        """Initialize wrapper with base agent."""
+        self._base_agent = base_agent
+
+    @property
+    def name(self) -> str:
+        """Delegate to base agent."""
+        return self._base_agent.name
+
+    @property
+    def description(self) -> str:
+        """Delegate to base agent."""
+        return self._base_agent.description
+
+    def __getattr__(self, name: str) -> Any:
+        """Delegate all other attribute access to base agent."""
+        return getattr(self._base_agent, name)
+
+    async def run_agent(self, input_data: dict[str, Any]) -> AsyncGenerator[Any, None]:
+        """Run the agent and then simulate step execution."""
+        import logging
+        import uuid
+
+        logger = logging.getLogger(__name__)
+        logger.info("TaskStepsAgentWithExecution.run_agent() called - wrapper is active")
+
+        # First, run the base agent to generate the plan - buffer text messages
+        final_state: dict[str, Any] | None = None
+        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_agent(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
+                    logger.info(f"Captured STATE_SNAPSHOT event with state: {final_state}")
+                    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  # type: ignore
+            chat_client = chat_agent.chat_client  # type: ignore
+
+            # Build messages for summary call
+            from agent_framework._types import ChatMessage, TextContent
+
+            original_messages = input_data.get("messages", [])
+
+            # Convert to ChatMessage objects if needed
+            messages: list[ChatMessage] = []
+            for msg in original_messages:
+                if isinstance(msg, dict):
+                    content_str = msg.get("content", "")
+                    if isinstance(content_str, str):
+                        messages.append(
+                            ChatMessage(
+                                role=msg.get("role", "user"),
+                                contents=[TextContent(text=content_str)],
+                            )
+                        )
+                elif isinstance(msg, ChatMessage):
+                    messages.append(msg)
+
+            # Add completion message
+            messages.append(
+                ChatMessage(
+                    role="user",
+                    contents=[
+                        TextContent(
+                            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 chat_client.get_streaming_response(messages=messages):
+                    # 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
+
+
+# Export the wrapped agent
+task_steps_agent_wrapped = TaskStepsAgentWithExecution(task_steps_agent)

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

@@ -0,0 +1,119 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Example agent demonstrating Tool-based Generative UI (Feature 5)."""
+
+from typing import Any
+
+from agent_framework import ChatAgent, ai_function
+from agent_framework.azure import AzureOpenAIChatClient
+
+from agent_framework_ag_ui import AgentFrameworkAgent
+
+
+@ai_function
+def generate_haiku(english: list[str], japanese: list[str], image_name: str | None, gradient: str) -> str:
+    """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.
+
+    Args:
+        english: English haiku lines (exactly 3 lines)
+        japanese: Japanese haiku lines (exactly 3 lines)
+        image_name: 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: CSS gradient string for background (e.g., "linear-gradient(135deg, #667eea 0%, #764ba2 100%)")
+
+    Returns:
+        Haiku metadata for frontend rendering
+    """
+    return f"Haiku generated with image: {image_name}"
+
+
+@ai_function
+def create_chart(chart_type: str, data_points: list[dict[str, Any]], title: str) -> str:
+    """Create an interactive chart (FRONTEND_RENDER).
+
+    This tool creates chart specifications for frontend rendering.
+    The frontend should render this as an interactive chart component.
+
+    Args:
+        chart_type: Type of chart (bar, line, pie, scatter)
+        data_points: Data points for the chart
+        title: Chart title
+
+    Returns:
+        Chart specification for frontend rendering
+    """
+    return f"Chart '{title}' created with {len(data_points)} data points"
+
+
+@ai_function
+def display_timeline(events: list[dict[str, Any]], start_date: str, end_date: str) -> str:
+    """Display an interactive timeline (FRONTEND_RENDER).
+
+    This tool creates timeline specifications for frontend rendering.
+    The frontend should render this as an interactive timeline component.
+
+    Args:
+        events: Events to display on the timeline
+        start_date: Timeline start date
+        end_date: Timeline end date
+
+    Returns:
+        Timeline specification for frontend rendering
+    """
+    return f"Timeline created with {len(events)} events from {start_date} to {end_date}"
+
+
+@ai_function
+def show_comparison_table(items: list[dict[str, Any]], columns: list[str]) -> str:
+    """Show a comparison table (FRONTEND_RENDER).
+
+    This tool creates table specifications for frontend rendering.
+    The frontend should render this as an interactive comparison table.
+
+    Args:
+        items: Items to compare
+        columns: Column names
+
+    Returns:
+        Table specification for frontend rendering
+    """
+    return f"Comparison table created with {len(items)} items and {len(columns)} columns"
+
+
+# Create the UI generator agent using tool-based approach with forced tool usage
+agent = ChatAgent(
+    name="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).
+    """,
+    chat_client=AzureOpenAIChatClient(),
+    tools=[generate_haiku, create_chart, display_timeline, show_comparison_table],
+    # Force tool usage - the LLM MUST call a tool, cannot respond with plain text
+    chat_options={"tool_choice": "required"},
+)
+
+ui_generator_agent = 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,71 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Weather agent example demonstrating backend tool rendering."""
+
+from typing import Any
+
+from agent_framework import ChatAgent, ai_function
+from agent_framework.azure import AzureOpenAIChatClient
+
+
+@ai_function
+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,
+    }
+
+
+@ai_function
+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)
+
+
+# Create the weather agent
+weather_agent = ChatAgent(
+    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."
+    ),
+    chat_client=AzureOpenAIChatClient(),
+    tools=[get_weather, get_forecast],
+)

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/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""API endpoints for AG-UI examples."""

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

@@ -0,0 +1,22 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Backend tool rendering endpoint."""
+
+from fastapi import FastAPI
+
+from agent_framework_ag_ui import add_agent_framework_fastapi_endpoint
+
+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.
+    """
+    add_agent_framework_fastapi_endpoint(
+        app,
+        weather_agent,
+        "/backend_tool_rendering",
+    )

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

@@ -0,0 +1,129 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Example FastAPI server with AG-UI endpoints."""
+
+import logging
+import os
+
+import uvicorn
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+
+from agent_framework_ag_ui import add_agent_framework_fastapi_endpoint
+
+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 agent as simple_agent
+from ..agents.task_steps_agent import task_steps_agent_wrapped as task_steps_agent  # Custom wrapper
+from ..agents.ui_generator_agent import ui_generator_agent
+from ..agents.weather_agent import weather_agent
+
+# 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=["*"],
+)
+
+# Agentic Chat - basic chat agent
+add_agent_framework_fastapi_endpoint(
+    app=app,
+    agent=simple_agent,
+    path="/agentic_chat",
+)
+
+# Backend Tool Rendering - agent with tools
+add_agent_framework_fastapi_endpoint(
+    app=app,
+    agent=weather_agent,
+    path="/backend_tool_rendering",
+)
+
+# Shared State - recipe agent with structured output
+add_agent_framework_fastapi_endpoint(
+    app=app,
+    agent=recipe_agent,
+    path="/shared_state",
+)
+
+# Predictive State Updates - document writer with predictive state
+add_agent_framework_fastapi_endpoint(
+    app=app,
+    agent=document_writer_agent,
+    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,
+    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,  # type: ignore[arg-type]
+    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,
+    path="/tool_based_generative_ui",
+)
+
+
+def main():
+    """Run the server."""
+    port = int(os.getenv("PORT", "8888"))
+    host = os.getenv("HOST", "127.0.0.1")
+
+    # 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,461 @@
+# 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.
+
+### ChatAgent Integration (`client_with_agent.py`)
+
+Best practice example using `ChatAgent` 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 `ChatAgent` 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?"
+    ↓
+ChatAgent 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')
+    ↓
+ChatAgent 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 ChatAgent'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 ChatAgent wrapper, client-side tools work:
+- Tools passed in ChatOptions execute locally
+- Server can also have its own tools
+- Hybrid execution works automatically
+
+## 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 ChatAgent
+from agent_framework.azure import AzureOpenAIChatClient
+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")
+deployment_name = os.environ.get("AZURE_OPENAI_DEPLOYMENT_NAME")
+api_key = os.environ.get("AZURE_OPENAI_API_KEY")
+
+if not endpoint:
+    raise ValueError("AZURE_OPENAI_ENDPOINT environment variable is required")
+if not deployment_name:
+    raise ValueError("AZURE_OPENAI_DEPLOYMENT_NAME environment variable is required")
+if not api_key:
+    raise ValueError("AZURE_OPENAI_API_KEY environment variable is required")
+
+# Create the AI agent
+agent = ChatAgent(
+    name="AGUIAssistant",
+    instructions="You are a helpful assistant.",
+    chat_client=AzureOpenAIChatClient(
+        endpoint=endpoint,
+        deployment_name=deployment_name,
+        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
+- **`ChatAgent`**: 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**: `AzureOpenAIChatClient` can read from environment variables (`AZURE_OPENAI_ENDPOINT`, `AZURE_OPENAI_CHAT_DEPLOYMENT_NAME`, `AZURE_OPENAI_API_KEY`) or accept parameters directly
+
+**Alternative (simpler)**: Use environment variables only:
+
+```python
+# No need to read environment variables manually
+agent = ChatAgent(
+    name="AGUIAssistant",
+    instructions="You are a helpful assistant.",
+    chat_client=AzureOpenAIChatClient(),  # 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_CHAT_DEPLOYMENT_NAME="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 import TextContent
+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_streaming_response(message, metadata=metadata):
+                    # 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 isinstance(content, TextContent) 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_streaming_response()` for real-time streaming or `get_response()` for non-streaming
+- **Context Manager**: Use `async with` for automatic cleanup of HTTP connections
+- **Standard Interface**: Works with all Agent Framework patterns (ChatAgent, 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,72 @@
+# 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 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
+
+                async for update in client.get_streaming_response(message, metadata=metadata):
+                    # 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
+                    from agent_framework import TextContent
+
+                    for content in update.contents:
+                        if isinstance(content, TextContent) 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,235 @@
+# 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
+"""
+
+import asyncio
+import os
+
+from agent_framework import ai_function
+
+from agent_framework_ag_ui import AGUIChatClient
+
+
+@ai_function
+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}")
+
+
+@ai_function
+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)
+
+    async for update in client.get_streaming_response("Tell me a short joke", metadata=metadata):
+        if not thread_id and update.additional_properties:
+            thread_id = update.additional_properties.get("thread_id")
+
+        from agent_framework import TextContent
+
+        for content in update.contents:
+            if isinstance(content, TextContent) and content.text:
+                print(content.text, end="", flush=True)
+
+    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("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 ChatAgent 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 ChatAgent wrapper with tools
+    - See client_with_agent.py for the hybrid pattern
+    - ChatAgent 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(
+        "What's the weather in Seattle?", tools=[get_weather, calculate], metadata=metadata
+    )
+
+    print(f"Assistant: {response.text}")
+
+    # Show tool calls if any
+    from agent_framework import FunctionCallContent
+
+    tool_called = False
+    for message in response.messages:
+        for content in message.contents:
+            if isinstance(content, FunctionCallContent):
+                print(f"\n[Tool Called: {content.name}]")
+                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("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("What's my name?", 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 ChatAgent for history management]")
+
+    # Third turn
+    print("\nUser: Can you also tell me what 10 * 5 is?\n")
+    response3 = await client.get_response(
+        "Can you also tell me what 10 * 5 is?", 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,186 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Example showing ChatAgent with AGUIChatClient for hybrid tool execution.
+
+This demonstrates the HYBRID pattern matching .NET AGUIClient implementation:
+
+1. AgentThread Pattern (like .NET):
+   - Create thread with agent.get_new_thread()
+   - Pass thread to agent.run_stream() on each turn
+   - Thread automatically maintains conversation history via message_store
+
+2. Hybrid Tool Execution:
+   - AGUIChatClient has @use_function_invocation decorator
+   - 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: thread maintains state, tools execute on appropriate side.
+"""
+
+import asyncio
+import logging
+import os
+
+from agent_framework import ChatAgent, FunctionCallContent, FunctionResultContent, TextContent, ai_function
+
+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__)
+
+
+@ai_function(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 ChatAgent + AGUIChatClient hybrid tool execution.
+
+    This matches the .NET pattern from Program.cs where:
+    - AIAgent agent = chatClient.CreateAIAgent(tools: [...])
+    - AgentThread thread = agent.GetNewThread()
+    - RunStreamingAsync(messages, thread)
+
+    Python equivalent:
+    - agent = ChatAgent(chat_client=AGUIChatClient(...), tools=[...])
+    - thread = agent.get_new_thread()  # Creates thread with message_store
+    - agent.run_stream(message, thread=thread)  # Thread accumulates history
+    """
+    server_url = os.environ.get("AGUI_SERVER_URL", "http://127.0.0.1:5100/")
+
+    print("=" * 70)
+    print("ChatAgent + AGUIChatClient: Hybrid Tool Execution")
+    print("=" * 70)
+    print(f"\nServer: {server_url}")
+    print("\nThis example demonstrates:")
+    print("  1. AgentThread maintains conversation state (like .NET)")
+    print("  2. Client-side tools execute locally via @use_function_invocation")
+    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 ChatAgent for conversation history management
+            agent = ChatAgent(
+                name="remote_assistant",
+                instructions="You are a helpful assistant. Remember user information across the conversation.",
+                chat_client=remote_client,
+                tools=[get_weather],
+            )
+
+            # Create a thread to maintain conversation state (like .NET AgentThread)
+            thread = agent.get_new_thread()
+
+            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_stream("My name is Alice and I live in Seattle", thread=thread):
+                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_stream("What's my name?", thread=thread):
+                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_stream("Where do I live?", thread=thread):
+                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_stream("What's the weather forecast for today in Seattle?", thread=thread):
+                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_stream("What time zone is Seattle in?", thread=thread):
+                if chunk.text:
+                    print(chunk.text, end="", flush=True)
+            print("\n")
+
+            # Show thread state
+            if thread.message_store:
+
+                def _preview_for_message(m) -> str:
+                    # Prefer plain text when present
+                    if getattr(m, "text", ""):
+                        t = m.text
+                        return (t[:60] + "...") if len(t) > 60 else t
+                    # Build from contents when no direct text
+                    parts: list[str] = []
+                    for c in getattr(m, "contents", []) or []:
+                        if isinstance(c, FunctionCallContent):
+                            args = c.arguments
+                            if isinstance(args, dict):
+                                try:
+                                    import json as _json
+
+                                    args_str = _json.dumps(args)
+                                except Exception:
+                                    args_str = str(args)
+                            else:
+                                args_str = str(args or "{}")
+                            parts.append(f"tool_call {c.name} {args_str}")
+                        elif isinstance(c, FunctionResultContent):
+                            parts.append(f"tool_result[{c.call_id}]: {str(c.result)[:40]}")
+                        elif isinstance(c, TextContent):
+                            if c.text:
+                                parts.append(c.text)
+                        else:
+                            typename = getattr(c, "type", c.__class__.__name__)
+                            parts.append(f"<{typename}>")
+                    preview = " | ".join(parts) if parts else ""
+                    return (preview[:60] + "...") if len(preview) > 60 else preview
+
+                messages = await thread.message_store.list_messages()
+                print(f"\n[THREAD STATE] {len(messages)} messages in thread's message_store")
+                for i, msg in enumerate(messages[-6:], 1):  # Show last 6
+                    role = msg.role.value if hasattr(msg.role, "value") else str(msg.role)
+                    text_preview = _preview_for_message(msg)
+                    print(f"  {i}. [{role}]: {text_preview}")
+
+    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,81 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""AG-UI server example with server-side tools."""
+
+import logging
+import os
+
+from agent_framework import ChatAgent, ai_function
+from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
+from agent_framework.azure import AzureOpenAIChatClient
+from dotenv import load_dotenv
+from fastapi import FastAPI
+
+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")
+deployment_name = os.environ.get("AZURE_OPENAI_CHAT_DEPLOYMENT_NAME")
+
+if not endpoint:
+    raise ValueError("AZURE_OPENAI_ENDPOINT environment variable is required")
+if not deployment_name:
+    raise ValueError("AZURE_OPENAI_CHAT_DEPLOYMENT_NAME environment variable is required")
+
+
+# Server-side tool (executes on server)
+@ai_function(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 @use_function_invocation on AGUIChatClient will execute it client-side.
+# This matches the .NET AG-UI hybrid execution pattern.
+agent = ChatAgent(
+    name="AGUIAssistant",
+    instructions="You are a helpful assistant. Use get_weather for weather and get_time_zone for time zones.",
+    chat_client=AzureOpenAIChatClient(
+        endpoint=endpoint,
+        deployment_name=deployment_name,
+    ),
+    tools=[get_time_zone],  # ONLY server-side tools
+)
+
+# 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, log_level="debug", access_log=True)

python/packages/ag-ui/pyproject.toml

@@ -0,0 +1,74 @@
+[project]
+name = "agent-framework-ag-ui"
+version = "1.0.0b251111"
+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",
+    "ag-ui-protocol>=0.1.9",
+    "fastapi>=0.115.0",
+    "uvicorn>=0.30.0"
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=8.0.0",
+    "pytest-asyncio>=0.24.0",
+    "httpx>=0.27.0",
+]
+
+[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"]
+pythonpath = ["."]
+
+[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]
+exclude = ["tests", "examples"]
+typeCheckingMode = "basic"
+
+[tool.poe]
+executor.type = "uv"
+include = "../../shared_tasks.toml"
+
+[tool.poe.tasks]
+mypy = "mypy --config-file $POE_ROOT/pyproject.toml agent_framework_ag_ui"
+test = "pytest --cov=agent_framework_ag_ui --cov-report=term-missing:skip-covered tests"

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

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

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

@@ -0,0 +1,577 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Comprehensive tests for AgentFrameworkAgent (_agent.py)."""
+
+import json
+
+import pytest
+from agent_framework import ChatAgent, TextContent
+from agent_framework._types import ChatResponseUpdate
+
+
+async def test_agent_initialization_basic():
+    """Test basic agent initialization without state schema."""
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
+
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    wrapper = AgentFrameworkAgent(agent=agent)
+
+    assert wrapper.name == "test_agent"
+    assert wrapper.agent == agent
+    assert wrapper.config.state_schema == {}
+    assert wrapper.config.predict_state_config == {}
+
+
+async def test_agent_initialization_with_state_schema():
+    """Test agent initialization with state_schema."""
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
+
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    state_schema = {"document": {"type": "string"}}
+    wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
+
+    assert wrapper.config.state_schema == state_schema
+
+
+async def test_agent_initialization_with_predict_state_config():
+    """Test agent initialization with predict_state_config."""
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
+
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    predict_config = {"document": {"tool": "write_doc", "tool_argument": "content"}}
+    wrapper = AgentFrameworkAgent(agent=agent, predict_state_config=predict_config)
+
+    assert wrapper.config.predict_state_config == predict_config
+
+
+async def test_run_started_event_emission():
+    """Test RunStartedEvent is emitted at start of run."""
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
+
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    wrapper = AgentFrameworkAgent(agent=agent)
+
+    input_data = {"messages": [{"role": "user", "content": "Hi"}]}
+
+    events = []
+    async for event in wrapper.run_agent(input_data):
+        events.append(event)
+
+    # First event should be RunStartedEvent
+    assert events[0].type == "RUN_STARTED"
+    assert events[0].run_id is not None
+    assert events[0].thread_id is not None
+
+
+async def test_predict_state_custom_event_emission():
+    """Test PredictState CustomEvent is emitted when predict_state_config is present."""
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
+
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    predict_config = {
+        "document": {"tool": "write_doc", "tool_argument": "content"},
+        "summary": {"tool": "summarize", "tool_argument": "text"},
+    }
+    wrapper = AgentFrameworkAgent(agent=agent, predict_state_config=predict_config)
+
+    input_data = {"messages": [{"role": "user", "content": "Hi"}]}
+
+    events = []
+    async for event in wrapper.run_agent(input_data):
+        events.append(event)
+
+    # Find PredictState event
+    predict_events = [e for e in events if e.type == "CUSTOM" and e.name == "PredictState"]
+    assert len(predict_events) == 1
+
+    predict_value = predict_events[0].value
+    assert len(predict_value) == 2
+    assert {"state_key": "document", "tool": "write_doc", "tool_argument": "content"} in predict_value
+    assert {"state_key": "summary", "tool": "summarize", "tool_argument": "text"} in predict_value
+
+
+async def test_initial_state_snapshot_with_schema():
+    """Test initial StateSnapshotEvent emission when state_schema present."""
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
+
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    state_schema = {"document": {"type": "string"}}
+    wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
+
+    input_data = {
+        "messages": [{"role": "user", "content": "Hi"}],
+        "state": {"document": "Initial content"},
+    }
+
+    events = []
+    async for event in wrapper.run_agent(input_data):
+        events.append(event)
+
+    # Find StateSnapshotEvent
+    snapshot_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
+    assert len(snapshot_events) >= 1
+
+    # First snapshot should have initial state
+    assert snapshot_events[0].snapshot == {"document": "Initial content"}
+
+
+async def test_state_initialization_object_type():
+    """Test state initialization with object type in schema."""
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
+
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    state_schema = {"recipe": {"type": "object", "properties": {}}}
+    wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
+
+    input_data = {"messages": [{"role": "user", "content": "Hi"}]}
+
+    events = []
+    async for event in wrapper.run_agent(input_data):
+        events.append(event)
+
+    # Find StateSnapshotEvent
+    snapshot_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
+    assert len(snapshot_events) >= 1
+
+    # Should initialize as empty object
+    assert snapshot_events[0].snapshot == {"recipe": {}}
+
+
+async def test_state_initialization_array_type():
+    """Test state initialization with array type in schema."""
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
+
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    state_schema = {"steps": {"type": "array", "items": {}}}
+    wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
+
+    input_data = {"messages": [{"role": "user", "content": "Hi"}]}
+
+    events = []
+    async for event in wrapper.run_agent(input_data):
+        events.append(event)
+
+    # Find StateSnapshotEvent
+    snapshot_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
+    assert len(snapshot_events) >= 1
+
+    # Should initialize as empty array
+    assert snapshot_events[0].snapshot == {"steps": []}
+
+
+async def test_run_finished_event_emission():
+    """Test RunFinishedEvent is emitted at end of run."""
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
+
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    wrapper = AgentFrameworkAgent(agent=agent)
+
+    input_data = {"messages": [{"role": "user", "content": "Hi"}]}
+
+    events = []
+    async for event in wrapper.run_agent(input_data):
+        events.append(event)
+
+    # Last event should be RunFinishedEvent
+    assert events[-1].type == "RUN_FINISHED"
+
+
+async def test_tool_result_confirm_changes_accepted():
+    """Test confirm_changes tool result handling when accepted."""
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            yield ChatResponseUpdate(contents=[TextContent(text="Document updated")])
+
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    wrapper = AgentFrameworkAgent(
+        agent=agent,
+        state_schema={"document": {"type": "string"}},
+        predict_state_config={"document": {"tool": "write_doc", "tool_argument": "content"}},
+    )
+
+    # Simulate tool result message with acceptance
+    tool_result = {"accepted": True, "steps": []}
+    input_data = {
+        "messages": [
+            {
+                "role": "tool",  # Tool result from UI
+                "content": json.dumps(tool_result),
+                "toolCallId": "confirm_call_123",
+            }
+        ],
+        "state": {"document": "Updated content"},
+    }
+
+    events = []
+    async for event in wrapper.run_agent(input_data):
+        events.append(event)
+
+    # Should emit text message confirming acceptance
+    text_content_events = [e for e in events if e.type == "TEXT_MESSAGE_CONTENT"]
+    assert len(text_content_events) > 0
+    # Should contain confirmation message mentioning the state key or generic confirmation
+    confirmation_found = any(
+        "document" in e.delta.lower()
+        or "confirm" in e.delta.lower()
+        or "applied" in e.delta.lower()
+        or "changes" in e.delta.lower()
+        for e in text_content_events
+    )
+    assert confirmation_found, f"No confirmation in deltas: {[e.delta for e in text_content_events]}"
+
+
+async def test_tool_result_confirm_changes_rejected():
+    """Test confirm_changes tool result handling when rejected."""
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            yield ChatResponseUpdate(contents=[TextContent(text="OK")])
+
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    wrapper = AgentFrameworkAgent(agent=agent)
+
+    # Simulate tool result message with rejection
+    tool_result = {"accepted": False, "steps": []}
+    input_data = {
+        "messages": [
+            {
+                "role": "tool",
+                "content": json.dumps(tool_result),
+                "toolCallId": "confirm_call_123",
+            }
+        ],
+    }
+
+    events = []
+    async for event in wrapper.run_agent(input_data):
+        events.append(event)
+
+    # Should emit text message asking what to change
+    text_content_events = [e for e in events if e.type == "TEXT_MESSAGE_CONTENT"]
+    assert len(text_content_events) > 0
+    assert any("what would you like me to change" in e.delta.lower() for e in text_content_events)
+
+
+async def test_tool_result_function_approval_accepted():
+    """Test function approval tool result when steps are accepted."""
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            yield ChatResponseUpdate(contents=[TextContent(text="OK")])
+
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    wrapper = AgentFrameworkAgent(agent=agent)
+
+    # Simulate tool result with multiple steps
+    tool_result = {
+        "accepted": True,
+        "steps": [
+            {"id": "step1", "description": "Send email", "status": "enabled"},
+            {"id": "step2", "description": "Create calendar event", "status": "enabled"},
+        ],
+    }
+    input_data = {
+        "messages": [
+            {
+                "role": "tool",
+                "content": json.dumps(tool_result),
+                "toolCallId": "approval_call_123",
+            }
+        ],
+    }
+
+    events = []
+    async for event in wrapper.run_agent(input_data):
+        events.append(event)
+
+    # Should list enabled steps
+    text_content_events = [e for e in events if e.type == "TEXT_MESSAGE_CONTENT"]
+    assert len(text_content_events) > 0
+
+    # Concatenate all text content
+    full_text = "".join(e.delta for e in text_content_events)
+    assert "executing" in full_text.lower()
+    assert "2 approved steps" in full_text.lower()
+    assert "send email" in full_text.lower()
+    assert "create calendar event" in full_text.lower()
+
+
+async def test_tool_result_function_approval_rejected():
+    """Test function approval tool result when rejected."""
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            yield ChatResponseUpdate(contents=[TextContent(text="OK")])
+
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    wrapper = AgentFrameworkAgent(agent=agent)
+
+    # Simulate tool result rejection with steps
+    tool_result = {
+        "accepted": False,
+        "steps": [{"id": "step1", "description": "Send email", "status": "disabled"}],
+    }
+    input_data = {
+        "messages": [
+            {
+                "role": "tool",
+                "content": json.dumps(tool_result),
+                "toolCallId": "approval_call_123",
+            }
+        ],
+    }
+
+    events = []
+    async for event in wrapper.run_agent(input_data):
+        events.append(event)
+
+    # Should ask what to change about the plan
+    text_content_events = [e for e in events if e.type == "TEXT_MESSAGE_CONTENT"]
+    assert len(text_content_events) > 0
+    assert any("what would you like me to change about the plan" in e.delta.lower() for e in text_content_events)
+
+
+async def test_thread_metadata_tracking():
+    """Test that thread metadata includes ag_ui_thread_id and ag_ui_run_id."""
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    thread_metadata = {}
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            # Capture thread metadata from kwargs
+            nonlocal thread_metadata
+            if "thread" in kwargs:
+                thread_metadata = kwargs["thread"].metadata
+            yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
+
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    wrapper = AgentFrameworkAgent(agent=agent)
+
+    input_data = {
+        "messages": [{"role": "user", "content": "Hi"}],
+        "thread_id": "test_thread_123",
+        "run_id": "test_run_456",
+    }
+
+    events = []
+    async for event in wrapper.run_agent(input_data):
+        events.append(event)
+
+    # Check thread metadata was set
+    # Note: This test may need adjustment based on actual thread passing mechanism
+
+
+async def test_state_context_injection():
+    """Test that current state is injected into thread metadata."""
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    thread_metadata = {}
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            # Track if state context message was added
+            nonlocal thread_metadata
+            # In actual implementation, thread is passed and state is in metadata
+            yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
+
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    wrapper = AgentFrameworkAgent(
+        agent=agent,
+        state_schema={"document": {"type": "string"}},
+    )
+
+    input_data = {
+        "messages": [{"role": "user", "content": "Hi"}],
+        "state": {"document": "Test content"},
+    }
+
+    events = []
+    async for event in wrapper.run_agent(input_data):
+        events.append(event)
+
+    # State should be injected - this is validated by agent execution flow
+
+
+async def test_no_messages_provided():
+    """Test handling when no messages are provided."""
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            yield ChatResponseUpdate(contents=[TextContent(text="Hello")])
+
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    wrapper = AgentFrameworkAgent(agent=agent)
+
+    input_data = {"messages": []}
+
+    events = []
+    async for event in wrapper.run_agent(input_data):
+        events.append(event)
+
+    # Should emit RunStartedEvent and RunFinishedEvent only
+    assert len(events) == 2
+    assert events[0].type == "RUN_STARTED"
+    assert events[-1].type == "RUN_FINISHED"
+
+
+async def test_message_end_event_emission():
+    """Test TextMessageEndEvent is emitted for assistant messages."""
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            yield ChatResponseUpdate(contents=[TextContent(text="Hello world")])
+
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    wrapper = AgentFrameworkAgent(agent=agent)
+
+    input_data = {"messages": [{"role": "user", "content": "Hi"}]}
+
+    events = []
+    async for event in wrapper.run_agent(input_data):
+        events.append(event)
+
+    # Should have TextMessageEndEvent before RunFinishedEvent
+    end_events = [e for e in events if e.type == "TEXT_MESSAGE_END"]
+    assert len(end_events) == 1
+
+    # EndEvent should come before FinishedEvent
+    end_index = events.index(end_events[0])
+    finished_index = events.index([e for e in events if e.type == "RUN_FINISHED"][0])
+    assert end_index < finished_index
+
+
+async def test_error_handling_with_exception():
+    """Test that exceptions during agent execution are re-raised."""
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    class FailingChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            if False:
+                yield
+            raise RuntimeError("Simulated failure")
+
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=FailingChatClient())
+    wrapper = AgentFrameworkAgent(agent=agent)
+
+    input_data = {"messages": [{"role": "user", "content": "Hi"}]}
+
+    with pytest.raises(RuntimeError, match="Simulated failure"):
+        async for event in wrapper.run_agent(input_data):
+            pass
+
+
+async def test_json_decode_error_in_tool_result():
+    """Test handling of JSONDecodeError when parsing tool result."""
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            yield ChatResponseUpdate(contents=[TextContent(text="Fallback response")])
+
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    wrapper = AgentFrameworkAgent(agent=agent)
+
+    # Send invalid JSON as tool result
+    input_data = {
+        "messages": [
+            {
+                "role": "tool",
+                "content": "invalid json {not valid}",
+                "toolCallId": "call_123",
+            }
+        ],
+    }
+
+    events = []
+    async for event in wrapper.run_agent(input_data):
+        events.append(event)
+
+    # Should fall through to normal agent processing
+    text_events = [e for e in events if e.type == "TEXT_MESSAGE_CONTENT"]
+    assert len(text_events) > 0
+    assert text_events[0].delta == "Fallback response"
+
+
+async def test_suppressed_summary_with_document_state():
+    """Test suppressed summary uses document state for confirmation message."""
+    from agent_framework_ag_ui import AgentFrameworkAgent, DocumentWriterConfirmationStrategy
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            yield ChatResponseUpdate(contents=[TextContent(text="Response")])
+
+    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
+    wrapper = AgentFrameworkAgent(
+        agent=agent,
+        state_schema={"document": {"type": "string"}},
+        predict_state_config={"document": {"tool": "write_doc", "tool_argument": "content"}},
+        confirmation_strategy=DocumentWriterConfirmationStrategy(),
+    )
+
+    # Simulate confirmation with document state
+    tool_result = {"accepted": True, "steps": []}
+    input_data = {
+        "messages": [
+            {
+                "role": "tool",
+                "content": json.dumps(tool_result),
+                "toolCallId": "confirm_123",
+            }
+        ],
+        "state": {"document": "This is the beginning of a document. It contains important information."},
+    }
+
+    events = []
+    async for event in wrapper.run_agent(input_data):
+        events.append(event)
+
+    # Should generate fallback summary from document state
+    text_events = [e for e in events if e.type == "TEXT_MESSAGE_CONTENT"]
+    assert len(text_events) > 0
+    # Should contain some reference to the document
+    full_text = "".join(e.delta for e in text_events)
+    assert "written" in full_text.lower() or "document" in full_text.lower()

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

@@ -0,0 +1,124 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for backend tool rendering."""
+
+from ag_ui.core import (
+    TextMessageContentEvent,
+    TextMessageStartEvent,
+    ToolCallArgsEvent,
+    ToolCallEndEvent,
+    ToolCallResultEvent,
+    ToolCallStartEvent,
+)
+from agent_framework import AgentRunResponseUpdate, FunctionCallContent, FunctionResultContent, TextContent
+
+from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+
+async def test_tool_call_flow():
+    """Test complete tool call flow: call -> args -> end -> result."""
+    bridge = AgentFrameworkEventBridge(run_id="test-run", thread_id="test-thread")
+
+    # Step 1: Tool call starts
+    tool_call = FunctionCallContent(
+        call_id="weather-123",
+        name="get_weather",
+        arguments={"location": "Seattle"},
+    )
+
+    update1 = AgentRunResponseUpdate(contents=[tool_call])
+    events1 = await bridge.from_agent_run_update(update1)
+
+    # Should have: ToolCallStartEvent, ToolCallArgsEvent
+    assert len(events1) == 2
+    assert isinstance(events1[0], ToolCallStartEvent)
+    assert isinstance(events1[1], ToolCallArgsEvent)
+
+    start_event = events1[0]
+    assert start_event.tool_call_id == "weather-123"
+    assert start_event.tool_call_name == "get_weather"
+
+    args_event = events1[1]
+    assert "Seattle" in args_event.delta
+
+    # Step 2: Tool result comes back
+    tool_result = FunctionResultContent(
+        call_id="weather-123",
+        result="Weather in Seattle: Rainy, 52°F",
+    )
+
+    update2 = AgentRunResponseUpdate(contents=[tool_result])
+    events2 = await bridge.from_agent_run_update(update2)
+
+    # Should have: ToolCallEndEvent, ToolCallResultEvent, MessagesSnapshotEvent
+    assert len(events2) == 3
+    assert isinstance(events2[0], ToolCallEndEvent)
+    assert isinstance(events2[1], ToolCallResultEvent)
+
+    end_event = events2[0]
+    assert end_event.tool_call_id == "weather-123"
+
+    result_event = events2[1]
+    assert result_event.tool_call_id == "weather-123"
+    assert "Seattle" in result_event.content
+    assert "Rainy" in result_event.content
+
+
+async def test_text_with_tool_call():
+    """Test agent response with both text and tool calls."""
+    bridge = AgentFrameworkEventBridge(run_id="test-run", thread_id="test-thread")
+
+    # Agent says something then calls a tool
+    text_content = TextContent(text="Let me check the weather for you.")
+    tool_call = FunctionCallContent(
+        call_id="weather-456",
+        name="get_forecast",
+        arguments={"location": "San Francisco", "days": 3},
+    )
+
+    update = AgentRunResponseUpdate(contents=[text_content, tool_call])
+    events = await bridge.from_agent_run_update(update)
+
+    # Should have: TextMessageStart, TextMessageContent, ToolCallStart, ToolCallArgs
+    assert len(events) == 4
+
+    assert isinstance(events[0], TextMessageStartEvent)
+    assert isinstance(events[1], TextMessageContentEvent)
+    assert isinstance(events[2], ToolCallStartEvent)
+    assert isinstance(events[3], ToolCallArgsEvent)
+
+    text_event = events[1]
+    assert "check the weather" in text_event.delta
+
+    tool_start = events[2]
+    assert tool_start.tool_call_name == "get_forecast"
+
+
+async def test_multiple_tool_results():
+    """Test handling multiple tool results in sequence."""
+    bridge = AgentFrameworkEventBridge(run_id="test-run", thread_id="test-thread")
+
+    # Multiple tool results
+    results = [
+        FunctionResultContent(call_id="tool-1", result="Result 1"),
+        FunctionResultContent(call_id="tool-2", result="Result 2"),
+        FunctionResultContent(call_id="tool-3", result="Result 3"),
+    ]
+
+    update = AgentRunResponseUpdate(contents=results)
+    events = await bridge.from_agent_run_update(update)
+
+    # Should have 3 pairs of ToolCallEndEvent + ToolCallResultEvent = 6 events
+    assert len(events) == 6
+
+    # Verify the pattern: End, Result, End, Result, End, Result
+    for i in range(3):
+        end_idx = i * 2
+        result_idx = i * 2 + 1
+
+        assert isinstance(events[end_idx], ToolCallEndEvent)
+        assert isinstance(events[result_idx], ToolCallResultEvent)
+
+        assert events[end_idx].tool_call_id == f"tool-{i + 1}"
+        assert events[result_idx].tool_call_id == f"tool-{i + 1}"
+        assert f"Result {i + 1}" in events[result_idx].content

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

@@ -0,0 +1,317 @@
+"""Tests for AGUIChatClient."""
+
+import json
+
+from agent_framework import ChatMessage, ChatOptions, FunctionCallContent, Role, ai_function
+
+from agent_framework_ag_ui._client import AGUIChatClient, ServerFunctionCallContent
+
+
+class TestAGUIChatClient:
+    """Test suite for AGUIChatClient."""
+
+    async def test_client_initialization(self) -> None:
+        """Test client initialization."""
+        client = AGUIChatClient(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 AGUIChatClient(endpoint="http://localhost:8888/") as client:
+            assert client is not None
+
+    async def test_extract_state_from_messages_no_state(self) -> None:
+        """Test state extraction when no state is present."""
+        client = AGUIChatClient(endpoint="http://localhost:8888/")
+        messages = [
+            ChatMessage(role="user", text="Hello"),
+            ChatMessage(role="assistant", text="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 = AGUIChatClient(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")
+
+        from agent_framework import DataContent
+
+        messages = [
+            ChatMessage(role="user", text="Hello"),
+            ChatMessage(
+                role="user",
+                contents=[DataContent(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 = AGUIChatClient(endpoint="http://localhost:8888/")
+
+        invalid_json = "not valid json"
+        state_b64 = base64.b64encode(invalid_json.encode("utf-8")).decode("utf-8")
+
+        from agent_framework import DataContent
+
+        messages = [
+            ChatMessage(
+                role="user",
+                contents=[DataContent(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 = AGUIChatClient(endpoint="http://localhost:8888/")
+        messages = [
+            ChatMessage(role=Role.USER, text="What is the weather?"),
+            ChatMessage(role=Role.ASSISTANT, text="Let me check.", message_id="msg_123"),
+        ]
+
+        agui_messages = client._convert_messages_to_agui_format(messages)
+
+        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 = AGUIChatClient(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 = AGUIChatClient(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_streaming_response(self, 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, **kwargs):
+            for event in mock_events:
+                yield event
+
+        client = AGUIChatClient(endpoint="http://localhost:8888/")
+        monkeypatch.setattr(client._http_service, "post_run", mock_post_run)
+
+        messages = [ChatMessage(role="user", text="Test message")]
+        chat_options = ChatOptions()
+
+        updates = []
+        async for update in client._inner_get_streaming_response(messages=messages, chat_options=chat_options):
+            updates.append(update)
+
+        assert len(updates) == 4
+        assert updates[0].additional_properties["thread_id"] == "thread_1"
+        assert updates[1].contents[0].text == "Hello"
+        assert updates[2].contents[0].text == " world"
+
+    async def test_get_response_non_streaming(self, monkeypatch) -> None:
+        """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, **kwargs):
+            for event in mock_events:
+                yield event
+
+        client = AGUIChatClient(endpoint="http://localhost:8888/")
+        monkeypatch.setattr(client._http_service, "post_run", mock_post_run)
+
+        messages = [ChatMessage(role="user", text="Test message")]
+        chat_options = ChatOptions()
+
+        response = await client._inner_get_response(messages=messages, chat_options=chat_options)
+
+        assert response is not None
+        assert len(response.messages) > 0
+        assert "Complete response" in response.text
+
+    async def test_tool_handling(self, 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, @use_function_invocation decorator
+        intercepts and executes it locally. This matches .NET AG-UI implementation.
+        """
+        from agent_framework import ai_function
+
+        @ai_function
+        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, **kwargs):
+            # Client tool metadata should be sent to server
+            tools = kwargs.get("tools")
+            assert tools is not None
+            assert len(tools) == 1
+            assert tools[0]["name"] == "test_tool"
+            assert tools[0]["description"] == "Test tool."
+            assert "parameters" in tools[0]
+            for event in mock_events:
+                yield event
+
+        client = AGUIChatClient(endpoint="http://localhost:8888/")
+        monkeypatch.setattr(client._http_service, "post_run", mock_post_run)
+
+        messages = [ChatMessage(role="user", text="Test with tools")]
+        chat_options = ChatOptions(tools=[test_tool])
+
+        response = await client._inner_get_response(messages=messages, chat_options=chat_options)
+
+        assert response is not None
+
+    async def test_server_tool_calls_unwrapped_after_invocation(self, 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, **kwargs):
+            for event in mock_events:
+                yield event
+
+        client = AGUIChatClient(endpoint="http://localhost:8888/")
+        monkeypatch.setattr(client._http_service, "post_run", mock_post_run)
+
+        messages = [ChatMessage(role="user", text="Test server tool execution")]
+        chat_options = ChatOptions()
+
+        updates = []
+        async for update in client.get_streaming_response(messages, chat_options=chat_options):
+            updates.append(update)
+
+        function_calls = [
+            content for update in updates for content in update.contents if isinstance(content, FunctionCallContent)
+        ]
+        assert function_calls
+        assert function_calls[0].name == "get_time_zone"
+        assert not any(
+            isinstance(content, ServerFunctionCallContent) for update in updates for content in update.contents
+        )
+
+    async def test_server_tool_calls_not_executed_locally(self, monkeypatch) -> None:
+        """Server tools should not trigger local function invocation even when client tools exist."""
+
+        @ai_function
+        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, **kwargs):
+            for event in mock_events:
+                yield event
+
+        async def fake_auto_invoke(*args, **kwargs):
+            function_call = kwargs.get("function_call_content") or args[0]
+            raise AssertionError(f"Unexpected local execution of server tool: {getattr(function_call, 'name', '?')}")
+
+        monkeypatch.setattr("agent_framework._tools._auto_invoke_function", fake_auto_invoke)
+
+        client = AGUIChatClient(endpoint="http://localhost:8888/")
+        monkeypatch.setattr(client._http_service, "post_run", mock_post_run)
+
+        messages = [ChatMessage(role="user", text="Test server tool execution")]
+        chat_options = ChatOptions(tool_choice="auto", tools=[client_tool])
+
+        async for _ in client.get_streaming_response(messages, chat_options=chat_options):
+            pass
+
+    async def test_state_transmission(self, 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")
+
+        from agent_framework import DataContent
+
+        messages = [
+            ChatMessage(role="user", text="Hello"),
+            ChatMessage(
+                role="user",
+                contents=[DataContent(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, **kwargs):
+            assert kwargs.get("state") == state_data
+            for event in mock_events:
+                yield event
+
+        client = AGUIChatClient(endpoint="http://localhost:8888/")
+        monkeypatch.setattr(client._http_service, "post_run", mock_post_run)
+
+        chat_options = ChatOptions()
+
+        response = await client._inner_get_response(messages=messages, chat_options=chat_options)
+
+        assert response is not None

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

@@ -0,0 +1,275 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Comprehensive tests for all confirmation strategies."""
+
+import pytest
+
+from agent_framework_ag_ui._confirmation_strategies import (
+    ConfirmationStrategy,
+    DefaultConfirmationStrategy,
+    DocumentWriterConfirmationStrategy,
+    RecipeConfirmationStrategy,
+    TaskPlannerConfirmationStrategy,
+)
+
+
+@pytest.fixture
+def sample_steps():
+    """Sample steps for testing approval messages."""
+    return [
+        {"description": "Step 1: Do something", "status": "enabled"},
+        {"description": "Step 2: Do another thing", "status": "enabled"},
+        {"description": "Step 3: Disabled step", "status": "disabled"},
+    ]
+
+
+@pytest.fixture
+def all_enabled_steps():
+    """All steps enabled."""
+    return [
+        {"description": "Task A", "status": "enabled"},
+        {"description": "Task B", "status": "enabled"},
+        {"description": "Task C", "status": "enabled"},
+    ]
+
+
+@pytest.fixture
+def empty_steps():
+    """Empty steps list."""
+    return []
+
+
+class TestDefaultConfirmationStrategy:
+    """Tests for DefaultConfirmationStrategy."""
+
+    def test_on_approval_accepted_with_enabled_steps(self, sample_steps):
+        strategy = DefaultConfirmationStrategy()
+        message = strategy.on_approval_accepted(sample_steps)
+
+        assert "Executing 2 approved steps" in message
+        assert "Step 1: Do something" in message
+        assert "Step 2: Do another thing" in message
+        assert "Step 3" not in message  # Disabled step shouldn't appear
+        assert "All steps completed successfully!" in message
+
+    def test_on_approval_accepted_with_all_enabled(self, all_enabled_steps):
+        strategy = DefaultConfirmationStrategy()
+        message = strategy.on_approval_accepted(all_enabled_steps)
+
+        assert "Executing 3 approved steps" in message
+        assert "Task A" in message
+        assert "Task B" in message
+        assert "Task C" in message
+
+    def test_on_approval_accepted_with_empty_steps(self, empty_steps):
+        strategy = DefaultConfirmationStrategy()
+        message = strategy.on_approval_accepted(empty_steps)
+
+        assert "Executing 0 approved steps" in message
+        assert "All steps completed successfully!" in message
+
+    def test_on_approval_rejected(self, sample_steps):
+        strategy = DefaultConfirmationStrategy()
+        message = strategy.on_approval_rejected(sample_steps)
+
+        assert "No problem!" in message
+        assert "What would you like me to change" in message
+
+    def test_on_state_confirmed(self):
+        strategy = DefaultConfirmationStrategy()
+        message = strategy.on_state_confirmed()
+
+        assert "Changes confirmed" in message
+        assert "successfully" in message
+
+    def test_on_state_rejected(self):
+        strategy = DefaultConfirmationStrategy()
+        message = strategy.on_state_rejected()
+
+        assert "No problem!" in message
+        assert "What would you like me to change" in message
+
+
+class TestTaskPlannerConfirmationStrategy:
+    """Tests for TaskPlannerConfirmationStrategy."""
+
+    def test_on_approval_accepted_with_enabled_steps(self, sample_steps):
+        strategy = TaskPlannerConfirmationStrategy()
+        message = strategy.on_approval_accepted(sample_steps)
+
+        assert "Executing your requested tasks" in message
+        assert "1. Step 1: Do something" in message
+        assert "2. Step 2: Do another thing" in message
+        assert "Step 3" not in message
+        assert "All tasks completed successfully!" in message
+
+    def test_on_approval_accepted_with_all_enabled(self, all_enabled_steps):
+        strategy = TaskPlannerConfirmationStrategy()
+        message = strategy.on_approval_accepted(all_enabled_steps)
+
+        assert "Executing your requested tasks" in message
+        assert "1. Task A" in message
+        assert "2. Task B" in message
+        assert "3. Task C" in message
+
+    def test_on_approval_accepted_with_empty_steps(self, empty_steps):
+        strategy = TaskPlannerConfirmationStrategy()
+        message = strategy.on_approval_accepted(empty_steps)
+
+        assert "Executing your requested tasks" in message
+        assert "All tasks completed successfully!" in message
+
+    def test_on_approval_rejected(self, sample_steps):
+        strategy = TaskPlannerConfirmationStrategy()
+        message = strategy.on_approval_rejected(sample_steps)
+
+        assert "No problem!" in message
+        assert "revise the plan" in message
+
+    def test_on_state_confirmed(self):
+        strategy = TaskPlannerConfirmationStrategy()
+        message = strategy.on_state_confirmed()
+
+        assert "Tasks confirmed" in message
+        assert "ready to execute" in message
+
+    def test_on_state_rejected(self):
+        strategy = TaskPlannerConfirmationStrategy()
+        message = strategy.on_state_rejected()
+
+        assert "No problem!" in message
+        assert "adjust the task list" in message
+
+
+class TestRecipeConfirmationStrategy:
+    """Tests for RecipeConfirmationStrategy."""
+
+    def test_on_approval_accepted_with_enabled_steps(self, sample_steps):
+        strategy = RecipeConfirmationStrategy()
+        message = strategy.on_approval_accepted(sample_steps)
+
+        assert "Updating your recipe" in message
+        assert "1. Step 1: Do something" in message
+        assert "2. Step 2: Do another thing" in message
+        assert "Step 3" not in message
+        assert "Recipe updated successfully!" in message
+
+    def test_on_approval_accepted_with_all_enabled(self, all_enabled_steps):
+        strategy = RecipeConfirmationStrategy()
+        message = strategy.on_approval_accepted(all_enabled_steps)
+
+        assert "Updating your recipe" in message
+        assert "1. Task A" in message
+        assert "2. Task B" in message
+        assert "3. Task C" in message
+
+    def test_on_approval_accepted_with_empty_steps(self, empty_steps):
+        strategy = RecipeConfirmationStrategy()
+        message = strategy.on_approval_accepted(empty_steps)
+
+        assert "Updating your recipe" in message
+        assert "Recipe updated successfully!" in message
+
+    def test_on_approval_rejected(self, sample_steps):
+        strategy = RecipeConfirmationStrategy()
+        message = strategy.on_approval_rejected(sample_steps)
+
+        assert "No problem!" in message
+        assert "ingredients or steps" in message
+
+    def test_on_state_confirmed(self):
+        strategy = RecipeConfirmationStrategy()
+        message = strategy.on_state_confirmed()
+
+        assert "Recipe changes applied" in message
+        assert "successfully" in message
+
+    def test_on_state_rejected(self):
+        strategy = RecipeConfirmationStrategy()
+        message = strategy.on_state_rejected()
+
+        assert "No problem!" in message
+        assert "adjust in the recipe" in message
+
+
+class TestDocumentWriterConfirmationStrategy:
+    """Tests for DocumentWriterConfirmationStrategy."""
+
+    def test_on_approval_accepted_with_enabled_steps(self, sample_steps):
+        strategy = DocumentWriterConfirmationStrategy()
+        message = strategy.on_approval_accepted(sample_steps)
+
+        assert "Applying your edits" in message
+        assert "1. Step 1: Do something" in message
+        assert "2. Step 2: Do another thing" in message
+        assert "Step 3" not in message
+        assert "Document updated successfully!" in message
+
+    def test_on_approval_accepted_with_all_enabled(self, all_enabled_steps):
+        strategy = DocumentWriterConfirmationStrategy()
+        message = strategy.on_approval_accepted(all_enabled_steps)
+
+        assert "Applying your edits" in message
+        assert "1. Task A" in message
+        assert "2. Task B" in message
+        assert "3. Task C" in message
+
+    def test_on_approval_accepted_with_empty_steps(self, empty_steps):
+        strategy = DocumentWriterConfirmationStrategy()
+        message = strategy.on_approval_accepted(empty_steps)
+
+        assert "Applying your edits" in message
+        assert "Document updated successfully!" in message
+
+    def test_on_approval_rejected(self, sample_steps):
+        strategy = DocumentWriterConfirmationStrategy()
+        message = strategy.on_approval_rejected(sample_steps)
+
+        assert "No problem!" in message
+        assert "keep or modify" in message
+
+    def test_on_state_confirmed(self):
+        strategy = DocumentWriterConfirmationStrategy()
+        message = strategy.on_state_confirmed()
+
+        assert "Document edits applied!" in message
+
+    def test_on_state_rejected(self):
+        strategy = DocumentWriterConfirmationStrategy()
+        message = strategy.on_state_rejected()
+
+        assert "No problem!" in message
+        assert "change about the document" in message
+
+
+class TestConfirmationStrategyInterface:
+    """Tests for ConfirmationStrategy abstract base class."""
+
+    def test_cannot_instantiate_abstract_class(self):
+        """Verify ConfirmationStrategy is abstract and cannot be instantiated."""
+        with pytest.raises(TypeError):
+            ConfirmationStrategy()  # type: ignore
+
+    def test_all_strategies_implement_interface(self):
+        """Verify all concrete strategies implement the full interface."""
+        strategies = [
+            DefaultConfirmationStrategy(),
+            TaskPlannerConfirmationStrategy(),
+            RecipeConfirmationStrategy(),
+            DocumentWriterConfirmationStrategy(),
+        ]
+
+        sample_steps = [{"description": "Test", "status": "enabled"}]
+
+        for strategy in strategies:
+            # All should have these methods
+            assert callable(strategy.on_approval_accepted)
+            assert callable(strategy.on_approval_rejected)
+            assert callable(strategy.on_state_confirmed)
+            assert callable(strategy.on_state_rejected)
+
+            # All should return strings
+            assert isinstance(strategy.on_approval_accepted(sample_steps), str)
+            assert isinstance(strategy.on_approval_rejected(sample_steps), str)
+            assert isinstance(strategy.on_state_confirmed(), str)
+            assert isinstance(strategy.on_state_rejected(), str)

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

@@ -0,0 +1,243 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for document writer predictive state flow with confirm_changes."""
+
+from ag_ui.core import EventType
+from agent_framework import FunctionCallContent, FunctionResultContent, TextContent
+from agent_framework._types import AgentRunResponseUpdate
+
+from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+
+async def test_streaming_document_with_state_deltas():
+    """Test that streaming tool arguments emit progressive StateDeltaEvents."""
+    predict_config = {
+        "document": {"tool": "write_document_local", "tool_argument": "document"},
+    }
+
+    bridge = AgentFrameworkEventBridge(
+        run_id="test_run",
+        thread_id="test_thread",
+        predict_state_config=predict_config,
+    )
+
+    # Simulate streaming tool call - first chunk with name
+    tool_call_start = FunctionCallContent(
+        call_id="call_123",
+        name="write_document_local",
+        arguments='{"document":"Once',
+    )
+    update1 = AgentRunResponseUpdate(contents=[tool_call_start])
+    events1 = await bridge.from_agent_run_update(update1)
+
+    # Should have ToolCallStartEvent and ToolCallArgsEvent
+    assert any(e.type == EventType.TOOL_CALL_START for e in events1)
+    assert any(e.type == EventType.TOOL_CALL_ARGS for e in events1)
+
+    # Second chunk - incomplete JSON, should try partial extraction
+    tool_call_chunk2 = FunctionCallContent(
+        call_id="call_123",
+        name=None,  # Name only in first chunk
+        arguments=" upon a time",
+    )
+    update2 = AgentRunResponseUpdate(contents=[tool_call_chunk2])
+    events2 = await bridge.from_agent_run_update(update2)
+
+    # Should emit StateDeltaEvent with partial document
+    state_deltas = [e for e in events2 if e.type == EventType.STATE_DELTA]
+    assert len(state_deltas) >= 1
+
+    # Check JSON Patch format
+    delta = state_deltas[0]
+    assert isinstance(delta.delta, list)
+    assert len(delta.delta) > 0
+    assert delta.delta[0]["op"] == "replace"
+    assert delta.delta[0]["path"] == "/document"
+    assert "Once upon a time" in delta.delta[0]["value"]
+
+
+async def test_confirm_changes_emission():
+    """Test that confirm_changes tool call is emitted after predictive tool completion."""
+    predict_config = {
+        "document": {"tool": "write_document_local", "tool_argument": "document"},
+    }
+
+    current_state = {}
+
+    bridge = AgentFrameworkEventBridge(
+        run_id="test_run",
+        thread_id="test_thread",
+        predict_state_config=predict_config,
+        current_state=current_state,
+    )
+
+    # Set current tool name (simulating earlier tool call start)
+    bridge.current_tool_call_name = "write_document_local"
+    bridge.pending_state_updates = {"document": "A short story"}
+
+    # Tool result
+    tool_result = FunctionResultContent(
+        call_id="call_123",
+        result="Document written.",
+    )
+
+    update = AgentRunResponseUpdate(contents=[tool_result])
+    events = await bridge.from_agent_run_update(update)
+
+    # Should have: ToolCallEndEvent, ToolCallResultEvent, StateSnapshotEvent, confirm_changes sequence
+    assert any(e.type == EventType.TOOL_CALL_END for e in events)
+    assert any(e.type == EventType.TOOL_CALL_RESULT for e in events)
+    assert any(e.type == EventType.STATE_SNAPSHOT for e in events)
+
+    # Check for confirm_changes tool call
+    confirm_starts = [
+        e for e in events if e.type == EventType.TOOL_CALL_START and e.tool_call_name == "confirm_changes"
+    ]
+    assert len(confirm_starts) == 1
+
+    confirm_args = [e for e in events if e.type == EventType.TOOL_CALL_ARGS and e.delta == "{}"]
+    assert len(confirm_args) >= 1
+
+    confirm_ends = [e for e in events if e.type == EventType.TOOL_CALL_END]
+    # At least 2: one for write_document_local, one for confirm_changes
+    assert len(confirm_ends) >= 2
+
+    # Check that stop flag is set
+    assert bridge.should_stop_after_confirm is True
+
+
+async def test_text_suppression_before_confirm():
+    """Test that text messages are suppressed when confirm_changes is pending."""
+    predict_config = {
+        "document": {"tool": "write_document_local", "tool_argument": "document"},
+    }
+
+    bridge = AgentFrameworkEventBridge(
+        run_id="test_run",
+        thread_id="test_thread",
+        predict_state_config=predict_config,
+    )
+
+    # Set flag indicating we're waiting for confirmation
+    bridge.should_stop_after_confirm = True
+
+    # Text content that should be suppressed
+    text = TextContent(text="I have written a story about pirates.")
+    update = AgentRunResponseUpdate(contents=[text])
+
+    events = await bridge.from_agent_run_update(update)
+
+    # Should NOT emit TextMessageContentEvent
+    text_events = [e for e in events if e.type == EventType.TEXT_MESSAGE_CONTENT]
+    assert len(text_events) == 0
+
+    # But should save the text
+    assert bridge.suppressed_summary == "I have written a story about pirates."
+
+
+async def test_no_confirm_for_non_predictive_tools():
+    """Test that confirm_changes is NOT emitted for regular tool calls."""
+    predict_config = {
+        "document": {"tool": "write_document_local", "tool_argument": "document"},
+    }
+
+    current_state = {}
+
+    bridge = AgentFrameworkEventBridge(
+        run_id="test_run",
+        thread_id="test_thread",
+        predict_state_config=predict_config,
+        current_state=current_state,
+    )
+
+    # Different tool (not in predict_state_config)
+    bridge.current_tool_call_name = "get_weather"
+
+    tool_result = FunctionResultContent(
+        call_id="call_456",
+        result="Sunny, 72°F",
+    )
+
+    update = AgentRunResponseUpdate(contents=[tool_result])
+    events = await bridge.from_agent_run_update(update)
+
+    # Should NOT have confirm_changes
+    confirm_starts = [
+        e for e in events if e.type == EventType.TOOL_CALL_START and e.tool_call_name == "confirm_changes"
+    ]
+    assert len(confirm_starts) == 0
+
+    # Stop flag should NOT be set
+    assert bridge.should_stop_after_confirm is False
+
+
+async def test_state_delta_deduplication():
+    """Test that duplicate state values don't emit multiple StateDeltaEvents."""
+    predict_config = {
+        "document": {"tool": "write_document_local", "tool_argument": "document"},
+    }
+
+    bridge = AgentFrameworkEventBridge(
+        run_id="test_run",
+        thread_id="test_thread",
+        predict_state_config=predict_config,
+    )
+
+    # First tool call with document
+    tool_call1 = FunctionCallContent(
+        call_id="call_1",
+        name="write_document_local",
+        arguments='{"document":"Same text"}',
+    )
+    update1 = AgentRunResponseUpdate(contents=[tool_call1])
+    events1 = await bridge.from_agent_run_update(update1)
+
+    # Count state deltas
+    state_deltas_1 = [e for e in events1 if e.type == EventType.STATE_DELTA]
+    assert len(state_deltas_1) >= 1
+
+    # Second tool call with SAME document (shouldn't emit new delta)
+    bridge.current_tool_call_name = "write_document_local"
+    tool_call2 = FunctionCallContent(
+        call_id="call_2",
+        name=None,
+        arguments='{"document":"Same text"}',  # Identical content
+    )
+    update2 = AgentRunResponseUpdate(contents=[tool_call2])
+    events2 = await bridge.from_agent_run_update(update2)
+
+    # Should NOT emit state delta (same value)
+    state_deltas_2 = [e for e in events2 if e.type == EventType.STATE_DELTA]
+    assert len(state_deltas_2) == 0
+
+
+async def test_predict_state_config_multiple_fields():
+    """Test predictive state with multiple state fields."""
+    predict_config = {
+        "title": {"tool": "create_post", "tool_argument": "title"},
+        "content": {"tool": "create_post", "tool_argument": "body"},
+    }
+
+    bridge = AgentFrameworkEventBridge(
+        run_id="test_run",
+        thread_id="test_thread",
+        predict_state_config=predict_config,
+    )
+
+    # Tool call with both fields
+    tool_call = FunctionCallContent(
+        call_id="call_999",
+        name="create_post",
+        arguments='{"title":"My Post","body":"Post content"}',
+    )
+    update = AgentRunResponseUpdate(contents=[tool_call])
+    events = await bridge.from_agent_run_update(update)
+
+    # Should emit StateDeltaEvent for both fields
+    state_deltas = [e for e in events if e.type == EventType.STATE_DELTA]
+    assert len(state_deltas) >= 2
+
+    # Check both fields are present
+    paths = [delta.delta[0]["path"] for delta in state_deltas]
+    assert "/title" in paths
+    assert "/content" in paths

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

@@ -0,0 +1,242 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for FastAPI endpoint creation (_endpoint.py)."""
+
+import json
+from typing import Any
+
+from agent_framework import ChatAgent, TextContent
+from agent_framework._types import ChatResponseUpdate
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+from agent_framework_ag_ui._agent import AgentFrameworkAgent
+from agent_framework_ag_ui._endpoint import add_agent_framework_fastapi_endpoint
+
+
+class MockChatClient:
+    """Mock chat client for testing."""
+
+    def __init__(self, response_text: str = "Test response"):
+        self.response_text = response_text
+
+    async def get_streaming_response(self, messages: list[Any], chat_options: Any, **kwargs: Any):
+        """Mock streaming response."""
+        yield ChatResponseUpdate(contents=[TextContent(text=self.response_text)])
+
+
+async def test_add_endpoint_with_agent_protocol():
+    """Test adding endpoint with raw AgentProtocol."""
+    app = FastAPI()
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+
+    add_agent_framework_fastapi_endpoint(app, agent, path="/test-agent")
+
+    client = TestClient(app)
+    response = client.post("/test-agent", json={"messages": [{"role": "user", "content": "Hello"}]})
+
+    assert response.status_code == 200
+    assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
+
+
+async def test_add_endpoint_with_wrapped_agent():
+    """Test adding endpoint with pre-wrapped AgentFrameworkAgent."""
+    app = FastAPI()
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+    wrapped_agent = AgentFrameworkAgent(agent=agent, name="wrapped")
+
+    add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/wrapped-agent")
+
+    client = TestClient(app)
+    response = client.post("/wrapped-agent", json={"messages": [{"role": "user", "content": "Hello"}]})
+
+    assert response.status_code == 200
+    assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
+
+
+async def test_endpoint_with_state_schema():
+    """Test endpoint with state_schema parameter."""
+    app = FastAPI()
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+    state_schema = {"document": {"type": "string"}}
+
+    add_agent_framework_fastapi_endpoint(app, agent, path="/stateful", state_schema=state_schema)
+
+    client = TestClient(app)
+    response = client.post(
+        "/stateful", json={"messages": [{"role": "user", "content": "Hello"}], "state": {"document": ""}}
+    )
+
+    assert response.status_code == 200
+
+
+async def test_endpoint_with_predict_state_config():
+    """Test endpoint with predict_state_config parameter."""
+    app = FastAPI()
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+    predict_config = {"document": {"tool": "write_doc", "tool_argument": "content"}}
+
+    add_agent_framework_fastapi_endpoint(app, agent, path="/predictive", predict_state_config=predict_config)
+
+    client = TestClient(app)
+    response = client.post("/predictive", json={"messages": [{"role": "user", "content": "Hello"}]})
+
+    assert response.status_code == 200
+
+
+async def test_endpoint_request_logging():
+    """Test that endpoint logs request details."""
+    app = FastAPI()
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+
+    add_agent_framework_fastapi_endpoint(app, agent, path="/logged")
+
+    client = TestClient(app)
+    response = client.post(
+        "/logged",
+        json={
+            "messages": [{"role": "user", "content": "Test"}],
+            "run_id": "run-123",
+            "thread_id": "thread-456",
+        },
+    )
+
+    assert response.status_code == 200
+
+
+async def test_endpoint_event_streaming():
+    """Test that endpoint streams events correctly."""
+    app = FastAPI()
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient("Streamed response"))
+
+    add_agent_framework_fastapi_endpoint(app, agent, path="/stream")
+
+    client = TestClient(app)
+    response = client.post("/stream", json={"messages": [{"role": "user", "content": "Hello"}]})
+
+    assert response.status_code == 200
+
+    content = response.content.decode("utf-8")
+    lines = [line for line in content.split("\n") if line.strip()]
+
+    found_run_started = False
+    found_text_content = False
+    found_run_finished = False
+
+    for line in lines:
+        if line.startswith("data: "):
+            event_data = json.loads(line[6:])
+            if event_data.get("type") == "RUN_STARTED":
+                found_run_started = True
+            elif event_data.get("type") == "TEXT_MESSAGE_CONTENT":
+                found_text_content = True
+            elif event_data.get("type") == "RUN_FINISHED":
+                found_run_finished = True
+
+    assert found_run_started
+    assert found_text_content
+    assert found_run_finished
+
+
+async def test_endpoint_error_handling():
+    """Test endpoint error handling during request parsing."""
+    app = FastAPI()
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+
+    add_agent_framework_fastapi_endpoint(app, agent, path="/failing")
+
+    client = TestClient(app)
+
+    # Send invalid JSON to trigger parsing error before streaming
+    response = client.post("/failing", data="invalid json", headers={"content-type": "application/json"})
+
+    # The exception handler catches it and returns JSON error
+    assert response.status_code == 200
+    content = json.loads(response.content)
+    assert "error" in content
+    assert "Expecting value" in content["error"]
+
+
+async def test_endpoint_multiple_paths():
+    """Test adding multiple endpoints with different paths."""
+    app = FastAPI()
+    agent1 = ChatAgent(name="agent1", instructions="First agent", chat_client=MockChatClient("Response 1"))
+    agent2 = ChatAgent(name="agent2", instructions="Second agent", chat_client=MockChatClient("Response 2"))
+
+    add_agent_framework_fastapi_endpoint(app, agent1, path="/agent1")
+    add_agent_framework_fastapi_endpoint(app, agent2, path="/agent2")
+
+    client = TestClient(app)
+
+    response1 = client.post("/agent1", json={"messages": [{"role": "user", "content": "Hi"}]})
+    response2 = client.post("/agent2", json={"messages": [{"role": "user", "content": "Hi"}]})
+
+    assert response1.status_code == 200
+    assert response2.status_code == 200
+
+
+async def test_endpoint_default_path():
+    """Test endpoint with default path."""
+    app = FastAPI()
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+
+    add_agent_framework_fastapi_endpoint(app, agent)
+
+    client = TestClient(app)
+    response = client.post("/", json={"messages": [{"role": "user", "content": "Hello"}]})
+
+    assert response.status_code == 200
+
+
+async def test_endpoint_response_headers():
+    """Test that endpoint sets correct response headers."""
+    app = FastAPI()
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+
+    add_agent_framework_fastapi_endpoint(app, agent, path="/headers")
+
+    client = TestClient(app)
+    response = client.post("/headers", json={"messages": [{"role": "user", "content": "Test"}]})
+
+    assert response.status_code == 200
+    assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
+    assert "cache-control" in response.headers
+    assert response.headers["cache-control"] == "no-cache"
+
+
+async def test_endpoint_empty_messages():
+    """Test endpoint with empty messages list."""
+    app = FastAPI()
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+
+    add_agent_framework_fastapi_endpoint(app, agent, path="/empty")
+
+    client = TestClient(app)
+    response = client.post("/empty", json={"messages": []})
+
+    assert response.status_code == 200
+
+
+async def test_endpoint_complex_input():
+    """Test endpoint with complex input data."""
+    app = FastAPI()
+    agent = ChatAgent(name="test", instructions="Test agent", chat_client=MockChatClient())
+
+    add_agent_framework_fastapi_endpoint(app, agent, path="/complex")
+
+    client = TestClient(app)
+    response = client.post(
+        "/complex",
+        json={
+            "messages": [
+                {"role": "user", "content": "First message", "id": "msg-1"},
+                {"role": "assistant", "content": "Response", "id": "msg-2"},
+                {"role": "user", "content": "Follow-up", "id": "msg-3"},
+            ],
+            "run_id": "complex-run-123",
+            "thread_id": "complex-thread-456",
+            "state": {"custom_field": "value"},
+        },
+    )
+
+    assert response.status_code == 200

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

@@ -0,0 +1,287 @@
+"""Tests for AG-UI event converter."""
+
+from agent_framework import FinishReason, Role
+
+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.role == 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 == 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 == 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 = [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 == 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 = [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 == 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 == Role.ASSISTANT
+        assert update.finish_reason == FinishReason.STOP
+        assert update.additional_properties["thread_id"] == "thread_123"
+        assert update.additional_properties["run_id"] == "run_456"
+
+    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 == Role.ASSISTANT
+        assert update.finish_reason == FinishReason.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_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 = [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 = [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"

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

@@ -0,0 +1,659 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Comprehensive tests for AgentFrameworkEventBridge (_events.py)."""
+
+import json
+
+from agent_framework import (
+    AgentRunResponseUpdate,
+    FunctionApprovalRequestContent,
+    FunctionCallContent,
+    FunctionResultContent,
+    TextContent,
+)
+
+
+async def test_basic_text_message_conversion():
+    """Test basic TextContent to AG-UI events."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
+
+    update = AgentRunResponseUpdate(contents=[TextContent(text="Hello")])
+    events = await bridge.from_agent_run_update(update)
+
+    assert len(events) == 2
+    assert events[0].type == "TEXT_MESSAGE_START"
+    assert events[0].role == "assistant"
+    assert events[1].type == "TEXT_MESSAGE_CONTENT"
+    assert events[1].delta == "Hello"
+
+
+async def test_text_message_streaming():
+    """Test streaming TextContent with multiple chunks."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
+
+    update1 = AgentRunResponseUpdate(contents=[TextContent(text="Hello ")])
+    update2 = AgentRunResponseUpdate(contents=[TextContent(text="world")])
+
+    events1 = await bridge.from_agent_run_update(update1)
+    events2 = await bridge.from_agent_run_update(update2)
+
+    # First update: START + CONTENT
+    assert len(events1) == 2
+    assert events1[0].type == "TEXT_MESSAGE_START"
+    assert events1[1].delta == "Hello "
+
+    # Second update: just CONTENT (same message)
+    assert len(events2) == 1
+    assert events2[0].type == "TEXT_MESSAGE_CONTENT"
+    assert events2[0].delta == "world"
+
+    # Both content events should have same message_id
+    assert events1[1].message_id == events2[0].message_id
+
+
+async def test_skip_text_content_for_structured_outputs():
+    """Test that text content is skipped when skip_text_content=True."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread", skip_text_content=True)
+
+    update = AgentRunResponseUpdate(contents=[TextContent(text='{"result": "data"}')])
+    events = await bridge.from_agent_run_update(update)
+
+    # No events should be emitted
+    assert len(events) == 0
+
+
+async def test_tool_call_with_name():
+    """Test FunctionCallContent with name emits ToolCallStartEvent."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
+
+    update = AgentRunResponseUpdate(contents=[FunctionCallContent(name="search_web", call_id="call_123")])
+    events = await bridge.from_agent_run_update(update)
+
+    assert len(events) == 1
+    assert events[0].type == "TOOL_CALL_START"
+    assert events[0].tool_call_name == "search_web"
+    assert events[0].tool_call_id == "call_123"
+
+
+async def test_tool_call_streaming_args():
+    """Test streaming tool call arguments."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
+
+    # First chunk: name only
+    update1 = AgentRunResponseUpdate(contents=[FunctionCallContent(name="search_web", call_id="call_123")])
+    events1 = await bridge.from_agent_run_update(update1)
+
+    # Second chunk: arguments chunk 1 (name can be empty string for continuation)
+    update2 = AgentRunResponseUpdate(
+        contents=[FunctionCallContent(name="", call_id="call_123", arguments='{"query": "')]
+    )
+    events2 = await bridge.from_agent_run_update(update2)
+
+    # Third chunk: arguments chunk 2
+    update3 = AgentRunResponseUpdate(contents=[FunctionCallContent(name="", call_id="call_123", arguments='AI"}')])
+    events3 = await bridge.from_agent_run_update(update3)
+
+    # First update: ToolCallStartEvent
+    assert len(events1) == 1
+    assert events1[0].type == "TOOL_CALL_START"
+
+    # Second update: ToolCallArgsEvent
+    assert len(events2) == 1
+    assert events2[0].type == "TOOL_CALL_ARGS"
+    assert events2[0].delta == '{"query": "'
+
+    # Third update: ToolCallArgsEvent
+    assert len(events3) == 1
+    assert events3[0].type == "TOOL_CALL_ARGS"
+    assert events3[0].delta == 'AI"}'
+
+    # All should have same tool_call_id
+    assert events1[0].tool_call_id == events2[0].tool_call_id == events3[0].tool_call_id
+
+
+async def test_tool_result_with_dict():
+    """Test FunctionResultContent with dict result."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
+
+    result_data = {"status": "success", "count": 42}
+    update = AgentRunResponseUpdate(contents=[FunctionResultContent(call_id="call_123", result=result_data)])
+    events = await bridge.from_agent_run_update(update)
+
+    # Should emit ToolCallEndEvent + ToolCallResultEvent
+    assert len(events) == 2
+    assert events[0].type == "TOOL_CALL_END"
+    assert events[0].tool_call_id == "call_123"
+
+    assert events[1].type == "TOOL_CALL_RESULT"
+    assert events[1].tool_call_id == "call_123"
+    assert events[1].role == "tool"
+    # Result should be JSON-serialized
+    assert json.loads(events[1].content) == result_data
+
+
+async def test_tool_result_with_string():
+    """Test FunctionResultContent with string result."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
+
+    update = AgentRunResponseUpdate(contents=[FunctionResultContent(call_id="call_123", result="Search complete")])
+    events = await bridge.from_agent_run_update(update)
+
+    assert len(events) == 2
+    assert events[0].type == "TOOL_CALL_END"
+    assert events[1].type == "TOOL_CALL_RESULT"
+    assert events[1].content == "Search complete"
+
+
+async def test_tool_result_with_none():
+    """Test FunctionResultContent with None result."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
+
+    update = AgentRunResponseUpdate(contents=[FunctionResultContent(call_id="call_123", result=None)])
+    events = await bridge.from_agent_run_update(update)
+
+    assert len(events) == 2
+    assert events[0].type == "TOOL_CALL_END"
+    assert events[1].type == "TOOL_CALL_RESULT"
+    assert events[1].content == ""
+
+
+async def test_multiple_tool_results_in_sequence():
+    """Test multiple tool results processed sequentially."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
+
+    update = AgentRunResponseUpdate(
+        contents=[
+            FunctionResultContent(call_id="call_1", result="Result 1"),
+            FunctionResultContent(call_id="call_2", result="Result 2"),
+        ]
+    )
+    events = await bridge.from_agent_run_update(update)
+
+    # Each result emits: ToolCallEndEvent + ToolCallResultEvent = 4 events total
+    assert len(events) == 4
+    assert events[0].tool_call_id == "call_1"
+    assert events[1].tool_call_id == "call_1"
+    assert events[2].tool_call_id == "call_2"
+    assert events[3].tool_call_id == "call_2"
+
+
+async def test_function_approval_request_basic():
+    """Test FunctionApprovalRequestContent conversion."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
+
+    func_call = FunctionCallContent(
+        call_id="call_123",
+        name="send_email",
+        arguments={"to": "user@example.com", "subject": "Test"},
+    )
+    approval = FunctionApprovalRequestContent(
+        id="approval_001",
+        function_call=func_call,
+    )
+
+    update = AgentRunResponseUpdate(contents=[approval])
+    events = await bridge.from_agent_run_update(update)
+
+    # Should emit: ToolCallEndEvent + CustomEvent
+    assert len(events) == 2
+
+    # First: ToolCallEndEvent to close the tool call
+    assert events[0].type == "TOOL_CALL_END"
+    assert events[0].tool_call_id == "call_123"
+
+    # Second: CustomEvent with approval details
+    assert events[1].type == "CUSTOM"
+    assert events[1].name == "function_approval_request"
+    assert events[1].value["id"] == "approval_001"
+    assert events[1].value["function_call"]["name"] == "send_email"
+
+
+async def test_empty_predict_state_config():
+    """Test behavior with no predictive state configuration."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    bridge = AgentFrameworkEventBridge(
+        run_id="test_run",
+        thread_id="test_thread",
+        predict_state_config={},  # Empty config
+    )
+
+    # Tool call with arguments
+    update = AgentRunResponseUpdate(
+        contents=[
+            FunctionCallContent(name="write_doc", call_id="call_1", arguments='{"content": "test"}'),
+            FunctionResultContent(call_id="call_1", result="Done"),
+        ]
+    )
+    events = await bridge.from_agent_run_update(update)
+
+    # Should NOT emit StateDeltaEvent or confirm_changes
+    event_types = [e.type for e in events]
+    assert "STATE_DELTA" not in event_types
+    assert "STATE_SNAPSHOT" not in event_types
+
+    # Should have: ToolCallStart, ToolCallArgs, ToolCallEnd, ToolCallResult, MessagesSnapshot
+    # MessagesSnapshotEvent is emitted after tool results to track the conversation
+    assert event_types == [
+        "TOOL_CALL_START",
+        "TOOL_CALL_ARGS",
+        "TOOL_CALL_END",
+        "TOOL_CALL_RESULT",
+        "MESSAGES_SNAPSHOT",
+    ]
+
+
+async def test_tool_not_in_predict_state_config():
+    """Test tool that doesn't match any predict_state_config entry."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    bridge = AgentFrameworkEventBridge(
+        run_id="test_run",
+        thread_id="test_thread",
+        predict_state_config={
+            "document": {"tool": "write_document", "tool_argument": "content"},
+        },
+    )
+
+    # Different tool name
+    update = AgentRunResponseUpdate(
+        contents=[
+            FunctionCallContent(name="search_web", call_id="call_1", arguments='{"query": "AI"}'),
+            FunctionResultContent(call_id="call_1", result="Results"),
+        ]
+    )
+    events = await bridge.from_agent_run_update(update)
+
+    # Should NOT emit StateDeltaEvent or confirm_changes
+    event_types = [e.type for e in events]
+    assert "STATE_DELTA" not in event_types
+    assert "STATE_SNAPSHOT" not in event_types
+
+
+async def test_state_management_tracking():
+    """Test current_state and pending_state_updates tracking."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    initial_state = {"document": ""}
+    bridge = AgentFrameworkEventBridge(
+        run_id="test_run",
+        thread_id="test_thread",
+        predict_state_config={
+            "document": {"tool": "write_doc", "tool_argument": "content"},
+        },
+        current_state=initial_state,
+    )
+
+    # Streaming tool call
+    update1 = AgentRunResponseUpdate(
+        contents=[
+            FunctionCallContent(name="write_doc", call_id="call_1"),
+            FunctionCallContent(name="", call_id="call_1", arguments='{"content": "Hello"}'),
+        ]
+    )
+    await bridge.from_agent_run_update(update1)
+
+    # Check pending_state_updates was populated
+    assert "document" in bridge.pending_state_updates
+    assert bridge.pending_state_updates["document"] == "Hello"
+
+    # Tool result should update current_state
+    update2 = AgentRunResponseUpdate(contents=[FunctionResultContent(call_id="call_1", result="Done")])
+    await bridge.from_agent_run_update(update2)
+
+    # current_state should be updated
+    assert bridge.current_state["document"] == "Hello"
+
+    # pending_state_updates should be cleared
+    assert len(bridge.pending_state_updates) == 0
+
+
+async def test_wildcard_tool_argument():
+    """Test tool_argument='*' uses all arguments as state value."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    bridge = AgentFrameworkEventBridge(
+        run_id="test_run",
+        thread_id="test_thread",
+        predict_state_config={
+            "recipe": {"tool": "create_recipe", "tool_argument": "*"},
+        },
+        current_state={},
+    )
+
+    # Complete tool call with dict arguments
+    update = AgentRunResponseUpdate(
+        contents=[
+            FunctionCallContent(
+                name="create_recipe",
+                call_id="call_1",
+                arguments={"title": "Pasta", "ingredients": ["pasta", "sauce"]},
+            ),
+            FunctionResultContent(call_id="call_1", result="Created"),
+        ]
+    )
+    events = await bridge.from_agent_run_update(update)
+
+    # Find StateDeltaEvent
+    delta_events = [e for e in events if e.type == "STATE_DELTA"]
+    assert len(delta_events) > 0
+
+    # Value should be the entire arguments dict
+    delta = delta_events[0].delta[0]
+    assert delta["path"] == "/recipe"
+    assert delta["value"] == {"title": "Pasta", "ingredients": ["pasta", "sauce"]}
+
+
+async def test_run_lifecycle_events():
+    """Test RunStartedEvent and RunFinishedEvent creation."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
+
+    started = bridge.create_run_started_event()
+    assert started.type == "RUN_STARTED"
+    assert started.run_id == "test_run"
+    assert started.thread_id == "test_thread"
+
+    finished = bridge.create_run_finished_event(result={"status": "complete"})
+    assert finished.type == "RUN_FINISHED"
+    assert finished.run_id == "test_run"
+    assert finished.thread_id == "test_thread"
+    assert finished.result == {"status": "complete"}
+
+
+async def test_message_lifecycle_events():
+    """Test TextMessageStartEvent and TextMessageEndEvent creation."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
+
+    start = bridge.create_message_start_event("msg_123", role="assistant")
+    assert start.type == "TEXT_MESSAGE_START"
+    assert start.message_id == "msg_123"
+    assert start.role == "assistant"
+
+    end = bridge.create_message_end_event("msg_123")
+    assert end.type == "TEXT_MESSAGE_END"
+    assert end.message_id == "msg_123"
+
+
+async def test_state_event_creation():
+    """Test StateSnapshotEvent and StateDeltaEvent creation helpers."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
+
+    # StateSnapshotEvent
+    snapshot = bridge.create_state_snapshot_event({"document": "content"})
+    assert snapshot.type == "STATE_SNAPSHOT"
+    assert snapshot.snapshot == {"document": "content"}
+
+    # StateDeltaEvent with JSON Patch
+    delta = bridge.create_state_delta_event([{"op": "replace", "path": "/document", "value": "new content"}])
+    assert delta.type == "STATE_DELTA"
+    assert len(delta.delta) == 1
+    assert delta.delta[0]["op"] == "replace"
+    assert delta.delta[0]["path"] == "/document"
+    assert delta.delta[0]["value"] == "new content"
+
+
+async def test_state_snapshot_after_tool_result():
+    """Test StateSnapshotEvent emission after tool result with pending updates."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    bridge = AgentFrameworkEventBridge(
+        run_id="test_run",
+        thread_id="test_thread",
+        predict_state_config={
+            "document": {"tool": "write_doc", "tool_argument": "content"},
+        },
+        current_state={"document": ""},
+    )
+
+    # Tool call with streaming args
+    update1 = AgentRunResponseUpdate(
+        contents=[
+            FunctionCallContent(name="write_doc", call_id="call_1"),
+            FunctionCallContent(name="", call_id="call_1", arguments='{"content": "Test"}'),
+        ]
+    )
+    await bridge.from_agent_run_update(update1)
+
+    # Tool result should trigger StateSnapshotEvent
+    update2 = AgentRunResponseUpdate(contents=[FunctionResultContent(call_id="call_1", result="Done")])
+    events = await bridge.from_agent_run_update(update2)
+
+    # Should have: ToolCallEnd, ToolCallResult, StateSnapshot, ToolCallStart (confirm_changes), ToolCallArgs, ToolCallEnd
+    snapshot_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
+    assert len(snapshot_events) == 1
+    assert snapshot_events[0].snapshot["document"] == "Test"
+
+
+async def test_message_id_persistence_across_chunks():
+    """Test that message_id persists across multiple text chunks."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
+
+    # First chunk
+    update1 = AgentRunResponseUpdate(contents=[TextContent(text="Hello ")])
+    events1 = await bridge.from_agent_run_update(update1)
+    message_id = events1[0].message_id
+
+    # Second chunk
+    update2 = AgentRunResponseUpdate(contents=[TextContent(text="world")])
+    events2 = await bridge.from_agent_run_update(update2)
+
+    # Should use same message_id
+    assert events2[0].message_id == message_id
+    assert bridge.current_message_id == message_id
+
+
+async def test_tool_call_id_tracking():
+    """Test tool_call_id tracking across streaming chunks."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
+
+    # First chunk with name
+    update1 = AgentRunResponseUpdate(contents=[FunctionCallContent(name="search", call_id="call_1")])
+    await bridge.from_agent_run_update(update1)
+
+    assert bridge.current_tool_call_id == "call_1"
+    assert bridge.current_tool_call_name == "search"
+
+    # Second chunk with args but no name
+    update2 = AgentRunResponseUpdate(contents=[FunctionCallContent(name="", call_id="call_1", arguments='{"q":"AI"}')])
+    events2 = await bridge.from_agent_run_update(update2)
+
+    # Should still track same tool call
+    assert bridge.current_tool_call_id == "call_1"
+    assert events2[0].tool_call_id == "call_1"
+
+
+async def test_tool_name_reset_after_result():
+    """Test current_tool_call_name is reset after tool result."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    bridge = AgentFrameworkEventBridge(
+        run_id="test_run",
+        thread_id="test_thread",
+        predict_state_config={
+            "document": {"tool": "write_doc", "tool_argument": "content"},
+        },
+    )
+
+    # Tool call
+    update1 = AgentRunResponseUpdate(
+        contents=[
+            FunctionCallContent(name="write_doc", call_id="call_1"),
+            FunctionCallContent(name="", call_id="call_1", arguments='{"content": "Test"}'),
+        ]
+    )
+    await bridge.from_agent_run_update(update1)
+
+    assert bridge.current_tool_call_name == "write_doc"
+
+    # Tool result with predictive state (should trigger confirm_changes and reset)
+    update2 = AgentRunResponseUpdate(contents=[FunctionResultContent(call_id="call_1", result="Done")])
+    await bridge.from_agent_run_update(update2)
+
+    # Tool name should be reset
+    assert bridge.current_tool_call_name is None
+
+
+async def test_function_approval_with_wildcard_argument():
+    """Test function approval with wildcard * argument."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    bridge = AgentFrameworkEventBridge(
+        run_id="test_run",
+        thread_id="test_thread",
+        predict_state_config={
+            "payload": {"tool": "submit", "tool_argument": "*"},
+        },
+    )
+
+    approval_content = FunctionApprovalRequestContent(
+        id="approval_1",
+        function_call=FunctionCallContent(
+            name="submit", call_id="call_1", arguments='{"key1": "value1", "key2": "value2"}'
+        ),
+    )
+
+    update = AgentRunResponseUpdate(contents=[approval_content])
+    events = await bridge.from_agent_run_update(update)
+
+    # Should emit StateSnapshotEvent with entire parsed args as value
+    snapshot_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
+    assert len(snapshot_events) == 1
+    assert snapshot_events[0].snapshot["payload"] == {"key1": "value1", "key2": "value2"}
+
+
+async def test_function_approval_missing_argument():
+    """Test function approval when specified argument is not in parsed args."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    bridge = AgentFrameworkEventBridge(
+        run_id="test_run",
+        thread_id="test_thread",
+        predict_state_config={
+            "data": {"tool": "process", "tool_argument": "missing_field"},
+        },
+    )
+
+    approval_content = FunctionApprovalRequestContent(
+        id="approval_1",
+        function_call=FunctionCallContent(name="process", call_id="call_1", arguments='{"other_field": "value"}'),
+    )
+
+    update = AgentRunResponseUpdate(contents=[approval_content])
+    events = await bridge.from_agent_run_update(update)
+
+    # Should not emit StateSnapshotEvent since argument not found
+    snapshot_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
+    assert len(snapshot_events) == 0
+
+
+async def test_empty_predict_state_config_no_deltas():
+    """Test with empty predict_state_config (no predictive updates)."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread", predict_state_config={})
+
+    # Tool call with arguments
+    update = AgentRunResponseUpdate(
+        contents=[
+            FunctionCallContent(name="search", call_id="call_1"),
+            FunctionCallContent(name="", call_id="call_1", arguments='{"query": "test"}'),
+        ]
+    )
+    events = await bridge.from_agent_run_update(update)
+
+    # Should not emit any StateDeltaEvents
+    delta_events = [e for e in events if e.type == "STATE_DELTA"]
+    assert len(delta_events) == 0
+
+
+async def test_tool_with_no_matching_config():
+    """Test tool call for tool not in predict_state_config."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    bridge = AgentFrameworkEventBridge(
+        run_id="test_run",
+        thread_id="test_thread",
+        predict_state_config={"document": {"tool": "write_doc", "tool_argument": "content"}},
+    )
+
+    # Tool call for different tool
+    update = AgentRunResponseUpdate(
+        contents=[
+            FunctionCallContent(name="search_web", call_id="call_1"),
+            FunctionCallContent(name="", call_id="call_1", arguments='{"query": "test"}'),
+        ]
+    )
+    events = await bridge.from_agent_run_update(update)
+
+    # Should not emit StateDeltaEvents
+    delta_events = [e for e in events if e.type == "STATE_DELTA"]
+    assert len(delta_events) == 0
+
+
+async def test_tool_call_without_name_or_id():
+    """Test handling FunctionCallContent with no name and no call_id."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    bridge = AgentFrameworkEventBridge(run_id="test_run", thread_id="test_thread")
+
+    # This should not crash but log an error
+    update = AgentRunResponseUpdate(contents=[FunctionCallContent(name="", call_id="", arguments='{"arg": "val"}')])
+    events = await bridge.from_agent_run_update(update)
+
+    # Should emit ToolCallArgsEvent with generated ID
+    assert len(events) >= 1
+
+
+async def test_state_delta_count_logging():
+    """Test that state delta count increments and logs at intervals."""
+    from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+    bridge = AgentFrameworkEventBridge(
+        run_id="test_run",
+        thread_id="test_thread",
+        predict_state_config={"doc": {"tool": "write", "tool_argument": "text"}},
+    )
+
+    # Emit multiple state deltas with different content each time
+    for i in range(15):
+        update = AgentRunResponseUpdate(
+            contents=[
+                FunctionCallContent(name="", call_id="call_1", arguments=f'{{"text": "Content variation {i}"}}'),
+            ]
+        )
+        # Set the tool name to match config
+        bridge.current_tool_call_name = "write"
+        await bridge.from_agent_run_update(update)
+
+    # State delta count should have incremented (one per unique state update)
+    assert bridge.state_delta_count >= 1

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

@@ -0,0 +1,238 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for AGUIHttpService."""
+
+import json
+from unittest.mock import AsyncMock, Mock
+
+import httpx
+import pytest
+
+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_and_tools(mock_http_client):
+    """Test posting run with state and tools."""
+
+    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"}}]
+
+    async for _ in service.post_run(thread_id="thread_123", run_id="run_456", messages=[], state=state, tools=tools):
+        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
+
+
+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/test_human_in_the_loop.py

@@ -0,0 +1,96 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for human in the loop (function approval requests)."""
+
+from agent_framework import FunctionApprovalRequestContent, FunctionCallContent
+from agent_framework._types import AgentRunResponseUpdate
+
+from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+
+async def test_function_approval_request_emission():
+    """Test that CustomEvent is emitted for FunctionApprovalRequestContent."""
+    bridge = AgentFrameworkEventBridge(
+        run_id="test_run",
+        thread_id="test_thread",
+    )
+
+    # Create approval request
+    func_call = FunctionCallContent(
+        call_id="call_123",
+        name="send_email",
+        arguments={"to": "user@example.com", "subject": "Test"},
+    )
+    approval_request = FunctionApprovalRequestContent(
+        id="approval_001",
+        function_call=func_call,
+    )
+
+    update = AgentRunResponseUpdate(contents=[approval_request])
+    events = await bridge.from_agent_run_update(update)
+
+    # Should emit ToolCallEndEvent + CustomEvent for approval request
+    assert len(events) == 2
+
+    # First event: ToolCallEndEvent to close the tool call
+    assert events[0].type == "TOOL_CALL_END"
+    assert events[0].tool_call_id == "call_123"
+
+    # Second event: CustomEvent with approval details
+    event = events[1]
+    assert event.type == "CUSTOM"
+    assert event.name == "function_approval_request"
+    assert event.value["id"] == "approval_001"
+    assert event.value["function_call"]["call_id"] == "call_123"
+    assert event.value["function_call"]["name"] == "send_email"
+    assert event.value["function_call"]["arguments"]["to"] == "user@example.com"
+    assert event.value["function_call"]["arguments"]["subject"] == "Test"
+
+
+async def test_multiple_approval_requests():
+    """Test handling multiple approval requests in one update."""
+    bridge = AgentFrameworkEventBridge(
+        run_id="test_run",
+        thread_id="test_thread",
+    )
+
+    func_call_1 = FunctionCallContent(
+        call_id="call_1",
+        name="create_event",
+        arguments={"title": "Meeting"},
+    )
+    approval_1 = FunctionApprovalRequestContent(
+        id="approval_1",
+        function_call=func_call_1,
+    )
+
+    func_call_2 = FunctionCallContent(
+        call_id="call_2",
+        name="book_room",
+        arguments={"room": "Conference A"},
+    )
+    approval_2 = FunctionApprovalRequestContent(
+        id="approval_2",
+        function_call=func_call_2,
+    )
+
+    update = AgentRunResponseUpdate(contents=[approval_1, approval_2])
+    events = await bridge.from_agent_run_update(update)
+
+    # Should emit ToolCallEndEvent + CustomEvent for each approval (4 events total)
+    assert len(events) == 4
+
+    # Events should alternate: End, Custom, End, Custom
+    assert events[0].type == "TOOL_CALL_END"
+    assert events[0].tool_call_id == "call_1"
+
+    assert events[1].type == "CUSTOM"
+    assert events[1].name == "function_approval_request"
+    assert events[1].value["id"] == "approval_1"
+
+    assert events[2].type == "TOOL_CALL_END"
+    assert events[2].tool_call_id == "call_2"
+
+    assert events[3].type == "CUSTOM"
+    assert events[3].name == "function_approval_request"
+    assert events[3].value["id"] == "approval_2"

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

@@ -0,0 +1,280 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for message adapters."""
+
+import pytest
+from agent_framework import ChatMessage, FunctionCallContent, Role, TextContent
+
+from agent_framework_ag_ui._message_adapters import (
+    agent_framework_messages_to_agui,
+    agui_messages_to_agent_framework,
+    extract_text_from_contents,
+)
+
+
+@pytest.fixture
+def sample_agui_message():
+    """Create a sample AG-UI message."""
+    return {"role": "user", "content": "Hello", "id": "msg-123"}
+
+
+@pytest.fixture
+def sample_agent_framework_message():
+    """Create a sample Agent Framework message."""
+    return ChatMessage(role=Role.USER, contents=[TextContent(text="Hello")], message_id="msg-123")
+
+
+def test_agui_to_agent_framework_basic(sample_agui_message):
+    """Test converting AG-UI message to Agent Framework."""
+    messages = agui_messages_to_agent_framework([sample_agui_message])
+
+    assert len(messages) == 1
+    assert messages[0].role == Role.USER
+    assert messages[0].message_id == "msg-123"
+
+
+def test_agent_framework_to_agui_basic(sample_agent_framework_message):
+    """Test converting Agent Framework message to AG-UI."""
+    messages = agent_framework_messages_to_agui([sample_agent_framework_message])
+
+    assert len(messages) == 1
+    assert messages[0]["role"] == "user"
+    assert messages[0]["content"] == "Hello"
+    assert messages[0]["id"] == "msg-123"
+
+
+def test_agui_tool_result_to_agent_framework():
+    """Test converting AG-UI tool result message to Agent Framework."""
+    tool_result_message = {
+        "role": "tool",
+        "content": '{"accepted": true, "steps": []}',
+        "toolCallId": "call_123",
+        "id": "msg_456",
+    }
+
+    messages = agui_messages_to_agent_framework([tool_result_message])
+
+    assert len(messages) == 1
+    message = messages[0]
+
+    assert message.role == Role.USER
+
+    assert len(message.contents) == 1
+    assert isinstance(message.contents[0], TextContent)
+    assert message.contents[0].text == '{"accepted": true, "steps": []}'
+
+    assert message.additional_properties is not None
+    assert message.additional_properties.get("is_tool_result") is True
+    assert message.additional_properties.get("tool_call_id") == "call_123"
+
+
+def test_agui_multiple_messages_to_agent_framework():
+    """Test converting multiple AG-UI messages."""
+    messages_input = [
+        {"role": "user", "content": "First message", "id": "msg-1"},
+        {"role": "assistant", "content": "Second message", "id": "msg-2"},
+        {"role": "user", "content": "Third message", "id": "msg-3"},
+    ]
+
+    messages = agui_messages_to_agent_framework(messages_input)
+
+    assert len(messages) == 3
+    assert messages[0].role == Role.USER
+    assert messages[1].role == Role.ASSISTANT
+    assert messages[2].role == Role.USER
+
+
+def test_agui_empty_messages():
+    """Test handling of empty messages list."""
+    messages = agui_messages_to_agent_framework([])
+    assert len(messages) == 0
+
+
+def test_agui_function_approvals():
+    """Test converting function approvals from AG-UI to Agent Framework."""
+    agui_msg = {
+        "role": "user",
+        "function_approvals": [
+            {
+                "call_id": "call-1",
+                "name": "search",
+                "arguments": {"query": "test"},
+                "approved": True,
+                "id": "approval-1",
+            },
+            {
+                "call_id": "call-2",
+                "name": "update",
+                "arguments": {"value": 42},
+                "approved": False,
+                "id": "approval-2",
+            },
+        ],
+        "id": "msg-123",
+    }
+
+    messages = agui_messages_to_agent_framework([agui_msg])
+
+    assert len(messages) == 1
+    msg = messages[0]
+    assert msg.role == Role.USER
+    assert len(msg.contents) == 2
+
+    from agent_framework import FunctionApprovalResponseContent
+
+    assert isinstance(msg.contents[0], FunctionApprovalResponseContent)
+    assert msg.contents[0].approved is True
+    assert msg.contents[0].id == "approval-1"
+    assert msg.contents[0].function_call.name == "search"
+    assert msg.contents[0].function_call.call_id == "call-1"
+
+    assert isinstance(msg.contents[1], FunctionApprovalResponseContent)
+    assert msg.contents[1].approved is False
+
+
+def test_agui_system_role():
+    """Test converting system role messages."""
+    messages = agui_messages_to_agent_framework([{"role": "system", "content": "System prompt"}])
+
+    assert len(messages) == 1
+    assert messages[0].role == Role.SYSTEM
+
+
+def test_agui_non_string_content():
+    """Test handling non-string content."""
+    messages = agui_messages_to_agent_framework([{"role": "user", "content": {"nested": "object"}}])
+
+    assert len(messages) == 1
+    assert len(messages[0].contents) == 1
+    assert isinstance(messages[0].contents[0], TextContent)
+    assert "nested" in messages[0].contents[0].text
+
+
+def test_agui_message_without_id():
+    """Test message without ID field."""
+    messages = agui_messages_to_agent_framework([{"role": "user", "content": "No ID"}])
+
+    assert len(messages) == 1
+    assert messages[0].message_id is None
+
+
+def test_agui_with_tool_calls_to_agent_framework():
+    """Assistant message with tool_calls is converted to FunctionCallContent."""
+    agui_msg = {
+        "role": "assistant",
+        "content": "Calling tool",
+        "tool_calls": [
+            {
+                "id": "call-123",
+                "type": "function",
+                "function": {"name": "get_weather", "arguments": {"location": "Seattle"}},
+            }
+        ],
+        "id": "msg-789",
+    }
+
+    messages = agui_messages_to_agent_framework([agui_msg])
+
+    assert len(messages) == 1
+    msg = messages[0]
+    assert msg.role == Role.ASSISTANT
+    assert msg.message_id == "msg-789"
+    # First content is text, second is the function call
+    assert isinstance(msg.contents[0], TextContent)
+    assert msg.contents[0].text == "Calling tool"
+    assert isinstance(msg.contents[1], FunctionCallContent)
+    assert msg.contents[1].call_id == "call-123"
+    assert msg.contents[1].name == "get_weather"
+    assert msg.contents[1].arguments == {"location": "Seattle"}
+
+
+def test_agent_framework_to_agui_with_tool_calls():
+    """Test converting Agent Framework message with tool calls to AG-UI."""
+    msg = ChatMessage(
+        role=Role.ASSISTANT,
+        contents=[
+            TextContent(text="Calling tool"),
+            FunctionCallContent(call_id="call-123", name="search", arguments={"query": "test"}),
+        ],
+        message_id="msg-456",
+    )
+
+    messages = agent_framework_messages_to_agui([msg])
+
+    assert len(messages) == 1
+    agui_msg = messages[0]
+    assert agui_msg["role"] == "assistant"
+    assert agui_msg["content"] == "Calling tool"
+    assert "tool_calls" in agui_msg
+    assert len(agui_msg["tool_calls"]) == 1
+    assert agui_msg["tool_calls"][0]["id"] == "call-123"
+    assert agui_msg["tool_calls"][0]["type"] == "function"
+    assert agui_msg["tool_calls"][0]["function"]["name"] == "search"
+    assert agui_msg["tool_calls"][0]["function"]["arguments"] == {"query": "test"}
+
+
+def test_agent_framework_to_agui_multiple_text_contents():
+    """Test concatenating multiple text contents."""
+    msg = ChatMessage(
+        role=Role.ASSISTANT,
+        contents=[TextContent(text="Part 1 "), TextContent(text="Part 2")],
+    )
+
+    messages = agent_framework_messages_to_agui([msg])
+
+    assert len(messages) == 1
+    assert messages[0]["content"] == "Part 1 Part 2"
+
+
+def test_agent_framework_to_agui_no_message_id():
+    """Test message without message_id - should auto-generate ID."""
+    msg = ChatMessage(role=Role.USER, contents=[TextContent(text="Hello")])
+
+    messages = agent_framework_messages_to_agui([msg])
+
+    assert len(messages) == 1
+    assert "id" in messages[0]  # ID should be auto-generated
+    assert messages[0]["id"]  # ID should not be empty
+    assert len(messages[0]["id"]) > 0  # ID should be a valid string
+
+
+def test_agent_framework_to_agui_system_role():
+    """Test system role conversion."""
+    msg = ChatMessage(role=Role.SYSTEM, contents=[TextContent(text="System")])
+
+    messages = agent_framework_messages_to_agui([msg])
+
+    assert len(messages) == 1
+    assert messages[0]["role"] == "system"
+
+
+def test_extract_text_from_contents():
+    """Test extracting text from contents list."""
+    contents = [TextContent(text="Hello "), TextContent(text="World")]
+
+    result = extract_text_from_contents(contents)
+
+    assert result == "Hello World"
+
+
+def test_extract_text_from_empty_contents():
+    """Test extracting text from empty contents."""
+    result = extract_text_from_contents([])
+
+    assert result == ""
+
+
+class CustomTextContent:
+    """Custom content with text attribute."""
+
+    def __init__(self, text: str):
+        self.text = text
+
+
+def test_extract_text_from_custom_contents():
+    """Test extracting text from custom content objects."""
+    contents = [CustomTextContent(text="Custom "), TextContent(text="Mixed")]
+
+    result = extract_text_from_contents(contents)
+
+    assert result == "Custom Mixed"

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

@@ -0,0 +1,82 @@
+"""Tests for AG-UI orchestrators."""
+
+from collections.abc import AsyncGenerator
+from types import SimpleNamespace
+from typing import Any
+
+from agent_framework import AgentRunResponseUpdate, TextContent, ai_function
+from agent_framework._tools import FunctionInvocationConfiguration
+
+from agent_framework_ag_ui._agent import AgentConfig
+from agent_framework_ag_ui._orchestrators import DefaultOrchestrator, ExecutionContext
+
+
+@ai_function
+def server_tool() -> str:
+    """Server-executable tool."""
+    return "server"
+
+
+class DummyAgent:
+    """Minimal agent stub to capture run_stream parameters."""
+
+    def __init__(self) -> None:
+        self.chat_options = SimpleNamespace(tools=[server_tool], response_format=None)
+        self.tools = [server_tool]
+        self.chat_client = SimpleNamespace(
+            function_invocation_configuration=FunctionInvocationConfiguration(),
+        )
+        self.seen_tools: list[Any] | None = None
+
+    async def run_stream(
+        self,
+        messages: list[Any],
+        *,
+        thread: Any,
+        tools: list[Any] | None = None,
+    ) -> AsyncGenerator[AgentRunResponseUpdate, None]:
+        self.seen_tools = tools
+        yield AgentRunResponseUpdate(contents=[TextContent(text="ok")], role="assistant")
+
+
+async def test_default_orchestrator_merges_client_tools() -> None:
+    """Client tool declarations are merged with server tools before running agent."""
+
+    agent = DummyAgent()
+    orchestrator = DefaultOrchestrator()
+
+    input_data = {
+        "messages": [
+            {
+                "role": "user",
+                "content": [{"type": "input_text", "text": "Hello"}],
+            }
+        ],
+        "tools": [
+            {
+                "name": "get_weather",
+                "description": "Client weather lookup.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {"location": {"type": "string"}},
+                    "required": ["location"],
+                },
+            }
+        ],
+    }
+
+    context = ExecutionContext(
+        input_data=input_data,
+        agent=agent,
+        config=AgentConfig(),
+    )
+
+    events = []
+    async for event in orchestrator.run(context):
+        events.append(event)
+
+    assert agent.seen_tools is not None
+    tool_names = [getattr(tool, "name", "?") for tool in agent.seen_tools]
+    assert "server_tool" in tool_names
+    assert "get_weather" in tool_names
+    assert agent.chat_client.function_invocation_configuration.additional_tools

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

@@ -0,0 +1,109 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for shared state management."""
+
+import pytest
+from ag_ui.core import StateSnapshotEvent
+from agent_framework import ChatAgent, TextContent
+from agent_framework._types import ChatResponseUpdate
+
+from agent_framework_ag_ui._agent import AgentFrameworkAgent
+from agent_framework_ag_ui._events import AgentFrameworkEventBridge
+
+
+@pytest.fixture
+def mock_agent():
+    """Create a mock agent for testing."""
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            yield ChatResponseUpdate(contents=[TextContent(text="Hello!")])
+
+    return ChatAgent(
+        name="test_agent",
+        instructions="Test agent",
+        chat_client=MockChatClient(),
+    )
+
+
+def test_state_snapshot_event():
+    """Test creating state snapshot events."""
+    bridge = AgentFrameworkEventBridge(run_id="test-run", thread_id="test-thread")
+
+    state = {
+        "recipe": {
+            "name": "Chocolate Chip Cookies",
+            "ingredients": ["flour", "sugar", "chocolate chips"],
+            "instructions": ["Mix ingredients", "Bake at 350°F"],
+            "servings": 24,
+        }
+    }
+
+    event = bridge.create_state_snapshot_event(state)
+
+    assert isinstance(event, StateSnapshotEvent)
+    assert event.snapshot == state
+    assert event.snapshot["recipe"]["name"] == "Chocolate Chip Cookies"
+    assert len(event.snapshot["recipe"]["ingredients"]) == 3
+
+
+def test_state_delta_event():
+    """Test creating state delta events using JSON Patch format."""
+    bridge = AgentFrameworkEventBridge(run_id="test-run", thread_id="test-thread")
+
+    # JSON Patch operations (RFC 6902)
+    delta = [
+        {"op": "add", "path": "/recipe/ingredients/-", "value": "vanilla extract"},
+        {"op": "replace", "path": "/recipe/servings", "value": 30},
+    ]
+
+    event = bridge.create_state_delta_event(delta)
+
+    assert event.delta == delta
+    assert len(event.delta) == 2
+    assert event.delta[0]["op"] == "add"
+    assert event.delta[1]["op"] == "replace"
+
+
+async def test_agent_with_initial_state(mock_agent):
+    """Test agent emits state snapshot when initial state provided."""
+    state_schema = {"recipe": {"type": "object", "properties": {"name": {"type": "string"}}}}
+
+    agent = AgentFrameworkAgent(
+        agent=mock_agent,
+        state_schema=state_schema,
+    )
+
+    initial_state = {"recipe": {"name": "Test Recipe"}}
+
+    input_data = {
+        "messages": [{"role": "user", "content": "Hello"}],
+        "state": initial_state,
+    }
+
+    events = []
+    async for event in agent.run_agent(input_data):
+        events.append(event)
+
+    # Should have RunStartedEvent, StateSnapshotEvent, RunFinishedEvent at minimum
+    snapshot_events = [e for e in events if isinstance(e, StateSnapshotEvent)]
+    assert len(snapshot_events) == 1
+    assert snapshot_events[0].snapshot == initial_state
+
+
+async def test_agent_without_state_schema(mock_agent):
+    """Test agent doesn't emit state events without state schema."""
+    agent = AgentFrameworkAgent(agent=mock_agent)
+
+    input_data = {
+        "messages": [{"role": "user", "content": "Hello"}],
+        "state": {"some": "state"},
+    }
+
+    events = []
+    async for event in agent.run_agent(input_data):
+        events.append(event)
+
+    # Should NOT have any StateSnapshotEvent
+    snapshot_events = [e for e in events if isinstance(e, StateSnapshotEvent)]
+    assert len(snapshot_events) == 0

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

@@ -0,0 +1,257 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for structured output handling in _agent.py."""
+
+import json
+from typing import Any
+
+from agent_framework import ChatAgent, ChatOptions, TextContent
+from agent_framework._types import ChatResponseUpdate
+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():
+    """Test structured output processing with recipe state."""
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            # Simulate structured output
+            yield ChatResponseUpdate(
+                contents=[TextContent(text='{"recipe": {"name": "Pasta"}, "message": "Here is your recipe"}')]
+            )
+
+    agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
+    agent.chat_options = ChatOptions(response_format=RecipeOutput)
+
+    wrapper = AgentFrameworkAgent(
+        agent=agent,
+        state_schema={"recipe": {"type": "object"}},
+    )
+
+    input_data = {"messages": [{"role": "user", "content": "Make pasta"}]}
+
+    events = []
+    async for event in wrapper.run_agent(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():
+    """Test structured output processing with steps state."""
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            steps_data = {
+                "steps": [
+                    {"id": "1", "description": "Step 1", "status": "pending"},
+                    {"id": "2", "description": "Step 2", "status": "pending"},
+                ]
+            }
+            yield ChatResponseUpdate(contents=[TextContent(text=json.dumps(steps_data))])
+
+    agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
+    agent.chat_options = ChatOptions(response_format=StepsOutput)
+
+    wrapper = AgentFrameworkAgent(
+        agent=agent,
+        state_schema={"steps": {"type": "array"}},
+    )
+
+    input_data = {"messages": [{"role": "user", "content": "Do steps"}]}
+
+    events = []
+    async for event in wrapper.run_agent(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():
+    """Test structured output when response fields don't match state_schema keys."""
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            # Response has "data" field but schema expects "result" field
+            yield ChatResponseUpdate(contents=[TextContent(text='{"data": {"key": "value"}}')])
+
+    agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
+    agent.chat_options = ChatOptions(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 = []
+    async for event in wrapper.run_agent(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():
+    """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
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            yield ChatResponseUpdate(contents=[TextContent(text='{"data": {"key": "value"}, "info": "processed"}')])
+
+    agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
+    agent.chat_options = ChatOptions(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 = []
+    async for event in wrapper.run_agent(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():
+    """Test that structured output path is skipped when no response_format."""
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            yield ChatResponseUpdate(contents=[TextContent(text="Regular text")])
+
+    agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
+    # No response_format set
+
+    wrapper = AgentFrameworkAgent(agent=agent)
+
+    input_data = {"messages": [{"role": "user", "content": "Hi"}]}
+
+    events = []
+    async for event in wrapper.run_agent(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():
+    """Test structured output that includes a message field."""
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            output_data = {"recipe": {"name": "Salad"}, "message": "Fresh salad recipe ready"}
+            yield ChatResponseUpdate(contents=[TextContent(text=json.dumps(output_data))])
+
+    agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
+    agent.chat_options = ChatOptions(response_format=RecipeOutput)
+
+    wrapper = AgentFrameworkAgent(
+        agent=agent,
+        state_schema={"recipe": {"type": "object"}},
+    )
+
+    input_data = {"messages": [{"role": "user", "content": "Make salad"}]}
+
+    events = []
+    async for event in wrapper.run_agent(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():
+    """Test that empty updates don't trigger structured output processing."""
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    class MockChatClient:
+        async def get_streaming_response(self, messages, chat_options, **kwargs):
+            # Return nothing
+            if False:
+                yield
+
+    agent = ChatAgent(name="test", instructions="Test", chat_client=MockChatClient())
+    agent.chat_options = ChatOptions(response_format=RecipeOutput)
+
+    wrapper = AgentFrameworkAgent(agent=agent)
+
+    input_data = {"messages": [{"role": "user", "content": "Test"}]}
+
+    events = []
+    async for event in wrapper.run_agent(input_data):
+        events.append(event)
+
+    # Should only have start and end events
+    assert len(events) == 2  # RunStarted, RunFinished

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

@@ -0,0 +1,145 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for type definitions in _types.py."""
+
+from agent_framework_ag_ui._types import AgentState, 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]

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

@@ -0,0 +1,305 @@
+# 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 = {"a": 1, "b": 2}
+    update = {"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 = {"x": 10, "y": 20}
+    update = {}
+
+    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 = {}
+    update = {"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 = {"recipe": {"name": "Cake", "ingredients": ["flour", "sugar"]}}
+    update = {"other": "value"}
+
+    result = merge_state(current, update)
+
+    result["recipe"]["ingredients"].append("eggs")
+
+    assert "eggs" not in current["recipe"]["ingredients"]
+    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 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_convert_tools_to_agui_format_with_ai_function():
+    """Test converting AIFunction to AG-UI format."""
+    from agent_framework import ai_function
+
+    from agent_framework_ag_ui._utils import convert_tools_to_agui_format
+
+    @ai_function
+    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 ai_function
+
+    from agent_framework_ag_ui._utils import convert_tools_to_agui_format
+
+    @ai_function
+    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 ai_function
+
+    from agent_framework_ag_ui._utils import convert_tools_to_agui_format
+
+    @ai_function
+    def tool1(x: int) -> int:
+        """Tool 1."""
+        return x
+
+    @ai_function
+    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"

python/packages/anthropic/pyproject.toml

@@ -4,7 +4,7 @@ description = "Anthropic integration for Microsoft Agent Framework."
 authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
 readme = "README.md"
 requires-python = ">=3.10"
-version = "1.0.0b251104"
+version = "1.0.0b251111"
 license-files = ["LICENSE"]
 urls.homepage = "https://aka.ms/agent-framework"
 urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
@@ -19,6 +19,7 @@ classifiers = [
   "Programming Language :: Python :: 3.11",
   "Programming Language :: Python :: 3.12",
   "Programming Language :: Python :: 3.13",
+  "Programming Language :: Python :: 3.14",
   "Typing :: Typed",
 ]
 dependencies = [

python/packages/azure-ai/agent_framework_azure_ai/__init__.py

@@ -2,7 +2,8 @@
 
 import importlib.metadata
 
-from ._chat_client import AzureAIAgentClient, AzureAISettings
+from ._chat_client import AzureAIAgentClient
+from ._shared import AzureAISettings
 
 try:
     __version__ = importlib.metadata.version(__name__)

python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py

@@ -40,9 +40,9 @@ from agent_framework import (
     use_chat_middleware,
     use_function_invocation,
 )
-from agent_framework._pydantic import AFBaseSettings
 from agent_framework.exceptions import ServiceInitializationError, ServiceResponseException
 from agent_framework.observability import use_observability
+from azure.ai.agents.aio import AgentsClient
 from azure.ai.agents.models import (
     Agent,
     AgentsNamedToolChoice,
@@ -85,11 +85,11 @@ from azure.ai.agents.models import (
     ToolDefinition,
     ToolOutput,
 )
-from azure.ai.projects.aio import AIProjectClient
 from azure.core.credentials_async import AsyncTokenCredential
-from azure.core.exceptions import HttpResponseError, ResourceNotFoundError
 from pydantic import ValidationError
 
+from ._shared import AzureAISettings
+
 if sys.version_info >= (3, 11):
     from typing import Self  # pragma: no cover
 else:
@@ -99,47 +99,6 @@ else:
 logger = get_logger("agent_framework.azure")
 
 
-class AzureAISettings(AFBaseSettings):
-    """Azure AI Project settings.
-
-    The settings are first loaded from environment variables with the prefix 'AZURE_AI_'.
-    If the environment variables are not found, the settings can be loaded from a .env file
-    with the encoding 'utf-8'. If the settings are not found in the .env file, the settings
-    are ignored; however, validation will fail alerting that the settings are missing.
-
-    Keyword Args:
-        project_endpoint: The Azure AI Project endpoint URL.
-            Can be set via environment variable AZURE_AI_PROJECT_ENDPOINT.
-        model_deployment_name: The name of the model deployment to use.
-            Can be set via environment variable AZURE_AI_MODEL_DEPLOYMENT_NAME.
-        env_file_path: If provided, the .env settings are read from this file path location.
-        env_file_encoding: The encoding of the .env file, defaults to 'utf-8'.
-
-    Examples:
-        .. code-block:: python
-
-            from agent_framework_azure_ai import AzureAISettings
-
-            # Using environment variables
-            # Set AZURE_AI_PROJECT_ENDPOINT=https://your-project.cognitiveservices.azure.com
-            # Set AZURE_AI_MODEL_DEPLOYMENT_NAME=gpt-4
-            settings = AzureAISettings()
-
-            # Or passing parameters directly
-            settings = AzureAISettings(
-                project_endpoint="https://your-project.cognitiveservices.azure.com", model_deployment_name="gpt-4"
-            )
-
-            # Or loading from a .env file
-            settings = AzureAISettings(env_file_path="path/to/.env")
-    """
-
-    env_prefix: ClassVar[str] = "AZURE_AI_"
-
-    project_endpoint: str | None = None
-    model_deployment_name: str | None = None
-
-
 TAzureAIAgentClient = TypeVar("TAzureAIAgentClient", bound="AzureAIAgentClient")
 
 
@@ -154,13 +113,14 @@ class AzureAIAgentClient(BaseChatClient):
     def __init__(
         self,
         *,
-        project_client: AIProjectClient | None = None,
+        agents_client: AgentsClient | None = None,
         agent_id: str | None = None,
         agent_name: str | None = None,
         thread_id: str | None = None,
         project_endpoint: str | None = None,
         model_deployment_name: str | None = None,
         async_credential: AsyncTokenCredential | None = None,
+        should_cleanup_agent: bool = True,
         env_file_path: str | None = None,
         env_file_encoding: str | None = None,
         **kwargs: Any,
@@ -168,19 +128,22 @@ class AzureAIAgentClient(BaseChatClient):
         """Initialize an Azure AI Agent client.
 
         Keyword Args:
-            project_client: An existing AIProjectClient to use. If not provided, one will be created.
-            agent_id: The ID of an existing agent to use. If not provided and project_client is provided,
-                a new agent will be created (and deleted after the request). If neither project_client
+            agents_client: An existing AgentsClient to use. If not provided, one will be created.
+            agent_id: The ID of an existing agent to use. If not provided and agents_client is provided,
+                a new agent will be created (and deleted after the request). If neither agents_client
                 nor agent_id is provided, both will be created and managed automatically.
             agent_name: The name to use when creating new agents.
             thread_id: Default thread ID to use for conversations. Can be overridden by
                 conversation_id property when making a request.
             project_endpoint: The Azure AI Project endpoint URL.
                 Can also be set via environment variable AZURE_AI_PROJECT_ENDPOINT.
-                Ignored when a project_client is passed.
+                Ignored when a agents_client is passed.
             model_deployment_name: The model deployment name to use for agent creation.
                 Can also be set via environment variable AZURE_AI_MODEL_DEPLOYMENT_NAME.
             async_credential: Azure async credential to use for authentication.
+            should_cleanup_agent: Whether to cleanup (delete) agents created by this client when
+                the client is closed or context is exited. Defaults to True. Only affects agents
+                created by this client instance; existing agents passed via agent_id are never deleted.
             env_file_path: Path to environment file for loading settings.
             env_file_encoding: Encoding of the environment file.
             kwargs: Additional keyword arguments passed to the parent class.
@@ -217,9 +180,9 @@ class AzureAIAgentClient(BaseChatClient):
         except ValidationError as ex:
             raise ServiceInitializationError("Failed to create Azure AI settings.", ex) from ex
 
-        # If no project_client is provided, create one
+        # If no agents_client is provided, create one
         should_close_client = False
-        if project_client is None:
+        if agents_client is None:
             if not azure_ai_settings.project_endpoint:
                 raise ServiceInitializationError(
                     "Azure AI project endpoint is required. Set via 'project_endpoint' parameter "
@@ -234,8 +197,8 @@ class AzureAIAgentClient(BaseChatClient):
 
             # Use provided credential
             if not async_credential:
-                raise ServiceInitializationError("Azure credential is required when project_client is not provided.")
-            project_client = AIProjectClient(
+                raise ServiceInitializationError("Azure credential is required when agents_client is not provided.")
+            agents_client = AgentsClient(
                 endpoint=azure_ai_settings.project_endpoint,
                 credential=async_credential,
                 user_agent=AGENT_FRAMEWORK_USER_AGENT,
@@ -246,37 +209,17 @@ class AzureAIAgentClient(BaseChatClient):
         super().__init__(**kwargs)
 
         # Initialize instance variables
-        self.project_client = project_client
+        self.agents_client = agents_client
         self.credential = async_credential
         self.agent_id = agent_id
         self.agent_name = agent_name
         self.model_id = azure_ai_settings.model_deployment_name
         self.thread_id = thread_id
-        self._should_delete_agent = False  # Track whether we should delete the agent
+        self.should_cleanup_agent = should_cleanup_agent  # Track whether we should delete the agent
+        self._agent_created = False  # Track whether agent was created inside this class
         self._should_close_client = should_close_client  # Track whether we should close client connection
         self._agent_definition: Agent | None = None  # Cached definition for existing agent
 
-    async def setup_azure_ai_observability(self, enable_sensitive_data: bool | None = None) -> None:
-        """Use this method to setup tracing in your Azure AI Project.
-
-        This will take the connection string from the project project_client.
-        It will override any connection string that is set in the environment variables.
-        It will disable any OTLP endpoint that might have been set.
-        """
-        try:
-            conn_string = await self.project_client.telemetry.get_application_insights_connection_string()
-        except ResourceNotFoundError:
-            logger.warning(
-                "No Application Insights connection string found for the Azure AI Project, "
-                "please call setup_observability() manually."
-            )
-            return
-        from agent_framework.observability import setup_observability
-
-        setup_observability(
-            applicationinsights_connection_string=conn_string, enable_sensitive_data=enable_sensitive_data
-        )
-
     async def __aenter__(self) -> "Self":
         """Async context manager entry."""
         return self
@@ -286,7 +229,7 @@ class AzureAIAgentClient(BaseChatClient):
         await self.close()
 
     async def close(self) -> None:
-        """Close the project_client and clean up any agents we created."""
+        """Close the agents_client and clean up any agents we created."""
         await self._cleanup_agent_if_needed()
         await self._close_client_if_needed()
 
@@ -298,7 +241,7 @@ class AzureAIAgentClient(BaseChatClient):
             settings: A dictionary of settings for the service.
         """
         return cls(
-            project_client=settings.get("project_client"),
+            agents_client=settings.get("agents_client"),
             agent_id=settings.get("agent_id"),
             thread_id=settings.get("thread_id"),
             project_endpoint=settings.get("project_endpoint"),
@@ -306,6 +249,7 @@ class AzureAIAgentClient(BaseChatClient):
             agent_name=settings.get("agent_name"),
             credential=settings.get("credential"),
             env_file_path=settings.get("env_file_path"),
+            should_cleanup_agent=settings.get("should_cleanup_agent", True),
         )
 
     async def _inner_get_response(
@@ -374,14 +318,17 @@ class AzureAIAgentClient(BaseChatClient):
                 args["instructions"] = run_options["instructions"]
             if "response_format" in run_options:
                 args["response_format"] = run_options["response_format"]
+
             if "temperature" in run_options:
                 args["temperature"] = run_options["temperature"]
             if "top_p" in run_options:
                 args["top_p"] = run_options["top_p"]
-            created_agent = await self.project_client.agents.create_agent(**args)
+
+            created_agent = await self.agents_client.create_agent(**args)
+
             self.agent_id = str(created_agent.id)
             self._agent_definition = created_agent
-            self._should_delete_agent = True
+            self._agent_created = True
 
         return self.agent_id
 
@@ -422,7 +369,7 @@ class AzureAIAgentClient(BaseChatClient):
                 args["tool_outputs"] = tool_outputs
             if tool_approvals:
                 args["tool_approvals"] = tool_approvals
-            await self.project_client.agents.runs.submit_tool_outputs_stream(**args)  # type: ignore[reportUnknownMemberType]
+            await self.agents_client.runs.submit_tool_outputs_stream(**args)  # type: ignore[reportUnknownMemberType]
             # Pass the handler to the stream to continue processing
             stream = handler  # type: ignore
             final_thread_id = thread_run.thread_id
@@ -432,7 +379,7 @@ class AzureAIAgentClient(BaseChatClient):
 
             # Now create a new run and stream the results.
             run_options.pop("conversation_id", None)
-            stream = await self.project_client.agents.runs.stream(  # type: ignore[reportUnknownMemberType]
+            stream = await self.agents_client.runs.stream(  # type: ignore[reportUnknownMemberType]
                 final_thread_id, agent_id=agent_id, **run_options
             )
 
@@ -443,9 +390,7 @@ class AzureAIAgentClient(BaseChatClient):
         if thread_id is None:
             return None
 
-        async for run in self.project_client.agents.runs.list(
-            thread_id=thread_id, limit=1, order=ListSortOrder.DESCENDING
-        ):  # type: ignore[reportUnknownMemberType]
+        async for run in self.agents_client.runs.list(thread_id=thread_id, limit=1, order=ListSortOrder.DESCENDING):  # type: ignore[reportUnknownMemberType]
             if run.status not in [
                 RunStatus.COMPLETED,
                 RunStatus.CANCELLED,
@@ -462,12 +407,12 @@ class AzureAIAgentClient(BaseChatClient):
         if thread_id is not None:
             if thread_run is not None:
                 # There was an active run; we need to cancel it before starting a new run.
-                await self.project_client.agents.runs.cancel(thread_id, thread_run.id)
+                await self.agents_client.runs.cancel(thread_id, thread_run.id)
 
             return thread_id
 
         # No thread ID was provided, so create a new thread.
-        thread = await self.project_client.agents.threads.create(
+        thread = await self.agents_client.threads.create(
             tool_resources=run_options.get("tool_resources"), metadata=run_options.get("metadata")
         )
         thread_id = thread.id
@@ -476,7 +421,7 @@ class AzureAIAgentClient(BaseChatClient):
         # once fixed, in the function above, readd:
         # `messages=run_options.pop("additional_messages")`
         for msg in run_options.pop("additional_messages", []):
-            await self.project_client.agents.messages.create(
+            await self.agents_client.messages.create(
                 thread_id=thread_id, role=msg.role, content=msg.content, metadata=msg.metadata
             )
         # and remove until here.
@@ -709,21 +654,21 @@ class AzureAIAgentClient(BaseChatClient):
         return []
 
     async def _close_client_if_needed(self) -> None:
-        """Close project_client session if we created it."""
+        """Close agents_client session if we created it."""
         if self._should_close_client:
-            await self.project_client.close()
+            await self.agents_client.close()
 
     async def _cleanup_agent_if_needed(self) -> None:
         """Clean up the agent if we created it."""
-        if self._should_delete_agent and self.agent_id is not None:
-            await self.project_client.agents.delete_agent(self.agent_id)
+        if self._agent_created and self.should_cleanup_agent and self.agent_id is not None:
+            await self.agents_client.delete_agent(self.agent_id)
             self.agent_id = None
-            self._should_delete_agent = False
+            self._agent_created = False
 
     async def _load_agent_definition_if_needed(self) -> Agent | None:
         """Load and cache agent details if not already loaded."""
         if self._agent_definition is None and self.agent_id is not None:
-            self._agent_definition = await self.project_client.agents.get_agent(self.agent_id)
+            self._agent_definition = await self.agents_client.get_agent(self.agent_id)
         return self._agent_definition
 
     def _prepare_tool_choice(self, chat_options: ChatOptions) -> None:
@@ -913,59 +858,34 @@ class AzureAIAgentClient(BaseChatClient):
                         config_args["market"] = market
                     if set_lang := additional_props.get("set_lang"):
                         config_args["set_lang"] = set_lang
-                    # Bing Grounding (support both connection_id and connection_name)
+                    # Bing Grounding
                     connection_id = additional_props.get("connection_id") or os.getenv("BING_CONNECTION_ID")
-                    connection_name = additional_props.get("connection_name") or os.getenv("BING_CONNECTION_NAME")
                     # Custom Bing Search
-                    custom_connection_name = additional_props.get("custom_connection_name") or os.getenv(
-                        "BING_CUSTOM_CONNECTION_NAME"
+                    custom_connection_id = additional_props.get("custom_connection_id") or os.getenv(
+                        "BING_CUSTOM_CONNECTION_ID"
                     )
-                    custom_configuration_name = additional_props.get("custom_instance_name") or os.getenv(
+                    custom_instance_name = additional_props.get("custom_instance_name") or os.getenv(
                         "BING_CUSTOM_INSTANCE_NAME"
                     )
                     bing_search: BingGroundingTool | BingCustomSearchTool | None = None
-                    if (
-                        (connection_id or connection_name)
-                        and not custom_connection_name
-                        and not custom_configuration_name
-                    ):
+                    if (connection_id) and not custom_connection_id and not custom_instance_name:
                         if connection_id:
                             conn_id = connection_id
-                        elif connection_name:
-                            try:
-                                bing_connection = await self.project_client.connections.get(name=connection_name)
-                            except HttpResponseError as err:
-                                raise ServiceInitializationError(
-                                    f"Bing connection '{connection_name}' not found in the Azure AI Project.",
-                                    err,
-                                ) from err
-                            else:
-                                conn_id = bing_connection.id
                         else:
-                            raise ServiceInitializationError("Neither connection_id nor connection_name provided.")
+                            raise ServiceInitializationError("Parameter connection_id is not provided.")
                         bing_search = BingGroundingTool(connection_id=conn_id, **config_args)
-                    if custom_connection_name and custom_configuration_name:
-                        try:
-                            bing_custom_connection = await self.project_client.connections.get(
-                                name=custom_connection_name
-                            )
-                        except HttpResponseError as err:
-                            raise ServiceInitializationError(
-                                f"Bing custom connection '{custom_connection_name}' not found in the Azure AI Project.",
-                                err,
-                            ) from err
-                        else:
-                            bing_search = BingCustomSearchTool(
-                                connection_id=bing_custom_connection.id,
-                                instance_name=custom_configuration_name,
-                                **config_args,
-                            )
+                    if custom_connection_id and custom_instance_name:
+                        bing_search = BingCustomSearchTool(
+                            connection_id=custom_connection_id,
+                            instance_name=custom_instance_name,
+                            **config_args,
+                        )
                     if not bing_search:
                         raise ServiceInitializationError(
-                            "Bing search tool requires either 'connection_id' or 'connection_name' for Bing Grounding "
-                            "or both 'custom_connection_name' and 'custom_instance_name' for Custom Bing Search. "
+                            "Bing search tool requires either 'connection_id' for Bing Grounding "
+                            "or both 'custom_connection_id' and 'custom_instance_name' for Custom Bing Search. "
                             "These can be provided via additional_properties or environment variables: "
-                            "'BING_CONNECTION_ID', 'BING_CONNECTION_NAME', 'BING_CUSTOM_CONNECTION_NAME', "
+                            "'BING_CONNECTION_ID', 'BING_CUSTOM_CONNECTION_ID', "
                             "'BING_CUSTOM_INSTANCE_NAME'"
                         )
                     tool_definitions.extend(bing_search.definitions)
@@ -1056,4 +976,4 @@ class AzureAIAgentClient(BaseChatClient):
         Returns:
             The service URL for the chat client, or None if not set.
         """
-        return self.project_client._config.endpoint
+        return self.agents_client._config.endpoint  # type: ignore

python/packages/azure-ai/agent_framework_azure_ai/_shared.py

@@ -0,0 +1,46 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from typing import ClassVar
+
+from agent_framework._pydantic import AFBaseSettings
+
+
+class AzureAISettings(AFBaseSettings):
+    """Azure AI Project settings.
+
+    The settings are first loaded from environment variables with the prefix 'AZURE_AI_'.
+    If the environment variables are not found, the settings can be loaded from a .env file
+    with the encoding 'utf-8'. If the settings are not found in the .env file, the settings
+    are ignored; however, validation will fail alerting that the settings are missing.
+
+    Keyword Args:
+        project_endpoint: The Azure AI Project endpoint URL.
+            Can be set via environment variable AZURE_AI_PROJECT_ENDPOINT.
+        model_deployment_name: The name of the model deployment to use.
+            Can be set via environment variable AZURE_AI_MODEL_DEPLOYMENT_NAME.
+        env_file_path: If provided, the .env settings are read from this file path location.
+        env_file_encoding: The encoding of the .env file, defaults to 'utf-8'.
+
+    Examples:
+        .. code-block:: python
+
+            from agent_framework.azure import AzureAISettings
+
+            # Using environment variables
+            # Set AZURE_AI_PROJECT_ENDPOINT=https://your-project.cognitiveservices.azure.com
+            # Set AZURE_AI_MODEL_DEPLOYMENT_NAME=gpt-4
+            settings = AzureAISettings()
+
+            # Or passing parameters directly
+            settings = AzureAISettings(
+                project_endpoint="https://your-project.cognitiveservices.azure.com", model_deployment_name="gpt-4"
+            )
+
+            # Or loading from a .env file
+            settings = AzureAISettings(env_file_path="path/to/.env")
+    """
+
+    env_prefix: ClassVar[str] = "AZURE_AI_"
+
+    project_endpoint: str | None = None
+    model_deployment_name: str | None = None

python/packages/azure-ai/pyproject.toml

@@ -4,7 +4,7 @@ description = "Azure AI Foundry integration for Microsoft Agent Framework."
 authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
 readme = "README.md"
 requires-python = ">=3.10"
-version = "1.0.0b251104"
+version = "1.0.0b251111"
 license-files = ["LICENSE"]
 urls.homepage = "https://aka.ms/agent-framework"
 urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
@@ -19,12 +19,14 @@ classifiers = [
   "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",
     "azure-ai-projects >= 1.0.0b11",
     "azure-ai-agents == 1.2.0b5",
+    "aiohttp",
 ]
 
 [tool.uv]

python/packages/azure-ai/tests/conftest.py

@@ -44,31 +44,30 @@ def azure_ai_unit_test_env(monkeypatch, exclude_list, override_env_param_dict):
 
 
 @fixture
-def mock_ai_project_client() -> MagicMock:
-    """Fixture that provides a mock AIProjectClient."""
+def mock_agents_client() -> MagicMock:
+    """Fixture that provides a mock AgentsClient."""
     mock_client = MagicMock()
 
     # Mock agents property
-    mock_client.agents = MagicMock()
-    mock_client.agents.create_agent = AsyncMock()
-    mock_client.agents.delete_agent = AsyncMock()
+    mock_client.create_agent = AsyncMock()
+    mock_client.delete_agent = AsyncMock()
 
     # Mock agent creation response
     mock_agent = MagicMock()
     mock_agent.id = "test-agent-id"
-    mock_client.agents.create_agent.return_value = mock_agent
+    mock_client.create_agent.return_value = mock_agent
 
     # Mock threads property
-    mock_client.agents.threads = MagicMock()
-    mock_client.agents.threads.create = AsyncMock()
-    mock_client.agents.messages.create = AsyncMock()
+    mock_client.threads = MagicMock()
+    mock_client.threads.create = AsyncMock()
+    mock_client.messages.create = AsyncMock()
 
     # Mock runs property
-    mock_client.agents.runs = MagicMock()
-    mock_client.agents.runs.list = AsyncMock()
-    mock_client.agents.runs.cancel = AsyncMock()
-    mock_client.agents.runs.stream = AsyncMock()
-    mock_client.agents.runs.submit_tool_outputs_stream = AsyncMock()
+    mock_client.runs = MagicMock()
+    mock_client.runs.list = AsyncMock()
+    mock_client.runs.cancel = AsyncMock()
+    mock_client.runs.stream = AsyncMock()
+    mock_client.runs.submit_tool_outputs_stream = AsyncMock()
 
     return mock_client
 

python/packages/chatkit/.gitignore

@@ -0,0 +1,3 @@
+chatkit-python
+openai-chatkit-advanced-samples
+chatkit-js

python/packages/chatkit/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/chatkit/README.md

@@ -0,0 +1,87 @@
+# Agent Framework and ChatKit Integration
+
+This package provides an integration layer between Microsoft Agent Framework
+and [OpenAI ChatKit (Python)](https://github.com/openai/chatkit-python/).
+Specifically, it mirrors the [Agent SDK integration](https://github.com/openai/chatkit-python/blob/main/docs/server.md#agents-sdk-integration), and provides the following helpers:
+
+- `stream_agent_response`: A helper to convert a streamed `AgentRunResponseUpdate`
+  from a Microsoft Agent Framework agent that implements `AgentProtocol` to ChatKit events.
+- `ThreadItemConverter`: A extendable helper class to convert ChatKit thread items to
+  `ChatMessage` objects that can be consumed by an Agent Framework agent.
+- `simple_to_agent_input`: A helper function that uses the default implementation
+  of `ThreadItemConverter` to convert a ChatKit thread to a list of `ChatMessage`,
+  useful for getting started quickly.
+
+## Installation
+
+```bash
+pip install agent-framework-chatkit --pre
+```
+
+This will install `agent-framework-core` and `openai-chatkit` as dependencies.
+
+## Example Usage
+
+Here's a minimal example showing how to integrate Agent Framework with ChatKit:
+
+```python
+from collections.abc import AsyncIterator
+from typing import Any
+
+from azure.identity import AzureCliCredential
+from fastapi import FastAPI, Request
+from fastapi.responses import Response, StreamingResponse
+
+from agent_framework import ChatAgent
+from agent_framework.azure import AzureOpenAIChatClient
+from agent_framework.chatkit import simple_to_agent_input, stream_agent_response
+
+from chatkit.server import ChatKitServer
+from chatkit.types import ThreadMetadata, UserMessageItem, ThreadStreamEvent
+
+# You'll need to implement a Store - see the sample for a SQLiteStore implementation
+from your_store import YourStore  # type: ignore[import-not-found]  # Replace with your Store implementation
+
+# Define your agent with tools
+agent = ChatAgent(
+    chat_client=AzureOpenAIChatClient(credential=AzureCliCredential()),
+    instructions="You are a helpful assistant.",
+    tools=[],  # Add your tools here
+)
+
+# Create a ChatKit server that uses your agent
+class MyChatKitServer(ChatKitServer[dict[str, Any]]):
+    async def respond(
+        self,
+        thread: ThreadMetadata,
+        input_user_message: UserMessageItem | None,
+        context: dict[str, Any],
+    ) -> AsyncIterator[ThreadStreamEvent]:
+        if input_user_message is None:
+            return
+
+        # Convert ChatKit message to Agent Framework format
+        agent_messages = await simple_to_agent_input(input_user_message)
+
+        # Run the agent and stream responses
+        response_stream = agent.run_stream(agent_messages)
+
+        # Convert agent responses back to ChatKit events
+        async for event in stream_agent_response(response_stream, thread.id):
+            yield event
+
+# Set up FastAPI endpoint
+app = FastAPI()
+chatkit_server = MyChatKitServer(YourStore())  # type: ignore[misc]
+
+@app.post("/chatkit")
+async def chatkit_endpoint(request: Request):
+    result = await chatkit_server.process(await request.body(), {"request": request})
+
+    if hasattr(result, '__aiter__'):  # Streaming
+        return StreamingResponse(result, media_type="text/event-stream")  # type: ignore[arg-type]
+    else:  # Non-streaming
+        return Response(content=result.json, media_type="application/json")  # type: ignore[union-attr]
+```
+
+For a complete end-to-end example with a full frontend, see the [weather agent sample](../../samples/demos/chatkit-integration/README.md).

python/packages/chatkit/agent_framework_chatkit/__init__.py

@@ -0,0 +1,25 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Agent Framework and ChatKit Integration.
+
+This package provides an integration layer between Microsoft Agent Framework
+and OpenAI ChatKit (Python). It mirrors the Agent SDK integration and provides
+helpers to convert between Agent Framework and ChatKit types.
+"""
+
+import importlib.metadata
+
+from ._converter import ThreadItemConverter, simple_to_agent_input
+from ._streaming import stream_agent_response
+
+try:
+    __version__ = importlib.metadata.version(__name__)
+except importlib.metadata.PackageNotFoundError:
+    __version__ = "0.0.0"  # Fallback for development mode
+
+__all__ = [
+    "ThreadItemConverter",
+    "__version__",
+    "simple_to_agent_input",
+    "stream_agent_response",
+]

python/packages/chatkit/agent_framework_chatkit/_converter.py

@@ -0,0 +1,603 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Converter utilities for converting ChatKit thread items to Agent Framework messages."""
+
+import logging
+import sys
+from collections.abc import Awaitable, Callable, Sequence
+
+if sys.version_info >= (3, 11):
+    from typing import assert_never
+else:
+    from typing_extensions import assert_never
+
+from agent_framework import (
+    ChatMessage,
+    DataContent,
+    FunctionCallContent,
+    FunctionResultContent,
+    Role,
+    TextContent,
+    UriContent,
+)
+from chatkit.types import (
+    AssistantMessageItem,
+    Attachment,
+    ClientToolCallItem,
+    EndOfTurnItem,
+    HiddenContextItem,
+    ImageAttachment,
+    TaskItem,
+    ThreadItem,
+    UserMessageItem,
+    UserMessageTagContent,
+    UserMessageTextContent,
+    WidgetItem,
+    WorkflowItem,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class ThreadItemConverter:
+    """Helper class to convert ChatKit thread items to Agent Framework ChatMessage objects.
+
+    This class provides a base implementation for converting ChatKit thread items
+    to Agent Framework messages. It can be extended to handle attachments,
+    @-mentions, hidden context items, and custom thread item formats.
+
+    Args:
+        attachment_data_fetcher: Optional async function to fetch attachment binary data.
+            If provided, it should take an attachment ID and return the binary data as bytes.
+            If not provided, attachments will be converted to UriContent using available URLs.
+    """
+
+    def __init__(
+        self,
+        attachment_data_fetcher: Callable[[str], Awaitable[bytes]] | None = None,
+    ) -> None:
+        """Initialize the converter.
+
+        Args:
+            attachment_data_fetcher: Optional async function to fetch attachment data by ID.
+        """
+        self.attachment_data_fetcher = attachment_data_fetcher
+
+    async def user_message_to_input(
+        self, item: UserMessageItem, is_last_message: bool = True
+    ) -> ChatMessage | list[ChatMessage] | None:
+        """Convert a ChatKit UserMessageItem to Agent Framework ChatMessage(s).
+
+        This method is called internally by `to_agent_input()`. Override this method
+        to customize how user messages are converted.
+
+        Args:
+            item: The ChatKit user message item to convert.
+            is_last_message: Whether this is the last message in the thread (used for quoted_text handling).
+
+        Returns:
+            A ChatMessage, list of messages, or None to skip.
+
+        Note:
+            Instead of calling this method directly, use `to_agent_input()` which handles
+            all ThreadItem types and provides proper message ordering.
+        """
+        # Extract text content from the user message
+        text_content = ""
+        if item.content:
+            for content_part in item.content:
+                if isinstance(content_part, UserMessageTextContent):
+                    text_content += content_part.text
+
+        # Convert attachments to DataContent or UriContent
+        data_contents: list[DataContent | UriContent] = []
+        if item.attachments:
+            for attachment in item.attachments:
+                content = await self.attachment_to_message_content(attachment)
+                if content is not None:
+                    data_contents.append(content)
+
+        # Create the message with text and attachments
+        if not text_content.strip() and not data_contents:
+            return None
+
+        # If only text and no attachments, use text parameter for simplicity
+        if text_content.strip() and not data_contents:
+            user_message = ChatMessage(role=Role.USER, text=text_content.strip())
+        else:
+            # Build contents list with both text and attachments
+            contents: list[TextContent | DataContent | UriContent] = []
+            if text_content.strip():
+                contents.append(TextContent(text=text_content.strip()))
+            contents.extend(data_contents)
+            user_message = ChatMessage(role=Role.USER, contents=contents)
+
+        # Handle quoted text if this is the last message
+        messages = [user_message]
+        if item.quoted_text and is_last_message:
+            quoted_context = ChatMessage(
+                role=Role.USER,
+                text=f"The user is referring to this in particular:\n{item.quoted_text}",
+            )
+            # Prepend quoted context before the main message
+            messages.insert(0, quoted_context)
+
+        return messages
+
+    async def attachment_to_message_content(self, attachment: Attachment) -> DataContent | UriContent | None:
+        """Convert a ChatKit attachment to Agent Framework content.
+
+        This method is called internally by `user_message_to_input()` to handle attachments.
+        Override this method to customize attachment handling for your storage backend.
+
+        The default implementation provides two strategies:
+        1. If an attachment_data_fetcher was provided, it fetches the binary data
+           and creates a DataContent object
+        2. Otherwise, for ImageAttachment with preview_url, it creates a UriContent object
+
+        For FileAttachment without a data fetcher, returns None (attachment is skipped).
+
+        Args:
+            attachment: The ChatKit attachment to convert (FileAttachment or ImageAttachment).
+
+        Returns:
+            DataContent if binary data is available, UriContent if only URL is available,
+            or None if the attachment cannot be converted.
+
+        Note:
+            Instead of calling this method directly, use `to_agent_input()` which handles
+            all ThreadItem types including attachments within user messages.
+
+        Examples:
+            .. code-block:: python
+
+                # With data fetcher
+                async def fetch_data(attachment_id: str) -> bytes:
+                    return await my_storage.get_file(attachment_id)
+
+
+                converter = ThreadItemConverter(attachment_data_fetcher=fetch_data)
+                messages = await converter.to_agent_input(thread_items)
+
+                # Without data fetcher (uses URLs for images)
+                converter = ThreadItemConverter()
+                messages = await converter.to_agent_input(thread_items)
+        """
+        # If we have a data fetcher, use it to get binary data
+        if self.attachment_data_fetcher is not None:
+            try:
+                data = await self.attachment_data_fetcher(attachment.id)
+                return DataContent(data=data, media_type=attachment.mime_type)
+            except Exception as e:
+                # If fetch fails, fall through to URL-based approach
+                logger.debug(f"Failed to fetch attachment data for {attachment.id}: {e}")
+
+        # For ImageAttachment, try to use preview_url
+        if isinstance(attachment, ImageAttachment) and attachment.preview_url:
+            return UriContent(uri=str(attachment.preview_url), media_type=attachment.mime_type)
+
+        # For FileAttachment without data fetcher, skip the attachment
+        # Subclasses can override this method to provide custom handling
+        return None
+
+    def hidden_context_to_input(self, item: HiddenContextItem) -> ChatMessage | list[ChatMessage] | None:
+        """Convert a ChatKit HiddenContextItem to Agent Framework ChatMessage(s).
+
+        This method is called internally by `to_agent_input()`. Override this method
+        to customize how hidden context is converted.
+
+        The default implementation wraps the hidden context in XML tags and returns
+        a system message. This allows the model to distinguish hidden context from
+        regular conversation.
+
+        Args:
+            item: The ChatKit hidden context item to convert.
+
+        Returns:
+            A ChatMessage with system role, a list of messages, or None to skip.
+
+        Note:
+            Instead of calling this method directly, use `to_agent_input()` which handles
+            all ThreadItem types and provides proper message ordering.
+
+        Examples:
+            .. code-block:: python
+
+                # Default behavior
+                converter = ThreadItemConverter()
+                hidden_item = HiddenContextItem(
+                    id="ctx_1",
+                    thread_id="thread_1",
+                    created_at=datetime.now(),
+                    content="User's email: user@example.com",
+                )
+                message = converter.hidden_context_to_input(hidden_item)
+                # Returns: ChatMessage(role=SYSTEM, text="<HIDDEN_CONTEXT>User's email: ...</HIDDEN_CONTEXT>")
+        """
+        return ChatMessage(role=Role.SYSTEM, text=f"<HIDDEN_CONTEXT>{item.content}</HIDDEN_CONTEXT>")
+
+    def tag_to_message_content(self, tag: UserMessageTagContent) -> TextContent:
+        """Convert a ChatKit tag (@-mention) to Agent Framework content.
+
+        This method is called internally by `user_message_to_input()` to handle tags.
+        Override this method to customize tag conversion for your application.
+
+        The default implementation extracts the tag's display name and wraps it in
+        XML tags to provide context to the model about the @-mention.
+
+        Args:
+            tag: The ChatKit tag content to convert.
+
+        Returns:
+            TextContent with the tag information.
+
+        Note:
+            Instead of calling this method directly, use `to_agent_input()` which handles
+            all ThreadItem types including tags within user messages.
+
+        Examples:
+            .. code-block:: python
+
+                # Default behavior
+                converter = ThreadItemConverter()
+                tag = UserMessageTagContent(
+                    type="input_tag", id="tag_1", text="john", data={"name": "John Doe"}, interactive=False
+                )
+                content = converter.tag_to_message_content(tag)
+                # Returns: TextContent(text="<TAG>Name:John Doe</TAG>")
+        """
+        name = getattr(tag.data, "name", tag.text if hasattr(tag, "text") else "unknown")
+        return TextContent(text=f"<TAG>Name:{name}</TAG>")
+
+    def task_to_input(self, item: TaskItem) -> ChatMessage | list[ChatMessage] | None:
+        """Convert a ChatKit TaskItem to Agent Framework ChatMessage(s).
+
+        This method is called internally by `to_agent_input()`. Override this method
+        to customize how tasks are converted.
+
+        The default implementation converts custom tasks with title/content into
+        a user message explaining what task was displayed to the user.
+
+        Args:
+            item: The ChatKit task item to convert.
+
+        Returns:
+            A ChatMessage, a list of messages, or None to skip the task.
+
+        Note:
+            Instead of calling this method directly, use `to_agent_input()` which handles
+            all ThreadItem types and provides proper message ordering.
+
+        Examples:
+            .. code-block:: python
+
+                # Task with both title and content
+                from chatkit.types import Task
+
+                task_item = TaskItem(
+                    id="task_1",
+                    thread_id="thread_1",
+                    created_at=datetime.now(),
+                    task=Task(type="custom", title="Data Analysis", content="Analyzed sales data"),
+                )
+                message = converter.task_to_input(task_item)
+                # Returns message explaining the task was performed
+        """
+        if item.task.type != "custom" or (not item.task.title and not item.task.content):
+            return None
+
+        title = item.task.title or ""
+        content = item.task.content or ""
+        task_text = f"{title}: {content}" if title and content else title or content
+        text = (
+            f"A message was displayed to the user that the following task was performed:\n<Task>\n{task_text}\n</Task>"
+        )
+
+        return ChatMessage(role=Role.USER, text=text)
+
+    def workflow_to_input(self, item: WorkflowItem) -> ChatMessage | list[ChatMessage] | None:
+        """Convert a ChatKit WorkflowItem to Agent Framework ChatMessage(s).
+
+        This method is called internally by `to_agent_input()`. Override this method
+        to customize how workflows are converted.
+
+        The default implementation converts each custom task in the workflow into
+        a separate user message explaining what tasks were performed.
+
+        Args:
+            item: The ChatKit workflow item to convert.
+
+        Returns:
+            A list of ChatMessages (one per task), a single message, or None to skip.
+
+        Note:
+            Instead of calling this method directly, use `to_agent_input()` which handles
+            all ThreadItem types and provides proper message ordering.
+
+        Examples:
+            .. code-block:: python
+
+                # Workflow with multiple tasks
+                from chatkit.types import Workflow, Task
+
+                workflow_item = WorkflowItem(
+                    id="wf_1",
+                    thread_id="thread_1",
+                    created_at=datetime.now(),
+                    workflow=Workflow(
+                        type="custom",
+                        tasks=[
+                            Task(type="custom", title="Step 1", content="Gathered data"),
+                            Task(type="custom", title="Step 2", content="Analyzed results"),
+                        ],
+                    ),
+                )
+                messages = converter.workflow_to_input(workflow_item)
+                # Returns list of messages for each task
+        """
+        messages: list[ChatMessage] = []
+        for task in item.workflow.tasks:
+            if task.type != "custom" or (not task.title and not task.content):
+                continue
+
+            title = task.title or ""
+            content = task.content or ""
+            task_text = f"{title}: {content}" if title and content else title or content
+            text = (
+                "A message was displayed to the user that the following task was performed:\n"
+                f"<Task>\n{task_text}\n</Task>"
+            )
+
+            messages.append(ChatMessage(role=Role.USER, text=text))
+
+        return messages if messages else None
+
+    def widget_to_input(self, item: WidgetItem) -> ChatMessage | list[ChatMessage] | None:
+        """Convert a ChatKit WidgetItem to Agent Framework ChatMessage(s).
+
+        This method is called internally by `to_agent_input()`. Override this method
+        to customize how widgets are converted.
+
+        The default implementation converts the widget to a JSON representation
+        and includes it in a user message, allowing the model to understand what
+        UI element was displayed to the user.
+
+        Args:
+            item: The ChatKit widget item to convert.
+
+        Returns:
+            A ChatMessage describing the widget, or None to skip.
+
+        Note:
+            Instead of calling this method directly, use `to_agent_input()` which handles
+            all ThreadItem types and provides proper message ordering.
+
+        Examples:
+            .. code-block:: python
+
+                # Widget item
+                from chatkit.widgets import Card, Text
+
+                widget_item = WidgetItem(
+                    id="widget_1",
+                    thread_id="thread_1",
+                    created_at=datetime.now(),
+                    widget=Card(children=[Text(value="Hello")]),
+                )
+                message = converter.widget_to_input(widget_item)
+                # Returns message with JSON representation of the widget
+        """
+        try:
+            widget_json = item.widget.model_dump_json(exclude_unset=True, exclude_none=True)
+            text = f"The following graphical UI widget (id: {item.id}) was displayed to the user:{widget_json}"
+            return ChatMessage(role=Role.USER, text=text)
+        except Exception:
+            # If JSON serialization fails, skip the widget
+            return None
+
+    async def assistant_message_to_input(self, item: AssistantMessageItem) -> ChatMessage | list[ChatMessage] | None:
+        """Convert a ChatKit AssistantMessageItem to Agent Framework ChatMessage(s).
+
+        The default implementation extracts text from all content parts and creates
+        an assistant message.
+
+        Args:
+            item: The ChatKit assistant message item to convert.
+
+        Returns:
+            A ChatMessage with assistant role, or None to skip.
+
+        Note:
+            Instead of calling this method directly, use `to_agent_input()` which handles
+            all ThreadItem types and provides proper message ordering.
+        """
+        # Extract text from all content parts
+        text_parts = [content.text for content in item.content]
+        if not text_parts:
+            return None
+
+        return ChatMessage(role=Role.ASSISTANT, text="".join(text_parts))
+
+    async def client_tool_call_to_input(self, item: ClientToolCallItem) -> ChatMessage | list[ChatMessage] | None:
+        """Convert a ChatKit ClientToolCallItem to Agent Framework ChatMessage(s).
+
+        The default implementation converts completed tool calls into function call
+        and result content.
+
+        Args:
+            item: The ChatKit client tool call item to convert.
+
+        Returns:
+            A list containing function call and result messages, or None for pending calls.
+
+        Note:
+            Instead of calling this method directly, use `to_agent_input()` which handles
+            all ThreadItem types and provides proper message ordering.
+        """
+        if item.status == "pending":
+            # Skip pending tool calls - they cannot be sent to the model
+            return None
+
+        import json
+
+        # Create function call message
+        function_call_msg = ChatMessage(
+            role=Role.ASSISTANT,
+            contents=[
+                FunctionCallContent(
+                    call_id=item.call_id,
+                    name=item.name,
+                    arguments=json.dumps(item.arguments),
+                )
+            ],
+        )
+
+        # Create function result message
+        function_result_msg = ChatMessage(
+            role=Role.TOOL,
+            contents=[
+                FunctionResultContent(
+                    call_id=item.call_id,
+                    result=json.dumps(item.output) if item.output is not None else "",
+                )
+            ],
+        )
+
+        return [function_call_msg, function_result_msg]
+
+    async def end_of_turn_to_input(self, item: EndOfTurnItem) -> ChatMessage | list[ChatMessage] | None:
+        """Convert a ChatKit EndOfTurnItem to Agent Framework ChatMessage(s).
+
+        The default implementation skips end-of-turn markers as they are only UI hints.
+
+        Args:
+            item: The ChatKit end-of-turn item to convert.
+
+        Returns:
+            None (end-of-turn items are not converted).
+
+        Note:
+            Instead of calling this method directly, use `to_agent_input()` which handles
+            all ThreadItem types and provides proper message ordering.
+        """
+        # End-of-turn is only used for UI hints - skip it
+        return None
+
+    async def _thread_item_to_input_item(
+        self,
+        item: ThreadItem,
+        is_last_message: bool = True,
+    ) -> list[ChatMessage]:
+        """Internal method to convert a single ThreadItem to ChatMessage(s).
+
+        Args:
+            item: The thread item to convert.
+            is_last_message: Whether this is the last item in the thread.
+
+        Returns:
+            A list of ChatMessage objects (may be empty).
+        """
+        match item:
+            case UserMessageItem():
+                out = await self.user_message_to_input(item, is_last_message) or []
+                return out if isinstance(out, list) else [out]
+            case AssistantMessageItem():
+                out = await self.assistant_message_to_input(item) or []
+                return out if isinstance(out, list) else [out]
+            case ClientToolCallItem():
+                out = await self.client_tool_call_to_input(item) or []
+                return out if isinstance(out, list) else [out]
+            case EndOfTurnItem():
+                out = await self.end_of_turn_to_input(item) or []
+                return out if isinstance(out, list) else [out]
+            case WidgetItem():
+                out = self.widget_to_input(item) or []
+                return out if isinstance(out, list) else [out]
+            case WorkflowItem():
+                out = self.workflow_to_input(item) or []
+                return out if isinstance(out, list) else [out]
+            case TaskItem():
+                out = self.task_to_input(item) or []
+                return out if isinstance(out, list) else [out]
+            case HiddenContextItem():
+                out = self.hidden_context_to_input(item) or []
+                return out if isinstance(out, list) else [out]
+            case _:
+                assert_never(item)
+
+    async def to_agent_input(
+        self,
+        thread_items: Sequence[ThreadItem] | ThreadItem,
+    ) -> list[ChatMessage]:
+        """Convert ChatKit thread items to Agent Framework ChatMessages.
+
+        This is the main entry point for converting ChatKit thread items. It handles
+        all ThreadItem types (UserMessageItem, AssistantMessageItem, TaskItem, etc.)
+        and calls the appropriate conversion method for each.
+
+        Args:
+            thread_items: A single ThreadItem or a sequence of ThreadItems to convert.
+
+        Returns:
+            A list of ChatMessage objects that can be sent to an Agent Framework agent.
+
+        Examples:
+            .. code-block:: python
+
+                from agent_framework_chatkit import ThreadItemConverter
+
+                converter = ThreadItemConverter()
+
+                # Convert a single thread item
+                messages = await converter.to_agent_input(user_message_item)
+
+                # Convert multiple thread items
+                messages = await converter.to_agent_input([user_message_item, assistant_message_item, task_item])
+
+                # Use with agent
+                from agent_framework import ChatAgent
+
+                agent = ChatAgent(...)
+                response = await agent.run_stream(messages)
+        """
+        thread_items = list(thread_items) if isinstance(thread_items, Sequence) else [thread_items]
+
+        output: list[ChatMessage] = []
+        for item in thread_items:
+            output.extend(
+                await self._thread_item_to_input_item(
+                    item,
+                    is_last_message=item is thread_items[-1],
+                )
+            )
+        return output
+
+
+# Default converter instance
+_DEFAULT_CONVERTER = ThreadItemConverter()
+
+
+async def simple_to_agent_input(thread_items: Sequence[ThreadItem] | ThreadItem) -> list[ChatMessage]:
+    """Helper function that uses the default ThreadItemConverter.
+
+    This function provides a quick way to get started with ChatKit integration
+    without needing to create a custom ThreadItemConverter instance.
+
+    Args:
+        thread_items: A single ThreadItem or a sequence of ThreadItems to convert.
+
+    Returns:
+        A list of ChatMessage objects that can be sent to an Agent Framework agent.
+
+    Examples:
+        .. code-block:: python
+
+            from agent_framework_chatkit import simple_to_agent_input
+
+            # Convert a single item
+            messages = await simple_to_agent_input(user_message_item)
+
+            # Convert multiple items
+            messages = await simple_to_agent_input([user_message_item, assistant_message_item, task_item])
+    """
+    return await _DEFAULT_CONVERTER.to_agent_input(thread_items)

python/packages/chatkit/agent_framework_chatkit/_streaming.py

@@ -0,0 +1,104 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Streaming utilities for converting Agent Framework responses to ChatKit events."""
+
+import uuid
+from collections.abc import AsyncIterable, AsyncIterator, Callable
+from datetime import datetime
+
+from agent_framework import AgentRunResponseUpdate, TextContent
+from chatkit.types import (
+    AssistantMessageContent,
+    AssistantMessageContentPartTextDelta,
+    AssistantMessageItem,
+    ThreadItemAddedEvent,
+    ThreadItemDoneEvent,
+    ThreadItemUpdated,
+    ThreadStreamEvent,
+)
+
+
+async def stream_agent_response(
+    response_stream: AsyncIterable[AgentRunResponseUpdate],
+    thread_id: str,
+    generate_id: Callable[[str], str] | None = None,
+) -> AsyncIterator[ThreadStreamEvent]:
+    """Convert a streamed AgentRunResponseUpdate from Agent Framework to ChatKit events.
+
+    This helper function takes a stream of AgentRunResponseUpdate objects from
+    a Microsoft Agent Framework agent and converts them to ChatKit ThreadStreamEvent
+    objects that can be consumed by the ChatKit UI.
+
+    The function supports real-time token-by-token streaming by emitting
+    ThreadItemUpdated events with AssistantMessageContentPartTextDelta for each
+    text chunk as it arrives from the agent.
+
+    Args:
+        response_stream: An async iterable of AgentRunResponseUpdate objects
+                        from an Agent Framework agent.
+        thread_id: The ChatKit thread ID for the conversation.
+        generate_id: Optional function to generate IDs for ChatKit items.
+                    If not provided, simple incremental IDs will be used.
+
+    Yields:
+        ThreadStreamEvent: ChatKit events representing the agent's response,
+                          including incremental text deltas for streaming display.
+    """
+    # Use provided ID generator or create default one
+    if generate_id is None:
+
+        def _default_id_generator(item_type: str) -> str:
+            return f"{item_type}_{uuid.uuid4().hex[:8]}"
+
+        message_id = _default_id_generator("msg")
+    else:
+        message_id = generate_id("msg")
+
+    # Track if we've started the message
+    message_started = False
+    accumulated_text = ""
+    content_index = 0
+
+    async for update in response_stream:
+        # Start the assistant message if not already started
+        if not message_started:
+            assistant_message = AssistantMessageItem(
+                id=message_id,
+                thread_id=thread_id,
+                type="assistant_message",
+                content=[],
+                created_at=datetime.now(),
+            )
+
+            yield ThreadItemAddedEvent(type="thread.item.added", item=assistant_message)
+            message_started = True
+
+        # Process the update content
+        if update.contents:
+            for content in update.contents:
+                # Handle text content - only TextContent has a text attribute
+                if isinstance(content, TextContent) and content.text is not None:
+                    # Yield incremental text delta for streaming display
+                    yield ThreadItemUpdated(
+                        type="thread.item.updated",
+                        item_id=message_id,
+                        update=AssistantMessageContentPartTextDelta(
+                            content_index=content_index,
+                            delta=content.text,
+                        ),
+                    )
+                    accumulated_text += content.text
+
+    # Finalize the message
+    if message_started:
+        final_message = AssistantMessageItem(
+            id=message_id,
+            thread_id=thread_id,
+            type="assistant_message",
+            content=[AssistantMessageContent(type="output_text", text=accumulated_text, annotations=[])]
+            if accumulated_text
+            else [],
+            created_at=datetime.now(),
+        )
+
+        yield ThreadItemDoneEvent(type="thread.item.done", item=final_message)

python/packages/chatkit/pyproject.toml

@@ -0,0 +1,89 @@
+[project]
+name = "agent-framework-chatkit"
+description = "OpenAI ChatKit integration for Microsoft Agent Framework."
+authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
+readme = "README.md"
+requires-python = ">=3.10"
+version = "1.0.0b251111"
+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",
+  "Typing :: Typed",
+]
+dependencies = [
+    "agent-framework-core",
+    "openai-chatkit>=1.1.0,<2.0.0",
+]
+
+[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 = []
+timeout = 120
+
+[tool.ruff]
+extend = "../../pyproject.toml"
+
+[tool.ruff.lint]
+ignore = ["RUF029"]
+
+[tool.coverage.run]
+omit = [
+    "**/__init__.py"
+]
+
+[tool.pyright]
+extend = "../../pyproject.toml"
+exclude = ['tests', 'chatkit-python', 'openai-chatkit-advanced-samples']
+
+[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_chatkit"]
+exclude_dirs = ["tests"]
+
+[tool.poe]
+executor.type = "uv"
+include = "../../shared_tasks.toml"
+[tool.poe.tasks]
+mypy = "mypy --config-file $POE_ROOT/pyproject.toml agent_framework_chatkit"
+test = "pytest --cov=agent_framework_chatkit --cov-report=term-missing:skip-covered tests"
+
+[build-system]
+requires = ["flit-core >= 3.11,<4.0"]
+build-backend = "flit_core.buildapi"

python/packages/chatkit/tests/__init__.py

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

python/packages/chatkit/tests/test_converter.py

@@ -0,0 +1,426 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for ChatKit to Agent Framework converter utilities."""
+
+from unittest.mock import Mock
+
+import pytest
+from agent_framework import ChatMessage, Role, TextContent
+from chatkit.types import UserMessageTextContent
+
+from agent_framework_chatkit import ThreadItemConverter, simple_to_agent_input
+
+
+class TestThreadItemConverter:
+    """Tests for ThreadItemConverter class."""
+
+    @pytest.fixture
+    def converter(self):
+        """Create a ThreadItemConverter instance for testing."""
+        return ThreadItemConverter()
+
+    async def test_to_agent_input_none(self, converter):
+        """Test converting empty list returns empty list."""
+        result = await converter.to_agent_input([])
+        assert result == []
+
+    async def test_to_agent_input_with_text(self, converter):
+        """Test converting user message with text content."""
+        from datetime import datetime
+
+        from chatkit.types import UserMessageItem
+
+        input_item = UserMessageItem(
+            id="msg_1",
+            thread_id="thread_1",
+            created_at=datetime.now(),
+            type="user_message",
+            content=[UserMessageTextContent(text="Hello, how can you help me?")],
+            attachments=[],
+            inference_options={},
+        )
+
+        result = await converter.to_agent_input(input_item)
+
+        assert len(result) == 1
+        assert isinstance(result[0], ChatMessage)
+        assert result[0].role == Role.USER
+        assert result[0].text == "Hello, how can you help me?"
+
+    async def test_to_agent_input_empty_text(self, converter):
+        """Test converting user message with empty or whitespace-only text."""
+        from datetime import datetime
+
+        from chatkit.types import UserMessageItem
+
+        input_item = UserMessageItem(
+            id="msg_1",
+            thread_id="thread_1",
+            created_at=datetime.now(),
+            type="user_message",
+            content=[UserMessageTextContent(text="   ")],
+            attachments=[],
+            inference_options={},
+        )
+
+        result = await converter.to_agent_input(input_item)
+        assert result == []
+
+    async def test_to_agent_input_no_content(self, converter):
+        """Test converting user message with no content."""
+        from datetime import datetime
+
+        from chatkit.types import UserMessageItem
+
+        input_item = UserMessageItem(
+            id="msg_1",
+            thread_id="thread_1",
+            created_at=datetime.now(),
+            type="user_message",
+            content=[],
+            attachments=[],
+            inference_options={},
+        )
+
+        result = await converter.to_agent_input(input_item)
+        assert result == []
+
+    async def test_to_agent_input_multiple_content_parts(self, converter):
+        """Test converting user message with multiple text content parts."""
+        from datetime import datetime
+
+        from chatkit.types import UserMessageItem
+
+        input_item = UserMessageItem(
+            id="msg_1",
+            thread_id="thread_1",
+            created_at=datetime.now(),
+            type="user_message",
+            content=[
+                UserMessageTextContent(text="Hello "),
+                UserMessageTextContent(text="world!"),
+            ],
+            attachments=[],
+            inference_options={},
+        )
+
+        result = await converter.to_agent_input(input_item)
+
+        assert len(result) == 1
+        assert result[0].text == "Hello world!"
+
+    def test_hidden_context_to_input(self, converter):
+        """Test converting hidden context item to ChatMessage."""
+        hidden_item = Mock()
+        hidden_item.content = "This is hidden context information"
+
+        result = converter.hidden_context_to_input(hidden_item)
+
+        assert isinstance(result, ChatMessage)
+        assert result.role == Role.SYSTEM
+        assert result.text == "<HIDDEN_CONTEXT>This is hidden context information</HIDDEN_CONTEXT>"
+
+    def test_tag_to_message_content(self, converter):
+        """Test converting tag to message content."""
+        from chatkit.types import UserMessageTagContent
+
+        tag = UserMessageTagContent(
+            type="input_tag",
+            id="tag_1",
+            text="john",
+            data={"name": "John Doe"},
+            interactive=False,
+        )
+
+        result = converter.tag_to_message_content(tag)
+        assert isinstance(result, TextContent)
+        # Since data is a dict, getattr won't work, so it will fall back to text
+        assert result.text == "<TAG>Name:john</TAG>"
+
+    def test_tag_to_message_content_no_name(self, converter):
+        """Test converting tag with no name to message content."""
+        from chatkit.types import UserMessageTagContent
+
+        tag = UserMessageTagContent(
+            type="input_tag",
+            id="tag_2",
+            text="jane",
+            data={},
+            interactive=False,
+        )
+
+        result = converter.tag_to_message_content(tag)
+        assert isinstance(result, TextContent)
+        assert result.text == "<TAG>Name:jane</TAG>"
+
+    async def test_attachment_to_message_content_file_without_fetcher(self, converter):
+        """Test that FileAttachment without data fetcher returns None."""
+        from chatkit.types import FileAttachment
+
+        attachment = FileAttachment(
+            id="file_123",
+            name="document.pdf",
+            mime_type="application/pdf",
+            type="file",
+        )
+
+        result = await converter.attachment_to_message_content(attachment)
+        assert result is None
+
+    async def test_attachment_to_message_content_image_with_preview_url(self, converter):
+        """Test that ImageAttachment with preview_url creates UriContent."""
+        from agent_framework import UriContent
+        from chatkit.types import ImageAttachment
+
+        attachment = ImageAttachment(
+            id="img_123",
+            name="photo.jpg",
+            mime_type="image/jpeg",
+            type="image",
+            preview_url="https://example.com/photo.jpg",
+        )
+
+        result = await converter.attachment_to_message_content(attachment)
+        assert isinstance(result, UriContent)
+        assert result.uri == "https://example.com/photo.jpg"
+        assert result.media_type == "image/jpeg"
+
+    async def test_attachment_to_message_content_with_data_fetcher(self):
+        """Test attachment conversion with data fetcher."""
+        from agent_framework import DataContent
+        from chatkit.types import FileAttachment
+
+        # Mock data fetcher
+        async def fetch_data(attachment_id: str) -> bytes:
+            return b"file content data"
+
+        converter = ThreadItemConverter(attachment_data_fetcher=fetch_data)
+
+        attachment = FileAttachment(
+            id="file_123",
+            name="document.pdf",
+            mime_type="application/pdf",
+            type="file",
+        )
+
+        result = await converter.attachment_to_message_content(attachment)
+        assert isinstance(result, DataContent)
+        assert result.media_type == "application/pdf"
+
+    async def test_to_agent_input_with_image_attachment(self):
+        """Test converting user message with text and image attachment."""
+        from datetime import datetime
+
+        from agent_framework import UriContent
+        from chatkit.types import ImageAttachment, UserMessageItem
+
+        attachment = ImageAttachment(
+            id="img_123",
+            name="photo.jpg",
+            mime_type="image/jpeg",
+            type="image",
+            preview_url="https://example.com/photo.jpg",
+        )
+
+        input_item = UserMessageItem(
+            id="msg_1",
+            thread_id="thread_1",
+            created_at=datetime.now(),
+            type="user_message",
+            content=[UserMessageTextContent(text="Check out this photo!")],
+            attachments=[attachment],
+            inference_options={},
+        )
+
+        converter = ThreadItemConverter()
+        result = await converter.to_agent_input(input_item)
+
+        assert len(result) == 1
+        message = result[0]
+        assert message.role == Role.USER
+        assert len(message.contents) == 2
+
+        # First content should be text
+        assert isinstance(message.contents[0], TextContent)
+        assert message.contents[0].text == "Check out this photo!"
+
+        # Second content should be UriContent for the image
+        assert isinstance(message.contents[1], UriContent)
+        assert message.contents[1].uri == "https://example.com/photo.jpg"
+        assert message.contents[1].media_type == "image/jpeg"
+
+    async def test_to_agent_input_with_file_attachment_and_fetcher(self):
+        """Test converting user message with file attachment using data fetcher."""
+        from datetime import datetime
+
+        from agent_framework import DataContent
+        from chatkit.types import FileAttachment, UserMessageItem
+
+        attachment = FileAttachment(
+            id="file_123",
+            name="report.pdf",
+            mime_type="application/pdf",
+            type="file",
+        )
+
+        input_item = UserMessageItem(
+            id="msg_1",
+            thread_id="thread_1",
+            created_at=datetime.now(),
+            type="user_message",
+            content=[UserMessageTextContent(text="Here's the document")],
+            attachments=[attachment],
+            inference_options={},
+        )
+
+        # Create converter with data fetcher
+        async def fetch_data(attachment_id: str) -> bytes:
+            return b"PDF content data"
+
+        converter = ThreadItemConverter(attachment_data_fetcher=fetch_data)
+        result = await converter.to_agent_input(input_item)
+
+        assert len(result) == 1
+        message = result[0]
+        assert len(message.contents) == 2
+
+        # First content should be text
+        assert isinstance(message.contents[0], TextContent)
+
+        # Second content should be DataContent for the file
+        assert isinstance(message.contents[1], DataContent)
+        assert message.contents[1].media_type == "application/pdf"
+
+    def test_task_to_input(self, converter):
+        """Test converting TaskItem to ChatMessage."""
+        from datetime import datetime
+
+        from chatkit.types import CustomTask, TaskItem
+
+        task_item = TaskItem(
+            id="task_1",
+            thread_id="thread_1",
+            created_at=datetime.now(),
+            type="task",
+            task=CustomTask(type="custom", title="Analysis", content="Analyzed the data"),
+        )
+
+        result = converter.task_to_input(task_item)
+        assert isinstance(result, ChatMessage)
+        assert result.role == Role.USER
+        assert "Analysis: Analyzed the data" in result.text
+        assert "<Task>" in result.text
+
+    def test_task_to_input_no_custom_task(self, converter):
+        """Test that non-custom tasks return None."""
+        from datetime import datetime
+
+        from chatkit.types import TaskItem, ThoughtTask
+
+        task_item = TaskItem(
+            id="task_1",
+            thread_id="thread_1",
+            created_at=datetime.now(),
+            type="task",
+            task=ThoughtTask(type="thought", title="Think", content="Thinking..."),
+        )
+
+        result = converter.task_to_input(task_item)
+        assert result is None
+
+    def test_workflow_to_input(self, converter):
+        """Test converting WorkflowItem to ChatMessages."""
+        from datetime import datetime
+
+        from chatkit.types import CustomTask, Workflow, WorkflowItem
+
+        workflow_item = WorkflowItem(
+            id="wf_1",
+            thread_id="thread_1",
+            created_at=datetime.now(),
+            type="workflow",
+            workflow=Workflow(
+                type="custom",
+                tasks=[
+                    CustomTask(type="custom", title="Step 1", content="First step"),
+                    CustomTask(type="custom", title="Step 2", content="Second step"),
+                ],
+            ),
+        )
+
+        result = converter.workflow_to_input(workflow_item)
+        assert isinstance(result, list)
+        assert len(result) == 2
+        assert all(isinstance(msg, ChatMessage) for msg in result)
+        assert "Step 1: First step" in result[0].text
+        assert "Step 2: Second step" in result[1].text
+
+    def test_workflow_to_input_empty(self, converter):
+        """Test that workflows with no custom tasks return None."""
+        from datetime import datetime
+
+        from chatkit.types import Workflow, WorkflowItem
+
+        workflow_item = WorkflowItem(
+            id="wf_1",
+            thread_id="thread_1",
+            created_at=datetime.now(),
+            type="workflow",
+            workflow=Workflow(type="custom", tasks=[]),
+        )
+
+        result = converter.workflow_to_input(workflow_item)
+        assert result is None
+
+    def test_widget_to_input(self, converter):
+        """Test converting WidgetItem to ChatMessage."""
+        from datetime import datetime
+
+        from chatkit.types import WidgetItem
+        from chatkit.widgets import Card, Text
+
+        widget_item = WidgetItem(
+            id="widget_1",
+            thread_id="thread_1",
+            created_at=datetime.now(),
+            type="widget",
+            widget=Card(key="card1", children=[Text(value="Hello")]),
+        )
+
+        result = converter.widget_to_input(widget_item)
+        assert isinstance(result, ChatMessage)
+        assert result.role == Role.USER
+        assert "widget_1" in result.text
+        assert "graphical UI widget" in result.text
+
+
+class TestSimpleToAgentInput:
+    """Tests for simple_to_agent_input helper function."""
+
+    async def test_simple_to_agent_input_empty_list(self):
+        """Test simple conversion with empty list."""
+        result = await simple_to_agent_input([])
+        assert result == []
+
+    async def test_simple_to_agent_input_with_text(self):
+        """Test simple conversion with text content."""
+        from datetime import datetime
+
+        from chatkit.types import UserMessageItem
+
+        input_item = UserMessageItem(
+            id="msg_1",
+            thread_id="thread_1",
+            created_at=datetime.now(),
+            type="user_message",
+            content=[UserMessageTextContent(text="Test message")],
+            attachments=[],
+            inference_options={},
+        )
+
+        result = await simple_to_agent_input(input_item)
+
+        assert len(result) == 1
+        assert isinstance(result[0], ChatMessage)
+        assert result[0].role == Role.USER
+        assert result[0].text == "Test message"

python/packages/chatkit/tests/test_streaming.py

@@ -0,0 +1,142 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for Agent Framework to ChatKit streaming utilities."""
+
+from unittest.mock import Mock
+
+from agent_framework import AgentRunResponseUpdate, Role, TextContent
+from chatkit.types import (
+    ThreadItemAddedEvent,
+    ThreadItemDoneEvent,
+    ThreadItemUpdated,
+)
+
+from agent_framework_chatkit import stream_agent_response
+
+
+class TestStreamAgentResponse:
+    """Tests for stream_agent_response function."""
+
+    async def test_stream_empty_response(self):
+        """Test streaming empty response."""
+
+        async def empty_stream():
+            return
+            yield  # Make it a generator
+
+        events = []
+        async for event in stream_agent_response(empty_stream(), thread_id="test_thread"):
+            events.append(event)
+
+        assert len(events) == 0
+
+    async def test_stream_single_text_update(self):
+        """Test streaming single text update."""
+
+        async def single_update_stream():
+            yield AgentRunResponseUpdate(role=Role.ASSISTANT, contents=[TextContent(text="Hello world")])
+
+        events = []
+        async for event in stream_agent_response(single_update_stream(), thread_id="test_thread"):
+            events.append(event)
+
+        # Should have: item_added, item_updated (delta), item_done
+        assert len(events) == 3
+
+        # Check event types
+        assert isinstance(events[0], ThreadItemAddedEvent)
+        assert isinstance(events[1], ThreadItemUpdated)
+        assert isinstance(events[2], ThreadItemDoneEvent)
+
+        # Check delta event
+        assert events[1].update.delta == "Hello world"
+
+        # Check final message content
+        assert len(events[2].item.content) == 1
+        assert events[2].item.content[0].text == "Hello world"
+
+    async def test_stream_multiple_text_updates(self):
+        """Test streaming multiple text updates."""
+
+        async def multiple_updates_stream():
+            yield AgentRunResponseUpdate(role=Role.ASSISTANT, contents=[TextContent(text="Hello ")])
+            yield AgentRunResponseUpdate(role=Role.ASSISTANT, contents=[TextContent(text="world!")])
+
+        events = []
+        async for event in stream_agent_response(multiple_updates_stream(), thread_id="test_thread"):
+            events.append(event)
+
+        # Should have: item_added, item_updated (delta 1), item_updated (delta 2), item_done
+        assert len(events) == 4
+
+        # Check event types
+        assert isinstance(events[0], ThreadItemAddedEvent)
+        assert isinstance(events[1], ThreadItemUpdated)
+        assert isinstance(events[2], ThreadItemUpdated)
+        assert isinstance(events[3], ThreadItemDoneEvent)
+
+        # Check delta events
+        assert events[1].update.delta == "Hello "
+        assert events[2].update.delta == "world!"
+
+        # Check final accumulated text
+        final_message_event = events[-1]
+        assert isinstance(final_message_event, ThreadItemDoneEvent)
+        assert final_message_event.item.content[0].text == "Hello world!"
+
+    async def test_stream_with_custom_id_generator(self):
+        """Test streaming with custom ID generator."""
+
+        def custom_id_generator(item_type: str) -> str:
+            return f"custom_{item_type}_123"
+
+        async def single_update_stream():
+            yield AgentRunResponseUpdate(role=Role.ASSISTANT, contents=[TextContent(text="Test")])
+
+        events = []
+        async for event in stream_agent_response(
+            single_update_stream(), thread_id="test_thread", generate_id=custom_id_generator
+        ):
+            events.append(event)
+
+        # Check that custom IDs are used
+        message_added_event = events[0]
+        assert message_added_event.item.id == "custom_msg_123"
+
+    async def test_stream_empty_content_updates(self):
+        """Test streaming updates with empty content."""
+
+        async def empty_content_stream():
+            yield AgentRunResponseUpdate(role=Role.ASSISTANT, contents=[])
+            yield AgentRunResponseUpdate(role=Role.ASSISTANT, contents=None)
+
+        events = []
+        async for event in stream_agent_response(empty_content_stream(), thread_id="test_thread"):
+            events.append(event)
+
+        # Should have item_added and item_done
+        assert len(events) == 2
+        assert isinstance(events[0], ThreadItemAddedEvent)
+        assert isinstance(events[1], ThreadItemDoneEvent)
+
+        # Final message should have empty content
+        assert len(events[1].item.content) == 0
+
+    async def test_stream_non_text_content(self):
+        """Test streaming updates with non-text content."""
+        # Mock a content object without text attribute
+        non_text_content = Mock()
+        # Don't set text attribute
+        del non_text_content.text
+
+        async def non_text_stream():
+            yield AgentRunResponseUpdate(role=Role.ASSISTANT, contents=[non_text_content])
+
+        events = []
+        async for event in stream_agent_response(non_text_stream(), thread_id="test_thread"):
+            events.append(event)
+
+        # Should have item_added and item_done, but no content since no text
+        assert len(events) == 2
+        assert isinstance(events[0], ThreadItemAddedEvent)
+        assert isinstance(events[1], ThreadItemDoneEvent)

python/packages/copilotstudio/pyproject.toml

@@ -4,7 +4,7 @@ description = "Copilot Studio integration for Microsoft Agent Framework."
 authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
 readme = "README.md"
 requires-python = ">=3.10"
-version = "1.0.0b251104"
+version = "1.0.0b251111"
 license-files = ["LICENSE"]
 urls.homepage = "https://aka.ms/agent-framework"
 urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
@@ -19,6 +19,7 @@ classifiers = [
   "Programming Language :: Python :: 3.11",
   "Programming Language :: Python :: 3.12",
   "Programming Language :: Python :: 3.13",
+  "Programming Language :: Python :: 3.14",
   "Typing :: Typed",
 ]
 dependencies = [

python/packages/core/agent_framework/_agents.py

@@ -587,9 +587,11 @@ class ChatAgent(BaseAgent):
         name: str | None = None,
         description: str | None = None,
         chat_message_store_factory: Callable[[], ChatMessageStoreProtocol] | None = None,
-        conversation_id: str | None = None,
         context_providers: ContextProvider | list[ContextProvider] | AggregateContextProvider | None = None,
         middleware: Middleware | list[Middleware] | None = None,
+        # chat option params
+        allow_multiple_tool_calls: bool | None = None,
+        conversation_id: str | None = None,
         frequency_penalty: float | None = None,
         logit_bias: dict[str | int, float] | None = None,
         max_tokens: int | None = None,
@@ -630,15 +632,17 @@ class ChatAgent(BaseAgent):
             description: A brief description of the agent's purpose.
             chat_message_store_factory: Factory function to create an instance of ChatMessageStoreProtocol.
                 If not provided, the default in-memory store will be used.
-            conversation_id: The conversation ID for service-managed threads.
-                Cannot be used together with chat_message_store_factory.
             context_providers: The collection of multiple context providers to include during agent invocation.
             middleware: List of middleware to intercept agent and function invocations.
+            allow_multiple_tool_calls: Whether to allow multiple tool calls in a single response.
+            conversation_id: The conversation ID for service-managed threads.
+                Cannot be used together with chat_message_store_factory.
             frequency_penalty: The frequency penalty to use.
             logit_bias: The logit bias to use.
             max_tokens: The maximum number of tokens to generate.
             metadata: Additional metadata to include in the request.
             model_id: The model_id to use for the agent.
+                This overrides the model_id set in the chat client if it contains one.
             presence_penalty: The presence penalty to use.
             response_format: The format of the response.
             seed: The random seed to use.
@@ -687,7 +691,8 @@ class ChatAgent(BaseAgent):
         self._local_mcp_tools = [tool for tool in normalized_tools if isinstance(tool, MCPTool)]
         agent_tools = [tool for tool in normalized_tools if not isinstance(tool, MCPTool)]
         self.chat_options = ChatOptions(
-            model_id=model_id,
+            model_id=model_id or (str(chat_client.model_id) if hasattr(chat_client, "model_id") else None),
+            allow_multiple_tool_calls=allow_multiple_tool_calls,
             conversation_id=conversation_id,
             frequency_penalty=frequency_penalty,
             instructions=instructions,
@@ -758,6 +763,7 @@ class ChatAgent(BaseAgent):
         messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
         *,
         thread: AgentThread | None = None,
+        allow_multiple_tool_calls: bool | None = None,
         frequency_penalty: float | None = None,
         logit_bias: dict[str | int, float] | None = None,
         max_tokens: int | None = None,
@@ -793,6 +799,7 @@ class ChatAgent(BaseAgent):
 
         Keyword Args:
             thread: The thread to use for the agent.
+            allow_multiple_tool_calls: Whether to allow multiple tool calls in a single response.
             frequency_penalty: The frequency penalty to use.
             logit_bias: The logit bias to use.
             max_tokens: The maximum number of tokens to generate.
@@ -844,6 +851,7 @@ class ChatAgent(BaseAgent):
         co = run_chat_options & ChatOptions(
             model_id=model_id,
             conversation_id=thread.service_thread_id,
+            allow_multiple_tool_calls=allow_multiple_tool_calls,
             frequency_penalty=frequency_penalty,
             logit_bias=logit_bias,
             max_tokens=max_tokens,
@@ -887,6 +895,7 @@ class ChatAgent(BaseAgent):
         messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
         *,
         thread: AgentThread | None = None,
+        allow_multiple_tool_calls: bool | None = None,
         frequency_penalty: float | None = None,
         logit_bias: dict[str | int, float] | None = None,
         max_tokens: int | None = None,
@@ -922,6 +931,7 @@ class ChatAgent(BaseAgent):
 
         Keyword Args:
             thread: The thread to use for the agent.
+            allow_multiple_tool_calls: Whether to allow multiple tool calls in a single response.
             frequency_penalty: The frequency penalty to use.
             logit_bias: The logit bias to use.
             max_tokens: The maximum number of tokens to generate.
@@ -971,6 +981,7 @@ class ChatAgent(BaseAgent):
 
         co = run_chat_options & ChatOptions(
             conversation_id=thread.service_thread_id,
+            allow_multiple_tool_calls=allow_multiple_tool_calls,
             frequency_penalty=frequency_penalty,
             logit_bias=logit_bias,
             max_tokens=max_tokens,

python/packages/core/agent_framework/_clients.py

@@ -19,7 +19,7 @@ from ._middleware import (
 )
 from ._serialization import SerializationMixin
 from ._threads import ChatMessageStoreProtocol
-from ._tools import ToolProtocol
+from ._tools import FUNCTION_INVOKING_CHAT_CLIENT_MARKER, FunctionInvocationConfiguration, ToolProtocol
 from ._types import ChatMessage, ChatOptions, ChatResponse, ChatResponseUpdate, ToolMode, prepare_messages
 
 if TYPE_CHECKING:
@@ -224,7 +224,7 @@ def _merge_chat_options(
     stop: str | Sequence[str] | None = None,
     store: bool | None = None,
     temperature: float | None = None,
-    tool_choice: ToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = "auto",
+    tool_choice: ToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = None,
     tools: list[ToolProtocol | dict[str, Any] | Callable[..., Any]] | None = None,
     top_p: float | None = None,
     user: str | None = None,
@@ -357,6 +357,10 @@ class BaseChatClient(SerializationMixin, ABC):
 
         self.middleware = middleware
 
+        self.function_invocation_configuration = (
+            FunctionInvocationConfiguration() if hasattr(self.__class__, FUNCTION_INVOKING_CHAT_CLIENT_MARKER) else None
+        )
+
     def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]:
         """Convert the instance to a dictionary.
 
@@ -492,7 +496,7 @@ class BaseChatClient(SerializationMixin, ABC):
         stop: str | Sequence[str] | None = None,
         store: bool | None = None,
         temperature: float | None = None,
-        tool_choice: ToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = "auto",
+        tool_choice: ToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = None,
         tools: ToolProtocol
         | Callable[..., Any]
         | MutableMapping[str, Any]
@@ -591,7 +595,7 @@ class BaseChatClient(SerializationMixin, ABC):
         stop: str | Sequence[str] | None = None,
         store: bool | None = None,
         temperature: float | None = None,
-        tool_choice: ToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = "auto",
+        tool_choice: ToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = None,
         tools: ToolProtocol
         | Callable[..., Any]
         | MutableMapping[str, Any]
@@ -718,6 +722,8 @@ class BaseChatClient(SerializationMixin, ABC):
         chat_message_store_factory: Callable[[], ChatMessageStoreProtocol] | None = None,
         context_providers: ContextProvider | list[ContextProvider] | AggregateContextProvider | None = None,
         middleware: Middleware | list[Middleware] | None = None,
+        allow_multiple_tool_calls: bool | None = None,
+        conversation_id: str | None = None,
         frequency_penalty: float | None = None,
         logit_bias: dict[str | int, float] | None = None,
         max_tokens: int | None = None,
@@ -755,6 +761,8 @@ class BaseChatClient(SerializationMixin, ABC):
                 If not provided, the default in-memory store will be used.
             context_providers: Context providers to include during agent invocation.
             middleware: List of middleware to intercept agent and function invocations.
+            allow_multiple_tool_calls: Whether to allow multiple tool calls per agent turn.
+            conversation_id: The conversation ID to associate with the agent's messages.
             frequency_penalty: The frequency penalty to use.
             logit_bias: The logit bias to use.
             max_tokens: The maximum number of tokens to generate.
@@ -805,6 +813,8 @@ class BaseChatClient(SerializationMixin, ABC):
             chat_message_store_factory=chat_message_store_factory,
             context_providers=context_providers,
             middleware=middleware,
+            allow_multiple_tool_calls=allow_multiple_tool_calls,
+            conversation_id=conversation_id,
             frequency_penalty=frequency_penalty,
             logit_bias=logit_bias,
             max_tokens=max_tokens,

python/packages/core/agent_framework/_mcp.py

@@ -19,7 +19,7 @@ from mcp.client.websocket import websocket_client
 from mcp.shared.context import RequestContext
 from mcp.shared.exceptions import McpError
 from mcp.shared.session import RequestResponder
-from pydantic import BaseModel, create_model
+from pydantic import BaseModel, Field, create_model
 
 from ._tools import AIFunction, HostedMCPSpecificApproval
 from ._types import ChatMessage, Contents, DataContent, Role, TextContent, UriContent
@@ -224,13 +224,20 @@ def _get_input_model_from_mcp_tool(tool: types.Tool) -> type[BaseModel]:
         prop_details = json.loads(prop_details) if isinstance(prop_details, str) else prop_details
 
         python_type = resolve_type(prop_details)
+        description = prop_details.get("description", "")
 
         # Create field definition for create_model
         if prop_name in required:
-            field_definitions[prop_name] = (python_type, ...)
+            field_definitions[prop_name] = (
+                (python_type, Field(description=description)) if description else (python_type, ...)
+            )
         else:
             default_value = prop_details.get("default", None)
-            field_definitions[prop_name] = (python_type, default_value)
+            field_definitions[prop_name] = (
+                (python_type, Field(default=default_value, description=description))
+                if description
+                else (python_type, default_value)
+            )
 
     return create_model(f"{tool.name}_input", **field_definitions)
 

python/packages/core/agent_framework/_tools.py

@@ -71,6 +71,7 @@ logger = get_logger()
 __all__ = [
     "FUNCTION_INVOKING_CHAT_CLIENT_MARKER",
     "AIFunction",
+    "FunctionInvocationConfiguration",
     "HostedCodeInterpreterTool",
     "HostedFileSearchTool",
     "HostedMCPSpecificApproval",
@@ -84,7 +85,8 @@ __all__ = [
 
 logger = get_logger()
 FUNCTION_INVOKING_CHAT_CLIENT_MARKER: Final[str] = "__function_invoking_chat_client__"
-DEFAULT_MAX_ITERATIONS: Final[int] = 10
+DEFAULT_MAX_ITERATIONS: Final[int] = 40
+DEFAULT_MAX_CONSECUTIVE_ERRORS_PER_REQUEST: Final[int] = 3
 TChatClient = TypeVar("TChatClient", bound="ChatClientProtocol")
 # region Helpers
 
@@ -156,34 +158,19 @@ def _parse_inputs(
 # region Tools
 @runtime_checkable
 class ToolProtocol(Protocol):
-    """Represents a generic tool that can be specified to an AI service.
+    """Represents a generic tool.
 
     This protocol defines the interface that all tools must implement to be compatible
-    with the agent framework.
+    with the agent framework. It is implemented by various tool classes such as HostedMCPTool,
+    HostedWebSearchTool, and AIFunction's. A AIFunction is usually created by the `ai_function` decorator.
+
+    Since each connector needs to parse tools differently, users can pass a dict to
+    specify a service-specific tool when no abstraction is available.
 
     Attributes:
         name: The name of the tool.
         description: A description of the tool, suitable for use in describing the purpose to a model.
         additional_properties: Additional properties associated with the tool.
-
-    Examples:
-        .. code-block:: python
-
-            from agent_framework import ToolProtocol
-
-
-            class CustomTool:
-                def __init__(self, name: str, description: str) -> None:
-                    self.name = name
-                    self.description = description
-                    self.additional_properties = None
-
-                def __str__(self) -> str:
-                    return f"CustomTool(name={self.name})"
-
-
-            # Tool now implements ToolProtocol
-            tool: ToolProtocol = CustomTool("my_tool", "Does something useful")
     """
 
     name: str
@@ -201,22 +188,11 @@ class ToolProtocol(Protocol):
 class BaseTool(SerializationMixin):
     """Base class for AI tools, providing common attributes and methods.
 
-    This class provides the foundation for creating custom tools with serialization support.
+    Used as the base class for the various tools in the agent framework, such as HostedMCPTool,
+    HostedWebSearchTool, and AIFunction.
 
-    Examples:
-        .. code-block:: python
-
-            from agent_framework import BaseTool
-
-
-            class MyCustomTool(BaseTool):
-                def __init__(self, name: str, custom_param: str) -> None:
-                    super().__init__(name=name, description="My custom tool")
-                    self.custom_param = custom_param
-
-
-            tool = MyCustomTool(name="custom", custom_param="value")
-            print(tool)  # MyCustomTool(name=custom, description=My custom tool)
+    Since each connector needs to parse tools differently, this class is not exposed directly to end users.
+    In most cases, users can pass a dict to specify a service-specific tool when no abstraction is available.
     """
 
     DEFAULT_EXCLUDE: ClassVar[set[str]] = {"additional_properties"}
@@ -551,6 +527,10 @@ def _default_histogram() -> Histogram:
 TClass = TypeVar("TClass", bound="SerializationMixin")
 
 
+class EmptyInputModel(BaseModel):
+    """An empty input model for functions with no parameters."""
+
+
 class AIFunction(BaseTool, Generic[ArgsT, ReturnT]):
     """A tool that wraps a Python function to make it callable by AI models.
 
@@ -602,8 +582,10 @@ class AIFunction(BaseTool, Generic[ArgsT, ReturnT]):
         name: str,
         description: str = "",
         approval_mode: Literal["always_require", "never_require"] | None = None,
+        max_invocations: int | None = None,
+        max_invocation_exceptions: int | None = None,
         additional_properties: dict[str, Any] | None = None,
-        func: Callable[..., Awaitable[ReturnT] | ReturnT],
+        func: Callable[..., Awaitable[ReturnT] | ReturnT] | None = None,
         input_model: type[ArgsT] | Mapping[str, Any] | None = None,
         **kwargs: Any,
     ) -> None:
@@ -614,6 +596,10 @@ class AIFunction(BaseTool, Generic[ArgsT, ReturnT]):
             description: A description of the function.
             approval_mode: Whether or not approval is required to run this tool.
                 Default is that approval is not needed.
+            max_invocations: The maximum number of times this function can be invoked.
+                If None, there is no limit. Should be at least 1.
+            max_invocation_exceptions: The maximum number of exceptions allowed during invocations.
+                If None, there is no limit. Should be at least 1.
             additional_properties: Additional properties to set on the function.
             func: The function to wrap.
             input_model: The Pydantic model that defines the input parameters for the function.
@@ -630,21 +616,56 @@ class AIFunction(BaseTool, Generic[ArgsT, ReturnT]):
         self.func = func
         self.input_model = self._resolve_input_model(input_model)
         self.approval_mode = approval_mode or "never_require"
+        if max_invocations is not None and max_invocations < 1:
+            raise ValueError("max_invocations must be at least 1 or None.")
+        if max_invocation_exceptions is not None and max_invocation_exceptions < 1:
+            raise ValueError("max_invocation_exceptions must be at least 1 or None.")
+        self.max_invocations = max_invocations
+        self.invocation_count = 0
+        self.max_invocation_exceptions = max_invocation_exceptions
+        self.invocation_exception_count = 0
         self._invocation_duration_histogram = _default_histogram()
         self.type: Literal["ai_function"] = "ai_function"
 
+    @property
+    def declaration_only(self) -> bool:
+        """Indicate whether the function is declaration only (i.e., has no implementation)."""
+        return self.func is None
+
     def _resolve_input_model(self, input_model: type[ArgsT] | Mapping[str, Any] | None) -> type[ArgsT]:
-        if input_model:
-            if inspect.isclass(input_model) and issubclass(input_model, BaseModel):
-                return input_model
-            if isinstance(input_model, Mapping):
-                return cast(type[ArgsT], _create_model_from_json_schema(self.name, input_model))
-            raise TypeError("input_model must be a Pydantic BaseModel subclass or a JSON schema dict.")
-        return cast(type[ArgsT], _create_input_model_from_func(self.func, self.name))
+        """Resolve the input model for the function."""
+        if input_model is None:
+            if self.func is None:
+                return cast(type[ArgsT], EmptyInputModel)
+            return cast(type[ArgsT], _create_input_model_from_func(func=self.func, name=self.name))
+        if inspect.isclass(input_model) and issubclass(input_model, BaseModel):
+            return input_model
+        if isinstance(input_model, Mapping):
+            return cast(type[ArgsT], _create_model_from_json_schema(self.name, input_model))
+        raise TypeError("input_model must be a Pydantic BaseModel subclass or a JSON schema dict.")
 
     def __call__(self, *args: Any, **kwargs: Any) -> ReturnT | Awaitable[ReturnT]:
         """Call the wrapped function with the provided arguments."""
-        return self.func(*args, **kwargs)
+        if self.func is None:
+            raise ToolException(f"Function '{self.name}' is declaration only and cannot be invoked.")
+        if self.max_invocations is not None and self.invocation_count >= self.max_invocations:
+            raise ToolException(
+                f"Function '{self.name}' has reached its maximum invocation limit, you can no longer use this tool."
+            )
+        if (
+            self.max_invocation_exceptions is not None
+            and self.invocation_exception_count >= self.max_invocation_exceptions
+        ):
+            raise ToolException(
+                f"Function '{self.name}' has reached its maximum exception limit, "
+                f"you tried to use this tool too many times and it kept failing."
+            )
+        self.invocation_count += 1
+        try:
+            return self.func(*args, **kwargs)
+        except Exception:
+            self.invocation_exception_count += 1
+            raise
 
     async def invoke(
         self,
@@ -664,6 +685,8 @@ class AIFunction(BaseTool, Generic[ArgsT, ReturnT]):
         Raises:
             TypeError: If arguments is not an instance of the expected input model.
         """
+        if self.declaration_only:
+            raise ToolException(f"Function '{self.name}' is declaration only and cannot be invoked.")
         global OBSERVABILITY_SETTINGS
         from .observability import OBSERVABILITY_SETTINGS
 
@@ -833,7 +856,7 @@ def _parse_annotation(annotation: Any) -> Any:
     return annotation
 
 
-def _create_input_model_from_func(func: Callable[..., Any], tool_name: str) -> type[BaseModel]:
+def _create_input_model_from_func(func: Callable[..., Any], name: str) -> type[BaseModel]:
     """Create a Pydantic model from a function's signature."""
     sig = inspect.signature(func)
     fields = {
@@ -844,7 +867,7 @@ def _create_input_model_from_func(func: Callable[..., Any], tool_name: str) -> t
         for pname, param in sig.parameters.items()
         if pname not in {"self", "cls"}
     }
-    return create_model(f"{tool_name}_input", **fields)  # type: ignore[call-overload, no-any-return]
+    return create_model(f"{name}_input", **fields)  # type: ignore[call-overload, no-any-return]
 
 
 # Map JSON Schema types to Pydantic types
@@ -907,6 +930,8 @@ def ai_function(
     name: str | None = None,
     description: str | None = None,
     approval_mode: Literal["always_require", "never_require"] | None = None,
+    max_invocations: int | None = None,
+    max_invocation_exceptions: int | None = None,
     additional_properties: dict[str, Any] | None = None,
 ) -> AIFunction[Any, ReturnT]: ...
 
@@ -918,6 +943,8 @@ def ai_function(
     name: str | None = None,
     description: str | None = None,
     approval_mode: Literal["always_require", "never_require"] | None = None,
+    max_invocations: int | None = None,
+    max_invocation_exceptions: int | None = None,
     additional_properties: dict[str, Any] | None = None,
 ) -> Callable[[Callable[..., ReturnT | Awaitable[ReturnT]]], AIFunction[Any, ReturnT]]: ...
 
@@ -928,6 +955,8 @@ def ai_function(
     name: str | None = None,
     description: str | None = None,
     approval_mode: Literal["always_require", "never_require"] | None = None,
+    max_invocations: int | None = None,
+    max_invocation_exceptions: int | None = None,
     additional_properties: dict[str, Any] | None = None,
 ) -> AIFunction[Any, ReturnT] | Callable[[Callable[..., ReturnT | Awaitable[ReturnT]]], AIFunction[Any, ReturnT]]:
     """Decorate a function to turn it into a AIFunction that can be passed to models and executed automatically.
@@ -940,6 +969,22 @@ def ai_function(
     with a string description as the second argument. You can also use Pydantic's
     ``Field`` class for more advanced configuration.
 
+    Args:
+        func: The function to decorate.
+
+    Keyword Args:
+        name: The name of the function. If not provided, the function's ``__name__``
+            attribute will be used.
+        description: A description of the function. If not provided, the function's
+            docstring will be used.
+        approval_mode: Whether or not approval is required to run this tool.
+            Default is that approval is not needed.
+        max_invocations: The maximum number of times this function can be invoked.
+            If None, there is no limit, should be at least 1.
+        max_invocation_exceptions: The maximum number of exceptions allowed during invocations.
+            If None, there is no limit, should be at least 1.
+        additional_properties: Additional properties to set on the function.
+
     Note:
         When approval_mode is set to "always_require", the function will not be executed
         until explicit approval is given, this only applies to the auto-invocation flow.
@@ -997,6 +1042,8 @@ def ai_function(
                 name=tool_name,
                 description=tool_desc,
                 approval_mode=approval_mode,
+                max_invocations=max_invocations,
+                max_invocation_exceptions=max_invocation_exceptions,
                 additional_properties=additional_properties or {},
                 func=f,
             )
@@ -1009,10 +1056,123 @@ def ai_function(
 # region Function Invoking Chat Client
 
 
+class FunctionInvocationConfiguration(SerializationMixin):
+    """Configuration for function invocation in chat clients.
+
+    This class is created automatically on every chat client that supports function invocation.
+    This means that for most cases you can just alter the attributes on the instance, rather then creating a new one.
+
+    Example:
+        .. code-block:: python
+            from agent_framework.openai import OpenAIChatClient
+
+            # Create an OpenAI chat client
+            client = OpenAIChatClient(api_key="your_api_key")
+
+            # Disable function invocation
+            client.function_invocation_config.enabled = False
+
+            # Set maximum iterations to 10
+            client.function_invocation_config.max_iterations = 10
+
+            # Enable termination on unknown function calls
+            client.function_invocation_config.terminate_on_unknown_calls = True
+
+            # Add additional tools for function execution
+            client.function_invocation_config.additional_tools = [my_custom_tool]
+
+            # Enable detailed error information in function results
+            client.function_invocation_config.include_detailed_errors = True
+
+            # You can also create a new configuration instance if needed
+            new_config = FunctionInvocationConfiguration(
+                enabled=True,
+                max_iterations=20,
+                terminate_on_unknown_calls=False,
+                additional_tools=[another_tool],
+                include_detailed_errors=False,
+            )
+
+            # and then assign it to the client
+            client.function_invocation_config = new_config
+
+
+    Attributes:
+        enabled: Whether function invocation is enabled.
+            When this is set to False, the client will not attempt to invoke any functions,
+            because the tool mode will be set to None.
+        max_iterations: Maximum number of function invocation iterations.
+            Each request to this client might end up making multiple requests to the model. Each time the model responds
+            with a function call request, this client might perform that invocation and send the results back to the
+            model in a new request. This property limits the number of times such a roundtrip is performed. The value
+            must be at least one, as it includes the initial request.
+            If you want to fully disable function invocation, use the ``enabled`` property.
+            The default is 40.
+        max_consecutive_errors_per_request: Maximum consecutive errors allowed per request.
+            The maximum number of consecutive function call errors allowed before stopping
+            further function calls for the request.
+            The default is 3.
+        terminate_on_unknown_calls: Whether to terminate on unknown function calls.
+            When False, call requests to any tools that aren't available to the client
+            will result in a response message automatically being created and returned to the inner client stating that
+            the tool couldn't be found. This behavior can help in cases where a model hallucinates a function, but it's
+            problematic if the model has been made aware of the existence of tools outside of the normal mechanisms, and
+            requests one of those. ``additional_tools`` can be used to help with that. But if instead the consumer wants
+            to know about all function call requests that the client can't handle, this can be set to True. Upon
+            receiving a request to call a function that the client doesn't know about, it will terminate the function
+            calling loop and return the response, leaving the handling of the function call requests to the consumer of
+            the client.
+        additional_tools: Additional tools to include for function execution.
+            These will not impact the requests sent by the client, which will pass through the
+            ``tools`` unmodified. However, if the inner client requests the invocation of a tool
+            that was not in ``ChatOptions.tools``, this ``additional_tools`` collection will also be consulted to look
+            for a corresponding tool. This is useful when the service might have been pre-configured to be aware of
+            certain tools that aren't also sent on each individual request. These tools are treated the same as
+            ``declaration_only`` tools and will be returned to the user.
+        include_detailed_errors: Whether to include detailed error information in function results.
+            When set to True, detailed error information such as exception type and message
+            will be included in the function result content when a function invocation fails.
+            When False, only a generic error message will be included.
+
+
+    """
+
+    def __init__(
+        self,
+        enabled: bool = True,
+        max_iterations: int = DEFAULT_MAX_ITERATIONS,
+        max_consecutive_errors_per_request: int = DEFAULT_MAX_CONSECUTIVE_ERRORS_PER_REQUEST,
+        terminate_on_unknown_calls: bool = False,
+        additional_tools: Sequence[ToolProtocol] | None = None,
+        include_detailed_errors: bool = False,
+    ) -> None:
+        """Initialize FunctionInvocationConfiguration.
+
+        Args:
+            enabled: Whether function invocation is enabled.
+            max_iterations: Maximum number of function invocation iterations.
+            max_consecutive_errors_per_request: Maximum consecutive errors allowed per request.
+            terminate_on_unknown_calls: Whether to terminate on unknown function calls.
+            additional_tools: Additional tools to include for function execution.
+            include_detailed_errors: Whether to include detailed error information in function results.
+        """
+        self.enabled = enabled
+        if max_iterations < 1:
+            raise ValueError("max_iterations must be at least 1.")
+        self.max_iterations = max_iterations
+        if max_consecutive_errors_per_request < 0:
+            raise ValueError("max_consecutive_errors_per_request must be 0 or more.")
+        self.max_consecutive_errors_per_request = max_consecutive_errors_per_request
+        self.terminate_on_unknown_calls = terminate_on_unknown_calls
+        self.additional_tools = additional_tools or []
+        self.include_detailed_errors = include_detailed_errors
+
+
 async def _auto_invoke_function(
     function_call_content: "FunctionCallContent | FunctionApprovalResponseContent",
     custom_args: dict[str, Any] | None = None,
     *,
+    config: FunctionInvocationConfiguration,
     tool_map: dict[str, AIFunction[BaseModel, Any]],
     sequence_index: int | None = None,
     request_index: int | None = None,
@@ -1025,6 +1185,7 @@ async def _auto_invoke_function(
         custom_args: Additional custom arguments to merge with parsed arguments.
 
     Keyword Args:
+        config: The function invocation configuration.
         tool_map: A mapping of tool names to AIFunction instances.
         sequence_index: The index of the function call in the sequence.
         request_index: The index of the request iteration.
@@ -1037,29 +1198,33 @@ async def _auto_invoke_function(
         KeyError: If the requested function is not found in the tool map.
     """
     from ._types import (
-        FunctionApprovalRequestContent,
-        FunctionApprovalResponseContent,
-        FunctionCallContent,
         FunctionResultContent,
     )
 
+    # Note: The scenarios for approval_mode="always_require", declaration_only, and
+    # terminate_on_unknown_calls are all handled in _try_execute_function_calls before
+    # this function is called. This function only handles the actual execution of approved,
+    # non-declaration-only functions.
+
     tool: AIFunction[BaseModel, Any] | None = None
-    if isinstance(function_call_content, FunctionCallContent):
+    if function_call_content.type == "function_call":
         tool = tool_map.get(function_call_content.name)
+        # Tool should exist because _try_execute_function_calls validates this
         if tool is None:
-            raise KeyError(f"No tool or function named '{function_call_content.name}'")
-        if tool.approval_mode == "always_require":
-            return FunctionApprovalRequestContent(id=function_call_content.call_id, function_call=function_call_content)
+            exc = KeyError(f'Function "{function_call_content.name}" not found.')
+            return FunctionResultContent(
+                call_id=function_call_content.call_id,
+                result=f'Error: Requested function "{function_call_content.name}" not found.',
+                exception=exc,
+            )
     else:
-        if isinstance(function_call_content, FunctionApprovalResponseContent):
-            if function_call_content.approved:
-                tool = tool_map.get(function_call_content.function_call.name)
-                if tool is None:
-                    # we assume it is a hosted tool
-                    return function_call_content
-                function_call_content = function_call_content.function_call
-            else:
-                raise ToolException("Unapproved tool cannot be executed.")
+        # Note: Unapproved tools (approved=False) are handled in _replace_approval_contents_with_results
+        # and never reach this function, so we only handle approved=True cases here.
+        tool = tool_map.get(function_call_content.function_call.name)
+        if tool is None:
+            # we assume it is a hosted tool
+            return function_call_content
+        function_call_content = function_call_content.function_call
 
     parsed_args: dict[str, Any] = dict(function_call_content.parse_arguments() or {})
 
@@ -1068,10 +1233,10 @@ async def _auto_invoke_function(
     try:
         args = tool.input_model.model_validate(merged_args)
     except ValidationError as exc:
-        return FunctionResultContent(
-            call_id=function_call_content.call_id,
-            exception=exc,
-        )
+        message = "Error: Argument parsing failed."
+        if config.include_detailed_errors:
+            message = f"{message} Exception: {exc}"
+        return FunctionResultContent(call_id=function_call_content.call_id, result=message, exception=exc)
     if not middleware_pipeline or (
         not hasattr(middleware_pipeline, "has_middlewares") and not middleware_pipeline.has_middlewares
     ):
@@ -1086,10 +1251,10 @@ async def _auto_invoke_function(
                 result=function_result,
             )
         except Exception as exc:
-            return FunctionResultContent(
-                call_id=function_call_content.call_id,
-                exception=exc,
-            )
+            message = "Error: Function failed."
+            if config.include_detailed_errors:
+                message = f"{message} Exception: {exc}"
+            return FunctionResultContent(call_id=function_call_content.call_id, result=message, exception=exc)
     # Execute through middleware pipeline if available
     from ._middleware import FunctionInvocationContext
 
@@ -1117,10 +1282,10 @@ async def _auto_invoke_function(
             result=function_result,
         )
     except Exception as exc:
-        return FunctionResultContent(
-            call_id=function_call_content.call_id,
-            exception=exc,
-        )
+        message = "Error: Function failed."
+        if config.include_detailed_errors:
+            message = f"{message} Exception: {exc}"
+        return FunctionResultContent(call_id=function_call_content.call_id, result=message, exception=exc)
 
 
 def _get_tool_map(
@@ -1141,7 +1306,7 @@ def _get_tool_map(
     return ai_function_list
 
 
-async def _execute_function_calls(
+async def _try_execute_function_calls(
     custom_args: dict[str, Any],
     attempt_idx: int,
     function_calls: Sequence["FunctionCallContent"] | Sequence["FunctionApprovalResponseContent"],
@@ -1149,6 +1314,7 @@ async def _execute_function_calls(
     | Callable[..., Any] \
     | MutableMapping[str, Any] \
     | Sequence[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]",
+    config: FunctionInvocationConfiguration,
     middleware_pipeline: Any = None,  # Optional MiddlewarePipeline to avoid circular imports
 ) -> Sequence["Contents"]:
     """Execute multiple function calls concurrently.
@@ -1158,22 +1324,33 @@ async def _execute_function_calls(
         attempt_idx: The index of the current attempt iteration.
         function_calls: A sequence of FunctionCallContent to execute.
         tools: The tools available for execution.
+        config: Configuration for function invocation.
         middleware_pipeline: Optional middleware pipeline to apply during execution.
 
     Returns:
-        A list of Contents containing the results of each function call.
+        A list of Contents containing the results of each function call,
+        or the approval requests if any function requires approval,
+        or the original function calls if any are declaration only.
     """
     from ._types import FunctionApprovalRequestContent, FunctionCallContent
 
     tool_map = _get_tool_map(tools)
     approval_tools = [tool_name for tool_name, tool in tool_map.items() if tool.approval_mode == "always_require"]
+    declaration_only = [tool_name for tool_name, tool in tool_map.items() if tool.declaration_only]
+    additional_tool_names = [tool.name for tool in config.additional_tools] if config.additional_tools else []
     # check if any are calling functions that need approval
     # if so, we return approval request for all
     approval_needed = False
+    declaration_only_flag = False
     for fcc in function_calls:
         if isinstance(fcc, FunctionCallContent) and fcc.name in approval_tools:
             approval_needed = True
             break
+        if isinstance(fcc, FunctionCallContent) and (fcc.name in declaration_only or fcc.name in additional_tool_names):
+            declaration_only_flag = True
+            break
+        if config.terminate_on_unknown_calls and isinstance(fcc, FunctionCallContent) and fcc.name not in tool_map:
+            raise KeyError(f'Error: Requested function "{fcc.name}" not found.')
     if approval_needed:
         # approval can only be needed for Function Call Contents, not Approval Responses.
         return [
@@ -1181,6 +1358,9 @@ async def _execute_function_calls(
             for fcc in function_calls
             if isinstance(fcc, FunctionCallContent)
         ]
+    if declaration_only_flag:
+        # return the declaration only tools to the user, since we cannot execute them.
+        return [fcc for fcc in function_calls if isinstance(fcc, FunctionCallContent)]
 
     # Run all function calls concurrently
     return await asyncio.gather(*[
@@ -1191,6 +1371,7 @@ async def _execute_function_calls(
             sequence_index=seq_idx,
             request_index=attempt_idx,
             middleware_pipeline=middleware_pipeline,
+            config=config,
         )
         for seq_idx, function_call in enumerate(function_calls)
     ])
@@ -1334,17 +1515,23 @@ def _handle_function_calls_response(
             # because the underlying function may not preserve it in kwargs
             stored_middleware_pipeline = kwargs.get("_function_middleware_pipeline")
 
-            # Get max_iterations from instance additional_properties or class attribute
-            instance_max_iterations: int = DEFAULT_MAX_ITERATIONS
-            if hasattr(self, "additional_properties") and self.additional_properties:
-                instance_max_iterations = self.additional_properties.get("max_iterations", DEFAULT_MAX_ITERATIONS)
-            elif hasattr(self.__class__, "MAX_ITERATIONS"):
-                instance_max_iterations = getattr(self.__class__, "MAX_ITERATIONS", DEFAULT_MAX_ITERATIONS)
+            # Get the config for function invocation (not part of ChatClientProtocol, hence getattr)
+            config: FunctionInvocationConfiguration | None = getattr(self, "function_invocation_configuration", None)
+            if not config:
+                # Default config if not set
+                config = FunctionInvocationConfiguration()
 
+            errors_in_a_row: int = 0
             prepped_messages = prepare_messages(messages)
             response: "ChatResponse | None" = None
             fcc_messages: "list[ChatMessage]" = []
-            for attempt_idx in range(instance_max_iterations):
+
+            # If tools are provided but tool_choice is not set, default to "auto" for function invocation
+            tools = _extract_tools(kwargs)
+            if tools and kwargs.get("tool_choice") is None:
+                kwargs["tool_choice"] = "auto"
+
+            for attempt_idx in range(config.max_iterations if config.enabled else 0):
                 fcc_todo = _collect_approval_responses(prepped_messages)
                 if fcc_todo:
                     tools = _extract_tools(kwargs)
@@ -1352,13 +1539,29 @@ def _handle_function_calls_response(
                     approved_responses = [resp for resp in fcc_todo.values() if resp.approved]
                     approved_function_results: list[Contents] = []
                     if approved_responses:
-                        approved_function_results = await _execute_function_calls(
+                        approved_function_results = await _try_execute_function_calls(
                             custom_args=kwargs,
                             attempt_idx=attempt_idx,
                             function_calls=approved_responses,
                             tools=tools,  # type: ignore
                             middleware_pipeline=stored_middleware_pipeline,
+                            config=config,
                         )
+                        if any(
+                            fcr.exception is not None
+                            for fcr in approved_function_results
+                            if isinstance(fcr, FunctionResultContent)
+                        ):
+                            errors_in_a_row += 1
+                            # no need to reset the counter here, since this is the start of a new attempt.
+                        if errors_in_a_row >= config.max_consecutive_errors_per_request:
+                            logger.warning(
+                                "Maximum consecutive function call errors reached (%d). "
+                                "Stopping further function calls for this request.",
+                                config.max_consecutive_errors_per_request,
+                            )
+                            # break out of the loop and do the fallback response
+                            break
                     _replace_approval_contents_with_results(prepped_messages, fcc_todo, approved_function_results)
 
                 response = await func(self, messages=prepped_messages, **kwargs)
@@ -1381,15 +1584,15 @@ def _handle_function_calls_response(
                 if function_calls and tools:
                     # Use the stored middleware pipeline instead of extracting from kwargs
                     # because kwargs may have been modified by the underlying function
-                    function_call_results: list[Contents] = await _execute_function_calls(
+                    function_call_results: list[Contents] = await _try_execute_function_calls(
                         custom_args=kwargs,
                         attempt_idx=attempt_idx,
                         function_calls=function_calls,
                         tools=tools,  # type: ignore
                         middleware_pipeline=stored_middleware_pipeline,
+                        config=config,
                     )
-
-                    # Check if we have approval requests in the results
+                    # Check if we have approval requests or function calls (not results) in the results
                     if any(isinstance(fccr, FunctionApprovalRequestContent) for fccr in function_call_results):
                         # Add approval requests to the existing assistant message (with tool_calls)
                         # instead of creating a separate tool message
@@ -1402,6 +1605,26 @@ def _handle_function_calls_response(
                             result_message = ChatMessage(role="assistant", contents=function_call_results)
                             response.messages.append(result_message)
                         return response
+                    if any(isinstance(fccr, FunctionCallContent) for fccr in function_call_results):
+                        # the function calls are already in the response, so we just continue
+                        return response
+
+                    if any(
+                        fcr.exception is not None
+                        for fcr in function_call_results
+                        if isinstance(fcr, FunctionResultContent)
+                    ):
+                        errors_in_a_row += 1
+                        if errors_in_a_row >= config.max_consecutive_errors_per_request:
+                            logger.warning(
+                                "Maximum consecutive function call errors reached (%d). "
+                                "Stopping further function calls for this request.",
+                                config.max_consecutive_errors_per_request,
+                            )
+                            # break out of the loop and do the fallback response
+                            break
+                    else:
+                        errors_in_a_row = 0
 
                     # add a single ChatMessage to the response with the results
                     result_message = ChatMessage(role="tool", contents=function_call_results)
@@ -1482,16 +1705,16 @@ def _handle_function_calls_streaming_response(
             # because the underlying function may not preserve it in kwargs
             stored_middleware_pipeline = kwargs.get("_function_middleware_pipeline")
 
-            # Get max_iterations from instance additional_properties or class attribute
-            instance_max_iterations: int = DEFAULT_MAX_ITERATIONS
-            if hasattr(self, "additional_properties") and self.additional_properties:
-                instance_max_iterations = self.additional_properties.get("max_iterations", DEFAULT_MAX_ITERATIONS)
-            elif hasattr(self.__class__, "MAX_ITERATIONS"):
-                instance_max_iterations = getattr(self.__class__, "MAX_ITERATIONS", DEFAULT_MAX_ITERATIONS)
+            # Get the config for function invocation (not part of ChatClientProtocol, hence getattr)
+            config: FunctionInvocationConfiguration | None = getattr(self, "function_invocation_configuration", None)
+            if not config:
+                # Default config if not set
+                config = FunctionInvocationConfiguration()
 
+            errors_in_a_row: int = 0
             prepped_messages = prepare_messages(messages)
             fcc_messages: "list[ChatMessage]" = []
-            for attempt_idx in range(instance_max_iterations):
+            for attempt_idx in range(config.max_iterations if config.enabled else 0):
                 fcc_todo = _collect_approval_responses(prepped_messages)
                 if fcc_todo:
                     tools = _extract_tools(kwargs)
@@ -1499,13 +1722,21 @@ def _handle_function_calls_streaming_response(
                     approved_responses = [resp for resp in fcc_todo.values() if resp.approved]
                     approved_function_results: list[Contents] = []
                     if approved_responses:
-                        approved_function_results = await _execute_function_calls(
+                        approved_function_results = await _try_execute_function_calls(
                             custom_args=kwargs,
                             attempt_idx=attempt_idx,
                             function_calls=approved_responses,
                             tools=tools,  # type: ignore
                             middleware_pipeline=stored_middleware_pipeline,
+                            config=config,
                         )
+                        if any(
+                            fcr.exception is not None
+                            for fcr in approved_function_results
+                            if isinstance(fcr, FunctionResultContent)
+                        ):
+                            errors_in_a_row += 1
+                            # no need to reset the counter here, since this is the start of a new attempt.
                     _replace_approval_contents_with_results(prepped_messages, fcc_todo, approved_function_results)
 
                 all_updates: list["ChatResponseUpdate"] = []
@@ -1551,15 +1782,16 @@ def _handle_function_calls_streaming_response(
                 if function_calls and tools:
                     # Use the stored middleware pipeline instead of extracting from kwargs
                     # because kwargs may have been modified by the underlying function
-                    function_call_results: list[Contents] = await _execute_function_calls(
+                    function_call_results: list[Contents] = await _try_execute_function_calls(
                         custom_args=kwargs,
                         attempt_idx=attempt_idx,
                         function_calls=function_calls,
                         tools=tools,  # type: ignore
                         middleware_pipeline=stored_middleware_pipeline,
+                        config=config,
                     )
 
-                    # Check if we have approval requests in the results
+                    # Check if we have approval requests or function calls (not results) in the results
                     if any(isinstance(fccr, FunctionApprovalRequestContent) for fccr in function_call_results):
                         # Add approval requests to the existing assistant message (with tool_calls)
                         # instead of creating a separate tool message
@@ -1575,6 +1807,26 @@ def _handle_function_calls_streaming_response(
                             yield ChatResponseUpdate(contents=function_call_results, role="assistant")
                             response.messages.append(result_message)
                         return
+                    if any(isinstance(fccr, FunctionCallContent) for fccr in function_call_results):
+                        # the function calls were already yielded.
+                        return
+
+                    if any(
+                        fcr.exception is not None
+                        for fcr in function_call_results
+                        if isinstance(fcr, FunctionResultContent)
+                    ):
+                        errors_in_a_row += 1
+                        if errors_in_a_row >= config.max_consecutive_errors_per_request:
+                            logger.warning(
+                                "Maximum consecutive function call errors reached (%d). "
+                                "Stopping further function calls for this request.",
+                                config.max_consecutive_errors_per_request,
+                            )
+                            # break out of the loop and do the fallback response
+                            break
+                    else:
+                        errors_in_a_row = 0
 
                     # add a single ChatMessage to the response with the results
                     result_message = ChatMessage(role="tool", contents=function_call_results)
@@ -1648,10 +1900,6 @@ def use_function_invocation(
     if getattr(chat_client, FUNCTION_INVOKING_CHAT_CLIENT_MARKER, False):
         return chat_client
 
-    # Set MAX_ITERATIONS as a class variable if not already set
-    if not hasattr(chat_client, "MAX_ITERATIONS"):
-        chat_client.MAX_ITERATIONS = DEFAULT_MAX_ITERATIONS  # type: ignore
-
     try:
         chat_client.get_response = _handle_function_calls_response(  # type: ignore
             func=chat_client.get_response,  # type: ignore

python/packages/core/agent_framework/_types.py

@@ -1050,6 +1050,50 @@ class DataContent(BaseContent):
     def has_top_level_media_type(self, top_level_media_type: Literal["application", "audio", "image", "text"]) -> bool:
         return _has_top_level_media_type(self.media_type, top_level_media_type)
 
+    @staticmethod
+    def detect_image_format_from_base64(image_base64: str) -> str:
+        """Detect image format from base64 data by examining the binary header.
+
+        Args:
+            image_base64: Base64 encoded image data
+
+        Returns:
+            Image format as string (png, jpeg, webp, gif) with png as fallback
+        """
+        try:
+            # Constants for image format detection
+            # ~75 bytes of binary data should be enough to detect most image formats
+            FORMAT_DETECTION_BASE64_CHARS = 100
+
+            # Decode a small portion to detect format
+            decoded_data = base64.b64decode(image_base64[:FORMAT_DETECTION_BASE64_CHARS])
+            if decoded_data.startswith(b"\x89PNG"):
+                return "png"
+            if decoded_data.startswith(b"\xff\xd8\xff"):
+                return "jpeg"
+            if decoded_data.startswith(b"RIFF") and b"WEBP" in decoded_data[:12]:
+                return "webp"
+            if decoded_data.startswith(b"GIF87a") or decoded_data.startswith(b"GIF89a"):
+                return "gif"
+            return "png"  # Default fallback
+        except Exception:
+            return "png"  # Fallback if decoding fails
+
+    @classmethod
+    def create_data_uri_from_base64(cls, image_base64: str) -> tuple[str, str]:
+        """Create a data URI and media type from base64 image data.
+
+        Args:
+            image_base64: Base64 encoded image data
+
+        Returns:
+            Tuple of (data_uri, media_type)
+        """
+        format_type = cls.detect_image_format_from_base64(image_base64)
+        uri = f"data:image/{format_type};base64,{image_base64}"
+        media_type = f"image/{format_type}"
+        return uri, media_type
+
 
 class UriContent(BaseContent):
     """Represents a URI content.

python/packages/core/agent_framework/_workflows/_agent_executor.py

@@ -2,11 +2,14 @@
 
 import logging
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, cast
+
+from agent_framework import FunctionApprovalRequestContent, FunctionApprovalResponseContent
 
 from .._agents import AgentProtocol, ChatAgent
 from .._threads import AgentThread
 from .._types import AgentRunResponse, AgentRunResponseUpdate, ChatMessage
+from ._checkpoint_encoding import decode_checkpoint_value, encode_checkpoint_value
 from ._conversation_state import encode_chat_messages
 from ._events import (
     AgentRunEvent,
@@ -14,6 +17,7 @@ from ._events import (
 )
 from ._executor import Executor, handler
 from ._message_utils import normalize_messages_input
+from ._request_info_mixin import response_handler
 from ._workflow_context import WorkflowContext
 
 logger = logging.getLogger(__name__)
@@ -83,6 +87,8 @@ class AgentExecutor(Executor):
         super().__init__(exec_id)
         self._agent = agent
         self._agent_thread = agent_thread or self._agent.get_new_thread()
+        self._pending_agent_requests: dict[str, FunctionApprovalRequestContent] = {}
+        self._pending_responses_to_agent: list[FunctionApprovalResponseContent] = []
         self._output_response = output_response
         self._cache: list[ChatMessage] = []
 
@@ -93,50 +99,6 @@ class AgentExecutor(Executor):
             return [AgentRunResponse]
         return []
 
-    async def _run_agent_and_emit(self, ctx: WorkflowContext[AgentExecutorResponse, AgentRunResponse]) -> None:
-        """Execute the underlying agent, emit events, and enqueue response.
-
-        Checks ctx.is_streaming() to determine whether to emit incremental AgentRunUpdateEvent
-        events (streaming mode) or a single AgentRunEvent (non-streaming mode).
-        """
-        if ctx.is_streaming():
-            # Streaming mode: emit incremental updates
-            updates: list[AgentRunResponseUpdate] = []
-            async for update in self._agent.run_stream(
-                self._cache,
-                thread=self._agent_thread,
-            ):
-                updates.append(update)
-                await ctx.add_event(AgentRunUpdateEvent(self.id, update))
-
-            if isinstance(self._agent, ChatAgent):
-                response_format = self._agent.chat_options.response_format
-                response = AgentRunResponse.from_agent_run_response_updates(
-                    updates,
-                    output_format_type=response_format,
-                )
-            else:
-                response = AgentRunResponse.from_agent_run_response_updates(updates)
-        else:
-            # Non-streaming mode: use run() and emit single event
-            response = await self._agent.run(
-                self._cache,
-                thread=self._agent_thread,
-            )
-            await ctx.add_event(AgentRunEvent(self.id, response))
-
-        if self._output_response:
-            await ctx.yield_output(response)
-
-        # Always construct a full conversation snapshot from inputs (cache)
-        # plus agent outputs (agent_run_response.messages). Do not mutate
-        # response.messages so AgentRunEvent remains faithful to the raw output.
-        full_conversation: list[ChatMessage] = list(self._cache) + list(response.messages)
-
-        agent_response = AgentExecutorResponse(self.id, response, full_conversation=full_conversation)
-        await ctx.send_message(agent_response)
-        self._cache.clear()
-
     @handler
     async def run(
         self, request: AgentExecutorRequest, ctx: WorkflowContext[AgentExecutorResponse, AgentRunResponse]
@@ -192,6 +154,31 @@ class AgentExecutor(Executor):
         self._cache = normalize_messages_input(messages)
         await self._run_agent_and_emit(ctx)
 
+    @response_handler
+    async def handle_user_input_response(
+        self,
+        original_request: FunctionApprovalRequestContent,
+        response: FunctionApprovalResponseContent,
+        ctx: WorkflowContext[AgentExecutorResponse, AgentRunResponse],
+    ) -> None:
+        """Handle user input responses for function approvals during agent execution.
+
+        This will hold the executor's execution until all pending user input requests are resolved.
+
+        Args:
+            original_request: The original function approval request sent by the agent.
+            response: The user's response to the function approval request.
+            ctx: The workflow context for emitting events and outputs.
+        """
+        self._pending_responses_to_agent.append(response)
+        self._pending_agent_requests.pop(original_request.id, None)
+
+        if not self._pending_agent_requests:
+            # All pending requests have been resolved; resume agent execution
+            self._cache = normalize_messages_input(ChatMessage(role="user", contents=self._pending_responses_to_agent))
+            self._pending_responses_to_agent.clear()
+            await self._run_agent_and_emit(ctx)
+
     async def snapshot_state(self) -> dict[str, Any]:
         """Capture current executor state for checkpointing.
 
@@ -226,6 +213,8 @@ class AgentExecutor(Executor):
         return {
             "cache": encode_chat_messages(self._cache),
             "agent_thread": serialized_thread,
+            "pending_agent_requests": encode_checkpoint_value(self._pending_agent_requests),
+            "pending_responses_to_agent": encode_checkpoint_value(self._pending_responses_to_agent),
         }
 
     async def restore_state(self, state: dict[str, Any]) -> None:
@@ -258,7 +247,109 @@ class AgentExecutor(Executor):
         else:
             self._agent_thread = self._agent.get_new_thread()
 
+        pending_requests_payload = state.get("pending_agent_requests")
+        if pending_requests_payload:
+            self._pending_agent_requests = decode_checkpoint_value(pending_requests_payload)
+
+        pending_responses_payload = state.get("pending_responses_to_agent")
+        if pending_responses_payload:
+            self._pending_responses_to_agent = decode_checkpoint_value(pending_responses_payload)
+
     def reset(self) -> None:
         """Reset the internal cache of the executor."""
         logger.debug("AgentExecutor %s: Resetting cache", self.id)
         self._cache.clear()
+
+    async def _run_agent_and_emit(self, ctx: WorkflowContext[AgentExecutorResponse, AgentRunResponse]) -> None:
+        """Execute the underlying agent, emit events, and enqueue response.
+
+        Checks ctx.is_streaming() to determine whether to emit incremental AgentRunUpdateEvent
+        events (streaming mode) or a single AgentRunEvent (non-streaming mode).
+        """
+        if ctx.is_streaming():
+            # Streaming mode: emit incremental updates
+            response = await self._run_agent_streaming(cast(WorkflowContext, ctx))
+        else:
+            # Non-streaming mode: use run() and emit single event
+            response = await self._run_agent(cast(WorkflowContext, ctx))
+
+        if response is None:
+            # Agent did not complete (e.g., waiting for user input); do not emit response
+            logger.info("AgentExecutor %s: Agent did not complete, awaiting user input", self.id)
+            return
+
+        if self._output_response:
+            await ctx.yield_output(response)
+
+        # Always construct a full conversation snapshot from inputs (cache)
+        # plus agent outputs (agent_run_response.messages). Do not mutate
+        # response.messages so AgentRunEvent remains faithful to the raw output.
+        full_conversation: list[ChatMessage] = list(self._cache) + list(response.messages)
+
+        agent_response = AgentExecutorResponse(self.id, response, full_conversation=full_conversation)
+        await ctx.send_message(agent_response)
+        self._cache.clear()
+
+    async def _run_agent(self, ctx: WorkflowContext) -> AgentRunResponse | None:
+        """Execute the underlying agent in non-streaming mode.
+
+        Args:
+            ctx: The workflow context for emitting events.
+
+        Returns:
+            The complete AgentRunResponse, or None if waiting for user input.
+        """
+        response = await self._agent.run(
+            self._cache,
+            thread=self._agent_thread,
+        )
+        await ctx.add_event(AgentRunEvent(self.id, response))
+
+        # Handle any user input requests
+        if response.user_input_requests:
+            for user_input_request in response.user_input_requests:
+                self._pending_agent_requests[user_input_request.id] = user_input_request
+                await ctx.request_info(user_input_request, FunctionApprovalResponseContent)
+            return None
+
+        return response
+
+    async def _run_agent_streaming(self, ctx: WorkflowContext) -> AgentRunResponse | None:
+        """Execute the underlying agent in streaming mode and collect the full response.
+
+        Args:
+            ctx: The workflow context for emitting events.
+
+        Returns:
+            The complete AgentRunResponse, or None if waiting for user input.
+        """
+        updates: list[AgentRunResponseUpdate] = []
+        user_input_requests: list[FunctionApprovalRequestContent] = []
+        async for update in self._agent.run_stream(
+            self._cache,
+            thread=self._agent_thread,
+        ):
+            updates.append(update)
+            await ctx.add_event(AgentRunUpdateEvent(self.id, update))
+
+            if update.user_input_requests:
+                user_input_requests.extend(update.user_input_requests)
+
+        # Build the final AgentRunResponse from the collected updates
+        if isinstance(self._agent, ChatAgent):
+            response_format = self._agent.chat_options.response_format
+            response = AgentRunResponse.from_agent_run_response_updates(
+                updates,
+                output_format_type=response_format,
+            )
+        else:
+            response = AgentRunResponse.from_agent_run_response_updates(updates)
+
+        # Handle any user input requests after the streaming completes
+        if user_input_requests:
+            for user_input_request in user_input_requests:
+                self._pending_agent_requests[user_input_request.id] = user_input_request
+                await ctx.request_info(user_input_request, FunctionApprovalResponseContent)
+            return None
+
+        return response

python/packages/core/agent_framework/_workflows/_handoff.py

@@ -85,8 +85,8 @@ def _clone_chat_agent(agent: ChatAgent) -> ChatAgent:
     # so we need to recombine them here to pass the complete tools list to the constructor.
     # This makes sure MCP tools are preserved when cloning agents for handoff workflows.
     all_tools = list(options.tools) if options.tools else []
-    if agent._local_mcp_tools:
-        all_tools.extend(agent._local_mcp_tools)
+    if agent._local_mcp_tools:  # type: ignore
+        all_tools.extend(agent._local_mcp_tools)  # type: ignore
 
     return ChatAgent(
         chat_client=agent.chat_client,
@@ -133,6 +133,14 @@ class _ConversationWithUserInput:
     full_conversation: list[ChatMessage] = field(default_factory=lambda: [])  # type: ignore[misc]
 
 
+@dataclass
+class _ConversationForUserInput:
+    """Internal message from coordinator to gateway specifying which agent will receive the response."""
+
+    conversation: list[ChatMessage]
+    next_agent_id: str
+
+
 class _AutoHandoffMiddleware(FunctionMiddleware):
     """Intercept handoff tool invocations and short-circuit execution with synthetic results."""
 
@@ -275,6 +283,7 @@ class _HandoffCoordinator(BaseGroupChatOrchestrator):
         termination_condition: Callable[[list[ChatMessage]], bool | Awaitable[bool]],
         id: str,
         handoff_tool_targets: Mapping[str, str] | None = None,
+        return_to_previous: bool = False,
     ) -> None:
         """Create a coordinator that manages routing between specialists and the user."""
         super().__init__(id)
@@ -284,6 +293,8 @@ class _HandoffCoordinator(BaseGroupChatOrchestrator):
         self._input_gateway_id = input_gateway_id
         self._termination_condition = termination_condition
         self._handoff_tool_targets = {k.lower(): v for k, v in (handoff_tool_targets or {}).items()}
+        self._return_to_previous = return_to_previous
+        self._current_agent_id: str | None = None  # Track the current agent handling conversation
 
     def _get_author_name(self) -> str:
         """Get the coordinator name for orchestrator-generated messages."""
@@ -293,7 +304,7 @@ class _HandoffCoordinator(BaseGroupChatOrchestrator):
     async def handle_agent_response(
         self,
         response: AgentExecutorResponse,
-        ctx: WorkflowContext[AgentExecutorRequest | list[ChatMessage], list[ChatMessage]],
+        ctx: WorkflowContext[AgentExecutorRequest | list[ChatMessage], list[ChatMessage] | _ConversationForUserInput],
     ) -> None:
         """Process an agent's response and determine whether to route, request input, or terminate."""
         # Hydrate coordinator state (and detect new run) using checkpointable executor state
@@ -329,6 +340,9 @@ class _HandoffCoordinator(BaseGroupChatOrchestrator):
         # Check for handoff from ANY agent (starting agent or specialist)
         target = self._resolve_specialist(response.agent_run_response, conversation)
         if target is not None:
+            # Update current agent when handoff occurs
+            self._current_agent_id = target
+            logger.info(f"Handoff detected: {source} -> {target}. Routing control to specialist '{target}'.")
             await self._persist_state(ctx)
             # Clean tool-related content before sending to next agent
             cleaned = clean_conversation_for_handoff(conversation)
@@ -340,10 +354,15 @@ class _HandoffCoordinator(BaseGroupChatOrchestrator):
         if not is_starting_agent and source not in self._specialist_ids:
             raise RuntimeError(f"HandoffCoordinator received response from unknown executor '{source}'.")
 
+        # Update current agent when they respond without handoff
+        self._current_agent_id = source
+        logger.info(
+            f"Agent '{source}' responded without handoff. "
+            f"Requesting user input. Return-to-previous: {self._return_to_previous}"
+        )
         await self._persist_state(ctx)
 
         if await self._check_termination():
-            logger.info("Handoff workflow termination condition met. Ending conversation.")
             # Clean the output conversation for display
             cleaned_output = clean_conversation_for_handoff(conversation)
             await ctx.yield_output(cleaned_output)
@@ -352,7 +371,13 @@ class _HandoffCoordinator(BaseGroupChatOrchestrator):
         # Clean conversation before sending to gateway for user input request
         # This removes tool messages that shouldn't be shown to users
         cleaned_for_display = clean_conversation_for_handoff(conversation)
-        await ctx.send_message(cleaned_for_display, target_id=self._input_gateway_id)
+
+        # The awaiting_agent_id is the agent that just responded and is awaiting user input
+        # This is the source of the current response
+        next_agent_id = source
+
+        message_to_gateway = _ConversationForUserInput(conversation=cleaned_for_display, next_agent_id=next_agent_id)
+        await ctx.send_message(message_to_gateway, target_id=self._input_gateway_id)  # type: ignore[arg-type]
 
     @handler
     async def handle_user_input(
@@ -367,14 +392,26 @@ class _HandoffCoordinator(BaseGroupChatOrchestrator):
 
         # Check termination before sending to agent
         if await self._check_termination():
-            logger.info("Handoff workflow termination condition met. Ending conversation.")
             await ctx.yield_output(list(self._conversation))
             return
 
-        # Clean before sending to starting agent
+        # Determine routing target based on return-to-previous setting
+        target_agent_id = self._starting_agent_id
+        if self._return_to_previous and self._current_agent_id:
+            # Route back to the current agent that's handling the conversation
+            target_agent_id = self._current_agent_id
+            logger.info(
+                f"Return-to-previous enabled: routing user input to current agent '{target_agent_id}' "
+                f"(bypassing coordinator '{self._starting_agent_id}')"
+            )
+        else:
+            logger.info(f"Routing user input to coordinator '{target_agent_id}'")
+        # Note: Stack is only used for specialist-to-specialist handoffs, not user input routing
+
+        # Clean before sending to target agent
         cleaned = clean_conversation_for_handoff(self._conversation)
         request = AgentExecutorRequest(messages=cleaned, should_respond=True)
-        await ctx.send_message(request, target_id=self._starting_agent_id)
+        await ctx.send_message(request, target_id=target_agent_id)
 
     def _resolve_specialist(self, agent_response: AgentRunResponse, conversation: list[ChatMessage]) -> str | None:
         """Resolve the specialist executor id requested by the agent response, if any."""
@@ -444,22 +481,27 @@ class _HandoffCoordinator(BaseGroupChatOrchestrator):
     def _snapshot_pattern_metadata(self) -> dict[str, Any]:
         """Serialize pattern-specific state.
 
-        Handoff has no additional metadata beyond base conversation state.
+        Includes the current agent for return-to-previous routing.
 
         Returns:
-            Empty dict (no pattern-specific state)
+            Dict containing current agent if return-to-previous is enabled
         """
+        if self._return_to_previous:
+            return {
+                "current_agent_id": self._current_agent_id,
+            }
         return {}
 
     def _restore_pattern_metadata(self, metadata: dict[str, Any]) -> None:
         """Restore pattern-specific state.
 
-        Handoff has no additional metadata beyond base conversation state.
+        Restores the current agent for return-to-previous routing.
 
         Args:
-            metadata: Pattern-specific state dict (ignored)
+            metadata: Pattern-specific state dict
         """
-        pass
+        if self._return_to_previous and "current_agent_id" in metadata:
+            self._current_agent_id = metadata["current_agent_id"]
 
     def _restore_conversation_from_state(self, state: Mapping[str, Any]) -> list[ChatMessage]:
         """Rehydrate the coordinator's conversation history from checkpointed state.
@@ -507,8 +549,21 @@ class _UserInputGateway(Executor):
         self._prompt = prompt or "Provide your next input for the conversation."
 
     @handler
-    async def request_input(self, conversation: list[ChatMessage], ctx: WorkflowContext) -> None:
+    async def request_input(self, message: _ConversationForUserInput, ctx: WorkflowContext) -> None:
         """Emit a `HandoffUserInputRequest` capturing the conversation snapshot."""
+        if not message.conversation:
+            raise ValueError("Handoff workflow requires non-empty conversation before requesting user input.")
+        request = HandoffUserInputRequest(
+            conversation=list(message.conversation),
+            awaiting_agent_id=message.next_agent_id,
+            prompt=self._prompt,
+            source_executor_id=self.id,
+        )
+        await ctx.request_info(request, object)
+
+    @handler
+    async def request_input_legacy(self, conversation: list[ChatMessage], ctx: WorkflowContext) -> None:
+        """Legacy handler for backward compatibility - emit user input request with starting agent."""
         if not conversation:
             raise ValueError("Handoff workflow requires non-empty conversation before requesting user input.")
         request = HandoffUserInputRequest(
@@ -558,7 +613,7 @@ def _as_user_messages(payload: Any) -> list[ChatMessage]:
 
 
 def _default_termination_condition(conversation: list[ChatMessage]) -> bool:
-    """Default termination: stop after 10 user messages to prevent infinite loops."""
+    """Default termination: stop after 10 user messages."""
     user_message_count = sum(1 for msg in conversation if msg.role == Role.USER)
     return user_message_count >= 10
 
@@ -743,6 +798,7 @@ class HandoffBuilder:
         )
         self._auto_register_handoff_tools: bool = True
         self._handoff_config: dict[str, list[str]] = {}  # Maps agent_id -> [target_agent_ids]
+        self._return_to_previous: bool = False
 
         if participants:
             self.participants(participants)
@@ -1198,6 +1254,77 @@ class HandoffBuilder:
         self._termination_condition = condition
         return self
 
+    def enable_return_to_previous(self, enabled: bool = True) -> "HandoffBuilder":
+        """Enable direct return to the current agent after user input, bypassing the coordinator.
+
+        When enabled, after a specialist responds without requesting another handoff, user input
+        routes directly back to that same specialist instead of always routing back to the
+        coordinator agent for re-evaluation.
+
+        This is useful when a specialist needs multiple turns with the user to gather information
+        or resolve an issue, avoiding unnecessary coordinator involvement while maintaining context.
+
+        Flow Comparison:
+
+        **Default (disabled):**
+            User -> Coordinator -> Specialist -> User -> Coordinator -> Specialist -> ...
+
+        **With return_to_previous (enabled):**
+            User -> Coordinator -> Specialist -> User -> Specialist -> ...
+
+        Args:
+            enabled: Whether to enable return-to-previous routing. Default is True.
+
+        Returns:
+            Self for method chaining.
+
+        Example:
+
+        .. code-block:: python
+
+            workflow = (
+                HandoffBuilder(participants=[triage, technical_support, billing])
+                .set_coordinator("triage")
+                .add_handoff(triage, [technical_support, billing])
+                .enable_return_to_previous()  # Enable direct return routing
+                .build()
+            )
+
+            # Flow: User asks question
+            # -> Triage routes to Technical Support
+            # -> Technical Support asks clarifying question
+            # -> User provides more info
+            # -> Routes back to Technical Support (not Triage)
+            # -> Technical Support continues helping
+
+        Multi-tier handoff example:
+
+        .. code-block:: python
+
+            workflow = (
+                HandoffBuilder(participants=[triage, specialist_a, specialist_b])
+                .set_coordinator("triage")
+                .add_handoff(triage, [specialist_a, specialist_b])
+                .add_handoff(specialist_a, specialist_b)
+                .enable_return_to_previous()
+                .build()
+            )
+
+            # Flow: User asks question
+            # -> Triage routes to Specialist A
+            # -> Specialist A hands off to Specialist B
+            # -> Specialist B asks clarifying question
+            # -> User provides more info
+            # -> Routes back to Specialist B (who is currently handling the conversation)
+
+        Note:
+            This feature routes to whichever agent most recently responded, whether that's
+            the coordinator or a specialist. The conversation continues with that agent until
+            they either hand off to another agent or the termination condition is met.
+        """
+        self._return_to_previous = enabled
+        return self
+
     def build(self) -> Workflow:
         """Construct the final Workflow instance from the configured builder.
 
@@ -1326,6 +1453,7 @@ class HandoffBuilder:
                 termination_condition=self._termination_condition,
                 id="handoff-coordinator",
                 handoff_tool_targets=handoff_tool_targets,
+                return_to_previous=self._return_to_previous,
             )
 
         wiring = _GroupChatConfig(

python/packages/core/agent_framework/ag_ui/__init__.py

@@ -0,0 +1,35 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import importlib
+from typing import Any
+
+PACKAGE_NAME = "agent_framework_ag_ui"
+PACKAGE_EXTRA = "ag-ui"
+_IMPORTS = [
+    "__version__",
+    "AgentFrameworkAgent",
+    "add_agent_framework_fastapi_endpoint",
+    "AGUIChatClient",
+    "AGUIEventConverter",
+    "AGUIHttpService",
+    "ConfirmationStrategy",
+    "DefaultConfirmationStrategy",
+    "TaskPlannerConfirmationStrategy",
+    "RecipeConfirmationStrategy",
+    "DocumentWriterConfirmationStrategy",
+]
+
+
+def __getattr__(name: str) -> Any:
+    if name in _IMPORTS:
+        try:
+            return getattr(importlib.import_module(PACKAGE_NAME), name)
+        except ModuleNotFoundError as exc:
+            raise ModuleNotFoundError(
+                f"The '{PACKAGE_EXTRA}' extra is not installed, please do `pip install agent-framework-{PACKAGE_EXTRA}`"
+            ) from exc
+    raise AttributeError(f"Module {PACKAGE_NAME} has no attribute {name}.")
+
+
+def __dir__() -> list[str]:
+    return _IMPORTS

python/packages/core/agent_framework/ag_ui/__init__.pyi

@@ -0,0 +1,29 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from agent_framework_ag_ui import (
+    AgentFrameworkAgent,
+    AGUIChatClient,
+    AGUIEventConverter,
+    AGUIHttpService,
+    ConfirmationStrategy,
+    DefaultConfirmationStrategy,
+    DocumentWriterConfirmationStrategy,
+    RecipeConfirmationStrategy,
+    TaskPlannerConfirmationStrategy,
+    __version__,
+    add_agent_framework_fastapi_endpoint,
+)
+
+__all__ = [
+    "AGUIChatClient",
+    "AGUIEventConverter",
+    "AGUIHttpService",
+    "AgentFrameworkAgent",
+    "ConfirmationStrategy",
+    "DefaultConfirmationStrategy",
+    "DocumentWriterConfirmationStrategy",
+    "RecipeConfirmationStrategy",
+    "TaskPlannerConfirmationStrategy",
+    "__version__",
+    "add_agent_framework_fastapi_endpoint",
+]

python/packages/core/agent_framework/chatkit/__init__.py

@@ -0,0 +1,23 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import importlib
+from typing import Any
+
+PACKAGE_NAME = "agent_framework_chatkit"
+PACKAGE_EXTRA = "chatkit"
+_IMPORTS = ["__version__", "ThreadItemConverter", "simple_to_agent_input", "stream_agent_response"]
+
+
+def __getattr__(name: str) -> Any:
+    if name in _IMPORTS:
+        try:
+            return getattr(importlib.import_module(PACKAGE_NAME), name)
+        except ModuleNotFoundError as exc:
+            raise ModuleNotFoundError(
+                f"The '{PACKAGE_EXTRA}' extra is not installed, please do `pip install agent-framework-{PACKAGE_EXTRA}`"
+            ) from exc
+    raise AttributeError(f"Module {PACKAGE_NAME} has no attribute {name}.")
+
+
+def __dir__() -> list[str]:
+    return _IMPORTS

python/packages/core/agent_framework/chatkit/__init__.pyi

@@ -0,0 +1,10 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from agent_framework_chatkit import (
+    ThreadItemConverter,
+    __version__,
+    simple_to_agent_input,
+    stream_agent_response,
+)
+
+__all__ = ["ThreadItemConverter", "__version__", "simple_to_agent_input", "stream_agent_response"]

python/packages/core/agent_framework/microsoft/__init__.py

@@ -14,9 +14,11 @@ _IMPORTS: dict[str, tuple[str, list[str]]] = {
     "PurviewAppLocation": ("agent_framework_purview", ["microsoft-purview", "purview"]),
     "PurviewLocationType": ("agent_framework_purview", ["microsoft-purview", "purview"]),
     "PurviewAuthenticationError": ("agent_framework_purview", ["microsoft-purview", "purview"]),
+    "PurviewPaymentRequiredError": ("agent_framework_purview", ["microsoft-purview", "purview"]),
     "PurviewRateLimitError": ("agent_framework_purview", ["microsoft-purview", "purview"]),
     "PurviewRequestError": ("agent_framework_purview", ["microsoft-purview", "purview"]),
     "PurviewServiceError": ("agent_framework_purview", ["microsoft-purview", "purview"]),
+    "CacheProvider": ("agent_framework_purview", ["microsoft-purview", "purview"]),
 }
 
 

python/packages/core/agent_framework/microsoft/__init__.pyi

@@ -2,10 +2,12 @@
 
 from agent_framework_copilotstudio import CopilotStudioAgent, __version__, acquire_token
 from agent_framework_purview import (
+    CacheProvider,
     PurviewAppLocation,
     PurviewAuthenticationError,
     PurviewChatPolicyMiddleware,
     PurviewLocationType,
+    PurviewPaymentRequiredError,
     PurviewPolicyMiddleware,
     PurviewRateLimitError,
     PurviewRequestError,
@@ -14,11 +16,13 @@ from agent_framework_purview import (
 )
 
 __all__ = [
+    "CacheProvider",
     "CopilotStudioAgent",
     "PurviewAppLocation",
     "PurviewAuthenticationError",
     "PurviewChatPolicyMiddleware",
     "PurviewLocationType",
+    "PurviewPaymentRequiredError",
     "PurviewPolicyMiddleware",
     "PurviewRateLimitError",
     "PurviewRequestError",

python/packages/core/agent_framework/observability.py

@@ -846,6 +846,7 @@ def _trace_get_response(
                 kwargs.get("model_id")
                 or (chat_options.model_id if (chat_options := kwargs.get("chat_options")) else None)
                 or getattr(self, "model_id", None)
+                or "unknown"
             )
             service_url = str(
                 service_url_func()
@@ -933,6 +934,7 @@ def _trace_get_streaming_response(
                 kwargs.get("model_id")
                 or (chat_options.model_id if (chat_options := kwargs.get("chat_options")) else None)
                 or getattr(self, "model_id", None)
+                or "unknown"
             )
             service_url = str(
                 service_url_func()
@@ -1324,7 +1326,10 @@ def _get_span(
     attributes: dict[str, Any],
     span_name_attribute: str,
 ) -> Generator["trace.Span", Any, Any]:
-    """Start a span for a agent run."""
+    """Start a span for a agent run.
+
+    Note: `attributes` must contain the `span_name_attribute` key.
+    """
     span = get_tracer().start_span(f"{attributes[OtelAttr.OPERATION]} {attributes[span_name_attribute]}")
     span.set_attributes(attributes)
     with trace.use_span(
@@ -1353,7 +1358,8 @@ def _get_span_attributes(**kwargs: Any) -> dict[str, Any]:
         attributes[SpanAttributes.LLM_SYSTEM] = system_name
     if provider_name := kwargs.get("provider_name"):
         attributes[OtelAttr.PROVIDER_NAME] = provider_name
-    attributes[SpanAttributes.LLM_REQUEST_MODEL] = kwargs.get("model", "unknown")
+    if model_id := kwargs.get("model", chat_options.model_id):
+        attributes[SpanAttributes.LLM_REQUEST_MODEL] = model_id
     if service_url := kwargs.get("service_url"):
         attributes[OtelAttr.ADDRESS] = service_url
     if conversation_id := kwargs.get("conversation_id", chat_options.conversation_id):

python/packages/core/agent_framework/openai/_assistants_client.py

@@ -502,8 +502,6 @@ class OpenAIAssistantsClient(OpenAIConfigMixin, BaseChatClient):
                     tool_outputs = []
                 if function_result_content.result:
                     output = prepare_function_call_results(function_result_content.result)
-                elif function_result_content.exception:
-                    output = "Error: " + str(function_result_content.exception)
                 else:
                     output = "No output received."
                 tool_outputs.append(ToolOutput(tool_call_id=call_id, output=output))

python/packages/core/agent_framework/openai/_chat_client.py

@@ -380,11 +380,6 @@ class OpenAIBaseChatClient(OpenAIBase, BaseChatClient):
                     args["tool_call_id"] = content.call_id
                     if content.result is not None:
                         args["content"] = prepare_function_call_results(content.result)
-                    elif content.exception is not None:
-                        # Send the exception message to the model
-                        # Otherwise we won't have any channels to talk to OpenAI
-                        # TODO(yuge): This should ideally be customizable
-                        args["content"] = "Error: " + str(content.exception)
                 case _:
                     if "content" not in args:
                         args["content"] = []

python/packages/core/agent_framework/openai/_responses_client.py

@@ -293,6 +293,14 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
                             # Map the parameter name and remove the old one
                             mapped_tool[api_param] = mapped_tool.pop(user_param)
 
+                    # Validate partial_images parameter for streaming image generation
+                    # OpenAI API requires partial_images to be between 0-3 (inclusive) for image_generation tool
+                    # Reference: https://platform.openai.com/docs/api-reference/responses/create#responses_create-tools-image_generation_tool-partial_images
+                    if "partial_images" in mapped_tool:
+                        partial_images = mapped_tool["partial_images"]
+                        if not isinstance(partial_images, int) or partial_images < 0 or partial_images > 3:
+                            raise ValueError("partial_images must be an integer between 0 and 3 (inclusive).")
+
                     response_tools.append(mapped_tool)
                 else:
                     response_tools.append(tool_dict)
@@ -501,8 +509,6 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
                 }
                 if content.result:
                     args["output"] = prepare_function_call_results(content.result)
-                if content.exception:
-                    args["output"] = "Error: " + str(content.exception)
                 return args
             case FunctionApprovalRequestContent():
                 return {
@@ -697,29 +703,8 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
                         uri = item.result
                         media_type = None
                         if not uri.startswith("data:"):
-                            # Raw base64 string - convert to proper data URI format
-                            # Detect format from base64 data
-                            import base64
-
-                            try:
-                                # Decode a small portion to detect format
-                                decoded_data = base64.b64decode(uri[:100])  # First ~75 bytes should be enough
-                                if decoded_data.startswith(b"\x89PNG"):
-                                    format_type = "png"
-                                elif decoded_data.startswith(b"\xff\xd8\xff"):
-                                    format_type = "jpeg"
-                                elif decoded_data.startswith(b"RIFF") and b"WEBP" in decoded_data[:12]:
-                                    format_type = "webp"
-                                elif decoded_data.startswith(b"GIF87a") or decoded_data.startswith(b"GIF89a"):
-                                    format_type = "gif"
-                                else:
-                                    # Default to png if format cannot be detected
-                                    format_type = "png"
-                            except Exception:
-                                # Fallback to png if decoding fails
-                                format_type = "png"
-                            uri = f"data:image/{format_type};base64,{uri}"
-                            media_type = f"image/{format_type}"
+                            # Raw base64 string - convert to proper data URI format using helper
+                            uri, media_type = DataContent.create_data_uri_from_base64(uri)
                         else:
                             # Parse media type from existing data URI
                             try:
@@ -935,6 +920,25 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
                             raw_representation=event,
                         )
                     )
+            case "response.image_generation_call.partial_image":
+                # Handle streaming partial image generation
+                image_base64 = event.partial_image_b64
+                partial_index = event.partial_image_index
+
+                # Use helper function to create data URI from base64
+                uri, media_type = DataContent.create_data_uri_from_base64(image_base64)
+
+                contents.append(
+                    DataContent(
+                        uri=uri,
+                        media_type=media_type,
+                        additional_properties={
+                            "partial_image_index": partial_index,
+                            "is_partial_image": True,
+                        },
+                        raw_representation=event,
+                    )
+                )
             case _:
                 logger.debug("Unparsed event of type: %s: %s", event.type, event)
 

python/packages/core/agent_framework/openai/_shared.py

@@ -17,7 +17,7 @@ from openai.types.chat import ChatCompletion, ChatCompletionChunk
 from openai.types.images_response import ImagesResponse
 from openai.types.responses.response import Response
 from openai.types.responses.response_stream_event import ResponseStreamEvent
-from packaging import version
+from packaging.version import parse
 from pydantic import SecretStr
 
 from .._logging import get_logger
@@ -58,8 +58,8 @@ def _check_openai_version_for_callable_api_key() -> None:
     If the version is too old, raise a ServiceInitializationError with helpful message.
     """
     try:
-        current_version = version.parse(openai.__version__)
-        min_required_version = version.parse("1.106.0")
+        current_version = parse(openai.__version__)
+        min_required_version = parse("1.106.0")
 
         if current_version < min_required_version:
             raise ServiceInitializationError(

python/packages/core/pyproject.toml

@@ -4,7 +4,7 @@ description = "Microsoft Agent Framework for building AI Agents with Python. Thi
 authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
 readme = "README.md"
 requires-python = ">=3.10"
-version = "1.0.0b251104"
+version = "1.0.0b251111"
 license-files = ["LICENSE"]
 urls.homepage = "https://aka.ms/agent-framework"
 urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
@@ -19,6 +19,7 @@ classifiers = [
   "Programming Language :: Python :: 3.11",
   "Programming Language :: Python :: 3.12",
   "Programming Language :: Python :: 3.13",
+  "Programming Language :: Python :: 3.14",
   "Typing :: Typed",
 ]
 dependencies = [
@@ -32,21 +33,23 @@ dependencies = [
     "opentelemetry-exporter-otlp-proto-grpc>=1.36.0",
     "opentelemetry-semantic-conventions-ai>=0.4.13",
     # connectors and functions
-    "openai>=1.99.0,<2",
+    "openai>=1.99.0",
     "azure-identity>=1,<2",
     "mcp[ws]>=1.13",
+    "packaging>=24.1",
 ]
 
 [project.optional-dependencies]
 all = [
     "agent-framework-a2a",
+    "agent-framework-ag-ui",
+    "agent-framework-anthropic",
     "agent-framework-azure-ai",
     "agent-framework-copilotstudio",
-    "agent-framework-mem0",
-    "agent-framework-redis",
     "agent-framework-devui",
+    "agent-framework-mem0",
     "agent-framework-purview",
-    "agent-framework-anthropic",
+    "agent-framework-redis",
 ]
 
 [tool.uv]

python/packages/core/tests/core/test_observability.py

@@ -279,6 +279,45 @@ async def test_chat_client_streaming_observability(
         assert span.attributes[OtelAttr.OUTPUT_MESSAGES] is not None
 
 
+async def test_chat_client_without_model_id_observability(mock_chat_client, span_exporter: InMemorySpanExporter):
+    """Test telemetry shouldn't fail when the model_id is not provided for unknown reason."""
+    client = use_observability(mock_chat_client)()
+    messages = [ChatMessage(role=Role.USER, text="Test")]
+    span_exporter.clear()
+    response = await client.get_response(messages=messages)
+
+    assert response is not None
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    span = spans[0]
+
+    assert span.name == "chat unknown"
+    assert span.attributes[OtelAttr.OPERATION.value] == OtelAttr.CHAT_COMPLETION_OPERATION
+    assert span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "unknown"
+
+
+async def test_chat_client_streaming_without_model_id_observability(
+    mock_chat_client, span_exporter: InMemorySpanExporter
+):
+    """Test streaming telemetry shouldn't fail when the model_id is not provided for unknown reason."""
+    client = use_observability(mock_chat_client)()
+    messages = [ChatMessage(role=Role.USER, text="Test")]
+    span_exporter.clear()
+    # Collect all yielded updates
+    updates = []
+    async for update in client.get_streaming_response(messages=messages):
+        updates.append(update)
+
+    # Verify we got the expected updates, this shouldn't be dependent on otel
+    assert len(updates) == 2
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    span = spans[0]
+    assert span.name == "chat unknown"
+    assert span.attributes[OtelAttr.OPERATION.value] == OtelAttr.CHAT_COMPLETION_OPERATION
+    assert span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "unknown"
+
+
 def test_prepend_user_agent_with_none_value():
     """Test prepend user agent with None value in headers."""
     headers = {"User-Agent": None}
@@ -368,6 +407,7 @@ def mock_chat_agent():
             self.name = "test_agent"
             self.display_name = "Test Agent"
             self.description = "Test agent description"
+            self.chat_options = ChatOptions(model_id="TestModel")
 
         async def run(self, messages=None, *, thread=None, **kwargs):
             return AgentRunResponse(
@@ -405,7 +445,7 @@ async def test_agent_instrumentation_enabled(
     assert span.attributes[OtelAttr.AGENT_ID] == "test_agent_id"
     assert span.attributes[OtelAttr.AGENT_NAME] == "Test Agent"
     assert span.attributes[OtelAttr.AGENT_DESCRIPTION] == "Test agent description"
-    assert span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "unknown"
+    assert span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "TestModel"
     assert span.attributes[OtelAttr.INPUT_TOKENS] == 15
     assert span.attributes[OtelAttr.OUTPUT_TOKENS] == 25
     if enable_sensitive_data:
@@ -433,7 +473,7 @@ async def test_agent_streaming_response_with_diagnostics_enabled_via_decorator(
     assert span.attributes[OtelAttr.AGENT_ID] == "test_agent_id"
     assert span.attributes[OtelAttr.AGENT_NAME] == "Test Agent"
     assert span.attributes[OtelAttr.AGENT_DESCRIPTION] == "Test agent description"
-    assert span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "unknown"
+    assert span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "TestModel"
     if enable_sensitive_data:
         assert span.attributes.get(OtelAttr.OUTPUT_MESSAGES) is not None  # Streaming, so no usage yet
 

python/packages/core/tests/core/test_tools.py

@@ -63,6 +63,26 @@ def test_ai_function_decorator_without_args():
     assert test_tool(1, 2) == 3
 
 
+def test_ai_function_without_args():
+    """Test the ai_function decorator."""
+
+    @ai_function
+    def test_tool() -> int:
+        """A simple function that adds two numbers."""
+        return 1 + 2
+
+    assert isinstance(test_tool, ToolProtocol)
+    assert isinstance(test_tool, AIFunction)
+    assert test_tool.name == "test_tool"
+    assert test_tool.description == "A simple function that adds two numbers."
+    assert test_tool.parameters() == {
+        "properties": {},
+        "title": "test_tool_input",
+        "type": "object",
+    }
+    assert test_tool() == 3
+
+
 async def test_ai_function_decorator_with_async():
     """Test the ai_function decorator with an async function."""