Merge branch 'strands-agents:main' into session-persistence

Author: Unshure

pyproject.toml

@@ -90,7 +90,7 @@ writer = [
 ]
 
 a2a = [
-    "a2a-sdk[sql]>=0.2.11",
+    "a2a-sdk[sql]>=0.2.11,<1.0.0",
     "uvicorn>=0.34.2",
     "httpx>=0.28.1",
     "fastapi>=0.115.12",
@@ -136,7 +136,7 @@ all = [
     "opentelemetry-exporter-otlp-proto-http>=1.30.0,<2.0.0",
 
     # a2a
-    "a2a-sdk[sql]>=0.2.11",
+    "a2a-sdk[sql]>=0.2.11,<1.0.0",
     "uvicorn>=0.34.2",
     "httpx>=0.28.1",
     "fastapi>=0.115.12",

src/strands/models/mistral.py

@@ -90,11 +90,9 @@ def __init__(
 
         logger.debug("config=<%s> | initializing", self.config)
 
-        client_args = client_args or {}
+        self.client_args = client_args or {}
         if api_key:
-            client_args["api_key"] = api_key
-
-        self.client = mistralai.Mistral(**client_args)
+            self.client_args["api_key"] = api_key
 
     @override
     def update_config(self, **model_config: Unpack[MistralConfig]) -> None:  # type: ignore
@@ -421,67 +419,70 @@ async def stream(
             logger.debug("got response from model")
             if not self.config.get("stream", True):
                 # Use non-streaming API
-                response = await self.client.chat.complete_async(**request)
-                for event in self._handle_non_streaming_response(response):
-                    yield self.format_chunk(event)
+                async with mistralai.Mistral(**self.client_args) as client:
+                    response = await client.chat.complete_async(**request)
+                    for event in self._handle_non_streaming_response(response):
+                        yield self.format_chunk(event)
+
                 return
 
             # Use the streaming API
-            stream_response = await self.client.chat.stream_async(**request)
+            async with mistralai.Mistral(**self.client_args) as client:
+                stream_response = await client.chat.stream_async(**request)
 
-            yield self.format_chunk({"chunk_type": "message_start"})
+                yield self.format_chunk({"chunk_type": "message_start"})
 
-            content_started = False
-            tool_calls: dict[str, list[Any]] = {}
-            accumulated_text = ""
+                content_started = False
+                tool_calls: dict[str, list[Any]] = {}
+                accumulated_text = ""
 
-            async for chunk in stream_response:
-                if hasattr(chunk, "data") and hasattr(chunk.data, "choices") and chunk.data.choices:
-                    choice = chunk.data.choices[0]
+                async for chunk in stream_response:
+                    if hasattr(chunk, "data") and hasattr(chunk.data, "choices") and chunk.data.choices:
+                        choice = chunk.data.choices[0]
 
-                    if hasattr(choice, "delta"):
-                        delta = choice.delta
+                        if hasattr(choice, "delta"):
+                            delta = choice.delta
 
-                        if hasattr(delta, "content") and delta.content:
-                            if not content_started:
-                                yield self.format_chunk({"chunk_type": "content_start", "data_type": "text"})
-                                content_started = True
+                            if hasattr(delta, "content") and delta.content:
+                                if not content_started:
+                                    yield self.format_chunk({"chunk_type": "content_start", "data_type": "text"})
+                                    content_started = True
 
-                            yield self.format_chunk(
-                                {"chunk_type": "content_delta", "data_type": "text", "data": delta.content}
-                            )
-                            accumulated_text += delta.content
+                                yield self.format_chunk(
+                                    {"chunk_type": "content_delta", "data_type": "text", "data": delta.content}
+                                )
+                                accumulated_text += delta.content
 
-                        if hasattr(delta, "tool_calls") and delta.tool_calls:
-                            for tool_call in delta.tool_calls:
-                                tool_id = tool_call.id
-                                tool_calls.setdefault(tool_id, []).append(tool_call)
+                            if hasattr(delta, "tool_calls") and delta.tool_calls:
+                                for tool_call in delta.tool_calls:
+                                    tool_id = tool_call.id
+                                    tool_calls.setdefault(tool_id, []).append(tool_call)
 
-                    if hasattr(choice, "finish_reason") and choice.finish_reason:
-                        if content_started:
-                            yield self.format_chunk({"chunk_type": "content_stop", "data_type": "text"})
+                        if hasattr(choice, "finish_reason") and choice.finish_reason:
+                            if content_started:
+                                yield self.format_chunk({"chunk_type": "content_stop", "data_type": "text"})
 
-                        for tool_deltas in tool_calls.values():
-                            yield self.format_chunk(
-                                {"chunk_type": "content_start", "data_type": "tool", "data": tool_deltas[0]}
-                            )
+                            for tool_deltas in tool_calls.values():
+                                yield self.format_chunk(
+                                    {"chunk_type": "content_start", "data_type": "tool", "data": tool_deltas[0]}
+                                )
 
-                            for tool_delta in tool_deltas:
-                                if hasattr(tool_delta.function, "arguments"):
-                                    yield self.format_chunk(
-                                        {
-                                            "chunk_type": "content_delta",
-                                            "data_type": "tool",
-                                            "data": tool_delta.function.arguments,
-                                        }
-                                    )
+                                for tool_delta in tool_deltas:
+                                    if hasattr(tool_delta.function, "arguments"):
+                                        yield self.format_chunk(
+                                            {
+                                                "chunk_type": "content_delta",
+                                                "data_type": "tool",
+                                                "data": tool_delta.function.arguments,
+                                            }
+                                        )
 
-                            yield self.format_chunk({"chunk_type": "content_stop", "data_type": "tool"})
+                                yield self.format_chunk({"chunk_type": "content_stop", "data_type": "tool"})
 
-                        yield self.format_chunk({"chunk_type": "message_stop", "data": choice.finish_reason})
+                            yield self.format_chunk({"chunk_type": "message_stop", "data": choice.finish_reason})
 
-                        if hasattr(chunk, "usage"):
-                            yield self.format_chunk({"chunk_type": "metadata", "data": chunk.usage})
+                            if hasattr(chunk, "usage"):
+                                yield self.format_chunk({"chunk_type": "metadata", "data": chunk.usage})
 
         except Exception as e:
             if "rate" in str(e).lower() or "429" in str(e):
@@ -518,7 +519,8 @@ async def structured_output(
         formatted_request["tool_choice"] = "any"
         formatted_request["parallel_tool_calls"] = False
 
-        response = await self.client.chat.complete_async(**formatted_request)
+        async with mistralai.Mistral(**self.client_args) as client:
+            response = await client.chat.complete_async(**formatted_request)
 
         if response.choices and response.choices[0].message.tool_calls:
             tool_call = response.choices[0].message.tool_calls[0]

src/strands/models/ollama.py

@@ -68,14 +68,12 @@ def __init__(
             ollama_client_args: Additional arguments for the Ollama client.
             **model_config: Configuration options for the Ollama model.
         """
+        self.host = host
+        self.client_args = ollama_client_args or {}
         self.config = OllamaModel.OllamaConfig(**model_config)
 
         logger.debug("config=<%s> | initializing", self.config)
 
-        ollama_client_args = ollama_client_args if ollama_client_args is not None else {}
-
-        self.client = ollama.AsyncClient(host, **ollama_client_args)
-
     @override
     def update_config(self, **model_config: Unpack[OllamaConfig]) -> None:  # type: ignore
         """Update the Ollama Model configuration with the provided arguments.
@@ -306,7 +304,8 @@ async def stream(
         logger.debug("invoking model")
         tool_requested = False
 
-        response = await self.client.chat(**request)
+        client = ollama.AsyncClient(self.host, **self.client_args)
+        response = await client.chat(**request)
 
         logger.debug("got response from model")
         yield self.format_chunk({"chunk_type": "message_start"})
@@ -346,7 +345,9 @@ async def structured_output(
         formatted_request = self.format_request(messages=prompt)
         formatted_request["format"] = output_model.model_json_schema()
         formatted_request["stream"] = False
-        response = await self.client.chat(**formatted_request)
+
+        client = ollama.AsyncClient(self.host, **self.client_args)
+        response = await client.chat(**formatted_request)
 
         try:
             content = response.message.content.strip()

src/strands/tools/mcp/mcp_client.py

@@ -16,13 +16,14 @@
 from concurrent import futures
 from datetime import timedelta
 from types import TracebackType
-from typing import Any, Callable, Coroutine, Dict, List, Optional, TypeVar, Union
+from typing import Any, Callable, Coroutine, Dict, Optional, TypeVar, Union
 
 from mcp import ClientSession, ListToolsResult
 from mcp.types import CallToolResult as MCPCallToolResult
 from mcp.types import ImageContent as MCPImageContent
 from mcp.types import TextContent as MCPTextContent
 
+from ...types import PaginatedList
 from ...types.exceptions import MCPClientInitializationError
 from ...types.media import ImageFormat
 from ...types.tools import ToolResult, ToolResultContent, ToolResultStatus
@@ -140,7 +141,7 @@ async def _set_close_event() -> None:
         self._background_thread = None
         self._session_id = uuid.uuid4()
 
-    def list_tools_sync(self) -> List[MCPAgentTool]:
+    def list_tools_sync(self, pagination_token: Optional[str] = None) -> PaginatedList[MCPAgentTool]:
         """Synchronously retrieves the list of available tools from the MCP server.
 
         This method calls the asynchronous list_tools method on the MCP session
@@ -154,14 +155,14 @@ def list_tools_sync(self) -> List[MCPAgentTool]:
             raise MCPClientInitializationError(CLIENT_SESSION_NOT_RUNNING_ERROR_MESSAGE)
 
         async def _list_tools_async() -> ListToolsResult:
-            return await self._background_thread_session.list_tools()
+            return await self._background_thread_session.list_tools(cursor=pagination_token)
 
         list_tools_response: ListToolsResult = self._invoke_on_background_thread(_list_tools_async()).result()
         self._log_debug_with_thread("received %d tools from MCP server", len(list_tools_response.tools))
 
         mcp_tools = [MCPAgentTool(tool, self) for tool in list_tools_response.tools]
         self._log_debug_with_thread("successfully adapted %d MCP tools", len(mcp_tools))
-        return mcp_tools
+        return PaginatedList[MCPAgentTool](mcp_tools, token=list_tools_response.nextCursor)
 
     def call_tool_sync(
         self,

src/strands/types/__init__.py

@@ -1 +1,5 @@
 """SDK type definitions."""
+
+from .collections import PaginatedList
+
+__all__ = ["PaginatedList"]

src/strands/types/collections.py

@@ -0,0 +1,23 @@
+"""Generic collection types for the Strands SDK."""
+
+from typing import Generic, List, Optional, TypeVar
+
+T = TypeVar("T")
+
+
+class PaginatedList(list, Generic[T]):
+    """A generic list-like object that includes a pagination token.
+
+    This maintains backwards compatibility by inheriting from list,
+    so existing code that expects List[T] will continue to work.
+    """
+
+    def __init__(self, data: List[T], token: Optional[str] = None):
+        """Initialize a PaginatedList with data and an optional pagination token.
+
+        Args:
+            data: The list of items to store.
+            token: Optional pagination token for retrieving additional items.
+        """
+        super().__init__(data)
+        self.pagination_token = token

tests/strands/models/test_mistral.py

@@ -11,7 +11,9 @@
 @pytest.fixture
 def mistral_client():
     with unittest.mock.patch.object(strands.models.mistral.mistralai, "Mistral") as mock_client_cls:
-        yield mock_client_cls.return_value
+        mock_client = unittest.mock.AsyncMock()
+        mock_client_cls.return_value.__aenter__.return_value = mock_client
+        yield mock_client
 
 
 @pytest.fixture
@@ -25,9 +27,7 @@ def max_tokens():
 
 
 @pytest.fixture
-def model(mistral_client, model_id, max_tokens):
-    _ = mistral_client
-
+def model(model_id, max_tokens):
     return MistralModel(model_id=model_id, max_tokens=max_tokens)
 
 

tests/strands/models/test_ollama.py

@@ -26,9 +26,7 @@ def host():
 
 
 @pytest.fixture
-def model(ollama_client, model_id, host):
-    _ = ollama_client
-
+def model(model_id, host):
     return OllamaModel(host, model_id=model_id)
 
 

tests/strands/tools/mcp/test_mcp_client.py

@@ -71,10 +71,11 @@ def test_list_tools_sync(mock_transport, mock_session):
     with MCPClient(mock_transport["transport_callable"]) as client:
         tools = client.list_tools_sync()
 
-        mock_session.list_tools.assert_called_once()
+        mock_session.list_tools.assert_called_once_with(cursor=None)
 
         assert len(tools) == 1
         assert tools[0].tool_name == "test_tool"
+        assert tools.pagination_token is None
 
 
 def test_list_tools_sync_session_not_active():
@@ -85,6 +86,34 @@ def test_list_tools_sync_session_not_active():
         client.list_tools_sync()
 
 
+def test_list_tools_sync_with_pagination_token(mock_transport, mock_session):
+    """Test that list_tools_sync correctly passes pagination token and returns next cursor."""
+    mock_tool = MCPTool(name="test_tool", description="A test tool", inputSchema={"type": "object", "properties": {}})
+    mock_session.list_tools.return_value = ListToolsResult(tools=[mock_tool], nextCursor="next_page_token")
+
+    with MCPClient(mock_transport["transport_callable"]) as client:
+        tools = client.list_tools_sync(pagination_token="current_page_token")
+
+        mock_session.list_tools.assert_called_once_with(cursor="current_page_token")
+        assert len(tools) == 1
+        assert tools[0].tool_name == "test_tool"
+        assert tools.pagination_token == "next_page_token"
+
+
+def test_list_tools_sync_without_pagination_token(mock_transport, mock_session):
+    """Test that list_tools_sync works without pagination token and handles missing next cursor."""
+    mock_tool = MCPTool(name="test_tool", description="A test tool", inputSchema={"type": "object", "properties": {}})
+    mock_session.list_tools.return_value = ListToolsResult(tools=[mock_tool])  # No nextCursor
+
+    with MCPClient(mock_transport["transport_callable"]) as client:
+        tools = client.list_tools_sync()
+
+        mock_session.list_tools.assert_called_once_with(cursor=None)
+        assert len(tools) == 1
+        assert tools[0].tool_name == "test_tool"
+        assert tools.pagination_token is None
+
+
 @pytest.mark.parametrize("is_error,expected_status", [(False, "success"), (True, "error")])
 def test_call_tool_sync_status(mock_transport, mock_session, is_error, expected_status):
     """Test that call_tool_sync correctly handles success and error results."""