"""Tests for kind-based Anthropic Messages dispatch in the agent loop.""" from __future__ import annotations from typing import TYPE_CHECKING, Any, Literal, cast, overload import pytest from msgspec import UNSET from plyngent.agent import ( AssistantMessageEvent, ReasoningDeltaEvent, TextDeltaEvent, ToolCallEvent, ToolRegistry, ToolResultEvent, UsageEvent, run_chat_loop, tool, ) from plyngent.agent.messages_dispatch import dispatch_messages from plyngent.lmproto.anthropic.model import ( AnthropicContentBlockDelta, AnthropicContentBlockStart, AnthropicContentBlockStop, AnthropicMessageDelta, AnthropicMessageResponse, AnthropicMessagesParam, AnthropicMessageStart, AnthropicMessageStop, AnthropicRawContentBlock, AnthropicResponseText, AnthropicResponseToolUse, AnthropicStreamEvent, AnthropicToolResultContent, AnthropicUsage, AnthropicUserMessage, ) from plyngent.lmproto.openai_compatible.model import ( AssistantChatMessage, ChatCompletionsParam, UserChatMessage, ) if TYPE_CHECKING: from collections.abc import AsyncIterator, Sequence # What the API answers a request whose thinking block it does not take (a gateway # paraphrase of the documented wording, which names adaptive thinking either way). _THINKING_REFUSED_400: dict[str, str] = { "enabled": ( 'messages HTTP 400: {"error":{"type":"error","message":"claude-opus-5-5 requires adaptive ' 'thinking; omit thinking or use thinking.type=adaptive and output_config.effort"}}' ), "disabled": ( 'messages HTTP 400: {"error":{"type":"error","message":"\\"thinking.type.disabled\\" is not ' 'supported for this model. Use \\"thinking.type.adaptive\\" and \\"output_config.effort\\" to ' 'control thinking behavior."}}' ), } def _message_response( *, text: str = "hello", tool_uses: list[tuple[str, str, dict[str, Any]]] | None = None, stop_reason: str = "end_turn", usage: AnthropicUsage | None = None, ) -> AnthropicMessageResponse: content: list[Any] = [] if text: content.append(AnthropicResponseText(text=text)) if tool_uses: for call_id, name, inp in tool_uses: content.append(AnthropicResponseToolUse(id=call_id, name=name, input=inp)) stop_reason = "tool_use" return AnthropicMessageResponse( id="msg_test", model="claude-test", content=content, stop_reason=stop_reason, usage=usage or AnthropicUsage(input_tokens=10, output_tokens=5), ) class ScriptedMessagesClient: """Fake AnthropicClient that scripts ``messages()`` results.""" kind: str = "messages" calls: list[AnthropicMessagesParam] # Mirrors AnthropicClient: explicit thinking budget (0 = derive from effort), # and the thinking types the dispatch found the API refusing. thinking_budget_tokens: int refused_thinking_types: dict[str, set[str]] _non_stream: list[AnthropicMessageResponse] _stream_events: list[list[AnthropicStreamEvent]] _refusals: frozenset[str] _error: str | None def __init__( self, *, non_stream: Sequence[AnthropicMessageResponse] | None = None, stream_events: Sequence[Sequence[AnthropicStreamEvent]] | None = None, refusals: Sequence[str] = (), error: str | None = None, ) -> None: self.calls = [] self.thinking_budget_tokens = 0 self.refused_thinking_types = {} self._non_stream = list(non_stream or ()) self._stream_events = [list(events) for events in (stream_events or ())] self._refusals = frozenset(refusals) self._error = error def refuses_thinking_type(self, model: str, thinking_type: str) -> bool: return thinking_type in self.refused_thinking_types.get(model, ()) def note_refused_thinking_type(self, model: str, thinking_type: str) -> None: self.refused_thinking_types.setdefault(model, set()).add(thinking_type) def _reject(self, param: AnthropicMessagesParam) -> None: """Raise what the API raises for a request it refuses to run.""" if self._error is not None: raise RuntimeError(self._error) if param.thinking is not UNSET and param.thinking.type in self._refusals: raise RuntimeError(_THINKING_REFUSED_400[param.thinking.type]) @overload async def messages( self, param: AnthropicMessagesParam, *, stream: Literal[False] = False ) -> AnthropicMessageResponse: ... @overload async def messages( self, param: AnthropicMessagesParam, *, stream: Literal[True] ) -> AsyncIterator[AnthropicStreamEvent]: ... async def messages( self, param: AnthropicMessagesParam, *, stream: bool = False ) -> AnthropicMessageResponse | AsyncIterator[AnthropicStreamEvent]: self.calls.append(param) if stream: if not self._stream_events: msg = "no more scripted stream events" raise RuntimeError(msg) async def _gen() -> AsyncIterator[AnthropicStreamEvent]: # The real client reads the error body on the first iteration, # so a refused request consumes no scripted events. self._reject(param) for event in self._stream_events.pop(0): yield event return _gen() self._reject(param) if not self._non_stream: msg = "no more scripted messages" raise RuntimeError(msg) return self._non_stream.pop(0) async def test_dispatch_messages_non_stream_text() -> None: client = ScriptedMessagesClient(non_stream=[_message_response(text="hi there")]) param = ChatCompletionsParam( model="claude-test", messages=[UserChatMessage(content="hello")], ) result = await dispatch_messages(cast("Any", client), param, stream=False) assert result.choices[0].message.content == "hi there" assert result.choices[0].finish_reason == "stop" assert len(client.calls) == 1 assert client.calls[0].max_tokens == 8192 async def test_dispatch_messages_stream_text_and_usage() -> None: events: list[AnthropicStreamEvent] = [ AnthropicMessageStart( message=AnthropicMessageResponse( id="msg_s", model="claude-test", content=[], usage=AnthropicUsage(input_tokens=8, output_tokens=0), ) ), AnthropicContentBlockStart( index=0, content_block=AnthropicRawContentBlock(type="text", text=""), ), AnthropicContentBlockDelta( index=0, delta=AnthropicRawContentBlock(type="text_delta", text="stream"), ), AnthropicContentBlockDelta( index=0, delta=AnthropicRawContentBlock(type="text_delta", text="ed"), ), AnthropicMessageDelta( delta={"stop_reason": "end_turn"}, usage=AnthropicUsage(input_tokens=0, output_tokens=4), ), AnthropicMessageStop(), ] client = ScriptedMessagesClient(stream_events=[events]) param = ChatCompletionsParam( model="claude-test", messages=[UserChatMessage(content="hi")], ) stream = await dispatch_messages(cast("Any", client), param, stream=True) chunks = [chunk async for chunk in cast("Any", stream)] texts = [c.choices[0].delta.content for c in chunks if c.choices and isinstance(c.choices[0].delta.content, str)] assert texts == ["stream", "ed"] assert any(c.choices and c.choices[0].finish_reason == "stop" for c in chunks) assert any(getattr(c, "usage", UNSET) not in (UNSET, None) for c in chunks) async def test_run_chat_loop_messages_kind_text() -> None: client = ScriptedMessagesClient(non_stream=[_message_response(text="done")]) messages: list[Any] = [UserChatMessage(content="go")] events = [ event async for event in run_chat_loop( cast("Any", client), messages, model="claude-test", stream=False, ) ] assert any(isinstance(e, TextDeltaEvent) and e.content == "done" for e in events) assert any(isinstance(e, AssistantMessageEvent) for e in events) assert any(isinstance(e, UsageEvent) and e.usage.prompt_tokens == 10 for e in events) assert isinstance(messages[-1], AssistantChatMessage) assert messages[-1].content == "done" # No thinking configured → the request carries no thinking block. assert client.calls[0].thinking is UNSET async def test_run_chat_loop_passes_reasoning_effort_and_client_thinking_budget() -> None: client = ScriptedMessagesClient(non_stream=[_message_response(text="done")]) client.thinking_budget_tokens = 12000 events = [ event async for event in run_chat_loop( cast("Any", client), [UserChatMessage(content="go")], model="claude-test", stream=False, reasoning_effort="high", ) ] assert any(isinstance(e, TextDeltaEvent) for e in events) thinking = client.calls[0].thinking assert thinking is not UNSET assert thinking.type == "enabled" # The client's explicit budget wins over the level's derived one. assert thinking.budget_tokens == 12000 async def test_run_chat_loop_reasoning_effort_none_disables_thinking() -> None: client = ScriptedMessagesClient(non_stream=[_message_response(text="done")]) _ = [ event async for event in run_chat_loop( cast("Any", client), [UserChatMessage(content="go")], model="claude-test", stream=False, reasoning_effort="none", ) ] thinking = client.calls[0].thinking assert thinking is not UNSET assert thinking.type == "disabled" async def test_run_chat_loop_messages_kind_tool_round() -> None: first = _message_response( text="", tool_uses=[("call_1", "add", {"a": 1, "b": 2})], ) second = _message_response(text="sum is 3") client = ScriptedMessagesClient(non_stream=[first, second]) @tool(register=False) def add(a: int, b: int) -> str: return str(a + b) registry = ToolRegistry([add]) messages: list[Any] = [UserChatMessage(content="add 1 and 2")] events = [ event async for event in run_chat_loop( cast("Any", client), messages, model="claude-test", tools=registry, stream=False, ) ] assert any(isinstance(e, ToolCallEvent) for e in events) assert any(isinstance(e, ToolResultEvent) and "3" in e.message.content for e in events) assert any(isinstance(e, AssistantMessageEvent) and e.message.content == "sum is 3" for e in events) assert len(client.calls) == 2 # Second turn must include tool_result in a user message. second_msgs = client.calls[1].messages assert any( isinstance(m, AnthropicUserMessage) and isinstance(m.content, list) and any(isinstance(b, AnthropicToolResultContent) and b.tool_use_id == "call_1" for b in m.content) for m in second_msgs ) async def test_stream_tool_call_deltas() -> None: events: list[AnthropicStreamEvent] = [ AnthropicMessageStart( message=AnthropicMessageResponse( id="msg_s", model="claude-test", content=[], usage=AnthropicUsage(input_tokens=5, output_tokens=0), ) ), AnthropicContentBlockStart( index=0, content_block=AnthropicRawContentBlock(type="tool_use", id="tu_9", name="add"), ), AnthropicContentBlockDelta( index=0, delta=AnthropicRawContentBlock(type="input_json_delta", partial_json='{"a":'), ), AnthropicContentBlockDelta( index=0, delta=AnthropicRawContentBlock(type="input_json_delta", partial_json="1}"), ), AnthropicMessageDelta( delta={"stop_reason": "tool_use"}, usage=AnthropicUsage(input_tokens=0, output_tokens=6), ), AnthropicMessageStop(), ] client = ScriptedMessagesClient(stream_events=[events]) @tool(register=False) def add(a: int) -> str: return str(a) registry = ToolRegistry([add]) # Second streamed round after tool execution. client._stream_events.append( [ AnthropicMessageStart( message=AnthropicMessageResponse( id="msg_2", model="claude-test", content=[], usage=AnthropicUsage(input_tokens=12, output_tokens=0), ) ), AnthropicContentBlockDelta( index=0, delta=AnthropicRawContentBlock(type="text_delta", text="ok"), ), AnthropicMessageDelta( delta={"stop_reason": "end_turn"}, usage=AnthropicUsage(input_tokens=0, output_tokens=1), ), AnthropicMessageStop(), ] ) messages: list[Any] = [UserChatMessage(content="go")] events_out = [ event async for event in run_chat_loop( cast("Any", client), messages, model="claude-test", tools=registry, stream=True, ) ] assert any(isinstance(e, ToolCallEvent) for e in events_out) assert any(isinstance(e, ToolResultEvent) for e in events_out) assert any(isinstance(e, AssistantMessageEvent) and e.message.content == "ok" for e in events_out) async def test_dispatch_stream_missing_terminal_yields_no_finish() -> None: """Stream with payload but no message_stop/stop_reason yields no finish chunk.""" events: list[AnthropicStreamEvent] = [ AnthropicMessageStart( message=AnthropicMessageResponse( id="msg_s", model="claude-test", content=[], usage=AnthropicUsage(input_tokens=5, output_tokens=0), ) ), AnthropicContentBlockDelta( index=0, delta=AnthropicRawContentBlock(type="text_delta", text="partial"), ), # no message_delta stop_reason, no message_stop ] client = ScriptedMessagesClient(stream_events=[events]) param = ChatCompletionsParam(model="claude-test", messages=[UserChatMessage(content="x")]) stream = await dispatch_messages(cast("Any", client), param, stream=True) chunks = [chunk async for chunk in cast("Any", stream)] # Only the text delta; no finish_reason / usage terminal chunk. assert len(chunks) == 1 assert chunks[0].choices and chunks[0].choices[0].delta.content == "partial" assert not any(c.choices and c.choices[0].finish_reason not in (None, UNSET) for c in chunks) async def test_dispatch_stream_tool_calls_without_stop_yields_no_finish() -> None: """Tool-call deltas without a stop signal must NOT fabricate finish_reason.""" events: list[AnthropicStreamEvent] = [ AnthropicMessageStart( message=AnthropicMessageResponse( id="msg_s", model="claude-test", content=[], usage=AnthropicUsage(input_tokens=5, output_tokens=0), ) ), AnthropicContentBlockStart( index=0, content_block=AnthropicRawContentBlock(type="tool_use", id="tu_1", name="add"), ), # stream dies here: no message_delta stop_reason, no message_stop ] client = ScriptedMessagesClient(stream_events=[events]) param = ChatCompletionsParam(model="claude-test", messages=[UserChatMessage(content="x")]) stream = await dispatch_messages(cast("Any", client), param, stream=True) chunks = [chunk async for chunk in cast("Any", stream)] # tool-call start delta only; no fabricated finish_reason=tool_calls assert not any(c.choices and c.choices[0].finish_reason not in (None, UNSET) for c in chunks) def _search_stream_events(*, text: bool = True) -> list[AnthropicStreamEvent]: """The stream a DeepSeek search turn sends: thinking, search, results, text.""" events: list[AnthropicStreamEvent] = [ AnthropicMessageStart( message=AnthropicMessageResponse( id="msg_s", model="deepseek-flash", content=[], usage=AnthropicUsage(input_tokens=164, output_tokens=0), ) ), AnthropicContentBlockStart( index=0, content_block=AnthropicRawContentBlock(type="thinking", thinking=""), ), AnthropicContentBlockDelta(index=0, delta=AnthropicRawContentBlock(type="thinking_delta", thinking="Search ")), AnthropicContentBlockDelta( index=0, delta=AnthropicRawContentBlock(type="thinking_delta", thinking="the docs."), ), AnthropicContentBlockDelta(index=0, delta=AnthropicRawContentBlock(type="signature_delta", signature="sig-1")), AnthropicContentBlockStop(index=0), AnthropicContentBlockStart( index=1, content_block=AnthropicRawContentBlock(type="server_tool_use", id="call_00_1", name="web_search", input={}), ), AnthropicContentBlockDelta( index=1, delta=AnthropicRawContentBlock(type="input_json_delta", partial_json='{"query":'), ), AnthropicContentBlockDelta( index=1, delta=AnthropicRawContentBlock(type="input_json_delta", partial_json='"DeepSeek-V4"}'), ), AnthropicContentBlockStop(index=1), AnthropicContentBlockStart( index=2, content_block=AnthropicRawContentBlock( type="web_search_tool_result", tool_use_id="call_00_1", content=[ { "type": "web_search_result", "url": "https://api-docs.deepseek.com", "encrypted_content": "sd918tZXSvVp", } ], ), ), AnthropicContentBlockStop(index=2), ] if text: events.extend( [ AnthropicContentBlockStart(index=3, content_block=AnthropicRawContentBlock(type="text", text="")), AnthropicContentBlockDelta( index=3, delta=AnthropicRawContentBlock(type="text_delta", text="DeepSeek-V4 is out."), ), AnthropicContentBlockStop(index=3), ] ) events.extend( [ AnthropicMessageDelta( delta={"stop_reason": "end_turn"}, usage=AnthropicUsage(input_tokens=0, output_tokens=12), ), AnthropicMessageStop(), ] ) return events async def test_search_turn_streams_its_text_without_a_local_tool_call() -> None: """A server tool's arguments must not become a call for the registry. DeepSeek streams a ``server_tool_use`` block whose ``input_json_delta`` fragments carry the search query. Reading those as a local call's arguments asks the registry for a tool the *server* already ran — and would have the agent execute ``web_search`` (or a nameless tool) itself. """ client = ScriptedMessagesClient(stream_events=[_search_stream_events()]) @tool(register=False) def echo(text: str) -> str: return text events = [ event async for event in run_chat_loop( cast("Any", client), [UserChatMessage(content="what is DeepSeek-V4?")], model="deepseek-flash", tools=ToolRegistry([echo]), stream=True, ) ] assert not any(isinstance(event, ToolCallEvent) for event in events) assert not any(isinstance(event, ToolResultEvent) for event in events) assert any(isinstance(event, TextDeltaEvent) and event.content == "DeepSeek-V4 is out." for event in events) # The chain of thought streams as reasoning, fragment by fragment. assert [event.content for event in events if isinstance(event, ReasoningDeltaEvent)] == ["Search ", "the docs."] # ...and the round keeps the search blocks, rebuilt, for the next request. assistant = next(event.message for event in events if isinstance(event, AssistantMessageEvent)) assert assistant.reasoning_content == "Search the docs." assert assistant.provider_blocks == [ {"type": "thinking", "thinking": "Search the docs.", "signature": "sig-1"}, { "type": "server_tool_use", "id": "call_00_1", "name": "web_search", "input": {"query": "DeepSeek-V4"}, }, { "type": "web_search_tool_result", "tool_use_id": "call_00_1", "content": [ { "type": "web_search_result", "url": "https://api-docs.deepseek.com", "encrypted_content": "sd918tZXSvVp", } ], }, ] # One round: the search answered inside the response, so nothing was sent back. assert len(client.calls) == 1 async def test_search_only_round_is_not_an_empty_completion() -> None: """A search turn with no text yet still produced payload. A budget-capped or ``max_uses_exceeded`` search can end the turn on its blocks alone; calling that an empty completion failed a round the model had actually used. """ client = ScriptedMessagesClient(stream_events=[_search_stream_events(text=False)]) events = [ event async for event in run_chat_loop( cast("Any", client), [UserChatMessage(content="what is DeepSeek-V4?")], model="deepseek-flash", stream=True, ) ] assistant = next(event.message for event in events if isinstance(event, AssistantMessageEvent)) assert assistant.content is None assert assistant.provider_blocks is not UNSET assert len(assistant.provider_blocks) == 3 async def test_dispatch_messages_sends_provider_tools_as_server_tools() -> None: """Hosted tools reach the Anthropic request in that surface's own shape.""" client = ScriptedMessagesClient(non_stream=[_message_response(text="ok")]) param = ChatCompletionsParam(model="deepseek-flash", messages=[UserChatMessage(content="hi")]) _ = await dispatch_messages( cast("Any", client), param, provider_tools=[{"type": "web_search", "max_uses": 3}], ) assert client.calls[0].tools == [{"type": "web_search_20250305", "name": "web_search", "max_uses": 3}] async def test_dispatch_messages_without_provider_tools_sends_no_server_tool() -> None: client = ScriptedMessagesClient(non_stream=[_message_response(text="ok")]) param = ChatCompletionsParam(model="deepseek-flash", messages=[UserChatMessage(content="hi")]) _ = await dispatch_messages(cast("Any", client), param, provider_tools=[]) assert client.calls[0].tools is UNSET async def test_search_turn_reports_the_prompt_the_delta_counts() -> None: """A search turn's real prompt size only shows up on ``message_delta``. DeepSeek's compat opens with the question alone (``message_start`` sees 165 tokens) and reports what the search then pulled in on the delta (12468); keeping the smaller one leaves the CLI's context gauge — and the auto-compact threshold — that far short of what the model was sent. """ events: list[AnthropicStreamEvent] = [ AnthropicMessageStart( message=AnthropicMessageResponse( id="msg_s", model="deepseek-flash", content=[], usage=AnthropicUsage(input_tokens=165, output_tokens=0), ) ), AnthropicContentBlockStart(index=0, content_block=AnthropicRawContentBlock(type="text", text="")), AnthropicContentBlockDelta(index=0, delta=AnthropicRawContentBlock(type="text_delta", text="done")), AnthropicMessageDelta( delta={"stop_reason": "end_turn"}, usage=AnthropicUsage(input_tokens=12468, output_tokens=198), ), AnthropicMessageStop(), ] client = ScriptedMessagesClient(stream_events=[events]) param = ChatCompletionsParam(model="deepseek-flash", messages=[UserChatMessage(content="hi")]) stream = await dispatch_messages(cast("Any", client), param, stream=True) chunks = [chunk async for chunk in cast("Any", stream)] usage = next(getattr(chunk, "usage", None) for chunk in chunks if getattr(chunk, "usage", None)) assert isinstance(usage, dict) assert usage["prompt_tokens"] == 12468 assert usage["completion_tokens"] == 198 def _thinking_param(model: str = "claude-opus-5-5", effort: str = "high") -> ChatCompletionsParam: return ChatCompletionsParam( model=model, messages=[UserChatMessage(content="hi")], reasoning_effort=cast("Any", effort), ) async def test_dispatch_messages_retries_with_adaptive_thinking() -> None: client = ScriptedMessagesClient(non_stream=[_message_response(text="done")], refusals=("enabled",)) result = await dispatch_messages(cast("Any", client), _thinking_param(), stream=False) assert result.choices[0].message.content == "done" assert len(client.calls) == 2 refused, retried = client.calls assert refused.thinking is not UNSET assert refused.thinking.type == "enabled" assert retried.thinking is not UNSET assert retried.thinking.type == "adaptive" assert retried.output_config is not UNSET assert retried.output_config.effort == "high" async def test_dispatch_messages_remembers_the_refusing_model() -> None: client = ScriptedMessagesClient( non_stream=[_message_response(text="one"), _message_response(text="two")], refusals=("enabled",), ) param = _thinking_param() _ = await dispatch_messages(cast("Any", client), param, stream=False) _ = await dispatch_messages(cast("Any", client), param, stream=False) # Only the model's first request pays the refused attempt. assert [call.thinking.type for call in client.calls] == ["enabled", "adaptive", "adaptive"] async def test_dispatch_messages_drops_a_refused_disabled_block() -> None: """A model whose thinking is always on answers ``disabled`` with a 400.""" client = ScriptedMessagesClient( non_stream=[_message_response(text="one"), _message_response(text="two")], refusals=("enabled", "disabled"), ) param = _thinking_param(effort="none") result = await dispatch_messages(cast("Any", client), param, stream=False) assert result.choices[0].message.content == "one" assert len(client.calls) == 2 assert client.calls[0].thinking.type == "disabled" # No block at all: there is nothing left to ask such a model for. assert client.calls[1].thinking is UNSET _ = await dispatch_messages(cast("Any", client), param, stream=False) assert client.calls[2].thinking is UNSET async def test_dispatch_messages_stream_retries_with_adaptive_thinking() -> None: events: list[AnthropicStreamEvent] = [ AnthropicContentBlockStart(index=0, content_block=AnthropicRawContentBlock(type="text", text="")), AnthropicContentBlockDelta(index=0, delta=AnthropicRawContentBlock(type="text_delta", text="streamed")), AnthropicMessageDelta(delta={"stop_reason": "end_turn"}, usage=AnthropicUsage(output_tokens=3)), AnthropicMessageStop(), ] client = ScriptedMessagesClient(stream_events=[events], refusals=("enabled",)) stream = await dispatch_messages(cast("Any", client), _thinking_param(), stream=True) chunks = [chunk async for chunk in cast("Any", stream)] texts = [c.choices[0].delta.content for c in chunks if c.choices and isinstance(c.choices[0].delta.content, str)] assert texts == ["streamed"] assert len(client.calls) == 2 assert client.calls[1].thinking.type == "adaptive" async def test_dispatch_messages_does_not_retry_other_failures() -> None: client = ScriptedMessagesClient(error='messages HTTP 400: {"error":{"message":"max_tokens: required"}}') with pytest.raises(RuntimeError, match="max_tokens"): _ = await dispatch_messages(cast("Any", client), _thinking_param(), stream=False) assert len(client.calls) == 1 async def test_dispatch_messages_does_not_retry_a_request_without_thinking() -> None: client = ScriptedMessagesClient(error=_THINKING_REFUSED_400["enabled"]) param = ChatCompletionsParam(model="claude-opus-5-5", messages=[UserChatMessage(content="hi")]) with pytest.raises(RuntimeError, match="adaptive thinking"): _ = await dispatch_messages(cast("Any", client), param, stream=False) assert len(client.calls) == 1