Python: Context providers abstraction and Mem0 implementation (#631)

* Added context provider abstractions

* Added mem0 implementation

* Example and small fixes

* Added unit tests for agent

* Added unit tests for mem0 provider

* Updated README

* Small doc updates

* Update python/packages/mem0/agent_framework_mem0/_provider.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Small fixes in tests

* Renaming based on PR feedback

* Small fixes

* Added tests for AggregateContextProvider

* Small improvements

* More improvements based on PR feedback

* Small constant update

* Added more examples

* Added README for Mem0 examples

* Small updates to API

* Updated initialization logic

* Updates for context manager

* Updated Context class

* Dependency update

* Revert changes

* Fixed tests

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Chris <66376200+crickman@users.noreply.github.com>

python/packages/main/agent_framework/__init__.py

@@ -12,6 +12,7 @@ from ._agents import *  # noqa: F403
 from ._clients import *  # noqa: F403
 from ._logging import *  # noqa: F403
 from ._mcp import *  # noqa: F403
+from ._memory import *  # noqa: F403
 from ._threads import *  # noqa: F403
 from ._tools import *  # noqa: F403
 from ._types import *  # noqa: F403

python/packages/main/agent_framework/_agents.py

@@ -3,7 +3,6 @@
 import sys
 from collections.abc import AsyncIterable, Callable, MutableMapping, Sequence
 from contextlib import AbstractAsyncContextManager, AsyncExitStack
-from itertools import chain
 from typing import Any, ClassVar, Literal, Protocol, TypeVar, runtime_checkable
 from uuid import uuid4
 
@@ -12,6 +11,7 @@ from pydantic import BaseModel, Field, PrivateAttr
 from ._clients import BaseChatClient, ChatClientProtocol
 from ._logging import get_logger
 from ._mcp import MCPTool
+from ._memory import AggregateContextProvider, Context, ContextProvider
 from ._pydantic import AFBaseModel
 from ._threads import AgentThread, ChatMessageStore, deserialize_thread_state, thread_on_new_messages
 from ._tools import FUNCTION_INVOKING_CHAT_CLIENT_MARKER, ToolProtocol
@@ -137,12 +137,13 @@ class BaseAgent(AFBaseModel):
        name: The name of the agent, can be None.
        description: The description of the agent.
        display_name: The display name of the agent, which is either the name or id.
-
+       context_providers: The collection of multiple context providers to include during agent invocation.
     """
 
     id: str = Field(default_factory=lambda: str(uuid4()))
     name: str | None = None
     description: str | None = None
+    context_providers: AggregateContextProvider | None = None
 
     async def _notify_thread_of_new_messages(
         self, thread: AgentThread, new_messages: ChatMessage | Sequence[ChatMessage]
@@ -214,6 +215,7 @@ class ChatAgent(BaseAgent):
         user: str | None = None,
         additional_properties: dict[str, Any] | None = None,
         chat_message_store_factory: Callable[[], ChatMessageStore] | None = None,
+        context_providers: ContextProvider | list[ContextProvider] | AggregateContextProvider | None = None,
         **kwargs: Any,
     ) -> None:
         """Create a ChatAgent.
@@ -248,6 +250,7 @@ class ChatAgent(BaseAgent):
             additional_properties: additional properties to include in the request.
             chat_message_store_factory: factory function to create an instance of ChatMessageStore. If not provided,
                 the default in-memory store will be used.
+            context_providers: The collection of multiple context providers to include during agent invocation.
             kwargs: any additional keyword arguments.
                 Unused, can be used by subclasses of this Agent.
         """
@@ -258,6 +261,8 @@ class ChatAgent(BaseAgent):
 
         kwargs.update(additional_properties or {})
 
+        aggregate_context_providers = self._prepare_context_providers(context_providers)
+
         # We ignore the MCP Servers here and store them separately,
         # we add their functions to the tools list at runtime
         normalized_tools = [] if tools is None else tools if isinstance(tools, list) else [tools]
@@ -266,6 +271,7 @@ class ChatAgent(BaseAgent):
         args: dict[str, Any] = {
             "chat_client": chat_client,
             "chat_message_store_factory": chat_message_store_factory,
+            "context_providers": aggregate_context_providers,
             "chat_options": ChatOptions(
                 ai_model_id=model,
                 frequency_penalty=frequency_penalty,
@@ -301,12 +307,16 @@ class ChatAgent(BaseAgent):
     async def __aenter__(self) -> "Self":
         """Async context manager entry.
 
-        If either the chat_client or the local_mcp_tools are context managers,
+        If any of the chat_client, local_mcp_tools, or context_providers are context managers,
         they will be entered into the async exit stack to ensure proper cleanup.
 
         This list might be extended in the future.
         """
-        for context_manager in chain([self.chat_client], self._local_mcp_tools):
+        context_managers = [self.chat_client, *self._local_mcp_tools]
+        if self.context_providers:
+            context_managers.append(self.context_providers)
+
+        for context_manager in context_managers:
             if isinstance(context_manager, AbstractAsyncContextManager):
                 await self._async_exit_stack.enter_async_context(context_manager)
         return self
@@ -388,7 +398,10 @@ class ChatAgent(BaseAgent):
                 will only be passed to functions that are called.
         """
         input_messages = self._normalize_messages(messages)
-        thread, thread_messages = await self._prepare_thread_and_messages(thread=thread, input_messages=input_messages)
+        context = await self.context_providers.model_invoking(input_messages) if self.context_providers else None
+        thread, thread_messages = await self._prepare_thread_and_messages(
+            thread=thread, context=context, input_messages=input_messages
+        )
         agent_name = self._get_agent_name()
 
         # Resolve final tool list (runtime provided tools + local MCP server tools)
@@ -441,6 +454,10 @@ class ChatAgent(BaseAgent):
         await self._notify_thread_of_new_messages(thread, input_messages)
         await self._notify_thread_of_new_messages(thread, response.messages)
 
+        if self.context_providers:
+            await self.context_providers.thread_created(response.conversation_id)
+            await self.context_providers.messages_adding(thread.service_thread_id, input_messages + response.messages)
+
         return AgentRunResponse(
             messages=response.messages,
             response_id=response.response_id,
@@ -509,7 +526,10 @@ class ChatAgent(BaseAgent):
 
         """
         input_messages = self._normalize_messages(messages)
-        thread, thread_messages = await self._prepare_thread_and_messages(thread=thread, input_messages=input_messages)
+        context = await self.context_providers.model_invoking(input_messages) if self.context_providers else None
+        thread, thread_messages = await self._prepare_thread_and_messages(
+            thread=thread, context=context, input_messages=input_messages
+        )
         agent_name = self._get_agent_name()
         response_updates: list[ChatResponseUpdate] = []
 
@@ -575,6 +595,10 @@ class ChatAgent(BaseAgent):
         await self._notify_thread_of_new_messages(thread, input_messages)
         await self._notify_thread_of_new_messages(thread, response.messages)
 
+        if self.context_providers:
+            await self.context_providers.thread_created(response.conversation_id)
+            await self.context_providers.messages_adding(thread.service_thread_id, input_messages + response.messages)
+
     def get_new_thread(self) -> AgentThread:
         message_store: ChatMessageStore | None = None
 
@@ -617,12 +641,14 @@ class ChatAgent(BaseAgent):
         self,
         *,
         thread: AgentThread | None,
+        context: Context | None,
         input_messages: list[ChatMessage] | None = None,
     ) -> tuple[AgentThread, list[ChatMessage]]:
         """Prepare the messages for agent execution.
 
         Args:
             thread: The conversation thread.
+            context: Context to include in messages.
             input_messages: Messages to process.
 
         Returns:
@@ -636,6 +662,8 @@ class ChatAgent(BaseAgent):
         messages: list[ChatMessage] = []
         if self.instructions:
             messages.append(ChatMessage(role=Role.SYSTEM, text=self.instructions))
+        if context and context.contents:
+            messages.append(ChatMessage(role=Role.SYSTEM, contents=context.contents))
         if thread.message_store:
             messages.extend(await thread.message_store.list_messages() or [])
         messages.extend(input_messages or [])
@@ -658,3 +686,18 @@ class ChatAgent(BaseAgent):
 
     def _get_agent_name(self) -> str:
         return self.name or "UnnamedAgent"
+
+    def _prepare_context_providers(
+        self,
+        context_providers: ContextProvider | list[ContextProvider] | AggregateContextProvider | None = None,
+    ) -> AggregateContextProvider | None:
+        if not context_providers:
+            return None
+
+        if isinstance(context_providers, AggregateContextProvider):
+            return context_providers
+
+        if isinstance(context_providers, ContextProvider):
+            return AggregateContextProvider([context_providers])
+
+        return AggregateContextProvider(context_providers)

python/packages/main/agent_framework/_clients.py

@@ -9,6 +9,7 @@ from pydantic import BaseModel, Field
 
 from ._logging import get_logger
 from ._mcp import MCPTool
+from ._memory import AggregateContextProvider, ContextProvider
 from ._pydantic import AFBaseModel
 from ._threads import ChatMessageStore
 from ._tools import ToolProtocol
@@ -463,6 +464,7 @@ class BaseChatClient(AFBaseModel, ABC):
         | list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
         | None = None,
         chat_message_store_factory: Callable[[], ChatMessageStore] | None = None,
+        context_providers: ContextProvider | list[ContextProvider] | AggregateContextProvider | None = None,
         **kwargs: Any,
     ) -> "ChatAgent":
         """Create an agent with the given name and instructions.
@@ -473,6 +475,7 @@ class BaseChatClient(AFBaseModel, ABC):
             tools: Optional list of tools to associate with the agent.
             chat_message_store_factory: Factory function to create an instance of ChatMessageStore. If not provided,
                 the default in-memory store will be used.
+            context_providers: Context providers to include during agent invocation.
             **kwargs: Additional keyword arguments to pass to the agent.
                 See ChatAgent for all the available options.
 
@@ -487,6 +490,7 @@ class BaseChatClient(AFBaseModel, ABC):
             instructions=instructions,
             tools=tools,
             chat_message_store_factory=chat_message_store_factory,
+            context_providers=context_providers,
             **kwargs,
         )
 

python/packages/main/agent_framework/_memory.py

@@ -0,0 +1,188 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import sys
+from abc import ABC, abstractmethod
+from collections.abc import MutableSequence, Sequence
+from contextlib import AsyncExitStack
+from types import TracebackType
+
+from ._pydantic import AFBaseModel
+from ._types import ChatMessage, Contents
+
+if sys.version_info >= (3, 11):
+    from typing import Self  # pragma: no cover
+else:
+    from typing_extensions import Self  # pragma: no cover
+
+# region Context
+
+
+class Context(AFBaseModel):
+    """A class containing any context that should be provided to the AI model as supplied by an ContextProvider.
+
+    Each ContextProvider has the ability to provide its own context for each invocation.
+    The Context class contains the additional context supplied by the ContextProvider.
+    This context will be combined with context supplied by other providers before being passed to the AI model.
+    This context is per invocation, and will not be stored as part of the chat history.
+    """
+
+    contents: list[Contents] | None = None
+    """
+    Any content to pass to the AI model in addition to any other prompts
+    that it may already have (in the case of an agent), or chat history that may already exist.
+    """
+
+
+# region ContextProvider
+
+
+class ContextProvider(AFBaseModel, ABC):
+    """Base class for all context providers.
+
+    A context provider is a component that can be used to enhance the AI's context management.
+    It can listen to changes in the conversation and provide additional context to the AI model
+    just before invocation.
+    """
+
+    async def thread_created(self, thread_id: str | None) -> None:
+        """Called just after a new thread is created.
+
+        Implementers can use this method to do any operations required at the creation of a new thread.
+        For example, checking long term storage for any data that is relevant
+        to the current session based on the input text.
+
+        Args:
+            thread_id: The ID of the new thread.
+        """
+        pass
+
+    async def messages_adding(self, thread_id: str | None, new_messages: ChatMessage | Sequence[ChatMessage]) -> None:
+        """Called just before messages are added to the chat by any participant.
+
+        Inheritors can use this method to update their context based on new messages.
+
+        Args:
+            thread_id: The ID of the thread for the new message.
+            new_messages: New messages to add.
+        """
+        pass
+
+    @abstractmethod
+    async def model_invoking(self, messages: ChatMessage | MutableSequence[ChatMessage]) -> Context:
+        """Called just before the Model/Agent/etc. is invoked.
+
+        Implementers can load any additional context required at this time,
+        and they should return any context that should be passed to the agent.
+
+        Args:
+            messages: The most recent messages that the agent is being invoked with.
+        """
+        pass
+
+    async def __aenter__(self) -> "Self":
+        """Async context manager entry.
+
+        Override this method to perform any setup operations when the context provider is entered.
+
+        Returns:
+            Self for chaining.
+        """
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Async context manager exit.
+
+        Override this method to perform any cleanup operations when the context provider is exited.
+
+        Args:
+            exc_type: Exception type if an exception occurred, None otherwise.
+            exc_val: Exception value if an exception occurred, None otherwise.
+            exc_tb: Exception traceback if an exception occurred, None otherwise.
+        """
+        pass
+
+
+# region AggregateContextProvider
+
+
+class AggregateContextProvider(ContextProvider):
+    """A ContextProvider that contains multiple context providers.
+
+    It delegates events to multiple context providers and aggregates responses from those events before returning.
+    """
+
+    providers: list[ContextProvider]
+    """List of registered context providers."""
+
+    def __init__(self, context_providers: Sequence[ContextProvider] | None = None) -> None:
+        """Initialize AggregateContextProvider with context providers.
+
+        Args:
+            context_providers: Context providers to add.
+        """
+        super().__init__(providers=list(context_providers or []))  # type: ignore
+        self._exit_stack: AsyncExitStack | None = None
+
+    def add(self, context_provider: ContextProvider) -> None:
+        """Adds new context provider.
+
+        Args:
+            context_provider: Context provider to add.
+        """
+        self.providers.append(context_provider)
+
+    async def thread_created(self, thread_id: str | None = None) -> None:
+        await asyncio.gather(*[x.thread_created(thread_id) for x in self.providers])
+
+    async def messages_adding(self, thread_id: str | None, new_messages: ChatMessage | Sequence[ChatMessage]) -> None:
+        await asyncio.gather(*[x.messages_adding(thread_id, new_messages) for x in self.providers])
+
+    async def model_invoking(self, messages: ChatMessage | MutableSequence[ChatMessage]) -> Context:
+        sub_contexts = await asyncio.gather(*[x.model_invoking(messages) for x in self.providers])
+        combined_context = Context()
+        # Flatten the list of lists and filter out None values
+        all_contents = []
+        for ctx in sub_contexts:
+            if ctx.contents:
+                all_contents.extend(ctx.contents)
+
+        combined_context.contents = all_contents if all_contents else None
+        return combined_context
+
+    async def __aenter__(self) -> "Self":
+        """Enter async context manager and set up all providers.
+
+        Returns:
+            Self for chaining.
+        """
+        self._exit_stack = AsyncExitStack()
+        await self._exit_stack.__aenter__()
+
+        # Enter all context providers
+        for provider in self.providers:
+            await self._exit_stack.enter_async_context(provider)
+
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Exit async context manager and clean up all providers.
+
+        Args:
+            exc_type: Exception type if an exception occurred, None otherwise.
+            exc_val: Exception value if an exception occurred, None otherwise.
+            exc_tb: Exception traceback if an exception occurred, None otherwise.
+        """
+        if self._exit_stack is not None:
+            await self._exit_stack.__aexit__(exc_type, exc_val, exc_tb)
+            self._exit_stack = None

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

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

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

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

python/packages/main/agent_framework/openai/_responses_client.py

@@ -455,7 +455,7 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
                         )
                         response_tools.append(
                             WebSearchToolParam(
-                                type="web_search",
+                                type="web_search_preview",
                                 user_location=WebSearchUserLocation(
                                     type="approximate",
                                     city=location.get("city", None),

python/packages/main/pyproject.toml

@@ -23,7 +23,7 @@ classifiers = [
   "Typing :: Typed",
 ]
 dependencies = [
-    "openai>=1.103.0",
+    "openai>=1.99.0",
     "pydantic>=2.11.7",
     "pydantic-settings>=2.10.1",
     "typing-extensions>=4.14.0",

python/packages/main/tests/main/conftest.py

@@ -151,11 +151,19 @@ class MockBaseChatClient(BaseChatClient):
         logger.debug(f"Running base chat client inner, with: {messages=}, {chat_options=}, {kwargs=}")
         if not self.run_responses:
             return ChatResponse(messages=ChatMessage(role="assistant", text=f"test response - {messages[0].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...")
+                messages=ChatMessage(
+                    role="assistant",
+                    text="I broke out of the function invocation loop...",
+                ),
+                conversation_id=response.conversation_id,
             )
-        return self.run_responses.pop(0)
+
+        return response
 
     @override
     async def _inner_get_streaming_response(

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

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from collections.abc import AsyncIterable
+from collections.abc import AsyncIterable, MutableSequence, Sequence
 from uuid import uuid4
 
 from pytest import raises
@@ -15,10 +15,12 @@ from agent_framework import (
     ChatMessage,
     ChatMessageList,
     ChatResponse,
+    Contents,
     HostedCodeInterpreterTool,
     Role,
     TextContent,
 )
+from agent_framework._memory import AggregateContextProvider, Context, ContextProvider
 from agent_framework.exceptions import AgentExecutionException
 
 
@@ -100,6 +102,7 @@ async def test_chat_client_agent_prepare_thread_and_messages(chat_client: ChatCl
 
     _, result_messages = await agent._prepare_thread_and_messages(  # type: ignore[reportPrivateUsage]
         thread=thread,
+        context=Context(),
         input_messages=[ChatMessage(role=Role.USER, text="Test")],
     )
 
@@ -184,3 +187,210 @@ async def test_chat_client_agent_author_name_is_used_from_response(chat_client_b
     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):
+    context_contents: list[Contents] | None = None
+    thread_created_called: bool = False
+    messages_adding_called: bool = False
+    model_invoking_called: bool = False
+    thread_created_thread_id: str | None = None
+    messages_adding_thread_id: str | None = None
+    new_messages: list[ChatMessage] = []
+
+    def __init__(self, contents: list[Contents] | None = None) -> None:
+        super().__init__()
+        self.context_contents = contents
+        self.thread_created_called = False
+        self.messages_adding_called = False
+        self.model_invoking_called = False
+        self.thread_created_thread_id = None
+        self.messages_adding_thread_id = None
+        self.new_messages = []
+
+    async def thread_created(self, thread_id: str | None) -> None:
+        self.thread_created_called = True
+        self.thread_created_thread_id = thread_id
+
+    async def messages_adding(self, thread_id: str | None, new_messages: ChatMessage | Sequence[ChatMessage]) -> None:
+        self.messages_adding_called = True
+        self.messages_adding_thread_id = thread_id
+        if isinstance(new_messages, ChatMessage):
+            self.new_messages.append(new_messages)
+        else:
+            self.new_messages.extend(new_messages)
+
+    async def model_invoking(self, messages: ChatMessage | MutableSequence[ChatMessage]) -> Context:
+        self.model_invoking_called = True
+        return Context(contents=self.context_contents)
+
+
+async def test_chat_agent_context_providers_model_invoking(chat_client: ChatClientProtocol) -> None:
+    """Test that context providers' model_invoking is called during agent run."""
+    mock_provider = MockContextProvider(contents=[TextContent("Test context instructions")])
+    agent = ChatAgent(chat_client=chat_client, context_providers=mock_provider)
+
+    await agent.run("Hello")
+
+    assert mock_provider.model_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' messages_adding 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.messages_adding_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(contents=[TextContent("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
+    context = Context(contents=[TextContent("Context-specific instructions")])
+    _, messages = await agent._prepare_thread_and_messages(  # type: ignore[reportPrivateUsage]
+        thread=None, context=context, input_messages=[ChatMessage(role=Role.USER, text="Hello")]
+    )
+
+    # Should have agent instructions, context instructions, and user message
+    assert len(messages) == 3
+    assert messages[0].role == Role.SYSTEM
+    assert messages[0].text == "Agent instructions"
+    assert messages[1].role == Role.SYSTEM
+    assert messages[1].text == "Context-specific instructions"
+    assert messages[2].role == Role.USER
+    assert messages[2].text == "Hello"
+
+
+async def test_chat_agent_context_instructions_without_agent_instructions(chat_client: ChatClientProtocol) -> None:
+    """Test that AI context instructions work when agent has no instructions."""
+    agent = ChatAgent(chat_client=chat_client)  # No instructions
+    context = Context(contents=[TextContent("Context-only instructions")])
+
+    _, messages = await agent._prepare_thread_and_messages(  # type: ignore[reportPrivateUsage]
+        thread=None, context=context, input_messages=[ChatMessage(role=Role.USER, text="Hello")]
+    )
+
+    # Should have context instructions and user message only
+    assert len(messages) == 2
+    assert messages[0].role == Role.SYSTEM
+    assert messages[0].text == "Context-only instructions"
+    assert messages[1].role == Role.USER
+    assert messages[1].text == "Hello"
+
+
+async def test_chat_agent_no_context_instructions(chat_client: ChatClientProtocol) -> None:
+    """Test behavior when AI context has no instructions."""
+    agent = ChatAgent(chat_client=chat_client, instructions="Agent instructions")
+    context = Context()  # No instructions
+
+    _, messages = await agent._prepare_thread_and_messages(  # type: ignore[reportPrivateUsage]
+        thread=None, context=context, input_messages=[ChatMessage(role=Role.USER, text="Hello")]
+    )
+
+    # Should have agent instructions and user message only
+    assert len(messages) == 2
+    assert messages[0].role == Role.SYSTEM
+    assert messages[0].text == "Agent instructions"
+    assert messages[1].role == Role.USER
+    assert messages[1].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(contents=[TextContent("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.model_invoking_called
+    assert mock_provider.thread_created_called
+    assert mock_provider.messages_adding_called
+
+
+async def test_chat_agent_multiple_context_providers(chat_client: ChatClientProtocol) -> None:
+    """Test that multiple context providers work together."""
+    provider1 = MockContextProvider(contents=[TextContent("First provider instructions")])
+    provider2 = MockContextProvider(contents=[TextContent("Second provider instructions")])
+
+    agent = ChatAgent(chat_client=chat_client, context_providers=[provider1, provider2])
+
+    await agent.run("Hello")
+
+    # Both providers should be called
+    assert provider1.model_invoking_called
+    assert provider1.thread_created_called
+    assert provider1.messages_adding_called
+
+    assert provider2.model_invoking_called
+    assert provider2.thread_created_called
+    assert provider2.messages_adding_called
+
+
+async def test_chat_agent_aggregate_context_provider_combines_instructions() -> None:
+    """Test that AggregateContextProvider combines instructions from multiple providers."""
+    provider1 = MockContextProvider(contents=[TextContent("First instruction")])
+    provider2 = MockContextProvider(contents=[TextContent("Second instruction")])
+
+    aggregate = AggregateContextProvider()
+    aggregate.providers.append(provider1)
+    aggregate.providers.append(provider2)
+
+    # Test model_invoking combines instructions
+    result = await aggregate.model_invoking([ChatMessage(role=Role.USER, text="Test")])
+
+    assert result.contents
+    assert isinstance(result.contents[0], TextContent)
+    assert isinstance(result.contents[1], TextContent)
+    assert result.contents[0].text == "First instruction"
+    assert result.contents[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 = AgentThread(service_thread_id="existing-thread-id")
+    await agent.run("Hello", thread=thread)
+
+    # messages_adding should be called with the service thread ID from response
+    assert mock_provider.messages_adding_called
+    assert mock_provider.messages_adding_thread_id == "service-thread-123"  # Updated thread ID from response

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

@@ -0,0 +1,302 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from collections.abc import MutableSequence, Sequence
+from unittest.mock import AsyncMock, Mock
+
+from agent_framework import ChatMessage, Contents, Role, TextContent
+from agent_framework._memory import AggregateContextProvider, Context, ContextProvider
+
+
+class MockContextProvider(ContextProvider):
+    """Mock ContextProvider for testing."""
+
+    context_contents: list[Contents] | None = None
+    thread_created_called: bool = False
+    messages_adding_called: bool = False
+    model_invoking_called: bool = False
+    thread_created_thread_id: str | None = None
+    messages_adding_thread_id: str | None = None
+    messages_adding_new_messages: ChatMessage | Sequence[ChatMessage] | None = None
+    model_invoking_messages: ChatMessage | MutableSequence[ChatMessage] | None = None
+
+    def __init__(self, context_contents: list[Contents] | None = None) -> None:
+        super().__init__()
+        self.context_contents = context_contents
+        self.thread_created_called = False
+        self.messages_adding_called = False
+        self.model_invoking_called = False
+        self.thread_created_thread_id = None
+        self.messages_adding_thread_id = None
+        self.messages_adding_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 messages_adding(self, thread_id: str | None, new_messages: ChatMessage | Sequence[ChatMessage]) -> None:
+        """Track messages_adding calls."""
+        self.messages_adding_called = True
+        self.messages_adding_thread_id = thread_id
+        self.messages_adding_new_messages = new_messages
+
+    async def model_invoking(self, messages: ChatMessage | MutableSequence[ChatMessage]) -> Context:
+        """Track model_invoking calls and return context."""
+        self.model_invoking_called = True
+        self.model_invoking_messages = messages
+        context = Context()
+        context.contents = self.context_contents
+        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([TextContent("instructions1")])
+        provider2 = MockContextProvider([TextContent("instructions1")])
+        providers = [provider1, provider2]
+
+        aggregate = AggregateContextProvider(providers)
+        assert len(aggregate.providers) == 2
+        assert aggregate.providers[0] is provider1
+        assert aggregate.providers[1] is provider2
+
+    def test_add_provider(self) -> None:
+        """Test adding a provider."""
+        aggregate = AggregateContextProvider()
+        provider = MockContextProvider([TextContent("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([TextContent("instructions1")])
+        provider2 = MockContextProvider([TextContent("instructions2")])
+
+        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([TextContent("instructions1")])
+        provider2 = MockContextProvider([TextContent("instructions2")])
+        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([TextContent("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 messages_adding with no providers."""
+        aggregate = AggregateContextProvider()
+        message = ChatMessage(text="Hello", role=Role.USER)
+
+        # Should not raise an exception
+        await aggregate.messages_adding("thread-123", message)
+
+    async def test_messages_adding_with_single_message(self) -> None:
+        """Test messages_adding with a single message."""
+        provider1 = MockContextProvider([TextContent("instructions1")])
+        provider2 = MockContextProvider([TextContent("instructions2")])
+        aggregate = AggregateContextProvider([provider1, provider2])
+
+        thread_id = "thread-123"
+        message = ChatMessage(text="Hello", role=Role.USER)
+        await aggregate.messages_adding(thread_id, message)
+
+        assert provider1.messages_adding_called
+        assert provider1.messages_adding_thread_id == thread_id
+        assert provider1.messages_adding_new_messages == message
+        assert provider2.messages_adding_called
+        assert provider2.messages_adding_thread_id == thread_id
+        assert provider2.messages_adding_new_messages == message
+
+    async def test_messages_adding_with_message_sequence(self) -> None:
+        """Test messages_adding with a sequence of messages."""
+        provider = MockContextProvider([TextContent("instructions")])
+        aggregate = AggregateContextProvider([provider])
+
+        thread_id = "thread-123"
+        messages = [
+            ChatMessage(text="Hello", role=Role.USER),
+            ChatMessage(text="Hi there", role=Role.ASSISTANT),
+        ]
+        await aggregate.messages_adding(thread_id, messages)
+
+        assert provider.messages_adding_called
+        assert provider.messages_adding_thread_id == thread_id
+        assert provider.messages_adding_new_messages == messages
+
+    async def test_model_invoking_with_no_providers(self) -> None:
+        """Test model_invoking with no providers."""
+        aggregate = AggregateContextProvider()
+        message = ChatMessage(text="Hello", role=Role.USER)
+
+        context = await aggregate.model_invoking(message)
+
+        assert isinstance(context, Context)
+        assert not context.contents
+
+    async def test_model_invoking_with_single_provider(self) -> None:
+        """Test model_invoking with a single provider."""
+        provider = MockContextProvider([TextContent("Test instructions")])
+        aggregate = AggregateContextProvider([provider])
+
+        message = ChatMessage(text="Hello", role=Role.USER)
+        context = await aggregate.model_invoking(message)
+
+        assert provider.model_invoking_called
+        assert provider.model_invoking_messages == message
+        assert isinstance(context, Context)
+
+        assert context.contents
+        assert isinstance(context.contents[0], TextContent)
+        assert context.contents[0].text == "Test instructions"
+
+    async def test_model_invoking_with_multiple_providers(self) -> None:
+        """Test model_invoking combines contexts from multiple providers."""
+        provider1 = MockContextProvider([TextContent("Instructions 1")])
+        provider2 = MockContextProvider([TextContent("Instructions 2")])
+        provider3 = MockContextProvider([TextContent("Instructions 3")])
+        aggregate = AggregateContextProvider([provider1, provider2, provider3])
+
+        messages = [ChatMessage(text="Hello", role=Role.USER)]
+        context = await aggregate.model_invoking(messages)
+
+        assert provider1.model_invoking_called
+        assert provider1.model_invoking_messages == messages
+        assert provider2.model_invoking_called
+        assert provider2.model_invoking_messages == messages
+        assert provider3.model_invoking_called
+        assert provider3.model_invoking_messages == messages
+
+        assert isinstance(context, Context)
+
+        assert context.contents
+        assert isinstance(context.contents[0], TextContent)
+        assert isinstance(context.contents[1], TextContent)
+        assert isinstance(context.contents[2], TextContent)
+        assert context.contents[0].text == "Instructions 1"
+        assert context.contents[1].text == "Instructions 2"
+        assert context.contents[2].text == "Instructions 3"
+
+    async def test_model_invoking_with_none_instructions(self) -> None:
+        """Test model_invoking filters out None instructions."""
+        provider1 = MockContextProvider([TextContent("Instructions 1")])
+        provider2 = MockContextProvider(None)  # None instructions
+        provider3 = MockContextProvider([TextContent("Instructions 3")])
+        aggregate = AggregateContextProvider([provider1, provider2, provider3])
+
+        message = ChatMessage(text="Hello", role=Role.USER)
+        context = await aggregate.model_invoking(message)
+
+        assert isinstance(context, Context)
+        assert context.contents
+        assert isinstance(context.contents[0], TextContent)
+        assert isinstance(context.contents[1], TextContent)
+        assert context.contents[0].text == "Instructions 1"
+        assert context.contents[1].text == "Instructions 3"
+
+    async def test_model_invoking_with_all_none_instructions(self) -> None:
+        """Test model_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.model_invoking(message)
+
+        assert isinstance(context, Context)
+        assert not context.contents
+
+    async def test_model_invoking_with_mutable_sequence(self) -> None:
+        """Test model_invoking with MutableSequence of messages."""
+        provider = MockContextProvider([TextContent("Test instructions")])
+        aggregate = AggregateContextProvider([provider])
+
+        messages = [ChatMessage(text="Hello", role=Role.USER)]
+        context = await aggregate.model_invoking(messages)
+
+        assert provider.model_invoking_called
+        assert provider.model_invoking_messages == messages
+        assert isinstance(context, Context)
+        assert context.contents
+        assert isinstance(context.contents[0], TextContent)
+        assert context.contents[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.messages_adding = AsyncMock()
+        provider1.model_invoking = AsyncMock(return_value=Context(contents=[TextContent("Test 1")]))
+
+        provider2 = Mock(spec=ContextProvider)
+        provider2.thread_created = AsyncMock()
+        provider2.messages_adding = AsyncMock()
+        provider2.model_invoking = AsyncMock(return_value=Context(contents=[TextContent("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 messages_adding
+        message = ChatMessage(text="Hello", role=Role.USER)
+        await aggregate.messages_adding("thread-123", message)
+        provider1.messages_adding.assert_called_once_with("thread-123", message)
+        provider2.messages_adding.assert_called_once_with("thread-123", message)
+
+        # Test model_invoking
+        context = await aggregate.model_invoking(message)
+        provider1.model_invoking.assert_called_once_with(message)
+        provider2.model_invoking.assert_called_once_with(message)
+        assert context.contents
+        assert isinstance(context.contents[0], TextContent)
+        assert isinstance(context.contents[1], TextContent)
+        assert context.contents[0].text == "Test 1"
+        assert context.contents[1].text == "Test 2"