Python: [BREAKING] cleanup of thread API and serialization (#893)

* cleanup of threads and serialization

* fix for sliding window

* fix redis test

* updated from comments

* updated context provider and threads

* updated lock

* add asyncio default

* fix redis tests

* fix tests

* fix tests

* renamed to invoking

* fixed tests

* fix for instructions

python/packages/lab/pyproject.toml

@@ -130,7 +130,9 @@ test-tau2 = "pytest tau2/tests --cov=agent_framework_lab_tau2 --cov-report=term-
 [tool.pytest.ini_options]
 pythonpath = ["."]
 addopts = "--strict-markers --strict-config"
+asyncio_mode = "auto"
+asyncio_default_fixture_loop_scope = "function"
 markers = [
   "unit: marks tests as unit tests",
   "integration: marks tests as integration tests",
-]
+]

python/packages/lab/tau2/agent_framework_lab_tau2/_sliding_window.py

@@ -5,13 +5,12 @@ from collections.abc import Sequence
 from typing import Any
 
 import tiktoken
-from agent_framework._threads import ChatMessageList
-from agent_framework._types import ChatMessage, Role
+from agent_framework import ChatMessage, ChatMessageStore, Role
 from loguru import logger
 
 
-class SlidingWindowChatMessageList(ChatMessageList):
-    """A token-aware sliding window implementation of ChatMessageList.
+class SlidingWindowChatMessageStore(ChatMessageStore):
+    """A token-aware sliding window implementation of ChatMessageStore.
 
     Maintains two message lists: complete history and truncated window.
     Automatically removes oldest messages when token limit is exceeded.
@@ -25,8 +24,8 @@ class SlidingWindowChatMessageList(ChatMessageList):
         system_message: str | None = None,
         tool_definitions: Any | None = None,
     ):
-        super().__init__(messages)
-        self._truncated_messages = self._messages.copy()  # Separate truncated view
+        super().__init__(messages=messages)
+        self.truncated_messages = self.messages.copy()
         self.max_tokens = max_tokens
         self.system_message = system_message  # Included in token count
         self.tool_definitions = tool_definitions
@@ -36,25 +35,25 @@ class SlidingWindowChatMessageList(ChatMessageList):
     async def add_messages(self, messages: Sequence[ChatMessage]) -> None:
         await super().add_messages(messages)
 
-        self._truncated_messages = self._messages.copy()
+        self.truncated_messages = self.messages.copy()
         self.truncate_messages()
 
     async def list_messages(self) -> list[ChatMessage]:
         """Get the current list of messages, which may be truncated."""
-        return self._truncated_messages
+        return self.truncated_messages
 
     async def list_all_messages(self) -> list[ChatMessage]:
         """Get all messages from the store including the truncated ones."""
-        return self._messages
+        return self.messages
 
     def truncate_messages(self) -> None:
-        while len(self._truncated_messages) > 0 and self.get_token_count() > self.max_tokens:
+        while len(self.truncated_messages) > 0 and self.get_token_count() > self.max_tokens:
             logger.warning("Messages exceed max tokens. Truncating oldest message.")
-            self._truncated_messages.pop(0)
+            self.truncated_messages.pop(0)
         # Remove leading tool messages
-        while len(self._truncated_messages) > 0 and self._truncated_messages[0].role == Role.TOOL:
+        while len(self.truncated_messages) > 0 and self.truncated_messages[0].role == Role.TOOL:
             logger.warning("Removing leading tool message because tool result cannot be the first message.")
