[None][fix] Anthropic Messages: sign the thinking blocks the server generates

tensorrt_llm/serve/anthropic_adapter.py

@@ -24,6 +24,8 @@
 chat templates, or the engine.
 """
 
+import base64
+import hashlib
 import json
 import re
 import traceback
@@ -50,6 +52,7 @@
     AnthropicMessagesResponse,
     AnthropicMessageStartEvent,
     AnthropicMessageStopEvent,
+    AnthropicSignatureDelta,
     AnthropicStopReason,
     AnthropicTextBlock,
     AnthropicTextDelta,
@@ -612,6 +615,21 @@ def convert_usage(usage: Optional[UsageInfo]) -> AnthropicUsage:
     return anthropic_usage
 
 
+def thinking_signature(thinking: str) -> str:
+    """The signature of a thinking block this server generated.
+
+    Anthropic clients replay a thinking block in later turns only together
+    with its signature, and some turn an unsigned one into a text block -- the
+    pi coding agent does -- so the model then reads its own reasoning back as
+    an answer it gave. Anthropic signs with an encrypted proof of origin; this
+    signature is only the base64 SHA-256 digest of the text. It proves nothing
+    and the server does not check it on replay. It exists so that a thinking
+    block round-trips as thinking.
+    """
+    digest = hashlib.sha256(thinking.encode("utf-8")).digest()
+    return base64.b64encode(digest).decode("ascii")
+
+
 def convert_chat_response(chat_response: ChatCompletionResponse) -> AnthropicMessagesResponse:
     """Translate a non-streaming chat completion into an Anthropic message."""
     content: List[Any] = []
@@ -623,7 +641,9 @@ def convert_chat_response(chat_response: ChatCompletionResponse) -> AnthropicMes
         message = choice.message
         reasoning = message.reasoning_content or message.reasoning
         if reasoning:
-            content.append(AnthropicThinkingBlock(thinking=reasoning))
+            content.append(
+                AnthropicThinkingBlock(thinking=reasoning, signature=thinking_signature(reasoning))
+            )
         if message.content and message.content.strip():
             content.append(AnthropicTextBlock(text=message.content))
         for tool_call in message.tool_calls:
@@ -723,6 +743,8 @@ def __init__(self, model: str):
         self.message_started = False
         self.block_index = -1
         self.open_block_type: Optional[str] = None
+        # Text of the open thinking block, signed when the block closes.
+        self.thinking_parts: List[str] = []
         # Upstream index of the tool call whose block is currently open.
         self.open_tool_key: Optional[Any] = None
         # Per-upstream-index tool call state. Insertion order is preserved and
@@ -750,10 +772,23 @@ def __init__(self, model: str):
     def _close_block(self) -> List[str]:
         if self.open_block_type is None:
             return []
-        event = AnthropicContentBlockStopEvent(index=self.block_index)
+        frames = []
+        if self.open_block_type == "thinking":
+            # As Anthropic streams it: the signature is the thinking block's
+            # last delta, once its whole text is known.
+            signature = thinking_signature("".join(self.thinking_parts))
+            self.thinking_parts = []
+            frames.append(
+                anthropic_sse(
+                    AnthropicContentBlockDeltaEvent(
+                        index=self.block_index, delta=AnthropicSignatureDelta(signature=signature)
+                    )
+                )
+            )
+        frames.append(anthropic_sse(AnthropicContentBlockStopEvent(index=self.block_index)))
         self.open_block_type = None
         self.open_tool_key = None
-        return [anthropic_sse(event)]
+        return frames
 
     def _open_block(self, block: Any, block_type: str) -> List[str]:
         if block_type != "text":
@@ -780,7 +815,11 @@ def _ensure_block(self, block_type: str) -> List[str]:
         if block_type == "text":
             frames.extend(self._open_block(AnthropicTextBlock(text=""), "text"))
         elif block_type == "thinking":
-            frames.extend(self._open_block(AnthropicThinkingBlock(thinking=""), "thinking"))
+            # The signature arrives as the block's last delta; the opening block
+            # carries it empty, as a thinking block always has the field.
+            frames.extend(
+                self._open_block(AnthropicThinkingBlock(thinking="", signature=""), "thinking")
+            )
         else:
             raise ValueError(f"unexpected block type {block_type}")
         return frames
@@ -939,6 +978,7 @@ def process_chunk(self, chunk: ChatCompletionStreamResponse) -> List[str]:
             reasoning_delta = delta.reasoning_content or delta.reasoning
             if reasoning_delta:
                 frames.extend(self._ensure_block("thinking"))
+                self.thinking_parts.append(reasoning_delta)
                 frames.append(
                     anthropic_sse(
                         AnthropicContentBlockDeltaEvent(

tests/unittest/llmapi/apps/test_anthropic_adapter.py

@@ -19,6 +19,8 @@
 """
 
 import asyncio
+import base64
+import hashlib
 import json
 
 import pytest
