chore: import upstream snapshot with attribution

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@@ -0,0 +1,123 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for AgentFrameworkWorkflow wrapper behavior."""
+
+from __future__ import annotations
+
+from typing import Any, cast
+
+import pytest
+from agent_framework import Workflow, WorkflowBuilder, WorkflowContext, executor
+
+from agent_framework_ag_ui import AgentFrameworkWorkflow
+
+
+async def _run(agent: AgentFrameworkWorkflow, payload: dict[str, Any]) -> list[Any]:
+    return [event async for event in agent.run(payload)]
+
+
+def _interrupts_from_finished(event: Any) -> list[dict[str, Any]]:
+    dumped = event.model_dump(by_alias=True, exclude_none=True)
+    assert "interrupt" not in dumped
+    outcome = dumped.get("outcome")
+    assert isinstance(outcome, dict)
+    assert outcome.get("type") == "interrupt"
+    interrupts = outcome.get("interrupts")
+    assert isinstance(interrupts, list)
+    return cast(list[dict[str, Any]], interrupts)
+
+
+async def test_workflow_wrapper_rejects_workflow_and_factory_at_once() -> None:
+    """Workflow wrapper should reject ambiguous workflow source configuration."""
+
+    @executor(id="start")
+    async def start(message: Any, ctx: WorkflowContext) -> None:
+        del message
+        await ctx.yield_output("ok")  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+
+    workflow = WorkflowBuilder(start_executor=start).build()
+    with pytest.raises(ValueError, match="workflow_factory"):
+        AgentFrameworkWorkflow(workflow=workflow, workflow_factory=lambda _thread_id: workflow)
+
+
+async def test_workflow_wrapper_factory_is_thread_scoped() -> None:
+    """Thread-scoped workflow factories should isolate workflow instances by thread id."""
+
+    @executor(id="requester")
+    async def requester(message: Any, ctx: WorkflowContext) -> None:
+        del message
+        await ctx.request_info({"message": "Choose an option", "options": ["a", "b"]}, dict, request_id="choice")
+
+    factory_calls: dict[str, int] = {}
+
+    def workflow_factory(thread_id: str) -> Workflow:
+        factory_calls[thread_id] = factory_calls.get(thread_id, 0) + 1
+        return WorkflowBuilder(start_executor=requester).build()
+
+    agent = AgentFrameworkWorkflow(workflow_factory=workflow_factory)
+
+    first_events = await _run(
+        agent,
+        {
+            "thread_id": "thread-a",
+            "messages": [{"role": "user", "content": "start"}],
+        },
+    )
+    first_finished = [event for event in first_events if event.type == "RUN_FINISHED"][0]
+    first_interrupt = _interrupts_from_finished(first_finished)
+    assert first_interrupt[0]["id"] == "choice"
+    assert factory_calls["thread-a"] == 1
+
+    second_events = await _run(
+        agent,
+        {
+            "thread_id": "thread-a",
+            "messages": [],
+            "resume": {"interrupts": [{"id": "choice", "value": {"selection": "a"}}]},
+        },
+    )
+    second_types = [event.type for event in second_events]
+    assert "RUN_ERROR" not in second_types
+    second_finished = [event for event in second_events if event.type == "RUN_FINISHED"][0].model_dump(
+        by_alias=True, exclude_none=True
+    )
+    assert "outcome" not in second_finished
+    assert factory_calls["thread-a"] == 1
+
+    third_events = await _run(
+        agent,
+        {
+            "thread_id": "thread-b",
+            "messages": [{"role": "user", "content": "start"}],
+        },
+    )
+    third_finished = [event for event in third_events if event.type == "RUN_FINISHED"][0]
+    third_interrupt = _interrupts_from_finished(third_finished)
+    assert third_interrupt[0]["id"] == "choice"
+    assert factory_calls["thread-b"] == 1
+
+    agent.clear_thread_workflow("thread-a")
+    await _run(
+        agent,
+        {
+            "thread_id": "thread-a",
+            "messages": [{"role": "user", "content": "restart"}],
+        },
+    )
+    assert factory_calls["thread-a"] == 2
+
+
+async def test_workflow_wrapper_without_workflow_raises_not_implemented() -> None:
+    """Without workflow/workflow_factory, run should raise NotImplementedError."""
+    agent = AgentFrameworkWorkflow()
+
+    with pytest.raises(NotImplementedError, match="No workflow is attached"):
+        _ = [event async for event in agent.run({"messages": [{"role": "user", "content": "start"}]})]
+
+
+async def test_workflow_wrapper_factory_return_type_is_validated() -> None:
+    """Factory outputs must be Workflow instances."""
+    agent = AgentFrameworkWorkflow(workflow_factory=lambda _thread_id: cast(Any, object()))
+
+    with pytest.raises(TypeError, match="workflow_factory must return a Workflow instance"):
+        _ = [event async for event in agent.run({"thread_id": "thread-a", "messages": []})]