Python: Added thread to AgentRunContext (#1732)

* Added thread to agent run context

* Added sample

* Update python/samples/getting_started/middleware/thread_behavior_middleware.py

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

* Small fix

---------

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

python/packages/core/agent_framework/_middleware.py

@@ -16,6 +16,7 @@ if TYPE_CHECKING:
 
     from ._agents import AgentProtocol
     from ._clients import ChatClientProtocol
+    from ._threads import AgentThread
     from ._tools import AIFunction
     from ._types import ChatOptions, ChatResponse, ChatResponseUpdate
 
@@ -61,6 +62,7 @@ class AgentRunContext(SerializationMixin):
     Attributes:
         agent: The agent being invoked.
         messages: The messages being sent to the agent.
+        thread: The agent thread for this invocation, if any.
         is_streaming: Whether this is a streaming invocation.
         metadata: Metadata dictionary for sharing data between agent middleware.
         result: Agent execution result. Can be observed after calling ``next()``
@@ -81,6 +83,7 @@ class AgentRunContext(SerializationMixin):
                 async def process(self, context: AgentRunContext, next):
                     print(f"Agent: {context.agent.name}")
                     print(f"Messages: {len(context.messages)}")
+                    print(f"Thread: {context.thread}")
                     print(f"Streaming: {context.is_streaming}")
 
                     # Store metadata
@@ -93,12 +96,13 @@ class AgentRunContext(SerializationMixin):
                     print(f"Result: {context.result}")
     """
 
-    INJECTABLE: ClassVar[set[str]] = {"agent", "result"}
+    INJECTABLE: ClassVar[set[str]] = {"agent", "thread", "result"}
 
     def __init__(
         self,
         agent: "AgentProtocol",
         messages: list[ChatMessage],
+        thread: "AgentThread | None" = None,
         is_streaming: bool = False,
         metadata: dict[str, Any] | None = None,
         result: AgentRunResponse | AsyncIterable[AgentRunResponseUpdate] | None = None,
@@ -110,6 +114,7 @@ class AgentRunContext(SerializationMixin):
         Args:
             agent: The agent being invoked.
             messages: The messages being sent to the agent.
+            thread: The agent thread for this invocation, if any.
             is_streaming: Whether this is a streaming invocation.
             metadata: Metadata dictionary for sharing data between agent middleware.
             result: Agent execution result.
@@ -118,6 +123,7 @@ class AgentRunContext(SerializationMixin):
         """
         self.agent = agent
         self.messages = messages
+        self.thread = thread
         self.is_streaming = is_streaming
         self.metadata = metadata if metadata is not None else {}
         self.result = result
@@ -1222,6 +1228,7 @@ def use_agent_middleware(agent_class: type[TAgent]) -> type[TAgent]:
             context = AgentRunContext(
                 agent=self,  # type: ignore[arg-type]
                 messages=normalized_messages,
+                thread=thread,
                 is_streaming=False,
                 kwargs=kwargs,
             )
@@ -1269,6 +1276,7 @@ def use_agent_middleware(agent_class: type[TAgent]) -> type[TAgent]:
             context = AgentRunContext(
                 agent=self,  # type: ignore[arg-type]
                 messages=normalized_messages,
+                thread=thread,
                 is_streaming=True,
                 kwargs=kwargs,
             )

python/packages/core/tests/core/test_middleware.py

@@ -56,6 +56,20 @@ class TestAgentRunContext:
         assert context.is_streaming is True
         assert context.metadata == metadata
 
+    def test_init_with_thread(self, mock_agent: AgentProtocol) -> None:
+        """Test AgentRunContext initialization with thread parameter."""
+        from agent_framework import AgentThread
+
+        messages = [ChatMessage(role=Role.USER, text="test")]
+        thread = AgentThread()
+        context = AgentRunContext(agent=mock_agent, messages=messages, thread=thread)
+
+        assert context.agent is mock_agent
+        assert context.messages == messages
+        assert context.thread is thread
+        assert context.is_streaming is False
+        assert context.metadata == {}
+
 
 class TestFunctionInvocationContext:
     """Test cases for FunctionInvocationContext."""
@@ -337,6 +351,61 @@ class TestAgentMiddlewarePipeline:
         assert context.terminate
         assert execution_order == ["handler_start", "handler_end"]
 
+    async def test_execute_with_thread_in_context(self, mock_agent: AgentProtocol) -> None:
+        """Test pipeline execution properly passes thread to middleware."""
+        from agent_framework import AgentThread
+
+        captured_thread = None
+
+        class ThreadCapturingMiddleware(AgentMiddleware):
+            async def process(
+                self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
+            ) -> None:
+                nonlocal captured_thread
+                captured_thread = context.thread
+                await next(context)
+
+        middleware = ThreadCapturingMiddleware()
+        pipeline = AgentMiddlewarePipeline([middleware])
+        messages = [ChatMessage(role=Role.USER, text="test")]
+        thread = AgentThread()
+        context = AgentRunContext(agent=mock_agent, messages=messages, thread=thread)
+
+        expected_response = AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
+
+        async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
+            return expected_response
+
+        result = await pipeline.execute(mock_agent, messages, context, final_handler)
+        assert result == expected_response
+        assert captured_thread is thread
+
+    async def test_execute_with_no_thread_in_context(self, mock_agent: AgentProtocol) -> None:
+        """Test pipeline execution when no thread is provided."""
+        captured_thread = "not_none"  # Use string to distinguish from None
+
+        class ThreadCapturingMiddleware(AgentMiddleware):
+            async def process(
+                self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
+            ) -> None:
+                nonlocal captured_thread
+                captured_thread = context.thread
+                await next(context)
+
+        middleware = ThreadCapturingMiddleware()
+        pipeline = AgentMiddlewarePipeline([middleware])
+        messages = [ChatMessage(role=Role.USER, text="test")]
+        context = AgentRunContext(agent=mock_agent, messages=messages, thread=None)
+
+        expected_response = AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
+
+        async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
+            return expected_response
+
+        result = await pipeline.execute(mock_agent, messages, context, final_handler)
+        assert result == expected_response
+        assert captured_thread is None
+
 
 class TestFunctionMiddlewarePipeline:
     """Test cases for FunctionMiddlewarePipeline."""

python/packages/core/tests/core/test_middleware_with_agent.py

@@ -1442,6 +1442,119 @@ class TestMiddlewareDecoratorLogic:
         assert test_function_middleware._middleware_type == MiddlewareType.FUNCTION  # type: ignore[attr-defined]
 
 
+class TestChatAgentThreadBehavior:
+    """Test cases for thread behavior in AgentRunContext across multiple runs."""
+
+    async def test_agent_run_context_thread_behavior_across_multiple_runs(self, chat_client: "MockChatClient") -> None:
+        """Test that AgentRunContext.thread property behaves correctly across multiple agent runs."""
+        thread_states: list[dict[str, Any]] = []
+
+        class ThreadTrackingMiddleware(AgentMiddleware):
+            async def process(
+                self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
+            ) -> None:
+                # Capture state before next() call
+                thread_messages = []
+                if context.thread and context.thread.message_store:
+                    thread_messages = await context.thread.message_store.list_messages()
+
+                before_state = {
+                    "before_next": True,
+                    "messages_count": len(context.messages),
+                    "thread_count": len(thread_messages),
+                    "messages_text": [msg.text for msg in context.messages if msg.text],
+                    "thread_messages_text": [msg.text for msg in thread_messages if msg.text],
+                }
+                thread_states.append(before_state)
+
+                await next(context)
+
+                # Capture state after next() call
+                thread_messages_after = []
+                if context.thread and context.thread.message_store:
+                    thread_messages_after = await context.thread.message_store.list_messages()
+
+                after_state = {
+                    "before_next": False,
+                    "messages_count": len(context.messages),
+                    "thread_count": len(thread_messages_after),
+                    "messages_text": [msg.text for msg in context.messages if msg.text],
+                    "thread_messages_text": [msg.text for msg in thread_messages_after if msg.text],
+                }
+                thread_states.append(after_state)
+
+        # Import the ChatMessageStore to configure the agent with a message store factory
+        from agent_framework import ChatMessageStore
+
+        # Create ChatAgent with thread tracking middleware and a message store factory
+        middleware = ThreadTrackingMiddleware()
+        agent = ChatAgent(chat_client=chat_client, middleware=[middleware], chat_message_store_factory=ChatMessageStore)
+
+        # Create a thread that will persist messages between runs
+        thread = agent.get_new_thread()
+
+        # First run
+        first_messages = [ChatMessage(role=Role.USER, text="first message")]
+        first_response = await agent.run(first_messages, thread=thread)
+
+        # Verify first response
+        assert first_response is not None
+        assert len(first_response.messages) > 0
+
+        # Second run - use the same thread
+        second_messages = [ChatMessage(role=Role.USER, text="second message")]
+        second_response = await agent.run(second_messages, thread=thread)
+
+        # Verify second response
+        assert second_response is not None
+        assert len(second_response.messages) > 0
+
+        # Verify we captured states for both runs (before and after next() for each)
+        assert len(thread_states) == 4
+
+        # First run - before next()
+        first_before = thread_states[0]
+        assert first_before["before_next"] is True
+        assert first_before["messages_count"] == 1
+        assert first_before["thread_count"] == 0  # Thread is empty before first run
+        assert first_before["messages_text"] == ["first message"]
+        assert first_before["thread_messages_text"] == []
+
+        # First run - after next()
+        first_after = thread_states[1]
+        assert first_after["before_next"] is False
+        assert first_after["messages_count"] == 1  # Input messages unchanged
+        assert first_after["thread_count"] == 2  # Input + response
+        assert first_after["messages_text"] == ["first message"]
+        # Thread should contain input + response
+        assert "first message" in first_after["thread_messages_text"]
+        assert "test response" in " ".join(first_after["thread_messages_text"])
+
+        # Second run - before next()
+        second_before = thread_states[2]
+        assert second_before["before_next"] is True
+        assert second_before["messages_count"] == 1  # Only current run input
+        assert second_before["thread_count"] == 2  # Previous run history (input + response)
+        assert second_before["messages_text"] == ["second message"]
+        # Thread should contain previous run history but not current input yet
+        assert "first message" in second_before["thread_messages_text"]
+        assert "test response" in " ".join(second_before["thread_messages_text"])
+        assert "second message" not in second_before["thread_messages_text"]
+
+        # Second run - after next()
+        second_after = thread_states[3]
+        assert second_after["before_next"] is False
+        assert second_after["messages_count"] == 1  # Input messages unchanged
+        assert second_after["thread_count"] == 4  # Previous history + current input + current response
+        assert second_after["messages_text"] == ["second message"]
+        # Thread should contain: first input + first response + second input + second response
+        assert "first message" in second_after["thread_messages_text"]
+        assert "second message" in second_after["thread_messages_text"]
+        # Should have two "test response" entries (one for each run)
+        response_count = sum(1 for text in second_after["thread_messages_text"] if "test response" in text)
+        assert response_count == 2
+
+
 class TestChatAgentChatMiddleware:
     """Test cases for chat middleware integration with ChatAgent."""
 

python/packages/lab/gaia/agent_framework_lab_gaia/gaia.py

@@ -293,7 +293,7 @@ class GAIA:
         print(f"Downloading GAIA dataset to {self.data_dir}...")
         from huggingface_hub import snapshot_download
 
-        local_dir = snapshot_download(
+        local_dir = snapshot_download(  # type: ignore
             repo_id="gaia-benchmark/GAIA",
             repo_type="dataset",
             token=token,

python/samples/README.md

@@ -182,6 +182,7 @@ This directory contains samples demonstrating the capabilities of Microsoft Agen
 | [`getting_started/middleware/middleware_termination.py`](./getting_started/middleware/middleware_termination.py) | Middleware termination example |
 | [`getting_started/middleware/override_result_with_middleware.py`](./getting_started/middleware/override_result_with_middleware.py) | Override result with middleware example |
 | [`getting_started/middleware/shared_state_middleware.py`](./getting_started/middleware/shared_state_middleware.py) | Shared state middleware example |
+| [`getting_started/middleware/thread_behavior_middleware.py`](./getting_started/middleware/thread_behavior_middleware.py) | Thread behavior middleware example demonstrating how to track conversation state across multiple agent runs |
 
 ## Multimodal Input
 

python/samples/getting_started/middleware/README.md

@@ -13,6 +13,7 @@ This folder contains examples demonstrating various middleware patterns with the
 | [`exception_handling_with_middleware.py`](exception_handling_with_middleware.py) | Demonstrates how to use middleware for centralized exception handling in function calls. Shows how to catch exceptions from functions, provide graceful error responses, and override function results when errors occur to provide user-friendly messages. |
 | [`override_result_with_middleware.py`](override_result_with_middleware.py) | Shows how to use middleware to intercept and modify function results after execution, supporting both regular and streaming agent responses. Demonstrates result filtering, formatting, enhancement, and custom streaming response generation. |
 | [`shared_state_middleware.py`](shared_state_middleware.py) | Demonstrates how to implement function-based middleware within a class to share state between multiple middleware functions. Shows how middleware can work together by sharing state, including call counting and result enhancement. |
+| [`thread_behavior_middleware.py`](thread_behavior_middleware.py) | Demonstrates how middleware can access and track thread state across multiple agent runs. Shows how `AgentRunContext.thread` behaves differently before and after the `next()` call, how conversation history accumulates in threads, and timing of thread message updates. Essential for understanding conversation flow in middleware. |
 | [`agent_and_run_level_middleware.py`](agent_and_run_level_middleware.py) | Explains the difference between agent-level middleware (applied to ALL runs of the agent) and run-level middleware (applied to specific runs only). Shows security validation, performance monitoring, and context-specific middleware patterns. |
 | [`chat_middleware.py`](chat_middleware.py) | Demonstrates how to use chat middleware to observe and override inputs sent to AI models. Shows how to intercept chat requests, log and modify input messages, and override entire responses before they reach the underlying AI service. |
 
@@ -42,4 +43,4 @@ This folder contains examples demonstrating various middleware patterns with the
 
 - **Termination**: Use `context.terminate` to stop execution early
 - **Result Override**: Modify or replace function/agent results
-- **Streaming Support**: Handle both regular and streaming responses
+- **Streaming Support**: Handle both regular and streaming responses

python/samples/getting_started/middleware/thread_behavior_middleware.py

@@ -0,0 +1,99 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+from collections.abc import Awaitable, Callable
+from typing import Annotated
+
+from agent_framework import (
+    AgentRunContext,
+    ChatMessageStore,
+)
+from agent_framework.azure import AzureOpenAIChatClient
+from azure.identity import AzureCliCredential
+from pydantic import Field
+
+"""
+Thread Behavior Middleware Example
+
+This sample demonstrates how middleware can access and track thread state across multiple agent runs.
+The example shows:
+
+- How AgentRunContext.thread property behaves across multiple runs
+- How middleware can access conversation history through the thread
+- The timing of when thread messages are populated (before vs after next() call)
+- How to track thread state changes across runs
+
+Key behaviors demonstrated:
+1. First run: context.messages is populated, context.thread is initially empty (before next())
+2. After next(): thread contains input message + response from agent
+3. Second run: context.messages contains only current input, thread contains previous history
+4. After next(): thread contains full conversation history (all previous + current messages)
+"""
+
+
+def get_weather(
+    location: Annotated[str, Field(description="The location to get the weather for.")],
+) -> str:
+    """Get the weather for a given location."""
+    from random import randint
+
+    conditions = ["sunny", "cloudy", "rainy", "stormy"]
+    return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C."
+
+
+async def thread_tracking_middleware(
+    context: AgentRunContext,
+    next: Callable[[AgentRunContext], Awaitable[None]],
+) -> None:
+    """Middleware that tracks and logs thread behavior across runs."""
+    thread_messages = []
+    if context.thread and context.thread.message_store:
+        thread_messages = await context.thread.message_store.list_messages()
+
+    print(f"[Middleware pre-execution] Current input messages: {len(context.messages)}")
+    print(f"[Middleware pre-execution] Thread history messages: {len(thread_messages)}")
+
+    # Call next to execute the agent
+    await next(context)
+
+    # Check thread state after agent execution
+    updated_thread_messages = []
+    if context.thread and context.thread.message_store:
+        updated_thread_messages = await context.thread.message_store.list_messages()
+
+    print(f"[Middleware post-execution] Updated thread messages: {len(updated_thread_messages)}")
+
+
+async def main() -> None:
+    """Example demonstrating thread behavior in middleware across multiple runs."""
+    print("=== Thread Behavior Middleware Example ===")
+
+    # For authentication, run `az login` command in terminal or replace AzureCliCredential with preferred
+    # authentication option.
+    agent = AzureOpenAIChatClient(credential=AzureCliCredential()).create_agent(
+        name="WeatherAgent",
+        instructions="You are a helpful weather assistant.",
+        tools=get_weather,
+        middleware=thread_tracking_middleware,
+        # Configure agent with message store factory to persist conversation history
+        chat_message_store_factory=ChatMessageStore,
+    )
+
+    # Create a thread that will persist messages between runs
+    thread = agent.get_new_thread()
+
+    print("\nFirst Run:")
+    query1 = "What's the weather like in Tokyo?"
+    print(f"User: {query1}")
+    result1 = await agent.run(query1, thread=thread)
+    print(f"Agent: {result1.text}")
+
+    print("\nSecond Run:")
+    query2 = "How about in London?"
+    print(f"User: {query2}")
+    result2 = await agent.run(query2, thread=thread)
+    print(f"Agent: {result2.text}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())