@@ -53,6 +55,11 @@
 MODEL = "test-model"
 
 
+def thinking_signature(thinking: str) -> str:
+    """The signature the server is expected to put on a thinking block."""
+    return base64.b64encode(hashlib.sha256(thinking.encode("utf-8")).digest()).decode("ascii")
+
+
 def make_request(**overrides) -> AnthropicMessagesRequest:
     payload = {
         "model": MODEL,
@@ -749,6 +756,18 @@ def test_reasoning_becomes_thinking_block():
     assert resp.content[1].type == "text"
 
 
+def test_thinking_block_is_signed():
+    """Clients replay a thinking block as thinking only when it carries a signature.
+
+    Without one, the pi coding agent sends the block back as text, and the
+    model then reads its reasoning as an answer it gave.
+    """
+    message = ChatMessage(role="assistant", content="answer", reasoning_content="step by step")
+    resp = convert_chat_response(make_chat_response(message))
+    assert resp.content[0].signature == thinking_signature("step by step")
+    assert resp.model_dump(exclude_none=True)["content"][0]["signature"]
+
+
 def test_empty_content_gets_placeholder_text_block():
     resp = convert_chat_response(make_chat_response(ChatMessage(role="assistant", content=None)))
     assert len(resp.content) == 1
@@ -1191,6 +1210,100 @@ def test_stream_thinking_then_text():
     assert starts == ["thinking", "text"]
 
 
+def test_stream_thinking_block_ends_with_its_signature():
+    events = run_reframer(
+        [
+            chunk({"reasoning_content": "first "}),
+            chunk({"reasoning_content": "second"}),
+            chunk({"content": "answer"}),
+            chunk({}, finish_reason="stop"),
+        ]
+    )
+    assert_event_invariants(events)
+    thinking = [
+        (name, payload)
+        for name, payload in events
+        if name in ("content_block_delta", "content_block_stop") and payload["index"] == 0
+    ]
+    assert [name for name, _ in thinking] == [
+        "content_block_delta",
+        "content_block_delta",
+        "content_block_delta",
+        "content_block_stop",
+    ]
+    assert thinking[2][1]["delta"] == {
+        "type": "signature_delta",
+        "signature": thinking_signature("first second"),
+    }
+    signed = [
+        p
+        for n, p in events
+        if n == "content_block_delta" and p["delta"]["type"] == "signature_delta"
+    ]
+    assert len(signed) == 1, "only the thinking block is signed"
+
+
+def test_stream_thinking_block_opens_with_an_empty_signature():
+    """A thinking block always has a signature field, even before its last delta fills it."""
+    events = run_reframer(
+        [
+            chunk({"reasoning_content": "first"}),
+            chunk({"content": "answer"}),
+            chunk({}, finish_reason="stop"),
+        ]
+    )
+    assert_event_invariants(events)
+    starts = [p["content_block"] for n, p in events if n == "content_block_start"]
+    assert starts[0] == {"type": "thinking", "thinking": "", "signature": ""}
+    assert "signature" not in starts[1]
+
+
+def test_stream_thinking_left_open_is_signed_when_the_stream_ends():
+    events = run_reframer(
+        [
+            chunk({"reasoning_content": "unfinished"}),
+            chunk({}, finish_reason="length"),
+        ]
+    )
+    assert_event_invariants(events)
+    deltas = [p["delta"] for n, p in events if n == "content_block_delta"]
+    assert deltas[-1] == {"type": "signature_delta", "signature": thinking_signature("unfinished")}
+
+
+def test_stream_error_signs_the_open_thinking_block_before_the_error():
+    reframer = AnthropicStreamReframer(model=MODEL)
+    frames = reframer.process_chunk(chunk({"reasoning_content": "partial"}))
+    frames += reframer.error("engine failed")
+    events = parse_frames(frames)
+    (delta_name, delta), (stop_name, stop), (error_name, error) = events[-3:]
+    assert [delta_name, stop_name, error_name] == [
+        "content_block_delta",
+        "content_block_stop",
+        "error",
+    ]
+    assert delta["delta"] == {"type": "signature_delta", "signature": thinking_signature("partial")}
+    assert stop["index"] == delta["index"]
+    assert error["error"]["message"] == "engine failed"
+
+
+def test_stream_each_thinking_block_is_signed_with_its_own_text():
+    events = run_reframer(
+        [
+            chunk({"reasoning_content": "one"}),
+            chunk({"content": "text"}),
+            chunk({"reasoning_content": "two"}),
+            chunk({}, finish_reason="stop"),
+        ]
+    )
+    assert_event_invariants(events)
+    signatures = [
+        p["delta"]["signature"]
+        for n, p in events
+        if n == "content_block_delta" and p["delta"]["type"] == "signature_delta"
+    ]
+    assert signatures == [thinking_signature("one"), thinking_signature("two")]
+
+
 def test_stream_empty_generation_still_valid():
     events = run_reframer([chunk({}, finish_reason="stop")])
     assert_event_invariants(events)