Python: [BREAKING] Main to core (#983)

* removed pydantic from types

* fix assistants client

* Remove Pydantic usage from workflow code.

* updated lock and test fixes

* moved main to core, and setup meta package

* updated versions

* updated lock

* fixed agents dependency

* added retry to merge tests

---------

Co-authored-by: Evan Mattson <evan.mattson@microsoft.com>

python/packages/core/.vscode/launch.json

@@ -0,0 +1,15 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Python Debugger: Current File",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "${file}",
+            "console": "integratedTerminal"
+        }
+    ]
+}

python/packages/core/LICENSE

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

python/packages/core/README.md

@@ -0,0 +1,233 @@
+# Get Started with Microsoft Agent Framework
+
+Highlights
+
+- Flexible Agent Framework: build, orchestrate, and deploy AI agents and multi-agent systems
+- Multi-Agent Orchestration: Group chat, sequential, concurrent, and handoff patterns
+- Plugin Ecosystem: Extend with native functions, OpenAPI, Model Context Protocol (MCP), and more
+- LLM Support: OpenAI, Azure OpenAI, Azure AI, and more
+- Runtime Support: In-process and distributed agent execution
+- Multimodal: Text, vision, and function calling
+- Cross-Platform: .NET and Python implementations
+
+## Quick Install
+
+```bash
+pip install agent-framework[all]
+# Optional: Add Azure AI integration
+pip install agent-framework-azure-ai
+# Optional: Both
+pip install agent-framework-azure-ai agent-framework-copilotstudio
+```
+
+Supported Platforms:
+
+- Python: 3.10+
+- OS: Windows, macOS, Linux
+
+## 1. Setup API Keys
+
+Set as environment variables, or create a .env file at your project root:
+
+```bash
+OPENAI_API_KEY=sk-...
+OPENAI_CHAT_MODEL_ID=...
+OPENAI_RESPONSES_MODEL_ID=...
+...
+AZURE_OPENAI_API_KEY=...
+AZURE_OPENAI_ENDPOINT=...
+AZURE_OPENAI_CHAT_DEPLOYMENT_NAME=...
+...
+AZURE_AI_PROJECT_ENDPOINT=...
+AZURE_AI_MODEL_DEPLOYMENT_NAME=...
+```
+
+You can also override environment variables by explicitly passing configuration parameters to the chat client constructor:
+
+```python
+from agent_framework.azure import AzureOpenAIChatClient
+
+chat_client = AzureOpenAIChatClient(
+    api_key="",
+    endpoint="",
+    deployment_name="",
+    api_version="",
+)
+```
+
+See the following [setup guide](https://github.com/microsoft/agent-framework/tree/main/python/samples/getting_started) for more information.
+
+## 2. Create a Simple Agent
+
+Create agents and invoke them directly:
+
+```python
+import asyncio
+from agent_framework import ChatAgent
+from agent_framework.openai import OpenAIChatClient
+
+async def main():
+    agent = ChatAgent(
+        chat_client=OpenAIChatClient(),
+        instructions="""
+        1) A robot may not injure a human being...
+        2) A robot must obey orders given it by human beings...
+        3) A robot must protect its own existence...
+
+        Give me the TLDR in exactly 5 words.
+        """
+    )
+
+    result = await agent.run("Summarize the Three Laws of Robotics")
+    print(result)
+
+asyncio.run(main())
+# Output: Protect humans, obey, self-preserve, prioritized.
+```
+
+## 3. Directly Use Chat Clients (No Agent Required)
+
+You can use the chat client classes directly for advanced workflows:
+
+```python
+import asyncio
+from agent_framework.openai import OpenAIChatClient
+from agent_framework import ChatMessage, Role
+
+async def main():
+    client = OpenAIChatClient()
+
+    messages = [
+        ChatMessage(role=Role.SYSTEM, text="You are a helpful assistant."),
+        ChatMessage(role=Role.USER, text="Write a haiku about Agent Framework.")
+    ]
+
+    response = await client.get_response(messages)
+    print(response.messages[0].text)
+
+    """
+    Output:
+
+    Agents work in sync,
+    Framework threads through each task—
+    Code sparks collaboration.
+    """
+
+asyncio.run(main())
+```
+
+## 4. Build an Agent with Tools and Functions
+
+Enhance your agent with custom tools and function calling:
+
+```python
+import asyncio
+from typing import Annotated
+from random import randint
+from pydantic import Field
+from agent_framework import ChatAgent
+from agent_framework.openai import OpenAIChatClient
+
+
+def get_weather(
+    location: Annotated[str, Field(description="The location to get the weather for.")],
+) -> str:
+    """Get the weather for a given location."""
+    conditions = ["sunny", "cloudy", "rainy", "stormy"]
+    return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C."
+
+
+def get_menu_specials() -> str:
+    """Get today's menu specials."""
+    return """
+    Special Soup: Clam Chowder
+    Special Salad: Cobb Salad
+    Special Drink: Chai Tea
+    """
+
+
+async def main():
+    agent = ChatAgent(
+        chat_client=OpenAIChatClient(),
+        instructions="You are a helpful assistant that can provide weather and restaurant information.",
+        tools=[get_weather, get_menu_specials]
+    )
+
+    response = await agent.run("What's the weather in Amsterdam and what are today's specials?")
+    print(response)
+
+    # Output:
+    # The weather in Amsterdam is sunny with a high of 22°C. Today's specials include
+    # Clam Chowder soup, Cobb Salad, and Chai Tea as the special drink.
+
+asyncio.run(main())
+```
+
+You can explore additional agent samples [here](https://github.com/microsoft/agent-framework/tree/main/python/samples/getting_started/agents).
+
+## 5. Multi-Agent Orchestration
+
+Coordinate multiple agents to collaborate on complex tasks using orchestration patterns:
+
+```python
+import asyncio
+from agent_framework import ChatAgent
+from agent_framework.openai import OpenAIChatClient
+
+
+async def main():
+    # Create specialized agents
+    writer = ChatAgent(
+        chat_client=OpenAIChatClient(),
+        name="Writer",
+        instructions="You are a creative content writer. Generate and refine slogans based on feedback."
+    )
+
+    reviewer = ChatAgent(
+        chat_client=OpenAIChatClient(),
+        name="Reviewer",
+        instructions="You are a critical reviewer. Provide detailed feedback on proposed slogans."
+    )
+
+    # Sequential workflow: Writer creates, Reviewer provides feedback
+    task = "Create a slogan for a new electric SUV that is affordable and fun to drive."
+
+    # Step 1: Writer creates initial slogan
+    initial_result = await writer.run(task)
+    print(f"Writer: {initial_result}")
+
+    # Step 2: Reviewer provides feedback
+    feedback_request = f"Please review this slogan: {initial_result}"
+    feedback = await reviewer.run(feedback_request)
+    print(f"Reviewer: {feedback}")
+
+    # Step 3: Writer refines based on feedback
+    refinement_request = f"Please refine this slogan based on the feedback: {initial_result}\nFeedback: {feedback}"
+    final_result = await writer.run(refinement_request)
+    print(f"Final Slogan: {final_result}")
+
+    # Example Output:
+    # Writer: "Charge Forward: Affordable Adventure Awaits!"
+    # Reviewer: "Good energy, but 'Charge Forward' is overused in EV marketing..."
+    # Final Slogan: "Power Up Your Adventure: Premium Feel, Smart Price!"
+
+if __name__ == "__main__":
+    asyncio.run(main())
+```
+
+**Note**: Advanced orchestration patterns like GroupChat, Sequential, and Concurrent orchestrations are coming soon.
+
+## More Examples & Samples
+
+- [Getting Started with Agents](https://github.com/microsoft/agent-framework/tree/main/python/samples/getting_started/agents): Basic agent creation and tool usage
+- [Chat Client Examples](https://github.com/microsoft/agent-framework/tree/main/python/samples/getting_started/chat_client): Direct chat client usage patterns
+- [Azure AI Integration](https://github.com/microsoft/agent-framework/tree/main/python/packages/azure-ai): Azure AI integration
+- [.NET Orchestration Samples](https://github.com/microsoft/agent-framework/tree/main/dotnet/samples/GettingStarted/Orchestration): Advanced multi-agent patterns (.NET)
+
+## Agent Framework Documentation
+
+- [Agent Framework Repository](https://github.com/microsoft/agent-framework)
+- [Python Package Documentation](https://github.com/microsoft/agent-framework/tree/main/python)
+- [.NET Package Documentation](https://github.com/microsoft/agent-framework/tree/main/dotnet)
+- [Design Documents](https://github.com/microsoft/agent-framework/tree/main/docs/design)
+- Learn docs are coming soon.

python/packages/core/agent_framework/__init__.py

@@ -0,0 +1,23 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import importlib
+import importlib.metadata
+from typing import Final
+
+try:
+    _version = importlib.metadata.version(__name__)
+except importlib.metadata.PackageNotFoundError:
+    _version = "0.0.0"  # Fallback for development mode
+__version__: Final[str] = _version
+
+from ._agents import *  # noqa: F403
+from ._clients import *  # noqa: F403
+from ._logging import *  # noqa: F403
+from ._mcp import *  # noqa: F403
+from ._memory import *  # noqa: F403
+from ._middleware import *  # noqa: F403
+from ._telemetry import *  # noqa: F403
+from ._threads import *  # noqa: F403
+from ._tools import *  # noqa: F403
+from ._types import *  # noqa: F403
+from ._workflow import *  # noqa: F403

python/packages/core/agent_framework/_agents.py

@@ -0,0 +1,856 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import inspect
+import sys
+from collections.abc import AsyncIterable, Awaitable, Callable, MutableMapping, Sequence
+from contextlib import AbstractAsyncContextManager, AsyncExitStack
+from copy import copy
+from itertools import chain
+from typing import Any, ClassVar, Literal, Protocol, TypeVar, cast, runtime_checkable
+from uuid import uuid4
+
+from pydantic import BaseModel, Field, create_model
+
+from ._clients import BaseChatClient, ChatClientProtocol
+from ._logging import get_logger
+from ._mcp import MCPTool
+from ._memory import AggregateContextProvider, Context, ContextProvider
+from ._middleware import Middleware, use_agent_middleware
+from ._threads import AgentThread, ChatMessageStoreProtocol
+from ._tools import FUNCTION_INVOKING_CHAT_CLIENT_MARKER, AIFunction, ToolProtocol
+from ._types import (
+    AgentRunResponse,
+    AgentRunResponseUpdate,
+    ChatMessage,
+    ChatOptions,
+    ChatResponse,
+    ChatResponseUpdate,
+    Role,
+    ToolMode,
+)
+from .exceptions import AgentExecutionException
+from .observability import use_agent_observability
+
+if sys.version_info >= (3, 12):
+    from typing import override  # type: ignore # pragma: no cover
+else:
+    from typing_extensions import override  # type: ignore[import] # pragma: no cover
+if sys.version_info >= (3, 11):
+    from typing import Self  # pragma: no cover
+else:
+    from typing_extensions import Self  # pragma: no cover
+
+logger = get_logger("agent_framework")
+
+TThreadType = TypeVar("TThreadType", bound="AgentThread")
+
+__all__ = ["AgentProtocol", "BaseAgent", "ChatAgent"]
+
+
+# region Agent Protocol
+
+
+@runtime_checkable
+class AgentProtocol(Protocol):
+    """A protocol for an agent that can be invoked."""
+
+    @property
+    def id(self) -> str:
+        """Returns the ID of the agent."""
+        ...
+
+    @property
+    def name(self) -> str | None:
+        """Returns the name of the agent."""
+        ...
+
+    @property
+    def display_name(self) -> str:
+        """Returns the display name of the agent."""
+        ...
+
+    @property
+    def description(self) -> str | None:
+        """Returns the description of the agent."""
+        ...
+
+    async def run(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
+        *,
+        thread: AgentThread | None = None,
+        **kwargs: Any,
+    ) -> AgentRunResponse:
+        """Get a response from the agent.
+
+        This method returns the final result of the agent's execution
+        as a single AgentRunResponse object. The caller is blocked until
+        the final result is available.
+
+        Note: For streaming responses, use the run_stream method, which returns
+        intermediate steps and the final result as a stream of AgentRunResponseUpdate
+        objects. Streaming only the final result is not feasible because the timing of
+        the final result's availability is unknown, and blocking the caller until then
+        is undesirable in streaming scenarios.
+
+        Args:
+            messages: The message(s) to send to the agent.
+            thread: The conversation thread associated with the message(s).
+            kwargs: Additional keyword arguments.
+
+        Returns:
+            An agent response item.
+        """
+        ...
+
+    def run_stream(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
+        *,
+        thread: AgentThread | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterable[AgentRunResponseUpdate]:
+        """Run the agent as a stream.
+
+        This method will return the intermediate steps and final results of the
+        agent's execution as a stream of AgentRunResponseUpdate objects to the caller.
+
+        Note: An AgentRunResponseUpdate object contains a chunk of a message.
+
+        Args:
+            messages: The message(s) to send to the agent.
+            thread: The conversation thread associated with the message(s).
+            kwargs: Additional keyword arguments.
+
+        Yields:
+            An agent response item.
+        """
+        ...
+
+    def get_new_thread(self, **kwargs: Any) -> AgentThread:
+        """Creates a new conversation thread for the agent."""
+        ...
+
+
+# region BaseAgent
+
+
+class BaseAgent:
+    """Base class for all Agent Framework agents.
+
+    Attributes:
+       id: The unique identifier of the agent  If no id is provided,
+           a new UUID will be generated.
+       name: The name of the agent, can be None.
+       description: The description of the agent.
+       display_name: The display name of the agent, which is either the name or id.
+       context_providers: The collection of multiple context providers to include during agent invocation.
+       middleware: List of middleware to intercept agent and function invocations.
+    """
+
+    def __init__(
+        self,
+        id: str | None = None,
+        name: str | None = None,
+        description: str | None = None,
+        context_providers: ContextProvider | Sequence[ContextProvider] | None = None,
+        middleware: Middleware | Sequence[Middleware] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Base class for all Agent Framework agents.
+
+        Args:
+            id: The unique identifier of the agent  If no id is provided,
+                a new UUID will be generated.
+            name: The name of the agent, can be None.
+            description: The description of the agent.
+            display_name: The display name of the agent, which is either the name or id.
+            context_providers: The collection of multiple context providers to include during agent invocation.
+            middleware: List of middleware to intercept agent and function invocations.
+            kwargs: will be stored in `additional_properties`
+        """
+        if id is None:
+            id = str(uuid4())
+        self.id = id
+        self.name = name
+        self.description = description
+        self.context_provider = self._prepare_context_providers(context_providers)
+        if middleware is None or isinstance(middleware, Sequence):
+            self.middleware: list[Middleware] | None = cast(list[Middleware], middleware) if middleware else None
+        else:
+            self.middleware = [middleware]
+        self.additional_properties = kwargs
+
+    async def _notify_thread_of_new_messages(
+        self,
+        thread: AgentThread,
+        input_messages: ChatMessage | Sequence[ChatMessage],
+        response_messages: ChatMessage | Sequence[ChatMessage],
+    ) -> None:
+        """Notify the thread of new messages.
+
+        This also calls the invoked method of a potential context provider on the thread.
+        """
+        if isinstance(input_messages, ChatMessage) or len(input_messages) > 0:
+            await thread.on_new_messages(input_messages)
+        if isinstance(response_messages, ChatMessage) or len(response_messages) > 0:
+            await thread.on_new_messages(response_messages)
+        if thread.context_provider:
+            await thread.context_provider.invoked(input_messages, response_messages)
+
+    @property
+    def display_name(self) -> str:
+        """Returns the display name of the agent.
+
+        This is the name if present, otherwise the id.
+        """
+        return self.name or self.id
+
+    def get_new_thread(self, **kwargs: Any) -> AgentThread:
+        """Returns AgentThread instance that is compatible with the agent."""
+        return AgentThread(**kwargs, context_provider=self.context_provider)
+
+    async def deserialize_thread(self, serialized_thread: Any, **kwargs: Any) -> AgentThread:
+        """Deserializes the thread."""
+        thread: AgentThread = self.get_new_thread()
+        await thread.deserialize(serialized_thread, **kwargs)
+        return thread
+
+    def as_tool(
+        self,
+        *,
+        name: str | None = None,
+        description: str | None = None,
+        arg_name: str = "task",
+        arg_description: str | None = None,
+        stream_callback: Callable[[AgentRunResponseUpdate], None]
+        | Callable[[AgentRunResponseUpdate], Awaitable[None]]
+        | None = None,
+    ) -> AIFunction[BaseModel, str]:
+        """Create an AIFunction tool that wraps this agent.
+
+        Args:
+            name: The name for the tool. If None, uses the agent's name.
+            description: The description for the tool. If None, uses the agent's description or empty string.
+            arg_name: The name of the function argument (default: "task").
+            arg_description: The description for the function argument.
+                If None, defaults to "Input for {self.display_name}".
+            stream_callback: Optional callback for streaming responses. If provided, uses run_stream.
+
+        Returns:
+            An AIFunction that can be used as a tool by other agents.
+        """
+        # Verify that self implements AgentProtocol
+        if not isinstance(self, AgentProtocol):
+            raise TypeError(f"Agent {self.__class__.__name__} must implement AgentProtocol to be used as a tool")
+
+        tool_name = name or self.name
+        if tool_name is None:
+            raise ValueError("Agent tool name cannot be None. Either provide a name parameter or set the agent's name.")
+        tool_description = description or self.description or ""
+        argument_description = arg_description or f"Task for {tool_name}"
+
+        # Create dynamic input model with the specified argument name
+        field_info = Field(..., description=argument_description)
+        input_model = create_model(f"{name or self.name or 'agent'}_task", **{arg_name: (str, field_info)})  # type: ignore[call-overload]
+
+        # Check if callback is async once, outside the wrapper
+        is_async_callback = stream_callback is not None and inspect.iscoroutinefunction(stream_callback)
+
+        async def agent_wrapper(**kwargs: Any) -> str:
+            """Wrapper function that calls the agent."""
+            # Extract the input from kwargs using the specified arg_name
+            input_text = kwargs.get(arg_name, "")
+
+            if stream_callback is None:
+                # Use non-streaming mode
+                return (await self.run(input_text)).text
+
+            # Use streaming mode - accumulate updates and create final response
+            response_updates: list[AgentRunResponseUpdate] = []
+            async for update in self.run_stream(input_text):
+                response_updates.append(update)
+                if is_async_callback:
+                    await stream_callback(update)  # type: ignore[misc]
+                else:
+                    stream_callback(update)
+
+            # Create final text from accumulated updates
+            return AgentRunResponse.from_agent_run_response_updates(response_updates).text
+
+        return AIFunction(
+            name=tool_name,
+            description=tool_description,
+            func=agent_wrapper,
+            input_model=input_model,
+        )
+
+    def _normalize_messages(
+        self,
+        messages: str | ChatMessage | Sequence[str] | Sequence[ChatMessage] | None = None,
+    ) -> list[ChatMessage]:
+        if messages is None:
+            return []
+
+        if isinstance(messages, str):
+            return [ChatMessage(role=Role.USER, text=messages)]
+
+        if isinstance(messages, ChatMessage):
+            return [messages]
+
+        return [ChatMessage(role=Role.USER, text=msg) if isinstance(msg, str) else msg for msg in messages]
+
+    def _prepare_context_providers(
+        self,
+        context_providers: ContextProvider | Sequence[ContextProvider] | None = None,
+    ) -> AggregateContextProvider | None:
+        if not context_providers:
+            return None
+
+        if isinstance(context_providers, AggregateContextProvider):
+            return context_providers
+
+        return AggregateContextProvider(context_providers)
+
+
+# region ChatAgent
+
+
+@use_agent_middleware
+@use_agent_observability
+class ChatAgent(BaseAgent):
+    """A Chat Client Agent."""
+
+    AGENT_SYSTEM_NAME: ClassVar[str] = "microsoft.agent_framework"
+
+    def __init__(
+        self,
+        chat_client: ChatClientProtocol,
+        instructions: str | None = None,
+        *,
+        id: str | None = None,
+        name: str | None = None,
+        description: str | None = None,
+        chat_message_store_factory: Callable[[], ChatMessageStoreProtocol] | None = None,
+        context_providers: ContextProvider | list[ContextProvider] | AggregateContextProvider | None = None,
+        middleware: Middleware | list[Middleware] | None = None,
+        frequency_penalty: float | None = None,
+        logit_bias: dict[str | int, float] | None = None,
+        max_tokens: int | None = None,
+        metadata: dict[str, Any] | None = None,
+        model: str | None = None,
+        presence_penalty: float | None = None,
+        response_format: type[BaseModel] | None = None,
+        seed: int | None = None,
+        stop: str | Sequence[str] | None = None,
+        store: bool | None = None,
+        temperature: float | None = None,
+        tool_choice: ToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = "auto",
+        tools: ToolProtocol
+        | Callable[..., Any]
+        | MutableMapping[str, Any]
+        | Sequence[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
+        | None = None,
+        top_p: float | None = None,
+        user: str | None = None,
+        request_kwargs: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Create a ChatAgent.
+
+        Remarks:
+            The set of attributes from frequency_penalty to additional_properties are used to
+            call the chat client, they can also be passed to both run methods.
+            When both are set, the ones passed to the run methods take precedence.
+
+        Args:
+            chat_client: The chat client to use for the agent.
+            instructions: Optional instructions for the agent.
+            These will be put into the messages sent to the chat client service as a system message.
+            id: The unique identifier for the agent, will be created automatically if not provided.
+            name: The name of the agent.
+            description: A brief description of the agent's purpose.
+            chat_message_store_factory: factory function to create an instance of ChatMessageStoreProtocol.
+                If not provided, the default in-memory store will be used.
+            context_providers: The collection of multiple context providers to include during agent invocation.
+            middleware: List of middleware to intercept agent and function invocations.
+            frequency_penalty: the frequency penalty to use.
+            logit_bias: the logit bias to use.
+            max_tokens: The maximum number of tokens to generate.
+            metadata: additional metadata to include in the request.
+            model: The model to use for the agent.
+            presence_penalty: the presence penalty to use.
+            response_format: the format of the response.
+            seed: the random seed to use.
+            stop: the stop sequence(s) for the request.
+            store: whether to store the response.
+            temperature: the sampling temperature to use.
+            tool_choice: the tool choice for the request.
+            tools: the tools to use for the request.
+            top_p: the nucleus sampling probability to use.
+            user: the user to associate with the request.
+            request_kwargs: a dictionary of other values that will be passed through
+                to the chat_client `get_response` and `get_streaming_response` methods.
+            kwargs: any additional keyword arguments. Will be stored as `additional_properties`
+        """
+        if not hasattr(chat_client, FUNCTION_INVOKING_CHAT_CLIENT_MARKER) and isinstance(chat_client, BaseChatClient):
+            logger.warning(
+                "The provided chat client does not support function invoking, this might limit agent capabilities."
+            )
+
+        super().__init__(
+            id=id,
+            name=name,
+            description=description,
+            context_providers=context_providers,
+            middleware=middleware,
+            **kwargs,
+        )
+        self.chat_client = chat_client
+        self.chat_message_store_factory = chat_message_store_factory
+
+        # We ignore the MCP Servers here and store them separately,
+        # we add their functions to the tools list at runtime
+        normalized_tools: list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]] = (  # type:ignore[reportUnknownVariableType]
+            [] if tools is None else tools if isinstance(tools, list) else [tools]  # type: ignore[list-item]
+        )
+        self._local_mcp_tools = [tool for tool in normalized_tools if isinstance(tool, MCPTool)]
+        agent_tools = [tool for tool in normalized_tools if not isinstance(tool, MCPTool)]
+        self.chat_options = ChatOptions(
+            model_id=model,
+            frequency_penalty=frequency_penalty,
+            instructions=instructions,
+            logit_bias=logit_bias,
+            max_tokens=max_tokens,
+            metadata=metadata,
+            presence_penalty=presence_penalty,
+            response_format=response_format,
+            seed=seed,
+            stop=stop,
+            store=store,
+            temperature=temperature,
+            tool_choice=tool_choice,
+            tools=agent_tools,
+            top_p=top_p,
+            user=user,
+            additional_properties=request_kwargs or {},  # type: ignore
+        )
+        self._async_exit_stack = AsyncExitStack()
+        self._update_agent_name()
+
+    async def __aenter__(self) -> "Self":
+        """Async context manager entry.
+
+        If any of the chat_client, local_mcp_tools, or context_providers are context managers,
+        they will be entered into the async exit stack to ensure proper cleanup.
+
+        This list might be extended in the future.
+        """
+        for context_manager in chain([self.chat_client], self._local_mcp_tools):
+            if isinstance(context_manager, AbstractAsyncContextManager):
+                await self._async_exit_stack.enter_async_context(context_manager)
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: Any,
+    ) -> None:
+        """Async context manager exit.
+
+        Close the async exit stack to ensure all context managers are exited properly.
+        """
+        await self._async_exit_stack.aclose()
+
+    def _update_agent_name(self) -> None:
+        """Update the agent name in a chat client.
+
+        Checks if there is a agent name, the implementation
+        should check if there is already a agent name defined, and if not
+        set it to this value.
+        """
+        if hasattr(self.chat_client, "_update_agent_name") and callable(self.chat_client._update_agent_name):  # type: ignore[reportAttributeAccessIssue, attr-defined]
+            self.chat_client._update_agent_name(self.name)  # type: ignore[reportAttributeAccessIssue, attr-defined]
+
+    async def run(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
+        *,
+        thread: AgentThread | None = None,
+        frequency_penalty: float | None = None,
+        logit_bias: dict[str | int, float] | None = None,
+        max_tokens: int | None = None,
+        metadata: dict[str, Any] | None = None,
+        model: str | None = None,
+        presence_penalty: float | None = None,
+        response_format: type[BaseModel] | None = None,
+        seed: int | None = None,
+        stop: str | Sequence[str] | None = None,
+        store: bool | None = None,
+        temperature: float | None = None,
+        tool_choice: ToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = None,
+        tools: ToolProtocol
+        | Callable[..., Any]
+        | MutableMapping[str, Any]
+        | list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
+        | None = None,
+        top_p: float | None = None,
+        user: str | None = None,
+        additional_properties: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> AgentRunResponse:
+        """Run the agent with the given messages and options.
+
+        Remarks:
+            Since you won't always call the agent.run directly, but it get's called
+            through orchestration, it is advised to set your default values for
+            all the chat client parameters in the agent constructor.
+            If both parameters are used, the ones passed to the run methods take precedence.
+
+        Args:
+            messages: The messages to process.
+            thread: The thread to use for the agent.
+            frequency_penalty: the frequency penalty to use.
+            logit_bias: the logit bias to use.
+            max_tokens: The maximum number of tokens to generate.
+            metadata: additional metadata to include in the request.
+            model: The model to use for the agent.
+            presence_penalty: the presence penalty to use.
+            response_format: the format of the response.
+            seed: the random seed to use.
+            stop: the stop sequence(s) for the request.
+            store: whether to store the response.
+            temperature: the sampling temperature to use.
+            tool_choice: the tool choice for the request.
+            tools: the tools to use for the request.
+            top_p: the nucleus sampling probability to use.
+            user: the user to associate with the request.
+            additional_properties: additional properties to include in the request.
+            kwargs: Additional keyword arguments for the agent.
+                will only be passed to functions that are called.
+        """
+        input_messages = self._normalize_messages(messages)
+        thread, run_chat_options, thread_messages = await self._prepare_thread_and_messages(
+            thread=thread, input_messages=input_messages
+        )
+        normalized_tools: list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]] = (  # type:ignore[reportUnknownVariableType]
+            [] if tools is None else tools if isinstance(tools, list) else [tools]
+        )
+        agent_name = self._get_agent_name()
+
+        # Resolve final tool list (runtime provided tools + local MCP server tools)
+        final_tools: list[ToolProtocol | Callable[..., Any] | dict[str, Any]] = []
+        # Normalize tools argument to a list without mutating the original parameter
+        for tool in normalized_tools:
+            if isinstance(tool, MCPTool):
+                if not tool.is_connected:
+                    await self._async_exit_stack.enter_async_context(tool)
+                final_tools.extend(tool.functions)  # type: ignore
+            else:
+                final_tools.append(tool)  # type: ignore
+
+        for mcp_server in self._local_mcp_tools:
+            if not mcp_server.is_connected:
+                await self._async_exit_stack.enter_async_context(mcp_server)
+            final_tools.extend(mcp_server.functions)
+        response = await self.chat_client.get_response(
+            messages=thread_messages,
+            chat_options=run_chat_options
+            & ChatOptions(
+                model_id=model,
+                conversation_id=thread.service_thread_id,
+                frequency_penalty=frequency_penalty,
+                logit_bias=logit_bias,
+                max_tokens=max_tokens,
+                metadata=metadata,
+                presence_penalty=presence_penalty,
+                response_format=response_format,
+                seed=seed,
+                stop=stop,
+                store=store,
+                temperature=temperature,
+                tool_choice=tool_choice,
+                tools=final_tools,
+                top_p=top_p,
+                user=user,
+                additional_properties=additional_properties or {},
+            ),
+            **kwargs,
+        )
+
+        await self._update_thread_with_type_and_conversation_id(thread, response.conversation_id)
+
+        # Ensure that the author name is set for each message in the response.
+        for message in response.messages:
+            if message.author_name is None:
+                message.author_name = agent_name
+
+        # Only notify the thread of new messages if the chatResponse was successful
+        # to avoid inconsistent messages state in the thread.
+        await self._notify_thread_of_new_messages(thread, input_messages, response.messages)
+        return AgentRunResponse(
+            messages=response.messages,
+            response_id=response.response_id,
+            created_at=response.created_at,
+            usage_details=response.usage_details,
+            value=response.value,
+            raw_representation=response,
+            additional_properties=response.additional_properties,
+        )
+
+    async def run_stream(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
+        *,
+        thread: AgentThread | None = None,
+        frequency_penalty: float | None = None,
+        logit_bias: dict[str | int, float] | None = None,
+        max_tokens: int | None = None,
+        metadata: dict[str, Any] | None = None,
+        model: str | None = None,
+        presence_penalty: float | None = None,
+        response_format: type[BaseModel] | None = None,
+        seed: int | None = None,
+        stop: str | Sequence[str] | None = None,
+        store: bool | None = None,
+        temperature: float | None = None,
+        tool_choice: ToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = None,
+        tools: ToolProtocol
+        | Callable[..., Any]
+        | MutableMapping[str, Any]
+        | list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
+        | None = None,
+        top_p: float | None = None,
+        user: str | None = None,
+        additional_properties: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterable[AgentRunResponseUpdate]:
+        """Stream the agent with the given messages and options.
+
+        Remarks:
+            Since you won't always call the agent.run_stream directly, but it get's called
+            through orchestration, it is advised to set your default values for
+            all the chat client parameters in the agent constructor.
+            If both parameters are used, the ones passed to the run methods take precedence.
+
+        Args:
+            messages: The messages to process.
+            thread: The thread to use for the agent.
+            frequency_penalty: the frequency penalty to use.
+            logit_bias: the logit bias to use.
+            max_tokens: The maximum number of tokens to generate.
+            metadata: additional metadata to include in the request.
+            model: The model to use for the agent.
+            presence_penalty: the presence penalty to use.
+            response_format: the format of the response.
+            seed: the random seed to use.
+            stop: the stop sequence(s) for the request.
+            store: whether to store the response.
+            temperature: the sampling temperature to use.
+            tool_choice: the tool choice for the request.
+            tools: the tools to use for the request.
+            top_p: the nucleus sampling probability to use.
+            user: the user to associate with the request.
+            additional_properties: additional properties to include in the request.
+            kwargs: any additional keyword arguments.
+                will only be passed to functions that are called.
+
+        """
+        input_messages = self._normalize_messages(messages)
+        thread, run_chat_options, thread_messages = await self._prepare_thread_and_messages(
+            thread=thread, input_messages=input_messages
+        )
+        agent_name = self._get_agent_name()
+        response_updates: list[ChatResponseUpdate] = []
+
+        # Resolve final tool list (runtime provided tools + local MCP server tools)
+        final_tools: list[ToolProtocol | MutableMapping[str, Any] | Callable[..., Any]] = []
+        normalized_tools: list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]] = (  # type: ignore[reportUnknownVariableType]
+            [] if tools is None else tools if isinstance(tools, list) else [tools]
+        )
+        # Normalize tools argument to a list without mutating the original parameter
+        for tool in normalized_tools:
+            if isinstance(tool, MCPTool):
+                if not tool.is_connected:
+                    await self._async_exit_stack.enter_async_context(tool)
+                final_tools.extend(tool.functions)  # type: ignore
+            else:
+                final_tools.append(tool)
+
+        for mcp_server in self._local_mcp_tools:
+            if not mcp_server.is_connected:
+                await self._async_exit_stack.enter_async_context(mcp_server)
+            final_tools.extend(mcp_server.functions)
+
+        async for update in self.chat_client.get_streaming_response(
+            messages=thread_messages,
+            chat_options=run_chat_options
+            & ChatOptions(
+                conversation_id=thread.service_thread_id,
+                frequency_penalty=frequency_penalty,
+                logit_bias=logit_bias,
+                max_tokens=max_tokens,
+                metadata=metadata,
+                model_id=model,
+                presence_penalty=presence_penalty,
+                response_format=response_format,
+                seed=seed,
+                stop=stop,
+                store=store,
+                temperature=temperature,
+                tool_choice=tool_choice,
+                tools=final_tools,  # type: ignore[reportArgumentType]
+                top_p=top_p,
+                user=user,
+                additional_properties=additional_properties or {},
+            ),
+            **kwargs,
+        ):
+            response_updates.append(update)
+
+            if update.author_name is None:
+                update.author_name = agent_name
+
+            yield AgentRunResponseUpdate(
+                contents=update.contents,
+                role=update.role,
+                author_name=update.author_name,
+                response_id=update.response_id,
+                message_id=update.message_id,
+                created_at=update.created_at,
+                additional_properties=update.additional_properties,
+                raw_representation=update,
+            )
+
+        response = ChatResponse.from_chat_response_updates(response_updates)
+        await self._update_thread_with_type_and_conversation_id(thread, response.conversation_id)
+        await self._notify_thread_of_new_messages(thread, input_messages, response.messages)
+
+    @override
+    def get_new_thread(
+        self,
+        *,
+        service_thread_id: str | None = None,
+        **kwargs: Any,
+    ) -> AgentThread:
+        """Get a new conversation thread for the agent.
+
+        If you supply a service_thread_id, the thread will be marked as service managed.
+
+        If you don't supply a service_thread_id but have a chat_message_store_factory configured on the agent,
+        that factory will be used to create a message store for the thread and the thread will be
+        managed locally.
+
+        When neither is present, the thread will be created without a service ID or message store,
+        this will be updated based on usage, when you run the agent with this thread.
+        If you run with store=True, the response will respond with a thread_id and that will be set.
+        Otherwise a messages store is created from the default factory.
+
+        Args:
+            service_thread_id: Optional service managed thread ID.
+            kwargs: not used at present.
+        """
+        if service_thread_id is not None:
+            return AgentThread(
+                service_thread_id=service_thread_id,
+                context_provider=self.context_provider,
+            )
+        if self.chat_message_store_factory is not None:
+            return AgentThread(
+                message_store=self.chat_message_store_factory(),
+                context_provider=self.context_provider,
+            )
+        return AgentThread(context_provider=self.context_provider)
+
+    async def _update_thread_with_type_and_conversation_id(
+        self, thread: AgentThread, response_conversation_id: str | None
+    ) -> None:
+        """Update thread with storage type and conversation ID.
+
+        Args:
+            thread: The thread to update.
+            response_conversation_id: The conversation ID from the response, if any.
+
+        Raises:
+            AgentExecutionException: If conversation ID is missing for service-managed thread.
+        """
+        if response_conversation_id is None and thread.service_thread_id is not None:
+            # We were passed a thread that is service managed, but we got no conversation id back from the chat client,
+            # meaning the service doesn't support service managed threads,
+            # so the thread cannot be used with this service.
+            raise AgentExecutionException(
+                "Service did not return a valid conversation id when using a service managed thread."
+            )
+
+        if response_conversation_id is not None:
+            # If we got a conversation id back from the chat client, it means that the service
+            # supports server side thread storage so we should update the thread with the new id.
+            thread.service_thread_id = response_conversation_id
+            if thread.context_provider:
+                await thread.context_provider.thread_created(thread.service_thread_id)
+        elif thread.message_store is None and self.chat_message_store_factory is not None:
+            # If the service doesn't use service side thread storage (i.e. we got no id back from invocation), and
+            # the thread has no message_store yet, and we have a custom messages store, we should update the thread
+            # with the custom message_store so that it has somewhere to store the chat history.
+            thread.message_store = self.chat_message_store_factory()
+
+    async def _prepare_thread_and_messages(
+        self,
+        *,
+        thread: AgentThread | None,
+        input_messages: list[ChatMessage] | None = None,
+    ) -> tuple[AgentThread, ChatOptions, list[ChatMessage]]:
+        """Prepare the messages for agent execution.
+
+        Also updates the chat_options of the agent, with
+
+        Args:
+            thread: The conversation thread.
+            input_messages: Messages to process.
+
+        Returns:
+            The validated thread and normalized messages.
+
+        Raises:
+            AgentExecutionException: If the thread is not of the expected type.
+        """
+        chat_options = copy(self.chat_options) if self.chat_options else ChatOptions()
+        thread = thread or self.get_new_thread()
+        if thread.service_thread_id and thread.context_provider:
+            await thread.context_provider.thread_created(thread.service_thread_id)
+        thread_messages: list[ChatMessage] = []
+        if thread.message_store:
+            thread_messages.extend(await thread.message_store.list_messages() or [])
+        context: Context | None = None
+        if self.context_provider:
+            async with self.context_provider:
+                context = await self.context_provider.invoking(input_messages or [])
+                if context:
+                    if context.messages:
+                        thread_messages.extend(context.messages)
+                    if context.tools:
+                        if chat_options.tools is not None:
+                            chat_options.tools.extend(context.tools)
+                        else:
+                            chat_options.tools = list(context.tools)
+                    if context.instructions:
+                        chat_options.instructions = (
+                            context.instructions
+                            if not chat_options.instructions
+                            else f"{chat_options.instructions}\n{context.instructions}"
+                        )
+        thread_messages.extend(input_messages or [])
+        if (
+            thread.service_thread_id
+            and chat_options.conversation_id
+            and thread.service_thread_id != chat_options.conversation_id
+        ):
+            raise AgentExecutionException(
+                "The conversation_id set on the agent is different from the one set on the thread, "
+                "only one ID can be used for a run."
+            )
+        return thread, chat_options, thread_messages
+
+    def _get_agent_name(self) -> str:
+        return self.name or "UnnamedAgent"

python/packages/core/agent_framework/_clients.py

@@ -0,0 +1,528 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+from abc import ABC, abstractmethod
+from collections.abc import AsyncIterable, Callable, MutableMapping, MutableSequence, Sequence
+from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeVar, runtime_checkable
+
+from pydantic import BaseModel, Field
+
+from ._logging import get_logger
+from ._mcp import MCPTool
+from ._memory import AggregateContextProvider, ContextProvider
+from ._middleware import (
+    ChatMiddleware,
+    ChatMiddlewareCallable,
+    FunctionMiddleware,
+    FunctionMiddlewareCallable,
+    Middleware,
+)
+from ._pydantic import AFBaseModel
+from ._threads import ChatMessageStoreProtocol
+from ._tools import ToolProtocol
+from ._types import (
+    ChatMessage,
+    ChatOptions,
+    ChatResponse,
+    ChatResponseUpdate,
+    ToolMode,
+)
+
+if TYPE_CHECKING:
+    from ._agents import ChatAgent
+
+
+TInput = TypeVar("TInput", contravariant=True)
+TEmbedding = TypeVar("TEmbedding")
+TBaseChatClient = TypeVar("TBaseChatClient", bound="BaseChatClient")
+
+logger = get_logger()
+
+__all__ = [
+    "BaseChatClient",
+    "ChatClientProtocol",
+]
+
+
+# region ChatClientProtocol Protocol
+
+
+@runtime_checkable
+class ChatClientProtocol(Protocol):
+    """A protocol for a chat client that can generate responses."""
+
+    @property
+    def additional_properties(self) -> dict[str, Any]:
+        """Get additional properties associated with the client."""
+        ...
+
+    async def get_response(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage],
+        *,
+        frequency_penalty: float | None = None,
+        logit_bias: dict[str | int, float] | None = None,
+        max_tokens: int | None = None,
+        metadata: dict[str, Any] | None = None,
+        model: str | None = None,
+        presence_penalty: float | None = None,
+        response_format: type[BaseModel] | None = None,
+        seed: int | None = None,
+        stop: str | Sequence[str] | None = None,
+        store: bool | None = None,
+        temperature: float | None = None,
+        tool_choice: ToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = "auto",
+        tools: ToolProtocol
+        | Callable[..., Any]
+        | MutableMapping[str, Any]
+        | list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
+        | None = None,
+        top_p: float | None = None,
+        user: str | None = None,
+        additional_properties: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> ChatResponse:
+        """Sends input and returns the response.
+
+        Args:
+            messages: The sequence of input messages to send.
+            response_format: the format of the response.
+            frequency_penalty: the frequency penalty to use.
+            logit_bias: the logit bias to use.
+            max_tokens: The maximum number of tokens to generate.
+            metadata: additional metadata to include in the request.
+            model: The model to use for the agent.
+            presence_penalty: the presence penalty to use.
+            seed: the random seed to use.
+            stop: the stop sequence(s) for the request.
+            store: whether to store the response.
+            temperature: the sampling temperature to use.
+            tool_choice: the tool choice for the request.
+            tools: the tools to use for the request.
+            top_p: the nucleus sampling probability to use.
+            user: the user to associate with the request.
+            additional_properties: additional properties to include in the request
+            kwargs: any additional keyword arguments,
+                will only be passed to functions that are called.
+
+        Returns:
+            The response messages generated by the client.
+
+        Raises:
+            ValueError: If the input message sequence is `None`.
+        """
+        ...
+
+    def get_streaming_response(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage],
+        *,
+        frequency_penalty: float | None = None,
+        logit_bias: dict[str | int, float] | None = None,
+        max_tokens: int | None = None,
+        metadata: dict[str, Any] | None = None,
+        model: str | None = None,
+        presence_penalty: float | None = None,
+        response_format: type[BaseModel] | None = None,
+        seed: int | None = None,
+        stop: str | Sequence[str] | None = None,
+        store: bool | None = None,
+        temperature: float | None = None,
+        tool_choice: ToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = "auto",
+        tools: ToolProtocol
+        | Callable[..., Any]
+        | MutableMapping[str, Any]
+        | list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
+        | None = None,
+        top_p: float | None = None,
+        user: str | None = None,
+        additional_properties: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterable[ChatResponseUpdate]:
+        """Sends input messages and streams the response.
+
+        Args:
+            messages: The sequence of input messages to send.
+            frequency_penalty: the frequency penalty to use.
+            logit_bias: the logit bias to use.
+            max_tokens: The maximum number of tokens to generate.
+            metadata: additional metadata to include in the request.
+            model: The model to use for the agent.
+            presence_penalty: the presence penalty to use.
+            response_format: the format of the response.
+            seed: the random seed to use.
+            stop: the stop sequence(s) for the request.
+            store: whether to store the response.
+            temperature: the sampling temperature to use.
+            tool_choice: the tool choice for the request.
+            tools: the tools to use for the request.
+            top_p: the nucleus sampling probability to use.
+            user: the user to associate with the request.
+            additional_properties: additional properties to include in the request
+            kwargs: any additional keyword arguments,
+                will only be passed to functions that are called.
+
+        Yields:
+            An async iterable of chat response updates containing the content of the response messages
+            generated by the client.
+
+        Raises:
+            ValueError: If the input message sequence is `None`.
+        """
+        ...
+
+
+# region ChatClientBase
+
+
+def prepare_messages(messages: str | ChatMessage | list[str] | list[ChatMessage]) -> list[ChatMessage]:
+    """Turn the allowed input into a list of chat messages."""
+    if isinstance(messages, str):
+        return [ChatMessage(role="user", text=messages)]
+    if isinstance(messages, ChatMessage):
+        return [messages]
+    return_messages: list[ChatMessage] = []
+    for msg in messages:
+        if isinstance(msg, str):
+            msg = ChatMessage(role="user", text=msg)
+        return_messages.append(msg)
+    return return_messages
+
+
+class BaseChatClient(AFBaseModel, ABC):
+    """Base class for chat clients."""
+
+    additional_properties: dict[str, Any] = Field(default_factory=dict)
+    middleware: (
+        ChatMiddleware
+        | ChatMiddlewareCallable
+        | FunctionMiddleware
+        | FunctionMiddlewareCallable
+        | list[ChatMiddleware | ChatMiddlewareCallable | FunctionMiddleware | FunctionMiddlewareCallable]
+        | None
+    ) = None
+    OTEL_PROVIDER_NAME: str = "unknown"
+    # This is used for OTel setup, should be overridden in subclasses
+
+    def prepare_messages(
+        self, messages: str | ChatMessage | list[str] | list[ChatMessage], chat_options: ChatOptions
+    ) -> MutableSequence[ChatMessage]:
+        """Turn the allowed input into a list of chat messages."""
+        if chat_options.instructions:
+            system_msg = ChatMessage(role="system", text=chat_options.instructions)
+            return [system_msg, *prepare_messages(messages)]
+        return prepare_messages(messages)
+
+    def _filter_internal_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]:
+        """Filter out internal framework parameters that shouldn't be passed to chat client implementations.
+
+        Args:
+            kwargs: The original kwargs dictionary.
+
+        Returns:
+            A filtered kwargs dictionary without internal parameters.
+        """
+        return {k: v for k, v in kwargs.items() if not k.startswith("_")}
+
+    @staticmethod
+    def _normalize_tools(
+        tools: ToolProtocol
+        | MutableMapping[str, Any]
+        | Callable[..., Any]
+        | list[ToolProtocol | MutableMapping[str, Any] | Callable[..., Any]]
+        | None = None,
+    ) -> list[ToolProtocol | dict[str, Any] | Callable[..., Any]]:
+        """Normalize the tools input to a list of tools."""
+        final_tools: list[ToolProtocol | dict[str, Any] | Callable[..., Any]] = []
+        if not tools:
+            return final_tools
+        for tool in tools if isinstance(tools, list) else [tools]:  # type: ignore[reportUnknownType]
+            if isinstance(tool, MCPTool):
+                final_tools.extend(tool.functions)  # type: ignore
+                continue
+            final_tools.append(tool)  # type: ignore
+        return final_tools
+
+    # region Internal methods to be implemented by the derived classes
+
+    @abstractmethod
+    async def _inner_get_response(
+        self,
+        *,
+        messages: MutableSequence[ChatMessage],
+        chat_options: ChatOptions,
+        **kwargs: Any,
+    ) -> ChatResponse:
+        """Send a chat request to the AI service.
+
+        Args:
+            messages: The chat messages to send.
+            chat_options: The options for the request.
+            kwargs: Any additional keyword arguments.
+
+        Returns:
+            The chat response contents representing the response(s).
+        """
+
+    @abstractmethod
+    async def _inner_get_streaming_response(
+        self,
+        *,
+        messages: MutableSequence[ChatMessage],
+        chat_options: ChatOptions,
+        **kwargs: Any,
+    ) -> AsyncIterable[ChatResponseUpdate]:
+        """Send a streaming chat request to the AI service.
+
+        Args:
+            messages: The chat messages to send.
+            chat_options: The chat_options for the request.
+            kwargs: Any additional keyword arguments.
+
+        Yields:
+            ChatResponseUpdate: The streaming chat message contents.
+        """
+        # Below is needed for mypy: https://mypy.readthedocs.io/en/stable/more_types.html#asynchronous-iterators
+        if False:
+            yield
+        await asyncio.sleep(0)  # pragma: no cover
+        # This is a no-op, but it allows the method to be async and return an AsyncIterable.
+        # The actual implementation should yield ChatResponseUpdate instances as needed.
+
+    # endregion
+
+    # region Public method
+
+    async def get_response(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage],
+        *,
+        frequency_penalty: float | None = None,
+        logit_bias: dict[str | int, float] | None = None,
+        max_tokens: int | None = None,
+        metadata: dict[str, Any] | None = None,
+        model: str | None = None,
+        presence_penalty: float | None = None,
+        response_format: type[BaseModel] | None = None,
+        seed: int | None = None,
+        stop: str | Sequence[str] | None = None,
+        store: bool | None = None,
+        temperature: float | None = None,
+        tool_choice: ToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = "auto",
+        tools: ToolProtocol
+        | Callable[..., Any]
+        | MutableMapping[str, Any]
+        | list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
+        | None = None,
+        top_p: float | None = None,
+        user: str | None = None,
+        additional_properties: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> ChatResponse:
+        """Get a response from a chat client.
+
+        Args:
+            messages: the message or messages to send to the model
+            frequency_penalty: the frequency penalty to use.
+            logit_bias: the logit bias to use.
+            max_tokens: The maximum number of tokens to generate.
+            metadata: additional metadata to include in the request.
+            model: The model to use for the agent.
+            presence_penalty: the presence penalty to use.
+            response_format: the format of the response.
+            seed: the random seed to use.
+            stop: the stop sequence(s) for the request.
+            store: whether to store the response.
+            temperature: the sampling temperature to use.
+            tool_choice: the tool choice for the request.
+            tools: the tools to use for the request.
+            top_p: the nucleus sampling probability to use.
+            user: the user to associate with the request.
+            additional_properties: additional properties to include in the request.
+            kwargs: any additional keyword arguments,
+                will only be passed to functions that are called.
+
+        Returns:
+            A chat response from the model.
+        """
+        # Should we merge chat options instead of ignoring the input params?
+        if "chat_options" in kwargs:
+            chat_options = kwargs.pop("chat_options")
+            if not isinstance(chat_options, ChatOptions):
+                raise TypeError("chat_options must be an instance of ChatOptions")
+        else:
+            chat_options = ChatOptions(
+                model_id=model,
+                frequency_penalty=frequency_penalty,
+                logit_bias=logit_bias,
+                max_tokens=max_tokens,
+                metadata=metadata,
+                presence_penalty=presence_penalty,
+                response_format=response_format,
+                seed=seed,
+                stop=stop,
+                store=store,
+                temperature=temperature,
+                top_p=top_p,
+                tool_choice=tool_choice,
+                tools=self._normalize_tools(tools),  # type: ignore
+                user=user,
+                additional_properties=additional_properties or {},
+            )
+        prepped_messages = self.prepare_messages(messages, chat_options)
+        self._prepare_tool_choice(chat_options=chat_options)
+
+        filtered_kwargs = self._filter_internal_kwargs(kwargs)
+        return await self._inner_get_response(messages=prepped_messages, chat_options=chat_options, **filtered_kwargs)
+
+    async def get_streaming_response(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage],
+        *,
+        frequency_penalty: float | None = None,
+        logit_bias: dict[str | int, float] | None = None,
+        max_tokens: int | None = None,
+        metadata: dict[str, Any] | None = None,
+        model: str | None = None,
+        presence_penalty: float | None = None,
+        response_format: type[BaseModel] | None = None,
+        seed: int | None = None,
+        stop: str | Sequence[str] | None = None,
+        store: bool | None = None,
+        temperature: float | None = None,
+        tool_choice: ToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = "auto",
+        tools: ToolProtocol
+        | Callable[..., Any]
+        | MutableMapping[str, Any]
+        | list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
+        | None = None,
+        top_p: float | None = None,
+        user: str | None = None,
+        additional_properties: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterable[ChatResponseUpdate]:
+        """Get a streaming response from a chat client.
+
+        Args:
+            messages: the message or messages to send to the model
+            frequency_penalty: the frequency penalty to use
+            logit_bias: the logit bias to use
+            max_tokens: The maximum number of tokens to generate.
+            metadata: additional metadata to include in the request.
+            model: The model to use for the agent.
+            presence_penalty: the presence penalty to use.
+            response_format: the format of the response.
+            seed: the random seed to use.
+            stop: the stop sequence(s) for the request.
+            store: whether to store the response.
+            temperature: the sampling temperature to use.
+            tool_choice: the tool choice for the request.
+            tools: the tools to use for the request.
+            top_p: the nucleus sampling probability to use.
+            user: the user to associate with the request.
+            additional_properties: additional properties to include in the request
+            kwargs: any additional keyword arguments
+
+        Yields:
+            A stream representing the response(s) from the LLM.
+        """
+        # Should we merge chat options instead of ignoring the input params?
+        if "chat_options" in kwargs:
+            chat_options = kwargs.pop("chat_options")
+            if not isinstance(chat_options, ChatOptions):
+                raise TypeError("chat_options must be an instance of ChatOptions")
+        else:
+            chat_options = ChatOptions(
+                model_id=model,
+                frequency_penalty=frequency_penalty,
+                logit_bias=logit_bias,
+                max_tokens=max_tokens,
+                metadata=metadata,
+                presence_penalty=presence_penalty,
+                response_format=response_format,
+                seed=seed,
+                stop=stop,
+                store=store,
+                temperature=temperature,
+                top_p=top_p,
+                tool_choice=tool_choice,
+                tools=self._normalize_tools(tools),
+                user=user,
+                additional_properties=additional_properties or {},
+            )
+        prepped_messages = self.prepare_messages(messages, chat_options)
+        self._prepare_tool_choice(chat_options=chat_options)
+
+        filtered_kwargs = self._filter_internal_kwargs(kwargs)
+        async for update in self._inner_get_streaming_response(
+            messages=prepped_messages, chat_options=chat_options, **filtered_kwargs
+        ):
+            yield update
+
+    def _prepare_tool_choice(self, chat_options: ChatOptions) -> None:
+        """Prepare the tools and tool choice for the chat options.
+
+        This function should be overridden by subclasses to customize tool handling.
+        Because it currently parses only AIFunctions.
+        """
+        chat_tool_mode = chat_options.tool_choice
+        if chat_tool_mode is None or chat_tool_mode == ToolMode.NONE or chat_tool_mode == "none":
+            chat_options.tools = None
+            chat_options.tool_choice = ToolMode.NONE.mode
+            return
+        if not chat_options.tools:
+            chat_options.tool_choice = ToolMode.NONE.mode
+        else:
+            chat_options.tool_choice = chat_tool_mode.mode if isinstance(chat_tool_mode, ToolMode) else chat_tool_mode
+
+    def service_url(self) -> str:
+        """Get the URL of the service.
+
+        Override this in the subclass to return the proper URL.
+        If the service does not have a URL, return None.
+        """
+        return "Unknown"
+
+    def create_agent(
+        self,
+        *,
+        name: str | None = None,
+        instructions: str | None = None,
+        tools: ToolProtocol
+        | Callable[..., Any]
+        | MutableMapping[str, Any]
+        | list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
+        | None = None,
+        chat_message_store_factory: Callable[[], ChatMessageStoreProtocol] | None = None,
+        context_providers: ContextProvider | list[ContextProvider] | AggregateContextProvider | None = None,
+        middleware: Middleware | list[Middleware] | None = None,
+        **kwargs: Any,
+    ) -> "ChatAgent":
+        """Create an agent with the given name and instructions.
+
+        Args:
+            name: The name of the agent.
+            instructions: The instructions for the agent.
+            tools: Optional list of tools to associate with the agent.
+            chat_message_store_factory: Factory function to create an instance of ChatMessageStoreProtocol.
+                If not provided, the default in-memory store will be used.
+            context_providers: Context providers to include during agent invocation.
+            middleware: List of middleware to intercept agent and function invocations.
+            **kwargs: Additional keyword arguments to pass to the agent.
+                See ChatAgent for all the available options.
+
+        Returns:
+            An instance of ChatAgent.
+        """
+        from ._agents import ChatAgent
+
+        return ChatAgent(
+            chat_client=self,
+            name=name,
+            instructions=instructions,
+            tools=tools,
+            chat_message_store_factory=chat_message_store_factory,
+            context_providers=context_providers,
+            middleware=middleware,
+            **kwargs,
+        )

python/packages/core/agent_framework/_logging.py

@@ -0,0 +1,26 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import logging
+
+from .exceptions import AgentFrameworkException
+
+logging.basicConfig(
+    format="[%(asctime)s - %(pathname)s:%(lineno)d - %(levelname)s] %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+)
+
+__all__ = ["get_logger"]
+
+
+def get_logger(name: str = "agent_framework") -> logging.Logger:
+    """Get a logger with the specified name, defaulting to 'agent_framework'.
+
+    Args:
+        name (str): The name of the logger. Defaults to 'agent_framework'.
+
+    Returns:
+        logging.Logger: The configured logger instance.
+    """
+    if not name.startswith("agent_framework"):
+        raise AgentFrameworkException("Logger name must start with 'agent_framework'.")
+    return logging.getLogger(name)

python/packages/core/agent_framework/_mcp.py

@@ -0,0 +1,709 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import json
+import logging
+import re
+import sys
+from abc import abstractmethod
+from contextlib import AsyncExitStack, _AsyncGeneratorContextManager  # type: ignore
+from datetime import timedelta
+from functools import partial
+from typing import TYPE_CHECKING, Any
+
+from mcp import types
+from mcp.client.session import ClientSession
+from mcp.client.stdio import StdioServerParameters, stdio_client
+from mcp.client.streamable_http import streamablehttp_client
+from mcp.client.websocket import websocket_client
+from mcp.shared.context import RequestContext
+from mcp.shared.exceptions import McpError
+from mcp.shared.session import RequestResponder
+from pydantic import BaseModel, create_model
+
+from ._tools import AIFunction
+from ._types import ChatMessage, Contents, DataContent, Role, TextContent, UriContent
+from .exceptions import ToolException, ToolExecutionException
+
+if sys.version_info >= (3, 11):
+    from typing import Self  # pragma: no cover
+else:
+    from typing_extensions import Self  # pragma: no cover
+
+if TYPE_CHECKING:
+    from ._clients import ChatClientProtocol
+
+logger = logging.getLogger(__name__)
+
+# region: Helpers
+
+LOG_LEVEL_MAPPING: dict[types.LoggingLevel, int] = {
+    "debug": logging.DEBUG,
+    "info": logging.INFO,
+    "notice": logging.INFO,
+    "warning": logging.WARNING,
+    "error": logging.ERROR,
+    "critical": logging.CRITICAL,
+    "alert": logging.CRITICAL,
+    "emergency": logging.CRITICAL,
+}
+
+__all__ = [
+    "MCPStdioTool",
+    "MCPStreamableHTTPTool",
+    "MCPWebsocketTool",
+]
+
+
+def _mcp_prompt_message_to_chat_message(
+    mcp_type: types.PromptMessage | types.SamplingMessage,
+) -> ChatMessage:
+    """Convert a MCP container type to a Agent Framework type."""
+    return ChatMessage(
+        role=Role(value=mcp_type.role),
+        contents=[_mcp_type_to_ai_content(mcp_type.content)],  # type: ignore[call-arg]
+        raw_representation=mcp_type,
+    )
+
+
+def _mcp_call_tool_result_to_ai_contents(
+    mcp_type: types.CallToolResult,
+) -> list[Contents]:
+    """Convert a MCP container type to a Agent Framework type."""
+    return [_mcp_type_to_ai_content(item) for item in mcp_type.content]
+
+
+def _mcp_type_to_ai_content(
+    mcp_type: types.ImageContent | types.TextContent | types.AudioContent | types.EmbeddedResource | types.ResourceLink,
+) -> Contents:
+    """Convert a MCP type to a Agent Framework type."""
+    match mcp_type:
+        case types.TextContent():
+            return TextContent(text=mcp_type.text, raw_representation=mcp_type)
+        case types.ImageContent() | types.AudioContent():
+            return DataContent(uri=mcp_type.data, media_type=mcp_type.mimeType, raw_representation=mcp_type)
+        case types.ResourceLink():
+            return UriContent(
+                uri=str(mcp_type.uri), media_type=mcp_type.mimeType or "application/json", raw_representation=mcp_type
+            )
+        case _:
+            match mcp_type.resource:
+                case types.TextResourceContents():
+                    return TextContent(
+                        text=mcp_type.resource.text,
+                        raw_representation=mcp_type,
+                        additional_properties=mcp_type.annotations.model_dump() if mcp_type.annotations else None,
+                    )
+                case types.BlobResourceContents():
+                    return DataContent(
+                        uri=mcp_type.resource.blob,
+                        media_type=mcp_type.resource.mimeType,
+                        raw_representation=mcp_type,
+                        additional_properties=mcp_type.annotations.model_dump() if mcp_type.annotations else None,
+                    )
+
+
+def _ai_content_to_mcp_types(
+    content: Contents,
+) -> types.TextContent | types.ImageContent | types.AudioContent | types.EmbeddedResource | types.ResourceLink | None:
+    """Convert a BaseContent type to a MCP type."""
+    match content:
+        case TextContent():
+            return types.TextContent(type="text", text=content.text)
+        case DataContent():
+            if content.media_type and content.media_type.startswith("image/"):
+                return types.ImageContent(type="image", data=content.uri, mimeType=content.media_type)
+            if content.media_type and content.media_type.startswith("audio/"):
+                return types.AudioContent(type="audio", data=content.uri, mimeType=content.media_type)
+            if content.media_type and content.media_type.startswith("application/"):
+                return types.EmbeddedResource(
+                    type="resource",
+                    resource=types.BlobResourceContents(
+                        blob=content.uri,
+                        mimeType=content.media_type,
+                        # uri's are not limited in MCP but they have to be set.
+                        # the uri of data content, contains the data uri, which
+                        # is not the uri meant here, UriContent would match this.
+                        uri=content.additional_properties.get("uri", "af://binary")
+                        if content.additional_properties
+                        else "af://binary",  # type: ignore[reportArgumentType]
+                    ),
+                )
+            return None
+        case UriContent():
+            return types.ResourceLink(
+                type="resource_link",
+                uri=content.uri,  # type: ignore[reportArgumentType]
+                mimeType=content.media_type,
+                name=content.additional_properties.get("name", "Unknown")
+                if content.additional_properties
+                else "Unknown",
+            )
+        case _:
+            return None
+
+
+def _chat_message_to_mcp_types(
+    content: ChatMessage,
+) -> list[types.TextContent | types.ImageContent | types.AudioContent | types.EmbeddedResource | types.ResourceLink]:
+    """Convert a ChatMessage to a list of MCP types."""
+    messages: list[
+        types.TextContent | types.ImageContent | types.AudioContent | types.EmbeddedResource | types.ResourceLink
+    ] = []
+    for item in content.contents:
+        mcp_content = _ai_content_to_mcp_types(item)
+        if mcp_content:
+            messages.append(mcp_content)
+    return messages
+
+
+def _get_input_model_from_mcp_prompt(prompt: types.Prompt) -> type[BaseModel]:
+    """Creates a Pydantic model from a prompt's parameters."""
+    # Check if 'arguments' is missing or empty
+    if not prompt.arguments:
+        return create_model(f"{prompt.name}_input")
+
+    field_definitions: dict[str, Any] = {}
+    for prompt_argument in prompt.arguments:
+        # For prompts, all arguments are typically required and string type
+        # unless specified otherwise in the prompt argument
+        python_type = str  # Default type for prompt arguments
+
+        # Create field definition for create_model
+        if prompt_argument.required:
+            field_definitions[prompt_argument.name] = (python_type, ...)
+        else:
+            field_definitions[prompt_argument.name] = (python_type, None)
+
+    return create_model(f"{prompt.name}_input", **field_definitions)
+
+
+def _get_input_model_from_mcp_tool(tool: types.Tool) -> type[BaseModel]:
+    """Creates a Pydantic model from a tools parameters."""
+    properties = tool.inputSchema.get("properties", None)
+    required = tool.inputSchema.get("required", [])
+    # Check if 'properties' is missing or not a dictionary
+    if not properties:
+        return create_model(f"{tool.name}_input")
+
+    field_definitions: dict[str, Any] = {}
+    for prop_name, prop_details in properties.items():
+        prop_details = json.loads(prop_details) if isinstance(prop_details, str) else prop_details
+
+        # Map JSON Schema types to Python types
+        json_type = prop_details.get("type", "string")
+        python_type: type = str  # default
+        if json_type == "integer":
+            python_type = int
+        elif json_type == "number":
+            python_type = float
+        elif json_type == "boolean":
+            python_type = bool
+        elif json_type == "array":
+            python_type = list
+        elif json_type == "object":
+            python_type = dict
+
+        # Create field definition for create_model
+        if prop_name in required:
+            field_definitions[prop_name] = (python_type, ...)
+        else:
+            default_value = prop_details.get("default", None)
+            field_definitions[prop_name] = (python_type, default_value)
+
+    return create_model(f"{tool.name}_input", **field_definitions)
+
+
+def _normalize_mcp_name(name: str) -> str:
+    """Normalize MCP tool/prompt names to allowed identifier pattern (A-Za-z0-9_.-)."""
+    return re.sub(r"[^A-Za-z0-9_.-]", "-", name)
+
+
+# region: MCP Plugin
+
+
+class MCPTool:
+    """Main MCP class, to initialize use one of the subclasses."""
+
+    def __init__(
+        self,
+        name: str,
+        description: str | None = None,
+        additional_properties: dict[str, Any] | None = None,
+        load_tools: bool = True,
+        load_prompts: bool = True,
+        session: ClientSession | None = None,
+        request_timeout: int | None = None,
+        chat_client: "ChatClientProtocol | None" = None,
+    ) -> None:
+        """Initialize the MCP Plugin Base."""
+        self.name = name
+        self.description = description or ""
+        self.additional_properties = additional_properties
+        self.load_tools_flag = load_tools
+        self.load_prompts_flag = load_prompts
+        self._exit_stack = AsyncExitStack()
+        self.session = session
+        self.request_timeout = request_timeout
+        self.chat_client = chat_client
+        self.functions: list[AIFunction[Any, Any]] = []
+        self.is_connected: bool = False
+
+    def __str__(self) -> str:
+        return f"MCPTool(name={self.name}, description={self.description})"
+
+    async def connect(self) -> None:
+        """Connect to the MCP server."""
+        if not self.session:
+            try:
+                transport = await self._exit_stack.enter_async_context(self.get_mcp_client())
+            except Exception as ex:
+                await self._exit_stack.aclose()
+                raise ToolException(
+                    "Failed to connect to the MCP server. Please check your configuration.", inner_exception=ex
+                ) from ex
+            try:
+                session = await self._exit_stack.enter_async_context(
+                    ClientSession(
+                        read_stream=transport[0],
+                        write_stream=transport[1],
+                        read_timeout_seconds=timedelta(seconds=self.request_timeout) if self.request_timeout else None,
+                        message_handler=self.message_handler,
+                        logging_callback=self.logging_callback,
+                        sampling_callback=self.sampling_callback,
+                    )
+                )
+            except Exception as ex:
+                await self._exit_stack.aclose()
+                raise ToolException(
+                    message="Failed to create a session. Please check your configuration.", inner_exception=ex
+                ) from ex
+            await session.initialize()
+            self.session = session
+        elif self.session._request_id == 0:  # type: ignore[reportPrivateUsage]
+            # If the session is not initialized, we need to reinitialize it
+            await self.session.initialize()
+        logger.debug("Connected to MCP server: %s", self.session)
+        self.is_connected = True
+        if self.load_tools_flag:
+            await self.load_tools()
+        if self.load_prompts_flag:
+            await self.load_prompts()
+
+        if logger.level != logging.NOTSET:
+            try:
+                await self.session.set_logging_level(
+                    next(level for level, value in LOG_LEVEL_MAPPING.items() if value == logger.level)
+                )
+            except Exception as exc:
+                logger.warning("Failed to set log level to %s", logger.level, exc_info=exc)
+
+    async def sampling_callback(
+        self, context: RequestContext[ClientSession, Any], params: types.CreateMessageRequestParams
+    ) -> types.CreateMessageResult | types.ErrorData:
+        """Callback function for sampling.
+
+        This function is called when the MCP server needs to get a message completed.
+
+        This is a simple version of this function, it can be overridden to allow more complex sampling.
+        It get's added to the session at initialization time, so overriding it is the best way to do this.
+        """
+        if not self.chat_client:
+            return types.ErrorData(
+                code=types.INTERNAL_ERROR,
+                message="No chat client available. Please set a chat client.",
+            )
+        logger.debug("Sampling callback called with params: %s", params)
+        messages: list[ChatMessage] = []
+        for msg in params.messages:
+            messages.append(_mcp_prompt_message_to_chat_message(msg))
+        try:
+            response = await self.chat_client.get_response(
+                messages,
+                temperature=params.temperature,
+                max_tokens=params.maxTokens,
+                stop=params.stopSequences,
+            )
+        except Exception as ex:
+            return types.ErrorData(
+                code=types.INTERNAL_ERROR,
+                message=f"Failed to get chat message content: {ex}",
+            )
+        if not response or not response.messages:
+            return types.ErrorData(
+                code=types.INTERNAL_ERROR,
+                message="Failed to get chat message content.",
+            )
+        mcp_contents = _chat_message_to_mcp_types(response.messages[0])
+        # grab the first content that is of type TextContent or ImageContent
+        mcp_content = next(
+            (content for content in mcp_contents if isinstance(content, (types.TextContent, types.ImageContent))),
+            None,
+        )
+        if not mcp_content:
+            return types.ErrorData(
+                code=types.INTERNAL_ERROR,
+                message="Failed to get right content types from the response.",
+            )
+        return types.CreateMessageResult(
+            role="assistant",
+            content=mcp_content,
+            model=response.model_id or "unknown",
+        )
+
+    async def logging_callback(self, params: types.LoggingMessageNotificationParams) -> None:
+        """Callback function for logging.
+
+        This function is called when the MCP Server sends a log message.
+        By default it will log the message to the logger with the level set in the params.
+
+        Please subclass the MCP*Plugin and override this function if you want to adapt the behavior.
+        """
+        logger.log(LOG_LEVEL_MAPPING[params.level], params.data)
+
+    async def message_handler(
+        self,
+        message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
+    ) -> None:
+        """Handle messages from the MCP server.
+
+        By default this function will handle exceptions on the server, by logging those.
+
+        And it will trigger a reload of the tools and prompts when the list changed notification is received.
+
+        If you want to extend this behavior you can subclass the MCPPlugin and override this function,
+        if you want to keep the default behavior, make sure to call `super().message_handler(message)`.
+        """
+        if isinstance(message, Exception):
+            logger.error("Error from MCP server: %s", message, exc_info=message)
+            return
+        if isinstance(message, types.ServerNotification):
+            match message.root.method:
+                case "notifications/tools/list_changed":
+                    await self.load_tools()
+                case "notifications/prompts/list_changed":
+                    await self.load_prompts()
+                case _:
+                    logger.debug("Unhandled notification: %s", message.root.method)
+
+    async def load_prompts(self) -> None:
+        """Load prompts from the MCP server."""
+        if not self.session:
+            raise ToolExecutionException("MCP server not connected, please call connect() before using this method.")
+        try:
+            prompt_list = await self.session.list_prompts()
+        except Exception as exc:
+            logger.info(
+                "Prompt could not be loaded, you can exclude trying to load, by setting: load_prompts=False",
+                exc_info=exc,
+            )
+            prompt_list = None
+        for prompt in prompt_list.prompts if prompt_list else []:
+            local_name = _normalize_mcp_name(prompt.name)
+            input_model = _get_input_model_from_mcp_prompt(prompt)
+            func: AIFunction[BaseModel, list[ChatMessage]] = AIFunction(
+                func=partial(self.get_prompt, prompt.name),
+                name=local_name,
+                description=prompt.description or "",
+                input_model=input_model,
+            )
+            self.functions.append(func)
+
+    async def load_tools(self) -> None:
+        """Load tools from the MCP server."""
+        if not self.session:
+            raise ToolExecutionException("MCP server not connected, please call connect() before using this method.")
+        try:
+            tool_list = await self.session.list_tools()
+        except Exception as exc:
+            logger.info(
+                "Tools could not be loaded, you can exclude trying to load, by setting: load_tools=False",
+                exc_info=exc,
+            )
+            tool_list = None
+        for tool in tool_list.tools if tool_list else []:
+            local_name = _normalize_mcp_name(tool.name)
+            input_model = _get_input_model_from_mcp_tool(tool)
+            # Create AIFunctions out of each tool
+            func: AIFunction[BaseModel, list[Contents]] = AIFunction(
+                func=partial(self.call_tool, tool.name),
+                name=local_name,
+                description=tool.description or "",
+                input_model=input_model,
+            )
+            self.functions.append(func)
+
+    async def close(self) -> None:
+        """Disconnect from the MCP server."""
+        await self._exit_stack.aclose()
+        self.session = None
+        self.is_connected = False
+
+    @abstractmethod
+    def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
+        """Get an MCP client."""
+        pass
+
+    async def call_tool(self, tool_name: str, **kwargs: Any) -> list[Contents]:
+        """Call a tool with the given arguments."""
+        if not self.session:
+            raise ToolExecutionException("MCP server not connected, please call connect() before using this method.")
+        if not self.load_tools_flag:
+            raise ToolExecutionException(
+                "Tools are not loaded for this server, please set load_tools=True in the constructor."
+            )
+        try:
+            return _mcp_call_tool_result_to_ai_contents(await self.session.call_tool(tool_name, arguments=kwargs))
+        except McpError as mcp_exc:
+            raise ToolExecutionException(mcp_exc.error.message, inner_exception=mcp_exc) from mcp_exc
+        except Exception as ex:
+            raise ToolExecutionException(f"Failed to call tool '{tool_name}'.", inner_exception=ex) from ex
+
+    async def get_prompt(self, prompt_name: str, **kwargs: Any) -> list[ChatMessage]:
+        """Call a prompt with the given arguments."""
+        if not self.session:
+            raise ToolExecutionException("MCP server not connected, please call connect() before using this method.")
+        if not self.load_prompts_flag:
+            raise ToolExecutionException(
+                "Prompts are not loaded for this server, please set load_prompts=True in the constructor."
+            )
+        try:
+            prompt_result = await self.session.get_prompt(prompt_name, arguments=kwargs)
+            return [_mcp_prompt_message_to_chat_message(message) for message in prompt_result.messages]
+        except McpError as mcp_exc:
+            raise ToolExecutionException(mcp_exc.error.message, inner_exception=mcp_exc) from mcp_exc
+        except Exception as ex:
+            raise ToolExecutionException(f"Failed to call prompt '{prompt_name}'.", inner_exception=ex) from ex
+
+    async def __aenter__(self) -> Self:
+        """Enter the context manager."""
+        try:
+            await self.connect()
+            return self
+        except ToolException:
+            raise
+        except Exception as ex:
+            await self._exit_stack.aclose()
+            raise ToolExecutionException("Failed to enter context manager.", inner_exception=ex) from ex
+
+    async def __aexit__(
+        self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: Any
+    ) -> None:
+        """Exit the context manager."""
+        await self.close()
+
+
+# region: MCP Plugin Implementations
+
+
+class MCPStdioTool(MCPTool):
+    """MCP stdio server configuration."""
+
+    def __init__(
+        self,
+        name: str,
+        command: str,
+        *,
+        load_tools: bool = True,
+        load_prompts: bool = True,
+        request_timeout: int | None = None,
+        session: ClientSession | None = None,
+        description: str | None = None,
+        additional_properties: dict[str, Any] | None = None,
+        args: list[str] | None = None,
+        env: dict[str, str] | None = None,
+        encoding: str | None = None,
+        chat_client: "ChatClientProtocol | None" = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize the MCP stdio plugin.
+
+        The arguments are used to create a StdioServerParameters object.
+        Which is then used to create a stdio client.
+        see mcp.client.stdio.stdio_client and mcp.client.stdio.stdio_server_parameters
+        for more details.
+
+        Args:
+            name: The name of the plugin.
+            command: The command to run the MCP server.
+            load_tools: Whether to load tools from the MCP server.
+            load_prompts: Whether to load prompts from the MCP server.
+            request_timeout: The default timeout used for all requests.
+            session: The session to use for the MCP connection.
+            description: The description of the plugin.
+            additional_properties: Additional properties.
+            args: The arguments to pass to the command.
+            env: The environment variables to set for the command.
+            encoding: The encoding to use for the command output.
+            chat_client: The chat client to use for sampling.
+            kwargs: Any extra arguments to pass to the stdio client.
+
+        """
+        super().__init__(
+            name=name,
+            description=description,
+            additional_properties=additional_properties,
+            session=session,
+            chat_client=chat_client,
+            load_tools=load_tools,
+            load_prompts=load_prompts,
+            request_timeout=request_timeout,
+        )
+        self.command = command
+        self.args = args or []
+        self.env = env
+        self.encoding = encoding
+        self._client_kwargs = kwargs
+
+    def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
+        """Get an MCP stdio client."""
+        args: dict[str, Any] = {
+            "command": self.command,
+            "args": self.args,
+            "env": self.env,
+        }
+        if self.encoding:
+            args["encoding"] = self.encoding
+        if self._client_kwargs:
+            args.update(self._client_kwargs)
+        return stdio_client(server=StdioServerParameters(**args))
+
+
+class MCPStreamableHTTPTool(MCPTool):
+    """MCP streamable http server configuration."""
+
+    def __init__(
+        self,
+        name: str,
+        url: str,
+        *,
+        load_tools: bool = True,
+        load_prompts: bool = True,
+        request_timeout: int | None = None,
+        session: ClientSession | None = None,
+        description: str | None = None,
+        additional_properties: dict[str, Any] | None = None,
+        headers: dict[str, Any] | None = None,
+        timeout: float | None = None,
+        sse_read_timeout: float | None = None,
+        terminate_on_close: bool | None = None,
+        chat_client: "ChatClientProtocol | None" = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize the MCP streamable http plugin.
+
+        The arguments are used to create a streamable http client.
+        see mcp.client.streamable_http.streamablehttp_client for more details.
+
+        Any extra arguments passed to the constructor will be passed to the
+        streamable http client constructor.
+
+        Args:
+            name: The name of the plugin.
+            url: The URL of the MCP server.
+            load_tools: Whether to load tools from the MCP server.
+            load_prompts: Whether to load prompts from the MCP server.
+            request_timeout: The default timeout used for all requests.
+            session: The session to use for the MCP connection.
+            description: The description of the plugin.
+            additional_properties: Additional properties.
+            headers: The headers to send with the request.
+            timeout: The timeout for the request.
+            sse_read_timeout: The timeout for reading from the SSE stream.
+            terminate_on_close: Close the transport when the MCP client is terminated.
+            chat_client: The chat client to use for sampling.
+            kwargs: Any extra arguments to pass to the sse client.
+        """
+        super().__init__(
+            name=name,
+            description=description,
+            additional_properties=additional_properties,
+            session=session,
+            chat_client=chat_client,
+            load_tools=load_tools,
+            load_prompts=load_prompts,
+            request_timeout=request_timeout,
+        )
+        self.url = url
+        self.headers = headers or {}
+        self.timeout = timeout
+        self.sse_read_timeout = sse_read_timeout
+        self.terminate_on_close = terminate_on_close
+        self._client_kwargs = kwargs
+
+    def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
+        """Get an MCP streamable http client."""
+        args: dict[str, Any] = {
+            "url": self.url,
+        }
+        if self.headers:
+            args["headers"] = self.headers
+        if self.timeout is not None:
+            args["timeout"] = self.timeout
+        if self.sse_read_timeout is not None:
+            args["sse_read_timeout"] = self.sse_read_timeout
+        if self.terminate_on_close is not None:
+            args["terminate_on_close"] = self.terminate_on_close
+        if self._client_kwargs:
+            args.update(self._client_kwargs)
+        return streamablehttp_client(**args)
+
+
+class MCPWebsocketTool(MCPTool):
+    """MCP websocket server configuration."""
+
+    def __init__(
+        self,
+        name: str,
+        url: str,
+        *,
+        load_tools: bool = True,
+        load_prompts: bool = True,
+        request_timeout: int | None = None,
+        session: ClientSession | None = None,
+        description: str | None = None,
+        additional_properties: dict[str, Any] | None = None,
+        chat_client: "ChatClientProtocol | None" = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize the MCP websocket plugin.
+
+                The arguments are used to create a websocket client.
+        see mcp.client.websocket.websocket_client for more details.
+
+        Any extra arguments passed to the constructor will be passed to the
+        websocket client constructor.
+
+        Args:
+            name: The name of the plugin.
+            url: The URL of the MCP server.
+            load_tools: Whether to load tools from the MCP server.
+            load_prompts: Whether to load prompts from the MCP server.
+            request_timeout: The default timeout used for all requests.
+            session: The session to use for the MCP connection.
+            description: The description of the plugin.
+            additional_properties: Additional properties.
+            chat_client: The chat client to use for sampling.
+            kwargs: Any extra arguments to pass to the websocket client.
+
+        """
+        super().__init__(
+            name=name,
+            description=description,
+            additional_properties=additional_properties,
+            session=session,
+            chat_client=chat_client,
+            load_tools=load_tools,
+            load_prompts=load_prompts,
+            request_timeout=request_timeout,
+        )
+        self.url = url
+        self._client_kwargs = kwargs
+
+    def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
+        """Get an MCP websocket client."""
+        args: dict[str, Any] = {
+            "url": self.url,
+        }
+        if self._client_kwargs:
+            args.update(self._client_kwargs)
+        return websocket_client(**args)

python/packages/core/agent_framework/_memory.py

@@ -0,0 +1,241 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import sys
+from abc import ABC, abstractmethod
+from collections.abc import MutableSequence, Sequence
+from contextlib import AsyncExitStack
+from types import TracebackType
+from typing import Any, Final, cast
+
+from ._tools import ToolProtocol
+from ._types import ChatMessage
+
+if sys.version_info >= (3, 12):
+    from typing import override  # type: ignore # pragma: no cover
+else:
+    from typing_extensions import override  # type: ignore[import] # pragma: no cover
+if sys.version_info >= (3, 11):
+    from typing import Self  # pragma: no cover
+else:
+    from typing_extensions import Self  # pragma: no cover
+
+# region Context
+
+__all__ = ["AggregateContextProvider", "Context", "ContextProvider"]
+
+
+class Context:
+    """A class containing any context that should be provided to the AI model as supplied by an ContextProvider.
+
+    Each ContextProvider has the ability to provide its own context for each invocation.
+    The Context class contains the additional context supplied by the ContextProvider.
+    This context will be combined with context supplied by other providers before being passed to the AI model.
+    This context is per invocation, and will not be stored as part of the chat history.
+    """
+
+    def __init__(
+        self,
+        instructions: str | None = None,
+        messages: Sequence[ChatMessage] | None = None,
+        tools: Sequence[ToolProtocol] | None = None,
+    ):
+        """Create a new Context object.
+
+        Args:
+            instructions: Instructions to provide to the AI model.
+            messages: a list of messages.
+            tools: a list of tools to provide to this run.
+        """
+        self.instructions = instructions
+        self.messages: Sequence[ChatMessage] = messages or []
+        self.tools: Sequence[ToolProtocol] = tools or []
+
+
+# region ContextProvider
+
+
+class ContextProvider(ABC):
+    """Base class for all context providers.
+
+    A context provider is a component that can be used to enhance the AI's context management.
+    It can listen to changes in the conversation and provide additional context to the AI model
+    just before invocation.
+
+    It also has a default memory prompt that can be used by all providers.
+    """
+
+    # Default prompt to be used by all context providers when assembling memories/instructions
+    DEFAULT_CONTEXT_PROMPT: Final[str] = "## Memories\nConsider the following memories when answering user questions:"
+
+    async def thread_created(self, thread_id: str | None) -> None:
+        """Called just after a new thread is created.
+
+        Implementers can use this method to do any operations required at the creation of a new thread.
+        For example, checking long term storage for any data that is relevant
+        to the current session based on the input text.
+
+        Args:
+            thread_id: The ID of the new thread.
+        """
+        pass
+
+    async def invoked(
+        self,
+        request_messages: ChatMessage | Sequence[ChatMessage],
+        response_messages: ChatMessage | Sequence[ChatMessage] | None = None,
+        invoke_exception: Exception | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called after the agent has received a response from the underlying inference service.
+
+        You can inspect the request and response messages, and update the state of the context provider
+
+        Args:
+            request_messages: messages that were sent to the model/agent
+            response_messages: messages that were returned by the model/agent
+            invoke_exception: exception that was thrown, if any.
+            kwargs: not used at present.
+        """
+        pass
+
+    @abstractmethod
+    async def invoking(self, messages: ChatMessage | MutableSequence[ChatMessage], **kwargs: Any) -> Context:
+        """Called just before the Model/Agent/etc. is invoked.
+
+        Implementers can load any additional context required at this time,
+        and they should return any context that should be passed to the agent.
+
+        Args:
+            messages: The most recent messages that the agent is being invoked with.
+            kwargs: not used at present.
+        """
+        pass
+
+    async def __aenter__(self) -> "Self":
+        """Async context manager entry.
+
+        Override this method to perform any setup operations when the context provider is entered.
+
+        Returns:
+            Self for chaining.
+        """
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Async context manager exit.
+
+        Override this method to perform any cleanup operations when the context provider is exited.
+
+        Args:
+            exc_type: Exception type if an exception occurred, None otherwise.
+            exc_val: Exception value if an exception occurred, None otherwise.
+            exc_tb: Exception traceback if an exception occurred, None otherwise.
+        """
+        pass
+
+
+# region AggregateContextProvider
+
+
+class AggregateContextProvider(ContextProvider):
+    """A ContextProvider that contains multiple context providers.
+
+    It delegates events to multiple context providers and aggregates responses from those events before returning.
+    """
+
+    def __init__(self, context_providers: ContextProvider | Sequence[ContextProvider] | None = None) -> None:
+        """Initialize the AggregateContextProvider with context providers.
+
+        Args:
+            context_providers: Context providers to add.
+        """
+        if isinstance(context_providers, ContextProvider):
+            self.providers = [context_providers]
+        else:
+            self.providers = cast(list[ContextProvider], context_providers) or []
+        self._exit_stack: AsyncExitStack | None = None
+
+    def add(self, context_provider: ContextProvider) -> None:
+        """Adds new context provider.
+
+        Args:
+            context_provider: Context provider to add.
+        """
+        self.providers.append(context_provider)
+
+    @override
+    async def thread_created(self, thread_id: str | None = None) -> None:
+        await asyncio.gather(*[x.thread_created(thread_id) for x in self.providers])
+
+    @override
+    async def invoking(self, messages: ChatMessage | MutableSequence[ChatMessage], **kwargs: Any) -> Context:
+        contexts = await asyncio.gather(*[provider.invoking(messages, **kwargs) for provider in self.providers])
+        instructions: str = ""
+        return_messages: list[ChatMessage] = []
+        tools: list[ToolProtocol] = []
+        for ctx in contexts:
+            if ctx.instructions:
+                instructions += ctx.instructions
+            if ctx.messages:
+                return_messages.extend(ctx.messages)
+            if ctx.tools:
+                tools.extend(ctx.tools)
+        return Context(instructions=instructions, messages=return_messages, tools=tools)
+
+    @override
+    async def invoked(
+        self,
+        request_messages: ChatMessage | Sequence[ChatMessage],
+        response_messages: ChatMessage | Sequence[ChatMessage] | None = None,
+        invoke_exception: Exception | None = None,
+        **kwargs: Any,
+    ) -> None:
+        await asyncio.gather(*[
+            x.invoked(
+                request_messages=request_messages,
+                response_messages=response_messages,
+                invoke_exception=invoke_exception,
+                **kwargs,
+            )
+            for x in self.providers
+        ])
+
+    @override
+    async def __aenter__(self) -> "Self":
+        """Enter async context manager and set up all providers.
+
+        Returns:
+            Self for chaining.
+        """
+        self._exit_stack = AsyncExitStack()
+        await self._exit_stack.__aenter__()
+
+        # Enter all context providers
+        for provider in self.providers:
+            await self._exit_stack.enter_async_context(provider)
+
+        return self
+
+    @override
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Exit async context manager and clean up all providers.
+
+        Args:
+            exc_type: Exception type if an exception occurred, None otherwise.
+            exc_val: Exception value if an exception occurred, None otherwise.
+            exc_tb: Exception traceback if an exception occurred, None otherwise.
+        """
+        if self._exit_stack is not None:
+            await self._exit_stack.__aexit__(exc_type, exc_val, exc_tb)
+            self._exit_stack = None

python/packages/core/agent_framework/_pydantic.py

@@ -0,0 +1,74 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+
+from typing import Annotated, Any, ClassVar, TypeVar
+
+from pydantic import BaseModel, ConfigDict, Field, UrlConstraints
+from pydantic.networks import AnyUrl
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+HTTPsUrl = Annotated[AnyUrl, UrlConstraints(max_length=2083, allowed_schemes=["https"])]
+
+__all__ = ["AFBaseModel", "AFBaseSettings", "HTTPsUrl"]
+
+
+class AFBaseModel(BaseModel):
+    """Base class for all pydantic models in the Agent Framework."""
+
+    model_config = ConfigDict(populate_by_name=True, arbitrary_types_allowed=True, validate_assignment=True)
+
+
+TSettings = TypeVar("TSettings", bound="AFBaseSettings")
+
+
+class AFBaseSettings(BaseSettings):
+    """Base class for all settings classes in the Agent Framework.
+
+    A subclass creates it's fields and overrides the env_prefix class variable
+    with the prefix for the environment variables.
+
+    In the case where a value is specified for the same Settings field in multiple ways,
+    the selected value is determined as follows (in descending order of priority):
+    - Arguments passed to the Settings class initializer.
+    - Environment variables, e.g. my_prefix_special_function as described above.
+    - Variables loaded from a dotenv (.env) file.
+    - Variables loaded from the secrets directory.
+    - The default field values for the Settings model.
+    """
+
+    env_prefix: ClassVar[str] = ""
+    env_file_path: str | None = Field(default=None, exclude=True)
+    env_file_encoding: str | None = Field(default="utf-8", exclude=True)
+
+    model_config = SettingsConfigDict(
+        extra="ignore",
+        case_sensitive=False,
+    )
+
+    def __init__(
+        self,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize the settings class."""
+        # Remove any None values from the kwargs so that defaults are used.
+        kwargs = {k: v for k, v in kwargs.items() if v is not None}
+        super().__init__(**kwargs)
+
+    def __new__(cls: type["TSettings"], *args: Any, **kwargs: Any) -> "TSettings":
+        """Override the __new__ method to set the env_prefix."""
+        # for both, if supplied but None, set to default
+        if "env_file_encoding" in kwargs and kwargs["env_file_encoding"] is not None:
+            env_file_encoding = kwargs["env_file_encoding"]
+        else:
+            env_file_encoding = "utf-8"
+        if "env_file_path" in kwargs and kwargs["env_file_path"] is not None:
+            env_file_path = kwargs["env_file_path"]
+        else:
+            env_file_path = ".env"
+        cls.model_config.update(  # type: ignore
+            env_prefix=cls.env_prefix,
+            env_file=env_file_path,
+            env_file_encoding=env_file_encoding,
+        )
+        cls.model_rebuild()
+        return super().__new__(cls)  # type: ignore[return-value]

python/packages/core/agent_framework/_telemetry.py

@@ -0,0 +1,59 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import os
+from typing import Any, Final
+
+from . import __version__ as version_info
+from ._logging import get_logger
+
+logger = get_logger()
+
+__all__ = [
+    "AGENT_FRAMEWORK_USER_AGENT",
+    "APP_INFO",
+    "USER_AGENT_KEY",
+    "USER_AGENT_TELEMETRY_DISABLED_ENV_VAR",
+    "prepend_agent_framework_to_user_agent",
+]
+
+# Note that if this environment variable does not exist, user agent telemetry is enabled.
+USER_AGENT_TELEMETRY_DISABLED_ENV_VAR = "AGENT_FRAMEWORK_USER_AGENT_DISABLED"
+IS_TELEMETRY_ENABLED = os.environ.get(USER_AGENT_TELEMETRY_DISABLED_ENV_VAR, "false").lower() not in ["true", "1"]
+
+APP_INFO = (
+    {
+        "agent-framework-version": f"python/{version_info}",  # type: ignore[has-type]
+    }
+    if IS_TELEMETRY_ENABLED
+    else None
+)
+USER_AGENT_KEY: Final[str] = "User-Agent"
+HTTP_USER_AGENT: Final[str] = "agent-framework-python"
+AGENT_FRAMEWORK_USER_AGENT = f"{HTTP_USER_AGENT}/{version_info}"  # type: ignore[has-type]
+
+
+def prepend_agent_framework_to_user_agent(headers: dict[str, Any] | None = None) -> dict[str, Any]:
+    """Prepend "agent-framework" to the User-Agent in the headers.
+
+    When user agent telemetry is disabled, through the AZURE_TELEMETRY_DISABLED environment variable,
+    the User-Agent header will not include the agent-framework information, it will be sent back as is,
+    or as a empty dict when None is passed.
+
+    Args:
+        headers: The existing headers dictionary.
+
+    Returns:
+        A new dict with "User-Agent" set to "agent-framework-python/{version}" if headers is None.
+        The modified headers dictionary with "agent-framework-python/{version}" prepended to the User-Agent.
+    """
+    if not IS_TELEMETRY_ENABLED:
+        return headers or {}
+    if not headers:
+        return {USER_AGENT_KEY: AGENT_FRAMEWORK_USER_AGENT}
+    headers[USER_AGENT_KEY] = (
+        f"{AGENT_FRAMEWORK_USER_AGENT} {headers[USER_AGENT_KEY]}"
+        if USER_AGENT_KEY in headers
+        else AGENT_FRAMEWORK_USER_AGENT
+    )
+
+    return headers

python/packages/core/agent_framework/_threads.py

@@ -0,0 +1,355 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from collections.abc import Sequence
+from typing import Any, Protocol, TypeVar
+
+from pydantic import model_validator
+
+from ._memory import AggregateContextProvider
+from ._pydantic import AFBaseModel
+from ._types import ChatMessage
+from .exceptions import AgentThreadException
+
+__all__ = ["AgentThread", "ChatMessageStore", "ChatMessageStoreProtocol"]
+
+
+class ChatMessageStoreProtocol(Protocol):
+    """Defines methods for storing and retrieving chat messages associated with a specific thread.
+
+    Implementations of this protocol are responsible for managing the storage of chat messages,
+    including handling large volumes of data by truncating or summarizing messages as necessary.
+    """
+
+    async def list_messages(self) -> list[ChatMessage]:
+        """Gets all the messages from the store that should be used for the next agent invocation.
+
+        Messages are returned in ascending chronological order, with the oldest message first.
+
+        If the messages stored in the store become very large, it is up to the store to
+        truncate, summarize or otherwise limit the number of messages returned.
+
+        When using implementations of ChatMessageStoreProtocol, a new one should be created for each thread
+        since they may contain state that is specific to a thread.
+        """
+        ...
+
+    async def add_messages(self, messages: Sequence[ChatMessage]) -> None:
+        """Adds messages to the store."""
+        ...
+
+    @classmethod
+    async def deserialize(cls, serialized_store_state: Any, **kwargs: Any) -> "ChatMessageStoreProtocol":
+        """Creates a new instance of the store from previously serialized state.
+
+        This method, together with serialize_state can be used to save and load messages from a persistent store
+        if this store only has messages in memory.
+        """
+        ...
+
+    async def update_from_state(self, serialized_store_state: Any, **kwargs: Any) -> None:
+        """Update the current ChatMessageStore instance from serialized state data.
+
+        Args:
+            serialized_store_state: Previously serialized state data containing messages.
+            **kwargs: Additional arguments for deserialization.
+        """
+        ...
+
+    async def serialize(self, **kwargs: Any) -> Any:
+        """Serializes the current object's state.
+
+        This method, together with deserialize can be used to save and load messages from a persistent store
+        if this store only has messages in memory.
+        """
+        ...
+
+
+class AgentThreadState(AFBaseModel):
+    """State model for serializing and deserializing thread information.
+
+    Attributes:
+        service_thread_id: Optional ID of the thread managed by the agent service.
+        chat_message_store_state: Optional serialized state of the chat message store.
+    """
+
+    service_thread_id: str | None = None
+    chat_message_store_state: Any | None = None
+
+    @model_validator(mode="before")
+    def validate_only_one(cls, values: dict[str, Any]) -> dict[str, Any]:
+        if (
+            isinstance(values, dict)
+            and values.get("service_thread_id") is not None
+            and values.get("chat_message_store_state") is not None
+        ):
+            raise AgentThreadException("Only one of service_thread_id or chat_message_store_state may be set.")
+        return values
+
+
+class ChatMessageStoreState(AFBaseModel):
+    """State model for serializing and deserializing chat message store data.
+
+    Attributes:
+        messages: List of chat messages stored in the message store.
+    """
+
+    messages: list[ChatMessage]
+
+
+TChatMessageStore = TypeVar("TChatMessageStore", bound="ChatMessageStore")
+
+
+class ChatMessageStore:
+    """An in-memory implementation of ChatMessageStoreProtocol that stores messages in a list.
+
+    This implementation provides a simple, list-based storage for chat messages
+    with support for serialization and deserialization. It implements all the
+    required methods of the ChatMessageStoreProtocol protocol.
+
+    The store maintains messages in memory and provides methods to serialize
+    and deserialize the state for persistence purposes.
+
+    Args:
+        messages: Optional initial list of ChatMessage objects to populate the store.
+    """
+
+    def __init__(self, messages: Sequence[ChatMessage] | None = None):
+        """Create a ChatMessageStore for use in a thread.
+
+        Args:
+            messages: The messages to store.
+        """
+        self.messages = list(messages) if messages else []
+
+    async def add_messages(self, messages: Sequence[ChatMessage]) -> None:
+        """Add messages to the store.
+
+        Args:
+            messages: Sequence of ChatMessage objects to add to the store.
+        """
+        self.messages.extend(messages)
+
+    async def list_messages(self) -> list[ChatMessage]:
+        """Get all messages from the store in chronological order.
+
+        Returns:
+            List of ChatMessage objects, ordered from oldest to newest.
+        """
+        return self.messages
+
+    @classmethod
+    async def deserialize(
+        cls: type[TChatMessageStore], serialized_store_state: Any, **kwargs: Any
+    ) -> TChatMessageStore:
+        """Create a new ChatMessageStore instance from serialized state data.
+
+        Args:
+            serialized_store_state: Previously serialized state data containing messages.
+            **kwargs: Additional arguments for deserialization.
+
+        Returns:
+            A new ChatMessageStore instance populated with messages from the serialized state.
+        """
+        state = ChatMessageStoreState.model_validate(serialized_store_state, **kwargs)
+        if state.messages:
+            return cls(messages=state.messages)
+        return cls()
+
+    async def update_from_state(self, serialized_store_state: Any, **kwargs: Any) -> None:
+        """Update the current ChatMessageStore instance from serialized state data.
+
+        Args:
+            serialized_store_state: Previously serialized state data containing messages.
+            **kwargs: Additional arguments for deserialization.
+        """
+        if not serialized_store_state:
+            return
+        state = ChatMessageStoreState.model_validate(serialized_store_state, **kwargs)
+        if state.messages:
+            self.messages = state.messages
+
+    async def serialize(self, **kwargs: Any) -> Any:
+        """Serialize the current store state for persistence.
+
+        Args:
+            **kwargs: Additional arguments for serialization.
+
+        Returns:
+            Serialized state data that can be used with deserialize_state.
+        """
+        state = ChatMessageStoreState(messages=self.messages)
+        return state.model_dump(**kwargs)
+
+
+TAgentThread = TypeVar("TAgentThread", bound="AgentThread")
+
+
+class AgentThread:
+    """The Agent thread class, this can represent both a locally managed thread or a thread managed by the service."""
+
+    def __init__(
+        self,
+        *,
+        service_thread_id: str | None = None,
+        message_store: ChatMessageStoreProtocol | None = None,
+        context_provider: AggregateContextProvider | None = None,
+    ) -> None:
+        """Initialize an AgentThread, do not use this method manually, always use: agent.get_new_thread().
+
+        Args:
+            service_thread_id: Optional ID of the thread managed by the agent service.
+            message_store: Optional ChatMessageStore implementation for managing chat messages.
+            context_provider: Optional ContextProvider for the thread.
+
+        Note:
+            Either service_thread_id or message_store may be set, but not both.
+        """
+        if service_thread_id is not None and message_store is not None:
+            raise AgentThreadException("Only the service_thread_id or message_store may be set, but not both.")
+
+        self._service_thread_id = service_thread_id
+        self._message_store = message_store
+        self.context_provider = context_provider
+
+    @property
+    def is_initialized(self) -> bool:
+        """Indicates if the thread is initialized.
+
+        This means either the service_thread_id or the message_store is set.
+        """
+        return self._service_thread_id is not None or self._message_store is not None
+
+    @property
+    def service_thread_id(self) -> str | None:
+        """Gets the ID of the current thread to support cases where the thread is owned by the agent service."""
+        return self._service_thread_id
+
+    @service_thread_id.setter
+    def service_thread_id(self, service_thread_id: str | None) -> None:
+        """Sets the ID of the current thread to support cases where the thread is owned by the agent service.
+
+        Note that either service_thread_id or message_store may be set, but not both.
+        """
+        if service_thread_id is None:
+            return
+
+        if self._message_store is not None:
+            raise AgentThreadException(
+                "Only the service_thread_id or message_store may be set, "
+                "but not both and switching from one to another is not supported."
+            )
+        self._service_thread_id = service_thread_id
+
+    @property
+    def message_store(self) -> ChatMessageStoreProtocol | None:
+        """Gets the ChatMessageStoreProtocol used by this thread."""
+        return self._message_store
+
+    @message_store.setter
+    def message_store(self, message_store: ChatMessageStoreProtocol | None) -> None:
+        """Sets the ChatMessageStoreProtocol used by this thread.
+
+        Note that either service_thread_id or message_store may be set, but not both.
+        """
+        if message_store is None:
+            return
+
+        if self._service_thread_id is not None:
+            raise AgentThreadException(
+                "Only the service_thread_id or message_store may be set, "
+                "but not both and switching from one to another is not supported."
+            )
+
+        self._message_store = message_store
+
+    async def on_new_messages(self, new_messages: ChatMessage | Sequence[ChatMessage]) -> None:
+        """Invoked when a new message has been contributed to the chat by any participant."""
+        if self._service_thread_id is not None:
+            # If the thread messages are stored in the service there is nothing to do here,
+            # since invoking the service should already update the thread.
+            return
+        if self._message_store is None:
+            # If there is no conversation id, and no store we can
+            # create a default in memory store.
+            self._message_store = ChatMessageStore()
+        # If a store has been provided, we need to add the messages to the store.
+        if isinstance(new_messages, ChatMessage):
+            new_messages = [new_messages]
+        await self._message_store.add_messages(new_messages)
+
+    async def serialize(self, **kwargs: Any) -> dict[str, Any]:
+        """Serializes the current object's state.
+
+        Args:
+            **kwargs: Arguments for serialization.
+        """
+        chat_message_store_state = None
+        if self._message_store is not None:
+            chat_message_store_state = await self._message_store.serialize(**kwargs)
+
+        state = AgentThreadState(
+            service_thread_id=self._service_thread_id, chat_message_store_state=chat_message_store_state
+        )
+        return state.model_dump()
+
+    @classmethod
+    async def deserialize(
+        cls: type[TAgentThread],
+        serialized_thread_state: dict[str, Any],
+        *,
+        message_store: ChatMessageStoreProtocol | None = None,
+        **kwargs: Any,
+    ) -> TAgentThread:
+        """Deserializes the state from a dictionary into a new AgentThread instance.
+
+        Args:
+            serialized_thread_state: The serialized thread state as a dictionary.
+            message_store: Optional ChatMessageStoreProtocol to use for managing messages.
+                If not provided, a new ChatMessageStore will be created if needed.
+            **kwargs: Additional arguments for deserialization.
+
+        Returns:
+            A new AgentThread instance with properties set from the serialized state.
+        """
+        state = AgentThreadState.model_validate(serialized_thread_state)
+
+        if state.service_thread_id is not None:
+            return cls(service_thread_id=state.service_thread_id)
+
+        # If we don't have any ChatMessageStoreProtocol state return here.
+        if state.chat_message_store_state is None:
+            return cls()
+
+        if message_store is not None:
+            try:
+                await message_store.update_from_state(state.chat_message_store_state, **kwargs)
+            except Exception as ex:
+                raise AgentThreadException("Failed to deserialize the provided message store.") from ex
+            return cls(message_store=message_store)
+        try:
+            message_store = await ChatMessageStore.deserialize(state.chat_message_store_state, **kwargs)
+        except Exception as ex:
+            raise AgentThreadException("Failed to deserialize the message store.") from ex
+        return cls(message_store=message_store)
+
+    async def update_from_thread_state(
+        self,
+        serialized_thread_state: dict[str, Any],
+        **kwargs: Any,
+    ) -> None:
+        """Deserializes the state from a dictionary into the thread properties."""
+        state = AgentThreadState.model_validate(serialized_thread_state)
+
+        if state.service_thread_id is not None:
+            self.service_thread_id = state.service_thread_id
+            # Since we have an ID, we should not have a chat message store and we can return here.
+            return
+        # If we don't have any ChatMessageStoreProtocol state return here.
+        if state.chat_message_store_state is None:
+            return
+        if self.message_store is not None:
+            await self.message_store.update_from_state(state.chat_message_store_state, **kwargs)
+            # If we don't have a chat message store yet, create an in-memory one.
+            return
+        # Create the message store from the default.
+        self.message_store = await ChatMessageStore.deserialize(state.chat_message_store_state, **kwargs)  # type: ignore

python/packages/core/agent_framework/_tools.py

@@ -0,0 +1,954 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import inspect
+import json
+import sys
+from collections.abc import AsyncIterable, Awaitable, Callable, Collection, MutableMapping, Sequence
+from functools import wraps
+from time import perf_counter, time_ns
+from typing import (
+    TYPE_CHECKING,
+    Annotated,
+    Any,
+    Final,
+    Generic,
+    Literal,
+    Protocol,
+    TypeVar,
+    get_args,
+    get_origin,
+    runtime_checkable,
+)
+
+from opentelemetry.metrics import Histogram
+from pydantic import AnyUrl, BaseModel, Field, PrivateAttr, ValidationError, create_model, field_validator
+
+from ._logging import get_logger
+from ._pydantic import AFBaseModel
+from .exceptions import ChatClientInitializationError, ToolException
+from .observability import (
+    OPERATION_DURATION_BUCKET_BOUNDARIES,
+    OtelAttr,
+    capture_exception,  # type: ignore
+    get_function_span,
+    get_function_span_attributes,
+    get_meter,
+)
+
+if TYPE_CHECKING:
+    from ._clients import ChatClientProtocol
+    from ._types import (
+        ChatMessage,
+        ChatResponse,
+        ChatResponseUpdate,
+        Contents,
+        FunctionCallContent,
+    )
+
+if sys.version_info >= (3, 12):
+    from typing import TypedDict  # pragma: no cover
+else:
+    from typing_extensions import TypedDict  # pragma: no cover
+
+logger = get_logger()
+
+__all__ = [
+    "FUNCTION_INVOKING_CHAT_CLIENT_MARKER",
+    "AIFunction",
+    "HostedCodeInterpreterTool",
+    "HostedFileSearchTool",
+    "HostedMCPSpecificApproval",
+    "HostedMCPTool",
+    "HostedWebSearchTool",
+    "ToolProtocol",
+    "ai_function",
+    "use_function_invocation",
+]
+
+
+logger = get_logger()
+FUNCTION_INVOKING_CHAT_CLIENT_MARKER: Final[str] = "__function_invoking_chat_client__"
+DEFAULT_MAX_ITERATIONS: Final[int] = 10
+TChatClient = TypeVar("TChatClient", bound="ChatClientProtocol")
+# region Helpers
+
+ArgsT = TypeVar("ArgsT", bound=BaseModel)
+ReturnT = TypeVar("ReturnT")
+
+
+class _NoOpHistogram:
+    def record(self, *args: Any, **kwargs: Any) -> None:  # pragma: no cover - trivial
+        return None
+
+
+_NOOP_HISTOGRAM = _NoOpHistogram()
+
+
+def _parse_inputs(
+    inputs: "Contents | dict[str, Any] | str | list[Contents | dict[str, Any] | str] | None",
+) -> list["Contents"]:
+    """Parse the inputs for a tool, ensuring they are of type Contents."""
+    if inputs is None:
+        return []
+
+    from ._types import BaseContent, DataContent, HostedFileContent, HostedVectorStoreContent, UriContent
+
+    parsed_inputs: list["Contents"] = []
+    if not isinstance(inputs, list):
+        inputs = [inputs]
+    for input_item in inputs:
+        if isinstance(input_item, str):
+            # If it's a string, we assume it's a URI or similar identifier.
+            # Convert it to a UriContent or similar type as needed.
+            parsed_inputs.append(UriContent(uri=input_item, media_type="text/plain"))
+        elif isinstance(input_item, dict):
+            # If it's a dict, we assume it contains properties for a specific content type.
+            # we check if the required keys are present to determine the type.
+            # for instance, if it has "uri" and "media_type", we treat it as UriContent.
+            # if is only has uri, then we treat it as DataContent.
+            # etc.
+            if "uri" in input_item:
+                parsed_inputs.append(
+                    UriContent(**input_item) if "media_type" in input_item else DataContent(**input_item)
+                )
+            elif "file_id" in input_item:
+                parsed_inputs.append(HostedFileContent(**input_item))
+            elif "vector_store_id" in input_item:
+                parsed_inputs.append(HostedVectorStoreContent(**input_item))
+            elif "data" in input_item:
+                parsed_inputs.append(DataContent(**input_item))
+            else:
+                raise ValueError(f"Unsupported input type: {input_item}")
+        elif isinstance(input_item, BaseContent):
+            parsed_inputs.append(input_item)
+        else:
+            raise TypeError(f"Unsupported input type: {type(input_item).__name__}. Expected Contents or dict.")
+    return parsed_inputs
+
+
+# region Tools
+@runtime_checkable
+class ToolProtocol(Protocol):
+    """Represents a generic tool that can be specified to an AI service.
+
+    Parameters:
+        name: The name of the tool.
+        description: A description of the tool.
+        additional_properties: Additional properties associated with the tool.
+    """
+
+    name: str
+    """The name of the tool."""
+    description: str
+    """A description of the tool, suitable for use in describing the purpose to a model."""
+    additional_properties: dict[str, Any] | None
+    """Additional properties associated with the tool."""
+
+    def __str__(self) -> str:
+        """Return a string representation of the tool."""
+        ...
+
+
+class BaseTool(AFBaseModel):
+    """Base class for AI tools, providing common attributes and methods.
+
+    Args:
+        name: The name of the tool.
+        description: A description of the tool.
+        additional_properties: Additional properties associated with the tool.
+    """
+
+    name: str = Field(..., kw_only=False)
+    description: str = ""
+    additional_properties: dict[str, Any] | None = None
+
+    def __str__(self) -> str:
+        """Return a string representation of the tool."""
+        if self.description:
+            return f"{self.__class__.__name__}(name={self.name}, description={self.description})"
+        return f"{self.__class__.__name__}(name={self.name})"
+
+
+class HostedCodeInterpreterTool(BaseTool):
+    """Represents a hosted tool that can be specified to an AI service to enable it to execute generated code.
+
+    This tool does not implement code interpretation itself. It serves as a marker to inform a service
+    that it is allowed to execute generated code if the service is capable of doing so.
+    """
+
+    inputs: list[Any] = Field(default_factory=list)
+
+    def __init__(
+        self,
+        *,
+        inputs: "Contents | dict[str, Any] | str | list[Contents | dict[str, Any] | str] | None" = None,
+        description: str | None = None,
+        additional_properties: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize the HostedCodeInterpreterTool.
+
+        Args:
+            inputs: A list of contents that the tool can accept as input. Defaults to None.
+                This should mostly be HostedFileContent or HostedVectorStoreContent.
+                Can also be DataContent, depending on the service used.
+                When supplying a list, it can contain:
+                - Contents instances
+                - dicts with properties for Contents (e.g., {"uri": "http://example.com", "media_type": "text/html"})
+                - strings (which will be converted to UriContent with media_type "text/plain").
+                If None, defaults to an empty list.
+            description: A description of the tool.
+            additional_properties: Additional properties associated with the tool.
+            **kwargs: Additional keyword arguments to pass to the base class.
+        """
+        args: dict[str, Any] = {
+            "name": "code_interpreter",
+        }
+        if inputs:
+            args["inputs"] = _parse_inputs(inputs)
+        if description is not None:
+            args["description"] = description
+        if additional_properties is not None:
+            args["additional_properties"] = additional_properties
+        if "name" in kwargs:
+            raise ValueError("The 'name' argument is reserved for the HostedCodeInterpreterTool and cannot be set.")
+        super().__init__(**args, **kwargs)
+
+
+class HostedWebSearchTool(BaseTool):
+    """Represents a web search tool that can be specified to an AI service to enable it to perform web searches."""
+
+    def __init__(
+        self,
+        description: str | None = None,
+        additional_properties: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ):
+        """Initialize a HostedWebSearchTool.
+
+        Args:
+            description: A description of the tool.
+            additional_properties: Additional properties associated with the tool
+                (e.g., {"user_location": {"city": "Seattle", "country": "US"}}).
+            **kwargs: Additional keyword arguments to pass to the base class.
+                if additional_properties is not provided, any kwargs will be added to additional_properties.
+        """
+        args: dict[str, Any] = {
+            "name": "web_search",
+        }
+        if additional_properties is not None:
+            args["additional_properties"] = additional_properties
+        elif kwargs:
+            args["additional_properties"] = kwargs
+        if description is not None:
+            args["description"] = description
+        super().__init__(**args)
+
+
+class HostedMCPSpecificApproval(TypedDict, total=False):
+    """Represents the `specific` mode for a hosted tool.
+
+    When using this mode, the user must specify which tools always or never require approval.
+    This is represented as a dictionary with two optional keys:
+    - `always_require_approval`: A sequence of tool names that always require approval.
+    - `never_require_approval`: A sequence of tool names that never require approval.
+
+    """
+
+    always_require_approval: Collection[str] | None
+    never_require_approval: Collection[str] | None
+
+
+class HostedMCPTool(BaseTool):
+    """Represents a MCP tool that is managed and executed by the service."""
+
+    url: AnyUrl
+    approval_mode: Literal["always_require", "never_require"] | HostedMCPSpecificApproval | None = None
+    allowed_tools: set[str] | None = None
+    headers: dict[str, str] | None = None
+
+    def __init__(
+        self,
+        *,
+        name: str,
+        description: str | None = None,
+        url: AnyUrl | str,
+        approval_mode: Literal["always_require", "never_require"] | HostedMCPSpecificApproval | None = None,
+        allowed_tools: Collection[str] | None = None,
+        headers: dict[str, str] | None = None,
+        additional_properties: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Create a hosted MCP tool.
+
+        Args:
+            name: The name of the tool.
+            description: A description of the tool.
+            url: The URL of the tool.
+            approval_mode: The approval mode for the tool. This can be:
+                - "always_require": The tool always requires approval before use.
+                - "never_require": The tool never requires approval before use.
+                - A dict with keys `always_require_approval` or `never_require_approval`,
+                  followed by a sequence of strings with the names of the relevant tools.
+            allowed_tools: A list of tools that are allowed to use this tool.
+            headers: Headers to include in requests to the tool.
+            additional_properties: Additional properties to include in the tool definition.
+            **kwargs: Additional keyword arguments to pass to the base class.
+        """
+        args: dict[str, Any] = {
+            "name": name,
+            "url": url,
+        }
+        if allowed_tools is not None:
+            args["allowed_tools"] = allowed_tools
+        if approval_mode is not None:
+            args["approval_mode"] = approval_mode
+        if headers is not None:
+            args["headers"] = headers
+        if description is not None:
+            args["description"] = description
+        if additional_properties is not None:
+            args["additional_properties"] = additional_properties
+        try:
+            super().__init__(**args, **kwargs)
+        except ValidationError as err:
+            raise ToolException(f"Error initializing HostedMCPTool: {err}", inner_exception=err) from err
+
+    @field_validator("approval_mode")
+    def validate_approval_mode(cls, approval_mode: str | dict[str, Any] | None) -> str | dict[str, Any] | None:
+        """Validate the approval_mode field to ensure it is one of the accepted values."""
+        if approval_mode is None or not isinstance(approval_mode, dict):
+            return approval_mode
+        # Validate that the dict has sets
+        for key, value in approval_mode.items():
+            if not isinstance(value, set):
+                approval_mode[key] = set(value)  # Convert to set if it's a list or other collection
+        return approval_mode
+
+
+class HostedFileSearchTool(BaseTool):
+    """Represents a file search tool that can be specified to an AI service to enable it to perform file searches."""
+
+    inputs: list[Any] | None = None
+    max_results: int | None = None
+
+    def __init__(
+        self,
+        inputs: "Contents | dict[str, Any] | str | list[Contents | dict[str, Any] | str] | None" = None,
+        max_results: int | None = None,
+        description: str | None = None,
+        additional_properties: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ):
+        """Initialize a FileSearchTool.
+
+        Args:
+            inputs: A list of contents that the tool can accept as input. Defaults to None.
+                This should be one or more HostedVectorStoreContents.
+                When supplying a list, it can contain:
+                - Contents instances
+                - dicts with properties for Contents (e.g., {"uri": "http://example.com", "media_type": "text/html"})
+                - strings (which will be converted to UriContent with media_type "text/plain").
+                If None, defaults to an empty list.
+            max_results: The maximum number of results to return from the file search.
+                If None, max limit is applied.
+            description: A description of the tool.
+            additional_properties: Additional properties associated with the tool.
+            **kwargs: Additional keyword arguments to pass to the base class.
+        """
+        args: dict[str, Any] = {
+            "name": "file_search",
+        }
+        if inputs:
+            args["inputs"] = _parse_inputs(inputs)
+        if max_results:
+            args["max_results"] = max_results
+        if description is not None:
+            args["description"] = description
+        if additional_properties is not None:
+            args["additional_properties"] = additional_properties
+        if "name" in kwargs:
+            raise ValueError("The 'name' argument is reserved for the HostedFileSearchTool and cannot be set.")
+        super().__init__(**args, **kwargs)
+
+
+def _default_histogram() -> Histogram:
+    """Get the default histogram for function invocation duration."""
+    from .observability import OBSERVABILITY_SETTINGS  # local import to avoid circulars
+
+    if not OBSERVABILITY_SETTINGS.ENABLED:  # type: ignore[name-defined]
+        return _NOOP_HISTOGRAM  # type: ignore[return-value]
+    meter = get_meter()
+    try:
+        return meter.create_histogram(
+            name=OtelAttr.MEASUREMENT_FUNCTION_INVOCATION_DURATION,
+            unit=OtelAttr.DURATION_UNIT,
+            description="Measures the duration of a function's execution",
+            explicit_bucket_boundaries_advisory=OPERATION_DURATION_BUCKET_BOUNDARIES,
+        )
+    except TypeError:
+        return meter.create_histogram(
+            name=OtelAttr.MEASUREMENT_FUNCTION_INVOCATION_DURATION,
+            unit=OtelAttr.DURATION_UNIT,
+            description="Measures the duration of a function's execution",
+        )
+
+
+class AIFunction(BaseTool, Generic[ArgsT, ReturnT]):
+    """A AITool that is callable as code.
+
+    Args:
+        name: The name of the function.
+        description: A description of the function.
+        additional_properties: Additional properties to set on the function.
+        func: The function to wrap. If None, returns a decorator.
+        input_model: The Pydantic model that defines the input parameters for the function.
+    """
+
+    func: Callable[..., Awaitable[ReturnT] | ReturnT]
+    input_model: type[ArgsT]
+    _invocation_duration_histogram: Histogram = PrivateAttr(default_factory=_default_histogram)
+
+    def __call__(self, *args: Any, **kwargs: Any) -> ReturnT | Awaitable[ReturnT]:
+        """Call the wrapped function with the provided arguments."""
+        return self.func(*args, **kwargs)
+
+    async def invoke(
+        self,
+        *,
+        arguments: ArgsT | None = None,
+        **kwargs: Any,
+    ) -> ReturnT:
+        """Run the AI function with the provided arguments as a Pydantic model.
+
+        Args:
+            arguments: A Pydantic model instance containing the arguments for the function.
+            kwargs: keyword arguments to pass to the function, will not be used if `arguments` is provided.
+        """
+        global OBSERVABILITY_SETTINGS
+        from .observability import OBSERVABILITY_SETTINGS
+
+        tool_call_id = kwargs.pop("tool_call_id", None)
+        if arguments is not None:
+            if not isinstance(arguments, self.input_model):
+                raise TypeError(f"Expected {self.input_model.__name__}, got {type(arguments).__name__}")
+            kwargs = arguments.model_dump(exclude_none=True)
+        if not OBSERVABILITY_SETTINGS.ENABLED:  # type: ignore[name-defined]
+            logger.info(f"Function name: {self.name}")
+            logger.debug(f"Function arguments: {kwargs}")
+            res = self.__call__(**kwargs)
+            result = await res if inspect.isawaitable(res) else res
+            logger.info(f"Function {self.name} succeeded.")
+            logger.debug(f"Function result: {result or 'None'}")
+            return result  # type: ignore[reportReturnType]
+
+        attributes = get_function_span_attributes(self, tool_call_id=tool_call_id)
+        if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED:  # type: ignore[name-defined]
+            attributes.update({
+                OtelAttr.TOOL_ARGUMENTS: arguments.model_dump_json()
+                if arguments
+                else json.dumps(kwargs)
+                if kwargs
+                else "None"
+            })
+        with get_function_span(attributes=attributes) as span:
+            attributes[OtelAttr.MEASUREMENT_FUNCTION_TAG_NAME] = self.name
+            logger.info(f"Function name: {self.name}")
+            if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED:  # type: ignore[name-defined]
+                logger.debug(f"Function arguments: {kwargs}")
+            start_time_stamp = perf_counter()
+            end_time_stamp: float | None = None
+            try:
+                res = self.__call__(**kwargs)
+                result = await res if inspect.isawaitable(res) else res
+                end_time_stamp = perf_counter()
+            except Exception as exception:
+                end_time_stamp = perf_counter()
+                attributes[OtelAttr.ERROR_TYPE] = type(exception).__name__
+                capture_exception(span=span, exception=exception, timestamp=time_ns())
+                logger.error(f"Function failed. Error: {exception}")
+                raise
+            else:
+                logger.info(f"Function {self.name} succeeded.")
+                if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED:  # type: ignore[name-defined]
+                    try:
+                        json_result = json.dumps(result)
+                    except (TypeError, OverflowError):
+                        span.set_attribute(OtelAttr.TOOL_RESULT, "<non-serializable result>")
+                        logger.debug("Function result: <non-serializable result>")
+                    else:
+                        span.set_attribute(OtelAttr.TOOL_RESULT, json_result)
+                        logger.debug(f"Function result: {json_result}")
+                return result  # type: ignore[reportReturnType]
+            finally:
+                duration = (end_time_stamp or perf_counter()) - start_time_stamp
+                span.set_attribute(OtelAttr.MEASUREMENT_FUNCTION_INVOCATION_DURATION, duration)
+                self._invocation_duration_histogram.record(duration, attributes=attributes)
+                logger.info("Function duration: %fs", duration)
+
+    def parameters(self) -> dict[str, Any]:
+        """Create the json schema of the parameters."""
+        return self.input_model.model_json_schema()
+
+    def to_json_schema_spec(self) -> dict[str, Any]:
+        """Convert a AIFunction to the JSON Schema function specification format."""
+        return {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "parameters": self.parameters(),
+            },
+        }
+
+
+# region AI Function Decorator
+
+
+def _parse_annotation(annotation: Any) -> Any:
+    """Parse a type annotation and return the corresponding type.
+
+    If the second annotation (after the type) is a string, then we convert that to a pydantic Field description.
+    The rest are returned as-is, allowing for multiple annotations.
+    """
+    origin = get_origin(annotation)
+    if origin is not None:
+        args = get_args(annotation)
+        # For other generics, return the origin type (e.g., list for List[int])
+        if len(args) > 1 and isinstance(args[1], str):
+            # Create a new Annotated type with the updated Field
+            args_list = list(args)
+            if len(args_list) == 2:
+                return Annotated[args_list[0], Field(description=args_list[1])]
+            return Annotated[args_list[0], Field(description=args_list[1]), tuple(args_list[2:])]
+    return annotation
+
+
+def ai_function(
+    func: Callable[..., ReturnT | Awaitable[ReturnT]] | None = None,
+    *,
+    name: str | None = None,
+    description: str | None = None,
+    additional_properties: dict[str, Any] | None = None,
+) -> AIFunction[Any, ReturnT]:
+    """Decorate a function to turn it into a AIFunction that can be passed to models and executed automatically.
+
+    This function will create a Pydantic model from the function's signature,
+    which will be used to validate the arguments passed to the function.
+    And will be used to generate the JSON schema for the function's parameters.
+    In order to add descriptions to parameters, in your function signature,
+    use the `Annotated` type from `typing` and the `Field` class from `pydantic`:
+
+    Example:
+
+        .. code-block:: python
+
+            from typing import Annotated
+            from pydantic import Field
+
+
+            def ai_function_example(
+                arg1: Annotated[str, Field(description="The first argument")],
+                arg2: Annotated[int, Field(description="The second argument")],
+            ) -> str:
+                # An example function that takes two arguments and returns a string.
+                return f"arg1: {arg1}, arg2: {arg2}"
+
+    Args:
+        func: The function to wrap. If None, returns a decorator.
+        name: The name of the tool. Defaults to the function's name.
+        description: A description of the tool. Defaults to the function's docstring.
+        additional_properties: Additional properties to set on the tool.
+
+    """
+
+    def decorator(func: Callable[..., ReturnT | Awaitable[ReturnT]]) -> AIFunction[Any, ReturnT]:
+        @wraps(func)
+        def wrapper(f: Callable[..., ReturnT | Awaitable[ReturnT]]) -> AIFunction[Any, ReturnT]:
+            tool_name: str = name or getattr(f, "__name__", "unknown_function")  # type: ignore[assignment]
+            tool_desc: str = description or (f.__doc__ or "")
+            sig = inspect.signature(f)
+            fields = {
+                pname: (
+                    _parse_annotation(param.annotation) if param.annotation is not inspect.Parameter.empty else str,
+                    param.default if param.default is not inspect.Parameter.empty else ...,
+                )
+                for pname, param in sig.parameters.items()
+                if pname not in {"self", "cls"}
+            }
+            input_model: Any = create_model(f"{tool_name}_input", **fields)  # type: ignore[call-overload]
+            if not issubclass(input_model, BaseModel):
+                raise TypeError(f"Input model for {tool_name} must be a subclass of BaseModel, got {input_model}")
+
+            return AIFunction[Any, ReturnT](
+                name=tool_name,
+                description=tool_desc,
+                additional_properties=additional_properties or {},
+                func=f,
+                input_model=input_model,
+            )
+
+        return wrapper(func)
+
+    return decorator(func) if func else decorator  # type: ignore[reportReturnType, return-value]
+
+
+# region Function Invoking Chat Client
+
+
+async def _auto_invoke_function(
+    function_call_content: "FunctionCallContent",
+    custom_args: dict[str, Any] | None = None,
+    *,
+    tool_map: dict[str, AIFunction[BaseModel, Any]],
+    sequence_index: int | None = None,
+    request_index: int | None = None,
+    middleware_pipeline: Any = None,  # Optional MiddlewarePipeline
+) -> "Contents":
+    """Invoke a function call requested by the agent, applying filters that are defined in the agent."""
+    from ._types import FunctionResultContent
+
+    tool: AIFunction[BaseModel, Any] | None = tool_map.get(function_call_content.name)
+    if tool is None:
+        raise KeyError(f"No tool or function named '{function_call_content.name}'")
+
+    parsed_args: dict[str, Any] = dict(function_call_content.parse_arguments() or {})
+
+    # Merge with user-supplied args; right-hand side dominates, so parsed args win on conflicts.
+    merged_args: dict[str, Any] = (custom_args or {}) | parsed_args
+    args = tool.input_model.model_validate(merged_args)
+    exception = None
+
+    # Execute through middleware pipeline if available
+    if middleware_pipeline and hasattr(middleware_pipeline, "has_middlewares") and middleware_pipeline.has_middlewares:
+        from ._middleware import FunctionInvocationContext
+
+        middleware_context = FunctionInvocationContext(
+            function=tool,
+            arguments=args,
+            kwargs=custom_args or {},
+        )
+
+        async def final_function_handler(context_obj: Any) -> Any:
+            return await tool.invoke(
+                arguments=context_obj.arguments,
+                tool_call_id=function_call_content.call_id,
+            )
+
+        try:
+            function_result = await middleware_pipeline.execute(
+                function=tool,
+                arguments=args,
+                context=middleware_context,
+                final_handler=final_function_handler,
+            )
+        except Exception as ex:
+            exception = ex
+            function_result = None
+    else:
+        # No middleware - execute directly
+        try:
+            function_result = await tool.invoke(
+                arguments=args,
+                tool_call_id=function_call_content.call_id,
+            )  # type: ignore[arg-type]
+        except Exception as ex:
+            exception = ex
+            function_result = None
+
+    return FunctionResultContent(
+        call_id=function_call_content.call_id,
+        exception=exception,
+        result=function_result,
+    )
+
+
+def _get_tool_map(
+    tools: "ToolProtocol \
+    | Callable[..., Any] \
+    | MutableMapping[str, Any] \
+    | list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]",
+) -> dict[str, AIFunction[Any, Any]]:
+    ai_function_list: dict[str, AIFunction[Any, Any]] = {}
+    for tool in tools if isinstance(tools, list) else [tools]:
+        if isinstance(tool, AIFunction):
+            ai_function_list[tool.name] = tool
+            continue
+        if callable(tool):
+            # Convert to AITool if it's a function or callable
+            ai_tool = ai_function(tool)
+            ai_function_list[ai_tool.name] = ai_tool
+    return ai_function_list
+
+
+async def execute_function_calls(
+    custom_args: dict[str, Any],
+    attempt_idx: int,
+    function_calls: Sequence["FunctionCallContent"],
+    tools: "ToolProtocol \
+    | Callable[..., Any] \
+    | MutableMapping[str, Any] \
+    | list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]",
+    middleware_pipeline: Any = None,  # Optional MiddlewarePipeline to avoid circular imports
+) -> list["Contents"]:
+    tool_map = _get_tool_map(tools)
+    # Run all function calls concurrently
+    return await asyncio.gather(*[
+        _auto_invoke_function(
+            function_call_content=function_call,
+            custom_args=custom_args,
+            tool_map=tool_map,
+            sequence_index=seq_idx,
+            request_index=attempt_idx,
+            middleware_pipeline=middleware_pipeline,
+        )
+        for seq_idx, function_call in enumerate(function_calls)
+    ])
+
+
+def update_conversation_id(kwargs: dict[str, Any], conversation_id: str | None) -> None:
+    """Update kwargs with conversation id."""
+    if conversation_id is None:
+        return
+    if "chat_options" in kwargs:
+        kwargs["chat_options"].conversation_id = conversation_id
+    else:
+        kwargs["conversation_id"] = conversation_id
+
+
+def _handle_function_calls_response(
+    func: Callable[..., Awaitable["ChatResponse"]],
+    *,
+    max_iterations: int = 10,
+) -> Callable[..., Awaitable["ChatResponse"]]:
+    """Decorate the get_response method to enable function calls.
+
+    Args:
+        func: The get_response method to decorate.
+        max_iterations: The maximum number of function call iterations to perform.
+
+    """
+
+    def decorator(
+        func: Callable[..., Awaitable["ChatResponse"]],
+    ) -> Callable[..., Awaitable["ChatResponse"]]:
+        """Inner decorator."""
+
+        @wraps(func)
+        async def function_invocation_wrapper(
+            self: "ChatClientProtocol",
+            messages: "str | ChatMessage | list[str] | list[ChatMessage]",
+            **kwargs: Any,
+        ) -> "ChatResponse":
+            from ._clients import prepare_messages
+            from ._middleware import extract_and_merge_function_middleware
+            from ._types import ChatMessage, ChatOptions, FunctionCallContent, FunctionResultContent
+
+            # Extract and merge function middleware from chat client with kwargs pipeline
+            extract_and_merge_function_middleware(self, kwargs)
+
+            # Extract the middleware pipeline before calling the underlying function
+            # because the underlying function may not preserve it in kwargs
+            stored_middleware_pipeline = kwargs.get("_function_middleware_pipeline")
+
+            prepped_messages = prepare_messages(messages)
+            response: "ChatResponse | None" = None
+            fcc_messages: "list[ChatMessage]" = []
+            for attempt_idx in range(max_iterations):
+                response = await func(self, messages=prepped_messages, **kwargs)
+                # if there are function calls, we will handle them first
+                function_results = {
+                    it.call_id for it in response.messages[0].contents if isinstance(it, FunctionResultContent)
+                }
+                function_calls = [
+                    it
+                    for it in response.messages[0].contents
+                    if isinstance(it, FunctionCallContent) and it.call_id not in function_results
+                ]
+
+                if response.conversation_id is not None:
+                    update_conversation_id(kwargs, response.conversation_id)
+                    prepped_messages = []
+
+                tools = kwargs.get("tools")
+                if not tools and (chat_options := kwargs.get("chat_options")) and isinstance(chat_options, ChatOptions):
+                    tools = chat_options.tools
+                if function_calls and tools:
+                    # Use the stored middleware pipeline instead of extracting from kwargs
+                    # because kwargs may have been modified by the underlying function
+                    middleware_pipeline = stored_middleware_pipeline
+                    function_results = await execute_function_calls(
+                        custom_args=kwargs,
+                        attempt_idx=attempt_idx,
+                        function_calls=function_calls,
+                        tools=tools,  # type: ignore
+                        middleware_pipeline=middleware_pipeline,
+                    )
+                    # add a single ChatMessage to the response with the results
+                    result_message = ChatMessage(role="tool", contents=function_results)  # type: ignore[call-overload]
+                    response.messages.append(result_message)
+                    # response should contain 2 messages after this,
+                    # one with function call contents
+                    # and one with function result contents
+                    # the amount and call_id's should match
+                    # this runs in every but the first run
+                    # we need to keep track of all function call messages
+                    fcc_messages.extend(response.messages)
+                    # and add them as additional context to the messages
+                    if getattr(kwargs.get("chat_options"), "store", False):
+                        prepped_messages.clear()
+                        prepped_messages.append(result_message)
+                    else:
+                        prepped_messages.extend(response.messages)
+                    continue
+                # If we reach this point, it means there were no function calls to handle,
+                # we'll add the previous function call and responses
+                # to the front of the list, so that the final response is the last one
+                # TODO (eavanvalkenburg): control this behavior?
+                if fcc_messages:
+                    for msg in reversed(fcc_messages):
+                        response.messages.insert(0, msg)
+                return response
+
+            # Failsafe: give up on tools, ask model for plain answer
+            kwargs["tool_choice"] = "none"
+            response = await func(self, messages=prepped_messages, **kwargs)
+            if fcc_messages:
+                for msg in reversed(fcc_messages):
+                    response.messages.insert(0, msg)
+            return response
+
+        return function_invocation_wrapper  # type: ignore
+
+    return decorator(func)
+
+
+def _handle_function_calls_streaming_response(
+    func: Callable[..., AsyncIterable["ChatResponseUpdate"]],
+    *,
+    max_iterations: int = 10,
+) -> Callable[..., AsyncIterable["ChatResponseUpdate"]]:
+    """Decorate the get_streaming_response method to handle function calls.
+
+    Args:
+        func: The get_streaming_response method to decorate.
+        max_iterations: The maximum number of function call iterations to perform.
+
+    """
+
+    def decorator(
+        func: Callable[..., AsyncIterable["ChatResponseUpdate"]],
+    ) -> Callable[..., AsyncIterable["ChatResponseUpdate"]]:
+        """Inner decorator."""
+
+        @wraps(func)
+        async def streaming_function_invocation_wrapper(
+            self: "ChatClientProtocol",
+            messages: "str | ChatMessage | list[str] | list[ChatMessage]",
+            **kwargs: Any,
+        ) -> AsyncIterable["ChatResponseUpdate"]:
+            """Wrap the inner get streaming response method to handle tool calls."""
+            from ._clients import prepare_messages
+            from ._middleware import extract_and_merge_function_middleware
+            from ._types import ChatMessage, ChatOptions, ChatResponse, ChatResponseUpdate, FunctionCallContent
+
+            # Extract and merge function middleware from chat client with kwargs pipeline
+            extract_and_merge_function_middleware(self, kwargs)
+
+            # Extract the middleware pipeline before calling the underlying function
+            # because the underlying function may not preserve it in kwargs
+            stored_middleware_pipeline = kwargs.get("_function_middleware_pipeline")
+
+            prepped_messages = prepare_messages(messages)
+            for attempt_idx in range(max_iterations):
+                all_updates: list["ChatResponseUpdate"] = []
+                async for update in func(self, messages=prepped_messages, **kwargs):
+                    all_updates.append(update)
+                    yield update
+
+                # efficient check for FunctionCallContent in the updates
+                # if there is at least one, this stops and continuous
+                # if there are no FCC's then it returns
+                if not any(isinstance(item, FunctionCallContent) for upd in all_updates for item in upd.contents):
+                    return
+
+                # Now combining the updates to create the full response.
+                # Depending on the prompt, the message may contain both function call
+                # content and others
+
+                response: "ChatResponse" = ChatResponse.from_chat_response_updates(all_updates)
+                # add the response message to the previous messages
+                prepped_messages.append(response.messages[0])
+                # get the fccs
+                function_calls = [
+                    item for item in response.messages[0].contents if isinstance(item, FunctionCallContent)
+                ]
+
+                # When conversation id is present, it means that messages are hosted on the server.
+                # In this case, we need to update kwargs with conversation id and also clear messages
+                if response.conversation_id is not None:
+                    update_conversation_id(kwargs, response.conversation_id)
+                    prepped_messages = []
+
+                tools = kwargs.get("tools")
+                if not tools and (chat_options := kwargs.get("chat_options")) and isinstance(chat_options, ChatOptions):
+                    tools = chat_options.tools
+
+                if function_calls and tools:
+                    # Use the stored middleware pipeline instead of extracting from kwargs
+                    # because kwargs may have been modified by the underlying function
+                    middleware_pipeline = stored_middleware_pipeline
+                    function_results = await execute_function_calls(
+                        custom_args=kwargs,
+                        attempt_idx=attempt_idx,
+                        function_calls=function_calls,
+                        tools=tools,  # type: ignore[reportArgumentType]
+                        middleware_pipeline=middleware_pipeline,
+                    )
+                    function_result_msg = ChatMessage(role="tool", contents=function_results)
+                    yield ChatResponseUpdate(contents=function_results, role="tool")
+                    response.messages.append(function_result_msg)
+                    prepped_messages.append(function_result_msg)
+                    continue
+
+            # Failsafe: give up on tools, ask model for plain answer
+            kwargs["tool_choice"] = "none"
+            async for update in func(self, messages=prepped_messages, **kwargs):
+                yield update
+
+        return streaming_function_invocation_wrapper
+
+    return decorator(func)
+
+
+def use_function_invocation(
+    chat_client: type[TChatClient],
+) -> type[TChatClient]:
+    """Class decorator that enables tool calling for a chat client."""
+    if getattr(chat_client, FUNCTION_INVOKING_CHAT_CLIENT_MARKER, False):
+        return chat_client
+
+    max_iterations = DEFAULT_MAX_ITERATIONS
+
+    try:
+        chat_client.get_response = _handle_function_calls_response(  # type: ignore
+            func=chat_client.get_response,  # type: ignore
+            max_iterations=max_iterations,
+        )
+    except AttributeError as ex:
+        raise ChatClientInitializationError(
+            f"Chat client {chat_client.__name__} does not have a get_response method, cannot apply function invocation."
+        ) from ex
+    try:
+        chat_client.get_streaming_response = _handle_function_calls_streaming_response(  # type: ignore
+            func=chat_client.get_streaming_response,
+            max_iterations=max_iterations,
+        )
+    except AttributeError as ex:
+        raise ChatClientInitializationError(
+            f"Chat client {chat_client.__name__} does not have a get_streaming_response method, "
+            "cannot apply function invocation."
+        ) from ex
+    setattr(chat_client, FUNCTION_INVOKING_CHAT_CLIENT_MARKER, True)
+    return chat_client

python/packages/core/agent_framework/_workflow/README.md

@@ -0,0 +1,15 @@
+# Get Started with Microsoft Agent Framework Workflow
+
+Workflow capabilities now ship with the core `agent-framework` package.
+
+```bash
+pip install agent-framework
+```
+
+Optional visualization support is still available via the `viz` extra:
+
+```bash
+pip install agent-framework[viz]
+```
+
+See the [project README](https://github.com/microsoft/agent-framework/tree/main/python/README.md) for more information.

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

@@ -0,0 +1,182 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from ._agent import WorkflowAgent
+from ._checkpoint import (
+    CheckpointStorage,
+    FileCheckpointStorage,
+    InMemoryCheckpointStorage,
+    WorkflowCheckpoint,
+)
+from ._concurrent import ConcurrentBuilder
+from ._const import (
+    DEFAULT_MAX_ITERATIONS,
+)
+from ._edge import (
+    Case,
+    Default,
+    Edge,
+    FanInEdgeGroup,
+    FanOutEdgeGroup,
+    SingleEdgeGroup,
+    SwitchCaseEdgeGroup,
+    SwitchCaseEdgeGroupCase,
+    SwitchCaseEdgeGroupDefault,
+)
+from ._edge_runner import create_edge_runner
+from ._events import (
+    AgentRunEvent,
+    AgentRunUpdateEvent,
+    ExecutorCompletedEvent,
+    ExecutorEvent,
+    ExecutorFailedEvent,
+    ExecutorInvokedEvent,
+    RequestInfoEvent,
+    WorkflowErrorDetails,
+    WorkflowEvent,
+    WorkflowEventSource,
+    WorkflowFailedEvent,
+    WorkflowLifecycleEvent,
+    WorkflowOutputEvent,
+    WorkflowRunState,
+    WorkflowStartedEvent,
+    WorkflowStatusEvent,
+)
+from ._executor import (
+    AgentExecutor,
+    AgentExecutorRequest,
+    AgentExecutorResponse,
+    Executor,
+    RequestInfoExecutor,
+    RequestInfoMessage,
+    RequestResponse,
+    handler,
+)
+from ._function_executor import FunctionExecutor, executor
+from ._magentic import (
+    MagenticAgentDeltaEvent,
+    MagenticAgentExecutor,
+    MagenticAgentMessageEvent,
+    MagenticBuilder,
+    MagenticCallbackEvent,
+    MagenticCallbackMode,
+    MagenticContext,
+    MagenticFinalResultEvent,
+    MagenticManagerBase,
+    MagenticOrchestratorExecutor,
+    MagenticOrchestratorMessageEvent,
+    MagenticPlanReviewDecision,
+    MagenticPlanReviewReply,
+    MagenticPlanReviewRequest,
+    MagenticProgressLedger,
+    MagenticProgressLedgerItem,
+    MagenticRequestMessage,
+    MagenticResponseMessage,
+    MagenticStartMessage,
+    StandardMagenticManager,
+)
+from ._runner import Runner
+from ._runner_context import (
+    InProcRunnerContext,
+    Message,
+    RunnerContext,
+)
+from ._sequential import SequentialBuilder
+from ._shared_state import SharedState
+from ._validation import (
+    EdgeDuplicationError,
+    ExecutorDuplicationError,
+    GraphConnectivityError,
+    TypeCompatibilityError,
+    ValidationTypeEnum,
+    WorkflowValidationError,
+    validate_workflow_graph,
+)
+from ._viz import WorkflowViz
+from ._workflow import Workflow, WorkflowBuilder, WorkflowRunResult
+from ._workflow_context import WorkflowContext
+from ._workflow_executor import WorkflowExecutor
+
+__all__ = [
+    "DEFAULT_MAX_ITERATIONS",
+    "AgentExecutor",
+    "AgentExecutorRequest",
+    "AgentExecutorResponse",
+    "AgentRunEvent",
+    "AgentRunUpdateEvent",
+    "Case",
+    "CheckpointStorage",
+    "ConcurrentBuilder",
+    "Default",
+    "Edge",
+    "EdgeDuplicationError",
+    "Executor",
+    "ExecutorCompletedEvent",
+    "ExecutorDuplicationError",
+    "ExecutorEvent",
+    "ExecutorFailedEvent",
+    "ExecutorInvokedEvent",
+    "FanInEdgeGroup",
+    "FanOutEdgeGroup",
+    "FileCheckpointStorage",
+    "FunctionExecutor",
+    "GraphConnectivityError",
+    "InMemoryCheckpointStorage",
+    "InProcRunnerContext",
+    "MagenticAgentDeltaEvent",
+    "MagenticAgentExecutor",
+    "MagenticAgentMessageEvent",
+    "MagenticBuilder",
+    "MagenticCallbackEvent",
+    "MagenticCallbackMode",
+    "MagenticContext",
+    "MagenticFinalResultEvent",
+    "MagenticManagerBase",
+    "MagenticOrchestratorExecutor",
+    "MagenticOrchestratorMessageEvent",
+    "MagenticPlanReviewDecision",
+    "MagenticPlanReviewReply",
+    "MagenticPlanReviewRequest",
+    "MagenticProgressLedger",
+    "MagenticProgressLedgerItem",
+    "MagenticRequestMessage",
+    "MagenticResponseMessage",
+    "MagenticStartMessage",
+    "Message",
+    "RequestInfoEvent",
+    "RequestInfoExecutor",
+    "RequestInfoMessage",
+    "RequestResponse",
+    "Runner",
+    "RunnerContext",
+    "SequentialBuilder",
+    "SharedState",
+    "SingleEdgeGroup",
+    "StandardMagenticManager",
+    "SwitchCaseEdgeGroup",
+    "SwitchCaseEdgeGroupCase",
+    "SwitchCaseEdgeGroupDefault",
+    "TypeCompatibilityError",
+    "ValidationTypeEnum",
+    "Workflow",
+    "WorkflowAgent",
+    "WorkflowBuilder",
+    "WorkflowCheckpoint",
+    "WorkflowContext",
+    "WorkflowErrorDetails",
+    "WorkflowEvent",
+    "WorkflowEventSource",
+    "WorkflowExecutor",
+    "WorkflowFailedEvent",
+    "WorkflowLifecycleEvent",
+    "WorkflowOutputEvent",
+    "WorkflowRunResult",
+    "WorkflowRunState",
+    "WorkflowStartedEvent",
+    "WorkflowStatusEvent",
+    "WorkflowValidationError",
+    "WorkflowViz",
+    "create_edge_runner",
+    "executor",
+    "handler",
+    "validate_workflow_graph",
+]

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

@@ -0,0 +1,180 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from ._agent import WorkflowAgent
+from ._checkpoint import (
+    CheckpointStorage,
+    FileCheckpointStorage,
+    InMemoryCheckpointStorage,
+    WorkflowCheckpoint,
+)
+from ._concurrent import ConcurrentBuilder
+from ._const import DEFAULT_MAX_ITERATIONS
+from ._edge import (
+    Case,
+    Default,
+    Edge,
+    FanInEdgeGroup,
+    FanOutEdgeGroup,
+    SingleEdgeGroup,
+    SwitchCaseEdgeGroup,
+    SwitchCaseEdgeGroupCase,
+    SwitchCaseEdgeGroupDefault,
+)
+from ._edge_runner import create_edge_runner
+from ._events import (
+    AgentRunEvent,
+    AgentRunUpdateEvent,
+    ExecutorCompletedEvent,
+    ExecutorEvent,
+    ExecutorFailedEvent,
+    ExecutorInvokedEvent,
+    RequestInfoEvent,
+    WorkflowErrorDetails,
+    WorkflowEvent,
+    WorkflowEventSource,
+    WorkflowFailedEvent,
+    WorkflowLifecycleEvent,
+    WorkflowOutputEvent,
+    WorkflowRunState,
+    WorkflowStartedEvent,
+    WorkflowStatusEvent,
+)
+from ._executor import (
+    AgentExecutor,
+    AgentExecutorRequest,
+    AgentExecutorResponse,
+    Executor,
+    RequestInfoExecutor,
+    RequestInfoMessage,
+    RequestResponse,
+    handler,
+)
+from ._function_executor import FunctionExecutor, executor
+from ._magentic import (
+    MagenticAgentDeltaEvent,
+    MagenticAgentExecutor,
+    MagenticAgentMessageEvent,
+    MagenticBuilder,
+    MagenticCallbackEvent,
+    MagenticCallbackMode,
+    MagenticContext,
+    MagenticFinalResultEvent,
+    MagenticManagerBase,
+    MagenticOrchestratorExecutor,
+    MagenticOrchestratorMessageEvent,
+    MagenticPlanReviewDecision,
+    MagenticPlanReviewReply,
+    MagenticPlanReviewRequest,
+    MagenticProgressLedger,
+    MagenticProgressLedgerItem,
+    MagenticRequestMessage,
+    MagenticResponseMessage,
+    MagenticStartMessage,
+    StandardMagenticManager,
+)
+from ._runner import Runner
+from ._runner_context import (
+    InProcRunnerContext,
+    Message,
+    RunnerContext,
+)
+from ._sequential import SequentialBuilder
+from ._shared_state import SharedState
+from ._validation import (
+    EdgeDuplicationError,
+    ExecutorDuplicationError,
+    GraphConnectivityError,
+    TypeCompatibilityError,
+    ValidationTypeEnum,
+    WorkflowValidationError,
+    validate_workflow_graph,
+)
+from ._viz import WorkflowViz
+from ._workflow import Workflow, WorkflowBuilder, WorkflowRunResult
+from ._workflow_context import WorkflowContext
+from ._workflow_executor import WorkflowExecutor
+
+__all__ = [
+    "DEFAULT_MAX_ITERATIONS",
+    "AgentExecutor",
+    "AgentExecutorRequest",
+    "AgentExecutorResponse",
+    "AgentRunEvent",
+    "AgentRunUpdateEvent",
+    "Case",
+    "CheckpointStorage",
+    "ConcurrentBuilder",
+    "Default",
+    "Edge",
+    "EdgeDuplicationError",
+    "Executor",
+    "ExecutorCompletedEvent",
+    "ExecutorDuplicationError",
+    "ExecutorEvent",
+    "ExecutorFailedEvent",
+    "ExecutorInvokedEvent",
+    "FanInEdgeGroup",
+    "FanOutEdgeGroup",
+    "FileCheckpointStorage",
+    "FunctionExecutor",
+    "GraphConnectivityError",
+    "InMemoryCheckpointStorage",
+    "InProcRunnerContext",
+    "MagenticAgentDeltaEvent",
+    "MagenticAgentExecutor",
+    "MagenticAgentMessageEvent",
+    "MagenticBuilder",
+    "MagenticCallbackEvent",
+    "MagenticCallbackMode",
+    "MagenticContext",
+    "MagenticFinalResultEvent",
+    "MagenticManagerBase",
+    "MagenticOrchestratorExecutor",
+    "MagenticOrchestratorMessageEvent",
+    "MagenticPlanReviewDecision",
+    "MagenticPlanReviewReply",
+    "MagenticPlanReviewRequest",
+    "MagenticProgressLedger",
+    "MagenticProgressLedgerItem",
+    "MagenticRequestMessage",
+    "MagenticResponseMessage",
+    "MagenticStartMessage",
+    "Message",
+    "RequestInfoEvent",
+    "RequestInfoExecutor",
+    "RequestInfoMessage",
+    "RequestResponse",
+    "Runner",
+    "RunnerContext",
+    "SequentialBuilder",
+    "SharedState",
+    "SingleEdgeGroup",
+    "StandardMagenticManager",
+    "SwitchCaseEdgeGroup",
+    "SwitchCaseEdgeGroupCase",
+    "SwitchCaseEdgeGroupDefault",
+    "TypeCompatibilityError",
+    "ValidationTypeEnum",
+    "Workflow",
+    "WorkflowAgent",
+    "WorkflowBuilder",
+    "WorkflowCheckpoint",
+    "WorkflowContext",
+    "WorkflowErrorDetails",
+    "WorkflowEvent",
+    "WorkflowEventSource",
+    "WorkflowExecutor",
+    "WorkflowFailedEvent",
+    "WorkflowLifecycleEvent",
+    "WorkflowOutputEvent",
+    "WorkflowRunResult",
+    "WorkflowRunState",
+    "WorkflowStartedEvent",
+    "WorkflowStatusEvent",
+    "WorkflowValidationError",
+    "WorkflowViz",
+    "create_edge_runner",
+    "executor",
+    "handler",
+    "validate_workflow_graph",
+]

python/packages/core/agent_framework/_workflow/_agent.py

@@ -0,0 +1,472 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import json
+import logging
+import uuid
+from collections.abc import AsyncIterable, Sequence
+from dataclasses import dataclass
+from datetime import datetime
+from typing import TYPE_CHECKING, Any, ClassVar, TypedDict, cast
+
+from agent_framework import (
+    AgentRunResponse,
+    AgentRunResponseUpdate,
+    AgentThread,
+    BaseAgent,
+    ChatMessage,
+    FunctionCallContent,
+    FunctionResultContent,
+    Role,
+    TextContent,
+    UsageDetails,
+)
+
+from ..exceptions import AgentExecutionException
+from ._events import (
+    AgentRunUpdateEvent,
+    RequestInfoEvent,
+    WorkflowEvent,
+)
+
+if TYPE_CHECKING:
+    from ._workflow import Workflow
+
+logger = logging.getLogger(__name__)
+
+
+class WorkflowAgent(BaseAgent):
+    """An `Agent` subclass that wraps a workflow and exposes it as an agent."""
+
+    # Class variable for the request info function name
+    REQUEST_INFO_FUNCTION_NAME: ClassVar[str] = "request_info"
+
+    @dataclass
+    class RequestInfoFunctionArgs:
+        request_id: str
+        data: Any
+
+        def to_dict(self) -> dict[str, Any]:
+            return {"request_id": self.request_id, "data": self.data}
+
+        def to_json(self) -> str:
+            return json.dumps(self.to_dict())
+
+        @classmethod
+        def from_dict(cls, payload: dict[str, Any]) -> "WorkflowAgent.RequestInfoFunctionArgs":
+            return cls(request_id=payload.get("request_id", ""), data=payload.get("data"))
+
+        @classmethod
+        def from_json(cls, raw: str) -> "WorkflowAgent.RequestInfoFunctionArgs":
+            data = json.loads(raw)
+            if not isinstance(data, dict):
+                raise ValueError("RequestInfoFunctionArgs JSON payload must decode to a mapping")
+            return cls.from_dict(data)
+
+    def __init__(
+        self,
+        workflow: "Workflow",
+        *,
+        id: str | None = None,
+        name: str | None = None,
+        description: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize the WorkflowAgent.
+
+        Args:
+            workflow: The workflow to wrap as an agent.
+            id: Unique identifier for the agent. If None, will be generated.
+            name: Optional name for the agent.
+            description: Optional description of the agent.
+            **kwargs: Additional keyword arguments passed to BaseAgent.
+        """
+        if id is None:
+            id = f"WorkflowAgent_{uuid.uuid4().hex[:8]}"
+        # Initialize with standard BaseAgent parameters first
+        # Validate the workflow's start executor can handle agent-facing message inputs
+        try:
+            start_executor = workflow.get_start_executor()
+        except KeyError as exc:  # Defensive: workflow lacks a configured entry point
+            raise ValueError("Workflow's start executor is not defined.") from exc
+
+        if list[ChatMessage] not in start_executor.input_types:
+            raise ValueError("Workflow's start executor cannot handle list[ChatMessage]")
+
+        super().__init__(id=id, name=name, description=description, **kwargs)
+        self._workflow: "Workflow" = workflow
+        self._pending_requests: dict[str, RequestInfoEvent] = {}
+
+    @property
+    def workflow(self) -> "Workflow":
+        return self._workflow
+
+    @property
+    def pending_requests(self) -> dict[str, RequestInfoEvent]:
+        return self._pending_requests
+
+    async def run(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
+        *,
+        thread: AgentThread | None = None,
+        **kwargs: Any,
+    ) -> AgentRunResponse:
+        """Get a response from the workflow agent (non-streaming).
+
+        This method collects all streaming updates and merges them into a single response.
+
+        Args:
+            messages: The message(s) to send to the workflow.
+            thread: The conversation thread. If None, a new thread will be created.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            The final workflow response as an AgentRunResponse.
+        """
+        # Collect all streaming updates
+        response_updates: list[AgentRunResponseUpdate] = []
+        input_messages = self._normalize_messages(messages)
+        thread = thread or self.get_new_thread()
+        response_id = str(uuid.uuid4())
+
+        async for update in self._run_stream_impl(input_messages, response_id):
+            response_updates.append(update)
+
+        # Convert updates to final response.
+        response = self.merge_updates(response_updates, response_id)
+
+        # Notify thread of new messages (both input and response messages)
+        await self._notify_thread_of_new_messages(thread, input_messages, response.messages)
+
+        return response
+
+    async def run_stream(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
+        *,
+        thread: AgentThread | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterable[AgentRunResponseUpdate]:
+        """Stream response updates from the workflow agent.
+
+        Args:
+            messages: The message(s) to send to the workflow.
+            thread: The conversation thread. If None, a new thread will be created.
+            **kwargs: Additional keyword arguments.
+
+        Yields:
+            AgentRunResponseUpdate objects representing the workflow execution progress.
+        """
+        input_messages = self._normalize_messages(messages)
+        thread = thread or self.get_new_thread()
+        response_updates: list[AgentRunResponseUpdate] = []
+        response_id = str(uuid.uuid4())
+
+        async for update in self._run_stream_impl(input_messages, response_id):
+            response_updates.append(update)
+            yield update
+
+        # Convert updates to final response.
+        response = self.merge_updates(response_updates, response_id)
+
+        # Notify thread of new messages (both input and response messages)
+        await self._notify_thread_of_new_messages(thread, input_messages, response.messages)
+
+    async def _run_stream_impl(
+        self,
+        input_messages: list[ChatMessage],
+        response_id: str,
+    ) -> AsyncIterable[AgentRunResponseUpdate]:
+        """Internal implementation of streaming execution.
+
+        Args:
+            input_messages: Normalized input messages to process.
+            response_id: The unique response ID for this workflow execution.
+
+        Yields:
+            AgentRunResponseUpdate objects representing the workflow execution progress.
+        """
+        # Determine the event stream based on whether we have function responses
+        if bool(self.pending_requests):
+            # This is a continuation - use send_responses_streaming to send function responses back
+            logger.info(f"Continuing workflow to address {len(self.pending_requests)} requests")
+
+            # Extract function responses from input messages, and ensure that
+            # only function responses are present in messages if there is any
+            # pending request.
+            function_responses = self._extract_function_responses(input_messages)
+
+            # Pop pending requests if fulfilled.
+            for request_id in list(self.pending_requests.keys()):
+                if request_id in function_responses:
+                    self.pending_requests.pop(request_id)
+
+            # NOTE: It is possible that some pending requests are not fulfilled,
+            # and we will let the workflow to handle this -- the agent does not
+            # have an opinion on this.
+            event_stream = self.workflow.send_responses_streaming(function_responses)
+        else:
+            # Execute workflow with streaming (initial run or no function responses)
+            # Pass the new input messages directly to the workflow
+            event_stream = self.workflow.run_stream(input_messages)
+
+        # Process events from the stream
+        async for event in event_stream:
+            # Convert workflow event to agent update
+            update = self._convert_workflow_event_to_agent_update(response_id, event)
+            if update:
+                yield update
+
+    def _normalize_messages(
+        self,
+        messages: str | ChatMessage | Sequence[str] | Sequence[ChatMessage] | None = None,
+    ) -> list[ChatMessage]:
+        """Normalize input messages to a list of ChatMessage objects."""
+        if messages is None:
+            return []
+
+        if isinstance(messages, str):
+            return [ChatMessage(role=Role.USER, contents=[TextContent(text=messages)])]
+
+        if isinstance(messages, ChatMessage):
+            return [messages]
+
+        normalized: list[ChatMessage] = []
+        for msg in messages:
+            if isinstance(msg, str):
+                normalized.append(ChatMessage(role=Role.USER, contents=[TextContent(text=msg)]))
+            elif isinstance(msg, ChatMessage):
+                normalized.append(msg)
+        return normalized
+
+    def _convert_workflow_event_to_agent_update(
+        self,
+        response_id: str,
+        event: WorkflowEvent,
+    ) -> AgentRunResponseUpdate | None:
+        """Convert a workflow event to an AgentRunResponseUpdate.
+
+        Only AgentRunUpdateEvent and RequestInfoEvent are processed and the rest
+        are not relevant. Returns None if the event is not relevant.
+        """
+        match event:
+            case AgentRunUpdateEvent(data=update):
+                # Direct pass-through of update in an agent streaming event
+                if update:
+                    return cast(AgentRunResponseUpdate, update)
+                return None
+
+            case RequestInfoEvent(request_id=request_id):
+                # Store the pending request for later correlation
+                self.pending_requests[request_id] = event
+
+                # Convert to function call content
+                # TODO(ekzhu): update this to FunctionApprovalRequestContent
+                # monitor: https://github.com/microsoft/agent-framework/issues/285
+                function_call = FunctionCallContent(
+                    call_id=request_id,
+                    name=self.REQUEST_INFO_FUNCTION_NAME,
+                    arguments=self.RequestInfoFunctionArgs(request_id=request_id, data=event.data).to_dict(),
+                )
+                return AgentRunResponseUpdate(
+                    contents=[function_call],
+                    role=Role.ASSISTANT,
+                    author_name=self.name,
+                    response_id=response_id,
+                    message_id=str(uuid.uuid4()),
+                    created_at=datetime.now().strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
+                )
+            case _:
+                # Ignore non-agent workflow events
+                pass
+        # We only care about the above two events and discard the rest.
+        return None
+
+    def _extract_function_responses(self, input_messages: list[ChatMessage]) -> dict[str, Any]:
+        """Extract function responses from input messages."""
+        function_responses: dict[str, Any] = {}
+        for message in input_messages:
+            for content in message.contents:
+                # TODO(ekzhu): update this to FunctionApprovalResponseContent
+                # monitor: https://github.com/microsoft/agent-framework/issues/285
+                if isinstance(content, FunctionResultContent):
+                    request_id = content.call_id
+                    # Check if we have a pending request for this call_id
+                    if request_id in self.pending_requests:
+                        response_data = content.result if hasattr(content, "result") else str(content)
+                        function_responses[request_id] = response_data
+                    elif bool(self.pending_requests):
+                        # Function result for unknown request when we have pending requests - this is an error
+                        raise AgentExecutionException(
+                            "Only FunctionResultContent for pending requests is allowed in input messages "
+                            "when there are pending requests."
+                        )
+                else:
+                    if bool(self.pending_requests):
+                        # Non-function content when we have pending requests - this is an error
+                        raise AgentExecutionException(
+                            "Only FunctionResultContent is allowed in input messages when there are pending requests."
+                        )
+        return function_responses
+
+    class _ResponseState(TypedDict):
+        """State for grouping response updates by message_id."""
+
+        by_msg: dict[str, list[AgentRunResponseUpdate]]
+        dangling: list[AgentRunResponseUpdate]
+
+    @staticmethod
+    def merge_updates(updates: list[AgentRunResponseUpdate], response_id: str) -> AgentRunResponse:
+        """Merge streaming updates into a single AgentRunResponse.
+
+        Behavior:
+        - Group updates by response_id; within each response_id, group by message_id and keep a dangling bucket for
+          updates without message_id.
+        - Convert each group (per message and dangling) into an intermediate AgentRunResponse via
+          AgentRunResponse.from_agent_run_response_updates, then sort by created_at and merge.
+        - Append messages from updates without any response_id at the end (global dangling), while aggregating metadata.
+
+        Args:
+            updates: The list of AgentRunResponseUpdate objects to merge.
+            response_id: The response identifier to set on the returned AgentRunResponse.
+
+        Returns:
+            An AgentRunResponse with messages in processing order and aggregated metadata.
+        """
+        # PHASE 1: GROUP UPDATES BY RESPONSE_ID AND MESSAGE_ID
+        states: dict[str, WorkflowAgent._ResponseState] = {}
+        global_dangling: list[AgentRunResponseUpdate] = []
+
+        for u in updates:
+            if u.response_id:
+                state = states.setdefault(u.response_id, {"by_msg": {}, "dangling": []})
+                by_msg = state["by_msg"]
+                dangling = state["dangling"]
+                if u.message_id:
+                    by_msg.setdefault(u.message_id, []).append(u)
+                else:
+                    dangling.append(u)
+            else:
+                global_dangling.append(u)
+
+        # HELPER FUNCTIONS
+        def _parse_dt(value: str | None) -> tuple[int, datetime | str | None]:
+            if not value:
+                return (1, None)
+            v = value
+            if v.endswith("Z"):
+                v = v[:-1] + "+00:00"
+            try:
+                return (0, datetime.fromisoformat(v))
+            except Exception:
+                return (0, v)
+
+        def _sum_usage(a: UsageDetails | None, b: UsageDetails | None) -> UsageDetails | None:
+            if a is None:
+                return b
+            if b is None:
+                return a
+            return a + b
+
+        def _merge_responses(current: AgentRunResponse | None, incoming: AgentRunResponse) -> AgentRunResponse:
+            if current is None:
+                return incoming
+            raw_list: list[object] = []
+
+            def _add_raw(value: object) -> None:
+                if isinstance(value, list):
+                    raw_list.extend(cast(list[object], value))
+                else:
+                    raw_list.append(value)
+
+            if current.raw_representation is not None:
+                _add_raw(current.raw_representation)
+            if incoming.raw_representation is not None:
+                _add_raw(incoming.raw_representation)
+            return AgentRunResponse(
+                messages=(current.messages or []) + (incoming.messages or []),
+                response_id=current.response_id or incoming.response_id,
+                created_at=incoming.created_at or current.created_at,
+                usage_details=_sum_usage(current.usage_details, incoming.usage_details),
+                raw_representation=raw_list if raw_list else None,
+                additional_properties=incoming.additional_properties or current.additional_properties,
+            )
+
+        # PHASE 2: CONVERT GROUPED UPDATES TO RESPONSES AND MERGE
+        final_messages: list[ChatMessage] = []
+        merged_usage: UsageDetails | None = None
+        latest_created_at: str | None = None
+        merged_additional_properties: dict[str, Any] | None = None
+        raw_representations: list[object] = []
+
+        for grouped_response_id in states:
+            state = states[grouped_response_id]
+            by_msg = state["by_msg"]
+            dangling = state["dangling"]
+
+            per_message_responses: list[AgentRunResponse] = []
+            for _, msg_updates in by_msg.items():
+                if msg_updates:
+                    per_message_responses.append(AgentRunResponse.from_agent_run_response_updates(msg_updates))
+            if dangling:
+                per_message_responses.append(AgentRunResponse.from_agent_run_response_updates(dangling))
+
+            per_message_responses.sort(key=lambda r: _parse_dt(r.created_at))
+
+            aggregated: AgentRunResponse | None = None
+            for resp in per_message_responses:
+                if resp.response_id and grouped_response_id and resp.response_id != grouped_response_id:
+                    resp.response_id = grouped_response_id
+                aggregated = _merge_responses(aggregated, resp)
+
+            if aggregated:
+                final_messages.extend(aggregated.messages)
+                if aggregated.usage_details:
+                    merged_usage = _sum_usage(merged_usage, aggregated.usage_details)
+                if aggregated.created_at and (
+                    not latest_created_at or _parse_dt(aggregated.created_at) > _parse_dt(latest_created_at)
+                ):
+                    latest_created_at = aggregated.created_at
+                if aggregated.additional_properties:
+                    if merged_additional_properties is None:
+                        merged_additional_properties = {}
+                    merged_additional_properties.update(aggregated.additional_properties)
+                raw_value = aggregated.raw_representation
+                if raw_value:
+                    cast_value = cast(object | list[object], raw_value)
+                    if isinstance(cast_value, list):
+                        raw_representations.extend(cast(list[object], cast_value))
+                    else:
+                        raw_representations.append(cast_value)
+
+        # PHASE 3: HANDLE GLOBAL DANGLING UPDATES (NO RESPONSE_ID)
+        if global_dangling:
+            flattened = AgentRunResponse.from_agent_run_response_updates(global_dangling)
+            final_messages.extend(flattened.messages)
+            if flattened.usage_details:
+                merged_usage = _sum_usage(merged_usage, flattened.usage_details)
+            if flattened.created_at and (
+                not latest_created_at or _parse_dt(flattened.created_at) > _parse_dt(latest_created_at)
+            ):
+                latest_created_at = flattened.created_at
+            if flattened.additional_properties:
+                if merged_additional_properties is None:
+                    merged_additional_properties = {}
+                merged_additional_properties.update(flattened.additional_properties)
+            flat_raw = flattened.raw_representation
+            if flat_raw:
+                cast_flat = cast(object | list[object], flat_raw)
+                if isinstance(cast_flat, list):
+                    raw_representations.extend(cast(list[object], cast_flat))
+                else:
+                    raw_representations.append(cast_flat)
+
+        # PHASE 4: CONSTRUCT FINAL RESPONSE WITH INPUT RESPONSE_ID
+        return AgentRunResponse(
+            messages=final_messages,
+            response_id=response_id,
+            created_at=latest_created_at,
+            usage_details=merged_usage,
+            raw_representation=raw_representations if raw_representations else None,
+            additional_properties=merged_additional_properties,
+        )

python/packages/core/agent_framework/_workflow/_checkpoint.py

@@ -0,0 +1,200 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import json
+import logging
+import os
+import uuid
+from collections.abc import Mapping
+from dataclasses import asdict, dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Protocol
+
+from ._const import DEFAULT_MAX_ITERATIONS
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(slots=True)
+class WorkflowCheckpoint:
+    """Represents a complete checkpoint of workflow state."""
+
+    checkpoint_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    workflow_id: str = ""
+    timestamp: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
+
+    # Core workflow state
+    messages: dict[str, list[dict[str, Any]]] = field(default_factory=dict)  # type: ignore[misc]
+    shared_state: dict[str, Any] = field(default_factory=dict)  # type: ignore[misc]
+    executor_states: dict[str, dict[str, Any]] = field(default_factory=dict)  # type: ignore[misc]
+
+    # Runtime state
+    iteration_count: int = 0
+    max_iterations: int = DEFAULT_MAX_ITERATIONS
+
+    # Metadata
+    metadata: dict[str, Any] = field(default_factory=dict)  # type: ignore[misc]
+    version: str = "1.0"
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: Mapping[str, Any]) -> "WorkflowCheckpoint":
+        return cls(**data)
+
+
+class CheckpointStorage(Protocol):
+    """Protocol for checkpoint storage backends."""
+
+    async def save_checkpoint(self, checkpoint: WorkflowCheckpoint) -> str:
+        """Save a checkpoint and return its ID."""
+        ...
+
+    async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None:
+        """Load a checkpoint by ID."""
+        ...
+
+    async def list_checkpoint_ids(self, workflow_id: str | None = None) -> list[str]:
+        """List checkpoint IDs. If workflow_id is provided, filter by that workflow."""
+        ...
+
+    async def list_checkpoints(self, workflow_id: str | None = None) -> list[WorkflowCheckpoint]:
+        """List checkpoint objects. If workflow_id is provided, filter by that workflow."""
+        ...
+
+    async def delete_checkpoint(self, checkpoint_id: str) -> bool:
+        """Delete a checkpoint by ID."""
+        ...
+
+
+class InMemoryCheckpointStorage:
+    """In-memory checkpoint storage for testing and development."""
+
+    def __init__(self) -> None:
+        """Initialize the memory storage."""
+        self._checkpoints: dict[str, WorkflowCheckpoint] = {}
+
+    async def save_checkpoint(self, checkpoint: WorkflowCheckpoint) -> str:
+        """Save a checkpoint and return its ID."""
+        self._checkpoints[checkpoint.checkpoint_id] = checkpoint
+        logger.debug(f"Saved checkpoint {checkpoint.checkpoint_id} to memory")
+        return checkpoint.checkpoint_id
+
+    async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None:
+        """Load a checkpoint by ID."""
+        checkpoint = self._checkpoints.get(checkpoint_id)
+        if checkpoint:
+            logger.debug(f"Loaded checkpoint {checkpoint_id} from memory")
+        return checkpoint
+
+    async def list_checkpoint_ids(self, workflow_id: str | None = None) -> list[str]:
+        """List checkpoint IDs. If workflow_id is provided, filter by that workflow."""
+        if workflow_id is None:
+            return list(self._checkpoints.keys())
+        return [cp.checkpoint_id for cp in self._checkpoints.values() if cp.workflow_id == workflow_id]
+
+    async def list_checkpoints(self, workflow_id: str | None = None) -> list[WorkflowCheckpoint]:
+        """List checkpoint objects. If workflow_id is provided, filter by that workflow."""
+        if workflow_id is None:
+            return list(self._checkpoints.values())
+        return [cp for cp in self._checkpoints.values() if cp.workflow_id == workflow_id]
+
+    async def delete_checkpoint(self, checkpoint_id: str) -> bool:
+        """Delete a checkpoint by ID."""
+        if checkpoint_id in self._checkpoints:
+            del self._checkpoints[checkpoint_id]
+            logger.debug(f"Deleted checkpoint {checkpoint_id} from memory")
+            return True
+        return False
+
+
+class FileCheckpointStorage:
+    """File-based checkpoint storage for persistence."""
+
+    def __init__(self, storage_path: str | Path):
+        """Initialize the file storage."""
+        self.storage_path = Path(storage_path)
+        self.storage_path.mkdir(parents=True, exist_ok=True)
+        logger.info(f"Initialized file checkpoint storage at {self.storage_path}")
+
+    async def save_checkpoint(self, checkpoint: WorkflowCheckpoint) -> str:
+        """Save a checkpoint and return its ID."""
+        file_path = self.storage_path / f"{checkpoint.checkpoint_id}.json"
+        checkpoint_dict = asdict(checkpoint)
+
+        def _write_atomic() -> None:
+            tmp_path = file_path.with_suffix(".json.tmp")
+            with open(tmp_path, "w") as f:
+                json.dump(checkpoint_dict, f, indent=2, ensure_ascii=False)
+            os.replace(tmp_path, file_path)
+
+        await asyncio.to_thread(_write_atomic)
+
+        logger.info(f"Saved checkpoint {checkpoint.checkpoint_id} to {file_path}")
+        return checkpoint.checkpoint_id
+
+    async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None:
+        """Load a checkpoint by ID."""
+        file_path = self.storage_path / f"{checkpoint_id}.json"
+
+        if not file_path.exists():
+            return None
+
+        def _read() -> dict[str, Any]:
+            with open(file_path) as f:
+                return json.load(f)  # type: ignore[no-any-return]
+
+        checkpoint_dict = await asyncio.to_thread(_read)
+
+        checkpoint = WorkflowCheckpoint(**checkpoint_dict)
+        logger.info(f"Loaded checkpoint {checkpoint_id} from {file_path}")
+        return checkpoint
+
+    async def list_checkpoint_ids(self, workflow_id: str | None = None) -> list[str]:
+        """List checkpoint IDs. If workflow_id is provided, filter by that workflow."""
+
+        def _list_ids() -> list[str]:
+            checkpoint_ids: list[str] = []
+            for file_path in self.storage_path.glob("*.json"):
+                try:
+                    with open(file_path) as f:
+                        data = json.load(f)
+                    if workflow_id is None or data.get("workflow_id") == workflow_id:
+                        checkpoint_ids.append(data.get("checkpoint_id", file_path.stem))
+                except Exception as e:
+                    logger.warning(f"Failed to read checkpoint file {file_path}: {e}")
+            return checkpoint_ids
+
+        return await asyncio.to_thread(_list_ids)
+
+    async def list_checkpoints(self, workflow_id: str | None = None) -> list[WorkflowCheckpoint]:
+        """List checkpoint objects. If workflow_id is provided, filter by that workflow."""
+
+        def _list_checkpoints() -> list[WorkflowCheckpoint]:
+            checkpoints: list[WorkflowCheckpoint] = []
+            for file_path in self.storage_path.glob("*.json"):
+                try:
+                    with open(file_path) as f:
+                        data = json.load(f)
+                    if workflow_id is None or data.get("workflow_id") == workflow_id:
+                        checkpoints.append(WorkflowCheckpoint.from_dict(data))
+                except Exception as e:
+                    logger.warning(f"Failed to read checkpoint file {file_path}: {e}")
+            return checkpoints
+
+        return await asyncio.to_thread(_list_checkpoints)
+
+    async def delete_checkpoint(self, checkpoint_id: str) -> bool:
+        """Delete a checkpoint by ID."""
+        file_path = self.storage_path / f"{checkpoint_id}.json"
+
+        def _delete() -> bool:
+            if file_path.exists():
+                file_path.unlink()
+                logger.info(f"Deleted checkpoint {checkpoint_id} from {file_path}")
+                return True
+            return False
+
+        return await asyncio.to_thread(_delete)

python/packages/core/agent_framework/_workflow/_concurrent.py

@@ -0,0 +1,324 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import inspect
+import logging
+from collections.abc import Callable, Sequence
+from typing import Any
+
+from typing_extensions import Never
+
+from agent_framework import AgentProtocol, ChatMessage, Role
+
+from ._checkpoint import CheckpointStorage
+from ._executor import AgentExecutorRequest, AgentExecutorResponse, Executor, handler
+from ._workflow import Workflow, WorkflowBuilder
+from ._workflow_context import WorkflowContext
+
+logger = logging.getLogger(__name__)
+
+"""Concurrent builder for agent-only fan-out/fan-in workflows.
+
+This module provides a high-level, agent-focused API to quickly assemble a
+parallel workflow with:
+- a default dispatcher that broadcasts the input to all agent participants
+- a default aggregator that combines all agent conversations and completes the workflow
+
+Notes:
+- Participants should be AgentProtocol instances or Executors.
+- A custom aggregator can be provided as:
+  - an Executor instance (it should handle list[AgentExecutorResponse],
+    yield output), or
+  - a callback function with signature:
+        def cb(results: list[AgentExecutorResponse]) -> Any | None
+        def cb(results: list[AgentExecutorResponse], ctx: WorkflowContext) -> Any | None
+    The callback is wrapped in _CallbackAggregator.
+    If the callback returns a non-None value, _CallbackAggregator yields that as output.
+    If it returns None, the callback may have already yielded an output via ctx, so no further action is taken.
+"""
+
+
+class _DispatchToAllParticipants(Executor):
+    """Broadcasts input to all downstream participants (via fan-out edges)."""
+
+    @handler
+    async def from_request(self, request: AgentExecutorRequest, ctx: WorkflowContext[AgentExecutorRequest]) -> None:
+        # No explicit target: edge routing delivers to all connected participants.
+        await ctx.send_message(request)
+
+    @handler
+    async def from_str(self, prompt: str, ctx: WorkflowContext[AgentExecutorRequest]) -> None:
+        request = AgentExecutorRequest(messages=[ChatMessage(Role.USER, text=prompt)], should_respond=True)
+        await ctx.send_message(request)
+
+    @handler
+    async def from_message(self, message: ChatMessage, ctx: WorkflowContext[AgentExecutorRequest]) -> None:  # type: ignore[name-defined]
+        request = AgentExecutorRequest(messages=[message], should_respond=True)
+        await ctx.send_message(request)
+
+    @handler
+    async def from_messages(self, messages: list[ChatMessage], ctx: WorkflowContext[AgentExecutorRequest]) -> None:  # type: ignore[name-defined]
+        request = AgentExecutorRequest(messages=list(messages), should_respond=True)
+        await ctx.send_message(request)
+
+
+class _AggregateAgentConversations(Executor):
+    """Aggregates agent responses and completes with combined ChatMessages.
+
+    Emits a list[ChatMessage] shaped as:
+      [ single_user_prompt?, agent1_final_assistant, agent2_final_assistant, ... ]
+
+    - Extracts a single user prompt (first user message seen across results).
+    - For each result, selects the final assistant message (prefers agent_run_response.messages).
+    - Avoids duplicating the same user message per agent.
+    """
+
+    @handler
+    async def aggregate(
+        self, results: list[AgentExecutorResponse], ctx: WorkflowContext[Never, list[ChatMessage]]
+    ) -> None:
+        if not results:
+            logger.error("Concurrent aggregator received empty results list")
+            raise ValueError("Aggregation failed: no results provided")
+
+        def _is_role(msg: Any, role: Role) -> bool:
+            r = getattr(msg, "role", None)
+            if r is None:
+                return False
+            # Normalize both r and role to lowercase strings for comparison
+            r_str = str(r).lower() if isinstance(r, str) or hasattr(r, "__str__") else r
+            role_str = getattr(role, "value", None)
+            if role_str is None:
+                role_str = str(role)
+            role_str = role_str.lower()
+            return r_str == role_str
+
+        prompt_message: ChatMessage | None = None
+        assistant_replies: list[ChatMessage] = []
+
+        for r in results:
+            resp_messages = list(getattr(r.agent_run_response, "messages", []) or [])
+            conv = r.full_conversation if r.full_conversation is not None else resp_messages
+
+            logger.debug(
+                f"Aggregating executor {getattr(r, 'executor_id', '<unknown>')}: "
+                f"{len(resp_messages)} response msgs, {len(conv)} conversation msgs"
+            )
+
+            # Capture a single user prompt (first encountered across any conversation)
+            if prompt_message is None:
+                found_user = next((m for m in conv if _is_role(m, Role.USER)), None)
+                if found_user is not None:
+                    prompt_message = found_user
+
+            # Pick the final assistant message from the response; fallback to conversation search
+            final_assistant = next((m for m in reversed(resp_messages) if _is_role(m, Role.ASSISTANT)), None)
+            if final_assistant is None:
+                final_assistant = next((m for m in reversed(conv) if _is_role(m, Role.ASSISTANT)), None)
+
+            if final_assistant is not None:
+                assistant_replies.append(final_assistant)
+            else:
+                logger.warning(
+                    f"No assistant reply found for executor {getattr(r, 'executor_id', '<unknown>')}; skipping"
+                )
+
+        if not assistant_replies:
+            logger.error(f"Aggregation failed: no assistant replies found across {len(results)} results")
+            raise RuntimeError("Aggregation failed: no assistant replies found")
+
+        output: list[ChatMessage] = []
+        if prompt_message is not None:
+            output.append(prompt_message)
+        else:
+            logger.warning("No user prompt found in any conversation; emitting assistants only")
+        output.extend(assistant_replies)
+
+        await ctx.yield_output(output)
+
+
+class _CallbackAggregator(Executor):
+    """Wraps a Python callback as an aggregator.
+
+    Accepts either an async or sync callback with one of the signatures:
+      - (results: list[AgentExecutorResponse]) -> Any | None
+      - (results: list[AgentExecutorResponse], ctx: WorkflowContext[Any]) -> Any | None
+
+    Notes:
+    - Async callbacks are awaited directly.
+    - Sync callbacks are executed via asyncio.to_thread to avoid blocking the event loop.
+    - If the callback returns a non-None value, it is yielded as an output.
+    """
+
+    def __init__(self, callback: Callable[..., Any], id: str | None = None) -> None:
+        derived_id = getattr(callback, "__name__", "") or ""
+        if not derived_id or derived_id == "<lambda>":
+            derived_id = f"{type(self).__name__}_unnamed"
+        super().__init__(id or derived_id)
+        self._callback = callback
+        self._param_count = len(inspect.signature(callback).parameters)
+
+    @handler
+    async def aggregate(self, results: list[AgentExecutorResponse], ctx: WorkflowContext[Never, Any]) -> None:
+        # Call according to provided signature, always non-blocking for sync callbacks
+        if self._param_count >= 2:
+            if inspect.iscoroutinefunction(self._callback):
+                ret = await self._callback(results, ctx)  # type: ignore[misc]
+            else:
+                ret = await asyncio.to_thread(self._callback, results, ctx)
+        else:
+            if inspect.iscoroutinefunction(self._callback):
+                ret = await self._callback(results)  # type: ignore[misc]
+            else:
+                ret = await asyncio.to_thread(self._callback, results)
+
+        # If the callback returned a value, finalize the workflow with it
+        if ret is not None:
+            await ctx.yield_output(ret)
+
+
+class ConcurrentBuilder:
+    r"""High-level builder for concurrent agent workflows.
+
+    - `participants([...])` accepts a list of AgentProtocol (recommended) or Executor.
+    - `build()` wires: dispatcher -> fan-out -> participants -> fan-in -> aggregator.
+    - `with_custom_aggregator(...)` overrides the default aggregator with an Executor or callback.
+
+    Usage:
+    ```python
+    from agent_framework import ConcurrentBuilder
+
+    # Minimal: use default aggregator (returns list[ChatMessage])
+    workflow = ConcurrentBuilder().participants([agent1, agent2, agent3]).build()
+
+
+    # Custom aggregator via callback (sync or async). The callback receives
+    # list[AgentExecutorResponse] and its return value becomes the workflow's output.
+    def summarize(results):
+        return " | ".join(r.agent_run_response.messages[-1].text for r in results)
+
+
+    workflow = ConcurrentBuilder().participants([agent1, agent2, agent3]).with_custom_aggregator(summarize).build()
+
+
+    # Enable checkpoint persistence so runs can resume
+    workflow = ConcurrentBuilder().participants([agent1, agent2, agent3]).with_checkpointing(storage).build()
+    ```
+    """
+
+    def __init__(self) -> None:
+        self._participants: list[AgentProtocol | Executor] = []
+        self._aggregator: Executor | None = None
+        self._checkpoint_storage: CheckpointStorage | None = None
+
+    def participants(self, participants: Sequence[AgentProtocol | Executor]) -> "ConcurrentBuilder":
+        r"""Define the parallel participants for this concurrent workflow.
+
+        Accepts AgentProtocol instances (e.g., created by a chat client) or Executor
+        instances. Each participant is wired as a parallel branch using fan-out edges
+        from an internal dispatcher.
+
+        Raises:
+            ValueError: if `participants` is empty or contains duplicates
+            TypeError: if any entry is not AgentProtocol or Executor
+
+        Example:
+        ```python
+        wf = ConcurrentBuilder().participants([researcher_agent, marketer_agent, legal_agent]).build()
+
+        # Mixing agent(s) and executor(s) is supported
+        wf2 = ConcurrentBuilder().participants([researcher_agent, my_custom_executor]).build()
+        ```
+        """
+        if not participants:
+            raise ValueError("participants cannot be empty")
+
+        # Defensive duplicate detection
+        seen_agent_ids: set[int] = set()
+        seen_executor_ids: set[str] = set()
+        for p in participants:
+            if isinstance(p, Executor):
+                if p.id in seen_executor_ids:
+                    raise ValueError(f"Duplicate executor participant detected: id '{p.id}'")
+                seen_executor_ids.add(p.id)
+            elif isinstance(p, AgentProtocol):
+                pid = id(p)
+                if pid in seen_agent_ids:
+                    raise ValueError("Duplicate agent participant detected (same agent instance provided twice)")
+                seen_agent_ids.add(pid)
+            else:
+                raise TypeError(f"participants must be AgentProtocol or Executor instances; got {type(p).__name__}")
+
+        self._participants = list(participants)
+        return self
+
+    def with_aggregator(self, aggregator: Executor | Callable[..., Any]) -> "ConcurrentBuilder":
+        r"""Override the default aggregator with an Executor or a callback.
+
+        - Executor: must handle `list[AgentExecutorResponse]` and
+            yield output using `ctx.yield_output(...)` and add a
+          output and the workflow becomes idle.
+        - Callback: sync or async callable with one of the signatures:
+          `(results: list[AgentExecutorResponse]) -> Any | None` or
+          `(results: list[AgentExecutorResponse], ctx: WorkflowContext) -> Any | None`.
+          If the callback returns a non-None value, it becomes the workflow's output.
+
+        Example:
+        ```python
+        # Callback-based aggregator (string result)
+        async def summarize(results):
+            return " | ".join(r.agent_run_response.messages[-1].text for r in results)
+
+
+        wf = ConcurrentBuilder().participants([a1, a2, a3]).with_custom_aggregator(summarize).build()
+        ```
+        """
+        if isinstance(aggregator, Executor):
+            self._aggregator = aggregator
+        elif callable(aggregator):
+            self._aggregator = _CallbackAggregator(aggregator)
+        else:
+            raise TypeError("aggregator must be an Executor or a callable")
+        return self
+
+    def with_checkpointing(self, checkpoint_storage: CheckpointStorage) -> "ConcurrentBuilder":
+        """Enable checkpoint persistence using the provided storage backend."""
+        self._checkpoint_storage = checkpoint_storage
+        return self
+
+    def build(self) -> Workflow:
+        r"""Build and validate the concurrent workflow.
+
+        Wiring pattern:
+        - Dispatcher (internal) fans out the input to all `participants`
+        - Fan-in aggregator collects `AgentExecutorResponse` objects
+        - Aggregator yields output and the workflow becomes idle. The output is either:
+          - list[ChatMessage] (default aggregator: one user + one assistant per agent)
+          - custom payload from the provided callback/executor
+
+        Returns:
+            Workflow: a ready-to-run workflow instance
+
+        Raises:
+            ValueError: if no participants were defined
+
+        Example:
+        ```python
+        workflow = ConcurrentBuilder().participants([agent1, agent2]).build()
+        ```
+        """
+        if not self._participants:
+            raise ValueError("No participants provided. Call .participants([...]) first.")
+
+        dispatcher = _DispatchToAllParticipants(id="dispatcher")
+        aggregator = self._aggregator or _AggregateAgentConversations(id="aggregator")
+
+        builder = WorkflowBuilder()
+        builder.set_start_executor(dispatcher)
+        builder.add_fan_out_edges(dispatcher, list(self._participants))
+        builder.add_fan_in_edges(list(self._participants), aggregator)
+
+        if self._checkpoint_storage is not None:
+            builder = builder.with_checkpointing(self._checkpoint_storage)
+
+        return builder.build()

python/packages/core/agent_framework/_workflow/_const.py

@@ -0,0 +1,3 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+DEFAULT_MAX_ITERATIONS = 100  # Default maximum iterations for workflow execution.

python/packages/core/agent_framework/_workflow/_edge.py

@@ -0,0 +1,867 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import logging
+import uuid
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass, field
+from typing import Any, ClassVar
+
+from ._executor import Executor
+from ._model_utils import DictConvertible, encode_value
+
+logger = logging.getLogger(__name__)
+
+
+def _extract_function_name(func: Callable[..., Any]) -> str:
+    """Map a Python callable to a concise, human-focused identifier.
+
+    The workflow graph persists references to callables by recording only an
+    identifier. This helper inspects standard callable metadata and picks a
+    stable value so that serialized representations remain intelligible when
+    they are later rendered in logs or reconstructed during deserialization.
+
+    Example:
+        ```python
+        def threshold(value: float) -> bool:
+            return value > 0.5
+
+
+        assert _extract_function_name(threshold) == "threshold"
+        ```
+    """
+    if hasattr(func, "__name__"):
+        name = func.__name__
+        return name if name != "<lambda>" else "<lambda>"
+    return "<callable>"
+
+
+def _missing_callable(name: str) -> Callable[..., Any]:
+    """Create a defensive placeholder for callables that cannot be restored.
+
+    When a workflow is deserialized in an environment that lacks the original
+    Python callable, we install a proxy that fails loudly. Surfacing the error
+    at invocation time preserves a clean separation between I/O concerns and
+    runtime execution, while making it obvious which callable needs to be
+    re-registered.
+
+    Example:
+        ```python
+        guard = _missing_callable("transform_price")
+        try:
+            guard()
+        except RuntimeError as exc:
+            assert "transform_price" in str(exc)
+        ```
+    """
+
+    def _raise(*_: Any, **__: Any) -> Any:
+        raise RuntimeError(f"Callable '{name}' is unavailable after serialization")
+
+    return _raise
+
+
+@dataclass(init=False)
+class Edge(DictConvertible):
+    """Model a directed, optionally-conditional hand-off between two executors.
+
+    Each `Edge` captures the minimal metadata required to move a message from
+    one executor to another inside the workflow graph. It optionally embeds a
+    boolean predicate that decides if the edge should be taken at runtime. By
+    serialising the edge down to primitives we can reconstruct the topology of
+    a workflow irrespective of the original Python process.
+
+    Example:
+        ```python
+        edge = Edge(source_id="ingest", target_id="score", condition=lambda payload: payload["ready"])
+        assert edge.should_route({"ready": True}) is True
+        assert edge.should_route({"ready": False}) is False
+        ```
+    """
+
+    ID_SEPARATOR: ClassVar[str] = "->"
+
+    source_id: str
+    target_id: str
+    condition_name: str | None
+    _condition: Callable[[Any], bool] | None = field(default=None, repr=False, compare=False)
+
+    def __init__(
+        self,
+        source_id: str,
+        target_id: str,
+        condition: Callable[[Any], bool] | None = None,
+        *,
+        condition_name: str | None = None,
+    ) -> None:
+        """Initialize a fully-specified edge between two workflow executors.
+
+        Parameters
+        ----------
+        source_id:
+            Canonical identifier of the upstream executor instance.
+        target_id:
+            Canonical identifier of the downstream executor instance.
+        condition:
+            Optional predicate that receives the message payload and returns
+            `True` when the edge should be traversed. When omitted, the edge is
+            considered unconditionally active.
+        condition_name:
+            Optional override that pins a human-friendly name for the condition
+            when the callable cannot be introspected (for example after
+            deserialization).
+
+        Example:
+            ```python
+            edge = Edge("fetch", "parse", condition=lambda data: data.is_valid)
+            assert edge.source_id == "fetch"
+            assert edge.target_id == "parse"
+            ```
+        """
+        if not source_id:
+            raise ValueError("Edge source_id must be a non-empty string")
+        if not target_id:
+            raise ValueError("Edge target_id must be a non-empty string")
+        self.source_id = source_id
+        self.target_id = target_id
+        self._condition = condition
+        self.condition_name = _extract_function_name(condition) if condition is not None else condition_name
+
+    @property
+    def id(self) -> str:
+        """Return the stable identifier used to reference this edge.
+
+        The identifier combines the source and target executor identifiers with
+        a deterministic separator. This allows other graph structures such as
+        adjacency lists or visualisations to refer to an edge without carrying
+        the full object.
+
+        Example:
+            ```python
+            edge = Edge("reader", "writer")
+            assert edge.id == "reader->writer"
+            ```
+        """
+        return f"{self.source_id}{self.ID_SEPARATOR}{self.target_id}"
+
+    def should_route(self, data: Any) -> bool:
+        """Evaluate the edge predicate against an incoming payload.
+
+        When the edge was defined without an explicit predicate the method
+        returns `True`, signalling an unconditional routing rule. Otherwise the
+        user-supplied callable decides whether the message should proceed along
+        this edge. Any exception raised by the callable is deliberately allowed
+        to surface to the caller to avoid masking logic bugs.
+
+        Example:
+            ```python
+            edge = Edge("stage1", "stage2", condition=lambda payload: payload["score"] > 0.8)
+            assert edge.should_route({"score": 0.9}) is True
+            assert edge.should_route({"score": 0.4}) is False
+            ```
+        """
+        if self._condition is None:
+            return True
+        return self._condition(data)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Produce a JSON-serialisable view of the edge metadata.
+
+        The representation includes the source and target executor identifiers
+        plus the condition name when it is known. Serialisation intentionally
+        omits the live callable to keep payloads transport-friendly.
+
+        Example:
+            ```python
+            edge = Edge("reader", "writer", condition=lambda payload: payload["ok"])
+            snapshot = edge.to_dict()
+            assert snapshot == {"source_id": "reader", "target_id": "writer", "condition_name": "<lambda>"}
+            ```
+        """
+        payload = {"source_id": self.source_id, "target_id": self.target_id}
+        if self.condition_name is not None:
+            payload["condition_name"] = self.condition_name
+        return payload
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "Edge":
+        """Reconstruct an `Edge` from its serialised dictionary form.
+
+        The deserialised edge will lack the executable predicate because we do
+        not attempt to hydrate Python callables from storage. Instead, the
+        stored `condition_name` is preserved so that downstream consumers can
+        detect missing callables and re-register them where appropriate.
+
+        Example:
+            ```python
+            payload = {"source_id": "reader", "target_id": "writer", "condition_name": "is_ready"}
+            edge = Edge.from_dict(payload)
+            assert edge.source_id == "reader"
+            assert edge.condition_name == "is_ready"
+            ```
+        """
+        return cls(
+            source_id=data["source_id"],
+            target_id=data["target_id"],
+            condition=None,
+            condition_name=data.get("condition_name"),
+        )
+
+
+@dataclass
+class Case:
+    """Runtime wrapper combining a switch-case predicate with its target.
+
+    Each `Case` couples a boolean predicate with the executor that should
+    handle the message when the predicate evaluates to `True`. The runtime
+    keeps this lightweight container separate from the serialisable
+    `SwitchCaseEdgeGroupCase` so that execution can operate with live callables
+    without polluting persisted state.
+
+    Example:
+        ```python
+        class JsonExecutor(Executor):
+            def __init__(self) -> None:
+                super().__init__(id="json", defer_discovery=True)
+
+
+        processor = JsonExecutor()
+        case = Case(condition=lambda payload: payload["kind"] == "json", target=processor)
+        assert case.target.id == "json"
+        ```
+    """
+
+    condition: Callable[[Any], bool]
+    target: Executor
+
+
+@dataclass
+class Default:
+    """Runtime representation of the default branch in a switch-case group.
+
+    The default branch is invoked only when no other case predicates match. In
+    practice it is guaranteed to exist so that routing never produces an empty
+    target.
+
+    Example:
+        ```python
+        class DeadLetterExecutor(Executor):
+            def __init__(self) -> None:
+                super().__init__(id="dead_letter", defer_discovery=True)
+
+
+        fallback = Default(target=DeadLetterExecutor())
+        assert fallback.target.id == "dead_letter"
+        ```
+    """
+
+    target: Executor
+
+
+@dataclass(init=False)
+class EdgeGroup(DictConvertible):
+    """Bundle edges that share a common routing semantics under a single id.
+
+    The workflow runtime manipulates `EdgeGroup` instances rather than raw
+    edges so it can reason about higher-order routing behaviours such as
+    fan-out, fan-in, switch-case, and other graph patterns. The base class stores the
+    identifying information and handles serialisation duties so specialised
+    groups need only maintain their additional state.
+
+    Example:
+        ```python
+        group = EdgeGroup([Edge("source", "sink")])
+        assert group.source_executor_ids == ["source"]
+        ```
+    """
+
+    id: str
+    type: str
+    edges: list[Edge]
+
+    from builtins import type as builtin_type
+
+    _TYPE_REGISTRY: ClassVar[dict[str, builtin_type["EdgeGroup"]]] = {}
+
+    def __init__(
+        self,
+        edges: Sequence[Edge] | None = None,
+        *,
+        id: str | None = None,
+        type: str | None = None,
+    ) -> None:
+        """Construct an edge group shell around a set of `Edge` instances.
+
+        Parameters
+        ----------
+        edges:
+            Sequence of edges that participate in this group. When omitted we
+            start from an empty list so subclasses can append later.
+        id:
+            Stable identifier for the group. Defaults to a random UUID so
+            serialised graphs remain uniquely addressable.
+        type:
+            Logical discriminator used to recover the appropriate subclass when
+            de-serialising.
+
+        Example:
+            ```python
+            edges = [Edge("validate", "persist")]
+            group = EdgeGroup(edges, id="stage", type="Custom")
+            assert group.to_dict()["type"] == "Custom"
+            ```
+        """
+        self.id = id or f"{self.__class__.__name__}/{uuid.uuid4()}"
+        self.type = type or self.__class__.__name__
+        self.edges = list(edges) if edges is not None else []
+
+    @property
+    def source_executor_ids(self) -> list[str]:
+        """Return the deduplicated list of upstream executor ids.
+
+        The property preserves order-of-first-appearance so the caller can rely
+        on deterministic iteration when reconstructing graph topology.
+
+        Example:
+            ```python
+            group = EdgeGroup([Edge("read", "write"), Edge("read", "archive")])
+            assert group.source_executor_ids == ["read"]
+            ```
+        """
+        return list(dict.fromkeys(edge.source_id for edge in self.edges))
+
+    @property
+    def target_executor_ids(self) -> list[str]:
+        """Return the ordered, deduplicated list of downstream executor ids.
+
+        Example:
+            ```python
+            group = EdgeGroup([Edge("read", "write"), Edge("read", "archive")])
+            assert group.target_executor_ids == ["write", "archive"]
+            ```
+        """
+        return list(dict.fromkeys(edge.target_id for edge in self.edges))
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialise the group metadata and contained edges into primitives.
+
+        The payload captures each edge through its own `to_dict` call, enabling
+        round-tripping through formats such as JSON without leaking Python
+        objects.
+
+        Example:
+            ```python
+            group = EdgeGroup([Edge("read", "write")])
+            snapshot = group.to_dict()
+            assert snapshot["edges"][0]["source_id"] == "read"
+            ```
+        """
+        return {
+            "id": self.id,
+            "type": self.type,
+            "edges": [edge.to_dict() for edge in self.edges],
+        }
+
+    @classmethod
+    def register(cls, subclass: builtin_type["EdgeGroup"]) -> builtin_type["EdgeGroup"]:
+        """Register a subclass so deserialisation can recover the right type.
+
+        Registration is typically performed via the decorator syntax applied to
+        each concrete edge group. The registry stores classes by their
+        `__name__`, which must therefore remain stable across versions when
+        persisted workflows are in circulation.
+
+        Example:
+            ```python
+            @EdgeGroup.register
+            class CustomGroup(EdgeGroup):
+                pass
+
+
+            assert EdgeGroup._TYPE_REGISTRY["CustomGroup"] is CustomGroup
+            ```
+        """
+        cls._TYPE_REGISTRY[subclass.__name__] = subclass
+        return subclass
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "EdgeGroup":
+        """Hydrate the correct `EdgeGroup` subclass from serialised state.
+
+        The method inspects the `type` field, allocates the corresponding class
+        without executing subclass `__init__`, and then manually restores any
+        subtype-specific attributes. This keeps deserialisation deterministic
+        even for complex group types that configure additional runtime
+        callables.
+
+        Example:
+            ```python
+            payload = {"type": "EdgeGroup", "edges": [{"source_id": "a", "target_id": "b"}]}
+            group = EdgeGroup.from_dict(payload)
+            assert isinstance(group, EdgeGroup)
+            ```
+        """
+        group_type = data.get("type", "EdgeGroup")
+        target_cls = cls._TYPE_REGISTRY.get(group_type, EdgeGroup)
+        edges = [Edge.from_dict(entry) for entry in data.get("edges", [])]
+
+        obj = target_cls.__new__(target_cls)  # type: ignore[misc]
+        EdgeGroup.__init__(obj, edges=edges, id=data.get("id"), type=group_type)
+
+        # Handle FanOutEdgeGroup-specific attributes
+        if isinstance(obj, FanOutEdgeGroup):
+            obj.selection_func_name = data.get("selection_func_name")  # type: ignore[attr-defined]
+            obj._selection_func = (  # type: ignore[attr-defined]
+                None
+                if obj.selection_func_name is None  # type: ignore[attr-defined]
+                else _missing_callable(obj.selection_func_name)  # type: ignore[attr-defined]
+            )
+            obj._target_ids = [edge.target_id for edge in obj.edges]  # type: ignore[attr-defined]
+
+        # Handle SwitchCaseEdgeGroup-specific attributes
+        if isinstance(obj, SwitchCaseEdgeGroup):
+            cases_payload = data.get("cases", [])
+            restored_cases: list[SwitchCaseEdgeGroupCase | SwitchCaseEdgeGroupDefault] = []
+            for case_data in cases_payload:
+                case_type = case_data.get("type")
+                if case_type == "Default":
+                    restored_cases.append(SwitchCaseEdgeGroupDefault.from_dict(case_data))
+                else:
+                    restored_cases.append(SwitchCaseEdgeGroupCase.from_dict(case_data))
+            obj.cases = restored_cases  # type: ignore[attr-defined]
+            obj._selection_func = _missing_callable("switch_case_selection")  # type: ignore[attr-defined]
+
+        return obj
+
+
+@EdgeGroup.register
+@dataclass(init=False)
+class SingleEdgeGroup(EdgeGroup):
+    """Convenience wrapper for a solitary edge, keeping the group API uniform."""
+
+    def __init__(
+        self,
+        source_id: str,
+        target_id: str,
+        condition: Callable[[Any], bool] | None = None,
+        *,
+        id: str | None = None,
+    ) -> None:
+        """Create a one-to-one edge group between two executors.
+
+        Example:
+            ```python
+            group = SingleEdgeGroup("ingest", "validate")
+            assert group.edges[0].source_id == "ingest"
+            ```
+        """
+        edge = Edge(source_id=source_id, target_id=target_id, condition=condition)
+        super().__init__([edge], id=id, type=self.__class__.__name__)
+
+
+@EdgeGroup.register
+@dataclass(init=False)
+class FanOutEdgeGroup(EdgeGroup):
+    """Represent a broadcast-style edge group with optional selection logic.
+
+    A fan-out forwards a message produced by a single source executor to one
+    or more downstream executors. At runtime we may further narrow the targets
+    by executing a `selection_func` that inspects the payload and returns the
+    subset of ids that should receive the message.
+    """
+
+    selection_func_name: str | None
+    _selection_func: Callable[[Any, list[str]], list[str]] | None
+    _target_ids: list[str]
+
+    def __init__(
+        self,
+        source_id: str,
+        target_ids: Sequence[str],
+        selection_func: Callable[[Any, list[str]], list[str]] | None = None,
+        *,
+        selection_func_name: str | None = None,
+        id: str | None = None,
+    ) -> None:
+        """Create a fan-out mapping from a single source to many targets.
+
+        Parameters
+        ----------
+        source_id:
+            Identifier of the upstream executor broadcasting the message.
+        target_ids:
+            Ordered set of downstream executor identifiers that may receive the
+            message. At least two targets are required to preserve the fan-out
+            semantics.
+        selection_func:
+            Optional callable that returns the subset of `target_ids` that
+            should be active for a given payload. The callable receives the
+            original message plus a copy of all configured target ids.
+        selection_func_name:
+            Static identifier used when persisting the fan-out. Needed when the
+            callable cannot be introspected or is unavailable during
+            deserialisation.
+        id:
+            Stable identifier for the group; defaults to an autogenerated UUID.
+
+        Example:
+            ```python
+            def choose_targets(message: dict[str, Any], available: list[str]) -> list[str]:
+                return [target for target in available if message.get(target)]
+
+
+            group = FanOutEdgeGroup("sensor", ["db", "cache"], selection_func=choose_targets)
+            assert group.selection_func is choose_targets
+            ```
+        """
+        if len(target_ids) <= 1:
+            raise ValueError("FanOutEdgeGroup must contain at least two targets.")
+
+        edges = [Edge(source_id=source_id, target_id=target) for target in target_ids]
+        super().__init__(edges, id=id, type=self.__class__.__name__)
+
+        self._target_ids = list(target_ids)
+        self._selection_func = selection_func
+        self.selection_func_name = (
+            _extract_function_name(selection_func) if selection_func is not None else selection_func_name
+        )
+
+    @property
+    def target_ids(self) -> list[str]:
+        """Return a shallow copy of the configured downstream executor ids.
+
+        The list is defensively copied to prevent callers from mutating the
+        internal state while still providing deterministic ordering.
+
+        Example:
+            ```python
+            group = FanOutEdgeGroup("node", ["alpha", "beta"])
+            assert group.target_ids == ["alpha", "beta"]
+            ```
+        """
+        return list(self._target_ids)
+
+    @property
+    def selection_func(self) -> Callable[[Any, list[str]], list[str]] | None:
+        """Expose the runtime callable used to select active fan-out targets.
+
+        When no selection function was supplied the property returns `None`,
+        signalling that all targets must receive the payload.
+
+        Example:
+            ```python
+            group = FanOutEdgeGroup("source", ["x", "y"], selection_func=None)
+            assert group.selection_func is None
+            ```
+        """
+        return self._selection_func
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialise the fan-out group while preserving selection metadata.
+
+        In addition to the base `EdgeGroup` payload we embed the human-friendly
+        name of the selection function. The callable itself is not persisted.
+
+        Example:
+            ```python
+            group = FanOutEdgeGroup("source", ["a", "b"], selection_func=lambda *_: ["a"])
+            snapshot = group.to_dict()
+            assert snapshot["selection_func_name"] == "<lambda>"
+            ```
+        """
+        payload = super().to_dict()
+        payload["selection_func_name"] = self.selection_func_name
+        return payload
+
+
+@EdgeGroup.register
+@dataclass(init=False)
+class FanInEdgeGroup(EdgeGroup):
+    """Represent a converging set of edges that feed a single downstream executor.
+
+    Fan-in groups are typically used when multiple upstream stages independently
+    produce messages that should all arrive at the same downstream processor.
+    """
+
+    def __init__(self, source_ids: Sequence[str], target_id: str, *, id: str | None = None) -> None:
+        """Build a fan-in mapping that merges several sources into one target.
+
+        Parameters
+        ----------
+        source_ids:
+            Sequence of upstream executor identifiers contributing messages.
+        target_id:
+            Downstream executor that receives every message emitted by the
+            sources.
+        id:
+            Optional explicit identifier for the edge group.
+
+        Example:
+            ```python
+            group = FanInEdgeGroup(["parser", "enricher"], target_id="writer")
+            assert group.to_dict()["edges"][0]["target_id"] == "writer"
+            ```
+        """
+        if len(source_ids) <= 1:
+            raise ValueError("FanInEdgeGroup must contain at least two sources.")
+
+        edges = [Edge(source_id=source, target_id=target_id) for source in source_ids]
+        super().__init__(edges, id=id, type=self.__class__.__name__)
+
+
+@dataclass(init=False)
+class SwitchCaseEdgeGroupCase(DictConvertible):
+    """Persistable description of a single conditional branch in a switch-case.
+
+    Unlike the runtime `Case` object this serialisable variant stores only the
+    target identifier and a descriptive name for the predicate. When the
+    underlying callable is unavailable during deserialisation we substitute a
+    proxy placeholder that fails loudly, ensuring the missing dependency is
+    immediately visible.
+    """
+
+    target_id: str
+    condition_name: str | None
+    type: str
+    _condition: Callable[[Any], bool] = field(repr=False, compare=False)
+
+    def __init__(
+        self,
+        condition: Callable[[Any], bool] | None,
+        target_id: str,
+        *,
+        condition_name: str | None = None,
+    ) -> None:
+        """Record the routing metadata for a conditional case branch.
+
+        Parameters
+        ----------
+        condition:
+            Optional live predicate. When omitted we fall back to a placeholder
+            that raises at runtime to highlight missing registrations.
+        target_id:
+            Identifier of the executor that should handle messages when the
+            predicate succeeds.
+        condition_name:
+            Human-friendly label for the predicate used for diagnostics and
+            on-disk persistence.
+
+        Example:
+            ```python
+            case = SwitchCaseEdgeGroupCase(lambda payload: payload["type"] == "csv", target_id="csv_handler")
+            assert case.condition_name == "<lambda>"
+            ```
+        """
+        if not target_id:
+            raise ValueError("SwitchCaseEdgeGroupCase requires a target_id")
+        self.target_id = target_id
+        self.type = "Case"
+        if condition is not None:
+            self._condition = condition
+            self.condition_name = _extract_function_name(condition)
+        else:
+            safe_name = condition_name or "<missing_condition>"
+            self._condition = _missing_callable(safe_name)
+            self.condition_name = condition_name
+
+    @property
+    def condition(self) -> Callable[[Any], bool]:
+        """Return the predicate associated with this case.
+
+        The placeholder installed during deserialisation raises a
+        `RuntimeError` when invoked so that workflow authors are forced to
+        provide the missing callable explicitly.
+
+        Example:
+            ```python
+            case = SwitchCaseEdgeGroupCase(None, target_id="missing", condition_name="needs_registration")
+            guard = case.condition
+            try:
+                guard({})
+            except RuntimeError:
+                pass
+            ```
+        """
+        return self._condition
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialise the case metadata without the executable predicate.
+
+        Example:
+            ```python
+            case = SwitchCaseEdgeGroupCase(lambda _: True, target_id="handler")
+            assert case.to_dict()["target_id"] == "handler"
+            ```
+        """
+        payload = {"target_id": self.target_id, "type": self.type}
+        if self.condition_name is not None:
+            payload["condition_name"] = self.condition_name
+        return payload
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "SwitchCaseEdgeGroupCase":
+        """Instantiate a case from its serialised dictionary payload.
+
+        Example:
+            ```python
+            payload = {"target_id": "handler", "condition_name": "is_ready"}
+            case = SwitchCaseEdgeGroupCase.from_dict(payload)
+            assert case.target_id == "handler"
+            ```
+        """
+        return cls(
+            condition=None,
+            target_id=data["target_id"],
+            condition_name=data.get("condition_name"),
+        )
+
+
+@dataclass(init=False)
+class SwitchCaseEdgeGroupDefault(DictConvertible):
+    """Persistable descriptor for the fallback branch of a switch-case group.
+
+    The default branch is guaranteed to exist and is invoked when every other
+    case predicate fails to match the payload.
+    """
+
+    target_id: str
+    type: str
+
+    def __init__(self, target_id: str) -> None:
+        """Point the default branch toward the given executor identifier.
+
+        Example:
+            ```python
+            fallback = SwitchCaseEdgeGroupDefault(target_id="dead_letter")
+            assert fallback.target_id == "dead_letter"
+            ```
+        """
+        if not target_id:
+            raise ValueError("SwitchCaseEdgeGroupDefault requires a target_id")
+        self.target_id = target_id
+        self.type = "Default"
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialise the default branch metadata for persistence or logging.
+
+        Example:
+            ```python
+            fallback = SwitchCaseEdgeGroupDefault("dead_letter")
+            assert fallback.to_dict()["type"] == "Default"
+            ```
+        """
+        return {"target_id": self.target_id, "type": self.type}
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "SwitchCaseEdgeGroupDefault":
+        """Recreate the default branch from its persisted form.
+
+        Example:
+            ```python
+            payload = {"target_id": "dead_letter", "type": "Default"}
+            fallback = SwitchCaseEdgeGroupDefault.from_dict(payload)
+            assert fallback.target_id == "dead_letter"
+            ```
+        """
+        return cls(target_id=data["target_id"])
+
+
+@EdgeGroup.register
+@dataclass(init=False)
+class SwitchCaseEdgeGroup(FanOutEdgeGroup):
+    """Fan-out variant that mimics a traditional switch/case control flow.
+
+    Each case inspects the message payload and decides whether it should handle
+    the message. Exactly one case-or the default branch-returns a target at
+    runtime, preserving single-dispatch semantics.
+    """
+
+    cases: list[SwitchCaseEdgeGroupCase | SwitchCaseEdgeGroupDefault]
+
+    def __init__(
+        self,
+        source_id: str,
+        cases: Sequence[SwitchCaseEdgeGroupCase | SwitchCaseEdgeGroupDefault],
+        *,
+        id: str | None = None,
+    ) -> None:
+        """Configure a switch/case routing structure for a single source executor.
+
+        Parameters
+        ----------
+        source_id:
+            Identifier of the executor producing the message to be routed.
+        cases:
+            Ordered sequence of case descriptors concluding with a
+            `SwitchCaseEdgeGroupDefault`. Ordering matters because the runtime
+            evaluates each branch sequentially until one matches.
+        id:
+            Optional explicit identifier for the edge group.
+
+        Example:
+            ```python
+            cases = [
+                SwitchCaseEdgeGroupCase(lambda payload: payload["kind"] == "csv", target_id="process_csv"),
+                SwitchCaseEdgeGroupDefault(target_id="process_default"),
+            ]
+            group = SwitchCaseEdgeGroup("router", cases)
+            encoded = group.to_dict()
+            assert encoded["cases"][0]["type"] == "Case"
+            ```
+        """
+        if len(cases) < 2:
+            raise ValueError("SwitchCaseEdgeGroup must contain at least two cases (including the default case).")
+
+        default_cases = [case for case in cases if isinstance(case, SwitchCaseEdgeGroupDefault)]
+        if len(default_cases) != 1:
+            raise ValueError("SwitchCaseEdgeGroup must contain exactly one default case.")
+
+        if not isinstance(cases[-1], SwitchCaseEdgeGroupDefault):
+            logger.warning(
+                "Default case in the switch-case edge group is not the last case. "
+                "This may result in unexpected behavior."
+            )
+
+        def selection_func(message: Any, targets: list[str]) -> list[str]:
+            for case in cases:
+                if isinstance(case, SwitchCaseEdgeGroupDefault):
+                    return [case.target_id]
+                try:
+                    if case.condition(message):
+                        return [case.target_id]
+                except Exception as exc:  # pragma: no cover - defensive logging
+                    logger.warning("Error evaluating condition for case %s: %s", case.target_id, exc)
+            raise RuntimeError("No matching case found in SwitchCaseEdgeGroup")
+
+        target_ids = [case.target_id for case in cases]
+        # Call FanOutEdgeGroup constructor directly to avoid type checking issues
+        edges = [Edge(source_id=source_id, target_id=target) for target in target_ids]
+        EdgeGroup.__init__(self, edges, id=id, type=self.__class__.__name__)
+
+        # Initialize FanOutEdgeGroup-specific attributes
+        self._target_ids = list(target_ids)  # type: ignore[attr-defined]
+        self._selection_func = selection_func  # type: ignore[attr-defined]
+        self.selection_func_name = None  # type: ignore[attr-defined]
+        self.cases = list(cases)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialise the switch-case group, capturing all case descriptors.
+
+        Each case is converted using `encode_value` to respect dataclass
+        semantics as well as any nested serialisable structures.
+
+        Example:
+            ```python
+            group = SwitchCaseEdgeGroup(
+                "router",
+                [
+                    SwitchCaseEdgeGroupCase(lambda _: True, target_id="handler"),
+                    SwitchCaseEdgeGroupDefault(target_id="fallback"),
+                ],
+            )
+            snapshot = group.to_dict()
+            assert len(snapshot["cases"]) == 2
+            ```
+        """
+        payload = super().to_dict()
+        payload["cases"] = [encode_value(case) for case in self.cases]
+        return payload

python/packages/core/agent_framework/_workflow/_edge_runner.py

@@ -0,0 +1,385 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import logging
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from collections.abc import Callable
+from typing import Any, cast
+
+from ..observability import EdgeGroupDeliveryStatus, OtelAttr, create_edge_group_processing_span
+from ._edge import Edge, EdgeGroup, FanInEdgeGroup, FanOutEdgeGroup, SingleEdgeGroup, SwitchCaseEdgeGroup
+from ._executor import Executor
+from ._runner_context import Message, RunnerContext
+from ._shared_state import SharedState
+
+logger = logging.getLogger(__name__)
+
+
+class EdgeRunner(ABC):
+    """Abstract base class for edge runners that handle message delivery."""
+
+    def __init__(self, edge_group: EdgeGroup, executors: dict[str, Executor]) -> None:
+        """Initialize the edge runner with an edge group and executor map.
+
+        Args:
+            edge_group: The edge group to run.
+            executors: Map of executor IDs to executor instances.
+        """
+        self._edge_group = edge_group
+        self._executors = executors
+
+    @abstractmethod
+    async def send_message(self, message: Message, shared_state: SharedState, ctx: RunnerContext) -> bool:
+        """Send a message through the edge group.
+
+        Args:
+            message: The message to send.
+            shared_state: The shared state to use for holding data.
+            ctx: The context for the runner.
+
+        Returns:
+            bool: True if the message was processed successfully,
+                False if the target executor cannot handle the message.
+        """
+        raise NotImplementedError
+
+    def _can_handle(self, executor_id: str, message_data: Any) -> bool:
+        """Check if an executor can handle the given message data."""
+        if executor_id not in self._executors:
+            return False
+        return self._executors[executor_id].can_handle(message_data)
+
+    async def _execute_on_target(
+        self,
+        target_id: str,
+        source_ids: list[str],
+        message: Message,
+        shared_state: SharedState,
+        ctx: RunnerContext,
+    ) -> None:
+        """Execute a message on a target executor with trace context."""
+        if target_id not in self._executors:
+            raise RuntimeError(f"Target executor {target_id} not found.")
+
+        target_executor = self._executors[target_id]
+
+        # Execute with trace context parameters
+        await target_executor.execute(
+            message.data,
+            source_ids,  # source_executor_ids
+            shared_state,  # shared_state
+            ctx,  # runner_context
+            trace_contexts=message.trace_contexts,  # Pass trace contexts
+            source_span_ids=message.source_span_ids,  # Pass source span IDs for linking
+        )
+
+
+class SingleEdgeRunner(EdgeRunner):
+    """Runner for single edge groups."""
+
+    def __init__(self, edge_group: SingleEdgeGroup, executors: dict[str, Executor]) -> None:
+        super().__init__(edge_group, executors)
+        self._edge = edge_group.edges[0]
+
+    async def send_message(self, message: Message, shared_state: SharedState, ctx: RunnerContext) -> bool:
+        """Send a message through the single edge."""
+        should_execute = False
+        target_id = None
+        source_id = None
+        with create_edge_group_processing_span(
+            self._edge_group.__class__.__name__,
+            edge_group_id=self._edge_group.id,
+            message_source_id=message.source_id,
+            message_target_id=message.target_id,
+            source_trace_contexts=message.trace_contexts,
+            source_span_ids=message.source_span_ids,
+        ) as span:
+            try:
+                if message.target_id and message.target_id != self._edge.target_id:
+                    span.set_attributes({
+                        OtelAttr.EDGE_GROUP_DELIVERED: False,
+                        OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TARGET_MISMATCH.value,
+                    })
+                    return False
+
+                if self._can_handle(self._edge.target_id, message.data):
+                    if self._edge.should_route(message.data):
+                        span.set_attributes({
+                            OtelAttr.EDGE_GROUP_DELIVERED: True,
+                            OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DELIVERED.value,
+                        })
+                        should_execute = True
+                        target_id = self._edge.target_id
+                        source_id = self._edge.source_id
+                    else:
+                        span.set_attributes({
+                            OtelAttr.EDGE_GROUP_DELIVERED: False,
+                            OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_CONDITION_FALSE.value,
+                        })
+                        # Return True here because message was processed, just condition failed
+                        return True
+                else:
+                    span.set_attributes({
+                        OtelAttr.EDGE_GROUP_DELIVERED: False,
+                        OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH.value,
+                    })
+                    return False
+            except Exception as e:
+                span.set_attributes({
+                    OtelAttr.EDGE_GROUP_DELIVERED: False,
+                    OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.EXCEPTION.value,
+                })
+                raise e
+
+        # Execute outside the span
+        if should_execute and target_id and source_id:
+            await self._execute_on_target(target_id, [source_id], message, shared_state, ctx)
+            return True
+
+        return False
+
+
+class FanOutEdgeRunner(EdgeRunner):
+    """Runner for fan-out edge groups."""
+
+    def __init__(self, edge_group: FanOutEdgeGroup, executors: dict[str, Executor]) -> None:
+        super().__init__(edge_group, executors)
+        self._edges = edge_group.edges
+        self._target_ids = edge_group.target_executor_ids
+        self._target_map = {edge.target_id: edge for edge in self._edges}
+        self._selection_func = cast(
+            Callable[[Any, list[str]], list[str]] | None, getattr(edge_group, "selection_func", None)
+        )
+
+    async def send_message(self, message: Message, shared_state: SharedState, ctx: RunnerContext) -> bool:
+        """Send a message through all edges in the fan-out edge group."""
+        deliverable_edges = []
+        single_target_edge = None
+        # Process routing logic within span
+        with create_edge_group_processing_span(
+            self._edge_group.__class__.__name__,
+            edge_group_id=self._edge_group.id,
+            message_source_id=message.source_id,
+            message_target_id=message.target_id,
+            source_trace_contexts=message.trace_contexts,
+            source_span_ids=message.source_span_ids,
+        ) as span:
+            try:
+                selection_results = (
+                    self._selection_func(message.data, self._target_ids) if self._selection_func else self._target_ids
+                )
+                if not self._validate_selection_result(selection_results):
+                    span.set_attributes({
+                        OtelAttr.EDGE_GROUP_DELIVERED: False,
+                        OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.EXCEPTION.value,
+                    })
+                    raise RuntimeError(
+                        f"Invalid selection result: {selection_results}. "
+                        f"Expected selections to be a subset of valid target executor IDs: {self._target_ids}."
+                    )
+
+                if message.target_id:
+                    # If the target ID is specified and the selection result contains it, send the message to that edge
+                    if message.target_id in selection_results:
+                        edge = self._target_map.get(message.target_id)
+                        if edge and self._can_handle(edge.target_id, message.data):
+                            if edge.should_route(message.data):
+                                span.set_attributes({
+                                    OtelAttr.EDGE_GROUP_DELIVERED: True,
+                                    OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DELIVERED.value,
+                                })
+                                single_target_edge = edge
+                            else:
+                                span.set_attributes({
+                                    OtelAttr.EDGE_GROUP_DELIVERED: False,
+                                    OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_CONDITION_FALSE.value,  # noqa: E501
+                                })
+                                # For targeted messages with condition failure, return True (message was processed)
+                                return True
+                        else:
+                            span.set_attributes({
+                                OtelAttr.EDGE_GROUP_DELIVERED: False,
+                                OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH.value,  # noqa: E501
+                            })
+                            # For targeted messages that can't be handled, return False
+                            return False
+                    else:
+                        span.set_attributes({
+                            OtelAttr.EDGE_GROUP_DELIVERED: False,
+                            OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TARGET_MISMATCH.value,
+                        })
+                        # For targeted messages not in selection, return False
+                        return False
+                else:
+                    # If no target ID, send the message to the selected targets
+                    for target_id in selection_results:
+                        edge = self._target_map[target_id]
+                        if self._can_handle(edge.target_id, message.data) and edge.should_route(message.data):
+                            deliverable_edges.append(edge)
+
+                    if len(deliverable_edges) > 0:
+                        span.set_attributes({
+                            OtelAttr.EDGE_GROUP_DELIVERED: True,
+                            OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DELIVERED.value,
+                        })
+                    else:
+                        span.set_attributes({
+                            OtelAttr.EDGE_GROUP_DELIVERED: False,
+                            OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH.value,
+                        })
+
+            except Exception as e:
+                span.set_attributes({
+                    OtelAttr.EDGE_GROUP_DELIVERED: False,
+                    OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.EXCEPTION.value,
+                })
+                raise e
+
+        # Execute outside the span
+        if single_target_edge:
+            await self._execute_on_target(
+                single_target_edge.target_id, [single_target_edge.source_id], message, shared_state, ctx
+            )
+            return True
+
+        if deliverable_edges:
+
+            async def send_to_edge(edge: Edge) -> bool:
+                await self._execute_on_target(edge.target_id, [edge.source_id], message, shared_state, ctx)
+                return True
+
+            tasks = [send_to_edge(edge) for edge in deliverable_edges]
+            results = await asyncio.gather(*tasks)
+            return any(results)
+
+        # If we get here, it's a broadcast message with no deliverable edges
+        return False
+
+    def _validate_selection_result(self, selection_results: list[str]) -> bool:
+        """Validate the selection results to ensure all IDs are valid target executor IDs."""
+        return all(result in self._target_ids for result in selection_results)
+
+
+class FanInEdgeRunner(EdgeRunner):
+    """Runner for fan-in edge groups."""
+
+    def __init__(self, edge_group: FanInEdgeGroup, executors: dict[str, Executor]) -> None:
+        super().__init__(edge_group, executors)
+        self._edges = edge_group.edges
+        # Buffer to hold messages before sending them to the target executor
+        # Key is the source executor ID, value is a list of messages
+        self._buffer: dict[str, list[Message]] = defaultdict(list)
+
+    async def send_message(self, message: Message, shared_state: SharedState, ctx: RunnerContext) -> bool:
+        """Send a message through all edges in the fan-in edge group."""
+        execution_data: dict[str, Any] | None = None
+        with create_edge_group_processing_span(
+            self._edge_group.__class__.__name__,
+            edge_group_id=self._edge_group.id,
+            message_source_id=message.source_id,
+            message_target_id=message.target_id,
+            source_trace_contexts=message.trace_contexts,
+            source_span_ids=message.source_span_ids,
+        ) as span:
+            try:
+                if message.target_id and message.target_id != self._edges[0].target_id:
+                    span.set_attributes({
+                        OtelAttr.EDGE_GROUP_DELIVERED: False,
+                        OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TARGET_MISMATCH.value,
+                    })
+                    return False
+
+                # Check if target can handle list of message data (fan-in aggregates multiple messages)
+                if self._can_handle(self._edges[0].target_id, [message.data]):
+                    # If the edge can handle the data, buffer the message
+                    self._buffer[message.source_id].append(message)
+                    span.set_attributes({
+                        OtelAttr.EDGE_GROUP_DELIVERED: True,
+                        OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.BUFFERED.value,
+                    })
+                else:
+                    # If the edge cannot handle the data, return False
+                    span.set_attributes({
+                        OtelAttr.EDGE_GROUP_DELIVERED: False,
+                        OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH.value,
+                    })
+                    return False
+
+                if self._is_ready_to_send():
+                    # If all edges in the group have data, prepare for execution
+                    messages_to_send = [msg for edge in self._edges for msg in self._buffer[edge.source_id]]
+                    self._buffer.clear()
+                    # Send aggregated data to target
+                    aggregated_data = [msg.data for msg in messages_to_send]
+
+                    # Collect all trace contexts and source span IDs for fan-in linking
+                    trace_contexts = [msg.trace_context for msg in messages_to_send if msg.trace_context]
+                    source_span_ids = [msg.source_span_id for msg in messages_to_send if msg.source_span_id]
+
+                    # Create a new Message object for the aggregated data
+                    aggregated_message = Message(
+                        data=aggregated_data,
+                        source_id=self._edge_group.__class__.__name__,  # This won't be used in self._execute_on_target.
+                        trace_contexts=trace_contexts,
+                        source_span_ids=source_span_ids,
+                    )
+                    span.set_attributes({
+                        OtelAttr.EDGE_GROUP_DELIVERED: True,
+                        OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DELIVERED.value,
+                    })
+
+                    # Store execution data for later
+                    execution_data = {
+                        "target_id": self._edges[0].target_id,
+                        "source_ids": [edge.source_id for edge in self._edges],
+                        "message": aggregated_message,
+                    }
+
+            except Exception as e:
+                span.set_attributes({
+                    OtelAttr.EDGE_GROUP_DELIVERED: False,
+                    OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.EXCEPTION.value,
+                })
+                raise e
+
+        # Execute outside the span if needed
+        if execution_data:
+            await self._execute_on_target(
+                execution_data["target_id"], execution_data["source_ids"], execution_data["message"], shared_state, ctx
+            )
+            return True
+
+        return True  # Return True for buffered messages (waiting for more)
+
+    def _is_ready_to_send(self) -> bool:
+        """Check if all edges in the group have data to send."""
+        return all(self._buffer[edge.source_id] for edge in self._edges)
+
+
+class SwitchCaseEdgeRunner(FanOutEdgeRunner):
+    """Runner for switch-case edge groups (inherits from FanOutEdgeRunner)."""
+
+    def __init__(self, edge_group: SwitchCaseEdgeGroup, executors: dict[str, Executor]) -> None:
+        super().__init__(edge_group, executors)
+
+
+def create_edge_runner(edge_group: EdgeGroup, executors: dict[str, Executor]) -> EdgeRunner:
+    """Factory function to create the appropriate edge runner for an edge group.
+
+    Args:
+        edge_group: The edge group to create a runner for.
+        executors: Map of executor IDs to executor instances.
+
+    Returns:
+        The appropriate EdgeRunner instance.
+    """
+    if isinstance(edge_group, SingleEdgeGroup):
+        return SingleEdgeRunner(edge_group, executors)
+    if isinstance(edge_group, SwitchCaseEdgeGroup):
+        return SwitchCaseEdgeRunner(edge_group, executors)
+    if isinstance(edge_group, FanOutEdgeGroup):
+        return FanOutEdgeRunner(edge_group, executors)
+    if isinstance(edge_group, FanInEdgeGroup):
+        return FanInEdgeRunner(edge_group, executors)
+    raise ValueError(f"Unsupported edge group type: {type(edge_group)}")

python/packages/core/agent_framework/_workflow/_events.py

@@ -0,0 +1,330 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import traceback as _traceback
+from collections.abc import Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass
+from enum import Enum
+from typing import TYPE_CHECKING, Any, TypeAlias
+
+from agent_framework import AgentRunResponse, AgentRunResponseUpdate
+
+if TYPE_CHECKING:
+    from ._executor import RequestInfoMessage
+
+
+class WorkflowEventSource(str, Enum):
+    """Identifies whether a workflow event came from the framework or an executor.
+
+    Use `FRAMEWORK` for events emitted by built-in orchestration paths—even when the
+    code that raises them lives in runner-related modules—and `EXECUTOR` for events
+    surfaced by developer-provided executor implementations.
+    """
+
+    FRAMEWORK = "FRAMEWORK"  # Framework-owned orchestration, regardless of module location
+    EXECUTOR = "EXECUTOR"  # User-supplied executor code and callbacks
+
+
+_event_origin_context: ContextVar[WorkflowEventSource] = ContextVar(
+    "workflow_event_origin", default=WorkflowEventSource.EXECUTOR
+)
+
+
+def _current_event_origin() -> WorkflowEventSource:
+    """Return the origin to associate with newly created workflow events."""
+    return _event_origin_context.get()
+
+
+@contextmanager
+def _framework_event_origin() -> Iterator[None]:  # pyright: ignore[reportUnusedFunction]
+    """Temporarily mark subsequently created events as originating from the framework (internal)."""
+    token = _event_origin_context.set(WorkflowEventSource.FRAMEWORK)
+    try:
+        yield
+    finally:
+        _event_origin_context.reset(token)
+
+
+class WorkflowEvent:
+    """Base class for workflow events."""
+
+    def __init__(self, data: Any | None = None):
+        """Initialize the workflow event with optional data."""
+        self.data = data
+        self.origin = _current_event_origin()
+
+    def __repr__(self) -> str:
+        """Return a string representation of the workflow event."""
+        data_repr = self.data if self.data is not None else "None"
+        return f"{self.__class__.__name__}(origin={self.origin}, data={data_repr})"
+
+
+class WorkflowStartedEvent(WorkflowEvent):
+    """Built-in lifecycle event emitted when a workflow run begins."""
+
+    ...
+
+
+class WorkflowWarningEvent(WorkflowEvent):
+    """Executor-origin event signaling a warning surfaced by user code."""
+
+    def __init__(self, data: str):
+        """Initialize the workflow warning event with optional data and warning message."""
+        super().__init__(data)
+
+    def __repr__(self) -> str:
+        """Return a string representation of the workflow warning event."""
+        return f"{self.__class__.__name__}(message={self.data}, origin={self.origin})"
+
+
+class WorkflowErrorEvent(WorkflowEvent):
+    """Executor-origin event signaling an error surfaced by user code."""
+
+    def __init__(self, data: Exception):
+        """Initialize the workflow error event with optional data and error message."""
+        super().__init__(data)
+
+    def __repr__(self) -> str:
+        """Return a string representation of the workflow error event."""
+        return f"{self.__class__.__name__}(exception={self.data}, origin={self.origin})"
+
+
+class WorkflowRunState(str, Enum):
+    """Run-level state of a workflow execution.
+
+    Semantics:
+      - STARTED: Run has been initiated and the workflow context has been created.
+        This is an initial state before any meaningful work is performed. In this
+        codebase we emit a dedicated `WorkflowStartedEvent` for telemetry, and
+        typically advance the status directly to `IN_PROGRESS`. Consumers may
+        still rely on `STARTED` for state machines that need an explicit pre-work
+        phase.
+
+      - IN_PROGRESS: The workflow is actively executing (e.g., the initial
+        message has been delivered to the start executor or a superstep is
+        running). This status is emitted at the beginning of a run and can be
+        followed by other statuses as the run progresses.
+
+      - IN_PROGRESS_PENDING_REQUESTS: Active execution while one or more
+        request-for-information operations are outstanding. New work may still
+        be scheduled while requests are in flight.
+
+      - IDLE: The workflow is quiescent with no outstanding requests and no more
+        work to do. This is the normal terminal state for workflows that have
+        finished executing, potentially having produced outputs along the way.
+
+      - IDLE_WITH_PENDING_REQUESTS: The workflow is paused awaiting external
+        input (e.g., emitted a `RequestInfoEvent`). This is a non-terminal
+        state; the workflow can resume when responses are supplied.
+
+      - FAILED: Terminal state indicating an error surfaced. Accompanied by a
+        `WorkflowFailedEvent` with structured error details.
+
+      - CANCELLED: Terminal state indicating the run was cancelled by a caller
+        or orchestrator. Not currently emitted by default runner paths but
+        included for integrators/orchestrators that support cancellation.
+    """
+
+    STARTED = "STARTED"  # Explicit pre-work phase (rarely emitted as status; see note above)
+    IN_PROGRESS = "IN_PROGRESS"  # Active execution is underway
+    IN_PROGRESS_PENDING_REQUESTS = "IN_PROGRESS_PENDING_REQUESTS"  # Active execution with outstanding requests
+    IDLE = "IDLE"  # No active work and no outstanding requests
+    IDLE_WITH_PENDING_REQUESTS = "IDLE_WITH_PENDING_REQUESTS"  # Paused awaiting external responses
+    FAILED = "FAILED"  # Finished with an error
+    CANCELLED = "CANCELLED"  # Finished due to cancellation
+
+
+class WorkflowStatusEvent(WorkflowEvent):
+    """Built-in lifecycle event emitted for workflow run state transitions."""
+
+    def __init__(
+        self,
+        state: WorkflowRunState,
+        data: Any | None = None,
+    ):
+        """Initialize the workflow status event with a new state and optional data.
+
+        Args:
+            state: The new state of the workflow run.
+            data: Optional additional data associated with the state change.
+        """
+        super().__init__(data)
+        self.state = state
+
+    def __repr__(self) -> str:  # pragma: no cover - representation only
+        return f"{self.__class__.__name__}(state={self.state}, data={self.data!r}, origin={self.origin})"
+
+
+@dataclass
+class WorkflowErrorDetails:
+    """Structured error information to surface in error events/results."""
+
+    error_type: str
+    message: str
+    traceback: str | None = None
+    executor_id: str | None = None
+    extra: dict[str, Any] | None = None
+
+    @classmethod
+    def from_exception(
+        cls,
+        exc: BaseException,
+        *,
+        executor_id: str | None = None,
+        extra: dict[str, Any] | None = None,
+    ) -> "WorkflowErrorDetails":
+        tb = None
+        try:
+            tb = "".join(_traceback.format_exception(type(exc), exc, exc.__traceback__))
+        except Exception:
+            tb = None
+        return cls(
+            error_type=exc.__class__.__name__,
+            message=str(exc),
+            traceback=tb,
+            executor_id=executor_id,
+            extra=extra,
+        )
+
+
+class WorkflowFailedEvent(WorkflowEvent):
+    """Built-in lifecycle event emitted when a workflow run terminates with an error."""
+
+    def __init__(
+        self,
+        details: WorkflowErrorDetails,
+        data: Any | None = None,
+    ):
+        super().__init__(data)
+        self.details = details
+
+    def __repr__(self) -> str:  # pragma: no cover - representation only
+        return f"{self.__class__.__name__}(details={self.details}, data={self.data!r}, origin={self.origin})"
+
+
+class RequestInfoEvent(WorkflowEvent):
+    """Event triggered when a workflow executor requests external information."""
+
+    def __init__(
+        self,
+        request_id: str,
+        source_executor_id: str,
+        request_type: type,
+        request_data: "RequestInfoMessage",
+    ):
+        """Initialize the request info event.
+
+        Args:
+            request_id: Unique identifier for the request.
+            source_executor_id: ID of the executor that made the request.
+            request_type: Type of the request (e.g., a specific data type).
+            request_data: The data associated with the request.
+        """
+        super().__init__(request_data)
+        self.request_id = request_id
+        self.source_executor_id = source_executor_id
+        self.request_type = request_type
+
+    def __repr__(self) -> str:
+        """Return a string representation of the request info event."""
+        return (
+            f"{self.__class__.__name__}("
+            f"request_id={self.request_id}, "
+            f"source_executor_id={self.source_executor_id}, "
+            f"request_type={self.request_type.__name__}, "
+            f"data={self.data})"
+        )
+
+
+class WorkflowOutputEvent(WorkflowEvent):
+    """Event triggered when a workflow executor yields output."""
+
+    def __init__(
+        self,
+        data: Any,
+        source_executor_id: str,
+    ):
+        """Initialize the workflow output event.
+
+        Args:
+            data: The output yielded by the executor.
+            source_executor_id: ID of the executor that yielded the output.
+        """
+        super().__init__(data)
+        self.source_executor_id = source_executor_id
+
+    def __repr__(self) -> str:
+        """Return a string representation of the workflow output event."""
+        return f"{self.__class__.__name__}(data={self.data}, source_executor_id={self.source_executor_id})"
+
+
+class ExecutorEvent(WorkflowEvent):
+    """Base class for executor events."""
+
+    def __init__(self, executor_id: str, data: Any | None = None):
+        """Initialize the executor event with an executor ID and optional data."""
+        super().__init__(data)
+        self.executor_id = executor_id
+
+    def __repr__(self) -> str:
+        """Return a string representation of the executor event."""
+        return f"{self.__class__.__name__}(executor_id={self.executor_id}, data={self.data})"
+
+
+class ExecutorInvokedEvent(ExecutorEvent):
+    """Event triggered when an executor handler is invoked."""
+
+    def __repr__(self) -> str:
+        """Return a string representation of the executor handler invoke event."""
+        return f"{self.__class__.__name__}(executor_id={self.executor_id}, data={self.data})"
+
+
+class ExecutorCompletedEvent(ExecutorEvent):
+    """Event triggered when an executor handler is completed."""
+
+    def __repr__(self) -> str:
+        """Return a string representation of the executor handler complete event."""
+        return f"{self.__class__.__name__}(executor_id={self.executor_id}, data={self.data})"
+
+
+class ExecutorFailedEvent(ExecutorEvent):
+    """Event triggered when an executor handler raises an error."""
+
+    def __init__(
+        self,
+        executor_id: str,
+        details: WorkflowErrorDetails,
+    ):
+        super().__init__(executor_id, details)
+        self.details = details
+
+    def __repr__(self) -> str:  # pragma: no cover - representation only
+        return f"{self.__class__.__name__}(executor_id={self.executor_id}, details={self.details})"
+
+
+class AgentRunUpdateEvent(ExecutorEvent):
+    """Event triggered when an agent is streaming messages."""
+
+    def __init__(self, executor_id: str, data: AgentRunResponseUpdate | None = None):
+        """Initialize the agent streaming event."""
+        super().__init__(executor_id, data)
+
+    def __repr__(self) -> str:
+        """Return a string representation of the agent streaming event."""
+        return f"{self.__class__.__name__}(executor_id={self.executor_id}, messages={self.data})"
+
+
+class AgentRunEvent(ExecutorEvent):
+    """Event triggered when an agent run is completed."""
+
+    def __init__(self, executor_id: str, data: AgentRunResponse | None = None):
+        """Initialize the agent run event."""
+        super().__init__(executor_id, data)
+
+    def __repr__(self) -> str:
+        """Return a string representation of the agent run event."""
+        return f"{self.__class__.__name__}(executor_id={self.executor_id}, data={self.data})"
+
+
+WorkflowLifecycleEvent: TypeAlias = WorkflowStartedEvent | WorkflowStatusEvent | WorkflowFailedEvent

python/packages/core/agent_framework/_workflow/_function_executor.py

@@ -0,0 +1,166 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Function-based Executor and decorator utilities.
+
+This module provides:
+- FunctionExecutor: an Executor subclass that wraps a user-defined function
+  with signature (message) or (message, ctx: WorkflowContext[T]). Both sync and async functions are supported.
+  Synchronous functions are executed in a thread pool using asyncio.to_thread() to avoid blocking the event loop.
+- executor decorator: converts such a function into a ready-to-use Executor instance
+  with proper type validation and handler registration.
+"""
+
+import asyncio
+from collections.abc import Awaitable, Callable
+from typing import Any, overload
+
+from ._executor import Executor
+from ._workflow_context import WorkflowContext, validate_function_signature
+
+
+class FunctionExecutor(Executor):
+    """Executor that wraps a user-defined function.
+
+    This executor allows users to define simple functions (both sync and async) and use them
+    as workflow executors without needing to create full executor classes.
+
+    Synchronous functions are executed in a thread pool using asyncio.to_thread() to avoid
+    blocking the event loop.
+    """
+
+    @staticmethod
+    def _validate_function(func: Callable[..., Any]) -> tuple[type, Any, list[type[Any]], list[type[Any]]]:
+        """Validate that the function has the correct signature for an executor.
+
+        Args:
+            func: The function to validate (can be sync or async)
+
+        Returns:
+            Tuple of (message_type, ctx_annotation, output_types, workflow_output_types)
+
+        Raises:
+            ValueError: If the function signature is incorrect
+        """
+        return validate_function_signature(func, "Function")
+
+    def __init__(self, func: Callable[..., Any], id: str | None = None):
+        """Initialize the FunctionExecutor with a user-defined function.
+
+        Args:
+            func: The function to wrap as an executor (can be sync or async)
+            id: Optional executor ID. If None, uses the function name.
+        """
+        # Validate function signature and extract types
+        message_type, ctx_annotation, output_types, workflow_output_types = self._validate_function(func)
+
+        # Determine if function has WorkflowContext parameter
+        has_context = ctx_annotation is not None
+        is_async = asyncio.iscoroutinefunction(func)
+
+        # Initialize parent WITHOUT calling _discover_handlers yet
+        # We'll manually set up the attributes first
+        executor_id = str(id or getattr(func, "__name__", "FunctionExecutor"))
+        kwargs = {"type": "FunctionExecutor"}
+
+        super().__init__(id=executor_id, defer_discovery=True, **kwargs)
+        self._handlers = {}
+        self._handler_specs = []
+
+        # Store the original function and whether it has context
+        self._original_func = func
+        self._has_context = has_context
+        self._is_async = is_async
+
+        # Create a wrapper function that always accepts both message and context
+        if has_context and is_async:
+            # Async function with context - already has the right signature
+            wrapped_func: Callable[[Any, WorkflowContext[Any]], Awaitable[Any]] = func  # type: ignore
+        elif has_context and not is_async:
+            # Sync function with context - wrap to make async using thread pool
+            async def wrapped_func(message: Any, ctx: WorkflowContext[Any]) -> Any:
+                # Call the sync function with both parameters in a thread
+                return await asyncio.to_thread(func, message, ctx)  # type: ignore
+
+        elif not has_context and is_async:
+            # Async function without context - wrap to ignore context
+            async def wrapped_func(message: Any, ctx: WorkflowContext[Any]) -> Any:
+                # Call the async function with just the message
+                return await func(message)  # type: ignore
+
+        else:
+            # Sync function without context - wrap to make async and ignore context using thread pool
+            async def wrapped_func(message: Any, ctx: WorkflowContext[Any]) -> Any:
+                # Call the sync function with just the message in a thread
+                return await asyncio.to_thread(func, message)  # type: ignore
+
+        # Now register our instance handler
+        self._register_instance_handler(
+            name=func.__name__,
+            func=wrapped_func,
+            message_type=message_type,
+            ctx_annotation=ctx_annotation,
+            output_types=output_types,
+            workflow_output_types=workflow_output_types,
+        )
+
+        # Now we can safely call _discover_handlers (it won't find any class-level handlers)
+        self._discover_handlers()
+
+        if not self._handlers:
+            raise ValueError(
+                f"FunctionExecutor {self.__class__.__name__} failed to register handler for {func.__name__}"
+            )
+
+
+@overload
+def executor(func: Callable[..., Any]) -> FunctionExecutor: ...
+
+
+@overload
+def executor(*, id: str | None = None) -> Callable[[Callable[..., Any]], FunctionExecutor]: ...
+
+
+def executor(
+    func: Callable[..., Any] | None = None, *, id: str | None = None
+) -> Callable[[Callable[..., Any]], FunctionExecutor] | FunctionExecutor:
+    """Decorator that converts a function into a FunctionExecutor instance.
+
+    Supports both synchronous and asynchronous functions. Synchronous functions
+    are executed in a thread pool to avoid blocking the event loop.
+
+    Usage:
+
+    .. code-block:: python
+
+        # With arguments (async function):
+        @executor(id="upper_case")
+        async def to_upper(text: str, ctx: WorkflowContext[str]):
+            await ctx.send_message(text.upper())
+
+
+        # Without parentheses (sync function - runs in thread pool):
+        @executor
+        def process_data(data: str):
+            # Process data without sending messages
+            return data.upper()
+
+
+        # Sync function with context (runs in thread pool):
+        @executor
+        def sync_with_context(data: int, ctx: WorkflowContext[int]):
+            # Note: sync functions can still use context
+            return data * 2
+
+    Returns:
+        An Executor instance that can be wired into a Workflow.
+    """
+
+    def wrapper(func: Callable[..., Any]) -> FunctionExecutor:
+        return FunctionExecutor(func, id=id)
+
+    # If func is provided, this means @executor was used without parentheses
+    if func is not None:
+        return wrapper(func)
+
+    # Otherwise, return the wrapper for @executor() or @executor(id="...")
+    return wrapper

python/packages/core/agent_framework/_workflow/_model_utils.py

@@ -0,0 +1,51 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import copy
+import sys
+from typing import Any, TypeVar
+
+if sys.version_info >= (3, 11):
+    from typing import Self  # pragma: no cover
+else:
+    from typing_extensions import Self  # pragma: no cover
+
+TModel = TypeVar("TModel", bound="DictConvertible")
+
+
+class DictConvertible:
+    """Mixin providing conversion helpers for plain Python models."""
+
+    def to_dict(self) -> dict[str, Any]:
+        raise NotImplementedError
+
+    @classmethod
+    def from_dict(cls: type[TModel], data: dict[str, Any]) -> TModel:
+        return cls(**data)  # type: ignore[arg-type]
+
+    def clone(self, *, deep: bool = True) -> Self:
+        return copy.deepcopy(self) if deep else copy.copy(self)  # type: ignore[return-value]
+
+    def to_json(self) -> str:
+        import json
+
+        return json.dumps(self.to_dict())
+
+    @classmethod
+    def from_json(cls: type[TModel], raw: str) -> TModel:
+        import json
+
+        data = json.loads(raw)
+        if not isinstance(data, dict):
+            raise ValueError("JSON payload must decode to a mapping")
+        return cls.from_dict(data)
+
+
+def encode_value(value: Any) -> Any:
+    """Recursively encode values for JSON-friendly serialization."""
+    if isinstance(value, DictConvertible):
+        return value.to_dict()
+    if isinstance(value, dict):
+        return {k: encode_value(v) for k, v in value.items()}  # type: ignore[misc]
+    if isinstance(value, (list, tuple, set)):
+        return [encode_value(v) for v in value]  # type: ignore[misc]
+    return value

python/packages/core/agent_framework/_workflow/_runner.py

@@ -0,0 +1,447 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import logging
+from collections import defaultdict
+from collections.abc import AsyncGenerator, Sequence
+from typing import TYPE_CHECKING, Any
+
+from ._checkpoint import CheckpointStorage, WorkflowCheckpoint
+from ._edge import EdgeGroup
+from ._edge_runner import EdgeRunner, create_edge_runner
+from ._events import WorkflowEvent, WorkflowOutputEvent, _framework_event_origin
+from ._executor import Executor
+from ._runner_context import (
+    _DATACLASS_MARKER,  # type: ignore
+    _MODEL_MARKER,  # type: ignore
+    CheckpointState,
+    Message,
+    RunnerContext,
+    _decode_checkpoint_value,  # type: ignore
+)
+from ._shared_state import SharedState
+
+if TYPE_CHECKING:
+    from ._executor import RequestInfoExecutor
+
+logger = logging.getLogger(__name__)
+
+
+class Runner:
+    """A class to run a workflow in Pregel supersteps."""
+
+    def __init__(
+        self,
+        edge_groups: Sequence[EdgeGroup],
+        executors: dict[str, Executor],
+        shared_state: SharedState,
+        ctx: RunnerContext,
+        max_iterations: int = 100,
+        workflow_id: str | None = None,
+    ) -> None:
+        """Initialize the runner with edges, shared state, and context.
+
+        Args:
+            edge_groups: The edge groups of the workflow.
+            executors: Map of executor IDs to executor instances.
+            shared_state: The shared state for the workflow.
+            ctx: The runner context for the workflow.
+            max_iterations: The maximum number of iterations to run.
+            workflow_id: The workflow ID for checkpointing.
+        """
+        self._executors = executors
+        self._edge_runners = [create_edge_runner(group, executors) for group in edge_groups]
+        self._edge_runner_map = self._parse_edge_runners(self._edge_runners)
+        self._ctx = ctx
+        self._iteration = 0
+        self._max_iterations = max_iterations
+        self._shared_state = shared_state
+        self._workflow_id = workflow_id
+        self._running = False
+        self._resumed_from_checkpoint = False  # Track whether we resumed
+        self.graph_signature_hash: str | None = None
+
+        # Set workflow ID in context if provided
+        if workflow_id:
+            self._ctx.set_workflow_id(workflow_id)
+
+    @property
+    def context(self) -> RunnerContext:
+        """Get the workflow context."""
+        return self._ctx
+
+    def mark_resumed(self, iteration: int | None = None, max_iterations: int | None = None) -> None:
+        """Mark the runner as having resumed from a checkpoint.
+
+        Optionally set the current iteration and max iterations.
+        """
+        self._resumed_from_checkpoint = True
+        if iteration is not None:
+            self._iteration = iteration
+        if max_iterations is not None:
+            self._max_iterations = max_iterations
+
+    async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]:
+        """Run the workflow until no more messages are sent."""
+        if self._running:
+            raise RuntimeError("Runner is already running.")
+
+        self._running = True
+        try:
+            # Emit any events already produced prior to entering loop
+            if await self._ctx.has_events():
+                logger.info("Yielding pre-loop events")
+                for event in await self._ctx.drain_events():
+                    yield event
+
+            # Create first checkpoint if there are messages from initial execution
+            if await self._ctx.has_messages() and self._ctx.has_checkpointing():
+                if not self._resumed_from_checkpoint:
+                    logger.info("Creating checkpoint after initial execution")
+                    await self._create_checkpoint_if_enabled("after_initial_execution")
+                else:
+                    logger.info("Skipping 'after_initial_execution' checkpoint because we resumed from a checkpoint")
+
+            # Initialize context with starting iteration state
+            await self._update_context_with_shared_state()
+
+            while self._iteration < self._max_iterations:
+                logger.info(f"Starting superstep {self._iteration + 1}")
+
+                # Run iteration concurrently with live event streaming: we poll
+                # for new events while the iteration coroutine progresses.
+                iteration_task = asyncio.create_task(self._run_iteration())
+                while not iteration_task.done():
+                    try:
+                        # Wait briefly for any new event; timeout allows progress checks
+                        event = await asyncio.wait_for(self._ctx.next_event(), timeout=0.05)
+                        yield event
+                    except asyncio.TimeoutError:
+                        # Periodically continue to let iteration advance
+                        continue
+
+                # Propagate errors from iteration, but first surface any pending events
+                try:
+                    await iteration_task
+                except Exception:
+                    # Make sure failure-related events (like ExecutorFailedEvent) are surfaced
+                    if await self._ctx.has_events():
+                        for event in await self._ctx.drain_events():
+                            yield event
+                    raise
+                self._iteration += 1
+
+                # Drain any straggler events emitted at tail end
+                if await self._ctx.has_events():
+                    for event in await self._ctx.drain_events():
+                        yield event
+
+                # Update context with current iteration state immediately
+                await self._update_context_with_shared_state()
+
+                logger.info(f"Completed superstep {self._iteration}")
+
+                # Create checkpoint after each superstep iteration
+                await self._create_checkpoint_if_enabled(f"superstep_{self._iteration}")
+
+                if not await self._ctx.has_messages():
+                    break
+
+            if self._iteration >= self._max_iterations and await self._ctx.has_messages():
+                raise RuntimeError(f"Runner did not converge after {self._max_iterations} iterations.")
+
+            logger.info(f"Workflow completed after {self._iteration} supersteps")
+            self._iteration = 0
+            self._resumed_from_checkpoint = False  # Reset resume flag for next run
+        finally:
+            self._running = False
+
+    async def _run_iteration(self) -> None:
+        async def _deliver_messages(source_executor_id: str, messages: list[Message]) -> None:
+            """Outer loop to concurrently deliver messages from all sources to their targets."""
+
+            async def _deliver_message_inner(edge_runner: EdgeRunner, message: Message) -> bool:
+                """Inner loop to deliver a single message through an edge runner."""
+                return await edge_runner.send_message(message, self._shared_state, self._ctx)
+
+            def _normalize_message_payload(message: Message) -> None:
+                data = message.data
+                if not isinstance(data, dict):
+                    return
+                if _MODEL_MARKER not in data and _DATACLASS_MARKER not in data:
+                    return
+                try:
+                    decoded = _decode_checkpoint_value(data)
+                except Exception as exc:  # pragma: no cover - defensive
+                    logger.debug("Failed to decode checkpoint payload during delivery: %s", exc)
+                    return
+                message.data = decoded
+
+            # Route all messages through normal workflow edges
+            associated_edge_runners = self._edge_runner_map.get(source_executor_id, [])
+            for message in messages:
+                _normalize_message_payload(message)
+                # Deliver a message through all edge runners associated with the source executor concurrently.
+                tasks = [_deliver_message_inner(edge_runner, message) for edge_runner in associated_edge_runners]
+                if not tasks:
+                    # No outgoing edges. If this is an AgentExecutorResponse, treat it as an
+                    # intentional terminal emission and emit a WorkflowOutputEvent here.
+                    # (Previously this relied on the executor to emit, but AgentExecutor only
+                    # sends an AgentExecutorResponse message; centralized completion keeps the
+                    # contract consistent with other executors.)
+                    try:  # Local import to avoid circular dependencies at module import time.
+                        from ._executor import AgentExecutorResponse  # type: ignore
+
+                        if isinstance(message.data, AgentExecutorResponse):
+                            final_messages = message.data.agent_run_response.messages
+                            final_text = final_messages[-1].text if final_messages else "(no content)"
+                            with _framework_event_origin():
+                                # TODO(moonbox3): does user expect this event to contain the final text?
+                                output_event = WorkflowOutputEvent(data=final_text, source_executor_id="<Runner>")
+                            await self._ctx.add_event(output_event)
+                            continue  # Terminal handled
+                    except Exception as exc:  # pragma: no cover - defensive
+                        logger.debug("Suppressed exception during terminal message type check: %s", exc)
+                    # Otherwise keep prior behavior (emit warning for unexpected undelivered message).
+                    logger.warning(
+                        f"Message {message} could not be delivered (no outgoing edges). "
+                        "Add a downstream executor or remove the send if this is unexpected."
+                    )
+                    continue
+                results = await asyncio.gather(*tasks)
+                if not any(results):
+                    # Outgoing edges exist but none accepted the message. If this is an
+                    # AgentExecutorResponse, treat as natural terminal and emit completion.
+                    try:
+                        from ._executor import AgentExecutorResponse  # type: ignore
+
+                        if isinstance(message.data, AgentExecutorResponse):
+                            # Emit a single completion event with final text (best-effort extraction)
+                            final_messages = message.data.agent_run_response.messages
+                            final_text = final_messages[-1].text if final_messages else "(no content)"
+                            with _framework_event_origin():
+                                # TODO(moonbox3): does user expect this event to contain the final text?
+                                output_event = WorkflowOutputEvent(data=final_text, source_executor_id="<Runner>")
+                            await self._ctx.add_event(output_event)
+                            continue
+                    except Exception as exc:  # pragma: no cover
+                        logger.debug("Terminal completion emission failed: %s", exc)
+
+                    logger.warning(
+                        f"Message {message} could not be delivered. "
+                        "This may be due to type incompatibility or no matching targets."
+                    )
+
+        messages = await self._ctx.drain_messages()
+        tasks = [_deliver_messages(source_executor_id, messages) for source_executor_id, messages in messages.items()]
+        await asyncio.gather(*tasks)
+
+    async def _create_checkpoint_if_enabled(self, checkpoint_type: str) -> str | None:
+        """Create a checkpoint if checkpointing is enabled and attach a label and metadata."""
+        if not self._ctx.has_checkpointing():
+            return None
+
+        try:
+            # Auto-snapshot executor states
+            await self._auto_snapshot_executor_states()
+            await self._update_context_with_shared_state()
+            checkpoint_category = "initial" if checkpoint_type == "after_initial_execution" else "superstep"
+            metadata = {
+                "superstep": self._iteration,
+                "checkpoint_type": checkpoint_category,
+            }
+            if self.graph_signature_hash:
+                metadata["graph_signature"] = self.graph_signature_hash
+            checkpoint_id = await self._ctx.create_checkpoint(metadata=metadata)
+            logger.info(f"Created {checkpoint_type} checkpoint: {checkpoint_id}")
+            return checkpoint_id
+        except Exception as e:
+            logger.warning(f"Failed to create {checkpoint_type} checkpoint: {e}")
+            return None
+
+    async def _auto_snapshot_executor_states(self) -> None:
+        """Populate executor state by calling snapshot hooks on executors if available.
+
+        Convention:
+          - If an executor defines an async or sync method `snapshot_state(self) -> dict`, use it.
+          - Else if it has a plain attribute `state` that is a dict, use that.
+        Only JSON-serializable dicts should be provided by executors.
+        """
+        for exec_id, executor in self._executors.items():
+            state_dict: dict[str, Any] | None = None
+            snapshot = getattr(executor, "snapshot_state", None)
+            try:
+                if callable(snapshot):
+                    maybe = snapshot()
+                    if asyncio.iscoroutine(maybe):  # type: ignore[arg-type]
+                        maybe = await maybe  # type: ignore[assignment]
+                    if isinstance(maybe, dict):
+                        state_dict = maybe  # type: ignore[assignment]
+                else:
+                    state_attr = getattr(executor, "state", None)
+                    if isinstance(state_attr, dict):
+                        state_dict = state_attr  # type: ignore[assignment]
+            except Exception as ex:  # pragma: no cover
+                logger.debug(f"Executor {exec_id} snapshot_state failed: {ex}")
+            if state_dict is not None:
+                try:
+                    await self._ctx.set_state(exec_id, state_dict)
+                except Exception as ex:  # pragma: no cover
+                    logger.debug(f"Failed to persist state for executor {exec_id}: {ex}")
+
+    async def _update_context_with_shared_state(self) -> None:
+        if not self._ctx.has_checkpointing():
+            return
+
+        try:
+            current_state = await self._ctx.get_checkpoint_state()
+
+            shared_state_data = {}
+            async with self._shared_state.hold():
+                if hasattr(self._shared_state, "_state"):
+                    shared_state_data = dict(self._shared_state._state)  # type: ignore[attr-defined]
+
+            current_state["shared_state"] = shared_state_data
+            current_state["iteration_count"] = self._iteration
+            current_state["max_iterations"] = self._max_iterations
+
+            await self._ctx.set_checkpoint_state(current_state)
+        except Exception as e:
+            logger.warning(f"Failed to update context with shared state: {e}")
+
+    async def restore_from_checkpoint(
+        self,
+        checkpoint_id: str,
+        checkpoint_storage: CheckpointStorage | None = None,
+    ) -> bool:
+        """Restore workflow state from a checkpoint.
+
+        Args:
+            checkpoint_id: The ID of the checkpoint to restore from
+            checkpoint_storage: Optional storage to load checkpoints from when the
+                runner context itself is not configured with checkpointing.
+
+        Returns:
+            True if restoration was successful, False otherwise
+        """
+        try:
+            checkpoint: WorkflowCheckpoint | None
+            if self._ctx.has_checkpointing():
+                checkpoint = await self._ctx.load_checkpoint(checkpoint_id)
+            elif checkpoint_storage is not None:
+                checkpoint = await checkpoint_storage.load_checkpoint(checkpoint_id)
+            else:
+                logger.warning("Context does not support checkpointing and no external storage was provided")
+                return False
+
+            if not checkpoint:
+                logger.error(f"Checkpoint {checkpoint_id} not found")
+                return False
+
+            graph_hash = getattr(self, "graph_signature_hash", None)
+            checkpoint_hash = (checkpoint.metadata or {}).get("graph_signature")
+            if graph_hash and checkpoint_hash and graph_hash != checkpoint_hash:
+                raise ValueError(
+                    "Workflow graph has changed since the checkpoint was created. "
+                    "Please rebuild the original workflow before resuming."
+                )
+            if graph_hash and not checkpoint_hash:
+                logger.warning(
+                    "Checkpoint %s does not include graph signature metadata; skipping topology validation.",
+                    checkpoint_id,
+                )
+
+            state = self._checkpoint_to_state(checkpoint)
+            await self._ctx.set_checkpoint_state(state)
+            if checkpoint.workflow_id:
+                self._ctx.set_workflow_id(checkpoint.workflow_id)
+            self._workflow_id = checkpoint.workflow_id
+
+            await self._restore_shared_state_from_context()
+            self.mark_resumed(
+                iteration=checkpoint.iteration_count,
+                max_iterations=checkpoint.max_iterations,
+            )
+            logger.info(f"Successfully restored workflow from checkpoint: {checkpoint_id}")
+            return True
+        except ValueError:
+            raise
+        except Exception as e:
+            logger.error(f"Failed to restore from checkpoint {checkpoint_id}: {e}")
+            return False
+
+    async def _restore_shared_state_from_context(self) -> None:
+        try:
+            restored_state = await self._ctx.get_checkpoint_state()
+
+            shared_state_data = restored_state.get("shared_state", {})
+            if shared_state_data and hasattr(self._shared_state, "_state"):
+                async with self._shared_state.hold():
+                    self._shared_state._state.clear()  # type: ignore[attr-defined]
+                    self._shared_state._state.update(shared_state_data)  # type: ignore[attr-defined]
+
+            self._iteration = restored_state.get("iteration_count", 0)
+            self._max_iterations = restored_state.get("max_iterations", self._max_iterations)
+
+        except Exception as e:
+            logger.warning(f"Failed to restore shared state from context: {e}")
+
+    @staticmethod
+    def _checkpoint_to_state(checkpoint: WorkflowCheckpoint) -> CheckpointState:
+        return {
+            "messages": checkpoint.messages,
+            "shared_state": checkpoint.shared_state,
+            "executor_states": checkpoint.executor_states,
+            "iteration_count": checkpoint.iteration_count,
+            "max_iterations": checkpoint.max_iterations,
+        }
+
+    def _parse_edge_runners(self, edge_runners: list[EdgeRunner]) -> dict[str, list[EdgeRunner]]:
+        """Parse the edge runners of the workflow into a mapping where each source executor ID maps to its edge runners.
+
+        Args:
+            edge_runners: A list of edge runners in the workflow.
+
+        Returns:
+            A dictionary mapping each source executor ID to a list of edge runners.
+        """
+        parsed: defaultdict[str, list[EdgeRunner]] = defaultdict(list)
+        for runner in edge_runners:
+            # Accessing protected attribute (_edge_group) intentionally for internal wiring.
+            for source_executor_id in runner._edge_group.source_executor_ids:  # type: ignore[attr-defined]
+                parsed[source_executor_id].append(runner)
+
+        return parsed
+
+    def _find_request_info_executor(self) -> "RequestInfoExecutor | None":
+        """Find the RequestInfoExecutor instance in this workflow.
+
+        Returns:
+            The RequestInfoExecutor instance if found, None otherwise.
+        """
+        from ._executor import RequestInfoExecutor
+
+        for executor in self._executors.values():
+            if isinstance(executor, RequestInfoExecutor):
+                return executor
+        return None
+
+    def _is_message_to_request_info_executor(self, msg: "Message") -> bool:
+        """Check if message targets any RequestInfoExecutor in this workflow.
+
+        Args:
+            msg: The message to check.
+
+        Returns:
+            True if the message targets a RequestInfoExecutor, False otherwise.
+        """
+        from ._executor import RequestInfoExecutor
+
+        if not msg.target_id:
+            return False
+
+        # Check all executors to see if target_id matches a RequestInfoExecutor
+        for executor in self._executors.values():
+            if executor.id == msg.target_id and isinstance(executor, RequestInfoExecutor):
+                return True
+        return False

python/packages/core/agent_framework/_workflow/_runner_context.py

@@ -0,0 +1,643 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import contextlib
+import importlib
+import logging
+import sys
+import uuid
+from copy import copy
+from dataclasses import dataclass, fields, is_dataclass
+from typing import Any, Protocol, TypedDict, TypeVar, cast, runtime_checkable
+
+from ._checkpoint import CheckpointStorage, WorkflowCheckpoint
+from ._const import DEFAULT_MAX_ITERATIONS
+from ._events import AgentRunUpdateEvent, WorkflowEvent
+from ._shared_state import SharedState
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+@dataclass
+class Message:
+    """A class representing a message in the workflow."""
+
+    data: Any
+    source_id: str
+    target_id: str | None = None
+
+    # OpenTelemetry trace context fields for message propagation
+    # These are plural to support fan-in scenarios where multiple messages are aggregated
+    trace_contexts: list[dict[str, str]] | None = None  # W3C Trace Context headers from multiple sources
+    source_span_ids: list[str] | None = None  # Publishing span IDs for linking from multiple sources
+
+    # Backward compatibility properties
+    @property
+    def trace_context(self) -> dict[str, str] | None:
+        """Get the first trace context for backward compatibility."""
+        return self.trace_contexts[0] if self.trace_contexts else None
+
+    @property
+    def source_span_id(self) -> str | None:
+        """Get the first source span ID for backward compatibility."""
+        return self.source_span_ids[0] if self.source_span_ids else None
+
+
+class CheckpointState(TypedDict):
+    messages: dict[str, list[dict[str, Any]]]
+    shared_state: dict[str, Any]
+    executor_states: dict[str, dict[str, Any]]
+    iteration_count: int
+    max_iterations: int
+
+
+# Checkpoint serialization helpers
+_MODEL_MARKER = "__af_model__"
+_DATACLASS_MARKER = "__af_dataclass__"
+_AF_MARKER = "__af__"
+
+# Guards to prevent runaway recursion while encoding arbitrary user data
+_MAX_ENCODE_DEPTH = 100
+_CYCLE_SENTINEL = "<cycle>"
+
+
+def _instantiate_checkpoint_dataclass(cls: type[Any], payload: Any) -> Any | None:
+    if not isinstance(cls, type):
+        logger.debug(f"Checkpoint decoder received non-type dataclass reference: {cls!r}")
+        return None
+
+    if isinstance(payload, dict):
+        try:
+            return cls(**payload)  # type: ignore[arg-type]
+        except TypeError as exc:
+            logger.debug(f"Checkpoint decoder could not call {cls.__name__}(**payload): {exc}")
+        except Exception as exc:
+            logger.warning(f"Checkpoint decoder encountered unexpected error calling {cls.__name__}(**payload): {exc}")
+        try:
+            instance = object.__new__(cls)
+        except Exception as exc:
+            logger.debug(f"Checkpoint decoder could not allocate {cls.__name__} without __init__: {exc}")
+            return None
+        for key, val in payload.items():  # type: ignore[attr-defined]
+            try:
+                setattr(instance, key, val)  # type: ignore[arg-type]
+            except Exception as exc:
+                logger.debug(f"Checkpoint decoder could not set attribute {key} on {cls.__name__}: {exc}")
+        return instance
+
+    try:
+        return cls(payload)  # type: ignore[call-arg]
+    except TypeError as exc:
+        logger.debug(f"Checkpoint decoder could not call {cls.__name__}({payload!r}): {exc}")
+    except Exception as exc:
+        logger.warning(f"Checkpoint decoder encountered unexpected error calling {cls.__name__}({payload!r}): {exc}")
+    return None
+
+
+def _supports_model_protocol(obj: object) -> bool:
+    """Detect objects that expose dictionary serialization hooks."""
+    try:
+        obj_type: type[Any] = type(obj)
+    except Exception:
+        return False
+
+    has_to_dict = hasattr(obj, "to_dict") and callable(getattr(obj, "to_dict", None))  # type: ignore[arg-type]
+    has_from_dict = hasattr(obj_type, "from_dict") and callable(getattr(obj_type, "from_dict", None))
+
+    has_to_json = hasattr(obj, "to_json") and callable(getattr(obj, "to_json", None))  # type: ignore[arg-type]
+    has_from_json = hasattr(obj_type, "from_json") and callable(getattr(obj_type, "from_json", None))
+
+    return (has_to_dict and has_from_dict) or (has_to_json and has_from_json)
+
+
+def _import_qualified_name(qualname: str) -> type[Any] | None:
+    if ":" not in qualname:
+        return None
+    module_name, class_name = qualname.split(":", 1)
+    module = sys.modules.get(module_name)
+    if module is None:
+        module = importlib.import_module(module_name)
+    attr: Any = module
+    for part in class_name.split("."):
+        attr = getattr(attr, part)
+    return attr if isinstance(attr, type) else None
+
+
+def _encode_checkpoint_value(value: Any) -> Any:
+    """Recursively encode values into JSON-serializable structures.
+
+    - Objects exposing to_dict/to_json -> { _MODEL_MARKER: "module:Class", value: encoded }
+    - dataclass instances -> { _DATACLASS_MARKER: "module:Class", value: {field: encoded} }
+    - dict -> encode keys as str and values recursively
+    - list/tuple/set -> list of encoded items
+    - other -> returned as-is if already JSON-serializable
+
+    Includes cycle and depth protection to avoid infinite recursion.
+    """
+
+    def _enc(v: Any, stack: set[int], depth: int) -> Any:
+        # Depth guard
+        if depth > _MAX_ENCODE_DEPTH:
+            logger.debug(f"Max encode depth reached at depth={depth} for type={type(v)}")
+            return "<max_depth>"
+
+        # Structured model handling (objects exposing to_dict/to_json)
+        if _supports_model_protocol(v):
+            cls = cast(type[Any], type(v))  # type: ignore
+            try:
+                if hasattr(v, "to_dict") and callable(getattr(v, "to_dict", None)):
+                    raw = v.to_dict()  # type: ignore[attr-defined]
+                    strategy = "to_dict"
+                elif hasattr(v, "to_json") and callable(getattr(v, "to_json", None)):
+                    serialized = v.to_json()  # type: ignore[attr-defined]
+                    if isinstance(serialized, (bytes, bytearray)):
+                        try:
+                            serialized = serialized.decode()
+                        except Exception:
+                            serialized = serialized.decode(errors="replace")
+                    raw = serialized
+                    strategy = "to_json"
+                else:
+                    raise AttributeError("Structured model lacks serialization hooks")
+                return {
+                    _MODEL_MARKER: f"{cls.__module__}:{cls.__name__}",
+                    "strategy": strategy,
+                    "value": _enc(raw, stack, depth + 1),
+                }
+            except Exception as exc:  # best-effort fallback
+                logger.debug(f"Structured model serialization failed for {cls}: {exc}")
+                return str(v)
+
+        # Dataclasses (instances only)
+        if is_dataclass(v) and not isinstance(v, type):
+            oid = id(v)
+            if oid in stack:
+                logger.debug("Cycle detected while encoding dataclass instance")
+                return _CYCLE_SENTINEL
+            stack.add(oid)
+            try:
+                # type(v) already narrows sufficiently; cast was redundant
+                dc_cls: type[Any] = type(v)
+                field_values: dict[str, Any] = {}
+                for f in fields(v):  # type: ignore[arg-type]
+                    field_values[f.name] = _enc(getattr(v, f.name), stack, depth + 1)
+                return {
+                    _DATACLASS_MARKER: f"{dc_cls.__module__}:{dc_cls.__name__}",
+                    "value": field_values,
+                }
+            finally:
+                stack.remove(oid)
+
+        # Collections
+        if isinstance(v, dict):
+            v_dict = cast("dict[object, object]", v)
+            oid = id(v_dict)
+            if oid in stack:
+                logger.debug("Cycle detected while encoding dict")
+                return _CYCLE_SENTINEL
+            stack.add(oid)
+            try:
+                json_dict: dict[str, Any] = {}
+                for k_any, val_any in v_dict.items():  # type: ignore[assignment]
+                    k_str: str = str(k_any)
+                    json_dict[k_str] = _enc(val_any, stack, depth + 1)
+                return json_dict
+            finally:
+                stack.remove(oid)
+
+        if isinstance(v, (list, tuple, set)):
+            iterable_v = cast("list[object] | tuple[object, ...] | set[object]", v)
+            oid = id(iterable_v)
+            if oid in stack:
+                logger.debug("Cycle detected while encoding iterable")
+                return _CYCLE_SENTINEL
+            stack.add(oid)
+            try:
+                seq: list[object] = list(iterable_v)
+                encoded_list: list[Any] = []
+                for item in seq:
+                    encoded_list.append(_enc(item, stack, depth + 1))
+                return encoded_list
+            finally:
+                stack.remove(oid)
+
+        # Primitives (or unknown objects): ensure JSON-serializable
+        if isinstance(v, (str, int, float, bool)) or v is None:
+            return v
+        # Fallback: stringify unknown objects to avoid JSON serialization errors
+        try:
+            return str(v)
+        except Exception:
+            return f"<{type(v).__name__}>"
+
+    return _enc(value, set(), 0)
+
+
+def _decode_checkpoint_value(value: Any) -> Any:
+    """Recursively decode values previously encoded by _encode_checkpoint_value."""
+    if isinstance(value, dict):
+        value_dict = cast(dict[str, Any], value)  # encoded form always uses string keys
+        # Structured model marker handling
+        if _MODEL_MARKER in value_dict and "value" in value_dict:
+            type_key: str | None = value_dict.get(_MODEL_MARKER)  # type: ignore[assignment]
+            strategy: str | None = value_dict.get("strategy")  # type: ignore[assignment]
+            raw_encoded: Any = value_dict.get("value")
+            decoded_payload = _decode_checkpoint_value(raw_encoded)
+            if isinstance(type_key, str):
+                try:
+                    cls = _import_qualified_name(type_key)
+                except Exception as exc:
+                    logger.debug(f"Failed to import structured model {type_key}: {exc}")
+                    cls = None
+
+                if cls is not None:
+                    if strategy == "to_dict" and hasattr(cls, "from_dict"):
+                        with contextlib.suppress(Exception):
+                            return cls.from_dict(decoded_payload)
+                    if strategy == "to_json" and hasattr(cls, "from_json"):
+                        if isinstance(decoded_payload, (str, bytes, bytearray)):
+                            with contextlib.suppress(Exception):
+                                return cls.from_json(decoded_payload)
+                        if isinstance(decoded_payload, dict) and hasattr(cls, "from_dict"):
+                            with contextlib.suppress(Exception):
+                                return cls.from_dict(decoded_payload)
+            return decoded_payload
+        # Dataclass marker handling
+        if _DATACLASS_MARKER in value_dict and "value" in value_dict:
+            type_key_dc: str | None = value_dict.get(_DATACLASS_MARKER)  # type: ignore[assignment]
+            raw_dc: Any = value_dict.get("value")
+            decoded_raw = _decode_checkpoint_value(raw_dc)
+            if isinstance(type_key_dc, str):
+                try:
+                    module_name, class_name = type_key_dc.split(":", 1)
+                    module = sys.modules.get(module_name)
+                    if module is None:
+                        module = importlib.import_module(module_name)
+                    cls_dc: Any = getattr(module, class_name)
+                    constructed = _instantiate_checkpoint_dataclass(cls_dc, decoded_raw)
+                    if constructed is not None:
+                        return constructed
+                except Exception as exc:
+                    logger.debug(f"Failed to decode dataclass {type_key_dc}: {exc}; returning raw value")
+            return decoded_raw
+
+        # Regular dict: decode recursively
+        decoded: dict[str, Any] = {}
+        for k_any, v_any in value_dict.items():
+            decoded[k_any] = _decode_checkpoint_value(v_any)
+        return decoded
+    if isinstance(value, list):
+        # After isinstance check, treat value as list[Any] for decoding
+        value_list: list[Any] = value  # type: ignore[assignment]
+        return [_decode_checkpoint_value(v_any) for v_any in value_list]
+    return value
+
+
+@runtime_checkable
+class RunnerContext(Protocol):
+    """Protocol for the execution context used by the runner.
+
+    A single context that supports messaging, events, and optional checkpointing.
+    If checkpoint storage is not configured, checkpoint methods may raise.
+    """
+
+    async def send_message(self, message: Message) -> None:
+        """Send a message from the executor to the context.
+
+        Args:
+            message: The message to be sent.
+        """
+        ...
+
+    async def drain_messages(self) -> dict[str, list[Message]]:
+        """Drain all messages from the context.
+
+        Returns:
+            A dictionary mapping executor IDs to lists of messages.
+        """
+        ...
+
+    async def has_messages(self) -> bool:
+        """Check if there are any messages in the context.
+
+        Returns:
+            True if there are messages, False otherwise.
+        """
+        ...
+
+    async def add_event(self, event: WorkflowEvent) -> None:
+        """Add an event to the execution context.
+
+        Args:
+            event: The event to be added.
+        """
+        ...
+
+    async def drain_events(self) -> list[WorkflowEvent]:
+        """Drain all events from the context.
+
+        Returns:
+            A list of events that were added to the context.
+        """
+        ...
+
+    async def has_events(self) -> bool:
+        """Check if there are any events in the context.
+
+        Returns:
+            True if there are events, False otherwise.
+        """
+        ...
+
+    async def next_event(self) -> WorkflowEvent:  # pragma: no cover - interface only
+        """Wait for and return the next event emitted by the workflow run."""
+        ...
+
+    async def set_state(self, executor_id: str, state: dict[str, Any]) -> None:
+        """Set the state for a specific executor.
+
+        Args:
+            executor_id: The ID of the executor whose state is being set.
+            state: The state data to be set for the executor.
+        """
+        ...
+
+    async def get_state(self, executor_id: str) -> dict[str, Any] | None:
+        """Get the state for a specific executor.
+
+        Args:
+            executor_id: The ID of the executor whose state is being retrieved.
+
+        Returns:
+            The state data for the executor, or None if not found.
+        """
+        ...
+
+    # Checkpointing capability
+    def has_checkpointing(self) -> bool:
+        """Check if the context supports checkpointing.
+
+        Returns:
+            True if checkpointing is supported, False otherwise.
+        """
+        ...
+
+    # Checkpointing APIs (optional, enabled by storage)
+    def set_workflow_id(self, workflow_id: str) -> None:
+        """Set the workflow ID for the context."""
+        ...
+
+    def reset_for_new_run(self, workflow_shared_state: SharedState | None = None) -> None:
+        """Reset the context for a new workflow run."""
+        ...
+
+    async def create_checkpoint(self, metadata: dict[str, Any] | None = None) -> str:
+        """Create a checkpoint of the current workflow state.
+
+        Args:
+            metadata: Optional metadata to associate with the checkpoint.
+        """
+        ...
+
+    async def restore_from_checkpoint(self, checkpoint_id: str) -> bool:
+        """Restore the context from a checkpoint.
+
+        Args:
+            checkpoint_id: The ID of the checkpoint to restore from.
+
+        Returns:
+            True if the restoration was successful, False otherwise.
+        """
+        ...
+
+    async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None:
+        """Load a checkpoint without mutating the current context state."""
+        ...
+
+    async def get_checkpoint_state(self) -> CheckpointState:
+        """Get the current state of the context suitable for checkpointing."""
+        ...
+
+    async def set_checkpoint_state(self, state: CheckpointState) -> None:
+        """Set the state of the context from a checkpoint.
+
+        Args:
+            state: The state data to set for the context.
+        """
+        ...
+
+
+class InProcRunnerContext:
+    """In-process execution context for local execution and optional checkpointing."""
+
+    def __init__(self, checkpoint_storage: CheckpointStorage | None = None):
+        """Initialize the in-process execution context.
+
+        Args:
+            checkpoint_storage: Optional storage to enable checkpointing.
+        """
+        self._messages: dict[str, list[Message]] = {}
+        # Event queue for immediate streaming of events (e.g., AgentRunUpdateEvent)
+        self._event_queue: asyncio.Queue[WorkflowEvent] = asyncio.Queue()
+
+        # Checkpointing configuration/state
+        self._checkpoint_storage = checkpoint_storage
+        self._workflow_id: str | None = None
+        self._shared_state: dict[str, Any] = {}
+        self._executor_states: dict[str, dict[str, Any]] = {}
+        self._iteration_count: int = 0
+        self._max_iterations: int = 100
+
+    async def send_message(self, message: Message) -> None:
+        self._messages.setdefault(message.source_id, [])
+        self._messages[message.source_id].append(message)
+
+    async def drain_messages(self) -> dict[str, list[Message]]:
+        messages = copy(self._messages)
+        self._messages.clear()
+        return messages
+
+    async def has_messages(self) -> bool:
+        return bool(self._messages)
+
+    async def add_event(self, event: WorkflowEvent) -> None:
+        """Add an event to the context immediately.
+
+        Events are enqueued so runners can stream them in real time instead of
+        waiting for superstep boundaries.
+        """
+        # Filter out empty AgentRunUpdateEvent updates to avoid emitting None/empty chunks
+        try:
+            if isinstance(event, AgentRunUpdateEvent):
+                update = getattr(event, "data", None)
+                # Skip if no update payload
+                if not update:
+                    return
+                # Robust emptiness check: allow either top-level text or any text-bearing content
+                text_val = getattr(update, "text", None)
+                contents = getattr(update, "contents", None)
+                has_text_content = False
+                if contents:
+                    for c in contents:
+                        if getattr(c, "text", None):
+                            has_text_content = True
+                            break
+                if not (text_val or has_text_content):
+                    return
+        except Exception as exc:  # pragma: no cover - defensive logging path
+            # Best-effort filtering only; never block event delivery on filtering errors
+            logger.debug(f"Error while filtering event {event!r}: {exc}", exc_info=True)
+
+        await self._event_queue.put(event)
+
+    async def drain_events(self) -> list[WorkflowEvent]:
+        """Drain all currently queued events without blocking for new ones."""
+        events: list[WorkflowEvent] = []
+        while True:
+            try:
+                events.append(self._event_queue.get_nowait())
+            except asyncio.QueueEmpty:  # type: ignore[attr-defined]
+                break
+        return events
+
+    async def has_events(self) -> bool:
+        return not self._event_queue.empty()
+
+    async def next_event(self) -> WorkflowEvent:
+        """Wait for and return the next event.
+
+        Used by the runner to interleave event emission with ongoing iteration work.
+        """
+        return await self._event_queue.get()
+
+    async def set_state(self, executor_id: str, state: dict[str, Any]) -> None:
+        self._executor_states[executor_id] = state
+
+    async def get_state(self, executor_id: str) -> dict[str, Any] | None:
+        return self._executor_states.get(executor_id)
+
+    def has_checkpointing(self) -> bool:
+        return self._checkpoint_storage is not None
+
+    def set_workflow_id(self, workflow_id: str) -> None:
+        self._workflow_id = workflow_id
+
+    def reset_for_new_run(self, workflow_shared_state: SharedState | None = None) -> None:
+        self._messages.clear()
+        # Clear any pending events (best-effort) by recreating the queue
+        self._event_queue = asyncio.Queue()
+        self._shared_state.clear()
+        self._executor_states.clear()
+        self._iteration_count = 0
+        if workflow_shared_state is not None and hasattr(workflow_shared_state, "_state"):
+            workflow_shared_state._state.clear()  # type: ignore[attr-defined]
+
+    async def create_checkpoint(self, metadata: dict[str, Any] | None = None) -> str:
+        if not self._checkpoint_storage:
+            raise ValueError("Checkpoint storage not configured")
+
+        wf_id = self._workflow_id or str(uuid.uuid4())
+        self._workflow_id = wf_id
+        state = await self.get_checkpoint_state()
+
+        checkpoint = WorkflowCheckpoint(
+            workflow_id=wf_id,
+            messages=state["messages"],
+            shared_state=state.get("shared_state", {}),
+            executor_states=state.get("executor_states", {}),
+            iteration_count=state.get("iteration_count", 0),
+            max_iterations=state.get("max_iterations", DEFAULT_MAX_ITERATIONS),
+            metadata=metadata or {},
+        )
+        checkpoint_id = await self._checkpoint_storage.save_checkpoint(checkpoint)
+        logger.info(f"Created checkpoint {checkpoint_id} for workflow {wf_id}'")
+        return checkpoint_id
+
+    async def restore_from_checkpoint(self, checkpoint_id: str) -> bool:
+        if not self._checkpoint_storage:
+            raise ValueError("Checkpoint storage not configured")
+
+        checkpoint = await self._checkpoint_storage.load_checkpoint(checkpoint_id)
+        if not checkpoint:
+            logger.error(f"Checkpoint {checkpoint_id} not found")
+            return False
+
+        state: CheckpointState = {
+            "messages": checkpoint.messages,
+            "shared_state": checkpoint.shared_state,
+            "executor_states": checkpoint.executor_states,
+            "iteration_count": checkpoint.iteration_count,
+            "max_iterations": checkpoint.max_iterations,
+        }
+        await self.set_checkpoint_state(state)
+        self._workflow_id = checkpoint.workflow_id
+        logger.info(f"Restored state from checkpoint {checkpoint_id}'")
+        return True
+
+    async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None:
+        if not self._checkpoint_storage:
+            raise ValueError("Checkpoint storage not configured")
+        return await self._checkpoint_storage.load_checkpoint(checkpoint_id)
+
+    async def get_checkpoint_state(self) -> CheckpointState:
+        serializable_messages: dict[str, list[dict[str, Any]]] = {}
+        for source_id, message_list in self._messages.items():
+            serializable_messages[source_id] = [
+                {
+                    "data": _encode_checkpoint_value(msg.data),
+                    "source_id": msg.source_id,
+                    "target_id": msg.target_id,
+                    "trace_contexts": msg.trace_contexts,
+                    "source_span_ids": msg.source_span_ids,
+                }
+                for msg in message_list
+            ]
+        return {
+            "messages": serializable_messages,
+            "shared_state": _encode_checkpoint_value(self._shared_state),
+            "executor_states": _encode_checkpoint_value(self._executor_states),
+            "iteration_count": self._iteration_count,
+            "max_iterations": self._max_iterations,
+        }
+
+    async def set_checkpoint_state(self, state: CheckpointState) -> None:
+        self._messages.clear()
+        messages_data = state.get("messages", {})
+        for source_id, message_list in messages_data.items():
+            self._messages[source_id] = [
+                Message(
+                    data=_decode_checkpoint_value(msg.get("data")),
+                    source_id=msg.get("source_id", ""),
+                    target_id=msg.get("target_id"),
+                    trace_contexts=msg.get("trace_contexts"),
+                    source_span_ids=msg.get("source_span_ids"),
+                )
+                for msg in message_list
+            ]
+        # Restore shared_state
+        decoded_shared_raw = _decode_checkpoint_value(state.get("shared_state", {}))
+        if isinstance(decoded_shared_raw, dict):
+            self._shared_state = cast(dict[str, Any], decoded_shared_raw)
+        else:  # fallback to empty dict if corrupted
+            self._shared_state = {}
+
+        # Restore executor_states ensuring value types are dicts
+        decoded_exec_raw = _decode_checkpoint_value(state.get("executor_states", {}))
+        if isinstance(decoded_exec_raw, dict):
+            typed_exec: dict[str, dict[str, Any]] = {}
+            for k_raw, v_raw in decoded_exec_raw.items():  # type: ignore[assignment]
+                if isinstance(k_raw, str) and isinstance(v_raw, dict):
+                    # Filter inner dict to string keys only (best-effort)
+                    inner: dict[str, Any] = {}
+                    for inner_k, inner_v in v_raw.items():  # type: ignore[assignment]
+                        if isinstance(inner_k, str):
+                            inner[inner_k] = inner_v
+                    typed_exec[k_raw] = inner
+            self._executor_states = typed_exec
+        else:
+            self._executor_states = {}
+
+        self._iteration_count = state.get("iteration_count", 0)
+        self._max_iterations = state.get("max_iterations", 100)

python/packages/core/agent_framework/_workflow/_sequential.py

@@ -0,0 +1,198 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Sequential builder for agent/executor workflows with shared conversation context.
+
+This module provides a high-level, agent-focused API to assemble a sequential
+workflow where:
+- Participants are a sequence of AgentProtocol instances or Executors
+- A shared conversation context (list[ChatMessage]) is passed along the chain
+- Agents append their assistant messages to the context
+- Custom executors can transform or summarize and return a refined context
+- The workflow finishes with the final context produced by the last participant
+
+Typical wiring:
+    input -> _InputToConversation -> participant1 -> (agent? -> _ResponseToConversation) -> ... -> participantN -> _EndWithConversation
+
+Notes:
+- Participants can mix AgentProtocol and Executor objects
+- Agents are auto-wrapped by WorkflowBuilder as AgentExecutor
+- AgentExecutor produces AgentExecutorResponse; _ResponseToConversation converts this to list[ChatMessage]
+- Non-agent executors must define a handler that consumes `list[ChatMessage]` and sends back
+  the updated `list[ChatMessage]` via their workflow context
+
+Why include the small internal adapter executors?
+- Input normalization ("input-conversation"): ensures the workflow always starts with a
+  `list[ChatMessage]` regardless of whether callers pass a `str`, a single `ChatMessage`,
+  or a list. This keeps the first hop strongly typed and avoids boilerplate in participants.
+- Agent response adaptation ("to-conversation:<participant>"): agents (via AgentExecutor)
+  emit `AgentExecutorResponse`. The adapter converts that to a `list[ChatMessage]`
+  using `full_conversation` so original prompts aren't lost when chaining.
+- Result output ("end"): yields the final conversation list and the workflow becomes idle
+  giving a consistent terminal payload shape for both agents and custom executors.
+
+These adapters are first-class executors by design so they are type-checked at edges,
+observable (ExecutorInvoke/Completed events), and easily testable/reusable. Their IDs are
+deterministic and self-describing (for example, "to-conversation:writer") to reduce event-log
+confusion and to mirror how the concurrent builder uses explicit dispatcher/aggregator nodes.
+"""  # noqa: E501
+
+import logging
+from collections.abc import Sequence
+from typing import Any
+
+from agent_framework import AgentProtocol, ChatMessage, Role
+
+from ._checkpoint import CheckpointStorage
+from ._executor import (
+    AgentExecutor,
+    AgentExecutorResponse,
+    Executor,
+    handler,
+)
+from ._workflow import Workflow, WorkflowBuilder
+from ._workflow_context import WorkflowContext
+
+logger = logging.getLogger(__name__)
+
+
+class _InputToConversation(Executor):
+    """Normalizes initial input into a list[ChatMessage] conversation."""
+
+    @handler
+    async def from_str(self, prompt: str, ctx: WorkflowContext[list[ChatMessage]]) -> None:
+        await ctx.send_message([ChatMessage(Role.USER, text=prompt)])
+
+    @handler
+    async def from_message(self, message: ChatMessage, ctx: WorkflowContext[list[ChatMessage]]) -> None:  # type: ignore[name-defined]
+        await ctx.send_message([message])
+
+    @handler
+    async def from_messages(self, messages: list[ChatMessage], ctx: WorkflowContext[list[ChatMessage]]) -> None:  # type: ignore[name-defined]
+        # Make a copy to avoid mutation downstream
+        await ctx.send_message(list(messages))
+
+
+class _ResponseToConversation(Executor):
+    """Converts AgentExecutorResponse to list[ChatMessage] conversation for chaining."""
+
+    @handler
+    async def convert(self, response: AgentExecutorResponse, ctx: WorkflowContext[list[ChatMessage]]) -> None:
+        # Always use full_conversation; AgentExecutor guarantees it is populated.
+        if response.full_conversation is None:  # Defensive: indicates a contract violation
+            raise RuntimeError("AgentExecutorResponse.full_conversation missing. AgentExecutor must populate it.")
+        await ctx.send_message(list(response.full_conversation))
+
+
+class _EndWithConversation(Executor):
+    """Terminates the workflow by emitting the final conversation context."""
+
+    @handler
+    async def end(self, conversation: list[ChatMessage], ctx: WorkflowContext[Any, list[ChatMessage]]) -> None:
+        await ctx.yield_output(list(conversation))
+
+
+class SequentialBuilder:
+    r"""High-level builder for sequential agent/executor workflows with shared context.
+
+    - `participants([...])` accepts a list of AgentProtocol (recommended) or Executor
+    - The workflow wires participants in order, passing a list[ChatMessage] down the chain
+    - Agents append their assistant messages to the conversation
+    - Custom executors can transform/summarize and return a list[ChatMessage]
+    - The final output is the conversation produced by the last participant
+
+    Usage:
+    ```python
+    from agent_framework import SequentialBuilder
+
+    workflow = SequentialBuilder().participants([agent1, agent2, summarizer_exec]).build()
+
+    # Enable checkpoint persistence
+    workflow = SequentialBuilder().participants([agent1, agent2]).with_checkpointing(storage).build()
+    ```
+    """
+
+    def __init__(self) -> None:
+        self._participants: list[AgentProtocol | Executor] = []
+        self._checkpoint_storage: CheckpointStorage | None = None
+
+    def participants(self, participants: Sequence[AgentProtocol | Executor]) -> "SequentialBuilder":
+        """Define the ordered participants for this sequential workflow.
+
+        Accepts AgentProtocol instances (auto-wrapped as AgentExecutor) or Executor instances.
+        Raises if empty or duplicates are provided for clarity.
+        """
+        if not participants:
+            raise ValueError("participants cannot be empty")
+
+        # Defensive duplicate detection
+        seen_agent_ids: set[int] = set()
+        seen_executor_ids: set[str] = set()
+        for p in participants:
+            if isinstance(p, Executor):
+                if p.id in seen_executor_ids:
+                    raise ValueError(f"Duplicate executor participant detected: id '{p.id}'")
+                seen_executor_ids.add(p.id)
+            else:
+                # Treat non-Executor as agent-like (AgentProtocol). Structural checks can be brittle at runtime.
+                pid = id(p)
+                if pid in seen_agent_ids:
+                    raise ValueError("Duplicate agent participant detected (same agent instance provided twice)")
+                seen_agent_ids.add(pid)
+
+        self._participants = list(participants)
+        return self
+
+    def with_checkpointing(self, checkpoint_storage: CheckpointStorage) -> "SequentialBuilder":
+        """Enable checkpointing for the built workflow using the provided storage."""
+        self._checkpoint_storage = checkpoint_storage
+        return self
+
+    def build(self) -> Workflow:
+        """Build and validate the sequential workflow.
+
+        Wiring pattern:
+        - _InputToConversation normalizes the initial input into list[ChatMessage]
+        - For each participant in order:
+            - If Agent (or AgentExecutor): pass conversation to the agent, then convert response
+              to conversation via _ResponseToConversation
+            - Else (custom Executor): pass conversation directly to the executor
+        - _EndWithConversation yields the final conversation and the workflow becomes idle
+        """
+        if not self._participants:
+            raise ValueError("No participants provided. Call .participants([...]) first.")
+
+        # Internal nodes
+        input_conv = _InputToConversation(id="input-conversation")
+        end = _EndWithConversation(id="end")
+
+        builder = WorkflowBuilder()
+        builder.set_start_executor(input_conv)
+
+        # Start of the chain is the input normalizer
+        prior: Executor | AgentProtocol = input_conv
+
+        for p in self._participants:
+            # Agent-like branch: either explicitly an AgentExecutor or any non-AgentExecutor
+            if not (isinstance(p, Executor) and not isinstance(p, AgentExecutor)):
+                # input conversation -> (agent) -> response -> conversation
+                builder.add_edge(prior, p)
+                # Give the adapter a deterministic, self-describing id
+                label: str
+                label = p.id if isinstance(p, Executor) else getattr(p, "name", None) or p.__class__.__name__
+                resp_to_conv = _ResponseToConversation(id=f"to-conversation:{label}")
+                builder.add_edge(p, resp_to_conv)
+                prior = resp_to_conv
+            elif isinstance(p, Executor):
+                # Custom executor operates on list[ChatMessage]
+                builder.add_edge(prior, p)
+                prior = p
+            else:  # pragma: no cover - defensive
+                raise TypeError(f"Unsupported participant type: {type(p).__name__}")
+
+        # Terminate with the final conversation
+        builder.add_edge(prior, end)
+
+        if self._checkpoint_storage is not None:
+            builder = builder.with_checkpointing(self._checkpoint_storage)
+
+        return builder.build()

python/packages/core/agent_framework/_workflow/_shared_state.py

@@ -0,0 +1,69 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from typing import Any
+
+
+class SharedState:
+    """A class to manage shared state in a workflow."""
+
+    def __init__(self) -> None:
+        """Initialize the shared state."""
+        self._state: dict[str, Any] = {}
+        self._shared_state_lock = asyncio.Lock()
+
+    async def set(self, key: str, value: Any) -> None:
+        """Set a value in the shared state."""
+        async with self._shared_state_lock:
+            await self.set_within_hold(key, value)
+
+    async def get(self, key: str) -> Any:
+        """Get a value from the shared state."""
+        async with self._shared_state_lock:
+            return await self.get_within_hold(key)
+
+    async def has(self, key: str) -> bool:
+        """Check if a key exists in the shared state."""
+        async with self._shared_state_lock:
+            return await self.has_within_hold(key)
+
+    async def delete(self, key: str) -> None:
+        """Delete a key from the shared state."""
+        async with self._shared_state_lock:
+            await self.delete_within_hold(key)
+
+    @asynccontextmanager
+    async def hold(self) -> AsyncIterator["SharedState"]:
+        """Context manager to hold the shared state lock for multiple operations.
+
+        Usage:
+            async with shared_state.hold():
+                await shared_state.set_within_hold("key", value)
+                value = await shared_state.get_within_hold("key")
+        """
+        async with self._shared_state_lock:
+            yield self
+
+    # Unsafe methods that don't acquire locks (for use within hold() context)
+    async def set_within_hold(self, key: str, value: Any) -> None:
+        """Set a value without acquiring the lock (unsafe - use within hold() context)."""
+        self._state[key] = value
+
+    async def get_within_hold(self, key: str) -> Any:
+        """Get a value without acquiring the lock (unsafe - use within hold() context)."""
+        if key not in self._state:
+            raise KeyError(f"Key '{key}' not found in shared state.")
+        return self._state[key]
+
+    async def has_within_hold(self, key: str) -> bool:
+        """Check if a key exists without acquiring the lock (unsafe - use within hold() context)."""
+        return key in self._state
+
+    async def delete_within_hold(self, key: str) -> None:
+        """Delete a key without acquiring the lock (unsafe - use within hold() context)."""
+        if key in self._state:
+            del self._state[key]
+        else:
+            raise KeyError(f"Key '{key}' not found in shared state.")

python/packages/core/agent_framework/_workflow/_typing_utils.py

@@ -0,0 +1,155 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import logging
+from collections.abc import Mapping
+from dataclasses import fields, is_dataclass
+from types import UnionType
+from typing import Any, Union, get_args, get_origin
+
+logger = logging.getLogger(__name__)
+
+
+def _coerce_to_type(value: Any, target_type: type) -> Any | None:
+    """Best-effort conversion of value into target_type."""
+    if isinstance(value, target_type):
+        return value
+
+    # Convert dataclass instances or objects with __dict__ into dict first
+    if not isinstance(value, dict):
+        if is_dataclass(value):
+            value = {f.name: getattr(value, f.name) for f in fields(value)}
+        else:
+            value_dict = getattr(value, "__dict__", None)
+            if isinstance(value_dict, dict):
+                value = dict(value_dict)
+
+    if isinstance(value, dict):
+        ctor_kwargs: dict[str, Any] = dict(value)
+
+        if is_dataclass(target_type):
+            field_names = {f.name for f in fields(target_type)}
+            ctor_kwargs = {k: v for k, v in value.items() if k in field_names}
+
+        try:
+            return target_type(**ctor_kwargs)  # type: ignore[arg-type]
+        except TypeError as exc:
+            logger.debug(f"_coerce_to_type could not call {target_type.__name__}(**..): {exc}")
+        except Exception as exc:  # pragma: no cover - unexpected constructor failure
+            logger.warning(
+                f"_coerce_to_type encountered unexpected error calling {target_type.__name__} constructor: {exc}"
+            )
+        try:
+            instance: Any = object.__new__(target_type)
+        except Exception as exc:  # pragma: no cover - pathological type
+            logger.debug(f"_coerce_to_type could not allocate {target_type.__name__} without __init__: {exc}")
+            return None
+        for key, val in value.items():
+            try:
+                setattr(instance, key, val)
+            except Exception as exc:
+                logger.debug(
+                    f"_coerce_to_type could not set {target_type.__name__}.{key} during fallback assignment: {exc}"
+                )
+                continue
+        return instance
+
+    return None
+
+
+def is_instance_of(data: Any, target_type: type | UnionType | Any) -> bool:
+    """Check if the data is an instance of the target type.
+
+    Args:
+        data (Any): The data to check.
+        target_type (type): The type to check against.
+
+    Returns:
+        bool: True if data is an instance of target_type, False otherwise.
+    """
+    # Case 0: target_type is Any - always return True
+    if target_type is Any:
+        return True
+
+    origin = get_origin(target_type)
+    args = get_args(target_type)
+
+    # Case 1: origin is None, meaning target_type is not a generic type
+    if origin is None:
+        return isinstance(data, target_type)
+
+    # Case 2: target_type is Optional[T] or Union[T1, T2, ...]
+    # Optional[T] is really just as Union[T, None]
+    if origin is UnionType:
+        return any(is_instance_of(data, arg) for arg in args)
+
+    # Case 2b: Handle typing.Union (legacy Union syntax)
+    if origin is Union:
+        return any(is_instance_of(data, arg) for arg in args)
+
+    # Case 3: target_type is a generic type
+    if origin in [list, set]:
+        return isinstance(data, origin) and (
+            not args or all(any(is_instance_of(item, arg) for arg in args) for item in data)
+        )  # type: ignore
+
+    # Case 4: target_type is a tuple
+    if origin is tuple:
+        if len(args) == 2 and args[1] is Ellipsis:  # Tuple[T, ...] case
+            element_type = args[0]
+            return isinstance(data, tuple) and all(is_instance_of(item, element_type) for item in data)
+        if len(args) == 1 and args[0] is Ellipsis:  # Tuple[...] case
+            return isinstance(data, tuple)
+        if len(args) == 0:
+            return isinstance(data, tuple)
+        return (
+            isinstance(data, tuple)
+            and len(data) == len(args)  # type: ignore
+            and all(is_instance_of(item, arg) for item, arg in zip(data, args, strict=False))  # type: ignore
+        )
+
+    # Case 5: target_type is a dict
+    if origin is dict:
+        return isinstance(data, dict) and (
+            not args
+            or all(
+                is_instance_of(key, args[0]) and is_instance_of(value, args[1])
+                for key, value in data.items()  # type: ignore
+            )
+        )
+
+    # Case 6: target_type is RequestResponse[T, U] - validate generic parameters
+    if origin and hasattr(origin, "__name__") and origin.__name__ == "RequestResponse":
+        if not isinstance(data, origin):
+            return False
+        # Validate generic parameters for RequestResponse[TRequest, TResponse]
+        if len(args) >= 2:
+            request_type, response_type = args[0], args[1]
+            # Check if the original_request matches TRequest and data matches TResponse
+            if (
+                hasattr(data, "original_request")
+                and data.original_request is not None
+                and not is_instance_of(data.original_request, request_type)
+            ):
+                # Checkpoint decoding can leave original_request as a plain mapping. In that
+                # case we coerce it back into the expected request type so downstream handlers
+                # and validators still receive a fully typed RequestResponse instance.
+                original_request = data.original_request
+                if isinstance(original_request, Mapping):
+                    coerced = _coerce_to_type(dict(original_request), request_type)
+                    if coerced is None or not isinstance(coerced, request_type):
+                        return False
+                    data.original_request = coerced
+                else:
+                    return False
+            if hasattr(data, "data") and data.data is not None and not is_instance_of(data.data, response_type):
+                return False
+        return True
+
+    # Case 7: Other custom generic classes - check origin type only
+    # For generic classes, we check if data is an instance of the origin type
+    # We don't validate the generic parameters at runtime since that's handled by type system
+    if origin and hasattr(origin, "__name__"):
+        return isinstance(data, origin)
+
+    # Fallback: if we reach here, we assume data is an instance of the target_type
+    return isinstance(data, target_type)

python/packages/core/agent_framework/_workflow/_validation.py

@@ -0,0 +1,602 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import inspect
+import logging
+from collections import defaultdict
+from collections.abc import Sequence
+from enum import Enum
+from types import UnionType
+from typing import Any, Union, get_args, get_origin
+
+from ._edge import Edge, EdgeGroup, FanInEdgeGroup
+from ._executor import Executor, RequestInfoExecutor
+
+logger = logging.getLogger(__name__)
+
+# Track cycle signatures we've already reported to avoid spamming logs when workflows
+# with intentional feedback loops are constructed multiple times in the same process.
+_LOGGED_CYCLE_SIGNATURES: set[tuple[str, ...]] = set()
+
+
+# region Enums and Base Classes
+class ValidationTypeEnum(Enum):
+    """Enumeration of workflow validation types."""
+
+    EDGE_DUPLICATION = "EDGE_DUPLICATION"
+    EXECUTOR_DUPLICATION = "EXECUTOR_DUPLICATION"
+    TYPE_COMPATIBILITY = "TYPE_COMPATIBILITY"
+    GRAPH_CONNECTIVITY = "GRAPH_CONNECTIVITY"
+    HANDLER_OUTPUT_ANNOTATION = "HANDLER_OUTPUT_ANNOTATION"
+    INTERCEPTOR_CONFLICT = "INTERCEPTOR_CONFLICT"
+
+
+class WorkflowValidationError(Exception):
+    """Base exception for workflow validation errors."""
+
+    def __init__(self, message: str, validation_type: ValidationTypeEnum):
+        super().__init__(message)
+        self.message = message
+        self.validation_type = validation_type
+
+    def __str__(self) -> str:
+        return f"[{self.validation_type.value}] {self.message}"
+
+
+class EdgeDuplicationError(WorkflowValidationError):
+    """Exception raised when duplicate edges are detected in the workflow."""
+
+    def __init__(self, edge_id: str):
+        super().__init__(
+            message=f"Duplicate edge detected: {edge_id}. Each edge in the workflow must be unique.",
+            validation_type=ValidationTypeEnum.EDGE_DUPLICATION,
+        )
+        self.edge_id = edge_id
+
+
+class ExecutorDuplicationError(WorkflowValidationError):
+    """Exception raised when duplicate executor identifiers are detected."""
+
+    def __init__(self, executor_id: str):
+        super().__init__(
+            message=(
+                f"Duplicate executor id detected: '{executor_id}'. Executor ids must be globally unique within a "
+                "workflow."
+            ),
+            validation_type=ValidationTypeEnum.EXECUTOR_DUPLICATION,
+        )
+        self.executor_id = executor_id
+
+
+class TypeCompatibilityError(WorkflowValidationError):
+    """Exception raised when type incompatibility is detected between connected executors."""
+
+    def __init__(
+        self,
+        source_executor_id: str,
+        target_executor_id: str,
+        source_types: list[type[Any]],
+        target_types: list[type[Any]],
+    ):
+        # Use a placeholder for incompatible types - will be computed in WorkflowGraphValidator
+        super().__init__(
+            message=f"Type incompatibility between executors '{source_executor_id}' -> '{target_executor_id}'. "
+            f"Source executor outputs types {[str(t) for t in source_types]} but target executor "
+            f"can only handle types {[str(t) for t in target_types]}.",
+            validation_type=ValidationTypeEnum.TYPE_COMPATIBILITY,
+        )
+        self.source_executor_id = source_executor_id
+        self.target_executor_id = target_executor_id
+        self.source_types = source_types
+        self.target_types = target_types
+
+
+class GraphConnectivityError(WorkflowValidationError):
+    """Exception raised when graph connectivity issues are detected."""
+
+    def __init__(self, message: str):
+        super().__init__(message, validation_type=ValidationTypeEnum.GRAPH_CONNECTIVITY)
+
+
+class InterceptorConflictError(WorkflowValidationError):
+    """Exception raised when multiple executors intercept the same request type from the same sub-workflow."""
+
+    def __init__(self, message: str):
+        super().__init__(message, validation_type=ValidationTypeEnum.INTERCEPTOR_CONFLICT)
+
+
+# endregion
+
+
+# region Workflow Graph Validator
+class WorkflowGraphValidator:
+    """Validator for workflow graphs.
+
+    This validator performs multiple validation checks:
+    1. Edge duplication validation
+    2. Type compatibility validation between connected executors
+    3. Graph connectivity validation
+    """
+
+    def __init__(self) -> None:
+        self._edges: list[Edge] = []
+        self._executors: dict[str, Executor] = {}
+        self._duplicate_executor_ids: set[str] = set()
+        self._start_executor_ref: Executor | str | None = None
+
+    # region Core Validation Methods
+    def validate_workflow(
+        self,
+        edge_groups: Sequence[EdgeGroup],
+        executors: dict[str, Executor],
+        start_executor: Executor | str,
+        *,
+        duplicate_executor_ids: Sequence[str] | None = None,
+    ) -> None:
+        """Validate the entire workflow graph.
+
+        Args:
+            edge_groups: list of edge groups in the workflow
+            executors: Map of executor IDs to executor instances
+            start_executor: The starting executor (can be instance or ID)
+            duplicate_executor_ids: Optional list of known duplicate executor IDs to pre-populate
+
+        Raises:
+            WorkflowValidationError: If any validation fails
+        """
+        self._executors = executors
+        self._edges = [edge for group in edge_groups for edge in group.edges]
+        self._edge_groups = edge_groups
+        self._duplicate_executor_ids = set(duplicate_executor_ids or [])
+        self._start_executor_ref = start_executor
+
+        # If only the start executor exists, add it to the executor map
+        # Handle the special case where the workflow consists of only a single executor and no edges.
+        # In this scenario, the executor map will be empty because there are no edge groups to reference executors.
+        # Adding the start executor to the map ensures that single-executor workflows (without any edges) are supported,
+        # allowing validation and execution to proceed for workflows that do not require inter-executor communication.
+        if not self._executors and start_executor and isinstance(start_executor, Executor):
+            self._executors[start_executor.id] = start_executor
+
+        # Validate that start_executor exists in the graph
+        # It should because we check for it in the WorkflowBuilder
+        # but we do it here for completeness.
+        start_executor_id = start_executor.id if isinstance(start_executor, Executor) else start_executor
+        if start_executor_id not in self._executors:
+            raise GraphConnectivityError(f"Start executor '{start_executor_id}' is not present in the workflow graph")
+
+        # Additional presence verification:
+        # A start executor that is only injected via the builder (present in the executors map)
+        # but not referenced by any edge while other executors ARE referenced indicates a
+        # configuration error: the chosen start node is effectively disconnected / unknown to the
+        # defined graph topology. For single-node workflows (no edges) we allow the start executor
+        # to stand alone (handled above when we inject it into the map). We perform this refined
+        # check only when there is at least one edge group defined.
+        if self._edges:  # Only evaluate when the workflow defines edges
+            edge_executor_ids: set[str] = set()
+            for _e in self._edges:
+                edge_executor_ids.add(_e.source_id)
+                edge_executor_ids.add(_e.target_id)
+            if start_executor_id not in edge_executor_ids:
+                raise GraphConnectivityError(
+                    f"Start executor '{start_executor_id}' is not present in the workflow graph"
+                )
+
+        # Run all checks
+        self._validate_executor_id_uniqueness(start_executor_id)
+        self._validate_edge_duplication()
+        self._validate_handler_output_annotations()
+        self._validate_type_compatibility()
+        self._validate_graph_connectivity(start_executor_id)
+        self._validate_self_loops()
+        self._validate_dead_ends()
+        self._validate_cycles()
+
+    def _validate_handler_output_annotations(self) -> None:
+        """Validate that each handler's ctx parameter is annotated with WorkflowContext[T].
+
+        Note: This validation is now primarily handled at handler registration time
+        via the unified validation functions in _workflow_context.py when the @handler
+        decorator is applied. This method is kept minimal for any edge cases.
+        """
+        # The comprehensive validation is already done during handler registration:
+        # 1. @handler decorator calls validate_function_signature()
+        # 2. FunctionExecutor constructor calls validate_function_signature()
+        # 3. Both use validate_workflow_context_annotation() for WorkflowContext validation
+        #
+        # All executors in the workflow must have gone through one of these paths,
+        # so redundant validation here is unnecessary and has been removed.
+        pass
+
+    # endregion
+
+    def _validate_executor_id_uniqueness(self, start_executor_id: str) -> None:
+        """Ensure executor identifiers are unique throughout the workflow graph."""
+        duplicates: set[str] = set(self._duplicate_executor_ids)
+
+        id_counts: defaultdict[str, int] = defaultdict(int)
+        for key, executor in self._executors.items():
+            id_counts[executor.id] += 1
+            if key != executor.id:
+                duplicates.add(executor.id)
+
+        duplicates.update({executor_id for executor_id, count in id_counts.items() if count > 1})
+
+        if isinstance(self._start_executor_ref, Executor):
+            mapped = self._executors.get(start_executor_id)
+            if mapped is not None and mapped is not self._start_executor_ref:
+                duplicates.add(start_executor_id)
+
+        if duplicates:
+            raise ExecutorDuplicationError(sorted(duplicates)[0])
+
+    # region Edge and Type Validation
+    def _validate_edge_duplication(self) -> None:
+        """Validate that there are no duplicate edges in the workflow.
+
+        Raises:
+            EdgeDuplicationError: If duplicate edges are found
+        """
+        seen_edge_ids: set[str] = set()
+
+        for edge in self._edges:
+            edge_id = edge.id
+            if edge_id in seen_edge_ids:
+                raise EdgeDuplicationError(edge_id)
+            seen_edge_ids.add(edge_id)
+
+    def _validate_type_compatibility(self) -> None:
+        """Validate type compatibility between connected executors.
+
+        This checks that the output types of source executors are compatible
+        with the input types expected by target executors.
+
+        Raises:
+            TypeCompatibilityError: If type incompatibility is detected
+        """
+        for edge_group in self._edge_groups:
+            for edge in edge_group.edges:
+                self._validate_edge_type_compatibility(edge, edge_group)
+
+    def _validate_edge_type_compatibility(self, edge: Edge, edge_group: EdgeGroup) -> None:
+        """Validate type compatibility for a specific edge.
+
+        This checks that the output types of the source executor are compatible
+        with the input types expected by the target executor.
+
+        Args:
+            edge: The edge to validate
+            edge_group: The edge group containing this edge
+
+        Raises:
+            TypeCompatibilityError: If type incompatibility is detected
+        """
+        source_executor = self._executors[edge.source_id]
+        target_executor = self._executors[edge.target_id]
+
+        # Get output types from source executor
+        source_output_types = list(source_executor.output_types)
+
+        # Get input types from target executor
+        target_input_types = target_executor.input_types
+
+        # If either executor has no type information, log warning and skip validation
+        # This allows for dynamic typing scenarios but warns about reduced validation coverage
+        if not source_output_types or not target_input_types:
+            # Suppress warnings for RequestInfoExecutor where dynamic typing is expected
+            if not source_output_types and not isinstance(source_executor, RequestInfoExecutor):
+                logger.warning(
+                    f"Executor '{source_executor.id}' has no output type annotations. "
+                    f"Type compatibility validation will be skipped for edges from this executor. "
+                    f"Consider adding WorkflowContext[T] generics in handlers for better validation."
+                )
+            if not target_input_types and not isinstance(target_executor, RequestInfoExecutor):
+                logger.warning(
+                    f"Executor '{target_executor.id}' has no input type annotations. "
+                    f"Type compatibility validation will be skipped for edges to this executor. "
+                    f"Consider adding type annotations to message handler parameters for better validation."
+                )
+            return
+
+        # Check if any source output type is compatible with any target input type
+        compatible = False
+        compatible_pairs: list[tuple[type[Any], type[Any]]] = []
+
+        for source_type in source_output_types:
+            for target_type in target_input_types:
+                if isinstance(edge_group, FanInEdgeGroup):
+                    # If the edge is part of an edge group, the target expects a list of data types
+                    if self._is_type_compatible(list[source_type], target_type):  # type: ignore[valid-type]
+                        compatible = True
+                        compatible_pairs.append((list[source_type], target_type))  # type: ignore[valid-type]
+                else:
+                    if self._is_type_compatible(source_type, target_type):
+                        compatible = True
+                        compatible_pairs.append((source_type, target_type))
+
+        # Log successful type compatibility for debugging
+        if compatible:
+            logger.debug(
+                f"Type compatibility validated for edge '{source_executor.id}' -> '{target_executor.id}'. "
+                f"Compatible type pairs: {[(str(s), str(t)) for s, t in compatible_pairs]}"
+            )
+
+        if not compatible:
+            # Enhanced error with more detailed information
+            raise TypeCompatibilityError(
+                source_executor.id,
+                target_executor.id,
+                source_output_types,
+                target_input_types,
+            )
+
+    # endregion
+
+    # region Graph Connectivity Validation
+    def _validate_graph_connectivity(self, start_executor_id: str) -> None:
+        """Validate graph connectivity and detect potential issues.
+
+        This performs several checks:
+        - Detects unreachable executors from the start node
+        - Detects isolated executors (no incoming or outgoing edges)
+        - Warns about potential infinite loops
+
+        Args:
+            start_executor_id: The ID of the starting executor
+
+        Raises:
+            GraphConnectivityError: If connectivity issues are detected
+        """
+        # Build adjacency list for the graph
+        graph: dict[str, list[str]] = defaultdict(list)
+        all_executors = set(self._executors.keys())
+
+        for edge in self._edges:
+            graph[edge.source_id].append(edge.target_id)
+
+        # Find reachable nodes from start
+        reachable = self._find_reachable_nodes(graph, start_executor_id)
+
+        # Check for unreachable executors
+        unreachable = all_executors - reachable
+        if unreachable:
+            raise GraphConnectivityError(
+                f"The following executors are unreachable from the start executor '{start_executor_id}': "
+                f"{sorted(unreachable)}. This may indicate a disconnected workflow graph."
+            )
+
+        # Check for isolated executors (no edges)
+        isolated_executors: list[str] = []
+        for executor_id in all_executors:
+            has_incoming = any(edge.target_id == executor_id for edge in self._edges)
+            has_outgoing = any(edge.source_id == executor_id for edge in self._edges)
+
+            if not has_incoming and not has_outgoing and executor_id != start_executor_id:
+                isolated_executors.append(executor_id)
+
+        if isolated_executors:
+            raise GraphConnectivityError(
+                f"The following executors are isolated (no incoming or outgoing edges): "
+                f"{sorted(isolated_executors)}. Isolated executors will never be executed."
+            )
+
+    def _find_reachable_nodes(self, graph: dict[str, list[str]], start: str) -> set[str]:
+        """Find all nodes reachable from the start node using DFS.
+
+        Args:
+            graph: Adjacency list representation of the graph
+            start: Starting node ID
+
+        Returns:
+            Set of reachable node IDs
+        """
+        visited: set[str] = set()
+        stack = [start]
+
+        while stack:
+            node = stack.pop()
+            if node not in visited:
+                visited.add(node)
+                stack.extend(graph[node])
+
+        return visited
+
+    # endregion
+
+    # region Additional Validation Scenarios
+    def _validate_self_loops(self) -> None:
+        """Detect and log self-loops (edges from executor to itself).
+
+        Self-loops might indicate recursive processing which could be intentional
+        but should be highlighted for review.
+        """
+        self_loops = [edge for edge in self._edges if edge.source_id == edge.target_id]
+
+        for edge in self_loops:
+            logger.warning(
+                f"Self-loop detected: Executor '{edge.source_id}' connects to itself. "
+                f"This may cause infinite recursion if not properly handled with conditions."
+            )
+
+    def _validate_dead_ends(self) -> None:
+        """Identify executors that have no outgoing edges (potential dead ends).
+
+        These might be intentional final nodes or could indicate missing connections.
+        """
+        executors_with_outgoing = {edge.source_id for edge in self._edges}
+        all_executor_ids = set(self._executors.keys())
+        dead_ends = all_executor_ids - executors_with_outgoing
+
+        if dead_ends:
+            logger.info(
+                f"Dead-end executors detected (no outgoing edges): {sorted(dead_ends)}. "
+                f"Verify these are intended as final nodes in the workflow."
+            )
+
+    def _validate_cycles(self) -> None:
+        """Detect cycles in the workflow graph.
+
+        Cycles might be intentional for iterative processing but should be flagged
+        for review to ensure proper termination conditions exist. We surface each
+        distinct cycle group only once per process to avoid noisy, repeated warnings
+        when rebuilding the same workflow.
+        """
+        # Build adjacency list (ensure every executor appears even if it has no outgoing edges)
+        graph: dict[str, list[str]] = defaultdict(list)
+        for edge in self._edges:
+            graph[edge.source_id].append(edge.target_id)
+            graph.setdefault(edge.target_id, [])
+        for executor_id in self._executors:
+            graph.setdefault(executor_id, [])
+
+        # Tarjan's algorithm to locate strongly-connected components that form cycles
+        index: dict[str, int] = {}
+        lowlink: dict[str, int] = {}
+        on_stack: set[str] = set()
+        stack: list[str] = []
+        current_index = 0
+        cycle_components: list[list[str]] = []
+
+        def strongconnect(node: str) -> None:
+            nonlocal current_index
+
+            index[node] = current_index
+            lowlink[node] = current_index
+            current_index += 1
+            stack.append(node)
+            on_stack.add(node)
+
+            for neighbor in graph[node]:
+                if neighbor not in index:
+                    strongconnect(neighbor)
+                    lowlink[node] = min(lowlink[node], lowlink[neighbor])
+                elif neighbor in on_stack:
+                    lowlink[node] = min(lowlink[node], index[neighbor])
+
+            if lowlink[node] == index[node]:
+                component: list[str] = []
+                while True:
+                    member = stack.pop()
+                    on_stack.discard(member)
+                    component.append(member)
+                    if member == node:
+                        break
+
+                # A strongly connected component represents a cycle if it has more than one
+                # node or if a single node references itself directly.
+                if len(component) > 1 or any(member in graph[member] for member in component):
+                    cycle_components.append(component)
+
+        for executor_id in graph:
+            if executor_id not in index:
+                strongconnect(executor_id)
+
+        if not cycle_components:
+            return
+
+        unseen_components: list[list[str]] = []
+        for component in cycle_components:
+            signature = tuple(sorted(component))
+            if signature in _LOGGED_CYCLE_SIGNATURES:
+                continue
+            _LOGGED_CYCLE_SIGNATURES.add(signature)
+            unseen_components.append(component)
+
+        if not unseen_components:
+            # All cycles already reported in this process; keep noise low but retain traceability.
+            logger.debug(
+                "Cycle detected in workflow graph but previously reported. Components: %s",
+                [sorted(component) for component in cycle_components],
+            )
+            return
+
+        def _format_cycle(component: list[str]) -> str:
+            if not component:
+                return ""
+            ordered = list(component)
+            ordered.append(component[0])
+            return " -> ".join(ordered)
+
+        formatted_cycles = ", ".join(_format_cycle(component) for component in unseen_components)
+        logger.warning(
+            "Cycle detected in the workflow graph involving: %s. Ensure termination or iteration limits exist.",
+            formatted_cycles,
+        )
+
+    # endregion
+
+    # region Type Compatibility Utilities
+    @staticmethod
+    def _is_type_compatible(source_type: type[Any], target_type: type[Any]) -> bool:
+        """Check if source_type is compatible with target_type."""
+        # Handle Any type
+        if source_type is Any or target_type is Any:
+            return True
+
+        # Handle exact match
+        if source_type == target_type:
+            return True
+
+        # Handle inheritance
+        try:
+            if inspect.isclass(source_type) and inspect.isclass(target_type):
+                return issubclass(source_type, target_type)
+        except TypeError:
+            # Handle generic types that can't be used with issubclass
+            pass
+
+        # Handle Union types
+        source_origin = get_origin(source_type)
+        target_origin = get_origin(target_type)
+
+        if target_origin in (Union, UnionType):
+            target_args = get_args(target_type)
+            return any(WorkflowGraphValidator._is_type_compatible(source_type, arg) for arg in target_args)
+
+        if source_origin in (Union, UnionType):
+            source_args = get_args(source_type)
+            return all(WorkflowGraphValidator._is_type_compatible(arg, target_type) for arg in source_args)
+
+        # Handle generic types
+        if source_origin is not None and target_origin is not None and source_origin == target_origin:
+            source_args = get_args(source_type)
+            target_args = get_args(target_type)
+            if len(source_args) == len(target_args):
+                return all(
+                    WorkflowGraphValidator._is_type_compatible(s_arg, t_arg)
+                    for s_arg, t_arg in zip(source_args, target_args, strict=True)
+                )
+
+        # No other special compatibility cases
+        return False
+
+    # endregion
+
+
+# endregion
+
+
+def validate_workflow_graph(
+    edge_groups: Sequence[EdgeGroup],
+    executors: dict[str, Executor],
+    start_executor: Executor | str,
+    *,
+    duplicate_executor_ids: Sequence[str] | None = None,
+) -> None:
+    """Convenience function to validate a workflow graph.
+
+    Args:
+        edge_groups: list of edge groups in the workflow
+        executors: Map of executor IDs to executor instances
+        start_executor: The starting executor (can be instance or ID)
+        duplicate_executor_ids: Optional list of known duplicate executor IDs to pre-populate
+
+    Raises:
+        WorkflowValidationError: If any validation fails
+    """
+    validator = WorkflowGraphValidator()
+    validator.validate_workflow(
+        edge_groups,
+        executors,
+        start_executor,
+        duplicate_executor_ids=duplicate_executor_ids,
+    )

python/packages/core/agent_framework/_workflow/_viz.py

@@ -0,0 +1,350 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import hashlib
+import re
+import tempfile
+import uuid
+from pathlib import Path
+from typing import Literal
+
+from ._edge import FanInEdgeGroup
+from ._workflow import Workflow
+
+# Import of WorkflowExecutor is performed lazily inside methods to avoid cycles
+
+"""Workflow visualization module using graphviz."""
+
+
+class WorkflowViz:
+    """A class for visualizing workflows using graphviz."""
+
+    def __init__(self, workflow: Workflow):
+        """Initialize the WorkflowViz with a workflow.
+
+        Args:
+            workflow: The workflow to visualize.
+        """
+        self._workflow = workflow
+
+    def to_digraph(self) -> str:
+        """Export the workflow as a DOT format digraph string.
+
+        Returns:
+            A string representation of the workflow in DOT format.
+        """
+        lines = ["digraph Workflow {"]
+        lines.append("  rankdir=TD;")  # Top to bottom layout
+        lines.append("  node [shape=box, style=filled, fillcolor=lightblue];")
+        lines.append("  edge [color=black, arrowhead=vee];")
+        lines.append("")
+
+        # Emit the top-level workflow nodes/edges
+        self._emit_workflow_digraph(self._workflow, lines, indent="  ")
+
+        # Emit sub-workflows hosted by WorkflowExecutor as nested clusters
+        self._emit_sub_workflows_digraph(self._workflow, lines, indent="  ")
+
+        lines.append("}")
+        return "\n".join(lines)
+
+    def export(self, format: Literal["svg", "png", "pdf", "dot"] = "svg", filename: str | None = None) -> str:
+        """Export the workflow visualization to a file or return the file path.
+
+        Args:
+            format: The output format. Supported formats: 'svg', 'png', 'pdf', 'dot'.
+            filename: Optional filename to save the output. If None, creates a temporary file.
+
+        Returns:
+            The path to the saved file.
+
+        Raises:
+            ImportError: If graphviz is not installed.
+            ValueError: If an unsupported format is specified.
+        """
+        # Validate format first
+        if format not in ["svg", "png", "pdf", "dot"]:
+            raise ValueError(f"Unsupported format: {format}. Supported formats: svg, png, pdf, dot")
+
+        if format == "dot":
+            content = self.to_digraph()
+            if filename:
+                with open(filename, "w", encoding="utf-8") as f:
+                    f.write(content)
+                return filename
+            # Create temporary file for dot format
+            with tempfile.NamedTemporaryFile(mode="w", suffix=".dot", delete=False, encoding="utf-8") as temp_file:
+                temp_file.write(content)
+                return temp_file.name
+
+        try:
+            import graphviz  # type: ignore
+        except ImportError as e:
+            raise ImportError(
+                "viz extra is required for export. Install it with: pip install agent-framework[viz]. "
+                "You also need to install graphviz separately. E.g., sudo apt-get install graphviz on Debian/Ubuntu "
+                "or brew install graphviz on macOS. See https://graphviz.org/download/ for details."
+            ) from e
+
+        # Create a temporary graphviz Source object
+        dot_content = self.to_digraph()
+        source = graphviz.Source(dot_content)
+
+        try:
+            if filename:
+                # Save to specified file
+                output_path = Path(filename)
+                if output_path.suffix and output_path.suffix[1:] != format:
+                    raise ValueError(f"File extension {output_path.suffix} doesn't match format {format}")
+
+                # Remove extension if present since graphviz.render() adds it
+                base_name = str(output_path.with_suffix(""))
+                source.render(base_name, format=format, cleanup=True)
+
+                # Return the actual filename with extension
+                return f"{base_name}.{format}"
+            # Create temporary file
+            with tempfile.NamedTemporaryFile(suffix=f".{format}", delete=False) as temp_file:
+                temp_path = Path(temp_file.name)
+                base_name = str(temp_path.with_suffix(""))
+
+            source.render(base_name, format=format, cleanup=True)
+            return f"{base_name}.{format}"
+        except graphviz.backend.execute.ExecutableNotFound as e:
+            raise ImportError(
+                "The graphviz executables are not found. The graphviz Python package is installed, but the "
+                "graphviz executables (dot, neato, etc.) are not available on your system's PATH. "
+                "Install graphviz executables: sudo apt-get install graphviz on Debian/Ubuntu, "
+                "brew install graphviz on macOS, or download from https://graphviz.org/download/ for other platforms."
+            ) from e
+
+    def save_svg(self, filename: str) -> str:
+        """Convenience method to save as SVG.
+
+        Args:
+            filename: The filename to save the SVG file.
+
+        Returns:
+            The path to the saved SVG file.
+        """
+        return self.export(format="svg", filename=filename)
+
+    def save_png(self, filename: str) -> str:
+        """Convenience method to save as PNG.
+
+        Args:
+            filename: The filename to save the PNG file.
+
+        Returns:
+            The path to the saved PNG file.
+        """
+        return self.export(format="png", filename=filename)
+
+    def save_pdf(self, filename: str) -> str:
+        """Convenience method to save as PDF.
+
+        Args:
+            filename: The filename to save the PDF file.
+
+        Returns:
+            The path to the saved PDF file.
+        """
+        return self.export(format="pdf", filename=filename)
+
+    def to_mermaid(self) -> str:
+        """Export the workflow as a Mermaid flowchart string.
+
+        Returns:
+            A string representation of the workflow in Mermaid flowchart syntax.
+        """
+
+        def _san(s: str) -> str:
+            """Sanitize an ID for Mermaid (alphanumeric and underscore, start with letter)."""
+            s2 = re.sub(r"[^0-9A-Za-z_]", "_", s)
+            if not s2 or not s2[0].isalpha():
+                s2 = f"n_{s2}"
+            return s2
+
+        lines: list[str] = ["flowchart TD"]
+
+        # Emit top-level workflow
+        self._emit_workflow_mermaid(self._workflow, lines, indent="  ")
+
+        # Emit sub-workflows as Mermaid subgraphs
+        self._emit_sub_workflows_mermaid(self._workflow, lines, indent="  ")
+
+        return "\n".join(lines)
+
+    # region Private helpers
+
+    def _fan_in_digest(self, target: str, sources: list[str]) -> str:
+        sources_sorted = sorted(sources)
+        return hashlib.sha256((target + "|" + "|".join(sources_sorted)).encode("utf-8")).hexdigest()[:8]
+
+    def _compute_fan_in_descriptors(self, wf: Workflow | None = None) -> list[tuple[str, list[str], str]]:
+        """Return list of (node_id, sources, target) for fan-in groups.
+
+        node_id is DOT-oriented: fan_in::target::digest
+        """
+        result: list[tuple[str, list[str], str]] = []
+        workflow = wf or self._workflow
+        for group in workflow.edge_groups:
+            if isinstance(group, FanInEdgeGroup):
+                target = group.target_executor_ids[0]
+                sources = list(group.source_executor_ids)
+                digest = self._fan_in_digest(target, sources)
+                node_id = f"fan_in::{target}::{digest}"
+                result.append((node_id, sorted(sources), target))
+        return result
+
+    def _compute_normal_edges(self, wf: Workflow | None = None) -> list[tuple[str, str, bool]]:
+        """Return list of (source_id, target_id, is_conditional) for non-fan-in groups."""
+        edges: list[tuple[str, str, bool]] = []
+        workflow = wf or self._workflow
+        for group in workflow.edge_groups:
+            if isinstance(group, FanInEdgeGroup):
+                continue
+            for edge in group.edges:
+                is_cond = getattr(edge, "_condition", None) is not None
+                edges.append((edge.source_id, edge.target_id, is_cond))
+        return edges
+
+    # endregion
+
+    # region Internal emitters (DOT)
+
+    def _emit_workflow_digraph(self, wf: Workflow, lines: list[str], indent: str, ns: str | None = None) -> None:
+        """Emit DOT nodes/edges for the given workflow.
+
+        If ns (namespace) is provided, node ids are prefixed with f"{ns}/" for uniqueness,
+        but labels remain the original executor ids.
+        """
+
+        def map_id(x: str) -> str:
+            return f"{ns}/{x}" if ns else x
+
+        # Nodes
+        start_executor_id = wf.start_executor_id
+        lines.append(
+            f'{indent}"{map_id(start_executor_id)}" [fillcolor=lightgreen, label="{start_executor_id}\\n(Start)"];'
+        )
+        for executor_id in wf.executors:
+            if executor_id != start_executor_id:
+                lines.append(f'{indent}"{map_id(executor_id)}" [label="{executor_id}"];')
+
+        # Fan-in nodes
+        fan_in_nodes = self._compute_fan_in_descriptors(wf)
+        if fan_in_nodes:
+            lines.append("")
+            for node_id, _, _ in fan_in_nodes:
+                lines.append(f'{indent}"{map_id(node_id)}" [shape=ellipse, fillcolor=lightgoldenrod, label="fan-in"];')
+
+        # Fan-in edges
+        for node_id, sources, target in fan_in_nodes:
+            for src in sources:
+                lines.append(f'{indent}"{map_id(src)}" -> "{map_id(node_id)}";')
+            lines.append(f'{indent}"{map_id(node_id)}" -> "{map_id(target)}";')
+
+        # Normal edges
+        for src, tgt, is_cond in self._compute_normal_edges(wf):
+            edge_attr = ' [style=dashed, label="conditional"]' if is_cond else ""
+            lines.append(f'{indent}"{map_id(src)}" -> "{map_id(tgt)}"{edge_attr};')
+
+    def _emit_sub_workflows_digraph(self, wf: Workflow, lines: list[str], indent: str) -> None:
+        """Emit DOT subgraphs for any WorkflowExecutor instances found in the workflow."""
+        # Lazy import to avoid any potential import cycles
+        try:
+            from ._workflow_executor import WorkflowExecutor  # type: ignore
+        except ImportError:  # pragma: no cover - best-effort; if unavailable, skip subgraphs
+            return
+
+        for exec_id, exec_obj in wf.executors.items():
+            if isinstance(exec_obj, WorkflowExecutor) and hasattr(exec_obj, "workflow") and exec_obj.workflow:
+                subgraph_id = f"cluster_{uuid.uuid5(uuid.NAMESPACE_OID, exec_id).hex[:8]}"
+                lines.append(f"{indent}subgraph {subgraph_id} {{")
+                lines.append(f'{indent}  label="sub-workflow: {exec_id}";')
+                lines.append(f"{indent}  style=dashed;")
+
+                # Emit the nested workflow inside this cluster using a namespace
+                ns = exec_id
+                self._emit_workflow_digraph(exec_obj.workflow, lines, indent=f"{indent}  ", ns=ns)
+
+                # Recurse into deeper nested sub-workflows
+                self._emit_sub_workflows_digraph(exec_obj.workflow, lines, indent=f"{indent}  ")
+
+                lines.append(f"{indent}}}")
+
+    # endregion
+
+    # region Internal emitters (Mermaid)
+
+    def _emit_workflow_mermaid(self, wf: Workflow, lines: list[str], indent: str, ns: str | None = None) -> None:
+        def _san(s: str) -> str:
+            s2 = re.sub(r"[^0-9A-Za-z_]", "_", s)
+            if not s2 or not s2[0].isalpha():
+                s2 = f"n_{s2}"
+            return s2
+
+        def map_id(x: str) -> str:
+            if ns:
+                return f"{_san(ns)}__{_san(x)}"
+            return _san(x)
+
+        # Nodes
+        start_executor_id = wf.start_executor_id
+        lines.append(f'{indent}{map_id(start_executor_id)}["{start_executor_id} (Start)"];')
+        for executor_id in wf.executors:
+            if executor_id == start_executor_id:
+                continue
+            lines.append(f'{indent}{map_id(executor_id)}["{executor_id}"];')
+
+        # Fan-in nodes
+        fan_in_nodes_dot = self._compute_fan_in_descriptors(wf)
+        fan_in_nodes: list[tuple[str, list[str], str]] = []
+        for dot_node_id, sources, target in fan_in_nodes_dot:
+            digest = dot_node_id.split("::")[-1]
+            base = f"{target}__{digest}"
+            fan_node_id = f"fan_in__{_san(ns) + '__' if ns else ''}{_san(base)}"
+            fan_in_nodes.append((fan_node_id, sources, target))
+
+        for fan_node_id, _, _ in fan_in_nodes:
+            # Keep this line without trailing semicolon to match existing tests
+            lines.append(f"{indent}{fan_node_id}((fan-in))")
+
+        # Fan-in edges
+        for fan_node_id, sources, target in fan_in_nodes:
+            for s in sources:
+                lines.append(f"{indent}{map_id(s)} --> {fan_node_id};")
+            lines.append(f"{indent}{fan_node_id} --> {map_id(target)};")
+
+        # Normal edges
+        for src, tgt, is_cond in self._compute_normal_edges(wf):
+            s = map_id(src)
+            t = map_id(tgt)
+            if is_cond:
+                lines.append(f"{indent}{s} -. conditional .-> {t};")
+            else:
+                lines.append(f"{indent}{s} --> {t};")
+
+    def _emit_sub_workflows_mermaid(self, wf: Workflow, lines: list[str], indent: str) -> None:
+        try:
+            from ._workflow_executor import WorkflowExecutor  # type: ignore
+        except ImportError:  # pragma: no cover
+            return
+
+        def _san(s: str) -> str:
+            s2 = re.sub(r"[^0-9A-Za-z_]", "_", s)
+            if not s2 or not s2[0].isalpha():
+                s2 = f"n_{s2}"
+            return s2
+
+        for exec_id, exec_obj in wf.executors.items():
+            if isinstance(exec_obj, WorkflowExecutor) and hasattr(exec_obj, "workflow") and exec_obj.workflow:
+                sg_id = _san(exec_id)
+                lines.append(f"{indent}subgraph {sg_id}")
+                # Render nested workflow within this subgraph using namespacing
+                self._emit_workflow_mermaid(exec_obj.workflow, lines, indent=f"{indent}  ", ns=exec_id)
+                # Recurse into deeper sub-workflows
+                self._emit_sub_workflows_mermaid(exec_obj.workflow, lines, indent=f"{indent}  ")
+                lines.append(f"{indent}end")
+
+    # endregion

python/packages/core/agent_framework/_workflow/_workflow_context.py

@@ -0,0 +1,447 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import inspect
+import logging
+from collections.abc import Callable
+from types import UnionType
+from typing import Any, Generic, Union, cast, get_args, get_origin
+
+from opentelemetry.propagate import inject
+from opentelemetry.trace import SpanKind
+from typing_extensions import Never, TypeVar
+
+from ..observability import OtelAttr, create_workflow_span
+from ._events import (
+    WorkflowEvent,
+    WorkflowEventSource,
+    WorkflowFailedEvent,
+    WorkflowLifecycleEvent,
+    WorkflowOutputEvent,
+    WorkflowStartedEvent,
+    WorkflowStatusEvent,
+    WorkflowWarningEvent,
+    _framework_event_origin,
+)
+from ._runner_context import Message, RunnerContext
+from ._shared_state import SharedState
+
+T_Out = TypeVar("T_Out", default=Never)
+T_W_Out = TypeVar("T_W_Out", default=Never)
+
+
+logger = logging.getLogger(__name__)
+
+
+def infer_output_types_from_ctx_annotation(ctx_annotation: Any) -> tuple[list[type[Any]], list[type[Any]]]:
+    """Infer message types and workflow output types from the WorkflowContext generic parameters.
+
+    Examples:
+    - WorkflowContext -> ([], [])
+    - WorkflowContext[str] -> ([str], [])
+    - WorkflowContext[str, int] -> ([str], [int])
+    - WorkflowContext[str | int, bool | int] -> ([str, int], [bool, int])
+    - WorkflowContext[Union[str, int], Union[bool, int]] -> ([str, int], [bool, int])
+    - WorkflowContext[Any] -> ([Any], [])
+    - WorkflowContext[Any, Any] -> ([Any], [Any])
+    - WorkflowContext[Never, Never] -> ([], [])
+    - WorkflowContext[Never, int] -> ([], [int])
+
+    Returns:
+        Tuple of (message_types, workflow_output_types)
+    """
+    # If no annotation or not parameterized, return empty lists
+    try:
+        origin = get_origin(ctx_annotation)
+    except Exception:
+        origin = None
+
+    # If annotation is unsubscripted WorkflowContext, nothing to infer
+    if origin is None:
+        return [], []
+
+    # Expecting WorkflowContext[T_Out, T_W_Out]
+    if origin is not WorkflowContext:
+        return [], []
+
+    args = list(get_args(ctx_annotation))
+    if not args:
+        return [], []
+
+    # WorkflowContext[T_Out] -> message_types from T_Out, no workflow output types
+    if len(args) == 1:
+        t = args[0]
+        t_origin = get_origin(t)
+        if t is Any:
+            return [cast(type[Any], Any)], []
+
+        if t_origin in (Union, UnionType):
+            message_types = [arg for arg in get_args(t) if arg is not Any and arg is not Never]
+            return message_types, []
+
+        if t is Never:
+            return [], []
+        return [t], []
+
+    # WorkflowContext[T_Out, T_W_Out] -> message_types from T_Out, workflow_output_types from T_W_Out
+    t_out, t_w_out = args[:2]  # Take first two args in case there are more
+
+    # Process T_Out for message_types
+    message_types = []
+    t_out_origin = get_origin(t_out)
+    if t_out is Any:
+        message_types = [cast(type[Any], Any)]
+    elif t_out is not Never:
+        if t_out_origin in (Union, UnionType):
+            message_types = [arg for arg in get_args(t_out) if arg is not Any and arg is not Never]
+        else:
+            message_types = [t_out]
+
+    # Process T_W_Out for workflow_output_types
+    workflow_output_types = []
+    t_w_out_origin = get_origin(t_w_out)
+    if t_w_out is Any:
+        workflow_output_types = [cast(type[Any], Any)]
+    elif t_w_out is not Never:
+        if t_w_out_origin in (Union, UnionType):
+            workflow_output_types = [arg for arg in get_args(t_w_out) if arg is not Any and arg is not Never]
+        else:
+            workflow_output_types = [t_w_out]
+
+    return message_types, workflow_output_types
+
+
+def _is_workflow_context_type(annotation: Any) -> bool:
+    """Check if an annotation represents WorkflowContext, WorkflowContext[T], or WorkflowContext[T, U]."""
+    origin = get_origin(annotation)
+    if origin is WorkflowContext:
+        return True
+    # Also handle the case where the raw class is used
+    return annotation is WorkflowContext
+
+
+def validate_workflow_context_annotation(
+    annotation: Any,
+    parameter_name: str,
+    context_description: str,
+) -> tuple[list[type[Any]], list[type[Any]]]:
+    """Validate a WorkflowContext annotation and return inferred types.
+
+    Args:
+        annotation: The type annotation to validate
+        parameter_name: Name of the parameter (for error messages)
+        context_description: Description of the context (e.g., "Function func1", "Handler method")
+
+    Returns:
+        Tuple of (output_types, workflow_output_types)
+
+    Raises:
+        ValueError: If the annotation is invalid
+    """
+    if annotation == inspect.Parameter.empty:
+        raise ValueError(
+            f"{context_description} {parameter_name} must have a WorkflowContext, "
+            f"WorkflowContext[T] or WorkflowContext[T, U] type annotation, "
+            f"where T is output message type and U is workflow output type"
+        )
+
+    if not _is_workflow_context_type(annotation):
+        raise ValueError(
+            f"{context_description} {parameter_name} must be annotated as "
+            f"WorkflowContext, WorkflowContext[T], or WorkflowContext[T, U], "
+            f"got {annotation}"
+        )
+
+    # Validate type arguments for WorkflowContext[T] or WorkflowContext[T, U]
+    type_args = get_args(annotation)
+
+    if len(type_args) > 2:
+        raise ValueError(
+            f"{context_description} {parameter_name} must have at most 2 type arguments, "
+            "WorkflowContext, WorkflowContext[T], or WorkflowContext[T, U], "
+            f"got {len(type_args)} arguments"
+        )
+
+    if type_args:
+        # Helper function to check if a value is a valid type annotation
+        def _is_type_like(x: Any) -> bool:
+            """Check if a value is a type-like entity (class, type, or typing construct)."""
+            return isinstance(x, type) or get_origin(x) is not None or x is Never
+
+        for i, type_arg in enumerate(type_args):
+            param_description = "T_Out" if i == 0 else "T_W_Out"
+
+            # Allow Any explicitly
+            if type_arg is Any:
+                continue
+
+            # Check if it's a union type and validate each member
+            union_origin = get_origin(type_arg)
+            if union_origin in (Union, UnionType):
+                union_members = get_args(type_arg)
+                invalid_members = [m for m in union_members if not _is_type_like(m) and m is not Any]
+                if invalid_members:
+                    raise ValueError(
+                        f"{context_description} {parameter_name} {param_description} "
+                        f"contains invalid type entries: {invalid_members}. "
+                        f"Use proper types or typing generics"
+                    )
+            else:
+                # Check if it's a valid type
+                if not _is_type_like(type_arg):
+                    raise ValueError(
+                        f"{context_description} {parameter_name} {param_description} "
+                        f"contains invalid type entry: {type_arg}. "
+                        f"Use proper types or typing generics"
+                    )
+
+    return infer_output_types_from_ctx_annotation(annotation)
+
+
+def validate_function_signature(
+    func: Callable[..., Any], context_description: str
+) -> tuple[type, Any, list[type[Any]], list[type[Any]]]:
+    """Validate function signature for executor functions.
+
+    Args:
+        func: The function to validate
+        context_description: Description for error messages (e.g., "Function", "Handler method")
+
+    Returns:
+        Tuple of (message_type, ctx_annotation, output_types, workflow_output_types)
+
+    Raises:
+        ValueError: If the function signature is invalid
+    """
+    signature = inspect.signature(func)
+    params = list(signature.parameters.values())
+
+    # Determine expected parameter count based on context
+    expected_counts: tuple[int, ...]
+    if context_description.startswith("Function"):
+        # Function executor: (message) or (message, ctx)
+        expected_counts = (1, 2)
+        param_description = "(message: T) or (message: T, ctx: WorkflowContext[U])"
+    else:
+        # Handler method: (self, message, ctx)
+        expected_counts = (3,)
+        param_description = "(self, message: T, ctx: WorkflowContext[U])"
+
+    if len(params) not in expected_counts:
+        raise ValueError(
+            f"{context_description} {func.__name__} must have {param_description}. Got {len(params)} parameters."
+        )
+
+    # Extract message parameter (index 0 for functions, index 1 for methods)
+    message_param_idx = 0 if context_description.startswith("Function") else 1
+    message_param = params[message_param_idx]
+
+    # Check message parameter has type annotation
+    if message_param.annotation == inspect.Parameter.empty:
+        raise ValueError(f"{context_description} {func.__name__} must have a type annotation for the message parameter")
+
+    message_type = message_param.annotation
+
+    # Check if there's a context parameter
+    ctx_param_idx = message_param_idx + 1
+    if len(params) > ctx_param_idx:
+        ctx_param = params[ctx_param_idx]
+        output_types, workflow_output_types = validate_workflow_context_annotation(
+            ctx_param.annotation, f"parameter '{ctx_param.name}'", context_description
+        )
+        ctx_annotation = ctx_param.annotation
+    else:
+        # No context parameter (only valid for function executors)
+        if not context_description.startswith("Function"):
+            raise ValueError(f"{context_description} {func.__name__} must have a WorkflowContext parameter")
+        output_types, workflow_output_types = [], []
+        ctx_annotation = None
+
+    return message_type, ctx_annotation, output_types, workflow_output_types
+
+
+_FRAMEWORK_LIFECYCLE_EVENT_TYPES: tuple[type[WorkflowEvent], ...] = cast(
+    tuple[type[WorkflowEvent], ...],
+    tuple(get_args(WorkflowLifecycleEvent))
+    or (
+        WorkflowStartedEvent,
+        WorkflowStatusEvent,
+        WorkflowFailedEvent,
+    ),
+)
+
+
+class WorkflowContext(Generic[T_Out, T_W_Out]):
+    """Execution context that enables executors to interact with workflows and other executors.
+
+    ## Overview
+    WorkflowContext provides a controlled interface for executors to send messages, yield outputs,
+    manage state, and interact with the broader workflow ecosystem. It enforces type safety through
+    generic parameters while preventing direct access to internal runtime components.
+
+    ## Type Parameters
+    The context is parameterized to enforce type safety for different operations:
+
+    ### WorkflowContext (no parameters)
+    For executors that only perform side effects without sending messages or yielding outputs:
+    ```python
+    async def log_handler(message: str, ctx: WorkflowContext) -> None:
+        print(f"Received: {message}")  # Only side effects
+    ```
+
+    ### WorkflowContext[T_Out]
+    Enables sending messages of type T_Out to other executors:
+    ```python
+    async def processor(message: str, ctx: WorkflowContext[int]) -> None:
+        result = len(message)
+        await ctx.send_message(result)  # Send int to downstream executors
+    ```
+
+    ### WorkflowContext[T_Out, T_W_Out]
+    Enables both sending messages (T_Out) and yielding workflow outputs (T_W_Out):
+    ```python
+    async def dual_output(message: str, ctx: WorkflowContext[int, str]) -> None:
+        await ctx.send_message(42)  # Send int message
+        await ctx.yield_output("complete")  # Yield str workflow output
+    ```
+
+    ### Union Types
+    Multiple types can be specified using union notation:
+    ```python
+    async def flexible(message: str, ctx: WorkflowContext[int | str, bool | dict]) -> None:
+        await ctx.send_message("text")  # or send 42
+        await ctx.yield_output(True)  # or yield {"status": "done"}
+    ```
+    """
+
+    def __init__(
+        self,
+        executor_id: str,
+        source_executor_ids: list[str],
+        shared_state: SharedState,
+        runner_context: RunnerContext,
+        trace_contexts: list[dict[str, str]] | None = None,
+        source_span_ids: list[str] | None = None,
+    ):
+        """Initialize the executor context with the given workflow context.
+
+        Args:
+            executor_id: The unique identifier of the executor that this context belongs to.
+            source_executor_ids: The IDs of the source executors that sent messages to this executor.
+                This is a list to support fan_in scenarios where multiple sources send aggregated
+                messages to the same executor.
+            shared_state: The shared state for the workflow.
+            runner_context: The runner context that provides methods to send messages and events.
+            trace_contexts: Optional trace contexts from multiple sources for OpenTelemetry propagation.
+            source_span_ids: Optional source span IDs from multiple sources for linking (not for nesting).
+        """
+        self._executor_id = executor_id
+        self._source_executor_ids = source_executor_ids
+        self._runner_context = runner_context
+        self._shared_state = shared_state
+
+        # Store trace contexts and source span IDs for linking (supporting multiple sources)
+        self._trace_contexts = trace_contexts or []
+        self._source_span_ids = source_span_ids or []
+
+        if not self._source_executor_ids:
+            raise ValueError("source_executor_ids cannot be empty. At least one source executor ID is required.")
+
+    async def send_message(self, message: T_Out, target_id: str | None = None) -> None:
+        """Send a message to the workflow context.
+
+        Args:
+            message: The message to send. This must conform to the output type(s) declared on this context.
+            target_id: The ID of the target executor to send the message to.
+                       If None, the message will be sent to all target executors.
+        """
+        global OBSERVABILITY_SETTINGS
+        from ..observability import OBSERVABILITY_SETTINGS
+
+        # Create publishing span (inherits current trace context automatically)
+        attributes: dict[str, str] = {OtelAttr.MESSAGE_TYPE: type(message).__name__}
+        if target_id:
+            attributes[OtelAttr.MESSAGE_DESTINATION_EXECUTOR_ID] = target_id
+        with create_workflow_span(OtelAttr.MESSAGE_SEND_SPAN, attributes, kind=SpanKind.PRODUCER) as span:
+            # Create Message wrapper
+            msg = Message(data=message, source_id=self._executor_id, target_id=target_id)
+
+            # Inject current trace context if tracing enabled
+            if OBSERVABILITY_SETTINGS.ENABLED and span and span.is_recording():  # type: ignore[name-defined]
+                trace_context: dict[str, str] = {}
+                inject(trace_context)  # Inject current trace context for message propagation
+
+                msg.trace_contexts = [trace_context]
+                msg.source_span_ids = [format(span.get_span_context().span_id, "016x")]
+
+            await self._runner_context.send_message(msg)
+
+    async def yield_output(self, output: T_W_Out) -> None:
+        """Set the output of the workflow.
+
+        Args:
+            output: The output to yield. This must conform to the workflow output type(s)
+                    declared on this context.
+        """
+        with _framework_event_origin():
+            event = WorkflowOutputEvent(data=output, source_executor_id=self._executor_id)
+        await self._runner_context.add_event(event)
+
+    async def add_event(self, event: WorkflowEvent) -> None:
+        """Add an event to the workflow context."""
+        if event.origin == WorkflowEventSource.EXECUTOR and isinstance(event, _FRAMEWORK_LIFECYCLE_EVENT_TYPES):
+            event_name = event.__class__.__name__
+            warning_msg = (
+                f"Executor '{self._executor_id}' attempted to emit {event_name}, "
+                "which is reserved for framework lifecycle notifications. The "
+                "event was ignored."
+            )
+            logger.warning(warning_msg)
+            await self._runner_context.add_event(WorkflowWarningEvent(warning_msg))
+            return
+        await self._runner_context.add_event(event)
+
+    async def get_shared_state(self, key: str) -> Any:
+        """Get a value from the shared state."""
+        return await self._shared_state.get(key)
+
+    async def set_shared_state(self, key: str, value: Any) -> None:
+        """Set a value in the shared state."""
+        await self._shared_state.set(key, value)
+
+    def get_source_executor_id(self) -> str:
+        """Get the ID of the source executor that sent the message to this executor.
+
+        Raises:
+            RuntimeError: If there are multiple source executors, this method raises an error.
+        """
+        if len(self._source_executor_ids) > 1:
+            raise RuntimeError(
+                "Cannot get source executor ID when there are multiple source executors. "
+                "Access the full list via the source_executor_ids property instead."
+            )
+        return self._source_executor_ids[0]
+
+    @property
+    def source_executor_ids(self) -> list[str]:
+        """Get the IDs of the source executors that sent messages to this executor."""
+        return self._source_executor_ids
+
+    @property
+    def shared_state(self) -> SharedState:
+        """Get the shared state."""
+        return self._shared_state
+
+    async def set_state(self, state: dict[str, Any]) -> None:
+        """Persist this executor's state into the checkpointable context.
+
+        Executors call this with a JSON-serializable dict capturing the minimal
+        state needed to resume. It replaces any previously stored state.
+        """
+        if hasattr(self._runner_context, "set_state"):
+            await self._runner_context.set_state(self._executor_id, state)  # type: ignore[arg-type]
+
+    async def get_state(self) -> dict[str, Any] | None:
+        """Retrieve previously persisted state for this executor, if any."""
+        if hasattr(self._runner_context, "get_state"):
+            return await self._runner_context.get_state(self._executor_id)  # type: ignore[return-value]
+        return None

python/packages/core/agent_framework/_workflow/_workflow_executor.py

@@ -0,0 +1,659 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import contextlib
+import inspect
+import logging
+import uuid
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from ._workflow import Workflow
+
+from ._events import (
+    RequestInfoEvent,
+    WorkflowErrorEvent,
+    WorkflowFailedEvent,
+    WorkflowRunState,
+)
+from ._executor import (
+    Executor,
+    RequestInfoExecutor,
+    RequestInfoMessage,
+    RequestResponse,
+    handler,
+)
+from ._typing_utils import is_instance_of
+from ._workflow_context import WorkflowContext
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ExecutionContext:
+    """Context for tracking a single sub-workflow execution."""
+
+    execution_id: str
+    collected_responses: dict[str, Any]  # request_id -> response_data
+    expected_response_count: int
+    pending_requests: dict[str, Any]  # request_id -> original request data
+
+
+class WorkflowExecutor(Executor):
+    """An executor that wraps a workflow to enable hierarchical workflow composition.
+
+    ## Overview
+    WorkflowExecutor makes a workflow behave as a single executor within a parent workflow,
+    enabling nested workflow architectures. It handles the complete lifecycle of sub-workflow
+    execution including event processing, output forwarding, and request/response coordination
+    between parent and child workflows.
+
+    ## Execution Model
+    When invoked, WorkflowExecutor:
+    1. Starts the wrapped workflow with the input message
+    2. Runs the sub-workflow to completion or until it needs external input
+    3. Processes the sub-workflow's complete event stream after execution
+    4. Forwards outputs to the parent workflow's event stream
+    5. Handles external requests by routing them to the parent workflow
+    6. Accumulates responses and resumes sub-workflow execution
+
+    ## Event Stream Processing
+    WorkflowExecutor processes events after sub-workflow completion:
+
+    ### Output Forwarding
+    All outputs from the sub-workflow are automatically forwarded to the parent:
+    ```python
+    # Sub-workflow yields outputs
+    await ctx.yield_output("sub-workflow result")
+
+    # WorkflowExecutor forwards to parent via ctx.send_message()
+    # Parent receives the output as a regular message
+    ```
+
+    ### Request/Response Coordination
+    When sub-workflows need external information:
+    ```python
+    # Sub-workflow makes request
+    request = MyDataRequest(query="user info")
+    # RequestInfoExecutor emits RequestInfoEvent
+
+    # WorkflowExecutor sets source_executor_id and forwards to parent
+    request.source_executor_id = "child_workflow_executor_id"
+    # Parent workflow can handle via @handler for RequestInfoMessage subclasses,
+    # or directly forward to external source via a RequestInfoExecutor in the parent
+    # workflow.
+    ```
+
+    ### State Management
+    WorkflowExecutor maintains execution state across request/response cycles:
+    - Tracks pending requests by request_id
+    - Accumulates responses until all expected responses are received
+    - Resumes sub-workflow execution with complete response batch
+    - Handles concurrent executions and multiple pending requests
+
+    ## Type System Integration
+    WorkflowExecutor inherits its type signature from the wrapped workflow:
+
+    ### Input Types
+    Matches the wrapped workflow's start executor input types:
+    ```python
+    # If sub-workflow accepts str, WorkflowExecutor accepts str
+    workflow_executor = WorkflowExecutor(my_workflow, id="wrapper")
+    assert workflow_executor.input_types == my_workflow.input_types
+    ```
+
+    ### Output Types
+    Combines sub-workflow outputs with request coordination types:
+    ```python
+    # Includes all sub-workflow output types
+    # Plus RequestInfoMessage if sub-workflow can make requests
+    output_types = workflow.output_types + [RequestInfoMessage]  # if applicable
+    ```
+
+    ## Error Handling
+    WorkflowExecutor propagates sub-workflow failures:
+    - Captures WorkflowFailedEvent from sub-workflow
+    - Converts to WorkflowErrorEvent in parent context
+    - Provides detailed error information including sub-workflow ID
+
+    ## Concurrent Execution Support
+    WorkflowExecutor fully supports multiple concurrent sub-workflow executions:
+
+    ### Per-Execution State Isolation
+    Each sub-workflow invocation creates an isolated ExecutionContext:
+    ```python
+    # Multiple concurrent invocations are supported
+    workflow_executor = WorkflowExecutor(my_workflow, id="concurrent_executor")
+
+    # Each invocation gets its own execution context
+    # Execution 1: processes input_1 independently
+    # Execution 2: processes input_2 independently
+    # No state interference between executions
+    ```
+
+    ### Request/Response Coordination
+    Responses are correctly routed to the originating execution:
+    - Each execution tracks its own pending requests and expected responses
+    - Request-to-execution mapping ensures responses reach the correct sub-workflow
+    - Response accumulation is isolated per execution
+    - Automatic cleanup when execution completes
+
+    ### Memory Management
+    - Unlimited concurrent executions supported
+    - Each execution has unique UUID-based identification
+    - Cleanup of completed execution contexts
+    - Thread-safe state management for concurrent access
+
+    ### Important Considerations
+    **Shared Workflow Instance**: All concurrent executions use the same underlying workflow instance.
+    For proper isolation, ensure that:
+    - The wrapped workflow and its executors are stateless
+    - Executors use WorkflowContext state management instead of instance variables
+    - Any shared state is managed through WorkflowContext.get_shared_state/set_shared_state
+    ```python
+    # Good: Stateless executor using context state
+    class StatelessExecutor(Executor):
+        @handler
+        async def process(self, data: str, ctx: WorkflowContext[str]) -> None:
+            # Use context state instead of instance variables
+            state = await ctx.get_state() or {}
+            state["processed"] = data
+            await ctx.set_state(state)
+
+
+    # Avoid: Stateful executor with instance variables
+    class StatefulExecutor(Executor):
+        def __init__(self):
+            super().__init__(id="stateful")
+            self.data = []  # This will be shared across concurrent executions!
+    ```
+
+    ## Integration with Parent Workflows
+    Parent workflows can intercept sub-workflow requests:
+    ```python
+    class ParentExecutor(Executor):
+        @handler
+        async def handle_request(
+            self,
+            request: MyRequestType,  # Subclass of RequestInfoMessage
+            ctx: WorkflowContext[RequestResponse[RequestInfoMessage, Any] | RequestInfoMessage],
+        ) -> None:
+            # Handle request locally or forward to external source
+            if self.can_handle_locally(request):
+                # Send response back to sub-workflow
+                response = RequestResponse(data="local result", original_request=request, request_id=request.request_id)
+                await ctx.send_message(response, target_id=request.source_executor_id)
+            else:
+                # Forward to external handler
+                await ctx.send_message(request)
+    ```
+
+    ## Implementation Notes
+    - Sub-workflows run to completion before processing their results
+    - Event processing is atomic - all outputs are forwarded before requests
+    - Response accumulation ensures sub-workflows receive complete response batches
+    - Execution state is maintained for proper resumption after external requests
+    - Concurrent executions are fully isolated and do not interfere with each other
+    """
+
+    def __init__(self, workflow: "Workflow", id: str, **kwargs: Any):
+        """Initialize the WorkflowExecutor.
+
+        Args:
+            workflow: The workflow to execute as a sub-workflow.
+            id: Unique identifier for this executor.
+            **kwargs: Additional keyword arguments passed to the parent constructor.
+        """
+        super().__init__(id, **kwargs)
+        self.workflow = workflow
+
+        # Track execution contexts for concurrent sub-workflow executions
+        self._execution_contexts: dict[str, ExecutionContext] = {}  # execution_id -> ExecutionContext
+        # Map request_id to execution_id for response routing
+        self._request_to_execution: dict[str, str] = {}  # request_id -> execution_id
+        self._active_executions: int = 0  # Count of active sub-workflow executions
+        self._state_loaded: bool = False
+
+    @property
+    def input_types(self) -> list[type[Any]]:
+        """Get the input types based on the underlying workflow's input types plus WorkflowExecutor-specific types.
+
+        Returns:
+            A list of input types that the WorkflowExecutor can accept.
+        """
+        input_types = list(self.workflow.input_types)
+
+        # WorkflowExecutor can also handle RequestResponse for sub-workflow responses
+        if RequestResponse not in input_types:
+            input_types.append(RequestResponse)
+
+        return input_types
+
+    @property
+    def output_types(self) -> list[type[Any]]:
+        """Get the output types based on the underlying workflow's output types.
+
+        Returns:
+            A list of output types that the underlying workflow can produce.
+            Includes specific RequestInfoMessage subtypes if the sub-workflow contains RequestInfoExecutor.
+        """
+        output_types = list(self.workflow.output_types)
+
+        # Check if the sub-workflow contains a RequestInfoExecutor
+        # If so, collect the specific RequestInfoMessage subtypes from all executors
+        has_request_info_executor = any(
+            isinstance(executor, RequestInfoExecutor) for executor in self.workflow.executors.values()
+        )
+
+        if has_request_info_executor:
+            # Collect all RequestInfoMessage subtypes from executor output types
+            for executor in self.workflow.executors.values():
+                for output_type in executor.output_types:
+                    # Check if this is a RequestInfoMessage subclass
+                    if (
+                        inspect.isclass(output_type)
+                        and issubclass(output_type, RequestInfoMessage)
+                        and output_type not in output_types
+                    ):
+                        output_types.append(output_type)
+
+        return output_types
+
+    def to_dict(self) -> dict[str, Any]:
+        data = super().to_dict()
+        data["workflow"] = self.workflow.to_dict()
+        return data
+
+    def can_handle(self, message: Any) -> bool:
+        """Override can_handle to only accept messages that the wrapped workflow can handle.
+
+        This prevents the WorkflowExecutor from accepting messages that should go to other
+        executors (like RequestInfoExecutor).
+        """
+        # Always handle RequestResponse (for the handle_response handler)
+        if isinstance(message, RequestResponse):
+            return True
+
+        # For other messages, only handle if the wrapped workflow can accept them as input
+        return any(is_instance_of(message, input_type) for input_type in self.workflow.input_types)
+
+    @handler  # No output_types - can send any completion data type
+    async def process_workflow(self, input_data: object, ctx: WorkflowContext[Any]) -> None:
+        """Execute the sub-workflow with raw input data.
+
+        This handler starts a new sub-workflow execution. When the sub-workflow
+        needs external information, it pauses and sends a request to the parent.
+
+        Args:
+            input_data: The input data to send to the sub-workflow.
+            ctx: The workflow context from the parent.
+        """
+        # Skip RequestResponse - it has a specific handler
+        if isinstance(input_data, RequestResponse):
+            logger.debug(f"WorkflowExecutor {self.id} ignoring input of type {type(input_data)}")
+            return
+
+        await self._ensure_state_loaded(ctx)
+
+        # Create execution context for this sub-workflow run
+        execution_id = str(uuid.uuid4())
+        execution_context = ExecutionContext(
+            execution_id=execution_id,
+            collected_responses={},
+            expected_response_count=0,
+            pending_requests={},
+        )
+        self._execution_contexts[execution_id] = execution_context
+
+        # Track this execution
+        self._active_executions += 1
+
+        logger.debug(f"WorkflowExecutor {self.id} starting sub-workflow {self.workflow.id} execution {execution_id}")
+
+        try:
+            # Run the sub-workflow and collect all events
+            result = await self.workflow.run(input_data)
+
+            logger.debug(
+                f"WorkflowExecutor {self.id} sub-workflow {self.workflow.id} "
+                f"execution {execution_id} completed with {len(result)} events"
+            )
+
+            # Process the workflow result using shared logic
+            await self._process_workflow_result(result, execution_context, ctx)
+        finally:
+            # Clean up execution context if it's completed (no pending requests)
+            if execution_id in self._execution_contexts:
+                exec_ctx = self._execution_contexts[execution_id]
+                if not exec_ctx.pending_requests:
+                    del self._execution_contexts[execution_id]
+                    self._active_executions -= 1
+
+    async def _process_workflow_result(
+        self, result: Any, execution_context: ExecutionContext, ctx: WorkflowContext[Any]
+    ) -> None:
+        """Process the result from a workflow execution.
+
+        This method handles the common logic for processing outputs, request info events,
+        and final states that is shared between process_workflow and handle_response.
+
+        Args:
+            result: The workflow execution result.
+            execution_context: The execution context for this sub-workflow run.
+            ctx: The workflow context.
+        """
+        # Collect all events from the workflow
+        request_info_events = result.get_request_info_events()
+        outputs = result.get_outputs()
+        final_state = result.get_final_state()
+        logger.debug(
+            f"WorkflowExecutor {self.id} processing workflow result with "
+            f"{len(outputs)} outputs and {len(request_info_events)} request info events, "
+            f"final state: {final_state}"
+        )
+
+        # Process outputs
+        for output in outputs:
+            await ctx.send_message(output)
+
+        # Process request info events
+        for event in request_info_events:
+            # Track the pending request in execution context
+            execution_context.pending_requests[event.request_id] = event.data
+            # Map request to execution for response routing
+            self._request_to_execution[event.request_id] = execution_context.execution_id
+            # Set source_executor_id for response routing and send to parent
+            if not isinstance(event.data, RequestInfoMessage):
+                raise TypeError(f"Expected RequestInfoMessage, got {type(event.data)}")
+            # Set the source_executor_id to this WorkflowExecutor's ID for response routing
+            event.data.source_executor_id = self.id
+            await ctx.send_message(event.data)
+
+        # Update expected response count for this execution
+        execution_context.expected_response_count = len(request_info_events)
+
+        # Handle final state
+        if final_state == WorkflowRunState.FAILED:
+            # Find the WorkflowFailedEvent.
+            failed_events = [e for e in result if isinstance(e, WorkflowFailedEvent)]
+            if failed_events:
+                failed_event = failed_events[0]
+                error_type = failed_event.details.error_type
+                error_message = failed_event.details.message
+                exception = Exception(
+                    f"Sub-workflow {self.workflow.id} failed with error: {error_type} - {error_message}"
+                )
+                error_event = WorkflowErrorEvent(
+                    data=exception,
+                )
+                await ctx.add_event(error_event)
+                self._active_executions -= 1
+        elif final_state == WorkflowRunState.IDLE:
+            # Sub-workflow is idle - nothing more to do now
+            logger.debug(f"Sub-workflow {self.workflow.id} is idle with {self._active_executions} active executions")
+            self._active_executions -= 1  # Treat idle as completion for now
+        elif final_state == WorkflowRunState.CANCELLED:
+            # Sub-workflow was cancelled - treat as completion
+            logger.debug(
+                f"Sub-workflow {self.workflow.id} was cancelled with {self._active_executions} active executions"
+            )
+            self._active_executions -= 1
+        elif final_state == WorkflowRunState.IN_PROGRESS_PENDING_REQUESTS:
+            # Sub-workflow is still running with pending requests
+            logger.debug(
+                f"Sub-workflow {self.workflow.id} is still in progress with {len(request_info_events)} "
+                f"pending requests with {self._active_executions} active executions"
+            )
+        elif final_state == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS:
+            # Sub-workflow is idle but has pending requests
+            logger.debug(
+                f"Sub-workflow {self.workflow.id} is idle with pending requests: "
+                f"{len(request_info_events)} with {self._active_executions} active executions"
+            )
+        else:
+            raise RuntimeError(f"Unexpected final state: {final_state}")
+
+        await self._persist_execution_state(ctx)
+
+    @handler
+    async def handle_response(
+        self,
+        response: RequestResponse[RequestInfoMessage, Any],
+        ctx: WorkflowContext[Any],
+    ) -> None:
+        """Handle response from parent for a forwarded request.
+
+        This handler accumulates responses and only resumes the sub-workflow
+        when all expected responses have been received for that execution.
+
+        Args:
+            response: The response to a previous request.
+            ctx: The workflow context.
+        """
+        await self._ensure_state_loaded(ctx)
+
+        # Find the execution context for this request
+        execution_id = self._request_to_execution.get(response.request_id)
+        if not execution_id or execution_id not in self._execution_contexts:
+            logger.warning(
+                f"WorkflowExecutor {self.id} received response for unknown request_id: {response.request_id}, ignoring"
+            )
+            return
+
+        execution_context = self._execution_contexts[execution_id]
+
+        # Check if we have this pending request in the execution context
+        if response.request_id not in execution_context.pending_requests:
+            logger.warning(
+                f"WorkflowExecutor {self.id} received response for unknown request_id: "
+                f"{response.request_id} in execution {execution_id}, ignoring"
+            )
+            return
+
+        # Remove the request from pending list and request mapping
+        execution_context.pending_requests.pop(response.request_id, None)
+        self._request_to_execution.pop(response.request_id, None)
+
+        # Accumulate the response in this execution's context
+        execution_context.collected_responses[response.request_id] = response.data
+
+        await self._persist_execution_state(ctx)
+
+        # Check if we have all expected responses for this execution
+        if len(execution_context.collected_responses) < execution_context.expected_response_count:
+            logger.debug(
+                f"WorkflowExecutor {self.id} execution {execution_id} waiting for more responses: "
+                f"{len(execution_context.collected_responses)}/{execution_context.expected_response_count} received"
+            )
+            return  # Wait for more responses
+
+        # Send all collected responses to the sub-workflow
+        responses_to_send = dict(execution_context.collected_responses)
+        execution_context.collected_responses.clear()  # Clear for next batch
+
+        try:
+            # Resume the sub-workflow with all collected responses
+            result = await self.workflow.send_responses(responses_to_send)
+
+            # Process the workflow result using shared logic
+            await self._process_workflow_result(result, execution_context, ctx)
+        finally:
+            # Clean up execution context if it's completed (no pending requests)
+            if not execution_context.pending_requests:
+                del self._execution_contexts[execution_id]
+                self._active_executions -= 1
+
+    async def _ensure_state_loaded(self, ctx: WorkflowContext[Any]) -> None:
+        if self._state_loaded:
+            return
+
+        state: dict[str, Any] | None = None
+        try:
+            state = await ctx.get_state()
+        except Exception:
+            state = None
+
+        if isinstance(state, dict) and state:
+            with contextlib.suppress(Exception):
+                self.restore_state(state)
+                self._state_loaded = True
+        else:
+            self._state_loaded = True
+
+    def restore_state(self, state: dict[str, Any]) -> None:
+        """Restore pending request bookkeeping from a checkpoint snapshot."""
+        self._execution_contexts = {}
+        self._request_to_execution = {}
+
+        executions_payload = state.get("executions")
+        if isinstance(executions_payload, Mapping) and executions_payload:
+            for execution_id, payload in executions_payload.items():
+                if not isinstance(execution_id, str) or not isinstance(payload, Mapping):
+                    continue
+
+                pending_ids_raw = payload.get("pending_request_ids", [])
+                if not isinstance(pending_ids_raw, list):
+                    continue
+                pending_ids = [rid for rid in pending_ids_raw if isinstance(rid, str)]
+
+                expected = payload.get("expected_response_count", len(pending_ids))
+                try:
+                    expected_count = int(expected)
+                except (TypeError, ValueError):
+                    expected_count = len(pending_ids)
+
+                collected_ids_raw = payload.get("collected_response_ids", [])
+                collected: dict[str, Any] = {}
+                if isinstance(collected_ids_raw, list):
+                    for rid in collected_ids_raw:
+                        if isinstance(rid, str):
+                            collected[rid] = None
+
+                exec_ctx = ExecutionContext(
+                    execution_id=execution_id,
+                    collected_responses=collected,
+                    expected_response_count=expected_count,
+                    pending_requests={rid: None for rid in pending_ids},
+                )
+
+                if exec_ctx.pending_requests or exec_ctx.collected_responses:
+                    self._execution_contexts[execution_id] = exec_ctx
+                    for rid in exec_ctx.pending_requests:
+                        self._request_to_execution[rid] = execution_id
+        else:
+            pending_ids = state.get("pending_request_ids", [])
+            if isinstance(pending_ids, list):
+                pending = [rid for rid in pending_ids if isinstance(rid, str)]
+                if pending:
+                    try:
+                        expected = int(state.get("expected_response_count", len(pending)))
+                    except (TypeError, ValueError):
+                        expected = len(pending)
+
+                    execution_id = str(uuid.uuid4())
+                    exec_ctx = ExecutionContext(
+                        execution_id=execution_id,
+                        collected_responses={},
+                        expected_response_count=expected,
+                        pending_requests={rid: None for rid in pending},
+                    )
+                    self._execution_contexts[execution_id] = exec_ctx
+                    for rid in pending:
+                        self._request_to_execution[rid] = execution_id
+
+        try:
+            self._active_executions = int(state.get("active_executions", len(self._execution_contexts)))
+        except (TypeError, ValueError):
+            self._active_executions = len(self._execution_contexts)
+
+        helper_states = state.get("request_info_executor_states", {})
+        restored_request_data: dict[str, RequestInfoMessage] = {}
+        if isinstance(helper_states, Mapping):
+            for exec_id, helper_state in helper_states.items():
+                helper_executor = self.workflow.executors.get(exec_id)
+                if not isinstance(helper_executor, RequestInfoExecutor) or not isinstance(helper_state, Mapping):
+                    continue
+                with contextlib.suppress(Exception):
+                    helper_executor.restore_state(dict(helper_state))
+                    for req_id, event in getattr(helper_executor, "_request_events", {}).items():  # type: ignore[attr-defined]
+                        if (
+                            isinstance(req_id, str)
+                            and isinstance(event, RequestInfoEvent)
+                            and isinstance(event.data, RequestInfoMessage)
+                        ):
+                            restored_request_data[req_id] = event.data
+
+        if restored_request_data:
+            for req_id, data in restored_request_data.items():
+                execution_id = self._request_to_execution.get(req_id)
+                if execution_id and execution_id in self._execution_contexts:
+                    self._execution_contexts[execution_id].pending_requests[req_id] = data
+
+        for execution_id, exec_ctx in self._execution_contexts.items():
+            for req_id in exec_ctx.pending_requests:
+                self._request_to_execution.setdefault(req_id, execution_id)
+
+        request_map = state.get("request_to_execution")
+        if isinstance(request_map, Mapping):
+            for req_id, execution_id in request_map.items():
+                if (
+                    isinstance(req_id, str)
+                    and isinstance(execution_id, str)
+                    and execution_id in self._execution_contexts
+                ):
+                    self._request_to_execution.setdefault(req_id, execution_id)
+
+        self._state_loaded = True
+
+    def _build_state_snapshot(self) -> dict[str, Any]:
+        executions: dict[str, Any] = {}
+        pending_request_ids: list[str] = []
+
+        for execution_id, exec_ctx in self._execution_contexts.items():
+            if not exec_ctx.pending_requests and not exec_ctx.collected_responses:
+                continue
+
+            request_ids = list(exec_ctx.pending_requests.keys())
+            pending_request_ids.extend(request_ids)
+
+            summary: dict[str, Any] = {
+                "pending_request_ids": request_ids,
+                "expected_response_count": exec_ctx.expected_response_count,
+            }
+
+            if exec_ctx.collected_responses:
+                summary["collected_response_ids"] = list(exec_ctx.collected_responses.keys())
+
+            executions[execution_id] = summary
+
+        helper_states: dict[str, Any] = {}
+        for exec_id, executor in self.workflow.executors.items():
+            if isinstance(executor, RequestInfoExecutor):
+                with contextlib.suppress(Exception):
+                    snapshot = executor.snapshot_state()
+                    if snapshot:
+                        helper_states[exec_id] = snapshot
+
+        has_state = bool(executions or helper_states or self._request_to_execution)
+        if not has_state:
+            return {}
+
+        state: dict[str, Any] = {
+            "executions": executions,
+            "request_to_execution": dict(self._request_to_execution),
+            "pending_request_ids": pending_request_ids,
+            "active_executions": self._active_executions,
+        }
+
+        if helper_states:
+            state["request_info_executor_states"] = helper_states
+
+        return state
+
+    async def _persist_execution_state(self, ctx: WorkflowContext[Any]) -> None:
+        snapshot = self._build_state_snapshot()
+        try:
+            await ctx.set_state(snapshot)
+        except Exception as exc:  # pragma: no cover - transport specific
+            logger.warning(f"WorkflowExecutor {self.id} failed to persist state: {exc}")

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

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

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

@@ -0,0 +1,5 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from agent_framework_a2a import A2AAgent, __version__
+
+__all__ = ["A2AAgent", "__version__"]

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

@@ -0,0 +1,33 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+
+import importlib
+from typing import Any
+
+_IMPORTS: dict[str, tuple[str, list[str]]] = {
+    "AzureAIAgentClient": ("agent_framework_azure_ai", ["azure_ai", "azure"]),
+    "AzureOpenAIAssistantsClient": ("agent_framework.azure._assistants_client", []),
+    "AzureOpenAIChatClient": ("agent_framework.azure._chat_client", []),
+    "AzureAISettings": ("agent_framework_azure_ai", ["azure_ai", "azure"]),
+    "AzureOpenAISettings": ("agent_framework.azure._shared", []),
+    "AzureOpenAIResponsesClient": ("agent_framework.azure._responses_client", []),
+    "get_entra_auth_token": ("agent_framework.azure._entra_id_authentication", []),
+}
+
+
+def __getattr__(name: str) -> Any:
+    if name in _IMPORTS:
+        package_name, package_extra = _IMPORTS[name]
+        try:
+            return getattr(importlib.import_module(package_name), name)
+        except ModuleNotFoundError as exc:
+            raise ModuleNotFoundError(
+                f"The {' or '.join(package_extra)} extra is not installed, "
+                f"please use `pip install agent-framework[{package_extra[0]}]`, "
+                "or update your requirements.txt or pyproject.toml file."
+            ) from exc
+    raise AttributeError(f"Module `azure` has no attribute {name}.")
+
+
+def __dir__() -> list[str]:
+    return list(_IMPORTS.keys())

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

@@ -0,0 +1,19 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from agent_framework_azure_ai import AzureAIAgentClient, AzureAISettings
+
+from agent_framework.azure._assistants_client import AzureOpenAIAssistantsClient
+from agent_framework.azure._chat_client import AzureOpenAIChatClient
+from agent_framework.azure._entra_id_authentication import get_entra_auth_token
+from agent_framework.azure._responses_client import AzureOpenAIResponsesClient
+from agent_framework.azure._shared import AzureOpenAISettings
+
+__all__ = [
+    "AzureAIAgentClient",
+    "AzureAISettings",
+    "AzureOpenAIAssistantsClient",
+    "AzureOpenAIChatClient",
+    "AzureOpenAIResponsesClient",
+    "AzureOpenAISettings",
+    "get_entra_auth_token",
+]

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

@@ -0,0 +1,135 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Any, ClassVar
+
+from openai.lib.azure import AsyncAzureADTokenProvider, AsyncAzureOpenAI
+from pydantic import SecretStr, ValidationError
+from pydantic.networks import AnyUrl
+
+from ..exceptions import ServiceInitializationError
+from ..openai import OpenAIAssistantsClient
+from ._shared import AzureOpenAISettings
+
+if TYPE_CHECKING:
+    from azure.core.credentials import TokenCredential
+
+__all__ = ["AzureOpenAIAssistantsClient"]
+
+
+class AzureOpenAIAssistantsClient(OpenAIAssistantsClient):
+    """Azure OpenAI Assistants client."""
+
+    DEFAULT_AZURE_API_VERSION: ClassVar[str] = "2024-05-01-preview"
+
+    def __init__(
+        self,
+        deployment_name: str | None = None,
+        assistant_id: str | None = None,
+        assistant_name: str | None = None,
+        thread_id: str | None = None,
+        api_key: str | None = None,
+        endpoint: str | None = None,
+        base_url: str | None = None,
+        api_version: str | None = None,
+        ad_token: str | None = None,
+        ad_token_provider: AsyncAzureADTokenProvider | None = None,
+        token_endpoint: str | None = None,
+        credential: "TokenCredential | None" = None,
+        default_headers: Mapping[str, str] | None = None,
+        async_client: AsyncAzureOpenAI | None = None,
+        env_file_path: str | None = None,
+        env_file_encoding: str | None = None,
+    ) -> None:
+        """Initialize an Azure OpenAI Assistants client.
+
+        Args:
+            deployment_name: The Azure OpenAI deployment name for the model to use.
+            assistant_id: The ID of an Azure OpenAI assistant to use.
+                If not provided, a new assistant will be created (and deleted after the request).
+            assistant_name: The name to use when creating new assistants.
+            thread_id: Default thread ID to use for conversations. Can be overridden by
+                conversation_id property, when making a request.
+                If not provided, a new thread will be created (and deleted after the request).
+            api_key: The optional API key to use. If provided will override,
+                the env vars or .env file value.
+            endpoint: The optional deployment endpoint. If provided will override the value
+                in the env vars or .env file.
+            base_url: The optional deployment base_url. If provided will override the value
+                in the env vars or .env file.
+            api_version: The optional deployment api version. If provided will override the value
+                in the env vars or .env file.
+            ad_token: The Azure Active Directory token. (Optional)
+            ad_token_provider: The Azure Active Directory token provider. (Optional)
+            token_endpoint: The token endpoint to request an Azure token. (Optional)
+            credential: The Azure credential to use for authentication. (Optional)
+            default_headers: The default headers mapping of string keys to
+                string values for HTTP requests. (Optional)
+            async_client: An existing client to use. (Optional)
+            env_file_path: Use the environment settings file as a fallback
+                to environment variables. (Optional)
+            env_file_encoding: The encoding of the environment settings file. (Optional)
+        """
+        try:
+            azure_openai_settings = AzureOpenAISettings(
+                api_key=SecretStr(api_key) if api_key else None,
+                base_url=AnyUrl(base_url) if base_url else None,
+                endpoint=AnyUrl(endpoint) if endpoint else None,
+                chat_deployment_name=deployment_name,
+                api_version=api_version,
+                env_file_path=env_file_path,
+                env_file_encoding=env_file_encoding,
+                token_endpoint=token_endpoint,
+                default_api_version=self.DEFAULT_AZURE_API_VERSION,
+            )
+        except ValidationError as ex:
+            raise ServiceInitializationError("Failed to create Azure OpenAI settings.", ex) from ex
+
+        if not azure_openai_settings.chat_deployment_name:
+            raise ServiceInitializationError(
+                "Azure OpenAI deployment name is required. Set via 'deployment_name' parameter "
+                "or 'AZURE_OPENAI_CHAT_DEPLOYMENT_NAME' environment variable."
+            )
+
+        # Handle authentication: try API key first, then AD token, then Entra ID
+        if (
+            not async_client
+            and not azure_openai_settings.api_key
+            and not ad_token
+            and not ad_token_provider
+            and azure_openai_settings.token_endpoint
+            and credential
+        ):
+            ad_token = azure_openai_settings.get_azure_auth_token(credential)
+
+        if not async_client and not azure_openai_settings.api_key and not ad_token and not ad_token_provider:
+            raise ServiceInitializationError("The Azure OpenAI API key, ad_token, or ad_token_provider is required.")
+
+        # Create Azure client if not provided
+        if not async_client:
+            client_params: dict[str, Any] = {
+                "api_version": azure_openai_settings.api_version,
+                "default_headers": default_headers,
+            }
+
+            if azure_openai_settings.api_key:
+                client_params["api_key"] = azure_openai_settings.api_key.get_secret_value()
+            elif ad_token:
+                client_params["azure_ad_token"] = ad_token
+            elif ad_token_provider:
+                client_params["azure_ad_token_provider"] = ad_token_provider
+
+            if azure_openai_settings.base_url:
+                client_params["base_url"] = str(azure_openai_settings.base_url)
+            elif azure_openai_settings.endpoint:
+                client_params["azure_endpoint"] = str(azure_openai_settings.endpoint)
+
+            async_client = AsyncAzureOpenAI(**client_params)
+
+        super().__init__(
+            ai_model_id=azure_openai_settings.chat_deployment_name,
+            assistant_id=assistant_id,
+            assistant_name=assistant_name,
+            thread_id=thread_id,
+            async_client=async_client,  # type: ignore[reportArgumentType]
+        )

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

@@ -0,0 +1,194 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import json
+import logging
+import sys
+from collections.abc import Mapping
+from typing import Any, TypeVar
+
+from azure.core.credentials import TokenCredential
+from openai.lib.azure import AsyncAzureADTokenProvider, AsyncAzureOpenAI
+from openai.types.chat.chat_completion import Choice
+from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
+from pydantic import SecretStr, ValidationError
+from pydantic.networks import AnyUrl
+
+from agent_framework import (
+    ChatResponse,
+    ChatResponseUpdate,
+    CitationAnnotation,
+    TextContent,
+    use_chat_middleware,
+    use_function_invocation,
+)
+from agent_framework.exceptions import ServiceInitializationError
+from agent_framework.observability import use_observability
+from agent_framework.openai._chat_client import OpenAIBaseChatClient
+
+from ._shared import (
+    AzureOpenAIConfigMixin,
+    AzureOpenAISettings,
+)
+
+if sys.version_info >= (3, 12):
+    from typing import override  # type: ignore # pragma: no cover
+else:
+    from typing_extensions import override  # type: ignore[import] # pragma: no cover
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+TChatResponse = TypeVar("TChatResponse", ChatResponse, ChatResponseUpdate)
+TAzureOpenAIChatClient = TypeVar("TAzureOpenAIChatClient", bound="AzureOpenAIChatClient")
+
+
+@use_function_invocation
+@use_observability
+@use_chat_middleware
+class AzureOpenAIChatClient(AzureOpenAIConfigMixin, OpenAIBaseChatClient):
+    """Azure OpenAI Chat completion class."""
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        deployment_name: str | None = None,
+        endpoint: str | None = None,
+        base_url: str | None = None,
+        api_version: str | None = None,
+        ad_token: str | None = None,
+        ad_token_provider: AsyncAzureADTokenProvider | None = None,
+        token_endpoint: str | None = None,
+        credential: TokenCredential | None = None,
+        default_headers: Mapping[str, str] | None = None,
+        async_client: AsyncAzureOpenAI | None = None,
+        env_file_path: str | None = None,
+        env_file_encoding: str | None = None,
+        instruction_role: str | None = None,
+    ) -> None:
+        """Initialize an AzureChatCompletion service.
+
+        Args:
+            api_key: The optional api key. If provided, will override the value in the
+                env vars or .env file.
+            deployment_name: The optional deployment. If provided, will override the value
+                (chat_deployment_name) in the env vars or .env file.
+            endpoint: The optional deployment endpoint. If provided will override the value
+                in the env vars or .env file.
+            base_url: The optional deployment base_url. If provided will override the value
+                in the env vars or .env file.
+            api_version: The optional deployment api version. If provided will override the value
+                in the env vars or .env file.
+            ad_token: The Azure Active Directory token. (Optional)
+            ad_token_provider: The Azure Active Directory token provider. (Optional)
+            token_endpoint: The token endpoint to request an Azure token. (Optional)
+            credential: The Azure credential for authentication. (Optional)
+            default_headers: The default headers mapping of string keys to
+                string values for HTTP requests. (Optional)
+            async_client: An existing client to use. (Optional)
+            env_file_path: Use the environment settings file as a fallback to using env vars.
+            env_file_encoding: The encoding of the environment settings file, defaults to 'utf-8'.
+            instruction_role: The role to use for 'instruction' messages, for example, summarization
+                prompts could use `developer` or `system`. (Optional)
+        """
+        try:
+            # Filter out any None values from the arguments
+            azure_openai_settings = AzureOpenAISettings(
+                api_key=SecretStr(api_key) if api_key else None,
+                base_url=AnyUrl(base_url) if base_url else None,
+                endpoint=AnyUrl(endpoint) if endpoint else None,
+                chat_deployment_name=deployment_name,
+                api_version=api_version,
+                env_file_path=env_file_path,
+                env_file_encoding=env_file_encoding,
+                token_endpoint=token_endpoint,
+            )
+        except ValidationError as exc:
+            raise ServiceInitializationError(f"Failed to validate settings: {exc}") from exc
+
+        if not azure_openai_settings.chat_deployment_name:
+            raise ServiceInitializationError(
+                "Azure OpenAI deployment name is required. Set via 'deployment_name' parameter "
+                "or 'AZURE_OPENAI_CHAT_DEPLOYMENT_NAME' environment variable."
+            )
+
+        super().__init__(
+            deployment_name=azure_openai_settings.chat_deployment_name,
+            endpoint=azure_openai_settings.endpoint,
+            base_url=azure_openai_settings.base_url,
+            api_version=azure_openai_settings.api_version,  # type: ignore
+            api_key=azure_openai_settings.api_key.get_secret_value() if azure_openai_settings.api_key else None,
+            ad_token=ad_token,
+            ad_token_provider=ad_token_provider,
+            token_endpoint=azure_openai_settings.token_endpoint,
+            credential=credential,
+            default_headers=default_headers,
+            client=async_client,
+            instruction_role=instruction_role,
+        )
+
+    @classmethod
+    def from_dict(cls: type[TAzureOpenAIChatClient], settings: dict[str, Any]) -> TAzureOpenAIChatClient:
+        """Initialize an Azure OpenAI service from a dictionary of settings.
+
+        Args:
+            settings: A dictionary of settings for the service.
+                should contain keys: service_id, and optionally:
+                ad_auth, ad_token_provider, default_headers
+        """
+        return cls(
+            api_key=settings.get("api_key"),
+            deployment_name=settings.get("deployment_name"),
+            endpoint=settings.get("endpoint"),
+            base_url=settings.get("base_url"),
+            api_version=settings.get("api_version"),
+            ad_token=settings.get("ad_token"),
+            ad_token_provider=settings.get("ad_token_provider"),
+            default_headers=settings.get("default_headers"),
+            env_file_path=settings.get("env_file_path"),
+        )
+
+    @override
+    def _parse_text_from_choice(self, choice: Choice | ChunkChoice) -> TextContent | None:
+        """Parse the choice into a TextContent object.
+
+        Overwritten from OpenAIBaseChatClient to deal with Azure On Your Data function.
+        For docs see:
+        https://learn.microsoft.com/en-us/azure/ai-foundry/openai/references/on-your-data?tabs=python#context
+        """
+        message = choice.message if isinstance(choice, Choice) else choice.delta
+        if hasattr(message, "refusal") and message.refusal:
+            return TextContent(text=message.refusal, raw_representation=choice)
+        if not message.content:
+            return None
+        text_content = TextContent(text=message.content, raw_representation=choice)
+        if not message.model_extra or "context" not in message.model_extra:
+            return text_content
+
+        context: dict[str, Any] | str = message.context  # type: ignore[assignment, union-attr]
+        if isinstance(context, str):
+            try:
+                context = json.loads(context)
+            except json.JSONDecodeError:
+                logger.warning("Context is not a valid JSON string, ignoring context.")
+                return text_content
+        if not isinstance(context, dict):
+            logger.warning("Context is not a valid dictionary, ignoring context.")
+            return text_content
+        # `all_retrieved_documents` is currently not used, but can be retrieved
+        # through the raw_representation in the text content.
+        if intent := context.get("intent"):
+            text_content.additional_properties = {"intent": intent}
+        if citations := context.get("citations"):
+            text_content.annotations = []
+            for citation in citations:
+                text_content.annotations.append(
+                    CitationAnnotation(
+                        title=citation.get("title", ""),
+                        url=citation.get("url", ""),
+                        snippet=citation.get("content", ""),
+                        file_id=citation.get("filepath", ""),
+                        tool_name="Azure-on-your-Data",
+                        additional_properties={"chunk_id": citation.get("chunk_id", "")},
+                        raw_representation=citation,
+                    )
+                )
+        return text_content

python/packages/core/agent_framework/azure/_entra_id_authentication.py

@@ -0,0 +1,76 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+from azure.core.exceptions import ClientAuthenticationError
+
+from ..exceptions import ServiceInvalidAuthError
+
+if TYPE_CHECKING:
+    from azure.core.credentials import TokenCredential
+    from azure.core.credentials_async import AsyncTokenCredential
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+def get_entra_auth_token(
+    credential: "TokenCredential",
+    token_endpoint: str,
+    **kwargs: Any,
+) -> str | None:
+    """Retrieve a Microsoft Entra Auth Token for a given token endpoint.
+
+    The token endpoint may be specified as an environment variable, via the .env
+    file or as an argument. If the token endpoint is not provided, the default is None.
+
+    Args:
+        credential: The Azure credential to use for authentication.
+        token_endpoint: The token endpoint to use to retrieve the authentication token.
+        **kwargs: Additional keyword arguments to pass to the token retrieval method.
+
+    Returns:
+        The Azure token or None if the token could not be retrieved.
+    """
+    if not token_endpoint:
+        raise ServiceInvalidAuthError(
+            "A token endpoint must be provided either in settings, as an environment variable, or as an argument."
+        )
+
+    try:
+        auth_token = credential.get_token(token_endpoint, **kwargs)
+    except ClientAuthenticationError as ex:
+        logger.error(f"Failed to retrieve Azure token for the specified endpoint: `{token_endpoint}`, with error: {ex}")
+        return None
+
+    return auth_token.token if auth_token else None
+
+
+async def get_entra_auth_token_async(
+    credential: "AsyncTokenCredential", token_endpoint: str, **kwargs: Any
+) -> str | None:
+    """Retrieve a async Microsoft Entra Auth Token for a given token endpoint.
+
+    The token endpoint may be specified as an environment variable, via the .env
+    file or as an argument. If the token endpoint is not provided, the default is None.
+
+    Args:
+        credential: The async Azure credential to use for authentication.
+        token_endpoint: The token endpoint to use to retrieve the authentication token.
+        **kwargs: Additional keyword arguments to pass to the token retrieval method.
+
+    Returns:
+        The Azure token or None if the token could not be retrieved.
+    """
+    if not token_endpoint:
+        raise ServiceInvalidAuthError(
+            "A token endpoint must be provided either in settings, as an environment variable, or as an argument."
+        )
+
+    try:
+        auth_token = await credential.get_token(token_endpoint, **kwargs)
+    except ClientAuthenticationError as ex:
+        logger.error(f"Failed to retrieve Azure token for the specified endpoint: `{token_endpoint}`, with error: {ex}")
+        return None
+
+    return auth_token.token if auth_token else None

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

@@ -0,0 +1,136 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from collections.abc import Mapping
+from typing import Any, TypeVar
+from urllib.parse import urljoin
+
+from azure.core.credentials import TokenCredential
+from openai.lib.azure import AsyncAzureADTokenProvider, AsyncAzureOpenAI
+from pydantic import SecretStr, ValidationError
+from pydantic.networks import AnyUrl
+
+from agent_framework import use_chat_middleware, use_function_invocation
+from agent_framework.exceptions import ServiceInitializationError
+from agent_framework.observability import use_observability
+from agent_framework.openai._responses_client import OpenAIBaseResponsesClient
+
+from ._shared import (
+    AzureOpenAIConfigMixin,
+    AzureOpenAISettings,
+)
+
+TAzureOpenAIResponsesClient = TypeVar("TAzureOpenAIResponsesClient", bound="AzureOpenAIResponsesClient")
+
+
+@use_observability
+@use_function_invocation
+@use_chat_middleware
+class AzureOpenAIResponsesClient(AzureOpenAIConfigMixin, OpenAIBaseResponsesClient):
+    """Azure Responses completion class."""
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        deployment_name: str | None = None,
+        endpoint: str | None = None,
+        base_url: str | None = None,
+        api_version: str | None = None,
+        ad_token: str | None = None,
+        ad_token_provider: AsyncAzureADTokenProvider | None = None,
+        token_endpoint: str | None = None,
+        credential: TokenCredential | None = None,
+        default_headers: Mapping[str, str] | None = None,
+        async_client: AsyncAzureOpenAI | None = None,
+        env_file_path: str | None = None,
+        env_file_encoding: str | None = None,
+        instruction_role: str | None = None,
+    ) -> None:
+        """Initialize an AzureResponses service.
+
+        Args:
+            api_key: The optional api key. If provided, will override the value in the
+                env vars or .env file.
+            deployment_name: The optional deployment. If provided, will override the value
+                (responses_deployment_name) in the env vars or .env file.
+            endpoint: The optional deployment endpoint. If provided will override the value
+                in the env vars or .env file.
+            base_url: The optional deployment base_url. If provided will override the value
+                in the env vars or .env file. Currently, the base_url must end with "/openai/v1/"
+            api_version: The optional deployment api version. If provided will override the value
+                in the env vars or .env file. Currently, the api_version must be "preview".
+            ad_token: The Azure Active Directory token. (Optional)
+            ad_token_provider: The Azure Active Directory token provider. (Optional)
+            token_endpoint: The token endpoint to request an Azure token. (Optional)
+            credential: The Azure credential for authentication. (Optional)
+            default_headers: The default headers mapping of string keys to
+                string values for HTTP requests. (Optional)
+            async_client: An existing client to use. (Optional)
+            env_file_path: Use the environment settings file as a fallback to using env vars.
+            env_file_encoding: The encoding of the environment settings file, defaults to 'utf-8'.
+            instruction_role: The role to use for 'instruction' messages, for example, summarization
+                prompts could use `developer` or `system`. (Optional)
+        """
+        try:
+            # Filter out any None values from the arguments
+            azure_openai_settings = AzureOpenAISettings(
+                api_key=SecretStr(api_key) if api_key else None,
+                base_url=AnyUrl(base_url) if base_url else None,
+                endpoint=AnyUrl(endpoint) if endpoint else None,
+                responses_deployment_name=deployment_name,
+                api_version=api_version,
+                env_file_path=env_file_path,
+                env_file_encoding=env_file_encoding,
+                token_endpoint=token_endpoint,
+                default_api_version="preview",
+            )
+            # TODO(peterychang): This is a temporary hack to ensure that the base_url is set correctly
+            # while this feature is in preview.
+            # But we should only do this if we're on azure. Private deployments may not need this.
+            if (
+                not azure_openai_settings.base_url
+                and azure_openai_settings.endpoint
+                and str(azure_openai_settings.endpoint).rstrip("/").endswith("openai.azure.com")
+            ):
+                azure_openai_settings.base_url = AnyUrl(urljoin(str(azure_openai_settings.endpoint), "/openai/v1/"))
+        except ValidationError as exc:
+            raise ServiceInitializationError(f"Failed to validate settings: {exc}") from exc
+
+        if not azure_openai_settings.responses_deployment_name:
+            raise ServiceInitializationError(
+                "Azure OpenAI deployment name is required. Set via 'deployment_name' parameter "
+                "or 'AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME' environment variable."
+            )
+
+        super().__init__(
+            deployment_name=azure_openai_settings.responses_deployment_name,
+            endpoint=azure_openai_settings.endpoint,
+            base_url=azure_openai_settings.base_url,
+            api_version=azure_openai_settings.api_version,  # type: ignore
+            api_key=azure_openai_settings.api_key.get_secret_value() if azure_openai_settings.api_key else None,
+            ad_token=ad_token,
+            ad_token_provider=ad_token_provider,
+            token_endpoint=azure_openai_settings.token_endpoint,
+            credential=credential,
+            default_headers=default_headers,
+            client=async_client,
+            instruction_role=instruction_role,
+        )
+
+    @classmethod
+    def from_dict(cls: type[TAzureOpenAIResponsesClient], settings: dict[str, Any]) -> TAzureOpenAIResponsesClient:
+        """Initialize an Open AI service from a dictionary of settings.
+
+        Args:
+            settings: A dictionary of settings for the service.
+        """
+        return cls(
+            api_key=settings.get("api_key"),
+            deployment_name=settings.get("deployment_name"),
+            endpoint=settings.get("endpoint"),
+            base_url=settings.get("base_url"),
+            api_version=settings.get("api_version"),
+            ad_token=settings.get("ad_token"),
+            ad_token_provider=settings.get("ad_token_provider"),
+            default_headers=settings.get("default_headers"),
+            env_file_path=settings.get("env_file_path"),
+        )

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

@@ -0,0 +1,242 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import logging
+import sys
+from collections.abc import Awaitable, Callable, Mapping
+from copy import copy
+from typing import Any, ClassVar, Final
+
+from azure.core.credentials import TokenCredential
+from openai.lib.azure import AsyncAzureOpenAI
+from pydantic import ConfigDict, SecretStr, model_validator, validate_call
+
+from .._pydantic import AFBaseSettings, HTTPsUrl
+from .._telemetry import APP_INFO, USER_AGENT_KEY, prepend_agent_framework_to_user_agent
+from ..exceptions import ServiceInitializationError
+from ..openai._shared import OpenAIBase
+from ._entra_id_authentication import get_entra_auth_token
+
+if sys.version_info >= (3, 11):
+    from typing import Self  # pragma: no cover
+else:
+    from typing_extensions import Self  # pragma: no cover
+
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+DEFAULT_AZURE_API_VERSION: Final[str] = "2024-10-21"
+DEFAULT_AZURE_TOKEN_ENDPOINT: Final[str] = "https://cognitiveservices.azure.com/.default"  # noqa: S105
+
+
+class AzureOpenAISettings(AFBaseSettings):
+    """AzureOpenAI model settings.
+
+    The settings are first loaded from environment variables with the prefix 'AZURE_OPENAI_'.
+    If the environment variables are not found, the settings can be loaded from a .env file
+    with the encoding 'utf-8'. If the settings are not found in the .env file, the settings
+    are ignored; however, validation will fail alerting that the settings are missing.
+
+    Args:
+        endpoint: The endpoint of the Azure deployment. This value
+            can be found in the Keys & Endpoint section when examining
+            your resource from the Azure portal, the endpoint should end in openai.azure.com.
+            If both base_url and endpoint are supplied, base_url will be used.
+            (Env var AZURE_OPENAI_ENDPOINT)
+        chat_deployment_name: The name of the Azure Chat deployment. This value
+            will correspond to the custom name you chose for your deployment
+            when you deployed a model. This value can be found under
+            Resource Management > Deployments in the Azure portal or, alternatively,
+            under Management > Deployments in Azure AI Foundry.
+            (Env var AZURE_OPENAI_CHAT_DEPLOYMENT_NAME)
+        responses_deployment_name: The name of the Azure Responses deployment. This value
+            will correspond to the custom name you chose for your deployment
+            when you deployed a model. This value can be found under
+            Resource Management > Deployments in the Azure portal or, alternatively,
+            under Management > Deployments in Azure AI Foundry.
+            (Env var AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME)
+        api_key: The API key for the Azure deployment. This value can be
+            found in the Keys & Endpoint section when examining your resource in
+            the Azure portal. You can use either KEY1 or KEY2.
+            (Env var AZURE_OPENAI_API_KEY)
+        api_version: The API version to use. The default value is `default_api_version`.
+            (Env var AZURE_OPENAI_API_VERSION)
+        base_url: The url of the Azure deployment. This value
+            can be found in the Keys & Endpoint section when examining
+            your resource from the Azure portal, the base_url consists of the endpoint,
+            followed by /openai/deployments/{deployment_name}/,
+            use endpoint if you only want to supply the endpoint.
+            (Env var AZURE_OPENAI_BASE_URL)
+        token_endpoint: The token endpoint to use to retrieve the authentication token.
+            The default value is `default_token_endpoint`.
+            (Env var AZURE_OPENAI_TOKEN_ENDPOINT)
+        default_api_version: The default API version to use if not specified.
+            The default value is "2024-10-21".
+        default_token_endpoint: The default token endpoint to use if not specified.
+            The default value is "https://cognitiveservices.azure.com/.default".
+        env_file_path: The path to the .env file to load settings from.
+        env_file_encoding: The encoding of the .env file, defaults to 'utf-8'.
+    """
+
+    env_prefix: ClassVar[str] = "AZURE_OPENAI_"
+
+    chat_deployment_name: str | None = None
+    responses_deployment_name: str | None = None
+    endpoint: HTTPsUrl | None = None
+    base_url: HTTPsUrl | None = None
+    api_key: SecretStr | None = None
+    api_version: str | None = None
+    token_endpoint: str | None = None
+    default_api_version: str = DEFAULT_AZURE_API_VERSION
+    default_token_endpoint: str = DEFAULT_AZURE_TOKEN_ENDPOINT
+
+    def get_azure_auth_token(
+        self, credential: "TokenCredential", token_endpoint: str | None = None, **kwargs: Any
+    ) -> str | None:
+        """Retrieve a Microsoft Entra Auth Token for a given token endpoint for the use with Azure OpenAI.
+
+        The required role for the token is `Cognitive Services OpenAI Contributor`.
+        The token endpoint may be specified as an environment variable, via the .env
+        file or as an argument. If the token endpoint is not provided, the default is None.
+        The `token_endpoint` argument takes precedence over the `token_endpoint` attribute.
+
+        Args:
+            credential: The Azure AD credential to use.
+            token_endpoint: The token endpoint to use. Defaults to `https://cognitiveservices.azure.com/.default`.
+            **kwargs: Additional keyword arguments to pass to the token retrieval method.
+
+        Returns:
+            The Azure token or None if the token could not be retrieved.
+
+        Raises:
+            ServiceInitializationError: If the token endpoint is not provided.
+        """
+        endpoint_to_use = token_endpoint or self.token_endpoint or self.default_token_endpoint
+        return get_entra_auth_token(credential, endpoint_to_use, **kwargs)
+
+    @model_validator(mode="after")
+    def _validate_fields(self) -> Self:
+        self.api_version = self.api_version or self.default_api_version
+        self.token_endpoint = self.token_endpoint or self.default_token_endpoint
+        return self
+
+
+class AzureOpenAIConfigMixin(OpenAIBase):
+    """Internal class for configuring a connection to an Azure OpenAI service."""
+
+    OTEL_PROVIDER_NAME: ClassVar[str] = "azure_openai"  # type: ignore[reportIncompatibleVariableOverride, misc]
+
+    @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
+    def __init__(
+        self,
+        deployment_name: str,
+        endpoint: HTTPsUrl | None = None,
+        base_url: HTTPsUrl | None = None,
+        api_version: str = DEFAULT_AZURE_API_VERSION,
+        api_key: str | None = None,
+        ad_token: str | None = None,
+        ad_token_provider: Callable[[], str | Awaitable[str]] | None = None,
+        token_endpoint: str | None = None,
+        credential: TokenCredential | None = None,
+        default_headers: Mapping[str, str] | None = None,
+        client: AsyncAzureOpenAI | None = None,
+        instruction_role: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Internal class for configuring a connection to an Azure OpenAI service.
+
+        The `validate_call` decorator is used with a configuration that allows arbitrary types.
+        This is necessary for types like `HTTPsUrl` and `OpenAIModelTypes`.
+
+        Args:
+            deployment_name: Name of the deployment.
+            ai_model_type: The type of OpenAI model to deploy.
+            endpoint: The specific endpoint URL for the deployment.
+            base_url: The base URL for Azure services.
+            api_version: Azure API version. Defaults to the defined DEFAULT_AZURE_API_VERSION.
+            api_key: API key for Azure services.
+            ad_token: Azure AD token for authentication.
+            ad_token_provider: A callable or coroutine function providing Azure AD tokens.
+            token_endpoint: Azure AD token endpoint use to get the token.
+            credential: Azure credential for authentication.
+            default_headers: Default headers for HTTP requests.
+            client: An existing client to use.
+            instruction_role: The role to use for 'instruction' messages, for example, summarization
+                prompts could use `developer` or `system`.
+            kwargs: Additional keyword arguments.
+
+        """
+        # Merge APP_INFO into the headers if it exists
+        merged_headers = dict(copy(default_headers)) if default_headers else {}
+        if APP_INFO:
+            merged_headers.update(APP_INFO)
+            merged_headers = prepend_agent_framework_to_user_agent(merged_headers)
+        if not client:
+            # If the client is None, the api_key is none, the ad_token is none, and the ad_token_provider is none,
+            # then we will attempt to get the ad_token using the default endpoint specified in the Azure OpenAI
+            # settings.
+            if not api_key and not ad_token_provider and not ad_token and token_endpoint and credential:
+                ad_token = get_entra_auth_token(credential, token_endpoint)
+
+            if not api_key and not ad_token and not ad_token_provider:
+                raise ServiceInitializationError(
+                    "Please provide either api_key, ad_token or ad_token_provider or a client."
+                )
+
+            if not endpoint and not base_url:
+                raise ServiceInitializationError("Please provide an endpoint or a base_url")
+
+            args: dict[str, Any] = {
+                "default_headers": merged_headers,
+            }
+            if api_version:
+                args["api_version"] = api_version
+            if ad_token:
+                args["azure_ad_token"] = ad_token
+            if ad_token_provider:
+                args["azure_ad_token_provider"] = ad_token_provider
+            if api_key:
+                args["api_key"] = api_key
+            if base_url:
+                args["base_url"] = str(base_url)
+            if endpoint and not base_url:
+                args["azure_endpoint"] = str(endpoint)
+            if deployment_name:
+                args["azure_deployment"] = deployment_name
+            if "websocket_base_url" in kwargs:
+                args["websocket_base_url"] = kwargs.pop("websocket_base_url")
+
+            client = AsyncAzureOpenAI(**args)
+        args = {
+            "ai_model_id": deployment_name,
+            "client": client,
+        }
+        if instruction_role:
+            args["instruction_role"] = instruction_role
+        super().__init__(**args, **kwargs)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert the configuration to a dictionary."""
+        client_settings = {
+            "base_url": str(self.client.base_url),
+            "api_version": self.client._custom_query["api-version"],  # type: ignore
+            "api_key": self.client.api_key,
+            "ad_token": getattr(self.client, "_azure_ad_token", None),
+            "ad_token_provider": getattr(self.client, "_azure_ad_token_provider", None),
+            "default_headers": {k: v for k, v in self.client.default_headers.items() if k != USER_AGENT_KEY},
+        }
+        base = self.model_dump(
+            exclude={
+                "prompt_tokens",
+                "completion_tokens",
+                "total_tokens",
+                "api_type",
+                "org_id",
+                "service_id",
+                "client",
+            },
+            by_alias=True,
+            exclude_none=True,
+        )
+        base.update(client_settings)
+        return base

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

@@ -0,0 +1,35 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import importlib
+from typing import Any
+
+PACKAGE_NAME = "agent_framework_devui"
+PACKAGE_EXTRA = "devui"
+_IMPORTS = [
+    "AgentFrameworkRequest",
+    "DevServer",
+    "DiscoveryResponse",
+    "EntityInfo",
+    "OpenAIError",
+    "OpenAIResponse",
+    "ResponseStreamEvent",
+    "main",
+    "serve",
+    "__version__",
+]
+
+
+def __getattr__(name: str) -> Any:
+    if name in _IMPORTS:
+        try:
+            return getattr(importlib.import_module(PACKAGE_NAME), name)
+        except ModuleNotFoundError as exc:
+            raise ModuleNotFoundError(
+                f"The '{PACKAGE_EXTRA}' extra is not installed, "
+                f"please do `pip install agent-framework[{PACKAGE_EXTRA}]`"
+            ) from exc
+    raise AttributeError(f"Module {PACKAGE_NAME} has no attribute {name}.")
+
+
+def __dir__() -> list[str]:
+    return _IMPORTS

python/packages/core/agent_framework/exceptions.py

@@ -0,0 +1,142 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import logging
+from typing import Any, Literal
+
+logger = logging.getLogger("agent_framework")
+
+
+class AgentFrameworkException(Exception):
+    """Base exceptions for the Agent Framework.
+
+    Automatically logs the message as debug.
+    """
+
+    def __init__(
+        self,
+        message: str,
+        inner_exception: Exception | None = None,
+        log_level: Literal[0] | Literal[10] | Literal[20] | Literal[30] | Literal[40] | Literal[50] | None = 10,
+        *args: Any,
+        **kwargs: Any,
+    ):
+        """Create an AgentFrameworkException.
+
+        This emits a debug log (by default), with the inner_exception if provided.
+        """
+        if log_level is not None:
+            logger.log(log_level, message, exc_info=inner_exception)
+        if inner_exception:
+            super().__init__(message, inner_exception, *args)  # type: ignore
+        super().__init__(message, *args)  # type: ignore
+
+
+class AgentException(AgentFrameworkException):
+    """Base class for all agent exceptions."""
+
+    pass
+
+
+class AgentExecutionException(AgentException):
+    """An error occurred while executing the agent."""
+
+    pass
+
+
+class AgentInitializationError(AgentException):
+    """An error occurred while initializing the agent."""
+
+    pass
+
+
+class AgentThreadException(AgentException):
+    """An error occurred while managing the agent thread."""
+
+    pass
+
+
+class ChatClientException(AgentFrameworkException):
+    """An error occurred while dealing with a chat client."""
+
+    pass
+
+
+class ChatClientInitializationError(ChatClientException):
+    """An error occurred while initializing the chat client."""
+
+    pass
+
+
+# region Service Exceptions
+
+
+class ServiceException(AgentFrameworkException):
+    """Base class for all service exceptions."""
+
+    pass
+
+
+class ServiceInitializationError(ServiceException):
+    """An error occurred while initializing the service."""
+
+    pass
+
+
+class ServiceResponseException(ServiceException):
+    """Base class for all service response exceptions."""
+
+    pass
+
+
+class ServiceContentFilterException(ServiceResponseException):
+    """An error was raised by the content filter of the service."""
+
+    pass
+
+
+class ServiceInvalidAuthError(ServiceException):
+    """An error occurred while authenticating the service."""
+
+    pass
+
+
+class ServiceInvalidExecutionSettingsError(ServiceResponseException):
+    """An error occurred while validating the execution settings of the service."""
+
+    pass
+
+
+class ServiceInvalidRequestError(ServiceResponseException):
+    """An error occurred while validating the request to the service."""
+
+    pass
+
+
+class ServiceInvalidResponseError(ServiceResponseException):
+    """An error occurred while validating the response from the service."""
+
+    pass
+
+
+class ToolException(AgentFrameworkException):
+    """An error occurred while executing a tool."""
+
+    pass
+
+
+class ToolExecutionException(ToolException):
+    """An error occurred while executing a tool."""
+
+    pass
+
+
+class AdditionItemMismatch(AgentFrameworkException):
+    """An error occurred while adding two types."""
+
+    pass
+
+
+class MiddlewareException(AgentFrameworkException):
+    """An error occurred during middleware execution."""
+
+    pass

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

@@ -0,0 +1,4 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# This makes agent_framework.lab a namespace package
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)

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

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

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

@@ -0,0 +1,5 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from agent_framework_mem0 import Mem0Provider, __version__
+
+__all__ = ["Mem0Provider", "__version__"]

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

@@ -0,0 +1,30 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import importlib
+from typing import Any
+
+PACKAGE_NAME = "agent_framework_copilotstudio"
+PACKAGE_EXTRA = ["microsoft-copilotstudio", "copilotstudio"]
+_IMPORTS: dict[str, tuple[str, list[str]]] = {
+    "CopilotStudioAgent": ("agent_framework_copilotstudio", ["microsoft-copilotstudio", "copilotstudio"]),
+    "__version__": ("agent_framework_copilotstudio", ["microsoft-copilotstudio", "copilotstudio"]),
+    "acquire_token": ("agent_framework_copilotstudio", ["microsoft-copilotstudio", "copilotstudio"]),
+}
+
+
+def __getattr__(name: str) -> Any:
+    if name in _IMPORTS:
+        package_name, package_extra = _IMPORTS[name]
+        try:
+            return getattr(importlib.import_module(package_name), name)
+        except ModuleNotFoundError as exc:
+            raise ModuleNotFoundError(
+                f"The {' or '.join(package_extra)} extra is not installed, "
+                f"please use `pip install agent-framework[{package_extra[0]}]`, "
+                "or update your requirements.txt or pyproject.toml file."
+            ) from exc
+    raise AttributeError(f"Module `azure` has no attribute {name}.")
+
+
+def __dir__() -> list[str]:
+    return list(_IMPORTS.keys())

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

@@ -0,0 +1,5 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from agent_framework_copilotstudio import CopilotStudioAgent, __version__, acquire_token
+
+__all__ = ["CopilotStudioAgent", "__version__", "acquire_token"]

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

@@ -0,0 +1,8 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+
+from ._assistants_client import *  # noqa: F403
+from ._chat_client import *  # noqa: F403
+from ._exceptions import *  # noqa: F403
+from ._responses_client import *  # noqa: F403
+from ._shared import *  # noqa: F403

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

@@ -0,0 +1,498 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import json
+import sys
+from collections.abc import AsyncIterable, Mapping, MutableMapping, MutableSequence
+from typing import Any
+
+from openai import AsyncOpenAI
+from openai.types.beta.threads import (
+    ImageURLContentBlockParam,
+    ImageURLParam,
+    MessageContentPartParam,
+    MessageDeltaEvent,
+    Run,
+    TextContentBlockParam,
+    TextDeltaBlock,
+)
+from openai.types.beta.threads.run_create_params import AdditionalMessage
+from openai.types.beta.threads.run_submit_tool_outputs_params import ToolOutput
+from openai.types.beta.threads.runs import RunStep
+from pydantic import Field, PrivateAttr, SecretStr, ValidationError
+
+from .._clients import BaseChatClient
+from .._middleware import use_chat_middleware
+from .._tools import AIFunction, HostedCodeInterpreterTool, HostedFileSearchTool, use_function_invocation
+from .._types import (
+    ChatMessage,
+    ChatOptions,
+    ChatResponse,
+    ChatResponseUpdate,
+    Contents,
+    FunctionCallContent,
+    FunctionResultContent,
+    Role,
+    TextContent,
+    ToolMode,
+    UriContent,
+    UsageContent,
+    UsageDetails,
+)
+from ..exceptions import ServiceInitializationError
+from ..observability import use_observability
+from ._shared import OpenAIConfigMixin, OpenAISettings
+
+if sys.version_info >= (3, 11):
+    from typing import Self  # pragma: no cover
+else:
+    from typing_extensions import Self  # pragma: no cover
+
+
+__all__ = ["OpenAIAssistantsClient"]
+
+
+@use_function_invocation
+@use_observability
+@use_chat_middleware
+class OpenAIAssistantsClient(OpenAIConfigMixin, BaseChatClient):
+    """OpenAI Assistants client."""
+
+    assistant_id: str | None = Field(default=None)
+    assistant_name: str | None = Field(default=None)
+    thread_id: str | None = Field(default=None)
+    _should_delete_assistant: bool = PrivateAttr(default=False)  # Track whether we should delete the assistant
+
+    def __init__(
+        self,
+        ai_model_id: str | None = None,
+        assistant_id: str | None = None,
+        assistant_name: str | None = None,
+        thread_id: str | None = None,
+        api_key: str | None = None,
+        org_id: str | None = None,
+        default_headers: Mapping[str, str] | None = None,
+        async_client: AsyncOpenAI | None = None,
+        env_file_path: str | None = None,
+        env_file_encoding: str | None = None,
+    ) -> None:
+        """Initialize an OpenAI Assistants client.
+
+        Args:
+            ai_model_id: OpenAI model name, see
+                https://platform.openai.com/docs/models
+            assistant_id: The ID of an OpenAI assistant to use.
+                If not provided, a new assistant will be created (and deleted after the request).
+            assistant_name: The name to use when creating new assistants.
+            thread_id: Default thread ID to use for conversations. Can be overridden by
+                conversation_id property, when making a request.
+                If not provided, a new thread will be created (and deleted after the request).
+            api_key: The optional API key to use. If provided will override,
+                the env vars or .env file value.
+            org_id: The optional org ID to use. If provided will override,
+                the env vars or .env file value.
+            default_headers: The default headers mapping of string keys to
+                string values for HTTP requests. (Optional)
+            async_client: An existing client to use. (Optional)
+            env_file_path: Use the environment settings file as a fallback
+                to environment variables. (Optional)
+            env_file_encoding: The encoding of the environment settings file. (Optional)
+        """
+        try:
+            openai_settings = OpenAISettings(
+                api_key=SecretStr(api_key) if api_key else None,
+                org_id=org_id,
+                chat_model_id=ai_model_id,
+                env_file_path=env_file_path,
+                env_file_encoding=env_file_encoding,
+            )
+        except ValidationError as ex:
+            raise ServiceInitializationError("Failed to create OpenAI settings.", ex) from ex
+
+        if not async_client and not openai_settings.api_key:
+            raise ServiceInitializationError(
+                "OpenAI API key is required. Set via 'api_key' parameter or 'OPENAI_API_KEY' environment variable."
+            )
+        if not openai_settings.chat_model_id:
+            raise ServiceInitializationError(
+                "OpenAI model ID is required. "
+                "Set via 'model_id' parameter or 'OPENAI_CHAT_MODEL_ID' environment variable."
+            )
+
+        super().__init__(
+            ai_model_id=openai_settings.chat_model_id,
+            assistant_id=assistant_id,  # type: ignore[reportCallIssue]
+            assistant_name=assistant_name,  # type: ignore[reportCallIssue]
+            thread_id=thread_id,  # type: ignore[reportCallIssue]
+            api_key=openai_settings.api_key.get_secret_value() if openai_settings.api_key else None,
+            org_id=openai_settings.org_id,
+            default_headers=default_headers,
+            client=async_client,
+        )
+
+    async def __aenter__(self) -> "Self":
+        """Async context manager entry."""
+        return self
+
+    async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None:
+        """Async context manager exit - clean up any assistants we created."""
+        await self.close()
+
+    async def close(self) -> None:
+        """Clean up any assistants we created."""
+        if self._should_delete_assistant and self.assistant_id is not None:
+            await self.client.beta.assistants.delete(self.assistant_id)
+            self.assistant_id = None
+            self._should_delete_assistant = False
+
+    async def _inner_get_response(
+        self,
+        *,
+        messages: MutableSequence[ChatMessage],
+        chat_options: ChatOptions,
+        **kwargs: Any,
+    ) -> ChatResponse:
+        return await ChatResponse.from_chat_response_generator(
+            updates=self._inner_get_streaming_response(messages=messages, chat_options=chat_options, **kwargs),
+            output_format_type=chat_options.response_format,
+        )
+
+    async def _inner_get_streaming_response(
+        self,
+        *,
+        messages: MutableSequence[ChatMessage],
+        chat_options: ChatOptions,
+        **kwargs: Any,
+    ) -> AsyncIterable[ChatResponseUpdate]:
+        # Extract necessary state from messages and options
+        run_options, tool_results = self._prepare_options(messages, chat_options, **kwargs)
+
+        # Get the thread ID
+        thread_id: str | None = (
+            chat_options.conversation_id
+            if chat_options.conversation_id is not None
+            else run_options.get("conversation_id", self.thread_id)
+        )
+
+        if thread_id is None and tool_results is not None:
+            raise ValueError("No thread ID was provided, but chat messages includes tool results.")
+
+        # Determine which assistant to use and create if needed
+        assistant_id = await self._get_assistant_id_or_create()
+
+        # Create the streaming response
+        stream, thread_id = await self._create_assistant_stream(thread_id, assistant_id, run_options, tool_results)
+
+        # Process and yield each update from the stream
+        async for update in self._process_stream_events(stream, thread_id):
+            yield update
+
+    async def _get_assistant_id_or_create(self) -> str:
+        """Determine which assistant to use and create if needed.
+
+        Returns:
+            str: The assistant_id to use.
+        """
+        # If no assistant is provided, create a temporary assistant
+        if self.assistant_id is None:
+            created_assistant = await self.client.beta.assistants.create(
+                name=self.assistant_name, model=self.ai_model_id
+            )
+
+            self.assistant_id = created_assistant.id
+            self._should_delete_assistant = True
+
+        return self.assistant_id
+
+    async def _create_assistant_stream(
+        self,
+        thread_id: str | None,
+        assistant_id: str,
+        run_options: dict[str, Any],
+        tool_results: list[FunctionResultContent] | None,
+    ) -> tuple[Any, str]:
+        """Create the assistant stream for processing.
+
+        Returns:
+            tuple: (stream, final_thread_id)
+        """
+        # Get any active run for this thread
+        thread_run = await self._get_active_thread_run(thread_id)
+
+        tool_run_id, tool_outputs = self._convert_function_results_to_tool_output(tool_results)
+
+        if thread_run is not None and tool_run_id is not None and tool_run_id == thread_run.id and tool_outputs:
+            # There's an active run and we have tool results to submit, so submit the results.
+            stream = self.client.beta.threads.runs.submit_tool_outputs_stream(  # type: ignore[reportDeprecated]
+                run_id=tool_run_id, thread_id=thread_run.thread_id, tool_outputs=tool_outputs
+            )
+            final_thread_id = thread_run.thread_id
+        else:
+            # Handle thread creation or cancellation
+            final_thread_id = await self._prepare_thread(thread_id, thread_run, run_options)
+
+            # Now create a new run and stream the results.
+            stream = self.client.beta.threads.runs.stream(  # type: ignore[reportDeprecated]
+                assistant_id=assistant_id, thread_id=final_thread_id, **run_options
+            )
+
+        return stream, final_thread_id
+
+    async def _get_active_thread_run(self, thread_id: str | None) -> Run | None:
+        """Get any active run for the given thread."""
+        if thread_id is None:
+            return None
+
+        async for run in self.client.beta.threads.runs.list(thread_id=thread_id, limit=1, order="desc"):  # type: ignore[reportDeprecated]
+            if run.status not in ["completed", "cancelled", "failed", "expired"]:
+                return run
+        return None
+
+    async def _prepare_thread(self, thread_id: str | None, thread_run: Run | None, run_options: dict[str, Any]) -> str:
+        """Prepare the thread for a new run, creating or cleaning up as needed."""
+        if thread_id is None:
+            # No thread ID was provided, so create a new thread.
+            thread = await self.client.beta.threads.create(  # type: ignore[reportDeprecated]
+                messages=run_options["additional_messages"],
+                tool_resources=run_options.get("tool_resources"),
+                metadata=run_options.get("metadata"),
+            )
+            run_options["additional_messages"] = []
+            run_options.pop("tool_resources", None)
+            return thread.id
+
+        if thread_run is not None:
+            # There was an active run; we need to cancel it before starting a new run.
+            await self.client.beta.threads.runs.cancel(run_id=thread_run.id, thread_id=thread_id)  # type: ignore[reportDeprecated]
+
+        return thread_id
+
+    async def _process_stream_events(self, stream: Any, thread_id: str) -> AsyncIterable[ChatResponseUpdate]:
+        response_id: str | None = None
+
+        async with stream as response_stream:
+            async for response in response_stream:
+                if response.event == "thread.run.created":
+                    yield ChatResponseUpdate(
+                        contents=[],
+                        conversation_id=thread_id,
+                        message_id=response_id,
+                        raw_representation=response.data,
+                        response_id=response_id,
+                        role=Role.ASSISTANT,
+                    )
+                elif response.event == "thread.run.step.created" and isinstance(response.data, RunStep):
+                    response_id = response.data.run_id
+                elif response.event == "thread.message.delta" and isinstance(response.data, MessageDeltaEvent):
+                    delta = response.data.delta
+                    role = Role.USER if delta.role == "user" else Role.ASSISTANT
+
+                    for delta_block in delta.content or []:
+                        if isinstance(delta_block, TextDeltaBlock) and delta_block.text and delta_block.text.value:
+                            yield ChatResponseUpdate(
+                                role=role,
+                                text=delta_block.text.value,
+                                conversation_id=thread_id,
+                                message_id=response_id,
+                                raw_representation=response.data,
+                                response_id=response_id,
+                            )
+                elif response.event == "thread.run.requires_action" and isinstance(response.data, Run):
+                    contents = self._create_function_call_contents(response.data, response_id)
+                    if contents:
+                        yield ChatResponseUpdate(
+                            role=Role.ASSISTANT,
+                            contents=contents,
+                            conversation_id=thread_id,
+                            message_id=response_id,
+                            raw_representation=response.data,
+                            response_id=response_id,
+                        )
+                elif (
+                    response.event == "thread.run.completed"
+                    and isinstance(response.data, Run)
+                    and response.data.usage is not None
+                ):
+                    usage = response.data.usage
+                    usage_content = UsageContent(
+                        UsageDetails(
+                            input_token_count=usage.prompt_tokens,
+                            output_token_count=usage.completion_tokens,
+                            total_token_count=usage.total_tokens,
+                        )
+                    )
+                    yield ChatResponseUpdate(
+                        role=Role.ASSISTANT,
+                        contents=[usage_content],
+                        conversation_id=thread_id,
+                        message_id=response_id,
+                        raw_representation=response.data,
+                        response_id=response_id,
+                    )
+                else:
+                    yield ChatResponseUpdate(
+                        contents=[],
+                        conversation_id=thread_id,
+                        message_id=response_id,
+                        raw_representation=response.data,
+                        response_id=response_id,
+                        role=Role.ASSISTANT,
+                    )
+
+    def _create_function_call_contents(self, event_data: Run, response_id: str | None) -> list[Contents]:
+        """Create function call contents from a tool action event."""
+        contents: list[Contents] = []
+
+        if event_data.required_action is not None:
+            for tool_call in event_data.required_action.submit_tool_outputs.tool_calls:
+                call_id = json.dumps([response_id, tool_call.id])
+                function_name = tool_call.function.name
+                function_arguments = json.loads(tool_call.function.arguments)
+                contents.append(FunctionCallContent(call_id=call_id, name=function_name, arguments=function_arguments))
+
+        return contents
+
+    def _prepare_options(
+        self,
+        messages: MutableSequence[ChatMessage],
+        chat_options: ChatOptions | None,
+        **kwargs: Any,
+    ) -> tuple[dict[str, Any], list[FunctionResultContent] | None]:
+        run_options: dict[str, Any] = {**kwargs}
+
+        if chat_options is not None:
+            run_options["max_completion_tokens"] = chat_options.max_tokens
+            run_options["model"] = chat_options.model_id
+            run_options["top_p"] = chat_options.top_p
+            run_options["temperature"] = chat_options.temperature
+
+            if chat_options.allow_multiple_tool_calls is not None:
+                run_options["parallel_tool_calls"] = chat_options.allow_multiple_tool_calls
+
+            if chat_options.tool_choice is not None:
+                tool_definitions: list[MutableMapping[str, Any]] = []
+                if chat_options.tool_choice != "none" and chat_options.tools is not None:
+                    for tool in chat_options.tools:
+                        if isinstance(tool, AIFunction):
+                            tool_definitions.append(tool.to_json_schema_spec())  # type: ignore[reportUnknownArgumentType]
+                        elif isinstance(tool, HostedCodeInterpreterTool):
+                            tool_definitions.append({"type": "code_interpreter"})
+                        elif isinstance(tool, HostedFileSearchTool):
+                            params: dict[str, Any] = {
+                                "type": "file_search",
+                            }
+                            if tool.max_results is not None:
+                                params["max_num_results"] = tool.max_results
+                            tool_definitions.append(params)
+                        elif isinstance(tool, MutableMapping):
+                            tool_definitions.append(tool)
+
+                if len(tool_definitions) > 0:
+                    run_options["tools"] = tool_definitions
+
+                if chat_options.tool_choice == "none" or chat_options.tool_choice == "auto":
+                    run_options["tool_choice"] = chat_options.tool_choice
+                elif (
+                    isinstance(chat_options.tool_choice, ToolMode)
+                    and chat_options.tool_choice == "required"
+                    and chat_options.tool_choice.required_function_name is not None
+                ):
+                    run_options["tool_choice"] = {
+                        "type": "function",
+                        "function": {"name": chat_options.tool_choice.required_function_name},
+                    }
+
+            if chat_options.response_format is not None:
+                run_options["response_format"] = {
+                    "type": "json_schema",
+                    "json_schema": chat_options.response_format.model_json_schema(),
+                }
+
+        instructions: list[str] = [chat_options.instructions] if chat_options and chat_options.instructions else []
+        tool_results: list[FunctionResultContent] | None = None
+
+        additional_messages: list[AdditionalMessage] | None = None
+
+        # System/developer messages are turned into instructions,
+        # since there is no such message roles in OpenAI Assistants.
+        # All other messages are added 1:1.
+        for chat_message in messages:
+            if chat_message.role.value in ["system", "developer"]:
+                for text_content in [content for content in chat_message.contents if isinstance(content, TextContent)]:
+                    instructions.append(text_content.text)
+
+                continue
+
+            message_contents: list[MessageContentPartParam] = []
+
+            for content in chat_message.contents:
+                if isinstance(content, TextContent):
+                    message_contents.append(TextContentBlockParam(type="text", text=content.text))
+                elif isinstance(content, UriContent) and content.has_top_level_media_type("image"):
+                    message_contents.append(
+                        ImageURLContentBlockParam(type="image_url", image_url=ImageURLParam(url=content.uri))
+                    )
+                elif isinstance(content, FunctionResultContent):
+                    if tool_results is None:
+                        tool_results = []
+                    tool_results.append(content)
+
+            if len(message_contents) > 0:
+                if additional_messages is None:
+                    additional_messages = []
+                additional_messages.append(
+                    AdditionalMessage(
+                        role="assistant" if chat_message.role == Role.ASSISTANT else "user",
+                        content=message_contents,
+                    )
+                )
+
+        if additional_messages is not None:
+            run_options["additional_messages"] = additional_messages
+
+        if len(instructions) > 0:
+            run_options["instructions"] = "".join(instructions)
+
+        return run_options, tool_results
+
+    def _convert_function_results_to_tool_output(
+        self,
+        tool_results: list[FunctionResultContent] | None,
+    ) -> tuple[str | None, list[ToolOutput] | None]:
+        run_id: str | None = None
+        tool_outputs: list[ToolOutput] | None = None
+
+        if tool_results:
+            for function_result_content in tool_results:
+                # When creating the FunctionCallContent, we created it with a CallId == [runId, callId].
+                # We need to extract the run ID and ensure that the ToolOutput we send back to Azure
+                # is only the call ID.
+                run_and_call_ids: list[str] = json.loads(function_result_content.call_id)
+
+                if (
+                    not run_and_call_ids
+                    or len(run_and_call_ids) != 2
+                    or not run_and_call_ids[0]
+                    or not run_and_call_ids[1]
+                    or (run_id is not None and run_id != run_and_call_ids[0])
+                ):
+                    continue
+
+                run_id = run_and_call_ids[0]
+                call_id = run_and_call_ids[1]
+
+                if tool_outputs is None:
+                    tool_outputs = []
+                tool_outputs.append(ToolOutput(tool_call_id=call_id, output=str(function_result_content.result)))
+
+        return run_id, tool_outputs
+
+    def _update_agent_name(self, agent_name: str | None) -> None:
+        """Update the agent name in the chat client.
+
+        Args:
+            agent_name: The new name for the agent.
+        """
+        # This is a no-op in the base class, but can be overridden by subclasses
+        # to update the agent name in the client.
+        if agent_name and not self.assistant_name:
+            self.assistant_name = agent_name

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

@@ -0,0 +1,540 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import json
+import sys
+from collections.abc import AsyncIterable, Mapping, MutableMapping, MutableSequence, Sequence
+from datetime import datetime
+from itertools import chain
+from typing import Any, TypeVar
+
+from openai import AsyncOpenAI, BadRequestError
+from openai.lib._parsing._completions import type_to_response_format_param
+from openai.types import CompletionUsage
+from openai.types.chat.chat_completion import ChatCompletion, Choice
+from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
+from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
+from openai.types.chat.chat_completion_message_custom_tool_call import ChatCompletionMessageCustomToolCall
+from pydantic import BaseModel, SecretStr, ValidationError
+
+from .._clients import BaseChatClient
+from .._logging import get_logger
+from .._middleware import use_chat_middleware
+from .._tools import AIFunction, HostedWebSearchTool, ToolProtocol, use_function_invocation
+from .._types import (
+    ChatMessage,
+    ChatOptions,
+    ChatResponse,
+    ChatResponseUpdate,
+    Contents,
+    DataContent,
+    FinishReason,
+    FunctionCallContent,
+    FunctionResultContent,
+    Role,
+    TextContent,
+    UriContent,
+    UsageContent,
+    UsageDetails,
+)
+from ..exceptions import (
+    ServiceInitializationError,
+    ServiceInvalidRequestError,
+    ServiceResponseException,
+)
+from ..observability import use_observability
+from ._exceptions import OpenAIContentFilterException
+from ._shared import OpenAIBase, OpenAIConfigMixin, OpenAISettings, prepare_function_call_results
+
+if sys.version_info >= (3, 12):
+    from typing import override  # type: ignore # pragma: no cover
+else:
+    from typing_extensions import override  # type: ignore[import] # pragma: no cover
+
+__all__ = ["OpenAIChatClient"]
+
+logger = get_logger("agent_framework.openai")
+
+
+# region Base Client
+class OpenAIBaseChatClient(OpenAIBase, BaseChatClient):
+    """OpenAI Chat completion class."""
+
+    async def _inner_get_response(
+        self,
+        *,
+        messages: MutableSequence[ChatMessage],
+        chat_options: ChatOptions,
+        **kwargs: Any,
+    ) -> ChatResponse:
+        options_dict = self._prepare_options(messages, chat_options)
+        try:
+            return self._create_chat_response(
+                await self.client.chat.completions.create(stream=False, **options_dict), chat_options
+            )
+        except BadRequestError as ex:
+            if ex.code == "content_filter":
+                raise OpenAIContentFilterException(
+                    f"{type(self)} service encountered a content error: {ex}",
+                    inner_exception=ex,
+                ) from ex
+            raise ServiceResponseException(
+                f"{type(self)} service failed to complete the prompt: {ex}",
+                inner_exception=ex,
+            ) from ex
+        except Exception as ex:
+            raise ServiceResponseException(
+                f"{type(self)} service failed to complete the prompt: {ex}",
+                inner_exception=ex,
+            ) from ex
+
+    async def _inner_get_streaming_response(
+        self,
+        *,
+        messages: MutableSequence[ChatMessage],
+        chat_options: ChatOptions,
+        **kwargs: Any,
+    ) -> AsyncIterable[ChatResponseUpdate]:
+        options_dict = self._prepare_options(messages, chat_options)
+        options_dict["stream_options"] = {"include_usage": True}
+        try:
+            async for chunk in await self.client.chat.completions.create(stream=True, **options_dict):
+                if len(chunk.choices) == 0 and chunk.usage is None:
+                    continue
+                yield self._create_chat_response_update(chunk)
+        except BadRequestError as ex:
+            if ex.code == "content_filter":
+                raise OpenAIContentFilterException(
+                    f"{type(self)} service encountered a content error: {ex}",
+                    inner_exception=ex,
+                ) from ex
+            raise ServiceResponseException(
+                f"{type(self)} service failed to complete the prompt: {ex}",
+                inner_exception=ex,
+            ) from ex
+        except Exception as ex:
+            raise ServiceResponseException(
+                f"{type(self)} service failed to complete the prompt: {ex}",
+                inner_exception=ex,
+            ) from ex
+
+    # region content creation
+
+    def _chat_to_tool_spec(self, tools: Sequence[ToolProtocol | MutableMapping[str, Any]]) -> list[dict[str, Any]]:
+        chat_tools: list[dict[str, Any]] = []
+        for tool in tools:
+            if isinstance(tool, ToolProtocol):
+                match tool:
+                    case AIFunction():
+                        chat_tools.append(tool.to_json_schema_spec())
+                    case _:
+                        logger.debug("Unsupported tool passed (type: %s), ignoring", type(tool))
+            else:
+                chat_tools.append(tool if isinstance(tool, dict) else dict(tool))
+        return chat_tools
+
+    def _process_web_search_tool(
+        self, tools: Sequence[ToolProtocol | MutableMapping[str, Any]]
+    ) -> dict[str, Any] | None:
+        for tool in tools:
+            if isinstance(tool, HostedWebSearchTool):
+                # Web search tool requires special handling
+                return (
+                    {
+                        "user_location": {
+                            "approximate": tool.additional_properties.get("user_location", None),
+                            "type": "approximate",
+                        }
+                    }
+                    if tool.additional_properties and "user_location" in tool.additional_properties
+                    else {}
+                )
+
+        return None
+
+    def _prepare_options(self, messages: MutableSequence[ChatMessage], chat_options: ChatOptions) -> dict[str, Any]:
+        # Preprocess web search tool if it exists
+        options_dict = chat_options.to_provider_settings()
+        instructions = options_dict.pop("instructions", None)
+        if instructions:
+            messages = [ChatMessage(role="system", text=instructions), *messages]
+        if messages and "messages" not in options_dict:
+            options_dict["messages"] = self._prepare_chat_history_for_request(messages)
+        if "messages" not in options_dict:
+            raise ServiceInvalidRequestError("Messages are required for chat completions")
+        if chat_options.tools is not None:
+            web_search_options = self._process_web_search_tool(chat_options.tools)
+            if web_search_options:
+                options_dict["web_search_options"] = web_search_options
+            options_dict["tools"] = self._chat_to_tool_spec(chat_options.tools)
+        if not options_dict.get("tools", None):
+            options_dict.pop("tools", None)
+            options_dict.pop("parallel_tool_calls", None)
+            options_dict.pop("tool_choice", None)
+
+        if "model" not in options_dict:
+            options_dict["model"] = self.ai_model_id
+        if (
+            chat_options.response_format
+            and isinstance(chat_options.response_format, type)
+            and issubclass(chat_options.response_format, BaseModel)
+        ):
+            options_dict["response_format"] = type_to_response_format_param(chat_options.response_format)
+        return options_dict
+
+    def _create_chat_response(self, response: ChatCompletion, chat_options: ChatOptions) -> "ChatResponse":
+        """Create a chat message content object from a choice."""
+        response_metadata = self._get_metadata_from_chat_response(response)
+        messages: list[ChatMessage] = []
+        finish_reason: FinishReason | None = None
+        for choice in response.choices:
+            response_metadata.update(self._get_metadata_from_chat_choice(choice))
+            if choice.finish_reason:
+                finish_reason = FinishReason(value=choice.finish_reason)
+            contents: list[Contents] = []
+            if text_content := self._parse_text_from_choice(choice):
+                contents.append(text_content)
+            if parsed_tool_calls := [tool for tool in self._get_tool_calls_from_chat_choice(choice)]:
+                contents.extend(parsed_tool_calls)
+            messages.append(ChatMessage(role="assistant", contents=contents))
+        return ChatResponse(
+            response_id=response.id,
+            created_at=datetime.fromtimestamp(response.created).strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
+            usage_details=self._usage_details_from_openai(response.usage) if response.usage else None,
+            messages=messages,
+            model_id=response.model,
+            additional_properties=response_metadata,
+            finish_reason=finish_reason,
+            response_format=chat_options.response_format,
+        )
+
+    def _create_chat_response_update(
+        self,
+        chunk: ChatCompletionChunk,
+    ) -> ChatResponseUpdate:
+        """Create a streaming chat message content object from a choice."""
+        chunk_metadata = self._get_metadata_from_streaming_chat_response(chunk)
+        if chunk.usage:
+            return ChatResponseUpdate(
+                role=Role.ASSISTANT,
+                contents=[UsageContent(details=self._usage_details_from_openai(chunk.usage), raw_representation=chunk)],
+                model_id=chunk.model,
+                additional_properties=chunk_metadata,
+                response_id=chunk.id,
+                message_id=chunk.id,
+            )
+        contents: list[Contents] = []
+        finish_reason: FinishReason | None = None
+        for choice in chunk.choices:
+            chunk_metadata.update(self._get_metadata_from_chat_choice(choice))
+            contents.extend(self._get_tool_calls_from_chat_choice(choice))
+            if choice.finish_reason:
+                finish_reason = FinishReason(value=choice.finish_reason)
+
+            if text_content := self._parse_text_from_choice(choice):
+                contents.append(text_content)
+        return ChatResponseUpdate(
+            created_at=datetime.fromtimestamp(chunk.created).strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
+            contents=contents,
+            role=Role.ASSISTANT,
+            model_id=chunk.model,
+            additional_properties=chunk_metadata,
+            finish_reason=finish_reason,
+            raw_representation=chunk,
+            response_id=chunk.id,
+            message_id=chunk.id,
+        )
+
+    def _usage_details_from_openai(self, usage: CompletionUsage) -> UsageDetails:
+        details = UsageDetails(
+            input_token_count=usage.prompt_tokens,
+            output_token_count=usage.completion_tokens,
+            total_token_count=usage.total_tokens,
+        )
+        if usage.completion_tokens_details:
+            if tokens := usage.completion_tokens_details.accepted_prediction_tokens:
+                details["completion/accepted_prediction_tokens"] = tokens
+            if tokens := usage.completion_tokens_details.audio_tokens:
+                details["completion/audio_tokens"] = tokens
+            if tokens := usage.completion_tokens_details.reasoning_tokens:
+                details["completion/reasoning_tokens"] = tokens
+            if tokens := usage.completion_tokens_details.rejected_prediction_tokens:
+                details["completion/rejected_prediction_tokens"] = tokens
+        if usage.prompt_tokens_details:
+            if tokens := usage.prompt_tokens_details.audio_tokens:
+                details["prompt/audio_tokens"] = tokens
+            if tokens := usage.prompt_tokens_details.cached_tokens:
+                details["prompt/cached_tokens"] = tokens
+        return details
+
+    def _parse_text_from_choice(self, choice: Choice | ChunkChoice) -> TextContent | None:
+        """Parse the choice into a TextContent object."""
+        message = choice.message if isinstance(choice, Choice) else choice.delta
+        if message.content:
+            return TextContent(text=message.content, raw_representation=choice)
+        if hasattr(message, "refusal") and message.refusal:
+            return TextContent(text=message.refusal, raw_representation=choice)
+        return None
+
+    def _get_metadata_from_chat_response(self, response: ChatCompletion) -> dict[str, Any]:
+        """Get metadata from a chat response."""
+        return {
+            "system_fingerprint": response.system_fingerprint,
+        }
+
+    def _get_metadata_from_streaming_chat_response(self, response: ChatCompletionChunk) -> dict[str, Any]:
+        """Get metadata from a streaming chat response."""
+        return {
+            "system_fingerprint": response.system_fingerprint,
+        }
+
+    def _get_metadata_from_chat_choice(self, choice: Choice | ChunkChoice) -> dict[str, Any]:
+        """Get metadata from a chat choice."""
+        return {
+            "logprobs": getattr(choice, "logprobs", None),
+        }
+
+    def _get_tool_calls_from_chat_choice(self, choice: Choice | ChunkChoice) -> list[Contents]:
+        """Get tool calls from a chat choice."""
+        resp: list[Contents] = []
+        content = choice.message if isinstance(choice, Choice) else choice.delta
+        if content and content.tool_calls:
+            for tool in content.tool_calls:
+                if not isinstance(tool, ChatCompletionMessageCustomToolCall) and tool.function:
+                    # ignoring tool.custom
+                    fcc = FunctionCallContent(
+                        call_id=tool.id if tool.id else "",
+                        name=tool.function.name if tool.function.name else "",
+                        arguments=tool.function.arguments if tool.function.arguments else "",
+                        raw_representation=tool.function,
+                    )
+                    resp.append(fcc)
+
+        # When you enable asynchronous content filtering in Azure OpenAI, you may receive empty deltas
+        return resp
+
+    def _prepare_chat_history_for_request(
+        self,
+        chat_messages: Sequence[ChatMessage],
+        role_key: str = "role",
+        content_key: str = "content",
+    ) -> list[dict[str, Any]]:
+        """Prepare the chat history for a request.
+
+        Allowing customization of the key names for role/author, and optionally overriding the role.
+
+        Role.TOOL messages need to be formatted different than system/user/assistant messages:
+            They require a "tool_call_id" and (function) "name" key, and the "metadata" key should
+            be removed. The "encoding" key should also be removed.
+
+        Override this method to customize the formatting of the chat history for a request.
+
+        Args:
+            chat_messages: The chat history to prepare.
+            role_key: The key name for the role/author.
+            content_key: The key name for the content/message.
+
+        Returns:
+            prepared_chat_history (Any): The prepared chat history for a request.
+        """
+        list_of_list = [self._openai_chat_message_parser(message) for message in chat_messages]
+        # Flatten the list of lists into a single list
+        return list(chain.from_iterable(list_of_list))
+
+    # region Parsers
+
+    def _openai_chat_message_parser(self, message: ChatMessage) -> list[dict[str, Any]]:
+        """Parse a chat message into the openai format."""
+        all_messages: list[dict[str, Any]] = []
+        for content in message.contents:
+            args: dict[str, Any] = {
+                "role": message.role.value if isinstance(message.role, Role) else message.role,
+            }
+            if message.additional_properties:
+                args["metadata"] = message.additional_properties
+            match content:
+                case FunctionCallContent():
+                    if all_messages and "tool_calls" in all_messages[-1]:
+                        # If the last message already has tool calls, append to it
+                        all_messages[-1]["tool_calls"].append(self._openai_content_parser(content))
+                    else:
+                        args["tool_calls"] = [self._openai_content_parser(content)]  # type: ignore
+                case FunctionResultContent():
+                    args["tool_call_id"] = content.call_id
+                    if content.result is not None:
+                        args["content"] = prepare_function_call_results(content.result)
+                    elif content.exception is not None:
+                        # Send the exception message to the model
+                        # Otherwise we won't have any channels to talk to OpenAI
+                        # TODO(yuge): This should ideally be customizable
+                        args["content"] = "Error: " + str(content.exception)
+                case _:
+                    if "content" not in args:
+                        args["content"] = []
+                    # this is a list to allow multi-modal content
+                    args["content"].append(self._openai_content_parser(content))  # type: ignore
+            if "content" in args or "tool_calls" in args:
+                all_messages.append(args)
+        return all_messages
+
+    def _openai_content_parser(self, content: Contents) -> dict[str, Any]:
+        """Parse contents into the openai format."""
+        match content:
+            case FunctionCallContent():
+                args = json.dumps(content.arguments) if isinstance(content.arguments, Mapping) else content.arguments
+                return {
+                    "id": content.call_id,
+                    "type": "function",
+                    "function": {"name": content.name, "arguments": args},
+                }
+            case FunctionResultContent():
+                return {
+                    "tool_call_id": content.call_id,
+                    "content": content.result,
+                }
+            case DataContent() | UriContent() if content.has_top_level_media_type("image"):
+                return {
+                    "type": "image_url",
+                    "image_url": {"url": content.uri},
+                }
+            case DataContent() | UriContent() if content.has_top_level_media_type("audio"):
+                if content.media_type and "wav" in content.media_type:
+                    audio_format = "wav"
+                elif content.media_type and "mp3" in content.media_type:
+                    audio_format = "mp3"
+                else:
+                    # Fallback to default to_dict for unsupported audio formats
+                    return content.to_dict(exclude_none=True)
+
+                # Extract base64 data from data URI
+                audio_data = content.uri
+                if audio_data.startswith("data:"):
+                    # Extract just the base64 part after "data:audio/format;base64,"
+                    audio_data = audio_data.split(",", 1)[-1]
+
+                return {
+                    "type": "input_audio",
+                    "input_audio": {
+                        "data": audio_data,
+                        "format": audio_format,
+                    },
+                }
+            case DataContent() | UriContent() if content.media_type and content.media_type.startswith("application/"):
+                if content.media_type == "application/pdf":
+                    if content.uri.startswith("data:"):
+                        filename = (
+                            getattr(content, "filename", None)
+                            or content.additional_properties.get("filename", "document.pdf")
+                            if hasattr(content, "additional_properties") and content.additional_properties
+                            else "document.pdf"
+                        )
+                        return {
+                            "type": "file",
+                            "file": {
+                                "file_data": content.uri,  # Send full data URI
+                                "filename": filename,
+                            },
+                        }
+
+                    return content.to_dict(exclude_none=True)
+
+                return content.to_dict(exclude_none=True)
+            case _:
+                return content.to_dict(exclude_none=True)
+
+    @override
+    def service_url(self) -> str:
+        """Get the URL of the service.
+
+        Override this in the subclass to return the proper URL.
+        If the service does not have a URL, return None.
+        """
+        return str(self.client.base_url) if self.client else "Unknown"
+
+
+# region Public client
+
+TOpenAIChatClient = TypeVar("TOpenAIChatClient", bound="OpenAIChatClient")
+
+
+@use_function_invocation
+@use_observability
+@use_chat_middleware
+class OpenAIChatClient(OpenAIConfigMixin, OpenAIBaseChatClient):
+    """OpenAI Chat completion class."""
+
+    def __init__(
+        self,
+        ai_model_id: str | None = None,
+        api_key: str | None = None,
+        org_id: str | None = None,
+        default_headers: Mapping[str, str] | None = None,
+        async_client: AsyncOpenAI | None = None,
+        instruction_role: str | None = None,
+        base_url: str | None = None,
+        env_file_path: str | None = None,
+        env_file_encoding: str | None = None,
+    ) -> None:
+        """Initialize an OpenAIChatCompletion service.
+
+        Args:
+            ai_model_id: OpenAI model name, see
+                https://platform.openai.com/docs/models
+            api_key: The optional API key to use. If provided will override,
+                the env vars or .env file value.
+            org_id: The optional org ID to use. If provided will override,
+                the env vars or .env file value.
+            default_headers: The default headers mapping of string keys to
+                string values for HTTP requests. (Optional)
+            async_client: An existing client to use. (Optional)
+            instruction_role: The role to use for 'instruction' messages, for example,
+                "system" or "developer". If not provided, the default is "system".
+            base_url: The optional base URL to use. If provided will override
+                the standard value for a OpenAI connector,
+                the env vars or .env file value.
+            env_file_path: Use the environment settings file as a fallback
+                to environment variables. (Optional)
+            env_file_encoding: The encoding of the environment settings file. (Optional)
+        """
+        try:
+            openai_settings = OpenAISettings(
+                api_key=SecretStr(api_key) if api_key else None,
+                base_url=base_url,
+                org_id=org_id,
+                chat_model_id=ai_model_id,
+                env_file_path=env_file_path,
+                env_file_encoding=env_file_encoding,
+            )
+        except ValidationError as ex:
+            raise ServiceInitializationError("Failed to create OpenAI settings.", ex) from ex
+
+        if not async_client and not openai_settings.api_key:
+            raise ServiceInitializationError(
+                "OpenAI API key is required. Set via 'api_key' parameter or 'OPENAI_API_KEY' environment variable."
+            )
+        if not openai_settings.chat_model_id:
+            raise ServiceInitializationError(
+                "OpenAI model ID is required. "
+                "Set via 'ai_model_id' parameter or 'OPENAI_CHAT_MODEL_ID' environment variable."
+            )
+
+        super().__init__(
+            ai_model_id=openai_settings.chat_model_id,
+            api_key=openai_settings.api_key.get_secret_value() if openai_settings.api_key else None,
+            base_url=openai_settings.base_url if openai_settings.base_url else None,
+            org_id=openai_settings.org_id,
+            default_headers=default_headers,
+            client=async_client,
+            instruction_role=instruction_role,
+        )
+
+    @classmethod
+    def from_dict(cls: type[TOpenAIChatClient], settings: dict[str, Any]) -> TOpenAIChatClient:
+        """Initialize an Open AI Chat Client from a dictionary of settings.
+
+        Args:
+            settings: A dictionary of settings for the service.
+        """
+        return cls(**settings)
+
+
+# endregion

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

@@ -0,0 +1,92 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any
+
+from openai import BadRequestError
+
+from ..exceptions import ServiceContentFilterException
+
+__all__ = ["ContentFilterResultSeverity", "OpenAIContentFilterException"]
+
+
+class ContentFilterResultSeverity(Enum):
+    """The severity of the content filter result."""
+
+    HIGH = "high"
+    MEDIUM = "medium"
+    SAFE = "safe"
+    LOW = "low"
+
+
+@dataclass
+class ContentFilterResult:
+    """The result of a content filter check."""
+
+    filtered: bool = False
+    detected: bool = False
+    severity: ContentFilterResultSeverity = ContentFilterResultSeverity.SAFE
+
+    @classmethod
+    def from_inner_error_result(cls, inner_error_results: dict[str, Any]) -> "ContentFilterResult":
+        """Creates a ContentFilterResult from the inner error results.
+
+        Args:
+            key (str): The key to get the inner error result from.
+            inner_error_results (Dict[str, Any]): The inner error results.
+
+        Returns:
+            ContentFilterResult: The ContentFilterResult.
+        """
+        return cls(
+            filtered=inner_error_results.get("filtered", False),
+            detected=inner_error_results.get("detected", False),
+            severity=ContentFilterResultSeverity(
+                inner_error_results.get("severity", ContentFilterResultSeverity.SAFE.value)
+            ),
+        )
+
+
+class ContentFilterCodes(Enum):
+    """Content filter codes."""
+
+    RESPONSIBLE_AI_POLICY_VIOLATION = "ResponsibleAIPolicyViolation"
+
+
+@dataclass
+class OpenAIContentFilterException(ServiceContentFilterException):
+    """AI exception for an error from Azure OpenAI's content filter."""
+
+    # The parameter that caused the error.
+    param: str | None
+
+    # The error code specific to the content filter.
+    content_filter_code: ContentFilterCodes
+
+    # The results of the different content filter checks.
+    content_filter_result: dict[str, ContentFilterResult]
+
+    def __init__(
+        self,
+        message: str,
+        inner_exception: BadRequestError,
+    ) -> None:
+        """Initializes a new instance of the ContentFilterAIException class.
+
+        Args:
+            message (str): The error message.
+            inner_exception (Exception): The inner exception.
+        """
+        super().__init__(message)
+
+        self.param = inner_exception.param
+        if inner_exception.body is not None and isinstance(inner_exception.body, dict):
+            inner_error = inner_exception.body.get("innererror", {})  # type: ignore
+            self.content_filter_code = ContentFilterCodes(
+                inner_error.get("code", ContentFilterCodes.RESPONSIBLE_AI_POLICY_VIOLATION.value)  # type: ignore
+            )
+            self.content_filter_result = {
+                key: ContentFilterResult.from_inner_error_result(values)  # type: ignore
+                for key, values in inner_error.get("content_filter_result", {}).items()  # type: ignore
+            }

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

@@ -0,0 +1,198 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import json
+import logging
+from collections.abc import Mapping
+from copy import copy
+from typing import Annotated, Any, ClassVar, Union
+
+from openai import (
+    AsyncOpenAI,
+    AsyncStream,
+    _legacy_response,  # type: ignore
+)
+from openai.types import Completion
+from openai.types.audio import Transcription
+from openai.types.chat import ChatCompletion, ChatCompletionChunk
+from openai.types.images_response import ImagesResponse
+from openai.types.responses.response import Response
+from openai.types.responses.response_stream_event import ResponseStreamEvent
+from pydantic import ConfigDict, Field, SecretStr, validate_call
+from pydantic.types import StringConstraints
+
+from .._logging import get_logger
+from .._pydantic import AFBaseModel, AFBaseSettings
+from .._telemetry import APP_INFO, USER_AGENT_KEY, prepend_agent_framework_to_user_agent
+from .._types import ChatOptions, Contents
+from ..exceptions import ServiceInitializationError
+
+logger: logging.Logger = get_logger("agent_framework.openai")
+
+
+RESPONSE_TYPE = Union[
+    ChatCompletion,
+    Completion,
+    AsyncStream[ChatCompletionChunk],
+    AsyncStream[Completion],
+    list[Any],
+    ImagesResponse,
+    Response,
+    AsyncStream[ResponseStreamEvent],
+    Transcription,
+    _legacy_response.HttpxBinaryResponseContent,
+]
+
+OPTION_TYPE = Union[ChatOptions, dict[str, Any]]
+
+
+__all__ = [
+    "OpenAISettings",
+]
+
+
+def _prepare_function_call_results_as_dumpable(content: Contents | Any | list[Contents | Any]) -> Any:
+    if isinstance(content, list):
+        # Particularly deal with lists of Content
+        return [_prepare_function_call_results_as_dumpable(item) for item in content]
+    if isinstance(content, dict):
+        return {k: _prepare_function_call_results_as_dumpable(v) for k, v in content.items()}
+    if hasattr(content, "to_dict"):
+        return content.to_dict(exclude={"raw_representation", "additional_properties"})
+    return content
+
+
+def prepare_function_call_results(content: Contents | Any | list[Contents | Any]) -> str | list[str]:
+    """Prepare the values of the function call results."""
+    if isinstance(content, Contents):
+        # For BaseContent objects, use to_dict and serialize to JSON
+        return json.dumps(content.to_dict(exclude={"raw_representation", "additional_properties"}))
+
+    dumpable = _prepare_function_call_results_as_dumpable(content)
+    if isinstance(dumpable, str):
+        return dumpable
+    # fallback
+    return json.dumps(dumpable)
+
+
+class OpenAISettings(AFBaseSettings):
+    """OpenAI environment settings.
+
+    The settings are first loaded from environment variables with the prefix 'OPENAI_'.
+    If the environment variables are not found, the settings can be loaded from a .env file with the
+    encoding 'utf-8'. If the settings are not found in the .env file, the settings are ignored;
+    however, validation will fail alerting that the settings are missing.
+
+    Args:
+        api_key: OpenAI API key, see https://platform.openai.com/account/api-keys
+            (Env var OPENAI_API_KEY)
+        base_url: The base URL for the OpenAI API.
+            (Env var OPENAI_BASE_URL)
+        org_id: This is usually optional unless your account belongs to multiple organizations.
+            (Env var OPENAI_ORG_ID)
+        chat_model_id: The OpenAI chat model ID to use, for example, gpt-3.5-turbo or gpt-4.
+            (Env var OPENAI_CHAT_MODEL_ID)
+        responses_model_id: The OpenAI responses model ID to use, for example, gpt-4o or o1.
+            (Env var OPENAI_RESPONSES_MODEL_ID)
+        env_file_path: The path to the .env file to load settings from.
+        env_file_encoding: The encoding of the .env file, defaults to 'utf-8'.
+    """
+
+    env_prefix: ClassVar[str] = "OPENAI_"
+
+    api_key: SecretStr | None = None
+    base_url: str | None = None
+    org_id: str | None = None
+    chat_model_id: str | None = None
+    responses_model_id: str | None = None
+
+
+class OpenAIBase(AFBaseModel):
+    """Base class for OpenAI Clients."""
+
+    client: AsyncOpenAI
+    ai_model_id: Annotated[str, StringConstraints(strip_whitespace=True, min_length=1)]
+
+
+class OpenAIConfigMixin(OpenAIBase):
+    """Internal class for configuring a connection to an OpenAI service."""
+
+    OTEL_PROVIDER_NAME: ClassVar[str] = "openai"  # type: ignore[reportIncompatibleVariableOverride, misc]
+
+    @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
+    def __init__(
+        self,
+        ai_model_id: str = Field(min_length=1),
+        api_key: str | None = Field(min_length=1),
+        org_id: str | None = None,
+        default_headers: Mapping[str, str] | None = None,
+        client: AsyncOpenAI | None = None,
+        instruction_role: str | None = None,
+        base_url: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize a client for OpenAI services.
+
+        This constructor sets up a client to interact with OpenAI's API, allowing for
+        different types of AI model interactions, like chat or text completion.
+
+        Args:
+            ai_model_id: OpenAI model identifier. Must be non-empty.
+                Default to a preset value.
+            api_key: OpenAI API key for authentication.
+                Must be non-empty. (Optional)
+            org_id: OpenAI organization ID. This is optional
+                unless the account belongs to multiple organizations.
+            default_headers: Default headers
+                for HTTP requests. (Optional)
+            client: An existing OpenAI client, optional.
+            instruction_role: The role to use for 'instruction'
+                messages, for example, summarization prompts could use `developer` or `system`. (Optional)
+            base_url: The optional base URL to use. If provided will override the standard value for a OpenAI connector.
+                Will not be used when supplying a custom client.
+            kwargs: Additional keyword arguments.
+
+        """
+        # Merge APP_INFO into the headers if it exists
+        merged_headers = dict(copy(default_headers)) if default_headers else {}
+        if APP_INFO:
+            merged_headers.update(APP_INFO)
+            merged_headers = prepend_agent_framework_to_user_agent(merged_headers)
+
+        if not client:
+            if not api_key:
+                raise ServiceInitializationError("Please provide an api_key")
+            args: dict[str, Any] = {"api_key": api_key, "default_headers": merged_headers}
+            if org_id:
+                args["organization"] = org_id
+            if base_url:
+                args["base_url"] = base_url
+            client = AsyncOpenAI(**args)
+        args = {
+            "ai_model_id": ai_model_id,
+            "client": client,
+        }
+        if instruction_role:
+            args["instruction_role"] = instruction_role
+        super().__init__(**args, **kwargs)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Create a dict of the service settings."""
+        client_settings = {
+            "api_key": self.client.api_key,
+            "default_headers": {k: v for k, v in self.client.default_headers.items() if k != USER_AGENT_KEY},
+        }
+        if self.client.organization:
+            client_settings["org_id"] = self.client.organization
+        base = self.model_dump(
+            exclude={
+                "prompt_tokens",
+                "completion_tokens",
+                "total_tokens",
+                "api_type",
+                "client",
+            },
+            by_alias=True,
+            exclude_none=True,
+        )
+        base.update(client_settings)
+        return base

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

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

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

@@ -0,0 +1,5 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from agent_framework_redis import RedisChatMessageStore, RedisProvider, __version__
+
+__all__ = ["RedisChatMessageStore", "RedisProvider", "__version__"]

python/packages/core/pyproject.toml

@@ -0,0 +1,115 @@
+[project]
+name = "agent-framework-core"
+description = "Microsoft Agent Framework for building AI Agents with Python. This is the core package that has all the core abstractions and implementations."
+authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
+readme = "README.md"
+requires-python = ">=3.10"
+version = "1.0.0-b251001"
+license-files = ["LICENSE"]
+urls.homepage = "https://learn.microsoft.com/en-us/semantic-kernel/overview/"
+urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
+urls.release_notes = "https://github.com/microsoft/agent-framework/releases?q=tag%3Apython-1&expanded=true"
+urls.issues = "https://github.com/microsoft/agent-framework/issues"
+classifiers = [
+  "License :: OSI Approved :: MIT License",
+  "Development Status :: 5 - Production/Stable",
+  "Intended Audience :: Developers",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+  "Programming Language :: Python :: 3.13",
+  "Framework :: Pydantic :: 2",
+  "Typing :: Typed",
+]
+dependencies = [
+    "openai>=1.99.0",
+    "pydantic>=2,<3",
+    "pydantic-settings>=2,<3",
+    "typing-extensions",
+    "opentelemetry-api>=1.24",
+    "opentelemetry-sdk>=1.24",
+    "mcp[ws]>=1.13",
+    "azure-monitor-opentelemetry>=1.7.0",
+    "azure-monitor-opentelemetry-exporter>=1.0.0b41",
+    "opentelemetry-exporter-otlp-proto-grpc>=1.36.0",
+    "opentelemetry-semantic-conventions-ai>=0.4.13",
+    "aiofiles>=24.1.0",
+    "azure-identity>=1,<2",
+]
+
+[project.optional-dependencies]
+viz = [
+    "graphviz>=0.20.0"
+]
+
+[tool.uv]
+prerelease = "if-necessary-or-explicit"
+environments = [
+    "sys_platform == 'darwin'",
+    "sys_platform == 'linux'",
+    "sys_platform == 'win32'"
+]
+
+[tool.uv-dynamic-versioning]
+fallback-version = "0.0.0"
+
+[tool.pytest.ini_options]
+testpaths = [
+    'tests',
+    'packages/core/tests',
+    'packages/a2a/tests',
+    'packages/azure-ai/tests',
+    'packages/copilotstudio/tests',
+    'packages/mem0/tests',
+    'packages/runtime/tests'
+]
+addopts = "-ra -q -r fEX"
+asyncio_mode = "auto"
+asyncio_default_fixture_loop_scope = "function"
+filterwarnings = []
+timeout = 120
+
+[tool.coverage.run]
+omit = [
+    "**/__init__.py"
+]
+
+[tool.ruff]
+extend = "../../pyproject.toml"
+
+[tool.pyright]
+extend = "../../pyproject.toml"
+exclude = ['tests']
+
+[tool.mypy]
+plugins = ['pydantic.mypy']
+strict = true
+python_version = "3.10"
+ignore_missing_imports = true
+disallow_untyped_defs = true
+no_implicit_optional = true
+check_untyped_defs = true
+warn_return_any = true
+show_error_codes = true
+warn_unused_ignores = false
+disallow_incomplete_defs = true
+disallow_untyped_decorators = true
+
+[tool.bandit]
+targets = ["agent_framework"]
+exclude_dirs = ["tests"]
+
+[tool.poe]
+executor.type = "uv"
+include = "../../shared_tasks.toml"
+[tool.poe.tasks]
+mypy = "mypy --config-file $POE_ROOT/pyproject.toml agent_framework"
+test = "pytest --cov=agent_framework --cov-report=term-missing:skip-covered tests"
+
+[tool.flit.module]
+name = "agent_framework"
+
+[build-system]
+requires = ["flit-core >= 3.11,<4.0"]
+build-backend = "flit_core.buildapi"

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

@@ -0,0 +1,62 @@
+# Copyright (c) Microsoft. All rights reserved.
+from typing import Any
+
+from pytest import fixture
+
+from agent_framework import ChatMessage
+
+
+# region: Connector Settings fixtures
+@fixture
+def exclude_list(request: Any) -> list[str]:
+    """Fixture that returns a list of environment variables to exclude."""
+    return request.param if hasattr(request, "param") else []
+
+
+@fixture
+def override_env_param_dict(request: Any) -> dict[str, str]:
+    """Fixture that returns a dict of environment variables to override."""
+    return request.param if hasattr(request, "param") else {}
+
+
+# These two fixtures are used for multiple things, also non-connector tests
+@fixture()
+def azure_openai_unit_test_env(monkeypatch, exclude_list, override_env_param_dict):  # type: ignore
+    """Fixture to set environment variables for AzureOpenAISettings."""
+
+    if exclude_list is None:
+        exclude_list = []
+
+    if override_env_param_dict is None:
+        override_env_param_dict = {}
+
+    env_vars = {
+        "AZURE_OPENAI_ENDPOINT": "https://test-endpoint.com",
+        "AZURE_OPENAI_CHAT_DEPLOYMENT_NAME": "test_chat_deployment",
+        "AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME": "test_chat_deployment",
+        "AZURE_OPENAI_TEXT_DEPLOYMENT_NAME": "test_text_deployment",
+        "AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME": "test_embedding_deployment",
+        "AZURE_OPENAI_TEXT_TO_IMAGE_DEPLOYMENT_NAME": "test_text_to_image_deployment",
+        "AZURE_OPENAI_AUDIO_TO_TEXT_DEPLOYMENT_NAME": "test_audio_to_text_deployment",
+        "AZURE_OPENAI_TEXT_TO_AUDIO_DEPLOYMENT_NAME": "test_text_to_audio_deployment",
+        "AZURE_OPENAI_REALTIME_DEPLOYMENT_NAME": "test_realtime_deployment",
+        "AZURE_OPENAI_API_KEY": "test_api_key",
+        "AZURE_OPENAI_API_VERSION": "2023-03-15-preview",
+        "AZURE_OPENAI_BASE_URL": "https://test_text_deployment.test-base-url.com",
+        "AZURE_OPENAI_TOKEN_ENDPOINT": "https://test-token-endpoint.com",
+    }
+
+    env_vars.update(override_env_param_dict)  # type: ignore
+
+    for key, value in env_vars.items():
+        if key in exclude_list:
+            monkeypatch.delenv(key, raising=False)  # type: ignore
+            continue
+        monkeypatch.setenv(key, value)  # type: ignore
+
+    return env_vars
+
+
+@fixture(scope="function")
+def chat_history() -> list[ChatMessage]:
+    return []

python/packages/core/tests/azure/test_azure_assistants_client.py

@@ -0,0 +1,723 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import os
+from typing import Annotated
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from azure.identity import AzureCliCredential
+from pydantic import Field
+
+from agent_framework import (
+    AgentRunResponse,
+    AgentRunResponseUpdate,
+    AgentThread,
+    ChatAgent,
+    ChatClientProtocol,
+    ChatMessage,
+    ChatResponse,
+    ChatResponseUpdate,
+    HostedCodeInterpreterTool,
+    TextContent,
+)
+from agent_framework.azure import AzureOpenAIAssistantsClient
+from agent_framework.exceptions import ServiceInitializationError
+
+skip_if_azure_integration_tests_disabled = pytest.mark.skipif(
+    os.getenv("RUN_INTEGRATION_TESTS", "false").lower() != "true"
+    or os.getenv("AZURE_OPENAI_ENDPOINT", "") in ("", "https://test-endpoint.com"),
+    reason="No real AZURE_OPENAI_ENDPOINT provided; skipping integration tests."
+    if os.getenv("RUN_INTEGRATION_TESTS", "false").lower() == "true"
+    else "Integration tests are disabled.",
+)
+
+
+def create_test_azure_assistants_client(
+    mock_async_azure_openai: MagicMock,
+    deployment_name: str | None = None,
+    assistant_id: str | None = None,
+    assistant_name: str | None = None,
+    thread_id: str | None = None,
+    should_delete_assistant: bool = False,
+) -> AzureOpenAIAssistantsClient:
+    """Helper function to create AzureOpenAIAssistantsClient instances for testing, bypassing Pydantic validation."""
+    return AzureOpenAIAssistantsClient.model_construct(
+        ai_model_id=deployment_name or "test_chat_deployment",
+        assistant_id=assistant_id,
+        assistant_name=assistant_name,
+        thread_id=thread_id,
+        api_key="test-api-key",
+        endpoint="https://test-endpoint.com",
+        client=mock_async_azure_openai,
+        _should_delete_assistant=should_delete_assistant,
+    )
+
+
+@pytest.fixture
+def mock_async_azure_openai() -> MagicMock:
+    """Mock AsyncAzureOpenAI client."""
+    mock_client = MagicMock()
+
+    # Mock beta.assistants
+    mock_client.beta.assistants.create = AsyncMock(return_value=MagicMock(id="test-assistant-id"))
+    mock_client.beta.assistants.delete = AsyncMock()
+
+    # Mock beta.threads
+    mock_client.beta.threads.create = AsyncMock(return_value=MagicMock(id="test-thread-id"))
+    mock_client.beta.threads.delete = AsyncMock()
+
+    # Mock beta.threads.runs
+    mock_client.beta.threads.runs.create = AsyncMock(return_value=MagicMock(id="test-run-id"))
+    mock_client.beta.threads.runs.retrieve = AsyncMock()
+    mock_client.beta.threads.runs.submit_tool_outputs = AsyncMock()
+
+    # Mock beta.threads.messages
+    mock_client.beta.threads.messages.create = AsyncMock()
+    mock_client.beta.threads.messages.list = AsyncMock(return_value=MagicMock(data=[]))
+
+    return mock_client
+
+
+def test_azure_assistants_client_init_with_client(mock_async_azure_openai: MagicMock) -> None:
+    """Test AzureOpenAIAssistantsClient initialization with existing client."""
+    chat_client = create_test_azure_assistants_client(
+        mock_async_azure_openai,
+        deployment_name="test_chat_deployment",
+        assistant_id="existing-assistant-id",
+        thread_id="test-thread-id",
+    )
+
+    assert chat_client.client is mock_async_azure_openai
+    assert chat_client.ai_model_id == "test_chat_deployment"
+    assert chat_client.assistant_id == "existing-assistant-id"
+    assert chat_client.thread_id == "test-thread-id"
+    assert not chat_client._should_delete_assistant  # type: ignore
+    assert isinstance(chat_client, ChatClientProtocol)
+
+
+def test_azure_assistants_client_init_auto_create_client(
+    azure_openai_unit_test_env: dict[str, str],
+    mock_async_azure_openai: MagicMock,
+) -> None:
+    """Test AzureOpenAIAssistantsClient initialization with auto-created client."""
+    chat_client = AzureOpenAIAssistantsClient.model_construct(
+        ai_model_id=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
+        assistant_id=None,
+        assistant_name="TestAssistant",
+        thread_id=None,
+        api_key=azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
+        endpoint=azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"],
+        client=mock_async_azure_openai,
+        _should_delete_assistant=False,
+    )
+
+    assert chat_client.client is mock_async_azure_openai
+    assert chat_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
+    assert chat_client.assistant_id is None
+    assert chat_client.assistant_name == "TestAssistant"
+    assert not chat_client._should_delete_assistant  # type: ignore
+
+
+def test_azure_assistants_client_init_validation_fail() -> None:
+    """Test AzureOpenAIAssistantsClient initialization with validation failure."""
+    with pytest.raises(ServiceInitializationError):
+        # Force failure by providing invalid deployment name type - this should cause validation to fail
+        AzureOpenAIAssistantsClient(deployment_name=123, api_key="valid-key")  # type: ignore
+
+
+@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]], indirect=True)
+def test_azure_assistants_client_init_missing_deployment_name(azure_openai_unit_test_env: dict[str, str]) -> None:
+    """Test AzureOpenAIAssistantsClient initialization with missing deployment name."""
+    with pytest.raises(ServiceInitializationError):
+        AzureOpenAIAssistantsClient(
+            api_key=azure_openai_unit_test_env.get("AZURE_OPENAI_API_KEY", "test-key"), env_file_path="nonexistent.env"
+        )
+
+
+def test_azure_assistants_client_init_with_default_headers(azure_openai_unit_test_env: dict[str, str]) -> None:
+    """Test AzureOpenAIAssistantsClient initialization with default headers."""
+    default_headers = {"X-Unit-Test": "test-guid"}
+
+    chat_client = AzureOpenAIAssistantsClient(
+        deployment_name="test_chat_deployment",
+        api_key=azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
+        endpoint=azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"],
+        default_headers=default_headers,
+    )
+
+    assert chat_client.ai_model_id == "test_chat_deployment"
+    assert isinstance(chat_client, ChatClientProtocol)
+
+    # Assert that the default header we added is present in the client's default headers
+    for key, value in default_headers.items():
+        assert key in chat_client.client.default_headers
+        assert chat_client.client.default_headers[key] == value
+
+
+async def test_azure_assistants_client_get_assistant_id_or_create_existing_assistant(
+    mock_async_azure_openai: MagicMock,
+) -> None:
+    """Test _get_assistant_id_or_create when assistant_id is already provided."""
+    chat_client = create_test_azure_assistants_client(mock_async_azure_openai, assistant_id="existing-assistant-id")
+
+    assistant_id = await chat_client._get_assistant_id_or_create()  # type: ignore
+
+    assert assistant_id == "existing-assistant-id"
+    assert not chat_client._should_delete_assistant  # type: ignore
+    mock_async_azure_openai.beta.assistants.create.assert_not_called()
+
+
+async def test_azure_assistants_client_get_assistant_id_or_create_create_new(
+    mock_async_azure_openai: MagicMock,
+) -> None:
+    """Test _get_assistant_id_or_create when creating a new assistant."""
+    chat_client = create_test_azure_assistants_client(
+        mock_async_azure_openai, deployment_name="test_chat_deployment", assistant_name="TestAssistant"
+    )
+
+    assistant_id = await chat_client._get_assistant_id_or_create()  # type: ignore
+
+    assert assistant_id == "test-assistant-id"
+    assert chat_client._should_delete_assistant  # type: ignore
+    mock_async_azure_openai.beta.assistants.create.assert_called_once()
+
+
+async def test_azure_assistants_client_aclose_should_not_delete(
+    mock_async_azure_openai: MagicMock,
+) -> None:
+    """Test close when assistant should not be deleted."""
+    chat_client = create_test_azure_assistants_client(
+        mock_async_azure_openai, assistant_id="assistant-to-keep", should_delete_assistant=False
+    )
+
+    await chat_client.close()  # type: ignore
+
+    # Verify assistant deletion was not called
+    mock_async_azure_openai.beta.assistants.delete.assert_not_called()
+    assert not chat_client._should_delete_assistant  # type: ignore
+
+
+async def test_azure_assistants_client_aclose_should_delete(mock_async_azure_openai: MagicMock) -> None:
+    """Test close method calls cleanup."""
+    chat_client = create_test_azure_assistants_client(
+        mock_async_azure_openai, assistant_id="assistant-to-delete", should_delete_assistant=True
+    )
+
+    await chat_client.close()
+
+    # Verify assistant deletion was called
+    mock_async_azure_openai.beta.assistants.delete.assert_called_once_with("assistant-to-delete")
+    assert not chat_client._should_delete_assistant  # type: ignore
+
+
+async def test_azure_assistants_client_async_context_manager(mock_async_azure_openai: MagicMock) -> None:
+    """Test async context manager functionality."""
+    chat_client = create_test_azure_assistants_client(
+        mock_async_azure_openai, assistant_id="assistant-to-delete", should_delete_assistant=True
+    )
+
+    # Test context manager
+    async with chat_client:
+        pass  # Just test that we can enter and exit
+
+    # Verify cleanup was called on exit
+    mock_async_azure_openai.beta.assistants.delete.assert_called_once_with("assistant-to-delete")
+
+
+def test_azure_assistants_client_serialize(azure_openai_unit_test_env: dict[str, str]) -> None:
+    """Test serialization of AzureOpenAIAssistantsClient."""
+    default_headers = {"X-Unit-Test": "test-guid"}
+
+    # Test basic initialization and to_dict
+    chat_client = AzureOpenAIAssistantsClient(
+        deployment_name="test_chat_deployment",
+        assistant_id="test-assistant-id",
+        assistant_name="TestAssistant",
+        thread_id="test-thread-id",
+        api_key=azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
+        endpoint=azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"],
+        default_headers=default_headers,
+    )
+
+    dumped_settings = chat_client.to_dict()
+
+    assert dumped_settings["ai_model_id"] == "test_chat_deployment"
+    assert dumped_settings["assistant_id"] == "test-assistant-id"
+    assert dumped_settings["assistant_name"] == "TestAssistant"
+    assert dumped_settings["thread_id"] == "test-thread-id"
+    assert dumped_settings["api_key"] == azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"]
+
+    # Assert that the default header we added is present in the dumped_settings default headers
+    for key, value in default_headers.items():
+        assert key in dumped_settings["default_headers"]
+        assert dumped_settings["default_headers"][key] == value
+    # Assert that the 'User-Agent' header is not present in the dumped_settings default headers
+    assert "User-Agent" not in dumped_settings["default_headers"]
+
+
+def get_weather(
+    location: Annotated[str, Field(description="The location to get the weather for.")],
+) -> str:
+    """Get the weather for a given location."""
+    return f"The weather in {location} is sunny with a high of 25°C."
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_assistants_client_get_response() -> None:
+    """Test Azure Assistants Client response."""
+    async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as azure_assistants_client:
+        assert isinstance(azure_assistants_client, ChatClientProtocol)
+
+        messages: list[ChatMessage] = []
+        messages.append(
+            ChatMessage(
+                role="user",
+                text="The weather in Seattle is currently sunny with a high of 25°C. "
+                "It's a beautiful day for outdoor activities.",
+            )
+        )
+        messages.append(ChatMessage(role="user", text="What's the weather like today?"))
+
+        # Test that the client can be used to get a response
+        response = await azure_assistants_client.get_response(messages=messages)
+
+        assert response is not None
+        assert isinstance(response, ChatResponse)
+        assert any(word in response.text.lower() for word in ["sunny", "25", "weather", "seattle"])
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_assistants_client_get_response_tools() -> None:
+    """Test Azure Assistants Client response with tools."""
+    async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as azure_assistants_client:
+        assert isinstance(azure_assistants_client, ChatClientProtocol)
+
+        messages: list[ChatMessage] = []
+        messages.append(ChatMessage(role="user", text="What's the weather like in Seattle?"))
+
+        # Test that the client can be used to get a response
+        response = await azure_assistants_client.get_response(
+            messages=messages,
+            tools=[get_weather],
+            tool_choice="auto",
+        )
+
+        assert response is not None
+        assert isinstance(response, ChatResponse)
+        assert any(word in response.text.lower() for word in ["sunny", "25", "weather"])
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_assistants_client_streaming() -> None:
+    """Test Azure Assistants Client streaming response."""
+    async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as azure_assistants_client:
+        assert isinstance(azure_assistants_client, ChatClientProtocol)
+
+        messages: list[ChatMessage] = []
+        messages.append(
+            ChatMessage(
+                role="user",
+                text="The weather in Seattle is currently sunny with a high of 25°C. "
+                "It's a beautiful day for outdoor activities.",
+            )
+        )
+        messages.append(ChatMessage(role="user", text="What's the weather like today?"))
+
+        # Test that the client can be used to get a response
+        response = azure_assistants_client.get_streaming_response(messages=messages)
+
+        full_message: str = ""
+        async for chunk in response:
+            assert chunk is not None
+            assert isinstance(chunk, ChatResponseUpdate)
+            for content in chunk.contents:
+                if isinstance(content, TextContent) and content.text:
+                    full_message += content.text
+
+        assert any(word in full_message.lower() for word in ["sunny", "25", "weather", "seattle"])
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_assistants_client_streaming_tools() -> None:
+    """Test Azure Assistants Client streaming response with tools."""
+    async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as azure_assistants_client:
+        assert isinstance(azure_assistants_client, ChatClientProtocol)
+
+        messages: list[ChatMessage] = []
+        messages.append(ChatMessage(role="user", text="What's the weather like in Seattle?"))
+
+        # Test that the client can be used to get a response
+        response = azure_assistants_client.get_streaming_response(
+            messages=messages,
+            tools=[get_weather],
+            tool_choice="auto",
+        )
+        full_message: str = ""
+        async for chunk in response:
+            assert chunk is not None
+            assert isinstance(chunk, ChatResponseUpdate)
+            for content in chunk.contents:
+                if isinstance(content, TextContent) and content.text:
+                    full_message += content.text
+
+        assert any(word in full_message.lower() for word in ["sunny", "25", "weather"])
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_assistants_client_with_existing_assistant() -> None:
+    """Test Azure Assistants Client with existing assistant ID."""
+    # First create an assistant to use in the test
+    async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as temp_client:
+        # Get the assistant ID by triggering assistant creation
+        messages = [ChatMessage(role="user", text="Hello")]
+        await temp_client.get_response(messages=messages)
+        assistant_id = temp_client.assistant_id
+
+        # Now test using the existing assistant
+        async with AzureOpenAIAssistantsClient(
+            assistant_id=assistant_id, credential=AzureCliCredential()
+        ) as azure_assistants_client:
+            assert isinstance(azure_assistants_client, ChatClientProtocol)
+            assert azure_assistants_client.assistant_id == assistant_id
+
+            messages = [ChatMessage(role="user", text="What can you do?")]
+
+            # Test that the client can be used to get a response
+            response = await azure_assistants_client.get_response(messages=messages)
+
+            assert response is not None
+            assert isinstance(response, ChatResponse)
+            assert len(response.text) > 0
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_assistants_agent_basic_run():
+    """Test ChatAgent basic run functionality with AzureOpenAIAssistantsClient."""
+    async with ChatAgent(
+        chat_client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
+    ) as agent:
+        # Run a simple query
+        response = await agent.run("Hello! Please respond with 'Hello World' exactly.")
+
+        # Validate response
+        assert isinstance(response, AgentRunResponse)
+        assert response.text is not None
+        assert len(response.text) > 0
+        assert "Hello World" in response.text
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_assistants_agent_basic_run_streaming():
+    """Test ChatAgent basic streaming functionality with AzureOpenAIAssistantsClient."""
+    async with ChatAgent(
+        chat_client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
+    ) as agent:
+        # Run streaming query
+        full_message: str = ""
+        async for chunk in agent.run_stream("Please respond with exactly: 'This is a streaming response test.'"):
+            assert chunk is not None
+            assert isinstance(chunk, AgentRunResponseUpdate)
+            if chunk.text:
+                full_message += chunk.text
+
+        # Validate streaming response
+        assert len(full_message) > 0
+        assert "streaming response test" in full_message.lower()
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_assistants_agent_thread_persistence():
+    """Test ChatAgent thread persistence across runs with AzureOpenAIAssistantsClient."""
+    async with ChatAgent(
+        chat_client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
+        instructions="You are a helpful assistant with good memory.",
+    ) as agent:
+        # Create a new thread that will be reused
+        thread = agent.get_new_thread()
+
+        # First message - establish context
+        first_response = await agent.run(
+            "Remember this number: 42. What number did I just tell you to remember?", thread=thread
+        )
+        assert isinstance(first_response, AgentRunResponse)
+        assert "42" in first_response.text
+
+        # Second message - test conversation memory
+        second_response = await agent.run(
+            "What number did I tell you to remember in my previous message?", thread=thread
+        )
+        assert isinstance(second_response, AgentRunResponse)
+        assert "42" in second_response.text
+
+        # Verify thread has been populated with conversation ID
+        assert thread.service_thread_id is not None
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_assistants_agent_existing_thread_id():
+    """Test ChatAgent with existing thread ID to continue conversations across agent instances."""
+    # First, create a conversation and capture the thread ID
+    existing_thread_id = None
+
+    async with ChatAgent(
+        chat_client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
+        instructions="You are a helpful weather agent.",
+        tools=[get_weather],
+    ) as agent:
+        # Start a conversation and get the thread ID
+        thread = agent.get_new_thread()
+        response1 = await agent.run("What's the weather in Paris?", thread=thread)
+
+        # Validate first response
+        assert isinstance(response1, AgentRunResponse)
+        assert response1.text is not None
+        assert any(word in response1.text.lower() for word in ["weather", "paris"])
+
+        # The thread ID is set after the first response
+        existing_thread_id = thread.service_thread_id
+        assert existing_thread_id is not None
+
+    # Now continue with the same thread ID in a new agent instance
+
+    async with ChatAgent(
+        chat_client=AzureOpenAIAssistantsClient(thread_id=existing_thread_id, credential=AzureCliCredential()),
+        instructions="You are a helpful weather agent.",
+        tools=[get_weather],
+    ) as agent:
+        # Create a thread with the existing ID
+        thread = AgentThread(service_thread_id=existing_thread_id)
+
+        # Ask about the previous conversation
+        response2 = await agent.run("What was the last city I asked about?", thread=thread)
+
+        # Validate that the agent remembers the previous conversation
+        assert isinstance(response2, AgentRunResponse)
+        assert response2.text is not None
+        # Should reference Paris from the previous conversation
+        assert "paris" in response2.text.lower()
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_assistants_agent_code_interpreter():
+    """Test ChatAgent with code interpreter through AzureOpenAIAssistantsClient."""
+
+    async with ChatAgent(
+        chat_client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
+        instructions="You are a helpful assistant that can write and execute Python code.",
+        tools=[HostedCodeInterpreterTool()],
+    ) as agent:
+        # Request code execution
+        response = await agent.run("Write Python code to calculate the factorial of 5 and show the result.")
+
+        # Validate response
+        assert isinstance(response, AgentRunResponse)
+        assert response.text is not None
+        # Factorial of 5 is 120
+        assert "120" in response.text or "factorial" in response.text.lower()
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_assistants_client_agent_level_tool_persistence():
+    """Test that agent-level tools persist across multiple runs with Azure Assistants Client."""
+
+    async with ChatAgent(
+        chat_client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
+        instructions="You are a helpful assistant that uses available tools.",
+        tools=[get_weather],  # Agent-level tool
+    ) as agent:
+        # First run - agent-level tool should be available
+        first_response = await agent.run("What's the weather like in Chicago?")
+
+        assert isinstance(first_response, AgentRunResponse)
+        assert first_response.text is not None
+        # Should use the agent-level weather tool
+        assert any(term in first_response.text.lower() for term in ["chicago", "sunny", "72"])
+
+        # Second run - agent-level tool should still be available (persistence test)
+        second_response = await agent.run("What's the weather in Miami?")
+
+        assert isinstance(second_response, AgentRunResponse)
+        assert second_response.text is not None
+        # Should use the agent-level weather tool again
+        assert any(term in second_response.text.lower() for term in ["miami", "sunny", "72"])
+
+
+def test_azure_assistants_client_entra_id_authentication() -> None:
+    """Test Entra ID authentication path with credential."""
+    mock_credential = MagicMock()
+
+    with (
+        patch("agent_framework.azure._assistants_client.AzureOpenAISettings") as mock_settings_class,
+        patch("agent_framework.azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client,
+        patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
+    ):
+        mock_settings = MagicMock()
+        mock_settings.chat_deployment_name = "test-deployment"
+        mock_settings.api_key = None  # No API key to trigger Entra ID path
+        mock_settings.token_endpoint = "https://login.microsoftonline.com/test"
+        mock_settings.get_azure_auth_token.return_value = "entra-token-12345"
+        mock_settings.api_version = "2024-05-01-preview"
+        mock_settings.endpoint = "https://test-endpoint.openai.azure.com"
+        mock_settings.base_url = None
+        mock_settings_class.return_value = mock_settings
+
+        client = AzureOpenAIAssistantsClient(
+            deployment_name="test-deployment",
+            api_key="placeholder-key",
+            endpoint="https://test-endpoint.openai.azure.com",
+            credential=mock_credential,
+            token_endpoint="https://login.microsoftonline.com/test",
+        )
+
+        # Verify Entra ID token was requested
+        mock_settings.get_azure_auth_token.assert_called_once_with(mock_credential)
+
+        # Verify client was created with the token
+        mock_azure_client.assert_called_once()
+        call_args = mock_azure_client.call_args[1]
+        assert call_args["azure_ad_token"] == "entra-token-12345"
+
+        assert client is not None
+        assert isinstance(client, AzureOpenAIAssistantsClient)
+
+
+def test_azure_assistants_client_no_authentication_error() -> None:
+    """Test authentication validation error when no auth provided."""
+    with patch("agent_framework.azure._assistants_client.AzureOpenAISettings") as mock_settings_class:
+        mock_settings = MagicMock()
+        mock_settings.chat_deployment_name = "test-deployment"
+        mock_settings.api_key = None  # No API key
+        mock_settings.token_endpoint = None  # No token endpoint
+        mock_settings_class.return_value = mock_settings
+
+        # Test missing authentication raises error
+        with pytest.raises(ServiceInitializationError, match="API key, ad_token, or ad_token_provider is required"):
+            AzureOpenAIAssistantsClient(
+                deployment_name="test-deployment",
+                endpoint="https://test-endpoint.openai.azure.com",
+                # No authentication provided at all
+            )
+
+
+def test_azure_assistants_client_ad_token_authentication() -> None:
+    """Test ad_token authentication client parameter path."""
+    with (
+        patch("agent_framework.azure._assistants_client.AzureOpenAISettings") as mock_settings_class,
+        patch("agent_framework.azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client,
+        patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
+    ):
+        mock_settings = MagicMock()
+        mock_settings.chat_deployment_name = "test-deployment"
+        mock_settings.api_key = None  # No API key
+        mock_settings.api_version = "2024-05-01-preview"
+        mock_settings.endpoint = "https://test-endpoint.openai.azure.com"
+        mock_settings.base_url = None
+        mock_settings_class.return_value = mock_settings
+
+        client = AzureOpenAIAssistantsClient(
+            deployment_name="test-deployment",
+            endpoint="https://test-endpoint.openai.azure.com",
+            ad_token="test-ad-token-12345",
+        )
+
+        # ad_token path
+        mock_azure_client.assert_called_once()
+        call_args = mock_azure_client.call_args[1]
+        assert call_args["azure_ad_token"] == "test-ad-token-12345"
+
+        assert client is not None
+        assert isinstance(client, AzureOpenAIAssistantsClient)
+
+
+def test_azure_assistants_client_ad_token_provider_authentication() -> None:
+    """Test ad_token_provider authentication client parameter path."""
+    from openai.lib.azure import AsyncAzureADTokenProvider
+
+    mock_token_provider = MagicMock(spec=AsyncAzureADTokenProvider)
+
+    with (
+        patch("agent_framework.azure._assistants_client.AzureOpenAISettings") as mock_settings_class,
+        patch("agent_framework.azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client,
+        patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
+    ):
+        mock_settings = MagicMock()
+        mock_settings.chat_deployment_name = "test-deployment"
+        mock_settings.api_key = None  # No API key
+        mock_settings.api_version = "2024-05-01-preview"
+        mock_settings.endpoint = "https://test-endpoint.openai.azure.com"
+        mock_settings.base_url = None
+        mock_settings_class.return_value = mock_settings
+
+        client = AzureOpenAIAssistantsClient(
+            deployment_name="test-deployment",
+            endpoint="https://test-endpoint.openai.azure.com",
+            ad_token_provider=mock_token_provider,
+        )
+
+        # ad_token_provider path
+        mock_azure_client.assert_called_once()
+        call_args = mock_azure_client.call_args[1]
+        assert call_args["azure_ad_token_provider"] is mock_token_provider
+
+        assert client is not None
+        assert isinstance(client, AzureOpenAIAssistantsClient)
+
+
+def test_azure_assistants_client_base_url_configuration() -> None:
+    """Test base_url client parameter path."""
+    with (
+        patch("agent_framework.azure._assistants_client.AzureOpenAISettings") as mock_settings_class,
+        patch("agent_framework.azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client,
+        patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
+    ):
+        mock_settings = MagicMock()
+        mock_settings.chat_deployment_name = "test-deployment"
+        mock_settings.api_key.get_secret_value.return_value = "test-api-key"
+        mock_settings.base_url = "https://custom-base-url.com"
+        mock_settings.endpoint = None  # No endpoint, should use base_url
+        mock_settings.api_version = "2024-05-01-preview"
+        mock_settings_class.return_value = mock_settings
+
+        client = AzureOpenAIAssistantsClient(
+            deployment_name="test-deployment", api_key="test-api-key", base_url="https://custom-base-url.com"
+        )
+
+        # base_url path
+        mock_azure_client.assert_called_once()
+        call_args = mock_azure_client.call_args[1]
+        assert call_args["base_url"] == "https://custom-base-url.com"
+        assert "azure_endpoint" not in call_args
+
+        assert client is not None
+        assert isinstance(client, AzureOpenAIAssistantsClient)
+
+
+def test_azure_assistants_client_azure_endpoint_configuration() -> None:
+    """Test azure_endpoint client parameter path."""
+    with (
+        patch("agent_framework.azure._assistants_client.AzureOpenAISettings") as mock_settings_class,
+        patch("agent_framework.azure._assistants_client.AsyncAzureOpenAI") as mock_azure_client,
+        patch("agent_framework.openai.OpenAIAssistantsClient.__init__", return_value=None),
+    ):
+        mock_settings = MagicMock()
+        mock_settings.chat_deployment_name = "test-deployment"
+        mock_settings.api_key.get_secret_value.return_value = "test-api-key"
+        mock_settings.base_url = None  # No base_url
+        mock_settings.endpoint = "https://test-endpoint.openai.azure.com"
+        mock_settings.api_version = "2024-05-01-preview"
+        mock_settings_class.return_value = mock_settings
+
+        client = AzureOpenAIAssistantsClient(
+            deployment_name="test-deployment",
+            api_key="test-api-key",
+            endpoint="https://test-endpoint.openai.azure.com",
+        )
+
+        # azure_endpoint path
+        mock_azure_client.assert_called_once()
+        call_args = mock_azure_client.call_args[1]
+        assert call_args["azure_endpoint"] == "https://test-endpoint.openai.azure.com"
+        assert "base_url" not in call_args
+
+        assert client is not None
+        assert isinstance(client, AzureOpenAIAssistantsClient)

python/packages/core/tests/azure/test_azure_chat_client.py

@@ -0,0 +1,835 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import json
+import os
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import openai
+import pytest
+from azure.identity import AzureCliCredential
+from httpx import Request, Response
+from openai import AsyncAzureOpenAI, AsyncStream
+from openai.resources.chat.completions import AsyncCompletions as AsyncChatCompletions
+from openai.types.chat import ChatCompletion, ChatCompletionChunk
+from openai.types.chat.chat_completion import Choice
+from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
+from openai.types.chat.chat_completion_chunk import ChoiceDelta as ChunkChoiceDelta
+from openai.types.chat.chat_completion_message import ChatCompletionMessage
+
+from agent_framework import (
+    AgentRunResponse,
+    AgentRunResponseUpdate,
+    BaseChatClient,
+    ChatAgent,
+    ChatClientProtocol,
+    ChatMessage,
+    ChatResponse,
+    ChatResponseUpdate,
+    TextContent,
+    ai_function,
+)
+from agent_framework._telemetry import USER_AGENT_KEY
+from agent_framework.azure import AzureOpenAIChatClient
+from agent_framework.exceptions import ServiceInitializationError, ServiceResponseException
+from agent_framework.openai import (
+    ContentFilterResultSeverity,
+    OpenAIContentFilterException,
+)
+
+# region Service Setup
+
+skip_if_azure_integration_tests_disabled = pytest.mark.skipif(
+    os.getenv("RUN_INTEGRATION_TESTS", "false").lower() != "true"
+    or os.getenv("AZURE_OPENAI_ENDPOINT", "") in ("", "https://test-endpoint.com"),
+    reason="No real AZURE_OPENAI_ENDPOINT provided; skipping integration tests."
+    if os.getenv("RUN_INTEGRATION_TESTS", "false").lower() == "true"
+    else "Integration tests are disabled.",
+)
+
+
+def test_init(azure_openai_unit_test_env: dict[str, str]) -> None:
+    # Test successful initialization
+    azure_chat_client = AzureOpenAIChatClient()
+
+    assert azure_chat_client.client is not None
+    assert isinstance(azure_chat_client.client, AsyncAzureOpenAI)
+    assert azure_chat_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
+    assert isinstance(azure_chat_client, BaseChatClient)
+
+
+def test_init_client(azure_openai_unit_test_env: dict[str, str]) -> None:
+    # Test successful initialization with client
+    client = MagicMock(spec=AsyncAzureOpenAI)
+    azure_chat_client = AzureOpenAIChatClient(async_client=client)
+
+    assert azure_chat_client.client is not None
+    assert isinstance(azure_chat_client.client, AsyncAzureOpenAI)
+
+
+def test_init_base_url(azure_openai_unit_test_env: dict[str, str]) -> None:
+    # Custom header for testing
+    default_headers = {"X-Unit-Test": "test-guid"}
+
+    azure_chat_client = AzureOpenAIChatClient(
+        default_headers=default_headers,
+    )
+
+    assert azure_chat_client.client is not None
+    assert isinstance(azure_chat_client.client, AsyncAzureOpenAI)
+    assert azure_chat_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
+    assert isinstance(azure_chat_client, BaseChatClient)
+    for key, value in default_headers.items():
+        assert key in azure_chat_client.client.default_headers
+        assert azure_chat_client.client.default_headers[key] == value
+
+
+@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_BASE_URL"]], indirect=True)
+def test_init_endpoint(azure_openai_unit_test_env: dict[str, str]) -> None:
+    azure_chat_client = AzureOpenAIChatClient()
+
+    assert azure_chat_client.client is not None
+    assert isinstance(azure_chat_client.client, AsyncAzureOpenAI)
+    assert azure_chat_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
+    assert isinstance(azure_chat_client, BaseChatClient)
+
+
+@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]], indirect=True)
+def test_init_with_empty_deployment_name(azure_openai_unit_test_env: dict[str, str]) -> None:
+    with pytest.raises(ServiceInitializationError):
+        AzureOpenAIChatClient(
+            env_file_path="test.env",
+        )
+
+
+@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_ENDPOINT", "AZURE_OPENAI_BASE_URL"]], indirect=True)
+def test_init_with_empty_endpoint_and_base_url(azure_openai_unit_test_env: dict[str, str]) -> None:
+    with pytest.raises(ServiceInitializationError):
+        AzureOpenAIChatClient(
+            env_file_path="test.env",
+        )
+
+
+@pytest.mark.parametrize("override_env_param_dict", [{"AZURE_OPENAI_ENDPOINT": "http://test.com"}], indirect=True)
+def test_init_with_invalid_endpoint(azure_openai_unit_test_env: dict[str, str]) -> None:
+    with pytest.raises(ServiceInitializationError):
+        AzureOpenAIChatClient()
+
+
+@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_BASE_URL"]], indirect=True)
+def test_serialize(azure_openai_unit_test_env: dict[str, str]) -> None:
+    default_headers = {"X-Test": "test"}
+
+    settings = {
+        "deployment_name": azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
+        "endpoint": azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"],
+        "api_key": azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
+        "api_version": azure_openai_unit_test_env["AZURE_OPENAI_API_VERSION"],
+        "default_headers": default_headers,
+        "env_file_path": "test.env",
+    }
+
+    azure_chat_client = AzureOpenAIChatClient.from_dict(settings)
+    dumped_settings = azure_chat_client.to_dict()
+    assert dumped_settings["ai_model_id"] == settings["deployment_name"]
+    assert str(settings["endpoint"]) in str(dumped_settings["base_url"])
+    assert str(settings["deployment_name"]) in str(dumped_settings["base_url"])
+    assert settings["api_key"] == dumped_settings["api_key"]
+    assert settings["api_version"] == dumped_settings["api_version"]
+
+    # Assert that the default header we added is present in the dumped_settings default headers
+    for key, value in default_headers.items():
+        assert key in dumped_settings["default_headers"]
+        assert dumped_settings["default_headers"][key] == value
+
+    # Assert that the 'User-agent' header is not present in the dumped_settings default headers
+    assert USER_AGENT_KEY not in dumped_settings["default_headers"]
+
+
+# endregion
+# region CMC
+
+
+@pytest.fixture
+def mock_chat_completion_response() -> ChatCompletion:
+    return ChatCompletion(
+        id="test_id",
+        choices=[
+            Choice(index=0, message=ChatCompletionMessage(content="test", role="assistant"), finish_reason="stop")
+        ],
+        created=0,
+        model="test",
+        object="chat.completion",
+    )
+
+
+@pytest.fixture
+def mock_streaming_chat_completion_response() -> AsyncStream[ChatCompletionChunk]:
+    content = ChatCompletionChunk(
+        id="test_id",
+        choices=[ChunkChoice(index=0, delta=ChunkChoiceDelta(content="test", role="assistant"), finish_reason="stop")],
+        created=0,
+        model="test",
+        object="chat.completion.chunk",
+    )
+    stream = MagicMock(spec=AsyncStream)
+    stream.__aiter__.return_value = [content]
+    return stream
+
+
+@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
+async def test_cmc(
+    mock_create: AsyncMock,
+    azure_openai_unit_test_env: dict[str, str],
+    chat_history: list[ChatMessage],
+    mock_chat_completion_response: ChatCompletion,
+) -> None:
+    mock_create.return_value = mock_chat_completion_response
+    chat_history.append(ChatMessage(text="hello world", role="user"))
+
+    azure_chat_client = AzureOpenAIChatClient()
+    await azure_chat_client.get_response(
+        messages=chat_history,
+    )
+    mock_create.assert_awaited_once_with(
+        model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
+        stream=False,
+        messages=azure_chat_client._prepare_chat_history_for_request(chat_history),  # type: ignore
+    )
+
+
+@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
+async def test_cmc_with_logit_bias(
+    mock_create: AsyncMock,
+    azure_openai_unit_test_env: dict[str, str],
+    chat_history: list[ChatMessage],
+    mock_chat_completion_response: ChatCompletion,
+) -> None:
+    mock_create.return_value = mock_chat_completion_response
+    prompt = "hello world"
+    chat_history.append(ChatMessage(text=prompt, role="user"))
+
+    token_bias: dict[str | int, float] = {"1": -100}
+
+    azure_chat_client = AzureOpenAIChatClient()
+
+    await azure_chat_client.get_response(messages=chat_history, logit_bias=token_bias)
+
+    mock_create.assert_awaited_once_with(
+        model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
+        messages=azure_chat_client._prepare_chat_history_for_request(chat_history),  # type: ignore
+        stream=False,
+        logit_bias=token_bias,
+    )
+
+
+@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
+async def test_cmc_with_stop(
+    mock_create: AsyncMock,
+    azure_openai_unit_test_env: dict[str, str],
+    chat_history: list[ChatMessage],
+    mock_chat_completion_response: ChatCompletion,
+) -> None:
+    mock_create.return_value = mock_chat_completion_response
+    prompt = "hello world"
+    chat_history.append(ChatMessage(text=prompt, role="user"))
+
+    stop = ["!"]
+
+    azure_chat_client = AzureOpenAIChatClient()
+
+    await azure_chat_client.get_response(messages=chat_history, stop=stop)
+
+    mock_create.assert_awaited_once_with(
+        model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
+        messages=azure_chat_client._prepare_chat_history_for_request(chat_history),  # type: ignore
+        stream=False,
+        stop=stop,
+    )
+
+
+@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
+async def test_azure_on_your_data(
+    mock_create: AsyncMock,
+    azure_openai_unit_test_env: dict[str, str],
+    chat_history: list[ChatMessage],
+    mock_chat_completion_response: ChatCompletion,
+) -> None:
+    mock_chat_completion_response.choices = [
+        Choice(
+            index=0,
+            message=ChatCompletionMessage(
+                content="test",
+                role="assistant",
+                context={  # type: ignore
+                    "citations": [
+                        {
+                            "content": "test content",
+                            "title": "test title",
+                            "url": "test url",
+                            "filepath": "test filepath",
+                            "chunk_id": "test chunk_id",
+                        }
+                    ],
+                    "intent": "query used",
+                },
+            ),
+            finish_reason="stop",
+        )
+    ]
+    mock_create.return_value = mock_chat_completion_response
+    prompt = "hello world"
+    messages_in = chat_history
+    chat_history.append(ChatMessage(text=prompt, role="user"))
+    messages_out: list[ChatMessage] = []
+    messages_out.append(ChatMessage(text=prompt, role="user"))
+
+    expected_data_settings = {
+        "data_sources": [
+            {
+                "type": "AzureCognitiveSearch",
+                "parameters": {
+                    "indexName": "test_index",
+                    "endpoint": "https://test-endpoint-search.com",
+                    "key": "test_key",
+                },
+            }
+        ]
+    }
+
+    azure_chat_client = AzureOpenAIChatClient()
+
+    content = await azure_chat_client.get_response(
+        messages=messages_in,
+        additional_properties={"extra_body": expected_data_settings},
+    )
+    assert len(content.messages) == 1
+    assert len(content.messages[0].contents) == 1
+    assert isinstance(content.messages[0].contents[0], TextContent)
+    assert len(content.messages[0].contents[0].annotations) == 1
+    assert content.messages[0].contents[0].annotations[0].title == "test title"
+    assert content.messages[0].contents[0].text == "test"
+
+    mock_create.assert_awaited_once_with(
+        model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
+        messages=azure_chat_client._prepare_chat_history_for_request(messages_out),  # type: ignore
+        stream=False,
+        extra_body=expected_data_settings,
+    )
+
+
+@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
+async def test_azure_on_your_data_string(
+    mock_create: AsyncMock,
+    azure_openai_unit_test_env: dict[str, str],
+    chat_history: list[ChatMessage],
+    mock_chat_completion_response: ChatCompletion,
+) -> None:
+    mock_chat_completion_response.choices = [
+        Choice(
+            index=0,
+            message=ChatCompletionMessage(
+                content="test",
+                role="assistant",
+                context=json.dumps({  # type: ignore
+                    "citations": [
+                        {
+                            "content": "test content",
+                            "title": "test title",
+                            "url": "test url",
+                            "filepath": "test filepath",
+                            "chunk_id": "test chunk_id",
+                        }
+                    ],
+                    "intent": "query used",
+                }),
+            ),
+            finish_reason="stop",
+        )
+    ]
+    mock_create.return_value = mock_chat_completion_response
+    prompt = "hello world"
+    messages_in = chat_history
+    messages_in.append(ChatMessage(text=prompt, role="user"))
+    messages_out: list[ChatMessage] = []
+    messages_out.append(ChatMessage(text=prompt, role="user"))
+
+    expected_data_settings = {
+        "data_sources": [
+            {
+                "type": "AzureCognitiveSearch",
+                "parameters": {
+                    "indexName": "test_index",
+                    "endpoint": "https://test-endpoint-search.com",
+                    "key": "test_key",
+                },
+            }
+        ]
+    }
+
+    azure_chat_client = AzureOpenAIChatClient()
+
+    content = await azure_chat_client.get_response(
+        messages=messages_in,
+        additional_properties={"extra_body": expected_data_settings},
+    )
+    assert len(content.messages) == 1
+    assert len(content.messages[0].contents) == 1
+    assert isinstance(content.messages[0].contents[0], TextContent)
+    assert len(content.messages[0].contents[0].annotations) == 1
+    assert content.messages[0].contents[0].annotations[0].title == "test title"
+    assert content.messages[0].contents[0].text == "test"
+
+    mock_create.assert_awaited_once_with(
+        model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
+        messages=azure_chat_client._prepare_chat_history_for_request(messages_out),  # type: ignore
+        stream=False,
+        extra_body=expected_data_settings,
+    )
+
+
+@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
+async def test_azure_on_your_data_fail(
+    mock_create: AsyncMock,
+    azure_openai_unit_test_env: dict[str, str],
+    chat_history: list[ChatMessage],
+    mock_chat_completion_response: ChatCompletion,
+) -> None:
+    mock_chat_completion_response.choices = [
+        Choice(
+            index=0,
+            message=ChatCompletionMessage(
+                content="test",
+                role="assistant",
+                context="not a dictionary",  # type: ignore
+            ),
+            finish_reason="stop",
+        )
+    ]
+    mock_create.return_value = mock_chat_completion_response
+    prompt = "hello world"
+    messages_in = chat_history
+    messages_in.append(ChatMessage(text=prompt, role="user"))
+    messages_out: list[ChatMessage] = []
+    messages_out.append(ChatMessage(text=prompt, role="user"))
+
+    expected_data_settings = {
+        "data_sources": [
+            {
+                "type": "AzureCognitiveSearch",
+                "parameters": {
+                    "indexName": "test_index",
+                    "endpoint": "https://test-endpoint-search.com",
+                    "key": "test_key",
+                },
+            }
+        ]
+    }
+
+    azure_chat_client = AzureOpenAIChatClient()
+
+    content = await azure_chat_client.get_response(
+        messages=messages_in,
+        additional_properties={"extra_body": expected_data_settings},
+    )
+    assert len(content.messages) == 1
+    assert len(content.messages[0].contents) == 1
+    assert isinstance(content.messages[0].contents[0], TextContent)
+    assert content.messages[0].contents[0].text == "test"
+
+    mock_create.assert_awaited_once_with(
+        model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
+        messages=azure_chat_client._prepare_chat_history_for_request(messages_out),  # type: ignore
+        stream=False,
+        extra_body=expected_data_settings,
+    )
+
+
+CONTENT_FILTERED_ERROR_MESSAGE = (
+    "The response was filtered due to the prompt triggering Azure OpenAI's content management policy. Please "
+    "modify your prompt and retry. To learn more about our content filtering policies please read our "
+    "documentation: https://go.microsoft.com/fwlink/?linkid=2198766"
+)
+CONTENT_FILTERED_ERROR_FULL_MESSAGE = (
+    "Error code: 400 - {'error': {'message': \"%s\", 'type': null, 'param': 'prompt', 'code': 'content_filter', "
+    "'status': 400, 'innererror': {'code': 'ResponsibleAIPolicyViolation', 'content_filter_result': {'hate': "
+    "{'filtered': True, 'severity': 'high'}, 'self_harm': {'filtered': False, 'severity': 'safe'}, 'sexual': "
+    "{'filtered': False, 'severity': 'safe'}, 'violence': {'filtered': False, 'severity': 'safe'}}}}}"
+) % CONTENT_FILTERED_ERROR_MESSAGE
+
+
+@patch.object(AsyncChatCompletions, "create")
+async def test_content_filtering_raises_correct_exception(
+    mock_create: AsyncMock,
+    azure_openai_unit_test_env: dict[str, str],
+    chat_history: list[ChatMessage],
+) -> None:
+    prompt = "some prompt that would trigger the content filtering"
+    chat_history.append(ChatMessage(text=prompt, role="user"))
+
+    test_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
+    assert test_endpoint is not None
+    mock_create.side_effect = openai.BadRequestError(
+        CONTENT_FILTERED_ERROR_FULL_MESSAGE,
+        response=Response(400, request=Request("POST", test_endpoint)),
+        body={
+            "message": CONTENT_FILTERED_ERROR_MESSAGE,
+            "type": None,
+            "param": "prompt",
+            "code": "content_filter",
+            "status": 400,
+            "innererror": {
+                "code": "ResponsibleAIPolicyViolation",
+                "content_filter_result": {
+                    "hate": {"filtered": True, "severity": "high"},
+                    "self_harm": {"filtered": False, "severity": "safe"},
+                    "sexual": {"filtered": False, "severity": "safe"},
+                    "violence": {"filtered": False, "severity": "safe"},
+                },
+            },
+        },
+    )
+
+    azure_chat_client = AzureOpenAIChatClient()
+
+    with pytest.raises(OpenAIContentFilterException, match="service encountered a content error") as exc_info:
+        await azure_chat_client.get_response(
+            messages=chat_history,
+        )
+
+    content_filter_exc = exc_info.value
+    assert content_filter_exc.param == "prompt"
+    assert content_filter_exc.content_filter_result["hate"].filtered
+    assert content_filter_exc.content_filter_result["hate"].severity == ContentFilterResultSeverity.HIGH
+
+
+@patch.object(AsyncChatCompletions, "create")
+async def test_content_filtering_without_response_code_raises_with_default_code(
+    mock_create: AsyncMock,
+    azure_openai_unit_test_env: dict[str, str],
+    chat_history: list[ChatMessage],
+) -> None:
+    prompt = "some prompt that would trigger the content filtering"
+    chat_history.append(ChatMessage(text=prompt, role="user"))
+
+    test_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
+    assert test_endpoint is not None
+    mock_create.side_effect = openai.BadRequestError(
+        CONTENT_FILTERED_ERROR_FULL_MESSAGE,
+        response=Response(400, request=Request("POST", test_endpoint)),
+        body={
+            "message": CONTENT_FILTERED_ERROR_MESSAGE,
+            "type": None,
+            "param": "prompt",
+            "code": "content_filter",
+            "status": 400,
+            "innererror": {
+                "content_filter_result": {
+                    "hate": {"filtered": True, "severity": "high"},
+                    "self_harm": {"filtered": False, "severity": "safe"},
+                    "sexual": {"filtered": False, "severity": "safe"},
+                    "violence": {"filtered": False, "severity": "safe"},
+                },
+            },
+        },
+    )
+
+    azure_chat_client = AzureOpenAIChatClient()
+
+    with pytest.raises(OpenAIContentFilterException, match="service encountered a content error"):
+        await azure_chat_client.get_response(
+            messages=chat_history,
+        )
+
+
+@patch.object(AsyncChatCompletions, "create")
+async def test_bad_request_non_content_filter(
+    mock_create: AsyncMock,
+    azure_openai_unit_test_env: dict[str, str],
+    chat_history: list[ChatMessage],
+) -> None:
+    prompt = "some prompt that would trigger the content filtering"
+    chat_history.append(ChatMessage(text=prompt, role="user"))
+
+    test_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
+    assert test_endpoint is not None
+    mock_create.side_effect = openai.BadRequestError(
+        "The request was bad.", response=Response(400, request=Request("POST", test_endpoint)), body={}
+    )
+
+    azure_chat_client = AzureOpenAIChatClient()
+
+    with pytest.raises(ServiceResponseException, match="service failed to complete the prompt"):
+        await azure_chat_client.get_response(
+            messages=chat_history,
+        )
+
+
+@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
+async def test_get_streaming(
+    mock_create: AsyncMock,
+    azure_openai_unit_test_env: dict[str, str],
+    chat_history: list[ChatMessage],
+    mock_streaming_chat_completion_response: AsyncStream[ChatCompletionChunk],
+) -> None:
+    mock_create.return_value = mock_streaming_chat_completion_response
+    chat_history.append(ChatMessage(text="hello world", role="user"))
+
+    azure_chat_client = AzureOpenAIChatClient()
+    async for msg in azure_chat_client.get_streaming_response(
+        messages=chat_history,
+    ):
+        assert msg is not None
+        assert msg.message_id is not None
+        assert msg.response_id is not None
+    mock_create.assert_awaited_once_with(
+        model=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
+        stream=True,
+        messages=azure_chat_client._prepare_chat_history_for_request(chat_history),  # type: ignore
+        # NOTE: The `stream_options={"include_usage": True}` is explicitly enforced in
+        # `OpenAIChatCompletionBase._inner_get_streaming_response`.
+        # To ensure consistency, we align the arguments here accordingly.
+        stream_options={"include_usage": True},
+    )
+
+
+@ai_function
+def get_story_text() -> str:
+    """Returns a story about Emily and David."""
+    return (
+        "Emily and David, two passionate scientists, met during a research expedition to Antarctica. "
+        "Bonded by their love for the natural world and shared curiosity, they uncovered a "
+        "groundbreaking phenomenon in glaciology that could potentially reshape our understanding "
+        "of climate change."
+    )
+
+
+@ai_function
+def get_weather(location: str) -> str:
+    """Get the current weather for a location."""
+    return f"The weather in {location} is sunny and 72°F."
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_openai_chat_client_response() -> None:
+    """Test Azure OpenAI chat completion responses."""
+    azure_chat_client = AzureOpenAIChatClient(credential=AzureCliCredential())
+    assert isinstance(azure_chat_client, ChatClientProtocol)
+
+    messages: list[ChatMessage] = []
+    messages.append(
+        ChatMessage(
+            role="user",
+            text="Emily and David, two passionate scientists, met during a research expedition to Antarctica. "
+            "Bonded by their love for the natural world and shared curiosity, they uncovered a "
+            "groundbreaking phenomenon in glaciology that could potentially reshape our understanding "
+            "of climate change.",
+        )
+    )
+    messages.append(ChatMessage(role="user", text="who are Emily and David?"))
+
+    # Test that the client can be used to get a response
+    response = await azure_chat_client.get_response(messages=messages)
+
+    assert response is not None
+    assert isinstance(response, ChatResponse)
+    # Check for any relevant keywords that indicate the AI understood the context
+    assert any(
+        word in response.text.lower() for word in ["scientists", "research", "antarctica", "glaciology", "climate"]
+    )
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_openai_chat_client_response_tools() -> None:
+    """Test AzureOpenAI chat completion responses."""
+    azure_chat_client = AzureOpenAIChatClient(credential=AzureCliCredential())
+    assert isinstance(azure_chat_client, ChatClientProtocol)
+
+    messages: list[ChatMessage] = []
+    messages.append(ChatMessage(role="user", text="who are Emily and David?"))
+
+    # Test that the client can be used to get a response
+    response = await azure_chat_client.get_response(
+        messages=messages,
+        tools=[get_story_text],
+        tool_choice="auto",
+    )
+
+    assert response is not None
+    assert isinstance(response, ChatResponse)
+    assert "scientists" in response.text
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_openai_chat_client_streaming() -> None:
+    """Test Azure OpenAI chat completion responses."""
+    azure_chat_client = AzureOpenAIChatClient(credential=AzureCliCredential())
+    assert isinstance(azure_chat_client, ChatClientProtocol)
+
+    messages: list[ChatMessage] = []
+    messages.append(
+        ChatMessage(
+            role="user",
+            text="Emily and David, two passionate scientists, met during a research expedition to Antarctica. "
+            "Bonded by their love for the natural world and shared curiosity, they uncovered a "
+            "groundbreaking phenomenon in glaciology that could potentially reshape our understanding "
+            "of climate change.",
+        )
+    )
+    messages.append(ChatMessage(role="user", text="who are Emily and David?"))
+
+    # Test that the client can be used to get a response
+    response = azure_chat_client.get_streaming_response(messages=messages)
+
+    full_message: str = ""
+    async for chunk in response:
+        assert chunk is not None
+        assert isinstance(chunk, ChatResponseUpdate)
+        assert chunk.message_id is not None
+        assert chunk.response_id is not None
+        for content in chunk.contents:
+            if isinstance(content, TextContent) and content.text:
+                full_message += content.text
+
+    assert "scientists" in full_message
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_openai_chat_client_streaming_tools() -> None:
+    """Test AzureOpenAI chat completion responses."""
+    azure_chat_client = AzureOpenAIChatClient(credential=AzureCliCredential())
+    assert isinstance(azure_chat_client, ChatClientProtocol)
+
+    messages: list[ChatMessage] = []
+    messages.append(ChatMessage(role="user", text="who are Emily and David?"))
+
+    # Test that the client can be used to get a response
+    response = azure_chat_client.get_streaming_response(
+        messages=messages,
+        tools=[get_story_text],
+        tool_choice="auto",
+    )
+    full_message: str = ""
+    async for chunk in response:
+        assert chunk is not None
+        assert isinstance(chunk, ChatResponseUpdate)
+        for content in chunk.contents:
+            if isinstance(content, TextContent) and content.text:
+                full_message += content.text
+
+    assert "scientists" in full_message
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_openai_chat_client_agent_basic_run():
+    """Test Azure OpenAI chat client agent basic run functionality with AzureOpenAIChatClient."""
+    async with ChatAgent(
+        chat_client=AzureOpenAIChatClient(credential=AzureCliCredential()),
+    ) as agent:
+        # Test basic run
+        response = await agent.run("Hello! Please respond with 'Hello World' exactly.")
+
+        assert isinstance(response, AgentRunResponse)
+        assert response.text is not None
+        assert len(response.text) > 0
+        assert "hello world" in response.text.lower()
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_openai_chat_client_agent_basic_run_streaming():
+    """Test Azure OpenAI chat client agent basic streaming functionality with AzureOpenAIChatClient."""
+    async with ChatAgent(
+        chat_client=AzureOpenAIChatClient(credential=AzureCliCredential()),
+    ) as agent:
+        # Test streaming run
+        full_text = ""
+        async for chunk in agent.run_stream("Please respond with exactly: 'This is a streaming response test.'"):
+            assert isinstance(chunk, AgentRunResponseUpdate)
+            if chunk.text:
+                full_text += chunk.text
+
+        assert len(full_text) > 0
+        assert "streaming response test" in full_text.lower()
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_openai_chat_client_agent_thread_persistence():
+    """Test Azure OpenAI chat client agent thread persistence across runs with AzureOpenAIChatClient."""
+    async with ChatAgent(
+        chat_client=AzureOpenAIChatClient(credential=AzureCliCredential()),
+        instructions="You are a helpful assistant with good memory.",
+    ) as agent:
+        # Create a new thread that will be reused
+        thread = agent.get_new_thread()
+
+        # First interaction
+        response1 = await agent.run("My name is Alice. Remember this.", thread=thread)
+
+        assert isinstance(response1, AgentRunResponse)
+        assert response1.text is not None
+
+        # Second interaction - test memory
+        response2 = await agent.run("What is my name?", thread=thread)
+
+        assert isinstance(response2, AgentRunResponse)
+        assert response2.text is not None
+        assert "alice" in response2.text.lower()
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_openai_chat_client_agent_existing_thread():
+    """Test Azure OpenAI chat client agent with existing thread to continue conversations across agent instances."""
+    # First conversation - capture the thread
+    preserved_thread = None
+
+    async with ChatAgent(
+        chat_client=AzureOpenAIChatClient(credential=AzureCliCredential()),
+        instructions="You are a helpful assistant with good memory.",
+    ) as first_agent:
+        # Start a conversation and capture the thread
+        thread = first_agent.get_new_thread()
+        first_response = await first_agent.run("My name is Alice. Remember this.", thread=thread)
+
+        assert isinstance(first_response, AgentRunResponse)
+        assert first_response.text is not None
+
+        # Preserve the thread for reuse
+        preserved_thread = thread
+
+    # Second conversation - reuse the thread in a new agent instance
+    if preserved_thread:
+        async with ChatAgent(
+            chat_client=AzureOpenAIChatClient(credential=AzureCliCredential()),
+            instructions="You are a helpful assistant with good memory.",
+        ) as second_agent:
+            # Reuse the preserved thread
+            second_response = await second_agent.run("What is my name?", thread=preserved_thread)
+
+            assert isinstance(second_response, AgentRunResponse)
+            assert second_response.text is not None
+            assert "alice" in second_response.text.lower()
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_chat_client_agent_level_tool_persistence():
+    """Test that agent-level tools persist across multiple runs with Azure Chat Client."""
+
+    async with ChatAgent(
+        chat_client=AzureOpenAIChatClient(credential=AzureCliCredential()),
+        instructions="You are a helpful assistant that uses available tools.",
+        tools=[get_weather],  # Agent-level tool
+    ) as agent:
+        # First run - agent-level tool should be available
+        first_response = await agent.run("What's the weather like in Chicago?")
+
+        assert isinstance(first_response, AgentRunResponse)
+        assert first_response.text is not None
+        # Should use the agent-level weather tool
+        assert any(term in first_response.text.lower() for term in ["chicago", "sunny", "72"])
+
+        # Second run - agent-level tool should still be available (persistence test)
+        second_response = await agent.run("What's the weather in Miami?")
+
+        assert isinstance(second_response, AgentRunResponse)
+        assert second_response.text is not None
+        # Should use the agent-level weather tool again
+        assert any(term in second_response.text.lower() for term in ["miami", "sunny", "72"])

python/packages/core/tests/azure/test_azure_responses_client.py

@@ -0,0 +1,624 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import os
+from typing import Annotated
+
+import pytest
+from azure.identity import AzureCliCredential
+from pydantic import BaseModel
+
+from agent_framework import (
+    AgentRunResponse,
+    AgentRunResponseUpdate,
+    AgentThread,
+    ChatAgent,
+    ChatClientProtocol,
+    ChatMessage,
+    ChatResponse,
+    ChatResponseUpdate,
+    HostedCodeInterpreterTool,
+    HostedFileSearchTool,
+    HostedMCPTool,
+    HostedVectorStoreContent,
+    TextContent,
+    ai_function,
+)
+from agent_framework.azure import AzureOpenAIResponsesClient
+from agent_framework.exceptions import ServiceInitializationError
+
+skip_if_azure_integration_tests_disabled = pytest.mark.skipif(
+    os.getenv("RUN_INTEGRATION_TESTS", "false").lower() != "true"
+    or os.getenv("AZURE_OPENAI_ENDPOINT", "") in ("", "https://test-endpoint.com"),
+    reason="No real AZURE_OPENAI_ENDPOINT provided; skipping integration tests."
+    if os.getenv("RUN_INTEGRATION_TESTS", "false").lower() == "true"
+    else "Integration tests are disabled.",
+)
+
+
+class OutputStruct(BaseModel):
+    """A structured output for testing purposes."""
+
+    location: str
+    weather: str
+
+
+@ai_function
+async def get_weather(location: Annotated[str, "The location as a city name"]) -> str:
+    """Get the current weather in a given location."""
+    # Implementation of the tool to get weather
+    return f"The weather in {location} is sunny and 72°F."
+
+
+async def create_vector_store(client: AzureOpenAIResponsesClient) -> tuple[str, HostedVectorStoreContent]:
+    """Create a vector store with sample documents for testing."""
+    file = await client.client.files.create(
+        file=("todays_weather.txt", b"The weather today is sunny with a high of 75F."), purpose="assistants"
+    )
+    vector_store = await client.client.vector_stores.create(
+        name="knowledge_base",
+        expires_after={"anchor": "last_active_at", "days": 1},
+    )
+    result = await client.client.vector_stores.files.create_and_poll(vector_store_id=vector_store.id, file_id=file.id)
+    if result.last_error is not None:
+        raise Exception(f"Vector store file processing failed with status: {result.last_error.message}")
+
+    return file.id, HostedVectorStoreContent(vector_store_id=vector_store.id)
+
+
+async def delete_vector_store(client: AzureOpenAIResponsesClient, file_id: str, vector_store_id: str) -> None:
+    """Delete the vector store after tests."""
+
+    await client.client.vector_stores.delete(vector_store_id=vector_store_id)
+    await client.client.files.delete(file_id=file_id)
+
+
+def test_init(azure_openai_unit_test_env: dict[str, str]) -> None:
+    # Test successful initialization
+    azure_responses_client = AzureOpenAIResponsesClient()
+
+    assert azure_responses_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]
+    assert isinstance(azure_responses_client, ChatClientProtocol)
+
+
+def test_init_validation_fail() -> None:
+    # Test successful initialization
+    with pytest.raises(ServiceInitializationError):
+        AzureOpenAIResponsesClient(api_key="34523", deployment_name={"test": "dict"})  # type: ignore
+
+
+def test_init_ai_model_id_constructor(azure_openai_unit_test_env: dict[str, str]) -> None:
+    # Test successful initialization
+    ai_model_id = "test_model_id"
+    azure_responses_client = AzureOpenAIResponsesClient(deployment_name=ai_model_id)
+
+    assert azure_responses_client.ai_model_id == ai_model_id
+    assert isinstance(azure_responses_client, ChatClientProtocol)
+
+
+def test_init_with_default_header(azure_openai_unit_test_env: dict[str, str]) -> None:
+    default_headers = {"X-Unit-Test": "test-guid"}
+
+    # Test successful initialization
+    azure_responses_client = AzureOpenAIResponsesClient(
+        default_headers=default_headers,
+    )
+
+    assert azure_responses_client.ai_model_id == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]
+    assert isinstance(azure_responses_client, ChatClientProtocol)
+
+    # Assert that the default header we added is present in the client's default headers
+    for key, value in default_headers.items():
+        assert key in azure_responses_client.client.default_headers
+        assert azure_responses_client.client.default_headers[key] == value
+
+
+@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]], indirect=True)
+def test_init_with_empty_model_id(azure_openai_unit_test_env: dict[str, str]) -> None:
+    with pytest.raises(ServiceInitializationError):
+        AzureOpenAIResponsesClient(
+            env_file_path="test.env",
+        )
+
+
+def test_serialize(azure_openai_unit_test_env: dict[str, str]) -> None:
+    default_headers = {"X-Unit-Test": "test-guid"}
+
+    settings = {
+        "ai_model_id": azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"],
+        "api_key": azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
+        "default_headers": default_headers,
+    }
+
+    azure_responses_client = AzureOpenAIResponsesClient.from_dict(settings)
+    dumped_settings = azure_responses_client.to_dict()
+    assert dumped_settings["ai_model_id"] == azure_openai_unit_test_env["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"]
+    assert dumped_settings["api_key"] == azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"]
+    # Assert that the default header we added is present in the dumped_settings default headers
+    for key, value in default_headers.items():
+        assert key in dumped_settings["default_headers"]
+        assert dumped_settings["default_headers"][key] == value
+    # Assert that the 'User-Agent' header is not present in the dumped_settings default headers
+    assert "User-Agent" not in dumped_settings["default_headers"]
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_responses_client_response() -> None:
+    """Test azure responses client responses."""
+    azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
+
+    assert isinstance(azure_responses_client, ChatClientProtocol)
+
+    messages: list[ChatMessage] = []
+    messages.append(
+        ChatMessage(
+            role="user",
+            text="Emily and David, two passionate scientists, met during a research expedition to Antarctica. "
+            "Bonded by their love for the natural world and shared curiosity, they uncovered a "
+            "groundbreaking phenomenon in glaciology that could potentially reshape our understanding "
+            "of climate change.",
+        )
+    )
+    messages.append(ChatMessage(role="user", text="who are Emily and David?"))
+
+    # Test that the client can be used to get a response
+    response = await azure_responses_client.get_response(messages=messages)
+
+    assert response is not None
+    assert isinstance(response, ChatResponse)
+    assert "scientists" in response.text
+
+    messages.clear()
+    messages.append(ChatMessage(role="user", text="The weather in New York is sunny"))
+    messages.append(ChatMessage(role="user", text="What is the weather in New York?"))
+
+    # Test that the client can be used to get a structured response
+    structured_response = await azure_responses_client.get_response(  # type: ignore[reportAssignmentType]
+        messages=messages,
+        response_format=OutputStruct,
+    )
+
+    assert structured_response is not None
+    assert isinstance(structured_response, ChatResponse)
+    assert isinstance(structured_response.value, OutputStruct)
+    assert structured_response.value.location == "New York"
+    assert "sunny" in structured_response.value.weather.lower()
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_responses_client_response_tools() -> None:
+    """Test azure responses client tools."""
+    azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
+
+    assert isinstance(azure_responses_client, ChatClientProtocol)
+
+    messages: list[ChatMessage] = []
+    messages.append(ChatMessage(role="user", text="What is the weather in New York?"))
+
+    # Test that the client can be used to get a response
+    response = await azure_responses_client.get_response(
+        messages=messages,
+        tools=[get_weather],
+        tool_choice="auto",
+    )
+
+    assert response is not None
+    assert isinstance(response, ChatResponse)
+    assert "sunny" in response.text
+
+    messages.clear()
+    messages.append(ChatMessage(role="user", text="What is the weather in Seattle?"))
+
+    # Test that the client can be used to get a response
+    structured_response: ChatResponse = await azure_responses_client.get_response(  # type: ignore[reportAssignmentType]
+        messages=messages,
+        tools=[get_weather],
+        tool_choice="auto",
+        response_format=OutputStruct,
+    )
+
+    assert structured_response is not None
+    assert isinstance(structured_response, ChatResponse)
+    assert isinstance(structured_response.value, OutputStruct)
+    assert "Seattle" in structured_response.value.location
+    assert "sunny" in structured_response.value.weather.lower()
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_responses_client_streaming() -> None:
+    """Test Azure azure responses client streaming responses."""
+    azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
+
+    assert isinstance(azure_responses_client, ChatClientProtocol)
+
+    messages: list[ChatMessage] = []
+    messages.append(
+        ChatMessage(
+            role="user",
+            text="Emily and David, two passionate scientists, met during a research expedition to Antarctica. "
+            "Bonded by their love for the natural world and shared curiosity, they uncovered a "
+            "groundbreaking phenomenon in glaciology that could potentially reshape our understanding "
+            "of climate change.",
+        )
+    )
+    messages.append(ChatMessage(role="user", text="who are Emily and David?"))
+
+    # Test that the client can be used to get a response
+    response = azure_responses_client.get_streaming_response(messages=messages)
+
+    full_message: str = ""
+    async for chunk in response:
+        assert chunk is not None
+        assert isinstance(chunk, ChatResponseUpdate)
+        for content in chunk.contents:
+            if isinstance(content, TextContent) and content.text:
+                full_message += content.text
+
+    assert "scientists" in full_message
+
+    messages.clear()
+    messages.append(ChatMessage(role="user", text="The weather in Seattle is sunny"))
+    messages.append(ChatMessage(role="user", text="What is the weather in Seattle?"))
+
+    structured_response = await ChatResponse.from_chat_response_generator(
+        azure_responses_client.get_streaming_response(
+            messages=messages,
+            response_format=OutputStruct,
+        ),
+        output_format_type=OutputStruct,
+    )
+    assert structured_response is not None
+    assert isinstance(structured_response, ChatResponse)
+    assert isinstance(structured_response.value, OutputStruct)
+    assert "Seattle" in structured_response.value.location
+    assert "sunny" in structured_response.value.weather.lower()
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_responses_client_streaming_tools() -> None:
+    """Test azure responses client streaming tools."""
+    azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
+
+    assert isinstance(azure_responses_client, ChatClientProtocol)
+
+    messages: list[ChatMessage] = [ChatMessage(role="user", text="What is the weather in Seattle?")]
+
+    # Test that the client can be used to get a response
+    response = azure_responses_client.get_streaming_response(
+        messages=messages,
+        tools=[get_weather],
+        tool_choice="auto",
+    )
+    full_message: str = ""
+    async for chunk in response:
+        assert chunk is not None
+        assert isinstance(chunk, ChatResponseUpdate)
+        for content in chunk.contents:
+            if isinstance(content, TextContent) and content.text:
+                full_message += content.text
+
+    assert "sunny" in full_message
+
+    messages.clear()
+    messages.append(ChatMessage(role="user", text="What is the weather in Seattle?"))
+
+    structured_response = azure_responses_client.get_streaming_response(
+        messages=messages,
+        tools=[get_weather],
+        tool_choice="auto",
+        response_format=OutputStruct,
+    )
+    full_message = ""
+    async for chunk in structured_response:
+        assert chunk is not None
+        assert isinstance(chunk, ChatResponseUpdate)
+        for content in chunk.contents:
+            if isinstance(content, TextContent) and content.text:
+                full_message += content.text
+
+    output = OutputStruct.model_validate_json(full_message)
+    assert "Seattle" in output.location
+    assert "sunny" in output.weather.lower()
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_responses_client_agent_basic_run():
+    """Test Azure Responses Client agent basic run functionality with AzureOpenAIResponsesClient."""
+    agent = AzureOpenAIResponsesClient(credential=AzureCliCredential()).create_agent(
+        instructions="You are a helpful assistant.",
+    )
+
+    # Test basic run
+    response = await agent.run("Hello! Please respond with 'Hello World' exactly.")
+
+    assert isinstance(response, AgentRunResponse)
+    assert response.text is not None
+    assert len(response.text) > 0
+    assert "hello world" in response.text.lower()
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_responses_client_agent_basic_run_streaming():
+    """Test Azure Responses Client agent basic streaming functionality with AzureOpenAIResponsesClient."""
+    async with ChatAgent(
+        chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
+    ) as agent:
+        # Test streaming run
+        full_text = ""
+        async for chunk in agent.run_stream("Please respond with exactly: 'This is a streaming response test.'"):
+            assert isinstance(chunk, AgentRunResponseUpdate)
+            if chunk.text:
+                full_text += chunk.text
+
+        assert len(full_text) > 0
+        assert "streaming response test" in full_text.lower()
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_responses_client_agent_thread_persistence():
+    """Test Azure Responses Client agent thread persistence across runs with AzureOpenAIResponsesClient."""
+    async with ChatAgent(
+        chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
+        instructions="You are a helpful assistant with good memory.",
+    ) as agent:
+        # Create a new thread that will be reused
+        thread = agent.get_new_thread()
+
+        # First interaction
+        first_response = await agent.run("My favorite programming language is Python. Remember this.", thread=thread)
+
+        assert isinstance(first_response, AgentRunResponse)
+        assert first_response.text is not None
+
+        # Second interaction - test memory
+        second_response = await agent.run("What is my favorite programming language?", thread=thread)
+
+        assert isinstance(second_response, AgentRunResponse)
+        assert second_response.text is not None
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_responses_client_agent_thread_storage_with_store_true():
+    """Test Azure Responses Client agent with store=True to verify service_thread_id is returned."""
+    async with ChatAgent(
+        chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
+        instructions="You are a helpful assistant.",
+    ) as agent:
+        # Create a new thread
+        thread = AgentThread()
+
+        # Initially, service_thread_id should be None
+        assert thread.service_thread_id is None
+
+        # Run with store=True to store messages on Azure/OpenAI side
+        response = await agent.run(
+            "Hello! Please remember that my name is Alex.",
+            thread=thread,
+            store=True,
+        )
+
+        # Validate response
+        assert isinstance(response, AgentRunResponse)
+        assert response.text is not None
+        assert len(response.text) > 0
+
+        # After store=True, service_thread_id should be populated
+        assert thread.service_thread_id is not None
+        assert isinstance(thread.service_thread_id, str)
+        assert len(thread.service_thread_id) > 0
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_responses_client_agent_existing_thread():
+    """Test Azure Responses Client agent with existing thread to continue conversations across agent instances."""
+    # First conversation - capture the thread
+    preserved_thread = None
+
+    async with ChatAgent(
+        chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
+        instructions="You are a helpful assistant with good memory.",
+    ) as first_agent:
+        # Start a conversation and capture the thread
+        thread = first_agent.get_new_thread()
+        first_response = await first_agent.run("My hobby is photography. Remember this.", thread=thread)
+
+        assert isinstance(first_response, AgentRunResponse)
+        assert first_response.text is not None
+
+        # Preserve the thread for reuse
+        preserved_thread = thread
+
+    # Second conversation - reuse the thread in a new agent instance
+    if preserved_thread:
+        async with ChatAgent(
+            chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
+            instructions="You are a helpful assistant with good memory.",
+        ) as second_agent:
+            # Reuse the preserved thread
+            second_response = await second_agent.run("What is my hobby?", thread=preserved_thread)
+
+            assert isinstance(second_response, AgentRunResponse)
+            assert second_response.text is not None
+            assert "photography" in second_response.text.lower()
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_responses_client_agent_hosted_code_interpreter_tool():
+    """Test Azure Responses Client agent with HostedCodeInterpreterTool through AzureOpenAIResponsesClient."""
+    async with ChatAgent(
+        chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
+        instructions="You are a helpful assistant that can execute Python code.",
+        tools=[HostedCodeInterpreterTool()],
+    ) as agent:
+        # Test code interpreter functionality
+        response = await agent.run("Calculate the sum of numbers from 1 to 10 using Python code.")
+
+        assert isinstance(response, AgentRunResponse)
+        assert response.text is not None
+        assert len(response.text) > 0
+        # Should contain calculation result (sum of 1-10 = 55) or code execution content
+        contains_relevant_content = any(
+            term in response.text.lower() for term in ["55", "sum", "code", "python", "calculate", "10"]
+        )
+        assert contains_relevant_content or len(response.text.strip()) > 10
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_responses_client_agent_level_tool_persistence():
+    """Test that agent-level tools persist across multiple runs with Azure Responses Client."""
+
+    async with ChatAgent(
+        chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
+        instructions="You are a helpful assistant that uses available tools.",
+        tools=[get_weather],  # Agent-level tool
+    ) as agent:
+        # First run - agent-level tool should be available
+        first_response = await agent.run("What's the weather like in Chicago?")
+
+        assert isinstance(first_response, AgentRunResponse)
+        assert first_response.text is not None
+        # Should use the agent-level weather tool
+        assert any(term in first_response.text.lower() for term in ["chicago", "sunny", "72"])
+
+        # Second run - agent-level tool should still be available (persistence test)
+        second_response = await agent.run("What's the weather in Miami?")
+
+        assert isinstance(second_response, AgentRunResponse)
+        assert second_response.text is not None
+        # Should use the agent-level weather tool again
+        assert any(term in second_response.text.lower() for term in ["miami", "sunny", "72"])
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_responses_client_agent_chat_options_run_level() -> None:
+    """Integration test for comprehensive ChatOptions parameter coverage with Azure Response Agent."""
+    async with ChatAgent(
+        chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
+        instructions="You are a helpful assistant.",
+    ) as agent:
+        response = await agent.run(
+            "Provide a brief, helpful response.",
+            max_tokens=100,
+            temperature=0.7,
+            top_p=0.9,
+            seed=123,
+            user="comprehensive-test-user",
+            tools=[get_weather],
+            tool_choice="auto",
+        )
+
+        assert isinstance(response, AgentRunResponse)
+        assert response.text is not None
+        assert len(response.text) > 0
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_responses_client_agent_chat_options_agent_level() -> None:
+    """Integration test for comprehensive ChatOptions parameter coverage with Azure Response Agent."""
+    async with ChatAgent(
+        chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
+        instructions="You are a helpful assistant.",
+        max_tokens=100,
+        temperature=0.7,
+        top_p=0.9,
+        seed=123,
+        user="comprehensive-test-user",
+        tools=[get_weather],
+        tool_choice="auto",
+    ) as agent:
+        response = await agent.run(
+            "Provide a brief, helpful response.",
+        )
+
+        assert isinstance(response, AgentRunResponse)
+        assert response.text is not None
+        assert len(response.text) > 0
+
+
+@skip_if_azure_integration_tests_disabled
+async def test_azure_responses_client_agent_hosted_mcp_tool() -> None:
+    """Integration test for HostedMCPTool with Azure Response Agent using Microsoft Learn MCP."""
+
+    mcp_tool = HostedMCPTool(
+        name="Microsoft Learn MCP",
+        url="https://learn.microsoft.com/api/mcp",
+        description="A Microsoft Learn MCP server for documentation questions",
+        approval_mode="never_require",
+    )
+
+    async with ChatAgent(
+        chat_client=AzureOpenAIResponsesClient(credential=AzureCliCredential()),
+        instructions="You are a helpful assistant that can help with microsoft documentation questions.",
+        tools=[mcp_tool],
+    ) as agent:
+        response = await agent.run(
+            "How to create an Azure storage account using az cli?",
+            max_tokens=200,
+        )
+
+        assert isinstance(response, AgentRunResponse)
+        assert response.text is not None
+        assert len(response.text) > 0
+        # Should contain Azure-related content since it's asking about Azure CLI
+        assert any(term in response.text.lower() for term in ["azure", "storage", "account", "cli"])
+
+
+@skip_if_azure_integration_tests_disabled
+@pytest.mark.skip(reason="File search requires API key auth, subscription only allows token auth")
+async def test_azure_responses_client_file_search() -> None:
+    """Test Azure responses client with file search tool."""
+    azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
+
+    assert isinstance(azure_responses_client, ChatClientProtocol)
+
+    file_id, vector_store = await create_vector_store(azure_responses_client)
+    # Test that the client will use the web search tool
+    response = await azure_responses_client.get_response(
+        messages=[
+            ChatMessage(
+                role="user",
+                text="What is the weather today? Do a file search to find the answer.",
+            )
+        ],
+        tools=[HostedFileSearchTool(inputs=vector_store)],
+        tool_choice="auto",
+    )
+
+    await delete_vector_store(azure_responses_client, file_id, vector_store.vector_store_id)
+    assert "sunny" in response.text.lower()
+    assert "75" in response.text
+
+
+@skip_if_azure_integration_tests_disabled
+@pytest.mark.skip(reason="File search requires API key auth, subscription only allows token auth")
+async def test_azure_responses_client_file_search_streaming() -> None:
+    """Test Azure responses client with file search tool and streaming."""
+    azure_responses_client = AzureOpenAIResponsesClient(credential=AzureCliCredential())
+
+    assert isinstance(azure_responses_client, ChatClientProtocol)
+
+    file_id, vector_store = await create_vector_store(azure_responses_client)
+    # Test that the client will use the web search tool
+    response = azure_responses_client.get_streaming_response(
+        messages=[
+            ChatMessage(
+                role="user",
+                text="What is the weather today? Do a file search to find the answer.",
+            )
+        ],
+        tools=[HostedFileSearchTool(inputs=vector_store)],
+        tool_choice="auto",
+    )
+
+    assert response is not None
+    full_message: str = ""
+    async for chunk in response:
+        assert chunk is not None
+        assert isinstance(chunk, ChatResponseUpdate)
+        for content in chunk.contents:
+            if isinstance(content, TextContent) and content.text:
+                full_message += content.text
+
+    await delete_vector_store(azure_responses_client, file_id, vector_store.vector_store_id)
+
+    assert "sunny" in full_message.lower()
+    assert "75" in full_message

python/packages/core/tests/azure/test_entra_id_authentication.py

@@ -0,0 +1,156 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+from azure.core.exceptions import ClientAuthenticationError
+
+from agent_framework.azure._entra_id_authentication import (
+    get_entra_auth_token,
+    get_entra_auth_token_async,
+)
+from agent_framework.exceptions import ServiceInvalidAuthError
+
+
+@pytest.fixture
+def mock_credential() -> MagicMock:
+    """Mock synchronous TokenCredential."""
+    mock_cred = MagicMock()
+    # Create a mock token object with a .token attribute
+    mock_token = MagicMock()
+    mock_token.token = "test-access-token-12345"
+    mock_cred.get_token.return_value = mock_token
+    return mock_cred
+
+
+@pytest.fixture
+def mock_async_credential() -> MagicMock:
+    """Mock asynchronous AsyncTokenCredential."""
+    mock_cred = MagicMock()
+    # Create a mock token object with a .token attribute
+    mock_token = MagicMock()
+    mock_token.token = "test-async-access-token-12345"
+    mock_cred.get_token = AsyncMock(return_value=mock_token)
+    return mock_cred
+
+
+def test_get_entra_auth_token_success(mock_credential: MagicMock) -> None:
+    """Test successful token retrieval with sync function."""
+
+    token_endpoint = "https://test-endpoint.com/.default"
+
+    result = get_entra_auth_token(mock_credential, token_endpoint)
+
+    # Assert - check the results
+    assert result == "test-access-token-12345"
+    mock_credential.get_token.assert_called_once_with(token_endpoint)
+
+
+async def test_get_entra_auth_token_async_success(mock_async_credential: MagicMock) -> None:
+    """Test successful token retrieval with async function."""
+
+    token_endpoint = "https://test-endpoint.com/.default"
+
+    result = await get_entra_auth_token_async(mock_async_credential, token_endpoint)
+
+    # Assert - check the results
+    assert result == "test-async-access-token-12345"
+    mock_async_credential.get_token.assert_called_once_with(token_endpoint)
+
+
+def test_get_entra_auth_token_missing_endpoint(mock_credential: MagicMock) -> None:
+    """Test that missing token endpoint raises ServiceInvalidAuthError."""
+    # Test with empty string
+    with pytest.raises(ServiceInvalidAuthError, match="A token endpoint must be provided"):
+        get_entra_auth_token(mock_credential, "")
+
+    # Test with None
+    with pytest.raises(ServiceInvalidAuthError, match="A token endpoint must be provided"):
+        get_entra_auth_token(mock_credential, None)  # type: ignore
+
+
+async def test_get_entra_auth_token_async_missing_endpoint(mock_async_credential: MagicMock) -> None:
+    """Test that missing token endpoint raises ServiceInvalidAuthError in async function."""
+    # Test with empty string
+    with pytest.raises(ServiceInvalidAuthError, match="A token endpoint must be provided"):
+        await get_entra_auth_token_async(mock_async_credential, "")
+
+    # Test with None
+    with pytest.raises(ServiceInvalidAuthError, match="A token endpoint must be provided"):
+        await get_entra_auth_token_async(mock_async_credential, None)  # type: ignore
+
+
+def test_get_entra_auth_token_auth_failure(mock_credential: MagicMock) -> None:
+    """Test that Azure authentication failure returns None."""
+
+    mock_credential.get_token.side_effect = ClientAuthenticationError("Auth failed")
+    token_endpoint = "https://test-endpoint.com/.default"
+
+    result = get_entra_auth_token(mock_credential, token_endpoint)
+
+    # Assert - should return None on auth failure
+    assert result is None
+    mock_credential.get_token.assert_called_once_with(token_endpoint)
+
+
+async def test_get_entra_auth_token_async_auth_failure(mock_async_credential: MagicMock) -> None:
+    """Test that Azure authentication failure returns None in async function."""
+
+    mock_async_credential.get_token.side_effect = ClientAuthenticationError("Auth failed")
+    token_endpoint = "https://test-endpoint.com/.default"
+
+    result = await get_entra_auth_token_async(mock_async_credential, token_endpoint)
+
+    # Assert - should return None on auth failure
+    assert result is None
+    mock_async_credential.get_token.assert_called_once_with(token_endpoint)
+
+
+def test_get_entra_auth_token_none_token_response(mock_credential: MagicMock) -> None:
+    """Test that None token response returns None."""
+    mock_credential.get_token.return_value = None
+    token_endpoint = "https://test-endpoint.com/.default"
+
+    result = get_entra_auth_token(mock_credential, token_endpoint)
+
+    # Assert
+    assert result is None
+    mock_credential.get_token.assert_called_once_with(token_endpoint)
+
+
+async def test_get_entra_auth_token_async_none_token_response(mock_async_credential: MagicMock) -> None:
+    """Test that None token response returns None in async function."""
+    mock_async_credential.get_token.return_value = None
+    token_endpoint = "https://test-endpoint.com/.default"
+
+    result = await get_entra_auth_token_async(mock_async_credential, token_endpoint)
+
+    # Assert
+    assert result is None
+    mock_async_credential.get_token.assert_called_once_with(token_endpoint)
+
+
+def test_get_entra_auth_token_with_kwargs(mock_credential: MagicMock) -> None:
+    """Test that kwargs are passed through to get_token."""
+
+    token_endpoint = "https://test-endpoint.com/.default"
+    extra_kwargs = {"scopes": ["read", "write"], "tenant_id": "test-tenant"}
+
+    result = get_entra_auth_token(mock_credential, token_endpoint, **extra_kwargs)
+
+    # Assert
+    assert result == "test-access-token-12345"
+    mock_credential.get_token.assert_called_once_with(token_endpoint, **extra_kwargs)
+
+
+async def test_get_entra_auth_token_async_with_kwargs(mock_async_credential: MagicMock) -> None:
+    """Test that kwargs are passed through to async get_token."""
+
+    token_endpoint = "https://test-endpoint.com/.default"
+    extra_kwargs = {"scopes": ["read", "write"], "tenant_id": "test-tenant"}
+
+    result = await get_entra_auth_token_async(mock_async_credential, token_endpoint, **extra_kwargs)
+
+    # Assert
+    assert result == "test-async-access-token-12345"
+    mock_async_credential.get_token.assert_called_once_with(token_endpoint, **extra_kwargs)

python/packages/core/tests/conftest.py

@@ -0,0 +1,71 @@
+# 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_otel(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_otel: bool, enable_sensitive_data: bool) -> Generator[SpanExporter]:
+    """Fixture to remove environment variables for ObservabilitySettings."""
+
+    env_vars = [
+        "ENABLE_OTEL",
+        "ENABLE_SENSITIVE_DATA",
+        "OTLP_ENDPOINT",
+        "APPLICATIONINSIGHTS_CONNECTION_STRING",
+    ]
+
+    for key in env_vars:
+        monkeypatch.delenv(key, raising=False)  # type: ignore
+    monkeypatch.setenv("ENABLE_OTEL", str(enable_otel))  # type: ignore
+    if not enable_otel:
+        # 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(env_file_path="test.env")
+    observability_settings._configure()  # pyright: ignore[reportPrivateUsage]
+    monkeypatch.setattr(observability, "OBSERVABILITY_SETTINGS", observability_settings, raising=False)  # type: ignore
+
+    with (
+        patch("agent_framework.observability.OBSERVABILITY_SETTINGS", observability_settings),
+        patch("agent_framework.observability.setup_observability"),
+    ):
+        exporter = InMemorySpanExporter()
+        if enable_otel or enable_sensitive_data:
+            tracer_provider = trace.get_tracer_provider()
+            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/main/conftest.py

@@ -0,0 +1,262 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import logging
+import sys
+from collections.abc import AsyncIterable, MutableSequence
+from typing import Any
+from unittest.mock import patch
+from uuid import uuid4
+
+from pydantic import BaseModel, Field
+from pytest import fixture
+
+from agent_framework import (
+    AgentProtocol,
+    AgentRunResponse,
+    AgentRunResponseUpdate,
+    AgentThread,
+    BaseChatClient,
+    ChatMessage,
+    ChatOptions,
+    ChatResponse,
+    ChatResponseUpdate,
+    Role,
+    TextContent,
+    ToolProtocol,
+    ai_function,
+    use_chat_middleware,
+    use_function_invocation,
+)
+
+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[ChatMessage]:
+    return []
+
+
+# region Tools
+
+
+@fixture
+def ai_tool() -> ToolProtocol:
+    """Returns a generic ToolProtocol."""
+
+    class GenericTool(BaseModel):
+        name: str
+        description: str
+        additional_properties: dict[str, Any] | None = None
+
+        def parameters(self) -> dict[str, Any]:
+            """Return the parameters of the tool as a JSON schema."""
+            return {
+                "name": {"type": "string"},
+            }
+
+    return GenericTool(name="generic_tool", description="A generic tool")
+
+
+@fixture
+def ai_function_tool() -> ToolProtocol:
+    """Returns a executable ToolProtocol."""
+
+    @ai_function
+    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) -> None:
+        self.additional_properties: dict[str, Any] = {}
+        self.call_count: int = 0
+        self.responses: list[ChatResponse] = []
+        self.streaming_responses: list[list[ChatResponseUpdate]] = []
+
+    async def get_response(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage],
+        **kwargs: Any,
+    ) -> 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=ChatMessage(role="assistant", text="test response"))
+
+    async def get_streaming_response(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage],
+        **kwargs: Any,
+    ) -> 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(text=TextContent(text="test streaming response "), role="assistant")
+            yield ChatResponseUpdate(contents=[TextContent(text="another update")], role="assistant")
+
+
+@use_chat_middleware
+class MockBaseChatClient(BaseChatClient):
+    """Mock implementation of the BaseChatClient."""
+
+    run_responses: list[ChatResponse] = Field(default_factory=list)
+    streaming_responses: list[list[ChatResponseUpdate]] = Field(default_factory=list)
+    call_count: int = Field(default=0)
+
+    @override
+    async def _inner_get_response(
+        self,
+        *,
+        messages: MutableSequence[ChatMessage],
+        chat_options: ChatOptions,
+        **kwargs: Any,
+    ) -> ChatResponse:
+        """Send a chat request to the AI service.
+
+        Args:
+            messages: The chat messages to send.
+            chat_options: The options for the request.
+            kwargs: Any additional keyword arguments.
+
+        Returns:
+            The chat response contents representing the response(s).
+        """
+        logger.debug(f"Running base chat client inner, with: {messages=}, {chat_options=}, {kwargs=}")
+        self.call_count += 1
+        if not self.run_responses:
+            return ChatResponse(messages=ChatMessage(role="assistant", text=f"test response - {messages[-1].text}"))
+
+        response = self.run_responses.pop(0)
+
+        if chat_options.tool_choice == "none":
+            return ChatResponse(
+                messages=ChatMessage(
+                    role="assistant",
+                    text="I broke out of the function invocation loop...",
+                ),
+                conversation_id=response.conversation_id,
+            )
+
+        return response
+
+    @override
+    async def _inner_get_streaming_response(
+        self,
+        *,
+        messages: MutableSequence[ChatMessage],
+        chat_options: ChatOptions,
+        **kwargs: Any,
+    ) -> AsyncIterable[ChatResponseUpdate]:
+        logger.debug(f"Running base chat client inner stream, with: {messages=}, {chat_options=}, {kwargs=}")
+        if not self.streaming_responses:
+            yield ChatResponseUpdate(text=f"update - {messages[0].text}", role="assistant")
+            return
+        if chat_options.tool_choice == "none":
+            yield ChatResponseUpdate(text="I broke out of the function invocation loop...", role="assistant")
+            return
+        response = self.streaming_responses.pop(0)
+        for update in response:
+            yield update
+        await asyncio.sleep(0)
+
+
+@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 chat_client(enable_function_calling: bool, max_iterations: int) -> MockChatClient:
+    if enable_function_calling:
+        with patch("agent_framework._tools.DEFAULT_MAX_ITERATIONS", max_iterations):
+            return use_function_invocation(MockChatClient)()
+    return MockChatClient()
+
+
+@fixture
+def chat_client_base(enable_function_calling: bool, max_iterations: int) -> MockBaseChatClient:
+    if enable_function_calling:
+        with patch("agent_framework._tools.DEFAULT_MAX_ITERATIONS", max_iterations):
+            return use_function_invocation(MockBaseChatClient)()
+    return MockBaseChatClient()
+
+
+# region Agents
+class MockAgentThread(AgentThread):
+    pass
+
+
+# Mock Agent implementation for testing
+class MockAgent(AgentProtocol):
+    @property
+    def id(self) -> str:
+        return str(uuid4())
+
+    @property
+    def name(self) -> str | None:
+        """Returns the name of the agent."""
+        return "Name"
+
+    @property
+    def display_name(self) -> str:
+        """Returns the name of the agent."""
+        return "Display Name"
+
+    @property
+    def description(self) -> str | None:
+        return "Description"
+
+    async def run(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
+        *,
+        thread: AgentThread | None = None,
+        **kwargs: Any,
+    ) -> AgentRunResponse:
+        logger.debug(f"Running mock agent, with: {messages=}, {thread=}, {kwargs=}")
+        return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, contents=[TextContent("Response")])])
+
+    async def run_stream(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
+        *,
+        thread: AgentThread | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterable[AgentRunResponseUpdate]:
+        logger.debug(f"Running mock agent stream, with: {messages=}, {thread=}, {kwargs=}")
+        yield AgentRunResponseUpdate(contents=[TextContent("Response")])
+
+    def get_new_thread(self) -> AgentThread:
+        return MockAgentThread()
+
+
+@fixture
+def agent_thread() -> AgentThread:
+    return MockAgentThread()
+
+
+@fixture
+def agent() -> AgentProtocol:
+    return MockAgent()

python/packages/core/tests/main/test_agents.py

@@ -0,0 +1,508 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from collections.abc import AsyncIterable, MutableSequence, Sequence
+from typing import Any
+from uuid import uuid4
+
+from pytest import raises
+
+from agent_framework import (
+    AgentProtocol,
+    AgentRunResponse,
+    AgentRunResponseUpdate,
+    AgentThread,
+    AggregateContextProvider,
+    ChatAgent,
+    ChatClientProtocol,
+    ChatMessage,
+    ChatMessageStore,
+    ChatResponse,
+    Context,
+    ContextProvider,
+    HostedCodeInterpreterTool,
+    Role,
+    TextContent,
+)
+from agent_framework.exceptions import AgentExecutionException
+
+
+def test_agent_thread_type(agent_thread: AgentThread) -> None:
+    assert isinstance(agent_thread, AgentThread)
+
+
+def test_agent_type(agent: AgentProtocol) -> None:
+    assert isinstance(agent, AgentProtocol)
+
+
+async def test_agent_run(agent: AgentProtocol) -> None:
+    response = await agent.run("test")
+    assert response.messages[0].role == Role.ASSISTANT
+    assert response.messages[0].text == "Response"
+
+
+async def test_agent_run_streaming(agent: AgentProtocol) -> None:
+    async def collect_updates(updates: AsyncIterable[AgentRunResponseUpdate]) -> list[AgentRunResponseUpdate]:
+        return [u async for u in updates]
+
+    updates = await collect_updates(agent.run_stream(messages="test"))
+    assert len(updates) == 1
+    assert updates[0].text == "Response"
+
+
+def test_chat_client_agent_type(chat_client: ChatClientProtocol) -> None:
+    chat_client_agent = ChatAgent(chat_client=chat_client)
+    assert isinstance(chat_client_agent, AgentProtocol)
+
+
+async def test_chat_client_agent_init(chat_client: ChatClientProtocol) -> None:
+    agent_id = str(uuid4())
+    agent = ChatAgent(chat_client=chat_client, id=agent_id, description="Test")
+
+    assert agent.id == agent_id
+    assert agent.name is None
+    assert agent.description == "Test"
+    assert agent.display_name == agent_id  # Display name defaults to id if name is None
+
+
+async def test_chat_client_agent_init_with_name(chat_client: ChatClientProtocol) -> None:
+    agent_id = str(uuid4())
+    agent = ChatAgent(chat_client=chat_client, id=agent_id, name="Test Agent", description="Test")
+
+    assert agent.id == agent_id
+    assert agent.name == "Test Agent"
+    assert agent.description == "Test"
+    assert agent.display_name == "Test Agent"  # Display name is the name if present
+
+
+async def test_chat_client_agent_run(chat_client: ChatClientProtocol) -> None:
+    agent = ChatAgent(chat_client=chat_client)
+
+    result = await agent.run("Hello")
+
+    assert result.text == "test response"
+
+
+async def test_chat_client_agent_run_streaming(chat_client: ChatClientProtocol) -> None:
+    agent = ChatAgent(chat_client=chat_client)
+
+    result = await AgentRunResponse.from_agent_response_generator(agent.run_stream("Hello"))
+
+    assert result.text == "test streaming response another update"
+
+
+async def test_chat_client_agent_get_new_thread(chat_client: ChatClientProtocol) -> None:
+    agent = ChatAgent(chat_client=chat_client)
+    thread = agent.get_new_thread()
+
+    assert isinstance(thread, AgentThread)
+
+
+async def test_chat_client_agent_prepare_thread_and_messages(chat_client: ChatClientProtocol) -> None:
+    agent = ChatAgent(chat_client=chat_client)
+    message = ChatMessage(role=Role.USER, text="Hello")
+    thread = AgentThread(message_store=ChatMessageStore(messages=[message]))
+
+    _, _, result_messages = await agent._prepare_thread_and_messages(  # type: ignore[reportPrivateUsage]
+        thread=thread,
+        input_messages=[ChatMessage(role=Role.USER, text="Test")],
+    )
+
+    assert len(result_messages) == 2
+    assert result_messages[0] == message
+    assert result_messages[1].text == "Test"
+
+
+async def test_chat_client_agent_update_thread_id(chat_client_base: ChatClientProtocol) -> None:
+    mock_response = ChatResponse(
+        messages=[ChatMessage(role=Role.ASSISTANT, contents=[TextContent("test response")])],
+        conversation_id="123",
+    )
+    chat_client_base.run_responses = [mock_response]
+    agent = ChatAgent(
+        chat_client=chat_client_base,
+        tools=HostedCodeInterpreterTool(),
+    )
+    thread = agent.get_new_thread()
+
+    result = await agent.run("Hello", thread=thread)
+    assert result.text == "test response"
+
+    assert thread.service_thread_id == "123"
+
+
+async def test_chat_client_agent_update_thread_messages(chat_client: ChatClientProtocol) -> None:
+    agent = ChatAgent(chat_client=chat_client)
+    thread = agent.get_new_thread()
+
+    result = await agent.run("Hello", thread=thread)
+    assert result.text == "test response"
+
+    assert thread.service_thread_id is None
+    assert thread.message_store is not None
+
+    chat_messages: list[ChatMessage] = await thread.message_store.list_messages()
+
+    assert chat_messages is not None
+    assert len(chat_messages) == 2
+    assert chat_messages[0].text == "Hello"
+    assert chat_messages[1].text == "test response"
+
+
+async def test_chat_client_agent_update_thread_conversation_id_missing(chat_client: ChatClientProtocol) -> None:
+    agent = ChatAgent(chat_client=chat_client)
+    thread = AgentThread(service_thread_id="123")
+
+    with raises(AgentExecutionException, match="Service did not return a valid conversation id"):
+        await agent._update_thread_with_type_and_conversation_id(thread, None)  # type: ignore[reportPrivateUsage]
+
+
+async def test_chat_client_agent_default_author_name(chat_client: ChatClientProtocol) -> None:
+    # Name is not specified here, so default name should be used
+    agent = ChatAgent(chat_client=chat_client)
+
+    result = await agent.run("Hello")
+    assert result.text == "test response"
+    assert result.messages[0].author_name == "UnnamedAgent"
+
+
+async def test_chat_client_agent_author_name_as_agent_name(chat_client: ChatClientProtocol) -> None:
+    # Name is specified here, so it should be used as author name
+    agent = ChatAgent(chat_client=chat_client, name="TestAgent")
+
+    result = await agent.run("Hello")
+    assert result.text == "test response"
+    assert result.messages[0].author_name == "TestAgent"
+
+
+async def test_chat_client_agent_author_name_is_used_from_response(chat_client_base: ChatClientProtocol) -> None:
+    chat_client_base.run_responses = [
+        ChatResponse(
+            messages=[
+                ChatMessage(role=Role.ASSISTANT, contents=[TextContent("test response")], author_name="TestAuthor")
+            ]
+        )
+    ]
+
+    agent = ChatAgent(chat_client=chat_client_base, tools=HostedCodeInterpreterTool())
+
+    result = await agent.run("Hello")
+    assert result.text == "test response"
+    assert result.messages[0].author_name == "TestAuthor"
+
+
+# Mock context provider for testing
+class MockContextProvider(ContextProvider):
+    def __init__(self, messages: list[ChatMessage] | None = None) -> None:
+        self.context_messages = messages
+        self.thread_created_called = False
+        self.invoked_called = False
+        self.invoking_called = False
+        self.thread_created_thread_id = None
+        self.invoked_thread_id = None
+        self.new_messages: list[ChatMessage] = []
+
+    async def thread_created(self, thread_id: str | None) -> None:
+        self.thread_created_called = True
+        self.thread_created_thread_id = thread_id
+
+    async def invoked(
+        self,
+        request_messages: ChatMessage | Sequence[ChatMessage],
+        response_messages: ChatMessage | Sequence[ChatMessage] | None = None,
+        invoke_exception: Any = None,
+        **kwargs: Any,
+    ) -> None:
+        self.invoked_called = True
+        if isinstance(request_messages, ChatMessage):
+            self.new_messages.append(request_messages)
+        else:
+            self.new_messages.extend(request_messages)
+        if isinstance(response_messages, ChatMessage):
+            self.new_messages.append(response_messages)
+        else:
+            self.new_messages.extend(response_messages)
+
+    async def invoking(self, messages: ChatMessage | MutableSequence[ChatMessage], **kwargs: Any) -> Context:
+        self.invoking_called = True
+        return Context(messages=self.context_messages)
+
+
+async def test_chat_agent_context_providers_model_invoking(chat_client: ChatClientProtocol) -> None:
+    """Test that context providers' invoking is called during agent run."""
+    mock_provider = MockContextProvider(messages=[ChatMessage(role=Role.SYSTEM, text="Test context instructions")])
+    agent = ChatAgent(chat_client=chat_client, context_providers=mock_provider)
+
+    await agent.run("Hello")
+
+    assert mock_provider.invoking_called
+
+
+async def test_chat_agent_context_providers_thread_created(chat_client_base: ChatClientProtocol) -> None:
+    """Test that context providers' thread_created is called during agent run."""
+    mock_provider = MockContextProvider()
+    chat_client_base.run_responses = [
+        ChatResponse(
+            messages=[ChatMessage(role=Role.ASSISTANT, contents=[TextContent("test response")])],
+            conversation_id="test-thread-id",
+        )
+    ]
+
+    agent = ChatAgent(chat_client=chat_client_base, context_providers=mock_provider)
+
+    await agent.run("Hello")
+
+    assert mock_provider.thread_created_called
+    assert mock_provider.thread_created_thread_id == "test-thread-id"
+
+
+async def test_chat_agent_context_providers_messages_adding(chat_client: ChatClientProtocol) -> None:
+    """Test that context providers' invoked is called during agent run."""
+    mock_provider = MockContextProvider()
+    agent = ChatAgent(chat_client=chat_client, context_providers=mock_provider)
+
+    await agent.run("Hello")
+
+    assert mock_provider.invoked_called
+    # Should be called with both input and response messages
+    assert len(mock_provider.new_messages) >= 2
+
+
+async def test_chat_agent_context_instructions_in_messages(chat_client: ChatClientProtocol) -> None:
+    """Test that AI context instructions are included in messages."""
+    mock_provider = MockContextProvider(messages=[ChatMessage(role="system", text="Context-specific instructions")])
+    agent = ChatAgent(chat_client=chat_client, instructions="Agent instructions", context_providers=mock_provider)
+
+    # We need to test the _prepare_thread_and_messages method directly
+    _, _, messages = await agent._prepare_thread_and_messages(  # type: ignore[reportPrivateUsage]
+        thread=None, input_messages=[ChatMessage(role=Role.USER, text="Hello")]
+    )
+
+    # Should have context instructions, and user message
+    assert len(messages) == 2
+    assert messages[0].role == Role.SYSTEM
+    assert messages[0].text == "Context-specific instructions"
+    assert messages[1].role == Role.USER
+    assert messages[1].text == "Hello"
+    # instructions system message is added by a chat_client
+
+
+async def test_chat_agent_no_context_instructions(chat_client: ChatClientProtocol) -> None:
+    """Test behavior when AI context has no instructions."""
+    mock_provider = MockContextProvider()
+    agent = ChatAgent(chat_client=chat_client, instructions="Agent instructions", context_providers=mock_provider)
+
+    _, _, messages = await agent._prepare_thread_and_messages(  # type: ignore[reportPrivateUsage]
+        thread=None, input_messages=[ChatMessage(role=Role.USER, text="Hello")]
+    )
+
+    # Should have agent instructions and user message only
+    assert len(messages) == 1
+    assert messages[0].role == Role.USER
+    assert messages[0].text == "Hello"
+
+
+async def test_chat_agent_run_stream_context_providers(chat_client: ChatClientProtocol) -> None:
+    """Test that context providers work with run_stream method."""
+    mock_provider = MockContextProvider(messages=[ChatMessage(role=Role.SYSTEM, text="Stream context instructions")])
+    agent = ChatAgent(chat_client=chat_client, context_providers=mock_provider)
+
+    # Collect all stream updates
+    updates: list[AgentRunResponseUpdate] = []
+    async for update in agent.run_stream("Hello"):
+        updates.append(update)
+
+    # Verify context provider was called
+    assert mock_provider.invoking_called
+    # no conversation id is created, so no need to thread_create to be called.
+    assert not mock_provider.thread_created_called
+    assert mock_provider.invoked_called
+
+
+async def test_chat_agent_multiple_context_providers(chat_client: ChatClientProtocol) -> None:
+    """Test that multiple context providers work together."""
+    provider1 = MockContextProvider(messages=[ChatMessage(role=Role.SYSTEM, text="First provider instructions")])
+    provider2 = MockContextProvider(messages=[ChatMessage(role=Role.SYSTEM, text="Second provider instructions")])
+
+    agent = ChatAgent(chat_client=chat_client, context_providers=[provider1, provider2])
+
+    await agent.run("Hello")
+
+    # Both providers should be called
+    assert provider1.invoking_called
+    assert not provider1.thread_created_called
+    assert provider1.invoked_called
+
+    assert provider2.invoking_called
+    assert not provider2.thread_created_called
+    assert provider2.invoked_called
+
+
+async def test_chat_agent_aggregate_context_provider_combines_instructions() -> None:
+    """Test that AggregateContextProvider combines instructions from multiple providers."""
+    provider1 = MockContextProvider(messages=[ChatMessage(role=Role.SYSTEM, text="First instruction")])
+    provider2 = MockContextProvider(messages=[ChatMessage(role=Role.SYSTEM, text="Second instruction")])
+
+    aggregate = AggregateContextProvider()
+    aggregate.providers.append(provider1)
+    aggregate.providers.append(provider2)
+
+    # Test invoking combines instructions
+    result = await aggregate.invoking([ChatMessage(role=Role.USER, text="Test")])
+
+    assert result.messages
+    assert isinstance(result.messages[0], ChatMessage)
+    assert isinstance(result.messages[1], ChatMessage)
+    assert result.messages[0].text == "First instruction"
+    assert result.messages[1].text == "Second instruction"
+
+
+async def test_chat_agent_context_providers_with_thread_service_id(chat_client_base: ChatClientProtocol) -> None:
+    """Test context providers with service-managed thread."""
+    mock_provider = MockContextProvider()
+    chat_client_base.run_responses = [
+        ChatResponse(
+            messages=[ChatMessage(role=Role.ASSISTANT, contents=[TextContent("test response")])],
+            conversation_id="service-thread-123",
+        )
+    ]
+
+    agent = ChatAgent(chat_client=chat_client_base, context_providers=mock_provider)
+
+    # Use existing service-managed thread
+    thread = agent.get_new_thread(service_thread_id="existing-thread-id")
+    await agent.run("Hello", thread=thread)
+
+    # invoked should be called with the service thread ID from response
+    assert mock_provider.invoked_called
+
+
+# Tests for as_tool method
+async def test_chat_agent_as_tool_basic(chat_client: ChatClientProtocol) -> None:
+    """Test basic as_tool functionality."""
+    agent = ChatAgent(chat_client=chat_client, name="TestAgent", description="Test agent for as_tool")
+
+    tool = agent.as_tool()
+
+    assert tool.name == "TestAgent"
+    assert tool.description == "Test agent for as_tool"
+    assert hasattr(tool, "func")
+    assert hasattr(tool, "input_model")
+
+
+async def test_chat_agent_as_tool_custom_parameters(chat_client: ChatClientProtocol) -> None:
+    """Test as_tool with custom parameters."""
+    agent = ChatAgent(chat_client=chat_client, name="TestAgent", description="Original description")
+
+    tool = agent.as_tool(
+        name="CustomTool",
+        description="Custom description",
+        arg_name="query",
+        arg_description="Custom input description",
+    )
+
+    assert tool.name == "CustomTool"
+    assert tool.description == "Custom description"
+
+    # Check that the input model has the custom field name
+    schema = tool.input_model.model_json_schema()
+    assert "query" in schema["properties"]
+    assert schema["properties"]["query"]["description"] == "Custom input description"
+
+
+async def test_chat_agent_as_tool_defaults(chat_client: ChatClientProtocol) -> None:
+    """Test as_tool with default parameters."""
+    agent = ChatAgent(
+        chat_client=chat_client,
+        name="TestAgent",
+        # No description provided
+    )
+
+    tool = agent.as_tool()
+
+    assert tool.name == "TestAgent"
+    assert tool.description == ""  # Should default to empty string
+
+    # Check default input field
+    schema = tool.input_model.model_json_schema()
+    assert "task" in schema["properties"]
+    assert "Task for TestAgent" in schema["properties"]["task"]["description"]
+
+
+async def test_chat_agent_as_tool_no_name(chat_client: ChatClientProtocol) -> None:
+    """Test as_tool when agent has no name (should raise ValueError)."""
+    agent = ChatAgent(chat_client=chat_client)  # No name provided
+
+    # Should raise ValueError since agent has no name
+    with raises(ValueError, match="Agent tool name cannot be None"):
+        agent.as_tool()
+
+
+async def test_chat_agent_as_tool_function_execution(chat_client: ChatClientProtocol) -> None:
+    """Test that the generated AIFunction can be executed."""
+    agent = ChatAgent(chat_client=chat_client, name="TestAgent", description="Test agent")
+
+    tool = agent.as_tool()
+
+    # Test function execution
+    result = await tool.invoke(arguments=tool.input_model(task="Hello"))
+
+    # Should return the agent's response text
+    assert isinstance(result, str)
+    assert result == "test response"  # From mock chat client
+
+
+async def test_chat_agent_as_tool_with_stream_callback(chat_client: ChatClientProtocol) -> None:
+    """Test as_tool with stream callback functionality."""
+    agent = ChatAgent(chat_client=chat_client, name="StreamingAgent")
+
+    # Collect streaming updates
+    collected_updates: list[AgentRunResponseUpdate] = []
+
+    def stream_callback(update: AgentRunResponseUpdate) -> None:
+        collected_updates.append(update)
+
+    tool = agent.as_tool(stream_callback=stream_callback)
+
+    # Execute the tool
+    result = await tool.invoke(arguments=tool.input_model(task="Hello"))
+
+    # Should have collected streaming updates
+    assert len(collected_updates) > 0
+    assert isinstance(result, str)
+    # Result should be concatenation of all streaming updates
+    expected_text = "".join(update.text for update in collected_updates)
+    assert result == expected_text
+
+
+async def test_chat_agent_as_tool_with_custom_arg_name(chat_client: ChatClientProtocol) -> None:
+    """Test as_tool with custom argument name."""
+    agent = ChatAgent(chat_client=chat_client, name="CustomArgAgent")
+
+    tool = agent.as_tool(arg_name="prompt", arg_description="Custom prompt input")
+
+    # Test that the custom argument name works
+    result = await tool.invoke(arguments=tool.input_model(prompt="Test prompt"))
+    assert result == "test response"
+
+
+async def test_chat_agent_as_tool_with_async_stream_callback(chat_client: ChatClientProtocol) -> None:
+    """Test as_tool with async stream callback functionality."""
+    agent = ChatAgent(chat_client=chat_client, name="AsyncStreamingAgent")
+
+    # Collect streaming updates using an async callback
+    collected_updates: list[AgentRunResponseUpdate] = []
+
+    async def async_stream_callback(update: AgentRunResponseUpdate) -> None:
+        collected_updates.append(update)
+
+    tool = agent.as_tool(stream_callback=async_stream_callback)
+
+    # Execute the tool
+    result = await tool.invoke(arguments=tool.input_model(task="Hello"))
+
+    # Should have collected streaming updates
+    assert len(collected_updates) > 0
+    assert isinstance(result, str)
+    # Result should be concatenation of all streaming updates
+    expected_text = "".join(update.text for update in collected_updates)
+    assert result == expected_text

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

@@ -0,0 +1,163 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import sys
+
+from agent_framework import (
+    BaseChatClient,
+    ChatClientProtocol,
+    ChatMessage,
+    ChatResponse,
+    ChatResponseUpdate,
+    FunctionCallContent,
+    FunctionResultContent,
+    Role,
+    TextContent,
+    ai_function,
+)
+
+if sys.version_info >= (3, 12):
+    pass  # type: ignore
+else:
+    pass  # type: ignore[import]
+
+
+def test_chat_client_type(chat_client: ChatClientProtocol):
+    assert isinstance(chat_client, ChatClientProtocol)
+
+
+async def test_chat_client_get_response(chat_client: ChatClientProtocol):
+    response = await chat_client.get_response(ChatMessage(role="user", text="Hello"))
+    assert response.text == "test response"
+    assert response.messages[0].role == Role.ASSISTANT
+
+
+async def test_chat_client_get_streaming_response(chat_client: ChatClientProtocol):
+    async for update in chat_client.get_streaming_response(ChatMessage(role="user", text="Hello")):
+        assert update.text == "test streaming response " or update.text == "another update"
+        assert update.role == Role.ASSISTANT
+
+
+def test_base_client(chat_client_base: ChatClientProtocol):
+    assert isinstance(chat_client_base, BaseChatClient)
+    assert isinstance(chat_client_base, ChatClientProtocol)
+
+
+async def test_base_client_get_response(chat_client_base: ChatClientProtocol):
+    response = await chat_client_base.get_response(ChatMessage(role="user", text="Hello"))
+    assert response.messages[0].role == Role.ASSISTANT
+    assert response.messages[0].text == "test response - Hello"
+
+
+async def test_base_client_get_streaming_response(chat_client_base: ChatClientProtocol):
+    async for update in chat_client_base.get_streaming_response(ChatMessage(role="user", text="Hello")):
+        assert update.text == "update - Hello" or update.text == "another update"
+
+
+async def test_base_client_with_function_calling(chat_client_base: ChatClientProtocol):
+    exec_counter = 0
+
+    @ai_function(name="test_function")
+    def ai_func(arg1: str) -> str:
+        nonlocal exec_counter
+        exec_counter += 1
+        return f"Processed {arg1}"
+
+    chat_client_base.run_responses = [
+        ChatResponse(
+            messages=ChatMessage(
+                role="assistant",
+                contents=[FunctionCallContent(call_id="1", name="test_function", arguments='{"arg1": "value1"}')],
+            )
+        ),
+        ChatResponse(messages=ChatMessage(role="assistant", text="done")),
+    ]
+    response = await chat_client_base.get_response("hello", tool_choice="auto", tools=[ai_func])
+    assert exec_counter == 1
+    assert len(response.messages) == 3
+    assert response.messages[0].role == Role.ASSISTANT
+    assert isinstance(response.messages[0].contents[0], FunctionCallContent)
+    assert response.messages[0].contents[0].name == "test_function"
+    assert response.messages[0].contents[0].arguments == '{"arg1": "value1"}'
+    assert response.messages[0].contents[0].call_id == "1"
+    assert response.messages[1].role == Role.TOOL
+    assert isinstance(response.messages[1].contents[0], FunctionResultContent)
+    assert response.messages[1].contents[0].call_id == "1"
+    assert response.messages[1].contents[0].result == "Processed value1"
+    assert response.messages[2].role == Role.ASSISTANT
+    assert response.messages[2].text == "done"
+
+
+async def test_base_client_with_function_calling_resets(chat_client_base: ChatClientProtocol):
+    exec_counter = 0
+
+    @ai_function(name="test_function")
+    def ai_func(arg1: str) -> str:
+        nonlocal exec_counter
+        exec_counter += 1
+        return f"Processed {arg1}"
+
+    chat_client_base.run_responses = [
+        ChatResponse(
+            messages=ChatMessage(
+                role="assistant",
+                contents=[FunctionCallContent(call_id="1", name="test_function", arguments='{"arg1": "value1"}')],
+            )
+        ),
+        ChatResponse(
+            messages=ChatMessage(
+                role="assistant",
+                contents=[FunctionCallContent(call_id="2", name="test_function", arguments='{"arg1": "value1"}')],
+            )
+        ),
+        ChatResponse(messages=ChatMessage(role="assistant", text="done")),
+    ]
+    response = await chat_client_base.get_response("hello", tool_choice="auto", tools=[ai_func])
+    assert exec_counter == 2
+    assert len(response.messages) == 5
+    assert response.messages[0].role == Role.ASSISTANT
+    assert response.messages[1].role == Role.TOOL
+    assert response.messages[2].role == Role.ASSISTANT
+    assert response.messages[3].role == Role.TOOL
+    assert response.messages[4].role == Role.ASSISTANT
+    assert isinstance(response.messages[0].contents[0], FunctionCallContent)
+    assert isinstance(response.messages[1].contents[0], FunctionResultContent)
+    assert isinstance(response.messages[2].contents[0], FunctionCallContent)
+    assert isinstance(response.messages[3].contents[0], FunctionResultContent)
+
+
+async def test_base_client_with_streaming_function_calling(chat_client_base: ChatClientProtocol):
+    exec_counter = 0
+
+    @ai_function(name="test_function")
+    def ai_func(arg1: str) -> str:
+        nonlocal exec_counter
+        exec_counter += 1
+        return f"Processed {arg1}"
+
+    chat_client_base.streaming_responses = [
+        [
+            ChatResponseUpdate(
+                contents=[FunctionCallContent(call_id="1", name="test_function", arguments='{"arg1":')],
+                role="assistant",
+            ),
+            ChatResponseUpdate(
+                contents=[FunctionCallContent(call_id="1", name="test_function", arguments='"value1"}')],
+                role="assistant",
+            ),
+        ],
+        [
+            ChatResponseUpdate(
+                contents=[TextContent(text="Processed value1")],
+                role="assistant",
+            )
+        ],
+    ]
+    updates = []
+    async for update in chat_client_base.get_streaming_response("hello", tool_choice="auto", tools=[ai_func]):
+        updates.append(update)
+    assert len(updates) == 4  # two updates with the function call, the function result and the final text
+    assert updates[0].contents[0].call_id == "1"
+    assert updates[1].contents[0].call_id == "1"
+    assert updates[2].contents[0].call_id == "1"
+    assert updates[3].text == "Processed value1"
+    assert exec_counter == 1

python/packages/core/tests/main/test_logging.py

@@ -0,0 +1,39 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+
+import pytest
+
+from agent_framework import get_logger
+from agent_framework.exceptions import AgentFrameworkException
+
+
+def test_get_logger():
+    """Test that the logger is created with the correct name."""
+    logger = get_logger()
+    assert logger.name == "agent_framework"
+
+
+def test_get_logger_custom_name():
+    """Test that the logger can be created with a custom name."""
+    custom_name = "agent_framework.custom"
+    logger = get_logger(custom_name)
+    assert logger.name == custom_name
+
+
+def test_get_logger_invalid_name():
+    """Test that an exception is raised for an invalid logger name."""
+    with pytest.raises(AgentFrameworkException):
+        get_logger("invalid_name")
+
+
+def test_log(caplog):
+    """Test that the logger can log messages and adheres to the expected format."""
+    logger = get_logger()
+    with caplog.at_level("DEBUG"):
+        logger.debug("This is a debug message")
+        assert len(caplog.records) == 1
+        record = caplog.records[0]
+        assert record.levelname == "DEBUG"
+        assert record.message == "This is a debug message"
+        assert record.name == "agent_framework"
+        assert record.pathname.endswith("test_logging.py")

python/packages/core/tests/main/test_mcp.py

@@ -0,0 +1,536 @@
+# Copyright (c) Microsoft. All rights reserved.
+# type: ignore[reportPrivateUsage]
+import os
+from contextlib import _AsyncGeneratorContextManager  # type: ignore
+from typing import Any
+from unittest.mock import AsyncMock, Mock
+
+import pytest
+from mcp import types
+from mcp.client.session import ClientSession
+from mcp.shared.exceptions import McpError
+from pydantic import AnyUrl, ValidationError
+
+from agent_framework import (
+    ChatMessage,
+    DataContent,
+    MCPStdioTool,
+    MCPStreamableHTTPTool,
+    MCPWebsocketTool,
+    Role,
+    TextContent,
+    ToolProtocol,
+    UriContent,
+)
+from agent_framework._mcp import (
+    MCPTool,
+    _ai_content_to_mcp_types,
+    _chat_message_to_mcp_types,
+    _get_input_model_from_mcp_prompt,
+    _get_input_model_from_mcp_tool,
+    _mcp_call_tool_result_to_ai_contents,
+    _mcp_prompt_message_to_chat_message,
+    _mcp_type_to_ai_content,
+    _normalize_mcp_name,
+)
+from agent_framework.exceptions import ToolExecutionException
+
+# Integration test skip condition
+skip_if_mcp_integration_tests_disabled = pytest.mark.skipif(
+    os.getenv("RUN_INTEGRATION_TESTS", "false").lower() != "true" or os.getenv("LOCAL_MCP_URL", "") == "",
+    reason="No LOCAL_MCP_URL provided; skipping integration tests."
+    if os.getenv("RUN_INTEGRATION_TESTS", "false").lower() == "true"
+    else "Integration tests are disabled.",
+)
+
+
+# Helper function tests
+def test_normalize_mcp_name():
+    """Test MCP name normalization."""
+    assert _normalize_mcp_name("valid_name") == "valid_name"
+    assert _normalize_mcp_name("name-with-dashes") == "name-with-dashes"
+    assert _normalize_mcp_name("name.with.dots") == "name.with.dots"
+    assert _normalize_mcp_name("name with spaces") == "name-with-spaces"
+    assert _normalize_mcp_name("name@with#special$chars") == "name-with-special-chars"
+    assert _normalize_mcp_name("name/with\\slashes") == "name-with-slashes"
+
+
+def test_mcp_prompt_message_to_ai_content():
+    """Test conversion from MCP prompt message to AI content."""
+    mcp_message = types.PromptMessage(role="user", content=types.TextContent(type="text", text="Hello, world!"))
+    ai_content = _mcp_prompt_message_to_chat_message(mcp_message)
+
+    assert isinstance(ai_content, ChatMessage)
+    assert ai_content.role.value == "user"
+    assert len(ai_content.contents) == 1
+    assert isinstance(ai_content.contents[0], TextContent)
+    assert ai_content.contents[0].text == "Hello, world!"
+    assert ai_content.raw_representation == mcp_message
+
+
+def test_mcp_call_tool_result_to_ai_contents():
+    """Test conversion from MCP tool result to AI contents."""
+    mcp_result = types.CallToolResult(
+        content=[
+            types.TextContent(type="text", text="Result text"),
+            types.ImageContent(type="image", data="data:image/png;base64,xyz", mimeType="image/png"),
+        ]
+    )
+    ai_contents = _mcp_call_tool_result_to_ai_contents(mcp_result)
+
+    assert len(ai_contents) == 2
+    assert isinstance(ai_contents[0], TextContent)
+    assert ai_contents[0].text == "Result text"
+    assert isinstance(ai_contents[1], DataContent)
+    assert ai_contents[1].uri == "data:image/png;base64,xyz"
+    assert ai_contents[1].media_type == "image/png"
+
+
+def test_mcp_content_types_to_ai_content_text():
+    """Test conversion of MCP text content to AI content."""
+    mcp_content = types.TextContent(type="text", text="Sample text")
+    ai_content = _mcp_type_to_ai_content(mcp_content)
+
+    assert isinstance(ai_content, TextContent)
+    assert ai_content.text == "Sample text"
+    assert ai_content.raw_representation == mcp_content
+
+
+def test_mcp_content_types_to_ai_content_image():
+    """Test conversion of MCP image content to AI content."""
+    mcp_content = types.ImageContent(type="image", data="data:image/jpeg;base64,abc", mimeType="image/jpeg")
+    ai_content = _mcp_type_to_ai_content(mcp_content)
+
+    assert isinstance(ai_content, DataContent)
+    assert ai_content.uri == "data:image/jpeg;base64,abc"
+    assert ai_content.media_type == "image/jpeg"
+    assert ai_content.raw_representation == mcp_content
+
+
+def test_mcp_content_types_to_ai_content_audio():
+    """Test conversion of MCP audio content to AI content."""
+    mcp_content = types.AudioContent(type="audio", data="data:audio/wav;base64,def", mimeType="audio/wav")
+    ai_content = _mcp_type_to_ai_content(mcp_content)
+
+    assert isinstance(ai_content, DataContent)
+    assert ai_content.uri == "data:audio/wav;base64,def"
+    assert ai_content.media_type == "audio/wav"
+    assert ai_content.raw_representation == mcp_content
+
+
+def test_mcp_content_types_to_ai_content_resource_link():
+    """Test conversion of MCP resource link to AI content."""
+    mcp_content = types.ResourceLink(
+        type="resource_link",
+        uri=AnyUrl("https://example.com/resource"),
+        name="test_resource",
+        mimeType="application/json",
+    )
+    ai_content = _mcp_type_to_ai_content(mcp_content)
+
+    assert isinstance(ai_content, UriContent)
+    assert ai_content.uri == "https://example.com/resource"
+    assert ai_content.media_type == "application/json"
+    assert ai_content.raw_representation == mcp_content
+
+
+def test_mcp_content_types_to_ai_content_embedded_resource_text():
+    """Test conversion of MCP embedded text resource to AI content."""
+    text_resource = types.TextResourceContents(
+        uri=AnyUrl("file://test.txt"), mimeType="text/plain", text="Embedded text content"
+    )
+    mcp_content = types.EmbeddedResource(type="resource", resource=text_resource)
+    ai_content = _mcp_type_to_ai_content(mcp_content)
+
+    assert isinstance(ai_content, TextContent)
+    assert ai_content.text == "Embedded text content"
+    assert ai_content.raw_representation == mcp_content
+
+
+def test_mcp_content_types_to_ai_content_embedded_resource_blob():
+    """Test conversion of MCP embedded blob resource to AI content."""
+    # Use a proper data URI in the blob field since that's what the MCP implementation expects
+    blob_resource = types.BlobResourceContents(
+        uri=AnyUrl("file://test.bin"),
+        mimeType="application/octet-stream",
+        blob="data:application/octet-stream;base64,dGVzdCBkYXRh",
+    )
+    mcp_content = types.EmbeddedResource(type="resource", resource=blob_resource)
+    ai_content = _mcp_type_to_ai_content(mcp_content)
+
+    assert isinstance(ai_content, DataContent)
+    assert ai_content.uri == "data:application/octet-stream;base64,dGVzdCBkYXRh"
+    assert ai_content.media_type == "application/octet-stream"
+    assert ai_content.raw_representation == mcp_content
+
+
+def test_ai_content_to_mcp_content_types_text():
+    """Test conversion of AI text content to MCP content."""
+    ai_content = TextContent(text="Sample text")
+    mcp_content = _ai_content_to_mcp_types(ai_content)
+
+    assert isinstance(mcp_content, types.TextContent)
+    assert mcp_content.type == "text"
+    assert mcp_content.text == "Sample text"
+
+
+def test_ai_content_to_mcp_content_types_data_image():
+    """Test conversion of AI data content to MCP content."""
+    ai_content = DataContent(uri="data:image/png;base64,xyz", media_type="image/png")
+    mcp_content = _ai_content_to_mcp_types(ai_content)
+
+    assert isinstance(mcp_content, types.ImageContent)
+    assert mcp_content.type == "image"
+    assert mcp_content.data == "data:image/png;base64,xyz"
+    assert mcp_content.mimeType == "image/png"
+
+
+def test_ai_content_to_mcp_content_types_data_audio():
+    """Test conversion of AI data content to MCP content."""
+    ai_content = DataContent(uri="data:audio/mpeg;base64,xyz", media_type="audio/mpeg")
+    mcp_content = _ai_content_to_mcp_types(ai_content)
+
+    assert isinstance(mcp_content, types.AudioContent)
+    assert mcp_content.type == "audio"
+    assert mcp_content.data == "data:audio/mpeg;base64,xyz"
+    assert mcp_content.mimeType == "audio/mpeg"
+
+
+def test_ai_content_to_mcp_content_types_data_binary():
+    """Test conversion of AI data content to MCP content."""
+    ai_content = DataContent(uri="data:application/octet-stream;base64,xyz", media_type="application/octet-stream")
+    mcp_content = _ai_content_to_mcp_types(ai_content)
+
+    assert isinstance(mcp_content, types.EmbeddedResource)
+    assert mcp_content.type == "resource"
+    assert mcp_content.resource.blob == "data:application/octet-stream;base64,xyz"
+    assert mcp_content.resource.mimeType == "application/octet-stream"
+
+
+def test_ai_content_to_mcp_content_types_uri():
+    """Test conversion of AI URI content to MCP content."""
+    ai_content = UriContent(uri="https://example.com/resource", media_type="application/json")
+    mcp_content = _ai_content_to_mcp_types(ai_content)
+
+    assert isinstance(mcp_content, types.ResourceLink)
+    assert mcp_content.type == "resource_link"
+    assert str(mcp_content.uri) == "https://example.com/resource"
+    assert mcp_content.mimeType == "application/json"
+
+
+def test_chat_message_to_mcp_types():
+    message = ChatMessage(
+        role="user",
+        contents=[TextContent(text="test"), DataContent(uri="data:image/png;base64,xyz", media_type="image/png")],
+    )
+    mcp_contents = _chat_message_to_mcp_types(message)
+    assert len(mcp_contents) == 2
+    assert isinstance(mcp_contents[0], types.TextContent)
+    assert isinstance(mcp_contents[1], types.ImageContent)
+
+
+def test_get_input_model_from_mcp_tool():
+    """Test creation of input model from MCP tool."""
+    tool = types.Tool(
+        name="test_tool",
+        description="A test tool",
+        inputSchema={
+            "type": "object",
+            "properties": {"param1": {"type": "string"}, "param2": {"type": "number"}},
+            "required": ["param1"],
+        },
+    )
+    model = _get_input_model_from_mcp_tool(tool)
+
+    # Create an instance to verify the model works
+    instance = model(param1="test", param2=42)
+    assert instance.param1 == "test"
+    assert instance.param2 == 42
+
+    # Test validation
+    with pytest.raises(ValidationError):  # Missing required param1
+        model(param2=42)
+
+
+def test_get_input_model_from_mcp_prompt():
+    """Test creation of input model from MCP prompt."""
+    prompt = types.Prompt(
+        name="test_prompt",
+        description="A test prompt",
+        arguments=[
+            types.PromptArgument(name="arg1", description="First argument", required=True),
+            types.PromptArgument(name="arg2", description="Second argument", required=False),
+        ],
+    )
+    model = _get_input_model_from_mcp_prompt(prompt)
+
+    # Create an instance to verify the model works
+    instance = model(arg1="test", arg2="optional")
+    assert instance.arg1 == "test"
+    assert instance.arg2 == "optional"
+
+    # Test validation
+    with pytest.raises(ValidationError):  # Missing required arg1
+        model(arg2="optional")
+
+
+# MCPTool tests
+async def test_local_mcp_server_initialization():
+    """Test MCPTool initialization."""
+    server = MCPTool(name="test_server")
+    assert isinstance(server, ToolProtocol)
+    assert server.name == "test_server"
+    assert server.session is None
+    assert server.functions == []
+
+
+async def test_local_mcp_server_context_manager():
+    """Test MCPTool as context manager."""
+
+    class TestServer(MCPTool):
+        async def connect(self):
+            # Mock connection
+            self.session = Mock(spec=ClientSession)
+
+        def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
+            return None
+
+    server = TestServer(name="test_server")
+    async with server:
+        assert server.session is not None
+
+    assert server.session is None
+
+
+async def test_local_mcp_server_load_functions():
+    """Test loading functions from MCP server."""
+
+    class TestServer(MCPTool):
+        async def connect(self):
+            self.session = Mock(spec=ClientSession)
+            # Mock tools list response
+            self.session.list_tools = AsyncMock(
+                return_value=types.ListToolsResult(
+                    tools=[
+                        types.Tool(
+                            name="test_tool",
+                            description="Test tool",
+                            inputSchema={
+                                "type": "object",
+                                "properties": {"param": {"type": "string"}},
+                                "required": ["param"],
+                            },
+                        )
+                    ]
+                )
+            )
+
+        def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
+            return None
+
+    server = TestServer(name="test_server")
+    assert isinstance(server, ToolProtocol)
+    async with server:
+        await server.load_tools()
+        assert len(server.functions) == 1
+        assert server.functions[0].name == "test_tool"
+
+
+async def test_local_mcp_server_load_prompts():
+    """Test loading prompts from MCP server."""
+
+    class TestServer(MCPTool):
+        async def connect(self):
+            self.session = Mock(spec=ClientSession)
+            # Mock prompts list response
+            self.session.list_prompts = AsyncMock(
+                return_value=types.ListPromptsResult(
+                    prompts=[
+                        types.Prompt(
+                            name="test_prompt",
+                            description="Test prompt",
+                            arguments=[types.PromptArgument(name="arg", description="Test arg", required=True)],
+                        )
+                    ]
+                )
+            )
+
+        def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
+            return None
+
+    server = TestServer(name="test_server")
+    async with server:
+        await server.load_prompts()
+        assert len(server.functions) == 1
+        assert server.functions[0].name == "test_prompt"
+
+
+async def test_local_mcp_server_function_execution():
+    """Test function execution through MCP server."""
+
+    class TestServer(MCPTool):
+        async def connect(self):
+            self.session = Mock(spec=ClientSession)
+            self.session.list_tools = AsyncMock(
+                return_value=types.ListToolsResult(
+                    tools=[
+                        types.Tool(
+                            name="test_tool",
+                            description="Test tool",
+                            inputSchema={
+                                "type": "object",
+                                "properties": {"param": {"type": "string"}},
+                                "required": ["param"],
+                            },
+                        )
+                    ]
+                )
+            )
+            self.session.call_tool = AsyncMock(
+                return_value=types.CallToolResult(
+                    content=[types.TextContent(type="text", text="Tool executed successfully")]
+                )
+            )
+
+        def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
+            return None
+
+    server = TestServer(name="test_server")
+    async with server:
+        await server.load_tools()
+        func = server.functions[0]
+        result = await func.invoke(param="test_value")
+
+        assert len(result) == 1
+        assert isinstance(result[0], TextContent)
+        assert result[0].text == "Tool executed successfully"
+
+
+async def test_local_mcp_server_function_execution_error():
+    """Test function execution error handling."""
+
+    class TestServer(MCPTool):
+        async def connect(self):
+            self.session = Mock(spec=ClientSession)
+            self.session.list_tools = AsyncMock(
+                return_value=types.ListToolsResult(
+                    tools=[
+                        types.Tool(
+                            name="test_tool",
+                            description="Test tool",
+                            inputSchema={
+                                "type": "object",
+                                "properties": {"param": {"type": "string"}},
+                                "required": ["param"],
+                            },
+                        )
+                    ]
+                )
+            )
+            # Mock a tool call that raises an MCP error
+            self.session.call_tool = AsyncMock(
+                side_effect=McpError(types.ErrorData(code=-1, message="Tool execution failed"))
+            )
+
+        def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
+            return None
+
+    server = TestServer(name="test_server")
+    async with server:
+        await server.load_tools()
+        func = server.functions[0]
+
+        with pytest.raises(ToolExecutionException):
+            await func.invoke(param="test_value")
+
+
+async def test_local_mcp_server_prompt_execution():
+    """Test prompt execution through MCP server."""
+
+    class TestMCPTool(MCPTool):
+        async def connect(self):
+            self.session = Mock(spec=ClientSession)
+            self.session.list_prompts = AsyncMock(
+                return_value=types.ListPromptsResult(
+                    prompts=[
+                        types.Prompt(
+                            name="test_prompt",
+                            description="Test prompt",
+                            arguments=[types.PromptArgument(name="arg", description="Test arg", required=True)],
+                        )
+                    ]
+                )
+            )
+            self.session.get_prompt = AsyncMock(
+                return_value=types.GetPromptResult(
+                    description="Generated prompt",
+                    messages=[
+                        types.PromptMessage(role="user", content=types.TextContent(type="text", text="Test message"))
+                    ],
+                )
+            )
+
+        def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
+            return None
+
+    server = TestMCPTool(name="test_server")
+    async with server:
+        await server.load_prompts()
+        prompt = server.functions[0]
+        result = await prompt.invoke(arg="test_value")
+
+        assert len(result) == 1
+        assert isinstance(result[0], ChatMessage)
+        assert result[0].role == Role.USER
+        assert len(result[0].contents) == 1
+        assert result[0].contents[0].text == "Test message"
+
+
+# Server implementation tests
+def test_local_mcp_stdio_tool_init():
+    """Test MCPStdioTool initialization."""
+    tool = MCPStdioTool(name="test", command="echo", args=["hello"])
+    assert tool.name == "test"
+    assert tool.command == "echo"
+    assert tool.args == ["hello"]
+
+
+def test_local_mcp_websocket_tool_init():
+    """Test MCPWebsocketTool initialization."""
+    tool = MCPWebsocketTool(name="test", url="ws://localhost:8080")
+    assert tool.name == "test"
+    assert tool.url == "ws://localhost:8080"
+
+
+def test_local_mcp_streamable_http_tool_init():
+    """Test MCPStreamableHTTPTool initialization."""
+    tool = MCPStreamableHTTPTool(name="test", url="http://localhost:8080")
+    assert tool.name == "test"
+    assert tool.url == "http://localhost:8080"
+
+
+# Integration test
+@skip_if_mcp_integration_tests_disabled
+async def test_streamable_http_integration():
+    """Test MCP StreamableHTTP integration."""
+    url = os.environ.get("LOCAL_MCP_URL", "")
+    if not url.startswith("http"):
+        pytest.skip("LOCAL_MCP_URL is not an HTTP URL")
+
+    tool = MCPStreamableHTTPTool(name="integration_test", url=url)
+
+    async with tool:
+        # Test that we can connect and load tools
+        assert tool.session is not None
+        assert isinstance(tool.functions, list)
+
+        # If there are functions available, try to get information about one
+        assert tool.functions, "The MCP server should have at least one function."
+
+        func = tool.functions[0]
+
+        assert hasattr(func, "name")
+        assert hasattr(func, "description")
+
+        result = await func.invoke(query="What is Agent Framework?")
+        assert result[0].text is not None

python/packages/core/tests/main/test_memory.py

@@ -0,0 +1,296 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from collections.abc import MutableSequence
+from typing import Any
+from unittest.mock import AsyncMock, Mock
+
+from agent_framework import ChatMessage, Role, TextContent
+from agent_framework._memory import AggregateContextProvider, Context, ContextProvider
+
+
+class MockContextProvider(ContextProvider):
+    """Mock ContextProvider for testing."""
+
+    def __init__(self, messages: list[ChatMessage] | None = None) -> None:
+        self.context_messages = messages
+        self.thread_created_called = False
+        self.invoked_called = False
+        self.invoking_called = False
+        self.thread_created_thread_id = None
+        self.new_messages = None
+        self.model_invoking_messages = None
+
+    async def thread_created(self, thread_id: str | None) -> None:
+        """Track thread_created calls."""
+        self.thread_created_called = True
+        self.thread_created_thread_id = thread_id
+
+    async def invoked(
+        self,
+        request_messages: Any,
+        response_messages: Any | None = None,
+        invoke_exception: Exception | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Track invoked calls."""
+        self.invoked_called = True
+        self.new_messages = request_messages
+
+    async def invoking(self, messages: ChatMessage | MutableSequence[ChatMessage], **kwargs: Any) -> Context:
+        """Track invoking calls and return context."""
+        self.invoking_called = True
+        self.model_invoking_messages = messages
+        context = Context()
+        context.messages = self.context_messages
+        return context
+
+
+class TestAggregateContextProvider:
+    """Tests for AggregateContextProvider class."""
+
+    def test_init_with_no_providers(self) -> None:
+        """Test initialization with no providers."""
+        aggregate = AggregateContextProvider()
+        assert aggregate.providers == []
+
+    def test_init_with_none_providers(self) -> None:
+        """Test initialization with None providers."""
+        aggregate = AggregateContextProvider(None)
+        assert aggregate.providers == []
+
+    def test_init_with_providers(self) -> None:
+        """Test initialization with providers."""
+        provider1 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 1")])
+        provider2 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 2")])
+        provider3 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 3")])
+        providers = [provider1, provider2, provider3]
+
+        aggregate = AggregateContextProvider(providers)
+        assert len(aggregate.providers) == 3
+        assert aggregate.providers[0] is provider1
+        assert aggregate.providers[1] is provider2
+        assert aggregate.providers[2] is provider3
+
+    def test_add_provider(self) -> None:
+        """Test adding a provider."""
+        aggregate = AggregateContextProvider()
+        provider = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions")])
+
+        aggregate.add(provider)
+        assert len(aggregate.providers) == 1
+        assert aggregate.providers[0] is provider
+
+    def test_add_multiple_providers(self) -> None:
+        """Test adding multiple providers."""
+        aggregate = AggregateContextProvider()
+        provider1 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 1")])
+        provider2 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 2")])
+
+        aggregate.add(provider1)
+        aggregate.add(provider2)
+
+        assert len(aggregate.providers) == 2
+        assert aggregate.providers[0] is provider1
+        assert aggregate.providers[1] is provider2
+
+    async def test_thread_created_with_no_providers(self) -> None:
+        """Test thread_created with no providers."""
+        aggregate = AggregateContextProvider()
+
+        # Should not raise an exception
+        await aggregate.thread_created("thread-123")
+
+    async def test_thread_created_with_providers(self) -> None:
+        """Test thread_created calls all providers."""
+        provider1 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 1")])
+        provider2 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 2")])
+        aggregate = AggregateContextProvider([provider1, provider2])
+
+        thread_id = "thread-123"
+        await aggregate.thread_created(thread_id)
+
+        assert provider1.thread_created_called
+        assert provider1.thread_created_thread_id == thread_id
+        assert provider2.thread_created_called
+        assert provider2.thread_created_thread_id == thread_id
+
+    async def test_thread_created_with_none_thread_id(self) -> None:
+        """Test thread_created with None thread_id."""
+        provider = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions")])
+        aggregate = AggregateContextProvider([provider])
+
+        await aggregate.thread_created(None)
+
+        assert provider.thread_created_called
+        assert provider.thread_created_thread_id is None
+
+    async def test_messages_adding_with_no_providers(self) -> None:
+        """Test invoked with no providers."""
+        aggregate = AggregateContextProvider()
+        message = ChatMessage(text="Hello", role=Role.USER)
+
+        # Should not raise an exception
+        await aggregate.invoked(message)
+
+    async def test_messages_adding_with_single_message(self) -> None:
+        """Test invoked with a single message."""
+        provider1 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 1")])
+        provider2 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 2")])
+        aggregate = AggregateContextProvider([provider1, provider2])
+
+        message = ChatMessage(text="Hello", role=Role.USER)
+        await aggregate.invoked(message)
+
+        assert provider1.invoked_called
+        assert provider1.new_messages == message
+        assert provider2.invoked_called
+        assert provider2.new_messages == message
+
+    async def test_messages_adding_with_message_sequence(self) -> None:
+        """Test invoked with a sequence of messages."""
+        provider = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions")])
+        aggregate = AggregateContextProvider([provider])
+
+        messages = [
+            ChatMessage(text="Hello", role=Role.USER),
+            ChatMessage(text="Hi there", role=Role.ASSISTANT),
+        ]
+        await aggregate.invoked(messages)
+
+        assert provider.invoked_called
+        assert provider.new_messages == messages
+
+    async def test_model_invoking_with_no_providers(self) -> None:
+        """Test invoking with no providers."""
+        aggregate = AggregateContextProvider()
+        message = ChatMessage(text="Hello", role=Role.USER)
+
+        context = await aggregate.invoking(message)
+
+        assert isinstance(context, Context)
+        assert not context.messages
+
+    async def test_model_invoking_with_single_provider(self) -> None:
+        """Test invoking with a single provider."""
+        provider = MockContextProvider(messages=[ChatMessage(role="user", text="Test instructions")])
+        aggregate = AggregateContextProvider([provider])
+
+        message = [ChatMessage(text="Hello", role=Role.USER)]
+        context = await aggregate.invoking(message)
+
+        assert provider.invoking_called
+        assert provider.model_invoking_messages == message
+        assert isinstance(context, Context)
+
+        assert context.messages
+        assert isinstance(context.messages[0].contents[0], TextContent)
+        assert context.messages[0].text == "Test instructions"
+
+    async def test_model_invoking_with_multiple_providers(self) -> None:
+        """Test invoking combines contexts from multiple providers."""
+        provider1 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 1")])
+        provider2 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 2")])
+        provider3 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 3")])
+        aggregate = AggregateContextProvider([provider1, provider2, provider3])
+
+        messages = [ChatMessage(text="Hello", role=Role.USER)]
+        context = await aggregate.invoking(messages)
+
+        assert provider1.invoking_called
+        assert provider1.model_invoking_messages == messages
+        assert provider2.invoking_called
+        assert provider2.model_invoking_messages == messages
+        assert provider3.invoking_called
+        assert provider3.model_invoking_messages == messages
+
+        assert isinstance(context, Context)
+
+        assert context.messages
+        assert isinstance(context.messages[0].contents[0], TextContent)
+        assert isinstance(context.messages[1].contents[0], TextContent)
+        assert isinstance(context.messages[2].contents[0], TextContent)
+        assert context.messages[0].text == "Instructions 1"
+        assert context.messages[1].text == "Instructions 2"
+        assert context.messages[2].text == "Instructions 3"
+
+    async def test_model_invoking_with_none_instructions(self) -> None:
+        """Test invoking filters out None instructions."""
+        provider1 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 1")])
+        provider2 = MockContextProvider(messages=None)  # None instructions
+        provider3 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 3")])
+        aggregate = AggregateContextProvider([provider1, provider2, provider3])
+
+        message = ChatMessage(text="Hello", role=Role.USER)
+        context = await aggregate.invoking(message)
+
+        assert isinstance(context, Context)
+        assert context.messages
+        assert isinstance(context.messages[0].contents[0], TextContent)
+        assert isinstance(context.messages[1].contents[0], TextContent)
+        assert context.messages[0].text == "Instructions 1"
+        assert context.messages[1].text == "Instructions 3"
+
+    async def test_model_invoking_with_all_none_instructions(self) -> None:
+        """Test invoking when all providers return None instructions."""
+        provider1 = MockContextProvider(None)
+        provider2 = MockContextProvider(None)
+        aggregate = AggregateContextProvider([provider1, provider2])
+
+        message = ChatMessage(text="Hello", role=Role.USER)
+        context = await aggregate.invoking(message)
+
+        assert isinstance(context, Context)
+        assert not context.messages
+
+    async def test_model_invoking_with_mutable_sequence(self) -> None:
+        """Test invoking with MutableSequence of messages."""
+        provider = MockContextProvider(messages=[ChatMessage(role="user", text="Test instructions")])
+        aggregate = AggregateContextProvider([provider])
+
+        messages = [ChatMessage(text="Hello", role=Role.USER)]
+        context = await aggregate.invoking(messages)
+
+        assert provider.invoking_called
+        assert provider.model_invoking_messages == messages
+        assert isinstance(context, Context)
+        assert context.messages
+        assert isinstance(context.messages[0].contents[0], TextContent)
+        assert context.messages[0].text == "Test instructions"
+
+    async def test_async_methods_concurrent_execution(self) -> None:
+        """Test that async methods execute providers concurrently."""
+        # Use AsyncMock to verify concurrent execution
+        provider1 = Mock(spec=ContextProvider)
+        provider1.thread_created = AsyncMock()
+        provider1.invoked = AsyncMock()
+        provider1.invoking = AsyncMock(return_value=Context(messages=[ChatMessage(role="user", text="Test 1")]))
+
+        provider2 = Mock(spec=ContextProvider)
+        provider2.thread_created = AsyncMock()
+        provider2.invoked = AsyncMock()
+        provider2.invoking = AsyncMock(return_value=Context(messages=[ChatMessage(role="user", text="Test 2")]))
+
+        aggregate = AggregateContextProvider([provider1, provider2])
+
+        # Test thread_created
+        await aggregate.thread_created("thread-123")
+        provider1.thread_created.assert_called_once_with("thread-123")
+        provider2.thread_created.assert_called_once_with("thread-123")
+
+        # Test invoked
+        message = ChatMessage(text="Hello", role=Role.USER)
+        await aggregate.invoked(message)
+        provider1.invoked.assert_called_once_with(
+            request_messages=message, response_messages=None, invoke_exception=None
+        )
+        provider2.invoked.assert_called_once_with(
+            request_messages=message, response_messages=None, invoke_exception=None
+        )
+
+        # Test invoking
+        context = await aggregate.invoking(message)
+        provider1.invoking.assert_called_once_with(message)
+        provider2.invoking.assert_called_once_with(message)
+        assert context.messages
+        assert context.messages[0].text == "Test 1"
+        assert context.messages[1].text == "Test 2"

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

@@ -0,0 +1,463 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from collections.abc import AsyncIterable, Awaitable, Callable
+from typing import Any
+from unittest.mock import MagicMock
+
+import pytest
+from pydantic import BaseModel, Field
+
+from agent_framework import (
+    AgentProtocol,
+    AgentRunResponse,
+    AgentRunResponseUpdate,
+    ChatAgent,
+    ChatMessage,
+    Role,
+    TextContent,
+)
+from agent_framework._middleware import (
+    AgentMiddleware,
+    AgentMiddlewarePipeline,
+    AgentRunContext,
+    FunctionInvocationContext,
+    FunctionMiddleware,
+    FunctionMiddlewarePipeline,
+)
+from agent_framework._tools import AIFunction
+
+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: AgentProtocol) -> None:
+        """Test that agent middleware can override response for non-streaming execution."""
+        override_response = AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="overridden response")])
+
+        class ResponseOverrideMiddleware(AgentMiddleware):
+            async def process(
+                self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
+            ) -> None:
+                # Execute the pipeline first, then override the response
+                await next(context)
+                context.result = override_response
+
+        middleware = ResponseOverrideMiddleware()
+        pipeline = AgentMiddlewarePipeline([middleware])
+        messages = [ChatMessage(role=Role.USER, text="test")]
+        context = AgentRunContext(agent=mock_agent, messages=messages)
+
+        handler_called = False
+
+        async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
+            nonlocal handler_called
+            handler_called = True
+            return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="original response")])
+
+        result = await pipeline.execute(mock_agent, messages, 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"
+        # Verify original handler was called since middleware called next()
+        assert handler_called
+
+    async def test_agent_middleware_response_override_streaming(self, mock_agent: AgentProtocol) -> None:
+        """Test that agent middleware can override response for streaming execution."""
+
+        async def override_stream() -> AsyncIterable[AgentRunResponseUpdate]:
+            yield AgentRunResponseUpdate(contents=[TextContent(text="overridden")])
+            yield AgentRunResponseUpdate(contents=[TextContent(text=" stream")])
+
+        class StreamResponseOverrideMiddleware(AgentMiddleware):
+            async def process(
+                self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
+            ) -> None:
+                # Execute the pipeline first, then override the response stream
+                await next(context)
+                context.result = override_stream()
+
+        middleware = StreamResponseOverrideMiddleware()
+        pipeline = AgentMiddlewarePipeline([middleware])
+        messages = [ChatMessage(role=Role.USER, text="test")]
+        context = AgentRunContext(agent=mock_agent, messages=messages)
+
+        async def final_handler(ctx: AgentRunContext) -> AsyncIterable[AgentRunResponseUpdate]:
+            yield AgentRunResponseUpdate(contents=[TextContent(text="original")])
+
+        updates: list[AgentRunResponseUpdate] = []
+        async for update in pipeline.execute_stream(mock_agent, messages, context, final_handler):
+            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: AIFunction[Any, Any]) -> None:
+        """Test that function middleware can override result."""
+        override_result = "overridden function result"
+
+        class ResultOverrideMiddleware(FunctionMiddleware):
+            async def process(
+                self,
+                context: FunctionInvocationContext,
+                next: Callable[[FunctionInvocationContext], Awaitable[None]],
+            ) -> None:
+                # Execute the pipeline first, then override the result
+                await next(context)
+                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(mock_function, arguments, 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 ChatAgent integration."""
+        mock_chat_client = MockChatClient()
+
+        class ChatAgentResponseOverrideMiddleware(AgentMiddleware):
+            async def process(
+                self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
+            ) -> None:
+                # Always call next() first to allow execution
+                await next(context)
+                # Then conditionally override based on content
+                if any("special" in msg.text for msg in context.messages if msg.text):
+                    context.result = AgentRunResponse(
+                        messages=[ChatMessage(role=Role.ASSISTANT, text="Special response from middleware!")]
+                    )
+
+        # Create ChatAgent with override middleware
+        middleware = ChatAgentResponseOverrideMiddleware()
+        agent = ChatAgent(chat_client=mock_chat_client, middleware=[middleware])
+
+        # Test override case
+        override_messages = [ChatMessage(role=Role.USER, text="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 = [ChatMessage(role=Role.USER, text="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 ChatAgent integration."""
+        mock_chat_client = MockChatClient()
+
+        async def custom_stream() -> AsyncIterable[AgentRunResponseUpdate]:
+            yield AgentRunResponseUpdate(contents=[TextContent(text="Custom")])
+            yield AgentRunResponseUpdate(contents=[TextContent(text=" streaming")])
+            yield AgentRunResponseUpdate(contents=[TextContent(text=" response!")])
+
+        class ChatAgentStreamOverrideMiddleware(AgentMiddleware):
+            async def process(
+                self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
+            ) -> None:
+                # Always call next() first to allow execution
+                await next(context)
+                # Then conditionally override based on content
+                if any("custom stream" in msg.text for msg in context.messages if msg.text):
+                    context.result = custom_stream()
+
+        # Create ChatAgent with override middleware
+        middleware = ChatAgentStreamOverrideMiddleware()
+        agent = ChatAgent(chat_client=mock_chat_client, middleware=[middleware])
+
+        # Test streaming override case
+        override_messages = [ChatMessage(role=Role.USER, text="Give me a custom stream")]
+        override_updates: list[AgentRunResponseUpdate] = []
+        async for update in agent.run_stream(override_messages):
+            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 = [ChatMessage(role=Role.USER, text="Normal streaming request")]
+        normal_updates: list[AgentRunResponseUpdate] = []
+        async for update in agent.run_stream(normal_messages):
+            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: AgentProtocol) -> None:
+        """Test that when agent middleware conditionally doesn't call next(), no execution happens."""
+
+        class ConditionalNoNextMiddleware(AgentMiddleware):
+            async def process(
+                self, context: AgentRunContext, next: Callable[[AgentRunContext], 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 next(context)
+                # Otherwise, don't call next() - no execution should happen
+
+        middleware = ConditionalNoNextMiddleware()
+        pipeline = AgentMiddlewarePipeline([middleware])
+
+        handler_called = False
+
+        async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
+            nonlocal handler_called
+            handler_called = True
+            return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="executed response")])
+
+        # Test case where next() is NOT called
+        no_execute_messages = [ChatMessage(role=Role.USER, text="Don't run this")]
+        no_execute_context = AgentRunContext(agent=mock_agent, messages=no_execute_messages)
+        no_execute_result = await pipeline.execute(mock_agent, no_execute_messages, no_execute_context, final_handler)
+
+        # When middleware doesn't call next(), result should be empty AgentRunResponse
+        assert no_execute_result is not None
+        assert isinstance(no_execute_result, AgentRunResponse)
+        assert no_execute_result.messages == []  # Empty response
+        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_messages = [ChatMessage(role=Role.USER, text="Please execute this")]
+        execute_context = AgentRunContext(agent=mock_agent, messages=execute_messages)
+        execute_result = await pipeline.execute(mock_agent, execute_messages, execute_context, final_handler)
+
+        assert execute_result is not None
+        assert execute_result.messages[0].text == "executed response"
+        assert handler_called
+
+    async def test_function_middleware_conditional_no_next(self, mock_function: AIFunction[Any, Any]) -> None:
+        """Test that when function middleware conditionally doesn't call next(), no execution happens."""
+
+        class ConditionalNoNextFunctionMiddleware(FunctionMiddleware):
+            async def process(
+                self,
+                context: FunctionInvocationContext,
+                next: Callable[[FunctionInvocationContext], Awaitable[None]],
+            ) -> None:
+                # Only call next() if argument name contains "execute"
+                args = context.arguments
+                assert isinstance(args, FunctionTestArgs)
+                if "execute" in args.name:
+                    await next(context)
+                # 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(mock_function, no_execute_args, 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(mock_function, execute_args, 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: AgentProtocol) -> None:
+        """Test that middleware can observe response after execution."""
+        observed_responses: list[AgentRunResponse] = []
+
+        class ObservabilityMiddleware(AgentMiddleware):
+            async def process(
+                self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
+            ) -> None:
+                # Context should be empty before next()
+                assert context.result is None
+
+                # Call next to execute
+                await next(context)
+
+                # Context should now contain the response for observability
+                assert context.result is not None
+                assert isinstance(context.result, AgentRunResponse)
+                observed_responses.append(context.result)
+
+        middleware = ObservabilityMiddleware()
+        pipeline = AgentMiddlewarePipeline([middleware])
+        messages = [ChatMessage(role=Role.USER, text="test")]
+        context = AgentRunContext(agent=mock_agent, messages=messages)
+
+        async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
+            return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="executed response")])
+
+        result = await pipeline.execute(mock_agent, messages, 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: AIFunction[Any, Any]) -> None:
+        """Test that middleware can observe function result after execution."""
+        observed_results: list[str] = []
+
+        class ObservabilityMiddleware(FunctionMiddleware):
+            async def process(
+                self,
+                context: FunctionInvocationContext,
+                next: Callable[[FunctionInvocationContext], Awaitable[None]],
+            ) -> None:
+                # Context should be empty before next()
+                assert context.result is None
+
+                # Call next to execute
+                await next(context)
+
+                # 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(mock_function, arguments, 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: AgentProtocol) -> None:
+        """Test that middleware can override response after observing execution."""
+
+        class PostExecutionOverrideMiddleware(AgentMiddleware):
+            async def process(
+                self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
+            ) -> None:
+                # Call next to execute first
+                await next(context)
+
+                # Now observe and conditionally override
+                assert context.result is not None
+                assert isinstance(context.result, AgentRunResponse)
+
+                if "modify" in context.result.messages[0].text:
+                    # Override after observing
+                    context.result = AgentRunResponse(
+                        messages=[ChatMessage(role=Role.ASSISTANT, text="modified after execution")]
+                    )
+
+        middleware = PostExecutionOverrideMiddleware()
+        pipeline = AgentMiddlewarePipeline([middleware])
+        messages = [ChatMessage(role=Role.USER, text="test")]
+        context = AgentRunContext(agent=mock_agent, messages=messages)
+
+        async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
+            return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response to modify")])
+
+        result = await pipeline.execute(mock_agent, messages, context, final_handler)
+
+        # Verify response was modified after execution
+        assert result is not None
+        assert result.messages[0].text == "modified after execution"
+
+    async def test_function_middleware_post_execution_override(self, mock_function: AIFunction[Any, Any]) -> None:
+        """Test that middleware can override function result after observing execution."""
+
+        class PostExecutionOverrideMiddleware(FunctionMiddleware):
+            async def process(
+                self,
+                context: FunctionInvocationContext,
+                next: Callable[[FunctionInvocationContext], Awaitable[None]],
+            ) -> None:
+                # Call next to execute first
+                await next(context)
+
+                # 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(mock_function, arguments, context, final_handler)
+
+        # Verify result was modified after execution
+        assert result == "modified after execution"
+
+
+@pytest.fixture
+def mock_agent() -> AgentProtocol:
+    """Mock agent for testing."""
+    agent = MagicMock(spec=AgentProtocol)
+    agent.name = "test_agent"
+    return agent
+
+
+@pytest.fixture
+def mock_function() -> AIFunction[Any, Any]:
+    """Mock function for testing."""
+    function = MagicMock(spec=AIFunction[Any, Any])
+    function.name = "test_function"
+    return function

python/packages/core/tests/main/test_middleware_with_chat.py

@@ -0,0 +1,436 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+from agent_framework import (
+    ChatAgent,
+    ChatContext,
+    ChatMessage,
+    ChatMiddleware,
+    ChatResponse,
+    FunctionCallContent,
+    FunctionInvocationContext,
+    Role,
+    chat_middleware,
+    function_middleware,
+    use_function_invocation,
+)
+
+from .conftest import MockBaseChatClient
+
+
+class TestChatMiddleware:
+    """Test cases for chat middleware functionality."""
+
+    async def test_class_based_chat_middleware(self, chat_client_base: "MockBaseChatClient") -> None:
+        """Test class-based chat middleware with ChatClient."""
+        execution_order: list[str] = []
+
+        class LoggingChatMiddleware(ChatMiddleware):
+            async def process(
+                self,
+                context: ChatContext,
+                next: Callable[[ChatContext], Awaitable[None]],
+            ) -> None:
+                execution_order.append("chat_middleware_before")
+                await next(context)
+                execution_order.append("chat_middleware_after")
+
+        # Add middleware to chat client
+        chat_client_base.middleware = [LoggingChatMiddleware()]
+
+        # Execute chat client directly
+        messages = [ChatMessage(role=Role.USER, text="test message")]
+        response = await chat_client_base.get_response(messages)
+
+        # Verify response
+        assert response is not None
+        assert len(response.messages) > 0
+        assert response.messages[0].role == Role.ASSISTANT
+
+        # Verify middleware execution order
+        assert execution_order == ["chat_middleware_before", "chat_middleware_after"]
+
+    async def test_function_based_chat_middleware(self, chat_client_base: "MockBaseChatClient") -> None:
+        """Test function-based chat middleware with ChatClient."""
+        execution_order: list[str] = []
+
+        @chat_middleware
+        async def logging_chat_middleware(context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
+            execution_order.append("function_middleware_before")
+            await next(context)
+            execution_order.append("function_middleware_after")
+
+        # Add middleware to chat client
+        chat_client_base.middleware = [logging_chat_middleware]
+
+        # Execute chat client directly
+        messages = [ChatMessage(role=Role.USER, text="test message")]
+        response = await chat_client_base.get_response(messages)
+
+        # Verify response
+        assert response is not None
+        assert len(response.messages) > 0
+        assert response.messages[0].role == Role.ASSISTANT
+
+        # Verify middleware execution order
+        assert execution_order == ["function_middleware_before", "function_middleware_after"]
+
+    async def test_chat_middleware_can_modify_messages(self, chat_client_base: "MockBaseChatClient") -> None:
+        """Test that chat middleware can modify messages before sending to model."""
+
+        @chat_middleware
+        async def message_modifier_middleware(
+            context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]
+        ) -> None:
+            # Modify the first message by adding a prefix
+            if context.messages and len(context.messages) > 0:
+                original_text = context.messages[0].text or ""
+                context.messages[0] = ChatMessage(role=context.messages[0].role, text=f"MODIFIED: {original_text}")
+            await next(context)
+
+        # Add middleware to chat client
+        chat_client_base.middleware = [message_modifier_middleware]
+
+        # Execute chat client
+        messages = [ChatMessage(role=Role.USER, text="test message")]
+        response = await chat_client_base.get_response(messages)
+
+        # Verify that the message was modified (MockChatClient echoes back the input)
+        assert response is not None
+        assert len(response.messages) > 0
+        # The mock client should receive the modified message
+        assert "MODIFIED: test message" in response.messages[0].text
+
+    async def test_chat_middleware_can_override_response(self, chat_client_base: "MockBaseChatClient") -> None:
+        """Test that chat middleware can override the response."""
+
+        @chat_middleware
+        async def response_override_middleware(
+            context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]
+        ) -> None:
+            # Override the response without calling next()
+            context.result = ChatResponse(
+                messages=[ChatMessage(role=Role.ASSISTANT, text="Middleware overridden response")],
+                response_id="middleware-response-123",
+            )
+            context.terminate = True
+
+        # Add middleware to chat client
+        chat_client_base.middleware = [response_override_middleware]
+
+        # Execute chat client
+        messages = [ChatMessage(role=Role.USER, text="test message")]
+        response = await chat_client_base.get_response(messages)
+
+        # Verify that the response was overridden
+        assert response is not None
+        assert len(response.messages) > 0
+        assert response.messages[0].text == "Middleware overridden response"
+        assert response.response_id == "middleware-response-123"
+
+    async def test_multiple_chat_middleware_execution_order(self, chat_client_base: "MockBaseChatClient") -> None:
+        """Test that multiple chat middleware execute in the correct order."""
+        execution_order: list[str] = []
+
+        @chat_middleware
+        async def first_middleware(context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
+            execution_order.append("first_before")
+            await next(context)
+            execution_order.append("first_after")
+
+        @chat_middleware
+        async def second_middleware(context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
+            execution_order.append("second_before")
+            await next(context)
+            execution_order.append("second_after")
+
+        # Add middleware to chat client (order should be preserved)
+        chat_client_base.middleware = [first_middleware, second_middleware]
+
+        # Execute chat client
+        messages = [ChatMessage(role=Role.USER, text="test message")]
+        response = await chat_client_base.get_response(messages)
+
+        # Verify response
+        assert response is not None
+
+        # Verify middleware execution order (nested execution)
+        expected_order = ["first_before", "second_before", "second_after", "first_after"]
+        assert execution_order == expected_order
+
+    async def test_chat_agent_with_chat_middleware(self) -> None:
+        """Test ChatAgent with chat middleware specified at agent level."""
+        execution_order: list[str] = []
+
+        @chat_middleware
+        async def agent_level_chat_middleware(
+            context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]
+        ) -> None:
+            execution_order.append("agent_chat_middleware_before")
+            await next(context)
+            execution_order.append("agent_chat_middleware_after")
+
+        chat_client = MockBaseChatClient()
+
+        # Create ChatAgent with chat middleware
+        agent = ChatAgent(chat_client=chat_client, middleware=[agent_level_chat_middleware])
+
+        # Execute the agent
+        messages = [ChatMessage(role=Role.USER, text="test message")]
+        response = await agent.run(messages)
+
+        # Verify response
+        assert response is not None
+        assert len(response.messages) > 0
+        assert response.messages[0].role == Role.ASSISTANT
+
+        # Verify middleware execution order
+        assert execution_order == ["agent_chat_middleware_before", "agent_chat_middleware_after"]
+
+    async def test_chat_agent_with_multiple_chat_middleware(self, chat_client_base: "MockBaseChatClient") -> None:
+        """Test that ChatAgent can have multiple chat middleware."""
+        execution_order: list[str] = []
+
+        @chat_middleware
+        async def first_middleware(context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
+            execution_order.append("first_before")
+            await next(context)
+            execution_order.append("first_after")
+
+        @chat_middleware
+        async def second_middleware(context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
+            execution_order.append("second_before")
+            await next(context)
+            execution_order.append("second_after")
+
+        # Create ChatAgent with multiple chat middleware
+        agent = ChatAgent(chat_client=chat_client_base, middleware=[first_middleware, second_middleware])
+
+        # Execute the agent
+        messages = [ChatMessage(role=Role.USER, text="test message")]
+        response = await agent.run(messages)
+
+        # Verify response
+        assert response is not None
+
+        # Verify both middleware executed (nested execution order)
+        expected_order = ["first_before", "second_before", "second_after", "first_after"]
+        assert execution_order == expected_order
+
+    async def test_chat_middleware_with_streaming(self, chat_client_base: "MockBaseChatClient") -> None:
+        """Test chat middleware with streaming responses."""
+        execution_order: list[str] = []
+
+        @chat_middleware
+        async def streaming_middleware(context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
+            execution_order.append("streaming_before")
+            # Verify it's a streaming context
+            assert context.is_streaming is True
+            await next(context)
+            execution_order.append("streaming_after")
+
+        # Add middleware to chat client
+        chat_client_base.middleware = [streaming_middleware]
+
+        # Execute streaming response
+        messages = [ChatMessage(role=Role.USER, text="test message")]
+        updates: list[object] = []
+        async for update in chat_client_base.get_streaming_response(messages):
+            updates.append(update)
+
+        # Verify we got updates
+        assert len(updates) > 0
+
+        # Verify middleware executed
+        assert execution_order == ["streaming_before", "streaming_after"]
+
+    async def test_run_level_middleware_isolation(self, chat_client_base: "MockBaseChatClient") -> None:
+        """Test that run-level middleware is isolated and doesn't persist across calls."""
+        execution_count = {"count": 0}
+
+        @chat_middleware
+        async def counting_middleware(context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
+            execution_count["count"] += 1
+            await next(context)
+
+        # First call with run-level middleware
+        messages = [ChatMessage(role=Role.USER, text="first message")]
+        response1 = await chat_client_base.get_response(messages, middleware=[counting_middleware])
+        assert response1 is not None
+        assert execution_count["count"] == 1
+
+        # Second call WITHOUT run-level middleware - should not execute the middleware
+        messages = [ChatMessage(role=Role.USER, text="second message")]
+        response2 = await chat_client_base.get_response(messages)
+        assert response2 is not None
+        assert execution_count["count"] == 1  # Should still be 1, not 2
+
+        # Third call with run-level middleware again - should execute
+        messages = [ChatMessage(role=Role.USER, text="third message")]
+        response3 = await chat_client_base.get_response(messages, middleware=[counting_middleware])
+        assert response3 is not None
+        assert execution_count["count"] == 2  # Should be 2 now
+
+    async def test_chat_client_middleware_can_access_and_override_custom_kwargs(
+        self, chat_client_base: "MockBaseChatClient"
+    ) -> None:
+        """Test that chat client middleware can access and override custom parameters like temperature."""
+        captured_kwargs: dict[str, Any] = {}
+        modified_kwargs: dict[str, Any] = {}
+
+        @chat_middleware
+        async def kwargs_middleware(context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
+            # Capture the original kwargs
+            captured_kwargs.update(context.kwargs)
+
+            # Modify some kwargs
+            context.kwargs["temperature"] = 0.9
+            context.kwargs["max_tokens"] = 500
+            context.kwargs["new_param"] = "added_by_middleware"
+
+            # Store modified kwargs for verification
+            modified_kwargs.update(context.kwargs)
+
+            await next(context)
+
+        # Add middleware to chat client
+        chat_client_base.middleware = [kwargs_middleware]
+
+        # Execute chat client with custom parameters
+        messages = [ChatMessage(role=Role.USER, text="test message")]
+        response = await chat_client_base.get_response(
+            messages, temperature=0.7, max_tokens=100, custom_param="test_value"
+        )
+
+        # Verify response
+        assert response is not None
+        assert len(response.messages) > 0
+
+        assert captured_kwargs["temperature"] == 0.7
+        assert captured_kwargs["max_tokens"] == 100
+        assert captured_kwargs["custom_param"] == "test_value"
+
+        # Verify middleware could modify the kwargs
+        assert modified_kwargs["temperature"] == 0.9
+        assert modified_kwargs["max_tokens"] == 500
+        assert modified_kwargs["new_param"] == "added_by_middleware"
+        assert modified_kwargs["custom_param"] == "test_value"  # Should still be there
+
+    async def test_function_middleware_registration_on_chat_client(self) -> None:
+        """Test function middleware registered on ChatClient is executed during function calls."""
+        execution_order: list[str] = []
+
+        @function_middleware
+        async def test_function_middleware(
+            context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
+        ) -> None:
+            execution_order.append(f"function_middleware_before_{context.function.name}")
+            await next(context)
+            execution_order.append(f"function_middleware_after_{context.function.name}")
+
+        # Define a simple tool function
+        def sample_tool(location: str) -> str:
+            """Get weather for a location."""
+            return f"Weather in {location}: sunny"
+
+        # Create function-invocation enabled chat client
+        chat_client = use_function_invocation(MockBaseChatClient)()
+
+        # Set function middleware directly on the chat client
+        chat_client.middleware = [test_function_middleware]
+
+        # Prepare responses that will trigger function invocation
+        function_call_response = ChatResponse(
+            messages=[
+                ChatMessage(
+                    role=Role.ASSISTANT,
+                    contents=[
+                        FunctionCallContent(
+                            call_id="call_1",
+                            name="sample_tool",
+                            arguments={"location": "San Francisco"},
+                        )
+                    ],
+                )
+            ]
+        )
+        final_response = ChatResponse(
+            messages=[ChatMessage(role=Role.ASSISTANT, text="Based on the weather data, it's sunny!")]
+        )
+
+        chat_client.run_responses = [function_call_response, final_response]
+
+        # Execute the chat client directly with tools - this should trigger function invocation and middleware
+        messages = [ChatMessage(role=Role.USER, text="What's the weather in San Francisco?")]
+        response = await chat_client.get_response(messages, tools=[sample_tool])
+
+        # Verify response
+        assert response is not None
+        assert len(response.messages) > 0
+        assert chat_client.call_count == 2  # Two calls: function call + final response
+
+        # Verify function middleware was executed
+        assert execution_order == [
+            "function_middleware_before_sample_tool",
+            "function_middleware_after_sample_tool",
+        ]
+
+    async def test_run_level_function_middleware(self) -> None:
+        """Test that function middleware passed to get_response method is also invoked."""
+        execution_order: list[str] = []
+
+        @function_middleware
+        async def run_level_function_middleware(
+            context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
+        ) -> None:
+            execution_order.append("run_level_function_middleware_before")
+            await next(context)
+            execution_order.append("run_level_function_middleware_after")
+
+        # Define a simple tool function
+        def sample_tool(location: str) -> str:
+            """Get weather for a location."""
+            return f"Weather in {location}: sunny"
+
+        # Create function-invocation enabled chat client
+        chat_client = use_function_invocation(MockBaseChatClient)()
+
+        # Prepare responses that will trigger function invocation
+        function_call_response = ChatResponse(
+            messages=[
+                ChatMessage(
+                    role=Role.ASSISTANT,
+                    contents=[
+                        FunctionCallContent(
+                            call_id="call_2",
+                            name="sample_tool",
+                            arguments={"location": "New York"},
+                        )
+                    ],
+                )
+            ]
+        )
+        final_response = ChatResponse(
+            messages=[ChatMessage(role=Role.ASSISTANT, text="The weather information has been retrieved!")]
+        )
+
+        chat_client.run_responses = [function_call_response, final_response]
+
+        # Execute the chat client directly with run-level middleware and tools
+        messages = [ChatMessage(role=Role.USER, text="What's the weather in New York?")]
+        response = await chat_client.get_response(
+            messages, tools=[sample_tool], middleware=[run_level_function_middleware]
+        )
+
+        # Verify response
+        assert response is not None
+        assert len(response.messages) > 0
+        assert chat_client.call_count == 2  # Two calls: function call + final response
+
+        # Verify run-level function middleware was executed once (during function invocation)
+        assert execution_order == [
+            "run_level_function_middleware_before",
+            "run_level_function_middleware_after",
+        ]

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

@@ -0,0 +1,469 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import logging
+from collections.abc import MutableSequence
+from typing import Any
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+from opentelemetry.semconv_ai import SpanAttributes
+from opentelemetry.trace import StatusCode
+
+from agent_framework import (
+    AGENT_FRAMEWORK_USER_AGENT,
+    AgentProtocol,
+    AgentRunResponse,
+    AgentThread,
+    BaseChatClient,
+    ChatMessage,
+    ChatOptions,
+    ChatResponse,
+    ChatResponseUpdate,
+    Role,
+    UsageDetails,
+    prepend_agent_framework_to_user_agent,
+)
+from agent_framework.exceptions import AgentInitializationError, ChatClientInitializationError
+from agent_framework.observability import (
+    OPEN_TELEMETRY_AGENT_MARKER,
+    OPEN_TELEMETRY_CHAT_CLIENT_MARKER,
+    ROLE_EVENT_MAP,
+    ChatMessageListTimestampFilter,
+    OtelAttr,
+    get_function_span,
+    use_agent_observability,
+    use_observability,
+)
+
+# region Test constants
+
+
+def test_role_event_map():
+    """Test that ROLE_EVENT_MAP contains expected mappings."""
+    assert ROLE_EVENT_MAP["system"] == OtelAttr.SYSTEM_MESSAGE
+    assert ROLE_EVENT_MAP["user"] == OtelAttr.USER_MESSAGE
+    assert ROLE_EVENT_MAP["assistant"] == OtelAttr.ASSISTANT_MESSAGE
+    assert ROLE_EVENT_MAP["tool"] == OtelAttr.TOOL_MESSAGE
+
+
+def test_enum_values():
+    """Test that OtelAttr enum has expected values."""
+    assert OtelAttr.OPERATION == "gen_ai.operation.name"
+    assert SpanAttributes.LLM_SYSTEM == "gen_ai.system"
+    assert SpanAttributes.LLM_REQUEST_MODEL == "gen_ai.request.model"
+    assert OtelAttr.CHAT_COMPLETION_OPERATION == "chat"
+    assert OtelAttr.TOOL_EXECUTION_OPERATION == "execute_tool"
+    assert OtelAttr.AGENT_INVOKE_OPERATION == "invoke_agent"
+
+
+# region Test ChatMessageListTimestampFilter
+
+
+def test_filter_without_index_key():
+    """Test filter method when record doesn't have INDEX_KEY."""
+    log_filter = ChatMessageListTimestampFilter()
+    record = logging.LogRecord(
+        name="test", level=logging.INFO, pathname="", lineno=0, msg="test message", args=(), exc_info=None
+    )
+    original_created = record.created
+
+    result = log_filter.filter(record)
+
+    assert result is True
+    assert record.created == original_created
+
+
+def test_filter_with_index_key():
+    """Test filter method when record has INDEX_KEY."""
+    log_filter = ChatMessageListTimestampFilter()
+    record = logging.LogRecord(
+        name="test", level=logging.INFO, pathname="", lineno=0, msg="test message", args=(), exc_info=None
+    )
+    original_created = record.created
+
+    # Add the index key
+    setattr(record, ChatMessageListTimestampFilter.INDEX_KEY, 5)
+
+    result = log_filter.filter(record)
+
+    assert result is True
+    # Should increment by 5 microseconds (5 * 1e-6)
+    assert record.created == original_created + 5 * 1e-6
+
+
+def test_index_key_constant():
+    """Test that INDEX_KEY constant is correctly defined."""
+    assert ChatMessageListTimestampFilter.INDEX_KEY == "chat_message_index"
+
+
+# region Test get_function_span
+
+
+def test_start_span_basic(span_exporter: InMemorySpanExporter):
+    """Test starting a span with basic function info."""
+    # Create a mock function
+    mock_function = Mock()
+    mock_function.name = "test_function"
+    mock_function.description = "Test function description"
+    attributes = {
+        OtelAttr.OPERATION: OtelAttr.TOOL_EXECUTION_OPERATION,
+        OtelAttr.TOOL_NAME: "test_function",
+        OtelAttr.TOOL_DESCRIPTION: "Test function description",
+        OtelAttr.TOOL_TYPE: "function",
+    }
+    span_exporter.clear()
+    with get_function_span(attributes) as function_span:
+        assert function_span is not None
+        function_span.set_attribute("test_attr", "test_value")
+
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    span = spans[0]
+    assert span.name == "execute_tool test_function"
+    assert span.attributes["test_attr"] == "test_value"
+    assert span.attributes[OtelAttr.OPERATION.value] == OtelAttr.TOOL_EXECUTION_OPERATION
+    assert span.attributes[OtelAttr.TOOL_NAME] == "test_function"
+    assert span.attributes[OtelAttr.TOOL_DESCRIPTION] == "Test function description"
+
+
+def test_start_span_with_tool_call_id(span_exporter: InMemorySpanExporter):
+    """Test starting a span with tool_call_id."""
+
+    tool_call_id = "test_call_123"
+    attributes = {
+        OtelAttr.OPERATION: OtelAttr.TOOL_EXECUTION_OPERATION,
+        OtelAttr.TOOL_NAME: "test_function",
+        OtelAttr.TOOL_DESCRIPTION: "Test function",
+        OtelAttr.TOOL_TYPE: "function",
+        OtelAttr.TOOL_CALL_ID: tool_call_id,
+    }
+
+    span_exporter.clear()
+    with get_function_span(attributes) as function_span:
+        assert function_span is not None
+        function_span.set_attribute("test_attr", "test_value")
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    span = spans[0]
+    assert span.name == "execute_tool test_function"
+    assert span.attributes["test_attr"] == "test_value"
+    assert span.attributes[OtelAttr.TOOL_CALL_ID] == tool_call_id
+    # Verify all attributes
+    assert span.attributes[OtelAttr.OPERATION.value] == OtelAttr.TOOL_EXECUTION_OPERATION
+    assert span.attributes[OtelAttr.TOOL_NAME] == "test_function"
+    assert span.attributes[OtelAttr.TOOL_DESCRIPTION] == "Test function"
+    assert span.attributes[OtelAttr.TOOL_TYPE] == "function"
+
+
+# region Test use_observability decorator
+
+
+def test_decorator_with_valid_class():
+    """Test that decorator works with a valid BaseChatClient-like class."""
+
+    # Create a mock class with the required methods
+    class MockChatClient:
+        async def get_response(self, messages, **kwargs):
+            return Mock()
+
+        async def get_streaming_response(self, messages, **kwargs):
+            async def gen():
+                yield Mock()
+
+            return gen()
+
+    # Apply the decorator
+    decorated_class = use_observability(MockChatClient)
+    assert hasattr(decorated_class, OPEN_TELEMETRY_CHAT_CLIENT_MARKER)
+
+
+def test_decorator_with_missing_methods():
+    """Test that decorator handles classes missing required methods gracefully."""
+
+    class MockChatClient:
+        OTEL_PROVIDER_NAME = "test_provider"
+
+    # Apply the decorator - should not raise an error
+    with pytest.raises(ChatClientInitializationError):
+        use_observability(MockChatClient)
+
+
+def test_decorator_with_partial_methods():
+    """Test decorator when only one method is present."""
+
+    class MockChatClient:
+        OTEL_PROVIDER_NAME = "test_provider"
+
+        async def get_response(self, messages, **kwargs):
+            return Mock()
+
+    with pytest.raises(ChatClientInitializationError):
+        use_observability(MockChatClient)
+
+
+# region Test telemetry decorator with mock client
+
+
+@pytest.fixture
+def mock_chat_client():
+    """Create a mock chat client for testing."""
+
+    class MockChatClient(BaseChatClient):
+        def service_url(self):
+            return "https://test.example.com"
+
+        async def _inner_get_response(
+            self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+        ):
+            return ChatResponse(
+                messages=[ChatMessage(role=Role.ASSISTANT, text="Test response")],
+                usage_details=UsageDetails(input_token_count=10, output_token_count=20),
+                finish_reason=None,
+            )
+
+        async def _inner_get_streaming_response(
+            self, *, messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+        ):
+            yield ChatResponseUpdate(text="Hello", role=Role.ASSISTANT)
+            yield ChatResponseUpdate(text=" world", role=Role.ASSISTANT)
+
+    return MockChatClient
+
+
+@pytest.mark.parametrize("enable_sensitive_data", [True, False], indirect=True)
+async def test_chat_client_observability(mock_chat_client, span_exporter: InMemorySpanExporter, enable_sensitive_data):
+    """Test that when diagnostics are enabled, telemetry is applied."""
+    client = use_observability(mock_chat_client)()
+
+    messages = [ChatMessage(role=Role.USER, text="Test message")]
+    span_exporter.clear()
+    response = await client.get_response(messages=messages, model="Test")
+    assert response is not None
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    span = spans[0]
+    assert span.name == "chat Test"
+    assert span.attributes[OtelAttr.OPERATION.value] == OtelAttr.CHAT_COMPLETION_OPERATION
+    assert span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "Test"
+    assert span.attributes[OtelAttr.INPUT_TOKENS] == 10
+    assert span.attributes[OtelAttr.OUTPUT_TOKENS] == 20
+    if enable_sensitive_data:
+        assert span.attributes[OtelAttr.INPUT_MESSAGES] is not None
+        assert span.attributes[OtelAttr.OUTPUT_MESSAGES] is not None
+
+
+@pytest.mark.parametrize("enable_sensitive_data", [True, False], indirect=True)
+async def test_chat_client_streaming_observability(
+    mock_chat_client, span_exporter: InMemorySpanExporter, enable_sensitive_data
+):
+    """Test streaming telemetry through the use_observability decorator."""
+    client = use_observability(mock_chat_client)()
+    messages = [ChatMessage(role=Role.USER, text="Test")]
+    span_exporter.clear()
+    # Collect all yielded updates
+    updates = []
+    async for update in client.get_streaming_response(messages=messages, model="Test"):
+        updates.append(update)
+
+    # Verify we got the expected updates, this shouldn't be dependent on otel
+    assert len(updates) == 2
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    span = spans[0]
+    assert span.name == "chat Test"
+    assert span.attributes[OtelAttr.OPERATION.value] == OtelAttr.CHAT_COMPLETION_OPERATION
+    assert span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "Test"
+    if enable_sensitive_data:
+        assert span.attributes[OtelAttr.INPUT_MESSAGES] is not None
+        assert span.attributes[OtelAttr.OUTPUT_MESSAGES] is not None
+
+
+def test_prepend_user_agent_with_none_value():
+    """Test prepend user agent with None value in headers."""
+    headers = {"User-Agent": None}
+    result = prepend_agent_framework_to_user_agent(headers)
+
+    # Should handle None gracefully
+    assert "User-Agent" in result
+    assert AGENT_FRAMEWORK_USER_AGENT in str(result["User-Agent"])
+
+
+# region Test use_agent_observability decorator
+
+
+def test_agent_decorator_with_valid_class():
+    """Test that agent decorator works with a valid ChatAgent-like class."""
+
+    # Create a mock class with the required methods
+    class MockChatClientAgent:
+        AGENT_SYSTEM_NAME = "test_agent_system"
+
+        def __init__(self):
+            self.id = "test_agent_id"
+            self.name = "test_agent"
+            self.display_name = "Test Agent"
+            self.description = "Test agent description"
+
+        async def run(self, messages=None, *, thread=None, **kwargs):
+            return Mock()
+
+        async def run_stream(self, messages=None, *, thread=None, **kwargs):
+            async def gen():
+                yield Mock()
+
+            return gen()
+
+        def get_new_thread(self) -> AgentThread:
+            return AgentThread()
+
+    # Apply the decorator
+    decorated_class = use_agent_observability(MockChatClientAgent)
+
+    assert hasattr(decorated_class, OPEN_TELEMETRY_AGENT_MARKER)
+
+
+def test_agent_decorator_with_missing_methods():
+    """Test that agent decorator handles classes missing required methods gracefully."""
+
+    class MockAgent:
+        AGENT_SYSTEM_NAME = "test_agent_system"
+
+    # Apply the decorator - should not raise an error
+    with pytest.raises(AgentInitializationError):
+        use_agent_observability(MockAgent)
+
+
+def test_agent_decorator_with_partial_methods():
+    """Test agent decorator when only one method is present."""
+    from agent_framework.observability import use_agent_observability
+
+    class MockAgent:
+        AGENT_SYSTEM_NAME = "test_agent_system"
+
+        def __init__(self):
+            self.id = "test_agent_id"
+            self.name = "test_agent"
+            self.display_name = "Test Agent"
+
+        async def run(self, messages=None, *, thread=None, **kwargs):
+            return Mock()
+
+    with pytest.raises(AgentInitializationError):
+        use_agent_observability(MockAgent)
+
+
+# region Test agent telemetry decorator with mock agent
+
+
+@pytest.fixture
+def mock_chat_agent():
+    """Create a mock chat client agent for testing."""
+
+    class MockChatClientAgent:
+        AGENT_SYSTEM_NAME = "test_agent_system"
+
+        def __init__(self):
+            self.id = "test_agent_id"
+            self.name = "test_agent"
+            self.display_name = "Test Agent"
+            self.description = "Test agent description"
+
+        async def run(self, messages=None, *, thread=None, **kwargs):
+            return AgentRunResponse(
+                messages=[ChatMessage(role=Role.ASSISTANT, text="Agent response")],
+                usage_details=UsageDetails(input_token_count=15, output_token_count=25),
+                response_id="test_response_id",
+                raw_representation=Mock(finish_reason=Mock(value="stop")),
+            )
+
+        async def run_stream(self, messages=None, *, thread=None, **kwargs):
+            from agent_framework import AgentRunResponseUpdate
+
+            yield AgentRunResponseUpdate(text="Hello", role=Role.ASSISTANT)
+            yield AgentRunResponseUpdate(text=" from agent", role=Role.ASSISTANT)
+
+    return MockChatClientAgent
+
+
+@pytest.mark.parametrize("enable_sensitive_data", [True, False], indirect=True)
+async def test_agent_instrumentation_enabled(
+    mock_chat_agent: AgentProtocol, span_exporter: InMemorySpanExporter, enable_sensitive_data
+):
+    """Test that when agent diagnostics are enabled, telemetry is applied."""
+
+    agent = use_agent_observability(mock_chat_agent)()
+
+    span_exporter.clear()
+    response = await agent.run("Test message")
+    assert response is not None
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    span = spans[0]
+    assert span.name == "invoke_agent Test Agent"
+    assert span.attributes[OtelAttr.OPERATION.value] == OtelAttr.AGENT_INVOKE_OPERATION
+    assert span.attributes[OtelAttr.AGENT_ID] == "test_agent_id"
+    assert span.attributes[OtelAttr.AGENT_NAME] == "Test Agent"
+    assert span.attributes[OtelAttr.AGENT_DESCRIPTION] == "Test agent description"
+    assert span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "unknown"
+    assert span.attributes[OtelAttr.INPUT_TOKENS] == 15
+    assert span.attributes[OtelAttr.OUTPUT_TOKENS] == 25
+    if enable_sensitive_data:
+        assert span.attributes[OtelAttr.OUTPUT_MESSAGES] is not None
+
+
+@pytest.mark.parametrize("enable_sensitive_data", [True, False], indirect=True)
+async def test_agent_streaming_response_with_diagnostics_enabled_via_decorator(
+    mock_chat_agent: AgentProtocol, span_exporter: InMemorySpanExporter, enable_sensitive_data
+):
+    """Test agent streaming telemetry through the use_agent_observability decorator."""
+    agent = use_agent_observability(mock_chat_agent)()
+    span_exporter.clear()
+    updates = []
+    async for update in agent.run_stream("Test message"):
+        updates.append(update)
+
+    # Verify we got the expected updates
+    assert len(updates) == 2
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    span = spans[0]
+    assert span.name == "invoke_agent Test Agent"
+    assert span.attributes[OtelAttr.OPERATION.value] == OtelAttr.AGENT_INVOKE_OPERATION
+    assert span.attributes[OtelAttr.AGENT_ID] == "test_agent_id"
+    assert span.attributes[OtelAttr.AGENT_NAME] == "Test Agent"
+    assert span.attributes[OtelAttr.AGENT_DESCRIPTION] == "Test agent description"
+    assert span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "unknown"
+    if enable_sensitive_data:
+        assert span.attributes.get(OtelAttr.OUTPUT_MESSAGES) is not None  # Streaming, so no usage yet
+
+
+async def test_agent_run_with_exception_handling(mock_chat_agent: AgentProtocol):
+    """Test agent run with exception handling."""
+
+    async def run_with_error(self, messages=None, *, thread=None, **kwargs):
+        raise RuntimeError("Agent run error")
+
+    mock_chat_agent.run = run_with_error
+
+    agent = use_agent_observability(mock_chat_agent)()
+
+    from opentelemetry.trace import Span
+
+    with (
+        patch("agent_framework.observability._get_span") as mock_get_span,
+    ):
+        mock_span = MagicMock(spec=Span)
+        # Ensure the patched context manager returns mock_span when entered
+        mock_get_span.return_value.__enter__.return_value = mock_span
+        # Should raise the exception and call error handler
+        with pytest.raises(RuntimeError, match="Agent run error"):
+            await agent.run("Test message")
+
+        # Verify error was recorded
+        # Check that both error attributes were set on the span
+        mock_span.set_attribute.assert_called_with(OtelAttr.ERROR_TYPE, "RuntimeError")
+        mock_span.record_exception.assert_called_once()
+        mock_span.set_status.assert_called_once_with(
+            status=StatusCode.ERROR, description=repr(RuntimeError("Agent run error"))
+        )

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

@@ -0,0 +1,98 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from unittest.mock import patch
+
+from agent_framework import (
+    AGENT_FRAMEWORK_USER_AGENT,
+    USER_AGENT_KEY,
+    USER_AGENT_TELEMETRY_DISABLED_ENV_VAR,
+    prepend_agent_framework_to_user_agent,
+)
+
+# 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 = {}
+    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

python/packages/core/tests/main/test_threads.py

@@ -0,0 +1,398 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from collections.abc import Sequence
+from typing import Any
+
+import pytest
+
+from agent_framework import AgentThread, ChatMessage, ChatMessageStore, Role
+from agent_framework._threads import AgentThreadState, ChatMessageStoreState
+from agent_framework.exceptions import AgentThreadException
+
+
+class MockChatMessageStore:
+    """Mock implementation of ChatMessageStoreProtocol for testing."""
+
+    def __init__(self, messages: list[ChatMessage] | None = None) -> None:
+        self._messages = messages or []
+        self._serialize_calls = 0
+        self._deserialize_calls = 0
+
+    async def list_messages(self) -> list[ChatMessage]:
+        return self._messages
+
+    async def add_messages(self, messages: Sequence[ChatMessage]) -> None:
+        self._messages.extend(messages)
+
+    async def serialize(self, **kwargs: Any) -> Any:
+        self._serialize_calls += 1
+        return {"messages": [msg.__dict__ for msg in self._messages], "kwargs": kwargs}
+
+    async def update_from_state(self, serialized_store_state: Any, **kwargs: Any) -> None:
+        self._deserialize_calls += 1
+        if serialized_store_state and "messages" in serialized_store_state:
+            self._messages = serialized_store_state["messages"]
+
+    @classmethod
+    async def deserialize(cls, serialized_store_state: Any, **kwargs: Any) -> "MockChatMessageStore":
+        instance = cls()
+        await instance.update_from_state(serialized_store_state, **kwargs)
+        return instance
+
+
+@pytest.fixture
+def sample_messages() -> list[ChatMessage]:
+    """Fixture providing sample chat messages for testing."""
+    return [
+        ChatMessage(role=Role.USER, text="Hello", message_id="msg1"),
+        ChatMessage(role=Role.ASSISTANT, text="Hi there!", message_id="msg2"),
+        ChatMessage(role=Role.USER, text="How are you?", message_id="msg3"),
+    ]
+
+
+@pytest.fixture
+def sample_message() -> ChatMessage:
+    """Fixture providing a single sample chat message for testing."""
+    return ChatMessage(role=Role.USER, text="Test message", message_id="test1")
+
+
+class TestAgentThread:
+    """Test cases for AgentThread class."""
+
+    def test_init_with_no_parameters(self) -> None:
+        """Test AgentThread initialization with no parameters."""
+        thread = AgentThread()
+        assert thread.service_thread_id is None
+        assert thread.message_store is None
+
+    def test_init_with_service_thread_id(self) -> None:
+        """Test AgentThread initialization with service_thread_id."""
+        service_thread_id = "test-conversation-123"
+        thread = AgentThread(service_thread_id=service_thread_id)
+        assert thread.service_thread_id == service_thread_id
+        assert thread.message_store is None
+
+    def test_init_with_message_store(self) -> None:
+        """Test AgentThread initialization with message_store."""
+        store = ChatMessageStore()
+        thread = AgentThread(message_store=store)
+        assert thread.service_thread_id is None
+        assert thread.message_store is store
+
+    def test_service_thread_id_property_setter(self) -> None:
+        """Test service_thread_id property setter."""
+        thread = AgentThread()
+        service_thread_id = "test-conversation-456"
+
+        thread.service_thread_id = service_thread_id
+        assert thread.service_thread_id == service_thread_id
+
+    def test_service_thread_id_setter_with_existing_message_store_raises_error(self) -> None:
+        """Test that setting service_thread_id when message_store exists raises AgentThreadException."""
+        store = ChatMessageStore()
+        thread = AgentThread(message_store=store)
+
+        with pytest.raises(AgentThreadException, match="Only the service_thread_id or message_store may be set"):
+            thread.service_thread_id = "test-conversation-789"
+
+    def test_service_thread_id_setter_with_none_values(self) -> None:
+        """Test service_thread_id setter with None values does nothing."""
+        thread = AgentThread()
+        thread.service_thread_id = None  # Should not raise error
+        assert thread.service_thread_id is None
+
+    def test_message_store_property_setter(self) -> None:
+        """Test message_store property setter."""
+        thread = AgentThread()
+        store = ChatMessageStore()
+
+        thread.message_store = store
+        assert thread.message_store is store
+
+    def test_message_store_setter_with_existing_service_thread_id_raises_error(self) -> None:
+        """Test that setting message_store when service_thread_id exists raises AgentThreadException."""
+        service_thread_id = "test-conversation-999"
+        thread = AgentThread(service_thread_id=service_thread_id)
+        store = ChatMessageStore()
+
+        with pytest.raises(AgentThreadException, match="Only the service_thread_id or message_store may be set"):
+            thread.message_store = store
+
+    def test_message_store_setter_with_none_values(self) -> None:
+        """Test message_store setter with None values does nothing."""
+        thread = AgentThread()
+        thread.message_store = None  # Should not raise error
+        assert thread.message_store is None
+
+    async def test_get_messages_with_message_store(self, sample_messages: list[ChatMessage]) -> None:
+        """Test get_messages when message_store is set."""
+        store = ChatMessageStore(sample_messages)
+        thread = AgentThread(message_store=store)
+
+        assert thread.message_store is not None
+
+        messages: list[ChatMessage] = await thread.message_store.list_messages()
+
+        assert messages is not None
+        assert len(messages) == 3
+        assert messages[0].text == "Hello"
+        assert messages[1].text == "Hi there!"
+        assert messages[2].text == "How are you?"
+
+    async def test_get_messages_with_no_message_store(self) -> None:
+        """Test get_messages when no message_store is set."""
+        thread = AgentThread()
+
+        assert thread.message_store is None
+
+    async def test_on_new_messages_with_service_thread_id(self, sample_message: ChatMessage) -> None:
+        """Test _on_new_messages when service_thread_id is set (should do nothing)."""
+        thread = AgentThread(service_thread_id="test-conv")
+
+        await thread.on_new_messages(sample_message)
+
+        # Should not create a message store
+        assert thread.message_store is None
+
+    async def test_on_new_messages_single_message_creates_store(self, sample_message: ChatMessage) -> None:
+        """Test _on_new_messages with single message creates ChatMessageStore."""
+        thread = AgentThread()
+
+        await thread.on_new_messages(sample_message)
+
+        assert thread.message_store is not None
+        assert isinstance(thread.message_store, ChatMessageStore)
+        messages = await thread.message_store.list_messages()
+        assert len(messages) == 1
+        assert messages[0].text == "Test message"
+
+    async def test_on_new_messages_multiple_messages(self, sample_messages: list[ChatMessage]) -> None:
+        """Test _on_new_messages with multiple messages."""
+        thread = AgentThread()
+
+        await thread.on_new_messages(sample_messages)
+
+        assert thread.message_store is not None
+        messages = await thread.message_store.list_messages()
+        assert len(messages) == 3
+
+    async def test_on_new_messages_with_existing_store(self, sample_message: ChatMessage) -> None:
+        """Test _on_new_messages adds to existing message store."""
+        initial_messages = [ChatMessage(role=Role.USER, text="Initial", message_id="init1")]
+        store = ChatMessageStore(initial_messages)
+        thread = AgentThread(message_store=store)
+
+        await thread.on_new_messages(sample_message)
+
+        assert thread.message_store is not None
+        messages = await thread.message_store.list_messages()
+        assert len(messages) == 2
+        assert messages[0].text == "Initial"
+        assert messages[1].text == "Test message"
+
+    async def test_deserialize_with_service_thread_id(self) -> None:
+        """Test _deserialize with service_thread_id."""
+        serialized_data = {"service_thread_id": "test-conv-123", "chat_message_store_state": None}
+
+        thread = await AgentThread.deserialize(serialized_data)
+
+        assert thread.service_thread_id == "test-conv-123"
+        assert thread.message_store is None
+
+    async def test_deserialize_with_store_state(self, sample_messages: list[ChatMessage]) -> None:
+        """Test _deserialize with chat_message_store_state."""
+        store_state = {"messages": sample_messages}
+        serialized_data = {"service_thread_id": None, "chat_message_store_state": store_state}
+
+        thread = await AgentThread.deserialize(serialized_data)
+
+        assert thread.service_thread_id is None
+        assert thread.message_store is not None
+        assert isinstance(thread.message_store, ChatMessageStore)
+
+    async def test_deserialize_with_no_state(self) -> None:
+        """Test _deserialize with no state."""
+        thread = AgentThread()
+        serialized_data = {"service_thread_id": None, "chat_message_store_state": None}
+
+        await thread.deserialize(serialized_data)
+
+        assert thread.service_thread_id is None
+        assert thread.message_store is None
+
+    async def test_deserialize_with_existing_store(self) -> None:
+        """Test _deserialize with existing message store."""
+        store = MockChatMessageStore()
+        thread = AgentThread(message_store=store)
+        serialized_data: dict[str, Any] = {"service_thread_id": None, "chat_message_store_state": {"messages": []}}
+
+        await thread.update_from_thread_state(serialized_data)
+
+        assert store._deserialize_calls == 1  # pyright: ignore[reportPrivateUsage]
+
+    async def test_serialize_with_service_thread_id(self) -> None:
+        """Test serialize with service_thread_id."""
+        thread = AgentThread(service_thread_id="test-conv-456")
+
+        result = await thread.serialize()
+
+        assert result["service_thread_id"] == "test-conv-456"
+        assert result["chat_message_store_state"] is None
+
+    async def test_serialize_with_message_store(self) -> None:
+        """Test serialize with message_store."""
+        store = MockChatMessageStore()
+        thread = AgentThread(message_store=store)
+
+        result = await thread.serialize()
+
+        assert result["service_thread_id"] is None
+        assert result["chat_message_store_state"] is not None
+        assert store._serialize_calls == 1  # pyright: ignore[reportPrivateUsage]
+
+    async def test_serialize_with_no_state(self) -> None:
+        """Test serialize with no state."""
+        thread = AgentThread()
+
+        result = await thread.serialize()
+
+        assert result["service_thread_id"] is None
+        assert result["chat_message_store_state"] is None
+
+    async def test_serialize_with_kwargs(self) -> None:
+        """Test serialize passes kwargs to message store."""
+        store = MockChatMessageStore()
+        thread = AgentThread(message_store=store)
+
+        await thread.serialize(custom_param="test_value")
+
+        assert store._serialize_calls == 1  # pyright: ignore[reportPrivateUsage]
+
+
+class TestChatMessageList:
+    """Test cases for ChatMessageStore class."""
+
+    def test_init_empty(self) -> None:
+        """Test ChatMessageStore initialization with no messages."""
+        store = ChatMessageStore()
+        assert len(store.messages) == 0
+
+    def test_init_with_messages(self, sample_messages: list[ChatMessage]) -> None:
+        """Test ChatMessageStore initialization with messages."""
+        store = ChatMessageStore(sample_messages)
+        assert len(store.messages) == 3
+
+    async def test_add_messages(self, sample_messages: list[ChatMessage]) -> None:
+        """Test adding messages to the store."""
+        store = ChatMessageStore()
+
+        await store.add_messages(sample_messages)
+
+        assert len(store.messages) == 3
+        messages = await store.list_messages()
+        assert messages[0].text == "Hello"
+
+    async def test_get_messages(self, sample_messages: list[ChatMessage]) -> None:
+        """Test getting messages from the store."""
+        store = ChatMessageStore(sample_messages)
+
+        messages = await store.list_messages()
+
+        assert len(messages) == 3
+        assert messages[0].message_id == "msg1"
+
+    async def test_serialize_state(self, sample_messages: list[ChatMessage]) -> None:
+        """Test serializing store state."""
+        store = ChatMessageStore(sample_messages)
+
+        result = await store.serialize()
+
+        assert "messages" in result
+        assert len(result["messages"]) == 3
+
+    async def test_serialize_state_empty(self) -> None:
+        """Test serializing empty store state."""
+        store = ChatMessageStore()
+
+        result = await store.serialize()
+
+        assert "messages" in result
+        assert len(result["messages"]) == 0
+
+    async def test_deserialize_state(self, sample_messages: list[ChatMessage]) -> None:
+        """Test deserializing store state."""
+        store = ChatMessageStore()
+        state_data = {"messages": sample_messages}
+
+        await store.update_from_state(state_data)
+
+        messages = await store.list_messages()
+        assert len(messages) == 3
+        assert messages[0].text == "Hello"
+
+    async def test_deserialize_state_none(self) -> None:
+        """Test deserializing None state."""
+        store = ChatMessageStore()
+
+        await store.update_from_state(None)
+
+        assert len(store.messages) == 0
+
+    async def test_deserialize_state_empty(self) -> None:
+        """Test deserializing empty state."""
+        store = ChatMessageStore()
+
+        await store.update_from_state({})
+
+        assert len(store.messages) == 0
+
+
+class TestStoreState:
+    """Test cases for ChatMessageStoreState class."""
+
+    def test_init(self, sample_messages: list[ChatMessage]) -> None:
+        """Test ChatMessageStoreState initialization."""
+        state = ChatMessageStoreState(messages=sample_messages)
+
+        assert len(state.messages) == 3
+        assert state.messages[0].text == "Hello"
+
+    def test_init_empty(self) -> None:
+        """Test ChatMessageStoreState initialization with empty messages."""
+        state = ChatMessageStoreState(messages=[])
+
+        assert len(state.messages) == 0
+
+
+class TestThreadState:
+    """Test cases for AgentThreadState class."""
+
+    def test_init_with_service_thread_id(self) -> None:
+        """Test AgentThreadState initialization with service_thread_id."""
+        state = AgentThreadState(service_thread_id="test-conv-123")
+
+        assert state.service_thread_id == "test-conv-123"
+        assert state.chat_message_store_state is None
+
+    def test_init_with_chat_message_store_state(self) -> None:
+        """Test AgentThreadState initialization with chat_message_store_state."""
+        store_data: dict[str, Any] = {"messages": []}
+        state = AgentThreadState(chat_message_store_state=store_data)
+
+        assert state.service_thread_id is None
+        assert state.chat_message_store_state == store_data
+
+    def test_init_with_both(self) -> None:
+        """Test AgentThreadState initialization with both parameters."""
+        store_data: dict[str, Any] = {"messages": []}
+        with pytest.raises(
+            AgentThreadException, match="Only one of service_thread_id or chat_message_store_state may be set"
+        ):
+            AgentThreadState(service_thread_id="test-conv-123", chat_message_store_state=store_data)
+
+    def test_init_defaults(self) -> None:
+        """Test AgentThreadState initialization with defaults."""
+        state = AgentThreadState()
+
+        assert state.service_thread_id is None
+        assert state.chat_message_store_state is None

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

@@ -0,0 +1,618 @@
+# Copyright (c) Microsoft. All rights reserved.
+from typing import Any
+from unittest.mock import Mock
+
+import pytest
+from opentelemetry import trace
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+from pydantic import BaseModel
+
+from agent_framework import (
+    AIFunction,
+    HostedCodeInterpreterTool,
+    HostedMCPTool,
+    ToolProtocol,
+    ai_function,
+)
+from agent_framework._tools import _parse_inputs
+from agent_framework.exceptions import ToolException
+from agent_framework.observability import OtelAttr
+
+# region AIFunction and ai_function decorator tests
+
+
+def test_ai_function_decorator():
+    """Test the ai_function decorator."""
+
+    @ai_function(name="test_tool", description="A test tool")
+    def test_tool(x: int, y: int) -> int:
+        """A simple function that adds two numbers."""
+        return x + y
+
+    assert isinstance(test_tool, ToolProtocol)
+    assert isinstance(test_tool, AIFunction)
+    assert test_tool.name == "test_tool"
+    assert test_tool.description == "A test tool"
+    assert test_tool.parameters() == {
+        "properties": {"x": {"title": "X", "type": "integer"}, "y": {"title": "Y", "type": "integer"}},
+        "required": ["x", "y"],
+        "title": "test_tool_input",
+        "type": "object",
+    }
+    assert test_tool(1, 2) == 3
+
+
+def test_ai_function_decorator_without_args():
+    """Test the ai_function decorator."""
+
+    @ai_function
+    def test_tool(x: int, y: int) -> int:
+        """A simple function that adds two numbers."""
+        return x + y
+
+    assert isinstance(test_tool, ToolProtocol)
+    assert isinstance(test_tool, AIFunction)
+    assert test_tool.name == "test_tool"
+    assert test_tool.description == "A simple function that adds two numbers."
+    assert test_tool.parameters() == {
+        "properties": {"x": {"title": "X", "type": "integer"}, "y": {"title": "Y", "type": "integer"}},
+        "required": ["x", "y"],
+        "title": "test_tool_input",
+        "type": "object",
+    }
+    assert test_tool(1, 2) == 3
+
+
+async def test_ai_function_decorator_with_async():
+    """Test the ai_function decorator with an async function."""
+
+    @ai_function(name="async_test_tool", description="An async test tool")
+    async def async_test_tool(x: int, y: int) -> int:
+        """An async function that adds two numbers."""
+        return x + y
+
+    assert isinstance(async_test_tool, ToolProtocol)
+    assert isinstance(async_test_tool, AIFunction)
+    assert async_test_tool.name == "async_test_tool"
+    assert async_test_tool.description == "An async test tool"
+    assert async_test_tool.parameters() == {
+        "properties": {"x": {"title": "X", "type": "integer"}, "y": {"title": "Y", "type": "integer"}},
+        "required": ["x", "y"],
+        "title": "async_test_tool_input",
+        "type": "object",
+    }
+    assert (await async_test_tool(1, 2)) == 3
+
+
+async def test_ai_function_invoke_telemetry_enabled(span_exporter: InMemorySpanExporter):
+    """Test the ai_function invoke method with telemetry enabled."""
+
+    @ai_function(
+        name="telemetry_test_tool",
+        description="A test tool for telemetry",
+    )
+    def telemetry_test_tool(x: int, y: int) -> int:
+        """A function that adds two numbers for telemetry testing."""
+        return x + y
+
+    # Mock the histogram
+    mock_histogram = Mock()
+    telemetry_test_tool._invocation_duration_histogram = mock_histogram
+    span_exporter.clear()
+    # Call invoke
+    result = await telemetry_test_tool.invoke(x=1, y=2, tool_call_id="test_call_id")
+
+    # Verify result
+    assert result == 3
+
+    # Verify telemetry calls
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    span = spans[0]
+    assert OtelAttr.TOOL_EXECUTION_OPERATION.value in span.name
+    assert "telemetry_test_tool" in span.name
+    assert span.attributes[OtelAttr.TOOL_NAME] == "telemetry_test_tool"
+    assert span.attributes[OtelAttr.TOOL_CALL_ID] == "test_call_id"
+    assert span.attributes[OtelAttr.TOOL_TYPE] == "function"
+    assert span.attributes[OtelAttr.TOOL_DESCRIPTION] == "A test tool for telemetry"
+    assert span.attributes[OtelAttr.TOOL_ARGUMENTS] == '{"x": 1, "y": 2}'
+    assert span.attributes[OtelAttr.TOOL_RESULT] == "3"
+
+    # Verify histogram was called with correct attributes
+    mock_histogram.record.assert_called_once()
+    call_args = mock_histogram.record.call_args
+    assert call_args[0][0] > 0  # duration should be positive
+    attributes = call_args[1]["attributes"]
+    assert attributes[OtelAttr.MEASUREMENT_FUNCTION_TAG_NAME] == "telemetry_test_tool"
+    assert attributes[OtelAttr.TOOL_CALL_ID] == "test_call_id"
+
+
+@pytest.mark.parametrize("enable_sensitive_data", [False], indirect=True)
+async def test_ai_function_invoke_telemetry_sensitive_disabled(span_exporter: InMemorySpanExporter):
+    """Test the ai_function invoke method with telemetry enabled."""
+
+    @ai_function(
+        name="telemetry_test_tool",
+        description="A test tool for telemetry",
+    )
+    def telemetry_test_tool(x: int, y: int) -> int:
+        """A function that adds two numbers for telemetry testing."""
+        return x + y
+
+    # Mock the histogram
+    mock_histogram = Mock()
+    telemetry_test_tool._invocation_duration_histogram = mock_histogram
+    span_exporter.clear()
+    # Call invoke
+    result = await telemetry_test_tool.invoke(x=1, y=2, tool_call_id="test_call_id")
+
+    # Verify result
+    assert result == 3
+
+    # Verify telemetry calls
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    span = spans[0]
+    assert OtelAttr.TOOL_EXECUTION_OPERATION.value in span.name
+    assert "telemetry_test_tool" in span.name
+    assert span.attributes[OtelAttr.TOOL_NAME] == "telemetry_test_tool"
+    assert span.attributes[OtelAttr.TOOL_CALL_ID] == "test_call_id"
+    assert span.attributes[OtelAttr.TOOL_TYPE] == "function"
+    assert span.attributes[OtelAttr.TOOL_DESCRIPTION] == "A test tool for telemetry"
+    assert OtelAttr.TOOL_ARGUMENTS not in span.attributes
+    assert OtelAttr.TOOL_RESULT not in span.attributes
+
+    # Verify histogram was called with correct attributes
+    mock_histogram.record.assert_called_once()
+    call_args = mock_histogram.record.call_args
+    assert call_args[0][0] > 0  # duration should be positive
+    attributes = call_args[1]["attributes"]
+    assert attributes[OtelAttr.MEASUREMENT_FUNCTION_TAG_NAME] == "telemetry_test_tool"
+    assert attributes[OtelAttr.TOOL_CALL_ID] == "test_call_id"
+
+
+async def test_ai_function_invoke_telemetry_with_pydantic_args(span_exporter: InMemorySpanExporter):
+    """Test the ai_function invoke method with Pydantic model arguments."""
+
+    @ai_function(
+        name="pydantic_test_tool",
+        description="A test tool with Pydantic args",
+    )
+    def pydantic_test_tool(x: int, y: int) -> int:
+        """A function that adds two numbers using Pydantic args."""
+        return x + y
+
+    # Create arguments as Pydantic model instance
+    args_model = pydantic_test_tool.input_model(x=5, y=10)
+
+    mock_histogram = Mock()
+    pydantic_test_tool._invocation_duration_histogram = mock_histogram
+    span_exporter.clear()
+    # Call invoke with Pydantic model
+    result = await pydantic_test_tool.invoke(arguments=args_model, tool_call_id="pydantic_call")
+
+    # Verify result
+    assert result == 15
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    span = spans[0]
+    assert OtelAttr.TOOL_EXECUTION_OPERATION.value in span.name
+    assert "pydantic_test_tool" in span.name
+    assert span.attributes[OtelAttr.TOOL_NAME] == "pydantic_test_tool"
+    assert span.attributes[OtelAttr.TOOL_CALL_ID] == "pydantic_call"
+    assert span.attributes[OtelAttr.TOOL_TYPE] == "function"
+    assert span.attributes[OtelAttr.TOOL_DESCRIPTION] == "A test tool with Pydantic args"
+    assert span.attributes[OtelAttr.TOOL_ARGUMENTS] == '{"x":5,"y":10}'
+
+
+async def test_ai_function_invoke_telemetry_with_exception(span_exporter: InMemorySpanExporter):
+    """Test the ai_function invoke method with telemetry when an exception occurs."""
+
+    @ai_function(
+        name="exception_test_tool",
+        description="A test tool that raises an exception",
+    )
+    def exception_test_tool(x: int, y: int) -> int:
+        """A function that raises an exception for telemetry testing."""
+        raise ValueError("Test exception for telemetry")
+
+    mock_histogram = Mock()
+    exception_test_tool._invocation_duration_histogram = mock_histogram
+    span_exporter.clear()
+    # Call invoke and expect exception
+    with pytest.raises(ValueError, match="Test exception for telemetry"):
+        await exception_test_tool.invoke(x=1, y=2, tool_call_id="exception_call")
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    span = spans[0]
+    assert OtelAttr.TOOL_EXECUTION_OPERATION.value in span.name
+    assert "exception_test_tool" in span.name
+    assert span.attributes[OtelAttr.TOOL_NAME] == "exception_test_tool"
+    assert span.attributes[OtelAttr.TOOL_CALL_ID] == "exception_call"
+    assert span.attributes[OtelAttr.TOOL_TYPE] == "function"
+    assert span.attributes[OtelAttr.TOOL_DESCRIPTION] == "A test tool that raises an exception"
+    assert span.attributes[OtelAttr.TOOL_ARGUMENTS] == '{"x": 1, "y": 2}'
+    assert span.attributes[OtelAttr.ERROR_TYPE] == ValueError.__name__
+    assert span.status.status_code == trace.StatusCode.ERROR
+
+    # Verify histogram was called with error attributes
+    mock_histogram.record.assert_called_once()
+    call_args = mock_histogram.record.call_args
+    attributes = call_args[1]["attributes"]
+    assert attributes[OtelAttr.ERROR_TYPE] == ValueError.__name__
+
+
+async def test_ai_function_invoke_telemetry_async_function(span_exporter: InMemorySpanExporter):
+    """Test the ai_function invoke method with telemetry on async function."""
+
+    @ai_function(
+        name="async_telemetry_test",
+        description="An async test tool for telemetry",
+    )
+    async def async_telemetry_test(x: int, y: int) -> int:
+        """An async function for telemetry testing."""
+        return x * y
+
+    mock_histogram = Mock()
+    async_telemetry_test._invocation_duration_histogram = mock_histogram
+    span_exporter.clear()
+    # Call invoke
+    result = await async_telemetry_test.invoke(x=3, y=4, tool_call_id="async_call")
+
+    # Verify result
+    assert result == 12
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    span = spans[0]
+    assert OtelAttr.TOOL_EXECUTION_OPERATION.value in span.name
+    assert "async_telemetry_test" in span.name
+    assert span.attributes[OtelAttr.TOOL_NAME] == "async_telemetry_test"
+    assert span.attributes[OtelAttr.TOOL_CALL_ID] == "async_call"
+    assert span.attributes[OtelAttr.TOOL_TYPE] == "function"
+    assert span.attributes[OtelAttr.TOOL_DESCRIPTION] == "An async test tool for telemetry"
+    assert span.attributes[OtelAttr.TOOL_ARGUMENTS] == '{"x": 3, "y": 4}'
+
+    # Verify histogram recording
+    mock_histogram.record.assert_called_once()
+    call_args = mock_histogram.record.call_args
+    attributes = call_args[1]["attributes"]
+    assert attributes[OtelAttr.MEASUREMENT_FUNCTION_TAG_NAME] == "async_telemetry_test"
+
+
+async def test_ai_function_invoke_invalid_pydantic_args():
+    """Test the ai_function invoke method with invalid Pydantic model arguments."""
+
+    @ai_function(name="invalid_args_test", description="A test tool for invalid args")
+    def invalid_args_test(x: int, y: int) -> int:
+        """A function for testing invalid Pydantic args."""
+        return x + y
+
+    # Create a different Pydantic model
+    class WrongModel(BaseModel):
+        a: str
+        b: str
+
+    wrong_args = WrongModel(a="hello", b="world")
+
+    # Call invoke with wrong model type
+    with pytest.raises(TypeError, match="Expected invalid_args_test_input, got WrongModel"):
+        await invalid_args_test.invoke(arguments=wrong_args)
+
+
+# region HostedCodeInterpreterTool and _parse_inputs
+
+
+def test_hosted_code_interpreter_tool_default():
+    """Test HostedCodeInterpreterTool with default parameters."""
+    tool = HostedCodeInterpreterTool()
+
+    assert tool.name == "code_interpreter"
+    assert tool.inputs == []
+    assert tool.description == ""
+    assert tool.additional_properties is None
+    assert str(tool) == "HostedCodeInterpreterTool(name=code_interpreter)"
+
+
+def test_hosted_code_interpreter_tool_with_description():
+    """Test HostedCodeInterpreterTool with description and additional properties."""
+    tool = HostedCodeInterpreterTool(
+        description="A test code interpreter",
+        additional_properties={"version": "1.0", "language": "python"},
+    )
+
+    assert tool.name == "code_interpreter"
+    assert tool.description == "A test code interpreter"
+    assert tool.additional_properties == {"version": "1.0", "language": "python"}
+
+
+def test_parse_inputs_none():
+    """Test _parse_inputs with None input."""
+    result = _parse_inputs(None)
+    assert result == []
+
+
+def test_parse_inputs_string():
+    """Test _parse_inputs with string input."""
+    from agent_framework import UriContent
+
+    result = _parse_inputs("http://example.com")
+    assert len(result) == 1
+    assert isinstance(result[0], UriContent)
+    assert result[0].uri == "http://example.com"
+    assert result[0].media_type == "text/plain"
+
+
+def test_parse_inputs_list_of_strings():
+    """Test _parse_inputs with list of strings."""
+    from agent_framework import UriContent
+
+    inputs = ["http://example.com", "https://test.org"]
+    result = _parse_inputs(inputs)
+
+    assert len(result) == 2
+    assert all(isinstance(item, UriContent) for item in result)
+    assert result[0].uri == "http://example.com"
+    assert result[1].uri == "https://test.org"
+    assert all(item.media_type == "text/plain" for item in result)
+
+
+def test_parse_inputs_uri_dict():
+    """Test _parse_inputs with URI dictionary."""
+    from agent_framework import UriContent
+
+    input_dict = {"uri": "http://example.com", "media_type": "application/json"}
+    result = _parse_inputs(input_dict)
+
+    assert len(result) == 1
+    assert isinstance(result[0], UriContent)
+    assert result[0].uri == "http://example.com"
+    assert result[0].media_type == "application/json"
+
+
+def test_parse_inputs_hosted_file_dict():
+    """Test _parse_inputs with hosted file dictionary."""
+    from agent_framework import HostedFileContent
+
+    input_dict = {"file_id": "file-123"}
+    result = _parse_inputs(input_dict)
+
+    assert len(result) == 1
+    assert isinstance(result[0], HostedFileContent)
+    assert result[0].file_id == "file-123"
+
+
+def test_parse_inputs_hosted_vector_store_dict():
+    """Test _parse_inputs with hosted vector store dictionary."""
+    from agent_framework import HostedVectorStoreContent
+
+    input_dict = {"vector_store_id": "vs-789"}
+    result = _parse_inputs(input_dict)
+
+    assert len(result) == 1
+    assert isinstance(result[0], HostedVectorStoreContent)
+    assert result[0].vector_store_id == "vs-789"
+
+
+def test_parse_inputs_data_dict():
+    """Test _parse_inputs with data dictionary."""
+    from agent_framework import DataContent
+
+    input_dict = {"data": b"test data", "media_type": "application/octet-stream"}
+    result = _parse_inputs(input_dict)
+
+    assert len(result) == 1
+    assert isinstance(result[0], DataContent)
+    assert result[0].uri == "data:application/octet-stream;base64,dGVzdCBkYXRh"
+    assert result[0].media_type == "application/octet-stream"
+
+
+def test_parse_inputs_ai_contents_instance():
+    """Test _parse_inputs with Contents instance."""
+    from agent_framework import TextContent
+
+    text_content = TextContent(text="Hello, world!")
+    result = _parse_inputs(text_content)
+
+    assert len(result) == 1
+    assert isinstance(result[0], TextContent)
+    assert result[0].text == "Hello, world!"
+
+
+def test_parse_inputs_mixed_list():
+    """Test _parse_inputs with mixed input types."""
+    from agent_framework import HostedFileContent, TextContent, UriContent
+
+    inputs = [
+        "http://example.com",  # string
+        {"uri": "https://test.org", "media_type": "text/html"},  # URI dict
+        {"file_id": "file-456"},  # hosted file dict
+        TextContent(text="Hello"),  # Contents instance
+    ]
+
+    result = _parse_inputs(inputs)
+
+    assert len(result) == 4
+    assert isinstance(result[0], UriContent)
+    assert result[0].uri == "http://example.com"
+    assert isinstance(result[1], UriContent)
+    assert result[1].uri == "https://test.org"
+    assert result[1].media_type == "text/html"
+    assert isinstance(result[2], HostedFileContent)
+    assert result[2].file_id == "file-456"
+    assert isinstance(result[3], TextContent)
+    assert result[3].text == "Hello"
+
+
+def test_parse_inputs_unsupported_dict():
+    """Test _parse_inputs with unsupported dictionary format."""
+    input_dict = {"unsupported_key": "value"}
+
+    with pytest.raises(ValueError, match="Unsupported input type"):
+        _parse_inputs(input_dict)
+
+
+def test_parse_inputs_unsupported_type():
+    """Test _parse_inputs with unsupported input type."""
+    with pytest.raises(TypeError, match="Unsupported input type: int"):
+        _parse_inputs(123)
+
+
+def test_hosted_code_interpreter_tool_with_string_input():
+    """Test HostedCodeInterpreterTool with string input."""
+    from agent_framework import UriContent
+
+    tool = HostedCodeInterpreterTool(inputs="http://example.com")
+
+    assert len(tool.inputs) == 1
+    assert isinstance(tool.inputs[0], UriContent)
+    assert tool.inputs[0].uri == "http://example.com"
+
+
+def test_hosted_code_interpreter_tool_with_dict_inputs():
+    """Test HostedCodeInterpreterTool with dictionary inputs."""
+    from agent_framework import HostedFileContent, UriContent
+
+    inputs = [{"uri": "http://example.com", "media_type": "text/html"}, {"file_id": "file-123"}]
+
+    tool = HostedCodeInterpreterTool(inputs=inputs)
+
+    assert len(tool.inputs) == 2
+    assert isinstance(tool.inputs[0], UriContent)
+    assert tool.inputs[0].uri == "http://example.com"
+    assert tool.inputs[0].media_type == "text/html"
+    assert isinstance(tool.inputs[1], HostedFileContent)
+    assert tool.inputs[1].file_id == "file-123"
+
+
+def test_hosted_code_interpreter_tool_with_ai_contents():
+    """Test HostedCodeInterpreterTool with Contents instances."""
+    from agent_framework import DataContent, TextContent
+
+    inputs = [TextContent(text="Hello, world!"), DataContent(data=b"test", media_type="text/plain")]
+
+    tool = HostedCodeInterpreterTool(inputs=inputs)
+
+    assert len(tool.inputs) == 2
+    assert isinstance(tool.inputs[0], TextContent)
+    assert tool.inputs[0].text == "Hello, world!"
+    assert isinstance(tool.inputs[1], DataContent)
+    assert tool.inputs[1].media_type == "text/plain"
+
+
+def test_hosted_code_interpreter_tool_with_single_input():
+    """Test HostedCodeInterpreterTool with single input (not in list)."""
+    from agent_framework import HostedFileContent
+
+    input_dict = {"file_id": "file-single"}
+    tool = HostedCodeInterpreterTool(inputs=input_dict)
+
+    assert len(tool.inputs) == 1
+    assert isinstance(tool.inputs[0], HostedFileContent)
+    assert tool.inputs[0].file_id == "file-single"
+
+
+def test_hosted_code_interpreter_tool_with_unknown_input():
+    """Test HostedCodeInterpreterTool with single unknown input."""
+    with pytest.raises(ValueError, match="Unsupported input type"):
+        HostedCodeInterpreterTool(inputs={"hosted_file": "file-single"})
+
+
+# region HostedMCPTool tests
+
+
+def test_hosted_mcp_tool_with_other_fields():
+    """Test creating a HostedMCPTool with a specific approval dict, headers and additional properties."""
+    tool = HostedMCPTool(
+        name="mcp-tool",
+        url="https://mcp.example",
+        description="A test MCP tool",
+        headers={"x": "y"},
+        additional_properties={"p": 1},
+    )
+
+    assert tool.name == "mcp-tool"
+    # pydantic AnyUrl preserves as string-like
+    assert str(tool.url).startswith("https://")
+    assert tool.headers == {"x": "y"}
+    assert tool.additional_properties == {"p": 1}
+    assert tool.description == "A test MCP tool"
+
+
+@pytest.mark.parametrize(
+    "approval_mode",
+    [
+        "always_require",
+        "never_require",
+        {
+            "always_require_approval": {"toolA"},
+            "never_require_approval": {"toolB"},
+        },
+        {
+            "always_require_approval": ["toolA"],
+            "never_require_approval": ("toolB",),
+        },
+    ],
+    ids=["always_require", "never_require", "specific", "specific_with_parsing"],
+)
+def test_hosted_mcp_tool_with_approval_mode(approval_mode: str | dict[str, Any]):
+    """Test creating a HostedMCPTool with a specific approval dict, headers and additional properties."""
+    tool = HostedMCPTool(name="mcp-tool", url="https://mcp.example", approval_mode=approval_mode)
+
+    assert tool.name == "mcp-tool"
+    # pydantic AnyUrl preserves as string-like
+    assert str(tool.url).startswith("https://")
+    if not isinstance(approval_mode, dict):
+        assert tool.approval_mode == approval_mode
+    else:
+        # approval_mode parsed to sets
+        assert isinstance(tool.approval_mode["always_require_approval"], set)
+        assert isinstance(tool.approval_mode["never_require_approval"], set)
+        assert "toolA" in tool.approval_mode["always_require_approval"]
+        assert "toolB" in tool.approval_mode["never_require_approval"]
+
+
+def test_hosted_mcp_tool_invalid_approval_mode_raises():
+    """Invalid approval_mode string should raise ServiceInitializationError."""
+    with pytest.raises(ToolException):
+        HostedMCPTool(name="bad", url="https://x", approval_mode="invalid_mode")
+
+
+@pytest.mark.parametrize(
+    "tools",
+    [
+        {"toolA", "toolB"},
+        ("toolA", "toolB"),
+        ["toolA", "toolB"],
+        ["toolA", "toolB", "toolA"],
+    ],
+    ids=[
+        "set",
+        "tuple",
+        "list",
+        "list_with_duplicates",
+    ],
+)
+def test_hosted_mcp_tool_with_allowed_tools(tools: list[str] | tuple[str, ...] | set[str]):
+    """Test creating a HostedMCPTool with a list of allowed tools."""
+    tool = HostedMCPTool(
+        name="mcp-tool",
+        url="https://mcp.example",
+        allowed_tools=tools,
+    )
+
+    assert tool.name == "mcp-tool"
+    # pydantic AnyUrl preserves as string-like
+    assert str(tool.url).startswith("https://")
+    # approval_mode parsed to set
+    assert isinstance(tool.allowed_tools, set)
+    assert tool.allowed_tools == {"toolA", "toolB"}
+
+
+def test_hosted_mcp_tool_with_dict_of_allowed_tools():
+    """Test creating a HostedMCPTool with a dict of allowed tools."""
+    with pytest.raises(ToolException):
+        HostedMCPTool(
+            name="mcp-tool",
+            url="https://mcp.example",
+            allowed_tools={"toolA": "Tool A", "toolC": "Tool C"},
+        )

python/packages/core/tests/main/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/openai/conftest.py

@@ -0,0 +1,51 @@
+# Copyright (c) Microsoft. All rights reserved.
+from typing import Any
+
+from pytest import fixture
+
+
+# region Connector Settings fixtures
+@fixture
+def exclude_list(request: Any) -> list[str]:
+    """Fixture that returns a list of environment variables to exclude."""
+    return request.param if hasattr(request, "param") else []
+
+
+@fixture
+def override_env_param_dict(request: Any) -> dict[str, str]:
+    """Fixture that returns a dict of environment variables to override."""
+    return request.param if hasattr(request, "param") else {}
+
+
+@fixture()
+def openai_unit_test_env(monkeypatch, exclude_list, override_env_param_dict):  # type: ignore
+    """Fixture to set environment variables for OpenAISettings."""
+
+    if exclude_list is None:
+        exclude_list = []
+
+    if override_env_param_dict is None:
+        override_env_param_dict = {}
+
+    env_vars = {
+        "OPENAI_API_KEY": "test-dummy-key",
+        "OPENAI_ORG_ID": "test_org_id",
+        "OPENAI_RESPONSES_MODEL_ID": "test_responses_model_id",
+        "OPENAI_CHAT_MODEL_ID": "test_chat_model_id",
+        "OPENAI_TEXT_MODEL_ID": "test_text_model_id",
+        "OPENAI_EMBEDDING_MODEL_ID": "test_embedding_model_id",
+        "OPENAI_TEXT_TO_IMAGE_MODEL_ID": "test_text_to_image_model_id",
+        "OPENAI_AUDIO_TO_TEXT_MODEL_ID": "test_audio_to_text_model_id",
+        "OPENAI_TEXT_TO_AUDIO_MODEL_ID": "test_text_to_audio_model_id",
+        "OPENAI_REALTIME_MODEL_ID": "test_realtime_model_id",
+    }
+
+    env_vars.update(override_env_param_dict)  # type: ignore
+
+    for key, value in env_vars.items():
+        if key in exclude_list:
+            monkeypatch.delenv(key, raising=False)  # type: ignore
+            continue
+        monkeypatch.setenv(key, value)  # type: ignore
+
+    return env_vars

python/packages/core/tests/openai/test_openai_chat_client.py

@@ -0,0 +1,775 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import os
+from typing import Annotated
+from unittest.mock import MagicMock, patch
+
+import pytest
+from openai import BadRequestError
+
+from agent_framework import (
+    AgentRunResponse,
+    AgentRunResponseUpdate,
+    ChatAgent,
+    ChatClientProtocol,
+    ChatMessage,
+    ChatOptions,
+    ChatResponse,
+    ChatResponseUpdate,
+    DataContent,
+    FunctionResultContent,
+    HostedWebSearchTool,
+    TextContent,
+    ToolProtocol,
+    ai_function,
+)
+from agent_framework.exceptions import ServiceInitializationError, ServiceResponseException
+from agent_framework.openai import OpenAIChatClient
+from agent_framework.openai._exceptions import OpenAIContentFilterException
+from agent_framework.openai._shared import prepare_function_call_results
+
+skip_if_openai_integration_tests_disabled = pytest.mark.skipif(
+    os.getenv("RUN_INTEGRATION_TESTS", "false").lower() != "true"
+    or os.getenv("OPENAI_API_KEY", "") in ("", "test-dummy-key"),
+    reason="No real OPENAI_API_KEY provided; skipping integration tests."
+    if os.getenv("RUN_INTEGRATION_TESTS", "false").lower() == "true"
+    else "Integration tests are disabled.",
+)
+
+
+def test_init(openai_unit_test_env: dict[str, str]) -> None:
+    # Test successful initialization
+    open_ai_chat_completion = OpenAIChatClient()
+
+    assert open_ai_chat_completion.ai_model_id == openai_unit_test_env["OPENAI_CHAT_MODEL_ID"]
+    assert isinstance(open_ai_chat_completion, ChatClientProtocol)
+
+
+def test_init_validation_fail() -> None:
+    # Test successful initialization
+    with pytest.raises(ServiceInitializationError):
+        OpenAIChatClient(api_key="34523", ai_model_id={"test": "dict"})  # type: ignore
+
+
+def test_init_ai_model_id_constructor(openai_unit_test_env: dict[str, str]) -> None:
+    # Test successful initialization
+    ai_model_id = "test_model_id"
+    open_ai_chat_completion = OpenAIChatClient(ai_model_id=ai_model_id)
+
+    assert open_ai_chat_completion.ai_model_id == ai_model_id
+    assert isinstance(open_ai_chat_completion, ChatClientProtocol)
+
+
+def test_init_with_default_header(openai_unit_test_env: dict[str, str]) -> None:
+    default_headers = {"X-Unit-Test": "test-guid"}
+
+    # Test successful initialization
+    open_ai_chat_completion = OpenAIChatClient(
+        default_headers=default_headers,
+    )
+
+    assert open_ai_chat_completion.ai_model_id == openai_unit_test_env["OPENAI_CHAT_MODEL_ID"]
+    assert isinstance(open_ai_chat_completion, ChatClientProtocol)
+
+    # Assert that the default header we added is present in the client's default headers
+    for key, value in default_headers.items():
+        assert key in open_ai_chat_completion.client.default_headers
+        assert open_ai_chat_completion.client.default_headers[key] == value
+
+
+def test_init_base_url(openai_unit_test_env: dict[str, str]) -> None:
+    # Test successful initialization
+    open_ai_chat_completion = OpenAIChatClient(base_url="http://localhost:1234/v1")
+    assert str(open_ai_chat_completion.client.base_url) == "http://localhost:1234/v1/"
+
+
+def test_init_base_url_from_settings_env() -> None:
+    """Test that base_url from OpenAISettings environment variable is properly used."""
+    # Set environment variable for base_url
+    with patch.dict(
+        os.environ,
+        {
+            "OPENAI_API_KEY": "dummy",
+            "OPENAI_CHAT_MODEL_ID": "gpt-5",
+            "OPENAI_BASE_URL": "https://custom-openai-endpoint.com/v1",
+        },
+    ):
+        client = OpenAIChatClient()
+        assert client.ai_model_id == "gpt-5"
+        assert str(client.client.base_url) == "https://custom-openai-endpoint.com/v1/"
+
+
+@pytest.mark.parametrize("exclude_list", [["OPENAI_CHAT_MODEL_ID"]], indirect=True)
+def test_init_with_empty_model_id(openai_unit_test_env: dict[str, str]) -> None:
+    with pytest.raises(ServiceInitializationError):
+        OpenAIChatClient(
+            env_file_path="test.env",
+        )
+
+
+@pytest.mark.parametrize("exclude_list", [["OPENAI_API_KEY"]], indirect=True)
+def test_init_with_empty_api_key(openai_unit_test_env: dict[str, str]) -> None:
+    ai_model_id = "test_model_id"
+
+    with pytest.raises(ServiceInitializationError):
+        OpenAIChatClient(
+            ai_model_id=ai_model_id,
+            env_file_path="test.env",
+        )
+
+
+def test_serialize(openai_unit_test_env: dict[str, str]) -> None:
+    default_headers = {"X-Unit-Test": "test-guid"}
+
+    settings = {
+        "ai_model_id": openai_unit_test_env["OPENAI_CHAT_MODEL_ID"],
+        "api_key": openai_unit_test_env["OPENAI_API_KEY"],
+        "default_headers": default_headers,
+    }
+
+    open_ai_chat_completion = OpenAIChatClient.from_dict(settings)
+    dumped_settings = open_ai_chat_completion.to_dict()
+    assert dumped_settings["ai_model_id"] == openai_unit_test_env["OPENAI_CHAT_MODEL_ID"]
+    assert dumped_settings["api_key"] == openai_unit_test_env["OPENAI_API_KEY"]
+    # Assert that the default header we added is present in the dumped_settings default headers
+    for key, value in default_headers.items():
+        assert key in dumped_settings["default_headers"]
+        assert dumped_settings["default_headers"][key] == value
+    # Assert that the 'User-Agent' header is not present in the dumped_settings default headers
+    assert "User-Agent" not in dumped_settings["default_headers"]
+
+
+def test_serialize_with_org_id(openai_unit_test_env: dict[str, str]) -> None:
+    settings = {
+        "ai_model_id": openai_unit_test_env["OPENAI_CHAT_MODEL_ID"],
+        "api_key": openai_unit_test_env["OPENAI_API_KEY"],
+        "org_id": openai_unit_test_env["OPENAI_ORG_ID"],
+    }
+
+    open_ai_chat_completion = OpenAIChatClient.from_dict(settings)
+    dumped_settings = open_ai_chat_completion.to_dict()
+    assert dumped_settings["ai_model_id"] == openai_unit_test_env["OPENAI_CHAT_MODEL_ID"]
+    assert dumped_settings["api_key"] == openai_unit_test_env["OPENAI_API_KEY"]
+    assert dumped_settings["org_id"] == openai_unit_test_env["OPENAI_ORG_ID"]
+    # Assert that the 'User-Agent' header is not present in the dumped_settings default headers
+    assert "User-Agent" not in dumped_settings["default_headers"]
+
+
+async def test_content_filter_exception_handling(openai_unit_test_env: dict[str, str]) -> None:
+    """Test that content filter errors are properly handled."""
+    client = OpenAIChatClient()
+    messages = [ChatMessage(role="user", text="test message")]
+
+    # Create a mock BadRequestError with content_filter code
+    mock_response = MagicMock()
+    mock_error = BadRequestError(
+        message="Content filter error", response=mock_response, body={"error": {"code": "content_filter"}}
+    )
+    mock_error.code = "content_filter"
+
+    # Mock the client to raise the content filter error
+    with (
+        patch.object(client.client.chat.completions, "create", side_effect=mock_error),
+        pytest.raises(OpenAIContentFilterException),
+    ):
+        await client._inner_get_response(messages=messages, chat_options=ChatOptions())  # type: ignore
+
+
+def test_unsupported_tool_handling(openai_unit_test_env: dict[str, str]) -> None:
+    """Test that unsupported tool types are handled correctly."""
+    client = OpenAIChatClient()
+
+    # Create a mock ToolProtocol that's not an AIFunction
+    unsupported_tool = MagicMock(spec=ToolProtocol)
+    unsupported_tool.__class__.__name__ = "UnsupportedAITool"
+
+    # This should ignore the unsupported ToolProtocol and return empty list
+    result = client._chat_to_tool_spec([unsupported_tool])  # type: ignore
+    assert result == []
+
+    # Also test with a non-ToolProtocol that should be converted to dict
+    dict_tool = {"type": "function", "name": "test"}
+    result = client._chat_to_tool_spec([dict_tool])  # type: ignore
+    assert result == [dict_tool]
+
+
+@ai_function
+def get_story_text() -> str:
+    """Returns a story about Emily and David."""
+    return (
+        "Emily and David, two passionate scientists, met during a research expedition to Antarctica. "
+        "Bonded by their love for the natural world and shared curiosity, they uncovered a "
+        "groundbreaking phenomenon in glaciology that could potentially reshape our understanding "
+        "of climate change."
+    )
+
+
+@ai_function
+def get_weather(location: str) -> str:
+    """Get the current weather for a location."""
+    return f"The weather in {location} is sunny and 72°F."
+
+
+@skip_if_openai_integration_tests_disabled
+async def test_openai_chat_completion_response() -> None:
+    """Test OpenAI chat completion responses."""
+    openai_chat_client = OpenAIChatClient()
+
+    assert isinstance(openai_chat_client, ChatClientProtocol)
+
+    messages: list[ChatMessage] = []
+    messages.append(
+        ChatMessage(
+            role="user",
+            text="Emily and David, two passionate scientists, met during a research expedition to Antarctica. "
+            "Bonded by their love for the natural world and shared curiosity, they uncovered a "
+            "groundbreaking phenomenon in glaciology that could potentially reshape our understanding "
+            "of climate change.",
+        )
+    )
+    messages.append(ChatMessage(role="user", text="who are Emily and David?"))
+
+    # Test that the client can be used to get a response
+    response = await openai_chat_client.get_response(messages=messages)
+
+    assert response is not None
+    assert isinstance(response, ChatResponse)
+    assert "scientists" in response.text
+
+
+@skip_if_openai_integration_tests_disabled
+async def test_openai_chat_completion_response_tools() -> None:
+    """Test OpenAI chat completion responses."""
+    openai_chat_client = OpenAIChatClient()
+
+    assert isinstance(openai_chat_client, ChatClientProtocol)
+
+    messages: list[ChatMessage] = []
+    messages.append(ChatMessage(role="user", text="who are Emily and David?"))
+
+    # Test that the client can be used to get a response
+    response = await openai_chat_client.get_response(
+        messages=messages,
+        tools=[get_story_text],
+        tool_choice="auto",
+    )
+
+    assert response is not None
+    assert isinstance(response, ChatResponse)
+    assert "scientists" in response.text
+
+
+@skip_if_openai_integration_tests_disabled
+async def test_openai_chat_client_streaming() -> None:
+    """Test Azure OpenAI chat completion responses."""
+    openai_chat_client = OpenAIChatClient()
+
+    assert isinstance(openai_chat_client, ChatClientProtocol)
+
+    messages: list[ChatMessage] = []
+    messages.append(
+        ChatMessage(
+            role="user",
+            text="Emily and David, two passionate scientists, met during a research expedition to Antarctica. "
+            "Bonded by their love for the natural world and shared curiosity, they uncovered a "
+            "groundbreaking phenomenon in glaciology that could potentially reshape our understanding "
+            "of climate change.",
+        )
+    )
+    messages.append(ChatMessage(role="user", text="who are Emily and David?"))
+
+    # Test that the client can be used to get a response
+    response = openai_chat_client.get_streaming_response(messages=messages)
+
+    full_message: str = ""
+    async for chunk in response:
+        assert chunk is not None
+        assert isinstance(chunk, ChatResponseUpdate)
+        assert chunk.message_id is not None
+        assert chunk.response_id is not None
+        for content in chunk.contents:
+            if isinstance(content, TextContent) and content.text:
+                full_message += content.text
+
+    assert "scientists" in full_message
+
+
+@skip_if_openai_integration_tests_disabled
+async def test_openai_chat_client_streaming_tools() -> None:
+    """Test AzureOpenAI chat completion responses."""
+    openai_chat_client = OpenAIChatClient()
+
+    assert isinstance(openai_chat_client, ChatClientProtocol)
+
+    messages: list[ChatMessage] = []
+    messages.append(ChatMessage(role="user", text="who are Emily and David?"))
+
+    # Test that the client can be used to get a response
+    response = openai_chat_client.get_streaming_response(
+        messages=messages,
+        tools=[get_story_text],
+        tool_choice="auto",
+    )
+    full_message: str = ""
+    async for chunk in response:
+        assert chunk is not None
+        assert isinstance(chunk, ChatResponseUpdate)
+        for content in chunk.contents:
+            if isinstance(content, TextContent) and content.text:
+                full_message += content.text
+
+    assert "scientists" in full_message
+
+
+@skip_if_openai_integration_tests_disabled
+async def test_openai_chat_client_web_search() -> None:
+    # Currently only a select few models support web search tool calls
+    openai_chat_client = OpenAIChatClient(ai_model_id="gpt-4o-search-preview")
+
+    assert isinstance(openai_chat_client, ChatClientProtocol)
+
+    # Test that the client will use the web search tool
+    response = await openai_chat_client.get_response(
+        messages=[
+            ChatMessage(
+                role="user",
+                text="Who are the main characters of Kpop Demon Hunters? Do a web search to find the answer.",
+            )
+        ],
+        tools=[HostedWebSearchTool()],
+        tool_choice="auto",
+    )
+
+    assert response is not None
+    assert isinstance(response, ChatResponse)
+    assert "Rumi" in response.text
+    assert "Mira" in response.text
+    assert "Zoey" in response.text
+
+    # Test that the client will use the web search tool with location
+    additional_properties = {
+        "user_location": {
+            "country": "US",
+            "city": "Seattle",
+        }
+    }
+    response = await openai_chat_client.get_response(
+        messages=[ChatMessage(role="user", text="What is the current weather? Do not ask for my current location.")],
+        tools=[HostedWebSearchTool(additional_properties=additional_properties)],
+        tool_choice="auto",
+    )
+    assert response.text is not None
+
+
+@skip_if_openai_integration_tests_disabled
+async def test_openai_chat_client_web_search_streaming() -> None:
+    openai_chat_client = OpenAIChatClient(ai_model_id="gpt-4o-search-preview")
+
+    assert isinstance(openai_chat_client, ChatClientProtocol)
+
+    # Test that the client will use the web search tool
+    response = openai_chat_client.get_streaming_response(
+        messages=[
+            ChatMessage(
+                role="user",
+                text="Who are the main characters of Kpop Demon Hunters? Do a web search to find the answer.",
+            )
+        ],
+        tools=[HostedWebSearchTool()],
+        tool_choice="auto",
+    )
+
+    assert response is not None
+    full_message: str = ""
+    async for chunk in response:
+        assert chunk is not None
+        assert isinstance(chunk, ChatResponseUpdate)
+        for content in chunk.contents:
+            if isinstance(content, TextContent) and content.text:
+                full_message += content.text
+    assert "Rumi" in full_message
+    assert "Mira" in full_message
+    assert "Zoey" in full_message
+
+    # Test that the client will use the web search tool with location
+    additional_properties = {
+        "user_location": {
+            "country": "US",
+            "city": "Seattle",
+        }
+    }
+    response = openai_chat_client.get_streaming_response(
+        messages=[ChatMessage(role="user", text="What is the current weather? Do not ask for my current location.")],
+        tools=[HostedWebSearchTool(additional_properties=additional_properties)],
+        tool_choice="auto",
+    )
+    assert response is not None
+    full_message: str = ""
+    async for chunk in response:
+        assert chunk is not None
+        assert isinstance(chunk, ChatResponseUpdate)
+        for content in chunk.contents:
+            if isinstance(content, TextContent) and content.text:
+                full_message += content.text
+    assert full_message is not None
+
+
+@skip_if_openai_integration_tests_disabled
+async def test_openai_chat_client_agent_basic_run():
+    """Test OpenAI chat client agent basic run functionality with OpenAIChatClient."""
+    async with ChatAgent(
+        chat_client=OpenAIChatClient(ai_model_id="gpt-4o-search-preview"),
+    ) as agent:
+        # Test basic run
+        response = await agent.run("Hello! Please respond with 'Hello World' exactly.")
+
+        assert isinstance(response, AgentRunResponse)
+        assert response.text is not None
+        assert len(response.text) > 0
+        assert "hello world" in response.text.lower()
+
+
+@skip_if_openai_integration_tests_disabled
+async def test_openai_chat_client_agent_basic_run_streaming():
+    """Test OpenAI chat client agent basic streaming functionality with OpenAIChatClient."""
+    async with ChatAgent(
+        chat_client=OpenAIChatClient(ai_model_id="gpt-4o-search-preview"),
+    ) as agent:
+        # Test streaming run
+        full_text = ""
+        async for chunk in agent.run_stream("Please respond with exactly: 'This is a streaming response test.'"):
+            assert isinstance(chunk, AgentRunResponseUpdate)
+            if chunk.text:
+                full_text += chunk.text
+
+        assert len(full_text) > 0
+        assert "streaming response test" in full_text.lower()
+
+
+@skip_if_openai_integration_tests_disabled
+async def test_openai_chat_client_agent_thread_persistence():
+    """Test OpenAI chat client agent thread persistence across runs with OpenAIChatClient."""
+    async with ChatAgent(
+        chat_client=OpenAIChatClient(ai_model_id="gpt-4o-search-preview"),
+        instructions="You are a helpful assistant with good memory.",
+    ) as agent:
+        # Create a new thread that will be reused
+        thread = agent.get_new_thread()
+
+        # First interaction
+        response1 = await agent.run("My name is Alice. Remember this.", thread=thread)
+
+        assert isinstance(response1, AgentRunResponse)
+        assert response1.text is not None
+
+        # Second interaction - test memory
+        response2 = await agent.run("What is my name?", thread=thread)
+
+        assert isinstance(response2, AgentRunResponse)
+        assert response2.text is not None
+        assert "alice" in response2.text.lower()
+
+
+@skip_if_openai_integration_tests_disabled
+async def test_openai_chat_client_agent_existing_thread():
+    """Test OpenAI chat client agent with existing thread to continue conversations across agent instances."""
+    # First conversation - capture the thread
+    preserved_thread = None
+
+    async with ChatAgent(
+        chat_client=OpenAIChatClient(ai_model_id="gpt-4o-search-preview"),
+        instructions="You are a helpful assistant with good memory.",
+    ) as first_agent:
+        # Start a conversation and capture the thread
+        thread = first_agent.get_new_thread()
+        first_response = await first_agent.run("My name is Alice. Remember this.", thread=thread)
+
+        assert isinstance(first_response, AgentRunResponse)
+        assert first_response.text is not None
+
+        # Preserve the thread for reuse
+        preserved_thread = thread
+
+    # Second conversation - reuse the thread in a new agent instance
+    if preserved_thread:
+        async with ChatAgent(
+            chat_client=OpenAIChatClient(ai_model_id="gpt-4o-search-preview"),
+            instructions="You are a helpful assistant with good memory.",
+        ) as second_agent:
+            # Reuse the preserved thread
+            second_response = await second_agent.run("What is my name?", thread=preserved_thread)
+
+            assert isinstance(second_response, AgentRunResponse)
+            assert second_response.text is not None
+            assert "alice" in second_response.text.lower()
+
+
+@skip_if_openai_integration_tests_disabled
+async def test_openai_chat_client_agent_level_tool_persistence():
+    """Test that agent-level tools persist across multiple runs with OpenAI Chat Client."""
+
+    async with ChatAgent(
+        chat_client=OpenAIChatClient(ai_model_id="gpt-4.1"),
+        instructions="You are a helpful assistant that uses available tools.",
+        tools=[get_weather],  # Agent-level tool
+    ) as agent:
+        # First run - agent-level tool should be available
+        first_response = await agent.run("What's the weather like in Chicago?")
+
+        assert isinstance(first_response, AgentRunResponse)
+        assert first_response.text is not None
+        # Should use the agent-level weather tool
+        assert any(term in first_response.text.lower() for term in ["chicago", "sunny", "72"])
+
+        # Second run - agent-level tool should still be available (persistence test)
+        second_response = await agent.run("What's the weather in Miami?")
+
+        assert isinstance(second_response, AgentRunResponse)
+        assert second_response.text is not None
+        # Should use the agent-level weather tool again
+        assert any(term in second_response.text.lower() for term in ["miami", "sunny", "72"])
+
+
+@skip_if_openai_integration_tests_disabled
+async def test_openai_chat_client_run_level_tool_isolation():
+    """Test that run-level tools are isolated to specific runs and don't persist with OpenAI Chat Client."""
+    # Counter to track how many times the weather tool is called
+    call_count = 0
+
+    @ai_function
+    async def get_weather_with_counter(location: Annotated[str, "The location as a city name"]) -> str:
+        """Get the current weather in a given location."""
+        nonlocal call_count
+        call_count += 1
+        return f"The weather in {location} is sunny and 72°F."
+
+    async with ChatAgent(
+        chat_client=OpenAIChatClient(ai_model_id="gpt-4.1"),
+        instructions="You are a helpful assistant.",
+    ) as agent:
+        # First run - use run-level tool
+        first_response = await agent.run(
+            "What's the weather like in Chicago?",
+            tools=[get_weather_with_counter],  # Run-level tool
+        )
+
+        assert isinstance(first_response, AgentRunResponse)
+        assert first_response.text is not None
+        # Should use the run-level weather tool (call count should be 1)
+        assert call_count == 1
+        assert any(term in first_response.text.lower() for term in ["chicago", "sunny", "72"])
+
+        # Second run - run-level tool should NOT persist (key isolation test)
+        second_response = await agent.run("What's the weather like in Miami?")
+
+        assert isinstance(second_response, AgentRunResponse)
+        assert second_response.text is not None
+        # Should NOT use the weather tool since it was only run-level in previous call
+        # Call count should still be 1 (no additional calls)
+        assert call_count == 1
+
+
+async def test_exception_message_includes_original_error_details() -> None:
+    """Test that exception messages include original error details in the new format."""
+    client = OpenAIChatClient(ai_model_id="test-model", api_key="test-key")
+    messages = [ChatMessage(role="user", text="test message")]
+
+    mock_response = MagicMock()
+    original_error_message = "Invalid API request format"
+    mock_error = BadRequestError(
+        message=original_error_message,
+        response=mock_response,
+        body={"error": {"code": "invalid_request", "message": original_error_message}},
+    )
+    mock_error.code = "invalid_request"
+
+    with (
+        patch.object(client.client.chat.completions, "create", side_effect=mock_error),
+        pytest.raises(ServiceResponseException) as exc_info,
+    ):
+        await client._inner_get_response(messages=messages, chat_options=ChatOptions())  # type: ignore
+
+    exception_message = str(exc_info.value)
+    assert "service failed to complete the prompt:" in exception_message
+    assert original_error_message in exception_message
+
+
+def test_chat_response_content_order_text_before_tool_calls(openai_unit_test_env: dict[str, str]):
+    """Test that text content appears before tool calls in ChatResponse contents."""
+    # Import locally to avoid break other tests when the import changes
+    from openai.types.chat.chat_completion import ChatCompletion, Choice
+    from openai.types.chat.chat_completion_message import ChatCompletionMessage
+    from openai.types.chat.chat_completion_message_tool_call import ChatCompletionMessageToolCall, Function
+
+    # Create a mock OpenAI response with both text and tool calls
+    mock_response = ChatCompletion(
+        id="test-response",
+        object="chat.completion",
+        created=1234567890,
+        model="gpt-4o-mini",
+        choices=[
+            Choice(
+                index=0,
+                message=ChatCompletionMessage(
+                    role="assistant",
+                    content="I'll help you with that calculation.",
+                    tool_calls=[
+                        ChatCompletionMessageToolCall(
+                            id="call-123",
+                            type="function",
+                            function=Function(name="calculate", arguments='{"x": 5, "y": 3}'),
+                        )
+                    ],
+                ),
+                finish_reason="tool_calls",
+            )
+        ],
+    )
+
+    client = OpenAIChatClient()
+    response = client._create_chat_response(mock_response, ChatOptions())
+
+    # Verify we have both text and tool call content
+    assert len(response.messages) == 1
+    message = response.messages[0]
+    assert len(message.contents) == 2
+
+    # Verify text content comes first, tool call comes second
+    assert message.contents[0].type == "text"
+    assert message.contents[0].text == "I'll help you with that calculation."
+    assert message.contents[1].type == "function_call"
+    assert message.contents[1].name == "calculate"
+
+
+def test_function_result_falsy_values_handling(openai_unit_test_env: dict[str, str]):
+    """Test that falsy values (like empty list) in function result are properly handled."""
+    client = OpenAIChatClient()
+
+    # Test with empty list (falsy but not None)
+    message_with_empty_list = ChatMessage(role="tool", contents=[FunctionResultContent(call_id="call-123", result=[])])
+
+    openai_messages = client._openai_chat_message_parser(message_with_empty_list)
+    assert len(openai_messages) == 1
+    assert openai_messages[0]["content"] == "[]"  # Empty list should be JSON serialized
+
+    # Test with empty string (falsy but not None)
+    message_with_empty_string = ChatMessage(
+        role="tool", contents=[FunctionResultContent(call_id="call-456", result="")]
+    )
+
+    openai_messages = client._openai_chat_message_parser(message_with_empty_string)
+    assert len(openai_messages) == 1
+    assert openai_messages[0]["content"] == ""  # Empty string should be preserved
+
+    # Test with False (falsy but not None)
+    message_with_false = ChatMessage(role="tool", contents=[FunctionResultContent(call_id="call-789", result=False)])
+
+    openai_messages = client._openai_chat_message_parser(message_with_false)
+    assert len(openai_messages) == 1
+    assert openai_messages[0]["content"] == "false"  # False should be JSON serialized
+
+
+def test_function_result_exception_handling(openai_unit_test_env: dict[str, str]):
+    """Test that exceptions in function result are properly handled.
+
+    Feel free to remove this test in case there's another new behavior.
+    """
+    client = OpenAIChatClient()
+
+    # Test with exception (no result)
+    test_exception = ValueError("Test error message")
+    message_with_exception = ChatMessage(
+        role="tool", contents=[FunctionResultContent(call_id="call-123", exception=test_exception)]
+    )
+
+    openai_messages = client._openai_chat_message_parser(message_with_exception)
+    assert len(openai_messages) == 1
+    assert openai_messages[0]["content"] == "Error: Test error message"
+    assert openai_messages[0]["tool_call_id"] == "call-123"
+
+
+def test_prepare_function_call_results_string_passthrough():
+    """Test that string values are passed through directly without JSON encoding."""
+    result = prepare_function_call_results("simple string")
+    assert result == "simple string"
+    assert isinstance(result, str)
+
+
+def test_openai_content_parser_data_content_image(openai_unit_test_env: dict[str, str]) -> None:
+    """Test _openai_content_parser converts DataContent with image media type to OpenAI format."""
+    client = OpenAIChatClient()
+
+    # Test DataContent with image media type
+    image_data_content = DataContent(
+        uri="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==",
+        media_type="image/png",
+    )
+
+    result = client._openai_content_parser(image_data_content)  # type: ignore
+
+    # Should convert to OpenAI image_url format
+    assert result["type"] == "image_url"
+    assert result["image_url"]["url"] == image_data_content.uri
+
+    # Test DataContent with non-image media type should use default model_dump
+    text_data_content = DataContent(uri="data:text/plain;base64,SGVsbG8gV29ybGQ=", media_type="text/plain")
+
+    result = client._openai_content_parser(text_data_content)  # type: ignore
+
+    # Should use default model_dump format
+    assert result["type"] == "data"
+    assert result["uri"] == text_data_content.uri
+    assert result["media_type"] == "text/plain"
+
+    # Test DataContent with audio media type
+    audio_data_content = DataContent(
+        uri="data:audio/wav;base64,UklGRjBEAABXQVZFZm10IBAAAAABAAEAQB8AAEAfAAABAAgAZGF0YQwEAAAAAAAAAAAA",
+        media_type="audio/wav",
+    )
+
+    result = client._openai_content_parser(audio_data_content)  # type: ignore
+
+    # Should convert to OpenAI input_audio format
+    assert result["type"] == "input_audio"
+    # Data should contain just the base64 part, not the full data URI
+    assert result["input_audio"]["data"] == "UklGRjBEAABXQVZFZm10IBAAAAABAAEAQB8AAEAfAAABAAgAZGF0YQwEAAAAAAAAAAAA"
+    assert result["input_audio"]["format"] == "wav"
+
+    # Test DataContent with MP3 audio
+    mp3_data_content = DataContent(uri="data:audio/mp3;base64,//uQAAAAWGluZwAAAA8AAAACAAACcQ==", media_type="audio/mp3")
+
+    result = client._openai_content_parser(mp3_data_content)  # type: ignore
+
+    # Should convert to OpenAI input_audio format with mp3
+    assert result["type"] == "input_audio"
+    # Data should contain just the base64 part, not the full data URI
+    assert result["input_audio"]["data"] == "//uQAAAAWGluZwAAAA8AAAACAAACcQ=="
+    assert result["input_audio"]["format"] == "mp3"
+
+    # Test DataContent with PDF file
+    pdf_data_content = DataContent(
+        uri="data:application/pdf;base64,JVBERi0xLjQKJcfsj6IKNSAwIG9iago8PC9UeXBlL0NhdGFsb2cvUGFnZXMgMiAwIFI+PgplbmRvYmoKMiAwIG9iago8PC9UeXBlL1BhZ2VzL0tpZHNbMyAwIFJdL0NvdW50IDE+PgplbmRvYmoKMyAwIG9iago8PC9UeXBlL1BhZ2UvTWVkaWFCb3ggWzAgMCA2MTIgNzkyXS9QYXJlbnQgMiAwIFIvUmVzb3VyY2VzPDwvRm9udDw8L0YxIDQgMCBSPj4+Pi9Db250ZW50cyA1IDAgUj4+CmVuZG9iago0IDAgb2JqCjw8L1R5cGUvRm9udC9TdWJ0eXBlL1R5cGUxL0Jhc2VGb250L0hlbHZldGljYT4+CmVuZG9iago1IDAgb2JqCjw8L0xlbmd0aCA0ND4+CnN0cmVhbQpCVApxCjcwIDUwIFRECi9GMSA4IFRmCihIZWxsbyBXb3JsZCEpIFRqCkVUCmVuZHN0cmVhbQplbmRvYmoKeHJlZgowIDYKMDAwMDAwMDAwMCA2NTUzNSBmIAowMDAwMDAwMDA5IDAwMDAwIG4gCjAwMDAwMDAwNTggMDAwMDAgbiAKMDAwMDAwMDExNSAwMDAwMCBuIAowMDAwMDAwMjQ1IDAwMDAwIG4gCjAwMDAwMDAzMDcgMDAwMDAgbiAKdHJhaWxlcgo8PC9TaXplIDYvUm9vdCAxIDAgUj4+CnN0YXJ0eHJlZgo0MDUKJSVFT0Y=",
+        media_type="application/pdf",
+    )
+
+    result = client._openai_content_parser(pdf_data_content)  # type: ignore
+
+    # Should convert to OpenAI file format
+    assert result["type"] == "file"
+    assert result["file"]["filename"] == "document.pdf"
+    assert "file_data" in result["file"]
+    # Base64 data should be the full data URI (OpenAI requirement)
+    assert result["file"]["file_data"].startswith("data:application/pdf;base64,")
+
+    # Test DataContent with PDF and custom filename
+    pdf_with_filename = DataContent(
+        uri="data:application/pdf;base64,JVBERi0xLjQ=",
+        media_type="application/pdf",
+        additional_properties={"filename": "report.pdf"},
+    )
+
+    result = client._openai_content_parser(pdf_with_filename)  # type: ignore
+
+    # Should use custom filename
+    assert result["type"] == "file"
+    assert result["file"]["filename"] == "report.pdf"

python/packages/core/tests/openai/test_openai_chat_client_base.py

@@ -0,0 +1,354 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from copy import deepcopy
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from openai import AsyncStream
+from openai.resources.chat.completions import AsyncCompletions as AsyncChatCompletions
+from openai.types.chat import ChatCompletion, ChatCompletionChunk
+from openai.types.chat.chat_completion import Choice
+from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
+from openai.types.chat.chat_completion_chunk import ChoiceDelta as ChunkChoiceDelta
+from openai.types.chat.chat_completion_message import ChatCompletionMessage
+from pydantic import BaseModel
+
+from agent_framework import ChatMessage, ChatResponseUpdate
+from agent_framework.exceptions import (
+    ServiceResponseException,
+)
+from agent_framework.openai import OpenAIChatClient
+
+
+async def mock_async_process_chat_stream_response(_):
+    mock_content = MagicMock(spec=ChatResponseUpdate)
+    yield mock_content, None
+
+
+@pytest.fixture(scope="function")
+def chat_history() -> list[ChatMessage]:
+    return []
+
+
+@pytest.fixture
+def mock_chat_completion_response() -> ChatCompletion:
+    return ChatCompletion(
+        id="test_id",
+        choices=[
+            Choice(index=0, message=ChatCompletionMessage(content="test", role="assistant"), finish_reason="stop")
+        ],
+        created=0,
+        model="test",
+        object="chat.completion",
+    )
+
+
+@pytest.fixture
+def mock_streaming_chat_completion_response() -> AsyncStream[ChatCompletionChunk]:
+    content = ChatCompletionChunk(
+        id="test_id",
+        choices=[ChunkChoice(index=0, delta=ChunkChoiceDelta(content="test", role="assistant"), finish_reason="stop")],
+        created=0,
+        model="test",
+        object="chat.completion.chunk",
+    )
+    stream = MagicMock(spec=AsyncStream)
+    stream.__aiter__.return_value = [content]
+    return stream
+
+
+# region Chat Message Content
+
+
+@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
+async def test_cmc(
+    mock_create: AsyncMock,
+    chat_history: list[ChatMessage],
+    mock_chat_completion_response: ChatCompletion,
+    openai_unit_test_env: dict[str, str],
+):
+    mock_create.return_value = mock_chat_completion_response
+    chat_history.append(ChatMessage(role="user", text="hello world"))
+
+    openai_chat_completion = OpenAIChatClient()
+    await openai_chat_completion.get_response(
+        messages=chat_history,
+    )
+    mock_create.assert_awaited_once_with(
+        model=openai_unit_test_env["OPENAI_CHAT_MODEL_ID"],
+        stream=False,
+        messages=openai_chat_completion._prepare_chat_history_for_request(chat_history),  # type: ignore
+    )
+
+
+@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
+async def test_cmc_chat_options(
+    mock_create: AsyncMock,
+    chat_history: list[ChatMessage],
+    mock_chat_completion_response: ChatCompletion,
+    openai_unit_test_env: dict[str, str],
+):
+    mock_create.return_value = mock_chat_completion_response
+    chat_history.append(ChatMessage(role="user", text="hello world"))
+
+    openai_chat_completion = OpenAIChatClient()
+    await openai_chat_completion.get_response(
+        messages=chat_history,
+    )
+    mock_create.assert_awaited_once_with(
+        model=openai_unit_test_env["OPENAI_CHAT_MODEL_ID"],
+        stream=False,
+        messages=openai_chat_completion._prepare_chat_history_for_request(chat_history),  # type: ignore
+    )
+
+
+@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
+async def test_cmc_no_fcc_in_response(
+    mock_create: AsyncMock,
+    chat_history: list[ChatMessage],
+    mock_chat_completion_response: ChatCompletion,
+    openai_unit_test_env: dict[str, str],
+):
+    mock_create.return_value = mock_chat_completion_response
+    chat_history.append(ChatMessage(role="user", text="hello world"))
+    orig_chat_history = deepcopy(chat_history)
+
+    openai_chat_completion = OpenAIChatClient()
+    await openai_chat_completion.get_response(
+        messages=chat_history,
+        arguments={},
+    )
+    mock_create.assert_awaited_once_with(
+        model=openai_unit_test_env["OPENAI_CHAT_MODEL_ID"],
+        stream=False,
+        messages=openai_chat_completion._prepare_chat_history_for_request(orig_chat_history),  # type: ignore
+    )
+
+
+@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
+async def test_cmc_structured_output_no_fcc(
+    mock_create: AsyncMock,
+    chat_history: list[ChatMessage],
+    mock_chat_completion_response: ChatCompletion,
+    openai_unit_test_env: dict[str, str],
+):
+    mock_create.return_value = mock_chat_completion_response
+    chat_history.append(ChatMessage(role="user", text="hello world"))
+
+    # Define a mock response format
+    class Test(BaseModel):
+        name: str
+
+    openai_chat_completion = OpenAIChatClient()
+    await openai_chat_completion.get_response(
+        messages=chat_history,
+        response_format=Test,
+    )
+    mock_create.assert_awaited_once()
+
+
+@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
+async def test_scmc_chat_options(
+    mock_create: AsyncMock,
+    chat_history: list[ChatMessage],
+    mock_streaming_chat_completion_response: AsyncStream[ChatCompletionChunk],
+    openai_unit_test_env: dict[str, str],
+):
+    mock_create.return_value = mock_streaming_chat_completion_response
+    chat_history.append(ChatMessage(role="user", text="hello world"))
+
+    openai_chat_completion = OpenAIChatClient()
+    async for msg in openai_chat_completion.get_streaming_response(
+        messages=chat_history,
+    ):
+        assert isinstance(msg, ChatResponseUpdate)
+        assert msg.message_id is not None
+        assert msg.response_id is not None
+    mock_create.assert_awaited_once_with(
+        model=openai_unit_test_env["OPENAI_CHAT_MODEL_ID"],
+        stream=True,
+        stream_options={"include_usage": True},
+        messages=openai_chat_completion._prepare_chat_history_for_request(chat_history),  # type: ignore
+    )
+
+
+@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock, side_effect=Exception)
+async def test_cmc_general_exception(
+    mock_create: AsyncMock,
+    chat_history: list[ChatMessage],
+    mock_chat_completion_response: ChatCompletion,
+    openai_unit_test_env: dict[str, str],
+):
+    mock_create.return_value = mock_chat_completion_response
+    chat_history.append(ChatMessage(role="user", text="hello world"))
+
+    openai_chat_completion = OpenAIChatClient()
+    with pytest.raises(ServiceResponseException):
+        await openai_chat_completion.get_response(
+            messages=chat_history,
+        )
+
+
+# region Streaming
+
+
+@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
+async def test_get_streaming(
+    mock_create: AsyncMock,
+    chat_history: list[ChatMessage],
+    openai_unit_test_env: dict[str, str],
+):
+    content1 = ChatCompletionChunk(
+        id="test_id",
+        choices=[],
+        created=0,
+        model="test",
+        object="chat.completion.chunk",
+    )
+    content2 = ChatCompletionChunk(
+        id="test_id",
+        choices=[ChunkChoice(index=0, delta=ChunkChoiceDelta(content="test", role="assistant"), finish_reason="stop")],
+        created=0,
+        model="test",
+        object="chat.completion.chunk",
+    )
+    stream = MagicMock(spec=AsyncStream)
+    stream.__aiter__.return_value = [content1, content2]
+    mock_create.return_value = stream
+    chat_history.append(ChatMessage(role="user", text="hello world"))
+    orig_chat_history = deepcopy(chat_history)
+
+    openai_chat_completion = OpenAIChatClient()
+    async for msg in openai_chat_completion.get_streaming_response(
+        messages=chat_history,
+    ):
+        assert isinstance(msg, ChatResponseUpdate)
+    mock_create.assert_awaited_once_with(
+        model=openai_unit_test_env["OPENAI_CHAT_MODEL_ID"],
+        stream=True,
+        stream_options={"include_usage": True},
+        messages=openai_chat_completion._prepare_chat_history_for_request(orig_chat_history),  # type: ignore
+    )
+
+
+@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
+async def test_get_streaming_singular(
+    mock_create: AsyncMock,
+    chat_history: list[ChatMessage],
+    openai_unit_test_env: dict[str, str],
+):
+    content1 = ChatCompletionChunk(
+        id="test_id",
+        choices=[],
+        created=0,
+        model="test",
+        object="chat.completion.chunk",
+    )
+    content2 = ChatCompletionChunk(
+        id="test_id",
+        choices=[ChunkChoice(index=0, delta=ChunkChoiceDelta(content="test", role="assistant"), finish_reason="stop")],
+        created=0,
+        model="test",
+        object="chat.completion.chunk",
+    )
+    stream = MagicMock(spec=AsyncStream)
+    stream.__aiter__.return_value = [content1, content2]
+    mock_create.return_value = stream
+    chat_history.append(ChatMessage(role="user", text="hello world"))
+    orig_chat_history = deepcopy(chat_history)
+
+    openai_chat_completion = OpenAIChatClient()
+    async for msg in openai_chat_completion.get_streaming_response(
+        messages=chat_history,
+    ):
+        assert isinstance(msg, ChatResponseUpdate)
+    mock_create.assert_awaited_once_with(
+        model=openai_unit_test_env["OPENAI_CHAT_MODEL_ID"],
+        stream=True,
+        stream_options={"include_usage": True},
+        messages=openai_chat_completion._prepare_chat_history_for_request(orig_chat_history),  # type: ignore
+    )
+
+
+@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
+async def test_get_streaming_structured_output_no_fcc(
+    mock_create: AsyncMock,
+    chat_history: list[ChatMessage],
+    openai_unit_test_env: dict[str, str],
+):
+    content1 = ChatCompletionChunk(
+        id="test_id",
+        choices=[],
+        created=0,
+        model="test",
+        object="chat.completion.chunk",
+    )
+    content2 = ChatCompletionChunk(
+        id="test_id",
+        choices=[ChunkChoice(index=0, delta=ChunkChoiceDelta(content="test", role="assistant"), finish_reason="stop")],
+        created=0,
+        model="test",
+        object="chat.completion.chunk",
+    )
+    stream = MagicMock(spec=AsyncStream)
+    stream.__aiter__.return_value = [content1, content2]
+    mock_create.return_value = stream
+    chat_history.append(ChatMessage(role="user", text="hello world"))
+
+    # Define a mock response format
+    class Test(BaseModel):
+        name: str
+
+    openai_chat_completion = OpenAIChatClient()
+    async for msg in openai_chat_completion.get_streaming_response(
+        messages=chat_history,
+        response_format=Test,
+    ):
+        assert isinstance(msg, ChatResponseUpdate)
+    mock_create.assert_awaited_once()
+
+
+@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
+async def test_get_streaming_no_fcc_in_response(
+    mock_create: AsyncMock,
+    chat_history: list[ChatMessage],
+    mock_streaming_chat_completion_response: ChatCompletion,
+    openai_unit_test_env: dict[str, str],
+):
+    mock_create.return_value = mock_streaming_chat_completion_response
+    chat_history.append(ChatMessage(role="user", text="hello world"))
+    orig_chat_history = deepcopy(chat_history)
+
+    openai_chat_completion = OpenAIChatClient()
+    [
+        msg
+        async for msg in openai_chat_completion.get_streaming_response(
+            messages=chat_history,
+        )
+    ]
+    mock_create.assert_awaited_once_with(
+        model=openai_unit_test_env["OPENAI_CHAT_MODEL_ID"],
+        stream=True,
+        stream_options={"include_usage": True},
+        messages=openai_chat_completion._prepare_chat_history_for_request(orig_chat_history),  # type: ignore
+    )
+
+
+@patch.object(AsyncChatCompletions, "create", new_callable=AsyncMock)
+async def test_get_streaming_no_stream(
+    mock_create: AsyncMock,
+    chat_history: list[ChatMessage],
+    openai_unit_test_env: dict[str, str],
+    mock_chat_completion_response: ChatCompletion,  # AsyncStream[ChatCompletionChunk]?
+):
+    mock_create.return_value = mock_chat_completion_response
+    chat_history.append(ChatMessage(role="user", text="hello world"))
+
+    openai_chat_completion = OpenAIChatClient()
+    with pytest.raises(ServiceResponseException):
+        [
+            msg
+            async for msg in openai_chat_completion.get_streaming_response(
+                messages=chat_history,
+            )
+        ]

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

@@ -0,0 +1,334 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import json
+import tempfile
+from datetime import datetime, timezone
+from pathlib import Path
+
+from agent_framework import (
+    FileCheckpointStorage,
+    InMemoryCheckpointStorage,
+    WorkflowCheckpoint,
+)
+
+
+def test_workflow_checkpoint_default_values():
+    checkpoint = WorkflowCheckpoint()
+
+    assert checkpoint.checkpoint_id != ""
+    assert checkpoint.workflow_id == ""
+    assert checkpoint.timestamp != ""
+    assert checkpoint.messages == {}
+    assert checkpoint.shared_state == {}
+    assert checkpoint.executor_states == {}
+    assert checkpoint.iteration_count == 0
+    assert checkpoint.max_iterations == 100
+    assert checkpoint.metadata == {}
+    assert checkpoint.version == "1.0"
+
+
+def test_workflow_checkpoint_custom_values():
+    custom_timestamp = datetime.now(timezone.utc).isoformat()
+    checkpoint = WorkflowCheckpoint(
+        checkpoint_id="test-checkpoint-123",
+        workflow_id="test-workflow-456",
+        timestamp=custom_timestamp,
+        messages={"executor1": [{"data": "test"}]},
+        shared_state={"key": "value"},
+        executor_states={"executor1": {"state": "active"}},
+        iteration_count=5,
+        max_iterations=50,
+        metadata={"test": True},
+        version="2.0",
+    )
+
+    assert checkpoint.checkpoint_id == "test-checkpoint-123"
+    assert checkpoint.workflow_id == "test-workflow-456"
+    assert checkpoint.timestamp == custom_timestamp
+    assert checkpoint.messages == {"executor1": [{"data": "test"}]}
+    assert checkpoint.shared_state == {"key": "value"}
+    assert checkpoint.executor_states == {"executor1": {"state": "active"}}
+    assert checkpoint.iteration_count == 5
+    assert checkpoint.max_iterations == 50
+    assert checkpoint.metadata == {"test": True}
+    assert checkpoint.version == "2.0"
+
+
+async def test_memory_checkpoint_storage_save_and_load():
+    storage = InMemoryCheckpointStorage()
+    checkpoint = WorkflowCheckpoint(workflow_id="test-workflow", messages={"executor1": [{"data": "hello"}]})
+
+    # Save checkpoint
+    saved_id = await storage.save_checkpoint(checkpoint)
+    assert saved_id == checkpoint.checkpoint_id
+
+    # Load checkpoint
+    loaded_checkpoint = await storage.load_checkpoint(checkpoint.checkpoint_id)
+    assert loaded_checkpoint is not None
+    assert loaded_checkpoint.checkpoint_id == checkpoint.checkpoint_id
+    assert loaded_checkpoint.workflow_id == checkpoint.workflow_id
+    assert loaded_checkpoint.messages == checkpoint.messages
+
+
+async def test_memory_checkpoint_storage_load_nonexistent():
+    storage = InMemoryCheckpointStorage()
+
+    result = await storage.load_checkpoint("nonexistent-id")
+    assert result is None
+
+
+async def test_memory_checkpoint_storage_list_checkpoints():
+    storage = InMemoryCheckpointStorage()
+
+    # Create checkpoints for different workflows
+    checkpoint1 = WorkflowCheckpoint(workflow_id="workflow-1")
+    checkpoint2 = WorkflowCheckpoint(workflow_id="workflow-1")
+    checkpoint3 = WorkflowCheckpoint(workflow_id="workflow-2")
+
+    await storage.save_checkpoint(checkpoint1)
+    await storage.save_checkpoint(checkpoint2)
+    await storage.save_checkpoint(checkpoint3)
+
+    # Test list_checkpoint_ids for workflow-1
+    workflow1_checkpoint_ids = await storage.list_checkpoint_ids("workflow-1")
+    assert len(workflow1_checkpoint_ids) == 2
+    assert checkpoint1.checkpoint_id in workflow1_checkpoint_ids
+    assert checkpoint2.checkpoint_id in workflow1_checkpoint_ids
+
+    # Test list_checkpoints for workflow-1 (returns objects)
+    workflow1_checkpoints = await storage.list_checkpoints("workflow-1")
+    assert len(workflow1_checkpoints) == 2
+    assert all(isinstance(cp, WorkflowCheckpoint) for cp in workflow1_checkpoints)
+    assert {cp.checkpoint_id for cp in workflow1_checkpoints} == {checkpoint1.checkpoint_id, checkpoint2.checkpoint_id}
+
+    # Test list_checkpoint_ids for workflow-2
+    workflow2_checkpoint_ids = await storage.list_checkpoint_ids("workflow-2")
+    assert len(workflow2_checkpoint_ids) == 1
+    assert checkpoint3.checkpoint_id in workflow2_checkpoint_ids
+
+    # Test list_checkpoints for workflow-2 (returns objects)
+    workflow2_checkpoints = await storage.list_checkpoints("workflow-2")
+    assert len(workflow2_checkpoints) == 1
+    assert workflow2_checkpoints[0].checkpoint_id == checkpoint3.checkpoint_id
+
+    # Test list_checkpoint_ids for non-existent workflow
+    empty_checkpoint_ids = await storage.list_checkpoint_ids("nonexistent-workflow")
+    assert len(empty_checkpoint_ids) == 0
+
+    # Test list_checkpoints for non-existent workflow
+    empty_checkpoints = await storage.list_checkpoints("nonexistent-workflow")
+    assert len(empty_checkpoints) == 0
+
+    # Test list_checkpoint_ids without workflow filter (all checkpoints)
+    all_checkpoint_ids = await storage.list_checkpoint_ids()
+    assert len(all_checkpoint_ids) == 3
+    expected_ids = {checkpoint1.checkpoint_id, checkpoint2.checkpoint_id, checkpoint3.checkpoint_id}
+    assert expected_ids.issubset(set(all_checkpoint_ids))
+
+    # Test list_checkpoints without workflow filter (all checkpoints)
+    all_checkpoints = await storage.list_checkpoints()
+    assert len(all_checkpoints) == 3
+    assert all(isinstance(cp, WorkflowCheckpoint) for cp in all_checkpoints)
+
+
+async def test_memory_checkpoint_storage_delete():
+    storage = InMemoryCheckpointStorage()
+    checkpoint = WorkflowCheckpoint(workflow_id="test-workflow")
+
+    # Save checkpoint
+    await storage.save_checkpoint(checkpoint)
+    assert await storage.load_checkpoint(checkpoint.checkpoint_id) is not None
+
+    # Delete checkpoint
+    result = await storage.delete_checkpoint(checkpoint.checkpoint_id)
+    assert result is True
+
+    # Verify deletion
+    assert await storage.load_checkpoint(checkpoint.checkpoint_id) is None
+
+    # Try to delete again
+    result = await storage.delete_checkpoint(checkpoint.checkpoint_id)
+    assert result is False
+
+
+async def test_file_checkpoint_storage_save_and_load():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+        checkpoint = WorkflowCheckpoint(
+            workflow_id="test-workflow",
+            messages={"executor1": [{"data": "hello", "source_id": "test", "target_id": None}]},
+            shared_state={"key": "value"},
+        )
+
+        # Save checkpoint
+        saved_id = await storage.save_checkpoint(checkpoint)
+        assert saved_id == checkpoint.checkpoint_id
+
+        # Verify file was created
+        file_path = Path(temp_dir) / f"{checkpoint.checkpoint_id}.json"
+        assert file_path.exists()
+
+        # Load checkpoint
+        loaded_checkpoint = await storage.load_checkpoint(checkpoint.checkpoint_id)
+        assert loaded_checkpoint is not None
+        assert loaded_checkpoint.checkpoint_id == checkpoint.checkpoint_id
+        assert loaded_checkpoint.workflow_id == checkpoint.workflow_id
+        assert loaded_checkpoint.messages == checkpoint.messages
+        assert loaded_checkpoint.shared_state == checkpoint.shared_state
+
+
+async def test_file_checkpoint_storage_load_nonexistent():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+
+        result = await storage.load_checkpoint("nonexistent-id")
+        assert result is None
+
+
+async def test_file_checkpoint_storage_list_checkpoints():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+
+        # Create checkpoints for different workflows
+        checkpoint1 = WorkflowCheckpoint(workflow_id="workflow-1")
+        checkpoint2 = WorkflowCheckpoint(workflow_id="workflow-1")
+        checkpoint3 = WorkflowCheckpoint(workflow_id="workflow-2")
+
+        await storage.save_checkpoint(checkpoint1)
+        await storage.save_checkpoint(checkpoint2)
+        await storage.save_checkpoint(checkpoint3)
+
+        # Test list_checkpoint_ids for workflow-1
+        workflow1_checkpoint_ids = await storage.list_checkpoint_ids("workflow-1")
+        assert len(workflow1_checkpoint_ids) == 2
+        assert checkpoint1.checkpoint_id in workflow1_checkpoint_ids
+        assert checkpoint2.checkpoint_id in workflow1_checkpoint_ids
+
+        # Test list_checkpoints for workflow-1 (returns objects)
+        workflow1_checkpoints = await storage.list_checkpoints("workflow-1")
+        assert len(workflow1_checkpoints) == 2
+        assert all(isinstance(cp, WorkflowCheckpoint) for cp in workflow1_checkpoints)
+        checkpoint_ids = {cp.checkpoint_id for cp in workflow1_checkpoints}
+        assert checkpoint_ids == {checkpoint1.checkpoint_id, checkpoint2.checkpoint_id}
+
+        # Test list_checkpoint_ids for workflow-2
+        workflow2_checkpoint_ids = await storage.list_checkpoint_ids("workflow-2")
+        assert len(workflow2_checkpoint_ids) == 1
+        assert checkpoint3.checkpoint_id in workflow2_checkpoint_ids
+
+        # Test list_checkpoints for workflow-2 (returns objects)
+        workflow2_checkpoints = await storage.list_checkpoints("workflow-2")
+        assert len(workflow2_checkpoints) == 1
+        assert workflow2_checkpoints[0].checkpoint_id == checkpoint3.checkpoint_id
+
+        # Test list all checkpoints
+        all_checkpoint_ids = await storage.list_checkpoint_ids()
+        assert len(all_checkpoint_ids) == 3
+
+        all_checkpoints = await storage.list_checkpoints()
+        assert len(all_checkpoints) == 3
+        assert all(isinstance(cp, WorkflowCheckpoint) for cp in all_checkpoints)
+
+
+async def test_file_checkpoint_storage_delete():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+        checkpoint = WorkflowCheckpoint(workflow_id="test-workflow")
+
+        # Save checkpoint
+        await storage.save_checkpoint(checkpoint)
+        file_path = Path(temp_dir) / f"{checkpoint.checkpoint_id}.json"
+        assert file_path.exists()
+
+        # Delete checkpoint
+        result = await storage.delete_checkpoint(checkpoint.checkpoint_id)
+        assert result is True
+        assert not file_path.exists()
+
+        # Try to delete again
+        result = await storage.delete_checkpoint(checkpoint.checkpoint_id)
+        assert result is False
+
+
+async def test_file_checkpoint_storage_directory_creation():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        nested_path = Path(temp_dir) / "nested" / "checkpoint" / "storage"
+        storage = FileCheckpointStorage(nested_path)
+
+        # Directory should be created
+        assert nested_path.exists()
+        assert nested_path.is_dir()
+
+        # Should be able to save checkpoints
+        checkpoint = WorkflowCheckpoint(workflow_id="test")
+        await storage.save_checkpoint(checkpoint)
+
+        file_path = nested_path / f"{checkpoint.checkpoint_id}.json"
+        assert file_path.exists()
+
+
+async def test_file_checkpoint_storage_corrupted_file():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+
+        # Create a corrupted JSON file
+        corrupted_file = Path(temp_dir) / "corrupted.json"
+        with open(corrupted_file, "w") as f:  # noqa: ASYNC230
+            f.write("{ invalid json }")
+
+        # list_checkpoints should handle the corrupted file gracefully
+        checkpoints = await storage.list_checkpoints("any-workflow")
+        assert checkpoints == []
+
+
+async def test_file_checkpoint_storage_json_serialization():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+
+        # Create checkpoint with complex nested data
+        checkpoint = WorkflowCheckpoint(
+            workflow_id="complex-workflow",
+            messages={"executor1": [{"data": {"nested": {"value": 42}}, "source_id": "test", "target_id": None}]},
+            shared_state={"list": [1, 2, 3], "dict": {"a": "b", "c": {"d": "e"}}, "bool": True, "null": None},
+            executor_states={"executor1": {"state": "active", "config": {"timeout": 30, "retries": 3}}},
+        )
+
+        # Save and load
+        await storage.save_checkpoint(checkpoint)
+        loaded = await storage.load_checkpoint(checkpoint.checkpoint_id)
+
+        assert loaded is not None
+        assert loaded.messages == checkpoint.messages
+        assert loaded.shared_state == checkpoint.shared_state
+        assert loaded.executor_states == checkpoint.executor_states
+
+        # Verify the JSON file is properly formatted
+        file_path = Path(temp_dir) / f"{checkpoint.checkpoint_id}.json"
+        with open(file_path) as f:  # noqa: ASYNC230
+            data = json.load(f)
+
+        assert data["messages"]["executor1"][0]["data"]["nested"]["value"] == 42
+        assert data["shared_state"]["list"] == [1, 2, 3]
+        assert data["shared_state"]["bool"] is True
+        assert data["shared_state"]["null"] is None
+
+
+def test_checkpoint_storage_protocol_compliance():
+    # This test ensures both implementations have all required methods
+    memory_storage = InMemoryCheckpointStorage()
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        file_storage = FileCheckpointStorage(temp_dir)
+
+        for storage in [memory_storage, file_storage]:
+            # Test that all protocol methods exist and are callable
+            assert hasattr(storage, "save_checkpoint")
+            assert callable(storage.save_checkpoint)
+            assert hasattr(storage, "load_checkpoint")
+            assert callable(storage.load_checkpoint)
+            assert hasattr(storage, "list_checkpoint_ids")
+            assert callable(storage.list_checkpoint_ids)
+            assert hasattr(storage, "list_checkpoints")
+            assert callable(storage.list_checkpoints)
+            assert hasattr(storage, "delete_checkpoint")
+            assert callable(storage.delete_checkpoint)

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

@@ -0,0 +1,50 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from dataclasses import dataclass
+from typing import Any, cast
+
+from agent_framework._workflow._executor import RequestInfoMessage, RequestResponse
+from agent_framework._workflow._runner_context import _decode_checkpoint_value, _encode_checkpoint_value  # type: ignore
+from agent_framework._workflow._typing_utils import is_instance_of
+
+
+@dataclass(kw_only=True)
+class SampleRequest(RequestInfoMessage):
+    prompt: str
+
+
+def test_decode_dataclass_with_nested_request() -> None:
+    original = RequestResponse[SampleRequest, str](
+        data="approve",
+        original_request=SampleRequest(request_id="abc", prompt="prompt"),
+        request_id="abc",
+    )
+
+    encoded = _encode_checkpoint_value(original)
+    decoded = cast(RequestResponse[SampleRequest, str], _decode_checkpoint_value(encoded))
+
+    assert isinstance(decoded, RequestResponse)
+    assert decoded.data == "approve"
+    assert decoded.request_id == "abc"
+    assert isinstance(decoded.original_request, SampleRequest)
+    assert decoded.original_request.prompt == "prompt"
+
+
+def test_is_instance_of_coerces_request_response_original_request_dict() -> None:
+    response = RequestResponse[SampleRequest, str](
+        data="approve",
+        original_request=SampleRequest(request_id="req-1", prompt="prompt"),
+        request_id="req-1",
+    )
+
+    # Simulate checkpoint decode fallback leaving a dict
+    response.original_request = cast(
+        Any,
+        {
+            "request_id": "req-1",
+            "prompt": "prompt",
+        },
+    )
+
+    assert is_instance_of(response, RequestResponse[SampleRequest, str])
+    assert isinstance(response.original_request, SampleRequest)

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

@@ -0,0 +1,74 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import pytest
+from typing_extensions import Never
+
+from agent_framework import WorkflowBuilder, WorkflowContext, WorkflowRunState, WorkflowStatusEvent, handler
+from agent_framework._workflow._checkpoint import InMemoryCheckpointStorage
+from agent_framework._workflow._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:
+        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).set_start_executor(start).add_edge(start, finish)
+    builder = builder.with_checkpointing(checkpoint_storage=storage)
+    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_stream("hello")]  # noqa: F841
+
+    checkpoints = await storage.list_checkpoints()
+    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(ValueError, match="Workflow graph has changed"):
+        _ = [
+            event
+            async for event in mismatched_workflow.run_stream_from_checkpoint(
+                target_checkpoint.checkpoint_id,
+                checkpoint_storage=storage,
+            )
+        ]
+
+
+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_stream("hello")]  # noqa: F841
+
+    checkpoints = sorted(await storage.list_checkpoints(), 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_stream_from_checkpoint(
+            target_checkpoint.checkpoint_id,
+            checkpoint_storage=storage,
+        )
+    ]
+
+    assert any(isinstance(event, WorkflowStatusEvent) and event.state == WorkflowRunState.IDLE for event in events)

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

@@ -0,0 +1,209 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from typing import Any, cast
+
+import pytest
+
+from agent_framework import (
+    AgentExecutorRequest,
+    AgentExecutorResponse,
+    AgentRunResponse,
+    ChatMessage,
+    ConcurrentBuilder,
+    Executor,
+    Role,
+    WorkflowContext,
+    WorkflowOutputEvent,
+    WorkflowRunState,
+    WorkflowStatusEvent,
+    handler,
+)
+from agent_framework._workflow._checkpoint import InMemoryCheckpointStorage
+
+
+class _FakeAgentExec(Executor):
+    """Test executor that mimics an agent by emitting an AgentExecutorResponse.
+
+    It takes the incoming AgentExecutorRequest, produces a single assistant message
+    with the configured reply text, and sends an AgentExecutorResponse that includes
+    full_conversation (the original user prompt followed by the assistant message).
+    """
+
+    def __init__(self, id: str, reply_text: str) -> None:
+        super().__init__(id)
+        self._reply_text = reply_text
+
+    @handler
+    async def run(self, request: AgentExecutorRequest, ctx: WorkflowContext[AgentExecutorResponse]) -> None:
+        response = AgentRunResponse(messages=ChatMessage(Role.ASSISTANT, text=self._reply_text))
+        full_conversation = list(request.messages) + list(response.messages)
+        await ctx.send_message(AgentExecutorResponse(self.id, response, full_conversation=full_conversation))
+
+
+def test_concurrent_builder_rejects_empty_participants() -> None:
+    with pytest.raises(ValueError):
+        ConcurrentBuilder().participants([])
+
+
+def test_concurrent_builder_rejects_duplicate_executors() -> None:
+    a = _FakeAgentExec("dup", "A")
+    b = _FakeAgentExec("dup", "B")  # same executor id
+    with pytest.raises(ValueError):
+        ConcurrentBuilder().participants([a, b])
+
+
+async def test_concurrent_default_aggregator_emits_single_user_and_assistants() -> None:
+    # Three synthetic agent executors
+    e1 = _FakeAgentExec("agentA", "Alpha")
+    e2 = _FakeAgentExec("agentB", "Beta")
+    e3 = _FakeAgentExec("agentC", "Gamma")
+
+    wf = ConcurrentBuilder().participants([e1, e2, e3]).build()
+
+    completed = False
+    output: list[ChatMessage] | None = None
+    async for ev in wf.run_stream("prompt: hello world"):
+        if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
+            completed = True
+        elif isinstance(ev, WorkflowOutputEvent):
+            output = cast(list[ChatMessage], ev.data)
+        if completed and output is not None:
+            break
+
+    assert completed
+    assert output is not None
+    messages: list[ChatMessage] = output
+
+    # Expect one user message + one assistant message per participant
+    assert len(messages) == 1 + 3
+    assert messages[0].role == Role.USER
+    assert "hello world" in messages[0].text
+
+    assistant_texts = {m.text for m in messages[1:]}
+    assert assistant_texts == {"Alpha", "Beta", "Gamma"}
+    assert all(m.role == Role.ASSISTANT for m in messages[1:])
+
+
+async def test_concurrent_custom_aggregator_callback_is_used() -> None:
+    # Two synthetic agent executors for brevity
+    e1 = _FakeAgentExec("agentA", "One")
+    e2 = _FakeAgentExec("agentB", "Two")
+
+    async def summarize(results: list[AgentExecutorResponse]) -> str:
+        texts: list[str] = []
+        for r in results:
+            msgs: list[ChatMessage] = r.agent_run_response.messages
+            texts.append(msgs[-1].text if msgs else "")
+        return " | ".join(sorted(texts))
+
+    wf = ConcurrentBuilder().participants([e1, e2]).with_aggregator(summarize).build()
+
+    completed = False
+    output: str | None = None
+    async for ev in wf.run_stream("prompt: custom"):
+        if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
+            completed = True
+        elif isinstance(ev, WorkflowOutputEvent):
+            output = cast(str, ev.data)
+        if completed and output is not None:
+            break
+
+    assert completed
+    assert output is not None
+    # Custom aggregator returns a string payload
+    assert isinstance(output, str)
+    assert output == "One | Two"
+
+
+async def test_concurrent_custom_aggregator_sync_callback_is_used() -> None:
+    e1 = _FakeAgentExec("agentA", "One")
+    e2 = _FakeAgentExec("agentB", "Two")
+
+    # Sync callback with ctx parameter (should run via asyncio.to_thread)
+    def summarize_sync(results: list[AgentExecutorResponse], _ctx: WorkflowContext[Any]) -> str:  # type: ignore[unused-argument]
+        texts: list[str] = []
+        for r in results:
+            msgs: list[ChatMessage] = r.agent_run_response.messages
+            texts.append(msgs[-1].text if msgs else "")
+        return " | ".join(sorted(texts))
+
+    wf = ConcurrentBuilder().participants([e1, e2]).with_aggregator(summarize_sync).build()
+
+    completed = False
+    output: str | None = None
+    async for ev in wf.run_stream("prompt: custom sync"):
+        if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
+            completed = True
+        elif isinstance(ev, WorkflowOutputEvent):
+            output = cast(str, ev.data)
+        if completed and output is not None:
+            break
+
+    assert completed
+    assert output is not None
+    assert isinstance(output, str)
+    assert output == "One | Two"
+
+
+def test_concurrent_custom_aggregator_uses_callback_name_for_id() -> None:
+    e1 = _FakeAgentExec("agentA", "One")
+    e2 = _FakeAgentExec("agentB", "Two")
+
+    def summarize(results: list[AgentExecutorResponse]) -> str:  # type: ignore[override]
+        return str(len(results))
+
+    wf = ConcurrentBuilder().participants([e1, e2]).with_aggregator(summarize).build()
+
+    assert "summarize" in wf.executors
+    aggregator = wf.executors["summarize"]
+    assert aggregator.id == "summarize"
+
+
+async def test_concurrent_checkpoint_resume_round_trip() -> None:
+    storage = InMemoryCheckpointStorage()
+
+    participants = (
+        _FakeAgentExec("agentA", "Alpha"),
+        _FakeAgentExec("agentB", "Beta"),
+        _FakeAgentExec("agentC", "Gamma"),
+    )
+
+    wf = ConcurrentBuilder().participants(list(participants)).with_checkpointing(storage).build()
+
+    baseline_output: list[ChatMessage] | None = None
+    async for ev in wf.run_stream("checkpoint concurrent"):
+        if isinstance(ev, WorkflowOutputEvent):
+            baseline_output = ev.data  # type: ignore[assignment]
+        if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
+            break
+
+    assert baseline_output is not None
+
+    checkpoints = await storage.list_checkpoints()
+    assert checkpoints
+    checkpoints.sort(key=lambda cp: cp.timestamp)
+    resume_checkpoint = next(
+        (cp for cp in checkpoints if (cp.metadata or {}).get("checkpoint_type") == "superstep"),
+        checkpoints[-1],
+    )
+
+    resumed_participants = (
+        _FakeAgentExec("agentA", "Alpha"),
+        _FakeAgentExec("agentB", "Beta"),
+        _FakeAgentExec("agentC", "Gamma"),
+    )
+    wf_resume = ConcurrentBuilder().participants(list(resumed_participants)).with_checkpointing(storage).build()
+
+    resumed_output: list[ChatMessage] | None = None
+    async for ev in wf_resume.run_stream_from_checkpoint(resume_checkpoint.checkpoint_id):
+        if isinstance(ev, WorkflowOutputEvent):
+            resumed_output = ev.data  # type: ignore[assignment]
+        if isinstance(ev, WorkflowStatusEvent) and ev.state in (
+            WorkflowRunState.IDLE,
+            WorkflowRunState.IDLE_WITH_PENDING_REQUESTS,
+        ):
+            break
+
+    assert resumed_output is not None
+    assert [m.role for m in resumed_output] == [m.role for m in baseline_output]
+    assert [m.text for m in resumed_output] == [m.text for m in baseline_output]