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/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)