Python: Keep pending Anthropic server tool calls as server_tool_use

python/packages/anthropic/agent_framework_anthropic/_chat_client.py

@@ -5,7 +5,7 @@
 import asyncio
 import logging
 import sys
-from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, Sequence
+from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, Sequence, Set
 from dataclasses import dataclass, field
 from typing import Any, ClassVar, Final, Generic, Literal, TypedDict, cast
 
@@ -901,9 +901,11 @@ def _prepare_messages_for_anthropic(self, messages: Sequence[Message]) -> list[d
         else:
             msgs = list(messages)
 
+        pending_hosted_calls = self._pending_hosted_function_calls(msgs)
+
         result: list[dict[str, Any]] = []
         for msg in msgs:
-            result.extend(self._prepare_message_groups_for_anthropic(msg))
+            result.extend(self._prepare_message_groups_for_anthropic(msg, pending_hosted_calls))
 
         # Anthropic requires the conversation to end with a user message.
         # Append a synthetic user turn so chained agent outputs work as
@@ -913,9 +915,43 @@ def _prepare_messages_for_anthropic(self, messages: Sequence[Message]) -> list[d
 
         return result
 
-    def _prepare_message_groups_for_anthropic(self, message: Message) -> list[dict[str, Any]]:
+    @staticmethod
+    def _pending_hosted_function_calls(messages: Sequence[Message]) -> set[int]:
+        """Return the identities (``id()``) of hosted function calls that have no result yet.
+
+        A hosted (server tool) call without a result is still running on Anthropic's side, for example
+        when Claude called a local tool in parallel. It has to be sent back as ``server_tool_use``, or
+        Anthropic expects a local ``tool_result`` for it.
+
+        A ``call_id`` is not unique across a transcript: a later occurrence may reuse it. Calls and results
+        are therefore correlated in transcript order. A result settles the open call occurrence with the
+        same ``call_id`` whose ``Content.id`` it carries, and otherwise the earliest open occurrence with
+        that ``call_id``. A result never settles a call that appears after it.
+        """
+        open_calls_by_call_id: dict[str, list[Content]] = {}
+        for message in messages:
+            for content in message.contents:
+                if not content.call_id:
+                    continue
+                if content.type == "function_call":
+                    open_calls_by_call_id.setdefault(content.call_id, []).append(content)
+                elif content.type == "function_result" and (open_calls := open_calls_by_call_id.get(content.call_id)):
+                    settled_index = next(
+                        (
+                            index
+                            for index, call in enumerate(open_calls)
+                            if content.id is not None and call.id == content.id
+                        ),
+                        0,
+                    )
+                    open_calls.pop(settled_index)
+        return {id(call) for open_calls in open_calls_by_call_id.values() for call in open_calls if call.hosted}
+
+    def _prepare_message_groups_for_anthropic(
+        self, message: Message, pending_hosted_calls: Set[int] | None = None
+    ) -> list[dict[str, Any]]:
         """Prepare a Message and split Anthropic content blocks into valid role groups."""
-        prepared_message = self._prepare_message_for_anthropic(message)
+        prepared_message = self._prepare_message_for_anthropic(message, pending_hosted_calls)
         content = prepared_message.get("content")
         if not isinstance(content, list):
             return [prepared_message]
@@ -964,11 +1000,17 @@ def _message_has_tool_use(self, message: dict[str, Any]) -> bool:
 
         return False
 
-    def _prepare_message_for_anthropic(self, message: Message) -> dict[str, Any]:
+    def _prepare_message_for_anthropic(
+        self, message: Message, pending_hosted_calls: Set[int] | None = None
+    ) -> dict[str, Any]:
         """Prepare a Message for the Anthropic client.
 
         Args:
             message: The Message to convert.
+            pending_hosted_calls: Identities (``id()``) of the hosted (server tool) call contents that have
+                no result yet, from ``_pending_hosted_function_calls``. These are sent as ``server_tool_use``
+                blocks. An identity rather than a ``call_id`` is used because a later occurrence may reuse
+                the ``call_id`` of a call that already has its result.
 
         Returns:
             A dictionary representing the message in Anthropic format.
@@ -1001,8 +1043,9 @@ def _prepare_message_for_anthropic(self, message: Message) -> dict[str, Any]:
                     else:
                         logger.debug(f"Ignoring unsupported data content media type: {content.media_type} for now")
                 case "function_call":
+                    pending_hosted_call = content.hosted and id(content) in (pending_hosted_calls or ())
                     a_content.append({
-                        "type": "tool_use",
+                        "type": "server_tool_use" if pending_hosted_call else "tool_use",
                         "id": content.call_id,
                         "name": content.name,
                         "input": content.parse_arguments(),

python/packages/anthropic/tests/test_anthropic_client.py

@@ -39,6 +39,7 @@
 from anthropic.types.beta import (
     BetaMessage,
     BetaMessageDeltaUsage,
+    BetaServerToolUseBlock,
     BetaTextBlock,
     BetaToolUseBlock,
     BetaUsage,
@@ -825,6 +826,161 @@ def test_prepare_messages_for_anthropic_without_system(
     assert result[2]["content"] == "Continue"
 
 
+def _parallel_local_and_server_tool_call_contents(client: AnthropicClient) -> list[Content]:
+    """Parse a response in which Claude calls a local tool and a server tool in parallel."""
+    return client._parse_contents_from_anthropic(
+        [
+            BetaToolUseBlock(type="tool_use", id="toolu_local", name="local_value", input={}),
+            BetaServerToolUseBlock(
+                type="server_tool_use", id="srvtoolu_search", name="web_search", input={"query": "release"}
+            ),
+        ],
+        _AnthropicRequestState(),
+    )
+
+
+@pytest.mark.parametrize("reload_history", [False, True], ids=["in_memory", "reloaded"])
+def test_prepare_messages_preserves_parallel_server_tool_use(
+    mock_anthropic_client: MagicMock,
+    reload_history: bool,
+) -> None:
+    """A pending server tool call is sent back as server_tool_use next to the local tool result."""
+    client = create_test_anthropic_client(mock_anthropic_client)
+    messages = [
+        Message(role="user", contents=["Call local_value and search the web."]),
+        Message(role="assistant", contents=_parallel_local_and_server_tool_call_contents(client)),
+        Message(role="tool", contents=[Content.from_function_result(call_id="toolu_local", result="local-ok")]),
+    ]
+    if reload_history:
+        messages = [Message.from_dict(message.to_dict()) for message in messages]
+
+    result = client._prepare_messages_for_anthropic(messages)
+
+    assert [message["role"] for message in result] == ["user", "assistant", "user"]
+    assert result[1]["content"] == [
+        {"type": "tool_use", "id": "toolu_local", "name": "local_value", "input": {}},
+        {"type": "server_tool_use", "id": "srvtoolu_search", "name": "web_search", "input": {"query": "release"}},
+    ]
+    assert [(block["type"], block["tool_use_id"]) for block in result[2]["content"]] == [("tool_result", "toolu_local")]
+
+
+def test_prepare_messages_keeps_completed_server_tool_call_paired_with_its_result(
+    mock_anthropic_client: MagicMock,
+) -> None:
+    """A server tool call that already has its result keeps the existing tool_use/tool_result pairing."""
+    client = create_test_anthropic_client(mock_anthropic_client)
+    contents = _parallel_local_and_server_tool_call_contents(client)
+    search_result = Content.from_function_result(call_id="srvtoolu_search", result="search results")
+    messages = [
+        Message(role="user", contents=["Call local_value and search the web."]),
+        Message(role="assistant", contents=contents),
+        Message(role="tool", contents=[Content.from_function_result(call_id="toolu_local", result="local-ok")]),
+        Message(role="assistant", contents=[search_result, Content.from_text("Both done.")]),
+    ]
+
+    result = client._prepare_messages_for_anthropic(messages)
+
+    assert [block["type"] for block in result[1]["content"]] == ["tool_use", "tool_use"]
+    tool_result_ids = [
+        block["tool_use_id"]
+        for message in result
+        if message["role"] == "user" and isinstance(message["content"], list)
+        for block in message["content"]
+        if block["type"] == "tool_result"
+    ]
+    assert tool_result_ids == ["toolu_local", "srvtoolu_search"]
+
+
+@pytest.mark.parametrize("reload_history", [False, True], ids=["in_memory", "reloaded"])
+def test_prepare_messages_keeps_later_server_tool_call_pending_when_call_id_is_reused(
+    mock_anthropic_client: MagicMock,
+    reload_history: bool,
+) -> None:
+    """An earlier result for a reused call_id does not settle a later pending server tool call."""
+    client = create_test_anthropic_client(mock_anthropic_client)
+    messages = [
+        Message(role="user", contents=["Search the web."]),
+        Message(
+            role="assistant",
+            contents=[
+                Content.from_function_call(
+                    call_id="srvtoolu_search", name="web_search", arguments={"query": "first"}, hosted=True
+                ),
+                Content.from_function_result(call_id="srvtoolu_search", result="first results"),
+                Content.from_text("First search done."),
+            ],
+        ),
+        Message(role="user", contents=["Call local_value and search again."]),
+        Message(role="assistant", contents=_parallel_local_and_server_tool_call_contents(client)),
+        Message(role="tool", contents=[Content.from_function_result(call_id="toolu_local", result="local-ok")]),
+    ]
+    if reload_history:
+        messages = [Message.from_dict(message.to_dict()) for message in messages]
+
+    result = client._prepare_messages_for_anthropic(messages)
+
+    blocks = [
+        (block["type"], block.get("id") or block.get("tool_use_id"))
+        for m in result
+        for block in m["content"]
+        if isinstance(m["content"], list) and block["type"] != "text"
+    ]
+    assert blocks == [
+        ("tool_use", "srvtoolu_search"),
+        ("tool_result", "srvtoolu_search"),
+        ("tool_use", "toolu_local"),
+        ("server_tool_use", "srvtoolu_search"),
+        ("tool_result", "toolu_local"),
+    ]
+
+
+def test_prepare_messages_settles_the_server_tool_call_occurrence_named_by_the_result_id(
+    mock_anthropic_client: MagicMock,
+) -> None:
+    """A result carrying Content.id settles that occurrence, not the earliest one sharing its call_id."""
+    client = create_test_anthropic_client(mock_anthropic_client)
+    messages = [
+        Message(role="user", contents=["Search the web twice."]),
+        Message(
+            role="assistant",
+            contents=[
+                Content.from_function_call(
+                    call_id="srvtoolu_search", id="occurrence_1", name="web_search", arguments={}, hosted=True
+                ),
+                Content.from_function_call(
+                    call_id="srvtoolu_search", id="occurrence_2", name="web_search", arguments={}, hosted=True
+                ),
+                Content.from_function_result(call_id="srvtoolu_search", id="occurrence_2", result="results"),
+            ],
+        ),
+    ]
+
+    result = client._prepare_messages_for_anthropic(messages)
+
+    assert [block["type"] for block in result[1]["content"]] == ["server_tool_use", "tool_use"]
+
+
+def test_prepare_messages_does_not_settle_a_server_tool_call_with_an_earlier_result(
+    mock_anthropic_client: MagicMock,
+) -> None:
+    """A result that precedes every call with its call_id leaves a later server tool call pending."""
+    client = create_test_anthropic_client(mock_anthropic_client)
+    messages = [
+        Message(role="user", contents=["Search the web."]),
+        Message(role="tool", contents=[Content.from_function_result(call_id="srvtoolu_search", result="stale")]),
+        Message(
+            role="assistant",
+            contents=[
+                Content.from_function_call(call_id="srvtoolu_search", name="web_search", arguments={}, hosted=True)
+            ],
+        ),
+    ]
+
+    result = client._prepare_messages_for_anthropic(messages)
+
+    assert [block["type"] for block in result[-1]["content"]] == ["server_tool_use"]
+
+
 def test_prepare_messages_for_anthropic_does_not_append_after_tool_use(
     mock_anthropic_client: MagicMock,
 ) -> None: