chore: import upstream snapshot with attribution

python/packages/core/tests/__init__.py

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

python/packages/core/tests/conftest.py

@@ -0,0 +1,89 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from collections.abc import Generator
+from typing import Any
+from unittest.mock import patch
+
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExporter
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+from pytest import fixture
+
+
+@fixture
+def enable_instrumentation(request: Any) -> bool:
+    """Fixture that returns a boolean indicating if Otel is enabled."""
+    return request.param if hasattr(request, "param") else True
+
+
+@fixture
+def enable_sensitive_data(request: Any) -> bool:
+    """Fixture that returns a boolean indicating if sensitive data is enabled."""
+    return request.param if hasattr(request, "param") else True
+
+
+@fixture
+def span_exporter(monkeypatch, enable_instrumentation: bool, enable_sensitive_data: bool) -> Generator[SpanExporter]:
+    """Fixture to remove environment variables for ObservabilitySettings."""
+
+    env_vars = [
+        "ENABLE_INSTRUMENTATION",
+        "ENABLE_SENSITIVE_DATA",
+        "ENABLE_CONSOLE_EXPORTERS",
+        "OTEL_EXPORTER_OTLP_ENDPOINT",
+        "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT",
+        "OTEL_EXPORTER_OTLP_METRICS_ENDPOINT",
+        "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT",
+        "OTEL_EXPORTER_OTLP_PROTOCOL",
+        "OTEL_EXPORTER_OTLP_HEADERS",
+        "OTEL_EXPORTER_OTLP_TRACES_HEADERS",
+        "OTEL_EXPORTER_OTLP_METRICS_HEADERS",
+        "OTEL_EXPORTER_OTLP_LOGS_HEADERS",
+        "OTEL_SERVICE_NAME",
+        "OTEL_SERVICE_VERSION",
+        "OTEL_RESOURCE_ATTRIBUTES",
+    ]
+
+    for key in env_vars:
+        monkeypatch.delenv(key, raising=False)  # type: ignore
+    monkeypatch.setenv("ENABLE_INSTRUMENTATION", str(enable_instrumentation))  # type: ignore
+    if not enable_instrumentation:
+        # we overwrite sensitive data for tests
+        enable_sensitive_data = False
+    monkeypatch.setenv("ENABLE_SENSITIVE_DATA", str(enable_sensitive_data))  # type: ignore
+    import importlib
+
+    from opentelemetry import trace
+
+    import agent_framework.observability as observability
+
+    # Reload the module to ensure a clean state for tests, then create a
+    # fresh ObservabilitySettings instance and patch the module attribute.
+    importlib.reload(observability)
+
+    # recreate observability settings with values from above and no file.
+    observability_settings = observability.ObservabilitySettings()
+
+    # Configure providers manually without calling _configure() to avoid OTLP imports
+    if enable_instrumentation or enable_sensitive_data:
+        from opentelemetry.sdk.trace import TracerProvider
+
+        tracer_provider = TracerProvider(resource=observability.create_resource())
+        trace.set_tracer_provider(tracer_provider)
+
+    monkeypatch.setattr(observability, "OBSERVABILITY_SETTINGS", observability_settings, raising=False)  # type: ignore
+
+    with (
+        patch("agent_framework.observability.OBSERVABILITY_SETTINGS", observability_settings),
+        patch("agent_framework.observability.configure_otel_providers"),
+    ):
+        exporter = InMemorySpanExporter()
+        if enable_instrumentation or enable_sensitive_data:
+            tracer_provider = trace.get_tracer_provider()  # type: ignore[assignment]
+            if not hasattr(tracer_provider, "add_span_processor"):
+                raise RuntimeError("Tracer provider does not support adding span processors.")
+
+            tracer_provider.add_span_processor(SimpleSpanProcessor(exporter))  # type: ignore
+
+        yield exporter
+        # Clean up
+        exporter.clear()

python/packages/core/tests/core/__init__.py

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

python/packages/core/tests/core/conftest.py

@@ -0,0 +1,330 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import logging
+import sys
+import warnings
+from collections.abc import AsyncIterable, Awaitable, MutableSequence, Sequence
+from typing import Any, Generic
+from typing import TypedDict as TypedDict  # noqa: F401  # pydantic mypy plugin needs TypedDict in module scope
+from unittest.mock import patch
+from uuid import uuid4
+
+from pytest import fixture
+
+warnings.filterwarnings(
+    "ignore",
+    message=r"\[SKILLS\].*",
+    category=FutureWarning,
+)
+
+from agent_framework import (  # noqa: E402
+    AgentResponse,
+    AgentResponseUpdate,
+    AgentSession,
+    BaseChatClient,
+    ChatMiddlewareLayer,
+    ChatResponse,
+    ChatResponseUpdate,
+    Content,
+    FunctionInvocationLayer,
+    FunctionTool,
+    Message,
+    ResponseStream,
+    SupportsAgentRun,
+    tool,
+)
+from agent_framework._clients import OptionsCoT  # noqa: E402
+from agent_framework.observability import ChatTelemetryLayer  # noqa: E402
+
+if sys.version_info >= (3, 12):
+    from typing import override  # type: ignore
+else:
+    from typing_extensions import override  # type: ignore[import]
+# region Chat History
+
+logger = logging.getLogger(__name__)
+
+
+@fixture(scope="function")
+def chat_history() -> list[Message]:
+    return []
+
+
+# region Tools
+
+
+@fixture
+def ai_tool() -> FunctionTool:
+    """Returns a generic FunctionTool."""
+
+    @tool
+    def generic_tool(name: str) -> str:
+        """A generic tool that echoes the name."""
+        return f"Hello, {name}"
+
+    return generic_tool
+
+
+@fixture
+def tool_tool() -> FunctionTool:
+    """Returns a executable FunctionTool."""
+
+    @tool(approval_mode="never_require")
+    def simple_function(x: int, y: int) -> int:
+        """A simple function that adds two numbers."""
+        return x + y
+
+    return simple_function
+
+
+# region Chat Clients
+class MockChatClient:
+    """Simple implementation of a chat client."""
+
+    def __init__(self, **kwargs: Any) -> None:
+        self.additional_properties: dict[str, Any] = {}
+        self.call_count: int = 0
+        self.responses: list[ChatResponse] = []
+        self.streaming_responses: list[list[ChatResponseUpdate]] = []
+        super().__init__(**kwargs)
+
+    def get_response(
+        self,
+        messages: str | Message | list[str] | list[Message],
+        *,
+        stream: bool = False,
+        options: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
+        options = options or {}
+        if stream:
+            return self._get_streaming_response(messages=messages, options=options, **kwargs)
+
+        async def _get() -> ChatResponse:
+            logger.debug(f"Running custom chat client, with: {messages=}, {kwargs=}")
+            self.call_count += 1
+            if self.responses:
+                return self.responses.pop(0)
+            return ChatResponse(messages=Message(role="assistant", contents=["test response"]))
+
+        return _get()
+
+    def _get_streaming_response(
+        self,
+        *,
+        messages: str | Message | list[str] | list[Message],
+        options: dict[str, Any],
+        **kwargs: Any,
+    ) -> ResponseStream[ChatResponseUpdate, ChatResponse]:
+        async def _stream() -> AsyncIterable[ChatResponseUpdate]:
+            logger.debug(f"Running custom chat client stream, with: {messages=}, {kwargs=}")
+            self.call_count += 1
+            if self.streaming_responses:
+                for update in self.streaming_responses.pop(0):
+                    yield update
+            else:
+                yield ChatResponseUpdate(contents=[Content.from_text("test streaming response ")], role="assistant")
+                yield ChatResponseUpdate(contents=[Content.from_text("another update")], role="assistant")
+
+        def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse:
+            return ChatResponse.from_updates(updates, output_format_type=options.get("response_format"))
+
+        return ResponseStream(_stream(), finalizer=_finalize)
+
+
+class MockBaseChatClient(
+    FunctionInvocationLayer[OptionsCoT],
+    ChatMiddlewareLayer[OptionsCoT],
+    ChatTelemetryLayer[OptionsCoT],
+    BaseChatClient[OptionsCoT],
+    Generic[OptionsCoT],
+):
+    """Mock implementation of a full-featured ChatClient."""
+
+    def __init__(self, **kwargs: Any):
+        super().__init__(middleware=[], **kwargs)
+        self.run_responses: list[ChatResponse] = []
+        self.streaming_responses: list[list[ChatResponseUpdate]] = []
+        self.call_count: int = 0
+
+    @override
+    def _inner_get_response(  # pyrefly: ignore[bad-override]  # ty: ignore[invalid-method-override]
+        self,
+        *,
+        messages: MutableSequence[Message],  # type: ignore[override]
+        stream: bool,
+        options: dict[str, Any],  # type: ignore[override]
+        **kwargs: Any,
+    ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
+        """Send a chat request to the AI service.
+
+        Args:
+            messages: The chat messages to send.
+            stream: Whether to stream the response.
+            options: The options dict for the request.
+            kwargs: Any additional keyword arguments.
+
+        Returns:
+            The chat response or ResponseStream.
+        """
+        if stream:
+            return self._get_streaming_response(messages=messages, options=options, **kwargs)
+
+        async def _get() -> ChatResponse:
+            return await self._get_non_streaming_response(messages=messages, options=options, **kwargs)
+
+        return _get()
+
+    async def _get_non_streaming_response(
+        self,
+        *,
+        messages: MutableSequence[Message],
+        options: dict[str, Any],
+        **kwargs: Any,
+    ) -> ChatResponse:
+        """Get a non-streaming response."""
+        logger.debug(f"Running base chat client inner, with: {messages=}, {options=}, {kwargs=}")
+        self.call_count += 1
+        if not self.run_responses:
+            return ChatResponse(messages=Message(role="assistant", contents=[f"test response - {messages[-1].text}"]))
+
+        response = self.run_responses.pop(0)
+
+        if options.get("tool_choice") == "none":
+            return ChatResponse(
+                messages=Message(
+                    role="assistant",
+                    contents=["I broke out of the function invocation loop..."],
+                ),
+                conversation_id=response.conversation_id,
+            )
+
+        return response
+
+    def _get_streaming_response(
+        self,
+        *,
+        messages: MutableSequence[Message],
+        options: dict[str, Any],
+        **kwargs: Any,
+    ) -> ResponseStream[ChatResponseUpdate, ChatResponse]:
+        """Get a streaming response."""
+
+        async def _stream() -> AsyncIterable[ChatResponseUpdate]:
+            logger.debug(f"Running base chat client inner stream, with: {messages=}, {options=}, {kwargs=}")
+            self.call_count += 1
+            if not self.streaming_responses:
+                yield ChatResponseUpdate(
+                    contents=[Content.from_text(f"update - {messages[0].text}")], role="assistant", finish_reason="stop"
+                )
+                return
+            if options.get("tool_choice") == "none":
+                yield ChatResponseUpdate(
+                    contents=[Content.from_text("I broke out of the function invocation loop...")],
+                    role="assistant",
+                    finish_reason="stop",
+                )
+                return
+            response = self.streaming_responses.pop(0)
+            for update in response:
+                yield update
+            await asyncio.sleep(0)
+
+        def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse:
+            return ChatResponse.from_updates(updates, output_format_type=options.get("response_format"))
+
+        return ResponseStream(_stream(), finalizer=_finalize)
+
+
+@fixture
+def enable_function_calling(request: Any) -> bool:
+    return request.param if hasattr(request, "param") else True
+
+
+@fixture
+def max_iterations(request: Any) -> int:
+    return request.param if hasattr(request, "param") else 2
+
+
+@fixture
+def client(enable_function_calling: bool, max_iterations: int) -> MockChatClient:
+    if enable_function_calling:
+        with patch("agent_framework._tools.DEFAULT_MAX_ITERATIONS", max_iterations):
+            return type("FunctionInvokingMockChatClient", (FunctionInvocationLayer, MockChatClient), {})()
+    return MockChatClient()
+
+
+@fixture
+def chat_client_base(enable_function_calling: bool, max_iterations: int) -> MockBaseChatClient:
+    with patch("agent_framework._tools.DEFAULT_MAX_ITERATIONS", max_iterations):
+        client = MockBaseChatClient()
+    if not enable_function_calling:
+        client.function_invocation_configuration["enabled"] = False
+    return client
+
+
+# region Agents
+class MockAgentSession(AgentSession):
+    pass
+
+
+# Mock Agent implementation for testing
+class MockAgent(SupportsAgentRun):
+    @property
+    def id(self) -> str:  # type: ignore[override]  # pyrefly: ignore[bad-override]
+        return str(uuid4())
+
+    @property
+    def name(self) -> str | None:  # type: ignore[override]  # pyrefly: ignore[bad-override]
+        """Returns the name of the agent."""
+        return "Name"
+
+    @property
+    def description(self) -> str | None:  # type: ignore[override]  # pyrefly: ignore[bad-override]
+        return "Description"
+
+    def run(  # type: ignore[override]  # pyrefly: ignore[bad-override]  # ty: ignore[invalid-method-override]
+        self,
+        messages: str | Message | list[str] | list[Message] | None = None,
+        *,
+        session: AgentSession | None = None,
+        stream: bool = False,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse] | AsyncIterable[AgentResponseUpdate]:
+        if stream:
+            return self._run_stream_impl(messages=messages, session=session, **kwargs)
+        return self._run_impl(messages=messages, session=session, **kwargs)
+
+    async def _run_impl(
+        self,
+        messages: str | Message | list[str] | list[Message] | None = None,
+        *,
+        session: AgentSession | None = None,
+        **kwargs: Any,
+    ) -> AgentResponse:
+        logger.debug(f"Running mock agent, with: {messages=}, {session=}, {kwargs=}")
+        return AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("Response")])])
+
+    async def _run_stream_impl(
+        self,
+        messages: str | Message | list[str] | list[Message] | None = None,
+        *,
+        session: AgentSession | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterable[AgentResponseUpdate]:
+        logger.debug(f"Running mock agent stream, with: {messages=}, {session=}, {kwargs=}")
+        yield AgentResponseUpdate(contents=[Content.from_text("Response")])
+
+    def create_session(self) -> AgentSession:  # type: ignore[override]  # pyrefly: ignore[bad-override]  # ty: ignore[invalid-method-override]
+        return MockAgentSession()
+
+
+@fixture
+def agent_session() -> AgentSession:
+    return MockAgentSession()
+
+
+@fixture
+def agent() -> SupportsAgentRun:
+    return MockAgent()  # type: ignore[abstract]  # pyrefly: ignore[bad-instantiation]

python/packages/core/tests/core/test_as_tool_kwargs_propagation.py

@@ -0,0 +1,371 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for kwargs propagation through as_tool() method."""
+
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+from agent_framework import Agent, ChatResponse, Content, Message, agent_middleware
+from agent_framework._middleware import AgentContext, FunctionInvocationContext
+
+from .conftest import MockChatClient
+
+
+class TestAsToolKwargsPropagation:
+    """Test cases for kwargs propagation through as_tool() delegation."""
+
+    @staticmethod
+    def _build_context(
+        tool: Any,
+        *,
+        task: str,
+        runtime_kwargs: dict[str, Any] | None = None,
+    ) -> FunctionInvocationContext:
+        return FunctionInvocationContext(
+            function=tool,
+            arguments={"task": task},
+            kwargs=runtime_kwargs,
+        )
+
+    async def test_as_tool_forwards_runtime_kwargs(self, client: MockChatClient) -> None:
+        """Test that runtime kwargs are forwarded through as_tool() to sub-agent tools."""
+        captured_kwargs: dict[str, Any] = {}
+        captured_function_invocation_kwargs: dict[str, Any] = {}
+
+        @agent_middleware
+        async def capture_middleware(context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
+            captured_kwargs.update(context.kwargs)
+            captured_function_invocation_kwargs.update(context.function_invocation_kwargs)
+            await call_next()
+
+        # Setup mock response
+        client.responses = [
+            ChatResponse(messages=[Message(role="assistant", contents=["Response from sub-agent"])]),
+        ]
+
+        # Create sub-agent with middleware
+        sub_agent = Agent(
+            client=client,  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+            name="sub_agent",
+            middleware=[capture_middleware],
+        )
+
+        # Create tool from sub-agent
+        tool = sub_agent.as_tool(name="delegate", arg_name="task")
+
+        # Directly invoke the tool with explicit runtime context (simulating agent execution).
+        _ = await tool.invoke(
+            context=self._build_context(
+                tool,
+                task="Test delegation",
+                runtime_kwargs={
+                    "api_token": "secret-xyz-123",
+                    "user_id": "user-456",
+                    "session_id": "session-789",
+                },
+            ),
+        )
+
+        assert captured_kwargs == {}
+        assert captured_function_invocation_kwargs["api_token"] == "secret-xyz-123"
+        assert captured_function_invocation_kwargs["user_id"] == "user-456"
+        assert captured_function_invocation_kwargs["session_id"] == "session-789"
+
+    async def test_as_tool_forwards_context_kwargs_verbatim(self, client: MockChatClient) -> None:
+        """Test that runtime kwargs are forwarded exactly from FunctionInvocationContext.kwargs."""
+        captured_function_invocation_kwargs: dict[str, Any] = {}
+
+        @agent_middleware
+        async def capture_middleware(context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
+            captured_function_invocation_kwargs.update(context.function_invocation_kwargs)
+            await call_next()
+
+        # Setup mock response
+        client.responses = [
+            ChatResponse(messages=[Message(role="assistant", contents=["Response from sub-agent"])]),
+        ]
+
+        sub_agent = Agent(
+            client=client,  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+            name="sub_agent",
+            middleware=[capture_middleware],
+        )
+
+        tool = sub_agent.as_tool(arg_name="custom_task")
+
+        # Invoke tool with both the arg_name field and additional kwargs
+        await tool.invoke(
+            context=FunctionInvocationContext(
+                function=tool,
+                arguments={"custom_task": "Test task"},
+                kwargs={
+                    "api_token": "token-123",
+                    "custom_task": "should_be_excluded",
+                },
+            )
+        )
+
+        assert captured_function_invocation_kwargs["custom_task"] == "should_be_excluded"
+        assert captured_function_invocation_kwargs["api_token"] == "token-123"
+
+    async def test_as_tool_nested_delegation_propagates_kwargs(self, client: MockChatClient) -> None:
+        """Test that runtime kwargs propagate through multiple levels of delegation (A -> B -> C)."""
+        captured_function_invocation_kwargs_list: list[dict[str, Any]] = []
+
+        @agent_middleware
+        async def capture_middleware(context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
+            captured_function_invocation_kwargs_list.append(dict(context.function_invocation_kwargs))
+            await call_next()
+
+        # Setup mock responses to trigger nested tool invocation: B calls tool C, then completes.
+        client.responses = [
+            ChatResponse(
+                messages=[
+                    Message(
+                        role="assistant",
+                        contents=[
+                            Content.from_function_call(
+                                call_id="call_c_1",
+                                name="call_c",
+                                arguments='{"task": "Please execute agent_c"}',
+                            )
+                        ],
+                    )
+                ]
+            ),
+            ChatResponse(messages=[Message(role="assistant", contents=["Response from agent_c"])]),
+            ChatResponse(messages=[Message(role="assistant", contents=["Response from agent_b"])]),
+        ]
+
+        # Create agent C (bottom level)
+        agent_c = Agent(
+            client=client,  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+            name="agent_c",
+            middleware=[capture_middleware],
+        )
+
+        # Create agent B (middle level) - delegates to C
+        agent_b = Agent(
+            client=client,  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+            name="agent_b",
+            tools=[agent_c.as_tool(name="call_c")],
+            middleware=[capture_middleware],
+        )
+
+        # Create tool from B for direct invocation
+        tool_b = agent_b.as_tool(name="call_b")
+
+        # Invoke tool B with kwargs - should propagate to both B and C
+        await tool_b.invoke(
+            context=self._build_context(
+                tool_b,
+                task="Test cascade",
+                runtime_kwargs={
+                    "trace_id": "trace-abc-123",
+                    "tenant_id": "tenant-xyz",
+                },
+            ),
+        )
+
+        assert len(captured_function_invocation_kwargs_list) >= 1
+        assert captured_function_invocation_kwargs_list[0].get("trace_id") == "trace-abc-123"
+        assert captured_function_invocation_kwargs_list[0].get("tenant_id") == "tenant-xyz"
+
+    async def test_as_tool_streaming_mode_forwards_kwargs(self, client: MockChatClient) -> None:
+        """Test that runtime kwargs are forwarded in streaming mode."""
+        captured_kwargs: dict[str, Any] = {}
+        captured_function_invocation_kwargs: dict[str, Any] = {}
+
+        @agent_middleware
+        async def capture_middleware(context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
+            captured_kwargs.update(context.kwargs)
+            captured_function_invocation_kwargs.update(context.function_invocation_kwargs)
+            await call_next()
+
+        # Setup mock streaming responses
+        from agent_framework import ChatResponseUpdate
+
+        client.streaming_responses = [
+            [ChatResponseUpdate(contents=[Content.from_text(text="Streaming response")], role="assistant")],
+        ]
+
+        sub_agent = Agent(
+            client=client,  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+            name="sub_agent",
+            middleware=[capture_middleware],
+        )
+
+        captured_updates: list[Any] = []
+
+        async def stream_callback(update: Any) -> None:
+            captured_updates.append(update)
+
+        tool = sub_agent.as_tool(stream_callback=stream_callback)
+
+        # Invoke tool with kwargs while streaming callback is active
+        await tool.invoke(
+            context=self._build_context(
+                tool,
+                task="Test streaming",
+                runtime_kwargs={"api_key": "streaming-key-999"},
+            ),
+        )
+
+        assert captured_kwargs == {}
+        assert captured_function_invocation_kwargs["api_key"] == "streaming-key-999"
+        assert len(captured_updates) == 1
+
+    async def test_as_tool_empty_kwargs_still_works(self, client: MockChatClient) -> None:
+        """Test that as_tool works correctly when no extra kwargs are provided."""
+        # Setup mock response
+        client.responses = [
+            ChatResponse(messages=[Message(role="assistant", contents=["Response from agent"])]),
+        ]
+
+        sub_agent = Agent(
+            client=client,  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+            name="sub_agent",
+        )
+
+        tool = sub_agent.as_tool()
+
+        # Invoke without any extra kwargs - should work without errors
+        result = await tool.invoke(arguments={"task": "Simple task"})
+
+        # Verify tool executed successfully
+        assert result is not None
+
+    async def test_as_tool_kwargs_with_chat_options(self, client: MockChatClient) -> None:
+        """Test that runtime kwargs are forwarded only via function_invocation_kwargs."""
+        captured_kwargs: dict[str, Any] = {}
+        captured_function_invocation_kwargs: dict[str, Any] = {}
+
+        @agent_middleware
+        async def capture_middleware(context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
+            captured_kwargs.update(context.kwargs)
+            captured_function_invocation_kwargs.update(context.function_invocation_kwargs)
+            await call_next()
+
+        # Setup mock response
+        client.responses = [
+            ChatResponse(messages=[Message(role="assistant", contents=["Response with options"])]),
+        ]
+
+        sub_agent = Agent(
+            client=client,  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+            name="sub_agent",
+            middleware=[capture_middleware],
+        )
+
+        tool = sub_agent.as_tool()
+
+        # Invoke with various kwargs
+        await tool.invoke(
+            context=self._build_context(
+                tool,
+                task="Test with options",
+                runtime_kwargs={
+                    "temperature": 0.8,
+                    "max_tokens": 500,
+                    "custom_param": "custom_value",
+                },
+            ),
+        )
+
+        assert captured_kwargs == {}
+        assert captured_function_invocation_kwargs["temperature"] == 0.8
+        assert captured_function_invocation_kwargs["max_tokens"] == 500
+        assert captured_function_invocation_kwargs["custom_param"] == "custom_value"
+
+    async def test_as_tool_kwargs_isolated_per_invocation(self, client: MockChatClient) -> None:
+        """Test that runtime kwargs are isolated per invocation and don't leak between calls."""
+        first_call_function_invocation_kwargs: dict[str, Any] = {}
+        second_call_function_invocation_kwargs: dict[str, Any] = {}
+        call_count = 0
+
+        @agent_middleware
+        async def capture_middleware(context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
+            nonlocal call_count
+            call_count += 1
+            if call_count == 1:
+                first_call_function_invocation_kwargs.update(context.function_invocation_kwargs)
+            elif call_count == 2:
+                second_call_function_invocation_kwargs.update(context.function_invocation_kwargs)
+            await call_next()
+
+        # Setup mock responses for both calls
+        client.responses = [
+            ChatResponse(messages=[Message(role="assistant", contents=["First response"])]),
+            ChatResponse(messages=[Message(role="assistant", contents=["Second response"])]),
+        ]
+
+        sub_agent = Agent(
+            client=client,  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+            name="sub_agent",
+            middleware=[capture_middleware],
+        )
+
+        tool = sub_agent.as_tool()
+
+        # First call with specific kwargs
+        await tool.invoke(
+            context=self._build_context(
+                tool,
+                task="First task",
+                runtime_kwargs={"session_id": "session-1", "api_token": "token-1"},
+            ),
+        )
+
+        # Second call with different kwargs
+        await tool.invoke(
+            context=self._build_context(
+                tool,
+                task="Second task",
+                runtime_kwargs={"session_id": "session-2", "api_token": "token-2"},
+            ),
+        )
+
+        assert first_call_function_invocation_kwargs.get("session_id") == "session-1"
+        assert first_call_function_invocation_kwargs.get("api_token") == "token-1"
+
+        assert second_call_function_invocation_kwargs.get("session_id") == "session-2"
+        assert second_call_function_invocation_kwargs.get("api_token") == "token-2"
+
+    async def test_as_tool_forwards_conversation_id_from_context_kwargs(self, client: MockChatClient) -> None:
+        """Test that conversation_id is forwarded when explicitly present in runtime context kwargs."""
+        captured_function_invocation_kwargs: dict[str, Any] = {}
+
+        @agent_middleware
+        async def capture_middleware(context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
+            captured_function_invocation_kwargs.update(context.function_invocation_kwargs)
+            await call_next()
+
+        # Setup mock response
+        client.responses = [
+            ChatResponse(messages=[Message(role="assistant", contents=["Response from sub-agent"])]),
+        ]
+
+        sub_agent = Agent(
+            client=client,  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+            name="sub_agent",
+            middleware=[capture_middleware],
+        )
+
+        tool = sub_agent.as_tool(name="delegate", arg_name="task")
+
+        # Invoke tool with conversation_id in kwargs (simulating parent's conversation state)
+        await tool.invoke(
+            context=self._build_context(
+                tool,
+                task="Test delegation",
+                runtime_kwargs={
+                    "conversation_id": "conv-parent-456",
+                    "api_token": "secret-xyz-123",
+                    "user_id": "user-456",
+                },
+            ),
+        )
+
+        assert captured_function_invocation_kwargs.get("conversation_id") == "conv-parent-456"
+        assert captured_function_invocation_kwargs.get("api_token") == "secret-xyz-123"
+        assert captured_function_invocation_kwargs.get("user_id") == "user-456"

python/packages/core/tests/core/test_azure_namespace.py

@@ -0,0 +1,10 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from agent_framework_azure_cosmos import CosmosHistoryProvider
+
+import agent_framework.azure as azure
+
+
+def test_azure_namespace_exposes_cosmos_history_provider() -> None:
+    assert azure.CosmosHistoryProvider is CosmosHistoryProvider
+    assert "CosmosHistoryProvider" in dir(azure)

python/packages/core/tests/core/test_clients.py

@@ -0,0 +1,510 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+
+from typing import Any
+from unittest.mock import patch
+
+import pytest
+
+from agent_framework import (
+    GROUP_ANNOTATION_KEY,
+    GROUP_TOKEN_COUNT_KEY,
+    BaseChatClient,
+    ChatResponse,
+    ChatResponseUpdate,
+    Content,
+    Message,
+    SlidingWindowStrategy,
+    SupportsChatGetResponse,
+    ToolResultCompactionStrategy,
+    TruncationStrategy,
+    tool,
+)
+
+
+class _FixedTokenizer:
+    def __init__(self, token_count: int) -> None:
+        self.token_count = token_count
+
+    def count_tokens(self, text: str) -> int:
+        return self.token_count
+
+
+def test_chat_client_type(client: SupportsChatGetResponse):
+    assert isinstance(client, SupportsChatGetResponse)
+
+
+async def test_chat_client_get_response(client: SupportsChatGetResponse):
+    response = await client.get_response([Message(role="user", contents=["Hello"])])
+    assert response.text == "test response"
+    assert response.messages[0].role == "assistant"
+
+
+async def test_chat_client_get_response_streaming(client: SupportsChatGetResponse):
+    async for update in client.get_response([Message(role="user", contents=["Hello"])], stream=True):
+        assert update.text == "test streaming response " or update.text == "another update"
+        assert update.role == "assistant"
+
+
+def test_base_client(chat_client_base: SupportsChatGetResponse):
+    assert isinstance(chat_client_base, BaseChatClient)
+    assert isinstance(chat_client_base, SupportsChatGetResponse)
+
+
+def test_base_client_rejects_direct_additional_properties(chat_client_base: SupportsChatGetResponse) -> None:
+    with pytest.raises(TypeError):
+        type(chat_client_base)(legacy_key="legacy-value")  # type: ignore[call-arg]  # pyrefly: ignore[bad-instantiation, unexpected-keyword]  # ty: ignore[unknown-argument]
+
+
+def test_base_client_as_agent_uses_explicit_additional_properties(chat_client_base: SupportsChatGetResponse) -> None:
+    agent = chat_client_base.as_agent(additional_properties={"team": "core"})  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+    assert agent.additional_properties == {"team": "core"}
+
+
+def test_base_client_as_agent_rejects_function_invocation_configuration(
+    chat_client_base: SupportsChatGetResponse,
+) -> None:
+    with pytest.raises(
+        TypeError,
+        match=r"as_agent\(\) got an unexpected keyword argument 'function_invocation_configuration'",
+    ):
+        chat_client_base.as_agent(function_invocation_configuration={"enabled": False})  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+
+async def test_base_client_get_response_uses_explicit_client_kwargs(chat_client_base: SupportsChatGetResponse) -> None:
+    async def fake_inner_get_response(**kwargs):
+        assert kwargs["trace_id"] == "trace-123"
+        assert "function_invocation_kwargs" not in kwargs
+        return ChatResponse(messages=[Message(role="assistant", contents=["ok"])])
+
+    with patch.object(
+        chat_client_base,
+        "_inner_get_response",
+        side_effect=fake_inner_get_response,
+    ) as mock_inner_get_response:
+        await chat_client_base.get_response(
+            [Message(role="user", contents=["hello"])],
+            function_invocation_kwargs={"tool_request_id": "tool-123"},
+            client_kwargs={"trace_id": "trace-123"},
+        )
+        mock_inner_get_response.assert_called_once()
+
+
+async def test_base_client_get_response(chat_client_base: SupportsChatGetResponse):
+    response = await chat_client_base.get_response([Message(role="user", contents=["Hello"])])
+    assert response.messages[0].role == "assistant"
+    assert response.messages[0].text == "test response - Hello"
+
+
+async def test_base_client_get_response_streaming(chat_client_base: SupportsChatGetResponse):
+    async for update in chat_client_base.get_response([Message(role="user", contents=["Hello"])], stream=True):
+        assert update.text == "update - Hello" or update.text == "another update"
+
+
+async def test_base_client_applies_compaction_before_non_streaming_inner_call(
+    chat_client_base: SupportsChatGetResponse,
+):
+    chat_client_base.function_invocation_configuration["enabled"] = False  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    chat_client_base.compaction_strategy = TruncationStrategy(max_n=1, compact_to=1)  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    captured_roles: list[list[str]] = []
+    original = chat_client_base._get_non_streaming_response  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+    async def _capture(
+        *,
+        messages: list[Message],
+        options: dict[str, Any],
+        **kwargs: Any,
+    ) -> ChatResponse:
+        captured_roles.append([message.role for message in messages])
+        return await original(messages=messages, options=options, **kwargs)
+
+    chat_client_base._get_non_streaming_response = _capture  # type: ignore[attr-defined, method-assign]  # ty: ignore[unresolved-attribute]
+    await chat_client_base.get_response([
+        Message(role="user", contents=["Hello"]),
+        Message(role="assistant", contents=["Previous response"]),
+    ])
+    assert captured_roles == [["assistant"]]
+
+
+async def test_base_client_applies_compaction_before_streaming_inner_call(
+    chat_client_base: SupportsChatGetResponse,
+):
+    chat_client_base.function_invocation_configuration["enabled"] = False  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    chat_client_base.compaction_strategy = TruncationStrategy(max_n=1, compact_to=1)  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    captured_roles: list[list[str]] = []
+    original = chat_client_base._get_streaming_response  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+    def _capture(
+        *,
+        messages: list[Message],
+        options: dict[str, Any],
+        **kwargs: Any,
+    ):
+        captured_roles.append([message.role for message in messages])
+        return original(messages=messages, options=options, **kwargs)
+
+    chat_client_base._get_streaming_response = _capture  # type: ignore[attr-defined, method-assign]  # ty: ignore[unresolved-attribute]
+    async for _ in chat_client_base.get_response(
+        [
+            Message(role="user", contents=["Hello"]),
+            Message(role="assistant", contents=["Previous response"]),
+        ],
+        stream=True,
+    ):
+        pass
+    assert captured_roles == [["assistant"]]
+
+
+async def test_base_client_per_call_compaction_override_applies_before_inner_call(
+    chat_client_base: SupportsChatGetResponse,
+) -> None:
+    chat_client_base.function_invocation_configuration["enabled"] = False  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    captured_roles: list[list[str]] = []
+    original = chat_client_base._get_non_streaming_response  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+    async def _capture(
+        *,
+        messages: list[Message],
+        options: dict[str, Any],
+        **kwargs: Any,
+    ) -> ChatResponse:
+        captured_roles.append([message.role for message in messages])
+        return await original(messages=messages, options=options, **kwargs)
+
+    chat_client_base._get_non_streaming_response = _capture  # type: ignore[attr-defined, method-assign]  # ty: ignore[unresolved-attribute]
+    await chat_client_base.get_response(
+        [
+            Message(role="user", contents=["Hello"]),
+            Message(role="assistant", contents=["Previous response"]),
+        ],
+        compaction_strategy=TruncationStrategy(max_n=1, compact_to=1),
+    )
+    assert captured_roles == [["assistant"]]
+
+
+async def test_base_client_per_call_tokenizer_override_annotates_messages(
+    chat_client_base: SupportsChatGetResponse,
+) -> None:
+    chat_client_base.function_invocation_configuration["enabled"] = False  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    captured_token_counts: list[list[int | None]] = []
+    original = chat_client_base._get_non_streaming_response  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+    async def _capture(
+        *,
+        messages: list[Message],
+        options: dict[str, Any],
+        **kwargs: Any,
+    ) -> ChatResponse:
+        captured_token_counts.append([
+            group.get(GROUP_TOKEN_COUNT_KEY) if isinstance(group, dict) else None
+            for group in (message.additional_properties.get(GROUP_ANNOTATION_KEY) for message in messages)
+        ])
+        return await original(messages=messages, options=options, **kwargs)
+
+    chat_client_base._get_non_streaming_response = _capture  # type: ignore[attr-defined, method-assign]  # ty: ignore[unresolved-attribute]
+    await chat_client_base.get_response(
+        [
+            Message(role="user", contents=["Hello"]),
+            Message(role="assistant", contents=["Previous response"]),
+        ],
+        compaction_strategy=SlidingWindowStrategy(keep_last_groups=2),
+        tokenizer=_FixedTokenizer(17),
+    )
+    assert captured_token_counts == [[17, 17]]
+
+
+async def test_base_client_per_call_tokenizer_override_without_strategy_annotates_messages(
+    chat_client_base: SupportsChatGetResponse,
+) -> None:
+    chat_client_base.function_invocation_configuration["enabled"] = False  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    captured_token_counts: list[list[int | None]] = []
+    original = chat_client_base._get_non_streaming_response  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+    async def _capture(
+        *,
+        messages: list[Message],
+        options: dict[str, Any],
+        **kwargs: Any,
+    ) -> ChatResponse:
+        captured_token_counts.append([
+            group.get(GROUP_TOKEN_COUNT_KEY) if isinstance(group, dict) else None
+            for group in (message.additional_properties.get(GROUP_ANNOTATION_KEY) for message in messages)
+        ])
+        return await original(messages=messages, options=options, **kwargs)
+
+    chat_client_base._get_non_streaming_response = _capture  # type: ignore[attr-defined, method-assign]  # ty: ignore[unresolved-attribute]
+    await chat_client_base.get_response(
+        [
+            Message(role="user", contents=["Hello"]),
+            Message(role="assistant", contents=["Previous response"]),
+        ],
+        tokenizer=_FixedTokenizer(17),
+    )
+    assert captured_token_counts == [[17, 17]]
+
+
+async def test_base_client_default_tokenizer_without_strategy_annotates_messages(
+    chat_client_base: SupportsChatGetResponse,
+) -> None:
+    chat_client_base.function_invocation_configuration["enabled"] = False  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    chat_client_base.tokenizer = _FixedTokenizer(19)  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    captured_token_counts: list[list[int | None]] = []
+    original = chat_client_base._get_non_streaming_response  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+    async def _capture(
+        *,
+        messages: list[Message],
+        options: dict[str, Any],
+        **kwargs: Any,
+    ) -> ChatResponse:
+        captured_token_counts.append([
+            group.get(GROUP_TOKEN_COUNT_KEY) if isinstance(group, dict) else None
+            for group in (message.additional_properties.get(GROUP_ANNOTATION_KEY) for message in messages)
+        ])
+        return await original(messages=messages, options=options, **kwargs)
+
+    chat_client_base._get_non_streaming_response = _capture  # type: ignore[attr-defined, method-assign]  # ty: ignore[unresolved-attribute]
+    await chat_client_base.get_response([
+        Message(role="user", contents=["Hello"]),
+        Message(role="assistant", contents=["Previous response"]),
+    ])
+    assert captured_token_counts == [[19, 19]]
+
+
+def _tool_call_response(call_id: str, location: str) -> ChatResponse:
+    return ChatResponse(
+        messages=Message(
+            role="assistant",
+            contents=[
+                Content.from_function_call(
+                    call_id=call_id,
+                    name="lookup_weather",
+                    arguments=f'{{"location": "{location}"}}',
+                )
+            ],
+        ),
+        response_id=f"resp_{call_id}",
+    )
+
+
+def _is_tool_result_summary(message: Message) -> bool:
+    text = message.text or ""
+    return message.role == "assistant" and text.startswith("[Tool results:")
+
+
+async def test_function_loop_persists_inserted_summaries_across_iterations(
+    chat_client_base: SupportsChatGetResponse,
+) -> None:
+    # Regression test for #4991: compaction inserts summary messages and excludes the
+    # originals. Across tool-loop iterations the exclusion flags persisted (shared Message
+    # objects) but the inserted summaries were dropped (they only lived on a throwaway copy),
+    # so older tool groups were silently lost with no summary representing them.
+    chat_client_base.function_invocation_configuration["enabled"] = True  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    chat_client_base.function_invocation_configuration["max_iterations"] = 3  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    chat_client_base.compaction_strategy = ToolResultCompactionStrategy(keep_last_tool_call_groups=1)  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+    @tool(name="lookup_weather", approval_mode="never_require")
+    def lookup_weather(location: str) -> str:
+        return f"Weather in {location}: sunny"
+
+    chat_client_base.run_responses = [  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+        _tool_call_response("call_1", "London"),
+        _tool_call_response("call_2", "Paris"),
+        _tool_call_response("call_3", "Tokyo"),
+    ]
+
+    captured_inputs: list[list[Message]] = []
+    original = chat_client_base._get_non_streaming_response  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+    async def _capture(
+        *,
+        messages: list[Message],
+        options: dict[str, Any],
+        **kwargs: Any,
+    ) -> ChatResponse:
+        captured_inputs.append(list(messages))
+        return await original(messages=messages, options=options, **kwargs)
+
+    chat_client_base._get_non_streaming_response = _capture  # type: ignore[attr-defined, method-assign]  # ty: ignore[unresolved-attribute]
+
+    await chat_client_base.get_response(
+        [Message(role="user", contents=["What is the weather in London?"])],
+        options={"tools": [lookup_weather]},  # type: ignore[typeddict-unknown-key]
+    )
+
+    # The final model call should represent every compacted tool group with a summary.
+    # Two older tool groups get collapsed (London, Paris) while the last (Tokyo) is kept.
+    final_input = captured_inputs[-1]
+    summaries = [message for message in final_input if _is_tool_result_summary(message)]
+    summary_text = " ".join(message.text or "" for message in summaries)
+
+    assert len(summaries) == 2, [message.text for message in final_input]
+    assert "London" in summary_text
+    assert "Paris" in summary_text
+
+
+def _tool_call_update(call_id: str, location: str) -> list[ChatResponseUpdate]:
+    return [
+        ChatResponseUpdate(
+            contents=[
+                Content.from_function_call(
+                    call_id=call_id,
+                    name="lookup_weather",
+                    arguments=f'{{"location": "{location}"}}',
+                )
+            ],
+            role="assistant",
+            finish_reason="stop",
+            response_id=f"resp_{call_id}",
+        )
+    ]
+
+
+async def test_function_loop_persists_inserted_summaries_across_iterations_streaming(
+    chat_client_base: SupportsChatGetResponse,
+) -> None:
+    # Streaming counterpart of the #4991 regression test: the summary persistence fix in
+    # ``_prepare_messages_for_model_call`` must cover the streaming tool loop too.
+    chat_client_base.function_invocation_configuration["enabled"] = True  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    chat_client_base.function_invocation_configuration["max_iterations"] = 3  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    chat_client_base.compaction_strategy = ToolResultCompactionStrategy(keep_last_tool_call_groups=1)  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+    @tool(name="lookup_weather", approval_mode="never_require")
+    def lookup_weather(location: str) -> str:
+        return f"Weather in {location}: sunny"
+
+    chat_client_base.streaming_responses = [  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+        _tool_call_update("call_1", "London"),
+        _tool_call_update("call_2", "Paris"),
+        _tool_call_update("call_3", "Tokyo"),
+    ]
+
+    captured_inputs: list[list[Message]] = []
+    original = chat_client_base._get_streaming_response  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+    def _capture(
+        *,
+        messages: list[Message],
+        options: dict[str, Any],
+        **kwargs: Any,
+    ):
+        captured_inputs.append(list(messages))
+        return original(messages=messages, options=options, **kwargs)
+
+    chat_client_base._get_streaming_response = _capture  # type: ignore[attr-defined, method-assign]  # ty: ignore[unresolved-attribute]
+
+    stream = chat_client_base.get_response(
+        [Message(role="user", contents=["What is the weather in London?"])],
+        stream=True,
+        options={"tools": [lookup_weather]},  # type: ignore[typeddict-unknown-key]
+    )
+    async for _ in stream:
+        pass
+
+    final_input = captured_inputs[-1]
+    summaries = [message for message in final_input if _is_tool_result_summary(message)]
+    summary_text = " ".join(message.text or "" for message in summaries)
+
+    assert len(summaries) == 2, [message.text for message in final_input]
+    assert "London" in summary_text
+    assert "Paris" in summary_text
+
+
+async def test_function_loop_compaction_conversation_id_mode_does_not_resend_history(
+    chat_client_base: SupportsChatGetResponse,
+) -> None:
+    # In conversation-id mode the server owns prior context, so the tool loop clears
+    # ``prepped_messages`` and only sends the latest message. Compaction must not fight that
+    # by re-inserting summaries or re-sending earlier turns.
+    chat_client_base.function_invocation_configuration["enabled"] = True  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    chat_client_base.function_invocation_configuration["max_iterations"] = 3  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    chat_client_base.compaction_strategy = ToolResultCompactionStrategy(keep_last_tool_call_groups=1)  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+    @tool(name="lookup_weather", approval_mode="never_require")
+    def lookup_weather(location: str) -> str:
+        return f"Weather in {location}: sunny"
+
+    def _conversation_tool_call(call_id: str, location: str) -> ChatResponse:
+        response = _tool_call_response(call_id, location)
+        response.conversation_id = "conv_1"
+        return response
+
+    chat_client_base.run_responses = [  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+        _conversation_tool_call("call_1", "London"),
+        _conversation_tool_call("call_2", "Paris"),
+        _conversation_tool_call("call_3", "Tokyo"),
+    ]
+
+    captured_inputs: list[list[Message]] = []
+    original = chat_client_base._get_non_streaming_response  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+    async def _capture(
+        *,
+        messages: list[Message],
+        options: dict[str, Any],
+        **kwargs: Any,
+    ) -> ChatResponse:
+        captured_inputs.append(list(messages))
+        return await original(messages=messages, options=options, **kwargs)
+
+    chat_client_base._get_non_streaming_response = _capture  # type: ignore[attr-defined, method-assign]  # ty: ignore[unresolved-attribute]
+
+    await chat_client_base.get_response(
+        [Message(role="user", contents=["What is the weather in London?"])],
+        options={"tools": [lookup_weather]},  # type: ignore[typeddict-unknown-key]
+    )
+
+    # After the conversation id is established the loop only forwards the latest message,
+    # so subsequent model calls never receive the full history or summary messages.
+    for sent in captured_inputs[1:]:
+        assert len(sent) <= 1, [message.text for message in sent]
+        assert not any(_is_tool_result_summary(message) for message in sent)
+
+
+def test_base_client_as_agent_does_not_copy_client_compaction_defaults(
+    chat_client_base: SupportsChatGetResponse,
+) -> None:
+    strategy = TruncationStrategy(max_n=1, compact_to=1)
+    tokenizer = _FixedTokenizer(11)
+    chat_client_base.compaction_strategy = strategy  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    chat_client_base.tokenizer = tokenizer  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+    agent = chat_client_base.as_agent(name="shared-client-agent")  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+    assert agent.compaction_strategy is None  # type: ignore[attr-defined]
+    assert agent.tokenizer is None  # type: ignore[attr-defined]
+
+
+async def test_chat_client_instructions_handling(chat_client_base: SupportsChatGetResponse):
+    instructions = "You are a helpful assistant."
+
+    async def fake_inner_get_response(**kwargs):
+        return ChatResponse(messages=[Message(role="assistant", contents=["ok"])])
+
+    with patch.object(
+        chat_client_base,
+        "_inner_get_response",
+        side_effect=fake_inner_get_response,
+    ) as mock_inner_get_response:
+        await chat_client_base.get_response(
+            [Message(role="user", contents=["hello"])], options={"instructions": instructions}
+        )
+        mock_inner_get_response.assert_called_once()
+        _, kwargs = mock_inner_get_response.call_args
+        messages = kwargs.get("messages", [])
+        assert len(messages) == 1
+        assert messages[0].role == "user"
+        assert messages[0].text == "hello"
+
+        from agent_framework._types import prepend_instructions_to_messages
+
+        appended_messages = prepend_instructions_to_messages(
+            [Message(role="user", contents=["hello"])],
+            instructions,
+        )
+        assert len(appended_messages) == 2
+        assert appended_messages[0].role == "system"
+        assert appended_messages[0].text == "You are a helpful assistant."
+        assert appended_messages[1].role == "user"
+        assert appended_messages[1].text == "hello"

python/packages/core/tests/core/test_docstrings.py

@@ -0,0 +1,244 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from agent_framework._docstrings import apply_layered_docstring, build_layered_docstring, insert_docstring_block
+
+# -- Helpers: stub functions with various docstring shapes --
+
+
+def _source_with_full_docstring(x: int) -> int:
+    """Do something useful.
+
+    Args:
+        x: The input value.
+
+    Keyword Args:
+        timeout: Max seconds to wait.
+
+    Returns:
+        The computed result.
+    """
+    return x
+
+
+def _source_with_args_only(x: int) -> int:
+    """Do something useful.
+
+    Args:
+        x: The input value.
+
+    Returns:
+        The computed result.
+    """
+    return x
+
+
+def _source_no_sections() -> None:
+    """A plain summary with no Google-style sections."""
+
+
+def _source_with_attributes() -> None:
+    """A documented object.
+
+    Attributes:
+        value: A documented attribute.
+    """
+
+
+def _source_no_docstring() -> None:
+    pass
+
+
+def _target_stub() -> None:
+    pass
+
+
+# -- build_layered_docstring tests --
+
+
+def test_build_returns_none_when_source_has_no_docstring() -> None:
+    result = build_layered_docstring(_source_no_docstring)
+    assert result is None
+
+
+def test_build_returns_original_when_no_extra_kwargs() -> None:
+    result = build_layered_docstring(_source_with_full_docstring)
+    assert result is not None
+    assert "Do something useful." in result
+    assert "Keyword Args:" in result
+
+
+def test_build_returns_original_when_extra_kwargs_empty() -> None:
+    result = build_layered_docstring(_source_with_full_docstring, extra_keyword_args={})
+    assert result is not None
+    assert result == build_layered_docstring(_source_with_full_docstring)
+
+
+def test_build_appends_to_existing_keyword_args_section() -> None:
+    result = build_layered_docstring(
+        _source_with_full_docstring,
+        extra_keyword_args={"retries": "Number of retries."},
+    )
+    assert result is not None
+    assert "timeout: Max seconds to wait." in result
+    assert "retries: Number of retries." in result
+    # Both should be under Keyword Args
+    lines = result.splitlines()
+    kw_index = next(i for i, line in enumerate(lines) if line == "Keyword Args:")
+    ret_index = next(i for i, line in enumerate(lines) if line == "Returns:")
+    retries_index = next(i for i, line in enumerate(lines) if "retries:" in line)
+    assert kw_index < retries_index < ret_index
+
+
+def test_build_inserts_keyword_args_after_args_section() -> None:
+    result = build_layered_docstring(
+        _source_with_args_only,
+        extra_keyword_args={"verbose": "Enable verbose output."},
+    )
+    assert result is not None
+    assert "Keyword Args:" in result
+    assert "verbose: Enable verbose output." in result
+    lines = result.splitlines()
+    args_index = next(i for i, line in enumerate(lines) if line == "Args:")
+    kw_index = next(i for i, line in enumerate(lines) if line == "Keyword Args:")
+    ret_index = next(i for i, line in enumerate(lines) if line == "Returns:")
+    assert args_index < kw_index < ret_index
+
+
+def test_build_inserts_keyword_args_in_docstring_with_no_sections() -> None:
+    result = build_layered_docstring(
+        _source_no_sections,
+        extra_keyword_args={"debug": "Enable debug mode."},
+    )
+    assert result is not None
+    assert "A plain summary" in result
+    assert "Keyword Args:" in result
+    assert "debug: Enable debug mode." in result
+
+
+def test_build_handles_multiline_descriptions() -> None:
+    result = build_layered_docstring(
+        _source_with_args_only,
+        extra_keyword_args={
+            "config": "The configuration object.\nMust be a valid mapping.\nDefaults to empty.",
+        },
+    )
+    assert result is not None
+    lines = result.splitlines()
+    config_line = next(line for line in lines if "config:" in line)
+    assert "The configuration object." in config_line
+    # Continuation lines should be indented
+    config_idx = lines.index(config_line)
+    assert "Must be a valid mapping." in lines[config_idx + 1]
+    assert "Defaults to empty." in lines[config_idx + 2]
+
+
+def test_build_preserves_multiple_extra_kwargs_order() -> None:
+    result = build_layered_docstring(
+        _source_with_args_only,
+        extra_keyword_args={
+            "alpha": "First.",
+            "beta": "Second.",
+            "gamma": "Third.",
+        },
+    )
+    assert result is not None
+    lines = result.splitlines()
+    alpha_idx = next(i for i, line in enumerate(lines) if "alpha:" in line)
+    beta_idx = next(i for i, line in enumerate(lines) if "beta:" in line)
+    gamma_idx = next(i for i, line in enumerate(lines) if "gamma:" in line)
+    assert alpha_idx < beta_idx < gamma_idx
+
+
+# -- insert_docstring_block tests --
+
+
+def test_insert_docstring_block_before_args_section() -> None:
+    result = insert_docstring_block(
+        _source_with_args_only.__doc__,
+        block="""\
+        .. warning:: Experimental
+
+            This API is experimental.
+        """,
+    )
+    assert result is not None
+    lines = result.splitlines()
+    warning_index = next(i for i, line in enumerate(lines) if line == ".. warning:: Experimental")
+    args_index = next(i for i, line in enumerate(lines) if line == "Args:")
+    assert warning_index < args_index
+
+
+def test_insert_docstring_block_before_attributes_section() -> None:
+    result = insert_docstring_block(
+        _source_with_attributes.__doc__,
+        block="""\
+        .. warning:: Experimental
+
+            This API is experimental.
+        """,
+    )
+    assert result is not None
+    lines = result.splitlines()
+    warning_index = next(i for i, line in enumerate(lines) if line == ".. warning:: Experimental")
+    attributes_index = next(i for i, line in enumerate(lines) if line == "Attributes:")
+    assert warning_index < attributes_index
+
+
+def test_insert_docstring_block_appends_when_no_sections() -> None:
+    result = insert_docstring_block(
+        _source_no_sections.__doc__,
+        block="""\
+        .. note:: Release candidate
+
+            This API is nearly final.
+        """,
+    )
+    assert result is not None
+    assert result.endswith("This API is nearly final.")
+    assert ".. note:: Release candidate" in result
+
+
+def test_insert_docstring_block_returns_block_for_missing_docstring() -> None:
+    result = insert_docstring_block(
+        _source_no_docstring.__doc__,
+        block="""\
+        .. warning:: Experimental
+
+            This API is experimental.
+        """,
+    )
+    assert result == ".. warning:: Experimental\n\n    This API is experimental."
+
+
+# -- apply_layered_docstring tests --
+
+
+def test_apply_sets_docstring_on_target() -> None:
+    def target() -> None:
+        pass
+
+    apply_layered_docstring(target, _source_with_full_docstring)
+    assert target.__doc__ is not None
+    assert "Do something useful." in target.__doc__
+
+
+def test_apply_with_extra_kwargs() -> None:
+    def target() -> None:
+        pass
+
+    apply_layered_docstring(
+        target,
+        _source_with_args_only,
+        extra_keyword_args={"flag": "A boolean flag."},
+    )
+    assert target.__doc__ is not None
+    assert "flag: A boolean flag." in target.__doc__
+    assert "Keyword Args:" in target.__doc__
+
+
+def test_apply_sets_none_when_source_has_no_docstring() -> None:
+    def target() -> None:
+        """Original."""
+
+    apply_layered_docstring(target, _source_no_docstring)
+    assert target.__doc__ is None

python/packages/core/tests/core/test_embedding_client.py

@@ -0,0 +1,104 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import pytest
+
+from agent_framework import (
+    BaseEmbeddingClient,
+    Embedding,
+    EmbeddingGenerationOptions,
+    GeneratedEmbeddings,
+    SupportsGetEmbeddings,
+)
+
+
+class MockEmbeddingClient(BaseEmbeddingClient):
+    """A simple mock embedding client for testing."""
+
+    async def get_embeddings(
+        self,
+        values: Sequence[str],
+        *,
+        options: EmbeddingGenerationOptions | None = None,
+    ) -> GeneratedEmbeddings[list[float]]:
+        return GeneratedEmbeddings(
+            [Embedding(vector=[0.1, 0.2, 0.3], model="mock-model") for _ in values],
+            usage={"prompt_tokens": len(values), "total_tokens": len(values)},  # type: ignore[arg-type]
+        )
+
+
+# --- BaseEmbeddingClient tests ---
+
+
+async def test_base_get_embeddings() -> None:
+    client = MockEmbeddingClient()
+    result = await client.get_embeddings(["hello", "world"])
+    assert len(result) == 2
+    assert result[0].vector == [0.1, 0.2, 0.3]
+    assert result[0].model == "mock-model"
+
+
+async def test_base_get_embeddings_with_options() -> None:
+    client = MockEmbeddingClient()
+    options: EmbeddingGenerationOptions = {"model": "test", "dimensions": 3}
+    result = await client.get_embeddings(["hello"], options=options)
+    assert len(result) == 1
+
+
+async def test_base_get_embeddings_usage() -> None:
+    client = MockEmbeddingClient()
+    result = await client.get_embeddings(["a", "b", "c"])
+    assert result.usage is not None
+    assert result.usage["prompt_tokens"] == 3  # type: ignore[typeddict-item]
+
+
+def test_base_additional_properties_default() -> None:
+    client = MockEmbeddingClient()
+    assert client.additional_properties == {}
+
+
+def test_base_additional_properties_custom() -> None:
+    client = MockEmbeddingClient(additional_properties={"key": "value"})
+    assert client.additional_properties == {"key": "value"}
+
+
+def test_base_embedding_client_rejects_unknown_kwargs() -> None:
+    with pytest.raises(TypeError):
+        MockEmbeddingClient(legacy_key="value")  # type: ignore[call-arg]  # ty: ignore[unknown-argument]
+
+
+# --- SupportsGetEmbeddings protocol tests ---
+
+
+def test_mock_client_satisfies_protocol() -> None:
+    client = MockEmbeddingClient()
+    assert isinstance(client, SupportsGetEmbeddings)
+
+
+def test_plain_class_satisfies_protocol() -> None:
+    """A plain class with the right signature should satisfy the protocol."""
+
+    class PlainEmbeddingClient:
+        additional_properties: dict = {}
+
+        async def get_embeddings(self, values, *, options=None):
+            return GeneratedEmbeddings()
+
+    client = PlainEmbeddingClient()
+    assert isinstance(client, SupportsGetEmbeddings)
+
+
+def test_wrong_class_does_not_satisfy_protocol() -> None:
+    """A class without get_embeddings should not satisfy the protocol."""
+
+    class NotAnEmbeddingClient:
+        additional_properties: dict = {}
+
+        async def generate(self, values):
+            pass
+
+    client = NotAnEmbeddingClient()
+    assert not isinstance(client, SupportsGetEmbeddings)

python/packages/core/tests/core/test_embedding_types.py

@@ -0,0 +1,182 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from datetime import datetime
+
+from agent_framework import Embedding, EmbeddingGenerationOptions, GeneratedEmbeddings
+
+# --- Embedding tests ---
+
+
+def test_embedding_basic_construction() -> None:
+    embedding = Embedding(vector=[0.1, 0.2, 0.3])
+    assert embedding.vector == [0.1, 0.2, 0.3]
+    assert embedding.model is None
+    assert embedding.created_at is None
+    assert embedding.additional_properties == {}
+
+
+def test_embedding_construction_with_metadata() -> None:
+    now = datetime.now()
+    embedding = Embedding(
+        vector=[0.1, 0.2],
+        model="text-embedding-3-small",
+        created_at=now,
+        additional_properties={"key": "value"},
+    )
+    assert embedding.model == "text-embedding-3-small"
+    assert embedding.created_at == now
+    assert embedding.additional_properties == {"key": "value"}
+
+
+def test_embedding_dimensions_computed_from_list() -> None:
+    embedding = Embedding(vector=[0.1, 0.2, 0.3])
+    assert embedding.dimensions == 3
+
+
+def test_embedding_dimensions_computed_from_tuple() -> None:
+    embedding = Embedding(vector=(0.1, 0.2, 0.3, 0.4))
+    assert embedding.dimensions == 4
+
+
+def test_embedding_dimensions_computed_from_bytes() -> None:
+    embedding = Embedding(vector=b"\x00\x01\x02")
+    assert embedding.dimensions == 3
+
+
+def test_embedding_dimensions_explicit_overrides_computed() -> None:
+    embedding = Embedding(vector=[0.1, 0.2, 0.3], dimensions=1536)
+    assert embedding.dimensions == 1536
+
+
+def test_embedding_dimensions_none_for_unknown_type() -> None:
+    embedding = Embedding(vector="not a list")  # type: ignore[arg-type]
+    assert embedding.dimensions is None
+
+
+def test_embedding_dimensions_explicit_with_unknown_type() -> None:
+    embedding = Embedding(vector="not a list", dimensions=100)  # type: ignore[arg-type]
+    assert embedding.dimensions == 100
+
+
+def test_embedding_empty_vector() -> None:
+    embedding = Embedding(vector=[])  # type: ignore[var-annotated]
+    assert embedding.dimensions == 0
+
+
+def test_embedding_int_vector() -> None:
+    embedding = Embedding(vector=[1, 2, 3])
+    assert embedding.vector == [1, 2, 3]
+    assert embedding.dimensions == 3
+
+
+# --- GeneratedEmbeddings tests ---
+
+
+def test_generated_basic_construction() -> None:
+    embeddings = GeneratedEmbeddings()
+    assert len(embeddings) == 0
+    assert embeddings.options is None
+    assert embeddings.usage is None
+    assert embeddings.additional_properties == {}
+
+
+def test_generated_construction_with_embeddings() -> None:
+    items = [Embedding(vector=[0.1, 0.2]), Embedding(vector=[0.3, 0.4])]
+    embeddings = GeneratedEmbeddings(items)
+    assert len(embeddings) == 2
+    assert embeddings[0].vector == [0.1, 0.2]
+    assert embeddings[1].vector == [0.3, 0.4]
+
+
+def test_generated_construction_with_usage() -> None:
+    usage = {"prompt_tokens": 10, "total_tokens": 10}
+    embeddings = GeneratedEmbeddings(
+        [
+            Embedding(
+                vector=[0.1],
+                model="test-model",
+            )
+        ],
+        usage=usage,  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+    )
+    assert embeddings.usage == usage
+    assert embeddings.usage["prompt_tokens"] == 10  # type: ignore[index, typeddict-item]  # pyrefly: ignore[unsupported-operation]  # ty: ignore[not-subscriptable]
+
+
+def test_generated_construction_with_additional_properties() -> None:
+    embeddings = GeneratedEmbeddings(
+        additional_properties={"model": "test"},
+    )
+    assert embeddings.additional_properties == {"model": "test"}
+
+
+def test_generated_construction_with_options() -> None:
+    opts: EmbeddingGenerationOptions = {"model": "text-embedding-3-small", "dimensions": 256}
+    embeddings = GeneratedEmbeddings(
+        [Embedding(vector=[0.1])],
+        options=opts,
+    )
+    assert embeddings.options is not None
+    assert embeddings.options["model"] == "text-embedding-3-small"
+    assert embeddings.options["dimensions"] == 256
+
+
+def test_generated_list_behavior_iteration() -> None:
+    items = [Embedding(vector=[float(i)]) for i in range(5)]
+    embeddings = GeneratedEmbeddings(items)
+    vectors = [e.vector for e in embeddings]
+    assert vectors == [[0.0], [1.0], [2.0], [3.0], [4.0]]
+
+
+def test_generated_list_behavior_indexing() -> None:
+    items = [Embedding(vector=[0.1]), Embedding(vector=[0.2])]
+    embeddings = GeneratedEmbeddings(items)
+    assert embeddings[0].vector == [0.1]
+    assert embeddings[-1].vector == [0.2]
+
+
+def test_generated_list_behavior_slicing() -> None:
+    items = [Embedding(vector=[float(i)]) for i in range(5)]
+    embeddings = GeneratedEmbeddings(items)
+    sliced = embeddings[1:3]
+    assert len(sliced) == 2
+
+
+def test_generated_list_behavior_append() -> None:
+    embeddings = GeneratedEmbeddings()
+    embeddings.append(Embedding(vector=[0.1]))
+    assert len(embeddings) == 1
+
+
+def test_generated_none_embeddings_creates_empty_list() -> None:
+    embeddings = GeneratedEmbeddings(None)
+    assert len(embeddings) == 0
+
+
+# --- EmbeddingGenerationOptions tests ---
+
+
+def test_options_empty() -> None:
+    options: EmbeddingGenerationOptions = {}
+    assert "model" not in options
+
+
+def test_options_with_model() -> None:
+    options: EmbeddingGenerationOptions = {"model": "text-embedding-3-small"}
+    assert options["model"] == "text-embedding-3-small"
+
+
+def test_options_with_dimensions() -> None:
+    options: EmbeddingGenerationOptions = {"dimensions": 1536}
+    assert options["dimensions"] == 1536
+
+
+def test_options_with_all_fields() -> None:
+    options: EmbeddingGenerationOptions = {
+        "model": "text-embedding-3-small",
+        "dimensions": 1536,
+    }
+    assert options["model"] == "text-embedding-3-small"
+    assert options["dimensions"] == 1536

python/packages/core/tests/core/test_exceptions.py

@@ -0,0 +1,27 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for AgentFrameworkException inner_exception handling."""
+
+from agent_framework import AgentFrameworkException
+
+
+def test_exception_with_inner_exception():
+    """When inner_exception is provided, it should be set as the second arg."""
+    inner = ValueError("inner error")
+    exc = AgentFrameworkException("test message", inner_exception=inner)
+    assert exc.args[0] == "test message"
+    assert exc.args[1] is inner
+
+
+def test_exception_without_inner_exception():
+    """When inner_exception is None, args should only contain the message."""
+    exc = AgentFrameworkException("test message")
+    assert exc.args == ("test message",)
+    assert len(exc.args) == 1
+
+
+def test_exception_inner_exception_none_explicit():
+    """When inner_exception is explicitly None, args should only contain the message."""
+    exc = AgentFrameworkException("test message", inner_exception=None)
+    assert exc.args == ("test message",)
+    assert len(exc.args) == 1

python/packages/core/tests/core/test_feature_stage.py

@@ -0,0 +1,461 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import inspect
+import warnings
+from collections.abc import Generator
+from enum import Enum
+from typing import Protocol, runtime_checkable
+
+import pytest
+
+from agent_framework import ExperimentalFeature as PublicExperimentalFeature
+from agent_framework import ReleaseCandidateFeature as PublicReleaseCandidateFeature
+from agent_framework._feature_stage import (
+    _WARNED_FEATURES,
+    ExperimentalWarning,
+    _feature_stage,
+    experimental,
+    release_candidate,
+)
+from agent_framework._feature_stage import (
+    ExperimentalFeature as InternalExperimentalFeature,
+)
+from agent_framework._feature_stage import (
+    ReleaseCandidateFeature as InternalReleaseCandidateFeature,
+)
+
+
+class AlternateExperimentalFeature(str, Enum):
+    EXPERIMENTAL_FEATURE = "EXPERIMENTAL_FEATURE"
+    SHARED_FEATURE = "SHARED_EXPERIMENTAL_FEATURE"
+    ALTERNATE_FEATURE = "ALTERNATE_EXPERIMENTAL_FEATURE"
+
+
+class InvalidStageFeature(str, Enum):
+    LOWERCASE = "skills"
+
+
+class NonStringFeature(Enum):
+    INTEGER = 1
+
+
+class HelperReleaseCandidateFeature(str, Enum):
+    RC_FEATURE = "RC_FEATURE"
+
+
+@pytest.fixture(autouse=True)
+def clear_feature_warning_state() -> Generator[None]:  # type: ignore[misc]  # pyrefly: ignore[bad-return]
+    _WARNED_FEATURES.clear()
+    yield
+    _WARNED_FEATURES.clear()
+
+
+def test_feature_enums_are_exposed_from_root() -> None:
+    assert PublicExperimentalFeature is InternalExperimentalFeature
+    assert PublicReleaseCandidateFeature is InternalReleaseCandidateFeature
+
+
+def test_experimental_decorator_accepts_feature_enum() -> None:
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+
+        @experimental(feature_id=AlternateExperimentalFeature.EXPERIMENTAL_FEATURE)  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+        def skill_function() -> None:
+            pass
+
+    assert not caught
+
+    with warnings.catch_warnings(record=True) as caught:
+        skill_function()
+
+    assert len(caught) == 1
+    assert f"[{AlternateExperimentalFeature.EXPERIMENTAL_FEATURE.value}]" in str(caught[0].message)
+    assert "skill_function" in str(caught[0].message)
+    assert skill_function.__feature_id__ == AlternateExperimentalFeature.EXPERIMENTAL_FEATURE.value  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+
+def test_experimental_function_warns_on_call_and_not_on_definition() -> None:
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+
+        @experimental(feature_id=AlternateExperimentalFeature.EXPERIMENTAL_FEATURE)  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+        def my_function(value: int) -> int:
+            """Double the input.
+
+            Args:
+                value: Value to double.
+
+            Returns:
+                The doubled value.
+            """
+            return value * 2
+
+    assert not caught
+
+    with warnings.catch_warnings(record=True) as caught:
+        assert my_function(3) == 6
+        assert my_function(4) == 8
+
+    assert len(caught) == 1
+    assert f"[{AlternateExperimentalFeature.EXPERIMENTAL_FEATURE.value}]" in str(caught[0].message)
+    assert "my_function" in str(caught[0].message)
+    assert my_function.__feature_stage__ == "experimental"  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert my_function.__feature_id__ == AlternateExperimentalFeature.EXPERIMENTAL_FEATURE.value  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert my_function.__doc__ is not None
+    lines = my_function.__doc__.splitlines()
+    warning_index = next(i for i, line in enumerate(lines) if line == ".. warning:: Experimental")
+    args_index = next(i for i, line in enumerate(lines) if line == "Args:")
+    assert warning_index < args_index
+
+
+def test_experimental_class_warns_on_instantiation_and_not_on_definition() -> None:
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+
+        @experimental(feature_id=AlternateExperimentalFeature.EXPERIMENTAL_FEATURE)  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+        class ExperimentalClass:
+            """An experimental class.
+
+            Args:
+                value: Value to store.
+            """
+
+            def __init__(self, value: int) -> None:
+                self.value = value
+
+    assert not caught
+
+    with warnings.catch_warnings(record=True) as caught:
+        instantiation_line = inspect.currentframe().f_lineno + 1  # type: ignore[union-attr]  # ty: ignore[unresolved-attribute]
+        instance = ExperimentalClass(4)
+        second_instance = ExperimentalClass(5)
+
+    assert len(caught) == 1
+    assert f"[{AlternateExperimentalFeature.EXPERIMENTAL_FEATURE.value}]" in str(caught[0].message)
+    assert "ExperimentalClass" in str(caught[0].message)
+    assert caught[0].filename == __file__
+    assert caught[0].lineno == instantiation_line
+    assert instance.value == 4
+    assert second_instance.value == 5
+    assert ExperimentalClass.__feature_stage__ == "experimental"  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert ExperimentalClass.__feature_id__ == AlternateExperimentalFeature.EXPERIMENTAL_FEATURE.value  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+
+def test_experimental_abc_subclass_warning_points_at_user_file() -> None:
+    """Subclassing an experimental ABC must report the warning at the user's
+    ``class Sub(...):`` line, not at internal abc.py / <frozen abc> frames.
+
+    Regression: previously the fixed ``stacklevel=3`` landed inside abc.py for
+    ABC-driven class creation, surfacing ``<frozen abc>:106`` to users.
+    """
+    from abc import ABC, abstractmethod
+
+    @experimental(feature_id=AlternateExperimentalFeature.EXPERIMENTAL_FEATURE)  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+    class ExperimentalABC(ABC):
+        @abstractmethod
+        def do(self) -> int: ...
+
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        subclass_line = inspect.currentframe().f_lineno + 1  # type: ignore[union-attr]  # ty: ignore[unresolved-attribute]
+
+        class Concrete(ExperimentalABC):
+            def do(self) -> int:
+                return 1
+
+    assert len(caught) == 1
+    assert caught[0].filename == __file__
+    # __init_subclass__ fires at the end of the class body, so the lineno
+    # points somewhere inside the Concrete class definition rather than at
+    # the ``class Concrete`` header itself. The key behaviour we want to
+    # guarantee is that it is in the *user* file at all (not abc.py).
+    assert subclass_line <= caught[0].lineno <= subclass_line + 5
+    assert issubclass(caught[0].category, ExperimentalWarning)
+    assert Concrete().do() == 1
+
+
+def test_experimental_runtime_checkable_protocol_keeps_protocol_runtime_checks() -> None:
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+
+        @runtime_checkable
+        @experimental(feature_id=AlternateExperimentalFeature.EXPERIMENTAL_FEATURE)  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+        class ExampleProtocol(Protocol):
+            """A protocol used for runtime checks.
+
+            Returns:
+                Nothing.
+            """
+
+            def __call__(self, value: int) -> int: ...
+
+    assert not caught
+
+    def implementation(value: int) -> int:
+        return value
+
+    assert isinstance(implementation, ExampleProtocol)
+    assert ExampleProtocol.__doc__ is not None
+    assert ".. warning:: Experimental" in ExampleProtocol.__doc__
+    assert getattr(ExampleProtocol, "__feature_stage__", None) is None
+    assert getattr(ExampleProtocol, "__feature_id__", None) is None
+
+
+def test_experimental_warning_is_emitted_once_per_feature() -> None:
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+
+        @experimental(feature_id=AlternateExperimentalFeature.SHARED_FEATURE)  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+        def first() -> None:
+            pass
+
+        @experimental(feature_id=AlternateExperimentalFeature.SHARED_FEATURE)  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+        class Second:
+            pass
+
+    assert not caught
+
+    with warnings.catch_warnings(record=True) as caught:
+        first()
+        Second()
+
+    assert first is not None
+    assert Second is not None
+    assert len(caught) == 1
+    assert f"[{AlternateExperimentalFeature.SHARED_FEATURE.value}]" in str(caught[0].message)
+    assert "first" in str(caught[0].message)
+
+
+def test_release_candidate_internal_helper_adds_metadata_without_runtime_warning() -> None:
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+
+        @_feature_stage(
+            stage="release_candidate",
+            feature_id=HelperReleaseCandidateFeature.RC_FEATURE,
+            docstring_block="""\
+.. note:: Release candidate
+
+    This API is in release-candidate stage and may receive
+    minor refinements before it is considered generally available.
+""",
+            warning_category=None,
+        )
+        class ReleaseCandidateClass:
+            """A release-candidate class.
+
+            Args:
+                value: Value to store.
+            """
+
+            def __init__(self, value: int) -> None:
+                self.value = value
+
+    assert not caught
+
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        instance = ReleaseCandidateClass(5)
+
+    assert instance.value == 5
+    assert not caught
+    assert ReleaseCandidateClass.__feature_stage__ == "release_candidate"  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert ReleaseCandidateClass.__feature_id__ == HelperReleaseCandidateFeature.RC_FEATURE.value  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert ReleaseCandidateClass.__doc__ is not None
+    assert ".. note:: Release candidate" in ReleaseCandidateClass.__doc__
+
+
+def test_experimental_property_warns_on_access_and_not_on_definition() -> None:
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+
+        class Example:
+            @property
+            @experimental(feature_id=AlternateExperimentalFeature.EXPERIMENTAL_FEATURE)  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+            def value(self) -> int:
+                """Return the value.
+
+                Returns:
+                    The stored value.
+                """
+                return 1
+
+    assert not caught
+
+    with warnings.catch_warnings(record=True) as caught:
+        assert Example().value == 1
+        assert Example().value == 1
+
+    assert len(caught) == 1
+    assert f"[{AlternateExperimentalFeature.EXPERIMENTAL_FEATURE.value}]" in str(caught[0].message)
+    assert "Example.value" in str(caught[0].message)
+    assert Example.value.__doc__ is not None
+    lines = Example.value.__doc__.splitlines()
+    warning_index = next(i for i, line in enumerate(lines) if line == ".. warning:: Experimental")
+    returns_index = next(i for i, line in enumerate(lines) if line == "Returns:")
+    assert warning_index < returns_index
+
+
+def test_experimental_staticmethod_warns_when_decorator_wraps_descriptor() -> None:
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+
+        class Example:
+            @experimental(feature_id=AlternateExperimentalFeature.EXPERIMENTAL_FEATURE)  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+            @staticmethod
+            def value() -> int:
+                """Return the value.
+
+                Returns:
+                    The stored value.
+                """
+                return 1
+
+    assert not caught
+
+    with warnings.catch_warnings(record=True) as caught:
+        assert Example.value() == 1
+        assert Example.value() == 1
+
+    assert len(caught) == 1
+    assert f"[{AlternateExperimentalFeature.EXPERIMENTAL_FEATURE.value}]" in str(caught[0].message)
+    assert "Example.value" in str(caught[0].message)
+    assert Example.value.__feature_id__ == AlternateExperimentalFeature.EXPERIMENTAL_FEATURE.value  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert Example.value.__doc__ is not None
+    lines = Example.value.__doc__.splitlines()
+    warning_index = next(i for i, line in enumerate(lines) if line == ".. warning:: Experimental")
+    returns_index = next(i for i, line in enumerate(lines) if line == "Returns:")
+    assert warning_index < returns_index
+
+
+def test_experimental_classmethod_warns_when_decorator_wraps_descriptor() -> None:
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+
+        class Example:
+            @experimental(feature_id=AlternateExperimentalFeature.EXPERIMENTAL_FEATURE)  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+            @classmethod
+            def value(cls) -> int:
+                """Return the value.
+
+                Returns:
+                    The stored value.
+                """
+                return 1
+
+    assert not caught
+
+    with warnings.catch_warnings(record=True) as caught:
+        assert Example.value() == 1
+        assert Example.value() == 1
+
+    assert len(caught) == 1
+    assert f"[{AlternateExperimentalFeature.EXPERIMENTAL_FEATURE.value}]" in str(caught[0].message)
+    assert "Example.value" in str(caught[0].message)
+    assert Example.value.__func__.__feature_id__ == AlternateExperimentalFeature.EXPERIMENTAL_FEATURE.value  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert Example.value.__doc__ is not None
+    lines = Example.value.__doc__.splitlines()
+    warning_index = next(i for i, line in enumerate(lines) if line == ".. warning:: Experimental")
+    returns_index = next(i for i, line in enumerate(lines) if line == "Returns:")
+    assert warning_index < returns_index
+
+
+def test_feature_id_allows_lowercase_values() -> None:
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+
+        @_feature_stage(
+            stage="experimental",
+            feature_id=InvalidStageFeature.LOWERCASE,
+            docstring_block=".. warning:: Experimental",
+            warning_category=ExperimentalWarning,
+        )
+        def lowercase_feature() -> None:
+            pass
+
+    assert not caught
+
+    with warnings.catch_warnings(record=True) as caught:
+        lowercase_feature()
+
+    assert len(caught) == 1
+    assert "[skills]" in str(caught[0].message)
+    assert "lowercase_feature" in str(caught[0].message)
+    assert lowercase_feature.__feature_id__ == "skills"  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+
+def test_experimental_decorator_allows_string_feature_id_at_runtime() -> None:
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+
+        @experimental(feature_id="STRING_FEATURE")  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+        def skill_function() -> None:
+            pass
+
+    assert not caught
+
+    with warnings.catch_warnings(record=True) as caught:
+        skill_function()
+
+    assert len(caught) == 1
+    assert "[STRING_FEATURE]" in str(caught[0].message)
+    assert "skill_function" in str(caught[0].message)
+    assert skill_function.__feature_id__ == "STRING_FEATURE"  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+
+def test_experimental_decorator_allows_other_enum_values_at_runtime() -> None:
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+
+        @experimental(feature_id=AlternateExperimentalFeature.ALTERNATE_FEATURE)  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+        def my_function() -> None:
+            pass
+
+    assert not caught
+
+    with warnings.catch_warnings(record=True) as caught:
+        my_function()
+
+    assert len(caught) == 1
+    assert f"[{AlternateExperimentalFeature.ALTERNATE_FEATURE.value}]" in str(caught[0].message)
+    assert "my_function" in str(caught[0].message)
+    assert my_function.__feature_id__ == AlternateExperimentalFeature.ALTERNATE_FEATURE.value  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+
+def test_release_candidate_decorator_allows_string_feature_id_at_runtime() -> None:
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+
+        @release_candidate(feature_id="RC_FEATURE")  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+        class ReleaseCandidateClass:
+            """A release-candidate class."""
+
+    assert not caught
+    assert ReleaseCandidateClass.__feature_stage__ == "release_candidate"  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert ReleaseCandidateClass.__feature_id__ == "RC_FEATURE"  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+
+def test_feature_id_stringifies_non_string_enum_values() -> None:
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+
+        @_feature_stage(
+            stage="experimental",
+            feature_id=NonStringFeature.INTEGER,
+            docstring_block=".. warning:: Experimental",
+            warning_category=ExperimentalWarning,
+        )
+        def numeric_feature() -> None:
+            pass
+
+    assert not caught
+
+    with warnings.catch_warnings(record=True) as caught:
+        numeric_feature()
+
+    assert len(caught) == 1
+    assert "[1]" in str(caught[0].message)
+    assert "numeric_feature" in str(caught[0].message)
+    assert numeric_feature.__feature_id__ == "1"  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]

python/packages/core/tests/core/test_foundry_namespace.py

@@ -0,0 +1,24 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import pytest
+from agent_framework_foundry import FoundryChatClient, FoundryMemoryProvider
+from agent_framework_foundry_local import FoundryLocalClient
+
+import agent_framework.azure as azure
+import agent_framework.foundry as foundry
+
+
+def test_foundry_namespace_exposes_cloud_and_local_symbols() -> None:
+    assert foundry.FoundryChatClient is FoundryChatClient
+    assert foundry.FoundryMemoryProvider is FoundryMemoryProvider
+    assert foundry.FoundryLocalClient is FoundryLocalClient
+    assert "FoundryChatClient" in dir(foundry)
+    assert "FoundryLocalClient" in dir(foundry)
+
+
+def test_azure_namespace_no_longer_exposes_foundry_symbols() -> None:
+    assert "FoundryChatClient" not in dir(azure)
+    assert "FoundryLocalClient" not in dir(azure)
+
+    with pytest.raises(AttributeError, match="Module `azure` has no attribute FoundryChatClient\\."):
+        _ = azure.FoundryChatClient  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]

python/packages/core/tests/core/test_harness_background_agents.py

@@ -0,0 +1,544 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+import pytest
+
+from agent_framework import (
+    AgentResponse,
+    AgentSession,
+    BackgroundAgentsProvider,
+    BackgroundTaskInfo,
+    BackgroundTaskStatus,
+    Message,
+    ServiceSessionId,
+)
+from agent_framework._sessions import SessionContext
+
+# Suppress "coroutine was never awaited" warnings from task cancellation in tests.
+# This occurs when cancelling tasks that wrap coroutines through _run_agent().
+pytestmark = pytest.mark.filterwarnings("ignore::RuntimeWarning:asyncio")
+
+# --- Test Helpers ---
+
+
+class _FakeAgent:
+    """Minimal agent stub for testing background agent delegation."""
+
+    def __init__(
+        self,
+        name: str,
+        description: str | None = None,
+        *,
+        response_text: str = "done",
+        delay: float = 0.0,
+        should_fail: bool = False,
+    ):
+        self.id = f"agent-{name}"
+        self.name = name
+        self.description = description
+        self._response_text = response_text
+        self._delay = delay
+        self._should_fail = should_fail
+
+    def create_session(self, *, session_id: str | None = None) -> AgentSession:
+        return AgentSession(session_id=session_id)
+
+    def get_session(
+        self,
+        service_session_id: str | ServiceSessionId,
+        *,
+        session_id: str | None = None,
+    ) -> AgentSession:
+        return AgentSession(service_session_id=service_session_id, session_id=session_id)
+
+    async def run(
+        self, messages: Any = None, *, stream: bool = False, session: Any = None, **kwargs: Any
+    ) -> AgentResponse[Any]:
+        if self._delay > 0:
+            await asyncio.sleep(self._delay)
+        if self._should_fail:
+            raise RuntimeError("Agent execution failed")
+        return AgentResponse(messages=[Message(role="assistant", contents=[self._response_text])])
+
+
+def _make_provider(*agents: _FakeAgent) -> BackgroundAgentsProvider:
+    """Create a provider with given agents."""
+    return BackgroundAgentsProvider(agents)  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+
+
+def _make_session() -> AgentSession:
+    """Create a session for testing."""
+    return AgentSession()
+
+
+async def _get_tools(provider: BackgroundAgentsProvider, session: AgentSession) -> dict[str, Any]:
+    """Run before_run and return tools by name."""
+    context = SessionContext(input_messages=[])
+    await provider.before_run(agent=None, session=session, context=context, state={})
+    tools_by_name: dict[str, Any] = {}
+    for t in context.tools:
+        tools_by_name[t.name if hasattr(t, "name") else str(t)] = t
+    return tools_by_name
+
+
+async def _invoke_tool(tool_obj: Any, **kwargs: Any) -> str:
+    """Invoke a FunctionTool and return the raw result string."""
+    return await tool_obj.invoke(arguments=kwargs, skip_parsing=True)
+
+
+# --- Constructor Tests ---
+
+
+def test_constructor_requires_at_least_one_agent() -> None:
+    """Should reject empty agent list."""
+    with pytest.raises(ValueError, match="At least one background agent"):
+        BackgroundAgentsProvider([])
+
+
+def test_constructor_requires_agent_names() -> None:
+    """Should reject agents with no name."""
+    agent = _FakeAgent("")
+    with pytest.raises(ValueError, match="non-empty name"):
+        BackgroundAgentsProvider([agent])  # type: ignore[list-item]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+
+
+def test_constructor_rejects_duplicate_names() -> None:
+    """Should reject duplicate agent names (case-insensitive)."""
+    agent1 = _FakeAgent("Research")
+    agent2 = _FakeAgent("research")
+    with pytest.raises(ValueError, match="Duplicate background agent name"):
+        BackgroundAgentsProvider([agent1, agent2])  # type: ignore[list-item]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+
+
+def test_constructor_valid_agents() -> None:
+    """Should succeed with valid unique agents."""
+    provider = BackgroundAgentsProvider([_FakeAgent("Alpha"), _FakeAgent("Beta")])  # type: ignore[list-item]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+    assert provider.source_id == "background_agents"
+
+
+def test_constructor_custom_source_id() -> None:
+    """Should accept custom source_id."""
+    provider = BackgroundAgentsProvider([_FakeAgent("Agent1")], source_id="custom_bg")  # type: ignore[list-item]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+    assert provider.source_id == "custom_bg"
+
+
+# --- Tool Injection Tests ---
+
+
+async def test_before_run_injects_six_tools() -> None:
+    """before_run should inject exactly 6 tools."""
+    provider = _make_provider(_FakeAgent("Worker"))
+    tools = await _get_tools(provider, _make_session())
+    assert len(tools) == 6
+    expected_names = {
+        "background_agents_start_task",
+        "background_agents_wait_for_first_completion",
+        "background_agents_get_task_results",
+        "background_agents_get_all_tasks",
+        "background_agents_continue_task",
+        "background_agents_clear_completed_task",
+    }
+    assert set(tools.keys()) == expected_names
+
+
+async def test_before_run_injects_instructions() -> None:
+    """before_run should inject instructions mentioning agent names."""
+    provider = _make_provider(_FakeAgent("ResearchBot", "Does research"))
+    context = SessionContext(input_messages=[])
+    session = _make_session()
+    await provider.before_run(agent=None, session=session, context=context, state={})
+    all_instructions = " ".join(context.instructions)
+    assert "ResearchBot" in all_instructions
+    assert "Does research" in all_instructions
+
+
+# --- Start Task Tests ---
+
+
+async def test_start_task_success() -> None:
+    """Should start a task and return confirmation."""
+    provider = _make_provider(_FakeAgent("Worker", response_text="result"))
+    session = _make_session()
+    tools = await _get_tools(provider, session)
+
+    result = await _invoke_tool(
+        tools["background_agents_start_task"],
+        agent_name="Worker",
+        input="do something",
+        description="test task",
+    )
+    assert "task 1 started" in result.lower()
+    assert "Worker" in result
+
+
+async def test_start_task_unknown_agent() -> None:
+    """Should return error for unknown agent name."""
+    provider = _make_provider(_FakeAgent("Worker"))
+    session = _make_session()
+    tools = await _get_tools(provider, session)
+
+    result = await _invoke_tool(
+        tools["background_agents_start_task"],
+        agent_name="NonExistent",
+        input="do something",
+        description="test",
+    )
+    assert "Error" in result
+    assert "NonExistent" in result
+
+
+async def test_start_task_increments_ids() -> None:
+    """Task IDs should increment sequentially."""
+    provider = _make_provider(_FakeAgent("Worker"))
+    session = _make_session()
+    tools = await _get_tools(provider, session)
+
+    r1 = await _invoke_tool(
+        tools["background_agents_start_task"],
+        agent_name="Worker",
+        input="task 1",
+        description="first",
+    )
+    r2 = await _invoke_tool(
+        tools["background_agents_start_task"],
+        agent_name="Worker",
+        input="task 2",
+        description="second",
+    )
+    assert "task 1 started" in r1.lower()
+    assert "task 2 started" in r2.lower()
+
+
+# --- Get All Tasks Tests ---
+
+
+async def test_get_all_tasks_empty() -> None:
+    """Should return 'No tasks.' when no tasks exist."""
+    provider = _make_provider(_FakeAgent("Worker"))
+    session = _make_session()
+    tools = await _get_tools(provider, session)
+
+    result = await _invoke_tool(tools["background_agents_get_all_tasks"])
+    assert "No tasks" in result
+
+
+async def test_get_all_tasks_shows_tasks() -> None:
+    """Should list all tasks with status and description."""
+    provider = _make_provider(_FakeAgent("Worker"))
+    session = _make_session()
+    tools = await _get_tools(provider, session)
+
+    await _invoke_tool(
+        tools["background_agents_start_task"],
+        agent_name="Worker",
+        input="hello",
+        description="my task",
+    )
+    result = await _invoke_tool(tools["background_agents_get_all_tasks"])
+    assert "my task" in result
+    assert "Worker" in result
+
+
+# --- Wait for Completion Tests ---
+
+
+async def test_wait_for_first_completion() -> None:
+    """Should wait and return when a task completes."""
+    provider = _make_provider(_FakeAgent("Fast", response_text="fast result", delay=0.01))
+    session = _make_session()
+    tools = await _get_tools(provider, session)
+
+    await _invoke_tool(
+        tools["background_agents_start_task"],
+        agent_name="Fast",
+        input="go",
+        description="fast task",
+    )
+    result = await _invoke_tool(
+        tools["background_agents_wait_for_first_completion"],
+        task_ids=[1],
+    )
+    assert "finished" in result.lower()
+    assert "completed" in result.lower()
+
+
+async def test_wait_empty_task_ids() -> None:
+    """Should return error for empty task_ids."""
+    provider = _make_provider(_FakeAgent("Worker"))
+    session = _make_session()
+    tools = await _get_tools(provider, session)
+
+    result = await _invoke_tool(
+        tools["background_agents_wait_for_first_completion"],
+        task_ids=[],
+    )
+    assert "Error" in result
+
+
+async def test_wait_no_running_tasks() -> None:
+    """Should return error when no specified tasks are running."""
+    provider = _make_provider(_FakeAgent("Worker"))
+    session = _make_session()
+    tools = await _get_tools(provider, session)
+
+    result = await _invoke_tool(
+        tools["background_agents_wait_for_first_completion"],
+        task_ids=[999],
+    )
+    assert "Error" in result or "not running" in result.lower()
+
+
+# --- Get Task Results Tests ---
+
+
+async def test_get_task_results_completed() -> None:
+    """Should return result text for completed task."""
+    provider = _make_provider(_FakeAgent("Worker", response_text="the answer", delay=0.01))
+    session = _make_session()
+    tools = await _get_tools(provider, session)
+
+    await _invoke_tool(
+        tools["background_agents_start_task"],
+        agent_name="Worker",
+        input="query",
+        description="test",
+    )
+    # Wait for completion.
+    await _invoke_tool(
+        tools["background_agents_wait_for_first_completion"],
+        task_ids=[1],
+    )
+    result = await _invoke_tool(
+        tools["background_agents_get_task_results"],
+        task_id=1,
+    )
+    assert result == "the answer"
+
+
+async def test_get_task_results_running() -> None:
+    """Should indicate task is still running."""
+    provider = _make_provider(_FakeAgent("Slow", delay=10.0))
+    session = _make_session()
+    tools = await _get_tools(provider, session)
+
+    await _invoke_tool(
+        tools["background_agents_start_task"],
+        agent_name="Slow",
+        input="query",
+        description="slow task",
+    )
+    try:
+        result = await _invoke_tool(
+            tools["background_agents_get_task_results"],
+            task_id=1,
+        )
+        assert "still running" in result.lower()
+    finally:
+        runtime = provider._get_runtime(session)
+        for task in list(runtime.in_flight_tasks.values()):
+            task.cancel()
+        await asyncio.gather(*runtime.in_flight_tasks.values(), return_exceptions=True)
+
+
+async def test_get_task_results_failed() -> None:
+    """Should return error text for failed task."""
+    provider = _make_provider(_FakeAgent("Broken", should_fail=True, delay=0.01))
+    session = _make_session()
+    tools = await _get_tools(provider, session)
+
+    await _invoke_tool(
+        tools["background_agents_start_task"],
+        agent_name="Broken",
+        input="query",
+        description="will fail",
+    )
+    await _invoke_tool(
+        tools["background_agents_wait_for_first_completion"],
+        task_ids=[1],
+    )
+    result = await _invoke_tool(
+        tools["background_agents_get_task_results"],
+        task_id=1,
+    )
+    assert "failed" in result.lower()
+
+
+async def test_get_task_results_not_found() -> None:
+    """Should return error for non-existent task."""
+    provider = _make_provider(_FakeAgent("Worker"))
+    session = _make_session()
+    tools = await _get_tools(provider, session)
+
+    result = await _invoke_tool(
+        tools["background_agents_get_task_results"],
+        task_id=999,
+    )
+    assert "Error" in result
+
+
+# --- Continue Task Tests ---
+
+
+async def test_continue_task_after_completion() -> None:
+    """Should be able to continue a completed task."""
+    provider = _make_provider(_FakeAgent("Worker", response_text="first result", delay=0.01))
+    session = _make_session()
+    tools = await _get_tools(provider, session)
+
+    await _invoke_tool(
+        tools["background_agents_start_task"],
+        agent_name="Worker",
+        input="first input",
+        description="continuable",
+    )
+    await _invoke_tool(
+        tools["background_agents_wait_for_first_completion"],
+        task_ids=[1],
+    )
+    result = await _invoke_tool(
+        tools["background_agents_continue_task"],
+        task_id=1,
+        text="follow up",
+    )
+    assert "continued" in result.lower()
+
+
+async def test_continue_task_still_running() -> None:
+    """Should return error if task is still running."""
+    provider = _make_provider(_FakeAgent("Slow", delay=10.0))
+    session = _make_session()
+    tools = await _get_tools(provider, session)
+
+    await _invoke_tool(
+        tools["background_agents_start_task"],
+        agent_name="Slow",
+        input="input",
+        description="running",
+    )
+    try:
+        result = await _invoke_tool(
+            tools["background_agents_continue_task"],
+            task_id=1,
+            text="follow up",
+        )
+        assert "still running" in result.lower()
+    finally:
+        runtime = provider._get_runtime(session)
+        for task in list(runtime.in_flight_tasks.values()):
+            task.cancel()
+        await asyncio.gather(*runtime.in_flight_tasks.values(), return_exceptions=True)
+
+
+async def test_continue_task_not_found() -> None:
+    """Should return error for non-existent task."""
+    provider = _make_provider(_FakeAgent("Worker"))
+    session = _make_session()
+    tools = await _get_tools(provider, session)
+
+    result = await _invoke_tool(
+        tools["background_agents_continue_task"],
+        task_id=999,
+        text="hello",
+    )
+    assert "Error" in result
+
+
+# --- Clear Task Tests ---
+
+
+async def test_clear_completed_task() -> None:
+    """Should clear a completed task."""
+    provider = _make_provider(_FakeAgent("Worker", response_text="done", delay=0.01))
+    session = _make_session()
+    tools = await _get_tools(provider, session)
+
+    await _invoke_tool(
+        tools["background_agents_start_task"],
+        agent_name="Worker",
+        input="task",
+        description="clearable",
+    )
+    await _invoke_tool(
+        tools["background_agents_wait_for_first_completion"],
+        task_ids=[1],
+    )
+    result = await _invoke_tool(
+        tools["background_agents_clear_completed_task"],
+        task_id=1,
+    )
+    assert "cleared" in result.lower()
+
+    # Verify task is gone.
+    all_tasks = await _invoke_tool(tools["background_agents_get_all_tasks"])
+    assert "No tasks" in all_tasks
+
+
+async def test_clear_running_task_error() -> None:
+    """Should return error when clearing a running task."""
+    provider = _make_provider(_FakeAgent("Slow", delay=10.0))
+    session = _make_session()
+    tools = await _get_tools(provider, session)
+
+    await _invoke_tool(
+        tools["background_agents_start_task"],
+        agent_name="Slow",
+        input="task",
+        description="still going",
+    )
+    try:
+        result = await _invoke_tool(
+            tools["background_agents_clear_completed_task"],
+            task_id=1,
+        )
+        assert "still running" in result.lower()
+    finally:
+        runtime = provider._get_runtime(session)
+        for task in list(runtime.in_flight_tasks.values()):
+            task.cancel()
+        await asyncio.gather(*runtime.in_flight_tasks.values(), return_exceptions=True)
+
+
+async def test_clear_not_found() -> None:
+    """Should return error for non-existent task."""
+    provider = _make_provider(_FakeAgent("Worker"))
+    session = _make_session()
+    tools = await _get_tools(provider, session)
+
+    result = await _invoke_tool(
+        tools["background_agents_clear_completed_task"],
+        task_id=999,
+    )
+    assert "Error" in result
+
+
+# --- BackgroundTaskInfo Tests ---
+
+
+def test_task_info_serialization() -> None:
+    """BackgroundTaskInfo should round-trip through to_dict/from_dict."""
+    info = BackgroundTaskInfo(
+        id=1,
+        agent_name="Worker",
+        description="test task",
+        status=BackgroundTaskStatus.COMPLETED,
+        result_text="hello",
+    )
+    data = info.to_dict()
+    restored = BackgroundTaskInfo.from_dict(data)
+    assert restored.id == 1
+    assert restored.agent_name == "Worker"
+    assert restored.status == BackgroundTaskStatus.COMPLETED
+    assert restored.result_text == "hello"
+    assert restored.error_text is None
+
+
+def test_task_status_enum_values() -> None:
+    """BackgroundTaskStatus should have expected values."""
+    assert BackgroundTaskStatus.RUNNING == "running"
+    assert BackgroundTaskStatus.COMPLETED == "completed"
+    assert BackgroundTaskStatus.FAILED == "failed"
+    assert BackgroundTaskStatus.LOST == "lost"

python/packages/core/tests/core/test_harness_file_memory.py

@@ -0,0 +1,479 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import json
+import re
+
+import pytest
+
+from agent_framework import (
+    AgentSession,
+    Content,
+    FileMemoryProvider,
+    FunctionTool,
+    InMemoryAgentFileStore,
+)
+from agent_framework._harness._file_memory import (
+    _MAX_INDEX_ENTRIES,
+    _MEMORY_INDEX_FILE_NAME,
+    DEFAULT_FILE_MEMORY_INSTRUCTIONS,
+    DEFAULT_FILE_MEMORY_SOURCE_ID,
+    _combine_paths,
+    _description_file_name,
+    _is_internal_file,
+)
+from agent_framework._sessions import SessionContext
+
+
+def _tool_by_name(tools: list[object], name: str) -> object:
+    """Return the tool with the requested name from a prepared tool list."""
+    for tool in tools:
+        if getattr(tool, "name", None) == name:
+            return tool
+    raise AssertionError(f"Tool {name!r} was not found.")
+
+
+def _text(result: list[Content]) -> str:
+    """Return the first content item's text (memory tools always emit text)."""
+    return result[0].text or ""
+
+
+async def _prepare(
+    provider: FileMemoryProvider, *, session_id: str = "session-1"
+) -> tuple[SessionContext, dict[str, FunctionTool]]:
+    """Run ``before_run`` against a fresh session context and return tools by name."""
+    session = AgentSession(session_id=session_id)
+    context = SessionContext(session_id=session_id, input_messages=[])
+    await provider.before_run(agent=None, session=session, context=context, state={})
+    tools: dict[str, FunctionTool] = {tool.name: tool for tool in context.tools}
+    return context, tools
+
+
+def test_description_file_name_replaces_extension() -> None:
+    """The description sidecar replaces a known extension and appends otherwise."""
+    assert _description_file_name("notes.md") == "notes_description.md"
+    assert _description_file_name("data.json") == "data_description.md"
+    assert _description_file_name("noext") == "noext_description.md"
+    # Leading-dot files have no stem, so the suffix is appended.
+    assert _description_file_name(".hidden") == ".hidden_description.md"
+
+
+def test_is_internal_file_detects_sidecars_and_index() -> None:
+    """Internal files are description sidecars and the memory index, case-insensitively."""
+    assert _is_internal_file("notes_description.md")
+    assert _is_internal_file("NOTES_DESCRIPTION.MD")
+    assert _is_internal_file(_MEMORY_INDEX_FILE_NAME)
+    assert _is_internal_file("Memories.md")
+    assert not _is_internal_file("notes.md")
+    assert not _is_internal_file("description.md")
+
+
+def test_combine_paths_joins_with_forward_slash() -> None:
+    """Working-folder paths join with a single forward slash and tolerate empties."""
+    assert _combine_paths("session-1", "notes.md") == "session-1/notes.md"
+    assert _combine_paths("session-1/", "/notes.md") == "session-1/notes.md"
+    assert _combine_paths("", "notes.md") == "notes.md"
+    assert _combine_paths("session-1", "") == "session-1"
+
+
+async def test_provider_registers_tools_and_instructions() -> None:
+    """``before_run`` should register all tools and the default instructions."""
+    provider = FileMemoryProvider(store=InMemoryAgentFileStore())
+    context, tools = await _prepare(provider)
+
+    expected = {
+        "file_memory_write",
+        "file_memory_read",
+        "file_memory_delete",
+        "file_memory_ls",
+        "file_memory_grep",
+        "file_memory_replace",
+        "file_memory_replace_lines",
+    }
+    assert set(tools) >= expected
+    assert all(t.approval_mode == "never_require" for t in context.tools)  # type: ignore[attr-defined]
+    assert any(DEFAULT_FILE_MEMORY_INSTRUCTIONS in chunk for chunk in context.instructions)
+
+
+async def test_provider_uses_default_source_id() -> None:
+    """The default source id should match the public constant."""
+    provider = FileMemoryProvider(store=InMemoryAgentFileStore())
+    assert provider.source_id == DEFAULT_FILE_MEMORY_SOURCE_ID
+
+
+async def test_save_read_delete_round_trip() -> None:
+    """The tools should drive a save/read/list/delete flow with index maintenance."""
+    store = InMemoryAgentFileStore()
+    provider = FileMemoryProvider(store=store)
+    _, tools = await _prepare(provider)
+
+    save = tools["file_memory_write"]
+    read = tools["file_memory_read"]
+    delete = tools["file_memory_delete"]
+    list_files = tools["file_memory_ls"]
+
+    saved = await save.invoke(arguments={"file_name": "plan.md", "content": "step 1"})
+    assert "plan.md" in _text(saved) and "written" in _text(saved)
+
+    read_back = await read.invoke(arguments={"file_name": "plan.md"})
+    assert _text(read_back) == "step 1"
+
+    # Overwrite is allowed (no overwrite flag needed).
+    await save.invoke(arguments={"file_name": "plan.md", "content": "step 2"})
+    assert _text(await read.invoke(arguments={"file_name": "plan.md"})) == "step 2"
+
+    listed = json.loads(_text(await list_files.invoke()))
+    assert listed == [{"name": "plan.md", "type": "file", "description": None}]
+
+    deleted = await delete.invoke(arguments={"file_name": "plan.md"})
+    assert "deleted" in _text(deleted)
+    missing = await read.invoke(arguments={"file_name": "plan.md"})
+    assert "not found" in _text(missing)
+    missing_delete = await delete.invoke(arguments={"file_name": "plan.md"})
+    assert "not found" in _text(missing_delete)
+
+
+async def test_description_sidecar_is_written_and_listed() -> None:
+    """Saving with a description writes a sidecar and surfaces it in listings."""
+    store = InMemoryAgentFileStore()
+    provider = FileMemoryProvider(store=store, scope="user-1")
+    _, tools = await _prepare(provider)
+    save = tools["file_memory_write"]
+    list_files = tools["file_memory_ls"]
+
+    result = await save.invoke(
+        arguments={"file_name": "arch.md", "content": "big content", "description": "system architecture"}
+    )
+    assert "with description" in _text(result)
+
+    sidecar = await store.read(_combine_paths("user-1", "arch_description.md"))
+    assert sidecar == "system architecture"
+
+    listed = json.loads(_text(await list_files.invoke()))
+    assert listed == [{"name": "arch.md", "type": "file", "description": "system architecture"}]
+
+    # Re-saving without a description removes the sidecar.
+    await save.invoke(arguments={"file_name": "arch.md", "content": "big content"})
+    assert await store.read(_combine_paths("user-1", "arch_description.md")) is None
+    listed_again = json.loads(_text(await list_files.invoke()))
+    assert listed_again == [{"name": "arch.md", "type": "file", "description": None}]
+
+
+async def test_delete_removes_sidecar() -> None:
+    """Deleting a file also removes its companion description sidecar."""
+    store = InMemoryAgentFileStore()
+    provider = FileMemoryProvider(store=store, scope="user-1")
+    _, tools = await _prepare(provider)
+
+    await tools["file_memory_write"].invoke(arguments={"file_name": "arch.md", "content": "x", "description": "desc"})
+    assert await store.read(_combine_paths("user-1", "arch_description.md")) == "desc"
+
+    await tools["file_memory_delete"].invoke(arguments={"file_name": "arch.md"})
+    assert await store.read(_combine_paths("user-1", "arch_description.md")) is None
+
+
+async def test_index_is_rebuilt_and_injected_on_next_run() -> None:
+    """Saved memories should be summarized in the index and injected as a context message."""
+    store = InMemoryAgentFileStore()
+    provider = FileMemoryProvider(store=store, scope="user-1")
+    _, tools = await _prepare(provider)
+
+    await tools["file_memory_write"].invoke(
+        arguments={"file_name": "arch.md", "content": "x", "description": "architecture"}
+    )
+    await tools["file_memory_write"].invoke(arguments={"file_name": "todo.md", "content": "y"})
+
+    index = await store.read(_combine_paths("user-1", _MEMORY_INDEX_FILE_NAME))
+    assert index is not None
+    assert "# Memory Index" in index
+    assert "- **arch.md**: architecture" in index
+    assert "- **todo.md**" in index
+
+    # A subsequent run injects the index as a user context message.
+    session = AgentSession(session_id="ignored")
+    context = SessionContext(session_id="ignored", input_messages=[])
+    await provider.before_run(agent=None, session=session, context=context, state={})
+    injected = context.context_messages.get(DEFAULT_FILE_MEMORY_SOURCE_ID, [])
+    assert len(injected) == 1
+    assert injected[0].role == "user"
+    assert "arch.md" in injected[0].text
+
+
+async def test_list_and_search_hide_internal_files() -> None:
+    """Listing and search must hide description sidecars and the memory index."""
+    store = InMemoryAgentFileStore()
+    provider = FileMemoryProvider(store=store, scope="user-1")
+    _, tools = await _prepare(provider)
+
+    await tools["file_memory_write"].invoke(
+        arguments={"file_name": "arch.md", "content": "architecture text", "description": "architecture"}
+    )
+
+    listed = json.loads(_text(await tools["file_memory_ls"].invoke()))
+    assert [e["name"] for e in listed] == ["arch.md"]
+
+    # The description text lives in an internal sidecar, so a regex matching it
+    # must not return the sidecar (only the memory file itself).
+    found = json.loads(_text(await tools["file_memory_grep"].invoke(arguments={"regex_pattern": "architecture"})))
+    names = [e["file_name"] for e in found]
+    assert "arch.md" in names
+    assert all(not _is_internal_file(name) for name in names)
+
+
+async def test_scope_isolates_memories_across_sessions() -> None:
+    """Two sessions sharing a store should not see each other's memories by default."""
+    store = InMemoryAgentFileStore()
+    provider = FileMemoryProvider(store=store)
+
+    _, tools_a = await _prepare(provider, session_id="session-a")
+    await tools_a["file_memory_write"].invoke(arguments={"file_name": "a.md", "content": "from a"})
+
+    _, tools_b = await _prepare(provider, session_id="session-b")
+    listed_b = json.loads(_text(await tools_b["file_memory_ls"].invoke()))
+    assert listed_b == []
+
+    # The original session still sees its own memory.
+    _, tools_a2 = await _prepare(provider, session_id="session-a")
+    listed_a = json.loads(_text(await tools_a2["file_memory_ls"].invoke()))
+    assert [e["name"] for e in listed_a] == ["a.md"]
+
+
+async def test_explicit_scope_shares_memories_across_sessions() -> None:
+    """An explicit scope groups memories regardless of session id."""
+    store = InMemoryAgentFileStore()
+    provider = FileMemoryProvider(store=store, scope="shared")
+
+    _, tools_a = await _prepare(provider, session_id="session-a")
+    await tools_a["file_memory_write"].invoke(arguments={"file_name": "shared.md", "content": "v"})
+
+    _, tools_b = await _prepare(provider, session_id="session-b")
+    listed_b = json.loads(_text(await tools_b["file_memory_ls"].invoke()))
+    assert [e["name"] for e in listed_b] == ["shared.md"]
+
+
+async def test_save_rejects_reserved_internal_names() -> None:
+    """Saving a file whose name collides with an internal file must be rejected."""
+    provider = FileMemoryProvider(store=InMemoryAgentFileStore())
+    _, tools = await _prepare(provider)
+    save = tools["file_memory_write"]
+
+    reserved = await save.invoke(arguments={"file_name": _MEMORY_INDEX_FILE_NAME, "content": "x"})
+    assert "reserved" in _text(reserved)
+
+    sidecar = await save.invoke(arguments={"file_name": "notes_description.md", "content": "x"})
+    assert "reserved" in _text(sidecar)
+
+
+async def test_tools_surface_path_validation_errors() -> None:
+    """Path traversal and rooted paths should be reported as tool messages, not raised."""
+    provider = FileMemoryProvider(store=InMemoryAgentFileStore())
+    _, tools = await _prepare(provider)
+
+    bad_save = await tools["file_memory_write"].invoke(arguments={"file_name": "../escape.md", "content": "x"})
+    assert "Could not write" in _text(bad_save)
+
+    bad_read = await tools["file_memory_read"].invoke(arguments={"file_name": "/rooted.md"})
+    assert "Could not read" in _text(bad_read)
+
+    bad_delete = await tools["file_memory_delete"].invoke(arguments={"file_name": "../escape.md"})
+    assert "Could not delete" in _text(bad_delete)
+
+
+async def test_provider_accepts_custom_instructions() -> None:
+    """Custom instructions override the default banner."""
+    provider = FileMemoryProvider(store=InMemoryAgentFileStore(), instructions="custom memory banner")
+    context, _ = await _prepare(provider)
+    assert "custom memory banner" in context.instructions
+    assert all(DEFAULT_FILE_MEMORY_INSTRUCTIONS not in chunk for chunk in context.instructions)
+
+
+def test_file_memory_provider_is_experimental() -> None:
+    """The provider should be marked experimental under the harness feature."""
+    assert getattr(FileMemoryProvider, "__feature_stage__", None) == "experimental"
+
+
+async def test_tools_reject_nested_paths() -> None:
+    """Memory files must stay flat; nested names are rejected/undiscoverable."""
+    store = InMemoryAgentFileStore()
+    provider = FileMemoryProvider(store=store)
+    _, tools = await _prepare(provider)
+
+    saved = await tools["file_memory_write"].invoke(arguments={"file_name": "notes/plan.md", "content": "x"})
+    assert "subdirectory" in _text(saved)
+    # Nothing should have been written for the nested name.
+    assert await store.list_children("") == []
+
+    # Backslash separators are normalized to "/" and rejected the same way.
+    saved_backslash = await tools["file_memory_write"].invoke(arguments={"file_name": "notes\\plan.md", "content": "x"})
+    assert "subdirectory" in _text(saved_backslash)
+
+    # Reading/deleting a nested name reports a clean "not found" message.
+    read_back = await tools["file_memory_read"].invoke(arguments={"file_name": "notes/plan.md"})
+    assert "not found" in _text(read_back)
+    deleted = await tools["file_memory_delete"].invoke(arguments={"file_name": "notes/plan.md"})
+    assert "not found" in _text(deleted)
+
+
+async def test_index_caps_entries_at_max() -> None:
+    """The rebuilt ``memories.md`` index lists at most ``_MAX_INDEX_ENTRIES`` files."""
+    store = InMemoryAgentFileStore()
+    provider = FileMemoryProvider(store=store, scope="user-1")
+    _, tools = await _prepare(provider)
+    save = tools["file_memory_write"]
+
+    total = _MAX_INDEX_ENTRIES + 5
+    for i in range(total):
+        await save.invoke(arguments={"file_name": f"memory-{i:03d}.md", "content": "x"})
+
+    index = await store.read(_combine_paths("user-1", _MEMORY_INDEX_FILE_NAME))
+    assert index is not None
+    entry_lines = [line for line in index.splitlines() if line.startswith("- ")]
+    assert len(entry_lines) == _MAX_INDEX_ENTRIES
+
+
+async def test_tools_surface_store_value_errors() -> None:
+    """``ValueError`` raised by the store is returned as a tool message, not raised."""
+
+    class _ValueErrorStore(InMemoryAgentFileStore):
+        async def write(self, path: str, content: str, *, overwrite: bool = True) -> None:
+            raise ValueError("boom-write")
+
+        async def read(self, path: str) -> str | None:
+            raise ValueError("boom-read")
+
+        async def delete(self, path: str) -> bool:
+            raise ValueError("boom-delete")
+
+    provider = FileMemoryProvider(store=_ValueErrorStore())
+    _, tools = await _prepare(provider)
+
+    saved = await tools["file_memory_write"].invoke(arguments={"file_name": "plan.md", "content": "x"})
+    assert "Could not write" in _text(saved) and "boom-write" in _text(saved)
+
+    read_back = await tools["file_memory_read"].invoke(arguments={"file_name": "plan.md"})
+    assert "Could not read" in _text(read_back) and "boom-read" in _text(read_back)
+
+    deleted = await tools["file_memory_delete"].invoke(arguments={"file_name": "plan.md"})
+    assert "Could not delete" in _text(deleted) and "boom-delete" in _text(deleted)
+
+
+async def test_before_run_skips_injection_when_index_unreadable() -> None:
+    """A failing index read must not crash the run; injection is simply skipped."""
+
+    class _UnreadableIndexStore(InMemoryAgentFileStore):
+        async def read(self, path: str) -> str | None:
+            if path.endswith(_MEMORY_INDEX_FILE_NAME):
+                raise ValueError("corrupt index")
+            return await super().read(path)
+
+    store = _UnreadableIndexStore()
+    # Seed an index so before_run attempts to read it.
+    await store.write(_combine_paths("user-1", _MEMORY_INDEX_FILE_NAME), "# Memory Index\n")
+    provider = FileMemoryProvider(store=store, scope="user-1")
+
+    session = AgentSession(session_id="s-1")
+    context = SessionContext(session_id="s-1", input_messages=[])
+    # Should not raise despite the unreadable index.
+    await provider.before_run(agent=None, session=session, context=context, state={})
+    assert context.context_messages.get(DEFAULT_FILE_MEMORY_SOURCE_ID, []) == []
+
+
+async def test_search_propagates_invalid_regex() -> None:
+    """An invalid regex from the model is surfaced as a raised error so it can retry."""
+    provider = FileMemoryProvider(store=InMemoryAgentFileStore())
+    _, tools = await _prepare(provider)
+
+    with pytest.raises(re.error):
+        await tools["file_memory_grep"].invoke(arguments={"regex_pattern": "[unclosed"})
+
+
+async def test_memory_replace() -> None:
+    """``file_memory_replace`` substitutes text and enforces match-count rules."""
+    provider = FileMemoryProvider(store=InMemoryAgentFileStore())
+    _, tools = await _prepare(provider)
+    await tools["file_memory_write"].invoke(arguments={"file_name": "a.md", "content": "foo bar foo"})
+
+    multi = await tools["file_memory_replace"].invoke(
+        arguments={"file_name": "a.md", "old_string": "foo", "new_string": "baz"}
+    )
+    assert "2 times" in _text(multi)
+
+    done = await tools["file_memory_replace"].invoke(
+        arguments={"file_name": "a.md", "old_string": "foo", "new_string": "baz", "replace_all": True}
+    )
+    assert "2 occurrence" in _text(done)
+    assert _text(await tools["file_memory_read"].invoke(arguments={"file_name": "a.md"})) == "baz bar baz"
+
+    # Unique single occurrence with the default replace_all=False -> replaces exactly one.
+    await tools["file_memory_write"].invoke(arguments={"file_name": "u.md", "content": "alpha beta gamma"})
+    single = await tools["file_memory_replace"].invoke(
+        arguments={"file_name": "u.md", "old_string": "beta", "new_string": "BETA"}
+    )
+    assert "1 occurrence" in _text(single)
+    assert _text(await tools["file_memory_read"].invoke(arguments={"file_name": "u.md"})) == "alpha BETA gamma"
+
+    # Internal files (the memories.md index and *_description.md sidecars) are reserved.
+    reserved = await tools["file_memory_replace"].invoke(
+        arguments={"file_name": "memories.md", "old_string": "x", "new_string": "y"}
+    )
+    assert "reserved for internal use" in _text(reserved)
+    reserved_desc = await tools["file_memory_replace"].invoke(
+        arguments={"file_name": "a_description.md", "old_string": "x", "new_string": "y"}
+    )
+    assert "reserved for internal use" in _text(reserved_desc)
+
+
+async def test_memory_replace_lines() -> None:
+    """``file_memory_replace_lines`` applies literal 1-based line edits and rejects bad input."""
+    provider = FileMemoryProvider(store=InMemoryAgentFileStore())
+    _, tools = await _prepare(provider)
+
+    async def write(content: str) -> None:
+        await tools["file_memory_write"].invoke(arguments={"file_name": "a.md", "content": content})
+
+    async def current() -> str:
+        return _text(await tools["file_memory_read"].invoke(arguments={"file_name": "a.md"}))
+
+    # Literal replacement: the caller supplies the trailing newline.
+    await write("one\ntwo\nthree")
+    done = await tools["file_memory_replace_lines"].invoke(
+        arguments={"file_name": "a.md", "edits": [{"line_number": 2, "new_line": "TWO\n"}]}
+    )
+    assert "1 line" in _text(done)
+    assert await current() == "one\nTWO\nthree"
+
+    # Empty new_line deletes a line; embedded newlines expand one line into several.
+    await write("a\nb\nc\n")
+    await tools["file_memory_replace_lines"].invoke(
+        arguments={
+            "file_name": "a.md",
+            "edits": [{"line_number": 1, "new_line": ""}, {"line_number": 2, "new_line": "b1\nb2\n"}],
+        }
+    )
+    assert await current() == "b1\nb2\nc\n"
+
+    oor = await tools["file_memory_replace_lines"].invoke(
+        arguments={"file_name": "a.md", "edits": [{"line_number": 9, "new_line": "x"}]}
+    )
+    assert "out of range" in _text(oor)
+
+    # Internal files (the memories.md index and *_description.md sidecars) are reserved.
+    reserved = await tools["file_memory_replace_lines"].invoke(
+        arguments={"file_name": "memories.md", "edits": [{"line_number": 1, "new_line": "x"}]}
+    )
+    assert "reserved for internal use" in _text(reserved)
+
+    # Empty edits list -> failure surfaced to the caller.
+    await write("one\ntwo")
+    empty = await tools["file_memory_replace_lines"].invoke(arguments={"file_name": "a.md", "edits": []})
+    assert "At least one line edit" in _text(empty)
+
+    # Duplicate line numbers -> failure surfaced to the caller.
+    dup = await tools["file_memory_replace_lines"].invoke(
+        arguments={
+            "file_name": "a.md",
+            "edits": [{"line_number": 1, "new_line": "x"}, {"line_number": 1, "new_line": "y"}],
+        }
+    )
+    assert "Duplicate" in _text(dup)

python/packages/core/tests/core/test_harness_memory.py

@@ -0,0 +1,775 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import asyncio
+import json
+from collections.abc import Mapping, Sequence
+from datetime import datetime, timedelta, timezone
+from typing import Any
+
+import pytest
+
+from agent_framework import (
+    DEFAULT_MEMORY_SOURCE_ID,
+    Agent,
+    AgentSession,
+    ChatOptions,
+    ChatResponse,
+    Content,
+    ExperimentalFeature,
+    FileHistoryProvider,
+    MemoryContextProvider,
+    MemoryFileStore,
+    MemoryIndexEntry,
+    MemoryStore,
+    MemoryTopicRecord,
+    Message,
+)
+
+
+def _no_store_options() -> ChatOptions:
+    return {"store": False}
+
+
+def _tool_by_name(tools: list[object], name: str) -> object:
+    """Return the tool with the requested name from a prepared tool list."""
+    for tool in tools:
+        if getattr(tool, "name", None) == name:
+            return tool
+    raise AssertionError(f"Tool {name!r} was not found.")
+
+
+class _MemoryHarnessClient:
+    """Deterministic chat client used by the memory harness tests."""
+
+    additional_properties: dict[str, Any]
+
+    def __init__(
+        self,
+        *,
+        extraction_payload: dict[str, Any] | None = None,
+        consolidation_payload: dict[str, Any] | None = None,
+        default_text: str = "Assistant reply.",
+    ) -> None:
+        self.additional_properties = {}
+        self.extraction_payload = extraction_payload or {
+            "memories": [
+                {
+                    "topic": "preferences",
+                    "memory": "Prefers concise answers.",
+                }
+            ]
+        }
+        self.consolidation_payload = consolidation_payload or {
+            "summary": "Prefers concise answers.",
+            "memories": ["Prefers concise answers."],
+        }
+        self.default_text = default_text
+        self.calls: list[str] = []
+
+    async def get_response(
+        self,
+        messages: Sequence[Message],
+        *,
+        stream: bool = False,
+        options: Mapping[str, Any] | None = None,
+        compaction_strategy: object | None = None,
+        tokenizer: object | None = None,
+        function_invocation_kwargs: Mapping[str, Any] | None = None,
+        client_kwargs: Mapping[str, Any] | None = None,
+    ) -> ChatResponse[Any]:
+        del options, compaction_strategy, tokenizer, function_invocation_kwargs, client_kwargs
+        assert not stream
+        system_text = messages[0].text if messages and messages[0].role == "system" else ""
+        if "extract durable memory candidates" in system_text.lower():
+            self.calls.append("extract")
+            return ChatResponse(messages=[Message(role="assistant", contents=[json.dumps(self.extraction_payload)])])
+        if "consolidate one topic memory file" in system_text.lower():
+            self.calls.append("consolidate")
+            return ChatResponse(messages=[Message(role="assistant", contents=[json.dumps(self.consolidation_payload)])])
+        self.calls.append("agent")
+        return ChatResponse(messages=[Message(role="assistant", contents=[self.default_text])])
+
+
+def test_memory_index_entry_round_trips_and_trims_pointer_lines() -> None:
+    """Memory index entries should preserve value equality and trim pointer lines."""
+    raw_entry = {
+        "topic": "Architecture Decisions",
+        "slug": "architecture-decisions",
+        "summary": (
+            "PostgreSQL was chosen because it keeps the relational model while supporting flexible JSONB fields."
+        ),
+        "updated_at": "2026-04-21T10:00:00+00:00",
+    }
+
+    entry = MemoryIndexEntry.from_dict(raw_entry)
+
+    assert entry == MemoryIndexEntry(**raw_entry)
+    assert entry.to_dict() == raw_entry
+    assert len(entry.to_pointer_line(max_length=80)) <= 80
+    assert "MemoryIndexEntry(" in repr(entry)
+
+
+def test_memory_topic_record_round_trips_through_dict_and_markdown() -> None:
+    """Topic memory records should preserve their structured content and markdown form."""
+    raw_record = {
+        "topic": "preferences",
+        "slug": "preferences",
+        "summary": "Prefers concise answers.",
+        "memories": ["Prefers concise answers.", "Prefers aisle seats."],
+        "updated_at": "2026-04-21T10:05:00+00:00",
+        "session_ids": ["session-1", "session-2"],
+    }
+
+    record = MemoryTopicRecord.from_dict(raw_record)
+    reparsed_record = MemoryTopicRecord.from_markdown(record.to_markdown())
+
+    assert record == MemoryTopicRecord(**raw_record)  # type: ignore[arg-type]
+    assert record.to_dict() == raw_record
+    assert reparsed_record == record
+    assert "MemoryTopicRecord(" in repr(record)
+
+
+async def test_memory_file_store_writes_topics_index_state_and_transcripts(tmp_path) -> None:
+    """The file-backed memory store should manage topics, ``MEMORY.md``, state, and transcript search."""
+    store = MemoryFileStore(
+        tmp_path,
+        kind="memories",
+        owner_prefix="user_",
+        owner_state_key="owner_id",
+        dumps=lambda value: json.dumps(value, separators=(",", ":"), sort_keys=True),
+        loads=json.loads,
+    )
+    session = AgentSession(session_id="session-1")
+    session.state["owner_id"] = "alice"
+    updated_at = datetime(2026, 4, 21, tzinfo=timezone.utc).replace(microsecond=0).isoformat()
+
+    preferences_record = MemoryTopicRecord(
+        topic="preferences",
+        summary="Prefers concise answers.",
+        memories=["Prefers concise answers.", "Prefers aisle seats."],
+        updated_at=updated_at,
+        session_ids=["session-1"],
+    )
+    travel_record = MemoryTopicRecord(
+        topic="travel",
+        summary="Planning a Norway trip.",
+        memories=["Visit Oslo in June."],
+        updated_at=updated_at,
+        session_ids=["session-1"],
+    )
+
+    store.write_topic(session, preferences_record, source_id=DEFAULT_MEMORY_SOURCE_ID)
+    store.write_topic(session, travel_record, source_id=DEFAULT_MEMORY_SOURCE_ID)
+    entries = store.rebuild_index(
+        session,
+        source_id=DEFAULT_MEMORY_SOURCE_ID,
+        line_limit=200,
+        line_length=150,
+    )
+
+    assert [entry.topic for entry in entries] == ["preferences", "travel"]
+    assert "preferences" in store.get_index_text(
+        session,
+        source_id=DEFAULT_MEMORY_SOURCE_ID,
+        line_limit=200,
+        line_length=150,
+    )
+
+    assert store.read_state(session, source_id=DEFAULT_MEMORY_SOURCE_ID) == {
+        "last_consolidated_at": None,
+        "sessions_since_consolidation": [],
+    }
+    store.write_state(
+        session,
+        {
+            "last_consolidated_at": updated_at,
+            "sessions_since_consolidation": ["session-1"],
+        },
+        source_id=DEFAULT_MEMORY_SOURCE_ID,
+    )
+    assert store.read_state(
+        session,
+        source_id=DEFAULT_MEMORY_SOURCE_ID,
+    )["sessions_since_consolidation"] == ["session-1"]
+
+    history_provider = FileHistoryProvider(
+        store.get_transcripts_directory(session, source_id=DEFAULT_MEMORY_SOURCE_ID),
+        dumps=lambda value: json.dumps(value, separators=(",", ":"), sort_keys=True),
+        loads=json.loads,
+    )
+    await history_provider.save_messages(
+        session.session_id,
+        [
+            Message(role="user", contents=["I prefer aisle seats."]),
+            Message(role="assistant", contents=["Recorded."]),
+        ],
+    )
+
+    assert store.search_transcripts(session, source_id=DEFAULT_MEMORY_SOURCE_ID, query="aisle") == [
+        {
+            "session_id": "session-1",
+            "line_number": 1,
+            "role": "user",
+            "text": "I prefer aisle seats.",
+        }
+    ]
+
+
+def test_memory_file_store_rejects_owner_path_traversal(tmp_path) -> None:
+    """Owner IDs with path traversal segments should not escape ``base_path``."""
+    session = AgentSession(session_id="session-1")
+    session.state["owner_id"] = "../escape"
+    store = MemoryFileStore(tmp_path, owner_state_key="owner_id")
+    record = MemoryTopicRecord(
+        topic="preferences",
+        summary="Prefers concise answers.",
+        memories=["Prefers concise answers."],
+        updated_at=datetime(2026, 4, 21, tzinfo=timezone.utc).replace(microsecond=0).isoformat(),
+    )
+
+    with pytest.raises(ValueError, match="path traversal"):
+        store.write_topic(session, record, source_id=DEFAULT_MEMORY_SOURCE_ID)
+
+    assert not (tmp_path.parent / "escape").exists()
+
+
+async def test_memory_file_store_namespaces_topics_state_and_transcripts_by_source_id(tmp_path) -> None:
+    """Providers sharing one file store should not collide when they use different source IDs."""
+    session = AgentSession(session_id="session-1")
+    session.state["owner_id"] = "alice"
+    store = MemoryFileStore(
+        tmp_path,
+        owner_state_key="owner_id",
+        dumps=lambda value: json.dumps(value, separators=(",", ":"), sort_keys=True),
+        loads=json.loads,
+    )
+    updated_at = datetime(2026, 4, 21, tzinfo=timezone.utc).replace(microsecond=0).isoformat()
+
+    store.write_topic(
+        session,
+        MemoryTopicRecord(
+            topic="preferences",
+            summary="Source A summary.",
+            memories=["Source A memory."],
+            updated_at=updated_at,
+        ),
+        source_id="source-a",
+    )
+    store.write_topic(
+        session,
+        MemoryTopicRecord(
+            topic="preferences",
+            summary="Source B summary.",
+            memories=["Source B memory."],
+            updated_at=updated_at,
+        ),
+        source_id="source-b",
+    )
+    store.write_state(
+        session, {"last_consolidated_at": updated_at, "sessions_since_consolidation": ["a"]}, source_id="source-a"
+    )
+    store.write_state(
+        session, {"last_consolidated_at": None, "sessions_since_consolidation": ["b"]}, source_id="source-b"
+    )
+
+    await FileHistoryProvider(store.get_transcripts_directory(session, source_id="source-a")).save_messages(
+        "session-1", [Message(role="user", contents=["Source A transcript."])]
+    )
+    await FileHistoryProvider(store.get_transcripts_directory(session, source_id="source-b")).save_messages(
+        "session-1", [Message(role="user", contents=["Source B transcript."])]
+    )
+
+    assert store.get_topic(session, source_id="source-a", topic="preferences").memories == ["Source A memory."]
+    assert store.get_topic(session, source_id="source-b", topic="preferences").memories == ["Source B memory."]
+    assert store.read_state(session, source_id="source-a")["sessions_since_consolidation"] == ["a"]
+    assert store.read_state(session, source_id="source-b")["sessions_since_consolidation"] == ["b"]
+    assert (
+        store.search_transcripts(session, source_id="source-a", query="transcript")[0]["text"] == "Source A transcript."
+    )
+    assert (
+        store.search_transcripts(session, source_id="source-b", query="transcript")[0]["text"] == "Source B transcript."
+    )
+
+
+async def test_memory_context_provider_does_not_rewrite_unchanged_index(tmp_path) -> None:
+    """A second before-run pass with unchanged memories should preserve ``MEMORY.md`` mtime."""
+    session = AgentSession(session_id="session-1")
+    session.state["owner_id"] = "alice"
+    store = MemoryFileStore(tmp_path, owner_state_key="owner_id")
+    agent = Agent(
+        client=_MemoryHarnessClient(),  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+        context_providers=[MemoryContextProvider(store=store)],
+        default_options=_no_store_options(),
+    )
+
+    await agent._prepare_session_and_messages(  # pyright: ignore[reportPrivateUsage]
+        session=session,
+        input_messages=[Message(role="user", contents=["Current question"])],
+    )
+    index_path = next(tmp_path.rglob("MEMORY.md"))
+    first_mtime_ns = index_path.stat().st_mtime_ns
+
+    await agent._prepare_session_and_messages(  # pyright: ignore[reportPrivateUsage]
+        session=session,
+        input_messages=[Message(role="user", contents=["Current question"])],
+    )
+
+    assert index_path.stat().st_mtime_ns == first_mtime_ns
+
+
+async def test_memory_context_provider_tools_and_automation(tmp_path) -> None:
+    """The memory provider should expose tools and automate extraction plus consolidation."""
+    session = AgentSession(session_id="session-1")
+    session.state["owner_id"] = "alice"
+    store = MemoryFileStore(
+        tmp_path,
+        kind="memories",
+        owner_prefix="user_",
+        owner_state_key="owner_id",
+        dumps=lambda value: json.dumps(value, separators=(",", ":"), sort_keys=True),
+        loads=json.loads,
+    )
+    provider = MemoryContextProvider(
+        store=store,
+        consolidation_min_sessions=1,
+        consolidation_interval=timedelta(0),
+    )
+    agent = Agent(
+        client=_MemoryHarnessClient(),  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+        context_providers=[provider],
+        default_options=_no_store_options(),
+    )
+
+    _, options = await agent._prepare_session_and_messages(  # pyright: ignore[reportPrivateUsage]
+        session=session,
+        input_messages=[Message(role="user", contents=["Remember this."])],
+    )
+    tools = options["tools"]
+    assert isinstance(tools, list)
+
+    write_memory = _tool_by_name(tools, "write_memory")
+    list_memory_topics = _tool_by_name(tools, "list_memory_topics")
+    search_memory_transcripts = _tool_by_name(tools, "search_memory_transcripts")
+    consolidate_memories = _tool_by_name(tools, "consolidate_memories")
+
+    write_result = await write_memory.invoke(arguments={"topic": "travel", "memory": "Visit Oslo in June."})  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    created_topic = json.loads(write_result[0].text)
+    assert created_topic["topic"] == "travel"
+
+    list_result = await list_memory_topics.invoke()  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert [entry["topic"] for entry in json.loads(list_result[0].text)] == ["travel"]
+
+    await agent.run("Please remember that I prefer concise answers.", session=session)
+
+    serialized_session = session.to_dict()
+    assert serialized_session["state"][DEFAULT_MEMORY_SOURCE_ID] == {"owner_id": "alice"}
+
+    preferences_topic = store.get_topic(session, source_id=DEFAULT_MEMORY_SOURCE_ID, topic="preferences")
+    assert preferences_topic.summary == "Prefers concise answers."
+    assert preferences_topic.memories == ["Prefers concise answers."]
+
+    transcript_search_result = await search_memory_transcripts.invoke(arguments={"query": "concise", "limit": 5})  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    search_payload = json.loads(transcript_search_result[0].text)
+    assert search_payload[0]["role"] == "user"
+    assert "concise answers" in search_payload[0]["text"]
+
+    consolidate_result = await consolidate_memories.invoke()  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert json.loads(consolidate_result[0].text)["consolidated_topics"] >= 1
+
+
+async def test_memory_context_provider_injects_recent_turns(tmp_path) -> None:
+    """The memory provider should inject only the configured recent transcript turns."""
+    session = AgentSession(session_id="session-1")
+    session.state["owner_id"] = "alice"
+    store = MemoryFileStore(
+        tmp_path,
+        kind="memories",
+        owner_prefix="user_",
+        owner_state_key="owner_id",
+        dumps=lambda value: json.dumps(value, separators=(",", ":"), sort_keys=True),
+        loads=json.loads,
+    )
+    provider = MemoryContextProvider(store=store, recent_turns=2)
+    provider_state = store.export_provider_state(session)
+    await provider.save_messages(
+        session.session_id,
+        [
+            Message(role="user", contents=["First question"]),
+            Message(role="assistant", contents=["First answer"]),
+            Message(role="user", contents=["Second question"]),
+            Message(role="assistant", contents=["Second answer"]),
+            Message(role="user", contents=["Third question"]),
+            Message(role="assistant", contents=["Third answer"]),
+        ],
+        state=provider_state,
+    )
+    agent = Agent(
+        client=_MemoryHarnessClient(),  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+        context_providers=[provider],
+        default_options=_no_store_options(),
+    )
+
+    session_context, _ = await agent._prepare_session_and_messages(  # pyright: ignore[reportPrivateUsage]
+        session=session,
+        input_messages=[Message(role="user", contents=["Current question"])],
+    )
+    prepared_messages = session_context.get_messages(include_input=True)
+
+    assert [message.text for message in prepared_messages[:4]] == [
+        "Second question",
+        "Second answer",
+        "Third question",
+        "Third answer",
+    ]
+    assert "First question" not in [message.text for message in prepared_messages]
+    assert "### MEMORY.md" in prepared_messages[4].text
+    assert prepared_messages[-1].text == "Current question"
+
+
+async def test_memory_context_provider_recent_turns_can_skip_tool_call_groups(tmp_path) -> None:
+    """Recent-turn loading should follow compaction grouping and optionally skip tool-call groups."""
+    session = AgentSession(session_id="session-1")
+    session.state["owner_id"] = "alice"
+    store = MemoryFileStore(
+        tmp_path,
+        kind="memories",
+        owner_prefix="user_",
+        owner_state_key="owner_id",
+        dumps=lambda value: json.dumps(value, separators=(",", ":"), sort_keys=True),
+        loads=json.loads,
+    )
+    provider_state = store.export_provider_state(session)
+    await MemoryContextProvider(store=store).save_messages(
+        session.session_id,
+        [
+            Message(role="user", contents=["First question"]),
+            Message(role="assistant", contents=["First answer"]),
+            Message(role="user", contents=["Second question"]),
+            Message(role="assistant", contents=[Content.from_text_reasoning(text="Let me check that.")]),
+            Message(
+                role="assistant",
+                contents=[
+                    Content.from_function_call(call_id="call-1", name="lookup_answer", arguments='{"topic":"second"}')
+                ],
+            ),
+            Message(role="tool", contents=[Content.from_function_result(call_id="call-1", result="Tool result")]),
+            Message(role="assistant", contents=["Second final answer"]),
+            Message(role="user", contents=["Third question"]),
+            Message(role="assistant", contents=["Third answer"]),
+        ],
+        state=provider_state,
+    )
+    with_tools_agent = Agent(
+        client=_MemoryHarnessClient(),  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+        context_providers=[MemoryContextProvider(store=store, recent_turns=2, load_tool_turns=True)],
+        default_options=_no_store_options(),
+    )
+    without_tools_agent = Agent(
+        client=_MemoryHarnessClient(),  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+        context_providers=[MemoryContextProvider(store=store, recent_turns=2, load_tool_turns=False)],
+        default_options=_no_store_options(),
+    )
+
+    with_tools_context, _ = await with_tools_agent._prepare_session_and_messages(  # pyright: ignore[reportPrivateUsage]
+        session=session,
+        input_messages=[Message(role="user", contents=["Current question"])],
+    )
+    without_tools_context, _ = await without_tools_agent._prepare_session_and_messages(  # pyright: ignore[reportPrivateUsage]
+        session=session,
+        input_messages=[Message(role="user", contents=["Current question"])],
+    )
+    with_tools_messages = with_tools_context.get_messages(include_input=True)
+    without_tools_messages = without_tools_context.get_messages(include_input=True)
+
+    assert [message.text for message in without_tools_messages[:4]] == [
+        "Second question",
+        "Second final answer",
+        "Third question",
+        "Third answer",
+    ]
+    assert not any(message.role == "tool" for message in without_tools_messages)
+    assert not any(
+        any(content.type == "function_call" for content in message.contents) for message in without_tools_messages
+    )
+    assert not any(
+        any(content.type == "text_reasoning" for content in message.contents) for message in without_tools_messages
+    )
+
+    assert with_tools_messages[0].text == "Second question"
+    assert with_tools_messages[1].contents[0].type == "text_reasoning"
+    assert with_tools_messages[2].contents[0].type == "function_call"
+    assert with_tools_messages[3].role == "tool"
+    assert with_tools_messages[3].contents[0].type == "function_result"
+    assert with_tools_messages[4].text == "Second final answer"
+
+
+async def test_memory_context_provider_uses_explicit_consolidation_client(tmp_path) -> None:
+    """The memory provider should use the explicit consolidation client when one is configured."""
+    session = AgentSession(session_id="session-1")
+    session.state["owner_id"] = "alice"
+    store = MemoryFileStore(
+        tmp_path,
+        kind="memories",
+        owner_prefix="user_",
+        owner_state_key="owner_id",
+        dumps=lambda value: json.dumps(value, separators=(",", ":"), sort_keys=True),
+        loads=json.loads,
+    )
+    main_client = _MemoryHarnessClient()
+    consolidation_client = _MemoryHarnessClient(
+        consolidation_payload={
+            "summary": "Consolidated by the cheaper client.",
+            "memories": ["Visit Oslo in June."],
+        }
+    )
+    provider = MemoryContextProvider(
+        store=store,
+        consolidation_client=consolidation_client,  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+    )
+    agent = Agent(
+        client=main_client,  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+        context_providers=[provider],
+        default_options=_no_store_options(),
+    )
+
+    _, options = await agent._prepare_session_and_messages(  # pyright: ignore[reportPrivateUsage]
+        session=session,
+        input_messages=[Message(role="user", contents=["Remember this."])],
+    )
+    tools = options["tools"]
+    assert isinstance(tools, list)
+
+    write_memory = _tool_by_name(tools, "write_memory")
+    consolidate_memories = _tool_by_name(tools, "consolidate_memories")
+
+    await write_memory.invoke(arguments={"topic": "travel", "memory": "Visit Oslo in June."})  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    await consolidate_memories.invoke()  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+    travel_topic = store.get_topic(session, source_id=DEFAULT_MEMORY_SOURCE_ID, topic="travel")
+    assert travel_topic.summary == "Consolidated by the cheaper client."
+    assert main_client.calls == []
+    assert consolidation_client.calls == ["consolidate"]
+
+
+async def test_memory_context_provider_preserves_concurrent_writes_to_same_topic(tmp_path) -> None:
+    """Concurrent writes to one topic should preserve every memory line."""
+    session = AgentSession(session_id="session-1")
+    session.state["owner_id"] = "alice"
+    store = MemoryFileStore(tmp_path, owner_state_key="owner_id")
+    provider = MemoryContextProvider(store=store)
+    agent = Agent(client=_MemoryHarnessClient(), context_providers=[provider], default_options=_no_store_options())  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+
+    _, options = await agent._prepare_session_and_messages(  # pyright: ignore[reportPrivateUsage]
+        session=session,
+        input_messages=[Message(role="user", contents=["Remember these."])],
+    )
+    tools = options["tools"]
+    assert isinstance(tools, list)
+    write_memory = _tool_by_name(tools, "write_memory")
+    memories = [f"Concurrent memory {index}." for index in range(20)]
+
+    await asyncio.gather(
+        *(write_memory.invoke(arguments={"topic": "preferences", "memory": memory}) for memory in memories)  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    )
+
+    topic = store.get_topic(session, source_id=DEFAULT_MEMORY_SOURCE_ID, topic="preferences")
+    assert sorted(topic.memories) == sorted(memories)
+
+
+def test_memory_harness_classes_are_marked_experimental() -> None:
+    """Memory harness public classes should expose HARNESS experimental metadata."""
+    assert MemoryIndexEntry.__feature_id__ == ExperimentalFeature.HARNESS.value  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert MemoryTopicRecord.__feature_id__ == ExperimentalFeature.HARNESS.value  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert MemoryStore.__feature_id__ == ExperimentalFeature.HARNESS.value  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert MemoryFileStore.__feature_id__ == ExperimentalFeature.HARNESS.value  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert MemoryContextProvider.__feature_id__ == ExperimentalFeature.HARNESS.value  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert ".. warning:: Experimental" in MemoryContextProvider.__doc__  # type: ignore[operator]  # pyrefly: ignore[not-iterable]  # ty: ignore[unsupported-operator]
+
+
+def test_memory_topic_record_round_trips_when_text_contains_section_markers() -> None:
+    """Embedded ``## Summary``/``## Memories`` markers must not be re-interpreted as headings."""
+    record = MemoryTopicRecord(
+        topic="weird",
+        summary="Multi line summary.\n## Summary\nstill summary",
+        memories=[
+            "## Memories pretend",
+            "Real memory.",
+            "  ## Memories nested",
+        ],
+        updated_at="2026-04-21T10:00:00+00:00",
+        session_ids=["session-1"],
+    )
+
+    reparsed = MemoryTopicRecord.from_markdown(record.to_markdown())
+
+    assert reparsed.summary == record.summary
+    assert reparsed.memories == record.memories
+
+
+async def test_memory_file_store_atomic_write_preserves_prior_topic_on_failure(tmp_path, monkeypatch) -> None:
+    """If ``os.replace`` fails mid-write, the previous topic file must remain intact."""
+    from agent_framework._harness import _memory as memory_module
+
+    store = MemoryFileStore(tmp_path, owner_state_key="owner_id")
+    session = AgentSession(session_id="session-1")
+    session.state["owner_id"] = "alice"
+    original = MemoryTopicRecord(
+        topic="preferences",
+        summary="Prefers concise answers.",
+        memories=["Prefers concise answers."],
+        updated_at="2026-04-21T10:00:00+00:00",
+        session_ids=["session-1"],
+    )
+    store.write_topic(session, original, source_id=DEFAULT_MEMORY_SOURCE_ID)
+
+    real_replace = memory_module.os.replace
+
+    def _boom(*args: object, **kwargs: object) -> None:
+        raise OSError("simulated disk-full")
+
+    monkeypatch.setattr(memory_module.os, "replace", _boom)
+    with pytest.raises(OSError, match="simulated disk-full"):
+        store.write_topic(
+            session,
+            MemoryTopicRecord(
+                topic="preferences",
+                summary="Updated.",
+                memories=["Updated."],
+                updated_at="2026-04-21T11:00:00+00:00",
+                session_ids=["session-1"],
+            ),
+            source_id=DEFAULT_MEMORY_SOURCE_ID,
+        )
+
+    monkeypatch.setattr(memory_module.os, "replace", real_replace)
+    surviving = store.get_topic(session, source_id=DEFAULT_MEMORY_SOURCE_ID, topic="preferences")
+    assert surviving.summary == "Prefers concise answers."
+    # Temp file should not be left behind.
+    topics_dir = surviving_dir = tmp_path
+    leftover = [path for path in topics_dir.rglob("*.tmp.*")]
+    assert leftover == []
+    del surviving_dir
+
+
+async def test_memory_file_store_does_not_mkdir_on_pure_read_paths(tmp_path) -> None:
+    """List/read calls on a never-written session should not create any directories."""
+    store = MemoryFileStore(tmp_path, owner_state_key="owner_id")
+    session = AgentSession(session_id="session-1")
+    session.state["owner_id"] = "alice"
+
+    assert store.list_topics(session, source_id=DEFAULT_MEMORY_SOURCE_ID) == []
+    assert store.read_state(session, source_id=DEFAULT_MEMORY_SOURCE_ID) == {
+        "last_consolidated_at": None,
+        "sessions_since_consolidation": [],
+    }
+    assert store.search_transcripts(session, source_id=DEFAULT_MEMORY_SOURCE_ID, query="anything") == []
+
+    # tmp_path itself was passed in by pytest so it exists; assert no children were created.
+    assert list(tmp_path.iterdir()) == []
+
+
+class _RaisingMemoryClient:
+    """Chat client that raises a transient error for every consolidation request."""
+
+    additional_properties: dict[str, Any]
+
+    def __init__(self) -> None:
+        from agent_framework.exceptions import ChatClientException
+
+        self.additional_properties = {}
+        self.error_class = ChatClientException
+        self.calls: list[str] = []
+
+    async def get_response(
+        self,
+        messages: Sequence[Message],
+        *,
+        stream: bool = False,
+        options: Mapping[str, Any] | None = None,
+        compaction_strategy: object | None = None,
+        tokenizer: object | None = None,
+        function_invocation_kwargs: Mapping[str, Any] | None = None,
+        client_kwargs: Mapping[str, Any] | None = None,
+    ) -> ChatResponse[Any]:
+        del messages, stream, options, compaction_strategy, tokenizer
+        del function_invocation_kwargs, client_kwargs
+        self.calls.append("call")
+        raise self.error_class("simulated transient failure")
+
+
+class _ProgrammerErrorMemoryClient:
+    """Chat client whose ``get_response`` raises a non-transient programmer error."""
+
+    additional_properties: dict[str, Any]
+
+    def __init__(self) -> None:
+        self.additional_properties = {}
+
+    async def get_response(self, *args: object, **kwargs: object) -> ChatResponse[Any]:
+        del args, kwargs
+        raise AttributeError("misconfigured client")
+
+
+async def test_memory_consolidation_transient_failure_preserves_state(tmp_path) -> None:
+    """A transient consolidation failure must not advance the maintenance window."""
+    session = AgentSession(session_id="session-1")
+    session.state["owner_id"] = "alice"
+    store = MemoryFileStore(tmp_path, owner_state_key="owner_id")
+    raising_client = _RaisingMemoryClient()
+    provider = MemoryContextProvider(store=store, consolidation_client=raising_client)  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+    pre_state = {
+        "last_consolidated_at": "2026-04-20T09:00:00+00:00",
+        "sessions_since_consolidation": ["queued-session"],
+    }
+    store.write_state(session, pre_state, source_id=DEFAULT_MEMORY_SOURCE_ID)
+    store.write_topic(
+        session,
+        MemoryTopicRecord(
+            topic="preferences",
+            summary="Prefers concise answers.",
+            memories=["Prefers concise answers."],
+            updated_at="2026-04-21T10:00:00+00:00",
+            session_ids=["session-1"],
+        ),
+        source_id=DEFAULT_MEMORY_SOURCE_ID,
+    )
+
+    consolidated_count = await provider._run_consolidation(  # pyright: ignore[reportPrivateUsage]
+        client=raising_client,  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+        session=session,
+        force=True,
+        now=datetime(2026, 4, 22, tzinfo=timezone.utc),
+    )
+
+    assert consolidated_count == 0
+    assert raising_client.calls == ["call"]
+    assert store.read_state(session, source_id=DEFAULT_MEMORY_SOURCE_ID) == pre_state
+    surviving = store.get_topic(session, source_id=DEFAULT_MEMORY_SOURCE_ID, topic="preferences")
+    assert surviving.summary == "Prefers concise answers."
+
+
+async def test_memory_extraction_propagates_programmer_errors(tmp_path) -> None:
+    """Non-transient errors from the chat client must surface so misconfigurations fail loudly."""
+    session = AgentSession(session_id="session-1")
+    session.state["owner_id"] = "alice"
+    store = MemoryFileStore(tmp_path, owner_state_key="owner_id")
+    provider = MemoryContextProvider(store=store)
+    bad_client = _ProgrammerErrorMemoryClient()
+
+    from agent_framework import AgentResponse
+    from agent_framework._sessions import SessionContext
+
+    context = SessionContext(
+        input_messages=[Message(role="user", contents=["q"])],
+    )
+    context._response = AgentResponse(messages=[Message(role="assistant", contents=["a"])])  # pyright: ignore[reportPrivateUsage]
+
+    with pytest.raises(AttributeError, match="misconfigured client"):
+        await provider._extract_memories(  # pyright: ignore[reportPrivateUsage]
+            client=bad_client,  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+            session=session,
+            context=context,
+            now=datetime(2026, 4, 22, tzinfo=timezone.utc),
+        )

python/packages/core/tests/core/test_harness_mode.py

@@ -0,0 +1,303 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import json
+
+import pytest
+
+from agent_framework import (
+    DEFAULT_MODE_SOURCE_ID,
+    Agent,
+    AgentModeProvider,
+    AgentSession,
+    ExperimentalFeature,
+    Message,
+    SupportsChatGetResponse,
+    get_agent_mode,
+    set_agent_mode,
+)
+
+
+def _tool_by_name(tools: list[object], name: str) -> object:
+    """Return the tool with the requested name from a prepared tool list."""
+    for tool in tools:
+        if getattr(tool, "name", None) == name:
+            return tool
+    raise AssertionError(f"Tool {name!r} was not found.")
+
+
+def test_get_and_set_agent_mode_manage_session_state() -> None:
+    """Mode helpers should initialize session state, normalize values, and validate modes."""
+    session = AgentSession(session_id="session-1")
+
+    assert get_agent_mode(session) == "plan"
+    assert session.state[DEFAULT_MODE_SOURCE_ID] == {"current_mode": "plan"}
+    assert set_agent_mode(session, " execute ") == "execute"
+    assert get_agent_mode(session) == "execute"
+
+    custom_session = AgentSession(session_id="session-2")
+    assert (
+        get_agent_mode(
+            custom_session,
+            default_mode="draft",
+            available_modes=("draft", "final"),
+        )
+        == "draft"
+    )
+
+    with pytest.raises(ValueError, match="Invalid mode"):
+        set_agent_mode(session, "ship")
+
+
+def test_agent_mode_helpers_reject_non_dict_provider_state() -> None:
+    """Mode helpers should not overwrite unrelated non-dict session state."""
+    session = AgentSession(session_id="session-1")
+    session.state[DEFAULT_MODE_SOURCE_ID] = "unrelated state"
+
+    with pytest.raises(TypeError, match="source_id 'agent_mode'.*str"):
+        get_agent_mode(session)
+
+    assert session.state[DEFAULT_MODE_SOURCE_ID] == "unrelated state"
+
+
+def test_agent_mode_context_provider_validates_configuration_and_is_experimental() -> None:
+    """Mode provider should validate configuration and expose HARNESS experimental metadata."""
+    with pytest.raises(ValueError, match="at least one mode"):
+        AgentModeProvider(mode_descriptions={})
+
+    with pytest.raises(ValueError, match="Invalid mode"):
+        AgentModeProvider(default_mode="ship")
+
+    assert AgentModeProvider.__feature_id__ == ExperimentalFeature.HARNESS.value  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert get_agent_mode.__feature_id__ == ExperimentalFeature.HARNESS.value  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert set_agent_mode.__feature_id__ == ExperimentalFeature.HARNESS.value  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert ".. warning:: Experimental" in AgentModeProvider.__doc__  # type: ignore[operator]  # pyrefly: ignore[not-iterable]  # ty: ignore[unsupported-operator]
+    assert get_agent_mode.__doc__ is not None
+    assert ".. warning:: Experimental" in get_agent_mode.__doc__
+    assert set_agent_mode.__doc__ is not None
+    assert ".. warning:: Experimental" in set_agent_mode.__doc__
+
+
+async def test_external_read_with_provider_config_preserves_nondefault_mode(
+    chat_client_base: SupportsChatGetResponse,
+) -> None:
+    """A pre-run external mode read must honor the provider's configured default, not the built-in one.
+
+    Regression test for the harness console bug where ``configure_run_options`` read the mode with a
+    bare ``get_agent_mode(session)`` before the agent ran. Because ``get_agent_mode`` persists the
+    resolved default into session state, the built-in ``plan`` default was stored and the provider —
+    configured with ``default_mode="execute"`` — then read back ``plan``, so the agent ran in plan
+    mode while the console showed execute. Threading the provider's configuration into the read keeps
+    the two in sync.
+    """
+    provider = AgentModeProvider(default_mode="execute")
+
+    # A bare read (the original buggy call) would resolve and persist the built-in ``plan`` default,
+    # which does not match the provider's configured ``execute`` default.
+    poisoned_session = AgentSession(session_id="poisoned")
+    assert get_agent_mode(poisoned_session) == "plan"
+    agent = Agent(client=chat_client_base, context_providers=[provider])
+    _, poisoned_options = await agent._prepare_session_and_messages(  # pyright: ignore[reportPrivateUsage]
+        session=poisoned_session,
+        input_messages=[Message(role="user", contents=["Go"])],
+    )
+    assert "You are currently operating in the plan mode." in poisoned_options["instructions"]
+
+    # Reading with the provider's own configuration resolves and persists ``execute``, so the
+    # provider injects execute-mode instructions on the run.
+    session = AgentSession(session_id="configured")
+    assert (
+        get_agent_mode(
+            session,
+            source_id=provider.source_id,
+            default_mode=provider.default_mode,
+            available_modes=provider.available_modes,
+        )
+        == "execute"
+    )
+    _, options = await agent._prepare_session_and_messages(  # pyright: ignore[reportPrivateUsage]
+        session=session,
+        input_messages=[Message(role="user", contents=["Go"])],
+    )
+    assert "You are currently operating in the execute mode." in options["instructions"]
+
+
+async def test_agent_mode_context_provider_normalizes_custom_modes(
+    chat_client_base: SupportsChatGetResponse,
+) -> None:
+    """Mode provider should accept differently-cased custom modes and display configured names."""
+    session = AgentSession(session_id="session-1")
+    provider = AgentModeProvider(
+        default_mode="Draft", mode_descriptions={"Draft": "Draft it.", "Final": "Finalize it."}
+    )
+    agent = Agent(client=chat_client_base, context_providers=[provider])
+
+    _, options = await agent._prepare_session_and_messages(  # pyright: ignore[reportPrivateUsage]
+        session=session,
+        input_messages=[Message(role="user", contents=["Start drafting"])],
+    )
+    instructions = options["instructions"]
+    assert isinstance(instructions, str)
+    assert "#### Draft" in instructions
+    assert "Draft it." in instructions
+    assert "#### Final" in instructions
+    assert "Finalize it." in instructions
+    assert "You are currently operating in the draft mode." in instructions
+
+    assert (
+        get_agent_mode(session, source_id=provider.source_id, default_mode="Draft", available_modes=("Draft", "Final"))
+        == "draft"
+    )
+    assert set_agent_mode(session, "draft", source_id=provider.source_id, available_modes=("Draft", "Final")) == "draft"
+    assert (
+        get_agent_mode(session, source_id=provider.source_id, default_mode="Draft", available_modes=("Draft", "Final"))
+        == "draft"
+    )
+
+
+async def test_agent_mode_context_provider_serializes_tool_outputs_as_json(
+    chat_client_base: SupportsChatGetResponse,
+) -> None:
+    """Mode tools should serialize JSON correctly for mode names with quotes."""
+    session = AgentSession(session_id="session-1")
+    mode_name = 'edit "preview"'
+    provider = AgentModeProvider(default_mode=mode_name, mode_descriptions={mode_name: "Preview edits."})
+    agent = Agent(client=chat_client_base, context_providers=[provider])
+
+    _, options = await agent._prepare_session_and_messages(  # pyright: ignore[reportPrivateUsage]
+        session=session,
+        input_messages=[Message(role="user", contents=["Preview edits"])],
+    )
+    tools = options["tools"]
+    assert isinstance(tools, list)
+    get_mode_tool = _tool_by_name(tools, "mode_get")
+    set_mode_tool = _tool_by_name(tools, "mode_set")
+
+    initial_mode = await get_mode_tool.invoke()  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert json.loads(initial_mode[0].text) == {"mode": mode_name}
+
+    set_result = await set_mode_tool.invoke(arguments={"mode": mode_name})  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert json.loads(set_result[0].text) == {"mode": mode_name, "message": f"Mode changed to '{mode_name}'."}
+
+
+async def test_agent_mode_context_provider_updates_agent_mode(
+    chat_client_base: SupportsChatGetResponse,
+) -> None:
+    """Mode provider tools should read and write session-backed mode state."""
+    session = AgentSession(session_id="session-1")
+    provider = AgentModeProvider()
+    agent = Agent(client=chat_client_base, context_providers=[provider])
+
+    _, options = await agent._prepare_session_and_messages(  # pyright: ignore[reportPrivateUsage]
+        session=session,
+        input_messages=[Message(role="user", contents=["Start planning"])],
+    )
+    tools = options["tools"]
+    assert isinstance(tools, list)
+    instructions = options["instructions"]
+    assert isinstance(instructions, str)
+    assert "## Agent Mode" in instructions
+    assert "Use the mode_set tool to switch between modes as your work progresses." in instructions
+    assert "ask clarifying questions, discuss options, and get user approval before proceeding" in instructions
+    assert "If you encounter ambiguity" in instructions
+    assert "You are currently operating in the plan mode." in instructions
+
+    get_mode_tool = _tool_by_name(tools, "mode_get")
+    set_mode_tool = _tool_by_name(tools, "mode_set")
+
+    initial_mode = await get_mode_tool.invoke()  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert json.loads(initial_mode[0].text) == {"mode": "plan"}
+
+    set_result = await set_mode_tool.invoke(arguments={"mode": "execute"})  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert json.loads(set_result[0].text) == {"mode": "execute", "message": "Mode changed to 'execute'."}
+    assert get_agent_mode(session, source_id=provider.source_id) == "execute"
+    assert set_agent_mode(session, "plan", source_id=provider.source_id) == "plan"
+
+
+def test_default_mode_falls_back_to_first_available_mode() -> None:
+    """When ``default_mode`` is omitted, helpers and provider should use the first configured mode."""
+    session = AgentSession(session_id="session-1")
+
+    assert get_agent_mode(session, available_modes=("draft", "final")) == "draft"
+
+    provider = AgentModeProvider(mode_descriptions={"Draft": "Draft it.", "Final": "Finalize it."})
+    assert provider.default_mode == "draft"
+
+
+def test_get_agent_mode_falls_back_when_stored_mode_not_in_available_modes() -> None:
+    """A previously persisted mode that is no longer configured should be reset to the default."""
+    session = AgentSession(session_id="session-1")
+    set_agent_mode(session, "execute")
+    assert session.state[DEFAULT_MODE_SOURCE_ID]["current_mode"] == "execute"
+
+    # Reconfigure with a smaller mode set that no longer includes "execute".
+    current = get_agent_mode(session, default_mode="draft", available_modes=("draft", "final"))
+    assert current == "draft"
+    assert session.state[DEFAULT_MODE_SOURCE_ID]["current_mode"] == "draft"
+
+
+def test_set_agent_mode_records_previous_mode_for_external_change_notification() -> None:
+    """External mode changes via ``set_agent_mode`` should record the previous mode for notification."""
+    session = AgentSession(session_id="session-1")
+    set_agent_mode(session, "plan")
+    set_agent_mode(session, "execute")
+
+    assert session.state[DEFAULT_MODE_SOURCE_ID]["current_mode"] == "execute"
+    assert session.state[DEFAULT_MODE_SOURCE_ID]["previous_mode_for_notification"] == "plan"
+
+
+def test_set_agent_mode_no_op_does_not_record_previous_mode() -> None:
+    """Setting the same mode should not queue a notification."""
+    session = AgentSession(session_id="session-1")
+    set_agent_mode(session, "plan")
+    set_agent_mode(session, "plan")
+
+    assert "previous_mode_for_notification" not in session.state[DEFAULT_MODE_SOURCE_ID]
+
+
+async def test_agent_mode_provider_injects_user_message_after_external_change(
+    chat_client_base: SupportsChatGetResponse,
+) -> None:
+    """External mode changes should inject a user message announcing the switch on the next run."""
+    session = AgentSession(session_id="session-1")
+    provider = AgentModeProvider()
+    agent = Agent(client=chat_client_base, context_providers=[provider])
+
+    # First run: agent uses mode_set tool to switch to execute. The tool path must NOT queue a
+    # notification because the agent already saw its own tool call in the chat history.
+    _, first_options = await agent._prepare_session_and_messages(  # pyright: ignore[reportPrivateUsage]
+        session=session,
+        input_messages=[Message(role="user", contents=["Plan first."])],
+    )
+    set_mode_tool = _tool_by_name(first_options["tools"], "mode_set")
+    await set_mode_tool.invoke(arguments={"mode": "execute"})  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert "previous_mode_for_notification" not in session.state[provider.source_id]
+
+    # Now an external caller (e.g., a /mode slash command) switches the mode back to plan.
+    set_agent_mode(session, "plan", source_id=provider.source_id)
+    assert session.state[provider.source_id]["previous_mode_for_notification"] == "execute"
+
+    # Next run: the provider should inject a user message announcing the change and clear the flag.
+    second_context, second_options = await agent._prepare_session_and_messages(  # pyright: ignore[reportPrivateUsage]
+        session=session,
+        input_messages=[Message(role="user", contents=["Carry on."])],
+    )
+    instructions = second_options["instructions"]
+    assert isinstance(instructions, str)
+    assert "You are currently operating in the plan mode." in instructions
+
+    notification_messages = [message for message in second_context.context_messages.get(provider.source_id, [])]
+    assert len(notification_messages) == 1
+    assert notification_messages[0].role == "user"
+    assert "Mode changed" in notification_messages[0].text
+    assert '"execute"' in notification_messages[0].text
+    assert '"plan"' in notification_messages[0].text
+    assert "previous_mode_for_notification" not in session.state[provider.source_id]
+
+    # Third run with no further external change must not re-inject the notification.
+    third_context, _ = await agent._prepare_session_and_messages(  # pyright: ignore[reportPrivateUsage]
+        session=session,
+        input_messages=[Message(role="user", contents=["Status?"])],
+    )
+    assert third_context.context_messages.get(provider.source_id, []) == []

python/packages/core/tests/core/test_harness_todo.py

@@ -0,0 +1,377 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+from pathlib import Path
+
+import pytest
+
+from agent_framework import (
+    Agent,
+    AgentSession,
+    ExperimentalFeature,
+    Message,
+    SupportsChatGetResponse,
+    TodoFileStore,
+    TodoInput,
+    TodoItem,
+    TodoProvider,
+    TodoSessionStore,
+    TodoStore,
+)
+
+
+def _tool_by_name(tools: list[object], name: str) -> object:
+    """Return the tool with the requested name from a prepared tool list."""
+    for tool in tools:
+        if getattr(tool, "name", None) == name:
+            return tool
+    raise AssertionError(f"Tool {name!r} was not found.")
+
+
+def test_todo_item_round_trips_with_value_equality() -> None:
+    """Todo items should support value equality and JSON serialization."""
+    raw_item = {
+        "id": 1,
+        "title": "Write tests",
+        "description": "Cover the harness",
+        "is_complete": False,
+    }
+
+    item = TodoItem.from_dict(raw_item)
+
+    assert item == TodoItem(**raw_item)  # type: ignore[arg-type]
+    assert item.to_dict() == raw_item
+    assert json.loads(item.to_json()) == raw_item
+    assert "TodoItem(" in repr(item)
+
+
+def test_todo_input_round_trips_and_validates() -> None:
+    """Todo input should trim titles and reject invalid payloads."""
+    todo_input = TodoInput.from_dict({"title": "  Write tests  ", "description": "Cover the harness"})
+
+    assert todo_input.title == "Write tests"
+    assert todo_input.to_dict() == {"title": "Write tests", "description": "Cover the harness"}
+    assert json.loads(todo_input.to_json()) == {"title": "Write tests", "description": "Cover the harness"}
+
+    with pytest.raises(ValueError, match="non-empty string"):
+        TodoInput(title="   ")
+
+    with pytest.raises(ValueError, match="description must be a string or null"):
+        TodoInput.from_dict({"title": "Write tests", "description": 123})
+
+
+async def test_todo_session_store_initializes_and_round_trips_state() -> None:
+    """Session-backed todo storage should initialize and persist todo state."""
+    session = AgentSession(session_id="session-1")
+    store = TodoSessionStore()
+
+    items, next_id = await store.load_state(session, source_id="todo")
+    assert items == []
+    assert next_id == 1
+    assert session.state["todo"] == {}
+
+    todo_item = TodoItem(id=1, title="Ship feature", description="Use session storage")
+    await store.save_state(session, [todo_item], next_id=2, source_id="todo")
+
+    loaded_items, loaded_next_id = await store.load_state(session, source_id="todo")
+    assert loaded_items == [todo_item]
+    assert loaded_next_id == 2
+    assert await store.load_items(session, source_id="todo") == [todo_item]
+
+
+async def test_todo_file_store_round_trips_state(tmp_path: Path) -> None:
+    """Todo file storage should persist one JSON state file per owner and session."""
+    session = AgentSession(session_id="session-1")
+    session.state["owner_id"] = "alice"
+    store = TodoFileStore(
+        tmp_path,
+        kind="todos",
+        owner_prefix="user_",
+        owner_state_key="owner_id",
+    )
+
+    await store.save_state(
+        session,
+        [TodoItem(id=1, title="Ship feature", description="Use file storage")],
+        next_id=2,
+        source_id="todo",
+    )
+
+    items, next_id = await store.load_state(session, source_id="todo")
+    assert items == [TodoItem(id=1, title="Ship feature", description="Use file storage", is_complete=False)]
+    assert next_id == 2
+
+    state_path = tmp_path / "user_alice" / "todos" / "session-1" / "todos.todo.json"
+    assert state_path.exists()
+    assert json.loads(state_path.read_text(encoding="utf-8")) == {
+        "items": [{"id": 1, "title": "Ship feature", "description": "Use file storage", "is_complete": False}],
+        "next_id": 2,
+    }
+
+    with pytest.raises(RuntimeError, match="owner_id"):
+        await store.load_state(AgentSession(session_id="missing-owner"), source_id="todo")
+
+
+async def test_todo_file_store_load_does_not_create_directories(tmp_path: Path) -> None:
+    """Loading from a never-written session must not create empty directories on disk."""
+    session = AgentSession(session_id="session-1")
+    store = TodoFileStore(tmp_path)
+
+    items, next_id = await store.load_state(session, source_id="todo")
+    assert items == []
+    assert next_id == 1
+    assert list(tmp_path.iterdir()) == []  # noqa: ASYNC240
+
+
+async def test_todo_file_store_writes_state_atomically(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    """A crash between writing the temp file and renaming must not corrupt existing state."""
+    session = AgentSession(session_id="session-1")
+    store = TodoFileStore(tmp_path)
+
+    await store.save_state(session, [TodoItem(id=1, title="Initial")], next_id=2, source_id="todo")
+    state_path = tmp_path / "session-1" / "todos.todo.json"
+    original_contents = state_path.read_text(encoding="utf-8")
+
+    def _boom(*args: object, **kwargs: object) -> None:
+        raise OSError("disk full")
+
+    monkeypatch.setattr(os, "replace", _boom)
+
+    with pytest.raises(OSError, match="disk full"):
+        await store.save_state(session, [TodoItem(id=2, title="Replacement")], next_id=3, source_id="todo")
+
+    # Original file is untouched, no temp leftovers.
+    assert state_path.read_text(encoding="utf-8") == original_contents
+    assert sorted(p.name for p in state_path.parent.iterdir()) == [state_path.name]
+
+
+async def test_todo_session_store_rejects_non_mapping_items() -> None:
+    """Session-backed todo storage should report malformed item entries clearly."""
+    session = AgentSession(session_id="session-1")
+    session.state["todo"] = {"items": [{"id": 1, "title": "Good"}, "bad"], "next_id": 2}
+    store = TodoSessionStore()
+
+    with pytest.raises(ValueError, match="index 1.*str"):
+        await store.load_state(session, source_id="todo")
+
+
+async def test_todo_session_store_rejects_malformed_state_types() -> None:
+    """Session-backed todo storage should raise for malformed top-level state, mirroring TodoFileStore."""
+    session = AgentSession(session_id="session-1")
+    session.state["todo"] = "not a dict"
+    store = TodoSessionStore()
+
+    with pytest.raises(ValueError, match="must be a dict"):
+        await store.load_state(session, source_id="todo")
+
+    session.state["todo"] = {"items": "not a list", "next_id": 1}
+    with pytest.raises(ValueError, match="non-list 'items'"):
+        await store.load_state(session, source_id="todo")
+
+    session.state["todo"] = {"items": [], "next_id": "1"}
+    with pytest.raises(ValueError, match="non-integer 'next_id'"):
+        await store.load_state(session, source_id="todo")
+
+
+async def test_todo_stores_clamp_next_id_to_avoid_collisions(tmp_path: Path) -> None:
+    """Both stores should clamp ``next_id`` to ``max(item.id) + 1`` to prevent ID collisions."""
+    session_a = AgentSession(session_id="session-a")
+    session_a.state["todo"] = {"items": [{"id": 5, "title": "Seeded"}], "next_id": 1}
+
+    session_store = TodoSessionStore()
+    items, next_id = await session_store.load_state(session_a, source_id="todo")
+    assert next_id == 6  # clamped over the stored next_id of 1
+    assert items == [TodoItem(id=5, title="Seeded")]
+
+    session_b = AgentSession(session_id="session-b")
+    file_store = TodoFileStore(tmp_path)
+    state_path = tmp_path / "session-b" / "todos.todo.json"
+    state_path.parent.mkdir(parents=True)
+    state_path.write_text(json.dumps({"items": [{"id": 7, "title": "Seeded"}], "next_id": 1}) + "\n", encoding="utf-8")
+    items, next_id = await file_store.load_state(session_b, source_id="todo")
+    assert next_id == 8
+    assert items == [TodoItem(id=7, title="Seeded")]
+
+
+async def test_todo_provider_evicts_locks_when_session_is_garbage_collected() -> None:
+    """The provider should not retain mutation locks for sessions that have been GC'd."""
+    import gc
+
+    provider = TodoProvider()
+    session = AgentSession(session_id="session-1")
+    provider._mutation_lock(session)  # pyright: ignore[reportPrivateUsage]
+    assert len(provider._mutation_locks) == 1  # pyright: ignore[reportPrivateUsage]
+
+    del session
+    gc.collect()
+    assert len(provider._mutation_locks) == 0  # pyright: ignore[reportPrivateUsage]
+
+
+async def test_todo_file_store_rejects_session_path_traversal(tmp_path: Path) -> None:
+    """File-backed todo storage should not write outside its base path for malicious session IDs."""
+    session = AgentSession(session_id="../escape")
+    store = TodoFileStore(tmp_path)
+
+    with pytest.raises(ValueError, match="session_id.*path separators"):
+        await store.save_state(session, [TodoItem(id=1, title="Escape")], next_id=2, source_id="todo")
+
+    assert list(tmp_path.rglob("*")) == []  # noqa: ASYNC240
+
+
+async def test_todo_file_store_namespaces_state_by_source_id(tmp_path: Path) -> None:
+    """File-backed todo storage should isolate providers that share a session."""
+    session = AgentSession(session_id="session-1")
+    store = TodoFileStore(tmp_path)
+
+    await store.save_state(session, [TodoItem(id=1, title="First source")], next_id=2, source_id="first")
+    await store.save_state(session, [TodoItem(id=1, title="Second source")], next_id=2, source_id="second")
+
+    first_items, _ = await store.load_state(session, source_id="first")
+    second_items, _ = await store.load_state(session, source_id="second")
+
+    assert first_items == [TodoItem(id=1, title="First source")]
+    assert second_items == [TodoItem(id=1, title="Second source")]
+    assert (tmp_path / "session-1" / "todos.first.json").exists()
+    assert (tmp_path / "session-1" / "todos.second.json").exists()
+
+
+async def test_todo_provider_runs_with_file_store(tmp_path: Path, chat_client_base: SupportsChatGetResponse) -> None:
+    """The provider should drive the full add/list flow when backed by ``TodoFileStore``."""
+    session = AgentSession(session_id="session-1")
+    provider = TodoProvider(store=TodoFileStore(tmp_path))
+    agent = Agent(client=chat_client_base, context_providers=[provider])
+
+    _, options = await agent._prepare_session_and_messages(  # pyright: ignore[reportPrivateUsage]
+        session=session,
+        input_messages=[Message(role="user", contents=["Track this work"])],
+    )
+    tools = options["tools"]
+    assert isinstance(tools, list)
+
+    add_todos = _tool_by_name(tools, "todos_add")
+    get_all_todos = _tool_by_name(tools, "todos_get_all")
+
+    await add_todos.invoke(arguments={"todos": [{"title": "Persist me"}]})  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    state_path = tmp_path / "session-1" / "todos.todo.json"
+    assert state_path.exists()
+    persisted = json.loads(state_path.read_text(encoding="utf-8"))
+    assert persisted["items"] == [{"id": 1, "title": "Persist me", "description": None, "is_complete": False}]
+    assert persisted["next_id"] == 2
+
+    get_all_result = await get_all_todos.invoke()  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert json.loads(get_all_result[0].text) == [
+        {"id": 1, "title": "Persist me", "description": None, "is_complete": False}
+    ]
+
+
+async def test_todo_provider_tools_manage_session_state(
+    chat_client_base: SupportsChatGetResponse,
+) -> None:
+    """Todo provider tools should add, complete, remove, and list session-backed todos."""
+    session = AgentSession(session_id="session-1")
+    provider = TodoProvider()
+    agent = Agent(client=chat_client_base, context_providers=[provider])
+
+    _, options = await agent._prepare_session_and_messages(  # pyright: ignore[reportPrivateUsage]
+        session=session,
+        input_messages=[Message(role="user", contents=["Track this work"])],
+    )
+    tools = options["tools"]
+    assert isinstance(tools, list)
+
+    add_todos = _tool_by_name(tools, "todos_add")
+    complete_todos = _tool_by_name(tools, "todos_complete")
+    remove_todos = _tool_by_name(tools, "todos_remove")
+    get_remaining_todos = _tool_by_name(tools, "todos_get_remaining")
+    get_all_todos = _tool_by_name(tools, "todos_get_all")
+
+    add_result = await add_todos.invoke(  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+        arguments={
+            "todos": [
+                {"title": "  Write tests  ", "description": "  Cover stores  "},
+                {"title": "Ship feature"},
+            ]
+        }
+    )
+    assert json.loads(add_result[0].text) == [
+        {"id": 1, "title": "Write tests", "description": "Cover stores", "is_complete": False},
+        {"id": 2, "title": "Ship feature", "description": None, "is_complete": False},
+    ]
+
+    complete_result = await complete_todos.invoke(arguments={"items": [{"id": 1, "reason": "Tests written"}]})  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert json.loads(complete_result[0].text) == {"completed": 1}
+
+    remaining_result = await get_remaining_todos.invoke()  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert json.loads(remaining_result[0].text) == [
+        {"id": 2, "title": "Ship feature", "description": None, "is_complete": False}
+    ]
+
+    remove_result = await remove_todos.invoke(arguments={"ids": [2]})  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert json.loads(remove_result[0].text) == {"removed": 1}
+
+    get_all_result = await get_all_todos.invoke()  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert json.loads(get_all_result[0].text) == [
+        {"id": 1, "title": "Write tests", "description": "Cover stores", "is_complete": True}
+    ]
+
+
+async def test_todo_provider_serializes_concurrent_mutations(
+    chat_client_base: SupportsChatGetResponse,
+) -> None:
+    """Concurrent todo mutations should not duplicate IDs or lose updates."""
+    session = AgentSession(session_id="session-1")
+    provider = TodoProvider()
+    agent = Agent(client=chat_client_base, context_providers=[provider])
+
+    _, options = await agent._prepare_session_and_messages(  # pyright: ignore[reportPrivateUsage]
+        session=session,
+        input_messages=[Message(role="user", contents=["Track this work"])],
+    )
+    tools = options["tools"]
+    assert isinstance(tools, list)
+
+    add_todos = _tool_by_name(tools, "todos_add")
+    complete_todos = _tool_by_name(tools, "todos_complete")
+    get_all_todos = _tool_by_name(tools, "todos_get_all")
+
+    await add_todos.invoke(arguments={"todos": [{"title": f"Existing {index}"} for index in range(1, 6)]})  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+
+    await asyncio.gather(
+        add_todos.invoke(arguments={"todos": [{"title": "Add A1"}, {"title": "Add A2"}]}),  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+        add_todos.invoke(arguments={"todos": [{"title": "Add B1"}, {"title": "Add B2"}]}),  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+        complete_todos.invoke(arguments={"items": [{"id": i, "reason": "Done"} for i in range(1, 6)]}),  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    )
+
+    get_all_result = await get_all_todos.invoke()  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    payload = json.loads(get_all_result[0].text)
+    ids = [item["id"] for item in payload]
+
+    assert sorted(ids) == list(range(1, 10))
+    assert len(ids) == len(set(ids))
+    assert {item["title"] for item in payload} == {
+        "Existing 1",
+        "Existing 2",
+        "Existing 3",
+        "Existing 4",
+        "Existing 5",
+        "Add A1",
+        "Add A2",
+        "Add B1",
+        "Add B2",
+    }
+    assert {item["id"] for item in payload if item["is_complete"]} == {1, 2, 3, 4, 5}
+
+
+def test_todo_harness_classes_are_marked_experimental() -> None:
+    """Todo harness public classes should expose HARNESS experimental metadata."""
+    assert TodoStore.__feature_id__ == ExperimentalFeature.HARNESS.value  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert TodoItem.__feature_id__ == ExperimentalFeature.HARNESS.value  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert TodoInput.__feature_id__ == ExperimentalFeature.HARNESS.value  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert TodoSessionStore.__feature_id__ == ExperimentalFeature.HARNESS.value  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert TodoFileStore.__feature_id__ == ExperimentalFeature.HARNESS.value  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert TodoProvider.__feature_id__ == ExperimentalFeature.HARNESS.value  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    assert ".. warning:: Experimental" in TodoProvider.__doc__  # type: ignore[operator]  # pyrefly: ignore[not-iterable]  # ty: ignore[unsupported-operator]

python/packages/core/tests/core/test_harness_tool_approval.py

@@ -0,0 +1,823 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from agent_framework import (
+    DEFAULT_TOOL_APPROVAL_SOURCE_ID,
+    Agent,
+    AgentSession,
+    ChatResponse,
+    ChatResponseUpdate,
+    Content,
+    Message,
+    ToolApprovalMiddleware,
+    ToolApprovalState,
+    create_always_approve_tool_response,
+    create_always_approve_tool_with_arguments_response,
+    tool,
+)
+
+from .conftest import MockBaseChatClient
+
+
+def _approval_requests(messages: list[Message]) -> list[Content]:
+    return [
+        content for message in messages for content in message.contents if content.type == "function_approval_request"
+    ]
+
+
+def _function_call(request: Content) -> Content:
+    assert request.function_call is not None
+    return request.function_call
+
+
+async def test_mixed_batch_hides_already_approved_request_until_approval_replay(
+    chat_client_base: MockBaseChatClient,
+) -> None:
+    """Mixed batches should only show real approval requests when a session can store hidden requests."""
+    no_approval_calls = 0
+    approval_calls = 0
+
+    @tool(name="lookup_work_items", approval_mode="never_require")
+    def lookup_work_items(query: str) -> str:
+        nonlocal no_approval_calls
+        no_approval_calls += 1
+        return f"found {query}"
+
+    @tool(name="add_comment", approval_mode="always_require")
+    def add_comment(comment: str) -> str:
+        nonlocal approval_calls
+        approval_calls += 1
+        return f"added {comment}"
+
+    agent = Agent(client=chat_client_base, tools=[lookup_work_items, add_comment])
+    session = AgentSession(session_id="approval-session")
+    chat_client_base.run_responses = [
+        ChatResponse(
+            messages=Message(
+                role="assistant",
+                contents=[
+                    Content.from_function_call(
+                        call_id="call_lookup",
+                        name="lookup_work_items",
+                        arguments='{"query": "mine"}',
+                    ),
+                    Content.from_function_call(
+                        call_id="call_comment",
+                        name="add_comment",
+                        arguments='{"comment": "done"}',
+                    ),
+                ],
+            )
+        )
+    ]
+
+    first_response = await agent.run("update work item", session=session)
+
+    requests = _approval_requests(first_response.messages)
+    assert [_function_call(request).name for request in requests] == ["add_comment"]
+    assert no_approval_calls == 0
+    assert approval_calls == 0
+
+    chat_client_base.run_responses = [ChatResponse(messages=Message(role="assistant", contents=["complete"]))]
+    second_response = await agent.run(requests[0].to_function_approval_response(approved=True), session=session)
+
+    assert second_response.text == "complete"
+    assert no_approval_calls == 1
+    assert approval_calls == 1
+
+
+async def test_mixed_batch_accepts_restored_tool_approval_state(
+    chat_client_base: MockBaseChatClient,
+) -> None:
+    """Mixed-batch bypass should work when session state contains ToolApprovalState."""
+    safe_calls = 0
+    risky_calls = 0
+
+    @tool(name="safe_read", approval_mode="never_require")
+    def safe_read() -> str:
+        nonlocal safe_calls
+        safe_calls += 1
+        return "safe"
+
+    @tool(name="risky_write", approval_mode="always_require")
+    def risky_write() -> str:
+        nonlocal risky_calls
+        risky_calls += 1
+        return "risky"
+
+    agent = Agent(client=chat_client_base, tools=[safe_read, risky_write])
+    session = AgentSession(session_id="restored-state-session")
+    session.state[DEFAULT_TOOL_APPROVAL_SOURCE_ID] = ToolApprovalState()
+    chat_client_base.run_responses = [
+        ChatResponse(
+            messages=Message(
+                role="assistant",
+                contents=[
+                    Content.from_function_call(call_id="call_safe", name="safe_read", arguments="{}"),
+                    Content.from_function_call(call_id="call_risky", name="risky_write", arguments="{}"),
+                ],
+            )
+        )
+    ]
+
+    first_response = await agent.run("read and write", session=session)
+    requests = _approval_requests(first_response.messages)
+
+    assert [_function_call(request).name for request in requests] == ["risky_write"]
+    assert safe_calls == 0
+    assert risky_calls == 0
+
+    chat_client_base.run_responses = [ChatResponse(messages=Message(role="assistant", contents=["done"]))]
+    final_response = await agent.run(requests[0].to_function_approval_response(approved=True), session=session)
+
+    assert final_response.text == "done"
+    assert safe_calls == 1
+    assert risky_calls == 1
+
+
+async def test_hidden_mixed_batch_requests_do_not_replay_on_unrelated_turn(
+    chat_client_base: MockBaseChatClient,
+) -> None:
+    """Stored hidden approvals should only replay when an approval response resumes the flow."""
+    safe_calls = 0
+    risky_calls = 0
+
+    @tool(name="safe_lookup", approval_mode="never_require")
+    def safe_lookup() -> str:
+        nonlocal safe_calls
+        safe_calls += 1
+        return "safe"
+
+    @tool(name="risky_update", approval_mode="always_require")
+    def risky_update() -> str:
+        nonlocal risky_calls
+        risky_calls += 1
+        return "risky"
+
+    agent = Agent(client=chat_client_base, tools=[safe_lookup, risky_update])
+    session = AgentSession(session_id="stale-hidden-session")
+    chat_client_base.run_responses = [
+        ChatResponse(
+            messages=Message(
+                role="assistant",
+                contents=[
+                    Content.from_function_call(call_id="call_safe", name="safe_lookup", arguments="{}"),
+                    Content.from_function_call(call_id="call_risky", name="risky_update", arguments="{}"),
+                ],
+            )
+        )
+    ]
+
+    first_response = await agent.run("lookup and update", session=session)
+    request = _approval_requests(first_response.messages)[0]
+
+    chat_client_base.run_responses = [ChatResponse(messages=Message(role="assistant", contents=["unrelated"]))]
+    unrelated_response = await agent.run("never mind, answer something else", session=session)
+
+    assert unrelated_response.text == "unrelated"
+    assert safe_calls == 0
+    assert risky_calls == 0
+
+    chat_client_base.run_responses = [ChatResponse(messages=Message(role="assistant", contents=["done"]))]
+    final_response = await agent.run(request.to_function_approval_response(approved=True), session=session)
+
+    assert final_response.text == "done"
+    assert safe_calls == 1
+    assert risky_calls == 1
+
+
+async def test_hidden_mixed_batch_requests_replay_only_for_matching_visible_approval(
+    chat_client_base: MockBaseChatClient,
+) -> None:
+    """Approving one mixed batch must not replay hidden calls from another abandoned batch."""
+    safe_a_calls = 0
+    safe_b_calls = 0
+    risky_a_calls = 0
+    risky_b_calls = 0
+
+    @tool(name="safe_a", approval_mode="never_require")
+    def safe_a() -> str:
+        nonlocal safe_a_calls
+        safe_a_calls += 1
+        return "safe-a"
+
+    @tool(name="safe_b", approval_mode="never_require")
+    def safe_b() -> str:
+        nonlocal safe_b_calls
+        safe_b_calls += 1
+        return "safe-b"
+
+    @tool(name="risky_a", approval_mode="always_require")
+    def risky_a() -> str:
+        nonlocal risky_a_calls
+        risky_a_calls += 1
+        return "risky-a"
+
+    @tool(name="risky_b", approval_mode="always_require")
+    def risky_b() -> str:
+        nonlocal risky_b_calls
+        risky_b_calls += 1
+        return "risky-b"
+
+    agent = Agent(client=chat_client_base, tools=[safe_a, safe_b, risky_a, risky_b])
+    session = AgentSession(session_id="grouped-hidden-session")
+    chat_client_base.run_responses = [
+        ChatResponse(
+            messages=Message(
+                role="assistant",
+                contents=[
+                    Content.from_function_call(call_id="call_safe_a", name="safe_a", arguments="{}"),
+                    Content.from_function_call(call_id="call_risky_a", name="risky_a", arguments="{}"),
+                ],
+            )
+        )
+    ]
+
+    first_response = await agent.run("batch a", session=session)
+    assert [_function_call(request).name for request in _approval_requests(first_response.messages)] == ["risky_a"]
+
+    chat_client_base.run_responses = [
+        ChatResponse(
+            messages=Message(
+                role="assistant",
+                contents=[
+                    Content.from_function_call(call_id="call_safe_b", name="safe_b", arguments="{}"),
+                    Content.from_function_call(call_id="call_risky_b", name="risky_b", arguments="{}"),
+                ],
+            )
+        )
+    ]
+
+    second_response = await agent.run("batch b", session=session)
+    second_request = _approval_requests(second_response.messages)[0]
+
+    chat_client_base.run_responses = [ChatResponse(messages=Message(role="assistant", contents=["done"]))]
+    final_response = await agent.run(second_request.to_function_approval_response(approved=True), session=session)
+
+    assert final_response.text == "done"
+    assert safe_a_calls == 0
+    assert risky_a_calls == 0
+    assert safe_b_calls == 1
+    assert risky_b_calls == 1
+
+
+async def test_tool_approval_middleware_queues_multiple_approval_requests(
+    chat_client_base: MockBaseChatClient,
+) -> None:
+    """The opt-in middleware should present multiple unresolved approvals one at a time."""
+    first_calls = 0
+    second_calls = 0
+
+    @tool(name="first_tool", approval_mode="always_require")
+    def first_tool() -> str:
+        nonlocal first_calls
+        first_calls += 1
+        return "first"
+
+    @tool(name="second_tool", approval_mode="always_require")
+    def second_tool() -> str:
+        nonlocal second_calls
+        second_calls += 1
+        return "second"
+
+    agent = Agent(
+        client=chat_client_base,
+        tools=[first_tool, second_tool],
+        middleware=[ToolApprovalMiddleware()],
+    )
+    session = AgentSession(session_id="queue-session")
+    chat_client_base.run_responses = [
+        ChatResponse(
+            messages=Message(
+                role="assistant",
+                contents=[
+                    Content.from_function_call(call_id="call_first", name="first_tool", arguments="{}"),
+                    Content.from_function_call(call_id="call_second", name="second_tool", arguments="{}"),
+                ],
+            )
+        )
+    ]
+
+    first_response = await agent.run("call both", session=session)
+
+    first_requests = _approval_requests(first_response.messages)
+    assert [_function_call(request).name for request in first_requests] == ["first_tool"]
+    assert first_calls == 0
+    assert second_calls == 0
+
+    second_response = await agent.run(first_requests[0].to_function_approval_response(approved=True), session=session)
+
+    second_requests = _approval_requests(second_response.messages)
+    assert [_function_call(request).name for request in second_requests] == ["second_tool"]
+    assert first_calls == 0
+    assert second_calls == 0
+
+    chat_client_base.run_responses = [ChatResponse(messages=Message(role="assistant", contents=["done"]))]
+    final_response = await agent.run(second_requests[0].to_function_approval_response(approved=True), session=session)
+
+    assert final_response.text == "done"
+    assert first_calls == 1
+    assert second_calls == 1
+
+
+async def test_tool_approval_middleware_preserves_hidden_mixed_batch_requests(
+    chat_client_base: MockBaseChatClient,
+) -> None:
+    """Middleware state saves should not discard core hidden already-approved requests."""
+    lookup_calls = 0
+    write_calls = 0
+
+    @tool(name="lookup_records", approval_mode="never_require")
+    def lookup_records() -> str:
+        nonlocal lookup_calls
+        lookup_calls += 1
+        return "records"
+
+    @tool(name="write_record", approval_mode="always_require")
+    def write_record() -> str:
+        nonlocal write_calls
+        write_calls += 1
+        return "written"
+
+    agent = Agent(
+        client=chat_client_base,
+        tools=[lookup_records, write_record],
+        middleware=[ToolApprovalMiddleware()],
+    )
+    session = AgentSession(session_id="mixed-middleware-session")
+    chat_client_base.run_responses = [
+        ChatResponse(
+            messages=Message(
+                role="assistant",
+                contents=[
+                    Content.from_function_call(call_id="call_lookup", name="lookup_records", arguments="{}"),
+                    Content.from_function_call(call_id="call_write", name="write_record", arguments="{}"),
+                ],
+            )
+        )
+    ]
+
+    first_response = await agent.run("lookup and write", session=session)
+    request = _approval_requests(first_response.messages)[0]
+
+    chat_client_base.run_responses = [ChatResponse(messages=Message(role="assistant", contents=["done"]))]
+    second_response = await agent.run(request.to_function_approval_response(approved=True), session=session)
+
+    assert second_response.text == "done"
+    assert lookup_calls == 1
+    assert write_calls == 1
+
+
+async def test_tool_approval_middleware_auto_approval_rule_receives_function_call(
+    chat_client_base: MockBaseChatClient,
+) -> None:
+    """Heuristic auto-approval callbacks should receive function-call content and approve matching calls."""
+    auto_calls = 0
+    manual_calls = 0
+    seen_calls: list[tuple[str, str | None]] = []
+
+    @tool(name="auto_write", approval_mode="always_require")
+    def auto_write() -> str:
+        nonlocal auto_calls
+        auto_calls += 1
+        return "auto"
+
+    @tool(name="manual_write", approval_mode="always_require")
+    def manual_write() -> str:
+        nonlocal manual_calls
+        manual_calls += 1
+        return "manual"
+
+    async def auto_approve_auto_write(function_call: Content) -> bool:
+        seen_calls.append((function_call.type, function_call.name))
+        return function_call.name == "auto_write"
+
+    agent = Agent(
+        client=chat_client_base,
+        tools=[auto_write, manual_write],
+        middleware=[ToolApprovalMiddleware(auto_approval_rules=[auto_approve_auto_write])],
+    )
+    session = AgentSession(session_id="heuristic-session")
+    chat_client_base.run_responses = [
+        ChatResponse(
+            messages=Message(
+                role="assistant",
+                contents=[
+                    Content.from_function_call(call_id="call_auto", name="auto_write", arguments="{}"),
+                    Content.from_function_call(call_id="call_manual", name="manual_write", arguments="{}"),
+                ],
+            )
+        )
+    ]
+
+    first_response = await agent.run("write both", session=session)
+
+    requests = _approval_requests(first_response.messages)
+    assert [_function_call(request).name for request in requests] == ["manual_write"]
+    assert seen_calls == [("function_call", "auto_write"), ("function_call", "manual_write")]
+    assert auto_calls == 0
+    assert manual_calls == 0
+
+    chat_client_base.run_responses = [ChatResponse(messages=Message(role="assistant", contents=["done"]))]
+    final_response = await agent.run(requests[0].to_function_approval_response(approved=True), session=session)
+
+    assert final_response.text == "done"
+    assert auto_calls == 1
+    assert manual_calls == 1
+
+
+async def test_tool_approval_middleware_auto_approved_loops_share_function_call_budget(
+    chat_client_base: MockBaseChatClient,
+) -> None:
+    """Auto-approved re-entry should not reset max_function_calls."""
+    calls = 0
+
+    @tool(name="budgeted_tool", approval_mode="always_require")
+    def budgeted_tool(value: str) -> str:
+        nonlocal calls
+        calls += 1
+        return value
+
+    def auto_approve_budgeted_tool(function_call: Content) -> bool:
+        return function_call.name == "budgeted_tool"
+
+    chat_client_base.function_invocation_configuration["max_function_calls"] = 1
+    agent = Agent(
+        client=chat_client_base,
+        tools=[budgeted_tool],
+        middleware=[ToolApprovalMiddleware(auto_approval_rules=[auto_approve_budgeted_tool])],
+    )
+    session = AgentSession(session_id="shared-budget-session")
+    chat_client_base.run_responses = [
+        ChatResponse(
+            messages=Message(
+                role="assistant",
+                contents=[
+                    Content.from_function_call(
+                        call_id="call_first",
+                        name="budgeted_tool",
+                        arguments='{"value": "first"}',
+                    )
+                ],
+            )
+        ),
+        ChatResponse(
+            messages=Message(
+                role="assistant",
+                contents=[
+                    Content.from_function_call(
+                        call_id="call_second",
+                        name="budgeted_tool",
+                        arguments='{"value": "second"}',
+                    )
+                ],
+            )
+        ),
+    ]
+
+    response = await agent.run("call repeatedly", session=session)
+
+    assert response.text == "I broke out of the function invocation loop..."
+    assert calls == 1
+
+
+async def test_tool_approval_middleware_queues_streamed_approval_requests(
+    chat_client_base: MockBaseChatClient,
+) -> None:
+    """Streaming approval requests should also be queued one at a time."""
+    calls = 0
+
+    @tool(name="first_streamed_tool", approval_mode="always_require")
+    def first_streamed_tool() -> str:
+        nonlocal calls
+        calls += 1
+        return "first"
+
+    @tool(name="second_streamed_tool", approval_mode="always_require")
+    def second_streamed_tool() -> str:
+        nonlocal calls
+        calls += 1
+        return "second"
+
+    agent = Agent(
+        client=chat_client_base,
+        tools=[first_streamed_tool, second_streamed_tool],
+        middleware=[ToolApprovalMiddleware()],
+    )
+    session = AgentSession(session_id="stream-queue-session")
+    chat_client_base.streaming_responses = [
+        [
+            ChatResponseUpdate(
+                contents=[Content.from_function_call(call_id="call_first", name="first_streamed_tool", arguments="{}")],
+                role="assistant",
+            ),
+            ChatResponseUpdate(
+                contents=[
+                    Content.from_function_call(call_id="call_second", name="second_streamed_tool", arguments="{}")
+                ],
+                role="assistant",
+            ),
+        ]
+    ]
+
+    first_stream = agent.run("call both", stream=True, session=session)
+    first_updates = [update async for update in first_stream]
+    first_requests = [content for update in first_updates for content in update.user_input_requests]
+    assert [_function_call(request).name for request in first_requests] == ["first_streamed_tool"]
+    assert calls == 0
+
+    second_stream = agent.run(
+        first_requests[0].to_function_approval_response(approved=True),
+        stream=True,
+        session=session,
+    )
+    second_updates = [update async for update in second_stream]
+    second_requests = [content for update in second_updates for content in update.user_input_requests]
+    assert [_function_call(request).name for request in second_requests] == ["second_streamed_tool"]
+    assert calls == 0
+
+    chat_client_base.streaming_responses = [
+        [ChatResponseUpdate(contents=[Content.from_text("done")], role="assistant")]
+    ]
+    final_stream = agent.run(
+        second_requests[0].to_function_approval_response(approved=True),
+        stream=True,
+        session=session,
+    )
+    final_updates = [update async for update in final_stream]
+    final_response = await final_stream.get_final_response()
+
+    assert final_updates[-1].text == "done"
+    assert final_response.text == "done"
+    assert calls == 2
+
+
+async def test_tool_approval_middleware_always_approve_tool_rule(
+    chat_client_base: MockBaseChatClient,
+) -> None:
+    """An always-approve response should add a standing tool-level approval rule."""
+    calls = 0
+
+    @tool(name="dangerous_tool", approval_mode="always_require")
+    def dangerous_tool(value: str) -> str:
+        nonlocal calls
+        calls += 1
+        return value
+
+    agent = Agent(
+        client=chat_client_base,
+        tools=[dangerous_tool],
+        middleware=[ToolApprovalMiddleware()],
+    )
+    session = AgentSession(session_id="standing-rule-session")
+    chat_client_base.run_responses = [
+        ChatResponse(
+            messages=Message(
+                role="assistant",
+                contents=[
+                    Content.from_function_call(
+                        call_id="call_initial",
+                        name="dangerous_tool",
+                        arguments='{"value": "one"}',
+                    )
+                ],
+            )
+        )
+    ]
+
+    first_response = await agent.run("call once", session=session)
+    first_request = _approval_requests(first_response.messages)[0]
+
+    chat_client_base.run_responses = [ChatResponse(messages=Message(role="assistant", contents=["first done"]))]
+    await agent.run(create_always_approve_tool_response(first_request), session=session)
+
+    assert calls == 1
+
+    chat_client_base.run_responses = [
+        ChatResponse(
+            messages=Message(
+                role="assistant",
+                contents=[
+                    Content.from_function_call(
+                        call_id="call_auto",
+                        name="dangerous_tool",
+                        arguments='{"value": "two"}',
+                    )
+                ],
+            )
+        ),
+        ChatResponse(messages=Message(role="assistant", contents=["second done"])),
+    ]
+
+    second_response = await agent.run("call again", session=session)
+
+    assert second_response.text == "second done"
+    assert calls == 2
+
+
+async def test_tool_approval_middleware_standing_rules_include_hosted_server_boundary(
+    chat_client_base: MockBaseChatClient,
+) -> None:
+    """A standing hosted-tool rule should only match the same server_label."""
+    calls = 0
+
+    @tool(name="hosted_tool", approval_mode="always_require")
+    def hosted_tool() -> str:
+        nonlocal calls
+        calls += 1
+        return "hosted"
+
+    def hosted_call(call_id: str, server_label: str) -> Content:
+        return Content.from_function_call(
+            call_id=call_id,
+            name="hosted_tool",
+            arguments="{}",
+            additional_properties={"server_label": server_label},
+        )
+
+    agent = Agent(
+        client=chat_client_base,
+        tools=[hosted_tool],
+        middleware=[ToolApprovalMiddleware()],
+    )
+    session = AgentSession(session_id="hosted-boundary-session")
+    chat_client_base.run_responses = [
+        ChatResponse(messages=Message(role="assistant", contents=[hosted_call("call_initial", "server-a")]))
+    ]
+
+    first_response = await agent.run("call hosted a", session=session)
+    first_request = _approval_requests(first_response.messages)[0]
+
+    chat_client_base.run_responses = [ChatResponse(messages=Message(role="assistant", contents=["server a done"]))]
+    await agent.run(create_always_approve_tool_response(first_request), session=session)
+
+    assert calls == 0
+
+    chat_client_base.run_responses = [
+        ChatResponse(messages=Message(role="assistant", contents=[hosted_call("call_same_server", "server-a")])),
+        ChatResponse(messages=Message(role="assistant", contents=["same server done"])),
+    ]
+
+    same_server_response = await agent.run("call hosted a again", session=session)
+
+    assert same_server_response.text == "same server done"
+    assert _approval_requests(same_server_response.messages) == []
+    assert calls == 0
+
+    chat_client_base.run_responses = [
+        ChatResponse(messages=Message(role="assistant", contents=[hosted_call("call_other_server", "server-b")]))
+    ]
+
+    other_server_response = await agent.run("call hosted b", session=session)
+
+    requests = _approval_requests(other_server_response.messages)
+    assert [_function_call(request).additional_properties["server_label"] for request in requests] == ["server-b"]
+    assert calls == 0
+
+
+async def test_tool_approval_middleware_always_approve_tool_with_arguments_rule(
+    chat_client_base: MockBaseChatClient,
+) -> None:
+    """Argument-scoped always-approve rules should require exact argument matches."""
+    calls = 0
+
+    @tool(name="argument_scoped_tool", approval_mode="always_require")
+    def argument_scoped_tool(value: str) -> str:
+        nonlocal calls
+        calls += 1
+        return value
+
+    agent = Agent(
+        client=chat_client_base,
+        tools=[argument_scoped_tool],
+        middleware=[ToolApprovalMiddleware()],
+    )
+    session = AgentSession(session_id="argument-rule-session")
+    chat_client_base.run_responses = [
+        ChatResponse(
+            messages=Message(
+                role="assistant",
+                contents=[
+                    Content.from_function_call(
+                        call_id="call_initial",
+                        name="argument_scoped_tool",
+                        arguments='{"value": "same"}',
+                    )
+                ],
+            )
+        )
+    ]
+
+    first_response = await agent.run("call with same", session=session)
+    first_request = _approval_requests(first_response.messages)[0]
+
+    chat_client_base.run_responses = [ChatResponse(messages=Message(role="assistant", contents=["first done"]))]
+    await agent.run(create_always_approve_tool_with_arguments_response(first_request), session=session)
+
+    assert calls == 1
+
+    chat_client_base.run_responses = [
+        ChatResponse(
+            messages=Message(
+                role="assistant",
+                contents=[
+                    Content.from_function_call(
+                        call_id="call_same",
+                        name="argument_scoped_tool",
+                        arguments='{"value": "same"}',
+                    )
+                ],
+            )
+        ),
+        ChatResponse(messages=Message(role="assistant", contents=["same done"])),
+    ]
+
+    second_response = await agent.run("call with same again", session=session)
+
+    assert second_response.text == "same done"
+    assert calls == 2
+
+    chat_client_base.run_responses = [
+        ChatResponse(
+            messages=Message(
+                role="assistant",
+                contents=[
+                    Content.from_function_call(
+                        call_id="call_different",
+                        name="argument_scoped_tool",
+                        arguments='{"value": "different"}',
+                    )
+                ],
+            )
+        )
+    ]
+
+    third_response = await agent.run("call with different args", session=session)
+
+    requests = _approval_requests(third_response.messages)
+    assert [_function_call(request).arguments for request in requests] == ['{"value": "different"}']
+    assert calls == 2
+
+
+async def test_tool_approval_middleware_empty_arguments_rule_is_not_tool_wide(
+    chat_client_base: MockBaseChatClient,
+) -> None:
+    """An argument-scoped no-argument approval should not become a wildcard."""
+    calls = 0
+
+    @tool(name="optional_args_tool", approval_mode="always_require")
+    def optional_args_tool(value: str = "default") -> str:
+        nonlocal calls
+        calls += 1
+        return value
+
+    agent = Agent(
+        client=chat_client_base,
+        tools=[optional_args_tool],
+        middleware=[ToolApprovalMiddleware()],
+    )
+    session = AgentSession(session_id="empty-arguments-rule-session")
+    chat_client_base.run_responses = [
+        ChatResponse(
+            messages=Message(
+                role="assistant",
+                contents=[
+                    Content.from_function_call(
+                        call_id="call_empty",
+                        name="optional_args_tool",
+                        arguments="{}",
+                    )
+                ],
+            )
+        )
+    ]
+
+    first_response = await agent.run("call without args", session=session)
+    first_request = _approval_requests(first_response.messages)[0]
+
+    chat_client_base.run_responses = [ChatResponse(messages=Message(role="assistant", contents=["empty done"]))]
+    await agent.run(create_always_approve_tool_with_arguments_response(first_request), session=session)
+
+    assert calls == 1
+
+    chat_client_base.run_responses = [
+        ChatResponse(
+            messages=Message(
+                role="assistant",
+                contents=[
+                    Content.from_function_call(
+                        call_id="call_non_empty",
+                        name="optional_args_tool",
+                        arguments='{"value": "custom"}',
+                    )
+                ],
+            )
+        )
+    ]
+
+    second_response = await agent.run("call with args", session=session)
+
+    requests = _approval_requests(second_response.messages)
+    assert [_function_call(request).arguments for request in requests] == ['{"value": "custom"}']
+    assert calls == 1

python/packages/core/tests/core/test_hyperlight_namespace.py

@@ -0,0 +1,42 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import sys
+from types import ModuleType
+
+import pytest
+
+import agent_framework.hyperlight as hyperlight
+
+
+def test_hyperlight_namespace_dir_lists_lazy_exports() -> None:
+    names = dir(hyperlight)
+    for expected in (
+        "AllowedDomain",
+        "AllowedDomainInput",
+        "FileMount",
+        "FileMountInput",
+        "HyperlightCodeActProvider",
+        "HyperlightExecuteCodeTool",
+    ):
+        assert expected in names
+
+
+def test_hyperlight_namespace_lazy_loads_known_attribute(monkeypatch: pytest.MonkeyPatch) -> None:
+    sentinel = object()
+    fake_module = ModuleType("agent_framework_hyperlight")
+    fake_module.HyperlightCodeActProvider = sentinel  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+    monkeypatch.setitem(sys.modules, "agent_framework_hyperlight", fake_module)
+
+    assert hyperlight.HyperlightCodeActProvider is sentinel
+
+
+def test_hyperlight_namespace_unknown_attribute_raises_attribute_error() -> None:
+    with pytest.raises(AttributeError, match="Module `hyperlight` has no attribute DoesNotExist."):
+        _ = hyperlight.DoesNotExist  # type: ignore[attr-defined]
+
+
+def test_hyperlight_namespace_missing_package_raises_helpful_error(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setitem(sys.modules, "agent_framework_hyperlight", None)
+
+    with pytest.raises(ModuleNotFoundError, match="agent-framework-hyperlight"):
+        _ = hyperlight.HyperlightCodeActProvider

python/packages/core/tests/core/test_mcp_observability.py

@@ -0,0 +1,377 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for MCP client span instrumentation per OTel GenAI Semantic Conventions.
+
+See: https://opentelemetry.io/docs/specs/semconv/gen-ai/mcp/#client
+"""
+
+from __future__ import annotations
+
+from typing import Any
+from unittest.mock import AsyncMock, Mock
+
+import pytest
+from mcp import types
+from mcp.shared.exceptions import McpError
+from mcp.types import ErrorData
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+from opentelemetry.trace import SpanKind, StatusCode
+
+from agent_framework import MCPStdioTool, MCPStreamableHTTPTool, MCPWebsocketTool
+from agent_framework._mcp import MCPTool
+from agent_framework.exceptions import ToolExecutionException
+from agent_framework.observability import OtelAttr
+
+# region helpers
+
+
+def _make_connected_mcp_tool(
+    name: str = "test-mcp",
+    *,
+    supports_tools: bool = True,
+    supports_prompts: bool = True,
+) -> MCPTool:
+    """Create an MCPTool with a mocked session, ready for testing."""
+    tool = MCPTool(name=name)  # type: ignore[abstract]
+    tool.session = AsyncMock()
+    tool.is_connected = True
+    tool._supports_tools = supports_tools
+    tool._supports_prompts = supports_prompts
+    tool.load_tools_flag = True
+    tool.load_prompts_flag = True
+    return tool
+
+
+def _make_tool_list_result(
+    tools: list[dict[str, Any]] | None = None,
+) -> Mock:
+    """Create a mock ListToolsResult."""
+    if tools is None:
+        tools = [{"name": "get-weather", "description": "Get weather", "inputSchema": {"type": "object"}}]
+    result = Mock()
+    result.tools = [
+        types.Tool(name=t["name"], description=t.get("description", ""), inputSchema=t.get("inputSchema", {}))
+        for t in tools
+    ]
+    result.nextCursor = None
+    return result
+
+
+def _make_prompt_list_result(
+    prompts: list[dict[str, Any]] | None = None,
+) -> Mock:
+    """Create a mock ListPromptsResult."""
+    if prompts is None:
+        prompts = [{"name": "analyze-code", "description": "Analyze code"}]
+    result = Mock()
+    result.prompts = [
+        types.Prompt(name=p["name"], description=p.get("description", ""), arguments=None) for p in prompts
+    ]
+    result.nextCursor = None
+    return result
+
+
+def _make_call_tool_result(text: str = "result", is_error: bool = False) -> Mock:
+    """Create a mock CallToolResult."""
+    result = Mock()
+    result.isError = is_error
+    result.content = [types.TextContent(type="text", text=text)]
+    result.structuredContent = None
+    return result
+
+
+def _make_get_prompt_result(text: str = "prompt result") -> types.GetPromptResult:
+    """Create a mock GetPromptResult."""
+    return types.GetPromptResult(
+        description="test prompt",
+        messages=[
+            types.PromptMessage(
+                role="user",
+                content=types.TextContent(type="text", text=text),
+            )
+        ],
+    )
+
+
+# endregion
+
+
+# region initialize span
+
+
+async def test_mcp_initialize_span(span_exporter: InMemorySpanExporter):
+    """session.initialize() should produce an MCP CLIENT span named 'initialize'."""
+    tool = MCPTool(name="test-server")  # type: ignore[abstract]
+
+    mock_session_cls = AsyncMock()
+    init_result = Mock()
+    init_result.capabilities = None
+    init_result.protocolVersion = "2025-06-18"
+    mock_session_cls.initialize = AsyncMock(return_value=init_result)
+
+    # Create a mock transport context manager
+    mock_transport = AsyncMock()
+    mock_transport.__aenter__ = AsyncMock(return_value=(Mock(), Mock()))
+    mock_transport.__aexit__ = AsyncMock(return_value=False)
+
+    # Mock get_mcp_client and the session creation
+    tool.session = None
+    tool.load_tools_flag = False
+    tool.load_prompts_flag = False
+
+    span_exporter.clear()
+
+    with pytest.MonkeyPatch.context() as m:
+        m.setattr(tool, "get_mcp_client", lambda: mock_transport)
+
+        async def patched_connect(self_: Any, *, reset: bool = False, load_configured: bool = True) -> None:
+            # Simulate _connect_on_owner: create initialize span and call session.initialize()
+            from agent_framework._mcp import create_mcp_client_span
+            from agent_framework.observability import OtelAttr
+
+            with create_mcp_client_span("initialize", attributes=self_._mcp_base_span_attributes()) as init_span:
+                result = await mock_session_cls.initialize()
+                protocol_version = getattr(result, "protocolVersion", None)
+                if protocol_version:
+                    init_span.set_attribute(OtelAttr.MCP_PROTOCOL_VERSION, protocol_version)
+
+            self_.session = mock_session_cls
+            self_.is_connected = True
+
+        m.setattr(MCPTool, "_connect_on_owner", patched_connect)
+        await tool.connect()
+
+    mock_session_cls.initialize.assert_awaited_once()
+    spans = span_exporter.get_finished_spans()
+    init_spans = [s for s in spans if s.name == "initialize"]
+    assert len(init_spans) == 1
+    span = init_spans[0]
+    assert span.kind == SpanKind.CLIENT
+    assert span.attributes[OtelAttr.MCP_METHOD_NAME] == "initialize"  # type: ignore[index]  # pyrefly: ignore[unsupported-operation]  # ty: ignore[not-subscriptable]
+    assert span.attributes.get(OtelAttr.MCP_PROTOCOL_VERSION) == "2025-06-18"  # type: ignore[union-attr]  # ty: ignore[unresolved-attribute]
+
+
+# endregion
+
+
+# region tools/list span
+
+
+async def test_mcp_tools_list_span(span_exporter: InMemorySpanExporter):
+    """session.list_tools() should produce an MCP CLIENT span named 'tools/list'."""
+    tool = _make_connected_mcp_tool()
+    tool.session.list_tools = AsyncMock(return_value=_make_tool_list_result())  # type: ignore[method-assign, union-attr]  # ty: ignore[invalid-assignment]
+
+    span_exporter.clear()
+    await tool.load_tools()
+
+    spans = span_exporter.get_finished_spans()
+    list_spans = [s for s in spans if s.name == "tools/list"]
+    assert len(list_spans) == 1
+    span = list_spans[0]
+    assert span.kind == SpanKind.CLIENT
+    assert span.attributes[OtelAttr.MCP_METHOD_NAME] == "tools/list"  # type: ignore[index]  # pyrefly: ignore[unsupported-operation]  # ty: ignore[not-subscriptable]
+
+
+# endregion
+
+
+# region prompts/list span
+
+
+async def test_mcp_prompts_list_span(span_exporter: InMemorySpanExporter):
+    """session.list_prompts() should produce an MCP CLIENT span named 'prompts/list'."""
+    tool = _make_connected_mcp_tool()
+    tool.session.list_prompts = AsyncMock(return_value=_make_prompt_list_result())  # type: ignore[method-assign, union-attr]  # ty: ignore[invalid-assignment]
+
+    span_exporter.clear()
+    await tool.load_prompts()
+
+    spans = span_exporter.get_finished_spans()
+    list_spans = [s for s in spans if s.name == "prompts/list"]
+    assert len(list_spans) == 1
+    span = list_spans[0]
+    assert span.kind == SpanKind.CLIENT
+    assert span.attributes[OtelAttr.MCP_METHOD_NAME] == "prompts/list"  # type: ignore[index]  # pyrefly: ignore[unsupported-operation]  # ty: ignore[not-subscriptable]
+
+
+# endregion
+
+
+# region tools/call span
+
+
+async def test_mcp_tools_call_creates_client_span_when_no_parent(span_exporter: InMemorySpanExporter):
+    """Direct call_tool() without FunctionTool wrapper creates new MCP CLIENT span."""
+    tool = _make_connected_mcp_tool()
+    tool.session.call_tool = AsyncMock(return_value=_make_call_tool_result("hello"))  # type: ignore[method-assign, union-attr]  # ty: ignore[invalid-assignment]
+
+    span_exporter.clear()
+    result = await tool.call_tool("get-weather", city="Seattle")
+
+    assert result is not None
+    spans = span_exporter.get_finished_spans()
+    call_spans = [s for s in spans if "tools/call" in s.name]
+    assert len(call_spans) == 1
+    span = call_spans[0]
+    assert span.kind == SpanKind.CLIENT
+    assert span.name == "tools/call get-weather"
+    assert span.attributes[OtelAttr.MCP_METHOD_NAME] == "tools/call"  # type: ignore[index]  # pyrefly: ignore[unsupported-operation]  # ty: ignore[not-subscriptable]
+    assert span.attributes[OtelAttr.TOOL_NAME] == "get-weather"  # type: ignore[index]  # pyrefly: ignore[unsupported-operation]  # ty: ignore[not-subscriptable]
+
+
+async def test_mcp_tools_call_tool_error_sets_error_type(span_exporter: InMemorySpanExporter):
+    """When CallToolResult.isError is true, error.type should be 'tool_error' per MCP spec."""
+    tool = _make_connected_mcp_tool()
+    tool.session.call_tool = AsyncMock(return_value=_make_call_tool_result("bad input", is_error=True))  # type: ignore[method-assign, union-attr]  # ty: ignore[invalid-assignment]
+
+    span_exporter.clear()
+    with pytest.raises(ToolExecutionException):
+        await tool.call_tool("get-weather", city="invalid")
+
+    spans = span_exporter.get_finished_spans()
+    call_spans = [s for s in spans if "tools/call" in s.name]
+    assert len(call_spans) == 1
+    span = call_spans[0]
+    assert span.attributes.get(OtelAttr.ERROR_TYPE) == "tool_error"  # type: ignore[union-attr]  # ty: ignore[unresolved-attribute]
+    assert span.status.status_code == StatusCode.ERROR
+
+
+async def test_mcp_tools_call_mcp_error_sets_error_type(span_exporter: InMemorySpanExporter):
+    """When session.call_tool() raises McpError, error.type should be the exception class name."""
+    tool = _make_connected_mcp_tool()
+    tool.session.call_tool = AsyncMock(side_effect=McpError(ErrorData(code=-32600, message="invalid request")))  # type: ignore[method-assign, union-attr]  # ty: ignore[invalid-assignment]
+
+    span_exporter.clear()
+    with pytest.raises(ToolExecutionException):
+        await tool.call_tool("get-weather")
+
+    spans = span_exporter.get_finished_spans()
+    call_spans = [s for s in spans if "tools/call" in s.name]
+    assert len(call_spans) == 1
+    span = call_spans[0]
+    assert span.attributes.get(OtelAttr.ERROR_TYPE) == "McpError"  # type: ignore[union-attr]  # ty: ignore[unresolved-attribute]
+    assert span.status.status_code == StatusCode.ERROR
+
+
+# endregion
+
+
+# region prompts/get span
+
+
+async def test_mcp_prompts_get_creates_client_span(span_exporter: InMemorySpanExporter):
+    """get_prompt() should always create a new MCP CLIENT span (not enrich execute_tool)."""
+    tool = _make_connected_mcp_tool()
+    tool.session.get_prompt = AsyncMock(return_value=_make_get_prompt_result("code analysis"))  # type: ignore[method-assign, union-attr]  # ty: ignore[invalid-assignment]
+
+    span_exporter.clear()
+    result = await tool.get_prompt("analyze-code", language="python")
+
+    assert "code analysis" in result
+    spans = span_exporter.get_finished_spans()
+    prompt_spans = [s for s in spans if "prompts/get" in s.name]
+    assert len(prompt_spans) == 1
+    span = prompt_spans[0]
+    assert span.kind == SpanKind.CLIENT
+    assert span.name == "prompts/get analyze-code"
+    assert span.attributes[OtelAttr.MCP_METHOD_NAME] == "prompts/get"  # type: ignore[index]  # pyrefly: ignore[unsupported-operation]  # ty: ignore[not-subscriptable]
+    assert span.attributes[OtelAttr.PROMPT_NAME] == "analyze-code"  # type: ignore[index]  # pyrefly: ignore[unsupported-operation]  # ty: ignore[not-subscriptable]
+
+
+async def test_mcp_prompts_get_mcp_error_sets_error_type(span_exporter: InMemorySpanExporter):
+    """When session.get_prompt() raises McpError, the span should have error.type and ERROR status."""
+    tool = _make_connected_mcp_tool()
+    tool.session.get_prompt = AsyncMock(  # type: ignore[method-assign, union-attr]  # ty: ignore[invalid-assignment]
+        side_effect=McpError(ErrorData(code=-32602, message="prompt not found"))
+    )
+
+    span_exporter.clear()
+    with pytest.raises(ToolExecutionException):
+        await tool.get_prompt("missing-prompt")
+
+    spans = span_exporter.get_finished_spans()
+    prompt_spans = [s for s in spans if "prompts/get" in s.name]
+    assert len(prompt_spans) == 1
+    span = prompt_spans[0]
+    assert span.attributes.get(OtelAttr.ERROR_TYPE) == "McpError"  # type: ignore[union-attr]  # ty: ignore[unresolved-attribute]
+    assert span.status.status_code == StatusCode.ERROR
+
+
+# endregion
+
+
+# region transport attributes
+
+
+def test_mcp_stdio_tool_transport_attributes():
+    """MCPStdioTool should have network.transport='pipe'."""
+    tool = MCPStdioTool(name="test", command="python")
+    attrs = tool._mcp_base_span_attributes()
+    assert attrs[OtelAttr.NETWORK_TRANSPORT] == "pipe"
+    assert OtelAttr.ADDRESS not in attrs
+
+
+def test_mcp_http_tool_transport_attributes():
+    """MCPStreamableHTTPTool should have tcp transport and URL-based server address/port."""
+    tool = MCPStreamableHTTPTool(name="test", url="https://api.example.com:8443/mcp")
+    attrs = tool._mcp_base_span_attributes()
+    assert attrs[OtelAttr.NETWORK_TRANSPORT] == "tcp"
+    assert attrs[OtelAttr.NETWORK_PROTOCOL_NAME] == "http"
+    assert attrs[OtelAttr.ADDRESS] == "api.example.com"
+    assert attrs[OtelAttr.PORT] == 8443
+
+
+def test_mcp_http_tool_default_port():
+    """MCPStreamableHTTPTool should default to 443 for https."""
+    tool = MCPStreamableHTTPTool(name="test", url="https://api.example.com/mcp")
+    attrs = tool._mcp_base_span_attributes()
+    assert attrs[OtelAttr.PORT] == 443
+
+
+def test_mcp_http_tool_http_default_port():
+    """MCPStreamableHTTPTool should default to 80 for http."""
+    tool = MCPStreamableHTTPTool(name="test", url="http://localhost/mcp")
+    attrs = tool._mcp_base_span_attributes()
+    assert attrs[OtelAttr.PORT] == 80
+
+
+def test_mcp_websocket_tool_transport_attributes():
+    """MCPWebsocketTool should have tcp transport and URL-based server address/port."""
+    tool = MCPWebsocketTool(name="test", url="wss://ws.example.com:9090/mcp")
+    attrs = tool._mcp_base_span_attributes()
+    assert attrs[OtelAttr.NETWORK_TRANSPORT] == "tcp"
+    assert attrs[OtelAttr.NETWORK_PROTOCOL_NAME] == "websocket"
+    assert attrs[OtelAttr.ADDRESS] == "ws.example.com"
+    assert attrs[OtelAttr.PORT] == 9090
+
+
+def test_mcp_websocket_tool_default_port():
+    """MCPWebsocketTool should default to 443 for wss."""
+    tool = MCPWebsocketTool(name="test", url="wss://ws.example.com/mcp")
+    attrs = tool._mcp_base_span_attributes()
+    assert attrs[OtelAttr.PORT] == 443
+
+
+# endregion
+
+
+# region observability disabled
+
+
+@pytest.mark.parametrize("enable_instrumentation", [False], indirect=True)
+async def test_mcp_spans_not_created_when_observability_disabled(span_exporter: InMemorySpanExporter):
+    """No MCP spans should be created when observability is disabled."""
+    tool = _make_connected_mcp_tool()
+    tool.session.list_tools = AsyncMock(return_value=_make_tool_list_result())  # type: ignore[method-assign, union-attr]  # ty: ignore[invalid-assignment]
+    tool.session.call_tool = AsyncMock(return_value=_make_call_tool_result("ok"))  # type: ignore[method-assign, union-attr]  # ty: ignore[invalid-assignment]
+
+    span_exporter.clear()
+    await tool.load_tools()
+    await tool.call_tool("get-weather", city="Seattle")
+
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 0
+
+
+# endregion

python/packages/core/tests/core/test_mcp_skills.py

@@ -0,0 +1,661 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for MCP-based skills (MCPSkillsSource, MCPSkill, MCPSkillResource)."""
+
+from __future__ import annotations
+
+import base64
+import json
+from unittest.mock import AsyncMock
+
+import pytest
+from mcp.shared.exceptions import McpError
+from mcp.types import (
+    BlobResourceContents,
+    ErrorData,
+    ReadResourceResult,
+    TextResourceContents,
+)
+from pydantic import AnyUrl
+
+from agent_framework import MCPSkill, MCPSkillResource, MCPSkillsSource, SkillsSourceContext
+from agent_framework._skills import _parse_mcp_skill_index
+
+from .conftest import MockAgent
+
+# ---------------------------------------------------------------------------
+# Fixtures & helpers
+# ---------------------------------------------------------------------------
+
+
+# Shared context for exercising skill sources where the agent/session are irrelevant.
+_SOURCE_CTX = SkillsSourceContext(agent=MockAgent())  # type: ignore[abstract]  # pyrefly: ignore[bad-instantiation]
+
+SAMPLE_SKILL_MD = """\
+---
+name: unit-converter
+description: Convert between common units.
+---
+# Unit Converter
+
+Body content here.
+"""
+
+SAMPLE_SKILL_INDEX = json.dumps({
+    "$schema": "https://schemas.agentskills.io/discovery/0.2.0/schema.json",
+    "skills": [
+        {
+            "name": "unit-converter",
+            "type": "skill-md",
+            "description": "Convert between common units.",
+            "url": "skill://unit-converter/SKILL.md",
+        }
+    ],
+})
+
+
+def _make_text_result(text: str, uri: str = "skill://test") -> ReadResourceResult:
+    """Create a ReadResourceResult with a single TextResourceContents."""
+    return ReadResourceResult(contents=[TextResourceContents(uri=AnyUrl(uri), text=text, mimeType="text/markdown")])
+
+
+def _make_blob_result(
+    data: bytes,
+    uri: str = "skill://test",
+    mime_type: str = "application/octet-stream",
+) -> ReadResourceResult:
+    """Create a ReadResourceResult with a single BlobResourceContents."""
+    return ReadResourceResult(
+        contents=[BlobResourceContents(uri=AnyUrl(uri), blob=base64.b64encode(data).decode(), mimeType=mime_type)]
+    )
+
+
+def _make_empty_result() -> ReadResourceResult:
+    """Create a ReadResourceResult with no contents."""
+    return ReadResourceResult(contents=[])
+
+
+def _make_client(**read_resource_responses: ReadResourceResult) -> AsyncMock:
+    """Create a mock ClientSession whose read_resource returns different results per URI.
+
+    Args:
+        **read_resource_responses: Mapping of URI string to ReadResourceResult.
+            Any URI not in this mapping raises McpError with the MCP-spec
+            "Resource not found" code (-32002).
+    """
+    client = AsyncMock()
+
+    async def _read_resource(uri: AnyUrl) -> ReadResourceResult:
+        uri_str = str(uri)
+        if uri_str in read_resource_responses:
+            return read_resource_responses[uri_str]
+        raise McpError(error=ErrorData(code=-32002, message=f"Resource not found: {uri_str}"))
+
+    client.read_resource = AsyncMock(side_effect=_read_resource)
+    return client
+
+
+# ---------------------------------------------------------------------------
+# _parse_mcp_skill_index tests
+# ---------------------------------------------------------------------------
+
+
+class TestParseMCPSkillIndex:
+    """Tests for the _parse_mcp_skill_index helper."""
+
+    def test_parses_valid_index(self) -> None:
+        index = _parse_mcp_skill_index(SAMPLE_SKILL_INDEX)
+        assert index.schema == "https://schemas.agentskills.io/discovery/0.2.0/schema.json"
+        assert len(index.skills) == 1
+        assert index.skills[0].name == "unit-converter"
+        assert index.skills[0].type == "skill-md"
+        assert index.skills[0].url == "skill://unit-converter/SKILL.md"
+
+    def test_parses_empty_skills_array(self) -> None:
+        index = _parse_mcp_skill_index('{"$schema": "test", "skills": []}')
+        assert index.skills == []
+
+    def test_parses_missing_skills_key(self) -> None:
+        index = _parse_mcp_skill_index('{"$schema": "test"}')
+        assert index.skills == []
+
+    def test_raises_on_non_object(self) -> None:
+        with pytest.raises(ValueError, match="must be a JSON object"):
+            _parse_mcp_skill_index("[]")
+
+    def test_raises_on_invalid_json(self) -> None:
+        with pytest.raises(json.JSONDecodeError):
+            _parse_mcp_skill_index("not json")
+
+    def test_skips_non_dict_entries(self) -> None:
+        index = _parse_mcp_skill_index('{"skills": ["not-a-dict", {"name": "ok", "type": "skill-md"}]}')
+        assert len(index.skills) == 1
+        assert index.skills[0].name == "ok"
+
+
+# ---------------------------------------------------------------------------
+# MCPSkillResource tests
+# ---------------------------------------------------------------------------
+
+
+class TestMCPSkillsExperimentalStage:
+    """Tests confirming the MCP skills types remain experimental (MCP_SKILLS)."""
+
+    def test_docstrings_include_experimental_warning(self) -> None:
+        assert MCPSkillResource.__doc__ is not None
+        assert MCPSkill.__doc__ is not None
+        assert MCPSkillsSource.__doc__ is not None
+
+        assert ".. warning:: Experimental" in MCPSkillResource.__doc__
+        assert ".. warning:: Experimental" in MCPSkill.__doc__
+        assert ".. warning:: Experimental" in MCPSkillsSource.__doc__
+
+    def test_feature_metadata_is_set(self) -> None:
+        for cls in (MCPSkillResource, MCPSkill, MCPSkillsSource):
+            assert getattr(cls, "__feature_stage__", None) == "experimental"
+            assert getattr(cls, "__feature_id__", None) == "MCP_SKILLS"
+
+
+class TestMCPSkillResource:
+    """Tests for MCPSkillResource."""
+
+    @pytest.mark.asyncio
+    async def test_read_text_content(self) -> None:
+        result = _make_text_result("hello world")
+        resource = MCPSkillResource(name="test.md", result=result)
+        content = await resource.read()
+        assert content == "hello world"
+
+    @pytest.mark.asyncio
+    async def test_read_binary_content(self) -> None:
+        data = bytes([0x01, 0x02, 0x03, 0x04])
+        result = _make_blob_result(data)
+        resource = MCPSkillResource(name="icon.bin", result=result)
+        content = await resource.read()
+        assert content == data
+
+    @pytest.mark.asyncio
+    async def test_read_empty_returns_none(self) -> None:
+        result = _make_empty_result()
+        resource = MCPSkillResource(name="empty", result=result)
+        content = await resource.read()
+        assert content is None
+
+    @pytest.mark.asyncio
+    async def test_read_multiple_text_contents_joined(self) -> None:
+        result = ReadResourceResult(
+            contents=[
+                TextResourceContents(uri=AnyUrl("skill://a"), text="line1", mimeType="text/plain"),
+                TextResourceContents(uri=AnyUrl("skill://b"), text="line2", mimeType="text/plain"),
+            ]
+        )
+        resource = MCPSkillResource(name="multi", result=result)
+        content = await resource.read()
+        assert content == "line1\nline2"
+
+    @pytest.mark.asyncio
+    async def test_binary_takes_precedence_over_text(self) -> None:
+        data = b"\xff\xfe"
+        result = ReadResourceResult(
+            contents=[
+                TextResourceContents(uri=AnyUrl("skill://a"), text="text", mimeType="text/plain"),
+                BlobResourceContents(
+                    uri=AnyUrl("skill://b"),
+                    blob=base64.b64encode(data).decode(),
+                    mimeType="application/octet-stream",
+                ),
+            ]
+        )
+        resource = MCPSkillResource(name="mixed", result=result)
+        content = await resource.read()
+        # The implementation iterates all contents checking for BlobResourceContents
+        # first, so when both text and binary are present, binary is returned.
+        assert content == data
+
+
+# ---------------------------------------------------------------------------
+# MCPSkill tests
+# ---------------------------------------------------------------------------
+
+
+class TestMCPSkill:
+    """Tests for MCPSkill."""
+
+    @pytest.mark.asyncio
+    async def test_get_content_fetches_and_caches(self) -> None:
+        client = _make_client(**{"skill://unit-converter/SKILL.md": _make_text_result(SAMPLE_SKILL_MD)})
+        from agent_framework import SkillFrontmatter
+
+        fm = SkillFrontmatter(name="unit-converter", description="Convert between common units.")
+        skill = MCPSkill(frontmatter=fm, skill_md_uri="skill://unit-converter/SKILL.md", client=client)
+
+        content1 = await skill.get_content()
+        content2 = await skill.get_content()
+
+        assert "Body content here." in content1
+        assert content1 == content2
+        # Only one MCP call should be made (cached)
+        assert client.read_resource.call_count == 1
+
+    @pytest.mark.asyncio
+    async def test_get_content_raises_on_empty(self) -> None:
+        client = _make_client(**{"skill://empty/SKILL.md": _make_empty_result()})
+        from agent_framework import SkillFrontmatter
+
+        fm = SkillFrontmatter(name="empty-skill", description="Empty skill.")
+        skill = MCPSkill(frontmatter=fm, skill_md_uri="skill://empty/SKILL.md", client=client)
+
+        with pytest.raises(ValueError, match="no text content"):
+            await skill.get_content()
+
+    @pytest.mark.asyncio
+    async def test_get_resource_text(self) -> None:
+        client = _make_client(**{
+            "skill://unit-converter/SKILL.md": _make_text_result(SAMPLE_SKILL_MD),
+            "skill://unit-converter/references/checklist.md": _make_text_result("- check thing 1\n- check thing 2"),
+        })
+        from agent_framework import SkillFrontmatter
+
+        fm = SkillFrontmatter(name="unit-converter", description="Convert between common units.")
+        skill = MCPSkill(frontmatter=fm, skill_md_uri="skill://unit-converter/SKILL.md", client=client)
+
+        resource = await skill.get_resource("references/checklist.md")
+        assert resource is not None
+        content = await resource.read()
+        assert content == "- check thing 1\n- check thing 2"
+
+    @pytest.mark.asyncio
+    async def test_get_resource_binary(self) -> None:
+        data = bytes([0x01, 0x02, 0x03, 0x04])
+        client = _make_client(**{
+            "skill://unit-converter/SKILL.md": _make_text_result(SAMPLE_SKILL_MD),
+            "skill://unit-converter/assets/icon.bin": _make_blob_result(data),
+        })
+        from agent_framework import SkillFrontmatter
+
+        fm = SkillFrontmatter(name="unit-converter", description="Convert between common units.")
+        skill = MCPSkill(frontmatter=fm, skill_md_uri="skill://unit-converter/SKILL.md", client=client)
+
+        resource = await skill.get_resource("assets/icon.bin")
+        assert resource is not None
+        content = await resource.read()
+        assert content == data
+
+    @pytest.mark.asyncio
+    async def test_get_resource_unknown_returns_none(self) -> None:
+        client = _make_client(**{"skill://unit-converter/SKILL.md": _make_text_result(SAMPLE_SKILL_MD)})
+        from agent_framework import SkillFrontmatter
+
+        fm = SkillFrontmatter(name="unit-converter", description="Convert between common units.")
+        skill = MCPSkill(frontmatter=fm, skill_md_uri="skill://unit-converter/SKILL.md", client=client)
+
+        resource = await skill.get_resource("references/does-not-exist.md")
+        assert resource is None
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "name",
+        [
+            "../escape.md",
+            "references/../../escape.md",
+            "..",
+            "..\\escape.md",
+            "/etc/passwd",
+            "http://attacker.example.com/payload",
+        ],
+    )
+    async def test_get_resource_path_traversal_returns_none(self, name: str) -> None:
+        # Register a permissive mock that would happily return content for any URI,
+        # so the test fails unless the client-side validation rejects the name
+        # before issuing the read.
+        client = AsyncMock()
+        client.read_resource = AsyncMock(return_value=_make_text_result("should never be returned"))
+
+        from agent_framework import SkillFrontmatter
+
+        fm = SkillFrontmatter(name="unit-converter", description="Convert between common units.")
+        skill = MCPSkill(frontmatter=fm, skill_md_uri="skill://unit-converter/SKILL.md", client=client)
+
+        resource = await skill.get_resource(name)
+        assert resource is None
+        client.read_resource.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_get_resource_empty_name_returns_none(self) -> None:
+        client = _make_client()
+        from agent_framework import SkillFrontmatter
+
+        fm = SkillFrontmatter(name="test-skill", description="Test.")
+        skill = MCPSkill(frontmatter=fm, skill_md_uri="skill://test/SKILL.md", client=client)
+
+        assert await skill.get_resource("") is None
+        assert await skill.get_resource("   ") is None
+
+    @pytest.mark.asyncio
+    async def test_get_script_returns_none(self) -> None:
+        client = _make_client()
+        from agent_framework import SkillFrontmatter
+
+        fm = SkillFrontmatter(name="test-skill", description="Test.")
+        skill = MCPSkill(frontmatter=fm, skill_md_uri="skill://test/SKILL.md", client=client)
+
+        assert await skill.get_script("anything") is None
+
+    def test_compute_skill_root_uri_strips_suffix(self) -> None:
+        assert MCPSkill._compute_skill_root_uri("skill://unit-converter/SKILL.md") == "skill://unit-converter/"
+
+    def test_compute_skill_root_uri_trailing_slash(self) -> None:
+        assert MCPSkill._compute_skill_root_uri("skill://unit-converter/") == "skill://unit-converter/"
+
+    def test_compute_skill_root_uri_no_suffix_adds_slash(self) -> None:
+        assert MCPSkill._compute_skill_root_uri("skill://unit-converter") == "skill://unit-converter/"
+
+
+# ---------------------------------------------------------------------------
+# MCPSkillsSource tests
+# ---------------------------------------------------------------------------
+
+
+class TestMCPSkillsSource:
+    """Tests for MCPSkillsSource."""
+
+    @pytest.mark.asyncio
+    async def test_index_based_discovery_returns_skill(self) -> None:
+        client = _make_client(**{
+            "skill://index.json": _make_text_result(SAMPLE_SKILL_INDEX, uri="skill://index.json"),
+            "skill://unit-converter/SKILL.md": _make_text_result(SAMPLE_SKILL_MD),
+        })
+        source = MCPSkillsSource(client=client)
+        skills = await source.get_skills(_SOURCE_CTX)
+
+        assert len(skills) == 1
+        assert skills[0].frontmatter.name == "unit-converter"
+        assert skills[0].frontmatter.description == "Convert between common units."
+
+        # Content is fetched on demand, not during discovery
+        content = await skills[0].get_content()
+        assert "Body content here." in content
+
+    @pytest.mark.asyncio
+    async def test_no_index_returns_empty(self) -> None:
+        client = _make_client()  # No resources at all
+        source = MCPSkillsSource(client=client)
+        skills = await source.get_skills(_SOURCE_CTX)
+        assert skills == []
+
+    @pytest.mark.asyncio
+    async def test_does_not_read_skill_md_during_discovery(self) -> None:
+        # Index points to a skill, but SKILL.md is not registered on the server.
+        # Discovery should succeed because it only reads the index.
+        client = _make_client(**{"skill://index.json": _make_text_result(SAMPLE_SKILL_INDEX, uri="skill://index.json")})
+        source = MCPSkillsSource(client=client)
+        skills = await source.get_skills(_SOURCE_CTX)
+
+        assert len(skills) == 1
+        assert skills[0].frontmatter.name == "unit-converter"
+
+    @pytest.mark.asyncio
+    async def test_invalid_name_is_skipped(self) -> None:
+        index_json = json.dumps({
+            "$schema": "https://schemas.agentskills.io/discovery/0.2.0/schema.json",
+            "skills": [
+                {
+                    "name": "UnitConverter",  # Invalid: uppercase
+                    "type": "skill-md",
+                    "description": "Convert between common units.",
+                    "url": "skill://UnitConverter/SKILL.md",
+                }
+            ],
+        })
+        client = _make_client(**{"skill://index.json": _make_text_result(index_json, uri="skill://index.json")})
+        source = MCPSkillsSource(client=client)
+        skills = await source.get_skills(_SOURCE_CTX)
+        assert skills == []
+
+    @pytest.mark.asyncio
+    async def test_missing_required_fields_is_skipped(self) -> None:
+        index_json = json.dumps({
+            "$schema": "https://schemas.agentskills.io/discovery/0.2.0/schema.json",
+            "skills": [
+                {
+                    "name": "unit-converter",
+                    "type": "skill-md",
+                    # Missing description and url
+                }
+            ],
+        })
+        client = _make_client(**{"skill://index.json": _make_text_result(index_json, uri="skill://index.json")})
+        source = MCPSkillsSource(client=client)
+        skills = await source.get_skills(_SOURCE_CTX)
+        assert skills == []
+
+    @pytest.mark.asyncio
+    async def test_unsupported_type_is_skipped(self) -> None:
+        index_json = json.dumps({
+            "$schema": "https://schemas.agentskills.io/discovery/0.2.0/schema.json",
+            "skills": [
+                {
+                    "name": "some-skill",
+                    "type": "archive",
+                    "description": "Packaged skill.",
+                    "url": "skill://some-skill.tar.gz",
+                }
+            ],
+        })
+        client = _make_client(**{"skill://index.json": _make_text_result(index_json, uri="skill://index.json")})
+        source = MCPSkillsSource(client=client)
+        skills = await source.get_skills(_SOURCE_CTX)
+        assert skills == []
+
+    @pytest.mark.asyncio
+    async def test_template_type_is_skipped(self) -> None:
+        index_json = json.dumps({
+            "$schema": "https://schemas.agentskills.io/discovery/0.2.0/schema.json",
+            "skills": [
+                {
+                    "type": "mcp-resource-template",
+                    "description": "Per-product documentation skill",
+                    "url": "skill://docs/{product}/SKILL.md",
+                }
+            ],
+        })
+        client = _make_client(**{"skill://index.json": _make_text_result(index_json, uri="skill://index.json")})
+        source = MCPSkillsSource(client=client)
+        skills = await source.get_skills(_SOURCE_CTX)
+        assert skills == []
+
+    @pytest.mark.asyncio
+    async def test_empty_index_returns_empty(self) -> None:
+        client = _make_client(**{"skill://index.json": _make_text_result('{"skills": []}', uri="skill://index.json")})
+        source = MCPSkillsSource(client=client)
+        skills = await source.get_skills(_SOURCE_CTX)
+        assert skills == []
+
+    @pytest.mark.asyncio
+    async def test_malformed_index_json_returns_empty(self) -> None:
+        client = _make_client(**{"skill://index.json": _make_text_result("not valid json", uri="skill://index.json")})
+        source = MCPSkillsSource(client=client)
+        skills = await source.get_skills(_SOURCE_CTX)
+        assert skills == []
+
+    @pytest.mark.asyncio
+    async def test_sibling_text_resource(self) -> None:
+        client = _make_client(**{
+            "skill://index.json": _make_text_result(SAMPLE_SKILL_INDEX, uri="skill://index.json"),
+            "skill://unit-converter/SKILL.md": _make_text_result(SAMPLE_SKILL_MD),
+            "skill://unit-converter/references/checklist.md": _make_text_result("- check thing 1\n- check thing 2"),
+        })
+        source = MCPSkillsSource(client=client)
+        skill = (await source.get_skills(_SOURCE_CTX))[0]
+        resource = await skill.get_resource("references/checklist.md")
+        assert resource is not None
+        content = await resource.read()
+        assert content == "- check thing 1\n- check thing 2"
+
+    @pytest.mark.asyncio
+    async def test_sibling_binary_resource(self) -> None:
+        data = bytes([0x01, 0x02, 0x03, 0x04])
+        client = _make_client(**{
+            "skill://index.json": _make_text_result(SAMPLE_SKILL_INDEX, uri="skill://index.json"),
+            "skill://unit-converter/SKILL.md": _make_text_result(SAMPLE_SKILL_MD),
+            "skill://unit-converter/assets/icon.bin": _make_blob_result(data),
+        })
+        source = MCPSkillsSource(client=client)
+        skill = (await source.get_skills(_SOURCE_CTX))[0]
+        resource = await skill.get_resource("assets/icon.bin")
+        assert resource is not None
+        content = await resource.read()
+        assert content == data
+
+
+# ---------------------------------------------------------------------------
+# McpError code branching tests
+# ---------------------------------------------------------------------------
+
+
+class TestMCPSkillsSourceErrorCodeBranching:
+    """Tests that MCPSkillsSource and MCPSkill branch on McpError.error.code.
+
+    Only "not found" codes (RESOURCE_NOT_FOUND -32002, METHOD_NOT_FOUND -32601)
+    should be silently swallowed as "no skills available." Other McpError codes
+    and non-McpError exceptions must propagate so that auth failures, server
+    crashes, and connection drops are visible.
+    """
+
+    @pytest.mark.asyncio
+    async def test_index_method_not_found_returns_empty(self) -> None:
+        """METHOD_NOT_FOUND (-32601) -> server doesn't support resources/read."""
+        client = AsyncMock()
+        client.read_resource = AsyncMock(side_effect=McpError(error=ErrorData(code=-32601, message="Method not found")))
+        source = MCPSkillsSource(client=client)
+        skills = await source.get_skills(_SOURCE_CTX)
+        assert skills == []
+
+    @pytest.mark.asyncio
+    async def test_index_resource_not_found_returns_empty(self) -> None:
+        """MCP-spec "Resource not found" (-32002) -> server has no index."""
+        client = AsyncMock()
+        client.read_resource = AsyncMock(
+            side_effect=McpError(error=ErrorData(code=-32002, message="Resource not found"))
+        )
+        source = MCPSkillsSource(client=client)
+        skills = await source.get_skills(_SOURCE_CTX)
+        assert skills == []
+
+    @pytest.mark.asyncio
+    async def test_index_invalid_params_propagates(self) -> None:
+        """INVALID_PARAMS (-32602) is a real bug, must propagate (not "not found")."""
+        client = AsyncMock()
+        client.read_resource = AsyncMock(side_effect=McpError(error=ErrorData(code=-32602, message="Invalid params")))
+        source = MCPSkillsSource(client=client)
+        with pytest.raises(McpError):
+            await source.get_skills(_SOURCE_CTX)
+
+    @pytest.mark.asyncio
+    async def test_index_internal_error_propagates(self) -> None:
+        """INTERNAL_ERROR (-32603) must propagate, not silently return empty."""
+        client = AsyncMock()
+        client.read_resource = AsyncMock(side_effect=McpError(error=ErrorData(code=-32603, message="Internal error")))
+        source = MCPSkillsSource(client=client)
+        with pytest.raises(McpError):
+            await source.get_skills(_SOURCE_CTX)
+
+    @pytest.mark.asyncio
+    async def test_index_connection_closed_propagates(self) -> None:
+        """CONNECTION_CLOSED (-32000) must propagate."""
+        client = AsyncMock()
+        client.read_resource = AsyncMock(
+            side_effect=McpError(error=ErrorData(code=-32000, message="Connection closed"))
+        )
+        source = MCPSkillsSource(client=client)
+        with pytest.raises(McpError):
+            await source.get_skills(_SOURCE_CTX)
+
+    @pytest.mark.asyncio
+    async def test_index_generic_error_code_propagates(self) -> None:
+        """Generic handler error (code 0) must propagate."""
+        client = AsyncMock()
+        client.read_resource = AsyncMock(side_effect=McpError(error=ErrorData(code=0, message="Some handler error")))
+        source = MCPSkillsSource(client=client)
+        with pytest.raises(McpError):
+            await source.get_skills(_SOURCE_CTX)
+
+    @pytest.mark.asyncio
+    async def test_index_non_mcp_error_propagates(self) -> None:
+        """Non-McpError exceptions (connection drop, timeout) must propagate."""
+        client = AsyncMock()
+        client.read_resource = AsyncMock(side_effect=ConnectionError("connection lost"))
+        source = MCPSkillsSource(client=client)
+        with pytest.raises(ConnectionError):
+            await source.get_skills(_SOURCE_CTX)
+
+    @pytest.mark.asyncio
+    async def test_get_resource_internal_error_propagates(self) -> None:
+        """McpError with INTERNAL_ERROR on get_resource must propagate."""
+        from agent_framework import SkillFrontmatter
+
+        client = AsyncMock()
+        client.read_resource = AsyncMock(side_effect=McpError(error=ErrorData(code=-32603, message="Server crashed")))
+        fm = SkillFrontmatter(name="test-skill", description="Test.")
+        skill = MCPSkill(frontmatter=fm, skill_md_uri="skill://test/SKILL.md", client=client)
+        with pytest.raises(McpError):
+            await skill.get_resource("references/file.md")
+
+    @pytest.mark.asyncio
+    async def test_get_resource_not_found_returns_none(self) -> None:
+        """McpError with RESOURCE_NOT_FOUND (-32002) on get_resource returns None."""
+        from agent_framework import SkillFrontmatter
+
+        client = AsyncMock()
+        client.read_resource = AsyncMock(
+            side_effect=McpError(error=ErrorData(code=-32002, message="Resource not found"))
+        )
+        fm = SkillFrontmatter(name="test-skill", description="Test.")
+        skill = MCPSkill(frontmatter=fm, skill_md_uri="skill://test/SKILL.md", client=client)
+        result = await skill.get_resource("references/file.md")
+        assert result is None
+
+    @pytest.mark.asyncio
+    async def test_get_resource_connection_error_propagates(self) -> None:
+        """A plain ConnectionError on get_resource must propagate, not return None."""
+        from agent_framework import SkillFrontmatter
+
+        client = AsyncMock()
+        client.read_resource = AsyncMock(side_effect=ConnectionError("connection lost"))
+        fm = SkillFrontmatter(name="test-skill", description="Test.")
+        skill = MCPSkill(frontmatter=fm, skill_md_uri="skill://test/SKILL.md", client=client)
+        with pytest.raises(ConnectionError):
+            await skill.get_resource("references/file.md")
+
+    @pytest.mark.asyncio
+    async def test_get_resource_timeout_error_propagates(self) -> None:
+        """A TimeoutError on get_resource must propagate, not return None."""
+        from agent_framework import SkillFrontmatter
+
+        client = AsyncMock()
+        client.read_resource = AsyncMock(side_effect=TimeoutError("read timed out"))
+        fm = SkillFrontmatter(name="test-skill", description="Test.")
+        skill = MCPSkill(frontmatter=fm, skill_md_uri="skill://test/SKILL.md", client=client)
+        with pytest.raises(TimeoutError):
+            await skill.get_resource("references/file.md")
+
+    @pytest.mark.asyncio
+    async def test_get_resource_generic_mcp_error_propagates(self) -> None:
+        """McpError with a generic code (0) on get_resource must propagate."""
+        from agent_framework import SkillFrontmatter
+
+        client = AsyncMock()
+        client.read_resource = AsyncMock(side_effect=McpError(error=ErrorData(code=0, message="Handler error")))
+        fm = SkillFrontmatter(name="test-skill", description="Test.")
+        skill = MCPSkill(frontmatter=fm, skill_md_uri="skill://test/SKILL.md", client=client)
+        with pytest.raises(McpError):
+            await skill.get_resource("references/file.md")
+
+    @pytest.mark.asyncio
+    async def test_index_timeout_error_propagates(self) -> None:
+        """A TimeoutError reading skill://index.json must propagate."""
+        client = AsyncMock()
+        client.read_resource = AsyncMock(side_effect=TimeoutError("read timed out"))
+        source = MCPSkillsSource(client=client)
+        with pytest.raises(TimeoutError):
+            await source.get_skills(_SOURCE_CTX)

python/packages/core/tests/core/test_middleware_context_result.py

@@ -0,0 +1,450 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from collections.abc import AsyncIterable, Awaitable, Callable
+from unittest.mock import MagicMock
+
+import pytest
+from pydantic import BaseModel, Field
+
+from agent_framework import (
+    Agent,
+    AgentResponse,
+    AgentResponseUpdate,
+    Content,
+    Message,
+    ResponseStream,
+    SupportsAgentRun,
+)
+from agent_framework._middleware import (
+    AgentContext,
+    AgentMiddleware,
+    AgentMiddlewarePipeline,
+    FunctionInvocationContext,
+    FunctionMiddleware,
+    FunctionMiddlewarePipeline,
+)
+from agent_framework._tools import FunctionTool
+
+from .conftest import MockChatClient
+
+
+class FunctionTestArgs(BaseModel):
+    """Test arguments for function middleware tests."""
+
+    name: str = Field(description="Test name parameter")
+
+
+class TestResultOverrideMiddleware:
+    """Test cases for middleware result override functionality."""
+
+    async def test_agent_middleware_response_override_non_streaming(self, mock_agent: SupportsAgentRun) -> None:
+        """Test that agent middleware can override response for non-streaming execution."""
+        override_response = AgentResponse(messages=[Message(role="assistant", contents=["overridden response"])])
+
+        class ResponseOverrideMiddleware(AgentMiddleware):
+            async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
+                # Execute the pipeline first, then override the response
+                await call_next()
+                context.result = override_response
+
+        middleware = ResponseOverrideMiddleware()
+        pipeline = AgentMiddlewarePipeline(middleware)
+        messages = [Message(role="user", contents=["test"])]
+        context = AgentContext(agent=mock_agent, messages=messages)
+
+        handler_called = False
+
+        async def final_handler(ctx: AgentContext) -> AgentResponse:
+            nonlocal handler_called
+            handler_called = True
+            return AgentResponse(messages=[Message(role="assistant", contents=["original response"])])
+
+        result = await pipeline.execute(context, final_handler)
+
+        # Verify the overridden response is returned
+        assert result is not None
+        assert result == override_response
+        assert result.messages[0].text == "overridden response"  # ty: ignore[unresolved-attribute]  # type: ignore[union-attr]
+        # Verify original handler was called since middleware called next()
+        assert handler_called
+
+    async def test_agent_middleware_response_override_streaming(self, mock_agent: SupportsAgentRun) -> None:
+        """Test that agent middleware can override response for streaming execution."""
+
+        async def override_stream() -> AsyncIterable[AgentResponseUpdate]:
+            yield AgentResponseUpdate(contents=[Content.from_text(text="overridden")])
+            yield AgentResponseUpdate(contents=[Content.from_text(text=" stream")])
+
+        class StreamResponseOverrideMiddleware(AgentMiddleware):
+            async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
+                # Execute the pipeline first, then override the response stream
+                await call_next()
+                context.result = ResponseStream(override_stream())
+
+        middleware = StreamResponseOverrideMiddleware()
+        pipeline = AgentMiddlewarePipeline(middleware)
+        messages = [Message(role="user", contents=["test"])]
+        context = AgentContext(agent=mock_agent, messages=messages, stream=True)
+
+        async def final_handler(ctx: AgentContext) -> ResponseStream[AgentResponseUpdate, AgentResponse]:
+            async def _stream() -> AsyncIterable[AgentResponseUpdate]:
+                yield AgentResponseUpdate(contents=[Content.from_text(text="original")])
+
+            return ResponseStream(_stream())
+
+        updates: list[AgentResponseUpdate] = []
+        stream = await pipeline.execute(context, final_handler)  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+        async for update in stream:  # type: ignore[attr-defined, union-attr]  # pyrefly: ignore[not-iterable]  # ty: ignore[not-iterable]
+            updates.append(update)
+
+        # Verify the overridden response stream is returned
+        assert len(updates) == 2
+        assert updates[0].text == "overridden"
+        assert updates[1].text == " stream"
+
+    async def test_function_middleware_result_override(self, mock_function: FunctionTool) -> None:
+        """Test that function middleware can override result."""
+        override_result = "overridden function result"
+
+        class ResultOverrideMiddleware(FunctionMiddleware):
+            async def process(
+                self,
+                context: FunctionInvocationContext,
+                call_next: Callable[[], Awaitable[None]],
+            ) -> None:
+                # Execute the pipeline first, then override the result
+                await call_next()
+                context.result = override_result
+
+        middleware = ResultOverrideMiddleware()
+        pipeline = FunctionMiddlewarePipeline(middleware)
+        arguments = FunctionTestArgs(name="test")
+        context = FunctionInvocationContext(function=mock_function, arguments=arguments)
+
+        handler_called = False
+
+        async def final_handler(ctx: FunctionInvocationContext) -> str:
+            nonlocal handler_called
+            handler_called = True
+            return "original function result"
+
+        result = await pipeline.execute(context, final_handler)
+
+        # Verify the overridden result is returned
+        assert result == override_result
+        # Verify original handler was called since middleware called next()
+        assert handler_called
+
+    async def test_chat_agent_middleware_response_override(self) -> None:
+        """Test result override functionality with Agent integration."""
+        mock_chat_client = MockChatClient()
+
+        class ChatAgentResponseOverrideMiddleware(AgentMiddleware):
+            async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
+                # Always call next() first to allow execution
+                await call_next()
+                # Then conditionally override based on content
+                if any("special" in msg.text for msg in context.messages if msg.text):
+                    context.result = AgentResponse(
+                        messages=[Message(role="assistant", contents=["Special response from middleware!"])]
+                    )
+
+        # Create Agent with override middleware
+        middleware = ChatAgentResponseOverrideMiddleware()
+        agent = Agent(client=mock_chat_client, middleware=[middleware])  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+
+        # Test override case
+        override_messages = [Message(role="user", contents=["Give me a special response"])]
+        override_response = await agent.run(override_messages)
+        assert override_response.messages[0].text == "Special response from middleware!"
+        # Verify chat client was called since middleware called next()
+        assert mock_chat_client.call_count == 1
+
+        # Test normal case
+        normal_messages = [Message(role="user", contents=["Normal request"])]
+        normal_response = await agent.run(normal_messages)
+        assert normal_response.messages[0].text == "test response"
+        # Verify chat client was called for normal case
+        assert mock_chat_client.call_count == 2
+
+    async def test_chat_agent_middleware_streaming_override(self) -> None:
+        """Test streaming result override functionality with Agent integration."""
+        mock_chat_client = MockChatClient()
+
+        async def custom_stream() -> AsyncIterable[AgentResponseUpdate]:
+            yield AgentResponseUpdate(contents=[Content.from_text(text="Custom")])
+            yield AgentResponseUpdate(contents=[Content.from_text(text=" streaming")])
+            yield AgentResponseUpdate(contents=[Content.from_text(text=" response!")])
+
+        class ChatAgentStreamOverrideMiddleware(AgentMiddleware):
+            async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
+                # Check if we want to override BEFORE calling next to avoid creating unused streams
+                if any("custom stream" in msg.text for msg in context.messages if msg.text):
+                    context.result = ResponseStream(custom_stream())
+                    return  # Don't call next() - we're overriding the entire result
+                # Normal case - let the agent handle it
+                await call_next()
+
+        # Create Agent with override middleware
+        middleware = ChatAgentStreamOverrideMiddleware()
+        agent = Agent(client=mock_chat_client, middleware=[middleware])  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+
+        # Test streaming override case
+        override_messages = [Message(role="user", contents=["Give me a custom stream"])]
+        override_updates: list[AgentResponseUpdate] = []
+        async for update in agent.run(override_messages, stream=True):
+            override_updates.append(update)
+
+        assert len(override_updates) == 3
+        assert override_updates[0].text == "Custom"
+        assert override_updates[1].text == " streaming"
+        assert override_updates[2].text == " response!"
+
+        # Test normal streaming case
+        normal_messages = [Message(role="user", contents=["Normal streaming request"])]
+        normal_updates: list[AgentResponseUpdate] = []
+        async for update in agent.run(normal_messages, stream=True):
+            normal_updates.append(update)
+
+        assert len(normal_updates) == 2
+        assert normal_updates[0].text == "test streaming response "
+        assert normal_updates[1].text == "another update"
+
+    async def test_agent_middleware_conditional_no_next(self, mock_agent: SupportsAgentRun) -> None:
+        """Test that when agent middleware conditionally doesn't call next(), no execution happens."""
+
+        class ConditionalNoNextMiddleware(AgentMiddleware):
+            async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
+                # Only call next() if message contains "execute"
+                if any("execute" in msg.text for msg in context.messages if msg.text):
+                    await call_next()
+                # Otherwise, don't call next() - no execution should happen
+
+        middleware = ConditionalNoNextMiddleware()
+        pipeline = AgentMiddlewarePipeline(middleware)
+
+        handler_called = False
+
+        async def final_handler(ctx: AgentContext) -> AgentResponse:
+            nonlocal handler_called
+            handler_called = True
+            return AgentResponse(messages=[Message(role="assistant", contents=["executed response"])])
+
+        # Test case where next() is NOT called
+        no_execute_messages = [Message(role="user", contents=["Don't run this"])]
+        no_execute_context = AgentContext(agent=mock_agent, messages=no_execute_messages, stream=False)
+        no_execute_result = await pipeline.execute(no_execute_context, final_handler)
+
+        # When middleware doesn't call next(), result should be empty AgentResponse
+        assert no_execute_result is None
+        assert not handler_called
+
+        # Reset for next test
+        handler_called = False
+
+        # Test case where next() IS called
+        execute_messages = [Message(role="user", contents=["Please execute this"])]
+        execute_context = AgentContext(agent=mock_agent, messages=execute_messages, stream=False)
+        execute_result = await pipeline.execute(execute_context, final_handler)
+
+        assert execute_result is not None
+        assert execute_result.messages[0].text == "executed response"  # type: ignore[union-attr]  # ty: ignore[unresolved-attribute]
+        assert handler_called
+
+    async def test_function_middleware_conditional_no_next(self, mock_function: FunctionTool) -> None:
+        """Test that when function middleware conditionally doesn't call next(), no execution happens."""
+
+        class ConditionalNoNextFunctionMiddleware(FunctionMiddleware):
+            async def process(
+                self,
+                context: FunctionInvocationContext,
+                call_next: Callable[[], Awaitable[None]],
+            ) -> None:
+                # Only call next() if argument name contains "execute"
+                args = context.arguments
+                assert isinstance(args, FunctionTestArgs)
+                if "execute" in args.name:
+                    await call_next()
+                # Otherwise, don't call next() - no execution should happen
+
+        middleware = ConditionalNoNextFunctionMiddleware()
+        pipeline = FunctionMiddlewarePipeline(middleware)
+
+        handler_called = False
+
+        async def final_handler(ctx: FunctionInvocationContext) -> str:
+            nonlocal handler_called
+            handler_called = True
+            return "executed function result"
+
+        # Test case where next() is NOT called
+        no_execute_args = FunctionTestArgs(name="test_no_action")
+        no_execute_context = FunctionInvocationContext(function=mock_function, arguments=no_execute_args)
+        no_execute_result = await pipeline.execute(no_execute_context, final_handler)
+
+        # When middleware doesn't call next(), function result should be None (functions can return None)
+        assert no_execute_result is None
+        assert not handler_called
+        assert no_execute_context.result is None
+
+        # Reset for next test
+        handler_called = False
+
+        # Test case where next() IS called
+        execute_args = FunctionTestArgs(name="test_execute")
+        execute_context = FunctionInvocationContext(function=mock_function, arguments=execute_args)
+        execute_result = await pipeline.execute(execute_context, final_handler)
+
+        assert execute_result == "executed function result"
+        assert handler_called
+
+
+class TestResultObservability:
+    """Test cases for middleware result observability functionality."""
+
+    async def test_agent_middleware_response_observability(self, mock_agent: SupportsAgentRun) -> None:
+        """Test that middleware can observe response after execution."""
+        observed_responses: list[AgentResponse] = []
+
+        class ObservabilityMiddleware(AgentMiddleware):
+            async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
+                # Context should be empty before next()
+                assert context.result is None
+
+                # Call next to execute
+                await call_next()
+
+                # Context should now contain the response for observability
+                assert context.result is not None
+                assert isinstance(context.result, AgentResponse)
+                observed_responses.append(context.result)
+
+        middleware = ObservabilityMiddleware()
+        pipeline = AgentMiddlewarePipeline(middleware)
+        messages = [Message(role="user", contents=["test"])]
+        context = AgentContext(agent=mock_agent, messages=messages, stream=False)
+
+        async def final_handler(ctx: AgentContext) -> AgentResponse:
+            return AgentResponse(messages=[Message(role="assistant", contents=["executed response"])])
+
+        result = await pipeline.execute(context, final_handler)
+
+        # Verify response was observed
+        assert len(observed_responses) == 1
+        assert observed_responses[0].messages[0].text == "executed response"
+        assert result == observed_responses[0]
+
+    async def test_function_middleware_result_observability(self, mock_function: FunctionTool) -> None:
+        """Test that middleware can observe function result after execution."""
+        observed_results: list[str] = []
+
+        class ObservabilityMiddleware(FunctionMiddleware):
+            async def process(
+                self,
+                context: FunctionInvocationContext,
+                call_next: Callable[[], Awaitable[None]],
+            ) -> None:
+                # Context should be empty before next()
+                assert context.result is None
+
+                # Call next to execute
+                await call_next()
+
+                # Context should now contain the result for observability
+                assert context.result is not None
+                observed_results.append(context.result)
+
+        middleware = ObservabilityMiddleware()
+        pipeline = FunctionMiddlewarePipeline(middleware)
+        arguments = FunctionTestArgs(name="test")
+        context = FunctionInvocationContext(function=mock_function, arguments=arguments)
+
+        async def final_handler(ctx: FunctionInvocationContext) -> str:
+            return "executed function result"
+
+        result = await pipeline.execute(context, final_handler)
+
+        # Verify result was observed
+        assert len(observed_results) == 1
+        assert observed_results[0] == "executed function result"
+        assert result == observed_results[0]
+
+    async def test_agent_middleware_post_execution_override(self, mock_agent: SupportsAgentRun) -> None:
+        """Test that middleware can override response after observing execution."""
+
+        class PostExecutionOverrideMiddleware(AgentMiddleware):
+            async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
+                # Call next to execute first
+                await call_next()
+
+                # Now observe and conditionally override
+                assert context.result is not None
+                assert isinstance(context.result, AgentResponse)
+
+                if "modify" in context.result.messages[0].text:
+                    # Override after observing
+                    context.result = AgentResponse(
+                        messages=[Message(role="assistant", contents=["modified after execution"])]
+                    )
+
+        middleware = PostExecutionOverrideMiddleware()
+        pipeline = AgentMiddlewarePipeline(middleware)
+        messages = [Message(role="user", contents=["test"])]
+        context = AgentContext(agent=mock_agent, messages=messages, stream=False)
+
+        async def final_handler(ctx: AgentContext) -> AgentResponse:
+            return AgentResponse(messages=[Message(role="assistant", contents=["response to modify"])])
+
+        result = await pipeline.execute(context, final_handler)
+
+        # Verify response was modified after execution
+        assert result is not None
+        assert result.messages[0].text == "modified after execution"  # type: ignore[union-attr]  # ty: ignore[unresolved-attribute]
+
+    async def test_function_middleware_post_execution_override(self, mock_function: FunctionTool) -> None:
+        """Test that middleware can override function result after observing execution."""
+
+        class PostExecutionOverrideMiddleware(FunctionMiddleware):
+            async def process(
+                self,
+                context: FunctionInvocationContext,
+                call_next: Callable[[], Awaitable[None]],
+            ) -> None:
+                # Call next to execute first
+                await call_next()
+
+                # Now observe and conditionally override
+                assert context.result is not None
+
+                if "modify" in context.result:
+                    # Override after observing
+                    context.result = "modified after execution"
+
+        middleware = PostExecutionOverrideMiddleware()
+        pipeline = FunctionMiddlewarePipeline(middleware)
+        arguments = FunctionTestArgs(name="test")
+        context = FunctionInvocationContext(function=mock_function, arguments=arguments)
+
+        async def final_handler(ctx: FunctionInvocationContext) -> str:
+            return "result to modify"
+
+        result = await pipeline.execute(context, final_handler)
+
+        # Verify result was modified after execution
+        assert result == "modified after execution"
+
+
+@pytest.fixture
+def mock_agent() -> SupportsAgentRun:
+    """Mock agent for testing."""
+    agent = MagicMock(spec=SupportsAgentRun)
+    agent.name = "test_agent"
+    return agent
+
+
+@pytest.fixture
+def mock_function() -> FunctionTool:
+    """Mock function for testing."""
+    function = MagicMock(spec=FunctionTool)
+    function.name = "test_function"
+    return function

python/packages/core/tests/core/test_optional_dependencies.py

@@ -0,0 +1,181 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import sys
+
+import pytest
+
+import agent_framework
+import agent_framework.observability as observability
+from agent_framework import Agent
+
+
+def _hide_otel_sdk(monkeypatch: pytest.MonkeyPatch) -> None:
+    import builtins
+
+    real_import = builtins.__import__
+    for module_name in list(sys.modules):
+        if module_name == "opentelemetry.sdk" or module_name.startswith("opentelemetry.sdk."):
+            sys.modules.pop(module_name, None)
+
+    def _import_without_otel_sdk(
+        name: str,
+        globals_: dict[str, object] | None = None,
+        locals_: dict[str, object] | None = None,
+        fromlist: tuple[str, ...] = (),
+        level: int = 0,
+    ) -> object:
+        if name == "opentelemetry.sdk" or name.startswith("opentelemetry.sdk."):
+            raise ModuleNotFoundError(f"No module named '{name}'", name=name)
+        return real_import(name, globals_, locals_, fromlist, level)
+
+    monkeypatch.setattr(builtins, "__import__", _import_without_otel_sdk)
+
+
+def test_create_resource_requires_otel_sdk(monkeypatch: pytest.MonkeyPatch) -> None:
+    _hide_otel_sdk(monkeypatch)
+
+    with pytest.raises(ModuleNotFoundError, match="opentelemetry-sdk"):
+        observability.create_resource()
+
+
+def test_observability_settings_initializes_without_cached_resource(monkeypatch: pytest.MonkeyPatch) -> None:
+    _hide_otel_sdk(monkeypatch)
+
+    settings = observability.ObservabilitySettings()
+
+    assert not hasattr(settings, "_resource")
+
+
+def test_configure_otel_providers_requires_otel_sdk(monkeypatch: pytest.MonkeyPatch) -> None:
+    _hide_otel_sdk(monkeypatch)
+    for key in [
+        "OTEL_EXPORTER_OTLP_ENDPOINT",
+        "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT",
+        "OTEL_EXPORTER_OTLP_METRICS_ENDPOINT",
+        "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT",
+        "VS_CODE_EXTENSION_PORT",
+    ]:
+        monkeypatch.delenv(key, raising=False)
+
+    with pytest.raises(ModuleNotFoundError, match="opentelemetry-sdk"):
+        observability.configure_otel_providers()
+
+
+def test_agent_framework_mcp_exports_remain_importable_without_mcp(monkeypatch: pytest.MonkeyPatch) -> None:
+    import builtins
+
+    import agent_framework._mcp as mcp_module
+
+    real_import = builtins.__import__
+
+    def _import_without_mcp(
+        name: str,
+        globals_: dict[str, object] | None = None,
+        locals_: dict[str, object] | None = None,
+        fromlist: tuple[str, ...] = (),
+        level: int = 0,
+    ) -> object:
+        if name == "mcp" or name.startswith("mcp."):
+            raise ModuleNotFoundError("No module named 'mcp'")
+        return real_import(name, globals_, locals_, fromlist, level)
+
+    monkeypatch.setattr(builtins, "__import__", _import_without_mcp)
+    assert agent_framework.MCPStdioTool is mcp_module.MCPStdioTool
+
+    with pytest.raises(ModuleNotFoundError, match=r"Please install `mcp`\.$"):
+        agent_framework.MCPStdioTool(name="test", command="python").get_mcp_client()
+
+
+def test_mcp_streamable_http_tool_requires_mcp(monkeypatch: pytest.MonkeyPatch) -> None:
+    import builtins
+
+    real_import = builtins.__import__
+
+    def _import_without_mcp(
+        name: str,
+        globals_: dict[str, object] | None = None,
+        locals_: dict[str, object] | None = None,
+        fromlist: tuple[str, ...] = (),
+        level: int = 0,
+    ) -> object:
+        if name == "mcp" or name.startswith("mcp."):
+            raise ModuleNotFoundError("No module named 'mcp'")
+        return real_import(name, globals_, locals_, fromlist, level)
+
+    monkeypatch.setattr(builtins, "__import__", _import_without_mcp)
+
+    with pytest.raises(ModuleNotFoundError, match=r"Please install `mcp`\.$"):
+        agent_framework.MCPStreamableHTTPTool(name="test", url="https://example.com").get_mcp_client()
+
+
+def test_agent_as_mcp_server_requires_mcp(client, monkeypatch: pytest.MonkeyPatch) -> None:
+    import builtins
+
+    real_import = builtins.__import__
+
+    def _import_without_mcp(
+        name: str,
+        globals_: dict[str, object] | None = None,
+        locals_: dict[str, object] | None = None,
+        fromlist: tuple[str, ...] = (),
+        level: int = 0,
+    ) -> object:
+        if name == "mcp" or name.startswith("mcp."):
+            raise ModuleNotFoundError("No module named 'mcp'")
+        return real_import(name, globals_, locals_, fromlist, level)
+
+    monkeypatch.setattr(builtins, "__import__", _import_without_mcp)
+
+    agent = Agent(client=client)  # type: ignore[arg-type]
+
+    with pytest.raises(ModuleNotFoundError, match=r"Please install `mcp`\.$"):
+        agent.as_mcp_server()
+
+
+def test_mcp_websocket_tool_requires_ws_support(monkeypatch: pytest.MonkeyPatch) -> None:
+    import builtins
+
+    real_import = builtins.__import__
+
+    sys.modules.pop("mcp.client.websocket", None)
+
+    def _import_without_websocket_support(
+        name: str,
+        globals_: dict[str, object] | None = None,
+        locals_: dict[str, object] | None = None,
+        fromlist: tuple[str, ...] = (),
+        level: int = 0,
+    ) -> object:
+        if name == "mcp.client.websocket":
+            raise ModuleNotFoundError("No module named 'websockets'", name="websockets")
+        return real_import(name, globals_, locals_, fromlist, level)
+
+    monkeypatch.setattr(builtins, "__import__", _import_without_websocket_support)
+
+    with pytest.raises(ModuleNotFoundError, match=r"mcp\[ws\]"):
+        agent_framework.MCPWebsocketTool(name="test", url="wss://example.com").get_mcp_client()
+
+
+def test_mcp_websocket_tool_requires_mcp(monkeypatch: pytest.MonkeyPatch) -> None:
+    import builtins
+
+    real_import = builtins.__import__
+    sys.modules.pop("mcp.client.websocket", None)
+
+    def _import_without_mcp(
+        name: str,
+        globals_: dict[str, object] | None = None,
+        locals_: dict[str, object] | None = None,
+        fromlist: tuple[str, ...] = (),
+        level: int = 0,
+    ) -> object:
+        if name == "mcp.client.websocket":
+            raise ModuleNotFoundError("No module named 'mcp.client.websocket'", name="mcp.client.websocket")
+        return real_import(name, globals_, locals_, fromlist, level)
+
+    monkeypatch.setattr(builtins, "__import__", _import_without_mcp)
+
+    with pytest.raises(ModuleNotFoundError, match=r"agent-framework-core\[mcp\]|mcp\[ws\]"):
+        agent_framework.MCPWebsocketTool(name="test", url="wss://example.com").get_mcp_client()

python/packages/core/tests/core/test_root_exports.py

@@ -0,0 +1,44 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import ast
+from pathlib import Path
+
+import agent_framework
+
+
+def _stub_all() -> set[str]:
+    stub_path = Path(agent_framework.__file__).with_suffix(".pyi")
+    module = ast.parse(stub_path.read_text(encoding="utf-8"))
+    for node in module.body:
+        if isinstance(node, ast.Assign):
+            for target in node.targets:
+                if isinstance(target, ast.Name) and target.id == "__all__":
+                    return set(ast.literal_eval(node.value))
+    raise AssertionError("__all__ not found in agent_framework root stub")
+
+
+def test_root_all_matches_stub_all() -> None:
+    assert set(agent_framework.__all__) == _stub_all()
+
+
+def test_root_star_import_loads_representative_symbols() -> None:
+    namespace: dict[str, object] = {}
+    exec("from agent_framework import *", namespace)
+
+    assert namespace["Agent"] is agent_framework.Agent
+    assert namespace["Message"] is agent_framework.Message
+    assert namespace["tool"] is agent_framework.tool
+    assert namespace["FileStoreEntry"] is agent_framework.FileStoreEntry
+    assert namespace["SkillsSourceContext"] is agent_framework.SkillsSourceContext
+
+
+def test_root_from_import_representative_symbols() -> None:
+    from agent_framework import Agent, FileStoreEntry, Message, SkillsSourceContext, tool
+
+    assert Agent is agent_framework.Agent
+    assert Message is agent_framework.Message
+    assert tool is agent_framework.tool
+    assert FileStoreEntry is agent_framework.FileStoreEntry
+    assert SkillsSourceContext is agent_framework.SkillsSourceContext

python/packages/core/tests/core/test_serializable_mixin.py

@@ -0,0 +1,529 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for SerializationMixin functionality."""
+
+import logging
+from typing import Any
+
+from agent_framework._serialization import SerializationMixin
+
+
+class TestSerializationMixin:
+    """Test SerializationMixin serialization, deserialization, and dependency injection."""
+
+    def test_basic_serialization(self):
+        """Test basic to_dict and from_dict functionality."""
+
+        class TestClass(SerializationMixin):
+            def __init__(self, value: str, number: int):
+                self.value = value
+                self.number = number
+
+        obj = TestClass(value="test", number=42)
+        data = obj.to_dict()
+
+        assert data["type"] == "test_class"
+        assert data["value"] == "test"
+        assert data["number"] == 42
+
+        restored = TestClass.from_dict(data)
+        assert restored.value == "test"
+        assert restored.number == 42
+
+    def test_injectable_dependency_no_warning(self, caplog):
+        """Test that injectable dependencies don't trigger debug logging."""
+
+        class TestClass(SerializationMixin):
+            INJECTABLE = {"client"}
+
+            def __init__(self, value: str, client: Any = None):
+                self.value = value
+                self.client = client
+
+        mock_client = "mock_client_instance"
+
+        with caplog.at_level(logging.DEBUG):
+            obj = TestClass.from_dict(
+                {"type": "test_class", "value": "test"},
+                dependencies={"test_class": {"client": mock_client}},
+            )
+
+        assert obj.value == "test"
+        assert obj.client == mock_client
+        # No debug message should be logged for injectable dependency
+        assert not any("is not in INJECTABLE set" in record.message for record in caplog.records)
+
+    def test_non_injectable_dependency_logs_debug(self, caplog):
+        """Test that non-injectable dependencies trigger debug logging."""
+
+        class TestClass(SerializationMixin):
+            INJECTABLE = {"client"}
+
+            def __init__(self, value: str, other: Any = None):
+                self.value = value
+                self.other = other
+
+        mock_other = "mock_other_instance"
+
+        with caplog.at_level(logging.DEBUG):
+            obj = TestClass.from_dict(
+                {"type": "test_class", "value": "test"},
+                dependencies={"test_class": {"other": mock_other}},
+            )
+
+        assert obj.value == "test"
+        assert obj.other == mock_other
+        # Debug message should be logged for non-injectable dependency
+        debug_messages = [record.message for record in caplog.records if record.levelname == "DEBUG"]
+        assert any("is not in INJECTABLE set" in msg for msg in debug_messages)
+        assert any("other" in msg for msg in debug_messages)
+        assert any("client" in msg for msg in debug_messages)  # Should mention available injectable
+
+    def test_multiple_dependencies_mixed_injectable(self, caplog):
+        """Test with both injectable and non-injectable dependencies."""
+
+        class TestClass(SerializationMixin):
+            INJECTABLE = {"client", "logger"}
+
+            def __init__(
+                self,
+                value: str,
+                client: Any = None,
+                logger: Any = None,
+                other: Any = None,
+            ):
+                self.value = value
+                self.client = client
+                self.logger = logger
+                self.other = other
+
+        mock_client = "mock_client"
+        mock_logger = "mock_logger"
+        mock_other = "mock_other"
+
+        with caplog.at_level(logging.DEBUG):
+            obj = TestClass.from_dict(
+                {"type": "test_class", "value": "test"},
+                dependencies={
+                    "test_class": {
+                        "client": mock_client,
+                        "logger": mock_logger,
+                        "other": mock_other,
+                    }
+                },
+            )
+
+        assert obj.value == "test"
+        assert obj.client == mock_client
+        assert obj.logger == mock_logger
+        assert obj.other == mock_other
+
+        # Only 'other' should trigger debug logging
+        debug_messages = [record.message for record in caplog.records if record.levelname == "DEBUG"]
+        assert any("other" in msg and "is not in INJECTABLE set" in msg for msg in debug_messages)
+        # 'client' and 'logger' should not be mentioned as non-injectable dependencies
+        assert not any("Dependency 'client'" in msg and "is not in INJECTABLE set" in msg for msg in debug_messages)
+        assert not any("Dependency 'logger'" in msg and "is not in INJECTABLE set" in msg for msg in debug_messages)
+
+    def test_no_injectable_set_defined(self, caplog):
+        """Test behavior when INJECTABLE is not defined (empty set default)."""
+
+        class TestClass(SerializationMixin):
+            def __init__(self, value: str, client: Any = None):
+                self.value = value
+                self.client = client
+
+        mock_client = "mock_client"
+
+        with caplog.at_level(logging.DEBUG):
+            obj = TestClass.from_dict(
+                {"type": "test_class", "value": "test"},
+                dependencies={"test_class": {"client": mock_client}},
+            )
+
+        assert obj.value == "test"
+        assert obj.client == mock_client
+        # Should log debug message since INJECTABLE is empty by default
+        debug_messages = [record.message for record in caplog.records if record.levelname == "DEBUG"]
+        assert any("client" in msg and "is not in INJECTABLE set" in msg for msg in debug_messages)
+
+    def test_default_exclude_serialization(self):
+        """Test that DEFAULT_EXCLUDE fields are not included in to_dict()."""
+
+        class TestClass(SerializationMixin):
+            DEFAULT_EXCLUDE = {"secret"}
+
+            def __init__(self, value: str, secret: str):
+                self.value = value
+                self.secret = secret
+
+        obj = TestClass(value="test", secret="hidden")
+        data = obj.to_dict()
+
+        assert "value" in data
+        assert "secret" not in data
+        assert data["value"] == "test"
+
+    def test_roundtrip_with_injectable_dependency(self):
+        """Test full roundtrip serialization/deserialization with injectable dependency."""
+
+        class TestClass(SerializationMixin):
+            INJECTABLE = {"client"}
+            DEFAULT_EXCLUDE = {"client"}
+
+            def __init__(self, value: str, number: int, client: Any = None):
+                self.value = value
+                self.number = number
+                self.client = client
+
+        mock_client = "mock_client"
+        obj = TestClass(value="test", number=42, client=mock_client)
+
+        # Serialize
+        data = obj.to_dict()
+        assert data["value"] == "test"
+        assert data["number"] == 42
+        assert "client" not in data  # Excluded from serialization
+
+        # Deserialize with dependency injection
+        restored = TestClass.from_dict(data, dependencies={"test_class": {"client": mock_client}})
+        assert restored.value == "test"
+        assert restored.number == 42
+        assert restored.client == mock_client
+
+    def test_exclude_none_in_to_dict(self):
+        """Test that exclude_none parameter removes None values from to_dict()."""
+
+        class TestClass(SerializationMixin):
+            def __init__(self, value: str, optional: str | None = None):
+                self.value = value
+                self.optional = optional
+
+        obj = TestClass(value="test", optional=None)
+        data = obj.to_dict(exclude_none=True)
+
+        assert data["value"] == "test"
+        assert "optional" not in data
+
+    def test_to_dict_with_nested_serialization_protocol(self):
+        """Test to_dict handles nested SerializationProtocol objects."""
+
+        class InnerClass(SerializationMixin):
+            def __init__(self, inner_value: str):
+                self.inner_value = inner_value
+
+        class OuterClass(SerializationMixin):
+            def __init__(self, outer_value: str, inner: Any = None):
+                self.outer_value = outer_value
+                self.inner = inner
+
+        inner = InnerClass(inner_value="inner_test")
+        outer = OuterClass(outer_value="outer_test", inner=inner)
+        data = outer.to_dict()
+
+        assert data["outer_value"] == "outer_test"
+        assert data["inner"]["inner_value"] == "inner_test"
+
+    def test_to_dict_with_list_of_serialization_protocol(self):
+        """Test to_dict handles lists containing SerializationProtocol objects."""
+
+        class ItemClass(SerializationMixin):
+            def __init__(self, name: str):
+                self.name = name
+
+        class ContainerClass(SerializationMixin):
+            def __init__(self, items: list):
+                self.items = items
+
+        items = [ItemClass(name="item1"), ItemClass(name="item2")]
+        container = ContainerClass(items=items)
+        data = container.to_dict()
+
+        assert len(data["items"]) == 2
+        assert data["items"][0]["name"] == "item1"
+        assert data["items"][1]["name"] == "item2"
+
+    def test_to_dict_skips_non_serializable_in_list(self, caplog):
+        """Test to_dict skips non-serializable items in lists with debug logging."""
+
+        class NonSerializable:
+            pass
+
+        class TestClass(SerializationMixin):
+            def __init__(self, items: list):
+                self.items = items
+
+        obj = TestClass(items=["serializable", NonSerializable()])
+
+        with caplog.at_level(logging.DEBUG):
+            data = obj.to_dict()
+
+        # Should only contain the serializable item
+        assert len(data["items"]) == 1
+        assert data["items"][0] == "serializable"
+
+    def test_to_dict_with_dict_containing_serialization_protocol(self):
+        """Test to_dict handles dicts containing SerializationProtocol values."""
+
+        class ItemClass(SerializationMixin):
+            def __init__(self, name: str):
+                self.name = name
+
+        class ContainerClass(SerializationMixin):
+            def __init__(self, items_dict: dict):
+                self.items_dict = items_dict
+
+        items = {"a": ItemClass(name="item1"), "b": ItemClass(name="item2")}
+        container = ContainerClass(items_dict=items)
+        data = container.to_dict()
+
+        assert data["items_dict"]["a"]["name"] == "item1"
+        assert data["items_dict"]["b"]["name"] == "item2"
+
+    def test_to_dict_with_datetime_in_dict(self):
+        """Test to_dict converts datetime objects in dicts to strings."""
+        from datetime import datetime
+
+        class TestClass(SerializationMixin):
+            def __init__(self, metadata: dict):
+                self.metadata = metadata
+
+        now = datetime(2025, 1, 27, 12, 0, 0)
+        obj = TestClass(metadata={"created_at": now})
+        data = obj.to_dict()
+
+        assert isinstance(data["metadata"]["created_at"], str)
+
+    def test_to_dict_skips_non_serializable_in_dict(self, caplog):
+        """Test to_dict skips non-serializable values in dicts with debug logging."""
+
+        class NonSerializable:
+            pass
+
+        class TestClass(SerializationMixin):
+            def __init__(self, metadata: dict):
+                self.metadata = metadata
+
+        obj = TestClass(metadata={"valid": "value", "invalid": NonSerializable()})
+
+        with caplog.at_level(logging.DEBUG):
+            data = obj.to_dict()
+
+        assert data["metadata"]["valid"] == "value"
+        assert "invalid" not in data["metadata"]
+
+    def test_to_dict_skips_non_serializable_attributes(self, caplog):
+        """Test to_dict skips non-serializable top-level attributes."""
+
+        class TestClass(SerializationMixin):
+            def __init__(self, value: str, func: Any = None):
+                self.value = value
+                self.func = func
+
+        obj = TestClass(value="test", func=lambda x: x)
+
+        with caplog.at_level(logging.DEBUG):
+            data = obj.to_dict()
+
+        assert data["value"] == "test"
+        assert "func" not in data
+
+    def test_from_dict_without_type_in_data(self):
+        """Test from_dict uses class TYPE when no type field in data."""
+
+        class TestClass(SerializationMixin):
+            TYPE = "my_custom_type"
+
+            def __init__(self, value: str):
+                self.value = value
+
+        # Data without 'type' field - class TYPE should be used for type identifier
+        data = {"value": "test"}
+
+        obj = TestClass.from_dict(data)
+        assert obj.value == "test"
+
+        # Verify to_dict includes the type
+        out = obj.to_dict()
+        assert out["type"] == "my_custom_type"
+
+    def test_from_json(self):
+        """Test from_json deserializes JSON string."""
+
+        class TestClass(SerializationMixin):
+            def __init__(self, value: str):
+                self.value = value
+
+        json_str = '{"type": "test_class", "value": "test_value"}'
+        obj = TestClass.from_json(json_str)
+
+        assert obj.value == "test_value"
+
+    def test_get_type_identifier_with_instance_type(self):
+        """Test _get_type_identifier uses instance 'type' attribute."""
+
+        class TestClass(SerializationMixin):
+            def __init__(self, value: str):
+                self.value = value
+                self.type = "custom_type"
+
+        obj = TestClass(value="test")
+        data = obj.to_dict()
+
+        assert data["type"] == "custom_type"
+
+    def test_get_type_identifier_with_class_TYPE(self):
+        """Test _get_type_identifier uses class TYPE constant."""
+
+        class TestClass(SerializationMixin):
+            TYPE = "class_level_type"
+
+            def __init__(self, value: str):
+                self.value = value
+
+        obj = TestClass(value="test")
+        data = obj.to_dict()
+
+        assert data["type"] == "class_level_type"
+
+    def test_instance_specific_dependency_injection(self):
+        """Test instance-specific dependency injection with field:name format."""
+
+        class TestClass(SerializationMixin):
+            INJECTABLE = {"config"}
+
+            def __init__(self, name: str, config: Any = None):
+                self.name = name
+                self.config = config
+
+        dependencies = {
+            "test_class": {
+                "name:special_instance": {"config": "special_config"},
+            }
+        }
+
+        # This should match the instance-specific dependency
+        obj = TestClass.from_dict({"type": "test_class", "name": "special_instance"}, dependencies=dependencies)
+
+        assert obj.name == "special_instance"
+        assert obj.config == "special_config"
+
+    def test_dependency_dict_merging(self):
+        """Test that dict dependencies are merged with existing dict kwargs."""
+
+        class TestClass(SerializationMixin):
+            INJECTABLE = {"options"}
+
+            def __init__(self, value: str, options: dict | None = None):
+                self.value = value
+                self.options = options or {}
+
+        # Existing options in data
+        data = {"type": "test_class", "value": "test", "options": {"existing": "value"}}
+        # Additional options from dependencies
+        dependencies = {"test_class": {"options": {"injected": "option"}}}
+
+        obj = TestClass.from_dict(data, dependencies=dependencies)
+
+        assert obj.options["existing"] == "value"
+        assert obj.options["injected"] == "option"
+
+    def test_deepcopy_preserves_shallow_copy_fields_by_reference(self):
+        """Test that deepcopy keeps _SHALLOW_COPY_FIELDS fields as shallow references."""
+        import copy
+
+        class NonCopyable:
+            def __deepcopy__(self, memo):
+                raise TypeError("cannot deepcopy")
+
+        class TestClass(SerializationMixin):
+            _SHALLOW_COPY_FIELDS = {"raw_representation", "other_opaque"}
+
+            def __init__(self, items: list, raw_representation: Any = None, other_opaque: Any = None):
+                self.items = items
+                self.raw_representation = raw_representation
+                self.other_opaque = other_opaque
+
+        raw = NonCopyable()
+        opaque = NonCopyable()
+        original_items = ["a", "b"]
+        obj = TestClass(items=original_items, raw_representation=raw, other_opaque=opaque)
+        cloned = copy.deepcopy(obj)
+
+        # _SHALLOW_COPY_FIELDS fields should be the same object (shallow copy)
+        assert cloned.raw_representation is raw
+        assert cloned.other_opaque is opaque
+        # Normal attributes should be independent copies
+        assert cloned.items is not original_items
+        assert cloned.items == ["a", "b"]
+
+    def test_deepcopy_deep_copies_non_shallow_copy_fields(self):
+        """Test that deepcopy fully copies fields not in _SHALLOW_COPY_FIELDS."""
+        import copy
+
+        class TestClass(SerializationMixin):
+            _SHALLOW_COPY_FIELDS = {"raw_representation"}
+
+            def __init__(self, items: list, raw_representation: Any = None):
+                self.items = items
+                self.raw_representation = raw_representation
+
+        original_list = ["a", "b"]
+        obj = TestClass(items=original_list, raw_representation="raw")
+        cloned = copy.deepcopy(obj)
+
+        # list should be a new object
+        assert cloned.items is not original_list
+        assert cloned.items == ["a", "b"]
+        # raw_representation should be the same object
+        assert cloned.raw_representation is obj.raw_representation
+
+    def test_deepcopy_deep_copies_default_exclude_fields(self):
+        """Test that DEFAULT_EXCLUDE fields are deep-copied unless also in _SHALLOW_COPY_FIELDS."""
+        import copy
+
+        class TestClass(SerializationMixin):
+            DEFAULT_EXCLUDE = {"additional_properties"}
+
+            def __init__(self, items: list, additional_properties: dict | None = None):
+                self.items = items
+                self.additional_properties = additional_properties or {}
+
+        original_props = {"key": "value"}
+        obj = TestClass(items=["a"], additional_properties=original_props)
+        cloned = copy.deepcopy(obj)
+
+        # DEFAULT_EXCLUDE field should be deep-copied (independent copy)
+        assert cloned.additional_properties is not original_props
+        assert cloned.additional_properties == {"key": "value"}
+
+    def test_deepcopy_shallow_copy_fields_override_default_exclude(self):
+        """Test that _SHALLOW_COPY_FIELDS controls deepcopy independently of DEFAULT_EXCLUDE."""
+        import copy
+
+        class NonCopyable:
+            def __deepcopy__(self, memo):
+                raise TypeError("cannot deepcopy")
+
+        class TestClass(SerializationMixin):
+            DEFAULT_EXCLUDE = {"opaque", "additional_properties"}
+            _SHALLOW_COPY_FIELDS = {"opaque"}
+
+            def __init__(self, items: list, opaque: Any = None, additional_properties: dict | None = None):
+                self.items = items
+                self.opaque = opaque
+                self.additional_properties = additional_properties or {}
+
+        opaque = NonCopyable()
+        original_props = {"key": "value"}
+        obj = TestClass(items=["a"], opaque=opaque, additional_properties=original_props)
+        cloned = copy.deepcopy(obj)
+
+        # Field in both DEFAULT_EXCLUDE and _SHALLOW_COPY_FIELDS: shallow-copied
+        assert cloned.opaque is opaque
+        # Field in DEFAULT_EXCLUDE only: deep-copied
+        assert cloned.additional_properties is not original_props
+        assert cloned.additional_properties == {"key": "value"}
+        # Normal field: deep-copied
+        assert cloned.items is not obj.items
+        assert cloned.items == ["a"]

python/packages/core/tests/core/test_sessions.py

@@ -0,0 +1,798 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import json
+import threading
+import time
+from collections.abc import Awaitable, Callable, Sequence
+from pathlib import Path
+from typing import Any
+
+import pytest
+
+from agent_framework import (
+    AgentContext,
+    AgentSession,
+    ChatContext,
+    ContextProvider,
+    ExperimentalFeature,
+    FileHistoryProvider,
+    HistoryProvider,
+    InMemoryHistoryProvider,
+    Message,
+    SessionContext,
+    agent_middleware,
+    chat_middleware,
+)
+from agent_framework._sessions import LOCAL_HISTORY_CONVERSATION_ID, is_local_history_conversation_id
+from agent_framework.exceptions import MiddlewareException
+
+# ---------------------------------------------------------------------------
+# SessionContext tests
+# ---------------------------------------------------------------------------
+
+
+class TestSessionContext:
+    def test_init_defaults(self) -> None:
+        ctx = SessionContext(input_messages=[])
+        assert ctx.session_id is None
+        assert ctx.service_session_id is None
+        assert ctx.input_messages == []
+        assert ctx.context_messages == {}
+        assert ctx.instructions == []
+        assert ctx.tools == []
+        assert ctx.response is None
+        assert ctx.options == {}
+        assert ctx.metadata == {}
+
+    def test_extend_messages_creates_key(self) -> None:
+        ctx = SessionContext(input_messages=[])
+        msg = Message(role="user", contents=["hello"])
+        ctx.extend_messages("rag", [msg])
+        assert "rag" in ctx.context_messages
+        assert len(ctx.context_messages["rag"]) == 1
+        assert ctx.context_messages["rag"][0].text == "hello"
+
+    def test_extend_messages_appends_to_existing(self) -> None:
+        ctx = SessionContext(input_messages=[])
+        msg1 = Message(role="user", contents=["first"])
+        msg2 = Message(role="user", contents=["second"])
+        ctx.extend_messages("src", [msg1])
+        ctx.extend_messages("src", [msg2])
+        assert len(ctx.context_messages["src"]) == 2
+
+    def test_extend_messages_preserves_source_order(self) -> None:
+        ctx = SessionContext(input_messages=[])
+        ctx.extend_messages("a", [Message(role="user", contents=["a"])])
+        ctx.extend_messages("b", [Message(role="user", contents=["b"])])
+        ctx.extend_messages("c", [Message(role="user", contents=["c"])])
+        assert list(ctx.context_messages.keys()) == ["a", "b", "c"]
+
+    def test_extend_messages_sets_attribution(self) -> None:
+        ctx = SessionContext(input_messages=[])
+        msg = Message(role="system", contents=["context"])
+        ctx.extend_messages("rag", [msg])
+        stored = ctx.context_messages["rag"][0]
+        assert stored.additional_properties["_attribution"] == {"source_id": "rag"}
+        # Original message is not mutated
+        assert "_attribution" not in msg.additional_properties
+
+    def test_extend_messages_does_not_overwrite_existing_attribution(self) -> None:
+        ctx = SessionContext(input_messages=[])
+        msg = Message(
+            role="system", contents=["context"], additional_properties={"_attribution": {"source_id": "custom"}}
+        )
+        ctx.extend_messages("rag", [msg])
+        stored = ctx.context_messages["rag"][0]
+        assert stored.additional_properties["_attribution"] == {"source_id": "custom"}
+
+    def test_extend_messages_copies_messages(self) -> None:
+        ctx = SessionContext(input_messages=[])
+        msg = Message(role="user", contents=["hello"])
+        ctx.extend_messages("src", [msg])
+        stored = ctx.context_messages["src"][0]
+        assert stored is not msg
+        assert stored.text == "hello"
+        # Mutating stored copy does not affect original
+        stored.additional_properties["extra"] = True
+        assert "extra" not in msg.additional_properties
+
+    def test_extend_messages_sender_sets_source_type(self) -> None:
+        class MyProvider:
+            source_id = "rag"
+
+        ctx = SessionContext(input_messages=[])
+        msg = Message(role="system", contents=["ctx"])
+        ctx.extend_messages(MyProvider(), [msg])
+        stored = ctx.context_messages["rag"][0]
+        assert stored.additional_properties["_attribution"] == {"source_id": "rag", "source_type": "MyProvider"}
+
+    def test_extend_instructions_string(self) -> None:
+        ctx = SessionContext(input_messages=[])
+        ctx.extend_instructions("sys", "Be helpful")
+        assert ctx.instructions == ["Be helpful"]
+
+    def test_extend_instructions_sequence(self) -> None:
+        ctx = SessionContext(input_messages=[])
+        ctx.extend_instructions("sys", ["Be helpful", "Be concise"])
+        assert ctx.instructions == ["Be helpful", "Be concise"]
+
+    def test_extend_middleware_creates_key_and_appends(self) -> None:
+        ctx = SessionContext(input_messages=[])
+
+        @chat_middleware
+        async def first_middleware(context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None:
+            await call_next()
+
+        @chat_middleware
+        async def second_middleware(context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None:
+            await call_next()
+
+        ctx.extend_middleware("rag", first_middleware)
+        ctx.extend_middleware("rag", [second_middleware])
+
+        assert ctx.middleware["rag"] == [first_middleware, second_middleware]
+        assert ctx.get_middleware() == [first_middleware, second_middleware]
+
+    def test_extend_middleware_preserves_source_order(self) -> None:
+        ctx = SessionContext(input_messages=[])
+
+        @chat_middleware
+        async def first_middleware(context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None:
+            await call_next()
+
+        @chat_middleware
+        async def second_middleware(context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None:
+            await call_next()
+
+        ctx.extend_middleware("a", first_middleware)
+        ctx.extend_middleware("b", second_middleware)
+
+        assert list(ctx.middleware.keys()) == ["a", "b"]
+        assert ctx.get_middleware() == [first_middleware, second_middleware]
+
+    def test_extend_middleware_rejects_agent_middleware(self) -> None:
+        ctx = SessionContext(input_messages=[])
+
+        @agent_middleware
+        async def provider_agent_middleware(context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
+            await call_next()
+
+        with pytest.raises(MiddlewareException, match="Context providers may only add chat or function middleware"):
+            ctx.extend_middleware("rag", provider_agent_middleware)
+
+    def test_get_messages_all(self) -> None:
+        ctx = SessionContext(input_messages=[])
+        ctx.extend_messages("a", [Message(role="user", contents=["a"])])
+        ctx.extend_messages("b", [Message(role="user", contents=["b"])])
+        result = ctx.get_messages()
+        assert len(result) == 2
+        assert result[0].text == "a"
+        assert result[1].text == "b"
+
+    def test_get_messages_filter_sources(self) -> None:
+        ctx = SessionContext(input_messages=[])
+        ctx.extend_messages("a", [Message(role="user", contents=["a"])])
+        ctx.extend_messages("b", [Message(role="user", contents=["b"])])
+        result = ctx.get_messages(sources=["a"])  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+        assert len(result) == 1
+        assert result[0].text == "a"
+
+    def test_get_messages_exclude_sources(self) -> None:
+        ctx = SessionContext(input_messages=[])
+        ctx.extend_messages("a", [Message(role="user", contents=["a"])])
+        ctx.extend_messages("b", [Message(role="user", contents=["b"])])
+        result = ctx.get_messages(exclude_sources=["a"])  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+        assert len(result) == 1
+        assert result[0].text == "b"
+
+    def test_get_messages_include_input(self) -> None:
+        input_msg = Message(role="user", contents=["input"])
+        ctx = SessionContext(input_messages=[input_msg])
+        ctx.extend_messages("a", [Message(role="user", contents=["context"])])
+        result = ctx.get_messages(include_input=True)
+        assert len(result) == 2
+        assert result[1].text == "input"
+
+    def test_get_messages_include_response(self) -> None:
+        from agent_framework import AgentResponse
+
+        ctx = SessionContext(input_messages=[])
+        ctx._response = AgentResponse(messages=[Message(role="assistant", contents=["reply"])])
+        result = ctx.get_messages(include_response=True)
+        assert len(result) == 1
+        assert result[0].text == "reply"
+
+    def test_response_readonly(self) -> None:
+        ctx = SessionContext(input_messages=[])
+        assert ctx.response is None
+        # Can set via _response internally
+        from agent_framework import AgentResponse
+
+        resp = AgentResponse(messages=[])
+        ctx._response = resp
+        assert ctx.response is resp
+
+    def test_local_history_conversation_id_sentinel(self) -> None:
+        assert is_local_history_conversation_id(LOCAL_HISTORY_CONVERSATION_ID) is True
+        assert is_local_history_conversation_id("some_other_id") is False
+
+
+# ---------------------------------------------------------------------------
+# ContextProvider tests
+# ---------------------------------------------------------------------------
+
+
+class TestContextProvider:
+    def test_source_id_required(self) -> None:
+        provider = ContextProvider(source_id="test")
+        assert provider.source_id == "test"
+
+    async def test_before_run_is_noop(self) -> None:
+        provider = ContextProvider(source_id="test")
+        session = AgentSession()
+        ctx = SessionContext(input_messages=[])
+        # Should not raise
+        await provider.before_run(agent=None, session=session, context=ctx, state={})  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+
+    async def test_after_run_is_noop(self) -> None:
+        provider = ContextProvider(source_id="test")
+        session = AgentSession()
+        ctx = SessionContext(input_messages=[])
+        await provider.after_run(agent=None, session=session, context=ctx, state={})  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+
+
+# ---------------------------------------------------------------------------
+# HistoryProvider tests
+# ---------------------------------------------------------------------------
+
+
+class ConcreteHistoryProvider(HistoryProvider):
+    """Concrete test implementation."""
+
+    def __init__(self, source_id: str, stored_messages: list[Message] | None = None, **kwargs) -> None:
+        super().__init__(source_id, **kwargs)
+        self.stored: list[Message] = []
+        self._stored_messages = stored_messages or []
+
+    async def get_messages(self, session_id: str | None, *, state=None, **kwargs) -> list[Message]:
+        return list(self._stored_messages)
+
+    async def save_messages(self, session_id: str | None, messages: Sequence[Message], *, state=None, **kwargs) -> None:
+        self.stored.extend(messages)
+
+
+class TestHistoryProviderBase:
+    def test_default_flags(self) -> None:
+        provider = ConcreteHistoryProvider("mem")
+        assert provider.load_messages is True
+        assert provider.store_outputs is True
+        assert provider.store_inputs is True
+        assert provider.store_context_messages is False
+        assert provider.store_context_from is None
+
+    def test_custom_flags(self) -> None:
+        provider = ConcreteHistoryProvider(
+            "audit",
+            load_messages=False,
+            store_inputs=False,
+            store_context_messages=True,
+            store_context_from={"rag"},
+        )
+        assert provider.load_messages is False
+        assert provider.store_inputs is False
+        assert provider.store_context_messages is True
+        assert provider.store_context_from == {"rag"}
+
+    async def test_before_run_loads_messages(self) -> None:
+        msgs = [Message(role="user", contents=["history"])]
+        provider = ConcreteHistoryProvider("mem", stored_messages=msgs)
+        session = AgentSession()
+        ctx = SessionContext(session_id="s1", input_messages=[])
+        await provider.before_run(agent=None, session=session, context=ctx, state={})  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+        assert len(ctx.context_messages["mem"]) == 1
+        assert ctx.context_messages["mem"][0].text == "history"
+
+    async def test_after_run_stores_inputs_and_responses(self) -> None:
+        from agent_framework import AgentResponse
+
+        provider = ConcreteHistoryProvider("mem")
+        session = AgentSession()
+        input_msg = Message(role="user", contents=["hello"])
+        resp_msg = Message(role="assistant", contents=["hi"])
+        ctx = SessionContext(session_id="s1", input_messages=[input_msg])
+        ctx._response = AgentResponse(messages=[resp_msg])
+        await provider.after_run(agent=None, session=session, context=ctx, state={})  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+        assert len(provider.stored) == 2
+        assert provider.stored[0].text == "hello"
+        assert provider.stored[1].text == "hi"
+
+    async def test_after_run_stores_coalesced_code_interpreter_chunks(self) -> None:
+        from agent_framework import AgentResponse, AgentResponseUpdate, Content
+
+        provider = ConcreteHistoryProvider("mem", store_inputs=False)
+        updates = [
+            AgentResponseUpdate(
+                role="assistant",
+                contents=[
+                    Content.from_code_interpreter_tool_result(
+                        call_id="ci_123",
+                        outputs=[],
+                    )
+                ],
+            ),
+            AgentResponseUpdate(
+                contents=[
+                    Content.from_code_interpreter_tool_call(
+                        call_id="ci_123",
+                        inputs=[Content.from_text(text="import")],
+                        additional_properties={"sequence_number": 1},
+                    )
+                ],
+            ),
+            AgentResponseUpdate(
+                contents=[
+                    Content.from_code_interpreter_tool_call(
+                        call_id="ci_123",
+                        inputs=[Content.from_text(text=" pandas")],
+                        additional_properties={"sequence_number": 2},
+                    )
+                ],
+            ),
+            AgentResponseUpdate(
+                contents=[
+                    Content.from_code_interpreter_tool_call(
+                        call_id="ci_123",
+                        inputs=[Content.from_text(text="import pandas as pd")],
+                        additional_properties={"sequence_number": 3},
+                    )
+                ],
+            ),
+        ]
+        ctx = SessionContext(session_id="s1", input_messages=[Message(role="user", contents=["make a sheet"])])
+        ctx._response = AgentResponse.from_updates(updates)
+
+        await provider.after_run(agent=None, session=AgentSession(), context=ctx, state={})  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+
+        assert len(provider.stored) == 1
+        stored_contents = provider.stored[0].contents
+        calls = [content for content in stored_contents if content.type == "code_interpreter_tool_call"]
+        results = [content for content in stored_contents if content.type == "code_interpreter_tool_result"]
+        assert len(calls) == 1
+        assert len(results) == 1
+        assert calls[0].inputs is not None
+        assert len(calls[0].inputs) == 1
+        assert calls[0].inputs[0].text == "import pandas as pd"
+
+    async def test_after_run_skips_inputs_when_disabled(self) -> None:
+        from agent_framework import AgentResponse
+
+        provider = ConcreteHistoryProvider("mem", store_inputs=False)
+        ctx = SessionContext(session_id="s1", input_messages=[Message(role="user", contents=["hello"])])
+        ctx._response = AgentResponse(messages=[Message(role="assistant", contents=["hi"])])
+        await provider.after_run(agent=None, session=AgentSession(), context=ctx, state={})  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+        assert len(provider.stored) == 1
+        assert provider.stored[0].text == "hi"
+
+    async def test_after_run_skips_responses_when_disabled(self) -> None:
+        from agent_framework import AgentResponse
+
+        provider = ConcreteHistoryProvider("mem", store_outputs=False)
+        ctx = SessionContext(session_id="s1", input_messages=[Message(role="user", contents=["hello"])])
+        ctx._response = AgentResponse(messages=[Message(role="assistant", contents=["hi"])])
+        await provider.after_run(agent=None, session=AgentSession(), context=ctx, state={})  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+        assert len(provider.stored) == 1
+        assert provider.stored[0].text == "hello"
+
+    async def test_after_run_stores_context_messages(self) -> None:
+        from agent_framework import AgentResponse
+
+        provider = ConcreteHistoryProvider("audit", load_messages=False, store_context_messages=True)
+        ctx = SessionContext(session_id="s1", input_messages=[Message(role="user", contents=["hello"])])
+        ctx.extend_messages("rag", [Message(role="system", contents=["context"])])
+        ctx._response = AgentResponse(messages=[Message(role="assistant", contents=["hi"])])
+        await provider.after_run(agent=None, session=AgentSession(), context=ctx, state={})  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+        # Should store: context from rag + input + response
+        texts = [m.text for m in provider.stored]
+        assert "context" in texts
+        assert "hello" in texts
+        assert "hi" in texts
+
+    async def test_after_run_stores_context_from_specific_sources(self) -> None:
+        from agent_framework import AgentResponse
+
+        provider = ConcreteHistoryProvider(
+            "audit", load_messages=False, store_context_messages=True, store_context_from={"rag"}
+        )
+        ctx = SessionContext(session_id="s1", input_messages=[])
+        ctx.extend_messages("rag", [Message(role="system", contents=["rag-context"])])
+        ctx.extend_messages("other", [Message(role="system", contents=["other-context"])])
+        ctx._response = AgentResponse(messages=[])
+        await provider.after_run(agent=None, session=AgentSession(), context=ctx, state={})  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+        texts = [m.text for m in provider.stored]
+        assert "rag-context" in texts
+        assert "other-context" not in texts
+
+
+# ---------------------------------------------------------------------------
+# AgentSession tests
+# ---------------------------------------------------------------------------
+
+
+class TestAgentSession:
+    def test_auto_generates_session_id(self) -> None:
+        session = AgentSession()
+        assert session.session_id is not None
+        assert len(session.session_id) > 0
+
+    def test_custom_session_id(self) -> None:
+        session = AgentSession(session_id="custom-123")
+        assert session.session_id == "custom-123"
+
+    def test_state_starts_empty(self) -> None:
+        session = AgentSession()
+        assert session.state == {}
+
+    def test_service_session_id(self) -> None:
+        session = AgentSession(service_session_id="svc-456")
+        assert session.service_session_id == "svc-456"
+
+    def test_service_session_id_accepts_structured_mapping(self) -> None:
+        service_session_id = {"context_id": "ctx-123", "task_id": "task-456", "task_state": "working"}
+        session = AgentSession(service_session_id=service_session_id)
+        assert session.service_session_id == service_session_id
+
+    def test_to_dict(self) -> None:
+        session = AgentSession(session_id="s1", service_session_id="svc1")
+        session.state = {"key": "value"}
+        d = session.to_dict()
+        assert d["type"] == "session"
+        assert d["session_id"] == "s1"
+        assert d["service_session_id"] == "svc1"
+        assert d["state"] == {"key": "value"}
+
+    def test_from_dict(self) -> None:
+        data = {
+            "type": "session",
+            "session_id": "s1",
+            "service_session_id": "svc1",
+            "state": {"key": "value"},
+        }
+        session = AgentSession.from_dict(data)
+        assert session.session_id == "s1"
+        assert session.service_session_id == "svc1"
+        assert session.state == {"key": "value"}
+
+    def test_roundtrip(self) -> None:
+        session = AgentSession(session_id="rt-1")
+        session.state = {"messages": ["a", "b"], "count": 42}
+        json_str = json.dumps(session.to_dict())
+        restored = AgentSession.from_dict(json.loads(json_str))
+        assert restored.session_id == "rt-1"
+        assert restored.state == {"messages": ["a", "b"], "count": 42}
+
+    def test_roundtrip_with_structured_service_session_id(self) -> None:
+        service_session_id = {"context_id": "ctx-123", "task_id": "task-456", "task_state": "working"}
+        session = AgentSession(session_id="rt-2", service_session_id=service_session_id)
+        json_str = json.dumps(session.to_dict())
+        restored = AgentSession.from_dict(json.loads(json_str))
+        assert restored.session_id == "rt-2"
+        assert restored.service_session_id == service_session_id
+
+    def test_from_dict_missing_state(self) -> None:
+        data = {"session_id": "s1"}
+        session = AgentSession.from_dict(data)
+        assert session.state == {}
+
+
+# ---------------------------------------------------------------------------
+# InMemoryHistoryProvider tests
+# ---------------------------------------------------------------------------
+
+
+class TestInMemoryHistoryProvider:
+    async def test_empty_state_returns_no_messages(self) -> None:
+        provider = InMemoryHistoryProvider()
+        session = AgentSession()
+        ctx = SessionContext(session_id="s1", input_messages=[])
+        await provider.before_run(  # type: ignore[arg-type]
+            agent=None,  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+            session=session,
+            context=ctx,
+            state=session.state.setdefault(provider.source_id, {}),
+        )
+        assert ctx.context_messages.get(provider.source_id, []) == []
+
+    async def test_stores_and_loads_messages(self) -> None:
+        from agent_framework import AgentResponse
+
+        provider = InMemoryHistoryProvider()
+        session = AgentSession()
+
+        # First run: send input, get response
+        input_msg = Message(role="user", contents=["hello"])
+        resp_msg = Message(role="assistant", contents=["hi there"])
+        ctx1 = SessionContext(session_id="s1", input_messages=[input_msg])
+        await provider.before_run(  # type: ignore[arg-type]
+            agent=None,  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+            session=session,
+            context=ctx1,
+            state=session.state.setdefault(provider.source_id, {}),
+        )
+        ctx1._response = AgentResponse(messages=[resp_msg])
+        await provider.after_run(  # type: ignore[arg-type]
+            agent=None,  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+            session=session,
+            context=ctx1,
+            state=session.state.setdefault(provider.source_id, {}),
+        )
+
+        # Second run: should load previous messages
+        ctx2 = SessionContext(session_id="s1", input_messages=[Message(role="user", contents=["again"])])
+        await provider.before_run(  # type: ignore[arg-type]
+            agent=None,  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+            session=session,
+            context=ctx2,
+            state=session.state.setdefault(provider.source_id, {}),
+        )
+        loaded = ctx2.context_messages.get(provider.source_id, [])
+        assert len(loaded) == 2
+        assert loaded[0].text == "hello"
+        assert loaded[1].text == "hi there"
+
+    async def test_state_is_serializable(self) -> None:
+        from agent_framework import AgentResponse
+
+        provider = InMemoryHistoryProvider()
+        session = AgentSession()
+
+        input_msg = Message(role="user", contents=["test"])
+        ctx = SessionContext(session_id="s1", input_messages=[input_msg])
+        await provider.before_run(  # type: ignore[arg-type]
+            agent=None,  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+            session=session,
+            context=ctx,
+            state=session.state.setdefault(provider.source_id, {}),
+        )
+        ctx._response = AgentResponse(messages=[Message(role="assistant", contents=["reply"])])
+        await provider.after_run(  # type: ignore[arg-type]
+            agent=None,  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+            session=session,
+            context=ctx,
+            state=session.state.setdefault(provider.source_id, {}),
+        )
+
+        # State contains Message objects (not dicts)
+        assert isinstance(session.state[provider.source_id]["messages"][0], Message)
+
+        # to_dict() serializes them via SerializationProtocol
+        session_dict = session.to_dict()
+        json_str = json.dumps(session_dict)
+        assert json_str  # no error
+
+        # Round-trip through session serialization restores Message objects
+        restored = AgentSession.from_dict(json.loads(json_str))
+        assert isinstance(restored.state[provider.source_id]["messages"][0], Message)
+        assert restored.state[provider.source_id]["messages"][0].text == "test"
+        assert restored.state[provider.source_id]["messages"][1].text == "reply"
+
+    async def test_source_id_attribution(self) -> None:
+        provider = InMemoryHistoryProvider("custom-source")
+        assert provider.source_id == "custom-source"
+        ctx = SessionContext(session_id="s1", input_messages=[])
+        ctx.extend_messages("custom-source", [Message(role="user", contents=["test"])])
+        assert "custom-source" in ctx.context_messages
+
+
+class TestFileHistoryProvider:
+    def test_is_marked_experimental(self) -> None:
+        assert FileHistoryProvider.__feature_stage__ == "experimental"  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+        assert FileHistoryProvider.__feature_id__ == ExperimentalFeature.FILE_HISTORY.value  # type: ignore[attr-defined]  # ty: ignore[unresolved-attribute]
+        assert FileHistoryProvider.__doc__ is not None
+        assert ".. warning:: Experimental" in FileHistoryProvider.__doc__
+
+    async def test_stores_and_loads_messages(self, tmp_path: Path) -> None:
+        from agent_framework import AgentResponse
+
+        provider = FileHistoryProvider(tmp_path)
+        session = AgentSession(session_id="s1")
+
+        input_message = Message(role="user", contents=["hello"])
+        response_message = Message(role="assistant", contents=["hi there"])
+        first_context = SessionContext(session_id=session.session_id, input_messages=[input_message])
+
+        await provider.before_run(  # type: ignore[arg-type]
+            agent=None,  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+            session=session,
+            context=first_context,
+            state={},
+        )
+        first_context._response = AgentResponse(messages=[response_message])
+        await provider.after_run(  # type: ignore[arg-type]
+            agent=None,  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+            session=session,
+            context=first_context,
+            state={},
+        )
+
+        session_file = provider._session_file_path(session.session_id)
+        assert session_file.name == "s1.jsonl"
+        assert session_file.exists()
+        raw_lines = (await asyncio.to_thread(session_file.read_text, encoding="utf-8")).splitlines()
+        assert len(raw_lines) == 2
+        payloads = [json.loads(line) for line in raw_lines]
+        assert all(payload["type"] == "message" for payload in payloads)
+        assert all("session_id" not in payload for payload in payloads)
+
+        second_context = SessionContext(
+            session_id=session.session_id, input_messages=[Message(role="user", contents=["again"])]
+        )
+        await provider.before_run(  # type: ignore[arg-type]
+            agent=None,  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+            session=session,
+            context=second_context,
+            state={},
+        )
+        loaded = second_context.context_messages.get(provider.source_id, [])
+        assert len(loaded) == 2
+        assert loaded[0].text == "hello"
+        assert loaded[1].text == "hi there"
+
+    def test_creates_storage_directory(self, tmp_path: Path) -> None:
+        nested_path = tmp_path / "nested" / "history"
+        provider = FileHistoryProvider(nested_path)
+        assert provider.storage_path == nested_path
+        assert nested_path.exists()
+        assert nested_path.is_dir()
+
+    async def test_uses_encoded_filename_for_unsafe_session_id(self, tmp_path: Path) -> None:
+        provider = FileHistoryProvider(tmp_path)
+        unsafe_session_id = "../unsafe/session"
+
+        await provider.save_messages(unsafe_session_id, [Message(role="user", contents=["hello"])])
+
+        session_file = provider._session_file_path(unsafe_session_id)
+        assert session_file.parent == provider.storage_path
+        assert session_file.name.startswith("~session-")
+        assert session_file.suffix == ".jsonl"
+        assert session_file.exists()
+        jsonl_files = await asyncio.to_thread(
+            lambda: sorted(path.name for path in provider.storage_path.glob("*.jsonl"))
+        )
+        assert jsonl_files == [session_file.name]
+
+    async def test_allows_custom_serializers_returning_bytes(self, tmp_path: Path) -> None:
+        calls: list[str] = []
+
+        def dumps(payload: object) -> bytes:
+            calls.append("dumps")
+            return json.dumps(payload).encode("utf-8")
+
+        def loads(payload: str | bytes) -> object:
+            calls.append("loads")
+            if isinstance(payload, bytes):
+                payload = payload.decode("utf-8")
+            return json.loads(payload)
+
+        provider = FileHistoryProvider(tmp_path, dumps=dumps, loads=loads)
+
+        await provider.save_messages("custom-serializer", [Message(role="user", contents=["hello"])])
+        loaded = await provider.get_messages("custom-serializer")
+
+        assert calls == ["dumps", "loads"]
+        assert len(loaded) == 1
+        assert loaded[0].text == "hello"
+
+    async def test_invalid_jsonl_line_raises(self, tmp_path: Path) -> None:
+        provider = FileHistoryProvider(tmp_path)
+        await asyncio.to_thread(provider._session_file_path("broken").write_text, "{not-json}\n", encoding="utf-8")
+
+        with pytest.raises(ValueError, match="Failed to deserialize history line 1"):
+            await provider.get_messages("broken")
+
+    async def test_missing_session_file_returns_empty_messages(self, tmp_path: Path) -> None:
+        provider = FileHistoryProvider(tmp_path)
+
+        loaded = await provider.get_messages("missing")
+
+        assert loaded == []
+
+    async def test_none_session_id_uses_default_jsonl_file(self, tmp_path: Path) -> None:
+        provider = FileHistoryProvider(tmp_path)
+
+        await provider.save_messages(None, [Message(role="user", contents=["hello"])])
+
+        session_file = provider._session_file_path(None)
+        assert session_file.name == "default.jsonl"
+        loaded = await provider.get_messages(None)
+        assert [message.text for message in loaded] == ["hello"]
+
+    async def test_non_mapping_jsonl_line_raises(self, tmp_path: Path) -> None:
+        provider = FileHistoryProvider(tmp_path)
+        await asyncio.to_thread(provider._session_file_path("non-mapping").write_text, "[1, 2, 3]\n", encoding="utf-8")
+
+        with pytest.raises(ValueError, match="did not deserialize to a mapping"):
+            await provider.get_messages("non-mapping")
+
+    async def test_skip_excluded_omits_excluded_messages(self, tmp_path: Path) -> None:
+        provider = FileHistoryProvider(tmp_path, skip_excluded=True)
+
+        await provider.save_messages(
+            "skip-excluded",
+            [
+                Message(role="user", contents=["keep"]),
+                Message(role="assistant", contents=["skip"], additional_properties={"_excluded": True}),
+            ],
+        )
+
+        loaded = await provider.get_messages("skip-excluded")
+
+        assert [message.text for message in loaded] == ["keep"]
+
+    async def test_serializer_must_return_single_line_json(self, tmp_path: Path) -> None:
+        def dumps(payload: object) -> str:
+            return json.dumps(payload, indent=2)
+
+        provider = FileHistoryProvider(tmp_path, dumps=dumps)
+
+        with pytest.raises(ValueError, match="single-line JSON"):
+            await provider.save_messages("pretty-json", [Message(role="user", contents=["hello"])])
+
+    async def test_concurrent_writes_for_same_session_are_locked(
+        self,
+        tmp_path: Path,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        provider = FileHistoryProvider(tmp_path)
+        session_id = "shared-session"
+        file_path = provider._session_file_path(session_id)
+        real_open = Path.open
+        write_started = threading.Event()
+        active_writes = 0
+        overlap_detected = False
+
+        class _TrackingFile:
+            def __init__(self, wrapped: Any) -> None:
+                self._wrapped = wrapped
+
+            def __enter__(self) -> "_TrackingFile":  # type: ignore[name-defined]
+                self._wrapped.__enter__()
+                return self
+
+            def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+                self._wrapped.__exit__(exc_type, exc_val, exc_tb)
+
+            def write(self, data: str) -> int:
+                nonlocal active_writes, overlap_detected
+                write_started.set()
+                active_writes += 1
+                overlap_detected = overlap_detected or active_writes > 1
+                try:
+                    time.sleep(0.05)
+                    return int(self._wrapped.write(data))
+                finally:
+                    active_writes -= 1
+
+            def __getattr__(self, name: str) -> Any:
+                return getattr(self._wrapped, name)
+
+        def tracked_open(path: Path, *args: Any, **kwargs: Any) -> Any:
+            handle = real_open(path, *args, **kwargs)
+            if path == file_path and args and args[0] == "a":
+                return _TrackingFile(handle)
+            return handle
+
+        monkeypatch.setattr(Path, "open", tracked_open)
+
+        first_save = asyncio.create_task(provider.save_messages(session_id, [Message(role="user", contents=["first"])]))
+        started = await asyncio.to_thread(write_started.wait, 1.0)
+        assert started
+
+        second_save = asyncio.create_task(
+            provider.save_messages(session_id, [Message(role="assistant", contents=["second"])])
+        )
+        await asyncio.gather(first_save, second_save)
+
+        assert not overlap_detected
+        loaded = await provider.get_messages(session_id)
+        assert [message.text for message in loaded] == ["first", "second"]

python/packages/core/tests/core/test_settings.py

@@ -0,0 +1,348 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for load_settings() function."""
+
+import os
+import tempfile
+from typing import TypedDict
+
+import pytest
+
+from agent_framework import SecretString, load_settings
+
+
+class SimpleSettings(TypedDict, total=False):
+    api_key: str | None
+    timeout: int | None
+    enabled: bool | None
+    rate_limit: float | None
+
+
+class RequiredFieldSettings(TypedDict, total=False):
+    name: str | None
+    optional_field: str | None
+
+
+class SecretSettings(TypedDict, total=False):
+    api_key: SecretString | None
+    username: str | None
+
+
+class ExclusiveSettings(TypedDict, total=False):
+    source_a: str | None
+    source_b: str | None
+    other: str | None
+
+
+class TestLoadSettingsBasic:
+    """Test basic load_settings functionality."""
+
+    def test_fields_are_none_when_unset(self) -> None:
+        settings = load_settings(SimpleSettings, env_prefix="TEST_APP_")
+
+        assert settings["api_key"] is None
+        assert settings["timeout"] is None
+        assert settings["enabled"] is None
+        assert settings["rate_limit"] is None
+
+    def test_overrides(self) -> None:
+        settings = load_settings(SimpleSettings, env_prefix="TEST_APP_", timeout=60, enabled=False)
+
+        assert settings["timeout"] == 60
+        assert settings["enabled"] is False
+
+    def test_none_overrides_are_filtered(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setenv("TEST_APP_TIMEOUT", "120")
+
+        settings = load_settings(SimpleSettings, env_prefix="TEST_APP_", timeout=None)
+
+        # timeout=None is filtered, so env var wins
+        assert settings["timeout"] == 120
+
+    def test_env_vars(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setenv("TEST_APP_API_KEY", "test-key-123")
+        monkeypatch.setenv("TEST_APP_TIMEOUT", "120")
+        monkeypatch.setenv("TEST_APP_ENABLED", "false")
+
+        settings = load_settings(SimpleSettings, env_prefix="TEST_APP_")
+
+        assert settings["api_key"] == "test-key-123"
+        assert settings["timeout"] == 120
+        assert settings["enabled"] is False
+
+    def test_overrides_beat_env_vars(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setenv("TEST_APP_TIMEOUT", "120")
+
+        settings = load_settings(SimpleSettings, env_prefix="TEST_APP_", timeout=60)
+
+        assert settings["timeout"] == 60
+
+    def test_no_prefix(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setenv("API_KEY", "no-prefix-key")
+
+        settings = load_settings(SimpleSettings, api_key=None)
+
+        assert settings["api_key"] == "no-prefix-key"
+
+
+class TestDotenvFile:
+    """Test .env file loading."""
+
+    def test_load_from_dotenv(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.delenv("TEST_APP_API_KEY", raising=False)
+        monkeypatch.delenv("TEST_APP_TIMEOUT", raising=False)
+
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".env", delete=False) as f:
+            f.write("TEST_APP_API_KEY=dotenv-key\n")
+            f.write("TEST_APP_TIMEOUT=90\n")
+            f.flush()
+            env_path = f.name
+
+        try:
+            settings = load_settings(SimpleSettings, env_prefix="TEST_APP_", env_file_path=env_path)
+
+            assert settings["api_key"] == "dotenv-key"
+            assert settings["timeout"] == 90
+        finally:
+            os.unlink(env_path)
+
+    def test_dotenv_overrides_env_vars_when_env_file_path_is_set(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setenv("TEST_APP_API_KEY", "real-env-key")
+
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".env", delete=False) as f:
+            f.write("TEST_APP_API_KEY=dotenv-key\n")
+            f.flush()
+            env_path = f.name
+
+        try:
+            settings = load_settings(SimpleSettings, env_prefix="TEST_APP_", env_file_path=env_path)
+
+            assert settings["api_key"] == "dotenv-key"
+        finally:
+            os.unlink(env_path)
+
+    def test_env_vars_are_used_when_env_file_path_is_not_set(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setenv("TEST_APP_API_KEY", "real-env-key")
+        settings = load_settings(SimpleSettings, env_prefix="TEST_APP_")
+
+        assert settings["api_key"] == "real-env-key"
+
+    def test_overrides_beat_dotenv_and_env_vars(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setenv("TEST_APP_TIMEOUT", "120")
+
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".env", delete=False) as f:
+            f.write("TEST_APP_TIMEOUT=90\n")
+            f.flush()
+            env_path = f.name
+
+        try:
+            settings = load_settings(SimpleSettings, env_prefix="TEST_APP_", env_file_path=env_path, timeout=60)
+
+            assert settings["timeout"] == 60
+        finally:
+            os.unlink(env_path)
+
+    def test_missing_dotenv_file_raises(self) -> None:
+        with pytest.raises(FileNotFoundError):
+            load_settings(SimpleSettings, env_prefix="TEST_APP_", env_file_path="/nonexistent/.env")
+
+
+class TestSecretString:
+    """Test SecretString type handling."""
+
+    def test_secretstring_from_env(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setenv("SECRET_API_KEY", "secret-value")
+
+        settings = load_settings(SecretSettings, env_prefix="SECRET_")
+
+        assert isinstance(settings["api_key"], SecretString)
+        assert settings["api_key"] == "secret-value"
+
+    def test_secretstring_from_override(self) -> None:
+        settings = load_settings(SecretSettings, env_prefix="SECRET_", api_key="kwarg-secret")
+
+        assert isinstance(settings["api_key"], SecretString)
+        assert settings["api_key"] == "kwarg-secret"
+
+    def test_secretstring_masked_in_repr(self) -> None:
+        s = SecretString("my-secret")
+        assert "my-secret" not in repr(s)
+        assert "**********" in repr(s)
+
+    def test_get_secret_value_compat(self) -> None:
+        s = SecretString("my-secret")
+
+        assert s.get_secret_value() == "my-secret"
+        assert isinstance(s.get_secret_value(), str)
+
+
+class TestTypeCoercion:
+    """Test type coercion from string values."""
+
+    def test_int_coercion(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setenv("TEST_APP_TIMEOUT", "42")
+
+        settings = load_settings(SimpleSettings, env_prefix="TEST_APP_")
+
+        assert settings["timeout"] == 42
+        assert isinstance(settings["timeout"], int)
+
+    def test_float_coercion(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setenv("TEST_APP_RATE_LIMIT", "2.5")
+
+        settings = load_settings(SimpleSettings, env_prefix="TEST_APP_")
+
+        assert settings["rate_limit"] == 2.5
+        assert isinstance(settings["rate_limit"], float)
+
+    def test_bool_coercion_true_values(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        for true_val in ["true", "True", "TRUE", "1", "yes", "on"]:
+            monkeypatch.setenv("TEST_APP_ENABLED", true_val)
+            settings = load_settings(SimpleSettings, env_prefix="TEST_APP_")
+            assert settings["enabled"] is True, f"Failed for {true_val}"
+
+    def test_bool_coercion_false_values(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        for false_val in ["false", "False", "FALSE", "0", "no", "off"]:
+            monkeypatch.setenv("TEST_APP_ENABLED", false_val)
+            settings = load_settings(SimpleSettings, env_prefix="TEST_APP_")
+            assert settings["enabled"] is False, f"Failed for {false_val}"
+
+
+class TestRequiredFields:
+    """Test required field validation."""
+
+    def test_required_field_provided(self) -> None:
+        settings = load_settings(
+            RequiredFieldSettings,
+            env_prefix="TEST_",
+            required_fields=["name"],
+            name="my-app",
+        )
+
+        assert settings["name"] == "my-app"
+        assert settings["optional_field"] is None
+
+    def test_required_field_from_env(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setenv("TEST_NAME", "env-app")
+
+        settings = load_settings(RequiredFieldSettings, env_prefix="TEST_", required_fields=["name"])
+
+        assert settings["name"] == "env-app"
+
+    def test_required_field_missing_raises(self) -> None:
+        from agent_framework.exceptions import SettingNotFoundError
+
+        with pytest.raises(SettingNotFoundError, match="Required setting 'name'"):
+            load_settings(RequiredFieldSettings, env_prefix="TEST_", required_fields=["name"])
+
+    def test_without_required_fields_param_allows_none(self) -> None:
+        settings = load_settings(RequiredFieldSettings, env_prefix="TEST_")
+
+        assert settings["name"] is None
+
+
+class TestOverrideTypeValidation:
+    """Test override type validation."""
+
+    def test_invalid_type_raises(self) -> None:
+
+        with pytest.raises(ValueError, match="Invalid type for setting 'api_key'"):
+            load_settings(SimpleSettings, env_prefix="TEST_", api_key={"bad": "type"})
+
+    def test_valid_types_accepted(self) -> None:
+        settings = load_settings(SimpleSettings, env_prefix="TEST_", timeout=42, enabled=True)
+
+        assert settings["timeout"] == 42
+        assert settings["enabled"] is True
+
+    def test_str_accepted_for_secretstring(self) -> None:
+        settings = load_settings(SecretSettings, env_prefix="TEST_", api_key="plain-string")
+
+        assert isinstance(settings["api_key"], SecretString)
+        assert settings["api_key"] == "plain-string"
+
+
+class TestMutuallyExclusive:
+    """Test mutually exclusive field validation via tuple entries in required_fields."""
+
+    def test_exactly_one_set_passes(self) -> None:
+        settings = load_settings(
+            ExclusiveSettings,
+            env_prefix="TEST_",
+            required_fields=[("source_a", "source_b")],
+            source_a="value-a",
+        )
+
+        assert settings["source_a"] == "value-a"
+        assert settings["source_b"] is None
+
+    def test_none_set_raises(self) -> None:
+        from agent_framework.exceptions import SettingNotFoundError
+
+        with pytest.raises(SettingNotFoundError, match="none was set"):
+            load_settings(
+                ExclusiveSettings,
+                env_prefix="TEST_",
+                required_fields=[("source_a", "source_b")],
+            )
+
+    def test_both_set_raises(self) -> None:
+        from agent_framework.exceptions import SettingNotFoundError
+
+        with pytest.raises(SettingNotFoundError, match="multiple were set"):
+            load_settings(
+                ExclusiveSettings,
+                env_prefix="TEST_",
+                required_fields=[("source_a", "source_b")],
+                source_a="a",
+                source_b="b",
+            )
+
+    def test_env_var_counts_as_set(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setenv("TEST_SOURCE_B", "env-b")
+
+        settings = load_settings(
+            ExclusiveSettings,
+            env_prefix="TEST_",
+            required_fields=[("source_a", "source_b")],
+        )
+
+        assert settings["source_b"] == "env-b"
+
+    def test_env_var_and_override_both_set_raises(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        from agent_framework.exceptions import SettingNotFoundError
+
+        monkeypatch.setenv("TEST_SOURCE_B", "env-b")
+
+        with pytest.raises(SettingNotFoundError, match="multiple were set"):
+            load_settings(
+                ExclusiveSettings,
+                env_prefix="TEST_",
+                required_fields=[("source_a", "source_b")],
+                source_a="a",
+            )
+
+    def test_other_fields_unaffected(self) -> None:
+        settings = load_settings(
+            ExclusiveSettings,
+            env_prefix="TEST_",
+            required_fields=[("source_a", "source_b")],
+            source_a="a",
+            other="extra",
+        )
+
+        assert settings["source_a"] == "a"
+        assert settings["other"] == "extra"
+
+    def test_mixed_required_and_exclusive(self) -> None:
+        settings = load_settings(
+            ExclusiveSettings,
+            env_prefix="TEST_",
+            required_fields=["other", ("source_a", "source_b")],
+            source_b="b",
+            other="required-val",
+        )
+
+        assert settings["other"] == "required-val"
+        assert settings["source_b"] == "b"
+        assert settings["source_a"] is None

python/packages/core/tests/core/test_telemetry.py

@@ -0,0 +1,289 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import os
+from unittest.mock import MagicMock, patch
+
+import agent_framework._telemetry as _telemetry_mod
+from agent_framework import (
+    AGENT_FRAMEWORK_USER_AGENT,
+    USER_AGENT_KEY,
+    USER_AGENT_TELEMETRY_DISABLED_ENV_VAR,
+    prepend_agent_framework_to_user_agent,
+)
+from agent_framework._telemetry import (
+    _FOUNDRY_HOSTING_ENV_VAR,
+    _HOSTED_USER_AGENT_PREFIX,
+    _add_user_agent_prefix,
+    _detect_hosted_environment,
+)
+
+# region Test constants
+
+
+def test_telemetry_disabled_env_var():
+    """Test that the telemetry disabled environment variable is correctly defined."""
+    assert USER_AGENT_TELEMETRY_DISABLED_ENV_VAR == "AGENT_FRAMEWORK_USER_AGENT_DISABLED"
+
+
+def test_user_agent_key():
+    """Test that the user agent key is correctly defined."""
+    assert USER_AGENT_KEY == "User-Agent"
+
+
+def test_agent_framework_user_agent_format():
+    """Test that the agent framework user agent is correctly formatted."""
+    assert AGENT_FRAMEWORK_USER_AGENT.startswith("agent-framework-python/")
+
+
+def test_app_info_when_telemetry_enabled():
+    """Test that APP_INFO is set when telemetry is enabled."""
+    with patch("agent_framework._telemetry.IS_TELEMETRY_ENABLED", True):
+        import importlib
+
+        import agent_framework._telemetry
+
+        importlib.reload(agent_framework._telemetry)
+        from agent_framework import APP_INFO
+
+        assert APP_INFO is not None
+        assert "agent-framework-version" in APP_INFO
+        assert APP_INFO["agent-framework-version"].startswith("python/")
+
+
+def test_app_info_when_telemetry_disabled():
+    """Test that APP_INFO is None when telemetry is disabled."""
+    # Test the logic directly since APP_INFO is set at module import time
+    with patch("agent_framework._telemetry.IS_TELEMETRY_ENABLED", False):
+        # Simulate the module's logic for APP_INFO
+        test_app_info = (
+            {
+                "agent-framework-version": "python/test",
+            }
+            if False  # This simulates IS_TELEMETRY_ENABLED being False
+            else None
+        )
+        assert test_app_info is None
+
+
+# region Test prepend_agent_framework_to_user_agent
+
+
+def test_prepend_to_existing_user_agent():
+    """Test prepending to existing User-Agent header."""
+    headers = {"User-Agent": "existing-agent/1.0"}
+    result = prepend_agent_framework_to_user_agent(headers)
+
+    assert "User-Agent" in result
+    assert result["User-Agent"].startswith("agent-framework-python/")
+    assert "existing-agent/1.0" in result["User-Agent"]
+
+
+def test_prepend_to_empty_headers():
+    """Test prepending to headers without User-Agent."""
+    headers = {"Content-Type": "application/json"}
+    result = prepend_agent_framework_to_user_agent(headers)
+
+    assert "User-Agent" in result
+    assert result["User-Agent"] == AGENT_FRAMEWORK_USER_AGENT
+    assert "Content-Type" in result
+
+
+def test_prepend_to_empty_dict():
+    """Test prepending to empty headers dict."""
+    headers: dict[str, str] = {}
+    result = prepend_agent_framework_to_user_agent(headers)
+
+    assert "User-Agent" in result
+    assert result["User-Agent"] == AGENT_FRAMEWORK_USER_AGENT
+
+
+def test_modifies_original_dict():
+    """Test that the function modifies the original headers dict."""
+    headers = {"Other-Header": "value"}
+    result = prepend_agent_framework_to_user_agent(headers)
+
+    assert result is headers  # Same object
+    assert "User-Agent" in headers
+
+
+# region Test _add_user_agent_prefix
+
+
+def test_add_user_agent_prefix_adds_prefix():
+    """Test that _add_user_agent_prefix permanently adds a prefix."""
+    _telemetry_mod._user_agent_prefixes.clear()
+    _add_user_agent_prefix("test-host")
+    result = prepend_agent_framework_to_user_agent()
+    assert result["User-Agent"].startswith("test-host/")
+    assert AGENT_FRAMEWORK_USER_AGENT in result["User-Agent"]
+    _telemetry_mod._user_agent_prefixes.clear()
+
+
+def test_add_user_agent_prefix_ignores_duplicates():
+    """Test that duplicate prefixes are not added."""
+    _telemetry_mod._user_agent_prefixes.clear()
+    _add_user_agent_prefix("test-host")
+    _add_user_agent_prefix("test-host")
+    result = prepend_agent_framework_to_user_agent()
+    assert result["User-Agent"].count("test-host") == 1
+    _telemetry_mod._user_agent_prefixes.clear()
+
+
+def test_add_user_agent_prefix_ignores_empty():
+    """Test that empty strings are not added as prefixes."""
+    _telemetry_mod._user_agent_prefixes.clear()
+    _add_user_agent_prefix("")
+    result = prepend_agent_framework_to_user_agent()
+    assert result["User-Agent"] == AGENT_FRAMEWORK_USER_AGENT
+    _telemetry_mod._user_agent_prefixes.clear()
+
+
+def test_add_user_agent_prefix_multiple():
+    """Test that multiple prefixes compose correctly."""
+    _telemetry_mod._user_agent_prefixes.clear()
+    _add_user_agent_prefix("outer")
+    _add_user_agent_prefix("inner")
+    result = prepend_agent_framework_to_user_agent()
+    assert "outer" in result["User-Agent"]
+    assert "inner" in result["User-Agent"]
+    _telemetry_mod._user_agent_prefixes.clear()
+
+
+# region Test _detect_hosted_environment
+
+
+def test_detect_hosted_env_var_truthy_adds_prefix():
+    """Test that a truthy FOUNDRY_HOSTING_ENVIRONMENT env var adds the prefix."""
+    _telemetry_mod._user_agent_prefixes.clear()
+    _telemetry_mod._hosted_env_detected = False
+    with patch.dict("os.environ", {_FOUNDRY_HOSTING_ENV_VAR: "production"}):
+        _detect_hosted_environment()
+    assert _HOSTED_USER_AGENT_PREFIX in _telemetry_mod._user_agent_prefixes
+    _telemetry_mod._user_agent_prefixes.clear()
+    _telemetry_mod._hosted_env_detected = False
+
+
+def test_detect_hosted_env_var_empty_skips_prefix():
+    """Test that an empty FOUNDRY_HOSTING_ENVIRONMENT env var does NOT add the prefix."""
+    _telemetry_mod._user_agent_prefixes.clear()
+    _telemetry_mod._hosted_env_detected = False
+    with patch.dict("os.environ", {_FOUNDRY_HOSTING_ENV_VAR: ""}):
+        _detect_hosted_environment()
+    assert _HOSTED_USER_AGENT_PREFIX not in _telemetry_mod._user_agent_prefixes
+    _telemetry_mod._user_agent_prefixes.clear()
+    _telemetry_mod._hosted_env_detected = False
+
+
+def test_detect_hosted_env_var_set_skips_agent_config_fallback():
+    """Test that when the env var is set, AgentConfig is never consulted even if import would fail."""
+    _telemetry_mod._user_agent_prefixes.clear()
+    _telemetry_mod._hosted_env_detected = False
+    import builtins
+
+    real_import = builtins.__import__
+
+    def _block_agentconfig(name: str, *args, **kwargs):  # type: ignore[no-untyped-def]
+        if "agentserver" in name:
+            raise AssertionError("AgentConfig should not be imported when env var is set")
+        return real_import(name, *args, **kwargs)
+
+    with (
+        patch.dict("os.environ", {_FOUNDRY_HOSTING_ENV_VAR: "prod"}),
+        patch("builtins.__import__", side_effect=_block_agentconfig),
+    ):
+        _detect_hosted_environment()
+    assert _HOSTED_USER_AGENT_PREFIX in _telemetry_mod._user_agent_prefixes
+    _telemetry_mod._user_agent_prefixes.clear()
+    _telemetry_mod._hosted_env_detected = False
+
+
+def _mock_agent_config(*, is_hosted: bool) -> MagicMock:
+    """Create a mock azure.ai.agentserver.core module with AgentConfig."""
+    mock_config = MagicMock()
+    mock_config.is_hosted = is_hosted
+    mock_module = MagicMock()
+    mock_module.AgentConfig.from_env.return_value = mock_config
+    return mock_module
+
+
+def test_detect_hosted_fallback_agent_config_is_hosted():
+    """Test that AgentConfig fallback adds the prefix when is_hosted is True."""
+    _telemetry_mod._user_agent_prefixes.clear()
+    _telemetry_mod._hosted_env_detected = False
+    env = {k: v for k, v in os.environ.items() if k != _FOUNDRY_HOSTING_ENV_VAR}
+    mock_module = _mock_agent_config(is_hosted=True)
+    mock_spec = MagicMock()
+    with (
+        patch.dict("os.environ", env, clear=True),
+        patch.dict("sys.modules", {"azure.ai.agentserver.core": mock_module}),
+        patch("importlib.util.find_spec", return_value=mock_spec),
+    ):
+        _detect_hosted_environment()
+    assert _HOSTED_USER_AGENT_PREFIX in _telemetry_mod._user_agent_prefixes
+    _telemetry_mod._user_agent_prefixes.clear()
+    _telemetry_mod._hosted_env_detected = False
+
+
+def test_detect_hosted_fallback_agent_config_not_hosted():
+    """Test that AgentConfig fallback does NOT add the prefix when is_hosted is False."""
+    _telemetry_mod._user_agent_prefixes.clear()
+    _telemetry_mod._hosted_env_detected = False
+    mock_module = _mock_agent_config(is_hosted=False)
+    mock_spec = MagicMock()
+    env = {k: v for k, v in os.environ.items() if k != _FOUNDRY_HOSTING_ENV_VAR}
+    with (
+        patch.dict("os.environ", env, clear=True),
+        patch.dict("sys.modules", {"azure.ai.agentserver.core": mock_module}),
+        patch("importlib.util.find_spec", return_value=mock_spec),
+    ):
+        _detect_hosted_environment()
+    assert _HOSTED_USER_AGENT_PREFIX not in _telemetry_mod._user_agent_prefixes
+    _telemetry_mod._user_agent_prefixes.clear()
+    _telemetry_mod._hosted_env_detected = False
+
+
+def test_detect_hosted_fallback_import_error():
+    """Test that ImportError from AgentConfig is silently handled."""
+    _telemetry_mod._user_agent_prefixes.clear()
+    _telemetry_mod._hosted_env_detected = False
+    env = {k: v for k, v in os.environ.items() if k != _FOUNDRY_HOSTING_ENV_VAR}
+    with patch.dict("os.environ", env, clear=True):
+        # The real import may succeed or fail depending on the environment;
+        # force the ImportError path by making the import raise.
+        import builtins
+
+        real_import = builtins.__import__
+
+        def _block_agentconfig(name: str, *args, **kwargs):  # type: ignore[no-untyped-def]
+            if "agentserver" in name:
+                raise ImportError("mocked")
+            return real_import(name, *args, **kwargs)
+
+        with patch("builtins.__import__", side_effect=_block_agentconfig):
+            _detect_hosted_environment()
+    assert _HOSTED_USER_AGENT_PREFIX not in _telemetry_mod._user_agent_prefixes
+    _telemetry_mod._user_agent_prefixes.clear()
+    _telemetry_mod._hosted_env_detected = False
+
+
+# region Test module-level auto-detection
+
+
+def test_lazy_detection_on_get_user_agent():
+    """Test that get_user_agent() lazily detects the hosted environment.
+
+    Since detection is deferred to the first ``get_user_agent()`` call,
+    this verifies the prefix is included without any explicit call to
+    ``_detect_hosted_environment()`` by consumer code.
+    """
+    _telemetry_mod._user_agent_prefixes.clear()
+    _telemetry_mod._hosted_env_detected = False
+    with patch.dict("os.environ", {_FOUNDRY_HOSTING_ENV_VAR: "production"}):
+        user_agent = _telemetry_mod.get_user_agent()
+
+    assert _HOSTED_USER_AGENT_PREFIX in _telemetry_mod._user_agent_prefixes
+    assert user_agent.startswith(f"{_HOSTED_USER_AGENT_PREFIX}/")
+
+    # Clean up
+    _telemetry_mod._user_agent_prefixes.clear()
+    _telemetry_mod._hosted_env_detected = False

python/packages/core/tests/core/test_tools_future_annotations.py

@@ -0,0 +1,134 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for @tool with PEP 563 (from __future__ import annotations).
+
+When ``from __future__ import annotations`` is active, all annotations
+become strings.  _resolve_input_model must resolve them via
+typing.get_type_hints() before passing them to Pydantic's create_model.
+"""
+
+from __future__ import annotations
+
+from pydantic import BaseModel
+
+from agent_framework import tool
+from agent_framework._middleware import FunctionInvocationContext
+
+
+class SearchConfig(BaseModel):
+    max_results: int = 10
+
+
+def test_tool_with_context_parameter():
+    """FunctionInvocationContext parameter is excluded from schema under PEP 563."""
+
+    @tool
+    def get_weather(location: str, ctx: FunctionInvocationContext) -> str:
+        """Get the weather for a given location."""
+        return f"Weather in {location}"
+
+    params = get_weather.parameters()
+    assert "ctx" not in params.get("properties", {})
+    assert "location" in params["properties"]
+
+
+def test_tool_with_context_parameter_first():
+    """FunctionInvocationContext as the first parameter is excluded under PEP 563."""
+
+    @tool
+    def get_weather(ctx: FunctionInvocationContext, location: str) -> str:
+        """Get the weather for a given location."""
+        return f"Weather in {location}"
+
+    params = get_weather.parameters()
+    assert "ctx" not in params.get("properties", {})
+    assert "location" in params["properties"]
+
+
+def test_tool_with_optional_param():
+    """Optional[int] is resolved to the actual type, not left as a string."""
+
+    @tool
+    def search(query: str, limit: int | None = None) -> str:
+        """Search for something."""
+        return query
+
+    params = search.parameters()
+    assert params["properties"]["query"]["type"] == "string"
+    limit_schema = params["properties"]["limit"]
+    limit_types = {t["type"] for t in limit_schema["anyOf"]}
+    assert limit_types == {"integer", "null"}
+
+
+def test_tool_with_optional_param_and_context():
+    """Optional param + FunctionInvocationContext both work under PEP 563."""
+
+    @tool
+    def search(query: str, limit: int | None = None, ctx: FunctionInvocationContext | None = None) -> str:
+        """Search for something."""
+        return query
+
+    params = search.parameters()
+    assert params["properties"]["query"]["type"] == "string"
+    limit_schema = params["properties"]["limit"]
+    limit_types = {t["type"] for t in limit_schema["anyOf"]}
+    assert limit_types == {"integer", "null"}
+    assert "ctx" not in params.get("properties", {})
+
+
+def test_tool_with_optional_custom_type():
+    """Optional[CustomType] is resolved under PEP 563 (original bug pattern)."""
+
+    @tool
+    def search(query: str, config: SearchConfig | None = None) -> str:
+        """Search for something."""
+        return query
+
+    params = search.parameters()
+    assert params["properties"]["query"]["type"] == "string"
+    config_schema = params["properties"]["config"]
+    config_types = [t.get("type") for t in config_schema["anyOf"]]
+    assert "null" in config_types
+
+
+def test_tool_with_unresolvable_forward_ref():
+    """Fallback to raw annotations when get_type_hints() fails."""
+    import types
+
+    # Build a function in an isolated namespace so get_type_hints() cannot resolve
+    # the forward reference, exercising the except-branch fallback.
+    ns: dict = {}
+    exec(
+        "def greet(name: str = 'world') -> str:\n    '''Greet someone.'''\n    return f'Hello {name}'\n",
+        ns,
+    )
+    func = ns["greet"]
+    # Place the function in a throwaway module so get_type_hints() will fail on
+    # any non-builtin forward ref while still having a valid __module__.
+    mod = types.ModuleType("_phantom")
+    func.__module__ = mod.__name__
+
+    t = tool(func)
+    params = t.parameters()
+    assert params["properties"]["name"]["type"] == "string"
+
+
+async def test_tool_invoke_with_context():
+    """Full invocation with FunctionInvocationContext under PEP 563."""
+
+    @tool
+    def get_weather(location: str, ctx: FunctionInvocationContext) -> str:
+        """Get the weather for a given location."""
+        user = ctx.kwargs.get("user", "anon")
+        return f"Weather in {location} for {user}"
+
+    params = get_weather.parameters()
+    assert "ctx" not in params.get("properties", {})
+
+    context = FunctionInvocationContext(
+        function=get_weather,
+        arguments=get_weather.input_model(location="Seattle"),  # type: ignore[misc, operator]  # pyrefly: ignore[not-callable]  # ty: ignore[call-non-callable]
+        kwargs={"user": "test_user"},
+    )
+    result = await get_weather.invoke(context=context)
+    assert result[0].text == "Weather in Seattle for test_user"

python/packages/core/tests/core/utils.py

@@ -0,0 +1,12 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+
+from copy import deepcopy
+from unittest.mock import MagicMock
+
+
+class CopyingMock(MagicMock):
+    def __call__(self, *args, **kwargs):
+        args = deepcopy(args)
+        kwargs = deepcopy(kwargs)
+        return super().__call__(*args, **kwargs)

python/packages/core/tests/workflow/__init__.py

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

python/packages/core/tests/workflow/test_agent_executor.py

@@ -0,0 +1,869 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from collections.abc import AsyncIterable, Awaitable
+from typing import Any, Literal, overload
+
+import pytest
+
+from agent_framework import (
+    AgentExecutor,
+    AgentResponse,
+    AgentResponseUpdate,
+    AgentRunInputs,
+    AgentSession,
+    BaseAgent,
+    Content,
+    Message,
+    ResponseStream,
+    WorkflowBuilder,
+    WorkflowEvent,
+    WorkflowRunState,
+)
+from agent_framework._workflows._agent_executor import AgentExecutorResponse
+from agent_framework._workflows._checkpoint import InMemoryCheckpointStorage
+from agent_framework._workflows._const import GLOBAL_KWARGS_KEY
+
+
+class _CountingAgent(BaseAgent):
+    """Agent that echoes messages with a counter to verify session state persistence."""
+
+    def __init__(self, **kwargs: Any):
+        super().__init__(**kwargs)
+        self.call_count = 0
+
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
+    def run(
+        self,
+        messages: AgentRunInputs | None = None,
+        *,
+        stream: bool = False,
+        session: AgentSession | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
+        self.call_count += 1
+        if stream:
+
+            async def _stream() -> AsyncIterable[AgentResponseUpdate]:
+                yield AgentResponseUpdate(
+                    contents=[Content.from_text(text=f"Response #{self.call_count}: {self.name}")]
+                )
+
+            return ResponseStream(_stream(), finalizer=AgentResponse.from_updates)
+
+        async def _run() -> AgentResponse:
+            return AgentResponse(messages=[Message("assistant", [f"Response #{self.call_count}: {self.name}"])])
+
+        return _run()
+
+
+class _StreamingHookAgent(BaseAgent):
+    """Agent that exposes whether its streaming result hook was executed."""
+
+    def __init__(self, **kwargs: Any):
+        super().__init__(**kwargs)
+        self.result_hook_called = False
+
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
+    def run(
+        self,
+        messages: AgentRunInputs | None = None,
+        *,
+        stream: bool = False,
+        session: AgentSession | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
+        if stream:
+
+            async def _stream() -> AsyncIterable[AgentResponseUpdate]:
+                yield AgentResponseUpdate(
+                    contents=[Content.from_text(text="hook test")],
+                    role="assistant",
+                )
+
+            async def _mark_result_hook_called(
+                response: AgentResponse,
+            ) -> AgentResponse:
+                self.result_hook_called = True
+                return response
+
+            return ResponseStream(_stream(), finalizer=AgentResponse.from_updates).with_result_hook(
+                _mark_result_hook_called
+            )
+
+        async def _run() -> AgentResponse:
+            return AgentResponse(messages=[Message("assistant", ["hook test"])])
+
+        return _run()
+
+
+async def test_agent_executor_streaming_finalizes_stream_and_runs_result_hooks() -> None:
+    """AgentExecutor should call get_final_response() so stream result hooks execute."""
+    agent = _StreamingHookAgent(id="hook_agent", name="HookAgent")
+    executor = AgentExecutor(agent, id="hook_exec")
+    workflow = WorkflowBuilder(start_executor=executor).build()
+
+    output_events: list[Any] = []
+    async for event in workflow.run("run hook test", stream=True):
+        if event.type == "output":
+            output_events.append(event)
+
+    assert output_events
+    assert agent.result_hook_called
+
+
+async def test_agent_executor_checkpoint_stores_and_restores_state() -> None:
+    """Test that workflow checkpoint stores AgentExecutor's cache and session states and restores them correctly."""
+    storage = InMemoryCheckpointStorage()
+
+    # Create two agents to form a two-step workflow
+    initial_agent_a = _CountingAgent(id="agent_a", name="AgentA")
+    initial_agent_b = _CountingAgent(id="agent_b", name="AgentB")
+    initial_session = AgentSession()
+
+    # Add some initial messages to the session state to verify session state persistence
+    initial_messages = [
+        Message(role="user", contents=["Initial message 1"]),
+        Message(role="assistant", contents=["Initial response 1"]),
+    ]
+    initial_session.state["history"] = {"messages": initial_messages}
+
+    # Create AgentExecutors — first executor gets the custom session
+    exec_a = AgentExecutor(initial_agent_a, id="exec_a", session=initial_session)
+    exec_b = AgentExecutor(initial_agent_b, id="exec_b")
+
+    # Build two-executor workflow with checkpointing enabled
+    wf = WorkflowBuilder(start_executor=exec_a, checkpoint_storage=storage).add_edge(exec_a, exec_b).build()
+
+    # Run the workflow with a user message
+    first_run_output: AgentExecutorResponse | None = None
+    async for ev in wf.run("First workflow run", stream=True):
+        if ev.type == "output":
+            first_run_output = ev.data  # type: ignore[assignment]
+        if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
+            break
+
+    assert first_run_output is not None
+    assert initial_agent_a.call_count == 1
+
+    # Verify checkpoint was created
+    checkpoints = await storage.list_checkpoints(workflow_name=wf.name)
+    assert len(checkpoints) >= 2, "Expected at least 2 checkpoints: one after exec_a and one after exec_b."
+
+    # Get the first checkpoint that contains exec_a's state (taken after exec_a completes,
+    # before exec_b runs)
+    checkpoints.sort(key=lambda cp: cp.timestamp)
+    restore_checkpoint = next(
+        cp for cp in checkpoints if "_executor_state" in cp.state and "exec_a" in cp.state["_executor_state"]
+    )
+
+    # Verify checkpoint contains executor state with both cache and session
+    executor_states = restore_checkpoint.state["_executor_state"]
+    assert isinstance(executor_states, dict)
+    assert exec_a.id in executor_states
+
+    executor_state = executor_states[exec_a.id]  # type: ignore[index]
+    assert "cache" in executor_state, "Checkpoint should store executor cache state"
+    assert "agent_session" in executor_state, "Checkpoint should store executor session state"
+
+    # Verify session state structure
+    session_state = executor_state["agent_session"]  # type: ignore[index]
+    assert "session_id" in session_state, "Session state should include session_id"
+    assert "state" in session_state, "Session state should include state dict"
+
+    # Verify checkpoint contains pending requests from agents and responses to be sent
+    assert "pending_agent_requests" in executor_state
+    assert "pending_responses_to_agent" in executor_state
+
+    # Create new agents and executors for restoration
+    # This simulates starting from a fresh state and restoring from checkpoint
+    restored_agent_a = _CountingAgent(id="agent_a", name="AgentA")
+    restored_agent_b = _CountingAgent(id="agent_b", name="AgentB")
+    restored_session = AgentSession()
+    restored_exec_a = AgentExecutor(restored_agent_a, id="exec_a", session=restored_session)
+    restored_exec_b = AgentExecutor(restored_agent_b, id="exec_b")
+
+    # Verify the restored agents start with a fresh state
+    assert restored_agent_a.call_count == 0
+    assert restored_agent_b.call_count == 0
+
+    # Build new workflow with the restored executors
+    wf_resume = (
+        WorkflowBuilder(start_executor=restored_exec_a, checkpoint_storage=storage)
+        .add_edge(restored_exec_a, restored_exec_b)
+        .build()
+    )
+
+    # Resume from checkpoint — exec_a already ran, so exec_b should run and produce output
+    resumed_output: AgentExecutorResponse | None = None
+    async for ev in wf_resume.run(checkpoint_id=restore_checkpoint.checkpoint_id, stream=True):
+        if ev.type == "output":
+            resumed_output = ev.data  # type: ignore[assignment]
+        if ev.type == "status" and ev.state in (
+            WorkflowRunState.IDLE,
+            WorkflowRunState.IDLE_WITH_PENDING_REQUESTS,
+        ):
+            break
+
+    assert resumed_output is not None
+
+    # Verify the restored executor's session state was restored
+    restored_session_obj = restored_exec_a._session  # pyright: ignore[reportPrivateUsage]
+    assert restored_session_obj is not None
+    assert restored_session_obj.session_id == initial_session.session_id
+
+
+async def test_agent_executor_save_and_restore_state_directly() -> None:
+    """Test AgentExecutor's on_checkpoint_save and on_checkpoint_restore methods directly."""
+    # Create agent with session containing state
+    agent = _CountingAgent(id="direct_test_agent", name="DirectTestAgent")
+    session = AgentSession()
+
+    # Add messages to session state
+    session_messages = [
+        Message(role="user", contents=["Message in session 1"]),
+        Message(role="assistant", contents=["Session response 1"]),
+        Message(role="user", contents=["Message in session 2"]),
+    ]
+    session.state["history"] = {"messages": session_messages}
+
+    executor = AgentExecutor(agent, session=session)
+
+    # Add messages to executor cache
+    cache_messages = [
+        Message(role="user", contents=["Cached user message"]),
+        Message(role="assistant", contents=["Cached assistant response"]),
+    ]
+    executor._cache = list(cache_messages)  # pyright: ignore[reportPrivateUsage]
+
+    # Snapshot the state
+    state = await executor.on_checkpoint_save()
+
+    # Verify snapshot contains both cache and session
+    assert "cache" in state
+    assert "agent_session" in state
+
+    # Verify session state structure
+    session_state = state["agent_session"]  # type: ignore[index]
+    assert "session_id" in session_state
+    assert "state" in session_state
+
+    # Create new executor to restore into
+    new_agent = _CountingAgent(id="direct_test_agent", name="DirectTestAgent")
+    new_session = AgentSession()
+    new_executor = AgentExecutor(new_agent, session=new_session)
+
+    # Verify new executor starts empty
+    assert len(new_executor._cache) == 0  # pyright: ignore[reportPrivateUsage]
+    assert len(new_session.state) == 0
+
+    # Restore state
+    await new_executor.on_checkpoint_restore(state)
+
+    # Verify cache is restored
+    restored_cache = new_executor._cache  # pyright: ignore[reportPrivateUsage]
+    assert len(restored_cache) == len(cache_messages)
+    assert restored_cache[0].text == "Cached user message"
+    assert restored_cache[1].text == "Cached assistant response"
+
+    # Verify session was restored with correct session_id
+    restored_session = new_executor._session  # pyright: ignore[reportPrivateUsage]
+    assert restored_session.session_id == session.session_id
+
+
+async def test_prepare_agent_run_args_extracts_invocation_kwargs() -> None:
+    """_prepare_agent_run_args extracts function_invocation_kwargs and client_kwargs."""
+    agent = _CountingAgent(id="test_agent", name="TestAgent")
+    executor = AgentExecutor(agent, id="test_exec")
+
+    raw: dict[str, Any] = {
+        "function_invocation_kwargs": {"__global__": {"key": "fi_val"}},
+        "client_kwargs": {"__global__": {"key": "ci_val"}},
+    }
+    fi_kwargs, ci_kwargs = executor._prepare_agent_run_args(raw)  # pyright: ignore[reportPrivateUsage]
+    assert fi_kwargs == {"key": "fi_val"}
+    assert ci_kwargs == {"key": "ci_val"}
+
+
+async def test_prepare_agent_run_args_returns_none_when_no_kwargs() -> None:
+    """_prepare_agent_run_args returns None for both when raw dict has no invocation kwargs."""
+    agent = _CountingAgent(id="test_agent", name="TestAgent")
+    executor = AgentExecutor(agent, id="test_exec")
+
+    fi_kwargs, ci_kwargs = executor._prepare_agent_run_args({})  # pyright: ignore[reportPrivateUsage]
+    assert fi_kwargs is None
+    assert ci_kwargs is None
+
+
+class _NonCopyableRaw:
+    """Simulates an LLM SDK response object that cannot be deep-copied (e.g., proto/gRPC)."""
+
+    def __deepcopy__(self, memo: dict) -> Any:
+        raise TypeError("Cannot deepcopy this object")
+
+
+class _AgentWithRawRepr(BaseAgent):
+    """Agent that returns responses with a non-copyable raw_representation."""
+
+    def __init__(self, raw: Any, **kwargs: Any):
+        super().__init__(**kwargs)
+        self._raw = raw
+
+    def run(
+        self,
+        messages: str | Message | list[str] | list[Message] | None = None,
+        *,
+        stream: bool = False,
+        session: AgentSession | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+        async def _run() -> AgentResponse:
+            return AgentResponse(
+                messages=[Message("assistant", [f"reply from {self.name}"])],
+                raw_representation=self._raw,
+            )
+
+        return _run()
+
+
+async def test_agent_executor_workflow_with_non_copyable_raw_representation() -> None:
+    """Workflow should complete when AgentResponse contains a raw_representation that cannot be deep-copied."""
+    raw = _NonCopyableRaw()
+
+    agent_a = _AgentWithRawRepr(raw=raw, id="a", name="AgentA")
+    agent_b = _CountingAgent(id="b", name="AgentB")
+
+    exec_a = AgentExecutor(agent_a, id="exec_a")  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+    exec_b = AgentExecutor(agent_b, id="exec_b")
+
+    workflow = WorkflowBuilder(start_executor=exec_a).add_edge(exec_a, exec_b).build()
+    events = await workflow.run("hello")
+
+    completed = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_completed"]
+    completed_a = [e for e in completed if e.executor_id == "exec_a"]
+
+    assert len(completed_a) == 1
+    assert completed_a[0].data is not None
+
+    # The yielded AgentResponse should preserve its raw_representation reference
+    agent_responses = [d for d in completed_a[0].data if isinstance(d, AgentResponse)]
+    assert len(agent_responses) > 0
+    assert agent_responses[0].text == "reply from AgentA"
+    assert agent_responses[0].raw_representation is raw
+
+
+# ---------------------------------------------------------------------------
+# Context mode tests
+# ---------------------------------------------------------------------------
+
+
+class _MessageCapturingAgent(BaseAgent):
+    """Agent that records the messages it received and returns a configurable reply."""
+
+    def __init__(self, *, reply_text: str = "reply", **kwargs: Any):
+        super().__init__(**kwargs)
+        self.reply_text = reply_text
+        self.last_messages: list[Message] = []
+
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
+    def run(
+        self,
+        messages: AgentRunInputs | None = None,
+        *,
+        stream: bool = False,
+        session: AgentSession | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
+        captured: list[Message] = []
+        if messages:
+            for m in messages:  # type: ignore[union-attr]  # ty: ignore[not-iterable]
+                if isinstance(m, Message):
+                    captured.append(m)
+                elif isinstance(m, str):
+                    captured.append(Message("user", [m]))
+        self.last_messages = captured
+
+        if stream:
+
+            async def _stream() -> AsyncIterable[AgentResponseUpdate]:
+                yield AgentResponseUpdate(contents=[Content.from_text(text=self.reply_text)])
+
+            return ResponseStream(_stream(), finalizer=AgentResponse.from_updates)
+
+        async def _run() -> AgentResponse:
+            return AgentResponse(messages=[Message("assistant", [self.reply_text])])
+
+        return _run()
+
+
+def test_context_mode_custom_requires_context_filter() -> None:
+    """context_mode='custom' without context_filter must raise ValueError."""
+    agent = _CountingAgent(id="a", name="A")
+    with pytest.raises(ValueError, match="context_filter must be provided"):
+        AgentExecutor(agent, context_mode="custom")
+
+
+def test_context_mode_custom_with_filter_succeeds() -> None:
+    """context_mode='custom' with a context_filter should not raise."""
+    agent = _CountingAgent(id="a", name="A")
+    executor = AgentExecutor(agent, context_mode="custom", context_filter=lambda msgs: msgs[-1:])
+    assert executor._context_mode == "custom"  # pyright: ignore[reportPrivateUsage]
+    assert executor._context_filter is not None  # pyright: ignore[reportPrivateUsage]
+
+
+def test_context_mode_defaults_to_full() -> None:
+    """Default context_mode should be 'full'."""
+    agent = _CountingAgent(id="a", name="A")
+    executor = AgentExecutor(agent)
+    assert executor._context_mode == "full"  # pyright: ignore[reportPrivateUsage]
+
+
+def test_context_mode_invalid_value_raises() -> None:
+    """Invalid context_mode value should raise ValueError."""
+    agent = _CountingAgent(id="a", name="A")
+    with pytest.raises(ValueError, match="context_mode must be one of"):
+        AgentExecutor(agent, context_mode="invalid_mode")  # type: ignore
+
+
+async def test_from_response_context_mode_full_passes_full_conversation() -> None:
+    """context_mode='full' (default) should pass full_conversation to the second agent."""
+    first = _MessageCapturingAgent(id="first", name="First", reply_text="first reply")
+    second = _MessageCapturingAgent(id="second", name="Second", reply_text="second reply")
+
+    exec_a = AgentExecutor(first, id="exec_a")
+    exec_b = AgentExecutor(second, id="exec_b", context_mode="full")
+
+    wf = WorkflowBuilder(start_executor=exec_a).add_edge(exec_a, exec_b).build()
+
+    async for ev in wf.run("hello", stream=True):
+        if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
+            break
+
+    # Second agent should see full conversation: [user("hello"), assistant("first reply")]
+    seen = second.last_messages
+    assert len(seen) == 2
+    assert seen[0].role == "user" and "hello" in (seen[0].text or "")
+    assert seen[1].role == "assistant" and "first reply" in (seen[1].text or "")
+
+
+async def test_from_response_context_mode_last_agent_passes_only_agent_messages() -> None:
+    """context_mode='last_agent' should pass only the previous agent's response messages."""
+    first = _MessageCapturingAgent(id="first", name="First", reply_text="first reply")
+    second = _MessageCapturingAgent(id="second", name="Second", reply_text="second reply")
+
+    exec_a = AgentExecutor(first, id="exec_a")
+    exec_b = AgentExecutor(second, id="exec_b", context_mode="last_agent")
+
+    wf = WorkflowBuilder(start_executor=exec_a).add_edge(exec_a, exec_b).build()
+
+    async for ev in wf.run("hello", stream=True):
+        if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
+            break
+
+    # Second agent should see only the assistant message from first: [assistant("first reply")]
+    seen = second.last_messages
+    assert len(seen) == 1
+    assert seen[0].role == "assistant" and "first reply" in (seen[0].text or "")
+
+
+async def test_from_response_context_mode_custom_uses_filter() -> None:
+    """context_mode='custom' should invoke context_filter on full_conversation."""
+    first = _MessageCapturingAgent(id="first", name="First", reply_text="first reply")
+    second = _MessageCapturingAgent(id="second", name="Second", reply_text="second reply")
+
+    # Custom filter: keep only user messages
+    def only_user_messages(msgs: list[Message]) -> list[Message]:
+        return [m for m in msgs if m.role == "user"]
+
+    exec_a = AgentExecutor(first, id="exec_a")
+    exec_b = AgentExecutor(second, id="exec_b", context_mode="custom", context_filter=only_user_messages)
+
+    wf = WorkflowBuilder(start_executor=exec_a).add_edge(exec_a, exec_b).build()
+
+    async for ev in wf.run("hello", stream=True):
+        if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
+            break
+
+    # Second agent should see only user messages: [user("hello")]
+    seen = second.last_messages
+    assert len(seen) == 1
+    assert seen[0].role == "user" and "hello" in (seen[0].text or "")
+
+
+async def test_checkpoint_save_does_not_include_context_mode() -> None:
+    """on_checkpoint_save should not include context_mode in the saved state."""
+    agent = _CountingAgent(id="a", name="A")
+    executor = AgentExecutor(agent, context_mode="last_agent")
+
+    state = await executor.on_checkpoint_save()
+
+    assert "context_mode" not in state
+    assert "cache" in state
+    assert "agent_session" in state
+
+
+async def test_checkpoint_restore_works_without_context_mode_in_state() -> None:
+    """on_checkpoint_restore should succeed when state does not contain context_mode."""
+    agent = _CountingAgent(id="a", name="A")
+    executor = AgentExecutor(agent, context_mode="last_agent")
+
+    # Simulate a checkpoint state without context_mode (as saved by the new code)
+    state: dict[str, Any] = {
+        "cache": [Message(role="user", contents=["cached msg"])],
+        "full_conversation": [],
+        "agent_session": AgentSession().to_dict(),
+        "pending_agent_requests": {},
+        "pending_responses_to_agent": [],
+    }
+
+    await executor.on_checkpoint_restore(state)
+
+    cache = executor._cache  # pyright: ignore[reportPrivateUsage]
+    assert len(cache) == 1
+    assert cache[0].text == "cached msg"
+    # context_mode should remain as configured in the constructor, not changed by restore
+    assert executor._context_mode == "last_agent"  # pyright: ignore[reportPrivateUsage]
+
+
+# ---------------------------------------------------------------------------
+# Per-executor kwargs resolution tests
+# ---------------------------------------------------------------------------
+
+
+async def test_resolve_executor_kwargs_returns_global_kwargs() -> None:
+    """_resolve_executor_kwargs with the global kwargs key returns the global kwargs."""
+    agent = _CountingAgent(id="a", name="A")
+    executor = AgentExecutor(agent, id="exec_a")
+
+    resolved = {GLOBAL_KWARGS_KEY: {"tool_param": "value"}}
+    result = executor._resolve_executor_kwargs(resolved)  # pyright: ignore[reportPrivateUsage]
+    assert result == {"tool_param": "value"}
+
+
+async def test_resolve_executor_kwargs_returns_per_executor_kwargs() -> None:
+    """_resolve_executor_kwargs with matching executor ID returns that executor's kwargs."""
+    agent = _CountingAgent(id="a", name="A")
+    executor = AgentExecutor(agent, id="exec_a")
+
+    resolved = {"exec_a": {"my_param": 42}, "exec_b": {"other_param": 99}}
+    result = executor._resolve_executor_kwargs(resolved)  # pyright: ignore[reportPrivateUsage]
+    assert result == {"my_param": 42}
+
+
+async def test_resolve_executor_kwargs_returns_none_for_unmatched_per_executor() -> None:
+    """_resolve_executor_kwargs returns None when per-executor dict has no matching ID."""
+    agent = _CountingAgent(id="a", name="A")
+    executor = AgentExecutor(agent, id="exec_c")
+
+    resolved = {"exec_a": {"my_param": 42}, "exec_b": {"other_param": 99}}
+    result = executor._resolve_executor_kwargs(resolved)  # pyright: ignore[reportPrivateUsage]
+    assert result is None
+
+
+async def test_resolve_executor_kwargs_returns_none_for_none_input() -> None:
+    """_resolve_executor_kwargs returns None when input is None."""
+    agent = _CountingAgent(id="a", name="A")
+    executor = AgentExecutor(agent, id="exec_a")
+
+    result = executor._resolve_executor_kwargs(None)  # pyright: ignore[reportPrivateUsage]
+    assert result is None
+
+
+async def test_resolve_executor_kwargs_prefers_executor_id_over_global() -> None:
+    """_resolve_executor_kwargs prefers executor-specific entry over __global__."""
+    agent = _CountingAgent(id="a", name="A")
+    executor = AgentExecutor(agent, id="exec_a")
+
+    # Dict has both a per-executor entry and a global entry
+    resolved = {"exec_a": {"specific": True}, GLOBAL_KWARGS_KEY: {"global": True}}
+    result = executor._resolve_executor_kwargs(resolved)  # pyright: ignore[reportPrivateUsage]
+    assert result == {"specific": True}
+
+
+async def test_prepare_agent_run_args_extracts_function_invocation_kwargs() -> None:
+    """_prepare_agent_run_args extracts function_invocation_kwargs from the state dict."""
+    agent = _CountingAgent(id="a", name="A")
+    executor = AgentExecutor(agent, id="exec_a")
+
+    raw: dict[str, Any] = {
+        "function_invocation_kwargs": {GLOBAL_KWARGS_KEY: {"tool_key": "tool_val"}},
+    }
+    fi_kwargs, client_kwargs = executor._prepare_agent_run_args(raw)  # pyright: ignore[reportPrivateUsage]
+    assert fi_kwargs == {"tool_key": "tool_val"}
+    assert client_kwargs is None
+
+
+async def test_prepare_agent_run_args_extracts_client_kwargs() -> None:
+    """_prepare_agent_run_args extracts client_kwargs from the state dict."""
+    agent = _CountingAgent(id="a", name="A")
+    executor = AgentExecutor(agent, id="exec_a")
+
+    raw: dict[str, Any] = {
+        "client_kwargs": {GLOBAL_KWARGS_KEY: {"model": "gpt-4"}},
+    }
+    fi_kwargs, client_kwargs = executor._prepare_agent_run_args(raw)  # pyright: ignore[reportPrivateUsage]
+    assert fi_kwargs is None
+    assert client_kwargs == {"model": "gpt-4"}
+
+
+async def test_prepare_agent_run_args_per_executor_resolution() -> None:
+    """_prepare_agent_run_args resolves per-executor function_invocation_kwargs using self.id."""
+    agent = _CountingAgent(id="a", name="A")
+    executor = AgentExecutor(agent, id="exec_a")
+
+    raw: dict[str, Any] = {
+        "function_invocation_kwargs": {
+            "exec_a": {"my_tool_key": "my_val"},
+            "exec_b": {"other_tool_key": "other_val"},
+        },
+    }
+    fi_kwargs, _ = executor._prepare_agent_run_args(raw)  # pyright: ignore[reportPrivateUsage]
+    assert fi_kwargs == {"my_tool_key": "my_val"}
+
+
+async def test_prepare_agent_run_args_per_executor_no_match() -> None:
+    """_prepare_agent_run_args returns None for function_invocation_kwargs when executor ID not found."""
+    agent = _CountingAgent(id="a", name="A")
+    executor = AgentExecutor(agent, id="exec_c")
+
+    raw: dict[str, Any] = {
+        "function_invocation_kwargs": {
+            "exec_a": {"my_tool_key": "my_val"},
+            "exec_b": {"other_tool_key": "other_val"},
+        },
+    }
+    fi_kwargs, _ = executor._prepare_agent_run_args(raw)  # pyright: ignore[reportPrivateUsage]
+    assert fi_kwargs is None
+
+
+async def test_resolve_executor_kwargs_empty_per_executor_does_not_fallback_to_global() -> None:
+    """An explicit empty per-executor dict should not fall through to global kwargs."""
+    agent = _CountingAgent(id="a", name="A")
+    executor = AgentExecutor(agent, id="exec_a")
+
+    # Per-executor entry for exec_a is empty, but global has values.
+    # The empty dict should be honoured (no fallback to global).
+    resolved = {"exec_a": {}, GLOBAL_KWARGS_KEY: {"global_key": "global_val"}}  # type: ignore[var-annotated]
+    result = executor._resolve_executor_kwargs(resolved)  # pyright: ignore[reportPrivateUsage]
+    assert result == {}
+
+
+# region Tool approval emission
+
+
+class _ApprovalEmittingAgent(BaseAgent):
+    """Agent that returns a single ``function_approval_request`` Content.
+
+    Used to verify that ``AgentExecutor`` does *not* surface the approval
+    payload via both an ``output`` event and a ``request_info`` event in the
+    same superstep — only the ``request_info`` event must carry it.
+    """
+
+    def __init__(
+        self,
+        *,
+        approval_request_id: str = "apr_1",
+        tool_name: str = "delete_file",
+        tool_arguments: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ):
+        super().__init__(**kwargs)
+        self._approval_request_id = approval_request_id
+        self._tool_name = tool_name
+        self._tool_arguments: dict[str, Any] = tool_arguments or {"path": "/tmp/secret.txt"}
+        self.run_count = 0
+
+    def _build_approval_content(self) -> Content:
+        function_call = Content.from_function_call(
+            call_id=self._approval_request_id,
+            name=self._tool_name,
+            arguments=self._tool_arguments,
+        )
+        return Content.from_function_approval_request(id=self._approval_request_id, function_call=function_call)
+
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
+
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
+    def run(
+        self,
+        messages: AgentRunInputs | None = None,
+        *,
+        stream: bool = False,
+        session: AgentSession | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
+        self.run_count += 1
+        approval = self._build_approval_content()
+
+        if stream:
+
+            async def _stream() -> AsyncIterable[AgentResponseUpdate]:
+                yield AgentResponseUpdate(contents=[approval], role="assistant")
+
+            return ResponseStream(_stream(), finalizer=AgentResponse.from_updates)
+
+        async def _run() -> AgentResponse:
+            return AgentResponse(messages=[Message("assistant", [approval])])
+
+        return _run()
+
+
+def _has_approval_payload(event: WorkflowEvent[Any]) -> bool:
+    """Return True if the event's data carries a ``function_approval_request`` content."""
+    data: Any = event.data
+
+    def _contents_of(value: Any) -> list[Content]:
+        if isinstance(value, AgentResponseUpdate):
+            return list(value.contents)
+        if isinstance(value, AgentResponse):
+            return [c for m in value.messages for c in m.contents]
+        if isinstance(value, AgentExecutorResponse):
+            return [c for m in value.agent_response.messages for c in m.contents]
+        if isinstance(value, Message):
+            return list(value.contents)
+        if isinstance(value, Content):
+            return [value]
+        return []
+
+    return any(c.type == "function_approval_request" for c in _contents_of(data))
+
+
+async def test_agent_executor_does_not_double_emit_approval_non_streaming() -> None:
+    """Non-streaming: approval payload must only appear in the ``request_info`` event.
+
+    Regression test for the bug where ``AgentExecutor._run_agent`` first
+    ``yield_output``-ed the response (carrying the approval Content) and then
+    additionally emitted a ``request_info`` event for the same payload.
+    """
+    agent = _ApprovalEmittingAgent(id="approve_agent", name="ApproveAgent", approval_request_id="apr_ns_1")
+    executor = AgentExecutor(agent, id="approve_exec")
+    workflow = WorkflowBuilder(start_executor=executor).build()
+
+    request_info_events: list[WorkflowEvent[Any]] = []
+    output_events: list[WorkflowEvent[Any]] = []
+
+    for event in await workflow.run("please delete it"):
+        if event.type == "request_info":
+            request_info_events.append(event)
+        elif event.type == "output":
+            output_events.append(event)
+
+    assert len(request_info_events) == 1
+    assert _has_approval_payload(request_info_events[0])
+    # The approval payload must not also be surfaced as a workflow output.
+    assert not any(_has_approval_payload(e) for e in output_events)
+    assert agent.run_count == 1
+
+
+async def test_agent_executor_does_not_double_emit_approval_streaming() -> None:
+    """Streaming: per-update approval payload must not be ``yield_output``-ed."""
+    agent = _ApprovalEmittingAgent(id="approve_agent_s", name="ApproveAgentS", approval_request_id="apr_st_1")
+    executor = AgentExecutor(agent, id="approve_exec_s")
+    workflow = WorkflowBuilder(start_executor=executor).build()
+
+    request_info_events: list[WorkflowEvent[Any]] = []
+    output_events: list[WorkflowEvent[Any]] = []
+
+    async for event in workflow.run("please delete it", stream=True):
+        if event.type == "request_info":
+            request_info_events.append(event)
+        elif event.type == "output":
+            output_events.append(event)
+
+    assert len(request_info_events) == 1
+    assert _has_approval_payload(request_info_events[0])
+    assert not any(_has_approval_payload(e) for e in output_events)
+    assert agent.run_count == 1
+
+
+async def test_agent_executor_request_info_uses_user_input_request_id() -> None:
+    """``ctx.request_info`` must register the request under the agent's approval id.
+
+    This makes the workflow's pending-request id round-trip with the
+    ``function_approval_response.id`` the caller echoes back, so
+    ``Workflow._send_responses_internal`` can look it up directly.
+    """
+    agent = _ApprovalEmittingAgent(id="approve_agent_id", name="ApproveAgentId", approval_request_id="apr_match")
+    executor = AgentExecutor(agent, id="approve_exec_id")
+    workflow = WorkflowBuilder(start_executor=executor).build()
+
+    request_info_events: list[WorkflowEvent[Any]] = []
+    async for event in workflow.run("please delete it", stream=True):
+        if event.type == "request_info":
+            request_info_events.append(event)
+
+    assert len(request_info_events) == 1
+    assert request_info_events[0].request_id == "apr_match"
+
+
+# endregion Tool approval emission

python/packages/core/tests/workflow/test_agent_executor_tool_calls.py

@@ -0,0 +1,606 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for AgentExecutor handling of tool calls and results in streaming mode."""
+
+from collections.abc import AsyncIterable, Awaitable, Mapping, Sequence
+from typing import Any, Literal, overload
+
+from typing_extensions import Never
+
+from agent_framework import (
+    Agent,
+    AgentExecutor,
+    AgentExecutorResponse,
+    AgentResponse,
+    AgentResponseUpdate,
+    AgentRunInputs,
+    AgentSession,
+    BaseAgent,
+    ChatResponse,
+    ChatResponseUpdate,
+    Content,
+    FunctionTool,
+    Message,
+    ResponseStream,
+    WorkflowBuilder,
+    WorkflowContext,
+    WorkflowEvent,
+    executor,
+    tool,
+)
+from agent_framework._clients import BaseChatClient
+from agent_framework._tools import FunctionInvocationLayer
+
+
+class _ToolCallingAgent(BaseAgent):
+    """Mock agent that simulates tool calls and results in streaming mode."""
+
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
+
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
+    def run(
+        self,
+        messages: AgentRunInputs | None = None,
+        *,
+        stream: bool = False,
+        session: AgentSession | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
+        if stream:
+            return ResponseStream(self._run_stream_impl(), finalizer=AgentResponse.from_updates)
+
+        async def _run() -> AgentResponse[Any]:
+            return AgentResponse(messages=[Message("assistant", ["done"])])
+
+        return _run()
+
+    async def _run_stream_impl(self) -> AsyncIterable[AgentResponseUpdate]:
+        """Simulate streaming with tool calls and results."""
+        # First update: some text
+        yield AgentResponseUpdate(
+            contents=[Content.from_text(text="Let me search for that...")],
+            role="assistant",
+        )
+
+        # Second update: tool call (no text!)
+        yield AgentResponseUpdate(
+            contents=[
+                Content.from_function_call(
+                    call_id="call_123",
+                    name="search",
+                    arguments={"query": "weather"},
+                )
+            ],
+            role="assistant",
+        )
+
+        # Third update: tool result (no text!)
+        yield AgentResponseUpdate(
+            contents=[
+                Content.from_function_result(
+                    call_id="call_123",
+                    result={"temperature": 72, "condition": "sunny"},
+                )
+            ],
+            role="tool",
+        )
+
+        # Fourth update: final text response
+        yield AgentResponseUpdate(
+            contents=[Content.from_text(text="The weather is sunny, 72°F.")],
+            role="assistant",
+        )
+
+
+async def test_agent_executor_emits_tool_calls_in_streaming_mode() -> None:
+    """Test that AgentExecutor emits updates containing FunctionCallContent and FunctionResultContent."""
+    # Arrange
+    agent = _ToolCallingAgent(id="tool_agent", name="ToolAgent")
+    agent_exec = AgentExecutor(agent, id="tool_exec")
+
+    workflow = WorkflowBuilder(start_executor=agent_exec).build()
+
+    # Act: run in streaming mode
+    events: list[WorkflowEvent[AgentResponseUpdate]] = []
+    async for event in workflow.run("What's the weather?", stream=True):
+        if event.type == "output" and isinstance(event.data, AgentResponseUpdate):
+            events.append(event)
+
+    # Assert: we should receive 4 events (text, function call, function result, text)
+    assert len(events) == 4, f"Expected 4 events, got {len(events)}"
+
+    # First event: text update
+    assert events[0].data is not None
+    assert events[0].data.contents[0].type == "text"
+    assert events[0].data.contents[0].text is not None
+    assert "Let me search" in events[0].data.contents[0].text
+
+    # Second event: function call
+    assert events[1].data is not None
+    assert events[1].data.contents[0].type == "function_call"
+    func_call = events[1].data.contents[0]
+    assert func_call.call_id == "call_123"
+    assert func_call.name == "search"
+
+    # Third event: function result
+    assert events[2].data is not None
+    assert events[2].data.contents[0].type == "function_result"
+    func_result = events[2].data.contents[0]
+    assert func_result.call_id == "call_123"
+
+    # Fourth event: final text
+    assert events[3].data is not None
+    assert events[3].data.contents[0].type == "text"
+    assert events[3].data.contents[0].text is not None
+    assert "sunny" in events[3].data.contents[0].text
+
+
+@tool(approval_mode="always_require")
+def mock_tool_requiring_approval(query: str) -> str:
+    """Mock tool that requires approval before execution."""
+    return f"Executed tool with query: {query}"
+
+
+class MockChatClient(FunctionInvocationLayer[Any], BaseChatClient[Any]):
+    """Simple implementation of a chat client with function invocation support.
+
+    This mock uses the proper layer hierarchy:
+    - FunctionInvocationLayer.get_response intercepts calls and handles tool invocation
+    - BaseChatClient.get_response prepares messages and calls _inner_get_response
+    - _inner_get_response provides the actual mock responses
+    """
+
+    def __init__(self, parallel_request: bool = False) -> None:
+        FunctionInvocationLayer.__init__(self)
+        BaseChatClient.__init__(self)
+        self._iteration: int = 0
+        self._parallel_request: bool = parallel_request
+
+    def _inner_get_response(
+        self,
+        *,
+        messages: Sequence[Message],
+        stream: bool,
+        options: Mapping[str, Any],
+        **kwargs: Any,
+    ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
+        """Provide mock responses for the function invocation layer."""
+        if stream:
+            return self._build_response_stream(self._stream_response())
+
+        async def _get_response() -> ChatResponse:
+            return self._create_response()
+
+        return _get_response()
+
+    def _create_response(self) -> ChatResponse:
+        """Create a mock response based on iteration count."""
+        if self._iteration == 0:
+            if self._parallel_request:
+                response = ChatResponse(
+                    messages=Message(
+                        "assistant",
+                        [
+                            Content.from_function_call(
+                                call_id="1", name="mock_tool_requiring_approval", arguments='{"query": "test"}'
+                            ),
+                            Content.from_function_call(
+                                call_id="2", name="mock_tool_requiring_approval", arguments='{"query": "test"}'
+                            ),
+                        ],
+                    )
+                )
+            else:
+                response = ChatResponse(
+                    messages=Message(
+                        "assistant",
+                        [
+                            Content.from_function_call(
+                                call_id="1", name="mock_tool_requiring_approval", arguments='{"query": "test"}'
+                            )
+                        ],
+                    )
+                )
+        else:
+            response = ChatResponse(messages=Message("assistant", ["Tool executed successfully."]))
+
+        self._iteration += 1
+        return response
+
+    async def _stream_response(self) -> AsyncIterable[ChatResponseUpdate]:
+        """Generate mock streaming responses."""
+        if self._iteration == 0:
+            if self._parallel_request:
+                yield ChatResponseUpdate(
+                    contents=[
+                        Content.from_function_call(
+                            call_id="1", name="mock_tool_requiring_approval", arguments='{"query": "test"}'
+                        ),
+                        Content.from_function_call(
+                            call_id="2", name="mock_tool_requiring_approval", arguments='{"query": "test"}'
+                        ),
+                    ],
+                    role="assistant",
+                )
+            else:
+                yield ChatResponseUpdate(
+                    contents=[
+                        Content.from_function_call(
+                            call_id="1", name="mock_tool_requiring_approval", arguments='{"query": "test"}'
+                        )
+                    ],
+                    role="assistant",
+                )
+        else:
+            yield ChatResponseUpdate(contents=[Content.from_text(text="Tool executed ")], role="assistant")
+            yield ChatResponseUpdate(contents=[Content.from_text(text="successfully.")], role="assistant")
+
+        self._iteration += 1
+
+
+@executor(id="test_executor")
+async def test_executor(agent_executor_response: AgentExecutorResponse, ctx: WorkflowContext[Never, str]) -> None:  # type: ignore[valid-type]
+    await ctx.yield_output(agent_executor_response.agent_response.text)
+
+
+async def test_agent_executor_tool_call_with_approval() -> None:
+    """Test that AgentExecutor handles tool calls requiring approval."""
+    # Arrange
+    agent = Agent(
+        client=MockChatClient(),
+        name="ApprovalAgent",
+        tools=[mock_tool_requiring_approval],
+    )
+
+    workflow = WorkflowBuilder(start_executor=agent, output_from=[test_executor]).add_edge(agent, test_executor).build()
+
+    # Act
+    events = await workflow.run("Invoke tool requiring approval")
+
+    # Assert
+    assert len(events.get_request_info_events()) == 1
+    approval_request = events.get_request_info_events()[0]
+    assert approval_request.data.type == "function_approval_request"
+    assert approval_request.data.function_call.name == "mock_tool_requiring_approval"
+    assert approval_request.data.function_call.arguments == '{"query": "test"}'
+
+    # Act
+    events = await workflow.run(
+        responses={approval_request.request_id: approval_request.data.to_function_approval_response(True)}
+    )
+
+    # Assert
+    final_response = events.get_outputs()
+    assert len(final_response) == 1
+    assert final_response[0] == "Tool executed successfully."
+
+
+async def test_agent_executor_tool_call_with_approval_streaming() -> None:
+    """Test that AgentExecutor handles tool calls requiring approval in streaming mode."""
+    # Arrange
+    agent = Agent(
+        client=MockChatClient(),
+        name="ApprovalAgent",
+        tools=[mock_tool_requiring_approval],
+    )
+
+    workflow = WorkflowBuilder(start_executor=agent).add_edge(agent, test_executor).build()
+
+    # Act
+    request_info_events: list[WorkflowEvent] = []
+    async for event in workflow.run("Invoke tool requiring approval", stream=True):
+        if event.type == "request_info":
+            request_info_events.append(event)
+
+    # Assert
+    assert len(request_info_events) == 1
+    approval_request = request_info_events[0]
+    assert approval_request.data.type == "function_approval_request"
+    assert approval_request.data.function_call.name == "mock_tool_requiring_approval"
+    assert approval_request.data.function_call.arguments == '{"query": "test"}'
+
+    # Act
+    output: str | None = None
+    async for event in workflow.run(
+        stream=True, responses={approval_request.request_id: approval_request.data.to_function_approval_response(True)}
+    ):
+        if event.type == "output":
+            output = event.data
+
+    # Assert
+    assert output is not None
+    assert output == "Tool executed successfully."
+
+
+async def test_agent_executor_parallel_tool_call_with_approval() -> None:
+    """Test that AgentExecutor handles parallel tool calls requiring approval."""
+    # Arrange
+    agent = Agent(
+        client=MockChatClient(parallel_request=True),
+        name="ApprovalAgent",
+        tools=[mock_tool_requiring_approval],
+    )
+
+    workflow = WorkflowBuilder(start_executor=agent, output_from=[test_executor]).add_edge(agent, test_executor).build()
+
+    # Act
+    events = await workflow.run("Invoke tool requiring approval")
+
+    # Assert
+    assert len(events.get_request_info_events()) == 2
+    for approval_request in events.get_request_info_events():
+        assert approval_request.data.type == "function_approval_request"
+        assert approval_request.data.function_call.name == "mock_tool_requiring_approval"
+        assert approval_request.data.function_call.arguments == '{"query": "test"}'
+
+    # Act
+    responses = {
+        approval_request.request_id: approval_request.data.to_function_approval_response(True)  # type: ignore
+        for approval_request in events.get_request_info_events()
+    }
+    events = await workflow.run(responses=responses)
+
+    # Assert
+    final_response = events.get_outputs()
+    assert len(final_response) == 1
+    assert final_response[0] == "Tool executed successfully."
+
+
+async def test_agent_executor_parallel_tool_call_with_approval_streaming() -> None:
+    """Test that AgentExecutor handles parallel tool calls requiring approval in streaming mode."""
+    # Arrange
+    agent = Agent(
+        client=MockChatClient(parallel_request=True),
+        name="ApprovalAgent",
+        tools=[mock_tool_requiring_approval],
+    )
+
+    workflow = WorkflowBuilder(start_executor=agent).add_edge(agent, test_executor).build()
+
+    # Act
+    request_info_events: list[WorkflowEvent] = []
+    async for event in workflow.run("Invoke tool requiring approval", stream=True):
+        if event.type == "request_info":
+            request_info_events.append(event)
+
+    # Assert
+    assert len(request_info_events) == 2
+    for approval_request in request_info_events:
+        assert approval_request.data.type == "function_approval_request"
+        assert approval_request.data.function_call.name == "mock_tool_requiring_approval"
+        assert approval_request.data.function_call.arguments == '{"query": "test"}'
+
+    # Act
+    responses = {
+        approval_request.request_id: approval_request.data.to_function_approval_response(True)  # type: ignore
+        for approval_request in request_info_events
+    }
+
+    output: str | None = None
+    async for event in workflow.run(stream=True, responses=responses):
+        if event.type == "output":
+            output = event.data
+
+    # Assert
+    assert output is not None
+    assert output == "Tool executed successfully."
+
+
+# --- Declaration-only tool tests ---
+
+declaration_only_tool = FunctionTool(
+    name="client_side_tool",
+    func=None,
+    description="A client-side tool that the framework cannot execute.",
+    input_model={"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]},
+)
+
+
+class DeclarationOnlyMockChatClient(FunctionInvocationLayer[Any], BaseChatClient[Any]):
+    """Mock chat client that calls a declaration-only tool on first iteration."""
+
+    def __init__(self, parallel_request: bool = False) -> None:
+        FunctionInvocationLayer.__init__(self)
+        BaseChatClient.__init__(self)
+        self._iteration: int = 0
+        self._parallel_request: bool = parallel_request
+
+    def _inner_get_response(
+        self,
+        *,
+        messages: Sequence[Message],
+        stream: bool,
+        options: Mapping[str, Any],
+        **kwargs: Any,
+    ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
+        if stream:
+            return self._build_response_stream(self._stream_response())
+
+        async def _get_response() -> ChatResponse:
+            return self._create_response()
+
+        return _get_response()
+
+    def _create_response(self) -> ChatResponse:
+        if self._iteration == 0:
+            if self._parallel_request:
+                response = ChatResponse(
+                    messages=Message(
+                        "assistant",
+                        [
+                            Content.from_function_call(
+                                call_id="1", name="client_side_tool", arguments='{"query": "test"}'
+                            ),
+                            Content.from_function_call(
+                                call_id="2", name="client_side_tool", arguments='{"query": "test2"}'
+                            ),
+                        ],
+                    )
+                )
+            else:
+                response = ChatResponse(
+                    messages=Message(
+                        "assistant",
+                        [
+                            Content.from_function_call(
+                                call_id="1", name="client_side_tool", arguments='{"query": "test"}'
+                            )
+                        ],
+                    )
+                )
+        else:
+            response = ChatResponse(messages=Message("assistant", ["Tool executed successfully."]))
+
+        self._iteration += 1
+        return response
+
+    async def _stream_response(self) -> AsyncIterable[ChatResponseUpdate]:
+        if self._iteration == 0:
+            if self._parallel_request:
+                yield ChatResponseUpdate(
+                    contents=[
+                        Content.from_function_call(call_id="1", name="client_side_tool", arguments='{"query": "test"}'),
+                        Content.from_function_call(
+                            call_id="2", name="client_side_tool", arguments='{"query": "test2"}'
+                        ),
+                    ],
+                    role="assistant",
+                )
+            else:
+                yield ChatResponseUpdate(
+                    contents=[
+                        Content.from_function_call(call_id="1", name="client_side_tool", arguments='{"query": "test"}')
+                    ],
+                    role="assistant",
+                )
+        else:
+            yield ChatResponseUpdate(contents=[Content.from_text(text="Tool executed ")], role="assistant")
+            yield ChatResponseUpdate(contents=[Content.from_text(text="successfully.")], role="assistant")
+
+        self._iteration += 1
+
+
+async def test_agent_executor_declaration_only_tool_emits_request_info() -> None:
+    """Test that AgentExecutor emits request_info when agent calls a declaration-only tool."""
+    agent = Agent(
+        client=DeclarationOnlyMockChatClient(),
+        name="DeclarationOnlyAgent",
+        tools=[declaration_only_tool],
+    )
+
+    workflow = WorkflowBuilder(start_executor=agent, output_from=[test_executor]).add_edge(agent, test_executor).build()
+
+    # Act
+    events = await workflow.run("Use the client side tool")
+
+    # Assert - workflow should pause with a request_info event
+    request_info_events = events.get_request_info_events()
+    assert len(request_info_events) == 1
+    request = request_info_events[0]
+    assert request.data.type == "function_call"
+    assert request.data.name == "client_side_tool"
+    assert request.data.call_id == "1"
+
+    # Act - provide the function result to resume the workflow
+    events = await workflow.run(
+        responses={
+            request.request_id: Content.from_function_result(call_id=request.data.call_id, result="client result")
+        }
+    )
+
+    # Assert - workflow should complete
+    final_response = events.get_outputs()
+    assert len(final_response) == 1
+    assert final_response[0] == "Tool executed successfully."
+
+
+async def test_agent_executor_declaration_only_tool_emits_request_info_streaming() -> None:
+    """Test that AgentExecutor emits request_info for declaration-only tools in streaming mode."""
+    agent = Agent(
+        client=DeclarationOnlyMockChatClient(),
+        name="DeclarationOnlyAgent",
+        tools=[declaration_only_tool],
+    )
+
+    workflow = WorkflowBuilder(start_executor=agent).add_edge(agent, test_executor).build()
+
+    # Act
+    request_info_events: list[WorkflowEvent] = []
+    async for event in workflow.run("Use the client side tool", stream=True):
+        if event.type == "request_info":
+            request_info_events.append(event)
+
+    # Assert
+    assert len(request_info_events) == 1
+    request = request_info_events[0]
+    assert request.data.type == "function_call"
+    assert request.data.name == "client_side_tool"
+    assert request.data.call_id == "1"
+
+    # Act - provide the function result
+    output: str | None = None
+    async for event in workflow.run(
+        stream=True,
+        responses={
+            request.request_id: Content.from_function_result(call_id=request.data.call_id, result="client result")
+        },
+    ):
+        if event.type == "output":
+            output = event.data
+
+    # Assert
+    assert output is not None
+    assert output == "Tool executed successfully."
+
+
+async def test_agent_executor_parallel_declaration_only_tool_emits_request_info() -> None:
+    """Test that AgentExecutor emits request_info for parallel declaration-only tool calls."""
+    agent = Agent(
+        client=DeclarationOnlyMockChatClient(parallel_request=True),
+        name="DeclarationOnlyAgent",
+        tools=[declaration_only_tool],
+    )
+
+    workflow = WorkflowBuilder(start_executor=agent, output_from=[test_executor]).add_edge(agent, test_executor).build()
+
+    # Act
+    events = await workflow.run("Use the client side tool")
+
+    # Assert - should get 2 request_info events
+    request_info_events = events.get_request_info_events()
+    assert len(request_info_events) == 2
+    for req in request_info_events:
+        assert req.data.type == "function_call"
+        assert req.data.name == "client_side_tool"
+
+    # Act - provide both function results
+    responses = {
+        req.request_id: Content.from_function_result(call_id=req.data.call_id, result=f"result for {req.data.call_id}")
+        for req in request_info_events
+    }
+    events = await workflow.run(responses=responses)
+
+    # Assert - workflow should complete
+    final_response = events.get_outputs()
+    assert len(final_response) == 1
+    assert final_response[0] == "Tool executed successfully."

python/packages/core/tests/workflow/test_agent_run_event_typing.py

@@ -0,0 +1,38 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for WorkflowEvent[T] generic type annotations."""
+
+from agent_framework import AgentResponse, AgentResponseUpdate, Message
+from agent_framework._workflows._events import WorkflowEvent
+
+
+def test_workflow_event_with_agent_response_data_type() -> None:
+    """Verify WorkflowEvent[AgentResponse].data is typed as AgentResponse."""
+    response = AgentResponse(messages=[Message(role="assistant", contents=["Hello"])])
+    event: WorkflowEvent[AgentResponse] = WorkflowEvent("intermediate", executor_id="test", data=response)
+
+    # This assignment should pass type checking without a cast
+    data: AgentResponse = event.data
+    assert data is not None
+    assert data.text == "Hello"
+
+
+def test_workflow_event_with_agent_response_update_data_type() -> None:
+    """Verify WorkflowEvent[AgentResponseUpdate].data is typed as AgentResponseUpdate."""
+    update = AgentResponseUpdate()
+    event: WorkflowEvent[AgentResponseUpdate] = WorkflowEvent("intermediate", executor_id="test", data=update)
+
+    # This assignment should pass type checking without a cast
+    data: AgentResponseUpdate = event.data
+    assert data is not None
+
+
+def test_workflow_event_repr() -> None:
+    """Verify WorkflowEvent.__repr__ uses consistent format."""
+    response = AgentResponse(messages=[Message(role="assistant", contents=["Hello"])])
+    event: WorkflowEvent[AgentResponse] = WorkflowEvent("intermediate", executor_id="test", data=response)
+
+    repr_str = repr(event)
+    assert "WorkflowEvent" in repr_str
+    assert "executor_id='test'" in repr_str
+    assert "data=" in repr_str

python/packages/core/tests/workflow/test_agent_utils.py

@@ -0,0 +1,86 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from collections.abc import Awaitable
+from typing import Any, Literal, overload
+
+from agent_framework import (
+    AgentResponse,
+    AgentResponseUpdate,
+    AgentRunInputs,
+    AgentSession,
+    ResponseStream,
+    ServiceSessionId,
+)
+from agent_framework._workflows._agent_utils import resolve_agent_id
+
+
+class MockAgent:
+    """Mock agent for testing agent utilities."""
+
+    def __init__(self, agent_id: str, name: str | None = None) -> None:
+        self.id: str = agent_id
+        self.name: str | None = name
+        self.description: str | None = None
+
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+    def run(  # type: ignore[empty-body]
+        self,
+        messages: AgentRunInputs | None = None,
+        *,
+        stream: bool = False,
+        session: AgentSession | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...  # ty: ignore[empty-body]
+
+    def create_session(self, **kwargs: Any) -> AgentSession:  # type: ignore[empty-body]  # ty: ignore[empty-body]
+        """Creates a new conversation session for the agent."""
+        ...
+
+    def get_session(self, *, service_session_id: str | ServiceSessionId, **kwargs: Any) -> AgentSession:
+        return AgentSession()
+
+
+def test_resolve_agent_id_with_name() -> None:
+    """Test that resolve_agent_id returns name when agent has a name."""
+    agent = MockAgent(agent_id="agent-123", name="MyAgent")
+    result = resolve_agent_id(agent)  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+    assert result == "MyAgent"
+
+
+def test_resolve_agent_id_without_name() -> None:
+    """Test that resolve_agent_id returns id when agent has no name."""
+    agent = MockAgent(agent_id="agent-456", name=None)
+    result = resolve_agent_id(agent)  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+    assert result == "agent-456"
+
+
+def test_resolve_agent_id_with_empty_name() -> None:
+    """Test that resolve_agent_id returns id when agent has empty string name."""
+    agent = MockAgent(agent_id="agent-789", name="")
+    result = resolve_agent_id(agent)  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+    assert result == "agent-789"
+
+
+def test_resolve_agent_id_prefers_name_over_id() -> None:
+    """Test that resolve_agent_id prefers name over id when both are set."""
+    agent = MockAgent(agent_id="agent-abc", name="PreferredName")
+    result = resolve_agent_id(agent)  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+    assert result == "PreferredName"
+    assert result != "agent-abc"

python/packages/core/tests/workflow/test_checkpoint_decode.py

@@ -0,0 +1,251 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import Any, cast
+
+import pytest
+
+from agent_framework import WorkflowCheckpointException
+from agent_framework._workflows._checkpoint_encoding import (
+    _TYPE_MARKER,  # type: ignore
+    decode_checkpoint_value,
+    encode_checkpoint_value,
+)
+
+
+@dataclass
+class SampleRequest:
+    """Sample request message for testing checkpoint encoding/decoding."""
+
+    request_id: str
+    prompt: str
+
+
+@dataclass
+class SampleResponse:
+    """Sample response message for testing checkpoint encoding/decoding."""
+
+    data: str
+    original_request: SampleRequest
+    request_id: str
+
+
+# --- Tests for round-trip encode/decode ---
+
+
+def test_roundtrip_simple_dataclass() -> None:
+    """Test encoding and decoding of a simple dataclass."""
+    original = SampleRequest(request_id="test-123", prompt="test prompt")
+
+    encoded = encode_checkpoint_value(original)
+    decoded = cast(SampleRequest, decode_checkpoint_value(encoded))
+
+    assert isinstance(decoded, SampleRequest)
+    assert decoded.request_id == "test-123"
+    assert decoded.prompt == "test prompt"
+
+
+def test_roundtrip_dataclass_with_nested_request() -> None:
+    """Test that dataclass with nested dataclass fields can be encoded and decoded correctly."""
+    original = SampleResponse(
+        data="approve",
+        original_request=SampleRequest(request_id="abc", prompt="prompt"),
+        request_id="abc",
+    )
+
+    encoded = encode_checkpoint_value(original)
+    decoded = cast(SampleResponse, decode_checkpoint_value(encoded))
+
+    assert isinstance(decoded, SampleResponse)
+    assert decoded.data == "approve"
+    assert decoded.request_id == "abc"
+    assert isinstance(decoded.original_request, SampleRequest)
+    assert decoded.original_request.prompt == "prompt"
+    assert decoded.original_request.request_id == "abc"
+
+
+def test_roundtrip_nested_structures() -> None:
+    """Test encoding and decoding of complex nested structures."""
+    nested_data = {
+        "requests": [
+            SampleRequest(request_id="req-1", prompt="first prompt"),
+            SampleRequest(request_id="req-2", prompt="second prompt"),
+        ],
+        "responses": {
+            "req-1": SampleResponse(
+                data="first response",
+                original_request=SampleRequest(request_id="req-1", prompt="first prompt"),
+                request_id="req-1",
+            ),
+        },
+    }
+
+    encoded = encode_checkpoint_value(nested_data)
+    decoded = decode_checkpoint_value(encoded)
+
+    assert isinstance(decoded, dict)
+    assert "requests" in decoded
+    assert "responses" in decoded
+
+    requests = cast(list[Any], decoded["requests"])
+    assert isinstance(requests, list)
+    assert len(requests) == 2
+    assert all(isinstance(req, SampleRequest) for req in requests)
+    first_request = cast(SampleRequest, requests[0])
+    second_request = cast(SampleRequest, requests[1])
+    assert first_request.request_id == "req-1"
+    assert second_request.request_id == "req-2"
+
+    responses = cast(dict[str, Any], decoded["responses"])
+    assert isinstance(responses, dict)
+    assert "req-1" in responses
+    response = cast(SampleResponse, responses["req-1"])
+    assert isinstance(response, SampleResponse)
+    assert response.data == "first response"
+    assert isinstance(response.original_request, SampleRequest)
+    assert response.original_request.request_id == "req-1"
+
+
+def test_roundtrip_datetime() -> None:
+    """Test round-trip encoding/decoding of datetime objects."""
+    original = datetime(2024, 5, 4, 12, 30, 45, tzinfo=timezone.utc)
+
+    encoded = encode_checkpoint_value(original)
+    decoded = decode_checkpoint_value(encoded)
+
+    assert isinstance(decoded, datetime)
+    assert decoded == original
+
+
+def test_roundtrip_primitives() -> None:
+    """Test that primitive types round-trip unchanged."""
+    for value in ["hello", 42, 3.14, True, False, None]:
+        assert decode_checkpoint_value(encode_checkpoint_value(value)) == value
+
+
+def test_roundtrip_dict_with_mixed_values() -> None:
+    """Test round-trip of a dict containing both primitives and complex types."""
+    original = {
+        "name": "test",
+        "request": SampleRequest(request_id="r1", prompt="p1"),
+        "count": 5,
+    }
+
+    encoded = encode_checkpoint_value(original)
+    decoded = decode_checkpoint_value(encoded)
+
+    assert decoded["name"] == "test"
+    assert decoded["count"] == 5
+    assert isinstance(decoded["request"], SampleRequest)
+    assert decoded["request"].request_id == "r1"
+
+
+# --- Tests for decode primitives ---
+
+
+def test_decode_string() -> None:
+    """Test decoding a string passes through unchanged."""
+    assert decode_checkpoint_value("hello") == "hello"
+
+
+def test_decode_integer() -> None:
+    """Test decoding an integer passes through unchanged."""
+    assert decode_checkpoint_value(42) == 42
+
+
+def test_decode_none() -> None:
+    """Test decoding None passes through unchanged."""
+    assert decode_checkpoint_value(None) is None
+
+
+# --- Tests for decode collections ---
+
+
+def test_decode_plain_dict() -> None:
+    """Test decoding a plain dictionary with primitive values."""
+    data = {"a": 1, "b": "two"}
+    assert decode_checkpoint_value(data) == {"a": 1, "b": "two"}
+
+
+def test_decode_plain_list() -> None:
+    """Test decoding a plain list with primitive values."""
+    data = [1, "two", 3.0]
+    assert decode_checkpoint_value(data) == [1, "two", 3.0]
+
+
+# --- Tests for type verification ---
+
+
+def test_decode_raises_on_type_mismatch() -> None:
+    """Test that decoding raises WorkflowCheckpointException when type doesn't match."""
+    # Encode a SampleRequest but tamper with the type marker
+    encoded = encode_checkpoint_value(SampleRequest(request_id="r1", prompt="p1"))
+    assert isinstance(encoded, dict)
+    encoded[_TYPE_MARKER] = "nonexistent.module:FakeClass"
+
+    with pytest.raises(WorkflowCheckpointException, match="Type mismatch"):
+        decode_checkpoint_value(encoded)
+
+
+class NotADataclass:  # noqa: B903
+    """A regular class that is not a dataclass."""
+
+    def __init__(self, value: str) -> None:
+        self.value = value
+
+
+def test_roundtrip_regular_class() -> None:
+    """Test that regular (non-dataclass) objects can be round-tripped via pickle."""
+    original = NotADataclass(value="test_value")
+
+    encoded = encode_checkpoint_value(original)
+    decoded = cast(NotADataclass, decode_checkpoint_value(encoded))
+
+    assert isinstance(decoded, NotADataclass)
+    assert decoded.value == "test_value"
+
+
+def test_roundtrip_tuple() -> None:
+    """Test that tuples preserve their type through encode/decode roundtrip."""
+    original = (1, "two", 3.0)
+
+    encoded = encode_checkpoint_value(original)
+    decoded = decode_checkpoint_value(encoded)
+
+    assert isinstance(decoded, tuple)
+    assert decoded == original
+
+
+def test_roundtrip_set() -> None:
+    """Test that sets preserve their type through encode/decode roundtrip."""
+    original = {1, 2, 3}
+
+    encoded = encode_checkpoint_value(original)
+    decoded = decode_checkpoint_value(encoded)
+
+    assert isinstance(decoded, set)
+    assert decoded == original
+
+
+def test_roundtrip_nested_tuple_in_dict() -> None:
+    """Test that tuples nested inside dicts preserve their type."""
+    original = {"items": (1, 2, 3), "name": "test"}
+
+    encoded = encode_checkpoint_value(original)
+    decoded = decode_checkpoint_value(encoded)
+
+    assert isinstance(decoded["items"], tuple)
+    assert decoded["items"] == (1, 2, 3)
+    assert decoded["name"] == "test"
+
+
+def test_roundtrip_set_in_list() -> None:
+    """Test that sets nested inside lists preserve their type."""
+    original = [{"tags": {1, 2, 3}}]
+
+    encoded = encode_checkpoint_value(original)
+    decoded = decode_checkpoint_value(encoded)
+
+    assert isinstance(decoded[0]["tags"], set)
+    assert decoded[0]["tags"] == {1, 2, 3}

python/packages/core/tests/workflow/test_checkpoint_encode.py

@@ -0,0 +1,315 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import json
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import Any, cast
+
+from agent_framework._workflows._checkpoint_encoding import (
+    _PICKLE_MARKER,  # pyright: ignore[reportPrivateUsage]
+    _TYPE_MARKER,  # pyright: ignore[reportPrivateUsage]
+    encode_checkpoint_value,
+)
+
+
+@dataclass
+class SimpleDataclass:
+    """A simple dataclass for testing encoding."""
+
+    name: str
+    value: int
+
+
+@dataclass
+class NestedDataclass:
+    """A dataclass with nested dataclass field."""
+
+    outer_name: str
+    inner: SimpleDataclass
+
+
+class ModelWithToDict:
+    """A class that implements to_dict/from_dict protocol."""
+
+    def __init__(self, data: str) -> None:
+        self.data = data
+
+    def to_dict(self) -> dict[str, Any]:
+        return {"data": self.data}
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> "ModelWithToDict":
+        return cls(data=d["data"])
+
+
+class UnknownObject:
+    """A class that doesn't support any serialization protocol."""
+
+    def __init__(self, value: str) -> None:
+        self.value = value
+
+    def __str__(self) -> str:
+        return f"UnknownObject({self.value})"
+
+
+# --- Tests for primitive encoding (pass-through) ---
+
+
+def test_encode_string() -> None:
+    """Test encoding a string value."""
+    assert encode_checkpoint_value("hello") == "hello"
+
+
+def test_encode_integer() -> None:
+    """Test encoding an integer value."""
+    assert encode_checkpoint_value(42) == 42
+
+
+def test_encode_float() -> None:
+    """Test encoding a float value."""
+    assert encode_checkpoint_value(3.14) == 3.14
+
+
+def test_encode_boolean_true() -> None:
+    """Test encoding a True boolean value."""
+    assert encode_checkpoint_value(True) is True
+
+
+def test_encode_boolean_false() -> None:
+    """Test encoding a False boolean value."""
+    assert encode_checkpoint_value(False) is False
+
+
+def test_encode_none() -> None:
+    """Test encoding a None value."""
+    assert encode_checkpoint_value(None) is None
+
+
+# --- Tests for collection encoding ---
+
+
+def test_encode_empty_dict() -> None:
+    """Test encoding an empty dictionary."""
+    assert encode_checkpoint_value({}) == {}
+
+
+def test_encode_simple_dict() -> None:
+    """Test encoding a simple dictionary with primitive values."""
+    data = {"name": "test", "count": 5, "active": True}
+    result = encode_checkpoint_value(data)
+    assert result == {"name": "test", "count": 5, "active": True}
+
+
+def test_encode_dict_with_non_string_keys() -> None:
+    """Test encoding a dictionary with non-string keys (converted to strings)."""
+    data = {1: "one", 2: "two"}
+    result = encode_checkpoint_value(data)
+    assert result == {"1": "one", "2": "two"}
+
+
+def test_encode_empty_list() -> None:
+    """Test encoding an empty list."""
+    assert encode_checkpoint_value([]) == []
+
+
+def test_encode_simple_list() -> None:
+    """Test encoding a simple list with primitive values."""
+    data = [1, 2, 3, "four"]
+    result = encode_checkpoint_value(data)
+    assert result == [1, 2, 3, "four"]
+
+
+def test_encode_tuple() -> None:
+    """Test encoding a tuple (pickled to preserve type)."""
+    data = (1, 2, 3)
+    result = encode_checkpoint_value(data)
+    assert isinstance(result, dict)
+    assert _PICKLE_MARKER in result
+    assert _TYPE_MARKER in result
+
+
+def test_encode_set() -> None:
+    """Test encoding a set (pickled to preserve type)."""
+    data = {1, 2, 3}
+    result = encode_checkpoint_value(data)
+    assert isinstance(result, dict)
+    assert _PICKLE_MARKER in result
+    assert _TYPE_MARKER in result
+
+
+def test_encode_nested_dict() -> None:
+    """Test encoding a nested dictionary structure."""
+    data = {"outer": {"inner": {"value": 42}}}
+    result = encode_checkpoint_value(data)
+    assert result == {"outer": {"inner": {"value": 42}}}
+
+
+def test_encode_list_of_dicts() -> None:
+    """Test encoding a list containing dictionaries."""
+    data = [{"a": 1}, {"b": 2}]
+    result = encode_checkpoint_value(data)
+    assert result == [{"a": 1}, {"b": 2}]
+
+
+# --- Tests for non-JSON-native types (pickled) ---
+
+
+def test_encode_simple_dataclass() -> None:
+    """Test encoding a simple dataclass produces a pickled entry."""
+    obj = SimpleDataclass(name="test", value=42)
+    result = encode_checkpoint_value(obj)
+
+    assert isinstance(result, dict)
+    assert _PICKLE_MARKER in result
+    assert _TYPE_MARKER in result
+    assert isinstance(result[_PICKLE_MARKER], str)  # base64 string
+
+
+def test_encode_nested_dataclass() -> None:
+    """Test encoding a dataclass with nested dataclass fields."""
+    inner = SimpleDataclass(name="inner", value=10)
+    outer = NestedDataclass(outer_name="outer", inner=inner)
+    result = encode_checkpoint_value(outer)
+
+    assert isinstance(result, dict)
+    assert _PICKLE_MARKER in result
+    assert _TYPE_MARKER in result
+
+
+def test_encode_list_of_dataclasses() -> None:
+    """Test encoding a list containing dataclass instances."""
+    data = [
+        SimpleDataclass(name="first", value=1),
+        SimpleDataclass(name="second", value=2),
+    ]
+    result = encode_checkpoint_value(data)
+
+    assert isinstance(result, list)
+    result_list = cast(list[Any], result)
+    assert len(result_list) == 2
+    for item in result_list:
+        assert _PICKLE_MARKER in item
+
+
+def test_encode_dict_with_dataclass_values() -> None:
+    """Test encoding a dictionary with dataclass values."""
+    data = {
+        "item1": SimpleDataclass(name="first", value=1),
+        "item2": SimpleDataclass(name="second", value=2),
+    }
+    result = encode_checkpoint_value(data)
+
+    assert isinstance(result, dict)
+    assert _PICKLE_MARKER in result["item1"]
+    assert _PICKLE_MARKER in result["item2"]
+
+
+def test_encode_model_with_to_dict() -> None:
+    """Test encoding an object with to_dict is pickled (not using to_dict)."""
+    obj = ModelWithToDict(data="test_data")
+    result = encode_checkpoint_value(obj)
+
+    assert isinstance(result, dict)
+    assert _PICKLE_MARKER in result
+
+
+def test_encode_unknown_object() -> None:
+    """Test that arbitrary objects are pickled."""
+    obj = UnknownObject(value="test")
+    result = encode_checkpoint_value(obj)
+
+    assert isinstance(result, dict)
+    assert _PICKLE_MARKER in result
+
+
+def test_encode_datetime() -> None:
+    """Test that datetime objects are pickled."""
+    dt = datetime(2024, 5, 4, 12, 30, 45, tzinfo=timezone.utc)
+    result = encode_checkpoint_value(dt)
+
+    assert isinstance(result, dict)
+    assert _PICKLE_MARKER in result
+
+
+# --- Tests for type marker ---
+
+
+def test_encode_type_marker_records_type_info() -> None:
+    """Test that encoded objects include correct type information."""
+    obj = SimpleDataclass(name="test", value=42)
+    result = encode_checkpoint_value(obj)
+
+    type_key = result[_TYPE_MARKER]
+    assert "SimpleDataclass" in type_key
+
+
+def test_encode_type_marker_uses_module_qualname_format() -> None:
+    """Test that type marker uses module:qualname format."""
+    obj = SimpleDataclass(name="test", value=42)
+    result = encode_checkpoint_value(obj)
+
+    type_key = result[_TYPE_MARKER]
+    assert ":" in type_key
+    module, qualname = type_key.split(":")
+    assert module  # non-empty module
+    assert qualname == "SimpleDataclass"
+
+
+# --- Tests for JSON serializability ---
+
+
+def test_encode_result_is_json_serializable() -> None:
+    """Test that encoded output is fully JSON-serializable."""
+    data = {
+        "dc": SimpleDataclass(name="test", value=42),
+        "model": ModelWithToDict(data="test"),
+        "dt": datetime.now(timezone.utc),
+        "nested": [SimpleDataclass(name="n", value=1)],
+    }
+
+    result = encode_checkpoint_value(data)
+    # Should not raise
+    json_str = json.dumps(result)
+    assert isinstance(json_str, str)
+
+
+# --- Tests for mixed complex structures ---
+
+
+def test_encode_complex_mixed_structure() -> None:
+    """Test encoding a complex structure with mixed types."""
+    data = {
+        "string_value": "hello",
+        "int_value": 42,
+        "float_value": 3.14,
+        "bool_value": True,
+        "none_value": None,
+        "list_value": [1, 2, 3],
+        "nested_dict": {"a": 1, "b": 2},
+        "dataclass_value": SimpleDataclass(name="test", value=100),
+    }
+
+    result = encode_checkpoint_value(data)
+
+    # Primitives and collections pass through
+    assert result["string_value"] == "hello"
+    assert result["int_value"] == 42
+    assert result["float_value"] == 3.14
+    assert result["bool_value"] is True
+    assert result["none_value"] is None
+    assert result["list_value"] == [1, 2, 3]
+    assert result["nested_dict"] == {"a": 1, "b": 2}
+    # Dataclass is pickled
+    assert _PICKLE_MARKER in result["dataclass_value"]
+
+
+def test_encode_preserves_dict_with_pickle_marker_key() -> None:
+    """Test that regular dicts containing _PICKLE_MARKER key are recursively encoded."""
+    data = {
+        _PICKLE_MARKER: "some_value",
+        "other_key": "test",
+    }
+    result = encode_checkpoint_value(data)
+    assert _PICKLE_MARKER in result
+    assert result[_PICKLE_MARKER] == "some_value"
+    assert result["other_key"] == "test"

python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py

@@ -0,0 +1,265 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for restricted checkpoint deserialization.
+
+These tests verify that persisted checkpoint loading uses a restricted
+unpickler by default:
+- Arbitrary callables are blocked during deserialization
+- __reduce__ payloads cannot execute code during deserialization
+- FileCheckpointStorage accepts allowed_checkpoint_types for extension
+- User-defined types are blocked unless explicitly allowed
+- Built-in safe types and framework types are always allowed
+"""
+
+import base64
+import os
+import pickle
+import tempfile
+from dataclasses import dataclass
+from datetime import datetime, timezone
+
+import pytest
+
+from agent_framework import WorkflowCheckpointException
+from agent_framework._workflows._checkpoint import FileCheckpointStorage
+from agent_framework._workflows._checkpoint_encoding import (
+    _PICKLE_MARKER,
+    _TYPE_MARKER,
+    decode_checkpoint_value,
+    encode_checkpoint_value,
+)
+
+
+class MaliciousPayload:
+    """A class whose __reduce__ executes code during unpickling."""
+
+    def __reduce__(self):
+        return (os.getpid, ())
+
+
+def test_restricted_decode_blocks_arbitrary_callable():
+    """Restricted decoding blocks arbitrary module-level callables."""
+    pickled = pickle.dumps(os.getpid, protocol=pickle.HIGHEST_PROTOCOL)
+    encoded_b64 = base64.b64encode(pickled).decode("ascii")
+
+    checkpoint_value = {
+        _PICKLE_MARKER: encoded_b64,
+        _TYPE_MARKER: "builtins:builtin_function_or_method",
+    }
+
+    with pytest.raises(WorkflowCheckpointException, match="deserialization blocked"):
+        decode_checkpoint_value(checkpoint_value, allowed_types=frozenset())
+
+
+def test_restricted_decode_blocks_reduce_payload():
+    """__reduce__-based payloads are blocked before code can execute."""
+    payload = MaliciousPayload()
+    pickled = pickle.dumps(payload, protocol=pickle.HIGHEST_PROTOCOL)
+    encoded_b64 = base64.b64encode(pickled).decode("ascii")
+
+    checkpoint_value = {
+        _PICKLE_MARKER: encoded_b64,
+        _TYPE_MARKER: f"{MaliciousPayload.__module__}:{MaliciousPayload.__qualname__}",
+    }
+
+    with pytest.raises(WorkflowCheckpointException, match="deserialization blocked"):
+        decode_checkpoint_value(checkpoint_value, allowed_types=frozenset())
+
+
+def test_restricted_decode_prevents_code_execution():
+    """Restricted deserialization prevents __reduce__ code from running."""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        marker_file = os.path.join(tmpdir, "checkpoint_test_marker")
+
+        payload_bytes = pickle.dumps(
+            type(
+                "Exploit",
+                (),
+                {
+                    "__reduce__": lambda self: (
+                        eval,
+                        (f"open({marker_file!r}, 'w').write('pwned')",),
+                    )
+                },
+            )(),
+            protocol=pickle.HIGHEST_PROTOCOL,
+        )
+        encoded_b64 = base64.b64encode(payload_bytes).decode("ascii")
+
+        checkpoint_value = {
+            _PICKLE_MARKER: encoded_b64,
+            _TYPE_MARKER: "builtins:int",
+        }
+        with pytest.raises(WorkflowCheckpointException, match="deserialization blocked"):
+            decode_checkpoint_value(checkpoint_value, allowed_types=frozenset())
+
+        assert not os.path.exists(marker_file), (
+            "Restricted unpickler should have prevented code execution, but the marker file was created."
+        )
+
+
+def test_file_checkpoint_storage_accepts_allowed_types():
+    """FileCheckpointStorage.__init__ accepts allowed_checkpoint_types."""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        storage = FileCheckpointStorage(
+            tmpdir,
+            allowed_checkpoint_types=["some.module:SomeType"],
+        )
+        assert storage is not None
+
+
+@dataclass
+class _AllowedTestState:
+    """Test dataclass that will be explicitly allowed."""
+
+    name: str
+    value: int
+
+
+def test_restricted_decode_blocks_unlisted_user_type():
+    """User-defined types are blocked when not in allowed_checkpoint_types."""
+    original = _AllowedTestState(name="test", value=42)
+    encoded = encode_checkpoint_value(original)
+
+    with pytest.raises(WorkflowCheckpointException, match="deserialization blocked"):
+        decode_checkpoint_value(encoded, allowed_types=frozenset())
+
+
+def test_restricted_decode_allows_listed_user_type():
+    """User-defined types are allowed when listed in allowed_types."""
+    original = _AllowedTestState(name="test", value=42)
+    encoded = encode_checkpoint_value(original)
+
+    type_key = f"{_AllowedTestState.__module__}:{_AllowedTestState.__qualname__}"
+    decoded = decode_checkpoint_value(encoded, allowed_types=frozenset({type_key}))
+
+    assert isinstance(decoded, _AllowedTestState)
+    assert decoded.name == "test"
+    assert decoded.value == 42
+
+
+def test_restricted_decode_allows_builtin_safe_types():
+    """Built-in safe types (datetime, set, etc.) are always allowed."""
+    test_values = [
+        datetime(2025, 1, 1, tzinfo=timezone.utc),
+        {1, 2, 3},
+        frozenset({4, 5, 6}),
+        (1, "two", 3.0),
+        complex(1, 2),
+    ]
+    for original in test_values:
+        encoded = encode_checkpoint_value(original)
+        decoded = decode_checkpoint_value(encoded, allowed_types=frozenset())
+        assert decoded == original
+
+
+def test_unrestricted_decode_allows_arbitrary_types():
+    """Without allowed_types, decode_checkpoint_value remains unrestricted."""
+    original = _AllowedTestState(name="test", value=42)
+    encoded = encode_checkpoint_value(original)
+
+    decoded = decode_checkpoint_value(encoded)
+
+    assert isinstance(decoded, _AllowedTestState)
+    assert decoded.name == "test"
+
+
+async def test_file_storage_blocks_unlisted_user_type():
+    """FileCheckpointStorage blocks user types not in allowed_checkpoint_types."""
+    from agent_framework import WorkflowCheckpoint
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # Save with a storage that allows the type
+        type_key = f"{_AllowedTestState.__module__}:{_AllowedTestState.__qualname__}"
+        save_storage = FileCheckpointStorage(tmpdir, allowed_checkpoint_types=[type_key])
+
+        checkpoint = WorkflowCheckpoint(
+            workflow_name="test",
+            graph_signature_hash="hash",
+            state={"data": _AllowedTestState(name="test", value=1)},
+        )
+        await save_storage.save(checkpoint)
+
+        # Load with a storage that does NOT allow the type
+        load_storage = FileCheckpointStorage(tmpdir)
+        with pytest.raises(WorkflowCheckpointException, match="deserialization blocked"):
+            await load_storage.load(checkpoint.checkpoint_id)
+
+
+async def test_file_storage_allows_listed_user_type():
+    """FileCheckpointStorage allows user types listed in allowed_checkpoint_types."""
+    from agent_framework import WorkflowCheckpoint
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        type_key = f"{_AllowedTestState.__module__}:{_AllowedTestState.__qualname__}"
+        storage = FileCheckpointStorage(tmpdir, allowed_checkpoint_types=[type_key])
+
+        checkpoint = WorkflowCheckpoint(
+            workflow_name="test",
+            graph_signature_hash="hash",
+            state={"data": _AllowedTestState(name="allowed", value=99)},
+        )
+        await storage.save(checkpoint)
+        loaded = await storage.load(checkpoint.checkpoint_id)
+
+        assert isinstance(loaded.state["data"], _AllowedTestState)
+        assert loaded.state["data"].name == "allowed"
+        assert loaded.state["data"].value == 99
+
+
+def test_restricted_unpickler_raises_pickle_error():
+    """_RestrictedUnpickler.find_class raises pickle.UnpicklingError, not a framework exception."""
+    from agent_framework._workflows._checkpoint_encoding import _RestrictedUnpickler
+
+    pickled = pickle.dumps(os.getpid, protocol=pickle.HIGHEST_PROTOCOL)
+
+    unpickler = _RestrictedUnpickler(pickled, frozenset())
+    with pytest.raises(pickle.UnpicklingError, match="deserialization blocked"):
+        unpickler.load()
+
+
+def test_restricted_decode_allows_openai_types():
+    """OpenAI SDK types are always allowed during restricted deserialization."""
+    from openai.types.chat.chat_completion import ChatCompletion, Choice
+    from openai.types.chat.chat_completion_message import ChatCompletionMessage
+    from openai.types.completion_usage import CompletionUsage
+
+    completion = ChatCompletion(
+        id="chatcmpl-test",
+        choices=[
+            Choice(
+                finish_reason="stop",
+                index=0,
+                message=ChatCompletionMessage(role="assistant", content="hello"),
+            )
+        ],
+        created=1700000000,
+        model="gpt-4",
+        object="chat.completion",
+        usage=CompletionUsage(completion_tokens=1, prompt_tokens=1, total_tokens=2),
+    )
+    encoded = encode_checkpoint_value(completion)
+    decoded = decode_checkpoint_value(encoded, allowed_types=frozenset())
+
+    assert isinstance(decoded, ChatCompletion)
+    assert decoded.id == "chatcmpl-test"
+    assert decoded.choices[0].message.content == "hello"
+
+
+def test_restricted_decode_allows_openai_response_types():
+    """OpenAI Responses API types are always allowed during restricted deserialization."""
+    from openai.types.responses.response_usage import InputTokensDetails, OutputTokensDetails, ResponseUsage
+
+    usage = ResponseUsage(
+        input_tokens=10,
+        output_tokens=20,
+        total_tokens=30,
+        input_tokens_details=InputTokensDetails(cached_tokens=0, cache_write_tokens=0),
+        output_tokens_details=OutputTokensDetails(reasoning_tokens=0),
+    )
+    encoded = encode_checkpoint_value(usage)
+    decoded = decode_checkpoint_value(encoded, allowed_types=frozenset())
+
+    assert isinstance(decoded, ResponseUsage)
+    assert decoded.input_tokens == 10
+    assert decoded.output_tokens == 20

python/packages/core/tests/workflow/test_checkpoint_validation.py

@@ -0,0 +1,168 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import pytest
+from typing_extensions import Never
+
+from agent_framework import (
+    WorkflowBuilder,
+    WorkflowCheckpointException,
+    WorkflowContext,
+    WorkflowExecutor,
+    WorkflowRunState,
+    handler,
+)
+from agent_framework._workflows._checkpoint import InMemoryCheckpointStorage
+from agent_framework._workflows._executor import Executor
+
+
+class StartExecutor(Executor):
+    @handler
+    async def run(self, message: str, ctx: WorkflowContext[str]) -> None:
+        await ctx.send_message(message, target_id="finish")
+
+
+class FinishExecutor(Executor):
+    @handler
+    async def finish(self, message: str, ctx: WorkflowContext[Never, str]) -> None:  # zuban: ignore
+        await ctx.yield_output(message)
+
+
+def build_workflow(storage: InMemoryCheckpointStorage, finish_id: str = "finish"):
+    start = StartExecutor(id="start")
+    finish = FinishExecutor(id=finish_id)
+
+    builder = WorkflowBuilder(max_iterations=3, start_executor=start, checkpoint_storage=storage).add_edge(
+        start, finish
+    )
+    return builder.build()
+
+
+async def test_resume_fails_when_graph_mismatch() -> None:
+    storage = InMemoryCheckpointStorage()
+    workflow = build_workflow(storage, finish_id="finish")
+
+    # Run once to create checkpoints
+    _ = [event async for event in workflow.run("hello", stream=True)]  # noqa: F841
+
+    checkpoints = await storage.list_checkpoints(workflow_name=workflow.name)
+    assert checkpoints, "expected at least one checkpoint to be created"
+    target_checkpoint = checkpoints[-1]
+
+    # Build a structurally different workflow (different finish executor id)
+    mismatched_workflow = build_workflow(storage, finish_id="finish_alt")
+
+    with pytest.raises(WorkflowCheckpointException, match="Workflow graph has changed"):
+        _ = [
+            event
+            async for event in mismatched_workflow.run(
+                checkpoint_id=target_checkpoint.checkpoint_id,
+                checkpoint_storage=storage,
+                stream=True,
+            )
+        ]
+
+
+async def test_resume_succeeds_when_graph_matches() -> None:
+    storage = InMemoryCheckpointStorage()
+    workflow = build_workflow(storage, finish_id="finish")
+    _ = [event async for event in workflow.run("hello", stream=True)]  # noqa: F841
+
+    checkpoints = sorted(await storage.list_checkpoints(workflow_name=workflow.name), key=lambda c: c.timestamp)
+    target_checkpoint = checkpoints[0]
+
+    resumed_workflow = build_workflow(storage, finish_id="finish")
+
+    events = [
+        event
+        async for event in resumed_workflow.run(
+            checkpoint_id=target_checkpoint.checkpoint_id,
+            checkpoint_storage=storage,
+            stream=True,
+        )
+    ]
+
+    assert any(event.type == "status" and event.state == WorkflowRunState.IDLE for event in events)
+
+
+# -- Sub-workflow checkpoint validation tests --
+
+
+class SubStartExecutor(Executor):
+    @handler
+    async def run(self, message: str, ctx: WorkflowContext[str]) -> None:
+        await ctx.send_message(message)
+
+
+class SubFinishExecutor(Executor):
+    @handler
+    async def finish(self, message: str, ctx: WorkflowContext[Never, str]) -> None:  # zuban: ignore
+        await ctx.yield_output(message)
+
+
+def build_sub_workflow(sub_finish_id: str = "sub_finish"):
+    sub_start = SubStartExecutor(id="sub_start")
+    sub_finish = SubFinishExecutor(id=sub_finish_id)
+    return WorkflowBuilder(start_executor=sub_start).add_edge(sub_start, sub_finish).build()
+
+
+def build_parent_workflow(storage: InMemoryCheckpointStorage, sub_finish_id: str = "sub_finish"):
+    sub_workflow = build_sub_workflow(sub_finish_id=sub_finish_id)
+    sub_executor = WorkflowExecutor(sub_workflow, id="sub_wf", allow_direct_output=True)
+
+    start = StartExecutor(id="start")
+    finish = FinishExecutor(id="finish")
+
+    builder = (
+        WorkflowBuilder(max_iterations=3, start_executor=start, checkpoint_storage=storage)
+        .add_edge(start, sub_executor)
+        .add_edge(sub_executor, finish)
+    )
+    return builder.build()
+
+
+async def test_resume_succeeds_when_sub_workflow_matches() -> None:
+    storage = InMemoryCheckpointStorage()
+    workflow = build_parent_workflow(storage, sub_finish_id="sub_finish")
+
+    _ = [event async for event in workflow.run("hello", stream=True)]
+
+    checkpoints = await storage.list_checkpoints(workflow_name=workflow.name)
+    assert checkpoints, "expected at least one checkpoint to be created"
+    target_checkpoint = checkpoints[-1]
+
+    resumed_workflow = build_parent_workflow(storage, sub_finish_id="sub_finish")
+
+    events = [
+        event
+        async for event in resumed_workflow.run(
+            checkpoint_id=target_checkpoint.checkpoint_id,
+            checkpoint_storage=storage,
+            stream=True,
+        )
+    ]
+
+    assert any(event.type == "status" and event.state == WorkflowRunState.IDLE for event in events)
+
+
+async def test_resume_fails_when_sub_workflow_changes() -> None:
+    storage = InMemoryCheckpointStorage()
+    workflow = build_parent_workflow(storage, sub_finish_id="sub_finish")
+
+    _ = [event async for event in workflow.run("hello", stream=True)]
+
+    checkpoints = await storage.list_checkpoints(workflow_name=workflow.name)
+    assert checkpoints, "expected at least one checkpoint to be created"
+    target_checkpoint = checkpoints[-1]
+
+    # Build parent with a structurally different sub-workflow (different executor id inside)
+    mismatched_workflow = build_parent_workflow(storage, sub_finish_id="sub_finish_alt")
+
+    with pytest.raises(WorkflowCheckpointException, match="Workflow graph has changed"):
+        _ = [
+            event
+            async for event in mismatched_workflow.run(
+                checkpoint_id=target_checkpoint.checkpoint_id,
+                checkpoint_storage=storage,
+                stream=True,
+            )
+        ]

python/packages/core/tests/workflow/test_executor_future.py

@@ -0,0 +1,124 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from typing import Any
+
+import pytest
+from pydantic import BaseModel
+
+from agent_framework import Executor, WorkflowContext, handler
+
+
+class MyTypeA(BaseModel):
+    pass
+
+
+class MyTypeB(BaseModel):
+    pass
+
+
+class MyTypeC(BaseModel):
+    pass
+
+
+class TestExecutorFutureAnnotations:
+    """Test suite for Executor with from __future__ import annotations."""
+
+    def test_handler_decorator_future_annotations(self):
+        """Test @handler decorator works with stringified annotations (issue #3898)."""
+
+        class MyExecutor(Executor):
+            @handler
+            async def example(self, input: str, ctx: WorkflowContext[MyTypeA, MyTypeB]) -> None:
+                pass
+
+        exec_instance = MyExecutor(id="test")
+        assert str in exec_instance._handlers  # pyright: ignore[reportPrivateUsage]
+        spec = exec_instance._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert spec["message_type"] is str
+        assert spec["output_types"] == [MyTypeA]
+        assert spec["workflow_output_types"] == [MyTypeB]
+
+    def test_handler_decorator_future_annotations_single_type_arg(self):
+        """Test @handler with single type argument and future annotations."""
+
+        class MyExecutor(Executor):
+            @handler
+            async def example(self, input: int, ctx: WorkflowContext[MyTypeA]) -> None:
+                pass
+
+        exec_instance = MyExecutor(id="test")
+        assert int in exec_instance._handlers  # pyright: ignore[reportPrivateUsage]
+        spec = exec_instance._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert spec["message_type"] is int
+        assert spec["output_types"] == [MyTypeA]
+
+    def test_handler_decorator_future_annotations_complex(self):
+        """Test @handler with complex type annotations and future annotations."""
+
+        class MyExecutor(Executor):
+            @handler
+            async def example(self, data: dict[str, Any], ctx: WorkflowContext[list[str]]) -> None:
+                pass
+
+        exec_instance = MyExecutor(id="test")
+        spec = exec_instance._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert spec["message_type"] == dict[str, Any]
+        assert spec["output_types"] == [list[str]]
+
+    def test_handler_decorator_future_annotations_bare_context(self):
+        """Test @handler with bare WorkflowContext and future annotations."""
+
+        class MyExecutor(Executor):
+            @handler
+            async def example(self, input: str, ctx: WorkflowContext) -> None:
+                pass
+
+        exec_instance = MyExecutor(id="test")
+        assert str in exec_instance._handlers  # pyright: ignore[reportPrivateUsage]
+        spec = exec_instance._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert spec["output_types"] == []
+        assert spec["workflow_output_types"] == []
+
+    def test_handler_decorator_future_annotations_explicit_types(self):
+        """Test @handler with explicit type parameters under future annotations."""
+
+        class MyExecutor(Executor):
+            @handler(input=str, output=MyTypeA)
+            async def example(self, input, ctx) -> None:  # type: ignore[no-untyped-def]
+                pass
+
+        exec_instance = MyExecutor(id="test")
+        assert str in exec_instance._handlers  # pyright: ignore[reportPrivateUsage]
+        spec = exec_instance._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert spec["message_type"] is str
+        assert spec["output_types"] == [MyTypeA]
+
+    def test_handler_decorator_future_annotations_union_context(self):
+        """Test @handler with union type context annotations and future annotations."""
+
+        class MyExecutor(Executor):
+            @handler
+            async def example(self, input: str, ctx: WorkflowContext[MyTypeA | MyTypeB, MyTypeC]) -> None:
+                pass
+
+        exec_instance = MyExecutor(id="test")
+        assert str in exec_instance._handlers  # pyright: ignore[reportPrivateUsage]
+        spec = exec_instance._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert spec["output_types"] == [MyTypeA, MyTypeB]
+        assert spec["workflow_output_types"] == [MyTypeC]
+
+    def test_handler_unresolvable_annotation_raises(self):
+        """Test that an unresolvable forward-reference annotation raises ValueError.
+
+        When get_type_hints fails (e.g. NameError for NonExistentType), the code falls back
+        to raw string annotations. The ctx parameter's raw string annotation is then not
+        recognised as a valid WorkflowContext type, so a ValueError is still raised.
+        """
+        with pytest.raises(ValueError):
+
+            class Bad(Executor):  # pyright: ignore[reportUnusedClass]
+                @handler  # pyright: ignore[reportUnknownArgumentType]
+                async def example(self, input: NonExistentType, ctx: WorkflowContext[MyTypeA, MyTypeB]) -> None:  # type: ignore[name-defined]  # ty: ignore[unresolved-reference]  # noqa: F821
+                    pass

python/packages/core/tests/workflow/test_full_conversation.py

@@ -0,0 +1,574 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from collections.abc import AsyncIterable, Awaitable
+from typing import Any, Literal, overload
+
+import pytest
+from pydantic import PrivateAttr
+from typing_extensions import Never
+
+from agent_framework import (
+    AgentExecutor,
+    AgentExecutorRequest,
+    AgentExecutorResponse,
+    AgentResponse,
+    AgentResponseUpdate,
+    AgentRunInputs,
+    AgentSession,
+    BaseAgent,
+    Content,
+    Executor,
+    Message,
+    ResponseStream,
+    ServiceSessionId,
+    WorkflowBuilder,
+    WorkflowContext,
+    WorkflowRunState,
+    executor,
+    handler,
+)
+from agent_framework.orchestrations import SequentialBuilder
+
+
+class _SimpleAgent(BaseAgent):
+    """Agent that returns a single assistant message."""
+
+    def __init__(self, *, reply_text: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self._reply_text = reply_text
+
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+    def run(
+        self,
+        messages: AgentRunInputs | None = None,
+        *,
+        stream: bool = False,
+        session: AgentSession | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
+        if stream:
+
+            async def _stream() -> AsyncIterable[AgentResponseUpdate]:
+                yield AgentResponseUpdate(contents=[Content.from_text(text=self._reply_text)])
+
+            return ResponseStream(_stream(), finalizer=AgentResponse.from_updates)
+
+        async def _run() -> AgentResponse:
+            return AgentResponse(messages=[Message("assistant", [self._reply_text])])
+
+        return _run()
+
+
+class _ToolHistoryAgent(BaseAgent):
+    """Agent that emits tool-call internals plus a final assistant summary."""
+
+    def __init__(self, *, summary_text: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self._summary_text = summary_text
+
+    def _messages(self) -> list[Message]:
+        return [
+            Message(
+                role="assistant",
+                contents=[
+                    Content.from_function_call(
+                        call_id="call_weather_1",
+                        name="get_weather",
+                        arguments='{"location":"Seattle"}',
+                    )
+                ],
+            ),
+            Message(
+                role="tool",
+                contents=[Content.from_function_result(call_id="call_weather_1", result="Sunny, 72F")],
+            ),
+            Message(role="assistant", contents=[Content.from_text(text=self._summary_text)]),
+        ]
+
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+    def run(
+        self,
+        messages: AgentRunInputs | None = None,
+        *,
+        stream: bool = False,
+        session: AgentSession | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
+        if stream:
+
+            async def _stream() -> AsyncIterable[AgentResponseUpdate]:
+                yield AgentResponseUpdate(
+                    contents=[
+                        Content.from_function_call(
+                            call_id="call_weather_1",
+                            name="get_weather",
+                            arguments='{"location":"Seattle"}',
+                        )
+                    ],
+                    role="assistant",
+                )
+                yield AgentResponseUpdate(
+                    contents=[Content.from_function_result(call_id="call_weather_1", result="Sunny, 72F")],
+                    role="tool",
+                )
+                yield AgentResponseUpdate(contents=[Content.from_text(text=self._summary_text)], role="assistant")
+
+            return ResponseStream(_stream(), finalizer=AgentResponse.from_updates)
+
+        async def _run() -> AgentResponse:
+            return AgentResponse(messages=self._messages())
+
+        return _run()
+
+
+class _CaptureFullConversation(Executor):
+    """Captures AgentExecutorResponse.full_conversation and completes the workflow."""
+
+    @handler
+    async def capture(self, response: AgentExecutorResponse, ctx: WorkflowContext[Never, dict[str, Any]]) -> None:  # type: ignore[valid-type]
+        full = response.full_conversation
+        # The AgentExecutor contract guarantees full_conversation is populated.
+        assert full is not None
+        payload = {
+            "length": len(full),
+            "roles": [m.role for m in full],
+            "texts": [m.text for m in full],
+        }
+        await ctx.yield_output(payload)
+        pass
+
+
+async def test_agent_executor_populates_full_conversation_non_streaming() -> None:
+    # Arrange: AgentExecutor will be non-streaming when using workflow.run()
+    agent = _SimpleAgent(id="agent1", name="A", reply_text="agent-reply")
+    agent_exec = AgentExecutor(agent, id="agent1-exec")
+    capturer = _CaptureFullConversation(id="capture")
+
+    wf = WorkflowBuilder(start_executor=agent_exec, output_from=[capturer]).add_edge(agent_exec, capturer).build()
+
+    # Act: use run() to test non-streaming mode
+    result = await wf.run("hello world")
+
+    # Extract output from run result
+    outputs = result.get_outputs()
+    assert len(outputs) == 1
+    payload = outputs[0]
+
+    # Assert: full_conversation contains [user("hello world"), assistant("agent-reply")]
+    assert isinstance(payload, dict)
+    assert payload["length"] == 2
+    assert payload["roles"][0] == "user" and "hello world" in (payload["texts"][0] or "")
+    assert payload["roles"][1] == "assistant" and "agent-reply" in (payload["texts"][1] or "")
+
+
+class _CaptureAgent(BaseAgent):
+    """Streaming-capable agent that records the messages it received."""
+
+    _last_messages: list[Message] = PrivateAttr(default_factory=list)  # type: ignore
+
+    def __init__(self, *, reply_text: str, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self._reply_text = reply_text
+
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+    def run(
+        self,
+        messages: AgentRunInputs | None = None,
+        *,
+        stream: bool = False,
+        session: AgentSession | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
+        # Normalize and record messages for verification
+        norm: list[Message] = []
+        if messages:
+            for m in messages:  # type: ignore[iteration-over-optional, union-attr]  # ty: ignore[not-iterable]
+                if isinstance(m, Message):
+                    norm.append(m)
+                elif isinstance(m, str):
+                    norm.append(Message("user", [m]))
+        self._last_messages = norm
+
+        if stream:
+
+            async def _stream() -> AsyncIterable[AgentResponseUpdate]:
+                yield AgentResponseUpdate(contents=[Content.from_text(text=self._reply_text)])
+
+            return ResponseStream(_stream(), finalizer=AgentResponse.from_updates)
+
+        async def _run() -> AgentResponse:
+            return AgentResponse(messages=[Message("assistant", [self._reply_text])])
+
+        return _run()
+
+
+async def test_sequential_adapter_uses_full_conversation() -> None:
+    # Arrange: two streaming agents; the second records what it receives
+    a1 = _CaptureAgent(id="agent1", name="A1", reply_text="A1 reply")
+    a2 = _CaptureAgent(id="agent2", name="A2", reply_text="A2 reply")
+
+    wf = SequentialBuilder(participants=[a1, a2]).build()
+
+    # Act
+    async for ev in wf.run("hello seq", stream=True):
+        if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
+            break
+
+    # Assert: second agent should have seen the user prompt and A1's assistant reply
+    seen = a2._last_messages  # pyright: ignore[reportPrivateUsage]
+    assert len(seen) == 2
+    assert seen[0].role == "user" and "hello seq" in (seen[0].text or "")
+    assert seen[1].role == "assistant" and "A1 reply" in (seen[1].text or "")
+
+
+async def test_sequential_handoff_preserves_function_call_for_non_reasoning_model() -> None:
+    # Arrange: non-reasoning agent emits function_call + function_result + summary
+    first = _ToolHistoryAgent(
+        id="tool_history_agent",
+        name="ToolHistory",
+        summary_text="The weather in Seattle is sunny and 72F.",
+    )
+    second = _CaptureAgent(id="capture_agent", name="Capture", reply_text="Captured")
+    wf = SequentialBuilder(participants=[first, second]).build()
+
+    # Act
+    result = await wf.run("Check weather and continue")
+
+    # Assert workflow completed
+    outputs = result.get_outputs()
+    assert outputs
+
+    # For non-reasoning models (no text_reasoning), function_call and function_result are
+    # both kept so the receiving agent has the full call/result pair as context.
+    seen = second._last_messages  # pyright: ignore[reportPrivateUsage]
+    assert len(seen) == 4  # user, assistant(function_call), tool(function_result), assistant(summary)
+    assert seen[0].role == "user"
+    assert "Check weather and continue" in (seen[0].text or "")
+    assert seen[1].role == "assistant"
+    assert any(content.type == "function_call" for content in seen[1].contents)
+    assert seen[2].role == "tool"
+    assert any(content.type == "function_result" for content in seen[2].contents)
+    assert seen[3].role == "assistant"
+    assert "Seattle is sunny" in (seen[3].text or "")
+    # No text_reasoning should appear (non-reasoning model)
+    assert all(content.type != "text_reasoning" for msg in seen for content in msg.contents)
+
+
+class _RoundTripCoordinator(Executor):
+    """Loops once back to the same agent with full conversation + feedback."""
+
+    def __init__(self, *, target_agent_id: str, id: str = "round_trip_coordinator") -> None:
+        super().__init__(id=id)
+        self._target_agent_id = target_agent_id
+        self._seen = 0
+
+    @handler
+    async def handle_response(
+        self,
+        response: AgentExecutorResponse,
+        ctx: WorkflowContext[AgentExecutorRequest, dict[str, Any]],
+    ) -> None:
+        self._seen += 1
+        if self._seen == 1:
+            assert response.full_conversation is not None
+            await ctx.send_message(
+                AgentExecutorRequest(
+                    messages=list(response.full_conversation) + [Message(role="user", contents=["apply feedback"])],
+                    should_respond=True,
+                ),
+                target_id=self._target_agent_id,
+            )
+            return
+
+        assert response.full_conversation is not None
+        await ctx.yield_output({
+            "roles": [m.role for m in response.full_conversation],
+            "texts": [m.text for m in response.full_conversation],
+        })
+
+
+async def test_agent_executor_full_conversation_round_trip_does_not_duplicate_history() -> None:
+    """When full history is replayed, AgentExecutor should not duplicate prior turns."""
+    agent = _SimpleAgent(id="writer_agent", name="Writer", reply_text="draft reply")
+    agent_exec = AgentExecutor(agent, id="writer_agent")
+    coordinator = _RoundTripCoordinator(target_agent_id="writer_agent")
+
+    wf = (
+        WorkflowBuilder(start_executor=agent_exec, output_from=[coordinator])
+        .add_edge(agent_exec, coordinator)
+        .add_edge(coordinator, agent_exec)
+        .build()
+    )
+
+    result = await wf.run("initial prompt")
+    outputs = result.get_outputs()
+    assert len(outputs) == 1
+    payload = outputs[0]
+    assert isinstance(payload, dict)
+
+    # Expected conversation after one loop:
+    # user(initial), assistant(first reply), user(feedback), assistant(second reply)
+    assert payload["roles"] == ["user", "assistant", "user", "assistant"]
+    assert payload["texts"][0] == "initial prompt"
+    assert payload["texts"][1] == "draft reply"
+    assert payload["texts"][2] == "apply feedback"
+    assert payload["texts"][3] == "draft reply"
+
+
+class _SessionIdCapturingAgent(BaseAgent):
+    """Records service_session_id of the session at run() time."""
+
+    _captured_service_session_id: str | ServiceSessionId | None = PrivateAttr(default="NOT_CAPTURED")
+
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+    def run(
+        self,
+        messages: AgentRunInputs | None = None,
+        *,
+        stream: bool = False,
+        session: AgentSession | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
+        self._captured_service_session_id = session.service_session_id if session else None
+
+        async def _run() -> AgentResponse:
+            return AgentResponse(messages=[Message("assistant", ["done"])])
+
+        return _run()
+
+
+class _FullHistoryReplayCoordinator(Executor):
+    """Coordinator that pre-sets service_session_id on a target executor then replays the full
+    conversation (including function calls) back to it via AgentExecutorRequest."""
+
+    def __init__(self, *, target_exec: AgentExecutor, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self._target_exec = target_exec
+
+    @handler
+    async def handle(
+        self,
+        response: AgentExecutorResponse,
+        ctx: WorkflowContext[AgentExecutorRequest, Any],
+    ) -> None:
+        full_conv = list(response.full_conversation or response.agent_response.messages)
+        full_conv.append(Message(role="user", contents=["follow-up"]))
+        # Simulate a prior run: the target executor has a stored previous_response_id.
+        self._target_exec._session.service_session_id = "resp_PREVIOUS_RUN"  # pyright: ignore[reportPrivateUsage]
+        await ctx.send_message(
+            AgentExecutorRequest(messages=full_conv, should_respond=True),
+            target_id=self._target_exec.id,
+        )
+
+
+@pytest.mark.xfail(
+    reason=(
+        "Tracks the executor-layer half of #3295: AgentExecutor should clear service_session_id "
+        "when handed a full prior conversation. The wire-level 'Duplicate item' API error is "
+        "already closed by the chat-client strip in #3295; this xfail covers the defense-in-depth "
+        "follow-up that makes the executor wiring reflect intent."
+    ),
+    strict=True,
+)
+async def test_run_request_with_full_history_clears_service_session_id() -> None:
+    """Replaying a full conversation (including function calls) via AgentExecutorRequest must
+    clear service_session_id so the API does not receive both previous_response_id and the
+    same function-call items in input — which would cause a 'Duplicate item' API error."""
+    tool_agent = _ToolHistoryAgent(id="tool_agent", name="ToolAgent", summary_text="Done.")
+    tool_exec = AgentExecutor(tool_agent, id="tool_agent")
+
+    spy_agent = _SessionIdCapturingAgent(id="spy_agent", name="SpyAgent")
+    spy_exec = AgentExecutor(spy_agent, id="spy_agent")
+
+    coordinator = _FullHistoryReplayCoordinator(id="coord", target_exec=spy_exec)
+
+    wf = (
+        WorkflowBuilder(start_executor=tool_exec, output_from=[coordinator])
+        .add_edge(tool_exec, coordinator)
+        .add_edge(coordinator, spy_exec)
+        .build()
+    )
+
+    result = await wf.run("initial prompt")
+    assert result.get_outputs() is not None
+
+    # The spy agent must have seen service_session_id=None (cleared before run).
+    # Without the fix, it would see "resp_PREVIOUS_RUN" and the API would raise
+    # "Duplicate item found" because the same function-call IDs appear in both
+    # previous_response_id (server-stored) and the explicit input messages.
+    assert spy_agent._captured_service_session_id is None  # pyright: ignore[reportPrivateUsage]
+
+
+async def test_from_response_preserves_service_session_id() -> None:
+    """from_response hands off a prior agent's full conversation to the next executor.
+    The receiving executor's service_session_id is preserved so the API can continue
+    the conversation using previous_response_id."""
+    tool_agent = _ToolHistoryAgent(id="tool_agent2", name="ToolAgent", summary_text="Done.")
+    tool_exec = AgentExecutor(tool_agent, id="tool_agent2")
+
+    spy_agent = _SessionIdCapturingAgent(id="spy_agent2", name="SpyAgent")
+    spy_exec = AgentExecutor(spy_agent, id="spy_agent2")
+    # Simulate a prior run on the spy executor.
+    spy_exec._session.service_session_id = "resp_PREVIOUS_RUN"  # pyright: ignore[reportPrivateUsage]
+
+    wf = WorkflowBuilder(start_executor=tool_exec, output_from=[spy_exec]).add_edge(tool_exec, spy_exec).build()
+
+    result = await wf.run("start")
+    assert result.get_outputs() is not None
+
+    assert spy_agent._captured_service_session_id == "resp_PREVIOUS_RUN"  # pyright: ignore[reportPrivateUsage]
+
+
+@executor(
+    id="upper_case_executor",
+    input=AgentExecutorResponse,
+    output=AgentExecutorResponse,
+    workflow_output=str,
+)
+async def _upper_case_executor(
+    response: AgentExecutorResponse,
+    ctx: WorkflowContext[AgentExecutorResponse, str],
+) -> None:
+    upper_text = response.agent_response.text.upper()
+    await ctx.send_message(response.with_text(upper_text))
+    await ctx.yield_output(upper_text)
+
+
+async def test_with_text_preserves_full_conversation_through_custom_executor() -> None:
+    """Custom executor using with_text must preserve the full conversation chain."""
+    # Mirrors the reproduction from issue #5246:
+    # agent1 ("User likes sky red") -> agent2 ("User likes sky blue") -> upper_case -> agent3 ("User likes sky green")
+    agent1 = AgentExecutor(
+        _SimpleAgent(id="agent1", name="ContextAgent1", reply_text="User likes sky red"), id="agent1"
+    )
+    agent2 = AgentExecutor(
+        _SimpleAgent(id="agent2", name="ContextAgent2", reply_text="User likes sky blue"), id="agent2"
+    )
+    agent3 = AgentExecutor(
+        _SimpleAgent(id="agent3", name="ContextAgent3", reply_text="User likes sky green"), id="agent3"
+    )
+    capturer = _CaptureFullConversation(id="capture")
+
+    wf = (
+        WorkflowBuilder(start_executor=agent1, output_from=[capturer])
+        .add_chain([agent1, agent2, _upper_case_executor, agent3, capturer])
+        .build()
+    )
+
+    result = await wf.run("")
+    payload = next(o for o in result.get_outputs() if isinstance(o, dict))
+
+    # The final agent must see the full conversation: user, agent1, UPPER(agent2), agent3
+    assert payload["roles"] == ["user", "assistant", "assistant", "assistant"]
+    assert payload["texts"][1] == "User likes sky red"
+    assert payload["texts"][2] == "USER LIKES SKY BLUE"
+    assert payload["texts"][3] == "User likes sky green"
+
+
+async def test_with_text_does_not_mutate_original() -> None:
+    """with_text returns a new instance; the original must be unmodified."""
+    original = AgentExecutorResponse(
+        executor_id="test_exec",
+        agent_response=AgentResponse(messages=[Message("assistant", ["original reply"])]),
+        full_conversation=[Message("user", ["prompt"]), Message("assistant", ["original reply"])],
+    )
+
+    new = original.with_text("transformed reply")
+
+    assert new is not original
+    assert new.agent_response.text == "transformed reply"
+    assert new.full_conversation[-1].text == "transformed reply"
+    assert new.full_conversation[-1].role == "assistant"
+    # Original unchanged
+    assert original.agent_response.text == "original reply"
+    assert original.full_conversation[-1].text == "original reply"
+
+
+async def test_with_text_strips_multi_message_agent_turn() -> None:
+    """When the agent turn has multiple messages (tool calls), with_text strips all of them."""
+    tool_call = Message("assistant", ["<tool_call>"])
+    tool_result = Message("tool", ["<result>"])
+    final_reply = Message("assistant", ["actual answer"])
+    user_msg = Message("user", ["question"])
+
+    original = AgentExecutorResponse(
+        executor_id="exec",
+        agent_response=AgentResponse(messages=[tool_call, tool_result, final_reply]),
+        full_conversation=[user_msg, tool_call, tool_result, final_reply],
+    )
+
+    new = original.with_text("summarised answer")
+
+    # Only the pre-agent-turn messages should remain, plus the replacement
+    assert len(new.full_conversation) == 2
+    assert new.full_conversation[0].text == "question"
+    assert new.full_conversation[1].text == "summarised answer"
+    assert new.agent_response.text == "summarised answer"

python/packages/core/tests/workflow/test_function_executor.py

@@ -0,0 +1,968 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from dataclasses import dataclass
+from typing import Any, TypeVar
+
+import pytest
+from typing_extensions import Never
+
+from agent_framework import (
+    FunctionExecutor,
+    WorkflowBuilder,
+    WorkflowContext,
+    WorkflowMessage,
+    executor,
+)
+
+
+# Module-level types for string forward reference tests
+@dataclass
+class FuncExecForwardRefMessage:
+    content: str
+
+
+@dataclass
+class FuncExecForwardRefTypeA:
+    value: str
+
+
+@dataclass
+class FuncExecForwardRefTypeB:
+    value: int
+
+
+@dataclass
+class FuncExecForwardRefResponse:
+    result: str
+
+
+class TestFunctionExecutor:
+    """Test suite for FunctionExecutor and @executor decorator."""
+
+    def test_function_executor_basic(self):
+        """Test basic FunctionExecutor creation and validation."""
+
+        async def process_string(text: str, ctx: WorkflowContext[str]) -> None:
+            await ctx.send_message(text.upper())
+
+        func_exec = FunctionExecutor(process_string)
+
+        # Check that handler was registered
+        assert len(func_exec._handlers) == 1  # pyright: ignore[reportPrivateUsage]
+        assert str in func_exec._handlers  # pyright: ignore[reportPrivateUsage]
+
+        # Check handler spec was created
+        assert len(func_exec._handler_specs) == 1  # pyright: ignore[reportPrivateUsage]
+        spec = func_exec._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert spec["name"] == "process_string"
+        assert spec["message_type"] is str
+        assert spec["output_types"] == [str]
+
+    def test_executor_decorator(self):
+        """Test @executor decorator creates proper FunctionExecutor."""
+
+        @executor(id="test_executor")
+        async def process_int(value: int, ctx: WorkflowContext[int]) -> None:
+            await ctx.send_message(value * 2)
+
+        assert isinstance(process_int, FunctionExecutor)
+        assert process_int.id == "test_executor"
+        assert int in process_int._handlers  # pyright: ignore[reportPrivateUsage]
+
+        # Check spec
+        spec = process_int._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert spec["message_type"] is int
+        assert spec["output_types"] == [int]
+
+    def test_executor_decorator_without_id(self):
+        """Test @executor decorator uses function name as default ID."""
+
+        @executor
+        async def my_function(data: dict[str, Any], ctx: WorkflowContext[Any]) -> None:
+            await ctx.send_message(data)
+
+        assert my_function.id == "my_function"
+
+    def test_executor_decorator_without_parentheses(self):
+        """Test @executor decorator works without parentheses."""
+
+        @executor
+        async def no_parens_function(data: str, ctx: WorkflowContext[str]) -> None:
+            await ctx.send_message(data.upper())
+
+        assert isinstance(no_parens_function, FunctionExecutor)
+        assert no_parens_function.id == "no_parens_function"
+        assert str in no_parens_function._handlers  # pyright: ignore[reportPrivateUsage]
+
+        # Also test with single parameter function
+        @executor
+        async def simple_no_parens(value: int):
+            return value * 2
+
+        assert isinstance(simple_no_parens, FunctionExecutor)
+        assert simple_no_parens.id == "simple_no_parens"
+        assert int in simple_no_parens._handlers  # pyright: ignore[reportPrivateUsage]
+
+    def test_union_output_types(self):
+        """Test that union output types are properly inferred for both messages and workflow outputs."""
+
+        @executor
+        async def multi_output(text: str, ctx: WorkflowContext[str | int]) -> None:
+            if text.isdigit():
+                await ctx.send_message(int(text))
+            else:
+                await ctx.send_message(text.upper())
+
+        spec = multi_output._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert set(spec["output_types"]) == {str, int}
+        assert spec["workflow_output_types"] == []  # No workflow outputs defined
+
+        # Test union types for workflow outputs too
+        @executor
+        async def multi_workflow_output(data: str, ctx: WorkflowContext[Never, str | int | bool]) -> None:  # type: ignore[valid-type]
+            if data.isdigit():
+                await ctx.yield_output(int(data))
+            elif data.lower() in ("true", "false"):
+                await ctx.yield_output(data.lower() == "true")
+            else:
+                await ctx.yield_output(data.upper())
+
+        workflow_spec = multi_workflow_output._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert workflow_spec["output_types"] == []  # None means no message outputs
+        assert set(workflow_spec["workflow_output_types"]) == {str, int, bool}
+
+    def test_none_output_type(self):
+        """Test WorkflowContext produces empty output types."""
+
+        @executor
+        async def no_output(data: Any, ctx: WorkflowContext) -> None:
+            # This executor doesn't send any messages
+            pass
+
+        spec = no_output._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert spec["output_types"] == []
+        assert spec["workflow_output_types"] == []  # No workflow outputs defined
+
+    def test_any_output_type(self):
+        """Test WorkflowContext[Any] and WorkflowContext[Any, Any] produce Any output types."""
+
+        @executor
+        async def any_output(data: str, ctx: WorkflowContext[Any]) -> None:
+            await ctx.send_message("result")
+
+        spec = any_output._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert spec["output_types"] == [Any]
+        assert spec["workflow_output_types"] == []  # No workflow outputs defined
+
+        # Test both parameters as Any
+        @executor
+        async def any_both_output(data: str, ctx: WorkflowContext[Any, Any]) -> None:
+            await ctx.send_message("message")
+            await ctx.yield_output("workflow_output")
+
+        both_spec = any_both_output._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert both_spec["output_types"] == [Any]
+        assert both_spec["workflow_output_types"] == [Any]
+
+    def test_validation_errors(self):
+        """Test various validation errors in function signatures."""
+
+        # Wrong number of parameters (now accepts 1 or 2, so 0 or 3+ should fail)
+        async def no_params() -> None:
+            pass
+
+        with pytest.raises(
+            ValueError, match="must have \\(message: T\\) or \\(message: T, ctx: WorkflowContext\\[U\\]\\)"
+        ):
+            FunctionExecutor(no_params)  # type: ignore
+
+        async def too_many_params(data: str, ctx: WorkflowContext[str], extra: int) -> None:
+            pass
+
+        with pytest.raises(
+            ValueError, match="must have \\(message: T\\) or \\(message: T, ctx: WorkflowContext\\[U\\]\\)"
+        ):
+            FunctionExecutor(too_many_params)  # type: ignore
+
+        # Missing message type annotation
+        async def no_msg_type(data, ctx: WorkflowContext[str]) -> None:  # type: ignore
+            pass
+
+        with pytest.raises(ValueError, match="type annotation for the message"):
+            FunctionExecutor(no_msg_type)  # type: ignore
+
+        # Missing ctx annotation (only for 2-parameter functions)
+        async def no_ctx_type(data: str, ctx) -> None:  # type: ignore
+            pass
+
+        with pytest.raises(ValueError, match="must have a WorkflowContext"):
+            FunctionExecutor(no_ctx_type)  # type: ignore
+
+        # Wrong ctx type
+        async def wrong_ctx_type(data: str, ctx: str) -> None:  # type: ignore
+            pass
+
+        with pytest.raises(ValueError, match="must be annotated as WorkflowContext"):
+            FunctionExecutor(wrong_ctx_type)  # type: ignore
+
+        # Unparameterized WorkflowContext is now allowed
+        async def unparameterized_ctx(data: str, ctx: WorkflowContext) -> None:  # type: ignore
+            pass
+
+        # This should now succeed since unparameterized WorkflowContext is allowed
+        executor = FunctionExecutor(unparameterized_ctx)
+        assert executor.output_types == []  # Unparameterized has no inferred types
+        assert executor.workflow_output_types == []  # No workflow output types
+
+    async def test_execution_in_workflow(self):
+        """Test that FunctionExecutor works properly in a workflow."""
+
+        @executor(id="upper")
+        async def to_upper(text: str, ctx: WorkflowContext[str]) -> None:
+            result = text.upper()
+            await ctx.send_message(result)
+
+        @executor(id="reverse")
+        async def reverse_text(text: str, ctx: WorkflowContext[Any, str]) -> None:
+            result = text[::-1]
+            await ctx.yield_output(result)
+
+        # Verify type inference for both executors
+        upper_spec = to_upper._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert upper_spec["output_types"] == [str]
+        assert upper_spec["workflow_output_types"] == []  # No workflow outputs
+
+        reverse_spec = reverse_text._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert reverse_spec["output_types"] == [Any]  # First parameter is Any
+        assert reverse_spec["workflow_output_types"] == [str]  # Second parameter is str
+
+        workflow = WorkflowBuilder(start_executor=to_upper).add_edge(to_upper, reverse_text).build()
+
+        # Run workflow
+        events = await workflow.run("hello world")
+        outputs = events.get_outputs()
+
+        # Assert that we got the expected output
+        assert len(outputs) == 1
+        assert outputs[0] == "DLROW OLLEH"
+
+    def test_can_handle_method(self):
+        """Test that can_handle method works with instance handlers."""
+
+        @executor
+        async def string_processor(text: str, ctx: WorkflowContext[str]) -> None:
+            await ctx.send_message(text)
+
+        assert string_processor.can_handle(WorkflowMessage(data="hello", source_id="Mock"))
+        assert not string_processor.can_handle(WorkflowMessage(data=123, source_id="Mock"))
+        assert not string_processor.can_handle(WorkflowMessage(data=[], source_id="Mock"))
+
+    def test_duplicate_handler_registration(self):
+        """Test that registering duplicate handlers raises an error."""
+
+        async def first_handler(text: str, ctx: WorkflowContext[str]) -> None:
+            await ctx.send_message(text)
+
+        func_exec = FunctionExecutor(first_handler)
+
+        # Try to register another handler for the same type
+        async def second_handler(message: str, ctx: WorkflowContext[str]) -> None:
+            await ctx.send_message(message)
+
+        with pytest.raises(ValueError, match="Handler for type .* already registered"):
+            func_exec._register_instance_handler(  # pyright: ignore[reportPrivateUsage]
+                name="second",
+                func=second_handler,
+                message_type=str,
+                ctx_annotation=WorkflowContext[str],
+                output_types=[str],
+                workflow_output_types=[],
+            )
+
+    def test_complex_type_annotations(self):
+        """Test with complex type annotations like List[str], Dict[str, int], etc."""
+
+        @executor
+        async def process_list(items: list[str], ctx: WorkflowContext[dict[str, int]]) -> None:
+            result = {item: len(item) for item in items}
+            await ctx.send_message(result)
+
+        spec = process_list._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert spec["message_type"] == list[str]
+        assert spec["output_types"] == [dict[str, int]]
+
+    def test_single_parameter_function(self):
+        """Test FunctionExecutor with single-parameter functions."""
+
+        @executor(id="simple_processor")
+        async def process_simple(text: str):
+            return text.upper()
+
+        assert isinstance(process_simple, FunctionExecutor)
+        assert process_simple.id == "simple_processor"
+        assert str in process_simple._handlers  # pyright: ignore[reportPrivateUsage]
+
+        # Check spec - single parameter functions have no output types since they can't send messages
+        spec = process_simple._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert spec["message_type"] is str
+        assert spec["output_types"] == []
+        assert spec["ctx_annotation"] is None
+
+    def test_single_parameter_validation(self):
+        """Test validation for single-parameter functions."""
+
+        # Valid single-parameter function
+        async def valid_single(data: int):
+            return data * 2
+
+        func_exec = FunctionExecutor(valid_single)
+        assert int in func_exec._handlers  # pyright: ignore[reportPrivateUsage]
+
+        # Single parameter with missing type annotation should still fail
+        async def no_annotation(data):  # type: ignore
+            pass
+
+        with pytest.raises(ValueError, match="type annotation for the message"):
+            FunctionExecutor(no_annotation)  # type: ignore
+
+    def test_single_parameter_can_handle(self):
+        """Test that single-parameter functions work with can_handle method."""
+
+        @executor
+        async def int_processor(value: int):
+            return value * 2
+
+        assert int_processor.can_handle(WorkflowMessage(data=42, source_id="mock"))
+        assert not int_processor.can_handle(WorkflowMessage(data="hello", source_id="mock"))
+        assert not int_processor.can_handle(WorkflowMessage(data=[], source_id="mock"))
+
+    async def test_single_parameter_execution(self):
+        """Test that single-parameter functions can be executed properly."""
+
+        @executor(id="double")
+        async def double_value(value: int):
+            return value * 2
+
+        # Since single-parameter functions can't send messages,
+        # they're typically used as terminal nodes or for side effects
+        WorkflowBuilder(start_executor=double_value).build()
+
+        # For testing purposes, we can check that the handler is registered correctly
+        assert double_value.can_handle(WorkflowMessage(data=5, source_id="mock"))
+        assert int in double_value._handlers  # pyright: ignore[reportPrivateUsage]
+
+    def test_sync_function_basic(self):
+        """Test basic synchronous function support."""
+
+        @executor(id="sync_processor")
+        def process_sync(text: str):
+            return text.upper()
+
+        assert isinstance(process_sync, FunctionExecutor)
+        assert process_sync.id == "sync_processor"
+        assert str in process_sync._handlers  # pyright: ignore[reportPrivateUsage]
+
+        # Check spec - sync single parameter functions have no output types
+        spec = process_sync._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert spec["message_type"] is str
+        assert spec["output_types"] == []
+        assert spec["ctx_annotation"] is None
+
+    def test_sync_function_with_context(self):
+        """Test synchronous function with WorkflowContext."""
+
+        @executor
+        def sync_with_ctx(value: int, ctx: WorkflowContext[int]):
+            # Sync functions can still use context
+            return value * 2
+
+        assert isinstance(sync_with_ctx, FunctionExecutor)
+        assert sync_with_ctx.id == "sync_with_ctx"
+        assert int in sync_with_ctx._handlers  # pyright: ignore[reportPrivateUsage]
+
+        # Check spec - sync functions with context can infer output types
+        spec = sync_with_ctx._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert spec["message_type"] is int
+        assert spec["output_types"] == [int]
+
+    def test_sync_function_can_handle(self):
+        """Test that sync functions work with can_handle method."""
+
+        @executor
+        def string_handler(text: str):
+            return text.strip()
+
+        assert string_handler.can_handle(WorkflowMessage(data="hello", source_id="mock"))
+        assert not string_handler.can_handle(WorkflowMessage(data=123, source_id="mock"))
+        assert not string_handler.can_handle(WorkflowMessage(data=[], source_id="mock"))
+
+    def test_sync_function_validation(self):
+        """Test validation for synchronous functions."""
+
+        # Valid sync function with one parameter
+        def valid_sync(data: str):
+            return data.upper()
+
+        func_exec = FunctionExecutor(valid_sync)
+        assert str in func_exec._handlers  # pyright: ignore[reportPrivateUsage]
+
+        # Valid sync function with two parameters
+        def valid_sync_with_ctx(data: int, ctx: WorkflowContext[str]):
+            return str(data)
+
+        func_exec2 = FunctionExecutor(valid_sync_with_ctx)
+        assert int in func_exec2._handlers  # pyright: ignore[reportPrivateUsage]
+
+        # Sync function with missing type annotation should still fail
+        def no_annotation(data):  # pyright: ignore[reportUnknownVariableType]  # type: ignore
+            return data  # pyright: ignore[reportUnknownVariableType]
+
+        with pytest.raises(ValueError, match="type annotation for the message"):
+            FunctionExecutor(no_annotation)  # type: ignore
+
+    def test_mixed_sync_async_decorator(self):
+        """Test that both sync and async functions work with decorator."""
+
+        @executor
+        def sync_func(data: str):
+            return data.lower()
+
+        @executor
+        async def async_func(data: str):
+            return data.upper()
+
+        # Both should be FunctionExecutor instances
+        assert isinstance(sync_func, FunctionExecutor)
+        assert isinstance(async_func, FunctionExecutor)
+
+        # Both should handle strings
+        assert sync_func.can_handle(WorkflowMessage(data="test", source_id="mock"))
+        assert async_func.can_handle(WorkflowMessage(data="test", source_id="mock"))
+
+        # Both should be different instances
+        assert sync_func is not async_func
+
+    async def test_sync_function_in_workflow(self):
+        """Test that sync functions work properly in a workflow context."""
+
+        @executor(id="sync_upper")
+        def to_upper_sync(text: str, ctx: WorkflowContext[str]):
+            return text.upper()
+            # Note: For the test, we'll use a sync send mechanism
+            # In practice, the wrapper handles the async conversion
+
+        @executor(id="async_reverse")
+        async def reverse_async(text: str, ctx: WorkflowContext[Any, str]):
+            result = text[::-1]
+            await ctx.yield_output(result)
+
+        # Verify type inference for sync and async functions
+        sync_spec = to_upper_sync._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert sync_spec["output_types"] == [str]
+        assert sync_spec["workflow_output_types"] == []  # No workflow outputs
+
+        async_spec = reverse_async._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert async_spec["output_types"] == [Any]  # First parameter is Any
+        assert async_spec["workflow_output_types"] == [str]  # Second parameter is str
+
+        # Verify the executors can handle their input types
+        assert to_upper_sync.can_handle(WorkflowMessage(data="hello", source_id="mock"))
+        assert reverse_async.can_handle(WorkflowMessage(data="HELLO", source_id="mock"))
+
+        # For integration testing, we mainly verify that the handlers are properly registered
+        # and the functions are wrapped correctly
+        assert str in to_upper_sync._handlers  # pyright: ignore[reportPrivateUsage]
+        assert str in reverse_async._handlers  # pyright: ignore[reportPrivateUsage]
+
+    async def test_sync_function_thread_execution(self):
+        """Test that sync functions run in thread pool and don't block the event loop."""
+        import threading
+        import time
+
+        _ = threading.get_ident()
+        execution_thread_id = None
+
+        @executor
+        def blocking_function(data: str):
+            nonlocal execution_thread_id
+            execution_thread_id = threading.get_ident()  # type: ignore[assignment]
+            # Simulate some CPU-bound work
+            time.sleep(0.01)  # Small sleep to verify thread execution
+            return data.upper()
+
+        # Verify the function is wrapped and registered
+        assert str in blocking_function._handlers  # pyright: ignore[reportPrivateUsage]
+
+        # For a more complete test, we'd need to create a full workflow context,
+        # but for now we can verify that the function was properly wrapped
+        # and that sync functions store the correct metadata
+        assert not blocking_function._is_async  # pyright: ignore[reportPrivateUsage]
+        assert not blocking_function._has_context  # pyright: ignore[reportPrivateUsage]
+
+        # The actual thread execution test would require a full workflow setup,
+        # but the important thing is that asyncio.to_thread is used in the wrapper
+
+    def test_executor_rejects_staticmethod(self):
+        """Test that @executor decorator properly rejects @staticmethod with clear error."""
+        with pytest.raises(ValueError) as exc_info:
+
+            class Example:  # pyright: ignore[reportUnusedClass]
+                @executor
+                @staticmethod
+                async def bad_handler(data: str) -> str:
+                    return data.upper()
+
+        assert "cannot be used with @staticmethod" in str(exc_info.value)
+        assert "@handler on instance methods" in str(exc_info.value)
+
+    def test_executor_rejects_classmethod(self):
+        """Test that @executor decorator properly rejects @classmethod with clear error."""
+        with pytest.raises(ValueError) as exc_info:
+
+            class Example:  # pyright: ignore[reportUnusedClass]
+                @executor
+                @classmethod
+                async def bad_handler(cls, data: str) -> str:  # type: ignore[operator]
+                    return data.upper()
+
+        assert "cannot be used with @classmethod" in str(exc_info.value)
+        assert "@handler on instance methods" in str(exc_info.value)
+
+    async def test_async_staticmethod_detection_behavior(self):
+        """Document the behavior of inspect.iscoroutinefunction with staticmethod descriptors.
+
+        This test explains why the unwrapping is necessary when decorators are stacked.
+        """
+        import asyncio
+        import inspect
+
+        # When @staticmethod is applied, it creates a descriptor
+        async def my_async_func():
+            await asyncio.sleep(0.001)
+            return "done"
+
+        # Apply staticmethod (what happens with innermost decorator)
+        static_wrapped = staticmethod(my_async_func)
+
+        # Direct check on descriptor object fails (this is the bug)
+        assert not inspect.iscoroutinefunction(static_wrapped)
+        assert isinstance(static_wrapped, staticmethod)
+
+        # But unwrapping __func__ reveals the async function
+        unwrapped = static_wrapped.__func__
+        assert inspect.iscoroutinefunction(unwrapped)
+
+        # When accessed via class attribute, Python's descriptor protocol
+        # automatically unwraps it, so it works:
+        class C:
+            async_static = static_wrapped
+
+        assert inspect.iscoroutinefunction(C.async_static)  # Works via descriptor protocol
+
+
+class TestExecutorExplicitTypes:
+    """Test suite for @executor decorator with explicit input_type and output_type parameters."""
+
+    def test_executor_with_explicit_input_type(self):
+        """Test that explicit input_type takes precedence over introspection."""
+
+        @executor(input=str)
+        async def process(message, ctx: WorkflowContext) -> None:  # type: ignore[no-untyped-def]
+            pass
+
+        # Handler should be registered for str (explicit)
+        assert str in process._handlers  # pyright: ignore[reportPrivateUsage]
+        assert len(process._handlers) == 1  # pyright: ignore[reportPrivateUsage]
+
+        # Can handle str messages
+        assert process.can_handle(WorkflowMessage(data="hello", source_id="mock"))
+        # Cannot handle int messages
+        assert not process.can_handle(WorkflowMessage(data=42, source_id="mock"))
+
+    def test_executor_with_explicit_output_type(self):
+        """Test that explicit output_type takes precedence over introspection."""
+
+        @executor(output=int)
+        async def process(message: str, ctx: WorkflowContext[str]) -> None:
+            pass
+
+        # Handler spec should have int as output type (explicit), not str (introspected)
+        spec = process._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert spec["output_types"] == [int]
+
+        # Executor output_types property should reflect explicit type
+        assert int in process.output_types
+        assert str not in process.output_types
+
+    def test_executor_with_explicit_input_and_output_types(self):
+        """Test that both explicit input_type and output_type work together."""
+
+        @executor(id="explicit_both", input=dict, output=list)
+        async def process(message, ctx: WorkflowContext) -> None:  # type: ignore[no-untyped-def]
+            pass
+
+        # Handler should be registered for dict (explicit input type)
+        assert dict in process._handlers  # pyright: ignore[reportPrivateUsage]
+        assert len(process._handlers) == 1  # pyright: ignore[reportPrivateUsage]
+
+        # Output type should be list (explicit)
+        spec = process._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert spec["output_types"] == [list]
+
+        # Verify can_handle
+        assert process.can_handle(WorkflowMessage(data={"key": "value"}, source_id="mock"))
+        assert not process.can_handle(WorkflowMessage(data="string", source_id="mock"))
+
+    def test_executor_with_explicit_union_input_type(self):
+        """Test that explicit union input_type is handled correctly."""
+
+        @executor(input=str | int)
+        async def process(message, ctx: WorkflowContext) -> None:  # type: ignore[no-untyped-def]
+            pass
+
+        # Handler should be registered for the union type
+        assert len(process._handlers) == 1  # pyright: ignore[reportPrivateUsage]
+
+        # Can handle both str and int messages
+        assert process.can_handle(WorkflowMessage(data="hello", source_id="mock"))
+        assert process.can_handle(WorkflowMessage(data=42, source_id="mock"))
+        # Cannot handle float
+        assert not process.can_handle(WorkflowMessage(data=3.14, source_id="mock"))
+
+    def test_executor_with_explicit_union_output_type(self):
+        """Test that explicit union output_type is normalized to a list."""
+
+        @executor(output=str | int | bool)
+        async def process(message: Any, ctx: WorkflowContext) -> None:
+            pass
+
+        # Output types should be a list with all union members
+        assert set(process.output_types) == {str, int, bool}
+
+    def test_executor_explicit_types_precedence_over_introspection(self):
+        """Test that explicit types always take precedence over introspected types."""
+
+        # Introspection would give: input=str, output=[int]
+        # Explicit gives: input=bytes, output=[float]
+        @executor(input=bytes, output=float)
+        async def process(message: str, ctx: WorkflowContext[int]) -> None:
+            pass
+
+        # Should use explicit input type (bytes), not introspected (str)
+        assert bytes in process._handlers  # pyright: ignore[reportPrivateUsage]
+        assert str not in process._handlers  # pyright: ignore[reportPrivateUsage]
+
+        # Should use explicit output type (float), not introspected (int)
+        assert float in process.output_types
+        assert int not in process.output_types
+
+    def test_executor_fallback_to_introspection_when_no_explicit_types(self):
+        """Test that introspection is used when no explicit types are provided."""
+
+        @executor
+        async def process(message: str, ctx: WorkflowContext[int]) -> None:
+            pass
+
+        # Should use introspected types
+        assert str in process._handlers  # pyright: ignore[reportPrivateUsage]
+        assert int in process.output_types
+
+    def test_executor_partial_explicit_types(self):
+        """Test that partial explicit types work (only input_type or only output_type)."""
+
+        # Only explicit input_type, introspect output_type
+        @executor(input=bytes)
+        async def process_input(message: str, ctx: WorkflowContext[int]) -> None:
+            pass
+
+        assert bytes in process_input._handlers  # pyright: ignore[reportPrivateUsage]  # Explicit
+        assert int in process_input.output_types  # Introspected
+
+        # Only explicit output_type, introspect input_type
+        @executor(output=float)
+        async def process_output(message: str, ctx: WorkflowContext[int]) -> None:
+            pass
+
+        assert str in process_output._handlers  # pyright: ignore[reportPrivateUsage]  # Introspected
+        assert float in process_output.output_types  # Explicit
+        assert int not in process_output.output_types  # Not introspected when explicit provided
+
+    def test_executor_explicit_input_type_allows_no_message_annotation(self):
+        """Test that explicit input_type allows function without message type annotation."""
+
+        @executor(input=str)
+        async def process(message, ctx: WorkflowContext) -> None:  # type: ignore[no-untyped-def]
+            pass
+
+        # Should work with explicit input_type
+        assert str in process._handlers  # pyright: ignore[reportPrivateUsage]
+        assert process.can_handle(WorkflowMessage(data="hello", source_id="mock"))
+
+    def test_executor_explicit_types_with_id(self):
+        """Test that explicit types work together with id parameter."""
+
+        @executor(id="custom_id", input=bytes, output=int)
+        async def process(message, ctx: WorkflowContext) -> None:  # type: ignore[no-untyped-def]
+            pass
+
+        assert process.id == "custom_id"
+        assert bytes in process._handlers  # pyright: ignore[reportPrivateUsage]
+        assert int in process.output_types
+
+    def test_executor_explicit_types_with_single_param_function(self):
+        """Test that explicit input_type works with single-parameter functions."""
+
+        @executor(input=str)
+        async def process(message):  # type: ignore[no-untyped-def]
+            return message.upper()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+
+        # Should work with explicit input_type
+        assert str in process._handlers  # pyright: ignore[reportPrivateUsage]
+        assert process.can_handle(WorkflowMessage(data="hello", source_id="mock"))
+        assert not process.can_handle(WorkflowMessage(data=42, source_id="mock"))
+
+    def test_executor_explicit_types_with_sync_function(self):
+        """Test that explicit types work with synchronous functions."""
+
+        @executor(input=int, output=str)
+        def process(message, ctx: WorkflowContext) -> None:  # type: ignore[no-untyped-def]
+            pass
+
+        assert int in process._handlers  # pyright: ignore[reportPrivateUsage]
+        assert str in process.output_types
+
+    def test_function_executor_constructor_with_explicit_types(self):
+        """Test FunctionExecutor constructor with explicit input_type and output_type."""
+
+        async def process(message, ctx: WorkflowContext) -> None:  # type: ignore[no-untyped-def]
+            pass
+
+        func_exec = FunctionExecutor(process, id="test", input=dict, output=list)  # pyright: ignore[reportUnknownArgumentType]
+
+        assert dict in func_exec._handlers  # pyright: ignore[reportPrivateUsage]
+        spec = func_exec._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert spec["message_type"] is dict
+        assert spec["output_types"] == [list]
+
+    def test_executor_explicit_union_types_via_typing_union(self):
+        """Test that Union[] syntax also works for explicit types."""
+        from typing import Union
+
+        @executor(input=Union[str, int], output=Union[bool, float])  # type: ignore[call-overload]
+        async def process(message, ctx: WorkflowContext) -> None:  # type: ignore[no-untyped-def]
+            pass
+
+        # Can handle both str and int
+        assert process.can_handle(WorkflowMessage(data="hello", source_id="mock"))
+        assert process.can_handle(WorkflowMessage(data=42, source_id="mock"))
+
+        # Output types should include both
+        assert set(process.output_types) == {bool, float}
+
+    def test_executor_with_string_forward_reference_input_type(self):
+        """Test that string forward references work for input_type."""
+
+        @executor(input="FuncExecForwardRefMessage")
+        async def process(message, ctx: WorkflowContext) -> None:  # type: ignore[no-untyped-def]
+            pass
+
+        # Should resolve the string to the actual type
+        assert FuncExecForwardRefMessage in process._handlers  # pyright: ignore[reportPrivateUsage]
+        assert process.can_handle(WorkflowMessage(data=FuncExecForwardRefMessage("hello"), source_id="mock"))
+
+    def test_executor_with_string_forward_reference_union(self):
+        """Test that string forward references work with union types."""
+
+        @executor(input="FuncExecForwardRefTypeA | FuncExecForwardRefTypeB")
+        async def process(message, ctx: WorkflowContext) -> None:  # type: ignore[no-untyped-def]
+            pass
+
+        # Should handle both types
+        assert process.can_handle(WorkflowMessage(data=FuncExecForwardRefTypeA("hello"), source_id="mock"))
+        assert process.can_handle(WorkflowMessage(data=FuncExecForwardRefTypeB(42), source_id="mock"))
+
+    def test_executor_with_string_forward_reference_output_type(self):
+        """Test that string forward references work for output_type."""
+
+        @executor(input=str, output="FuncExecForwardRefResponse")
+        async def process(message, ctx: WorkflowContext) -> None:  # type: ignore[no-untyped-def]
+            pass
+
+        # Should resolve the string output type
+        assert FuncExecForwardRefResponse in process.output_types
+
+    def test_executor_with_explicit_workflow_output_type(self):
+        """Test that explicit workflow_output_type takes precedence over introspection."""
+
+        @executor(workflow_output=bool)
+        async def process(message: str, ctx: WorkflowContext[int]) -> None:
+            pass
+
+        # Handler spec should have bool as workflow_output_type (explicit)
+        spec = process._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert spec["workflow_output_types"] == [bool]
+
+        # Executor workflow_output_types property should reflect explicit type
+        assert bool in process.workflow_output_types
+        # output_types should still come from introspection (int from WorkflowContext[int])
+        assert int in process.output_types
+
+    def test_executor_with_explicit_workflow_output_type_precedence(self):
+        """Test that explicit workflow_output_type overrides introspected WorkflowContext second param."""
+
+        @executor(workflow_output=str)
+        async def process(message: int, ctx: WorkflowContext[int, bool]) -> None:
+            pass
+
+        # workflow_output_types should be str (explicit), not bool (introspected from ctx)
+        assert str in process.workflow_output_types
+        assert bool not in process.workflow_output_types
+
+    def test_executor_with_all_explicit_types(self):
+        """Test that all three explicit type parameters work together."""
+        from typing import Any
+
+        @executor(input=str, output=int, workflow_output=bool)
+        async def process(message: Any, ctx: WorkflowContext) -> None:
+            pass
+
+        # Check input type
+        assert str in process._handlers  # pyright: ignore[reportPrivateUsage]
+        assert process.can_handle(WorkflowMessage(data="hello", source_id="mock"))
+
+        # Check output_type
+        assert int in process.output_types
+
+        # Check workflow_output_type
+        assert bool in process.workflow_output_types
+
+    def test_executor_with_union_workflow_output_type(self):
+        """Test that union types work for workflow_output_type."""
+
+        @executor(workflow_output=str | int)
+        async def process(message: str, ctx: WorkflowContext) -> None:
+            pass
+
+        # Should include both types from union
+        assert str in process.workflow_output_types
+        assert int in process.workflow_output_types
+
+    def test_executor_with_string_forward_reference_workflow_output_type(self):
+        """Test that string forward references work for workflow_output_type."""
+
+        @executor(input=str, workflow_output="FuncExecForwardRefResponse")
+        async def process(message, ctx: WorkflowContext) -> None:  # type: ignore[no-untyped-def]
+            pass
+
+        # Should resolve the string workflow_output_type
+        assert FuncExecForwardRefResponse in process.workflow_output_types
+
+    def test_executor_with_string_forward_reference_union_workflow_output_type(self):
+        """Test that string forward reference union types work for workflow_output_type."""
+
+        @executor(input=str, workflow_output="FuncExecForwardRefTypeA | FuncExecForwardRefTypeB")
+        async def process(message, ctx: WorkflowContext) -> None:  # type: ignore[no-untyped-def]
+            pass
+
+        # Should resolve both types from string union
+        assert FuncExecForwardRefTypeA in process.workflow_output_types
+        assert FuncExecForwardRefTypeB in process.workflow_output_types
+
+    def test_executor_fallback_to_introspection_for_workflow_output_type(self):
+        """Test that workflow_output_type falls back to introspection when not explicitly provided."""
+
+        @executor
+        async def process(message: str, ctx: WorkflowContext[int, bool]) -> None:
+            pass
+
+        # Should use introspected types from WorkflowContext[int, bool]
+        assert int in process.output_types
+        assert bool in process.workflow_output_types
+
+    def test_function_executor_constructor_with_workflow_output_type(self):
+        """Test FunctionExecutor constructor accepts workflow_output_type parameter."""
+
+        async def my_func(message: str, ctx: WorkflowContext) -> None:
+            pass
+
+        exec_instance = FunctionExecutor(
+            my_func,
+            id="test_constructor",
+            input=str,
+            output=int,
+            workflow_output=bool,
+        )
+
+        assert str in exec_instance._handlers  # pyright: ignore[reportPrivateUsage]
+        assert int in exec_instance.output_types
+        assert bool in exec_instance.workflow_output_types
+
+
+# region Tests for unresolved TypeVar rejection in function executor registration
+
+_FT = TypeVar("_FT")
+
+
+class TestFunctionExecutorTypeVarRejection:
+    """Tests that FunctionExecutor rejects unresolved TypeVar in message annotations."""
+
+    def test_function_executor_rejects_unresolved_typevar(self):
+        """Test that FunctionExecutor raises ValueError for unresolved TypeVar message annotation."""
+
+        def echo(message: _FT) -> _FT:
+            return message
+
+        with pytest.raises(ValueError, match="unresolved TypeVar"):
+            FunctionExecutor(echo, id="echo")
+
+    def test_function_executor_rejects_typevar_with_context(self):
+        """Test that FunctionExecutor raises ValueError for TypeVar even with WorkflowContext."""
+
+        async def echo(message: _FT, ctx: WorkflowContext) -> None:
+            pass
+
+        with pytest.raises(ValueError, match="unresolved TypeVar"):
+            FunctionExecutor(echo, id="echo")
+
+    def test_function_executor_explicit_input_bypasses_typevar_check(self):
+        """Test that explicit input= parameter bypasses TypeVar detection."""
+
+        async def echo(message: _FT, ctx: WorkflowContext) -> None:
+            pass
+
+        exec_instance = FunctionExecutor(echo, id="echo", input=str, output=str)
+        assert str in exec_instance.input_types
+
+    def test_function_executor_allows_concrete_types(self):
+        """Test that FunctionExecutor works normally with concrete type annotations."""
+
+        async def handle(message: str, ctx: WorkflowContext[str]) -> None:
+            pass
+
+        exec_instance = FunctionExecutor(handle, id="concrete")
+        assert str in exec_instance.input_types
+
+    def test_function_executor_error_recommends_explicit_types(self):
+        """Test that error message recommends @executor(input=..., output=...)."""
+
+        def echo(message: _FT) -> _FT:
+            return message
+
+        with pytest.raises(ValueError, match=r"@executor\(input=<concrete_type>, output=<concrete_type>\)"):
+            FunctionExecutor(echo, id="echo")
+
+
+# endregion: Tests for unresolved TypeVar rejection in function executor registration
+
+
+_FBT = TypeVar("_FBT", bound=str)
+
+
+def test_function_executor_rejects_bounded_typevar_in_message_annotation():
+    """Test that FunctionExecutor raises ValueError for a bounded TypeVar in message annotation."""
+
+    async def process(message: _FBT, ctx: WorkflowContext) -> None:
+        await ctx.send_message(message)  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+
+    with pytest.raises(ValueError, match="unresolved TypeVar"):
+        FunctionExecutor(process, id="bounded")

python/packages/core/tests/workflow/test_function_executor_future.py

@@ -0,0 +1,58 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from typing import Any
+
+import pytest
+
+from agent_framework import FunctionExecutor, WorkflowContext, executor
+
+
+class TestFunctionExecutorFutureAnnotations:
+    """Test suite for FunctionExecutor with from __future__ import annotations."""
+
+    def test_executor_decorator_future_annotations(self):
+        """Test @executor decorator works with stringified annotations."""
+
+        @executor(id="future_test")
+        async def process_future(value: int, ctx: WorkflowContext[int]) -> None:
+            await ctx.send_message(value * 2)
+
+        assert isinstance(process_future, FunctionExecutor)
+        assert process_future.id == "future_test"
+        assert int in process_future._handlers  # pyright: ignore[reportPrivateUsage]
+
+        # Check spec
+        spec = process_future._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert spec["message_type"] is int
+        assert spec["output_types"] == [int]
+
+    def test_executor_decorator_future_annotations_complex(self):
+        """Test @executor decorator works with complex stringified annotations."""
+
+        @executor
+        async def process_complex(data: dict[str, Any], ctx: WorkflowContext[list[str]]) -> None:
+            await ctx.send_message(["done"])
+
+        assert isinstance(process_complex, FunctionExecutor)
+        spec = process_complex._handler_specs[0]  # pyright: ignore[reportPrivateUsage]
+        assert spec["message_type"] == dict[str, Any]
+        assert spec["output_types"] == [list[str]]
+
+    def test_handler_unresolvable_annotation_raises(self):
+        """Test that an unresolvable forward-reference annotation raises ValueError.
+
+        When get_type_hints fails (e.g. NameError for NonExistentType), the code falls back
+        to raw string annotations. The ctx parameter's raw string annotation is then not
+        recognised as a valid WorkflowContext type, so a ValueError is still raised.
+        """
+        with pytest.raises(ValueError):
+            FunctionExecutor(
+                _func_with_bad_annotation,  # pyright: ignore[reportUnknownArgumentType]
+                id="bad",
+            )
+
+
+async def _func_with_bad_annotation(message: NonExistentType, ctx: WorkflowContext[int]) -> None:  # type: ignore[name-defined]  # ty: ignore[unresolved-reference]  # noqa: F821
+    pass

python/packages/core/tests/workflow/test_output_designation.py

@@ -0,0 +1,137 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for the ``OutputDesignation`` value type and the ``Workflow.is_terminal_executor``
+public predicate that delegates to it.
+
+The states the value type encodes:
+- Omitted-selection compatibility: ``outputs=None`` -> every executor is terminal.
+- Explicit: disjoint ``outputs`` and ``intermediates`` sets classify listed executors,
+  and hide unlisted executors.
+"""
+
+from __future__ import annotations
+
+import pytest
+from typing_extensions import Never
+
+from agent_framework import (
+    Message,
+    WorkflowBuilder,
+    WorkflowContext,
+    WorkflowValidationError,
+    executor,
+)
+from agent_framework._workflows._runner_context import InProcRunnerContext
+from agent_framework._workflows._workflow import OutputDesignation, Workflow
+
+# ---------------------------------------------------------------------------
+# OutputDesignation value type
+# ---------------------------------------------------------------------------
+
+
+def test_omitted_selection_designation_marks_every_executor_as_terminal() -> None:
+    designation = OutputDesignation()  # designated defaults to None
+    assert designation.outputs is None
+    assert designation.is_terminal("anything")
+    assert designation.is_terminal("else")
+    assert designation.classify("anything") == "output"
+
+
+def test_strict_empty_designation_marks_no_executor_as_terminal() -> None:
+    designation = OutputDesignation(outputs=frozenset())
+    assert designation.outputs == frozenset()
+    assert not designation.is_terminal("anything")
+    assert not designation.is_terminal("else")
+    assert designation.classify("anything") is None
+
+
+def test_strict_designated_set_only_terminal_for_members() -> None:
+    designation = OutputDesignation(outputs=frozenset({"alpha", "beta"}), intermediates=frozenset({"gamma"}))
+    assert designation.is_terminal("alpha")
+    assert designation.is_terminal("beta")
+    assert not designation.is_terminal("gamma")
+    assert designation.is_intermediate("gamma")
+    assert designation.classify("alpha") == "output"
+    assert designation.classify("gamma") == "intermediate"
+    assert designation.classify("delta") is None
+
+
+def test_designation_is_frozen() -> None:
+    from dataclasses import FrozenInstanceError
+
+    designation = OutputDesignation(outputs=frozenset({"alpha"}))
+    with pytest.raises(FrozenInstanceError):
+        designation.outputs = frozenset({"beta"})  # type: ignore[misc]  # ty: ignore[invalid-assignment]
+
+
+# ---------------------------------------------------------------------------
+# Workflow.is_terminal_executor delegates to the designation
+# ---------------------------------------------------------------------------
+
+
+@executor
+async def _emit_one(messages: list[Message], ctx: WorkflowContext[Never, str]) -> None:  # type: ignore[valid-type]
+    await ctx.yield_output("hello")
+
+
+@executor
+async def _downstream(message: str, ctx: WorkflowContext[Never, str]) -> None:  # type: ignore[valid-type]
+    await ctx.yield_output("downstream")
+
+
+def test_is_terminal_executor_omitted_selection_returns_true_for_any_id() -> None:
+    """Omitted-selection compatibility behavior: every executor is terminal."""
+    import warnings
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", DeprecationWarning)
+        workflow = WorkflowBuilder(start_executor=_emit_one).build()
+    assert workflow.is_terminal_executor(_emit_one.id)
+    assert workflow.is_terminal_executor("anything-else")
+
+
+def test_is_intermediate_executor_explicit_list_returns_true_only_for_designated() -> None:
+    """Explicit mode tracks intermediate-designated executors separately."""
+    workflow = WorkflowBuilder(start_executor=_emit_one, intermediate_output_from=[_emit_one]).build()
+    assert not workflow.is_terminal_executor(_emit_one.id)
+    assert not workflow.is_terminal_executor("nope")
+    assert workflow.is_intermediate_executor(_emit_one.id)
+    assert not workflow.is_intermediate_executor("nope")
+
+
+def test_is_terminal_executor_strict_list_returns_true_only_for_designated() -> None:
+    """Strict mode with a designated list: only listed executors are terminal."""
+    workflow = (
+        WorkflowBuilder(start_executor=_emit_one, output_from=[_emit_one]).add_edge(_emit_one, _downstream).build()
+    )
+    assert workflow.is_terminal_executor(_emit_one.id)
+    assert not workflow.is_terminal_executor(_downstream.id)
+
+
+def test_get_output_executors_throws_when_designation_references_missing_executor() -> None:
+    workflow = Workflow(
+        [],
+        {_emit_one.id: _emit_one},
+        _emit_one,
+        InProcRunnerContext(),
+        "test",
+        output_from=["missing"],
+    )
+
+    with pytest.raises(WorkflowValidationError, match="Output executor 'missing' is not present"):
+        workflow.get_output_executors()
+
+
+def test_get_intermediate_executors_throws_when_designation_references_missing_executor() -> None:
+    workflow = Workflow(
+        [],
+        {_emit_one.id: _emit_one},
+        _emit_one,
+        InProcRunnerContext(),
+        "test",
+        output_from=[],
+        intermediate_output_from=["missing"],
+    )
+
+    with pytest.raises(WorkflowValidationError, match="Intermediate executor 'missing' is not present"):
+        workflow.get_intermediate_executors()

python/packages/core/tests/workflow/test_output_executors_contract.py

@@ -0,0 +1,287 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for the explicit output/intermediate selection contract on WorkflowBuilder."""
+
+from __future__ import annotations
+
+import warnings
+from typing import Any
+
+import pytest
+from typing_extensions import Never
+
+from agent_framework import (
+    Message,
+    WorkflowBuilder,
+    WorkflowContext,
+    WorkflowValidationError,
+    executor,
+)
+
+
+@executor
+async def _emit_one(messages: list[Message], ctx: WorkflowContext[Never, str]) -> None:  # type: ignore[valid-type]
+    await ctx.yield_output("hello")
+
+
+@executor
+async def _start(messages: list[Message], ctx: WorkflowContext[str, str]) -> None:
+    await ctx.yield_output("from-start")
+    await ctx.send_message("downstream")
+
+
+@executor
+async def _downstream(message: str, ctx: WorkflowContext[Never, str]) -> None:  # type: ignore[valid-type]
+    await ctx.yield_output("from-downstream")
+
+
+def test_designation_unset_emits_deprecation_warning() -> None:
+    """State A: WorkflowBuilder built without explicit designation warns."""
+    with pytest.warns(DeprecationWarning, match="output_from or intermediate_output_from") as warning_info:
+        WorkflowBuilder(start_executor=_emit_one).build()
+    assert str(warning_info[0].message) == (
+        "WorkflowBuilder built without explicit output_from or intermediate_output_from; "
+        "every yield_output produces type='output' for compatibility. Pass output_from='all', "
+        "output_from=[...], or intermediate_output_from=[...] to opt into explicit designation - "
+        "explicit designation will be required in a future version."
+    )
+
+
+@pytest.mark.asyncio
+async def test_designation_unset_preserves_compatibility_all_output_behavior() -> None:
+    """Omitted designation keeps compatibility all-output behavior while warning."""
+    with pytest.warns(DeprecationWarning, match="output_from or intermediate_output_from"):
+        workflow = WorkflowBuilder(start_executor=_start).add_edge(_start, _downstream).build()
+
+    result = await workflow.run([Message(role="user", contents=["hi"])])
+
+    assert result.get_outputs() == ["from-start", "from-downstream"]
+    assert result.get_intermediate_outputs() == []
+
+
+@pytest.mark.asyncio
+async def test_output_from_all_emits_all_outputs_without_omitted_selection_warning() -> None:
+    """Explicit all-output designation emits every executor payload without omitted-selection warning."""
+    with warnings.catch_warnings():
+        warnings.simplefilter("error", DeprecationWarning)
+        workflow = WorkflowBuilder(start_executor=_start, output_from="all").add_edge(_start, _downstream).build()
+
+    result = await workflow.run([Message(role="user", contents=["hi"])])
+
+    assert result.get_outputs() == ["from-start", "from-downstream"]
+    assert result.get_intermediate_outputs() == []
+
+
+@pytest.mark.asyncio
+async def test_output_from_all_with_empty_intermediate_list_is_valid() -> None:
+    """Explicit all-output plus an empty intermediate list is a concrete no-intermediate selection."""
+    with warnings.catch_warnings():
+        warnings.simplefilter("error", DeprecationWarning)
+        workflow = (
+            WorkflowBuilder(start_executor=_start, output_from="all", intermediate_output_from=[])
+            .add_edge(_start, _downstream)
+            .build()
+        )
+
+    result = await workflow.run([Message(role="user", contents=["hi"])])
+
+    assert result.get_outputs() == ["from-start", "from-downstream"]
+    assert result.get_intermediate_outputs() == []
+
+
+@pytest.mark.asyncio
+async def test_intermediate_output_from_all_other_marks_non_outputs_as_intermediate() -> None:
+    """All-other intermediate designation classifies every non-output executor yield as intermediate."""
+    workflow = (
+        WorkflowBuilder(
+            start_executor=_start,
+            output_from=[_downstream],
+            intermediate_output_from="all_other",
+        )
+        .add_edge(_start, _downstream)
+        .build()
+    )
+
+    result = await workflow.run([Message(role="user", contents=["hi"])])
+
+    assert result.get_outputs() == ["from-downstream"]
+    assert result.get_intermediate_outputs() == ["from-start"]
+
+
+@pytest.mark.asyncio
+async def test_all_other_streaming_events_mark_non_outputs_as_intermediate() -> None:
+    """All-other emits intermediate events while streaming, not just in collected results."""
+    workflow = (
+        WorkflowBuilder(
+            start_executor=_start,
+            output_from=[_downstream],
+            intermediate_output_from="all_other",
+        )
+        .add_edge(_start, _downstream)
+        .build()
+    )
+    outputs: list[str] = []
+    intermediates: list[str] = []
+
+    async for event in workflow.run([Message(role="user", contents=["hi"])], stream=True):
+        if event.type == "output":
+            outputs.append(event.data)
+        elif event.type == "intermediate":
+            intermediates.append(event.data)
+
+    assert outputs == ["from-downstream"]
+    assert intermediates == ["from-start"]
+
+
+def test_all_other_expands_to_concrete_intermediate_executor_selection_at_build_time() -> None:
+    """The runner receives concrete executor IDs after all-other expansion."""
+    workflow = (
+        WorkflowBuilder(
+            start_executor=_start,
+            output_from=[_downstream],
+            intermediate_output_from="all_other",
+        )
+        .add_edge(_start, _downstream)
+        .build()
+    )
+
+    assert {executor.id for executor in workflow.get_output_executors()} == {_downstream.id}
+    assert {executor.id for executor in workflow.get_intermediate_executors()} == {_start.id}
+    assert workflow.is_intermediate_executor(_start.id)
+    assert not workflow.is_intermediate_executor(_downstream.id)
+
+
+@pytest.mark.asyncio
+async def test_all_other_with_omitted_output_from_emits_only_intermediate_outputs() -> None:
+    """All-other intermediate designation opts out of omitted-selection all-output behavior."""
+    workflow = (
+        WorkflowBuilder(
+            start_executor=_start,
+            intermediate_output_from="all_other",
+        )
+        .add_edge(_start, _downstream)
+        .build()
+    )
+
+    result = await workflow.run([Message(role="user", contents=["hi"])])
+
+    assert result.get_outputs() == []
+    assert result.get_intermediate_outputs() == ["from-start", "from-downstream"]
+
+
+@pytest.mark.asyncio
+async def test_all_other_with_empty_output_from_emits_only_intermediate_outputs() -> None:
+    """All-other intermediate designation treats an empty output list as selecting no workflow outputs."""
+    workflow = (
+        WorkflowBuilder(
+            start_executor=_start,
+            output_from=[],
+            intermediate_output_from="all_other",
+        )
+        .add_edge(_start, _downstream)
+        .build()
+    )
+
+    result = await workflow.run([Message(role="user", contents=["hi"])])
+
+    assert result.get_outputs() == []
+    assert result.get_intermediate_outputs() == ["from-start", "from-downstream"]
+
+
+@pytest.mark.asyncio
+async def test_all_other_with_output_from_all_expands_to_empty_intermediate_selection() -> None:
+    """All-other is empty when every output-capable executor is already selected as workflow output."""
+    workflow = (
+        WorkflowBuilder(
+            start_executor=_start,
+            output_from="all",
+            intermediate_output_from="all_other",
+        )
+        .add_edge(_start, _downstream)
+        .build()
+    )
+
+    result = await workflow.run([Message(role="user", contents=["hi"])])
+
+    assert result.get_outputs() == ["from-start", "from-downstream"]
+    assert result.get_intermediate_outputs() == []
+
+
+@pytest.mark.asyncio
+async def test_intermediate_output_from_all_routes_every_yield_to_intermediate() -> None:
+    """``intermediate_output_from="all"`` designates every output-capable executor as intermediate."""
+    workflow = (
+        WorkflowBuilder(start_executor=_start, intermediate_output_from="all").add_edge(_start, _downstream).build()
+    )
+
+    result = await workflow.run([Message(role="user", contents=["hi"])])
+
+    assert result.get_outputs() == []
+    assert result.get_intermediate_outputs() == ["from-start", "from-downstream"]
+
+
+def test_output_from_all_other_is_rejected() -> None:
+    """The all-other literal is only valid for intermediate output selection."""
+    with pytest.raises(ValueError, match="output_from.*all_other"):
+        WorkflowBuilder(start_executor=_emit_one, output_from="all_other")  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+
+
+@pytest.mark.parametrize(
+    ("output_from", "intermediate_output_from"),
+    [([_emit_one], None), (None, [_emit_one]), ([], [_emit_one])],
+    ids=["output_list", "intermediate_list", "empty_output_with_intermediate"],
+)
+def test_explicit_designation_with_executor_does_not_warn(output_from, intermediate_output_from) -> None:
+    """State B: any explicit designation with at least one executor opts into explicit mode without warning."""
+    with warnings.catch_warnings():
+        warnings.simplefilter("error", DeprecationWarning)
+        WorkflowBuilder(
+            start_executor=_emit_one,
+            output_from=output_from,
+            intermediate_output_from=intermediate_output_from,
+        ).build()
+
+
+@pytest.mark.parametrize(
+    ("output_from", "intermediate_output_from"),
+    [([], None), (None, []), ([], [])],
+    ids=["empty_output", "empty_intermediate", "both_empty"],
+)
+def test_empty_explicit_designation_fails(output_from, intermediate_output_from) -> None:
+    """State C: explicit mode needs at least one output or intermediate executor."""
+    with pytest.raises(WorkflowValidationError, match="at least one output or intermediate executor"):
+        WorkflowBuilder(
+            start_executor=_emit_one,
+            output_from=output_from,
+            intermediate_output_from=intermediate_output_from,
+        ).build()
+
+
+def test_passing_both_output_executors_and_output_from_raises_type_error() -> None:
+    """State D: supplying a deprecated alias and the canonical kwarg is unambiguous user error."""
+    with pytest.raises(TypeError, match="Cannot pass multiple workflow output selection parameters"):
+        WorkflowBuilder(
+            start_executor=_emit_one,
+            output_executors=[_emit_one],
+            output_from=[_emit_one],
+        )
+
+
+def test_intermediate_executors_builder_parameter_is_not_public() -> None:
+    """The branch-only intermediate_executors builder parameter is not supported."""
+    builder_type: Any = WorkflowBuilder
+    with pytest.raises(TypeError, match="unexpected keyword argument 'intermediate_executors'"):
+        builder_type(
+            start_executor=_emit_one,
+            intermediate_executors=[_emit_one],
+        )
+
+
+def test_final_output_from_builder_parameter_is_not_public() -> None:
+    """The branch-only final_output_from builder parameter is not supported."""
+    builder_type: Any = WorkflowBuilder
+    with pytest.raises(TypeError, match="unexpected keyword argument 'final_output_from'"):
+        builder_type(
+            start_executor=_emit_one,
+            final_output_from=[_emit_one],
+        )

python/packages/core/tests/workflow/test_request_info_and_response.py

@@ -0,0 +1,324 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from dataclasses import dataclass
+
+from agent_framework import (
+    WorkflowBuilder,
+    WorkflowContext,
+    WorkflowEvent,
+    WorkflowRunState,
+    handler,
+    response_handler,
+)
+from agent_framework._workflows._executor import Executor
+from agent_framework._workflows._request_info_mixin import RequestInfoMixin
+
+
+@dataclass
+class UserApprovalRequest:
+    """A request for user approval with context."""
+
+    prompt: str
+    context: str
+    request_id: str = ""
+
+    def __post_init__(self):
+        if not self.request_id:
+            import uuid
+
+            self.request_id = str(uuid.uuid4())
+
+
+@dataclass
+class CalculationRequest:
+    """A request for a complex calculation."""
+
+    operation: str
+    operands: list[float]
+    request_id: str = ""
+
+    def __post_init__(self):
+        if not self.request_id:
+            import uuid
+
+            self.request_id = str(uuid.uuid4())
+
+
+class ApprovalRequiredExecutor(Executor, RequestInfoMixin):
+    """Executor that requires approval before proceeding."""
+
+    def __init__(self, id: str):
+        super().__init__(id=id)
+        self.approval_received = False
+        self.final_result = None
+
+    @handler
+    async def start_process(self, message: str, ctx: WorkflowContext) -> None:
+        """Start a process that requires approval."""
+        # Request approval from external system
+        approval_request = UserApprovalRequest(
+            prompt=f"Please approve the operation: {message}",
+            context="This is a critical operation that requires human approval.",
+        )
+        await ctx.request_info(approval_request, bool)
+
+    @response_handler
+    async def handle_approval_response(
+        self, original_request: UserApprovalRequest, approved: bool, ctx: WorkflowContext[str]
+    ) -> None:
+        """Handle the approval response."""
+        self.approval_received = True
+
+        if approved:
+            self.final_result = f"Operation approved: {original_request.prompt}"  # type: ignore[assignment]
+            await ctx.send_message(f"APPROVED: {original_request.context}")
+        else:
+            self.final_result = "Operation denied by user"  # type: ignore[assignment]
+            await ctx.send_message("DENIED: Operation was not approved")
+
+
+class CalculationExecutor(Executor, RequestInfoMixin):
+    """Executor that delegates complex calculations to external services."""
+
+    def __init__(self, id: str):
+        super().__init__(id=id)
+        self.calculations_performed: list[tuple[str, list[float], float]] = []
+
+    @handler
+    async def process_calculation(self, message: str, ctx: WorkflowContext[str]) -> None:
+        """Process a calculation request."""
+        # Parse the message to extract operation
+        parts = message.split()
+        if len(parts) >= 3:
+            operation = parts[0]
+            try:
+                operands = [float(x) for x in parts[1:]]
+                calc_request = CalculationRequest(operation=operation, operands=operands)
+                await ctx.request_info(calc_request, float)
+            except ValueError:
+                await ctx.send_message("Invalid calculation format")
+        else:
+            await ctx.send_message("Insufficient parameters for calculation")
+
+    @response_handler
+    async def handle_calculation_response(
+        self, original_request: CalculationRequest, result: float, ctx: WorkflowContext[str]
+    ) -> None:
+        """Handle the calculation response."""
+        self.calculations_performed.append((original_request.operation, original_request.operands, result))
+        operands_str = ", ".join(map(str, original_request.operands))
+        await ctx.send_message(f"Calculation complete: {original_request.operation}({operands_str}) = {result}")
+
+
+class MultiRequestExecutor(Executor, RequestInfoMixin):
+    """Executor that makes multiple requests and waits for all responses."""
+
+    def __init__(self, id: str):
+        super().__init__(id=id)
+        self.responses_received: list[tuple[str, bool | float]] = []
+
+    @handler
+    async def start_multi_request(self, message: str, ctx: WorkflowContext) -> None:
+        """Start multiple requests simultaneously."""
+        # Request approval
+        approval_request = UserApprovalRequest(
+            prompt="Approve batch operation", context="Multiple operations will be performed"
+        )
+        await ctx.request_info(approval_request, bool)
+
+        # Request calculation
+        calc_request = CalculationRequest(operation="multiply", operands=[10.0, 5.0])
+        await ctx.request_info(calc_request, float)
+
+    @response_handler
+    async def handle_approval_response(
+        self, original_request: UserApprovalRequest, approved: bool, ctx: WorkflowContext[str]
+    ) -> None:
+        """Handle approval response."""
+        self.responses_received.append(("approval", approved))
+        await self._check_completion(ctx)
+
+    @response_handler
+    async def handle_calculation_response(
+        self, original_request: CalculationRequest, result: float, ctx: WorkflowContext[str]
+    ) -> None:
+        """Handle calculation response."""
+        self.responses_received.append(("calculation", result))
+        await self._check_completion(ctx)
+
+    async def _check_completion(self, ctx: WorkflowContext[str]) -> None:
+        """Check if all responses are received and send final result."""
+        if len(self.responses_received) == 2:
+            approval_result = next((r[1] for r in self.responses_received if r[0] == "approval"), None)
+            calc_result = next((r[1] for r in self.responses_received if r[0] == "calculation"), None)
+
+            if approval_result and calc_result is not None:
+                await ctx.send_message(f"All operations complete. Calculation result: {calc_result}")
+            else:
+                await ctx.send_message("Operations completed with mixed results")
+
+
+class OutputCollector(Executor):
+    """Simple executor that collects outputs for testing."""
+
+    def __init__(self, id: str):
+        super().__init__(id=id)
+        self.collected_outputs: list[str] = []
+
+    @handler
+    async def collect_output(self, message: str, ctx: WorkflowContext) -> None:
+        """Collect output messages."""
+        self.collected_outputs.append(message)
+
+
+class TestRequestInfoAndResponse:
+    """Test cases for end-to-end request info and response handling at the workflow level."""
+
+    async def test_approval_workflow(self):
+        """Test end-to-end workflow with approval request."""
+        executor = ApprovalRequiredExecutor(id="approval_executor")
+        workflow = WorkflowBuilder(start_executor=executor).build()
+
+        # First run the workflow until it emits a request
+        request_info_event: WorkflowEvent | None = None
+        async for event in workflow.run("test operation", stream=True):
+            if event.type == "request_info":
+                request_info_event = event
+
+        assert request_info_event is not None
+        assert isinstance(request_info_event.data, UserApprovalRequest)
+        assert request_info_event.data.prompt == "Please approve the operation: test operation"
+
+        # Send response and continue workflow
+        completed = False
+        async for event in workflow.run(stream=True, responses={request_info_event.request_id: True}):
+            if event.type == "status" and event.state == WorkflowRunState.IDLE:
+                completed = True
+
+        assert completed
+        assert executor.approval_received is True
+        assert executor.final_result == "Operation approved: Please approve the operation: test operation"
+
+    async def test_calculation_workflow(self):
+        """Test end-to-end workflow with calculation request."""
+        executor = CalculationExecutor(id="calc_executor")
+        workflow = WorkflowBuilder(start_executor=executor).build()
+
+        # First run the workflow until it emits a calculation request
+        request_info_event: WorkflowEvent | None = None
+        async for event in workflow.run("multiply 15.5 2.0", stream=True):
+            if event.type == "request_info":
+                request_info_event = event
+
+        assert request_info_event is not None
+        assert isinstance(request_info_event.data, CalculationRequest)
+        assert request_info_event.data.operation == "multiply"
+        assert request_info_event.data.operands == [15.5, 2.0]
+
+        # Send response with calculated result
+        calculated_result = 31.0
+        completed = False
+        async for event in workflow.run(stream=True, responses={request_info_event.request_id: calculated_result}):
+            if event.type == "status" and event.state == WorkflowRunState.IDLE:
+                completed = True
+
+        assert completed
+        assert len(executor.calculations_performed) == 1
+        assert executor.calculations_performed[0] == ("multiply", [15.5, 2.0], calculated_result)
+
+    async def test_multiple_requests_workflow(self):
+        """Test workflow with multiple concurrent requests."""
+        executor = MultiRequestExecutor(id="multi_executor")
+        workflow = WorkflowBuilder(start_executor=executor).build()
+
+        # Collect all request events by running the full stream
+        request_events: list[WorkflowEvent] = []
+        async for event in workflow.run("start batch", stream=True):
+            if event.type == "request_info":
+                request_events.append(event)
+
+        assert len(request_events) == 2
+
+        # Find the approval and calculation requests
+        approval_event: WorkflowEvent | None = next(
+            (e for e in request_events if isinstance(e.data, UserApprovalRequest)), None
+        )
+        calc_event: WorkflowEvent | None = next(
+            (e for e in request_events if isinstance(e.data, CalculationRequest)), None
+        )
+
+        assert approval_event is not None
+        assert calc_event is not None
+
+        # Send responses for both requests
+        responses = {approval_event.request_id: True, calc_event.request_id: 50.0}
+        completed = False
+        async for event in workflow.run(stream=True, responses=responses):
+            if event.type == "status" and event.state == WorkflowRunState.IDLE:
+                completed = True
+
+        assert completed
+        assert len(executor.responses_received) == 2
+
+    async def test_denied_approval_workflow(self):
+        """Test workflow when approval is denied."""
+        executor = ApprovalRequiredExecutor(id="approval_executor")
+        workflow = WorkflowBuilder(start_executor=executor).build()
+
+        # First run the workflow until it emits a request
+        request_info_event: WorkflowEvent | None = None
+        async for event in workflow.run("sensitive operation", stream=True):
+            if event.type == "request_info":
+                request_info_event = event
+
+        assert request_info_event is not None
+
+        # Deny the request
+        completed = False
+        async for event in workflow.run(stream=True, responses={request_info_event.request_id: False}):
+            if event.type == "status" and event.state == WorkflowRunState.IDLE:
+                completed = True
+
+        assert completed
+        assert executor.approval_received is True
+        assert executor.final_result == "Operation denied by user"
+
+    async def test_workflow_state_with_pending_requests(self):
+        """Test workflow state when waiting for responses."""
+        executor = ApprovalRequiredExecutor(id="approval_executor")
+        workflow = WorkflowBuilder(start_executor=executor).build()
+
+        # Run workflow until idle with pending requests
+        request_info_event: WorkflowEvent | None = None
+        idle_with_pending = False
+        async for event in workflow.run("test operation", stream=True):
+            if event.type == "request_info":
+                request_info_event = event
+            elif event.type == "status" and event.state == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS:
+                idle_with_pending = True
+
+        assert request_info_event is not None
+        assert idle_with_pending
+
+        # Continue with response
+        completed = False
+        async for event in workflow.run(stream=True, responses={request_info_event.request_id: True}):
+            if event.type == "status" and event.state == WorkflowRunState.IDLE:
+                completed = True
+
+        assert completed
+
+    async def test_invalid_calculation_input(self):
+        """Test workflow handling of invalid calculation input."""
+        executor = CalculationExecutor(id="calc_executor")
+        workflow = WorkflowBuilder(start_executor=executor).build()
+
+        # Send invalid input (no numbers)
+        completed = False
+        async for event in workflow.run("invalid input", stream=True):
+            if event.type == "status" and event.state == WorkflowRunState.IDLE:
+                completed = True
+
+        assert completed
+        # Should not have any calculations performed due to invalid input
+        assert len(executor.calculations_performed) == 0

python/packages/core/tests/workflow/test_request_info_event_rehydrate.py

@@ -0,0 +1,387 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import json
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+
+from agent_framework import (
+    FileCheckpointStorage,
+    InMemoryCheckpointStorage,
+    InProcRunnerContext,
+    WorkflowBuilder,
+    WorkflowRunState,
+)
+from agent_framework._workflows._checkpoint_encoding import (
+    _PICKLE_MARKER,  # type: ignore
+    encode_checkpoint_value,
+)
+from agent_framework._workflows._events import WorkflowEvent
+from agent_framework._workflows._state import State
+
+from .test_request_info_and_response import (
+    ApprovalRequiredExecutor,
+    CalculationRequest,
+    MultiRequestExecutor,
+    UserApprovalRequest,
+)
+
+
+@dataclass
+class MockRequest: ...
+
+
+@dataclass(kw_only=True)
+class SimpleApproval:
+    prompt: str = ""
+    draft: str = ""
+    iteration: int = 0
+
+
+@dataclass(slots=True)
+class SlottedApproval:
+    note: str = ""
+
+
+@dataclass
+class TimedApproval:
+    issued_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+
+
+async def test_rehydrate_request_info_event() -> None:
+    """Rehydration should succeed for valid request info events."""
+    request_info_event = WorkflowEvent.request_info(
+        request_id="request-123",
+        source_executor_id="review_gateway",
+        request_data=MockRequest(),
+        response_type=bool,
+    )
+
+    runner_context = InProcRunnerContext(InMemoryCheckpointStorage())
+    await runner_context.add_request_info_event(request_info_event)
+
+    checkpoint_id = await runner_context.create_checkpoint("test_name", "test_hash", State(), None, iteration_count=1)
+    checkpoint = await runner_context.load_checkpoint(checkpoint_id)
+
+    assert checkpoint is not None
+    assert checkpoint.pending_request_info_events
+    assert "request-123" in checkpoint.pending_request_info_events
+    assert checkpoint.pending_request_info_events["request-123"].request_type is MockRequest
+
+    # Rehydrate the context
+    await runner_context.apply_checkpoint(checkpoint)
+
+    pending_requests = await runner_context.get_pending_request_info_events()
+    assert "request-123" in pending_requests
+    rehydrated_event = pending_requests["request-123"]
+    assert rehydrated_event.request_id == "request-123"
+    assert rehydrated_event.source_executor_id == "review_gateway"
+    assert rehydrated_event.request_type is MockRequest
+    assert rehydrated_event.response_type is bool
+    assert isinstance(rehydrated_event.data, MockRequest)
+
+
+async def test_request_info_event_serializes_non_json_payloads() -> None:
+    req_1 = WorkflowEvent.request_info(
+        request_id="req-1",
+        source_executor_id="source",
+        request_data=TimedApproval(issued_at=datetime(2024, 5, 4, 12, 30, 45)),
+        response_type=bool,
+    )
+    req_2 = WorkflowEvent.request_info(
+        request_id="req-2",
+        source_executor_id="source",
+        request_data=SlottedApproval(note="slot-based"),
+        response_type=bool,
+    )
+
+    runner_context = InProcRunnerContext(InMemoryCheckpointStorage())
+    await runner_context.add_request_info_event(req_1)
+    await runner_context.add_request_info_event(req_2)
+
+    checkpoint_id = await runner_context.create_checkpoint("test_name", "test_hash", State(), None, iteration_count=1)
+    checkpoint = await runner_context.load_checkpoint(checkpoint_id)
+
+    # Should be JSON serializable despite datetime/slots
+    serialized = json.dumps(encode_checkpoint_value(checkpoint))
+    assert isinstance(serialized, str)
+
+    # Verify the structure contains pickled data for the request data fields
+    deserialized = json.loads(serialized)
+    assert _PICKLE_MARKER in deserialized  # checkpoint itself is pickled
+
+    # Verify we can rehydrate the checkpoint correctly
+    await runner_context.apply_checkpoint(checkpoint)
+    pending = await runner_context.get_pending_request_info_events()
+
+    assert "req-1" in pending
+    rehydrated_1 = pending["req-1"]
+    assert isinstance(rehydrated_1.data, TimedApproval)
+    assert rehydrated_1.data.issued_at == datetime(2024, 5, 4, 12, 30, 45)
+
+    assert "req-2" in pending
+    rehydrated_2 = pending["req-2"]
+    assert isinstance(rehydrated_2.data, SlottedApproval)
+    assert rehydrated_2.data.note == "slot-based"
+
+
+async def test_checkpoint_with_pending_request_info_events():
+    """Test that request info events are properly serialized in checkpoints and can be restored."""
+    import tempfile
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        # Use file-based storage to test full serialization
+        storage = FileCheckpointStorage(
+            temp_dir,
+            allowed_checkpoint_types=[
+                "tests.workflow.test_request_info_and_response:UserApprovalRequest",
+                "tests.workflow.test_request_info_and_response:CalculationRequest",
+                "tests.workflow.test_request_info_event_rehydrate:MockRequest",
+                "tests.workflow.test_request_info_event_rehydrate:SimpleApproval",
+                "tests.workflow.test_request_info_event_rehydrate:SlottedApproval",
+                "tests.workflow.test_request_info_event_rehydrate:TimedApproval",
+            ],
+        )
+
+        # Create workflow with checkpointing enabled
+        executor = ApprovalRequiredExecutor(id="approval_executor")
+        workflow = WorkflowBuilder(start_executor=executor, checkpoint_storage=storage).build()
+
+        # Step 1: Run workflow to completion to ensure checkpoints are created
+        request_info_event: WorkflowEvent | None = None
+        async for event in workflow.run("checkpoint test operation", stream=True):
+            if event.type == "request_info":
+                request_info_event = event
+
+        # Verify request was emitted
+        assert request_info_event is not None
+        assert isinstance(request_info_event.data, UserApprovalRequest)
+        assert request_info_event.data.prompt == "Please approve the operation: checkpoint test operation"
+        assert request_info_event.source_executor_id == "approval_executor"
+
+        # Step 2: List checkpoints to find the one with our pending request
+        checkpoints = await storage.list_checkpoints(workflow_name=workflow.name)
+        assert len(checkpoints) > 0, "No checkpoints were created during workflow execution"
+
+        # Find the checkpoint with our pending request
+        checkpoint_with_request = None
+        for checkpoint in checkpoints:
+            if request_info_event.request_id in checkpoint.pending_request_info_events:
+                checkpoint_with_request = checkpoint
+                break
+
+        assert checkpoint_with_request is not None, "No checkpoint found with pending request info event"
+
+        # Step 3: Verify the pending request info event was properly serialized
+        serialized_event = checkpoint_with_request.pending_request_info_events[request_info_event.request_id]
+        assert serialized_event.data
+        assert serialized_event.request_type is UserApprovalRequest
+        assert serialized_event.request_id == request_info_event.request_id
+        assert serialized_event.source_executor_id == "approval_executor"
+
+        # Step 4: Create a fresh workflow and restore from checkpoint
+        new_executor = ApprovalRequiredExecutor(id="approval_executor")
+        restored_workflow = WorkflowBuilder(start_executor=new_executor, checkpoint_storage=storage).build()
+
+        # Step 5: Resume from checkpoint and verify the request can be continued
+        completed = False
+        restored_request_event: WorkflowEvent | None = None
+        async for event in restored_workflow.run(checkpoint_id=checkpoint_with_request.checkpoint_id, stream=True):
+            # Should re-emit the pending request info event
+            if event.type == "request_info" and event.request_id == request_info_event.request_id:
+                restored_request_event = event
+            elif event.type == "status" and event.state == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS:
+                completed = True
+
+        assert completed, "Workflow should reach idle with pending requests state after restoration"
+        assert restored_request_event is not None, "Restored request info event should be emitted"
+
+        # Verify the restored event matches the original
+        assert restored_request_event.source_executor_id == request_info_event.source_executor_id
+        assert isinstance(restored_request_event.data, UserApprovalRequest)
+        assert restored_request_event.data.prompt == request_info_event.data.prompt
+        assert restored_request_event.data.context == request_info_event.data.context
+
+        # Step 6: Provide response to the restored request and complete the workflow
+        final_completed = False
+        async for event in restored_workflow.run(
+            stream=True,
+            responses={
+                request_info_event.request_id: True  # Approve the request
+            },
+        ):
+            if event.type == "status" and event.state == WorkflowRunState.IDLE:
+                final_completed = True
+
+        assert final_completed, "Workflow should complete after providing response to restored request"
+
+        # Step 7: Verify the executor state was properly restored and response was processed
+        assert new_executor.approval_received is True
+        expected_result = "Operation approved: Please approve the operation: checkpoint test operation"
+        assert new_executor.final_result == expected_result
+
+
+async def test_checkpoint_restore_with_responses_does_not_reemit_handled_requests():
+    """Test that request_info events are not re-emitted when responses are provided with checkpoint restore.
+
+    When calling run(checkpoint_id=..., responses=...), the workflow restores from a checkpoint
+    that contains pending request_info events. Because responses are provided for those events,
+    they should NOT be re-emitted in the event stream - they are considered "handled".
+
+    Note: The workflow's internal state tracking still sees the request_info events (before filtering),
+    so the final status may be IDLE_WITH_PENDING_REQUESTS even though the requests were handled.
+    The key behavior we're testing is that the CALLER doesn't see the request_info events.
+    """
+    import tempfile
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        # Use file-based storage to test full serialization
+        storage = FileCheckpointStorage(
+            temp_dir,
+            allowed_checkpoint_types=[
+                "tests.workflow.test_request_info_and_response:UserApprovalRequest",
+                "tests.workflow.test_request_info_and_response:CalculationRequest",
+                "tests.workflow.test_request_info_event_rehydrate:MockRequest",
+                "tests.workflow.test_request_info_event_rehydrate:SimpleApproval",
+                "tests.workflow.test_request_info_event_rehydrate:SlottedApproval",
+                "tests.workflow.test_request_info_event_rehydrate:TimedApproval",
+            ],
+        )
+
+        # Create workflow with checkpointing enabled
+        executor = ApprovalRequiredExecutor(id="approval_executor")
+        workflow = WorkflowBuilder(start_executor=executor, checkpoint_storage=storage).build()
+
+        # Step 1: Run workflow until it emits a request_info event
+        request_info_event: WorkflowEvent | None = None
+        async for event in workflow.run("test pending request suppression", stream=True):
+            if event.type == "request_info":
+                request_info_event = event
+
+        assert request_info_event is not None
+        request_id = request_info_event.request_id
+
+        # Step 2: Find the checkpoint with the pending request
+        checkpoints = await storage.list_checkpoints(workflow_name=workflow.name)
+        checkpoint_with_request = None
+        for checkpoint in checkpoints:
+            if request_id in checkpoint.pending_request_info_events:
+                checkpoint_with_request = checkpoint
+                break
+
+        assert checkpoint_with_request is not None
+
+        # Step 3: Create a fresh workflow and restore from checkpoint WITH responses in one call
+        new_executor = ApprovalRequiredExecutor(id="approval_executor")
+        restored_workflow = WorkflowBuilder(start_executor=new_executor, checkpoint_storage=storage).build()
+
+        # Track all emitted events
+        emitted_events: list[WorkflowEvent] = []
+        async for event in restored_workflow.run(
+            checkpoint_id=checkpoint_with_request.checkpoint_id,
+            responses={request_id: True},  # Provide response for the pending request
+            stream=True,
+        ):
+            emitted_events.append(event)
+
+        # Step 4: Verify the request_info event was NOT re-emitted to the caller
+        reemitted_request_info_events = [
+            e for e in emitted_events if e.type == "request_info" and e.request_id == request_id
+        ]
+        assert len(reemitted_request_info_events) == 0, (
+            f"request_info event should NOT be re-emitted when response is provided. "
+            f"Found {len(reemitted_request_info_events)} request_info events with request_id={request_id}"
+        )
+
+        # Step 5: Verify the response was processed by checking executor state
+        assert new_executor.approval_received is True, "Response should have been processed by the executor"
+        assert new_executor.final_result == (
+            "Operation approved: Please approve the operation: test pending request suppression"
+        )
+
+
+async def test_checkpoint_restore_with_partial_responses_reemits_unhandled_requests():
+    """Test that only unhandled request_info events are re-emitted when partial responses are provided.
+
+    When calling run(checkpoint_id=..., responses=...) with responses for only some of the
+    pending requests, only the unhandled request_info events should be re-emitted.
+    """
+    import tempfile
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(
+            temp_dir,
+            allowed_checkpoint_types=[
+                "tests.workflow.test_request_info_and_response:UserApprovalRequest",
+                "tests.workflow.test_request_info_and_response:CalculationRequest",
+                "tests.workflow.test_request_info_event_rehydrate:MockRequest",
+                "tests.workflow.test_request_info_event_rehydrate:SimpleApproval",
+                "tests.workflow.test_request_info_event_rehydrate:SlottedApproval",
+                "tests.workflow.test_request_info_event_rehydrate:TimedApproval",
+            ],
+        )
+
+        # Create workflow with multiple requests
+        executor = MultiRequestExecutor(id="multi_executor")
+        workflow = WorkflowBuilder(start_executor=executor, checkpoint_storage=storage).build()
+
+        # Step 1: Run workflow until it emits multiple request_info events
+        request_events: list[WorkflowEvent] = []
+        async for event in workflow.run("start batch", stream=True):
+            if event.type == "request_info":
+                request_events.append(event)
+
+        assert len(request_events) == 2
+
+        # Find the approval and calculation requests
+        approval_event = next((e for e in request_events if isinstance(e.data, UserApprovalRequest)), None)
+        calc_event = next((e for e in request_events if isinstance(e.data, CalculationRequest)), None)
+        assert approval_event is not None
+        assert calc_event is not None
+
+        # Step 2: Find the checkpoint with pending requests
+        checkpoints = await storage.list_checkpoints(workflow_name=workflow.name)
+        checkpoint_with_requests = None
+        for checkpoint in checkpoints:
+            has_approval = approval_event.request_id in checkpoint.pending_request_info_events
+            has_calc = calc_event.request_id in checkpoint.pending_request_info_events
+            if has_approval and has_calc:
+                checkpoint_with_requests = checkpoint
+                break
+
+        assert checkpoint_with_requests is not None
+
+        # Step 3: Restore from checkpoint with ONLY the approval response (not the calculation)
+        new_executor = MultiRequestExecutor(id="multi_executor")
+        restored_workflow = WorkflowBuilder(start_executor=new_executor, checkpoint_storage=storage).build()
+
+        emitted_events: list[WorkflowEvent] = []
+        async for event in restored_workflow.run(
+            checkpoint_id=checkpoint_with_requests.checkpoint_id,
+            responses={approval_event.request_id: True},  # Only respond to approval
+            stream=True,
+        ):
+            emitted_events.append(event)
+
+        # Step 4: Verify the approval request_info was NOT re-emitted
+        reemitted_approval_events = [
+            e for e in emitted_events if e.type == "request_info" and e.request_id == approval_event.request_id
+        ]
+        assert len(reemitted_approval_events) == 0, (
+            "Approval request_info should NOT be re-emitted since response was provided"
+        )
+
+        # Step 5: Verify the calculation request_info WAS re-emitted (no response provided)
+        reemitted_calc_events = [
+            e for e in emitted_events if e.type == "request_info" and e.request_id == calc_event.request_id
+        ]
+        assert len(reemitted_calc_events) == 1, (
+            "Calculation request_info SHOULD be re-emitted since no response was provided"
+        )
+
+        # Step 6: Verify workflow is in IDLE_WITH_PENDING_REQUESTS state (calc still pending)
+        status_events = [e for e in emitted_events if e.type == "status"]
+        final_status = status_events[-1] if status_events else None
+        assert final_status is not None
+        assert final_status.state == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS, (
+            f"Workflow should be IDLE_WITH_PENDING_REQUESTS, got {final_status.state}"
+        )

python/packages/core/tests/workflow/test_request_info_mixin.py

@@ -0,0 +1,954 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import inspect
+from typing import Any
+
+import pytest
+
+from agent_framework._workflows._executor import Executor, handler
+from agent_framework._workflows._request_info_mixin import response_handler
+from agent_framework._workflows._workflow_context import WorkflowContext
+
+
+class TestRequestInfoMixin:
+    """Test cases for RequestInfoMixin functionality."""
+
+    def test_request_info_mixin_initialization(self):
+        """Test that RequestInfoMixin can be initialized."""
+
+        class TestExecutor(Executor):
+            def __init__(self):
+                super().__init__(id="test")
+
+            @handler
+            async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
+                pass
+
+        executor = TestExecutor()
+        # After calling _discover_response_handlers, it should have the attributes
+        assert hasattr(executor, "_response_handlers")
+        assert hasattr(executor, "_response_handler_specs")
+        assert hasattr(executor, "is_request_response_capable")
+        assert executor.is_request_response_capable is False
+
+    def test_response_handler_decorator_creates_metadata(self):
+        """Test that the response_handler decorator creates proper metadata."""
+
+        @response_handler
+        async def test_handler(self: Any, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
+            """Test handler docstring."""
+            pass
+
+        # Check that the decorator preserves function attributes
+        assert test_handler.__name__ == "test_handler"
+        assert test_handler.__doc__ == "Test handler docstring."
+        assert hasattr(test_handler, "_response_handler_spec")
+
+        # Check the spec attributes
+        spec = test_handler._response_handler_spec  # type: ignore[reportAttributeAccessIssue]
+        assert spec["name"] == "test_handler"  # ty: ignore[not-subscriptable]
+        assert spec["response_type"] is int  # ty: ignore[not-subscriptable]
+        assert spec["request_type"] is str  # ty: ignore[not-subscriptable]
+
+    def test_response_handler_with_workflow_context_types(self):
+        """Test response handler with different WorkflowContext type parameters."""
+
+        @response_handler
+        async def handler_with_output_types(
+            self: Any, original_request: str, response: int, ctx: WorkflowContext[str, bool]
+        ) -> None:
+            pass
+
+        spec = handler_with_output_types._response_handler_spec  # type: ignore[attr-defined, reportAttributeAccessIssue]  # ty: ignore[unresolved-attribute]
+        assert "output_types" in spec
+        assert "workflow_output_types" in spec
+
+    def test_response_handler_preserves_signature(self):
+        """Test that response_handler preserves the original function signature."""
+
+        async def original_handler(self: Any, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
+            pass
+
+        decorated = response_handler(original_handler)
+
+        # Check that signature is preserved
+        original_sig = inspect.signature(original_handler)
+        decorated_sig = inspect.signature(decorated)
+
+        # Both should have the same parameter names and types
+        assert list(original_sig.parameters.keys()) == list(decorated_sig.parameters.keys())
+
+    def test_executor_with_response_handlers(self):
+        """Test an executor with valid response handlers."""
+
+        class TestExecutor(Executor):
+            def __init__(self):
+                super().__init__(id="test_executor")
+
+            @handler
+            async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
+                pass
+
+            @response_handler
+            async def handle_string_response(
+                self, original_request: str, response: int, ctx: WorkflowContext[str]
+            ) -> None:
+                pass
+
+            @response_handler
+            async def handle_dict_response(
+                self, original_request: dict[str, Any], response: bool, ctx: WorkflowContext[bool]
+            ) -> None:
+                pass
+
+        executor = TestExecutor()
+
+        # Should be request-response capable
+        assert executor.is_request_response_capable is True
+
+        # Should have registered handlers
+        response_handlers = executor._response_handlers  # type: ignore[reportAttributeAccessIssue]
+        assert len(response_handlers) == 2
+        assert (str, int) in response_handlers
+        assert (dict[str, Any], bool) in response_handlers
+
+    def test_executor_without_response_handlers(self):
+        """Test an executor without response handlers."""
+
+        class PlainExecutor(Executor):
+            def __init__(self):
+                super().__init__(id="plain_executor")
+
+            @handler
+            async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
+                pass
+
+        executor = PlainExecutor()
+
+        # Should not be request-response capable
+        assert executor.is_request_response_capable is False
+
+        # Should have empty handlers
+        response_handlers = executor._response_handlers  # type: ignore[reportAttributeAccessIssue]
+        assert len(response_handlers) == 0
+
+    def test_duplicate_response_handlers_raise_error(self):
+        """Test that duplicate response handlers for the same message type raise an error."""
+
+        class DuplicateExecutor(Executor):
+            def __init__(self):
+                super().__init__(id="duplicate_executor")
+
+            @handler
+            async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
+                pass
+
+            @response_handler
+            async def handle_first(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
+                pass
+
+            @response_handler
+            async def handle_second(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
+                pass
+
+        with pytest.raises(
+            ValueError,
+            match="Duplicate response handler for request type <class 'str'> and response type <class 'int'>",
+        ):
+            DuplicateExecutor()
+
+    async def test_response_handler_function_callable(self):
+        """Test that response handlers can actually be called."""
+
+        class TestExecutor(Executor):
+            def __init__(self):
+                super().__init__(id="test_executor")
+                self.handled_request = None
+                self.handled_response = None
+
+            @handler
+            async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
+                pass
+
+            @response_handler
+            async def handle_response(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
+                self.handled_request = original_request  # type: ignore[assignment]
+                self.handled_response = response  # type: ignore[assignment]
+
+        executor = TestExecutor()
+
+        # Get the handler
+        response_handler_func = executor._response_handlers[(str, int)]  # type: ignore[reportAttributeAccessIssue]
+
+        # Create a mock context - we'll just use None since the handler doesn't use it
+        await response_handler_func("test_request", 42, None)  # type: ignore[arg-type, reportArgumentType]  # ty: ignore[invalid-argument-type]
+
+        assert executor.handled_request == "test_request"
+        assert executor.handled_response == 42
+
+    def test_inheritance_with_response_handlers(self):
+        """Test that response handlers work correctly with inheritance."""
+
+        class BaseExecutor(Executor):
+            def __init__(self):
+                super().__init__(id="base_executor")
+
+            @handler
+            async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
+                pass
+
+            @response_handler
+            async def base_handler(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
+                pass
+
+        class ChildExecutor(BaseExecutor):
+            def __init__(self):
+                super().__init__()
+                self.id = "child_executor"
+
+            @response_handler
+            async def child_handler(self, original_request: str, response: bool, ctx: WorkflowContext[str]) -> None:
+                pass
+
+        child = ChildExecutor()
+
+        # Should have both handlers
+        response_handlers = child._response_handlers  # type: ignore[reportAttributeAccessIssue]
+        assert len(response_handlers) == 2
+        assert (str, int) in response_handlers
+        assert (str, bool) in response_handlers
+        assert child.is_request_response_capable is True
+
+    def test_response_handler_spec_attributes(self):
+        """Test that response handler specs contain expected attributes."""
+
+        class TestExecutor(Executor):
+            def __init__(self):
+                super().__init__(id="test_executor")
+
+            @handler
+            async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
+                pass
+
+            @response_handler
+            async def test_handler(self, original_request: str, response: int, ctx: WorkflowContext[str, bool]) -> None:
+                pass
+
+        executor = TestExecutor()
+
+        specs = executor._response_handler_specs  # type: ignore[reportAttributeAccessIssue]
+        assert len(specs) == 1
+
+        spec = specs[0]
+        assert spec["name"] == "test_handler"
+        assert spec["request_type"] is str
+        assert spec["response_type"] is int
+        assert "output_types" in spec
+        assert "workflow_output_types" in spec
+        assert "ctx_annotation" in spec
+
+    def test_multiple_discovery_calls_raise_error(self):
+        """Test that multiple calls to _discover_response_handlers raise an error for duplicates."""
+
+        class TestExecutor(Executor):
+            def __init__(self):
+                super().__init__(id="test_executor")
+
+            @handler
+            async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
+                pass
+
+            @response_handler
+            async def test_handler(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
+                pass
+
+        executor = TestExecutor()
+
+        # First call should work fine
+        first_handlers = len(executor._response_handlers)  # type: ignore[reportAttributeAccessIssue]
+
+        # Second call should raise an error due to duplicate registration
+        with pytest.raises(
+            ValueError,
+            match="Duplicate response handler for request type <class 'str'> and response type <class 'int'>",
+        ):
+            executor._discover_response_handlers()  # type: ignore[reportAttributeAccessIssue]
+
+        # Handlers count should remain the same
+        assert first_handlers == 1
+
+    def test_non_callable_attributes_ignored(self):
+        """Test that non-callable attributes are ignored during discovery."""
+
+        class TestExecutor(Executor):
+            def __init__(self):
+                super().__init__(id="test_executor")
+
+            some_variable = "not_a_function"
+            another_attr = 42
+
+            @handler
+            async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
+                pass
+
+            @response_handler
+            async def valid_handler(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
+                pass
+
+        executor = TestExecutor()
+
+        # Should only have one handler despite other attributes
+        response_handlers = executor._response_handlers  # type: ignore[reportAttributeAccessIssue]
+        assert len(response_handlers) == 1
+        assert (str, int) in response_handlers
+
+    async def test_same_request_type_different_response_types(self):
+        """Test that handlers with same request type but different response types are distinct."""
+
+        class TestExecutor(Executor):
+            def __init__(self):
+                super().__init__(id="test_executor")
+                self.str_int_handler_called = False
+                self.str_bool_handler_called = False
+                self.str_dict_handler_called = False
+
+            @handler
+            async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
+                pass
+
+            @response_handler
+            async def handle_str_int(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
+                self.str_int_handler_called = True
+
+            @response_handler
+            async def handle_str_bool(self, original_request: str, response: bool, ctx: WorkflowContext[str]) -> None:
+                self.str_bool_handler_called = True
+
+            @response_handler
+            async def handle_str_dict(
+                self, original_request: str, response: dict[str, Any], ctx: WorkflowContext[str]
+            ) -> None:
+                self.str_dict_handler_called = True
+
+        executor = TestExecutor()
+
+        # Should have three distinct handlers
+        response_handlers = executor._response_handlers  # type: ignore[reportAttributeAccessIssue]
+        assert len(response_handlers) == 3
+        assert (str, int) in response_handlers
+        assert (str, bool) in response_handlers
+        assert (str, dict[str, Any]) in response_handlers
+
+        # Test that each handler can be found correctly
+        str_int_handler = executor._find_response_handler("test", 42)  # pyright: ignore[reportPrivateUsage]
+        str_bool_handler = executor._find_response_handler("test", True)  # pyright: ignore[reportPrivateUsage]
+        str_dict_handler = executor._find_response_handler("test", {"key": "value"})  # pyright: ignore[reportPrivateUsage]
+
+        assert str_int_handler is not None
+        assert str_bool_handler is not None
+        assert str_dict_handler is not None
+
+        # Test that handlers are called correctly
+        await str_int_handler(42, None)  # type: ignore[reportArgumentType]
+        await str_bool_handler(True, None)  # type: ignore[reportArgumentType]
+        await str_dict_handler({"key": "value"}, None)  # type: ignore[reportArgumentType]
+
+        assert executor.str_int_handler_called
+        assert executor.str_bool_handler_called
+        assert executor.str_dict_handler_called
+
+    async def test_different_request_types_same_response_type(self):
+        """Test that handlers with different request types but same response type are distinct."""
+
+        class TestExecutor(Executor):
+            def __init__(self):
+                super().__init__(id="test_executor")
+                self.str_int_handler_called = False
+                self.dict_int_handler_called = False
+                self.list_int_handler_called = False
+
+            @handler
+            async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
+                pass
+
+            @response_handler
+            async def handle_str_int(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
+                self.str_int_handler_called = True
+
+            @response_handler
+            async def handle_dict_int(
+                self, original_request: dict[str, Any], response: int, ctx: WorkflowContext[str]
+            ) -> None:
+                self.dict_int_handler_called = True
+
+            @response_handler
+            async def handle_list_int(
+                self, original_request: list[str], response: int, ctx: WorkflowContext[str]
+            ) -> None:
+                self.list_int_handler_called = True
+
+        executor = TestExecutor()
+
+        # Should have three distinct handlers
+        response_handlers = executor._response_handlers  # type: ignore[reportAttributeAccessIssue]
+        assert len(response_handlers) == 3
+        assert (str, int) in response_handlers
+        assert (dict[str, Any], int) in response_handlers
+        assert (list[str], int) in response_handlers
+
+        # Test that each handler can be found correctly
+        str_int_handler = executor._find_response_handler("test", 42)  # pyright: ignore[reportPrivateUsage]
+        dict_int_handler = executor._find_response_handler({"key": "value"}, 42)  # pyright: ignore[reportPrivateUsage]
+        list_int_handler = executor._find_response_handler(["test"], 42)  # pyright: ignore[reportPrivateUsage]
+
+        assert str_int_handler is not None
+        assert dict_int_handler is not None
+        assert list_int_handler is not None
+
+        # Test that handlers are called correctly
+        await str_int_handler(42, None)  # type: ignore[reportArgumentType]
+        await dict_int_handler(42, None)  # type: ignore[reportArgumentType]
+        await list_int_handler(42, None)  # type: ignore[reportArgumentType]
+
+        assert executor.str_int_handler_called
+        assert executor.dict_int_handler_called
+        assert executor.list_int_handler_called
+
+    def test_complex_type_combinations(self):
+        """Test response handlers with complex type combinations."""
+
+        class CustomRequest:
+            pass
+
+        class CustomResponse:
+            pass
+
+        class TestExecutor(Executor):
+            def __init__(self):
+                super().__init__(id="test_executor")
+                self.custom_custom_called = False
+                self.custom_str_called = False
+                self.str_custom_called = False
+
+            @handler
+            async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
+                pass
+
+            @response_handler
+            async def handle_custom_custom(
+                self, original_request: CustomRequest, response: CustomResponse, ctx: WorkflowContext[str]
+            ) -> None:
+                self.custom_custom_called = True
+
+            @response_handler
+            async def handle_custom_str(
+                self, original_request: CustomRequest, response: str, ctx: WorkflowContext[str]
+            ) -> None:
+                self.custom_str_called = True
+
+            @response_handler
+            async def handle_str_custom(
+                self, original_request: str, response: CustomResponse, ctx: WorkflowContext[str]
+            ) -> None:
+                self.str_custom_called = True
+
+        executor = TestExecutor()
+
+        # Should have three distinct handlers
+        response_handlers = executor._response_handlers  # type: ignore[reportAttributeAccessIssue]
+        assert len(response_handlers) == 3
+        assert (CustomRequest, CustomResponse) in response_handlers
+        assert (CustomRequest, str) in response_handlers
+        assert (str, CustomResponse) in response_handlers
+
+        # Test that each handler can be found correctly
+        custom_request = CustomRequest()
+        custom_response = CustomResponse()
+
+        custom_custom_handler = executor._find_response_handler(custom_request, custom_response)  # pyright: ignore[reportPrivateUsage]
+        custom_str_handler = executor._find_response_handler(custom_request, "test")  # pyright: ignore[reportPrivateUsage]
+        str_custom_handler = executor._find_response_handler("test", custom_response)  # pyright: ignore[reportPrivateUsage]
+
+        assert custom_custom_handler is not None
+        assert custom_str_handler is not None
+        assert str_custom_handler is not None
+
+    def test_handler_key_uniqueness(self):
+        """Test that handler keys (request_type, response_type) are truly unique."""
+
+        class TestExecutor(Executor):
+            def __init__(self):
+                super().__init__(id="test_executor")
+
+            @handler
+            async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
+                pass
+
+            @response_handler
+            async def handle1(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
+                pass
+
+            @response_handler
+            async def handle2(self, original_request: int, response: str, ctx: WorkflowContext[str]) -> None:
+                pass
+
+            @response_handler
+            async def handle3(self, original_request: str, response: str, ctx: WorkflowContext[str]) -> None:
+                pass
+
+            @response_handler
+            async def handle4(self, original_request: int, response: int, ctx: WorkflowContext[str]) -> None:
+                pass
+
+        executor = TestExecutor()
+
+        # Should have four distinct handlers based on different combinations
+        response_handlers = executor._response_handlers  # type: ignore[reportAttributeAccessIssue]
+        assert len(response_handlers) == 4
+
+        # Verify all expected combinations exist
+        expected_keys = {
+            (str, int),  # handle1
+            (int, str),  # handle2
+            (str, str),  # handle3
+            (int, int),  # handle4
+        }
+
+        actual_keys = set(response_handlers.keys())
+        assert actual_keys == expected_keys
+
+    def test_no_false_matches_with_similar_types(self):
+        """Test that handlers don't match with similar but different types."""
+
+        class TestExecutor(Executor):
+            def __init__(self):
+                super().__init__(id="test_executor")
+
+            @handler
+            async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
+                pass
+
+            @response_handler
+            async def handle_str_int(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
+                pass
+
+            @response_handler
+            async def handle_list_str_float(
+                self, original_request: list[str], response: float, ctx: WorkflowContext[str]
+            ) -> None:
+                pass
+
+        executor = TestExecutor()
+
+        # Test that wrong combinations don't match
+        assert executor._find_response_handler("test", 3.14) is None  # pyright: ignore[reportPrivateUsage]  # str request, float response - no handler
+        assert executor._find_response_handler(["test"], 42) is None  # pyright: ignore[reportPrivateUsage]  # list request, int response - no handler
+        assert executor._find_response_handler(42, "test") is None  # pyright: ignore[reportPrivateUsage]  # int request, str response - no handler
+
+        # Test that correct combinations do match
+        assert executor._find_response_handler("test", 42) is not None  # pyright: ignore[reportPrivateUsage]  # str request, int response - has handler
+        assert executor._find_response_handler(["test"], 3.14) is not None  # pyright: ignore[reportPrivateUsage]  # list request, float response - has handler
+
+    def test_is_request_supported_with_exact_matches(self):
+        """Test is_request_supported with exact type matches."""
+
+        class TestExecutor(Executor):
+            def __init__(self):
+                super().__init__(id="test_executor")
+
+            @handler
+            async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
+                pass
+
+            @response_handler
+            async def handle_str_int(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
+                pass
+
+            @response_handler
+            async def handle_dict_bool(
+                self, original_request: dict[str, Any], response: bool, ctx: WorkflowContext[str]
+            ) -> None:
+                pass
+
+        executor = TestExecutor()
+
+        # Test exact matches
+        assert executor.is_request_supported(str, int) is True
+        assert executor.is_request_supported(str, bool) is True  # bool and int are compatible
+        assert executor.is_request_supported(dict[str, Any], bool) is True
+
+        # Test non-matches
+        assert executor.is_request_supported(int, str) is False
+        assert executor.is_request_supported(list[str], int) is False
+
+    def test_is_request_supported_without_handlers(self):
+        """Test is_request_supported when no handlers are registered."""
+
+        class TestExecutor(Executor):
+            def __init__(self):
+                super().__init__(id="test_executor")
+
+            @handler
+            async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
+                pass
+
+        executor = TestExecutor()
+
+        # Should return False for any type combination
+        assert executor.is_request_supported(str, int) is False
+        assert executor.is_request_supported(dict[str, Any], bool) is False
+        assert executor.is_request_supported(int, str) is False
+
+    def test_is_request_supported_before_discovery(self):
+        """Test is_request_supported before response handlers are discovered."""
+
+        class TestExecutor(Executor):
+            def __init__(self):
+                super().__init__(id="test_executor", defer_discovery=True)
+
+            @handler
+            async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
+                pass
+
+            @response_handler
+            async def handle_str_int(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
+                pass
+
+        executor = TestExecutor()
+        # Don't call _discover_response_handlers()
+
+        # Should return False when _response_handlers attribute doesn't exist
+        assert executor.is_request_supported(str, int) is False
+        assert executor.is_request_supported(dict[str, Any], bool) is False
+
+    def test_is_request_supported_with_compatible_types(self):
+        """Test is_request_supported with type-compatible scenarios."""
+
+        class BaseRequest:
+            pass
+
+        class DerivedRequest(BaseRequest):
+            pass
+
+        class BaseResponse:
+            pass
+
+        class DerivedResponse(BaseResponse):
+            pass
+
+        class TestExecutor(Executor):
+            def __init__(self):
+                super().__init__(id="test_executor")
+
+            @handler
+            async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
+                pass
+
+            @response_handler
+            async def handle_base_base(
+                self, original_request: BaseRequest, response: BaseResponse, ctx: WorkflowContext[str]
+            ) -> None:
+                pass
+
+            @response_handler
+            async def handle_str_int(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
+                pass
+
+        executor = TestExecutor()
+
+        # Test exact matches
+        assert executor.is_request_supported(BaseRequest, BaseResponse) is True
+        assert executor.is_request_supported(str, int) is True
+
+        # Test compatible derived types (depends on is_type_compatible implementation)
+        # These should return True if the type compatibility function supports inheritance
+        result_derived_request = executor.is_request_supported(DerivedRequest, BaseResponse)
+        result_derived_response = executor.is_request_supported(BaseRequest, DerivedResponse)
+        result_both_derived = executor.is_request_supported(DerivedRequest, DerivedResponse)
+
+        # The actual result depends on the is_type_compatible implementation
+        # We'll just assert that the method doesn't raise an exception
+        assert isinstance(result_derived_request, bool)
+        assert isinstance(result_derived_response, bool)
+        assert isinstance(result_both_derived, bool)
+
+    def test_is_request_supported_with_multiple_handlers(self):
+        """Test is_request_supported when multiple handlers are registered."""
+
+        class TestExecutor(Executor):
+            def __init__(self):
+                super().__init__(id="test_executor")
+
+            @handler
+            async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
+                pass
+
+            @response_handler
+            async def handle_str_int(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
+                pass
+
+            @response_handler
+            async def handle_str_bool(self, original_request: str, response: bool, ctx: WorkflowContext[str]) -> None:
+                pass
+
+            @response_handler
+            async def handle_dict_str(
+                self, original_request: dict[str, Any], response: str, ctx: WorkflowContext[str]
+            ) -> None:
+                pass
+
+            @response_handler
+            async def handle_list_float(
+                self, original_request: list[str], response: float, ctx: WorkflowContext[str]
+            ) -> None:
+                pass
+
+        executor = TestExecutor()
+
+        # Test all registered combinations
+        assert executor.is_request_supported(str, int) is True
+        assert executor.is_request_supported(str, bool) is True
+        assert executor.is_request_supported(dict[str, Any], str) is True
+        assert executor.is_request_supported(list[str], float) is True
+
+        # Test combinations that don't exist
+        assert executor.is_request_supported(str, float) is False
+        assert executor.is_request_supported(int, str) is False
+        assert executor.is_request_supported(dict[str, Any], int) is False
+        assert executor.is_request_supported(list[str], bool) is False
+
+    def test_is_request_supported_with_complex_types(self):
+        """Test is_request_supported with complex generic types."""
+
+        class TestExecutor(Executor):
+            def __init__(self):
+                super().__init__(id="test_executor")
+
+            @handler
+            async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
+                pass
+
+            @response_handler
+            async def handle_dict_list(
+                self, original_request: dict[str, Any], response: list[int], ctx: WorkflowContext[str]
+            ) -> None:
+                pass
+
+            @response_handler
+            async def handle_list_dict(
+                self, original_request: list[str], response: dict[str, bool], ctx: WorkflowContext[str]
+            ) -> None:
+                pass
+
+        executor = TestExecutor()
+
+        # Test complex type matches
+        assert executor.is_request_supported(dict[str, Any], list[int]) is True
+        assert executor.is_request_supported(list[str], dict[str, bool]) is True
+
+        # Test non-matches with similar but different complex types
+        assert executor.is_request_supported(dict[str, Any], list[str]) is False
+        assert executor.is_request_supported(list[int], dict[str, bool]) is False
+        assert executor.is_request_supported(dict[int, Any], list[int]) is False
+
+    def test_is_request_supported_with_inheritance(self):
+        """Test is_request_supported with inherited response handlers."""
+
+        class BaseExecutor(Executor):
+            def __init__(self):
+                super().__init__(id="base_executor")
+
+            @handler
+            async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
+                pass
+
+            @response_handler
+            async def base_handler(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
+                pass
+
+        class ChildExecutor(BaseExecutor):
+            def __init__(self):
+                super().__init__()
+                self.id = "child_executor"
+
+            @response_handler
+            async def child_handler(self, original_request: str, response: bool, ctx: WorkflowContext[str]) -> None:
+                pass
+
+        child = ChildExecutor()
+
+        # Should support both inherited and child-defined handlers
+        assert child.is_request_supported(str, int) is True  # From base class
+        assert child.is_request_supported(str, bool) is True  # From child class
+
+        # Should not support unregistered combinations
+        assert child.is_request_supported(str, str) is False
+        assert child.is_request_supported(int, str) is False
+
+
+class TestResponseHandlerExplicitTypes:
+    """Test cases for response_handler with explicit type parameters."""
+
+    def test_response_handler_with_explicit_types(self):
+        """Test response_handler with explicit request and response types."""
+
+        @response_handler(request=str, response=int)
+        async def test_handler(self: Any, original_request: Any, response: Any, ctx: WorkflowContext) -> None:
+            pass
+
+        spec = test_handler._response_handler_spec  # type: ignore[attr-defined, reportAttributeAccessIssue]  # ty: ignore[unresolved-attribute]
+        assert spec["name"] == "test_handler"
+        assert spec["request_type"] is str
+        assert spec["response_type"] is int
+
+    def test_response_handler_with_explicit_output_types(self):
+        """Test response_handler with explicit output and workflow_output types."""
+
+        @response_handler(request=str, response=int, output=bool, workflow_output=float)
+        async def test_handler(self: Any, original_request: Any, response: Any, ctx: WorkflowContext) -> None:
+            pass
+
+        spec = test_handler._response_handler_spec  # type: ignore[attr-defined, reportAttributeAccessIssue]  # ty: ignore[unresolved-attribute]
+        assert spec["request_type"] is str
+        assert spec["response_type"] is int
+        assert bool in spec["output_types"]
+        assert float in spec["workflow_output_types"]
+
+    def test_response_handler_with_union_types(self):
+        """Test response_handler with union types."""
+
+        @response_handler(request=str | int, response=bool | float)  # pyright: ignore[reportArgumentType]
+        async def test_handler(self: Any, original_request: Any, response: Any, ctx: WorkflowContext) -> None:
+            pass
+
+        spec = test_handler._response_handler_spec  # type: ignore[attr-defined, reportAttributeAccessIssue]  # ty: ignore[unresolved-attribute]
+        assert spec["request_type"] == str | int
+        assert spec["response_type"] == bool | float
+
+    def test_response_handler_with_string_forward_references(self):
+        """Test response_handler with string forward references."""
+
+        @response_handler(request="str", response="int")
+        async def test_handler(self: Any, original_request: Any, response: Any, ctx: WorkflowContext) -> None:
+            pass
+
+        spec = test_handler._response_handler_spec  # type: ignore[attr-defined, reportAttributeAccessIssue]  # ty: ignore[unresolved-attribute]
+        assert spec["request_type"] is str
+        assert spec["response_type"] is int
+
+    def test_response_handler_explicit_missing_request_raises_error(self):
+        """Test that using explicit types without request raises an error."""
+        with pytest.raises(ValueError, match="must specify 'request' type"):
+
+            @response_handler(response=int)
+            async def test_handler(self: Any, original_request: Any, response: Any, ctx: WorkflowContext) -> None:  # pyright: ignore[reportUnusedFunction]
+                pass
+
+    def test_response_handler_explicit_missing_response_raises_error(self):
+        """Test that using explicit types without response raises an error."""
+        with pytest.raises(ValueError, match="must specify 'response' type"):
+
+            @response_handler(request=str)
+            async def test_handler(self: Any, original_request: Any, response: Any, ctx: WorkflowContext) -> None:  # pyright: ignore[reportUnusedFunction]
+                pass
+
+    def test_response_handler_explicit_only_output_raises_error(self):
+        """Test that using only output without request/response raises an error."""
+        with pytest.raises(ValueError, match="must specify 'request' type"):
+
+            @response_handler(output=bool)
+            async def test_handler(self: Any, original_request: Any, response: Any, ctx: WorkflowContext) -> None:  # pyright: ignore[reportUnusedFunction]
+                pass
+
+    def test_executor_with_explicit_response_handlers(self):
+        """Test an executor with explicit type response handlers."""
+
+        class TestExecutor(Executor):
+            def __init__(self):
+                super().__init__(id="test_executor")
+
+            @handler
+            async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
+                pass
+
+            @response_handler(request=str, response=int, output=bool)
+            async def handle_explicit(self, original_request: Any, response: Any, ctx: WorkflowContext) -> None:
+                pass
+
+        executor = TestExecutor()
+
+        # Should be request-response capable
+        assert executor.is_request_response_capable is True
+
+        # Should have registered handler
+        response_handlers = executor._response_handlers  # type: ignore[reportAttributeAccessIssue]
+        assert len(response_handlers) == 1
+        assert (str, int) in response_handlers
+
+        # Check specs
+        specs = executor._response_handler_specs  # type: ignore[reportAttributeAccessIssue]
+        assert len(specs) == 1
+        assert specs[0]["request_type"] is str
+        assert specs[0]["response_type"] is int
+        assert bool in specs[0]["output_types"]
+
+    def test_response_handler_explicit_callable(self):
+        """Test that explicit type response handlers can be called."""
+
+        class TestExecutor(Executor):
+            def __init__(self):
+                super().__init__(id="test_executor")
+                self.handled_request = None
+                self.handled_response = None
+
+            @handler
+            async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
+                pass
+
+            @response_handler(request=str, response=int)
+            async def handle_response(self, original_request: Any, response: Any, ctx: WorkflowContext) -> None:
+                self.handled_request = original_request
+                self.handled_response = response
+
+        executor = TestExecutor()
+
+        # Get the handler
+        response_handler_func = executor._response_handlers[(str, int)]  # type: ignore[reportAttributeAccessIssue]
+
+        # Call the handler
+        asyncio.run(response_handler_func("test_request", 42, None))  # type: ignore[arg-type, reportArgumentType]  # ty: ignore[invalid-argument-type]
+
+        assert executor.handled_request == "test_request"
+        assert executor.handled_response == 42
+
+    def test_mixed_introspection_and_explicit_handlers(self):
+        """Test executor with both introspection and explicit type handlers."""
+
+        class TestExecutor(Executor):
+            def __init__(self):
+                super().__init__(id="test_executor")
+
+            @handler
+            async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
+                pass
+
+            # Introspection-based handler
+            @response_handler
+            async def handle_introspection(
+                self, original_request: str, response: int, ctx: WorkflowContext[str]
+            ) -> None:
+                pass
+
+            # Explicit type handler
+            @response_handler(request=dict, response=bool)
+            async def handle_explicit(self, original_request: Any, response: Any, ctx: WorkflowContext) -> None:
+                pass
+
+        executor = TestExecutor()
+
+        # Should have both handlers
+        response_handlers = executor._response_handlers  # type: ignore[reportAttributeAccessIssue]
+        assert len(response_handlers) == 2
+        assert (str, int) in response_handlers
+        assert (dict, bool) in response_handlers

python/packages/core/tests/workflow/test_serialization.py

@@ -0,0 +1,846 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import json
+from typing import Any
+
+import pytest
+
+from agent_framework import Executor, WorkflowBuilder, WorkflowContext, handler
+from agent_framework._workflows._const import INTERNAL_SOURCE_ID
+from agent_framework._workflows._edge import (
+    Case,
+    Default,
+    Edge,
+    FanInEdgeGroup,
+    FanOutEdgeGroup,
+    InternalEdgeGroup,
+    SingleEdgeGroup,
+    SwitchCaseEdgeGroup,
+    SwitchCaseEdgeGroupCase,
+    SwitchCaseEdgeGroupDefault,
+)
+from agent_framework._workflows._workflow_executor import (
+    WorkflowExecutor,
+)
+
+
+class SampleExecutor(Executor):
+    """Sample executor for serialization testing."""
+
+    @handler
+    async def handle_str(self, message: str, ctx: WorkflowContext[str]) -> None:
+        """Handle string messages."""
+        await ctx.send_message(f"Processed: {message}")
+
+
+class SampleAggregator(Executor):
+    """Sample aggregator executor that can handle lists of messages."""
+
+    @handler
+    async def handle_str_list(self, messages: list[str], ctx: WorkflowContext[str]) -> None:
+        """Handle list of string messages for fan-in aggregation."""
+        combined = " | ".join(messages)
+        await ctx.send_message(f"Aggregated: {combined}")
+
+
+class TestSerializationWorkflowClasses:
+    """Test serialization of workflow classes."""
+
+    def test_executor_serialization(self) -> None:
+        """Test that Executor can be serialized and has correct fields, including type."""
+        executor = SampleExecutor(id="test-executor")
+
+        # Test to_dict
+        data = executor.to_dict()
+        assert data["id"] == "test-executor"
+
+        # Test type field
+        assert "type" in data, "Executor should have 'type' field"
+        assert data["type"] == "SampleExecutor", f"Expected type 'SampleExecutor', got {data['type']}"
+
+        # Test model_dump_json
+        json_str = executor.to_json()
+        parsed = json.loads(json_str)
+        assert parsed["id"] == "test-executor"
+
+        # Test type field in JSON
+        assert "type" in parsed, "JSON should have 'type' field"
+        assert parsed["type"] == "SampleExecutor", "JSON should preserve type field"
+
+    def test_edge_serialization(self) -> None:
+        """Test that Edge can be serialized and has correct fields."""
+        # Test edge without condition
+        edge = Edge(source_id="source", target_id="target")
+
+        # Test to_dict
+        data = edge.to_dict()
+        assert data["source_id"] == "source"
+        assert data["target_id"] == "target"
+        assert "condition_name" not in data or data["condition_name"] is None
+
+        # Test model_dump_json
+        json_str = json.dumps(edge.to_dict())
+        parsed = json.loads(json_str)
+        assert parsed["source_id"] == "source"
+        assert parsed["target_id"] == "target"
+        assert "condition_name" not in parsed or parsed["condition_name"] is None
+
+    def test_edge_serialization_with_named_condition(self) -> None:
+        """Test that Edge with named function condition serializes condition_name correctly."""
+
+        def is_positive(x: int) -> bool:
+            return x > 0
+
+        edge = Edge(source_id="source", target_id="target", condition=is_positive)
+
+        # Test to_dict
+        data = edge.to_dict()
+        assert data["source_id"] == "source"
+        assert data["target_id"] == "target"
+        assert data["condition_name"] == "is_positive"
+
+        # Test model_dump_json
+        json_str = json.dumps(edge.to_dict())
+        parsed = json.loads(json_str)
+        assert parsed["source_id"] == "source"
+        assert parsed["target_id"] == "target"
+        assert parsed["condition_name"] == "is_positive"
+
+    def test_edge_serialization_with_lambda_condition(self) -> None:
+        """Test that Edge with lambda condition serializes condition_name as '<lambda>'."""
+        edge = Edge(source_id="source", target_id="target", condition=lambda x: x > 0)
+
+        # Test to_dict
+        data = edge.to_dict()
+        assert data["source_id"] == "source"
+        assert data["target_id"] == "target"
+        assert data["condition_name"] == "<lambda>"
+
+        # Test model_dump_json
+        json_str = json.dumps(edge.to_dict())
+        parsed = json.loads(json_str)
+        assert parsed["source_id"] == "source"
+        assert parsed["target_id"] == "target"
+        assert parsed["condition_name"] == "<lambda>"
+
+    def test_single_edge_group_serialization(self) -> None:
+        """Test that SingleEdgeGroup can be serialized and has correct fields, including edges and type."""
+        edge_group = SingleEdgeGroup(source_id="source", target_id="target")
+
+        # Test to_dict
+        data = edge_group.to_dict()
+        assert "id" in data
+        assert data["id"].startswith("SingleEdgeGroup/")
+
+        # Test type field
+        assert "type" in data, "SingleEdgeGroup should have 'type' field"
+        assert data["type"] == "SingleEdgeGroup", f"Expected type 'SingleEdgeGroup', got {data['type']}"
+
+        # Verify edges field is present and contains the edge
+        assert "edges" in data, "SingleEdgeGroup should have 'edges' field"
+        assert len(data["edges"]) == 1, "SingleEdgeGroup should have exactly one edge"
+        edge = data["edges"][0]
+        assert "source_id" in edge, "Edge should have source_id"
+        assert "target_id" in edge, "Edge should have target_id"
+        assert edge["source_id"] == "source", f"Expected source_id 'source', got {edge['source_id']}"
+        assert edge["target_id"] == "target", f"Expected target_id 'target', got {edge['target_id']}"
+
+        # Test model_dump_json
+        json_str = json.dumps(edge_group.to_dict())
+        parsed = json.loads(json_str)
+        assert "id" in parsed
+        assert parsed["id"].startswith("SingleEdgeGroup/")
+
+        # Test type field in JSON
+        assert "type" in parsed, "JSON should have 'type' field"
+        assert parsed["type"] == "SingleEdgeGroup", "JSON should preserve type field"
+
+        # Verify edges are preserved in JSON
+        assert "edges" in parsed, "JSON should have 'edges' field"
+        assert len(parsed["edges"]) == 1, "JSON should have exactly one edge"
+        json_edge = parsed["edges"][0]
+        assert json_edge["source_id"] == "source", "JSON should preserve edge source_id"
+        assert json_edge["target_id"] == "target", "JSON should preserve edge target_id"
+
+    def test_fan_out_edge_group_serialization(self) -> None:
+        """Test that FanOutEdgeGroup can be serialized and has correct fields, including edges and type."""
+        edge_group = FanOutEdgeGroup(source_id="source", target_ids=["target1", "target2"])
+
+        # Test to_dict
+        data = edge_group.to_dict()
+        assert "id" in data
+        assert data["id"].startswith("FanOutEdgeGroup/")
+
+        # Test type field
+        assert "type" in data, "FanOutEdgeGroup should have 'type' field"
+        assert data["type"] == "FanOutEdgeGroup", f"Expected type 'FanOutEdgeGroup', got {data['type']}"
+
+        # Test selection_func_name field (should be None when no selection function is provided)
+        assert "selection_func_name" in data, "FanOutEdgeGroup should have 'selection_func_name' field"
+        assert data["selection_func_name"] is None, (
+            "selection_func_name should be None when no selection function is provided"
+        )
+
+        # Verify edges field is present and contains the correct edges
+        assert "edges" in data, "FanOutEdgeGroup should have 'edges' field"
+        assert len(data["edges"]) == 2, "FanOutEdgeGroup should have exactly two edges"
+
+        edges = data["edges"]
+        sources = [edge["source_id"] for edge in edges]
+        targets = [edge["target_id"] for edge in edges]
+
+        assert all(source == "source" for source in sources), f"All edges should have source 'source', got {sources}"
+        assert set(targets) == {"target1", "target2"}, f"Expected targets {{'target1', 'target2'}}, got {set(targets)}"
+
+        # Test model_dump_json
+        json_str = json.dumps(edge_group.to_dict())
+        parsed = json.loads(json_str)
+        assert "id" in parsed
+        assert parsed["id"].startswith("FanOutEdgeGroup/")
+
+        # Test type field in JSON
+        assert "type" in parsed, "JSON should have 'type' field"
+        assert parsed["type"] == "FanOutEdgeGroup", "JSON should preserve type field"
+
+        # Test selection_func_name field in JSON
+        assert "selection_func_name" in parsed, "JSON should have 'selection_func_name' field"
+        assert parsed["selection_func_name"] is None, (
+            "JSON selection_func_name should be None when no selection function is provided"
+        )
+
+        # Verify edges are preserved in JSON
+        assert "edges" in parsed, "JSON should have 'edges' field"
+        assert len(parsed["edges"]) == 2, "JSON should have exactly two edges"
+        json_edges = parsed["edges"]
+        json_sources = [edge["source_id"] for edge in json_edges]
+        json_targets = [edge["target_id"] for edge in json_edges]
+
+        assert all(source == "source" for source in json_sources), "JSON should preserve edge sources"
+        assert set(json_targets) == {"target1", "target2"}, "JSON should preserve edge targets"
+
+    def test_fan_out_edge_group_serialization_with_selection_func(self) -> None:
+        """Test that FanOutEdgeGroup with named selection function serializes selection_func_name correctly."""
+
+        def custom_selector(data: Any, targets: list[str]) -> list[str]:
+            """Custom selection function for testing."""
+            return targets[:1]  # Select only the first target
+
+        edge_group = FanOutEdgeGroup(
+            source_id="source", target_ids=["target1", "target2"], selection_func=custom_selector
+        )
+
+        # Test to_dict
+        data = edge_group.to_dict()
+        assert "selection_func_name" in data, "FanOutEdgeGroup should have 'selection_func_name' field"
+        assert data["selection_func_name"] == "custom_selector", (
+            f"Expected selection_func_name 'custom_selector', got {data['selection_func_name']}"
+        )
+
+        # Test model_dump_json
+        json_str = json.dumps(edge_group.to_dict())
+        parsed = json.loads(json_str)
+        assert "selection_func_name" in parsed, "JSON should have 'selection_func_name' field"
+        assert parsed["selection_func_name"] == "custom_selector", "JSON should preserve selection_func_name"
+
+    def test_fan_out_edge_group_serialization_with_lambda_selection_func(self) -> None:
+        """Test that FanOutEdgeGroup with lambda selection function serializes selection_func_name as '<lambda>'."""
+        edge_group = FanOutEdgeGroup(
+            source_id="source", target_ids=["target1", "target2"], selection_func=lambda data, targets: targets[:1]
+        )
+
+        # Test to_dict
+        data = edge_group.to_dict()
+        assert "selection_func_name" in data, "FanOutEdgeGroup should have 'selection_func_name' field"
+        assert data["selection_func_name"] == "<lambda>", (
+            f"Expected selection_func_name '<lambda>', got {data['selection_func_name']}"
+        )
+
+        # Test model_dump_json
+        json_str = json.dumps(edge_group.to_dict())
+        parsed = json.loads(json_str)
+        assert "selection_func_name" in parsed, "JSON should have 'selection_func_name' field"
+        assert parsed["selection_func_name"] == "<lambda>", "JSON should preserve selection_func_name as '<lambda>'"
+
+    def test_fan_in_edge_group_serialization(self) -> None:
+        """Test that FanInEdgeGroup can be serialized and has correct fields, including edges and type."""
+        edge_group = FanInEdgeGroup(source_ids=["source1", "source2"], target_id="target")
+
+        # Test to_dict
+        data = edge_group.to_dict()
+        assert "id" in data
+        assert data["id"].startswith("FanInEdgeGroup/")
+
+        # Test type field
+        assert "type" in data, "FanInEdgeGroup should have 'type' field"
+        assert data["type"] == "FanInEdgeGroup", f"Expected type 'FanInEdgeGroup', got {data['type']}"
+
+        # Verify edges field is present and contains the correct edges
+        assert "edges" in data, "FanInEdgeGroup should have 'edges' field"
+        assert len(data["edges"]) == 2, "FanInEdgeGroup should have exactly two edges"
+
+        edges = data["edges"]
+        sources = [edge["source_id"] for edge in edges]
+        targets = [edge["target_id"] for edge in edges]
+
+        assert set(sources) == {"source1", "source2"}, f"Expected sources {{'source1', 'source2'}}, got {set(sources)}"
+        assert all(target == "target" for target in targets), f"All edges should have target 'target', got {targets}"
+
+        # Test model_dump_json
+        json_str = json.dumps(edge_group.to_dict())
+        parsed = json.loads(json_str)
+        assert "id" in parsed
+        assert parsed["id"].startswith("FanInEdgeGroup/")
+
+        # Test type field in JSON
+        assert "type" in parsed, "JSON should have 'type' field"
+        assert parsed["type"] == "FanInEdgeGroup", "JSON should preserve type field"
+
+        # Verify edges are preserved in JSON
+        assert "edges" in parsed, "JSON should have 'edges' field"
+        assert len(parsed["edges"]) == 2, "JSON should have exactly two edges"
+        json_edges = parsed["edges"]
+        json_sources = [edge["source_id"] for edge in json_edges]
+        json_targets = [edge["target_id"] for edge in json_edges]
+
+        assert set(json_sources) == {"source1", "source2"}, "JSON should preserve edge sources"
+        assert all(target == "target" for target in json_targets), "JSON should preserve edge targets"
+
+    def test_switch_case_edge_group_serialization(self) -> None:
+        """Test that SwitchCaseEdgeGroup can be serialized and has correct fields, including edges and type."""
+        cases = [
+            SwitchCaseEdgeGroupCase(condition=lambda x: x > 0, target_id="positive"),
+            SwitchCaseEdgeGroupDefault(target_id="default"),
+        ]
+        edge_group = SwitchCaseEdgeGroup(source_id="source", cases=cases)  # type: ignore[arg-type]
+
+        # Test to_dict
+        data = edge_group.to_dict()
+        assert "id" in data
+        assert data["id"].startswith("SwitchCaseEdgeGroup/")
+
+        # Test type field
+        assert "type" in data, "SwitchCaseEdgeGroup should have 'type' field"
+        assert data["type"] == "SwitchCaseEdgeGroup", f"Expected type 'SwitchCaseEdgeGroup', got {data['type']}"
+
+        # Test cases field
+        assert "cases" in data, "SwitchCaseEdgeGroup should have 'cases' field"
+        assert len(data["cases"]) == 2, "SwitchCaseEdgeGroup should have exactly two cases"
+
+        cases_data = data["cases"]
+        # Check first case (SwitchCaseEdgeGroupCase)
+        case_obj = cases_data[0]
+        assert "target_id" in case_obj, "SwitchCaseEdgeGroupCase should have 'target_id' field"
+        assert "condition_name" in case_obj, "SwitchCaseEdgeGroupCase should have 'condition_name' field"
+        assert "type" in case_obj, "SwitchCaseEdgeGroupCase should have 'type' field"
+        assert case_obj["target_id"] == "positive", f"Expected target_id 'positive', got {case_obj['target_id']}"
+        assert case_obj["condition_name"] == "<lambda>", (
+            f"Expected condition_name '<lambda>', got {case_obj['condition_name']}"
+        )
+        assert case_obj["type"] == "Case", f"Expected type 'Case', got {case_obj['type']}"
+
+        # Check default case (SwitchCaseEdgeGroupDefault)
+        default_obj = cases_data[1]
+        assert "target_id" in default_obj, "SwitchCaseEdgeGroupDefault should have 'target_id' field"
+        assert "type" in default_obj, "SwitchCaseEdgeGroupDefault should have 'type' field"
+        assert default_obj["target_id"] == "default", f"Expected target_id 'default', got {default_obj['target_id']}"
+        assert default_obj["type"] == "Default", f"Expected type 'Default', got {default_obj['type']}"
+
+        # Verify edges field is present and contains the correct edges
+        assert "edges" in data, "SwitchCaseEdgeGroup should have 'edges' field"
+        assert len(data["edges"]) == 2, "SwitchCaseEdgeGroup should have exactly two edges"
+
+        edges = data["edges"]
+        sources = [edge["source_id"] for edge in edges]
+        targets = [edge["target_id"] for edge in edges]
+
+        assert all(source == "source" for source in sources), f"All edges should have source 'source', got {sources}"
+        assert set(targets) == {"positive", "default"}, (
+            f"Expected targets {{'positive', 'default'}}, got {set(targets)}"
+        )
+
+        # Check condition_name field in edges - SwitchCaseEdgeGroup edges don't have conditions
+        # because the conditional logic is implemented in the selection_func at the group level
+        condition_names = [edge.get("condition_name") for edge in edges]
+        assert all(name is None for name in condition_names), (
+            "SwitchCaseEdgeGroup edges should not have condition_name since conditions are handled at group level"
+        )
+
+        # Test model_dump_json
+        json_str = json.dumps(edge_group.to_dict())
+        parsed = json.loads(json_str)
+        assert "id" in parsed
+        assert parsed["id"].startswith("SwitchCaseEdgeGroup/")
+
+        # Test type field in JSON
+        assert "type" in parsed, "JSON should have 'type' field"
+        assert parsed["type"] == "SwitchCaseEdgeGroup", "JSON should preserve type field"
+
+        # Test cases field in JSON
+        assert "cases" in parsed, "JSON should have 'cases' field"
+        assert len(parsed["cases"]) == 2, "JSON should have exactly two cases"
+
+        json_cases = parsed["cases"]
+        json_case_obj = json_cases[0]
+        assert json_case_obj["target_id"] == "positive", "JSON should preserve case target_id"
+        assert json_case_obj["condition_name"] == "<lambda>", "JSON should preserve case condition_name"
+        assert json_case_obj["type"] == "Case", "JSON should preserve case type"
+
+        json_default_obj = json_cases[1]
+        assert json_default_obj["target_id"] == "default", "JSON should preserve default target_id"
+        assert json_default_obj["type"] == "Default", "JSON should preserve default type"
+
+        # Verify edges are preserved in JSON
+        assert "edges" in parsed, "JSON should have 'edges' field"
+        assert len(parsed["edges"]) == 2, "JSON should have exactly two edges"
+        json_edges = parsed["edges"]
+        json_sources = [edge["source_id"] for edge in json_edges]
+        json_targets = [edge["target_id"] for edge in json_edges]
+
+        assert all(source == "source" for source in json_sources), "JSON should preserve edge sources"
+        assert set(json_targets) == {"positive", "default"}, "JSON should preserve edge targets"
+
+        # Check condition_name field in JSON edges - should be None for SwitchCaseEdgeGroup
+        json_condition_names = [edge.get("condition_name") for edge in json_edges]
+        assert all(name is None for name in json_condition_names), (
+            "JSON SwitchCaseEdgeGroup edges should not have condition_name"
+        )
+
+    def test_nested_workflow_executor_serialization(self) -> None:
+        """Test complete serialization of deeply nested WorkflowExecutors (subworkflows within subworkflows).
+
+        This test verifies that nested WorkflowExecutor objects are fully serialized with their
+        complete workflow structures, including deeply nested workflows and all their executors.
+        """
+        # Create innermost workflow
+        inner_executor = SampleExecutor(id="inner-exec")
+        inner_workflow = WorkflowBuilder(max_iterations=10, start_executor=inner_executor).build()
+
+        # Create middle workflow with WorkflowExecutor
+        inner_workflow_executor = WorkflowExecutor(workflow=inner_workflow, id="inner-workflow-exec")
+        middle_executor = SampleExecutor(id="middle-exec")
+        middle_workflow = (
+            WorkflowBuilder(max_iterations=20, start_executor=middle_executor)
+            .add_edge(middle_executor, inner_workflow_executor)
+            .build()
+        )
+
+        # Create outer workflow with nested WorkflowExecutor
+        middle_workflow_executor = WorkflowExecutor(workflow=middle_workflow, id="middle-workflow-exec")
+        outer_executor = SampleExecutor(id="outer-exec")
+        outer_workflow = (
+            WorkflowBuilder(max_iterations=30, start_executor=outer_executor)
+            .add_edge(outer_executor, middle_workflow_executor)
+            .build()
+        )
+
+        # Test serialization of the nested structure
+        data = outer_workflow.to_dict()
+
+        # Verify outer structure
+        assert data["start_executor_id"] == "outer-exec"
+        assert data["max_iterations"] == 30
+        assert "outer-exec" in data["executors"]
+        assert "middle-workflow-exec" in data["executors"]
+
+        # Verify middle WorkflowExecutor is present with full nested workflow serialization
+        middle_exec_data = data["executors"]["middle-workflow-exec"]
+        assert middle_exec_data["type"] == "WorkflowExecutor"
+        assert middle_exec_data["id"] == "middle-workflow-exec"
+
+        # Verify the nested workflow is fully serialized
+        assert "workflow" in middle_exec_data, "WorkflowExecutor should include nested workflow in serialization"
+        middle_workflow_data = middle_exec_data["workflow"]
+        assert "start_executor_id" in middle_workflow_data
+        assert "executors" in middle_workflow_data
+        assert "max_iterations" in middle_workflow_data
+        assert middle_workflow_data["start_executor_id"] == "middle-exec"
+        assert middle_workflow_data["max_iterations"] == 20
+
+        # Verify the deeply nested executors are present
+        assert "middle-exec" in middle_workflow_data["executors"]
+        assert "inner-workflow-exec" in middle_workflow_data["executors"]
+
+        # Verify the innermost WorkflowExecutor is also fully serialized
+        inner_workflow_exec_data = middle_workflow_data["executors"]["inner-workflow-exec"]
+        assert inner_workflow_exec_data["type"] == "WorkflowExecutor"
+        assert "workflow" in inner_workflow_exec_data, "Deeply nested WorkflowExecutor should also include its workflow"
+        innermost_workflow_data = inner_workflow_exec_data["workflow"]
+        assert "start_executor_id" in innermost_workflow_data
+        assert "executors" in innermost_workflow_data
+        assert "max_iterations" in innermost_workflow_data
+        assert innermost_workflow_data["start_executor_id"] == "inner-exec"
+        assert innermost_workflow_data["max_iterations"] == 10
+        assert "inner-exec" in innermost_workflow_data["executors"]
+
+        # Test JSON serialization preserves the complete nested structure
+        json_str = outer_workflow.to_json()
+        parsed = json.loads(json_str)
+
+        # Verify the complete structure is preserved in JSON
+        middle_exec_json = parsed["executors"]["middle-workflow-exec"]
+        assert middle_exec_json["type"] == "WorkflowExecutor"
+        assert middle_exec_json["id"] == "middle-workflow-exec"
+
+        # Verify nested workflow is present in JSON
+        assert "workflow" in middle_exec_json, "JSON serialization should include nested workflow"
+        middle_workflow_json = middle_exec_json["workflow"]
+        assert middle_workflow_json["start_executor_id"] == "middle-exec"
+        assert middle_workflow_json["max_iterations"] == 20
+        assert "middle-exec" in middle_workflow_json["executors"]
+        assert "inner-workflow-exec" in middle_workflow_json["executors"]
+
+        # Verify deeply nested structure in JSON
+        inner_workflow_exec_json = middle_workflow_json["executors"]["inner-workflow-exec"]
+        assert inner_workflow_exec_json["type"] == "WorkflowExecutor"
+        assert "workflow" in inner_workflow_exec_json, "Deeply nested WorkflowExecutor should be in JSON"
+        innermost_workflow_json = inner_workflow_exec_json["workflow"]
+        assert innermost_workflow_json["start_executor_id"] == "inner-exec"
+        assert innermost_workflow_json["max_iterations"] == 10
+        assert "inner-exec" in innermost_workflow_json["executors"]
+
+        # Test that WorkflowExecutor also serializes correctly when accessed directly
+        direct_middle_data = middle_workflow_executor.to_dict()
+        assert "workflow" in direct_middle_data
+        assert direct_middle_data["type"] == "WorkflowExecutor"
+        assert "executors" in direct_middle_data["workflow"]
+        assert "inner-workflow-exec" in direct_middle_data["workflow"]["executors"]
+
+    def test_switch_case_edge_group_serialization_with_named_condition(self) -> None:
+        """Test that SwitchCaseEdgeGroup with named condition function serializes condition_name correctly."""
+
+        def is_positive(x: int) -> bool:
+            return x > 0
+
+        cases = [
+            SwitchCaseEdgeGroupCase(condition=is_positive, target_id="positive"),
+            SwitchCaseEdgeGroupDefault(target_id="default"),
+        ]
+        edge_group = SwitchCaseEdgeGroup(source_id="source", cases=cases)  # type: ignore[arg-type]
+
+        # Test to_dict
+        data = edge_group.to_dict()
+        assert "cases" in data, "SwitchCaseEdgeGroup should have 'cases' field"
+
+        cases_data = data["cases"]
+        case_obj = cases_data[0]
+        assert case_obj["condition_name"] == "is_positive", (
+            f"Expected condition_name 'is_positive', got {case_obj['condition_name']}"
+        )
+
+        # Test model_dump_json
+        json_str = json.dumps(edge_group.to_dict())
+        parsed = json.loads(json_str)
+        json_cases = parsed["cases"]
+        json_case_obj = json_cases[0]
+        assert json_case_obj["condition_name"] == "is_positive", "JSON should preserve named condition_name"
+
+    def test_workflow_serialization(self) -> None:
+        """Test that Workflow can be serialized and has correct fields, including edges."""
+        executor1 = SampleExecutor(id="executor1")
+        executor2 = SampleExecutor(id="executor2")
+
+        workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
+
+        # Test model_dump
+        data = workflow.to_dict()
+        assert "edge_groups" in data
+        assert "executors" in data
+        assert "start_executor_id" in data
+        assert "max_iterations" in data
+        assert "id" in data
+
+        assert data["start_executor_id"] == "executor1"
+        assert "executor1" in data["executors"]
+        assert "executor2" in data["executors"]
+
+        # Verify edge groups contain edges
+        edge_groups = data["edge_groups"]
+
+        single_edge_groups = [SingleEdgeGroup.from_dict(eg) for eg in edge_groups if eg["type"] == "SingleEdgeGroup"]
+        internal_edge_groups = [
+            InternalEdgeGroup.from_dict(eg) for eg in edge_groups if eg["type"] == "InternalEdgeGroup"
+        ]
+
+        assert len(single_edge_groups) == 1, "Should have exactly one SingleEdgeGroup for the added edge"
+        assert len(internal_edge_groups) == 2, (
+            "Should have exactly two (one per executor) InternalEdgeGroups for request/response handling"
+        )
+
+        for edge_group in single_edge_groups:
+            assert len(edge_group.edges) == 1, "Should have exactly one edge"
+
+            edge = edge_group.edges[0]
+
+            assert edge.source_id == "executor1", f"Expected source_id 'executor1', got {edge.source_id}"
+            assert edge.target_id == "executor2", f"Expected target_id 'executor2', got {edge.target_id}"
+
+        for edge_group in internal_edge_groups:
+            assert len(edge_group.edges) == 1, "Each InternalEdgeGroup should have exactly one edge"
+
+            edge = edge_group.edges[0]
+
+            assert edge.source_id == INTERNAL_SOURCE_ID(edge.target_id)
+            assert edge.target_id in [executor1.id, executor2.id]
+
+        # Test model_dump_json
+        json_str = workflow.to_json()
+        parsed = json.loads(json_str)
+        assert parsed["start_executor_id"] == "executor1"
+        assert "executor1" in parsed["executors"]
+        assert "executor2" in parsed["executors"]
+
+        # Verify edges are preserved in JSON serialization
+        json_edge_groups = parsed["edge_groups"]
+        assert len(json_edge_groups) == 1 + 2, "JSON should have exactly one SingleEdgeGroup and two InternalEdgeGroups"
+
+        for json_edge_group in json_edge_groups:
+            assert "edges" in json_edge_group, "JSON edge group should contain 'edges' field"
+            assert len(json_edge_group["edges"]) == 1, "Each JSON edge group should have exactly one edge"
+            if json_edge_group["type"] == "SingleEdgeGroup":
+                json_edge = json_edge_group["edges"][0]
+                assert json_edge["source_id"] == "executor1", "JSON should preserve edge source_id"
+                assert json_edge["target_id"] == "executor2", "JSON should preserve edge target_id"
+            elif json_edge_group["type"] == "InternalEdgeGroup":
+                json_edge = json_edge_group["edges"][0]
+                assert json_edge["source_id"] == INTERNAL_SOURCE_ID(json_edge["target_id"])
+                assert json_edge["target_id"] in [executor1.id, executor2.id]
+            else:
+                pytest.fail(f"Unexpected edge group type: {json_edge_group['type']}")
+
+    def test_workflow_serialization_excludes_non_serializable_fields(self) -> None:
+        """Test that non-serializable fields are excluded from serialization."""
+        executor1 = SampleExecutor(id="executor1")
+        executor2 = SampleExecutor(id="executor2")
+
+        workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
+
+        # Test model_dump - should not include private runtime objects
+        data = workflow.to_dict()
+
+        # These private runtime fields should not be in the serialized data
+        assert "_runner_context" not in data
+        assert "_state" not in data
+        assert "_runner" not in data
+
+    def test_workflow_name_description_serialization(self) -> None:
+        """Test that workflow name and description are serialized correctly."""
+        # Test 1: With name and description
+        workflow1 = WorkflowBuilder(
+            name="Test Pipeline",
+            description="Test workflow description",
+            start_executor=SampleExecutor(id="e1"),
+        ).build()
+
+        assert workflow1.name == "Test Pipeline"
+        assert workflow1.description == "Test workflow description"
+
+        data1 = workflow1.to_dict()
+        assert data1["name"] == "Test Pipeline"
+        assert data1["description"] == "Test workflow description"
+
+        # Test JSON serialization
+        json_str1 = workflow1.to_json()
+        parsed1 = json.loads(json_str1)
+        assert parsed1["name"] == "Test Pipeline"
+        assert parsed1["description"] == "Test workflow description"
+
+        # Test 2: Without name and description (defaults)
+        workflow2 = WorkflowBuilder(start_executor=SampleExecutor(id="e2")).build()
+
+        assert workflow2.name is not None
+        assert workflow2.description is None
+
+        data2 = workflow2.to_dict()
+        assert "description" not in data2  # Should not include None values
+
+        # Test 3: With only name (no description)
+        workflow3 = WorkflowBuilder(name="Named Only", start_executor=SampleExecutor(id="e3")).build()
+
+        assert workflow3.name == "Named Only"
+        assert workflow3.description is None
+
+        data3 = workflow3.to_dict()
+        assert data3["name"] == "Named Only"
+        assert "description" not in data3
+
+    def test_executor_field_validation(self) -> None:
+        """Test that Executor field validation works correctly."""
+        # Valid executor
+        executor = SampleExecutor(id="valid-id")
+        assert executor.id == "valid-id"
+
+        with pytest.raises(ValueError):
+            SampleExecutor(id="")
+
+    def test_edge_field_validation(self) -> None:
+        """Test that Edge field validation works correctly."""
+        # Valid edge
+        edge = Edge(source_id="source", target_id="target")
+        assert edge.source_id == "source"
+        assert edge.target_id == "target"
+
+        # Test validation failure for empty source_id
+        with pytest.raises(ValueError):
+            Edge(source_id="", target_id="target")
+
+        # Test validation failure for empty target_id
+        with pytest.raises(ValueError):
+            Edge(source_id="source", target_id="")
+
+
+def test_comprehensive_edge_groups_workflow_serialization() -> None:
+    """Test serialization of a workflow that uses all edge group types: SwitchCase, FanOut, and FanIn."""
+    # Create executors for a comprehensive workflow
+    router = SampleExecutor(id="router")
+    processor_a = SampleExecutor(id="proc_a")
+    processor_b = SampleExecutor(id="proc_b")
+    fanout_hub = SampleExecutor(id="fanout_hub")
+    parallel_1 = SampleExecutor(id="parallel_1")
+    parallel_2 = SampleExecutor(id="parallel_2")
+    aggregator = SampleAggregator(id="aggregator")
+
+    # Build workflow with all three edge group types
+    workflow = (
+        WorkflowBuilder(start_executor=router)
+        # 1. SwitchCaseEdgeGroup: Conditional routing
+        .add_switch_case_edge_group(
+            router,
+            [
+                Case(condition=lambda msg: len(str(msg)) < 10, target=processor_a),
+                Default(target=processor_b),
+            ],
+        )
+        # 2. Direct edges
+        .add_edge(processor_a, fanout_hub)
+        .add_edge(processor_b, fanout_hub)
+        # 3. FanOutEdgeGroup: One-to-many distribution
+        .add_fan_out_edges(fanout_hub, [parallel_1, parallel_2])
+        # 4. FanInEdgeGroup: Many-to-one aggregation
+        .add_fan_in_edges([parallel_1, parallel_2], aggregator)
+        .build()
+    )
+
+    # Test workflow serialization
+    data = workflow.to_dict()
+
+    # Verify basic workflow structure
+    assert "edge_groups" in data
+    assert "executors" in data
+    assert "start_executor_id" in data
+    assert data["start_executor_id"] == "router"
+
+    # Verify all executors are present
+    expected_executors = {"router", "proc_a", "proc_b", "fanout_hub", "parallel_1", "parallel_2", "aggregator"}
+    assert set(data["executors"].keys()) == expected_executors
+
+    # Verify edge groups contain all three types
+    edge_groups = data["edge_groups"]
+    edge_group_types = [eg.get("id", "").split("/")[0] for eg in edge_groups]
+
+    # Should have: SwitchCaseEdgeGroup, SingleEdgeGroup (x2), FanOutEdgeGroup, FanInEdgeGroup
+    assert "SwitchCaseEdgeGroup" in edge_group_types, f"Expected SwitchCaseEdgeGroup in {edge_group_types}"
+    assert "FanOutEdgeGroup" in edge_group_types, f"Expected FanOutEdgeGroup in {edge_group_types}"
+    assert "FanInEdgeGroup" in edge_group_types, f"Expected FanInEdgeGroup in {edge_group_types}"
+    assert "SingleEdgeGroup" in edge_group_types, f"Expected SingleEdgeGroup in {edge_group_types}"
+
+    # Test JSON serialization
+    json_str = workflow.to_json()
+    parsed = json.loads(json_str)
+
+    # Verify JSON structure matches model_dump
+    assert parsed["start_executor_id"] == "router"
+    assert set(parsed["executors"].keys()) == expected_executors
+    assert len(parsed["edge_groups"]) == len(edge_groups)
+
+    # Verify that serialization excludes non-serializable fields
+    assert "_runner_context" not in data
+    assert "_state" not in data
+    assert "_runner" not in data
+
+    # Test that we can identify each edge group type by examining their structure
+    switch_case_groups = [eg for eg in edge_groups if eg.get("id", "").startswith("SwitchCaseEdgeGroup/")]
+    fan_out_groups = [eg for eg in edge_groups if eg.get("id", "").startswith("FanOutEdgeGroup/")]
+    fan_in_groups = [eg for eg in edge_groups if eg.get("id", "").startswith("FanInEdgeGroup/")]
+    single_groups = [eg for eg in edge_groups if eg.get("id", "").startswith("SingleEdgeGroup/")]
+
+    assert len(switch_case_groups) == 1, f"Expected 1 SwitchCaseEdgeGroup, got {len(switch_case_groups)}"
+    assert len(fan_out_groups) == 1, f"Expected 1 FanOutEdgeGroup, got {len(fan_out_groups)}"
+    assert len(fan_in_groups) == 1, f"Expected 1 FanInEdgeGroup, got {len(fan_in_groups)}"
+    assert len(single_groups) == 2, f"Expected 2 SingleEdgeGroups, got {len(single_groups)}"
+
+    # The key validation is that all edge group types are present and serializable
+    # Individual edge group fields may vary based on implementation,
+    # but each should have at least an 'id' field that identifies its type and 'edges' field
+    for group_type, groups in [
+        ("SwitchCaseEdgeGroup", switch_case_groups),
+        ("FanOutEdgeGroup", fan_out_groups),
+        ("FanInEdgeGroup", fan_in_groups),
+        ("SingleEdgeGroup", single_groups),
+    ]:
+        for group in groups:
+            assert "id" in group, f"{group_type} should have 'id' field"
+            assert group["id"].startswith(f"{group_type}/"), f"{group_type} id should start with '{group_type}/'"
+            assert "edges" in group, f"{group_type} should have 'edges' field"
+            assert isinstance(group["edges"], list), f"{group_type} 'edges' should be a list"
+            assert len(group["edges"]) > 0, f"{group_type} should have at least one edge"
+
+            # Verify each edge has required fields
+            for edge in group["edges"]:
+                assert "source_id" in edge, f"{group_type} edge should have 'source_id'"
+                assert "target_id" in edge, f"{group_type} edge should have 'target_id'"
+                assert isinstance(edge["source_id"], str), f"{group_type} edge source_id should be string"
+                assert isinstance(edge["target_id"], str), f"{group_type} edge target_id should be string"
+                assert len(edge["source_id"]) > 0, f"{group_type} edge source_id should not be empty"
+                assert len(edge["target_id"]) > 0, f"{group_type} edge target_id should not be empty"
+
+    # Verify specific edge group edge counts
+    assert len(switch_case_groups[0]["edges"]) == 2, "SwitchCaseEdgeGroup should have 2 edges (proc_a and proc_b)"
+    assert len(fan_out_groups[0]["edges"]) == 2, "FanOutEdgeGroup should have 2 edges (parallel_1 and parallel_2)"
+    assert len(fan_in_groups[0]["edges"]) == 2, "FanInEdgeGroup should have 2 edges (from parallel_1 and parallel_2)"
+    for single_group in single_groups:
+        assert len(single_group["edges"]) == 1, "Each SingleEdgeGroup should have exactly 1 edge"
+
+
+def test_to_dict_preserves_compatibility_wire_keys_for_output_designation() -> None:
+    """to_dict() must emit the compatibility wire keys regardless of the Python kwarg names.
+
+    The Python API renamed ``output_executors`` -> ``output_from`` and
+    uses ``intermediate_output_from`` for intermediate selection, but the serialized
+    dict must keep the old keys so existing checkpoints stay readable. This is a
+    regression guard against accidental renames of the wire format.
+    """
+
+    class _Yielder(Executor):
+        @handler
+        async def handle(self, message: str, ctx: WorkflowContext[str, str]) -> None:
+            await ctx.yield_output(message)
+            await ctx.send_message(message)
+
+    class _Terminal(Executor):
+        @handler
+        async def handle(self, message: str, ctx: WorkflowContext[str, str]) -> None:
+            await ctx.yield_output(f"final: {message}")
+
+    start = _Yielder(id="start")
+    progress = _Yielder(id="progress")
+    final = _Terminal(id="final")
+
+    workflow = (
+        WorkflowBuilder(
+            start_executor=start,
+            output_from=[final],
+            intermediate_output_from=[progress],
+        )
+        .add_edge(start, progress)
+        .add_edge(progress, final)
+        .build()
+    )
+
+    d = workflow.to_dict()
+
+    assert "output_executors" in d, "wire key 'output_executors' must be preserved"
+    assert "intermediate_executors" in d, "wire key 'intermediate_executors' must be preserved"
+    assert "output_from" not in d, "new Python kwarg name must NOT leak into the wire format"
+    assert "intermediate_output_from" not in d, "new Python kwarg name must NOT leak into the wire format"
+    assert d["output_executors"] == ["final"]
+    assert d["intermediate_executors"] == ["progress"]

python/packages/core/tests/workflow/test_state.py

@@ -0,0 +1,303 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Unit tests for the State class superstep caching behavior."""
+
+import pytest
+
+from agent_framework._workflows._state import State
+
+
+class TestStateBasicOperations:
+    """Tests for basic State get/set/has/delete operations."""
+
+    def test_set_and_get(self) -> None:
+        state = State()
+        state.set("key", "value")
+        assert state.get("key") == "value"
+
+    def test_get_with_default(self) -> None:
+        state = State()
+        assert state.get("missing") is None
+        assert state.get("missing", "default") == "default"
+
+    def test_has_returns_true_for_existing_key(self) -> None:
+        state = State()
+        state.set("key", "value")
+        assert state.has("key") is True
+
+    def test_has_returns_false_for_missing_key(self) -> None:
+        state = State()
+        assert state.has("missing") is False
+
+    def test_delete_existing_key(self) -> None:
+        state = State()
+        state.set("key", "value")
+        state.commit()
+        state.delete("key")
+        state.commit()
+        assert state.has("key") is False
+        assert state.get("key") is None
+
+    def test_delete_missing_key_raises(self) -> None:
+        state = State()
+        with pytest.raises(KeyError, match="Key 'missing' not found"):
+            state.delete("missing")
+
+    def test_clear(self) -> None:
+        state = State()
+        state.set("key1", "value1")
+        state.commit()
+        state.set("key2", "value2")
+        state.clear()
+        assert state.get("key1") is None
+        assert state.get("key2") is None
+
+
+class TestSuperstepCaching:
+    """Tests for superstep caching semantics - pending vs committed state."""
+
+    def test_set_writes_to_pending_not_committed(self) -> None:
+        state = State()
+        state.set("key", "value")
+
+        # Value is in pending
+        assert "key" in state._pending  # pyright: ignore[reportPrivateUsage]
+        # Value is NOT in committed
+        assert "key" not in state._committed  # pyright: ignore[reportPrivateUsage]
+        # But get() still returns it
+        assert state.get("key") == "value"
+
+    def test_commit_moves_pending_to_committed(self) -> None:
+        state = State()
+        state.set("key", "value")
+
+        # Before commit: in pending, not committed
+        assert "key" in state._pending  # pyright: ignore[reportPrivateUsage]
+        assert "key" not in state._committed  # pyright: ignore[reportPrivateUsage]
+
+        state.commit()
+
+        # After commit: in committed, pending cleared
+        assert "key" not in state._pending  # pyright: ignore[reportPrivateUsage]
+        assert "key" in state._committed  # pyright: ignore[reportPrivateUsage]
+        assert state.get("key") == "value"
+
+    def test_discard_clears_pending_without_committing(self) -> None:
+        state = State()
+        state.set("existing", "original")
+        state.commit()
+
+        # Make a pending change
+        state.set("existing", "modified")
+        state.set("new_key", "new_value")
+
+        # Discard pending changes
+        state.discard()
+
+        # Original value is preserved, new key never committed
+        assert state.get("existing") == "original"
+        assert state.get("new_key") is None
+
+    def test_pending_overrides_committed_on_get(self) -> None:
+        state = State()
+        state.set("key", "committed_value")
+        state.commit()
+
+        state.set("key", "pending_value")
+
+        # get() returns pending value, not committed
+        assert state.get("key") == "pending_value"
+        # But committed still has old value
+        assert state._committed["key"] == "committed_value"  # pyright: ignore[reportPrivateUsage]
+
+    def test_multiple_sets_before_commit(self) -> None:
+        state = State()
+        state.set("key", "value1")
+        state.set("key", "value2")
+        state.set("key", "value3")
+
+        # Only final value is in pending
+        assert state.get("key") == "value3"
+
+        state.commit()
+        assert state.get("key") == "value3"
+
+
+class TestDeleteWithSuperstepCaching:
+    """Tests for delete behavior with superstep caching."""
+
+    def test_delete_pending_only_key(self) -> None:
+        state = State()
+        state.set("key", "value")
+        # Key only in pending, not committed
+        assert "key" in state._pending  # pyright: ignore[reportPrivateUsage]
+        assert "key" not in state._committed  # pyright: ignore[reportPrivateUsage]
+
+        state.delete("key")
+
+        # Should be removed from pending
+        assert "key" not in state._pending  # pyright: ignore[reportPrivateUsage]
+        assert state.get("key") is None
+        assert state.has("key") is False
+
+    def test_delete_committed_key_marks_for_deletion(self) -> None:
+        state = State()
+        state.set("key", "value")
+        state.commit()
+
+        state.delete("key")
+
+        # Key should be marked for deletion in pending (sentinel)
+        assert "key" in state._pending  # pyright: ignore[reportPrivateUsage]
+        # get() should return default (not the sentinel!)
+        assert state.get("key") is None
+        assert state.get("key", "default") == "default"
+        # has() should return False
+        assert state.has("key") is False
+        # But committed still has it until commit()
+        assert "key" in state._committed  # pyright: ignore[reportPrivateUsage]
+
+    def test_delete_committed_key_removed_on_commit(self) -> None:
+        state = State()
+        state.set("key", "value")
+        state.commit()
+
+        state.delete("key")
+        state.commit()
+
+        # Now it should be gone from committed too
+        assert "key" not in state._committed  # pyright: ignore[reportPrivateUsage]
+        assert "key" not in state._pending  # pyright: ignore[reportPrivateUsage]
+
+    def test_delete_key_in_both_pending_and_committed(self) -> None:
+        """Test delete when key exists in both pending (modified) and committed."""
+        state = State()
+        state.set("key", "original")
+        state.commit()
+
+        # Modify the key (now in both pending and committed)
+        state.set("key", "modified")
+        assert state._pending["key"] == "modified"  # pyright: ignore[reportPrivateUsage]
+        assert state._committed["key"] == "original"  # pyright: ignore[reportPrivateUsage]
+
+        # Delete should mark for deletion from committed
+        state.delete("key")
+
+        # Should be marked for deletion
+        assert state.get("key") is None
+        assert state.has("key") is False
+
+        # After commit, key should be fully removed
+        state.commit()
+        assert "key" not in state._committed  # pyright: ignore[reportPrivateUsage]
+        assert "key" not in state._pending  # pyright: ignore[reportPrivateUsage]
+
+    def test_discard_after_delete_restores_committed_value(self) -> None:
+        state = State()
+        state.set("key", "value")
+        state.commit()
+
+        state.delete("key")
+        # Key appears deleted
+        assert state.has("key") is False
+
+        state.discard()
+        # After discard, committed value is restored
+        assert state.has("key") is True
+        assert state.get("key") == "value"
+
+
+class TestFailureScenarios:
+    """Tests simulating failure scenarios - pending changes should not leak to committed."""
+
+    def test_failure_before_commit_preserves_committed_state(self) -> None:
+        """Simulate executor failure - pending changes should not affect committed state."""
+        state = State()
+        state.set("key1", "original1")
+        state.set("key2", "original2")
+        state.commit()
+
+        # Superstep starts - make some changes
+        state.set("key1", "modified1")
+        state.set("key3", "new_value")
+        state.delete("key2")
+
+        # Simulate failure - we call discard() instead of commit()
+        state.discard()
+
+        # All original values should be intact
+        assert state.get("key1") == "original1"
+        assert state.get("key2") == "original2"
+        assert state.get("key3") is None
+
+    def test_no_partial_commits(self) -> None:
+        """Ensure commit is atomic - either all changes apply or none."""
+        state = State()
+        state.set("key1", "value1")
+        state.set("key2", "value2")
+        state.set("key3", "value3")
+
+        # Before commit - nothing in committed
+        assert len(state._committed) == 0  # pyright: ignore[reportPrivateUsage]
+
+        state.commit()
+
+        # After commit - all three values committed together
+        assert state._committed == {"key1": "value1", "key2": "value2", "key3": "value3"}  # pyright: ignore[reportPrivateUsage]
+
+    def test_repeated_supersteps_are_isolated(self) -> None:
+        """Test that each superstep's changes are isolated until committed."""
+        state = State()
+
+        # Superstep 1
+        state.set("counter", 1)
+        state.commit()
+        assert state.get("counter") == 1
+
+        # Superstep 2
+        state.set("counter", 2)
+        state.set("temp", "should_be_discarded")
+        state.discard()  # Simulate failure
+        assert state.get("counter") == 1  # Reverted to superstep 1 value
+        assert state.get("temp") is None
+
+        # Superstep 3
+        state.set("counter", 3)
+        state.commit()
+        assert state.get("counter") == 3
+
+
+class TestExportImport:
+    """Tests for state serialization (export/import)."""
+
+    def test_export_returns_committed_only(self) -> None:
+        state = State()
+        state.set("committed_key", "committed_value")
+        state.commit()
+        state.set("pending_key", "pending_value")
+
+        exported = state.export_state()
+
+        # Only committed state is exported
+        assert exported == {"committed_key": "committed_value"}
+        assert "pending_key" not in exported
+
+    def test_import_merges_into_committed(self) -> None:
+        state = State()
+        state.set("existing", "original")
+        state.commit()
+
+        state.import_state({"imported": "value", "existing": "overwritten"})
+
+        assert state.get("imported") == "value"
+        assert state.get("existing") == "overwritten"
+
+    def test_import_does_not_affect_pending(self) -> None:
+        state = State()
+        state.set("pending_key", "pending_value")
+
+        state.import_state({"imported": "value"})
+
+        # Pending is still there
+        assert state.get("pending_key") == "pending_value"
+        assert "pending_key" in state._pending  # pyright: ignore[reportPrivateUsage]

python/packages/core/tests/workflow/test_strict_mode_event_labeling.py

@@ -0,0 +1,118 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for the runner's explicit output selection event labeling."""
+
+from __future__ import annotations
+
+import warnings
+from typing import Any
+
+import pytest
+from typing_extensions import Never
+
+from agent_framework import (
+    Message,
+    WorkflowBuilder,
+    WorkflowContext,
+    executor,
+)
+
+
+@executor
+async def _start(messages: list[Message], ctx: WorkflowContext[str, str]) -> None:
+    await ctx.yield_output("from-start")
+    await ctx.send_message("downstream")
+
+
+@executor
+async def _downstream(message: str, ctx: WorkflowContext[Never, str]) -> None:  # type: ignore[valid-type]
+    await ctx.yield_output("from-downstream")
+
+
+def _input_msg() -> list[Message]:
+    return [Message(role="user", contents=["hi"])]
+
+
+@pytest.mark.asyncio
+async def test_strict_mode_designated_executor_emits_output_events() -> None:
+    """Output-designated executor yields produce type='output' events."""
+    workflow = WorkflowBuilder(start_executor=_start, output_from=[_start]).add_edge(_start, _downstream).build()
+    output_events: list[Any] = []
+    intermediate_events: list[Any] = []
+    async for event in workflow.run(_input_msg(), stream=True):
+        if event.type == "output":
+            output_events.append(event)
+        elif event.type == "intermediate":
+            intermediate_events.append(event)
+
+    assert any(ev.data == "from-start" for ev in output_events), "designated executor's yield is type='output'"
+    assert intermediate_events == []
+    assert all(ev.data != "from-downstream" for ev in output_events), "unlisted executor yield is hidden"
+
+
+@pytest.mark.asyncio
+async def test_intermediate_designated_executor_emits_intermediate_events() -> None:
+    """Intermediate-designated executor yields produce type='intermediate' events."""
+    workflow = (
+        WorkflowBuilder(start_executor=_start, intermediate_output_from=[_downstream])
+        .add_edge(_start, _downstream)
+        .build()
+    )
+    output_events: list[Any] = []
+    intermediate_events: list[Any] = []
+    async for event in workflow.run(_input_msg(), stream=True):
+        if event.type == "output":
+            output_events.append(event)
+        elif event.type == "intermediate":
+            intermediate_events.append(event)
+
+    assert len(output_events) == 0
+    assert {ev.data for ev in intermediate_events} == {"from-downstream"}
+
+
+@pytest.mark.asyncio
+async def test_omitted_selection_keeps_all_yields_as_output() -> None:
+    """Omitted output selection preserves today's behavior: all yields are type='output'."""
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", DeprecationWarning)
+        workflow = WorkflowBuilder(start_executor=_start).add_edge(_start, _downstream).build()
+    output_events: list[Any] = []
+    intermediate_events: list[Any] = []
+    async for event in workflow.run(_input_msg(), stream=True):
+        if event.type == "output":
+            output_events.append(event)
+        elif event.type == "intermediate":
+            intermediate_events.append(event)
+
+    assert {ev.data for ev in output_events} == {"from-start", "from-downstream"}
+    assert len(intermediate_events) == 0
+
+
+@pytest.mark.asyncio
+async def test_strict_mode_get_outputs_returns_only_designated() -> None:
+    """WorkflowRunResult.get_outputs() returns only output-designated payloads."""
+    workflow = (
+        WorkflowBuilder(
+            start_executor=_start,
+            output_from=[_downstream],
+            intermediate_output_from=[_start],
+        )
+        .add_edge(_start, _downstream)
+        .build()
+    )
+    result = await workflow.run(_input_msg())
+    assert result.get_outputs() == ["from-downstream"]
+    assert result.get_intermediate_outputs() == ["from-start"]
+
+
+@pytest.mark.asyncio
+async def test_hidden_yields_remain_in_executor_completion_events() -> None:
+    """Hidden yield_output payloads stay available through executor_completed observability."""
+    workflow = WorkflowBuilder(start_executor=_start, output_from=[_downstream]).add_edge(_start, _downstream).build()
+    result = await workflow.run(_input_msg())
+    assert result.get_outputs() == ["from-downstream"]
+    assert result.get_intermediate_outputs() == []
+    assert not any(event.type in {"output", "intermediate"} and event.data == "from-start" for event in result)
+    completed = [event for event in result if event.type == "executor_completed" and event.executor_id == _start.id]
+    assert completed
+    assert completed[0].data == ["downstream", "from-start"]

python/packages/core/tests/workflow/test_sub_workflow.py

@@ -0,0 +1,691 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from dataclasses import dataclass, field
+from typing import Any
+from uuid import uuid4
+
+from typing_extensions import Never
+
+from agent_framework import (
+    Executor,
+    SubWorkflowRequestMessage,
+    SubWorkflowResponseMessage,
+    Workflow,
+    WorkflowBuilder,
+    WorkflowContext,
+    WorkflowEvent,
+    WorkflowExecutor,
+    handler,
+    response_handler,
+)
+from agent_framework._workflows._checkpoint import InMemoryCheckpointStorage
+
+
+# Test message types
+@dataclass
+class EmailValidationRequest:
+    """Request to validate an email address."""
+
+    email: str
+
+
+@dataclass
+class DomainCheckRequest:
+    """Request to check if a domain is approved."""
+
+    id: str = field(default_factory=lambda: str(uuid4()))
+    domain: str = ""
+    email: str = ""  # Include original email for correlation
+
+
+@dataclass
+class ValidationResult:
+    """Result of email validation."""
+
+    email: str
+    is_valid: bool
+    reason: str
+
+
+class Coordinator(Executor):
+    """Coordinator executor in the parent workflow for simple sub-workflow tests."""
+
+    def __init__(self, cache: dict[str, bool] | None = None) -> None:
+        super().__init__(id="basic_parent")
+        self.result: ValidationResult | None = None
+        self.cache: dict[str, bool] = dict(cache) if cache is not None else {}
+        self._pending_sub_workflow_requests: dict[str, SubWorkflowRequestMessage] = {}
+
+    @handler
+    async def start(self, email: str, ctx: WorkflowContext[EmailValidationRequest]) -> None:
+        request = EmailValidationRequest(email=email)
+        await ctx.send_message(request)
+
+    @handler
+    async def handle_domain_request(
+        self,
+        sub_workflow_request: SubWorkflowRequestMessage,
+        ctx: WorkflowContext[SubWorkflowResponseMessage],
+    ) -> None:
+        """Handle requests from sub-workflows with optional caching."""
+        if not isinstance(sub_workflow_request.source_event.data, DomainCheckRequest):
+            raise ValueError("Unexpected request type")
+
+        domain_request = sub_workflow_request.source_event.data
+
+        if domain_request.domain in self.cache:
+            # Return cached result
+            await ctx.send_message(sub_workflow_request.create_response(self.cache[domain_request.domain]))
+        else:
+            # Not in cache, forward to external
+            self._pending_sub_workflow_requests[domain_request.id] = sub_workflow_request
+            await ctx.request_info(domain_request, bool)
+
+    @response_handler
+    async def handle_domain_response(
+        self,
+        original_request: DomainCheckRequest,
+        is_approved: bool,
+        ctx: WorkflowContext[SubWorkflowResponseMessage],
+    ) -> None:
+        """Handle domain check response with correlation and send the response back to the sub-workflow."""
+        if original_request.id not in self._pending_sub_workflow_requests:
+            raise ValueError("No pending sub-workflow request for the given domain check response")
+
+        sub_workflow_request = self._pending_sub_workflow_requests.pop(original_request.id)
+        await ctx.send_message(sub_workflow_request.create_response(is_approved))
+
+    @handler
+    async def collect(self, result: ValidationResult, ctx: WorkflowContext) -> None:
+        self.result = result
+
+
+class EmailFormatValidator(Executor):
+    """Validates the format of an email address."""
+
+    def __init__(self):
+        super().__init__(id="email_format_validator")
+
+    @handler
+    async def validate(
+        self, request: EmailValidationRequest, ctx: WorkflowContext[DomainCheckRequest, ValidationResult]
+    ) -> None:
+        """Validate email format and extract domain."""
+        email = request.email
+        if "@" not in email:
+            result = ValidationResult(email=email, is_valid=False, reason="Invalid email format")
+            await ctx.yield_output(result)
+            return
+
+        domain = email.split("@")[1]
+        domain_check = DomainCheckRequest(domain=domain, email=email)
+        await ctx.send_message(domain_check)
+
+
+class EmailDomainValidator(Executor):
+    """Validates email addresses in a sub-workflow."""
+
+    def __init__(self):
+        super().__init__(id="email_domain_validator")
+
+    @handler
+    async def validate_request(
+        self, request: DomainCheckRequest, ctx: WorkflowContext[DomainCheckRequest, ValidationResult]
+    ) -> None:
+        """Validate an email address."""
+        domain = request.domain
+
+        if not domain:
+            result = ValidationResult(email=request.email, is_valid=False, reason="Invalid email format")
+            await ctx.yield_output(result)
+            return
+
+        # Request domain check from external source
+        await ctx.request_info(request, bool)
+
+    @response_handler
+    async def handle_domain_response(
+        self,
+        original_request: DomainCheckRequest,
+        is_approved: bool,
+        ctx: WorkflowContext[Never, ValidationResult],  # type: ignore[valid-type]
+    ) -> None:
+        """Handle domain check response with correlation."""
+        # Use the original email from the correlated response
+        result = ValidationResult(
+            email=original_request.email,
+            is_valid=is_approved,
+            reason="Domain approved" if is_approved else "Domain not approved",
+        )
+        await ctx.yield_output(result)
+
+
+# Test helper functions
+def create_email_validation_workflow() -> Workflow:
+    """Create a standard email validation workflow."""
+    email_format_validator = EmailFormatValidator()
+    email_domain_validator = EmailDomainValidator()
+
+    return (
+        WorkflowBuilder(start_executor=email_format_validator)
+        .add_edge(email_format_validator, email_domain_validator)
+        .build()
+    )
+
+
+async def test_basic_sub_workflow() -> None:
+    """Test basic sub-workflow execution without interception."""
+    # Create sub-workflow
+    validation_workflow = create_email_validation_workflow()
+
+    # Create parent workflow without interception
+    parent = Coordinator()
+    workflow_executor = WorkflowExecutor(validation_workflow, "email_validation_workflow")
+
+    main_workflow = (
+        WorkflowBuilder(start_executor=parent)
+        .add_edge(parent, workflow_executor)
+        .add_edge(workflow_executor, parent)
+        .build()
+    )
+
+    # Run workflow with mocked external response
+    result = await main_workflow.run("test@example.com")
+
+    # Get request event and respond
+    request_events = result.get_request_info_events()
+    assert len(request_events) == 1
+    assert isinstance(request_events[0].data, DomainCheckRequest)
+    assert request_events[0].data.domain == "example.com"
+
+    # Send response through the main workflow
+    await main_workflow.run(
+        responses={
+            request_events[0].request_id: True  # Domain is approved
+        }
+    )
+
+    # Check result
+    assert parent.result is not None
+    assert parent.result.email == "test@example.com"
+    assert parent.result.is_valid is True
+
+
+async def test_sub_workflow_with_interception():
+    """Test sub-workflow with parent interception and conditional forwarding."""
+    # Create sub-workflow
+    validation_workflow = create_email_validation_workflow()
+
+    # Create parent workflow with interception cache
+    parent = Coordinator(cache={"example.com": True, "internal.org": True})
+    workflow_executor = WorkflowExecutor(validation_workflow, "email_workflow")
+
+    main_workflow = (
+        WorkflowBuilder(start_executor=parent)
+        .add_edge(parent, workflow_executor)
+        .add_edge(workflow_executor, parent)
+        .build()
+    )
+
+    # Test 1: Email with cached domain (intercepted)
+    result = await main_workflow.run("user@example.com")
+    request_events = result.get_request_info_events()
+    assert len(request_events) == 0  # No external requests, handled from cache
+    assert parent.result is not None
+    assert parent.result.email == "user@example.com"
+    assert parent.result.is_valid is True
+
+    # Test 2: Email with unknown domain (forwarded to external)
+    parent.result = None
+    result = await main_workflow.run("user@unknown.com")
+    request_events = result.get_request_info_events()
+    assert len(request_events) == 1  # Forwarded to external
+    assert isinstance(request_events[0].data, DomainCheckRequest)
+    assert request_events[0].data.domain == "unknown.com"
+
+    # Send external response
+    await main_workflow.run(
+        responses={
+            request_events[0].request_id: False  # Domain not approved
+        }
+    )
+    assert parent.result is not None
+    assert parent.result.email == "user@unknown.com"
+    assert parent.result.is_valid is False
+
+    # Test 3: Another cached domain
+    parent.result = None
+    result = await main_workflow.run("user@internal.org")
+    request_events = result.get_request_info_events()
+    assert len(request_events) == 0  # Handled from cache
+    assert parent.result is not None
+    assert parent.result.is_valid is True
+
+
+async def test_workflow_scoped_interception() -> None:
+    """Test interception scoped to specific sub-workflows."""
+
+    class MultiWorkflowParent(Executor):
+        """Parent handling multiple sub-workflows."""
+
+        def __init__(self) -> None:
+            super().__init__(id="multi_parent")
+            self.results: dict[str, ValidationResult] = {}
+            self._pending_sub_workflow_requests: dict[str, SubWorkflowRequestMessage] = {}
+
+        @handler
+        async def start(self, data: dict[str, str], ctx: WorkflowContext[EmailValidationRequest]) -> None:
+            # Send to different sub-workflows
+            await ctx.send_message(EmailValidationRequest(email=data["email1"]), target_id="workflow_a")
+            await ctx.send_message(EmailValidationRequest(email=data["email2"]), target_id="workflow_b")
+
+        @handler
+        async def handle_domain_request(
+            self,
+            sub_workflow_request: SubWorkflowRequestMessage,
+            ctx: WorkflowContext[SubWorkflowResponseMessage],
+        ) -> None:
+            """Handle requests from sub-workflows with optional caching."""
+            if not isinstance(sub_workflow_request.source_event.data, DomainCheckRequest):
+                raise ValueError("Unexpected request type")
+
+            domain_request = sub_workflow_request.source_event.data
+
+            if sub_workflow_request.executor_id == "workflow_a" and domain_request.domain == "strict.com":
+                # Strict rules for workflow A
+                await ctx.send_message(
+                    sub_workflow_request.create_response(True), target_id=sub_workflow_request.executor_id
+                )
+                return
+            if sub_workflow_request.executor_id == "workflow_b" and domain_request.domain.endswith(".com"):
+                # Lenient rules for workflow B
+                await ctx.send_message(
+                    sub_workflow_request.create_response(True), target_id=sub_workflow_request.executor_id
+                )
+                return
+
+            # Unknown source, forward to external
+            self._pending_sub_workflow_requests[domain_request.id] = sub_workflow_request
+            await ctx.request_info(domain_request, bool)
+
+        @response_handler
+        async def handle_domain_response(
+            self,
+            original_request: DomainCheckRequest,
+            is_approved: bool,
+            ctx: WorkflowContext[SubWorkflowResponseMessage],
+        ) -> None:
+            """Handle domain check response with correlation and send the response back to the sub-workflow."""
+            if original_request.id not in self._pending_sub_workflow_requests:
+                raise ValueError("No pending sub-workflow request for the given domain check response")
+
+            sub_workflow_request = self._pending_sub_workflow_requests.pop(original_request.id)
+            await ctx.send_message(
+                sub_workflow_request.create_response(is_approved), target_id=sub_workflow_request.executor_id
+            )
+
+        @handler
+        async def collect(self, result: ValidationResult, ctx: WorkflowContext) -> None:
+            self.results[result.email] = result
+
+    # Create two identical sub-workflows
+    workflow_a = create_email_validation_workflow()
+    workflow_b = create_email_validation_workflow()
+
+    parent = MultiWorkflowParent()
+    executor_a = WorkflowExecutor(workflow_a, "workflow_a")
+    executor_b = WorkflowExecutor(workflow_b, "workflow_b")
+
+    main_workflow = (
+        WorkflowBuilder(start_executor=parent)
+        .add_edge(parent, executor_a)
+        .add_edge(parent, executor_b)
+        .add_edge(executor_a, parent)
+        .add_edge(executor_b, parent)
+        .build()
+    )
+
+    # Run test
+    result = await main_workflow.run({"email1": "user@strict.com", "email2": "user@random.com"})
+
+    # Workflow A should handle strict.com
+    # Workflow B should handle any .com domain
+    request_events = result.get_request_info_events()
+    assert len(request_events) == 0  # Both handled internally
+
+    assert len(parent.results) == 2
+    assert parent.results["user@strict.com"].is_valid is True
+    assert parent.results["user@random.com"].is_valid is True
+
+
+async def test_concurrent_sub_workflow_execution() -> None:
+    """Test that WorkflowExecutor can handle multiple concurrent invocations properly."""
+
+    class ConcurrentProcessor(Executor):
+        """Processor that sends multiple concurrent requests to the same sub-workflow."""
+
+        def __init__(self) -> None:
+            super().__init__(id="concurrent_processor")
+            self.results: list[ValidationResult] = []
+            self._pending_sub_workflow_requests: dict[str, SubWorkflowRequestMessage] = {}
+
+        @handler
+        async def start(self, emails: list[str], ctx: WorkflowContext[EmailValidationRequest]) -> None:
+            """Send multiple concurrent requests to the same sub-workflow."""
+            # Send all requests concurrently to the same workflow executor
+            for email in emails:
+                request = EmailValidationRequest(email=email)
+                await ctx.send_message(request)
+
+        @handler
+        async def handle_domain_request(
+            self,
+            sub_workflow_request: SubWorkflowRequestMessage,
+            ctx: WorkflowContext[SubWorkflowResponseMessage],
+        ) -> None:
+            """Handle requests from sub-workflows with optional caching."""
+            if not isinstance(sub_workflow_request.source_event.data, DomainCheckRequest):
+                raise ValueError("Unexpected request type")
+
+            domain_request = sub_workflow_request.source_event.data
+            self._pending_sub_workflow_requests[domain_request.id] = sub_workflow_request
+            await ctx.request_info(domain_request, bool)
+
+        @response_handler
+        async def handle_domain_response(
+            self,
+            original_request: DomainCheckRequest,
+            is_approved: bool,
+            ctx: WorkflowContext[SubWorkflowResponseMessage],
+        ) -> None:
+            """Handle domain check response with correlation and send the response back to the sub-workflow."""
+            if original_request.id not in self._pending_sub_workflow_requests:
+                raise ValueError("No pending sub-workflow request for the given domain check response")
+
+            sub_workflow_request = self._pending_sub_workflow_requests.pop(original_request.id)
+            await ctx.send_message(sub_workflow_request.create_response(is_approved))
+
+        @handler
+        async def collect_result(self, result: ValidationResult, ctx: WorkflowContext) -> None:
+            """Collect results from concurrent executions."""
+            self.results.append(result)
+
+    # Create sub-workflow for email validation
+    validation_workflow = create_email_validation_workflow()
+
+    # Create parent workflow
+    processor = ConcurrentProcessor()
+    workflow_executor = WorkflowExecutor(validation_workflow, "email_workflow")
+
+    main_workflow = (
+        WorkflowBuilder(start_executor=processor)
+        .add_edge(processor, workflow_executor)
+        .add_edge(workflow_executor, processor)
+        .build()
+    )
+
+    # Test concurrent execution with multiple emails
+    emails = [
+        "user1@domain1.com",
+        "user2@domain2.com",
+        "user3@domain3.com",
+        "user4@domain4.com",
+        "user5@domain5.com",
+    ]
+
+    result = await main_workflow.run(emails)
+
+    # Each email should generate one external request
+    request_events = result.get_request_info_events()
+    assert len(request_events) == len(emails)
+
+    # Verify each request corresponds to the correct domain
+    domains_requested = {event.data.domain for event in request_events}  # type: ignore[union-attr]
+    expected_domains = {f"domain{i}.com" for i in range(1, 6)}
+    assert domains_requested == expected_domains
+
+    # Send responses for all requests (approve all domains)
+    responses = {event.request_id: True for event in request_events}
+    await main_workflow.run(responses=responses)
+
+    # All results should be collected
+    assert len(processor.results) == len(emails)
+
+    # Verify each email was processed correctly
+    result_emails = {result.email for result in processor.results}
+    expected_emails = set(emails)
+    assert result_emails == expected_emails
+
+    # All should be valid since we approved all domains
+    for result_obj in processor.results:
+        assert result_obj.is_valid is True
+        assert result_obj.reason == "Domain approved"
+
+    # Verify that concurrent executions were properly isolated
+    # (This is implicitly tested by the fact that we got correct results for all emails)
+
+
+# region Checkpoint-related message types and executors for sub-workflow tests
+
+
+@dataclass
+class CheckpointRequest:
+    """Request in a two-step checkpoint test."""
+
+    prompt: str
+    id: str = field(default_factory=lambda: str(uuid4()))
+
+
+class TwoStepSubWorkflowExecutor(Executor):
+    """Sub-workflow executor that makes two sequential requests."""
+
+    def __init__(self) -> None:
+        super().__init__(id="two_step_executor")
+        self._responses: list[str] = []
+
+    @handler
+    async def handle_start(self, msg: str, ctx: WorkflowContext) -> None:
+        await ctx.request_info(
+            request_data=CheckpointRequest(prompt=f"First request for: {msg}"),
+            response_type=str,
+        )
+
+    @response_handler
+    async def handle_response(
+        self,
+        original_request: CheckpointRequest,
+        response: str,
+        ctx: WorkflowContext[Never, bool],  # type: ignore[valid-type]
+    ) -> None:
+        self._responses.append(response)
+        if len(self._responses) == 1:
+            # First response received, make second request
+            await ctx.request_info(
+                request_data=CheckpointRequest(prompt="Second request"),
+                response_type=str,
+            )
+        else:
+            # Second response received, yield final output
+            await ctx.yield_output(True)
+
+    async def on_checkpoint_save(self) -> dict[str, Any]:
+        return {"responses": self._responses}
+
+    async def on_checkpoint_restore(self, state: dict[str, Any]) -> None:
+        self._responses = state.get("responses", [])
+
+
+class CheckpointTestCoordinator(Executor):
+    """Coordinator for checkpoint sub-workflow tests."""
+
+    def __init__(self) -> None:
+        super().__init__(id="checkpoint_coordinator")
+        self._pending_requests: dict[str, SubWorkflowRequestMessage] = {}
+
+    @handler
+    async def start(self, value: str, ctx: WorkflowContext[str]) -> None:
+        await ctx.send_message(value)
+
+    @handler
+    async def handle_sub_workflow_request(
+        self,
+        request: SubWorkflowRequestMessage,
+        ctx: WorkflowContext,
+    ) -> None:
+        data = request.source_event.data
+        if isinstance(data, CheckpointRequest):
+            self._pending_requests[data.id] = request
+            await ctx.request_info(data, str)
+
+    @response_handler
+    async def handle_response(
+        self,
+        original_request: CheckpointRequest,
+        response: str,
+        ctx: WorkflowContext[SubWorkflowResponseMessage],
+    ) -> None:
+        sub_request = self._pending_requests.pop(original_request.id, None)
+        if sub_request is None:
+            raise ValueError(f"No pending request for ID: {original_request.id}")
+        await ctx.send_message(sub_request.create_response(response))
+
+    async def on_checkpoint_save(self) -> dict[str, Any]:
+        return {"pending_requests": self._pending_requests}
+
+    async def on_checkpoint_restore(self, state: dict[str, Any]) -> None:
+        self._pending_requests = state.get("pending_requests", {})
+
+
+def _build_checkpoint_test_workflow(storage: InMemoryCheckpointStorage) -> Workflow:
+    """Build the main workflow with checkpointing for testing."""
+    two_step_executor = TwoStepSubWorkflowExecutor()
+    sub_workflow = WorkflowBuilder(start_executor=two_step_executor).build()
+    sub_workflow_executor = WorkflowExecutor(sub_workflow, id="sub_workflow_executor")
+
+    coordinator = CheckpointTestCoordinator()
+    return (
+        WorkflowBuilder(start_executor=coordinator, checkpoint_storage=storage)
+        .add_edge(coordinator, sub_workflow_executor)
+        .add_edge(sub_workflow_executor, coordinator)
+        .build()
+    )
+
+
+async def test_sub_workflow_checkpoint_restore_no_duplicate_requests() -> None:
+    """Test that resuming a sub-workflow from checkpoint does not emit duplicate requests.
+
+    This test verifies the fix for an issue where after checkpoint restore, when a response
+    is sent to a sub-workflow, duplicate RequestInfoEvents were emitted. The bug occurred
+    because checkpoint rehydration re-added RequestInfoEvents to the event queue, and when
+    the workflow was resumed, those events were emitted again along with any new requests.
+
+    The fix ensures that already-handled requests are filtered out from the result when
+    the sub-workflow is resumed with responses.
+    """
+    storage = InMemoryCheckpointStorage()
+
+    # Step 1: Run workflow until first request
+    workflow1 = _build_checkpoint_test_workflow(storage)
+
+    first_request_id: str | None = None
+    async for event in workflow1.run("test_value", stream=True):
+        if event.type == "request_info":
+            first_request_id = event.request_id
+
+    assert first_request_id is not None
+
+    # Get checkpoint
+    checkpoints = await storage.list_checkpoints(workflow_name=workflow1.name)
+    checkpoint_id = max(checkpoints, key=lambda cp: cp.iteration_count).checkpoint_id
+
+    # Step 2: Resume workflow from checkpoint
+    workflow2 = _build_checkpoint_test_workflow(storage)
+
+    resumed_first_request_id: str | None = None
+    async for event in workflow2.run(checkpoint_id=checkpoint_id, stream=True):
+        if event.type == "request_info":
+            resumed_first_request_id = event.request_id
+
+    assert resumed_first_request_id is not None
+    assert resumed_first_request_id == first_request_id
+
+    request_events: list[WorkflowEvent] = []
+    async for event in workflow2.run(stream=True, responses={resumed_first_request_id: "first_answer"}):
+        if event.type == "request_info":
+            request_events.append(event)
+
+    # Key assertion: Only the second request should be received, not a duplicate of the first
+    assert len(request_events) == 1
+    assert request_events[0].data.prompt == "Second request"
+
+
+async def test_sub_workflow_intermediate_outputs_propagate_to_parent() -> None:
+    """A child workflow's intermediate emissions must bubble up through the parent.
+
+    Regression guard for the bug where WorkflowExecutor._process_workflow_result only
+    forwarded result.get_outputs() and silently dropped result.get_intermediate_outputs().
+    The forwarded event must carry the WorkflowExecutor's own id as the source so outer
+    callers don't have to know the child's internal executor layout, and it must keep
+    type='intermediate' regardless of how the parent designates the WorkflowExecutor.
+    """
+
+    class _ProgressEmitter(Executor):
+        def __init__(self) -> None:
+            super().__init__(id="progress_emitter")
+
+        @handler
+        async def run(self, message: str, ctx: WorkflowContext[str, str]) -> None:
+            await ctx.yield_output(f"progress: {message}")
+            await ctx.send_message(message)
+
+    class _Finalizer(Executor):
+        def __init__(self) -> None:
+            super().__init__(id="finalizer")
+
+        @handler
+        async def run(self, message: str, ctx: WorkflowContext[Never, str]) -> None:  # type: ignore[valid-type]
+            await ctx.yield_output(f"final: {message}")
+
+    progress = _ProgressEmitter()
+    finalizer = _Finalizer()
+    child = (
+        WorkflowBuilder(
+            start_executor=progress,
+            output_from=[finalizer],
+            intermediate_output_from=[progress],
+        )
+        .add_edge(progress, finalizer)
+        .build()
+    )
+
+    sub = WorkflowExecutor(child, id="sub")
+
+    class _ParentSink(Executor):
+        def __init__(self) -> None:
+            super().__init__(id="parent_sink")
+            self.received: list[str] = []
+
+        @handler
+        async def run(self, message: str, ctx: WorkflowContext[Never, str]) -> None:  # type: ignore[valid-type]
+            self.received.append(message)
+            await ctx.yield_output(message)
+
+    sink = _ParentSink()
+    parent = WorkflowBuilder(start_executor=sub, output_from=[sink]).add_edge(sub, sink).build()
+
+    intermediate_events: list[WorkflowEvent[Any]] = []
+    output_events: list[WorkflowEvent[Any]] = []
+    async for event in parent.run("hello", stream=True):
+        if event.type == "intermediate":
+            intermediate_events.append(event)
+        elif event.type == "output":
+            output_events.append(event)
+
+    # The child's intermediate emission bubbled up labeled with the WorkflowExecutor id,
+    # not the child's internal executor id.
+    assert len(intermediate_events) == 1, [(e.executor_id, e.data) for e in intermediate_events]
+    assert intermediate_events[0].executor_id == "sub"
+    assert intermediate_events[0].data == "progress: hello"
+
+    # The parent's own terminal output is unaffected.
+    assert any(e.executor_id == "parent_sink" and e.data == "final: hello" for e in output_events)

python/packages/core/tests/workflow/test_typevar_validation.py

@@ -0,0 +1,258 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for unresolved TypeVar detection during handler/executor registration."""
+
+from typing import TypeVar
+
+import pytest
+from typing_extensions import Never
+
+from agent_framework import (
+    Executor,
+    FunctionExecutor,
+    WorkflowContext,
+    executor,
+    handler,
+)
+from agent_framework._workflows._typing_utils import contains_typevar, is_typevar
+
+T = TypeVar("T")
+U = TypeVar("U")
+
+
+class TestIsTypevarHelper:
+    """Tests for the runtime-safe is_typevar helper."""
+
+    def test_detects_typing_typevar(self):
+        """is_typevar should detect TypeVar from typing module."""
+        import typing
+
+        tv = typing.TypeVar("tv")
+        assert is_typevar(tv)
+
+    def test_detects_typing_extensions_typevar(self):
+        """is_typevar should detect TypeVar from typing_extensions module."""
+        import typing_extensions
+
+        tv = typing_extensions.TypeVar("tv")
+        assert is_typevar(tv)
+
+    def test_rejects_concrete_types(self):
+        """is_typevar should return False for concrete types."""
+        assert not is_typevar(str)
+        assert not is_typevar(int)
+        assert not is_typevar(None)
+        assert not is_typevar(Never)
+
+    def test_rejects_non_types(self):
+        """is_typevar should return False for non-type values."""
+        assert not is_typevar("hello")
+        assert not is_typevar(42)
+        assert not is_typevar([])
+
+    def test_contains_typevar_detects_nested_typevars(self):
+        """contains_typevar should detect TypeVar nested in typing constructs."""
+        assert contains_typevar(list[T])  # type: ignore[misc, valid-type]
+        assert contains_typevar(dict[str, T])  # type: ignore[misc, valid-type]
+        assert contains_typevar(str | list[T])  # type: ignore[misc, valid-type]
+
+    def test_contains_typevar_rejects_concrete_nested_types(self):
+        """contains_typevar should return False for concrete nested types."""
+        assert not contains_typevar(list[str])
+        assert not contains_typevar(dict[str, int])
+        assert not contains_typevar(str | None)
+
+
+class TestHandlerTypeVarValidation:
+    """Tests for @handler decorator rejecting unresolved TypeVars."""
+
+    def test_handler_explicit_input_typevar_raises(self):
+        """@handler(input=T) with a TypeVar should raise ValueError."""
+        with pytest.raises(ValueError, match="unresolved TypeVar"):
+
+            class _Bad(Executor):  # pyright: ignore[reportUnusedClass]
+                @handler(input=T)  # type: ignore[arg-type, call-overload]  # ty: ignore[invalid-argument-type]
+                async def handle(self, message, ctx: WorkflowContext[str]) -> None:  # type: ignore[no-untyped-def]
+                    pass
+
+    def test_handler_explicit_output_typevar_raises(self):
+        """@handler(input=str, output=T) with a TypeVar should raise ValueError."""
+        with pytest.raises(ValueError, match="unresolved TypeVar"):
+
+            class _Bad(Executor):  # pyright: ignore[reportUnusedClass]
+                @handler(input=str, output=T)  # type: ignore[arg-type, call-overload]  # ty: ignore[invalid-argument-type]
+                async def handle(self, message: str, ctx: WorkflowContext[str]) -> None:
+                    pass
+
+    def test_handler_explicit_workflow_output_typevar_raises(self):
+        """@handler(input=str, workflow_output=T) should raise ValueError."""
+        with pytest.raises(ValueError, match="unresolved TypeVar"):
+
+            class _Bad(Executor):  # pyright: ignore[reportUnusedClass]
+                @handler(input=str, workflow_output=T)  # type: ignore[arg-type, call-overload]  # ty: ignore[invalid-argument-type]
+                async def handle(self, message: str, ctx: WorkflowContext[str]) -> None:
+                    pass
+
+    def test_handler_explicit_nested_input_typevar_raises(self):
+        """@handler(input=list[T]) should raise ValueError."""
+        with pytest.raises(ValueError, match="unresolved TypeVar"):
+
+            class _Bad(Executor):  # pyright: ignore[reportUnusedClass]
+                @handler(input=list[T])  # type: ignore[arg-type, call-overload, misc, valid-type]
+                async def handle(self, message, ctx: WorkflowContext[str]) -> None:  # type: ignore[no-untyped-def]
+                    pass
+
+    def test_handler_introspected_typevar_raises(self):
+        """@handler with TypeVar in message annotation should raise ValueError."""
+        with pytest.raises(ValueError, match="unresolved TypeVar"):
+
+            class _Bad(Executor):  # pyright: ignore[reportUnusedClass]
+                @handler  # type: ignore[arg-type]
+                async def handle(self, message: T, ctx: WorkflowContext[str]) -> None:  # type: ignore[valid-type]
+                    pass
+
+    def test_handler_introspected_nested_typevar_raises(self):
+        """@handler with TypeVar nested in message annotation should raise ValueError."""
+        with pytest.raises(ValueError, match="unresolved TypeVar"):
+
+            class _Bad(Executor):  # pyright: ignore[reportUnusedClass]
+                @handler  # type: ignore[arg-type]
+                async def handle(self, message: list[T], ctx: WorkflowContext[str]) -> None:  # type: ignore[valid-type]
+                    pass
+
+    def test_handler_concrete_types_work(self):
+        """@handler with concrete types should succeed."""
+
+        class Good(Executor):
+            @handler(input=str, output=str)
+            async def handle(self, message: str, ctx: WorkflowContext[str]) -> None:
+                pass
+
+        assert Good is not None
+
+
+class TestExecutorTypeVarValidation:
+    """Tests for @executor decorator rejecting unresolved TypeVars."""
+
+    def test_executor_explicit_input_typevar_raises(self):
+        """@executor(input=T) with a TypeVar should raise ValueError."""
+        with pytest.raises(ValueError, match="unresolved TypeVar"):
+
+            @executor(input=T)  # type: ignore[arg-type, call-overload]  # ty: ignore[invalid-argument-type]
+            async def bad_func(message, ctx: WorkflowContext[str]) -> None:  # type: ignore[no-untyped-def]
+                pass
+
+    def test_executor_explicit_output_typevar_raises(self):
+        """@executor(input=str, output=T) with a TypeVar should raise ValueError."""
+        with pytest.raises(ValueError, match="unresolved TypeVar"):
+
+            @executor(input=str, output=T)  # type: ignore[arg-type, call-overload]  # ty: ignore[invalid-argument-type]
+            async def bad_func(message: str, ctx: WorkflowContext[str]) -> None:
+                pass
+
+    def test_executor_explicit_nested_input_typevar_raises(self):
+        """@executor(input=list[T]) should raise ValueError."""
+        with pytest.raises(ValueError, match="unresolved TypeVar"):
+
+            @executor(input=list[T])  # type: ignore[arg-type, call-overload, misc, valid-type]
+            async def bad_func(message, ctx: WorkflowContext[str]) -> None:  # type: ignore[no-untyped-def]
+                pass
+
+    def test_executor_introspected_typevar_raises(self):
+        """@executor with TypeVar in message annotation should raise ValueError."""
+        with pytest.raises(ValueError, match="unresolved TypeVar"):
+            FunctionExecutor(self._make_typevar_func())  # type: ignore[arg-type]
+
+    def test_executor_introspected_nested_typevar_raises(self):
+        """@executor with TypeVar nested in message annotation should raise ValueError."""
+        with pytest.raises(ValueError, match="unresolved TypeVar"):
+            FunctionExecutor(self._make_nested_typevar_func())  # type: ignore[arg-type]
+
+    def test_executor_concrete_types_work(self):
+        """@executor with concrete types should succeed."""
+
+        @executor(input=str, output=str)
+        async def good_func(message: str, ctx: WorkflowContext[str]) -> None:
+            pass
+
+        assert good_func is not None
+
+    @staticmethod
+    def _make_typevar_func():
+        """Create a function with TypeVar annotation for testing."""
+
+        async def func(message: T, ctx: WorkflowContext[str]) -> None:  # type: ignore[valid-type]
+            pass
+
+        return func
+
+    @staticmethod
+    def _make_nested_typevar_func():
+        """Create a function with nested TypeVar annotation for testing."""
+
+        async def func(message: list[T], ctx: WorkflowContext[str]) -> None:  # type: ignore[valid-type]
+            pass
+
+        return func
+
+
+class TestWorkflowContextTypeVarValidation:
+    """Tests for WorkflowContext[T] rejecting unresolved TypeVars."""
+
+    def test_context_direct_typevar_raises(self):
+        """WorkflowContext[T] with a TypeVar should raise ValueError."""
+        with pytest.raises(ValueError, match="unresolved TypeVar"):
+
+            @executor(id="bad")
+            async def bad_func(message: str, ctx: WorkflowContext[T]) -> None:  # type: ignore[valid-type]
+                pass
+
+    def test_context_union_typevar_raises(self):
+        """WorkflowContext[T | str] with a TypeVar in union should raise ValueError."""
+        with pytest.raises(ValueError, match="unresolved TypeVar"):
+
+            @executor(id="bad")
+            async def bad_func(message: str, ctx: WorkflowContext[T | str]) -> None:  # type: ignore[valid-type]
+                pass
+
+    def test_context_nested_typevar_raises(self):
+        """WorkflowContext[list[T]] with a nested TypeVar should raise ValueError."""
+        with pytest.raises(ValueError, match="unresolved TypeVar"):
+
+            @executor(id="bad")
+            async def bad_func(message: str, ctx: WorkflowContext[list[T]]) -> None:  # type: ignore[valid-type]
+                pass
+
+    def test_context_workflow_output_typevar_raises(self):
+        """WorkflowContext[str, T] with a TypeVar should raise ValueError."""
+        with pytest.raises(ValueError, match="unresolved TypeVar"):
+
+            @executor(id="bad")
+            async def bad_func(message: str, ctx: WorkflowContext[str, T]) -> None:  # type: ignore[valid-type]
+                pass
+
+    def test_context_nested_workflow_output_typevar_raises(self):
+        """WorkflowContext[str, dict[str, T]] with a nested TypeVar should raise ValueError."""
+        with pytest.raises(ValueError, match="unresolved TypeVar"):
+
+            @executor(id="bad")
+            async def bad_func(message: str, ctx: WorkflowContext[str, dict[str, T]]) -> None:  # type: ignore[valid-type]
+                pass
+
+    def test_context_concrete_types_work(self):
+        """WorkflowContext[str] with concrete types should succeed."""
+
+        @executor(id="good")
+        async def good_func(message: str, ctx: WorkflowContext[str]) -> None:
+            pass
+
+        assert good_func is not None
+
+    def test_context_class_handler_typevar_raises(self):
+        """Class-based handler with WorkflowContext[T] should raise ValueError."""
+        with pytest.raises(ValueError, match="unresolved TypeVar"):
+
+            class _Bad(Executor):  # pyright: ignore[reportUnusedClass]
+                @handler  # pyright: ignore[reportUnknownArgumentType]
+                async def handle(self, message: str, ctx: WorkflowContext[T]) -> None:  # type: ignore[valid-type]
+                    pass

python/packages/core/tests/workflow/test_typing_utils.py

@@ -0,0 +1,492 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from dataclasses import dataclass
+from typing import Any, Generic, Optional, TypeVar, Union
+
+import pytest
+
+from agent_framework import WorkflowEvent
+from agent_framework._workflows._typing_utils import (
+    deserialize_type,
+    is_instance_of,
+    is_type_compatible,
+    normalize_type_to_list,
+    resolve_type_annotation,
+    serialize_type,
+    try_coerce_to_type,
+)
+
+# region: normalize_type_to_list tests
+
+
+def test_normalize_type_to_list_single_type() -> None:
+    """Test normalize_type_to_list with single types."""
+    assert normalize_type_to_list(str) == [str]
+    assert normalize_type_to_list(int) == [int]
+    assert normalize_type_to_list(float) == [float]
+    assert normalize_type_to_list(bool) == [bool]
+    assert normalize_type_to_list(list) == [list]
+    assert normalize_type_to_list(dict) == [dict]
+
+
+def test_normalize_type_to_list_none() -> None:
+    """Test normalize_type_to_list with None returns empty list."""
+    assert normalize_type_to_list(None) == []
+
+
+def test_normalize_type_to_list_union_pipe_syntax() -> None:
+    """Test normalize_type_to_list with union types using | syntax."""
+    result = normalize_type_to_list(str | int)  # pyright: ignore[reportArgumentType]
+    assert set(result) == {str, int}
+
+    result = normalize_type_to_list(str | int | bool)  # pyright: ignore[reportArgumentType]
+    assert set(result) == {str, int, bool}
+
+
+def test_normalize_type_to_list_union_typing_syntax() -> None:
+    """Test normalize_type_to_list with Union[] from typing module."""
+    result = normalize_type_to_list(Union[str, int])  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType]
+    assert set(result) == {str, int}
+
+    result = normalize_type_to_list(Union[str, int, bool])  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType]
+    assert set(result) == {str, int, bool}
+
+
+def test_normalize_type_to_list_optional() -> None:
+    """Test normalize_type_to_list with Optional types (Union[T, None])."""
+    # Optional[str] is Union[str, None]
+    result = normalize_type_to_list(Optional[str])  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType]
+    assert str in result
+    assert type(None) in result
+    assert len(result) == 2
+
+    # str | None is equivalent
+    result = normalize_type_to_list(str | None)  # pyright: ignore[reportArgumentType]
+    assert str in result
+    assert type(None) in result
+    assert len(result) == 2
+
+
+def test_normalize_type_to_list_custom_types() -> None:
+    """Test normalize_type_to_list with custom class types."""
+
+    @dataclass
+    class CustomMessage:
+        content: str
+
+    result = normalize_type_to_list(CustomMessage)
+    assert result == [CustomMessage]
+
+    result = normalize_type_to_list(CustomMessage | str)  # pyright: ignore[reportArgumentType]
+    assert set(result) == {CustomMessage, str}
+
+
+# endregion: normalize_type_to_list tests
+
+
+# region: resolve_type_annotation tests
+
+
+def test_resolve_type_annotation_none() -> None:
+    """Test resolve_type_annotation with None returns None."""
+    assert resolve_type_annotation(None) is None
+
+
+def test_resolve_type_annotation_actual_types() -> None:
+    """Test resolve_type_annotation passes through actual types unchanged."""
+    assert resolve_type_annotation(str) is str
+    assert resolve_type_annotation(int) is int
+    assert resolve_type_annotation(str | int) == str | int  # pyright: ignore[reportArgumentType]
+
+
+def test_resolve_type_annotation_string_builtin() -> None:
+    """Test resolve_type_annotation resolves string references to builtin types."""
+    result = resolve_type_annotation("str", {"str": str})
+    assert result is str
+
+    result = resolve_type_annotation("int", {"int": int})
+    assert result is int
+
+
+def test_resolve_type_annotation_string_union() -> None:
+    """Test resolve_type_annotation resolves string union types."""
+    result = resolve_type_annotation("str | int", {"str": str, "int": int})
+    assert result == str | int
+
+
+def test_resolve_type_annotation_string_custom_type() -> None:
+    """Test resolve_type_annotation resolves string references to custom types."""
+
+    @dataclass
+    class MyCustomType:
+        value: int
+
+    result = resolve_type_annotation("MyCustomType", {"MyCustomType": MyCustomType})
+    assert result is MyCustomType
+
+    result = resolve_type_annotation("MyCustomType | str", {"MyCustomType": MyCustomType, "str": str})
+    assert set(result.__args__) == {MyCustomType, str}  # type: ignore[union-attr]  # ty: ignore[unresolved-attribute]
+
+
+def test_resolve_type_annotation_string_typing_union() -> None:
+    """Test resolve_type_annotation resolves Union[] syntax in strings."""
+    result = resolve_type_annotation("Union[str, int]", {"str": str, "int": int})
+    assert set(result.__args__) == {str, int}  # type: ignore[union-attr]  # ty: ignore[unresolved-attribute]
+
+
+def test_resolve_type_annotation_string_optional() -> None:
+    """Test resolve_type_annotation resolves Optional[] syntax in strings."""
+    result = resolve_type_annotation("Optional[str]", {"str": str})
+    assert str in result.__args__  # type: ignore[union-attr]  # ty: ignore[unresolved-attribute]
+    assert type(None) in result.__args__  # type: ignore[union-attr]  # ty: ignore[unresolved-attribute]
+
+
+def test_resolve_type_annotation_unresolvable_raises() -> None:
+    """Test resolve_type_annotation raises NameError for unresolvable types."""
+    with pytest.raises(NameError, match="Could not resolve type annotation"):
+        resolve_type_annotation("NonExistentType", {})
+
+
+# endregion: resolve_type_annotation tests
+
+
+def test_basic_types() -> None:
+    """Test basic built-in types."""
+    assert is_instance_of(5, int)
+    assert is_instance_of("hello", str)
+    assert is_instance_of(None, type(None))
+
+
+def test_union_types() -> None:
+    """Test union types (|) and optional types."""
+    assert is_instance_of(5, int | str)
+    assert is_instance_of("hello", int | str)
+    assert is_instance_of(5, Union[int, str])
+    assert not is_instance_of(5.0, int | str)
+
+
+def test_list_types() -> None:
+    """Test list types with various element types."""
+    assert is_instance_of([], list)
+    assert is_instance_of([1, 2, 3], list)
+    assert is_instance_of([1, 2, 3], list[int])
+    assert is_instance_of([1, 2, 3], list[int | str])
+    assert is_instance_of([1, "a", 3], list[int | str])
+    assert is_instance_of([1, "a", 3], list[Union[int, str]])
+    assert not is_instance_of([1, 2.0, 3], dict)
+    assert not is_instance_of([1, 2.0, 3], list[int | str])
+
+
+def test_tuple_types() -> None:
+    """Test tuple types with fixed and variable lengths."""
+    assert is_instance_of((1, "a"), tuple)
+    assert is_instance_of((1, "a"), tuple[int, str])
+    assert is_instance_of((1, "a", 3), tuple[int | str, ...])
+    assert is_instance_of((1, 2.0, "a"), tuple[...])  # type: ignore
+    assert not is_instance_of((1, 2.0, 3), tuple[int | str, ...])
+    assert not is_instance_of((1, 2.0, 3), dict)
+
+
+def test_dict_types() -> None:
+    """Test dictionary types with typed keys and values."""
+    assert is_instance_of({"key": "value"}, dict)
+    assert is_instance_of({"key": "value"}, dict[str, str])
+    assert is_instance_of({"key": 5, "another_key": "value"}, dict[str, int | str])
+    assert not is_instance_of({"key": 5, "another_key": 3.0}, dict[str, int | str])
+    assert not is_instance_of({"key": 5, "another_key": 3.0}, list)
+
+
+def test_set_types() -> None:
+    """Test set types with various element types."""
+    assert is_instance_of({1, 2, 3}, set)
+    assert is_instance_of({1, 2, 3}, set[int])
+    assert is_instance_of({1, 2, 3}, set[int | str])
+    assert is_instance_of({1, "a", 3}, set[int | str])
+    assert is_instance_of({1, "a", 3}, set[Union[int, str]])
+    assert is_instance_of(set(), set[int])
+    assert not is_instance_of({1, 2.0, 3}, set[int | str])
+    assert not is_instance_of({1, 2, 3}, list)
+    assert not is_instance_of({1, 2, 3}, dict)
+
+
+def test_any_type() -> None:
+    """Test Any type - should accept all values."""
+    assert is_instance_of(5, Any)
+    assert is_instance_of("hello", Any)
+    assert is_instance_of([1, 2, 3], Any)
+
+
+def test_nested_types() -> None:
+    """Test complex nested type structures."""
+    assert is_instance_of([{"key": [1, 2]}, {"another_key": [3]}], list[dict[str, list[int]]])
+    assert not is_instance_of([{"key": [1, 2]}, {"another_key": [3.0]}], list[dict[str, list[int]]])
+
+
+def test_custom_type() -> None:
+    """Test custom object type checking."""
+
+    @dataclass
+    class CustomClass:
+        value: int
+
+    instance = CustomClass(10)
+    assert is_instance_of(instance, CustomClass)
+    assert not is_instance_of(instance, dict)
+
+
+def test_custom_generic_type() -> None:
+    """Test custom generic type checking."""
+
+    T = TypeVar("T")
+    U = TypeVar("U")
+
+    class CustomClass(Generic[T, U]):
+        def __init__(self, request: T, response: U, extra: Any | None = None) -> None:
+            self.request = request
+            self.response = response
+            self.extra = extra
+
+    instance = CustomClass[int, str](request=5, response="response")
+
+    assert is_instance_of(instance, CustomClass[int, str])
+    # Generic parameters are not strictly enforced at runtime
+    assert is_instance_of(instance, CustomClass[str, str])
+
+
+def test_edge_cases() -> None:
+    """Test edge cases and unusual scenarios."""
+    assert is_instance_of([], list[int])  # Empty list should be valid
+    assert is_instance_of((), tuple[int, ...])  # Empty tuple should be valid
+    assert is_instance_of({}, dict[str, int])  # Empty dict should be valid
+    assert is_instance_of(None, int | None)  # Optional type with None
+    assert not is_instance_of(5, str | None)  # Optional type without matching type
+
+
+def test_serialize_type() -> None:
+    """Test serialization of types to strings."""
+    # Test built-in types
+    assert serialize_type(int) == "builtins.int"
+    assert serialize_type(str) == "builtins.str"
+    assert serialize_type(float) == "builtins.float"
+    assert serialize_type(bool) == "builtins.bool"
+    assert serialize_type(list) == "builtins.list"
+    assert serialize_type(dict) == "builtins.dict"
+    assert serialize_type(tuple) == "builtins.tuple"
+    assert serialize_type(set) == "builtins.set"
+
+    # Test custom class
+    @dataclass
+    class TestClass:
+        value: int
+
+    # The custom class will be in the test module
+    expected = f"{TestClass.__module__}.{TestClass.__qualname__}"
+    assert serialize_type(TestClass) == expected
+
+
+def test_deserialize_type() -> None:
+    """Test deserialization of type strings back to types."""
+    # Test built-in types
+    assert deserialize_type("builtins.int") is int
+    assert deserialize_type("builtins.str") is str
+    assert deserialize_type("builtins.float") is float
+    assert deserialize_type("builtins.bool") is bool
+    assert deserialize_type("builtins.list") is list
+    assert deserialize_type("builtins.dict") is dict
+    assert deserialize_type("builtins.tuple") is tuple
+    assert deserialize_type("builtins.set") is set
+
+
+def test_serialize_deserialize_roundtrip() -> None:
+    """Test that serialization and deserialization are inverse operations."""
+    # Test built-in types
+    types_to_test = [int, str, float, bool, list, dict, tuple, set]
+
+    for type_to_test in types_to_test:
+        serialized = serialize_type(type_to_test)
+        deserialized = deserialize_type(serialized)
+        assert deserialized is type_to_test
+
+    # Test agent framework type roundtrip
+
+    serialized = serialize_type(WorkflowEvent)
+    deserialized = deserialize_type(serialized)
+    assert deserialized is WorkflowEvent
+
+    # Verify we can instantiate the deserialized type via factory method
+    instance = WorkflowEvent.request_info(
+        request_id="request-123",
+        source_executor_id="executor_1",
+        request_data="test",
+        response_type=str,
+    )
+    assert isinstance(instance, WorkflowEvent)
+    assert instance.type == "request_info"
+
+
+def test_deserialize_type_error_handling() -> None:
+    """Test error handling in deserialize_type function."""
+    import pytest
+
+    # Test with non-existent module
+    with pytest.raises(ModuleNotFoundError):
+        deserialize_type("nonexistent.module.Type")
+
+    # Test with non-existent type in existing module
+    with pytest.raises(AttributeError):
+        deserialize_type("builtins.NonExistentType")
+
+
+def test_type_compatibility_basic() -> None:
+    """Test basic type compatibility scenarios."""
+    # Exact type match
+    assert is_type_compatible(str, str)
+    assert is_type_compatible(int, int)
+
+    # bool is a subtype of int
+    assert is_type_compatible(bool, int)
+
+    # Any compatibility
+    assert is_type_compatible(str, Any)
+    assert is_type_compatible(list[int], Any)
+
+    # Subclass compatibility
+    class Animal:
+        pass
+
+    class Dog(Animal):
+        pass
+
+    assert is_type_compatible(Dog, Animal)
+    assert not is_type_compatible(Animal, Dog)
+
+
+def test_type_compatibility_unions() -> None:
+    """Test type compatibility with Union types."""
+    # Source matches target union member
+    assert is_type_compatible(str, Union[str, int])
+    assert is_type_compatible(int, Union[str, int])
+    assert not is_type_compatible(float, Union[str, int])
+
+    # Source union - all members must be compatible with target
+    assert is_type_compatible(Union[str, int], Union[str, int, float])
+    assert not is_type_compatible(Union[str, int, bytes], Union[str, int])
+
+
+def test_type_compatibility_collections() -> None:
+    """Test type compatibility with collection types."""
+
+    # List compatibility - key use case
+    @dataclass
+    class Message:
+        text: str
+
+    assert is_type_compatible(list[Message], list[Union[str, Message]])
+    assert is_type_compatible(list[str], list[Union[str, Message]])
+    assert not is_type_compatible(list[Union[str, Message]], list[Message])
+
+    # Dict compatibility
+    assert is_type_compatible(dict[str, int], dict[str, Union[int, float]])
+    assert not is_type_compatible(dict[str, Union[int, float]], dict[str, int])
+
+    # Set compatibility
+    assert is_type_compatible(set[str], set[Union[str, int]])
+    assert not is_type_compatible(set[Union[str, int]], set[str])
+
+
+def test_type_compatibility_tuples() -> None:
+    """Test type compatibility with tuple types."""
+    # Fixed length tuples
+    assert is_type_compatible(tuple[str, int], tuple[Union[str, bytes], Union[int, float]])
+    assert not is_type_compatible(tuple[str, int], tuple[str, int, bool])  # Different lengths
+
+    # Variable length tuples
+    assert is_type_compatible(tuple[str, ...], tuple[Union[str, bytes], ...])
+    assert is_type_compatible(tuple[str, int, bool], tuple[Union[str, int, bool], ...])
+    assert not is_type_compatible(tuple[str, ...], tuple[str, int])  # Variable to fixed
+
+
+def test_type_compatibility_complex() -> None:
+    """Test complex nested type compatibility."""
+
+    @dataclass
+    class Message:
+        content: str
+
+    # Complex nested structure
+    source = list[dict[str, Message]]
+    target = list[dict[Union[str, bytes], Union[str, Message]]]
+    assert is_type_compatible(source, target)
+
+    # Incompatible nested structure
+    incompatible_target = list[dict[Union[str, bytes], int]]
+    assert not is_type_compatible(source, incompatible_target)
+
+
+# region: try_coerce_to_type tests
+
+
+def test_coerce_already_correct_type() -> None:
+    """Values already matching the target type are returned as-is."""
+    assert try_coerce_to_type(42, int) == 42
+    assert try_coerce_to_type("hello", str) == "hello"
+    assert try_coerce_to_type(True, bool) is True
+
+
+def test_coerce_int_to_float() -> None:
+    """JSON integers should be coercible to float."""
+    result = try_coerce_to_type(1, float)
+    assert result == 1.0
+    assert isinstance(result, float)
+
+
+def test_coerce_dict_to_dataclass() -> None:
+    """Dicts (from JSON) should be coercible to dataclasses."""
+
+    @dataclass
+    class Point:
+        x: int
+        y: int
+
+    result = try_coerce_to_type({"x": 1, "y": 2}, Point)
+    assert isinstance(result, Point)
+    assert result.x == 1
+    assert result.y == 2
+
+
+def test_coerce_dict_to_dataclass_bad_keys_returns_original() -> None:
+    """Dicts with wrong keys should return the original dict, not raise."""
+
+    @dataclass
+    class Point:
+        x: int
+        y: int
+
+    original = {"a": 1, "b": 2}
+    result = try_coerce_to_type(original, Point)
+    assert result is original
+
+
+def test_coerce_non_concrete_target_returns_original() -> None:
+    """Union and other non-concrete types should return the original value."""
+    result = try_coerce_to_type(42, int | str)
+    assert result == 42
+
+    result = try_coerce_to_type({"x": 1}, Union[str, int])
+    assert result == {"x": 1}
+
+
+def test_coerce_unrelated_types_returns_original() -> None:
+    """Coercion between unrelated types should return the original value."""
+    assert try_coerce_to_type("hello", int) == "hello"
+    assert try_coerce_to_type(3.14, str) == 3.14
+    assert try_coerce_to_type([1, 2], dict) == [1, 2]
+
+
+def test_coerce_any_returns_original() -> None:
+    """Any target type should accept any value without coercion."""
+    assert try_coerce_to_type(42, Any) == 42
+    assert try_coerce_to_type({"k": "v"}, Any) == {"k": "v"}
+
+
+# endregion: try_coerce_to_type tests

python/packages/core/tests/workflow/test_validation.py

@@ -0,0 +1,736 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import logging
+from typing import Any
+
+import pytest
+
+from agent_framework import (
+    EdgeDuplicationError,
+    Executor,
+    GraphConnectivityError,
+    TypeCompatibilityError,
+    ValidationTypeEnum,
+    WorkflowBuilder,
+    WorkflowContext,
+    WorkflowValidationError,
+    handler,
+    validate_workflow_graph,
+)
+from agent_framework._workflows._edge import SingleEdgeGroup
+
+
+class StringExecutor(Executor):
+    @handler
+    async def handle_string(self, message: str, ctx: WorkflowContext[str]) -> None:
+        await ctx.send_message(message.upper())
+
+
+class StringAggregator(Executor):
+    """A mock executor that aggregates results from multiple executors."""
+
+    @handler
+    async def mock_handler(self, messages: list[str], ctx: WorkflowContext[str]) -> None:
+        # This mock simply returns the data incremented by 1
+        await ctx.send_message("Aggregated: " + ", ".join(messages))
+
+
+class IntExecutor(Executor):
+    @handler
+    async def handle_int(self, message: int, ctx: WorkflowContext[int]) -> None:
+        await ctx.send_message(message * 2)
+
+
+class AnyExecutor(Executor):
+    @handler
+    async def handle_any(self, message: Any, ctx: WorkflowContext[Any]) -> None:
+        await ctx.send_message(f"Processed: {message}")
+
+
+class NoOutputTypesExecutor(Executor):
+    @handler
+    async def handle_message(self, message: str, ctx: WorkflowContext) -> None:
+        await ctx.send_message("processed")  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+
+
+class MultiTypeExecutor(Executor):
+    @handler
+    async def handle_string(self, message: str, ctx: WorkflowContext[str]) -> None:
+        await ctx.send_message(f"String: {message}")
+
+    @handler
+    async def handle_int(self, message: int, ctx: WorkflowContext[str]) -> None:
+        await ctx.send_message(f"Int: {message}")
+
+
+def test_valid_workflow_passes_validation():
+    executor1 = StringExecutor(id="string_executor")
+    executor2 = StringExecutor(id="string_executor_2")
+
+    # Create a valid workflow
+    workflow = (
+        WorkflowBuilder(start_executor=executor1)
+        .add_edge(executor1, executor2)
+        .build()  # This should not raise any exceptions
+    )
+
+    assert workflow is not None
+
+
+def test_duplicate_executor_ids_fail_validation():
+    executor1 = StringExecutor(id="dup")
+    executor2 = IntExecutor(id="dup")
+
+    with pytest.raises(ValueError) as exc_info:
+        (WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build())
+
+    assert str(exc_info.value) == "Duplicate executor ID 'dup' detected in workflow."
+
+
+def test_edge_duplication_validation_fails():
+    executor1 = StringExecutor(id="executor1")
+    executor2 = StringExecutor(id="executor2")
+
+    with pytest.raises(EdgeDuplicationError) as exc_info:
+        WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).add_edge(executor1, executor2).build()
+
+    assert "executor1->executor2" in str(exc_info.value)
+    assert exc_info.value.validation_type == ValidationTypeEnum.EDGE_DUPLICATION
+
+
+def test_type_compatibility_validation_fails():
+    string_executor = StringExecutor(id="string_executor")
+    int_executor = IntExecutor(id="int_executor")
+
+    with pytest.raises(TypeCompatibilityError) as exc_info:
+        WorkflowBuilder(start_executor=string_executor).add_edge(string_executor, int_executor).build()
+
+    error = exc_info.value
+    assert error.source_executor_id == "string_executor"
+    assert error.target_executor_id == "int_executor"
+    assert error.validation_type == ValidationTypeEnum.TYPE_COMPATIBILITY
+
+
+def test_type_compatibility_with_any_type_passes():
+    string_executor = StringExecutor(id="string_executor")
+    any_executor = AnyExecutor(id="any_executor")
+
+    # This should not raise an exception
+    workflow = WorkflowBuilder(start_executor=string_executor).add_edge(string_executor, any_executor).build()
+
+    assert workflow is not None
+
+
+def test_type_compatibility_with_no_output_types():
+    no_output_executor = NoOutputTypesExecutor(id="no_output")
+    string_executor = StringExecutor(id="string_executor")
+
+    # This should pass validation since no output types are specified
+    workflow = WorkflowBuilder(start_executor=no_output_executor).add_edge(no_output_executor, string_executor).build()
+
+    assert workflow is not None
+
+
+def test_multi_type_executor_compatibility():
+    string_executor = StringExecutor(id="string_executor")
+    multi_type_executor = MultiTypeExecutor(id="multi_type")
+
+    # String executor outputs strings, multi-type can handle strings
+    workflow = WorkflowBuilder(start_executor=string_executor).add_edge(string_executor, multi_type_executor).build()
+
+    assert workflow is not None
+
+
+def test_graph_connectivity_unreachable_executors():
+    executor1 = StringExecutor(id="executor1")
+    executor2 = StringExecutor(id="executor2")
+    executor3 = StringExecutor(id="executor3")  # This will be unreachable
+
+    with pytest.raises(GraphConnectivityError) as exc_info:
+        WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).add_edge(executor3, executor2).build()
+
+    assert "unreachable" in str(exc_info.value).lower()
+    assert "executor3" in str(exc_info.value)
+    assert exc_info.value.validation_type == ValidationTypeEnum.GRAPH_CONNECTIVITY
+
+
+def test_graph_connectivity_isolated_executors():
+    executor1 = StringExecutor(id="executor1")
+    executor2 = StringExecutor(id="executor2")
+    executor3 = StringExecutor(id="executor3")  # This will be isolated
+
+    # Create edges that include an isolated executor (self-loop that's not connected to main graph)
+    edge_groups = [
+        SingleEdgeGroup(executor1.id, executor2.id),
+        SingleEdgeGroup(executor3.id, executor3.id),
+    ]  # Self-loop to include in graph
+
+    executors: dict[str, Executor] = {executor1.id: executor1, executor2.id: executor2, executor3.id: executor3}
+
+    with pytest.raises(GraphConnectivityError) as exc_info:
+        validate_workflow_graph(edge_groups, executors, executor1, [])
+
+    assert "unreachable" in str(exc_info.value).lower()
+    assert "executor3" in str(exc_info.value)
+
+
+def test_disconnected_start_executor_not_in_graph():
+    executor1 = StringExecutor(id="executor1")
+    executor2 = StringExecutor(id="executor2")
+    executor3 = StringExecutor(id="executor3")  # Not in graph
+
+    with pytest.raises(GraphConnectivityError) as exc_info:
+        WorkflowBuilder(start_executor=executor3).add_edge(executor1, executor2).build()
+
+    assert "The following executors are unreachable from the start executor 'executor3'" in str(exc_info.value)
+
+
+def test_missing_start_executor():
+    with pytest.raises(TypeError):
+        WorkflowBuilder()  # type: ignore[call-arg]  # ty: ignore[missing-argument]
+
+
+def test_workflow_validation_error_base_class():
+    error = WorkflowValidationError("Test message", ValidationTypeEnum.EDGE_DUPLICATION)
+    assert str(error) == "[EDGE_DUPLICATION] Test message"
+    assert error.message == "Test message"
+    assert error.validation_type == ValidationTypeEnum.EDGE_DUPLICATION
+
+
+def test_complex_workflow_validation():
+    # Create a workflow with multiple paths
+    executor1 = StringExecutor(id="executor1")
+    executor2 = MultiTypeExecutor(id="executor2")
+    executor3 = StringExecutor(id="executor3")
+    executor4 = AnyExecutor(id="executor4")
+
+    workflow = (
+        WorkflowBuilder(start_executor=executor1)
+        .add_edge(executor1, executor2)  # str -> MultiType (compatible)
+        .add_edge(executor2, executor3)  # MultiType -> str (compatible)
+        .add_edge(executor2, executor4)  # MultiType -> Any (compatible)
+        .add_edge(executor3, executor4)  # str -> Any (compatible)
+        .build()
+    )
+
+    assert workflow is not None
+
+
+def test_type_compatibility_inheritance():
+    class BaseExecutor(Executor):
+        @handler
+        async def handle_base(self, message: str, ctx: WorkflowContext[str]) -> None:
+            await ctx.send_message("base")
+
+    class DerivedExecutor(Executor):
+        @handler
+        async def handle_derived(self, message: str, ctx: WorkflowContext[str]) -> None:
+            await ctx.send_message("derived")
+
+    base_executor = BaseExecutor(id="base")
+    derived_executor = DerivedExecutor(id="derived")
+
+    # This should pass since both handle str
+    workflow = WorkflowBuilder(start_executor=base_executor).add_edge(base_executor, derived_executor).build()
+
+    assert workflow is not None
+
+
+def test_direct_validation_function():
+    executor1 = StringExecutor(id="executor1")
+    executor2 = StringExecutor(id="executor2")
+    edge_groups = [SingleEdgeGroup(executor1.id, executor2.id)]
+    executors: dict[str, Executor] = {executor1.id: executor1, executor2.id: executor2}
+
+    # This should not raise any exceptions
+    validate_workflow_graph(edge_groups, executors, executor1, [])
+
+    # Test with invalid start executor
+    executor3 = StringExecutor(id="executor3")
+    with pytest.raises(GraphConnectivityError):
+        validate_workflow_graph(edge_groups, executors, executor3, [])
+
+
+def test_fan_out_validation():
+    source = StringExecutor(id="source")
+    target1 = StringExecutor(id="target1")
+    target2 = AnyExecutor(id="target2")
+
+    workflow = WorkflowBuilder(start_executor=source).add_fan_out_edges(source, [target1, target2]).build()
+
+    assert workflow is not None
+
+
+def test_fan_in_validation():
+    start_executor = StringExecutor(id="start")
+    source1 = StringExecutor(id="source1")
+    source2 = StringExecutor(id="source2")
+    target = StringAggregator(id="target")
+
+    # Create a proper fan-in by having a start executor that connects to both sources
+    workflow = (
+        WorkflowBuilder(start_executor=start_executor)
+        .add_edge(start_executor, source1)  # Start connects to source1
+        .add_edge(start_executor, source2)  # Start connects to source2
+        .add_fan_in_edges([source1, source2], target)  # Both sources fan-in to target
+        .build()
+    )
+
+    assert workflow is not None
+
+
+def test_chain_validation():
+    executor1 = StringExecutor(id="executor1")
+    executor2 = StringExecutor(id="executor2")
+    executor3 = AnyExecutor(id="executor3")
+
+    workflow = WorkflowBuilder(start_executor=executor1).add_chain([executor1, executor2, executor3]).build()
+
+    assert workflow is not None
+
+
+def test_logging_for_missing_output_types(caplog: Any) -> None:
+    caplog.set_level(logging.WARNING)
+
+    # Create executor without output types
+    no_output_executor = NoOutputTypesExecutor(id="no_output")
+    string_executor = StringExecutor(id="string_executor")
+
+    # This should trigger a warning log
+    workflow = WorkflowBuilder(start_executor=no_output_executor).add_edge(no_output_executor, string_executor).build()
+
+    assert workflow is not None
+    assert "has no output type annotations" in caplog.text
+    assert "Consider adding WorkflowContext[T] generics" in caplog.text
+
+
+def test_logging_for_missing_input_types(caplog: Any) -> None:
+    caplog.set_level(logging.WARNING)
+
+    class NoInputTypesExecutor(Executor):
+        # Handler without type annotation for input parameter
+        async def handle_message(self, message: Any, ctx: WorkflowContext[Any]) -> None:
+            await ctx.send_message("processed")
+
+        def _discover_handlers(self) -> None:
+            # Override to manually register handler without type info
+            self._handlers[str] = self.handle_message
+
+    string_executor = StringExecutor(id="string_executor")
+    no_input_executor = NoInputTypesExecutor(id="no_input")
+
+    # This should pass since NoInputTypesExecutor has no proper input types
+    workflow = WorkflowBuilder(start_executor=string_executor).add_edge(string_executor, no_input_executor).build()
+
+    assert workflow is not None
+
+
+def test_self_loop_detection_warning(caplog: Any) -> None:
+    caplog.set_level(logging.WARNING)
+
+    executor = StringExecutor(id="self_loop_executor")
+
+    # Create a self-loop
+    workflow = WorkflowBuilder(start_executor=executor).add_edge(executor, executor).build()
+
+    assert workflow is not None
+    assert "Self-loop detected" in caplog.text
+    assert "may cause infinite recursion" in caplog.text
+
+
+def test_handler_validation_basic(caplog: Any) -> None:
+    caplog.set_level(logging.WARNING)
+
+    # Test basic handler validation - ensure the validation code runs without errors
+    start_executor = StringExecutor(id="start")
+    target_executor = StringExecutor(id="target")
+
+    workflow = WorkflowBuilder(start_executor=start_executor).add_edge(start_executor, target_executor).build()
+
+    assert workflow is not None
+    # Just ensure the validation runs without errors
+
+
+def test_dead_end_detection(caplog: Any) -> None:
+    caplog.set_level(logging.INFO)
+
+    executor1 = StringExecutor(id="executor1")
+    executor2 = StringExecutor(id="executor2")  # This will be a dead end
+
+    workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
+
+    assert workflow is not None
+    assert "Dead-end executors detected" in caplog.text
+    assert "executor2" in caplog.text
+    assert "Verify these are intended as final nodes" in caplog.text
+
+
+def test_successful_type_compatibility_logging(caplog: Any) -> None:
+    caplog.set_level(logging.DEBUG)
+
+    executor1 = StringExecutor(id="executor1")
+    executor2 = StringExecutor(id="executor2")
+
+    workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
+
+    assert workflow is not None
+    assert "Type compatibility validated for edge" in caplog.text
+    assert "Compatible type pairs" in caplog.text
+
+
+def test_multiple_dead_ends_detection(caplog: Any) -> None:
+    caplog.set_level(logging.INFO)
+
+    executor1 = StringExecutor(id="executor1")
+    executor2 = StringExecutor(id="executor2")  # Dead end
+    executor3 = StringExecutor(id="executor3")  # Dead end
+
+    workflow = (
+        WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).add_edge(executor1, executor3).build()
+    )
+
+    assert workflow is not None
+    assert "Dead-end executors detected" in caplog.text
+    assert "executor2" in caplog.text and "executor3" in caplog.text
+
+
+def test_single_executor_workflow(caplog: Any) -> None:
+    caplog.set_level(logging.INFO)
+
+    # Test workflow with minimal structure
+    executor1 = StringExecutor(id="executor1")
+    executor2 = StringExecutor(id="executor2")
+
+    # Create a simple two-executor workflow to avoid graph validation issues
+    workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
+
+    assert workflow is not None
+    # Should detect executor2 as dead end
+    assert "Dead-end executors detected" in caplog.text
+
+
+def test_enhanced_type_compatibility_error_details():
+    string_executor = StringExecutor(id="string_executor")
+    int_executor = IntExecutor(id="int_executor")
+
+    with pytest.raises(TypeCompatibilityError) as exc_info:
+        WorkflowBuilder(start_executor=string_executor).add_edge(string_executor, int_executor).build()
+
+    error = exc_info.value
+    # Verify enhanced error contains detailed type information
+    assert "Source executor outputs types" in str(error)
+    assert "target executor can only handle types" in str(error)
+    assert error.source_types is not None
+    assert error.target_types is not None
+
+
+def test_union_type_compatibility_validation() -> None:
+    class UnionOutputExecutor(Executor):
+        @handler
+        async def handle_message(self, message: str, ctx: WorkflowContext[str | int]) -> None:
+            await ctx.send_message("output")
+
+    class UnionInputExecutor(Executor):
+        @handler
+        async def handle_message(self, message: str, ctx: WorkflowContext[str]) -> None:
+            await ctx.send_message("processed")
+
+    union_output = UnionOutputExecutor(id="union_output")
+    union_input = UnionInputExecutor(id="union_input")
+
+    # This should pass validation due to type compatibility (str)
+    workflow = WorkflowBuilder(start_executor=union_output).add_edge(union_output, union_input).build()
+
+    assert workflow is not None
+
+
+def test_generic_type_compatibility() -> None:
+    class ListOutputExecutor(Executor):
+        @handler
+        async def handle_message(self, message: str, ctx: WorkflowContext[list[str]]) -> None:
+            await ctx.send_message(["output"])
+
+    class ListInputExecutor(Executor):
+        @handler
+        async def handle_message(self, message: list[str], ctx: WorkflowContext[str]) -> None:
+            await ctx.send_message("processed")
+
+    list_output = ListOutputExecutor(id="list_output")
+    list_input = ListInputExecutor(id="list_input")
+
+    # This should pass validation for generic type compatibility
+    workflow = WorkflowBuilder(start_executor=list_output).add_edge(list_output, list_input).build()
+
+    assert workflow is not None
+
+
+def test_validation_enum_usage() -> None:
+    # Test that all validation types use the enum correctly
+    edge_error = EdgeDuplicationError("test->test")
+    assert edge_error.validation_type == ValidationTypeEnum.EDGE_DUPLICATION
+
+    type_error = TypeCompatibilityError("source", "target", [str], [int])
+    assert type_error.validation_type == ValidationTypeEnum.TYPE_COMPATIBILITY
+
+    graph_error = GraphConnectivityError("test message")
+    assert graph_error.validation_type == ValidationTypeEnum.GRAPH_CONNECTIVITY
+
+    # Test enum string representation
+    assert str(ValidationTypeEnum.EDGE_DUPLICATION) == "ValidationTypeEnum.EDGE_DUPLICATION"
+    assert ValidationTypeEnum.EDGE_DUPLICATION.value == "EDGE_DUPLICATION"
+
+
+def test_handler_ctx_missing_annotation_raises() -> None:
+    # Validation now happens at handler registration time, not workflow build time
+    with pytest.raises(ValueError) as exc:
+
+        class BadExecutor(Executor):  # pyright: ignore[reportUnusedClass]
+            @handler  # pyright: ignore[reportUnknownArgumentType]
+            async def handle(self, message: str, ctx) -> None:  # type: ignore[no-untyped-def]
+                pass
+
+    assert "must have a WorkflowContext" in str(exc.value)
+
+
+def test_handler_ctx_invalid_t_out_entries_raises() -> None:
+    # Validation now happens at handler registration time, not workflow build time
+    with pytest.raises(ValueError) as exc:
+
+        class BadExecutor(Executor):  # pyright: ignore[reportUnusedClass]
+            @handler  # pyright: ignore[reportUnknownArgumentType]
+            async def handle(self, message: str, ctx: WorkflowContext[123]) -> None:  # type: ignore[valid-type]  # ty: ignore[invalid-type-form]
+                pass
+
+    assert "invalid type entry" in str(exc.value)
+
+
+def test_handler_ctx_none_is_allowed() -> None:
+    class NoneExecutor(Executor):
+        @handler
+        async def handle(self, message: str, ctx: WorkflowContext) -> None:
+            # does not emit
+            return None
+
+    start = StringExecutor(id="s")
+    none_exec = NoneExecutor(id="n")
+
+    # Should build successfully
+    wf = WorkflowBuilder(start_executor=start).add_edge(start, none_exec).build()
+    assert wf is not None
+
+
+def test_handler_ctx_any_is_allowed_but_skips_type_checks(caplog: Any) -> None:
+    caplog.set_level(logging.WARNING)
+
+    class AnyOutExecutor(Executor):
+        @handler
+        async def handle(self, message: str, ctx: WorkflowContext[Any]) -> None:
+            return None
+
+    start = StringExecutor(id="s")
+    any_out = AnyOutExecutor(id="a")
+
+    # Builds; later edges from this executor will skip type compatibility when outputs are unspecified
+    wf = WorkflowBuilder(start_executor=start).add_edge(start, any_out).build()
+    assert wf is not None
+
+
+# region Output Validation Tests
+
+
+class OutputExecutor(Executor):
+    @handler
+    async def handle_string(self, message: str, ctx: WorkflowContext[str, str]) -> None:
+        pass
+
+
+def test_output_validation_with_valid_output_executors():
+    """Test that output validation passes when output executors exist and have output types."""
+    executor1 = OutputExecutor(id="executor1")
+    executor2 = OutputExecutor(id="executor2")
+
+    # Build workflow with valid output executors
+    workflow = WorkflowBuilder(start_executor=executor1, output_from=[executor2]).add_edge(executor1, executor2).build()
+
+    assert workflow is not None
+    assert {ex.id for ex in workflow.get_output_executors()} == {"executor2"}
+
+
+def test_output_validation_with_multiple_valid_output_executors():
+    """Test that output validation passes with multiple valid output executors."""
+    executor1 = OutputExecutor(id="executor1")
+    executor2 = OutputExecutor(id="executor2")
+    executor3 = OutputExecutor(id="executor3")
+
+    workflow = (
+        WorkflowBuilder(start_executor=executor1, output_from=[executor1, executor3])
+        .add_edge(executor1, executor2)
+        .add_edge(executor2, executor3)
+        .build()
+    )
+
+    assert workflow is not None
+    assert {ex.id for ex in workflow.get_output_executors()} == {"executor1", "executor3"}
+
+
+def test_output_validation_fails_for_nonexistent_executor():
+    """Test that output validation fails when an output executor doesn't exist in the graph."""
+    executor1 = OutputExecutor(id="executor1")
+    executor2 = OutputExecutor(id="executor2")
+    edge_groups = [SingleEdgeGroup(executor1.id, executor2.id)]
+    executors: dict[str, Executor] = {executor1.id: executor1, executor2.id: executor2}
+
+    # Directly test validation with a nonexistent output executor
+    with pytest.raises(WorkflowValidationError) as exc_info:
+        validate_workflow_graph(edge_groups, executors, executor1, ["nonexistent_executor"])
+
+    assert "not present in the workflow graph" in str(exc_info.value)
+    assert "nonexistent_executor" in str(exc_info.value)
+    assert exc_info.value.validation_type == ValidationTypeEnum.OUTPUT_VALIDATION
+
+
+def test_output_validation_fails_for_executor_without_output_types():
+    """Test that output validation fails when an output executor has no output type annotations."""
+    executor1 = OutputExecutor(id="executor1")
+    no_output_executor = NoOutputTypesExecutor(id="no_output")
+
+    with pytest.raises(WorkflowValidationError) as exc_info:
+        (
+            WorkflowBuilder(start_executor=executor1, output_from=[no_output_executor])
+            .add_edge(executor1, no_output_executor)
+            .build()
+        )
+
+    assert "must have output type annotations defined" in str(exc_info.value)
+    assert "no_output" in str(exc_info.value)
+    assert exc_info.value.validation_type == ValidationTypeEnum.OUTPUT_VALIDATION
+
+
+def test_output_validation_empty_explicit_designation_fails():
+    """Test that explicit mode rejects an empty output/intermediate designation."""
+    executor1 = OutputExecutor(id="executor1")
+    executor2 = OutputExecutor(id="executor2")
+
+    with pytest.raises(WorkflowValidationError) as exc_info:
+        WorkflowBuilder(start_executor=executor1, output_from=[]).add_edge(executor1, executor2).build()
+
+    assert "at least one output or intermediate executor" in str(exc_info.value)
+    assert exc_info.value.validation_type == ValidationTypeEnum.OUTPUT_VALIDATION
+
+
+def test_output_validation_with_valid_intermediate_executors():
+    """Test that output validation passes when intermediate executors exist and have output types."""
+    executor1 = OutputExecutor(id="executor1")
+    executor2 = OutputExecutor(id="executor2")
+
+    workflow = (
+        WorkflowBuilder(start_executor=executor1, intermediate_output_from=[executor1])
+        .add_edge(executor1, executor2)
+        .build()
+    )
+
+    assert workflow is not None
+    assert {ex.id for ex in workflow.get_intermediate_executors()} == {"executor1"}
+    assert workflow.is_intermediate_executor("executor1")
+    assert not workflow.is_terminal_executor("executor2")
+
+
+def test_output_validation_fails_for_designation_overlap():
+    """Test that an executor cannot be both terminal and intermediate."""
+    executor1 = OutputExecutor(id="executor1")
+
+    with pytest.raises(WorkflowValidationError) as exc_info:
+        WorkflowBuilder(
+            start_executor=executor1,
+            output_from=[executor1],
+            intermediate_output_from=[executor1],
+        ).build()
+
+    assert "both output and intermediate" in str(exc_info.value)
+    assert exc_info.value.validation_type == ValidationTypeEnum.OUTPUT_VALIDATION
+
+
+def test_output_validation_fails_for_duplicate_designation():
+    """Test that duplicate output or intermediate designation entries are rejected."""
+    executor1 = OutputExecutor(id="executor1")
+
+    with pytest.raises(WorkflowValidationError) as exc_info:
+        WorkflowBuilder(start_executor=executor1, output_from=[executor1, executor1]).build()
+
+    assert "Duplicate output executor designation" in str(exc_info.value)
+    assert exc_info.value.validation_type == ValidationTypeEnum.OUTPUT_VALIDATION
+
+
+def test_output_validation_fails_for_unknown_intermediate_executor():
+    """Test that intermediate designation rejects executors outside the workflow graph."""
+    executor1 = OutputExecutor(id="executor1")
+    executor2 = OutputExecutor(id="executor2")
+    missing = OutputExecutor(id="missing")
+
+    with pytest.raises(WorkflowValidationError) as exc_info:
+        (
+            WorkflowBuilder(start_executor=executor1, intermediate_output_from=[missing])
+            .add_edge(executor1, executor2)
+            .build()
+        )
+
+    assert "not present in the workflow graph" in str(exc_info.value)
+    assert "missing" in str(exc_info.value)
+    assert exc_info.value.validation_type == ValidationTypeEnum.OUTPUT_VALIDATION
+
+
+def test_output_validation_with_direct_validate_workflow_graph():
+    """Test _output_validation directly via validate_workflow_graph function."""
+    executor1 = OutputExecutor(id="executor1")
+    executor2 = OutputExecutor(id="executor2")
+    edge_groups = [SingleEdgeGroup(executor1.id, executor2.id)]
+    executors: dict[str, Executor] = {executor1.id: executor1, executor2.id: executor2}
+
+    # Valid output executors
+    validate_workflow_graph(edge_groups, executors, executor1, ["executor2"])
+
+    # Invalid output executor (doesn't exist)
+    with pytest.raises(WorkflowValidationError) as exc_info:
+        validate_workflow_graph(edge_groups, executors, executor1, ["nonexistent"])
+
+    assert "not present in the workflow graph" in str(exc_info.value)
+    assert exc_info.value.validation_type == ValidationTypeEnum.OUTPUT_VALIDATION
+
+
+def test_output_validation_with_no_output_types_via_direct_validation():
+    """Test _output_validation fails for executors without output types via direct validation."""
+    executor1 = OutputExecutor(id="executor1")
+    no_output_executor = NoOutputTypesExecutor(id="no_output")
+    edge_groups = [SingleEdgeGroup(executor1.id, no_output_executor.id)]
+    executors: dict[str, Executor] = {executor1.id: executor1, no_output_executor.id: no_output_executor}
+
+    # Should fail because no_output_executor has no output types
+    with pytest.raises(WorkflowValidationError) as exc_info:
+        validate_workflow_graph(edge_groups, executors, executor1, ["no_output"])
+
+    assert "must have output type annotations defined" in str(exc_info.value)
+    assert exc_info.value.validation_type == ValidationTypeEnum.OUTPUT_VALIDATION
+
+
+def test_output_validation_partial_invalid_list():
+    """Test that output validation fails if any executor in the list is invalid."""
+    executor1 = OutputExecutor(id="executor1")
+    executor2 = OutputExecutor(id="executor2")
+    edge_groups = [SingleEdgeGroup(executor1.id, executor2.id)]
+    executors: dict[str, Executor] = {executor1.id: executor1, executor2.id: executor2}
+
+    # First executor is valid, second doesn't exist - validation should fail
+    with pytest.raises(WorkflowValidationError) as exc_info:
+        validate_workflow_graph(edge_groups, executors, executor1, ["executor2", "nonexistent"])
+
+    assert "not present in the workflow graph" in str(exc_info.value)
+    assert "nonexistent" in str(exc_info.value)
+
+
+def test_output_validation_type_enum_value():
+    """Test that OUTPUT_VALIDATION is properly defined in ValidationTypeEnum."""
+    assert hasattr(ValidationTypeEnum, "OUTPUT_VALIDATION")
+    assert ValidationTypeEnum.OUTPUT_VALIDATION.value == "OUTPUT_VALIDATION"
+
+
+# endregion

python/packages/core/tests/workflow/test_viz.py

@@ -0,0 +1,410 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for the workflow visualization module."""
+
+from pathlib import Path
+from typing import Any
+
+import pytest
+
+from agent_framework import Executor, WorkflowBuilder, WorkflowContext, WorkflowExecutor, WorkflowViz, handler
+
+
+class MockExecutor(Executor):
+    """A mock executor for testing purposes."""
+
+    @handler
+    async def mock_handler(self, message: str, ctx: WorkflowContext) -> None:
+        """A mock handler that does nothing."""
+        pass
+
+
+class ListStrTargetExecutor(Executor):
+    """A mock executor that accepts a list of strings (for fan-in targets)."""
+
+    @handler
+    async def handle(self, message: list[str], ctx: WorkflowContext) -> None:
+        pass
+
+
+@pytest.fixture
+def basic_sub_workflow() -> dict[str, Any]:
+    """Fixture that creates a basic sub-workflow setup for testing."""
+    # Create a sub-workflow
+    sub_exec1 = MockExecutor(id="sub_exec1")
+    sub_exec2 = MockExecutor(id="sub_exec2")
+
+    sub_workflow = WorkflowBuilder(start_executor=sub_exec1).add_edge(sub_exec1, sub_exec2).build()
+
+    # Create a workflow executor that wraps the sub-workflow
+    workflow_executor = WorkflowExecutor(sub_workflow, id="workflow_executor_1")
+
+    # Create a main workflow that includes the workflow executor
+    main_exec = MockExecutor(id="main_executor")
+    final_exec = MockExecutor(id="final_executor")
+
+    main_workflow = (
+        WorkflowBuilder(start_executor=main_exec)
+        .add_edge(main_exec, workflow_executor)
+        .add_edge(workflow_executor, final_exec)
+        .build()
+    )
+
+    return {
+        "main_workflow": main_workflow,
+        "workflow_executor": workflow_executor,
+        "sub_workflow": sub_workflow,
+        "main_exec": main_exec,
+        "final_exec": final_exec,
+        "sub_exec1": sub_exec1,
+        "sub_exec2": sub_exec2,
+    }
+
+
+def test_workflow_viz_to_digraph():
+    """Test that WorkflowViz can generate a DOT digraph."""
+    # Create a simple workflow
+    executor1 = MockExecutor(id="executor1")
+    executor2 = MockExecutor(id="executor2")
+
+    workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
+
+    viz = WorkflowViz(workflow)
+    dot_content = viz.to_digraph()
+
+    # Check that the DOT content contains expected elements
+    assert "digraph Workflow {" in dot_content
+    assert '"executor1"' in dot_content
+    assert '"executor2"' in dot_content
+    assert '"executor1" -> "executor2"' in dot_content
+    assert "fillcolor=lightgreen" in dot_content  # Start executor styling
+    assert "(Start)" in dot_content
+
+
+def test_workflow_viz_export_dot():
+    """Test exporting workflow as DOT format."""
+    executor1 = MockExecutor(id="executor1")
+    executor2 = MockExecutor(id="executor2")
+
+    workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
+
+    viz = WorkflowViz(workflow)
+
+    # Test export without filename (returns temporary file path)
+    file_path = viz.export(format="dot")
+    assert file_path.endswith(".dot")
+
+    with open(file_path, encoding="utf-8") as f:
+        content = f.read()
+
+    assert "digraph Workflow {" in content
+    assert '"executor1" -> "executor2"' in content
+
+
+def test_workflow_viz_export_dot_with_filename(tmp_path: Path):
+    """Test exporting workflow as DOT format with specified filename."""
+    executor1 = MockExecutor(id="executor1")
+    executor2 = MockExecutor(id="executor2")
+
+    workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
+
+    viz = WorkflowViz(workflow)
+
+    # Test export with filename
+    output_file = tmp_path / "test_workflow.dot"
+    result_path = viz.export(format="dot", filename=str(output_file))
+
+    assert result_path == str(output_file)
+    assert output_file.exists()
+
+    content = output_file.read_text(encoding="utf-8")
+    assert "digraph Workflow {" in content
+    assert '"executor1" -> "executor2"' in content
+
+
+def test_workflow_viz_complex_workflow():
+    """Test visualization of a more complex workflow."""
+    executor1 = MockExecutor(id="start")
+    executor2 = MockExecutor(id="middle1")
+    executor3 = MockExecutor(id="middle2")
+    executor4 = MockExecutor(id="end")
+
+    workflow = (
+        WorkflowBuilder(start_executor=executor1)
+        .add_edge(executor1, executor2)
+        .add_edge(executor1, executor3)
+        .add_edge(executor2, executor4)
+        .add_edge(executor3, executor4)
+        .build()
+    )
+
+    viz = WorkflowViz(workflow)
+    dot_content = viz.to_digraph()
+
+    # Check all executors are present
+    assert '"start"' in dot_content
+    assert '"middle1"' in dot_content
+    assert '"middle2"' in dot_content
+    assert '"end"' in dot_content
+
+    # Check all edges are present
+    assert '"start" -> "middle1"' in dot_content
+    assert '"start" -> "middle2"' in dot_content
+    assert '"middle1" -> "end"' in dot_content
+    assert '"middle2" -> "end"' in dot_content
+
+    # Check start executor has special styling
+    assert "fillcolor=lightgreen" in dot_content
+
+
+@pytest.mark.skipif(True, reason="Requires graphviz to be installed")
+def test_workflow_viz_export_svg():
+    """Test exporting workflow as SVG format. Skipped unless graphviz is available."""
+    executor1 = MockExecutor(id="executor1")
+    executor2 = MockExecutor(id="executor2")
+
+    workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
+
+    viz = WorkflowViz(workflow)
+
+    try:
+        file_path = viz.export(format="svg")
+        assert file_path.endswith(".svg")
+    except ImportError:
+        pytest.skip("graphviz not available")
+
+
+def test_workflow_viz_unsupported_format():
+    """Test that unsupported formats raise ValueError."""
+    executor1 = MockExecutor(id="executor1")
+    executor2 = MockExecutor(id="executor2")
+
+    workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
+
+    viz = WorkflowViz(workflow)
+
+    with pytest.raises(ValueError, match="Unsupported format: invalid"):
+        viz.export(format="invalid")  # type: ignore
+
+
+def test_workflow_viz_graphviz_binary_not_found():
+    """Test that missing graphviz binary raises ImportError with helpful message."""
+    import unittest.mock
+
+    # Skip test if graphviz package is not available
+    pytest.importorskip("graphviz")
+
+    executor1 = MockExecutor(id="executor1")
+    executor2 = MockExecutor(id="executor2")
+
+    workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
+    viz = WorkflowViz(workflow)
+
+    # Mock graphviz.Source.render to raise ExecutableNotFound
+    with unittest.mock.patch("graphviz.Source") as mock_source_class:
+        mock_source = unittest.mock.MagicMock()
+        mock_source_class.return_value = mock_source
+
+        # Import the ExecutableNotFound exception for the test
+        from graphviz.backend.execute import ExecutableNotFound  # type: ignore[import-not-found]
+
+        mock_source.render.side_effect = ExecutableNotFound("failed to execute PosixPath('dot')")
+
+        # Test that the proper ImportError is raised with helpful message
+        with pytest.raises(ImportError, match="The graphviz executables are not found"):
+            viz.export(format="svg")
+
+
+def test_workflow_viz_conditional_edge():
+    """Test that conditional edges are rendered dashed with a label."""
+    start = MockExecutor(id="start")
+    mid = MockExecutor(id="mid")
+    end = MockExecutor(id="end")
+
+    # Condition that is never used during viz, but presence should mark the edge
+    def only_if_foo(msg: str) -> bool:  # pragma: no cover - simple predicate
+        return msg == "foo"
+
+    wf = WorkflowBuilder(start_executor=start).add_edge(start, mid, condition=only_if_foo).add_edge(mid, end).build()
+
+    dot = WorkflowViz(wf).to_digraph()
+
+    # Conditional edge should be dashed and labeled
+    assert '"start" -> "mid" [style=dashed, label="conditional"];' in dot
+    # Non-conditional edge should be plain
+    assert '"mid" -> "end"' in dot
+    assert '"mid" -> "end" [style=dashed' not in dot
+
+
+def test_workflow_viz_fan_in_edge_group():
+    """Test that fan-in edges render an intermediate node with label and routed edges."""
+    start = MockExecutor(id="start")
+    s1 = MockExecutor(id="s1")
+    s2 = MockExecutor(id="s2")
+    t = ListStrTargetExecutor(id="t")
+
+    # Build a connected workflow: start fans out to s1 and s2, which then fan-in to t
+    wf = WorkflowBuilder(start_executor=start).add_fan_out_edges(start, [s1, s2]).add_fan_in_edges([s1, s2], t).build()
+
+    dot = WorkflowViz(wf).to_digraph()
+
+    # There should be a single fan-in node with special styling and label
+    lines = [line.strip() for line in dot.splitlines()]
+    fan_in_lines = [line for line in lines if "shape=ellipse" in line and 'label="fan-in"' in line]
+    assert len(fan_in_lines) == 1
+
+    # Extract the intermediate node id from the line: "<id>" [shape=ellipse, ... label="fan-in"];
+    fan_in_line = fan_in_lines[0]
+    first_quote = fan_in_line.find('"')
+    second_quote = fan_in_line.find('"', first_quote + 1)
+    assert first_quote != -1 and second_quote != -1
+    fan_in_node_id = fan_in_line[first_quote + 1 : second_quote]
+    assert fan_in_node_id  # non-empty
+
+    # Edges should be routed through the intermediate node, not direct to target
+    assert f'"s1" -> "{fan_in_node_id}";' in dot
+    assert f'"s2" -> "{fan_in_node_id}";' in dot
+    assert f'"{fan_in_node_id}" -> "t";' in dot
+
+    # Ensure direct edges are not present
+    assert '"s1" -> "t"' not in dot
+    assert '"s2" -> "t"' not in dot
+
+
+def test_workflow_viz_to_mermaid_basic():
+    """Mermaid: basic workflow nodes and edge are present with start label."""
+    executor1 = MockExecutor(id="executor1")
+    executor2 = MockExecutor(id="executor2")
+
+    workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
+    mermaid = WorkflowViz(workflow).to_mermaid()
+
+    # Start node and normal node
+    assert 'executor1["executor1 (Start)"]' in mermaid
+    assert 'executor2["executor2"]' in mermaid
+    # Edge uses sanitized ids (same as ids here)
+    assert "executor1 --> executor2" in mermaid
+
+
+def test_workflow_viz_mermaid_conditional_edge():
+    """Mermaid: conditional edges are dotted with a label."""
+    start = MockExecutor(id="start")
+    mid = MockExecutor(id="mid")
+
+    def only_if_foo(msg: str) -> bool:  # pragma: no cover - simple predicate
+        return msg == "foo"
+
+    wf = WorkflowBuilder(start_executor=start).add_edge(start, mid, condition=only_if_foo).build()
+    mermaid = WorkflowViz(wf).to_mermaid()
+
+    assert "start -. conditional .-> mid" in mermaid
+
+
+def test_workflow_viz_mermaid_fan_in_edge_group():
+    """Mermaid: fan-in uses an intermediate node and routes edges via it."""
+    start = MockExecutor(id="start")
+    s1 = MockExecutor(id="s1")
+    s2 = MockExecutor(id="s2")
+    t = ListStrTargetExecutor(id="t")
+
+    wf = WorkflowBuilder(start_executor=start).add_fan_out_edges(start, [s1, s2]).add_fan_in_edges([s1, s2], t).build()
+
+    mermaid = WorkflowViz(wf).to_mermaid()
+    lines = [line.strip() for line in mermaid.splitlines()]
+    # Find the fan-in node (line ends with ((fan-in)))
+    fan_lines = [ln for ln in lines if ln.endswith("((fan-in))")]
+    assert len(fan_lines) == 1
+    fan_line = fan_lines[0]
+    # fan_in node is emitted as: <id>((fan-in)) -> extract <id>
+    token = fan_line.strip()
+    suffix = "((fan-in))"
+    assert token.endswith(suffix)
+    fan_node_id = token[: -len(suffix)]
+    assert fan_node_id
+
+    # Ensure routing via the intermediate node
+    assert f"s1 --> {fan_node_id}" in mermaid
+    assert f"s2 --> {fan_node_id}" in mermaid
+    assert f"{fan_node_id} --> t" in mermaid
+
+    # Ensure direct edges to target are not present
+    assert "s1 --> t" not in mermaid
+    assert "s2 --> t" not in mermaid
+
+
+def test_workflow_viz_sub_workflow_digraph(basic_sub_workflow: dict[str, Any]):
+    """Test that WorkflowViz can visualize sub-workflows in DOT format."""
+    main_workflow = basic_sub_workflow["main_workflow"]
+
+    viz = WorkflowViz(main_workflow)
+    dot_content = viz.to_digraph()
+
+    # Check that main workflow nodes are present
+    assert "main_executor" in dot_content
+    assert "workflow_executor_1" in dot_content
+    assert "final_executor" in dot_content
+
+    # Check that sub-workflow is rendered as a cluster
+    assert "subgraph cluster_" in dot_content
+    assert "sub-workflow: workflow_executor_1" in dot_content
+
+    # Check that sub-workflow nodes are namespaced
+    assert '"workflow_executor_1/sub_exec1"' in dot_content
+    assert '"workflow_executor_1/sub_exec2"' in dot_content
+
+    # Check that sub-workflow edges are present
+    assert '"workflow_executor_1/sub_exec1" -> "workflow_executor_1/sub_exec2"' in dot_content
+
+
+def test_workflow_viz_sub_workflow_mermaid(basic_sub_workflow: dict[str, Any]):
+    """Test that WorkflowViz can visualize sub-workflows in Mermaid format."""
+    main_workflow = basic_sub_workflow["main_workflow"]
+
+    viz = WorkflowViz(main_workflow)
+    mermaid_content = viz.to_mermaid()
+
+    # Check that main workflow nodes are present
+    assert "main_executor" in mermaid_content
+    assert "workflow_executor_1" in mermaid_content
+    assert "final_executor" in mermaid_content
+
+    # Check that sub-workflow is rendered as a subgraph
+    assert "subgraph workflow_executor_1" in mermaid_content
+    assert "end" in mermaid_content
+
+    # Check that sub-workflow nodes are namespaced properly for Mermaid
+    assert "workflow_executor_1__sub_exec1" in mermaid_content
+    assert "workflow_executor_1__sub_exec2" in mermaid_content
+
+
+def test_workflow_viz_nested_sub_workflows():
+    """Test visualization of deeply nested sub-workflows."""
+    # Create innermost sub-workflow
+    inner_exec = MockExecutor(id="inner_exec")
+    inner_workflow = WorkflowBuilder(start_executor=inner_exec).build()
+
+    # Create middle sub-workflow that contains the inner one
+    inner_workflow_executor = WorkflowExecutor(inner_workflow, id="inner_wf_exec")
+    middle_exec = MockExecutor(id="middle_exec")
+
+    middle_workflow = WorkflowBuilder(start_executor=middle_exec).add_edge(middle_exec, inner_workflow_executor).build()
+
+    # Create outer workflow
+    middle_workflow_executor = WorkflowExecutor(middle_workflow, id="middle_wf_exec")
+    outer_exec = MockExecutor(id="outer_exec")
+
+    outer_workflow = WorkflowBuilder(start_executor=outer_exec).add_edge(outer_exec, middle_workflow_executor).build()
+
+    viz = WorkflowViz(outer_workflow)
+    dot_content = viz.to_digraph()
+
+    # Check that all levels are present
+    assert "outer_exec" in dot_content
+    assert "middle_wf_exec" in dot_content
+    assert "inner_wf_exec" in dot_content
+
+    # Check for nested clusters
+    assert "subgraph cluster_" in dot_content
+    # Should have multiple subgraphs for nested structure
+    subgraph_count = dot_content.count("subgraph cluster_")
+    assert subgraph_count >= 2  # At least one for each level of nesting

python/packages/core/tests/workflow/test_workflow_agent_intermediate.py

@@ -0,0 +1,353 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for WorkflowAgent forwarding of intermediate workflow events.
+
+Covers:
+- type='intermediate' surfaces as AgentResponseUpdate without content-type rewriting
+- type='data' (compatibility alias via WorkflowEvent.emit) is forwarded
+- Message.additional_properties survives the intermediate translation path
+- Terminal yields keep using regular text content (backward compat)
+"""
+
+from __future__ import annotations
+
+import warnings
+
+import pytest
+from typing_extensions import Never
+
+from agent_framework import (
+    AgentResponse,
+    AgentResponseUpdate,
+    Content,
+    Message,
+    WorkflowBuilder,
+    WorkflowContext,
+    WorkflowEvent,
+    executor,
+)
+from agent_framework.exceptions import AgentInvalidRequestException
+
+
+@pytest.mark.asyncio
+async def test_workflow_agent_forwards_intermediate_events_without_content_rewrite() -> None:
+    """An intermediate yield from an intermediate-designated executor surfaces through as_agent
+    as an AgentResponseUpdate carrying its original content type."""
+
+    @executor
+    async def emit(messages: list[Message], ctx: WorkflowContext[str, str]) -> None:
+        await ctx.yield_output("intermediate progress")
+        await ctx.send_message("downstream")
+
+    @executor
+    async def terminal(message: str, ctx: WorkflowContext[Never, str]) -> None:  # type: ignore[valid-type]
+        await ctx.yield_output("FINAL")
+
+    workflow = (
+        WorkflowBuilder(
+            start_executor=emit,
+            output_from=[terminal],
+            intermediate_output_from=[emit],
+        )
+        .add_edge(emit, terminal)
+        .build()
+    )
+    agent = workflow.as_agent("test")
+
+    updates: list[AgentResponseUpdate] = []
+    async for update in agent.run("hi", stream=True):
+        updates.append(update)
+
+    text = " ".join(c.text for u in updates for c in u.contents if c.type == "text")  # type: ignore[misc]  # pyrefly: ignore[no-matching-overload]  # ty: ignore[no-matching-overload]
+    reasoning_text = " ".join(c.text for u in updates for c in u.contents if c.type == "text_reasoning")  # type: ignore[misc]  # pyrefly: ignore[no-matching-overload]  # ty: ignore[no-matching-overload]
+
+    assert "intermediate progress" in text
+    assert "FINAL" in text
+    assert reasoning_text == ""
+
+
+@pytest.mark.asyncio
+async def test_workflow_agent_text_accessor_includes_forwarded_intermediate_text() -> None:
+    """Intermediate text is forwarded as text until issue 5885 defines the final mapping."""
+
+    @executor
+    async def emit(messages: list[Message], ctx: WorkflowContext[str, str]) -> None:
+        await ctx.yield_output("invisible-progress")
+        await ctx.send_message("forward")
+
+    @executor
+    async def terminal(message: str, ctx: WorkflowContext[Never, str]) -> None:  # type: ignore[valid-type]
+        await ctx.yield_output("the-answer")
+
+    workflow = (
+        WorkflowBuilder(
+            start_executor=emit,
+            output_from=[terminal],
+            intermediate_output_from=[emit],
+        )
+        .add_edge(emit, terminal)
+        .build()
+    )
+    agent = workflow.as_agent("test")
+
+    response = await agent.run("hi")
+    assert isinstance(response, AgentResponse)
+    assert "invisible-progress" in response.text
+    assert "the-answer" in response.text
+
+
+@pytest.mark.asyncio
+async def test_workflow_agent_hidden_yields_do_not_surface_non_streaming() -> None:
+    """In explicit designation mode, unlisted executor yields stay out of agent responses."""
+
+    @executor
+    async def hidden(messages: list[Message], ctx: WorkflowContext[str, str]) -> None:
+        await ctx.yield_output("hidden-progress")
+        await ctx.send_message("forward")
+
+    @executor
+    async def terminal(message: str, ctx: WorkflowContext[Never, str]) -> None:  # type: ignore[valid-type]
+        await ctx.yield_output("visible-answer")
+
+    workflow = WorkflowBuilder(start_executor=hidden, output_from=[terminal]).add_edge(hidden, terminal).build()
+    agent = workflow.as_agent("test")
+
+    response = await agent.run("hi")
+    all_text = " ".join(c.text for m in response.messages for c in m.contents if hasattr(c, "text"))  # type: ignore[misc]  # pyrefly: ignore[no-matching-overload]  # ty: ignore[no-matching-overload]
+
+    assert response.text == "visible-answer"
+    assert "hidden-progress" not in all_text
+
+
+@pytest.mark.asyncio
+async def test_workflow_agent_hidden_yields_do_not_surface_streaming() -> None:
+    """In explicit designation mode, unlisted executor yields stay out of agent updates."""
+
+    @executor
+    async def hidden(messages: list[Message], ctx: WorkflowContext[str, str]) -> None:
+        await ctx.yield_output("hidden-progress")
+        await ctx.send_message("forward")
+
+    @executor
+    async def terminal(message: str, ctx: WorkflowContext[Never, str]) -> None:  # type: ignore[valid-type]
+        await ctx.yield_output("visible-answer")
+
+    workflow = WorkflowBuilder(start_executor=hidden, output_from=[terminal]).add_edge(hidden, terminal).build()
+    agent = workflow.as_agent("test")
+
+    updates: list[AgentResponseUpdate] = []
+    async for update in agent.run("hi", stream=True):
+        updates.append(update)
+
+    all_text = " ".join(c.text for u in updates for c in u.contents if hasattr(c, "text"))  # type: ignore[misc]  # pyrefly: ignore[no-matching-overload]  # ty: ignore[no-matching-overload]
+
+    assert "visible-answer" in all_text
+    assert "hidden-progress" not in all_text
+
+
+@pytest.mark.asyncio
+async def test_workflow_agent_data_event_emit_factory_still_forwarded() -> None:
+    """Even the deprecated WorkflowEvent.emit() / type='data' path is forwarded."""
+
+    @executor
+    async def emit_data_alias(messages: list[Message], ctx: WorkflowContext[Never, str]) -> None:  # type: ignore[valid-type]
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", DeprecationWarning)
+            await ctx.add_event(WorkflowEvent.emit("emit_data_alias", "data-alias-payload"))
+        await ctx.yield_output("DONE")
+
+    workflow = WorkflowBuilder(start_executor=emit_data_alias, output_from=[emit_data_alias]).build()
+    agent = workflow.as_agent("test")
+
+    updates: list[AgentResponseUpdate] = []
+    async for update in agent.run("hi", stream=True):
+        updates.append(update)
+
+    text = " ".join(c.text for u in updates for c in u.contents if c.type == "text")  # type: ignore[misc]  # pyrefly: ignore[no-matching-overload]  # ty: ignore[no-matching-overload]
+    assert "data-alias-payload" in text
+
+
+@pytest.mark.asyncio
+async def test_workflow_agent_intermediate_message_preserves_additional_properties() -> None:
+    """Message.additional_properties survives intermediate forwarding.
+
+    Producer-attached metadata (tracking_id, conversation_id, etc.) must not disappear
+    for messages flowing through intermediate-designated executors.
+    """
+
+    @executor
+    async def emit(messages: list[Message], ctx: WorkflowContext[str, AgentResponse]) -> None:
+        msg = Message(
+            role="assistant",
+            contents=[Content.from_text(text="hi")],
+            additional_properties={"tracking_id": "abc-123"},
+        )
+        await ctx.yield_output(AgentResponse(messages=[msg]))
+        await ctx.send_message("forward")
+
+    @executor
+    async def terminal(message: str, ctx: WorkflowContext[Never, str]) -> None:  # type: ignore[valid-type]
+        await ctx.yield_output("done")
+
+    workflow = (
+        WorkflowBuilder(
+            start_executor=emit,
+            output_from=[terminal],
+            intermediate_output_from=[emit],
+        )
+        .add_edge(emit, terminal)
+        .build()
+    )
+    agent = workflow.as_agent("test")
+
+    response = await agent.run("hi")
+    intermediate_msgs = [m for m in response.messages if any(c.type == "text" and c.text == "hi" for c in m.contents)]
+    assert intermediate_msgs, "expected at least one intermediate message in the response"
+    assert intermediate_msgs[0].additional_properties.get("tracking_id") == "abc-123"
+
+
+@pytest.mark.asyncio
+async def test_workflow_agent_terminal_text_stays_text_not_reasoning() -> None:
+    """A designated executor's text yield surfaces as Content.text."""
+
+    @executor
+    async def only(messages: list[Message], ctx: WorkflowContext[Never, str]) -> None:  # type: ignore[valid-type]
+        await ctx.yield_output("the-answer")
+
+    workflow = WorkflowBuilder(start_executor=only, output_from=[only]).build()
+    agent = workflow.as_agent("test")
+
+    response = await agent.run("hi")
+    assert response.text == "the-answer"
+    # No text_reasoning content because everything from `only` is terminal.
+    assert all(c.type != "text_reasoning" for m in response.messages for c in m.contents)
+
+
+@pytest.mark.asyncio
+async def test_workflow_agent_non_streaming_rejects_terminal_update() -> None:
+    """A terminal event carrying AgentResponseUpdate is streaming-only and invalid in run()."""
+
+    @executor
+    async def emit(messages: list[Message], ctx: WorkflowContext[Never, AgentResponseUpdate]) -> None:  # type: ignore[valid-type]
+        await ctx.yield_output(AgentResponseUpdate(contents=[Content.from_text(text="partial")], role="assistant"))
+
+    workflow = WorkflowBuilder(start_executor=emit, output_from=[emit]).build()
+    agent = workflow.as_agent("test")
+
+    with pytest.raises(AgentInvalidRequestException, match="AgentResponseUpdate"):
+        await agent.run("hi")
+
+
+@pytest.mark.asyncio
+async def test_workflow_agent_non_streaming_rejects_intermediate_update() -> None:
+    """An intermediate event carrying AgentResponseUpdate is streaming-only and invalid in run()."""
+
+    @executor
+    async def emit(messages: list[Message], ctx: WorkflowContext[str, AgentResponseUpdate]) -> None:
+        await ctx.yield_output(AgentResponseUpdate(contents=[Content.from_text(text="partial")], role="assistant"))
+        await ctx.send_message("forward")
+
+    @executor
+    async def terminal(message: str, ctx: WorkflowContext[Never, str]) -> None:  # type: ignore[valid-type]
+        await ctx.yield_output("FINAL")
+
+    workflow = (
+        WorkflowBuilder(
+            start_executor=emit,
+            output_from=[terminal],
+            intermediate_output_from=[emit],
+        )
+        .add_edge(emit, terminal)
+        .build()
+    )
+    agent = workflow.as_agent("test")
+
+    with pytest.raises(AgentInvalidRequestException, match="AgentResponseUpdate"):
+        await agent.run("hi")
+
+
+@pytest.mark.asyncio
+async def test_workflow_agent_streaming_update_payloads_preserve_classification() -> None:
+    """Streaming AgentResponseUpdate payloads preserve original content types."""
+
+    @executor
+    async def emit(messages: list[Message], ctx: WorkflowContext[str, AgentResponseUpdate]) -> None:
+        await ctx.yield_output(
+            AgentResponseUpdate(contents=[Content.from_text(text="intermediate-chunk")], role="assistant")
+        )
+        await ctx.send_message("forward")
+
+    @executor
+    async def terminal(message: str, ctx: WorkflowContext[Never, AgentResponseUpdate]) -> None:  # type: ignore[valid-type]
+        await ctx.yield_output(
+            AgentResponseUpdate(contents=[Content.from_text(text="terminal-chunk")], role="assistant")
+        )
+
+    workflow = (
+        WorkflowBuilder(
+            start_executor=emit,
+            output_from=[terminal],
+            intermediate_output_from=[emit],
+        )
+        .add_edge(emit, terminal)
+        .build()
+    )
+    agent = workflow.as_agent("test")
+
+    updates: list[AgentResponseUpdate] = []
+    async for update in agent.run("hi", stream=True):
+        updates.append(update)
+
+    text = " ".join(c.text for u in updates for c in u.contents if c.type == "text")  # type: ignore[misc]  # pyrefly: ignore[no-matching-overload]  # ty: ignore[no-matching-overload]
+    reasoning_text = " ".join(c.text for u in updates for c in u.contents if c.type == "text_reasoning")  # type: ignore[misc]  # pyrefly: ignore[no-matching-overload]  # ty: ignore[no-matching-overload]
+
+    assert "intermediate-chunk" in text
+    assert "terminal-chunk" in text
+    assert reasoning_text == ""
+
+
+@pytest.mark.asyncio
+async def test_workflow_agent_drops_orchestration_internal_events() -> None:
+    """Orchestration-internal event types (group_chat / handoff_sent / magentic_orchestrator)
+    must not surface through workflow.as_agent(). Their dataclass payloads would otherwise
+    be stringified by the generic fallback path and leak into response history."""
+
+    @executor
+    async def emit(messages: list[Message], ctx: WorkflowContext[Never, str]) -> None:  # type: ignore[valid-type]
+        # Construct typed orchestration-internal events directly to assert they get
+        # dropped at the agent boundary regardless of payload.
+        await ctx.add_event(WorkflowEvent("group_chat", data={"orchestrator": "details"}))  # type: ignore[arg-type]
+        await ctx.add_event(WorkflowEvent("handoff_sent", data={"target": "agent_b"}))  # type: ignore[arg-type]
+        await ctx.add_event(WorkflowEvent("magentic_orchestrator", data={"plan": "..."}))  # type: ignore[arg-type]
+        await ctx.yield_output("FINAL")
+
+    workflow = WorkflowBuilder(start_executor=emit, output_from=[emit]).build()
+    agent = workflow.as_agent("test")
+
+    response = await agent.run("hi")
+    all_text = " ".join(c.text for m in response.messages for c in m.contents if hasattr(c, "text"))  # type: ignore[misc]  # pyrefly: ignore[no-matching-overload]  # ty: ignore[no-matching-overload]
+    assert "orchestrator" not in all_text
+    assert "agent_b" not in all_text
+    assert "plan" not in all_text
+    assert response.text == "FINAL"
+
+
+@pytest.mark.asyncio
+async def test_workflow_agent_drops_orchestration_internal_events_streaming() -> None:
+    """Streaming counterpart — orchestration-internal events stay inside the workflow."""
+
+    @executor
+    async def emit(messages: list[Message], ctx: WorkflowContext[Never, str]) -> None:  # type: ignore[valid-type]
+        await ctx.add_event(WorkflowEvent("group_chat", data={"orchestrator": "details"}))  # type: ignore[arg-type]
+        await ctx.yield_output("FINAL")
+
+    workflow = WorkflowBuilder(start_executor=emit, output_from=[emit]).build()
+    agent = workflow.as_agent("test")
+
+    updates: list[AgentResponseUpdate] = []
+    async for update in agent.run("hi", stream=True):
+        updates.append(update)
+
+    all_text = " ".join(c.text for u in updates for c in u.contents if hasattr(c, "text"))  # type: ignore[misc]  # pyrefly: ignore[no-matching-overload]  # ty: ignore[no-matching-overload]
+    assert "orchestrator" not in all_text
+    assert "FINAL" in all_text

python/packages/core/tests/workflow/test_workflow_builder.py

@@ -0,0 +1,374 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from collections.abc import AsyncIterator, Awaitable
+from dataclasses import dataclass
+from typing import Any, Literal, overload
+
+import pytest
+
+from agent_framework import (
+    AgentExecutor,
+    AgentResponse,
+    AgentResponseUpdate,
+    AgentRunInputs,
+    AgentSession,
+    BaseAgent,
+    Case,
+    Default,
+    Executor,
+    Message,
+    ResponseStream,
+    WorkflowBuilder,
+    WorkflowContext,
+    WorkflowValidationError,
+    handler,
+)
+
+
+class DummyAgent(BaseAgent):
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
+
+    @overload
+    def run(
+        self,
+        messages: AgentRunInputs | None = ...,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = ...,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
+    def run(
+        self,
+        messages: AgentRunInputs | None = None,
+        *,
+        stream: bool = False,
+        session: AgentSession | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
+        if stream:
+            return ResponseStream[AgentResponseUpdate, AgentResponse[Any]](self._run_stream_impl())
+        return self._run_impl(messages)
+
+    async def _run_impl(self, messages: AgentRunInputs | None = None) -> AgentResponse:
+        norm: list[Message] = []
+        if messages:
+            for m in messages:  # type: ignore[union-attr]  # ty: ignore[not-iterable]
+                if isinstance(m, Message):
+                    norm.append(m)
+                elif isinstance(m, str):
+                    norm.append(Message(role="user", contents=[m]))
+        return AgentResponse(messages=norm)
+
+    async def _run_stream_impl(self) -> AsyncIterator[AgentResponseUpdate]:
+        # Minimal async generator
+        yield AgentResponseUpdate()
+
+
+def test_builder_accepts_agents_directly():
+    agent1 = DummyAgent(id="agent1", name="writer")
+    agent2 = DummyAgent(id="agent2", name="reviewer")
+
+    wf = WorkflowBuilder(start_executor=agent1).add_edge(agent1, agent2).build()
+
+    # Confirm auto-wrapped executors use agent names as IDs
+    assert wf.start_executor_id == "writer"
+    assert any(isinstance(e, AgentExecutor) and e.id in {"writer", "reviewer"} for e in wf.executors.values())
+
+
+@dataclass
+class MockMessage:
+    """A mock message for testing purposes."""
+
+    data: Any
+
+
+class MockExecutor(Executor):
+    """A mock executor for testing purposes."""
+
+    @handler
+    async def mock_handler(self, message: MockMessage, ctx: WorkflowContext[MockMessage, MockMessage]) -> None:
+        """A mock handler that does nothing."""
+        pass
+
+
+class MockAggregator(Executor):
+    """A mock executor that aggregates results from multiple executors."""
+
+    @handler
+    async def mock_handler(self, messages: list[MockMessage], ctx: WorkflowContext[MockMessage]) -> None:
+        # This mock simply returns the data incremented by 1
+        pass
+
+
+def test_workflow_builder_without_start_executor_throws():
+    """Test creating a workflow builder without a start executor."""
+    with pytest.raises(TypeError):
+        WorkflowBuilder()  # type: ignore[call-arg]  # ty: ignore[missing-argument]
+
+
+def test_workflow_builder_fluent_api():
+    """Test the fluent API of the workflow builder."""
+    executor_a = MockExecutor(id="executor_a")
+    executor_b = MockExecutor(id="executor_b")
+    executor_c = MockExecutor(id="executor_c")
+    executor_d = MockExecutor(id="executor_d")
+    executor_e = MockAggregator(id="executor_e")
+    executor_f = MockExecutor(id="executor_f")
+
+    workflow = (
+        WorkflowBuilder(max_iterations=5, start_executor=executor_a)
+        .add_edge(executor_a, executor_b)
+        .add_fan_out_edges(executor_b, [executor_c, executor_d])
+        .add_fan_in_edges([executor_c, executor_d], executor_e)
+        .add_chain([executor_e, executor_f])
+        .build()
+    )
+
+    assert len(workflow.edge_groups) == 4 + 6  # 4 defined edges + 6 internal edges for request-response handling
+    assert workflow.start_executor_id == executor_a.id
+    assert len(workflow.executors) == 6
+
+
+def test_add_agent_reuses_same_wrapper():
+    """Test that using the same agent instance multiple times reuses the same wrapper."""
+    reuse_agent = DummyAgent(id="agent_reuse", name="reuse_agent")
+    agent_a = DummyAgent(id="agent_a", name="agent_a")
+
+    builder = WorkflowBuilder(start_executor=reuse_agent)
+    # Use the same agent instance in add_edge - should reuse the same wrapper
+    builder.add_edge(reuse_agent, agent_a)
+    builder.add_edge(agent_a, reuse_agent)
+
+    workflow = builder.build()
+
+    # Verify only one executor exists for this agent
+    assert workflow.start_executor_id == "reuse_agent"
+    assert "reuse_agent" in workflow.executors
+    assert len([e for e in workflow.executors.values() if isinstance(e, AgentExecutor)]) == 2
+
+
+def test_add_agent_duplicate_id_raises_error():
+    """Test that adding agents with duplicate IDs raises an error."""
+    agent1 = DummyAgent(id="agent1", name="first")
+    agent2 = DummyAgent(id="agent2", name="first")  # Same name as agent1
+    builder = WorkflowBuilder(start_executor=agent1)
+
+    with pytest.raises(ValueError, match="Duplicate executor ID"):
+        builder.add_edge(agent1, agent2).build()
+
+
+def test_fan_out_edges_with_direct_instances():
+    """Test fan-out edges with direct executor instances."""
+    source = MockExecutor(id="Source")
+    target1 = MockExecutor(id="Target1")
+    target2 = MockExecutor(id="Target2")
+
+    workflow = WorkflowBuilder(start_executor=source).add_fan_out_edges(source, [target1, target2]).build()
+
+    assert "Source" in workflow.executors
+    assert "Target1" in workflow.executors
+    assert "Target2" in workflow.executors
+
+
+def test_fan_in_edges_with_direct_instances():
+    """Test fan-in edges with direct executor instances."""
+    source1 = MockExecutor(id="Source1")
+    source2 = MockExecutor(id="Source2")
+    aggregator = MockAggregator(id="Aggregator")
+
+    workflow = (
+        WorkflowBuilder(start_executor=source1)
+        .add_edge(source1, source2)
+        .add_fan_in_edges([source1, source2], aggregator)
+        .build()
+    )
+
+    assert "Source1" in workflow.executors
+    assert "Source2" in workflow.executors
+    assert "Aggregator" in workflow.executors
+
+
+def test_chain_with_direct_instances():
+    """Test add_chain with direct executor instances."""
+    step1 = MockExecutor(id="Step1")
+    step2 = MockExecutor(id="Step2")
+    step3 = MockExecutor(id="Step3")
+
+    workflow = WorkflowBuilder(start_executor=step1).add_chain([step1, step2, step3]).build()
+
+    assert "Step1" in workflow.executors
+    assert "Step2" in workflow.executors
+    assert "Step3" in workflow.executors
+    assert workflow.start_executor_id == "Step1"
+
+
+def test_add_edge_with_condition():
+    """Test adding edges with conditions using direct executor instances."""
+    source = MockExecutor(id="Source")
+    target = MockExecutor(id="Target")
+
+    def condition_func(msg: MockMessage) -> bool:
+        return msg.data > 0
+
+    workflow = WorkflowBuilder(start_executor=source).add_edge(source, target, condition=condition_func).build()
+
+    assert "Source" in workflow.executors
+    assert "Target" in workflow.executors
+
+
+def test_switch_case_with_agents():
+    """Test add_switch_case_edge_group with Case and Default edges using agents."""
+    router = DummyAgent(id="router_agent", name="router")
+    handler = DummyAgent(id="handler", name="handler")
+    fallback = DummyAgent(id="fallback_agent", name="fallback")
+
+    workflow = (
+        WorkflowBuilder(start_executor=router)
+        .add_switch_case_edge_group(
+            router,
+            [
+                Case(condition=lambda _: True, target=handler),
+                Default(target=fallback),
+            ],
+        )
+        .build()
+    )
+
+    # All three agents should be AgentExecutor wrappers
+    agent_executors = [e for e in workflow.executors.values() if isinstance(e, AgentExecutor)]
+    assert len(agent_executors) == 3
+
+
+# region with_output_from tests
+
+
+def test_with_output_from_returns_builder():
+    """Test that with_output_from returns the builder for method chaining."""
+    executor_a = MockExecutor(id="executor_a")
+    builder = WorkflowBuilder(output_from=[executor_a], start_executor=executor_a)
+
+    # Verify builder was created with output_from
+    assert builder._output_from == [executor_a]  # pyright: ignore[reportPrivateUsage]
+
+
+def test_with_output_from_with_executor_instances():
+    """Test with_output_from with direct executor instances."""
+    executor_a = MockExecutor(id="executor_a")
+    executor_b = MockExecutor(id="executor_b")
+
+    workflow = (
+        WorkflowBuilder(start_executor=executor_a, output_from=[executor_b]).add_edge(executor_a, executor_b).build()
+    )
+
+    # Verify that the workflow was built with the correct output executors
+    assert {ex.id for ex in workflow.get_output_executors()} == {"executor_b"}
+
+
+def test_with_output_from_with_agent_instances():
+    """Test with_output_from with agent instances."""
+    agent_a = DummyAgent(id="agent_a", name="writer")
+    agent_b = DummyAgent(id="agent_b", name="reviewer")
+
+    workflow = WorkflowBuilder(start_executor=agent_a, output_from=[agent_b]).add_edge(agent_a, agent_b).build()
+
+    # Verify that the workflow was built with the agent's name as output executor
+    assert {ex.id for ex in workflow.get_output_executors()} == {"reviewer"}
+
+
+def test_with_output_from_with_executor_instances_by_id():
+    """Test with_output_from with direct executor instances resolves to executor IDs."""
+    executor_a = MockExecutor(id="ExecutorA")
+    executor_b = MockExecutor(id="ExecutorB")
+
+    workflow = (
+        WorkflowBuilder(start_executor=executor_a, output_from=[executor_b]).add_edge(executor_a, executor_b).build()
+    )
+
+    assert {ex.id for ex in workflow.get_output_executors()} == {"ExecutorB"}
+
+
+def test_with_output_from_with_multiple_executors():
+    """Test with_output_from with multiple executors."""
+    executor_a = MockExecutor(id="executor_a")
+    executor_b = MockExecutor(id="executor_b")
+    executor_c = MockExecutor(id="executor_c")
+
+    workflow = (
+        WorkflowBuilder(start_executor=executor_a, output_from=[executor_a, executor_c])
+        .add_edge(executor_a, executor_b)
+        .add_edge(executor_b, executor_c)
+        .build()
+    )
+
+    # Verify that the workflow was built with both output executors
+    assert {ex.id for ex in workflow.get_output_executors()} == {"executor_a", "executor_c"}
+
+
+def test_with_output_from_can_be_set_to_different_value():
+    """Test that output_from can be set at construction time."""
+    executor_a = MockExecutor(id="executor_a")
+    executor_b = MockExecutor(id="executor_b")
+
+    workflow = (
+        WorkflowBuilder(start_executor=executor_a, output_from=[executor_b]).add_edge(executor_a, executor_b).build()
+    )
+
+    # Verify that the setting is applied
+    assert {ex.id for ex in workflow.get_output_executors()} == {"executor_b"}
+
+
+def test_with_output_from_with_agent_instances_resolves_name():
+    """Test with_output_from with agent instances resolves to agent names."""
+    agent_writer = DummyAgent(id="agent1", name="writer")
+    agent_reviewer = DummyAgent(id="agent2", name="reviewer")
+
+    workflow = (
+        WorkflowBuilder(start_executor=agent_writer, output_from=[agent_reviewer])
+        .add_edge(agent_writer, agent_reviewer)
+        .build()
+    )
+
+    assert {ex.id for ex in workflow.get_output_executors()} == {"reviewer"}
+
+
+def test_with_output_from_in_constructor():
+    """Test that output_from works correctly when set in the constructor."""
+    executor_a = MockExecutor(id="executor_a")
+    executor_b = MockExecutor(id="executor_b")
+    executor_c = MockExecutor(id="executor_c")
+
+    # Build workflow with output_from in the constructor
+    workflow = (
+        WorkflowBuilder(start_executor=executor_a, output_from=[executor_c])
+        .add_edge(executor_a, executor_b)
+        .add_edge(executor_b, executor_c)
+        .build()
+    )
+
+    # Verify that the setting persists through the chain
+    assert {ex.id for ex in workflow.get_output_executors()} == {"executor_c"}
+
+
+def test_with_output_from_with_invalid_executor_raises_validation_error():
+    """Test that with_output_from with an invalid executor raises an error."""
+    executor_a = MockExecutor(id="executor_a")
+
+    builder = WorkflowBuilder(start_executor=executor_a, output_from=[MockExecutor(id="executor_b")])
+
+    # Attempting to set output from an executor not in the workflow should raise an error
+    with pytest.raises(
+        WorkflowValidationError, match="Output executor 'executor_b' is not present in the workflow graph"
+    ):
+        builder.build()
+
+
+# endregion

python/packages/core/tests/workflow/test_workflow_context.py

@@ -0,0 +1,303 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, Any
+
+import pytest
+from typing_extensions import Never
+
+from agent_framework import (
+    Executor,
+    WorkflowBuilder,
+    WorkflowContext,
+    WorkflowEvent,
+    WorkflowRunState,
+    executor,
+    handler,
+)
+
+if TYPE_CHECKING:
+    from _pytest.logging import LogCaptureFixture
+
+    from agent_framework._workflows._runner_context import InProcRunnerContext
+
+
+class MockExecutor(Executor):
+    """Mock executor for testing."""
+
+    def __init__(self, id: str) -> None:
+        super().__init__(id=id)
+
+    @handler
+    async def handle_message(self, message: str, ctx: WorkflowContext[str]) -> None:
+        """Handle string messages."""
+        ...
+
+
+@asynccontextmanager
+async def make_context(
+    executor_id: str = "exec",
+) -> AsyncIterator[tuple[WorkflowContext[object], "InProcRunnerContext"]]:
+    from agent_framework._workflows._runner_context import InProcRunnerContext
+    from agent_framework._workflows._state import State
+
+    mock_executor = MockExecutor(executor_id)
+    runner_ctx = InProcRunnerContext()
+    state = State()
+    workflow_ctx: WorkflowContext[object] = WorkflowContext(
+        mock_executor,
+        ["source"],
+        state,
+        runner_ctx,
+    )
+    try:
+        yield workflow_ctx, runner_ctx
+    finally:
+        await asyncio.sleep(0)
+
+
+async def test_executor_cannot_emit_framework_lifecycle_event(caplog: "LogCaptureFixture") -> None:
+    async with make_context() as (ctx, runner_ctx):
+        caplog.clear()
+        with caplog.at_level("WARNING"):
+            await ctx.add_event(WorkflowEvent.status(state=WorkflowRunState.IN_PROGRESS))
+
+        events: list[WorkflowEvent] = await runner_ctx.drain_events()
+        assert len(events) == 1
+        assert events[0].type == "warning"
+        data = events[0].data
+        assert isinstance(data, str)
+        assert "reserved for framework lifecycle notifications" in data
+        assert any("attempted to emit" in message and "'status'" in message for message in list(caplog.messages))
+
+
+@pytest.mark.parametrize(
+    "event",
+    [
+        WorkflowEvent("output", executor_id="exec", data="output-payload"),
+        WorkflowEvent("intermediate", executor_id="exec", data="intermediate-payload"),
+    ],
+)
+async def test_executor_cannot_emit_output_selection_events(
+    event: WorkflowEvent[Any],
+    caplog: "LogCaptureFixture",
+) -> None:
+    async with make_context() as (ctx, runner_ctx):
+        caplog.clear()
+        with caplog.at_level("WARNING"):
+            await ctx.add_event(event)
+
+        events: list[WorkflowEvent] = await runner_ctx.drain_events()
+        assert len(events) == 1
+        assert events[0].type == "warning"
+        data = events[0].data
+        assert isinstance(data, str)
+        assert "reserved for ctx.yield_output()" in data
+        assert event.data not in [emitted.data for emitted in events]
+
+
+async def test_executor_emits_normal_event() -> None:
+    async with make_context() as (ctx, runner_ctx):
+        # Create a normal event to test event emission
+        await ctx.add_event(_TestEvent())
+
+        events: list[WorkflowEvent] = await runner_ctx.drain_events()
+        assert len(events) == 1
+        assert isinstance(events[0], _TestEvent)
+
+
+class _TestEvent(WorkflowEvent):
+    def __init__(self, data: Any = None) -> None:
+        super().__init__("test_event", data=data)  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+
+
+async def test_workflow_context_type_annotations_no_parameter() -> None:
+    # Test function-based executor
+    @executor(id="func1")
+    async def func1(text: str, ctx: WorkflowContext) -> None:
+        await ctx.add_event(_TestEvent())
+
+    wf = WorkflowBuilder(start_executor=func1).build()
+    events = await wf.run("hello")
+    test_events = [e for e in events if isinstance(e, _TestEvent)]
+    assert len(test_events) == 1
+
+    # Test class-based executor
+    class _exec1(Executor):
+        @handler
+        async def func1(self, text: str, ctx: WorkflowContext) -> None:
+            await ctx.add_event(_TestEvent())
+
+    executor1 = _exec1(id="exec1")
+
+    assert executor1.input_types == [str]
+    assert executor1.output_types == []
+    assert executor1.workflow_output_types == []
+
+    wf2 = WorkflowBuilder(start_executor=executor1).build()
+    events2 = await wf2.run("hello")
+    test_events2 = [e for e in events2 if isinstance(e, _TestEvent)]
+    assert len(test_events2) == 1
+
+
+async def test_workflow_context_type_annotations_message_type_parameter() -> None:
+    # Test function-based executor
+    @executor(id="func1")
+    async def func1(text: str, ctx: WorkflowContext[str]) -> None:
+        await ctx.send_message("world")
+
+    @executor(id="func2")
+    async def func2(text: str, ctx: WorkflowContext) -> None:
+        await ctx.add_event(_TestEvent(data=text))
+
+    wf = WorkflowBuilder(start_executor=func1).add_edge(func1, func2).build()
+    events = await wf.run("hello")
+    test_events = [e for e in events if isinstance(e, _TestEvent)]
+    assert len(test_events) == 1
+    assert test_events[0].data == "world"
+
+    # Test class-based executor
+    class _exec1(Executor):
+        @handler
+        async def func1(self, text: str, ctx: WorkflowContext[str]) -> None:
+            await ctx.send_message("world")
+
+    class _exec2(Executor):
+        @handler
+        async def func2(self, text: str, ctx: WorkflowContext) -> None:
+            await ctx.add_event(_TestEvent(data=text))
+
+    executor1 = _exec1(id="exec1")
+    executor2 = _exec2(id="exec2")
+
+    assert executor1.input_types == [str]
+    assert executor1.output_types == [str]
+    assert executor1.workflow_output_types == []
+    assert executor2.input_types == [str]
+    assert executor2.output_types == []
+    assert executor2.workflow_output_types == []
+
+    wf2 = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
+    events2 = await wf2.run("hello")
+    test_events2 = [e for e in events2 if isinstance(e, _TestEvent)]
+    assert len(test_events2) == 1
+    assert test_events2[0].data == "world"
+
+
+async def test_workflow_context_type_annotations_message_and_output_type_parameters() -> None:
+    # Test function-based executor
+    @executor(id="func1")
+    async def func1(text: str, ctx: WorkflowContext[str]) -> None:
+        await ctx.send_message("world")
+
+    @executor(id="func2")
+    async def func2(text: str, ctx: WorkflowContext[Never, str]) -> None:  # type: ignore[valid-type]
+        await ctx.add_event(_TestEvent(data=text))
+        await ctx.yield_output(text)
+
+    wf = WorkflowBuilder(start_executor=func1).add_edge(func1, func2).build()
+    events = await wf.run("hello")
+    outputs = events.get_outputs()
+    assert len(outputs) == 1
+    assert outputs[0] == "world"
+
+    # Test class-based executor
+    class _exec1(Executor):
+        @handler
+        async def func1(self, text: str, ctx: WorkflowContext[str]) -> None:
+            await ctx.send_message("world")
+
+    class _exec2(Executor):
+        @handler
+        async def func2(self, text: str, ctx: WorkflowContext[Never, str]) -> None:  # type: ignore[valid-type]
+            await ctx.add_event(_TestEvent(data=text))
+            await ctx.yield_output(text)
+
+    executor1 = _exec1(id="exec1")
+    executor2 = _exec2(id="exec2")
+
+    assert executor1.input_types == [str]
+    assert executor1.output_types == [str]
+    assert executor1.workflow_output_types == []
+    assert executor2.input_types == [str]
+    assert executor2.output_types == []
+    assert executor2.workflow_output_types == [str]
+
+    wf2 = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
+    events2 = await wf2.run("hello")
+    outputs2 = events2.get_outputs()
+    assert len(outputs2) == 1
+    assert outputs2[0] == "world"
+
+
+async def test_workflow_context_type_annotations_any() -> None:
+    class _exec1(Executor):
+        @handler
+        async def func1(self, text: str, ctx: WorkflowContext[Any]) -> None:
+            await ctx.add_event(_TestEvent())
+            await ctx.send_message(123)
+
+    executor1 = _exec1(id="exec1")
+    assert executor1.input_types == [str]
+    assert executor1.output_types == [Any]
+
+    class _exec2(Executor):
+        @handler
+        async def func2(self, number: int, ctx: WorkflowContext[Any, Any]) -> None:
+            await ctx.add_event(_TestEvent())
+            await ctx.send_message(456)
+            await ctx.yield_output(3.14)
+
+    executor2 = _exec2(id="exec2")
+    assert executor2.input_types == [int]
+    assert executor2.output_types == [Any]
+    assert executor2.workflow_output_types == [Any]
+
+
+async def test_workflow_context_missing_annotation_error() -> None:
+    """Test that missing WorkflowContext annotation raises appropriate error."""
+    import pytest
+
+    # Test function-based executor with missing ctx annotation
+    with pytest.raises(ValueError, match="must have a WorkflowContext"):
+
+        @executor(id="bad_func")
+        async def bad_func(text: str, ctx) -> None:  # type: ignore[no-untyped-def]
+            pass
+
+    # Test class-based executor with missing ctx annotation
+    with pytest.raises(ValueError, match="must have a WorkflowContext"):
+
+        class _BadExecutor(Executor):  # pyright: ignore[reportUnusedClass]
+            @handler  # pyright: ignore[reportUnknownArgumentType]
+            async def bad_handler(self, text: str, ctx) -> None:  # type: ignore[no-untyped-def]
+                pass
+
+
+async def test_workflow_context_invalid_type_parameter_error() -> None:
+    """Test that invalid type parameters like int values raise appropriate errors."""
+    import pytest
+
+    # Test function-based executor with invalid type parameter (int value instead of type)
+    with pytest.raises(ValueError, match="invalid type entry"):
+
+        @executor(id="bad_func")
+        async def bad_func(text: str, ctx: WorkflowContext[123]) -> None:  # type: ignore[valid-type]  # ty: ignore[invalid-type-form]
+            pass
+
+    # Test class-based executor with invalid type parameter
+    with pytest.raises(ValueError, match="invalid type entry"):
+
+        class _BadExecutor(Executor):  # pyright: ignore[reportUnusedClass]
+            @handler  # pyright: ignore[reportUnknownArgumentType]
+            async def bad_handler(self, text: str, ctx: WorkflowContext[456]) -> None:  # type: ignore[valid-type]  # ty: ignore[invalid-type-form]
+                pass
+
+    # Test two-parameter WorkflowContext with invalid workflow output type
+    with pytest.raises(ValueError, match="invalid type entry"):
+
+        @executor(id="bad_func2")
+        async def bad_func2(text: str, ctx: WorkflowContext[str, 789]) -> None:  # type: ignore[valid-type]  # ty: ignore[invalid-type-form]
+            pass

python/packages/core/tests/workflow/test_workflow_event_factories.py

@@ -0,0 +1,34 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for WorkflowEvent factory methods and WorkflowEvent.emit() deprecation."""
+
+from __future__ import annotations
+
+import warnings
+
+import pytest
+
+from agent_framework import AgentResponse, Message
+from agent_framework._workflows._events import WorkflowEvent
+
+
+def test_workflow_event_output_selection_factories_are_not_public() -> None:
+    """Callers should use ctx.yield_output(), not direct output/intermediate factories."""
+    assert not hasattr(WorkflowEvent, "output")
+    assert not hasattr(WorkflowEvent, "intermediate")
+
+
+def test_workflow_event_emit_emits_deprecation_warning() -> None:
+    """Calling WorkflowEvent.emit() raises a DeprecationWarning recommending the new path."""
+    response = AgentResponse(messages=[Message(role="assistant", contents=["x"])])
+    with pytest.warns(DeprecationWarning, match="yield_output"):
+        WorkflowEvent.emit(executor_id="t", data=response)
+
+
+def test_workflow_event_emit_still_returns_data_event() -> None:
+    """During the deprecation window, emit() still produces a type='data' event."""
+    response = AgentResponse(messages=[Message(role="assistant", contents=["x"])])
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", DeprecationWarning)
+        event = WorkflowEvent.emit(executor_id="t", data=response)
+    assert event.type == "data"

python/packages/core/tests/workflow/test_workflow_observability.py

@@ -0,0 +1,505 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from typing import Any, cast
+
+import pytest
+from opentelemetry import trace
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+
+from agent_framework import InMemoryCheckpointStorage, WorkflowBuilder
+from agent_framework._workflows._executor import Executor, handler
+from agent_framework._workflows._runner_context import InProcRunnerContext, MessageType, WorkflowMessage
+from agent_framework._workflows._state import State
+from agent_framework._workflows._workflow import Workflow
+from agent_framework._workflows._workflow_context import WorkflowContext
+from agent_framework.observability import (
+    OtelAttr,
+    create_processing_span,
+    create_workflow_span,
+)
+
+
+class MockExecutor(Executor):
+    """Mock executor for testing."""
+
+    def __init__(self, id: str = "mock_executor") -> None:
+        super().__init__(id=id)
+        # Use private field to avoid Pydantic validation
+        self._processed_messages: list[str] = []
+
+    @handler
+    async def handle_message(self, message: str, ctx: WorkflowContext[str]) -> None:
+        """Handle string messages."""
+        self._processed_messages.append(message)
+        await ctx.send_message(f"processed: {message}")
+
+    @property
+    def processed_messages(self) -> list[str]:
+        """Access to processed messages for testing."""
+        return self._processed_messages
+
+
+class SecondExecutor(Executor):
+    """Second executor for testing message chains."""
+
+    def __init__(self, id: str = "second_executor") -> None:
+        super().__init__(id=id)
+        # Use private field to avoid Pydantic validation
+        self._processed_messages: list[str] = []
+
+    @handler
+    async def handle_message(self, message: str, ctx: WorkflowContext) -> None:
+        """Handle string messages."""
+        self._processed_messages.append(message)
+
+    @property
+    def processed_messages(self) -> list[str]:
+        """Access to processed messages for testing."""
+        return self._processed_messages
+
+
+class ProcessingExecutor(Executor):
+    """Executor that processes and forwards messages with a custom prefix."""
+
+    def __init__(self, id: str, prefix: str = "processed") -> None:
+        super().__init__(id=id)
+        # Use private field to avoid Pydantic validation
+        self._processed_messages: list[str] = []
+        self._prefix = prefix
+
+    @handler
+    async def handle_message(self, message: str, ctx: WorkflowContext[str]) -> None:
+        """Handle string messages and send them forward with prefix."""
+        self._processed_messages.append(message)
+        await ctx.send_message(f"{self._prefix}: {message}")
+
+    @property
+    def processed_messages(self) -> list[str]:
+        return self._processed_messages
+
+
+class FanInAggregator(Executor):
+    """Fan-in aggregator that expects a list of inputs."""
+
+    def __init__(self, id: str = "aggregator") -> None:
+        super().__init__(id=id)
+        # Use private field to avoid Pydantic validation
+        self._processed_messages: list[Any] = []
+
+    @handler
+    async def handle_aggregated_data(self, messages: list[str], ctx: WorkflowContext) -> None:
+        # Process aggregated messages from fan-in
+        aggregated = f"aggregated: {', '.join(messages)}"
+        self._processed_messages.append(aggregated)
+
+    @property
+    def processed_messages(self) -> list[Any]:
+        """Access to processed messages for testing."""
+        return self._processed_messages
+
+
+async def test_span_creation_and_attributes(span_exporter: InMemorySpanExporter) -> None:
+    """Test creation and attributes of all span types (workflow, processing, sending)."""
+    # Create a mock workflow object
+    mock_workflow = cast(
+        Workflow,
+        type(
+            "MockWorkflow",
+            (),
+            {
+                "id": "test-workflow-123",
+                "max_iterations": 100,
+                "model_dump_json": lambda self: '{"id": "test-workflow-123", "type": "mock"}',  # pyright: ignore[reportUnknownLambdaType]
+            },
+        )(),
+    )
+
+    # Test all span types in nested context
+    with create_workflow_span(
+        OtelAttr.WORKFLOW_RUN_SPAN,
+        {
+            OtelAttr.WORKFLOW_ID: mock_workflow.id,
+        },
+    ) as workflow_span:
+        workflow_span.add_event(OtelAttr.WORKFLOW_STARTED)
+        sending_attributes: dict[str, str | int] = {
+            OtelAttr.MESSAGE_TYPE: "ResponseMessage",
+            OtelAttr.MESSAGE_DESTINATION_EXECUTOR_ID: "target-789",
+        }
+        with (
+            create_processing_span(
+                "executor-456", "TestExecutor", str(MessageType.STANDARD), "TestMessage"
+            ) as processing_span,
+            create_workflow_span(
+                OtelAttr.MESSAGE_SEND_SPAN, sending_attributes, kind=trace.SpanKind.PRODUCER
+            ) as sending_span,
+        ):
+            # Verify all spans are recording
+            assert workflow_span is not None and workflow_span.is_recording()
+            assert processing_span is not None and processing_span.is_recording()
+            assert sending_span is not None and sending_span.is_recording()
+
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 3
+
+    # Check workflow span
+    workflow_span = next(s for s in spans if s.name == "workflow.run")  # type: ignore[assignment, misc]
+    assert workflow_span.kind == trace.SpanKind.INTERNAL  # type: ignore[attr-defined]
+    assert workflow_span.attributes is not None  # type: ignore[attr-defined]
+    assert workflow_span.attributes.get(OtelAttr.WORKFLOW_ID) == "test-workflow-123"  # type: ignore[attr-defined]
+    assert workflow_span.events is not None  # type: ignore[attr-defined]
+    event_names = [event.name for event in workflow_span.events]  # type: ignore[attr-defined]
+    assert "workflow.started" in event_names
+
+    # Check processing span - span name uses format "executor.process {executor_id}"
+    processing_span = next(s for s in spans if s.name == "executor.process executor-456")  # type: ignore[assignment, misc]
+    assert processing_span.kind == trace.SpanKind.INTERNAL  # type: ignore[attr-defined]
+    assert processing_span.attributes is not None  # type: ignore[attr-defined]
+    assert processing_span.attributes.get("executor.id") == "executor-456"  # type: ignore[attr-defined]
+    assert processing_span.attributes.get("executor.type") == "TestExecutor"  # type: ignore[attr-defined]
+    assert processing_span.attributes.get("message.type") == str(MessageType.STANDARD)  # type: ignore[attr-defined]
+    assert processing_span.attributes.get("message.payload_type") == "TestMessage"  # type: ignore[attr-defined]
+
+    # Check sending span
+    sending_span = next(s for s in spans if s.name == "message.send")  # type: ignore[assignment, misc]
+    assert sending_span.kind == trace.SpanKind.PRODUCER  # type: ignore[attr-defined]
+    assert sending_span.attributes is not None  # type: ignore[attr-defined]
+    assert sending_span.attributes.get("message.type") == "ResponseMessage"  # type: ignore[attr-defined]
+    assert sending_span.attributes.get("message.destination_executor_id") == "target-789"  # type: ignore[attr-defined]
+
+
+async def test_trace_context_handling(span_exporter: InMemorySpanExporter) -> None:
+    """Test trace context propagation and handling in messages and executors."""
+    state = State()
+    ctx = InProcRunnerContext()
+    executor = MockExecutor("test-executor")
+
+    span_exporter.clear()
+
+    # Test trace context propagation in messages
+    workflow_ctx: WorkflowContext[str] = WorkflowContext(
+        executor,
+        ["source"],
+        state,
+        ctx,
+        trace_contexts=[{"traceparent": "00-12345678901234567890123456789012-1234567890123456-01"}],
+        source_span_ids=["1234567890123456"],
+    )
+
+    # Send a message (this should create a sending span and propagate trace context)
+    await workflow_ctx.send_message("test message")
+
+    # Check that message was created with trace context
+    messages = await ctx.drain_messages()
+    assert len(messages) == 1
+    message_list = list(messages.values())[0]
+    assert len(message_list) == 1
+    message = message_list[0]
+    assert message.trace_context is not None
+    assert message.source_span_id is not None
+
+    # Test executor trace context handling
+    await executor.execute(
+        "test message",
+        ["source"],  # source_executor_ids
+        state,  # state
+        ctx,  # runner_context
+        trace_contexts=[{"traceparent": "00-12345678901234567890123456789012-1234567890123456-01"}],
+        source_span_ids=["1234567890123456"],
+    )
+
+    # Check that spans were created with proper attributes
+    spans = span_exporter.get_finished_spans()
+    # Processing spans now use executor_id as the span name
+    processing_spans = [s for s in spans if s.attributes and s.attributes.get("executor.id") == "test-executor"]
+    sending_spans = [s for s in spans if s.name == "message.send"]
+
+    assert len(processing_spans) >= 1
+    assert len(sending_spans) >= 1
+
+    # Verify processing span attributes
+    processing_span = processing_spans[0]
+    assert (
+        processing_span.name == "executor.process test-executor"
+    )  # Span name uses format "executor.process {executor_id}"
+    assert processing_span.attributes is not None
+    assert processing_span.attributes.get("executor.id") == "test-executor"
+    assert processing_span.attributes.get("executor.type") == "MockExecutor"
+    assert processing_span.attributes.get("message.type") == str(MessageType.STANDARD)
+    assert processing_span.attributes.get("message.payload_type") == "str"
+
+
+@pytest.mark.parametrize("enable_instrumentation", [False], indirect=True)
+async def test_trace_context_disabled_when_tracing_disabled(
+    enable_instrumentation: bool, span_exporter: InMemorySpanExporter
+) -> None:
+    """Test that no trace context is added when tracing is disabled."""
+    # Tracing should be disabled by default
+    executor = MockExecutor("test-executor")
+    state = State()
+    ctx = InProcRunnerContext()
+
+    workflow_ctx: WorkflowContext[str] = WorkflowContext(
+        executor,
+        ["source"],
+        state,
+        ctx,
+    )
+
+    # Send a message
+    await workflow_ctx.send_message("test message")
+
+    # Check that message was created without trace context
+    messages = await ctx.drain_messages()
+    message = list(messages.values())[0][0]
+
+    # When tracing is disabled, trace_context should be None
+    assert message.trace_context is None
+    assert message.source_span_id is None
+
+
+async def test_end_to_end_workflow_tracing(span_exporter: InMemorySpanExporter) -> None:
+    """Test end-to-end tracing including workflow build, execution, and span linking with fan-in edges."""
+    # Create executors for fan-in scenario
+    executor1 = MockExecutor("executor1")
+    executor2 = ProcessingExecutor("executor2", "second")
+    executor3 = ProcessingExecutor("executor3", "third")
+    aggregator = FanInAggregator("aggregator")
+
+    # Create workflow with fan-in: executor1 -> [executor2, executor3] -> aggregator
+    workflow = (
+        WorkflowBuilder(start_executor=executor1)
+        .add_fan_out_edges(executor1, [executor2, executor3])
+        .add_fan_in_edges([executor2, executor3], aggregator)
+        .build()
+    )
+
+    # Verify build span was created
+    build_spans = [s for s in span_exporter.get_finished_spans() if s.name == "workflow.build"]
+    assert len(build_spans) == 1
+
+    build_span = build_spans[0]
+    assert build_span.attributes is not None
+    assert build_span.attributes.get(OtelAttr.WORKFLOW_ID) == workflow.id
+    assert build_span.attributes.get("workflow.definition") is not None
+    definition = build_span.attributes.get("workflow.definition")
+    assert definition == workflow.to_json()
+
+    # Check build events
+    assert build_span.events is not None
+    build_event_names = [event.name for event in build_span.events]
+    assert "build.started" in build_event_names
+    assert "build.validation_completed" in build_event_names
+    assert "build.completed" in build_event_names
+
+    # Clear spans to test workflow with name and description
+    span_exporter.clear()
+
+    # Test workflow with name and description - verify OTEL attributes
+    WorkflowBuilder(
+        name="Test Pipeline",
+        description="Test workflow description",
+        start_executor=MockExecutor("start"),
+    ).build()
+
+    build_spans_with_metadata = [s for s in span_exporter.get_finished_spans() if s.name == "workflow.build"]
+    assert len(build_spans_with_metadata) == 1
+    metadata_build_span = build_spans_with_metadata[0]
+    assert metadata_build_span.attributes is not None
+    assert metadata_build_span.attributes.get(OtelAttr.WORKFLOW_BUILDER_NAME) == "Test Pipeline"
+    assert metadata_build_span.attributes.get(OtelAttr.WORKFLOW_BUILDER_DESCRIPTION) == "Test workflow description"
+
+    # Clear spans to separate build from run tracing
+    span_exporter.clear()
+
+    # Run workflow (this should create run spans)
+    events: list[Any] = []
+    async for event in workflow.run("test input", stream=True):
+        events.append(event)
+
+    # Verify workflow executed correctly
+    assert len(executor1.processed_messages) == 1
+    assert executor1.processed_messages[0] == "test input"
+    assert len(executor2.processed_messages) == 1
+    assert executor2.processed_messages[0] == "processed: test input"
+    assert len(executor3.processed_messages) == 1
+    assert executor3.processed_messages[0] == "processed: test input"  # executor3 receives from executor1 via fan-out
+    assert len(aggregator.processed_messages) == 1
+    # The aggregator should receive both processed messages from executor2 and executor3
+    aggregated_msg = aggregator.processed_messages[0]
+    assert "second: processed: test input" in aggregated_msg
+    assert "third: processed: test input" in aggregated_msg
+
+    # Check run spans (build spans should not be present after clear)
+    spans = span_exporter.get_finished_spans()
+
+    # Should have workflow span, processing spans, and sending spans
+    # Processing spans now use executor_id as the span name, filter by executor.id attribute
+    workflow_spans = [s for s in spans if s.name == "workflow.run"]
+    processing_spans = [s for s in spans if s.attributes and s.attributes.get("executor.id") is not None]
+    sending_spans = [s for s in spans if s.name == "message.send"]
+    build_spans_after_run = [s for s in spans if s.name == "workflow.build"]
+
+    assert len(workflow_spans) == 1
+    assert len(processing_spans) >= 4  # executor1, executor2, executor3, aggregator
+    assert len(sending_spans) >= 3  # Messages sent between executors
+    assert len(build_spans_after_run) == 0  # No build spans should be present after clear
+
+    # Verify workflow span events
+    workflow_span = workflow_spans[0]
+    assert workflow_span.events is not None
+    event_names = [event.name for event in workflow_span.events]
+    assert "workflow.started" in event_names
+    assert "workflow.completed" in event_names
+
+    # Test fan-in span linking: find the aggregator's processing span
+    aggregator_spans = [s for s in processing_spans if s.attributes and s.attributes.get("executor.id") == "aggregator"]
+    assert len(aggregator_spans) == 1
+
+    aggregator_span = aggregator_spans[0]
+    # The aggregator span should have links to the source spans (from executor2 and executor3)
+    # This tests that FanInEdgeRunner properly handles multiple trace contexts and span IDs
+    assert aggregator_span.links is not None
+
+    # Find the sending spans from executor2 and executor3 by checking parent relationships
+    executor2_processing_spans = [
+        s for s in processing_spans if s.attributes and s.attributes.get("executor.id") == "executor2"
+    ]
+    executor3_processing_spans = [
+        s for s in processing_spans if s.attributes and s.attributes.get("executor.id") == "executor3"
+    ]
+
+    # Get span IDs from processing spans
+    executor2_processing_span_ids = {format(s.context.span_id, "016x") for s in executor2_processing_spans if s.context}
+    executor3_processing_span_ids = {format(s.context.span_id, "016x") for s in executor3_processing_spans if s.context}
+
+    executor2_sending_spans = [
+        s for s in sending_spans if s.parent and format(s.parent.span_id, "016x") in executor2_processing_span_ids
+    ]
+    executor3_sending_spans = [
+        s for s in sending_spans if s.parent and format(s.parent.span_id, "016x") in executor3_processing_span_ids
+    ]
+
+    # Verify that we have sending spans from both executors
+    assert len(executor2_sending_spans) >= 1, "Should have at least one sending span from executor2"
+    assert len(executor3_sending_spans) >= 1, "Should have at least one sending span from executor3"
+
+    # Verify that the aggregator span links point to the correct source spans
+    linked_span_ids = {link.context.span_id for link in aggregator_span.links}
+
+    # Should have links from both executor2 and executor3's sending spans
+    executor2_span_ids = {s.context.span_id for s in executor2_sending_spans if s.context}
+    executor3_span_ids = {s.context.span_id for s in executor3_sending_spans if s.context}
+
+    # At least one span from each executor should be linked
+    assert bool(linked_span_ids & executor2_span_ids), "Aggregator should link to executor2's sending span"
+    assert bool(linked_span_ids & executor3_span_ids), "Aggregator should link to executor3's sending span"
+
+    # Should have at least 2 links (one from each source executor)
+    assert len(aggregator_span.links) >= 2, f"Expected at least 2 links, got {len(aggregator_span.links)}"
+
+
+async def test_workflow_error_handling_in_tracing(span_exporter: InMemorySpanExporter) -> None:
+    """Test that workflow errors are properly recorded in traces."""
+
+    class FailingExecutor(Executor):
+        def __init__(self) -> None:
+            super().__init__(id="failing_executor")
+
+        @handler
+        async def handle_message(self, message: str, ctx: WorkflowContext) -> None:
+            raise ValueError("Test error")
+
+    failing_executor = FailingExecutor()
+    workflow = WorkflowBuilder(start_executor=failing_executor).build()
+
+    # Run workflow and expect error
+    with pytest.raises(ValueError, match="Test error"):
+        async for _ in workflow.run("test input", stream=True):
+            pass
+
+    spans = span_exporter.get_finished_spans()
+
+    # Find workflow span
+    workflow_spans = [s for s in spans if s.name == "workflow.run"]
+    assert len(workflow_spans) == 1
+
+    workflow_span = workflow_spans[0]
+
+    # Verify error event and status are recorded
+    assert workflow_span.events is not None
+    event_names = [event.name for event in workflow_span.events]
+    assert "workflow.started" in event_names
+    assert "workflow.error" in event_names
+    assert workflow_span.status.status_code.name == "ERROR"
+
+
+@pytest.mark.parametrize("enable_instrumentation", [False], indirect=True)
+async def test_message_trace_context_serialization(span_exporter: InMemorySpanExporter) -> None:
+    """Test that message trace context is properly serialized/deserialized."""
+    ctx = InProcRunnerContext(InMemoryCheckpointStorage())
+
+    # Create message with trace context
+    message = WorkflowMessage(
+        data="test",
+        source_id="source",
+        target_id="target",
+        trace_contexts=[{"traceparent": "00-trace-span-01"}],
+        source_span_ids=["span123"],
+    )
+
+    await ctx.send_message(message)
+
+    # Create a checkpoint that includes the message
+    checkpoint_id = await ctx.create_checkpoint("test_name", "test_hash", State(), None, 0)
+    checkpoint = await ctx.load_checkpoint(checkpoint_id)
+    assert checkpoint is not None
+
+    # Check serialized message includes trace context
+    serialized_msg = checkpoint.messages["source"][0]
+    assert serialized_msg.trace_contexts == [{"traceparent": "00-trace-span-01"}]
+    assert serialized_msg.source_span_ids == ["span123"]
+
+    # Test deserialization
+    await ctx.apply_checkpoint(checkpoint)
+    restored_messages = await ctx.drain_messages()
+
+    restored_msg = list(restored_messages.values())[0][0]
+    assert restored_msg.trace_context == {"traceparent": "00-trace-span-01"}  # Test backward compatibility
+    assert restored_msg.source_span_id == "span123"  # Test backward compatibility
+    assert restored_msg.trace_contexts == [{"traceparent": "00-trace-span-01"}]  # Test new format
+    assert restored_msg.source_span_ids == ["span123"]  # Test new format
+
+
+async def test_workflow_build_error_tracing(span_exporter: InMemorySpanExporter) -> None:
+    """Test that build errors are properly recorded in build spans."""
+
+    # Create a valid builder, then clear the start executor to trigger a build-time ValueError
+    builder = WorkflowBuilder(start_executor=MockExecutor(id="mock"))
+    builder._start_executor = None  # type: ignore[assignment]
+
+    with pytest.raises(ValueError):
+        builder.build()
+
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+
+    build_span = spans[0]
+    assert build_span.name == "workflow.build"
+
+    # Verify error status and events
+    assert build_span.status.status_code.name == "ERROR"
+    assert build_span.events is not None
+
+    event_names = [event.name for event in build_span.events]
+    assert "build.started" in event_names
+    assert "build.error" in event_names
+
+    # Check error event attributes
+    error_events = [event for event in build_span.events if event.name == "build.error"]
+    assert len(error_events) == 1
+
+    error_event = error_events[0]
+    assert error_event.attributes is not None
+    assert "starting executor" in str(error_event.attributes.get("build.error.message")).lower()
+    assert error_event.attributes.get("build.error.type") == "ValueError"

python/packages/core/tests/workflow/test_workflow_states.py

@@ -0,0 +1,220 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from typing import Any
+
+import pytest
+from typing_extensions import Never
+
+from agent_framework import (
+    Executor,
+    InProcRunnerContext,
+    Workflow,
+    WorkflowBuilder,
+    WorkflowContext,
+    WorkflowEvent,
+    WorkflowEventSource,
+    WorkflowRunResult,
+    WorkflowRunState,
+    handler,
+)
+from agent_framework._workflows._state import State
+
+
+class FailingExecutor(Executor):
+    """Executor that raises at runtime to test failure signaling."""
+
+    @handler
+    async def fail(self, msg: int, ctx: WorkflowContext) -> None:  # pragma: no cover - invoked via workflow
+        raise RuntimeError("boom")
+
+
+async def test_executor_failed_and_workflow_failed_events_streaming():
+    failing = FailingExecutor(id="f")
+    wf: Workflow = WorkflowBuilder(start_executor=failing).build()
+
+    events: list[object] = []
+    with pytest.raises(RuntimeError, match="boom"):
+        async for ev in wf.run(0, stream=True):
+            events.append(ev)
+
+    # executor_failed event (type='executor_failed') should be emitted before workflow failed event
+    executor_failed_events: list[WorkflowEvent[Any]] = [
+        e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_failed"
+    ]
+    assert executor_failed_events, "executor_failed event should be emitted when start executor fails"
+    assert executor_failed_events[0].executor_id == "f"
+    assert executor_failed_events[0].origin is WorkflowEventSource.FRAMEWORK
+
+    # Workflow-level failure and FAILED status should be surfaced
+    failed_events: list[WorkflowEvent[Any]] = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "failed"]
+    assert failed_events
+    assert all(e.origin is WorkflowEventSource.FRAMEWORK for e in failed_events)
+    status: list[WorkflowEvent[Any]] = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "status"]
+    assert status and status[-1].state == WorkflowRunState.FAILED
+    assert all(e.origin is WorkflowEventSource.FRAMEWORK for e in status)
+
+    # Verify executor_failed event comes before workflow failed event
+    executor_failed_idx = events.index(executor_failed_events[0])
+    workflow_failed_idx = events.index(failed_events[0])
+    assert executor_failed_idx < workflow_failed_idx, (
+        "executor_failed event should be emitted before workflow failed event"
+    )
+
+
+async def test_executor_failed_event_emitted_on_direct_execute():
+    failing = FailingExecutor(id="f")
+    ctx = InProcRunnerContext()
+    state = State()
+    with pytest.raises(RuntimeError, match="boom"):
+        await failing.execute(
+            0,
+            ["START"],
+            state,
+            ctx,
+        )
+    drained = await ctx.drain_events()
+    failed = [e for e in drained if isinstance(e, WorkflowEvent) and e.type == "executor_failed"]
+    assert failed
+    assert all(e.origin is WorkflowEventSource.FRAMEWORK for e in failed)
+
+
+class PassthroughExecutor(Executor):
+    """Executor that passes message to the next executor."""
+
+    @handler
+    async def passthrough(self, msg: int, ctx: WorkflowContext[int]) -> None:
+        await ctx.send_message(msg)
+
+
+async def test_executor_failed_event_from_second_executor_in_chain():
+    """Test that executor_failed event is emitted when a non-start executor fails."""
+    passthrough = PassthroughExecutor(id="passthrough")
+    failing = FailingExecutor(id="failing")
+    wf: Workflow = WorkflowBuilder(start_executor=passthrough).add_edge(passthrough, failing).build()
+
+    events: list[object] = []
+    with pytest.raises(RuntimeError, match="boom"):
+        async for ev in wf.run(0, stream=True):
+            events.append(ev)
+
+    # executor_failed event should be emitted for the failing executor
+    executor_failed_events: list[WorkflowEvent[Any]] = [
+        e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_failed"
+    ]
+    assert executor_failed_events, "executor_failed event should be emitted when second executor fails"
+    assert executor_failed_events[0].executor_id == "failing"
+    assert executor_failed_events[0].origin is WorkflowEventSource.FRAMEWORK
+
+    # Workflow-level failure should also be surfaced
+    failed_events: list[WorkflowEvent[Any]] = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "failed"]
+    assert failed_events
+    assert all(e.origin is WorkflowEventSource.FRAMEWORK for e in failed_events)
+
+    # Verify executor_failed event comes before workflow failed event
+    executor_failed_idx = events.index(executor_failed_events[0])
+    workflow_failed_idx = events.index(failed_events[0])
+    assert executor_failed_idx < workflow_failed_idx, (
+        "executor_failed event should be emitted before workflow failed event"
+    )
+
+
+class SimpleExecutor(Executor):
+    """Executor that does nothing, for testing."""
+
+    @handler
+    async def run(self, msg: str, ctx: WorkflowContext[str]) -> None:  # pragma: no cover
+        await ctx.send_message(msg)
+
+
+class Requester(Executor):
+    """Executor that always requests external info to test idle-with-requests state."""
+
+    @handler
+    async def ask(self, _: str, ctx: WorkflowContext) -> None:  # pragma: no cover
+        await ctx.request_info("Mock request data", str)
+
+
+async def test_idle_with_pending_requests_status_streaming():
+    simple_executor = SimpleExecutor(id="simple")
+    requester = Requester(id="req")
+    wf = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, requester).build()
+
+    events = [ev async for ev in wf.run("start", stream=True)]  # Consume stream fully
+
+    # Ensure a request was emitted
+    assert any(isinstance(e, WorkflowEvent) and e.type == "request_info" for e in events)
+    status_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "status"]
+    assert len(status_events) >= 3
+    assert status_events[-2].state == WorkflowRunState.IN_PROGRESS_PENDING_REQUESTS
+    assert status_events[-1].state == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS
+
+
+class Completer(Executor):
+    """Executor that completes immediately with provided data for testing."""
+
+    @handler
+    async def run(self, msg: str, ctx: WorkflowContext[Never, str]) -> None:  # pragma: no cover  # type: ignore[valid-type]
+        await ctx.yield_output(msg)
+
+
+async def test_completed_status_streaming():
+    c = Completer(id="c")
+    wf = WorkflowBuilder(start_executor=c).build()
+    events = [ev async for ev in wf.run("ok", stream=True)]  # no raise
+    # Last status should be IDLE
+    status = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "status"]
+    assert status and status[-1].state == WorkflowRunState.IDLE
+    assert all(e.origin is WorkflowEventSource.FRAMEWORK for e in status)
+
+
+async def test_started_and_completed_event_origins():
+    c = Completer(id="c-origin")
+    wf = WorkflowBuilder(start_executor=c).build()
+    events = [ev async for ev in wf.run("payload", stream=True)]
+
+    started = next(e for e in events if isinstance(e, WorkflowEvent) and e.type == "started")
+    assert started.origin is WorkflowEventSource.FRAMEWORK
+
+    # Check for IDLE status indicating completion
+    idle_status = next(
+        (e for e in events if isinstance(e, WorkflowEvent) and e.type == "status" and e.state == WorkflowRunState.IDLE),
+        None,
+    )
+    assert idle_status is not None
+    assert idle_status.origin is WorkflowEventSource.FRAMEWORK
+
+
+async def test_non_streaming_final_state_helpers():
+    # Completed case
+    c = Completer(id="c")
+    wf1 = WorkflowBuilder(start_executor=c).build()
+    result1: WorkflowRunResult = await wf1.run("done")
+    assert result1.get_final_state() == WorkflowRunState.IDLE
+
+    # Idle-with-pending-request case
+    simple_executor = SimpleExecutor(id="simple")
+    requester = Requester(id="req")
+    wf2 = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, requester).build()
+    result2: WorkflowRunResult = await wf2.run("start")
+    assert result2.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS
+
+
+async def test_run_includes_status_events_completed():
+    c = Completer(id="c2")
+    wf = WorkflowBuilder(start_executor=c).build()
+    result: WorkflowRunResult = await wf.run("ok")
+    timeline = result.status_timeline()
+    assert timeline, "Expected status timeline in non-streaming run() results"
+    assert timeline[-1].state == WorkflowRunState.IDLE
+
+
+async def test_run_includes_status_events_idle_with_requests():
+    simple_executor = SimpleExecutor(id="simple")
+    requester = Requester(id="req2")
+    wf = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, requester).build()
+    result: WorkflowRunResult = await wf.run("start")
+    timeline = result.status_timeline()
+    assert timeline, "Expected status timeline in non-streaming run() results"
+    assert len(timeline) >= 3
+    assert timeline[-2].state == WorkflowRunState.IN_PROGRESS_PENDING_REQUESTS
+    assert timeline[-1].state == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS

python/packages/core/tests/workflow/test_workflow_status.py

@@ -0,0 +1,149 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for the ``Workflow.status`` property."""
+
+from dataclasses import dataclass
+
+import pytest
+
+from agent_framework import (
+    Executor,
+    Workflow,
+    WorkflowBuilder,
+    WorkflowContext,
+    WorkflowEvent,
+    WorkflowRunState,
+    handler,
+    response_handler,
+)
+from agent_framework._workflows._executor import Executor as _Executor
+from agent_framework._workflows._request_info_mixin import RequestInfoMixin
+
+
+class PassThroughExecutor(Executor):
+    """Executor that yields its input as a workflow output and stops."""
+
+    @handler
+    async def passthrough(self, msg: str, ctx: WorkflowContext[str, str]) -> None:
+        await ctx.yield_output(msg)
+
+
+class FailingExecutor(Executor):
+    """Executor that raises at runtime to drive the FAILED status."""
+
+    @handler
+    async def fail(self, msg: int, ctx: WorkflowContext) -> None:  # pragma: no cover - invoked via workflow
+        raise RuntimeError("boom")
+
+
+@dataclass
+class _ApprovalRequest:
+    prompt: str
+    request_id: str = ""
+
+    def __post_init__(self) -> None:
+        if not self.request_id:
+            import uuid
+
+            self.request_id = str(uuid.uuid4())
+
+
+class ApprovalExecutor(_Executor, RequestInfoMixin):
+    """Executor that issues a single request_info call and finalizes on response."""
+
+    def __init__(self, id: str = "approval"):
+        super().__init__(id=id)
+
+    @handler
+    async def start(self, message: str, ctx: WorkflowContext[str, str]) -> None:
+        await ctx.request_info(_ApprovalRequest(prompt=message), bool)
+
+    @response_handler
+    async def on_response(
+        self, original_request: _ApprovalRequest, approved: bool, ctx: WorkflowContext[str, str]
+    ) -> None:
+        await ctx.yield_output(f"approved={approved}")
+
+
+def _build_passthrough_workflow() -> Workflow:
+    executor = PassThroughExecutor(id="p")
+    return WorkflowBuilder(start_executor=executor, output_from=[executor]).build()
+
+
+def _build_failing_workflow() -> Workflow:
+    # FailingExecutor has no workflow_output_types, so we leave designation
+    # implicit; the deprecation warning is filtered at call sites that need it.
+    return WorkflowBuilder(start_executor=FailingExecutor(id="f")).build()
+
+
+def _build_approval_workflow() -> Workflow:
+    executor = ApprovalExecutor(id="approval")
+    return WorkflowBuilder(start_executor=executor, output_from=[executor]).build()
+
+
+async def test_status_default_is_idle_before_first_run():
+    wf = _build_passthrough_workflow()
+    assert wf.status is WorkflowRunState.IDLE
+
+
+async def test_status_is_idle_after_successful_run():
+    wf = _build_passthrough_workflow()
+    await wf.run("hello")
+    assert wf.status is WorkflowRunState.IDLE
+
+
+async def test_status_is_failed_after_failure():
+    wf = _build_failing_workflow()
+    with pytest.raises(RuntimeError, match="boom"):
+        await wf.run(0)
+    assert wf.status is WorkflowRunState.FAILED
+
+
+async def test_status_transitions_during_streaming_run():
+    """Workflow.status mirrors the most recent emitted status event."""
+    wf = _build_passthrough_workflow()
+    observed: list[WorkflowRunState] = []
+
+    async for event in wf.run("hi", stream=True):
+        if isinstance(event, WorkflowEvent) and event.type == "status":
+            # By the time a status event surfaces to the consumer, the property
+            # must already reflect that state (updated in lockstep with emission).
+            assert wf.status == event.state
+            observed.append(event.state)  # type: ignore
+
+    # IN_PROGRESS must precede IDLE; both must appear.
+    assert WorkflowRunState.IN_PROGRESS in observed
+    assert observed[-1] is WorkflowRunState.IDLE
+    assert wf.status is WorkflowRunState.IDLE
+
+
+async def test_status_idle_with_pending_requests_then_resolves_to_idle():
+    wf = _build_approval_workflow()
+
+    request_event: WorkflowEvent | None = None
+    async for event in wf.run("please approve", stream=True):
+        if isinstance(event, WorkflowEvent) and event.type == "request_info":
+            request_event = event
+
+    assert request_event is not None
+    assert wf.status is WorkflowRunState.IDLE_WITH_PENDING_REQUESTS
+
+    async for _ in wf.run(stream=True, responses={request_event.request_id: True}):
+        pass
+
+    assert wf.status is WorkflowRunState.IDLE
+
+
+async def test_status_in_progress_pending_requests_observed_mid_run():
+    """While streaming, status reaches IN_PROGRESS_PENDING_REQUESTS after a request_info event."""
+    wf = _build_approval_workflow()
+    seen_states: list[WorkflowRunState] = []
+
+    async for event in wf.run("please approve", stream=True):
+        if isinstance(event, WorkflowEvent) and event.type == "status":
+            seen_states.append(event.state)  # type: ignore
+
+    assert WorkflowRunState.IN_PROGRESS in seen_states
+    assert WorkflowRunState.IN_PROGRESS_PENDING_REQUESTS in seen_states
+    assert seen_states[-1] is WorkflowRunState.IDLE_WITH_PENDING_REQUESTS
+    assert wf.status is WorkflowRunState.IDLE_WITH_PENDING_REQUESTS