chore: import upstream snapshot with attribution

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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