-            self._truncated_messages.pop(0)
+            self.truncated_messages.pop(0)
 
     def get_token_count(self) -> int:
         """Estimate token count for a list of messages using tiktoken.
@@ -72,7 +71,7 @@ class SlidingWindowChatMessageList(ChatMessageList):
             total_tokens += len(self.encoding.encode(self.system_message))
             total_tokens += 4  # Extra tokens for system message formatting
 
-        for msg in self._truncated_messages:
+        for msg in self.truncated_messages:
             # Add 4 tokens per message for role, formatting, etc.
             total_tokens += 4
 

python/packages/lab/tau2/agent_framework_lab_tau2/runner.py

@@ -29,7 +29,7 @@ from tau2.user.user_simulator import (  # type: ignore[import-untyped]
 from tau2.utils.utils import get_now  # type: ignore[import-untyped]
 
 from ._message_utils import flip_messages, log_messages
-from ._sliding_window import SlidingWindowChatMessageList
+from ._sliding_window import SlidingWindowChatMessageStore
 from ._tau2_utils import convert_agent_framework_messages_to_tau2_messages, convert_tau2_tool_to_ai_function
 
 # Agent instructions matching tau2's LLMAgent
@@ -196,7 +196,7 @@ class TaskRunner:
             instructions=assistant_system_prompt,
             tools=ai_functions,  # type: ignore
             temperature=self.assistant_sampling_temperature,
-            chat_message_store_factory=lambda: SlidingWindowChatMessageList(
+            chat_message_store_factory=lambda: SlidingWindowChatMessageStore(
                 system_message=assistant_system_prompt,
                 tool_definitions=[tool.openai_schema for tool in tools],
                 max_tokens=self.assistant_window_size,
@@ -352,7 +352,7 @@ class TaskRunner:
         # 2. The assistant's message store (not just the truncated window)
         # 3. The final user message (if any)
         assistant_executor = cast(AgentExecutor, self._assistant_executor)
-        message_store = cast(SlidingWindowChatMessageList, assistant_executor._agent_thread.message_store)
+        message_store = cast(SlidingWindowChatMessageStore, assistant_executor._agent_thread.message_store)
         full_conversation = [first_message] + await message_store.list_all_messages()
         if self._final_user_message is not None:
             full_conversation.extend(self._final_user_message)

python/packages/lab/tau2/tests/test_sliding_window.py

@@ -4,20 +4,19 @@
 
 from unittest.mock import patch
 
-import pytest
 from agent_framework._types import ChatMessage, FunctionCallContent, FunctionResultContent, Role, TextContent
-from agent_framework_lab_tau2._sliding_window import SlidingWindowChatMessageList
+from agent_framework_lab_tau2._sliding_window import SlidingWindowChatMessageStore
 
 
 def test_initialization_empty():
     """Test initializing with no messages."""
-    sliding_window = SlidingWindowChatMessageList(max_tokens=1000)
+    sliding_window = SlidingWindowChatMessageStore(max_tokens=1000)
 
     assert sliding_window.max_tokens == 1000
     assert sliding_window.system_message is None
     assert sliding_window.tool_definitions is None
-    assert len(sliding_window._messages) == 0
-    assert len(sliding_window._truncated_messages) == 0
+    assert len(sliding_window.messages) == 0
+    assert len(sliding_window.truncated_messages) == 0
 
 
 def test_initialization_with_parameters():
@@ -25,7 +24,7 @@ def test_initialization_with_parameters():
     system_msg = "You are a helpful assistant"
     tool_defs = [{"name": "test_tool", "description": "A test tool"}]
 
-    sliding_window = SlidingWindowChatMessageList(
+    sliding_window = SlidingWindowChatMessageStore(
         max_tokens=2000, system_message=system_msg, tool_definitions=tool_defs
     )
 
@@ -41,16 +40,15 @@ def test_initialization_with_messages():
         ChatMessage(role=Role.ASSISTANT, contents=[TextContent(text="Hi there!")]),
     ]
 
-    sliding_window = SlidingWindowChatMessageList(messages=messages, max_tokens=1000)
+    sliding_window = SlidingWindowChatMessageStore(messages=messages, max_tokens=1000)
 
-    assert len(sliding_window._messages) == 2
-    assert len(sliding_window._truncated_messages) == 2
+    assert len(sliding_window.messages) == 2
+    assert len(sliding_window.truncated_messages) == 2
 
 
-@pytest.mark.asyncio
 async def test_add_messages_simple():
     """Test adding messages without truncation."""
-    sliding_window = SlidingWindowChatMessageList(max_tokens=10000)  # Large limit
+    sliding_window = SlidingWindowChatMessageStore(max_tokens=10000)  # Large limit
 
     new_messages = [
         ChatMessage(role=Role.USER, contents=[TextContent(text="What's the weather?")]),
@@ -65,10 +63,9 @@ async def test_add_messages_simple():
     assert messages[1].text == "I can help with that."
 
 
-@pytest.mark.asyncio
 async def test_list_all_messages_vs_list_messages():
     """Test difference between list_all_messages and list_messages."""
-    sliding_window = SlidingWindowChatMessageList(max_tokens=50)  # Small limit to force truncation
+    sliding_window = SlidingWindowChatMessageStore(max_tokens=50)  # Small limit to force truncation
 
     # Add many messages to trigger truncation
     messages = [
@@ -89,8 +86,8 @@ async def test_list_all_messages_vs_list_messages():
 
 def test_get_token_count_basic():
     """Test basic token counting."""
-    sliding_window = SlidingWindowChatMessageList(max_tokens=1000)
-    sliding_window._truncated_messages = [ChatMessage(role=Role.USER, contents=[TextContent(text="Hello")])]
+    sliding_window = SlidingWindowChatMessageStore(max_tokens=1000)
+    sliding_window.truncated_messages = [ChatMessage(role=Role.USER, contents=[TextContent(text="Hello")])]
 
     token_count = sliding_window.get_token_count()
 
@@ -101,13 +98,13 @@ def test_get_token_count_basic():
 def test_get_token_count_with_system_message():
     """Test token counting includes system message."""
     system_msg = "You are a helpful assistant"
-    sliding_window = SlidingWindowChatMessageList(max_tokens=1000, system_message=system_msg)
+    sliding_window = SlidingWindowChatMessageStore(max_tokens=1000, system_message=system_msg)
 
     # Without messages
     token_count_empty = sliding_window.get_token_count()
 
     # Add a message
-    sliding_window._truncated_messages = [ChatMessage(role=Role.USER, contents=[TextContent(text="Hello")])]
+    sliding_window.truncated_messages = [ChatMessage(role=Role.USER, contents=[TextContent(text="Hello")])]
     token_count_with_message = sliding_window.get_token_count()
 
     # With message should be more tokens
@@ -119,8 +116,8 @@ def test_get_token_count_function_call():
     """Test token counting with function calls."""
     function_call = FunctionCallContent(call_id="call_123", name="test_function", arguments={"param": "value"})
 
-    sliding_window = SlidingWindowChatMessageList(max_tokens=1000)
-    sliding_window._truncated_messages = [ChatMessage(role=Role.ASSISTANT, contents=[function_call])]
+    sliding_window = SlidingWindowChatMessageStore(max_tokens=1000)
+    sliding_window.truncated_messages = [ChatMessage(role=Role.ASSISTANT, contents=[function_call])]
 
     token_count = sliding_window.get_token_count()
     assert token_count > 0
@@ -130,8 +127,8 @@ def test_get_token_count_function_result():
     """Test token counting with function results."""
     function_result = FunctionResultContent(call_id="call_123", result={"success": True, "data": "result"})
 
-    sliding_window = SlidingWindowChatMessageList(max_tokens=1000)
-    sliding_window._truncated_messages = [ChatMessage(role=Role.TOOL, contents=[function_result])]
+    sliding_window = SlidingWindowChatMessageStore(max_tokens=1000)
+    sliding_window.truncated_messages = [ChatMessage(role=Role.TOOL, contents=[function_result])]
 
     token_count = sliding_window.get_token_count()
     assert token_count > 0
@@ -140,7 +137,7 @@ def test_get_token_count_function_result():
 @patch("agent_framework_lab_tau2._sliding_window.logger")
 def test_truncate_messages_removes_old_messages(mock_logger):
     """Test that truncation removes old messages when token limit exceeded."""
-    sliding_window = SlidingWindowChatMessageList(max_tokens=20)  # Very small limit
+    sliding_window = SlidingWindowChatMessageStore(max_tokens=20)  # Very small limit
 
     # Create messages that will exceed the limit
     messages = [
@@ -155,11 +152,11 @@ def test_truncate_messages_removes_old_messages(mock_logger):
         ChatMessage(role=Role.USER, contents=[TextContent(text="Short msg")]),
     ]
 
-    sliding_window._truncated_messages = messages.copy()
+    sliding_window.truncated_messages = messages.copy()
     sliding_window.truncate_messages()
 
     # Should have fewer messages after truncation
-    assert len(sliding_window._truncated_messages) < len(messages)
+    assert len(sliding_window.truncated_messages) < len(messages)
 
     # Should have logged warnings
     assert mock_logger.warning.called
@@ -168,18 +165,18 @@ def test_truncate_messages_removes_old_messages(mock_logger):
 @patch("agent_framework_lab_tau2._sliding_window.logger")
 def test_truncate_messages_removes_leading_tool_messages(mock_logger):
     """Test that truncation removes leading tool messages."""
-    sliding_window = SlidingWindowChatMessageList(max_tokens=10000)  # Large limit
+    sliding_window = SlidingWindowChatMessageStore(max_tokens=10000)  # Large limit
 
     # Create messages starting with tool message
     tool_message = ChatMessage(role=Role.TOOL, contents=[FunctionResultContent(call_id="call_123", result="result")])
     user_message = ChatMessage(role=Role.USER, contents=[TextContent(text="Hello")])
 
-    sliding_window._truncated_messages = [tool_message, user_message]
+    sliding_window.truncated_messages = [tool_message, user_message]
     sliding_window.truncate_messages()
 
     # Tool message should be removed from the beginning
-    assert len(sliding_window._truncated_messages) == 1
-    assert sliding_window._truncated_messages[0].role == Role.USER
+    assert len(sliding_window.truncated_messages) == 1
+    assert sliding_window.truncated_messages[0].role == Role.USER
 
     # Should have logged warning about removing tool message
     mock_logger.warning.assert_called()
@@ -187,7 +184,7 @@ def test_truncate_messages_removes_leading_tool_messages(mock_logger):
 
 def test_estimate_any_object_token_count_dict():
     """Test token counting for dictionary objects."""
-    sliding_window = SlidingWindowChatMessageList(max_tokens=1000)
+    sliding_window = SlidingWindowChatMessageStore(max_tokens=1000)
 
     test_dict = {"key": "value", "number": 42}
     token_count = sliding_window.estimate_any_object_token_count(test_dict)
@@ -197,7 +194,7 @@ def test_estimate_any_object_token_count_dict():
 
 def test_estimate_any_object_token_count_string():
     """Test token counting for string objects."""
-    sliding_window = SlidingWindowChatMessageList(max_tokens=1000)
+    sliding_window = SlidingWindowChatMessageStore(max_tokens=1000)
 
     test_string = "This is a test string"
     token_count = sliding_window.estimate_any_object_token_count(test_string)
@@ -207,7 +204,7 @@ def test_estimate_any_object_token_count_string():
 
 def test_estimate_any_object_token_count_non_serializable():
     """Test token counting for non-JSON-serializable objects."""
-    sliding_window = SlidingWindowChatMessageList(max_tokens=1000)
+    sliding_window = SlidingWindowChatMessageStore(max_tokens=1000)
 
     # Create an object that can't be JSON serialized
     class CustomObject:
@@ -221,10 +218,9 @@ def test_estimate_any_object_token_count_non_serializable():
     assert token_count > 0
 
 
-@pytest.mark.asyncio
 async def test_real_world_scenario():
     """Test a realistic conversation scenario."""
-    sliding_window = SlidingWindowChatMessageList(
+    sliding_window = SlidingWindowChatMessageStore(
         max_tokens=30,
         system_message="You are a helpful assistant",  # Moderate limit
     )